From 9556ba0726319f1f40303fccb27ccbd7a5c25006 Mon Sep 17 00:00:00 2001 From: lesj0610 Date: Fri, 17 Jul 2026 15:26:26 +0900 Subject: [PATCH 01/20] [Bugfix] Stop forcing deprecated use_fast for Qwen3.5 Signed-off-by: lesj0610 --- .../multimodal/processing/test_qwen3_5.py | 41 +++++++++++++++++++ vllm/model_executor/models/qwen3_5.py | 9 ++++ 2 files changed, 50 insertions(+) create mode 100644 tests/models/multimodal/processing/test_qwen3_5.py diff --git a/tests/models/multimodal/processing/test_qwen3_5.py b/tests/models/multimodal/processing/test_qwen3_5.py new file mode 100644 index 000000000000..72b0ad396c2d --- /dev/null +++ b/tests/models/multimodal/processing/test_qwen3_5.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from unittest.mock import MagicMock + +import pytest +from transformers.models.qwen3_vl import Qwen3VLProcessor + +from vllm.model_executor.models.qwen3_5 import ( + Qwen3_5MoeProcessingInfo, + Qwen3_5ProcessingInfo, +) + + +@pytest.mark.parametrize( + "info_cls", + [Qwen3_5ProcessingInfo, Qwen3_5MoeProcessingInfo], +) +@pytest.mark.parametrize( + ("processor_kwargs", "expected_kwargs"), + [ + ({}, {}), + ({"use_fast": True}, {"use_fast": True}), + ({"use_fast": False}, {"use_fast": False}), + ({"use_fast": None}, {"use_fast": None}), + ({"max_pixels": 1024}, {"max_pixels": 1024}), + ], +) +def test_qwen3_5_processor_does_not_force_deprecated_use_fast( + info_cls, + processor_kwargs: dict[str, object], + expected_kwargs: dict[str, object], +) -> None: + ctx = MagicMock() + hf_processor = MagicMock(spec=Qwen3VLProcessor) + ctx.get_hf_processor.return_value = hf_processor + + info = info_cls(ctx) + + assert info.get_hf_processor(**processor_kwargs) is hf_processor + ctx.get_hf_processor.assert_called_once_with(Qwen3VLProcessor, **expected_kwargs) diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index a58c3c4dd718..c7bd3c0997e9 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -28,6 +28,7 @@ import torch from torch import nn +from transformers.models.qwen3_vl import Qwen3VLProcessor from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.decorators import support_torch_compile @@ -100,11 +101,19 @@ class Qwen3_5ProcessingInfo(Qwen3VLProcessingInfo): + def get_hf_processor(self, **kwargs: object) -> Qwen3VLProcessor: + # A top-level backend selector is also forwarded to the video processor. + # Let Transformers select Qwen's default image backend when none is given. + return self.ctx.get_hf_processor(Qwen3VLProcessor, **kwargs) + def get_hf_config(self): return self.ctx.get_hf_config(Qwen3_5Config) class Qwen3_5MoeProcessingInfo(Qwen3VLProcessingInfo): + def get_hf_processor(self, **kwargs: object) -> Qwen3VLProcessor: + return self.ctx.get_hf_processor(Qwen3VLProcessor, **kwargs) + def get_hf_config(self): return self.ctx.get_hf_config(Qwen3_5MoeConfig) From e0301afb2a1364fe594a52d9ce649264a62fa243 Mon Sep 17 00:00:00 2001 From: shaohuaxi Date: Tue, 15 Sep 2026 23:07:06 -0700 Subject: [PATCH 02/20] [Bugfix][Responses] Clean up MCP tool sessions once before closing them (#56988) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 子华 Co-authored-by: Codex --- .../responses/test_serving_responses.py | 84 ++++++++++++++++++- vllm/entrypoints/openai/responses/context.py | 8 +- 2 files changed, 90 insertions(+), 2 deletions(-) diff --git a/tests/entrypoints/openai/responses/test_serving_responses.py b/tests/entrypoints/openai/responses/test_serving_responses.py index 041e517e29dd..bd41b8cda282 100644 --- a/tests/entrypoints/openai/responses/test_serving_responses.py +++ b/tests/entrypoints/openai/responses/test_serving_responses.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from contextlib import AsyncExitStack +from contextlib import AsyncExitStack, asynccontextmanager from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock @@ -41,6 +41,7 @@ from vllm.entrypoints.openai.responses.context import ( ConversationContext, HarmonyContext, + ParsableContext, SimpleContext, ) from vllm.entrypoints.openai.responses.protocol import ( @@ -551,6 +552,87 @@ def test_responses_render_result_rejects_non_single_prompt(engine_inputs): class TestInitializeToolSessions: """Test class for _initialize_tool_sessions method""" + @pytest.fixture(params=[ParsableContext, HarmonyContext]) + def tool_context(self, request): + context = request.param.__new__(request.param) + context.available_tools = ["browser", "python"] + context._tool_sessions = {} + context.called_tools = set() + return context + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("available_tools", "called_tools"), + [ + ([], []), + (["browser"], ["browser"]), + (["browser", "python"], []), + (["browser", "python"], ["browser"]), + (["browser", "python"], ["python"]), + (["browser", "python"], ["browser", "python"]), + ], + ids=["empty", "single", "unused", "first", "last", "both"], + ) + async def test_mcp_sessions_cleaned_once_before_close( + self, tool_context, available_tools, called_tools + ): + events = [] + sessions = {} + + @asynccontextmanager + async def new_session(name, *_args): + session = MagicMock() + session.call_tool = AsyncMock( + side_effect=lambda *_: events.append(("cleanup", name)) + ) + sessions[name] = session + try: + yield session + finally: + events.append(("close", name)) + + tool_server = MagicMock(spec=ToolServer) + tool_server.new_session.side_effect = new_session + tool_context.available_tools = available_tools + async with AsyncExitStack() as stack: + await tool_context.init_tool_sessions(tool_server, stack, "req", {}) + # Reinitializing existing sessions must not duplicate cleanup either. + await tool_context.init_tool_sessions(tool_server, stack, "req", {}) + tool_context.called_tools.update(called_tools) + + for name, session in sessions.items(): + if name in called_tools: + session.call_tool.assert_awaited_once_with("cleanup_session", {}) + assert events.index(("cleanup", name)) < events.index(("close", name)) + else: + session.call_tool.assert_not_awaited() + assert [name for event, name in events if event == "close"] == list( + reversed(available_tools) + ) + + @pytest.mark.asyncio + async def test_mcp_partial_initialization_closes_opened_session(self, tool_context): + closed = [] + session = MagicMock(call_tool=AsyncMock()) + + @asynccontextmanager + async def new_session(name, *_args): + if name == "python": + raise RuntimeError("session initialization failed") + try: + yield session + finally: + closed.append(name) + + tool_server = MagicMock(spec=ToolServer) + tool_server.new_session.side_effect = new_session + with pytest.raises(RuntimeError, match="session initialization failed"): + async with AsyncExitStack() as stack: + await tool_context.init_tool_sessions(tool_server, stack, "req", {}) + + assert closed == ["browser"] + session.call_tool.assert_not_awaited() + @pytest_asyncio.fixture async def serving_responses_instance(self): """Create a real OpenAIServingResponses instance for testing""" diff --git a/vllm/entrypoints/openai/responses/context.py b/vllm/entrypoints/openai/responses/context.py index 2268bb0700a0..6b93d39b8803 100644 --- a/vllm/entrypoints/openai/responses/context.py +++ b/vllm/entrypoints/openai/responses/context.py @@ -565,6 +565,7 @@ async def init_tool_sessions( mcp_tools: dict[str, Mcp], ): if tool_server: + initialized_session = False for tool_name in self.available_tools: if tool_name in self._tool_sessions: continue @@ -577,6 +578,8 @@ async def init_tool_sessions( tool_server.new_session(tool_name, request_id, headers) ) self._tool_sessions[tool_name] = tool_session + initialized_session = True + if initialized_session: exit_stack.push_async_exit(self.cleanup_session) async def cleanup_session(self, *args, **kwargs) -> None: @@ -863,6 +866,7 @@ async def init_tool_sessions( mcp_tools: dict[str, Mcp], ): if tool_server: + initialized_session = False for tool_name in self.available_tools: if tool_name not in self._tool_sessions: tool_type = _map_tool_name_to_tool_type(tool_name) @@ -873,7 +877,9 @@ async def init_tool_sessions( tool_server.new_session(tool_name, request_id, headers) ) self._tool_sessions[tool_name] = tool_session - exit_stack.push_async_exit(self.cleanup_session) + initialized_session = True + if initialized_session: + exit_stack.push_async_exit(self.cleanup_session) async def call_container_tool( self, tool_session: Union["ClientSession", Tool], last_msg: Message From 0f8fa53acf43a9a6a7a1f7e6b196950c78cd6d7a Mon Sep 17 00:00:00 2001 From: AlexHuang Date: Wed, 16 Sep 2026 14:17:50 +0800 Subject: [PATCH 03/20] [Bugfix][KV Offload] Check cgroup memory before SHM allocation (#54014) Signed-off-by: Alex Signed-off-by: AlexHuang Signed-off-by: Alex Co-authored-by: OpenAI Codex Co-authored-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> --- tests/distributed/test_shm_broadcast.py | 35 +++- tests/utils_/test_cpu_resource_utils.py | 189 ++++++++++++++++++ .../cpu/test_shared_offload_region.py | 11 +- .../device_communicators/shm_broadcast.py | 40 ++-- .../model_loader/weight_utils.py | 8 +- vllm/utils/cpu_resource_utils.py | 89 ++++++++- .../kv_offload/cpu/shared_offload_region.py | 5 +- 7 files changed, 343 insertions(+), 34 deletions(-) create mode 100644 tests/utils_/test_cpu_resource_utils.py diff --git a/tests/distributed/test_shm_broadcast.py b/tests/distributed/test_shm_broadcast.py index c3833b6b9f82..ef7f523b78dc 100644 --- a/tests/distributed/test_shm_broadcast.py +++ b/tests/distributed/test_shm_broadcast.py @@ -694,14 +694,43 @@ def test_check_shm_free_space_raises_when_insufficient(tmp_path): def test_check_shm_free_space_passes_when_sufficient(tmp_path): - with mock.patch.object( - shm_broadcast.shutil, "disk_usage", return_value=_fake_disk_usage(512 << 20) + with ( + mock.patch.object( + shm_broadcast.shutil, + "disk_usage", + return_value=_fake_disk_usage(512 << 20), + ), + mock.patch.object(shm_broadcast, "check_cgroup_memory_available"), ): check_shm_free_space(240 << 20, shm_path=str(tmp_path)) def test_check_shm_free_space_skipped_when_path_missing(tmp_path): - check_shm_free_space(1 << 60, shm_path=str(tmp_path / "does-not-exist")) + with mock.patch.object(shm_broadcast, "check_cgroup_memory_available"): + check_shm_free_space(1 << 60, shm_path=str(tmp_path / "does-not-exist")) + + +def test_check_shm_free_space_checks_cgroup(tmp_path): + with ( + mock.patch.object( + shm_broadcast.shutil, + "disk_usage", + return_value=_fake_disk_usage(512 << 20), + ), + mock.patch.object( + shm_broadcast, "check_cgroup_memory_available" + ) as check_cgroup, + ): + check_shm_free_space( + 240 << 20, + shm_path=str(tmp_path), + allocation_name="SHM mmap", + ) + + check_cgroup.assert_called_once_with( + 240 << 20, + "SHM mmap", + ) def test_shm_ring_buffer_creation_checks_free_space(): diff --git a/tests/utils_/test_cpu_resource_utils.py b/tests/utils_/test_cpu_resource_utils.py new file mode 100644 index 000000000000..ff314d2231d2 --- /dev/null +++ b/tests/utils_/test_cpu_resource_utils.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cgroup memory readers and SHM allocation preflight checks.""" + +import sys +from io import StringIO +from unittest import mock + +import pytest + +from vllm.utils import cpu_resource_utils as cru +from vllm.utils.mem_constants import GiB_bytes + +_V2_LIMIT_PATH = "/sys/fs/cgroup/memory.max" +_V2_USAGE_PATH = "/sys/fs/cgroup/memory.current" +_V1_LIMIT_PATH = "/sys/fs/cgroup/memory/memory.limit_in_bytes" +_V1_USAGE_PATH = "/sys/fs/cgroup/memory/memory.usage_in_bytes" + + +@pytest.fixture(autouse=True) +def _clear_cgroup_cache(): + cru.get_cgroup_memory_limit.cache_clear() + yield + cru.get_cgroup_memory_limit.cache_clear() + + +def _stub_files(monkeypatch, files: dict): + """Stub ``open()`` for a fixed set of paths; ``None`` -> OSError.""" + real_open = open + + def fake_open(path, *args, **kwargs): + if path in files: + content = files[path] + if content is None: + raise OSError(f"no such file: {path}") + return StringIO(content) + return real_open(path, *args, **kwargs) + + monkeypatch.setattr("builtins.open", fake_open) + + +def test_cgroup_v2_limit_and_usage(monkeypatch): + monkeypatch.setattr(sys, "platform", "linux") + _stub_files( + monkeypatch, + { + _V2_LIMIT_PATH: f"{20 * GiB_bytes}\n", + _V2_USAGE_PATH: f"{5 * GiB_bytes}\n", + }, + ) + + assert cru.get_cgroup_memory_limit() == 20 * GiB_bytes + assert cru.get_cgroup_memory_usage() == 5 * GiB_bytes + + +def test_cgroup_v2_unlimited_falls_back_to_v1(monkeypatch): + monkeypatch.setattr(sys, "platform", "linux") + _stub_files( + monkeypatch, + { + _V2_LIMIT_PATH: "max\n", + _V1_LIMIT_PATH: f"{8 * GiB_bytes}\n", + _V1_USAGE_PATH: f"{1 * GiB_bytes}\n", + }, + ) + + assert cru.get_cgroup_memory_limit() == 8 * GiB_bytes + assert cru.get_cgroup_memory_usage() == 1 * GiB_bytes + + +def test_cgroup_v1_limit_and_usage(monkeypatch): + monkeypatch.setattr(sys, "platform", "linux") + _stub_files( + monkeypatch, + { + _V2_LIMIT_PATH: None, + _V1_LIMIT_PATH: f"{8 * GiB_bytes}\n", + _V1_USAGE_PATH: f"{2 * GiB_bytes}\n", + }, + ) + + assert cru.get_cgroup_memory_limit() == 8 * GiB_bytes + assert cru.get_cgroup_memory_usage() == 2 * GiB_bytes + + +def test_cgroup_v1_unlimited_sentinel_is_ignored(monkeypatch): + """An unlimited cgroup v1 sentinel must not be treated as a real limit.""" + monkeypatch.setattr(sys, "platform", "linux") + _stub_files( + monkeypatch, + {_V2_LIMIT_PATH: None, _V1_LIMIT_PATH: f"{(1 << 63) - 1}\n"}, + ) + + assert cru.get_cgroup_memory_limit() is None + assert cru.get_cgroup_memory_usage() is None + + +def test_cgroup_no_limit_files_present(monkeypatch): + monkeypatch.setattr(sys, "platform", "linux") + _stub_files(monkeypatch, {_V2_LIMIT_PATH: None, _V1_LIMIT_PATH: None}) + + assert cru.get_cgroup_memory_limit() is None + assert cru.get_cgroup_memory_usage() is None + + +def test_cgroup_skipped_on_non_linux(monkeypatch): + monkeypatch.setattr(sys, "platform", "darwin") + _stub_files(monkeypatch, {_V2_LIMIT_PATH: f"{1 * GiB_bytes}\n"}) + + assert cru.get_cgroup_memory_limit() is None + assert cru.get_cgroup_memory_usage() is None + + +def test_cgroup_limit_is_cached_until_explicitly_cleared(monkeypatch): + monkeypatch.setattr(sys, "platform", "linux") + files = { + _V2_LIMIT_PATH: f"{1 * GiB_bytes}\n", + _V2_USAGE_PATH: f"{512 << 20}\n", + } + _stub_files(monkeypatch, files) + + assert cru.get_cgroup_memory_limit() == 1 * GiB_bytes + files[_V2_LIMIT_PATH] = f"{2 * GiB_bytes}\n" + assert cru.get_cgroup_memory_limit() == 1 * GiB_bytes + + cru.get_cgroup_memory_limit.cache_clear() + assert cru.get_cgroup_memory_limit() == 2 * GiB_bytes + + +def test_cgroup_usage_is_read_without_cache(monkeypatch): + monkeypatch.setattr(sys, "platform", "linux") + files = { + _V2_LIMIT_PATH: f"{1 * GiB_bytes}\n", + _V2_USAGE_PATH: f"{512 << 20}\n", + } + _stub_files(monkeypatch, files) + + assert cru.get_cgroup_memory_usage() == 512 << 20 + files[_V2_USAGE_PATH] = f"{768 << 20}\n" + assert cru.get_cgroup_memory_usage() == 768 << 20 + + +def test_check_cgroup_memory_available_warns_on_low_headroom(monkeypatch): + monkeypatch.setattr(cru, "get_cgroup_memory_limit", lambda: 1 << 30) + monkeypatch.setattr(cru, "get_cgroup_memory_usage", lambda: 512 << 20) + + with ( + mock.patch.object(cru.logger, "debug") as log_debug, + mock.patch.object(cru.logger, "warning") as log_warning, + ): + cru.check_cgroup_memory_available(600 << 20, "mmap") + + log_debug.assert_not_called() + log_warning.assert_called_once() + assert "current headroom is below" in log_warning.call_args.args[-1] + + +def test_check_cgroup_memory_available_logs_success_at_debug(monkeypatch): + monkeypatch.setattr(cru, "get_cgroup_memory_limit", lambda: 1 << 30) + monkeypatch.setattr(cru, "get_cgroup_memory_usage", lambda: 512 << 20) + + with ( + mock.patch.object(cru.logger, "debug") as log_debug, + mock.patch.object(cru.logger, "warning") as log_warning, + ): + cru.check_cgroup_memory_available(256 << 20, "mmap") + + log_debug.assert_called_once() + log_warning.assert_not_called() + + +@pytest.mark.parametrize( + ("limit", "usage"), + [(None, 512 << 20), (1 << 30, None)], +) +def test_check_cgroup_memory_available_skips_without_snapshot( + monkeypatch, limit, usage +): + monkeypatch.setattr(cru, "get_cgroup_memory_limit", lambda: limit) + monkeypatch.setattr(cru, "get_cgroup_memory_usage", lambda: usage) + + with ( + mock.patch.object(cru.logger, "debug") as log_debug, + mock.patch.object(cru.logger, "warning") as log_warning, + ): + cru.check_cgroup_memory_available(1 << 60, "mmap") + + log_debug.assert_not_called() + log_warning.assert_not_called() diff --git a/tests/v1/kv_offload/cpu/test_shared_offload_region.py b/tests/v1/kv_offload/cpu/test_shared_offload_region.py index c448e76a9c51..4f22d122dbec 100644 --- a/tests/v1/kv_offload/cpu/test_shared_offload_region.py +++ b/tests/v1/kv_offload/cpu/test_shared_offload_region.py @@ -906,12 +906,13 @@ def test_insufficient_space_raises_clear_error(monkeypatch): monkeypatch.setattr(region.os, "open", mock_open) monkeypatch.setattr(region.os, "unlink", mock_unlink) monkeypatch.setattr(region.os, "close", mock_close) + mock_check = MagicMock( + side_effect=RuntimeError("Insufficient space in /dev/shm: 30 GB required.") + ) monkeypatch.setattr( region, "check_shm_free_space", - lambda *a, **kw: (_ for _ in ()).throw( - RuntimeError("Insufficient space in /dev/shm: 30 GB required.") - ), + mock_check, ) with pytest.raises(RuntimeError, match="Insufficient space"): @@ -925,6 +926,10 @@ def test_insufficient_space_raises_clear_error(monkeypatch): mock_unlink.assert_called_once_with(mmap_path) mock_close.assert_called_once_with(9999) + mock_check.assert_called_once_with( + 4 * PAGE_SIZE, + allocation_name="CPU KV offload shared region in /dev/shm", + ) def test_ftruncate_failure_cleans_up_creator(monkeypatch): diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index 6caacea6ec5d..2171f79fb7ce 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -34,6 +34,7 @@ from vllm.distributed.utils import StatelessProcessGroup, sched_yield from vllm.logger import init_logger from vllm.platforms import current_platform +from vllm.utils.cpu_resource_utils import check_cgroup_memory_available from vllm.utils.network_utils import ( get_ip, get_open_zmq_inproc_path, @@ -224,26 +225,37 @@ def notify(self): SHM_PATH = "/dev/shm" -def check_shm_free_space(required_bytes: int, shm_path: str = SHM_PATH) -> None: - """Raise if ``shm_path`` cannot fit a ``required_bytes`` shared segment. +def check_shm_free_space( + required_bytes: int, + shm_path: str = SHM_PATH, + *, + allocation_name: str = "shared-memory allocation", +) -> None: + """Raise if SHM cannot fit a shared segment and log cgroup headroom. Args: required_bytes: Size of the shared-memory segment to be created. - shm_path: Mount point backing POSIX shared memory; skipped if absent. + shm_path: Mount point backing POSIX shared memory; its filesystem + check is skipped if absent. + allocation_name: Human-readable name used in errors and logs. Raises: - RuntimeError: If ``required_bytes`` exceeds the free space. + RuntimeError: If the SHM filesystem has insufficient space. """ - if not os.path.isdir(shm_path): - return - free_bytes = shutil.disk_usage(shm_path).free - if required_bytes <= free_bytes: - return - mib = 1 << 20 - raise RuntimeError( - f"Insufficient space in {shm_path}: {required_bytes / mib:.0f} MiB " - f"required, {free_bytes / mib:.0f} MiB free. Increase {shm_path} " - "(e.g. --shm-size or --ipc=host)." + if os.path.isdir(shm_path): + free_bytes = shutil.disk_usage(shm_path).free + if required_bytes > free_bytes: + mib = 1 << 20 + raise RuntimeError( + f"Insufficient space in {shm_path} for {allocation_name}: " + f"{required_bytes / mib:.0f} MiB required, " + f"{free_bytes / mib:.0f} MiB free. Increase {shm_path} " + "(e.g. --shm-size or --ipc=host)." + ) + + check_cgroup_memory_available( + required_bytes, + allocation_name, ) diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index abe24fa06eca..d291f1929ca0 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -727,11 +727,15 @@ def _get_available_ram_bytes() -> int: host_available = psutil.virtual_memory().available - from vllm.utils.cpu_resource_utils import get_cgroup_memory_limit + from vllm.utils.cpu_resource_utils import ( + get_cgroup_memory_limit, + get_cgroup_memory_usage, + ) - cgroup_limit, cgroup_usage = get_cgroup_memory_limit() + cgroup_limit = get_cgroup_memory_limit() if cgroup_limit is None: return host_available + cgroup_usage = get_cgroup_memory_usage() cgroup_available = ( cgroup_limit if cgroup_usage is None else max(0, cgroup_limit - cgroup_usage) ) diff --git a/vllm/utils/cpu_resource_utils.py b/vllm/utils/cpu_resource_utils.py index 53e74610659b..a1e7ab0ee30d 100644 --- a/vllm/utils/cpu_resource_utils.py +++ b/vllm/utils/cpu_resource_utils.py @@ -12,8 +12,12 @@ import psutil import regex as re +from vllm.logger import init_logger + DEVICE_CONTROL_ENV_VAR = "CPU_VISIBLE_MEMORY_NODES" +logger = init_logger(__name__) + @dataclass class LogicalCPUInfo: @@ -63,21 +67,20 @@ def _read_int_file(path: str) -> int | None: @cache -def get_cgroup_memory_limit() -> tuple[int | None, int | None]: - """Return (limit, usage) in bytes from cgroup, or (None, None). +def get_cgroup_memory_limit() -> int | None: + """Return the cgroup memory limit in bytes, or None. - Supports both cgroup v2 (unified) and v1. Returns (None, None) when + Supports both cgroup v2 (unified) and v1. Returns None when not running under a constrained cgroup (e.g. bare metal, or limit reported as `max`/an unrealistically large value). """ if sys.platform != "linux": - return None, None + return None # cgroup v2 unified hierarchy v2_limit = _read_int_file("/sys/fs/cgroup/memory.max") if v2_limit is not None: - v2_usage = _read_int_file("/sys/fs/cgroup/memory.current") - return v2_limit, v2_usage + return v2_limit # cgroup v1 v1_limit = _read_int_file("/sys/fs/cgroup/memory/memory.limit_in_bytes") @@ -85,11 +88,74 @@ def get_cgroup_memory_limit() -> tuple[int | None, int | None]: # cgroup v1 reports a huge sentinel (close to PAGE_COUNTER_MAX) # when unlimited. Treat absurdly large values as "no limit". if v1_limit >= (1 << 62): - return None, None - v1_usage = _read_int_file("/sys/fs/cgroup/memory/memory.usage_in_bytes") - return v1_limit, v1_usage + return None + return v1_limit + + return None + - return None, None +def get_cgroup_memory_usage() -> int | None: + """Return the current cgroup memory usage in bytes, or None. + + The usage value is intentionally read on every call because cgroup + memory usage changes while the process is running. + """ + if sys.platform != "linux": + return None + + # cgroup v2 unified hierarchy + if _read_int_file("/sys/fs/cgroup/memory.max") is not None: + return _read_int_file("/sys/fs/cgroup/memory.current") + + # cgroup v1 + v1_limit = _read_int_file("/sys/fs/cgroup/memory/memory.limit_in_bytes") + if v1_limit is not None and v1_limit < (1 << 62): + return _read_int_file("/sys/fs/cgroup/memory/memory.usage_in_bytes") + + return None + + +def check_cgroup_memory_available( + required_bytes: int, + allocation_name: str, +) -> None: + """Log cgroup memory headroom for an upcoming allocation. + + Args: + required_bytes: Bytes required by the allocation. + allocation_name: Human-readable name used in log messages. + + Low headroom logs a warning, but does not reject the allocation because + cgroup usage can include reclaimable memory. If the cgroup limit or usage + cannot be read, the check is skipped. + """ + cgroup_limit = get_cgroup_memory_limit() + cgroup_usage = get_cgroup_memory_usage() + if cgroup_limit is None or cgroup_usage is None: + return + + cgroup_available = max(0, cgroup_limit - cgroup_usage) + mib = 1 << 20 + remaining_bytes = cgroup_available - required_bytes + log_fn = logger.debug if remaining_bytes >= 0 else logger.warning + status = ( + "current headroom meets the requested allocation" + if remaining_bytes >= 0 + else "current headroom is below the requested allocation; allocation " + "will still be attempted because cgroup usage may be reclaimable" + ) + log_fn( + "Cgroup memory preflight for %s: %.0f MiB required, %.0f MiB current " + "usage, %.0f MiB available under %.0f MiB limit, %.0f MiB remaining " + "after allocation based on current usage; %s.", + allocation_name, + required_bytes / mib, + cgroup_usage / mib, + cgroup_available / mib, + cgroup_limit / mib, + remaining_bytes / mib, + status, + ) def get_memory_affinity(pid: int = 0) -> list[int]: @@ -161,7 +227,8 @@ def get_memory_node_info(node_id: int = 0) -> MemoryNodeInfo: # would be applied to host RAM instead of the pod's limit. cgroup # does not expose per-NUMA-node limits, so we just clamp the totals # against the pod-wide limit here. - cgroup_limit, cgroup_usage = get_cgroup_memory_limit() + cgroup_limit = get_cgroup_memory_limit() + cgroup_usage = get_cgroup_memory_usage() if cgroup_limit is not None and cgroup_limit < total_memory: total_memory = cgroup_limit cgroup_available = cgroup_limit - (cgroup_usage or 0) diff --git a/vllm/v1/kv_offload/cpu/shared_offload_region.py b/vllm/v1/kv_offload/cpu/shared_offload_region.py index 46a1e7ae6460..284f1de68e8e 100644 --- a/vllm/v1/kv_offload/cpu/shared_offload_region.py +++ b/vllm/v1/kv_offload/cpu/shared_offload_region.py @@ -135,7 +135,10 @@ def __init__( self._creator = True if creator_memory_check is not None: creator_memory_check(self.total_size_bytes) - check_shm_free_space(self.total_size_bytes) + check_shm_free_space( + self.total_size_bytes, + allocation_name="CPU KV offload shared region in /dev/shm", + ) os.ftruncate(self.fd, self.total_size_bytes) logger.info( "Created mmap file %s (%.2f GB)", From 18553c5c9ff02acb84bd0f8698e90609e1aec2e9 Mon Sep 17 00:00:00 2001 From: freyfwt <110148731+freyfwt@users.noreply.github.com> Date: Wed, 16 Sep 2026 14:52:26 +0800 Subject: [PATCH 04/20] [Kernel] Use Murmur3 RNG for Gumbel sampling (#51367) Signed-off-by: freyfwt <110148731+freyfwt@users.noreply.github.com> Signed-off-by: freyfwt Co-authored-by: Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Jiangyun Zhu --- .../llm/test_struct_output_generate.py | 2 +- .../chat_completion/test_include_reasoning.py | 12 +- .../openai/completion/test_completion.py | 12 +- tests/v1/worker/test_gpu_gumbel_sample.py | 108 +++++++++++++++++- vllm/v1/worker/gpu/sample/gumbel.py | 108 ++++++++++++++++-- 5 files changed, 216 insertions(+), 26 deletions(-) diff --git a/tests/entrypoints/llm/test_struct_output_generate.py b/tests/entrypoints/llm/test_struct_output_generate.py index 0f55ad7adc3e..030b0207e4c7 100644 --- a/tests/entrypoints/llm/test_struct_output_generate.py +++ b/tests/entrypoints/llm/test_struct_output_generate.py @@ -948,7 +948,7 @@ def test_structured_output_batched_with_non_structured_outputs_requests( prompts = [structured_outputs_prompt, non_structured_outputs_prompt] sampling_params = [ SamplingParams( - temperature=1.0, + temperature=0, max_tokens=400, structured_outputs=StructuredOutputsParams(json=sample_json_schema), ), diff --git a/tests/entrypoints/openai/chat_completion/test_include_reasoning.py b/tests/entrypoints/openai/chat_completion/test_include_reasoning.py index 1d4e1ebc4055..5eeca438fe64 100644 --- a/tests/entrypoints/openai/chat_completion/test_include_reasoning.py +++ b/tests/entrypoints/openai/chat_completion/test_include_reasoning.py @@ -45,7 +45,8 @@ async def test_include_reasoning_true_non_streaming(client: openai.AsyncOpenAI): response = await client.chat.completions.create( model=MODEL_NAME, messages=MESSAGES, - max_tokens=200, + max_tokens=512, + temperature=0, extra_body={"include_reasoning": True}, ) @@ -63,7 +64,8 @@ async def test_include_reasoning_false_non_streaming(client: openai.AsyncOpenAI) response = await client.chat.completions.create( model=MODEL_NAME, messages=MESSAGES, - max_tokens=200, + max_tokens=512, + temperature=0, extra_body={"include_reasoning": False}, ) @@ -99,7 +101,8 @@ async def test_include_reasoning_true_streaming(client: openai.AsyncOpenAI): stream = await client.chat.completions.create( model=MODEL_NAME, messages=MESSAGES, - max_tokens=200, + max_tokens=512, + temperature=0, stream=True, extra_body={"include_reasoning": True}, ) @@ -130,7 +133,8 @@ async def test_include_reasoning_false_streaming(client: openai.AsyncOpenAI): stream = await client.chat.completions.create( model=MODEL_NAME, messages=MESSAGES, - max_tokens=200, + max_tokens=512, + temperature=0, stream=True, extra_body={"include_reasoning": False}, ) diff --git a/tests/entrypoints/openai/completion/test_completion.py b/tests/entrypoints/openai/completion/test_completion.py index af9f1def3614..1bc6260cdebc 100644 --- a/tests/entrypoints/openai/completion/test_completion.py +++ b/tests/entrypoints/openai/completion/test_completion.py @@ -353,7 +353,7 @@ async def test_parallel_no_streaming(client: openai.AsyncOpenAI, model_name: str prompt = "What is an LLM?" n = 3 - max_tokens = 50 # we want some to finish earlier than others + max_tokens = 50 # High temperature to maximize chance of unique completions. completion = await client.completions.create( @@ -371,16 +371,12 @@ async def test_parallel_no_streaming(client: openai.AsyncOpenAI, model_name: str num_completions = len(completion.choices) assert num_completions == n, f"Num completions {num_completions} but expected {n}." completion_repeats: dict[str, int] = {} - output_token_lengths = set() for idx, choice in enumerate(completion.choices): # Assert correct completion index & some finish reason. assert choice.index == idx, f"Index {choice.index} but expected {idx}." assert choice.finish_reason is not None, "None finish_reason is invalid." text = choice.text completion_repeats[text] = completion_repeats.get(text, 0) + 1 - output_token_lengths.add(len(choice.logprobs.tokens)) - # Assert subrequests finished at different times - assert len(output_token_lengths) > 1 # Assert `n` unique completions num_unique = len(completion_repeats) if num_unique != n: @@ -403,7 +399,7 @@ async def test_parallel_streaming(client: openai.AsyncOpenAI, model_name: str): prompt = "What is an LLM?" n = 3 - max_tokens = 50 # we want some to finish earlier than others + max_tokens = 50 stream = await client.completions.create( model=model_name, @@ -427,19 +423,15 @@ async def test_parallel_streaming(client: openai.AsyncOpenAI, model_name: str): f"Expected {n} completions with valid indices and finish_reason." ) completion_repeats: dict[str, int] = {} - chunk_lengths = set() for chunk in chunks: chunk_len = len(chunk) # Assert correct number of completion tokens - chunk_lengths.add(chunk_len) assert chunk_len <= max_tokens, ( f"max_tokens={max_tokens} but chunk len is {chunk_len}." ) text = "".join(chunk) completion_repeats[text] = completion_repeats.get(text, 0) + 1 print(text) - # Assert subrequests finished at different times - assert len(chunk_lengths) > 1 # Assert `n` unique completions num_unique = len(completion_repeats) if num_unique != n: diff --git a/tests/v1/worker/test_gpu_gumbel_sample.py b/tests/v1/worker/test_gpu_gumbel_sample.py index 10a09b9de4a6..a66a1395d2c7 100644 --- a/tests/v1/worker/test_gpu_gumbel_sample.py +++ b/tests/v1/worker/test_gpu_gumbel_sample.py @@ -21,7 +21,12 @@ 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 +from vllm.triton_utils import tl, triton +from vllm.v1.worker.gpu.sample.gumbel import ( + _uniform64_from_random53, + gumbel_sample, + murmur3_hash32, +) DEVICE = "cuda" VOCAB_SIZE = 200_000 @@ -33,6 +38,103 @@ Z_TOLERANCE = 10.0 +@triton.jit +def _murmur3_known_answer_kernel( + seeds_ptr, + positions_ptr, + offsets_ptr, + hashes_ptr, + domain_hashes_ptr, + numel: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + idx = tl.arange(0, BLOCK_SIZE) + mask = idx < numel + seed = tl.load(seeds_ptr + idx, mask=mask) + pos = tl.load(positions_ptr + idx, mask=mask) + offset = tl.load(offsets_ptr + idx, mask=mask) + hash32 = murmur3_hash32(seed, pos, offset).to(tl.int64) + domain_hash32 = murmur3_hash32(seed, pos, offset, domain=0x9E3779B9).to(tl.int64) + tl.store(hashes_ptr + idx, hash32, mask=mask) + tl.store(domain_hashes_ptr + idx, domain_hash32, mask=mask) + + +@triton.jit +def _uniform64_endpoint_kernel( + random53_ptr, + uniform_ptr, + numel: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + idx = tl.arange(0, BLOCK_SIZE) + mask = idx < numel + random53 = tl.load(random53_ptr + idx, mask=mask).to(tl.uint64) + tl.store(uniform_ptr + idx, _uniform64_from_random53(random53), mask=mask) + + +def test_murmur3_known_answers(): + seeds = torch.tensor( + [0, 1, -1, 0x123456789ABCDEF, -(1 << 63), (1 << 63) - 1], + dtype=torch.int64, + device=DEVICE, + ) + positions = torch.tensor( + [0, 2, 17, 123456789, 1 << 30, (1 << 30) + 7], + dtype=torch.int64, + device=DEVICE, + ) + offsets = torch.tensor( + [0, 3, 1023, 199999, 4096, 0xFFFFFFFF], + dtype=torch.int64, + device=DEVICE, + ) + hashes = torch.empty_like(seeds) + domain_hashes = torch.empty_like(seeds) + + _murmur3_known_answer_kernel[(1,)]( + seeds, + positions, + offsets, + hashes, + domain_hashes, + numel=seeds.numel(), + BLOCK_SIZE=8, + ) + + assert hashes.tolist() == [ + 0x8134CDF8, + 0x23B53F91, + 0xC8C13C14, + 0x5DDE974F, + 0xF1A9355B, + 0xD952F9AA, + ] + assert domain_hashes.tolist() == [ + 0xD28AD962, + 0xE3AAF818, + 0xD223C51C, + 0x0547D2FE, + 0x3AED160F, + 0x49E8EA7F, + ] + + +def test_uniform64_excludes_endpoints(): + random53 = torch.tensor([0, (1 << 53) - 1], dtype=torch.int64, device=DEVICE) + uniform = torch.empty(2, dtype=torch.float64, device=DEVICE) + + _uniform64_endpoint_kernel[(1,)]( + random53, + uniform, + numel=random53.numel(), + BLOCK_SIZE=2, + ) + + assert torch.isfinite(uniform).all() + assert ((uniform > 0.0) & (uniform < 1.0)).all() + assert uniform.tolist() == [2.0**-54, 1.0 - 2.0**-53] + + 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) @@ -180,7 +282,7 @@ def test_full_vocab_distribution_fidelity(): def test_drafting_uses_a_separate_noise_stream(): - """is_drafting salts the Philox offset: same inputs, different draws. + """is_drafting salts the RNG position: same inputs, different draws. The draft proposal and the residual resample after a rejection must be independent. They key noise by the same (seed, pos), so only the salt keeps @@ -189,7 +291,7 @@ def test_drafting_uses_a_separate_noise_stream(): tests/v1/spec_decode/test_rejection_sampler_utils.py for the distributional consequence. - Relocating the offset must not distort the draw either, so both streams are + Salting the position must not distort the draw either, so both streams are checked against the target's far-tail mass, which sits HEAD_LOG_GAP logits below the head -- the regime where fp32 Gumbel precision matters. """ diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index 4bf1f47280cf..722021236926 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/vllm/v1/worker/gpu/sample/gumbel.py @@ -2,10 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch -from vllm.triton_utils import HAS_TRITON, tl, tldevice, triton +from vllm.triton_utils import HAS_TRITON, tl, triton -# Smallest positive value produced by Triton's fp32 `tl.rand`. Used to clamp -# zero draws before the flipped Gumbel transform below. +# Smallest positive value produced by Triton's fp32 `tl.rand`. Used by the +# retained Philox helper for rejection sampling. # # Triton requires globals accessed from `@triton.jit` functions to be wrapped # in `tl.constexpr(...)`. We can only do that when Triton is actually @@ -87,6 +87,97 @@ def tl_rand32(seed, offset, includes_zero: tl.constexpr): return u +@triton.jit +def _murmur3_rotl32(value, shift: tl.constexpr): + return (value << shift) | (value >> (32 - shift)) + + +@triton.jit +def _murmur3_mix(h, key): + key *= 0xCC9E2D51 + key = _murmur3_rotl32(key, 15) + key *= 0x1B873593 + h ^= key + h = _murmur3_rotl32(h, 13) + return h * 5 + 0xE6546B64 + + +@triton.jit +def _murmur3_fmix32(h): + h ^= h >> 16 + h *= 0x85EBCA6B + h ^= h >> 13 + h *= 0xC2B2AE35 + return h ^ (h >> 16) + + +@triton.jit +def murmur3_hash32(seed, pos, offset, domain: tl.constexpr = 0): + seed = seed.to(tl.int64) + pos = pos.to(tl.int64) + offset = offset.to(tl.uint32) + # Keep the request-wide prefix scalar until the token offset is mixed in. + h = (seed ^ seed).to(tl.uint32) + h ^= domain + h = _murmur3_mix(h, (seed & 0xFFFFFFFF).to(tl.uint32)) + h = _murmur3_mix(h, ((seed >> 32) & 0xFFFFFFFF).to(tl.uint32)) + h = _murmur3_mix(h, (pos & 0xFFFFFFFF).to(tl.uint32)) + h = _murmur3_mix(h, offset) + return _murmur3_fmix32(h ^ 16) + + +@triton.jit +def murmur3_uniform32(seed, pos, offset): + random32 = murmur3_hash32(seed, pos, offset) + # Split the uint32 before converting to fp32 so backends without a native + # uint32-to-float conversion can still use all 32 source bits. Both 16-bit + # halves convert exactly; their sum is the correctly rounded fp32 value of + # (random32 + 0.5) * 2**-32. In particular, the u -> 0 winning tail keeps + # the full 32-bit source resolution instead of being truncated to 24 bits. + hi16 = (random32 >> 16).to(tl.int32) + lo16 = (random32 & 0xFFFF).to(tl.int32) + return ( + hi16.to(tl.float32) * 1.52587890625e-05 + + (lo16.to(tl.float32) + 0.5) * 2.3283064365386963e-10 + ) + + +@triton.jit +def _uniform64_from_random53(random53): + uniform = (random53.to(tl.float64) + 0.5) * 1.1102230246251565e-16 + # The largest midpoint rounds to 1.0 in fp64; keep the uniform open without + # relying on a near-one literal that Triton may materialize in fp32. + return tl.where(uniform == 1.0, uniform - 1.1102230246251565e-16, uniform) + + +@triton.jit +def murmur3_uniform64(seed, pos, offset): + lo = murmur3_hash32(seed, pos, offset).to(tl.uint64) + hi = murmur3_hash32(seed, pos, offset, domain=0x9E3779B9).to(tl.uint64) + random53 = ((hi << 32) | lo) >> 11 + return _uniform64_from_random53(random53) + + +@triton.jit +def _log1p_neg_stable(value): + # Preserve precision for the positive Gumbel tail without relying on a + # backend-specific libdevice log1p. The degree-8 series has absolute error + # below 6e-7 on [0, 0.25]; subtraction is well-conditioned elsewhere for + # the part of the distribution that can win the argmax. + polynomial = 1.0 / 8.0 + polynomial = 1.0 / 7.0 + value * polynomial + polynomial = 1.0 / 6.0 + value * polynomial + polynomial = 1.0 / 5.0 + value * polynomial + polynomial = 1.0 / 4.0 + value * polynomial + polynomial = 1.0 / 3.0 + value * polynomial + polynomial = 1.0 / 2.0 + value * polynomial + polynomial = 1.0 + value * polynomial + series = -value * polynomial + + direct = tl.log(tl.maximum(1.0 - value, 5.960464477539063e-08)) + return tl.where(value < 0.25, series, direct) + + @triton.jit def gumbel_noised_argmax( logits, @@ -117,14 +208,15 @@ def gumbel_noised_argmax( if temp != 0.0: if IS_DRAFTING: pos = pos + _DRAFT_NOISE_SALT - gumbel_seed = tl.randint(seed, pos) if USE_FP64: - u = tl_rand64(gumbel_seed, keys, includes_zero=False) + u = murmur3_uniform64(seed, pos, keys) gumbel_noise = -tl.log(-tl.log(u)) else: - u = tl_rand32(gumbel_seed, keys, includes_zero=False) - # log1p keeps the winning tail at u -> 0, where fp32 resolves it. - gumbel_noise = -tl.log(-tldevice.log1p(-u)) + u = murmur3_uniform32(seed, pos, keys) + # Draw the large-noise tail (which decides the argmax winner) from + # u -> 0, where fp32 has fine resolution. Avoid backend-specific + # log1p while preserving precision in the winning tail. + gumbel_noise = -tl.log(-_log1p_neg_stable(u)) logits = tl.where(mask, logits + gumbel_noise, float("-inf")) return tl.max(logits, axis=0, return_indices=True) From f30a195bbb15b920d9c2c40e6a3466d8961ab101 Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Wed, 16 Sep 2026 03:23:20 -0400 Subject: [PATCH 05/20] [Bugfix] Fix incorrect Mamba block allocation estimate that prevents request admission (#57050) Signed-off-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> --- .../core/test_single_type_kv_cache_manager.py | 57 +++++++++++++++++++ vllm/v1/core/single_type_kv_cache_manager.py | 8 +-- 2 files changed, 61 insertions(+), 4 deletions(-) diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index 14e46d7ba07c..f20a96f5cc31 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -179,6 +179,63 @@ def test_mamba_retirement_bounds_prefill_states(block_size, in_flight_chunks): assert pool.get_num_free_blocks() == initial_free +@pytest.mark.parametrize("block_size", [896, 1536]) +@pytest.mark.parametrize("num_speculative_blocks", [0, 1, 4]) +@pytest.mark.parametrize("prompt_tokens", [25121, 704547]) +def test_mamba_checkpoint_admission_matches_allocation( + block_size, num_speculative_blocks, prompt_tokens +): + """Checkpoint admission must match the subsequent physical allocation.""" + spec = MambaSpec( + block_size=block_size, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + num_speculative_blocks=num_speculative_blocks, + num_prefill_checkpoint_blocks=1, + prefill_checkpoint_alignment=64, + ) + pool = BlockPool( + num_gpu_blocks=2048, + enable_caching=True, + hash_block_size=128, + ) + manager = MambaManager( + spec, + block_pool=pool, + enable_caching=True, + kv_cache_group_id=0, + scheduler_block_size=block_size, + ) + request_id = "prefill" + computed_tokens = 23040 + + def estimate(num_tokens, total_computed_tokens, apply_admission_cap): + return manager.get_num_blocks_to_allocate( + request_id=request_id, + num_tokens=num_tokens, + new_computed_blocks=[], + total_computed_tokens=total_computed_tokens, + num_local_computed_tokens=total_computed_tokens, + num_tokens_main_model=num_tokens, + apply_admission_cap=apply_admission_cap, + ) + + estimate(computed_tokens, 0, False) + manager.allocate_new_blocks(request_id, computed_tokens, computed_tokens) + assert request_id in manager._allocated_block_reqs + + admission_estimate = estimate(prompt_tokens, computed_tokens, True) + allocation_estimate = estimate(prompt_tokens, computed_tokens, False) + assert request_id in manager._checkpoints + + free_before = pool.get_num_free_blocks() + manager.allocate_new_blocks(request_id, prompt_tokens, prompt_tokens) + allocated = free_before - pool.get_num_free_blocks() + + assert admission_estimate == allocation_estimate == allocated + + def get_sliding_window_manager( sliding_window_spec, block_pool, diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index fc5dc410485d..08eab7628768 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -1796,10 +1796,10 @@ def get_num_blocks_to_allocate( self._checkpoints.pop(request_id, None) if num_new_blocks > 0: blocks_allocated = request_id in self._allocated_block_reqs - if not (checkpoint_block and blocks_allocated): - num_new_blocks = 1 + int(has_partial_hit) + checkpoint_block - if not blocks_allocated: - num_new_blocks += self.num_speculative_blocks + physical_block_cap = 1 + int(has_partial_hit) + checkpoint_block + if not blocks_allocated or checkpoint_block: + physical_block_cap += self.num_speculative_blocks + num_new_blocks = min(num_new_blocks, physical_block_cap) num_evictable_computed_blocks = self._get_num_evictable_blocks( new_computed_blocks From ab35354c21cc3c36439e79e18070f85cfafc3af2 Mon Sep 17 00:00:00 2001 From: Zheng Gong Date: Wed, 16 Sep 2026 15:26:56 +0800 Subject: [PATCH 06/20] [ROCm][AITER] Skip AITER norm kernels when flattening to 2D would copy (#55991) Signed-off-by: Zheng Gong Co-authored-by: Cursor Co-authored-by: TJian --- tests/kernels/ir/test_aiter_norm_dispatch.py | 60 ++++++++++++++++++++ vllm/kernels/aiter_ops.py | 32 ++++++++++- 2 files changed, 90 insertions(+), 2 deletions(-) create mode 100644 tests/kernels/ir/test_aiter_norm_dispatch.py diff --git a/tests/kernels/ir/test_aiter_norm_dispatch.py b/tests/kernels/ir/test_aiter_norm_dispatch.py new file mode 100644 index 000000000000..19cf5d3112dc --- /dev/null +++ b/tests/kernels/ir/test_aiter_norm_dispatch.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the AITER norm support predicates. + +The AITER norm wrappers flatten >2-D activations with ``Tensor.reshape``, which +silently copies when the flattened shape is not expressible with the input's +strides. The support predicates must reject those inputs so dispatch falls +through to a provider that handles arbitrary strides. +""" + +import pytest +import torch + +from vllm._aiter_ops import is_aiter_found_and_supported +from vllm.kernels.aiter_ops import flatten_to_2d_is_free + +pytestmark = pytest.mark.skipif( + not is_aiter_found_and_supported(), + reason="Only test on ROCm with AITER installed and supported", +) + + +def _qkv_slice_by_head(num_tokens, num_q_heads, num_kv_heads, head_dim): + """q viewed per-head from a fused QKV projection, as Qwen3-style QK-norm does.""" + q_size, kv_size = num_q_heads * head_dim, num_kv_heads * head_dim + qkv = torch.empty(num_tokens, q_size + 2 * kv_size) + q = qkv.split([q_size, kv_size, kv_size], dim=-1)[0] + return q.view(*q.shape[:-1], num_q_heads, head_dim) + + +@pytest.mark.parametrize( + "x", + [ + torch.empty(8, 16), + torch.empty(16), + torch.empty(2, 4, 16), + torch.empty(2, 1, 16), + # A contiguous tensor stays flattenable after a leading-dim slice. + torch.empty(8, 4, 16)[2:6], + ], +) +def test_flattenable(x): + assert flatten_to_2d_is_free(x) + assert x.reshape(-1, x.shape[-1]).data_ptr() == x.data_ptr() + + +@pytest.mark.parametrize( + "x", + [ + _qkv_slice_by_head(8, 32, 8, 128), + # Non-unit last-dim stride. + torch.empty(4, 8, 32).transpose(-1, -2), + # Leading dims cannot be merged: a slice along the middle dim. + torch.empty(4, 8, 32)[:, :4], + ], +) +def test_not_flattenable(x): + assert not flatten_to_2d_is_free(x) + # reshape has to copy, which is exactly what the predicate guards against. + assert x.reshape(-1, x.shape[-1]).data_ptr() != x.data_ptr() diff --git a/vllm/kernels/aiter_ops.py b/vllm/kernels/aiter_ops.py index de8a53ae8b77..402c9cac1752 100644 --- a/vllm/kernels/aiter_ops.py +++ b/vllm/kernels/aiter_ops.py @@ -34,13 +34,38 @@ def is_aiter_found() -> bool: AITER_SUPPORTED = is_aiter_found() """Most kernels in this file are supported if AITER is installed.""" + +def flatten_to_2d_is_free(x: Tensor) -> bool: + """Whether ``x.reshape(-1, x.shape[-1])`` is a view with unit-stride rows. + + The AITER norm kernels only take dense 2D inputs, so the wrappers below + flatten the leading dims with ``Tensor.reshape``, which silently falls back + to ``contiguous()`` when the flattened shape is not expressible with the + existing strides, adding a whole-tensor device copy that costs more than + the kernel saves. A strided last dim flattens without a copy but not into a + dense buffer, so it is rejected too. + """ + if x.dim() <= 2: + return True + if x.stride(-1) != 1: + return False + expected_stride = x.size(-1) + for i in range(x.dim() - 2, -1, -1): + if x.size(i) != 1 and x.stride(i) != expected_stride: + return False + expected_stride *= x.size(i) + return True + + rms_no_var_16bit_only = ( lambda x, weight, epsilon, variance_size=None: variance_size is None and x.dtype in (torch.float16, torch.bfloat16) and (weight is None or weight.dtype == x.dtype) + and flatten_to_2d_is_free(x) ) """AITER rms_norm only supports float16 and bfloat16 acts, no var_size override, -and requires weight dtype to match x dtype.""" +requires weight dtype to match x dtype, and requires flattening the activation +to 2D to be free.""" @ir.ops.rms_norm.register_impl( @@ -79,10 +104,13 @@ def _rms_norm_fake(x: Tensor, weight: Tensor, variance_epsilon: float) -> Tensor lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None and x.dtype in (torch.float16, torch.bfloat16) and (weight is None or weight.dtype == x.dtype) + and flatten_to_2d_is_free(x) + and flatten_to_2d_is_free(x_residual) ) """ AITER fused_add_rms_norm only supports 16-bit activations and no var_size override. -Requires weight dtype to match x dtype. +Requires weight dtype to match x dtype, and flattening both the activation and the +residual to 2D to be free. """ From f12fe10c5c377ccd401b66c39df3ffb394fa067c Mon Sep 17 00:00:00 2001 From: MINJUN GIL Date: Wed, 16 Sep 2026 16:47:26 +0900 Subject: [PATCH 07/20] [Bugfix][KV Offload] Track cache recency once per request (#51787) Signed-off-by: mindungil Co-authored-by: Or Ozeri --- .../offloading_connector/test_scheduler.py | 71 +-- tests/v1/kv_offload/cpu/test_manager.py | 436 +++++++++++++++++- .../tiering/test_tiering_offloading.py | 61 ++- .../kv_connector/v1/offloading/scheduler.py | 70 ++- vllm/v1/kv_offload/base.py | 12 + vllm/v1/kv_offload/cpu/manager.py | 146 +++++- vllm/v1/kv_offload/cpu/policies/arc.py | 81 +++- vllm/v1/kv_offload/cpu/policies/base.py | 52 ++- vllm/v1/kv_offload/cpu/policies/lru.py | 142 ++++-- vllm/v1/kv_offload/tiering/manager.py | 18 +- 10 files changed, 918 insertions(+), 171 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 89b77ae43fa8..035fb6b0054f 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -258,6 +258,9 @@ def test_partial_tail_store_uses_attention_and_recurrent_cow_sources(): jobs = scheduler._build_partial_tail_store_jobs(output) assert len(jobs) == 1 + offered_keys = scheduler.manager.prepare_store.call_args.args[0] + req_context = scheduler._req_status["req"].req_context + assert all(req_context.get_offload_key_position(key) == 28 for key in offered_keys) [job_id] = jobs src_spec = jobs[job_id].src_spec assert isinstance(src_spec, GPULoadStoreSpec) @@ -380,6 +383,13 @@ def test_normal_store_excludes_align_mode_mamba_sources(): req_status.group_states[0].block_ids[:] = [11] req_status.group_states[1].block_ids[:] = [99] req_status.update_offload_keys() + for group_config, group_state in zip( + scheduler.config.kv_group_configs, req_status.group_states + ): + assert ( + req_status.req_context.get_offload_key_position(group_state.offload_keys[0]) + == group_config.tokens_per_chunk + ) scheduler.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) @@ -921,7 +931,6 @@ def test_offloading_connector(request_runner, async_scheduling: bool): runner.run(decoded_tokens=[0] * (tokens_per_chunk + 1)) # 1 more block (+ token for kicking off offloading) - # now check touch was called with all 6 blocks runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) @@ -929,9 +938,6 @@ def test_offloading_connector(request_runner, async_scheduling: bool): decoded_tokens=[0] * (tokens_per_chunk + 1), expected_stored=(15, 16, 17), ) - runner.manager.touch.assert_called() - block_hashes1 = list(runner.manager.touch.call_args.args[0]) - assert len(block_hashes1) == 6 # terminate request runner.run(decoded_tokens=[EOS_TOKEN_ID]) @@ -939,13 +945,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool): # create a new request differing only on the last token runner.new_request(token_ids=[0] * (tokens_per_chunk * 6 - 1) + [1]) runner.run(decoded_tokens=[0]) - runner.manager.touch.assert_called() - block_hashes2 = list(runner.manager.touch.call_args.args[0]) - assert len(block_hashes2) == 6 - - # verify hashes are the same, except for the last block - assert block_hashes1[:5] == block_hashes2[:5] - assert block_hashes1[5] != block_hashes2[5] + runner.manager.touch.assert_not_called() # terminate request runner.run( @@ -1353,14 +1353,7 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo generate_store_output(keys) ) runner.run(decoded_tokens=[0]) - # _touch called from get_num_new_matched_tokens (2 groups) and - # _get_reqs_to_store (2 groups) → 4 touch calls total. - touch_calls = runner.manager.touch.call_args_list - assert len(touch_calls) == 4 - assert len(touch_calls[0].args[0]) == 3 - assert len(touch_calls[1].args[0]) == 3 - assert len(touch_calls[2].args[0]) == 3 - assert len(touch_calls[3].args[0]) == 3 + runner.manager.touch.assert_not_called() # store 3 more block runner.manager.prepare_store.side_effect = lambda keys, req_context: ( @@ -1371,9 +1364,7 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo expected_stored=(0, 1, 2, 3, 4, 5), ) - # touch called from _get_reqs_to_store * 3 blocks, once for each group - touch_calls = runner.manager.touch.call_args_list - assert len(touch_calls) == 6 + runner.manager.touch.assert_not_called() # EOS lands in the last slot of the 7th block (offset 6). No forward pass # writes that slot, so the finishing step declines the block. @@ -1392,13 +1383,7 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo expected_loaded=((0, 0), (0, 1), (0, 2), (1, 1), (1, 2)), ) - # 2 touch calls from get_num_new_matched_tokens (2 groups) - touch_calls = runner.manager.touch.call_args_list - assert len(touch_calls) == 2 - # full attention group touched all 3 blocks - assert len(touch_calls[0].args[0]) == 3 - # sliding window group touched just the last 2 blocks - assert len(touch_calls[1].args[0]) == 2 + runner.manager.touch.assert_not_called() # 3 blocks are hit on GPU [0, 1, 2] # 1 block loaded [3,] @@ -1464,15 +1449,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool generate_store_output(keys) ) runner.run(decoded_tokens=[0]) - # _touch called from get_num_new_matched_tokens (2 groups) and - # _get_reqs_to_store (2 groups) → 4 touch calls total. - # Group 0 has 2 offload keys, group 1 has 1. - touch_calls = runner.manager.touch.call_args_list - assert len(touch_calls) == 4 - assert len(touch_calls[0].args[0]) == 2 - assert len(touch_calls[1].args[0]) == 1 - assert len(touch_calls[2].args[0]) == 2 - assert len(touch_calls[3].args[0]) == 1 + runner.manager.touch.assert_not_called() # Get to 31 tokens # No further blocks offloaded @@ -1482,11 +1459,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool # Group 0 blocks: [0, 1], ending_token_offset = 24 # Group 1 blocks: [0, 1], ending_token_offset = 32 runner.run(decoded_tokens=[0]) - # _get_reqs_to_store touch: only group 1 has a new block to store - touch_calls = runner.manager.touch.call_args_list - assert len(touch_calls) == 2 - assert len(touch_calls[0].args[0]) == 2 - assert len(touch_calls[1].args[0]) == 2 + runner.manager.touch.assert_not_called() # Get to 35 tokens # No further blocks offloaded @@ -1496,11 +1469,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool # Group 0 blocks: [0, 1, 2], ending_token_offset = 36 # Group 1 blocks: [0, 1], ending_token_offset = 32 runner.run(decoded_tokens=[0]) - # _get_reqs_to_store touch: only group 0 has a new block to store - touch_calls = runner.manager.touch.call_args_list - assert len(touch_calls) == 2 - assert len(touch_calls[0].args[0]) == 3 - assert len(touch_calls[1].args[0]) == 2 + runner.manager.touch.assert_not_called() # Get to 47 tokens # No further blocks offloaded @@ -1510,11 +1479,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool # Group 0 blocks: [0, 1, 2, 3], ending_token_offset = 4 # Group 1 blocks: [0, 1, 2], ending_token_offset = 48 runner.run(decoded_tokens=[0]) - # _get_reqs_to_store touch: both groups have a new block, each with 1 key - touch_calls = runner.manager.touch.call_args_list - assert len(touch_calls) == 2 - assert len(touch_calls[0].args[0]) == 4 - assert len(touch_calls[1].args[0]) == 3 + runner.manager.touch.assert_not_called() runner.run(decoded_tokens=[0], expected_stored=((0, 3), (1, 2))) diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index a79f480ff551..4fce65fb2a4f 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -23,6 +23,7 @@ ) from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy +from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy def make_req_context( @@ -430,31 +431,34 @@ def test_cpu_manager(): assert cpu_manager.lookup(to_key(0), _EMPTY_REQ_CTX) is LookupResult.MISS # prepare load [2, 3] - prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), _EMPTY_REQ_CTX) + load_ctx = make_req_context("load-2-3") + prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), load_ctx) verify_load_output(prepare_load_output, [1, 2]) # prepare store with no space ([2, 3] is being loaded) assert cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) is None # complete load [2, 3]. Load changes the eviction list, making 2, 3 recent. - cpu_manager.complete_load(to_keys([2, 3]), _EMPTY_REQ_CTX) + cpu_manager.complete_load(to_keys([2, 3]), load_ctx) + cpu_manager.on_request_finished(load_ctx) - # prepare store [6, 7, 8] -> evicts [4, 5, 2] (oldest) + # prepare store [6, 7, 8] -> evicts [4, 5, 3] (oldest). Within + # the accessed prefix [2, 3], the tail is less valuable than the head. prepare_store_output = cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) verify_store_output( prepare_store_output, ExpectedPrepareStoreOutput( keys_to_store=[6, 7, 8], - store_chunk_ids=[1, 0, 3], - evicted_keys=[4, 5, 2], + store_chunk_ids=[2, 0, 3], + evicted_keys=[4, 5, 3], ), ) # complete store [6, 7, 8] cpu_manager.complete_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) - # touch [3, 6, 7] (move to end of LRU order) - cpu_manager.touch(to_keys([3, 6, 7]), _EMPTY_REQ_CTX) + # touch [2, 6, 7] (move to end of LRU order) + cpu_manager.touch(to_keys([2, 6, 7]), _EMPTY_REQ_CTX) # prepare store [7, 9] -> evicts [8] (oldest following previous touch) prepare_store_output = cpu_manager.prepare_store(to_keys([9]), _EMPTY_REQ_CTX) @@ -477,7 +481,7 @@ def test_cpu_manager(): verify_events( cpu_manager.take_events(), expected_stores=({3, 4, 5}, {6, 7, 8}), - expected_evictions=({4, 5, 2}, {8}), + expected_evictions=({4, 5, 3}, {8}), ) @@ -518,6 +522,65 @@ def test_prepare_load_preserves_key_order(): manager.complete_load([key_a, key_b, key_c], _EMPTY_REQ_CTX) # order irrelevant +def test_lru_batch_eviction_failure_is_atomic(): + manager = make_cpu_manager(num_chunks=4, cache_policy="lru") + policy = manager._policy + assert isinstance(policy, LRUCachePolicy) + keys = to_keys([1, 2, 3, 4]) + assert manager.prepare_store(keys, _EMPTY_REQ_CTX) is not None + manager.complete_store(keys, _EMPTY_REQ_CTX) + + protected = set(keys[1:]) + assert policy.evict(2, protected) is None + assert all(manager.lookup(key, _EMPTY_REQ_CTX) is LookupResult.HIT for key in keys) + + evicted = policy.evict(1, protected) + assert evicted is not None + assert [key for key, _ in evicted] == [keys[0]] + + +def test_lru_repeated_pin_cycles_compact_lazy_heap_entries(): + manager = make_cpu_manager(num_chunks=1, cache_policy="lru") + policy = manager._policy + assert isinstance(policy, LRUCachePolicy) + key = to_key(1) + assert manager.prepare_store([key], _EMPTY_REQ_CTX) is not None + manager.complete_store([key], _EMPTY_REQ_CTX) + + for _ in range(128): + manager.prepare_load([key], _EMPTY_REQ_CTX) + manager.complete_load([key], _EMPTY_REQ_CTX) + + assert len(policy._heap) < 64 + evicted = policy.evict(1, set()) + assert evicted is not None + assert [evicted_key for evicted_key, _ in evicted] == [key] + + +def test_reset_discards_stale_request_access_classification(): + manager = make_cpu_manager(num_chunks=1, cache_policy="arc") + policy = manager._policy + assert isinstance(policy, ARCCachePolicy) + key = to_key(1) + resumed_ctx = make_req_context("resumed") + + assert manager.prepare_store([key], resumed_ctx) is not None + manager.complete_store([key], resumed_ctx) + manager.reset_cache() + + replacement_ctx = make_req_context("replacement") + assert manager.prepare_store([key], replacement_ctx) is not None + manager.complete_store([key], replacement_ctx) + manager.on_request_finished(replacement_ctx) + assert key in policy.t1 + + manager.prepare_load([key], resumed_ctx) + manager.complete_load([key], resumed_ctx) + manager.on_request_finished(resumed_ctx) + + assert key in policy.t2 + + class TestARCPolicy: """Unit tests for CPUOffloadingManager with ARC eviction policy.""" @@ -1042,6 +1105,48 @@ def test_filter_reused_manager_oversized_offer_makes_progress(): assert stored_keys == set(keys) +def test_store_threshold_does_not_hide_ready_reuse(): + manager = make_cpu_manager( + num_chunks=2, + cache_policy="lru", + store_threshold=2, + max_tracker_size=2, + ) + head, tail, noise_1, noise_2, replacement = to_keys(list(range(5))) + + seed_ctx = make_req_context("seed") + output = manager.prepare_store([head, tail], seed_ctx) + assert output is not None + assert not output.keys_to_store + output = manager.prepare_store([head, tail], seed_ctx) + assert output is not None + assert output.keys_to_store == [head, tail] + manager.complete_store([head, tail], seed_ctx) + manager.on_request_finished(seed_ctx) + + # Age both resident keys out of the admission tracker without storing the + # noise keys. Reusing the tail recreates its counter at one, below the + # threshold, but must still refresh its cache recency. + output = manager.prepare_store([noise_1, noise_2], make_req_context("noise")) + assert output is not None + assert not output.keys_to_store + reuse_ctx = make_req_context("reuse") + skipped_before_reuse = manager.stores_skipped_in_current_batch + output = manager.prepare_store([tail], reuse_ctx) + assert output is not None + assert not output.keys_to_store + assert manager.stores_skipped_in_current_batch == skipped_before_reuse + manager.on_request_finished(reuse_ctx) + + replacement_ctx = make_req_context("replacement") + output = manager.prepare_store([replacement], replacement_ctx) + assert output is not None + assert not output.keys_to_store + output = manager.prepare_store([replacement], replacement_ctx) + assert output is not None + assert output.evicted_keys == [head] + + def test_evictable_cache_chunk_count(): """ Verifies _num_evictable_cache_chunks is maintained correctly through the @@ -1156,3 +1261,318 @@ def spy_touch(keys: Iterable[OffloadKey], req_context: ReqContext) -> None: assert len(received) == 1 assert received[0][0] == keys assert received[0][1] is ctx + + +@pytest.mark.parametrize("finish_before_completion", [False, True]) +def test_request_finish_orders_lru_prefix_independent_of_store_completion( + finish_before_completion: bool, +): + """A transfer's completion order must not make the prefix head LRU.""" + manager = make_cpu_manager(num_chunks=4, cache_policy="lru") + prefix = to_keys([1, 2, 3, 4]) + ctx = make_req_context("prefix") + + output = manager.prepare_store(prefix, ctx) + assert output is not None + if finish_before_completion: + manager.on_request_finished(ctx) + manager.complete_store(set(prefix), ctx) + if not finish_before_completion: + manager.on_request_finished(ctx) + + output = manager.prepare_store(to_keys([5]), make_req_context("evict")) + assert output is not None + assert output.evicted_keys == to_keys([4]) + + +@pytest.mark.parametrize("finish_before_completion", [False, True]) +def test_request_finish_orders_lru_prefix_independent_of_load_completion( + finish_before_completion: bool, +): + manager = make_cpu_manager(num_chunks=4, cache_policy="lru") + prefix = to_keys([1, 2, 3, 4]) + + seed_ctx = make_req_context("seed") + output = manager.prepare_store(prefix, seed_ctx) + assert output is not None + manager.complete_store(prefix, seed_ctx) + + reuse_ctx = make_req_context("reuse") + manager.prepare_load(prefix, reuse_ctx) + if finish_before_completion: + manager.on_request_finished(reuse_ctx) + manager.complete_load(set(prefix), reuse_ctx) + if not finish_before_completion: + manager.on_request_finished(reuse_ctx) + + output = manager.prepare_store(to_keys([5]), make_req_context("evict")) + assert output is not None + assert output.evicted_keys == to_keys([4]) + + +def test_late_old_completion_does_not_override_newer_request_recency(): + manager = make_cpu_manager(num_chunks=4, cache_policy="lru") + old_keys = to_keys([1, 2]) + new_keys = to_keys([3, 4]) + + old_ctx = make_req_context("old") + assert manager.prepare_store(old_keys, old_ctx) is not None + manager.on_request_finished(old_ctx) + + new_ctx = make_req_context("new") + assert manager.prepare_store(new_keys, new_ctx) is not None + manager.complete_store(new_keys, new_ctx) + manager.on_request_finished(new_ctx) + + # Physical completion is late, but the logical access is still older. + manager.complete_store(old_keys, old_ctx) + + output = manager.prepare_store(to_keys([5]), make_req_context("evict")) + assert output is not None + assert output.evicted_keys == [old_keys[-1]] + + +def test_arc_counts_reuse_once_and_keeps_insertions_in_t1(): + manager = make_cpu_manager(num_chunks=4, cache_policy="arc") + policy = manager._policy + assert isinstance(policy, ARCCachePolicy) + prefix = to_keys([1, 2, 3, 4]) + + seed_ctx = make_req_context("seed") + assert manager.prepare_store(prefix, seed_ctx) is not None + manager.complete_store(prefix, seed_ctx) + # Repeated store offers from the same request must not become ARC hits. + assert manager.prepare_store(prefix, seed_ctx) is not None + # Nor should an internal read used to cascade the new chunks to a tier. + manager.prepare_load(prefix, seed_ctx) + manager.complete_load(prefix, seed_ctx) + manager.on_request_finished(seed_ctx) + + assert list(policy.t1) == list(reversed(prefix)) + assert not policy.t2 + + reuse_ctx = make_req_context("reuse") + manager.prepare_load(prefix, reuse_ctx) + manager.complete_load(prefix, reuse_ctx) + # A duplicate internal pin, as used by tiering cascades, is deduplicated. + manager.prepare_load(prefix, reuse_ctx) + manager.complete_load(prefix, reuse_ctx) + manager.on_request_finished(reuse_ctx) + + assert not policy.t1 + assert list(policy.t2) == list(reversed(prefix)) + + +def test_arc_reuse_wins_if_same_request_later_reinserts_key(): + manager = make_cpu_manager(num_chunks=1, cache_policy="arc") + policy = manager._policy + assert isinstance(policy, ARCCachePolicy) + key_1, key_2 = to_keys([1, 2]) + + seed_ctx = make_req_context("seed") + assert manager.prepare_store([key_1], seed_ctx) is not None + manager.complete_store([key_1], seed_ctx) + manager.on_request_finished(seed_ctx) + + reuse_ctx = make_req_context("reuse-and-reinsert") + manager.prepare_load([key_1], reuse_ctx) + manager.complete_load([key_1], reuse_ctx) + + other_ctx = make_req_context("evict-reused-key") + assert manager.prepare_store([key_2], other_ctx) is not None + manager.complete_store([key_2], other_ctx) + manager.on_request_finished(other_ctx) + + assert manager.prepare_store([key_1], reuse_ctx) is not None + manager.complete_store([key_1], reuse_ctx) + manager.on_request_finished(reuse_ctx) + + assert key_1 not in policy.t1 + assert list(policy.t2) == [key_1] + + +def test_request_finish_forwards_grouped_deduplicated_accesses(monkeypatch): + manager = make_cpu_manager(num_chunks=5) + existing = make_offload_key(b"existing", 0) + group_0_new = make_offload_key(b"group-0-new", 0) + group_1_head = make_offload_key(b"group-1-head", 1) + group_1_tail = make_offload_key(b"group-1-tail", 1) + + seed_ctx = make_req_context("seed") + assert manager.prepare_store([existing], seed_ctx) is not None + manager.complete_store([existing], seed_ctx) + + ctx = make_req_context("grouped") + # Group 1 is observed first, but policy finalization follows group index. + offered = [group_1_head, existing, group_0_new] + assert manager.prepare_store(offered, ctx) is not None + # Later offers extend a group's prefix and may repeat earlier keys. + assert ( + manager.prepare_store([group_0_new, group_1_head, group_1_tail], ctx) + is not None + ) + + received = [] + + def on_request_finished(key_groups, insertion_only_keys, reused_keys, req_context): + received.append((key_groups, insertion_only_keys, reused_keys, req_context)) + + monkeypatch.setattr(manager._policy, "on_request_finished", on_request_finished) + manager.on_request_finished(ctx) + manager.on_request_finished(ctx) + + assert received == [ + ( + ((existing, group_0_new), (group_1_head, group_1_tail)), + {group_0_new, group_1_head, group_1_tail}, + {existing}, + ctx, + ) + ] + + +def test_failed_store_records_access_once_and_uses_store_miss_hook(monkeypatch): + manager = make_cpu_manager(num_chunks=1) + resident, candidate = to_keys([1, 2]) + + seed_ctx = make_req_context("seed") + assert manager.prepare_store([resident], seed_ctx) is not None + manager.complete_store([resident], seed_ctx) + + pin_ctx = make_req_context("pin") + manager.prepare_load([resident], pin_ctx) + + record_calls = 0 + original_record = manager._record_request_cache_access + + def record_access(*args, **kwargs): + nonlocal record_calls + record_calls += 1 + return original_record(*args, **kwargs) + + store_miss_calls = [] + + def on_store_miss(keys, req_context): + store_miss_calls.append((list(keys), req_context)) + + finalized = [] + + def on_request_finished(key_groups, insertion_only_keys, reused_keys, req_context): + finalized.append((key_groups, insertion_only_keys, reused_keys, req_context)) + + monkeypatch.setattr(manager, "_record_request_cache_access", record_access) + monkeypatch.setattr(manager._policy, "on_store_miss", on_store_miss) + monkeypatch.setattr(manager._policy, "on_request_finished", on_request_finished) + + ctx = make_req_context("failed-store") + assert manager.prepare_store([candidate, resident], ctx) is None + assert record_calls == 1 + assert store_miss_calls == [([candidate], ctx)] + + manager.on_request_finished(ctx) + assert finalized == [ + (((candidate, resident),), set(), {resident}, ctx), + ] + + manager.complete_load([resident], pin_ctx) + + +def test_arc_ghost_hit_adapts_once_per_request_before_insertion(): + manager = make_cpu_manager(num_chunks=2, cache_policy="arc") + policy = manager._policy + assert isinstance(policy, ARCCachePolicy) + keys = to_keys([1, 2, 3]) + + seed_ctx = make_req_context("seed") + assert manager.prepare_store(keys[:2], seed_ctx) is not None + manager.complete_store(keys[:2], seed_ctx) + manager.on_request_finished(seed_ctx) + + evict_ctx = make_req_context("create-ghost") + output = manager.prepare_store([keys[2]], evict_ctx) + assert output is not None + [ghost_key] = output.evicted_keys + manager.complete_store([keys[2]], evict_ctx) + manager.on_request_finished(evict_ctx) + assert ghost_key in policy.b1 + + resident_keys = [key for key in keys if key != ghost_key] + pin_ctx = make_req_context("pin-residents") + manager.prepare_load(resident_keys, pin_ctx) + + ghost_ctx = make_req_context("ghost-hit") + target_before = policy.target_t1_size + assert manager.prepare_store([ghost_key], ghost_ctx) is None + target_after_first_offer = policy.target_t1_size + assert target_after_first_offer > target_before + assert manager.prepare_store([ghost_key], ghost_ctx) is None + assert policy.target_t1_size == target_after_first_offer + + manager.complete_load(resident_keys, pin_ctx) + output = manager.prepare_store([ghost_key], ghost_ctx) + assert output is not None + assert policy.target_t1_size == target_after_first_offer + manager.complete_store([ghost_key], ghost_ctx) + manager.on_request_finished(ghost_ctx) + assert ghost_key in policy.t1 + assert ghost_key not in policy.t2 + + +def test_arc_pending_chunk_is_not_a_frequency_hit(): + manager = make_cpu_manager(num_chunks=2, cache_policy="arc") + policy = manager._policy + assert isinstance(policy, ARCCachePolicy) + key = to_key(1) + + writer_ctx = make_req_context("writer") + assert manager.prepare_store([key], writer_ctx) is not None + + observer_ctx = make_req_context("pending-observer") + output = manager.prepare_store([key], observer_ctx) + assert output is not None and not output.keys_to_store + manager.on_request_finished(observer_ctx) + + chunk = policy.get(key) + assert chunk is not None and not chunk.is_ready + assert key in policy.t1 + assert key not in policy.t2 + + +def test_request_key_positions_override_store_observation_order(): + manager = make_cpu_manager(num_chunks=2, cache_policy="lru") + head, tail = to_keys([1, 2]) + ctx = make_req_context("out-of-order-store") + ctx.set_offload_key_position(head, 16) + ctx.set_offload_key_position(tail, 32) + + # A backfill or partial-tail path may offer a later key first. + assert manager.prepare_store([tail], ctx) is not None + assert manager.prepare_store([head], ctx) is not None + manager.on_request_finished(ctx) + manager.complete_store({head, tail}, ctx) + + output = manager.prepare_store([to_key(3)], make_req_context("evict-tail")) + assert output is not None + assert output.evicted_keys == [tail] + + +def test_request_key_positions_order_tails_across_kv_groups(): + manager = make_cpu_manager(num_chunks=4, cache_policy="lru") + group_0_head = make_offload_key(b"group-0-head", 0) + group_0_tail = make_offload_key(b"group-0-tail", 0) + group_1_head = make_offload_key(b"group-1-head", 1) + group_1_tail = make_offload_key(b"group-1-tail", 1) + keys = [group_0_head, group_0_tail, group_1_head, group_1_tail] + ctx = make_req_context("hybrid-groups") + for key in (group_0_head, group_1_head): + ctx.set_offload_key_position(key, 16) + for key in (group_0_tail, group_1_tail): + ctx.set_offload_key_position(key, 32) + + assert manager.prepare_store(keys, ctx) is not None + manager.complete_store(keys, ctx) + manager.on_request_finished(ctx) + + output = manager.prepare_store(to_keys([10, 11]), make_req_context("evict")) + assert output is not None + assert set(output.evicted_keys) == {group_0_tail, group_1_tail} diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index f737685621b4..771153ec4734 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -40,6 +40,7 @@ TierMatcher, make_offload_key, ) +from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy from vllm.v1.kv_offload.tiering.base import ( JobResult, SecondaryTierManager, @@ -753,14 +754,64 @@ def test_touch_propagates_to_all_tiers(self, manager_setup): # Touch chunks self.manager.touch(chunks, _CTX) - # Verify touch was called on primary tier (check LRU order) - primary_keys = list(self.primary_tier._policy.evictable_chunks.keys()) - assert primary_keys[-3:] == list(reversed(chunks)) + # Verify touch was called on the primary tier (head is most recent). + policy = self.primary_tier._policy + assert isinstance(policy, LRUCachePolicy) + ranks = policy._ranks + assert ranks[chunks[2]] < ranks[chunks[1]] < ranks[chunks[0]] # Verify touch was propagated to all secondary tiers self.secondary_tier1.touch.assert_called_once_with(chunks, _CTX) self.secondary_tier2.touch.assert_called_once_with(chunks, _CTX) + def test_request_order_survives_late_cascade_completion(self, manager_setup): + """Cascade completion must not replace request order with I/O order.""" + chunks = to_keys(range(5)) + self._start_request() + output = self.manager.prepare_store(chunks, _CTX) + assert output is not None + self.manager.complete_store(chunks, _CTX, success=True) + + # Finalize while both secondary-tier cascades still pin primary chunks. + self.manager.on_request_finished(_CTX) + self._simulate_on_schedule_end() + self._simulate_on_schedule_end() + + evict_ctx = ReqContext(req_id="evict") + self._start_request(evict_ctx) + output = self.manager.prepare_store(to_keys([5]), evict_ctx) + assert output is not None + assert output.evicted_keys == [chunks[-1]] + + def test_late_old_store_submission_does_not_override_newer_request_order( + self, manager_setup + ): + """Cascade submission after finish is not a new primary access.""" + old_blocks = to_keys(range(2)) + old_ctx = ReqContext(req_id="old") + self._start_request(old_ctx) + assert self.manager.prepare_store(old_blocks, old_ctx) is not None + self.manager.on_request_finished(old_ctx) + + new_blocks = to_keys(range(2, 5)) + new_ctx = ReqContext(req_id="new") + self._start_request(new_ctx) + assert self.manager.prepare_store(new_blocks, new_ctx) is not None + self.manager.complete_store(new_blocks, new_ctx, success=True) + self.manager.on_request_finished(new_ctx) + + # The old GPU->primary write lands after the newer request finished. + # Its cascade pins must not turn it into the more recent request. + self.manager.complete_store(old_blocks, old_ctx, success=True) + self._simulate_on_schedule_end() + self._simulate_on_schedule_end() + + evict_ctx = ReqContext(req_id="evict") + self._start_request(evict_ctx) + output = self.manager.prepare_store(to_keys([5]), evict_ctx) + assert output is not None + assert output.evicted_keys == [old_blocks[-1]] + def test_failed_store_no_cascade(self, manager_setup): """Test that failed GPU→primary store doesn't cascade.""" chunks = to_keys(range(3)) @@ -873,10 +924,10 @@ def test_complete_store_forwards_req_context_to_submit_store(self, manager_setup job_metadata = self.secondary_tier1.submit_store.call_args.args[0] assert job_metadata.req_context is ctx - def test_on_request_finished_delays_secondary_until_store_submitted( + def test_on_request_finished_delays_only_secondary_until_store_submitted( self, manager_setup ): - """Manager hook is eager; secondary hooks wait for cascade submission.""" + """Primary order commits immediately; secondary cleanup waits.""" chunks = to_keys(range(2)) ctx = ReqContext(req_id="req_delayed_secondary") calls: list[tuple[str, str]] = [] 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 5cc66e4bf0ff..b253a2399867 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -427,8 +427,11 @@ def update_offload_keys(self) -> None: None, group_config.hashes_per_chunk, ): - group_state.offload_keys.append( - make_offload_key(req_block_hash, group_config.group_idx) + key = make_offload_key(req_block_hash, group_config.group_idx) + group_state.offload_keys.append(key) + self.req_context.set_offload_key_position( + key, + len(group_state.offload_keys) * group_config.tokens_per_chunk, ) def update_block_id_groups( @@ -742,37 +745,6 @@ def _sliding_window_lookup( ) return None if defer_lookup or pending_in_window else consecutive_hits - def _touch(self, req_status: RequestOffloadState): - for group_config, group_state in zip( - self.config.kv_group_configs, req_status.group_states - ): - if group_config.sliding_window_size_in_chunks is None: - self.manager.touch(group_state.offload_keys, req_status.req_context) - else: - # Keep only chunks needed to hit the original request, plus - # decoded chunks. - chunks_to_skip = max( - 0, - group_state.num_hit_chunks - - group_config.sliding_window_size_in_chunks, - ) - self.manager.touch( - group_state.offload_keys[chunks_to_skip:], - req_status.req_context, - ) - if req_status.partial_tail_boundary is not None: - self.manager.touch( - tuple( - self._make_boundary_key( - req_status.req, - group.group_idx, - req_status.partial_tail_boundary, - ) - for group in self.config.kv_group_configs - ), - req_status.req_context, - ) - def _lookup_complete_chunks( self, req_status: RequestOffloadState, @@ -977,10 +949,16 @@ def _lookup_complete_chunks( return num_hit_tokens def _make_boundary_key( - self, request: Request, group_idx: int, boundary_tokens: int + self, + request: Request, + group_idx: int, + boundary_tokens: int, + req_context: ReqContext, ) -> OffloadKey: hash_idx = boundary_tokens // self.config.tokens_per_hash - 1 - return make_offload_key(request.block_hashes[hash_idx], group_idx) + key = make_offload_key(request.block_hashes[hash_idx], group_idx) + req_context.set_offload_key_position(key, boundary_tokens) + return key def _lookup( self, @@ -1015,7 +993,10 @@ def _lookup( boundary_keys = [] for group_config in self.config.kv_group_configs: key = self._make_boundary_key( - req_status.req, group_config.group_idx, boundary + req_status.req, + group_config.group_idx, + boundary, + req_status.req_context, ) boundary_keys.append(key) result = self.manager.lookup(key, req_status.req_context) @@ -1108,8 +1089,6 @@ def get_num_new_matched_tokens( self._maybe_observe_lookup_async_delay(req_status) req_status.update_num_hit_chunks(num_computed_tokens + (num_hit_tokens or 0)) - self._touch(req_status) - return num_hit_tokens, bool(num_hit_tokens) def update_state_after_alloc( @@ -1183,7 +1162,10 @@ def update_state_after_alloc( if partial_tail_boundary is not None: keys_to_load.append( self._make_boundary_key( - request, group_config.group_idx, partial_tail_boundary + request, + group_config.group_idx, + partial_tail_boundary, + req_status.req_context, ) ) @@ -1307,7 +1289,9 @@ def _build_aligned_boundary_store_jobs( ): continue - key = self._make_boundary_key(req, group_idx, boundary) + key = self._make_boundary_key( + req, group_idx, boundary, req_status.req_context + ) store_output = self.manager.prepare_store([key], req_status.req_context) if store_output is None: self._connector_stats.increase_counter( @@ -1402,7 +1386,9 @@ def _build_partial_tail_store_jobs( ): continue keys = [ - self._make_boundary_key(req, group.group_idx, boundary) + self._make_boundary_key( + req, group.group_idx, boundary, req_status.req_context + ) for group in self.config.kv_group_configs ] block_ids = [ @@ -1689,8 +1675,6 @@ def _build_store_jobs( req_status.advance_stored_idx(num_offloadable_tokens) continue - self._touch(req_status) - keys_to_store = set(store_output.keys_to_store) group_sizes: list[int] = [] diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 93f8e4bda4e9..988747e06e1e 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -96,6 +96,12 @@ class ReqContext: # kv_transfer_params once (in on_new_request) and read the result back # on later calls for the same request. _state: dict[type, Any] = field(default_factory=dict, repr=False, init=False) + # End-token position for each key in this request. The scheduler records + # these positions so managers can recover prefix order even when store + # calls arrive out of order (for example, SWA backfills). + _offload_key_positions: dict[OffloadKey, int] = field( + default_factory=dict, repr=False, init=False + ) def set_state(self, val: Any) -> None: self._state[type(val)] = val @@ -103,6 +109,12 @@ def set_state(self, val: Any) -> None: def get_state(self, cls: type[_T]) -> _T | None: return self._state.get(cls) + def set_offload_key_position(self, key: OffloadKey, end_token: int) -> None: + self._offload_key_positions[key] = end_token + + def get_offload_key_position(self, key: OffloadKey) -> int | None: + return self._offload_key_positions.get(key) + class LookupResult(Enum): """Result of OffloadingManager.lookup().""" diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index 67b469b6a394..9d85090f2d8f 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections import OrderedDict from collections.abc import Collection, Iterable +from dataclasses import dataclass, field from typing_extensions import override @@ -18,6 +19,7 @@ PrepareStoreOutput, ReqContext, RequestOffloadingContext, + get_offload_group_idx, ) from vllm.v1.kv_offload.cpu.common import ( CPULoadStoreSpec, @@ -27,6 +29,20 @@ from vllm.v1.kv_offload.cpu.policies.factory import CachePolicyFactory +@dataclass(slots=True) +class _RequestCacheAccess: + """Cache keys observed by one request, grouped in prefix order.""" + + owner: object + cache_generation: int + key_groups: dict[int, list[OffloadKey]] = field(default_factory=dict) + seen_keys: set[OffloadKey] = field(default_factory=set) + inserted_keys: set[OffloadKey] = field(default_factory=set) + reused_keys: set[OffloadKey] = field(default_factory=set) + store_miss_keys: set[OffloadKey] = field(default_factory=set) + finished: bool = False + + class CPUOffloadingManager(OffloadingManager): """ An OffloadingManager with a pluggable CachePolicy, resolved by name via @@ -66,6 +82,7 @@ def __init__( self.max_tracker_size: int = max_tracker_size self.stores_skipped_in_current_batch: int = 0 self.allocation_sizes_in_current_batch: list[int] = [] + self._cache_generation = 0 # Number of chunk references. It is ordered so can evict the LRU entry in O(1). self.counts: OrderedDict[OffloadKey, int] | None = ( @@ -124,10 +141,48 @@ def _record_accesses(self, keys: Collection[OffloadKey]) -> None: num_unprotected -= 1 self.counts[key] = 1 + def _get_request_cache_access(self, req_context: ReqContext) -> _RequestCacheAccess: + state = req_context.get_state(_RequestCacheAccess) + if ( + state is None + or state.owner is not self + or state.cache_generation != self._cache_generation + ): + state = _RequestCacheAccess( + owner=self, cache_generation=self._cache_generation + ) + req_context.set_state(state) + return state + + def _record_request_cache_access( + self, + keys: Iterable[OffloadKey], + req_context: ReqContext, + inserted_keys: Iterable[OffloadKey] = (), + reused_keys: Iterable[OffloadKey] = (), + ) -> None: + state = self._get_request_cache_access(req_context) + for key in keys: + if key in state.seen_keys: + continue + assert not state.finished, ( + "New cache keys observed after request finalization" + ) + group_idx = get_offload_group_idx(key) + state.key_groups.setdefault(group_idx, []).append(key) + state.seen_keys.add(key) + state.inserted_keys.update(inserted_keys) + # Re-reading a chunk inserted by this request is an internal transfer + # (for example, a tiering cascade), not a second cache access. + state.reused_keys.update( + key for key in reused_keys if key not in state.inserted_keys + ) + # --- OffloadingManager interface --- @override def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: + self._get_request_cache_access(req_context) return RequestOffloadingContext() @override @@ -144,6 +199,15 @@ def prepare_load( self, keys: Collection[OffloadKey], req_context: ReqContext, + ) -> LoadStoreSpec: + return self._prepare_load(keys, req_context, record_access=True) + + def _prepare_load( + self, + keys: Collection[OffloadKey], + req_context: ReqContext, + *, + record_access: bool, ) -> LoadStoreSpec: chunks = [] for key in keys: @@ -156,6 +220,8 @@ def prepare_load( assert self._num_evictable_cache_chunks >= 0 chunk.ref_cnt += 1 chunks.append(chunk) + if record_access: + self._record_request_cache_access(keys, req_context, reused_keys=keys) return self._get_load_store_spec(keys, chunks) @override @@ -181,13 +247,52 @@ def prepare_store( keys: Collection[OffloadKey], req_context: ReqContext, ) -> PrepareStoreOutput | None: + keys = list(keys) if self.counts is not None: - num_keys = len(keys) self._record_accesses(keys) - keys = [k for k in keys if self.counts.get(k, 0) >= self.store_threshold] - self.stores_skipped_in_current_batch += num_keys - len(keys) - # filter out chunks that are already stored - keys_to_store = [k for k in keys if self._policy.get(k) is None] + # Partition the original offer before admission filtering. Ready + # resident chunks are request reuses even if their tracker entry was + # aged out; the threshold applies only to new store candidates. + keys_to_store: list[OffloadKey] = [] + ready_existing_keys: list[OffloadKey] = [] + for key in keys: + chunk = self._policy.get(key) + if chunk is None: + keys_to_store.append(key) + else: + if chunk.is_ready: + ready_existing_keys.append(key) + + if self.counts is not None: + num_store_candidates = len(keys_to_store) + keys_to_store = [ + key + for key in keys_to_store + if self.counts.get(key, 0) >= self.store_threshold + ] + self.stores_skipped_in_current_batch += num_store_candidates - len( + keys_to_store + ) + + state = self._get_request_cache_access(req_context) + new_store_misses = [ + key for key in keys_to_store if key not in state.store_miss_keys + ] + if new_store_misses: + # ARC learns from B1/B2 before insert() removes the ghost entry. + # Deduplication makes this one policy observation per request. + self._policy.on_store_miss(new_store_misses, req_context) + state.store_miss_keys.update(new_store_misses) + + recorded_keys = set(keys_to_store) + recorded_keys.update(ready_existing_keys) + # Record once before any allocation-related early return. Store + # candidates are classified as insertions only after allocation succeeds. + self._record_request_cache_access( + (key for key in keys if key in recorded_keys), + req_context, + reused_keys=ready_existing_keys, + ) if not keys_to_store: return PrepareStoreOutput( @@ -239,6 +344,7 @@ def prepare_store( for key, chunk in zip(keys_to_store, chunks): self._policy.insert(key, chunk) self._num_write_pending_chunks += len(keys_to_store) + state.inserted_keys.update(keys_to_store) # build store specs for allocated chunks store_spec = self._get_load_store_spec(keys_to_store, chunks) @@ -284,6 +390,35 @@ def complete_store( ) ) + @override + def on_request_finished(self, req_context: ReqContext) -> None: + state = req_context.get_state(_RequestCacheAccess) + if ( + state is None + or state.owner is not self + or state.cache_generation != self._cache_generation + or state.finished + ): + return + state.finished = True + key_groups = [] + for group_idx in sorted(state.key_groups): + keys = state.key_groups[group_idx] + positions = { + key: position + for key in keys + if (position := req_context.get_offload_key_position(key)) is not None + } + if len(positions) == len(keys): + keys = sorted(keys, key=positions.__getitem__) + key_groups.append(tuple(keys)) + self._policy.on_request_finished( + tuple(key_groups), + state.inserted_keys - state.reused_keys, + state.reused_keys, + req_context, + ) + @override def reset_cache(self) -> None: # Clear ALL chunks unconditionally. The scheduler's _stale_job_threshold @@ -292,6 +427,7 @@ def reset_cache(self) -> None: # flushes in-flight load job IDs to the workers before any new stores # can begin, preventing a cross-direction data race on reused offload chunk IDs. self._policy.clear() + self._cache_generation += 1 self._num_evictable_cache_chunks = 0 self._num_write_pending_chunks = 0 diff --git a/vllm/v1/kv_offload/cpu/policies/arc.py b/vllm/v1/kv_offload/cpu/policies/arc.py index e413bd603eb4..882ee02219c2 100644 --- a/vllm/v1/kv_offload/cpu/policies/arc.py +++ b/vllm/v1/kv_offload/cpu/policies/arc.py @@ -1,12 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections import OrderedDict -from collections.abc import Iterable, Iterator +from collections.abc import Iterable, Iterator, Sequence from typing_extensions import override from vllm.v1.kv_offload.base import OffloadKey, ReqContext -from vllm.v1.kv_offload.cpu.policies.base import CachePolicy, ChunkStatus +from vllm.v1.kv_offload.cpu.policies.base import ( + CachePolicy, + ChunkStatus, + order_request_keys, +) class ARCCachePolicy(CachePolicy): @@ -24,12 +28,11 @@ class ARCCachePolicy(CachePolicy): Searches T1 and T2 for chunk hashes and counts consecutive hits until a miss or non-ready chunk is encountered. - 2. Cache touch (touch) - Adaptive Learning: - For each key (in reverse order): - - If in T1: Move to T2 (promotion from recent to frequent). - - If in T2: Move to MRU position (end of queue). - - If in B1 ghost list: Increase target_t1_size. - - If in B2 ghost list: Decrease target_t1_size. + 2. Request access - Adaptive Learning: + - Ready chunks reused by a request move from T1 to T2 once, or + move to the MRU end of T2. + - B1/B2 misses adjust target_t1_size once per request before + insertion removes their ghost entries. 3. Chunk eviction (evict) - Adaptive Replacement: Determines eviction source based on adaptive target: @@ -39,7 +42,7 @@ class ARCCachePolicy(CachePolicy): 4. Chunk insertion (insert): New chunks are always inserted into T1 and removed from B1/B2 if - present. Chunks may later be promoted to T2 during touch operations. + present. A later request reuse may promote them to T2. Adaptive Behavior: The algorithm self-tunes the recency vs. frequency trade-off: @@ -71,6 +74,26 @@ def remove(self, key: OffloadKey) -> None: if self.t1.pop(key, None) is None: self.t2.pop(key, None) + def _adapt_to_ghost_hit(self, key: OffloadKey) -> bool: + if key in self.b1: + delta = max(1, len(self.b2) / len(self.b1)) + self.target_t1_size = min(self.target_t1_size + delta, self.cache_capacity) + self.b1.move_to_end(key) + return True + if key in self.b2: + delta = max(1, len(self.b1) / len(self.b2)) + self.target_t1_size = max(self.target_t1_size - delta, 0) + self.b2.move_to_end(key) + return True + return False + + @override + def on_store_miss( + self, keys: Iterable[OffloadKey], req_context: ReqContext + ) -> None: + for key in reversed(list(keys)): + self._adapt_to_ghost_hit(key) + @override def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: for key in reversed(list(keys)): @@ -86,19 +109,33 @@ def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: elif key in self.t2: self.t2.move_to_end(key) - elif key in self.b1: - delta = max(1, len(self.b2) / len(self.b1)) - self.target_t1_size = min( - self.target_t1_size + delta, self.cache_capacity - ) - # move to MRU position (end) to keep it fresh in the ghost list - self.b1.move_to_end(key) - - elif key in self.b2: - delta = max(1, len(self.b1) / len(self.b2)) - self.target_t1_size = max(self.target_t1_size - delta, 0) - # move to MRU position (end) to keep it fresh in the ghost list - self.b2.move_to_end(key) + else: + self._adapt_to_ghost_hit(key) + + @override + def on_request_finished( + self, + key_groups: Sequence[Sequence[OffloadKey]], + insertion_only_keys: set[OffloadKey], + reused_keys: set[OffloadKey], + req_context: ReqContext, + ) -> None: + for key in reversed(order_request_keys(key_groups, req_context)): + if key in insertion_only_keys: + # A store is not a frequency hit. Preserve T1 membership + # while still restoring tail-to-head recency. + if key in self.t1: + self.t1.move_to_end(key) + elif key in self.t2: + self.t2.move_to_end(key) + continue + + # Ready chunks reused by this request count as one access. + if key in reused_keys and key in self.t1: + chunk = self.t1.pop(key) + self.t2[key] = chunk + elif key in reused_keys and key in self.t2: + self.t2.move_to_end(key) @override def clear(self) -> None: diff --git a/vllm/v1/kv_offload/cpu/policies/base.py b/vllm/v1/kv_offload/cpu/policies/base.py index 4059dd742f7c..3efffbd3961a 100644 --- a/vllm/v1/kv_offload/cpu/policies/base.py +++ b/vllm/v1/kv_offload/cpu/policies/base.py @@ -2,11 +2,26 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import ctypes from abc import ABC, abstractmethod -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from vllm.v1.kv_offload.base import OffloadKey, ReqContext +def order_request_keys( + key_groups: Sequence[Sequence[OffloadKey]], req_context: ReqContext +) -> list[OffloadKey]: + """Return one head-to-tail order across all KV cache groups.""" + keys = [key for group in key_groups for key in group] + positions = { + key: position + for key in keys + if (position := req_context.get_offload_key_position(key)) is not None + } + if len(positions) == len(keys): + keys.sort(key=positions.__getitem__) + return keys + + class ChunkStatus(ctypes.Structure): """ Offloading status for a single chunk of KV data. @@ -66,6 +81,41 @@ def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: req_context: Per-request context for the request touching these chunks. """ + def on_store_miss( + self, keys: Iterable[OffloadKey], req_context: ReqContext + ) -> None: + """Observe store misses before their cache entries are inserted. + + The default delegates to ``touch`` for compatibility with external + policies. Policies may override this hook when a miss has distinct + semantics, such as ARC adapting to a ghost-list hit. + """ + self.touch(keys, req_context) + + def on_request_finished( + self, + key_groups: Sequence[Sequence[OffloadKey]], + insertion_only_keys: set[OffloadKey], + reused_keys: set[OffloadKey], + req_context: ReqContext, + ) -> None: + """Apply one request-scoped cache access in prefix order. + + ``key_groups`` contains keys observed for each KV cache group in + head-to-tail order. ``insertion_only_keys`` and ``reused_keys`` + distinguish chunks only created by this request from ready chunks it + actually reused, which matters for policies such as ARC where reuse + changes frequency but insertion and pending observations do not. + + The default forwards one touch per group, preserving compatibility + for experimental out-of-tree policies while moving those touches to + request finalization. Policies that distinguish insertion from reuse + can override this hook and inspect the access classifications. + """ + del insertion_only_keys, reused_keys + for keys in key_groups: + self.touch(keys, req_context) + @abstractmethod def evict( self, n: int, protected: set[OffloadKey] diff --git a/vllm/v1/kv_offload/cpu/policies/lru.py b/vllm/v1/kv_offload/cpu/policies/lru.py index c1dc631d4fe9..5b5f65dec779 100644 --- a/vllm/v1/kv_offload/cpu/policies/lru.py +++ b/vllm/v1/kv_offload/cpu/policies/lru.py @@ -1,28 +1,74 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections import OrderedDict -from collections.abc import Iterable +import heapq +from collections.abc import Iterable, Sequence from typing_extensions import override from vllm.v1.kv_offload.base import OffloadKey, ReqContext -from vllm.v1.kv_offload.cpu.policies.base import CachePolicy, ChunkStatus +from vllm.v1.kv_offload.cpu.policies.base import ( + CachePolicy, + ChunkStatus, + order_request_keys, +) class LRUCachePolicy(CachePolicy): """ - LRU Caching policy that keeps a dedicated evictable list for fast eviction. + LRU caching policy with logical recency independent of transfer pinning. + + Evictable chunks live in a lazy-invalidating min-heap. A chunk's recency + can therefore be updated while it is pinned; when it later becomes + evictable, it enters the heap with the order assigned by the request + rather than its transfer-completion time. + A use is indicated by, - First time the key is added (store). - - Load job completion - - touch + - A request-scoped access. """ def __init__(self, cache_capacity: int): super().__init__(cache_capacity) - # Chunks with ref_cnt 0 (not participating in any loads/stores) ordered in LRU - self.evictable_chunks: OrderedDict[OffloadKey, None] = OrderedDict() self.chunks: dict[OffloadKey, ChunkStatus] = {} + self._ranks: dict[OffloadKey, int] = {} + self._evictable: set[OffloadKey] = set() + self._heap: list[tuple[int, OffloadKey]] = [] + self._next_rank = 0 + + def _assign_new_rank(self, key: OffloadKey) -> bool: + if key not in self.chunks: + return False + self._next_rank += 1 + self._ranks[key] = self._next_rank + return key in self._evictable + + def _push_evictable(self, key: OffloadKey) -> None: + heapq.heappush( + self._heap, + (self._ranks[key], key), + ) + + def _is_current(self, entry: tuple[int, OffloadKey]) -> bool: + rank, key = entry + return key in self._evictable and self._ranks.get(key) == rank + + def _maybe_compact_heap(self) -> None: + if len(self._heap) < max(64, 2 * len(self._evictable)): + return + self._heap = [(self._ranks[key], key) for key in self._evictable] + heapq.heapify(self._heap) + + def _update_recency(self, keys: Iterable[OffloadKey]) -> None: + updated_evictable = [key for key in keys if self._assign_new_rank(key)] + # Rebuilding is linear and substantially cheaper than k heap pushes + # for a long prefix. Small updates retain the incremental path. + if len(updated_evictable) >= max(64, len(self._evictable) // 4): + self._heap = [(self._ranks[key], key) for key in self._evictable] + heapq.heapify(self._heap) + else: + for key in updated_evictable: + self._push_evictable(key) + self._maybe_compact_heap() @override def get(self, key: OffloadKey) -> ChunkStatus | None: @@ -31,26 +77,40 @@ def get(self, key: OffloadKey) -> ChunkStatus | None: @override def insert(self, key: OffloadKey, chunk: ChunkStatus) -> None: self.chunks[key] = chunk + self._next_rank += 1 + self._ranks[key] = self._next_rank if chunk.ref_cnt == 0: - self.evictable_chunks[key] = None + self._evictable.add(key) + self._push_evictable(key) @override def remove(self, key: OffloadKey) -> None: del self.chunks[key] - self.evictable_chunks.pop(key, None) + self._ranks.pop(key, None) + self._evictable.discard(key) @override def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: - for key in reversed(list(keys)): - if key in self.evictable_chunks: - self.evictable_chunks.move_to_end(key) - # active chunks are untouched as they are non-evictable now. They - # will eventually reach the end of evictable_chunks when they finish. + self._update_recency(reversed(list(keys))) + + @override + def on_request_finished( + self, + key_groups: Sequence[Sequence[OffloadKey]], + insertion_only_keys: set[OffloadKey], + reused_keys: set[OffloadKey], + req_context: ReqContext, + ) -> None: + del insertion_only_keys, reused_keys + self._update_recency(reversed(order_request_keys(key_groups, req_context))) @override def clear(self) -> None: - self.evictable_chunks.clear() self.chunks.clear() + self._ranks.clear() + self._evictable.clear() + self._heap.clear() + self._next_rank = 0 @override def evict( @@ -59,22 +119,44 @@ def evict( if n == 0: return [] - candidates: list[tuple[OffloadKey, ChunkStatus]] = [] - for key, _ in self.evictable_chunks.items(): + selected: list[tuple[tuple[int, OffloadKey], ChunkStatus]] = [] + selected_keys: set[OffloadKey] = set() + deferred: list[tuple[int, OffloadKey]] = [] + while self._heap and len(selected) < n: + entry = heapq.heappop(self._heap) + if not self._is_current(entry): + continue + key = entry[1] + # Re-pinning without a recency change can leave an equivalent + # lazy entry behind. Select each cache key at most once. + if key in selected_keys: + continue if key in protected: + deferred.append(entry) continue - chunk = self.chunks[key] assert chunk.ref_cnt == 0 - candidates.append((key, chunk)) - if len(candidates) == n: - break - - if len(candidates) < n: + selected.append((entry, chunk)) + selected_keys.add(key) + + if len(selected) < n: + for entry, _ in selected: + heapq.heappush(self._heap, entry) + for entry in deferred: + heapq.heappush(self._heap, entry) return None - for key, _ in candidates: - del self.evictable_chunks[key] + + for entry in deferred: + heapq.heappush(self._heap, entry) + + candidates: list[tuple[OffloadKey, ChunkStatus]] = [] + for entry, chunk in selected: + key = entry[1] + self._evictable.remove(key) del self.chunks[key] + del self._ranks[key] + candidates.append((key, chunk)) + self._maybe_compact_heap() return candidates @override @@ -82,9 +164,11 @@ def mark_evictable(self, key: OffloadKey) -> None: # chunks can become evictable when, # store completes - i.e. ref_cnt -1 -> 0 # not in evictable list # all loads complete - i.e ref_cnt 1 -> 0 # not in evictable list - self.evictable_chunks[key] = None + self._evictable.add(key) + self._push_evictable(key) + self._maybe_compact_heap() @override def mark_non_evictable(self, key: OffloadKey) -> None: - # key must have been in the evictable list. - del self.evictable_chunks[key] + # key must have been evictable before it was pinned. + self._evictable.remove(key) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 0940aaeae07b..f19616828daa 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -110,13 +110,22 @@ def __init__( # read/write is for CPU<->secondary transfers, # load/store is for CPU<->GPU transfers. # These aliases avoid calling prepare_load inside a store path. - self.prepare_read = self.prepare_load self.complete_read = self.complete_load self.prepare_write = self.prepare_store self.complete_write = self.complete_store self._kv_memoryview = mmap_region.create_kv_memoryview() + def prepare_read( + self, keys: Collection[OffloadKey], req_context: ReqContext + ) -> LoadStoreSpec: + """Pin chunks for a CPU-to-secondary transfer. + + Cascade reads are implementation details of tiering, not additional + request accesses, so they must not alter request-scoped recency. + """ + return self._prepare_load(keys, req_context, record_access=False) + def get_kv_memoryview(self) -> memoryview: """Return the memoryview over the primary tier's KV cache buffer. @@ -781,11 +790,10 @@ def _maybe_finalize_request( req_id: str, exclude_tier_idx: int | None = None, ) -> None: - """Finalize secondary tiers once no more store cascades can be submitted. + """Finalize secondary tiers once no more cascades can be submitted. - Finalization means forwarding on_request_finished() to secondary tiers. - It is delayed until pending GPU->primary stores finish, since their - complete_store() callbacks may still submit primary->secondary stores. + Their finalization is delayed until pending GPU->primary stores + finish, since those callbacks may still submit secondary stores. """ state = self._req_state[req_id] if not state.is_finished: From b68408043d80faf9c99fc16c6aec1174336e9626 Mon Sep 17 00:00:00 2001 From: "Li, Jiang" Date: Wed, 16 Sep 2026 15:49:17 +0800 Subject: [PATCH 08/20] [Bugfix][CPU] Add per-tensor FP8 W8A16 kernel to fix Ministral crash on CPU (#56985) Signed-off-by: jiang1.li --- .../scripts/hardware_ci/run-cpu-test.sh | 21 +++- .../quantization/test_cpu_fp8_scaled_mm.py | 113 ++++++++++++++++++ tests/v1/e2e/test_cpu_spec_decode.py | 2 +- .../model_executor/kernels/linear/__init__.py | 2 + .../kernels/linear/scaled_mm/__init__.py | 2 + .../kernels/linear/scaled_mm/cpu.py | 100 ++++++++++++++++ 6 files changed, 233 insertions(+), 7 deletions(-) diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-test.sh index e8d0a24bc7d1..43555a19ab58 100644 --- a/.buildkite/scripts/hardware_ci/run-cpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test.sh @@ -83,14 +83,23 @@ rm -f "$build_log" # Run the image, setting --shm-size=4g for tensor parallel. Default to # HF_HUB_OFFLINE so a warm ~/.cache/huggingface doesn't hit the network; -# retry once online if the cache is missing something. vllm's get_config() -# wraps the raw huggingface_hub offline-mode errors in a generic ValueError -# (see transformers_utils/config.py), so match that message too or the -# fallback never triggers for a config-lookup cache miss. -OFFLINE_RETRY_PATTERN='huggingface_hub\.errors\.(LocalEntryNotFoundError|OfflineModeIsEnabled)|Invalid repository ID or local directory specified' +# retry once online if the cache is missing something. vllm wraps offline +# cache-miss errors in ways that don't preserve the raw huggingface_hub +# exception name: get_config() re-raises as a generic ValueError +# (transformers_utils/config.py), and default_loader.py raises its own +# RuntimeError when a snapshot dir has a config but no weight files. Match +# those messages too or the fallback never triggers for those cache misses. +# +# Also mount ~/.cache/vllm: multimodal test fixtures (e.g. VideoAsset, via +# vllm/assets/video.py) download into VLLM_ASSETS_CACHE (~/.cache/vllm/assets +# by default), a directory distinct from the HF hub cache above. Without a +# persistent mount there, those assets never survive past one container's +# lifetime, so every run re-fetches them from the network on the "offline" +# attempt and unconditionally falls back to the online retry. +OFFLINE_RETRY_PATTERN='huggingface_hub\.errors\.(LocalEntryNotFoundError|OfflineModeIsEnabled)|Invalid repository ID or local directory specified|Cannot find any model weights with' run_test() { local hf_offline=$1 - docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 -e HF_HUB_OFFLINE="$hf_offline" -e HF_DATASETS_OFFLINE="$hf_offline" --shm-size=4g "$IMAGE_NAME" \ + docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface -v ~/.cache/vllm:/root/.cache/vllm --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 -e HF_HUB_OFFLINE="$hf_offline" -e HF_DATASETS_OFFLINE="$hf_offline" --shm-size=4g "$IMAGE_NAME" \ timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}" } diff --git a/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py b/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py index 3154e2cb98bb..b70a1a8543ba 100644 --- a/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py +++ b/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py @@ -9,6 +9,15 @@ import torch from vllm import _custom_ops as ops +from vllm.model_executor.kernels.linear.scaled_mm.cpu import ( + CPUFp8PerTensorScaledMMLinearKernel, +) +from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import ( + FP8ScaledMMLinearLayerConfig, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8StaticTensorSym, +) from vllm.platforms import current_platform if not current_platform.is_cpu(): @@ -17,6 +26,9 @@ if not ops._supports_cpu_fp8_w8a16: pytest.skip("fp8_scaled_mm_cpu op not available", allow_module_level=True) +if not torch.cpu._is_amx_tile_supported(): + pytest.skip("requires AMX tile support", allow_module_level=True) + BLOCK_SIZE = [128, 128] @@ -160,3 +172,104 @@ def test_cpu_fp8_scaled_mm(M: int, N: int, K: int, use_bias: bool): assert kernel_out.dtype == out_dtype torch.testing.assert_close(kernel_out, ref_out, rtol=0.02, atol=0.01) + + +def quantize_weight_per_tensor_fp8( + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize weight [N, K] to FP8 with a single per-tensor scale.""" + fp8_max = torch.finfo(torch.float8_e4m3fn).max + scale = weight.abs().amax() / fp8_max + scale = torch.where(scale == 0, torch.ones_like(scale), scale) + q = (weight / scale).clamp(-fp8_max, fp8_max).to(torch.float8_e4m3fn) + return q, scale + + +def ref_fp8_per_tensor_scaled_mm( + x: torch.Tensor, + fp8_weight: torch.Tensor, + scale: torch.Tensor, + bias: torch.Tensor | None, + out_dtype: torch.dtype, +) -> torch.Tensor: + w_dq = fp8_weight.float() * scale + out = torch.mm(x.float(), w_dq.t()) + if bias is not None: + out = out + bias.float() + return out.to(out_dtype) + + +NK_SIZES_PER_TENSOR = [ + (32, 64), + (5120, 5120), +] + + +@pytest.mark.parametrize("M", [1, 64]) +@pytest.mark.parametrize("N,K", NK_SIZES_PER_TENSOR) +@pytest.mark.parametrize("use_bias", [False, True]) +def test_cpu_fp8_per_tensor_scaled_mm_kernel( + M: int, N: int, K: int, use_bias: bool, default_vllm_config +): + """CPUFp8PerTensorScaledMMLinearKernel correctness against float reference. + + Exercises the full kernel class (not just the raw op), including the + weight-orientation fixup: Fp8LinearMethod stores `layer.weight` as + [K, N] (torch._scaled_mm convention) before calling + process_weights_after_loading, so the kernel must transpose back to + [N, K] before VNNI-packing. + """ + torch.manual_seed(0) + out_dtype = torch.bfloat16 + + x = torch.randn(M, K, dtype=out_dtype) / (K**0.5) + w_f32 = torch.randn(N, K, dtype=torch.float32) / (K**0.5) + fp8_weight, scale = quantize_weight_per_tensor_fp8(w_f32) + bias = torch.randn(N, dtype=torch.float32) * 0.1 if use_bias else None + + ref_out = ref_fp8_per_tensor_scaled_mm(x, fp8_weight, scale, bias, out_dtype) + + config = FP8ScaledMMLinearLayerConfig( + weight_quant_key=kFp8StaticTensorSym, + activation_quant_key=kFp8StaticTensorSym, + weight_shape=(N, K), + input_dtype=out_dtype, + out_dtype=out_dtype, + ) + kernel = CPUFp8PerTensorScaledMMLinearKernel( + config, + layer_param_names=["weight", "weight_scale", "input_scale", "input_scale_ub"], + ) + + layer = torch.nn.Module() + # Fp8LinearMethod stores the weight transposed to [K, N] before calling + # process_weights_after_loading. + layer.register_parameter( + "weight", torch.nn.Parameter(fp8_weight.t().contiguous(), requires_grad=False) + ) + layer.register_parameter( + "weight_scale", torch.nn.Parameter(scale.clone(), requires_grad=False) + ) + kernel.process_weights_after_loading(layer) + + kernel_out = kernel.apply_weights(layer, x, bias) + + assert kernel_out.dtype == out_dtype + torch.testing.assert_close(kernel_out, ref_out, rtol=0.02, atol=0.01) + + +@pytest.mark.parametrize("n", [16, 48]) +def test_cpu_fp8_per_tensor_kernel_rejects_non_multiple_of_32_n(n: int): + """The AMX tinygemm kernel tiles N in chunks of 32 and cannot handle a + remainder tile smaller than that, so can_implement must reject any N + that isn't a multiple of 32 rather than let it crash the process. + """ + config = FP8ScaledMMLinearLayerConfig( + weight_quant_key=kFp8StaticTensorSym, + activation_quant_key=kFp8StaticTensorSym, + weight_shape=(n, 64), + input_dtype=torch.bfloat16, + out_dtype=torch.bfloat16, + ) + supported, _ = CPUFp8PerTensorScaledMMLinearKernel.can_implement(config) + assert not supported diff --git a/tests/v1/e2e/test_cpu_spec_decode.py b/tests/v1/e2e/test_cpu_spec_decode.py index 2d34c1165625..a7ba5580c33b 100644 --- a/tests/v1/e2e/test_cpu_spec_decode.py +++ b/tests/v1/e2e/test_cpu_spec_decode.py @@ -131,7 +131,7 @@ def test_ngram_spec_decode_matches_baseline(baseline_refs): id="dflash", ), pytest.param( - "Qwen/Qwen3.5-0.8B-Base", + "Qwen/Qwen3.5-0.8B", { "method": "mtp", "num_speculative_tokens": 3, diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index d3c41f97d8e5..c610e9b9c8f0 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -187,6 +187,7 @@ ) from vllm.model_executor.kernels.linear.scaled_mm.cpu import ( CPUFp8BlockScaledMMKernel, + CPUFp8PerTensorScaledMMLinearKernel, CPUInt8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.cutlass import ( @@ -438,6 +439,7 @@ def _resolve_backend_kernels( ChannelWiseTorchFP8ScaledMMLinearKernel, ], PlatformEnum.CPU: [ + CPUFp8PerTensorScaledMMLinearKernel, PerTensorTorchFP8ScaledMMLinearKernel, ChannelWiseTorchFP8ScaledMMLinearKernel, ], diff --git a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py index 39f9abd460e1..f4a062b7f2b3 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py @@ -9,6 +9,7 @@ ) from vllm.model_executor.kernels.linear.scaled_mm.cpu import ( CPUFp8BlockScaledMMKernel, + CPUFp8PerTensorScaledMMLinearKernel, CPUInt8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.cutlass import ( @@ -67,5 +68,6 @@ "ZentorchInt8ScaledMMLinearKernel", "Fp8BlockScaledMMLinearKernel", "CPUFp8BlockScaledMMKernel", + "CPUFp8PerTensorScaledMMLinearKernel", "XPUFp8BlockScaledMMKernel", ] diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py index 92ea073a1abb..e72dffcbbc50 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py @@ -21,10 +21,14 @@ FP8ScaledMMLinearLayerConfig, ) from .ScaledMMLinearKernel import ( + FP8ScaledMMLinearKernel, Int8ScaledMMLinearKernel, Int8ScaledMMLinearLayerConfig, ) +# BLOCK_K in csrc/cpu/sgl-kernels/gemm.h — the AMX kernel's fixed K-block size. +_BLOCK_SIZE_K = 128 + class CPUInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel): @classmethod @@ -330,3 +334,99 @@ def apply_block_scaled_mm( raise NotImplementedError( "CPUFp8BlockScaledMMKernel overrides apply_weights directly." ) + + +class CPUFp8PerTensorScaledMMLinearKernel(FP8ScaledMMLinearKernel): + """FP8 W8A16 per-tensor-scaled GEMM via AMX BRGEMM on CPU. + + Reuses the block-scaled AMX kernel (fp8_scaled_mm_cpu) with a single + synthetic block spanning the whole weight, so activations stay BF16/FP32 + — no FP8 activation quantization, unlike PerTensorTorchFP8ScaledMMLinearKernel. + """ + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_cpu(): + return False, "requires CPU platform." + if not torch.cpu._is_amx_tile_supported(): + return False, "requires AMX tile support (Sapphire Rapids or newer)." + if not ops._supports_cpu_fp8_w8a16: + return False, "fp8_scaled_mm_cpu op not available." + return True, None + + @classmethod + def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: + if not c.weight_quant_key.scale.group_shape.is_per_tensor(): + return False, "requires a per-tensor weight scale." + if not c.activation_quant_key.scale.group_shape.is_per_tensor(): + return False, "requires a per-tensor activation scale." + if c.out_dtype not in (torch.bfloat16, torch.float32): + return False, "Only bfloat16/float32 output dtype supported." + n = c.weight_shape[0] + if n % 32 != 0: + # The AMX tinygemm kernel tiles N in chunks of 32 and cannot + # handle a remainder tile smaller than that. + return False, f"requires weight output dim (N={n}) to be a multiple of 32." + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + w_name, w_s_name, _, _ = self.layer_param_names + # Fp8LinearMethod transposes weight to (K, N) for torch._scaled_mm-style + # kernels; convert_weight_packed/fp8_scaled_mm_cpu expect (N, K), as + # produced by CUDA nn.Linear-style checkpoints. + weight = getattr(layer, w_name).t().contiguous() + n, k = weight.shape + + packed_weight = torch.ops._C.convert_weight_packed(weight) + replace_parameter( + layer, w_name, torch.nn.Parameter(packed_weight, requires_grad=False) + ) + + # Synthesize a single-block "block scale" tensor from the scalar + # per-tensor weight scale so fp8_scaled_mm_cpu can be reused as-is. + weight_scale = getattr(layer, w_s_name) + self._block_size_n = -(-n // 32) * 32 # round up to a multiple of 32 + num_k_blocks = -(-k // _BLOCK_SIZE_K) + block_scale = weight_scale.reshape(()).expand(1, num_k_blocks).contiguous() + replace_parameter( + layer, w_s_name, torch.nn.Parameter(block_scale, requires_grad=False) + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + w_name, w_s_name, _, _ = self.layer_param_names + weight = getattr(layer, w_name) + weight_scale = getattr(layer, w_s_name) + + x_2d = x.reshape(-1, x.shape[-1]) if x.dim() > 2 else x + out = torch.ops._C.fp8_scaled_mm_cpu( + x_2d, + weight, + weight_scale, + [self._block_size_n, _BLOCK_SIZE_K], + bias, + x.dtype, + True, # is_vnni (weight already prepacked) + ) + return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out + + def apply_scaled_mm( + self, + *, + A: torch.Tensor, + B: torch.Tensor, + out_dtype: torch.dtype, + As: torch.Tensor, + Bs: torch.Tensor, + bias: torch.Tensor | None, + output_shape: list, + ) -> torch.Tensor: + raise NotImplementedError( + "CPUFp8PerTensorScaledMMLinearKernel overrides apply_weights directly." + ) From 9ca6dbba71329a7b08711e2acfa17b31e542553c Mon Sep 17 00:00:00 2001 From: Ama Senevirathne <97525823+amasen02@users.noreply.github.com> Date: Wed, 16 Sep 2026 13:26:50 +0530 Subject: [PATCH 09/20] fix(multimodal): tolerate malformed EXIF metadata during hashing (#56527) (#56576) Signed-off-by: amasen02 Co-authored-by: amasen02 --- tests/multimodal/test_hasher.py | 31 ++++++++++++++++++++++++++++++- vllm/multimodal/hasher.py | 29 +++++++++++++++++++---------- 2 files changed, 49 insertions(+), 11 deletions(-) diff --git a/tests/multimodal/test_hasher.py b/tests/multimodal/test_hasher.py index d0be316790d2..1b738eeeb89f 100644 --- a/tests/multimodal/test_hasher.py +++ b/tests/multimodal/test_hasher.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import contextlib import hashlib +import struct import uuid from io import BytesIO from pathlib import Path @@ -8,7 +10,7 @@ import numpy as np import pytest import torch -from PIL import Image, ImageDraw +from PIL import Image, ImageDraw, ImageOps from vllm.config.multimodal import MMHasherAlgorithm from vllm.multimodal.hasher import MultiModalHasher @@ -199,6 +201,33 @@ def test_hash_image_exif_id(): ) +def test_hash_image_malformed_exif(): + # Test that images with malformed EXIF headers (e.g. invalid TIFF header) + # do not raise an unhandled exception during hashing and fall back to image data. + buf = BytesIO() + Image.new("RGB", (64, 48)).save(buf, "JPEG") + jpg = buf.getvalue() + rest = jpg[2:] + rest = rest[2 + struct.unpack(">H", rest[2:4])[0] :] + payload = b"Exif\x00\x00XXXX\x00\x00\x00\x08" + bytes(32) + data = b"\xff\xd8\xff\xe1" + struct.pack(">H", len(payload) + 2) + payload + rest + + image = Image.open(BytesIO(data)) + with contextlib.suppress(Exception): + image = ImageOps.exif_transpose(image) + image.load() + + hasher = MultiModalHasher + # Should hash without raising SyntaxError or any other exception + hash_val = hasher.hash_kwargs("blake3", image=image) + assert isinstance(hash_val, str) and len(hash_val) > 0 + + # Also verify MediaWithBytes wrapping the image with malformed EXIF + media_item = MediaWithBytes(image, data) + hash_media = hasher.hash_kwargs("blake3", image=media_item) + assert isinstance(hash_media, str) and len(hash_media) > 0 + + def _rgba_png_bytes() -> bytes: image = Image.new("RGBA", (8, 8), (255, 0, 0, 128)) buf = BytesIO() diff --git a/vllm/multimodal/hasher.py b/vllm/multimodal/hasher.py index a432afdc9bbd..3750ae162bb0 100644 --- a/vllm/multimodal/hasher.py +++ b/vllm/multimodal/hasher.py @@ -48,6 +48,19 @@ def _get_hasher_factory( raise ValueError(f"Unsupported hash algorithm: {algorithm}") +def _get_image_id_bytes(image: Image.Image) -> bytes | None: + try: + exif = image.getexif() + image_id = exif.get(Image.ExifTags.Base.ImageID) + if isinstance(image_id, uuid.UUID): + return image_id.bytes + except Exception: + # Tolerate malformed EXIF metadata (e.g. invalid TIFF header) + # and fall back to serializing raw image data or bytes. + pass + return None + + class MultiModalHasher: @classmethod def serialize_item(cls, obj: object) -> Iterable[bytes | memoryview]: @@ -60,11 +73,9 @@ def serialize_item(cls, obj: object) -> Iterable[bytes | memoryview]: return (np.array(obj).tobytes(),) if isinstance(obj, Image.Image): - exif = obj.getexif() - if Image.ExifTags.Base.ImageID in exif and isinstance( - exif[Image.ExifTags.Base.ImageID], uuid.UUID - ): - return (exif[Image.ExifTags.Base.ImageID].bytes,) + image_id = _get_image_id_bytes(obj) + if image_id is not None: + return (image_id,) data = {"mode": obj.mode, "data": np.asarray(obj)} palette = obj.palette @@ -76,11 +87,9 @@ def serialize_item(cls, obj: object) -> Iterable[bytes | memoryview]: return cls.iter_item_to_bytes("image", data) if isinstance(obj, MediaWithBytes) and isinstance(obj.media, Image.Image): - exif = obj.media.getexif() - if Image.ExifTags.Base.ImageID in exif and isinstance( - exif[Image.ExifTags.Base.ImageID], uuid.UUID - ): - return (exif[Image.ExifTags.Base.ImageID].bytes,) + image_id = _get_image_id_bytes(obj.media) + if image_id is not None: + return (image_id,) if obj.io_config: return cls.iter_item_to_bytes( From 8be5205abbabf4c377c603d6c4180a99373f6415 Mon Sep 17 00:00:00 2001 From: dev <92231400+devtyagi3909@users.noreply.github.com> Date: Wed, 16 Sep 2026 13:47:22 +0530 Subject: [PATCH 10/20] [Parser] Fix: correct parser frontend handling of reasoning end and boundary tokens (#56635) Signed-off-by: DEV TYAGI --- tests/parser/engine/test_glm47_moe.py | 108 ++++++++++++++++++++++++++ vllm/parser/engine/parser_engine.py | 19 +++++ vllm/parser/glm47_moe.py | 14 ++-- 3 files changed, 135 insertions(+), 6 deletions(-) create mode 100644 tests/parser/engine/test_glm47_moe.py diff --git a/tests/parser/engine/test_glm47_moe.py b/tests/parser/engine/test_glm47_moe.py new file mode 100644 index 000000000000..58fe7aa28cbd --- /dev/null +++ b/tests/parser/engine/test_glm47_moe.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Reasoning-boundary tests for the GLM-4.7 parser. + +GLM can transition straight from reasoning into a tool call without ever +emitting ````, so ```` is an implicit reasoning +terminator. These cover that path plus the multi-turn prompts where an +earlier turn's markers must not be read as the current turn's state. +""" + +import pytest + +from tests.parser.engine.conftest import make_mock_tokenizer +from vllm.parser.glm47_moe import ( + THINK_END, + THINK_START, + TOOL_CALL_END, + TOOL_CALL_START, + Glm47MoeParser, +) + +THINK_S, THINK_E = 1000, 1001 +TOOL_S, TOOL_E = 1002, 1003 +ASSISTANT, OBSERVATION, USER = 1004, 1005, 1006 +TEXT = 42 # stand-in for an ordinary reasoning/content token + +VOCAB = { + THINK_START: THINK_S, + THINK_END: THINK_E, + TOOL_CALL_START: TOOL_S, + TOOL_CALL_END: TOOL_E, + "<|assistant|>": ASSISTANT, + "<|observation|>": OBSERVATION, + "<|user|>": USER, +} + + +@pytest.fixture +def parser(): + return Glm47MoeParser(make_mock_tokenizer(VOCAB)) + + +@pytest.fixture +def no_thinking_parser(): + return Glm47MoeParser( + make_mock_tokenizer(VOCAB), + chat_template_kwargs={"thinking": False}, + ) + + +class TestIsReasoningEnd: + def test_open_reasoning(self, parser): + assert not parser.is_reasoning_end([THINK_S, TEXT]) + + def test_think_end(self, parser): + assert parser.is_reasoning_end([THINK_S, TEXT, THINK_E, TEXT]) + + def test_tool_call_without_think_end(self, parser): + """Reasoning -> tool call with no ```` still ends reasoning.""" + assert parser.is_reasoning_end([THINK_S, TEXT, TOOL_S, TEXT]) + + def test_previous_turn_tool_call_ignored(self, parser): + """A finished tool call from an earlier turn says nothing about the + turn currently being generated.""" + prompt = [THINK_S, TEXT, TOOL_S, TEXT, TOOL_E, OBSERVATION, TEXT, ASSISTANT] + assert not parser.is_reasoning_end(prompt) + + def test_previous_turn_think_end_ignored(self, parser): + prompt = [THINK_S, TEXT, THINK_E, TEXT, USER, TEXT, ASSISTANT] + assert not parser.is_reasoning_end(prompt) + + def test_empty_input(self, parser): + assert not parser.is_reasoning_end([]) + + def test_thinking_disabled(self, no_thinking_parser): + assert no_thinking_parser.is_reasoning_end([THINK_S, TEXT]) + + +class TestExtractContentIds: + def test_think_end_wins_over_later_tool_call(self, parser): + """Content between ```` and a tool call must survive.""" + ids = [THINK_S, TEXT, THINK_E, 20, 21, TOOL_S, 30, TOOL_E] + assert parser.extract_content_ids(ids) == [20, 21, TOOL_S, 30, TOOL_E] + + def test_every_tool_call_after_think_end_kept(self, parser): + ids = [THINK_E, 20, TOOL_S, 30, TOOL_E, 21, TOOL_S, 31, TOOL_E] + assert parser.extract_content_ids(ids) == ids[1:] + + def test_falls_back_to_tool_call(self, parser): + """Without ````, content starts at the opener itself so the + tool parser still receives a well-formed call.""" + ids = [THINK_S, TEXT, TOOL_S, 30, TOOL_E] + assert parser.extract_content_ids(ids) == [TOOL_S, 30, TOOL_E] + + def test_falls_back_to_first_tool_call_of_turn(self, parser): + ids = [THINK_S, TEXT, TOOL_S, 30, TOOL_E, TOOL_S, 31, TOOL_E] + assert parser.extract_content_ids(ids) == ids[2:] + + def test_previous_turn_tool_call_ignored(self, parser): + ids = [TOOL_S, 30, TOOL_E, OBSERVATION, TEXT, ASSISTANT, THINK_S, TEXT] + assert parser.extract_content_ids(ids) == ids + + def test_no_markers_returns_input_ids(self, parser): + assert parser.extract_content_ids([20, 21]) == [20, 21] + + def test_thinking_disabled(self, no_thinking_parser): + ids = [THINK_S, TEXT, TOOL_S] + assert no_thinking_parser.extract_content_ids(ids) == ids diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index fa306a5da4a1..527c310c1485 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -672,11 +672,30 @@ def is_reasoning_end(self, input_ids: list[int]) -> bool: return not wait_for_reasoning def extract_content_ids(self, input_ids: list[int]) -> list[int]: + config = self.parser_engine_config + wait_for_reasoning = config.wait_for_reasoning + if wait_for_reasoning is None: + wait_for_reasoning = config.initial_state is ParserState.REASONING + if not wait_for_reasoning: + return input_ids + end_id = self._reasoning_end_token_id if end_id is not None: for i in range(len(input_ids) - 1, -1, -1): if input_ids[i] == end_id: return input_ids[i + 1 :] + + end_ids = self._reasoning_end_token_ids + if end_ids: + turn_start = 0 + boundary_ids = self._turn_boundary_token_ids + for i in range(len(input_ids) - 1, -1, -1): + if input_ids[i] in boundary_ids: + turn_start = i + 1 + break + for i in range(turn_start, len(input_ids)): + if input_ids[i] in end_ids: + return input_ids[i:] return input_ids def get_streaming_fallback_content( diff --git a/vllm/parser/glm47_moe.py b/vllm/parser/glm47_moe.py index 4a804a4c9c94..00e517f2ea68 100644 --- a/vllm/parser/glm47_moe.py +++ b/vllm/parser/glm47_moe.py @@ -41,6 +41,13 @@ ARG_VALUE_START = "" ARG_VALUE_END = "" +# Special tokens that delimit conversation turns in a rendered GLM prompt. +# Reasoning markers belonging to earlier turns must not be mistaken for the +# state of the turn currently being generated. +GLM_TURN_BOUNDARIES = frozenset( + ("<|system|>", "<|user|>", "<|assistant|>", "<|observation|>") +) + _ARG_RE = re.compile( r"(?P.*?)\s*" r"(?P.*?)", @@ -123,7 +130,6 @@ def glm47_moe_config(thinking: bool = True) -> ParserEngineConfig: return ParserEngineConfig( name="glm47_moe", initial_state=ParserState.REASONING if thinking else ParserState.CONTENT, - wait_for_reasoning=thinking, terminals={ **reasoning_terminals, "TOOL_START": TOOL_CALL_START, @@ -166,6 +172,7 @@ def glm47_moe_config(thinking: bool = True) -> ParserEngineConfig: ), **arg_tag_transitions, }, + turn_boundary_tokens=GLM_TURN_BOUNDARIES, arg_converter=_glm47_arg_converter, stream_arg_deltas=True, tool_args_json=False, @@ -207,11 +214,6 @@ def _handle_tool_end(self, event, deltas) -> None: self._tool_slots[idx].name = self._tool_slots[idx].name.strip() super()._handle_tool_end(event, deltas) - 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, From 0d8173d1537f60331084e12af6b72a6565978c9d Mon Sep 17 00:00:00 2001 From: wenjinhust <33047235+wenjinhust@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:25:49 +0800 Subject: [PATCH 11/20] [Bugfix][Frontend] Lazy-import model_hosting_container_standards to prevent log suppression (#51366) Signed-off-by: wenjinhust --- vllm/entrypoints/serve/lora/api_router.py | 30 ++++++++- .../entrypoints/serve/sagemaker/api_router.py | 61 ++++++++++++++++++- 2 files changed, 87 insertions(+), 4 deletions(-) diff --git a/vllm/entrypoints/serve/lora/api_router.py b/vllm/entrypoints/serve/lora/api_router.py index 95dff9792b2f..61436371c27d 100644 --- a/vllm/entrypoints/serve/lora/api_router.py +++ b/vllm/entrypoints/serve/lora/api_router.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import logging - -import model_hosting_container_standards.sagemaker as sagemaker_standards from fastapi import APIRouter, Depends, FastAPI, Request from fastapi.responses import JSONResponse, Response @@ -22,6 +21,12 @@ def attach_router(app: FastAPI): + """Attach the LoRA adapter load/unload endpoints to the API server. + + Does nothing when dynamic LoRA updating is disabled. Handler levels are + snapshotted and restored because importing + model_hosting_container_standards may reconfigure root logging. + """ if not envs.VLLM_ALLOW_RUNTIME_LORA_UPDATING: """If LoRA dynamic loading & unloading is not enabled, do nothing.""" return @@ -30,6 +35,23 @@ def attach_router(app: FastAPI): "This should ONLY be used for local development!" ) + snapshot = [ + (h, h.level) + for lg in (logging.getLogger(), logging.getLogger("vllm")) + for h in lg.handlers + ] + + try: + _attach_router(app) + finally: + for handler, level in snapshot: + handler.setLevel(level) + + +def _attach_router(app: FastAPI): + """Register the LoRA adapter load/unload routes on the app.""" + import model_hosting_container_standards.sagemaker as sagemaker_standards + @sagemaker_standards.register_load_adapter_handler( request_shape={ "lora_name": "body.name", @@ -40,6 +62,8 @@ def attach_router(app: FastAPI): ) @router.post("/v1/load_lora_adapter", dependencies=[Depends(validate_json_request)]) async def load_lora_adapter(request: LoadLoRAAdapterRequest, raw_request: Request): + """Handle POST /v1/load_lora_adapter: load a LoRA adapter into the + serving engine.""" handler: OpenAIServingModels = models(raw_request) response = await handler.load_lora_adapter(request) if isinstance(response, ErrorResponse): @@ -60,6 +84,8 @@ async def load_lora_adapter(request: LoadLoRAAdapterRequest, raw_request: Reques async def unload_lora_adapter( request: UnloadLoRAAdapterRequest, raw_request: Request ): + """Handle POST /v1/unload_lora_adapter: unload a LoRA adapter from + the serving engine.""" handler: OpenAIServingModels = models(raw_request) response = await handler.unload_lora_adapter(request) if isinstance(response, ErrorResponse): diff --git a/vllm/entrypoints/serve/sagemaker/api_router.py b/vllm/entrypoints/serve/sagemaker/api_router.py index 142256c26002..60576f289a0d 100644 --- a/vllm/entrypoints/serve/sagemaker/api_router.py +++ b/vllm/entrypoints/serve/sagemaker/api_router.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json +import logging from collections.abc import Awaitable, Callable from http import HTTPStatus from typing import Any -import model_hosting_container_standards.sagemaker as sagemaker_standards import pydantic from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request from fastapi.responses import JSONResponse, Response @@ -27,11 +27,55 @@ EndpointFn = Callable[[RequestType, Request], Awaitable[Any]] +def _snapshot_handler_levels() -> list[tuple[logging.Handler, int]]: + """Snapshot handler levels for loggers that may be affected by third-party + logging configuration side effects (e.g. model_hosting_container_standards + calling configure_root_logger() at import time).""" + loggers = [logging.getLogger(), logging.getLogger("vllm")] + seen: set[int] = set() + snapshot: list[tuple[logging.Handler, int]] = [] + for logger in loggers: + for handler in logger.handlers: + if id(handler) not in seen: + seen.add(id(handler)) + snapshot.append((handler, handler.level)) + return snapshot + + +def _restore_handler_levels( + snapshot: list[tuple[logging.Handler, int]], +) -> None: + """Restore handler levels from a snapshot.""" + for handler, level in snapshot: + handler.setLevel(level) + + def attach_router( app: FastAPI, supported_tasks: tuple["SupportedTask", ...], model_config: ModelConfig | None = None, ): + """Attach the SageMaker hosting endpoints to the API server. + + Handler levels are snapshotted and restored because importing + model_hosting_container_standards may reconfigure root logging. + """ + snapshot = _snapshot_handler_levels() + try: + _attach_router(app, supported_tasks, model_config) + finally: + _restore_handler_levels(snapshot) + + +def _attach_router( + app: FastAPI, + supported_tasks: tuple["SupportedTask", ...], + model_config: ModelConfig | None = None, +): + """Register the SageMaker hosting routes (/ping, /invocations) on the + app.""" + import model_hosting_container_standards.sagemaker as sagemaker_standards + router = APIRouter() # NOTE: Construct the TypeAdapters only once @@ -99,4 +143,17 @@ async def invocations(raw_request: Request): def sagemaker_standards_bootstrap(app: FastAPI) -> FastAPI: - return sagemaker_standards.bootstrap(app) + """Bootstrap the app with the SageMaker hosting standards. + + Handler levels are restored right after the import because importing + model_hosting_container_standards may reconfigure root logging, and + bootstrap must run with the original levels. + """ + snapshot = _snapshot_handler_levels() + try: + import model_hosting_container_standards.sagemaker as sagemaker_standards + + app = sagemaker_standards.bootstrap(app) + finally: + _restore_handler_levels(snapshot) + return app From f37c550bf6353d7d2a7289cbf256943e8c282fad Mon Sep 17 00:00:00 2001 From: JulienDarve <86800349+JulienDarve@users.noreply.github.com> Date: Wed, 16 Sep 2026 01:29:43 -0700 Subject: [PATCH 12/20] [Feature][Frontend] Expose effective attention block size for DCP (#56538) Co-authored-by: Bugen Zhao Co-authored-by: OpenAI Codex Signed-off-by: Julien Darve Signed-off-by: Bugen Zhao --- docs/serving/context_parallel_deployment.md | 9 +++++ rust/Cargo.toml | 2 +- rust/proto/README.md | 6 ++- rust/proto/control.proto | 2 + rust/src/engine-core-client/src/client.rs | 9 +++++ .../src/engine-core-client/src/mock_engine.rs | 1 + .../src/protocol/handshake.rs | 3 ++ .../engine-core-client/src/tests/client.rs | 6 +++ .../src/tests/python_compat.py | 2 + rust/src/server/src/grpc/control.rs | 1 + rust/src/server/src/grpc/tests.rs | 13 ++++++- tests/config/test_config_utils.py | 3 ++ tests/v1/core/test_kv_cache_utils.py | 38 +++++++++++++++++++ tests/v1/engine/test_engine_core_client.py | 31 +++++++++++++-- vllm/config/cache.py | 3 ++ vllm/v1/engine/__init__.py | 2 + vllm/v1/engine/core.py | 21 +++++++++- vllm/v1/engine/core_client.py | 18 +++++++-- 18 files changed, 159 insertions(+), 11 deletions(-) diff --git a/docs/serving/context_parallel_deployment.md b/docs/serving/context_parallel_deployment.md index 3df2b06635f5..cb9a1ece71a4 100644 --- a/docs/serving/context_parallel_deployment.md +++ b/docs/serving/context_parallel_deployment.md @@ -20,6 +20,15 @@ Both approaches are under active development. Due to the auto-regressive nature of decoding, every decoding step needs to compute a small amount of query tokens w.r.t. a large number of key/value tokens stored in the paged KV cache. The core of decode context parallel is how to shard the KV cache across GPUs. +`engine_client.vllm_config.cache_config.effective_attention_block_size` reports the +initialized full-attention block size in tokens, including DCP: a physical block of +16 tokens with DCP=4 represents 64 tokens. The same value is available through +gRPC `Control.GetServerInfo.effective_attention_block_size`. Physical block-size +fields retain their existing meaning. The new value is unavailable (`None` in +Python, absent in gRPC) if there is no common full-attention size, the scheduler +does not expose it, or an older engine omits it. gRPC also omits the value when +engine replicas disagree. Partial cache events carry their own actual size. + For a model with `H` kv-heads, a request with `T` tokens in the context needs to store `H * T` key/value tensors in the KV cache. 1. If one GPU can hold them all, and the performance is good enough, then no parallelization is needed. diff --git a/rust/Cargo.toml b/rust/Cargo.toml index bd9d99de623c..dde065e79044 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -147,7 +147,7 @@ vllm-llm = { path = "src/llm" } vllm-managed-engine = { path = "src/managed-engine" } vllm-metrics = { path = "src/metrics" } vllm-parser = { path = "src/parser" } -vllm-proto = { version = "0.2.0", path = "proto" } +vllm-proto = { path = "proto" } vllm-server = { path = "src/server" } vllm-text = { path = "src/text" } vllm-tokenizer = { path = "src/tokenizer" } diff --git a/rust/proto/README.md b/rust/proto/README.md index 3181cdc6a07b..b8a20b0c349b 100644 --- a/rust/proto/README.md +++ b/rust/proto/README.md @@ -2,6 +2,9 @@ This directory is the canonical source for vLLM's gRPC schema. +See [context parallel deployment](../../docs/serving/context_parallel_deployment.md) +for effective attention block-size metadata in Python and `Control.GetServerInfo`. + Schema updates are no longer published to the Buf Schema Registry. Rust consumers should use `vllm-proto` from crates.io; consumers in other languages can generate bindings from the `.proto` files in this directory. Buf still builds and lints @@ -34,8 +37,7 @@ On pull requests and releases, `cargo-semver-checks` compares the crate with its latest published version. Include any required version bump in the protocol change PR. This check becomes available after the first manual publication. -1. Update the crate version and the `vllm-proto` workspace dependency together, - and update `rust/Cargo.lock`. +1. Update the crate version in `rust/proto/Cargo.toml` and update `rust/Cargo.lock`. 2. Run `cargo publish --manifest-path rust/proto/Cargo.toml --locked --dry-run` and the frontend gRPC tests. Record the tested vLLM releases or revisions in the release notes; matching crate versions alone do not establish runtime compatibility. diff --git a/rust/proto/control.proto b/rust/proto/control.proto index 3050035fdcf6..d2816f68d347 100644 --- a/rust/proto/control.proto +++ b/rust/proto/control.proto @@ -50,6 +50,8 @@ message ServerInfo { uint64 max_batched_tokens = 9; uint32 max_loras = 10; RlCapabilities rl_capabilities = 11; + // Full-attention block size in tokens after initialization; absent if unavailable. + optional uint64 effective_attention_block_size = 14; } message RlCapabilities { diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index 63689ea3db83..4cd7420e561e 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -510,6 +510,15 @@ impl EngineCoreClient { self.engines.iter().map(|engine| engine.ready_response.num_gpu_blocks).sum() } + /// Return the effective attention block size if all engines report the same value. + pub fn effective_attention_block_size(&self) -> Option { + let size = self.ready_response().effective_attention_block_size?; + self.engines + .iter() + .all(|engine| engine.ready_response.effective_attention_block_size == Some(size)) + .then_some(size) + } + /// Return the minimum engine-reported `max_model_len` across all engines. /// /// This is the auto-fitted value after KV cache profiling and may differ diff --git a/rust/src/engine-core-client/src/mock_engine.rs b/rust/src/engine-core-client/src/mock_engine.rs index 77d72b525fbd..958fcf8aa1b3 100644 --- a/rust/src/engine-core-client/src/mock_engine.rs +++ b/rust/src/engine-core-client/src/mock_engine.rs @@ -72,6 +72,7 @@ pub fn default_ready_response() -> EngineCoreReadyResponse { weight_transfer_backend: None, enable_sleep_mode: false, supports_draft_weight_updates: false, + effective_attention_block_size: None, } } diff --git a/rust/src/engine-core-client/src/protocol/handshake.rs b/rust/src/engine-core-client/src/protocol/handshake.rs index c14c19a50373..52d897be52dc 100644 --- a/rust/src/engine-core-client/src/protocol/handshake.rs +++ b/rust/src/engine-core-client/src/protocol/handshake.rs @@ -104,6 +104,9 @@ pub struct EngineCoreReadyResponse { /// Whether the engine has a speculative draft model that can be updated. #[serde(default)] pub supports_draft_weight_updates: bool, + /// Full-attention block size in tokens after initialization, or unavailable. + #[serde(default)] + pub effective_attention_block_size: Option, } /// Frontend-owned ZMQ addresses that are sent to the engine during startup diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 27fbbe9abd93..3db573d2aadc 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -2846,6 +2846,11 @@ fn python_msgpack_fixtures_match_rust_encoding() { let ready_response: EngineCoreReadyResponse = rmp_serde::from_slice(&hex::decode(ready_response_hex).unwrap()).unwrap(); + let mut legacy_ready = serde_json::to_value(&ready_response).unwrap(); + legacy_ready.as_object_mut().unwrap().remove("effective_attention_block_size"); + let legacy_ready: EngineCoreReadyResponse = + rmp_serde::from_slice(&rmp_serde::to_vec_named(&legacy_ready).unwrap()).unwrap(); + assert!(legacy_ready.effective_attention_block_size.is_none()); assert!(ready_response.supports_lora); assert_eq!(ready_response.max_loras, 8); assert_eq!( @@ -2854,6 +2859,7 @@ fn python_msgpack_fixtures_match_rust_encoding() { ); assert!(ready_response.enable_sleep_mode); assert!(ready_response.supports_draft_weight_updates); + assert_eq!(ready_response.effective_attention_block_size, Some(64)); let kv_events_config = ready_response.kv_events_config.expect("KV events config should decode"); assert!(kv_events_config.enable_kv_cache_events); assert_eq!(kv_events_config.publisher, "zmq"); diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py index 20cec81808cc..0cebca972d4d 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -423,6 +423,7 @@ class EngineCoreReadyResponse: weight_transfer_backend: str | None = None enable_sleep_mode: bool = False supports_draft_weight_updates: bool = False + effective_attention_block_size: int | None = None ready_response = EngineCoreReadyResponse( @@ -446,6 +447,7 @@ class EngineCoreReadyResponse: weight_transfer_backend="nccl", enable_sleep_mode=True, supports_draft_weight_updates=True, + effective_attention_block_size=64, kv_events_config=KVEventsConfig( enable_kv_cache_events=True, publisher="zmq", diff --git a/rust/src/server/src/grpc/control.rs b/rust/src/server/src/grpc/control.rs index ddb74bf40f1d..de9c5ae65fb9 100644 --- a/rust/src/server/src/grpc/control.rs +++ b/rust/src/server/src/grpc/control.rs @@ -195,6 +195,7 @@ impl pb::control_server::Control for ControlServiceImpl { max_batched_tokens: ready.max_num_batched_tokens, max_loras: ready.max_loras, rl_capabilities: Some(self.rl_capabilities()), + effective_attention_block_size: self.client().effective_attention_block_size(), })) } diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index f5f74d557e69..a21e013323d4 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -1911,8 +1911,16 @@ async fn control_abort_resolves_external_id_and_empty_is_noop() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn control_reports_server_and_model_info() { + let mut ready = default_ready_response(); + ready.effective_attention_block_size = Some(64); let (generate_service, control_service, engine_health, _engine_task) = - setup_grpc_service(b"engine-grpc-info", default_stream_output_specs()).await; + setup_grpc_service_with_engine_script( + b"engine-grpc-info".to_vec(), + ready, + Arc::new(FakeTextBackend), + |_, _| boxed_test_future(async {}), + ) + .await; let (channel, server_task) = start_grpc_test_server( generate_service, control_service, @@ -1928,6 +1936,7 @@ async fn control_reports_server_and_model_info() { .expect("get server info") .into_inner(); assert_eq!(server.engine_version, "test-vllm-version"); + assert_eq!(server.effective_attention_block_size, Some(64)); assert_eq!(server.api_version, "vllm"); assert_eq!(server.instance_id, "test-instance"); assert_eq!(server.max_model_len, DEFAULT_MOCK_MAX_MODEL_LEN as u32); @@ -2195,6 +2204,7 @@ async fn control_aggregates_multi_engine_capacity() { ready_0.weight_transfer_backend = Some("nccl".to_string()); ready_0.enable_sleep_mode = true; ready_0.supports_draft_weight_updates = true; + ready_0.effective_attention_block_size = Some(64); let mut ready_1 = default_ready_response(); ready_1.max_model_len = 4_096; @@ -2248,6 +2258,7 @@ async fn control_aggregates_multi_engine_capacity() { .into_inner(); assert_eq!(server.max_model_len, 4_096); assert_eq!(server.total_kv_blocks, 30); + assert!(server.effective_attention_block_size.is_none()); let rl = server.rl_capabilities.expect("RL capabilities"); assert!(!rl.weight_transfer_enabled); assert!(rl.weight_transfer_backend.is_empty()); diff --git a/tests/config/test_config_utils.py b/tests/config/test_config_utils.py index 94060d1724a8..590e97c1eb54 100644 --- a/tests/config/test_config_utils.py +++ b/tests/config/test_config_utils.py @@ -215,6 +215,9 @@ def test_cache_config_hash_ignores_kv_cache_sizing_knobs(): base_hash = CacheConfig().compute_hash() assert CacheConfig(kv_cache_memory_bytes=1 << 30).compute_hash() == base_hash assert CacheConfig(gpu_memory_utilization=0.5).compute_hash() == base_hash + config = CacheConfig() + config.effective_attention_block_size = 64 + assert config.compute_hash() == base_hash def test_scheduler_config_hash_includes_max_num_seqs(): diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index c6354cdc0fa0..86e66e1941f1 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -279,6 +279,44 @@ def make_request( ) +@pytest.mark.parametrize("dcp", [1, 4]) +def test_effective_attention_block_size_matches_events(dcp): + from vllm.distributed.kv_events import BlockStored + from vllm.v1.engine.core import EngineCore + + config = KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(["attention"], new_kv_cache_spec()), + ], + ) + manager = KVCacheManager( + generate_scheduler_kv_cache_config([config]), + max_model_len=256, + scheduler_block_size=16 * dcp, + hash_block_size=16 * dcp, + dcp_world_size=dcp, + enable_kv_cache_events=True, + ) + core = EngineCore.__new__(EngineCore) + core.vllm_config = SimpleNamespace(cache_config=CacheConfig(block_size=16)) + core.scheduler = SimpleNamespace(kv_cache_manager=manager) + core._initialize_effective_attention_block_size() + block_size = core.vllm_config.cache_config.effective_attention_block_size + assert block_size == 16 * dcp + + request = make_request( + "block-size", list(range(64)), block_size=16 * dcp, hash_fn=sha256 + ) + assert manager.allocate_slots(request, 64) is not None + assert [ + event.block_size + for event in manager.take_events() + if isinstance(event, BlockStored) + ] == [block_size] + + def new_kv_cache_spec( block_size=16, num_kv_heads=2, diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index f97ed76c5c48..3f407457b22a 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -158,7 +158,10 @@ def setsockopt(self, *_args, **_kwargs): local_engines_only=False, enable_elastic_ep=False, ) - vllm_config = SimpleNamespace(parallel_config=parallel_config) + vllm_config = SimpleNamespace( + parallel_config=parallel_config, + cache_config=SimpleNamespace(effective_attention_block_size=None), + ) client = core_client_mod.MPClient( asyncio_mode=False, @@ -368,10 +371,15 @@ def test_dplb_finished_requests_release_inflight(): assert req.request_id not in client.reqs_in_flight -def test_apply_ready_response_syncs_block_size(): +@pytest.mark.parametrize( + ("effective_size", "other_size"), + [(None, None), (4224, 4224), (4224, 1056), (4224, None)], +) +def test_apply_ready_response_syncs_block_size(effective_size, other_size): import msgspec client = object.__new__(MPClient) + client._effective_attention_block_sizes = set() client.vllm_config = SimpleNamespace( cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0), model_config=SimpleNamespace(max_model_len=8192), @@ -399,14 +407,31 @@ def test_apply_ready_response_syncs_block_size(): max_loras=0, ) ) - client._apply_ready_response(payload) + fields = msgspec.msgpack.decode(payload) + if effective_size is None: + del fields["effective_attention_block_size"] + else: + fields["effective_attention_block_size"] = effective_size + client._apply_ready_response(msgspec.msgpack.encode(fields)) assert client.vllm_config.cache_config.block_size == 1056 + cache_config = client.vllm_config.cache_config + assert cache_config.effective_attention_block_size == effective_size + + fields["effective_attention_block_size"] = other_size + client._apply_ready_response(msgspec.msgpack.encode(fields)) + assert cache_config.effective_attention_block_size == ( + effective_size if effective_size == other_size else None + ) + + client._apply_ready_response(b"") + assert cache_config.effective_attention_block_size is None def test_apply_ready_response_syncs_mamba_block_size(): import msgspec client = object.__new__(MPClient) + client._effective_attention_block_sizes = set() client.vllm_config = SimpleNamespace( cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0), model_config=SimpleNamespace(max_model_len=8192), diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 9bd898731eb1..ce3199c52a25 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -211,6 +211,8 @@ class CacheConfig: """The number of blocks to allocate for CPU memory.""" # Set after KV cache initialization. + effective_attention_block_size: int | None = field(default=None, init=False) + """Full-attention block size in tokens, including DCP, or None if unavailable.""" kv_cache_size_tokens: int | None = field(default=None, init=False) """Per-DP-engine KV cache capacity in tokens (group-aware). Uses group-aware capacity since num_gpu_blocks * block_size can be wrong @@ -277,6 +279,7 @@ def compute_hash(self) -> str: # Post-init/derived counters "num_gpu_blocks", "num_cpu_blocks", + "effective_attention_block_size", "kv_cache_size_tokens", "kv_cache_max_concurrency", # WIP feature toggle not impacting compiled graph shape diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index 5ae9ee0cac83..7d68b6e7978b 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -102,6 +102,8 @@ class EngineCoreReadyResponse: weight_transfer_backend: str | None = None enable_sleep_mode: bool = False supports_draft_weight_updates: bool = False + # Full-attention block size in tokens after initialization, or unavailable. + effective_attention_block_size: int | None = None class EngineCoreRequest( diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 62b6ef880e29..b3d94de3e7db 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -85,7 +85,7 @@ EngineCoreSentinel, fault_tolerant_wrapper, ) -from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.kv_cache_interface import KVCacheConfig, is_full_attention_spec from vllm.v1.metrics.stats import SchedulerIterationDetails, SchedulerStats from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus @@ -168,6 +168,7 @@ def __init__( block_size=scheduler_block_size, hash_block_size=hash_block_size, ) + self._initialize_effective_attention_block_size() self.use_spec_decode = vllm_config.speculative_config is not None self.check_for_draft_tokens = ( self.use_spec_decode or vllm_config.model_config.is_diffusion @@ -382,6 +383,21 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig: ) return scheduler_kv_cache_config + def _initialize_effective_attention_block_size(self) -> None: + cache_config = self.vllm_config.cache_config + cache_config.effective_attention_block_size = None + cache_manager = getattr(self.scheduler, "kv_cache_manager", None) + if cache_manager is None: + return + block_sizes = { + manager.block_size + for manager in cache_manager.coordinator.single_type_managers + if is_full_attention_spec(manager.kv_cache_spec) + } + cache_config.effective_attention_block_size = ( + block_sizes.pop() if len(block_sizes) == 1 else None + ) + def get_supported_tasks(self) -> tuple[SupportedTask, ...]: supported_tasks = self.model_executor.supported_tasks self._log_pooler_config(supported_tasks) @@ -1649,6 +1665,9 @@ def _make_ready_response(self) -> EngineCoreReadyResponse: num_gpu_blocks=self.vllm_config.cache_config.num_gpu_blocks or 0, block_size=self.vllm_config.cache_config.block_size, mamba_block_size=self.vllm_config.cache_config.mamba_block_size, + effective_attention_block_size=( + self.vllm_config.cache_config.effective_attention_block_size + ), dp_stats_address=self.frontend_stats_publish_address, dtype=str(self.vllm_config.model_config.dtype).removeprefix("torch."), vllm_version=VLLM_VERSION, diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index c354ba699c75..b63e8016c470 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -566,6 +566,7 @@ def __init__( ): self.vllm_config = vllm_config self._renderer: BaseRenderer | None = renderer + self._effective_attention_block_sizes: set[int | None] = set() # ZMQ setup. sync_ctx = zmq.Context(io_threads=2) @@ -803,10 +804,21 @@ def monitor_engine_cores(): def _apply_ready_response(self, payload: bytes) -> None: """Decode an EngineCoreReadyResponse and sync any post-initialization config changes (e.g. auto-fitted max_model_len) back to the frontend.""" - if not payload: - return vllm_config = self.vllm_config - response = msgspec.msgpack.decode(payload, type=EngineCoreReadyResponse) + response = ( + msgspec.msgpack.decode(payload, type=EngineCoreReadyResponse) + if payload + else None + ) + self._effective_attention_block_sizes.add( + response.effective_attention_block_size if response is not None else None + ) + sizes = self._effective_attention_block_sizes + vllm_config.cache_config.effective_attention_block_size = ( + next(iter(sizes)) if len(sizes) == 1 else None + ) + if response is None: + return vllm_config.model_config.max_model_len = min( vllm_config.model_config.max_model_len, response.max_model_len ) From 35f6047c1a2df9bb26ebbefc2b4e64717010e0fe Mon Sep 17 00:00:00 2001 From: Jakub Byczkowski Date: Wed, 16 Sep 2026 10:54:02 +0200 Subject: [PATCH 13/20] [XPU][Bugfix] Fix Qwen2-Audio ValueError on audio clips longer than 30s (#56912) Signed-off-by: Jakub Byczkowski Signed-off-by: Kunshang Ji Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Kunshang Ji --- vllm/model_executor/models/qwen2_audio.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/model_executor/models/qwen2_audio.py b/vllm/model_executor/models/qwen2_audio.py index 74341575e744..639db2884706 100644 --- a/vllm/model_executor/models/qwen2_audio.py +++ b/vllm/model_executor/models/qwen2_audio.py @@ -253,6 +253,7 @@ def _preprocess_hf_mm_data( hf_processor_mm_kwargs = dict( **hf_processor_mm_kwargs, sampling_rate=feature_extractor.sampling_rate, + truncation=True, ) return mm_data, hf_processor_mm_kwargs From 0384e72693985d0eddae70bd330b9cb283ff4313 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Wed, 16 Sep 2026 04:54:10 -0400 Subject: [PATCH 14/20] [UX] Add thinking support to `vllm chat` (#57045) --- vllm/entrypoints/cli/openai.py | 24 +++++++++++++++++++----- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/vllm/entrypoints/cli/openai.py b/vllm/entrypoints/cli/openai.py index fa01ff534b4c..b4e286838524 100644 --- a/vllm/entrypoints/cli/openai.py +++ b/vllm/entrypoints/cli/openai.py @@ -46,21 +46,35 @@ def _interactive_cli(args: argparse.Namespace) -> tuple[str, OpenAI]: def _print_chat_stream(stream, stats: bool = False) -> str: - output = "" - start = time.perf_counter() + output: str = "" + in_reasoning: bool = False + start: float = time.perf_counter() ttft: float | None = None - completion_tokens = 0 + completion_tokens: int = 0 for chunk in stream: if chunk.usage is not None: completion_tokens = chunk.usage.completion_tokens if not chunk.choices: continue delta = chunk.choices[0].delta + reasoning = getattr(delta, "reasoning", None) or getattr( + delta, "reasoning_content", None + ) + if ttft is None and (reasoning or delta.content): + ttft = time.perf_counter() - start + if reasoning: + if not in_reasoning: + print("", flush=True) + in_reasoning = True + print(reasoning, end="", flush=True) if delta.content: - if ttft is None: - ttft = time.perf_counter() - start + if in_reasoning: + print("\n", flush=True) + in_reasoning = False output += delta.content print(delta.content, end="", flush=True) + if in_reasoning: + print("\n", end="", flush=True) print() if stats: _print_metrics(start, ttft, completion_tokens) From 03f67b3ad1e61d12c21d42e20e69cb5f99bdd82e Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Wed, 16 Sep 2026 17:08:42 +0800 Subject: [PATCH 15/20] [Perf][DSV4.1] Pad shared experts for native MegaMoE fusion (#56568) Signed-off-by: Canlin Co-authored-by: Codex --- tests/models/test_deepseek_v4_mega_moe.py | 98 ++++++++++++++++++----- vllm/models/deepseek_v4/nvidia/model.py | 53 ++++++++++-- 2 files changed, 128 insertions(+), 23 deletions(-) diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 0d3c2b658369..9a5723d233a1 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -367,17 +367,21 @@ def test_deepseek_v4_mega_moe_padding_preserves_weights(monkeypatch, intermediat [ pytest.param(128, 512, 128, True, False, id="aligned-block128"), pytest.param(128, 512, 128, True, True, id="aligned-block128-mxfp8"), - pytest.param(5120, 2304, 32, False, False, id="deepseek-v41-flash"), + pytest.param(5120, 2304, 32, True, False, id="deepseek-v41-flash"), + pytest.param(128, 2304, 128, True, False, id="padded-block128"), + pytest.param(5120, 2304, 32, True, True, id="deepseek-v41-flash-mxfp8"), + pytest.param(128, 2304, 128, False, False, id="padded-packed-scale-fallback"), ], ) def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights( monkeypatch, hidden_size, intermediate_size, block_size, fused, mxfp8 ): - """V4.1-Flash's routed padding must preserve the serial shared-expert fallback.""" + """Shared fusion preserves checkpoint channels and zeros padded channels.""" class FakeDeepGemm: transformed_dims: list[tuple[int, int]] = [] scale_inputs: list[tuple[int, ...]] = [] + scale_values: list[torch.Tensor] = [] @staticmethod def get_symm_buffer_for_mega_moe(*args, num_shared_experts=0, **kwargs): @@ -402,6 +406,7 @@ def fp8_fp4_mega_moe( @classmethod def transform_sf_into_required_layout(cls, sf, mn, k, *args, **kwargs): cls.scale_inputs.append(tuple(sf.shape)) + cls.scale_values.append(sf) return torch.empty((sf.shape[0], mn, k // 128), dtype=torch.int32) @classmethod @@ -429,7 +434,10 @@ def transform_weights_for_mega_moe(cls, l1_weights, l2_weights): def fp8_parameter(*shape): return torch.nn.Parameter( - torch.empty(*shape, dtype=torch.float8_e4m3fn), requires_grad=False + torch.empty(*shape, dtype=torch.uint8) + .random_(1, 119) + .view(torch.float8_e4m3fn), + requires_grad=False, ) def scale_parameter(*shape, dtype=torch.int32): @@ -460,21 +468,34 @@ def scale_parameter(*shape, dtype=torch.int32): linear.weight_block_size = (1, 32) linear.weight_scale = torch.nn.Parameter( linear.weight_scale_inv.view(torch.uint8) - .repeat_interleave(128, dim=0) - .repeat_interleave(4, dim=1), + .repeat_interleave(block_size, dim=0) + .repeat_interleave(block_size // 32, dim=1), requires_grad=False, ) del linear.weight_scale_inv + originals = [] + for linear in (shared_experts.gate_up_proj, shared_experts.down_proj): + scale = linear.weight_scale if mxfp8 else linear.weight_scale_inv + if not fused: + # Already packed scales cannot be padded as checkpoint UE8M0 bytes. + linear.weight_scale_inv = scale_parameter( + linear.weight.shape[0], linear.weight.shape[1] // 128 + ) + continue + scale.data.view(torch.uint8).random_(124, 130) + block_m, block_k = linear.weight_block_size + scale_1x32 = ( + experts._ue8m0_uint8_to_float(scale.view(torch.uint8)) + .repeat_interleave(block_m, dim=0) + .repeat_interleave(block_k // 32, dim=1) + ) + originals.append((linear.weight.view(torch.uint8).clone(), scale_1x32)) monkeypatch.setattr("vllm.utils.deep_gemm._import_deep_gemm", lambda: FakeDeepGemm) original_gate_up_ptr = shared_experts.gate_up_proj.weight.data_ptr() original_down_ptr = shared_experts.down_proj.weight.data_ptr() experts.finalize_weights(shared_experts) - assert FakeDeepGemm.scale_inputs[-2:] == [ - (1, 2 * intermediate_size, hidden_size // 32), - (1, hidden_size, intermediate_size // 32), - ] assert experts.has_fused_shared_experts is fused if not fused: assert experts.intermediate_size == 2560 @@ -490,15 +511,56 @@ def scale_parameter(*shape, dtype=torch.int32): return assert FakeDeepGemm.transformed_dims == [(3, 3), (2, 2)] - assert shared_experts.gate_up_proj.weight.data_ptr() != original_gate_up_ptr - assert ( - experts._transformed_shared_l1_weights[0].data_ptr() - == shared_experts.gate_up_proj.weight.data_ptr() - ) - assert ( - experts._transformed_shared_l2_weights[0].data_ptr() - == shared_experts.down_proj.weight.data_ptr() - ) + padded_size = experts.intermediate_size + assert FakeDeepGemm.scale_inputs[-2:] == [ + (1, 2 * padded_size, hidden_size // 32), + (1, hidden_size, padded_size // 32), + ] + for actual, original, fill in ( + ( + experts._transformed_shared_l1_weights[0].view(torch.uint8), + originals[0][0], + 0, + ), + (FakeDeepGemm.scale_values[-2][0], originals[0][1], 1), + ): + actual = actual.unflatten(0, (2, padded_size)) + assert torch.equal( + actual[:, :intermediate_size], original.unflatten(0, (2, intermediate_size)) + ) + assert torch.all(actual[:, intermediate_size:] == fill) + for actual, original, fill in ( + ( + experts._transformed_shared_l2_weights[0].view(torch.uint8), + originals[1][0], + 0, + ), + (FakeDeepGemm.scale_values[-1][0], originals[1][1], 1), + ): + assert torch.equal(actual[:, : original.shape[1]], original) + assert torch.all(actual[:, original.shape[1] :] == fill) + if padded_size != intermediate_size: + # Generic linear post-load hooks still receive checkpoint-shaped weights. + for linear, (weight, _) in zip( + (shared_experts.gate_up_proj, shared_experts.down_proj), originals + ): + assert torch.equal(linear.weight.view(torch.uint8), weight) + assert shared_experts.gate_up_proj.weight.data_ptr() == original_gate_up_ptr + assert shared_experts.down_proj.weight.data_ptr() == original_down_ptr + else: + assert shared_experts.gate_up_proj.weight.data_ptr() != original_gate_up_ptr + assert ( + experts._transformed_shared_l1_weights[0].data_ptr() + == shared_experts.gate_up_proj.weight.data_ptr() + ) + assert ( + experts._transformed_shared_l2_weights[0].data_ptr() + == shared_experts.down_proj.weight.data_ptr() + ) + + transformed_l1 = experts._transformed_shared_l1_weights + experts.finalize_weights(shared_experts) + assert experts._transformed_shared_l1_weights is transformed_l1 @pytest.mark.parametrize("fused", [False, True]) diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index b1379cc9eb45..4aba14cd03b1 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -204,6 +204,7 @@ def __init__( self.top_k = top_k self.hidden_size = hidden_size self.intermediate_size = intermediate_size + self.unpadded_intermediate_size = intermediate_size self.num_shared_experts = num_shared_experts self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens @@ -386,6 +387,17 @@ def _finalize_shared_expert_weights( # the generic linear post-load hook replaces the raw checkpoint scales # with its 128x128 DeepGEMM layout. checkpoint_scale_dtypes = (torch.float8_e8m0fnu, torch.uint8) + unpadded_size = self.unpadded_intermediate_size * self.num_shared_experts + padding = self.intermediate_size * self.num_shared_experts - unpadded_size + pad_weights = ( + padding > 0 + and gate_up_weight.dtype == torch.float8_e4m3fn + and down_weight.dtype == torch.float8_e4m3fn + and gate_up_weight.shape == (2 * unpadded_size, self.hidden_size) + and down_weight.shape == (self.hidden_size, unpadded_size) + and gate_up_scale.dtype in checkpoint_scale_dtypes + and down_scale.dtype in checkpoint_scale_dtypes + ) if ( gate_up_scale.dtype in checkpoint_scale_dtypes and down_scale.dtype in checkpoint_scale_dtypes @@ -396,6 +408,7 @@ def _finalize_shared_expert_weights( gate_up_scale, gate_up_weight.shape[0], gate_up_weight.shape[1], + padding=(padding, 0) if pad_weights else (0, 0), ) down_scale = self._prepare_shared_expert_scale( deep_gemm, @@ -403,12 +416,27 @@ def _finalize_shared_expert_weights( down_scale, down_weight.shape[0], down_weight.shape[1], + padding=(0, padding) if pad_weights else (0, 0), ) if gate_up_scale is None or down_scale is None: self.num_shared_experts = 0 return + if pad_weights: + # Pad gate/up separately so the SwiGLU split stays at the midpoint. + gate_up_weight = ( + torch.nn.functional.pad( + gate_up_weight.view(torch.uint8).unflatten(0, (2, unpadded_size)), + (0, 0, 0, padding), + ) + .flatten(0, 1) + .view(gate_up_weight.dtype) + ) + down_weight = torch.nn.functional.pad( + down_weight.view(torch.uint8), (0, padding) + ).view(down_weight.dtype) + shared_intermediate_size = self.intermediate_size * self.num_shared_experts expected_gate_up_shape = ( 2 * shared_intermediate_size, @@ -449,11 +477,11 @@ def _finalize_shared_expert_weights( # released instead of adding roughly 0.7 GiB per rank on DSV4-Flash. # The generic linear post-load hook may still repack the serial scales, # but this shared MLP is never called after native fusion is enabled. - gate_up.weight.data = transformed_l1[0] - self._transformed_shared_l1_weights = ( - gate_up.weight.data, - transformed_l1[1], - ) + # Padded weights need separate storage: the generic linear post-load + # hooks still expect the original checkpoint weight/scale shapes. + if not pad_weights: + gate_up.weight.data = transformed_l1[0] + self._transformed_shared_l1_weights = transformed_l1 self._transformed_shared_l2_weights = transformed_l2 def _prepare_shared_expert_scale( @@ -463,6 +491,8 @@ def _prepare_shared_expert_scale( scale: torch.Tensor, mn: int, k: int, + *, + padding: tuple[int, int] = (0, 0), ) -> torch.Tensor | None: block_size = getattr(linear, "weight_block_size", None) if block_size is None or len(block_size) != 2: @@ -497,6 +527,19 @@ def _prepare_shared_expert_scale( .repeat_interleave(block_k // 32, dim=1)[:mn, : k // 32] .contiguous() ) + pad_m, pad_k = padding + if pad_m: + scale_1x32 = torch.nn.functional.pad( + scale_1x32.unflatten(0, (2, mn // 2)), + (0, 0, 0, pad_m), + value=1.0, + ).flatten(0, 1) + mn += 2 * pad_m + if pad_k: + scale_1x32 = torch.nn.functional.pad( + scale_1x32, (0, pad_k // 32), value=1.0 + ) + k += pad_k # The grouped API is used with a singleton dimension to request the # MN-major, TMA-aligned packed-UE8M0 strides, then squeezed back to the # 2D layout required for a shared expert. From 82731931ae1a770ef0439a2afa88a54cb850cdfa Mon Sep 17 00:00:00 2001 From: Ganesh R Date: Wed, 16 Sep 2026 14:45:12 +0530 Subject: [PATCH 16/20] [Bugfix][CPU] Fall back to CpuPlatform when zentorch fails to import (#54923) Signed-off-by: R Co-authored-by: Cursor --- tests/test_zen_cpu_platform_detection.py | 52 +++++++++++++++++++++++- vllm/platforms/__init__.py | 9 ++++ 2 files changed, 59 insertions(+), 2 deletions(-) diff --git a/tests/test_zen_cpu_platform_detection.py b/tests/test_zen_cpu_platform_detection.py index 9f386f70cd98..0c99c6568a2b 100644 --- a/tests/test_zen_cpu_platform_detection.py +++ b/tests/test_zen_cpu_platform_detection.py @@ -65,7 +65,7 @@ def test_cpu_target_selects_cpu_platform_from_non_cpu_wheel( rocm_plugin.assert_not_called() -def test_platform_detection_logs_zentorch_import_failure(caplog): +def test_broken_zentorch_falls_back_to_cpu_platform(caplog): original_import = builtins.__import__ def import_with_broken_zentorch(name, *args, **kwargs): @@ -88,6 +88,54 @@ def import_with_broken_zentorch(name, *args, **kwargs): ): platform = resolve_current_platform_cls_qualname() + assert platform == "vllm.platforms.cpu.CpuPlatform" + assert "zentorch failed to import" in caplog.text + assert "OSError: incompatible shared library" in caplog.text + + +def test_zentorch_failure_other_than_oserror_is_not_recovered(caplog): + original_import = builtins.__import__ + + def import_with_broken_zentorch(name, *args, **kwargs): + if name == "zentorch": + raise RuntimeError("unknown zentorch failure") + return original_import(name, *args, **kwargs) + + with ( + patch("vllm.platforms.envs.VLLM_TARGET_DEVICE", "cuda"), + patch.dict( + "vllm.platforms.builtin_platform_plugins", + {"cpu": cpu_platform_plugin}, + clear=True, + ), + patch("vllm.platforms.load_plugins_by_group", return_value={}), + patch("vllm.platforms.vllm_version_matches_substr", return_value=True), + patch("vllm.platforms._is_amd_zen_cpu", return_value=True), + patch.object(builtins, "__import__", side_effect=import_with_broken_zentorch), + caplog.at_level(logging.DEBUG, logger="vllm.platforms"), + ): + platform = resolve_current_platform_cls_qualname() + + assert platform == "vllm.platforms.interface.UnspecifiedPlatform" + assert "RuntimeError: unknown zentorch failure" in caplog.text + + +def test_platform_plugin_failure_is_logged(caplog): + def failing_plugin(): + raise RuntimeError("plugin exploded") + + with ( + patch("vllm.platforms.envs.VLLM_TARGET_DEVICE", "cuda"), + patch.dict( + "vllm.platforms.builtin_platform_plugins", + {"cpu": failing_plugin}, + clear=True, + ), + patch("vllm.platforms.load_plugins_by_group", return_value={}), + caplog.at_level(logging.DEBUG, logger="vllm.platforms"), + ): + platform = resolve_current_platform_cls_qualname() + assert platform == "vllm.platforms.interface.UnspecifiedPlatform" assert "Platform plugin cpu failed during detection" in caplog.text - assert "OSError: incompatible shared library" in caplog.text + assert "RuntimeError: plugin exploded" in caplog.text diff --git a/vllm/platforms/__init__.py b/vllm/platforms/__init__.py index b34b0b9a7dd9..7cf0843669c3 100644 --- a/vllm/platforms/__init__.py +++ b/vllm/platforms/__init__.py @@ -217,6 +217,15 @@ def cpu_platform_plugin() -> str | None: "AMD Zen CPU detected but zentorch not installed, " "falling back to CpuPlatform." ) + except OSError: + # An ABI-mismatched build fails here with an undefined-symbol + # error; other failures are not known to be safe to recover from. + logger.warning( + "AMD Zen CPU detected but zentorch failed to import, falling " + "back to CpuPlatform. This usually means the zentorch build " + "does not match the installed torch version.", + exc_info=True, + ) return "vllm.platforms.cpu.CpuPlatform" From 5f203baedb29bd97f96243f2b83e73fb1bfec7f0 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Wed, 16 Sep 2026 10:16:32 +0100 Subject: [PATCH 17/20] [Docs] Split slash-combined docstring parameters (#57141) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) --- vllm/_custom_ops.py | 10 ++++++++-- .../attention/dsa/sparse_mqa_logits.py | 20 +++++++++++++------ vllm/utils/deep_gemm.py | 4 +++- 3 files changed, 25 insertions(+), 9 deletions(-) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 3311215016a7..99577766e7a8 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -1482,8 +1482,14 @@ def cutlass_w4a8_moe_mm( Cumulative token offsets problem_sizes: Per-expert (M, N, K) GEMM sizes used by the grouped GEMM launcher. - a/b/c/group_scale_strides: - Strides describing the memory layout of the input tensors. + a_strides: + Strides describing the memory layout of a_tensors. + b_strides: + Strides describing the memory layout of b_tensors. + c_strides: + Strides describing the memory layout of out_tensors. + group_scale_strides: + Strides describing the memory layout of b_group_scales. maybe_schedule: Optional override to choose a specific kernel or epilogue schedule. diff --git a/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py b/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py index 839c581583a7..138833ad149e 100644 --- a/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py +++ b/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py @@ -190,9 +190,11 @@ def candidate_blocks_to_sparse_indices( candidate_blocks: [rows, K] int32 request-local candidate block ids (-1 padded), in units of ``candidate_block_size`` positions. K must be a power of two (the in-kernel sort); rows may be strided. - row_ks/row_ke: [rows] int32 per-row K range; bounds are in the same + row_ks: [rows] int32 per-row K range start; bounds are in the same (packed-workspace) coordinates the sparse kernel iterates over. Pass zeros for the paged path, whose blocks are context-relative. + row_ke: [rows] int32 per-row K range end, in the same coordinates as + ``row_ks``. candidate_block_size: Positions per candidate block. sparse_block_kv: Positions per sparse block (8 or 16). out: Optional ``(sparse_indices, end)`` buffers to write into, @@ -322,12 +324,16 @@ def sparse_mqa_logits_prefill_chunk( k_scale: [total_kv] int32 packed UE8M0 K scales. weights: [rows, H] bf16 per-head weights; the sparse kernels take bf16 and do not fold the Q scale in. - cu_seqlen_ks/cu_seqlen_ke: [rows] int32 per-token K bounds in the - packed workspace. + cu_seqlen_ks: [rows] int32 per-token K start bounds in the packed + workspace. + cu_seqlen_ke: [rows] int32 per-token K end bounds in the packed + workspace. candidate_blocks: [rows, K] int32 request-local candidate block ids. topk_indices: [rows, topk_tokens] output buffer. - sparse_indices/end/col_indices: Caller-owned scratch, see - `candidate_blocks_to_sparse_indices` and `sparse_topk_remap`. + sparse_indices: Caller-owned scratch, see + `candidate_blocks_to_sparse_indices`. + end: Caller-owned scratch, see `candidate_blocks_to_sparse_indices`. + col_indices: Caller-owned scratch, see `sparse_topk_remap`. kernel_metadata: DeepGEMM schedule from a previous call with the same candidates and bounds (i.e. another indexer layer in the same step). When given, the candidate expansion is skipped and @@ -413,7 +419,9 @@ def sparse_mqa_logits_paged_decode( candidate_blocks: [rows, K] int32 candidate block ids. topk_indices: [rows, topk_tokens] output buffer. row_ks: [rows] int32 zeros (paged blocks are context-relative). - sparse_indices/end/col_indices: Caller-owned scratch. + sparse_indices: Caller-owned scratch. + end: Caller-owned scratch. + col_indices: Caller-owned scratch. kernel_metadata: See `sparse_mqa_logits_prefill_chunk`. Returns: diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 91e315afb9e7..6fef0072243f 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -726,8 +726,10 @@ def get_sparse_mqa_logits_metadata( need a host-side sync to pick one. Args: - cu_seqlen_ks/cu_seqlen_ke: Per-row K range bounds in the packed KV + cu_seqlen_ks: Per-row K range start bounds in the packed KV workspace, shape [num_q_tokens], dtype int32. + cu_seqlen_ke: Per-row K range end bounds in the packed KV workspace, + shape [num_q_tokens], dtype int32. num_kv_tokens: Total KV tokens in the packed workspace. sparse_kv_block_indices: Per-row candidate block ids, shape [num_q_tokens, num_max_sparse_blocks], dtype int32. Each row's From ceb87de065b9cbdaffd34a0e6aed86c99e772b95 Mon Sep 17 00:00:00 2001 From: Adababy Date: Wed, 16 Sep 2026 17:17:07 +0800 Subject: [PATCH 18/20] [Misc] Remove no-op self-assignments across vLLM (#55988) Signed-off-by: shaolila Co-authored-by: shaolila --- .../layers/fused_moe/routed_experts.py | 15 ++++++++------- vllm/model_executor/models/deepencoder.py | 4 +--- vllm/model_executor/models/mimo_v2_omni.py | 1 - vllm/model_executor/models/step3p5.py | 1 - 4 files changed, 9 insertions(+), 12 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 2dafb6c493d6..257e19b034a4 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -615,14 +615,15 @@ def weight_loader( # compressed-tensors checkpoints with packed weights are stored flipped # TODO (mgoin): check self.quant_method.quant_config.quant_format # against known CompressionFormat enum values that have this quality - if quant_method_name in ( - "CompressedTensorsWNA16MoEMethod", - "CompressedTensorsW4A16FlydslMoEMethod", + if ( + quant_method_name + in ( + "CompressedTensorsWNA16MoEMethod", + "CompressedTensorsW4A16FlydslMoEMethod", + ) + and is_transposed ): - if is_transposed: - loaded_weight = loaded_weight.t().contiguous() - else: - loaded_weight = loaded_weight + loaded_weight = loaded_weight.t().contiguous() if shard_id not in ("w1", "w2", "w3"): raise ValueError(f"shard_id must be ['w1','w2','w3'] but got {shard_id}.") diff --git a/vllm/model_executor/models/deepencoder.py b/vllm/model_executor/models/deepencoder.py index 189e8b55e510..d75e5bb77d2f 100644 --- a/vllm/model_executor/models/deepencoder.py +++ b/vllm/model_executor/models/deepencoder.py @@ -780,9 +780,7 @@ def forward( self, pixel_values: torch.Tensor, patch_embeds: torch.Tensor | None = None ) -> torch.Tensor: batch_size = pixel_values.shape[0] - if patch_embeds is not None: - patch_embeds = patch_embeds - else: + if patch_embeds is None: patch_embeds = self.patch_embedding(pixel_values) patch_embeds = patch_embeds.flatten(2).transpose(1, 2) diff --git a/vllm/model_executor/models/mimo_v2_omni.py b/vllm/model_executor/models/mimo_v2_omni.py index 93bd6a44b556..9281a1689ae1 100644 --- a/vllm/model_executor/models/mimo_v2_omni.py +++ b/vllm/model_executor/models/mimo_v2_omni.py @@ -964,7 +964,6 @@ def _apply_hf_processor_main( assert isinstance(va_item, (list, tuple)) and len(va_item) == 2 vid, audio_src = va_item va_item = VideoAudioInput(video=vid, audio=audio_src) - vid = vid # Convert video frames to (TCHW, timestamps) if needed if ( isinstance(vid, tuple) diff --git a/vllm/model_executor/models/step3p5.py b/vllm/model_executor/models/step3p5.py index 3a928aa59daa..03343a89a144 100644 --- a/vllm/model_executor/models/step3p5.py +++ b/vllm/model_executor/models/step3p5.py @@ -173,7 +173,6 @@ def __init__( rope_scaling = None if sliding_window is not None and enable_sliding_window: - sliding_window = sliding_window if swa_num_attention_heads is not None: num_heads = swa_num_attention_heads self.total_num_heads = swa_num_attention_heads From d6a1677d5504244c566eb900ca605cb5511f4ab4 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Wed, 16 Sep 2026 03:07:50 -0700 Subject: [PATCH 19/20] [Model][DSv4.1] FlashMLA mega attention and the NVFP4 compressed KV cache (#56935) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Yongye Zhu Signed-off-by: 云挚 Signed-off-by: Yongye Zhu Co-authored-by: Yongye Zhu Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Yunxiao Ning <73625538+foraxe@users.noreply.github.com> Co-authored-by: AI Assistant --- .../kernels/benchmark_dsv41_mega_attn.py | 272 ++++++++++ ...deepseek_v4_qnorm_rope_kv_insert_kernel.cu | 206 +++++--- csrc/libtorch_stable/ops.h | 3 +- csrc/libtorch_stable/torch_bindings.cpp | 9 +- tests/kernels/test_compressor_kv_cache.py | 139 ++++- tests/kernels/test_dsv41_mega_attn_layouts.py | 71 +++ ..._fused_deepseek_v4_qnorm_rope_kv_insert.py | 183 +++++++ vllm/models/deepseek_v41/amd/rocm.py | 3 +- vllm/models/deepseek_v41/attention.py | 204 +++++-- .../deepseek_v41/common/ops/cache_utils.py | 98 +++- .../common/ops/fused_compress_quant_cache.py | 84 ++- .../deepseek_v41/common/ops/fused_layout.py | 86 +++ vllm/models/deepseek_v41/nvidia/dspark.py | 14 + .../nvidia/flash_mla_mega_attn.py | 497 ++++++++++++++++++ .../deepseek_v41/nvidia/flashinfer_sparse.py | 6 +- vllm/models/deepseek_v41/nvidia/flashmla.py | 3 +- vllm/models/deepseek_v41/nvidia/model.py | 30 +- vllm/models/deepseek_v41/sparse_mla.py | 51 ++ vllm/v1/attention/backends/mla/sparse_swa.py | 5 +- vllm/v1/attention/backends/registry.py | 3 + 20 files changed, 1811 insertions(+), 156 deletions(-) create mode 100644 benchmarks/kernels/benchmark_dsv41_mega_attn.py create mode 100644 tests/kernels/test_dsv41_mega_attn_layouts.py create mode 100644 vllm/models/deepseek_v41/common/ops/fused_layout.py create mode 100644 vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py diff --git a/benchmarks/kernels/benchmark_dsv41_mega_attn.py b/benchmarks/kernels/benchmark_dsv41_mega_attn.py new file mode 100644 index 000000000000..1b3b5d9c3848 --- /dev/null +++ b/benchmarks/kernels/benchmark_dsv41_mega_attn.py @@ -0,0 +1,272 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FlashMLA mega attention vs the split-KV decode pipeline (DeepSeek V4.1). + +Times the full set of per-layer ops each path runs in a decode step, captured +in a CUDA graph -- eager timing is dominated by launch overhead and flatters +whichever path launches fewer kernels. + + mega: fused_qnorm_rope_kv_rope_quant_insert (Q pad + KV insert) + + mega kernel + split-KV: fused_qnorm_rope_kv_rope_quant_insert (Q RoPE + pad + KV insert) + + flash_mla_with_kvcache + fused_inv_rope_fp8_quant + +Both arms fuse their Q preparation into the KV insert, so that op is charged +to both; the mega arm's Q side is only a zero-pad (nothing at all when the +shard is already at the kernel head count) because its kernel does the Q RoPE, +the inverse RoPE and the FP8 cast. + +``--local-heads`` is the variable that decides the outcome: the mega kernel +absorbs work proportional to the live head count while paying for the padded +one, so it wins at 64 live heads and is a wash at 16 (TP4 on a 64-head model). + +Both arms use the same KV records so the comparison is kernel-for-kernel: +a V4.1 fp8 (528 B) sliding-window cache and a compressed cache that is either +V4.1 fp8 or V4.1 NVFP4 (288 B) -- FlashMLA's SM100 sparse decode reads the +NVFP4 compressed record in both the fused and unfused kernels. + +Run: .venv/bin/python benchmarks/kernels/benchmark_dsv41_mega_attn.py +""" + +import argparse + +import torch + +import vllm.v1.attention.ops.flashmla as fm +from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import ( + fused_inv_rope_fp8_quant, +) +from vllm.models.deepseek_v41.common.ops import quantize_and_insert_k_cache +from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( + rope_quant_insert, +) +from vllm.models.deepseek_v41.common.ops.fused_layout import permute_q_to_fused +from vllm.models.deepseek_v41.nvidia.flash_mla_mega_attn import ( + alloc_mega_attn_output, + is_flashmla_mega_attn_supported, +) +from vllm.utils.math_utils import round_up + +HEAD_DIM, ROPE_DIM = 512, 64 +V41_BYTES, V41_FP4_BYTES = 528, 288 +SWA_BLOCK, COMPRESSED_BLOCK = 32, 128 + + +def make_cos_sin_cache(max_pos, device): + inv_freq = 1.0 / ( + 10000 ** (torch.arange(0, ROPE_DIM, 2, device=device).float() / ROPE_DIM) + ) + freqs = torch.outer(torch.arange(max_pos, device=device).float(), inv_freq) + return torch.cat([freqs.cos(), freqs.sin()], -1) + + +def empty_paged_cache(num_rows, block_size, bytes_per_token, device): + """A zeroed paged cache view [num_blocks, block_size, bytes_per_token].""" + num_blocks = (num_rows + block_size - 1) // block_size + 1 + page = round_up(block_size * bytes_per_token, 512) + backing = torch.zeros(num_blocks, page, dtype=torch.uint8, device=device) + return backing, backing.as_strided( + (num_blocks, block_size, bytes_per_token), (page, bytes_per_token, 1) + ) + + +def fill_cache(view, bytes_per_token, cos_sin, device): + n = view.shape[0] * view.shape[1] + rows = torch.randn(n, HEAD_DIM, device=device, dtype=torch.bfloat16) + slots = torch.arange(n, dtype=torch.int64, device=device) + if bytes_per_token == V41_FP4_BYTES: + # A page-aligned cache has more slots than the cos_sin table has rows; + # wrap so the insert never indexes past it. + positions = slots % cos_sin.shape[0] + rope_quant_insert(rows, positions, cos_sin, view, slots, 1) + else: + quantize_and_insert_k_cache( + rows, + view.reshape(view.shape[0], -1), + slots, + block_size=view.shape[1], + bytes_per_token=bytes_per_token, + ) + + +def time_graph(fn, iters=100, warmup=5): + """Capture ``fn`` in a CUDA graph and time its replay.""" + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + fn() + torch.cuda.current_stream().wait_stream(s) + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + fn() + for _ in range(warmup): + g.replay() + torch.accelerator.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + g.replay() + end.record() + torch.accelerator.synchronize() + return start.elapsed_time(end) * 1000 / iters + + +def bench_decode(s_q, device, topk_extra, extra_bytes, local_heads): + padded_heads, scale = 64, HEAD_DIM**-0.5 + n_wv_group = padded_heads // 8 + cos_sin = make_cos_sin_cache(65536, device) + positions = torch.randint(0, 65536, (s_q,), device=device, dtype=torch.int64) + pos32 = positions.to(torch.int32) + + # Per-step inputs: the wq_b output (local heads) and this step's KV row. + q_local = torch.randn( + s_q, local_heads, HEAD_DIM, device=device, dtype=torch.bfloat16 + ) + q_fused = permute_q_to_fused(q_local) + kv = torch.randn(s_q, HEAD_DIM, device=device, dtype=torch.bfloat16) + + n_swa = s_q * 128 + 128 + swa_backing, swa = empty_paged_cache(n_swa, SWA_BLOCK, V41_BYTES, device) + fill_cache(swa, V41_BYTES, cos_sin, device) + _, extra = empty_paged_cache(65536, COMPRESSED_BLOCK, extra_bytes, device) + fill_cache(extra, extra_bytes, cos_sin, device) + + slot_mapping = torch.arange(s_q, dtype=torch.int64, device=device) + swa_idx = torch.randint(0, n_swa, (s_q, 128), device=device, dtype=torch.int32) + swa_len = torch.full((s_q,), 128, device=device, dtype=torch.int32) + ex_idx = torch.randint( + 0, 65536, (s_q, topk_extra), device=device, dtype=torch.int32 + ) + ex_len = torch.full((s_q,), topk_extra, device=device, dtype=torch.int32) + sink = torch.zeros(padded_heads, device=device) + sched = fm.FlashMLASchedMeta() + out = alloc_mega_attn_output(s_q, n_wv_group, device) + swa4 = swa.unsqueeze(-2) + extra4 = extra.unsqueeze(-2) + swa_2d = swa_backing + + def mega(): + # One launch zero-pads the fused-layout Q and inserts the SWA KV. + q_pad = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q_fused, + kv, + swa_2d, + slot_mapping, + positions, + cos_sin, + 0 if padded_heads == local_heads else padded_heads, + 1e-6, + SWA_BLOCK, + False, # apply_q_norm + True, # kv_mxfp8 + False, # apply_q_rope: the mega kernel rotates Q itself + True, # is_q_interleaved + ) + torch.ops._flashmla_C.fused_norm_rope_attn_rope_cast_decode( + q_fused if padded_heads == local_heads else q_pad, + swa4, + swa_idx, + scale, + HEAD_DIM, + sink, + swa_len, + extra4, + ex_idx, + ex_len, + False, + 0.0, + pos32, + False, + 64, + cos_sin, + n_wv_group, + 32, + True, + True, + True, + out.data, + out.scale, + ) + + def split_kv(): + q_pad = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q_local, + kv, + swa_2d, + slot_mapping, + positions, + cos_sin, + padded_heads, + 1e-6, + SWA_BLOCK, + False, + True, + ) + o = fm.flash_mla_with_kvcache( + q=q_pad.unsqueeze(1), + k_cache=swa4, + block_table=None, + head_dim_v=HEAD_DIM, + tile_scheduler_metadata=sched, + cache_seqlens=None, + is_fp8_kvcache=True, + indices=swa_idx.view(s_q, 1, -1), + topk_length=swa_len, + softmax_scale=scale, + attn_sink=sink, + extra_k_cache=extra4, + extra_indices_in_kvcache=ex_idx.view(s_q, 1, -1), + extra_topk_length=ex_len, + )[0] + fused_inv_rope_fp8_quant( + o.squeeze(1)[:, :local_heads], + positions, + cos_sin, + n_groups=local_heads // 8, + heads_per_group=8, + quant_group_size=32, + tma_aligned_scales=True, + ) + + return time_graph(mega), time_graph(split_kv) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--decode-s-q", + nargs="+", + type=int, + default=[1, 2, 4, 8, 16, 32, 64, 128, 256, 512], + ) + parser.add_argument("--topk-extra", type=int, default=512) + parser.add_argument("--local-heads", type=int, default=16) + parser.add_argument( + "--extra-bytes", + type=int, + choices=[V41_BYTES, V41_FP4_BYTES], + default=V41_BYTES, + help="compressed-cache record: 528 (V4.1 fp8) or 288 (V4.1 NVFP4)", + ) + args = parser.parse_args() + ok, reason = is_flashmla_mega_attn_supported() + if not ok: + raise SystemExit(reason) + device = torch.device("cuda") + torch.manual_seed(0) + print( + f"decode, cudagraph (topk_swa=128, topk_extra={args.topk_extra}, " + f"local_heads={args.local_heads}, extra={args.extra_bytes}B)" + ) + print(f"{'s_q':>6} {'mega_us':>9} {'split_kv_us':>12} {'speedup':>8}") + for s_q in args.decode_s_q: + m, s = bench_decode( + s_q, device, args.topk_extra, args.extra_bytes, args.local_heads + ) + print(f"{s_q:>6d} {m:9.1f} {s:12.1f} {s / m:7.2f}x") + + +if __name__ == "__main__": + main() diff --git a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu index 344b82fbda2c..5a5caacce7a5 100644 --- a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu +++ b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -3,7 +3,10 @@ * SPDX-FileCopyrightText: Copyright contributors to the vLLM project * * Horizontally-fused DeepseekV4-MLA kernel: - * - Q side: optional per-head RMSNorm + GPT-J RoPE on last ROPE_DIM + * - Q side: optional per-head RMSNorm (apply_q_norm) and GPT-J RoPE on the + * last ROPE_DIM (apply_q_rope), plus the zero-pad to + * q_head_padded, in the head-major layout or (is_q_interleaved) + * FlashMLA's mega-attention chunk-interleaved one * - KV side: GPT-J RoPE on last ROPE_DIM + UE8M0 FP8 quant on NoPE + paged * cache insert * @@ -174,6 +177,30 @@ __device__ __forceinline__ float warpSum(float val) { return val; } +// Offset of the 16 elements a lane owns inside a Q tensor of `num_heads` +// heads. The default layout is head-major: a head's 512 dims are contiguous. +// `Q_INTERLEAVED` is FlashMLA's mega-attention layout, 32 chunks of +// [num_heads, 16] per token, which changes the address only -- lane L still +// owns dims [16L, 16L+16) of its head either way, so `dim_base` stays the +// true dim offset and the norm and RoPE below are equally valid in both. +// What it does change is where the padding heads land: at the tail of every +// chunk rather than after the live heads. +template +__device__ __forceinline__ int64_t qElemOffset(int const tokenIdx, + int const slotIdx, + int const laneId, + int const dim_base, + int const num_heads) { + if constexpr (Q_INTERLEAVED) { + return static_cast(tokenIdx) * num_heads * kHeadDim + + static_cast(laneId) * num_heads * kElemsPerLane + + static_cast(slotIdx) * kElemsPerLane; + } else { + return (static_cast(tokenIdx) * num_heads + slotIdx) * kHeadDim + + dim_base; + } +} + // ──────────────────────────────────────────────────────────────────────────── // Per-slot inner pipeline // ──────────────────────────────────────────────────────────────────────────── @@ -192,10 +219,16 @@ __device__ __forceinline__ float warpSum(float val) { // + paged-cache // insert) // -// `kv_mxfp8` is grid-uniform and only steers the KV branch, so it stays a -// runtime argument rather than doubling every (head count, q-norm) -// instantiation of a kernel whose hot path is the Q branch. -template +// `kv_mxfp8` and `apply_q_rope` are grid-uniform, so they stay runtime +// arguments rather than multiplying every (head count, q-norm, q-layout) +// instantiation of a kernel whose hot path is the Q branch. `Q_INTERLEAVED` +// cannot: it picks the address arithmetic, which has to fold. +// +// The three Q knobs are independent. The mega-attention layer uses +// Q_INTERLEAVED with both norm and RoPE off, because its attention kernel +// applies those itself -- not because the layout forbids them. +template __device__ __forceinline__ void processDeepseekV4Slot( uint4 v0, uint4 v1, int const tokenIdx, int const slotIdx, int const dim_base, int const laneId, int const num_heads_q, @@ -203,7 +236,7 @@ __device__ __forceinline__ void processDeepseekV4Slot( uint8_t* __restrict__ k_cache, int64_t const* __restrict__ slot_mapping, int64_t const* __restrict__ position_ids, float const* __restrict__ cos_sin_cache, int const cache_block_size, - int const kv_block_stride, bool const kv_mxfp8) { + int const kv_block_stride, bool const kv_mxfp8, bool const apply_q_rope) { using Converter = vllm::_typeConvert; bool const isKV = (slotIdx == kNumHeadsQPadded); bool const isPadQ = !isKV && (slotIdx >= num_heads_q); @@ -213,10 +246,8 @@ __device__ __forceinline__ void processDeepseekV4Slot( // zero literal is correct. Matches the live-Q branch's vectorized store. if (isPadQ) { scalar_t_in* dst = - q_out + - (static_cast(tokenIdx) * kNumHeadsQPadded + slotIdx) * - kHeadDim + - dim_base; + q_out + qElemOffset(tokenIdx, slotIdx, laneId, dim_base, + kNumHeadsQPadded); uint4 const zero4 = {0u, 0u, 0u, 0u}; *reinterpret_cast(dst) = zero4; *reinterpret_cast(dst + 8) = zero4; @@ -264,7 +295,9 @@ __device__ __forceinline__ void processDeepseekV4Slot( // ── GPT-J RoPE on dims [NOPE_DIM, HEAD_DIM) ───────────────────────────── // All math in fp32. cos_sin_cache is loaded as fp32 (its native storage). - bool const is_rope_lane = dim_base >= kNopeDim; + // KV is always rotated; Q only when the caller has not already had it + // rotated elsewhere (the mega-attention kernel does its own). + bool const is_rope_lane = dim_base >= kNopeDim && (isKV || apply_q_rope); if (is_rope_lane) { int64_t const pos = position_ids[tokenIdx]; constexpr int kHalfRope = kRopeDim / 2; @@ -314,10 +347,8 @@ __device__ __forceinline__ void processDeepseekV4Slot( make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1])); } scalar_t_in* dst = - q_out + - (static_cast(tokenIdx) * kNumHeadsQPadded + slotIdx) * - kHeadDim + - dim_base; + q_out + qElemOffset(tokenIdx, slotIdx, laneId, dim_base, + kNumHeadsQPadded); *reinterpret_cast(dst) = out0; *reinterpret_cast(dst + 8) = out1; } else { @@ -432,14 +463,17 @@ __device__ __forceinline__ void processDeepseekV4Slot( // `kNumHeadsQPadded` is a template parameter (compile-time constant) so the // divisions in the grid math and the KV-sentinel comparison fold to fast // constant operations. The launch wrapper dispatches the runtime value to -// the matching instantiation. +// the matching instantiation. `kNumHeadsQPadded == 0` leaves one slot per +// token and is the KV-insert-only instantiation, for a Q that is already in +// the shape its attention kernel wants. // // With DP padding, q/kv/position_ids can have more rows than slot_mapping. // The live-Q and pad-Q branches cover all `num_tokens_full` rows (downstream // attention uses them). The KV branch only inserts the first // `num_tokens_insert` tokens (= slot_mapping length) into the paged cache. // -template +template __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel( scalar_t_in const* __restrict__ q_in, // [N, num_heads_q, 512] scalar_t_in* __restrict__ q_out, // [N, kNumHeadsQPadded, 512] @@ -454,7 +488,8 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel( int const num_heads_q, // live Q heads (input layout) int const cache_block_size, // tokens per paged-cache block int const kv_block_stride, // bytes per paged-cache block - bool const kv_mxfp8) { // V4.1 all-dims MXFP8 record + bool const kv_mxfp8, // V4.1 all-dims MXFP8 record + bool const apply_q_rope) { // rotate Q too, not just KV #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) // BF16 _typeConvert specialization is unavailable on pre-Ampere. The // DeepseekV4 kernel only runs with bf16 inputs in practice, so compile a @@ -497,20 +532,18 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel( if (isKV) { src_ptr = kv_in + static_cast(tokenIdx) * kHeadDim + dim_base; } else { - int64_t const q_row_offset = - (static_cast(tokenIdx) * num_heads_q + slotIdx) * - kHeadDim + - dim_base; - src_ptr = q_in + q_row_offset; + src_ptr = q_in + qElemOffset(tokenIdx, slotIdx, laneId, + dim_base, num_heads_q); } v0 = *reinterpret_cast(src_ptr); v1 = *reinterpret_cast(src_ptr + 8); } - processDeepseekV4Slot( + processDeepseekV4Slot( v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids, cos_sin_cache, cache_block_size, - kv_block_stride, kv_mxfp8); + kv_block_stride, kv_mxfp8, apply_q_rope); #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) cudaTriggerProgrammaticLaunchCompletion(); @@ -530,7 +563,8 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel( // Q branch (optional RMSNorm + RoPE, in place) head_slot == num_heads_q // KV branch (RoPE + UE8M0 quant + insert) // -template +template __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid( scalar_t_in const* __restrict__ q_in, scalar_t_in* __restrict__ q_out, scalar_t_in const* __restrict__ kv_in, uint8_t* __restrict__ k_cache, @@ -539,7 +573,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid( float const* __restrict__ cos_sin_cache, float const eps, int const num_tokens_full, int const num_tokens_insert, int const num_heads_q, int const cache_block_size, - int const kv_block_stride, bool const kv_mxfp8) { + int const kv_block_stride, bool const kv_mxfp8, bool const apply_q_rope) { #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) if constexpr (std::is_same_v) { return; @@ -569,11 +603,8 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid( if (s == kNumHeadsQPadded) { src = kv_in + static_cast(tokenIdx) * kHeadDim + dim_base; } else { - src = q_in + - (static_cast(tokenIdx) * num_heads_q + - static_cast(s)) * - kHeadDim + - dim_base; + src = q_in + qElemOffset(tokenIdx, s, laneId, dim_base, + num_heads_q); } va = *reinterpret_cast(src); vb = *reinterpret_cast(src + 8); @@ -594,10 +625,12 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid( load_slot(next_slot, v0_next, v1_next); } - processDeepseekV4Slot( + processDeepseekV4Slot( v0_curr, v1_curr, tokenIdx, curr_slot, dim_base, laneId, num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids, - cos_sin_cache, cache_block_size, kv_block_stride, kv_mxfp8); + cos_sin_cache, cache_block_size, kv_block_stride, kv_mxfp8, + apply_q_rope); // ── Buffer rotation: hand the prefetched LDGs to the next iter. v0_curr = v0_next; @@ -617,14 +650,15 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid( // ──────────────────────────────────────────────────────────────────────────── // Launch wrapper // ──────────────────────────────────────────────────────────────────────────── -template +template static void launchFusedDeepseekV4Templated( scalar_t_in const* q_in, scalar_t_in* q_out, scalar_t_in const* kv_in, uint8_t* k_cache, int64_t const* slot_mapping, int64_t const* position_ids, float const* cos_sin_cache, float const eps, int const num_tokens_full, int const num_tokens_insert, int const num_heads_q, int const cache_block_size, int const kv_block_stride, bool const kv_mxfp8, - cudaStream_t stream) { + bool const apply_q_rope, cudaStream_t stream) { constexpr int kBlockSize = 256; constexpr int kWarpsPerBlock = kBlockSize / 32; int64_t const total_warps = @@ -663,37 +697,40 @@ static void launchFusedDeepseekV4Templated( // grid instead. Only reachable above NUM_TOKEN_CUTOFF tokens in a single // insert, where it is worth ~2x. if (kNumHeadsQPadded == 0 || num_tokens_full < NUM_TOKEN_CUTOFF) { - cudaLaunchKernelEx(&config, - fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel< - scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>, - q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, - cos_sin_cache, eps, num_tokens_full, num_tokens_insert, - num_heads_q, cache_block_size, kv_block_stride, - kv_mxfp8); + cudaLaunchKernelEx( + &config, + fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel< + scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM, Q_INTERLEAVED>, + q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, + eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size, + kv_block_stride, kv_mxfp8, apply_q_rope); } else { config.gridDim = dim3(num_tokens_full); cudaLaunchKernelEx( &config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid< - scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>, + scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM, Q_INTERLEAVED>, q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size, - kv_block_stride, kv_mxfp8); + kv_block_stride, kv_mxfp8, apply_q_rope); } #else // ROCm: use standard kernel launch syntax (no PDL/stream serialization) // clang-format off fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel< - scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM> + scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM, Q_INTERLEAVED> <<>>( q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q, - cache_block_size, kv_block_stride, kv_mxfp8); + cache_block_size, kv_block_stride, kv_mxfp8, apply_q_rope); #endif } // Runtime dispatch into one of the precompiled `kNumHeadsQPadded` -// instantiations. Supported padded head counts: 8, 16, 32, 64, 128. +// instantiations. Supported padded head counts: 8, 16, 32, 64, 128, plus 0 +// for a KV-only launch that writes no Q at all. `apply_q_norm` and +// `is_q_interleaved` are template parameters, `apply_q_rope` is not (see +// processDeepseekV4Slot), so each head count carries four instantiations. template void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert( scalar_t_in const* q_in, scalar_t_in* q_out, scalar_t_in const* kv_in, @@ -702,22 +739,32 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert( int const num_tokens_full, int const num_tokens_insert, int const num_heads_q, int const num_heads_q_padded, int const cache_block_size, int const kv_block_stride, - bool const apply_q_norm, bool const kv_mxfp8, cudaStream_t stream) { -#define DISPATCH(N) \ - case N: \ - if (apply_q_norm) { \ - launchFusedDeepseekV4Templated( \ - q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \ - cos_sin_cache, eps, num_tokens_full, num_tokens_insert, \ - num_heads_q, cache_block_size, kv_block_stride, kv_mxfp8, stream); \ - } else { \ - launchFusedDeepseekV4Templated( \ - q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \ - cos_sin_cache, eps, num_tokens_full, num_tokens_insert, \ - num_heads_q, cache_block_size, kv_block_stride, kv_mxfp8, stream); \ - } \ + bool const apply_q_norm, bool const kv_mxfp8, bool const apply_q_rope, + bool const is_q_interleaved, cudaStream_t stream) { +#define LAUNCH(N, NORM, INTERLEAVED) \ + launchFusedDeepseekV4Templated( \ + q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, \ + eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size, \ + kv_block_stride, kv_mxfp8, apply_q_rope, stream) +#define DISPATCH(N) \ + case N: \ + if (is_q_interleaved) { \ + if (apply_q_norm) { \ + LAUNCH(N, true, true); \ + } else { \ + LAUNCH(N, false, true); \ + } \ + } else if (apply_q_norm) { \ + LAUNCH(N, true, false); \ + } else { \ + LAUNCH(N, false, false); \ + } \ return; + if (num_heads_q_padded == 0) { + LAUNCH(0, false, false); + return; + } switch (num_heads_q_padded) { DISPATCH(8) DISPATCH(16) @@ -729,9 +776,10 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert( "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert: " "unsupported num_heads_q_padded=", num_heads_q_padded, - " (compiled instantiations: 8, 16, 32, 64, 128)."); + " (compiled instantiations: 0, 8, 16, 32, 64, 128)."); } #undef DISPATCH +#undef LAUNCH } // ──────────────────────────────────────────────────────────────────────────── @@ -1059,10 +1107,15 @@ void fused_deepseek_v4_kv_rope_insert( // Zero query heads schedules only the existing KV branch, preserving its // RoPE rounding and quantization without allocating any query tensors. if (packed) { - vllm::deepseek_v4_fused_ops::launchFusedDeepseekV4Templated( + // With zero query heads no Q slot is scheduled, so neither Q knob is + // reachable: Q_INTERLEAVED picks addressing nothing uses, and apply_q_rope + // gates a branch only a Q slot takes (KV is rotated unconditionally). + vllm::deepseek_v4_fused_ops::launchFusedDeepseekV4Templated( nullptr, nullptr, input, cache, slots, positions, cos_sin, 0.0f, num_tokens, num_tokens, 0, static_cast(cache_block_size), - static_cast(k_cache.stride(0)), kv_mxfp8, stream); + static_cast(k_cache.stride(0)), kv_mxfp8, /*apply_q_rope=*/true, + stream); } else if (fp8) { vllm::deepseek_v4_fused_ops::launchFullCacheKernel( nullptr, nullptr, 0, 0, input, cache, slots, positions, cos_sin, @@ -1085,8 +1138,9 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::stable::Tensor const& slot_mapping, // [N] int64 torch::stable::Tensor const& position_ids, // [N] int64 torch::stable::Tensor const& cos_sin_cache, // [max_pos, rope_dim] bf16 - int64_t q_head_padded, // padded Q head count for output - double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8) { + int64_t q_head_padded, // padded Q head count for output, 0 for no Q + double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8, + bool apply_q_rope, bool is_q_interleaved) { STD_TORCH_CHECK(q_in.device().is_cuda() && q_in.is_contiguous(), "q_in must be contiguous CUDA"); STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(), @@ -1106,8 +1160,11 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]"); STD_TORCH_CHECK(q_in.scalar_type() == kv.scalar_type(), "q_in and kv dtype must match"); - STD_TORCH_CHECK(q_head_padded >= q_in.size(1), - "q_head_padded must be >= q_in.size(1) (num_heads_q)"); + // q_head_padded == 0 asks for the KV insert alone: Q is left untouched and + // an empty tensor comes back, for a caller whose Q is already in the shape + // its attention kernel reads. + STD_TORCH_CHECK(q_head_padded == 0 || q_head_padded >= q_in.size(1), + "q_head_padded must be 0 or >= q_in.size(1) (num_heads_q)"); STD_TORCH_CHECK(k_cache.scalar_type() == torch::headeronly::ScalarType::Byte, "k_cache must be uint8"); STD_TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64, @@ -1138,8 +1195,12 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( // Allocate the padded q output. The kernel writes every element (live // region gets optional RMSNorm+RoPE; pad region gets zeros), so `empty` is // safe. - auto q_out = torch::stable::new_empty( - q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.scalar_type()); + auto q_out = + q_head_padded == 0 + ? torch::stable::new_empty(q_in, {int64_t{0}}, q_in.scalar_type()) + : torch::stable::new_empty( + q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, + q_in.scalar_type()); VLLM_STABLE_DISPATCH_HALF_TYPES( q_in.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] { @@ -1155,7 +1216,8 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( cos_sin_cache.const_data_ptr(), static_cast(eps), num_tokens_full, num_tokens_insert, num_heads_q, num_heads_q_padded, cache_block_size_i, kv_block_stride, - apply_q_norm, kv_mxfp8, stream); + apply_q_norm, kv_mxfp8, apply_q_rope, is_q_interleaved, + stream); }); return q_out; } diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index bc08d90a4836..6b17815a07c7 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -271,7 +271,8 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, torch::stable::Tensor const& position_ids, torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded, - double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8); + double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8, + bool apply_q_rope, bool is_q_interleaved); void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( torch::stable::Tensor& q, torch::stable::Tensor const& kv, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 3918ab8498cd..a76852df89e5 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -425,6 +425,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "bool is_neox, Tensor position_ids, " "int forced_token_heads_per_warp=-1) -> ()"); + // q_head_padded is the padded Q head count of the returned tensor, or 0 to + // do the KV insert alone and return an empty tensor. The Q knobs are + // independent: apply_q_norm and apply_q_rope each drop that step for Q + // alone (KV is always rotated), and is_q_interleaved reads and writes Q in + // FlashMLA's mega-attention chunk-interleaved layout, which moves the + // padding heads to the tail of every head-dim chunk. ops.def( "fused_deepseek_v4_kv_rope_insert(" "Tensor kv, Tensor! k_cache, Tensor slot_mapping, Tensor position_ids, " @@ -436,7 +442,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor q_in, Tensor kv, Tensor! k_cache, " "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " "int q_head_padded, float eps, int cache_block_size, " - "bool apply_q_norm=True, bool kv_mxfp8=False) -> Tensor"); + "bool apply_q_norm=True, bool kv_mxfp8=False, bool apply_q_rope=True, " + "bool is_q_interleaved=False) -> Tensor"); // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate // FP8 tensor, and KV into a contiguous 512-wide token-strided cache. diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index c1de94be92fc..73c2027a7866 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -331,6 +331,19 @@ def run(): torch.testing.assert_close(cache_backing, expected_cache, rtol=0, atol=0) +def _rotate_rope_tail( + latent_row: torch.Tensor, cos_sin_row: torch.Tensor +) -> torch.Tensor: + """GPT-J RoPE of one latent row's last 64 dims, in fp32.""" + row = latent_row.float() + c, s = cos_sin_row.chunk(2) + rotated = row.clone() + even, odd = row[448::2], row[449::2] + rotated[448::2] = even * c - odd * s + rotated[449::2] = odd * c + even * s + return rotated + + def _mxfp8_record_reference( latent_row: torch.Tensor, cos_sin_row: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: @@ -339,12 +352,7 @@ def _mxfp8_record_reference( Mirrors FlashMLA's ``KVCacheLayout.V41_FP8Sparse`` quantizer: rotate the RoPE tail first, then scale every 32-dim tile, RoPE tiles included. """ - row = latent_row.float() - c, s = cos_sin_row.chunk(2) - rotated = row.clone() - even, odd = row[448::2], row[449::2] - rotated[448::2] = even * c - odd * s - rotated[449::2] = odd * c + even * s + rotated = _rotate_rope_tail(latent_row, cos_sin_row) quantized, scales = _ue8m0_reference(rotated, 32, 448.0) return quantized.view(torch.uint8), (scales.log2() + 127).to(torch.uint8) @@ -461,6 +469,125 @@ def test_v41_mxfp8_cache_round_trip(): assert (error <= tolerance.repeat_interleave(32, dim=-1)).all() +# e2m1 magnitudes indexed by the low 3 bits of a code; bit 3 is the sign. +_E2M1_MAGNITUDES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] + + +def _decode_nvfp4_row(packed: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + """Unpack one 256-byte e2m1 row and apply its 32 e4m3 tile scales.""" + mags = torch.tensor(_E2M1_MAGNITUDES, device=packed.device) + codes = torch.empty(512, dtype=torch.uint8, device=packed.device) + codes[0::2] = packed & 0xF # even element in the low nibble + codes[1::2] = packed >> 4 + vals = mags[(codes & 7).long()] * torch.where(codes >= 8, -1.0, 1.0) + return ( + vals.view(32, 16) * scales.view(torch.float8_e4m3fn).float().view(32, 1) + ).flatten() + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +@pytest.mark.parametrize("compress_ratio", [1, 2]) +def test_v41_rope_insert_nvfp4_record(compress_ratio: int): + """The 288-byte V4.1 NVFP4 record: 256 B of e2m1 pairs then 32 e4m3 scales. + + Scales are byte-exact against the reference (``amax / 6`` clamped to the + e4m3 range); values are checked as round-to-nearest on the e2m1 grid, whose + coarsest step is 2 between 4 and 6, so half a step is ``1.0 * scale``. + """ + from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( + rope_quant_insert, + ) + + torch.manual_seed(13) + device = "cuda" + cache_block = 64 + cache_stride = math.ceil(cache_block * 288 / 512) * 512 + num_tokens = 12 + + positions = torch.arange(num_tokens, dtype=torch.int64, device=device) + latent = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device) + angles = torch.randn(64, 32, device=device) + cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1) + + cache_backing = torch.full((2, cache_stride), 165, dtype=torch.uint8, device=device) + cache = cache_backing.as_strided((2, cache_block, 288), (cache_stride, 288, 1)) + slots = torch.arange(num_tokens, dtype=torch.int64, device=device) + slots[3] = -1 # covers the negative-slot skip + + rope_quant_insert(latent, positions, cos_sin, cache, slots, compress_ratio) + + untouched = torch.full_like(cache_backing, 165) + for t in range(num_tokens): + slot = slots[t].item() + pos = positions[t].item() + if slot < 0 or (pos + 1) % compress_ratio: + continue + page, row = divmod(slot, cache_block) + cs = cos_sin[pos // compress_ratio * compress_ratio] + rotated = _rotate_rope_tail(latent[t], cs) + amax = rotated.view(32, 16).abs().amax(-1) + scale = (amax / 6.0).clamp(2.0**-9, 448.0).to(torch.float8_e4m3fn) + + got_scales = cache_backing[page, cache_block * 256 + row * 32 :][:32] + torch.testing.assert_close(got_scales, scale.view(torch.uint8), rtol=0, atol=0) + decoded = _decode_nvfp4_row( + cache_backing[page, row * 256 : (row + 1) * 256], got_scales + ) + step = scale.float().repeat_interleave(16) + assert ((decoded - rotated).abs() <= step).all() + untouched[page, row * 256 : (row + 1) * 256] = cache_backing[ + page, row * 256 : (row + 1) * 256 + ] + untouched[page, cache_block * 256 + row * 32 :][:32] = got_scales + # Rows the kernel must not have touched, page padding included. + torch.testing.assert_close(cache_backing, untouched, rtol=0, atol=0) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +def test_v41_nvfp4_gather_matches_insert(): + """The NVFP4 gather dequantizes exactly what the insert kernel wrote.""" + from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache + from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( + rope_quant_insert, + ) + + torch.manual_seed(17) + device = "cuda" + block_size = 64 + num_tokens = 70 + num_blocks = 4 + page_bytes = math.ceil(block_size * 288 / 512) * 512 + + positions = torch.arange(num_tokens, dtype=torch.int64, device=device) + latent = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device) + angles = torch.randn(128, 32, device=device) + cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1) + backing = torch.zeros(num_blocks, page_bytes, dtype=torch.uint8, device=device) + cache = backing.as_strided((num_blocks, block_size, 288), (page_bytes, 288, 1)) + slots = torch.arange(num_tokens, dtype=torch.int64, device=device) + rope_quant_insert(latent, positions, cos_sin, cache, slots, 1) + + out = torch.zeros(1, num_tokens, 512, dtype=torch.bfloat16, device=device) + dequantize_and_gather_k_cache( + out, + cache, + seq_lens=torch.tensor([num_tokens], dtype=torch.int32, device=device), + gather_lens=None, + block_table=torch.arange(num_blocks, dtype=torch.int32, device=device).view( + 1, -1 + ), + block_size=block_size, + offset=0, + ) + for t in (0, 1, block_size, num_tokens - 1): + page, row = divmod(t, block_size) + expected = _decode_nvfp4_row( + backing[page, row * 256 : (row + 1) * 256], + backing[page, block_size * 256 + row * 32 :][:32], + ) + torch.testing.assert_close(out[0, t].float(), expected, rtol=0, atol=0) + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") @pytest.mark.parametrize("compress_ratio", [1, 2]) @pytest.mark.parametrize("store_fp8", [False, True]) diff --git a/tests/kernels/test_dsv41_mega_attn_layouts.py b/tests/kernels/test_dsv41_mega_attn_layouts.py new file mode 100644 index 000000000000..6f9ff56bd236 --- /dev/null +++ b/tests/kernels/test_dsv41_mega_attn_layouts.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Layout plumbing for DeepSeek V4.1 mega attention. + +The kernel's Q and O layouts are produced by permuting ``wq_b`` rows and +``wo_a`` columns once at load, so what has to hold is that the permuted GEMMs +agree with a reference that permutes the activation instead. +""" + +import torch + +from vllm.models.deepseek_v41.common.ops.fused_layout import ( + WV_GROUP_SIZE, + o_fused_chunk_permutation, + o_fused_permutation, + permute_q_to_fused, + permute_wo_a_, + permute_wq_b_, + q_fused_permutation, +) + +HEAD_DIM = 512 + + +def test_permuted_wq_b_gemm_matches_permuted_activation(): + """Permuting wq_b's rows makes the Q GEMM emit the fused layout directly.""" + torch.manual_seed(0) + num_heads, q_lora_rank, num_tokens = 16, 64, 5 + weight = torch.randn(num_heads * HEAD_DIM, q_lora_rank) + # One scale row per weight row, as an MXFP8 wq_b shard carries. + scale = torch.randint( + 0, 256, (num_heads * HEAD_DIM, q_lora_rank // 32), dtype=torch.uint8 + ) + qr = torch.randn(num_tokens, q_lora_rank) + orig_weight, orig_scale = weight.clone(), scale.clone() + + standard = (qr @ weight.T).view(num_tokens, num_heads, HEAD_DIM) + expected = permute_q_to_fused(standard) + + permute_wq_b_(weight, scale, num_heads) + fused = (qr @ weight.T).view(num_tokens, num_heads, HEAD_DIM) + torch.testing.assert_close(fused, expected) + # The scale has to follow its row, or dequant reads another row's scale. + perm = q_fused_permutation(num_heads, HEAD_DIM) + torch.testing.assert_close(weight, orig_weight[perm]) + torch.testing.assert_close(scale, orig_scale[perm]) + + +def test_permuted_wo_a_consumes_fused_output(): + """Permuting wo_a's input columns lets it read the kernel's O layout.""" + torch.manual_seed(1) + out_features, num_tokens = 32, 4 + in_features = WV_GROUP_SIZE * HEAD_DIM + weight = torch.randn(out_features, in_features) + scale = torch.randint(0, 256, (out_features, in_features // 32), dtype=torch.uint8) + standard_o = torch.randn(num_tokens, in_features) + orig_weight, orig_scale = weight.clone(), scale.clone() + + expected = standard_o @ weight.T + # The kernel emits the same values in the fused chunk order. + perm = o_fused_permutation(WV_GROUP_SIZE, HEAD_DIM) + fused_o = standard_o[:, perm] + + permute_wo_a_(weight, scale, WV_GROUP_SIZE) + torch.testing.assert_close(weight, orig_weight[:, perm]) + torch.testing.assert_close( + scale, orig_scale[:, o_fused_chunk_permutation(WV_GROUP_SIZE, HEAD_DIM)] + ) + # Unlike wq_b, this permutation sits inside the 4096-term reduction, so the + # summation order changes and the result differs in the last few ulps. + torch.testing.assert_close(fused_o @ weight.T, expected, rtol=1e-3, atol=1e-3) diff --git a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py index 68eabfb94e84..a0dcfc5bb7da 100644 --- a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py +++ b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py @@ -522,6 +522,189 @@ def _dequant(cache_2d): ) +# ── Test 2d: mega-attention Q layout (is_q_interleaved) ────────────────────── + +Q_CHUNK = 16 +NUM_Q_CHUNKS = HEAD_DIM // Q_CHUNK # 32 + + +@pytest.mark.parametrize("num_tokens", [1, 17, 2048]) +@pytest.mark.parametrize("n_heads,q_head_padded", [(16, 64), (64, 128), (64, 0)]) +def test_q_interleaved_pads_without_rope( + num_tokens: int, n_heads: int, q_head_padded: int +): + """``is_q_interleaved`` with no norm or RoPE only zero-pads Q. + + The mega-attention kernel norms and rotates Q itself, so every live + element must survive bit-exact while the padding heads -- which interleave + with the live ones inside each of the 32 head-dim chunks -- come back + zero. ``q_head_padded=0`` drops the Q pass entirely and returns nothing, + for a shard already at the kernel's head count. + """ + torch.manual_seed(5) + device = "cuda" + dtype = torch.bfloat16 + block_size = 64 + + q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype, device=device) + kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device) + positions = torch.arange(num_tokens, dtype=torch.int64, device=device) + cos_sin_cache = make_cos_sin_cache(4096, ROPE_DIM, torch.float32, device) + num_blocks = (num_tokens + block_size - 1) // block_size + 1 + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + k_cache = torch.zeros( + num_blocks, block_size * V41_HEAD_BYTES, dtype=torch.uint8, device=device + ) + + q_out = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q, + kv, + k_cache, + slot_mapping, + positions, + cos_sin_cache, + q_head_padded, + 1e-6, + block_size, + False, # apply_q_norm + True, # kv_mxfp8 + False, # apply_q_rope + True, # is_q_interleaved + ) + + # The KV half is unaffected by the Q mode: it must still be inserted. + assert k_cache.any() + + if q_head_padded == 0: + assert q_out.numel() == 0 + return + assert q_out.shape == (num_tokens, q_head_padded, HEAD_DIM) + got = q_out.view(num_tokens, NUM_Q_CHUNKS, q_head_padded, Q_CHUNK) + want = q.view(num_tokens, NUM_Q_CHUNKS, n_heads, Q_CHUNK) + torch.testing.assert_close(got[:, :, :n_heads], want, rtol=0, atol=0) + assert (got[:, :, n_heads:] == 0).all() + + +@pytest.mark.parametrize("n_heads", [8, 16, 32]) +def test_q_interleaved_survives_graph_replay(n_heads: int): + """A captured fused insert refills its padded Q buffer on every replay. + + The op allocates ``q_out`` itself, so the tensor a capture hands back has + to keep receiving the new padding when the graph is replayed over changed + input -- that is what lets the layer prepare Q inside the captured region + rather than in ``forward_mqa``. + """ + torch.manual_seed(6) + device = "cuda" + dtype = torch.bfloat16 + block_size, num_tokens, q_head_padded = 64, 7, 64 + + q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype, device=device) + kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device) + positions = torch.arange(num_tokens, dtype=torch.int64, device=device) + cos_sin_cache = make_cos_sin_cache(4096, ROPE_DIM, torch.float32, device) + num_blocks = (num_tokens + block_size - 1) // block_size + 1 + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + k_cache = torch.zeros( + num_blocks, block_size * V41_HEAD_BYTES, dtype=torch.uint8, device=device + ) + + def prepare(): + return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q, + kv, + k_cache, + slot_mapping, + positions, + cos_sin_cache, + q_head_padded, + 1e-6, + block_size, + False, # apply_q_norm + True, # kv_mxfp8 + False, # apply_q_rope + True, # is_q_interleaved + ) + + # Capture takes its warmup on a side stream. + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(3): + prepare() + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + q_out = prepare() + + for factor in (0.25, -0.75): + q.copy_(torch.randn_like(q) * factor) + graph.replay() + torch.accelerator.synchronize() + got = q_out.view(num_tokens, NUM_Q_CHUNKS, q_head_padded, Q_CHUNK) + want = q.view(num_tokens, NUM_Q_CHUNKS, n_heads, Q_CHUNK) + torch.testing.assert_close(got[:, :, :n_heads], want, rtol=0, atol=0) + assert (got[:, :, n_heads:] == 0).all() + + +@pytest.mark.parametrize("n_heads", [16, 64]) +@pytest.mark.parametrize("apply_q_norm", [False, True]) +@pytest.mark.parametrize("apply_q_rope", [False, True]) +def test_q_interleaved_is_orthogonal_to_norm_and_rope( + n_heads: int, apply_q_norm: bool, apply_q_rope: bool +): + """The Q layout composes with whatever norm and RoPE it is given. + + A lane owns the same 512-dim slice of its head in either layout -- only + the address differs -- so the head-major result permuted into the + interleaved layout must equal running the interleaved path on a permuted + input, bit for bit, for every norm/RoPE combination. + """ + from vllm.models.deepseek_v41.common.ops.fused_layout import permute_q_to_fused + + torch.manual_seed(7) + device = "cuda" + dtype = torch.bfloat16 + block_size, num_tokens, q_head_padded = 64, 5, 64 + + q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype, device=device) + kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device) + positions = torch.arange(num_tokens, dtype=torch.int64, device=device) + cos_sin_cache = make_cos_sin_cache(4096, ROPE_DIM, torch.float32, device) + num_blocks = (num_tokens + block_size - 1) // block_size + 1 + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + + def run(q_in, is_q_interleaved): + cache = torch.zeros( + num_blocks, block_size * V41_HEAD_BYTES, dtype=torch.uint8, device=device + ) + out = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q_in, + kv, + cache, + slot_mapping, + positions, + cos_sin_cache, + q_head_padded, + 1e-6, + block_size, + apply_q_norm, + True, # kv_mxfp8 + apply_q_rope, + is_q_interleaved, + ) + return out, cache + + head_major, cache_major = run(q, False) + interleaved, cache_inter = run(permute_q_to_fused(q), True) + + torch.testing.assert_close( + interleaved, permute_q_to_fused(head_major), rtol=0, atol=0 + ) + # The Q layout must not leak into the KV half. + torch.testing.assert_close(cache_inter, cache_major, rtol=0, atol=0) + + # ── Test 2b: DP padding (slot_mapping shorter than q/kv) ───────────────────── diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index 802a59e340fc..b3d80f8d2591 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -636,7 +636,8 @@ def _split_qkv_and_norm( transpose_scale=False, ) - def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + o = attn_out[:, : self.n_local_heads, :] # ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b. z = rocm_inv_rope_einsum( self.rotary_emb, diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 134328dbe908..752b335eafd8 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -44,6 +44,7 @@ VllmConfig, get_current_vllm_config, ) +from vllm.config.cache import CacheDType from vllm.distributed import get_tensor_model_parallel_world_size from vllm.forward_context import get_forward_context from vllm.logger import init_logger @@ -127,34 +128,40 @@ def _use_v41_mxfp8_kv_record() -> bool: def _resolve_dsv4_kv_cache_dtype( use_fp8_ds_mla_layout: bool, - kv_cache_dtype: str, + kv_cache_dtype: CacheDType, cache_config: CacheConfig | None, -) -> tuple[str, torch.dtype]: + packed_kv_cache_dtype: CacheDType = "fp8_ds_mla", +) -> tuple[CacheDType, torch.dtype]: """Map ``(layout, --kv-cache-dtype)`` to ``(cache_dtype_str, torch_dtype)``. Both layouts are paged; they differ in the per-token block format. The - ``fp8_ds_mla`` format is UE8M0 block-scaled fp8 packed as ``uint8`` (the - canonical ``fp8_ds_mla`` string is written back onto ``cache_config`` so the - page-size specs pick the 576B per-token slot). Plain-row backends store each - token's KV row in its element dtype: bf16 or per-tensor FP8 E4M3. + packed formats are ``uint8``-backed: ``fp8_ds_mla`` is UE8M0 block-scaled + fp8 throughout, ``nvfp4_ds_mla`` keeps that sliding-window record and + stores the compressed cache as NVFP4. An unspecific ``--kv-cache-dtype`` + (``auto`` / ``fp8``) resolves to ``packed_kv_cache_dtype``, the record this + layer's kernel prefers, and the canonical string is written back onto + ``cache_config`` so the page-size specs pick the right per-token slot. + Plain-row backends store each token's KV row in its element dtype: bf16 or + per-tensor FP8 E4M3. """ if use_fp8_ds_mla_layout: - # fp8_ds_mla block format: UE8M0 block-scaled fp8 packed as uint8. - if kv_cache_dtype == "auto": - kv_cache_dtype = "fp8" - if not kv_cache_dtype.startswith("fp8"): + if kv_cache_dtype in ("auto", "fp8"): + kv_cache_dtype = packed_kv_cache_dtype + elif not kv_cache_dtype.endswith("_ds_mla"): raise ValueError( - "DeepseekV4 fp8_ds_mla layout only supports fp8 " - f"kv-cache, got {kv_cache_dtype}. Please set " - "`--kv-cache-dtype fp8` or select a backend that supports " - "bfloat16 KV cache." + "DeepseekV4 packed KV layouts only support fp8 kv-cache, got " + f"{kv_cache_dtype}. Please set `--kv-cache-dtype fp8` or " + "select a backend that supports bfloat16 KV cache." ) - if kv_cache_dtype != "fp8_ds_mla": - if cache_config is not None: - cache_config.cache_dtype = "fp8_ds_mla" - kv_cache_dtype = "fp8_ds_mla" - logger.info_once("Using DeepSeek's fp8_ds_mla KV cache format.") + if cache_config is not None and cache_config.cache_dtype != kv_cache_dtype: + cache_config.cache_dtype = kv_cache_dtype + logger.info_once("Using DeepSeek's %s KV cache format.", kv_cache_dtype) return kv_cache_dtype, torch.uint8 + if kv_cache_dtype.endswith("_ds_mla"): + raise ValueError( + f"{kv_cache_dtype} is a packed FlashMLA DeepSeek V4.1 KV cache " + "format; the selected backend stores plain KV rows." + ) # Plain bf16 / per-tensor fp8 KV row (FlashInfer). if kv_cache_dtype.startswith("fp8"): @@ -189,6 +196,28 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): # path to pre-reserve that workspace. PREFILL_CHUNK_SIZE: ClassVar[int] = 4 + # ---- attention-interface contract, declared by the platform subclass ---- + + @property + def accepts_unnormed_unroped_query(self) -> bool: + """Whether ``forward_mqa``'s ``q`` is the raw ``wq_b`` output. + + True when the attention kernel applies the Q norm and RoPE itself, and + reads Q in its own chunk-interleaved layout, so the layer only + zero-pads Q to ``padded_heads`` and inserts KV before calling it. + """ + return False + + @property + def packed_kv_cache_dtype(self) -> CacheDType: + """The packed KV record this layer's kernel prefers. + + What an unspecific ``--kv-cache-dtype`` (``auto`` / ``fp8``) resolves + to. Mega attention overrides it: its kernel is the one that can read + an NVFP4 compressed cache. + """ + return "fp8_ds_mla" + @classmethod @abstractmethod def get_padded_num_q_heads(cls, num_heads: int) -> int: @@ -213,8 +242,19 @@ def forward_mqa( raise NotImplementedError @abstractmethod - def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: - """Inverse-RoPE + wo_a + wo_b output projection (platform-specific).""" + def _o_proj( + self, + attn_out: "torch.Tensor | QuantizedActivation", + positions: torch.Tensor, + ) -> torch.Tensor: + """Project whatever ``_alloc_attn_out`` produced through wo_a and wo_b. + + Takes the buffer whole, so each layer owns the shape it allocated: the + bf16 layers slice off their padding heads and apply the inverse RoPE, + while a layer whose attention kernel already did the inverse RoPE and + the FP8 cast gets a QuantizedActivation and has only wo_a and wo_b + left. + """ raise NotImplementedError def _uses_fp8_ds_mla_layout(self) -> bool: @@ -459,11 +499,27 @@ def __init__( # Resolve the kv-cache dtype from this backend's block format. The same # resolution drives the SWA cache tensor dtype below. self.kv_cache_dtype, self.kv_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype( - self._uses_fp8_ds_mla_layout(), cache_config.cache_dtype, cache_config + self._uses_fp8_ds_mla_layout(), + cache_config.cache_dtype, + cache_config, + self.packed_kv_cache_dtype, ) self.kv_mxfp8 = _use_v41_mxfp8_kv_record() - self.kv_bytes_per_token = 528 if self.kv_mxfp8 else 584 + self.swa_bytes_per_token = 528 if self.kv_mxfp8 else 584 + # nvfp4_ds_mla keeps the MXFP8 sliding-window record and stores the + # compressed cache as NVFP4 (256 B of e2m1 pairs + 32 e4m3 scales). + self.compressed_bytes_per_token = ( + 288 if self.kv_cache_dtype == "nvfp4_ds_mla" else self.swa_bytes_per_token + ) + # One alignment for every page in the block: the block stride is their + # sum, and 512 satisfies both TMA strides in play (512 for the V4.1 + # fp8 record, 256 for NVFP4). self.kv_page_alignment = 512 if self.kv_mxfp8 else 576 + if self.kv_cache_dtype == "nvfp4_ds_mla" and not self.kv_mxfp8: + raise ValueError( + "nvfp4_ds_mla needs the V4.1 KV records, which FlashMLA " + "decodes only on SM100." + ) self.swa_cache_layer = DeepseekV4SWACache( head_dim=self.head_dim, @@ -473,7 +529,7 @@ def __init__( cache_config=cache_config, backend_cls=self.swa_backend_cls, block_size=32, - packed_bytes_per_token=self.kv_bytes_per_token, + packed_bytes_per_token=self.swa_bytes_per_token, packed_page_alignment=self.kv_page_alignment, ) @@ -553,8 +609,12 @@ def __init__( block_size=self.swa_cache_layer.block_size, ) + # Every backend that gathers a chunk's KV through + # combine_topk_swa_indices needs its Triton kernel warmed, mega + # attention included -- it calls it from _forward_prefill_mega. if self.backend_cls.get_name() in ( "FLASHMLA_SPARSE_DSV41", + "FLASHMLA_MEGA_ATTN_DSV41", "ROCM_FLASHMLA_SPARSE_DSV4", ): from vllm.models.deepseek_v41.common.ops.cache_utils import ( @@ -569,14 +629,10 @@ def forward( hidden_states: torch.Tensor, llama_4_scaling: torch.Tensor | None = None, ) -> torch.Tensor: - # Pre-allocate attention output with FlashMLA-padded head count. - # The op writes into `o_padded`; we slice to n_local_heads after. - num_tokens = hidden_states.shape[0] - o_padded = torch.empty( - (num_tokens, self.padded_heads, self.head_dim), - dtype=hidden_states.dtype, - device=hidden_states.device, - ) + # The eager attention region writes into a caller-owned buffer + # (breakable_cudagraph needs in-place outputs); its shape and how it is + # projected afterwards follow the interface contract above. + attn_out = self._alloc_attn_out(hidden_states.shape[0], hidden_states) # Keep the attention input preparation in the captured graph. Only the # sparse indexer and MLA attention run in the eager break below. @@ -593,12 +649,19 @@ def forward( kv_score, indexer_weights, positions, - o_padded, + attn_out, ) - o = o_padded[:, : self.n_local_heads, :] + return self._o_proj(attn_out, positions) - # Inverse-RoPE + wo_a + wo_b output projection (platform-specific). - return self._o_proj(o, positions) + def _alloc_attn_out( + self, num_tokens: int, hidden_states: torch.Tensor + ) -> "torch.Tensor | QuantizedActivation": + """The buffer ``forward_mqa`` fills, per the interface contract.""" + return torch.empty( + (num_tokens, self.padded_heads, self.head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) @cached_property def _can_fuse_query_quant(self) -> bool: @@ -652,7 +715,7 @@ def _prepare_and_attn_eager( kv_score: torch.Tensor, indexer_weights: torch.Tensor, positions: torch.Tensor, - o_padded: torch.Tensor, + attn_out: "torch.Tensor | QuantizedActivation", ) -> None: """Wide eager region: the whole of ``_prepare_and_attn`` runs eagerly. @@ -667,7 +730,7 @@ def _prepare_and_attn_eager( kv_score, indexer_weights, positions, - o_padded, + attn_out, ) def _prepare_and_attn( @@ -679,7 +742,7 @@ def _prepare_and_attn( kv_score: torch.Tensor, indexer_weights: torch.Tensor, positions: torch.Tensor, - o_padded: torch.Tensor, + attn_out: "torch.Tensor | QuantizedActivation", ) -> None: """Attention input preparation followed by the sparse indexer and MLA. @@ -751,7 +814,7 @@ def prepare_indexer(): q, kv, positions, - o_padded, + attn_out, ) def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -858,16 +921,28 @@ def _fused_qnorm_rope_kv_insert( dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None ), ) -> torch.Tensor: + """Ready ``q`` for the attention kernel and publish this step's KV. + + One launch does both. With ``accepts_unnormed_unroped_query`` the + attention kernel norms and rotates Q itself and reads it in its own + chunk-interleaved layout, so the Q half of the launch is a zero-pad to + ``padded_heads`` -- and nothing at all once the shard is that wide. + """ if not isinstance(attn_metadata, dict): # Profile run: kernel doesn't fire; produce a padded tensor so # downstream FlashMLA gets the right shape. - if self.n_local_heads < self.padded_heads: - return F.pad( - q, - (0, 0, 0, self.padded_heads - self.n_local_heads), - value=0.0, - ) - return q + if self.n_local_heads >= self.padded_heads: + return q + if self.accepts_unnormed_unroped_query: + # Padding heads sit at the tail of every head-dim chunk in + # that layout, so no head-major pad of `q` reproduces it -- + # and nothing reads it on a profile run. + return q.new_zeros((q.shape[0], self.padded_heads, q.shape[2])) + return F.pad( + q, + (0, 0, 0, self.padded_heads - self.n_local_heads), + value=0.0, + ) swa_metadata = cast( "DeepseekSparseSWAMetadata | None", @@ -886,23 +961,38 @@ def _fused_qnorm_rope_kv_insert( if cache_dtype == torch.uint8: # fp8_ds_mla UE8M0 paged path. Horizontally fused: # Q side: GPT-J RoPE, zero-filling the padding head slots; the - # kernel allocates and returns the padded q tensor. + # kernel allocates and returns the padded q tensor. An + # interleaved Q skips the RoPE its attention kernel owns + # and keeps only the pad, which q_head_padded=0 drops + # too once the shard is already padded_heads wide. # KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert. swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1) - return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + pad_to = ( + 0 + if self.accepts_unnormed_unroped_query + and self.n_local_heads == self.padded_heads + else self.padded_heads + ) + q_padded = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( q, kv, swa_kv_cache_2d, swa_metadata.slot_mapping, positions, cos_sin_cache, - self.padded_heads, + pad_to, self.eps, swa_metadata.block_size, - False, + False, # apply_q_norm: qr is normed before wq_b self.kv_mxfp8, + not self.accepts_unnormed_unroped_query, # apply_q_rope + self.accepts_unnormed_unroped_query, # is_q_interleaved ) + return q if pad_to == 0 else q_padded + assert not self.accepts_unnormed_unroped_query, ( + "the chunk-interleaved Q layout only pairs with a packed KV record" + ) # Plain-row path: the [num_blocks, block_size, 512] cache stores the KV # row in its element dtype (no Q padding). bf16 rewrites q in place; # per-tensor fp8 writes a separately-allocated fp8 q and quantizes the @@ -959,7 +1049,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: # fp8_ds_mla is a UE8M0 block-scaled uint8 layout whose page rounds up # to the decode kernel's TMA stride; plain bf16 / per-tensor fp8 rows # use natural element-size pages. - uses_fp8_ds_mla_layout = self.kv_cache_dtype == "fp8_ds_mla" + uses_fp8_ds_mla_layout = self.kv_cache_dtype in ("fp8_ds_mla", "nvfp4_ds_mla") return MLAAttentionSpec( block_size=vllm_config.cache_config.block_size, num_kv_heads=1, @@ -972,7 +1062,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), # Packed record width; head_size stays semantic (512). state_content_bytes=( - self.kv_bytes_per_token if uses_fp8_ds_mla_layout else None + self.compressed_bytes_per_token if uses_fp8_ds_mla_layout else None ), ) @@ -1016,7 +1106,13 @@ def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: # head_dim already carries the fp8 scale padding # tokens_per_state=1 for V3.2, >1 for DeepseekV4; same cache layout. - uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla" + # nvfp4_ds_mla is packed too: its compressed record is NVFP4 but the + # sliding-window record stays the V4.1 MXFP8 one, so the indexer page + # takes the same alignment either way. + uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype in ( + "fp8_ds_mla", + "nvfp4_ds_mla", + ) page_alignment = ( 576 if uses_fp8_ds_mla_layout and not _use_v41_mxfp8_kv_record() else 512 ) diff --git a/vllm/models/deepseek_v41/common/ops/cache_utils.py b/vllm/models/deepseek_v41/common/ops/cache_utils.py index 51743016bf07..4a9da21c036b 100644 --- a/vllm/models/deepseek_v41/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v41/common/ops/cache_utils.py @@ -43,10 +43,16 @@ # dims each plus a pad byte. # V4.1 (528 B): all 512 dims as fp8 e4m3 (RoPE included), then 16 UE8M0 # scales of 32 dims each. FlashMLA's ``ModelType::V41``. +# V4.1 NVFP4 (288 B): 512 e2m1 values packed two per byte (even element in +# the low nibble), then 32 e4m3 scales of 16 dims each. +# FlashMLA's ``ModelType::V41_FP4``, compressed cache only. V4_BYTES_PER_TOKEN = 584 V41_BYTES_PER_TOKEN = 528 V41_QUANT_BLOCK = 32 V41_NUM_SCALES = 512 // V41_QUANT_BLOCK # 16 +V41_NVFP4_BYTES_PER_TOKEN = 288 +V41_NVFP4_QUANT_BLOCK = 16 +V41_NVFP4_NUM_SCALES = 512 // V41_NVFP4_QUANT_BLOCK # 32 @triton.jit @@ -496,6 +502,73 @@ def _dequantize_and_gather_k_mxfp8_kernel( tl.store(output_row_ptr + d, dequant.to(tl.bfloat16)) +@triton.jit +def _dequantize_and_gather_k_nvfp4_kernel( + out_ptr, + out_stride0, + out_stride1, + k_cache_ptr, + seq_lens_ptr, + block_table_ptr, + offset, + gather_lens_ptr, + max_blocks_per_seq: tl.constexpr, + head_dim: tl.constexpr, # 512 + scale_dim: tl.constexpr, # 32 + quant_block: tl.constexpr, # 16 + cache_block_size: tl.constexpr, + block_stride: tl.constexpr, +): + """Gather and dequantize V4.1 NVFP4 rows into a bf16 workspace.""" + batch_idx = tl.program_id(0) + worker_id = tl.program_id(1) + num_workers = tl.num_programs(1) + + seq_len = tl.load(seq_lens_ptr + batch_idx) + if gather_lens_ptr is not None: # noqa: SIM108 + gather_len = tl.load(gather_lens_ptr + batch_idx) + else: + gather_len = seq_len + start_pos = seq_len - gather_len + + packed_bytes: tl.constexpr = head_dim // 2 + d = tl.arange(0, head_dim) + for i in range(worker_id, gather_len, num_workers): + pos = start_pos + i + block_in_seq = pos // cache_block_size + pos_in_block = pos % cache_block_size + + physical_block_idx = tl.load( + block_table_ptr + batch_idx * max_blocks_per_seq + block_in_seq + ) + page = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride + + packed = tl.load( + page + pos_in_block * packed_bytes + tl.arange(0, packed_bytes) + ) + # Even element in the low nibble, odd in the high nibble. + codes = tl.interleave((packed & 0xF).to(tl.int32), (packed >> 4).to(tl.int32)) + # e2m1 magnitudes: 0, 0.5, 1, 1.5, 2, 3, 4, 6. + mag_code = codes & 7 + e = (mag_code >> 1).to(tl.float32) + m = (mag_code & 1).to(tl.float32) + mag = tl.where(mag_code < 2, m * 0.5, (1.0 + m * 0.5) * tl.exp2(e - 1.0)) + vals = tl.where(codes >= 8, -mag, mag) + + sf = tl.load( + page + + cache_block_size * packed_bytes + + pos_in_block * scale_dim + + tl.arange(0, scale_dim) + ) + scale = sf.to(tl.float8e4nv, bitcast=True).to(tl.float32) + tiles = tl.reshape(vals, (scale_dim, quant_block)) + dequant = tl.reshape(tiles * tl.reshape(scale, (scale_dim, 1)), (head_dim,)) + + output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1 + tl.store(output_row_ptr + d, dequant.to(tl.bfloat16)) + + def dequantize_and_gather_k_cache_triton( # [num_reqs, max_num_tokens, head_size] out: torch.Tensor, @@ -513,6 +586,25 @@ def dequantize_and_gather_k_cache_triton( ) -> None: num_reqs = seq_lens.shape[0] NUM_WORKERS = 128 + if k_cache.shape[-1] == V41_NVFP4_BYTES_PER_TOKEN: + _dequantize_and_gather_k_nvfp4_kernel[(num_reqs, NUM_WORKERS)]( + out, + out.stride(0), + out.stride(1), + k_cache, + seq_lens, + block_table, + offset, + gather_lens, + max_blocks_per_seq=block_table.shape[-1], + head_dim=512, + scale_dim=V41_NVFP4_NUM_SCALES, + quant_block=V41_NVFP4_QUANT_BLOCK, + cache_block_size=block_size, + block_stride=k_cache.stride(0), + ) + return + if k_cache.shape[-1] == V41_BYTES_PER_TOKEN: _dequantize_and_gather_k_mxfp8_kernel[(num_reqs, NUM_WORKERS)]( out, @@ -581,14 +673,16 @@ def dequantize_and_gather_k_cache( ) -> None: """Dequantize and gather a paged DSv4 K cache. - The record is read off ``k_cache.shape[-1]``; see the module header. + The record is read off ``k_cache.shape[-1]``; see the module header. Only + the fp8 records have a CuteDSL gather, so NVFP4 always takes the Triton + path. ``use_fnuz`` MUST match the encoder of the specific cache being read: ``False`` for ``compressed_k_cache`` (Triton encoder is OCP everywhere), ``current_platform.is_fp8_fnuz()`` for ``swa_k_cache`` (C++ encoder writes FNUZ on gfx942 and OCP on gfx950). """ - if has_cutedsl(): + if has_cutedsl() and k_cache.shape[-1] != V41_NVFP4_BYTES_PER_TOKEN: # lazily import, otherwise some tests fail due to CUDA driver init failure. from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import ( _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL, diff --git a/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py index 96f15849ddd3..0f7830cf58c0 100644 --- a/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py @@ -4,6 +4,7 @@ import torch +from vllm.models.deepseek_v4.common.ops.fused_indexer_q import _fp32x2_to_fp4x2 from vllm.platforms import current_platform from vllm.triton_utils import tl, triton @@ -232,13 +233,14 @@ def rope_quant_insert( The BF16 latent supplies both NoPE quantization and RoPE input. It is read only for valid slots at group boundaries. The cache dtype selects the - layout: ``uint8`` is the fp8_ds_mla paged layout, whose record the - per-token byte width names -- 584 B for V4 (576 value bytes and eight - segregated UE8M0 scale bytes, including one zero padding scale) or 528 B - for V4.1 (512 MXFP8 value bytes covering the RoPE dims too, then 16 UE8M0 - scales of 32 dims each). ``bfloat16`` and ``float8_e4m3fn`` are the plain - [448 NoPE | 64 RoPE] rows read by FlashInfer, the latter scaled by the - per-tensor ``fp8_scale``. + layout: ``uint8`` is a paged FlashMLA layout, whose record the per-token + byte width names -- 584 B for V4 (576 value bytes and eight segregated + UE8M0 scale bytes, including one zero padding scale), 528 B for V4.1 + (512 MXFP8 value bytes covering the RoPE dims too, then 16 UE8M0 scales of + 32 dims each), or 288 B for V4.1 NVFP4 (256 bytes of e2m1 pairs then 32 + e4m3 scales of 16 dims each), which only the compressed cache uses. + ``bfloat16`` and ``float8_e4m3fn`` are the plain [448 NoPE | 64 RoPE] rows + read by FlashInfer, the latter scaled by the per-tensor ``fp8_scale``. """ assert compress_ratio in (1, 2) assert latent.shape[1] == 512 and latent.dtype == torch.bfloat16 @@ -249,14 +251,14 @@ def rope_quant_insert( return launch_kwargs = {"launch_pdl": False} if current_platform.is_cuda() else {} if kv_cache.dtype == torch.uint8: - assert kv_cache.shape[-1] in (584, 528), ( + kernel = { + 584: _rope_quant_insert_kernel, + 528: _rope_quant_insert_mxfp8_kernel, + 288: _rope_quant_insert_nvfp4_kernel, + }.get(kv_cache.shape[-1]) + assert kernel is not None, ( f"unsupported paged KV record width {kv_cache.shape[-1]}" ) - kernel = ( - _rope_quant_insert_mxfp8_kernel - if kv_cache.shape[-1] == 528 - else _rope_quant_insert_kernel - ) kernel[(num_tokens,)]( latent, positions, @@ -402,6 +404,62 @@ def _rope_quant_insert_mxfp8_kernel( tl.store(scales + tl.arange(0, 16), encoded.to(tl.uint8)) +@triton.jit +def _rope_quant_insert_nvfp4_kernel( + latent, + positions, + cos_sin, + cache, + cache_slots, + COS_STRIDE: tl.constexpr, + CACHE_STRIDE: tl.constexpr, + CACHE_BLOCK: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + SANITIZE_CACHE_NANS: tl.constexpr, +): + """V4.1 NVFP4 record: RoPE, then e2m1 with one e4m3 scale per 16 dims. + + The scale is ``amax / 6`` (6 is e2m1's largest magnitude) clamped to the + e4m3 range, with no per-tensor scale on top. + """ + t = tl.program_id(0) + slot = tl.load(cache_slots + t) + if slot < 0: + return + position = tl.load(positions + t) + if (position + 1) % COMPRESS_RATIO != 0: + return + d = tl.arange(0, 512) + normed = tl.load(latent + t.to(tl.int64) * 512 + d).to(tl.float32) + + # NoPE pairs load (cos, sin) = (1, 0), so the rotation is the identity there. + even, odd = tl.split(tl.reshape(normed, (256, 2))) + pair = tl.arange(0, 256) - 224 + cs = cos_sin + (position // COMPRESS_RATIO * COMPRESS_RATIO) * COS_STRIDE + c = tl.load(cs + tl.maximum(pair, 0), pair >= 0, other=1.0).to(tl.float32) + s = tl.load(cs + 32 + tl.maximum(pair, 0), pair >= 0, other=0.0).to(tl.float32) + rotated = tl.interleave(even * c - odd * s, odd * c + even * s) + if SANITIZE_CACHE_NANS: + rotated = tl.where(rotated == rotated, rotated, 0.0) + + tiles = tl.reshape(rotated, (32, 16)) + amax = tl.max(tl.abs(tiles), 1) + # 2**-9 is the smallest normal e4m3 magnitude; 448 the largest. + scale = tl.clamp(amax * (1.0 / 6.0), 0.001953125, 448.0).to(tl.float8e4nv) + # Round-to-nearest division: Triton's default div.full misplaces values + # that land exactly on an e2m1 tie. + scaled = tl.math.div_rn(tiles, tl.reshape(scale.to(tl.float32), (32, 1))) + lo, hi = tl.split(tl.reshape(tl.reshape(scaled, (512,)), (256, 2))) + + page = cache + (slot // CACHE_BLOCK).to(tl.int64) * CACHE_STRIDE + tl.store( + page + (slot % CACHE_BLOCK) * 256 + tl.arange(0, 256), + _fp32x2_to_fp4x2(lo, hi), + ) + scales = page + CACHE_BLOCK * 256 + (slot % CACHE_BLOCK) * 32 + tl.store(scales + tl.arange(0, 32), scale.to(tl.uint8, bitcast=True)) + + @triton.jit def _rope_plain_insert_kernel( latent, diff --git a/vllm/models/deepseek_v41/common/ops/fused_layout.py b/vllm/models/deepseek_v41/common/ops/fused_layout.py new file mode 100644 index 000000000000..071afde6bca4 --- /dev/null +++ b/vllm/models/deepseek_v41/common/ops/fused_layout.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Weight permutations for FlashMLA's mega-attention kernel. + +The kernel reads Q with 16-element head-dim chunks interleaved across heads and +writes O with 32-element chunks interleaved across the 8 heads of a ``wo_a`` +group. Every permutation here satisfies ``fused = standard[perm]``, and is +applied once to ``wq_b`` rows and ``wo_a`` columns at load time so the +surrounding GEMMs produce and consume the kernel's layouts directly -- no +per-step shuffle. +""" + +import torch + +HEAD_DIM = 512 +Q_CHUNK = 16 +O_CHUNK = 32 +WV_GROUP_SIZE = 8 + + +def q_fused_permutation(num_heads: int, head_dim: int = HEAD_DIM) -> torch.Tensor: + """``fused[(d // 16) * (H * 16) + h * 16 + d % 16] = standard[h * D + d]``.""" + h = torch.arange(num_heads).view(num_heads, 1) + d = torch.arange(head_dim).view(1, head_dim) + fused_index = (d // Q_CHUNK) * (num_heads * Q_CHUNK) + h * Q_CHUNK + d % Q_CHUNK + perm = torch.empty(num_heads * head_dim, dtype=torch.long) + perm[fused_index.reshape(-1)] = torch.arange(num_heads * head_dim) + return perm + + +def o_fused_permutation( + heads_per_group: int = WV_GROUP_SIZE, head_dim: int = HEAD_DIM +) -> torch.Tensor: + """``fused[(c * G + h) * 32 + j] = standard[h * D + c * 32 + j]`` per group.""" + h = torch.arange(heads_per_group).view(-1, 1, 1) + c = torch.arange(head_dim // O_CHUNK).view(1, -1, 1) + j = torch.arange(O_CHUNK).view(1, 1, -1) + fused_index = (c * heads_per_group + h) * O_CHUNK + j + perm = torch.empty(heads_per_group * head_dim, dtype=torch.long) + perm[fused_index.reshape(-1)] = torch.arange(heads_per_group * head_dim) + return perm + + +def o_fused_chunk_permutation( + heads_per_group: int = WV_GROUP_SIZE, head_dim: int = HEAD_DIM +) -> torch.Tensor: + """Per-32-element-chunk form of :func:`o_fused_permutation`, for scales.""" + return o_fused_permutation(heads_per_group, head_dim)[::O_CHUNK] // O_CHUNK + + +def permute_q_to_fused(q: torch.Tensor) -> torch.Tensor: + """``[N, H, D]`` standard layout -> the same shape in the fused layout.""" + n, h, d = q.shape + perm = q_fused_permutation(h, d).to(q.device) + return q.reshape(n, h * d)[:, perm].view(n, h, d) + + +def _bytes_view(t: torch.Tensor) -> torch.Tensor: + """Byte view of a 1-byte-element tensor, so fp8/ue8m0 can be gathered.""" + return t.view(torch.uint8) if t.element_size() == 1 else t + + +def permute_wq_b_( + weight: torch.Tensor, weight_scale: torch.Tensor, num_local_heads: int +) -> None: + """Permute an MXFP8 ``wq_b`` shard's rows and per-row scales, in place.""" + head_dim = weight.shape[0] // num_local_heads + perm = q_fused_permutation(num_local_heads, head_dim).to(weight.device) + for t in (weight, weight_scale): + b = _bytes_view(t) + b.copy_(b[perm]) + + +def permute_wo_a_( + weight: torch.Tensor, + weight_scale: torch.Tensor, + heads_per_group: int = WV_GROUP_SIZE, +) -> None: + """Permute an MXFP8 ``wo_a`` shard's input columns and scales, in place.""" + head_dim = weight.shape[1] // heads_per_group + perm = o_fused_permutation(heads_per_group, head_dim).to(weight.device) + w = _bytes_view(weight) + w.copy_(w[:, perm]) + chunk_perm = o_fused_chunk_permutation(heads_per_group, head_dim) + s = _bytes_view(weight_scale) + s.copy_(s[:, chunk_perm.to(weight.device)]) diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index 214681f3a932..9a07189bc3bb 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/vllm/models/deepseek_v41/nvidia/dspark.py @@ -477,8 +477,22 @@ def _finalize_moe(self) -> None: for layer in self.model.layers: layer.ffn.finalize_mega_moe_weights() + def _finalize_attn(self) -> None: + """Run the attention backend's post-load weight step on the draft layers. + + They are ordinary DeepseekV4DecoderLayers, so they take the engine-wide + attention backend and owe it whatever the target model owes it -- mega + attention permutes wq_b / wo_a here and refuses to run without it. + Idempotent, like the target's. + """ + for layer in self.model.layers: + finalize = getattr(layer.attn, "finalize_loaded_weights", None) + if finalize is not None: + finalize() + def process_weights_after_loading(self) -> None: self._finalize_moe() + self._finalize_attn() def _remap_dspark_name(self, name: str) -> str | None: """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path. diff --git a/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py b/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py new file mode 100644 index 000000000000..233d2f80f903 --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py @@ -0,0 +1,497 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek V4.1 attention on FlashMLA's mega-attention kernel (SM100). + +One kernel does Q RoPE, sparse attention, the inverse RoPE of the output and +its FP8 cast, and writes straight into the buffer ``wo_a`` consumes. So the +layer declares ``accepts_unnormed_unroped_query`` -- the kernel RoPEs Q itself, +leaving the fused KV insert only the zero-pad to the kernel's head count -- and +its ``_alloc_attn_out`` / ``_o_proj`` pair speaks QuantizedActivation instead +of bf16, since the output needs no rotation or quantization here. + +``wq_b`` rows and ``wo_a`` columns are permuted once at load so the surrounding +GEMMs speak the kernel's chunk-interleaved layouts directly. A step's prefill +and decode segments write disjoint token ranges of one output buffer pair, so +a single ``wo_a`` einsum covers the whole step. +""" + +from dataclasses import replace +from typing import TYPE_CHECKING, cast + +import torch + +from vllm.config import VllmConfig +from vllm.config.cache import CacheDType +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic +from vllm.models.deepseek_v41.common.ops import ( + combine_topk_swa_indices, + compute_global_topk_indices_and_lens, + dequantize_and_gather_k_cache, +) +from vllm.models.deepseek_v41.common.ops.fused_layout import ( + WV_GROUP_SIZE, + permute_wo_a_, + permute_wq_b_, +) +from vllm.models.deepseek_v41.nvidia.flashmla import DeepseekV4FlashMLAAttention +from vllm.models.deepseek_v41.sparse_mla import ( + DeepseekV4FlashMLAMetadata, + FlashMLAMegaAttnBackend, +) +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import fp8_einsum, get_tma_aligned_size +from vllm.utils.math_utils import round_up +from vllm.v1.attention.ops.flashmla import is_flashmla_sparse_supported +from vllm.v1.worker.workspace import current_workspace_manager + +if TYPE_CHECKING: + from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata + +_HEAD_DIM_V = 512 +_ROPE_DIM = 64 +# The kernel always emits per-32 scales, whatever num_per_channels says. +_QUANT_GROUP = 32 +# Fixed kernel conventions: GPT-J (non-neox) RoPE, TMA-aligned col-major +# scales, round-to-nearest, packed ue8m0 -- i.e. what DeepGEMM's fp8_einsum +# consumes. Q norm is off: the V4.1 layer norms qr before wq_b. +_ROPE_ARGS = (False, _ROPE_DIM) +_SF_ARGS = (_QUANT_GROUP, True, True, True) + +# Decode always goes through the mega kernel: there is no batch size at which +# falling back to the split-KV decode pays, so there is no minimum-token +# threshold and no fallback path to maintain. What decides whether this layer +# is worth selecting at all is TP, not the batch size. +# +# Measured on one GB300 with benchmarks/kernels/benchmark_dsv41_mega_attn.py, +# CUDA graphs on, counting every per-step op each path runs (64-head model, +# topk_swa=128, topk_extra=512, NVFP4 compressed cache), as mega / split-KV +# microseconds: +# +# s_q 1 32 128 256 512 +# TP1 (64 live heads) 29/29 31/31 31/40 53/72 96/140 +# TP2 (32 live) 29/29 31/30 35/35 61/63 110/119 +# TP4 (16 live) 29/29 31/29 34/34 59/58 109/109 +# TP8 (8 live) 29/29 31/29 33/33 59/57 107/104 +# +# The kernel absorbs the inverse RoPE and the FP8 cast, whose cost scales with +# *live* heads, while its own cost scales with *padded* heads -- always 64, +# since FlashMLA's FP8 decode takes only h_q in {64, 128}. So it wins at TP1 +# (up to 1.45x, and TP1 is also where the fused insert skips the Q pad +# entirely) and from TP2 down the two cancel to within a few percent, with a +# fixed ~2 us penalty at small s_q where the padded-head cost dominates. +# Measure before assuming a win at high TP -- and measure under CUDA graphs: +# in eager mode launch overhead swamps both paths and overstates the fused one +# several-fold. + + +def is_flashmla_mega_attn_supported() -> tuple[bool, str | None]: + """Whether FlashMLA's fused mega-attention kernels are usable here.""" + is_available, maybe_reason = is_flashmla_sparse_supported() + if not is_available: + return False, maybe_reason + if not current_platform.is_device_capability_family(100): + return False, "FlashMLA mega attention requires sm_10x GPUs." + if not hasattr(torch.ops._flashmla_C, "fused_norm_rope_attn_rope_cast_decode"): + return False, "vllm._flashmla_C was built without the mega attention ops." + if not hasattr(torch.ops._flashmla_C, "permute_q_b_proj"): + return False, "vllm._flashmla_C predates the mega attention weight permutes." + return True, None + + +def alloc_mega_attn_output( + num_tokens: int, + n_wv_group: int, + device: torch.device, +) -> QuantizedActivation: + """Allocate the output buffer pair the mega-attention kernels write into. + + One pair per forward step: the prefill and decode segments fill disjoint + token ranges, so a single ``fp8_einsum`` over all ``num_tokens`` can + consume the result instead of one call per segment. + + The scale buffer is MN-major -- stride 1 along tokens, head-dim stride + ``ceil4(num_tokens)`` -- which is both what the kernel requires of a + caller-provided buffer and what DeepGEMM expects for this N. + """ + d = WV_GROUP_SIZE * _HEAD_DIM_V + data = torch.empty( + (num_tokens, n_wv_group, d), dtype=torch.float8_e4m3fn, device=device + ) + aligned = get_tma_aligned_size(num_tokens, torch.int32.itemsize) + scale = torch.empty( + (n_wv_group, d // (_QUANT_GROUP * 4), aligned), + dtype=torch.int32, + device=device, + ).permute(2, 0, 1)[:num_tokens] + return QuantizedActivation( + data=data, + scale=scale, + orig_dtype=torch.bfloat16, + orig_shape=data.shape, + quant_key=kMxfp8Dynamic, + ) + + +def _token_slice(out: QuantizedActivation, start: int, end: int) -> QuantizedActivation: + """The ``[start, end)`` token slice of a mega-attention output buffer.""" + data = out.data[start:end] + return replace(out, data=data, scale=out.scale[start:end], orig_shape=data.shape) + + +class DeepseekV4MegaAttnAttention(DeepseekV4FlashMLAAttention): + """FlashMLA mega-attention layer for DeepSeek V4.1 (SM100).""" + + backend_cls = FlashMLAMegaAttnBackend + + @classmethod + def is_available_for(cls, vllm_config: VllmConfig) -> bool: + """Whether this layer can serve the configured model on this device. + + The kernel is SM100-only and wants ``WV_GROUP_SIZE`` heads per ``wo_a`` + group, which stops holding once TP divides the head count far enough + (TP16 on a 64-head model leaves 4). ``__init__`` raises on both, so the + default-backend selector asks here first rather than turning an + unsupported topology into a startup crash. + """ + if not is_flashmla_mega_attn_supported()[0]: + return False + config = vllm_config.model_config.hf_config + config = getattr(config, "text_config", config) + n_heads = getattr(config, "num_attention_heads", 0) or 0 + n_groups = getattr(config, "o_groups", 0) or 0 + tp_size = vllm_config.parallel_config.tensor_parallel_size + if n_heads % tp_size or n_groups % tp_size: + return False + n_local_heads, n_local_groups = n_heads // tp_size, n_groups // tp_size + if not n_local_groups or n_local_heads % WV_GROUP_SIZE: + return False + return n_local_heads // n_local_groups == WV_GROUP_SIZE + + def __init__(self, vllm_config: VllmConfig, *args, **kwargs) -> None: + super().__init__(vllm_config, *args, **kwargs) + if self.n_local_heads % WV_GROUP_SIZE: + raise ValueError( + f"{self.prefix}: mega attention needs the local head count " + f"({self.n_local_heads}) to be a multiple of {WV_GROUP_SIZE}." + ) + if self.n_local_heads // self.n_local_groups != WV_GROUP_SIZE: + raise ValueError( + f"{self.prefix}: mega attention needs {WV_GROUP_SIZE} heads per " + "wo_a group." + ) + self.n_wv_group = self.padded_heads // WV_GROUP_SIZE + self._fused_layouts_ready = False + + # ---- interface contract ------------------------------------------------ + + @property + def accepts_unnormed_unroped_query(self) -> bool: + return True + + @property + def packed_kv_cache_dtype(self) -> CacheDType: + # This kernel is the only one that reads an NVFP4 compressed cache, so + # it is what an unspecific --kv-cache-dtype resolves to here. + return "nvfp4_ds_mla" + + def _alloc_attn_out( + self, num_tokens: int, hidden_states: torch.Tensor + ) -> QuantizedActivation: + return alloc_mega_attn_output(num_tokens, self.n_wv_group, hidden_states.device) + + def _o_proj( + self, attn_out: QuantizedActivation, positions: torch.Tensor + ) -> torch.Tensor: + """Grouped ``wo_a`` then ``wo_b`` over the kernel's quantized output. + + The inverse RoPE and the FP8 cast happened inside the attention + kernel, and ``wo_a`` is permuted for its output layout, so there is + nothing to rotate or quantize here and ``positions`` is unused. The + scale already is DeepGEMM's packed-ue8m0 MN-major layout, so the + einsum consumes the kernel's output with no repacking. + """ + groups = self.n_local_groups + z = torch.empty( + (attn_out.data.shape[0], groups, self.o_lora_rank), + dtype=torch.bfloat16, + device=attn_out.data.device, + ) + fp8_einsum( + "bhr,hdr->bhd", + (attn_out.data[:, :groups], attn_out.scale[:, :groups]), + (self.wo_a.weight, self.wo_a.weight_scale), + z, + recipe=self._einsum_recipe, + ) + return self.wo_b(z.flatten(1)) + + # ---- weights ----------------------------------------------------------- + + def finalize_loaded_weights(self) -> None: + """Permute wq_b rows / wo_a columns into the kernel's layouts. + + Idempotent: a second post-load pass must not permute twice. + """ + if self._fused_layouts_ready: + return + permute_wq_b_( + self.wq_b.weight.data, self.wq_b.weight_scale.data, self.n_local_heads + ) + permute_wo_a_( + self.wo_a.weight.data, + self.wo_a.weight_scale.data, + self.n_local_heads // self.n_local_groups, + ) + self._fused_layouts_ready = True + + # ---- forward ----------------------------------------------------------- + + def forward_mqa( + self, + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + output: QuantizedActivation, + ) -> None: + if not self._fused_layouts_ready: + raise RuntimeError( + f"{self.prefix}: wq_b / wo_a were never permuted for the mega " + "attention kernel; refusing to run with mismatched layouts." + ) + attn_metadata = get_forward_context().attn_metadata + if attn_metadata is None: + # Warmup dummy run: reserve the prefill workspace, produce zeros. + self._reserve_prefill_workspace(q) + output.data.zero_() + output.scale.zero_() + return + + assert isinstance(attn_metadata, dict) + # q was zero-padded to the kernel's head count by the fused KV insert. + assert q.shape[1] == self.padded_heads + flashmla_metadata = cast( + DeepseekV4FlashMLAMetadata | None, + attn_metadata.get(self.compressed_cache_prefix) + if self.compressed_cache_prefix is not None + else None, + ) + swa_metadata = cast( + "DeepseekSparseSWAMetadata", attn_metadata[self.swa_cache_layer.prefix] + ) + # The kernel takes int32 RoPE positions. + positions_int32 = positions.to(torch.int32) + num_decode_tokens = swa_metadata.num_decode_tokens + + if swa_metadata.num_prefills > 0: + self._forward_prefill_mega( + q[num_decode_tokens:], + positions_int32[num_decode_tokens:], + flashmla_metadata, + swa_metadata, + output, + num_decode_tokens, + ) + if swa_metadata.num_decodes > 0: + self._forward_decode_mega( + q[:num_decode_tokens], + positions_int32[:num_decode_tokens], + flashmla_metadata, + swa_metadata, + _token_slice(output, 0, num_decode_tokens), + ) + + def _reserve_prefill_workspace(self, q: torch.Tensor) -> None: + swa_only = self.compress_ratio == 0 + n = 0 if swa_only else -(-self.max_model_len // self.compress_ratio) + m = n + self.window_size + self.max_num_batched_tokens + if swa_only: + top_k = 0 + else: + assert self.topk_indices_buffer is not None + top_k = self.topk_indices_buffer.shape[-1] + combined_topk = round_up(top_k + self.window_size + self.max_image_tokens, 128) + current_workspace_manager().get_simultaneous( + ((self.PREFILL_CHUNK_SIZE, m, q.shape[-1]), torch.bfloat16), + ((self.max_num_batched_tokens, combined_topk), torch.int32), + ((self.max_num_batched_tokens,), torch.int32), + ) + + def _decode_compressed_kv_and_topk( + self, + flashmla_metadata: DeepseekV4FlashMLAMetadata | None, + swa_metadata: "DeepseekSparseSWAMetadata", + ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + """The compressed cache and its ``[n, topk]`` slot indices / lengths.""" + if self.compress_ratio == 0: + return None, None, None + assert flashmla_metadata is not None + assert swa_metadata.is_valid_token is not None + assert self.topk_indices_buffer is not None + num_decode_tokens = swa_metadata.num_decode_tokens + indices, lens = compute_global_topk_indices_and_lens( + self.topk_indices_buffer[:num_decode_tokens], + swa_metadata.token_to_req_indices, + flashmla_metadata.block_table[: swa_metadata.num_decodes], + flashmla_metadata.block_size // self.compress_ratio, + swa_metadata.is_valid_token[:num_decode_tokens], + ) + return self._compressed_kv_cache().unsqueeze(-2), indices, lens + + def _forward_decode_mega( + self, + q: torch.Tensor, + positions_int32: torch.Tensor, + flashmla_metadata: DeepseekV4FlashMLAMetadata | None, + swa_metadata: "DeepseekSparseSWAMetadata", + out: QuantizedActivation, + ) -> None: + """Mega attention over the paged quantized caches, in place into ``out``. + + The caches are ``[num_blocks, page, 1, bytes]``, the record taken from + ``bytes`` (528 V4.1 fp8, 288 V4.1 NVFP4; NVFP4 only as the compressed + cache beside an fp8 SWA one), and the indices ``[s_q, topk]`` int32 + slot ids (``block * page + offset``, ``-1`` invalid). + """ + extra_cache, extra_idx, extra_len = self._decode_compressed_kv_and_topk( + flashmla_metadata, swa_metadata + ) + assert swa_metadata.decode_swa_indices is not None + torch.ops._flashmla_C.fused_norm_rope_attn_rope_cast_decode( + q, + self.swa_cache_layer.kv_cache.unsqueeze(-2), + swa_metadata.decode_swa_indices.view(q.shape[0], -1), + self.scale, + _HEAD_DIM_V, + self.attn_sink, + swa_metadata.decode_swa_lens, + extra_cache, + extra_idx, + extra_len, + False, # enable_q_norm + 0.0, # rms_norm_eps + positions_int32, + *_ROPE_ARGS, + self.rotary_emb.cos_sin_cache, + self.n_wv_group, + *_SF_ARGS, + out.data, + out.scale, + ) + + def _forward_prefill_mega( + self, + q: torch.Tensor, + positions_int32: torch.Tensor, + flashmla_metadata: DeepseekV4FlashMLAMetadata | None, + swa_metadata: "DeepseekSparseSWAMetadata", + out: QuantizedActivation, + token_base: int, + ) -> None: + swa_only = self.compress_ratio == 0 + num_decodes = swa_metadata.num_decodes + num_decode_tokens = swa_metadata.num_decode_tokens + seq_lens = swa_metadata.prefill_seq_lens + gather_lens = swa_metadata.prefill_gather_lens + query_start_loc_cpu = swa_metadata.query_start_loc_cpu + query_start_loc = swa_metadata.query_start_loc + assert seq_lens is not None and gather_lens is not None + assert query_start_loc_cpu is not None and query_start_loc is not None + prefill_token_base = query_start_loc_cpu[num_decodes] + assert self.topk_indices_buffer is not None + topk_indices = self.topk_indices_buffer[num_decode_tokens:][ + : swa_metadata.num_prefill_tokens + ] + top_k = 0 if swa_only else topk_indices.shape[-1] + chunk_plan = swa_metadata.get_prefill_chunk_plan( + compress_ratio=self.compress_ratio, + prefill_chunk_size=self.PREFILL_CHUNK_SIZE, + has_compressed=not swa_only, + ) + assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0" + workspace_manager = current_workspace_manager() + combined_topk = round_up(top_k + self.window_size + self.max_image_tokens, 128) + for chunk_start, chunk_end, chunk_n, chunk_m in chunk_plan: + chunk_size = chunk_end - chunk_start + kv_ws, idx_ws, lens_ws = workspace_manager.get_simultaneous( + ((chunk_size, chunk_m, q.shape[-1]), torch.bfloat16), + ((self.max_num_batched_tokens, combined_topk), torch.int32), + ((self.max_num_batched_tokens,), torch.int32), + ) + if not swa_only: + assert flashmla_metadata is not None + dequantize_and_gather_k_cache( + kv_ws[:chunk_size], + self._compressed_kv_cache(), + seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio, + gather_lens=None, + block_table=flashmla_metadata.block_table[num_decodes:][ + chunk_start:chunk_end + ], + block_size=flashmla_metadata.block_size // self.compress_ratio, + offset=0, + ) + dequantize_and_gather_k_cache( + kv_ws[:chunk_size], + self.swa_cache_layer.kv_cache, + seq_lens=seq_lens[chunk_start:chunk_end], + gather_lens=gather_lens[chunk_start:chunk_end], + block_table=swa_metadata.block_table[num_decodes:][ + chunk_start:chunk_end + ], + block_size=swa_metadata.block_size, + offset=chunk_n, + ) + first, last = num_decodes + chunk_start, num_decodes + chunk_end + qs = int(query_start_loc_cpu[first] - prefill_token_base) + qe = int(query_start_loc_cpu[last] - prefill_token_base) + combined_indices, combined_lens = combine_topk_swa_indices( + topk_indices[qs:qe], + query_start_loc[first : last + 1], + seq_lens[chunk_start:chunk_end], + gather_lens[chunk_start:chunk_end], + self.window_size, + self.compress_ratio, + top_k, + chunk_m, + chunk_n, + out=(idx_ws[: qe - qs], lens_ws[: qe - qs]), + left_visible=( + swa_metadata.prefill_left_visible[ + num_decode_tokens + qs : num_decode_tokens + qe + ] + if swa_metadata.prefill_left_visible is not None + else None + ), + right_visible=( + swa_metadata.prefill_right_visible[ + num_decode_tokens + qs : num_decode_tokens + qe + ] + if swa_metadata.prefill_right_visible is not None + else None + ), + max_image_tokens=self.max_image_tokens, + ) + chunk_out = _token_slice(out, token_base + qs, token_base + qe) + # Mega attention over the gathered non-paged bf16 KV (RoPE already + # applied), in place into this chunk's token range. `indices` + # entries outside [0, s_kv) skip. + torch.ops._flashmla_C.fused_norm_rope_attn_rope_cast_fwd( + q[qs:qe], + kv_ws.view(-1, 1, q.shape[-1]), + combined_indices.unsqueeze(1), + self.scale, + _HEAD_DIM_V, + self.attn_sink, + combined_lens, + False, # enable_q_norm + 0.0, # rms_norm_eps + positions_int32[qs:qe], + *_ROPE_ARGS, + self.rotary_emb.cos_sin_cache, + self.n_wv_group, + *_SF_ARGS, + chunk_out.data, + chunk_out.scale, + ) diff --git a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py index 7bed50bde71a..0b43cfad190b 100644 --- a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py @@ -201,7 +201,8 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): def get_padded_num_q_heads(cls, num_heads: int) -> int: return _pad_to_supported_q_heads(num_heads) - def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + o = attn_out[:, : self.n_local_heads, :] return deep_gemm_fp8_o_proj( o, positions, @@ -562,7 +563,8 @@ def _as_sparse_cache(kv_cache: torch.Tensor) -> torch.Tensor: def get_padded_num_q_heads(cls, num_heads: int) -> int: return _pad_to_supported_q_heads(num_heads) - def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + o = attn_out[:, : self.n_local_heads, :] return deep_gemm_fp8_o_proj( o, positions, diff --git a/vllm/models/deepseek_v41/nvidia/flashmla.py b/vllm/models/deepseek_v41/nvidia/flashmla.py index ec834aa80dd9..47a7b35c76d6 100644 --- a/vllm/models/deepseek_v41/nvidia/flashmla.py +++ b/vllm/models/deepseek_v41/nvidia/flashmla.py @@ -58,7 +58,8 @@ def __init__(self, *args, **kwargs) -> None: self._o_proj_block_size ) - def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + o = attn_out[:, : self.n_local_heads, :] return deep_gemm_fp8_o_proj( o, positions, diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 9c9719e97729..21fa00e5c963 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -64,6 +64,9 @@ make_deepseek_v4_expert_params_mapping, ) from vllm.models.deepseek_v41.attention import DeepseekV4Attention +from vllm.models.deepseek_v41.nvidia.flash_mla_mega_attn import ( + DeepseekV4MegaAttnAttention, +) from vllm.models.deepseek_v41.nvidia.flashinfer_sparse import ( DeepseekV4FlashInferMLAAttention, DeepseekV4FlashInferSM120Attention, @@ -119,8 +122,9 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]: The generic CUDA backend selector does not instantiate DSv4 layers directly, so map generic sparse-MLA choices to the DSv4-specialized attention class. - Without an explicit backend, SM12 defaults to FlashInfer while the other - CUDA arches keep the FlashMLA path. + Without an explicit backend: SM12 takes FlashInfer, SM100 takes mega + attention where the topology allows it, and everything else keeps the + FlashMLA path. """ backend = vllm_config.attention_config.backend device_capability = current_platform.get_device_capability() @@ -140,6 +144,8 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]: if device_capability is not None and device_capability.major == 12: return DeepseekV4FlashInferSM120Attention return DeepseekV4FlashInferMLAAttention + if backend is AttentionBackendEnum.FLASHMLA_MEGA_ATTN_DSV41: + return DeepseekV4MegaAttnAttention if backend in ( AttentionBackendEnum.FLASHMLA_SPARSE, AttentionBackendEnum.FLASHMLA_SPARSE_DSV4, @@ -149,6 +155,14 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]: if device_capability is not None and device_capability.major == 12: return DeepseekV4FlashInferSM120Attention + # Mega attention is the SM100 default: it fuses Q RoPE, sparse attention, + # the output's inverse RoPE and its FP8 cast into one launch, and brings + # the 288 B NVFP4 compressed record -- the format the reference + # implementation itself stores. It declines topologies it cannot serve + # (non-SM100, TP that leaves fewer than WV_GROUP_SIZE heads per wo_a + # group, a build without the kernel), which then fall through to FlashMLA. + if DeepseekV4MegaAttnAttention.is_available_for(vllm_config): + return DeepseekV4MegaAttnAttention return DeepseekV4FlashMLAAttention @@ -936,6 +950,17 @@ def finalize_mega_moe_weights(self) -> None: for layer in islice(self.layers, self.start_layer, self.end_layer): layer.ffn.finalize_mega_moe_weights() + def finalize_mega_attn_weights(self) -> None: + """Permute wq_b / wo_a into FlashMLA's mega-attention layouts. + + A no-op for every other attention layer, and idempotent, so a second + post-load pass cannot permute twice. + """ + for layer in islice(self.layers, self.start_layer, self.end_layer): + finalize = getattr(layer.attn, "finalize_loaded_weights", None) + if finalize is not None: + finalize() + def finalize_mhc_broadcast_weights(self) -> None: if not get_pp_group().is_first_rank or self.start_layer >= self.end_layer: return @@ -1190,6 +1215,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: def process_weights_after_loading(self) -> None: self.model.finalize_mega_moe_weights() self.model.finalize_mhc_broadcast_weights() + self.model.finalize_mega_attn_weights() def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: return self.model.get_expert_mapping() diff --git a/vllm/models/deepseek_v41/sparse_mla.py b/vllm/models/deepseek_v41/sparse_mla.py index 6a0f09f952e4..3d69499d2a80 100644 --- a/vllm/models/deepseek_v41/sparse_mla.py +++ b/vllm/models/deepseek_v41/sparse_mla.py @@ -229,3 +229,54 @@ def get_name() -> str: @staticmethod def get_builder_cls() -> type[DeepseekV4FlashMLAMetadataBuilder]: return DeepseekV4FlashMLAMetadataBuilder + + +class FlashMLAMegaAttnBackend(DeepseekV4FlashMLABackend): + """FlashMLA's mega-attention kernel: Q RoPE + sparse attention + inverse + RoPE + FP8 cast of the output, in one launch. + + Same metadata and KV-cache geometry as ``FLASHMLA_SPARSE_DSV41`` -- what + differs is the attention layer's interface contract (it takes an unnormed, + unroped Q and returns an already-inverse-RoPE'd, quantized output) and the + extra ``nvfp4_ds_mla`` compressed-cache record only this kernel can read. + SM100 only; the kernel has no SM90 instantiation. + """ + + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "auto", + "fp8_ds_mla", + "fp8", # alias for fp8_ds_mla + "nvfp4_ds_mla", # V4.1 fp8 SWA cache + NVFP4 compressed cache + ] + + @staticmethod + def get_name() -> str: + return "FLASHMLA_MEGA_ATTN_DSV41" + + @staticmethod + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + return [128] + + @classmethod + def supports_compute_capability(cls, capability: DeviceCapability) -> bool: + return capability.major == 10 + + @classmethod + def supports_combination( + cls, + head_size: int, + dtype: torch.dtype, + kv_cache_dtype: CacheDType | None, + block_size: int | None, + use_mla: bool, + has_sink: bool, + use_sparse: bool, + use_mm_prefix: bool, + device_capability: DeviceCapability, + ) -> str | None: + # Imported here: the layer module imports this backend class. + from vllm.models.deepseek_v41.nvidia.flash_mla_mega_attn import ( + is_flashmla_mega_attn_supported, + ) + + return is_flashmla_mega_attn_supported()[1] diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index c4f3cd8b4bdd..700f328a37c0 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -112,7 +112,10 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: # fp8_ds_mla's UE8M0 paged layout rounds its page up to the decode # kernel's TMA stride; contiguous bf16/fp8 cache uses the natural # element-size page. - uses_fp8_ds_mla_layout = self.cache_config.cache_dtype == "fp8_ds_mla" + uses_fp8_ds_mla_layout = self.cache_config.cache_dtype in ( + "fp8_ds_mla", + "nvfp4_ds_mla", + ) return SlidingWindowMLASpec( block_size=self.block_size, num_kv_heads=1, diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 571566db1f25..dcaaffaaa8eb 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -114,6 +114,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): "vllm.models.deepseek_v41.nvidia.flashinfer_sparse." "DeepseekV4FlashInferMLASparseBackend" ) + FLASHMLA_MEGA_ATTN_DSV41 = ( + "vllm.models.deepseek_v41.sparse_mla.FlashMLAMegaAttnBackend" + ) B12X = "vllm.v1.attention.backends.b12x.B12xPagedAttentionBackend" FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend" FLASH_ATTN_MLA_SPARSE = ( From 118e17f5ff301f67f1f4f5bdae2c9ae5e50242c5 Mon Sep 17 00:00:00 2001 From: wangxiyuan Date: Wed, 16 Sep 2026 18:08:36 +0800 Subject: [PATCH 20/20] [Platform] Move env check function to platform interface (#48599) Signed-off-by: wangxiyuan Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/engine/arg_utils.py | 2 +- vllm/envs.py | 9 --------- vllm/platforms/interface.py | 18 ++++++++++++++++++ 3 files changed, 19 insertions(+), 10 deletions(-) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index b420e84f542b..6128bc625589 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -2051,7 +2051,7 @@ def create_engine_config( device_config = DeviceConfig(device=cast(Device, current_platform.device_type)) - envs.validate_environ(self.fail_on_environ_validation) + current_platform.validate_environ(self.fail_on_environ_validation) # Check if the model is a speculator and override model/tokenizer/config # BEFORE creating ModelConfig, so the config is created with the target model diff --git a/vllm/envs.py b/vllm/envs.py index bc4724e2e4a7..333d75b0ff0e 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -2249,15 +2249,6 @@ def is_set(name: str): raise AttributeError(f"module {__name__!r} has no attribute {name!r}") -def validate_environ(hard_fail: bool) -> None: - for env in os.environ: - if env.startswith("VLLM_") and env not in environment_variables: - if hard_fail: - raise ValueError(f"Unknown vLLM environment variable detected: {env}") - else: - logger.warning("Unknown vLLM environment variable detected: %s", env) - - def compile_factors() -> dict[str, object]: """Return env vars used for torch.compile cache keys. diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index bd5d98191201..98db974a02e4 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -1334,6 +1334,24 @@ def is_arch_support_pdl(cls) -> bool: """ return False + @classmethod + def validate_environ(cls, hard_fail: bool) -> None: + """ + Validate environment variables for the current platform. + """ + from vllm import envs + + for env in os.environ: + if env.startswith("VLLM_") and env not in envs.environment_variables: + if hard_fail: + raise ValueError( + f"Unknown vLLM environment variable detected: {env}" + ) + else: + logger.warning( + "Unknown vLLM environment variable detected: %s", env + ) + class UnspecifiedPlatform(Platform): _enum = PlatformEnum.UNSPECIFIED