From 8eac21a602d4959eb8478c430d7fe1ca5b57311f Mon Sep 17 00:00:00 2001 From: BadrBasowid <61441185+BadrBasowid@users.noreply.github.com> Date: Fri, 24 Jul 2026 22:16:17 +0800 Subject: [PATCH 001/185] [ROCM] Fix AITER Fused AllReduce RMSNorm for Transformers Backend (#49673) Signed-off-by: BadrBasowid --- tests/compile/fusions_e2e/test_tp2_ar_rms.py | 18 ++++++++++++++++-- vllm/_aiter_ops.py | 4 ++-- 2 files changed, 18 insertions(+), 4 deletions(-) diff --git a/tests/compile/fusions_e2e/test_tp2_ar_rms.py b/tests/compile/fusions_e2e/test_tp2_ar_rms.py index b18c41658fd..a969e91daa4 100644 --- a/tests/compile/fusions_e2e/test_tp2_ar_rms.py +++ b/tests/compile/fusions_e2e/test_tp2_ar_rms.py @@ -181,8 +181,13 @@ def test_tp2_ar_rms_fp4_fusions( @multi_gpu_test(num_gpus=2) @pytest.mark.parametrize( - "model_name, matches_fn, model_kwargs, hf_overrides", - [llama3_8b, qwen3_a3b, gpt_oss_20b], + "model_name, matches_fn, model_kwargs, hf_overrides, model_impl", + [ + (*llama3_8b, "auto"), + (*llama3_8b, "transformers"), + (*qwen3_a3b, "auto"), + (*gpt_oss_20b, "auto"), + ], ) @pytest.mark.parametrize( "attn_backend", @@ -202,17 +207,26 @@ def test_tp2_ar_rms_fusions( matches_fn: Callable[[int], Matches], model_kwargs: dict, hf_overrides: Callable[[int], dict], + model_impl: str, attn_backend: AttentionBackendCase, n_layers: int, custom_ops: str, inductor_graph_partition: bool, run_e2e_fusion_test, ): + if model_impl == "transformers" and not current_platform.is_rocm(): + pytest.skip("Transformers 3D AR+RMS regression is ROCm-only") + matches = matches_fn(n_layers) + if model_impl == "transformers": + # TODO(BadrBasowid): Match the vLLM backend's fusion count once the + # separate residual add and RMSNorm operations are fused. + matches = matches._replace(aiter_ar_rms_fusion=1) # Reduce size of model and skip weight loading time model_kwargs["hf_overrides"] = hf_overrides(n_layers) model_kwargs["load_format"] = "dummy" + model_kwargs["model_impl"] = model_impl model_kwargs["max_model_len"] = 1024 model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} model_kwargs["disable_custom_all_reduce"] = False diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 4a915c0dd68..e4b08e90aca 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -764,7 +764,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl( total_bytes = input_.numel() * input_.element_size() hidden_dim = input_.shape[-1] - token_num = input_.shape[0] + token_num = input_.numel() // hidden_dim if input_.dtype in (torch.bfloat16, torch.float16): pack_size = 16 // input_.element_size() hidden_ok = hidden_dim % pack_size == 0 and hidden_dim // pack_size <= 1024 @@ -824,7 +824,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl( total_bytes = input_.numel() * input_.element_size() hidden_dim = input_.shape[-1] - token_num = input_.shape[0] + token_num = input_.numel() // hidden_dim if input_.dtype in (torch.bfloat16, torch.float16): pack_size = 16 // input_.element_size() hidden_ok = hidden_dim % pack_size == 0 and hidden_dim // pack_size <= 1024 From 7b40fb96450e1deef40ebc383ff549bea35b9b8f Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Fri, 24 Jul 2026 10:16:22 -0400 Subject: [PATCH 002/185] [UX] Reject incompatible nested runtime overrides (#49247) Signed-off-by: Taneem Ibrahim Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/test_config.py | 21 +++++++++-- tests/v1/worker/test_gpu_model_runner.py | 2 +- vllm/config/utils.py | 46 ++++++++++++++++-------- vllm/v1/worker/gpu_model_runner.py | 10 +++--- 4 files changed, 56 insertions(+), 23 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index a785297997b..71e078ef3a2 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -406,26 +406,41 @@ class _TestNestedConfig: a: _TestConfigFields = field(default_factory=lambda: _TestConfigFields(a=0)) +@dataclass +class _TestDerivedConfigFields(_TestConfigFields): + pass + + def test_update_config(): # Simple update config1 = _TestConfigFields(a=0) new_config1 = update_config(config1, {"a": 42}) assert new_config1.a == 42 # Nonexistent field - with pytest.raises(AssertionError): + with pytest.raises(ValueError, match=r"_TestConfigFields\.nonexistent"): new_config1 = update_config(config1, {"nonexistent": 1}) # Nested update with dataclass config2 = _TestNestedConfig() new_inner_config = _TestConfigFields(a=1, c="new_value") new_config2 = update_config(config2, {"a": new_inner_config}) assert new_config2.a == new_inner_config + # Declared field type, not the live value's subtype, defines valid overrides + config_with_derived = _TestNestedConfig(a=_TestDerivedConfigFields(a=0)) + new_config2 = update_config(config_with_derived, {"a": new_inner_config}) + assert new_config2.a is new_inner_config + # Nested update with unrelated dataclass + with pytest.raises(ValueError, match=r"_TestNestedConfig\.a"): + update_config(config2, {"a": _TestNestedConfig()}) # Nested update with dict config3 = _TestNestedConfig() new_config3 = update_config(config3, {"a": {"c": "new_value"}}) assert new_config3.a.c == "new_value" # Nested update with invalid type - with pytest.raises(AssertionError): - new_config3 = update_config(config3, {"a": "new_value"}) + with pytest.raises(ValueError, match=r"_TestNestedConfig\.a"): + update_config(config3, {"a": "new_value"}) + # Invalid nested field preserves its full path + with pytest.raises(ValueError, match=r"_TestNestedConfig\.a\.nonexistent"): + update_config(config3, {"a": {"nonexistent": 1}}) @pytest.mark.parametrize( diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 57931669c54..79e3a60e981 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -831,7 +831,7 @@ def test_update_config(model_runner): model_runner.update_config({"load_config": {"load_format": "dummy"}}) assert model_runner.load_config.load_format == "dummy" # Raise error on non-existing config - with pytest.raises(AssertionError): + with pytest.raises(ValueError, match="do_not_exist_config"): model_runner.update_config({"do_not_exist_config": "dummy"}) diff --git a/vllm/config/utils.py b/vllm/config/utils.py index 3df0f7210f7..6ab346d1a03 100644 --- a/vllm/config/utils.py +++ b/vllm/config/utils.py @@ -13,7 +13,7 @@ import textwrap from collections.abc import Callable, Mapping, Sequence, Set from dataclasses import MISSING, field, fields, is_dataclass from itertools import pairwise -from typing import TYPE_CHECKING, Any, Protocol, TypeVar, cast, overload +from typing import TYPE_CHECKING, Any, Protocol, TypeVar, cast, get_type_hints, overload import torch from pydantic import ConfigDict @@ -227,22 +227,38 @@ class SupportsMetricsInfo(Protocol): def metrics_info(self) -> dict[str, str]: ... -def update_config(config: ConfigT, overrides: dict[str, Any]) -> ConfigT: - processed_overrides = {} +def update_config(config: ConfigT, overrides: Mapping[str, Any]) -> ConfigT: + return _update_config(config, overrides, type(config).__name__) + + +def _update_config( + config: ConfigT, overrides: Mapping[str, Any], config_path: str +) -> ConfigT: + processed_overrides: dict[str, Any] = {} + field_types = get_type_hints(type(config)) for field_name, value in overrides.items(): - assert hasattr(config, field_name), ( - f"{type(config)} has no field `{field_name}`" - ) + field_path = f"{config_path}.{field_name}" + if not hasattr(config, field_name): + raise ValueError(f"{field_path} is not a valid config field") + current_value = getattr(config, field_name) - if is_dataclass(current_value) and not is_dataclass(value): - assert isinstance(value, dict), ( - f"Overrides to {type(config)}.{field_name} must be a dict" - f" or {type(current_value)}, but got {type(value)}" - ) - value = update_config( - current_value, # type: ignore[type-var] - value, - ) + if is_dataclass(current_value): + expected_type = field_types[field_name] + if isinstance(value, Mapping): + value = _update_config( + current_value, # type: ignore[type-var] + value, + field_path, + ) + elif not isinstance(value, expected_type): + expected_type_name = getattr( + expected_type, "__name__", str(expected_type) + ) + raise ValueError( + f"Override for {field_path} must be a mapping or " + f"{expected_type_name}, got {type(value).__name__}" + ) + processed_overrides[field_name] = value return replace(config, **processed_overrides) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 40d2ec0e4d6..95221c86931 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -5273,10 +5273,12 @@ class GPUModelRunner( def update_config(self, overrides: dict[str, Any]) -> None: allowed_config_names = {"load_config", "model_config"} for config_name, config_overrides in overrides.items(): - assert config_name in allowed_config_names, ( - f"Config `{config_name}` not supported. " - f"Allowed configs: {allowed_config_names}" - ) + if config_name not in allowed_config_names: + allowed = ", ".join(sorted(allowed_config_names)) + raise ValueError( + f"Config override '{config_name}' is not supported. " + f"Supported configs: {allowed}" + ) config = getattr(self, config_name) new_config = update_config(config, config_overrides) setattr(self, config_name, new_config) From 453f01783d44abc3f325bcc6efadf4246ec18873 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Fri, 24 Jul 2026 10:16:39 -0400 Subject: [PATCH 003/185] [UX] Improve data-parallel launch validation (#49124) Signed-off-by: Taneem Ibrahim Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> --- vllm/engine/arg_utils.py | 75 ++++++++++++++++++++++++++++------------ 1 file changed, 52 insertions(+), 23 deletions(-) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index c947e2035cf..29ebb988e93 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -1955,12 +1955,21 @@ class EngineArgs: assert not headless or not self.data_parallel_hybrid_lb, ( "data_parallel_hybrid_lb is not applicable in headless mode" ) - assert not (self.data_parallel_hybrid_lb and self.data_parallel_external_lb), ( - "data_parallel_hybrid_lb and data_parallel_external_lb cannot both be True." - ) - assert self.data_parallel_backend == "mp" or self.nnodes == 1, ( - "nnodes > 1 is only supported with data_parallel_backend=mp" - ) + if self.data_parallel_hybrid_lb and self.data_parallel_external_lb: + raise ValueError( + "Invalid data-parallel launch options: " + "`--data-parallel-hybrid-lb` and " + "`--data-parallel-external-lb` cannot be enabled together. " + "Enable only one load-balancing mode." + ) + if self.nnodes > 1 and self.data_parallel_backend != "mp": + raise ValueError( + "Invalid data-parallel launch options: " + f"`--nnodes {self.nnodes}` requires " + "`--data-parallel-backend mp`; got " + f"`--data-parallel-backend {self.data_parallel_backend}`. " + "Use the MP backend or set `--nnodes 1`." + ) inferred_data_parallel_rank = 0 if self.nnodes > 1: world_size = ( @@ -1971,13 +1980,22 @@ class EngineArgs: world_size_within_dp = ( self.pipeline_parallel_size * self.tensor_parallel_size ) + if world_size % self.nnodes != 0: + raise ValueError( + "Invalid data-parallel launch options: " + f"`--nnodes {self.nnodes}` must evenly divide the total " + f"world size ({world_size}). Adjust `--nnodes`, " + "`--data-parallel-size`, `--pipeline-parallel-size`, or " + "`--tensor-parallel-size`." + ) + if not 0 <= self.node_rank < self.nnodes: + raise ValueError( + "Invalid data-parallel launch options: `--node-rank` must " + f"be between 0 and {self.nnodes - 1}; got " + f"`--node-rank {self.node_rank}`. Set it to this node's " + "zero-based index." + ) local_world_size = world_size // self.nnodes - assert world_size % self.nnodes == 0, ( - f"world_size={world_size} must be divisible by nnodes={self.nnodes}." - ) - assert self.node_rank < self.nnodes, ( - f"node_rank={self.node_rank} must be less than nnodes={self.nnodes}." - ) inferred_data_parallel_rank = ( self.node_rank * local_world_size ) // world_size_within_dp @@ -2008,14 +2026,21 @@ class EngineArgs: ) # Local DP rank = 1, use pure-external LB. if data_parallel_external_lb: - assert self.data_parallel_rank is not None, ( - "data_parallel_rank or node_rank must be specified if " - "data_parallel_external_lb is enable." - ) - assert self.data_parallel_size_local in (1, None), ( - "data_parallel_size_local must be 1 or None when data_parallel_rank " - "is set" - ) + if self.data_parallel_rank is None: + raise ValueError( + "Invalid data-parallel launch options: " + "`--data-parallel-external-lb` requires a data-parallel " + "rank. Set `--data-parallel-rank`, or set " + "`--data-parallel-size` greater than 1 and use `--nnodes` " + "with `--node-rank` so the rank can be inferred." + ) + if self.data_parallel_size_local not in (1, None): + raise ValueError( + "Invalid data-parallel launch options: an external " + "data-parallel rank requires `--data-parallel-size-local " + f"1`; got {self.data_parallel_size_local}. Set it to 1 or " + "omit it." + ) data_parallel_size_local = 1 # Use full external lb if we have local_size of 1. self.data_parallel_hybrid_lb = False @@ -2050,9 +2075,13 @@ class EngineArgs: self.node_rank, ) else: - assert not self.data_parallel_hybrid_lb, ( - "data_parallel_size_local must be set to use data_parallel_hybrid_lb." - ) + if self.data_parallel_hybrid_lb: + raise ValueError( + "Invalid data-parallel launch options: " + "`--data-parallel-hybrid-lb` requires " + "`--data-parallel-size-local`. Set it to the number of " + "data-parallel ranks on this node." + ) if self.data_parallel_backend == "ray" and ( envs.VLLM_RAY_DP_PACK_STRATEGY == "span" From d02df748bf9efd99022f1a062597dc3cb3808485 Mon Sep 17 00:00:00 2001 From: Thomas Fahrner Date: Fri, 24 Jul 2026 08:23:08 -0700 Subject: [PATCH 004/185] [Bugfix] Accept RFC 2397 parameters in base64 data URLs (#48973) Signed-off-by: Thomas Fahrner Co-authored-by: Claude Fable 5 --- tests/multimodal/media/test_connector.py | 38 ++++++++++++++++++++++++ vllm/multimodal/media/connector.py | 12 +++++--- 2 files changed, 46 insertions(+), 4 deletions(-) diff --git a/tests/multimodal/media/test_connector.py b/tests/multimodal/media/test_connector.py index 5cd22a8b18d..ddf6dbcef86 100644 --- a/tests/multimodal/media/test_connector.py +++ b/tests/multimodal/media/test_connector.py @@ -225,6 +225,44 @@ async def test_fetch_image_local_files_with_space_in_name(image_url: str): assert not ImageChops.difference(image_sync, image_async).getbbox() +@pytest.mark.asyncio +async def test_fetch_image_data_url_with_params(): + """RFC 2397 allows parameters between the mediatype and the base64 + marker; they must not be rejected or leak into the media type.""" + connector = MediaConnector() + + image = Image.new("RGB", (4, 4), color=(255, 0, 0)) + with NamedTemporaryFile(suffix=".png") as f: + image.save(f.name) + base64_image = base64.b64encode(f.read()).decode("utf-8") + + data_url = f"data:image/png;charset=utf-8;base64,{base64_image}" + image_sync = connector.fetch_image(data_url) + image_async = await connector.fetch_image_async(data_url) + assert _image_equals(image_sync, image_async) + + +def test_fetch_image_data_url_malformed(): + connector = MediaConnector() + + with pytest.raises(ValueError, match="missing ','"): + connector.fetch_image("data:image/png;base64") + + with pytest.raises(NotImplementedError, match="base64"): + connector.fetch_image("data:text/plain,hello") + + # ";base64" requires the ";"; here "base64" is a (bogus) media type. + with pytest.raises(NotImplementedError, match="base64"): + connector.fetch_image("data:base64,aGVsbG8=") + + # Strict RFC 2397 grammar: lowercase "base64", no whitespace. + with pytest.raises(NotImplementedError, match="base64"): + connector.fetch_image("data:image/png;BASE64,aGVsbG8=") + + with pytest.raises(NotImplementedError, match="base64"): + connector.fetch_image("data:image/png; base64,aGVsbG8=") + + @pytest.mark.asyncio async def test_fetch_image_error_conversion(): connector = MediaConnector() diff --git a/vllm/multimodal/media/connector.py b/vllm/multimodal/media/connector.py index bf9b7345ca0..656f537c5b1 100644 --- a/vllm/multimodal/media/connector.py +++ b/vllm/multimodal/media/connector.py @@ -302,14 +302,18 @@ class MediaConnector: media_io: MediaIO[_M], ) -> _M: # type: ignore[type-var] # Format per RFC 2397: - # data:[][;base64], - data_spec, data = url[5:].split(",", 1) - media_type, data_type = data_spec.split(";", 1) + # data:[][;=]*[;base64], + data_spec, sep, data = url[5:].partition(",") + if not sep: + msg = f"Invalid data URL {url[:32]!r}: missing ',' separator." + raise ValueError(msg) - if data_type != "base64": + media_type, sep, encoding = data_spec.rpartition(";") + if not sep or encoding != "base64": msg = "Only base64 data URLs are supported for now." raise NotImplementedError(msg) + media_type = media_type.partition(";")[0] return media_io.load_base64(media_type, data) def _load_file_url( From 866fea2b9900bf49d552c205d2eaac4716fb63ac Mon Sep 17 00:00:00 2001 From: Johnny-Liou <75152852+Johnny-Liou@users.noreply.github.com> Date: Fri, 24 Jul 2026 12:39:49 -0400 Subject: [PATCH 005/185] [Kernel] ReplaySSM: cache SSM inputs for faster Mamba2 standard decode (#48018) Signed-off-by: Johnny-Liou Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: tomeras91 <57313761+tomeras91@users.noreply.github.com> Co-authored-by: Cyrus Leung --- benchmarks/replayssm/e2e_decode_speedup.py | 267 +++++++ ...ayssm_prefill_decode_equivalence_mamba2.py | 244 ++++++ .../test_replayssm_standard_decode_mamba2.py | 673 +++++++++++++++++ tests/kernels/mamba/utils.py | 160 ++++ tests/models/test_registry.py | 16 + .../test_mamba_update_block_table.py | 2 + .../test_replayssm_metadata_builder.py | 256 +++++++ tests/v1/e2e/test_replayssm_decode.py | 138 ++++ vllm/config/cache.py | 11 + vllm/config/model.py | 4 + vllm/config/vllm.py | 33 +- vllm/engine/arg_utils.py | 8 + .../layers/mamba/mamba_mixer2.py | 120 ++- .../layers/mamba/mamba_utils.py | 34 + .../layers/mamba/ops/replayssm_config.py | 66 ++ ...tive_state_update_replayssm_output_only.py | 709 ++++++++++++++++++ vllm/model_executor/models/interfaces.py | 25 + vllm/model_executor/models/nemotron_h.py | 31 +- vllm/model_executor/models/registry.py | 3 + vllm/v1/attention/backend.py | 6 + vllm/v1/attention/backends/mamba_attn.py | 144 ++++ vllm/v1/worker/gpu_input_batch.py | 26 + vllm/v1/worker/gpu_model_runner.py | 9 + 23 files changed, 2954 insertions(+), 31 deletions(-) create mode 100644 benchmarks/replayssm/e2e_decode_speedup.py create mode 100644 tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py create mode 100644 tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py create mode 100644 tests/v1/attention/test_replayssm_metadata_builder.py create mode 100644 tests/v1/e2e/test_replayssm_decode.py create mode 100644 vllm/model_executor/layers/mamba/ops/replayssm_config.py create mode 100644 vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py diff --git a/benchmarks/replayssm/e2e_decode_speedup.py b/benchmarks/replayssm/e2e_decode_speedup.py new file mode 100644 index 00000000000..b7084a21315 --- /dev/null +++ b/benchmarks/replayssm/e2e_decode_speedup.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end autoregressive decode benchmark: ReplaySSM vs the standard SSM kernel. + +Loads a hybrid Mamba2 model, replicates one prompt across the batch, and times a +long greedy decode (CUDA graphs on) once with the standard kernel and once with +ReplaySSM, then reports the per-step / throughput speedup. The two modes run in +separate subprocesses so each gets a clean CUDA context. + +The FlashInfer FP4-MoE autotuner is disabled by default (it is unstable under +CUDA-graph capture on the pre-release Blackwell FP4 path); pass +--no-disable-flashinfer-autotune for non-FP4 models. + +Examples: + python e2e_decode_speedup.py --model-id nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 + python e2e_decode_speedup.py --dtype auto --buffer-len 16 \ + --model-id nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 # B300 NVFP4 +""" + +import argparse +import json +import os +import subprocess +import sys +import time + +DEFAULT_PROMPT = "My cat wrote all this CUDA code for a new language model and" + +MODE_LABEL = {"standard": "standard", "replayssm": "ReplaySSM"} + + +def parse_args(): + p = argparse.ArgumentParser( + description="E2E decode speedup: ReplaySSM vs the standard SSM kernel." + ) + p.add_argument("--model-id", default="nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16") + p.add_argument("--prompt", default=DEFAULT_PROMPT) + p.add_argument("--batch-size", type=int, default=256) + p.add_argument("--num-steps", type=int, default=1000) + p.add_argument("--warmup-steps", type=int, default=128) + p.add_argument("--repeats", type=int, default=1) + p.add_argument( + "--buffer-len", type=int, default=16, help="ReplaySSM input-buffer length." + ) + p.add_argument( + "--dtype", + default="bfloat16", + choices=["bfloat16", "float16", "float32", "auto"], + ) + p.add_argument("--gpu-memory-utilization", type=float, default=0.9) + p.add_argument("--max-model-len", type=int, default=None) + p.add_argument( + "--disable-flashinfer-autotune", + action=argparse.BooleanOptionalAction, + default=True, + help="Disable the FlashInfer FP4-MoE autotuner (default: on). " + "It is unstable under CUDA-graph capture on the " + "pre-release Blackwell FP4 path; pass " + "--no-disable-flashinfer-autotune for non-FP4 models.", + ) + p.add_argument( + "--mamba-ssm-cache-dtype", + default="auto", + choices=["auto", "float32", "float16", "bfloat16"], + help="SSM state dtype (both modes). 'auto' = config-driven; " + "'float32' = fp32 state, 'bfloat16' = s16 state.", + ) + p.add_argument( + "--baseline-ssm-config", + default="", + help="Pin the STANDARD baseline's SSM launch config as " + "'bsm,nw' via override_ssm_config (forces the in-process " + "engine so the override reaches the kernel). Empty = off.", + ) + p.add_argument( + "--worker", + choices=["standard", "replayssm"], + default=None, + help=argparse.SUPPRESS, + ) + return p.parse_args() + + +def resolve_max_model_len(args) -> int: + if args.max_model_len is not None: + return args.max_model_len + return args.num_steps + 256 + + +def run_worker(args): + # override_ssm_config is a module global; it only reaches the model if the + # engine runs in-process (default V1 spawns a separate EngineCore). Force it. + if args.worker == "standard" and args.baseline_ssm_config: + os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" + + import torch + + from vllm import LLM, SamplingParams + + mode = args.worker + max_model_len = resolve_max_model_len(args) + + llm_kwargs = dict( + model=args.model_id, + tensor_parallel_size=1, + dtype=args.dtype, + max_model_len=max_model_len, + trust_remote_code=True, + enable_prefix_caching=False, + enable_chunked_prefill=False, + max_num_seqs=args.batch_size, + max_num_batched_tokens=max(max_model_len, args.batch_size * 64), + enforce_eager=False, + disable_log_stats=True, + gpu_memory_utilization=args.gpu_memory_utilization, + # SSM state dtype (applies to both standard and ReplaySSM). + mamba_ssm_cache_dtype=args.mamba_ssm_cache_dtype, + ) + if args.disable_flashinfer_autotune: + # FP4-MoE autotuner is unstable under CUDA-graph capture on Blackwell; + # re-enable (--no-disable-flashinfer-autotune) only for non-FP4 models. + llm_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} + if mode == "replayssm": + llm_kwargs.update(use_replayssm=True, replayssm_buffer_len=args.buffer_len) + + _ssm_cm = None + if mode == "standard" and args.baseline_ssm_config: + from vllm.model_executor.layers.mamba.ops.mamba_ssm import override_ssm_config + + _bsm, _nw = (int(x) for x in args.baseline_ssm_config.split(",")) + _ssm_cm = override_ssm_config((_bsm, _nw)) + _ssm_cm.__enter__() # active through LLM() graph capture + decode + print( + f"[{mode}] override_ssm_config -> (BLOCK_SIZE_M={_bsm}, num_warps={_nw})", + flush=True, + ) + + llm = LLM(**llm_kwargs) + prompts = [args.prompt] * args.batch_size + + def timed_generate(n_tokens): + sp = SamplingParams( + n=1, + temperature=0.0, + ignore_eos=True, + min_tokens=n_tokens, + max_tokens=n_tokens, + ) + if torch.accelerator.is_available(): + torch.accelerator.synchronize() + t0 = time.perf_counter() + outs = llm.generate(prompts, sp, use_tqdm=False) + if torch.accelerator.is_available(): + torch.accelerator.synchronize() + elapsed = time.perf_counter() - t0 + produced = min(len(o.outputs[0].token_ids) for o in outs) + assert produced == n_tokens, f"expected {n_tokens} tokens, got {produced}" + return elapsed + + timed_generate(args.warmup_steps) + + best = None + for _ in range(args.repeats): + elapsed = timed_generate(args.num_steps) + tok_s = args.batch_size * args.num_steps / elapsed + per_step_ms = elapsed / args.num_steps * 1e3 + print( + f"[{mode}] {elapsed:.3f}s {tok_s:,.0f} tok/s {per_step_ms:.3f} ms/step", + flush=True, + ) + if best is None or elapsed < best["elapsed_s"]: + best = { + "mode": mode, + "elapsed_s": elapsed, + "tok_s": tok_s, + "per_step_ms": per_step_ms, + } + + print("RESULT_JSON " + json.dumps(best), flush=True) + if _ssm_cm is not None: + _ssm_cm.__exit__(None, None, None) + + +def run_one_mode(args, mode) -> dict: + cmd = [ + sys.executable, + __file__, + "--worker", + mode, + "--model-id", + args.model_id, + "--prompt", + args.prompt, + "--batch-size", + str(args.batch_size), + "--num-steps", + str(args.num_steps), + "--warmup-steps", + str(args.warmup_steps), + "--repeats", + str(args.repeats), + "--buffer-len", + str(args.buffer_len), + "--dtype", + args.dtype, + "--gpu-memory-utilization", + str(args.gpu_memory_utilization), + "--mamba-ssm-cache-dtype", + args.mamba_ssm_cache_dtype, + "--baseline-ssm-config", + args.baseline_ssm_config, + ] + cmd.append( + "--disable-flashinfer-autotune" + if args.disable_flashinfer_autotune + else "--no-disable-flashinfer-autotune" + ) + if args.max_model_len is not None: + cmd += ["--max-model-len", str(args.max_model_len)] + + result = None + proc = subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1 + ) + for line in proc.stdout: + sys.stdout.write(line) + sys.stdout.flush() + if line.startswith("RESULT_JSON "): + result = json.loads(line[len("RESULT_JSON ") :]) + proc.wait() + if proc.returncode != 0: + raise RuntimeError(f"mode '{mode}' worker exited with {proc.returncode}") + if result is None: + raise RuntimeError(f"mode '{mode}' produced no RESULT_JSON line") + return result + + +def main(): + args = parse_args() + if args.worker is not None: + run_worker(args) + return + + print( + f"model={args.model_id} batch_size={args.batch_size} " + f"steps={args.num_steps} buffer_len={args.buffer_len} dtype={args.dtype}" + ) + + std = run_one_mode(args, "standard") + fla = run_one_mode(args, "replayssm") + speedup = std["per_step_ms"] / fla["per_step_ms"] + + print() + header = f"{'mode':<10}{'ms/step':>12}{'tok/s':>16}{'wall (s)':>12}" + print(header) + print("-" * len(header)) + for r in (std, fla): + print( + f"{MODE_LABEL[r['mode']]:<10}{r['per_step_ms']:>12.3f}" + f"{r['tok_s']:>16,.0f}{r['elapsed_s']:>12.3f}" + ) + print("-" * len(header)) + print(f"speedup (standard / ReplaySSM, per step): {speedup:.3f}x") + + +if __name__ == "__main__": + main() diff --git a/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py b/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py new file mode 100644 index 00000000000..5319d8964d3 --- /dev/null +++ b/tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py @@ -0,0 +1,244 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Prefill == decode equivalence for the Mamba2 ReplaySSM kernels. + +The SSM recurrence is path-independent: vLLM's chunked prefill (the SSD kernel +``mamba_chunk_scan_combined_varlen``) and step-by-step decode must produce the +same per-position outputs and final state over a sequence. This file feeds one +set of inputs through the production dt flow (raw dt + softplus + a per-head +dt_bias, applied inside each kernel) to: + + * the exact fp32 step recurrence (``selective_state_update_ref``) -- the + ground truth, + * the chunked prefill kernel, + * the baseline decode kernel, + * the ReplaySSM output_only decode kernel, + +and checks all of them agree. Prefill (a chunked scan) and decode (a step +recurrence) are different code paths, so they differ numerically: the chunked +scan carries ~2e-2 (fp32) / ~4e-2 (bf16) vs the exact recurrence, far above the +near-exact decode. We therefore anchor every path on the exact recurrence at +SSD-level tolerances (the same regime as ``test_mamba_ssm_ssd.py``), which the +chunked scan sets, keyed off the activation dtype. + +State and activation/buffer precision are swept independently, including the +fp32-state + bf16-activation production config. +""" + +import pytest +import torch + +from tests.kernels.mamba.utils import selective_state_update_ref +from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_state_update +from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_output_only import ( # noqa: E501 + selective_state_update_replayssm_output_only, +) +from vllm.model_executor.layers.mamba.ops.ssd_combined import ( + mamba_chunk_scan_combined_varlen, +) +from vllm.utils.torch_utils import set_random_seed +from vllm.v1.attention.backends.mamba2_attn import compute_varlen_chunk_metadata + + +def _prefill_tolerances(act_dtype: torch.dtype) -> tuple[float, float]: + # The chunked prefill scan, not the decode, sets these: it carries ~2e-2 + # (fp32) / ~6e-2 (bf16) vs the exact recurrence, while ReplaySSM decode is + # near-exact (~2e-6 fp32). Keyed off the activation dtype (the outputs are + # in act_dtype). Same regime as test_mamba_ssm_ssd.py. + if act_dtype == torch.float32: + return 1e-2, 3e-2 + return 6e-2, 1e-1 + + +def _run_prefill_decode_equivalence( + *, + state_dtype: torch.dtype, + act_dtype: torch.dtype, + nheads: int, + headdim: int, + ngroups: int, + dstate: int, + seqlen: int, + chunk_size: int, + max_cache_len: int, + seed: int = 0, +) -> None: + """Prefill the whole sequence and decode it step by step; check both match + the exact fp32 recurrence (and each other). All paths use the production dt + flow (raw dt + softplus + a per-head dt_bias), so this also checks prefill + and decode apply the softplus/bias preprocessing consistently. ``state_dtype`` + is the recurrent-state precision; ``act_dtype`` the activation/buffer one.""" + device = "cuda" + rtol, atol = _prefill_tolerances(act_dtype) + set_random_seed(seed) + + # Production dt flow: raw dt + a per-head dt_bias (~-4 keeps softplus(dt + + # bias) small and well-conditioned). dt_bias is (nheads,) for prefill and + # (nheads, headdim) for the decode kernels/reference. + A = -torch.exp(torch.rand(nheads, device=device, dtype=act_dtype)) + dt = torch.randn(seqlen, nheads, device=device, dtype=act_dtype) + dt_bias = torch.rand(nheads, device=device, dtype=act_dtype) - 4 + dt_bias_hd = dt_bias.view(nheads, 1).expand(nheads, headdim) + X = torch.randn(seqlen, nheads, headdim, device=device, dtype=act_dtype) + B = torch.randn(seqlen, ngroups, dstate, device=device, dtype=act_dtype) + C = torch.randn(seqlen, ngroups, dstate, device=device, dtype=act_dtype) + A_bcast = A.view(nheads, 1, 1).expand(nheads, headdim, dstate) + + # Chunked prefill over the whole sequence (implicit batch=1, varlen). The + # kernel always returns the final state in fp32, so no state_dtype plumbing. + cu_seqlens = torch.tensor((0, seqlen), device=device).cumsum(0).to(torch.int32) + cu_chunk_seqlens, last_chunk_indices, seq_idx = compute_varlen_chunk_metadata( + cu_seqlens, chunk_size + ) + y_prefill = torch.empty(seqlen, nheads, headdim, device=device, dtype=act_dtype) + final_state_prefill = mamba_chunk_scan_combined_varlen( + X, + dt, + A, + B, + C, + chunk_size, + cu_seqlens=cu_seqlens, + cu_chunk_seqlens=cu_chunk_seqlens, + last_chunk_indices=last_chunk_indices, + seq_idx=seq_idx, + out=y_prefill, + D=None, + dt_bias=dt_bias, + dt_softplus=True, + ) + + # Step paths: exact fp32 recurrence (ground truth), baseline, ReplaySSM. + # State follows state_dtype; caches follow act_dtype (dt_cache is fp32). + state_ref = torch.zeros( + 1, nheads, headdim, dstate, device=device, dtype=torch.float32 + ) + state_base = torch.zeros( + 1, nheads, headdim, dstate, device=device, dtype=state_dtype + ) + state_dec = torch.zeros( + 1, nheads, headdim, dstate, device=device, dtype=state_dtype + ) + x_cache = torch.zeros( + 1, nheads, max_cache_len, headdim, device=device, dtype=act_dtype + ) + dt_cache = torch.zeros(1, nheads, max_cache_len, device=device, dtype=torch.float32) + B_cache = torch.zeros( + 1, ngroups, max_cache_len, dstate, device=device, dtype=act_dtype + ) + bc_pre = torch.empty(1, ngroups, max_cache_len, device=device, dtype=torch.float32) + write_pos = torch.zeros(1, dtype=torch.int32, device=device) + # No skip connection (D=0) on any path here; the D!=0 path is covered by the + # standard-decode suite. The baseline kernel needs a D tensor, not None. + D_zero = torch.zeros(nheads, headdim, device=device) + + y_ref = torch.empty(seqlen, nheads, headdim, device=device, dtype=torch.float32) + y_base = torch.empty(seqlen, nheads, headdim, device=device, dtype=act_dtype) + y_dec = torch.empty(seqlen, nheads, headdim, device=device, dtype=act_dtype) + for t in range(seqlen): + dt_t = dt[t].view(1, nheads, 1).expand(1, nheads, headdim) + is_flush = write_pos == max_cache_len - 1 + + y_ref[t] = selective_state_update_ref( + state_ref, + X[t : t + 1].float(), + dt_t.float(), + A_bcast.float(), + B[t : t + 1].float(), + C[t : t + 1].float(), + dt_bias=dt_bias_hd.float(), + dt_softplus=True, + )[0] + + out_b = torch.empty(1, nheads, headdim, device=device, dtype=act_dtype) + selective_state_update( + state_base, + X[t : t + 1], + dt_t, + A_bcast, + B[t : t + 1], + C[t : t + 1], + D=D_zero, + dt_bias=dt_bias_hd, + dt_softplus=True, + out=out_b, + ) + y_base[t] = out_b[0] + + out_d = torch.empty(1, nheads, headdim, device=device, dtype=act_dtype) + common = dict( + dt_bias=dt_bias_hd, + dt_softplus=True, + x_cache=x_cache, + dt_cache=dt_cache, + B_cache=B_cache, + write_pos=write_pos, + is_flush=is_flush, + max_cache_len=max_cache_len, + out=out_d, + ) + selective_state_update_replayssm_output_only( + state_dec, + X[t : t + 1], + dt_t, + A_bcast, + B[t : t + 1], + C[t : t + 1], + bc_pre=bc_pre, + **common, + ) + y_dec[t] = out_d[0] + + write_pos = torch.where(is_flush, torch.zeros_like(write_pos), write_pos + 1) + + # Every path computes the same recurrence; anchor each on the fp32 truth. + torch.testing.assert_close(y_prefill.float(), y_ref, rtol=rtol, atol=atol) + torch.testing.assert_close(y_base.float(), y_ref, rtol=rtol, atol=atol) + torch.testing.assert_close(y_dec.float(), y_ref, rtol=rtol, atol=atol) + # Headline: ReplaySSM decode matches the chunked prefill directly. + torch.testing.assert_close(y_dec.float(), y_prefill.float(), rtol=rtol, atol=atol) + # Final state too (the recurrence ends in state_ref after the loop). + torch.testing.assert_close( + final_state_prefill[0].float(), state_ref[0], rtol=rtol, atol=atol + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize( + "precision", + # fp32 state is the default; bf16/fp16 are reduced-footprint configs. fp16 + # appears as an activation dtype (s32_afp16, sfp16_afp16) and as a finer- + # mantissa state under bf16 activations (sfp16_a16). + [ + pytest.param((torch.float32, torch.float32), id="s32_a32"), + pytest.param((torch.float32, torch.bfloat16), id="s32_a16"), + pytest.param((torch.bfloat16, torch.bfloat16), id="s16_a16"), + pytest.param((torch.float32, torch.float16), id="s32_afp16"), + pytest.param((torch.float16, torch.float16), id="sfp16_afp16"), + pytest.param((torch.float16, torch.bfloat16), id="sfp16_a16"), + ], +) +@pytest.mark.parametrize( + "geometry", # (nheads, headdim, dstate, ngroups) + [ + pytest.param((8, 64, 64, 2), id="small"), + pytest.param((96, 80, 128, 8), id="nano4b"), + ], +) +def test_replayssm_prefill_decode_equivalence( + precision: tuple[torch.dtype, torch.dtype], + geometry: tuple[int, int, int, int], +): + state_dtype, act_dtype = precision + nheads, headdim, dstate, ngroups = geometry + _run_prefill_decode_equivalence( + state_dtype=state_dtype, + act_dtype=act_dtype, + nheads=nheads, + headdim=headdim, + ngroups=ngroups, + dstate=dstate, + seqlen=16, + chunk_size=8, + max_cache_len=4, + ) diff --git a/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py b/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py new file mode 100644 index 00000000000..588ce52dbeb --- /dev/null +++ b/tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py @@ -0,0 +1,673 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Standard (autoregressive) decode correctness for the Mamba2 ReplaySSM kernels. + +ReplaySSM caches the recent SSM inputs ``(x, dt, B)`` in a small ring buffer and +reconstructs / reads out the recurrent state on the fly, writing the full state +back to HBM only when the buffer flushes. This file checks that, over a +multi-step decode, the ReplaySSM output_only kernel reproduces the exact SSM +recurrence. + +For every step we assert, against trusted oracles driven one token at a time: + + * the ReplaySSM output matches its pure-PyTorch reference, which models the + kernel's exact arithmetic (including its bf16 reconstruction), at every + precision, + * when the state is fp32, the output also matches the baseline decode kernel + (``selective_state_update``); at bf16 state the baseline legitimately + differs -- it downcasts the fp32 state to bf16 every step, while ReplaySSM + accumulates a whole buffer in fp32 and is the more accurate path, + * when state and activations are both fp32, the output also matches the exact + elementwise ``selective_state_update_ref``, + * the cached inputs match the reference cache management, + * the checkpoint state matches the reference (and the baseline at fp32 state). + +State and activation/buffer precision are swept independently. Nemotron-3 +defaults to fp32 SSM state (``mamba_ssm_cache_dtype=float32``), but bf16 state is +also supported; the buffer dtype follows the activation dtype and ``dt_cache`` +is always fp32. +""" + +import pytest +import torch + +from tests.kernels.mamba.utils import ( + allocate_update_caches, + selective_state_update_ref, + selective_state_update_replayssm_output_only_ref, +) +from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_state_update +from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_output_only import ( # noqa: E501 + selective_state_update_replayssm_output_only, +) +from vllm.utils.torch_utils import set_random_seed +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + + +def _tolerances(dtype: torch.dtype) -> tuple[float, float]: + # fp32 demands true fp32 parity. A correct fp32 reconstruction (tl.dot with + # input_precision="tf32x3"/"ieee") matches the elementwise baseline to + # ~1e-5, while a TF32 reconstruction drifts to ~1e-2. atol=1e-3 sits between, + # so it flags TF32 degradation yet passes a correct fp32 kernel. bf16 stays + # loose (bf16 rounding dominates and is unaffected by the matmul precision). + if dtype == torch.float32: + return 1e-4, 1e-3 + return 6e-2, 2e-1 + + +def _tied_A(nheads: int, headdim: int, dstate: int, device: str) -> torch.Tensor: + A = -torch.rand(nheads, device=device) - 1.0 + return A.view(nheads, 1, 1).expand(nheads, headdim, dstate) + + +def _tied_dt( + batch: int, + nheads: int, + headdim: int, + device: str, + dtype: torch.dtype, +) -> torch.Tensor: + dt = torch.randn(batch, nheads, device=device, dtype=dtype) + return dt.unsqueeze(-1).expand(batch, nheads, headdim) + + +def _tied_dt_bias(nheads: int, headdim: int, device: str) -> torch.Tensor: + dt_bias = torch.rand(nheads, device=device) - 4.0 + return dt_bias.view(nheads, 1).expand(nheads, headdim) + + +def _run_standard_decode( + *, + state_dtype: torch.dtype, + act_dtype: torch.dtype, + batch: int, + nheads: int, + headdim: int, + ngroups: int, + dstate: int, + max_cache_len: int, + num_steps: int, + has_z: bool, + dt_softplus: bool, + use_dt_bias: bool, + desync_write_pos: bool = False, + seed: int = 0, +) -> None: + """Drive ``num_steps`` decode steps and check every path against the + trusted recurrence. ``state_dtype`` is the recurrent-state precision; + ``act_dtype`` is the activation/buffer precision.""" + device = "cuda" + both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32 + rtol, atol = _tolerances(torch.float32 if both_fp32 else torch.bfloat16) + set_random_seed(seed) + + # One state copy per path; all start identical at the same precision. + state0 = torch.randn( + batch, nheads, headdim, dstate, dtype=state_dtype, device=device + ) + state_anchor = state0.clone() + state_baseline = state0.clone() + state_cached = state0.clone() + state_ref = state0.clone() + + A = _tied_A(nheads, headdim, dstate, device) + dt_bias = _tied_dt_bias(nheads, headdim, device) if use_dt_bias else None + D = torch.randn(nheads, headdim, device=device) + + # Caches follow the activation dtype (dt_cache is forced to fp32 inside). + x_cache, dt_cache, B_cache, _ = allocate_update_caches( + batch, + nheads, + ngroups, + headdim, + dstate, + max_cache_len, + device, + act_dtype, + act_dtype, + ) + x_cache_ref, dt_cache_ref, B_cache_ref, _ = allocate_update_caches( + batch, + nheads, + ngroups, + headdim, + dstate, + max_cache_len, + device, + act_dtype, + act_dtype, + ) + bc_pre = torch.empty( + batch, ngroups, max_cache_len, device=device, dtype=torch.float32 + ) + + if desync_write_pos: + # Rows start at different ring positions so they flush on different + # steps, exercising per-row write-position handling. + write_pos = ( + torch.arange(batch, device=device, dtype=torch.int32) % max_cache_len + ) + else: + write_pos = torch.zeros(batch, dtype=torch.int32, device=device) + + for _ in range(num_steps): + x = torch.randn(batch, nheads, headdim, device=device, dtype=act_dtype) + dt = _tied_dt(batch, nheads, headdim, device, act_dtype) + B = torch.randn(batch, ngroups, dstate, device=device, dtype=act_dtype) + C = torch.randn(batch, ngroups, dstate, device=device, dtype=act_dtype) + z = torch.randn_like(x) if has_z else None + is_flush = write_pos == max_cache_len - 1 + + # Trusted recurrence (mutates state_anchor in place, returns output). + out_anchor = selective_state_update_ref( + state_anchor, + x, + dt, + A, + B, + C, + D=D, + z=z, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + ) + + # Upstream baseline decode kernel. + out_baseline = torch.empty_like(x) + selective_state_update( + state_baseline, + x, + dt, + A, + B, + C, + D=D, + z=z, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + out=out_baseline, + ) + + # ReplaySSM kernel under test + its pure-PyTorch reference. + out_cached = torch.empty_like(x) + common = dict( + D=D, + z=z, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + x_cache=x_cache, + dt_cache=dt_cache, + B_cache=B_cache, + write_pos=write_pos, + is_flush=is_flush, + max_cache_len=max_cache_len, + out=out_cached, + ) + selective_state_update_replayssm_output_only( + state_cached, x, dt, A, B, C, bc_pre=bc_pre, **common + ) + out_ref = selective_state_update_replayssm_output_only_ref( + state_ref, + x, + dt, + A, + B, + C, + D=D, + z=z, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + x_cache=x_cache_ref, + dt_cache=dt_cache_ref, + B_cache=B_cache_ref, + write_pos=write_pos, + max_cache_len=max_cache_len, + ) + + # The reference models the kernel's exact arithmetic (including its bf16 + # reconstruction), so the kernel must match it tightly at every + # precision. At fp32 this also flags any TF32 reconstruction drift. + torch.testing.assert_close(out_cached, out_ref, rtol=rtol, atol=atol) + # When the STATE is fp32 the baseline decode kernel is a valid oracle (it + # does not downcast the state per step), so ReplaySSM must match it. At + # bf16 state the baseline legitimately differs: it downcasts the fp32 + # state to bf16 every step, while ReplaySSM accumulates a whole buffer in + # fp32 and is the MORE accurate path -- so it is not a tight oracle there. + if state_dtype == torch.float32: + torch.testing.assert_close(out_cached, out_baseline, rtol=rtol, atol=atol) + # The exact elementwise reference is valid only when state AND + # activations are fp32; otherwise it downcasts the state (at readout for + # bf16 activations, or per step for bf16 state). + if both_fp32: + torch.testing.assert_close(out_cached, out_anchor, rtol=rtol, atol=atol) + + # Cached inputs match the reference cache management. + torch.testing.assert_close(x_cache, x_cache_ref, rtol=rtol, atol=atol) + torch.testing.assert_close(dt_cache, dt_cache_ref, rtol=rtol, atol=atol) + torch.testing.assert_close(B_cache, B_cache_ref, rtol=rtol, atol=atol) + + # Checkpoint state at flush matches the reference (and, when the state is + # fp32, the baseline kernel). + if bool(is_flush.any()): + torch.testing.assert_close( + state_cached[is_flush], state_ref[is_flush], rtol=rtol, atol=atol + ) + if state_dtype == torch.float32: + torch.testing.assert_close( + state_cached[is_flush], + state_baseline[is_flush], + rtol=rtol, + atol=atol, + ) + + write_pos = torch.where(is_flush, torch.zeros_like(write_pos), write_pos + 1) + + +# State/activation precisions. fp32 state is the default; bf16/fp16 are the +# reduced-footprint configs. fp16 appears both as an activation dtype (fully-fp16 +# model sfp16_afp16, or fp16 act over fp32 state s32_afp16) and as a state dtype +# under bf16 activations (sfp16_a16): fp16 has a finer mantissa than bf16 at the +# same 2 bytes, so it is a more accurate state at no extra footprint. We still +# skip fp16 state under fp32 activations (the unused low-state/high-act mix). +_PRECISIONS = [ + pytest.param((torch.float32, torch.float32), id="s32_a32"), + pytest.param((torch.float32, torch.bfloat16), id="s32_a16"), + pytest.param((torch.bfloat16, torch.bfloat16), id="s16_a16"), + pytest.param((torch.float32, torch.float16), id="s32_afp16"), + pytest.param((torch.float16, torch.float16), id="sfp16_afp16"), + pytest.param((torch.float16, torch.bfloat16), id="sfp16_a16"), +] +# Small synthetic shapes for the full axis sweep (compile fast). +_SMALL_GEOMETRIES = [ + pytest.param((8, 64, 64, 4), id="small"), + pytest.param((4, 64, 16, 1), id="tiny"), +] +# Production Mamba2 shapes (nheads, headdim, dstate, ngroups), TP=1. +_REAL_GEOMETRIES = [ + pytest.param((96, 80, 128, 8), id="nano4b"), + pytest.param((128, 64, 128, 8), id="super120b"), + pytest.param((256, 64, 128, 8), id="ultra550b"), +] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("precision", _PRECISIONS) +@pytest.mark.parametrize("max_cache_len", [1, 4, 16]) +@pytest.mark.parametrize("geometry", _SMALL_GEOMETRIES) +@pytest.mark.parametrize("has_z", [False, True]) +def test_replayssm_standard_decode_matches_reference( + precision: tuple[torch.dtype, torch.dtype], + max_cache_len: int, + geometry: tuple[int, int, int, int], + has_z: bool, +): + state_dtype, act_dtype = precision + nheads, headdim, dstate, ngroups = geometry + _run_standard_decode( + state_dtype=state_dtype, + act_dtype=act_dtype, + batch=4, + nheads=nheads, + headdim=headdim, + ngroups=ngroups, + dstate=dstate, + max_cache_len=max_cache_len, + num_steps=2 * max_cache_len + 1, + has_z=has_z, + dt_softplus=True, + use_dt_bias=True, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("precision", _PRECISIONS) +@pytest.mark.parametrize("geometry", _REAL_GEOMETRIES) +def test_replayssm_standard_decode_real_geometry( + precision: tuple[torch.dtype, torch.dtype], + geometry: tuple[int, int, int, int], +): + # Production Mamba2 shapes for the Nemotron-3 family (Nano-4B / Super-120B / + # Ultra-550B), at the production buffer length (8). All three precisions, + # including the bf16 state case. + state_dtype, act_dtype = precision + nheads, headdim, dstate, ngroups = geometry + _run_standard_decode( + state_dtype=state_dtype, + act_dtype=act_dtype, + batch=4, + nheads=nheads, + headdim=headdim, + ngroups=ngroups, + dstate=dstate, + max_cache_len=8, + num_steps=17, + has_z=True, + dt_softplus=True, + use_dt_bias=True, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize( + "precision", + [ + pytest.param((torch.float32, torch.float32), id="s32_a32"), + pytest.param((torch.float32, torch.bfloat16), id="s32_a16"), + pytest.param((torch.float32, torch.float16), id="s32_afp16"), + ], +) +def test_replayssm_standard_decode_desync_write_pos( + precision: tuple[torch.dtype, torch.dtype], +): + # Rows start at staggered ring positions, so they flush on different steps + # and hold genuinely different cached histories. + state_dtype, act_dtype = precision + _run_standard_decode( + state_dtype=state_dtype, + act_dtype=act_dtype, + batch=4, + nheads=8, + headdim=64, + ngroups=4, + dstate=64, + max_cache_len=4, + num_steps=12, + has_z=True, + dt_softplus=True, + use_dt_bias=True, + desync_write_pos=True, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("precision", _PRECISIONS) +@pytest.mark.parametrize("with_padding", [False, True]) +def test_replayssm_standard_decode_with_batch_indices( + precision: tuple[torch.dtype, torch.dtype], + with_padding: bool, +): + # Sparse state allocation via state_batch_indices, with NULL_BLOCK_ID + # padding rows. The pure-PyTorch references do not model the sparse + # gather, so the anchor here is the upstream baseline decode kernel. + state_dtype, act_dtype = precision + device = "cuda" + both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32 + rtol, atol = _tolerances(torch.float32 if both_fp32 else torch.bfloat16) + set_random_seed(0) + + batch = 3 + padding = 2 if with_padding else 0 + padded_batch = batch + padding + total_state_slots = 16 + nheads = 4 + ngroups = 2 + headdim = 64 + dstate = 16 + max_cache_len = 4 + num_steps = 2 * max_cache_len + + state = torch.randn( + total_state_slots, nheads, headdim, dstate, dtype=state_dtype, device=device + ) + state_baseline = state.clone() + state_cached = state.clone() + state_before = state.clone() + + state_indices = ( + torch.randperm(total_state_slots - 1, device=device)[:batch] + 1 + ).to(torch.int32) + state_batch_indices = torch.cat( + [ + state_indices, + torch.full((padding,), NULL_BLOCK_ID, dtype=torch.int32, device=device), + ] + ) + unused_states = torch.ones(total_state_slots, dtype=torch.bool, device=device) + unused_states[state_indices] = False + + A = _tied_A(nheads, headdim, dstate, device) + dt_bias = _tied_dt_bias(nheads, headdim, device) + D = torch.randn(nheads, headdim, device=device) + x_cache = torch.zeros( + total_state_slots, + nheads, + max_cache_len, + headdim, + device=device, + dtype=act_dtype, + ) + dt_cache = torch.zeros( + total_state_slots, nheads, max_cache_len, device=device, dtype=torch.float32 + ) + B_cache = torch.zeros( + total_state_slots, + ngroups, + max_cache_len, + dstate, + device=device, + dtype=act_dtype, + ) + bc_pre = torch.empty( + padded_batch, ngroups, max_cache_len, device=device, dtype=torch.float32 + ) + write_pos = torch.zeros(padded_batch, dtype=torch.int32, device=device) + + for _ in range(num_steps): + x = torch.randn(padded_batch, nheads, headdim, device=device, dtype=act_dtype) + dt = _tied_dt(padded_batch, nheads, headdim, device, act_dtype) + B = torch.randn(padded_batch, ngroups, dstate, device=device, dtype=act_dtype) + C = torch.randn(padded_batch, ngroups, dstate, device=device, dtype=act_dtype) + z = torch.randn_like(x) + is_flush = write_pos == max_cache_len - 1 + + out_baseline = torch.empty_like(x) + selective_state_update( + state_baseline, + x, + dt, + A, + B, + C, + D=D, + z=z, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=state_batch_indices, + out=out_baseline, + ) + + out_cached = torch.full_like(x, 42) + common = dict( + D=D, + z=z, + dt_bias=dt_bias, + dt_softplus=True, + x_cache=x_cache, + dt_cache=dt_cache, + B_cache=B_cache, + write_pos=write_pos, + is_flush=is_flush, + max_cache_len=max_cache_len, + state_batch_indices=state_batch_indices, + out=out_cached, + ) + selective_state_update_replayssm_output_only( + state_cached, x, dt, A, B, C, bc_pre=bc_pre, **common + ) + + torch.testing.assert_close( + out_cached[:batch], out_baseline[:batch], rtol=rtol, atol=atol + ) + if with_padding: + assert torch.equal( + out_cached[batch:], torch.full_like(out_cached[batch:], 42) + ) + + if bool(is_flush[:batch].all()): + torch.testing.assert_close( + state_cached[state_indices], + state_baseline[state_indices], + rtol=rtol, + atol=atol, + ) + + write_pos = torch.where(is_flush, torch.zeros_like(write_pos), write_pos + 1) + + assert torch.equal(state_cached[unused_states], state_before[unused_states]) + assert torch.equal(state_baseline[unused_states], state_before[unused_states]) + + +# Geometries with (nheads, ngroups) both divisible by the tp below. +_TP_GEOMETRIES = [ + pytest.param((8, 64, 64, 4), id="small"), + pytest.param((96, 80, 128, 8), id="nano4b"), +] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +@pytest.mark.parametrize("precision", _PRECISIONS) +@pytest.mark.parametrize("geometry", _TP_GEOMETRIES) +@pytest.mark.parametrize("tp", [2]) +def test_replayssm_standard_decode_tp_head_shard_equivalence( + precision: tuple[torch.dtype, torch.dtype], + geometry: tuple[int, int, int, int], + tp: int, +): + """Tensor-parallel correctness at the kernel boundary (single GPU). + + The kernel has no cross-rank communication, so head sharding must be exactly + separable: one run over all heads must equal concatenating ``tp`` independent + per-rank runs on ``nheads // tp`` heads and ``ngroups // tp`` groups. This + guards the per-rank divisors and group->head mapping the state-shape wiring + relies on. The real TP1==TP2 engine check lives in the v1/e2e suite.""" + state_dtype, act_dtype = precision + nheads, headdim, dstate, ngroups = geometry + assert nheads % tp == 0 and ngroups % tp == 0 + device = "cuda" + both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32 + rtol, atol = _tolerances(torch.float32 if both_fp32 else torch.bfloat16) + set_random_seed(0) + + batch = 4 + max_cache_len = 4 + num_steps = 2 * max_cache_len + 1 + nh_s = nheads // tp + ng_s = ngroups // tp + + # Shards slice the tied params; never .contiguous() -- that would drop the + # stride-0 broadcast the kernel's TIE_HDIM asserts require. + A = _tied_A(nheads, headdim, dstate, device) + dt_bias = _tied_dt_bias(nheads, headdim, device) + D = torch.randn(nheads, headdim, device=device) + + state_full = torch.randn( + batch, nheads, headdim, dstate, dtype=state_dtype, device=device + ) + x_cache, dt_cache, B_cache, _ = allocate_update_caches( + batch, + nheads, + ngroups, + headdim, + dstate, + max_cache_len, + device, + act_dtype, + act_dtype, + ) + bc_pre = torch.empty( + batch, ngroups, max_cache_len, device=device, dtype=torch.float32 + ) + + # Per-rank shards, each seeded from the matching head slice so all start equal. + shard_state = [] + shard_caches = [] + for r in range(tp): + h0 = r * nh_s + shard_state.append(state_full[:, h0 : h0 + nh_s].contiguous()) + xc, dtc, bc, _ = allocate_update_caches( + batch, + nh_s, + ng_s, + headdim, + dstate, + max_cache_len, + device, + act_dtype, + act_dtype, + ) + bcp = torch.empty( + batch, ng_s, max_cache_len, device=device, dtype=torch.float32 + ) + shard_caches.append((xc, dtc, bc, bcp)) + + write_pos = torch.zeros(batch, dtype=torch.int32, device=device) + for _ in range(num_steps): + x = torch.randn(batch, nheads, headdim, device=device, dtype=act_dtype) + dt = _tied_dt(batch, nheads, headdim, device, act_dtype) + B = torch.randn(batch, ngroups, dstate, device=device, dtype=act_dtype) + C = torch.randn(batch, ngroups, dstate, device=device, dtype=act_dtype) + z = torch.randn_like(x) + is_flush = write_pos == max_cache_len - 1 + + out_full = torch.empty_like(x) + selective_state_update_replayssm_output_only( + state_full, + x, + dt, + A, + B, + C, + D=D, + z=z, + dt_bias=dt_bias, + dt_softplus=True, + x_cache=x_cache, + dt_cache=dt_cache, + B_cache=B_cache, + bc_pre=bc_pre, + write_pos=write_pos, + is_flush=is_flush, + max_cache_len=max_cache_len, + out=out_full, + ) + + for r in range(tp): + h0 = r * nh_s + g0 = r * ng_s + xc, dtc, bc, bcp = shard_caches[r] + out_shard = torch.empty( + batch, nh_s, headdim, device=device, dtype=act_dtype + ) + selective_state_update_replayssm_output_only( + shard_state[r], + x[:, h0 : h0 + nh_s].contiguous(), + dt[:, h0 : h0 + nh_s], + A[h0 : h0 + nh_s], + B[:, g0 : g0 + ng_s].contiguous(), + C[:, g0 : g0 + ng_s].contiguous(), + D=D[h0 : h0 + nh_s].contiguous(), + z=z[:, h0 : h0 + nh_s].contiguous(), + dt_bias=dt_bias[h0 : h0 + nh_s], + dt_softplus=True, + x_cache=xc, + dt_cache=dtc, + B_cache=bc, + bc_pre=bcp, + write_pos=write_pos, + is_flush=is_flush, + max_cache_len=max_cache_len, + out=out_shard, + ) + + torch.testing.assert_close( + out_full[:, h0 : h0 + nh_s], out_shard, rtol=rtol, atol=atol + ) + if bool(is_flush.any()): + torch.testing.assert_close( + state_full[:, h0 : h0 + nh_s][is_flush], + shard_state[r][is_flush], + rtol=rtol, + atol=atol, + ) + + write_pos = torch.where(is_flush, torch.zeros_like(write_pos), write_pos + 1) diff --git a/tests/kernels/mamba/utils.py b/tests/kernels/mamba/utils.py index fb8a4b0a28e..b82c4d82c33 100644 --- a/tests/kernels/mamba/utils.py +++ b/tests/kernels/mamba/utils.py @@ -76,3 +76,163 @@ def selective_state_update_ref( if not has_heads: out = out.squeeze(1) return out + + +def selective_state_update_replayssm_output_only_ref( + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + z: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + x_cache: torch.Tensor | None = None, + dt_cache: torch.Tensor | None = None, + B_cache: torch.Tensor | None = None, + write_pos: torch.Tensor | None = None, + max_cache_len: int = 16, +) -> torch.Tensor: + """Pure-PyTorch cached-bc reference for validation.""" + has_heads = state.dim() > 3 + if state.dim() == 3: + state = state.unsqueeze(1) + if x.dim() == 2: + x = x.unsqueeze(1) + if dt.dim() == 2: + dt = dt.unsqueeze(1) + if A.dim() == 2: + A = A.unsqueeze(0) + if B.dim() == 2: + B = B.unsqueeze(1) + if C.dim() == 2: + C = C.unsqueeze(1) + if D is not None and D.dim() == 1: + D = D.unsqueeze(0) + if z is not None and z.dim() == 2: + z = z.unsqueeze(1) + if dt_bias is not None and dt_bias.dim() == 1: + dt_bias = dt_bias.unsqueeze(0) + + batch, nheads, dim, dstate = state.shape + assert x.shape == (batch, nheads, dim) + assert dt.shape == x.shape + assert A.shape == (nheads, dim, dstate) + ngroups = B.shape[1] + assert nheads % ngroups == 0, "nheads must be divisible by ngroups" + assert B.shape == (batch, ngroups, dstate) + assert C.shape == B.shape + + ratio = nheads // ngroups + + dt_val = dt[:, :, 0].float() + if dt_bias is not None: + dt_val = dt_val + dt_bias[:, 0].float() + if dt_softplus: + dt_val = F.softplus(dt_val) + A_val = A[:, 0, 0].float() + C_heads = C.repeat_interleave(ratio, dim=1) + out = torch.empty(batch, nheads, dim, device=x.device, dtype=torch.float32) + + assert x_cache is not None + assert dt_cache is not None + assert B_cache is not None + assert write_pos is not None + + for b in range(batch): + cache_len = int(write_pos[b].item()) + is_flush = cache_len == max_cache_len - 1 + n_steps = cache_len + 1 + + dt_all = torch.zeros(nheads, n_steps, device=x.device, dtype=torch.float32) + if cache_len > 0: + dt_all[:, :cache_len] = dt_cache[b, :, :cache_len] + dt_all[:, cache_len] = dt_val[b] + + cumsum = torch.cumsum(dt_all, dim=-1) + total = cumsum[:, -1] + dA_cumsum = A_val[:, None] * cumsum + dA_total = A_val * total + total_decay = torch.exp(dA_total) + scale = dt_all * torch.exp(dA_total[:, None] - dA_cumsum) + + x_all = torch.zeros(nheads, dim, n_steps, device=x.device, dtype=x.dtype) + if cache_len > 0: + x_all[..., :cache_len] = x_cache[b, :, :cache_len, :].permute(0, 2, 1) + x_all[..., cache_len] = x[b] + + B_all = torch.zeros(ngroups, n_steps, dstate, device=B.device, dtype=B.dtype) + if cache_len > 0: + B_all[:, :cache_len, :] = B_cache[b, :, :cache_len, :] + B_all[:, cache_len, :] = B[b] + + B_heads = B_all.repeat_interleave(ratio, dim=0) + C_heads_b = C_heads[b] + + if is_flush: + B_scaled = (B_heads.float() * scale[:, :, None]).to(B_heads.dtype) + delta = torch.einsum("hdk,hkn->hdn", x_all.float(), B_scaled.float()) + state_new = state[b].float() * total_decay[:, None, None] + delta + state[b].copy_(state_new.to(state.dtype)) + out[b] = torch.einsum("hdn,hn->hd", state_new, C_heads_b.float()) + else: + checkpoint_out = torch.einsum( + "hdn,hn->hd", state[b].float(), C_heads_b.float() + ) + checkpoint_out = checkpoint_out * total_decay[:, None] + BC = torch.einsum("hkn,hn->hk", B_heads.float(), C_heads_b.float()) + cache_out = torch.einsum("hdk,hk->hd", x_all.float(), scale * BC) + out[b] = checkpoint_out + cache_out + x_cache[b, :, cache_len, :] = x[b] + dt_cache[b, :, cache_len] = dt_val[b] + B_cache[b, :, cache_len, :] = B[b] + + if D is not None: + out = out + (x.float() * D[None]).to(out.dtype) + if z is not None: + out = out * F.silu(z.float()) + out = out.to(x.dtype) + if not has_heads: + out = out.squeeze(1) + return out + + +def allocate_update_caches( + batch: int, + nheads: int, + ngroups: int, + dim: int, + dstate: int, + max_cache_len: int, + device: torch.device, + x_dtype: torch.dtype, + B_dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Allocate dense reference caches for standalone validation.""" + x_cache = torch.zeros( + batch, + nheads, + max_cache_len, + dim, + device=device, + dtype=x_dtype, + ) + dt_cache = torch.zeros( + batch, + nheads, + max_cache_len, + device=device, + dtype=torch.float32, + ) + B_cache = torch.zeros( + batch, + ngroups, + max_cache_len, + dstate, + device=device, + dtype=B_dtype, + ) + write_pos = torch.zeros(batch, dtype=torch.int32, device=device) + return x_cache, dt_cache, B_cache, write_pos diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py index 1418704036c..45db16e8799 100644 --- a/tests/models/test_registry.py +++ b/tests/models/test_registry.py @@ -134,6 +134,22 @@ def test_registry_is_pp(model_arch, is_pp, init_cuda): ) +@create_new_process_for_each_test() +@pytest.mark.parametrize( + "model_arch,supported", + [ + # ReplaySSM is opt-in per model; only Nemotron-H sets the flag today. + ("NemotronHForCausalLM", True), + ("Mamba2ForCausalLM", False), + ("Zamba2ForCausalLM", False), + ], +) +def test_registry_supports_replayssm(model_arch, supported): + model_info = ModelRegistry._try_inspect_model_cls(model_arch) + assert model_info is not None + assert model_info.supports_replayssm is supported + + def test_lazy_modelinfo_package_hash_includes_submodules(tmp_path): package_dir = tmp_path / "model_package" package_dir.mkdir() diff --git a/tests/v1/attention/test_mamba_update_block_table.py b/tests/v1/attention/test_mamba_update_block_table.py index 99dcb09ab15..4ec13827020 100644 --- a/tests/v1/attention/test_mamba_update_block_table.py +++ b/tests/v1/attention/test_mamba_update_block_table.py @@ -42,6 +42,8 @@ def _make_vllm_config( cache_config=SimpleNamespace( block_size=block_size, mamba_cache_mode="all", + use_replayssm=False, + replayssm_buffer_len=16, ), compilation_config=SimpleNamespace( cudagraph_mode=CUDAGraphMode.FULL, diff --git a/tests/v1/attention/test_replayssm_metadata_builder.py b/tests/v1/attention/test_replayssm_metadata_builder.py new file mode 100644 index 00000000000..cbbfedee964 --- /dev/null +++ b/tests/v1/attention/test_replayssm_metadata_builder.py @@ -0,0 +1,256 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Mamba2 ReplaySSM decode write-position derivation in +BaseMambaAttentionMetadataBuilder: write_pos and is_flush computed from the +per-request ring origin (replayssm_decode_base) and num_computed. +""" + +from dataclasses import dataclass + +import pytest +import torch + +from tests.v1.attention.utils import ( + BatchSpec, + MockMambaBuilder, + create_common_attn_metadata, + create_vllm_config, +) +from vllm.v1.kv_cache_interface import MambaSpec + +BLOCK_SIZE = 16 +DEVICE = torch.device("cpu") + + +@dataclass +class ReplaySSMBuildCase: + """A decode batch and its expected per-row write_pos / is_flush. + + num_computed = seq_len - query_len; write_pos = + (num_computed - decode_base) % buffer_len; is_flush = write_pos == + buffer_len - 1 (or a forced one-token flush when num_computed < decode_base). + """ + + seq_lens: list[int] + query_lens: list[int] + is_prefilling: list[bool] + decode_base: list[int] + buffer_len: int + expected_write_pos: list[int] + expected_is_flush: list[int] + mamba_cache_mode: str = "none" + + +REPLAYSSM_BUILD_CASES = { + # decode_base == num_prompt (fresh request). + "fresh_decode": ReplaySSMBuildCase( + seq_lens=[106], + query_lens=[1], + is_prefilling=[False], + decode_base=[100], + buffer_len=16, + expected_write_pos=[5], + expected_is_flush=[0], + ), + # decode_base > num_prompt anchors write_pos at the resume point. + "resumed_reanchors_to_zero": ReplaySSMBuildCase( + seq_lens=[106], + query_lens=[1], + is_prefilling=[False], + decode_base=[105], + buffer_len=16, + expected_write_pos=[0], + expected_is_flush=[0], + ), + # write_pos == buffer_len - 1 flushes. + "flush_boundary": ReplaySSMBuildCase( + seq_lens=[116], + query_lens=[1], + is_prefilling=[False], + decode_base=[100], + buffer_len=16, + expected_write_pos=[15], + expected_is_flush=[1], + ), + # Resumed request landing on a flush boundary. + "resumed_flush_boundary": ReplaySSMBuildCase( + seq_lens=[121], + query_lens=[1], + is_prefilling=[False], + decode_base=[105], + buffer_len=16, + expected_write_pos=[15], + expected_is_flush=[1], + ), + # Per-row write_pos / is_flush are independent. + "mixed_rows": ReplaySSMBuildCase( + seq_lens=[104, 106, 216], + query_lens=[1, 1, 1], + is_prefilling=[False, False, False], + decode_base=[100, 105, 200], + buffer_len=16, + expected_write_pos=[3, 0, 15], + expected_is_flush=[0, 0, 1], + ), + # write_pos wraps within the buffer (6 % 4 == 2). + "small_buffer_wrap": ReplaySSMBuildCase( + seq_lens=[112], + query_lens=[1], + is_prefilling=[False], + decode_base=[105], + buffer_len=4, + expected_write_pos=[2], + expected_is_flush=[0], + ), + # Single-token prefill-as-decode still in the prompt (num_computed < + # decode_base): forced one-token flush. + "leftover_prompt_one_token_flush": ReplaySSMBuildCase( + seq_lens=[100], + query_lens=[1], + is_prefilling=[True], + decode_base=[100], + buffer_len=16, + expected_write_pos=[0], + expected_is_flush=[1], + ), + # Align mode (block_size 16). Past the first boundary the ring re-anchors at + # the block start: num_computed 117 -> block_start 112, write_pos 5 (vs 1 in + # none mode). + "align_reanchor_past_boundary": ReplaySSMBuildCase( + seq_lens=[118], + query_lens=[1], + is_prefilling=[False], + decode_base=[100], + buffer_len=16, + expected_write_pos=[5], + expected_is_flush=[0], + mamba_cache_mode="align", + ), + # First-block boundary (num_computed+1 == 112) forces a flush even though + # write_pos (11) != buffer_len - 1. + "align_first_block_boundary_flush": ReplaySSMBuildCase( + seq_lens=[112], + query_lens=[1], + is_prefilling=[False], + decode_base=[100], + buffer_len=16, + expected_write_pos=[11], + expected_is_flush=[1], + mamba_cache_mode="align", + ), + # First step of a new block re-anchors write_pos to 0. + "align_new_block_start_zero": ReplaySSMBuildCase( + seq_lens=[113], + query_lens=[1], + is_prefilling=[False], + decode_base=[100], + buffer_len=16, + expected_write_pos=[0], + expected_is_flush=[0], + mamba_cache_mode="align", + ), + # block_size % buffer_len == 0: a later boundary lands on write_pos == + # buffer_len - 1, so the boundary flush coincides with the natural flush. + "align_boundary_coincides_natural_flush": ReplaySSMBuildCase( + seq_lens=[128], + query_lens=[1], + is_prefilling=[False], + decode_base=[100], + buffer_len=16, + expected_write_pos=[15], + expected_is_flush=[1], + mamba_cache_mode="align", + ), + # block_size % buffer_len != 0 (buffer_len 6): the boundary step still flushes + # although write_pos (3) != buffer_len - 1. + "align_unaligned_buffer_forces_flush": ReplaySSMBuildCase( + seq_lens=[128], + query_lens=[1], + is_prefilling=[False], + decode_base=[100], + buffer_len=6, + expected_write_pos=[3], + expected_is_flush=[1], + mamba_cache_mode="align", + ), + # Per-row independence in align mode: partial-block / new-block / boundary. + "align_mixed_rows": ReplaySSMBuildCase( + seq_lens=[105, 113, 112], + query_lens=[1, 1, 1], + is_prefilling=[False, False, False], + decode_base=[100, 100, 100], + buffer_len=16, + expected_write_pos=[4, 0, 11], + expected_is_flush=[0, 0, 1], + mamba_cache_mode="align", + ), +} + + +def _make_mamba_spec(buffer_len: int) -> MambaSpec: + # Five-tensor ReplaySSM page; the builder only reads shapes[4][0] (bc groups). + return MambaSpec( + block_size=BLOCK_SIZE, + shapes=( + (1, 1), + (1, 1, 1), + (1, buffer_len, 1), + (1, buffer_len), + (1, buffer_len, 1), + ), + dtypes=(torch.float32,), + ) + + +def _create_replayssm_builder( + buffer_len: int, mamba_cache_mode: str = "none" +) -> MockMambaBuilder: + vllm_config = create_vllm_config( + model_name="Qwen/Qwen3.5-0.8B", block_size=BLOCK_SIZE + ) + # Set the flags after construction to skip validate_mamba_cached_kernel + # (it requires a Triton backend) on the mock model. + vllm_config.cache_config.use_replayssm = True + vllm_config.cache_config.replayssm_buffer_len = buffer_len + vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode + return MockMambaBuilder( + _make_mamba_spec(buffer_len), ["layer0"], vllm_config, DEVICE + ) + + +def _build(builder: MockMambaBuilder, case: ReplaySSMBuildCase): + batch = BatchSpec(seq_lens=case.seq_lens, query_lens=case.query_lens) + common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE).replace( + is_prefilling=torch.tensor(case.is_prefilling, dtype=torch.bool), + replayssm_decode_base_cpu=torch.tensor(case.decode_base, dtype=torch.int32), + ) + return builder.build(0, common) + + +@pytest.mark.parametrize( + "case", REPLAYSSM_BUILD_CASES.values(), ids=REPLAYSSM_BUILD_CASES.keys() +) +def test_replayssm_write_pos(case: ReplaySSMBuildCase): + builder = _create_replayssm_builder(case.buffer_len, case.mamba_cache_mode) + meta = _build(builder, case) + + assert meta.write_pos_d is not None + assert meta.is_flush_d is not None + n = len(case.expected_write_pos) + assert meta.write_pos_d[:n].tolist() == case.expected_write_pos + assert meta.is_flush_d[:n].tolist() == case.expected_is_flush + + +def test_resumed_request_differs_from_fresh(): + """Same token count, different decode_base: fresh (base 100) -> write_pos 5, + resumed (base 105) -> write_pos 0.""" + builder = _create_replayssm_builder(16) + batch = BatchSpec(seq_lens=[106, 106], query_lens=[1, 1]) + common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE).replace( + is_prefilling=torch.tensor([False, False]), + replayssm_decode_base_cpu=torch.tensor([100, 105], dtype=torch.int32), + ) + meta = builder.build(0, common) + + assert meta.write_pos_d.tolist()[:2] == [5, 0] + assert meta.is_flush_d.tolist()[:2] == [0, 0] diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py new file mode 100644 index 00000000000..4fc6768e744 --- /dev/null +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Engine-level parity: ReplaySSM standard decode vs the baseline SSM kernel.""" + +import pytest + +from vllm.v1.metrics.reader import Counter + +from ...models.utils import check_logprobs_close +from ...utils import large_gpu_mark, multi_gpu_test + +# Mamba2 (Nemotron-3) hybrid. +MAMBA2_MODEL = "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16" +MODELS = [ + pytest.param(MAMBA2_MODEL, marks=large_gpu_mark(min_gb=40)), +] + +PROMPTS = [ + "The capital of France is", + "Once upon a time, in a small village,", +] + + +def _check_replayssm_parity(vllm_runner, model_name, *, tensor_parallel_size=1): + # Compare logprobs, not greedy ids: ReplaySSM's fp arithmetic can flip a + # near-tie. Baseline and ReplaySSM run at the same TP, so TP numerics are + # common-mode and only ReplaySSM varies. + common = dict( + max_model_len=1024, + trust_remote_code=True, + enable_prefix_caching=False, + mamba_cache_mode="none", + tensor_parallel_size=tensor_parallel_size, + ) + with vllm_runner(model_name, **common) as llm: + baseline = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5) + with vllm_runner( + model_name, use_replayssm=True, replayssm_buffer_len=16, **common + ) as llm: + replay = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5) + + check_logprobs_close( + outputs_0_lst=baseline, + outputs_1_lst=replay, + name_0="baseline", + name_1="replayssm", + ) + + +@pytest.mark.parametrize("model_name", MODELS) +def test_replayssm_decode_matches_baseline(vllm_runner, model_name): + _check_replayssm_parity(vllm_runner, model_name) + + +@multi_gpu_test(num_gpus=2) +@pytest.mark.parametrize("model_name", [MAMBA2_MODEL]) +def test_replayssm_decode_matches_baseline_tp2(vllm_runner, model_name): + # Tensor-parallel correctness: ReplaySSM's caches and checkpoint state are + # sharded per rank, so TP2 decode must still match the baseline at TP2. + _check_replayssm_parity(vllm_runner, model_name, tensor_parallel_size=2) + + +# Prefix spans several mamba blocks; prefix caching only reuses full blocks. +_PC_SENTENCE = ( + "In a detailed survey of state space models, the authors compared many " + "architectures across a wide range of long-context language tasks and " + "measured their throughput, memory use, and accuracy in careful detail. " +) +_PC_PREFIX = _PC_SENTENCE * 120 +PREFIX_CACHING_PROMPTS = [ + _PC_PREFIX + "The most important conclusion was that", + _PC_PREFIX + "Surprisingly, the experiments showed that", + _PC_PREFIX + "The most important conclusion was that", +] + + +def _prefix_cache_hits(llm) -> int: + return sum( + m.value + for m in llm.llm.get_metrics() + if isinstance(m, Counter) and m.name == "vllm:prefix_cache_hits" + ) + + +def _check_replayssm_prefix_caching_parity( + vllm_runner, model_name, *, tensor_parallel_size=1 +): + # align mode materializes the exact SSM state at each block boundary, so + # ReplaySSM's cached prefixes must match the always-materialized baseline. + common = dict( + max_model_len=8192, + trust_remote_code=True, + enable_prefix_caching=True, + enable_chunked_prefill=True, + mamba_cache_mode="align", + disable_log_stats=False, # required for llm.get_metrics() + tensor_parallel_size=tensor_parallel_size, + ) + with vllm_runner(model_name, **common) as llm: + baseline = llm.generate_greedy_logprobs( + PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + ) + with vllm_runner( + model_name, use_replayssm=True, replayssm_buffer_len=16, **common + ) as llm: + # Prime the cache, then measure, so cache hits are deterministic. + llm.generate_greedy_logprobs( + PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + ) + replay = llm.generate_greedy_logprobs( + PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + ) + replay_hits = _prefix_cache_hits(llm) + + # Without real cache hits the cached path is never exercised. + assert replay_hits > 0, ( + "ReplaySSM align-mode run produced no prefix-cache hits; the shared " + "prefix may be shorter than one mamba block, so prefix caching is inert" + ) + check_logprobs_close( + outputs_0_lst=baseline, + outputs_1_lst=replay, + name_0="baseline_align_pc", + name_1="replayssm_align_pc", + ) + + +@pytest.mark.parametrize("model_name", MODELS) +def test_replayssm_prefix_caching_matches_baseline(vllm_runner, model_name): + _check_replayssm_prefix_caching_parity(vllm_runner, model_name) + + +@multi_gpu_test(num_gpus=2) +@pytest.mark.parametrize("model_name", [MAMBA2_MODEL]) +def test_replayssm_prefix_caching_matches_baseline_tp2(vllm_runner, model_name): + _check_replayssm_prefix_caching_parity( + vllm_runner, model_name, tensor_parallel_size=2 + ) diff --git a/vllm/config/cache.py b/vllm/config/cache.py index a628e7d7cdd..45c93624188 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -145,6 +145,17 @@ class CacheConfig: - "align": only cache the mamba state of the last token of each scheduler step and when the token is at position i * block_size. """ + replayssm_buffer_len: int = Field(default=16, gt=0) + """ReplaySSM history buffer length B: with use_replayssm, standard decode + caches recent SSM inputs in a size-B ring buffer and flushes the checkpoint + state to HBM every B steps. Default 16.""" + use_replayssm: bool = False + """Use the ReplaySSM Mamba2 decode kernel: cache recent SSM inputs and skip + the per-step full-state store, writing the checkpoint back only on flush. + Requires mamba_cache_mode 'none' or 'align' (prefix caching) and the Triton + mamba backend; standard (non-speculative) decode only. In align mode flushes + are most efficient when mamba_block_size is a multiple of replayssm_buffer_len, + but this is not required.""" # Will be set after profiling. num_gpu_blocks: int | None = field(default=None, init=False) diff --git a/vllm/config/model.py b/vllm/config/model.py index d9d2f57dc4b..0e7621c2b17 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1696,6 +1696,10 @@ class ModelConfig: def supports_mamba_prefix_caching(self) -> bool: return self._model_info.supports_mamba_prefix_caching + @property + def supports_replayssm(self) -> bool: + return self._model_info.supports_replayssm + @property def use_mla(self) -> bool: return self.is_deepseek_mla and not envs.VLLM_MLA_DISABLE diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index c05e35cbbcf..a7f185cad40 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -39,7 +39,7 @@ from .kv_events import KVEventsConfig from .kv_transfer import KVTransferConfig from .load import LoadConfig from .lora import LoRAConfig -from .mamba import MambaConfig +from .mamba import MambaBackendEnum, MambaConfig from .model import ModelConfig from .observability import ObservabilityConfig from .offload import OffloadConfig @@ -2308,6 +2308,37 @@ class VllmConfig: ) return self + @model_validator(mode="after") + def validate_mamba_cached_kernel(self) -> "VllmConfig": + if not self.cache_config.use_replayssm: + return self + # ReplaySSM adds a 3-tensor ring to the mamba state; only models that + # opt in (supports_replayssm) build a consistent shape on both the layer + # and config paths. Reject others so the mamba page size cannot desync. + if self.model_config is not None and not self.model_config.supports_replayssm: + raise ValueError( + "--use-replayssm is only supported for Nemotron-H models " + f"(got architecture {self.model_config.architecture!r})" + ) + if self.cache_config.mamba_cache_mode == "all": + raise ValueError( + "--use-replayssm supports prefix caching only in align mode; " + "pass --mamba-cache-mode align" + ) + if self.num_speculative_tokens > 0: + raise ValueError("--use-replayssm does not support speculative decoding") + if self.mamba_config.backend != MambaBackendEnum.TRITON: + raise ValueError("--use-replayssm requires --mamba-backend triton") + if ( + self.kv_transfer_config is not None + and self.kv_transfer_config.is_kv_transfer_instance + ): + raise ValueError( + "--use-replayssm is incompatible with KV connectors " + "(P/D disaggregation, KV cache offload)" + ) + return self + _current_vllm_config: VllmConfig | None = None _current_prefix: str | None = None diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 29ebb988e93..a14de27190e 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -695,6 +695,8 @@ class EngineArgs: mamba_block_size: int | None = get_field(CacheConfig, "mamba_block_size") prefix_match_unit: int | None = get_field(CacheConfig, "prefix_match_unit") mamba_cache_mode: MambaCacheMode = CacheConfig.mamba_cache_mode + replayssm_buffer_len: int = CacheConfig.replayssm_buffer_len + use_replayssm: bool = CacheConfig.use_replayssm mamba_backend: MambaBackendEnum = MambaBackendEnum.TRITON enable_mamba_cache_stochastic_rounding: bool = ( @@ -1201,6 +1203,10 @@ class EngineArgs: cache_group.add_argument( "--mamba-cache-mode", **cache_kwargs["mamba_cache_mode"] ) + cache_group.add_argument( + "--replayssm-buffer-len", **cache_kwargs["replayssm_buffer_len"] + ) + cache_group.add_argument("--use-replayssm", **cache_kwargs["use_replayssm"]) cache_group.add_argument( "--kv-offloading-size", **cache_kwargs["kv_offloading_size"] ) @@ -1910,6 +1916,8 @@ class EngineArgs: mamba_block_size=self.mamba_block_size, prefix_match_unit=self.prefix_match_unit, mamba_cache_mode=self.mamba_cache_mode, + replayssm_buffer_len=self.replayssm_buffer_len, + use_replayssm=self.use_replayssm, kv_offloading_size=self.kv_offloading_size, kv_offloading_backend=self.kv_offloading_backend, ) diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index a6524961ea9..7538a2b6b49 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -32,6 +32,9 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( causal_conv1d_update, ) from vllm.model_executor.layers.mamba.ops.layernorm_gated import rms_norm_gated +from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_output_only import ( # noqa: E501 + selective_state_update_replayssm_output_only, +) from vllm.model_executor.layers.mamba.ops.ssd_combined import ( mamba_chunk_scan_combined_varlen, ) @@ -494,12 +497,27 @@ class MambaMixer2(MambaBase, PluggableLayer): if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self - # The tuple is (conv_state, ssm_state) - self.kv_cache = (torch.tensor([]), torch.tensor([])) self.model_config = model_config self.cache_config = cache_config self.prefix = prefix + self.use_replayssm = ( + cache_config.use_replayssm if cache_config is not None else False + ) + self.replayssm_buffer_len = ( + cache_config.replayssm_buffer_len + if cache_config is not None and cache_config.use_replayssm + else None + ) + self.mamba_config = vllm_config.mamba_config + if self.use_replayssm and self.num_heads % self.tp_size != 0: + raise ValueError( + "--use-replayssm requires tensor-parallel heads to divide evenly" + ) + # The tuple is (conv_state, ssm_state); with the cached (ReplaySSM) decode + # kernel enabled it is (conv_state, ssm_state, x_cache, dt_cache, B_cache). + _n_state = 5 if self.use_replayssm else 2 + self.kv_cache = tuple(torch.tensor([]) for _ in range(_n_state)) self.num_spec = vllm_config.num_speculative_tokens if self.num_spec > 0: @@ -591,7 +609,7 @@ class MambaMixer2(MambaBase, PluggableLayer): # Triton's autotuner includes tensor dtypes in its cache key, # so state_dtype must match what real inference uses. - _, ssm_state_dtype = self.get_state_dtype() + ssm_state_dtype = self.get_state_dtype()[1] # SSD kernel autotune keys depend on dtype and head dimensions, # not on sequence length or batch size, so a single shape suffices. @@ -702,6 +720,10 @@ class MambaMixer2(MambaBase, PluggableLayer): else self.kv_cache[0].transpose(-1, -2) ) ssm_state = self.kv_cache[1] + if self.use_replayssm: + x_cache, dt_cache, B_cache = self.kv_cache[2:] + else: + x_cache = dt_cache = B_cache = None has_initial_states_p = attn_metadata.has_initial_states_p prep_initial_states = attn_metadata.prep_initial_states chunk_size = attn_metadata.chunk_size @@ -1027,37 +1049,78 @@ class MambaMixer2(MambaBase, PluggableLayer): # - mamba_cache_params.ssm_state's slots will be selected # using state_indices_tensor_d # NOTE: final output is an in-place update of out tensor - selective_state_update( - ssm_state, - hidden_states_d, - dt_d, - A_d, - B_d, - C_d, - D_d, - dt_bias, - dt_softplus=True, - state_batch_indices=state_indices_tensor_d_input, - dst_state_batch_indices=state_indices_tensor_d_output, - out=preallocated_ssm_out_d.view(num_decode_tokens, -1, self.head_dim), - num_accepted_tokens=num_accepted_tokens, - cu_seqlens=query_start_loc_d, - is_blackwell=self.is_blackwell, + preallocated_ssm_out_d = preallocated_ssm_out_d.view( + num_decode_tokens, -1, self.head_dim ) + if self.use_replayssm: + assert self.replayssm_buffer_len is not None + selective_state_update_replayssm_output_only( + ssm_state, + hidden_states_d, + dt_d, + A_d, + B_d, + C_d, + D_d, + dt_bias, + dt_softplus=True, + x_cache=x_cache, + dt_cache=dt_cache, + B_cache=B_cache, + bc_pre=attn_metadata.bc_pre_scratch, + write_pos=attn_metadata.write_pos_d, + is_flush=attn_metadata.is_flush_d, + max_cache_len=self.replayssm_buffer_len, + state_batch_indices=state_indices_tensor_d_input, + out=preallocated_ssm_out_d, + # Stochastic Rounding for the vanilla decode path is read + # from mamba_config inside ssu_dispatch; the replay kernel + # isn't a dispatch backend, so pass it here. + enable_stochastic_rounding=( + self.mamba_config.enable_stochastic_rounding + ), + cache_philox_rounds=( + self.mamba_config.stochastic_rounding_philox_rounds + ), + ) + else: + selective_state_update( + ssm_state, + hidden_states_d, + dt_d, + A_d, + B_d, + C_d, + D_d, + dt_bias, + dt_softplus=True, + state_batch_indices=state_indices_tensor_d_input, + dst_state_batch_indices=state_indices_tensor_d_output, + out=preallocated_ssm_out_d, + num_accepted_tokens=num_accepted_tokens, + cu_seqlens=query_start_loc_d, + is_blackwell=self.is_blackwell, + ) - def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]: + def get_state_dtype(self) -> tuple[torch.dtype, ...]: assert self.model_config is not None assert self.cache_config is not None - return MambaStateDtypeCalculator.mamba2_state_dtype( + base_dtype = MambaStateDtypeCalculator.mamba2_state_dtype( self.model_config.dtype, self.cache_config.mamba_cache_dtype, self.cache_config.mamba_ssm_cache_dtype, ) + if self.use_replayssm: + return MambaStateDtypeCalculator.append_replayssm_ring( + base_dtype, self.model_config.dtype + ) + return base_dtype - def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]: - return MambaStateShapeCalculator.mamba2_state_shape( + def get_state_shape(self) -> tuple[tuple[int, ...], ...]: + tp_world_size = get_tensor_model_parallel_world_size() + base_shape = MambaStateShapeCalculator.mamba2_state_shape( intermediate_size=self.intermediate_size, - tp_world_size=get_tensor_model_parallel_world_size(), + tp_world_size=tp_world_size, n_groups=self.n_groups, num_heads=self.num_heads, head_dim=self.head_dim, @@ -1065,6 +1128,15 @@ class MambaMixer2(MambaBase, PluggableLayer): conv_kernel=self.conv_kernel_size, num_spec=self.num_spec, ) + if self.use_replayssm: + assert self.replayssm_buffer_len is not None + return MambaStateShapeCalculator.append_replayssm_ring( + base_shape, + self.n_groups, + tp_world_size, + self.replayssm_buffer_len, + ) + return base_shape @property def mamba_type(self) -> MambaAttentionBackendEnum: diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 9e78b822280..d974da5c27d 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -80,6 +80,18 @@ class MambaStateDtypeCalculator: model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype ) + @classmethod + def append_replayssm_ring( + cls, + base_dtypes: tuple[torch.dtype, ...], + model_dtype: ModelDType | torch.dtype, + ) -> tuple[torch.dtype, ...]: + """Append the ReplaySSM ring dtypes to a base ``(conv, ssm)`` tuple: + ``(x_cache, dt_cache, B_cache)`` = ``(activation, fp32, activation)``. + """ + activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype) + return (*base_dtypes, activation_dtype, torch.float32, activation_dtype) + @classmethod def _mamba_state_dtype( cls, @@ -186,6 +198,28 @@ class MambaStateShapeCalculator: temporal_state_shape = (divide(num_heads, tp_world_size), head_dim, state_size) return conv_state_shape, temporal_state_shape + @classmethod + def append_replayssm_ring( + cls, + base_shapes: tuple[tuple[int, ...], ...], + n_groups: int, + tp_world_size: int, + replayssm_buffer_len: int, + ) -> tuple[tuple[int, ...], ...]: + """Append the ReplaySSM ring shapes (x_cache, dt_cache, B_cache) to a + base ``(conv, ssm)`` tuple. ``base_shapes[1]`` is the ssm shape + ``(nheads // tp, head_dim, state_size)``; B_cache uses the un-extended + ``n_groups``. + """ + local_nheads, head_dim, state_size = base_shapes[1] + local_ngroups = divide(n_groups, tp_world_size) + return ( + *base_shapes, + (local_nheads, replayssm_buffer_len, head_dim), + (local_nheads, replayssm_buffer_len), + (local_ngroups, replayssm_buffer_len, state_size), + ) + @classmethod def short_conv_state_shape( cls, diff --git a/vllm/model_executor/layers/mamba/ops/replayssm_config.py b/vllm/model_executor/layers/mamba/ops/replayssm_config.py new file mode 100644 index 00000000000..aff4fc1c3f2 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/replayssm_config.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Launch-config selection for the ReplaySSM Mamba2 output_only decode kernel. + +Mirrors ``mamba_ssm.py``: a hard-coded heuristic per kernel, plus an +``override`` context manager for benchmarks/tests/config sweeps. Hardware is +auto-detected (Blackwell vs not) so call sites need not thread it through. +""" + +import functools +from contextlib import contextmanager + +from vllm.platforms import current_platform +from vllm.triton_utils import triton + + +@functools.cache +def _is_blackwell() -> bool: + try: + return current_platform.is_device_capability_family(100) + except Exception: + return False + + +# Per-kernel overrides keyed by the kernel name passed to get_replayssm_config. +_overrides: dict[str, tuple] = {} + + +@contextmanager +def override_replayssm_config(kernel: str, config: tuple): + """Pin ``kernel``'s launch config for the duration of the context.""" + prev = _overrides.get(kernel) + _overrides[kernel] = config + try: + yield + finally: + if prev is None: + _overrides.pop(kernel, None) + else: + _overrides[kernel] = prev + + +def _dstate_tile(dstate: int, tile: int) -> int: + return max(16, min(tile, triton.next_power_of_2(dstate))) + + +def _mamba2_output_only(dstate, L, is_blackwell): + # (block_size_m, num_warps, nf_dstate_tile, fl_dstate_tile, num_stages); + # decoupled dstate tiling, serving-batch optimum, dtype-independent per device. + if is_blackwell: + return 64, 1, _dstate_tile(dstate, 32), _dstate_tile(dstate, 64), 2 + return 16, 1, _dstate_tile(dstate, 64), _dstate_tile(dstate, 128), 2 + + +def get_replayssm_config(kernel: str, **shape) -> tuple: + """Return the launch config for ``kernel`` (override > tuned default). + + kernel: "mamba2_output_only". ``shape`` carries the keying dims (dstate; + ``L`` for the buffer length, default 16); hardware is auto-detected. + """ + if kernel in _overrides: + return _overrides[kernel] + bw = _is_blackwell() + if kernel == "mamba2_output_only": + return _mamba2_output_only(shape["dstate"], shape.get("L", 16), bw) + raise ValueError(f"unknown ReplaySSM kernel config key: {kernel}") diff --git a/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py b/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py new file mode 100644 index 00000000000..c6046b06a50 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py @@ -0,0 +1,709 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# ruff: noqa: E501 + +import torch + +from vllm.model_executor.layers.mamba.ops.mamba_ssm import convert_rs_fp16x2, softplus +from vllm.model_executor.layers.mamba.ops.replayssm_config import get_replayssm_config +from vllm.triton_utils import tl, triton +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + + +@triton.heuristics( + { + "HAS_STATE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] + is not None + } +) +@triton.heuristics( + {"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])} +) +@triton.jit +def _replayssm_output_only_precompute_kernel( + B_ptr, + C_ptr, + B_cache_ptr, + write_pos_ptr, + is_flush_ptr, + bc_pre_ptr, + state_batch_indices_ptr, + null_block_id, + # Matrix dimensions + batch, + ngroups, + dstate, + # Input strides + stride_B_batch, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_group, + stride_C_dstate, + # Cache strides + stride_B_cache_batch, + stride_B_cache_group, + stride_B_cache_pos, + stride_B_cache_dstate, + stride_bc_pre_batch, + stride_bc_pre_group, + stride_bc_pre_pos, + stride_state_indices_batch, + stride_state_indices_T, + # Meta-parameters + MAX_CACHE_LEN: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + # heuristic-computed + BLOCK_SIZE_DSTATE: tl.constexpr, + HAS_STATE_BATCH_INDICES: tl.constexpr, +): + pid_b = tl.program_id(axis=0) + pid_g = tl.program_id(axis=1) + + # On flush steps the main kernel does not read bc_pre, so skip the work. + is_flush = tl.load(is_flush_ptr + pid_b) != 0 + if is_flush: + return + + if HAS_STATE_BATCH_INDICES: + state_batch_idx = tl.load( + state_batch_indices_ptr + + pid_b * stride_state_indices_batch + + 0 * stride_state_indices_T + ).to(tl.int64) + if state_batch_idx == null_block_id: + return + else: + state_batch_idx = pid_b + + offs_k = tl.arange(0, BLOCK_SIZE_K) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + + write_pos = tl.load(write_pos_ptr + pid_b).to(tl.int64) + + B_ptr += pid_b * stride_B_batch + pid_g * stride_B_group + C_ptr += pid_b * stride_C_batch + pid_g * stride_C_group + B_cache_ptr += state_batch_idx * stride_B_cache_batch + pid_g * stride_B_cache_group + bc_pre_ptr += pid_b * stride_bc_pre_batch + pid_g * stride_bc_pre_group + + B_cur = tl.load( + B_ptr + offs_n * stride_B_dstate, + mask=offs_n < dstate, + other=0.0, + ) + C = tl.load( + C_ptr + offs_n * stride_C_dstate, + mask=offs_n < dstate, + other=0.0, + ) + B_cache_ptrs = ( + B_cache_ptr + + offs_k[:, None] * stride_B_cache_pos + + offs_n[None, :] * stride_B_cache_dstate + ) + B_cache = tl.load( + B_cache_ptrs, + mask=(offs_k[:, None] < write_pos) & (offs_n[None, :] < dstate), + other=0.0, + ) + B_all = tl.where(offs_k[:, None] == write_pos, B_cur[None, :], B_cache) + bc = tl.sum(B_all.to(tl.float32) * C[None, :].to(tl.float32), axis=1) + + tl.store( + bc_pre_ptr + offs_k * stride_bc_pre_pos, + bc, + mask=(offs_k <= write_pos) & (offs_k < MAX_CACHE_LEN), + ) + + +@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) +@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) +@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) +@triton.heuristics( + { + "HAS_STATE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] + is not None + } +) +@triton.heuristics( + {"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])} +) +@triton.jit +def _replayssm_output_only_kernel( + # Pointers to matrices + state_ptr, + rand_seed_ptr, + x_ptr, + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + x_cache_ptr, + dt_cache_ptr, + B_cache_ptr, + bc_pre_ptr, + write_pos_ptr, + is_flush_ptr, + state_batch_indices_ptr, + null_block_id, + # Matrix dimensions + batch, + nheads, + dim, + dstate, + nheads_ngroups_ratio, + # State strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # Input strides + stride_x_batch, + stride_x_head, + stride_x_dim, + stride_dt_batch, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + stride_B_batch, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_head, + stride_out_dim, + # Cache strides + stride_x_cache_batch, + stride_x_cache_head, + stride_x_cache_dim, + stride_x_cache_pos, + stride_dt_cache_batch, + stride_dt_cache_head, + stride_dt_cache_pos, + stride_B_cache_batch, + stride_B_cache_group, + stride_B_cache_pos, + stride_B_cache_dstate, + stride_bc_pre_batch, + stride_bc_pre_group, + stride_bc_pre_pos, + stride_state_indices_batch, + stride_state_indices_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + MAX_CACHE_LEN: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_K_CACHE: tl.constexpr, + BLOCK_SIZE_K_DOT: tl.constexpr, + NF_DSTATE_TILE: tl.constexpr, + NF_NDS: tl.constexpr, + FL_DSTATE_TILE: tl.constexpr, + FL_NDS: tl.constexpr, + USE_RS_ROUNDING: tl.constexpr, + PHILOX_ROUNDS: tl.constexpr, + # heuristic-computed + BLOCK_SIZE_DSTATE: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_STATE_BATCH_INDICES: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + pid_b = tl.program_id(axis=1) + pid_h = tl.program_id(axis=2) + + # Resolve the physical state slot for this decode row; skip padded rows. + if HAS_STATE_BATCH_INDICES: + state_batch_idx = tl.load( + state_batch_indices_ptr + + pid_b * stride_state_indices_batch + + 0 * stride_state_indices_T + ).to(tl.int64) + if state_batch_idx == null_block_id: + return + else: + state_batch_idx = pid_b + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n_full = tl.arange(0, BLOCK_SIZE_DSTATE) + + # Buffer cursor (number of cached tokens so far) and the flush flag. + write_pos = tl.load(write_pos_ptr + pid_b).to(tl.int64) + is_flush = tl.load(is_flush_ptr + pid_b) != 0 + + # Advance every pointer to this (row, head, group). + state_ptr += state_batch_idx * stride_state_batch + pid_h * stride_state_head + x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head + dt_ptr += pid_b * stride_dt_batch + pid_h * stride_dt_head + B_ptr += pid_b * stride_B_batch + (pid_h // nheads_ngroups_ratio) * stride_B_group + C_ptr += pid_b * stride_C_batch + (pid_h // nheads_ngroups_ratio) * stride_C_group + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head + x_cache_ptr += state_batch_idx * stride_x_cache_batch + pid_h * stride_x_cache_head + dt_cache_ptr += ( + state_batch_idx * stride_dt_cache_batch + pid_h * stride_dt_cache_head + ) + B_cache_ptr += ( + state_batch_idx * stride_B_cache_batch + + (pid_h // nheads_ngroups_ratio) * stride_B_cache_group + ) + bc_pre_ptr += ( + pid_b * stride_bc_pre_batch + + (pid_h // nheads_ngroups_ratio) * stride_bc_pre_group + ) + + # Current-token dt (+ bias, softplus), scalar A, and current x. C, the + # checkpoint state S_0, and current-token B are read per dstate tile below. + dt_cur = tl.load(dt_ptr).to(tl.float32) + if HAS_DT_BIAS: + dt_cur += tl.load(dt_bias_ptr + pid_h * stride_dt_bias_head).to(tl.float32) + if DT_SOFTPLUS: + dt_cur = tl.where(dt_cur <= 20.0, softplus(dt_cur), dt_cur) + A = tl.load(A_ptr + pid_h * stride_A_head).to(tl.float32) + x_cur = tl.load(x_ptr + offs_m * stride_x_dim, mask=offs_m < dim, other=0.0) + + if not is_flush: + # Output-only route: read y without materializing the state, using the + # precomputed k^T q products (`bc`): + # y = total_decay * (S_0 q) + sum_j s_j (k_j^T q) v_j. + # Then append the current token to the buffer. + offs_k_cache = tl.arange(0, BLOCK_SIZE_K_CACHE) + # dt over the window (history + current token), then the decay weights. + dt_all_cache = tl.load( + dt_cache_ptr + offs_k_cache * stride_dt_cache_pos, + mask=offs_k_cache < write_pos, + other=0.0, + ).to(tl.float32) + dt_all_cache = tl.where(offs_k_cache == write_pos, dt_cur, dt_all_cache) + dA_cumsum_cache = A * tl.cumsum(dt_all_cache, axis=0) + dA_total_cache = A * tl.sum(dt_all_cache, axis=0) + total_decay_cache = tl.exp(dA_total_cache) + scale_cache = dt_all_cache * tl.exp(dA_total_cache - dA_cumsum_cache) + scale_cache = tl.where(offs_k_cache <= write_pos, scale_cache, 0.0) + + # Gather buffered x over the window (history + current token). + x_all_cache_ptrs = ( + x_cache_ptr + + offs_m[:, None] * stride_x_cache_dim + + offs_k_cache[None, :] * stride_x_cache_pos + ) + x_all_cache = tl.load( + x_all_cache_ptrs, + mask=(offs_m[:, None] < dim) & (offs_k_cache[None, :] < write_pos), + other=0.0, + ) + x_all_cache = tl.where( + offs_k_cache[None, :] == write_pos, x_cur[:, None], x_all_cache + ) + + # Decayed checkpoint readout sum_n S_0(m,n) q(n), streamed over NF dstate + # tiles so the (M, N) state slice is never held whole. + offs_nt = tl.arange(0, NF_DSTATE_TILE) + ck_acc = tl.zeros([BLOCK_SIZE_M], dtype=tl.float32) + for i in tl.static_range(NF_NDS): + offs_n = i * NF_DSTATE_TILE + offs_nt + nmask = offs_n < dstate + st = tl.load( + state_ptr + + offs_m[:, None] * stride_state_dim + + offs_n[None, :] * stride_state_dstate, + mask=(offs_m[:, None] < dim) & nmask[None, :], + other=0.0, + ) + c_chunk = tl.load( + C_ptr + offs_n * stride_C_dstate, mask=nmask, other=0.0 + ).to(tl.float32) + ck_acc += tl.sum(st.to(tl.float32) * c_chunk[None, :], axis=1) + checkpoint_out = ck_acc * total_decay_cache + bc_cache = tl.load( + bc_pre_ptr + offs_k_cache * stride_bc_pre_pos, + mask=offs_k_cache <= write_pos, + other=0.0, + ) + cache_out = tl.sum( + x_all_cache.to(tl.float32) * (scale_cache * bc_cache)[None, :], axis=1 + ) + out = checkpoint_out + cache_out + + # Append the current token (x, dt, B) into the buffer at write_pos. + tl.store( + x_cache_ptr + offs_m * stride_x_cache_dim + write_pos * stride_x_cache_pos, + x_cur, + mask=offs_m < dim, + ) + if pid_m == 0: + B_cur = tl.load( + B_ptr + offs_n_full * stride_B_dstate, + mask=offs_n_full < dstate, + other=0.0, + ) + tl.store(dt_cache_ptr + write_pos * stride_dt_cache_pos, dt_cur) + tl.store( + B_cache_ptr + + write_pos * stride_B_cache_pos + + offs_n_full * stride_B_cache_dstate, + B_cur, + mask=offs_n_full < dstate, + ) + else: + # Flush step: state route. Reconstruct the state from cached inputs, + # S_t = total_decay * S_0 + sum_j s_j (v_j k_j^T), persist it as the new + # checkpoint, then read y = S_t q -- streamed over FL dstate tiles. + offs_k_dot = tl.arange(0, BLOCK_SIZE_K_DOT) + dt_all_dot = tl.load( + dt_cache_ptr + offs_k_dot * stride_dt_cache_pos, + mask=offs_k_dot < write_pos, + other=0.0, + ).to(tl.float32) + dt_all_dot = tl.where(offs_k_dot == write_pos, dt_cur, dt_all_dot) + dA_cumsum_dot = A * tl.cumsum(dt_all_dot, axis=0) + dA_total_dot = A * tl.sum(dt_all_dot, axis=0) + total_decay_dot = tl.exp(dA_total_dot) + scale_dot = dt_all_dot * tl.exp(dA_total_dot - dA_cumsum_dot) + scale_dot = tl.where(offs_k_dot <= write_pos, scale_dot, 0.0) + + # Gather buffered x over the window (history + current token). + x_all_dot_ptrs = ( + x_cache_ptr + + offs_m[:, None] * stride_x_cache_dim + + offs_k_dot[None, :] * stride_x_cache_pos + ) + x_all_dot = tl.load( + x_all_dot_ptrs, + mask=(offs_m[:, None] < dim) & (offs_k_dot[None, :] < write_pos), + other=0.0, + ) + x_all_dot = tl.where( + offs_k_dot[None, :] == write_pos, x_cur[:, None], x_all_dot + ) + x_all_ty = x_all_dot.to(x_ptr.dtype.element_ty) + + # Distinct tile locals (_f) from the nf branch: differing tile widths + # would force a shape-mismatched merge at the if/else exit. + offs_nt_f = tl.arange(0, FL_DSTATE_TILE) + out = tl.zeros([BLOCK_SIZE_M], dtype=tl.float32) + for i in tl.static_range(FL_NDS): + offs_n_f = i * FL_DSTATE_TILE + offs_nt_f + nmask_f = offs_n_f < dstate + # Gather buffered B over the window (history + current token). + B_all_dot = tl.load( + B_cache_ptr + + offs_k_dot[:, None] * stride_B_cache_pos + + offs_n_f[None, :] * stride_B_cache_dstate, + mask=(offs_k_dot[:, None] < write_pos) & nmask_f[None, :], + other=0.0, + ) + B_cur_tile = tl.load( + B_ptr + offs_n_f * stride_B_dstate, mask=nmask_f, other=0.0 + ) + B_all_dot = tl.where( + offs_k_dot[:, None] == write_pos, B_cur_tile[None, :], B_all_dot + ) + # tf32x3 keeps fp32 parity with the elementwise baseline (plain tf32 on + # fp32 inputs drifts ~1e-2); bf16/fp16 inputs are unaffected by this flag. + B_scaled = (B_all_dot.to(tl.float32) * scale_dot[:, None]).to( + x_ptr.dtype.element_ty + ) + delta_state = tl.dot(x_all_ty, B_scaled, input_precision="tf32x3") + state_ptrs = ( + state_ptr + + offs_m[:, None] * stride_state_dim + + offs_n_f[None, :] * stride_state_dstate + ) + state_mask = (offs_m[:, None] < dim) & nmask_f[None, :] + st_f = tl.load(state_ptrs, mask=state_mask, other=0.0) + state_new = st_f.to(tl.float32) * total_decay_dot + delta_state.to( + tl.float32 + ) + if USE_RS_ROUNDING: + # Stochastic-round fp32->fp16 (Blackwell cvt.rs), mirroring the + # baseline. Only the flush step stores state, so this runs at 1/L + # the baseline's per-step rate. Absolute per-element offsets seed + # the RNG so each state element draws independently. + rand_seed = tl.load(rand_seed_ptr) + rand_offsets = ( + state_batch_idx * stride_state_batch + + pid_h * stride_state_head + + offs_m[:, None] * stride_state_dim + + offs_n_f[None, :] * stride_state_dstate + ) + if PHILOX_ROUNDS > 0: + rand = tl.randint(rand_seed, rand_offsets, PHILOX_ROUNDS) + else: + rand = tl.randint(rand_seed, rand_offsets) + tl.static_assert( + state_ptrs.dtype.element_ty == tl.float16, + "stochastic rounding requires an fp16 SSM state cache", + ) + state_store = convert_rs_fp16x2(state_new, rand) + else: + state_store = state_new.to(st_f.dtype) + tl.store(state_ptrs, state_store, mask=state_mask) + c_chunk_f = tl.load( + C_ptr + offs_n_f * stride_C_dstate, mask=nmask_f, other=0.0 + ).to(tl.float32) + out += tl.sum(state_new * c_chunk_f[None, :], axis=1) + + # Skip connection (D) and output gate (z). + if HAS_D: + D_ptr += pid_h * stride_D_head + D = tl.load(D_ptr + offs_m * stride_D_dim, mask=offs_m < dim, other=0.0).to( + tl.float32 + ) + out += x_cur.to(tl.float32) * D + if HAS_Z: + z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head + z = tl.load(z_ptr + offs_m * stride_z_dim, mask=offs_m < dim, other=0.0).to( + tl.float32 + ) + out *= z * tl.sigmoid(z) + + tl.store(out_ptr + offs_m * stride_out_dim, out, mask=offs_m < dim) + + +def selective_state_update_replayssm_output_only( + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + z: torch.Tensor | None = None, + dt_softplus: bool = False, + x_cache: torch.Tensor | None = None, + dt_cache: torch.Tensor | None = None, + B_cache: torch.Tensor | None = None, + bc_pre: torch.Tensor | None = None, + write_pos: torch.Tensor | None = None, + is_flush: torch.Tensor | None = None, + max_cache_len: int = 16, + state_batch_indices: torch.Tensor | None = None, + null_block_id: int = NULL_BLOCK_ID, + out: torch.Tensor | None = None, + enable_stochastic_rounding: bool = False, + cache_philox_rounds: int = 0, +) -> torch.Tensor: + """Cached-bc SSM update for vLLM's autoregressive Mamba2 decode path.""" + has_heads = state.dim() > 3 + if state.dim() == 3: + state = state.unsqueeze(1) + if x.dim() == 2: + x = x.unsqueeze(1) + if dt.dim() == 2: + dt = dt.unsqueeze(1) + if A.dim() == 2: + A = A.unsqueeze(0) + if B.dim() == 2: + B = B.unsqueeze(1) + if C.dim() == 2: + C = C.unsqueeze(1) + if D is not None and D.dim() == 1: + D = D.unsqueeze(0) + if z is not None and z.dim() == 2: + z = z.unsqueeze(1) + if dt_bias is not None and dt_bias.dim() == 1: + dt_bias = dt_bias.unsqueeze(0) + if out is not None and out.dim() == 2: + out = out.unsqueeze(1) + if state_batch_indices is not None and state_batch_indices.dim() == 1: + state_batch_indices = state_batch_indices.unsqueeze(1) + + _, nheads, dim, dstate = state.shape + batch = x.shape[0] + assert x.shape == (batch, nheads, dim) + assert dt.shape == x.shape + assert A.shape == (nheads, dim, dstate) + ngroups = B.shape[1] + assert nheads % ngroups == 0, "nheads must be divisible by ngroups" + assert B.shape == (batch, ngroups, dstate) + assert C.shape == B.shape + if D is not None: + assert D.shape == (nheads, dim) + if z is not None: + assert z.shape == x.shape + if dt_bias is not None: + assert dt_bias.shape == (nheads, dim) + assert out is not None and out.shape == x.shape + + assert A.stride(-1) == 0 and A.stride(-2) == 0, ( + "Cached kernel requires TIE_HDIM (A scalar per head)" + ) + assert dt.stride(-1) == 0, "Cached kernel requires TIE_HDIM (dt scalar per head)" + if dt_bias is not None: + assert dt_bias.stride(-1) == 0, ( + "Cached kernel requires TIE_HDIM (dt_bias scalar per head)" + ) + + assert x_cache is not None + assert dt_cache is not None + assert B_cache is not None + assert x_cache.shape[1:] == (nheads, max_cache_len, dim) + assert dt_cache.shape[1:] == (nheads, max_cache_len) + assert B_cache.shape[1:] == (ngroups, max_cache_len, dstate) + assert write_pos is not None and write_pos.shape[0] >= batch + assert write_pos.dtype == torch.int32 + assert is_flush is not None and is_flush.shape[0] >= batch + assert is_flush.dtype in (torch.bool, torch.int8) + assert bc_pre is not None + assert bc_pre.shape[0] >= batch and bc_pre.shape[1] >= ngroups + assert bc_pre.shape[2] == max_cache_len + assert bc_pre.dtype == torch.float32 + if state_batch_indices is not None: + assert state_batch_indices.shape[0] >= batch + assert state_batch_indices.shape[1] >= 1 + + block_size_k_cache = max(1, triton.next_power_of_2(max_cache_len)) + block_size_k_dot = max(16, block_size_k_cache) + block_size_m, num_warps, nf_tile, fl_tile, num_stages = get_replayssm_config( + "mamba2_output_only", dstate=dstate, L=max_cache_len + ) + bs_dstate = triton.next_power_of_2(dstate) + nf_dstate_tile = max(16, min(nf_tile, bs_dstate)) + nf_nds = triton.cdiv(bs_dstate, nf_dstate_tile) + fl_dstate_tile = max(16, min(fl_tile, bs_dstate)) + fl_nds = triton.cdiv(bs_dstate, fl_dstate_tile) + + grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) + z_strides = (z.stride(0), z.stride(1), z.stride(2)) if z is not None else (0, 0, 0) + state_indices_strides = ( + (state_batch_indices.stride(0), state_batch_indices.stride(1)) + if state_batch_indices is not None + else (0, 0) + ) + rand_seed = ( + torch.randint(0, 2**32, (1,), device=state.device) + if enable_stochastic_rounding + else None + ) + + with torch.accelerator.device_index(x.device.index): + # Both kernels always launch: the precompute kernel self-skips flush + # rows per row (the branch can't be hoisted out under CUDA graphs), so + # it is a no-op when every row is flushing. + _replayssm_output_only_precompute_kernel[(batch, ngroups)]( + B, + C, + B_cache, + write_pos, + is_flush, + bc_pre, + state_batch_indices, + null_block_id, + batch, + ngroups, + dstate, + B.stride(0), + B.stride(1), + B.stride(2), + C.stride(0), + C.stride(1), + C.stride(2), + B_cache.stride(0), + B_cache.stride(1), + B_cache.stride(2), + B_cache.stride(3), + bc_pre.stride(0), + bc_pre.stride(1), + bc_pre.stride(2), + state_indices_strides[0], + state_indices_strides[1], + max_cache_len, + block_size_k_cache, + num_warps=2, + ) + _replayssm_output_only_kernel[grid]( + state, + rand_seed, + x, + dt, + dt_bias, + A, + B, + C, + D, + z, + out, + x_cache, + dt_cache, + B_cache, + bc_pre, + write_pos, + is_flush, + state_batch_indices, + null_block_id, + batch, + nheads, + dim, + dstate, + nheads // ngroups, + state.stride(0), + state.stride(1), + state.stride(2), + state.stride(3), + x.stride(0), + x.stride(1), + x.stride(2), + dt.stride(0), + dt.stride(1), + dt_bias.stride(0) if dt_bias is not None else 0, + A.stride(0), + B.stride(0), + B.stride(1), + B.stride(2), + C.stride(0), + C.stride(1), + C.stride(2), + D.stride(0) if D is not None else 0, + D.stride(1) if D is not None else 0, + z_strides[0], + z_strides[1], + z_strides[2], + out.stride(0), + out.stride(1), + out.stride(2), + x_cache.stride(0), + x_cache.stride(1), + x_cache.stride(3), + x_cache.stride(2), + dt_cache.stride(0), + dt_cache.stride(1), + dt_cache.stride(2), + B_cache.stride(0), + B_cache.stride(1), + B_cache.stride(2), + B_cache.stride(3), + bc_pre.stride(0), + bc_pre.stride(1), + bc_pre.stride(2), + state_indices_strides[0], + state_indices_strides[1], + dt_softplus, + max_cache_len, + block_size_m, + block_size_k_cache, + block_size_k_dot, + nf_dstate_tile, + nf_nds, + fl_dstate_tile, + fl_nds, + enable_stochastic_rounding, + cache_philox_rounds, + num_warps=num_warps, + num_stages=num_stages, + ) + + if not has_heads: + out = out.squeeze(1) + return out diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 32534f8fbb7..df9c18ce534 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -984,6 +984,31 @@ def supports_mamba_prefix_caching( return getattr(model, "supports_mamba_prefix_caching", False) +@runtime_checkable +class SupportsReplaySSM(Protocol): + """The interface for models whose Mamba2 layers support ReplaySSM cached + standard decode. + + This is currently experimental. + """ + + supports_replayssm: ClassVar[Literal[True]] = True + + +@overload +def supports_replayssm(model: object) -> TypeIs[SupportsReplaySSM]: ... + + +@overload +def supports_replayssm(model: type[object]) -> TypeIs[type[SupportsReplaySSM]]: ... + + +def supports_replayssm( + model: type[object] | object, +) -> TypeIs[type[SupportsReplaySSM]] | TypeIs[SupportsReplaySSM]: + return getattr(model, "supports_replayssm", False) + + @runtime_checkable class SupportsCrossEncoding(Protocol): """The interface required for all models that support cross encoding.""" diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index 84f24094a3d..50dc5c02221 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -69,6 +69,7 @@ from vllm.model_executor.models.interfaces import ( SupportsMambaPrefixCaching, SupportsPP, SupportsQuant, + SupportsReplaySSM, ) from vllm.model_executor.models.utils import ( AutoWeightsLoader, @@ -707,6 +708,7 @@ class NemotronHForCausalLM( SupportsQuant, MixtureOfExperts, SupportsMambaPrefixCaching, + SupportsReplaySSM, ): # Relevant only if self.has_moe is True is_non_gated_moe: bool = True @@ -742,18 +744,24 @@ class NemotronHForCausalLM( def get_mamba_state_dtype_from_config( cls, vllm_config: "VllmConfig", - ) -> tuple[torch.dtype, torch.dtype]: - return MambaStateDtypeCalculator.mamba2_state_dtype( + ) -> tuple[torch.dtype, ...]: + cache_config = vllm_config.cache_config + base_dtype = MambaStateDtypeCalculator.mamba2_state_dtype( vllm_config.model_config.dtype, - vllm_config.cache_config.mamba_cache_dtype, - vllm_config.cache_config.mamba_ssm_cache_dtype, + cache_config.mamba_cache_dtype, + cache_config.mamba_ssm_cache_dtype, ) + if cache_config.use_replayssm: + return MambaStateDtypeCalculator.append_replayssm_ring( + base_dtype, vllm_config.model_config.dtype + ) + return base_dtype @classmethod def get_mamba_state_shape_from_config( cls, vllm_config: "VllmConfig", - ) -> tuple[tuple[int, int], tuple[int, int, int]]: + ) -> tuple[tuple[int, ...], ...]: """Calculate shapes for Mamba's convolutional and state caches. Args: @@ -763,12 +771,15 @@ class NemotronHForCausalLM( Tuple containing: - conv_state_shape: Shape for convolutional state cache - temporal_state_shape: Shape for state space model cache + - (when use_replayssm is enabled) the x_cache/dt_cache/B_cache + ring-buffer shapes """ parallel_config = vllm_config.parallel_config + cache_config = vllm_config.cache_config hf_config = vllm_config.model_config.hf_config intermediate_size = hf_config.mamba_num_heads * hf_config.mamba_head_dim - return MambaStateShapeCalculator.mamba2_state_shape( + base_shape = MambaStateShapeCalculator.mamba2_state_shape( intermediate_size=intermediate_size, tp_world_size=parallel_config.tensor_parallel_size, n_groups=hf_config.n_groups, @@ -778,6 +789,14 @@ class NemotronHForCausalLM( conv_kernel=hf_config.conv_kernel, num_spec=vllm_config.num_speculative_tokens, ) + if cache_config.use_replayssm: + return MambaStateShapeCalculator.append_replayssm_ring( + base_shape, + hf_config.n_groups, + parallel_config.tensor_parallel_size, + cache_config.replayssm_buffer_len, + ) + return base_shape @classmethod def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index b17a41834d2..a4fd44fc2cd 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -55,6 +55,7 @@ from .interfaces import ( supports_multimodal_encoder_tp_data, supports_multimodal_raw_input_only, supports_pp, + supports_replayssm, supports_transcription, ) from .interfaces_base import ( @@ -789,6 +790,7 @@ class _ModelInfo: is_hybrid: bool has_noops: bool supports_mamba_prefix_caching: bool + supports_replayssm: bool supports_transcription: bool supports_transcription_only: bool @@ -815,6 +817,7 @@ class _ModelInfo: is_attention_free=is_attention_free(model), is_hybrid=is_hybrid(model), supports_mamba_prefix_caching=supports_mamba_prefix_caching(model), + supports_replayssm=supports_replayssm(model), supports_transcription=supports_transcription(model), supports_transcription_only=( supports_transcription(model) and model.supports_transcription_only diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 6f23cff0b1e..b479aab5816 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -480,6 +480,11 @@ class CommonAttentionMetadata: fixed sliding window. None disables R-SWA. The attention backend copies this into its own persistent buffer and reads ``rswa_window`` from model config.""" + replayssm_decode_base_cpu: torch.Tensor | None = None + """(batch_size,) CPU ring origin for Mamba2 ReplaySSM decode: num_computed + at the current decode run's last full-state write. write_pos counts from + here, so a preemption-resumed request re-anchors past the prompt boundary.""" + # WARNING: Deprecated fields. Will be removed in a future release (v0.15.0) _seq_lens_cpu: torch.Tensor | None = None _num_computed_tokens_cpu: torch.Tensor | None = None @@ -591,6 +596,7 @@ class CommonAttentionMetadata: dcp_local_seq_lens_cpu=maybe_slice_reqs(self.dcp_local_seq_lens_cpu), is_prefilling=maybe_slice_reqs(self.is_prefilling), rswa_prefix_lens=maybe_slice_reqs(self.rswa_prefix_lens), + replayssm_decode_base_cpu=maybe_slice_reqs(self.replayssm_decode_base_cpu), ) diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index 6fce7d0dc7e..28928f28948 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -74,6 +74,12 @@ class BaseMambaAttentionMetadata: nums_dict: dict | None = None batch_ptr: torch.Tensor | None = None token_chunk_offset_ptr: torch.Tensor | None = None + # ReplaySSM standard decode: per-row ring cursor and flush flag, plus the + # per-step (decode_rows, ngroups, replayssm_buffer_len) fp32 scratch for the + # precomputed k^T q products. All None when use_replayssm is disabled. + write_pos_d: torch.Tensor | None = None + is_flush_d: torch.Tensor | None = None + bc_pre_scratch: torch.Tensor | None = None class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): @@ -98,6 +104,8 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): self.compilation_config = vllm_config.compilation_config self.num_spec_tokens: int = vllm_config.num_speculative_tokens self.use_spec_decode = self.num_spec_tokens > 0 + self.use_replayssm = vllm_config.cache_config.use_replayssm + self.replayssm_buffer_len = vllm_config.cache_config.replayssm_buffer_len assert isinstance(kv_cache_spec, MambaSpec) scheduler_config = vllm_config.scheduler_config @@ -158,6 +166,36 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): dtype=torch.int32, device=device, ) + # ReplaySSM standard-decode CUDA-graph buffers: per-row ring cursor, + # flush flag, and the k^T q precompute scratch. + if self.use_replayssm: + self.decode_write_pos_d: torch.Tensor = torch.empty( + (self.decode_cudagraph_max_bs,), + dtype=torch.int32, + device=device, + ) + self.decode_is_flush_d: torch.Tensor = torch.empty( + (self.decode_cudagraph_max_bs,), + dtype=torch.int8, + device=device, + ) + # B_cache shape = (ngroups, replayssm_buffer_len, dstate); the page + # layout is (conv_state, ssm_state, x_cache, dt_cache, B_cache). + bc_ngroups = kv_cache_spec.shapes[4][0] + bc_scratch_bs = max( + self.decode_cudagraph_max_bs, scheduler_config.max_num_seqs + ) + self.decode_bc_pre_scratch: torch.Tensor = torch.empty( + ( + bc_scratch_bs, + bc_ngroups, + self.replayssm_buffer_len, + ), + dtype=torch.float32, + device=device, + ) + else: + self.decode_bc_pre_scratch = None self._init_reorder_batch_threshold(1, self.use_spec_decode) if self.use_spec_decode: @@ -414,6 +452,8 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): has_prior_state = seq_lens_cpu > 1 prefill_to_decode = single_token_prefill_rows & has_prior_state if torch.any(prefill_to_decode).item(): + # ReplaySSM handles these rows as single-token flushes (see the + # write-position derivation below), same as the baseline decode path. is_prefilling = is_prefilling.clone() is_prefilling[prefill_to_decode] = False common_attn_metadata = common_attn_metadata.replace( @@ -444,6 +484,8 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): # for causal_conv1d nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None + write_pos_d = None + is_flush_d = None if self.vllm_config.cache_config.mamba_cache_mode == "all": num_computed_tokens = common_attn_metadata.compute_num_computed_tokens() @@ -530,6 +572,79 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): num_reqs - num_prefills : num_reqs ] + if self.use_replayssm and num_decodes > 0: + decode_base_cpu = common_attn_metadata.replayssm_decode_base_cpu + num_computed_tokens_cpu = common_attn_metadata._num_computed_tokens_cpu + if decode_base_cpu is None or num_computed_tokens_cpu is None: + raise ValueError( + "--use-replayssm requires CPU decode-base and " + "computed-token counts to derive decode write positions" + ) + num_computed_d = num_computed_tokens_cpu[:num_decodes] + decode_base_d = decode_base_cpu[:num_decodes] + align_mode = self.vllm_config.cache_config.mamba_cache_mode == "align" + block_size = self.kv_cache_spec.block_size + if align_mode: + # After a boundary the align copy leaves an exact checkpoint at + # the block start and the new block's ring restarts empty, so + # re-anchor there; max() keeps the prompt-end anchor for the + # first (partial) block. + effective_base = torch.maximum( + decode_base_d, (num_computed_d // block_size) * block_size + ) + else: + effective_base = decode_base_d + # write_pos counts decode steps since the ring's last full-state + # write (the anchor), so a resumed request re-anchors correctly. + decode_steps_cpu = num_computed_d - effective_base + query_lens_cpu = ( + common_attn_metadata.query_start_loc_cpu[1 : num_decodes + 1] + - common_attn_metadata.query_start_loc_cpu[:num_decodes] + ) + valid_decode_rows = query_lens_cpu > 0 + # A single-token prefill row replayed as decode (query_len==1 with + # prior state) has decode_steps < 0; force it to a one-token flush + # (write_pos=0, is_flush=1). The flush branch reads an empty history + # window, so it applies exactly one recurrence step off the checkpoint + # -- identical to the baseline decode kernel for that row. The split + # (treat_short_extends_as_decodes=False) admits only such rows here. + leftover_prompt = valid_decode_rows & (decode_steps_cpu < 0) + decode_steps_cpu = torch.where( + valid_decode_rows & ~leftover_prompt, + decode_steps_cpu, + torch.zeros_like(decode_steps_cpu), + ) + write_pos_cpu = torch.remainder(decode_steps_cpu, self.replayssm_buffer_len) + is_flush_cpu = ( + write_pos_cpu == self.replayssm_buffer_len - 1 + ) | leftover_prompt + if align_mode: + # Force a flush on the step completing a mamba block so the exact + # boundary state is materialized for prefix caching. + is_flush_cpu = is_flush_cpu | ( + valid_decode_rows + & ((num_computed_d + query_lens_cpu) % block_size == 0) + ) + is_flush_cpu = is_flush_cpu.to(torch.int8) + write_pos_d = async_tensor_h2d( + write_pos_cpu.to(torch.int32).tolist(), + dtype=torch.int32, + device=common_attn_metadata.query_start_loc.device, + ) + is_flush_d = async_tensor_h2d( + is_flush_cpu.tolist(), + dtype=torch.int8, + device=common_attn_metadata.query_start_loc.device, + ) + + bc_pre_scratch = None + if ( + self.use_replayssm + and self.decode_bc_pre_scratch is not None + and num_decodes > 0 + ): + bc_pre_scratch = self.decode_bc_pre_scratch[:num_decodes] + metadata = self.metadata_cls( num_prefills=num_prefills, num_prefill_tokens=num_prefill_tokens, @@ -539,6 +654,9 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): has_initial_states_p=has_initial_states_p, state_indices_tensor_p=state_indices_tensor_p, state_indices_tensor_d=state_indices_tensor_d, + write_pos_d=write_pos_d, + is_flush_d=is_flush_d, + bc_pre_scratch=bc_pre_scratch, num_accepted_tokens=num_accepted_tokens, query_start_loc_d=query_start_loc_d, block_idx_last_scheduled_token=block_idx_last_scheduled_token, @@ -573,6 +691,9 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): block_idx_last_scheduled_token_prev_step = ( metadata.block_idx_last_scheduled_token_prev_step ) + write_pos_d = metadata.write_pos_d + is_flush_d = metadata.is_flush_d + bc_pre_scratch = metadata.bc_pre_scratch if ( metadata.num_prefills == 0 and metadata.num_decodes <= self.decode_cudagraph_max_bs @@ -634,11 +755,34 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): ) block_idx_last_scheduled_token_prev_step[metadata.num_decodes :] = 0 + if self.use_replayssm: + assert write_pos_d is not None + assert is_flush_d is not None + self.decode_write_pos_d[: metadata.num_decodes].copy_( + write_pos_d[: metadata.num_decodes], + non_blocking=True, + ) + write_pos_d = self.decode_write_pos_d[:padded_bs] + write_pos_d[metadata.num_decodes :] = 0 + + self.decode_is_flush_d[: metadata.num_decodes].copy_( + is_flush_d[: metadata.num_decodes], + non_blocking=True, + ) + is_flush_d = self.decode_is_flush_d[:padded_bs] + is_flush_d[metadata.num_decodes :] = 0 + + if self.decode_bc_pre_scratch is not None: + bc_pre_scratch = self.decode_bc_pre_scratch[:padded_bs] + return replace( metadata, state_indices_tensor_d=state_indices_tensor_d, query_start_loc_d=query_start_loc_d, num_accepted_tokens=num_accepted_tokens, + write_pos_d=write_pos_d, + is_flush_d=is_flush_d, + bc_pre_scratch=bc_pre_scratch, block_idx_last_scheduled_token=block_idx_last_scheduled_token, block_idx_last_computed_token=block_idx_last_computed_token, block_idx_last_scheduled_token_prev_step=( diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py index baa502ee811..0b805fca765 100644 --- a/vllm/v1/worker/gpu_input_batch.py +++ b/vllm/v1/worker/gpu_input_batch.py @@ -106,6 +106,7 @@ class InputBatch: is_pooling_model: bool = False, cp_kv_cache_interleave_size: int = 1, reasoning_config: ReasoningConfig | None = None, + use_replayssm: bool = False, slot_mapping_modes: list[SlotMappingMode] | None = None, ): self.thinking_budget_state_holder = maybe_create_thinking_budget_state_holder( @@ -170,6 +171,17 @@ class InputBatch: ) self.num_computed_tokens_cpu = self.num_computed_tokens_cpu_tensor.numpy() + # Mamba2 ReplaySSM decode ring origin (num_computed at each request's + # last full-state write); populated only when the feature is on. + self.use_replayssm = use_replayssm + self.replayssm_decode_base_cpu_tensor = torch.zeros( + (max_num_reqs,), + device="cpu", + dtype=torch.int32, + pin_memory=PIN_MEMORY, + ) + self.replayssm_decode_base = self.replayssm_decode_base_cpu_tensor.numpy() + # Block table. self.block_table = MultiGroupBlockTable( max_num_reqs=max_num_reqs, @@ -377,6 +389,11 @@ class InputBatch: # Number of tokens without spec decode tokens. self.num_tokens_no_spec[req_index] = request.num_tokens + if self.use_replayssm: + # Ring origin = full context at (re)admission (prompt + any resumed + # output), so a resumed request re-anchors past the prompt. + self.replayssm_decode_base[req_index] = request.num_tokens + self.num_computed_tokens_cpu[req_index] = request.num_computed_tokens self.block_table.add_row(request.block_ids, req_index) @@ -597,6 +614,11 @@ class InputBatch: self.num_prompt_tokens[i2], self.num_prompt_tokens[i1], ) + if self.use_replayssm: + self.replayssm_decode_base[i1], self.replayssm_decode_base[i2] = ( + self.replayssm_decode_base[i2], + self.replayssm_decode_base[i1], + ) self.num_computed_tokens_cpu[i1], self.num_computed_tokens_cpu[i2] = ( self.num_computed_tokens_cpu[i2], self.num_computed_tokens_cpu[i1], @@ -754,6 +776,10 @@ class InputBatch: last_req_index ] self.num_prompt_tokens[empty_index] = self.num_prompt_tokens[last_req_index] + if self.use_replayssm: + self.replayssm_decode_base[empty_index] = self.replayssm_decode_base[ + last_req_index + ] self.num_computed_tokens_cpu[empty_index] = self.num_computed_tokens_cpu[ last_req_index ] diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 95221c86931..dd31ddab672 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -727,6 +727,7 @@ class GPUModelRunner( is_pooling_model=self.is_pooling_model, cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size, reasoning_config=self.vllm_config.reasoning_config, + use_replayssm=self.cache_config.use_replayssm, ) # Separate cuda stream for overlapping transfer of sampled token ids from @@ -2413,6 +2414,12 @@ class GPUModelRunner( if self.model_config.rswa_window is not None: rswa_prefix_lens = num_prompt_tokens_cpu + replayssm_decode_base_cpu = None + if self.cache_config.use_replayssm: + replayssm_decode_base_cpu = ( + self.input_batch.replayssm_decode_base_cpu_tensor[:num_reqs_padded] + ) + cm_base = CommonAttentionMetadata( query_start_loc=self.query_start_loc.gpu[: num_reqs_padded + 1], query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs_padded + 1], @@ -2420,6 +2427,7 @@ class GPUModelRunner( _seq_lens_cpu=seq_lens_cpu, _num_computed_tokens_cpu=num_computed_tokens_cpu, seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, + replayssm_decode_base_cpu=replayssm_decode_base_cpu, num_reqs=num_reqs_padded, num_actual_tokens=num_tokens_padded, max_query_len=max_query_len, @@ -7262,6 +7270,7 @@ class GPUModelRunner( is_pooling_model=self.is_pooling_model, cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size, reasoning_config=self.vllm_config.reasoning_config, + use_replayssm=self.cache_config.use_replayssm, slot_mapping_modes=slot_mapping_modes, ) From 41798069f318b88334efdb7778190ce48c1084a4 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Fri, 24 Jul 2026 13:41:08 -0500 Subject: [PATCH 006/185] [CI][AMD] Deprecate DinD for MI355 tests (#49257) Signed-off-by: Andreas Karatzas --- .buildkite/test-amd.yaml | 39 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index d3cdbc388f7..b8046fda9f5 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -3048,6 +3048,7 @@ steps: - label: Attention Benchmarks Smoke Test (B200-MI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_2 num_gpus: 2 working_dir: "/vllm-workspace/" @@ -3064,6 +3065,7 @@ steps: - label: Distributed Tests (2xH100-2xMI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_2 num_gpus: 2 optional: true @@ -3108,6 +3110,7 @@ steps: - label: Entrypoints Integration (API Server) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 optional: true fast_check: true @@ -3125,6 +3128,7 @@ steps: - label: Entrypoints Integration (API Server OpenAI - Part 1) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 fast_check: true optional: true @@ -3140,6 +3144,7 @@ steps: - label: Entrypoints Integration (API Server OpenAI - Part 2) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 fast_check: true optional: true @@ -3156,6 +3161,7 @@ steps: - label: Entrypoints Integration (API Server Generate) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 fast_check: true optional: true @@ -3176,6 +3182,7 @@ steps: - label: Entrypoints Integration (Speech to Text) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355] + dind: false agent_pool: mi355_1 fast_check: true working_dir: "/vllm-workspace/tests" @@ -3189,6 +3196,7 @@ steps: - label: Entrypoints Integration (Multimodal) timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355] + dind: false agent_pool: mi355_1 fast_check: true working_dir: "/vllm-workspace/tests" @@ -3202,6 +3210,7 @@ steps: - label: Entrypoints Integration (Pooling) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 fast_check: true working_dir: "/vllm-workspace/tests" @@ -3217,6 +3226,7 @@ steps: - label: GPQA Eval (GPT-OSS) (2xB200-2xMI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_2 num_gpus: 2 optional: true @@ -3239,6 +3249,7 @@ steps: - label: LM Eval Qwen3-5 Models (B200-MI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_2 num_gpus: 2 optional: true @@ -3261,6 +3272,7 @@ steps: - label: LM Eval Small Models (2xB200-2xMI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_2 num_gpus: 2 optional: true @@ -3280,6 +3292,7 @@ steps: - label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200-MI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_2 num_gpus: 2 working_dir: "/vllm-workspace" @@ -3300,6 +3313,7 @@ steps: - label: LM Eval Large Models (4xH100-4xMI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_4 num_gpus: 4 optional: true @@ -3322,6 +3336,7 @@ steps: - label: Examples # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 working_dir: "/vllm-workspace/examples" source_file_dependencies: @@ -3357,6 +3372,7 @@ steps: - label: Kernels (B200-MI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 working_dir: "/vllm-workspace/" source_file_dependencies: @@ -3382,6 +3398,7 @@ steps: - label: Kernels Attention Test %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 parallelism: 2 working_dir: "/vllm-workspace/tests" @@ -3399,6 +3416,7 @@ steps: - label: Kernels MoE Test %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 parallelism: 5 working_dir: "/vllm-workspace/tests" @@ -3419,6 +3437,7 @@ steps: - label: Kernels Quantization Test %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 parallelism: 2 working_dir: "/vllm-workspace/tests" @@ -3436,6 +3455,7 @@ steps: - label: Kernels FP8 MoE Test (2xH100-2xMI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_2 num_gpus: 2 working_dir: "/vllm-workspace/tests" @@ -3455,6 +3475,7 @@ steps: - label: Language Models Test (Extended Generation) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -3468,6 +3489,7 @@ steps: - label: Language Models Test (Extended Pooling) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" @@ -3480,6 +3502,7 @@ steps: - label: Language Models Test (PPL) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" @@ -3508,6 +3531,7 @@ steps: - label: Language Models Tests (Standard) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -3522,6 +3546,7 @@ steps: - label: Multi-Modal Models (Extended Generation 1) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" @@ -3536,6 +3561,7 @@ steps: - label: Multi-Modal Models (Extended Generation 3) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" @@ -3548,6 +3574,7 @@ steps: - label: Multi-Modal Models (Extended Pooling) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" @@ -3560,6 +3587,7 @@ steps: - label: "Multi-Modal Models (Standard) 1: qwen2" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" @@ -3573,6 +3601,7 @@ steps: - label: "Multi-Modal Models (Standard) 4: other + whisper" # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" @@ -3589,6 +3618,7 @@ steps: - label: Quantized Models Test # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -3605,6 +3635,7 @@ steps: - label: Quantization # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -3621,6 +3652,7 @@ steps: # - label: Quantized MoE Test (B200-MI355) # TBD # timeout_in_minutes: 180 # mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false # agent_pool: mi355_1 # working_dir: "/vllm-workspace/" # source_file_dependencies: @@ -3649,6 +3681,7 @@ steps: - label: V1 attention (B200-MI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -3665,6 +3698,7 @@ steps: - label: V1 Core + KV + Metrics # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" @@ -3691,6 +3725,7 @@ steps: - label: V1 Sample + Logits # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" @@ -3711,6 +3746,7 @@ steps: - label: V1 Spec Decode # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -3724,6 +3760,7 @@ steps: - label: Weight Loading Multiple GPU # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_2 num_gpus: 2 working_dir: "/vllm-workspace/tests" @@ -3736,6 +3773,7 @@ steps: - label: Weight Loading Multiple GPU - Large Models # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_2 working_dir: "/vllm-workspace/tests" num_gpus: 2 @@ -3751,6 +3789,7 @@ steps: - label: Regression # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false agent_pool: mi355_1 optional: true working_dir: "/vllm-workspace/tests" From 8c13ee5735f54fe7c653cf694be21bb462a99fed Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith Date: Fri, 24 Jul 2026 14:59:49 -0400 Subject: [PATCH 007/185] Add `sm_107` for Rubin (#49387) Signed-off-by: Tyler Michael Smith Co-authored-by: OpenAI Codex --- CMakeLists.txt | 19 ++++++++++++------- cmake/external_projects/deepgemm.cmake | 3 +++ cmake/external_projects/flashmla.cmake | 4 +++- cmake/external_projects/qutlass.cmake | 6 +++++- cmake/utils.cmake | 20 +++++++++++++++----- tests/test_cmake_utils.py | 23 +++++++++++++++++++++++ 6 files changed, 61 insertions(+), 14 deletions(-) create mode 100644 tests/test_cmake_utils.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 31984665660..d91113bd1c2 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -114,6 +114,11 @@ find_package(Torch REQUIRED) # Supported NVIDIA architectures. # This check must happen after find_package(Torch) because that's when CMAKE_CUDA_COMPILER_VERSION gets defined if(DEFINED CMAKE_CUDA_COMPILER_VERSION AND + CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.4) + # Rubin (10.7) can run SM100 family code, but CUDA 13.4 also supports + # targeting it directly. + set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.7;11.0;12.0") +elseif(DEFINED CMAKE_CUDA_COMPILER_VERSION AND CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0) # starting from CUDA 12.9 and Blackwell (10.0), we use family-specific targets (10.0f, 12.0f, etc) # to support the whole generation without specifying all sub-architectures @@ -420,7 +425,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS - "9.0a;10.0f;10.1f;10.3f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}") + "9.0a;10.0f;10.1f;10.3f;10.7f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}") else() cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}") @@ -695,7 +700,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later) if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;10.7f;11.0f;12.0f" "${CUDA_ARCHS}") else() cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}") endif() @@ -832,7 +837,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # The cutlass_scaled_mm kernels for Blackwell SM100 (c3x, i.e. CUTLASS 3.x) # require CUDA 12.8 or later if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}") else() cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}") endif() @@ -916,7 +921,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}") else() cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}") endif() @@ -941,7 +946,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # moe_data.cu is used by all CUTLASS MoE kernels. if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;10.7f;11.0f;12.0f" "${CUDA_ARCHS}") else() cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}") endif() @@ -998,7 +1003,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # SM10x/11x FP4 kernels. MXFP4 experts quantization is currently compiled # only in this block; SM12x has separate NVFP4 matmul/MoE kernels above. if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}") else() cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}") endif() @@ -1064,7 +1069,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # Runtime dispatch is gated in # vllm/v1/attention/backends/mla/cutlass_mla.py. if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(MLA_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}") else() cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}") endif() diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index 38d218d00ac..8c2f98d8fe6 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -68,6 +68,9 @@ endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4) + list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.7f") + endif() else() list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a") endif() diff --git a/cmake/external_projects/flashmla.cmake b/cmake/external_projects/flashmla.cmake index 56f7a83f678..4d8c2d74c20 100644 --- a/cmake/external_projects/flashmla.cmake +++ b/cmake/external_projects/flashmla.cmake @@ -60,6 +60,9 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) # CUDA 12.9 has introduced "Family-Specific Architecture Features" # this supports all compute_10x family list(APPEND SUPPORT_ARCHS "10.0f") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4) + list(APPEND SUPPORT_ARCHS "10.7f") + endif() elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) list(APPEND SUPPORT_ARCHS "10.0a") endif() @@ -188,4 +191,3 @@ else() add_custom_target(_flashmla_C) add_custom_target(_flashmla_extension_C) endif() - diff --git a/cmake/external_projects/qutlass.cmake b/cmake/external_projects/qutlass.cmake index 44b82f6be9e..cd3665c889e 100644 --- a/cmake/external_projects/qutlass.cmake +++ b/cmake/external_projects/qutlass.cmake @@ -55,7 +55,11 @@ message(STATUS "[QUTLASS] QuTLASS is available at ${qutlass_SOURCE_DIR}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0f" "${CUDA_ARCHS}") - cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4) + cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f;10.7f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}") + endif() else() cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}") cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0a;10.3a" "${CUDA_ARCHS}") diff --git a/cmake/utils.cmake b/cmake/utils.cmake index e3e766541df..14a94eebb22 100644 --- a/cmake/utils.cmake +++ b/cmake/utils.cmake @@ -396,14 +396,24 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR # match — e.g. SRC="12.0f" matches TGT="12.1a" since SM121 is in the SM12x # family. The output uses TGT's value to preserve the user's compilation flags. set(_CUDA_ARCHS) + # Resolve exact base matches before family fallbacks so a generic entry such + # as 10.0f cannot consume a 10.7 target that has a 10.7f source entry. + foreach(_arch ${_SRC_CUDA_ARCHS}) + if(_arch MATCHES "[af]$") + string(REGEX REPLACE "[af]$" "" _base "${_arch}") + if("${_base}" IN_LIST _TGT_CUDA_ARCHS) + list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}") + list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}") + list(APPEND _CUDA_ARCHS "${_arch}") + endif() + endif() + endforeach() + foreach(_arch ${_SRC_CUDA_ARCHS}) if(_arch MATCHES "[af]$") list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}") string(REGEX REPLACE "[af]$" "" _base "${_arch}") - if ("${_base}" IN_LIST TGT_CUDA_ARCHS) - list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}") - list(APPEND _CUDA_ARCHS "${_arch}") - elseif("${_base}a" IN_LIST _TGT_CUDA_ARCHS) + if("${_base}a" IN_LIST _TGT_CUDA_ARCHS) list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}a") list(APPEND _CUDA_ARCHS "${_base}a") elseif("${_base}f" IN_LIST _TGT_CUDA_ARCHS) @@ -487,7 +497,7 @@ endfunction() function(cuda_archs_sm90plus OUT_CUDA_ARCHS TGT_CUDA_ARCHS) if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(_archs "9.0a;10.0f;11.0f;12.0f" "${TGT_CUDA_ARCHS}") + cuda_archs_loose_intersection(_archs "9.0a;10.0f;10.7f;11.0f;12.0f" "${TGT_CUDA_ARCHS}") else() cuda_archs_loose_intersection(_archs "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${TGT_CUDA_ARCHS}") endif() diff --git a/tests/test_cmake_utils.py b/tests/test_cmake_utils.py new file mode 100644 index 00000000000..227ec231eb2 --- /dev/null +++ b/tests/test_cmake_utils.py @@ -0,0 +1,23 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import subprocess +from pathlib import Path + + +def test_exact_family_arch_precedes_generic_family_fallback(tmp_path: Path): + repo_root = Path(__file__).parents[1] + script = tmp_path / "test_cuda_archs.cmake" + script.write_text( + f""" +cmake_minimum_required(VERSION 3.26) +include("{repo_root / "cmake" / "utils.cmake"}") +cuda_archs_loose_intersection( + actual "10.0f;10.7f" "10.7") +if(NOT "${{actual}}" STREQUAL "10.7f") + message(FATAL_ERROR "Expected 10.7f, got '${{actual}}'") +endif() +""" + ) + + subprocess.run(["cmake", "-P", script], check=True) From 9863102ed94a91255cd6b0924027b3d3b985e913 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Fri, 24 Jul 2026 14:32:40 -0500 Subject: [PATCH 008/185] [CI] Reuse loaded config for cached tokenizer (#49509) Signed-off-by: Andreas Karatzas --- tests/tokenizers_/test_registry.py | 2 ++ vllm/tokenizers/registry.py | 7 +++++++ 2 files changed, 9 insertions(+) diff --git a/tests/tokenizers_/test_registry.py b/tests/tokenizers_/test_registry.py index 9635e9963b5..0a47426dc97 100644 --- a/tests/tokenizers_/test_registry.py +++ b/tests/tokenizers_/test_registry.py @@ -109,6 +109,8 @@ def test_cached_tokenizer_from_config_registers_local_config(tmp_path: Path): try: def fake_from_pretrained(path_or_repo_id: str, *args, **kwargs): + passed_config = kwargs.pop("config") + assert isinstance(passed_config, Qwen3_5MoeConfig) loaded_config = AutoConfig.from_pretrained( path_or_repo_id, trust_remote_code=False, diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py index e6c12ccc3bc..82df4339e85 100644 --- a/vllm/tokenizers/registry.py +++ b/vllm/tokenizers/registry.py @@ -18,6 +18,7 @@ from vllm.transformers_utils.repo_utils import ( ) from vllm.utils.import_utils import resolve_obj_by_qualname +from .hf import CachedHfTokenizer from .protocol import TokenizerLike if TYPE_CHECKING: @@ -233,6 +234,12 @@ def get_tokenizer( else: tokenizer_cls_ = tokenizer_cls + if config is not None and tokenizer_cls_ is CachedHfTokenizer: + # AutoTokenizer otherwise reloads config.json internally. Reuse the + # config that get_config just loaded successfully so a concurrent Hub + # cache refresh cannot invalidate the file between the two reads. + kwargs.setdefault("config", config) + tokenizer = tokenizer_cls_.from_pretrained(tokenizer_name, *args, **kwargs) if model_type in _MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: from vllm.tokenizers.hf import get_cached_tokenizer From 7e51939e25c34e6527a4971345c95dbefb029fc9 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Fri, 24 Jul 2026 14:33:53 -0500 Subject: [PATCH 009/185] [CI] Avoid unnecessary Hugging Face metadata requests (#49508) Signed-off-by: Andreas Karatzas --- vllm/config/model.py | 9 ++++++--- vllm/model_executor/model_loader/default_loader.py | 2 +- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/vllm/config/model.py b/vllm/config/model.py index 0e7621c2b17..1f62351c526 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -579,9 +579,9 @@ class ModelConfig: self.hf_text_config, "attention_chunk_size", None ) self.encoder_config = self._get_encoder_config() - self.hf_image_processor_config = get_hf_image_processor_config( - self.model, hf_token=self.hf_token, revision=self.revision - ) + # Image-processor metadata is only consumed by multimodal models. + # Probing it for text-only models causes avoidable Hub requests. + self.hf_image_processor_config: dict[str, Any] = {} architectures = self.architectures registry = self.registry @@ -703,6 +703,9 @@ class ModelConfig: # Init multimodal config if needed if self._model_info.supports_multimodal: + self.hf_image_processor_config = get_hf_image_processor_config( + self.model, hf_token=self.hf_token, revision=self.revision + ) if ( mm_encoder_tp_mode == "data" and not self._model_info.supports_multimodal_encoder_tp_data diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py index 3ea76f4d9b3..039a577f80d 100644 --- a/vllm/model_executor/model_loader/default_loader.py +++ b/vllm/model_executor/model_loader/default_loader.py @@ -220,7 +220,7 @@ class DefaultModelLoader(BaseModelLoader): # safetensors file. Using both breaks. # Here, we download the `model.safetensors.index.json` and filter # any files not found in the index. - if not is_local: + if not is_local and len(hf_weights_files) > 1: download_safetensors_index_file_from_hf( model_name_or_path, index_file, From c064fa52b6425b17dbc51fea2eaf4aac0f30863f Mon Sep 17 00:00:00 2001 From: Aarushi Jain <142941703+aarushjain29@users.noreply.github.com> Date: Fri, 24 Jul 2026 14:35:16 -0500 Subject: [PATCH 010/185] Fix GLM-4.1V video placeholder token ID handling. (#49484) Signed-off-by: aarushjain29 Co-authored-by: Andreas Karatzas --- .../multimodal/processing/test_glm4_1v.py | 5 +- vllm/model_executor/models/glm4_1v.py | 62 +++++++------------ 2 files changed, 24 insertions(+), 43 deletions(-) diff --git a/tests/models/multimodal/processing/test_glm4_1v.py b/tests/models/multimodal/processing/test_glm4_1v.py index 0cafa261a34..8a777832826 100644 --- a/tests/models/multimodal/processing/test_glm4_1v.py +++ b/tests/models/multimodal/processing/test_glm4_1v.py @@ -93,7 +93,6 @@ def test_processor_override( limit_mm_per_prompt={"video": 1}, ) processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) - tokenizer = processor.info.get_tokenizer() hf_processor_mm_kwargs = {"fps": fps} # Build the image str / prompt based on the number of images we pass @@ -112,8 +111,8 @@ def test_processor_override( # Ensure we have the right number of placeholders per num_crops size hf_processor = processor.info.get_hf_processor(**hf_processor_mm_kwargs) - video_token_id = tokenizer.convert_tokens_to_ids(hf_processor.video_token) - video_tok_count = processed_inputs["prompt_token_ids"].count(video_token_id) + image_token_id = hf_processor.image_token_id + video_tok_count = processed_inputs["prompt_token_ids"].count(image_token_id) grid_t, _, _ = processed_inputs["mm_kwargs"].get_data()["video_grid_thw"][0] assert grid_t == expected_grid_t diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py index 810d9de87b4..787e53e8df2 100644 --- a/vllm/model_executor/models/glm4_1v.py +++ b/vllm/model_executor/models/glm4_1v.py @@ -1403,6 +1403,11 @@ class Glm4vProcessingInfo(BaseProcessingInfo): timestamps_list = full_second_idxs[::2] return list(timestamps_list) + def _get_video_frame_embed_token_id(self, hf_processor: object) -> int: + if isinstance(hf_processor, Glm4vProcessor) or TRANSFORMERS_WITH_GA: + return hf_processor.image_token_id + return hf_processor.video_token_id + def _construct_video_placeholder( self, video_array: np.ndarray, @@ -1440,13 +1445,7 @@ class Glm4vProcessingInfo(BaseProcessingInfo): num_tokens_per_frame = int(H * W) // merge_length placeholder = [] placeholder.append(bov_token_id) - # Glm46VProcessor uses image_token_id for video frame embeddings; - # Glm4vProcessor uses video_token_id. - frame_embed_token_id = ( - hf_processor.video_token_id - if isinstance(hf_processor, Glm4vProcessor) or not TRANSFORMERS_WITH_GA - else hf_processor.image_token_id - ) + frame_embed_token_id = self._get_video_frame_embed_token_id(hf_processor) for frame_idx in frames_idx_token: placeholder.append(boi_token_id) placeholder.extend([frame_embed_token_id] * num_tokens_per_frame) @@ -1604,11 +1603,6 @@ class Glm4vMultiModalProcessor(BaseMultiModalProcessor[Glm4vProcessingInfo]): processor = self.info.get_hf_processor(**mm_kwargs) - # Glm46VProcessor and GLMGA handle image/video placeholders together - # via the direct path. Only Glm4vProcessor (GLM-4.1V) needs the - # split-video path because it uses image_token_id as the video - # placeholder. The direct path requires transformers >= 5.5.0 - # (Glm46VProcessor / GlmgaVideoProcessor support). use_direct_path = ( not isinstance(processor, Glm4vProcessor) and TRANSFORMERS_WITH_GA ) @@ -1630,6 +1624,8 @@ class Glm4vMultiModalProcessor(BaseMultiModalProcessor[Glm4vProcessingInfo]): ): video_grid_thw_lst = [] pixel_values_videos_lst = [] + frame_embed_token_id = self.info._get_video_frame_embed_token_id(processor) + swap_video_frame_tokens = frame_embed_token_id == processor.image_token_id for item in mm_data.pop("videos", []): video_array, metadata = item @@ -1650,9 +1646,10 @@ class Glm4vMultiModalProcessor(BaseMultiModalProcessor[Glm4vProcessingInfo]): tok_kwargs=tok_kwargs, ) input_ids = video_outputs.pop("input_ids") - input_ids[input_ids == processor.image_token_id] = ( - processor.video_token_id - ) + if swap_video_frame_tokens: + input_ids[input_ids == processor.image_token_id] = ( + processor.video_token_id + ) video_placeholder = processor.tokenizer.batch_decode(input_ids)[0] prompt = prompt.replace( "<|begin_of_video|><|video|><|end_of_video|>", @@ -1668,6 +1665,7 @@ class Glm4vMultiModalProcessor(BaseMultiModalProcessor[Glm4vProcessingInfo]): ) else: video_outputs = dict() + swap_video_frame_tokens = False processed_outputs = super()._call_hf_processor( prompt=prompt, @@ -1675,6 +1673,10 @@ class Glm4vMultiModalProcessor(BaseMultiModalProcessor[Glm4vProcessingInfo]): mm_kwargs=mm_kwargs, tok_kwargs=tok_kwargs, ) + if swap_video_frame_tokens: + input_ids = processed_outputs["input_ids"] + input_ids[input_ids == processor.video_token_id] = processor.image_token_id + processed_outputs["input_ids"] = input_ids combined_outputs = dict( processed_outputs, **video_outputs, @@ -1701,7 +1703,7 @@ class Glm4vMultiModalProcessor(BaseMultiModalProcessor[Glm4vProcessingInfo]): merge_length = image_processor.merge_size**2 - def get_image_replacement_glm4v(item_idx: int): + def get_image_replacement(item_idx: int): out_item = out_mm_kwargs["image"][item_idx] grid_thw = out_item["image_grid_thw"].data assert isinstance(grid_thw, torch.Tensor) @@ -1709,7 +1711,7 @@ class Glm4vMultiModalProcessor(BaseMultiModalProcessor[Glm4vProcessingInfo]): num_tokens = int(grid_thw.prod()) // merge_length return [hf_processor.image_token_id] * num_tokens - def get_video_replacement_glm4v(item_idx: int): + def get_video_replacement(item_idx: int): out_item = out_mm_kwargs["video"][item_idx] grid_thw = out_item["video_grid_thw"].data assert isinstance(grid_thw, torch.Tensor) @@ -1720,39 +1722,19 @@ class Glm4vMultiModalProcessor(BaseMultiModalProcessor[Glm4vProcessingInfo]): ) return PromptUpdateDetails.select_token_id( placeholder, - embed_token_id=hf_processor.video_token_id, + embed_token_id=self.info._get_video_frame_embed_token_id(hf_processor), ) - def get_video_replacement_glm46v(item_idx: int): - out_item = out_mm_kwargs["video"][item_idx] - grid_thw = out_item["video_grid_thw"].data - assert isinstance(grid_thw, torch.Tensor) - - video, metadata = mm_items["video"][item_idx] - placeholder = self.info._construct_video_placeholder( - video, metadata, grid_thw - ) - return PromptUpdateDetails.select_token_id( - placeholder, - embed_token_id=hf_processor.image_token_id, - ) - - is_glm46v = not isinstance(hf_processor, Glm4vProcessor) - return [ PromptReplacement( modality="image", target=hf_processor.image_token, - replacement=get_image_replacement_glm4v, + replacement=get_image_replacement, ), PromptReplacement( modality="video", target="<|begin_of_video|><|video|><|end_of_video|>", - replacement=( - get_video_replacement_glm46v - if is_glm46v and TRANSFORMERS_WITH_GA - else get_video_replacement_glm4v - ), + replacement=get_video_replacement, ), ] From e222c33f2f568b54b41d2ebfc44f6c85ea072737 Mon Sep 17 00:00:00 2001 From: djramic Date: Fri, 24 Jul 2026 21:57:49 +0200 Subject: [PATCH 011/185] [Bugfix] Register axk1 config to fix A.X-K1 init (#49727) Signed-off-by: Djordje Ramic --- vllm/transformers_utils/config.py | 1 + vllm/transformers_utils/configs/AXK1.py | 2 +- vllm/transformers_utils/model_arch_config_convertor.py | 4 ++-- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index fd76550d664..f900eedf8d3 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -72,6 +72,7 @@ class LazyConfigDict(dict): _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict( afmoe="AfmoeConfig", arctic="ArcticConfig", + axk1="AXK1Config", bagel="BagelConfig", umm="CheersConfig", chatglm="ChatGLMConfig", diff --git a/vllm/transformers_utils/configs/AXK1.py b/vllm/transformers_utils/configs/AXK1.py index 5c19a37324b..09a3a0f46ae 100644 --- a/vllm/transformers_utils/configs/AXK1.py +++ b/vllm/transformers_utils/configs/AXK1.py @@ -114,7 +114,7 @@ class AXK1Config(PretrainedConfig): The dropout ratio for the attention probabilities. """ - model_type = "AXK1" + model_type = "axk1" keys_to_ignore_at_inference = ["past_key_values"] def __init__( diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index bd146dff7dc..70bb4caa535 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -261,7 +261,7 @@ class ModelArchConfigConvertorBase: if not hasattr(self.hf_text_config, "model_type"): return False elif self.hf_text_config.model_type in ( - "AXK1", + "axk1", "deepseek_v2", "deepseek_v3", "deepseek_v32", @@ -290,7 +290,7 @@ class ModelArchConfigConvertorBase: return ( self.hf_text_config.model.model_type in ( - "AXK1", + "axk1", "deepseek_v2", "deepseek_v3", "deepseek_v32", From 5d8e90a96616c4fe339ff0b0c2a2d470f6eb24bf Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith Date: Fri, 24 Jul 2026 16:00:02 -0400 Subject: [PATCH 012/185] [WideEP] Update NCCL to 2.30.7 to enable DeepEPv2 in the vllm/vllm-openai image (#45321) Signed-off-by: Tyler Michael Smith Signed-off-by: Tyler Michael Smith Signed-off-by: Tyler Michael Smith Co-authored-by: Claude Co-authored-by: Ilya Markov Co-authored-by: OpenAI Codex Co-authored-by: Codex --- .buildkite/test_areas/misc.yaml | 2 +- .buildkite/test_areas/model_runner_v2.yaml | 2 +- docker/Dockerfile | 33 +++- docker/versions.json | 5 +- docs/serving/expert_parallel_deployment.md | 5 + tests/distributed/test_mnnvl_alltoall.py | 12 +- tests/kernels/moe/parallel_utils.py | 48 ++++-- tests/kernels/moe/test_deepep_v2_moe.py | 151 ++++++++++++------ tools/ep_kernels/README.md | 32 ++++ .../device_communicators/all2all.py | 27 ++++ .../device_communicators/pynccl_allocator.py | 8 +- .../device_communicators/pynccl_wrapper.py | 29 ++++ vllm/utils/nccl.py | 62 +++++++ 13 files changed, 338 insertions(+), 78 deletions(-) diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 1c2000a9e56..cadbaed5b05 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -212,7 +212,7 @@ steps: - vllm/multimodal - examples/ commands: - - pip install tensorizer # for tensorizer test + - pip install --no-deps tensorizer # for tensorizer test # for basic - python3 basic/offline_inference/chat.py - python3 basic/offline_inference/generate.py --model facebook/opt-125m diff --git a/.buildkite/test_areas/model_runner_v2.yaml b/.buildkite/test_areas/model_runner_v2.yaml index 3601aeee117..bc0b56f4ca2 100644 --- a/.buildkite/test_areas/model_runner_v2.yaml +++ b/.buildkite/test_areas/model_runner_v2.yaml @@ -41,7 +41,7 @@ steps: commands: - set -x - export VLLM_USE_V2_MODEL_RUNNER=1 - - pip install tensorizer # for tensorizer test + - pip install --no-deps tensorizer # for tensorizer test - python3 basic/offline_inference/chat.py # for basic - python3 basic/offline_inference/generate.py --model facebook/opt-125m #- python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 # TODO diff --git a/docker/Dockerfile b/docker/Dockerfile index f8c018ef8ab..c5d8f006c9c 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -25,6 +25,10 @@ ARG CUDA_VERSION=13.0.3 ARG PYTHON_VERSION=3.12 ARG UBUNTU_VERSION=22.04 +# DeepEPv2 requires NCCL >= 2.30.4 (GIN backend). +# This version is only used for CUDA 13+ builds; CUDA 12 falls back to +# the default NCCL version shipped with the base image. +ARG NCCL_VERSION=2.30.7 # By parameterizing the base images, we allow third-party to use their own # base images. One use case is hermetic builds with base images stored in @@ -477,10 +481,17 @@ WORKDIR /workspace # Build DeepEP wheels COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.sh # Defaults moved here from tools/ep_kernels/install_python_libraries.sh for centralized version management -ARG DEEPEP_COMMIT_HASH=73b6ea4 +ARG DEEPEP_COMMIT_HASH=d4f41e4e93 ARG NVSHMEM_VER +ARG NCCL_VERSION RUN --mount=type=cache,target=/opt/uv/cache \ mkdir -p /tmp/ep_kernels_workspace/dist && \ + CUDA_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) && \ + if [ "$CUDA_MAJOR" -ge 13 ] && [ -n "$NCCL_VERSION" ]; then \ + echo "nvidia-nccl-cu${CUDA_MAJOR}==${NCCL_VERSION}" \ + > /tmp/nccl-override.txt && \ + export UV_OVERRIDE=/tmp/nccl-override.txt; \ + fi && \ export TORCH_CUDA_ARCH_LIST='9.0a 10.0a' && \ /tmp/install_python_libraries.sh \ --workspace /tmp/ep_kernels_workspace \ @@ -644,6 +655,7 @@ FROM ${FINAL_BASE_IMAGE} AS vllm-base ARG CUDA_VERSION ARG PYTHON_VERSION +ARG NCCL_VERSION ARG DEADSNAKES_MIRROR_URL ARG DEADSNAKES_GPGKEY_URL ARG GET_PIP_URL @@ -696,7 +708,6 @@ RUN apt-get update -y \ # Install CUDA development tools for runtime JIT compilation # (FlashInfer, DeepGEMM, EP kernels all require compilation at runtime) RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \ - CUDA_VERSION_SHORT=$(echo $CUDA_VERSION | cut -d. -f1,2) && \ apt-get update -y && \ apt-get install -y --no-install-recommends --allow-change-held-packages \ cuda-nvcc-${CUDA_VERSION_DASH} \ @@ -709,12 +720,6 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \ libnuma-dev \ # numactl CLI for NUMA binding at runtime numactl && \ - # Fixes nccl_allocator requiring nccl.h at runtime - # https://github.com/vllm-project/vllm/blob/1336a1ea244fa8bfd7e72751cabbdb5b68a0c11a/vllm/distributed/device_communicators/pynccl_allocator.py#L22 - # NCCL packages don't use the cuda-MAJOR-MINOR naming convention, - # so we pin the version to match our CUDA version - NCCL_VER=$(apt-cache madison libnccl-dev | grep "+cuda${CUDA_VERSION_SHORT}" | head -1 | awk -F'|' '{gsub(/^ +| +$/, "", $2); print $2}') && \ - apt-get install -y --no-install-recommends --allow-change-held-packages libnccl-dev=${NCCL_VER} libnccl2=${NCCL_VER} && \ rm -rf /var/lib/apt/lists/* # Install uv for faster pip installs @@ -734,6 +739,18 @@ RUN mkdir -p "${UV_PYTHON_INSTALL_DIR}" "${UV_CACHE_DIR}" \ && chgrp -R 0 /opt/uv \ && chmod -R g+rwX,a+rX /opt/uv +# DeepEPv2 GIN requires NCCL >= 2.30.4 at both compile and runtime. torch pins +# an older version as a transitive dep; this override forces uv to use our +# pinned version whenever nvidia-nccl-cu* is resolved. Empty on CUDA 12 (no-op). +RUN CUDA_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) && \ + if [ "$CUDA_MAJOR" -ge 13 ]; then \ + echo "nvidia-nccl-cu${CUDA_MAJOR}==${NCCL_VERSION}" \ + > /etc/uv-overrides.txt; \ + else \ + touch /etc/uv-overrides.txt; \ + fi +ENV UV_OVERRIDE=/etc/uv-overrides.txt + # ---------------------------------------------------------------------- # Non-root support (opt-in) # ---------------------------------------------------------------------- diff --git a/docker/versions.json b/docker/versions.json index cc145da93bf..783d80aace4 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -10,6 +10,9 @@ "UBUNTU_VERSION": { "default": "22.04" }, + "NCCL_VERSION": { + "default": "2.30.7" + }, "BUILD_BASE_IMAGE": { "default": "nvidia/cuda:13.0.3-devel-ubuntu22.04" }, @@ -56,7 +59,7 @@ "default": "cuda" }, "DEEPEP_COMMIT_HASH": { - "default": "73b6ea4" + "default": "d4f41e4e93" }, "GIT_REPO_CHECK": { "default": "0" diff --git a/docs/serving/expert_parallel_deployment.md b/docs/serving/expert_parallel_deployment.md index 9a8bb68bc11..b348c5dd965 100644 --- a/docs/serving/expert_parallel_deployment.md +++ b/docs/serving/expert_parallel_deployment.md @@ -12,6 +12,11 @@ Before using EP, you need to install the necessary dependencies. We are actively 2. **Install DeepGEMM library**: Follow the [official instructions](https://github.com/deepseek-ai/DeepGEMM#installation). 3. **For disaggregated serving**: Install `gdrcopy` by running the [`install_gdrcopy.sh`](../../tools/install_gdrcopy.sh) script (e.g., `install_gdrcopy.sh "${GDRCOPY_OS_VERSION}" "12.8" "x64"`). You can find available OS versions [here](https://developer.download.nvidia.com/compute/redist/gdrcopy/CUDA%2012.8/). +!!! note "NCCL version (CUDA 13+)" + The `deepep_v2` backend requires NCCL >= 2.30.4. PyTorch ships an older + NCCL, so you must upgrade it before building or running DeepEP. See the + [EP kernels guide](../../tools/ep_kernels) for instructions. + ### Backend Selection Guide vLLM provides multiple communication backends for EP. Use `--all2all-backend` to select one: diff --git a/tests/distributed/test_mnnvl_alltoall.py b/tests/distributed/test_mnnvl_alltoall.py index fb2d9bb832e..f9b76ea1b03 100644 --- a/tests/distributed/test_mnnvl_alltoall.py +++ b/tests/distributed/test_mnnvl_alltoall.py @@ -73,7 +73,10 @@ def _spawn_workers(worker_fn, world_size, *, dp_size=None): err_queue.close() err_queue.join_thread() if errors: - pytest.fail("Worker(s) failed:\n" + "\n---\n".join(errors)) + combined = "\n---\n".join(errors) + if "NCCL GIN" in combined: + pytest.skip("NCCL GIN not available on this system") + pytest.fail("Worker(s) failed:\n" + combined) def _run_worker(rank, world_size, port, worker_fn, dp_size, dp_port, err_queue): @@ -916,8 +919,11 @@ def _deepep_v2_lifecycle_worker(rank, world_size): DeepEPV2All2AllManager, ) - cpu_group = get_ep_group().cpu_group - manager = DeepEPV2All2AllManager(cpu_group) + ep_group = get_ep_group() + manager = DeepEPV2All2AllManager( + ep_group.cpu_group, + device_group=ep_group.device_group, + ) assert manager.rank == rank assert manager.world_size == world_size diff --git a/tests/kernels/moe/parallel_utils.py b/tests/kernels/moe/parallel_utils.py index bb2f9efc7c4..35a2f198f5d 100644 --- a/tests/kernels/moe/parallel_utils.py +++ b/tests/kernels/moe/parallel_utils.py @@ -36,6 +36,10 @@ if has_deep_ep_v2(): P = ParamSpec("P") +class GINNotAvailableError(RuntimeError): + pass + + @dataclasses.dataclass class ProcessGroupInfo: world_size: int @@ -96,19 +100,27 @@ def parallel_launch( **kwargs: P.kwargs, ) -> None: assert not kwargs - spawn( - _worker_parallel_launch, - args=( - world_size, - world_size, - 0, - f"tcp://{os.getenv('LOCALHOST', 'localhost')}:{get_open_port()}", - worker, + try: + spawn( + _worker_parallel_launch, + args=( + world_size, + world_size, + 0, + f"tcp://{os.getenv('LOCALHOST', 'localhost')}:{get_open_port()}", + worker, + ) + + args, + nprocs=world_size, + join=True, ) - + args, - nprocs=world_size, - join=True, - ) + except Exception as exc: + # pytest.skip cannot propagate directly through torch.multiprocessing. + if "GINNotAvailableError" in str(exc): + import pytest + + pytest.skip("NCCL GIN not available (no IBGDA-capable hardware)") + raise ## DeepEP specific utils @@ -225,6 +237,18 @@ def make_deepep_v2_a2a( ): import deep_ep + from vllm.utils.nccl import query_nccl_gin_type + + # ElasticBuffer can segfault when GIN is unavailable. Initialize the + # lazy communicator and reject unsupported systems before entering DeepEP. + probe = torch.zeros(1, device=pgi.device) + torch.distributed.all_reduce(probe, group=pg) + gin_type = query_nccl_gin_type(pg) + if gin_type is None: + raise RuntimeError("Failed to determine NCCL GIN support") + if gin_type == 0: + raise GINNotAvailableError("NCCL GIN not available") + buffer = deep_ep.ElasticBuffer( group=pg, num_max_tokens_per_rank=v2_args.max_tokens_per_rank, diff --git a/tests/kernels/moe/test_deepep_v2_moe.py b/tests/kernels/moe/test_deepep_v2_moe.py index 93b7c136605..e4cb5bff09e 100644 --- a/tests/kernels/moe/test_deepep_v2_moe.py +++ b/tests/kernels/moe/test_deepep_v2_moe.py @@ -14,6 +14,7 @@ from torch.distributed import ProcessGroup from tests.kernels.moe.utils import make_dummy_moe_config, make_test_weights from tests.kernels.utils import torch_experts from vllm.config import VllmConfig, set_current_vllm_config +from vllm.forward_context import set_forward_context from vllm.model_executor.layers.fused_moe import TritonExperts from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( @@ -36,6 +37,14 @@ requires_deep_ep_v2 = pytest.mark.skipif( ) +def assert_fp8_close(actual: torch.Tensor, expected: torch.Tensor) -> None: + close = torch.isclose(actual, expected, atol=2e-1, rtol=2e-1) + close_fraction = close.float().mean().item() + assert close_fraction > 0.99, ( + f"Only {close_fraction:.1%} of FP8 outputs are within tolerance" + ) + + @dataclasses.dataclass class TestConfig: dtype: torch.dtype @@ -50,6 +59,7 @@ class TestConfig: class TestTensors: rank_tokens: torch.Tensor rank_token_scales: torch.Tensor | None + intermediate_scales: torch.Tensor | None topk: torch.Tensor topk_weights: torch.Tensor config: TestConfig @@ -63,6 +73,12 @@ class TestTensors: rank_tokens = ( torch.randn((config.m, config.k), device="cuda", dtype=token_dtype) / 10 ) + if config.dtype == torch.float8_e4m3fn: + rank_token_scales = torch.tensor(1 / 448, device="cuda") + intermediate_scales = torch.tensor(8 / 448, device="cuda") + else: + rank_token_scales = None + intermediate_scales = None topk = torch.stack( [ @@ -73,7 +89,8 @@ class TestTensors: topk_weights = torch.randn(topk.shape, dtype=torch.float32, device="cuda") return TestTensors( rank_tokens=rank_tokens, - rank_token_scales=None, + rank_token_scales=rank_token_scales, + intermediate_scales=intermediate_scales, topk=topk, topk_weights=topk_weights, config=config, @@ -124,7 +141,6 @@ def make_modular_kernel( mk = FusedMoEKernel( prepare_finalize=a2a, fused_experts=fused_experts, - inplace=False, ) return mk @@ -162,6 +178,7 @@ def deepep_v2_moe_impl( w2_scale=w2_scale, per_act_token_quant=per_act_token_quant, a1_scale=test_tensors.rank_token_scales, + a2_scale=test_tensors.intermediate_scales, ) hidden_size = test_tensors.rank_tokens.size(1) @@ -231,6 +248,8 @@ def _deep_ep_v2_moe( test_tensors.topk, w1_scale=w1_scale, w2_scale=w2_scale, + a1_scale=test_tensors.rank_token_scales, + a2_scale=test_tensors.intermediate_scales, quant_dtype=q_dtype, per_act_token_quant=per_act_token_quant, ) @@ -262,12 +281,15 @@ def _deep_ep_v2_moe( per_act_token_quant, ) - torch.testing.assert_close( - torch_combined, - deepep_combined, - atol=6e-2, - rtol=6e-2, - ) + if is_quantized: + assert_fp8_close(torch_combined, deepep_combined) + else: + torch.testing.assert_close( + torch_combined, + deepep_combined, + atol=6e-2, + rtol=6e-2, + ) MNKs = [ @@ -355,17 +377,29 @@ def _deep_ep_v2_moe_cudagraph( num_local_experts = config.num_experts // pgi.world_size hidden_size = config.k - # Create FP8 weights directly, then dequantize for bf16 reference. - w1_fp8 = torch.randn( - (config.num_experts, 2 * config.n, config.k), - device="cuda", - dtype=torch.bfloat16, - ).to(torch.float8_e4m3fn) - w2_fp8 = torch.randn( - (config.num_experts, config.k, config.n), - device="cuda", - dtype=torch.bfloat16, - ).to(torch.float8_e4m3fn) + # All ranks must use the same global weights before taking their EP slice. + w1_bf16 = ( + torch.randn( + (config.num_experts, 2 * config.n, config.k), + device="cuda", + dtype=torch.bfloat16, + ) + / 15 + ) + w2_bf16 = ( + torch.randn( + (config.num_experts, config.k, config.n), + device="cuda", + dtype=torch.bfloat16, + ) + / 15 + ) + torch.distributed.broadcast(w1_bf16, src=0, group=pg) + torch.distributed.broadcast(w2_bf16, src=0, group=pg) + + # Round-trip through FP8 before constructing the reference and kernel weights. + w1_fp8 = w1_bf16.to(torch.float8_e4m3fn) + w2_fp8 = w2_bf16.to(torch.float8_e4m3fn) w1_ref = w1_fp8.to(torch.bfloat16) w2_ref = w2_fp8.to(torch.bfloat16) @@ -385,21 +419,8 @@ def _deep_ep_v2_moe_cudagraph( backend="nccl", ) initialize_model_parallel(tensor_model_parallel_size=1) - # Reference MoE using dequantized bf16 weights - torch_combined = torch_experts( - test_tensors.rank_tokens, - w1_ref, - w2_ref, - test_tensors.topk_weights, - test_tensors.topk, - ) - - # Use the production pipeline: make_fused_moe_layer creates - # a FusedMoE layer, quantizes weights, runs - # process_weights_after_loading (TrtLLM W31 swap + BlockMajorK - # shuffle), and selects the kernel. - # Quantize weights using production helper, EP-slice, then - # convert to TrtLLM format. + # Mirror production weight processing: quantize, EP-slice, then + # convert to the TrtLLM BlockMajorK format. from tests.kernels.moe.test_moe_layer import _quantize_fp8_halves from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import ( TrtLlmFp8ExpertsModular, @@ -411,14 +432,32 @@ def _deep_ep_v2_moe_cudagraph( block_shape = [128, 128] qw = _quantize_fp8_halves(w1_ref, w2_ref, block_shape) + assert qw.w13_weight_scale is not None + assert qw.w2_weight_scale is not None + + # Reference MoE using the same blockwise FP8 quantization scheme as + # the production kernel. torch_experts quantizes activations before + # both GEMMs and dequantizes the operands for the reference matmuls. + reference_topk_weights = test_tensors.topk_weights.to(torch.bfloat16).to( + torch.float32 + ) + torch_combined = torch_experts( + test_tensors.rank_tokens, + qw.w13_weight, + qw.w2_weight, + reference_topk_weights, + test_tensors.topk, + w1_scale=qw.w13_weight_scale, + w2_scale=qw.w2_weight_scale, + quant_dtype=torch.float8_e4m3fn, + block_shape=block_shape, + ) # EP-slice before format conversion e_start = num_local_experts * pgi.rank e_end = e_start + num_local_experts w1_ep = qw.w13_weight[e_start:e_end] w2_ep = qw.w2_weight[e_start:e_end] - assert qw.w13_weight_scale is not None - assert qw.w2_weight_scale is not None w1_scale_ep = qw.w13_weight_scale[e_start:e_end] w2_scale_ep = qw.w2_weight_scale[e_start:e_end] @@ -452,11 +491,23 @@ def _deep_ep_v2_moe_cudagraph( w2_scale=w2_scale_ep, ) moe_config = make_dummy_moe_config( - num_experts=num_local_experts, + num_experts=config.num_experts, + num_local_experts=num_local_experts, experts_per_token=config.topk, hidden_dim=hidden_size, intermediate_size=config.n, ) + moe_parallel_config = dataclasses.replace( + moe_config.moe_parallel_config, + ep_size=pgi.world_size, + ep_rank=pgi.rank, + use_ep=True, + all2all_backend="deepep_v2", + ) + moe_config = dataclasses.replace( + moe_config, + moe_parallel_config=moe_parallel_config, + ) fused_experts = TrtLlmFp8ExpertsModular( moe_config=moe_config, quant_config=quant_config, @@ -480,21 +531,21 @@ def _deep_ep_v2_moe_cudagraph( mk_kernel = FusedMoEKernel( prepare_finalize=a2a, fused_experts=fused_experts, - inplace=False, ) - for _ in range(3): - out = mk_kernel.apply( - hidden_states=test_tensors.rank_tokens, - w1=w1_ep, - w2=w2_ep, - topk_weights=test_tensors.topk_weights, - topk_ids=test_tensors.topk, - activation=MoEActivation.SILU, - global_num_experts=config.num_experts, - expert_map=None, - apply_router_weight_on_input=False, - ) + with set_forward_context(None, vllm_cfg): + for _ in range(3): + out = mk_kernel.apply( + hidden_states=test_tensors.rank_tokens, + w1=w1_ep, + w2=w2_ep, + topk_weights=test_tensors.topk_weights, + topk_ids=test_tensors.topk, + activation=MoEActivation.SILU, + global_num_experts=config.num_experts, + expert_map=None, + apply_router_weight_on_input=False, + ) torch.testing.assert_close( torch_combined, diff --git a/tools/ep_kernels/README.md b/tools/ep_kernels/README.md index b4eabe18ca1..04bc3ef2768 100644 --- a/tools/ep_kernels/README.md +++ b/tools/ep_kernels/README.md @@ -11,6 +11,38 @@ Step 2 is necessary for multi-node deployment. All scripts accept a positional argument as workspace path for staging the build, defaulting to `$(pwd)/ep_kernels_workspace`. +## NCCL version requirement (CUDA 13+) + +DeepEPv2 uses the NCCL GIN (GPU-Initiated Networking) backend, which requires +NCCL >= 2.30.4 at both compile time and runtime. PyTorch 2.11 pins +`nvidia-nccl-cu13==2.28.9` as a transitive dependency, so you need to +override it. + +**With uv** (recommended): + +```bash +# Create an override file +echo "nvidia-nccl-cu13>=2.30.4" > /tmp/nccl-override.txt +export UV_OVERRIDE=/tmp/nccl-override.txt + +# All subsequent uv pip install commands will respect the override +uv pip install vllm +``` + +**With pip**: + +```bash +pip install vllm +pip install "nvidia-nccl-cu13>=2.30.4" --no-deps +``` + +The override / reinstall must happen before building DeepEP (for GIN device +headers) and must remain in place at runtime. You can verify with: + +```bash +python -c "from vllm.utils.import_utils import has_deep_ep_v2; print(has_deep_ep_v2())" +``` + ## Usage ```bash diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 33841de306e..679764f6a82 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -978,6 +978,7 @@ class DeepEPV2All2AllManager(All2AllManagerBase): self._device_group = device_group self.handle_cache = Cache() self._num_sms: int | None = None + self._gin_checked = False def _make_all2all_kwargs( self, @@ -1000,11 +1001,37 @@ class DeepEPV2All2AllManager(All2AllManagerBase): explicitly_destroy=True, ) + def _check_gin_support(self, group) -> None: + from vllm.utils.nccl import query_nccl_gin_type + + # ProcessGroupNCCL creates communicators lazily. Initialize this exact + # group before querying so a null comm pointer is not mistaken for + # missing GIN support. + probe = torch.zeros(1, device="cuda") + torch.distributed.all_reduce(probe, group=group) + + gin_type = query_nccl_gin_type(group) + if gin_type is None: + raise RuntimeError( + "DeepEPv2 communicator properties query failed; " + "networking capability could not be determined." + ) + if gin_type == 0: + raise RuntimeError( + "DeepEPv2 requires NCCL GIN (GPU-Initiated Networking). " + "This usually means IBGDA-capable InfiniBand NICs or drivers " + "are not available. See tools/ep_kernels/README.md for " + "requirements." + ) + def get_handle(self, kwargs): import deep_ep # type: ignore[import-not-found] num_experts = kwargs.pop("num_experts", 256) buffer_kwargs = self._make_all2all_kwargs(**kwargs) + if not self._gin_checked: + self._check_gin_support(buffer_kwargs["group"]) + self._gin_checked = True logger.debug("DeepEP v2 all2all args %s", buffer_kwargs) handle: deep_ep.ElasticBuffer = self.handle_cache.get_or_create( buffer_kwargs, deep_ep.ElasticBuffer diff --git a/vllm/distributed/device_communicators/pynccl_allocator.py b/vllm/distributed/device_communicators/pynccl_allocator.py index a5d7faceb2b..4898c6212f5 100644 --- a/vllm/distributed/device_communicators/pynccl_allocator.py +++ b/vllm/distributed/device_communicators/pynccl_allocator.py @@ -14,7 +14,7 @@ from vllm import envs from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator from vllm.logger import init_logger from vllm.platforms import current_platform -from vllm.utils.nccl import find_nccl_include_paths +from vllm.utils.nccl import find_nccl_include_paths, find_nccl_library_paths logger = init_logger(__name__) @@ -74,11 +74,15 @@ def compile_nccl_allocator(): out_dir = tempfile.gettempdir() nccl_allocator_libname = "nccl_allocator" nccl_include_paths = find_nccl_include_paths() + ldflags = ["-l:libnccl.so.2"] + nccl_lib_paths = find_nccl_library_paths() + if nccl_lib_paths: + ldflags = [f"-L{p}" for p in nccl_lib_paths] + ldflags load_inline( name=nccl_allocator_libname, cpp_sources=nccl_allocator_source, with_cuda=True, - extra_ldflags=["-lnccl"], + extra_ldflags=ldflags, verbose=envs.VLLM_LOGGING_LEVEL == "DEBUG", is_python_module=False, build_directory=out_dir, diff --git a/vllm/distributed/device_communicators/pynccl_wrapper.py b/vllm/distributed/device_communicators/pynccl_wrapper.py index 5ca8cc7c77f..78ee81e7b5e 100644 --- a/vllm/distributed/device_communicators/pynccl_wrapper.py +++ b/vllm/distributed/device_communicators/pynccl_wrapper.py @@ -51,6 +51,26 @@ class ncclUniqueId(ctypes.Structure): _fields_ = [("internal", ctypes.c_byte * 128)] +# NCCL 2.30+ ncclCommProperties_t. Only fields through ginType are read; +# trailing fields keep the layout aligned with NCCL's versioned structure. +class ncclCommProperties(ctypes.Structure): + _fields_ = [ + ("size", ctypes.c_size_t), + ("magic", ctypes.c_uint), + ("version", ctypes.c_uint), + ("rank", ctypes.c_int), + ("nRanks", ctypes.c_int), + ("cudaDev", ctypes.c_int), + ("nvmlDev", ctypes.c_int), + ("deviceApiSupport", ctypes.c_bool), + ("multimemSupport", ctypes.c_bool), + ("ginType", ctypes.c_int), + ("nLsaTeams", ctypes.c_int), + ("hostRmaSupport", ctypes.c_bool), + ("railedGinType", ctypes.c_int), + ] + + cudaStream_t = ctypes.c_void_p buffer_type = ctypes.c_void_p @@ -317,6 +337,12 @@ class NCCLLibrary: # ncclResult_t ncclCommWindowDeregister( # ncclComm_t comm, ncclWindow_t win); Function("ncclCommWindowDeregister", ncclResult_t, [ncclComm_t, ncclWindow_t]), + # Query runtime properties of a specific initialized communicator. + Function( + "ncclCommQueryProperties", + ncclResult_t, + [ncclComm_t, ctypes.POINTER(ncclCommProperties)], + ), ] # class attribute to store the mapping from the path to the library @@ -375,6 +401,9 @@ class NCCLLibrary: # Having an exception here on ROCm platform is # not allowed during graph capturing continue + elif func.name == "ncclCommQueryProperties": + # Optional on NCCL versions older than 2.29. + continue raise NCCLLibrary.path_to_dict_mapping[so_file] = _funcs self._funcs = NCCLLibrary.path_to_dict_mapping[so_file] diff --git a/vllm/utils/nccl.py b/vllm/utils/nccl.py index 4807bc076f8..aeadaf5d22d 100644 --- a/vllm/utils/nccl.py +++ b/vllm/utils/nccl.py @@ -3,6 +3,7 @@ from __future__ import annotations +import ctypes import importlib.util import os @@ -62,3 +63,64 @@ def find_nccl_include_paths() -> list[str] | None: out.append(p) seen.add(p) return out or None + + +def find_nccl_library_paths() -> list[str] | None: + """Return possible library paths containing `libnccl.so`. + + Looks inside the `nvidia-nccl-cuXX` pip package. + """ + paths: list[str] = [] + try: + spec = importlib.util.find_spec("nvidia.nccl") + if spec and (locs := getattr(spec, "submodule_search_locations", None)): + for loc in locs: + lib_dir = os.path.join(loc, "lib") + if os.path.isdir(lib_dir): + paths.append(lib_dir) + except Exception as e: + logger.debug("Failed to find nccl library path from nvidia.nccl package: %s", e) + return paths or None + + +def query_nccl_gin_type(group: torch.distributed.ProcessGroup) -> int | None: + """Return the GIN type for an initialized group, or ``None`` on failure.""" + from vllm.distributed.device_communicators.pynccl_wrapper import ( + NCCLLibrary, + ncclCommProperties, + ) + + try: + backend = group._get_backend(torch.device("cuda")) + # GIN is a property of this initialized communicator, not just the + # NCCL version. ncclCommQueryProperties requires its ncclComm_t. + comm_ptr = backend._comm_ptr() + if comm_ptr == 0: + return None + except Exception: + logger.warning( + "Failed to extract NCCL comm pointer from process group", + exc_info=True, + ) + return None + + try: + nccl = NCCLLibrary() + query_fn = nccl._funcs.get("ncclCommQueryProperties") + if query_fn is None: + return None + + props = ncclCommProperties() + ctypes.memset(ctypes.addressof(props), 0, ctypes.sizeof(props)) + props.size = ctypes.sizeof(props) + props.magic = 0xCAFEBEEF + props.version = nccl.ncclGetRawVersion() + result = query_fn(ctypes.c_void_p(comm_ptr), ctypes.byref(props)) + except Exception: + logger.warning("Failed to query NCCL communicator properties", exc_info=True) + return None + + if result != 0: + logger.warning("ncclCommQueryProperties returned error %d", result) + return None + return props.ginType From 2279575cd9caeb9174c7ec11a94a34af40903f4d Mon Sep 17 00:00:00 2001 From: Oxana Korzh Date: Fri, 24 Jul 2026 14:35:27 -0600 Subject: [PATCH 013/185] [AMD][Bugfix][EPLB] Fix elastic EP scaling accuracy on ROCm (#47206) Signed-off-by: okorzh Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: Andreas Karatzas --- vllm/distributed/eplb/eplb_state.py | 88 +++++++++++++++---- .../layers/fused_moe/routed_experts.py | 20 +++-- 2 files changed, 88 insertions(+), 20 deletions(-) diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index feacb03d28b..d34b844bd3d 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -46,6 +46,7 @@ from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator from vllm.distributed.utils import StatelessProcessGroup from vllm.logger import init_logger from vllm.model_executor.models.interfaces import MixtureOfExperts +from vllm.platforms import current_platform from .async_worker import start_async_worker from .eplb_communicator import EplbCommunicator, create_eplb_communicator @@ -825,23 +826,80 @@ class EplbState: eplb_model_state.physical_to_logical_map.cpu(), ) - # Update expert weights - rearrange_expert_weights_inplace( - eplb_model_state.physical_to_logical_map, - new_physical_to_logical_map, - eplb_model_state.model.expert_weights, - eplb_model_state.expert_buffer, - ep_group, - eplb_model_state.communicator, - is_profile, - rank_mapping, - ) + skip_rearrange = False + if ( + current_platform.is_rocm() + and not is_profile + and rank_mapping is None + and bool((eplb_model_state.physical_to_logical_map >= 0).all()) + ): + logical_loads = global_expert_load_window.float() + ep_size = ep_group.size() - if not is_profile: - _commit_eplb_maps( - eplb_model_state, - new_physical_to_logical_map=new_physical_to_logical_map, + def rank_load_imbalance( + mapping: torch.Tensor, + logical_loads: torch.Tensor = logical_loads, + ep_size: int = ep_size, + ) -> float: + mapping = mapping.to( + device=logical_loads.device, + dtype=torch.long, + ) + replica_counts = torch.zeros_like(logical_loads) + replica_counts.scatter_add_( + dim=1, + index=mapping, + src=torch.ones_like(mapping, dtype=logical_loads.dtype), + ) + loads_per_replica = torch.gather( + logical_loads / replica_counts.clamp_min(1), + dim=1, + index=mapping, + ) + loads_per_rank = loads_per_replica.reshape( + logical_loads.shape[0], ep_size, -1 + ).sum(dim=(0, 2)) + mean_load = loads_per_rank.mean() + if mean_load == 0: + return 1.0 + return (loads_per_rank.max() / mean_load).item() + + current_imbalance = rank_load_imbalance( + eplb_model_state.physical_to_logical_map ) + proposed_imbalance = rank_load_imbalance( + new_physical_to_logical_map + ) + relative_improvement = ( + current_imbalance - proposed_imbalance + ) / current_imbalance + skip_rearrange = relative_improvement < 0.05 + if skip_rearrange and is_main_rank: + logger.info( + "[EPLB] Skip rearrange: imbalance %.4f -> " + "%.4f (no material gain)", + current_imbalance, + proposed_imbalance, + ) + + if not skip_rearrange: + # Update expert weights + rearrange_expert_weights_inplace( + eplb_model_state.physical_to_logical_map, + new_physical_to_logical_map, + eplb_model_state.model.expert_weights, + eplb_model_state.expert_buffer, + ep_group, + eplb_model_state.communicator, + is_profile, + rank_mapping, + ) + + if not is_profile: + _commit_eplb_maps( + eplb_model_state, + new_physical_to_logical_map=new_physical_to_logical_map, + ) if is_main_rank: assert start_event is not None diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 8bf0123e57f..6f46283829d 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -233,17 +233,27 @@ class RoutedExperts(PluggableLayer): # Update local attributes from ExpertMapManager self.local_num_experts = self.expert_map_manager.local_num_experts self.expert_placement_strategy = self.expert_map_manager.placement_strategy - self.register_buffer("_expert_map", self.expert_map_manager.expert_map) - self.register_buffer("expert_mask", self.expert_map_manager.expert_mask) + self.register_buffer( + "_expert_map", self.expert_map_manager.expert_map, persistent=False + ) + self.register_buffer( + "expert_mask", self.expert_map_manager.expert_mask, persistent=False + ) # Get routing tables from ExpertMapManager routing_tables = self.expert_map_manager.routing_tables if routing_tables is not None: # Register routing tables as buffers for this layer global_to_physical, physical_to_global, local_global = routing_tables - self.register_buffer("expert_global_to_physical", global_to_physical) - self.register_buffer("expert_physical_to_global", physical_to_global) - self.register_buffer("expert_local_to_global", local_global) + self.register_buffer( + "expert_global_to_physical", global_to_physical, persistent=False + ) + self.register_buffer( + "expert_physical_to_global", physical_to_global, persistent=False + ) + self.register_buffer( + "expert_local_to_global", local_global, persistent=False + ) def _expert_routing_tables( self, From 972848f2764e2f79deb7853320ccf0094be452aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Elvir=20Crn=C4=8Devi=C4=87?= Date: Fri, 24 Jul 2026 22:38:59 +0200 Subject: [PATCH 014/185] [Bugfix] Support non-uniform page sizes in KVBlockZeroer (#49704) Signed-off-by: Elvir Crncevic Co-authored-by: Claude Opus 4.6 --- tests/v1/worker/test_kv_block_zeroer.py | 53 ++++++++++++++++-- .../warmup/qwen_triton_warmup.py | 16 +++--- vllm/v1/worker/utils.py | 54 +++++++++++-------- 3 files changed, 92 insertions(+), 31 deletions(-) diff --git a/tests/v1/worker/test_kv_block_zeroer.py b/tests/v1/worker/test_kv_block_zeroer.py index 365e4adacea..b212e3ae17b 100644 --- a/tests/v1/worker/test_kv_block_zeroer.py +++ b/tests/v1/worker/test_kv_block_zeroer.py @@ -20,9 +20,10 @@ def test_block_ids_are_not_overwritten_while_copy_is_in_flight(): zeroer.device = device zeroer._meta = ( torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device), - page_size_el, - page_size_el, - 1, + torch.tensor([page_size_el], dtype=torch.int64, device=device), + page_size_el // page_size_el, # max_chunks = 1 + page_size_el, # blk_size + 1, # n_segs ) stream = torch.cuda.Stream() @@ -39,3 +40,49 @@ def test_block_ids_are_not_overwritten_while_copy_is_in_flight(): assert torch.all(storage[1] == 0) assert torch.all(storage[2] == 0) assert torch.all(storage[3] == 1) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_non_uniform_page_sizes(): + """Two segments with different page sizes (e.g. MLA + DSA indexer).""" + device = torch.device("cuda") + num_blocks = 4 + page_size_a = 10496 # int32 elements + page_size_b = 2112 + + storage_a = torch.ones((num_blocks, page_size_a), dtype=torch.int32, device=device) + storage_b = torch.ones((num_blocks, page_size_b), dtype=torch.int32, device=device) + + zeroer = KVBlockZeroer.__new__(KVBlockZeroer) + zeroer.device = device + + seg_page_sizes = [page_size_a, page_size_b] + max_ps = max(seg_page_sizes) + + def largest_power_of_2_divisor(n): + return n & -n + + blk_size = min(min(largest_power_of_2_divisor(ps) for ps in seg_page_sizes), 1024) + + zeroer._meta = ( + torch.tensor( + [storage_a.data_ptr(), storage_b.data_ptr()], + dtype=torch.uint64, + device=device, + ), + torch.tensor(seg_page_sizes, dtype=torch.int64, device=device), + max_ps // blk_size, + blk_size, + 2, + ) + + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + zeroer.zero_block_ids([1, 2]) + stream.synchronize() + + for storage in (storage_a, storage_b): + assert torch.all(storage[0] == 1) + assert torch.all(storage[1] == 0) + assert torch.all(storage[2] == 0) + assert torch.all(storage[3] == 1) diff --git a/vllm/model_executor/warmup/qwen_triton_warmup.py b/vllm/model_executor/warmup/qwen_triton_warmup.py index 9b7d76207b4..8cbfa539b7e 100644 --- a/vllm/model_executor/warmup/qwen_triton_warmup.py +++ b/vllm/model_executor/warmup/qwen_triton_warmup.py @@ -36,7 +36,8 @@ _FLA_POST_CONV_WARMUP_LENGTHS = (1, 2, 16) @dataclass(frozen=True) class _ZeroKvWarmupConfig: - page_size_el: int + seg_page_sizes: torch.Tensor + max_chunks: int block_size: int n_segs: int @@ -160,9 +161,10 @@ def _zero_kv_warmup_config(runner: object) -> _ZeroKvWarmupConfig | None: if meta is None: return None - _, page_size_el, block_size, n_segs = meta + _, seg_page_sizes, max_chunks, block_size, n_segs = meta return _ZeroKvWarmupConfig( - page_size_el=int(page_size_el), + seg_page_sizes=seg_page_sizes, + max_chunks=int(max_chunks), block_size=int(block_size), n_segs=int(n_segs), ) @@ -185,8 +187,9 @@ def _warm_zero_kv_blocks_kernel( from vllm.v1.worker.utils import _zero_kv_blocks_kernel max_n_blocks = max(_ZERO_KV_N_BLOCKS) + max_page_size = int(config.seg_page_sizes.max().item()) scratch = torch.empty( - max_n_blocks * config.page_size_el, + max_n_blocks * max_page_size, dtype=torch.int32, device=device, ) @@ -198,13 +201,14 @@ def _warm_zero_kv_blocks_kernel( for n_blocks in _ZERO_KV_N_BLOCKS: block_ids = torch.arange(n_blocks, dtype=torch.int64, device=device) - grid = (n_blocks * config.n_segs * (config.page_size_el // config.block_size),) + grid = (n_blocks * config.n_segs * config.max_chunks,) _zero_kv_blocks_kernel[grid]( seg_addrs, + config.seg_page_sizes, block_ids, n_blocks, N_SEGS=config.n_segs, - PAGE_SIZE_EL=config.page_size_el, + MAX_CHUNKS=config.max_chunks, BLOCK_SIZE=config.block_size, ) diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index d83242b1606..d00974046fc 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -43,10 +43,11 @@ logger = init_logger(__name__) @triton.jit def _zero_kv_blocks_kernel( seg_addrs_ptr, + seg_page_sizes_ptr, block_ids_ptr, n_blocks, N_SEGS: tl.constexpr, - PAGE_SIZE_EL: tl.constexpr, + MAX_CHUNKS: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): """Zero KV cache blocks across all segments in a single launch. @@ -56,25 +57,33 @@ def _zero_kv_blocks_kernel( buffer. For backends where K/V is outermost (block_dim=1) there are two segments per buffer (one for K, one for V). + Segments may have different page sizes (e.g. models with multiple KV + cache groups like MLA + DSA indexer). Each segment's page size is + read from seg_page_sizes_ptr; programs whose chunk_index falls beyond + their segment's page size early-exit. + seg_addrs_ptr holds absolute byte addresses (int64) for each segment, allowing segments to live in different CUDA allocations. Programs are mapped as (block_index, seg_index, chunk_index). """ pid = tl.program_id(0) - chunks = PAGE_SIZE_EL // BLOCK_SIZE - work_per_block = N_SEGS * chunks + work_per_block = N_SEGS * MAX_CHUNKS block_index = pid // work_per_block if block_index >= n_blocks: return remainder = pid % work_per_block - seg_index = remainder // chunks - chunk_index = remainder % chunks + seg_index = remainder // MAX_CHUNKS + chunk_index = remainder % MAX_CHUNKS + page_size_el = tl.load(seg_page_sizes_ptr + seg_index) + if chunk_index >= page_size_el // BLOCK_SIZE: + return block_id = tl.load(block_ids_ptr + block_index) seg_addr = tl.load(seg_addrs_ptr + seg_index) ptr = tl.cast(seg_addr, tl.pointer_type(tl.int32)) offset = ( - block_id.to(tl.int64) * PAGE_SIZE_EL + chunk_index.to(tl.int64) * BLOCK_SIZE + block_id.to(tl.int64) * page_size_el.to(tl.int64) + + chunk_index.to(tl.int64) * BLOCK_SIZE ) cols = tl.arange(0, BLOCK_SIZE).to(tl.int64) tl.store(ptr + offset + cols, tl.zeros([BLOCK_SIZE], dtype=tl.int32)) @@ -104,19 +113,19 @@ class KVBlockZeroer: Block IDs from the scheduler reference logical blocks whose size may differ from the kernel block size (virtual block splitting). - PAGE_SIZE_EL accounts for this ratio so that - ``block_id * PAGE_SIZE_EL`` lands at the correct offset. + Each segment's page_size_el accounts for this ratio so that + ``block_id * page_size_el`` lands at the correct offset. Only AttentionSpec layers are processed; Mamba layers are skipped. """ self.device = device - self._meta: tuple[torch.Tensor, int, int, int] | None = None + self._meta: tuple[torch.Tensor, torch.Tensor, int, int, int] | None = None if runner_only_attn_layers is None: runner_only_attn_layers = set() seen_ptrs: set[int] = set() seg_addrs: list[int] = [] - page_size_el: int | None = None + seg_page_sizes: list[int] = [] for group in attn_groups_iter: spec = group.kv_cache_spec @@ -149,12 +158,6 @@ class KVBlockZeroer: assert cur_bytes % 4 == 0 kernel_block_el = cur_bytes // 4 cur_page_el = kernel_block_el * ratio - if page_size_el is None: - page_size_el = cur_page_el - else: - assert page_size_el == cur_page_el, ( - f"Non-uniform page sizes: {page_size_el} vs {cur_page_el}" - ) block_stride_bytes = cur_bytes outer_dims = [ @@ -166,15 +169,21 @@ class KVBlockZeroer: for outer in iprod(*(range(kv.shape[d]) for d in outer_dims)): off_bytes = sum(i * s for i, s in zip(outer, outer_strides)) seg_addrs.append(dp + off_bytes) + seg_page_sizes.append(cur_page_el) - if not seg_addrs or page_size_el is None: + if not seg_addrs: self._meta = None return - blk_size = min(largest_power_of_2_divisor(page_size_el), 1024) + max_page_size_el = max(seg_page_sizes) + blk_size = min( + min(largest_power_of_2_divisor(ps) for ps in seg_page_sizes), + 1024, + ) self._meta = ( torch.tensor(seg_addrs, dtype=torch.uint64, device=self.device), - page_size_el, + torch.tensor(seg_page_sizes, dtype=torch.int64, device=self.device), + max_page_size_el // blk_size, blk_size, len(seg_addrs), ) @@ -183,16 +192,17 @@ class KVBlockZeroer: """Zero the KV cache memory for the given block IDs.""" if not block_ids or self._meta is None: return - seg_addrs, page_size_el, blk_size, n_segs = self._meta + seg_addrs, seg_page_sizes, max_chunks, blk_size, n_segs = self._meta n_blocks = len(block_ids) idx = async_tensor_h2d(block_ids, device=self.device, dtype=torch.int64) - grid = (n_blocks * n_segs * (page_size_el // blk_size),) + grid = (n_blocks * n_segs * max_chunks,) _zero_kv_blocks_kernel[grid]( seg_addrs, + seg_page_sizes, idx, n_blocks, N_SEGS=n_segs, - PAGE_SIZE_EL=page_size_el, + MAX_CHUNKS=max_chunks, BLOCK_SIZE=blk_size, ) From 9e6746b3c7b4ec9bcc234db9a654df8eca5781ce Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Sat, 25 Jul 2026 04:45:09 +0800 Subject: [PATCH 015/185] [CI] Stabilize memory-sensitive compile and structured output tests (#49749) Signed-off-by: zjy0516 Co-authored-by: OpenAI Codex --- tests/compile/fullgraph/test_basic_correctness.py | 2 ++ tests/entrypoints/llm/test_struct_output_generate.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/tests/compile/fullgraph/test_basic_correctness.py b/tests/compile/fullgraph/test_basic_correctness.py index 35989dcde1d..7471564340d 100644 --- a/tests/compile/fullgraph/test_basic_correctness.py +++ b/tests/compile/fullgraph/test_basic_correctness.py @@ -64,6 +64,8 @@ class TestSetting: "bfloat16", "--max-model-len", "2048", + "--gpu-memory-utilization", + "0.98", ], pp_size=1, tp_size=1, diff --git a/tests/entrypoints/llm/test_struct_output_generate.py b/tests/entrypoints/llm/test_struct_output_generate.py index 3ece2723436..219ee7cd387 100644 --- a/tests/entrypoints/llm/test_struct_output_generate.py +++ b/tests/entrypoints/llm/test_struct_output_generate.py @@ -212,6 +212,7 @@ class CarDescription(BaseModel): PARAMS_MODELS_BACKENDS_TOKENIZER_MODE, ) def test_structured_output( + request: pytest.FixtureRequest, backend: str, tokenizer_mode: str, model_name: str, @@ -242,6 +243,7 @@ def test_structured_output( speculative_config=speculative_config, **platform_args, ) + request.addfinalizer(llm.llm_engine.engine_core.shutdown) # # Test 1: Generate JSON output based on a provided schema From 84d26b9ee3dde1991442539dcf5fc5b26aa08f7e Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Fri, 24 Jul 2026 21:54:24 +0100 Subject: [PATCH 016/185] [Model] Remove Plamo2 (#49729) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- .buildkite/test-amd.yaml | 6 +- .buildkite/test_areas/models_language.yaml | 2 - docs/design/custom_op.md | 2 - docs/models/supported_models.md | 1 - docs/usage/v1_guide.md | 2 +- tests/distributed/test_pipeline_parallel.py | 1 - .../models/language/generation/test_hybrid.py | 2 - tests/models/registry.py | 4 - tests/quantization/test_experts_int8.py | 2 +- vllm/config/compilation.py | 1 - vllm/config/model.py | 3 +- vllm/model_executor/models/plamo2.py | 992 ------------------ vllm/model_executor/models/registry.py | 2 +- 13 files changed, 6 insertions(+), 1014 deletions(-) delete mode 100644 vllm/model_executor/models/plamo2.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index b8046fda9f5..b67acbc40b3 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -81,10 +81,8 @@ # the above test.) Also run if model initialization test file is modified. # # * [Language Models Tests (Extra Standard) %N]: Shard slow subset of standard language models tests. Only run when model # # source is modified, or when specified test files are modified. # -# * [Language Models Tests (Hybrid) %N]: Install fast path packages for testing against transformers (mamba, conv1d) and to # -# run plamo2 model in vLLM. # -# * [Language Models Test (Extended Generation)]: Install fast path packages for testing against transformers (mamba, conv1d) # -# and to run plamo2 model in vLLM. # +# * [Language Models Tests (Hybrid) %N]: Install fast path packages for testing against transformers (mamba, conv1d). # +# * [Language Models Test (Extended Generation)]: Install fast path packages for testing against transformers (mamba, conv1d). # # * [Multi-Modal Models (Standard) 1-4]: # # - Do NOT remove `VLLM_WORKER_MULTIPROC_METHOD=spawn` setting as ROCm requires this for certain models to function. # # * [Transformers Nightly Models]: Whisper needs `VLLM_WORKER_MULTIPROC_METHOD=spawn` to avoid deadlock. # diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index 6124cb023dd..2d3f14d1715 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -63,7 +63,6 @@ steps: - tests/models/language/generation commands: # Install fast path packages for testing against transformers - # Note: also needed to run plamo2 model in vLLM - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0' - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' # Shard the hybrid language model tests that are numerically stable on Hopper. @@ -105,7 +104,6 @@ steps: - tests/models/language/generation commands: # Install fast path packages for testing against transformers - # Note: also needed to run plamo2 model in vLLM - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0' - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' diff --git a/docs/design/custom_op.md b/docs/design/custom_op.md index d2557a2281c..f99c8043fa0 100644 --- a/docs/design/custom_op.md +++ b/docs/design/custom_op.md @@ -122,8 +122,6 @@ For example: --8<-- "vllm/model_executor/layers/mamba/mamba_mixer2.py:mixer2_gated_rms_norm" ---8<-- "vllm/model_executor/models/plamo2.py:plamo2_mamba_mixer" - --8<-- "vllm/model_executor/layers/mamba/short_conv.py:short_conv" ``` diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 98f9cb5b65d..f1292d7a988 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -435,7 +435,6 @@ th { | `PhiForCausalLM` | Phi | `microsoft/phi-1_5`, `microsoft/phi-2`, etc. | ✅︎ | ✅︎ | | `Phi3ForCausalLM` | Phi-4, Phi-3 | `microsoft/Phi-4-mini-instruct`, `microsoft/Phi-4`, `microsoft/Phi-3-mini-4k-instruct`, `microsoft/Phi-3-mini-128k-instruct`, `microsoft/Phi-3-medium-128k-instruct`, etc. | ✅︎ | ✅︎ | | `PhiMoEForCausalLM` | Phi-3.5-MoE | `microsoft/Phi-3.5-MoE-instruct`, etc. | ✅︎ | ✅︎ | -| `Plamo2ForCausalLM` | PLaMo2 | `pfnet/plamo-2-1b`, `pfnet/plamo-2-8b`, etc. | ✅ | ✅︎ | | `Plamo3ForCausalLM` | PLaMo3 | `pfnet/plamo-3-nict-2b-base`, `pfnet/plamo-3-nict-8b-base`, etc. | ✅ | ✅︎ | | `Qwen2ForCausalLM` | QwQ, Qwen2 | `Qwen/QwQ-32B-Preview`, `Qwen/Qwen2-7B-Instruct`, `Qwen/Qwen2-7B`, etc. | ✅︎ | ✅︎ | | `Qwen2MoeForCausalLM` | Qwen2MoE | `Qwen/Qwen1.5-MoE-A2.7B`, `Qwen/Qwen1.5-MoE-A2.7B-Chat`, etc. | ✅︎ | ✅︎ | diff --git a/docs/usage/v1_guide.md b/docs/usage/v1_guide.md index 5613d5ba4e8..4ef2b0b5b6a 100644 --- a/docs/usage/v1_guide.md +++ b/docs/usage/v1_guide.md @@ -126,7 +126,7 @@ Models using selective state-space mechanisms instead of standard transformer at Models that use Mamba-2 and Mamba-1 layers (e.g., `Mamba2ForCausalLM`, `MambaForCausalLM`, `FalconMambaForCausalLM`) are supported. Hybrid models that combine Mamba-2 and Mamba-1 layers with standard attention layers are also supported (e.g., -`Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`, `JambaForCausalLM`, `Plamo2ForCausalLM`). +`Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`, `JambaForCausalLM`). Hybrid models with mechanisms different to Mamba are also supported (e.g, `Lfm2ForCausalLM`). diff --git a/tests/distributed/test_pipeline_parallel.py b/tests/distributed/test_pipeline_parallel.py index fd4a1d0f570..22e0321da60 100644 --- a/tests/distributed/test_pipeline_parallel.py +++ b/tests/distributed/test_pipeline_parallel.py @@ -121,7 +121,6 @@ TEXT_GENERATION_MODELS = { "ibm/PowerMoE-3b": PPTestSettings.fast(), "internlm/internlm2-chat-7b": PPTestSettings.fast(), "ai21labs/Jamba-tiny-dev": PPTestSettings.fast(), - "pfnet/plamo-2-1b": PPTestSettings.fast(), "pfnet/plamo-3-nict-2b-base": PPTestSettings.fast(), "meta-llama/Llama-3.2-1B-Instruct": PPTestSettings.detailed(), # Tests TransformersForCausalLM diff --git a/tests/models/language/generation/test_hybrid.py b/tests/models/language/generation/test_hybrid.py index 3fee0662bc6..768b26c8688 100644 --- a/tests/models/language/generation/test_hybrid.py +++ b/tests/models/language/generation/test_hybrid.py @@ -35,7 +35,6 @@ SSM_MODELS = [ HYBRID_MODELS = [ "ai21labs/Jamba-tiny-dev", - "pfnet/plamo-2-1b", "Zyphra/Zamba2-1.2B-instruct", "ibm-granite/granite-4.0-tiny-preview", "tiiuae/Falcon-H1-0.5B-Base", @@ -50,7 +49,6 @@ HYBRID_MODELS_REQUIRING_CHUNKED_PREFILL = { FULL_CUDA_GRAPH_MODELS = [ "ai21labs/Jamba-tiny-dev", - "pfnet/plamo-2-1b", "Zyphra/Zamba2-1.2B-instruct", ] diff --git a/tests/models/registry.py b/tests/models/registry.py index 4362e134c76..65cdc84dec8 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -484,10 +484,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "PhiMoEForCausalLM": _HfExamplesInfo( "microsoft/Phi-3.5-MoE-instruct", trust_remote_code=True ), - "Plamo2ForCausalLM": _HfExamplesInfo( - "pfnet/plamo-2-1b", - trust_remote_code=True, - ), "Plamo3ForCausalLM": _HfExamplesInfo( "pfnet/plamo-3-nict-2b-base", trust_remote_code=True, diff --git a/tests/quantization/test_experts_int8.py b/tests/quantization/test_experts_int8.py index 7cdb135fa07..6119c1c8468 100644 --- a/tests/quantization/test_experts_int8.py +++ b/tests/quantization/test_experts_int8.py @@ -12,7 +12,7 @@ from tests.quantization.utils import is_quant_method_supported from ..models.registry import HF_EXAMPLE_MODELS -MODELS = ["ai21labs/Jamba-tiny-random", "pfnet/plamo-2-1b"] +MODELS = ["ai21labs/Jamba-tiny-random"] @pytest.mark.skipif( diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 810e97c4cad..8142004748a 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -768,7 +768,6 @@ class CompilationConfig: "vllm::mamba_mixer", "vllm::short_conv", "vllm::linear_attention", - "vllm::plamo2_mamba_mixer", "vllm::qwen_gdn_attention_core", "vllm::gdn_attention_core_xpu", "vllm::olmo_hybrid_gdn_full_forward", diff --git a/vllm/config/model.py b/vllm/config/model.py index 1f62351c526..a65d5d5a94f 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1448,7 +1448,7 @@ class ModelConfig: """ Returns the mamba chunk size if it exists """ - # used by e.g. Bamba, FalconH1, Granite, PLaMo2 + # used by e.g. Bamba, FalconH1, Granite chunk_size = getattr(self.hf_text_config, "mamba_chunk_size", None) if chunk_size is None: # used by e.g. Mamba2, NemotronH, Zamba @@ -2022,7 +2022,6 @@ _FLOAT16_NOT_SUPPORTED_MODELS = { "gemma2": "Numerical instability. Please use bfloat16 or float32 instead.", "gemma3": "Numerical instability. Please use bfloat16 or float32 instead.", "gemma3_text": "Numerical instability. Please use bfloat16 or float32 instead.", - "plamo2": "Numerical instability. Please use bfloat16 or float32 instead.", "glm4": "Numerical instability. Please use bfloat16 or float32 instead.", } diff --git a/vllm/model_executor/models/plamo2.py b/vllm/model_executor/models/plamo2.py deleted file mode 100644 index 5fd925cf0be..00000000000 --- a/vllm/model_executor/models/plamo2.py +++ /dev/null @@ -1,992 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Inference-only PLaMo2 model.""" - -from collections.abc import Iterable -from itertools import islice -from typing import TYPE_CHECKING - -import torch -from torch import nn -from transformers import PretrainedConfig - -from vllm.compilation.decorators import support_torch_compile -from vllm.config import VllmConfig, get_current_vllm_config -from vllm.distributed import divide, get_tensor_model_parallel_world_size -from vllm.distributed.parallel_state import get_pp_group -from vllm.forward_context import ForwardContext, get_forward_context -from vllm.model_executor.custom_op import PluggableLayer -from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.linear import ( - ColumnParallelLinear, - MergedColumnParallelLinear, - QKVParallelLinear, - RowParallelLinear, -) -from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.mamba.abstract import MambaBase -from vllm.model_executor.layers.mamba.mamba_utils import ( - MambaStateCopyFunc, - MambaStateCopyFuncCalculator, - MambaStateDtypeCalculator, - MambaStateShapeCalculator, - is_conv_state_dim_first, -) -from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( - causal_conv1d_fn, - causal_conv1d_update, -) -from vllm.model_executor.layers.mamba.ops.ssd_combined import ( - mamba_chunk_scan_combined_varlen, -) -from vllm.model_executor.layers.mamba.ops.ssu_dispatch import selective_state_update -from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.model_executor.layers.rotary_embedding import get_rope -from vllm.model_executor.layers.vocab_parallel_embedding import ( - ParallelLMHead, - VocabParallelEmbedding, -) -from vllm.model_executor.model_loader.weight_utils import ( - composed_weight_loader, - default_weight_loader, - sharded_weight_loader, -) -from vllm.model_executor.models.interfaces import ( - HasInnerState, - IsHybrid, - SupportsLoRA, - SupportsPP, -) -from vllm.model_executor.models.utils import ( - AutoWeightsLoader, - is_pp_missing_parameter, - make_empty_intermediate_tensors_factory, - make_layers, - maybe_prefix, -) -from vllm.model_executor.utils import set_weight_attrs -from vllm.platforms import current_platform -from vllm.sequence import IntermediateTensors -from vllm.utils.torch_utils import direct_register_custom_op -from vllm.v1.attention.backend import AttentionMetadata -from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata -from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum - -# Only used for type hinting. -if TYPE_CHECKING: - - class Plamo2Config(PretrainedConfig): # type: ignore - model_type: str = "plamo2" - - hidden_size: int - num_hidden_layers: int - rms_norm_eps: float - # Attention - num_attention_heads: int - hidden_size_per_head: int - num_key_value_heads: int - # Mamba - mamba_d_state: int - mamba_d_conv: int - mamba_num_heads: int - mamba_step: int - # MLP - intermediate_size: int - # Tokenizer - vocab_size: int - - -def is_mamba(config: "Plamo2Config", i: int) -> bool: - assert config.mamba_step > 1 - - if config.num_hidden_layers <= (config.mamba_step // 2): - # use attention in last layer - return i != config.num_hidden_layers - 1 - return (i % config.mamba_step) != (config.mamba_step // 2) - - -# Adapted from: -# vllm.model_executor.layers.mamba.mamba_mixer2.MambaMixer2 -# transformers.models.mamba.modeling_mamba.MambaMixer -# --8<-- [start:plamo2_mamba_mixer] -@PluggableLayer.register("plamo2_mamba_mixer") -class Plamo2MambaMixer(MambaBase, PluggableLayer): - # --8<-- [end:plamo2_mamba_mixer] - - def __init__(self, vllm_config: VllmConfig, *, prefix: str = "", **kwargs) -> None: - super().__init__() - self.config = vllm_config.model_config.hf_config - self.cache_config = vllm_config.cache_config - self.model_config = vllm_config.model_config - self.quant_config = vllm_config.quant_config - self.is_lora_enabled = bool(vllm_config.lora_config) - self.hidden_size = self.config.hidden_size - self.ssm_state_size = self.config.mamba_d_state - self.conv_kernel_size = self.config.mamba_d_conv - self.intermediate_size = ( - self.config.mamba_num_heads * self.config.hidden_size_per_head - ) - self.tp_size = get_tensor_model_parallel_world_size() - self.head_dim = self.config.hidden_size_per_head - self.num_heads = self.config.mamba_num_heads - self.time_step_rank = max(64, self.hidden_size // 16) - self.conv1d = ColumnParallelLinear( - input_size=self.conv_kernel_size, - output_size=self.intermediate_size, - bias=False, - prefix=f"{prefix}.conv1d", - return_bias=False, - ) - # unsqueeze to fit conv1d weights shape into the linear weights shape. - # Can't do this in `weight_loader` since it already exists in - # `ColumnParallelLinear` and `set_weight_attrs` - # doesn't allow to override it - self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1) - - self.in_proj = MergedColumnParallelLinear( - self.hidden_size, - [self.intermediate_size] * 2, - bias=False, - quant_config=self.quant_config, - prefix=f"{prefix}.in_proj", - return_bias=False, - ) - # selective projection used to make dt, B and C input dependent - self.bcdt_proj = RowParallelLinear( - self.intermediate_size, - self.time_step_rank + self.ssm_state_size * 2, - bias=False, - quant_config=self.quant_config, - prefix=f"{prefix}.bcdt_proj", - return_bias=False, - ) - # time step projection (discretization) - - # In the forward we need to apply dt_proj without the bias, - # as the bias is added in the selective scan kernel. - self.dt_proj = ColumnParallelLinear( - self.time_step_rank, - self.num_heads, - bias=False, - quant_config=self.quant_config, - prefix=f"{prefix}.dt_proj", - return_bias=False, - ) - - self.A = nn.Parameter( - torch.empty( - divide(self.num_heads, self.tp_size), - dtype=torch.float32, - ) - ) - self.D = nn.Parameter(torch.ones(divide(self.num_heads, self.tp_size))) - self.dt_bias = nn.Parameter(torch.ones(divide(self.num_heads, self.tp_size))) - - set_weight_attrs(self.D, {"weight_loader": sharded_weight_loader(0)}) - a_weight_loader = composed_weight_loader( - sharded_weight_loader(0), lambda x: -torch.exp(x.float()) - ) - set_weight_attrs(self.A, {"weight_loader": a_weight_loader}) - set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)}) - - self.out_proj = RowParallelLinear( - self.intermediate_size, - self.hidden_size, - bias=False, - input_is_parallel=True, - quant_config=self.quant_config, - prefix=f"{prefix}.out_proj", - return_bias=False, - ) - # The activation function is fixed to SiLU. - self.activation = "silu" - - self.dt_norm = RMSNorm(self.time_step_rank, eps=self.config.rms_norm_eps) - self.B_norm = RMSNorm(self.ssm_state_size, eps=self.config.rms_norm_eps) - self.C_norm = RMSNorm(self.ssm_state_size, eps=self.config.rms_norm_eps) - - self.chunk_size = self.config.mamba_chunk_size - - compilation_config = get_current_vllm_config().compilation_config - if prefix in compilation_config.static_forward_context: - raise ValueError(f"Duplicate layer name: {prefix}") - compilation_config.static_forward_context[prefix] = self - # The tuple is (conv_state, ssm_state) - self.kv_cache = (torch.tensor([]), torch.tensor([])) - assert self.chunk_size != -1, "chunk_size must be set for v1" - - self.prefix = prefix - - def _project_ssm_parameters(self, hidden_states): - if self.is_lora_enabled: - # Lora kernel requires contiguous tensor. - ssm_parameters = self.bcdt_proj(hidden_states.contiguous()) - else: - ssm_parameters = self.bcdt_proj(hidden_states) - B, C, time_step = torch.split( - ssm_parameters, - [self.ssm_state_size, self.ssm_state_size, self.time_step_rank], - dim=-1, - ) - - # vllm._custom_ops.rms_norm requires contiguous input tensors. - time_step = self.dt_norm(time_step.contiguous()) - B = self.B_norm(B.contiguous()) - C = self.C_norm(C.contiguous()) - dt = self.dt_proj(time_step) - return B, C, dt - - def forward( - self, - hidden_states: torch.Tensor, - output: torch.Tensor, - **kwargs, - ): - torch.ops.vllm.plamo2_mamba_mixer( - hidden_states, - output, - self.prefix, - ) - - def forward_impl( - self, - hidden_states: torch.Tensor, - output: torch.Tensor, - **kwargs, - ): - forward_context = get_forward_context() - # attn_metadata contains metadata necessary for the mamba2 triton - # kernels to operate in continuous batching and in chunked prefill - # modes; they are computed at top-level model forward since they - # stay the same and reused for all mamba layers in the same iteration - attn_metadata: AttentionMetadata = forward_context.attn_metadata - - if attn_metadata is not None: - assert isinstance(attn_metadata, dict) - attn_metadata = attn_metadata[self.prefix] - assert isinstance(attn_metadata, Mamba2AttentionMetadata) - self_kv_cache = self.kv_cache - # conv_state = (..., dim, width-1) yet contiguous along 'dim' - # conv_state must be (..., dim, width-1) for the conv kernels. - # DS layout stores it that way directly; SD layout needs a transpose. - conv_state = ( - self_kv_cache[0] - if is_conv_state_dim_first() - else self_kv_cache[0].transpose(-1, -2) - ) - ssm_state = self_kv_cache[1] - state_indices_tensor_p = attn_metadata.state_indices_tensor_p - state_indices_tensor_d = attn_metadata.state_indices_tensor_d - has_initial_states_p = attn_metadata.has_initial_states_p - prep_initial_states = attn_metadata.prep_initial_states - chunk_size = attn_metadata.chunk_size - seq_idx_p = attn_metadata.seq_idx_p - query_start_loc_p = attn_metadata.query_start_loc_p - cu_chunk_seqlen_p = attn_metadata.cu_chunk_seqlen_p - last_chunk_indices_p = attn_metadata.last_chunk_indices_p - - # 1. Gated MLP's linear projection - projected_states = self.in_proj(hidden_states) - gate, hidden_states = projected_states.chunk(2, dim=-1) - - # 2. Convolution sequence transformation - conv_weights = self.conv1d.weight.view( - self.conv1d.weight.size(0), self.conv1d.weight.size(2) - ) - - if attn_metadata is None: - # profile run - hidden_states = ( - hidden_states.transpose(0, 1).clone().transpose(0, 1) - ).contiguous() - output[:] = self.out_proj(hidden_states) - return - - num_prefills = attn_metadata.num_prefills # request count - num_decodes = attn_metadata.num_decode_tokens # token count (=request) - num_prefill_tokens = attn_metadata.num_prefill_tokens # token count - has_prefill = num_prefills > 0 - has_decode = num_decodes > 0 - num_actual_tokens = num_prefill_tokens + num_decodes - - # Separate prefill and decode by splitting varlen input - # Split along token dimension - hidden_states_d, hidden_states_p = torch.split( - hidden_states[:num_actual_tokens], - [num_decodes, num_prefill_tokens], - dim=0, - ) - gate_d, gate_p = torch.split( - gate[:num_actual_tokens], [num_decodes, num_prefill_tokens], dim=0 - ) - # Preallocate output tensor to avoid memcpy cost for merging prefill - # and decode outputs - preallocated_ssm_out = torch.empty( - [ - num_prefill_tokens + num_decodes, - (self.num_heads // self.tp_size) * self.head_dim, - ], - dtype=hidden_states.dtype, - device=hidden_states.device, - ) - preallocated_ssm_out_d, preallocated_ssm_out_p = torch.split( - preallocated_ssm_out, - [num_decodes, num_prefill_tokens], - dim=0, - ) - - # Process prefill requests - if has_prefill: - # 2. Convolution sequence transformation - # - "cache_indices" updates the conv_state cache in positions - # pointed to by "state_indices_tensor_p" - x = hidden_states_p.transpose(0, 1) # this is the form that causal-conv see - hidden_states_p = causal_conv1d_fn( - x, - conv_weights, - self.conv1d.bias, - activation=self.activation, - conv_states=conv_state, - has_initial_state=has_initial_states_p, - cache_indices=state_indices_tensor_p, - metadata=attn_metadata, - query_start_loc=query_start_loc_p, - ) - hidden_states_p = hidden_states_p.transpose(0, 1) - hidden_states_p = hidden_states_p[:num_prefill_tokens] - # In some instances, the following `bcdt_proj` op - # requires contiguous inputs - # (e.g. if the Marlin kernel is used). - hidden_states_p = hidden_states_p.contiguous() - - B, C, dt = self._project_ssm_parameters(hidden_states_p) - - # 3. State Space Model sequence transformation - initial_states = None - if has_initial_states_p is not None and prep_initial_states: - # making a copy of the states - initial_states = torch.where( - has_initial_states_p[:, None, None, None], - ssm_state[state_indices_tensor_p], - 0, - ) - - varlen_state = mamba_chunk_scan_combined_varlen( - hidden_states_p.view( - num_prefill_tokens, self.num_heads // self.tp_size, self.head_dim - ), - dt, - self.A, - B.view(num_prefill_tokens, 1, -1), - C.view(num_prefill_tokens, 1, -1), - chunk_size=chunk_size, - D=self.D, - z=gate_p.view( - num_prefill_tokens, self.num_heads // self.tp_size, self.head_dim - ), - dt_bias=self.dt_bias, - seq_idx=seq_idx_p, - cu_seqlens=query_start_loc_p, - cu_chunk_seqlens=cu_chunk_seqlen_p, - last_chunk_indices=last_chunk_indices_p, - initial_states=initial_states, - dt_softplus=True, - dt_limit=(0.0, float("inf")), - out=preallocated_ssm_out_p.view(num_prefill_tokens, -1, self.head_dim), - state_dtype=ssm_state.dtype, - ) - - # update ssm states - # - varlen state is a (batch, nheads, headdim, dstate) tensor - ssm_state[state_indices_tensor_p] = varlen_state - - # Process decode requests - if has_decode: - # 2. Convolution sequence transformation - hidden_states_d = causal_conv1d_update( - hidden_states_d, - conv_state, - conv_weights, - self.conv1d.bias, - self.activation, - conv_state_indices=state_indices_tensor_d, - ) - - # ROCm: Ensure contiguous tensor for bcdt_proj linear layer. - # causal_conv1d_update returns a non-contiguous view (stride 8192 - # instead of 4096 for shape [batch, 4096]), causing incorrect GEMM - # results when batch > 1 on ROCm. - if current_platform.is_rocm(): - hidden_states_d = hidden_states_d.contiguous() - - B, C, dt = self._project_ssm_parameters(hidden_states_d) - - # 3. State Space Model sequence transformation - A = self.A[:, None, ...][:, :, None].expand( - -1, self.head_dim, self.config.mamba_d_state - ) - dt = dt[:, :, None].expand(-1, -1, self.head_dim) - dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim) - D = self.D[:, None, ...].expand(-1, self.head_dim) - B = B.unsqueeze(1) - C = C.unsqueeze(1) - hidden_states_d = hidden_states_d.view( - -1, self.num_heads // self.tp_size, self.head_dim - ) - - # - the hidden is reshaped into (bs, num_heads, head_dim) - # - ssm_state's slots will be selected - # using state_indices_tensor_d - - # NOTE: final output is an in-place update of out tensor - selective_state_update( - ssm_state, - hidden_states_d, - dt, - A, - B, - C, - D, - dt_bias, - z=gate_d.reshape(num_decodes, -1, self.head_dim), - dt_softplus=True, - state_batch_indices=state_indices_tensor_d, - out=preallocated_ssm_out_d.view(num_decodes, -1, self.head_dim), - ) - - # 4. Final linear projection - output[:num_actual_tokens] = self.out_proj(preallocated_ssm_out) - - def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]: - assert self.model_config is not None - assert self.cache_config is not None - return MambaStateDtypeCalculator.mamba2_state_dtype( - self.model_config.dtype, - self.cache_config.mamba_cache_dtype, - self.cache_config.mamba_ssm_cache_dtype, - ) - - def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]: - return MambaStateShapeCalculator.mamba2_state_shape( - intermediate_size=self.intermediate_size, - tp_world_size=get_tensor_model_parallel_world_size(), - n_groups=0, - num_heads=self.num_heads, - head_dim=self.head_dim, - state_size=self.ssm_state_size, - conv_kernel=self.conv_kernel_size, - ) - - @property - def mamba_type(self) -> MambaAttentionBackendEnum: - return MambaAttentionBackendEnum.MAMBA2 - - -def plamo2_mamba_mixer( - hidden_states: torch.Tensor, - output: torch.Tensor, - layer_name: str, -) -> None: - forward_context: ForwardContext = get_forward_context() - self = forward_context.no_compile_layers[layer_name] - self.forward_impl(hidden_states=hidden_states, output=output) - - -def plamo2_mamba_mixer_fake( - hidden_states: torch.Tensor, - output: torch.Tensor, - layer_name: str, -) -> None: - return - - -direct_register_custom_op( - op_name="plamo2_mamba_mixer", - op_func=plamo2_mamba_mixer, - mutates_args=["output"], - fake_impl=plamo2_mamba_mixer_fake, -) - - -class DenseMLP(nn.Module): - def __init__( - self, - config: "Plamo2Config", - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - super().__init__() - self.hidden_size = config.hidden_size - self.intermediate_size = config.intermediate_size - self.gate_up_proj = MergedColumnParallelLinear( - self.hidden_size, - [self.intermediate_size] * 2, - bias=False, - prefix=f"{prefix}.gate_up_proj", - quant_config=quant_config, - return_bias=False, - ) - self.act = SiluAndMul() - self.down_proj = RowParallelLinear( - self.intermediate_size, - self.hidden_size, - bias=False, - prefix=f"{prefix}.down_proj", - quant_config=quant_config, - return_bias=False, - ) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - h = self.gate_up_proj(hidden_states) - h = self.act(h) - return self.down_proj(h) - - -class Plamo2AttentionMixer(nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs) -> None: - super().__init__() - config = vllm_config.model_config.hf_config - cache_config = vllm_config.cache_config - quant_config = vllm_config.quant_config - self.hidden_size = config.hidden_size - tp_size = get_tensor_model_parallel_world_size() - self.total_num_heads = config.num_attention_heads - assert self.total_num_heads % tp_size == 0 - self.num_heads = self.total_num_heads // tp_size - self.total_num_kv_heads = config.num_key_value_heads - if self.total_num_kv_heads >= tp_size: - # Number of KV heads is greater than TP size, so we partition - # the KV heads across multiple tensor parallel GPUs. - assert self.total_num_kv_heads % tp_size == 0 - else: - # Number of KV heads is less than TP size, so we replicate - # the KV heads across multiple tensor parallel GPUs. - assert tp_size % self.total_num_kv_heads == 0 - self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) - self.head_dim = config.hidden_size_per_head - self.q_size = self.num_heads * self.head_dim - self.kv_size = self.num_kv_heads * self.head_dim - self.scaling = self.head_dim**-0.5 - - self.qkv_proj = QKVParallelLinear( - config.hidden_size, - self.head_dim, - self.total_num_heads, - self.total_num_kv_heads, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.qkv_proj", - ) - self.o_proj = RowParallelLinear( - self.total_num_heads * self.head_dim, - config.hidden_size, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.o_proj", - ) - - max_position = config.max_position_embeddings - if hasattr(vllm_config.model_config, "max_model_len") and isinstance( - vllm_config.model_config.max_model_len, int - ): - max_position = min(max_position, vllm_config.model_config.max_model_len) - - self.rotary_emb = get_rope( - self.head_dim, - max_position=max_position, - rope_parameters=config.rope_parameters, - ) - self.q_norm = RMSNorm(config.hidden_size_per_head, eps=config.rms_norm_eps) - self.q_norm.weight = torch.nn.Parameter( - torch.ones((self.num_heads, config.hidden_size_per_head)) - ) - set_weight_attrs( - self.q_norm.weight, {"weight_loader": sharded_weight_loader(0)} - ) - self.k_norm = RMSNorm(config.hidden_size_per_head, eps=config.rms_norm_eps) - self.k_norm.weight = torch.nn.Parameter( - torch.ones((self.num_kv_heads, config.hidden_size_per_head)) - ) - # Tensor-parallelism shards the K norm weights to the tp ranks - # in a head-wise manner. This approach does not work if there is only - # a single KV head, as is the case for PLaMo 2-1B. - if self.total_num_kv_heads != 1: - set_weight_attrs( - self.k_norm.weight, {"weight_loader": sharded_weight_loader(0)} - ) - - self.attn = Attention( - self.num_heads, - self.head_dim, - self.scaling, - num_kv_heads=self.num_kv_heads, - cache_config=cache_config, - prefix=f"{prefix}.attn", - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - **kwargs, - ) -> torch.Tensor: - qkv, _ = self.qkv_proj(hidden_states) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - - q_shape = q.shape - q = q.reshape(q_shape[:-1] + self.q_norm.weight.shape) - q = self.q_norm.forward_native(q).reshape(q_shape) - k_shape = k.shape - k = k.reshape(k_shape[:-1] + self.k_norm.weight.shape) - k = self.k_norm.forward_native(k).reshape(k_shape) - - q, k = self.rotary_emb(positions, q, k) - attn_output = self.attn(q, k, v) - output, _ = self.o_proj(attn_output) - return output - - -class Plamo2DecoderLayer(nn.Module): - def __init__( - self, vllm_config: VllmConfig, layer_idx: int, prefix: str = "", **kwargs - ) -> None: - super().__init__() - config = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - - self.is_mamba = is_mamba(config, layer_idx) - if self.is_mamba: - self.mixer = Plamo2MambaMixer( - vllm_config=vllm_config, prefix=f"{prefix}.mixer" - ) - else: - self.mixer = Plamo2AttentionMixer( - vllm_config=vllm_config, prefix=f"{prefix}.mixer" - ) - - self.mlp = DenseMLP( - config=config, quant_config=quant_config, prefix=f"{prefix}.mlp" - ) - self.pre_mixer_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_mixer_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.pre_mlp_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_mlp_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - residual: torch.Tensor | None, - **kwargs, - ): - if residual is None: - residual = hidden_states - hidden_states = self.pre_mixer_norm(hidden_states) - else: - hidden_states, residual = self.pre_mixer_norm(hidden_states, residual) - - if self.is_mamba: - # Plamo2MambaMixer writes output to this tensor - output = torch.empty_like(hidden_states) - mixer_kwargs = { - "output": output, - } - else: - mixer_kwargs = { - "positions": positions, - } - hidden_states = self.mixer( - hidden_states=hidden_states, - **mixer_kwargs, - ) - if self.is_mamba: - hidden_states = output - hidden_states = self.post_mixer_norm(hidden_states) - # Fully Connected - hidden_states, residual = self.pre_mlp_norm(hidden_states, residual) - hidden_states = self.mlp(hidden_states) - hidden_states = self.post_mlp_norm(hidden_states) - return hidden_states, residual - - -class Plamo2Decoder(torch.nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: - super().__init__() - config = vllm_config.model_config.hf_config - extra_kwargs = {"is_lora_enabled": bool(vllm_config.lora_config)} - - def get_layer(prefix: str): - layer_idx = int(prefix.rsplit(".", 1)[1]) - return Plamo2DecoderLayer( - vllm_config=vllm_config, - layer_idx=layer_idx, - prefix=prefix, - **extra_kwargs, - ) - - self.start_layer, self.end_layer, self.layers = make_layers( - config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers" - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - residual: torch.Tensor | None, - ) -> torch.Tensor: - for layer in islice(self.layers, self.start_layer, self.end_layer): - hidden_states, residual = layer( - positions=positions, - hidden_states=hidden_states, - residual=residual, - ) - return hidden_states, residual - - -@support_torch_compile -class Plamo2Model(torch.nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - - config = vllm_config.model_config.hf_config - - self.config = config - self.vocab_size = config.vocab_size - - self.embed_tokens = VocabParallelEmbedding( - self.vocab_size, - config.hidden_size, - prefix=f"{prefix}.embed_tokens", - ) - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( - ["hidden_states", "residual"], config.hidden_size - ) - self.layers = Plamo2Decoder(vllm_config=vllm_config, prefix=f"{prefix}.layers") - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor: - if get_pp_group().is_first_rank: - if inputs_embeds is not None: - hidden_states = inputs_embeds - else: - hidden_states = self.embed_input_ids(input_ids) - residual = None - else: - assert intermediate_tensors is not None - hidden_states = intermediate_tensors["hidden_states"] - residual = intermediate_tensors["residual"] - - hidden_states, residual = self.layers( - positions=positions, - hidden_states=hidden_states, - residual=residual, - ) - if not get_pp_group().is_last_rank: - return IntermediateTensors( - {"hidden_states": hidden_states, "residual": residual} - ) - hidden_states, _ = self.norm(hidden_states, residual) - return hidden_states - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - # Update the weight names to be compatible with the vllm version - # of the model. - # Do not change the order of the replacements. - replacements = { - # Rename incompatible weight names. - ".A_log": ".A", - ".B_norm_weight": ".B_norm.weight", - ".C_norm_weight": ".C_norm.weight", - ".dt_norm_weight": ".dt_norm.weight", - ".q_weight": ".q_norm.weight", - ".k_weight": ".k_norm.weight", - } - # Apply replacements based on the defined mappings - for old, new in replacements.items(): - if old in name: - name = name.replace(old, new) - - # Reshape the in_proj weights to match the shape expected - # by MergedColumnParallelLinear. - # This works both for unquantized weights and - # for quantized weights. - # In the quantized case, the weights are already transposed. - # Also, in addition to the quantized weights, - # the zero points and scales have to be reshaped as well. - # Packing should not be affected by this. - if ( - ".mixer.in_proj.weight" in name - or "mixer.in_proj.qweight" in name - or "mixer.in_proj.scales" in name - or "mixer.in_proj.qzeros" in name - ): - if "mixer.in_proj.weight" in name: - loaded_weight = loaded_weight.transpose(0, 1) - # for weight: - # loaded_weight.shape[0] == self.config.hidden_size - # for qweight: - # loaded_weight.shape[0] == self.config.hidden_size // param.pack_factor # noqa - # for scales and qzeros: - # loaded_weight.shape[0] == self.config.hidden_size // self.vllm_config.quant_config.group_size # noqa - loaded_weight = loaded_weight.reshape( - loaded_weight.shape[0], self.config.mamba_num_heads, -1 - ) - gate_weight, hidden_states_weight = loaded_weight.chunk(2, dim=-1) - gate_weight = gate_weight.reshape(loaded_weight.shape[0], -1) - hidden_states_weight = hidden_states_weight.reshape( - loaded_weight.shape[0], -1 - ) - loaded_weight = torch.cat([gate_weight, hidden_states_weight], dim=-1) - if "mixer.in_proj.weight" in name: - loaded_weight = loaded_weight.transpose(0, 1) - - # Offset parameter with vllm's RMSNorm haven't been supported yet. - if ".pre_mixer_norm" in name: - loaded_weight += 1.0 - elif ".post_mixer_norm" in name: - loaded_weight += 1.0 / 5 - elif ".pre_mlp_norm" in name: - loaded_weight += 1.0 - elif ".post_mlp_norm" in name: - loaded_weight += 1.0 / (5**1.5) - elif name == "norm.weight": - loaded_weight += 1.0 - - # Skip layers on other devices. - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - - -class Plamo2ForCausalLM( - torch.nn.Module, HasInnerState, SupportsLoRA, SupportsPP, IsHybrid -): - packed_modules_mapping = { - "qkv_proj": ["qkv_proj"], - "gate_up_proj": ["gate_up_proj"], - "in_proj": ["in_proj"], - } - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: - super().__init__() - config = vllm_config.model_config.hf_config - scheduler_config = vllm_config.scheduler_config - - self.config = config - self.vllm_config = vllm_config - self.model_config = vllm_config.model_config - self.scheduler_config = scheduler_config - - # ModelConfig.get_head_size assumes head_dim is set or calculated as - # hidden_size // num_attention_heads. However, this is not always - # the case for PLaMo2, as indicated by the FIXME comment. - self.config.head_dim = self.config.hidden_size_per_head - - self.model = Plamo2Model( - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") - ) - self.vocab_size = self.config.vocab_size - self.lm_head = ParallelLMHead( - self.vocab_size, - self.config.hidden_size, - prefix=f"{prefix}.lm_head", - ) - if self.config.tie_word_embeddings: - self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) - - self.logits_processor = LogitsProcessor( - config.vocab_size, self.config.vocab_size - ) - self.make_empty_intermediate_tensors = ( - self.model.make_empty_intermediate_tensors - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.model.embed_input_ids(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - **kwargs, - ): - hidden_states = self.model( - input_ids, positions, intermediate_tensors, inputs_embeds - ) - return hidden_states - - @classmethod - def get_mamba_state_dtype_from_config( - cls, - vllm_config: "VllmConfig", - ) -> tuple[torch.dtype, torch.dtype]: - return MambaStateDtypeCalculator.mamba2_state_dtype( - vllm_config.model_config.dtype, - vllm_config.cache_config.mamba_cache_dtype, - vllm_config.cache_config.mamba_ssm_cache_dtype, - ) - - @classmethod - def get_mamba_state_shape_from_config( - cls, - vllm_config: "VllmConfig", - ) -> tuple[tuple[int, int], tuple[int, int, int]]: - """Calculate shapes for Mamba's convolutional and state caches. - Args: - vllm_config: vLLM config - Returns: - Tuple containing: - - conv_state_shape: Shape for convolutional state cache - - temporal_state_shape: Shape for state space model cache - """ - parallel_config = vllm_config.parallel_config - hf_config = vllm_config.model_config.hf_config - intermediate_size = hf_config.mamba_num_heads * hf_config.hidden_size_per_head - - return MambaStateShapeCalculator.mamba2_state_shape( - intermediate_size=intermediate_size, - tp_world_size=parallel_config.tensor_parallel_size, - n_groups=0, - num_heads=hf_config.mamba_num_heads, - head_dim=hf_config.hidden_size_per_head, - state_size=hf_config.mamba_d_state, - conv_kernel=hf_config.mamba_d_conv, - ) - - @classmethod - def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: - return MambaStateCopyFuncCalculator.mamba2_state_copy_func() - - def compute_logits( - self, - hidden_states: torch.Tensor, - ) -> torch.Tensor | None: - logits = self.logits_processor(self.lm_head, hidden_states) - return logits - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) - return loader.load_weights(weights) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index a4fd44fc2cd..947e0d043dd 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -191,7 +191,6 @@ _TEXT_GENERATION_MODELS = { "PhiForCausalLM": ("phi", "PhiForCausalLM"), "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"), "PhiMoEForCausalLM": ("phimoe", "PhiMoEForCausalLM"), - "Plamo2ForCausalLM": ("plamo2", "Plamo2ForCausalLM"), "Plamo3ForCausalLM": ("plamo3", "Plamo3ForCausalLM"), "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"), "Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"), @@ -761,6 +760,7 @@ _PREVIOUSLY_SUPPORTED_MODELS = { "TeleChatForCausalLM": "0.25.0", "PersimmonForCausalLM": "0.25.0", "FuyuForCausalLM": "0.25.0", + "Plamo2ForCausalLM": "0.26.0", } _OOT_SUPPORTED_MODELS = { From 89f6aa3a9e56a9228f55e8d722f23594223187cf Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Fri, 24 Jul 2026 21:56:57 +0100 Subject: [PATCH 017/185] [KV Offload][CI] Fall back to buffered I/O without O_DIRECT; fix flaky api-server test (#49734) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- .../unit_tests/_api_server_spawn_workers.py | 14 +++++++ .../test_api_server_process_manager.py | 9 ++++- tests/v1/kv_offload/tiering/test_fs_tier.py | 37 ++++++++++++++++++ vllm/v1/kv_offload/tiering/fs/io.py | 38 ++++++++++++++++++- vllm/v1/kv_offload/tiering/fs/manager.py | 16 +++++++- 5 files changed, 109 insertions(+), 5 deletions(-) create mode 100644 tests/entrypoints/unit_tests/_api_server_spawn_workers.py diff --git a/tests/entrypoints/unit_tests/_api_server_spawn_workers.py b/tests/entrypoints/unit_tests/_api_server_spawn_workers.py new file mode 100644 index 00000000000..c8ff969a8a3 --- /dev/null +++ b/tests/entrypoints/unit_tests/_api_server_spawn_workers.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Spawn worker targets kept free of heavy imports. + +``multiprocessing`` with the ``spawn`` start method re-imports the module that +defines a process target in the child. Housing these stubs in a stdlib-only +module keeps child startup fast and deterministic, instead of paying a multi- +second ``import vllm`` before the child can run. +""" + + +def exit_before_report_worker(listen_address, sock, args, client_config=None): + """Exit immediately without touching ``actual_address_pipe``.""" + return diff --git a/tests/entrypoints/unit_tests/test_api_server_process_manager.py b/tests/entrypoints/unit_tests/test_api_server_process_manager.py index ada7a8797fa..902d0ce2d74 100644 --- a/tests/entrypoints/unit_tests/test_api_server_process_manager.py +++ b/tests/entrypoints/unit_tests/test_api_server_process_manager.py @@ -10,6 +10,9 @@ from unittest.mock import patch import pytest import zmq +from tests.entrypoints.unit_tests._api_server_spawn_workers import ( + exit_before_report_worker, +) from vllm.utils.network_utils import make_zmq_socket, split_zmq_path from vllm.v1.utils import ( APIServerProcessManager, @@ -228,6 +231,7 @@ def test_normal_completion(api_server_args): def test_external_process_monitoring(api_server_args): """Test that wait_for_completion_or_failure handles additional processes.""" global WORKER_RUNTIME_SECONDS + prev_worker_runtime = WORKER_RUNTIME_SECONDS WORKER_RUNTIME_SECONDS = 100 # Create and start the external process @@ -307,6 +311,7 @@ def test_external_process_monitoring(api_server_args): manager.shutdown() mock_coordinator.shutdown() time.sleep(0.2) + WORKER_RUNTIME_SECONDS = prev_worker_runtime @pytest.mark.timeout(60) @@ -383,10 +388,10 @@ def test_gather_actual_addresses_child_crash_before_report(): num_servers=num_servers, input_addresses=placeholder_inputs, output_addresses=placeholder_outputs, - # mock_run_api_server_worker exits without touching + # exit_before_report_worker exits without touching # ``actual_address_pipe`` — simulates a child that dies before # reporting its bound addresses. - target_server_fn=mock_run_api_server_worker, + target_server_fn=exit_before_report_worker, ) try: # Sentinel-first vs pipe-EOF-first both produce "reporting". diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index dcc92a4fa8b..589d7dda779 100644 --- a/tests/v1/kv_offload/tiering/test_fs_tier.py +++ b/tests/v1/kv_offload/tiering/test_fs_tier.py @@ -379,6 +379,43 @@ def test_store_load_data_integrity(fs_tier): ) +def test_store_load_roundtrip_without_o_direct(tmp_path, monkeypatch): + """Buffered fallback must round-trip data when O_DIRECT is unsupported. + + Simulates filesystems (e.g. overlayfs, some NFS) that reject O_DIRECT by + forcing the capability probe to report it unavailable. + """ + monkeypatch.setattr( + "vllm.v1.kv_offload.tiering.fs.manager.probe_o_direct", + lambda _dir: False, + ) + tensor = _page_aligned_rand_tensor(4, _BLOCK_ELEMENTS) + tier = FileSystemTierManager( + offloading_spec=_MOCK_OFFLOADING_SPEC, + primary_kv_view=memoryview(tensor.numpy()), + tier_type="fs", + root_dir=str(tmp_path), + n_read_threads=4, + n_write_threads=4, + ) + try: + assert tier._use_o_direct is False + + keys = [key(0), key(1)] + expected = tensor[:2].clone() + tier.submit_store(make_job(1, keys, [0, 1])) + assert all(r.success for r in drain(tier)) + + tensor[:2] = 0.0 + tier.submit_load(make_job(2, keys, [2, 3], is_promotion=True)) + assert all(r.success for r in drain(tier)) + + for i, bid in enumerate([2, 3]): + assert torch.allclose(tensor[bid], expected[i]) + finally: + tier.shutdown() + + def test_wait_idle_blocks_until_tasks_complete(): """wait_idle must not return while a task is still in flight.""" pool = DualQueueThreadPool(n_read_threads=1, n_write_threads=1) diff --git a/vllm/v1/kv_offload/tiering/fs/io.py b/vllm/v1/kv_offload/tiering/fs/io.py index c5a82a73c67..7785591fd2f 100644 --- a/vllm/v1/kv_offload/tiering/fs/io.py +++ b/vllm/v1/kv_offload/tiering/fs/io.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import contextlib import logging +import mmap import os import random import threading @@ -24,6 +26,34 @@ def _get_tmp_suffix() -> str: return _thread_local.tmp_suffix +def probe_o_direct(directory: str) -> bool: + """Return whether ``O_DIRECT`` I/O works in *directory*. + + ``O_DIRECT`` is unsupported on some filesystems (e.g. the overlayfs backing + a container ``/tmp``, older tmpfs, or some NFS mounts), where opening or + writing a file with it fails with ``EINVAL``. Probe once with an aligned + single-page write so callers can fall back to buffered I/O instead of + failing on every block. + """ + if not O_DIRECT: + return False + path = os.path.join(directory, f".o_direct_probe{_get_tmp_suffix()}") + page = mmap.mmap(-1, mmap.PAGESIZE) + try: + fd = os.open(path, os.O_CREAT | os.O_WRONLY | os.O_TRUNC | O_DIRECT, 0o644) + try: + os.write(fd, page) + finally: + os.close(fd) + return True + except OSError: + return False + finally: + page.close() + with contextlib.suppress(OSError): + os.remove(path) + + def _ensure_dirs(path: str) -> None: """Create parent directories of *path* if they don't exist.""" os.makedirs(os.path.dirname(path), exist_ok=True) @@ -34,6 +64,7 @@ def store_block( buffer: memoryview, offset: int, block_size: int, + use_o_direct: bool = True, ) -> None: """ Store callback: Writes to a temp file then atomically replaces the destination. @@ -49,10 +80,11 @@ def store_block( # Write block atomically. Cast to a flat byte view so the slice uses byte # indices; the raw memoryview may be multi-dimensional with itemsize > 1. view_slice = buffer.cast("B")[offset : offset + block_size] + o_direct = O_DIRECT if use_o_direct else 0 try: fd = os.open( tmp_path, - os.O_CREAT | os.O_EXCL | os.O_WRONLY | os.O_TRUNC | O_DIRECT, + os.O_CREAT | os.O_EXCL | os.O_WRONLY | os.O_TRUNC | o_direct, 0o644, ) try: @@ -77,14 +109,16 @@ def load_block( view: memoryview, offset: int, block_size: int, + use_o_direct: bool = True, ) -> None: """ Load callback: read one KV block from disk. Remove the file on failure. """ fd: int | None = None view_slice = view.cast("B")[offset : offset + block_size] + o_direct = O_DIRECT if use_o_direct else 0 try: - fd = os.open(source_path, os.O_RDONLY | O_DIRECT) + fd = os.open(source_path, os.O_RDONLY | o_direct) bytes_read = os.readv(fd, [view_slice]) if bytes_read < block_size: raise OSError(f"Short read: expected {block_size} bytes, read {bytes_read}") diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index f12ef2c6d4d..4fde24d760e 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -49,7 +49,7 @@ from vllm.v1.kv_offload.tiering.base import ( ScheduleEndContext, SecondaryTierManager, ) -from vllm.v1.kv_offload.tiering.fs.io import load_block, store_block +from vllm.v1.kv_offload.tiering.fs.io import load_block, probe_o_direct, store_block from vllm.v1.kv_offload.tiering.fs.thread_pool import DualQueueThreadPool if TYPE_CHECKING: @@ -168,6 +168,18 @@ class FileSystemTierManager(SecondaryTierManager): self.file_mapper.get_run_config(), f, indent=2, sort_keys=True ) + # Prefer O_DIRECT to bypass the page cache, but fall back to buffered + # I/O on filesystems that reject it (e.g. overlayfs, some NFS mounts) + # rather than failing every block. + self._use_o_direct = probe_o_direct(os.path.dirname(config_path)) + if not self._use_o_direct: + logger.warning( + "O_DIRECT is not supported at '%s'; falling back to buffered " + "I/O for the '%s' KV offload tier.", + root_dir, + tier_type, + ) + self._pool = DualQueueThreadPool( n_read_threads, n_write_threads, @@ -198,6 +210,7 @@ class FileSystemTierManager(SecondaryTierManager): self._primary_kv_view, int(bid) * self._block_size, self._block_size, + self._use_o_direct, ) for key, bid in zip(job_metadata.keys, job_metadata.block_ids) ) @@ -212,6 +225,7 @@ class FileSystemTierManager(SecondaryTierManager): self._primary_kv_view, int(bid) * self._block_size, self._block_size, + self._use_o_direct, ) for key, bid in zip(job_metadata.keys, job_metadata.block_ids) ) From 7513d071bd749c938181fac9e67ac103d0c52dc1 Mon Sep 17 00:00:00 2001 From: djramic Date: Fri, 24 Jul 2026 23:09:18 +0200 Subject: [PATCH 018/185] [ROCm][CI] Fix XPASS(strict) on mixed audio embeds test (#49733) Signed-off-by: Djordje Ramic Signed-off-by: Andreas Karatzas Co-authored-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- .../test_chat_completion_with_mixed_audio_embeds.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py b/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py index fd66d2d57fa..a417f860f61 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py @@ -17,6 +17,7 @@ from transformers import AutoConfig, AutoTokenizer from tests.utils import RemoteOpenAIServer from vllm.utils.serial_utils import tensor2base64 +from vllm.utils.torch_utils import is_torch_equal_or_newer QWEN2AUDIO_MODEL = "Qwen/Qwen2-Audio-7B-Instruct" @@ -148,6 +149,7 @@ def qwen2audio_aligned_content_and_embeds_b64() -> tuple[str, str]: False, id="text-then-audio_embeds", marks=pytest.mark.xfail( + condition=is_torch_equal_or_newer("2.12.0"), reason="torch 2.12 regression: prompt_embeds output diverges " "from raw-text when text precedes audio; " "https://github.com/pytorch/pytorch/issues/184431", From caa9cad31ea2e577ac34f8695014f81f53c89e26 Mon Sep 17 00:00:00 2001 From: Rohan Potdar Date: Fri, 24 Jul 2026 16:14:03 -0500 Subject: [PATCH 019/185] [ROCm][Docker] Drop MORI_GPU_ARCHS so MoRI autodetects the device arch (#49737) Signed-off-by: Rohan Potdar Co-authored-by: Claude Opus 4.8 (1M context) --- docker/Dockerfile.rocm_base | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index 85f48ead469..a22bad04155 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -30,7 +30,7 @@ ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib: ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151 ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} ENV AITER_ROCM_ARCH=gfx942;gfx950 -ENV MORI_GPU_ARCHS=gfx942;gfx950 +# Note: Do not set MORI_GPU_ARCHS here, it is automatically inferred at runtime # Required for RCCL in ROCm7.1 ENV HSA_NO_SCRATCH_RECLAIM=1 From 33c4f3551ce9b4dc75864f16c40496d8d64f8e9d Mon Sep 17 00:00:00 2001 From: Aarushi Jain <142941703+aarushjain29@users.noreply.github.com> Date: Fri, 24 Jul 2026 17:19:46 -0500 Subject: [PATCH 020/185] =?UTF-8?q?[ROCm][CI]=20Wait=20for=20ROCm=20VRAM?= =?UTF-8?q?=20to=20settle=20between=20compiled=20and=20eager=20LL=E2=80=A6?= =?UTF-8?q?=20(#49739)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: aarushjain29 --- tests/compile/test_dynamic_shapes_compilation.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/compile/test_dynamic_shapes_compilation.py b/tests/compile/test_dynamic_shapes_compilation.py index 96c3f49aba3..7f725c14f21 100644 --- a/tests/compile/test_dynamic_shapes_compilation.py +++ b/tests/compile/test_dynamic_shapes_compilation.py @@ -9,6 +9,7 @@ import pytest import torch from tests.models.utils import check_logprobs_close +from tests.utils import wait_for_rocm_memory_to_settle from vllm import LLM, SamplingParams from vllm.compilation.decorators import support_torch_compile from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config @@ -104,6 +105,7 @@ def test_dynamic_shapes_compilation( gc.collect() torch.accelerator.empty_cache() torch.accelerator.synchronize() + wait_for_rocm_memory_to_settle() eager_model = LLM(model=model_name, enforce_eager=True, max_model_len=1024) eager_outputs = [] From 213f681f8117b9026ca8189583d2855d70104401 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Fri, 24 Jul 2026 16:04:59 -0700 Subject: [PATCH 021/185] Revert "[Perf][GLM-5.2] Blackwell decode optimizations" (#49768) --- CMakeLists.txt | 17 - csrc/libtorch_stable/bf16_skinny_gemm.cu | 262 --------- .../libtorch_stable/bf16_skinny_gemm_entry.cu | 170 ------ csrc/libtorch_stable/dsv3_fused_a_gemm.cu | 148 ++--- .../quantization/fp4/nvfp4_quant_kernels.cu | 52 +- csrc/libtorch_stable/torch_bindings.cpp | 4 - recipes/glm5.2-ll-b300-tp8-mtp5.md | 41 -- tests/kernels/test_bf16_skinny_gemm.py | 70 --- .../test_fused_deepseek_v32_norm_rope.py | 77 --- tests/models/registry.py | 4 - vllm/_custom_ops.py | 25 - .../passes/fusion/allreduce_rms_fusion.py | 146 ++--- vllm/config/speculative.py | 8 +- vllm/cute_utils/__init__.py | 5 - vllm/cute_utils/cvt.py | 23 +- .../layers/fused_allreduce_gemma_rms_norm.py | 19 +- vllm/model_executor/models/registry.py | 3 +- vllm/models/deepseek_v32/__init__.py | 22 +- vllm/models/deepseek_v32/nvidia/attention.py | 84 +-- vllm/models/deepseek_v32/nvidia/kernels.py | 300 +--------- vllm/models/deepseek_v32/nvidia/model.py | 4 - vllm/models/deepseek_v32/nvidia/mtp.py | 82 +-- .../deepseek_v32/nvidia/ops/__init__.py | 2 - .../nvidia/ops/fused_q_cutedsl.py | 532 ------------------ .../backends/mla/flashinfer_mla_sparse.py | 31 - .../v1/attention/backends/mla/sparse_utils.py | 60 +- vllm/v1/spec_decode/llm_base_proposer.py | 5 +- .../spec_decode/autoregressive/speculator.py | 40 +- .../worker/gpu/spec_decode/mtp/speculator.py | 41 +- 29 files changed, 138 insertions(+), 2139 deletions(-) delete mode 100644 csrc/libtorch_stable/bf16_skinny_gemm.cu delete mode 100644 csrc/libtorch_stable/bf16_skinny_gemm_entry.cu delete mode 100644 recipes/glm5.2-ll-b300-tp8-mtp5.md delete mode 100644 tests/kernels/test_bf16_skinny_gemm.py delete mode 100644 vllm/models/deepseek_v32/nvidia/ops/__init__.py delete mode 100644 vllm/models/deepseek_v32/nvidia/ops/fused_q_cutedsl.py diff --git a/CMakeLists.txt b/CMakeLists.txt index d91113bd1c2..9bd73f18e70 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -732,23 +732,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "(requires SM90+ and CUDA >= 12.0).") endif() - # BF16 skinny GEMM (M<=32; weight-bandwidth-bound decode shapes). - # Requires SM90+. - cuda_archs_sm90plus(BF16_SKINNY_GEMM_ARCHS "${CUDA_ARCHS}") - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND BF16_SKINNY_GEMM_ARCHS) - set(BF16_SKINNY_GEMM_SRCS - "csrc/libtorch_stable/bf16_skinny_gemm_entry.cu" - "csrc/libtorch_stable/bf16_skinny_gemm.cu") - set_gencode_flags_for_srcs( - SRCS "${BF16_SKINNY_GEMM_SRCS}" - CUDA_ARCHS "${BF16_SKINNY_GEMM_ARCHS}") - list(APPEND VLLM_STABLE_EXT_SRC "${BF16_SKINNY_GEMM_SRCS}") - message(STATUS "Building bf16_skinny_gemm for archs: ${BF16_SKINNY_GEMM_ARCHS}") - else() - message(STATUS "Not building bf16_skinny_gemm as no compatible archs found " - "(requires SM90+ and CUDA >= 12.0).") - endif() - # Only build AllSpark kernels if we are building for at least some compatible archs. cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}") if (ALLSPARK_ARCHS) diff --git a/csrc/libtorch_stable/bf16_skinny_gemm.cu b/csrc/libtorch_stable/bf16_skinny_gemm.cu deleted file mode 100644 index fef90d29d61..00000000000 --- a/csrc/libtorch_stable/bf16_skinny_gemm.cu +++ /dev/null @@ -1,262 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project -// -// Skinny GEMM: activation(bf16) x weight(bf16)^T -> bf16, for decode-time -// M <= 32 with a large reduction dim. Replaces cuBLAS splitK (GEMM + -// splitKreduce) with a single block-per-output-column kernel; these shapes -// are weight-bandwidth-bound, so one coalesced pass over the weight at -// fp32 accumulation is optimal. Adapted from fp32_router_gemm.cu. -// -// First user: the DeepSeek-V32/GLM-5.2 MTP eh_proj (K=2*hidden=12288, -// N=hidden/TP), whose cuBLAS splitK pick costs ~34.6us vs the ~19us -// bandwidth floor per replicated read (and ~4us once column-parallel). - -#include -#include - -// --------------------------------------------------------------------------- -// Load helpers (8 x bf16 = one uint4 load, converted to fp32) -// --------------------------------------------------------------------------- - -namespace skinny { - -constexpr int VPT = 8; // bf16 values per thread per load - -__device__ __forceinline__ void load_bf16x8(__nv_bfloat16 const* ptr, - float* dst) { - uint4 v = *reinterpret_cast(ptr); - __nv_bfloat16 const* p = reinterpret_cast<__nv_bfloat16 const*>(&v); -#pragma unroll - for (int i = 0; i < VPT; i++) dst[i] = __bfloat162float(p[i]); -} - -// Streaming variant for the weight: each row is read exactly once across the -// whole grid, so bypass L2 residency (evict-first). Measured -1.4us at M=1. -__device__ __forceinline__ void load_bf16x8_cs(__nv_bfloat16 const* ptr, - float* dst) { - uint4 v = __ldcs(reinterpret_cast(ptr)); - __nv_bfloat16 const* p = reinterpret_cast<__nv_bfloat16 const*>(&v); -#pragma unroll - for (int i = 0; i < VPT; i++) dst[i] = __bfloat162float(p[i]); -} - -// --------------------------------------------------------------------------- -// Kernel: each block computes kNPB output columns for all kNumTokens rows. -// grid = kN / kNPB, block = kBlockSize threads. K is reduced VPT elements -// per thread per iteration; fp32 accumulation, warp butterfly + smem -// finalize, bf16 store. -// --------------------------------------------------------------------------- - -template -__global__ __launch_bounds__(kBlockSize, 1) void bf16_skinny_gemm_kernel( - __nv_bfloat16* out, __nv_bfloat16 const* mat_a, __nv_bfloat16 const* mat_b, - int64_t out_stride) { - constexpr int k_elems_per_iter = VPT * kBlockSize; - constexpr int k_iterations = kK / k_elems_per_iter; - static_assert(kK % k_elems_per_iter == 0); - constexpr int kWarpSize = 32; - constexpr int kNumWarps = kBlockSize / kWarpSize; - - int const n_base = blockIdx.x * kNPB; - int const tid = threadIdx.x; - int const warpId = tid / kWarpSize; - int const laneId = tid % kWarpSize; - - float acc[kNumTokens][kNPB] = {}; - __shared__ float sm_reduction[kNumTokens][kNPB][kNumWarps]; - - // Register prefetch (kPF > 0): W does not depend on the predecessor, so - // the first kPF iterations' weight chunks are loaded raw BEFORE the - // dependency sync; with a PDL-releasing producer (fused_eh_norm fires - // gdc_launch_dependents early) these DRAM round trips overlap the norm. - // Pair-measured on B300 (norm+gemm in one graph, full 6144x12288): - // M=1 pf2 28.43us vs pf0 29.00us; deeper prefetch or M >= 2 regresses - // (register pressure), hence the per-M selection in the launcher. - uint4 w_pre[kPF > 0 ? kPF : 1][kNPB]; -#pragma unroll - for (int pf = 0; pf < kPF; pf++) { -#pragma unroll - for (int n = 0; n < kNPB; n++) { - w_pre[pf][n] = - *reinterpret_cast(mat_b + (size_t)(n_base + n) * kK + - pf * k_elems_per_iter + tid * VPT); - } - } - -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - cudaGridDependencySynchronize(); -#endif - -#pragma unroll - for (int ki = 0; ki < k_iterations; ki++) { - int const k_base = ki * k_elems_per_iter + tid * VPT; - - float b_float[kNPB][VPT]; - if (ki < kPF) { -#pragma unroll - for (int n = 0; n < kNPB; n++) { - __nv_bfloat16 const* p = - reinterpret_cast<__nv_bfloat16 const*>(&w_pre[ki][n]); -#pragma unroll - for (int v = 0; v < VPT; v++) b_float[n][v] = __bfloat162float(p[v]); - } - } else { -#pragma unroll - for (int n = 0; n < kNPB; n++) { - load_bf16x8_cs(mat_b + (size_t)(n_base + n) * kK + k_base, b_float[n]); - } - } - -#pragma unroll - for (int m = 0; m < kNumTokens; m++) { - float a_float[VPT]; - load_bf16x8(mat_a + (size_t)m * kK + k_base, a_float); -#pragma unroll - for (int n = 0; n < kNPB; n++) { -#pragma unroll - for (int k = 0; k < VPT; k++) { - acc[m][n] += a_float[k] * b_float[n][k]; - } - } - } - } - - // Warp-level butterfly reduction -#pragma unroll - for (int m = 0; m < kNumTokens; m++) { -#pragma unroll - for (int n = 0; n < kNPB; n++) { - float sum = acc[m][n]; - sum += __shfl_xor_sync(0xffffffff, sum, 16); - sum += __shfl_xor_sync(0xffffffff, sum, 8); - sum += __shfl_xor_sync(0xffffffff, sum, 4); - sum += __shfl_xor_sync(0xffffffff, sum, 2); - sum += __shfl_xor_sync(0xffffffff, sum, 1); - if (laneId == 0) sm_reduction[m][n][warpId] = sum; - } - } - - __syncthreads(); - - // Parallel finalize: one thread per (m, n) output. - for (int idx = tid; idx < kNumTokens * kNPB; idx += kBlockSize) { - int const m = idx / kNPB; - int const n = idx % kNPB; - float final_sum = 0.0f; -#pragma unroll - for (int w = 0; w < kNumWarps; w++) final_sum += sm_reduction[m][n][w]; - out[(size_t)m * out_stride + n_base + n] = __float2bfloat16(final_sum); - } - - // Trigger after our stores: harmless hardening, not a guarantee — the - // trigger only permits dependent-launch scheduling and carries no memory - // visibility semantics, so a PDL consumer must still gridsync before - // reading our output. Every current consumer is a plain launch (full - // stream order); firing late just avoids an unnecessarily early launch - // window and matches fp32_router_gemm.cu. -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - cudaTriggerProgrammaticLaunchCompletion(); -#endif -} - -} // namespace skinny - -// --------------------------------------------------------------------------- -// Launcher -// --------------------------------------------------------------------------- - -template -void invokeBf16SkinnyGemm(__nv_bfloat16* output, __nv_bfloat16 const* mat_a, - __nv_bfloat16 const* mat_b, int64_t out_stride, - cudaStream_t stream) { - static_assert(kN % kNPB == 0); - // Weight prefetch depth: only M=1 measured a win (see kernel comment). - constexpr int kPF = (kNumTokens == 1) ? 2 : 0; - cudaLaunchConfig_t config; - config.gridDim = kN / kNPB; - config.blockDim = kBlockSize; - config.dynamicSmemBytes = 0; - config.stream = stream; - cudaLaunchAttribute attrs[1]; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = 1; - config.numAttrs = 1; - config.attrs = attrs; - cudaLaunchKernelEx(&config, - skinny::bf16_skinny_gemm_kernel, - output, mat_a, mat_b, out_stride); -} - -// --------------------------------------------------------------------------- -// Explicit instantiations. M = 1..32; (N, K) pairs: -// (768, 12288) eh_proj shard, TP8 -// (1536, 12288) eh_proj shard, TP4 -// (6144, 12288) eh_proj unsharded -// kNPB (B300 sweep, M=1): 6144 -> 2 (3072 blocks, 23.2us vs 25.3 at kNPB=8; -// narrow blocks minimize wave quantization); shards 768/1536 keep 4. -// --------------------------------------------------------------------------- - -#define INSTANTIATE(M, NPB, N, K) \ - template void invokeBf16SkinnyGemm<128, NPB, M, N, K>( \ - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int64_t, \ - cudaStream_t); - -#define INSTANTIATE_ALL_M(NPB, N, K) \ - INSTANTIATE(1, NPB, N, K) \ - INSTANTIATE(2, NPB, N, K) \ - INSTANTIATE(3, NPB, N, K) \ - INSTANTIATE(4, NPB, N, K) \ - INSTANTIATE(5, NPB, N, K) \ - INSTANTIATE(6, NPB, N, K) \ - INSTANTIATE(7, NPB, N, K) \ - INSTANTIATE(8, NPB, N, K) \ - INSTANTIATE(9, NPB, N, K) \ - INSTANTIATE(10, NPB, N, K) \ - INSTANTIATE(11, NPB, N, K) \ - INSTANTIATE(12, NPB, N, K) \ - INSTANTIATE(13, NPB, N, K) \ - INSTANTIATE(14, NPB, N, K) \ - INSTANTIATE(15, NPB, N, K) \ - INSTANTIATE(16, NPB, N, K) \ - INSTANTIATE(17, NPB, N, K) \ - INSTANTIATE(18, NPB, N, K) \ - INSTANTIATE(19, NPB, N, K) \ - INSTANTIATE(20, NPB, N, K) \ - INSTANTIATE(21, NPB, N, K) \ - INSTANTIATE(22, NPB, N, K) \ - INSTANTIATE(23, NPB, N, K) \ - INSTANTIATE(24, NPB, N, K) \ - INSTANTIATE(25, NPB, N, K) \ - INSTANTIATE(26, NPB, N, K) \ - INSTANTIATE(27, NPB, N, K) \ - INSTANTIATE(28, NPB, N, K) \ - INSTANTIATE(29, NPB, N, K) \ - INSTANTIATE(30, NPB, N, K) \ - INSTANTIATE(31, NPB, N, K) \ - INSTANTIATE(32, NPB, N, K) - -INSTANTIATE_ALL_M(4, 768, 12288) -INSTANTIATE_ALL_M(4, 1536, 12288) -INSTANTIATE_ALL_M(2, 6144, 12288) -// LL-mode (M<=8 wiring guard) backbone shapes, B300 sweep vs cuBLAS: -// q_b_proj (2048, 2048): 1.67x/1.29x/1.15x at M=4/6/8 (NPB=4 within -// 0.1us of per-M best) -// shared-expert gate_up (512, 6144): 1.95x/1.58x/1.40x at M=4/6/8 -// cuBLAS keeps qkv_a (2624,6144) and o_proj (6144,2048) — already at -// 3.4-3.8 TB/s there; the GEMV loses on activation re-reads. -INSTANTIATE_ALL_M(4, 2048, 2048) -INSTANTIATE_ALL_M(4, 512, 6144) -// fused_qkv_a (2624, 6144), 32MB: skinny wins ONLY at M<=2 (B300: M=1 -// 6.99us vs cuBLAS 9.21 = 1.32x, M=2 1.24x; M>=4 cuBLAS holds at 3.5TB/s -// and every alternative loses — cublasLt top-8 3.6TB/s wall, DeepGEMM -// 0.71x, wmma+cp.async custom 0.30x pending a TMA rewrite). -INSTANTIATE_ALL_M(4, 2624, 6144) -// DSv3.2 (TP8) siblings of the GLM shapes above, same dual-chip matrix: -// fused_qkv_a (2112, 7168), 30MB: wins M<=2 (M=1 1.30-1.34x) -// MTP eh_proj (7168, 14336), 205MB: wins M<=2 (M=1 1.12-1.16x) -INSTANTIATE_ALL_M(4, 2112, 7168) -INSTANTIATE_ALL_M(2, 7168, 14336) - -#undef INSTANTIATE_ALL_M -#undef INSTANTIATE diff --git a/csrc/libtorch_stable/bf16_skinny_gemm_entry.cu b/csrc/libtorch_stable/bf16_skinny_gemm_entry.cu deleted file mode 100644 index 60b9e8b60cd..00000000000 --- a/csrc/libtorch_stable/bf16_skinny_gemm_entry.cu +++ /dev/null @@ -1,170 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -#include -#include -#include - -#include "core/registration.h" -#include "libtorch_stable/torch_utils.h" - -#include -#include - -#include - -namespace { - -inline int getSMVersion() { - auto* props = get_device_prop(); - return props->major * 10 + props->minor; -} - -} // namespace - -static constexpr int SKINNY_MAX_TOKENS = 32; - -// Supported (N, K) pairs (must match the instantiations in -// bf16_skinny_gemm.cu): eh_proj shard TP8 / TP4 / unsharded. -static inline bool bf16_skinny_gemm_supported(int n, int k) { - if (k == 12288 && (n == 768 || n == 1536 || n == 6144)) return true; - // LL-mode backbone shapes (wire callers with an M <= 8 guard; the GEMV - // family loses to cuBLAS at larger M). - if (k == 2048 && n == 2048) return true; // q_b_proj (TP8) - if (k == 6144 && n == 2624) return true; // fused_qkv_a (wire M <= 2 only) - if (k == 7168 && n == 2112) return true; // DSv3.2 fused_qkv_a (M <= 2) - if (k == 14336 && n == 7168) return true; // DSv3.2 eh_proj (M <= 2) - if (k == 6144 && n == 512) return true; // shared-expert gate_up (TP8) - return false; -} - -// Forward declarations - template params must match bf16_skinny_gemm.cu -template -void invokeBf16SkinnyGemm(__nv_bfloat16* output, __nv_bfloat16 const* mat_a, - __nv_bfloat16 const* mat_b, int64_t out_stride, - cudaStream_t stream); - -template -struct SkinnyLoopUnroller { - static void unroll(int num_tokens, __nv_bfloat16* output, - __nv_bfloat16 const* mat_a, __nv_bfloat16 const* mat_b, - int64_t out_stride, cudaStream_t stream) { - if (num_tokens == kBegin) { - invokeBf16SkinnyGemm<128, kNPB, kBegin, kN, kK>(output, mat_a, mat_b, - out_stride, stream); - } else { - SkinnyLoopUnroller::unroll( - num_tokens, output, mat_a, mat_b, out_stride, stream); - } - } -}; - -template -struct SkinnyLoopUnroller { - static void unroll(int num_tokens, __nv_bfloat16* output, - __nv_bfloat16 const* mat_a, __nv_bfloat16 const* mat_b, - int64_t out_stride, cudaStream_t stream) { - if (num_tokens == kEnd) { - invokeBf16SkinnyGemm<128, kNPB, kEnd, kN, kK>(output, mat_a, mat_b, - out_stride, stream); - } else { - throw std::invalid_argument( - "bf16_skinny_gemm: num_tokens must be in [1, 32]"); - } - } -}; - -static void dispatchBf16SkinnyGemm(int n, int k, int num_tokens, - __nv_bfloat16* output, - __nv_bfloat16 const* mat_a, - __nv_bfloat16 const* mat_b, - int64_t out_stride, cudaStream_t stream) { - if (n == 768 && k == 12288) { - SkinnyLoopUnroller<4, 768, 12288, 1, SKINNY_MAX_TOKENS>::unroll( - num_tokens, output, mat_a, mat_b, out_stride, stream); - } else if (n == 1536 && k == 12288) { - SkinnyLoopUnroller<4, 1536, 12288, 1, SKINNY_MAX_TOKENS>::unroll( - num_tokens, output, mat_a, mat_b, out_stride, stream); - } else if (n == 6144 && k == 12288) { - SkinnyLoopUnroller<2, 6144, 12288, 1, SKINNY_MAX_TOKENS>::unroll( - num_tokens, output, mat_a, mat_b, out_stride, stream); - } else if (n == 2048 && k == 2048) { - SkinnyLoopUnroller<4, 2048, 2048, 1, SKINNY_MAX_TOKENS>::unroll( - num_tokens, output, mat_a, mat_b, out_stride, stream); - } else if (n == 2624 && k == 6144) { - SkinnyLoopUnroller<4, 2624, 6144, 1, SKINNY_MAX_TOKENS>::unroll( - num_tokens, output, mat_a, mat_b, out_stride, stream); - } else if (n == 2112 && k == 7168) { - SkinnyLoopUnroller<4, 2112, 7168, 1, SKINNY_MAX_TOKENS>::unroll( - num_tokens, output, mat_a, mat_b, out_stride, stream); - } else if (n == 7168 && k == 14336) { - SkinnyLoopUnroller<2, 7168, 14336, 1, SKINNY_MAX_TOKENS>::unroll( - num_tokens, output, mat_a, mat_b, out_stride, stream); - } else if (n == 512 && k == 6144) { - SkinnyLoopUnroller<4, 512, 6144, 1, SKINNY_MAX_TOKENS>::unroll( - num_tokens, output, mat_a, mat_b, out_stride, stream); - } else { - throw std::invalid_argument("bf16_skinny_gemm: unsupported (N, K) pair"); - } -} - -void bf16_skinny_gemm( - torch::stable::Tensor& output, // [num_tokens, N] bf16 - torch::stable::Tensor const& mat_a, // [num_tokens, K] bf16 - torch::stable::Tensor const& mat_b // [N, K] bf16 -) { - STD_TORCH_CHECK(output.dim() == 2 && mat_a.dim() == 2 && mat_b.dim() == 2); - STD_TORCH_CHECK(output.is_cuda() && mat_a.is_cuda() && mat_b.is_cuda(), - "bf16_skinny_gemm: all tensors must be CUDA tensors"); - STD_TORCH_CHECK(output.get_device_index() == mat_a.get_device_index() && - output.get_device_index() == mat_b.get_device_index(), - "bf16_skinny_gemm: all tensors must be on the same device"); - STD_TORCH_CHECK(mat_a.is_contiguous() && mat_b.is_contiguous(), - "bf16_skinny_gemm: inputs must be contiguous"); - // Output may be a column-slice view of a wider padded buffer: unit column - // stride, row stride >= N (rows must not overlap). - STD_TORCH_CHECK(output.stride(1) == 1, - "bf16_skinny_gemm: output columns must be contiguous"); - - const int num_tokens = mat_a.size(0); - const int n = mat_b.size(0); - const int k = mat_a.size(1); - - STD_TORCH_CHECK(output.size(0) == num_tokens && output.size(1) == n, - "bf16_skinny_gemm: output must be [num_tokens, N]"); - STD_TORCH_CHECK(mat_b.size(1) == k, - "bf16_skinny_gemm: mat_a and mat_b must share K"); - STD_TORCH_CHECK(bf16_skinny_gemm_supported(n, k), - "bf16_skinny_gemm: unsupported (N, K) pair"); - const int64_t out_stride = output.stride(0); - STD_TORCH_CHECK(num_tokens <= 1 || out_stride >= n, - "bf16_skinny_gemm: output rows overlap"); - STD_TORCH_CHECK(num_tokens >= 0 && num_tokens <= SKINNY_MAX_TOKENS, - "bf16_skinny_gemm: num_tokens must be in [0, 32]"); - STD_TORCH_CHECK( - mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16 && - mat_b.scalar_type() == torch::headeronly::ScalarType::BFloat16 && - output.scalar_type() == torch::headeronly::ScalarType::BFloat16, - "bf16_skinny_gemm: all tensors must be bfloat16"); - - // Empty batch (e.g. an empty rank at a DP/PP boundary): nothing to compute. - if (num_tokens == 0) { - return; - } - - const torch::stable::accelerator::DeviceGuard device_guard( - mat_a.get_device_index()); - STD_TORCH_CHECK(getSMVersion() >= 90, "bf16_skinny_gemm: requires SM90+"); - - auto stream = get_current_cuda_stream(mat_a.get_device_index()); - dispatchBf16SkinnyGemm( - n, k, num_tokens, - reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), out_stride, - stream); -} - -STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { - m.impl("bf16_skinny_gemm", TORCH_BOX(&bf16_skinny_gemm)); -} diff --git a/csrc/libtorch_stable/dsv3_fused_a_gemm.cu b/csrc/libtorch_stable/dsv3_fused_a_gemm.cu index e1ba2a8d0e1..585004c047b 100644 --- a/csrc/libtorch_stable/dsv3_fused_a_gemm.cu +++ b/csrc/libtorch_stable/dsv3_fused_a_gemm.cu @@ -391,7 +391,7 @@ struct MmaComputer { static constexpr int n_iter_cnt = (tile_n + 7) / 8; // Possible to have non-1 n_iter_cnt for ab_swap m16 case. - static_assert(m_iter_cnt == 1 || m_iter_cnt == 2); + static_assert(m_iter_cnt == 1); static_assert(n_iter_cnt == 1 || n_iter_cnt == 2); __device__ MmaComputer(bf16_t* gmem_c_local_, bf16_t* smem_a_, @@ -416,18 +416,13 @@ struct MmaComputer { public: __device__ void prepare() { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 - // Fragment addressing is per 16-row ldmatrix tile; m_iter selects the - // 16-row half within tile_m. #pragma unroll - for (int m = 0; m < m_iter_cnt; m++) { - #pragma unroll - for (int i = 0; i < k_phase_cnt; i++) { - int linear_idx = (lane_idx % 16) + (lane_idx / 16) * 128 + i * 256; - int m_idx = linear_idx % 16 + m * 16; - int k_idx = linear_idx / 16 + warp_k_offset_in_tile_k; - k_idx = apply_swizzle_343_on_elem_row_col(m_idx, k_idx); - a_smem_offsets[m][i] = m_idx * tile_k + k_idx; - } + for (int i = 0; i < k_phase_cnt; i++) { + int linear_idx = (lane_idx % 16) + (lane_idx / 16) * 128 + i * 256; + int m_idx = linear_idx % tile_m; + int k_idx = linear_idx / tile_m + warp_k_offset_in_tile_k; + k_idx = apply_swizzle_343_on_elem_row_col(m_idx, k_idx); + a_smem_offsets[0][i] = m_idx * tile_k + k_idx; } #pragma unroll for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) { @@ -451,14 +446,11 @@ struct MmaComputer { wait_barrier(smem_barrier + 0 + stage_idx * 2, phase_bit); #pragma unroll - for (int m = 0; m < m_iter_cnt; m++) { - #pragma unroll - for (int i = 0; i < k_phase_cnt; i++) { - int smem_offset = a_smem_offsets[m][i]; - bf16_t* smem_ptr_this_iter = - smem_a + stage_idx * tile_m * tile_k + smem_offset; - ldsm_x4(smem_ptr_this_iter, reinterpret_cast(a_reg[m][i])); - } + for (int i = 0; i < k_phase_cnt; i++) { + int smem_offset = a_smem_offsets[0][i]; + bf16_t* smem_ptr_this_iter = + smem_a + stage_idx * tile_m * tile_k + smem_offset; + ldsm_x4(smem_ptr_this_iter, reinterpret_cast(a_reg[0][i])); } #pragma unroll @@ -477,12 +469,9 @@ struct MmaComputer { for (int k_iter_idx = 0; k_iter_idx < k_phase_cnt; k_iter_idx++) { #pragma unroll for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) { - #pragma unroll - for (int m = 0; m < m_iter_cnt; m++) { - hmma_16_8_16_f32acc_bf16ab( - acc_reg[m][n_iter_idx], a_reg[m][k_iter_idx], - b_reg[n_iter_idx][k_iter_idx], acc_reg[m][n_iter_idx]); - } + hmma_16_8_16_f32acc_bf16ab( + acc_reg[0][n_iter_idx], a_reg[0][k_iter_idx], + b_reg[n_iter_idx][k_iter_idx], acc_reg[0][n_iter_idx]); } } ::arrive_barrier(smem_barrier + 1 + stage_idx * 2); @@ -497,14 +486,14 @@ struct MmaComputer { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt)); // reorganize the acc_reg - constexpr int thread_m = 2 * m_iter_cnt; + constexpr int thread_m = 2; constexpr int thread_n = 2 * n_iter_cnt; constexpr int cta_mma_n = n_iter_cnt * 8; float acc_reg_reorg[thread_m][thread_n]; for (int i = 0; i < thread_m; i++) { for (int j = 0; j < thread_n; j++) { - acc_reg_reorg[i][j] = acc_reg[i / 2][j / 2][(j % 2) + (i % 2) * 2]; + acc_reg_reorg[i][j] = acc_reg[0][j / 2][(j % 2) + (i * 2)]; } } @@ -525,8 +514,7 @@ struct MmaComputer { for (int m_idx_thread = 0; m_idx_thread < thread_m; m_idx_thread++) { #pragma unroll for (int n_idx_thread = 0; n_idx_thread < thread_n; n_idx_thread++) { - int m_idx = - (lane_idx / 4) + (m_idx_thread % 2) * 8 + (m_idx_thread / 2) * 16; + int m_idx = (lane_idx / 4) + m_idx_thread * 8; int n_idx = ((lane_idx % 4) * 2) + (n_idx_thread % 2) + (n_idx_thread / 2) * 8; smem_c[cosize_smem_c * warp_idx + smem_c_index_func(m_idx, n_idx)] = @@ -599,7 +587,7 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel( static_assert( tile_k == 128 || tile_k == 256 || tile_k == 512 || tile_k == 1024); // tile_k must be larger than 64 since 4 warp splitK. - static_assert(tile_m == 16 || tile_m == 32); + static_assert(tile_m == 16); constexpr int g2s_vec_bytes = 16; constexpr int a_elem_bytes = 2; constexpr int b_elem_bytes = 2; @@ -659,7 +647,7 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel( #endif } -template +template void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens, cudaStream_t const stream) { constexpr int gemm_m = kHdOut; // 2112 @@ -667,7 +655,7 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens, constexpr int gemm_k = kHdIn; // 7168 constexpr int batch_size = 1; std::swap(mat_a, mat_b); - constexpr int tile_m = kTileM; + constexpr int tile_m = 16; constexpr int tile_n = kTileN; // 8 or 16 constexpr int tile_k = std::max(256, 1024 / tile_n); // 256 constexpr int max_stage_cnt = @@ -714,44 +702,6 @@ template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>( __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens, cudaStream_t); -// GLM-5.2 fused_qkv_a (K=6144 -> N=2624). tile_m=32 so the grid is -// 2624/32 = 82 CTAs, a single wave on B300 (148 SMs); tile_m=16's 164 CTAs -// straddle two waves and drop to 2.9 TB/s vs 3.9-4.0 here. Beats cuBLAS -// at every decode M: 1.11-1.15x for M<=8 (tile_n=8), 1.12x at M=16 -// (tile_n=16). The M<=2 dispatch still belongs to bf16_skinny_gemm -// (4.5 TB/s). -template void invokeFusedAGemm<__nv_bfloat16, 6144, 2624, 8, 32>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens, - cudaStream_t); - -template void invokeFusedAGemm<__nv_bfloat16, 6144, 2624, 16, 32>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens, - cudaStream_t); - -// GLM-5.2 q_b_proj TP8 (K=2048 -> N=2048). tile_m=16 keeps 128 CTAs (already -// a single wave) and measures ahead of tile_m=32 here. Beats cuBLAS -// 1.79-1.87x at M=3..8 (tile_n=8) and 1.19-1.30x at M=9..16 (tile_n=16); -// M<=2 belongs to bf16_skinny_gemm, M>=20 to cuBLAS. -template void invokeFusedAGemm<__nv_bfloat16, 2048, 2048, 8>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens, - cudaStream_t); - -template void invokeFusedAGemm<__nv_bfloat16, 2048, 2048, 16>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens, - cudaStream_t); - -// DSv3.2 q_b_proj TP8 (K=1536 -> N=3072). tile_m=32 keeps 96 CTAs (single -// wave; tile_m=16's 192 straddle two). Beats cuBLAS 1.6-1.8x at M=1..8 and -// 1.05-1.18x at M=9..16 on B300/B200; the whole 1..16 range dispatches here -// (no skinny tier: K=1536 does not fit the GEMV's 128x8 K-step). -template void invokeFusedAGemm<__nv_bfloat16, 1536, 3072, 8, 32>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens, - cudaStream_t); - -template void invokeFusedAGemm<__nv_bfloat16, 1536, 3072, 16, 32>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens, - cudaStream_t); - void dsv3_fused_a_gemm(torch::stable::Tensor& output, torch::stable::Tensor const& mat_a, torch::stable::Tensor const& mat_b) { @@ -760,15 +710,12 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output, int const hd_in = mat_a.size(1); int const hd_out = mat_b.size(1); - bool const is_dsv3 = hd_in == 7168 && hd_out == 2112; - bool const is_glm = hd_in == 6144 && hd_out == 2624; - bool const is_glm_qb = hd_in == 2048 && hd_out == 2048; - bool const is_ds_qb = hd_in == 1536 && hd_out == 3072; + constexpr int kHdIn = 7168; + constexpr int kHdOut = 2112; STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16, "required 1 <= mat_a.shape[0] <= 16"); - STD_TORCH_CHECK(is_dsv3 || is_glm || is_glm_qb || is_ds_qb, - "supported (hd_in, hd_out): (7168, 2112), (6144, 2624), " - "(2048, 2048), (1536, 3072)"); + STD_TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168"); + STD_TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112"); STD_TORCH_CHECK(output.size(0) == num_tokens, "required output.shape[0] == mat_a.shape[0]"); STD_TORCH_CHECK(output.size(1) == hd_out, @@ -791,41 +738,18 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output, STD_TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90"); auto stream = get_current_cuda_stream(mat_a.get_device_index()); - auto* out = reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()); - auto* a = reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()); - auto* b = reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()); - if (is_dsv3) { - if (num_tokens <= 8) { - invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 8>(out, a, b, num_tokens, - stream); - } else { - invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>(out, a, b, num_tokens, - stream); - } - } else if (is_glm) { - if (num_tokens <= 8) { - invokeFusedAGemm<__nv_bfloat16, 6144, 2624, 8, 32>(out, a, b, num_tokens, - stream); - } else { - invokeFusedAGemm<__nv_bfloat16, 6144, 2624, 16, 32>(out, a, b, num_tokens, - stream); - } - } else if (is_glm_qb) { - if (num_tokens <= 8) { - invokeFusedAGemm<__nv_bfloat16, 2048, 2048, 8>(out, a, b, num_tokens, - stream); - } else { - invokeFusedAGemm<__nv_bfloat16, 2048, 2048, 16>(out, a, b, num_tokens, - stream); - } + if (num_tokens <= 8) { + invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 8>( + reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()), + reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), + reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens, + stream); } else { - if (num_tokens <= 8) { - invokeFusedAGemm<__nv_bfloat16, 1536, 3072, 8, 32>(out, a, b, num_tokens, - stream); - } else { - invokeFusedAGemm<__nv_bfloat16, 1536, 3072, 16, 32>(out, a, b, num_tokens, - stream); - } + invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 16>( + reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()), + reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), + reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens, + stream); } } diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu b/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu index f10c58eab75..f7c965dbc1b 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu +++ b/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu @@ -49,11 +49,6 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, "Vec size is not matched."); -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaGridDependencySynchronize(); - cudaTriggerProgrammaticLaunchCompletion(); -#endif - // Precompute SF layout parameter (constant for entire kernel). int32_t const numKTiles = (outputCols + 63) / 64; @@ -128,11 +123,6 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, "Vec size is not matched."); -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaGridDependencySynchronize(); - cudaTriggerProgrammaticLaunchCompletion(); -#endif - int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x; int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD; @@ -211,8 +201,6 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output, const torch::stable::accelerator::DeviceGuard device_guard( input.get_device_index()); auto stream = get_current_cuda_stream(input.get_device_index()); - auto* device_props = get_device_prop(); - int const sm_version = device_props->major * 10 + device_props->minor; int output_sf_n_unpadded = int(output_n / CVT_FP4_SF_VEC_SIZE); @@ -236,21 +224,10 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output, input.scalar_type(), "nvfp4_quant_kernel", [&] { using cuda_type = vllm::CUDATypeConverter::Type; auto input_ptr = static_cast(input.data_ptr()); - cudaLaunchConfig_t config = {}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 0; - config.stream = stream; - cudaLaunchAttribute attrs[1]; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = 1; - config.numAttrs = (sm_version >= 90) ? 1 : 0; - config.attrs = attrs; - cudaLaunchKernelEx(&config, vllm::cvt_fp16_to_fp4, - m, n, output_n, num_padded_cols, input_ptr, - input_sf_ptr, - reinterpret_cast(output_ptr), - reinterpret_cast(sf_out)); + vllm::cvt_fp16_to_fp4<<>>( + m, n, output_n, num_padded_cols, input_ptr, input_sf_ptr, + reinterpret_cast(output_ptr), + reinterpret_cast(sf_out)); }); } else { int num_packed_cols = output_n / CVT_FP4_ELTS_PER_THREAD; @@ -263,21 +240,12 @@ void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output, input.scalar_type(), "nvfp4_quant_kernel", [&] { using cuda_type = vllm::CUDATypeConverter::Type; auto input_ptr = static_cast(input.data_ptr()); - cudaLaunchConfig_t config = {}; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = 0; - config.stream = stream; - cudaLaunchAttribute attrs[1]; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = 1; - config.numAttrs = (sm_version >= 90) ? 1 : 0; - config.attrs = attrs; - cudaLaunchKernelEx( - &config, vllm::cvt_fp16_to_fp4_sf_major, m, n, - output_n, output_sf_n_unpadded, num_packed_cols, input_ptr, - input_sf_ptr, reinterpret_cast(output_ptr), - reinterpret_cast(sf_out)); + vllm::cvt_fp16_to_fp4_sf_major + <<>>( + m, n, output_n, output_sf_n_unpadded, num_packed_cols, + input_ptr, input_sf_ptr, + reinterpret_cast(output_ptr), + reinterpret_cast(sf_out)); }); } } diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 01441180623..0a475d02c6f 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -330,10 +330,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { // conditionally compiled so impl registration is in source file ops.def("fp32_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); - // BF16 skinny GEMM (M<=32, weight-BW-bound shapes, e.g. MTP eh_proj). - // conditionally compiled so impl registration is in source file - ops.def("bf16_skinny_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); - // reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel ops.def( "rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, " diff --git a/recipes/glm5.2-ll-b300-tp8-mtp5.md b/recipes/glm5.2-ll-b300-tp8-mtp5.md deleted file mode 100644 index a0deca9bd98..00000000000 --- a/recipes/glm5.2-ll-b300-tp8-mtp5.md +++ /dev/null @@ -1,41 +0,0 @@ -# GLM-5.2-NVFP4 — Low-Latency (LL) serving recipe · 8×B300 · TP8 · MTP=5 - -Single-node 8×B300 (sm_100/103), pure tensor parallel, NVFP4 weights, fp8 KV cache, -MTP=5 speculative decode, model runner v2. - -Validated on `glm5.2-LL` @ `8d407aee1`. - -## Versions - -| Component | Version | -| --- | --- | -| vLLM | `glm5.2-LL` @ `8d407aee1` | -| flashinfer-python | `0.6.15` | -| Model | `nvidia/GLM-5.2-NVFP4` (`modelopt_fp4`) | - -## Build (once) - -```bash -pip install 'flashinfer-python==0.6.15' - -CUDA_HOME=/usr/local/cuda-13.0 TORCH_CUDA_ARCH_LIST="10.0 10.3" MAX_JOBS=96 \ -VLLM_USE_PRECOMPILED=0 VLLM_USE_PRECOMPILED_RUST=0 \ -pip install -e . --no-deps --no-build-isolation -``` - -## Server - -```bash -export VLLM_USE_V2_MODEL_RUNNER=1 # model runner v2 -export VLLM_DEEP_GEMM_WARMUP=skip - -vllm serve nvidia/GLM-5.2-NVFP4 \ - --trust-remote-code \ - --tensor-parallel-size 8 \ - --quantization modelopt_fp4 --kv-cache-dtype fp8_e4m3 \ - --max-model-len 32768 --max-num-batched-tokens 16384 --max-num-seqs 256 \ - --no-enable-prefix-caching --gpu-memory-utilization 0.85 \ - --speculative-config '{"method":"mtp","num_speculative_tokens":5}' \ - --kernel-config '{"ir_op_priority":{"rms_norm":["vllm_c","native"],"fused_add_rms_norm":["vllm_c","native"]}}' \ - --host 0.0.0.0 --port 8000 -``` diff --git a/tests/kernels/test_bf16_skinny_gemm.py b/tests/kernels/test_bf16_skinny_gemm.py deleted file mode 100644 index f49e4b4db04..00000000000 --- a/tests/kernels/test_bf16_skinny_gemm.py +++ /dev/null @@ -1,70 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Numerical correctness for the bf16 skinny GEMM (decode-M GEMV) kernel. - -Covers every supported (N, K) shape across M = 0/1/2/8/16/32 (M=32 exercises -the tile_m=32 MMA path), plus strided (column-slice) output. M=0 must return -an empty result without launching the kernel (empty ranks at DP/PP -boundaries).""" - -import pytest -import torch - -from vllm import _custom_ops as ops -from vllm.platforms import current_platform - -# (N, K) pairs must match bf16_skinny_gemm_supported() in the entry .cu. -SHAPES = [ - (768, 12288), - (1536, 12288), - (6144, 12288), - (2048, 2048), - (2624, 6144), - (2112, 7168), - (7168, 14336), - (512, 6144), -] -MS = [0, 1, 2, 8, 16, 32] - -pytestmark = pytest.mark.skipif( - not current_platform.is_cuda() or not current_platform.has_device_capability(90), - reason="bf16_skinny_gemm requires CUDA SM90+", -) - - -def _rel_err(out: torch.Tensor, ref: torch.Tensor) -> float: - return (out.float() - ref).norm().item() / ref.norm().clamp_min(1e-6).item() - - -@pytest.mark.parametrize("n,k", SHAPES) -@pytest.mark.parametrize("m", MS) -def test_bf16_skinny_gemm_matches_reference(n: int, k: int, m: int): - torch.manual_seed(0) - x = torch.randn(m, k, dtype=torch.bfloat16, device="cuda") - w = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") - - out = ops.bf16_skinny_gemm(x, w) - assert out.shape == (m, n) - assert out.dtype == torch.bfloat16 - - if m == 0: # empty batch: nothing to check beyond the empty shape - return - ref = x.float() @ w.float().t() - rel = _rel_err(out, ref) - assert rel < 2e-2, f"rel err {rel:.4f} for (m,n,k)=({m},{n},{k})" - - -@pytest.mark.parametrize("n,k", [(2048, 2048), (512, 6144)]) -def test_bf16_skinny_gemm_strided_output(n: int, k: int): - """Output is a column-slice of a wider padded buffer (row stride > N).""" - torch.manual_seed(0) - m = 8 - x = torch.randn(m, k, dtype=torch.bfloat16, device="cuda") - w = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") - - wide = torch.empty(m, n + 64, dtype=torch.bfloat16, device="cuda") - view = wide[:, :n] - torch.ops._C.bf16_skinny_gemm(view, x, w) - - ref = x.float() @ w.float().t() - assert _rel_err(view, ref) < 2e-2 diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py index f2302a896b3..a27e67b6245 100644 --- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -49,32 +49,6 @@ pytestmark = pytest.mark.skipif( ) -def test_platform_capability_queries_are_constant_during_compile(monkeypatch): - monkeypatch.setattr( - K.current_platform, "has_device_capability", lambda capability: True - ) - monkeypatch.setattr(K, "has_cutedsl", lambda: True) - monkeypatch.setattr(K.current_platform, "is_arch_support_pdl", lambda: True) - - def capability_branches(x: torch.Tensor) -> torch.Tensor: - if K._can_use_fused_q_cutedsl() and K._is_arch_support_pdl(): - return x + 1 - return x - 1 - - compiled = torch.compile(capability_branches, backend="eager", fullgraph=True) - x = torch.zeros(1, device="cuda") - torch.testing.assert_close(compiled(x), torch.ones_like(x)) - - -def test_pdl_is_disabled_before_blackwell(monkeypatch): - monkeypatch.setattr( - K.current_platform, "has_device_capability", lambda capability: False - ) - monkeypatch.setattr(K.current_platform, "is_arch_support_pdl", lambda: True) - - assert not K._is_arch_support_pdl() - - # ── reference helpers ──────────────────────────────────────────────────────── @@ -318,57 +292,6 @@ def test_fused_norm_rope_no_indexer(num_tokens: int): assert (topk == 7).all(), "topk buffer should be untouched on shared layer" -def test_fused_norm_rope_profile_without_cache_compiles(): - """The cache-free profiling path must still compile and produce Q.""" - torch.manual_seed(2) - dev = "cuda" - num_tokens = 4 - pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) - q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16) - kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16) - k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16) - qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16) - kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16) - index_k = torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16) - index_w = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32) - index_b = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32) - cos_sin = make_cos_sin(64, ROPE_DIM, dev) - topk = torch.empty(num_tokens, 2048, device=dev, dtype=torch.int32) - - def profile_run( - q_c: torch.Tensor, - kv_c: torch.Tensor, - k_pe: torch.Tensor, - index_k: torch.Tensor, - ) -> torch.Tensor: - return K.fused_norm_rope( - pos, - q_c, - qw, - EPS, - kv_c, - kvw, - EPS, - k_pe, - cos_sin, - index_k, - index_w, - index_b, - EPS, - cos_sin, - topk, - slot_mapping=None, - indexer_k_cache=None, - mla_kv_cache=None, - has_indexer=True, - index_rope_interleave=False, - ) - - compiled = torch.compile(profile_run, fullgraph=True) - q_out = compiled(q_c, kv_c, k_pe, index_k) - assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm (profiling)") - - @pytest.mark.parametrize("num_tokens", [1, 4, 17, 512]) def test_fused_norm_rope_ds_mla(num_tokens: int): """fp8_ds_mla MLA cache layout (FlashMLA sparse, bf16-query path; SM90/SM100). diff --git a/tests/models/registry.py b/tests/models/registry.py index 65cdc84dec8..afbac216c93 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -1600,10 +1600,6 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = { speculative_model="luccafong/deepseek_mtp_draft_random", trust_remote_code=True, ), - "DeepseekV32MTPModel": _HfExamplesInfo( - "deepseek-ai/DeepSeek-V3.2-Exp", - speculative_model="deepseek-ai/DeepSeek-V3.2-Exp", - ), "DeepSeekV4MTPModel": _HfExamplesInfo( "deepseek-ai/DeepSeek-V4-Flash", speculative_model="deepseek-ai/DeepSeek-V4-Flash", diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 018fc1ac2b6..aa75c50a516 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2371,20 +2371,6 @@ def fp32_router_gemm( return output -def bf16_skinny_gemm( - x: torch.Tensor, - weight: torch.Tensor, -) -> torch.Tensor: - """Skinny bf16 GEMM (M<=32): x [M, K] @ weight [N, K]^T -> [M, N]. - - Single block-per-column kernel with fp32 accumulation; replaces cuBLAS - splitK for weight-bandwidth-bound decode shapes (e.g. MTP eh_proj). - """ - output = torch.empty((x.shape[0], weight.shape[0]), dtype=x.dtype, device=x.device) - torch.ops._C.bf16_skinny_gemm(output, x, weight) - return output - - if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "fp32_router_gemm"): @register_fake("_C::fp32_router_gemm") @@ -2396,17 +2382,6 @@ if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "fp32_router_gemm"): return -if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "bf16_skinny_gemm"): - - @register_fake("_C::bf16_skinny_gemm") - def bf16_skinny_gemm_fake( - output: torch.Tensor, - mat_a: torch.Tensor, - mat_b: torch.Tensor, - ) -> None: - return - - def topk_softmax( topk_weights: torch.Tensor, topk_ids: torch.Tensor, diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index 64ff3811787..4716e7012f3 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -202,7 +202,11 @@ if flashinfer_comm is not None: element_size = allreduce_in.element_size() current_tensor_size = num_tokens * hidden_size * element_size max_tensor_size = max_token_num * hidden_size * element_size - oversize = current_tensor_size > max_tensor_size + assert current_tensor_size <= max_tensor_size, ( + f"Current tensor size {current_tensor_size} is larger than " + f"max token num {max_token_num} * hidden size {hidden_size} * " + f"element size {element_size}" + ) curr_device = current_platform.get_device_capability() device_capability = curr_device.to_int() if curr_device is not None else None @@ -215,47 +219,6 @@ if flashinfer_comm is not None: get_workspace_fn = ( get_fi_ar_quant_workspace if is_quant_pattern else get_fi_ar_workspace ) - - orig_norm_out = norm_out - - def unfused_fallback() -> None: - # Fused-path equivalent honoring the op's output-buffer contract: - # with norm_out given, normed -> norm_out and the new residual -> - # the allreduce_in buffer; otherwise normed -> allreduce_in and - # the new residual -> the residual buffer. Runs fine during CUDA - # graph capture (the branch is resolved per captured shape). - # - # weight_bias mirrors the fused kernel's gamma offset (Gemma - # patterns pass 1.0 for rms_norm(x, weight + 1)); fold it into - # the gamma in fp32 like the reference before the plain rms_norm. - gamma = ( - rms_gamma - if weight_bias == 0.0 - else (rms_gamma.float() + weight_bias).to(rms_gamma.dtype) - ) - reduced = tensor_model_parallel_all_reduce(allreduce_in) - new_residual = reduced + residual - if orig_norm_out is None: - torch.ops._C.rms_norm(allreduce_in, new_residual, gamma, rms_eps) - residual.copy_(new_residual) - else: - torch.ops._C.rms_norm(orig_norm_out, new_residual, gamma, rms_eps) - allreduce_in.copy_(new_residual) - - if oversize: - assert not is_quant_pattern, ( - "flashinfer fused allreduce: input exceeds the workspace budget " - "and the unfused fallback does not cover quant patterns" - ) - logger.warning_once( - "fused allreduce+rmsnorm: %d tokens exceeds the workspace " - "budget (max_token_num=%d); using the unfused fallback for " - "this shape.", - num_tokens, - max_token_num, - ) - unfused_fallback() - return workspace = get_workspace_fn( world_size=world_size, rank=get_tensor_model_parallel_rank(), @@ -289,75 +252,36 @@ if flashinfer_comm is not None: if workspace.backend in ("trtllm", "mnnvl"): layout_code = flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4 - # vllm's one-shot size table is tuned for the trtllm backend. For - # mnnvl, flashinfer caps the workspace's one-shot region at its own - # MNNVL_ONE_SHOT_THRESHOLD (~4MB, e.g. 44 tokens at hidden 6144 bf16 - # TP8); forcing use_oneshot=True past that raises ValueError instead - # of using two-shot, so every decode capture between the one-shot cap - # and the fusion threshold silently lost fusion. Let flashinfer's - # AUTO heuristic pick the strategy (it matches its own workspace - # sizing), and keep trigger_completion_at_end=True since AUTO may - # select one-shot (see the Lamport NaN note below). - mnnvl_auto = workspace.backend == "mnnvl" - if mnnvl_auto: - use_oneshot = None - - try: - flashinfer_comm.allreduce_fusion( - input=allreduce_in, - workspace=workspace, - pattern=pattern_code, - launch_with_pdl=launch_with_pdl, - output=None, - residual_out=residual_out, - norm_out=norm_out, - quant_out=quant_out, - scale_out=scale_out, - residual_in=residual, - rms_gamma=rms_gamma, - rms_eps=rms_eps, - scale_factor=scale_factor, - layout_code=layout_code, - use_oneshot=use_oneshot, - fp32_acc=fp32_acc, - weight_bias=weight_bias, - # The one-shot Lamport all-reduce signals PDL completion before - # its output buffer is committed when trigger_completion_at_end - # is False, so the next PDL-launched kernel can read the - # uninitialized Lamport buffer and produce NaN. This only fires - # for num_tokens <= PDL_ADVANCE_LAUNCH_TOKENS (the batch=1 / - # spec-decode shapes, where the one-shot path is always - # selected). Complete at the end for the one-shot path; the - # two-shot path is synchronized and keeps the early completion. - # Related one-shot instability in the same kernel: - # flashinfer-ai/flashinfer#1223. - trigger_completion_at_end=True - if mnnvl_auto - else (use_oneshot or num_tokens > PDL_ADVANCE_LAUNCH_TOKENS), - ) - except ValueError as e: - # The globally cached workspace can be smaller than this call's - # budget (it is created once by whichever caller runs first, and - # the mnnvl workspace's is_buffer_size_sufficient() checks trtllm - # metadata, so vllm-side accounting cannot see the real capacity). - # flashinfer raises before launching anything, so falling back - # mid-capture is safe. Observed with spec-decode draft graphs - # whose token count (max_num_reqs * (1 + num_spec)) exceeds the - # workspace sized for the fusion threshold. - assert not is_quant_pattern, ( - "flashinfer fused allreduce failed and the unfused fallback " - f"does not cover quant patterns: {e}" - ) - logger.warning_once( - "flashinfer fused allreduce+rmsnorm unavailable for this " - "shape (num_tokens=%d, max_token_num=%d, world_size=%d): %s. " - "Using the unfused fallback.", - num_tokens, - max_token_num, - world_size, - e, - ) - unfused_fallback() + flashinfer_comm.allreduce_fusion( + input=allreduce_in, + workspace=workspace, + pattern=pattern_code, + launch_with_pdl=launch_with_pdl, + output=None, + residual_out=residual_out, + norm_out=norm_out, + quant_out=quant_out, + scale_out=scale_out, + residual_in=residual, + rms_gamma=rms_gamma, + rms_eps=rms_eps, + scale_factor=scale_factor, + layout_code=layout_code, + use_oneshot=use_oneshot, + fp32_acc=fp32_acc, + weight_bias=weight_bias, + # The one-shot Lamport all-reduce signals PDL completion before its + # output buffer is committed when trigger_completion_at_end is + # False, so the next PDL-launched kernel can read the uninitialized + # Lamport buffer and produce NaN. This only fires for + # num_tokens <= PDL_ADVANCE_LAUNCH_TOKENS (the batch=1 / spec-decode + # shapes, where the one-shot path is always selected). Complete at + # the end for the one-shot path; the two-shot path is synchronized + # and keeps the early completion. Related one-shot instability in + # the same kernel: flashinfer-ai/flashinfer#1223. + trigger_completion_at_end=(use_oneshot is True) + or num_tokens > PDL_ADVANCE_LAUNCH_TOKENS, + ) def call_trtllm_fused_allreduce_norm_fake( allreduce_in: torch.Tensor, diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index bd536800cc8..8ba55a2ec96 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -330,7 +330,6 @@ class SpeculativeConfig: @staticmethod def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig: initial_architecture = hf_config.architectures[0] - use_sparse_mtp = hf_config.model_type == "glm_moe_dsa" if hf_config.model_type in ( "deepseek_v3", "deepseek_v32", @@ -340,12 +339,7 @@ class SpeculativeConfig: if hf_config.model_type == "deepseek_mtp": n_predict = getattr(hf_config, "num_nextn_predict_layers", None) hf_config.update( - { - "n_predict": n_predict, - "architectures": [ - "DeepseekV32MTPModel" if use_sparse_mtp else "DeepSeekMTPModel" - ], - } + {"n_predict": n_predict, "architectures": ["DeepSeekMTPModel"]} ) if hf_config.model_type == "deepseek_v4": hf_config.model_type = "deepseek_mtp" diff --git a/vllm/cute_utils/__init__.py b/vllm/cute_utils/__init__.py index d5e536db90f..ca445284bf4 100644 --- a/vllm/cute_utils/__init__.py +++ b/vllm/cute_utils/__init__.py @@ -33,11 +33,6 @@ EVICT_NORMAL = Int64(0x1000000000000000) EVICT_FIRST = Int64(0x12F0000000000000) EVICT_LAST = Int64(0x14F0000000000000) -TORCH_TO_CUTE_DTYPE = { - torch.bfloat16: BFloat16, - torch.float32: Float32, -} - @dsl_user_op def recast_val(x, dtype, *, loc=None, ip=None): diff --git a/vllm/cute_utils/cvt.py b/vllm/cute_utils/cvt.py index 8bb4ad95198..4707f3f7421 100644 --- a/vllm/cute_utils/cvt.py +++ b/vllm/cute_utils/cvt.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from cutlass import BFloat16, Constexpr, Float32, Uint16, Uint32, cute +from cutlass import Constexpr, Float32, Uint32, cute from cutlass._mlir import ir from cutlass._mlir.dialects import llvm, vector from cutlass.cutlass_dsl import T, dsl_user_op @@ -38,11 +38,6 @@ def bf16x2_to_fp32x2(data, *, loc=None, ip=None) -> tuple[Float32, Float32]: ) elif isinstance(data, (cute.Tensor, cute.TensorSSA)): - # recast to 32-bit registers if needed - if data.element_type == BFloat16: - data = cute.recast_tensor(data, Uint32) - assert data.element_type == Uint32 - # NOTE: the output is always 1D size = cute.size(data.shape) out = cute.make_rmem_tensor(size * 2, Float32) @@ -89,22 +84,6 @@ def fp8x4_to_bf16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA: return cute.TensorSSA(vec, 2, Uint32) -@dsl_user_op -def fp32x2_to_fp8x2(a0: Float32, a1: Float32, *, loc=None, ip=None) -> Uint16: - out = llvm.inline_asm( - T.i16(), - [ - a0.ir_value(loc=loc, ip=ip), - a1.ir_value(loc=loc, ip=ip), - ], - "cvt.rn.satfinite.e4m3x2.f32 $0, $2, $1;", - "=h,f,f", - has_side_effects=False, - is_align_stack=False, - ) - return Uint16(out) - - @dsl_user_op def fp8x4_to_fp16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA: out = llvm.inline_asm( diff --git a/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py b/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py index 509befbcc14..e49e135b26a 100644 --- a/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py +++ b/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py @@ -43,7 +43,7 @@ try: if flashinfer_comm is not None else None ) -except (ImportError, AttributeError): +except ImportError: flashinfer_trtllm_fused_allreduce_norm = None # type: ignore[assignment] get_fi_ar_workspace = None # type: ignore[assignment] _AR_RESIDUAL_RMS_NORM = None @@ -52,23 +52,12 @@ except (ImportError, AttributeError): _FI_SUPPORTED_DTYPES = (torch.bfloat16, torch.float16) -@torch.compiler.assume_constant_result -def _fi_ar_max_size_mb() -> dict[int, float]: - """Flashinfer all-reduce fusion size table for the current device. - - Device capability is constant; marking the result constant keeps this out - of the traced graph. Otherwise the per-forward call reaches - ``current_platform.get_device_capability()`` inside ``torch.compile`` and - graph-breaks ("can't handle functions not implemented in python").""" - from vllm.config.compilation import PassConfig - - return PassConfig.default_fi_allreduce_fusion_max_size_mb() - - def _max_token_num(tp_size: int, hidden_size: int, dtype: torch.dtype) -> int | None: """Workspace token budget for flashinfer fused all-reduce, or None if the current world size / device is unsupported. Mirrors ``FlashInferAllReduce``.""" - max_size_mb = _fi_ar_max_size_mb().get(tp_size) + from vllm.config.compilation import PassConfig + + max_size_mb = PassConfig.default_fi_allreduce_fusion_max_size_mb().get(tp_size) if not max_size_mb: return None element_size = torch.tensor([], dtype=dtype).element_size() diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 947e0d043dd..cc2688e37fc 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -114,7 +114,7 @@ _TEXT_GENERATION_MODELS = { "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"), "Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"), "Glm4MoeLiteForCausalLM": ("glm4_moe_lite", "Glm4MoeLiteForCausalLM"), - "GlmMoeDsaForCausalLM": ("vllm.models.deepseek_v32", "GlmMoeDsaForCausalLM"), + "GlmMoeDsaForCausalLM": ("deepseek_v2", "GlmMoeDsaForCausalLM"), "GptOssForCausalLM": ("gpt_oss", "GptOssForCausalLM"), "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"), "GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"), @@ -628,7 +628,6 @@ _SPECULATIVE_DECODING_MODELS = { "Eagle3DeepseekV3ForCausalLM": ("deepseek_eagle3", "Eagle3DeepseekV2ForCausalLM"), "EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"), "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"), - "DeepseekV32MTPModel": ("vllm.models.deepseek_v32", "DeepseekV32MTP"), "DeepSeekV4MTPModel": ("vllm.models.deepseek_v4", "DeepSeekV4MTP"), "MiniMaxM3MTP": ("vllm.models.minimax_m3", "MiniMaxM3MTP"), "BailingMoeV25MTPModel": ("bailing_moe_mtp", "BailingMoeV25MTPModel"), diff --git a/vllm/models/deepseek_v32/__init__.py b/vllm/models/deepseek_v32/__init__.py index 6bd447bab2b..1b0aa64262f 100644 --- a/vllm/models/deepseek_v32/__init__.py +++ b/vllm/models/deepseek_v32/__init__.py @@ -6,31 +6,17 @@ DeepSeek V3.2 introduced the DeepSeek Sparse Attention (DSA) architecture: MLA + a "lightning indexer" that selects the top-k tokens for a sparse MLA attend. The same model code serves any DSA checkpoint, including GLM-5.2 (``glm_moe_dsa``), which reuses this architecture. - -The optimized kernels under ``nvidia/`` target the Blackwell (SM100) family. -Every other platform — ROCm, XPU, pre-SM100 CUDA (e.g. H100), CPU — falls back -to the generic ``deepseek_v2`` implementation, which already handles the DSA -(index_topk) architecture and is ``torch.compile``-friendly there. This matches -main's behavior on those platforms (no hard failure). """ from vllm.platforms import current_platform -if current_platform.is_cuda() and current_platform.is_device_capability_family(100): - from .nvidia.model import DeepseekV32ForCausalLM - from .nvidia.mtp import DeepseekV32MTP +if current_platform.is_rocm() or current_platform.is_xpu(): + raise NotImplementedError("deepseek_v32 currently supports NVIDIA SM100 only.") - # GLM-5.2 (glm_moe_dsa) reuses the same optimized DSA module on SM100. - GlmMoeDsaForCausalLM = DeepseekV32ForCausalLM -else: - from vllm.model_executor.models.deepseek_mtp import DeepSeekMTP as DeepseekV32MTP - from vllm.model_executor.models.deepseek_v2 import ( - DeepseekV3ForCausalLM as DeepseekV32ForCausalLM, - ) - from vllm.model_executor.models.deepseek_v2 import GlmMoeDsaForCausalLM +from .nvidia.model import DeepseekV32ForCausalLM +from .nvidia.mtp import DeepseekV32MTP __all__ = [ "DeepseekV32ForCausalLM", "DeepseekV32MTP", - "GlmMoeDsaForCausalLM", ] diff --git a/vllm/models/deepseek_v32/nvidia/attention.py b/vllm/models/deepseek_v32/nvidia/attention.py index 8e60e871896..dcf955ad59b 100644 --- a/vllm/models/deepseek_v32/nvidia/attention.py +++ b/vllm/models/deepseek_v32/nvidia/attention.py @@ -4,7 +4,6 @@ import torch import torch.nn as nn from transformers import DeepseekV2Config, DeepseekV3Config -import vllm._custom_ops as ops from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import CacheConfig, VllmConfig from vllm.distributed import get_tensor_model_parallel_world_size @@ -22,15 +21,17 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import ( per_token_group_quant_fp8, ) from vllm.model_executor.layers.rotary_embedding import get_rope -from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer +from vllm.model_executor.layers.sparse_attn_indexer import ( + SparseAttnIndexer, + sparse_attn_indexer, +) from vllm.model_executor.models.deepseek_v2 import ( DeepSeekV2FusedQkvAProjLinear, DeepseekV32IndexerCache, yarn_get_mscale, ) from vllm.model_executor.models.utils import extract_layer_index -from vllm.platforms import current_platform -from vllm.utils.torch_utils import _encode_layer_name, is_quantized_kv_cache +from vllm.utils.torch_utils import is_quantized_kv_cache from .kernels import fused_norm_rope, fused_q @@ -170,11 +171,6 @@ class DeepseekV32Attention(MLAAttention): ) -> None: quant_config = vllm_config.quant_config cache_config = vllm_config.cache_config - if cache_config is not None and cache_config.cache_dtype == "auto": - # This implementation requires an FP8 sparse cache. Start with the - # generic FP8 format; MLAAttention will canonicalize it for the - # selected sparse backend. - cache_config.cache_dtype = "fp8" hidden_size = config.hidden_size qk_nope_head_dim = config.qk_nope_head_dim @@ -335,37 +331,6 @@ class DeepseekV32Attention(MLAAttention): rope_parameters=config.rope_parameters, is_neox_style=False, ) - # Decode-M GEMM dispatch for the bf16 A/B projections, measured on - # B300 and B200 (graph replay + rotating weights vs fp64): - # bf16_skinny_gemm for M <= skinny_max, dsv3_fused_a_gemm up to - # M = 16, cuBLAS above. Blackwell-only: the boundaries and the - # fused_a tile choices are tied to the 148/152-SM single-wave - # geometry. DSv3.2 q_b has no skinny tier (K=1536 does not fit the - # GEMV's 128x8 K-step; fused_a still wins 1.6-1.8x there). - # Quantized linear methods may not register .weight until - # process_weights_after_loading (e.g. compressed-tensors NVFP4 - # registers weight_packed at init) — probe with getattr. - # (N, K) -> skinny_max - qkv_a_dispatch = {(2624, 6144): 2, (2112, 7168): 2} - q_b_dispatch = {(2048, 2048): 2, (3072, 1536): 0} - ok = ( - current_platform.is_device_capability_family(100) - and hasattr(torch.ops._C, "bf16_skinny_gemm") - and hasattr(torch.ops._C, "dsv3_fused_a_gemm") - ) - qkv_a_weight = getattr(self.fused_qkv_a_proj, "weight", None) - q_b_weight = getattr(self.q_b_proj, "weight", None) - self._qkv_a_skinny_max = ( - qkv_a_dispatch.get(tuple(qkv_a_weight.shape)) - if ok and qkv_a_weight is not None and qkv_a_weight.dtype == torch.bfloat16 - else None - ) - self._q_b_skinny_max = ( - q_b_dispatch.get(tuple(q_b_weight.shape)) - if ok and q_b_weight is not None and q_b_weight.dtype == torch.bfloat16 - else None - ) - # Lightning indexer uses its own RoPE; interleave maps to non-NeoX. self.indexer_rope_emb = get_rope( qk_rope_head_dim, @@ -374,32 +339,13 @@ class DeepseekV32Attention(MLAAttention): is_neox_style=not getattr(config, "indexer_rope_interleave", False), ) - def _decode_m_gemm( - self, x: torch.Tensor, weight: torch.Tensor, skinny_max: int - ) -> torch.Tensor: - """x @ weight.T for decode M <= 16.""" - if x.shape[0] == 0: - return torch.empty((0, weight.shape[0]), dtype=x.dtype, device=x.device) - if x.shape[0] <= skinny_max: - return ops.bf16_skinny_gemm(x, weight) - out = torch.empty((x.shape[0], weight.shape[0]), dtype=x.dtype, device=x.device) - ops.dsv3_fused_a_gemm(out, x, weight.t()) - return out - def forward( # type: ignore[override] self, positions: torch.Tensor, hidden_states: torch.Tensor, ) -> torch.Tensor: # Captured: A-projections (+ indexer A-GEMM on indexer layers). - if self._qkv_a_skinny_max is not None and hidden_states.shape[0] <= 16: - qkv_lora = self._decode_m_gemm( - hidden_states, - self.fused_qkv_a_proj.weight, - self._qkv_a_skinny_max, - ) - else: - qkv_lora = self.fused_qkv_a_proj(hidden_states)[0] + qkv_lora = self.fused_qkv_a_proj(hidden_states)[0] q_c, kv_c, k_pe = qkv_lora.split( [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 ) @@ -452,9 +398,7 @@ class DeepseekV32Attention(MLAAttention): assert isinstance(slot_mapping, dict) mla_slot = slot_mapping.get(self.layer_name) - # skip_topk (index_share_for_mtp_iteration reuse mode): don't run the - # indexer this step; keep the top-k already in the shared buffer. - if self.indexer is not None and not self.skip_topk: + if self.indexer is not None: has_indexer = True indexer_k_norm_w = self.indexer.k_norm.weight indexer_k_norm_bias = self.indexer.k_norm.bias @@ -507,16 +451,12 @@ class DeepseekV32Attention(MLAAttention): index_rope_interleave=self._index_rope_interleave, ) - if self._q_b_skinny_max is not None and q_c.shape[0] <= 16: - q = self._decode_m_gemm(q_c, self.q_b_proj.weight, self._q_b_skinny_max) - else: - q = self.q_b_proj(q_c)[0] - q = q.view(-1, self.num_local_heads, self.qk_head_dim) + q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim) q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) q_nope = q_nope.transpose(0, 1) ql_nope = torch.bmm(q_nope, self.W_UK_T).transpose(0, 1) - if self.indexer is not None and not self.skip_topk: + if self.indexer is not None: index_q = self.indexer.wq_b(q_c)[0] index_q = index_q.view(-1, self.indexer.n_head, self.indexer.head_dim) else: @@ -538,10 +478,10 @@ class DeepseekV32Attention(MLAAttention): quantize_mqa=self._fp8_query, ) - if self.indexer is not None and not self.skip_topk: - torch.ops.vllm.sparse_attn_indexer( + if self.indexer is not None: + sparse_attn_indexer( q_c, - _encode_layer_name(self.indexer.k_cache.prefix), + self.indexer.k_cache.prefix, self.indexer.k_cache.kv_cache, index_q_fp8, None, # q_scale folded into weights on the fp8 path diff --git a/vllm/models/deepseek_v32/nvidia/kernels.py b/vllm/models/deepseek_v32/nvidia/kernels.py index c931ac7cb41..ef7b14ef1c5 100644 --- a/vllm/models/deepseek_v32/nvidia/kernels.py +++ b/vllm/models/deepseek_v32/nvidia/kernels.py @@ -2,80 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch -from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.utils.import_utils import has_cutedsl -from vllm.utils.torch_utils import direct_register_custom_op # Cache of tiny 1-element dummy tensors (per device, dtype) reused by the # has_indexer=False path so the indexer args don't allocate every call. _DUMMY_CACHE: dict[tuple, torch.Tensor] = {} -@torch.compiler.assume_constant_result -def _can_use_fused_q_cutedsl() -> bool: - return current_platform.has_device_capability(100) and has_cutedsl() - - -@torch.compiler.assume_constant_result -def _is_arch_support_pdl() -> bool: - return ( - current_platform.has_device_capability(100) - and current_platform.is_arch_support_pdl() - ) - - -def _fused_q_cutedsl_impl( - positions: torch.Tensor, - q_pe: torch.Tensor, - rope_cache: torch.Tensor, - ql_nope: torch.Tensor, - q_scale: torch.Tensor, - mqa_output: torch.Tensor, - idx_q: torch.Tensor, - idx_rope_cache: torch.Tensor, - idx_weights: torch.Tensor, - idx_weights_softmax_scale: float, - idx_weights_head_scale: float, - idx_q_fp8: torch.Tensor, - idx_weights_out: torch.Tensor, - has_indexer: bool, - index_rope_interleave: bool, -) -> None: - from .ops.fused_q_cutedsl import fused_q_cutedsl - - fused_q_cutedsl( - positions, - q_pe, - rope_cache, - ql_nope, - q_scale, - mqa_output, - idx_q, - idx_rope_cache, - idx_weights, - idx_weights_softmax_scale, - idx_weights_head_scale, - idx_q_fp8, - idx_weights_out, - has_indexer=has_indexer, - index_rope_interleave=index_rope_interleave, - ) - - -def _fused_q_cutedsl_fake(*args, **kwargs) -> None: - pass - - -direct_register_custom_op( - op_name="fused_q_cutedsl", - op_func=_fused_q_cutedsl_impl, - mutates_args=["mqa_output", "idx_q_fp8", "idx_weights_out"], - fake_impl=_fused_q_cutedsl_fake, - dispatch_key="CUDA", -) - - def _dummy(shape: tuple, dtype: torch.dtype, device: torch.device) -> torch.Tensor: key = (shape, dtype, device) t = _DUMMY_CACHE.get(key) @@ -214,14 +147,9 @@ def _fused_norm_rope_kernel( TOPK_BLOCK_SIZE: tl.constexpr, HAS_INDEXER: tl.constexpr, INDEX_ROPE_INTERLEAVE: tl.constexpr, - USE_PDL: tl.constexpr, ): pid = tl.program_id(0) tok_idx = tl.program_id(1) - if USE_PDL: - tl.extra.cuda.gdc_wait() - tl.extra.cuda.gdc_launch_dependents() - if pid == 3: if not HAS_INDEXER: # Shared layer: reuse the previous indexer layer's top-k; do not @@ -238,17 +166,6 @@ def _fused_norm_rope_kernel( ) return - if pid == 2: - # Q RMS norm does not depend on cache availability. In particular, - # profiling uses a negative dummy slot mapping but still consumes Q. - q_block = tl.arange(0, Q_BLOCK_SIZE) - q_mask = q_block < Q_DIM - q_c = tl.load(q_c_ptr + tok_idx * q_c_stride + q_block, mask=q_mask, other=0.0) - q_c_rms_w = tl.load(q_rms_norm_w_ptr + q_block, mask=q_mask) - q_c = _rms_norm(q_c, q_c_rms_w, q_rms_eps, Q_DIM) - tl.store(q_c_out_ptr + tok_idx * q_c_out_stride + q_block, q_c, mask=q_mask) - return - if slot_mapping_ptr is None: # Memory profiling run. return @@ -257,7 +174,15 @@ def _fused_norm_rope_kernel( # Padding return - if pid == 1: + if pid == 2: + # Q RMS norm + q_block = tl.arange(0, Q_BLOCK_SIZE) + q_mask = q_block < Q_DIM + q_c = tl.load(q_c_ptr + tok_idx * q_c_stride + q_block, mask=q_mask, other=0.0) + q_c_rms_w = tl.load(q_rms_norm_w_ptr + q_block, mask=q_mask) + q_c = _rms_norm(q_c, q_c_rms_w, q_rms_eps, Q_DIM) + tl.store(q_c_out_ptr + tok_idx * q_c_out_stride + q_block, q_c, mask=q_mask) + elif pid == 1: # KV RMS Norm + KV RoPE + MLA concat_and_cache. # Merged so the normed kv_c and RoPE'd k_pe can be written # to the MLA KV cache directly without a separate kernel. @@ -286,6 +211,9 @@ def _fused_norm_rope_kernel( r2 = x2 * cos + x1 * sin # MLA concat_and_cache: write [kv_c_normed, k_pe_roped] to cache. + if mla_cache_entry_stride == 0: + return + mla_block_size = mla_cache_block_stride // mla_cache_entry_stride mla_block_idx = slot_idx // mla_block_size mla_block_off = slot_idx % mla_block_size @@ -435,7 +363,7 @@ def _fused_norm_rope_kernel( ) -def _fused_norm_rope_impl( +def fused_norm_rope( positions: torch.Tensor, q_c: torch.Tensor, q_rms_norm_w: torch.Tensor, @@ -537,17 +465,12 @@ def _fused_norm_rope_impl( # Dummy values — pid 2 will skip the MLA cache write because # slot_mapping is all -1. mla_kv_cache = torch.empty(0, dtype=torch.bfloat16, device=device) - # Torch's Triton HOP mutation analysis compiles every program-id - # branch even though the dummy slot mapping returns before cache - # access. Keep the unused strides nonzero so that analysis can lower - # the block-size calculation safely. - mla_block_stride = 1 - mla_entry_stride = 1 + mla_block_stride = 0 + mla_entry_stride = 0 mla_k_scale = torch.ones(1, dtype=torch.float32, device=device) if q_c_out is None: q_c_out = torch.empty_like(q_c) - use_pdl = _is_arch_support_pdl() _fused_norm_rope_kernel[(4, num_tokens)]( positions, # Q RMS norm @@ -606,157 +529,6 @@ def _fused_norm_rope_impl( TOPK_BLOCK_SIZE=1024, HAS_INDEXER=has_indexer, INDEX_ROPE_INTERLEAVE=index_rope_interleave, - USE_PDL=use_pdl, - launch_pdl=use_pdl, - ) - return q_c_out - - -def _fused_norm_rope_op( - positions: torch.Tensor, - q_c: torch.Tensor, - q_rms_norm_w: torch.Tensor, - q_rms_eps: float, - kv_c: torch.Tensor, - kv_rms_norm_w: torch.Tensor, - kv_rms_eps: float, - k_pe: torch.Tensor, - k_rope_cos_sin_cache: torch.Tensor, - index_k: torch.Tensor | None, - index_k_layer_norm_w: torch.Tensor | None, - index_k_layer_norm_bias: torch.Tensor | None, - index_k_layer_norm_eps: float, - index_k_rope_cos_sin_cache: torch.Tensor | None, - topk_indices_buffer: torch.Tensor, - slot_mapping: torch.Tensor | None, - indexer_k_cache: torch.Tensor | None, - mla_kv_cache: torch.Tensor | None, - mla_kv_cache_dtype: str, - mla_k_scale: torch.Tensor | None, - has_indexer: bool, - index_rope_interleave: bool, - q_c_out: torch.Tensor, -) -> None: - _fused_norm_rope_impl( - positions, - q_c, - q_rms_norm_w, - q_rms_eps, - kv_c, - kv_rms_norm_w, - kv_rms_eps, - k_pe, - k_rope_cos_sin_cache, - index_k, - index_k_layer_norm_w, - index_k_layer_norm_bias, - index_k_layer_norm_eps, - index_k_rope_cos_sin_cache, - topk_indices_buffer, - slot_mapping=slot_mapping, - indexer_k_cache=indexer_k_cache, - mla_kv_cache=mla_kv_cache, - mla_kv_cache_dtype=mla_kv_cache_dtype, - mla_k_scale=mla_k_scale, - has_indexer=has_indexer, - index_rope_interleave=index_rope_interleave, - q_c_out=q_c_out, - ) - - -def _fused_norm_rope_fake( - positions: torch.Tensor, - q_c: torch.Tensor, - q_rms_norm_w: torch.Tensor, - q_rms_eps: float, - kv_c: torch.Tensor, - kv_rms_norm_w: torch.Tensor, - kv_rms_eps: float, - k_pe: torch.Tensor, - k_rope_cos_sin_cache: torch.Tensor, - index_k: torch.Tensor | None, - index_k_layer_norm_w: torch.Tensor | None, - index_k_layer_norm_bias: torch.Tensor | None, - index_k_layer_norm_eps: float, - index_k_rope_cos_sin_cache: torch.Tensor | None, - topk_indices_buffer: torch.Tensor, - slot_mapping: torch.Tensor | None, - indexer_k_cache: torch.Tensor | None, - mla_kv_cache: torch.Tensor | None, - mla_kv_cache_dtype: str, - mla_k_scale: torch.Tensor | None, - has_indexer: bool, - index_rope_interleave: bool, - q_c_out: torch.Tensor, -) -> None: - return - - -direct_register_custom_op( - op_name="fused_norm_rope_deepseek_v32", - op_func=_fused_norm_rope_op, - mutates_args=[ - "topk_indices_buffer", - "indexer_k_cache", - "mla_kv_cache", - "q_c_out", - ], - fake_impl=_fused_norm_rope_fake, - dispatch_key="CUDA", -) - - -def fused_norm_rope( - positions: torch.Tensor, - q_c: torch.Tensor, - q_rms_norm_w: torch.Tensor, - q_rms_eps: float, - kv_c: torch.Tensor, - kv_rms_norm_w: torch.Tensor, - kv_rms_eps: float, - k_pe: torch.Tensor, - k_rope_cos_sin_cache: torch.Tensor, - index_k: torch.Tensor | None, - index_k_layer_norm_w: torch.Tensor | None, - index_k_layer_norm_bias: torch.Tensor | None, - index_k_layer_norm_eps: float, - index_k_rope_cos_sin_cache: torch.Tensor | None, - topk_indices_buffer: torch.Tensor, - slot_mapping: torch.Tensor | None = None, - indexer_k_cache: torch.Tensor | None = None, - mla_kv_cache: torch.Tensor | None = None, - mla_kv_cache_dtype: str = "auto", - mla_k_scale: torch.Tensor | None = None, - has_indexer: bool = True, - index_rope_interleave: bool = False, - q_c_out: torch.Tensor | None = None, -) -> torch.Tensor: - if q_c_out is None: - q_c_out = torch.empty_like(q_c) - torch.ops.vllm.fused_norm_rope_deepseek_v32( - positions, - q_c, - q_rms_norm_w, - q_rms_eps, - kv_c, - kv_rms_norm_w, - kv_rms_eps, - k_pe, - k_rope_cos_sin_cache, - index_k, - index_k_layer_norm_w, - index_k_layer_norm_bias, - index_k_layer_norm_eps, - index_k_rope_cos_sin_cache, - topk_indices_buffer, - slot_mapping, - indexer_k_cache, - mla_kv_cache, - mla_kv_cache_dtype, - mla_k_scale, - has_indexer, - index_rope_interleave, - q_c_out, ) return q_c_out @@ -810,14 +582,10 @@ def _fused_q_kernel( HAS_INDEXER: tl.constexpr, INDEX_ROPE_INTERLEAVE: tl.constexpr, QUANTIZE_MQA: tl.constexpr, - USE_PDL: tl.constexpr, ): pid = tl.program_id(0) tok_idx = tl.program_id(1) head_idx = tl.program_id(2) - if USE_PDL: - tl.extra.cuda.gdc_wait() - tl.extra.cuda.gdc_launch_dependents() if pid == 2: # ql_nope quantize + pack into the front of mqa_q_fp8. On the bf16 @@ -1005,12 +773,10 @@ def fused_q( ``(ql_nope, q_pe)`` tuple the backend expects. """ assert positions.ndim == 1 - assert positions.dtype == torch.int64 assert q_pe.ndim == 3 assert q_pe_cos_sin_cache.ndim == 2 assert ql_nope.ndim == 3 assert ql_nope.shape[:2] == q_pe.shape[:2] - assert q_scale.dtype == torch.float32 and q_scale.numel() == 1 num_tokens = positions.shape[0] num_q_heads = q_pe.shape[1] @@ -1028,18 +794,7 @@ def fused_q( assert index_weights is not None num_index_q_heads = index_q.shape[1] index_q_head_dim = index_q.shape[2] - use_cutedsl = False - if _can_use_fused_q_cutedsl(): - from .ops.fused_q_cutedsl import is_fused_q_cutedsl_supported - - use_cutedsl = is_fused_q_cutedsl_supported( - q_pe, - index_q, - ql_nope, - has_indexer=has_indexer, - quantize_mqa=quantize_mqa, - ) - grid_heads = max(mqa_grid_heads, num_index_q_heads if has_indexer else 1) + grid_heads = max(mqa_grid_heads, num_index_q_heads) if quantize_mqa: # fp8 path: pack [ql_nope; q_pe] into a single fp8 tensor. mqa_q_fp8 = torch.empty( @@ -1060,27 +815,6 @@ def fused_q( index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn) index_weights_out = torch.empty_like(index_weights, dtype=torch.float32) - if use_cutedsl: - torch.ops.vllm.fused_q_cutedsl( - positions, - q_pe, - q_pe_cos_sin_cache, - ql_nope, - q_scale, - mqa_q, - index_q, - index_q_cos_sin_cache, - index_weights, - index_weights_softmax_scale, - index_weights_head_scale, - index_q_fp8, - index_weights_out, - has_indexer=has_indexer, - index_rope_interleave=index_rope_interleave, - ) - return index_q_fp8, index_weights_out, mqa_q - - use_pdl = _is_arch_support_pdl() _fused_q_kernel[(3, num_tokens, grid_heads)]( positions, q_pe, @@ -1122,8 +856,6 @@ def fused_q( HAS_INDEXER=has_indexer, INDEX_ROPE_INTERLEAVE=index_rope_interleave, QUANTIZE_MQA=quantize_mqa, - USE_PDL=use_pdl, - launch_pdl=use_pdl, # num_warps=1 is optimal here: each program is a single 128-element # rope+quant, so the kernel is program-count/occupancy bound, not # per-program compute bound (swept 1/2/4/8 — 1 wins or ties everywhere). diff --git a/vllm/models/deepseek_v32/nvidia/model.py b/vllm/models/deepseek_v32/nvidia/model.py index 90a6a2cf456..353aedc8cee 100644 --- a/vllm/models/deepseek_v32/nvidia/model.py +++ b/vllm/models/deepseek_v32/nvidia/model.py @@ -6,7 +6,6 @@ from itertools import islice import torch -from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import ( get_pp_group, @@ -40,7 +39,6 @@ from vllm.model_executor.models.utils import ( sequence_parallel_chunk, ) from vllm.sequence import IntermediateTensors -from vllm.v1.attention.backends.mla.sparse_utils import register_phys_shadow from .attention import DeepseekV32Attention from .fused_ops import fused_allreduce_rms_norm @@ -179,7 +177,6 @@ class DeepseekV32DecoderLayer(torch.nn.Module): return hidden_states, residual -@support_torch_compile class DeepseekV32Model(torch.nn.Module): fall_back_to_pt_during_load = False @@ -203,7 +200,6 @@ class DeepseekV32Model(torch.nn.Module): dtype=torch.int32, device=self.device, ) - register_phys_shadow(topk_indices_buffer) if get_pp_group().is_first_rank: self.embed_tokens = VocabParallelEmbedding( diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index bafbf383ca6..c8f2fcc5ffe 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -6,9 +6,9 @@ from collections.abc import Callable, Iterable import torch import torch.nn as nn -import vllm._custom_ops as ops from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig +from vllm.distributed import tensor_model_parallel_all_reduce from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) @@ -40,12 +40,7 @@ from vllm.model_executor.models.utils import ( ) from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors -from vllm.v1.attention.backends.mla.sparse_utils import ( - phys_shadow, - register_phys_shadow, -) -from .fused_ops import fused_allreduce_rms_norm from .kernels import fused_eh_norm from .model import DeepseekV32DecoderLayer @@ -61,19 +56,6 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module): self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False) - # bf16 skinny GEMM for the eh_proj, B300+B200 measured: GLM-5.2 - # (6144, 12288) wins at M <= 3, DSv3.2 (7168, 14336) at M <= 2; - # cuBLAS holds above (the 151/205MB weights stream at ~5.7TB/s). - eh_dispatch = {(6144, 12288): 3, (7168, 14336): 2} - eh_weight = getattr(self.eh_proj, "weight", None) - self._eh_skinny_max = ( - eh_dispatch.get(tuple(eh_weight.shape), 0) - if current_platform.is_device_capability_family(100) - and hasattr(torch.ops._C, "bf16_skinny_gemm") - and eh_weight is not None - and eh_weight.dtype == torch.bfloat16 - else 0 - ) topk_indices_buffer = torch.empty( vllm_config.scheduler_config.max_num_batched_tokens, @@ -81,7 +63,6 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module): dtype=torch.int32, device=current_platform.device_type, ) - register_phys_shadow(topk_indices_buffer) self.shared_head = SharedHead( config=config, prefix=prefix, quant_config=quant_config ) @@ -110,13 +91,7 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module): self.hnorm.weight, self.enorm.variance_epsilon, ) - if ( - eh_input.shape[0] <= self._eh_skinny_max - and eh_input.dtype == torch.bfloat16 - ): - hidden_states = ops.bf16_skinny_gemm(eh_input, self.eh_proj.weight) - else: - hidden_states = self.eh_proj(eh_input) + hidden_states = self.eh_proj(eh_input) hidden_states, residual = self.mtp_block( positions=positions, hidden_states=hidden_states, residual=None ) @@ -126,6 +101,9 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module): positions.shape[0], is_sequence_parallel=self.mtp_block.use_sequence_parallel_moe, ) + if not self.mtp_block.use_sequence_parallel_moe: + # Without sequence parallelism, the MoE output is left un-reduced. + hidden_states = tensor_model_parallel_all_reduce(hidden_states) # Recycle the POST-final-norm hidden into the next draft step. The # residual-add is fused into the final RMSNorm so it is computed # exactly once, and the result is returned for both tuple positions: @@ -136,14 +114,7 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module): # is understood by both the V2 speculator (isinstance-tuple check) and # the legacy proposer (model_returns_tuple is True for the # DeepSeekMTPModel architecture). - if self.mtp_block.use_sequence_parallel_moe: - hidden_states, _ = self.shared_head.norm(hidden_states, residual) - else: - # The MoE output is left un-reduced; fuse its all-reduce into the - # final norm, as the main model does at layer boundaries. - hidden_states, _ = fused_allreduce_rms_norm( - hidden_states, residual, self.shared_head.norm - ) + hidden_states, _ = self.shared_head.norm(hidden_states, residual) return hidden_states, hidden_states @@ -186,15 +157,6 @@ class DeepseekV32MultiTokenPredictor(nn.Module): if self_attn is not None and hasattr(self_attn, "topk_indices_buffer"): topk_indices_buffer = self_attn.topk_indices_buffer topk_indices_buffer[:num_slots] = topk_indices_buffer[slot_ids] - # The physical-index shadow mirrors this buffer row-wise - # (shadow[j] == convert(logical[j])); re-arrange it the same - # way or skip_topk draft steps read the step-0 multi-token - # layout's rows for the wrong tokens/requests. - shadow = phys_shadow(topk_indices_buffer) - if shadow is not None: - phys, seq = shadow - phys[:num_slots] = phys[slot_ids] - seq[:num_slots] = seq[slot_ids] def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) @@ -230,19 +192,6 @@ class DeepseekV32MultiTokenPredictor(nn.Module): # second RMSNorm. return self.logits_processor(mtp_layer.shared_head.head, hidden_states) - def get_top_tokens( - self, - hidden_states: torch.Tensor, - spec_step_idx: int = 0, - ) -> torch.Tensor: - # Greedy draft token ids via vocab-parallel local argmax: skips the - # full-vocab logits all-gather (use_local_argmax_reduction). - current_step_idx = spec_step_idx % self.num_mtp_layers - mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] - return self.logits_processor.get_top_tokens( - mtp_layer.shared_head.head, hidden_states - ) - class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -291,25 +240,6 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts): ) -> torch.Tensor | None: return self.model.compute_logits(hidden_states, spec_step_idx) - def get_top_tokens( - self, - hidden_states: torch.Tensor, - spec_step_idx: int | None = None, - ) -> torch.Tensor: - # The V2 speculator calls this without spec_step_idx; that is only - # correct while a single MTP layer is cycled for every draft step - # (num_nextn_predict_layers == 1, as in GLM-5.2/DSV3.2). Guard the - # assumption so a future multi-layer MTP fails loudly instead of - # silently using step 0's head for every step. - if spec_step_idx is None: - assert self.model.num_mtp_layers == 1, ( - "get_top_tokens called without spec_step_idx on a " - "multi-layer MTP; thread the draft step index through " - "the speculator." - ) - spec_step_idx = 0 - return self.model.get_top_tokens(hidden_states, spec_step_idx) - def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: spec_layer_weight_names = [ "embed_tokens", diff --git a/vllm/models/deepseek_v32/nvidia/ops/__init__.py b/vllm/models/deepseek_v32/nvidia/ops/__init__.py deleted file mode 100644 index 208f01a7cb5..00000000000 --- a/vllm/models/deepseek_v32/nvidia/ops/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/deepseek_v32/nvidia/ops/fused_q_cutedsl.py b/vllm/models/deepseek_v32/nvidia/ops/fused_q_cutedsl.py deleted file mode 100644 index b98147a34a3..00000000000 --- a/vllm/models/deepseek_v32/nvidia/ops/fused_q_cutedsl.py +++ /dev/null @@ -1,532 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from functools import cache - -import cutlass -import cutlass.cute as cute -import torch -from cuda.bindings.driver import CUstream -from cutlass import BFloat16, Float8E4M3FN, Float32, Int64, Uint8, Uint16, Uint32 - -from vllm.cute_utils import TORCH_TO_CUTE_DTYPE, cvt - - -def _make_fake_tensor(dtype, shape, divisibility): - stride = tuple( - 1 if i == len(shape) - 1 else cute.sym_int64(divisibility=divisibility) - for i in range(len(shape)) - ) - return cute.runtime.make_fake_tensor( - dtype, - shape, - stride, - assumed_align=divisibility * dtype.width // 8, - ) - - -def is_fused_q_cutedsl_supported( - q_pe: torch.Tensor, - index_q: torch.Tensor | None, - ql_nope: torch.Tensor, - *, - has_indexer: bool, - quantize_mqa: bool, -) -> bool: - if not ( - quantize_mqa - and q_pe.dtype == ql_nope.dtype == torch.bfloat16 - and q_pe.shape[-1] == 64 - and ql_nope.shape[-1] == 512 - ): - return False - return not has_indexer or ( - index_q is not None - and index_q.dtype == torch.bfloat16 - and index_q.shape[1] % 16 == 0 - and index_q.shape[-1] == 128 - ) - - -def fused_q_cutedsl( - positions: torch.Tensor, - q_pe: torch.Tensor, - rope_cache: torch.Tensor, - ql_nope: torch.Tensor, - q_scale: torch.Tensor, - mqa_output: torch.Tensor, - idx_q: torch.Tensor, - idx_rope_cache: torch.Tensor, - idx_weights: torch.Tensor, - idx_weights_softmax_scale: float, - idx_weights_head_scale: float, - idx_q_fp8: torch.Tensor, - idx_weights_out: torch.Tensor, - has_indexer: bool = True, - index_rope_interleave: bool = True, -) -> None: - _, num_heads, rope_dim = q_pe.shape - _, _, nope_dim = ql_nope.shape - _, num_idx_heads, idx_dim = idx_q.shape - - if has_indexer: - idx_rope_type = TORCH_TO_CUTE_DTYPE[idx_rope_cache.dtype] - idx_weights_type = TORCH_TO_CUTE_DTYPE[idx_weights.dtype] - else: - idx_dim = num_idx_heads = 0 - idx_q = idx_rope_cache = idx_q_fp8 = None - idx_weights = idx_weights_out = None - idx_rope_type = idx_weights_type = None - - rope_type = TORCH_TO_CUTE_DTYPE[rope_cache.dtype] - compiled = FusedQKernel.compile( - rope_dim, - nope_dim, - num_heads, - rope_type, - idx_dim, - num_idx_heads, - idx_rope_type, - idx_weights_type, - index_rope_interleave, - ) - compiled( - positions, - q_pe, - rope_cache, - ql_nope, - q_scale.view(1), - mqa_output, - idx_q, - idx_rope_cache, - idx_weights, - idx_q_fp8, - idx_weights_out, - float(idx_weights_softmax_scale * idx_weights_head_scale), - ) - - -class FusedQKernel: - def __init__( - self, - rope_dim: int, - nope_dim: int, - num_heads: int, - idx_dim: int, - num_idx_heads: int, - index_rope_interleave: bool, - ) -> None: - assert rope_dim == 64 - assert nope_dim == 512 - assert idx_dim in (128, 0) - - self.rope_dim = rope_dim - self.nope_dim = nope_dim - self.num_heads = num_heads - self.idx_dim = idx_dim - self.num_idx_heads = num_idx_heads - self.index_rope_interleave = index_rope_interleave - - # mqa: rope_dim=64, nope_dim=512, num_heads=64/TP - # indexer: rope_dim=64, nope_dim=64, num_heads=32 - - self.num_warps = 4 - assert num_heads % self.num_warps == 0 - assert num_idx_heads % (4 * self.num_warps) == 0 - self.num_ctas_per_tok = num_heads // self.num_warps - self.num_ctas_per_idx_tok = num_idx_heads // (4 * self.num_warps) - - @cute.jit - def __call__( - self, - positions: cute.Tensor, - q_pe: cute.Tensor, - rope_cache: cute.Tensor, - ql_nope: cute.Tensor, - q_scale: cute.Tensor, - mqa_output: cute.Tensor, - idx_q: cute.Tensor, - idx_rope_cache: cute.Tensor, - idx_weights: cute.Tensor, - idx_q_fp8: cute.Tensor, - idx_weights_out: cute.Tensor, - weight_scale: Float32, - stream: CUstream, - ): - num_tokens = positions.shape[0] - if cutlass.const_expr(self.idx_dim == 0): - grid = (num_tokens, self.num_ctas_per_tok, 1) - else: - num_mqa_ctas = num_tokens * self.num_ctas_per_tok - num_idx_ctas = num_tokens * self.num_ctas_per_idx_tok - grid = (num_mqa_ctas + num_idx_ctas, 1, 1) - - self.kernel( - positions, - q_pe, - rope_cache, - ql_nope, - q_scale, - mqa_output, - idx_q, - idx_rope_cache, - idx_weights, - idx_q_fp8, - idx_weights_out, - weight_scale, - ).launch( - grid=grid, - block=(self.num_warps * 32, 1, 1), - stream=stream, - use_pdl=True, - ) - - @cute.kernel - def kernel( - self, - positions: cute.Tensor, - q_pe: cute.Tensor, - q_pe_rope_cache: cute.Tensor, - ql_nope: cute.Tensor, - q_scale: cute.Tensor, - mqa_output: cute.Tensor, - idx_q: cute.Tensor, - idx_q_rope_cache: cute.Tensor, - idx_weights: cute.Tensor, - idx_q_fp8: cute.Tensor, - idx_weights_out: cute.Tensor, - weight_scale: Float32, - ): - if cutlass.const_expr(self.idx_dim == 0): - token_id, group_id, _ = cute.arch.block_idx() - self.mqa( - positions, - q_pe, - q_pe_rope_cache, - ql_nope, - q_scale, - mqa_output, - token_id, - group_id, - ) - else: - # CTA-specialization - bid, _, _ = cute.arch.block_idx() - num_mqa_ctas = positions.shape[0] * self.num_ctas_per_tok - if bid < num_mqa_ctas: - self.mqa( - positions, - q_pe, - q_pe_rope_cache, - ql_nope, - q_scale, - mqa_output, - bid // self.num_ctas_per_tok, - bid % self.num_ctas_per_tok, - ) - else: - bid -= num_mqa_ctas - self.indexer( - positions, - idx_q, - idx_q_rope_cache, - idx_weights, - idx_q_fp8, - idx_weights_out, - weight_scale, - bid // self.num_ctas_per_idx_tok, - bid % self.num_ctas_per_idx_tok, - ) - - @cute.jit - def mqa( - self, - positions: cute.Tensor, - q_pe: cute.Tensor, - q_pe_rope_cache: cute.Tensor, - ql_nope: cute.Tensor, - q_scale: cute.Tensor, - mqa_output: cute.Tensor, - token_id, - group_id, - ): - tid, _, _ = cute.arch.thread_idx() - warp_id = cute.arch.make_warp_uniform(tid // 32) - lane_id = tid % 32 - head_id = group_id * self.num_warps + warp_id - - cute.arch.griddepcontrol_wait() - - pos = positions[token_id] - inv_scale = 1.0 / q_scale[0] - - cp_op = cute.nvgpu.CopyUniversalOp() - cp_32B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=256) - cp_16B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) - cp_4B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=32) - cp_2B = cute.make_copy_atom(cp_op, Uint8, num_bits_per_copy=16) - - ##### issue all loads asap ##### - rQ_nope_bf16 = cute.make_rmem_tensor(16, BFloat16) - rQ_rope_bf16 = cute.make_rmem_tensor(2, BFloat16) - - src_ql_nope = cute.local_tile( - ql_nope[token_id, head_id, None], (16,), (lane_id,) - ) - src_q_rope = cute.local_tile(q_pe[token_id, head_id, None], (2,), (lane_id,)) - cute.copy(cp_32B, src_ql_nope, rQ_nope_bf16) - cute.copy(cp_4B, src_q_rope, rQ_rope_bf16) - - rCos_raw = q_pe_rope_cache[pos, 0 + lane_id] - rSin_raw = q_pe_rope_cache[pos, 32 + lane_id] - - ##### process NoPE ##### - rQ_nope_f32 = cvt.bf16x2_to_fp32x2(rQ_nope_bf16).load() * inv_scale - rQ_nope_f8 = cute.make_rmem_tensor(16, Float8E4M3FN) - rQ_nope_f8.store(rQ_nope_f32.to(Float8E4M3FN)) - dst_Q_nope = cute.local_tile( - mqa_output[token_id, head_id, None], (16,), (lane_id,) - ) - cute.copy(cp_16B, rQ_nope_f8, dst_Q_nope) - - ##### process RoPE ###### - rQ_rope_f32 = cvt.bf16x2_to_fp32x2(rQ_rope_bf16) - rCos = rCos_raw.to(Float32) - rSin = rSin_raw.to(Float32) - r0 = (rQ_rope_f32[0] * rCos - rQ_rope_f32[1] * rSin) * inv_scale - r1 = (rQ_rope_f32[1] * rCos + rQ_rope_f32[0] * rSin) * inv_scale - - cute.arch.griddepcontrol_launch_dependents() - - # TensorSSA fp32->fp8 cvt has a bug. rely on direct PTX - rQ_rope_f8 = cute.make_rmem_tensor(2, Float8E4M3FN) - cute.recast_tensor(rQ_rope_f8, Uint16)[0] = cvt.fp32x2_to_fp8x2(r0, r1) - dst_Q_rope = cute.local_tile( - mqa_output[token_id, head_id, None], (2,), (256 + lane_id,) - ) - cute.copy(cp_2B, rQ_rope_f8, dst_Q_rope) - - @cute.jit - def indexer( - self, - positions: cute.Tensor, - idx_q: cute.Tensor, - idx_q_rope_cache: cute.Tensor, - idx_weights: cute.Tensor, - idx_q_fp8: cute.Tensor, - idx_weights_out: cute.Tensor, - weight_scale: Float32, - token_id, - group_id, - ): - tid, _, _ = cute.arch.thread_idx() - subwarp_id = tid // 8 - sublane_id = tid % 8 - head_id = group_id * (4 * self.num_warps) + subwarp_id - cute.arch.griddepcontrol_wait() - - pos = positions[token_id] - - cp_op = cute.nvgpu.CopyUniversalOp() - cp_16B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) - cp_8B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64) - cp_4B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=32) - - ##### issue all loads first ##### - rQ_rope_bf16 = cute.make_rmem_tensor(8, BFloat16) - if cutlass.const_expr(self.index_rope_interleave): - src_idx_q_rope = cute.local_tile( - idx_q[token_id, head_id, None], (8,), (sublane_id,) - ) - cute.copy(cp_16B, src_idx_q_rope, rQ_rope_bf16) - else: - src_idx_q_rope = cute.zipped_divide( - idx_q[token_id, head_id, None], (4,) - ) # (4,32) - cute.copy( - cp_8B, - src_idx_q_rope[None, 0 + sublane_id], - cute.local_tile(rQ_rope_bf16, (4,), (0,)), - ) - cute.copy( - cp_8B, - src_idx_q_rope[None, 8 + sublane_id], - cute.local_tile(rQ_rope_bf16, (4,), (1,)), - ) - - rQ_nope_bf16 = cute.make_rmem_tensor(8, BFloat16) - src_idx_q_nope = cute.local_tile( - idx_q[token_id, head_id, None], (8,), (8 + sublane_id,) - ) - cute.copy(cp_16B, src_idx_q_nope, rQ_nope_bf16) - - rCos_raw = cute.make_rmem_tensor(4, idx_q_rope_cache.element_type) - rSin_raw = cute.make_rmem_tensor(4, idx_q_rope_cache.element_type) - rope_cache_view = cute.zipped_divide(idx_q_rope_cache[pos, None], (4,)) - - if cutlass.const_expr(idx_q_rope_cache.element_type == Float32): - cute.copy(cp_16B, rope_cache_view[None, sublane_id], rCos_raw) - cute.copy(cp_16B, rope_cache_view[None, 8 + sublane_id], rSin_raw) - elif cutlass.const_expr(idx_q_rope_cache.element_type == BFloat16): - cute.copy(cp_8B, rope_cache_view[None, sublane_id], rCos_raw) - cute.copy(cp_8B, rope_cache_view[None, 8 + sublane_id], rSin_raw) - - # unpack to FP32 - rQ_rope_f32 = cvt.bf16x2_to_fp32x2(rQ_rope_bf16) - rQ_nope_f32 = cvt.bf16x2_to_fp32x2(rQ_nope_bf16) - if cutlass.const_expr(idx_q_rope_cache.element_type == Float32): - rCos = rCos_raw - rSin = rSin_raw - elif cutlass.const_expr(idx_q_rope_cache.element_type == BFloat16): - rCos = cvt.bf16x2_to_fp32x2(rCos_raw) - rSin = cvt.bf16x2_to_fp32x2(rSin_raw) - - # apply rope - for i in cutlass.range_constexpr(4): - if cutlass.const_expr(self.index_rope_interleave): - r0 = rQ_rope_f32[i * 2 + 0] * rCos[i] - rQ_rope_f32[i * 2 + 1] * rSin[i] - r1 = rQ_rope_f32[i * 2 + 1] * rCos[i] + rQ_rope_f32[i * 2 + 0] * rSin[i] - rQ_rope_f32[i * 2 + 0] = r0 - rQ_rope_f32[i * 2 + 1] = r1 - else: - r0 = rQ_rope_f32[0 + i] * rCos[i] - rQ_rope_f32[4 + i] * rSin[i] - r1 = rQ_rope_f32[4 + i] * rCos[i] + rQ_rope_f32[0 + i] * rSin[i] - rQ_rope_f32[0 + i] = r0 - rQ_rope_f32[4 + i] = r1 - - # amax - amax = Float32(1e-4) - for i in cutlass.range_constexpr(8): - amax = cute.arch.fmax(amax, cute.math.absf(rQ_rope_f32[i])) - amax = cute.arch.fmax(amax, cute.math.absf(rQ_nope_f32[i])) - - # warp reduction among 8 lanes - for i in cutlass.range_constexpr(3): - other = cute.arch.shuffle_sync_bfly(amax, 1 << i) - amax = cute.arch.fmax(amax, other) - - # compute scale from amax - # exp2(ceil(log2(scale))) via bit manipulation - scale = amax * (1.0 / 448.0) - bits = scale.bitcast(Uint32) - exp_bits = (bits + Uint32(0x007FFFFF)) & Uint32(0x7F800000) - scale = exp_bits.bitcast(Float32) - inv_scale = (Uint32(0x7F000000) - exp_bits).bitcast(Float32) - - for i in cutlass.range_constexpr(8): - rQ_nope_f32[i] *= inv_scale - rQ_rope_f32[i] *= inv_scale - - cute.arch.griddepcontrol_launch_dependents() - - # quantize and store - rQ_nope_f8 = cute.make_rmem_tensor(8, Float8E4M3FN) - rQ_nope_f8.store(rQ_nope_f32.load().to(Float8E4M3FN)) - dst_idx_q_nope = cute.local_tile( - idx_q_fp8[token_id, head_id, None], (8,), (8 + sublane_id,) - ) - cute.copy(cp_8B, rQ_nope_f8, dst_idx_q_nope) - - rQ_rope_f8 = cute.make_rmem_tensor(8, Float8E4M3FN) - rQ_rope_f8.store(rQ_rope_f32.load().to(Float8E4M3FN)) - if cutlass.const_expr(self.index_rope_interleave): - dst_idx_q_rope = cute.local_tile( - idx_q_fp8[token_id, head_id, None], (8,), (sublane_id,) - ) - cute.copy(cp_8B, rQ_rope_f8, dst_idx_q_rope) - else: - dst_idx_q_rope = cute.zipped_divide( - idx_q_fp8[token_id, head_id, None], (4,) - ) # (4,32) - cute.copy( - cp_4B, - cute.local_tile(rQ_rope_f8, (4,), (0,)), - dst_idx_q_rope[None, 0 + sublane_id], - ) - cute.copy( - cp_4B, - cute.local_tile(rQ_rope_f8, (4,), (1,)), - dst_idx_q_rope[None, 8 + sublane_id], - ) - - # scale indexer weights - if sublane_id == 0: - w = idx_weights[token_id, head_id].to(Float32) - idx_weights_out[token_id, head_id] = w * scale * weight_scale - - @cache - @staticmethod - def compile( - rope_dim: int, - nope_dim: int, - num_heads: int, - rope_type: type[cutlass.Numeric], - idx_dim: int, - num_idx_heads: int, - idx_rope_type: type[cutlass.Numeric] | None, - idx_weights_type: type[cutlass.Numeric] | None, - index_rope_interleave: bool, - ): - num_tokens = cute.sym_int() - max_pos = cute.sym_int() - - positions = _make_fake_tensor(Int64, (num_tokens,), divisibility=1) - q_pe = _make_fake_tensor( - BFloat16, (num_tokens, num_heads, rope_dim), divisibility=16 - ) - rope_cache = _make_fake_tensor(rope_type, (max_pos, rope_dim), divisibility=8) - ql_nope = _make_fake_tensor( - BFloat16, (num_tokens, num_heads, nope_dim), divisibility=16 - ) - q_scale = _make_fake_tensor(Float32, (1,), divisibility=4) - mqa_output = _make_fake_tensor( - Float8E4M3FN, - (num_tokens, num_heads, nope_dim + rope_dim), - divisibility=16, - ) - - if idx_rope_type is not None: - index_q = _make_fake_tensor( - BFloat16, (num_tokens, num_idx_heads, idx_dim), divisibility=16 - ) - index_rope_cache = _make_fake_tensor( - idx_rope_type, (max_pos, rope_dim), divisibility=8 - ) - index_weights = _make_fake_tensor( - idx_weights_type, (num_tokens, num_idx_heads), divisibility=8 - ) - index_q_fp8 = _make_fake_tensor( - Float8E4M3FN, (num_tokens, num_idx_heads, idx_dim), divisibility=16 - ) - index_weights_out = _make_fake_tensor( - Float32, (num_tokens, num_idx_heads), divisibility=4 - ) - else: - index_q = index_rope_cache = index_q_fp8 = None - index_weights = index_weights_out = None - - kernel = FusedQKernel( - rope_dim, - nope_dim, - num_heads, - idx_dim, - num_idx_heads, - index_rope_interleave, - ) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, - positions, - q_pe, - rope_cache, - ql_nope, - q_scale, - mqa_output, - index_q, - index_rope_cache, - index_weights, - index_q_fp8, - index_weights_out, - Float32(0.0), - stream, - options="--enable-tvm-ffi", - ) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 36df138a4ec..a6b9f4abb0d 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -28,7 +28,6 @@ from vllm.v1.attention.backend import ( MultipleOf, ) from vllm.v1.attention.backends.mla.sparse_utils import ( - phys_shadow, triton_convert_req_index_to_global_index, triton_filter_and_convert_dcp_index, ) @@ -405,36 +404,6 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]): NUM_TOPK_TOKENS=topk_indices.shape[1], return_valid_counts=True, ) - elif (shadow := phys_shadow(self.topk_indices_buffer)) is not None: - # skip_topk layers reuse the previous fresh layer's PHYSICAL - # indices: the logical top-k indices are shared (one buffer for - # the whole forward) and block_table/req_id are per-step constant, - # so re-converting them per layer produces identical output. This - # covers both the backbone (index_topk_freq > 1) and MTP draft - # steps 1+ (index_share_for_mtp_iteration sets skip_topk; the - # draft compact re-arranges the shadow together with the logical - # buffer). The fresh/skip split is fixed per captured graph, so - # the branch is cudagraph-safe. - phys_buf, seq_buf = shadow - wrote_fresh = getattr(layer, "indexer", None) is not None and not getattr( - layer, "skip_topk", False - ) - if wrote_fresh: - topk_indices_physical, seq_lens = ( - triton_convert_req_index_to_global_index( - attn_metadata.req_id_per_token[:num_actual_toks], - attn_metadata.block_table, - topk_indices, - BLOCK_SIZE=attn_metadata.block_size, - NUM_TOPK_TOKENS=topk_indices.shape[1], - return_valid_counts=True, - out=phys_buf[:num_actual_toks], - valid_counts_out=seq_buf[:num_actual_toks], - ) - ) - else: - topk_indices_physical = phys_buf[:num_actual_toks] - seq_lens = seq_buf[:num_actual_toks] else: topk_indices_physical, seq_lens = triton_convert_req_index_to_global_index( attn_metadata.req_id_per_token[:num_actual_toks], diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py index 3cc141e292b..681ab7cfd9f 100644 --- a/vllm/v1/attention/backends/mla/sparse_utils.py +++ b/vllm/v1/attention/backends/mla/sparse_utils.py @@ -128,8 +128,6 @@ def triton_convert_req_index_to_global_index( prefill_workspace_request_ids: torch.Tensor | None = None, prefill_workspace_starts: torch.Tensor | None = None, return_valid_counts: bool = False, - out: torch.Tensor | None = None, - valid_counts_out: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """ out[token_id, indice_id] = @@ -180,20 +178,14 @@ def triton_convert_req_index_to_global_index( req_id_c = req_id.contiguous() block_table_c = block_table.contiguous() token_indices_c = token_indices.contiguous() - # `out`/`valid_counts_out` allow writing into a stable pre-allocated buffer - # (used by the shared-physical-index cache to stay cudagraph-safe). - out = torch.empty_like(token_indices_c) if out is None else out + out = torch.empty_like(token_indices_c) # Allocate valid count buffer if needed (must be zero-initialized for atomics) valid_counts: torch.Tensor | None = None if return_valid_counts: - if valid_counts_out is None: - valid_counts = torch.zeros( - num_tokens, dtype=torch.int32, device=token_indices.device - ) - else: - valid_counts = valid_counts_out - valid_counts.zero_() + valid_counts = torch.zeros( + num_tokens, dtype=torch.int32, device=token_indices.device + ) # Strides in elements bt_stride0, bt_stride1 = block_table_c.stride() @@ -349,47 +341,3 @@ def triton_filter_and_convert_dcp_index( assert valid_counts is not None return out, valid_counts return out - - -# --- Physical-index shadow (DSA index_topk_freq > 1) ------------------------- -# GLM-5.2 uses index_topk_freq=4: only ~1/4 of layers write a fresh top-k into -# the single shared topk_indices_buffer; the rest reuse it. Since block_table -# and req_id_per_token are constant within a decode step, the physical indices -# (block_table lookup of the logical top-k) are identical across a freq-group. -# Fresh layers convert once into a STABLE shadow of the logical buffer and -# skip layers read it -- eliminating ~3/4 of the per-layer convert kernels. -# -# Invariant: shadow[j] == convert(logical[j]) row-wise. Whoever re-arranges -# the logical buffer (e.g. the MTP draft compact between the multi-token -# step-0 layout and the single-token steps-1+ layout) must apply the same -# gather to the shadow. -# -# Shadows are registered ONLY at buffer creation time (model init, never -# inside cudagraph capture) via register_phys_shadow; every other access is -# the read-only phys_shadow lookup, so addresses are stable across graph -# replays and no allocation can happen during capture. -import weakref as _weakref # noqa: E402 - -_PHYS_SHADOWS: dict[int, tuple[torch.Tensor, torch.Tensor]] = {} - - -def register_phys_shadow(topk_buf: torch.Tensor) -> None: - """Allocate the (physical-index, valid-count) shadow for a logical top-k - buffer. Call once where the buffer is created. The finalizer drops the - entry with the owning buffer, so a recycled id can never alias a stale - shadow.""" - key = id(topk_buf) - if key not in _PHYS_SHADOWS: - _PHYS_SHADOWS[key] = ( - torch.empty_like(topk_buf), - torch.empty(topk_buf.shape[0], dtype=torch.int32, device=topk_buf.device), - ) - _weakref.finalize(topk_buf, _PHYS_SHADOWS.pop, key, None) - - -def phys_shadow( - topk_buf: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor] | None: - """The registered shadow for this buffer, or None if it was never - registered. Never allocates; capture-safe.""" - return _PHYS_SHADOWS.get(id(topk_buf)) diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index bb1c59a8c2a..6ee442af51e 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -1009,9 +1009,8 @@ class SpecDecodeBaseProposer: # DeepSeek-family MTP (deepseek_mtp.py) recycles the post-final- # norm hidden, so its forward returns (logit_hidden, # recycle_hidden). Other MTP families return a single tensor. - return bool( - {"DeepSeekMTPModel", "DeepseekV32MTPModel"} - & set(self.draft_model_config.hf_config.architectures or []) + return "DeepSeekMTPModel" in ( + self.draft_model_config.hf_config.architectures or [] ) return self.method not in ("mtp", "draft_model", "dflash") diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index 9a14d132c55..b890574971c 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -48,18 +48,6 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): self.prefill_cudagraph_manager: SpeculatorCudaGraphManager | None = None self.decode_cudagraph_manager: SpeculatorCudaGraphManager | None = None - # Lifecycle hooks for model-specific optimizations. Subclasses override - # the ones they need. These fire in both `capture` and `propose` so that - # any state they toggle (e.g. attention flags baked into a CUDA graph) is - # identical at capture time and replay time. - def on_prefill_begin(self, num_reqs: int, num_tokens: int) -> None: ... - - def on_prefill_end(self, num_reqs: int, num_tokens: int) -> None: ... - - def on_multi_step_decode_begin(self, num_reqs: int) -> None: ... - - def on_multi_step_decode_end(self, num_reqs: int) -> None: ... - @property def advance_draft_positions(self) -> bool: """ @@ -110,9 +98,6 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): assert self.prefill_cudagraph_manager is not None if self.prefill_cudagraph_manager.use_breakable_cg: self.prefill_cudagraph_manager.init_breakable_cg_runner(self.model) - - self.on_prefill_begin(self.max_num_reqs, self.max_num_tokens) - self.prefill_cudagraph_manager.capture( self._prefill, self.model_state, @@ -123,13 +108,9 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): progress_bar_desc="Capturing prefill CUDA graphs", ) - self.on_prefill_end(self.max_num_reqs, self.max_num_tokens) - if self.num_speculative_steps == 1: return - self.on_multi_step_decode_begin(self.max_num_reqs) - # Capture the decode draft generation routine (model forward + # sample + update_draft_inputs) for a single # step. @@ -144,8 +125,6 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): progress_bar_desc="Capturing decode CUDA graphs", ) - self.on_multi_step_decode_end(self.max_num_reqs) - @torch.inference_mode() def propose( self, @@ -174,8 +153,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, is_profile: bool = False, ) -> torch.Tensor: - num_tokens = input_batch.num_tokens - num_tokens_padded = input_batch.num_tokens_after_padding + num_tokens = input_batch.num_tokens_after_padding num_reqs = input_batch.num_reqs max_query_len = input_batch.num_scheduled_tokens.max() max_seq_len = input_batch.seq_lens_cpu_upper_bound[:num_reqs].max().item() @@ -196,7 +174,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): ) else: hidden_states = last_hidden_states - self.hidden_states[:num_tokens_padded].copy_(hidden_states) + self.hidden_states[:num_tokens].copy_(hidden_states) self._copy_request_inputs( num_reqs, @@ -224,22 +202,20 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): num_reqs, # Use the actual number of tokens without padding added by # the target model during FULL cudagraph. - num_tokens, + input_batch.num_tokens, max_query_len, ) prefill_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp( self.prefill_cudagraph_manager, num_reqs, - num_tokens_padded, + num_tokens, uniform_token_count, dp_size=self.dp_size, dp_rank=self.dp_rank, need_eager=is_profile, ) - self._prepare_eplb_forward(num_tokens) - - self.on_prefill_begin(num_reqs, num_tokens) + self._prepare_eplb_forward(input_batch.num_tokens) if prefill_batch_desc.cg_mode == CUDAGraphMode.FULL: # Replay the full graph for draft prefill. @@ -259,8 +235,6 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): mm_inputs=mm_inputs, ) - self.on_prefill_end(num_reqs, num_tokens) - if self.num_speculative_steps == 1: # Early exit. return self.draft_tokens[:num_reqs, :1] @@ -288,8 +262,6 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): need_eager=is_profile, ) - self.on_multi_step_decode_begin(num_reqs) - # Generate the remaining num_speculative_steps - 1 draft tokens. self._multi_step_decode( num_reqs, @@ -299,8 +271,6 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): input_batch.seq_lens_cpu_upper_bound, ) - self.on_multi_step_decode_end(num_reqs) - return self.draft_tokens[:num_reqs] @torch.inference_mode() diff --git a/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py index 5f1cebc519b..4b9354f23e7 100644 --- a/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py @@ -1,10 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import torch import torch.nn as nn -from vllm.config import VllmConfig from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import ( AutoRegressiveSpeculator, ) @@ -12,46 +10,9 @@ from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model class MTPSpeculator(AutoRegressiveSpeculator): - def __init__(self, vllm_config: VllmConfig, device: torch.device): - super().__init__(vllm_config, device) - - spec_config = vllm_config.speculative_config - draft_hf_config = ( - spec_config.draft_model_config.hf_config - if spec_config is not None - else None - ) - # Detect index_share_for_mtp_iteration. When True, the proposer - # toggles skip_topk so step 0 computes MTP's own indices and - # steps 1+ reuse them. - self.share_mtp_topk_indices = getattr( - draft_hf_config, "index_share_for_mtp_iteration", False - ) - def load_draft_model( self, target_model: nn.Module, target_attn_layer_names: set[str], ) -> nn.Module: - draft_model = load_eagle_model(target_model, self.vllm_config) - self.share_mtp_topk_indices = self.share_mtp_topk_indices and hasattr( - draft_model.model, "set_skip_topk" - ) - return draft_model - - def on_prefill_end(self, num_reqs: int, num_tokens: int) -> None: - # Step 0 (prefill) wrote topk indices for every query token in the - # multi-token batch. Compact them down to each request's last token so - # steps 1+ can reuse them from the shared buffer. - if self.share_mtp_topk_indices and self.num_speculative_steps > 1: - self.model.model.compact_topk_indices(self.last_token_indices[:num_reqs]) - - def on_multi_step_decode_begin(self, num_reqs: int) -> None: - # Switch to reuse mode so draft steps 1+ skip the indexer op and read - # the indices that step 0 wrote into the shared buffer. - if self.share_mtp_topk_indices: - self.model.model.set_skip_topk(True) - - def on_multi_step_decode_end(self, num_reqs: int) -> None: - if self.share_mtp_topk_indices: - self.model.model.set_skip_topk(False) + return load_eagle_model(target_model, self.vllm_config) From 6a1acac3fe2924ca221fdf048b74f9a020dabddd Mon Sep 17 00:00:00 2001 From: liuzhenwei Date: Sat, 25 Jul 2026 08:36:53 +0800 Subject: [PATCH 022/185] [BUGFIX] Fix log capture in KV test (#49655) Signed-off-by: zhenwei-intel Co-authored-by: Kunshang Ji --- tests/v1/core/test_kv_cache_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 079694d505b..4ee2ea9ef4b 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -2183,7 +2183,7 @@ def test_mla_draft_prefers_standard_layout_when_pages_can_be_unified(): ) -def test_mla_with_incompatible_swa_uses_one_full_allocation_group(caplog): +def test_mla_with_incompatible_swa_uses_one_full_allocation_group(caplog_vllm): # Sparse MLA pages cannot be padded safely. Keeping the draft's attention # compute sliding-window while promoting only its allocation semantics lets # every layer share the target's block table and remain contiguous. @@ -2206,7 +2206,7 @@ def test_mla_with_incompatible_swa_uses_one_full_allocation_group(caplog): assert promoted_draft.block_size == 64 assert promoted_draft.sliding_window == draft.sliding_window assert specs["draft.0"] is draft - assert "attention compute is unchanged" in caplog.text + assert "attention compute is unchanged" in caplog_vllm.text def test_get_kv_cache_spec_kind_prefers_specific_attention_subclasses(): From 318b527cc2d1f672683407be05ea26a2cf1f3ea6 Mon Sep 17 00:00:00 2001 From: liuzhenwei Date: Sat, 25 Jul 2026 10:26:15 +0800 Subject: [PATCH 023/185] [XPU] add warning for xpu graph limitations (#49419) Signed-off-by: zhenwei-intel --- vllm/platforms/xpu.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index 80f99acc5ac..e673713330f 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -282,6 +282,16 @@ class XPUPlatform(Platform): "XPU Graph is disabled by environment variable, " "please set VLLM_XPU_ENABLE_XPU_GRAPH=1 to enable it." ) + else: + logger.warning_once( + "XPU Graph support is experimental and has known limitations: " + "(1) only single-GPU execution is supported; " + "(2) FLASH_ATTN supports PIECEWISE mode only; use TRITON_ATTN " + "for FULL mode; " + "(3) XPU Graph may increase device memory usage, " + "potentially causing OOM errors or leaving less memory " + "for the KV cache and reducing performance." + ) # Disable fusion passes not yet supported on XPU. from vllm.config.compilation import CompilationMode From 70052fb924fbd633580308562d43b9df838ddf6d Mon Sep 17 00:00:00 2001 From: Yifan Qiao Date: Fri, 24 Jul 2026 21:02:22 -0700 Subject: [PATCH 024/185] [Bugfix][KV Connector][Mooncake] Keep TP-sharded Mamba state out of the KV-head dedup (#49499) Signed-off-by: Yifan Qiao Co-authored-by: Claude --- .../unit/test_mooncake_store_hma_e2e.py | 3 +- .../unit/test_mooncake_store_worker.py | 226 +++++++++++++++++- .../kv_connector/v1/mooncake/store/worker.py | 175 ++++++++------ 3 files changed, 311 insertions(+), 93 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index 46a70654b1a..2ea776adf6f 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py @@ -173,7 +173,6 @@ def test_e2e_swa_plus_full_save_then_lookup_hits(): worker = _build_worker_with_dict_store(vllm_config, cfg, store) worker.tp_size = 1 worker.pp_size = 1 - worker.put_step = 1 worker.num_kv_head = 8 # Register kv_caches using mocked thread classes so register_kv_caches @@ -215,7 +214,7 @@ def test_e2e_swa_plus_full_save_then_lookup_hits(): block_size=worker.block_size, coord=worker.coord, tp_rank=worker.tp_rank, - put_step=worker.put_step, + group_put_steps=worker._group_tp_replication_factors, kv_role=worker.kv_role, ready_event=ready, enable_kv_event=False, diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index ec47bfa2385..7cddd6d4a7c 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -90,7 +90,7 @@ def _make_store_sending_thread( block_size=block_size, coord=coord, tp_rank=tp_rank, - put_step=put_step, + group_put_steps=[put_step] * len(token_databases), kv_role="kv_producer", ready_event=threading.Event(), replicate_config=replicate_config, @@ -527,6 +527,29 @@ def test_store_sending_thread_delta_strides_with_local_phase(): assert store.batch_put_from_multi_buffers.call_args.args[0] == keys +def test_tp_sharded_group_saves_every_block_on_every_rank(): + """Sharded ranks must write every block because peers hold different bytes.""" + store = MagicMock() + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) + store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys) + thread = _make_store_sending_thread(store, tp_rank=0, put_step=2) + thread.group_put_steps = [1] + + thread.add_stored_request("req-a") + thread._handle_request( + ReqMeta( + req_id="req-a", + token_len_chunk=64, + block_ids=([0, 1, 2, 3],), + block_hashes=[b"a0", b"a1", b"a2", b"a3"], + can_save=True, + ) + ) + + keys = store.batch_is_exist.call_args.args[0] + assert len(keys) == 4 + + def test_store_sending_thread_retries_skipped_range_after_pressure(): store = MagicMock() store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) @@ -1274,7 +1297,8 @@ def test_worker_put_striding_covers_every_rank_get_namespace( ] assert len(keys) == len(block_hashes) # PUT side: mirrors KVCacheStoreSendingThread's striding slice. - put_keys.update(keys[w.tp_rank % w.put_step :: w.put_step]) + put_step = w._group_tp_replication_factors[0] + put_keys.update(keys[w.tp_rank % put_step :: put_step]) # GET side: KVCacheStoreRecvingThread fetches every key. get_keys_per_rank[tp_rank] = set(keys) @@ -1613,6 +1637,15 @@ def _register_with_mocked_threads( worker.register_kv_caches(kv_caches) +def _refresh_group_tp_replication_factors( + worker: mooncake_store_worker.MooncakeStoreWorker, +) -> None: + worker._group_tp_replication_factors = ( + worker._compute_group_tp_replication_factors() + ) + worker._init_lookup_key_prefixes() + + def _make_bare_worker( *, num_gpu_blocks: int = 10, @@ -1630,11 +1663,9 @@ def _make_bare_worker( worker.cache_config.num_gpu_blocks = num_gpu_blocks worker.store = MagicMock() worker.store.register_buffer.return_value = 0 - worker.use_mla = False worker.kv_role = kv_role worker.block_size = block_size worker.tp_rank = 0 - worker.put_step = 1 worker.enable_kv_events = False worker.kv_send_thread = None worker.kv_recv_threads = [] @@ -1664,7 +1695,6 @@ def _make_bare_worker( worker.pcp_size = 1 worker.dcp_size = 1 worker.hash_block_size = block_size - worker.metadata = KeyMetadata("test-model", 0, 0, 0, 0) # Pre-build a single-group token_dbs so lookup-only tests don't have to # call register_kv_caches. worker.token_dbs = [ @@ -1679,7 +1709,7 @@ def _make_bare_worker( scheduler_block_size=block_size, hash_block_size=block_size, ) - worker._init_lookup_key_prefixes() + _refresh_group_tp_replication_factors(worker) return worker @@ -1688,9 +1718,8 @@ def test_lookup_key_prefixes_cover_dcp_rank_namespaces(): worker.tp_size = 4 worker.num_kv_head = 1 worker.dcp_size = 4 - worker._init_lookup_key_prefixes() + _refresh_group_tp_replication_factors(worker) - assert worker._lookup_expected_per_key == 4 assert worker._lookup_key_prefixes[0] == ( "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0", "test-model@tp_rank:1@pcp0@dcp1@pp_rank:0@group:0", @@ -1705,21 +1734,192 @@ def test_lookup_key_prefixes_cover_pcp_rank_namespaces(): worker.num_kv_head = 1 worker.pcp_size = 2 worker.dcp_size = 1 - worker._init_lookup_key_prefixes() + _refresh_group_tp_replication_factors(worker) - assert worker._lookup_expected_per_key == 2 assert worker._lookup_key_prefixes[0] == ( "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0", "test-model@tp_rank:0@pcp1@dcp0@pp_rank:0@group:0", ) +def test_lookup_key_prefixes_expand_tp_sharded_groups_per_rank(): + """Replicated attention needs one namespace; sharded Mamba needs every rank.""" + from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheGroupSpec, + MambaSpec, + ) + + worker = _make_bare_worker(block_size=16) + worker.tp_size = 2 + worker.num_kv_head = 1 + fa = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) + mamba = MambaSpec( + block_size=16, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + worker._kv_cache_groups = [ + KVCacheGroupSpec(["l0"], fa), + KVCacheGroupSpec(["l1"], mamba), + ] + worker.token_dbs = [ + ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=0), block_size=16 + ), + ChunkedTokenDatabase( + KeyMetadata("test-model", 1, 0, 0, 0, group_id=1), block_size=16 + ), + ] + _refresh_group_tp_replication_factors(worker) + + assert worker._lookup_key_prefixes[0] == ( + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0", + ) + assert worker._lookup_key_prefixes[1] == ( + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1", + "test-model@tp_rank:1@pcp0@dcp0@pp_rank:0@group:1", + ) + + +def test_group_tp_replication_factors_mixed_mla_gqa_mamba(): + from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheGroupSpec, + MambaSpec, + MLAAttentionSpec, + ) + + worker = _make_bare_worker(block_size=16) + worker.tp_size = 4 + worker.num_kv_head = 2 + mla = MLAAttentionSpec(block_size=16, num_kv_heads=1, head_size=64, dtype=None) + gqa = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) + mamba = MambaSpec( + block_size=16, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + worker._kv_cache_groups = [ + KVCacheGroupSpec(["l0"], mla), + KVCacheGroupSpec(["l1"], gqa), + KVCacheGroupSpec(["l2"], mamba), + ] + worker.token_dbs = [ + ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=g_idx), block_size=16 + ) + for g_idx in range(3) + ] + + _refresh_group_tp_replication_factors(worker) + assert worker._group_tp_replication_factors == (4, 2, 1) + assert worker._lookup_key_prefixes[0] == ( + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0", + ) + assert worker._lookup_key_prefixes[1] == ( + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1", + "test-model@tp_rank:1@pcp0@dcp0@pp_rank:0@group:1", + ) + assert worker._lookup_key_prefixes[2] == ( + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:2", + "test-model@tp_rank:1@pcp0@dcp0@pp_rank:0@group:2", + "test-model@tp_rank:2@pcp0@dcp0@pp_rank:0@group:2", + "test-model@tp_rank:3@pcp0@dcp0@pp_rank:0@group:2", + ) + + +@pytest.mark.parametrize("spec_order", [("mla", "gqa"), ("gqa", "mla")]) +def test_uniform_group_uses_common_inner_replication_factor(spec_order): + from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheGroupSpec, + MLAAttentionSpec, + UniformTypeKVCacheSpecs, + ) + + worker = _make_bare_worker(block_size=16) + worker.tp_size = 4 + worker.num_kv_head = 2 + specs_by_name = { + "mla": MLAAttentionSpec( + block_size=16, num_kv_heads=1, head_size=64, dtype=None + ), + "gqa": FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=64, dtype=None + ), + } + inner_specs = {name: specs_by_name[name] for name in spec_order} + uniform_spec = UniformTypeKVCacheSpecs( + block_size=16, + kv_cache_specs=inner_specs, + ) + worker._kv_cache_groups = [ + KVCacheGroupSpec(list(inner_specs), uniform_spec), + ] + + _refresh_group_tp_replication_factors(worker) + + assert worker._group_tp_replication_factors == (2,) + assert worker._lookup_key_prefixes[0] == ( + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0", + "test-model@tp_rank:1@pcp0@dcp0@pp_rank:0@group:0", + ) + + +def test_lookup_rejects_boundary_missing_one_mamba_shard(): + from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheGroupSpec, + MambaSpec, + ) + + worker = _make_bare_worker(block_size=16) + worker.tp_size = 2 + worker.num_kv_head = 1 + fa = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) + mamba = MambaSpec( + block_size=16, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + worker._kv_cache_groups = [ + KVCacheGroupSpec(["l0"], fa), + KVCacheGroupSpec(["l1"], mamba), + ] + worker.token_dbs = [ + ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=0), block_size=16 + ), + ChunkedTokenDatabase( + KeyMetadata("test-model", 1, 0, 0, 0, group_id=1), block_size=16 + ), + ] + worker.coord = mooncake_store_worker.MooncakeStoreCoordinator( + worker._kv_cache_groups, + scheduler_block_size=16, + hash_block_size=16, + ) + _refresh_group_tp_replication_factors(worker) + + worker.store.batch_is_exist.side_effect = lambda keys: [1] * len(keys) + assert worker.lookup(32, [b"h0", b"h1"]) == 32 + + worker.store.batch_is_exist.side_effect = lambda keys: [ + 0 if "tp_rank:1" in k and "group:1" in k else 1 for k in keys + ] + assert worker.lookup(32, [b"h0", b"h1"]) == 0 + + def test_lookup_requires_all_dcp_rank_namespaces(): worker = _make_bare_worker(block_size=16) worker.tp_size = 4 worker.num_kv_head = 1 worker.dcp_size = 4 - worker._init_lookup_key_prefixes() + _refresh_group_tp_replication_factors(worker) worker.store.batch_is_exist.return_value = [1, 1, 0, 1] assert worker.lookup(16, [b"a0"]) == 0 @@ -1850,7 +2050,7 @@ def test_lookup_checks_all_potential_swa_hit_boundaries(): hash_block_size=8, retention_interval=0, ) - worker._init_lookup_key_prefixes() + _refresh_group_tp_replication_factors(worker) # Candidate order: 3 full-attention chunks, then SWA chunks 3, 7, 11. # Only the first full chunk and the SWA chunk ending at token 32 exist, so # lookup should recover a 32-token external prefix hit. A sparse @@ -1911,7 +2111,7 @@ def test_lookup_applies_swa_mask_before_accessing_hashes(): hash_block_size=8, retention_interval=0, ) - worker._init_lookup_key_prefixes() + _refresh_group_tp_replication_factors(worker) block_hashes = _RecordingBlockHashes([f"h{i}".encode() for i in range(12)]) accessed_before_rpc: list[int] = [] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 03b2589c07f..cc49a3569d6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -66,7 +66,15 @@ from vllm.v1.core.kv_cache_utils import ( maybe_convert_block_hash, resolve_kv_cache_block_sizes, ) -from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec +from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupSpec, + KVCacheSpec, + MambaSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, + UniformTypeKVCacheSpecs, +) from .metrics import MooncakeStoreConnectorStats @@ -448,7 +456,7 @@ class KVCacheStoreSendingThread(KVTransferThread): token_databases: list[ChunkedTokenDatabase], block_size: int, tp_rank: int, - put_step: int, + group_put_steps: Sequence[int], kv_role: str, ready_event: threading.Event, enable_kv_event: bool = False, @@ -464,7 +472,8 @@ class KVCacheStoreSendingThread(KVTransferThread): name="KVCacheStoreSendingThread", record_operation=record_operation, ) - self.put_step = put_step + # Only ranks with identical group bytes may stripe PUTs (e.g., MLA). + self.group_put_steps = group_put_steps self.coord = coord self.kv_role = kv_role self.stored_requests: defaultdict[str, int] = defaultdict(int) @@ -568,13 +577,14 @@ class KVCacheStoreSendingThread(KVTransferThread): group_indices: list[int] = [] for g_idx, db in enumerate(self.token_databases): # Rotate the stride phase per group to balance load across ranks. - put_step_rank = (self.tp_rank + g_idx) % self.put_step + put_step = self.group_put_steps[g_idx] + put_step_rank = (self.tp_rank + g_idx) % put_step for start, end, block_hash in db.process_tokens( token_len, req_meta.block_hashes, mask_num=save_start, chunk_mask=store_masks[g_idx], - put_step=self.put_step, + put_step=put_step, put_step_rank=put_step_rank, ): starts.append(start) @@ -1002,44 +1012,7 @@ class MooncakeStoreWorker: ) self.num_layers = model_config.get_num_layers(parallel_config) - self.use_mla = False - if ( - hasattr(model_config, "use_mla") - and isinstance(model_config.use_mla, bool) - and model_config.use_mla - ): - self.use_mla = True - - if self.use_mla: - self.num_kv_head = 1 - else: - self.num_kv_head = model_config.get_total_num_kv_heads() - - if self.num_kv_head < self.tp_size and self.dcp_size <= 1: - # Dedup: TP ranks holding the same KV heads stripe PUTs across - # one shared key namespace. DCP splits the TP group, so with - # DCP>1 those ranks have different `@dcpN` namespaces and - # striping would leave keys unwritten (OBJECT_NOT_FOUND on - # GET). PCP is outer to TP (pcp_rank is constant within a TP - # group), so it needs no guard. - self.put_step = self.tp_size // self.num_kv_head - self.head_or_tp_rank = self.tp_rank // self.put_step - else: - self.head_or_tp_rank = self.tp_rank - self.put_step = 1 - - self.metadata = KeyMetadata( - model_name=model_config.model.rstrip("/").split("/")[-1], - tp_rank=self.head_or_tp_rank, - pcp_rank=self.pcp_rank, - dcp_rank=self.dcp_rank, - pp_rank=self.pp_rank, - cache_prefix=str( - vllm_config.kv_transfer_config.kv_connector_extra_config.get( - "cache_prefix", "" - ) - ), - ) + self.num_kv_head = model_config.get_total_num_kv_heads() # Initialize MooncakeDistributedStore with its own TransferEngine store_config = MooncakeStoreConfig.load_from_config() @@ -1156,10 +1129,32 @@ class MooncakeStoreWorker: retention_interval=envs.VLLM_PREFIX_CACHE_RETENTION_INTERVAL, ) # One ChunkedTokenDatabase per group; addresses populated in - # register_kv_caches once the kv-cache layout is known. + # register_kv_caches once the kv-cache layout is known. Each group's + # key namespace is its TP shard id: ranks holding identical bytes + # (MLA / shared GQA KV heads) share a namespace, TP-sharded Mamba + # state gets one namespace per rank. + metadata = KeyMetadata( + model_name=model_config.model.rstrip("/").split("/")[-1], + tp_rank=self.tp_rank, + pcp_rank=self.pcp_rank, + dcp_rank=self.dcp_rank, + pp_rank=self.pp_rank, + cache_prefix=str( + vllm_config.kv_transfer_config.kv_connector_extra_config.get( + "cache_prefix", "" + ) + ), + ) + self._group_tp_replication_factors: tuple[int, ...] = ( + self._compute_group_tp_replication_factors() + ) self.token_dbs: list[ChunkedTokenDatabase] = [ ChunkedTokenDatabase( - dataclasses.replace(self.metadata, group_id=g_idx), + dataclasses.replace( + metadata, + group_id=g_idx, + tp_rank=self.tp_rank // self._group_tp_replication_factors[g_idx], + ), g.kv_cache_spec.block_size, hash_block_size=self.hash_block_size, ) @@ -1167,27 +1162,50 @@ class MooncakeStoreWorker: ] self._init_lookup_key_prefixes() - def _init_lookup_key_prefixes(self) -> None: - """Prepare per-group key prefixes across parallel rank namespaces.""" - # (tp_rank, pcp_rank, dcp_rank, pp_rank) namespaces + def _spec_tp_replication_factor(self, spec: KVCacheSpec) -> int: if self.dcp_size > 1: - # DCP reuses the TP workers and splits each TP group into - # contiguous DCP groups, so dcp_rank == tp_rank % dcp_size. - # Store/load paths do not apply KV-head dedup under DCP - rank_namespaces = tuple( - (tp_rank, pcp_rank, tp_rank % self.dcp_size, pp_rank) + return 1 + inner_specs = ( + tuple(spec.kv_cache_specs.values()) + if isinstance(spec, UniformTypeKVCacheSpecs) + else (spec,) + ) + # Any rank-specific state makes the whole packed value rank-specific. + if any(isinstance(inner, MambaSpec) for inner in inner_specs): + return 1 + # A pure MLA packed value is replicated on every TP rank. + if all( + isinstance(inner, (MLAAttentionSpec, SlidingWindowMLASpec)) + for inner in inner_specs + ): + return self.tp_size + return max(1, self.tp_size // self.num_kv_head) + + def _compute_group_tp_replication_factors(self) -> tuple[int, ...]: + """Return the number of byte-identical TP replicas per cache group. + + DCP and Mamba use 1; MLA uses ``tp_size``; GQA uses + ``tp_size // num_kv_head``. + """ + return tuple( + self._spec_tp_replication_factor(group.kv_cache_spec) + for group in self._kv_cache_groups + ) + + def _init_lookup_key_prefixes(self) -> None: + def rank_namespaces(factor: int) -> tuple[tuple[int, int, int, int], ...]: + if self.dcp_size > 1: + # DCP is a TP subdivision: dcp_rank == tp_rank % dcp_size. + return tuple( + (tp_rank, pcp_rank, tp_rank % self.dcp_size, pp_rank) + for pcp_rank in range(self.pcp_size) + for tp_rank in range(self.tp_size) + for pp_rank in range(self.pp_size) + ) + return tuple( + (shard_rank, pcp_rank, 0, pp_rank) for pcp_rank in range(self.pcp_size) - for tp_rank in range(self.tp_size) - for pp_rank in range(self.pp_size) - ) - else: - # Without DCP, TP ranks that share a KV head write identical KV, so - # lookup only needs one TP namespace per unique KV head. - tp_count = min(self.tp_size, self.num_kv_head) - rank_namespaces = tuple( - (tp_rank, pcp_rank, 0, pp_rank) - for pcp_rank in range(self.pcp_size) - for tp_rank in range(tp_count) + for shard_rank in range(self.tp_size // factor) for pp_rank in range(self.pp_size) ) @@ -1200,11 +1218,12 @@ class MooncakeStoreWorker: dcp_rank=dcp_rank, pp_rank=pp_rank, ) - for tp_rank, pcp_rank, dcp_rank, pp_rank in rank_namespaces + for tp_rank, pcp_rank, dcp_rank, pp_rank in rank_namespaces( + self._group_tp_replication_factors[g_idx] + ) ) - for db in self.token_dbs + for g_idx, db in enumerate(self.token_dbs) ) - self._lookup_expected_per_key = len(rank_namespaces) def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None: """Register a cross-layers KV cache tensor. @@ -1297,7 +1316,7 @@ class MooncakeStoreWorker: self.token_dbs, self.block_size, self.tp_rank, - self.put_step, + self._group_tp_replication_factors, self.kv_role, ready_event_sending, self.enable_kv_events, @@ -1527,16 +1546,16 @@ class MooncakeStoreWorker: logger.error("Remote connection failed in lookup: %s", e) return 0 - # A (group, hash) is "present" only when every TP*PP rank has it. - ranks_per_candidate = self._lookup_expected_per_key - exists_set = { - (g_idx, hash_bytes) - for i, (g_idx, hash_bytes) in enumerate(candidate_meta) - if all( - res[i * ranks_per_candidate + j] == 1 - for j in range(ranks_per_candidate) - ) - } + # A (group, hash) is "present" only when every namespace that will be + # loaded has it (per-group count: sharded groups need every rank's + # shard, replicated groups one namespace per unique KV head). + exists_set = set() + pos = 0 + for g_idx, hash_bytes in candidate_meta: + count = len(self._lookup_key_prefixes[g_idx]) + if all(res[pos + j] == 1 for j in range(count)): + exists_set.add((g_idx, hash_bytes)) + pos += count cached_block_pool = ExternalCachedBlockPool( self.hash_block_size, From d9cd77419873b4765c7123d2ba003dc96219232d Mon Sep 17 00:00:00 2001 From: Aarushi Jain <142941703+aarushjain29@users.noreply.github.com> Date: Fri, 24 Jul 2026 23:14:37 -0500 Subject: [PATCH 025/185] [ROCm][CI] Force native compile caches onto local disk (#49763) Signed-off-by: aarushjain29 Co-authored-by: Andreas Karatzas --- .buildkite/scripts/hardware_ci/run-amd-test.sh | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index de3f443ad59..0f6bb2564c9 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -400,10 +400,10 @@ initialize_native_environment() { native_root="/tmp/vllm-native-${job_id}" TMPDIR="/tmp/vllm-${job_id_suffix}/tmp" VLLM_RPC_BASE_PATH="/tmp" - : "${TORCHINDUCTOR_CACHE_DIR:=${native_root}/cache/torchinductor}" - : "${TRITON_CACHE_DIR:=${native_root}/cache/triton}" - : "${VLLM_CACHE_ROOT:=${native_root}/cache/vllm}" - : "${XDG_CACHE_HOME:=${native_root}/cache/xdg}" + TORCHINDUCTOR_CACHE_DIR="${native_root}/cache/torchinductor" + TRITON_CACHE_DIR="${native_root}/cache/triton" + VLLM_CACHE_ROOT="${native_root}/cache/vllm" + XDG_CACHE_HOME="${native_root}/cache/xdg" : "${HF_HOME:=/home/buildkite-agent/huggingface}" : "${HF_HUB_DOWNLOAD_TIMEOUT:=300}" : "${HF_HUB_ETAG_TIMEOUT:=60}" @@ -419,6 +419,8 @@ initialize_native_environment() { "${XDG_CACHE_HOME}" \ "${HF_HOME}" || return 1 + echo "Native compile caches: VLLM_CACHE_ROOT=${VLLM_CACHE_ROOT} TORCHINDUCTOR_CACHE_DIR=${TORCHINDUCTOR_CACHE_DIR}" + if [[ "${VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE:-0}" == "1" ]]; then if ! command -v findmnt >/dev/null 2>&1; then echo "findmnt is required to verify the native Hugging Face cache mount" >&2 From aaaeda98dcc14af75cd5fbd26386751a0e8b79c9 Mon Sep 17 00:00:00 2001 From: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Date: Sat, 25 Jul 2026 00:15:19 -0400 Subject: [PATCH 026/185] [CI] fix compile test | refactor VLLM_DISABLE_COMPILE_CACHE for tests (#49770) Signed-off-by: Divakar Verma --- tests/compile/passes/test_fusion_attn.py | 5 +---- .../compile/passes/test_mla_attn_quant_fusion.py | 5 +---- tests/compile/test_compile_ranges.py | 11 +++++------ tests/conftest.py | 16 ++++++++++++++++ 4 files changed, 23 insertions(+), 14 deletions(-) diff --git a/tests/compile/passes/test_fusion_attn.py b/tests/compile/passes/test_fusion_attn.py index 76536f3387c..da335b91557 100644 --- a/tests/compile/passes/test_fusion_attn.py +++ b/tests/compile/passes/test_fusion_attn.py @@ -290,12 +290,9 @@ def test_attention_quant_pattern( model_class: type[AttentionQuantPatternModel], backend: AttentionBackendEnum, dist_init, - monkeypatch, - use_fresh_inductor_cache, + disable_vllm_compile_cache, ): """Test AttentionStaticQuantPattern fusion pass""" - monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1") - if backend == AttentionBackendEnum.FLASHINFER and ( not current_platform.is_device_capability((10, 0)) or not has_flashinfer() ): diff --git a/tests/compile/passes/test_mla_attn_quant_fusion.py b/tests/compile/passes/test_mla_attn_quant_fusion.py index 0a38ffca483..5b2ef5bfd95 100644 --- a/tests/compile/passes/test_mla_attn_quant_fusion.py +++ b/tests/compile/passes/test_mla_attn_quant_fusion.py @@ -419,8 +419,7 @@ def test_mla_attention_quant_pattern( model_class: type[MLAAttentionQuantPatternModel], backend: AttentionBackendEnum, dist_init, - monkeypatch, - use_fresh_inductor_cache, + disable_vllm_compile_cache, ): """Test MLA AttentionQuantPattern fusion pass""" if ( @@ -429,8 +428,6 @@ def test_mla_attention_quant_pattern( ): pytest.skip("NVFP4 is not supported on this GPU (requires SM 100+).") - monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1") - custom_ops_list = custom_ops.split(",") if custom_ops else [] device = torch.device(f"{DEVICE_TYPE}:0") diff --git a/tests/compile/test_compile_ranges.py b/tests/compile/test_compile_ranges.py index 9fd8e9577ba..4dfea42a6b4 100644 --- a/tests/compile/test_compile_ranges.py +++ b/tests/compile/test_compile_ranges.py @@ -66,7 +66,7 @@ class PostGradRangeChecker(InductorPass): return InductorPass.hash_dict(state) -def test_compile_ranges(use_fresh_inductor_cache): +def test_compile_ranges(disable_vllm_compile_cache): post_grad_range_checker = PostGradRangeChecker( [ Range(start=1, end=8), @@ -168,7 +168,7 @@ class PostGradStaticShapeChecker(InductorPass): return InductorPass.hash_dict(state) -def test_compile_sizes_produce_static_shapes(use_fresh_inductor_cache): +def test_compile_sizes_produce_static_shapes(disable_vllm_compile_cache): """Verify that compile_sizes entries are compiled with fully concrete shapes (no SymInts), while compile_ranges entries retain dynamic shapes.""" checker = PostGradStaticShapeChecker() @@ -209,10 +209,9 @@ def test_compile_sizes_produce_static_shapes(use_fresh_inductor_cache): ) -def test_inductor_cache_compile_ranges(monkeypatch, use_fresh_inductor_cache): - # To force multiple compilations, we disable the compile cache - monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1") - +def test_inductor_cache_compile_ranges(disable_vllm_compile_cache): + # disable_vllm_compile_cache sets VLLM_DISABLE_COMPILE_CACHE=1 to force + # multiple compilations by disabling vLLM's on-disk compile cache. post_grad_range_checker = PostGradRangeChecker( ranges=[ Range(start=1, end=8), diff --git a/tests/conftest.py b/tests/conftest.py index 47071167b56..406cb2ed2ea 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1768,6 +1768,22 @@ def use_fresh_inductor_cache(): yield +@pytest.fixture +def disable_vllm_compile_cache(monkeypatch, use_fresh_inductor_cache): + """ + Use a fresh inductor cache AND disable vLLM's on-disk torch.compile cache. + + This forces compilation (and any custom compile passes) to actually run + instead of being served from a warm cache left behind by previous runs + (e.g. on persistent CI agents). Use this for tests that inspect what + happens during compilation; use ``use_fresh_inductor_cache`` (or + ``fresh_vllm_cache``) instead when the vLLM compile cache must stay + enabled (e.g. cache save/load tests). + """ + monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1") + yield + + @pytest.fixture def fresh_vllm_cache(monkeypatch, use_fresh_inductor_cache): """Temporary VLLM_CACHE_ROOT combined with a fresh inductor cache.""" From 0ba2aa35a81dcc3246b26291368b53fa2389c7d7 Mon Sep 17 00:00:00 2001 From: Aarushi Jain <142941703+aarushjain29@users.noreply.github.com> Date: Fri, 24 Jul 2026 23:21:36 -0500 Subject: [PATCH 027/185] Stabilize GPU memory teardown between ROCm CI tests (#49242) Signed-off-by: aarushjain29 Signed-off-by: Andreas Karatzas Co-authored-by: Andreas Karatzas Co-authored-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- tests/conftest.py | 35 +++- tests/lora/test_qwenvl.py | 168 +++++++++--------- .../nixl_integration/run_accuracy_test.sh | 16 +- 3 files changed, 126 insertions(+), 93 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 406cb2ed2ea..db097fbc455 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1728,33 +1728,50 @@ def disable_deepgemm_ue8m0(monkeypatch): is_deep_gemm_e8m0_used.cache_clear() +def _should_clean_gpu_memory_between_tests() -> bool: + setting = os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY") + if setting == "1": + return True + if setting == "0": + return False + # ROCm reclaims VRAM lazily; default to waiting between tests on ROCm CI. + return current_platform.is_rocm() + + @pytest.fixture(autouse=True) def clean_gpu_memory_between_tests(): - if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1": + if not _should_clean_gpu_memory_between_tests(): yield return - # Wait for GPU memory to be cleared before starting the test import gc - from tests.utils import wait_for_gpu_memory_to_clear + from tests.utils import wait_for_gpu_memory_to_clear, wait_for_rocm_memory_to_settle num_gpus = torch.accelerator.device_count() - if num_gpus > 0: + + def _wait_for_settled_gpu_memory() -> None: + if num_gpus <= 0: + return try: - wait_for_gpu_memory_to_clear( - devices=list(range(num_gpus)), - threshold_ratio=0.1, - ) + if current_platform.is_rocm(): + wait_for_rocm_memory_to_settle() + else: + wait_for_gpu_memory_to_clear( + devices=list(range(num_gpus)), + threshold_ratio=0.1, + ) except ValueError as e: logger.info("Failed to clean GPU memory: %s", e) + _wait_for_settled_gpu_memory() + yield - # Clean up GPU memory after the test if torch.cuda.is_available(): torch.accelerator.empty_cache() gc.collect() + _wait_for_settled_gpu_memory() @pytest.fixture diff --git a/tests/lora/test_qwenvl.py b/tests/lora/test_qwenvl.py index a4a32278db0..3cbb534bdad 100644 --- a/tests/lora/test_qwenvl.py +++ b/tests/lora/test_qwenvl.py @@ -7,6 +7,7 @@ from packaging.version import Version from transformers import __version__ as TRANSFORMERS_VERSION import vllm +from tests.conftest import VllmRunner from vllm.assets.image import ImageAsset from vllm.lora.request import LoRARequest from vllm.platforms import current_platform @@ -56,24 +57,29 @@ class Qwen2VLTester: def __init__(self, config: TestConfig): self.config = config - self.llm = self._initialize_llm() - - def _initialize_llm(self) -> vllm.LLM: - """Initialize the LLM with given configuration""" - return vllm.LLM( - model=self.config.model_path, - max_num_seqs=self.config.max_num_seqs, + self._runner = VllmRunner( + model_name=config.model_path, + max_num_seqs=config.max_num_seqs, enable_lora=True, - max_loras=self.config.max_loras, - max_lora_rank=self.config.max_lora_rank, - enable_tower_connector_lora=self.config.enable_tower_connector_lora, - trust_remote_code=True, - gpu_memory_utilization=self.config.gpu_memory_utilization, - mm_processor_kwargs=self.config.mm_processor_kwargs, - mm_processor_cache_gb=self.config.mm_processor_cache_gb, - max_model_len=self.config.max_model_len, + max_loras=config.max_loras, + max_lora_rank=config.max_lora_rank, + enable_tower_connector_lora=config.enable_tower_connector_lora, + gpu_memory_utilization=config.gpu_memory_utilization, + mm_processor_kwargs=config.mm_processor_kwargs, + mm_processor_cache_gb=config.mm_processor_cache_gb, + max_model_len=config.max_model_len, ) + @property + def llm(self) -> vllm.LLM: + return self._runner.get_llm() + + def __enter__(self) -> "Qwen2VLTester": + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + self._runner.__exit__(exc_type, exc_value, traceback) + def run_test( self, images: list[ImageAsset], @@ -183,29 +189,29 @@ QWEN3VL_MODEL_PATH = "Qwen/Qwen3-VL-4B-Instruct" def test_qwen2vl_lora(qwen2vl_lora_files): """Test Qwen 2.0 VL model with LoRA""" config = TestConfig(model_path=QWEN2VL_MODEL_PATH, lora_path=qwen2vl_lora_files) - tester = Qwen2VLTester(config) - - # Test with different LoRA IDs - for lora_id in [1, 2]: - tester.run_test(TEST_IMAGES, expected_outputs=EXPECTED_OUTPUTS, lora_id=lora_id) + with Qwen2VLTester(config) as tester: + # Test with different LoRA IDs + for lora_id in [1, 2]: + tester.run_test( + TEST_IMAGES, expected_outputs=EXPECTED_OUTPUTS, lora_id=lora_id + ) def test_qwen2vl_lora_beam_search(qwen2vl_lora_files): """Test Qwen 2.0 VL model with LoRA through beam search.""" config = TestConfig(model_path=QWEN2VL_MODEL_PATH, lora_path=qwen2vl_lora_files) - tester = Qwen2VLTester(config) - - # Test with different LoRA IDs - for lora_id in [1, 2]: - # NOTE currently, we only test cherry blossom since stop sign - # output is slightly different for v1; - the root cause is likely - # independent of the intent of this test, which is to ensure beam - # search passes through lora through correctly. - tester.run_beam_search_test( - [ImageAsset("cherry_blossom")], - expected_outputs=EXPECTED_BEAM_SEARCH_OUTPUTS, - lora_id=lora_id, - ) + with Qwen2VLTester(config) as tester: + # Test with different LoRA IDs + for lora_id in [1, 2]: + # NOTE currently, we only test cherry blossom since stop sign + # output is slightly different for v1; - the root cause is likely + # independent of the intent of this test, which is to ensure beam + # search passes through lora through correctly. + tester.run_beam_search_test( + [ImageAsset("cherry_blossom")], + expected_outputs=EXPECTED_BEAM_SEARCH_OUTPUTS, + lora_id=lora_id, + ) @pytest.mark.skipif( @@ -214,11 +220,12 @@ def test_qwen2vl_lora_beam_search(qwen2vl_lora_files): def test_qwen25vl_lora(qwen25vl_lora_files): """Test Qwen 2.5 VL model with LoRA""" config = TestConfig(model_path=QWEN25VL_MODEL_PATH, lora_path=qwen25vl_lora_files) - tester = Qwen2VLTester(config) - - # Test with different LoRA IDs - for lora_id in [1, 2]: - tester.run_test(TEST_IMAGES, expected_outputs=EXPECTED_OUTPUTS, lora_id=lora_id) + with Qwen2VLTester(config) as tester: + # Test with different LoRA IDs + for lora_id in [1, 2]: + tester.run_test( + TEST_IMAGES, expected_outputs=EXPECTED_OUTPUTS, lora_id=lora_id + ) @pytest.mark.skipif( @@ -234,13 +241,13 @@ def test_qwen25vl_vision_lora(qwen25vl_vision_lora_files): mm_processor_cache_gb=0, enable_tower_connector_lora=True, ) - tester = Qwen2VLTester(config) - for lora_id in [1, 2]: - tester.run_test( - TEST_IMAGES, - expected_outputs=EXPECTED_OUTPUTS, - lora_id=lora_id, - ) + with Qwen2VLTester(config) as tester: + for lora_id in [1, 2]: + tester.run_test( + TEST_IMAGES, + expected_outputs=EXPECTED_OUTPUTS, + lora_id=lora_id, + ) def test_qwen3vl_vision_lora(qwen3vl_vision_lora_files): @@ -253,13 +260,13 @@ def test_qwen3vl_vision_lora(qwen3vl_vision_lora_files): mm_processor_cache_gb=0, enable_tower_connector_lora=True, ) - tester = Qwen2VLTester(config) - for lora_id in [1, 2]: - tester.run_test( - TEST_IMAGES, - expected_outputs=EXPECTED_OUTPUTS, - lora_id=lora_id, - ) + with Qwen2VLTester(config) as tester: + for lora_id in [1, 2]: + tester.run_test( + TEST_IMAGES, + expected_outputs=EXPECTED_OUTPUTS, + lora_id=lora_id, + ) def test_qwen2vl_multiple_lora_types( @@ -287,34 +294,33 @@ def test_qwen2vl_multiple_lora_types( mm_processor_cache_gb=0, enable_tower_connector_lora=True, ) - tester = Qwen2VLTester(config) + with Qwen2VLTester(config) as tester: + # Test 1: Language-only LoRA adapter + tester.config.lora_path = qwen2vl_language_lora_files + for lora_id in [1, 2]: + tester.run_test( + TEST_IMAGES, + expected_outputs=EXPECTED_OUTPUTS_LANGUAGE, + lora_id=lora_id, + lora_name="language_only", + ) - # Test 1: Language-only LoRA adapter - tester.config.lora_path = qwen2vl_language_lora_files - for lora_id in [1, 2]: - tester.run_test( - TEST_IMAGES, - expected_outputs=EXPECTED_OUTPUTS_LANGUAGE, - lora_id=lora_id, - lora_name="language_only", - ) + # Test 2: Vision tower + connector LoRA adapter + tester.config.lora_path = qwen2vl_vision_tower_connector_lora_files + for lora_id in [3, 4]: + tester.run_test( + TEST_IMAGES, + expected_outputs=EXPECTED_OUTPUTS_VISION, + lora_id=lora_id, + lora_name="vision_tower_connector", + ) - # Test 2: Vision tower + connector LoRA adapter - tester.config.lora_path = qwen2vl_vision_tower_connector_lora_files - for lora_id in [3, 4]: - tester.run_test( - TEST_IMAGES, - expected_outputs=EXPECTED_OUTPUTS_VISION, - lora_id=lora_id, - lora_name="vision_tower_connector", - ) - - # Test 3: Vision tower only LoRA adapter (no connector) - tester.config.lora_path = qwen2vl_vision_tower_lora_files - for lora_id in [5, 6]: - tester.run_test( - TEST_IMAGES, - expected_outputs=EXPECTED_OUTPUTS_VISION_NO_CONNECTOR, - lora_id=lora_id, - lora_name="vision_tower", - ) + # Test 3: Vision tower only LoRA adapter (no connector) + tester.config.lora_path = qwen2vl_vision_tower_lora_files + for lora_id in [5, 6]: + tester.run_test( + TEST_IMAGES, + expected_outputs=EXPECTED_OUTPUTS_VISION_NO_CONNECTOR, + lora_id=lora_id, + lora_name="vision_tower", + ) diff --git a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh index 0a45b6f54b8..22682e02bea 100755 --- a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh @@ -98,7 +98,7 @@ GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "") # Trap the SIGINT signal (triggered by Ctrl+C) -trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT +trap 'kill $(jobs -pr) 2>/dev/null || true' SIGINT SIGTERM EXIT # Waits for vLLM to start. wait_for_server() { @@ -110,10 +110,20 @@ wait_for_server() { } # Function to clean up previous instances +wait_for_gpu_memory_release() { + if [[ "$SMI_BIN" == *"rocm"* ]]; then + PYTHONPATH="${GIT_ROOT}" python3 -c "from tests.utils import wait_for_rocm_memory_to_settle; wait_for_rocm_memory_to_settle()" + fi +} + cleanup_instances() { echo "Cleaning up any running vLLM instances..." - pkill -f "vllm serve" || true + pkill -f "toy_proxy_server.py" || true + pkill -TERM -f "vllm serve" || true + sleep 3 + pkill -9 -f "vllm serve" || true sleep 2 + wait_for_gpu_memory_release } get_num_gpus() { @@ -132,6 +142,7 @@ get_num_gpus() { # Function to run tests for a specific model run_tests_for_model() { local model_name=$1 + cleanup_instances echo "================================" echo "Testing model: $model_name" echo "================================" @@ -303,7 +314,6 @@ run_tests_for_model() { # Clean up before running next model cleanup_instances - sleep 3 } # Run tests for each model From 94682b79f40239369bcefdbebee683a6bed90c85 Mon Sep 17 00:00:00 2001 From: Brandon Pelfrey Date: Fri, 24 Jul 2026 23:31:41 -0700 Subject: [PATCH 028/185] [multimodal] Make PyNvVideoCodec decoder concurrency configurable (#49753) Signed-off-by: Brandon Pelfrey Co-authored-by: OpenAI Codex --- docs/features/multimodal_inputs.md | 25 +++++++++-- tests/multimodal/media/test_video.py | 17 ++++++++ tests/multimodal/test_gpu_ipc_memory.py | 35 +++++++++++++-- tests/multimodal/test_video.py | 38 +++++++++++++++-- vllm/multimodal/gpu_ipc_memory.py | 57 +++++++++++++++++++++---- vllm/multimodal/media/video.py | 4 ++ vllm/multimodal/video.py | 36 ++++++++++++++-- 7 files changed, 189 insertions(+), 23 deletions(-) diff --git a/docs/features/multimodal_inputs.md b/docs/features/multimodal_inputs.md index df33cea0542..7f6c1aee760 100644 --- a/docs/features/multimodal_inputs.md +++ b/docs/features/multimodal_inputs.md @@ -818,16 +818,18 @@ Full example: [examples/generate/multimodal/openai_chat_completion_client_for_mu #### Video Decoding Backend -vLLM decodes video bytes into frames using a selectable decoding backend. Three +vLLM decodes video bytes into frames using a selectable decoding backend. Five backends are supported: - `opencv` (default): OpenCV-based decoder. - `pyav`: PyAV decoder. - `torchcodec`: TorchCodec (PyTorch-native) decoder. +- `pynvvideocodec`: NVIDIA NVDEC-based decoder. +- `deepstream`: NVIDIA DeepStream NVDEC-based decoder. -All three backends are ultimately backed by FFmpeg. `torchcodec` lets -you choose which FFmpeg version is used while `opencv` and `pyav` rely on -whichever FFmpeg build they were linked against. +The CPU backends are backed by FFmpeg. `torchcodec` lets you choose which FFmpeg +version is used while `opencv` and `pyav` rely on whichever FFmpeg build they +were linked against. Select the backend by passing the `backend` parameter via `--media-io-kwargs`: @@ -854,6 +856,21 @@ vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \ --media-io-kwargs '{"video": {"backend": "torchcodec", "seek_mode": "approximate", "num_ffmpeg_threads": 4}}' ``` +**PyNvVideoCodec-specific parameters:** + +- `hw_decoders`: Maximum number of concurrent hardware decoder slots retained + by each API server process. It must be a positive integer and defaults to `2`, + which is the recommended starting point for concurrent video workloads. + Because vLLM reserves GPU memory for these slots at startup, this value cannot + be overridden per request. Benchmark before increasing it because each + additional slot increases the GPU memory reservation. + +```bash +# Example: explicitly use the recommended 2 hardware decoders +vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \ + --media-io-kwargs '{"video": {"backend": "pynvvideocodec", "hw_decoders": 2}}' +``` + #### Video Frame Recovery For improved robustness when processing potentially corrupted or truncated video files, vLLM supports optional frame recovery using a dynamic window forward-scan approach. When enabled, if a target frame fails to load during sequential reading, the next successfully grabbed frame (before the next target frame) will be used in its place. diff --git a/tests/multimodal/media/test_video.py b/tests/multimodal/media/test_video.py index 671abd7b077..a5595121b2e 100644 --- a/tests/multimodal/media/test_video.py +++ b/tests/multimodal/media/test_video.py @@ -404,6 +404,23 @@ class TestMergeKwargsGpuBackendPolicy: ) assert result["backend"] == "pynvvideocodec" + def test_strips_request_level_hw_decoders_when_not_static(self): + result = VideoMediaIO.merge_kwargs( + default_kwargs={"video_backend": "pynvvideocodec"}, + runtime_kwargs={"hw_decoders": 4}, + ) + assert "hw_decoders" not in result + + def test_prevents_request_level_hw_decoders_override(self): + result = VideoMediaIO.merge_kwargs( + default_kwargs={ + "video_backend": "pynvvideocodec", + "hw_decoders": 2, + }, + runtime_kwargs={"hw_decoders": 4}, + ) + assert result["hw_decoders"] == 2 + @pytest.mark.parametrize("backend", ["opencv", "pyav", "torchcodec"]) def test_software_video_backend_passes_through(self, backend: str): result = VideoMediaIO.merge_kwargs( diff --git a/tests/multimodal/test_gpu_ipc_memory.py b/tests/multimodal/test_gpu_ipc_memory.py index 5c0f263af90..c98fd7673cb 100644 --- a/tests/multimodal/test_gpu_ipc_memory.py +++ b/tests/multimodal/test_gpu_ipc_memory.py @@ -18,7 +18,6 @@ from vllm.multimodal.gpu_ipc_memory import ( from vllm.multimodal.video import ( PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES, PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES, - PYNVVIDEOCODEC_MAX_RETAINED_DECODERS, PYNVVIDEOCODEC_VIDEO_BACKEND, ) from vllm.utils.mem_constants import GiB_bytes @@ -28,17 +27,26 @@ def _mm_config( *, mm_ipc_gpu_memory_gb: float = 0, video_backend: str | None = None, + hw_decoders: int | None = None, ) -> MultiModalConfig: - video_kwargs = {} if video_backend is None else {"video_backend": video_backend} + video_kwargs: dict[str, object] = ( + {} if video_backend is None else {"video_backend": video_backend} + ) + if hw_decoders is not None: + video_kwargs["hw_decoders"] = hw_decoders + return MultiModalConfig( mm_ipc_gpu_memory_gb=mm_ipc_gpu_memory_gb, media_io_kwargs={"video": video_kwargs} if video_kwargs else {}, ) -def _pynvvideocodec_decoder_budget(api_process_count: int = 1) -> int: +def _pynvvideocodec_decoder_budget( + api_process_count: int = 1, + hw_decoders: int = 2, +) -> int: return api_process_count * ( - PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES * PYNVVIDEOCODEC_MAX_RETAINED_DECODERS + PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES * hw_decoders + PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES ) @@ -240,3 +248,22 @@ def test_reserve_mm_ipc_gpu_memory_scales_decoder_budget_by_api_servers( _mm_config(), api_process_count=3, ) == available_bytes - _pynvvideocodec_decoder_budget(api_process_count=3) + + +def test_reserve_mm_ipc_gpu_memory_uses_configured_hw_decoders( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr( + multimodal_config_module.envs, + "VLLM_VIDEO_LOADER_BACKEND", + "opencv", + ) + available_bytes = 4 * GiB_bytes + mm_config = _mm_config( + video_backend=PYNVVIDEOCODEC_VIDEO_BACKEND, + hw_decoders=3, + ) + + assert reserve_mm_ipc_gpu_memory(available_bytes, mm_config) == ( + available_bytes - _pynvvideocodec_decoder_budget(hw_decoders=3) + ) diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index e10d6e1338e..f07a6d888e0 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -16,7 +16,6 @@ from transformers.video_utils import VideoMetadata from vllm.assets.base import get_vllm_public_assets from vllm.multimodal.video import ( PYNVVIDEOCODEC_DECODER_CACHE_SIZE, - PYNVVIDEOCODEC_MAX_RETAINED_DECODERS, PYNVVIDEOCODEC_VIDEO_BACKEND, VIDEO_LOADER_REGISTRY, DynamicVideoBackend, @@ -26,6 +25,7 @@ from vllm.multimodal.video import ( PyNvVideoCodecVideoBackend, Qwen2VLVideoBackend, Qwen3VLVideoBackend, + VideoBackend, VideoLoader, VideoSourceMetadata, VideoTargetMetadata, @@ -199,7 +199,11 @@ def test_pynvvideocodec_codec_uses_dynamic_sampling_strategy( assert metadata["frames_indices"] == [0, 9] -def test_pynvvideocodec_decoder_slots_are_bounded(monkeypatch: pytest.MonkeyPatch): +@pytest.mark.parametrize("hw_decoders", [1, 3]) +def test_pynvvideocodec_decoder_slots_are_bounded( + monkeypatch: pytest.MonkeyPatch, + hw_decoders: int, +): class FakeSlot: pass @@ -207,10 +211,13 @@ def test_pynvvideocodec_decoder_slots_are_bounded(monkeypatch: pytest.MonkeyPatc old_slots = PyNvVideoCodecVideoBackend._decoder_slots old_active_slots = PyNvVideoCodecVideoBackend._active_decoder_slots old_cond = PyNvVideoCodecVideoBackend._decoder_slot_cond + old_max_slots = PyNvVideoCodecVideoBackend._max_decoder_slots try: PyNvVideoCodecVideoBackend._decoder_slots = [] PyNvVideoCodecVideoBackend._active_decoder_slots = 0 PyNvVideoCodecVideoBackend._decoder_slot_cond = threading.Condition() + PyNvVideoCodecVideoBackend._max_decoder_slots = None + PyNvVideoCodecVideoBackend._configure_decoder_slots(hw_decoders) def fake_create_slot(cls): nonlocal create_count @@ -229,7 +236,7 @@ def test_pynvvideocodec_decoder_slots_are_bounded(monkeypatch: pytest.MonkeyPatc with ExitStack() as stack: retained_slots = [ stack.enter_context(PyNvVideoCodecVideoBackend._borrow_decoder_slot()) - for _ in range(PYNVVIDEOCODEC_MAX_RETAINED_DECODERS) + for _ in range(hw_decoders) ] def borrow_extra_slot(): @@ -246,11 +253,34 @@ def test_pynvvideocodec_decoder_slots_are_bounded(monkeypatch: pytest.MonkeyPatc assert not thread.is_alive() assert seen_slots[0] in retained_slots - assert create_count == PYNVVIDEOCODEC_MAX_RETAINED_DECODERS + assert create_count == hw_decoders finally: PyNvVideoCodecVideoBackend._decoder_slots = old_slots PyNvVideoCodecVideoBackend._active_decoder_slots = old_active_slots PyNvVideoCodecVideoBackend._decoder_slot_cond = old_cond + PyNvVideoCodecVideoBackend._max_decoder_slots = old_max_slots + + +def test_pynvvideocodec_decoder_slots_are_configured_once( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(PyNvVideoCodecVideoBackend, "_max_decoder_slots", None) + + PyNvVideoCodecVideoBackend._configure_decoder_slots(2) + PyNvVideoCodecVideoBackend._configure_decoder_slots(2) + + with pytest.raises(RuntimeError, match="already configured as 2, got 3"): + PyNvVideoCodecVideoBackend._configure_decoder_slots(3) + + +@pytest.mark.parametrize("hw_decoders", [0, -1, 1.5, True, "2"]) +def test_pynvvideocodec_rejects_invalid_hw_decoders(hw_decoders: object): + with pytest.raises(ValueError, match="hw_decoders must be a positive integer"): + VideoBackend.load_bytes( + b"fake video", + backend=PYNVVIDEOCODEC_VIDEO_BACKEND, + hw_decoders=hw_decoders, # type: ignore[arg-type] + ) def test_pynvvideocodec_decoder_slot_retains_simple_decoder(): diff --git a/vllm/multimodal/gpu_ipc_memory.py b/vllm/multimodal/gpu_ipc_memory.py index 02317490a84..58cb9888802 100644 --- a/vllm/multimodal/gpu_ipc_memory.py +++ b/vllm/multimodal/gpu_ipc_memory.py @@ -157,19 +157,43 @@ def reserve_mm_ipc_gpu_memory( mm_config: "MultiModalConfig | None", api_process_count: int = 1, ) -> int: - """Carve frontend multimodal GPU memory out of the KV cache. + """Return KV-cache memory remaining after frontend multimodal reservations. - Raw decoded frames are bounded by ``mm_ipc_gpu_memory_gb`` and acquired by - the frontend semaphore. Some decoders also keep persistent surfaces around; - reserve a fixed upper bound for those when a GPU backend is configured. + The reservation covers: + + * The total ``mm_ipc_gpu_memory_gb`` budget for transient decoded-frame + buffers. This budget is divided among API processes, so it is not + multiplied by ``api_process_count``. + * For GPU video backends, a fixed upper bound for each API process's + retained decoder surfaces and CUDA context. The PyNvVideoCodec surface + reservation scales with its configured ``hw_decoders`` value, and the + entire decoder reservation scales with ``api_process_count`` because + these resources are not shared between processes. + + Args: + available_kv_cache_memory_bytes: KV-cache capacity before reserving + memory for frontend multimodal processing. + mm_config: Multimodal configuration, or ``None`` when multimodal + processing is disabled. + api_process_count: Number of frontend API processes sharing the GPU. + Values below one are treated as one. + + Returns: + KV-cache capacity after subtracting the frontend reservation. + + Raises: + ValueError: If the reservation leaves no memory for the KV cache. """ if mm_config is None: return available_kv_cache_memory_bytes + from vllm import envs from vllm.multimodal.video import ( PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES, PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES, - PYNVVIDEOCODEC_MAX_RETAINED_DECODERS, + PYNVVIDEOCODEC_DEFAULT_HW_DECODERS, + PYNVVIDEOCODEC_VIDEO_BACKEND, + validate_pynvvideocodec_hw_decoders, ) raw_frame_reserved_bytes = int(mm_config.mm_ipc_gpu_memory_gb * GiB_bytes) @@ -177,8 +201,24 @@ def reserve_mm_ipc_gpu_memory( # context on the GPU, outside the worker memory pool. Reserve that footprint # per process so gpu_memory_utilization bounds total GPU usage across them. num_api_servers = max(1, api_process_count) + video_kwargs = mm_config.media_io_kwargs.get("video", {}) + video_loader_backend = ( + video_kwargs.get("video_backend") or envs.VLLM_VIDEO_LOADER_BACKEND + ) + codec_backend = video_kwargs.get("backend") + uses_pynvvideocodec = ( + video_loader_backend == PYNVVIDEOCODEC_VIDEO_BACKEND + or codec_backend == PYNVVIDEOCODEC_VIDEO_BACKEND + ) + hw_decoders = ( + validate_pynvvideocodec_hw_decoders( + video_kwargs.get("hw_decoders", PYNVVIDEOCODEC_DEFAULT_HW_DECODERS) + ) + if uses_pynvvideocodec + else 1 + ) per_server_decoder_bytes = ( - PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES * PYNVVIDEOCODEC_MAX_RETAINED_DECODERS + PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES * hw_decoders + PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES ) decoder_reserved_bytes = ( @@ -198,8 +238,9 @@ def reserve_mm_ipc_gpu_memory( f"({format_gib(raw_frame_reserved_bytes)} GiB raw-frame budget, " f"{format_gib(decoder_reserved_bytes)} GiB decoder cache budget), " f"but only {format_gib(available_kv_cache_memory_bytes)} GiB is " - "available for the KV cache. Reduce mm_ipc_gpu_memory_gb, use a " - "different video backend, or increase gpu_memory_utilization." + "available for the KV cache. Reduce mm_ipc_gpu_memory_gb or " + "hw_decoders, use a different video backend, or increase " + "gpu_memory_utilization." ) logger.info_once( "Reserving %s GiB of GPU memory for frontend multimodal decoding " diff --git a/vllm/multimodal/media/video.py b/vllm/multimodal/media/video.py index 45ea4c2fdf4..978cf06c23a 100644 --- a/vllm/multimodal/media/video.py +++ b/vllm/multimodal/media/video.py @@ -32,6 +32,10 @@ class VideoMediaIO(MediaIO[tuple[npt.NDArray, dict[str, Any]]]): runtime_kwargs: dict[str, Any] | None, ) -> dict[str, Any]: if runtime_kwargs: + # Decoder GPU memory is reserved from the startup value. + runtime_kwargs = dict(runtime_kwargs) + runtime_kwargs.pop("hw_decoders", None) + # Block request-level selection of GPU video backends that # were not configured (and VRAM-reserved) at startup. for key in ("video_backend", "backend"): diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index d289972341e..3e0969f1ad9 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -210,15 +210,25 @@ class VideoLoader: VIDEO_LOADER_REGISTRY = VideoLoaderRegistry() PYNVVIDEOCODEC_VIDEO_BACKEND: Literal["pynvvideocodec"] = "pynvvideocodec" -# Fixed upper bound reserved for persistent PyNvVideoCodec decoder surfaces. +# Per-decoder upper bound reserved for persistent PyNvVideoCodec surfaces. PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES = 128 * MiB_bytes PYNVVIDEOCODEC_DECODER_CACHE_SIZE = 2 -PYNVVIDEOCODEC_MAX_RETAINED_DECODERS = 1 +PYNVVIDEOCODEC_DEFAULT_HW_DECODERS = 2 # Per-API-server CUDA context and driver allocation, measured with # PyNvVideoCodec 2.0.4 on H100. PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES = int(1.8 * 1024 * MiB_bytes) +def validate_pynvvideocodec_hw_decoders(hw_decoders: object) -> int: + if ( + isinstance(hw_decoders, bool) + or not isinstance(hw_decoders, int) + or hw_decoders < 1 + ): + raise ValueError("hw_decoders must be a positive integer") + return hw_decoders + + class PyNvVideoCodecDecoderSlot: """A retained PyNv decoder slot and its CUDA stream. @@ -628,6 +638,7 @@ class PyNvVideoCodecVideoBackendMixin: _decoder_slots: ClassVar[list[PyNvVideoCodecDecoderSlot]] = [] _active_decoder_slots: ClassVar[int] = 0 _decoder_slot_cond: ClassVar[threading.Condition] = threading.Condition() + _max_decoder_slots: ClassVar[int | None] = None _DEVICE_INDEX: ClassVar[int] = 0 @classmethod @@ -651,6 +662,18 @@ class PyNvVideoCodecVideoBackendMixin: return PyNvVideoCodecDecoderSlot(torch.cuda.Stream(device=cls._DEVICE_INDEX)) + @classmethod + def _configure_decoder_slots(cls, hw_decoders: object) -> None: + hw_decoders = validate_pynvvideocodec_hw_decoders(hw_decoders) + with cls._decoder_slot_cond: + if cls._max_decoder_slots is None: + cls._max_decoder_slots = hw_decoders + elif cls._max_decoder_slots != hw_decoders: + raise RuntimeError( + "PyNvVideoCodec decoder count is already configured as " + f"{cls._max_decoder_slots}, got {hw_decoders}" + ) + @staticmethod @contextmanager def _torch_stream_context(stream): @@ -669,11 +692,14 @@ class PyNvVideoCodecVideoBackendMixin: def _borrow_decoder_slot(cls): create_slot = False with cls._decoder_slot_cond: + max_decoder_slots = cls._max_decoder_slots + if max_decoder_slots is None: + raise RuntimeError("PyNvVideoCodec decoder slots are not configured") while True: if cls._decoder_slots: slot = cls._decoder_slots.pop() break - if cls._active_decoder_slots < PYNVVIDEOCODEC_MAX_RETAINED_DECODERS: + if cls._active_decoder_slots < max_decoder_slots: cls._active_decoder_slots += 1 create_slot = True break @@ -999,6 +1025,7 @@ class VideoBackend( ] = "opencv", num_ffmpeg_threads: int = 0, seek_mode: Literal["exact", "approximate"] = "exact", + hw_decoders: int = PYNVVIDEOCODEC_DEFAULT_HW_DECODERS, **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: """Load sampled frames from raw video bytes. @@ -1025,6 +1052,8 @@ class VideoBackend( at the cost of relying on the file's metadata. See https://meta-pytorch.org/torchcodec/stable/generated_examples/decoding/approximate_mode.html for details. + hw_decoders: Maximum number of concurrent PyNvVideoCodec decoder + slots. Defaults to 2 and must be a positive integer. Returns: Tuple of ``(frames_array, metadata_dict)``. @@ -1088,6 +1117,7 @@ class VideoBackend( "frame_recovery is not supported for " f"`{PYNVVIDEOCODEC_VIDEO_BACKEND}` backend" ) + cls._configure_decoder_slots(hw_decoders) frames, source, frame_idx, valid = cls.decode_frames_pynvvideocodec( data, target, From 190be7dad2afa6684902324e0dffa2dc0229a364 Mon Sep 17 00:00:00 2001 From: Johnny-Liou <75152852+Johnny-Liou@users.noreply.github.com> Date: Sat, 25 Jul 2026 03:33:43 -0400 Subject: [PATCH 029/185] [Docs] Fix confusing docstring indentation in nemotron_h.py (#49781) Signed-off-by: Johnny-Liou --- vllm/model_executor/models/nemotron_h.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index 50dc5c02221..36d554f21f9 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -771,8 +771,7 @@ class NemotronHForCausalLM( Tuple containing: - conv_state_shape: Shape for convolutional state cache - temporal_state_shape: Shape for state space model cache - - (when use_replayssm is enabled) the x_cache/dt_cache/B_cache - ring-buffer shapes + - x_cache/dt_cache/B_cache ring-buffer shapes (use_replayssm only) """ parallel_config = vllm_config.parallel_config cache_config = vllm_config.cache_config From a82f1b388fe625038502eaa593690ed055fc4dd1 Mon Sep 17 00:00:00 2001 From: Agata Dobrzyniewicz <160237065+adobrzyn@users.noreply.github.com> Date: Sat, 25 Jul 2026 11:06:20 +0200 Subject: [PATCH 030/185] [Perf][V1] Skip LRU hash-split in free_blocks when prefix caching is off (#48017) Signed-off-by: Dobrzyniewicz, Agata Signed-off-by: Agata Dobrzyniewicz <160237065+adobrzyn@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 --- vllm/v1/core/block_pool.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py index b42c7662d2b..a64b7e32159 100644 --- a/vllm/v1/core/block_pool.py +++ b/vllm/v1/core/block_pool.py @@ -724,18 +724,20 @@ class BlockPool: ordered_blocks: A list of blocks to free ordered by their eviction priority. """ - # Identify blocks with hash (LRU cache) and without it (will never match in APC) + # Identify blocks with hash (LRU cache) and without it (never match APC) blocks_with_hash = [] blocks_without_hash = [] for block in ordered_blocks: block.ref_cnt -= 1 if block.ref_cnt == 0 and not block.is_null: - if block.block_hash is None: + # When caching is disabled we always append for better + # GPU cache locality from reusing recently used blocks + if block.block_hash is None and self.enable_caching: blocks_without_hash.append(block) else: blocks_with_hash.append(block) - # Blocks without hash always get evicted first - prepend them last to the tail + # Blocks without hash get evicted first - prepend them last to the tail self.free_block_queue.prepend_n(blocks_without_hash) self.free_block_queue.append_n(blocks_with_hash) From dbcc1cdd0a7b41a75f14e3e432cce52072862a6b Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Sat, 25 Jul 2026 10:50:08 +0100 Subject: [PATCH 031/185] [Model] Remove Ouro (#49786) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) --- docs/models/supported_models.md | 1 - tests/models/registry.py | 1 - vllm/model_executor/models/ouro.py | 448 ------------------------- vllm/model_executor/models/registry.py | 2 +- 4 files changed, 1 insertion(+), 451 deletions(-) delete mode 100644 vllm/model_executor/models/ouro.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index f1292d7a988..3ee942c175e 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -427,7 +427,6 @@ th { | `OlmoeForCausalLM` | OLMoE | `allenai/OLMoE-1B-7B-0924`, `allenai/OLMoE-1B-7B-0924-Instruct`, etc. | | ✅︎ | | `OPTForCausalLM` | OPT, OPT-IML | `facebook/opt-66b`, `facebook/opt-iml-max-30b`, etc. | ✅︎ | ✅︎ | | `OrionForCausalLM` | Orion | `OrionStarAI/Orion-14B-Base`, `OrionStarAI/Orion-14B-Chat`, etc. | | ✅︎ | -| `OuroForCausalLM` | ouro | `ByteDance/Ouro-1.4B`, `ByteDance/Ouro-2.6B`, etc. | ✅︎ | | | `PanguEmbeddedForCausalLM` | openPangu-Embedded-7B | `FreedomIntelligence/openPangu-Embedded-7B-V1.1` | ✅︎ | ✅︎ | | `PanguProMoEV2ForCausalLM` | openpangu-pro-moe-v2 | | ✅︎ | ✅︎ | | `PanguUltraMoEForCausalLM` | openpangu-ultra-moe-718b-model | `FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1` | ✅︎ | ✅︎ | diff --git a/tests/models/registry.py b/tests/models/registry.py index afbac216c93..2d8ecfb560f 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -461,7 +461,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "OrionForCausalLM": _HfExamplesInfo( "OrionStarAI/Orion-14B-Chat", trust_remote_code=True ), - "OuroForCausalLM": _HfExamplesInfo("ByteDance/Ouro-1.4B", trust_remote_code=True), "PanguEmbeddedForCausalLM": _HfExamplesInfo( "FreedomIntelligence/openPangu-Embedded-7B-V1.1", trust_remote_code=True ), diff --git a/vllm/model_executor/models/ouro.py b/vllm/model_executor/models/ouro.py deleted file mode 100644 index 527eeaa13bc..00000000000 --- a/vllm/model_executor/models/ouro.py +++ /dev/null @@ -1,448 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates -# Adapted from -# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/qwen2/modeling_qwen2.py -# Copyright 2024 The Qwen team. -# Copyright 2023 The vLLM team. -# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. -# -# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX -# and OPT implementations in this library. It has been modified from its -# original forms to accommodate minor architectural differences compared -# to GPT-NeoX and OPT used by the Meta AI team that trained the model. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Inference-only Ouro model compatible with HuggingFace weights.""" - -from collections.abc import Iterable -from typing import Any - -import torch -from torch import nn -from transformers import PretrainedConfig - -from vllm.compilation.decorators import support_torch_compile -from vllm.config import CacheConfig, VllmConfig -from vllm.distributed import get_tensor_model_parallel_world_size -from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.linear import ( - MergedColumnParallelLinear, - QKVParallelLinear, - RowParallelLinear, -) -from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.model_executor.layers.rotary_embedding import get_rope -from vllm.model_executor.layers.vocab_parallel_embedding import ( - ParallelLMHead, - VocabParallelEmbedding, -) -from vllm.sequence import IntermediateTensors -from vllm.v1.attention.backend import AttentionType - -from .interfaces import SupportsLoRA -from .utils import ( - AutoWeightsLoader, - WeightsMapper, - extract_layer_index, - make_empty_intermediate_tensors_factory, - make_layers, - maybe_prefix, -) - - -class OuroMLP(nn.Module): - def __init__( - self, - hidden_size: int, - intermediate_size: int, - hidden_act: str, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - super().__init__() - self.gate_up_proj = MergedColumnParallelLinear( - hidden_size, - [intermediate_size] * 2, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.gate_up_proj", - ) - self.down_proj = RowParallelLinear( - intermediate_size, - hidden_size, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.down_proj", - ) - if hidden_act != "silu": - raise ValueError( - f"Unsupported activation: {hidden_act}. Only silu is supported for now." - ) - self.act_fn = SiluAndMul() - - def forward(self, x): - gate_up, _ = self.gate_up_proj(x) - x = self.act_fn(gate_up) - x, _ = self.down_proj(x) - return x - - -class OuroAttention(nn.Module): - def __init__( - self, - config: PretrainedConfig, - hidden_size: int, - num_heads: int, - num_kv_heads: int, - max_position: int = 4096 * 32, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - attn_type: str = AttentionType.DECODER, - dual_chunk_attention_config: dict[str, Any] | None = None, - ) -> None: - super().__init__() - self.hidden_size = hidden_size - tp_size = get_tensor_model_parallel_world_size() - self.total_num_heads = num_heads - assert self.total_num_heads % tp_size == 0 - self.num_heads = self.total_num_heads // tp_size - self.total_num_kv_heads = num_kv_heads - if self.total_num_kv_heads >= tp_size: - # Number of KV heads is greater than TP size, so we partition - # the KV heads across multiple tensor parallel GPUs. - assert self.total_num_kv_heads % tp_size == 0 - else: - # Number of KV heads is less than TP size, so we replicate - # the KV heads across multiple tensor parallel GPUs. - assert tp_size % self.total_num_kv_heads == 0 - self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) - self.head_dim = hidden_size // self.total_num_heads - self.q_size = self.num_heads * self.head_dim - self.kv_size = self.num_kv_heads * self.head_dim - self.scaling = self.head_dim**-0.5 - self.dual_chunk_attention_config = dual_chunk_attention_config - - # Get total_ut_steps from config, default to 4 if not specified - total_ut_steps = getattr(config, "total_ut_steps", 4) - - # Use total number of hidden layers instead of hardcoded 24 - total_layers = config.num_hidden_layers - - self.qkv_proj = QKVParallelLinear( - hidden_size, - self.head_dim, - self.total_num_heads, - self.total_num_kv_heads, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.qkv_proj", - ) - self.o_proj = RowParallelLinear( - self.total_num_heads * self.head_dim, - hidden_size, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.o_proj", - ) - - self.rotary_emb = get_rope( - self.head_dim, - max_position=max_position, - rope_parameters=config.rope_parameters, - dual_chunk_attention_config=dual_chunk_attention_config, - ) - self.attn = nn.ModuleList() - for ut_step in range(total_ut_steps): - base_layer_idx = extract_layer_index(prefix) - unique_layer_idx = ut_step * total_layers + base_layer_idx - - unique_prefix = prefix.replace( - f"layers.{base_layer_idx}", f"layers.{unique_layer_idx}" - ) - - self.attn.append( - Attention( - self.num_heads, - self.head_dim, - self.scaling, - num_kv_heads=self.num_kv_heads, - cache_config=cache_config, - quant_config=quant_config, - attn_type=attn_type, - prefix=f"{unique_prefix}.attn", - **{ - "layer_idx": unique_layer_idx, - "dual_chunk_attention_config": dual_chunk_attention_config, - } - if dual_chunk_attention_config - else {}, - ) - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - current_ut: int, - ) -> torch.Tensor: - qkv, _ = self.qkv_proj(hidden_states) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q, k = self.rotary_emb(positions, q, k) - attn_output = self.attn[current_ut](q, k, v) - output, _ = self.o_proj(attn_output) - return output - - -class OuroDecoderLayer(nn.Module): - def __init__( - self, - config: PretrainedConfig, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - super().__init__() - self.hidden_size = config.hidden_size - dual_chunk_attention_config = getattr( - config, "dual_chunk_attention_config", None - ) - - if getattr(config, "is_causal", True): - attn_type = AttentionType.DECODER - else: - attn_type = AttentionType.ENCODER_ONLY - - self.self_attn = OuroAttention( - config=config, - hidden_size=self.hidden_size, - num_heads=config.num_attention_heads, - max_position=config.max_position_embeddings, - num_kv_heads=config.num_key_value_heads, - cache_config=cache_config, - quant_config=quant_config, - prefix=f"{prefix}.self_attn", - attn_type=attn_type, - dual_chunk_attention_config=dual_chunk_attention_config, - ) - self.mlp = OuroMLP( - hidden_size=self.hidden_size, - intermediate_size=config.intermediate_size, - hidden_act=config.hidden_act, - quant_config=quant_config, - prefix=f"{prefix}.mlp", - ) - self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.input_layernorm_2 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - - self.post_attention_layernorm = RMSNorm( - config.hidden_size, eps=config.rms_norm_eps - ) - self.post_attention_layernorm_2 = RMSNorm( - config.hidden_size, eps=config.rms_norm_eps - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - current_ut: int, - residual: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - if residual is None: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - else: - hidden_states, residual = self.input_layernorm(hidden_states, residual) - hidden_states = self.self_attn( - positions=positions, hidden_states=hidden_states, current_ut=current_ut - ) - hidden_states = self.input_layernorm_2(hidden_states) - - hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) - hidden_states = self.mlp(hidden_states) - hidden_states = self.post_attention_layernorm_2(hidden_states) - - return hidden_states, residual - - -@support_torch_compile( - dynamic_arg_dims={ - "input_ids": 0, - "positions": -1, - "intermediate_tensors": 0, - "inputs_embeds": 0, - } -) -class OuroModel(nn.Module): - def __init__( - self, - *, - vllm_config: VllmConfig, - prefix: str = "", - decoder_layer_type: type[nn.Module] = OuroDecoderLayer, - ): - super().__init__() - - config = vllm_config.model_config.hf_config - cache_config = vllm_config.cache_config - quant_config = vllm_config.quant_config - - # TODO (@robertgshaw2): see if this can be moved out - if cache_config.sliding_window is not None and hasattr( - config, "max_window_layers" - ): - assert config.max_window_layers == config.num_hidden_layers, ( - "Sliding window for some but all layers is not supported. " - "This model uses sliding window but `max_window_layers` = {} " - "is less than `num_hidden_layers` = {}. Please open an issue " - "to discuss this feature.".format( - config.max_window_layers, - config.num_hidden_layers, - ) - ) - - self.config = config - self.quant_config = quant_config - self.vocab_size = config.vocab_size - - self.embed_tokens = VocabParallelEmbedding( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=f"{prefix}.embed_tokens", - ) - - # Use the provided decoder layer type or default to OuroDecoderLayer - decoder_layer_type = decoder_layer_type or OuroDecoderLayer - self.start_layer, self.end_layer, self.layers = make_layers( - config.num_hidden_layers, - lambda prefix: decoder_layer_type( - config=config, - cache_config=cache_config, - quant_config=quant_config, - prefix=prefix, - ), - prefix=f"{prefix}.layers", - ) - - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( - ["hidden_states", "residual"], config.hidden_size - ) - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.early_exit_gate = RowParallelLinear(config.hidden_size, 1, bias=True) - - self.total_ut_steps = getattr(self.config, "total_ut_steps", 4) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: - if inputs_embeds is not None: - hidden_states = inputs_embeds - else: - hidden_states = self.embed_input_ids(input_ids) - - for current_ut in range(self.total_ut_steps): - residual = None - for layer in self.layers[self.start_layer : self.end_layer]: - hidden_states, residual = layer( - positions, hidden_states, current_ut, residual - ) - hidden_states, _ = self.norm(hidden_states, residual) - return hidden_states - - -class OuroForCausalLM(nn.Module, SupportsLoRA): - hf_to_vllm_mapper = WeightsMapper( - orig_to_new_stacked={ - # weight_name: (param_name, shard_id) - ".q_proj": (".qkv_proj", "q"), - ".k_proj": (".qkv_proj", "k"), - ".v_proj": (".qkv_proj", "v"), - ".gate_proj": (".gate_up_proj", 0), - ".up_proj": (".gate_up_proj", 1), - } - ) - packed_modules_mapping = { - "qkv_proj": ["q_proj", "k_proj", "v_proj"], - "gate_up_proj": ["gate_proj", "up_proj"], - } - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - config = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - - self.config = config - - self.quant_config = quant_config - self.model = OuroModel( - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") - ) - - if config.tie_word_embeddings: - self.lm_head = self.model.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) - - self.logits_processor = LogitsProcessor(config.vocab_size) - - self.make_empty_intermediate_tensors = ( - self.model.make_empty_intermediate_tensors - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.model.embed_input_ids(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: - hidden_states = self.model( - input_ids, positions, intermediate_tensors, inputs_embeds - ) - return hidden_states - - def compute_logits( - self, - hidden_states: torch.Tensor, - ) -> torch.Tensor | None: - logits = self.logits_processor(self.lm_head, hidden_states) - return logits - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index cc2688e37fc..229cf4fe87f 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -183,7 +183,6 @@ _TEXT_GENERATION_MODELS = { "OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"), "OPTForCausalLM": ("opt", "OPTForCausalLM"), "OrionForCausalLM": ("orion", "OrionForCausalLM"), - "OuroForCausalLM": ("ouro", "OuroForCausalLM"), "PanguEmbeddedForCausalLM": ("openpangu", "PanguEmbeddedForCausalLM"), "PanguProMoEV2ForCausalLM": ("openpangu", "PanguProMoEV2ForCausalLM"), "PanguUltraMoEForCausalLM": ("openpangu", "PanguUltraMoEForCausalLM"), @@ -760,6 +759,7 @@ _PREVIOUSLY_SUPPORTED_MODELS = { "PersimmonForCausalLM": "0.25.0", "FuyuForCausalLM": "0.25.0", "Plamo2ForCausalLM": "0.26.0", + "OuroForCausalLM": "0.26.0", } _OOT_SUPPORTED_MODELS = { From fe5145765f211c9b5abc446b42491f84895d70ba Mon Sep 17 00:00:00 2001 From: Tri Vo Date: Sat, 25 Jul 2026 17:06:14 +0700 Subject: [PATCH 032/185] [Core] Keep attention backends eligible for text-only serving of prefix-LM models (#48796) Signed-off-by: qtris123 Co-authored-by: Cyrus Leung --- tests/config/test_multimodal_config.py | 100 ++++++++++++++++++ tests/models/utils.py | 6 ++ vllm/config/model.py | 44 +++++++- .../model_arch_config_convertor.py | 22 ++-- 4 files changed, 165 insertions(+), 7 deletions(-) diff --git a/tests/config/test_multimodal_config.py b/tests/config/test_multimodal_config.py index 5260d7a40f7..befcc73b0b9 100644 --- a/tests/config/test_multimodal_config.py +++ b/tests/config/test_multimodal_config.py @@ -1,10 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from unittest.mock import MagicMock, patch + import pytest +from transformers import PretrainedConfig from vllm.config.model import ModelConfig from vllm.config.multimodal import MultiModalConfig +from vllm.transformers_utils.model_arch_config_convertor import ( + ModelArchConfigConvertorBase, +) from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -68,3 +74,97 @@ def test_mm_encoder_attn_dtype_hash_updates(tmp_path): ).compute_hash() assert base_hash != fp8_hash assert fp8_hash != fp8_static_hash + + +def _make_mm_prefix_model_config( + *, + language_model_only: bool = False, +) -> ModelConfig: + model_config = MagicMock(spec=ModelConfig) + model_config.multimodal_config = MultiModalConfig( + language_model_only=language_model_only + ) + # Bind real helper methods onto the mock. + model_config._supports_multimodal_for_mm_prefix = ( + ModelConfig._supports_multimodal_for_mm_prefix.__get__( + model_config, ModelConfig + ) + ) + return model_config + + +@pytest.mark.parametrize("supports_mm", [True, False]) +def test_supports_multimodal_for_mm_prefix_uses_registry(supports_mm: bool): + model_config = _make_mm_prefix_model_config() + + with patch( + "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs", + return_value=supports_mm, + ) as mocked: + assert model_config._supports_multimodal_for_mm_prefix() is supports_mm + mocked.assert_called_once_with(model_config) + + # Sticky cache — registry must not be consulted again. + with patch( + "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs", + side_effect=AssertionError("should use cache"), + ): + assert model_config._supports_multimodal_for_mm_prefix() is supports_mm + + +def test_supports_multimodal_for_mm_prefix_before_multimodal_config(): + model_config = _make_mm_prefix_model_config() + model_config.multimodal_config = None + + assert model_config._supports_multimodal_for_mm_prefix() is True + assert not hasattr(model_config, "_supports_multimodal_inputs_cached") + + +def test_language_model_only_disables_via_supports_multimodal_inputs(): + """language_model_only zeros all limits, so registry reports text-only.""" + model_config = _make_mm_prefix_model_config(language_model_only=True) + + with patch( + "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs", + return_value=False, + ): + assert model_config._supports_multimodal_for_mm_prefix() is False + + +def test_convertor_clears_mm_prefix_when_multimodal_disabled(): + hf_config = PretrainedConfig( + model_type="gemma3", + architectures=["Gemma3ForConditionalGeneration"], + ) + hf_config.is_mm_prefix_lm = True + convertor = ModelArchConfigConvertorBase(hf_config, hf_config) + + assert convertor.is_mm_prefix_lm(supports_multimodal=True) is True + assert convertor.is_mm_prefix_lm(supports_multimodal=False) is False + + enabled = convertor.convert(supports_multimodal=True) + disabled = convertor.convert(supports_multimodal=False) + assert enabled.is_mm_prefix_lm is True + assert disabled.is_mm_prefix_lm is False + + +def test_sticky_cache_survives_text_subconfig_regeneration(): + """with_hf_config deepcopies the cached decision onto text submodules.""" + model_config = _make_mm_prefix_model_config() + with patch( + "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs", + return_value=False, + ): + assert model_config._supports_multimodal_for_mm_prefix() is False + + # Simulate deepcopy onto a Gemma4ForCausalLM-like config that would + # otherwise fail registry lookup / return False incorrectly. + text_config = _make_mm_prefix_model_config() + text_config._supports_multimodal_inputs_cached = ( + model_config._supports_multimodal_inputs_cached + ) + with patch( + "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs", + side_effect=AssertionError("must not re-query registry"), + ): + assert text_config._supports_multimodal_for_mm_prefix() is False diff --git a/tests/models/utils.py b/tests/models/utils.py index 86fd80cd62a..f938a702231 100644 --- a/tests/models/utils.py +++ b/tests/models/utils.py @@ -507,6 +507,12 @@ def dummy_hf_overrides( hf_config = _hf_config hf_text_config = text_config + # Keep architecture conversion on the default multimodal path; this + # helper only needs HF-derived fields, not deployment MM limits. + @staticmethod + def _supports_multimodal_for_mm_prefix() -> bool: + return True + model_arch_config = ModelConfig.get_model_arch_config(DummyConfig) # Only set MoE related config when the model has MoE layers. # Otherwise all models detected as MoE by _get_transformers_backend_cls. diff --git a/vllm/config/model.py b/vllm/config/model.py index a65d5d5a94f..6b7825be08f 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -758,6 +758,10 @@ class ModelConfig: "disable the cache with --mm-processor-cache-gb 0." ) + # Rebuild after multimodal_config exists so text-only mm_prefix + # clearing is applied (and cached for later with_hf_config calls). + self.model_arch_config = self.get_model_arch_config() + if self.disable_sliding_window: # Set after get_and_verify_max_len to ensure that max_model_len # can be correctly capped to sliding window size @@ -770,6 +774,42 @@ class ModelConfig: self._verify_cuda_graph() self._verify_bnb_config() + def _supports_multimodal_for_mm_prefix(self) -> bool: + """Whether multimodal inputs can still appear for this deployment. + + This runs more than once per config: once early in ``__post_init__`` + (before ``multimodal_config`` exists), again after it is created, and + then for every ``get_model_arch_config`` regeneration -- notably + ``with_hf_config``, which deep-copies this ``ModelConfig`` and swaps + ``hf_config`` for a text-only submodule (e.g. ``Gemma4ForCausalLM``). + + The result is cached for correctness, not just to save work: on the + ``with_hf_config`` copy the submodule architecture has no registered + multimodal processor, so re-querying the registry would raise and be + treated as text-only, wrongly clearing ``is_mm_prefix_lm`` even when a + vision modality is still enabled (e.g. ``image=0`` but video allowed). + The deep-copied cache preserves the top-level decision instead. + """ + cached = getattr(self, "_supports_multimodal_inputs_cached", None) + if cached is not None: + return cached + + if self.multimodal_config is None: + # Early call before multimodal init — do not clear mm_prefix yet. + return True + + from vllm.multimodal import MULTIMODAL_REGISTRY + + supports_mm = MULTIMODAL_REGISTRY.supports_multimodal_inputs(self) + self._supports_multimodal_inputs_cached = supports_mm + if not supports_mm: + logger.info_once( + "Disabled mm_prefix attention mode because multimodal inputs " + "are configuration-disabled. Attention backends without " + "mm_prefix support may now be selected." + ) + return supports_mm + def get_model_arch_config( self, ) -> ModelArchitectureConfig: @@ -777,7 +817,9 @@ class ModelConfig: self.hf_config.model_type, ModelArchConfigConvertorBase ) convertor = convertor_cls(self.hf_config, self.hf_text_config) - return convertor.convert() + return convertor.convert( + supports_multimodal=self._supports_multimodal_for_mm_prefix() + ) @field_validator("tokenizer", "max_model_len", mode="wrap") @classmethod diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 70bb4caa535..e62b29295f0 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -300,8 +300,16 @@ class ModelArchConfigConvertorBase: ) return False - def is_mm_prefix_lm(self) -> bool: - """Whether to use bidirectional attention for mm positions.""" + def is_mm_prefix_lm(self, supports_multimodal: bool = True) -> bool: + """Whether to use bidirectional attention for mm positions. + + ``supports_multimodal`` is False when the deployment is configuration- + disabled for multimodal inputs (text-only serving). In that case + mm_prefix is unnecessary and must stay off so attention backends + without ``supports_mm_prefix()`` remain eligible. + """ + if not supports_multimodal: + return False if hasattr(self.hf_config, "is_mm_prefix_lm"): return bool(self.hf_config.is_mm_prefix_lm) # fallback to list of known models @@ -358,7 +366,7 @@ class ModelArchConfigConvertorBase: derived_max_model_len = tmp_max_len return derived_max_model_len, max_len_key - def convert(self) -> ModelArchitectureConfig: + def convert(self, supports_multimodal: bool = True) -> ModelArchitectureConfig: model_arch_config = ModelArchitectureConfig( architectures=self.get_architectures(), model_type=self.hf_config.model_type, @@ -372,7 +380,7 @@ class ModelArchConfigConvertorBase: num_experts=self.get_num_experts(), quantization_config=self.get_quantization_config(), is_deepseek_mla=self.is_deepseek_mla(), - is_mm_prefix_lm=self.is_mm_prefix_lm(), + is_mm_prefix_lm=self.is_mm_prefix_lm(supports_multimodal), rswa_window=self.rswa_window(), derived_max_model_len_and_key=self.derive_max_model_len_and_key(), ) @@ -401,7 +409,7 @@ class CohereAsrModelArchConfigConvertor(ModelArchConfigConvertorBase): ) return enc_num_kv_heads - def is_mm_prefix_lm(self) -> bool: + def is_mm_prefix_lm(self, supports_multimodal: bool = True) -> bool: return False @@ -584,7 +592,9 @@ class Gemma4MTPModelArchConfigConvertor(ModelArchConfigConvertorBase): class Gemma4ModelArchConfigConvertor(ModelArchConfigConvertorBase): - def is_mm_prefix_lm(self) -> bool: + def is_mm_prefix_lm(self, supports_multimodal: bool = True) -> bool: + if not supports_multimodal: + return False return ( getattr(self.hf_text_config, "use_bidirectional_attention", None) == "vision" From 0b1a8bb1f6386dd0b60105f63cf67337cb775a56 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Sat, 25 Jul 2026 11:10:26 +0100 Subject: [PATCH 033/185] [Bugfix][CI] Fix stale Mooncake lookup expectation broken by a merge race (#49802) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) --- tests/v1/kv_connector/unit/test_mooncake_store_worker.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index 7cddd6d4a7c..b9876b9827b 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -1905,13 +1905,15 @@ def test_lookup_rejects_boundary_missing_one_mamba_shard(): ) _refresh_group_tp_replication_factors(worker) + # 33 tokens for two 16-token blocks: the hit stops below the request end, + # so the full-hit re-derivation stays out of the shard accounting. worker.store.batch_is_exist.side_effect = lambda keys: [1] * len(keys) - assert worker.lookup(32, [b"h0", b"h1"]) == 32 + assert worker.lookup(33, [b"h0", b"h1"]) == 32 worker.store.batch_is_exist.side_effect = lambda keys: [ 0 if "tp_rank:1" in k and "group:1" in k else 1 for k in keys ] - assert worker.lookup(32, [b"h0", b"h1"]) == 0 + assert worker.lookup(33, [b"h0", b"h1"]) == 0 def test_lookup_requires_all_dcp_rank_namespaces(): From ca0defa3438832bb3fc5826d8bf2d4e572c57e3d Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Sat, 25 Jul 2026 12:18:39 +0100 Subject: [PATCH 034/185] Make bare `hugging_face` imports forbidden (#49726) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Andreas Karatzas --- tests/conftest.py | 8 +- tests/distributed/test_rocm_quick_reduce.py | 6 +- tests/entrypoints/conftest.py | 18 +++-- ...chat_completion_with_mixed_audio_embeds.py | 8 +- ...chat_completion_with_mixed_image_embeds.py | 6 +- .../chat_completion/test_default_mm_loras.py | 6 +- .../openai/chat_completion/test_chat.py | 4 +- tests/entrypoints/pooling/scoring/util.py | 4 +- tests/entrypoints/serve/sagemaker/conftest.py | 4 +- .../tool_parsers/test_hermes_tool_parser.py | 4 +- .../quantization/test_nvfp4_emulation.py | 4 +- tests/lora/conftest.py | 76 +++++++++++------- tests/lora/test_default_mm_loras.py | 4 +- .../test_runai_model_streamer_s3.py | 4 +- .../model_loader/test_sharded_state_loader.py | 4 +- tests/models/language/pooling/test_colbert.py | 5 +- .../multimodal/generation/test_phi4mm.py | 4 +- .../multimodal/pooling/test_intern_vit.py | 4 +- tests/models/multimodal/pooling/test_radio.py | 4 +- tests/models/quantization/test_gpt_oss.py | 3 +- .../gguf/test_gguf_plugin_multimodal.py | 6 +- .../test_filesystem_resolver.py | 6 +- tests/quantization/test_modelopt.py | 4 +- tests/quantization/test_quark.py | 3 +- tests/tool_use/conftest.py | 4 +- tests/tool_use/mistral/conftest.py | 4 +- tests/utils.py | 4 +- tools/pre_commit/check_forbidden_imports.py | 79 ++++++++++++++++++- .../model_executor/models/llava_onevision2.py | 15 ++-- 29 files changed, 199 insertions(+), 106 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index db097fbc455..2c5a907e557 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -31,7 +31,7 @@ import pytest import torch import torch.nn as nn import torch.nn.functional as F -from huggingface_hub import snapshot_download +from vllm.transformers_utils.repo_utils import hf_api from PIL import Image from transformers import ( AutoConfig, @@ -1496,7 +1496,7 @@ _dummy_gemma2_embedding_path = os.path.join(temp_dir, "dummy_gemma2_embedding") def dummy_opt_path(): json_path = os.path.join(_dummy_opt_path, "config.json") if not os.path.exists(_dummy_opt_path): - snapshot_download( + hf_api().snapshot_download( repo_id="facebook/opt-125m", local_dir=_dummy_opt_path, ignore_patterns=["*.bin", "*.bin.index.json", "*.pt", "*.h5", "*.msgpack"], @@ -1514,7 +1514,7 @@ def dummy_opt_path(): def dummy_llava_path(): json_path = os.path.join(_dummy_llava_path, "config.json") if not os.path.exists(_dummy_llava_path): - snapshot_download( + hf_api().snapshot_download( repo_id="llava-hf/llava-1.5-7b-hf", local_dir=_dummy_llava_path, ignore_patterns=[ @@ -1539,7 +1539,7 @@ def dummy_llava_path(): def dummy_gemma2_embedding_path(): json_path = os.path.join(_dummy_gemma2_embedding_path, "config.json") if not os.path.exists(_dummy_gemma2_embedding_path): - snapshot_download( + hf_api().snapshot_download( repo_id="BAAI/bge-multilingual-gemma2", local_dir=_dummy_gemma2_embedding_path, ignore_patterns=[ diff --git a/tests/distributed/test_rocm_quick_reduce.py b/tests/distributed/test_rocm_quick_reduce.py index 06933fb0af2..a7e91f236a2 100644 --- a/tests/distributed/test_rocm_quick_reduce.py +++ b/tests/distributed/test_rocm_quick_reduce.py @@ -15,12 +15,12 @@ from unittest.mock import patch import pytest import torch import torch.distributed as dist -from huggingface_hub import snapshot_download import vllm.envs as envs from vllm import LLM, SamplingParams from vllm.distributed import cleanup_dist_env_and_memory from vllm.platforms import current_platform +from vllm.transformers_utils.repo_utils import hf_api from vllm.utils.network_utils import get_open_port pytestmark = pytest.mark.skipif( @@ -210,11 +210,11 @@ def _log_prompt_summaries() -> None: @lru_cache(maxsize=1) def _get_model_path() -> str: try: - path = snapshot_download(repo_id=MODEL_NAME, local_files_only=True) + path = hf_api().snapshot_download(repo_id=MODEL_NAME, local_files_only=True) _log(f"using cached model snapshot: {path}") return path except Exception: - path = snapshot_download(repo_id=MODEL_NAME) + path = hf_api().snapshot_download(repo_id=MODEL_NAME) _log(f"downloaded model snapshot: {path}") return path diff --git a/tests/entrypoints/conftest.py b/tests/entrypoints/conftest.py index c2e9a1de318..27383f87ea7 100644 --- a/tests/entrypoints/conftest.py +++ b/tests/entrypoints/conftest.py @@ -190,30 +190,32 @@ number: "1" | "2" @pytest.fixture(scope="session") def qwen3_lora_files(): """Download Qwen3 LoRA files once per test session.""" - from huggingface_hub import snapshot_download + from vllm.transformers_utils.repo_utils import hf_api - return snapshot_download(repo_id="charent/self_cognition_Alice") + return hf_api().snapshot_download(repo_id="charent/self_cognition_Alice") @pytest.fixture(scope="session") def qwen3_meowing_lora_files(): """Download Qwen3 LoRA files once per test session.""" - from huggingface_hub import snapshot_download + from vllm.transformers_utils.repo_utils import hf_api - return snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Meow-LoRA") + return hf_api().snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Meow-LoRA") @pytest.fixture(scope="session") def qwen3_woofing_lora_files(): """Download Qwen3 LoRA files once per test session.""" - from huggingface_hub import snapshot_download + from vllm.transformers_utils.repo_utils import hf_api - return snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Woof-LoRA") + return hf_api().snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Woof-LoRA") @pytest.fixture(scope="session") def opt125_lora_files() -> str: """Download opt-125m LoRA files once per test session.""" - from huggingface_hub import snapshot_download + from vllm.transformers_utils.repo_utils import hf_api - return snapshot_download(repo_id="peft-internal-testing/opt-125m-dummy-lora") + return hf_api().snapshot_download( + repo_id="peft-internal-testing/opt-125m-dummy-lora" + ) diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py b/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py index a417f860f61..aac68334bf9 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py @@ -12,10 +12,10 @@ import pytest_asyncio import safetensors import torch import torch.nn as nn -from huggingface_hub import hf_hub_download from transformers import AutoConfig, AutoTokenizer from tests.utils import RemoteOpenAIServer +from vllm.transformers_utils.repo_utils import hf_api from vllm.utils.serial_utils import tensor2base64 from vllm.utils.torch_utils import is_torch_equal_or_newer @@ -126,11 +126,13 @@ def qwen2audio_aligned_content_and_embeds_b64() -> tuple[str, str]: content = "Describe this audio." tokenizer = AutoTokenizer.from_pretrained(QWEN2AUDIO_MODEL, trust_remote_code=True) - index_path = hf_hub_download(QWEN2AUDIO_MODEL, "model.safetensors.index.json") + index_path = hf_api().hf_hub_download( + QWEN2AUDIO_MODEL, "model.safetensors.index.json" + ) with open(index_path) as f: weight_map = json.load(f)["weight_map"] embed_key = next(k for k in weight_map if k.endswith("embed_tokens.weight")) - shard_path = hf_hub_download(QWEN2AUDIO_MODEL, weight_map[embed_key]) + shard_path = hf_api().hf_hub_download(QWEN2AUDIO_MODEL, weight_map[embed_key]) with safetensors.safe_open(shard_path, framework="pt", device="cpu") as f: embed_weight = f.get_tensor(embed_key) embed_layer = nn.Embedding.from_pretrained(embed_weight.to(QWEN2AUDIO_DTYPE)) diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py b/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py index dbbed3c4712..c6ce165cf39 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py @@ -13,12 +13,12 @@ import pytest_asyncio import safetensors import torch import torch.nn as nn -from huggingface_hub import hf_hub_download from transformers import AutoTokenizer from tests.utils import RemoteOpenAIServer from vllm.assets.image import ImageAsset from vllm.multimodal.utils import encode_image_url +from vllm.transformers_utils.repo_utils import hf_api from vllm.utils.serial_utils import tensor2base64 MODEL_NAME = "Qwen/Qwen2-VL-2B-Instruct" @@ -84,11 +84,11 @@ def aligned_content_and_embeds_b64() -> tuple[str, str]: content = "Describe this image." tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True) - index_path = hf_hub_download(MODEL_NAME, "model.safetensors.index.json") + index_path = hf_api().hf_hub_download(MODEL_NAME, "model.safetensors.index.json") with open(index_path) as f: weight_map = json.load(f)["weight_map"] embed_key = next(k for k in weight_map if k.endswith("embed_tokens.weight")) - shard_path = hf_hub_download(MODEL_NAME, weight_map[embed_key]) + shard_path = hf_api().hf_hub_download(MODEL_NAME, weight_map[embed_key]) with safetensors.safe_open(shard_path, framework="pt", device="cpu") as f: embed_weight = f.get_tensor(embed_key) embed_layer = nn.Embedding.from_pretrained(embed_weight.to(MODEL_DTYPE)) diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_default_mm_loras.py b/tests/entrypoints/multimodal/openai/chat_completion/test_default_mm_loras.py index e285c8d3139..eca23d0a48a 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_default_mm_loras.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_default_mm_loras.py @@ -6,17 +6,19 @@ import os import openai # use the official client for correctness check import pytest import pytest_asyncio -from huggingface_hub import snapshot_download from tests.conftest import AudioTestAssets from tests.utils import RemoteOpenAIServer +from vllm.transformers_utils.repo_utils import hf_api # NOTE - the tests in this module are currently analogous to test_chat, but are # separated to avoid OOM killing due to module-scoped servers, since we # need a multimodal model for these tests. # Contains a modality specific lora alongside the base model -MULTIMODAL_MODEL_NAME = snapshot_download("microsoft/Phi-4-multimodal-instruct") +MULTIMODAL_MODEL_NAME = hf_api().snapshot_download( + "microsoft/Phi-4-multimodal-instruct" +) AUDIO_LORA_PATH = os.path.join(MULTIMODAL_MODEL_NAME, "speech-lora") ACTIVE_MM_LORA_RESPONSE = "Spoken text: The first words I spoke in the original chronograph, a little piece of practical poetry. Mary had a little lamb, it slept with quite a snow, and everywhere that Mary went, the lamb was sure to go." # noqa: E501 diff --git a/tests/entrypoints/openai/chat_completion/test_chat.py b/tests/entrypoints/openai/chat_completion/test_chat.py index dbfb48f2351..32c72f1ef93 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_chat.py @@ -27,9 +27,9 @@ MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta" @pytest.fixture(scope="module") def zephyr_lora_files(): """Download zephyr LoRA files once per test session.""" - from huggingface_hub import snapshot_download + from vllm.transformers_utils.repo_utils import hf_api - return snapshot_download(repo_id="typeof/zephyr-7b-beta-lora") + return hf_api().snapshot_download(repo_id="typeof/zephyr-7b-beta-lora") @pytest.fixture(scope="module") diff --git a/tests/entrypoints/pooling/scoring/util.py b/tests/entrypoints/pooling/scoring/util.py index 8aab9cd1069..dbf4b3dc3ac 100644 --- a/tests/entrypoints/pooling/scoring/util.py +++ b/tests/entrypoints/pooling/scoring/util.py @@ -6,7 +6,6 @@ from io import BytesIO import pybase64 as base64 import torch import torch.nn.functional as F -from huggingface_hub import hf_hub_download from PIL import Image from safetensors.torch import load_file from transformers import AutoModel, AutoTokenizer @@ -18,6 +17,7 @@ from vllm.entrypoints.chat_utils import ( ) from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score +from vllm.transformers_utils.repo_utils import hf_api class ColBERTScoringHfRunner(torch.nn.Module): @@ -38,7 +38,7 @@ class ColBERTScoringHfRunner(torch.nn.Module): ).to(self.device) self.model.eval() - path = hf_hub_download(model_name, filename="model.safetensors") + path = hf_api().hf_hub_download(model_name, filename="model.safetensors") weights = load_file(path) self.linear_weight = weights[linear_weights_key].to(self.device).float() diff --git a/tests/entrypoints/serve/sagemaker/conftest.py b/tests/entrypoints/serve/sagemaker/conftest.py index d36c20ccd9a..0c76d25b423 100644 --- a/tests/entrypoints/serve/sagemaker/conftest.py +++ b/tests/entrypoints/serve/sagemaker/conftest.py @@ -21,9 +21,9 @@ HEADER_SAGEMAKER_NEW_SESSION_ID = "X-Amzn-SageMaker-New-Session-Id" @pytest.fixture(scope="session") def smollm2_lora_files(): """Download LoRA files once per test session.""" - from huggingface_hub import snapshot_download + from vllm.transformers_utils.repo_utils import hf_api - return snapshot_download(repo_id=LORA_ADAPTER_NAME_SMOLLM) + return hf_api().snapshot_download(repo_id=LORA_ADAPTER_NAME_SMOLLM) @pytest.fixture(scope="module") diff --git a/tests/entrypoints/tool_parsers/test_hermes_tool_parser.py b/tests/entrypoints/tool_parsers/test_hermes_tool_parser.py index 5d769c0fd88..f9a7bef2533 100644 --- a/tests/entrypoints/tool_parsers/test_hermes_tool_parser.py +++ b/tests/entrypoints/tool_parsers/test_hermes_tool_parser.py @@ -6,13 +6,13 @@ import json import openai import pytest import pytest_asyncio -from huggingface_hub import snapshot_download from typing_extensions import TypedDict from tests.utils import RemoteOpenAIServer from vllm.tool_parsers.abstract_tool_parser import ToolParser from vllm.tool_parsers.granite4_tool_parser import Granite4ToolParser from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser +from vllm.transformers_utils.repo_utils import hf_api LORA_MODEL = "minpeter/LoRA-Llama-3.2-1B-tool-vllm-ci" @@ -91,7 +91,7 @@ def server_config(request): config = CONFIGS[request.param] # download model and tokenizer using transformers - snapshot_download(config["model"]) + hf_api().snapshot_download(config["model"]) yield CONFIGS[request.param] diff --git a/tests/kernels/quantization/test_nvfp4_emulation.py b/tests/kernels/quantization/test_nvfp4_emulation.py index d2056fe01eb..eb057cebd91 100644 --- a/tests/kernels/quantization/test_nvfp4_emulation.py +++ b/tests/kernels/quantization/test_nvfp4_emulation.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import cast -import huggingface_hub import pytest import torch from safetensors import safe_open @@ -33,6 +32,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kNvfp4Static, ) from vllm.platforms import current_platform +from vllm.transformers_utils.repo_utils import hf_api from vllm.triton_utils import triton if current_platform.is_rocm(): @@ -56,7 +56,7 @@ MOE_MODEL_CONFIGS = { @pytest.fixture(scope="module") def loaded_model_files(): return { - model_id: huggingface_hub.snapshot_download( + model_id: hf_api().snapshot_download( repo_id=model_id, allow_patterns=config["shards"] ) for model_id, config in MOE_MODEL_CONFIGS.items() diff --git a/tests/lora/conftest.py b/tests/lora/conftest.py index dea54ed21ae..899b3129d62 100644 --- a/tests/lora/conftest.py +++ b/tests/lora/conftest.py @@ -9,7 +9,6 @@ from unittest.mock import MagicMock import pytest import torch import torch.nn as nn -from huggingface_hub import snapshot_download from vllm.distributed import ( cleanup_dist_env_and_memory, @@ -25,6 +24,7 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.model_executor.models.interfaces import SupportsLoRA from vllm.platforms import current_platform +from vllm.transformers_utils.repo_utils import hf_api @pytest.fixture() @@ -170,43 +170,43 @@ def dummy_model_gate_up(default_vllm_config) -> nn.Module: def mixtral_lora_files(): # Note: this module has incorrect adapter_config.json to test # https://github.com/vllm-project/vllm/pull/5909/files. - return snapshot_download(repo_id="SangBinCho/mixtral-lora") + return hf_api().snapshot_download(repo_id="SangBinCho/mixtral-lora") @pytest.fixture(scope="session") def chatglm3_lora_files(): - return snapshot_download(repo_id="jeeejeee/chatglm3-text2sql-spider") + return hf_api().snapshot_download(repo_id="jeeejeee/chatglm3-text2sql-spider") @pytest.fixture(scope="session") def baichuan_lora_files(): - return snapshot_download(repo_id="jeeejeee/baichuan7b-text2sql-spider") + return hf_api().snapshot_download(repo_id="jeeejeee/baichuan7b-text2sql-spider") @pytest.fixture(scope="session") def baichuan_zero_lora_files(): # all the lora_B weights are initialized to zero. - return snapshot_download(repo_id="jeeejeee/baichuan7b-zero-init") + return hf_api().snapshot_download(repo_id="jeeejeee/baichuan7b-zero-init") @pytest.fixture(scope="session") def baichuan_regex_lora_files(): - return snapshot_download(repo_id="jeeejeee/baichuan-7b-lora-zero-regex") + return hf_api().snapshot_download(repo_id="jeeejeee/baichuan-7b-lora-zero-regex") @pytest.fixture(scope="session") def ilama_lora_files(): - return snapshot_download(repo_id="jeeejeee/ilama-text2sql-spider") + return hf_api().snapshot_download(repo_id="jeeejeee/ilama-text2sql-spider") @pytest.fixture(scope="session") def minicpmv_lora_files(): - return snapshot_download(repo_id="jeeejeee/minicpmv25-lora-pokemon") + return hf_api().snapshot_download(repo_id="jeeejeee/minicpmv25-lora-pokemon") @pytest.fixture(scope="session") def qwen2vl_lora_files(): - return snapshot_download(repo_id="jeeejeee/qwen2-vl-lora-pokemon") + return hf_api().snapshot_download(repo_id="jeeejeee/qwen2-vl-lora-pokemon") @pytest.fixture(scope="session") @@ -217,74 +217,84 @@ def qwen25vl_base_huggingface_id(): @pytest.fixture(scope="session") def qwen25vl_lora_files(): - return snapshot_download(repo_id="jeeejeee/qwen25-vl-lora-pokemon") + return hf_api().snapshot_download(repo_id="jeeejeee/qwen25-vl-lora-pokemon") @pytest.fixture(scope="session") def qwen2vl_language_lora_files(): - return snapshot_download(repo_id="prashanth058/qwen2vl-flickr-lora-language") + return hf_api().snapshot_download( + repo_id="prashanth058/qwen2vl-flickr-lora-language" + ) @pytest.fixture(scope="session") def qwen2vl_vision_tower_connector_lora_files(): - return snapshot_download(repo_id="prashanth058/qwen2vl-flickr-lora-tower-connector") + return hf_api().snapshot_download( + repo_id="prashanth058/qwen2vl-flickr-lora-tower-connector" + ) @pytest.fixture(scope="session") def qwen2vl_vision_tower_lora_files(): - return snapshot_download(repo_id="prashanth058/qwen2vl-flickr-lora-tower") + return hf_api().snapshot_download(repo_id="prashanth058/qwen2vl-flickr-lora-tower") @pytest.fixture(scope="session") def qwen25vl_vision_lora_files(): - return snapshot_download(repo_id="EpochEcho/qwen2.5-3b-vl-lora-vision-connector") + return hf_api().snapshot_download( + repo_id="EpochEcho/qwen2.5-3b-vl-lora-vision-connector" + ) @pytest.fixture(scope="session") def qwen3vl_vision_lora_files(): - return snapshot_download(repo_id="EpochEcho/qwen3-4b-vl-lora-vision-connector") + return hf_api().snapshot_download( + repo_id="EpochEcho/qwen3-4b-vl-lora-vision-connector" + ) @pytest.fixture(scope="session") def qwen3_meowing_lora_files(): """Download Qwen3 Meow LoRA files once per test session.""" - return snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Meow-LoRA") + return hf_api().snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Meow-LoRA") @pytest.fixture(scope="session") def qwen3_woofing_lora_files(): """Download Qwen3 Woof LoRA files once per test session.""" - return snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Woof-LoRA") + return hf_api().snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Woof-LoRA") @pytest.fixture(scope="session") def tinyllama_lora_files(): - return snapshot_download(repo_id="jashing/tinyllama-colorist-lora") + return hf_api().snapshot_download(repo_id="jashing/tinyllama-colorist-lora") @pytest.fixture(scope="session") def deepseekv2_lora_files(): - return snapshot_download(repo_id="wuchen01/DeepSeek-V2-Lite-Chat-All-LoRA") + return hf_api().snapshot_download(repo_id="wuchen01/DeepSeek-V2-Lite-Chat-All-LoRA") @pytest.fixture(scope="session") def gptoss20b_lora_files(): - return snapshot_download(repo_id="jeeejeee/gpt-oss-20b-lora-adapter-text2sql") + return hf_api().snapshot_download( + repo_id="jeeejeee/gpt-oss-20b-lora-adapter-text2sql" + ) @pytest.fixture(scope="session") def qwen3moe_lora_files(): - return snapshot_download(repo_id="jeeejeee/qwen3-moe-text2sql-spider") + return hf_api().snapshot_download(repo_id="jeeejeee/qwen3-moe-text2sql-spider") @pytest.fixture(scope="session") def olmoe_lora_files(): - return snapshot_download(repo_id="jeeejeee/olmoe-instruct-text2sql-spider") + return hf_api().snapshot_download(repo_id="jeeejeee/olmoe-instruct-text2sql-spider") @pytest.fixture(scope="session") def qwen3_lora_files(): - return snapshot_download(repo_id="charent/self_cognition_Alice") + return hf_api().snapshot_download(repo_id="charent/self_cognition_Alice") @pytest.fixture(scope="session") @@ -295,32 +305,40 @@ def llama32_lora_huggingface_id(): @pytest.fixture(scope="session") def llama32_lora_files(llama32_lora_huggingface_id): - return snapshot_download(repo_id=llama32_lora_huggingface_id) + return hf_api().snapshot_download(repo_id=llama32_lora_huggingface_id) @pytest.fixture(scope="session") def whisper_lora_files(): - return snapshot_download(repo_id="chengyili2005/whisper-small-mandarin-lora") + return hf_api().snapshot_download( + repo_id="chengyili2005/whisper-small-mandarin-lora" + ) @pytest.fixture(scope="session") def qwen35_text_lora_files(): - return snapshot_download(repo_id="jeeejeee/qwen35-4b-text-only-sql-lora") + return hf_api().snapshot_download(repo_id="jeeejeee/qwen35-4b-text-only-sql-lora") @pytest.fixture(scope="session") def qwen35_vl_lora_files(): - return snapshot_download(repo_id="jeeejeee/qwen35-4b-all-linear-pokemon-lora") + return hf_api().snapshot_download( + repo_id="jeeejeee/qwen35-4b-all-linear-pokemon-lora" + ) @pytest.fixture(scope="session") def qwen36_moe_2d_lora_files(): - return snapshot_download(repo_id="jeeejeee/qwen36-35ba3b-2d-weights-poken-lora") + return hf_api().snapshot_download( + repo_id="jeeejeee/qwen36-35ba3b-2d-weights-poken-lora" + ) @pytest.fixture(scope="session") def qwen36_moe_3d_lora_files(): - return snapshot_download(repo_id="jeeejeee/qwen36-35ba3b-moe-all-linear-poken-lora") + return hf_api().snapshot_download( + repo_id="jeeejeee/qwen36-35ba3b-moe-all-linear-poken-lora" + ) @pytest.fixture diff --git a/tests/lora/test_default_mm_loras.py b/tests/lora/test_default_mm_loras.py index 19c910e2453..5e3c762f79b 100644 --- a/tests/lora/test_default_mm_loras.py +++ b/tests/lora/test_default_mm_loras.py @@ -8,15 +8,15 @@ import os import unittest.mock as mock import pytest -from huggingface_hub import snapshot_download from vllm.lora.request import LoRARequest from vllm.platforms import current_platform +from vllm.transformers_utils.repo_utils import hf_api from ..conftest import AudioTestAssets, VllmRunner from ..utils import create_new_process_for_each_test -MODEL_PATH = snapshot_download("microsoft/Phi-4-multimodal-instruct") +MODEL_PATH = hf_api().snapshot_download("microsoft/Phi-4-multimodal-instruct") AUDIO_LORA_PATH = os.path.join(MODEL_PATH, "speech-lora") IMAGE_LORA_PATH = os.path.join(MODEL_PATH, "vision-lora") diff --git a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_s3.py b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_s3.py index d60c9ba64cb..d7abc80edbf 100644 --- a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_s3.py +++ b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_s3.py @@ -3,10 +3,10 @@ from pathlib import Path -from huggingface_hub import snapshot_download from runai_model_streamer.safetensors_streamer.streamer_mock import StreamerPatcher from vllm.engine.arg_utils import EngineArgs +from vllm.transformers_utils.repo_utils import hf_api from .conftest import RunaiDummyExecutor @@ -25,7 +25,7 @@ def test_runai_model_loader_download_files_s3_mocked_with_patch( # Download model from HF mock_model_dir = f"{tmp_path}/gpt2" - snapshot_download(repo_id=test_model, local_dir=mock_model_dir) + hf_api().snapshot_download(repo_id=test_model, local_dir=mock_model_dir) monkeypatch.setattr( "vllm.transformers_utils.runai_utils.runai_list_safetensors", diff --git a/tests/model_executor/model_loader/test_sharded_state_loader.py b/tests/model_executor/model_loader/test_sharded_state_loader.py index a0b5a2a4aec..48672d4314d 100644 --- a/tests/model_executor/model_loader/test_sharded_state_loader.py +++ b/tests/model_executor/model_loader/test_sharded_state_loader.py @@ -9,11 +9,11 @@ from tempfile import TemporaryDirectory import pytest import torch -from huggingface_hub import snapshot_download from vllm import LLM, SamplingParams from vllm.model_executor.model_loader import ShardedStateLoader from vllm.platforms import current_platform +from vllm.transformers_utils.repo_utils import hf_api prompts = [ "Hello, my name is", @@ -52,7 +52,7 @@ def test_filter_subtensors(): @pytest.fixture(scope="module") def llama_3p2_1b_files(): - input_dir = snapshot_download( + input_dir = hf_api().snapshot_download( "meta-llama/Llama-3.2-1B-Instruct", ignore_patterns=["*.bin*", "original/*"] ) diff --git a/tests/models/language/pooling/test_colbert.py b/tests/models/language/pooling/test_colbert.py index 6c82ad8a9ca..810d7064d75 100644 --- a/tests/models/language/pooling/test_colbert.py +++ b/tests/models/language/pooling/test_colbert.py @@ -126,10 +126,11 @@ def _load_hf_model(model_name: str, hf_spec: dict, device: torch.device): def _load_projection_weight(model_name: str, hf_spec: dict, device: torch.device): """Download and return the ColBERT linear projection weight.""" - from huggingface_hub import hf_hub_download from safetensors.torch import load_file - path = hf_hub_download(model_name, filename=hf_spec["weights_file"]) + from vllm.transformers_utils.repo_utils import hf_api + + path = hf_api().hf_hub_download(model_name, filename=hf_spec["weights_file"]) weights = load_file(path) return weights[hf_spec["weights_key"]].to(device) diff --git a/tests/models/multimodal/generation/test_phi4mm.py b/tests/models/multimodal/generation/test_phi4mm.py index 1a4fb35a28a..5ab75e145ae 100644 --- a/tests/models/multimodal/generation/test_phi4mm.py +++ b/tests/models/multimodal/generation/test_phi4mm.py @@ -6,7 +6,6 @@ from collections.abc import Sequence import pytest import regex as re -from huggingface_hub import snapshot_download from transformers import AutoTokenizer from vllm.assets.image import ImageAsset @@ -14,6 +13,7 @@ from vllm.logprobs import SampleLogprobs from vllm.lora.request import LoRARequest from vllm.multimodal.image import convert_image_mode, rescale_image_size from vllm.multimodal.media.audio import load_audio +from vllm.transformers_utils.repo_utils import hf_api from ....conftest import ( IMAGE_ASSETS, @@ -35,7 +35,7 @@ HF_MULTIIMAGE_IMAGE_PROMPT = ( "<|user|>\n<|image_1|>\n<|image_2|>\nDescribe these images.<|end|>\n<|assistant|>\n" # noqa: E501 ) -model_path = snapshot_download("microsoft/Phi-4-multimodal-instruct") +model_path = hf_api().snapshot_download("microsoft/Phi-4-multimodal-instruct") # Since the vision-lora and speech-lora co-exist with the base model, # we have to manually specify the path of the lora weights. vision_lora_path = os.path.join(model_path, "vision-lora") diff --git a/tests/models/multimodal/pooling/test_intern_vit.py b/tests/models/multimodal/pooling/test_intern_vit.py index d7b67b8bdb6..72d2a7018bf 100644 --- a/tests/models/multimodal/pooling/test_intern_vit.py +++ b/tests/models/multimodal/pooling/test_intern_vit.py @@ -3,11 +3,11 @@ import pytest import torch import torch.nn as nn -from huggingface_hub import snapshot_download from transformers import AutoConfig, AutoModel, CLIPImageProcessor from vllm.distributed import cleanup_dist_env_and_memory from vllm.platforms import current_platform +from vllm.transformers_utils.repo_utils import hf_api from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from ....conftest import ImageTestAssets @@ -31,7 +31,7 @@ def run_intern_vit_test( *, dtype: str, ): - model = snapshot_download(model_id, allow_patterns=DOWNLOAD_PATTERN) + model = hf_api().snapshot_download(model_id, allow_patterns=DOWNLOAD_PATTERN) torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[dtype] img_processor = CLIPImageProcessor.from_pretrained(model) diff --git a/tests/models/multimodal/pooling/test_radio.py b/tests/models/multimodal/pooling/test_radio.py index fcab077fbba..7c039ab8b9c 100644 --- a/tests/models/multimodal/pooling/test_radio.py +++ b/tests/models/multimodal/pooling/test_radio.py @@ -3,13 +3,13 @@ import pytest import torch import torch.nn as nn -from huggingface_hub import snapshot_download from transformers import AutoConfig, AutoModel, CLIPImageProcessor from vllm.distributed import cleanup_dist_env_and_memory from vllm.model_executor.models.radio import RadioModel from vllm.platforms import current_platform from vllm.transformers_utils.configs.radio import RadioConfig +from vllm.transformers_utils.repo_utils import hf_api from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from ....conftest import ImageTestAssets @@ -28,7 +28,7 @@ def run_radio_test( *, dtype: str, ): - model = snapshot_download(model_id, allow_patterns=DOWNLOAD_PATTERN) + model = hf_api().snapshot_download(model_id, allow_patterns=DOWNLOAD_PATTERN) torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[dtype] img_processor = CLIPImageProcessor.from_pretrained(model) diff --git a/tests/models/quantization/test_gpt_oss.py b/tests/models/quantization/test_gpt_oss.py index 783f1773d21..f7a7b98c3ea 100644 --- a/tests/models/quantization/test_gpt_oss.py +++ b/tests/models/quantization/test_gpt_oss.py @@ -22,6 +22,7 @@ import pytest from packaging import version from vllm.platforms import current_platform +from vllm.transformers_utils.repo_utils import hf_api if current_platform.is_rocm(): from vllm.platforms.rocm import on_gfx950 @@ -47,7 +48,7 @@ QUARK_MXFP4_AVAILABLE = importlib.util.find_spec("quark") is not None and versio def has_huggingface_access(repo): try: - huggingface_hub.list_repo_refs(repo) + hf_api().list_repo_refs(repo) return True except huggingface_hub.errors.RepositoryNotFoundError: return False diff --git a/tests/plugins_tests/gguf/test_gguf_plugin_multimodal.py b/tests/plugins_tests/gguf/test_gguf_plugin_multimodal.py index cc7a021e981..56930a7d0bc 100644 --- a/tests/plugins_tests/gguf/test_gguf_plugin_multimodal.py +++ b/tests/plugins_tests/gguf/test_gguf_plugin_multimodal.py @@ -8,12 +8,12 @@ os.environ["TOKENIZERS_PARALLELISM"] = "true" from typing import Any, NamedTuple import pytest -from huggingface_hub import hf_hub_download from pytest import MarkDecorator from transformers import AutoModelForImageTextToText from vllm.assets.image import ImageAsset from vllm.multimodal.image import rescale_image_size +from vllm.transformers_utils.repo_utils import hf_api from vllm.utils.torch_utils import set_default_torch_num_threads from ...conftest import IMAGE_ASSETS, HfRunner, VllmRunner @@ -33,8 +33,8 @@ class GGUFMMTestConfig(NamedTuple): @property def gguf_model(self): - hf_hub_download(self.gguf_repo, filename=self.gguf_mmproj) - return hf_hub_download(self.gguf_repo, filename=self.gguf_backbone) + hf_api().hf_hub_download(self.gguf_repo, filename=self.gguf_mmproj) + return hf_api().hf_hub_download(self.gguf_repo, filename=self.gguf_backbone) # Common prompts aligned with test_common.py "gemma3" entry format diff --git a/tests/plugins_tests/lora_resolvers/test_filesystem_resolver.py b/tests/plugins_tests/lora_resolvers/test_filesystem_resolver.py index d4adf6f84cf..4ea8d8605c1 100644 --- a/tests/plugins_tests/lora_resolvers/test_filesystem_resolver.py +++ b/tests/plugins_tests/lora_resolvers/test_filesystem_resolver.py @@ -4,9 +4,9 @@ import os import shutil import pytest -from huggingface_hub import snapshot_download from vllm.plugins.lora_resolvers.filesystem_resolver import FilesystemResolver +from vllm.transformers_utils.repo_utils import hf_api MODEL_NAME = "Qwen/Qwen3-0.6B" LORA_NAME = "charent/self_cognition_Alice" @@ -22,12 +22,12 @@ def adapter_cache(request, tmpdir_factory): @pytest.fixture(scope="module") def qwen3_lora_files(): - return snapshot_download(repo_id=LORA_NAME) + return hf_api().snapshot_download(repo_id=LORA_NAME) @pytest.fixture(scope="module") def pa_files(): - return snapshot_download(repo_id=PA_NAME) + return hf_api().snapshot_download(repo_id=PA_NAME) @pytest.mark.asyncio diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index d2541dcc910..c8f8a429e37 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -42,12 +42,12 @@ def _skip(msg: str) -> NoReturn: def _snapshot_download_or_skip(model_id: str) -> str: try: - from huggingface_hub import snapshot_download + from vllm.transformers_utils.repo_utils import hf_api except Exception as e: # pragma: no cover _skip(f"huggingface_hub is required to download {model_id}: {e}") try: - return snapshot_download( + return hf_api().snapshot_download( repo_id=model_id, repo_type="model", # These checkpoints are already small; download full repo for simplicity. diff --git a/tests/quantization/test_quark.py b/tests/quantization/test_quark.py index 1d531fcfbfe..b31b8580e34 100644 --- a/tests/quantization/test_quark.py +++ b/tests/quantization/test_quark.py @@ -30,6 +30,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( is_layer_skipped, ) from vllm.platforms import current_platform +from vllm.transformers_utils.repo_utils import hf_api if current_platform.is_rocm(): from vllm.platforms.rocm import on_gfx942, on_gfx950 @@ -59,7 +60,7 @@ if QUARK_MXFP4_AVAILABLE: from quark.torch.quantization.config.config import FP4PerGroupSpec try: - huggingface_hub.list_repo_refs( + hf_api().list_repo_refs( "amd/Llama-3.3-70B-Instruct-WMXFP4-AMXFP4-KVFP8-Scale-UINT8-SQ" ) HF_HUB_AMD_ORG_ACCESS = True diff --git a/tests/tool_use/conftest.py b/tests/tool_use/conftest.py index ff9cdeeb737..909e12ddc82 100644 --- a/tests/tool_use/conftest.py +++ b/tests/tool_use/conftest.py @@ -3,10 +3,10 @@ import pytest import pytest_asyncio -from huggingface_hub import snapshot_download from tests.utils import RemoteOpenAIServer from vllm.platforms import current_platform +from vllm.transformers_utils.repo_utils import hf_api from .utils import ARGS, CONFIGS, ServerConfig @@ -47,7 +47,7 @@ def server_config(request): ) # download model and tokenizer using transformers - snapshot_download(config["model"]) + hf_api().snapshot_download(config["model"]) yield CONFIGS[request.param] diff --git a/tests/tool_use/mistral/conftest.py b/tests/tool_use/mistral/conftest.py index 9b0a6eb27fc..984b3a51416 100644 --- a/tests/tool_use/mistral/conftest.py +++ b/tests/tool_use/mistral/conftest.py @@ -3,10 +3,10 @@ import pytest import pytest_asyncio -from huggingface_hub import snapshot_download from tests.utils import RemoteOpenAIServer from vllm.platforms import current_platform +from vllm.transformers_utils.repo_utils import hf_api from .utils import ARGS, CONFIGS, ServerConfig @@ -22,7 +22,7 @@ def server_config(request): ) # download model and tokenizer using transformers - snapshot_download(config["model"]) + hf_api().snapshot_download(config["model"]) yield CONFIGS[request.param] diff --git a/tests/utils.py b/tests/utils.py index 645e988dbb8..90f12d444e2 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -33,7 +33,6 @@ import pytest import requests import torch import torch.nn.functional as F -from huggingface_hub import hf_hub_download from huggingface_hub.constants import HF_HUB_OFFLINE from openai.types.completion import Completion from typing_extensions import ParamSpec @@ -57,6 +56,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( from vllm.model_executor.model_loader import get_model_loader from vllm.platforms import current_platform from vllm.tokenizers import get_tokenizer +from vllm.transformers_utils.repo_utils import hf_api from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.utils.mem_constants import GB_bytes from vllm.utils.network_utils import get_open_port @@ -77,7 +77,7 @@ def prewarm_hf_cache(assets: list[tuple[str, str]]) -> None: return for repo_id, filename in assets: try: - hf_hub_download(repo_id=repo_id, filename=filename) + hf_api().hf_hub_download(repo_id=repo_id, filename=filename) except Exception as e: logger.warning( "Failed to prefetch %s/%s: %r. Tests depending on this asset may fail.", diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index a2fc173f035..a788cecc6ce 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/tools/pre_commit/check_forbidden_imports.py @@ -6,6 +6,13 @@ from dataclasses import dataclass, field import regex as re +# Hub entry points that must go through the vLLM-tagged repo_utils helpers. +_HF_NAMES = ( + r"HfApi|HfFileSystem|hf_hub_download|snapshot_download" + r"|list_repo_files|file_exists|try_to_load_from_cache" + r"|list_repo_refs|repo_exists" +) + @dataclass class ForbiddenImport: @@ -13,6 +20,7 @@ class ForbiddenImport: tip: str allowed_pattern: re.Pattern = re.compile(r"^$") # matches nothing by default allowed_files: set[str] = field(default_factory=set) + allowed_dirs: set[str] = field(default_factory=set) CHECK_IMPORTS = { @@ -83,6 +91,23 @@ CHECK_IMPORTS = { ), allowed_files={"vllm/triton_utils/importing.py"}, ), + "huggingface_hub repo API": ForbiddenImport( + # Catch `from huggingface_hub import `, including parenthesized, + # multi-line imports. + pattern=( + r"^\s*from\s+huggingface_hub\s+import\s*\([^)]*\b(?:" + _HF_NAMES + r")\b" + r"|" + r"^\s*from\s+huggingface_hub\s+import\b[^\n]*\b(?:" + _HF_NAMES + r")\b" + ), + tip=( + "Use the shared, vLLM-tagged helpers from " + "vllm.transformers_utils.repo_utils (e.g. hf_api(), hf_fs(), " + "list_repo_files, file_exists) instead of calling " + "huggingface_hub directly." + ), + allowed_files={"vllm/transformers_utils/repo_utils.py"}, + allowed_dirs={"examples/"}, + ), } @@ -95,11 +120,18 @@ def check_file(path: str) -> int: # Skip files that are allowed for this import if path in forbidden_import.allowed_files: continue + # Skip directories that are allowed for this import + if any(path.startswith(prefix) for prefix in forbidden_import.allowed_dirs): + continue # Search for forbidden imports for match in re.finditer(forbidden_import.pattern, content, re.MULTILINE): # Check if it's allowed if forbidden_import.allowed_pattern.match(match.group()): continue + # Skip matches inside a comment + line_start = content.rfind("\n", 0, match.start()) + 1 + if "#" in content[line_start : match.start()]: + continue # Calculate line number from match position line_num = content[: match.start() + 1].count("\n") + 1 print( @@ -119,7 +151,10 @@ def main(): def test_regex(): - test_cases = [ + def matches(rule: str, content: str) -> bool: + return bool(re.search(CHECK_IMPORTS[rule].pattern, content, re.MULTILINE)) + + pickle_cases = [ # Should match ("import pickle", True), ("import cloudpickle", True), @@ -139,11 +174,47 @@ def test_regex(): ("print('import pickle')", False), ("import pickleas as asdf", False), ] - for i, (line, should_match) in enumerate(test_cases): - result = bool(CHECK_IMPORTS["pickle/cloudpickle"].pattern.match(line)) + for i, (content, should_match) in enumerate(pickle_cases): + result = matches("pickle/cloudpickle", content) assert result == should_match, ( - f"Test case {i} failed: '{line}' (expected {should_match}, got {result})" + f"pickle case {i} failed: {content!r} " + f"(expected {should_match}, got {result})" ) + + hf_cases = [ + # Should match + ("from huggingface_hub import snapshot_download", True), + ("from huggingface_hub import hf_hub_download", True), + ("from huggingface_hub import HfApi", True), + ("from huggingface_hub import HfFileSystem", True), + ("from huggingface_hub import list_repo_files", True), + ("from huggingface_hub import try_to_load_from_cache", True), + (" from huggingface_hub import snapshot_download", True), + ("from huggingface_hub import PyTorchModelHubMixin, hf_hub_download", True), + ("from huggingface_hub import (snapshot_download)", True), + # Parenthesized multi-line import must not bypass the hook + ("from huggingface_hub import (\n snapshot_download,\n)", True), + ( + "from huggingface_hub import (\n PyTorchModelHubMixin,\n HfApi,\n)", + True, + ), + # Should not match + ("import huggingface_hub", False), + ("import huggingface_hub as hf", False), + ("from huggingface_hub import PyTorchModelHubMixin", False), + ("from huggingface_hub.constants import HF_HUB_CACHE", False), + ("from huggingface_hub.utils import EntryNotFoundError", False), + ("from vllm.transformers_utils.repo_utils import hf_api", False), + ("from huggingface_hub import (\n PyTorchModelHubMixin,\n)", False), + ("# from huggingface_hub import snapshot_download", False), + ] + for i, (content, should_match) in enumerate(hf_cases): + result = matches("huggingface_hub repo API", content) + assert result == should_match, ( + f"huggingface_hub case {i} failed: {content!r} " + f"(expected {should_match}, got {result})" + ) + print("All regex tests passed.") diff --git a/vllm/model_executor/models/llava_onevision2.py b/vllm/model_executor/models/llava_onevision2.py index 552e2a9e8f8..e17398cc8b7 100644 --- a/vllm/model_executor/models/llava_onevision2.py +++ b/vllm/model_executor/models/llava_onevision2.py @@ -24,8 +24,6 @@ from __future__ import annotations import hashlib import importlib -import json -import os from collections.abc import Callable, Iterable, Mapping, Sequence from functools import lru_cache from typing import ( @@ -39,7 +37,6 @@ import regex as re import torch import torch.nn as nn import torch.nn.functional as F -from huggingface_hub import hf_hub_download from PIL import Image from transformers import AutoProcessor, AutoTokenizer, BatchFeature from transformers.dynamic_module_utils import get_class_from_dynamic_module @@ -108,6 +105,7 @@ from vllm.multimodal.video import ( ) from vllm.sequence import IntermediateTensors from vllm.transformers_utils.processor import _merge_mm_kwargs +from vllm.transformers_utils.repo_utils import get_hf_file_to_dict from vllm.transformers_utils.utils import convert_model_repo_to_path from vllm.utils.tensor_schema import TensorSchema, TensorShape @@ -165,13 +163,10 @@ def _load_ov2_processor( # codec video backend keeps its configured defaults. codec_config: dict = {} try: - config_file = os.path.join(path, "preprocessor_config.json") - if not os.path.isfile(config_file): - config_file = hf_hub_download( - path, "preprocessor_config.json", revision=revision - ) - with open(config_file, encoding="utf-8") as f: - codec_config = json.load(f).get("codec", {}) or {} + preprocessor_config = get_hf_file_to_dict( + "preprocessor_config.json", path, revision + ) + codec_config = (preprocessor_config or {}).get("codec", {}) or {} except Exception: logger.debug("OV2: no codec defaults found in preprocessor_config.json") From b9b6306ebe0f5fa5046553667192e56bb3934af0 Mon Sep 17 00:00:00 2001 From: Harshal Janjani Date: Sat, 25 Jul 2026 15:20:58 +0400 Subject: [PATCH 035/185] =?UTF-8?q?feat[vLLM=20=C3=97=20v5]:=20Add=20audio?= =?UTF-8?q?=20support=20for=20the=20Transformers=20backend=20(#39330)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Harshal Janjani Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- docs/models/supported_models.md | 5 +- .../generation/test_transformers_audio.py | 133 +++++ .../multimodal/processing/test_common.py | 8 + .../processing/test_transformers_audio.py | 138 +++++ ...sformers.py => test_transformers_image.py} | 23 + tests/models/registry.py | 4 + vllm/model_executor/models/registry.py | 4 + .../models/transformers/multimodal.py | 519 +++++++++++++----- 8 files changed, 682 insertions(+), 152 deletions(-) create mode 100644 tests/models/multimodal/generation/test_transformers_audio.py create mode 100644 tests/models/multimodal/processing/test_transformers_audio.py rename tests/models/multimodal/processing/{test_transformers.py => test_transformers_image.py} (63%) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 3ee942c175e..566a3e1f955 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -19,7 +19,7 @@ vLLM also supports model implementations that are available in Transformers. We Currently, the Transformers modeling backend works for the following: -- Modalities: embedding models, language models and vision-language models* +- Modalities: embedding models, language models, vision-language models* and audio-language models - Architectures: encoder-only, decoder-only, mixture-of-experts - Attention types: full attention and/or sliding attention @@ -606,7 +606,8 @@ Some models are supported only via the [Transformers modeling backend](#transfor | Architecture | Models | Inputs | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) | | ------------ | ------ | ------ | ----------------- | --------------------------- | --------------------------------------- | -| `Emu3ForConditionalGeneration` | Emu3 | T + I | `BAAI/Emu3-Chat-hf` | ✅︎ | ✅︎ | +| `Emu3ForConditionalGeneration` | Emu3 | T + I+ | `BAAI/Emu3-Chat-hf` | ✅︎ | ✅︎ | +| `VibeVoiceAsrForConditionalGeneration` | VibeVoice-ASR | T + A+ | `microsoft/VibeVoice-ASR-HF` | ✅︎ | ✅︎ | ^ You need to set the architecture name via `--hf-overrides` to match the one in vLLM.
E Pre-computed embeddings can be inputted for this modality.
diff --git a/tests/models/multimodal/generation/test_transformers_audio.py b/tests/models/multimodal/generation/test_transformers_audio.py new file mode 100644 index 00000000000..91955798779 --- /dev/null +++ b/tests/models/multimodal/generation/test_transformers_audio.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +import pytest +from transformers import AutoModelForSeq2SeqLM + +from vllm.assets.audio import AudioAsset +from vllm.envs import disable_envs_cache +from vllm.lora.request import LoRARequest +from vllm.multimodal.audio import AudioResampler + +from ....conftest import HfRunner, VllmRunner +from ...utils import check_logprobs_close + +AUDIO_ASSET = AudioAsset("mary_had_lamb") + +AUDIO_MODEL_SETTINGS: dict[str, dict[str, Any]] = { + "ibm-granite/granite-speech-3.3-2b": { + "prompt": ( + "<|start_of_role|>system<|end_of_role|>" + "You are a helpful AI assistant<|end_of_text|>\n" + "<|start_of_role|>user<|end_of_role|>" + "<|audio|>can you transcribe the speech into a written format?" + "<|end_of_text|>\n" + "<|start_of_role|>assistant<|end_of_role|>" + ), + "audio_lora_path": "ibm-granite/granite-speech-3.3-2b", + }, + "nvidia/audio-flamingo-3-hf": { + "prompt": ( + "<|im_start|>system\n" + "You are a helpful assistant.<|im_end|>\n" + "<|im_start|>user\n" + "Transcribe the input speech.<|im_end|>\n" + "<|im_start|>assistant\n" + ), + "vllm_runner_kwargs": { + "gpu_memory_utilization": 0.85, + }, + }, + "microsoft/VibeVoice-ASR-HF": { + "prompt": ( + "<|im_start|>system\n" + "You are a helpful assistant that transcribes audio input " + "into text output in JSON format.<|im_end|>\n" + "<|im_start|>user\n" + "<|object_ref_start|><|box_start|><|object_ref_end|>\n" + "This is a 1.0 seconds audio, please transcribe it with " + "these keys: Start time, End time, Speaker ID, Content" + "<|im_end|>\n" + "<|im_start|>assistant\n" + ), + "sampling_rate": 24000, + "vllm_runner_kwargs": { + "max_num_batched_tokens": 2048, + "gpu_memory_utilization": 0.85, + }, + }, + "zai-org/GLM-ASR-Nano-2512": { + "prompt": ( + "<|user|>\n" + "<|begin_of_audio|><|pad|><|end_of_audio|><|user|>\n" + "Please transcribe this audio into text" + "<|assistant|>\n" + ), + }, +} + + +@pytest.mark.parametrize("model_id", list(AUDIO_MODEL_SETTINGS)) +def test_transformers_audio_generation( + hf_runner: type[HfRunner], + vllm_runner: type[VllmRunner], + monkeypatch, + model_id: str, +): + """Single-process workaround for V1 fork safety deadlock issue + (vllm-project/vllm/issues/17676). Running multiple audio models together + under pytest can cause (possibly flaky) hangs, so they are grouped under + the same config. Using VLLM_WORKER_MULTIPROC_METHOD=spawn avoids the + deadlock and allows worker processes to terminate cleanly, and release + GPU memory between test runs until the issue is fixed.""" + # TODO: Remove monkeypatch once + # https://github.com/vllm-project/vllm/issues/17676 is fixed. + disable_envs_cache() + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + + settings = AUDIO_MODEL_SETTINGS[model_id] + audio_lora_path = settings.get("audio_lora_path") + + audio, orig_sr = AUDIO_ASSET.audio_and_sample_rate + target_sr = settings.get("sampling_rate", orig_sr) + if orig_sr != target_sr: + audio = AudioResampler(target_sr=target_sr).resample(audio, orig_sr=orig_sr) + audio = (audio, target_sr) + + with vllm_runner( + model_id, + model_impl="transformers", + dtype="bfloat16", + max_model_len=2048, + enforce_eager=True, + limit_mm_per_prompt={"audio": 1}, + enable_lora=audio_lora_path is not None, + max_lora_rank=64, + **settings.get("vllm_runner_kwargs", {}), + ) as vllm_model: + lora_request = ( + LoRARequest("audio", 1, audio_lora_path) if audio_lora_path else None + ) + vllm_outputs = vllm_model.generate_greedy_logprobs( + [settings["prompt"]], + 128, + num_logprobs=10, + audios=[audio], + lora_request=lora_request, + ) + + with hf_runner( + model_id, dtype="bfloat16", auto_cls=AutoModelForSeq2SeqLM + ) as hf_model: + hf_outputs = hf_model.generate_greedy_logprobs_limit( + [settings["prompt"]], 128, num_logprobs=10, audios=[audio] + ) + + check_logprobs_close( + outputs_0_lst=hf_outputs, + outputs_1_lst=vllm_outputs, + name_0="hf", + name_1="vllm", + ) diff --git a/tests/models/multimodal/processing/test_common.py b/tests/models/multimodal/processing/test_common.py index 1ef39cfaa7b..9b586be4797 100644 --- a/tests/models/multimodal/processing/test_common.py +++ b/tests/models/multimodal/processing/test_common.py @@ -452,6 +452,14 @@ def test_processing_correctness( "audio placeholders from processed audio lengths. Its vLLM " "processor paths are covered by test_moss_audio.py." ) + # TODO: Remove when transformers 5.15.0 is released, which contains + # https://github.com/huggingface/transformers/pull/47483. + if model_id == "microsoft/VibeVoice-ASR-HF": + pytest.skip( + "VibeVoice ASR requires audio as a positional argument and hence " + "cannot pass the processing correctness test as is. Its generation " + "is covered by test_transformers_audio.py." + ) if model_id == "lmms-lab-encoder/LLaVA-OneVision-2-8B-Instruct": pytest.skip( "LLaVA-OneVision-2 video processing routes frames through custom " diff --git a/tests/models/multimodal/processing/test_transformers_audio.py b/tests/models/multimodal/processing/test_transformers_audio.py new file mode 100644 index 00000000000..3ff3e04379f --- /dev/null +++ b/tests/models/multimodal/processing/test_transformers_audio.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import numpy as np +import pytest + +from vllm.config import ModelConfig +from vllm.multimodal import MULTIMODAL_REGISTRY + +AUDIO_MODEL_SETTINGS = { + "ibm-granite/granite-speech-3.3-2b": { + "prompt": ( + "<|start_of_role|>system<|end_of_role|>" + "You are a helpful AI assistant<|end_of_text|>\n" + "<|start_of_role|>user<|end_of_role|>" + "<|audio|>can you transcribe the speech into a written format?" + "<|end_of_text|>\n" + "<|start_of_role|>assistant<|end_of_role|>" + ), + }, + "nvidia/audio-flamingo-3-hf": { + "prompt": ( + "<|im_start|>system\n" + "You are a helpful assistant.<|im_end|>\n" + "<|im_start|>user\n" + "Transcribe the input speech.<|im_end|>\n" + "<|im_start|>assistant\n" + ), + }, + "mistralai/Voxtral-Mini-3B-2507": { + "prompt": ("[INST][AUDIO]What can you tell me about this audio?[/INST]"), + }, + "microsoft/VibeVoice-ASR-HF": { + "prompt": ( + "<|im_start|>system\n" + "You are a helpful assistant that transcribes audio input " + "into text output in JSON format.<|im_end|>\n" + "<|im_start|>user\n" + "<|object_ref_start|><|box_start|><|object_ref_end|>\n" + "This is a 1.0 seconds audio, please transcribe it with " + "these keys: Start time, End time, Speaker ID, Content" + "<|im_end|>\n" + "<|im_start|>assistant\n" + ), + }, + "zai-org/GLM-ASR-Nano-2512": { + "prompt": ( + "<|user|>\n" + "<|begin_of_audio|><|pad|><|end_of_audio|><|user|>\n" + "Please transcribe this audio into text" + "<|assistant|>\n" + ), + }, +} + + +@pytest.mark.parametrize( + "model_id", + [ + "ibm-granite/granite-speech-3.3-2b", + "nvidia/audio-flamingo-3-hf", + pytest.param( + "mistralai/Voxtral-Mini-3B-2507", + marks=pytest.mark.xfail( + reason="MistralCommonBackend.encode does not produce the audio " + "placeholder token (ID 24) from raw text. apply_chat_template " + "yields token IDs with placeholders, but MultiModalProcessor." + "apply() decodes the prompt back to text and re-tokenizes, at " + "which point the placeholders are lost. Fix belongs in " + "mistral_common or in the Voxtral-specific path.", + strict=False, + ), + ), + "microsoft/VibeVoice-ASR-HF", + "zai-org/GLM-ASR-Nano-2512", + ], +) +def test_audio_multimodal_processor(model_id): + settings = AUDIO_MODEL_SETTINGS[model_id] + + model_config = ModelConfig( + model=model_id, + model_impl="transformers", + ) + + mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config) + + audio = np.zeros(16000, dtype=np.float32) + mm_data = {"audio": (audio, 16000)} + + result = mm_processor( + prompt=settings["prompt"], + mm_items=mm_processor.info.parse_mm_data(mm_data), + hf_processor_mm_kwargs={}, + ) + + assert "prompt_token_ids" in result + assert len(result["prompt_token_ids"]) > 0 + + mm_placeholders = result.get("mm_placeholders", {}) + assert "audio" in mm_placeholders, f"No audio placeholders found for {model_id}" + assert len(mm_placeholders["audio"]) == 1 + + placeholder = mm_placeholders["audio"][0] + assert placeholder.length > 0 + assert placeholder.offset >= 0 + + audio_items = result.get("mm_kwargs", {}).get("audio", []) + assert len(audio_items) == 1, f"Expected 1 audio item, got {len(audio_items)}" + item_keys = list(audio_items[0].keys()) + has_features = "input_features" in item_keys or "input_values" in item_keys + assert has_features, ( + f"No audio features (input_features/input_values) in {item_keys} for {model_id}" + ) + + +def test_audio_multiple_inputs(): + """Multiple audios per prompt are each detected as a separate placeholder + and multi-modal item by the Transformers backend.""" + model_id = "ibm-granite/granite-speech-3.3-2b" + model_config = ModelConfig(model=model_id, model_impl="transformers") + mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config) + + audio_token = mm_processor.info.get_hf_processor().audio_token + # One token per audio; the processor expands each to its placeholder run. + prompt = ( + "<|start_of_role|>user<|end_of_role|>" + f"{audio_token} and {audio_token} transcribe<|end_of_text|>\n" + ) + audios = [np.zeros(16000, dtype=np.float32), np.zeros(24000, dtype=np.float32)] + + result = mm_processor( + prompt=prompt, + mm_items=mm_processor.info.parse_mm_data({"audio": audios}), + hf_processor_mm_kwargs={}, + ) + + assert len(result["mm_placeholders"]["audio"]) == 2 + assert len(result["mm_kwargs"]["audio"]) == 2 diff --git a/tests/models/multimodal/processing/test_transformers.py b/tests/models/multimodal/processing/test_transformers_image.py similarity index 63% rename from tests/models/multimodal/processing/test_transformers.py rename to tests/models/multimodal/processing/test_transformers_image.py index a556b8f10af..2c31bcc6347 100644 --- a/tests/models/multimodal/processing/test_transformers.py +++ b/tests/models/multimodal/processing/test_transformers_image.py @@ -54,3 +54,26 @@ def test_multimodal_processor(model_id): str_processed_inputs["prompt_token_ids"] == ids_processed_inputs["prompt_token_ids"] ) + + +def test_image_multiple_inputs(): + """Multiple images per prompt are each detected as a separate placeholder + and multi-modal item by the Transformers backend.""" + model_id = "llava-hf/llava-onevision-qwen2-0.5b-ov-hf" + model_config = ModelConfig(model=model_id, model_impl="transformers") + mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config) + + image = ImageAsset("cherry_blossom").pil_image + prompt = ( + "<|im_start|>user \n and \n" + "What do these images show?<|im_end|><|im_start|>assistant\n" + ) + + result = mm_processor( + prompt=prompt, + mm_items=mm_processor.info.parse_mm_data({"image": [image, image]}), + hf_processor_mm_kwargs={}, + ) + + assert len(result["mm_placeholders"]["image"]) == 2 + assert len(result["mm_kwargs"]["image"]) == 2 diff --git a/tests/models/registry.py b/tests/models/registry.py index 2d8ecfb560f..f55daa51033 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -1393,6 +1393,10 @@ _MULTIMODAL_EXAMPLE_MODELS = { "fixie-ai/ultravox-v0_5-llama-3_2-1b", trust_remote_code=True, ), + "VibeVoiceAsrForConditionalGeneration": _HfExamplesInfo( + "microsoft/VibeVoice-ASR-HF", + min_transformers_version="5.13.0", + ), "VoxtralForConditionalGeneration": _HfExamplesInfo( "mistralai/Voxtral-Mini-3B-2507", tokenizer_mode="mistral", diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 229cf4fe87f..c130a7f4a2c 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -664,6 +664,10 @@ _TRANSFORMERS_SUPPORTED_MODELS = { "transformers", "TransformersMultiModalForCausalLM", ), + "VibeVoiceAsrForConditionalGeneration": ( + "transformers", + "TransformersMultiModalForCausalLM", + ), } _TRANSFORMERS_BACKEND_MODELS = { diff --git a/vllm/model_executor/models/transformers/multimodal.py b/vllm/model_executor/models/transformers/multimodal.py index ae203d6bdb6..5c6b0e286a1 100644 --- a/vllm/model_executor/models/transformers/multimodal.py +++ b/vllm/model_executor/models/transformers/multimodal.py @@ -17,7 +17,8 @@ """Transformers modeling backend mixin for multi-modal models.""" from collections.abc import Mapping -from typing import TYPE_CHECKING +from contextlib import nullcontext +from typing import TYPE_CHECKING, Any import torch from transformers import AutoModel @@ -27,6 +28,7 @@ from vllm.config.utils import getattr_iter from vllm.inputs import MultiModalDataDict, MultiModalInput, mm_input from vllm.logger import init_logger from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal +from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.multimodal import MultiModalKwargsItems from vllm.multimodal.inputs import ( MultiModalFeatureSpec, @@ -36,6 +38,7 @@ from vllm.multimodal.inputs import ( from vllm.multimodal.parse import ( ImageProcessorItems, MultiModalDataItems, + MultiModalDataParser, ) from vllm.multimodal.processing import ( BaseDummyInputsBuilder, @@ -59,11 +62,90 @@ _MODALITY_TO_TOKEN_TYPE_ID = {"image": 1, "video": 2, "audio": 3} class MultiModalProcessingInfo(BaseProcessingInfo): + def _get_audio_processor(self) -> Any: + # TODO: drop feature_extractor branch once huggingface/transformers#44394 lands. + return getattr_iter( + self.get_hf_processor(), ("audio_processor", "feature_extractor") + ) + + def _is_audio_model(self) -> bool: + return self._get_audio_processor() is not None + + def _is_image_model(self) -> bool: + return hasattr(self.get_hf_processor(), "image_processor") + + def _is_video_model(self) -> bool: + return hasattr(self.get_hf_processor(), "video_processor") + + def _get_audio_token_id(self) -> int: + processor = self.get_hf_processor() + if hasattr(processor, "audio_token_id"): + return processor.audio_token_id + config = self.get_hf_config() + val = getattr_iter(config, ("audio_token_id", "audio_token_index")) + if val is not None: + return val + if hasattr(processor, "audio_token"): + tokenizer = self.get_tokenizer() + vocab = tokenizer.get_vocab() + if processor.audio_token in vocab: + return vocab[processor.audio_token] + raise ValueError("Cannot find audio_token_id on processor or model config") + + def _get_audio_sampling_rate(self) -> float: + sub = self._get_audio_processor() + if sub is not None and hasattr(sub, "sampling_rate"): + return sub.sampling_rate + return 16000.0 + + def get_data_parser(self) -> MultiModalDataParser: + if self._is_audio_model(): + return MultiModalDataParser( + target_sr=self._get_audio_sampling_rate(), + expected_hidden_size=self._get_expected_hidden_size(), + ) + return super().get_data_parser() + def get_supported_mm_limits(self): - return {"image": None} + limits = {} + if self._is_audio_model(): + limits["audio"] = None + if self._is_image_model(): + limits["image"] = None + if not limits: + raise ValueError( + f"Unable to detect a supported modality on " + f"{type(self.get_hf_processor()).__name__}. " + "Checked `_is_audio_model` and `_is_image_model`." + ) + return limits def get_mm_max_tokens_per_item(self, seq_len, mm_counts): - return {"image": self.get_max_image_tokens()} + result = {} + if self._is_audio_model(): + result["audio"] = self.get_max_audio_tokens() + if self._is_image_model(): + result["image"] = self.get_max_image_tokens() + if not result: + raise ValueError( + f"Unable to detect a supported modality on " + f"{type(self.get_hf_processor()).__name__}. " + "Checked `_is_audio_model` and `_is_image_model`." + ) + return result + + def get_max_audio_tokens(self) -> int: + config = self.get_hf_config() + audio_config_names = ("audio_config", "encoder_config") + names = ("max_source_positions", "max_position_embeddings", "max_pos_emb") + audio_config = getattr_iter(config, audio_config_names, default=config) + val = getattr_iter(audio_config, names) + if val is not None: + return int(val) + raise ValueError( + f"Unable to get max input length from {type(audio_config).__name__}. " + f"The following attribute names were checked: {names}." + ) def get_max_image_tokens(self) -> int: width, height = self.get_max_image_size() @@ -82,14 +164,19 @@ class MultiModalProcessingInfo(BaseProcessingInfo): class MultiModalDummyInputsBuilder(BaseDummyInputsBuilder[MultiModalProcessingInfo]): def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: - num_images = mm_counts.get("image", 0) - - processor = self.info.get_hf_processor() - if "gemma3" in processor.__class__.__name__.lower(): - image_token = processor.boi_token - else: - image_token = getattr(processor, "image_token", "") - return image_token * num_images + text = "" + if self.info._is_audio_model() and (num_audios := mm_counts.get("audio", 0)): + processor = self.info.get_hf_processor() + audio_token = getattr(processor, "audio_token", "") + text += audio_token * num_audios + if self.info._is_image_model() and (num_images := mm_counts.get("image", 0)): + processor = self.info.get_hf_processor() + if "gemma3" in processor.__class__.__name__.lower(): + image_token = processor.boi_token + else: + image_token = getattr(processor, "image_token", "") + text += image_token * num_images + return text def get_dummy_mm_data( self, @@ -97,20 +184,28 @@ class MultiModalDummyInputsBuilder(BaseDummyInputsBuilder[MultiModalProcessingIn mm_counts: Mapping[str, int], mm_options: Mapping[str, "BaseDummyOptions"], ) -> MultiModalDataDict: - num_images = mm_counts.get("image", 0) - - target_width, target_height = self.info.get_max_image_size() - - image_overrides = mm_options.get("image") - - return { - "image": self._get_dummy_images( + data: MultiModalDataDict = {} + if self.info._is_audio_model() and (num_audios := mm_counts.get("audio", 0)): + sampling_rate = self.info._get_audio_sampling_rate() + sub = self.info._get_audio_processor() + chunk_length = getattr(sub, "chunk_length", None) if sub else None + if chunk_length is None: + chunk_length = 30 + audio_len = int(chunk_length * sampling_rate) + data["audio"] = self._get_dummy_audios( + length=audio_len, + num_audios=num_audios, + overrides=mm_options.get("audio"), + ) + if self.info._is_image_model() and (num_images := mm_counts.get("image", 0)): + target_width, target_height = self.info.get_max_image_size() + data["image"] = self._get_dummy_images( width=target_width, height=target_height, num_images=num_images, - overrides=image_overrides, - ), - } + overrides=mm_options.get("image"), + ) + return data class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]): @@ -142,21 +237,39 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]): ) -> Mapping[str, MultiModalFieldConfig]: # HF Processors always return a mask but vLLM doesn't need it hf_inputs.pop("attention_mask", None) - num_image_patches = hf_inputs.get("num_image_patches") - mm_fields = { - key: MultiModalFieldConfig.flat_from_sizes("image", num_image_patches) - for key in hf_inputs - } - mm_fields["image_embeds"] = MultiModalFieldConfig.flat_from_sizes( - "image", num_image_patches - ) - # Keep these as batched, as they always have batch size as first dim - mm_fields["image_grid_thw"] = MultiModalFieldConfig.batched("image") - mm_fields["video_grid_thw"] = MultiModalFieldConfig.batched("image") - mm_fields["num_image_patches"] = MultiModalFieldConfig.batched( - "image", keep_on_cpu=True - ) + mm_fields: dict[str, MultiModalFieldConfig] = {} + if self.info._is_audio_model(): + num_audio_tokens = hf_inputs.get("num_audio_tokens") + mm_fields.update( + { + key: MultiModalFieldConfig.flat_from_sizes( + "audio", num_audio_tokens + ) + for key in hf_inputs + } + ) + mm_fields["num_audio_tokens"] = MultiModalFieldConfig.batched("audio") + if self.info._is_image_model(): + num_image_patches = hf_inputs.get("num_image_patches") + mm_fields.update( + { + key: MultiModalFieldConfig.flat_from_sizes( + "image", num_image_patches + ) + for key in hf_inputs + } + ) + mm_fields["image_embeds"] = MultiModalFieldConfig.flat_from_sizes( + "image", num_image_patches + ) + + # Keep these as batched, as they always have batch size as first dim + mm_fields["image_grid_thw"] = MultiModalFieldConfig.batched("image") + mm_fields["video_grid_thw"] = MultiModalFieldConfig.batched("image") + mm_fields["num_image_patches"] = MultiModalFieldConfig.batched( + "image", keep_on_cpu=True + ) return mm_fields def _get_hf_mm_data( @@ -164,13 +277,93 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]): mm_items: MultiModalDataItems, ) -> tuple[Mapping[str, object], Mapping[str, object]]: """ - In contrast to the base class, this method always adds - `return_mm_token_type_ids` to the processor data + In contrast to the base class, this method requests + `return_mm_token_type_ids` and remaps the `audios` key to `audio` for + audio models. """ processor_data, passthrough_data = super()._get_hf_mm_data(mm_items) + if self.info._is_audio_model() and "audios" in processor_data: + processor_data["audio"] = processor_data.pop("audios") processor_data["return_mm_token_type_ids"] = True return processor_data, passthrough_data + def _apply_audio( + self, + prompt_ids: list[int], + processed_data: "BatchFeature", + ) -> dict[str, list[PlaceholderRange]]: + audio_token_id = self.info._get_audio_token_id() + prompt_tensor = torch.tensor(prompt_ids) + is_audio = prompt_tensor == audio_token_id + + if not is_audio.any(): + return {} + + padded = torch.cat([torch.tensor([False]), is_audio, torch.tensor([False])]) + transitions = padded.int().diff() + starts = torch.where(transitions == 1)[0] + ends = torch.where(transitions == -1)[0] + lengths = ends - starts + + ranges = [ + PlaceholderRange( + offset=s.item(), + length=ln.item(), + is_embed=torch.ones(ln.item(), dtype=torch.bool), + ) + for s, ln in zip(starts, lengths) + ] + processed_data["num_audio_tokens"] = lengths + return {"audio": ranges} + + def _apply_vision( + self, + prompt_ids: list[int], + processed_data: "BatchFeature", + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + mm_token_type_ids: torch.Tensor | None, + ) -> dict[str, list[PlaceholderRange]]: + if mm_token_type_ids is None: + return {} + + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + + # We can infer vLLM style placeholder from token type ids, if we split + # it for each input `mm_data`. + mm_positions = torch.where(mm_token_type_ids == 1)[1] + images = mm_items.get_items("image", ImageProcessorItems) + image_sizes = [] + for item_idx in range(len(images)): + image_size = images.get_image_size(item_idx) + image_sizes.append((image_size.height, image_size.width)) + + mm_tokens_per_modality = hf_processor._get_num_multimodal_tokens( + image_sizes=image_sizes, + **self.info.ctx.get_merged_mm_kwargs({}), + ) + + mm_placeholders: dict[str, list[PlaceholderRange]] = {} + split_sizes = mm_tokens_per_modality["num_image_tokens"] + if split_sizes: + chunked_mm_positions = torch.split(mm_positions, split_sizes) + mm_tokens = torch.tensor(prompt_ids)[mm_token_type_ids[0].bool()] + chunked_mm_tokens = torch.split(mm_tokens, split_sizes) + ranges = [ + PlaceholderRange( + offset=positions[0].item(), + length=positions.shape[0], + is_embed=(mm_tokens == hf_processor.image_token_id).bool(), + ) + for positions, mm_tokens in zip(chunked_mm_positions, chunked_mm_tokens) + ] + mm_placeholders = {"image": ranges} + + processed_data["num_image_patches"] = torch.tensor( + mm_tokens_per_modality["num_image_patches"] + ) + return mm_placeholders + def apply( self, inputs: ProcessorInputs, @@ -198,8 +391,8 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]): # Bypass cached processor and always apply to the full set of mm inputs # NOTE: we can't just set caching=False because base class method # transforms outputs to `MultiModalKwargs` which is not going to - # work for Transformers. We have a lot of logic tied to - # `mm_tokens_per_modality` below + # work for Transformers. The vision path has logic tied to + # `mm_tokens_per_modality` in _apply_vision() prompt_ids, processed_data, _ = self._apply_hf_processor_text_mm( prompt_text=prompt, mm_items=mm_items, @@ -207,52 +400,33 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]): tokenization_kwargs=tokenization_kwargs, ) + # Use overrides if provided; fallback to data-dependent hashing. + with timing_ctx.record("get_mm_hashes"): + mm_hashes = inputs.get_mm_hashes(self.info.model_id) + # For gemma3 we check `token_type_ids` as the key mm_token_type_ids = processed_data.pop("token_type_ids", None) mm_token_type_ids = processed_data.pop("mm_token_type_ids", mm_token_type_ids) - # We can infer vLLM style placeholder from token type ids, if we split - # it for each input `mm_data`. - mm_positions = torch.where(mm_token_type_ids == 1)[1] - images = mm_items.get_items("image", ImageProcessorItems) - image_sizes = [] - for item_idx in range(len(images)): - image_size = images.get_image_size(item_idx) - image_sizes.append((image_size.height, image_size.width)) - - mm_tokens_per_modality = hf_processor._get_num_multimodal_tokens( - image_sizes=image_sizes, - **self.info.ctx.get_merged_mm_kwargs({}), - ) - - mm_placeholders = {} - split_sizes = mm_tokens_per_modality["num_image_tokens"] - if split_sizes: - chunked_mm_positions = torch.split(mm_positions, split_sizes) - mm_tokens = torch.tensor(prompt_ids)[mm_token_type_ids[0].bool()] - chunked_mm_tokens = torch.split(mm_tokens, split_sizes) - ranges = [ - PlaceholderRange( - offset=positions[0].item(), - length=positions.shape[0], - is_embed=(mm_tokens == hf_processor.image_token_id).bool(), + mm_placeholders: dict[str, list[PlaceholderRange]] = {} + if self.info._is_audio_model(): + mm_placeholders.update(self._apply_audio(prompt_ids, processed_data)) + if self.info._is_image_model(): + mm_placeholders.update( + self._apply_vision( + prompt_ids, + processed_data, + mm_items, + hf_processor_mm_kwargs, + mm_token_type_ids, ) - for positions, mm_tokens in zip(chunked_mm_positions, chunked_mm_tokens) - ] - mm_placeholders = {"image": ranges} + ) - processed_data["num_image_patches"] = torch.tensor( - mm_tokens_per_modality["num_image_patches"] - ) mm_kwargs = MultiModalKwargsItems.from_hf_inputs( processed_data, self._get_mm_fields_config(processed_data, hf_processor_mm_kwargs), ) - # Use overrides if provided; fallback to data-dependent hashing. - with timing_ctx.record("get_mm_hashes"): - mm_hashes = inputs.get_mm_hashes(self.info.model_id) - return mm_input( prompt_token_ids=prompt_ids, mm_kwargs=mm_kwargs, @@ -364,7 +538,79 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE): return LanguageModel(self) - def embed_multimodal(self, **kwargs): + def get_mm_mapping(self) -> MultiModelKeys: + """ + Get the module prefix in multimodal models + """ + for name in ("language_model", "text_model"): + if getattr(self.model, name, None) is not None: + return MultiModelKeys.from_string_field(language_model=f"model.{name}") + raise ValueError( + "Could not locate the language model submodule for LoRA support" + ) + + def _split_embeddings( + self, embeddings: torch.Tensor, split_sizes: list[int] + ) -> list[torch.Tensor]: + total_expected = sum(split_sizes) + + # Flatten to 2D: [total_tokens, hidden_dim] + if embeddings.ndim == 3: + embeddings = embeddings.view(-1, embeddings.shape[-1]) + + total_tokens = embeddings.shape[0] + if total_tokens == total_expected: + # Direct match: split_sizes are actual token counts + token_split_sizes = split_sizes + elif total_expected > 0 and total_tokens % total_expected == 0: + # Uniform expansion: each item expands to N tokens + tokens_per_item = total_tokens // total_expected + token_split_sizes = [s * tokens_per_item for s in split_sizes] + elif total_expected > 0: + # Mismatch (profiling with dummy data) - pad/truncate + if total_tokens == 0: + raise ValueError( + "Encoder returned empty embeddings. " + f"Expected {total_expected} tokens from " + f"split_sizes={split_sizes}" + ) + if total_tokens < total_expected: + repeat_factor = (total_expected + total_tokens - 1) // total_tokens + embeddings = embeddings.repeat(repeat_factor, 1) + embeddings = embeddings[:total_expected] + token_split_sizes = split_sizes + else: + return [] + + return list(torch.split(embeddings, token_split_sizes, dim=0)) + + def _embed_audio(self, **kwargs) -> list[torch.Tensor] | None: + self.check_version("5.13.0", "audio models support") + input_features: torch.Tensor | None = kwargs.pop("input_features", None) + if input_features is None: + input_features = kwargs.pop("input_values", None) + if input_features is None: + return None + + num_audio_tokens = kwargs.pop("num_audio_tokens") + kwargs.pop("token_type_ids", None) + kwargs.pop("mm_token_type_ids", None) + + context = nullcontext() + if current_platform.is_rocm(): + context = torch.nn.attention.sdpa_kernel( + backends=[torch.nn.attention.SDPBackend.MATH] + ) + with context: + audio_output = self.model.get_audio_features( + input_features, return_dict=True, **kwargs + ) + audio_embeddings = audio_output.pooler_output + + split_sizes = num_audio_tokens.flatten().tolist() + return self._split_embeddings(audio_embeddings, split_sizes) + + def _embed_vision(self, **kwargs) -> list[torch.Tensor] | torch.Tensor | None: pixel_values: torch.Tensor | None = kwargs.pop("pixel_values", None) image_embeds: torch.Tensor | None = kwargs.pop("image_embeds", None) # Model might use `image_patches` instead of `pixel_values` @@ -379,84 +625,57 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE): num_image_patches = kwargs.pop("num_image_patches") - if pixel_values is not None: + context = nullcontext() + if current_platform.is_rocm(): # ROCm: Force math SDP backend for vision encoder to avoid accuracy issues # with flash_sdp and mem_efficient_sdp - if current_platform.is_rocm(): - # TODO: [ROCm] Fix accuracy issues with flash backend - logger.debug( - "ROCm platform detected. Forcing math SDP backend " - "for vision encoder. Currently ROCm platform has " - "accuracy issues with `flash_sdp` and" - "`mem_efficient_sdp` backends. See issue: " - "https://github.com/vllm-project/vllm/issues/30167" - ) - with torch.nn.attention.sdpa_kernel( - backends=[torch.nn.attention.SDPBackend.MATH] - ): - vision_embeddings = self.model.get_image_features( - pixel_values, **kwargs - ) - else: - vision_embeddings = self.model.get_image_features( - pixel_values, **kwargs - ) - - # Transformers `v5`, `self.get_image_features` returns a tuple - # containing the features and optionally attentions/hidden_states - # After v5 is settled, we can enable qwen3-vl with several outputs - # from `self.get_image_features` - if isinstance(vision_embeddings, tuple): - vision_embeddings = vision_embeddings[0] - elif isinstance(vision_embeddings, dict): - vision_embeddings = vision_embeddings.pooler_output - - if isinstance(vision_embeddings, torch.Tensor): - split_sizes = num_image_patches.flatten().tolist() - total_patches = sum(split_sizes) - - # Flatten to 2D: [total_tokens, hidden_dim] - if vision_embeddings.ndim == 3: - vision_embeddings = vision_embeddings.view( - -1, vision_embeddings.shape[-1] - ) - - total_tokens = vision_embeddings.shape[0] - if total_tokens == total_patches: - # Direct match: num_image_patches are actual token counts - # (e.g., Qwen2.5-VL style) - token_split_sizes = split_sizes - elif total_patches > 0 and total_tokens % total_patches == 0: - # Uniform expansion: each patch expands to N tokens - # (e.g., Idefics3 style) - tokens_per_patch = total_tokens // total_patches - token_split_sizes = [s * tokens_per_patch for s in split_sizes] - elif total_patches > 0: - # Mismatch (profiling with dummy data) - pad/truncate - if total_tokens == 0: - raise ValueError( - "Vision encoder returned empty embeddings. " - f"Expected {total_patches} patches from " - f"num_image_patches={split_sizes}" - ) - if total_tokens < total_patches: - repeat_factor = ( - total_patches + total_tokens - 1 - ) // total_tokens - vision_embeddings = vision_embeddings.repeat(repeat_factor, 1) - vision_embeddings = vision_embeddings[:total_patches] - token_split_sizes = split_sizes - else: - return [] - - return list(torch.split(vision_embeddings, token_split_sizes, dim=0)) - - return vision_embeddings - else: + # TODO: [ROCm] Fix accuracy issues with flash backend logger.debug( - "No pixel values or image embeddings provided for multimodal embedding." + "ROCm platform detected. Forcing math SDP backend " + "for vision encoder. Currently ROCm platform has " + "accuracy issues with `flash_sdp` and" + "`mem_efficient_sdp` backends. See issue: " + "https://github.com/vllm-project/vllm/issues/30167" ) - return None + context = torch.nn.attention.sdpa_kernel( + backends=[torch.nn.attention.SDPBackend.MATH] + ) + with context: + vision_embeddings = self.model.get_image_features(pixel_values, **kwargs) + + # Transformers `v5`, `self.get_image_features` returns a tuple + # containing the features and optionally attentions/hidden_states + # After v5 is settled, we can enable qwen3-vl with several outputs + # from `self.get_image_features` + if isinstance(vision_embeddings, tuple): + vision_embeddings = vision_embeddings[0] + elif isinstance(vision_embeddings, dict): + vision_embeddings = vision_embeddings.pooler_output + + if isinstance(vision_embeddings, torch.Tensor): + split_sizes = num_image_patches.flatten().tolist() + return self._split_embeddings(vision_embeddings, split_sizes) + + return vision_embeddings + + def embed_multimodal(self, **kwargs): + embeddings: tuple[torch.Tensor, ...] = () + if "input_features" in kwargs or "input_values" in kwargs: + audio_embeddings = self._embed_audio(**kwargs) + if audio_embeddings is not None: + embeddings += tuple(audio_embeddings) + if ( + "pixel_values" in kwargs + or "image_embeds" in kwargs + or "image_patches" in kwargs + ): + vision_embeddings = self._embed_vision(**kwargs) + if vision_embeddings is not None: + if isinstance(vision_embeddings, torch.Tensor): + embeddings += (vision_embeddings,) + else: + embeddings += tuple(vision_embeddings) + return embeddings def get_mrope_input_positions( self, From 9a504646985848c1f5785fa57b5f0afc297a29e6 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Sat, 25 Jul 2026 14:29:36 +0100 Subject: [PATCH 036/185] [CI] Stop flaky test from downloading model every time (#49800) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- .../instanttensor_loader/test_weight_utils.py | 36 +++++++++---------- 1 file changed, 16 insertions(+), 20 deletions(-) diff --git a/tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py b/tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py index 09e0a19d11b..3d15f6f1c55 100644 --- a/tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py +++ b/tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py @@ -2,9 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import glob -import tempfile -import huggingface_hub.constants import pytest import torch @@ -21,29 +19,27 @@ from vllm.platforms import current_platform reason="InstantTensor requires NVIDIA GPUs", ) def test_instanttensor_model_loader(): - with tempfile.TemporaryDirectory() as tmpdir: - huggingface_hub.constants.HF_HUB_OFFLINE = False - download_weights_from_hf( - "openai-community/gpt2", allow_patterns=["*.safetensors"], cache_dir=tmpdir - ) - safetensors = glob.glob(f"{tmpdir}/**/*.safetensors", recursive=True) - assert len(safetensors) > 0 + model_dir = download_weights_from_hf( + "openai-community/gpt2", cache_dir=None, allow_patterns=["*.safetensors"] + ) + safetensors = glob.glob(f"{model_dir}/*.safetensors") + assert len(safetensors) > 0 - instanttensor_tensors = {} - hf_safetensors_tensors = {} + instanttensor_tensors = {} + hf_safetensors_tensors = {} - for name, tensor in instanttensor_weights_iterator(safetensors, True): - instanttensor_tensors[name] = tensor.to("cpu") + for name, tensor in instanttensor_weights_iterator(safetensors, True): + instanttensor_tensors[name] = tensor.to("cpu") - for name, tensor in safetensors_weights_iterator(safetensors, True): - hf_safetensors_tensors[name] = tensor + for name, tensor in safetensors_weights_iterator(safetensors, True): + hf_safetensors_tensors[name] = tensor - assert len(instanttensor_tensors) == len(hf_safetensors_tensors) + assert len(instanttensor_tensors) == len(hf_safetensors_tensors) - for name, instanttensor_tensor in instanttensor_tensors.items(): - assert instanttensor_tensor.dtype == hf_safetensors_tensors[name].dtype - assert instanttensor_tensor.shape == hf_safetensors_tensors[name].shape - assert torch.all(instanttensor_tensor.eq(hf_safetensors_tensors[name])) + for name, instanttensor_tensor in instanttensor_tensors.items(): + assert instanttensor_tensor.dtype == hf_safetensors_tensors[name].dtype + assert instanttensor_tensor.shape == hf_safetensors_tensors[name].shape + assert torch.all(instanttensor_tensor.eq(hf_safetensors_tensors[name])) if __name__ == "__main__": From 1423569ff51418fdf2089ed295eeae86f2ea2cd9 Mon Sep 17 00:00:00 2001 From: mosya415 Date: Sat, 25 Jul 2026 16:34:43 +0300 Subject: [PATCH 037/185] [Bugfix][Tool Parser] Fix dropped streaming arguments in Jamba and InternLM2 parsers (#48852) Signed-off-by: mosya415 <263250241+mosya415@users.noreply.github.com> Co-authored-by: mosya415 <263250241+mosya415@users.noreply.github.com> --- .../test_internlm2_tool_parser.py | 43 +++++++++++++++++++ tests/tool_parsers/test_jamba_tool_parser.py | 31 +++++++++++++ vllm/tool_parsers/internlm2_tool_parser.py | 16 +++++-- vllm/tool_parsers/jamba_tool_parser.py | 17 ++++++-- 4 files changed, 99 insertions(+), 8 deletions(-) diff --git a/tests/tool_parsers/test_internlm2_tool_parser.py b/tests/tool_parsers/test_internlm2_tool_parser.py index 2e5069dbed9..7fd3860ef71 100644 --- a/tests/tool_parsers/test_internlm2_tool_parser.py +++ b/tests/tool_parsers/test_internlm2_tool_parser.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json from unittest.mock import MagicMock import pytest @@ -10,6 +11,7 @@ from tests.tool_parsers.common_tests import ( ToolParserTests, ) from vllm.tokenizers import TokenizerLike +from vllm.tool_parsers.internlm2_tool_parser import Internlm2ToolParser class TestInternLM2ToolParser(ToolParserTests): @@ -120,3 +122,44 @@ class TestInternLM2ToolParser(ToolParserTests): ), }, ) + + +def test_streaming_arguments_in_single_delta(default_tokenizer: TokenizerLike) -> None: + """Arguments arriving whole in one delta must not be dropped.""" + tokenizer_vocab = default_tokenizer.get_vocab() + default_tokenizer.get_vocab = MagicMock() + tokenizer_vocab.update( + { + "<|action_start|>": 92540, + "<|plugin|>": 92541, + "<|action_end|>": 92542, + } + ) + default_tokenizer.get_vocab.return_value = tokenizer_vocab + parser = Internlm2ToolParser(default_tokenizer) + + deltas = [ + '<|action_start|><|plugin|>{"name": "get_weather"', + ', "parameters": {"city": "Dallas", "state": "TX"}}<|action_end|>', + ] + + streamed = "" + current_text = "" + for delta_text in deltas: + previous_text = current_text + current_text += delta_text + delta_message = parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=delta_text, + previous_token_ids=[], + current_token_ids=[], + delta_token_ids=[], + request=None, + ) + if delta_message and delta_message.tool_calls: + arguments = delta_message.tool_calls[0].function.arguments + if arguments: + streamed += arguments + + assert json.loads(streamed) == {"city": "Dallas", "state": "TX"} diff --git a/tests/tool_parsers/test_jamba_tool_parser.py b/tests/tool_parsers/test_jamba_tool_parser.py index f0e7899c8aa..9eb7404209d 100644 --- a/tests/tool_parsers/test_jamba_tool_parser.py +++ b/tests/tool_parsers/test_jamba_tool_parser.py @@ -306,3 +306,34 @@ def test_extract_tool_calls_streaming( ) ] assert_tool_calls(actual_tool_calls, expected_tool_calls) + + +def test_extract_tool_calls_streaming_arguments_in_single_delta(jamba_tool_parser): + """Arguments delivered whole in one coarse delta must not be dropped.""" + deltas = [ + '[{"name": "get_current_weather"', + ",", + ' "arguments": {"city": "Dallas", "state": "TX"}}]', + "", + ] + + streamed_arguments = "" + current_text = "" + for delta_text in deltas: + previous_text = current_text + current_text += delta_text + delta_message = jamba_tool_parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=delta_text, + previous_token_ids=[], + current_token_ids=[], + delta_token_ids=[], + request=None, + ) + if delta_message and delta_message.tool_calls: + arguments = delta_message.tool_calls[0].function.arguments + if arguments: + streamed_arguments += arguments + + assert json.loads(streamed_arguments) == {"city": "Dallas", "state": "TX"} diff --git a/vllm/tool_parsers/internlm2_tool_parser.py b/vllm/tool_parsers/internlm2_tool_parser.py index f4aaeef71a0..7cbffcb1c0b 100644 --- a/vllm/tool_parsers/internlm2_tool_parser.py +++ b/vllm/tool_parsers/internlm2_tool_parser.py @@ -26,7 +26,7 @@ from vllm.tool_parsers.abstract_tool_parser import ( Tool, ToolParser, ) -from vllm.tool_parsers.utils import extract_intermediate_diff +from vllm.tool_parsers.utils import extract_intermediate_diff, is_complete_json logger = init_logger(__name__) @@ -146,9 +146,17 @@ class Internlm2ToolParser(ToolParser): elif cur_arguments and not prev_arguments: cur_arguments_json = json.dumps(cur_arguments, ensure_ascii=False) - arguments_delta = cur_arguments_json[ - : cur_arguments_json.index(delta_text) + len(delta_text) - ] + match_start = cur_arguments_json.find(delta_text) + if match_start != -1: + arguments_delta = cur_arguments_json[ + : match_start + len(delta_text) + ] + elif is_complete_json(parsable_arr): + # Complete in this delta: send whole, don't drop. + arguments_delta = cur_arguments_json + else: + # Still partial: wait for more text. + return None delta = DeltaMessage( tool_calls=[ DeltaToolCall( diff --git a/vllm/tool_parsers/jamba_tool_parser.py b/vllm/tool_parsers/jamba_tool_parser.py index dec3c88d934..193a51faa3d 100644 --- a/vllm/tool_parsers/jamba_tool_parser.py +++ b/vllm/tool_parsers/jamba_tool_parser.py @@ -24,7 +24,7 @@ from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser -from vllm.tool_parsers.utils import extract_intermediate_diff +from vllm.tool_parsers.utils import extract_intermediate_diff, is_complete_json from vllm.utils.mistral import is_mistral_tokenizer logger = init_logger(__name__) @@ -266,9 +266,18 @@ class JambaToolParser(ToolParser): cur_arguments_json = json.dumps(cur_arguments, ensure_ascii=False) logger.debug("finding %s in %s", new_text, cur_arguments_json) - arguments_delta = cur_arguments_json[ - : cur_arguments_json.index(new_text) + len(new_text) - ] + # `new_text` may not appear verbatim in the re-serialized JSON. + match_start = cur_arguments_json.find(new_text) + if match_start != -1: + arguments_delta = cur_arguments_json[ + : match_start + len(new_text) + ] + elif is_complete_json(parsable_arr): + # Complete in this delta: send whole, don't drop. + arguments_delta = cur_arguments_json + else: + # Still partial: wait for more text. + return None logger.debug( "First tokens in arguments received: %s", arguments_delta ) From d1a8ba63d9d2bb51ebf60dd5ea1463cf61c70cea Mon Sep 17 00:00:00 2001 From: "rongfu.leng" Date: Sat, 25 Jul 2026 21:45:27 +0800 Subject: [PATCH 038/185] =?UTF-8?q?[Bugfix][MiniMax-M3]=20Fix=20token-majo?= =?UTF-8?q?r=20top-k=20buffer=20handling=20in=20Triton=20=E2=80=A6=20(#491?= =?UTF-8?q?49)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: rongfu.leng --- vllm/models/minimax_m3/common/indexer.py | 7 +++++-- vllm/models/minimax_m3/common/sparse_attention.py | 9 ++++++++- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/vllm/models/minimax_m3/common/indexer.py b/vllm/models/minimax_m3/common/indexer.py index c66e7ce5267..fea22dbdacd 100644 --- a/vllm/models/minimax_m3/common/indexer.py +++ b/vllm/models/minimax_m3/common/indexer.py @@ -424,6 +424,9 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl): # (decode at [:, :nd], prefill at [:, nd:]) and return views into it; the # kernels' out= writes out[:, :total_q]. None -> allocate fresh. buf = self.topk_indices_buffer + buf_htk = ( + buf if buf is None or current_platform.is_rocm() else buf.transpose(0, 1) + ) decode_topk: torch.Tensor | None = None prefill_topk: torch.Tensor | None = None if index_md.num_decodes > 0: @@ -441,7 +444,7 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl): self.num_kv_heads, d.decode_query_len, d.max_decode_query_len, - out=buf, + out=buf_htk, ) if index_md.num_prefills > 0: p = index_md.prefill @@ -465,7 +468,7 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl): self.topk_blocks, self.init_blocks, self.local_blocks, - out=buf[:, nd:, :] if buf is not None else None, + out=buf_htk[:, nd:, :] if buf_htk is not None else None, ) return decode_topk, prefill_topk diff --git a/vllm/models/minimax_m3/common/sparse_attention.py b/vllm/models/minimax_m3/common/sparse_attention.py index f887f14b643..04aaef7c50f 100644 --- a/vllm/models/minimax_m3/common/sparse_attention.py +++ b/vllm/models/minimax_m3/common/sparse_attention.py @@ -384,7 +384,14 @@ class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl): nd = main_md.num_decode_tokens num_tokens = main_md.num_actual_tokens # Indexer top-k from the shared buffer: decode [:, :nd], prefill [:, nd:]. - topk = layer.topk_indices_buffer # type: ignore[attr-defined] + topk_buffer = layer.topk_indices_buffer # type: ignore[attr-defined] + assert topk_buffer is not None + + topk = ( + topk_buffer + if current_platform.is_rocm() + else topk_buffer[:num_tokens].transpose(0, 1) + ) assert topk is not None hd = self.head_size q = query[:num_tokens].view(-1, self.num_heads, hd) From 3e74c60b9c74b9bdb4a4adf9ebd010c7b22104db Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Sat, 25 Jul 2026 15:43:21 +0100 Subject: [PATCH 039/185] [Docs] Use `gen-files` for generated docs content (#49587) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- .gitignore | 3 - .markdownlint.yaml | 3 + .pre-commit-config.yaml | 4 - docs/.nav.yml | 4 +- docs/cli/.nav.yml | 16 +- docs/cli/bench/latency.md | 9 - docs/cli/bench/mm_processor.md | 55 --- docs/cli/bench/serve.md | 9 - docs/cli/bench/sweep/plot.md | 9 - docs/cli/bench/sweep/plot_pareto.md | 9 - docs/cli/bench/sweep/serve.md | 9 - docs/cli/bench/sweep/serve_workload.md | 9 - docs/cli/bench/throughput.md | 9 - docs/cli/chat.md | 5 - docs/cli/complete.md | 5 - docs/cli/json_tip.inc.md | 10 - docs/cli/launch/render.md | 22 - docs/cli/run-batch.md | 9 - docs/cli/serve.md | 9 - docs/configuration/engine_args.md | 10 +- docs/design/attention_backends.md | 100 +--- .../{hooks => gen_files}/generate_argparse.py | 231 +++++++-- .../gen_files/generate_attention_backends.py | 465 +++--------------- .../{hooks => gen_files}/generate_examples.py | 75 ++- .../{hooks => gen_files}/generate_metrics.py | 75 ++- docs/mkdocs/gen_files/generated_content.py | 56 +++ docs/mkdocs/hooks/url_schemes.py | 43 +- docs/pre_run_check.sh | 2 +- docs/usage/metrics.md | 8 +- mkdocs.yaml | 17 +- vllm/config/cache.py | 8 +- vllm/engine/arg_utils.py | 10 +- .../entrypoints/cli/benchmark/mm_processor.py | 49 +- vllm/entrypoints/cli/launch.py | 19 +- 34 files changed, 515 insertions(+), 861 deletions(-) delete mode 100644 docs/cli/bench/latency.md delete mode 100644 docs/cli/bench/mm_processor.md delete mode 100644 docs/cli/bench/serve.md delete mode 100644 docs/cli/bench/sweep/plot.md delete mode 100644 docs/cli/bench/sweep/plot_pareto.md delete mode 100644 docs/cli/bench/sweep/serve.md delete mode 100644 docs/cli/bench/sweep/serve_workload.md delete mode 100644 docs/cli/bench/throughput.md delete mode 100644 docs/cli/chat.md delete mode 100644 docs/cli/complete.md delete mode 100644 docs/cli/json_tip.inc.md delete mode 100644 docs/cli/launch/render.md delete mode 100644 docs/cli/run-batch.md delete mode 100644 docs/cli/serve.md rename docs/mkdocs/{hooks => gen_files}/generate_argparse.py (51%) rename tools/pre_commit/generate_attention_backend_docs.py => docs/mkdocs/gen_files/generate_attention_backends.py (81%) rename docs/mkdocs/{hooks => gen_files}/generate_examples.py (76%) rename docs/mkdocs/{hooks => gen_files}/generate_metrics.py (63%) create mode 100644 docs/mkdocs/gen_files/generated_content.py diff --git a/.gitignore b/.gitignore index 43787b8cd28..46418c03924 100644 --- a/.gitignore +++ b/.gitignore @@ -173,9 +173,6 @@ venv.bak/ # mkdocs documentation /site -docs/argparse -docs/examples/* -!docs/examples/README.md # mypy .mypy_cache/ diff --git a/.markdownlint.yaml b/.markdownlint.yaml index 937487f4736..9140af925ee 100644 --- a/.markdownlint.yaml +++ b/.markdownlint.yaml @@ -3,6 +3,9 @@ MD007: MD013: false MD024: siblings_only: true +MD025: + # Allow front matter title to be different from the first heading in the document. + front_matter_title: "" MD031: list_items: false MD033: false diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3d26a51bbac..a9f744bfa5c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -260,10 +260,6 @@ repos: files: ^docker/(Dockerfile|versions\.json)$ pass_filenames: false additional_dependencies: [dockerfile-parse] - - id: attention-backend-docs - name: Check attention backend documentation is up to date - entry: python tools/pre_commit/generate_attention_backend_docs.py --check - language: python - id: check-boolean-context-manager name: Check for boolean ops in with-statements entry: python tools/pre_commit/check_boolean_context_manager.py diff --git a/docs/.nav.yml b/docs/.nav.yml index 7d985fdeb58..a2fae5d655b 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -56,7 +56,9 @@ nav: - API Reference: - api/README.md - api/vllm - - CLI Reference: cli + - CLI Reference: + - cli/README.md + - vllm: cli - Community: - community/* - Governance: governance diff --git a/docs/cli/.nav.yml b/docs/cli/.nav.yml index 586685c5a10..1a6a5f69386 100644 --- a/docs/cli/.nav.yml +++ b/docs/cli/.nav.yml @@ -1,10 +1,8 @@ nav: - - README.md - - serve.md - - chat.md - - complete.md - - run-batch.md - - vllm bench: - - bench/**/*.md - - vllm launch: - - launch/**/*.md + - "*.md" + - bench: + - bench/*.md + - sweep: + - bench/sweep/*.md + - launch: + - launch/*.md diff --git a/docs/cli/bench/latency.md b/docs/cli/bench/latency.md deleted file mode 100644 index 9e1b9053397..00000000000 --- a/docs/cli/bench/latency.md +++ /dev/null @@ -1,9 +0,0 @@ -# vllm bench latency - -## JSON CLI Arguments - ---8<-- "docs/cli/json_tip.inc.md" - -## Arguments - ---8<-- "docs/generated/argparse/bench_latency.inc.md" diff --git a/docs/cli/bench/mm_processor.md b/docs/cli/bench/mm_processor.md deleted file mode 100644 index 26746ce12d8..00000000000 --- a/docs/cli/bench/mm_processor.md +++ /dev/null @@ -1,55 +0,0 @@ -# vllm bench mm-processor - -## Overview - -`vllm bench mm-processor` profiles the multimodal input processor pipeline of -vision-language models. It measures per-stage latency from the HuggingFace -processor through to the encoder forward pass, helping you identify -preprocessing bottlenecks and understand how different image resolutions or -item counts affect end-to-end request time. - -The benchmark supports two data sources: synthetic random multimodal inputs -(`random-mm`) and HuggingFace datasets (`hf`). Warmup requests are run before -measurement to ensure stable results. - -## Quick Start - -```bash -vllm bench mm-processor \ - --model Qwen/Qwen2-VL-7B-Instruct \ - --dataset-name random-mm \ - --num-prompts 50 \ - --random-input-len 300 \ - --random-output-len 40 \ - --random-mm-base-items-per-request 2 \ - --random-mm-limit-mm-per-prompt '{"image": 3, "video": 0}' \ - --random-mm-bucket-config '{(256, 256, 1): 0.7, (720, 1280, 1): 0.3}' -``` - -## Measured Stages - -| Stage | Description | -| ----- | ----------- | -| `get_mm_hashes_secs` | Time spent hashing multimodal inputs | -| `get_cache_missing_items_secs` | Time spent looking up the processor cache | -| `apply_hf_processor_secs` | Time spent in the HuggingFace processor | -| `merge_mm_kwargs_secs` | Time spent merging multimodal kwargs | -| `apply_prompt_updates_secs` | Time spent updating prompt tokens | -| `preprocessor_total_secs` | Total preprocessing time | -| `encoder_forward_secs` | Time spent in the encoder model forward pass | -| `num_encoder_calls` | Number of encoder invocations per request | - -The benchmark also reports end-to-end latency (TTFT + decode time) per -request. Use `--metric-percentiles` to select which percentiles to report -(default: p99) and `--output-json` to save results. - -For more examples (HF datasets, warmup, JSON output), see -[Benchmarking CLI — Multimodal Processor Benchmark](../../benchmarking/cli.md#multimodal-processor-benchmark). - -## JSON CLI Arguments - ---8<-- "docs/cli/json_tip.inc.md" - -## Arguments - ---8<-- "docs/generated/argparse/bench_mm_processor.inc.md" diff --git a/docs/cli/bench/serve.md b/docs/cli/bench/serve.md deleted file mode 100644 index 792c6e094b3..00000000000 --- a/docs/cli/bench/serve.md +++ /dev/null @@ -1,9 +0,0 @@ -# vllm bench serve - -## JSON CLI Arguments - ---8<-- "docs/cli/json_tip.inc.md" - -## Arguments - ---8<-- "docs/generated/argparse/bench_serve.inc.md" diff --git a/docs/cli/bench/sweep/plot.md b/docs/cli/bench/sweep/plot.md deleted file mode 100644 index d7dc65e6df6..00000000000 --- a/docs/cli/bench/sweep/plot.md +++ /dev/null @@ -1,9 +0,0 @@ -# vllm bench sweep plot - -## JSON CLI Arguments - ---8<-- "docs/cli/json_tip.inc.md" - -## Arguments - ---8<-- "docs/generated/argparse/bench_sweep_plot.inc.md" diff --git a/docs/cli/bench/sweep/plot_pareto.md b/docs/cli/bench/sweep/plot_pareto.md deleted file mode 100644 index 13dffd7f2b5..00000000000 --- a/docs/cli/bench/sweep/plot_pareto.md +++ /dev/null @@ -1,9 +0,0 @@ -# vllm bench sweep plot_pareto - -## JSON CLI Arguments - ---8<-- "docs/cli/json_tip.inc.md" - -## Arguments - ---8<-- "docs/generated/argparse/bench_sweep_plot_pareto.inc.md" diff --git a/docs/cli/bench/sweep/serve.md b/docs/cli/bench/sweep/serve.md deleted file mode 100644 index 6a8182feb40..00000000000 --- a/docs/cli/bench/sweep/serve.md +++ /dev/null @@ -1,9 +0,0 @@ -# vllm bench sweep serve - -## JSON CLI Arguments - ---8<-- "docs/cli/json_tip.inc.md" - -## Arguments - ---8<-- "docs/generated/argparse/bench_sweep_serve.inc.md" diff --git a/docs/cli/bench/sweep/serve_workload.md b/docs/cli/bench/sweep/serve_workload.md deleted file mode 100644 index 8c21788e8d9..00000000000 --- a/docs/cli/bench/sweep/serve_workload.md +++ /dev/null @@ -1,9 +0,0 @@ -# vllm bench sweep serve_workload - -## JSON CLI Arguments - ---8<-- "docs/cli/json_tip.inc.md" - -## Arguments - ---8<-- "docs/generated/argparse/bench_sweep_serve_workload.inc.md" diff --git a/docs/cli/bench/throughput.md b/docs/cli/bench/throughput.md deleted file mode 100644 index 66434c87819..00000000000 --- a/docs/cli/bench/throughput.md +++ /dev/null @@ -1,9 +0,0 @@ -# vllm bench throughput - -## JSON CLI Arguments - ---8<-- "docs/cli/json_tip.inc.md" - -## Arguments - ---8<-- "docs/generated/argparse/bench_throughput.inc.md" diff --git a/docs/cli/chat.md b/docs/cli/chat.md deleted file mode 100644 index 7b8e718f625..00000000000 --- a/docs/cli/chat.md +++ /dev/null @@ -1,5 +0,0 @@ -# vllm chat - -## Arguments - ---8<-- "docs/generated/argparse/chat.inc.md" diff --git a/docs/cli/complete.md b/docs/cli/complete.md deleted file mode 100644 index 65d953a7c04..00000000000 --- a/docs/cli/complete.md +++ /dev/null @@ -1,5 +0,0 @@ -# vllm complete - -## Arguments - ---8<-- "docs/generated/argparse/complete.inc.md" diff --git a/docs/cli/json_tip.inc.md b/docs/cli/json_tip.inc.md deleted file mode 100644 index 56c9cb2cc8e..00000000000 --- a/docs/cli/json_tip.inc.md +++ /dev/null @@ -1,10 +0,0 @@ - -When passing JSON CLI arguments, the following sets of arguments are equivalent: - -- `--json-arg '{"key1": "value1", "key2": {"key3": "value2"}}'` -- `--json-arg.key1 value1 --json-arg.key2.key3 value2` - -Additionally, list elements can be passed individually using `+`: - -- `--json-arg '{"key4": ["value3", "value4", "value5"]}'` -- `--json-arg.key4+ value3 --json-arg.key4+='value4,value5'` diff --git a/docs/cli/launch/render.md b/docs/cli/launch/render.md deleted file mode 100644 index 4d15e5f1162..00000000000 --- a/docs/cli/launch/render.md +++ /dev/null @@ -1,22 +0,0 @@ -# vllm launch render - -## Overview - -`vllm launch render` starts a GPU-less rendering server for preprocessing and -postprocessing only. - -```bash -vllm launch render meta-llama/Llama-3.2-1B-Instruct --port 8100 -``` - -This command reuses the standard serving parser, so model, frontend, -networking, and related CLI options follow the same conventions as -[`vllm serve`](../serve.md). - -## JSON CLI Arguments - ---8<-- "docs/cli/json_tip.inc.md" - -## Arguments - ---8<-- "docs/generated/argparse/launch_render.inc.md" diff --git a/docs/cli/run-batch.md b/docs/cli/run-batch.md deleted file mode 100644 index f2255e66373..00000000000 --- a/docs/cli/run-batch.md +++ /dev/null @@ -1,9 +0,0 @@ -# vllm run-batch - -## JSON CLI Arguments - ---8<-- "docs/cli/json_tip.inc.md" - -## Arguments - ---8<-- "docs/generated/argparse/run-batch.inc.md" diff --git a/docs/cli/serve.md b/docs/cli/serve.md deleted file mode 100644 index 0326fe29ec7..00000000000 --- a/docs/cli/serve.md +++ /dev/null @@ -1,9 +0,0 @@ -# vllm serve - -## JSON CLI Arguments - ---8<-- "docs/cli/json_tip.inc.md" - -## Arguments - ---8<-- "docs/generated/argparse/serve.inc.md" diff --git a/docs/configuration/engine_args.md b/docs/configuration/engine_args.md index b619cbf3db0..11886cca7be 100644 --- a/docs/configuration/engine_args.md +++ b/docs/configuration/engine_args.md @@ -11,12 +11,4 @@ Engine arguments control the behavior of the vLLM engine. The engine argument classes, [EngineArgs][vllm.engine.arg_utils.EngineArgs] and [AsyncEngineArgs][vllm.engine.arg_utils.AsyncEngineArgs], are a combination of the configuration classes defined in [vllm.config][]. Therefore, if you are interested in developer documentation, we recommend looking at these configuration classes as they are the source of truth for types, defaults and docstrings. ---8<-- "docs/cli/json_tip.inc.md" - -## `EngineArgs` - ---8<-- "docs/generated/argparse/engine_args.inc.md" - -## `AsyncEngineArgs` - ---8<-- "docs/generated/argparse/async_engine_args.inc.md" +--8<-- "gen:engine-args" diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index 9497f70e3bd..755ce8961e2 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -1,15 +1,9 @@ # Attention Backend Feature Support -This document is auto-generated by `tools/pre_commit/generate_attention_backend_docs.py`. -It shows the feature support for each registered attention backend -based on the checks in `AttentionBackend.validate_configuration()`. - -**Do not edit this file manually.** Run the following command to -regenerate it: - -```bash -python tools/pre_commit/generate_attention_backend_docs.py -``` +The priority and feature tables on this page are auto-generated from the +attention backend registry by +`docs/mkdocs/gen_files/generate_attention_backends.py`, based on the checks in +`AttentionBackend.validate_configuration()`. ## Setting the Attention Backend @@ -98,40 +92,11 @@ Priority is **1 = highest** (tried first). ### Standard Attention (MHA, MQA, GQA) -**Blackwell (SM 10.x):** - -| Priority | Backend | -| -------- | ------- | -| 1 | `FLASHINFER` | -| 2 | `FLASH_ATTN` | -| 3 | `TRITON_ATTN` | -| 4 | `FLEX_ATTENTION` | -| 5 | `TURBOQUANT` | - -**Ampere/Hopper (SM 8.x-9.x):** - -| Priority | Backend | -| -------- | ------- | -| 1 | `FLASH_ATTN` | -| 2 | `FLASHINFER` | -| 3 | `TRITON_ATTN` | -| 4 | `FLEX_ATTENTION` | -| 5 | `TURBOQUANT` | +--8<-- "gen:priority-standard" ### MLA Attention (DeepSeek-style) -**Blackwell (SM 10.x):** - -| Priority | Backend | -| -------- | ------- | -| 1 | `FLASHINFER_MLA` | -| 2 | `TOKENSPEED_MLA` | -| 3 | `CUTLASS_MLA` | -| 4 | `FLASH_ATTN_MLA` | -| 5 | `FLASHMLA` | -| 6 | `TRITON_MLA` | -| 7 | `FLASHINFER_MLA_SPARSE`**\*** | -| 8 | `FLASHMLA_SPARSE` | +--8<-- "gen:priority-mla" > **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise. > @@ -157,24 +122,7 @@ Priority is **1 = highest** (tried first). ## Standard Attention (MHA, MQA, GQA) Backends -| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. | -| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ | -| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ✅ | ❌ | ❌ | All | N/A | -| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ✅ | ❌ | ✅ | Decoder | 8.x-9.x | -| `FLASHINFER` | XQA† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 9.0 | -| `FLASHINFER` | trtllm-gen† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x | -| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 | -| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x | -| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 | -| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any | -| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any | -| `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 | -| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A | -| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A | -| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A | -| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int4_per_token_head`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ✅ | ❌ | All | Any | -| `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any | -| `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any | +--8<-- "gen:table-standard" > **†** FlashInfer Native is the regular FlashInfer path. XQA is the SM90 decode path exposed through FlashInfer's TRTLLM decode API. trtllm-gen is used on SM100 and supports sinks. Disable XQA/trtllm-gen via `--attention-config.use_trtllm_attention=0`. > @@ -188,9 +136,7 @@ automatic priority lists above. A lightning indexer scores KV blocks, the top-k blocks (plus fixed init/local blocks) are selected, and attention attends only to those blocks; index keys live in a separate side cache. -| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. | -| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ | -| `MINIMAX_M3_SPARSE` | bf16, fp16 | `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 128 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | Any | +--8<-- "gen:table-minimax" ## MLA (Multi-head Latent Attention) Backends @@ -203,12 +149,7 @@ To explicitly select a prefill backend, use Otherwise, the prefill backend is selected automatically at runtime based on hardware and configuration. -| Backend | Description | Dtypes | Compute Cap. | Notes | -| ------- | ----------- | ------ | ------------ | ----- | -| `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) (FA2/FA3/FA4) or (qk_nope_head_dim=64, qk_rope_head_dim=64, v_head_dim=128) (FA2/FA3/FA4) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) (FA2/FA3 only) | -| `TRTLLM_RAGGED` | TensorRT-LLM ragged attention | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) only | -| `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only | -| `TOKENSPEED_MLA` | | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only | +--8<-- "gen:table-mla-prefill" > **‡** Automatic selection tries FlashAttention first. On Blackwell > (SM100), the fallback order is TRT-LLM Ragged, FlashInfer, then @@ -219,22 +160,7 @@ hardware and configuration. MLA decode backends are selected using the standard `-ac.backend=` argument (e.g., `FLASHMLA`, `TRITON_MLA`). -| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. | -| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ | -| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x | -| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x | -| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x | -| `FLASHINFER_MLA_SPARSE_SM120` | bf16 | `auto`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 64, 256 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 12.x | -| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x | -| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x | -| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x | -| `FLASH_ATTN_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x | -| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | -| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A | -| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | -| `TOKENSPEED_MLA` | fp16, bf16 | `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x | -| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | Any | -| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | Any | +--8<-- "gen:table-mla-decode" ### DeepSeek V4 Decode Backends @@ -245,8 +171,4 @@ pipeline (compressor + SWA + indexer, 256-token blocks, head 512); default on NVIDIA is `FLASHINFER_MLA_SPARSE_DSV4` on SM12x and `FLASHMLA_SPARSE_DSV4` on other supported CUDA architectures. -| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. | -| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ | -| `FLASHINFER_MLA_SPARSE_DSV4` | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 256 | 512 | ✅ | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x, 12.x | -| `FLASHMLA_SPARSE_DSV4` | bf16 | `auto`, `fp8_ds_mla`, `fp8` | 256 | 512 | ✅ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x | -| `ROCM_FLASHMLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | +--8<-- "gen:table-mla-v4-decode" diff --git a/docs/mkdocs/hooks/generate_argparse.py b/docs/mkdocs/gen_files/generate_argparse.py similarity index 51% rename from docs/mkdocs/hooks/generate_argparse.py rename to docs/mkdocs/gen_files/generate_argparse.py index 4548e33f881..c0fa707a1cd 100644 --- a/docs/mkdocs/hooks/generate_argparse.py +++ b/docs/mkdocs/gen_files/generate_argparse.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import importlib.metadata import importlib.util +import inspect import logging import sys import textwrap @@ -10,17 +11,21 @@ from argparse import SUPPRESS, Action, HelpFormatter from collections.abc import Callable, Iterable from importlib.machinery import ModuleSpec from pathlib import Path -from typing import TYPE_CHECKING, Literal +from typing import TYPE_CHECKING from unittest.mock import MagicMock, patch +import mkdocs_gen_files +import regex as re from pydantic_core import core_schema logger = logging.getLogger("mkdocs") ROOT_DIR = Path(__file__).parent.parent.parent.parent -ARGPARSE_DOC_DIR = ROOT_DIR / "docs/generated/argparse" sys.path.insert(0, str(ROOT_DIR)) +sys.path.insert(0, str(Path(__file__).parent)) + +from generated_content import fill_markers # noqa: E402 def mock_if_no_torch(mock_module: str, mock: MagicMock): @@ -132,8 +137,8 @@ def auto_mock(module_name: str, attr: str, max_mocks: int = 100): bench_latency = auto_mock("vllm.benchmarks", "latency") -bench_mm_processor = auto_mock("vllm.benchmarks", "mm_processor") bench_serve = auto_mock("vllm.benchmarks", "serve") +bench_startup = auto_mock("vllm.benchmarks", "startup") bench_sweep_plot = auto_mock("vllm.benchmarks.sweep.plot", "SweepPlotArgs") bench_sweep_plot_pareto = auto_mock( "vllm.benchmarks.sweep.plot_pareto", "SweepPlotParetoArgs" @@ -142,12 +147,28 @@ bench_sweep_serve = auto_mock("vllm.benchmarks.sweep.serve", "SweepServeArgs") bench_sweep_serve_workload = auto_mock( "vllm.benchmarks.sweep.serve_workload", "SweepServeWorkloadArgs" ) +bench_sweep_startup = auto_mock("vllm.benchmarks.sweep.startup", "SweepStartupArgs") bench_throughput = auto_mock("vllm.benchmarks", "throughput") AsyncEngineArgs = auto_mock("vllm.engine.arg_utils", "AsyncEngineArgs") EngineArgs = auto_mock("vllm.engine.arg_utils", "EngineArgs") ChatCommand = auto_mock("vllm.entrypoints.cli.openai", "ChatCommand") CompleteCommand = auto_mock("vllm.entrypoints.cli.openai", "CompleteCommand") +BenchmarkSubcommand = auto_mock( + "vllm.entrypoints.cli.benchmark.main", "BenchmarkSubcommand" +) +import_bench_subcommands = auto_mock( + "vllm.entrypoints.cli.benchmark.main", "_import_bench_subcommand_modules" +) +BenchmarkSubcommandBase = auto_mock( + "vllm.entrypoints.cli.benchmark.base", "BenchmarkSubcommandBase" +) +BenchmarkMMProcessorSubcommand = auto_mock( + "vllm.entrypoints.cli.benchmark.mm_processor", "BenchmarkMMProcessorSubcommand" +) +LaunchSubcommandBase = auto_mock("vllm.entrypoints.cli.launch", "LaunchSubcommandBase") +launch_description = auto_mock("vllm.entrypoints.cli.launch", "DESCRIPTION") RenderSubcommand = auto_mock("vllm.entrypoints.cli.launch", "RenderSubcommand") +sweep_subcommands = auto_mock("vllm.benchmarks.sweep.cli", "SUBCOMMANDS") openai_cli_args = auto_mock("vllm.entrypoints.openai", "cli_args") openai_run_batch = auto_mock("vllm.entrypoints.openai", "run_batch") @@ -179,7 +200,7 @@ class MarkdownFormatter(HelpFormatter): def add_text(self, text: str): if text: - self._markdown_output.append(f"{text.strip()}\n\n") + self._markdown_output.append(f"{inspect.cleandoc(text)}\n\n") def add_usage(self, usage, actions, groups, prefix=None): pass @@ -241,49 +262,163 @@ def create_parser(add_cli_args, **kwargs) -> FlexibleArgumentParser: return _parser or parser -def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool): - logger.info("Generating argparse documentation") - logger.debug("Root directory: %s", ROOT_DIR.resolve()) - logger.debug("Output directory: %s", ARGPARSE_DOC_DIR.resolve()) - - # Create the ARGPARSE_DOC_DIR if it doesn't exist - if not ARGPARSE_DOC_DIR.exists(): - ARGPARSE_DOC_DIR.mkdir(parents=True) - - # Create parsers to document - parsers = { - # Engine args - "engine_args": create_parser(EngineArgs.add_cli_args), - "async_engine_args": create_parser( - AsyncEngineArgs.add_cli_args, async_args_only=True - ), - # CLI - "serve": create_parser(openai_cli_args.make_arg_parser), - "chat": create_parser(ChatCommand.add_cli_args), - "complete": create_parser(CompleteCommand.add_cli_args), - "launch_render": create_parser(RenderSubcommand.add_cli_args), - "run-batch": create_parser(openai_run_batch.make_arg_parser), - # Benchmark CLI - "bench_latency": create_parser(bench_latency.add_cli_args), - "bench_mm_processor": create_parser(bench_mm_processor.add_cli_args), - "bench_serve": create_parser(bench_serve.add_cli_args), - "bench_sweep_plot": create_parser(bench_sweep_plot.add_cli_args), - "bench_sweep_plot_pareto": create_parser(bench_sweep_plot_pareto.add_cli_args), - "bench_sweep_serve": create_parser(bench_sweep_serve.add_cli_args), - "bench_sweep_serve_workload": create_parser( - bench_sweep_serve_workload.add_cli_args - ), - "bench_throughput": create_parser(bench_throughput.add_cli_args), - } - - # Generate documentation for each parser - for stem, parser in parsers.items(): - doc_path = ARGPARSE_DOC_DIR / f"{stem}.inc.md" - # Specify encoding for building on Windows - with open(doc_path, "w", encoding="utf-8") as f: - f.write(super(type(parser), parser).format_help()) - logger.info("Argparse generated: %s", doc_path.relative_to(ROOT_DIR)) +def format_help(parser: FlexibleArgumentParser) -> str: + """Format a parser's help as markdown using `MarkdownFormatter`.""" + return super(type(parser), parser).format_help() -if __name__ == "__main__": - on_startup("build", False) +# Absolute docs URLs are kept in the help text because they are useful in the +# terminal. Wrap them as markdown links so the `url_schemes` hook can rewrite +# them into doc-relative links / cross-references at render time. +_DOCS_URL = re.compile(r"https://docs\.vllm\.ai/en/[^/\s]+/[^\s)>]+") + + +def linkify_docs_urls(text: str) -> str: + """Wrap bare docs.vllm.ai URLs in help text as markdown links.""" + return _DOCS_URL.sub(lambda m: f"[{m.group()}]({m.group()})", text) + + +logger.info("Generating argparse documentation") +logger.debug("Root directory: %s", ROOT_DIR.resolve()) + +# The JSON tip is always rendered immediately before generated argument content, +# and the generator is its only consumer, so it lives here rather than in a +# separate snippet file. (The runtime terminal equivalent is +# `FlexibleArgumentParser._json_tip` in vllm/utils/argparse_utils.py.) +JSON_TIP = """## JSON CLI Arguments + +When passing JSON CLI arguments, the following sets of arguments are equivalent: + +- `--json-arg '{"key1": "value1", "key2": {"key3": "value2"}}'` +- `--json-arg.key1 value1 --json-arg.key2.key3 value2` + +Additionally, list elements can be passed individually using `+`: + +- `--json-arg '{"key4": ["value3", "value4", "value5"]}'` +- `--json-arg.key4+ value3 --json-arg.key4+='value4,value5'` + +""" + +# Argument sections filled into `gen:` markers on handwritten pages +engine_args = create_parser(EngineArgs.add_cli_args) +async_engine_args = create_parser(AsyncEngineArgs.add_cli_args, async_args_only=True) +fill_markers( + "configuration/engine_args.md", + { + "engine-args": ( + f"{JSON_TIP}## `EngineArgs`\n\n" + f"{linkify_docs_urls(format_help(engine_args))}" + f"## `AsyncEngineArgs`\n\n" + f"{linkify_docs_urls(format_help(async_engine_args))}" + ) + }, +) + +# CLI reference pages generated entirely from their parser: page -> (parser, JSON tip) +pages = { + "cli/serve.md": (create_parser(openai_cli_args.make_arg_parser), True), + "cli/chat.md": (create_parser(ChatCommand.add_cli_args), False), + "cli/complete.md": (create_parser(CompleteCommand.add_cli_args), False), + "cli/run-batch.md": (create_parser(openai_run_batch.make_arg_parser), True), + "cli/launch/render.md": (create_parser(RenderSubcommand.add_cli_args), True), + "cli/bench/latency.md": (create_parser(bench_latency.add_cli_args), True), + # URL kept as `mm_processor` for back-compat; command name is `mm-processor` + "cli/bench/mm_processor.md": ( + create_parser(BenchmarkMMProcessorSubcommand.add_cli_args), + True, + ), + "cli/bench/serve.md": (create_parser(bench_serve.add_cli_args), True), + "cli/bench/startup.md": (create_parser(bench_startup.add_cli_args), True), + "cli/bench/throughput.md": (create_parser(bench_throughput.add_cli_args), True), + "cli/bench/sweep/plot.md": (create_parser(bench_sweep_plot.add_cli_args), True), + "cli/bench/sweep/plot_pareto.md": ( + create_parser(bench_sweep_plot_pareto.add_cli_args), + True, + ), + "cli/bench/sweep/serve.md": (create_parser(bench_sweep_serve.add_cli_args), True), + "cli/bench/sweep/serve_workload.md": ( + create_parser(bench_sweep_serve_workload.add_cli_args), + True, + ), + "cli/bench/sweep/startup.md": ( + create_parser(bench_sweep_startup.add_cli_args), + True, + ), +} + +# Command name for pages whose file stem differs (URL kept for back-compat). +COMMAND_NAMES = {"cli/bench/mm_processor.md": "mm-processor"} + +for doc_path, (parser, json_tip) in pages.items(): + segments = Path(doc_path).relative_to("cli").with_suffix("").parts + label = COMMAND_NAMES.get(doc_path, segments[-1]) + command = " ".join([*segments[:-1], label]) + # `title` frontmatter keeps the nav label to just this command's segment, + # while the H1 stays the full `vllm ...` command for the page heading. + content = f"---\ntitle: {label}\n---\n\n" + content += f"# vllm {command}\n\n" + if parser.description: + content += f"## Overview\n\n{parser.description}\n\n" + # Rendered above instead of at the top of the Arguments section + parser.description = None + if json_tip: + content += JSON_TIP + content += f"## Arguments\n\n{linkify_docs_urls(format_help(parser))}" + with mkdocs_gen_files.open(doc_path, "w") as f: + f.write(content) + logger.debug("CLI reference generated: %s", doc_path) + +logger.info("Total argparse docs generated: %d", len(pages) + 2) + + +# --- Bare subcommand (group) pages ------------------------------------------- +# Mirror `vllm --help`: an overview plus a table of child subcommands, +# each linked to its reference page. Children are read from the CLI registries +# so the listing can never drift from the actual subcommands. Each page is the +# `README.md` of its command directory so it becomes that section's index and is +# picked up by the existing nav globs. +import_bench_subcommands() # populate BenchmarkSubcommandBase.__subclasses__() +bench_subcommands = BenchmarkSubcommandBase.__subclasses__() +bench_children = [(cmd.name, cmd.help) for cmd in bench_subcommands] + +groups = { + "cli/bench/README.md": (BenchmarkSubcommand.help, bench_children), + "cli/launch/README.md": ( + launch_description, + [(cmd.name, cmd.help) for cmd in LaunchSubcommandBase.__subclasses__()], + ), + "cli/bench/sweep/README.md": ( + dict(bench_children).get("sweep"), + [(args.parser_name, args.parser_help) for args, _ in sweep_subcommands], + ), +} + +# Doc paths that exist, so we only link a child that has a reference page. +existing_pages = set(pages) | set(groups) + + +def child_link(group_doc: str, name: str) -> str | None: + group_dir = Path(group_doc).parent # cli/bench/README.md -> cli/bench + for stem in (name, name.replace("-", "_")): + # A leaf page (bench/latency.md) or a nested group index (sweep/README.md) + for candidate in (group_dir / f"{stem}.md", group_dir / stem / "README.md"): + if candidate.as_posix() in existing_pages: + return candidate.relative_to(group_dir).as_posix() + return None + + +for doc_path, (overview, children) in groups.items(): + title = "vllm " + Path(doc_path).parent.relative_to("cli").as_posix() + lines = [f"# {title.replace('/', ' ')}", ""] + if overview: + lines += ["## Overview", "", overview.strip(), ""] + lines += ["## Subcommands", "", "| Command | Description |", "| --- | --- |"] + for name, summary in children: + link = child_link(doc_path, name) + command = f"[`{name}`]({link})" if link else f"`{name}`" + lines.append(f"| {command} | {(summary or '').strip()} |") + with mkdocs_gen_files.open(doc_path, "w") as f: + f.write("\n".join(lines) + "\n") + logger.debug("CLI group reference generated: %s", doc_path) + +logger.info("CLI group reference pages generated: %d", len(groups)) diff --git a/tools/pre_commit/generate_attention_backend_docs.py b/docs/mkdocs/gen_files/generate_attention_backends.py similarity index 81% rename from tools/pre_commit/generate_attention_backend_docs.py rename to docs/mkdocs/gen_files/generate_attention_backends.py index 875d19a0f55..942165eaa94 100644 --- a/tools/pre_commit/generate_attention_backend_docs.py +++ b/docs/mkdocs/gen_files/generate_attention_backends.py @@ -9,33 +9,28 @@ based on the checks in AttentionBackend.validate_configuration(). This approach avoids requiring CUDA/ROCm/GPU libraries to be installed. -When used as a pre-commit hook, this script receives filenames as arguments -and only runs the check if any of the relevant files were modified. +It runs as an mkdocs-gen-files script, so the page is generated at docs build +time rather than being committed to the repository. """ -import argparse import ast -import fnmatch +import logging import sys from collections.abc import Callable from pathlib import Path from typing import Any +sys.path.insert(0, str(Path(__file__).parent)) + +from generated_content import fill_markers # noqa: E402 + +logger = logging.getLogger("mkdocs") + # --------------------------------------------------------------------------- # Constants and file paths # --------------------------------------------------------------------------- -REPO_ROOT = Path(__file__).parent.parent.parent - -RELEVANT_PATTERNS = [ - "vllm/v1/attention/backends/*.py", - "vllm/v1/attention/backends/**/*.py", - "vllm/models/minimax_m3/common/sparse_attention.py", - "vllm/model_executor/layers/attention/mla_attention.py", - "vllm/platforms/cuda.py", - "tools/pre_commit/generate_attention_backend_docs.py", - "docs/design/attention_backends.md", -] +REPO_ROOT = Path(__file__).parent.parent.parent.parent BACKENDS_DIR = REPO_ROOT / "vllm" / "v1" / "attention" / "backends" REGISTRY_FILE = BACKENDS_DIR / "registry.py" @@ -55,19 +50,6 @@ BACKEND_KV_DTYPE_EXCLUDES: dict[str, set[str]] = { } -def is_relevant_file(filepath: str) -> bool: - """Check if a file matches any of the relevant patterns.""" - path = Path(filepath) - if path.is_absolute(): - try: - path = path.relative_to(REPO_ROOT) - except ValueError: - return False - path_str = str(path) - - return any(fnmatch.fnmatch(path_str, pattern) for pattern in RELEVANT_PATTERNS) - - MLA_PREFILL_DIR = BACKENDS_DIR / "mla" / "prefill" MLA_PREFILL_REGISTRY_FILE = MLA_PREFILL_DIR / "registry.py" MLA_PREFILL_SELECTOR_FILE = MLA_PREFILL_DIR / "selector.py" @@ -960,7 +942,7 @@ def analyze_backend(backend_name: str, class_path: str) -> dict[str, Any] | None try: tree = ast.parse(file_path.read_text()) except Exception as e: - print(f" Warning: Could not parse {file_path}: {e}", file=sys.stderr) + logger.warning("Could not parse %s: %s", file_path, e) return None class_name = class_path.rsplit(".", 1)[1] @@ -1657,113 +1639,12 @@ def _render_table( return lines -def generate_markdown_table( - backends: list[dict[str, Any]], title: str, is_mla_table: bool = False -) -> str: - """Generate a titled markdown table from backend info.""" - if not backends: - return f"## {title}\n\nNo backends found.\n" - has_versions = any(b.get("version") for b in backends) - columns = _build_columns(is_mla_table, has_versions) - lines = [f"## {title}", ""] - lines.extend(_render_table(columns, backends)) - lines.append("") - return "\n".join(lines) - - -# --------------------------------------------------------------------------- -# Markdown section generators (usage, priority, legend, MLA) -# --------------------------------------------------------------------------- - - -def generate_usage_section() -> str: - """Generate the usage documentation section.""" - return """## Setting the Attention Backend - -### Command Line - -There are two ways to specify the backend from the command line: - -**Option 1: Using `--attention-backend` (simple)** - -```bash -vllm serve --attention-backend FLASH_ATTN -``` - -**Option 2: Using `--attention-config.backend` / `-ac.backend` (structured config)** - -```bash -# Dot notation -vllm serve --attention-config.backend FLASH_ATTN -vllm serve -ac.backend FLASH_ATTN - -# JSON format -vllm serve --attention-config '{"backend": "FLASH_ATTN"}' -vllm serve -ac '{"backend": "FLASH_ATTN"}' -``` - -> **Note:** `--attention-backend` and `--attention-config.backend` are mutually -> exclusive. Use one or the other, not both. - -### Python API - -Use `AttentionConfig` with the `LLM` class: - -```python -from vllm import LLM -from vllm.config import AttentionConfig -from vllm.v1.attention.backends.registry import AttentionBackendEnum - -# Method 1: Using AttentionConfig with enum -llm = LLM( - model="Qwen/Qwen3-0.6B", - attention_config=AttentionConfig(backend=AttentionBackendEnum.FLASH_ATTN), -) - -# Method 2: Using attention_backend parameter with string -llm = LLM( - model="Qwen/Qwen3-0.6B", - attention_backend="FLASH_ATTN", -) -``` - -## Backend Selection Behavior - -### Manual Selection - -When you explicitly set a backend via `--attention-backend` or `AttentionConfig`: - -1. The backend is **validated** against your configuration (model dtype, head - size, compute capability, etc.) -2. If the backend **doesn't support** your configuration, an error is raised - with the specific reason -3. If valid, the backend is used - -Example error when selecting an incompatible backend: - -```text -ValueError: Selected backend FLASHMLA is not valid for this configuration. -Reason: ['compute capability not supported'] -``` - -### Automatic Selection - -When no backend is specified (the default): - -1. vLLM iterates through backends in **priority order** (see tables below) -2. Each backend is validated against your configuration -3. The **first compatible backend** is selected -4. If no backend is compatible, an error is raised listing all backends and - their incompatibility reasons -""" - - def _priority_table( title: str, backends: list[str], annotations: dict[str, str] | None = None, ) -> list[str]: - """Generate a priority table for a list of backends.""" + """Render a priority table for a list of backends.""" def _fmt(b: str) -> str: suffix = annotations.get(b, "") if annotations else "" @@ -1779,102 +1660,38 @@ def _priority_table( ] -def generate_priority_section(priorities: dict[str, list[str]]) -> str: - """Generate the priority ranking section.""" - lines = [ - "## Backend Priority (CUDA)", - "", - "When no backend is explicitly selected, vLLM chooses the first", - "compatible backend from these priority-ordered lists.", - "", - "Priority is **1 = highest** (tried first).", - "", - "### Standard Attention (MHA, MQA, GQA)", - "", - ] - - sm100 = "Blackwell (SM 10.x)" - ampere = "Ampere/Hopper (SM 8.x-9.x)" - - if "standard_sm100" in priorities: - lines.extend(_priority_table(sm100, priorities["standard_sm100"])) - if "standard_default" in priorities: - lines.extend(_priority_table(ampere, priorities["standard_default"])) - - lines.extend(["### MLA Attention (DeepSeek-style)", ""]) - - mla_sm100_annotations = { - "FLASHINFER_MLA_SPARSE": "**\\***", - } - if "mla_sm100" in priorities: - lines.extend( - _priority_table(sm100, priorities["mla_sm100"], mla_sm100_annotations) - ) - if "mla_default" in priorities: - lines.extend(_priority_table(ampere, priorities["mla_default"])) - - if "mla_sm100" in priorities: - lines.append( - "> **\\*** For sparse MLA, FP8 KV cache always prefers " - "`FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` " - "is preferred for low query-head counts (<= 16), while " - "`FLASHMLA_SPARSE` is preferred otherwise." - ) - lines.append(">") - - lines.append( - "> **Note:** ROCm and CPU platforms have their own selection logic. " - "See the platform-specific documentation for details." - ) - lines.append("") - - return "\n".join(lines) +_SM100 = "Blackwell (SM 10.x)" +_AMPERE = "Ampere/Hopper (SM 8.x-9.x)" -def generate_legend() -> str: - """Generate a legend explaining the table columns.""" - return """## Legend - -| Column | Description | -| ------ | ----------- | -| **Dtypes** | Supported model data types (fp16, bf16, fp32) | -| **KV Dtypes** | Supported KV cache data types (`auto`, `fp8`, `fp8_e4m3`, etc.) | -| **Block Sizes** | Supported KV cache block sizes (%N means multiples of N) | -| **Head Sizes** | Supported attention head sizes | -| **Sink** | Attention sink support (for StreamingLLM) | -| **Non-Causal** | Non-causal (bidirectional) attention support for decoder models | -| **Sparse** | Sparse attention support (MLA only) | -| **MM Prefix** | Multimodal prefix full attention support | -| **DCP** | Decode Context Parallelism support (`--decode-context-parallel-size`) | -| **Attention Types** | Supported attention patterns (Decoder, Encoder, Enc-Dec) | -| **Compute Cap.** | Required CUDA compute capability (N/A for non-CUDA backends) | - -**Symbols:** ✅ = Supported, ❌ = Not supported -""" - - -def generate_mla_section( - prefill_backends: list[dict[str, Any]], - decode_backends: list[dict[str, Any]], - v4_decode_backends: list[dict[str, Any]] | None = None, +def _priority_block( + priorities: dict[str, list[str]], + sm100_key: str, + default_key: str, + sm100_annotations: dict[str, str] | None = None, ) -> str: - """Generate the complete MLA section with prefill and decode tables.""" + """Render whichever priority tables exist for one attention category.""" + lines: list[str] = [] + if sm100_key in priorities: + lines += _priority_table(_SM100, priorities[sm100_key], sm100_annotations) + if default_key in priorities: + lines += _priority_table(_AMPERE, priorities[default_key]) + return "\n".join(lines).strip() + + +def _feature_table(backends: list[dict[str, Any]], is_mla: bool) -> str: + """Render a backend feature table (header, separator, one row per backend).""" + has_versions = any(b.get("version") for b in backends) + columns = _build_columns(is_mla, has_versions) + return "\n".join(_render_table(columns, backends)) + + +def _mla_prefill_table(prefill_backends: list[dict[str, Any]]) -> str: + """Render the MLA prefill backend table.""" lines = [ - "## MLA (Multi-head Latent Attention) Backends", - "", - "MLA uses separate backends for prefill and decode phases.", - "", - "### Prefill Backends", - "", - "To explicitly select a prefill backend, use", - "`-ac.mla_prefill_backend=` (e.g., `FLASH_ATTN`, `FLASHINFER`).", - "Otherwise, the prefill backend is selected automatically at runtime based on", - "hardware and configuration.", - "", "| Backend | Description | Dtypes | Compute Cap. | Notes |", "| ------- | ----------- | ------ | ------------ | ----- |", ] - for backend in prefill_backends: row = "| `{}`{} | {} | {} | {} | {} |".format( backend["name"], @@ -1885,87 +1702,21 @@ def generate_mla_section( backend.get("notes", ""), ) lines.append(row.replace(" ", " ")) - - lines.extend( - [ - "", - "> **‡** Automatic selection tries FlashAttention first. On Blackwell", - "> (SM100), the fallback order is TRT-LLM Ragged, FlashInfer, then", - "> TokenSpeed MLA. On other GPUs, only FlashAttention is considered.", - "", - "### Decode Backends", - "", - "MLA decode backends are selected using the standard", - "`-ac.backend=` argument (e.g., `FLASHMLA`, `TRITON_MLA`).", - "", - ] - ) - - # Reuse data-driven table rendering for decode backends - columns = _build_columns(is_mla=True, has_versions=False) - lines.extend(_render_table(columns, decode_backends)) - - if v4_decode_backends: - lines.extend( - [ - "", - "### DeepSeek V4 Decode Backends", - "", - "DeepSeek V4 sparse MLA uses its own decode backends, selected via", - "`--attention-backend=` (e.g., `FLASHMLA_SPARSE_DSV4`,", - "`FLASHINFER_MLA_SPARSE_DSV4`). They share the V4 sparse-index", - "pipeline (compressor + SWA + indexer, 256-token blocks, head 512);", - "default on NVIDIA is `FLASHINFER_MLA_SPARSE_DSV4` on SM12x and", - "`FLASHMLA_SPARSE_DSV4` on other supported CUDA architectures.", - "", - ] - ) - lines.extend(_render_table(columns, v4_decode_backends)) - - lines.append("") return "\n".join(lines) -def generate_minimax_section(backends: list[dict[str, Any]]) -> str: - """Generate the MiniMax M3 sparse attention section.""" - lines = [ - "## MiniMax M3 Sparse Attention Backends", - "", - 'Block-sparse GQA backend used by MiniMax M3 sparse ("lightning indexer")', - "layers. It is wired in directly by the model and is not part of the", - "automatic priority lists above. A lightning indexer scores KV blocks, the", - "top-k blocks (plus fixed init/local blocks) are selected, and attention", - "attends only to those blocks; index keys live in a separate side cache.", - "", - ] - columns = _build_columns(is_mla=False, has_versions=False) - lines.extend(_render_table(columns, backends)) - lines.append("") - return "\n".join(lines) +def build_blocks() -> dict[str, str]: + """Build the generated table blocks keyed by their `gen:` marker name. - -# --------------------------------------------------------------------------- -# Top-level orchestration -# --------------------------------------------------------------------------- - - -def generate_docs() -> str: - """Generate the complete documentation.""" + Only the tables are generated here; the surrounding prose lives in the + handwritten ``docs/design/attention_backends.md`` page. + """ attention_backends_map = parse_registry() - - # Parse priority lists from cuda.py priorities = parse_cuda_priority_lists() - - # Parse FlashAttention FA2/FA3 feature differences fa_features = parse_flash_attn_features() - - # Parse FlashInfer TRTLLM feature differences (native vs TRTLLM on Blackwell) fi_features = parse_flashinfer_trtllm_features() - - # Parse MLA prefill backends mla_prefill_backends = parse_mla_prefill_backends() - # Collect backend info all_backends = [] for backend_name, class_path in attention_backends_map.items(): if backend_name in SKIP_BACKENDS: @@ -1973,17 +1724,14 @@ def generate_docs() -> str: info = analyze_backend(backend_name, class_path) if info: all_backends.append(info) - - # Expand backends into version variants if fa_features: all_backends = _expand_flash_attn_variants(all_backends, fa_features) if fi_features: all_backends = _expand_flashinfer_variants(all_backends, fi_features) - # DeepSeek V4 (*_DSV4) decode backends and MiniMax M3 sparse backends each - # get their own subsection rather than mixing into the main MLA / standard - # tables (the ROCm V4 backend isn't flagged is_mla by the AST heuristic, so - # filter purely on the name). + # DeepSeek V4 (*_DSV4) and MiniMax M3 sparse backends get their own tables + # rather than mixing into the main MLA / standard tables (the ROCm V4 backend + # isn't flagged is_mla by the AST heuristic, so filter purely on the name). def _is_v4(b: dict[str, Any]) -> bool: return b["name"].endswith("_DSV4") @@ -1999,112 +1747,21 @@ def generate_docs() -> str: if not b["is_mla"] and not _is_v4(b) and not _is_minimax(b) ] - # Generate documentation - script_path = "tools/pre_commit/generate_attention_backend_docs.py" - doc_lines = [ - "# Attention Backend Feature Support", - "", - f"This document is auto-generated by `{script_path}`.", - "It shows the feature support for each registered attention backend", - "based on the checks in `AttentionBackend.validate_configuration()`.", - "", - "**Do not edit this file manually.** Run the following command to", - "regenerate it:", - "", - "```bash", - f"python {script_path}", - "```", - "", - ] - - # Add usage documentation - doc_lines.append(generate_usage_section()) - - # Add priority section - doc_lines.append(generate_priority_section(priorities)) - - # Add legend and feature tables - doc_lines.append(generate_legend()) - standard_title = "Standard Attention (MHA, MQA, GQA) Backends" - doc_lines.append( - generate_markdown_table(non_mla_backends, standard_title, is_mla_table=False) - ) - # Add footnotes for version/variant distinctions (in table order) - footnotes = [] - if fi_features: - footnotes.append( - "> **†** FlashInfer Native is the regular FlashInfer path. XQA is the " - "SM90 decode path exposed through FlashInfer's TRTLLM decode API. " - "trtllm-gen is used on SM100 and supports sinks. Disable XQA/trtllm-gen " - "via `--attention-config.use_trtllm_attention=0`." - ) - if fa_features: - footnotes.append( - "> **\\*** Specify the FlashAttention version via " - "`--attention-config.flash_attn_version=2`, `3`, or `4`. " - "Default is FA4 on SM100+ (Blackwell), FA3 on SM90 (Hopper), " - "FA2 otherwise." - ) - if footnotes: - doc_lines.append("\n>\n".join(footnotes) + "\n") - - # Add MiniMax M3 sparse section (separate category after standard GQA) - if minimax_backends: - doc_lines.append(generate_minimax_section(minimax_backends)) - - # Add MLA section with prefill and decode backends - doc_lines.append( - generate_mla_section(mla_prefill_backends, mla_backends, v4_decode_backends) - ) - - return "\n".join(doc_lines) + mla_sm100_annotations = {"FLASHINFER_MLA_SPARSE": "**\\***"} + return { + "priority-standard": _priority_block( + priorities, "standard_sm100", "standard_default" + ), + "priority-mla": _priority_block( + priorities, "mla_sm100", "mla_default", mla_sm100_annotations + ), + "table-standard": _feature_table(non_mla_backends, is_mla=False), + "table-minimax": _feature_table(minimax_backends, is_mla=False), + "table-mla-prefill": _mla_prefill_table(mla_prefill_backends), + "table-mla-decode": _feature_table(mla_backends, is_mla=True), + "table-mla-v4-decode": _feature_table(v4_decode_backends, is_mla=True), + } -def main(): - parser = argparse.ArgumentParser( - description="Generate attention backend documentation table" - ) - parser.add_argument( - "--output", - "-o", - type=str, - default=str(REPO_ROOT / "docs" / "design" / "attention_backends.md"), - help="Output file path (default: docs/design/attention_backends.md)", - ) - parser.add_argument( - "--check", - action="store_true", - help="Check if the documentation is up to date (for pre-commit)", - ) - parser.add_argument( - "files", - nargs="*", - help="Files to check (passed by pre-commit). If none are relevant, skip.", - ) - args = parser.parse_args() - - if args.files and not any(is_relevant_file(f) for f in args.files): - sys.exit(0) - - output_path = Path(args.output) - new_content = generate_docs() - - if args.check: - needs_update = ( - not output_path.exists() or output_path.read_text() != new_content - ) - if needs_update: - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_text(new_content) - print(f"🔄 Regenerated: {output_path}") - sys.exit(1) - print(f"✅ Up to date: {output_path}") - sys.exit(0) - - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_text(new_content) - print(f"Generated: {output_path}") - - -if __name__ == "__main__": - main() +logger.info("Generating attention backend documentation") +fill_markers("design/attention_backends.md", build_blocks()) diff --git a/docs/mkdocs/hooks/generate_examples.py b/docs/mkdocs/gen_files/generate_examples.py similarity index 76% rename from docs/mkdocs/hooks/generate_examples.py rename to docs/mkdocs/gen_files/generate_examples.py index 07fbd7e4d55..c9c10018ecf 100644 --- a/docs/mkdocs/hooks/generate_examples.py +++ b/docs/mkdocs/gen_files/generate_examples.py @@ -5,16 +5,15 @@ import logging from dataclasses import dataclass from functools import cached_property from pathlib import Path -from typing import Literal +import mkdocs_awesome_nav.nav.directory as _nav_dir +import mkdocs_gen_files import regex as re logger = logging.getLogger("mkdocs") ROOT_DIR = Path(__file__).parent.parent.parent.parent -ROOT_DIR_RELATIVE = "../../../../.." EXAMPLE_DIR = ROOT_DIR / "examples" -EXAMPLE_DOC_DIR = ROOT_DIR / "docs/examples" def title(text: str) -> str: @@ -197,44 +196,38 @@ class Example: return content -def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool): - # Monkey-patch dirname_to_title in awesome-nav so that sub-directory names are - # title-cased (e.g. "Offline Inference" instead of "Offline inference"). - import mkdocs_awesome_nav.nav.directory as _nav_dir +# Monkey-patch dirname_to_title in awesome-nav so that sub-directory names are +# title-cased (e.g. "Offline Inference" instead of "Offline inference"). +_nav_dir.dirname_to_title = title +logger.info("Generating example documentation") +logger.debug("Root directory: %s", ROOT_DIR.resolve()) +logger.debug("Example directory: %s", EXAMPLE_DIR.resolve()) - _nav_dir.dirname_to_title = title - logger.info("Generating example documentation") - logger.debug("Root directory: %s", ROOT_DIR.resolve()) - logger.debug("Example directory: %s", EXAMPLE_DIR.resolve()) - logger.debug("Example document directory: %s", EXAMPLE_DOC_DIR.resolve()) +categories = sorted( + p for p in EXAMPLE_DIR.iterdir() if p.is_dir() and not p.name.startswith(".") +) - # Create the EXAMPLE_DOC_DIR if it doesn't exist - if not EXAMPLE_DOC_DIR.exists(): - EXAMPLE_DOC_DIR.mkdir(parents=True) +examples = [] +glob_patterns = ["*.py", "*.md", "*.sh"] +# Find categorised examples +for category in categories: + logger.info("Processing category: %s", category.stem) + globs = [category.glob(pattern) for pattern in glob_patterns] + for path in itertools.chain(*globs): + examples.append(Example(path, category.stem)) + # Find examples in subdirectories + globs = [category.glob(f"*/{pattern}") for pattern in glob_patterns] + for path in itertools.chain(*globs): + examples.append(Example(path.parent, category.stem)) - categories = sorted(p for p in EXAMPLE_DIR.iterdir() if p.is_dir()) - - examples = [] - glob_patterns = ["*.py", "*.md", "*.sh"] - # Find categorised examples - for category in categories: - logger.info("Processing category: %s", category.stem) - globs = [category.glob(pattern) for pattern in glob_patterns] - for path in itertools.chain(*globs): - examples.append(Example(path, category.stem)) - # Find examples in subdirectories - globs = [category.glob(f"*/{pattern}") for pattern in glob_patterns] - for path in itertools.chain(*globs): - examples.append(Example(path.parent, category.stem)) - - # Generate the example documentation - for example in sorted(examples, key=lambda e: e.path.stem): - example_name = f"{example.path.stem}.md" - doc_path = EXAMPLE_DOC_DIR / example.category / example_name - if not doc_path.parent.exists(): - doc_path.parent.mkdir(parents=True) - # Specify encoding for building on Windows - with open(doc_path, "w+", encoding="utf-8") as f: - f.write(example.generate()) - logger.debug("Example generated: %s", doc_path.relative_to(ROOT_DIR)) - logger.info("Total examples generated: %d", len(examples)) +# Generate the example documentation +for example in sorted(examples, key=lambda e: e.path.stem): + doc_path = f"examples/{example.category}/{example.path.stem}.md" + with mkdocs_gen_files.open(doc_path, "w") as f: + f.write(example.generate()) + if example.main_file is not None: + # Point the edit button at the example's source file + edit_path = Path("..") / example.main_file.relative_to(ROOT_DIR) + mkdocs_gen_files.set_edit_path(doc_path, str(edit_path)) + logger.debug("Example generated: %s", doc_path) +logger.info("Total examples generated: %d", len(examples)) diff --git a/docs/mkdocs/hooks/generate_metrics.py b/docs/mkdocs/gen_files/generate_metrics.py similarity index 63% rename from docs/mkdocs/hooks/generate_metrics.py rename to docs/mkdocs/gen_files/generate_metrics.py index 97282aaee7d..b952f53977a 100644 --- a/docs/mkdocs/hooks/generate_metrics.py +++ b/docs/mkdocs/gen_files/generate_metrics.py @@ -2,27 +2,28 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import ast import logging +import sys from pathlib import Path -from typing import Literal + +sys.path.insert(0, str(Path(__file__).parent)) + +from generated_content import fill_markers # noqa: E402 logger = logging.getLogger("mkdocs") ROOT_DIR = Path(__file__).parent.parent.parent.parent -DOCS_DIR = ROOT_DIR / "docs" -GENERATED_METRICS_DIR = DOCS_DIR / "generated" / "metrics" -# Files to scan for metric definitions - each will generate a separate table +# Files to scan for metric definitions - each fills a `gen:` marker in +# docs/usage/metrics.md with its table (the section heading and any preamble +# live in the tracked page next to the marker). METRIC_SOURCE_FILES = [ - {"path": "vllm/v1/metrics/loggers.py", "output": "general.inc.md"}, - { - "path": "vllm/v1/spec_decode/metrics.py", - "output": "spec_decode.inc.md", - }, + {"path": "vllm/v1/metrics/loggers.py", "key": "metrics-general"}, + {"path": "vllm/v1/spec_decode/metrics.py", "key": "metrics-spec-decode"}, { "path": "vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py", - "output": "nixl_connector.inc.md", + "key": "metrics-nixl", }, - {"path": "vllm/v1/metrics/perf.py", "output": "perf.inc.md"}, + {"path": "vllm/v1/metrics/perf.py", "key": "metrics-mfu"}, ] @@ -110,41 +111,27 @@ def generate_markdown_table(metrics: list[dict[str, str]]) -> str: return "\n".join(lines) + "\n" -def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool): - """Generate metrics documentation tables from source files.""" - logger.info("Generating metrics documentation") +logger.info("Generating metrics documentation") - # Create generated directory if it doesn't exist - GENERATED_METRICS_DIR.mkdir(parents=True, exist_ok=True) +blocks = {} +total_metrics = 0 +for source_config in METRIC_SOURCE_FILES: + source_path = source_config["path"] - total_metrics = 0 - for source_config in METRIC_SOURCE_FILES: - source_path = source_config["path"] - output_file = source_config["output"] + filepath = ROOT_DIR / source_path + if not filepath.exists(): + raise FileNotFoundError(f"Metrics source file not found: {filepath}") - filepath = ROOT_DIR / source_path - if not filepath.exists(): - raise FileNotFoundError(f"Metrics source file not found: {filepath}") + logger.debug("Extracting metrics from: %s", source_path) + metrics = extract_metrics_from_file(filepath) + logger.debug("Found %d metrics in %s", len(metrics), source_path) - logger.debug("Extracting metrics from: %s", source_path) - metrics = extract_metrics_from_file(filepath) - logger.debug("Found %d metrics in %s", len(metrics), source_path) + blocks[source_config["key"]] = generate_markdown_table(metrics).strip() + total_metrics += len(metrics) - # Generate and write the markdown table for this source - table_content = generate_markdown_table(metrics) - output_path = GENERATED_METRICS_DIR / output_file - with open(output_path, "w", encoding="utf-8") as f: - f.write(table_content) - - total_metrics += len(metrics) - logger.info( - "Generated metrics table: %s (%d metrics)", - output_path.relative_to(ROOT_DIR), - len(metrics), - ) - - logger.info( - "Total metrics generated: %d across %d files", - total_metrics, - len(METRIC_SOURCE_FILES), - ) +fill_markers("usage/metrics.md", blocks) +logger.info( + "Total metrics generated: %d across %d files", + total_metrics, + len(METRIC_SOURCE_FILES), +) diff --git a/docs/mkdocs/gen_files/generated_content.py b/docs/mkdocs/gen_files/generated_content.py new file mode 100644 index 00000000000..fa3f31b294b --- /dev/null +++ b/docs/mkdocs/gen_files/generated_content.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inline build-time generated content into existing docs pages. + +Source pages mark where generated content goes with a snippet-style marker, +`--8<-- "gen:"`, so the insertion point is explicit and readable. The +substitution happens here (at gen-files time, before mkdocs-gen-files shadows +the page), not via pymdownx.snippets, so the content can be generated at build +time without living in a real file on disk. + +The `gen:` prefix keeps these markers distinct from real pymdownx.snippets +includes, and `fill_markers` fails loudly if a marker is missing or left behind +(pymdownx.snippets would otherwise silently drop an unsubstituted marker). +""" + +from pathlib import Path + +import mkdocs_gen_files +import regex as re + +DOCS_DIR = Path(__file__).parent.parent.parent + +_MARKER = '--8<-- "gen:{key}"' +_ANY_MARKER = re.compile(r'--8<-- "gen:[^"]*"') + + +def fill_markers(doc_path: str, blocks: dict[str, str]) -> None: + """Replace `--8<-- "gen:"` markers in a docs page with generated content. + + Args: + doc_path: Docs-relative path of the source page to fill. + blocks: Mapping of marker key to the markdown to insert in its place. + + Raises: + FileNotFoundError: If the source page does not exist. + ValueError: If an expected marker is missing, or any `gen:` marker is + left unsubstituted after filling. + """ + source = DOCS_DIR / doc_path + if not source.exists(): + raise FileNotFoundError(f"Cannot fill markers in missing page: {doc_path}") + + text = source.read_text() + for key, content in blocks.items(): + marker = _MARKER.format(key=key) + if marker not in text: + raise ValueError(f"{doc_path}: missing marker {marker}") + text = text.replace(marker, content) + + if leftover := _ANY_MARKER.search(text): + raise ValueError(f"{doc_path}: unsubstituted marker {leftover.group()}") + + with mkdocs_gen_files.open(doc_path, "w") as f: + f.write(text) + # Keep the edit button pointing at the real source page + mkdocs_gen_files.set_edit_path(doc_path, doc_path) diff --git a/docs/mkdocs/hooks/url_schemes.py b/docs/mkdocs/hooks/url_schemes.py index e6faf95cd46..d208b9d6779 100644 --- a/docs/mkdocs/hooks/url_schemes.py +++ b/docs/mkdocs/hooks/url_schemes.py @@ -19,6 +19,7 @@ The on_page_markdown hook passes the current page context to the preprocessor be each page is converted. """ +import posixpath from pathlib import Path import regex as re @@ -38,18 +39,22 @@ TITLE = r"(?P[^\[\]<>]+?)" REPO = r"(?P<repo>.+?/.+?)" TYPE = r"(?P<type>issues|pull|projects)" NUMBER = r"(?P<number>\d+)" +VERSION = r"[^/\s]+" PATH = r"(?P<path>[^\s]+?)" FRAGMENT = r"(?P<fragment>#[^\s]+)?" -URL = f"https://github.com/{REPO}/{TYPE}/{NUMBER}{FRAGMENT}" +URL_GITHUB = f"https://github.com/{REPO}/{TYPE}/{NUMBER}{FRAGMENT}" RELATIVE = rf"(?!(https?|ftp)://|#){PATH}{FRAGMENT}" +URL_DOCS = f"https://docs.vllm.ai/en/{VERSION}/{PATH}{FRAGMENT}" # Common titles to use for GitHub links when none is provided in the link. TITLES = {"issues": "Issue ", "pull": "Pull Request ", "projects": "Project "} # Regex to match GitHub issue, PR, and project links with optional titles. -github_link = re.compile(rf"(\[{TITLE}\]\(|<){URL}(\)|>)") +github_link = re.compile(rf"(\[{TITLE}\]\(|<){URL_GITHUB}(\)|>)") # Regex to match relative file links with optional titles. relative_link = re.compile(rf"\[{TITLE}\]\({RELATIVE}\)") +# Regex to match absolute docs.vllm.ai links (should only exist in CLI). +docs_link = re.compile(rf"\[{TITLE}\]\({URL_DOCS}\)") class UrlSchemesPreprocessor(Preprocessor): @@ -61,7 +66,8 @@ class UrlSchemesPreprocessor(Preprocessor): def run(self, lines): page = self.ext.page - if page is None or getattr(page.file, "abs_src_path", None) is None: + files = self.ext.files + if page is None: return lines def replace_relative_link(match: re.Match) -> str: @@ -70,7 +76,7 @@ class UrlSchemesPreprocessor(Preprocessor): """ title = match.group("title") path = match.group("path") - path = (Path(page.file.abs_src_path).parent / path).resolve() + path = ((DOC_DIR / page.file.src_uri).parent / path).resolve() fragment = match.group("fragment") or "" # Check if the path exists and is outside the docs dir @@ -105,9 +111,36 @@ class UrlSchemesPreprocessor(Preprocessor): url = f"https://github.com/{repo}/{type}/{number}{fragment}" return f"[{gh_icon} {title}]({url})" + def replace_docs_link(match: re.Match) -> str: + """Rewrite absolute docs.vllm.ai links as doc-relative links.""" + title = match.group("title") + path = match.group("path").rstrip("/") + fragment = match.group("fragment") or "" + + # vllm.config.<Class> API reference -> mkdocstrings cross-reference + if path == "api/vllm/config" and re.fullmatch( + r"#vllm\.config\.\w+", fragment + ): + ident = fragment[1:] + return f"[`{ident}`][{ident}]" + + # Other docs pages -> link relative to the current page, but only + # when the target is a known docs page (real or generated); leave + # unknown/external URLs untouched. This is correct even when the same + # docstring is also rendered on its API reference page. + src = f"{path.removesuffix('.html')}.md" + if files.get_file_from_path(src) is None: + return match.group(0) + rel = posixpath.relpath(src, posixpath.dirname(page.file.src_uri)) + # Auto-wrapped bare URLs use the URL as their title; make it readable. + if title.startswith("http"): + title = path.removesuffix(".html") + return f"[{title}]({rel}{fragment})" + markdown = "\n".join(lines) markdown = github_link.sub(replace_github_link, markdown) markdown = relative_link.sub(replace_relative_link, markdown) + markdown = docs_link.sub(replace_docs_link, markdown) return markdown.split("\n") @@ -116,6 +149,7 @@ class UrlSchemesExtension(Extension): def __init__(self, **kwargs): self.page = None + self.files = None super().__init__(**kwargs) def extendMarkdown(self, md): @@ -138,4 +172,5 @@ def on_page_markdown( ) -> str: """Pass the current page context to the preprocessor.""" _ext.page = page + _ext.files = files return markdown diff --git a/docs/pre_run_check.sh b/docs/pre_run_check.sh index d55f8c8db12..d611e616071 100644 --- a/docs/pre_run_check.sh +++ b/docs/pre_run_check.sh @@ -28,7 +28,7 @@ DOCS_PATHS=( docs/ # Actual docs content examples/ # Examples are rendered in docs vllm/ # API & CLI reference - requirements/test/cuda.txt # CLI reference (see docs/mkdocs/hooks/generate_argparse.py) + requirements/test/cuda.txt # CLI reference (see docs/mkdocs/gen_files/generate_argparse.py) mkdocs.yaml # Affects build process .readthedocs.yaml # Affects build process requirements/docs.txt # Affects build process diff --git a/docs/usage/metrics.md b/docs/usage/metrics.md index 44c9c7cbfe5..305c21478fe 100644 --- a/docs/usage/metrics.md +++ b/docs/usage/metrics.md @@ -35,21 +35,21 @@ The following metrics are exposed: ## General Metrics ---8<-- "docs/generated/metrics/general.inc.md" +--8<-- "gen:metrics-general" ## Speculative Decoding Metrics ---8<-- "docs/generated/metrics/spec_decode.inc.md" +--8<-- "gen:metrics-spec-decode" ## NIXL KV Connector Metrics ---8<-- "docs/generated/metrics/nixl_connector.inc.md" +--8<-- "gen:metrics-nixl" ## Model Flops Utilization (MFU) Performance Metrics These metrics are available via `--enable-mfu-metrics`: ---8<-- "docs/generated/metrics/perf.inc.md" +--8<-- "gen:metrics-mfu" ## Deprecation Policy diff --git a/mkdocs.yaml b/mkdocs.yaml index a5c03c9e45c..f3b9deab787 100644 --- a/mkdocs.yaml +++ b/mkdocs.yaml @@ -3,7 +3,6 @@ site_url: !ENV READTHEDOCS_CANONICAL_URL repo_url: https://github.com/vllm-project/vllm edit_uri: edit/main/docs/ exclude_docs: | - argparse *.inc.md *.template.md theme: @@ -50,24 +49,22 @@ theme: hooks: - docs/mkdocs/hooks/remove_announcement.py - - docs/mkdocs/hooks/generate_examples.py - - docs/mkdocs/hooks/generate_argparse.py - - docs/mkdocs/hooks/generate_metrics.py - docs/mkdocs/hooks/url_schemes.py - docs/mkdocs/hooks/autoref_code.py plugins: - meta - search + - gen-files: + scripts: + - docs/mkdocs/gen_files/generate_examples.py + - docs/mkdocs/gen_files/generate_argparse.py + - docs/mkdocs/gen_files/generate_metrics.py + - docs/mkdocs/gen_files/generate_attention_backends.py - autorefs - awesome-nav - glightbox - - git-revision-date-localized: - # exclude autogenerated files - exclude: - - api/* - - examples/* - - generated/* + - git-revision-date-localized - minify: minify_html: true minify_js: true diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 45c93624188..d267e9af358 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -137,13 +137,13 @@ class CacheConfig: still be controlled by mamba_cache_dtype). If set to 'auto', the data type for the ssm state will be determined by mamba_cache_dtype.""" mamba_cache_mode: MambaCacheMode = "none" - """The cache strategy for Mamba layers. + """The cache strategy for Mamba layers: + - "none": set when prefix caching is disabled. - "all": cache the mamba state of all tokens at position i * block_size. This is - the default behavior (for models that support it) when prefix caching is - enabled. + the default behavior (for models that support it) when prefix caching is enabled. - "align": only cache the mamba state of the last token of each scheduler step and - when the token is at position i * block_size. + when the token is at position i * block_size. """ replayssm_buffer_len: int = Field(default=16, gt=0) """ReplaySSM history buffer length B: with use_replayssm, standard decode diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index a14de27190e..73a724c4427 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -249,17 +249,19 @@ def get_type_hints(type_hint: TypeHint) -> set[TypeHint]: NEEDS_HELP = ( any("--help" in arg for arg in sys.argv) # vllm SUBCOMMAND --help - or (argv0 := sys.argv[0]).endswith("mkdocs") # mkdocs SUBCOMMAND - or argv0.endswith("mkdocs/__main__.py") # python -m mkdocs SUBCOMMAND + or "mkdocs" in sys.modules # mkdocs SUBCOMMAND ) def _maybe_add_docs_url(cls: Any) -> str: """Generate API docs URL for a vllm config class.""" - if not cls.__module__.startswith("vllm.config"): + import vllm.config + + name = cls.__name__ + if getattr(vllm.config, name, None) is not cls: return "" version = f"v{VLLM_VERSION}" if "dev" not in VLLM_VERSION else "latest" - return f"\n\nAPI docs: https://docs.vllm.ai/en/{version}/api/vllm/config/#vllm.config.{cls.__name__}" + return f"\n\nAPI docs: https://docs.vllm.ai/en/{version}/api/vllm/config/#vllm.config.{name}" def _expand_json_human_readable_numbers(val: str) -> str: diff --git a/vllm/entrypoints/cli/benchmark/mm_processor.py b/vllm/entrypoints/cli/benchmark/mm_processor.py index 26b93aacdc5..de0d62e743a 100644 --- a/vllm/entrypoints/cli/benchmark/mm_processor.py +++ b/vllm/entrypoints/cli/benchmark/mm_processor.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import argparse +import inspect from vllm.benchmarks.mm_processor import add_cli_args, main from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase @@ -8,13 +9,59 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser class BenchmarkMMProcessorSubcommand(BenchmarkSubcommandBase): - """The `mm-processor` subcommand for `vllm bench`.""" + r"""`vllm bench mm-processor` profiles the multimodal input processor pipeline of + vision-language models. It measures per-stage latency from the HuggingFace + processor through to the encoder forward pass, helping you identify + preprocessing bottlenecks and understand how different image resolutions or + item counts affect end-to-end request time. + + The benchmark supports two data sources: synthetic random multimodal inputs + (`random-mm`) and HuggingFace datasets (`hf`). Warmup requests are run before + measurement to ensure stable results. + + ## Quick Start + + ```bash + vllm bench mm-processor \ + --model Qwen/Qwen2-VL-7B-Instruct \ + --dataset-name random-mm \ + --num-prompts 50 \ + --random-input-len 300 \ + --random-output-len 40 \ + --random-mm-base-items-per-request 2 \ + --random-mm-limit-mm-per-prompt '{"image": 3, "video": 0}' \ + --random-mm-bucket-config '{(256, 256, 1): 0.7, (720, 1280, 1): 0.3}' + ``` + + ## Measured Stages + + | Stage | Description | + | ----- | ----------- | + | `get_mm_hashes_secs` | Time spent hashing multimodal inputs | + | `get_cache_missing_items_secs` | Time spent looking up the processor cache | + | `apply_hf_processor_secs` | Time spent in the HuggingFace processor | + | `merge_mm_kwargs_secs` | Time spent merging multimodal kwargs | + | `apply_prompt_updates_secs` | Time spent updating prompt tokens | + | `preprocessor_total_secs` | Total preprocessing time | + | `encoder_forward_secs` | Time spent in the encoder model forward pass | + | `num_encoder_calls` | Number of encoder invocations per request | + + The benchmark also reports end-to-end latency (TTFT + decode time) per + request. Use `--metric-percentiles` to select which percentiles to report + (default: p99) and `--output-json` to save results. + + For more examples (HF datasets, warmup, JSON output), see the + [multimodal processor benchmark guide](https://docs.vllm.ai/en/latest/benchmarking/cli/#multimodal-processor-benchmark). + """ name = "mm-processor" help = "Benchmark multimodal processor latency across different configurations." @classmethod def add_cli_args(cls, parser: FlexibleArgumentParser) -> None: + # The class docstring is the page overview / `--help` description. + if cls.__doc__: + parser.description = inspect.cleandoc(cls.__doc__) add_cli_args(parser) @staticmethod diff --git a/vllm/entrypoints/cli/launch.py b/vllm/entrypoints/cli/launch.py index d13a0c67c83..91d4bf094b7 100644 --- a/vllm/entrypoints/cli/launch.py +++ b/vllm/entrypoints/cli/launch.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import argparse +import inspect import signal import uvloop @@ -36,9 +37,12 @@ class LaunchSubcommandBase(CLISubcommand): def add_cli_args(cls, parser: FlexibleArgumentParser) -> None: """Add the CLI arguments to the parser. - By default, adds the standard vLLM serving arguments. + By default, uses the subcommand's docstring as the description and adds + the standard vLLM serving arguments. Subclasses can override to add component-specific arguments. """ + if cls.__doc__: + parser.description = inspect.cleandoc(cls.__doc__) make_arg_parser(parser) @staticmethod @@ -47,7 +51,17 @@ class LaunchSubcommandBase(CLISubcommand): class RenderSubcommand(LaunchSubcommandBase): - """The `render` subcommand for `vllm launch`.""" + """`vllm launch render` starts a GPU-less rendering server for preprocessing + and postprocessing only. + + ```bash + vllm launch render meta-llama/Llama-3.2-1B-Instruct --port 8100 + ``` + + This command reuses the standard serving parser, so model, frontend, + networking, and related CLI options follow the same conventions as + [`vllm serve`](https://docs.vllm.ai/en/latest/cli/serve/). + """ name = "render" help = "Launch a GPU-less rendering server (preprocessing and postprocessing only)." @@ -93,7 +107,6 @@ class LaunchSubcommand(CLISubcommand): cmd_subparser = launch_subparsers.add_parser( cmd_cls.name, help=cmd_cls.help, - description=cmd_cls.help, usage=f"vllm {self.name} {cmd_cls.name} [options]", ) cmd_subparser.set_defaults(launch_command=cmd_cls.cmd) From 33ef67e9fb5f3f9c2dfe1aa95e9880d5ecd7b38b Mon Sep 17 00:00:00 2001 From: Canlin Guo <canlinguosdu@gmail.com> Date: Sat, 25 Jul 2026 22:48:12 +0800 Subject: [PATCH 040/185] [BugFix] Increase the max supported duration for MOSS-TD (#49403) Signed-off-by: Canlin Guo <canlinguosdu@gmail.com> --- vllm/model_executor/models/moss_transcribe_diarize.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/models/moss_transcribe_diarize.py b/vllm/model_executor/models/moss_transcribe_diarize.py index 5236f5c61be..789c789f0bb 100644 --- a/vllm/model_executor/models/moss_transcribe_diarize.py +++ b/vllm/model_executor/models/moss_transcribe_diarize.py @@ -67,6 +67,7 @@ from vllm.transformers_utils.processor import cached_processor_from_config from vllm.utils.tensor_schema import TensorSchema, TensorShape WHISPER_ENCODER_STRIDE = 2 +MAX_AUDIO_DURATION_S = 90 * 60 AUDIO_PLACEHOLDER = "<|audio_start|><|audio_pad|><|audio_end|>" @@ -145,9 +146,7 @@ def _compute_total_audio_tokens( def _get_max_audio_samples(feature_extractor: Any) -> int: - if hasattr(feature_extractor, "chunk_length"): - return int(feature_extractor.chunk_length * feature_extractor.sampling_rate) - return int(feature_extractor.n_samples) + return int(MAX_AUDIO_DURATION_S * feature_extractor.sampling_rate) def _as_audio_embedding_list(audio_embeds: object) -> list[torch.Tensor]: From 0b0bd2b5f6a7f15ef59621efe6e535b54c6348f9 Mon Sep 17 00:00:00 2001 From: fangyuchu <fangyuchu@qq.com> Date: Sat, 25 Jul 2026 23:49:09 +0800 Subject: [PATCH 041/185] [Feature] Add fault tolerance framework (simplified) for DP+EP external LB deployments (#44428) Signed-off-by: fangyuchu <fangyuchu@qq.com> Signed-off-by: a798347923 <2645302020@qq.com> Signed-off-by: TianZhuo <2770730562@qq.com> Signed-off-by: a798347923 <39047817+a798347923@users.noreply.github.com> Signed-off-by: 205150940 <112750056+205150940@users.noreply.github.com> Signed-off-by: w00689259 <wangzhuo66@huawei.com> Signed-off-by: zWaNg3 <37772915+zWaNg3@users.noreply.github.com> Signed-off-by: zWaNg3 <389750525@qq.com> Signed-off-by: yzchang-plus <1078477584@qq.com> Signed-off-by: Jade Zheng <zheng.shoujian@outlook.com> Co-authored-by: zWaNg3 <37772915+zWaNg3@users.noreply.github.com> Co-authored-by: a798347923 <2645302020@qq.com> Co-authored-by: TianZhuo <2770730562@qq.com> Co-authored-by: 205150940 <112750056+205150940@users.noreply.github.com> Co-authored-by: a798347923 <39047817+a798347923@users.noreply.github.com> Co-authored-by: w00689259 <wangzhuo66@huawei.com> Co-authored-by: zWaNg3 <389750525@qq.com> Co-authored-by: yzchang-plus <1078477584@qq.com> Co-authored-by: Jade Zheng <zheng.shoujian@outlook.com> --- .buildkite/test_areas/fault_tolerance.yaml | 26 ++ tests/test_config.py | 10 + tests/v1/fault_tolerance/__init__.py | 2 + .../test_fault_tolerance_e2e.py | 378 ++++++++++++++++++ vllm/config/__init__.py | 3 + vllm/config/fault_tolerance.py | 18 + vllm/config/parallel.py | 19 + .../device_communicators/all2all.py | 27 +- .../base_device_communicator.py | 22 +- vllm/engine/arg_utils.py | 31 ++ vllm/engine/protocol.py | 11 + vllm/entrypoints/openai/api_server.py | 7 + .../serve/fault_tolerance/__init__.py | 0 .../serve/fault_tolerance/api_router.py | 100 +++++ vllm/v1/engine/__init__.py | 8 + vllm/v1/engine/async_llm.py | 10 + vllm/v1/engine/core.py | 25 ++ vllm/v1/engine/core_client.py | 52 ++- vllm/v1/engine/utils.py | 3 - vllm/v1/fault_tolerance/__init__.py | 8 + .../fault_tolerance/engine_core_sentinel.py | 197 +++++++++ vllm/v1/fault_tolerance/utils.py | 17 + vllm/v1/worker/gpu/async_utils.py | 15 + vllm/v1/worker/gpu/model_runner.py | 8 + vllm/v1/worker/gpu_worker.py | 9 +- vllm/v1/worker/sentinel/__init__.py | 0 .../v1/worker/sentinel/gpu_worker_sentinel.py | 89 +++++ 27 files changed, 1088 insertions(+), 7 deletions(-) create mode 100644 .buildkite/test_areas/fault_tolerance.yaml create mode 100644 tests/v1/fault_tolerance/__init__.py create mode 100644 tests/v1/fault_tolerance/test_fault_tolerance_e2e.py create mode 100644 vllm/config/fault_tolerance.py create mode 100644 vllm/entrypoints/serve/fault_tolerance/__init__.py create mode 100644 vllm/entrypoints/serve/fault_tolerance/api_router.py create mode 100644 vllm/v1/fault_tolerance/__init__.py create mode 100644 vllm/v1/fault_tolerance/engine_core_sentinel.py create mode 100644 vllm/v1/fault_tolerance/utils.py create mode 100644 vllm/v1/worker/sentinel/__init__.py create mode 100644 vllm/v1/worker/sentinel/gpu_worker_sentinel.py diff --git a/.buildkite/test_areas/fault_tolerance.yaml b/.buildkite/test_areas/fault_tolerance.yaml new file mode 100644 index 00000000000..e2f700a8bd8 --- /dev/null +++ b/.buildkite/test_areas/fault_tolerance.yaml @@ -0,0 +1,26 @@ +group: Fault Tolerance +depends_on: + - image-build +steps: +- label: Fault Tolerance E2E (2xH100) + key: fault-tolerance-e2e-2xh100 + timeout_in_minutes: 35 + device: h100 + num_devices: 2 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/fault_tolerance/ + - vllm/v1/worker/sentinel/ + - vllm/entrypoints/serve/fault_tolerance/ + - vllm/distributed/elastic_ep/ + - vllm/distributed/device_communicators/ + - vllm/v1/engine/ + - vllm/v1/worker/ + - tests/v1/fault_tolerance/ + - tests/v1/distributed/test_external_lb_dp.py + commands: + # Base image has no nixl; install it or has_nixl_ep() skips the tests. + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh + # https://github.com/NVIDIA/nccl/issues/1838 + - export NCCL_CUMEM_HOST_ENABLE=0 + - pytest -v -s v1/fault_tolerance/test_fault_tolerance_e2e.py diff --git a/tests/test_config.py b/tests/test_config.py index 71e078ef3a2..1fc00a8f8a9 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1547,6 +1547,16 @@ def test_needs_dp_coordination( assert vllm_config.needs_dp_coordinator == expected_needs_coordinator +def test_fault_tolerance_requires_single_api_server(): + """Fault tolerance assumes one AsyncMPClient manages all engines, so it + is incompatible with API server scale-out (_api_process_count > 1).""" + with pytest.raises(ValueError, match="single API server"): + ParallelConfig(enable_fault_tolerance=True, _api_process_count=2) + + # Single API server (the FT-supported topology) is accepted. + ParallelConfig(enable_fault_tolerance=True, _api_process_count=1) + + def test_renderer_num_workers_with_mm_cache(): """Disallow renderer_num_workers > 1 when mm processor cache is enabled, since neither cache type is thread-safe.""" diff --git a/tests/v1/fault_tolerance/__init__.py b/tests/v1/fault_tolerance/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/tests/v1/fault_tolerance/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/v1/fault_tolerance/test_fault_tolerance_e2e.py b/tests/v1/fault_tolerance/test_fault_tolerance_e2e.py new file mode 100644 index 00000000000..f6d15343b56 --- /dev/null +++ b/tests/v1/fault_tolerance/test_fault_tolerance_e2e.py @@ -0,0 +1,378 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end tests for the elastic fault-tolerance framework. + +Requires nixl_ep FT hardware; gated behind ``has_nixl_ep()``. +""" + +import contextlib +import os +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from typing import Any + +import psutil +import pytest +import requests + +from tests.utils import RemoteOpenAIServer, multi_gpu_test +from vllm.utils.import_utils import has_nixl_ep + +MODEL_NAME = os.getenv("MODEL_NAME", "ibm-research/PowerMoE-3b") +DP_SIZE = 2 + +# Fault-detection timeout budget: +# - CPU: Gloo DP allreduce timeout (30s) detects the dead peer. +# - nixl_ep: kernel masks the dead rank after Buffer's default timeout_ms=30000 (30s). +# - Deadline (45s): slowest fallback (30s) + margin. +CPU_DISTRIBUTED_TIMEOUT_S = 30 +FAULT_DETECTION_DEADLINE_S = 45 + + +# Patches ``gpu.dp_utils.sync_cudagraph_and_dp_padding`` to raise on ``rank`` at +# a chosen step. Gated on VLLM_FT_TEST_INJECT_FAULT. +_FAULT_INJECT_SITECUSTOMIZE = """\ +import builtins +import os +import sys + +_SPEC = os.environ.get("VLLM_FT_TEST_INJECT_FAULT") +_MODULE = "vllm.v1.worker.gpu.dp_utils" +_ATTR = "sync_cudagraph_and_dp_padding" + +if _SPEC: + _f = dict(kv.split("=", 1) for kv in _SPEC.split(",")) + _RANK, _STEP = int(_f["rank"]), int(_f["step"]) + _steps = [0] + + def _patch(m): + import inspect + _orig = getattr(m, _ATTR) + _sig = inspect.signature(_orig) + def _wrapped(*args, **kwargs): + result = _orig(*args, **kwargs) + bound = _sig.bind(*args, **kwargs) + bound.apply_defaults() + dp_rank = bound.arguments.get("dp_rank") + if dp_rank == _RANK: + _steps[0] += 1 + if _steps[0] == _STEP: + raise RuntimeError( + "FT test fault injection (rank=%d step=%d)" % (_RANK, _STEP) + ) + return result + + setattr(m, _ATTR, _wrapped) + + _real_import = builtins.__import__ + + def _hook(name, *a, **k): + module = _real_import(name, *a, **k) + m = sys.modules.get(_MODULE) + # During vLLM's circular import the module lands in sys.modules before + # its functions are defined; hasattr guards against patching too early. + if ( + m is not None + and hasattr(m, _ATTR) + and not getattr(m, "_ft_patched", False) + ): + m._ft_patched = True + _patch(m) + return module + + builtins.__import__ = _hook +""" + + +def _install_fault_injection(monkeypatch, tmp_path, rank: int, step: int) -> None: + """Arrange for the DP-sync fn to raise on ``rank`` at serving ``step``. + + Writes a ``sitecustomize.py`` and prepends its dir to PYTHONPATH so every + vLLM subprocess picks it up; the fault spec is read from the environment. + """ + site_dir = tmp_path / "ft_inject" + site_dir.mkdir() + (site_dir / "sitecustomize.py").write_text(_FAULT_INJECT_SITECUSTOMIZE) + existing = os.environ.get("PYTHONPATH", "") + monkeypatch.setenv( + "PYTHONPATH", + str(site_dir) + (os.pathsep + existing if existing else ""), + ) + monkeypatch.setenv("VLLM_FT_TEST_INJECT_FAULT", f"rank={rank},step={step}") + + +def _ft_server_args() -> list[str]: + return [ + "--enforce-eager", + "--dtype", + "bfloat16", + "--max-model-len", + "2048", + "--max-num-seqs", + "128", + "--enable-expert-parallel", + "--all2all-backend", + "nixl_ep", + "--enable-fault-tolerance", + "--cpu-distributed-timeout-seconds", + str(CPU_DISTRIBUTED_TIMEOUT_S), + "--fault-tolerance-config", + '{"engine_recovery_timeout_sec": 120}', + ] + + +def _ft_manager(): + """Build the shared DP+EP fault-tolerant server topology (one engine/server).""" + from tests.v1.distributed.test_external_lb_dp import ExternalLBServerManager + + return ExternalLBServerManager( + MODEL_NAME, + DP_SIZE, + api_server_count=1, # FT requires a single API server per engine + base_server_args=_ft_server_args(), + tp_size=1, + ) + + +def _server_for_rank(servers, rank: int): + """Locate the server for a DP rank.""" + for server, sargs in servers: + if "--data-parallel-rank" in sargs: + idx = sargs.index("--data-parallel-rank") + if int(sargs[idx + 1]) == rank: + return server + raise AssertionError(f"no server found for DP rank {rank}") + + +def _complete(client): + """Issue the one standard completion the tests use everywhere.""" + return client.completions.create( + model=MODEL_NAME, + prompt="Hello, my name is", + max_tokens=5, + temperature=0.0, + timeout=10.0, + ) + + +def _in_parallel(fn, servers) -> list: + """Run ``fn(server)`` for all servers concurrently; return results in order.""" + with ThreadPoolExecutor(max_workers=len(servers)) as ex: + return list(ex.map(fn, servers)) + + +def _get_ft_status(server) -> dict: + resp = requests.get(server.url_for("fault_tolerance/status"), timeout=10) + resp.raise_for_status() + return resp.json() + + +def _assert_serving_and_healthy(servers) -> None: + """Wait until every engine is healthy, then serve one request per server.""" + healthy = _wait_for_engines( + list(servers), match_key="status", match_values={"healthy"} + ) + assert all(healthy), healthy + _in_parallel(lambda s: _complete(s.get_client()), servers) + + +def _apply_ft(server, instruction: str, params: dict | None = None) -> dict: + """POST an FT instruction; assert it is accepted (202) and return the body.""" + resp = requests.post( + server.url_for("fault_tolerance/apply"), + json={"instruction": instruction, "params": params or {}}, + timeout=10, + ) + assert resp.status_code == 202, resp.text + return resp.json() + + +def _kill_worker_process(server) -> None: + """SIGKILL only the worker proc, leaving EngineCore and API server alive.""" + workers = [ + p + for p in psutil.Process(server.proc.pid).children(recursive=True) + if "Worker" in " ".join(p.cmdline()) + ] + assert len(workers) == 1, f"expected 1 worker proc, found: {workers}" + workers[0].kill() + + +def _wait_for_engines( + servers: list[RemoteOpenAIServer], + match_key: str, + match_values: set[str], + deadline_s: int = FAULT_DETECTION_DEADLINE_S, +) -> list[dict[str, Any] | None]: + """Poll ``/fault_tolerance/status`` until each server's engine status matches. + + A server matches when its engine-status dict has ``match_key`` equal to + one of ``match_values``. Returns one engine-status dict per server. Servers still + unmatched after ``deadline_s`` get None. + """ + results: dict[int, dict[str, Any]] = {} + pending = dict(enumerate(servers)) + start = time.time() + while pending and time.time() - start < deadline_s: + for i, server in list(pending.items()): + with contextlib.suppress(Exception): + for engine_status in _get_ft_status(server)["engines"]: + if engine_status.get(match_key) in match_values: + results[i] = engine_status + del pending[i] + break + if pending: + time.sleep(1.0) + return [results.get(i) for i in range(len(servers))] + + +@contextlib.contextmanager +def _driving(*servers): + """Pump completions at each server in the background for the block's duration. + + Keeps every engine stepping into its failed component so a fault surfaces. + Errors are expected once faulted and are ignored. + """ + stop = threading.Event() + + def _drive(server): + client = server.get_client() + while not stop.is_set(): + with contextlib.suppress(Exception): + _complete(client) + time.sleep(0.2) + + threads = [threading.Thread(target=_drive, args=(s,), daemon=True) for s in servers] + for t in threads: + t.start() + try: + yield + finally: + stop.set() + for t in threads: + t.join(timeout=2) + + +def _wait_for_ft_apply_outcome(server, request_id: str, deadline_s: int) -> str | None: + """Wait until ``/fault_tolerance/status`` records the FT apply outcome.""" + engine_status = _wait_for_engines( + [server], + match_key="last_ft_request_id", + match_values={request_id}, + deadline_s=deadline_s, + )[0] + return engine_status.get("ft_error") if engine_status else None + + +@pytest.mark.skipif(not has_nixl_ep(), reason="Requires nixl_ep all2all backend") +@multi_gpu_test(num_gpus=2) +def test_injected_fault_retry_recovers_all_ranks(monkeypatch, tmp_path): + """An exception injected into the inference path drives full retry recovery. + + Injecting an exception into ``sync_cudagraph_and_dp_padding`` at a chosen + step on rank 1. + + - Rank 1 raises inside the busy loop and goes UNHEALTHY. + - Rank 0 detects the now-absent peer via the communication timeout and also + goes UNHEALTHY. + + Both being UNHEALTHY is the precondition for ``retry``. The fault is patched + into the DP-sync fn from the test (via a generated ``sitecustomize``). + """ + fault_step = int(os.getenv("FT_FAULT_STEP", "50")) + _install_fault_injection(monkeypatch, tmp_path, rank=1, step=fault_step) + + with _ft_manager() as servers: + assert len(servers) == DP_SIZE + rank0 = _server_for_rank(servers, 0) + rank1 = _server_for_rank(servers, 1) + + # 1. Both engines healthy and serving. + _assert_serving_and_healthy((rank0, rank1)) + + # 2. Drive both ranks so rank 1 accumulates execute_model steps and trips + # the injected fault; rank 0 then times out on the DP allreduce. + with _driving(rank0, rank1): + faulted = _wait_for_engines( + [rank0, rank1], match_key="status", match_values={"unhealthy"} + ) + + for rank, engine_status in enumerate(faulted): + assert engine_status is not None, ( + f"rank {rank} did not report UNHEALTHY within " + f"{FAULT_DETECTION_DEADLINE_S}s -- it likely hung" + ) + # The rank that raised carries the fault info from its own exception. + assert faulted[1] is not None + assert faulted[1].get("fault_info"), faulted[1] + + # 3. retry both engines. + for server in (rank0, rank1): + _apply_ft(server, "retry") + + # 4. Recovery completes: both engines return to healthy and serve again. + _assert_serving_and_healthy((rank0, rank1)) + + +@pytest.mark.skipif(not has_nixl_ep(), reason="Requires nixl_ep all2all backend") +@multi_gpu_test(num_gpus=2) +def test_worker_kill_survivor_unhealthy_and_dead_rejects_retry(): + """One worker kill surfaces two status transitions at once. + + SIGKILLing only rank 1's worker leaves both EngineCores alive, so the same + fault is seen two ways: + + - Survivor (rank 0): detects the dead peer via Gloo allreduce / nixl_ep + kernel timeout. Its own executor is fine, so ``on_fault`` marks it + UNHEALTHY with a ``fault_info``. + - Victim (rank 1): detects its own executor failure and marks itself DEAD. + + Recovery is gated on UNHEALTHY: the DEAD engine accepts ``retry`` at the + HTTP layer (202 = background dispatch) but rejects it in the engine, + recording the reason as ``ft_error``. + """ + with _ft_manager() as servers: + assert len(servers) == DP_SIZE + survivor = _server_for_rank(servers, 0) + victim = _server_for_rank(servers, 1) + + # 1. Confirm both engines are healthy and serving. + _assert_serving_and_healthy((survivor, victim)) + + # 2. Kill only the victim's worker; both EngineCores stay alive. + _kill_worker_process(victim) + + # 3. Drive both engines so each keeps stepping into the failed component. + with _driving(survivor, victim): + survivor_faulted, victim_faulted = _wait_for_engines( + [survivor, victim], + match_key="status", + match_values={"dead", "unhealthy"}, + ) + + assert survivor_faulted is not None, ( + "survivor did not report the peer fault within " + f"{FAULT_DETECTION_DEADLINE_S}s -- it likely hung" + ) + # The survivor's own executor is fine, so it must be UNHEALTHY, not DEAD. + assert survivor_faulted["status"] == "unhealthy", survivor_faulted + assert survivor_faulted.get("fault_info"), survivor_faulted + + assert victim_faulted is not None, ( + "victim did not report its worker's death within " + f"{FAULT_DETECTION_DEADLINE_S}s" + ) + assert victim_faulted["status"] == "dead", victim_faulted + + # 4. retry is accepted at the HTTP layer (202 = background dispatch)... + request_id = _apply_ft(victim, "retry")["request_id"] + + # 5. ...but the DEAD engine must reject it: recovery requires UNHEALTHY. + ft_error = _wait_for_ft_apply_outcome( + victim, request_id, FAULT_DETECTION_DEADLINE_S + ) + assert ft_error is not None, ( + "rejection was never recorded in /fault_tolerance/status" + ) + assert "status is DEAD" in ft_error, ft_error diff --git a/vllm/config/__init__.py b/vllm/config/__init__.py index 6070a3f8238..24a8b5a3171 100644 --- a/vllm/config/__init__.py +++ b/vllm/config/__init__.py @@ -13,6 +13,7 @@ from vllm.config.device import DeviceConfig from vllm.config.diffusion import DiffusionConfig from vllm.config.ec_manager_config import EncoderCacheManagerConfig from vllm.config.ec_transfer import ECTransferConfig +from vllm.config.fault_tolerance import FaultToleranceConfig from vllm.config.kernel import KernelConfig from vllm.config.kv_events import KVEventsConfig from vllm.config.kv_transfer import KVTransferConfig @@ -124,6 +125,8 @@ __all__ = [ "StructuredOutputsConfig", # From vllm.config.profiler "ProfilerConfig", + # From vllm.config.fault_tolerance + "FaultToleranceConfig", # From vllm.config.utils "ConfigType", "SupportsMetricsInfo", diff --git a/vllm/config/fault_tolerance.py b/vllm/config/fault_tolerance.py new file mode 100644 index 00000000000..7ed095d204f --- /dev/null +++ b/vllm/config/fault_tolerance.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +from vllm.config.utils import config + + +@config +class FaultToleranceConfig: + """Configuration for fault tolerance.""" + + engine_recovery_timeout_sec: int = 120 + """Timeout (in seconds) to wait for error handling instructions + before raising an exception. If the EngineCore encounters an + error, it waits up to this many seconds for vLLM to receive + instructions on how to handle the error and then recover from the fault. + If vLLM does not recover during this time, the original error is raised. + """ diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index ce038a99ea9..949eb298a17 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -13,6 +13,7 @@ from torch.distributed import ProcessGroup, ReduceOp, Store from typing_extensions import Self import vllm.envs as envs +from vllm.config.fault_tolerance import FaultToleranceConfig from vllm.config.utils import config from vllm.logger import init_logger from vllm.platforms import current_platform @@ -22,6 +23,7 @@ if TYPE_CHECKING: from ray.runtime_env import RuntimeEnv from ray.util.placement_group import PlacementGroup + from vllm.config.fault_tolerance import FaultToleranceConfig from vllm.v1.executor import Executor else: RuntimeEnv = Any @@ -393,6 +395,16 @@ class ParallelConfig: should only be set by API server scale-out. """ + enable_fault_tolerance: bool = False + """Enable fault tolerance for detailed error recovery, + such as scaling down fault DPEngineCore. + """ + + fault_tolerance_config: FaultToleranceConfig = Field( + default_factory=FaultToleranceConfig + ) + """The configurations for fault tolerance.""" + @field_validator("disable_nccl_for_dp_synchronization", mode="wrap") @classmethod def _skip_none_validation(cls, value: Any, handler: Callable) -> Any: @@ -445,6 +457,13 @@ class ParallelConfig: f"but found: {self._api_process_rank}" ) + if self.enable_fault_tolerance and self._api_process_count > 1: + raise ValueError( + "Fault tolerance requires a single API server process " + f"(--api-server-count=1), but got {self._api_process_count}. " + "The FT system assumes one AsyncMPClient manages all engines." + ) + if self.all2all_backend in ["pplx", "naive"]: logger.warning( "The '%s' all2all backend has been removed. " diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 679764f6a82..ee404f688a1 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -8,6 +8,7 @@ import torch import torch.distributed as dist import vllm.envs as envs +from vllm.config import get_current_vllm_config from vllm.distributed import get_dp_group, get_ep_group, get_pcp_group from vllm.distributed.utils import StatelessProcessGroup from vllm.forward_context import get_forward_context @@ -278,7 +279,9 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase): def __init__(self, cpu_group, tcp_store_group=None): super().__init__(cpu_group, tcp_store_group) - self.support_fault_tolerance = False # TODO: set to True when FT is supported. + self.support_fault_tolerance = ( + get_current_vllm_config().parallel_config.enable_fault_tolerance + ) def _make_all2all_kwargs( self, @@ -360,6 +363,16 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase): has_fault = (current != DeepEPLLAll2AllManager._last_mask).any() return has_fault + def clean_buffers(self) -> None: + buf = DeepEPLLAll2AllManager._buffer + if buf is None: + return + buf.get_local_buffer_tensor(dtype=torch.int8, use_rdma_buffer=True).zero_() + torch.accelerator.synchronize() + buf.low_latency_clean_mask_buffer() + torch.accelerator.synchronize() + DeepEPLLAll2AllManager._last_mask = None + @dataclass class _NixlEPBufferState: @@ -565,6 +578,18 @@ class NixlEPAll2AllManager(All2AllManagerBase): has_fault = (current != last).any() return has_fault + def clean_buffers(self) -> None: + if NixlEPAll2AllManager._buffer is None: + return + state = NixlEPAll2AllManager._buffer + state.buffer.get_local_buffer_tensor( + dtype=torch.int8, use_rdma_buffer=True + ).zero_() + torch.accelerator.synchronize() + state.buffer.clean_mask_buffer() + torch.accelerator.synchronize() + NixlEPAll2AllManager._last_mask = None + class FlashInferNVLinkTwoSidedManager(All2AllManagerBase): """ diff --git a/vllm/distributed/device_communicators/base_device_communicator.py b/vllm/distributed/device_communicators/base_device_communicator.py index 70f1fb5d62c..dc267143389 100644 --- a/vllm/distributed/device_communicators/base_device_communicator.py +++ b/vllm/distributed/device_communicators/base_device_communicator.py @@ -105,10 +105,30 @@ class All2AllManagerBase: raise NotImplementedError def query_active_mask(self) -> torch.Tensor: + """Return the all2all liveness mask for the EP ranks. + + Returns: + An int32 device tensor where 0 marks a live rank and 1 marks a + masked (dead/unreachable) rank. + """ raise NotImplementedError def query_fault(self) -> torch.Tensor: - """Returns has_fault scalar.""" + """Return a scalar bool tensor, True if a new fault appeared. + + Compares the current mask against the baseline recorded at the last + recovery point. + """ + raise NotImplementedError + + def clean_buffers(self) -> None: + """Reset this rank's RDMA buffers and all2all mask state (rank-local). + + Post-fault cleanup: a dispatch/combine that hit a dead peer or timed + out can leave partially-written or stale tokens in the RDMA receive + buffer, so it is zeroed to stop the next forward from reading that + contaminated data. + """ raise NotImplementedError def set_num_sms(self, num_sms: int): diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 73a724c4427..f6c65f95549 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -42,6 +42,7 @@ from vllm.config import ( DiffusionConfig, ECTransferConfig, EPLBConfig, + FaultToleranceConfig, KernelConfig, KVEventsConfig, KVTransferConfig, @@ -724,6 +725,11 @@ class EngineArgs: optimization_level: OptimizationLevel = VllmConfig.optimization_level performance_mode: PerformanceMode = VllmConfig.performance_mode + fault_tolerance_config: FaultToleranceConfig = get_field( + ParallelConfig, "fault_tolerance_config" + ) + enable_fault_tolerance: bool = ParallelConfig.enable_fault_tolerance + kv_offloading_size: float | None = CacheConfig.kv_offloading_size kv_offloading_backend: KVOffloadingBackend = CacheConfig.kv_offloading_backend tokens_only: bool = False @@ -756,6 +762,16 @@ class EngineArgs: self.weight_transfer_config = WeightTransferConfig( **self.weight_transfer_config ) + if isinstance(self.fault_tolerance_config, dict): + if not self.enable_fault_tolerance: + logger.warning( + "--fault-tolerance-config was passed. Fault tolerance is being " + "automatically enabled." + ) + self.enable_fault_tolerance = True + self.fault_tolerance_config = FaultToleranceConfig( + **self.fault_tolerance_config + ) if isinstance(self.ir_op_priority, dict): self.ir_op_priority = IrOpPriorityConfig(**self.ir_op_priority) @@ -1153,6 +1169,12 @@ class EngineArgs: parallel_group.add_argument( "--worker-extension-cls", **parallel_kwargs["worker_extension_cls"] ) + parallel_group.add_argument( + "--enable-fault-tolerance", **parallel_kwargs["enable_fault_tolerance"] + ) + parallel_group.add_argument( + "--fault-tolerance-config", **parallel_kwargs["fault_tolerance_config"] + ) # KV cache arguments cache_kwargs = get_kwargs(CacheConfig) @@ -1619,6 +1641,7 @@ class EngineArgs: def from_cli_args(cls, args: argparse.Namespace): # Get the list of attributes of this dataclass. attrs = [attr.name for attr in dataclasses.fields(cls)] + # Set the attributes from the parsed arguments. engine_args = cls( **{attr: getattr(args, attr) for attr in attrs if hasattr(args, attr)} @@ -2024,6 +2047,12 @@ class EngineArgs: data_parallel_external_lb = ( self.data_parallel_external_lb or self.data_parallel_rank is not None ) + if self.enable_fault_tolerance and not data_parallel_external_lb: + raise ValueError( + "Fault tolerance requires external load balancer mode " + "(--data-parallel-external-lb or --data-parallel-rank). " + "Internal LB mode is not supported." + ) if ( self.data_parallel_size > 1 and data_parallel_external_lb @@ -2181,6 +2210,8 @@ class EngineArgs: _api_process_count=self._api_process_count, _api_process_rank=self._api_process_rank, assigned_physical_gpu_ids=self._resolve_device_ids(), + enable_fault_tolerance=self.enable_fault_tolerance, + fault_tolerance_config=self.fault_tolerance_config, numa_bind=self.numa_bind, numa_bind_nodes=self.numa_bind_nodes, numa_bind_cpus=self.numa_bind_cpus, diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index c54123bea9e..ef3be178ac8 100644 --- a/vllm/engine/protocol.py +++ b/vllm/engine/protocol.py @@ -20,6 +20,7 @@ from vllm.sampling_params import SamplingParams from vllm.tasks import SupportedTask from vllm.v1.engine import EngineCoreRequest from vllm.v1.engine.input_processor import InputProcessor +from vllm.v1.fault_tolerance.utils import FaultToleranceRequest, FaultToleranceResult if TYPE_CHECKING: from vllm.v1.engine import PauseMode @@ -234,6 +235,16 @@ class EngineClient(ABC): """Perform a collective RPC call to the given path.""" raise NotImplementedError + async def handle_fault( + self, fault_tolerance_request: FaultToleranceRequest + ) -> FaultToleranceResult: + """send fault tolerance instruction to the engine""" + raise NotImplementedError + + async def get_status(self): + """Get fault tolerance status of all engines.""" + raise NotImplementedError + async def get_supported_tasks(self) -> tuple[SupportedTask, ...]: """Get supported tasks""" raise NotImplementedError diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index 59c7ee84cae..9103dd7fae9 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -269,6 +269,13 @@ def build_app( register_pooling_api_routers(app, supported_tasks, model_config) + if args.enable_fault_tolerance: + from vllm.entrypoints.serve.fault_tolerance.api_router import ( + register_fault_tolerance_api_router, + ) + + register_fault_tolerance_api_router(app) + # Endpoint plugins are attached last so their routes are registered after all core # routers. This runs even for the CPU only render server. A plugin eligible for # the `render` task still gets its routes registered. It receives diff --git a/vllm/entrypoints/serve/fault_tolerance/__init__.py b/vllm/entrypoints/serve/fault_tolerance/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm/entrypoints/serve/fault_tolerance/api_router.py b/vllm/entrypoints/serve/fault_tolerance/api_router.py new file mode 100644 index 00000000000..960833af215 --- /dev/null +++ b/vllm/entrypoints/serve/fault_tolerance/api_router.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json +import uuid +from http import HTTPStatus + +from fastapi import APIRouter, BackgroundTasks, Depends, FastAPI, HTTPException, Request +from fastapi.responses import JSONResponse + +from vllm.engine.protocol import EngineClient +from vllm.entrypoints.openai.engine.protocol import ErrorResponse +from vllm.entrypoints.serve.utils.api_utils import validate_json_request +from vllm.logger import init_logger +from vllm.v1.fault_tolerance.utils import FaultToleranceRequest + +logger = init_logger(__name__) + +router = APIRouter() + +_ALLOWED_INSTRUCTIONS = {"retry"} + + +def _validate_payload(body: dict) -> tuple[str, dict]: + if not isinstance(body, dict): + raise HTTPException(400, "Request body must be a JSON object.") + instruction = body.get("instruction") + if not instruction: + raise HTTPException(400, "'instruction' is required.") + if instruction not in _ALLOWED_INSTRUCTIONS: + raise HTTPException(400, f"Invalid instruction: '{instruction}'.") + params = body.get("params", {}) + if not isinstance(params, dict): + raise HTTPException(400, "'params' must be an object.") + return instruction, params + + +@router.post( + "/fault_tolerance/apply", + dependencies=[Depends(validate_json_request)], + responses={ + HTTPStatus.ACCEPTED.value: {"model": dict}, + HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, + }, +) +async def process_fault_tolerance_instruction( + raw_request: Request, background_tasks: BackgroundTasks +): + try: + body = await raw_request.json() + except json.JSONDecodeError as e: + raise HTTPException(400, "Invalid JSON format") from e + + instruction, params = _validate_payload(body) + ft_request = FaultToleranceRequest( + instruction=instruction, + params=params, + request_id=str(uuid.uuid4()), + ) + + client: EngineClient = raw_request.app.state.engine_client + # Recovery runs cross-rank collective ops that only complete once every rank + # has been dispatched. Run it in the background and return immediately so the + # orchestrator can dispatch to all ranks without blocking; completion is + # observed by polling GET /fault_tolerance/status. + background_tasks.add_task(_run_fault_recovery, client, ft_request) + return JSONResponse( + status_code=HTTPStatus.ACCEPTED.value, + content={ + "message": "Request accepted; poll /fault_tolerance/status for updates.", + "request_id": ft_request.request_id, + }, + background=background_tasks, + ) + + +async def _run_fault_recovery( + client: EngineClient, ft_request: FaultToleranceRequest +) -> None: + """Drive recovery to completion after the 202 response is sent.""" + try: + result = await client.handle_fault(ft_request) + except Exception: + logger.exception("[FT] Recovery dispatch failed.") + return + if not result.success: + logger.error( + "[FT] Recovery failed for request %s: %s", + ft_request.request_id, + result.reason, + ) + + +@router.get("/fault_tolerance/status") +async def get_status(raw_request: Request): + client: EngineClient = raw_request.app.state.engine_client + return JSONResponse(content=await client.get_status()) + + +def register_fault_tolerance_api_router(app: FastAPI): + app.include_router(router) diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index 919402a16ab..4ac27be5068 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -31,6 +31,8 @@ FINISH_REASON_STRINGS = ("stop", "length", "abort", "error", "repetition") EEP_NOTIFICATION_CALL_ID = -1 +FT_STATUS_CALL_ID = -2 + class EEPNotificationType(enum.Enum): NEW_CORE_ENGINES_INIT_READY = "NEW_CORE_ENGINES_INIT_READY" @@ -282,3 +284,9 @@ class ReconfigureRankType(enum.IntEnum): KEEP_CURRENT_RANK = -1 SHUTDOWN_CURRENT_RANK = -2 + + +class EngineStatusType(enum.IntEnum): + HEALTHY = 0 + DEAD = 1 + UNHEALTHY = 2 diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 93e02abf747..f1e7132339c 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -44,6 +44,7 @@ from vllm.v1.engine.input_processor import InputProcessor from vllm.v1.engine.output_processor import OutputProcessor, RequestOutputCollector from vllm.v1.engine.parallel_sampling import ParentRequest from vllm.v1.executor import Executor +from vllm.v1.fault_tolerance.utils import FaultToleranceRequest, FaultToleranceResult from vllm.v1.metrics.loggers import ( StatLoggerFactory, StatLoggerManager, @@ -1041,6 +1042,15 @@ class AsyncLLM(EngineClient): finally: set_scaling_elastic_ep(False) + async def handle_fault( + self, fault_tolerance_request: FaultToleranceRequest + ) -> FaultToleranceResult: + """send fault tolerance instruction to the engine""" + return await self.engine_core.handle_fault(fault_tolerance_request) + + async def get_status(self): + return await self.engine_core.get_status() + @property def is_running(self) -> bool: # Is None before the loop is started. diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 476f53d4611..8a62c8b2b1b 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -78,6 +78,11 @@ from vllm.v1.engine.utils import ( get_physical_gpu_ids_for_local_dp_rank, ) from vllm.v1.executor import Executor +from vllm.v1.fault_tolerance.engine_core_sentinel import ( + FT_UTILITY_METHOD, + EngineCoreSentinel, + fault_tolerant_wrapper, +) from vllm.v1.kv_cache_interface import KVCacheConfig, get_kv_cache_spec_kind from vllm.v1.metrics.stats import SchedulerIterationDetails, SchedulerStats from vllm.v1.outputs import ModelRunnerOutput @@ -1070,6 +1075,16 @@ class EngineCoreProc(EngineCore): internal_dp_balancing, ) + # Initialize fault tolerance settings. + self.enable_fault_tolerance = ( + vllm_config.parallel_config.enable_fault_tolerance + ) + if self.enable_fault_tolerance: + self.ft_sentinel = EngineCoreSentinel( + engine=self, + parallel_config=vllm_config.parallel_config, + ) + # Background Threads and Queues for IO. These enable us to # overlap ZMQ socket IO with GPU since they release the GIL, # and to overlap some serialization/deserialization with the @@ -1355,6 +1370,7 @@ class EngineCoreProc(EngineCore): """Returns true if shutdown has not been requested.""" return self.shutdown_state == EngineShutdownState.RUNNING + @fault_tolerant_wrapper def run_busy_loop(self): """Core busy loop of the EngineCore.""" while self._handle_shutdown(): @@ -1672,6 +1688,14 @@ class EngineCoreProc(EngineCore): except Exception: self._handle_request_preproc_error(req) continue + elif request_type == EngineCoreRequestType.UTILITY: + request = generic_decoder.decode(data_frames) + client_idx, call_id, method, args = request + if method == FT_UTILITY_METHOD: + self.ft_sentinel.handle_command( + client_idx, call_id, args[0] + ) + continue else: request = generic_decoder.decode(data_frames) @@ -2021,6 +2045,7 @@ class DPEngineCoreProc(EngineCoreProc): and self.step_counter % self.prefill_schedule_interval != 0 ) + @fault_tolerant_wrapper def run_busy_loop(self): """Core busy loop of the EngineCore for data parallel case.""" diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index bcb441e7564..f83e32096a9 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -16,6 +16,7 @@ from multiprocessing.queues import Queue from threading import Thread from typing import Any, TypeAlias, TypeVar +import msgspec import msgspec.msgpack import zmq import zmq.asyncio @@ -35,6 +36,7 @@ from vllm.utils.network_utils import ( ) from vllm.v1.engine import ( EEP_NOTIFICATION_CALL_ID, + FT_STATUS_CALL_ID, EEPNotificationType, EngineCoreOutputs, EngineCoreReadyResponse, @@ -56,6 +58,11 @@ from vllm.v1.engine.utils import ( launch_core_engines, ) from vllm.v1.executor import Executor +from vllm.v1.fault_tolerance.engine_core_sentinel import FT_UTILITY_METHOD +from vllm.v1.fault_tolerance.utils import ( + FaultToleranceRequest, + FaultToleranceResult, +) from vllm.v1.pool.late_interaction import get_late_interaction_engine_index from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder, bytestr @@ -272,6 +279,14 @@ class EngineCoreClient(ABC): ) -> list[_R]: raise NotImplementedError + async def handle_fault( + self, fault_tolerance_request: FaultToleranceRequest + ) -> FaultToleranceResult: + raise NotImplementedError + + async def get_status(self): + raise NotImplementedError + class InprocClient(EngineCoreClient): """ @@ -971,6 +986,14 @@ class AsyncMPClient(MPClient): self.client_count = client_count self.client_index = client_index self.outputs_queue = asyncio.Queue[EngineCoreOutputs | Exception]() + + # locally-cached engine status + self._engine_status: dict[int, dict] = {} + if self.vllm_config.parallel_config.enable_fault_tolerance: + self._engine_status = { + rank: {"id": rank, "status": "healthy"} + for rank in self.engine_ranks_managed + } try: # If we are running in an asyncio event loop, start the queue task. # Otherwise, it will be started lazily. If it is not started here, @@ -994,7 +1017,7 @@ class AsyncMPClient(MPClient): output_handler: ( Callable[[AsyncMPClient, EngineCoreOutputs], Awaitable[None]] | None ) = getattr(self.__class__, "process_engine_outputs", None) - _self_ref = weakref.ref(self) if output_handler else None + _self_ref = weakref.ref(self) output_socket = resources.output_socket assert output_socket is not None @@ -1025,6 +1048,14 @@ class AsyncMPClient(MPClient): asyncio.create_task( notification_callback_handler(_self, notification_data) ) + elif outputs.utility_output.call_id == FT_STATUS_CALL_ID: + _self = _self_ref() + if not _self: + return + if outputs.utility_output.result is not None: + _self._engine_status[outputs.engine_index] = ( + outputs.utility_output.result.result + ) else: _process_utility_output( outputs.utility_output, utility_results @@ -1196,6 +1227,25 @@ class AsyncMPClient(MPClient): "collective_rpc", method, timeout, args, kwargs ) + async def handle_fault( + self, ft_request: FaultToleranceRequest + ) -> FaultToleranceResult: + res = await self.call_utility_async(FT_UTILITY_METHOD, ft_request) + result = msgspec.convert(res, FaultToleranceResult) + if not result.success: + status = self._engine_status.get(self.engine_ranks_managed[0]) + if status is not None: + status["last_ft_request_id"] = result.request_id + status["ft_error"] = result.reason + return result + + async def get_status(self): + return { + "schema_version": 1, + "total_engines": len(self.engine_ranks_managed), + "engines": list(self._engine_status.values()), + } + class DPAsyncMPClient(AsyncMPClient): """Asyncio-compatible client for multi-proc, multi-engine (data parallel) diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index 093f065475a..db1896b0946 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -234,9 +234,6 @@ class CoreEngineProcManager: if exitcode != 0 and not self.manager_stopped.is_set(): self.failed_proc_name = proc.name if died_sentinels: - # Any engine exit currently triggers a shutdown. Future - # work (e.g., Elastic and fault-tolerant EP) will add finer-grained - # handling for different exit scenarios. break self.shutdown() diff --git a/vllm/v1/fault_tolerance/__init__.py b/vllm/v1/fault_tolerance/__init__.py new file mode 100644 index 00000000000..ee54b4eb2e9 --- /dev/null +++ b/vllm/v1/fault_tolerance/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from .engine_core_sentinel import EngineCoreSentinel, fault_tolerant_wrapper + +__all__ = [ + "EngineCoreSentinel", + "fault_tolerant_wrapper", +] diff --git a/vllm/v1/fault_tolerance/engine_core_sentinel.py b/vllm/v1/fault_tolerance/engine_core_sentinel.py new file mode 100644 index 00000000000..1d82dd9f743 --- /dev/null +++ b/vllm/v1/fault_tolerance/engine_core_sentinel.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""EngineCoreSentinel and fault_tolerant_wrapper for the engine core.""" + +import json +import threading +from collections.abc import Callable +from typing import TYPE_CHECKING + +import msgspec + +from vllm.config import set_current_vllm_config +from vllm.distributed import stateless_destroy_torch_distributed_process_group +from vllm.distributed.utils import stateless_init_torch_distributed_process_group +from vllm.logger import init_logger +from vllm.utils.network_utils import get_open_port +from vllm.v1.engine import ( + FT_STATUS_CALL_ID, + EngineCoreOutputs, + EngineStatusType, + UtilityOutput, +) +from vllm.v1.fault_tolerance.utils import FaultToleranceRequest, FaultToleranceResult +from vllm.v1.request import RequestStatus +from vllm.v1.serial_utils import UtilityResult, run_method + +if TYPE_CHECKING: + from vllm.v1.engine.core import EngineCoreProc + +logger = init_logger(__name__) + +FT_UTILITY_METHOD = "handle_fault_tolerance" + + +class EngineCoreSentinel: + """Manages fault tolerance state for a single engine core.""" + + def __init__(self, engine: "EngineCoreProc", parallel_config): + self.engine = engine + self.engine_index = engine.engine_index + self.parallel_config = parallel_config + ft_config = parallel_config.fault_tolerance_config + self.engine_recovery_timeout_sec = ft_config.engine_recovery_timeout_sec + + self.resumed = threading.Event() + self.resumed.set() + self.status_type = EngineStatusType.HEALTHY + self.fault_info: str | None = None + self._dp_reinit_epoch = 0 + + def handle_command(self, client_idx: int, call_id: int, ft_args: dict): + """Dispatch an FT command by instruction name.""" + ft_request = FaultToleranceRequest(**ft_args) + if self.status_type != EngineStatusType.UNHEALTHY: + reason = ( + f"[FT] Rejecting {ft_request.instruction} on engine " + f"{self.engine_index}: status is {self.status_type.name}" + ) + logger.warning(reason) + result = FaultToleranceResult( + request_id=ft_request.request_id, + success=False, + reason=reason, + ) + else: + try: + result = run_method(self, ft_request.instruction, (ft_request,), {}) + except Exception as e: + logger.exception("[FT] Instruction '%s' failed", ft_request.instruction) + result = FaultToleranceResult( + request_id=ft_request.request_id, success=False, reason=str(e) + ) + + uo = UtilityOutput(call_id) + uo.result = UtilityResult(msgspec.structs.asdict(result)) + self.engine.output_queue.put_nowait( + (client_idx, EngineCoreOutputs(utility_output=uo)) + ) + + def on_fault(self, exc: Exception): + """Called by the wrapper when the busy loop raises an exception.""" + self.resumed.clear() + logger.warning( + "[FT] Busy loop raised %s. Waiting for recovery.", type(exc).__name__ + ) + + engine = self.engine + aborted = engine.scheduler.finish_requests(None, RequestStatus.FINISHED_ABORTED) + engine._send_abort_outputs(aborted) + if engine.batch_queue is not None: + engine.batch_queue.clear() + if ( + hasattr(engine.model_executor, "is_failed") + and engine.model_executor.is_failed + ): + self.status_type = EngineStatusType.DEAD + else: + self.status_type = EngineStatusType.UNHEALTHY + self.fault_info = f"{type(exc).__name__}" + logger.info( + "[FT] Engine %d status -> %s:", + self.engine_index, + self.status_type.name, + exc_info=exc, + ) + self._push_status() + + def _push_status(self): + """Push current health to the client so it can refresh its cache.""" + payload = {"id": self.engine_index, "status": self.status_type.name.lower()} + if self.status_type == EngineStatusType.UNHEALTHY: + payload["fault_info"] = self.fault_info + outputs = EngineCoreOutputs( + utility_output=UtilityOutput( + call_id=FT_STATUS_CALL_ID, + result=UtilityResult(payload), + ) + ) + outputs.engine_index = self.engine_index + self.engine.output_queue.put_nowait((0, outputs)) + + def retry(self, ft_request: FaultToleranceRequest) -> FaultToleranceResult: + engine = self.engine + executor = engine.model_executor + + with set_current_vllm_config(engine.vllm_config): + ft_request.params.update(self._reinit_dp_group()) + if hasattr(engine, "step_counter"): + engine.step_counter = 0 + + executor.collective_rpc("handle_ft_command", args=(ft_request,)) + + self.status_type = EngineStatusType.HEALTHY + logger.info("[FT] Engine %d status -> HEALTHY", self.engine_index) + self.resumed.set() + self._push_status() + return FaultToleranceResult(request_id=ft_request.request_id, success=True) + + def _reinit_dp_group(self) -> dict: + """Reinit DP process group if in DP mode. Returns worker params.""" + engine = self.engine + if not hasattr(engine, "dp_group") or not hasattr(engine, "dp_store"): + return {} + + parallel_config = engine.vllm_config.parallel_config + worker_key = f"ft_worker_dp_ports_{self._dp_reinit_epoch}" + engine_key = f"ft_engine_dp_port_{self._dp_reinit_epoch}" + self._dp_reinit_epoch += 1 + + if parallel_config.data_parallel_rank == 0: + worker_ports = [get_open_port() for _ in range(parallel_config.world_size)] + engine_port = get_open_port() + engine.dp_store.set(worker_key, json.dumps(worker_ports).encode()) + engine.dp_store.set(engine_key, str(engine_port).encode()) + else: + worker_ports = json.loads(engine.dp_store.get(worker_key).decode()) + engine_port = int(engine.dp_store.get(engine_key).decode()) + + stateless_destroy_torch_distributed_process_group(engine.dp_group) + engine.dp_group, engine.dp_store = ( + stateless_init_torch_distributed_process_group( + parallel_config.data_parallel_master_ip, + engine_port, + parallel_config.data_parallel_rank, + parallel_config.data_parallel_size, + backend="gloo", + return_store=True, + ) + ) + return {"new_stateless_dp_group_ports": worker_ports} + + +def fault_tolerant_wrapper(busy_loop_func: Callable): + """Wrap the busy loop to catch faults and delegate recovery.""" + + def run_with_fault_tolerance(self: "EngineCoreProc"): + while True: + try: + busy_loop_func(self) + except SystemExit: + raise + except Exception as exc: + if not self.enable_fault_tolerance: + raise + self.ft_sentinel.on_fault(exc) + recovered = self.ft_sentinel.resumed.wait( + timeout=self.ft_sentinel.engine_recovery_timeout_sec + ) + if recovered: + continue + logger.error( + "[FT] No recovery within %ds timeout.", + self.ft_sentinel.engine_recovery_timeout_sec, + ) + raise + + return run_with_fault_tolerance diff --git a/vllm/v1/fault_tolerance/utils.py b/vllm/v1/fault_tolerance/utils.py new file mode 100644 index 00000000000..0c1b1689b01 --- /dev/null +++ b/vllm/v1/fault_tolerance/utils.py @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any + +import msgspec + + +class FaultToleranceResult(msgspec.Struct): + request_id: str + success: bool + reason: str | None = None + + +class FaultToleranceRequest(msgspec.Struct): + instruction: str + params: dict[str, Any] + request_id: str = "" diff --git a/vllm/v1/worker/gpu/async_utils.py b/vllm/v1/worker/gpu/async_utils.py index e4659104f49..4570d726734 100644 --- a/vllm/v1/worker/gpu/async_utils.py +++ b/vllm/v1/worker/gpu/async_utils.py @@ -5,6 +5,7 @@ import contextlib import numpy as np import torch +from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager from vllm.v1.outputs import AsyncModelRunnerOutput, LogprobsTensors, ModelRunnerOutput from vllm.v1.worker.gpu.sample.output import SamplerOutput @@ -17,6 +18,7 @@ class AsyncOutput(AsyncModelRunnerOutput): num_sampled_tokens: torch.Tensor, main_stream: torch.cuda.Stream, copy_stream: torch.cuda.Stream, + check_ep_fault: bool = False, ): # NOTE(woosuk): We must retain references to the GPU tensors, # as the copy operations are performed on a different CUDA stream than @@ -26,6 +28,7 @@ class AsyncOutput(AsyncModelRunnerOutput): self.num_sampled_tokens = num_sampled_tokens # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. self.copy_event = torch.cuda.Event(blocking=True) + self._has_fault: torch.Tensor | None = None with stream(copy_stream, main_stream): copy_stream.wait_stream(main_stream) @@ -44,6 +47,9 @@ class AsyncOutput(AsyncModelRunnerOutput): k: v.to_cpu_nonblocking() if v is not None else None for k, v in self.model_runner_output.prompt_logprobs_dict.items() } + if check_ep_fault: + has_fault = get_ep_all2all_manager().query_fault() + self._has_fault = has_fault.to("cpu", non_blocking=True) self.copy_event.record(copy_stream) def get_output(self) -> ModelRunnerOutput: @@ -67,6 +73,15 @@ class AsyncOutput(AsyncModelRunnerOutput): if self.logprobs_tensors is not None: self.model_runner_output.logprobs = self.logprobs_tensors.tolists() self.model_runner_output.prompt_logprobs_dict = self.prompt_logprobs_dict + + if self._has_fault is not None and self._has_fault.item(): + mask = get_ep_all2all_manager().query_active_mask() + raise RuntimeError( + "Fault detected in EP all2all communication: " + "one or more ranks timed out during dispatch/combine. " + f"Mask: {mask.cpu().tolist()}" + ) + return self.model_runner_output diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 3f86b5595cb..fdedbfb86d0 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -38,6 +38,7 @@ from vllm.distributed.parallel_state import ( ) from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( initialize_mamba_ssu_backend, ) @@ -175,6 +176,12 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.dp_size = self.parallel_config.data_parallel_size self.dp_rank = self.parallel_config.data_parallel_rank + # Detect EP all2all peer faults to prevent emitting corrupted output. + # Only meaningful for MoE + DP with an FT-capable all2all backend. + self.check_ep_fault = False + if self.dp_size > 1 and self.model_config.is_moe: + self.check_ep_fault = get_ep_all2all_manager().support_fault_tolerance + # Decode context parallelism. self.dcp_size = self.parallel_config.decode_context_parallel_size self.use_dcp = self.dcp_size > 1 @@ -1488,6 +1495,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): num_sampled_tokens=num_sampled, main_stream=self.main_stream, copy_stream=self.output_copy_stream, + check_ep_fault=self.check_ep_fault, ) mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index ca63e0a117c..7c3d0336688 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -70,6 +70,7 @@ from vllm.v1.outputs import ( ModelRunnerOutput, ) from vllm.v1.utils import compute_iteration_details, report_usage_stats +from vllm.v1.worker.sentinel.gpu_worker_sentinel import WorkerSentinel from vllm.v1.worker.startup_plan import ( maybe_apply_startup_plan, maybe_save_startup_plan, @@ -146,7 +147,9 @@ class Worker(WorkerBase): from vllm.distributed.elastic_ep.elastic_execute import ElasticEPScalingExecutor self.elastic_ep_executor = ElasticEPScalingExecutor(self) - + self.worker_sentinel: WorkerSentinel | None = None + if self.parallel_config.enable_fault_tolerance: + self.worker_sentinel = WorkerSentinel(worker=self) # Buffers saved before sleep self._sleep_saved_buffers: dict[str, torch.Tensor] = {} self._sleep_rebuild_draft_metadata_buffers = False @@ -414,6 +417,10 @@ class Worker(WorkerBase): # If usage stat is enabled, collect relevant info. report_usage_stats(self.vllm_config) + def handle_ft_command(self, ft_request): + assert self.worker_sentinel is not None + return self.worker_sentinel.handle_command(ft_request) + # FIXME(youkaichao & ywang96): Use TorchDispatchMode instead of memory pool # to hijack tensor allocation. def load_model(self, *, load_dummy_weights: bool = False) -> None: diff --git a/vllm/v1/worker/sentinel/__init__.py b/vllm/v1/worker/sentinel/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm/v1/worker/sentinel/gpu_worker_sentinel.py b/vllm/v1/worker/sentinel/gpu_worker_sentinel.py new file mode 100644 index 00000000000..80050cf9d39 --- /dev/null +++ b/vllm/v1/worker/sentinel/gpu_worker_sentinel.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import TYPE_CHECKING, cast + +import torch + +from vllm.config import set_current_vllm_config +from vllm.distributed import ( + get_dp_group, + stateless_destroy_torch_distributed_process_group, + stateless_init_torch_distributed_process_group, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager +from vllm.v1.fault_tolerance.utils import FaultToleranceRequest +from vllm.v1.serial_utils import run_method + +if TYPE_CHECKING: + from vllm.v1.worker.gpu.model_runner import GPUModelRunner as GPUModelRunnerV2 + from vllm.v1.worker.gpu_worker import Worker + +logger = init_logger(__name__) + +# All2all backends that support fault-tolerant timeout + rank masking, +# required for FT under DP+EP MoE deployments. +FT_BACKEND_SET = frozenset({"deepep_low_latency", "nixl_ep"}) + + +class WorkerSentinel: + """Holds FT state for a single worker (mask tensors, DP config). + + Methods are called via collective_rpc from EngineCoreSentinel. + """ + + def __init__(self, worker: "Worker"): + self.worker = worker + self.dp_rank = worker.parallel_config.data_parallel_rank + self.dp_size = worker.parallel_config.data_parallel_size + self.data_parallel_master_ip = worker.parallel_config.data_parallel_master_ip + all2all_backend = worker.parallel_config.all2all_backend + if all2all_backend not in FT_BACKEND_SET: + raise ValueError( + f"Fault tolerance requires an FT-capable all2all backend " + f"(one of {sorted(FT_BACKEND_SET)}), but got '{all2all_backend}'." + ) + + def handle_command(self, ft_request: FaultToleranceRequest): + """Dispatch an FT command by instruction name.""" + with set_current_vllm_config(self.worker.vllm_config): + return run_method(self, ft_request.instruction, (ft_request,), {}) + + def retry(self, ft_request: FaultToleranceRequest): + torch.accelerator.synchronize() + params = ft_request.params + self._clean_worker_state() + if self.dp_size > 1: + get_ep_all2all_manager().clean_buffers() + old_cpu_group = get_dp_group().cpu_group + stateless_destroy_torch_distributed_process_group(old_cpu_group) + world_size = self.worker.parallel_config.world_size + port = params["new_stateless_dp_group_ports"][self.worker.rank % world_size] + get_dp_group().cpu_group = stateless_init_torch_distributed_process_group( + self.data_parallel_master_ip, + port, + self.dp_rank, + self.dp_size, + backend="gloo", + ) + + def _clean_worker_state(self): + model_runner = self.worker.model_runner + model_runner.execute_model_state = None + if self.worker.use_v2_model_runner: + runner = cast("GPUModelRunnerV2", model_runner) + for req_id in list(runner.req_states.req_id_to_index): + runner._remove_request(req_id) + else: + model_runner.kv_connector_output = None + + input_batch = model_runner.input_batch + cached_req_ids = list(input_batch.req_id_to_index) + for req_id in cached_req_ids: + model_runner.requests.pop(req_id, None) + model_runner.num_prompt_logprobs.pop(req_id, None) + input_batch.remove_request(req_id) + + input_batch.condense() + input_batch.refresh_metadata() + input_batch.req_prompt_embeds.clear() From 2e0da2415052b50c9a964b09ccd5779567da9f82 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Sat, 25 Jul 2026 17:20:24 +0100 Subject: [PATCH 042/185] Mergify message not on cancelled (#45117) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- .github/mergify.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/mergify.yml b/.github/mergify.yml index 4e13588eea1..5615e78d0c2 100644 --- a/.github/mergify.yml +++ b/.github/mergify.yml @@ -19,6 +19,7 @@ pull_request_rules: description: Comment on PR when pre-commit check fails conditions: - check-failure=pre-commit + - -check-cancelled=pre-commit - -closed - -draft - or: From 7fe6d3c76b6f44b54ba629e9accf572adfe86f6a Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Sat, 25 Jul 2026 12:36:19 -0400 Subject: [PATCH 043/185] [Perf] Fix moe `reduce_scatter` perf regression by removing additional comm, 5% E2E throughput gain back. (#48763) Signed-off-by: yewentao256 <zhyanwentao@126.com> --- vllm/model_executor/models/deepseek_mtp.py | 28 +++------------------- vllm/models/deepseek_v32/nvidia/mtp.py | 22 ++++++++--------- 2 files changed, 13 insertions(+), 37 deletions(-) diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index 65f860a4a3f..35cfce3cb98 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/vllm/model_executor/models/deepseek_mtp.py @@ -44,25 +44,6 @@ from .utils import ( ) -def _restore_full_token_layout_if_needed( - hidden_states: torch.Tensor, - residual: torch.Tensor, - num_tokens: int, - is_sequence_parallel: bool = False, -) -> tuple[torch.Tensor, torch.Tensor]: - """Restore full token rows for the MTP proposer after SP MoE layers.""" - if not is_sequence_parallel and hidden_states.shape[0] == num_tokens: - return hidden_states, residual - - combined_states = torch.cat([hidden_states, residual], dim=-1) - combined_states = tensor_model_parallel_all_gather(combined_states, 0) - combined_states = combined_states[:num_tokens] - hidden_states, residual = combined_states.split( - [hidden_states.shape[-1], residual.shape[-1]], dim=-1 - ) - return hidden_states, residual - - class SharedHead(nn.Module): def __init__( self, @@ -142,13 +123,10 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module): hidden_states=hidden_states, residual=None, ) - hidden_states, residual = _restore_full_token_layout_if_needed( - hidden_states, - residual, - positions.shape[0], - is_sequence_parallel=self.mtp_block.use_sequence_parallel_moe, - ) hidden_states = residual + hidden_states # pre-final-norm (logits hidden) + if self.mtp_block.use_sequence_parallel_moe: + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) + hidden_states = hidden_states[: positions.shape[0]] # Recycle the post-final-norm hidden into the next draft step. # compute_logits applies shared_head (== final norm) to the pre-norm # element, so logits and the recycle each get exactly one final-norm. diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index c8f2fcc5ffe..7130bb5d199 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -8,7 +8,10 @@ import torch.nn as nn from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig -from vllm.distributed import tensor_model_parallel_all_reduce +from vllm.distributed import ( + tensor_model_parallel_all_gather, + tensor_model_parallel_all_reduce, +) from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) @@ -24,10 +27,7 @@ from vllm.model_executor.model_loader.weight_utils import ( default_weight_loader, maybe_remap_kv_scale_name, ) -from vllm.model_executor.models.deepseek_mtp import ( - SharedHead, - _restore_full_token_layout_if_needed, -) +from vllm.model_executor.models.deepseek_mtp import SharedHead from vllm.model_executor.models.deepseek_v2 import ( DeepseekV2MixtureOfExperts, DeepseekV2MoE, @@ -95,13 +95,8 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module): hidden_states, residual = self.mtp_block( positions=positions, hidden_states=hidden_states, residual=None ) - hidden_states, residual = _restore_full_token_layout_if_needed( - hidden_states, - residual, - positions.shape[0], - is_sequence_parallel=self.mtp_block.use_sequence_parallel_moe, - ) - if not self.mtp_block.use_sequence_parallel_moe: + is_sequence_parallel = self.mtp_block.use_sequence_parallel_moe + if not is_sequence_parallel: # Without sequence parallelism, the MoE output is left un-reduced. hidden_states = tensor_model_parallel_all_reduce(hidden_states) # Recycle the POST-final-norm hidden into the next draft step. The @@ -115,6 +110,9 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module): # the legacy proposer (model_returns_tuple is True for the # DeepSeekMTPModel architecture). hidden_states, _ = self.shared_head.norm(hidden_states, residual) + if is_sequence_parallel: + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) + hidden_states = hidden_states[: positions.shape[0]] return hidden_states, hidden_states From 26d725c334429dd86b3d1a9271fbb5e16e03c9ef Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Sat, 25 Jul 2026 17:54:15 +0100 Subject: [PATCH 044/185] [Model] Add VaultGemma via Transformers modeling backend (#49803) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- docs/models/supported_models.md | 1 + tests/models/registry.py | 1 + vllm/engine/arg_utils.py | 8 +++----- vllm/model_executor/models/qwen2.py | 4 ++-- vllm/model_executor/models/registry.py | 1 + vllm/transformers_utils/config.py | 10 ---------- 6 files changed, 8 insertions(+), 17 deletions(-) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 566a3e1f955..0836f11e267 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -464,6 +464,7 @@ Some models are supported only via the [Transformers modeling backend](#transfor | `Olmo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ | | `SmolLM3ForCausalLM` | SmolLM3 | `HuggingFaceTB/SmolLM3-3B` | ✅︎ | ✅︎ | | `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | ✅︎ | ✅︎ | +| `VaultGemmaForCausalLM` | VaultGemma | `google/vaultgemma-1b` | ✅︎ | ✅︎ | !!! note Currently, the ROCm version of vLLM supports Mistral and Mixtral only for context lengths up to 4096. diff --git a/tests/models/registry.py b/tests/models/registry.py index f55daa51033..00a84783bf1 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -557,6 +557,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "TeleFLMForCausalLM": _HfExamplesInfo( "CofeAI/FLM-2-52B-Instruct-2407", trust_remote_code=True ), + "VaultGemmaForCausalLM": _HfExamplesInfo("google/vaultgemma-1b"), "Zamba2ForCausalLM": _HfExamplesInfo("Zyphra/Zamba2-7B-instruct"), "MiMoForCausalLM": _HfExamplesInfo("XiaomiMiMo/MiMo-7B-RL", trust_remote_code=True), "MiMoV2FlashForCausalLM": _HfExamplesInfo( diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index f6c65f95549..b7c5f746545 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -102,10 +102,7 @@ from vllm.logger import init_logger, suppress_logging from vllm.platforms import CpuArchEnum, current_platform from vllm.plugins import load_general_plugins from vllm.ray.lazy_utils import is_in_ray_actor, is_ray_initialized -from vllm.transformers_utils.config import ( - is_interleaved, - maybe_override_with_speculators, -) +from vllm.transformers_utils.config import maybe_override_with_speculators from vllm.transformers_utils.repo_utils import get_model_path from vllm.transformers_utils.utils import is_cloud_storage from vllm.utils.argparse_utils import ( @@ -1908,7 +1905,8 @@ class EngineArgs: self._set_default_chunked_prefill_and_prefix_caching_args(model_config) self._set_default_reasoning_config_args() sliding_window: int | None = None - if not is_interleaved(model_config.hf_text_config): + layer_types = getattr(model_config.hf_text_config, "layer_types", None) + if layer_types is None or all(lt == "sliding_attention" for lt in layer_types): # Only set CacheConfig.sliding_window if the model is all sliding # window. Otherwise CacheConfig.sliding_window will override the # global layers in interleaved sliding window models. diff --git a/vllm/model_executor/models/qwen2.py b/vllm/model_executor/models/qwen2.py index 182b9758308..f1d5f23a264 100644 --- a/vllm/model_executor/models/qwen2.py +++ b/vllm/model_executor/models/qwen2.py @@ -55,7 +55,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.config import is_interleaved, set_default_rope_theta +from vllm.transformers_utils.config import set_default_rope_theta from vllm.v1.attention.backend import AttentionType from .interfaces import ( @@ -345,7 +345,7 @@ class Qwen2Model(nn.Module, EagleModelMixin): quant_config = vllm_config.quant_config # TODO (@robertgshaw2): see if this can be moved out - if is_interleaved(vllm_config.model_config.hf_text_config): + if len(set(getattr(config, "layer_types", []))) > 1: assert config.max_window_layers == config.num_hidden_layers, ( "Sliding window for some but all layers is not supported. " "This model uses sliding window but `max_window_layers` = {} " diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index c130a7f4a2c..1dcacc7936e 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -659,6 +659,7 @@ _TRANSFORMERS_SUPPORTED_MODELS = { "Olmo2ForCausalLM": ("transformers", "TransformersForCausalLM"), "SmolLM3ForCausalLM": ("transformers", "TransformersForCausalLM"), "Starcoder2ForCausalLM": ("transformers", "TransformersForCausalLM"), + "VaultGemmaForCausalLM": ("transformers", "TransformersForCausalLM"), # Multimodal models "Emu3ForConditionalGeneration": ( "transformers", diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index f900eedf8d3..6f8c25769ef 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -600,16 +600,6 @@ def is_encoder_decoder(config: PretrainedConfig) -> bool: return _is_encoder_decoder(config) or _is_encoder_decoder(config.get_text_config()) -def is_interleaved(config: PretrainedConfig) -> bool: - """ - Detect if the model with this config is used with interleaved attention. - """ - text_config = config.get_text_config() - if layer_types := getattr(text_config, "layer_types", None): - return len(set(layer_types)) > 1 - return False - - def _maybe_update_auto_config_kwargs(kwargs: dict[str, Any], model_type: str): """ Update kwargs for AutoConfig initialization based on model_type From 9321aff536c1a73f0a7fefc68aef1f2630a904cc Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Sat, 25 Jul 2026 18:56:32 +0100 Subject: [PATCH 045/185] [Bugfix] Wait for the linear bias before layerwise online processing (#49805) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> --- .../model_loader/test_reload.py | 54 +++++++++++++++++++ .../model_loader/reload/layerwise.py | 4 +- .../model_loader/reload/meta.py | 11 ++-- .../model_loader/reload/utils.py | 8 +-- 4 files changed, 67 insertions(+), 10 deletions(-) diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index 1480191083e..e390f8d25a7 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -10,9 +10,11 @@ from torch.nn.parameter import UninitializedParameter import vllm.model_executor.model_loader.reload.meta as reload_meta from vllm.model_executor.layers.linear import QKVParallelLinear +from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase from vllm.model_executor.model_loader.reload.layerwise import ( finalize_layerwise_reload, initialize_layerwise_reload, + initialize_online_processing, record_metadata_for_reloading, ) from vllm.model_executor.model_loader.reload.meta import ( @@ -278,6 +280,58 @@ def test_layerwise_reload_composed_loader_does_not_drop_params(monkeypatch): assert torch.equal(layer.D, loaded["D"]) +class _RecordingQuantMethod(QuantizeMethodBase): + """Records the layer's bias at the moment processing runs.""" + + uses_meta_device = True + + def __init__(self): + self.bias_at_process = None + + def create_weights(self, layer, *weight_args, **extra_weight_attrs): + pass + + def apply(self, layer, *args, **kwargs): + raise NotImplementedError + + def process_weights_after_loading(self, layer): + self.bias_at_process = layer.bias.detach().clone() + + +class _LateBiasLayer(torch.nn.Module): + """Mimics an online-quantized linear: `weight` is created on meta by + `create_weights()`, which wraps the loaders, and the linear base registers + `bias` afterwards.""" + + def __init__(self, quant_method): + super().__init__() + self.quant_method = quant_method + weight = torch.nn.Parameter(torch.empty(4, 2, device="meta")) + weight.weight_loader = default_weight_loader + self.register_parameter("weight", weight) + initialize_online_processing(self) + bias = torch.nn.Parameter(torch.zeros(4)) + bias.weight_loader = default_weight_loader + self.register_parameter("bias", bias) + + +def test_online_processing_waits_for_late_registered_bias(): + # Regression test: `bias` is skipped by the meta device paths, but it is + # still loaded by a weight loader. Excluding it from the processing trigger + # finalized the layer one load early, so the trailing bias was written into + # an already-processed layer (e.g. over FP8 Marlin's permuted bias). + quant_method = _RecordingQuantMethod() + layer = _LateBiasLayer(quant_method) + loaded_bias = torch.full((4,), 3.0) + + layer.weight.weight_loader(layer.weight, torch.full((4, 2), 2.0)) + assert quant_method.bias_at_process is None + + layer.bias.weight_loader(layer.bias, loaded_bias) + assert quant_method.bias_at_process is not None + assert torch.equal(quant_method.bias_at_process, loaded_bias) + + def test_layerwise_reload_skips_non_persistent_parameter_alias_buffers(monkeypatch): layer = _AliasedBufferLayer() model = torch.nn.Sequential(layer) diff --git a/vllm/model_executor/model_loader/reload/layerwise.py b/vllm/model_executor/model_loader/reload/layerwise.py index 92f454f9a5f..d6d4eb4b930 100644 --- a/vllm/model_executor/model_loader/reload/layerwise.py +++ b/vllm/model_executor/model_loader/reload/layerwise.py @@ -14,7 +14,7 @@ from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBa from vllm.model_executor.model_loader.weight_utils import default_weight_loader from .meta import ( - SKIP_TENSORS, + SKIP_LOAD_TENSORS, capture_layer_to_meta, get_numel_loaded, materialize_layer, @@ -140,7 +140,7 @@ def _wrap_parameters_weight_loader(layer: torch.nn.Module) -> None: """Wrap each parameter's weight loader.""" # Note that nested wrapping will occur for shared tensors for name, tensor in get_layer_tensors(layer).items(): - if name in SKIP_TENSORS: + if name in SKIP_LOAD_TENSORS: continue if _get_weight_loader(tensor).__name__ != "online_process_loader": tensor.weight_loader = make_online_process_loader(layer, name) diff --git a/vllm/model_executor/model_loader/reload/meta.py b/vllm/model_executor/model_loader/reload/meta.py index 1ccf2a34c60..213a9aaeaf8 100644 --- a/vllm/model_executor/model_loader/reload/meta.py +++ b/vllm/model_executor/model_loader/reload/meta.py @@ -20,20 +20,23 @@ __all__ = [ "get_numel_loaded", ] +# Modules whose tensors are never moved to, or materialized from, the meta device. SKIP_MODULES: set[str] = {"HadamardTransform"} -SKIP_TENSORS: set[str] = { +# Tensors never loaded by a weight loader, so the layerwise trigger ignores them. +SKIP_LOAD_TENSORS: set[str] = { "_expert_map", "expert_mask", "expert_global_to_physical", "expert_physical_to_global", "expert_local_to_global", "e_score_correction_bias", - # Built after create_weights(), so it is not tracked by the layerwise-reload - # trigger and would be re-materialized into uninitialized memory. Skip it. - "bias", } +# Tensors which are never moved to, or materialized from, the meta device. +# `bias` is built after create_weights(), so it is never on meta to begin with. +SKIP_TENSORS: set[str] = SKIP_LOAD_TENSORS | {"bias"} + def to_meta_tensor(tensor: torch.Tensor) -> torch.Tensor: """Convert a tensor to a meta tensor while preserving class and attributes.""" diff --git a/vllm/model_executor/model_loader/reload/utils.py b/vllm/model_executor/model_loader/reload/utils.py index f0078d0f9d8..b6d129a7940 100644 --- a/vllm/model_executor/model_loader/reload/utils.py +++ b/vllm/model_executor/model_loader/reload/utils.py @@ -33,15 +33,15 @@ def get_layer_params_buffers(layer: torch.nn.Module) -> LayerTensors: def get_layer_size(layer: torch.nn.Module) -> int: """Calculate total number of elements across loadable tensors in a layer. - Excludes SKIP_TENSORS (e.g. _expert_map) which are never moved to meta - device and never loaded via weight_loader during layerwise reload. + Excludes SKIP_LOAD_TENSORS (e.g. _expert_map) which are never loaded via + weight_loader during layerwise reload. """ - from .meta import SKIP_TENSORS + from .meta import SKIP_LOAD_TENSORS return sum( tensor.numel() for name, tensor in get_layer_tensors(layer).items() - if name not in SKIP_TENSORS + if name not in SKIP_LOAD_TENSORS ) From 6b0103d1c986819f5ce6ebf7d4edf055873ad7a8 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim <taneem.ibrahim@gmail.com> Date: Sat, 25 Jul 2026 14:36:50 -0400 Subject: [PATCH 046/185] [CI] Stabilize Pooling Rerank Equivalence Test (#49822) --- .../pooling/scoring/test_cross_encoder_online_vision.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py b/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py index 8c1d75cd762..aae69fe217d 100644 --- a/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py +++ b/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py @@ -473,7 +473,7 @@ async def test_rerank_api_instruction_field( async def test_rerank_api_instruction_field_matches_chat_template_kwargs( server: tuple[RemoteOpenAIServer, str], ): - remote_server, _ = server + remote_server, backend = server doc_list = [ document, @@ -514,4 +514,6 @@ async def test_rerank_api_instruction_field_matches_chat_template_kwargs( kwargs_scores = [ r.relevance_score for r in sorted(kwargs_rerank.results, key=lambda x: x.index) ] - assert field_scores == pytest.approx(kwargs_scores) + assert field_scores == pytest.approx( + kwargs_scores, rel=get_tol(backend), abs=get_abs_tol(backend) + ) From ee1d99636702b8ed9ad88c2c2b833d331dce01c1 Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith <tlrmchlsmth@gmail.com> Date: Sat, 25 Jul 2026 15:29:38 -0400 Subject: [PATCH 047/185] [Build] Fix for DeepEP manylinux pidfd sycall usage (#49814) --- tools/ep_kernels/install_python_libraries.sh | 27 ++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/tools/ep_kernels/install_python_libraries.sh b/tools/ep_kernels/install_python_libraries.sh index 94beef897d0..739f031c9ef 100755 --- a/tools/ep_kernels/install_python_libraries.sh +++ b/tools/ep_kernels/install_python_libraries.sh @@ -170,6 +170,33 @@ do_build() { sed -i "s|f'{nvshmem_dir}/include']|f'{nvshmem_dir}/include', '${CUDA_HOME}/include/cccl']|" "setup.py" fi + # DeepEPv2 requires Linux 5.6+ at runtime for pidfd_getfd (pidfd_open was + # added in Linux 5.3), but manylinux headers predate both definitions. + # DeepEP is built as a separate wheel for the vLLM container image and is + # not included in the vLLM wheel, so this does not change its manylinux ABI. + if [[ "$name" == "DeepEP" ]] && \ + ! grep -q "vLLM manylinux syscall compatibility" \ + csrc/kernels/backend/symmetric.hpp; then + sed -i '1i\ +// vLLM manylinux syscall compatibility\ +#if defined(__x86_64__) || defined(__aarch64__)\ +#ifndef SYS_pidfd_open\ +#ifdef __NR_pidfd_open\ +#define SYS_pidfd_open __NR_pidfd_open\ +#else\ +#define SYS_pidfd_open 434\ +#endif\ +#endif\ +#ifndef SYS_pidfd_getfd\ +#ifdef __NR_pidfd_getfd\ +#define SYS_pidfd_getfd __NR_pidfd_getfd\ +#else\ +#define SYS_pidfd_getfd 438\ +#endif\ +#endif\ +#endif' csrc/kernels/backend/symmetric.hpp + fi + if [ "$MODE" = "install" ]; then echo "Installing $name into environment" eval "$extra_env" uv pip install --no-build-isolation -vvv . From 70009fb9344d2a7ba642e68369e0a64a6252e8bc Mon Sep 17 00:00:00 2001 From: anthonsu <50185138+anthonsu@users.noreply.github.com> Date: Sat, 25 Jul 2026 13:02:09 -0700 Subject: [PATCH 048/185] [MM][CG] Support ViT CUDA Graph for Gemma-4 (#46837) Signed-off-by: Anthony Su <xsuanthony@gmail.com> Co-authored-by: Linkun Chen <github@lkchen.net> --- docs/design/cuda_graphs_multimodal.md | 1 + .../multimodal/vision_language_offline.py | 41 ++ .../generation/test_vit_cudagraph.py | 13 + vllm/model_executor/models/gemma4_mm.py | 422 +++++++++++++++++- 4 files changed, 476 insertions(+), 1 deletion(-) diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index 7eab425d6e2..7502186ae46 100644 --- a/docs/design/cuda_graphs_multimodal.md +++ b/docs/design/cuda_graphs_multimodal.md @@ -129,6 +129,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra | `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ | | `Gemma3ForConditionalGeneration` | `Gemma3` | ✅︎ | ❌︎ | ❌︎ | | `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ | ❌︎ | +| `Gemma4ForConditionalGeneration` | `Gemma-4` | ✅︎ | ✅︎ | ❌︎ | | `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ | ❌︎ | | `KimiVLForConditionalGeneration` | `Kimi-VL` | ✅︎ | ❌︎ | ❌︎ | | `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | ❌︎ | ❌︎ | diff --git a/examples/generate/multimodal/vision_language_offline.py b/examples/generate/multimodal/vision_language_offline.py index 661401046bd..d8e08835197 100644 --- a/examples/generate/multimodal/vision_language_offline.py +++ b/examples/generate/multimodal/vision_language_offline.py @@ -503,6 +503,45 @@ def run_gemma3n(questions: list[str], modality: str) -> ModelRequestData: ) +# Gemma 4 +def run_gemma4(questions: list[str], modality: str) -> ModelRequestData: + assert modality in ("image", "video") + model_name = "google/gemma-4-31B-it" + + # NOTE: Gemma-4-31B is a large model. Users running into Out-Of-Memory (OOM) + # errors might need to set `tensor_parallel_size` to > 1. + engine_args = EngineArgs( + model=model_name, + max_model_len=4096, + max_num_seqs=2, + limit_mm_per_prompt={modality: 1}, + ) + + if modality == "image": + prompts = [ + ( + "<bos><start_of_turn>user\n" + f"<|image|>\n{question}<end_of_turn>\n" + "<start_of_turn>model\n" + ) + for question in questions + ] + else: # video + prompts = [ + ( + "<bos><start_of_turn>user\n" + f"<|video|>\n{question}<end_of_turn>\n" + "<start_of_turn>model\n" + ) + for question in questions + ] + + return ModelRequestData( + engine_args=engine_args, + prompts=prompts, + ) + + # GLM-4v def run_glm4v(questions: list[str], modality: str) -> ModelRequestData: assert modality == "image" @@ -2303,6 +2342,7 @@ model_example_map = { "exaone4_5": run_exaone4_5, "gemma3": run_gemma3, "gemma3n": run_gemma3n, + "gemma4": run_gemma4, "glm4v": run_glm4v, "glm4_1v": run_glm4_1v, "glm4_5v": run_glm4_5v, @@ -2374,6 +2414,7 @@ MODELS_NEED_VIDEO_METADATA = [ MODELS_SUPPORT_VIT_CUDA_GRAPH = [ "llama4", + "gemma4", "qwen2_vl", "qwen2_5_vl", "qwen3_vl", diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index 387046f3e74..13fbe26880d 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/tests/models/multimodal/generation/test_vit_cudagraph.py @@ -249,6 +249,19 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = { }, skip=True, # TODO: Re-enable this once OOM issues are resolved on CI. ), + "gemma4": VitCudagraphTestConfig( + model="google/gemma-4-E2B-it", + image_prompt=( + "<bos><start_of_turn>user\n<|image|>\nWhat is in this image?<end_of_turn>\n" + "<start_of_turn>model\n" + ), + video_prompt=( + "<bos><start_of_turn>user\n<|video|>\nDescribe this video in one sentence." + "<end_of_turn>\n<start_of_turn>model\n" + ), + needs_video_metadata=True, + marks=[pytest.mark.core_model], + ), } diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index b0169cbc6dd..e158a785016 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -16,7 +16,7 @@ reason about temporal order. import math from collections.abc import Iterable, Mapping, Sequence -from typing import TYPE_CHECKING, Annotated, Any, Literal +from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal import numpy as np import torch @@ -72,6 +72,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape from .interfaces import ( MultiModalEmbeddings, SupportsEagle3, + SupportsEncoderCudaGraph, SupportsLoRA, SupportsMultiModal, SupportsPP, @@ -86,6 +87,12 @@ from .utils import ( if TYPE_CHECKING: from vllm.model_executor.layers.quantization import QuantizationConfig + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphCaptureInputs, + EncoderCudaGraphConfig, + EncoderCudaGraphReplayBuffers, + EncoderItemSpec, + ) logger = init_logger(__name__) @@ -980,7 +987,9 @@ class Gemma4ForConditionalGeneration( SupportsPP, SupportsLoRA, SupportsEagle3, + SupportsEncoderCudaGraph, ): + supports_encoder_cudagraph: ClassVar[Literal[True]] = True # Gemma4 clamps mm_prefix bidirectional ranges to the sliding window # in-kernel (HF's (causal OR blockwise) AND sliding_window). The model # runner reads this to keep image bidirectional ranges that exceed the @@ -1025,6 +1034,7 @@ class Gemma4ForConditionalGeneration( self.quant_config = quant_config self.multimodal_config = multimodal_config self.model_dtype = vllm_config.model_config.dtype + self.vllm_config = vllm_config # Only quantize towers when the quant method supports their # dimensions. BNB/torchao handle arbitrary sizes; other methods @@ -1525,6 +1535,416 @@ class Gemma4ForConditionalGeneration( return multimodal_embeddings + # ------------------------------------------------------------------ # + # EncoderCudaGraph protocol methods + # ------------------------------------------------------------------ # + + def get_encoder_cudagraph_config(self) -> "EncoderCudaGraphConfig": + from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig + + def pad_pixel_values(dst: torch.Tensor, src: torch.Tensor) -> None: + dst.zero_() + batch_size, num_patches = src.shape[0], src.shape[1] + dst[:batch_size, :num_patches].copy_(src) + + def pad_pixel_position_ids(dst: torch.Tensor, src: torch.Tensor) -> None: + dst.fill_(-1) + batch_size, num_patches = src.shape[0], src.shape[1] + dst[:batch_size, :num_patches].copy_(src) + + return EncoderCudaGraphConfig( + modalities=["image", "video"], + buffer_keys=[ + "pixel_values", + "pixel_position_ids", + "gather_indices", + ], + out_hidden_size=self.config.text_config.hidden_size, + max_frames_per_video=_VIDEO_MAX_FRAMES, + padding_logics={ + "pixel_values": pad_pixel_values, + "pixel_position_ids": pad_pixel_position_ids, + }, + ) + + def get_encoder_cudagraph_budget_range( + self, + vllm_config: VllmConfig, + ) -> tuple[int, int]: + min_budget = _SUPPORTED_SOFT_TOKENS[0] + max_budget = min( + vllm_config.scheduler_config.max_num_batched_tokens, + vllm_config.model_config.max_model_len, + ) + return (min_budget, max_budget) + + def get_input_modality(self, mm_kwargs: dict[str, Any]) -> str: + if "pixel_values" in mm_kwargs: + return "image" + elif "pixel_values_videos" in mm_kwargs: + return "video" + raise ValueError("Unsupported modality in mm_kwargs") + + def get_max_frames_per_video(self) -> int: + return _VIDEO_MAX_FRAMES + + def get_encoder_cudagraph_item_specs( + self, + mm_kwargs: dict[str, Any], + ) -> list["EncoderItemSpec"]: + from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec + + vision_cfg = self.vision_tower.config + pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2 + + modality = self.get_input_modality(mm_kwargs) + if modality == "image": + pixel_values = mm_kwargs["pixel_values"] + if isinstance(pixel_values, list): + return [ + EncoderItemSpec( + input_size=pv.shape[0], + output_tokens=pv.shape[0] // pool_ratio, + ) + for pv in pixel_values + ] + else: + return [ + EncoderItemSpec( + input_size=pixel_values.shape[1], + output_tokens=pixel_values.shape[1] // pool_ratio, + ) + for _ in range(pixel_values.shape[0]) + ] + elif modality == "video": + pixel_values_videos = mm_kwargs["pixel_values_videos"] + video_frame_counts = mm_kwargs["video_frame_counts"] + fc_list = ( + video_frame_counts.tolist() + if isinstance(video_frame_counts, torch.Tensor) + else list(video_frame_counts) + ) + np_patches = pixel_values_videos.shape[1] + return [ + EncoderItemSpec( + input_size=fc * np_patches, + output_tokens=fc * (np_patches // pool_ratio), + ) + for fc in fc_list + ] + raise ValueError(f"Unknown modality: {modality}") + + def select_encoder_cudagraph_items( + self, + mm_kwargs: dict[str, Any], + indices: list[int], + ) -> dict[str, Any]: + modality = self.get_input_modality(mm_kwargs) + if modality == "image": + pixel_values = mm_kwargs["pixel_values"] + pixel_position_ids = mm_kwargs["pixel_position_ids"] + if len(indices) == 0: + is_pv_list = isinstance(pixel_values, list) + is_pp_list = isinstance(pixel_position_ids, list) + return { + "pixel_values": ([] if is_pv_list else pixel_values[:0]), + "pixel_position_ids": ( + [] if is_pp_list else pixel_position_ids[:0] + ), + } + if isinstance(pixel_values, list): + return { + "pixel_values": [pixel_values[i] for i in indices], + "pixel_position_ids": [pixel_position_ids[i] for i in indices], + } + return { + "pixel_values": pixel_values[indices], + "pixel_position_ids": pixel_position_ids[indices], + } + elif modality == "video": + pixel_values_videos = mm_kwargs["pixel_values_videos"] + pixel_position_ids_videos = mm_kwargs["pixel_position_ids_videos"] + video_frame_counts = mm_kwargs["video_frame_counts"] + + if len(indices) == 0: + is_fc_tensor = isinstance(video_frame_counts, torch.Tensor) + return { + "pixel_values_videos": pixel_values_videos[:0], + "pixel_position_ids_videos": pixel_position_ids_videos[:0], + "video_frame_counts": ( + video_frame_counts[:0] if is_fc_tensor else [] + ), + } + + fc_list = ( + video_frame_counts.tolist() + if isinstance(video_frame_counts, torch.Tensor) + else list(video_frame_counts) + ) + cum_frames = [0] + for fc in fc_list: + cum_frames.append(cum_frames[-1] + fc) + + selected_pv = torch.cat( + [ + pixel_values_videos[cum_frames[i] : cum_frames[i + 1]] + for i in indices + ], + dim=0, + ) + selected_pp = torch.cat( + [ + pixel_position_ids_videos[cum_frames[i] : cum_frames[i + 1]] + for i in indices + ], + dim=0, + ) + selected_fc = ( + video_frame_counts[indices] + if isinstance(video_frame_counts, torch.Tensor) + else [video_frame_counts[i] for i in indices] + ) + return { + "pixel_values_videos": selected_pv, + "pixel_position_ids_videos": selected_pp, + "video_frame_counts": selected_fc, + } + raise ValueError(f"Unknown modality: {modality}") + + def prepare_encoder_cudagraph_capture_inputs( + self, + token_budget: int = 256, + max_batch_size: int = 4, + max_frames_per_batch: int = 1, + device: torch.device | str = "cpu", + dtype: torch.dtype | None = None, + path: str = "default", + **kwargs: Any, + ) -> "EncoderCudaGraphCaptureInputs": + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphCaptureInputs, + ) + + dtype = dtype or torch.float32 + + max_size = max(max_batch_size, max_frames_per_batch) + + vision_cfg = self.vision_tower.config + pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2 + + # Retrieve the model's actual configured maximum tokens: + configured_max_tokens = getattr( + self.config.vision_config, + "num_soft_tokens", + _SUPPORTED_SOFT_TOKENS[2], + ) + # Dynamically compute the slot capacity per item bounded by both the + # current graph budget and the user's maximum config: + per_item_output = min(token_budget, configured_max_tokens) + # Satisfy k^2 * per_item_output = per_item_patches + per_item_patches = per_item_output * pool_ratio + + patch_size = self.vision_tower.config.patch_size + num_channels = getattr(self.vision_tower.config, "num_channels", 3) + patch_pixels = (patch_size**2) * num_channels + + dummy_pixel_values = torch.zeros( + (max_size, per_item_patches, patch_pixels), + device=device, + dtype=dtype, + ) + dummy_pixel_position_ids = torch.full( + (max_size, per_item_patches, 2), + -1, + device=device, + dtype=torch.long, + ) + dummy_gather_indices = torch.zeros( + (token_budget,), + device=device, + dtype=torch.long, + ) + + return EncoderCudaGraphCaptureInputs( + values={ + "pixel_values": dummy_pixel_values, + "pixel_position_ids": dummy_pixel_position_ids, + "gather_indices": dummy_gather_indices, + } + ) + + def prepare_encoder_cudagraph_replay_buffers( + self, + mm_kwargs: dict[str, Any], + max_batch_size: int = 4, + max_frames_per_batch: int = 1, + path: str = "default", + **kwargs: Any, + ) -> "EncoderCudaGraphReplayBuffers": + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphReplayBuffers, + ) + + modality = self.get_input_modality(mm_kwargs) + if modality == "image": + pixel_values = mm_kwargs["pixel_values"] + pixel_position_ids = mm_kwargs["pixel_position_ids"] + elif modality == "video": + pixel_values = mm_kwargs["pixel_values_videos"] + pixel_position_ids = mm_kwargs["pixel_position_ids_videos"] + else: + raise ValueError(f"Unsupported modality: {modality}") + + if isinstance(pixel_values, list): + max_patches = max(pv.shape[0] for pv in pixel_values) + batch_size = len(pixel_values) + pv_tensor = torch.zeros( + (batch_size, max_patches, pixel_values[0].shape[1]), + dtype=pixel_values[0].dtype, + device=pixel_values[0].device, + ) + pp_tensor = torch.full( + (batch_size, max_patches, 2), + -1, + dtype=pixel_position_ids[0].dtype, + device=pixel_position_ids[0].device, + ) + for i, (pv, pp) in enumerate(zip(pixel_values, pixel_position_ids)): + pv_tensor[i, : pv.shape[0]].copy_(pv) + pp_tensor[i, : pp.shape[0]].copy_(pp) + pixel_values = pv_tensor + pixel_position_ids = pp_tensor + + item_specs = self.get_encoder_cudagraph_item_specs(mm_kwargs) + per_item_out_tokens = [spec.output_tokens for spec in item_specs] + total_tokens = sum(per_item_out_tokens) + + device = pixel_values.device + vision_cfg = self.vision_tower.config + pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2 + per_item_output = pixel_values.shape[1] // pool_ratio + + # ONLY allocate an array of exact size `total_tokens`. + # DO NOT pad it. The upstream Graph Manager handles the padding securely. + gather_indices = torch.zeros((total_tokens,), dtype=torch.long, device=device) + + if modality == "image": + dst_offset = 0 + for i, n_tok in enumerate(per_item_out_tokens): + safe_n_tok = min(n_tok, per_item_output) + src_start = i * per_item_output + src_end = src_start + safe_n_tok + gather_indices[dst_offset : dst_offset + safe_n_tok] = torch.arange( + src_start, src_end, dtype=torch.long, device=device + ) + dst_offset += safe_n_tok + elif modality == "video": + video_frame_counts = mm_kwargs["video_frame_counts"] + fc_list = ( + video_frame_counts.tolist() + if isinstance(video_frame_counts, torch.Tensor) + else list(video_frame_counts) + ) + vision_cfg = self.vision_tower.config + pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2 + np_patches = pixel_values.shape[1] + frame_output_tokens = np_patches // pool_ratio + safe_frame_output_tokens = min(frame_output_tokens, per_item_output) + + dst_offset = 0 + frame_idx = 0 + for fc in fc_list: + for f in range(fc): + src_start = (frame_idx + f) * per_item_output + src_end = src_start + safe_frame_output_tokens + gather_indices[ + dst_offset : dst_offset + safe_frame_output_tokens + ] = torch.arange( + src_start, src_end, dtype=torch.long, device=device + ) + dst_offset += safe_frame_output_tokens + frame_idx += fc + + return EncoderCudaGraphReplayBuffers( + values={ + "pixel_values": pixel_values, + "pixel_position_ids": pixel_position_ids, + "gather_indices": gather_indices, + } + ) + + def encoder_cudagraph_forward( + self, + inputs: dict[str, torch.Tensor], + path: str = "default", + **kwargs: Any, + ) -> torch.Tensor: + pixel_values = inputs["pixel_values"] + pixel_position_ids = inputs["pixel_position_ids"] + gather_indices = inputs["gather_indices"] + + pad_tensor = (pixel_position_ids == -1).all(dim=-1) + + vt = self.vision_tower + inputs_embeds = vt.patch_embedder( + pixel_values, + pixel_position_ids, + pad_tensor, + ).to(self.model_dtype) + + encoder_outputs = vt.encoder( + inputs_embeds=inputs_embeds, + attention_mask=~pad_tensor, + pixel_position_ids=pixel_position_ids, + ) + hidden_states = encoder_outputs.last_hidden_state + + pool_ratio = getattr(vt.config, "pooling_kernel_size", 2) ** 2 + per_item_output = pixel_values.shape[1] // pool_ratio + + pooled_states, _ = vt.pooler( + hidden_states=hidden_states, + pixel_position_ids=pixel_position_ids, + padding_positions=pad_tensor, + output_length=per_item_output, + ) + + if getattr(vt.config, "standardize", False): + pooled_states = (pooled_states - vt.std_bias) * vt.std_scale + + flat_pooled = pooled_states.reshape(-1, pooled_states.shape[-1]) + gathered_states = flat_pooled[gather_indices] + + # Cast to the projection layer's dtype to resolve mixed-precision crash + target_dtype = self.embed_vision.embedding_projection.weight.dtype + gathered_states = gathered_states.to(target_dtype) + + flat_proj_embs = self.embed_vision( + inputs_embeds=gathered_states.unsqueeze(0) + ).squeeze(0) + + return flat_proj_embs + + def encoder_eager_forward( + self, + mm_kwargs: dict[str, Any], + path: str = "default", + **kwargs: Any, + ) -> torch.Tensor: + modality = self.get_input_modality(mm_kwargs) + if modality == "image": + image_input = self._parse_and_validate_image_input(**mm_kwargs) + assert image_input is not None + embeddings = self._process_image_input(image_input) + elif modality == "video": + video_input = self._parse_and_validate_video_input(**mm_kwargs) + assert video_input is not None + embeddings = self._process_video_input(video_input) + else: + raise ValueError(f"Unsupported modality: {modality}") + + return torch.cat(embeddings, dim=0) + def embed_input_ids( self, input_ids: torch.Tensor, From dbd80cc031718cae22ec8999bc5a494bacc4e955 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim <taneem.ibrahim@gmail.com> Date: Sat, 25 Jul 2026 16:53:57 -0400 Subject: [PATCH 049/185] [UX] DCP Topology Validation (#49777) Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> --- vllm/config/model.py | 40 +++++++++++++++++++++++----------------- 1 file changed, 23 insertions(+), 17 deletions(-) diff --git a/vllm/config/model.py b/vllm/config/model.py index 6b7825be08f..20efeb03298 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1298,27 +1298,33 @@ class ModelConfig: decode_context_parallel_size = parallel_config.decode_context_parallel_size if decode_context_parallel_size > 1 and not self.use_mla: total_num_kv_heads = self.get_total_num_kv_heads() - assert tensor_parallel_size > total_num_kv_heads, ( - f"tensor parallel size {tensor_parallel_size} must be greater " - f"than total num kv heads {total_num_kv_heads} when enable " - f"decode context parallel for GQA/MQA" - ) + if tensor_parallel_size <= total_num_kv_heads: + raise ValueError( + "Decode context parallelism for GQA/MQA requires " + f"`--tensor-parallel-size` ({tensor_parallel_size}) to be " + "greater than the model's total number of KV heads " + f"({total_num_kv_heads}). Increase `--tensor-parallel-size` " + "or set `--decode-context-parallel-size 1`." + ) max_dcp_size = tensor_parallel_size // total_num_kv_heads - assert decode_context_parallel_size <= max_dcp_size, ( - f"decode context parallel size must less than or equal to " - f"(tensor parallel size {tensor_parallel_size} // total " - f"num kv heads {total_num_kv_heads}) = {max_dcp_size}, " - f"but got {decode_context_parallel_size}" - ) + if decode_context_parallel_size > max_dcp_size: + raise ValueError( + "`--decode-context-parallel-size` " + f"({decode_context_parallel_size}) exceeds the maximum " + f"supported value ({max_dcp_size}) for " + f"`--tensor-parallel-size` ({tensor_parallel_size}) and " + f"{total_num_kv_heads} model KV heads." + ) num_q_per_kv = total_num_attention_heads // total_num_kv_heads - assert num_q_per_kv % decode_context_parallel_size == 0, ( - f"Total number of q per kv attn heads ({num_q_per_kv})" - " must be divisible by dcp world size when enable " - "decode context parallel for GQA " - f"({parallel_config.decode_context_parallel_size})." - ) + if num_q_per_kv % decode_context_parallel_size != 0: + raise ValueError( + "The model's number of query heads per KV head " + f"({num_q_per_kv}) must be divisible by " + "`--decode-context-parallel-size` " + f"({decode_context_parallel_size}) for GQA/MQA." + ) # torch_shm uses a single IPC queue to rank 0; DP>1 is # incompatible because API servers can't know which From d30b1ecd1bdf7c3d92f3b444c4538efd8fbb40ac Mon Sep 17 00:00:00 2001 From: Rui Yin <2260891073@qq.com> Date: Sun, 26 Jul 2026 04:58:09 +0800 Subject: [PATCH 050/185] [Bugfix][KV Offloading] Defer request finalization until final store (#49671) Signed-off-by: Rui Yin <2260891073@qq.com> Co-authored-by: Or Ozeri <oro@il.ibm.com> --- .../offloading_connector/test_scheduler.py | 115 +++++++++++------- .../unit/offloading_connector/utils.py | 26 +++- .../kv_connector/v1/offloading/scheduler.py | 28 +++-- 3 files changed, 111 insertions(+), 58 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 29d206b426b..c65d02dda4b 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -68,7 +68,9 @@ def test_scheduler_reports_allocation_failure(request_runner): runner.run(decoded_tokens=[EOS_TOKEN_ID]) reduced = _reduce_kv_connector_stats(runner) - assert reduced[_ConnectorMetricName.ALLOCATION_FAILURE] == 1 + # Two attempts: once while running (block becomes full during prefill), + # once from finished_req_ids on the next step. + assert reduced[_ConnectorMetricName.ALLOCATION_FAILURE] == 2 @pytest.mark.parametrize("async_scheduling", [True, False]) @@ -76,7 +78,7 @@ def test_scheduler_reports_allocation_failure(request_runner): def test_last_block_offloaded_at_request_finish( request_runner, async_scheduling: bool, prompt_offset: int ): - """EOS fills the last block at request finish — verify req_status is kept alive. + """EOS fills the last block at request finish - verify the final block is stored. prompt = block_size + prompt_offset tokens → not a full block at schedule time, so _build_store_jobs creates no store job. After EOS, request_finished @@ -98,18 +100,16 @@ def test_last_block_offloaded_at_request_finish( generate_store_output(list(keys)) ) - # Run with one step (EOS) - runner.run( - decoded_tokens=[EOS_TOKEN_ID], - ) + if prompt_offset == -1: + # EOS fills the block, so a store job is created for block 0. + runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(0,)) + else: + # Block remains partial, so no store job is created. + runner.run(decoded_tokens=[EOS_TOKEN_ID]) cs = runner.connector_scheduler - # Verify req_status is kept alive for _build_store_jobs to process - # regardless of whether there are storable blocks - assert "0" in cs._req_status, ( - "req_status was deleted but should be kept alive " - "for _build_store_jobs to process finished_req_ids." - ) + # After the full run completes, req_status is cleaned up. + assert "0" not in cs._req_status @pytest.mark.parametrize("async_scheduling", [True, False]) @@ -569,14 +569,14 @@ def test_request_preemption(request_runner, async_scheduling: bool): @pytest.mark.parametrize("async_scheduling", [True, False]) -def test_on_request_finished_is_not_deferred_until_store_completion( +def test_on_request_finished_not_deferred_until_store_completion( request_runner, async_scheduling: bool ): - """on_request_finished fires when no more stores will be submitted. + """on_request_finished fires after the last prepare_store is submitted. - A request can finish while its GPU->primary store is still in flight. The - manager-level hook should not wait for that completion; complete_store may - still arrive afterward for already-submitted transfer jobs. + The manager contract guarantees no more submit-side calls (prepare_store) + after on_request_finished. However, complete_store callbacks for + already-submitted transfers may still arrive afterward. """ block_size = 4 blocks_per_chunk = 3 @@ -613,8 +613,9 @@ def test_on_request_finished_is_not_deferred_until_store_completion( complete_transfers=False, ) - # Finish the request while its stores are still in flight. The hook should - # fire immediately even though no complete_store has arrived yet. + # Finish the request while its stores are still in flight. The hook fires + # once the last prepare_store is issued (on the next schedule step), even + # though complete_store has not yet been called. runner.run( decoded_tokens=[EOS_TOKEN_ID], complete_transfers=False, @@ -624,8 +625,7 @@ def test_on_request_finished_is_not_deferred_until_store_completion( assert calls == [("on_request_finished", req_id)], calls - # Drain the stores afterward. The already-submitted complete_store calls - # are allowed to arrive after on_request_finished. + # Drain the stores afterward. complete_store is allowed after the hook. runner.run( decoded_tokens=[], complete_transfers=True, @@ -638,11 +638,50 @@ def test_on_request_finished_is_not_deferred_until_store_completion( finished_idx = calls.index(("on_request_finished", req_id)) store_indices = [i for i, c in enumerate(calls) if c == ("complete_store", req_id)] - # The request-level hook no longer waits for already-submitted transfers. + # complete_store arrives after on_request_finished, as allowed by the contract. assert store_indices, calls assert finished_idx < min(store_indices), calls +@pytest.mark.parametrize("async_scheduling", [True, False]) +def test_on_request_finished_fires_after_final_block_store( + request_runner, async_scheduling: bool +): + """on_request_finished fires after the final-block prepare_store at EOS. + + When EOS fills a partial block, request_finished() keeps req_status alive + so _build_store_jobs can create a store job for it on the next step. + """ + block_size = 4 + runner = request_runner( + block_size=block_size, + num_gpu_blocks=10, + async_scheduling=async_scheduling, + ) + + calls: list[tuple[str, str]] = [] + runner.manager.on_request_finished.side_effect = lambda req_context: calls.append( + ("on_request_finished", req_context.req_id) + ) + + def prepare_store(keys, req_context): + calls.append(("prepare_store", req_context.req_id)) + return generate_store_output(keys) + + runner.manager.prepare_store.side_effect = prepare_store + + runner.new_request(token_ids=[0] * (block_size - 1)) + runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(0,)) + + req_id = str(runner.req_id) + assert calls.count(("on_request_finished", req_id)) == 1, calls + + finished_idx = calls.index(("on_request_finished", req_id)) + prepare_indices = [i for i, c in enumerate(calls) if c == ("prepare_store", req_id)] + assert prepare_indices, calls + assert finished_idx > max(prepare_indices), calls + + @pytest.mark.parametrize("async_scheduling", [True, False]) def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: bool): block_size = 4 @@ -831,7 +870,10 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo touch_calls = runner.manager.touch.call_args_list assert len(touch_calls) == 6 - runner.run(decoded_tokens=[EOS_TOKEN_ID]) + # EOS fills the 7th block (offset 6). The extra schedule step processes + # finished_req_ids and stores block 6 for both groups before the request's + # GPU blocks are freed. + runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(6,)) runner.scheduler.reset_prefix_cache() @@ -844,19 +886,11 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo # Group 1 (sliding window, window=2): only the last 2 blocks # are within the window → loads blocks 1,2 expected_loaded=((0, 0), (0, 1), (0, 2), (1, 1), (1, 2)), - # The deferred store from the previous request's last block - # completes during this step, and its blocks are flushed because - # they were reallocated to the new request. - # Only block 1 (sliding window group) is stored — block 0's - # deferred store is flushed because it was reallocated. - expected_stored=((0, 1),), - expected_flushed=((0, 1),), ) - # 4 touch calls: 2 from get_num_new_matched_tokens (2 groups) - # + 2 from _get_reqs_to_store (2 groups) + # 2 touch calls from get_num_new_matched_tokens (2 groups) touch_calls = runner.manager.touch.call_args_list - assert len(touch_calls) == 4 + 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 @@ -1733,8 +1767,8 @@ def test_reset_cache(request_runner, async_scheduling: bool): def test_reset_cache_finalizes_finished_request_with_pending_store( request_runner, async_scheduling: bool ): - """reset_cache drops a finished request whose in-flight stores it discards - without calling on_request_finished twice. + """reset_cache fires on_request_finished for a finished request whose + in-flight stores it discards, exactly once. """ block_size = 4 blocks_per_chunk = 3 @@ -1770,16 +1804,15 @@ def test_reset_cache_finalizes_finished_request_with_pending_store( assert req_status.transfer_jobs, "expected an in-flight store before finish" assert any(job.is_store for job in cs._jobs.values()) - # Finish the request while its store is still in flight. request_finished - # fires the hook eagerly, but the entry stays tracked so later completions - # can still call complete_store(). + # Finish the request while its store is still in flight. The manager hook + # is deferred because the final store decision has not happened yet. req_status.req.status = RequestStatus.FINISHED_STOPPED cs.request_finished(req_status.req) - assert finalized == [req_id] + assert finalized == [] assert req_id in cs._req_status - # reset_cache discards the in-flight store and drops the state without a - # duplicate on_request_finished call. + # reset_cache discards both the in-flight and not-yet-prepared final stores, + # so it issues the deferred notification before dropping the state. cs.reset_cache() assert finalized == [req_id] assert req_id not in cs._req_status diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index b878e294a6e..c1979a0b1f7 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -482,8 +482,13 @@ class RequestRunner: # Strict-always-False frees the request immediately on EOS, but # the worker may still have a deferred store queued. In production # the next request's step drains it; in single-request tests we - # must keep stepping until the scheduler sees no in-flight jobs. - if not self.scheduler.requests and not self.connector_scheduler._jobs: + # must keep stepping until the scheduler sees no in-flight jobs + # and no pending finished_req_ids awaiting build_connector_meta. + if ( + not self.scheduler.requests + and not self.connector_scheduler._jobs + and not self.scheduler.finished_req_ids + ): break scheduler_output = self.scheduler.schedule() @@ -544,19 +549,30 @@ class RequestRunner: if ( prev_token_id == EOS_TOKEN_ID and prev_token_id != token_id - and (self.scheduler.requests or self.connector_scheduler._jobs) + and ( + self.scheduler.requests + or self.connector_scheduler._jobs + or self.scheduler.finished_req_ids + ) ): # continue for one more step to allow offloading to kick off continue if token_id is None: if self.async_scheduling: - # sample last token + # Flush the previous step's output. engine_outputs = self.scheduler.update_from_output( prev_scheduler_output, prev_model_runner_output ) self._record_kv_connector_stats(engine_outputs) - break + prev_model_runner_output = None + if self.scheduler.requests: + # Request still running, just exhausted decoded_tokens. + break + if not self.scheduler.finished_req_ids and ( + not complete_transfers or not self.connector_scheduler._jobs + ): + break self._parse_transfers() 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 2e1875ae73c..b53a365abd4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -278,6 +278,8 @@ class RequestOffloadState: # time.monotonic() of this request's first deferred offload lookup; # None once consumed (observed) or while no lookup is pending. deferred_lookup_start_time: float | None = None + # True once on_request_finished has been signaled to the manager. + finished_signaled: bool = False def __post_init__(self) -> None: self.group_states = tuple( @@ -481,13 +483,6 @@ class OffloadingConnectorScheduler: num = min(num, req_status.req.num_prompt_tokens) return num - def _maybe_cleanup_finished_req( - self, req_id: str, req_status: RequestOffloadState - ) -> None: - """Clean up req_status if finished and no in-flight jobs.""" - if req_status.req.is_finished() and not req_status.transfer_jobs: - del self._req_status[req_id] - def _maximal_prefix_lookup( self, keys: Iterable[OffloadKey], @@ -1039,7 +1034,6 @@ class OffloadingConnectorScheduler: if not new_offload_keys: req_status.advance_stored_idx(num_offloadable_tokens) - self._maybe_cleanup_finished_req(req_id, req_status) continue store_output = self.manager.prepare_store( @@ -1050,12 +1044,10 @@ class OffloadingConnectorScheduler: _ConnectorMetricName.ALLOCATION_FAILURE ) logger.warning("Request %s: cannot store chunks", req_id) - self._maybe_cleanup_finished_req(req_id, req_status) continue if not store_output.keys_to_store: req_status.advance_stored_idx(num_offloadable_tokens) - self._maybe_cleanup_finished_req(req_id, req_status) continue self._touch(req_status) @@ -1199,6 +1191,17 @@ class OffloadingConnectorScheduler: store_jobs=self._build_store_jobs(scheduler_output), jobs_to_flush=self._current_batch_jobs_to_flush, ) + + # All prepare_store calls for finished requests have been issued. + # Signal on_request_finished and clean up state where possible. + for req_id in scheduler_output.finished_req_ids or (): + req_status = self._req_status.get(req_id) + if req_status is None: + continue + req_status.finished_signaled = True + self.manager.on_request_finished(req_status.req_context) + if not req_status.transfer_jobs: + del self._req_status[req_id] self._current_batch_load_jobs = {} self._current_batch_jobs_to_flush = set() self._current_batch_allocated_block_ids = set() @@ -1289,7 +1292,7 @@ class OffloadingConnectorScheduler: del self._jobs[job_id] req_status.transfer_jobs.remove(job_id) - if not req_status.transfer_jobs and req_status.req.is_finished(): + if req_status.finished_signaled and not req_status.transfer_jobs: del self._req_status[job_status.req_id] def get_stats(self) -> OffloadingConnectorStats | None: @@ -1331,7 +1334,6 @@ class OffloadingConnectorScheduler: self.manager.on_request_finished(req_context) return False, None - self.manager.on_request_finished(req_status.req_context) self._maybe_observe_lookup_async_delay(req_status) # Update offload keys with final block hash so _build_store_jobs can @@ -1375,6 +1377,8 @@ class OffloadingConnectorScheduler: for req_id, status in list(self._req_status.items()): if status.req.is_finished(): + if not status.finished_signaled: + self.manager.on_request_finished(status.req_context) del self._req_status[req_id] # Reset offloading manager cache From 0111002323e226844a18e7f3e5834e9d0e1b6dcd Mon Sep 17 00:00:00 2001 From: Lena Onyshchenko <xonyshch@gmail.com> Date: Sun, 26 Jul 2026 02:50:07 +0200 Subject: [PATCH 051/185] [Kernel] TD operand loads for batched MoE GEMM (moe_mmk) on XPU (#46340) Signed-off-by: oonyshch <xonyshch@gmail.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com> --- tests/kernels/moe/test_batched_moe.py | 184 ++++++++++++++++++ vllm/config/kernel.py | 3 + .../layers/fused_moe/all2all_utils.py | 11 ++ .../fused_moe/experts/fused_batched_moe.py | 82 ++++++-- .../layers/fused_moe/oracle/unquantized.py | 2 + .../fused_moe/prepare_finalize/batched.py | 33 ++-- .../fused_moe/topk_weight_and_reduce.py | 23 ++- vllm/triton_utils/__init__.py | 12 +- vllm/triton_utils/tensor_descriptor.py | 14 ++ 9 files changed, 331 insertions(+), 33 deletions(-) create mode 100644 vllm/triton_utils/tensor_descriptor.py diff --git a/tests/kernels/moe/test_batched_moe.py b/tests/kernels/moe/test_batched_moe.py index b9fe8ceafcd..3c605ced688 100644 --- a/tests/kernels/moe/test_batched_moe.py +++ b/tests/kernels/moe/test_batched_moe.py @@ -351,3 +351,187 @@ def test_fused_moe_batched_experts( torch.testing.assert_close(batched_output, baseline_output, atol=3e-2, rtol=2e-2) torch.testing.assert_close(triton_output, batched_output, atol=2e-2, rtol=2e-2) + + +# USE_TD in moe_mmk + +# K % 8 == 0 (bf16, 16-byte alignment) +_TD_SHAPES = [ + (4, 1, 128, 256), + (8, 64, 512, 512), + (4, 256, 1024, 2048), +] + + +def _td_supported() -> bool: + return current_platform.is_xpu() or ( + current_platform.is_cuda() and current_platform.has_device_capability(90) + ) + + +def _run_td(A, B, num_expert_tokens, use_td: bool): + out = torch.zeros( + A.shape[0], + A.shape[1], + B.shape[1], + dtype=torch.bfloat16, + device=A.device, + ) + import triton + + from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import ( + batched_triton_kernel, + ) + + if use_td: + from vllm.triton_utils.allocation import set_triton_allocator + + set_triton_allocator(A.device) + + E = A.shape[0] + max_tokens = A.shape[1] + K = A.shape[2] + N = B.shape[1] + BM = BN = BK = 64 + grid = (E, triton.cdiv(max_tokens, BM) * triton.cdiv(N, BN)) + batched_triton_kernel[grid]( + A, + B, + out, + num_expert_tokens, + tl.bfloat16, + max_tokens, + K, + N, + None, + None, + None, + A.stride(0), + A.stride(1), + A.stride(2), + B.stride(0), + B.stride(2), + B.stride(1), + out.stride(0), + out.stride(1), + out.stride(2), + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + False, + False, + False, + BLOCK_M=BM, + BLOCK_N=BN, + BLOCK_K=BK, + USE_TD=use_td, + ) + return out + + +@pytest.mark.parametrize("num_experts,max_tokens_per_expert,K,N", _TD_SHAPES) +def test_batched_mm_td_matches_plain(num_experts, max_tokens_per_expert, K, N): + if not _td_supported(): + pytest.skip("TD requires XPU or CUDA sm_90+") + set_random_seed(42) + device = current_platform.device_type + A = ( + torch.randn( + num_experts, max_tokens_per_expert, K, device=device, dtype=torch.bfloat16 + ) + / 10 + ) + B = torch.randn(num_experts, N, K, device=device, dtype=torch.bfloat16) + num_expert_tokens = torch.randint( + 1, + max_tokens_per_expert + 1, + size=(num_experts,), + device=device, + dtype=torch.int32, + ) + + out_plain = _run_td(A, B, num_expert_tokens, False) + out_td = _run_td(A, B, num_expert_tokens, True) + + torch.testing.assert_close(out_td, out_plain, atol=6e-2, rtol=6e-2) + + +def test_batched_mm_td_zero_expert_tokens(): + if not _td_supported(): + pytest.skip("TD requires XPU or CUDA sm_90+") + set_random_seed(42) + device = current_platform.device_type + E, M, K, N = 8, 32, 256, 256 + A = torch.randn(E, M, K, device=device, dtype=torch.bfloat16) / 10 + B = torch.randn(E, N, K, device=device, dtype=torch.bfloat16) + num_expert_tokens = torch.zeros(E, device=device, dtype=torch.int32) + num_expert_tokens[::2] = M + + out_plain = _run_td(A, B, num_expert_tokens, False) + out_td = _run_td(A, B, num_expert_tokens, True) + + for e in range(E): + if num_expert_tokens[e].item() == 0: + assert out_plain[e].abs().max().item() == 0.0 + assert out_td[e].abs().max().item() == 0.0 + + torch.testing.assert_close(out_td, out_plain, atol=6e-2, rtol=6e-2) + + +# BatchedTritonExperts device enablement (XPU) +def test_batched_triton_experts_supports_current_device(): + from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import ( + BatchedTritonExperts, + ) + + if current_platform.is_xpu() or current_platform.is_cuda_alike(): + assert BatchedTritonExperts._supports_current_device() + else: + pytest.skip("No GPU device available") + + +@pytest.mark.parametrize("m,n,k,e,topk", [(32, 512, 512, 8, 2), (45, 1024, 128, 8, 1)]) +def test_batched_experts_end_to_end(m, n, k, e, topk): + """End-to-end BatchedTritonExperts via the reference (no-comms) + BatchedPrepareAndFinalize, validated against a torch reference. Exercises + the device-enablement path.""" + if not (current_platform.is_xpu() or current_platform.is_cuda_alike()): + pytest.skip("No GPU device available") + + from vllm.v1.worker.workspace import init_workspace_manager + + set_random_seed(7) + device = current_platform.device_type + init_workspace_manager(torch.device(f"{device}:0")) + + a = torch.randn((m, k), device=device, dtype=torch.bfloat16) / 10 + score = torch.randn((m, e), device=device, dtype=torch.bfloat16) + # w1 is the fused gate+up projection [E, 2N, K]; w2 is [E, K, N]. + w1 = torch.randn((e, 2 * n, k), device=device, dtype=torch.bfloat16) / 15 + w2 = torch.randn((e, k, n), device=device, dtype=torch.bfloat16) / 15 + + with set_current_vllm_config(vllm_config): + topk_weight, topk_ids, _ = fused_topk(a, score, topk, False) + + baseline_output = torch_experts(a, w1, w2, topk_weight, topk_ids) + triton_output = batched_moe(a, w1, w2, topk_weight, topk_ids) + + torch.testing.assert_close(triton_output, baseline_output, atol=3e-2, rtol=2e-2) + + +def test_batched_triton_backend_mapping(): + from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( + UnquantizedMoeBackend, + map_unquantized_backend, + ) + + assert ( + map_unquantized_backend("batched_triton") + == UnquantizedMoeBackend.BATCHED_TRITON + ) + assert map_unquantized_backend("triton") == UnquantizedMoeBackend.TRITON diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 16856dc3581..59195f307b8 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -122,6 +122,7 @@ class IrOpPriorityConfig: MoEBackend = Literal[ "auto", "triton", + "batched_triton", "deep_gemm", "deep_gemm_mega_moe", "cutlass", @@ -192,6 +193,8 @@ class KernelConfig: - "auto": Automatically select the best backend based on model and hardware - "triton": Use Triton-based fused MoE kernels + - "batched_triton": Use batched Triton experts (moe_mmk) on the batched + activation format ([E_local, max_num_tokens, K]) - "deep_gemm": Use DeepGEMM kernels (FP8 block-quantized only) - "deep_gemm_mega_moe": Use DeepGEMM mega MoE kernels - "cutlass": Use vLLM CUTLASS kernels diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 6af93bfde90..58c9c8d9f6d 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -19,6 +19,7 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import ( FusedMoEPrepareAndFinalize, ) from vllm.model_executor.layers.fused_moe.prepare_finalize import ( + BatchedPrepareAndFinalize, make_moe_prepare_and_finalize_naive_dp_ep, make_moe_prepare_and_finalize_no_dp_ep, ) @@ -140,6 +141,16 @@ def maybe_make_prepare_finalize( if not allow_new_interface: return None + # Opt-in XPU batched path: reorganize tokens into E x T x K locally + # (no all-to-all) so BatchedTritonExperts (moe_mmk TD) can run. + if current_platform.is_xpu() and moe.moe_backend == "batched_triton": + return BatchedPrepareAndFinalize( + max_num_tokens=moe.max_num_tokens, + num_local_experts=moe.num_local_experts, + num_dispatchers=1, + rank=moe.moe_parallel_config.ep_rank, + ) + # For DP/TP case, fall back to naive P/F. if moe.moe_parallel_config.dp_size > 1: logger.info_once( diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py index 21bda8e173f..609b28a9598 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py @@ -32,7 +32,15 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8StaticTensorSym, ) from vllm.platforms import current_platform -from vllm.triton_utils import tl, triton +from vllm.triton_utils import tl, triton, use_tensor_descriptor +from vllm.triton_utils.allocation import set_triton_allocator + + +def _is_capturing_or_compiling() -> bool: + # torch.cuda.is_current_stream_capturing() is unavailable on non-CUDA (XPU) torch. + return torch.compiler.is_compiling() or ( + current_platform.is_cuda_alike() and torch.cuda.is_current_stream_capturing() + ) @triton.jit @@ -71,9 +79,32 @@ def moe_mmk( use_w8a8: tl.constexpr, use_w8a16: tl.constexpr, per_act_token_quant: tl.constexpr, + # TD: a_base_ptr/b_base_ptr are the expert/CTA-offset bases of A[M,K]/B[N,K]. + a_base_ptr=None, + b_base_ptr=None, + M=0, + N=0, + stride_am: tl.int64 = 0, + stride_bn: tl.int64 = 0, + USE_TD: tl.constexpr = False, ): offs_k = tl.arange(0, BLOCK_K) + if USE_TD: + # make_tensor_descriptor requires the last (K) stride to be a + # compile-time 1; the launcher only enables USE_TD for K-contiguous A/B. + a_desc = tl.make_tensor_descriptor( + a_base_ptr, + shape=[M, K], + strides=[stride_am, 1], + block_shape=[BLOCK_M, BLOCK_K], + ) + b_desc = tl.make_tensor_descriptor( + b_base_ptr, + shape=[N, K], + strides=[stride_bn, 1], + block_shape=[BLOCK_N, BLOCK_K], + ) if use_w8a16: b_scale_ptrs = ( b_scale_ptr + expert_id * stride_bse + offs_n[None, :] * stride_bsn @@ -110,12 +141,17 @@ def moe_mmk( for k in range(0, tl.cdiv(K, BLOCK_K)): # Load the next block of A and B, generate a mask by checking the # K dimension. - a = tl.load( - a_ptrs, - mask=mask_m[:, None] & (offs_k[None, :] < K - k * BLOCK_K), - other=0.0, - ) - b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_K, other=0.0) + if USE_TD: + # B is [N, K]; tile is [BLOCK_N, BLOCK_K], transposed for dot. + a = a_desc.load([0, k * BLOCK_K]) + b = tl.trans(b_desc.load([0, k * BLOCK_K])) + else: + a = tl.load( + a_ptrs, + mask=mask_m[:, None] & (offs_k[None, :] < K - k * BLOCK_K), + other=0.0, + ) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_K, other=0.0) # We accumulate along the K dimension. if use_w8a16: accumulator = tl.dot(a, b.to(compute_type), acc=accumulator) @@ -193,6 +229,7 @@ def expert_triton_kernel( BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + USE_TD: tl.constexpr = False, ): offs_m = tl.arange(0, BLOCK_M) offs_n = tl.arange(0, BLOCK_N) % N @@ -238,6 +275,13 @@ def expert_triton_kernel( use_fp8_w8a8, use_int8_w8a16, per_act_token_quant, + a_ptr, + b_ptr, + M, + N, + stride_am, + stride_bn, + USE_TD, ) # store in C @@ -292,6 +336,7 @@ def batched_triton_kernel( BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + USE_TD: tl.constexpr = False, ): expert_id = tl.program_id(axis=0) e_num_tokens = tl.load(expert_num_tokens + expert_id) @@ -372,6 +417,7 @@ def batched_triton_kernel( BLOCK_M, BLOCK_N, BLOCK_K, + USE_TD, ) @@ -444,6 +490,18 @@ def invoke_moe_batched_triton_kernel( stride_asm = 0 stride_ask = 0 + use_td = ( + use_tensor_descriptor() + and A.stride(2) == 1 + and B.stride(2) == 1 + and (K * A.element_size()) % 16 == 0 + and (BLOCK_M & (BLOCK_M - 1)) == 0 + and (BLOCK_N & (BLOCK_N - 1)) == 0 + and (BLOCK_K & (BLOCK_K - 1)) == 0 + ) + if use_td: + set_triton_allocator(A.device) + batched_triton_kernel[grid]( A, B, @@ -485,6 +543,7 @@ def invoke_moe_batched_triton_kernel( BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, + USE_TD=use_td, ) @@ -616,10 +675,7 @@ class NaiveBatchedExperts(mk.FusedMoEExpertsModular): for expert in range(num_local_experts): # Indexing expert_num_tokens doesn't work w/cudagraphs or inductor - if ( - torch.compiler.is_compiling() - or torch.cuda.is_current_stream_capturing() - ): + if _is_capturing_or_compiling(): num = hidden_states.shape[1] else: num = int(expert_num_tokens[expert].item()) @@ -659,7 +715,7 @@ def batched_moe_kernel_quantize_input( per_act_token_quant: bool, block_shape: list[int] | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: - if torch.compiler.is_compiling() or torch.cuda.is_current_stream_capturing(): + if _is_capturing_or_compiling(): # Note: this does a bunch of extra work because expert_num_tokens is # ignored but it does support torch.compile + cudagraphs. hidden_dim = A.size(-1) @@ -745,7 +801,7 @@ class BatchedTritonExperts(mk.FusedMoEExpertsModular): @staticmethod def _supports_current_device() -> bool: - return current_platform.is_cuda_alike() + return current_platform.is_cuda_alike() or current_platform.is_xpu() @staticmethod def _supports_no_act_and_mul() -> bool: diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 5d4c7336313..8fa1a0c265c 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -151,6 +151,7 @@ def map_unquantized_backend(runner_backend: MoEBackend) -> UnquantizedMoeBackend """Map user's MoEBackend to UnquantizedMoeBackend.""" mapping = { "triton": UnquantizedMoeBackend.TRITON, + "batched_triton": UnquantizedMoeBackend.BATCHED_TRITON, "flashinfer_trtllm": UnquantizedMoeBackend.FLASHINFER_TRTLLM, "flashinfer_cutlass": UnquantizedMoeBackend.FLASHINFER_CUTLASS, "aiter": UnquantizedMoeBackend.AITER, @@ -233,6 +234,7 @@ def select_unquantized_moe_backend( activation_format = ( mk.FusedMoEActivationFormat.BatchedExperts if moe_config.moe_parallel_config.use_batched_activation_format + or moe_config.moe_backend == "batched_triton" else mk.FusedMoEActivationFormat.Standard ) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py index 943027717bb..257b5f04dc0 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py @@ -112,15 +112,28 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): a1_scale = normalize_scales_shape(quant_config.a1_scale) - for expert_id in range(first_expert, last_expert): - topks = torch.any(topk_ids == expert_id, dim=1).flatten() - rows = torch.count_nonzero(topks.flatten()) - if rows == 0: - continue - idx = expert_id - first_expert - tokens_per_expert[idx] = rows - rhs = a1[: topks.numel()][topks] - if quant_config.quant_dtype is not None: + if quant_config.quant_dtype is None: + # Vectorized dispatch: compute the [E_local, T] hit mask and + # per-expert slot offsets in one shot (single nonzero sync instead + # of a launch-bound per-expert Python loop). + local_ids = torch.arange(first_expert, last_expert, device=a1.device).view( + -1, 1, 1 + ) + hits = (topk_ids.unsqueeze(0) == local_ids).any(dim=2) # [E_local, T] + tokens_per_expert[:num_local_experts] = hits.sum(dim=1).to(torch.int32) + # slot of each token within its expert batch, preserving token order + slots = hits.to(torch.int32).cumsum(dim=1) - 1 # [E_local, T] + e_idx, t_idx = hits.nonzero(as_tuple=True) + b_a1[e_idx, slots[e_idx, t_idx]] = a1[t_idx].to(b_type) + else: + for expert_id in range(first_expert, last_expert): + topks = torch.any(topk_ids == expert_id, dim=1).flatten() + rows = torch.count_nonzero(topks.flatten()) + if rows == 0: + continue + idx = expert_id - first_expert + tokens_per_expert[idx] = rows + rhs = a1[: topks.numel()][topks] if a1_scale is not None: if quant_config.is_per_act_token: rhs_a1_scale = a1_scale[: topks.numel()][topks] @@ -140,8 +153,6 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): b_a1_scale[idx, :rows] = b_s[:rows] else: b_a1_scale[idx, : b_s.shape[0]] = b_s - else: - b_a1[idx, :rows, :] = rhs assert b_a1_scale is None or b_a1_scale.ndim == 3 diff --git a/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py b/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py index 81c84d64be0..f9543358ece 100644 --- a/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py +++ b/vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py @@ -164,13 +164,20 @@ class TopKWeightAndReduceNaiveBatched(mk.TopKWeightAndReduce): first_expert = num_local_experts * self.rank last_expert = first_expert + num_local_experts - for expert_id in range(first_expert, last_expert): - matching_tokens = topk_ids == expert_id - topks = torch.any(matching_tokens, dim=1).flatten() - rows = torch.count_nonzero(topks) - rhs = fused_expert_output[expert_id - first_expert, :rows, :] - if not apply_router_weight_on_input: - rhs.mul_(topk_weights[matching_tokens].view(rhs.size(0), 1)) - output[topks] = output[topks] + rhs + # Vectorized weighted scatter-add (single nonzero sync instead of a + # per-expert loop; mirrors the dispatch path). + local_ids = torch.arange( + first_expert, last_expert, device=topk_ids.device + ).view(-1, 1, 1) + matching = topk_ids.unsqueeze(0) == local_ids # [E_local, T, topk] + hits = matching.any(dim=2) # [E_local, T] + slots = hits.to(torch.int32).cumsum(dim=1) - 1 # [E_local, T] + e_idx, t_idx = hits.nonzero(as_tuple=True) + gathered = fused_expert_output[e_idx, slots[e_idx, t_idx], :] + if not apply_router_weight_on_input: + # weight of token t at expert e: the topk weight on the matching slot + weights = (matching * topk_weights.unsqueeze(0)).sum(dim=2) # [E_local, T] + gathered = gathered * weights[e_idx, t_idx].unsqueeze(1).to(gathered.dtype) + output.index_add_(0, t_idx, gathered.to(output.dtype)) return output diff --git a/vllm/triton_utils/__init__.py b/vllm/triton_utils/__init__.py index f4866a702dd..e20cb002cdc 100644 --- a/vllm/triton_utils/__init__.py +++ b/vllm/triton_utils/__init__.py @@ -17,7 +17,17 @@ else: tl = TritonLanguagePlaceholder() tldevice = TritonLanguagePlaceholder() +from vllm.triton_utils.tensor_descriptor import use_tensor_descriptor + LOG2E = 1.4426950408889634 LOGE2 = 0.6931471805599453 -__all__ = ["HAS_TRITON", "triton", "tl", "tldevice", "LOG2E", "LOGE2"] +__all__ = [ + "HAS_TRITON", + "triton", + "tl", + "tldevice", + "LOG2E", + "LOGE2", + "use_tensor_descriptor", +] diff --git a/vllm/triton_utils/tensor_descriptor.py b/vllm/triton_utils/tensor_descriptor.py new file mode 100644 index 00000000000..4472d39c9c1 --- /dev/null +++ b/vllm/triton_utils/tensor_descriptor.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +def use_tensor_descriptor(override: bool | None = None) -> bool: + """Tri-state VLLM_TRITON_USE_TD: unset=auto (on for XPU), 1/0=force on/off.""" + from vllm import envs + from vllm.platforms import current_platform + + if override is None: + override = envs.VLLM_TRITON_USE_TD + if override is not None: + return override + return current_platform.is_xpu() From 7a6a5b3667160ab6c05fbda0503b12a1f3bff32e Mon Sep 17 00:00:00 2001 From: Chang Guo <changg@nvidia.com> Date: Sat, 25 Jul 2026 17:51:21 -0700 Subject: [PATCH 052/185] [CI] Compute speech WER directly with jiwer (#49773) --- .../correctness/test_transcription_api_correctness.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py b/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py index af61ebc5264..62b1d367841 100644 --- a/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py +++ b/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py @@ -17,7 +17,7 @@ import pytest import soundfile import torch from datasets import Audio, load_dataset -from evaluate import load +from jiwer import wer from transformers.models.whisper.english_normalizer import EnglishTextNormalizer from vllm.benchmarks.datasets.datasets import ASRDataset @@ -202,8 +202,7 @@ def run_evaluation( # Compute WER predictions = [res[2] for res in results] references = [res[3] for res in results] - wer = load("wer") - wer_score = 100 * wer.compute(references=references, predictions=predictions) + wer_score = 100 * wer(references, predictions) print("WER:", wer_score) return wer_score @@ -302,8 +301,7 @@ def run_longform_evaluation( predictions = [res[2] for res in results] references = [res[3] for res in results] - wer = load("wer") - wer_score = 100 * wer.compute(references=references, predictions=predictions) + wer_score = 100 * wer(references, predictions) print("WER:", wer_score) return wer_score From b153ae6089e9ec3272c423340d2116da97b904ce Mon Sep 17 00:00:00 2001 From: liuzhenwei <zhenweiliu@habana.ai> Date: Sun, 26 Jul 2026 10:01:15 +0800 Subject: [PATCH 053/185] [XPU][CI] add heterogeneous TP UT (#49651) Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com> --- .buildkite/intel_jobs/misc_intel.yaml | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/.buildkite/intel_jobs/misc_intel.yaml b/.buildkite/intel_jobs/misc_intel.yaml index 047b00c49c7..5ce12269fc4 100644 --- a/.buildkite/intel_jobs/misc_intel.yaml +++ b/.buildkite/intel_jobs/misc_intel.yaml @@ -125,13 +125,13 @@ steps: pytest -v -s v1/kv_offload && pytest -v -s v1/kv_connector/unit/test_offloading_connector.py' -- label: NixlConnector PD accuracy (2 GPUs) +- label: NixlConnector PD accuracy (4 GPUs) timeout_in_minutes: 60 - num_devices: 2 + num_devices: 4 device: intel_gpu agent_tags: label: production - gpu: 2+ + gpu: 4+ mem: 16+ no_plugin: true working_dir: "." @@ -148,7 +148,10 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && - bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh' + bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh && + PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1 bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh && + PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh && + PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh' - label: Regression key: regression From 48ebd6f2f164057043edb8301771a17dcaebd61d Mon Sep 17 00:00:00 2001 From: "achyuthan.s" <113010327+Achyuthan-S@users.noreply.github.com> Date: Sun, 26 Jul 2026 06:24:27 +0400 Subject: [PATCH 054/185] [Bugfix][KVConnector] Disable cross-layer KV blocks for per-token-head quant (#49226) Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com> Co-authored-by: Or Ozeri <oro@il.ibm.com> --- vllm/v1/worker/kv_connector_model_runner_mixin.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/vllm/v1/worker/kv_connector_model_runner_mixin.py b/vllm/v1/worker/kv_connector_model_runner_mixin.py index c2c54e647df..7a640b875fa 100644 --- a/vllm/v1/worker/kv_connector_model_runner_mixin.py +++ b/vllm/v1/worker/kv_connector_model_runner_mixin.py @@ -155,6 +155,10 @@ class KVConnectorModelRunnerMixin: kv_cache_spec = attn_group.kv_cache_spec if not isinstance(kv_cache_spec, AttentionSpec): return False + # Per-token-head quant carves inline-scale views that assume per-layer + # contiguous KV buffers; the cross-layer layout breaks this and corrupts KV. + if kv_cache_spec.kv_quant_mode.is_per_token_head: + return False return kv_cache_spec.indexes_kv_by_block_stride @staticmethod From 1240c74c0a47473449cf0c3a9c2d87a1e159f73b Mon Sep 17 00:00:00 2001 From: Athrael Soju <athrael.soju@gmail.com> Date: Sun, 26 Jul 2026 06:08:07 +0200 Subject: [PATCH 055/185] [Bugfix] Respect declared attention contract for ColQwen3.5 retrievers (#49372) Signed-off-by: Athrael Soju <athrael.soju@gmail.com> --- .../multimodal/pooling/test_colqwen3_5.py | 111 ++++++++++++++++-- vllm/model_executor/models/config.py | 40 +++++-- 2 files changed, 135 insertions(+), 16 deletions(-) diff --git a/tests/models/multimodal/pooling/test_colqwen3_5.py b/tests/models/multimodal/pooling/test_colqwen3_5.py index 3513bd025b7..43914d819b8 100644 --- a/tests/models/multimodal/pooling/test_colqwen3_5.py +++ b/tests/models/multimodal/pooling/test_colqwen3_5.py @@ -7,6 +7,8 @@ ColBERT-style late interaction scoring (MaxSim). It produces per-token embeddings for both text and image inputs. """ +from types import SimpleNamespace + import pytest import torch @@ -154,12 +156,14 @@ def test_colqwen3_5_relevance_ordering( _run_relevance_test(vllm_runner, model, dtype=dtype) -def test_colqwen3_5_config_enables_bidirectional_attention() -> None: - """ColQwen3.5 retrieval must be served BIDIRECTIONAL (is_causal=False) so the - full_attention layers build with AttentionType.ENCODER_ONLY. This guards the - silent-causal regression (no GPU / model load needed).""" - from types import SimpleNamespace - +@pytest.mark.parametrize( + ("contract", "expected_is_causal"), + [("causal", True), ("bidirectional", False)], +) +def test_colqwen3_5_config_applies_declared_attention_contract( + contract: str, + expected_is_causal: bool, +) -> None: from vllm.model_executor.models.config import ( MODELS_CONFIG_MAP, ColQwen3_5Config, @@ -167,6 +171,97 @@ def test_colqwen3_5_config_enables_bidirectional_attention() -> None: assert MODELS_CONFIG_MAP["ColQwen3_5"] is ColQwen3_5Config - model_config = SimpleNamespace(hf_config=SimpleNamespace()) + hf_config = SimpleNamespace(retrieval_attention_contract=contract) + text_config = SimpleNamespace() + model_config = SimpleNamespace( + hf_config=hf_config, + hf_text_config=text_config, + ) ColQwen3_5Config.verify_and_update_model_config(model_config) - assert model_config.hf_config.is_causal is False + assert hf_config.is_causal is expected_is_causal + assert text_config.is_causal is expected_is_causal + + +@pytest.mark.parametrize( + "hf_config", + [ + SimpleNamespace(), + SimpleNamespace(retrieval_attention_contract="unsupported"), + SimpleNamespace( + retrieval_attention_contract="causal", + text_config=SimpleNamespace(retrieval_attention_contract="bidirectional"), + ), + ], +) +def test_colqwen3_5_config_rejects_invalid_attention_contract(hf_config) -> None: + from vllm.model_executor.models.config import ColQwen3_5Config + + text_config = getattr(hf_config, "text_config", SimpleNamespace()) + model_config = SimpleNamespace( + hf_config=hf_config, + hf_text_config=text_config, + ) + with pytest.raises(ValueError, match="retrieval_attention_contract"): + ColQwen3_5Config.verify_and_update_model_config(model_config) + + +def test_colqwen3_5_bidirectional_contract_builds_encoder_only_attention( + monkeypatch, +) -> None: + from vllm.model_executor.models import qwen3_next + from vllm.model_executor.models.config import ColQwen3_5Config + from vllm.v1.attention.backend import AttentionType + + hf_config = SimpleNamespace(retrieval_attention_contract="bidirectional") + text_config = SimpleNamespace( + hidden_size=256, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=128, + max_position_embeddings=4096, + rope_parameters={}, + rms_norm_eps=1e-6, + ) + model_config = SimpleNamespace( + hf_config=hf_config, + hf_text_config=text_config, + ) + ColQwen3_5Config.verify_and_update_model_config(model_config) + + captured = {} + + class FakeAttention(torch.nn.Module): + def __init__(self, *args, **kwargs) -> None: + super().__init__() + captured["attn_type"] = kwargs["attn_type"] + + monkeypatch.setattr(qwen3_next, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr( + qwen3_next, "QKVParallelLinear", lambda *args, **kwargs: torch.nn.Identity() + ) + monkeypatch.setattr( + qwen3_next, "RowParallelLinear", lambda *args, **kwargs: torch.nn.Identity() + ) + monkeypatch.setattr( + qwen3_next, + "get_rope", + lambda *args, **kwargs: SimpleNamespace(is_neox_style=False), + ) + monkeypatch.setattr( + qwen3_next, "Qwen3NextRMSNorm", lambda *args, **kwargs: torch.nn.Identity() + ) + monkeypatch.setattr(qwen3_next, "Attention", FakeAttention) + + qwen3_next.Qwen3NextAttention(text_config) + + assert captured["attn_type"] is AttentionType.ENCODER_ONLY + + +def test_colqwen3_5_encoder_only_attention_has_no_kv_cache_spec() -> None: + from vllm.model_executor.layers.attention import Attention + from vllm.v1.attention.backend import AttentionType + + attention = SimpleNamespace(attn_type=AttentionType.ENCODER_ONLY) + vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=16)) + + assert Attention.get_kv_cache_spec(attention, vllm_config) is None diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 96232237a06..3219a5edcda 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -769,17 +769,41 @@ class Qwen3_5ForConditionalGenerationConfig(VerifyAndUpdateConfig): class ColQwen3_5Config(Qwen3_5ForConditionalGenerationConfig): - """ColQwen3.5 (late-interaction retrieval) inherits Qwen3.5's mamba cache - handling and additionally serves BIDIRECTIONAL attention: ColPali-style - document/query encoding attends over the whole sequence, not causally. Set - is_causal=False so Qwen3NextAttention builds its full_attention layers with - AttentionType.ENCODER_ONLY (the linear_attention GatedDeltaNet layers are - unaffected). Generation arches keep the parent (causal) and are untouched. - """ + """Apply the attention contract declared by a ColQwen3.5 checkpoint.""" @staticmethod def verify_and_update_model_config(model_config: "ModelConfig") -> None: - model_config.hf_config.is_causal = False + configs = { + id(config): config + for config in ( + model_config.hf_config, + model_config.hf_text_config, + ) + } + declarations = [ + contract + for config in configs.values() + if (contract := getattr(config, "retrieval_attention_contract", None)) + is not None + ] + supported = {"causal", "bidirectional"} + if not declarations: + raise ValueError( + "ColQwen3.5 checkpoints must declare " + "retrieval_attention_contract as 'causal' or 'bidirectional'" + ) + if ( + any(not isinstance(contract, str) for contract in declarations) + or any(contract not in supported for contract in declarations) + or len(set(declarations)) != 1 + ): + raise ValueError( + "unsupported or conflicting ColQwen3.5 " + f"retrieval_attention_contract declarations: {declarations!r}" + ) + is_causal = declarations[0] == "causal" + for config in configs.values(): + config.is_causal = is_causal class SnowflakeGteNewModelConfig(VerifyAndUpdateConfig): From 7a29a3c54cc338b0103005978f38764e40299572 Mon Sep 17 00:00:00 2001 From: Jonguk Cheong <jdal3031@snu.ac.kr> Date: Sun, 26 Jul 2026 14:21:52 +0900 Subject: [PATCH 056/185] [Bugfix][KV Offload] Namespace persistent cache by model runner (#49440) Signed-off-by: Jonguk Cheong <jdal3031@snu.ac.kr> Co-authored-by: Or Ozeri <oro@il.ibm.com> --- tests/v1/kv_offload/test_file_mapper.py | 21 ++++++++++++++++++++- vllm/v1/kv_offload/file_mapper.py | 2 ++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/tests/v1/kv_offload/test_file_mapper.py b/tests/v1/kv_offload/test_file_mapper.py index 6c11f2d465f..c2c4427e184 100644 --- a/tests/v1/kv_offload/test_file_mapper.py +++ b/tests/v1/kv_offload/test_file_mapper.py @@ -85,7 +85,7 @@ def test_get_file_name_full_structure(): path = fm.get_file_name(key) expected_path = ( - "/tmp/cache/test-model_42b94bdc9933_r3/000/10_g2/0001020304050607.bin" + "/tmp/cache/test-model_de3bba26cf36_r3/000/10_g2/0001020304050607.bin" ) assert path == expected_path @@ -119,6 +119,7 @@ def test_get_run_config_fields(): } ], "inference_engine": "vllm", + "parallel_agnostic": False, } @@ -164,6 +165,7 @@ def test_parallel_agnostic_collapses_namespace_when_config_allows(): assert fm.fields["pcp_size"] == 1 assert fm.fields["dcp_size"] == 1 assert fm.rank == 0 + assert "parallel_agnostic" not in fm.fields def test_parallel_agnostic_ignored_when_config_disallows(): @@ -175,6 +177,7 @@ def test_parallel_agnostic_ignored_when_config_disallows(): ) assert fm.fields["tp_size"] == 2 assert fm.rank == 1 + assert fm.fields["parallel_agnostic"] is False def test_namespace_kept_without_parallel_agnostic_opt_in(): @@ -186,3 +189,19 @@ def test_namespace_kept_without_parallel_agnostic_opt_in(): ) assert fm.fields["tp_size"] == 2 assert fm.rank == 1 + assert fm.fields["parallel_agnostic"] is False + + +def test_parallel_agnostic_separates_persistent_layouts(): + agnostic = make_mapper_from_offloading_spec( + is_parallelism_agnostic=True, + parallel_agnostic=True, + ) + specific = make_mapper_from_offloading_spec( + is_parallelism_agnostic=False, + parallel_agnostic=True, + ) + + assert agnostic.base_path != specific.base_path + assert "parallel_agnostic" not in agnostic.fields + assert specific.fields["parallel_agnostic"] is False diff --git a/vllm/v1/kv_offload/file_mapper.py b/vllm/v1/kv_offload/file_mapper.py index b85d4d06979..8e8c19d53d6 100644 --- a/vllm/v1/kv_offload/file_mapper.py +++ b/vllm/v1/kv_offload/file_mapper.py @@ -58,6 +58,8 @@ class FileMapper: "kv_cache_groups": kv_cache_groups or [], "inference_engine": inference_engine, } + if not parallel_agnostic: + self.fields["parallel_agnostic"] = False self.base_path: str = self._compute_base_path(root_dir, self.fields) @classmethod From 7eca0e1a649e9fb85e88c101432d8bfdbb5b9bfe Mon Sep 17 00:00:00 2001 From: Chang Guo <changg@nvidia.com> Date: Sat, 25 Jul 2026 22:22:30 -0700 Subject: [PATCH 057/185] [KV Offload] Deduplicate replicated MLA KV in the shared CPU region (#48906) Signed-off-by: Change72 <changg@nvidia.com> Co-authored-by: OpenAI Codex <noreply@openai.com> --- .buildkite/test_areas/lm_eval.yaml | 4 +- tests/evals/gsm8k/test_gsm8k_offloading.py | 30 +- .../unit/offloading_connector/test_config.py | 587 +++++++++++++++ .../offloading_connector/test_scheduler.py | 47 ++ .../unit/offloading_connector/test_worker.py | 185 ++++- .../unit/offloading_connector/utils.py | 4 + tests/v1/kv_offload/cpu/test_gpu_worker.py | 12 +- tests/v1/kv_offload/test_factory.py | 707 ++++++------------ .../kv_connector/v1/offloading/config.py | 35 +- .../kv_connector/v1/offloading/worker.py | 11 + vllm/v1/kv_offload/base.py | 1 + vllm/v1/kv_offload/config.py | 6 + vllm/v1/kv_offload/cpu/spec.py | 10 +- vllm/v1/kv_offload/tiering/spec.py | 10 +- 14 files changed, 1143 insertions(+), 506 deletions(-) create mode 100644 tests/v1/kv_connector/unit/offloading_connector/test_config.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 7d0f987d109..f891d9fd24a 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -337,7 +337,7 @@ steps: - label: LM Eval KV-Offload (2xH100) key: kv-offload-medium - timeout_in_minutes: 30 + timeout_in_minutes: 45 device: h100 num_devices: 2 source_file_dependencies: @@ -347,7 +347,7 @@ steps: - vllm/v1/simple_kv_offload/ - tests/evals/gsm8k/test_gsm8k_offloading.py commands: - - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b" + - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b or deepseek-v2-lite" - label: LM Eval KV-Offload (4xH100) key: kv-offload-large diff --git a/tests/evals/gsm8k/test_gsm8k_offloading.py b/tests/evals/gsm8k/test_gsm8k_offloading.py index f652dcaf1bb..7ed97f7efcd 100644 --- a/tests/evals/gsm8k/test_gsm8k_offloading.py +++ b/tests/evals/gsm8k/test_gsm8k_offloading.py @@ -16,11 +16,13 @@ KV drops accuracy below threshold. The reload itself is not asserted, so a silently skipped reload (e.g. offloading disabled) would not be flagged. Covers both KV offloading connectors (OffloadingConnector and -SimpleCPUOffloadConnector) across four architecture families: +SimpleCPUOffloadConnector) across five architecture families: - Hybrid Mamba (NemotronH: attention + Mamba) - Heterogeneous head dim (Gemma 4) - Hybrid GDN (Qwen 3.5: attention + GatedDeltaNet) - Compressed attention (DeepSeek-V4-Flash: CSA) + - Pure MLA (DeepSeek-V2-Lite: TieringOffloadingSpec, TP=2, replicated + single-slot host layout) Usage: pytest -s -v evals/gsm8k/test_gsm8k_offloading.py @@ -47,14 +49,16 @@ NUM_FEWSHOT = 5 _OFFLOAD_SYNC_TIMEOUT = 60 -def _kv_transfer_config(connector: str, cpu_gib: int = 4) -> str: +def _kv_transfer_config( + connector: str, cpu_gib: int = 4, spec_name: str = "CPUOffloadingSpec" +) -> str: if connector == "OffloadingConnector": return json.dumps( { "kv_connector": "OffloadingConnector", "kv_role": "kv_both", "kv_connector_extra_config": { - "spec_name": "CPUOffloadingSpec", + "spec_name": spec_name, "cpu_bytes_to_use": cpu_gib << 30, "eviction_policy": "lru", }, @@ -115,8 +119,10 @@ class OffloadingModelConfig: accuracy_threshold: float tolerance: float = 0.05 extra_server_args: list[str] = field(default_factory=list) + env_dict: dict[str, str] = field(default_factory=dict) cpu_offload_gib: int = 4 startup_timeout: int = 600 + spec_name: str = "CPUOffloadingSpec" MODELS = [ @@ -165,6 +171,20 @@ MODELS = [ cpu_offload_gib=16, startup_timeout=1200, ), + OffloadingModelConfig( + id="offloading-deepseek-v2-lite-tiering-tp2", + model="deepseek-ai/DeepSeek-V2-Lite", + connector="OffloadingConnector", + # Baseline 0.360/0.345 over two runs (measured on 2xA100). + accuracy_threshold=0.35, + extra_server_args=[ + "--tensor-parallel-size", + "2", + ], + cpu_offload_gib=8, + startup_timeout=1200, + spec_name="TieringOffloadingSpec", + ), # ── SimpleCPUOffloadConnector ──────────────────────────────────── OffloadingModelConfig( id="simple-nemotron-h-8b", @@ -229,7 +249,7 @@ def test_gsm8k_offloading_correctness(cfg: OffloadingModelConfig): "--enable-prefix-caching", "--no-disable-hybrid-kv-cache-manager", "--kv-transfer-config", - _kv_transfer_config(cfg.connector, cfg.cpu_offload_gib), + _kv_transfer_config(cfg.connector, cfg.cpu_offload_gib, cfg.spec_name), "--trust-remote-code", "--disable-uvicorn-access-log", *cfg.extra_server_args, @@ -239,7 +259,7 @@ def test_gsm8k_offloading_correctness(cfg: OffloadingModelConfig): cfg.model, server_args, # /reset_prefix_cache requires dev mode. - env_dict={"VLLM_SERVER_DEV_MODE": "1"}, + env_dict={"VLLM_SERVER_DEV_MODE": "1", **cfg.env_dict}, max_wait_seconds=cfg.startup_timeout, ) as server: base_url = f"http://{server.host}:{server.port}" diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py new file mode 100644 index 00000000000..fc426ff318d --- /dev/null +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -0,0 +1,587 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for translating vLLM cache metadata to native offloading config.""" + +from typing import Any, cast +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from vllm.config import KVTransferConfig, ParallelConfig, VllmConfig +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import ( + build_offloading_config, +) +from vllm.platforms import current_platform +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + HiddenStateCacheSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheTensor, + MambaSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, + SlidingWindowSpec, + UniformTypeKVCacheSpecs, +) + + +def _make_vllm_config( + *, + extra_config: dict[str, Any] | None = None, + tensor_parallel_size: int = 1, + pipeline_parallel_size: int = 1, + prefill_context_parallel_size: int = 1, + decode_context_parallel_size: int = 1, +) -> VllmConfig: + config = MagicMock() + config.cache_config.block_size = 16 + config.cache_config.enable_prefix_caching = True + config.cache_config.prefix_match_unit = None + config.cache_config.cache_dtype = torch.float16 + config.model_config.model = "test-model" + config.model_config.use_mla = False + world_size = ( + tensor_parallel_size * pipeline_parallel_size * prefill_context_parallel_size + ) + with patch.object(current_platform, "device_count", return_value=world_size): + config.parallel_config = ParallelConfig( + tensor_parallel_size=tensor_parallel_size, + pipeline_parallel_size=pipeline_parallel_size, + prefill_context_parallel_size=prefill_context_parallel_size, + decode_context_parallel_size=decode_context_parallel_size, + ) + config.kv_events_config = None + config.use_v2_model_runner = False + config.kv_transfer_config = KVTransferConfig( + kv_connector="OffloadingConnector", + kv_role="kv_both", + kv_connector_extra_config=dict(extra_config or {}), + ) + return cast(VllmConfig, config) + + +def _make_kv_cache_config() -> KVCacheConfig: + num_blocks = 16 + kv_tensor = KVCacheTensor( + size=num_blocks * 8, + shared_by=["layer"], + block_stride=0, + ) + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[kv_tensor], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + FullAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + ], + ) + + +def _make_sizing_kv_cache_config(packed: bool) -> KVCacheConfig: + num_blocks = 4 + if packed: + kv_cache_tensors = [ + KVCacheTensor( + size=64, + shared_by=[layer_name], + block_stride=16, + ) + for layer_name in ("layer0", "layer1") + ] + else: + kv_cache_tensors = [ + KVCacheTensor(size=40, shared_by=["layer0"]), + KVCacheTensor(size=24, shared_by=["layer1"]), + ] + + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=kv_cache_tensors, + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer0", "layer1"], + FullAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + ], + ) + + +def _full_attention_spec(block_size: int = 16) -> FullAttentionSpec: + return FullAttentionSpec( + block_size=block_size, + num_kv_heads=4, + head_size=128, + dtype=torch.float32, + ) + + +def _mla_spec( + block_size: int = 16, + head_size: int = 512, + dtype: torch.dtype = torch.float32, +) -> MLAAttentionSpec: + return MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=head_size, + dtype=dtype, + ) + + +def _make_mla_kv_cache_config( + layer_names: list[str] | None = None, + head_size: int = 512, + dtype: torch.dtype = torch.float32, + num_blocks: int = 4, +) -> KVCacheConfig: + if layer_names is None: + layer_names = ["layer0", "layer1"] + spec = _mla_spec(head_size=head_size, dtype=dtype) + kv_cache_tensors = [ + KVCacheTensor( + size=spec.page_size_bytes * num_blocks, + shared_by=[layer_name], + ) + for layer_name in layer_names + ] + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=kv_cache_tensors, + kv_cache_groups=[KVCacheGroupSpec(layer_names, spec)], + ) + + +def _make_hybrid_kv_cache_config() -> KVCacheConfig: + return KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor(size=40, shared_by=["full_layer"]), + KVCacheTensor(size=24, shared_by=["mla_layer"]), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["full_layer"], _full_attention_spec(block_size=12)), + KVCacheGroupSpec(["mla_layer"], _mla_spec()), + ], + ) + + +def _parallelism_agnostic(kv_cache_groups: list[KVCacheGroupSpec]) -> bool: + config = _make_vllm_config() + kv_cache_config = KVCacheConfig( + num_blocks=0, + kv_cache_tensors=[], + kv_cache_groups=kv_cache_groups, + ) + return build_offloading_config( + config, kv_cache_config + ).parallel.is_parallelism_agnostic + + +def _replicated_layout( + kv_cache_config: KVCacheConfig, + *, + tensor_parallel_size: int = 4, + pipeline_parallel_size: int = 1, + prefill_context_parallel_size: int = 1, + decode_context_parallel_size: int = 1, + use_mla: bool = True, + use_v2_model_runner: bool = False, + distributed_executor_backend: Any = "mp", + nnodes: int = 1, + world_size: int | None = None, +) -> bool: + config = _make_vllm_config( + tensor_parallel_size=tensor_parallel_size, + pipeline_parallel_size=pipeline_parallel_size, + prefill_context_parallel_size=prefill_context_parallel_size, + decode_context_parallel_size=decode_context_parallel_size, + ) + config.model_config.use_mla = use_mla + config.use_v2_model_runner = use_v2_model_runner + config.parallel_config.distributed_executor_backend = distributed_executor_backend + config.parallel_config.nnodes = nnodes + if world_size is not None: + config.parallel_config.world_size = world_size + return build_offloading_config(config, kv_cache_config).replicated_layout + + +@pytest.mark.parametrize("packed", [False, True]) +def test_worker_kv_bytes_preserves_tensor_layout(packed: bool): + config = _make_vllm_config( + extra_config={"block_size": 32}, + tensor_parallel_size=3, + pipeline_parallel_size=2, + ) + + offloading_config = build_offloading_config( + config, _make_sizing_kv_cache_config(packed) + ) + + assert offloading_config.worker_kv_bytes_per_block == 16 + assert offloading_config.parallel.world_size == 6 + assert offloading_config.cache.blocks_per_chunk == 2 + + +def test_rejects_partially_packed_tensor_layout(): + kv_cache_config = _make_sizing_kv_cache_config(packed=False) + kv_cache_config.kv_cache_tensors[0].block_stride = 16 + + with pytest.raises(AssertionError): + build_offloading_config(_make_vllm_config(), kv_cache_config) + + +def test_zero_blocks_skips_tensor_layout_validation(): + kv_cache_config = _make_sizing_kv_cache_config(packed=False) + kv_cache_config.num_blocks = 0 + kv_cache_config.kv_cache_tensors[0].block_stride = 16 + + offloading_config = build_offloading_config(_make_vllm_config(), kv_cache_config) + + assert offloading_config.worker_kv_bytes_per_block == 0 + + +def test_prefill_context_parallelism_does_not_scale_group_blocks(): + config = _make_vllm_config( + extra_config={"block_size": 64}, + prefill_context_parallel_size=2, + ) + + offloading_config = build_offloading_config(config, _make_kv_cache_config()) + + assert tuple(group.tokens_per_block for group in offloading_config.groups) == (16,) + assert offloading_config.cache.tokens_per_hash == 16 + assert offloading_config.cache.blocks_per_chunk == 4 + + +def test_preserves_data_parallel_index(): + config = _make_vllm_config() + config.parallel_config.data_parallel_index = 2 + + offloading_config = build_offloading_config(config, _make_kv_cache_config()) + + assert offloading_config.parallel.data_parallel_index == 2 + + +def test_resolves_heterogeneous_hybrid_block_sizes(): + config = _make_vllm_config() + config.cache_config.block_size = 4 + + offloading_config = build_offloading_config(config, _make_hybrid_kv_cache_config()) + + assert tuple(group.tokens_per_block for group in offloading_config.groups) == ( + 12, + 16, + ) + assert offloading_config.cache.tokens_per_hash == 4 + assert offloading_config.cache.blocks_per_chunk == 1 + + +@pytest.mark.parametrize("world_size", [2, 4, 8]) +@pytest.mark.parametrize("use_v2_model_runner", [False, True], ids=["v1", "v2"]) +def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node( + world_size: int, + use_v2_model_runner: bool, +): + assert _replicated_layout( + _make_mla_kv_cache_config(), + tensor_parallel_size=world_size, + use_v2_model_runner=use_v2_model_runner, + ) + + +@pytest.mark.parametrize( + ("kv_cache_config", "case"), + [ + ( + KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=_mla_spec().page_size_bytes * 4, + shared_by=["layer"], + ) + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + SlidingWindowMLASpec( + block_size=16, + num_kv_heads=1, + head_size=512, + dtype=torch.float32, + sliding_window=128, + ), + ) + ], + ), + "sliding-window-mla", + ), + ( + KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=_mla_spec().page_size_bytes * 4, + shared_by=["layer"], + ) + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + HiddenStateCacheSpec( + block_size=16, + num_kv_heads=1, + head_size=512, + dtype=torch.float32, + ), + ) + ], + ), + "hidden-state", + ), + ( + KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=_mla_spec().page_size_bytes * 4, + shared_by=["layer0"], + ), + KVCacheTensor( + size=_mla_spec(head_size=256).page_size_bytes * 4, + shared_by=["layer1"], + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer0", "layer1"], + UniformTypeKVCacheSpecs( + block_size=16, + kv_cache_specs={ + "layer0": _mla_spec(), + "layer1": _mla_spec(head_size=256), + }, + ), + ) + ], + ), + "uniform-wrapper", + ), + ( + KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=_mla_spec().page_size_bytes * 4, + shared_by=["mla"], + ), + KVCacheTensor( + size=_full_attention_spec().page_size_bytes * 4, + shared_by=["full"], + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["mla"], _mla_spec()), + KVCacheGroupSpec(["full"], _full_attention_spec()), + ], + ), + "mla-full-hybrid", + ), + ( + KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=_mla_spec().page_size_bytes * 4, + shared_by=["mla"], + ), + KVCacheTensor(size=64 * 4, shared_by=["mamba"]), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["mla"], _mla_spec()), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=16, + shapes=((16, 1),), + dtypes=(torch.float32,), + ), + ), + ], + ), + "mla-mamba-hybrid", + ), + ( + KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=_mla_spec().page_size_bytes * 4, + shared_by=["layer0"], + ), + KVCacheTensor( + size=_mla_spec().page_size_bytes * 4, + shared_by=["layer1"], + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer0"], _mla_spec()), + KVCacheGroupSpec(["layer1"], _mla_spec()), + ], + ), + "multi-group-mla", + ), + ], + ids=[ + "sliding-window-mla", + "hidden-state", + "uniform-wrapper", + "mla-full-hybrid", + "mla-mamba-hybrid", + "multi-group-mla", + ], +) +def test_replicated_layout_excludes_unproven_cache_shapes( + kv_cache_config: KVCacheConfig, + case: str, +): + assert not _replicated_layout(kv_cache_config), case + + +def test_replicated_layout_rejects_bare_mla_with_mixed_page_accounting(): + num_blocks = 4 + main_spec = _mla_spec(head_size=512) + indexer_spec = _mla_spec(head_size=128, dtype=torch.uint8) + main_layers = [f"main_{i}" for i in range(61)] + indexer_layers = [f"indexer_{i}" for i in range(61)] + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=main_spec.page_size_bytes * len(main_layers) * num_blocks, + shared_by=main_layers, + ), + KVCacheTensor( + size=indexer_spec.page_size_bytes * len(indexer_layers) * num_blocks, + shared_by=indexer_layers, + ), + ], + kv_cache_groups=[KVCacheGroupSpec(main_layers + indexer_layers, main_spec)], + ) + + assert not _replicated_layout(kv_cache_config) + + +@pytest.mark.parametrize( + ("kwargs", "case"), + [ + ({"tensor_parallel_size": 1}, "tp1"), + ({"use_mla": False}, "use-mla-false"), + ({"pipeline_parallel_size": 2, "world_size": 4}, "pp2"), + ({"prefill_context_parallel_size": 2, "world_size": 4}, "pcp2"), + ({"decode_context_parallel_size": 2}, "dcp2"), + ({"world_size": 8}, "world-ne-tp"), + ({"distributed_executor_backend": "ray"}, "ray"), + ({"distributed_executor_backend": "uni"}, "uni"), + ({"distributed_executor_backend": type("DummyExecutor", (), {})}, "class"), + ({"nnodes": 2}, "multi-node"), + ], + ids=[ + "tp1", + "use-mla-false", + "pp2", + "pcp2", + "dcp2", + "world-ne-tp", + "ray", + "uni", + "class", + "multi-node", + ], +) +def test_replicated_layout_parallel_gate(kwargs: dict[str, Any], case: str): + assert not _replicated_layout(_make_mla_kv_cache_config(), **kwargs), case + + +def test_parallelism_agnostic_for_single_full_attention_group(): + assert _parallelism_agnostic([KVCacheGroupSpec(["l0"], _full_attention_spec())]) + + +@pytest.mark.parametrize( + "kv_cache_groups", + [ + [KVCacheGroupSpec(["l0"], _mla_spec(head_size=576))], + [ + KVCacheGroupSpec( + ["l0"], + SlidingWindowSpec( + block_size=16, + num_kv_heads=4, + head_size=128, + dtype=torch.float32, + sliding_window=128, + ), + ) + ], + [ + KVCacheGroupSpec(["l0"], _full_attention_spec()), + KVCacheGroupSpec(["l1"], _full_attention_spec()), + ], + ], +) +def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]): + assert not _parallelism_agnostic(kv_cache_groups) + + +def test_parallelism_agnostic_disabled_on_v2_model_runner(): + config = _make_vllm_config() + config.use_v2_model_runner = True + kv_cache_config = KVCacheConfig( + num_blocks=0, + kv_cache_tensors=[], + kv_cache_groups=[KVCacheGroupSpec(["l0"], _full_attention_spec())], + ) + + offloading_config = build_offloading_config(config, kv_cache_config) + + assert not offloading_config.parallel.is_parallelism_agnostic + + +def test_accepts_blocks_per_chunk_for_heterogeneous_groups(): + config = _make_vllm_config(extra_config={"blocks_per_chunk": 2}) + + offloading_config = build_offloading_config(config, _make_hybrid_kv_cache_config()) + + assert tuple(group.tokens_per_block for group in offloading_config.groups) == ( + 12, + 16, + ) + assert offloading_config.cache.blocks_per_chunk == 2 + + +def test_block_size_and_blocks_per_chunk_are_mutually_exclusive(): + config = _make_vllm_config(extra_config={"block_size": 64, "blocks_per_chunk": 2}) + + with pytest.raises(ValueError, match="Specify only one"): + build_offloading_config(config, _make_kv_cache_config()) + + +def test_blocks_per_chunk_must_be_positive(): + config = _make_vllm_config(extra_config={"blocks_per_chunk": 0}) + + with pytest.raises(ValueError, match="greater than 0"): + build_offloading_config(config, _make_kv_cache_config()) 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 c65d02dda4b..d1463aebb4b 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -14,6 +14,7 @@ from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import ( OffloadingConnectorMetadata, + OffloadingWorkerMetadata, ) from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( OffloadingConnectorStats, @@ -40,6 +41,7 @@ from vllm.v1.kv_offload.base import ( get_offload_block_hash, make_offload_key, ) +from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import RequestStatus @@ -1539,6 +1541,51 @@ def test_complete_store_called_per_job(request_runner, async_scheduling: bool): assert runner.manager.complete_store.call_count == 0 +@pytest.mark.parametrize("async_scheduling", [True, False]) +def test_complete_store_waits_for_all_worker_acks( + request_runner, async_scheduling: bool +): + tokens_per_block = 4 + blocks_per_chunk = 3 + tokens_per_chunk = tokens_per_block * blocks_per_chunk + runner = request_runner( + blocks_per_chunk=blocks_per_chunk, + block_size=tokens_per_block, + num_gpu_blocks=100, + async_scheduling=async_scheduling, + worker_count=3, + ) + runner.new_request(token_ids=[0] * tokens_per_chunk) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + runner.run(decoded_tokens=[0, 0], complete_transfers=False) + assert len(runner.connector_scheduler._jobs) == 1 + job_id = next(iter(runner.connector_scheduler._jobs)) + assert runner.connector_scheduler._jobs[job_id].pending_count == 3 + runner.manager.complete_store.reset_mock() + + runner.connector_scheduler.update_connector_output( + KVConnectorOutput( + kv_connector_worker_meta=OffloadingWorkerMetadata( + completed_jobs={job_id: 1} + ) + ) + ) + assert runner.manager.complete_store.call_count == 0 + assert runner.connector_scheduler._jobs[job_id].pending_count == 2 + + runner.connector_scheduler.update_connector_output( + KVConnectorOutput( + kv_connector_worker_meta=OffloadingWorkerMetadata( + completed_jobs={job_id: 2} + ) + ) + ) + assert runner.manager.complete_store.call_count == 1 + assert job_id not in runner.connector_scheduler._jobs + + @pytest.mark.parametrize("async_scheduling", [True, False]) def test_max_offload_tokens_validation(request_runner, async_scheduling: bool): """Validates max_offload_tokens: type coercion, boundary values, and capping. diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py index b4557bd7b62..f96fee36a85 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py @@ -6,6 +6,10 @@ from unittest.mock import MagicMock import pytest import torch +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import ( + OffloadingConnectorMetadata, + TransferJob, +) from vllm.platforms import current_platform from vllm.utils.torch_utils import get_dtype_size from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -22,7 +26,17 @@ from vllm.v1.kv_cache_interface import ( from vllm.v1.kv_offload.base import ( CanonicalKVCacheRef, CanonicalKVCaches, + GPULoadStoreSpec, + LoadStoreSpec, + OffloadingManager, OffloadingSpec, + OffloadingWorker, +) +from vllm.v1.kv_offload.config import ( + OffloadingCacheConfig, + OffloadingConfig, + OffloadingModelConfig, + OffloadingParallelConfig, ) NUM_BLOCKS = 10 @@ -87,7 +101,11 @@ def _allocate_and_reshape_kv_caches( set_kv_cache_layout(None) -def _make_worker(kv_cache_config: KVCacheConfig): +def _make_worker( + kv_cache_config: KVCacheConfig, + replicated_layout: bool = False, + rank: int = 0, +): """ Create an OffloadingConnectorWorker with mocked dependencies. """ @@ -96,6 +114,9 @@ def _make_worker(kv_cache_config: KVCacheConfig): ) spec = MagicMock(spec=OffloadingSpec) + spec.replicated_layout = replicated_layout + spec.config = MagicMock() + spec.config.parallel.rank = rank spec.get_worker.return_value = MagicMock() worker = OffloadingConnectorWorker( @@ -107,11 +128,173 @@ def _make_worker(kv_cache_config: KVCacheConfig): return worker, spec +def _store_metadata(job_id: int) -> OffloadingConnectorMetadata: + return OffloadingConnectorMetadata( + load_jobs={}, + store_jobs={ + job_id: TransferJob( + req_id="req", + src_spec=GPULoadStoreSpec([0], group_sizes=(1,), block_indices=(0,)), + dst_spec=LoadStoreSpec(), + ) + }, + ) + + +def _load_metadata(job_id: int) -> OffloadingConnectorMetadata: + return OffloadingConnectorMetadata( + load_jobs={ + job_id: TransferJob( + req_id="req", + src_spec=LoadStoreSpec(), + dst_spec=GPULoadStoreSpec([0], group_sizes=(1,), block_indices=(0,)), + ) + }, + store_jobs={}, + ) + + +def _empty_metadata() -> OffloadingConnectorMetadata: + return OffloadingConnectorMetadata(load_jobs={}, store_jobs={}) + + +def _offloading_config(rank: int = 0) -> OffloadingConfig: + return OffloadingConfig( + groups=(), + worker_kv_bytes_per_block=0, + enable_kv_cache_events=False, + extra_config={}, + engine_id="test-engine", + model=OffloadingModelConfig(name="test-model", dtype="float16"), + cache=OffloadingCacheConfig(tokens_per_hash=16, blocks_per_chunk=1), + parallel=OffloadingParallelConfig( + rank=rank, + world_size=2, + tp_size=2, + pp_size=1, + pcp_size=1, + dcp_size=1, + data_parallel_index=0, + is_parallelism_agnostic=False, + ), + ) + + +class BareExternalOffloadingSpec(OffloadingSpec): + def get_manager(self) -> OffloadingManager: + raise NotImplementedError + + def get_worker(self, kv_caches: CanonicalKVCaches) -> OffloadingWorker: + raise NotImplementedError + + # --------------------------------------------------------------------------- # Tests # --------------------------------------------------------------------------- +def test_prepare_store_kv_non_writer_marks_completed_without_submit(): + worker, _ = _make_worker( + KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]), + replicated_layout=True, + rank=1, + ) + + worker.prepare_store_kv(_store_metadata(7)) + worker.start_kv_transfers(_empty_metadata()) + + assert worker._unsubmitted_store_jobs == [] + assert worker.worker is not None + worker.worker.submit_store.assert_not_called() + meta = worker.build_connector_worker_meta() + assert meta is not None + assert meta.completed_jobs == {7: 1} + + +def test_prepare_store_kv_writer_submits_store(): + worker, _ = _make_worker( + KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]), + replicated_layout=True, + rank=0, + ) + + worker.prepare_store_kv(_store_metadata(8)) + assert worker.build_connector_worker_meta() is None + worker.start_kv_transfers(_empty_metadata()) + + assert worker.worker is not None + worker.worker.submit_store.assert_called_once() + + +def test_prepare_store_kv_non_replicated_rank_gt_zero_queues_store(): + worker, _ = _make_worker( + KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]), + replicated_layout=False, + rank=1, + ) + + worker.prepare_store_kv(_store_metadata(9)) + assert worker.build_connector_worker_meta() is None + assert len(worker._unsubmitted_store_jobs) == 1 + + worker.start_kv_transfers(_empty_metadata()) + + assert worker.worker is not None + worker.worker.submit_store.assert_called_once() + assert worker._unsubmitted_store_jobs == [] + + +def test_handle_preemptions_non_writer_acks_flushed_store(): + worker, _ = _make_worker( + KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]), + replicated_layout=True, + rank=1, + ) + metadata = _store_metadata(10) + metadata.jobs_to_flush = {10} + + worker.handle_preemptions(metadata) + + assert worker.worker is not None + worker.worker.submit_store.assert_not_called() + worker.worker.wait.assert_called_once_with({10}) + assert metadata.store_jobs == {} + meta = worker.build_connector_worker_meta() + assert meta is not None + assert meta.completed_jobs == {10: 1} + + +def test_start_kv_transfers_non_writer_still_submits_load(): + worker, _ = _make_worker( + KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]), + replicated_layout=True, + rank=1, + ) + + worker.start_kv_transfers(_load_metadata(10)) + + assert worker.worker is not None + worker.worker.submit_load.assert_called_once() + assert worker.build_connector_worker_meta() is None + + +def test_offloading_connector_worker_accepts_plugin_spec_default_layout(): + from vllm.distributed.kv_transfer.kv_connector.v1.offloading.worker import ( + OffloadingConnectorWorker, + ) + + spec = BareExternalOffloadingSpec(_offloading_config(rank=1)) + + OffloadingConnectorWorker( + spec=spec, + kv_cache_config=KVCacheConfig( + num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[] + ), + ) + + assert spec.replicated_layout is False + + @pytest.mark.parametrize("backend", ATTN_BACKENDS) def test_register_kv_caches(backend): """Test register_kv_caches with multiple groups covering all layer types. diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index c1979a0b1f7..9acf5f5d49f 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -179,6 +179,7 @@ class RequestRunner: async_scheduling: bool = True, kv_cache_groups: list[KVCacheGroupSpec] | None = None, extra_config_overrides: dict[str, Any] | None = None, + worker_count: int = 1, ): assert blocks_per_chunk == 1 or kv_cache_groups is None, ( "blocks_per_chunk > 1 requires all groups to have the same " @@ -198,6 +199,7 @@ class RequestRunner: disable_hybrid_kv_cache_manager=False, ) vllm_config.scheduler_config.async_scheduling = async_scheduling + vllm_config.parallel_config.world_size = worker_count extra_config: dict[str, Any] = { "spec_name": "MockOffloadingSpec", @@ -668,6 +670,7 @@ def request_runner(): blocks_per_chunk=1, kv_cache_groups=None, extra_config_overrides=None, + worker_count=1, ): runner = RequestRunner( block_size=block_size, @@ -676,6 +679,7 @@ def request_runner(): async_scheduling=async_scheduling, kv_cache_groups=kv_cache_groups, extra_config_overrides=extra_config_overrides, + worker_count=worker_count, ) runners.append(runner) return runner diff --git a/tests/v1/kv_offload/cpu/test_gpu_worker.py b/tests/v1/kv_offload/cpu/test_gpu_worker.py index 2f1ce67e9c4..4e5e94c2bcd 100644 --- a/tests/v1/kv_offload/cpu/test_gpu_worker.py +++ b/tests/v1/kv_offload/cpu/test_gpu_worker.py @@ -41,7 +41,10 @@ NUM_MAPPINGS_PER_GROUP = [2] @pytest.mark.parametrize("num_tensors", NUM_TENSORS) @pytest.mark.parametrize("seed", SEEDS) @pytest.mark.parametrize("device", DEVICES) -@pytest.mark.parametrize("use_shared_memory", [False, True]) +@pytest.mark.parametrize( + ("use_shared_memory", "replicated_layout"), + [(False, False), (True, False), (True, True)], +) @torch.inference_mode() def test_transfer( default_vllm_config, @@ -55,6 +58,7 @@ def test_transfer( seed: int, device: str, use_shared_memory: bool, + replicated_layout: bool, ) -> None: set_random_seed(seed) @@ -95,11 +99,15 @@ def test_transfer( gpu_page_size_bytes * num_tensors * blocks_per_chunk, SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT, ) + simulated_world_size = 2 + kv_bytes_per_block = ( + cpu_page_size if replicated_layout else cpu_page_size * simulated_world_size + ) mmap_region = SharedOffloadRegion( engine_id=str(uuid.uuid4()), num_blocks=num_cpu_blocks, rank=0, - kv_bytes_per_block=cpu_page_size, + kv_bytes_per_block=kv_bytes_per_block, cpu_page_size=cpu_page_size, ) diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py index d08dbee5765..6d57f1ce70f 100644 --- a/tests/v1/kv_offload/test_factory.py +++ b/tests/v1/kv_offload/test_factory.py @@ -1,35 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for OffloadingSpecFactory. +"""Unit tests for native offloading specs and their factory.""" -These tests verify: -1. Pre-registration integrity — registered module paths can actually import - and yield correct OffloadingSpec subclasses (CI sentinel against file moves). -2. End-to-end factory → spec construction with real configs. -3. Downstream collaboration — build_metric_definitions delegation. -4. Error paths — unregistered specs, missing config, duplicate registration. -""" - -from typing import cast -from unittest.mock import MagicMock, patch +from typing import Any +from unittest.mock import MagicMock import pytest -import torch -from vllm.config import KVTransferConfig, ParallelConfig, VllmConfig -from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import ( - build_offloading_config, -) -from vllm.platforms import current_platform -from vllm.v1.kv_cache_interface import ( - FullAttentionSpec, - KVCacheConfig, - KVCacheGroupSpec, - KVCacheTensor, - MLAAttentionSpec, - SlidingWindowSpec, -) from vllm.v1.kv_offload.base import ( CanonicalKVCaches, OffloadingHistogramMetadata, @@ -37,220 +14,81 @@ from vllm.v1.kv_offload.base import ( OffloadingSpec, OffloadingWorker, ) +from vllm.v1.kv_offload.config import ( + OffloadingCacheConfig, + OffloadingConfig, + OffloadingGroupConfig, + OffloadingModelConfig, + OffloadingParallelConfig, +) from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec from vllm.v1.kv_offload.factory import OffloadingSpecFactory from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec -# --------------------------------------------------------------------------- -# Fixtures -# --------------------------------------------------------------------------- - @pytest.fixture(autouse=True) def restore_registry(): - """Save and restore OffloadingSpecFactory._registry between tests.""" original = dict(OffloadingSpecFactory._registry) yield OffloadingSpecFactory._registry = original -def _get_extra_config(config: VllmConfig) -> dict: - assert config.kv_transfer_config is not None - return config.kv_transfer_config.kv_connector_extra_config - - -def _create_spec(config: VllmConfig, kv_cache_config: KVCacheConfig) -> OffloadingSpec: - return OffloadingSpecFactory.create_spec( - build_offloading_config(config, kv_cache_config) - ) - - -def _make_vllm_config( +def _make_offloading_config( + *, spec_name: str | None = "CPUOffloadingSpec", - cpu_bytes_to_use: int | None = None, - store_threshold: int = 0, - extra_config: dict | None = None, -): - """Build a real VllmConfig with kv_transfer_config set for offloading.""" - from vllm.config import ( - CacheConfig, - DeviceConfig, - ModelConfig, - SchedulerConfig, - VllmConfig, - ) - - model_config = ModelConfig( - model="facebook/opt-125m", - trust_remote_code=True, - dtype="float16", - seed=42, - ) - scheduler_config = SchedulerConfig( - max_num_seqs=16, - max_num_batched_tokens=64, - max_model_len=10000, - enable_chunked_prefill=True, - is_encoder_decoder=model_config.is_encoder_decoder, - ) - cache_config = CacheConfig( - block_size=16, - gpu_memory_utilization=0.9, - cache_dtype="auto", - enable_prefix_caching=True, - ) - - cfg = extra_config or {} + cpu_bytes_to_use: int | None = 65536, + worker_kv_bytes_per_block: int = 8, + groups: tuple[OffloadingGroupConfig, ...] | None = None, + tokens_per_hash: int = 16, + blocks_per_chunk: int = 1, + rank: int = 0, + world_size: int = 1, + tp_size: int | None = None, + pp_size: int = 1, + pcp_size: int = 1, + dcp_size: int = 1, + data_parallel_index: int = 0, + is_parallelism_agnostic: bool = False, + replicated_layout: bool = False, + extra_config: dict[str, Any] | None = None, +) -> OffloadingConfig: + normalized_extra_config = dict(extra_config or {}) + if spec_name is not None: + normalized_extra_config["spec_name"] = spec_name if cpu_bytes_to_use is not None: - cfg["cpu_bytes_to_use"] = cpu_bytes_to_use - cfg["spec_name"] = spec_name - if store_threshold > 0: - cfg["store_threshold"] = store_threshold + normalized_extra_config["cpu_bytes_to_use"] = cpu_bytes_to_use - kv_transfer_config = KVTransferConfig( - kv_connector="OffloadingConnector", - kv_role="kv_both", - kv_connector_extra_config=cfg, - ) - return VllmConfig( - scheduler_config=scheduler_config, - model_config=model_config, - cache_config=cache_config, - kv_transfer_config=kv_transfer_config, - device_config=DeviceConfig("cpu"), + if groups is None: + groups = (OffloadingGroupConfig(16, ("layer",)),) + + return OffloadingConfig( + groups=groups, + worker_kv_bytes_per_block=worker_kv_bytes_per_block, + enable_kv_cache_events=False, + extra_config=normalized_extra_config, + engine_id="test-engine", + model=OffloadingModelConfig(name="test-model", dtype="float16"), + cache=OffloadingCacheConfig( + tokens_per_hash=tokens_per_hash, + blocks_per_chunk=blocks_per_chunk, + ), + parallel=OffloadingParallelConfig( + rank=rank, + world_size=world_size, + tp_size=world_size if tp_size is None else tp_size, + pp_size=pp_size, + pcp_size=pcp_size, + dcp_size=dcp_size, + data_parallel_index=data_parallel_index, + is_parallelism_agnostic=is_parallelism_agnostic, + ), + replicated_layout=replicated_layout, ) -def _make_layout_vllm_config( - spec_name: str = "CPUOffloadingSpec", - cpu_bytes_to_use: int | None = None, - extra_config: dict | None = None, - tensor_parallel_size: int = 1, - pipeline_parallel_size: int = 1, - prefill_context_parallel_size: int = 1, - decode_context_parallel_size: int = 1, -) -> VllmConfig: - config = MagicMock() - config.cache_config.block_size = 16 - config.cache_config.enable_prefix_caching = True - config.cache_config.prefix_match_unit = None - config.cache_config.cache_dtype = torch.float16 - config.model_config.model = "test-model" - world_size = ( - tensor_parallel_size * pipeline_parallel_size * prefill_context_parallel_size - ) - with patch.object(current_platform, "device_count", return_value=world_size): - config.parallel_config = ParallelConfig( - tensor_parallel_size=tensor_parallel_size, - pipeline_parallel_size=pipeline_parallel_size, - prefill_context_parallel_size=prefill_context_parallel_size, - decode_context_parallel_size=decode_context_parallel_size, - ) - config.kv_events_config = None - config.use_v2_model_runner = False - - connector_extra_config = dict(extra_config or {}) - connector_extra_config["spec_name"] = spec_name - if cpu_bytes_to_use is not None: - connector_extra_config["cpu_bytes_to_use"] = cpu_bytes_to_use - config.kv_transfer_config = KVTransferConfig( - kv_connector="OffloadingConnector", - kv_role="kv_both", - kv_connector_extra_config=connector_extra_config, - ) - return cast(VllmConfig, config) - - -def _make_kv_cache_config(): - """Build a minimal KVCacheConfig with one KV cache tensor.""" - num_blocks = 16 - num_kv_heads = 1 - head_size = 1 - dtype = torch.float32 - page_size = 2 * num_kv_heads * head_size * torch.finfo(dtype).bits // 8 - kv_tensor = KVCacheTensor( - size=num_blocks * page_size, shared_by=["layer"], block_stride=0 - ) - return KVCacheConfig( - num_blocks=num_blocks, - kv_cache_tensors=[kv_tensor], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer"], - FullAttentionSpec( - block_size=16, - num_kv_heads=num_kv_heads, - head_size=head_size, - dtype=dtype, - ), - ) - ], - ) - - -def _make_sizing_kv_cache_config(packed: bool) -> KVCacheConfig: - num_blocks = 4 - if packed: - kv_cache_tensors = [ - KVCacheTensor( - size=64, - shared_by=[layer_name], - block_stride=16, - ) - for layer_name in ("layer0", "layer1") - ] - else: - kv_cache_tensors = [ - KVCacheTensor(size=40, shared_by=["layer0"]), - KVCacheTensor(size=24, shared_by=["layer1"]), - ] - - return KVCacheConfig( - num_blocks=num_blocks, - kv_cache_tensors=kv_cache_tensors, - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer0", "layer1"], - FullAttentionSpec( - block_size=16, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - ), - ) - ], - ) - - -def _make_hybrid_kv_cache_config() -> KVCacheConfig: - return KVCacheConfig( - num_blocks=4, - kv_cache_tensors=[ - KVCacheTensor(size=40, shared_by=["full_layer"]), - KVCacheTensor(size=24, shared_by=["mla_layer"]), - ], - kv_cache_groups=[ - KVCacheGroupSpec( - ["full_layer"], - FullAttentionSpec( - block_size=12, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - ), - ), - KVCacheGroupSpec( - ["mla_layer"], - MLAAttentionSpec( - block_size=16, - num_kv_heads=1, - head_size=576, - dtype=torch.float32, - ), - ), - ], - ) +def _create_spec(**kwargs: Any) -> OffloadingSpec: + return OffloadingSpecFactory.create_spec(_make_offloading_config(**kwargs)) class SingleArgExternalOffloadingSpec(OffloadingSpec): @@ -261,104 +99,60 @@ class SingleArgExternalOffloadingSpec(OffloadingSpec): raise NotImplementedError -# --------------------------------------------------------------------------- -# Pre-registration integrity (CI sentinel) -# --------------------------------------------------------------------------- - - def test_pre_registered_specs_can_be_imported(): - """If someone moves cpu/spec.py but forgets to update factory.py, CI fails.""" for name in OffloadingSpecFactory._registry: cls = OffloadingSpecFactory._registry[name]() assert issubclass(cls, OffloadingSpec) def test_cpu_spec_registered(): - """CPUOffloadingSpec is registered and importable.""" cls = OffloadingSpecFactory._registry["CPUOffloadingSpec"]() assert cls is CPUOffloadingSpec def test_tiering_spec_registered(): - """TieringOffloadingSpec is registered and importable.""" cls = OffloadingSpecFactory._registry["TieringOffloadingSpec"]() assert cls is TieringOffloadingSpec -# --------------------------------------------------------------------------- -# Normal path — get_spec_cls -# --------------------------------------------------------------------------- - - def test_get_spec_cls_returns_registered_class(): - """Registered spec_name returns correct class.""" - config = _make_vllm_config(spec_name="CPUOffloadingSpec") - spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config)) + spec_cls = OffloadingSpecFactory.get_spec_cls( + _make_offloading_config().extra_config + ) assert spec_cls is CPUOffloadingSpec -def test_get_spec_cls_default_to_cpu(): - """Default spec_name (absent from config) resolves to CPUOffloadingSpec.""" - config = _make_vllm_config(spec_name=None) - config.kv_transfer_config.kv_connector_extra_config.pop("spec_name", None) - spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config)) +def test_get_spec_cls_defaults_to_cpu(): + spec_cls = OffloadingSpecFactory.get_spec_cls( + _make_offloading_config(spec_name=None).extra_config + ) assert spec_cls is CPUOffloadingSpec -# --------------------------------------------------------------------------- -# End-to-end — create_spec -# --------------------------------------------------------------------------- - - -def test_create_cpu_offloading_spec_end_to_end(): - """Full factory → spec construction with real VllmConfig/KVCacheConfig. - - Verifies: - - cpu_bytes_to_use validation and num_blocks calculation - - block_size % tokens_per_hash assertion - - spec instance is CPUOffloadingSpec - """ - config = _make_vllm_config(cpu_bytes_to_use=65536) - kv_cache_config = _make_kv_cache_config() - spec = _create_spec(config, kv_cache_config) +def test_create_cpu_offloading_spec(): + spec = _create_spec() assert isinstance(spec, CPUOffloadingSpec) assert spec.num_blocks > 0 -@pytest.mark.parametrize("packed", [False, True]) -def test_cpu_spec_sizing_preserves_tensor_layout(packed: bool): - cpu_bytes_to_use = 1920 - config = _make_layout_vllm_config( - cpu_bytes_to_use=cpu_bytes_to_use, - extra_config={"block_size": 32}, - tensor_parallel_size=3, - pipeline_parallel_size=2, +def test_cpu_spec_sizes_normalized_worker_layout(): + spec = _create_spec( + cpu_bytes_to_use=1920, + worker_kv_bytes_per_block=16, + blocks_per_chunk=2, + world_size=6, + tp_size=3, + pp_size=2, ) - spec = _create_spec(config, _make_sizing_kv_cache_config(packed)) - assert isinstance(spec, CPUOffloadingSpec) assert spec.cpu_page_size_per_worker == 32 assert spec.kv_bytes_per_chunk == 192 - assert spec.num_blocks == cpu_bytes_to_use // 192 + assert spec.num_blocks == 10 -def test_cpu_spec_rejects_partially_packed_tensor_layout(): - config = _make_layout_vllm_config(cpu_bytes_to_use=65536) - kv_cache_config = _make_sizing_kv_cache_config(packed=False) - kv_cache_config.kv_cache_tensors[0].block_stride = 16 - - with pytest.raises(AssertionError): - _create_spec(config, kv_cache_config) - - -def test_cpu_spec_zero_blocks_skips_tensor_layout_validation(): - config = _make_layout_vllm_config(cpu_bytes_to_use=65536) - kv_cache_config = _make_sizing_kv_cache_config(packed=False) - kv_cache_config.num_blocks = 0 - kv_cache_config.kv_cache_tensors[0].block_stride = 16 - - spec = _create_spec(config, kv_cache_config) +def test_cpu_spec_zero_worker_bytes_produces_empty_cache(): + spec = _create_spec(worker_kv_bytes_per_block=0, world_size=4) assert isinstance(spec, CPUOffloadingSpec) assert spec.cpu_page_size_per_worker == 0 @@ -368,225 +162,209 @@ def test_cpu_spec_zero_blocks_skips_tensor_layout_validation(): def test_tiering_spec_aligns_row_size(): alignment = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT - cpu_bytes_to_use = alignment * 3 - config = _make_layout_vllm_config( + spec = _create_spec( spec_name="TieringOffloadingSpec", - cpu_bytes_to_use=cpu_bytes_to_use, - extra_config={"block_size": 32}, - tensor_parallel_size=3, - pipeline_parallel_size=2, + cpu_bytes_to_use=alignment * 3, + worker_kv_bytes_per_block=16, + blocks_per_chunk=2, + world_size=6, + tp_size=3, + pp_size=2, ) - spec = _create_spec(config, _make_sizing_kv_cache_config(packed=False)) - assert isinstance(spec, TieringOffloadingSpec) assert spec.cpu_page_size_per_worker == 32 assert spec.kv_bytes_per_chunk == alignment - assert spec.num_blocks == cpu_bytes_to_use // alignment + assert spec.num_blocks == 3 -def test_offloading_spec_kv_sharding_ignores_prefill_context_parallel(): - config = _make_layout_vllm_config( - cpu_bytes_to_use=65536, - extra_config={"block_size": 64}, - prefill_context_parallel_size=2, +@pytest.mark.parametrize("world_size", [2, 4, 8]) +def test_tiering_spec_replicated_sizing_removes_world_factor(world_size: int): + worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT + spec = _create_spec( + spec_name="TieringOffloadingSpec", + cpu_bytes_to_use=worker_kv_bytes_per_block * 8, + worker_kv_bytes_per_block=worker_kv_bytes_per_block, + world_size=world_size, + replicated_layout=True, ) - spec = _create_spec(config, _make_kv_cache_config()) - - assert spec.tokens_per_block == (16,) - assert spec.tokens_per_hash == 16 - assert spec.blocks_per_chunk == 4 + assert isinstance(spec, TieringOffloadingSpec) + assert spec.replicated_layout is True + assert spec.cpu_page_size_per_worker == worker_kv_bytes_per_block + assert spec.kv_bytes_per_chunk == worker_kv_bytes_per_block + assert spec.num_blocks == 8 -def test_offloading_config_preserves_data_parallel_index(): - config = _make_layout_vllm_config() - config.parallel_config.data_parallel_index = 2 +def test_tiering_spec_create_worker_uses_single_slot_for_replicated_layout(monkeypatch): + import vllm.v1.kv_offload.tiering.spec as tiering_spec_module - offloading_config = build_offloading_config(config, _make_kv_cache_config()) + worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT + spec = _create_spec( + spec_name="TieringOffloadingSpec", + cpu_bytes_to_use=worker_kv_bytes_per_block * 8, + worker_kv_bytes_per_block=worker_kv_bytes_per_block, + world_size=4, + replicated_layout=True, + ) + assert isinstance(spec, TieringOffloadingSpec) - assert offloading_config.parallel.data_parallel_index == 2 + region = MagicMock() + region_calls: list[dict[str, Any]] = [] + worker_calls: list[dict[str, Any]] = [] + + def fake_region_ctor(**kwargs): + region_calls.append(kwargs) + return region + + def fake_worker_ctor(**kwargs): + worker_calls.append(kwargs) + return MagicMock() + + monkeypatch.setattr(tiering_spec_module, "SharedOffloadRegion", fake_region_ctor) + monkeypatch.setattr(tiering_spec_module, "CPUOffloadingWorker", fake_worker_ctor) + monkeypatch.setattr( + tiering_spec_module.torch.accelerator, "current_device_index", lambda: 5 + ) + + kv_caches = MagicMock() + spec.create_worker(kv_caches) + + assert region_calls[0]["rank"] == 0 + assert region_calls[0]["kv_bytes_per_block"] == worker_kv_bytes_per_block + assert worker_calls[0]["kv_caches"] is kv_caches + assert worker_calls[0]["mmap_region"] is region -def test_offloading_spec_resolves_heterogeneous_hybrid_block_sizes(): - config = _make_layout_vllm_config(cpu_bytes_to_use=65536) - config.cache_config.block_size = 4 +def test_tiering_spec_create_worker_folds_device_index_for_sharded_layout(monkeypatch): + import vllm.v1.kv_offload.tiering.spec as tiering_spec_module - spec = _create_spec(config, _make_hybrid_kv_cache_config()) + spec = _create_spec( + spec_name="TieringOffloadingSpec", + worker_kv_bytes_per_block=4096, + world_size=4, + ) + assert isinstance(spec, TieringOffloadingSpec) + + region_calls: list[dict[str, Any]] = [] + + def fake_region_ctor(**kwargs): + region_calls.append(kwargs) + return MagicMock() + + monkeypatch.setattr(tiering_spec_module, "SharedOffloadRegion", fake_region_ctor) + monkeypatch.setattr(tiering_spec_module, "CPUOffloadingWorker", MagicMock()) + monkeypatch.setattr( + tiering_spec_module.torch.accelerator, + "current_device_index", + lambda: 5, + ) + + spec.create_worker(MagicMock()) + + assert region_calls[0]["rank"] == 1 + + +@pytest.mark.parametrize("world_size", [2, 4, 8]) +def test_cpu_spec_replicated_config_preserves_per_rank_sizing(world_size: int): + worker_kv_bytes_per_block = 4096 + spec = _create_spec( + cpu_bytes_to_use=worker_kv_bytes_per_block * world_size * 2, + worker_kv_bytes_per_block=worker_kv_bytes_per_block, + world_size=world_size, + replicated_layout=True, + ) + + assert isinstance(spec, CPUOffloadingSpec) + assert spec.replicated_layout is False + assert spec.cpu_page_size_per_worker == worker_kv_bytes_per_block + assert spec.kv_bytes_per_chunk == worker_kv_bytes_per_block * world_size + assert spec.num_blocks == 2 + + +def test_offloading_spec_has_replicated_layout_default(): + spec = SingleArgExternalOffloadingSpec(_make_offloading_config()) + assert spec.replicated_layout is False + + +def test_offloading_spec_uses_normalized_chunk_geometry(): + groups = ( + OffloadingGroupConfig(12, ("full_layer",)), + OffloadingGroupConfig(16, ("mla_layer",)), + ) + spec = _create_spec( + groups=groups, + tokens_per_hash=4, + blocks_per_chunk=2, + ) assert spec.tokens_per_block == (12, 16) assert spec.tokens_per_hash == 4 - assert spec.blocks_per_chunk == 1 + assert spec.blocks_per_chunk == 2 -def _full_attention_spec(block_size: int = 16) -> FullAttentionSpec: - return FullAttentionSpec( - block_size=block_size, num_kv_heads=4, head_size=128, dtype=torch.float32 - ) - - -def _parallelism_agnostic(kv_cache_groups: list[KVCacheGroupSpec]) -> bool: - config = _make_layout_vllm_config() - kv_cache_config = KVCacheConfig( - num_blocks=0, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups - ) - offloading_config = build_offloading_config(config, kv_cache_config) - return offloading_config.parallel.is_parallelism_agnostic - - -def test_parallelism_agnostic_for_single_full_attention_group(): - assert _parallelism_agnostic([KVCacheGroupSpec(["l0"], _full_attention_spec())]) - - -@pytest.mark.parametrize( - "kv_cache_groups", - [ - # MLA latent KV is replicated per rank, never head-sharded. - [ - KVCacheGroupSpec( - ["l0"], - MLAAttentionSpec( - block_size=16, num_kv_heads=1, head_size=576, dtype=torch.float32 - ), - ) - ], - # Sliding window is not full attention. - [ - KVCacheGroupSpec( - ["l0"], - SlidingWindowSpec( - block_size=16, - num_kv_heads=4, - head_size=128, - dtype=torch.float32, - sliding_window=128, - ), - ) - ], - # Hybrid model: more than one KV cache group. - [ - KVCacheGroupSpec(["l0"], _full_attention_spec()), - KVCacheGroupSpec(["l1"], _full_attention_spec()), - ], - ], -) -def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]): - assert not _parallelism_agnostic(kv_cache_groups) - - -def test_parallelism_agnostic_disabled_on_v2_model_runner(): - config = _make_layout_vllm_config() - config.use_v2_model_runner = True - kv_cache_config = KVCacheConfig( - num_blocks=0, - kv_cache_tensors=[], - kv_cache_groups=[KVCacheGroupSpec(["l0"], _full_attention_spec())], - ) - offloading_config = build_offloading_config(config, kv_cache_config) - assert not offloading_config.parallel.is_parallelism_agnostic - - -def test_create_dynamic_spec_receives_translated_config(): - config = _make_layout_vllm_config( +def test_create_dynamic_spec_receives_config(): + config = _make_offloading_config( spec_name="SingleArgExternalOffloadingSpec", - extra_config={ - "spec_module_path": "tests.v1.kv_offload.test_factory", - }, + extra_config={"spec_module_path": "tests.v1.kv_offload.test_factory"}, ) - kv_cache_config = _make_kv_cache_config() - offloading_config = build_offloading_config(config, kv_cache_config) - spec = OffloadingSpecFactory.create_spec(offloading_config) + spec = OffloadingSpecFactory.create_spec(config) assert isinstance(spec, SingleArgExternalOffloadingSpec) - assert spec.config is offloading_config - - -# --------------------------------------------------------------------------- -# Dynamic import via spec_module_path -# --------------------------------------------------------------------------- + assert spec.config is config def test_dynamic_load_via_spec_module_path(): - """External spec loaded via spec_module_path. - - This is how external projects (e.g., llm-d-kv-cache SharedStorageOffloadingSpec) - integrate with vLLM without being pre-registered in the factory. - The fallback path: registry miss → spec_module_path → importlib.import_module. - """ - config = _make_vllm_config(spec_name="CPUOffloadingSpec") - # Delete from registry to force the dynamic import path del OffloadingSpecFactory._registry["CPUOffloadingSpec"] - # spec_name not in registry → falls through to spec_module_path - config.kv_transfer_config.kv_connector_extra_config["spec_module_path"] = ( - "vllm.v1.kv_offload.cpu.spec" + config = _make_offloading_config( + extra_config={"spec_module_path": "vllm.v1.kv_offload.cpu.spec"} ) - spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config)) + + spec_cls = OffloadingSpecFactory.get_spec_cls(config.extra_config) + assert spec_cls is CPUOffloadingSpec -# --------------------------------------------------------------------------- -# Error paths -# --------------------------------------------------------------------------- - - def test_unregistered_spec_without_module_path_raises(): - """spec_name not in registry + no spec_module_path → ValueError.""" - config = _make_vllm_config(spec_name="NonexistentSpec") + config = _make_offloading_config(spec_name="NonexistentSpec") with pytest.raises(ValueError, match="Unsupported spec type"): - OffloadingSpecFactory.get_spec_cls(_get_extra_config(config)) + OffloadingSpecFactory.get_spec_cls(config.extra_config) - # create_spec should also fail (calls get_spec_cls internally) - kv_cache_config = _make_kv_cache_config() with pytest.raises(ValueError, match="Unsupported spec type"): - _create_spec(config, kv_cache_config) + OffloadingSpecFactory.create_spec(config) def test_cpu_spec_missing_cpu_bytes_to_use_raises(): - """CPUOffloadingSpec requires cpu_bytes_to_use → Exception.""" - config = _make_vllm_config(cpu_bytes_to_use=None) - config.kv_transfer_config.kv_connector_extra_config.pop("cpu_bytes_to_use", None) - kv_cache_config = _make_kv_cache_config() with pytest.raises(Exception, match="cpu_bytes_to_use must be specified"): - _create_spec(config, kv_cache_config) + _create_spec(cpu_bytes_to_use=None) def test_duplicate_registration_raises(): - """register_spec with existing name → ValueError.""" with pytest.raises(ValueError, match="is already registered"): OffloadingSpecFactory.register_spec( "CPUOffloadingSpec", "some.module", "SomeClass" ) -# --------------------------------------------------------------------------- -# Downstream collaboration — build_metric_definitions -# --------------------------------------------------------------------------- - - def test_build_metric_definitions_below_threshold(): - """store_threshold < 2 keeps stores_skipped disabled.""" from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics - config = _make_vllm_config(store_threshold=1) - spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config)) - metrics = spec_cls.build_metric_definitions( - config.kv_transfer_config.kv_connector_extra_config - ) + extra_config = {"store_threshold": 1} + spec_cls = OffloadingSpecFactory.get_spec_cls({"spec_name": "CPUOffloadingSpec"}) + metrics = spec_cls.build_metric_definitions(extra_config) + assert CPUOffloadingMetrics.STORES_SKIPPED not in metrics assert CPUOffloadingMetrics.CPU_ALLOCATION_SIZE in metrics def test_build_metric_definitions_allocation_size_histogram(): - """CPU allocation size is always reported as a histogram.""" from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics - config = _make_vllm_config(store_threshold=0) - spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config)) - metrics = spec_cls.build_metric_definitions( - config.kv_transfer_config.kv_connector_extra_config - ) + spec_cls = OffloadingSpecFactory.get_spec_cls({"spec_name": "CPUOffloadingSpec"}) + metrics = spec_cls.build_metric_definitions({}) metadata = metrics[CPUOffloadingMetrics.CPU_ALLOCATION_SIZE] + assert isinstance(metadata, OffloadingHistogramMetadata) assert metadata.buckets == ( 1, @@ -603,49 +381,10 @@ def test_build_metric_definitions_allocation_size_histogram(): def test_build_metric_definitions_returns_counter_at_threshold(): - """store_threshold >= 2 → returns stores_skipped counter definition.""" from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics - config = _make_vllm_config(store_threshold=2) - spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config)) - metrics = spec_cls.build_metric_definitions( - config.kv_transfer_config.kv_connector_extra_config - ) + extra_config = {"store_threshold": 2} + spec_cls = OffloadingSpecFactory.get_spec_cls({"spec_name": "CPUOffloadingSpec"}) + metrics = spec_cls.build_metric_definitions(extra_config) + assert CPUOffloadingMetrics.STORES_SKIPPED in metrics - - -def test_offloading_spec_accepts_blocks_per_chunk_for_heterogeneous_groups(): - config = _make_layout_vllm_config( - cpu_bytes_to_use=65536, - extra_config={"blocks_per_chunk": 2}, - ) - - spec = _create_spec(config, _make_hybrid_kv_cache_config()) - - assert spec.tokens_per_block == (12, 16) - assert spec.blocks_per_chunk == 2 - - -def test_block_size_and_blocks_per_chunk_are_mutually_exclusive(): - config = _make_layout_vllm_config( - cpu_bytes_to_use=65536, - extra_config={ - "block_size": 64, - "blocks_per_chunk": 2, - }, - ) - - with pytest.raises(ValueError, match="Specify only one"): - _create_spec(config, _make_kv_cache_config()) - - -def test_blocks_per_chunk_must_be_positive(): - config = _make_layout_vllm_config( - cpu_bytes_to_use=65536, - extra_config={ - "blocks_per_chunk": 0, - }, - ) - - with pytest.raises(ValueError, match="greater than 0"): - _create_spec(config, _make_kv_cache_config()) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index b807b480f02..2a5659c2a9b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -102,19 +102,39 @@ def build_offloading_config( ) worker_kv_bytes_per_block = total_gpu_kv_bytes // kv_cache_config.num_blocks - # Only a single non-MLA full-attention group is parallelism-invariant: - # MLA latent KV is replicated per rank (never head-sharded), and the V2 - # model runner's KV layout is not known to be parallelism-invariant. - single_group = ( + single_group_spec = ( kv_cache_config.kv_cache_groups[0].kv_cache_spec if len(kv_cache_config.kv_cache_groups) == 1 else None ) + replicated_layout = ( + vllm_config.model_config.use_mla + # Exact type: fail closed on wrappers and sliding-window variants. + and type(single_group_spec) is MLAAttentionSpec + # Page accounting: one MLA page per layer, no packed/mixed rows. + and worker_kv_bytes_per_block > 0 + and worker_kv_bytes_per_block + == single_group_spec.page_size_bytes + * len(kv_cache_config.kv_cache_groups[0].layer_names) + # Safe MVP boundary: TP-only, no other parallel axes. + and parallel_config.tensor_parallel_size > 1 + and parallel_config.pipeline_parallel_size == 1 + and parallel_config.prefill_context_parallel_size == 1 + and parallel_config.decode_context_parallel_size == 1 + and parallel_config.world_size == parallel_config.tensor_parallel_size + # Shared /dev/shm mmap layout is single-node mp only. + and parallel_config.distributed_executor_backend == "mp" + and parallel_config.nnodes_within_dp == 1 + ) + + # Only a single non-MLA full-attention group is parallelism-invariant: + # MLA latent KV is replicated per rank (never head-sharded), and the V2 + # model runner's KV layout is not known to be parallelism-invariant. is_parallelism_agnostic = ( not vllm_config.use_v2_model_runner - and single_group is not None - and isinstance(single_group, FullAttentionSpec) - and not isinstance(single_group, MLAAttentionSpec) + and single_group_spec is not None + and isinstance(single_group_spec, FullAttentionSpec) + and not isinstance(single_group_spec, MLAAttentionSpec) ) kv_events_config = vllm_config.kv_events_config @@ -144,4 +164,5 @@ def build_offloading_config( data_parallel_index=parallel_config.data_parallel_index, is_parallelism_agnostic=is_parallelism_agnostic, ), + replicated_layout=replicated_layout, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index a61ec917543..bbdcb8d0d69 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -45,6 +45,10 @@ class OffloadingConnectorWorker: self.spec = spec self.kv_cache_config = kv_cache_config self.worker: OffloadingWorker | None = None + # Non-writers still ack: pending_count waits for world_size per job. + self._is_store_writer = ( + not self.spec.replicated_layout or self.spec.config.parallel.rank == 0 + ) # job_id -> req_id for in-flight loads. self._load_jobs: dict[int, ReqId] = {} @@ -274,6 +278,9 @@ class OffloadingConnectorWorker: for job_id in kv_connector_metadata.jobs_to_flush: entry = kv_connector_metadata.store_jobs.pop(job_id, None) if entry is not None: + if not self._is_store_writer: + self._connector_worker_meta.mark_completed(job_id) + continue assert isinstance(entry.src_spec, GPULoadStoreSpec) self._unsubmitted_store_jobs.append( (job_id, entry.src_spec, entry.dst_spec) @@ -304,6 +311,10 @@ class OffloadingConnectorWorker: def prepare_store_kv(self, metadata: OffloadingConnectorMetadata): for job_id, entry in metadata.store_jobs.items(): + if not self._is_store_writer: + # Gate before queueing: no _unsubmitted_store_jobs entry. + self._connector_worker_meta.mark_completed(job_id) + continue # NOTE(orozery): defer the store to the beginning of the next # engine step, so that offloading starts AFTER transfers related # to token sampling, thereby avoiding delays to token generation. diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 60c0aa4374b..51f5162a990 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -496,6 +496,7 @@ class OffloadingSpec(ABC): ) self.config = config self.extra_config = config.extra_config + self.replicated_layout: bool = False self.kv_events_config = OffloadingKVEventsConfig( enable_kv_cache_events=config.enable_kv_cache_events, self_describing_kv_events=bool( diff --git a/vllm/v1/kv_offload/config.py b/vllm/v1/kv_offload/config.py index cd7b3ee2075..46ca489cc2b 100644 --- a/vllm/v1/kv_offload/config.py +++ b/vllm/v1/kv_offload/config.py @@ -68,3 +68,9 @@ class OffloadingConfig: model: OffloadingModelConfig cache: OffloadingCacheConfig parallel: OffloadingParallelConfig + # True when the offloaded bytes of every worker are expected to be + # byte-identical per block (pure-MLA model, single-node TP-only + # parallelism), enabling a single-copy host layout in backends that + # support it. Aggregate layout decision; per-layer replication metadata + # is planned for CanonicalKVCacheRef (#48408). + replicated_layout: bool = False diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index f6d1c29aff9..d755bfecbc4 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -24,6 +24,7 @@ from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager class CPUOffloadingSpec(OffloadingSpec): BLOCK_SIZE_ALIGNMENT = 1 + SUPPORTS_REPLICATED_LAYOUT = False @classmethod def build_metric_definitions( @@ -85,12 +86,16 @@ class CPUOffloadingSpec(OffloadingSpec): self.num_blocks = 0 self.kv_bytes_per_chunk = 0 self.cpu_page_size_per_worker = 0 + self.replicated_layout = ( + config.replicated_layout and self.SUPPORTS_REPLICATED_LAYOUT + ) if config.worker_kv_bytes_per_block > 0 and world_size > 0: - kv_bytes_per_block = config.worker_kv_bytes_per_block * world_size + num_copies = 1 if self.replicated_layout else world_size + kv_bytes_per_block = config.worker_kv_bytes_per_block * num_copies kv_bytes_per_chunk = kv_bytes_per_block * self.blocks_per_chunk # calculate cpu_page_size_per_worker - self.cpu_page_size_per_worker = kv_bytes_per_chunk // world_size + self.cpu_page_size_per_worker = kv_bytes_per_chunk // num_copies # calculate num_blocks aligned_kv_bytes_per_chunk = round_up( @@ -102,6 +107,7 @@ class CPUOffloadingSpec(OffloadingSpec): # kv_bytes_per_chunk. Note that this might contain # some padding. i.e. each offloaded block is of the form, # |--- W0-B0---|---- W1-B0---| ... |---- Wn-B0---| *** maybe-pad *** | + # or |--- B0 (single copy) ---| *** maybe-pad *** | self.kv_bytes_per_chunk = aligned_kv_bytes_per_chunk # scheduler-side diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 7aae3a2eb75..0d200a42390 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -71,6 +71,7 @@ class TieringOffloadingSpec(CPUOffloadingSpec): """ BLOCK_SIZE_ALIGNMENT = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT + SUPPORTS_REPLICATED_LAYOUT = True @classmethod @override @@ -229,10 +230,13 @@ class TieringOffloadingSpec(CPUOffloadingSpec): @override def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker: - # Fold the global physical device index into the replica-local - # [0, world_size) slot range. world_size = self.config.parallel.world_size - rank = torch.accelerator.current_device_index() % world_size + if self.replicated_layout: + rank = 0 + else: + # Fold the global physical device index into the replica-local + # [0, world_size) slot range. + rank = torch.accelerator.current_device_index() % world_size worker_mmap = SharedOffloadRegion( engine_id=self._engine_id, num_blocks=self.num_blocks, From 2e860de498a229018d8795afa34417f43b813d94 Mon Sep 17 00:00:00 2001 From: Nils Matteson <nilsmatteson@icloud.com> Date: Sat, 25 Jul 2026 23:24:01 -0600 Subject: [PATCH 058/185] [Doc] Add compile cache volume example to the Docker deployment page (#49782) Signed-off-by: Nils Matteson <nilsmatteson@icloud.com> --- docs/deployment/docker.md | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/docs/deployment/docker.md b/docs/deployment/docker.md index a8debf2cdb3..b5f3a8b049b 100644 --- a/docs/deployment/docker.md +++ b/docs/deployment/docker.md @@ -8,6 +8,26 @@ toc_depth: 2 --8<-- "docs/getting_started/installation/gpu.md:pre-built-images" +## Persist the compile cache across containers + +Mounting the Hugging Face cache keeps model weights across containers, but each +new container still starts with an empty `VLLM_CACHE_ROOT` (default +`~/.cache/vllm`) and recompiles the model's `torch.compile` artifacts. Mount a +named volume at that path to reuse the inductor, Triton, and AOT artifacts from +the second container onward: + +```bash +docker run --rm --gpus all \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + -v vllm-cache:/root/.cache/vllm \ + -p 8000:8000 \ + vllm/vllm-openai:latest \ + meta-llama/Llama-3.1-8B-Instruct +``` + +See [Faster Startup](../configuration/optimization.md#faster-startup) for the +mechanism and for what invalidates the cache. + ## Run as a non-root user The CUDA `vllm/vllm-openai` image runs as root by default for backward From 30b0714031ae24f77832099cf120576dae28b17c Mon Sep 17 00:00:00 2001 From: RED <64897097+LiuLi1998@users.noreply.github.com> Date: Sun, 26 Jul 2026 13:53:06 +0800 Subject: [PATCH 059/185] [Perf] DeepSeek-OCR-2 TTFT Optimize (#49531) Signed-off-by: RED <outofthewoods@qq.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com> --- vllm/model_executor/models/deepencoder2.py | 68 ++++++++++------------ 1 file changed, 32 insertions(+), 36 deletions(-) diff --git a/vllm/model_executor/models/deepencoder2.py b/vllm/model_executor/models/deepencoder2.py index fdec155d534..9e251c39b29 100644 --- a/vllm/model_executor/models/deepencoder2.py +++ b/vllm/model_executor/models/deepencoder2.py @@ -10,6 +10,8 @@ # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. +from functools import lru_cache + import torch import torch.nn as nn import transformers @@ -90,8 +92,6 @@ class CustomQwen2Decoder(PluggableLayer): return_dict=None, cache_position=None, ): - # token_type_ids - self._current_token_type_ids = token_type_ids causal_mask_mapping = { "full_attention": self._update_causal_mask( attention_mask, @@ -131,15 +131,12 @@ class CustomQwen2Decoder(PluggableLayer): input_tensor.shape[1], ) - token_type_ids = self._current_token_type_ids - # attention mask causal_mask = self._create_custom_4d_mask( sequence_length=sequence_length, dtype=dtype, device=device, batch_size=batch_size, - token_type_ids=token_type_ids, ) # padding mask @@ -150,44 +147,43 @@ class CustomQwen2Decoder(PluggableLayer): return causal_mask + @classmethod + @lru_cache(maxsize=8) + def compute_mask_base(cls, sequence_length, dtype, device): + # token_type_ids is the fixed pattern [0]*n_query + [1]*n_query, + # identical across the batch, so the mask depends only on + # sequence_length: img tokens (first half) attend to + # everything, txt tokens (second half) attend causally among + # themselves. lru_cache keeps one batch-invariant [1, 1, S, S] + # mask per (S, dtype, device). + min_dtype = torch.finfo(dtype).min + n_query = sequence_length // 2 + img = torch.arange(sequence_length, device=device) < n_query + txt = ~img + causal = torch.tril( + torch.ones( + sequence_length, + sequence_length, + dtype=torch.bool, + device=device, + ) + ) + allow = img[None, :] | (txt[:, None] & txt[None, :] & causal) + return torch.where( + allow, + torch.zeros((), dtype=dtype, device=device), + torch.full((), min_dtype, dtype=dtype, device=device), + )[None, None] + def _create_custom_4d_mask( self, sequence_length, dtype, device, batch_size, - token_type_ids, ): - min_dtype = torch.finfo(dtype).min - - masks = [] - for b in range(batch_size): - mask = torch.full( - (sequence_length, sequence_length), - fill_value=min_dtype, - dtype=dtype, - device=device, - ) - - type_ids = token_type_ids[b] - - image_positions = (type_ids == 0).nonzero(as_tuple=True)[0] - text_positions = (type_ids == 1).nonzero(as_tuple=True)[0] - - # non-casual - if len(image_positions) > 0: - mask[image_positions[:, None], image_positions] = 0.0 - - # causal - for i, text_pos in enumerate(text_positions): - if len(image_positions) > 0: - mask[text_pos, image_positions] = 0.0 - mask[text_pos, text_positions[: i + 1]] = 0.0 - - masks.append(mask) - - mask = torch.stack(masks, dim=0).unsqueeze(1) - return mask + base = self.compute_mask_base(sequence_length, dtype, device) + return base.expand(batch_size, -1, -1, -1) return CustomQwen2ModelInner(config) From 0164022c907014a27e81c514cd4dff6091904c2b Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim <taneem.ibrahim@gmail.com> Date: Sun, 26 Jul 2026 02:24:07 -0400 Subject: [PATCH 060/185] [CI] Fix speech correctness check rejecting improved WER (#49853) Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> --- .../correctness/test_transcription_api_correctness.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py b/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py index 62b1d367841..713fa48eee4 100644 --- a/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py +++ b/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py @@ -356,7 +356,12 @@ def test_wer_correctness( print(f"Expected WER: {expected_wer}, Actual WER: {wer}") if expected_wer: - torch.testing.assert_close(wer, expected_wer, atol=1e-1, rtol=1e-2) + wer_atol, wer_rtol = 1e-1, 1e-2 + max_wer = expected_wer + wer_atol + wer_rtol * abs(expected_wer) + assert wer <= max_wer, ( + f"WER {wer:.6f} exceeds maximum allowed {max_wer:.6f} " + f"(baseline {expected_wer:.6f})" + ) # 14-22mins of 6 audio samples of total ~115 mins and just 37MB. From 8d28b48d01b2ba56e962c7c57b894c6b4fcf8a35 Mon Sep 17 00:00:00 2001 From: Guan-Ming Chiu <105915352+guan404ming@users.noreply.github.com> Date: Sun, 26 Jul 2026 16:04:17 +0800 Subject: [PATCH 061/185] [Perf] Isolate MM preprocessing on its own executor (#49524) Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com> --- tests/test_config.py | 27 +++++++++++++++++++-------- vllm/config/model.py | 7 ++++--- vllm/renderers/base.py | 17 ++++++++--------- 3 files changed, 31 insertions(+), 20 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index 1fc00a8f8a9..8171950cdd0 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1558,20 +1558,31 @@ def test_fault_tolerance_requires_single_api_server(): def test_renderer_num_workers_with_mm_cache(): - """Disallow renderer_num_workers > 1 when mm processor cache is enabled, - since neither cache type is thread-safe.""" + """Disallow renderer_num_workers > 1 with the mm processor cache only for + pooling models, whose preprocessing runs on the renderer workers.""" mm_model = "Qwen/Qwen2-VL-2B-Instruct" - # Should raise: multi-worker + cache enabled (default cache_gb=4) + # Should raise: pooling + multi-worker + cache enabled (default cache_gb=4) with pytest.raises(ValueError, match="renderer-num-workers"): - ModelConfig(mm_model, renderer_num_workers=4) + ModelConfig(mm_model, runner="pooling", renderer_num_workers=4) - # Should raise: multi-worker + explicit cache size + # Should raise: pooling + multi-worker + explicit cache size with pytest.raises(ValueError, match="renderer-num-workers"): - ModelConfig(mm_model, renderer_num_workers=2, mm_processor_cache_gb=1.0) + ModelConfig( + mm_model, + runner="pooling", + renderer_num_workers=2, + mm_processor_cache_gb=1.0, + ) - # Should pass: multi-worker + cache disabled - config = ModelConfig(mm_model, renderer_num_workers=4, mm_processor_cache_gb=0) + # Should pass: pooling + multi-worker + cache disabled + config = ModelConfig( + mm_model, runner="pooling", renderer_num_workers=4, mm_processor_cache_gb=0 + ) + assert config.renderer_num_workers == 4 + + # Should pass: generate models preprocess on the dedicated mm executor + config = ModelConfig(mm_model, renderer_num_workers=4) assert config.renderer_num_workers == 4 # Should pass: single worker + cache enabled (default) diff --git a/vllm/config/model.py b/vllm/config/model.py index 20efeb03298..d64bd57e87d 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -748,12 +748,13 @@ class ModelConfig: if ( self.renderer_num_workers > 1 and self.multimodal_config.mm_processor_cache_gb > 0 + and self.runner_type == "pooling" ): raise ValueError( "Cannot use --renderer-num-workers > 1 with the " - "multimodal processor cache enabled. The cache is " - "not thread-safe and does not support concurrent " - "renderer workers. Please set " + "multimodal processor cache enabled for pooling models. " + "Pooling preprocessing runs on the renderer workers, and " + "the cache is not thread-safe. Please set " "--renderer-num-workers 1 (the default), or " "disable the cache with --mm-processor-cache-gb 0." ) diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py index bb02c115cae..b278a8c13dc 100644 --- a/vllm/renderers/base.py +++ b/vllm/renderers/base.py @@ -79,16 +79,15 @@ class BaseRenderer(ABC, Generic[_T]): self.tokenizer = tokenizer - # Shared thread pool executor for blocking tokenizer and - # multimodal preprocessing operations. The multimodal processor - # receives a deep-copied tokenizer (see #36557) so it is safe to - # run tokenization and MM preprocessing concurrently. + # Thread pool executor for blocking tokenizer operations. The + # multimodal processor receives a deep-copied tokenizer (see #36557) + # so it is safe to run tokenization and MM preprocessing concurrently. pool_workers = config.model_config.renderer_num_workers self._executor = ThreadPoolExecutor(max_workers=pool_workers) - # Multimodal preprocessing is always offloaded to the thread pool - # to keep the asyncio event loop responsive under concurrent load. - self._mm_executor: Executor = self._executor + # Separate single-worker executor so tokenization never queues behind + # MM preprocessing; must stay single-worker per #38418 (P0/P1 order). + self._mm_executor: Executor = ThreadPoolExecutor(max_workers=1) # Offload tokenization to the thread pool. The sync # ``_tokenize_prompt`` already encapsulates the unified ``__call__`` @@ -103,7 +102,7 @@ class BaseRenderer(ABC, Generic[_T]): self._readonly_mm_processor: BaseMultiModalProcessor | None = None self._mm_cache_stats: MultiModalCacheStats | None = None self._clear_mm_cache_async = make_async( - self.clear_mm_cache, executor=self._executor + self.clear_mm_cache, executor=self._mm_executor ) self._process_multimodal_async = make_async( self._process_multimodal, executor=self._mm_executor @@ -282,7 +281,7 @@ class BaseRenderer(ABC, Generic[_T]): self._clear_processor_cache(self._readonly_mm_processor) async def clear_mm_cache_async(self) -> None: - """Serialize clear_mm_cache through the shared executor to avoid + """Serialize clear_mm_cache through the multimodal executor to avoid races with concurrent process_inputs on the mm_processor_cache.""" await self._clear_mm_cache_async() From 21fd9e85a04f3a97da00eceddce91d8c6966b954 Mon Sep 17 00:00:00 2001 From: Guan-Ming Chiu <105915352+guan404ming@users.noreply.github.com> Date: Sun, 26 Jul 2026 16:39:25 +0800 Subject: [PATCH 062/185] [Model] Support top_k and top_p sampling for DiffusionGemma (#45429) Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com> --- vllm/model_executor/models/diffusion_gemma.py | 20 +++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 70566871e09..1a5026d513f 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -49,6 +49,7 @@ from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.platforms import current_platform from vllm.v1.outputs import LogprobsTensors +from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p from vllm.v1.worker.gpu.attn_utils import build_attn_metadata from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor, async_copy_to_gpu from vllm.v1.worker.gpu.input_batch import InputBatch @@ -1262,16 +1263,31 @@ class DiffusionSampler: valid_canvas_len_np.astype(np.int64), device=device ) + # Per-request top_k/top_p, mirroring the AR sampler. Masked tokens + # become -inf and survive the temperature scaling in the compiled + # step, so Gumbel sampling, probs, and entropy all see the filtered + # distribution. The committed argmax (always the top-1 token) is + # unaffected; only the canvas exploration is constrained. Applied + # before canvas padding so phantom positions stay uniform. + if num_decode > 0: + top_k, top_p = self.sampling_states.get_top_k_top_p( + decode_slots.repeat_interleave(valid_canvas_len), decode_slots_np + ) + if top_k is not None or top_p is not None: + logits = apply_top_k_top_p(logits.float(), top_k, top_p) + # Pad any truncated canvas back to CL so the uniform-CL sampler math # holds. Phantom (padded) positions are zeroed → uniform logits → high # entropy (no premature convergence) and argmax 0 (stable); they are - # never committed (num_sampled == real length). + # never committed (num_sampled == real length). masked_fill (not + # multiply) so -inf entries from top_k/top_p filtering above don't + # turn phantom rows into NaN. if num_decode > 0 and valid_canvas_len_np.min() < CL: ar = torch.arange(CL, device=device) starts = valid_canvas_len.cumsum(0) - valid_canvas_len # row offset per req valid = ar.unsqueeze(0) < valid_canvas_len.unsqueeze(1) # [num_decode, CL] src = (starts.unsqueeze(1) + ar.unsqueeze(0)).clamp_max(logits.shape[0] - 1) - logits = logits[src.reshape(-1)] * valid.reshape(-1, 1).to(logits.dtype) + logits = logits[src.reshape(-1)].masked_fill_(~valid.reshape(-1, 1), 0) # Clear once: the tiled loop below only scatters its own decode slots, # so it must not re-clear earlier tiles' writes. From da3a252fd13f51c22657bfc8650936f2fbb5b6f3 Mon Sep 17 00:00:00 2001 From: liranschour <liranschour@users.noreply.github.com> Date: Sun, 26 Jul 2026 11:45:47 +0300 Subject: [PATCH 063/185] [KVOffload][P2P] Generic P2P secondary tier: peer lookup and serving via ParentManager (#48021) Signed-off-by: Liran Schour <lirans@il.ibm.com> Signed-off-by: liranschour <liranschour@users.noreply.github.com> Co-authored-by: Or Ozeri <or@ozery.com> Co-authored-by: Or Ozeri <oro@il.ibm.com> --- docs/features/kv_offloading_usage.md | 67 +- .../tiering/p2p/p2p_connector_proxy.py | 29 +- .../v1/kv_offload/tiering/p2p/test_manager.py | 307 ++++- .../kv_offload/tiering/p2p/test_sessions.py | 1049 +++++++++++++++-- vllm/v1/kv_offload/base.py | 17 +- vllm/v1/kv_offload/tiering/p2p/manager.py | 326 +++-- .../tiering/p2p/session/__init__.py | 2 + .../kv_offload/tiering/p2p/session/client.py | 378 +++++- .../tiering/p2p/session/protocol.py | 100 +- .../kv_offload/tiering/p2p/session/server.py | 653 ++++++++-- .../kv_offload/tiering/p2p/session/session.py | 137 ++- 11 files changed, 2671 insertions(+), 394 deletions(-) diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index b5c9d68e268..13dbd5299d3 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -157,7 +157,7 @@ Object keys follow the same run-configuration digest scheme as the filesystem ti The P2P tier (`type: "p2p"`) shares completed KV blocks between vLLM instances over RDMA via NIXL. Each instance binds a control socket on `host:port` and exchanges blocks directly with peers — no shared filesystem required. -PYTHONHASHSEED environment variable must be set to the same fixed value on all nodes. +The `PYTHONHASHSEED` environment variable must be set to the same fixed value (e.g. `"0"`) on all nodes so that block content hashes match across instances (see [Cross-Process Sharing](#cross-process-sharing)). This is enforced: a P2P instance started without `PYTHONHASHSEED` set fails at startup, and each peer's value is verified during the connect handshake — a peer advertising a different `PYTHONHASHSEED` is rejected. | Key | Required | Default | Notes | | --- | --- | --- | --- | @@ -176,6 +176,71 @@ Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them on - `VLLM_P2P_SIDE_CHANNEL_HOST` (default `localhost`): address the P2P control socket binds to. It is used **verbatim** as both the bind address and the identity peers dial back — there is no auto-detection (this mirrors `VLLM_NIXL_SIDE_CHANNEL_HOST`). The default binds the loopback interface only, so peers on another host cannot reach it. **For any cross-host P2P deployment you must set this explicitly to the node's routable IP** (e.g. the pod IP) before launching `vllm serve` — otherwise remote peers will fail to connect. The NIXL agent name is a separate per-process identifier, so peers sharing a `host:port` never collide. - `VLLM_P2P_SIDE_CHANNEL_PORT` (default `5710`): base port for the P2P control socket. The port actually bound is `VLLM_P2P_SIDE_CHANNEL_PORT + data_parallel_index` — one socket per DP replica, matching NIXL (for DP=1 the offset is 0). The peer's port is passed as `remote_port` in `kv_transfer_params`; the router/EPP that selects the DP rank (e.g. via the `X-data-parallel-rank` header) computes `remote_port = base + rank`. The DP-index offset separates replicas *within* one deployment; two co-located *deployments* (a prefiller and a decoder on the same host) still need distinct base ports (e.g. decoder base `5711`) to avoid a bind collision. +#### Orchestration-Layer Protocol + +The P2P tier does not decide *which* peer to pull from — that is the orchestration layer's job (the router/EPP and its scheduler). The orchestrator drives every transfer through a request's `kv_transfer_params` dict: it picks the request's role, allocates a unique transaction ID, and supplies the remote peer's address. All block lookup, hash matching, and NIXL transfer happen at the tier level below; the orchestrator only sets the correct role keys and enforces the allowed combinations. + +Every vLLM instance is a symmetric **peer**. Per request it acts as a **consumer** (pulls KV blocks from a remote peer's CPU cache instead of computing locally) or a **producer** (serves blocks from its own CPU cache to remote consumers) — or both, on the same session, for different requests. Roles are chosen per request by the keys below; there are no fixed prefiller/decoder processes. + +Three role keys are defined, each mapping to a sub-dict. All are optional; a request with none of them uses the tier only as a local CPU cache. + +Each key names the **remote counterpart** this peer transfers with (not this +peer's own role), so the name reads as "the remote ___ I transfer with". + +| Key | Set on | Value fields | Meaning | +| --- | --- | --- | --- | +| `remote_decoder` | prefill producer request | `kv_request_id` | Peer computes KV and keeps it available in CPU cache for the remote decoder to pull. | +| `remote_prefiller` | decode consumer request | `kv_request_id`, `remote_host`, `remote_port` | Peer pulls KV from the remote prefiller at the given address (classic P/D disaggregation). | +| `remote_kv_source` | P2P consumer request | `kv_request_id`, `remote_host`, `remote_port` | Peer looks up and pulls whatever blocks the remote source currently holds in CPU cache. | + +Field semantics: + +- `kv_request_id` (str): unique transaction ID allocated by the orchestrator and pushed to every peer involved in the transfer; used to correlate the lookup, fetch, and transfer-done messages. The producer is implicit — it serves whatever block hashes it currently holds in its CPU cache for that ID. +- `remote_host` (str): IP/hostname of the remote peer's control socket to query. Must be the peer's routable node IP (see [Environment Variables](#environment-variables)). +- `remote_port` (int): the peer's bound control-socket port, i.e. `base + data_parallel_index` for the selected DP rank. + +Allowed and forbidden combinations: + +- **`remote_decoder` + `remote_kv_source`** is the only legal multi-key combination: a prefill producer may *also* act as a P2P consumer for the same request — skipping prefix prefill by pulling cached blocks from a source while still keeping its own computed blocks available for a downstream decoder. +- Forbidden: `remote_prefiller` + `remote_decoder` (contradictory roles), `remote_prefiller` + `remote_kv_source` (two competing fetch sources), and all three together. + +Minimal examples (values that would appear in the request's `kv_transfer_params`): + +```python +# Prefill producer — compute and keep KV for a remote decoder to pull +kv_transfer_params = {"remote_decoder": {"kv_request_id": "<unique-transfer-id>"}} + +# Decode consumer — pull KV from a specific prefiller (classic P/D) +kv_transfer_params = { + "remote_prefiller": { + "kv_request_id": "<unique-transfer-id>", + "remote_host": "<prefiller-node-ip>", + "remote_port": 5710, + } +} + +# P2P consumer — pull whatever the source already has cached +kv_transfer_params = { + "remote_kv_source": { + "kv_request_id": "<unique-transfer-id>", + "remote_host": "<source-node-ip>", + "remote_port": 5710, + } +} +``` + +Runtime handshake for a P2P (or P/D) pull, once the orchestrator has set the keys above: + +1. Both peers already have listener threads on their control sockets (see [Environment Variables](#environment-variables)). +2. **Lookup.** The consumer's tiering manager does per-block lookups; in P2P mode the tier returns `None` and registers the key. At `on_schedule_end` the consumer sends one **`LookupMsg`** (`kv_request_id` + block hashes) to the peer, per request, per step. +3. The producer matches those hashes against its local CPU cache and replies with a **`LookupRespMsg`** carrying the hit block hashes. +4. **Resolve.** Retried lookups now return hit / miss / in-flight. The consumer calls `submit_load` for hits only, allocating CPU slots only for hits. +5. The consumer sends a **`FetchMsg`** (`kv_request_id`, block hashes, destination block indexes). +6. The producer performs the **NIXL WRITE** transfer and sends **`TransferDone`** with a success status. +7. On `get_finished`, hits are loaded into GPU as ordinary cache hits; misses are recomputed by the engine. + +In classic **P/D mode** (`remote_prefiller` set, no `remote_kv_source`), the lookup phase (steps 2–4) is skipped: the decode consumer assumes the prefiller holds all of the request's blocks, so every block `lookup()` returns an immediate hit and the consumer jumps straight to the **`FetchMsg`** in step 5. The `LookupMsg`/`LookupRespMsg` round-trip only happens in P2P mode, where the consumer does not know in advance which blocks the peer has cached. + ## Tuning Tips - `cpu_bytes_to_use`: a bigger CPU tier means fewer trips to slower secondary tiers and a higher hit rate. The value is total across all workers, not per-worker. Leave headroom for the rest of the host workload. diff --git a/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py b/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py index 2f0df899023..630fcdef440 100644 --- a/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py +++ b/tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py @@ -81,6 +81,8 @@ async def lifespan(app: FastAPI): app.state.decode_dp_iterator = itertools.cycle(range(global_args.decoder_dp_size)) mode = "decoder-first" if global_args.decoder_first else "prefiller-first" + if global_args.p2p: + mode += ",p2p" pd_host = global_args.p2p_connector_host pd_port = global_args.p2p_connector_port n_pref = len(app.state.prefill_clients) @@ -145,6 +147,14 @@ def parse_args(): help="Send decode request before prefill so decoder is already " "waiting when KV blocks arrive (decoder-first mode)", ) + p.add_argument( + "--p2p", + action="store_true", + help="P2P mode: do not inject kv_transfer_params on the prefiller; " + "on the decoder, inject kv_transfer_params with a top-level " + "'remote_kv_source' block ({kv_request_id, remote_host, remote_port}) " + "instead of the default 'remote_prefiller' block.", + ) p.add_argument( "--prefiller-dp-size", type=int, @@ -206,11 +216,14 @@ def _auth_headers(request_id: str) -> dict: async def _prefill(client_info, endpoint, req_data, request_id, dp_rank=None): """Send a prefill-only request (max_tokens=1) to the prefiller.""" data = req_data.copy() - data["kv_transfer_params"] = { - "decode": { - "kv_request_id": request_id, - }, - } + if global_args.p2p: + data.pop("kv_transfer_params", None) + else: + data["kv_transfer_params"] = { + "remote_decoder": { + "kv_request_id": request_id, + }, + } data["stream"] = False data["max_tokens"] = 1 data.pop("max_completion_tokens", None) @@ -259,8 +272,9 @@ async def _handle_completions(api: str, request: Request): # Inject the prefiller's P2PConnector address so the decoder can pull # KV blocks from it. remote_port = base + prefill_rank targets the # replica that produced the KV (base+0 == base when dp=1). + decoder_key = "remote_kv_source" if global_args.p2p else "remote_prefiller" req_data["kv_transfer_params"] = { - "prefill": { + decoder_key: { "kv_request_id": request_id, "remote_host": global_args.p2p_connector_host, "remote_port": global_args.p2p_connector_port + prefill_rank, @@ -311,9 +325,10 @@ async def _handle_completions_decoder_first(api: str, request: Request): prefill_client = _get_next(request.app, "prefill") decode_client = _get_next(request.app, "decode") + decoder_key = "remote_kv_source" if global_args.p2p else "remote_prefiller" decode_data = req_data.copy() decode_data["kv_transfer_params"] = { - "prefill": { + decoder_key: { "kv_request_id": request_id, "remote_host": global_args.p2p_connector_host, "remote_port": global_args.p2p_connector_port + prefill_rank, diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index 01fd295302a..cd6e8f382a5 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -21,9 +21,11 @@ from vllm.v1.kv_offload.tiering.p2p import manager as manager_module from vllm.v1.kv_offload.tiering.p2p.manager import ( _UNBOUND_STORE_TIMEOUT_S, P2PSecondaryTierManager, + _annotate_req_context, ) from vllm.v1.kv_offload.tiering.p2p.session import ( LoadResult, + SessionCloseResult, SessionPollResult, StoreResult, ) @@ -33,15 +35,15 @@ from vllm.v1.kv_offload.tiering.p2p.session import ( # --------------------------------------------------------------------------- -def _prefill_kv_params( +def _remote_prefiller_kv_params( remote_host: str = "10.0.0.1", remote_port: int = 8000, kv_request_id: str = "req-1", ) -> dict: - """Decoder-side kv_transfer_params: ``prefill`` sub-dict carries + """Decoder-side kv_transfer_params: ``remote_prefiller`` sub-dict carries kv_request_id + remote_host + remote_port.""" return { - "prefill": { + "remote_prefiller": { "kv_request_id": kv_request_id, "remote_host": remote_host, "remote_port": remote_port, @@ -49,14 +51,34 @@ def _prefill_kv_params( } -def _decode_kv_params(kv_request_id: str = "req-1") -> dict: - """Prefiller-side kv_transfer_params: ``decode`` sub-dict carries +def _remote_kv_source_kv_params( + remote_host: str = "10.0.0.1", + remote_port: int = 8000, + kv_request_id: str = "req-1", +) -> dict: + """Symmetric-P2P consumer kv_transfer_params: ``remote_kv_source`` sub-dict has + the same shape as ``remote_prefiller`` (kv_request_id + remote_host + port).""" + return { + "remote_kv_source": { + "kv_request_id": kv_request_id, + "remote_host": remote_host, + "remote_port": remote_port, + }, + } + + +def _remote_decoder_kv_params(kv_request_id: str = "req-1") -> dict: + """Prefiller-side kv_transfer_params: ``remote_decoder`` sub-dict carries kv_request_id only.""" - return {"decode": {"kv_request_id": kv_request_id}} + return {"remote_decoder": {"kv_request_id": kv_request_id}} def _req_context(kv_params: dict | None = None) -> ReqContext: - return ReqContext(req_id="test", kv_transfer_params=kv_params) + ctx = ReqContext(req_id="test", kv_transfer_params=kv_params) + # Mirror on_new_request: parse the P2P routing state once and cache it, + # so lookup/submit_*/on_request_finished can read it back via get_state. + _annotate_req_context(ctx) + return ctx def _job_metadata( @@ -82,38 +104,78 @@ def _make_manager() -> P2PSecondaryTierManager: """Create a manager with stubbed __init__.""" mgr = P2PSecondaryTierManager.__new__(P2PSecondaryTierManager) mgr._local_id = "127.0.0.1:7777" + mgr._hash_seed = "0" mgr._finished_jobs = [] mgr._failed_req_ids = set() mgr._sessions = {} mgr._kv_to_session = {} mgr._unbound_stores = {} + mgr._failed_serve_ctxs = [] return mgr +def _init_offloading_spec() -> SimpleNamespace: + """Minimal offloading_spec for driving the real __init__.""" + return SimpleNamespace( + config=SimpleNamespace(parallel=SimpleNamespace(data_parallel_index=0)), + blocks_per_chunk=1, + ) + + # --------------------------------------------------------------------------- -# Tests for _remote_id_from_params +# Tests for __init__ PYTHONHASHSEED assertion # --------------------------------------------------------------------------- -class TestRemoteIdFromParams: +class TestInitHashSeedAssertion: + def test_missing_pythonhashseed_raises(self, monkeypatch): + """P2P instance refuses to start when PYTHONHASHSEED is unset.""" + monkeypatch.delenv("PYTHONHASHSEED", raising=False) + with pytest.raises(ValueError, match="PYTHONHASHSEED"): + P2PSecondaryTierManager( + offloading_spec=_init_offloading_spec(), + primary_kv_view=memoryview(bytearray(16)), + ) + + def test_pythonhashseed_set_succeeds(self, monkeypatch): + """With PYTHONHASHSEED set, __init__ records it for the handshake.""" + monkeypatch.setenv("PYTHONHASHSEED", "12345") + monkeypatch.setattr(manager_module, "NixlTransport", lambda *a, **k: object()) + monkeypatch.setattr(manager_module, "ZmqTransport", lambda *a, **k: object()) + monkeypatch.setattr( + manager_module.FileMapper, + "from_offloading_spec", + lambda **k: SimpleNamespace(get_run_config=lambda: {}), + ) + mgr = P2PSecondaryTierManager( + offloading_spec=_init_offloading_spec(), + primary_kv_view=memoryview(bytearray(16)), + ) + assert mgr._hash_seed == "12345" + + +# --------------------------------------------------------------------------- +# Tests for _peer_id_from_params +# --------------------------------------------------------------------------- + + +class TestPeerIdFromParams: def test_valid_params(self): - result = P2PSecondaryTierManager._remote_id_from_params( + result = manager_module._peer_id_from_params( {"remote_host": "10.0.0.1", "remote_port": 8000} ) assert result == "10.0.0.1:8000" def test_missing_host(self): - result = P2PSecondaryTierManager._remote_id_from_params({"remote_port": 8000}) + result = manager_module._peer_id_from_params({"remote_port": 8000}) assert result is None def test_missing_port(self): - result = P2PSecondaryTierManager._remote_id_from_params( - {"remote_host": "10.0.0.1"} - ) + result = manager_module._peer_id_from_params({"remote_host": "10.0.0.1"}) assert result is None def test_empty_dict(self): - result = P2PSecondaryTierManager._remote_id_from_params({}) + result = manager_module._peer_id_from_params({}) assert result is None @@ -130,35 +192,105 @@ class TestLookup: def test_lookup_returns_miss_without_required_fields(self): mgr = _make_manager() - ctx = _req_context(kv_params={"prefill": {"remote_host": "x"}}) + ctx = _req_context(kv_params={"remote_prefiller": {"remote_host": "x"}}) assert mgr.lookup(b"key", ctx) is LookupResult.MISS def test_lookup_returns_hit_for_valid_request(self): mgr = _make_manager() - ctx = _req_context(kv_params=_prefill_kv_params()) + ctx = _req_context(kv_params=_remote_prefiller_kv_params()) assert mgr.lookup(b"key", ctx) is LookupResult.HIT def test_lookup_returns_miss_for_failed_request(self): mgr = _make_manager() mgr._failed_req_ids.add("req-1") - ctx = _req_context(kv_params=_prefill_kv_params(kv_request_id="req-1")) + ctx = _req_context(kv_params=_remote_prefiller_kv_params(kv_request_id="req-1")) assert mgr.lookup(b"key", ctx) is LookupResult.MISS def test_lookup_returns_hit_for_different_request_id(self): mgr = _make_manager() mgr._failed_req_ids.add("req-1") - ctx = _req_context(kv_params=_prefill_kv_params(kv_request_id="req-2")) + ctx = _req_context(kv_params=_remote_prefiller_kv_params(kv_request_id="req-2")) assert mgr.lookup(b"key", ctx) is LookupResult.HIT def test_lookup_returns_miss_without_prefill_key(self): - """No ``prefill`` sub-dict means the request was not routed for + """No ``remote_prefiller`` sub-dict means the request was not routed for remote prefill — local prefill should run instead, so lookup() - returns MISS even when a stale ``decode`` block is present.""" + returns MISS even when a stale ``remote_decoder`` block is present.""" mgr = _make_manager() - ctx = _req_context(kv_params=_decode_kv_params()) + ctx = _req_context(kv_params=_remote_decoder_kv_params()) assert mgr.lookup(b"key", ctx) is LookupResult.MISS +# --------------------------------------------------------------------------- +# Tests for serve_external_requests +# --------------------------------------------------------------------------- + + +class _RecordingParent: + """Minimal ParentManager stub recording on_request_finished calls.""" + + def __init__(self) -> None: + self.finished: list[str] = [] + + def on_new_request(self, ctx): + from vllm.v1.kv_offload.base import RequestOffloadingContext + + return RequestOffloadingContext() + + def lookup(self, key, ctx): + return LookupResult.MISS + + def create_store_job(self, keys, ctx): + raise AssertionError("unreachable") + + def on_request_finished(self, ctx) -> None: + self.finished.append(ctx.req_id) + + +class _RecordingSession: + """Fake P2PSession that records the parent it was served with.""" + + def __init__(self) -> None: + self.served_with: list[object] = [] + + def serve_external_requests(self, parent) -> None: + self.served_with.append(parent) + + +class TestServeExternalRequests: + def test_flushes_failed_serve_ctxs_then_serves_each_session(self): + """serve_external_requests releases the failed serves left by + reaped sessions via parent.on_request_finished (clearing the + queue), then delegates to every live session with the same parent.""" + mgr = _make_manager() + ctx = ReqContext(req_id="p2p:peer:req-1:lu1") + mgr._failed_serve_ctxs = [ctx] + sess_a = _RecordingSession() + sess_b = _RecordingSession() + mgr._sessions = {"a": sess_a, "b": sess_b} # type: ignore[assignment] + + parent = _RecordingParent() + mgr.serve_external_requests(parent) # type: ignore[arg-type] + + # Failed serve released and queue cleared. + assert parent.finished == ["p2p:peer:req-1:lu1"] + assert mgr._failed_serve_ctxs == [] + # Every live session served with the same parent handle. + assert sess_a.served_with == [parent] + assert sess_b.served_with == [parent] + + def test_no_failed_serves_still_serves_sessions(self): + mgr = _make_manager() + sess = _RecordingSession() + mgr._sessions = {"a": sess} # type: ignore[assignment] + + parent = _RecordingParent() + mgr.serve_external_requests(parent) # type: ignore[arg-type] + + assert parent.finished == [] + assert sess.served_with == [parent] + + # --------------------------------------------------------------------------- # Tests for submit_store # --------------------------------------------------------------------------- @@ -166,16 +298,16 @@ class TestLookup: class TestSubmitStore: def test_no_decode_succeeds_immediately(self): - """Without a ``decode`` block, job succeeds immediately.""" + """Without a ``remote_decoder`` block, job succeeds immediately.""" mgr = _make_manager() job = _job_metadata(job_id=1, kv_params={}) mgr.submit_store(job) assert mgr._finished_jobs == [JobResult(job_id=1, success=True)] def test_missing_kv_request_id_fails(self): - """Missing kv_request_id inside ``decode`` fails the job.""" + """Missing kv_request_id inside ``remote_decoder`` fails the job.""" mgr = _make_manager() - params: dict = {"decode": {}} + params: dict = {"remote_decoder": {}} job = _job_metadata(job_id=1, kv_params=params) mgr.submit_store(job) assert mgr._finished_jobs == [JobResult(job_id=1, success=False)] @@ -189,7 +321,7 @@ class TestSubmitStore: job_id=1, keys=[b"k1", b"k2"], block_ids=[3, 4], - kv_params=_decode_kv_params(kv_request_id="req-1"), + kv_params=_remote_decoder_kv_params(kv_request_id="req-1"), ) mgr.submit_store(job) @@ -220,7 +352,7 @@ class TestSubmitStore: job_id=7, keys=[b"k1", b"k2"], block_ids=[3, 4], - kv_params=_decode_kv_params(kv_request_id="req-1"), + kv_params=_remote_decoder_kv_params(kv_request_id="req-1"), ) mgr.submit_store(job) @@ -233,9 +365,9 @@ class TestSubmitStore: def test_extra_top_level_keys_are_ignored(self): """Producer-side kv_transfer_params should not pre-create a session even when a stale caller still passes a top-level - ``remote_host``/``remote_port`` next to ``decode``.""" + ``remote_host``/``remote_port`` next to ``remote_decoder``.""" mgr = _make_manager() - params = _decode_kv_params() + params = _remote_decoder_kv_params() params["remote_host"] = "stale" params["remote_port"] = 12345 job = _job_metadata(job_id=1, kv_params=params) @@ -262,7 +394,7 @@ class TestSubmitLoad: """Empty key list succeeds immediately.""" mgr = _make_manager() job = _job_metadata( - job_id=1, keys=[], block_ids=[], kv_params=_prefill_kv_params() + job_id=1, keys=[], block_ids=[], kv_params=_remote_prefiller_kv_params() ) mgr.submit_load(job) assert mgr._finished_jobs == [JobResult(job_id=1, success=True)] @@ -270,7 +402,7 @@ class TestSubmitLoad: def test_no_session_fails(self): """No session for peer fails and marks request failed.""" mgr = _make_manager() - job = _job_metadata(job_id=1, kv_params=_prefill_kv_params()) + job = _job_metadata(job_id=1, kv_params=_remote_prefiller_kv_params()) mgr.submit_load(job) assert mgr._finished_jobs == [JobResult(job_id=1, success=False)] assert "req-1" in mgr._failed_req_ids @@ -286,7 +418,7 @@ class TestSubmitLoad: job_id=42, keys=[b"k1", b"k2"], block_ids=[5, 6], - kv_params=_prefill_kv_params(kv_request_id="req-42"), + kv_params=_remote_prefiller_kv_params(kv_request_id="req-42"), ) mgr.submit_load(job) @@ -294,6 +426,20 @@ class TestSubmitLoad: assert mgr._finished_jobs == [] assert "req-42" not in mgr._failed_req_ids + def test_missing_consumer_flag_fails(self): + """Peer fields present but neither do_remote_prefill nor + do_p2p_fetch is set — submit_load fails the job rather than + emit a stray FetchMsg.""" + mgr = _make_manager() + params = { + "remote_host": "10.0.0.1", + "remote_port": 8000, + "kv_request_id": "req-1", + } + job = _job_metadata(job_id=1, kv_params=params) + mgr.submit_load(job) + assert mgr._finished_jobs == [JobResult(job_id=1, success=False)] + # --------------------------------------------------------------------------- # Tests for on_request_finished @@ -308,7 +454,7 @@ class TestOnRequestFinished: def test_prunes_failed_req_ids(self): mgr = self._make_with_failed() - ctx = _req_context(kv_params=_prefill_kv_params(kv_request_id="req-1")) + ctx = _req_context(kv_params=_remote_prefiller_kv_params(kv_request_id="req-1")) mgr.on_request_finished(ctx) assert "req-1" not in mgr._failed_req_ids @@ -325,14 +471,26 @@ class TestOnRequestFinished: assert "req-1" in mgr._failed_req_ids def test_decoder_side_calls_session_finish_request(self): - """Decoder-side finish (``prefill`` set) still routes via peer_id + """Decoder-side finish (``remote_prefiller`` set) still routes via peer_id because the consumer addresses the producer it loaded from. The session's finish_request cancels the client-role load.""" mgr = _make_manager() peer_id = "10.0.0.1:8000" session = _FakeSession(peer_id=peer_id) mgr._sessions[peer_id] = session - ctx = _req_context(kv_params=_prefill_kv_params(kv_request_id="req-1")) + ctx = _req_context(kv_params=_remote_prefiller_kv_params(kv_request_id="req-1")) + mgr.on_request_finished(ctx) + assert session.finishes == ["req-1"] + + def test_p2p_consumer_side_calls_session_finish_request(self): + """Symmetric-P2P consumer finish (``remote_kv_source`` set) routes via peer_id + so the session drops any pending lookups (cancel_lookups) and + cancels any inbound load.""" + mgr = _make_manager() + peer_id = "10.0.0.1:8000" + session = _FakeSession(peer_id=peer_id) + mgr._sessions[peer_id] = session + ctx = _req_context(kv_params=_remote_kv_source_kv_params(kv_request_id="req-1")) mgr.on_request_finished(ctx) assert session.finishes == ["req-1"] @@ -342,7 +500,7 @@ class TestOnRequestFinished: mgr = _make_manager() bound = _FakeSession(peer_id="some-peer:1", connected=True) mgr._kv_to_session["req-1"] = bound # type: ignore[assignment] - ctx = _req_context(kv_params=_decode_kv_params(kv_request_id="req-1")) + ctx = _req_context(kv_params=_remote_decoder_kv_params(kv_request_id="req-1")) mgr.on_request_finished(ctx) assert bound.finishes == ["req-1"] assert "req-1" not in mgr._kv_to_session @@ -360,7 +518,7 @@ class TestOnRequestFinished: _UnboundStoreBatch(job_id=10, keys=[b"k"], block_ids=[0]), _UnboundStoreBatch(job_id=11, keys=[b"k2"], block_ids=[1]), ] - ctx = _req_context(kv_params=_decode_kv_params(kv_request_id="req-1")) + ctx = _req_context(kv_params=_remote_decoder_kv_params(kv_request_id="req-1")) mgr.on_request_finished(ctx) assert "req-1" in mgr._unbound_stores assert [b.job_id for b in mgr._unbound_stores["req-1"]] == [10, 11] @@ -378,11 +536,19 @@ class _FakeServerHalf: def __init__(self) -> None: self._inflight: dict[int, object] = {} + @property + def has_inflight_transfers(self) -> bool: + return bool(self._inflight) + class _FakeClientHalf: def __init__(self) -> None: self._inbound: dict[int, object] = {} + @property + def has_active_loads(self) -> bool: + return bool(self._inbound) + class _FakeSession: """Fake bidirectional session that returns canned poll() results.""" @@ -395,8 +561,10 @@ class _FakeSession: loads: list[LoadResult] | None = None, stores: list[StoreResult] | None = None, new_fetch_ids: list[str] | None = None, - close_loads: list[tuple[int, str]] | None = None, + close_jobs: list[int] | None = None, + close_req_ids: list[str] | None = None, close_stores: list[int] | None = None, + close_failed_serves: list[ReqContext] | None = None, ) -> None: self.peer_id = peer_id self.alive = alive @@ -405,18 +573,24 @@ class _FakeSession: self._loads = loads or [] self._stores = stores or [] self._new_fetch_ids = new_fetch_ids or [] - self._close_loads = close_loads or [] + self._close_jobs = close_jobs or [] + self._close_req_ids = close_req_ids or [] self._close_stores = close_stores or [] + self._close_failed_serves = close_failed_serves or [] self.requests: list[tuple[int, str]] = [] self.stores_added: list[tuple[str, list, object, int]] = [] self.attached: list[object] = [] self.finishes: list[str] = [] # Mirror P2PSession._server._inflight (transfer_id → handle) and - # P2PSession._client._inbound for the shutdown-drain and drain_jobs - # paths. Tests populate _server._inflight when needed. + # P2PSession._client.has_active_loads for the shutdown-drain and + # drain_jobs paths. Tests populate _server._inflight when needed. self._server = _FakeServerHalf() self._client = _FakeClientHalf() + @property + def has_pending_work(self) -> bool: + return self._client.has_active_loads or self._server.has_inflight_transfers + def poll(self): result = SessionPollResult( loads=self._loads, @@ -442,7 +616,12 @@ class _FakeSession: self.finishes.append(kv_request_id) def close(self): - return self._close_loads, self._close_stores + return SessionCloseResult( + failed_jobs=self._close_jobs, + failed_req_ids=self._close_req_ids, + failed_stores=self._close_stores, + failed_serves=self._close_failed_serves, + ) class TestGetFinished: @@ -484,7 +663,8 @@ class TestGetFinished: peer_id="dead:1234", alive=False, connected=True, - close_loads=[(20, "req-load")], + close_jobs=[20], + close_req_ids=["req-load"], close_stores=[10, 11], ) mgr._sessions["dead:1234"] = dead # type: ignore[assignment] @@ -498,6 +678,29 @@ class TestGetFinished: assert "dead:1234" not in mgr._sessions assert "req-load" in mgr._failed_req_ids + def test_reap_fails_probes(self): + """A reaped session's in-flight lookups land in _failed_req_ids so + the consumer's lookup() returns MISS instead of RETRY forever.""" + + class FakeData: + def remove_remote_peer(self, pid): + pass + + mgr = self._make() + mgr._data = FakeData() # type: ignore[assignment] + dead = _FakeSession( + peer_id="dead:1234", + alive=False, + connected=True, + close_req_ids=["req-probe-1", "req-probe-2"], + ) + mgr._sessions["dead:1234"] = dead # type: ignore[assignment] + + list(mgr.get_finished_jobs()) + assert "dead:1234" not in mgr._sessions + assert "req-probe-1" in mgr._failed_req_ids + assert "req-probe-2" in mgr._failed_req_ids + def test_unbound_store_kept_within_timeout(self): """Recently-parked unbound stores stay across a poll.""" mgr = self._make() @@ -537,7 +740,7 @@ class TestGetFinished: submitted_at stamp so the unbound-store sweep can age it out.""" mgr = _make_manager() job = _job_metadata( - job_id=1, kv_params=_decode_kv_params(kv_request_id="req-1") + job_id=1, kv_params=_remote_decoder_kv_params(kv_request_id="req-1") ) before = time.monotonic() mgr.submit_store(job) @@ -881,21 +1084,21 @@ class TestBidirectionalManager: b_loads_kv = "req-BtoA-load" # B loads, A serves a_decoder_params = { - "prefill": { + "remote_prefiller": { "kv_request_id": a_loads_kv, "remote_host": "B", "remote_port": 2, }, } b_decoder_params = { - "prefill": { + "remote_prefiller": { "kv_request_id": b_loads_kv, "remote_host": "A", "remote_port": 1, }, } - a_prefiller_params = {"decode": {"kv_request_id": b_loads_kv}} - b_prefiller_params = {"decode": {"kv_request_id": a_loads_kv}} + a_prefiller_params = {"remote_decoder": {"kv_request_id": b_loads_kv}} + b_prefiller_params = {"remote_decoder": {"kv_request_id": a_loads_kv}} # 1. Both sides open client-role sessions toward the peer. mgr_a.on_new_request(_req_context(a_decoder_params)) @@ -1065,7 +1268,8 @@ class TestPollOnce: peer_id=peer_dead, alive=False, connected=True, - close_loads=[(33, "req-33")], + close_jobs=[33], + close_req_ids=["req-33"], close_stores=[44], ) mgr._sessions[peer_dead] = dead # type: ignore[assignment] @@ -1309,13 +1513,13 @@ class TestConnectionDeathMidTransfer: mgr_a, mgr_b = _build_paired_managers() a_decoder_params = { - "prefill": { + "remote_prefiller": { "kv_request_id": "req-load", "remote_host": "B", "remote_port": 2, }, } - a_prefiller_params = {"decode": {"kv_request_id": "req-store"}} + a_prefiller_params = {"remote_decoder": {"kv_request_id": "req-store"}} # Open the outbound session A->B and submit one load + one store. mgr_a.on_new_request(_req_context(a_decoder_params)) @@ -1394,6 +1598,7 @@ class TestBindHostPortDefaults: identity (``host:port``) stays decoupled from the NIXL agent name (a uuid). """ + monkeypatch.setenv("PYTHONHASHSEED", "0") monkeypatch.setattr( manager_module, "FileMapper", diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py index d1c6ad0cda7..cab4ce6705a 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py +++ b/tests/v1/kv_offload/tiering/p2p/test_sessions.py @@ -13,9 +13,18 @@ completes its own load. from __future__ import annotations import time +from collections.abc import Sequence +import numpy as np import pytest +from vllm.v1.kv_offload.base import ( + LookupResult, + OffloadKey, + ReqContext, + RequestOffloadingContext, +) +from vllm.v1.kv_offload.tiering.base import JobMetadata from vllm.v1.kv_offload.tiering.p2p.session import ( LoadResult, P2PSession, @@ -33,6 +42,8 @@ from vllm.v1.kv_offload.tiering.p2p.session.protocol import ( ConnectMsg, DisconnectMsg, FetchMsg, + LookupMsg, + LookupRespMsg, TransferDoneMsg, ) from vllm.v1.kv_offload.tiering.p2p.session.server import ( @@ -48,6 +59,11 @@ from vllm.v1.kv_offload.tiering.p2p.session.session import ( # --------------------------------------------------------------------------- +# Shared PYTHONHASHSEED used by the session under test and the fake peer's +# ConnectMsg so the handshake succeeds unless a test overrides one side. +_DEFAULT_HASH_SEED = "0" + + class FakeDataTransport: """Minimal fake DataTransport for testing sessions.""" @@ -146,12 +162,16 @@ class FakeConnection: self._inbox: list[dict] = [] self._sent: list[dict] = [] self._closed = False + # When True, send() raises to simulate a broken/dead connection. + self.fail_send = False @property def alive(self) -> bool: return not self._closed def send(self, msg: dict) -> None: + if self.fail_send: + raise ConnectionError("simulated dead connection") self._sent.append(msg) def recv(self) -> list[dict]: @@ -173,6 +193,7 @@ def _peer_connect_msg( peer_id: str = "peer:8000", block_len: int = 4096, fingerprint: str | None = None, + hash_seed: str = _DEFAULT_HASH_SEED, ) -> dict: """Build a ConnectMsg as if the peer sent it.""" msg = { @@ -182,17 +203,82 @@ def _peer_connect_msg( ConnectMsg.BASE_ADDR: 0x2000, ConnectMsg.NUM_BLOCKS: 16, ConnectMsg.BLOCK_LEN: block_len, + ConnectMsg.HASH_SEED: hash_seed, } if fingerprint is not None: msg[ConnectMsg.CONFIG_FINGERPRINT] = fingerprint return msg +class FakeParent: + """Configurable :class:`ParentManager` for server-role tests. + + ``stored`` is the dict of ready blocks (key → primary block_id). + ``pending`` and ``retry`` script the first lookup() result for those + keys; subsequent lookups behave normally (a key that promised + HIT_PENDING / RETRY can later be promoted to HIT by adding it to + ``stored`` and removing it from ``pending``/``retry``). ``calls`` + captures every parent invocation in order for assertions. + + Injected per-step via ``session.serve_external_requests(parent)`` — + not held by the session, matching how ``TieringOffloadingManager`` + hands the tier a handle valid only for that call. + """ + + def __init__( + self, + stored: dict[OffloadKey, int] | None = None, + pending: set[OffloadKey] | None = None, + retry: set[OffloadKey] | None = None, + ) -> None: + self.stored: dict[OffloadKey, int] = dict(stored or {}) + self.pending: set[OffloadKey] = set(pending or ()) + self.retry: set[OffloadKey] = set(retry or ()) + self._next_job_id: int = 1000 + self.calls: list[tuple] = [] + + def on_new_request(self, ctx: ReqContext) -> RequestOffloadingContext: + self.calls.append(("on_new_request", ctx.req_id)) + return RequestOffloadingContext() + + def lookup(self, key: OffloadKey, ctx: ReqContext) -> LookupResult: + self.calls.append(("lookup", key, ctx.req_id)) + if key in self.pending: + return LookupResult.HIT_PENDING + if key in self.retry: + return LookupResult.RETRY + if key in self.stored: + return LookupResult.HIT + return LookupResult.MISS + + def create_store_job( + self, + keys: Sequence[OffloadKey], + ctx: ReqContext, + ) -> JobMetadata: + keys_list = list(keys) + self.calls.append(("create_store_job", tuple(keys_list), ctx.req_id)) + block_ids = np.array([self.stored[k] for k in keys_list], dtype=np.int32) + job_id = self._next_job_id + self._next_job_id += 1 + return JobMetadata( + job_id=job_id, + keys=keys_list, + block_ids=block_ids, + is_promotion=False, + req_context=ctx, + ) + + def on_request_finished(self, ctx: ReqContext) -> None: + self.calls.append(("on_request_finished", ctx.req_id)) + + def _make_session( conn: FakeConnection | None = None, transport: FakeDataTransport | None = None, peer_id: str = "peer:8000", local_id: str = "local:9000", + local_hash_seed: str = _DEFAULT_HASH_SEED, ) -> tuple[P2PSession, FakeConnection, FakeDataTransport]: if conn is None: conn = FakeConnection(peer_id=peer_id) @@ -203,11 +289,17 @@ def _make_session( local_id=local_id, transport=transport, # type: ignore[arg-type] local_block_len=transport.block_len, + local_hash_seed=local_hash_seed, conn=conn, # type: ignore[arg-type] ) return session, conn, transport +def _serve(session: P2PSession, parent: FakeParent) -> None: + """Resolve enqueued inbound lookups, as the manager does each step.""" + session.serve_external_requests(parent) # type: ignore[arg-type] + + def _activate( session: P2PSession, conn: FakeConnection, peer_id: str = "peer:8000" ) -> None: @@ -217,6 +309,42 @@ def _activate( session.poll() +# --- Accessors for the server role's consolidated per-request state. +# ServerRole keeps one _ServerRequestState per kv_request_id; entries are +# garbage-collected once fully idle, so "no outbound"/"no abort" reads as +# either a missing entry or a None field. These helpers paper over that. + + +def _srv_outbound(session: P2PSession, kv_request_id: str): + """Outbound serve state for a kv_request_id, or None (idle / GC'd).""" + st = session._server._requests.get(kv_request_id) + return st.outbound if st is not None else None + + +def _srv_lookups(session: P2PSession) -> list: + """Every parked inbound _ActiveLookup across all requests.""" + return [ + lu for st in session._server._requests.values() for lu in st.lookups.values() + ] + + +def _srv_abort_started(session: P2PSession, kv_request_id: str) -> float | None: + """Pending-abort start time for a kv_request_id, or None.""" + st = session._server._requests.get(kv_request_id) + return st.abort_started_at if st is not None else None + + +def _srv_inflight_count(session: P2PSession, kv_request_id: str) -> int: + """Inflight-transfer count tracked for a kv_request_id (0 if idle).""" + st = session._server._requests.get(kv_request_id) + return len(st.inflight_tids) if st is not None else 0 + + +def _srv_total_inflight(session: P2PSession) -> int: + """Sum of per-request inflight counts across all requests.""" + return sum(len(st.inflight_tids) for st in session._server._requests.values()) + + # --------------------------------------------------------------------------- # Connect / handshake # --------------------------------------------------------------------------- @@ -300,6 +428,29 @@ class TestConnectHandshake: session.poll() assert "peer:8000" in transport._remote_peers + def test_hash_seed_mismatch_marks_dead(self): + """Mismatched PYTHONHASHSEED rejects peer and marks connection dead.""" + session, conn, transport = _make_session(local_hash_seed="0") + conn.enqueue(_peer_connect_msg(hash_seed="12345")) # mismatch + session.poll() + assert "peer:8000" not in transport._remote_peers + assert not session.alive + assert not any(m[TYPE_KEY] == ConnectAckMsg.TYPE for m in conn._sent) + + def test_hash_seed_match_succeeds(self): + """Matching PYTHONHASHSEED registers the peer and acks.""" + session, conn, transport = _make_session(local_hash_seed="12345") + conn.enqueue(_peer_connect_msg(hash_seed="12345")) + session.poll() + assert "peer:8000" in transport._remote_peers + assert session.alive + assert any(m[TYPE_KEY] == ConnectAckMsg.TYPE for m in conn._sent) + + def test_hash_seed_advertised_in_connect_msg(self): + """Session advertises its own PYTHONHASHSEED in the ConnectMsg.""" + _, conn, _ = _make_session(local_hash_seed="777") + assert conn._sent[0][ConnectMsg.HASH_SEED] == "777" + # --------------------------------------------------------------------------- # Client-role flows @@ -316,7 +467,7 @@ class TestClientFlows: lookup = conn._sent[-1] assert lookup[TYPE_KEY] == FetchMsg.TYPE assert lookup[FetchMsg.KV_REQUEST_ID] == "req-1" - assert lookup[FetchMsg.BLOCK_HASHES] == [b"k1", b"k2"] + assert lookup[FetchMsg.KEYS] == [b"k1", b"k2"] assert lookup[FetchMsg.BLOCK_INDEXES] == [0, 1] def test_transfer_done_success(self): @@ -362,13 +513,46 @@ class TestClientFlows: assert abort[TYPE_KEY] == AbortFetchMsg.TYPE assert abort[AbortFetchMsg.KV_REQUEST_ID] == "req-1" + def test_active_loads_work_list_tracks_in_flight(self): + """collect_results / has_active_loads use the _active_loads work-list, + armed when a fetch is issued and discarded exactly when its load + clears — a probe-only request never enters it, and completion empties + it while the entry may briefly linger for GC.""" + session, conn, _ = _make_session() + _activate(session, conn) + client = session._client + + # A probe-only request has no in-flight load: not in _active_loads. + session.register_lookup("req-probe", b"hp") + assert client._active_loads == set() + assert client.has_active_loads is False + + # Issuing a fetch arms the work-list. + session.request_blocks( + job_id=1, kv_request_id="req-1", keys=[b"k"], block_ids=[0] + ) + assert client._active_loads == {"req-1"} + assert client.has_active_loads is True + + # Completion clears the load and discards it from the work-list. + conn.enqueue( + { + TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.KV_REQUEST_ID: "req-1", + TransferDoneMsg.SUCCESS: True, + } + ) + session.poll() + assert client._active_loads == set() + assert client.has_active_loads is False + def test_load_timeout_sends_abort(self): session, conn, _ = _make_session() _activate(session, conn) session.request_blocks( job_id=1, kv_request_id="req-1", keys=[b"k"], block_ids=[0] ) - session._client._inbound["req-1"].submitted_at = time.monotonic() - 60.0 + session._client._requests["req-1"].load.submitted_at = time.monotonic() - 60.0 session.poll() abort = conn._sent[-1] assert abort[TYPE_KEY] == AbortFetchMsg.TYPE @@ -376,7 +560,7 @@ class TestClientFlows: def test_load_abort_ack_timeout_surfaces_failure(self): """After load timeout sends AbortFetch, if no AbortAck arrives within _ABORT_ACK_TIMEOUT_S the request is surfaced as failed and removed - from _inbound — the engine cannot wait forever on a peer that won't + from _requests — the engine cannot wait forever on a peer that won't ack. """ session, conn, _ = _make_session() @@ -385,7 +569,7 @@ class TestClientFlows: job_id=7, kv_request_id="req-7", keys=[b"k"], block_ids=[0] ) # 1) Trip the load timeout to send AbortFetch and stamp aborted_at. - session._client._inbound["req-7"].submitted_at = ( + session._client._requests["req-7"].load.submitted_at = ( time.monotonic() - _LOAD_TIMEOUT_S - 1.0 ) loads = session.poll().loads @@ -395,32 +579,32 @@ class TestClientFlows: and m[AbortFetchMsg.KV_REQUEST_ID] == "req-7" for m in conn._sent ) - assert session._client._inbound["req-7"].aborted_at is not None + assert session._client._requests["req-7"].load.aborted_at is not None # 2) Now backdate aborted_at past the abort-ack timeout. No ack ever # arrived from the peer. - session._client._inbound["req-7"].aborted_at = ( + session._client._requests["req-7"].load.aborted_at = ( time.monotonic() - _ABORT_ACK_TIMEOUT_S - 1.0 ) loads = session.poll().loads assert loads == [LoadResult(job_id=7, kv_request_id="req-7", success=False)] - assert "req-7" not in session._client._inbound + assert "req-7" not in session._client._requests def test_load_abort_ack_clears_request(self): """After load timeout sends AbortFetch, an arriving AbortAckMsg from the peer surfaces the failure cleanly and removes the request from - _inbound — covers the on_abort_ack arrival path.""" + _requests — covers the on_abort_ack arrival path.""" session, conn, _ = _make_session() _activate(session, conn) session.request_blocks( job_id=8, kv_request_id="req-8", keys=[b"k"], block_ids=[0] ) - session._client._inbound["req-8"].submitted_at = ( + session._client._requests["req-8"].load.submitted_at = ( time.monotonic() - _LOAD_TIMEOUT_S - 1.0 ) # First poll: AbortFetch goes out. session.poll() - assert session._client._inbound["req-8"].aborted_at is not None + assert session._client._requests["req-8"].load.aborted_at is not None # Peer acks the abort. conn.enqueue( @@ -431,7 +615,652 @@ class TestClientFlows: ) loads = session.poll().loads assert loads == [LoadResult(job_id=8, kv_request_id="req-8", success=False)] - assert "req-8" not in session._client._inbound + assert "req-8" not in session._client._requests + + +# --------------------------------------------------------------------------- +# Symmetric-P2P lookup flow (do_p2p_fetch) +# --------------------------------------------------------------------------- + + +class TestLookupFlow: + """Consumer-side state machine for do_p2p_fetch lookups.""" + + def test_aggregate_flush_resolve_round_trip(self): + """register_lookup → flush sends one LookupMsg → response + resolves entries → register_lookup returns the cached bool on + every call, and repeat probes never re-issue a LookupMsg.""" + session, conn, _ = _make_session() + _activate(session, conn) + + # Aggregate two keys for the same kv_request_id; both return None. + assert session.register_lookup("req-1", b"hA") is None + assert session.register_lookup("req-1", b"hB") is None + + # Flush sends one LookupMsg with both keys. + sent_before = len(conn._sent) + session.flush_pending_lookups() + new = conn._sent[sent_before:] + assert len(new) == 1 + msg = new[0] + assert msg[TYPE_KEY] == LookupMsg.TYPE + assert msg[LookupMsg.KV_REQUEST_ID] == "req-1" + assert sorted(msg[LookupMsg.KEYS]) == [b"hA", b"hB"] + + # Idempotent re-flush sends nothing — the entries are now in-flight. + sent_before = len(conn._sent) + session.flush_pending_lookups() + assert conn._sent[sent_before:] == [] + + # While in-flight, register_lookup keeps returning None. + assert session.register_lookup("req-1", b"hA") is None + + # Peer answers: hA hit, hB miss. + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-1", + LookupRespMsg.KEYS: [b"hA", b"hB"], + LookupRespMsg.HITS: [True, False], + } + ) + session.poll() + + # register_lookup returns the resolved bool. + assert session.register_lookup("req-1", b"hA") is True + assert session.register_lookup("req-1", b"hB") is False + # The entry is cached, not popped: repeat probes keep returning the + # same result and never re-queue the key, so a flush sends nothing. + assert session.register_lookup("req-1", b"hA") is True + assert session.register_lookup("req-1", b"hB") is False + sent_before = len(conn._sent) + session.flush_pending_lookups() + assert [ + m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE + ] == [] + + def test_request_blocks_clears_probe_cache(self): + """A resolved HIT probe is popped when its fetch is issued, so a + re-scheduled request re-probes instead of trusting the stale True + (the served block is unpinned and may have been evicted).""" + session, conn, _ = _make_session() + _activate(session, conn) + + # Probe hA, flush, and let the peer resolve it to a HIT. + assert session.register_lookup("req-1", b"hA") is None + session.flush_pending_lookups() + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-1", + LookupRespMsg.KEYS: [b"hA"], + LookupRespMsg.HITS: [True], + } + ) + session.poll() + assert session.register_lookup("req-1", b"hA") is True + + # Fetch consumes the probe. + session.request_blocks( + job_id=1, kv_request_id="req-1", keys=[b"hA"], block_ids=[0] + ) + assert b"hA" not in session._client._requests["req-1"].probes + + # Re-scheduled probe of the same key is treated as brand-new: it + # returns None and re-queues, so a flush emits a fresh LookupMsg. + assert session.register_lookup("req-1", b"hA") is None + sent_before = len(conn._sent) + session.flush_pending_lookups() + fresh = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE] + assert len(fresh) == 1 + assert fresh[0][LookupMsg.KEYS] == [b"hA"] + + def test_flush_uses_work_list_not_full_scan(self): + """flush drains a work-list rather than scanning every live request. + + A request with no newly-registered keys is not revisited: after a + flush the work-list is empty, an idle re-flush sends nothing, and a + subsequent register re-arms exactly the one affected id — even while + an unrelated request stays live in ``_requests``. + """ + session, conn, _ = _make_session() + _activate(session, conn) + client = session._client + + # Two requests register keys; both are queued for flush. + session.register_lookup("req-A", b"hA") + session.register_lookup("req-B", b"hB") + assert client._flush_pending == {"req-A", "req-B"} + + # Flush drains the work-list even though both requests stay live. + session.flush_pending_lookups() + assert client._flush_pending == set() + assert set(client._requests) == {"req-A", "req-B"} + + # An idle re-flush visits nothing and sends no LookupMsg. + sent_before = len(conn._sent) + session.flush_pending_lookups() + assert [ + m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE + ] == [] + + # A new register re-arms only that id. + session.register_lookup("req-A", b"hA2") + assert client._flush_pending == {"req-A"} + + def test_separate_lookup_msg_per_kv_request_id(self): + """Hashes for different kv_request_ids flush as separate LookupMsgs.""" + session, conn, _ = _make_session() + _activate(session, conn) + + session.register_lookup("req-A", b"h1") + session.register_lookup("req-B", b"h2") + session.register_lookup("req-A", b"h3") + + sent_before = len(conn._sent) + session.flush_pending_lookups() + sent = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE] + assert len(sent) == 2 + by_req = {m[LookupMsg.KV_REQUEST_ID]: m[LookupMsg.KEYS] for m in sent} + assert sorted(by_req["req-A"]) == [b"h1", b"h3"] + assert by_req["req-B"] == [b"h2"] + + def test_multiple_lookup_msgs_across_steps(self): + """A request's block set may be discovered across scheduler steps: + each step that registers new keys flushes its own LookupMsg for + the same kv_request_id, carrying only the newly-probed keys.""" + session, conn, _ = _make_session() + _activate(session, conn) + + # Step 1: probe hA, hB. + session.register_lookup("req-1", b"hA") + session.register_lookup("req-1", b"hB") + sent_before = len(conn._sent) + session.flush_pending_lookups() + first = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE] + assert len(first) == 1 + assert first[0][LookupMsg.KV_REQUEST_ID] == "req-1" + assert sorted(first[0][LookupMsg.KEYS]) == [b"hA", b"hB"] + + # Step 2: a new key is discovered for the same request. The + # in-flight keys from step 1 are not re-sent; a second LookupMsg + # goes out carrying only the newly-probed key. + assert session.register_lookup("req-1", b"hA") is None # in-flight no-op + session.register_lookup("req-1", b"hC") + sent_before = len(conn._sent) + session.flush_pending_lookups() + second = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE] + assert len(second) == 1 + assert second[0][LookupMsg.KV_REQUEST_ID] == "req-1" + assert second[0][LookupMsg.KEYS] == [b"hC"] + + def test_split_response_resolves_across_messages(self): + """Producer may answer one LookupMsg's keys across multiple + LookupRespMsgs — pairs are self-describing so each lands.""" + session, conn, _ = _make_session() + _activate(session, conn) + + session.register_lookup("req-1", b"hA") + session.register_lookup("req-1", b"hB") + session.flush_pending_lookups() + + # Two responses, each carrying one of the two keys. + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-1", + LookupRespMsg.KEYS: [b"hA"], + LookupRespMsg.HITS: [True], + } + ) + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-1", + LookupRespMsg.KEYS: [b"hB"], + LookupRespMsg.HITS: [False], + } + ) + session.poll() + + assert session.register_lookup("req-1", b"hA") is True + assert session.register_lookup("req-1", b"hB") is False + + def test_finish_request_cancels_pending_lookups(self): + """finish_request drops every pending lookup for the kv_request_id.""" + session, conn, _ = _make_session() + _activate(session, conn) + + session.register_lookup("req-1", b"hA") + session.register_lookup("req-1", b"hB") + session.register_lookup("req-2", b"hC") + session.finish_request("req-1") + + # req-1 entries gone, req-2 untouched. + assert "req-1" not in session._client._requests + assert b"hC" in session._client._requests["req-2"].probes + + def test_finish_after_flushed_lookup_sends_empty_fetch(self): + """LookupMsg flushed but no FetchMsg sent (all-miss case) → + finish_request emits an empty FetchMsg so the peer can drop its + lookup state and call parent.on_request_finished.""" + session, conn, _ = _make_session() + _activate(session, conn) + + session.register_lookup("req-1", b"hA") + session.flush_pending_lookups() + sent_before = len(conn._sent) + + session.finish_request("req-1") + + fetches = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == FetchMsg.TYPE] + assert len(fetches) == 1 + assert fetches[0][FetchMsg.KV_REQUEST_ID] == "req-1" + assert fetches[0][FetchMsg.KEYS] == [] + assert fetches[0][FetchMsg.BLOCK_INDEXES] == [] + + def test_finish_without_flushed_lookup_sends_no_fetch(self): + """No LookupMsg was ever sent → finish_request must not emit an + empty FetchMsg (the peer has no state to release).""" + session, conn, _ = _make_session() + _activate(session, conn) + + # Register but never flush. + session.register_lookup("req-1", b"hA") + sent_before = len(conn._sent) + + session.finish_request("req-1") + + fetches = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == FetchMsg.TYPE] + assert fetches == [] + + def test_finish_after_real_fetch_sends_no_second_fetch(self): + """A real FetchMsg was already sent for the id → finish_request + must not emit a second (empty) FetchMsg.""" + session, conn, _ = _make_session() + _activate(session, conn) + + session.register_lookup("req-1", b"hA") + session.flush_pending_lookups() + # Resolve the probe to a HIT before fetching, as the manager only + # loads confirmed hits (an unresolved probe yields RETRY). + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-1", + LookupRespMsg.KEYS: [b"hA"], + LookupRespMsg.HITS: [True], + } + ) + session.poll() + session.request_blocks( + job_id=1, kv_request_id="req-1", keys=[b"hA"], block_ids=[7] + ) + sent_before = len(conn._sent) + + session.finish_request("req-1") + + fetches = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == FetchMsg.TYPE] + assert fetches == [] + + def test_server_lookup_deferred_until_serve_then_all_misses(self): + """``poll()`` only enqueues an inbound LookupMsg — no response is + sent until ``serve_external_requests``. With an all-miss parent + the aggregated LookupRespMsg carries the same keys and + ``hits=[False, ...]``.""" + session, conn, _ = _make_session() + _activate(session, conn) + + sent_before = len(conn._sent) + conn.enqueue( + { + TYPE_KEY: LookupMsg.TYPE, + LookupMsg.KV_REQUEST_ID: "req-1", + LookupMsg.KEYS: [b"hX", b"hY", b"hZ"], + } + ) + session.poll() + + # Dispatch alone must not answer — the parent handle is only valid + # during serve_external_requests. + assert [ + m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupRespMsg.TYPE + ] == [] + + _serve(session, FakeParent()) + + resps = [ + m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupRespMsg.TYPE + ] + assert len(resps) == 1 + resp = resps[0] + assert resp[LookupRespMsg.KV_REQUEST_ID] == "req-1" + assert resp[LookupRespMsg.KEYS] == [b"hX", b"hY", b"hZ"] + assert resp[LookupRespMsg.HITS] == [False, False, False] + + +# --------------------------------------------------------------------------- +# Server-side handling of inbound LookupMsg (ParentManager-driven) +# +# poll() only enqueues the LookupMsg; serve_external_requests(parent) +# resolves it. Tests follow the poll() → _serve() pattern. +# --------------------------------------------------------------------------- + + +def _send_lookup(conn: FakeConnection, kv_request_id: str, keys: list[bytes]): + conn.enqueue( + { + TYPE_KEY: LookupMsg.TYPE, + LookupMsg.KV_REQUEST_ID: kv_request_id, + LookupMsg.KEYS: list(keys), + } + ) + + +def _lookup_resps(conn: FakeConnection, since: int = 0) -> list[dict]: + return [m for m in conn._sent[since:] if m[TYPE_KEY] == LookupRespMsg.TYPE] + + +class TestServerLookupHandling: + def test_immediate_hits_create_one_store_job(self): + """All-HIT batch: one create_store_job call with all keys, one + LookupRespMsg with hits=[True]*N, on_request_finished fires at the + end of serve, and `available` is populated for the eventual fetch.""" + cb = FakeParent(stored={b"hA": 1, b"hB": 2, b"hC": 3}) + session, conn, _ = _make_session() + _activate(session, conn) + + sent_before = len(conn._sent) + _send_lookup(conn, "req-1", [b"hA", b"hB", b"hC"]) + session.poll() + _serve(session, cb) + + resps = _lookup_resps(conn, sent_before) + assert len(resps) == 1 + assert resps[0][LookupRespMsg.KEYS] == [b"hA", b"hB", b"hC"] + assert resps[0][LookupRespMsg.HITS] == [True, True, True] + + kinds = [c[0] for c in cb.calls] + assert kinds.count("create_store_job") == 1 + cs = next(c for c in cb.calls if c[0] == "create_store_job") + assert cs[1] == (b"hA", b"hB", b"hC") + assert cb.calls[-1][0] == "on_request_finished" + + # Hits are pinned in outbound state for the upcoming FetchMsg match. + assert set(_srv_outbound(session, "req-1").available) == { + b"hA", + b"hB", + b"hC", + } + + def test_all_misses_no_store_job_finish_fires(self): + """All-MISS batch: no create_store_job call; one LookupRespMsg + with hits=[False]*N; on_request_finished fires at end of serve.""" + cb = FakeParent() + session, conn, _ = _make_session() + _activate(session, conn) + + sent_before = len(conn._sent) + _send_lookup(conn, "req-1", [b"hA", b"hB"]) + session.poll() + _serve(session, cb) + + resps = _lookup_resps(conn, sent_before) + assert len(resps) == 1 + assert resps[0][LookupRespMsg.HITS] == [False, False] + assert all(c[0] != "create_store_job" for c in cb.calls) + assert cb.calls[-1][0] == "on_request_finished" + + def test_mixed_hit_miss_pending_defers_response_until_aggregate(self): + """HIT/MISS resolutions do not go out on first sight when any + key is still HIT_PENDING / RETRY. The lookup parks until every + key has settled (or the deadline fires), then one + LookupRespMsg carries all keys in wire order.""" + cb = FakeParent( + stored={b"hA": 1}, + pending={b"hB"}, + retry={b"hD"}, + ) + session, conn, _ = _make_session() + _activate(session, conn) + + sent_before = len(conn._sent) + _send_lookup(conn, "req-1", [b"hA", b"hB", b"hC", b"hD"]) + session.poll() + _serve(session, cb) + + # No LookupRespMsg yet — hB and hD are still pending. + assert _lookup_resps(conn, sent_before) == [] + # HIT is still pinned immediately so the eventual FetchMsg matches. + cs_calls = [c for c in cb.calls if c[0] == "create_store_job"] + assert len(cs_calls) == 1 + assert cs_calls[0][1] == (b"hA",) + # Lookup is parked; on_request_finished not yet called. + assert all(c[0] != "on_request_finished" for c in cb.calls) + assert len(_srv_lookups(session)) == 1 + + def test_pending_resolves_then_aggregate_response_fires(self): + """A HIT_PENDING key that becomes HIT on a later poll releases + the deferred aggregate response: one LookupRespMsg carrying + both keys in wire order, and one create_store_job call per + HIT (the second HIT is pinned when it resolves, not when the + response goes out).""" + cb = FakeParent(stored={b"hA": 1}, pending={b"hB"}) + session, conn, _ = _make_session() + _activate(session, conn) + + sent_before = len(conn._sent) + _send_lookup(conn, "req-1", [b"hA", b"hB"]) + session.poll() + _serve(session, cb) + # No response yet — hB still pending. + assert _lookup_resps(conn, sent_before) == [] + + # Promote hB. + cb.pending.discard(b"hB") + cb.stored[b"hB"] = 2 + + # Drive resolver via a second serve_external_requests. + _serve(session, cb) + + resps = _lookup_resps(conn, sent_before) + assert len(resps) == 1 + assert resps[0][LookupRespMsg.KEYS] == [b"hA", b"hB"] + assert resps[0][LookupRespMsg.HITS] == [True, True] + + cs_calls = [c for c in cb.calls if c[0] == "create_store_job"] + assert len(cs_calls) == 2 + assert cs_calls[0][1] == (b"hA",) + assert cs_calls[1][1] == (b"hB",) + # on_request_finished fires once after the aggregate resolve. + assert sum(1 for c in cb.calls if c[0] == "on_request_finished") == 1 + assert b"hA" in _srv_outbound(session, "req-1").available + assert b"hB" in _srv_outbound(session, "req-1").available + + def test_pending_timeout_replies_miss_no_store_job(self): + """A HIT_PENDING key that stays pending past the batch + ``deadline`` is force-MISS and never pinned; the deferred + aggregate response fires with hits=[False].""" + cb = FakeParent(pending={b"hA"}) + session, conn, _ = _make_session() + _activate(session, conn) + + sent_before = len(conn._sent) + _send_lookup(conn, "req-1", [b"hA"]) + session.poll() + _serve(session, cb) + # Initial serve: nothing immediate, lookup parked, no LookupRespMsg. + assert _lookup_resps(conn, sent_before) == [] + + # Forge the deadline into the past to trigger the timeout branch. + lookup = _srv_lookups(session)[0] + lookup.deadline = time.monotonic() - 0.1 + + _serve(session, cb) + + resps = _lookup_resps(conn, sent_before) + assert len(resps) == 1 + assert resps[0][LookupRespMsg.KEYS] == [b"hA"] + assert resps[0][LookupRespMsg.HITS] == [False] + assert all(c[0] != "create_store_job" for c in cb.calls) + assert sum(1 for c in cb.calls if c[0] == "on_request_finished") == 1 + + def test_finish_request_called_per_lookup_msg_not_per_kv_request_id(self): + """Two LookupMsgs for the same kv_request_id get distinct ctxs + and two on_request_finished calls (one per batch).""" + cb = FakeParent(stored={b"hA": 1, b"hB": 2}) + session, conn, _ = _make_session() + _activate(session, conn) + + _send_lookup(conn, "req-1", [b"hA"]) + session.poll() + _serve(session, cb) + _send_lookup(conn, "req-1", [b"hB"]) + session.poll() + _serve(session, cb) + + finish_calls = [c for c in cb.calls if c[0] == "on_request_finished"] + assert len(finish_calls) == 2 + # Distinct synthetic req_ids + assert finish_calls[0][1] != finish_calls[1][1] + # Both namespaced under the same kv_request_id + assert ":req-1:" in finish_calls[0][1] + assert ":req-1:" in finish_calls[1][1] + + def test_close_returns_open_batch_ctxs_as_failed_serves(self): + """Tearing the session down with a parked batch returns the + synthetic ctx as a failed serve (no parent handle at teardown) so + the manager can release the TieringManager's state on its next + serve.""" + cb = FakeParent(pending={b"hA"}) + session, conn, _ = _make_session() + _activate(session, conn) + + _send_lookup(conn, "req-1", [b"hA"]) + session.poll() + _serve(session, cb) + assert len(_srv_lookups(session)) == 1 + assert all(c[0] != "on_request_finished" for c in cb.calls) + + result = session.close() + + assert len(result.failed_serves) == 1 + assert ":req-1:" in result.failed_serves[0].req_id + # close() itself must not call the parent. + assert all(c[0] != "on_request_finished" for c in cb.calls) + + def test_wire_finish_drops_pending_batches_for_kv_request_id(self): + """``ServerRole.finish(kv_request_id)`` drops every parked batch + whose kv_request_id matches and queues its ctx for the next + serve's on_request_finished.""" + cb = FakeParent(pending={b"hA", b"hB"}) + session, conn, _ = _make_session() + _activate(session, conn) + + _send_lookup(conn, "req-1", [b"hA"]) + session.poll() + _serve(session, cb) + _send_lookup(conn, "req-2", [b"hB"]) + session.poll() + _serve(session, cb) + assert len(_srv_lookups(session)) == 2 + + session._server.finish("req-1") + + # req-1 batch dropped from parked lookups; its ctx queued for release. + remaining_kv_request_ids = {b.kv_request_id for b in _srv_lookups(session)} + assert remaining_kv_request_ids == {"req-2"} + queued = session._server._finished_lookup_ctxs + assert len(queued) == 1 + assert ":req-1:" in queued[0].req_id + + # The next serve fires on_request_finished exactly once for req-1. + _serve(session, cb) + finish_calls = [c for c in cb.calls if c[0] == "on_request_finished"] + assert len(finish_calls) == 1 + assert ":req-1:" in finish_calls[0][1] + + def test_incoming_fetch_drops_pending_lookups_for_kv_request_id(self): + """A peer FetchMsg terminates the lookup phase for its id: parked + lookups with matching kv_request_id are dropped and their + ctx queued for on_request_finished; other kv_request_ids untouched.""" + cb = FakeParent(pending={b"hA", b"hB"}) + session, conn, _ = _make_session() + _activate(session, conn) + + _send_lookup(conn, "req-1", [b"hA"]) + session.poll() + _serve(session, cb) + _send_lookup(conn, "req-2", [b"hB"]) + session.poll() + _serve(session, cb) + assert len(_srv_lookups(session)) == 2 + + # Empty FetchMsg: peer signals "lookup phase done" without asking + # for any blocks (the all-miss case). + conn.enqueue( + { + TYPE_KEY: FetchMsg.TYPE, + FetchMsg.KV_REQUEST_ID: "req-1", + FetchMsg.KEYS: [], + FetchMsg.BLOCK_INDEXES: [], + } + ) + session.poll() + + remaining_kv_request_ids = {lu.kv_request_id for lu in _srv_lookups(session)} + assert remaining_kv_request_ids == {"req-2"} + # Dispatch queues the ctx but does not call the parent yet. + queued = session._server._finished_lookup_ctxs + assert len(queued) == 1 + assert ":req-1:" in queued[0].req_id + + # The next serve fires on_request_finished exactly once for req-1. + _serve(session, cb) + finish_calls = [c for c in cb.calls if c[0] == "on_request_finished"] + assert len(finish_calls) == 1 + assert ":req-1:" in finish_calls[0][1] + + def test_lookup_then_fetch_round_trip_emits_store_result(self): + """End-to-end: lookup pins primary slots → fetch matches them → + NIXL transfer completes → StoreResult surfaces with the + create_store_job's job_id (the engine releases the pin).""" + cb = FakeParent(stored={b"hA": 7, b"hB": 8}) + session, conn, transport = _make_session() + _activate(session, conn) + + _send_lookup(conn, "req-1", [b"hA", b"hB"]) + session.poll() + _serve(session, cb) + cs = next(c for c in cb.calls if c[0] == "create_store_job") + # FakeParent issues monotonic job_ids starting at 1000. + expected_job_id = 1000 + + # Consumer issues FetchMsg on the resolved hits. + conn.enqueue( + { + TYPE_KEY: FetchMsg.TYPE, + FetchMsg.KV_REQUEST_ID: "req-1", + FetchMsg.KEYS: [b"hA", b"hB"], + FetchMsg.BLOCK_INDEXES: [20, 21], + } + ) + session.poll() + + # NIXL write_blocks called with our pinned local block_ids. + assert len(transport._transfers) == 1 + _, (_peer, local, remote) = next(iter(transport._transfers.items())) + assert local == [7, 8] + assert remote == [20, 21] + + # Drive the transport completion. + transport._poll_done.append(0) + result = session.poll() + + store_results = [s for s in result.stores if s.success] + assert any(s.job_id == expected_job_id for s in store_results) + # Sanity: kv mention in synthetic ctx. + assert cs[2].startswith("p2p:") # --------------------------------------------------------------------------- @@ -449,7 +1278,7 @@ class TestServerFlows: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"], + FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], } ) @@ -467,7 +1296,7 @@ class TestServerFlows: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], } ) @@ -485,7 +1314,7 @@ class TestServerFlows: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], } ) @@ -508,11 +1337,11 @@ class TestServerFlows: session.poll() ack = next(m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE) assert ack[AbortAckMsg.KV_REQUEST_ID] == "req-1" - assert "req-1" not in session._server._pending_aborts + assert _srv_abort_started(session, "req-1") is None def test_abort_fetch_defers_ack_when_cancel_pending(self): """If cancel(mode='wait') reports still-inflight tids, the ack is - deferred and the abort is parked in _pending_aborts.""" + deferred and the abort is parked (abort_started_at set).""" session, conn, transport = _make_session() _activate(session, conn) # Seed an inflight transfer for req-1 that the transport pretends @@ -533,7 +1362,7 @@ class TestServerFlows: session.poll() assert not any(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent) - assert "req-1" in session._server._pending_aborts + assert _srv_abort_started(session, "req-1") is not None # First attempt happens inside _on_abort_fetch; the per-tick # drain runs again at the end of poll() — both are wait-mode. assert all(mode == "wait" for _, mode in transport._cancel_calls) @@ -558,7 +1387,7 @@ class TestServerFlows: } ) session.poll() - assert "req-1" in session._server._pending_aborts + assert _srv_abort_started(session, "req-1") is not None # Backend finishes draining: transport.poll() will return tid as # DONE, and the next cancel(mode='wait') call sees it's gone. @@ -569,7 +1398,7 @@ class TestServerFlows: ack = next(m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE) assert ack[AbortAckMsg.KV_REQUEST_ID] == "req-1" - assert "req-1" not in session._server._pending_aborts + assert _srv_abort_started(session, "req-1") is None assert tid not in session._server._inflight def test_abort_fetch_force_cancels_after_timeout(self): @@ -591,9 +1420,9 @@ class TestServerFlows: } ) session.poll() - assert "req-1" in session._server._pending_aborts + assert _srv_abort_started(session, "req-1") is not None # Backdate past the drain deadline. - session._server._pending_aborts["req-1"] = ( + session._server._requests["req-1"].abort_started_at = ( time.monotonic() - _CANCEL_DRAIN_TIMEOUT_S - 1.0 ) # Even if the transport still claims it can't cancel, the @@ -604,7 +1433,7 @@ class TestServerFlows: ack = next(m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE) assert ack[AbortAckMsg.KV_REQUEST_ID] == "req-1" - assert "req-1" not in session._server._pending_aborts + assert _srv_abort_started(session, "req-1") is None assert tid not in session._server._inflight assert ([tid], "immediate") in transport._cancel_calls @@ -627,7 +1456,7 @@ class TestServerFlows: } ) session.poll() - first_started_at = session._server._pending_aborts["req-1"] + first_started_at = _srv_abort_started(session, "req-1") # Second AbortFetchMsg for the same kv_request_id while still # draining must not reset the deadline. @@ -638,7 +1467,7 @@ class TestServerFlows: } ) session.poll() - assert session._server._pending_aborts["req-1"] == first_started_at + assert _srv_abort_started(session, "req-1") == first_started_at # Now let the drain succeed and confirm exactly one ack ever. transport._cancel_still_inflight.discard(tid) @@ -669,7 +1498,7 @@ class TestServerFlows: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], } ) @@ -701,7 +1530,7 @@ class TestServerFlows: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], } ) @@ -743,12 +1572,12 @@ class TestFinishRequestServerSide: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], } ) session.poll() - assert "req-1" in session._server._outbound + assert _srv_outbound(session, "req-1") is not None session.finish_request("req-1") @@ -756,7 +1585,7 @@ class TestFinishRequestServerSide: assert msg is not None assert msg[TransferDoneMsg.KV_REQUEST_ID] == "req-1" assert msg[TransferDoneMsg.SUCCESS] is False - assert "req-1" not in session._server._outbound + assert _srv_outbound(session, "req-1") is None def test_with_inflight_defers_then_fires_on_last_transfer(self): """finish_request with inflight defers; last transfer fires the @@ -768,7 +1597,7 @@ class TestFinishRequestServerSide: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"], + FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], } ) @@ -780,8 +1609,8 @@ class TestFinishRequestServerSide: before = len(conn._sent) session.finish_request("req-1") assert len(conn._sent) == before - assert "req-1" in session._server._outbound - assert session._server._outbound["req-1"].finishing + assert _srv_outbound(session, "req-1") is not None + assert _srv_outbound(session, "req-1").finishing # Last inflight settles -> early-fail fires. tid = next(iter(transport._transfers)) @@ -792,7 +1621,7 @@ class TestFinishRequestServerSide: assert msg is not None assert msg[TransferDoneMsg.KV_REQUEST_ID] == "req-1" assert msg[TransferDoneMsg.SUCCESS] is False - assert "req-1" not in session._server._outbound + assert _srv_outbound(session, "req-1") is None def test_full_demand_satisfied_still_sends_success(self): """finish_request must not override a fully-satisfied transfer: @@ -803,7 +1632,7 @@ class TestFinishRequestServerSide: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], } ) @@ -831,13 +1660,13 @@ class TestFinishRequestServerSide: session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) # finish_request first — no demand received yet -> defer. session.finish_request("req-1") - assert "req-1" in session._server._outbound + assert _srv_outbound(session, "req-1") is not None # Fetch arrives now: demand fully satisfied by available. conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], } ) @@ -858,7 +1687,7 @@ class TestFinishRequestServerSide: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-2", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"], + FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], } ) @@ -870,7 +1699,7 @@ class TestFinishRequestServerSide: session.poll() msg = next(m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE) assert msg[TransferDoneMsg.SUCCESS] is False - assert "req-2" not in session._server._outbound + assert _srv_outbound(session, "req-2") is None def test_unknown_request_is_noop(self): session, conn, _ = _make_session() @@ -891,7 +1720,7 @@ class TestFinishRequestServerSide: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"demand"], + FetchMsg.KEYS: [b"demand"], FetchMsg.BLOCK_INDEXES: [5], } ) @@ -900,7 +1729,7 @@ class TestFinishRequestServerSide: # matches demand. Without the shortcut, job 42 sits in _store_jobs # for _STORE_TIMEOUT_S. session.add_stored_blocks("req-1", [b"unrelated"], [0], job_id=42) - assert session._server._outbound["req-1"].pending_job_ids == {42} + assert _srv_outbound(session, "req-1").pending_job_ids == {42} session.finish_request("req-1") @@ -908,7 +1737,7 @@ class TestFinishRequestServerSide: msg = self._last_transfer_done(conn) assert msg is not None assert msg[TransferDoneMsg.SUCCESS] is False - assert "req-1" not in session._server._outbound + assert _srv_outbound(session, "req-1") is None # Local store job surfaces on the next poll, success=False. stores = session.poll().stores @@ -925,7 +1754,7 @@ class TestFinishRequestServerSide: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], } ) @@ -933,7 +1762,7 @@ class TestFinishRequestServerSide: session.add_stored_blocks("req-1", [b"k1"], [0], job_id=7) # finish_request races with the inflight transfer. session.finish_request("req-1") - assert "req-1" in session._server._outbound # deferred + assert _srv_outbound(session, "req-1") is not None # deferred # Last inflight completes -> _finalize_outbound(success=True) fires. tid = next(iter(transport._transfers)) @@ -944,7 +1773,7 @@ class TestFinishRequestServerSide: assert msg is not None assert msg[TransferDoneMsg.SUCCESS] is True assert StoreResult(job_id=7, success=True) in stores - assert "req-1" not in session._server._outbound + assert _srv_outbound(session, "req-1") is None def test_write_blocks_failure_finalizes_with_failure(self): """write_blocks returning None must not leave the request hanging. @@ -962,7 +1791,7 @@ class TestFinishRequestServerSide: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], } ) @@ -973,7 +1802,7 @@ class TestFinishRequestServerSide: session.add_stored_blocks("req-1", [b"k1"], [0], job_id=42) # Outbound was finalized immediately (no other inflight). - assert "req-1" not in session._server._outbound + assert _srv_outbound(session, "req-1") is None # Peer notified with success=False. msg = next(m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE) assert msg[TransferDoneMsg.KV_REQUEST_ID] == "req-1" @@ -996,14 +1825,14 @@ class TestFinishRequestServerSide: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2", b"k3"], + FetchMsg.KEYS: [b"k1", b"k2", b"k3"], FetchMsg.BLOCK_INDEXES: [10, 11, 12], } ) session.poll() # Demand registered, no matches yet. assert session._server._inflight == {} - outbound = session._server._outbound["req-1"] + outbound = _srv_outbound(session, "req-1") assert outbound.remaining == 3 assert set(outbound.demanded.keys()) == {b"k1", b"k2", b"k3"} @@ -1024,7 +1853,7 @@ class TestFinishRequestServerSide: assert session._server._inflight == {} assert outbound.remaining == 2 # Not yet finalized — still 2 blocks demanded. - assert "req-1" in session._server._outbound + assert _srv_outbound(session, "req-1") is not None # Round 2: k2 and k3 arrive together. session.add_stored_blocks("req-1", [b"k2", b"k3"], [1, 2], job_id=200) @@ -1038,7 +1867,7 @@ class TestFinishRequestServerSide: stores = session.poll().stores assert StoreResult(job_id=200, success=True) in stores # _finalize_outbound fired — request gone, peer notified with success. - assert "req-1" not in session._server._outbound + assert _srv_outbound(session, "req-1") is None done = next(m for m in conn._sent if m.get(TYPE_KEY) == TransferDoneMsg.TYPE) assert done[TransferDoneMsg.KV_REQUEST_ID] == "req-1" assert done[TransferDoneMsg.SUCCESS] is True @@ -1057,7 +1886,7 @@ class TestFinishRequestServerSide: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"], + FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], } ) @@ -1067,7 +1896,7 @@ class TestFinishRequestServerSide: session.add_stored_blocks("req-1", [b"k1"], [0], job_id=100) assert len(session._server._inflight) == 1 tid_1 = next(iter(session._server._inflight)) - outbound = session._server._outbound["req-1"] + outbound = _srv_outbound(session, "req-1") assert outbound.remaining == 2 # decrement happens on completion assert outbound.finishing is False @@ -1078,7 +1907,7 @@ class TestFinishRequestServerSide: assert list(session._server._inflight.keys()) == [tid_1] # Marked finishing, but NOT finalized yet (transfer_1 still inflight). assert outbound.finishing is True - assert "req-1" in session._server._outbound + assert _srv_outbound(session, "req-1") is not None done_msgs = [m for m in conn._sent if m.get(TYPE_KEY) == TransferDoneMsg.TYPE] assert done_msgs == [] @@ -1090,7 +1919,7 @@ class TestFinishRequestServerSide: stores_first = session.poll().stores assert StoreResult(job_id=100, success=True) in stores_first # Outbound state cleaned up; peer notified with success=False. - assert "req-1" not in session._server._outbound + assert _srv_outbound(session, "req-1") is None done = next(m for m in conn._sent if m.get(TYPE_KEY) == TransferDoneMsg.TYPE) assert done[TransferDoneMsg.KV_REQUEST_ID] == "req-1" assert done[TransferDoneMsg.SUCCESS] is False @@ -1123,7 +1952,7 @@ class TestBidirectional: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-srv", - FetchMsg.BLOCK_HASHES: [b"served"], + FetchMsg.KEYS: [b"served"], FetchMsg.BLOCK_INDEXES: [7], } ) @@ -1178,6 +2007,7 @@ class TestPendingSession: local_id="local:9000", transport=transport, # type: ignore[arg-type] local_block_len=4096, + local_hash_seed=_DEFAULT_HASH_SEED, conn=None, ) session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) @@ -1197,6 +2027,7 @@ class TestPendingSession: local_id="local:9000", transport=transport, # type: ignore[arg-type] local_block_len=4096, + local_hash_seed=_DEFAULT_HASH_SEED, conn=None, ) conn = FakeConnection(peer_id="peer:8000") @@ -1218,13 +2049,16 @@ class TestPendingSession: local_id="local:9000", transport=transport, # type: ignore[arg-type] local_block_len=4096, + local_hash_seed=_DEFAULT_HASH_SEED, conn=None, ) session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) session.add_stored_blocks("req-2", [b"k2"], [1], job_id=2) - failed_loads, failed_stores = session.close() - assert failed_loads == [] - assert set(failed_stores) == {1, 2} + result = session.close() + assert result.failed_jobs == [] + assert result.failed_req_ids == [] + assert set(result.failed_stores) == {1, 2} + assert result.failed_serves == [] # --------------------------------------------------------------------------- @@ -1246,9 +2080,50 @@ class TestDisconnect: session.request_blocks(1, "req-1", [b"k"], [0]) session.request_blocks(2, "req-2", [b"k"], [0]) session.add_stored_blocks("req-srv", [b"k"], [0], job_id=10) - failed_loads, failed_stores = session.close() - assert set(failed_loads) == {(1, "req-1"), (2, "req-2")} - assert set(failed_stores) == {10} + result = session.close() + assert set(result.failed_jobs) == {1, 2} + assert set(result.failed_req_ids) == {"req-1", "req-2"} + assert set(result.failed_stores) == {10} + assert result.failed_serves == [] + + def test_send_failure_marks_connection_dead(self): + """A raising send must mark the connection dead, not silently drop + the message — otherwise the session lingers alive, is never reaped, + and in-flight lookups/loads toward the dead peer hang forever.""" + session, conn, _ = _make_session() + _activate(session, conn) + assert session.alive + + conn.fail_send = True + # request_blocks flushes a FetchMsg synchronously via _do_send. + session.request_blocks(1, "req-1", [b"k"], [0]) + + assert not session.alive + + def test_close_surfaces_inflight_lookups(self): + """close() reports kv_request_ids whose symmetric-P2P probe is still + unresolved; resolved probes are not reported (their answer is in).""" + session, conn, _ = _make_session() + _activate(session, conn) + + session.register_lookup("req-hit", b"hA") + session.register_lookup("req-inflight", b"hB") + session.flush_pending_lookups() + + # Only req-hit is answered; req-inflight stays in flight. + conn.enqueue( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: "req-hit", + LookupRespMsg.KEYS: [b"hA"], + LookupRespMsg.HITS: [True], + } + ) + session.poll() + + result = session.close() + assert result.failed_jobs == [] + assert result.failed_req_ids == ["req-inflight"] # --------------------------------------------------------------------------- @@ -1294,7 +2169,7 @@ class TestAdversarial: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-bad", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"], + FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [1], } ) @@ -1342,7 +2217,7 @@ class TestDispatchErrorHandling: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-bad", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"], + FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [1], } ) @@ -1372,7 +2247,7 @@ class TestDispatchErrorHandling: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [0], } ) @@ -1395,7 +2270,7 @@ class TestDispatchErrorHandling: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [0], } ) @@ -1422,7 +2297,7 @@ class TestDispatchErrorHandling: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1"], + FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [0], } ) @@ -1438,10 +2313,10 @@ class TestDispatchErrorHandling: class TestInflightPerReqInvariant: - """`_inflight_per_req` is the O(1) replacement for the previous - O(N) scan in `_has_inflight_for`. These tests check that every - mutation site keeps the counter in sync with `_inflight` and that - the lookup is correct under high fan-out. + """Per-request `inflight_tids` is the O(1) replacement for the + previous O(N) scan in `_has_inflight_for`. These tests check that + every mutation site keeps the set in sync with `_inflight` and + that the lookup is correct under high fan-out. """ def test_invariant_holds_through_lifecycle(self): @@ -1452,17 +2327,14 @@ class TestInflightPerReqInvariant: _activate(session, conn) def _invariant_holds() -> bool: - counted = sum(session._server._inflight_per_req.values()) - return counted == len(session._server._inflight) and all( - v > 0 for v in session._server._inflight_per_req.values() - ) + return _srv_total_inflight(session) == len(session._server._inflight) assert _invariant_holds() # Two requests, two blocks each, all dispatched in one batch. session.add_stored_blocks("req-A", [b"a1", b"a2"], [0, 1], job_id=10) session.add_stored_blocks("req-B", [b"b1", b"b2"], [2, 3], job_id=11) - for kv_id, hashes, indexes in ( + for kv_id, keys, indexes in ( ("req-A", [b"a1", b"a2"], [100, 101]), ("req-B", [b"b1", b"b2"], [102, 103]), ): @@ -1470,7 +2342,7 @@ class TestInflightPerReqInvariant: { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: kv_id, - FetchMsg.BLOCK_HASHES: hashes, + FetchMsg.KEYS: keys, FetchMsg.BLOCK_INDEXES: indexes, } ) @@ -1493,7 +2365,7 @@ class TestInflightPerReqInvariant: assert _invariant_holds() assert not session._server._has_inflight_for("req-A") - assert "req-A" not in session._server._inflight_per_req # entry was removed + assert _srv_inflight_count(session, "req-A") == 0 # entry drained assert session._server._has_inflight_for("req-B") # Complete req-B; counter must drain to empty. @@ -1508,7 +2380,7 @@ class TestInflightPerReqInvariant: assert _invariant_holds() assert session._server._inflight == {} - assert session._server._inflight_per_req == {} + assert _srv_total_inflight(session) == 0 def test_has_inflight_for_correct_with_many_requests(self): """Populate many inflight xfers across many ids; lookup must @@ -1524,9 +2396,7 @@ class TestInflightPerReqInvariant: tid, _InflightXfer(kv_request_id=kv_id, block_count=1, job_ids={tid}), ) - assert sum(session._server._inflight_per_req.values()) == len( - session._server._inflight - ) + assert _srv_total_inflight(session) == len(session._server._inflight) assert session._server._has_inflight_for("req-0") assert session._server._has_inflight_for("req-99") assert not session._server._has_inflight_for("req-missing") @@ -1539,7 +2409,7 @@ class TestInflightPerReqInvariant: ] for tid in tids_50: session._server._inflight_pop(tid) - assert "req-50" not in session._server._inflight_per_req + assert _srv_inflight_count(session, "req-50") == 0 assert not session._server._has_inflight_for("req-50") # Other ids unaffected. assert session._server._has_inflight_for("req-49") @@ -1559,6 +2429,7 @@ class TestConnectMsgValidation: ConnectMsg.BASE_ADDR: 0x1000, ConnectMsg.NUM_BLOCKS: 8, ConnectMsg.BLOCK_LEN: 4096, + ConnectMsg.HASH_SEED: "0", } def test_valid_message_passes(self): @@ -1600,13 +2471,25 @@ class TestConnectMsgValidation: with pytest.raises(ValueError, match="block_len"): ConnectMsg.validate(msg) + def test_missing_hash_seed(self): + msg = self._valid_msg() + del msg[ConnectMsg.HASH_SEED] + with pytest.raises(ValueError, match="hash_seed"): + ConnectMsg.validate(msg) + + def test_hash_seed_wrong_type(self): + msg = self._valid_msg() + msg[ConnectMsg.HASH_SEED] = 12345 # int, not str + with pytest.raises(ValueError, match="hash_seed"): + ConnectMsg.validate(msg) + class TestFetchMsgValidation: def _valid_msg(self) -> dict: return { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: "req-1", - FetchMsg.BLOCK_HASHES: [b"k1", b"k2"], + FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [0, 1], } diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 51f5162a990..cb90599cf81 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -6,9 +6,9 @@ Core abstractions for KV cache offloading in vLLM v1. from abc import ABC, abstractmethod from collections.abc import Collection, Iterable, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import Enum, auto -from typing import TYPE_CHECKING, Any, NamedTuple, NewType +from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeVar import numpy as np import torch @@ -45,10 +45,23 @@ def get_offload_group_idx(key: OffloadKey) -> int: return int.from_bytes(key[-4:], "big", signed=False) +_T = TypeVar("_T") + + @dataclass class ReqContext: req_id: str kv_transfer_params: dict[str, Any] | None = None + # Per-request scratch space keyed by value type, so a tier can parse + # 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) + + def set_state(self, val: Any) -> None: + self._state[type(val)] = val + + def get_state(self, cls: type[_T]) -> _T | None: + return self._state.get(cls) class LookupResult(Enum): diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index 95605fd1735..f1f557ce38e 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -8,6 +8,7 @@ Owns transports and a single bidirectional P2PSession per remote peer. from __future__ import annotations +import os import time import uuid from collections.abc import Iterable, Sequence @@ -23,12 +24,12 @@ from vllm.v1.kv_offload.base import ( OffloadKey, ReqContext, RequestOffloadingContext, + ScheduleEndContext, ) from vllm.v1.kv_offload.file_mapper import FileMapper from vllm.v1.kv_offload.tiering.base import ( JobMetadata, JobResult, - ScheduleEndContext, SecondaryTierManager, ) from vllm.v1.kv_offload.tiering.p2p.control import ControlTransport, ZmqTransport @@ -37,6 +38,7 @@ from vllm.v1.kv_offload.tiering.p2p.session import P2PSession if TYPE_CHECKING: from vllm.v1.kv_offload.base import OffloadingSpec + from vllm.v1.kv_offload.tiering.base import ParentManager from vllm.v1.kv_offload.tiering.p2p.control.base import ControlConnection logger = init_logger(__name__) @@ -59,24 +61,112 @@ _SHUTDOWN_DRAIN_TIMEOUT_S = 3.0 _DRAIN_SLEEP_S = 0.001 -def _prefill_params(kv_params: dict | None) -> dict | None: - """Return the ``prefill`` sub-dict, or None if absent. +def _remote_prefiller_params(kv_params: dict | None) -> dict | None: + """Return the ``remote_prefiller`` sub-dict, or None if absent. - Set on decoder requests; carries kv_request_id, remote_host, remote_port. + Set on decoder requests to name the remote prefiller they pull from; + carries kv_request_id, remote_host, remote_port. """ if not kv_params: return None - return kv_params.get("prefill") + return kv_params.get("remote_prefiller") -def _decode_params(kv_params: dict | None) -> dict | None: - """Return the ``decode`` sub-dict, or None if absent. +def _remote_decoder_params(kv_params: dict | None) -> dict | None: + """Return the ``remote_decoder`` sub-dict, or None if absent. - Set on prefiller requests; carries kv_request_id. + Set on prefiller requests to name the remote decoder they serve; + carries kv_request_id. """ if not kv_params: return None - return kv_params.get("decode") + return kv_params.get("remote_decoder") + + +def _remote_kv_source_params(kv_params: dict | None) -> dict | None: + """Return the ``remote_kv_source`` sub-dict, or None if absent. + + Set on symmetric-P2P consumer requests to name the remote source they + pull from; carries kv_request_id, remote_host, remote_port. + """ + if not kv_params: + return None + return kv_params.get("remote_kv_source") + + +def _peer_id_from_params(role_params: dict) -> str | None: + """Build ``host:port`` peer_id from a role-scoped sub-dict, or None.""" + host = role_params.get("remote_host") + port = role_params.get("remote_port") + if host and port: + return f"{host}:{port}" + return None + + +@dataclass(slots=True) +class P2PSourceInfo: + """Consumer side: this request fetches from a remote (prefiller or peer).""" + + kv_request_id: str + peer_id: str + do_probe: bool # False for remote_prefiller (PD), True for remote_kv_source + + +@dataclass(slots=True) +class P2PDestInfo: + """Producer side: a remote fetches this request's blocks from us. + + ``kv_request_id`` is None when the ``remote_decoder`` block is present + but malformed (no id); the block's presence still marks the request as + remote-decode, so submit_store must fail rather than store locally. + """ + + kv_request_id: str | None + + +def _parse_source(kv_params: dict | None) -> P2PSourceInfo | None: + """Parse the consumer sub-dict (PD ``remote_prefiller`` or symmetric + ``remote_kv_source``) into a ``P2PSourceInfo``, or None if absent/incomplete.""" + role = _remote_prefiller_params(kv_params) + do_probe = False + if role is None: + role = _remote_kv_source_params(kv_params) + do_probe = True + if not role: + return None + peer_id = _peer_id_from_params(role) + kv_request_id = role.get("kv_request_id") + if peer_id is None or not kv_request_id: + return None + return P2PSourceInfo( + kv_request_id=kv_request_id, + peer_id=peer_id, + do_probe=do_probe, + ) + + +def _parse_dest(kv_params: dict | None) -> P2PDestInfo | None: + """Parse the producer ``remote_decoder`` sub-dict into a ``P2PDestInfo``, + or None if the block is absent (not a remote-decode request).""" + role = _remote_decoder_params(kv_params) + if role is None: + return None + return P2PDestInfo(kv_request_id=role.get("kv_request_id") or None) + + +def _annotate_req_context(req_context: ReqContext) -> None: + """Parse kv_transfer_params once and cache the P2P routing state. + + Called from ``on_new_request``; later calls for the same request read + the cached ``P2PSourceInfo``/``P2PDestInfo`` via ``get_state`` instead + of re-parsing. + """ + source = _parse_source(req_context.kv_transfer_params) + if source is not None: + req_context.set_state(source) + dest = _parse_dest(req_context.kv_transfer_params) + if dest is not None: + req_context.set_state(dest) @dataclass @@ -157,6 +247,20 @@ class P2PSecondaryTierManager(SecondaryTierManager): **kwargs: Reserved for future tier-specific options. """ super().__init__(offloading_spec, primary_kv_view, tier_type) + # Block hashes chain from NONE_HASH, seeded from PYTHONHASHSEED + # (see init_none_hash in v1/core/kv_cache_utils.py). Peers with + # different seeds compute different hashes for identical content, so + # lookups silently miss and no KV crosses the wire. Require it here so + # a misconfigured P2P instance fails at startup rather than degrading + # silently; the value is also verified against each peer on handshake. + hash_seed = os.getenv("PYTHONHASHSEED") + if hash_seed is None: + raise ValueError( + "PYTHONHASHSEED must be set for P2P KV offload so that block " + "hashes match across instances. Set it to a fixed value (e.g. " + "PYTHONHASHSEED=0) on every P2P peer." + ) + self._hash_seed = hash_seed if host is None: host = envs.VLLM_P2P_SIDE_CHANNEL_HOST if port is None: @@ -210,6 +314,12 @@ class P2PSecondaryTierManager(SecondaryTierManager): # kv_request_ids that hit a transport/session failure; On load lookup() # rejects them so the request falls back to local prefill. self._failed_req_ids: set[str] = set() + # Synthetic lookup ctxs from reaped sessions still owing a + # ``parent.on_request_finished`` (the session's failed_serves). The + # dead session had no parent handle at teardown; these are flushed + # at the top of the next ``serve_external_requests`` where the + # handle is valid. + self._failed_serve_ctxs: list[ReqContext] = [] # ------------------------------------------------------------------ # SecondaryTierManager interface @@ -217,66 +327,81 @@ class P2PSecondaryTierManager(SecondaryTierManager): @override def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: - prefill = _prefill_params(req_context.kv_transfer_params) - if ( - not prefill - or not prefill.get("remote_host") - or not prefill.get("remote_port") - or not prefill.get("kv_request_id") - ): + source = req_context.get_state(P2PSourceInfo) + if source is None: + return LookupResult.MISS + if source.kv_request_id in self._failed_req_ids: return LookupResult.MISS - kv_request_id = prefill["kv_request_id"] - if kv_request_id in self._failed_req_ids: - return LookupResult.MISS + # Symmetric-P2P consumer (``remote_kv_source`` sub-dict): probe the + # peer asynchronously. First call registers the (kv_request_id, + # key) entry and returns RETRY; flush_pending_lookups() + # in on_schedule_end batches the LookupMsg; a later step's + # lookup() returns HIT/MISS once LookupRespMsg has arrived. + # PD path (``remote_prefiller`` sub-dict only) keeps the eager HIT. + if source.do_probe: + session = self._sessions.get(source.peer_id) + if session is None: + return LookupResult.MISS + result = session.register_lookup(source.kv_request_id, key) + if result is True: + return LookupResult.HIT + if result is False: + return LookupResult.MISS + return LookupResult.RETRY + + # PD consumer (we are the decoder): all kv blocks should be on the + # prefiller side. Return HIT immediately. return LookupResult.HIT @override def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: - """Open the outbound session toward the producer if needed. + """Parse kv_transfer_params once and open the outbound session. - On the decoder side (``prefill`` set), open a session toward the - producer at remote_host:remote_port so submit_load can issue - FetchMsg as soon as it fires. On the prefiller side, sessions - are created when the consumer's inbound connection arrives in - _accept_new_peers — submit_store no longer pre-creates anything. + Parses the P2P routing state onto ``req_context`` (cached for the + later lookup/submit/finish calls). On the consumer side + (``remote_prefiller`` for PD or ``remote_kv_source`` for symmetric + P2P), open a session toward the producer at remote_host:remote_port + so submit_load can issue FetchMsg as soon as it fires. On the + prefiller side, sessions are created when the consumer's inbound + connection arrives in _accept_new_peers — submit_store no longer + pre-creates anything. """ - prefill = _prefill_params(req_context.kv_transfer_params) - if prefill: - peer_id = self._remote_id_from_params(prefill) - if peer_id: - self._get_or_create_session(peer_id) + _annotate_req_context(req_context) + source = req_context.get_state(P2PSourceInfo) + if source is not None: + self._get_or_create_session(source.peer_id) return RequestOffloadingContext() @override def on_request_finished(self, req_context: ReqContext) -> None: """Cancels pending loads and prunes session-scoped state. - Decoder side (``prefill`` set): looks up the session by peer_id - because the producer's address is what addresses the client-role - load to cancel. Prefiller side (``decode`` set): looks up via - kv_request_id because peer_id is no longer carried on store-time + Consumer side (``remote_prefiller`` for PD or ``remote_kv_source`` + for symmetric-P2P): looks up the session by peer_id because the + producer's address is what addresses the client-role load to + cancel; also drops any pending symmetric-P2P lookup state via + ``session.finish_request``. + Prefiller side (``remote_decoder`` set): looks up via kv_request_id + because peer_id is no longer carried on store-time kv_transfer_params; if a session has bound the id, finish it. If no session has bound the id yet, this is a no-op: parked batches in `_unbound_stores` are left in place and cleaned up only by `_reap_unbound_stores` after `_UNBOUND_STORE_TIMEOUT_S`. """ - kv_params = req_context.kv_transfer_params - if not kv_params: - return - prefill = _prefill_params(kv_params) - decode = _decode_params(kv_params) - kv_request_id = (prefill or decode or {}).get("kv_request_id") + source = req_context.get_state(P2PSourceInfo) + dest = req_context.get_state(P2PDestInfo) + kv_request_id = source.kv_request_id if source is not None else None + if kv_request_id is None and dest is not None: + kv_request_id = dest.kv_request_id if not kv_request_id: return self._failed_req_ids.discard(kv_request_id) - if prefill: - peer_id = self._remote_id_from_params(prefill) - if peer_id: - session = self._sessions.get(peer_id) - if session is not None: - session.finish_request(kv_request_id) + if source is not None: + session = self._sessions.get(source.peer_id) + if session is not None: + session.finish_request(kv_request_id) return # Prefiller-side finish: identify the session via kv_request_id. @@ -293,24 +418,24 @@ class P2PSecondaryTierManager(SecondaryTierManager): assert len(keys) == len(block_ids) - kv_params = job_metadata.req_context.kv_transfer_params - decode = _decode_params(kv_params) + dest = job_metadata.req_context.get_state(P2PDestInfo) logger.debug( - "P2P %s: submit_store ENTRY job_id=%d blocks=%d decode=%s kv_request_id=%s", + "P2P %s: submit_store ENTRY job_id=%d blocks=%d " + "remote_decoder=%s kv_request_id=%s", self._local_id, job_id, len(block_ids), - decode is not None, - (decode or {}).get("kv_request_id"), + dest is not None, + dest.kv_request_id if dest is not None else None, ) - # Absent ``decode`` block => not a remote-decode request: succeed - # locally without parking. An empty/malformed dict is still a - # remote-decode signal and must fail the missing-id check below. - if decode is None: + # Absent ``remote_decoder`` block => not a remote-decode request: + # succeed locally without parking. An empty/malformed dict is still + # a remote-decode signal and must fail the missing-id check below. + if dest is None: self._finished_jobs.append(JobResult(job_id=job_id, success=True)) return - kv_request_id = decode.get("kv_request_id") + kv_request_id = dest.kv_request_id if not kv_request_id: logger.warning( "P2P %s: submit_store missing kv_request_id", @@ -350,32 +475,26 @@ class P2PSecondaryTierManager(SecondaryTierManager): keys = list(job_metadata.keys) block_ids = job_metadata.block_ids - prefill = _prefill_params(job_metadata.req_context.kv_transfer_params) + source = job_metadata.req_context.get_state(P2PSourceInfo) logger.debug( "P2P %s: submit_load ENTRY job_id=%d blocks=%d kv_request_id=%s peer=%s", self._local_id, job_id, len(block_ids), - (prefill or {}).get("kv_request_id"), - self._remote_id_from_params(prefill or {}), + source.kv_request_id if source is not None else None, + source.peer_id if source is not None else None, ) - if ( - not prefill - or not prefill.get("remote_host") - or not prefill.get("remote_port") - or not prefill.get("kv_request_id") - ): + if source is None: logger.debug( - "P2P %s: submit_load job_id=%d FAILED missing prefill params", + "P2P %s: submit_load job_id=%d FAILED missing consumer params", self._local_id, job_id, ) self._finished_jobs.append(JobResult(job_id=job_id, success=False)) return - kv_request_id = prefill["kv_request_id"] - peer_id = self._remote_id_from_params(prefill) - assert peer_id is not None # guaranteed by prefill checks above + kv_request_id = source.kv_request_id + peer_id = source.peer_id if not keys: logger.debug( @@ -442,10 +561,7 @@ class P2PSecondaryTierManager(SecondaryTierManager): warned = False while True: self._poll_once() - pending = any( - s._client._inbound or s._server._inflight - for s in self._sessions.values() - ) + pending = any(s.has_pending_work for s in self._sessions.values()) if not pending: return if not warned and time.monotonic() - start > 5.0: @@ -456,31 +572,45 @@ class P2PSecondaryTierManager(SecondaryTierManager): warned = True time.sleep(_DRAIN_SLEEP_S) + @override + def serve_external_requests(self, parent: ParentManager) -> None: + """Serve inbound peer lookups against the tiering manager. + + Called once per scheduler step (before this tier's + ``on_schedule_end``) with a ``parent`` handle valid only for the + duration of the call — the sole window in which the P2P server + role may query the tiering manager. First release bookkeeping for + the failed serves left by a reaped session, then let every live + session resolve its enqueued inbound LookupMsgs. + """ + if self._failed_serve_ctxs: + for ctx in self._failed_serve_ctxs: + parent.on_request_finished(ctx) + self._failed_serve_ctxs = [] + for session in self._sessions.values(): + session.serve_external_requests(parent) + @override def on_schedule_end(self, context: ScheduleEndContext) -> None: - return + # Flush any p2p lookups aggregated during this step. + # One LookupMsg per (peer, kv_request_id) with unsent entries; + # send-gating happens inside the session if not yet ready. + for session in self._sessions.values(): + session.flush_pending_lookups() # ------------------------------------------------------------------ # Internal # ------------------------------------------------------------------ - @staticmethod - def _remote_id_from_params(role_params: dict) -> str | None: - """Build peer_id from a role-scoped sub-dict (``prefill``/``p2p``).""" - host = role_params.get("remote_host") - port = role_params.get("remote_port") - if host and port: - return f"{host}:{port}" - return None - def _get_or_create_session(self, peer_id: str) -> P2PSession: """Return the existing session for peer_id, or open one outbound. - Decoder-side helper for on_new_request: when ``prefill`` is set, - the consumer must reach the producer at peer_id. If we already - have a session toward that peer (from a prior load or a - peer-initiated inbound), reuse it; otherwise open an outbound - ControlConnection and build a connected session. + Consumer-side helper for on_new_request: when ``remote_prefiller`` + (PD) or ``remote_kv_source`` (symmetric P2P) is set, the consumer must reach the + producer at peer_id. If we already have a session toward that + peer (from a prior load or a peer-initiated inbound), reuse it; + otherwise open an outbound ControlConnection and build a + connected session. """ session = self._sessions.get(peer_id) if session is not None: @@ -491,6 +621,7 @@ class P2PSecondaryTierManager(SecondaryTierManager): local_id=self._local_id, transport=self._data, local_block_len=self._data.block_len, + local_hash_seed=self._hash_seed, conn=conn, ) self._sessions[peer_id] = session @@ -512,6 +643,7 @@ class P2PSecondaryTierManager(SecondaryTierManager): local_id=self._local_id, transport=self._data, local_block_len=self._data.block_len, + local_hash_seed=self._hash_seed, conn=conn, ) logger.info( @@ -546,12 +678,19 @@ class P2PSecondaryTierManager(SecondaryTierManager): ] for kid in stale_kv_ids: del self._kv_to_session[kid] - failed_loads, failed_stores = session.close() - for job_id, kv_request_id in failed_loads: + close_result = session.close() + for job_id in close_result.failed_jobs: self._finished_jobs.append(JobResult(job_id=job_id, success=False)) - self._failed_req_ids.add(kv_request_id) - for job_id in failed_stores: + for job_id in close_result.failed_stores: self._finished_jobs.append(JobResult(job_id=job_id, success=False)) + # Fail every client-side request (in-flight loads plus unresolved + # symmetric-P2P probes) toward the dead peer so lookup() returns + # MISS (local prefill) instead of RETRY forever — even if a fresh + # session to the same peer is later opened by another request. + self._failed_req_ids.update(close_result.failed_req_ids) + # Release the TieringManager's per-request bookkeeping for the + # dead session's synthetic lookups on the next serve_external_requests. + self._failed_serve_ctxs.extend(close_result.failed_serves) self._data.remove_remote_peer(pid) logger.warning("P2P %s: peer %s down", self._local_id, pid) @@ -648,6 +787,9 @@ class P2PSecondaryTierManager(SecondaryTierManager): def shutdown(self) -> None: self._drain_inflight_for_shutdown() for session in self._sessions.values(): + # Orphan ctxs from close() are intentionally dropped: the manager + # is being torn down, so there is no next serve_external_requests + # to flush them and no TieringManager left to release. session.close() self._sessions.clear() self._kv_to_session.clear() diff --git a/vllm/v1/kv_offload/tiering/p2p/session/__init__.py b/vllm/v1/kv_offload/tiering/p2p/session/__init__.py index 82148e8340d..2efaf09e34e 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/__init__.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/__init__.py @@ -3,6 +3,7 @@ from vllm.v1.kv_offload.tiering.p2p.session.session import ( LoadResult, P2PSession, + SessionCloseResult, SessionPollResult, StoreResult, ) @@ -10,6 +11,7 @@ from vllm.v1.kv_offload.tiering.p2p.session.session import ( __all__ = [ "LoadResult", "P2PSession", + "SessionCloseResult", "SessionPollResult", "StoreResult", ] diff --git a/vllm/v1/kv_offload/tiering/p2p/session/client.py b/vllm/v1/kv_offload/tiering/p2p/session/client.py index f080cb79a36..ed91f7f9b00 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/client.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/client.py @@ -11,16 +11,19 @@ callback injected by the coordinator (which gates on ConnectAck). from __future__ import annotations +import enum import time from collections.abc import Callable, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING, NamedTuple from vllm.logger import init_logger +from vllm.v1.kv_offload.base import OffloadKey from vllm.v1.kv_offload.tiering.p2p.session.protocol import ( TYPE_KEY, AbortFetchMsg, FetchMsg, + LookupMsg, ) if TYPE_CHECKING: @@ -32,16 +35,60 @@ _LOAD_TIMEOUT_S = 30.0 _ABORT_ACK_TIMEOUT_S = 10.0 +class ClientPhase(enum.Enum): + """Lifecycle of a request's client-side lookup/fetch signalling. + + Advances monotonically. Only ``finish`` reads it, to decide + whether a terminal empty FetchMsg is owed to release the peer's + lookup state (owed only from ``PROBING``: a LookupMsg went out but no + FetchMsg has since closed the peer's lookup phase). + """ + + REGISTERED = enum.auto() # keys in probes/unsent, nothing sent yet + PROBING = enum.auto() # LookupMsg flushed, awaiting responses + FETCH_SENT = enum.auto() # FetchMsg sent (real or terminal empty) + + @dataclass -class _InboundRequestState: - """Client-role state for a single load request.""" +class _InboundLoadState: + """Client-role state for a single in-flight load request. + + Lives on ``_ClientRequestState.load`` for the duration of a fetch; + the owning kv_request_id is the dict key, so it isn't stored here. + """ job_id: int # opaque ID assigned by the manager to this load request - kv_request_id: str submitted_at: float aborted_at: float | None = None +@dataclass +class _ClientRequestState: + """Per-kv_request_id client-side state. + + One entry per kv_request_id we're driving. Lookup-phase fields are + used only by symmetric P2P (``do_p2p_fetch``); PD-only loads leave + ``probes``/``unsent`` empty and drive just ``phase`` and ``load``. An + entry is dropped once every field is idle — see ``ClientRole._maybe_prune``. + """ + + # -- Lookup phase (symmetric P2P only; untouched for PD) -- + # Probe outcome per OffloadKey: None while in-flight (registered/sent + # but unresolved), True/False once a LookupRespMsg lands. There is no + # timeout — finish (via finish_request) is guaranteed after + # the request's lookup() calls and clears every probe, so an + # unanswered probe simply stays None until then. + probes: dict[OffloadKey, bool | None] = field(default_factory=dict) + # OffloadKeys registered but not yet flushed onto the wire. Drained and + # cleared by the next flush_pending_lookups. + unsent: list[OffloadKey] = field(default_factory=list) + + # Monotonic lookup/fetch signalling phase; see ``ClientPhase``. + phase: ClientPhase = ClientPhase.REGISTERED + # Set while a fetch is in flight; cleared on completion/abort/timeout. + load: _InboundLoadState | None = None + + class LoadResult(NamedTuple): """Result from a session poll, client side.""" @@ -50,6 +97,21 @@ class LoadResult(NamedTuple): success: bool +class ClientCloseResult(NamedTuple): + """What the manager must fail when a peer session is torn down. + + ``failed_jobs`` is the ``job_id`` of every load still in flight; each + becomes a ``JobResult(success=False)``. ``failed_req_ids`` is every + kv_request_id whose lookup() would otherwise defer forever on the dead + peer — the in-flight loads plus any request holding an unresolved + symmetric-P2P probe. ``failed_jobs`` is the subset of ``failed_req_ids`` + that had a load job. + """ + + failed_jobs: list[int] + failed_req_ids: list[str] + + class ClientRole: """Client-side load state machine for one peer session. @@ -61,9 +123,55 @@ class ClientRole: def __init__(self, peer_id: str, send: Callable[[dict], None]) -> None: self._peer_id = peer_id self._send = send - self._inbound: dict[str, _InboundRequestState] = {} + # All per-kv_request_id state lives here. Entries are created + # lazily by request_blocks / register_lookup and dropped by + # _maybe_prune once every field is idle. + self._requests: dict[str, _ClientRequestState] = {} + # kv_request_ids with unsent lookup keys for the next flush to + # visit — the work-list that keeps flush_pending_lookups from + # scanning every request each scheduler step. Mirrors the server's + # _serve_pending. Populated by register_lookup, drained by + # flush_pending_lookups, and discarded on finish/close. + self._flush_pending: set[str] = set() + # kv_request_ids with a fetch in flight (``st.load is not None``) — + # the work-list collect_results walks for timeouts, and the + # has_active_loads predicate, instead of scanning every request. + # Kept in exact sync with ``st.load``: armed in request_blocks, + # discarded wherever load is cleared, and cleared on close. + self._active_loads: set[str] = set() self._completed_loads: list[LoadResult] = [] + # ------------------------------------------------------------------ + # State helpers + # ------------------------------------------------------------------ + + def _get_or_create_request(self, kv_request_id: str) -> _ClientRequestState: + """Get or create the state entry for a kv_request_id.""" + st = self._requests.get(kv_request_id) + if st is None: + st = _ClientRequestState() + self._requests[kv_request_id] = st + return st + + def _maybe_prune(self, kv_request_id: str) -> None: + """Drop the entry once it holds no live load or lookup state. + + The sticky ``phase`` is only read by ``finish``. A probe + clears when its fetch is issued (``request_blocks``) or when the + request finishes (``finish``/``close``); in the former case + ``load`` is set and keeps the entry alive, in the latter the phase + is no longer needed — so dropping on emptiness never loses a phase + still in use. + """ + st = self._requests.get(kv_request_id) + if st is not None and st.load is None and not st.probes and not st.unsent: + del self._requests[kv_request_id] + + @property + def has_active_loads(self) -> bool: + """True if any kv_request_id has a fetch in flight.""" + return bool(self._active_loads) + # ------------------------------------------------------------------ # Public API # ------------------------------------------------------------------ @@ -72,7 +180,7 @@ class ClientRole: self, job_id: JobId, kv_request_id: str, - keys: Sequence[bytes], + keys: Sequence[OffloadKey], block_ids: Sequence[int], send_ready: bool, ) -> None: @@ -86,44 +194,101 @@ class ClientRole: len(block_ids), send_ready, ) - self._inbound[kv_request_id] = _InboundRequestState( + st = self._get_or_create_request(kv_request_id) + st.load = _InboundLoadState( job_id=job_id, - kv_request_id=kv_request_id, submitted_at=time.monotonic(), ) + self._active_loads.add(kv_request_id) + st.phase = ClientPhase.FETCH_SENT self._send( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: kv_request_id, - FetchMsg.BLOCK_HASHES: list(keys), + FetchMsg.KEYS: list(keys), FetchMsg.BLOCK_INDEXES: [int(idx) for idx in block_ids], } ) + # Issuing the fetch ends this request's lookup phase, so drop all + # probe state. Once the peer serves this fetch both sides unpin, so + # the producer may evict the block; a stale cached True would + # otherwise let a re-scheduled lookup() return HIT without + # re-probing, pointing at a block the producer no longer holds. + # Clearing forces a fresh LookupMsg on re-schedule so the producer + # answers from current state. For a symmetric-P2P request (probes + # populated) every fetched block was a confirmed HIT; a PD-only load + # never probes, so probes is empty and the clear is a no-op. + if st.probes: + assert all(st.probes.get(key) is True for key in keys) + st.probes.clear() - def cancel(self, kv_request_id: str) -> None: - """Cancel a pending load. Sends AbortFetchMsg if still active.""" - req = self._inbound.pop(kv_request_id, None) - if req is not None and req.aborted_at is None: + def finish(self, kv_request_id: str) -> None: + """Finish a request: abort any in-flight load and release lookup state. + + Called from the session's ``finish_request``. The two branches are + mutually exclusive: ``load`` is set only by ``request_blocks``, which + also advances ``phase`` to ``FETCH_SENT``, and nothing moves it back + to ``PROBING`` — so a fetch in flight never coexists with the + ``PROBING`` phase. + + - Fetch in flight (``load`` set, phase ``FETCH_SENT``): send an + AbortFetchMsg unless the load is already aborting, then drop it. + - Outstanding lookups (``PROBING``): a LookupMsg was flushed but no + FetchMsg has closed the peer's lookup phase. Every FetchMsg the + server receives in p2p mode is its "request finished" signal (it + releases lookup state and fires ``cb.finish_request``); when the + client's lookups all missed no FetchMsg is otherwise sent, so emit + a terminal empty one purely to trigger those semantics. In + ``REGISTERED`` the peer never received a LookupMsg and in + ``FETCH_SENT`` a FetchMsg already closed the phase, so neither owes + a terminal FetchMsg. + + Then drop all probe/lookup state and prune the entry. + """ + st = self._requests.get(kv_request_id) + if st is None: + return + if st.load is not None: + if st.load.aborted_at is None: + self._send( + { + TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.KV_REQUEST_ID: kv_request_id, + } + ) + st.load = None + self._active_loads.discard(kv_request_id) + elif st.phase is ClientPhase.PROBING: + st.phase = ClientPhase.FETCH_SENT self._send( { - TYPE_KEY: AbortFetchMsg.TYPE, - AbortFetchMsg.KV_REQUEST_ID: kv_request_id, + TYPE_KEY: FetchMsg.TYPE, + FetchMsg.KV_REQUEST_ID: kv_request_id, + FetchMsg.KEYS: [], + FetchMsg.BLOCK_INDEXES: [], } ) + st.probes.clear() + st.unsent.clear() + self._flush_pending.discard(kv_request_id) + self._maybe_prune(kv_request_id) def on_transfer_done(self, kv_request_id: str, success: bool) -> None: """Handle a TransferDoneMsg from the peer.""" - req = self._inbound.pop(kv_request_id, None) - if req is not None: + st = self._requests.get(kv_request_id) + if st is not None and st.load is not None: self._completed_loads.append( LoadResult( - job_id=req.job_id, + job_id=st.load.job_id, kv_request_id=kv_request_id, success=success, ) ) + st.load = None + self._active_loads.discard(kv_request_id) + self._maybe_prune(kv_request_id) else: - # No matching _inbound entry: either a duplicate + # No matching in-flight load: either a duplicate # transfer_done from the peer (protocol violation) or a # benign race with a local cancel/abort/timeout that # already popped the entry. We don't track terminated ids, @@ -137,23 +302,26 @@ class ClientRole: def on_abort_ack(self, kv_request_id: str) -> None: """Handle an AbortAckMsg from the peer.""" - req = self._inbound.pop(kv_request_id, None) - if req is not None: + st = self._requests.get(kv_request_id) + if st is not None and st.load is not None: logger.warning( "P2PSession %s: load request %s (job_id=%d) timed out; " "load job completed with failure. If this recurs, ensure " "PYTHONHASHSEED is set to the same value on all nodes.", self._peer_id, kv_request_id, - req.job_id, + st.load.job_id, ) self._completed_loads.append( LoadResult( - job_id=req.job_id, + job_id=st.load.job_id, kv_request_id=kv_request_id, success=False, ) ) + st.load = None + self._active_loads.discard(kv_request_id) + self._maybe_prune(kv_request_id) else: # See on_transfer_done: same ambiguity (duplicate ack # vs. raced with local cancel/timeout that already popped). @@ -164,19 +332,141 @@ class ClientRole: kv_request_id, ) + # ------------------------------------------------------------------ + # Symmetric-P2P lookup (do_p2p_fetch=true) + # ------------------------------------------------------------------ + + def register_lookup(self, kv_request_id: str, key: bytes) -> bool | None: + """Register or resolve one (kv_request_id, key) probe. + + Idempotent across scheduler steps: + - First call: creates a pending entry, returns None. + - Subsequent calls while in-flight: returns None. + - Once a LookupRespMsg has resolved the entry: returns the cached + bool result on every call without popping it. + + A resolved entry is retained until its fetch is issued + (``request_blocks`` pops it) or the request finishes + (``finish`` clears all entries for the id). A request's + block set can be re-probed across steps, so popping on read would + make a repeat probe of an already-resolved key look brand-new and + re-queue it, emitting a redundant LookupMsg for an answer we + already hold. Keeping the entry until fetch makes repeat probes + free; clearing it at fetch forces a fresh probe if the request is + re-scheduled, since the block is unpinned once served. + """ + st = self._get_or_create_request(kv_request_id) + okey = OffloadKey(key) + if okey in st.probes: + return st.probes[okey] + st.probes[okey] = None + st.unsent.append(okey) + self._flush_pending.add(kv_request_id) + logger.debug( + "P2P LOOKUP client %s: REGISTER kv_request_id=%s key=%s (unsent=%d)", + self._peer_id, + kv_request_id, + key.hex()[:16], + len(st.unsent), + ) + return None + + def flush_pending_lookups(self) -> None: + """Send a LookupMsg for each kv_request_id with unsent entries. + + Called once per scheduler step from the manager's + ``on_schedule_end()``. A request's block set may be discovered + across several scheduler steps, so more than one LookupMsg can + go out per kv_request_id — one per step that registered new + keys. register_lookup() de-dups in-flight and already-resolved + (req_id, key) pairs, so each LookupMsg carries only the keys + first probed in that step. The peer's lookup phase for the id is + still closed by exactly one FetchMsg, which the client contract + guarantees is sent after every lookup for the id has resolved + (see request_blocks / finish). Send-gating is handled by + the injected ``_send`` callback (queues until ConnectAckMsg if + needed). + + Only requests that registered new keys since the last flush are + visited — the ``_flush_pending`` work-list avoids scanning every + live request each scheduler step. + """ + for req_id in self._flush_pending: + st = self._requests.get(req_id) + if st is None or not st.unsent: + continue + # Record that the peer now holds lookup state for this id so + # finish knows a terminal empty FetchMsg may be owed. + # Only promote from REGISTERED: once a fetch has gone out + # (FETCH_SENT) a later LookupMsg must not regress the phase, as + # no terminal FetchMsg is owed for an already-fetched request. + if st.phase is ClientPhase.REGISTERED: + st.phase = ClientPhase.PROBING + logger.debug( + "P2P LOOKUP client %s: SEND LookupMsg kv_request_id=%s keys=%d", + self._peer_id, + req_id, + len(st.unsent), + ) + self._send( + { + TYPE_KEY: LookupMsg.TYPE, + LookupMsg.KV_REQUEST_ID: req_id, + LookupMsg.KEYS: list(st.unsent), + } + ) + st.unsent = [] + self._flush_pending.clear() + + def on_lookup_resp( + self, + kv_request_id: str, + keys: Sequence[bytes], + hits: Sequence[bool], + ) -> None: + """Apply per-pair hit/miss results from a peer. + + Pairs that don't match a known entry (already cancelled or + never asked) are silently dropped — the producer is free to + split or coalesce responses. + """ + n_hit = sum(1 for hit in hits if hit) + logger.debug( + "P2P LOOKUP client %s: RECV LookupRespMsg kv_request_id=%s " + "keys=%d hits=%d misses=%d", + self._peer_id, + kv_request_id, + len(keys), + n_hit, + len(hits) - n_hit, + ) + st = self._requests.get(kv_request_id) + if st is None: + return + for h, hit in zip(keys, hits): + key = OffloadKey(h) + if key in st.probes: + st.probes[key] = hit + def collect_results(self) -> list[LoadResult]: - """Walk timeouts and drain completed loads. + """Walk load timeouts and drain completed loads. Active requests past ``_LOAD_TIMEOUT_S`` get an AbortFetchMsg sent and enter the aborting phase. Aborting requests past ``_ABORT_ACK_TIMEOUT_S`` are surfaced as failed loads. + + Lookups have no timeout: an unanswered probe stays None (RETRY) + until finish_request clears it — see ``_ClientRequestState.probes``. """ now = time.monotonic() to_remove: list[str] = [] - for req_id, req in self._inbound.items(): - if req.aborted_at is None: - if now - req.submitted_at >= _LOAD_TIMEOUT_S: - req.aborted_at = now + for req_id in self._active_loads: + st = self._requests[req_id] + assert st.load is not None + load = st.load + if load.aborted_at is None: + if now - load.submitted_at >= _LOAD_TIMEOUT_S: + load.aborted_at = now logger.warning( "P2PSession %s: %s timed out, sending abort", self._peer_id, @@ -189,11 +479,11 @@ class ClientRole: } ) else: - if now - req.aborted_at >= _ABORT_ACK_TIMEOUT_S: + if now - load.aborted_at >= _ABORT_ACK_TIMEOUT_S: to_remove.append(req_id) self._completed_loads.append( LoadResult( - job_id=req.job_id, + job_id=load.job_id, kv_request_id=req_id, success=False, ) @@ -204,15 +494,31 @@ class ClientRole: req_id, ) for req_id in to_remove: - self._inbound.pop(req_id) + self._requests[req_id].load = None + self._active_loads.discard(req_id) + self._maybe_prune(req_id) results = self._completed_loads self._completed_loads = [] return results - def close(self) -> list[tuple[int, str]]: - """Tear down. Returns ``(job_id, kv_request_id)`` for pending loads.""" - failed = [(req.job_id, req.kv_request_id) for req in self._inbound.values()] - self._inbound.clear() + def close(self) -> ClientCloseResult: + """Tear down, reporting work the dead peer can no longer complete. + + A request is failed if it has a load in flight (its job fails) or + holds an unresolved symmetric-P2P probe (its lookup() would defer + forever on an answer that can never arrive). See ``ClientCloseResult``. + """ + failed_jobs = [ + st.load.job_id for st in self._requests.values() if st.load is not None + ] + failed_req_ids = [ + req_id + for req_id, st in self._requests.items() + if st.load is not None or any(hit is None for hit in st.probes.values()) + ] + self._requests.clear() + self._flush_pending.clear() + self._active_loads.clear() self._completed_loads.clear() - return failed + return ClientCloseResult(failed_jobs=failed_jobs, failed_req_ids=failed_req_ids) diff --git a/vllm/v1/kv_offload/tiering/p2p/session/protocol.py b/vllm/v1/kv_offload/tiering/p2p/session/protocol.py index 3a2ab5fd88a..8988c7e79f3 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/protocol.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/protocol.py @@ -26,6 +26,14 @@ Block Transfer Flow (happy path) 1. Client sends FetchMsg with a kv_request_id and lists of block keys + remote indexes where it wants the data written. + In p2p mode FetchMsg is also the server-side "request finished" + signal for the id: no further ``cb.create_store_job`` will fire + (parked LookupMsg batches are popped, so pending-key resolution + cannot promote a HIT after this point), all server-side lookup + state for the id is released, and ``cb.finish_request`` fires on + each dropped batch. The client emits exactly one FetchMsg per + lookup-touched request, including an empty one when no blocks + end up being fetched. 2. Server matches requested blocks against locally stored blocks: - Blocks already available are transferred immediately via RDMA. - Blocks not yet available are recorded as "demanded" and @@ -119,6 +127,9 @@ class ConnectMsg: BLOCK_LEN: Size in bytes of each block (must match between peers). CONFIG_FINGERPRINT: SHA-256 prefix of the model configuration. Peers with different fingerprints are incompatible. + HASH_SEED: The peer's PYTHONHASHSEED. Block hashes chain from a seed + derived from it, so peers with different values compute different + hashes for identical content and must not exchange blocks. """ TYPE = "connect" @@ -128,6 +139,7 @@ class ConnectMsg: NUM_BLOCKS = "num_blocks" BLOCK_LEN = "block_len" CONFIG_FINGERPRINT = "config_fingerprint" + HASH_SEED = "hash_seed" @staticmethod def validate(msg: dict) -> None: @@ -137,6 +149,7 @@ class ConnectMsg: _require_non_neg_int(msg, ConnectMsg.BASE_ADDR) _require_pos_int(msg, ConnectMsg.NUM_BLOCKS) _require_pos_int(msg, ConnectMsg.BLOCK_LEN) + _require(msg, ConnectMsg.HASH_SEED, str) class ConnectAckMsg: @@ -165,37 +178,106 @@ class DisconnectMsg: class FetchMsg: - """Client → Server: request blocks by key. + """Client → Server: request blocks by key and close the lookup phase. + + In p2p mode FetchMsg is also the server-side "request finished" + signal for ``kv_request_id``: on receipt the server (a) fires no + further ``cb.create_store_job`` for this id — parked LookupMsg + batches are popped, so ``_resolve_pending_lookups`` cannot promote + a HIT_PENDING / RETRY key into a fresh pin after this point — and + (b) calls ``cb.finish_request(batch.ctx)`` on each dropped batch + so the TieringManager can release per-batch bookkeeping. In the + all-miss case the client emits an empty FetchMsg (``KEYS`` + and ``BLOCK_INDEXES`` both empty) purely to fire this signal. Fields: KV_REQUEST_ID: Identifies this block transfer request. - BLOCK_HASHES: List of block keys (OffloadKey bytes). - BLOCK_INDEXES: List of remote block indexes (same length as BLOCK_HASHES). + KEYS: List of block keys (OffloadKey bytes). May be empty. + BLOCK_INDEXES: List of remote block indexes (same length as KEYS). """ TYPE = "fetch" KV_REQUEST_ID = "kv_request_id" - BLOCK_HASHES = "block_hashes" + KEYS = "keys" BLOCK_INDEXES = "block_indexes" @staticmethod def validate(msg: dict) -> None: """Raise ValueError if any field has an invalid type or value.""" _require(msg, FetchMsg.KV_REQUEST_ID, str) - _require_list(msg, FetchMsg.BLOCK_HASHES) + _require_list(msg, FetchMsg.KEYS) _require_list(msg, FetchMsg.BLOCK_INDEXES) - hashes = msg[FetchMsg.BLOCK_HASHES] + keys = msg[FetchMsg.KEYS] indexes = msg[FetchMsg.BLOCK_INDEXES] - if len(hashes) != len(indexes): + if len(keys) != len(indexes): raise ValueError( - f"block_hashes/block_indexes length mismatch: " - f"{len(hashes)} vs {len(indexes)}" + f"keys/block_indexes length mismatch: {len(keys)} vs {len(indexes)}" ) for idx in indexes: if not isinstance(idx, int) or idx < 0: raise ValueError(f"block_indexes: invalid index {idx!r}") +class LookupMsg: + """Client → Server: probe which block keys the peer holds. + + Sent on the consumer side under symmetric P2P (do_p2p_fetch=true) + after the consumer has aggregated per-block lookups across a + scheduler step. The producer replies with one or more LookupRespMsg + covering the requested keys. + + Fields: + KV_REQUEST_ID: Identifies this lookup transaction. + KEYS: List of block keys (OffloadKey bytes) to probe. + """ + + TYPE = "lookup" + KV_REQUEST_ID = "kv_request_id" + KEYS = "keys" + + @staticmethod + def validate(msg: dict) -> None: + """Raise ValueError if any field has an invalid type or value.""" + _require(msg, LookupMsg.KV_REQUEST_ID, str) + _require_list(msg, LookupMsg.KEYS) + + +class LookupRespMsg: + """Server → Client: per-key hit/miss answer for a prior LookupMsg. + + Carries two parallel arrays of equal length so each (key, + hit) pair is self-describing. The producer is free to split or + coalesce responses across multiple LookupMsgs for the same + KV_REQUEST_ID — the consumer matches each pair back to its + pending entry by (KV_REQUEST_ID, key). + + Fields: + KV_REQUEST_ID: The lookup transaction this responds to. + KEYS: List of block keys answered by this message. + HITS: Parallel list of bools — True if the producer holds the + corresponding block, False otherwise. + """ + + TYPE = "lookup_resp" + KV_REQUEST_ID = "kv_request_id" + KEYS = "keys" + HITS = "hits" + + @staticmethod + def validate(msg: dict) -> None: + """Raise ValueError if any field has an invalid type or value.""" + _require(msg, LookupRespMsg.KV_REQUEST_ID, str) + _require_list(msg, LookupRespMsg.KEYS) + _require_list(msg, LookupRespMsg.HITS) + keys = msg[LookupRespMsg.KEYS] + hits = msg[LookupRespMsg.HITS] + if len(keys) != len(hits): + raise ValueError(f"keys/hits length mismatch: {len(keys)} vs {len(hits)}") + for hit in hits: + if not isinstance(hit, bool): + raise ValueError(f"hits: invalid value {hit!r}") + + class TransferDoneMsg: """Server → Client: all blocks transferred for a request. diff --git a/vllm/v1/kv_offload/tiering/p2p/session/server.py b/vllm/v1/kv_offload/tiering/p2p/session/server.py index 2a6c8c69244..3709ecd9c53 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/server.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/server.py @@ -18,26 +18,32 @@ coordinator's ``_dispatch_message`` can reuse its existing from __future__ import annotations import time -from collections.abc import Callable, Sequence +from collections.abc import Callable, Iterable, Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING, NamedTuple from vllm.logger import init_logger -from vllm.v1.kv_offload.base import OffloadKey +from vllm.v1.kv_offload.base import LookupResult, OffloadKey, ReqContext from vllm.v1.kv_offload.tiering.p2p.session.protocol import ( TYPE_KEY, AbortAckMsg, + LookupRespMsg, TransferDoneMsg, ) if TYPE_CHECKING: - from vllm.v1.kv_offload.tiering.base import JobId + from vllm.v1.kv_offload.tiering.base import JobId, ParentManager from vllm.v1.kv_offload.tiering.p2p.data import DataTransport logger = init_logger(__name__) _STORE_TIMEOUT_S = 30.0 _CANCEL_DRAIN_TIMEOUT_S = 10.0 +# Cap on time the server holds a HIT_PENDING / RETRY key from an inbound +# LookupMsg before falling back to MISS. Long enough that an in-flight +# primary write or a just-started promotion typically completes; short +# enough that the consumer doesn't sit idle on a stuck producer. +_LOOKUP_PENDING_TIMEOUT_S = 5.0 class StoreResult(NamedTuple): @@ -69,8 +75,8 @@ class _MatchResult(NamedTuple): class _OutboundRequestState: """Server-role state for a single peer fetch request. - The owning ``kv_request_id`` is the dict key in - ``ServerRole._outbound`` and is not duplicated on the value. + The owning ``kv_request_id`` is the ``ServerRole._requests`` dict key + and is not duplicated on the value. """ demand_received: bool = False @@ -90,7 +96,7 @@ class _OutboundRequestState: def add_stored_blocks( self, - block_hashes: Sequence[OffloadKey], + keys: Sequence[OffloadKey], block_ids: Sequence[int], job_id: int, ) -> _MatchResult: @@ -98,13 +104,13 @@ class _OutboundRequestState: self.pending_job_ids.add(job_id) local_idxs: list[int] = [] remote_idxs: list[int] = [] - for block_hash, local_idx in zip(block_hashes, block_ids): - remote_idx = self.demanded.pop(block_hash, None) + for key, local_idx in zip(keys, block_ids): + remote_idx = self.demanded.pop(key, None) if remote_idx is not None: local_idxs.append(local_idx) remote_idxs.append(remote_idx) else: - self.available[block_hash] = (job_id, local_idx) + self.available[key] = (job_id, local_idx) return _MatchResult( local_idxs=local_idxs, remote_idxs=remote_idxs, @@ -113,25 +119,25 @@ class _OutboundRequestState: def add_fetch_demand( self, - block_hashes: Sequence[OffloadKey], + keys: Sequence[OffloadKey], block_indexes: Sequence[int], ) -> _MatchResult: """Register the peer's fetch demand. Returns matched pairs.""" self.demand_received = True - self.remaining = len(block_hashes) + self.remaining = len(keys) local_idxs: list[int] = [] remote_idxs: list[int] = [] job_ids: set[int] = set() - for block_hash, remote_idx in zip(block_hashes, block_indexes): - stored_entry = self.available.pop(block_hash, None) + for key, remote_idx in zip(keys, block_indexes): + stored_entry = self.available.pop(key, None) if stored_entry is not None: stored_job_id, local_idx = stored_entry local_idxs.append(local_idx) remote_idxs.append(remote_idx) job_ids.add(stored_job_id) else: - self.demanded[block_hash] = remote_idx + self.demanded[key] = remote_idx return _MatchResult( local_idxs=local_idxs, remote_idxs=remote_idxs, @@ -139,6 +145,78 @@ class _OutboundRequestState: ) +@dataclass +class _ActiveLookup: + """In-flight state for one inbound LookupMsg. + + Aggregates per-key HIT/MISS resolutions and defers the single + outbound LookupRespMsg until every key has been resolved or the + ``deadline`` fires (remaining ``pending`` keys then force-MISS). + Exactly one LookupRespMsg is emitted per LookupMsg, carrying every + key in the original wire order. + """ + + lookup_id: int + kv_request_id: str + ctx: ReqContext + # Keys from the inbound LookupMsg, preserved in wire order so + # the aggregated response goes back in the same order. + keys: list[OffloadKey] = field(default_factory=list) + # Per-key resolution: True = HIT, False = MISS. A key is present + # here once definitively resolved; still-pending keys are only + # in ``pending``. + resolved: dict[OffloadKey, bool] = field(default_factory=dict) + # Keys still awaiting resolution (HIT_PENDING / RETRY from + # ``parent.lookup``); re-polled by ``_resolve_pending_lookups``. + pending: set[OffloadKey] = field(default_factory=set) + # Absolute deadline (``time.monotonic``). Once reached, remaining + # ``pending`` keys are force-resolved to MISS so the consumer + # can fall back instead of waiting on a stuck producer. + deadline: float = 0.0 + + +class _PendingLookup(NamedTuple): + """A raw inbound LookupMsg awaiting resolution. + + Enqueued by ``on_lookup`` during dispatch and drained by the next + ``serve_external_requests``. The deadline for any resulting + HIT_PENDING / RETRY key is measured from ``enqueued_at``. + """ + + keys: list[OffloadKey] + enqueued_at: float + + +@dataclass +class _ServerRequestState: + """Per-kv_request_id server-side state. + + Consolidates the outbound serve/transfer state, the symmetric-P2P + inbound-lookup state, abort bookkeeping, and the inflight-transfer + count for one kv_request_id. The owning id is the ``_requests`` dict + key and is not duplicated here. An entry is dropped once every field + is idle — see ``ServerRole._maybe_prune``. + """ + + # Outbound serve state (PD + symmetric producer side). None until the + # first add_stored_blocks / on_fetch; reset to None on finalize/abort. + outbound: _OutboundRequestState | None = None + # Raw inbound LookupMsgs not yet processed against the ParentManager. + pending_lookups: list[_PendingLookup] = field(default_factory=list) + # Per-LookupMsg state parked with HIT_PENDING / RETRY keys, keyed by + # the (globally unique) lookup_id and re-polled each serve. + lookups: dict[int, _ActiveLookup] = field(default_factory=dict) + # Start time of a pending abort drain (``time.monotonic``); None when + # no abort is in progress. + abort_started_at: float | None = None + # Transfer ids in ``ServerRole._inflight`` for this id. Kept in sync + # via _inflight_add / _inflight_pop so a non-empty set is an exact + # "has any inflight transfer" predicate and the abort drain can + # enumerate this request's transfers without scanning all of + # ``_inflight``. + inflight_tids: set[int] = field(default_factory=set) + + class ServerRole: """Server-side store/serve state machine for one peer session. @@ -157,22 +235,55 @@ class ServerRole: self._transport = transport self._send = send - self._outbound: dict[str, _OutboundRequestState] = {} + # All per-kv_request_id state lives here. Entries are created + # lazily and dropped by _maybe_prune once every field is idle. + self._requests: dict[str, _ServerRequestState] = {} + # kv_request_ids with lookup work (unprocessed pending_lookups or + # parked lookups) for the next serve to visit — the work-list that + # keeps serve_external_requests from scanning every request. + self._serve_pending: set[str] = set() # transfer_id → xfer. Mutate ONLY via _inflight_add / _inflight_pop - # so the per-request count below stays in sync. + # so the per-request inflight_tids stays in sync. self._inflight: dict[int, _InflightXfer] = {} - # kv_request_id → number of entries in _inflight for that id. - # Kept in sync with _inflight; entries that hit zero are removed - # so `kv_request_id in self._inflight_per_req` is an exact - # "has any inflight transfer" predicate (O(1) replacement for - # the previous O(N) scan). - self._inflight_per_req: dict[str, int] = {} self._store_jobs: dict[int, float] = {} # job_id → submitted_at - self._pending_aborts: dict[str, float] = {} # kv_request_id → start # StoreResults queued by _finalize_outbound for the next poll # tick to surface. Mirrors the deferred-result pattern used for # load timeouts. self._pending_store_results: list[StoreResult] = [] + # Synthetic lookup ctxs whose ``on_request_finished`` still needs to + # fire but which were closed outside a serve window (FetchMsg / local + # finish popped their parked lookup). Drained via + # ``parent.on_request_finished`` in ``serve_external_requests``. + self._finished_lookup_ctxs: list[ReqContext] = [] + self._lookup_id_counter: int = 0 + # kv_request_ids with a parked abort awaiting drain — work-list so + # drain_pending_aborts doesn't scan every request each poll tick. + self._parked_aborts: set[str] = set() + + # ------------------------------------------------------------------ + # State helpers + # ------------------------------------------------------------------ + + def _get_or_create_request(self, kv_request_id: str) -> _ServerRequestState: + """Get or create the state entry for a kv_request_id.""" + st = self._requests.get(kv_request_id) + if st is None: + st = _ServerRequestState() + self._requests[kv_request_id] = st + return st + + def _maybe_prune(self, kv_request_id: str) -> None: + """Drop the entry once it holds no live state.""" + st = self._requests.get(kv_request_id) + if ( + st is not None + and st.outbound is None + and not st.inflight_tids + and not st.lookups + and not st.pending_lookups + and st.abort_started_at is None + ): + del self._requests[kv_request_id] # ------------------------------------------------------------------ # Public API @@ -187,19 +298,41 @@ class ServerRole: ) -> None: """New blocks stored locally — match against pending fetch demand.""" self._store_jobs[job_id] = time.monotonic() - req = self._outbound.setdefault(kv_request_id, _OutboundRequestState()) - result = req.add_stored_blocks(keys, block_ids, job_id) - if result.local_idxs and req.demand_received: + st = self._get_or_create_request(kv_request_id) + if st.outbound is None: + st.outbound = _OutboundRequestState() + result = st.outbound.add_stored_blocks(keys, block_ids, job_id) + if result.local_idxs and st.outbound.demand_received: self._submit_transfer(kv_request_id, result) def on_fetch( self, kv_request_id: str, - block_hashes: Sequence[OffloadKey], + keys: Sequence[OffloadKey], block_indexes: Sequence[int], ) -> None: """Handle a FetchMsg from the peer. + In p2p mode FetchMsg is the server-side "request finished" + signal for ``kv_request_id``. Three consequences flow from that: + + - No further ``parent.create_store_job`` will fire for this id: + the request's parked ``lookups`` are popped here (via + ``_finish_inbound_lookups``) before the next + ``serve_external_requests`` runs ``_resolve_pending_lookups``, + so any HIT_PENDING / RETRY key that would otherwise later + promote to HIT and pin a slot is dropped instead. Any raw + not-yet-processed LookupMsg for this id is dropped too. + - All server-side lookup state for the id is cleaned up (the + ``lookups`` entries themselves). + - The synthetic ctx is queued for ``parent.on_request_finished`` + (fired by the next ``serve_external_requests``) so the + TieringManager can release per-lookup bookkeeping. + + The client contract guarantees exactly one FetchMsg per + lookup-touched request — including an empty one when no + blocks end up being fetched. + Raises ``ValueError`` on a duplicate fetch for the same ``kv_request_id``; the coordinator's dispatch loop turns that into a protocol-error disconnect. @@ -208,18 +341,28 @@ class ServerRole: "P2PSession %s: fetch RECEIVED kv_request_id=%s blocks=%d", self._peer_id, kv_request_id, - len(block_hashes), + len(keys), ) - existing = self._outbound.get(kv_request_id) + st = self._requests.get(kv_request_id) + existing = st.outbound if st is not None else None if existing is not None and existing.demand_received: # A second fetch for the same kv_request_id would overwrite # `remaining` and leak inflight bookkeeping. Treat as a # protocol violation. raise ValueError(f"duplicate fetch for kv_request_id={kv_request_id}") - req = self._outbound.setdefault(kv_request_id, _OutboundRequestState()) - result = req.add_fetch_demand(block_hashes, block_indexes) + st = self._get_or_create_request(kv_request_id) + if st.outbound is None: + st.outbound = _OutboundRequestState() + req = st.outbound + result = req.add_fetch_demand(keys, block_indexes) if result.local_idxs: self._submit_transfer(kv_request_id, result) + # Close the peer's request as far as the server's lookup phase + # is concerned: pop parked lookups so no further + # ``parent.create_store_job`` fires for this id, and queue their + # ctxs for ``parent.on_request_finished``. Done before the + # finalize path below so bookkeeping releases in-order. + self._finish_inbound_lookups(kv_request_id) # Prefiller-first mode: finish_request may have run before # fetch arrived. If so, finalize once we know what was # demanded — fully satisfied → success, else early-fail. @@ -230,9 +373,9 @@ class ServerRole: """Handle an AbortFetchMsg from the peer.""" # Abort for an unknown id may be a benign race/duplicate or a # real protocol violation; we don't track completed ids, so warn. - if kv_request_id not in self._outbound and not self._has_inflight_for( - kv_request_id - ): + st = self._requests.get(kv_request_id) + has_outbound = st is not None and st.outbound is not None + if not has_outbound and not self._has_inflight_for(kv_request_id): logger.warning( "P2PSession %s: abort_fetch for unknown kv_request_id=%s " "(no outbound or inflight state); benign race or stale", @@ -242,9 +385,285 @@ class ServerRole: # Idempotent: receiving AbortFetchMsg again before we've sent the # ack just triggers another drain attempt without resetting the # deadline. - self._pending_aborts.setdefault(kv_request_id, time.monotonic()) + st = self._get_or_create_request(kv_request_id) + if st.abort_started_at is None: + st.abort_started_at = time.monotonic() + self._parked_aborts.add(kv_request_id) self._drain_abort(kv_request_id) + def on_lookup( + self, + kv_request_id: str, + keys: Sequence[OffloadKey], + ) -> None: + """Enqueue a LookupMsg from a symmetric-P2P consumer. + + Dispatch runs during ``session.poll()`` where the + :class:`ParentManager` handle is not available, so this only + records the raw request. It is resolved — querying the tiering + manager and emitting the aggregated ``LookupRespMsg`` — by the + next :meth:`serve_external_requests`, the sole window in which + parent calls are valid. + """ + logger.debug( + "P2P LOOKUP server %s: RECV LookupMsg kv_request_id=%s keys=%d", + self._peer_id, + kv_request_id, + len(keys), + ) + self._get_or_create_request(kv_request_id).pending_lookups.append( + _PendingLookup(keys=list(keys), enqueued_at=time.monotonic()) + ) + self._serve_pending.add(kv_request_id) + + def serve_external_requests(self, parent: ParentManager) -> None: + """Resolve inbound peer lookups against the tiering manager. + + Called once per scheduler step with a ``parent`` handle valid + only for this call. Drains newly-enqueued LookupMsgs, re-polls + any parked HIT_PENDING / RETRY keys, and releases the + bookkeeping for lookups closed since the last serve. + """ + for kv_request_id in list(self._serve_pending): + st = self._requests.get(kv_request_id) + if st is None: + self._serve_pending.discard(kv_request_id) + continue + if st.pending_lookups: + pending = st.pending_lookups + st.pending_lookups = [] + for pl in pending: + self._process_inbound_lookup( + kv_request_id, pl.keys, pl.enqueued_at, parent + ) + self._resolve_pending_lookups(kv_request_id, parent) + st = self._requests.get(kv_request_id) + if st is None or (not st.pending_lookups and not st.lookups): + self._serve_pending.discard(kv_request_id) + self._maybe_prune(kv_request_id) + + if self._finished_lookup_ctxs: + for ctx in self._finished_lookup_ctxs: + parent.on_request_finished(ctx) + self._finished_lookup_ctxs = [] + + def _poll_lookup_keys( + self, + lookup: _ActiveLookup, + keys: Iterable[OffloadKey], + parent: ParentManager, + ) -> list[OffloadKey]: + """Poll ``keys`` against the tiering manager and pin any HITs. + + For each key not already definitively resolved, query + ``parent.lookup`` and record the outcome on ``lookup``: HIT / MISS + land in ``resolved`` and clear ``pending``; HIT_PENDING / RETRY park + in ``pending`` for a later serve. Newly-HIT keys are pinned in one + batch via :meth:`_pin_and_register_hits` and also returned (for + caller logging). + + Shared by the first-sighting pass (:meth:`_process_inbound_lookup`) + and the re-poll pass (:meth:`_resolve_pending_lookups`); callers pass + a de-duplicated ``keys`` collection. + """ + new_hits: list[OffloadKey] = [] + for h in keys: + if h in lookup.resolved: + continue + result = parent.lookup(h, lookup.ctx) + if result is LookupResult.HIT: + new_hits.append(h) + lookup.resolved[h] = True + lookup.pending.discard(h) + elif result is LookupResult.MISS: + lookup.resolved[h] = False + lookup.pending.discard(h) + else: + lookup.pending.add(h) + if new_hits: + self._pin_and_register_hits( + lookup.kv_request_id, new_hits, lookup.ctx, parent + ) + return new_hits + + def _process_inbound_lookup( + self, + kv_request_id: str, + keys: list[OffloadKey], + enqueued_at: float, + parent: ParentManager, + ) -> None: + """Resolve one enqueued LookupMsg against ``parent``. + + For each key, query the tiering manager via ``parent.lookup`` + and pin any HITs immediately via ``parent.create_store_job`` + (plumbed into the existing ``add_stored_blocks`` matching path so + the eventual FetchMsg finds them). HIT_PENDING / RETRY keys park + in :class:`_ActiveLookup` for re-polling by + :meth:`_resolve_pending_lookups`. + + The outbound ``LookupRespMsg`` is deferred until every key has + settled to HIT or MISS (or the batch-level ``deadline`` fires, + forcing any stragglers to MISS). One LookupRespMsg goes out per + LookupMsg — carrying every key in wire order — after which + ``parent.on_request_finished`` fires and the entry is dropped. + """ + self._lookup_id_counter += 1 + lookup_id = self._lookup_id_counter + ctx = ReqContext(req_id=f"p2p:{self._peer_id}:{kv_request_id}:lu{lookup_id}") + lookup = _ActiveLookup( + lookup_id=lookup_id, + kv_request_id=kv_request_id, + ctx=ctx, + keys=list(keys), + deadline=enqueued_at + _LOOKUP_PENDING_TIMEOUT_S, + ) + + # Open per-request bookkeeping for this synthetic ctx before the + # first lookup; released by ``on_request_finished`` once every + # key has settled. + parent.on_new_request(ctx) + + # dict.fromkeys de-duplicates keys within the LookupMsg while + # preserving wire order — each unique key is polled once. + hit_keys = self._poll_lookup_keys(lookup, dict.fromkeys(lookup.keys), parent) + + logger.debug( + "P2P LOOKUP server %s: RESOLVED kv_request_id=%s hits=%d misses=%d " + "pending=%d", + self._peer_id, + kv_request_id, + len(hit_keys), + sum(1 for v in lookup.resolved.values() if not v), + len(lookup.pending), + ) + + if lookup.pending: + self._get_or_create_request(kv_request_id).lookups[lookup_id] = lookup + else: + # Every key resolved on first sight — emit the aggregated + # response now and close the synthetic request. + self._finalize_lookup(lookup, parent) + + def _pin_and_register_hits( + self, + kv_request_id: str, + keys: list[OffloadKey], + ctx: ReqContext, + parent: ParentManager, + ) -> None: + """Pin primary slots for HIT keys and feed them into the + existing ``add_stored_blocks`` matching path. + + Caller has already confirmed every key is HIT (single-threaded + scheduler ⇒ no eviction race), so the JobMetadata returned by + ``parent.create_store_job`` carries parallel ``keys``/``block_ids`` + of length ``len(keys)``. + """ + meta = parent.create_store_job(keys, ctx) + self.add_stored_blocks( + kv_request_id, + list(meta.keys), + list(meta.block_ids), + meta.job_id, + ) + + def _resolve_pending_lookups( + self, kv_request_id: str, parent: ParentManager + ) -> None: + """Re-poll a request's deferred LookupMsg keys; finalize when ready. + + Walks every parked :class:`_ActiveLookup` for ``kv_request_id`` and + re-calls ``parent.lookup`` per still-pending key, moving HIT/MISS + results into ``resolved``. If ``deadline`` has passed, remaining + ``pending`` keys are force-resolved to MISS so the consumer + can fall back instead of waiting on a stuck producer. Newly-HIT + keys are pinned via ``parent.create_store_job`` in one call per + affected lookup. Lookups whose ``pending`` empties out get their + aggregated LookupRespMsg sent by :meth:`_finalize_lookup`. + """ + st = self._requests.get(kv_request_id) + if st is None or not st.lookups: + return + now = time.monotonic() + finished_lookups: list[int] = [] + for lookup_id, lookup in st.lookups.items(): + self._poll_lookup_keys(lookup, list(lookup.pending), parent) + + if lookup.pending and now >= lookup.deadline: + for h in lookup.pending: + lookup.resolved[h] = False + lookup.pending.clear() + + if not lookup.pending: + finished_lookups.append(lookup_id) + + for lookup_id in finished_lookups: + lookup = st.lookups.pop(lookup_id) + self._finalize_lookup(lookup, parent) + + def _finalize_lookup(self, lookup: _ActiveLookup, parent: ParentManager) -> None: + """Emit the aggregated LookupRespMsg and close the synthetic request. + + Called once per lookup when ``pending`` is empty — either every + key resolved to HIT or MISS, or the deadline forced remaining + stragglers to MISS. Preserves the wire order of the inbound + LookupMsg so the client can zip keys and hits positionally. + """ + if lookup.keys: + hits = [lookup.resolved[h] for h in lookup.keys] + n_hit = sum(1 for v in hits if v) + logger.debug( + "P2P LOOKUP server %s: SEND LookupRespMsg kv_request_id=%s " + "keys=%d hits=%d misses=%d", + self._peer_id, + lookup.kv_request_id, + len(lookup.keys), + n_hit, + len(hits) - n_hit, + ) + self._send( + { + TYPE_KEY: LookupRespMsg.TYPE, + LookupRespMsg.KV_REQUEST_ID: lookup.kv_request_id, + LookupRespMsg.KEYS: list(lookup.keys), + LookupRespMsg.HITS: hits, + } + ) + parent.on_request_finished(lookup.ctx) + + def _finish_inbound_lookups(self, kv_request_id: str) -> None: + """Close the server-side lookup phase for ``kv_request_id``. + + Pops every parked ``_ActiveLookup`` for this id (so + ``_resolve_pending_lookups`` cannot promote a HIT_PENDING / + RETRY key into a fresh ``parent.create_store_job`` after this + point) and queues each ``lookup.ctx`` for + ``parent.on_request_finished`` (fired by the next + ``serve_external_requests``, since no parent handle is available + during dispatch) so the TieringManager can release per-lookup + bookkeeping. Any still-unprocessed raw LookupMsg for this id is + dropped — it never got ``on_new_request``, so nothing is owed. + The aggregated LookupRespMsg is skipped — the client already + knows the request is over (it just sent a terminal FetchMsg, or + is finishing locally). + + Called on the two events that mean "no more lookup traffic for + ``kv_request_id`` is expected on this session": the terminal + FetchMsg from the peer (client contract: exactly one FetchMsg + per lookup-touched request, even if empty), and a local + ``finish``. Whichever fires second is a no-op. + """ + st = self._requests.get(kv_request_id) + if st is None: + return + st.pending_lookups.clear() + for lookup in st.lookups.values(): + self._finished_lookup_ctxs.append(lookup.ctx) + st.lookups.clear() + self._serve_pending.discard(kv_request_id) + self._maybe_prune(kv_request_id) + def finish(self, kv_request_id: str) -> None: """Mark an outbound request finishing. @@ -253,13 +672,23 @@ class ServerRole: peer to stop waiting (TransferDoneMsg success=False) instead of letting it hit _LOAD_TIMEOUT_S. + Also drops any in-flight lookups for this kv_request_id and + queues their ctxs for ``parent.on_request_finished`` so the + TieringManager can release per-request bookkeeping. (For + symmetric P2P this path is rarely hit since the producer has no + local request lifecycle for the consumer's id; this cleanup is + mostly active on the PD side.) + If the decoder hasn't sent fetch yet (no demand received), defer — on_fetch will finalize once demand arrives. If inflight transfers exist for this id, defer — the last completing transfer in collect_results will fire the message. """ - req = self._outbound.get(kv_request_id) + self._finish_inbound_lookups(kv_request_id) + + st = self._requests.get(kv_request_id) + req = st.outbound if st is not None else None if req is None: return req.finishing = True @@ -268,14 +697,19 @@ class ServerRole: if self._has_inflight_for(kv_request_id): return # Remaining > 0 here: if it had hit 0, the poll-done success - # branch would have already popped _outbound and we'd have + # branch would have already cleared outbound and we'd have # returned at `req is None` above. Helper derives success from # remaining and emits StoreResult(success=False) for any # leftover pending jobs. self._finalize_outbound(kv_request_id) def collect_results(self) -> list[StoreResult]: - """Drain timeouts, deferred results, and transport completions.""" + """Drain timeouts, deferred results, and transport completions. + + Inbound LookupMsg resolution (including re-polling HIT_PENDING / + RETRY keys) is NOT done here — it runs in + ``serve_external_requests`` where the ParentManager is available. + """ results: list[StoreResult] = self._timeout_pending_store_jobs() if self._pending_store_results: @@ -306,15 +740,9 @@ class ServerRole: tid, ) continue - req = self._outbound.get(xfer.kv_request_id) - for job_id in xfer.job_ids: - if self._store_jobs.pop(job_id, None) is None: - # Already reported (timeout, cancellation, etc.) — - # don't double-emit a contradictory success result. - continue - results.append(StoreResult(job_id=job_id, success=True)) - if req is not None: - req.pending_job_ids.discard(job_id) + results.extend(self._settle_xfer_jobs(xfer, success=True)) + st = self._requests.get(xfer.kv_request_id) + req = st.outbound if st is not None else None if req is not None and req.demand_received: req.remaining -= xfer.block_count assert req.remaining >= 0, ( @@ -324,6 +752,7 @@ class ServerRole: self._finalize_outbound(xfer.kv_request_id, success=True) elif req.finishing and not self._has_inflight_for(xfer.kv_request_id): self._finalize_outbound(xfer.kv_request_id, success=False) + self._maybe_prune(xfer.kv_request_id) failed_kv_request_ids: set[str] | None = None for tid in poll_result.failed: @@ -341,16 +770,11 @@ class ServerRole: if failed_kv_request_ids is None: failed_kv_request_ids = set() failed_kv_request_ids.add(xfer.kv_request_id) - req_for_xfer = self._outbound.get(xfer.kv_request_id) - for job_id in xfer.job_ids: - if self._store_jobs.pop(job_id, None) is None: - # Already reported (timeout, cancellation, etc.) — - # don't double-emit. - continue - results.append(StoreResult(job_id=job_id, success=False)) - if req_for_xfer is not None: - req_for_xfer.pending_job_ids.discard(job_id) - req = self._outbound.pop(xfer.kv_request_id, None) + results.extend(self._settle_xfer_jobs(xfer, success=False)) + st = self._requests.get(xfer.kv_request_id) + req = st.outbound if st is not None else None + if st is not None: + st.outbound = None if req is not None and req.demand_received: self._send( { @@ -359,6 +783,7 @@ class ServerRole: TransferDoneMsg.SUCCESS: False, } ) + self._maybe_prune(xfer.kv_request_id) # Cancel other inflight for the same failed kv_request_ids if failed_kv_request_ids: @@ -370,6 +795,8 @@ class ServerRole: for tid in ids_to_cancel: self._inflight_pop(tid) self._transport.cancel(ids_to_cancel) + for kv_request_id in failed_kv_request_ids: + self._maybe_prune(kv_request_id) return results @@ -384,54 +811,90 @@ class ServerRole: def drain_pending_aborts(self) -> None: """Re-attempt every parked abort once per poll tick.""" - if not self._pending_aborts: - return - for kv_request_id in list(self._pending_aborts): + for kv_request_id in list(self._parked_aborts): self._drain_abort(kv_request_id) - def close(self) -> list[int]: - """Tear down. Cancels inflight, returns failed store job ids.""" + def close(self) -> tuple[list[int], list[ReqContext]]: + """Tear down. Cancels inflight. + + Returns ``(failed_store_job_ids, failed_serves)`` where + ``failed_serves`` are synthetic lookup ctxs still owing a + ``parent.on_request_finished``. The session is going away with no + parent handle in hand, so the manager flushes these in its next + ``serve_external_requests``. + """ failed_stores = list(self._store_jobs.keys()) self._store_jobs.clear() if self._inflight: self._transport.cancel(list(self._inflight.keys())) self._inflight.clear() - self._inflight_per_req.clear() - self._outbound.clear() - self._pending_aborts.clear() self._pending_store_results.clear() - return failed_stores + # Surface every synthetic ctx still owing on_request_finished so + # the manager can release the TieringManager's per-request + # bookkeeping: parked lookups plus any already queued from a + # FetchMsg / finish that closed them before this teardown. + failed_serves = [ + lu.ctx for st in self._requests.values() for lu in st.lookups.values() + ] + failed_serves.extend(self._finished_lookup_ctxs) + self._requests.clear() + self._serve_pending.clear() + self._parked_aborts.clear() + self._finished_lookup_ctxs.clear() + return failed_stores, failed_serves + + @property + def has_inflight_transfers(self) -> bool: + """True if any outbound store transfer is still in flight.""" + return bool(self._inflight) # ------------------------------------------------------------------ # Internal — inflight bookkeeping # ------------------------------------------------------------------ def _has_inflight_for(self, kv_request_id: str) -> bool: - return kv_request_id in self._inflight_per_req + st = self._requests.get(kv_request_id) + return st is not None and bool(st.inflight_tids) def _inflight_add(self, tid: int, xfer: _InflightXfer) -> None: - """Insert an inflight transfer and bump the per-request count.""" + """Insert an inflight transfer and record it on the request.""" self._inflight[tid] = xfer - self._inflight_per_req[xfer.kv_request_id] = ( - self._inflight_per_req.get(xfer.kv_request_id, 0) + 1 - ) + self._get_or_create_request(xfer.kv_request_id).inflight_tids.add(tid) def _inflight_pop(self, tid: int) -> _InflightXfer | None: - """Pop an inflight transfer and decrement the per-request count. + """Pop an inflight transfer and drop it from the request's set. - Removes the per-request entry once the count hits zero so the - dict stays bounded and `_has_inflight_for` remains exact. + Callers are responsible for the ``_maybe_prune`` that may follow once + the request's other state has also cleared. """ xfer = self._inflight.pop(tid, None) if xfer is None: return None - new_count = self._inflight_per_req.get(xfer.kv_request_id, 0) - 1 - if new_count > 0: - self._inflight_per_req[xfer.kv_request_id] = new_count - else: - self._inflight_per_req.pop(xfer.kv_request_id, None) + st = self._requests.get(xfer.kv_request_id) + if st is not None: + st.inflight_tids.discard(tid) return xfer + def _settle_xfer_jobs( + self, xfer: _InflightXfer, success: bool + ) -> list[StoreResult]: + """Emit StoreResults for a completed transfer's store jobs. + + Pops each attached job from ``_store_jobs`` and clears it from the + request's pending set. A job already popped (via timeout, cancel, + etc.) is skipped so we never double-emit a contradictory result. + """ + results: list[StoreResult] = [] + st = self._requests.get(xfer.kv_request_id) + req = st.outbound if st is not None else None + for job_id in xfer.job_ids: + if self._store_jobs.pop(job_id, None) is None: + continue + results.append(StoreResult(job_id=job_id, success=success)) + if req is not None: + req.pending_job_ids.discard(job_id) + return results + # ------------------------------------------------------------------ # Internal — finalize / abort drain # ------------------------------------------------------------------ @@ -452,9 +915,10 @@ class ServerRole: The same flag is used for both the peer's TransferDoneMsg and the StoreResult(s) emitted for any leftover pending job_ids. """ - req = self._outbound.pop(kv_request_id, None) - if req is None: - return + st = self._requests[kv_request_id] + assert st.outbound is not None + req = st.outbound + st.outbound = None if success is None: success = req.remaining == 0 for job_id in req.pending_job_ids: @@ -469,6 +933,7 @@ class ServerRole: TransferDoneMsg.SUCCESS: success, } ) + self._maybe_prune(kv_request_id) def _drain_abort(self, kv_request_id: str) -> None: """One drain attempt for a pending abort. @@ -479,21 +944,15 @@ class ServerRole: inflight, or after ``_CANCEL_DRAIN_TIMEOUT_S`` falls back to ``mode="immediate"`` and acks anyway. """ - self._outbound.pop(kv_request_id, None) - ids = [ - tid - for tid, xfer in self._inflight.items() - if xfer.kv_request_id == kv_request_id - ] + st = self._requests[kv_request_id] + st.outbound = None + ids = list(st.inflight_tids) if not ids: self._finalize_abort(kv_request_id) return - started_at = self._pending_aborts.get(kv_request_id) - expired = ( - started_at is not None - and time.monotonic() - started_at >= _CANCEL_DRAIN_TIMEOUT_S - ) + assert st.abort_started_at is not None + expired = time.monotonic() - st.abort_started_at >= _CANCEL_DRAIN_TIMEOUT_S if expired: for tid in ids: self._inflight_pop(tid) @@ -521,13 +980,16 @@ class ServerRole: self._finalize_abort(kv_request_id) def _finalize_abort(self, kv_request_id: str) -> None: - self._pending_aborts.pop(kv_request_id, None) + st = self._requests[kv_request_id] + st.abort_started_at = None + self._parked_aborts.discard(kv_request_id) self._send( { TYPE_KEY: AbortAckMsg.TYPE, AbortAckMsg.KV_REQUEST_ID: kv_request_id, } ) + self._maybe_prune(kv_request_id) # ------------------------------------------------------------------ # Internal — transfers and store-job timeouts @@ -579,7 +1041,8 @@ class ServerRole: # nothing else is in flight, finalize now so the peer # and the local store jobs don't wait for finish_request # or for _STORE_TIMEOUT_S / _LOAD_TIMEOUT_S. - req = self._outbound.get(kv_request_id) + st = self._requests.get(kv_request_id) + req = st.outbound if st is not None else None if req is not None: req.finishing = True if not self._has_inflight_for(kv_request_id): diff --git a/vllm/v1/kv_offload/tiering/p2p/session/session.py b/vllm/v1/kv_offload/tiering/p2p/session/session.py index fc5f1feab58..7d19913b0a2 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/session.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/session.py @@ -34,6 +34,8 @@ from vllm.v1.kv_offload.tiering.p2p.session.protocol import ( ConnectMsg, DisconnectMsg, FetchMsg, + LookupMsg, + LookupRespMsg, TransferDoneMsg, ) from vllm.v1.kv_offload.tiering.p2p.session.server import ( @@ -42,7 +44,8 @@ from vllm.v1.kv_offload.tiering.p2p.session.server import ( ) if TYPE_CHECKING: - from vllm.v1.kv_offload.tiering.base import JobId + from vllm.v1.kv_offload.base import ReqContext + from vllm.v1.kv_offload.tiering.base import JobId, ParentManager from vllm.v1.kv_offload.tiering.p2p.data import DataTransport logger = init_logger(__name__) @@ -71,6 +74,22 @@ class SessionPollResult(NamedTuple): new_fetch_ids: list[str] +class SessionCloseResult(NamedTuple): + """Result of tearing down a P2PSession. + + `failed_jobs`/`failed_stores` are the in-flight jobs (client loads / + server stores) the manager must mark failed. `failed_req_ids` is every + client-side kv_request_id whose lookup() must fail (in-flight loads plus + unresolved probes); `failed_serves` is the server-side lookup state the + dead peer can no longer resolve. + """ + + failed_jobs: list[int] # client load job_ids + failed_req_ids: list[str] # client kv_request_ids (loads + probes) + failed_stores: list[int] # server store job_ids + failed_serves: list[ReqContext] # server-side lookup ctxs needing release + + class P2PSession: """Bidirectional session — coordinator over ClientRole + ServerRole. @@ -93,12 +112,14 @@ class P2PSession: local_id: str, transport: DataTransport, local_block_len: int, + local_hash_seed: str, conn: ControlConnection | None = None, ) -> None: self.peer_id = peer_id self._local_id = local_id self._transport = transport self._local_block_len = local_block_len + self._local_hash_seed = local_hash_seed self._conn: ControlConnection | None = None self._send_ready = False # True after the peer acked our ConnectMsg @@ -115,7 +136,11 @@ class P2PSession: self._new_fetch_ids: list[str] = [] self._client = ClientRole(peer_id=peer_id, send=self._send) - self._server = ServerRole(peer_id=peer_id, transport=transport, send=self._send) + self._server = ServerRole( + peer_id=peer_id, + transport=transport, + send=self._send, + ) if conn is not None: self.attach_connection(conn) @@ -139,6 +164,11 @@ class P2PSession: """True after the peer acked our ConnectMsg (we may send freely).""" return self._send_ready + @property + def has_pending_work(self) -> bool: + """True while inbound loads or outbound transfers are outstanding.""" + return self._client.has_active_loads or self._server.has_inflight_transfers + # ------------------------------------------------------------------ # Connection lifecycle # ------------------------------------------------------------------ @@ -162,7 +192,7 @@ class P2PSession: self, job_id: JobId, kv_request_id: str, - keys: Sequence[bytes], + keys: Sequence[OffloadKey], block_ids: Sequence[int], ) -> None: """Send fetch to the peer.""" @@ -183,13 +213,40 @@ class P2PSession: def finish_request(self, kv_request_id: str) -> None: """Called when the request is finishing locally. - Cancels any inbound load (client role) and finalizes any - outbound serving (server role) for this id. Roles that aren't - active for this id are silent no-ops. + Finishes the client role (aborts any inbound load and drops any + pending symmetric-P2P lookup state) and finalizes any outbound + serving (server role) for this id. Roles that aren't active for + this id are silent no-ops. """ - self._client.cancel(kv_request_id) + self._client.finish(kv_request_id) self._server.finish(kv_request_id) + def register_lookup(self, kv_request_id: str, key: bytes) -> bool | None: + """Register or resolve one (kv_request_id, key) probe. + + Called from the manager's lookup() for symmetric-P2P consumers + (``remote_kv_source`` sub-dict in kv_transfer_params). See + ``ClientRole.register_lookup`` for the state-machine contract. + """ + return self._client.register_lookup(kv_request_id, key) + + def flush_pending_lookups(self) -> None: + """Flush any aggregated symmetric-P2P lookups for this peer. + + Called once per scheduler step from the manager's + ``on_schedule_end()``. Send-gating is handled inside the + client's ``_send`` callback (queues until ConnectAckMsg). + """ + self._client.flush_pending_lookups() + + def serve_external_requests(self, parent: ParentManager) -> None: + """Resolve inbound peer lookups against the tiering manager. + + Delegates to the server role; the ``parent`` handle is valid + only for the duration of this call. + """ + self._server.serve_external_requests(parent) + def poll(self) -> SessionPollResult: """Process incoming messages, drive transfers, apply timeouts.""" if self._conn is None: @@ -214,14 +271,22 @@ class P2PSession: loads=loads, stores=stores, new_fetch_ids=new_fetch_ids ) - def close(self) -> tuple[list[tuple[int, str]], list[int]]: - """Shut down. Returns (failed_loads, failed_stores). + def close(self) -> SessionCloseResult: + """Shut down. - failed_loads: list of (job_id, kv_request_id) pairs. - failed_stores: list of job_ids. + failed_jobs: client load job_ids to fail. + failed_req_ids: client kv_request_ids to fail — in-flight loads plus + requests with an unresolved symmetric-P2P probe toward the + now-dead peer. The manager fails these so the consumer's lookup() + falls back to local prefill instead of deferring forever on an + answer that can never arrive. + failed_stores: server store job_ids to fail. + failed_serves: synthetic lookup ctxs still owing + ``parent.on_request_finished`` (the manager flushes these on + its next ``serve_external_requests``). """ - failed_loads = self._client.close() - failed_stores = self._server.close() + client_result = self._client.close() + failed_stores, failed_serves = self._server.close() if self._conn is not None: with contextlib.suppress(Exception): @@ -229,7 +294,12 @@ class P2PSession: self._conn.close() self._conn = None - return failed_loads, failed_stores + return SessionCloseResult( + failed_jobs=client_result.failed_jobs, + failed_req_ids=client_result.failed_req_ids, + failed_stores=failed_stores, + failed_serves=failed_serves, + ) # ------------------------------------------------------------------ # Message dispatch @@ -299,9 +369,9 @@ class P2PSession: elif msg_type == FetchMsg.TYPE: FetchMsg.validate(msg) kv_request_id = msg[FetchMsg.KV_REQUEST_ID] - block_hashes = [ + keys = [ OffloadKey(bh if isinstance(bh, bytes) else bytes(bh)) - for bh in msg[FetchMsg.BLOCK_HASHES] + for bh in msg[FetchMsg.KEYS] ] block_indexes = msg[FetchMsg.BLOCK_INDEXES] # Run the server-role state machine inline as today — @@ -310,7 +380,7 @@ class P2PSession: # the manager (after poll() returns) can replay any parked # submit_store batches; their add_stored_blocks calls hit # the demand recorded here and submit transfers immediately. - self._server.on_fetch(kv_request_id, block_hashes, block_indexes) + self._server.on_fetch(kv_request_id, keys, block_indexes) self._new_fetch_ids.append(kv_request_id) elif msg_type == AbortFetchMsg.TYPE: AbortFetchMsg.validate(msg) @@ -324,6 +394,23 @@ class P2PSession: elif msg_type == AbortAckMsg.TYPE: AbortAckMsg.validate(msg) self._client.on_abort_ack(msg[AbortAckMsg.KV_REQUEST_ID]) + elif msg_type == LookupMsg.TYPE: + LookupMsg.validate(msg) + kv_request_id = msg[LookupMsg.KV_REQUEST_ID] + keys = [ + OffloadKey(bh if isinstance(bh, bytes) else bytes(bh)) + for bh in msg[LookupMsg.KEYS] + ] + self._server.on_lookup(kv_request_id, keys) + elif msg_type == LookupRespMsg.TYPE: + LookupRespMsg.validate(msg) + kv_request_id = msg[LookupRespMsg.KV_REQUEST_ID] + keys = [ + OffloadKey(bh if isinstance(bh, bytes) else bytes(bh)) + for bh in msg[LookupRespMsg.KEYS] + ] + hits = msg[LookupRespMsg.HITS] + self._client.on_lookup_resp(kv_request_id, keys, hits) elif msg_type == DisconnectMsg.TYPE: if self._conn is not None: self._conn.mark_dead() @@ -361,6 +448,12 @@ class P2PSession: f"config fingerprint mismatch from {self.peer_id}: " f"remote={remote_fp!r}, local={local_fp!r}" ) + if msg[ConnectMsg.HASH_SEED] != self._local_hash_seed: + raise ValueError( + f"PYTHONHASHSEED mismatch from {self.peer_id}: " + f"remote={msg[ConnectMsg.HASH_SEED]!r}, " + f"local={self._local_hash_seed!r}" + ) self._transport.add_remote_peer( self.peer_id, agent_metadata=msg[ConnectMsg.AGENT_METADATA], @@ -410,6 +503,7 @@ class P2PSession: ConnectMsg.NUM_BLOCKS: self._transport.num_blocks, ConnectMsg.BLOCK_LEN: self._transport.block_len, ConnectMsg.CONFIG_FINGERPRINT: self._transport.config_fingerprint, + ConnectMsg.HASH_SEED: self._local_hash_seed, } ) @@ -433,8 +527,15 @@ class P2PSession: self._conn.send(msg) logger.debug("P2PSession %s: sent %s", self.peer_id, msg.get(TYPE_KEY)) except Exception: + # A send failure means the connection is broken. Swallowing it + # silently strands every in-flight lookup/load toward this peer: + # the session stays alive, is never reaped, and the consumer's + # lookup() keeps returning RETRY until the HTTP client times out. + # Mark the connection dead so the manager reaps the session on + # its next poll and surfaces the stranded work as failures. logger.warning( - "P2PSession %s: failed to send %s", + "P2PSession %s: send of %s failed — marking connection dead", self.peer_id, msg.get(TYPE_KEY), ) + self._conn.mark_dead() From 5559679229bc961848b121ccdeaa8fa5d79bec98 Mon Sep 17 00:00:00 2001 From: coltonottley <coltonottley@gmail.com> Date: Sun, 26 Jul 2026 05:53:32 -0600 Subject: [PATCH 064/185] [Bugfix][KV Offload] Bound unaligned SWA loads by physical GPU blocks (#49052) Signed-off-by: Colton Ottley <colton@ottleyengineering.com> Co-authored-by: Colton Ottley <colton@ottleyengineering.com> Co-authored-by: jasl <jasl9187@hotmail.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: Or Ozeri <oro@il.ibm.com> --- .../kv_transfer/kv_connector/v1/offloading/scheduler.py | 1 + 1 file changed, 1 insertion(+) 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 b53a365abd4..6ccc820f253 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -860,6 +860,7 @@ class OffloadingConnectorScheduler: num_pending_gpu_blocks <= group_config.sliding_window_size_in_chunks * self.config.blocks_per_chunk + + 1 ) num_chunks = cdiv(num_cached_tokens, tokens_per_chunk) From 0da6e7f3d6ed50fe1262cc2cf44066989a6b4cdd Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim <taneem.ibrahim@gmail.com> Date: Sun, 26 Jul 2026 08:42:14 -0400 Subject: [PATCH 065/185] [Bugfix] Reject contradictory custom-op directives (#49134) Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> --- tests/compile/test_config.py | 28 ++++++++++++++++++++++++++-- vllm/config/compilation.py | 31 ++++++++++++++++++++++++------- vllm/model_executor/custom_op.py | 11 +++++++++-- 3 files changed, 59 insertions(+), 11 deletions(-) diff --git a/tests/compile/test_config.py b/tests/compile/test_config.py index d822b68c503..c45907ee6e1 100644 --- a/tests/compile/test_config.py +++ b/tests/compile/test_config.py @@ -83,9 +83,33 @@ def test_copy_pass(): def test_custom_op(): # proper syntax _ = CompilationConfig(custom_ops=["+quant_fp8", "-silu_and_mul"]) + _ = CompilationConfig(custom_ops=["none", "+rms_norm"]) + _ = CompilationConfig(custom_ops=["+rms_norm", "+rms_norm"]) - with pytest.raises(ValueError, match="Invalid syntax '"): - _ = CompilationConfig(custom_ops=["quant_fp8"]) + for custom_ops in (["quant_fp8"], ["+"], ["-"]): + with pytest.raises(ValueError, match="Invalid syntax '"): + CompilationConfig(custom_ops=custom_ops) + + +@pytest.mark.parametrize( + ("custom_ops", "config_kwargs", "match"), + [ + (["all", "none"], {}, "can contain only one base mode"), + ( + ["none", "+rms_norm", "-rms_norm"], + {}, + "cannot both enable and disable.*rms_norm", + ), + ( + ["-rotary_embedding"], + {"pass_config": PassConfig(enable_qk_norm_rope_fusion=True)}, + "cannot both enable and disable.*rotary_embedding", + ), + ], +) +def test_reject_contradictory_custom_ops(custom_ops, config_kwargs, match): + with pytest.raises(ValueError, match=match): + CompilationConfig(custom_ops=custom_ops, **config_kwargs) # forked needed to workaround https://github.com/vllm-project/vllm/issues/21073 diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 8142004748a..fbc5dcff40e 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -904,10 +904,6 @@ class CompilationConfig: return handler(value) def __post_init__(self) -> None: - count_none = self.custom_ops.count("none") - count_all = self.custom_ops.count("all") - assert count_none + count_all <= 1, "Can only specify 'none' or 'all'" - # TODO(zou3519/luka): There are 2 issues with auto-functionalization V2: # 1. A bug in PyTorch, fixed in 2.7: # https://github.com/pytorch/pytorch/issues/147924 @@ -1000,13 +996,28 @@ class CompilationConfig: ) for op in self.custom_ops: - if op[0] not in {"+", "-"} and op not in {"all", "none"}: + if op not in {"all", "none"} and (len(op) < 2 or op[0] not in {"+", "-"}): raise ValueError( f"Invalid syntax '{op}' for custom op, " "must be 'all', 'none', '+op' or '-op' " "(where 'op' is the registered op name)" ) + base_modes = [op for op in self.custom_ops if op in {"all", "none"}] + if len(base_modes) > 1: + raise ValueError( + "custom_ops can contain only one base mode: 'all' or 'none'" + ) + + enabled_ops = {op[1:] for op in self.custom_ops if op.startswith("+")} + disabled_ops = {op[1:] for op in self.custom_ops if op.startswith("-")} + conflicting_ops = sorted(enabled_ops & disabled_ops) + if conflicting_ops: + raise ValueError( + "custom_ops cannot both enable and disable the same operation(s): " + f"{', '.join(conflicting_ops)}. Remove either the '+' or '-' directive" + ) + # Currently only eager and inductor backend are supported. # for piecewise compilation. Custom backends are not supported for # piecewise compilation. Update when more backends are supported. @@ -1343,10 +1354,16 @@ class CompilationConfig: ) def is_custom_op_enabled(self, op: str) -> bool: - if "all" in self.custom_ops: + count_all = self.custom_ops.count("all") + count_none = self.custom_ops.count("none") + if count_all + count_none != 1: + raise ValueError( + "custom_ops must contain exactly one base mode: 'all' or 'none'" + ) + + if count_all: return f"-{op}" not in self.custom_ops - assert "none" in self.custom_ops return f"+{op}" in self.custom_ops def resolve_cudagraph_mode_and_sizes( diff --git a/vllm/model_executor/custom_op.py b/vllm/model_executor/custom_op.py index a1514c9206b..bf98e46e854 100644 --- a/vllm/model_executor/custom_op.py +++ b/vllm/model_executor/custom_op.py @@ -284,7 +284,11 @@ class CustomOp(nn.Module): enabled = f"+{cls.name}" in custom_ops disabled = f"-{cls.name}" in custom_ops - assert not (enabled and disabled), f"Cannot enable and disable {cls.name}" + if enabled and disabled: + raise ValueError( + "custom_ops cannot both enable and disable the same operation: " + f"{cls.name}. Remove either the '+' or '-' directive" + ) return (CustomOp.default_on() or enabled) and not disabled @@ -299,7 +303,10 @@ class CustomOp(nn.Module): compilation_config = get_cached_compilation_config() count_none = compilation_config.custom_ops.count("none") count_all = compilation_config.custom_ops.count("all") - assert count_none + count_all == 1 + if count_none + count_all != 1: + raise ValueError( + "custom_ops must contain exactly one base mode: 'all' or 'none'" + ) return not count_none > 0 or count_all > 0 From 3f1d40960fb79e6f1314755abf2d43d142e33363 Mon Sep 17 00:00:00 2001 From: AlexHuang <jihui.huang@daocloud.io> Date: Sun, 26 Jul 2026 21:22:58 +0800 Subject: [PATCH 066/185] [KV Offload] Fix num_tokens_after_batch for different termination types (#49285) Signed-off-by: Alex <jihuihuang@example.com> Signed-off-by: Alex <jihui.huang@daocloud.io> Signed-off-by: Alex <alex.tech.lab@outlook.com> Signed-off-by: Alex <jihuihuang@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Or Ozeri <oro@il.ibm.com> --- .../kv_transfer/kv_connector/v1/offloading/scheduler.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) 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 6ccc820f253..f809dca5ba4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -50,7 +50,7 @@ from vllm.v1.kv_offload.base import ( make_offload_key, ) from vllm.v1.outputs import KVConnectorOutput -from vllm.v1.request import Request +from vllm.v1.request import Request, RequestStatus logger = init_logger(__name__) @@ -976,7 +976,9 @@ class OffloadingConnectorScheduler: continue req = req_status.req - if req.is_finished(): + if req.status is RequestStatus.FINISHED_ABORTED: + num_tokens_after_batch = req.num_computed_tokens + elif req.is_finished(): num_tokens_after_batch = req.num_tokens else: num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id] From 7154856f3dcb1d3fdd5a136f7d2c5987f22244f5 Mon Sep 17 00:00:00 2001 From: Daniel Socek <daniel.socek@intel.com> Date: Sun, 26 Jul 2026 10:22:42 -0400 Subject: [PATCH 067/185] [Bugfix] Fix handling 5D KV cache in kv_postprocess_layout_on_receive (#47791) Signed-off-by: Daniel Socek <daniel.socek@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com> --- vllm/distributed/kv_transfer/kv_connector/utils.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py index ac8744d8c46..eff9bec8ee9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/utils.py @@ -258,8 +258,12 @@ def kv_postprocess_layout_on_receive(cache, indices): This method corrects layout mismatches from direct memory copies by permuting the tensor dimensions. + 4D cache: - **Source Layout:** `[num_blocks, n_kv_head, block_size, head_dim]` - **Target Layout:** `[num_blocks, block_size, n_kv_head, head_dim]` + 5D cache: + - **Source Layout:** `[num_blocks, kv_dim, n_kv_head, block_size, head_dim]` + - **Target Layout:** `[num_blocks, kv_dim, block_size, n_kv_head, head_dim]` Implementation: - x = blocks_to_update.reshape(src_shape) # view local kv with sender layout @@ -270,7 +274,7 @@ def kv_postprocess_layout_on_receive(cache, indices): blocks_to_update = cache.index_select(0, indices) target_shape = list(blocks_to_update.shape) target_shape[0] = -1 - inv_order = [0, 2, 1, 3] + inv_order = [0, 1, 3, 2, 4] if blocks_to_update.ndim == 5 else [0, 2, 1, 3] src_shape = tuple(target_shape[i] for i in inv_order) blocks_to_update = cache.index_select(0, indices) permuted_blocks = blocks_to_update.reshape(src_shape).permute(*inv_order) From b68d7ef2622d2d22e964dd842381021865e942b8 Mon Sep 17 00:00:00 2001 From: Jonguk Cheong <jdal3031@snu.ac.kr> Date: Mon, 27 Jul 2026 02:59:09 +0900 Subject: [PATCH 068/185] [Bugfix][KV Offload] Namespace auto cache dtype by effective dtype (#49438) Signed-off-by: Jonguk Cheong <jdal3031@snu.ac.kr> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Or Ozeri <oro@il.ibm.com> --- .../kv_transfer/kv_connector/v1/offloading/config.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index 2a5659c2a9b..b86ebf96bb6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -138,6 +138,12 @@ def build_offloading_config( ) kv_events_config = vllm_config.kv_events_config + cache_dtype = ( + vllm_config.model_config.dtype + if vllm_config.cache_config.cache_dtype == "auto" + else vllm_config.cache_config.cache_dtype + ) + return OffloadingConfig( groups=groups, worker_kv_bytes_per_block=worker_kv_bytes_per_block, @@ -148,7 +154,7 @@ def build_offloading_config( engine_id=engine_id, model=OffloadingModelConfig( name=vllm_config.model_config.model, - dtype=str(vllm_config.cache_config.cache_dtype).replace("torch.", ""), + dtype=str(cache_dtype).removeprefix("torch."), ), cache=OffloadingCacheConfig( tokens_per_hash=tokens_per_hash, From b5b61c622c941426cf78a83672e878b5b829fe90 Mon Sep 17 00:00:00 2001 From: Schwinn Saereesitthipitak <17022745+galletas1712@users.noreply.github.com> Date: Sun, 26 Jul 2026 11:47:50 -0700 Subject: [PATCH 069/185] [Core][Distributed] Add process-checkpoint lifecycle hooks for communicators (starting with Flashinfer) (#46877) Signed-off-by: Schwinn Saereesitthipitak <schwinns@nvidia.com> --- tests/distributed/test_comm_ops.py | 40 ++++++++++++++++++ .../passes/fusion/allreduce_rms_fusion.py | 4 +- .../device_communicators/all2all.py | 14 +++++++ .../base_device_communicator.py | 21 ++++++++++ .../device_communicators/cuda_communicator.py | 16 +++++++ .../flashinfer_all_reduce.py | 42 ++++++++++++++++++- vllm/distributed/parallel_state.py | 36 ++++++++++++++++ .../layers/fused_allreduce_gemma_rms_norm.py | 2 +- vllm/v1/engine/async_llm.py | 6 +++ vllm/v1/worker/gpu_worker.py | 8 ++++ 10 files changed, 185 insertions(+), 4 deletions(-) diff --git a/tests/distributed/test_comm_ops.py b/tests/distributed/test_comm_ops.py index 48b007da664..19a095c7c6f 100644 --- a/tests/distributed/test_comm_ops.py +++ b/tests/distributed/test_comm_ops.py @@ -7,6 +7,7 @@ Run `pytest tests/distributed/test_comm_ops.py`. from collections.abc import Callable from typing import Any +from unittest.mock import Mock import pytest import ray @@ -19,6 +20,8 @@ from vllm.distributed import ( tensor_model_parallel_all_reduce, tensor_model_parallel_reduce_scatter, ) +from vllm.distributed.device_communicators import flashinfer_all_reduce +from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator from vllm.distributed.parallel_state import GroupCoordinator, TensorMetadata from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors @@ -278,6 +281,43 @@ def test_irecv_tensor_dict_send_allgather_postprocess_binds_keys( torch.testing.assert_close(td["b"], torch.ones(4, dtype=torch.int32)) +@pytest.mark.parametrize("aliased", [False, True]) +def test_cuda_communicator_checkpoints_flashinfer_workspaces( + monkeypatch: pytest.MonkeyPatch, + aliased: bool, +) -> None: + group = object() + normal_workspace = Mock() + quant_workspace = normal_workspace if aliased else Mock() + unique_workspaces = ( + [normal_workspace] if aliased else [normal_workspace, quant_workspace] + ) + + monkeypatch.setattr(flashinfer_all_reduce, "_fi_ar_workspace", normal_workspace) + monkeypatch.setattr( + flashinfer_all_reduce, "_fi_ar_quant_workspace", quant_workspace + ) + monkeypatch.setattr( + flashinfer_all_reduce, + "_fi_ar_workspace_groups", + {id(workspace): group for workspace in unique_workspaces}, + ) + monkeypatch.setattr( + flashinfer_all_reduce, "TorchDistBackend", lambda group: group, raising=False + ) + + communicator = CudaCommunicator.__new__(CudaCommunicator) + communicator.cpu_group = group + communicator.fi_ar_comm = None + communicator.all2all_manager = None + communicator.checkpoint_prepare() + communicator.checkpoint_restore() + + for workspace in unique_workspaces: + workspace.checkpoint_prepare.assert_called_once_with() + workspace.checkpoint_restore.assert_called_once_with(group) + + def test_async_intermediate_tensors_lazy_wait() -> None: work = _DummyWork() post_calls = {"n": 0} diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index 4716e7012f3..f7ff7df66cd 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -225,7 +225,7 @@ if flashinfer_comm is not None: max_token_num=max_token_num, hidden_dim=hidden_size, dtype=allreduce_in.dtype, - group=get_tp_group().device_group, + group=get_tp_group().cpu_group, ) assert workspace is not None, ( "Flashinfer allreduce workspace must be initialized when using flashinfer" @@ -996,7 +996,7 @@ class AllReduceFusionPass(VllmPatternMatcherPass): ) return self.hidden_dim = config.model_config.get_hidden_size() - self.group = get_tp_group().device_group + self.group = get_tp_group().cpu_group rank = get_tensor_model_parallel_rank() if flashinfer_comm is None: logger.warning( diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index ee404f688a1..5abe7568a29 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -885,6 +885,20 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase): self.mapping = None self.initialized = False + def checkpoint_prepare(self) -> None: + if self.initialized: + assert self.moe_alltoall is not None + self.moe_alltoall.checkpoint_prepare() + + def checkpoint_restore(self) -> None: + if self.initialized: + assert self.moe_alltoall is not None + from vllm.distributed.device_communicators.mnnvl_compat import ( + CustomCommunicator, + ) + + self.moe_alltoall.checkpoint_restore(CustomCommunicator(self.cpu_group)) + class MoriAll2AllManager(All2AllManagerBase): def __init__(self, cpu_group, all2all_backend: str): diff --git a/vllm/distributed/device_communicators/base_device_communicator.py b/vllm/distributed/device_communicators/base_device_communicator.py index dc267143389..45438a54691 100644 --- a/vllm/distributed/device_communicators/base_device_communicator.py +++ b/vllm/distributed/device_communicators/base_device_communicator.py @@ -7,8 +7,11 @@ import torch import torch.distributed as dist from torch.distributed import ProcessGroup +from vllm.logger import init_logger from vllm.utils import is_moe_layer +logger = init_logger(__name__) + class Cache: def __init__(self): @@ -137,6 +140,18 @@ class All2AllManagerBase: def max_sms_used(self) -> int | None: return None # None means it could use the whole GPU + def checkpoint_prepare(self) -> None: + logger.warning_once( + "%s.checkpoint_prepare is not implemented; skipping.", + type(self).__name__, + ) + + def checkpoint_restore(self) -> None: + logger.warning_once( + "%s.checkpoint_restore is not implemented; skipping.", + type(self).__name__, + ) + def combine(self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False): raise NotImplementedError @@ -216,6 +231,12 @@ class DeviceCommunicatorBase: dist.all_reduce(input_, group=self.device_group) return input_ + def checkpoint_prepare(self) -> None: + """Prepare reclaimable communicator state for checkpoint (default: no-op).""" + + def checkpoint_restore(self) -> None: + """Restore communicator state after checkpoint (default: no-op).""" + def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: if dim < 0: # Convert negative dim to positive. diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index b6138194775..fccc6ba60c3 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -571,6 +571,22 @@ class CudaCommunicator(DeviceCommunicatorBase): self.all2all_manager.destroy() self.all2all_manager = None # type: ignore[assignment] + def checkpoint_prepare(self) -> None: + # Only FlashInfer all-reduce and FlashInfer all2all are supported for now. + from .flashinfer_all_reduce import checkpoint_prepare_fi_ar_workspaces + + checkpoint_prepare_fi_ar_workspaces(self.cpu_group) + if self.all2all_manager is not None: + self.all2all_manager.checkpoint_prepare() + + def checkpoint_restore(self) -> None: + # Only FlashInfer all-reduce and FlashInfer all2all are supported for now. + from .flashinfer_all_reduce import checkpoint_restore_fi_ar_workspaces + + checkpoint_restore_fi_ar_workspaces(self.cpu_group) + if self.all2all_manager is not None: + self.all2all_manager.checkpoint_restore() + def all_gatherv( self, input_: torch.Tensor | list[torch.Tensor], diff --git a/vllm/distributed/device_communicators/flashinfer_all_reduce.py b/vllm/distributed/device_communicators/flashinfer_all_reduce.py index 881b265d242..8d3e8170924 100644 --- a/vllm/distributed/device_communicators/flashinfer_all_reduce.py +++ b/vllm/distributed/device_communicators/flashinfer_all_reduce.py @@ -6,6 +6,7 @@ import atexit import os import random import threading +from typing import Any import torch import torch.distributed as dist @@ -39,6 +40,7 @@ _fi_ar_workspace = None # allreduce backend or a fallback backend when the primary workspace is not # available on the current topology. _fi_ar_quant_workspace = None +_fi_ar_workspace_groups: dict[int, ProcessGroup] = {} def _create_workspace( @@ -81,6 +83,14 @@ def _create_workspace( return None finally: random.setstate(rng_state) + workspace_id = id(workspace) + workspace_group = _fi_ar_workspace_groups.get(workspace_id) + if workspace_group is not None and workspace_group is not group: + raise RuntimeError( + "FlashInfer returned an all-reduce workspace already associated " + "with a different process group" + ) + _fi_ar_workspace_groups[workspace_id] = group logger.debug( "Initialized FlashInfer All Reduce workspace: backend=%s, " "world_size=%d, rank=%d, max_token_num=%d, hidden_dim=%d, dtype=%s", @@ -111,7 +121,7 @@ def _resolve_fi_ar_backend() -> tuple[str, bool]: # Default to mnnvl for both single- and multi-node setups. The mnnvl # cudagraph hang that previously forced single-node to trtllm # (https://github.com/vllm-project/vllm/issues/35772) was fixed upstream in - # FlashInfer (>= 0.6.12, vLLM pins 0.6.13), so mnnvl is safe here. trtllm + # FlashInfer (>= 0.6.12, vLLM pins 0.6.15), so mnnvl is safe here. trtllm # does not support multi-node allreduce, so mnnvl is required there anyway. # mnnvl needs NVSwitch multicast; on single-node topologies without it, # fall back to trtllm so fused allreduce stays enabled. @@ -268,6 +278,36 @@ def destroy_fi_ar_workspace(): _fi_ar_quant_workspace.destroy() _fi_ar_workspace = _fi_ar_quant_workspace = None + _fi_ar_workspace_groups.clear() + + +def _fi_ar_workspaces_for_group(group: ProcessGroup) -> list[Any]: + workspaces = [_fi_ar_workspace] + if _fi_ar_quant_workspace is not _fi_ar_workspace: + workspaces.append(_fi_ar_quant_workspace) + + group_workspaces = [] + for workspace in workspaces: + if workspace is None: + continue + workspace_group = _fi_ar_workspace_groups.get(id(workspace)) + if workspace_group is None: + raise RuntimeError( + "FlashInfer all-reduce workspace process group was not retained" + ) + if workspace_group is group: + group_workspaces.append(workspace) + return group_workspaces + + +def checkpoint_prepare_fi_ar_workspaces(group: ProcessGroup) -> None: + for workspace in _fi_ar_workspaces_for_group(group): + workspace.checkpoint_prepare() + + +def checkpoint_restore_fi_ar_workspaces(group: ProcessGroup) -> None: + for workspace in _fi_ar_workspaces_for_group(group): + workspace.checkpoint_restore(TorchDistBackend(group=group)) atexit.register(destroy_fi_ar_workspace) diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index b3545d54b4d..4284a609d67 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -127,6 +127,28 @@ def _register_group(group: "GroupCoordinator") -> None: _groups[group.unique_name] = weakref.ref(group) +def _apply_to_device_comms( + action: Callable[[DeviceCommunicatorBase], None], +) -> None: + """Apply ``action`` to every group's device communicator. + + Walks the registered parallel groups and skips those without a device + communicator (absent at ``world_size == 1``). + """ + comms = [] + for group_ref in _groups.values(): + group = group_ref() + if group is None: + continue + dc = group.device_communicator + if dc is None: + continue + comms.append(dc) + + for dc in comms: + action(dc) + + def all_reduce(tensor: torch.Tensor, group_name: str) -> torch.Tensor: assert group_name in _groups, f"Group {group_name} is not found." group = _groups[group_name]() @@ -2028,6 +2050,20 @@ def prepare_communication_buffer_for_model(model: torch.nn.Module): _EPLB.prepare_communication_buffer_for_model(model) +def checkpoint_prepare_distributed_state() -> None: + """Prepare every device communicator for a process checkpoint.""" + torch.accelerator.synchronize() + _apply_to_device_comms(lambda comm: comm.checkpoint_prepare()) + torch.accelerator.synchronize() + + +def checkpoint_restore_distributed_state() -> None: + """Restore every device communicator after a process checkpoint.""" + torch.accelerator.synchronize() + _apply_to_device_comms(lambda comm: comm.checkpoint_restore()) + torch.accelerator.synchronize() + + def model_parallel_is_initialized(): """Check if tensor and pipeline parallel groups are initialized.""" return _TP is not None and _PP is not None diff --git a/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py b/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py index e49e135b26a..94d783e730a 100644 --- a/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py +++ b/vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py @@ -93,7 +93,7 @@ def _can_use_flashinfer(hidden_states: torch.Tensor, tp_size: int) -> tuple[bool max_token_num=max_token_num, hidden_dim=hidden_size, dtype=hidden_states.dtype, - group=get_tp_group().device_group, + group=get_tp_group().cpu_group, ) if workspace is None: return False, 0 diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index f1e7132339c..922a8aa5982 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -943,6 +943,12 @@ class AsyncLLM(EngineClient): if self.logger_manager is not None: self.logger_manager.record_sleep_state(0, 0) + async def checkpoint_prepare(self) -> None: + await self.collective_rpc("checkpoint_prepare") + + async def checkpoint_restore(self) -> None: + await self.collective_rpc("checkpoint_restore") + async def is_sleeping(self) -> bool: return await self.engine_core.is_sleeping_async() diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 7c3d0336688..556b1e6c7d9 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -41,6 +41,8 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import ( ) from vllm.distributed.parallel_state import ( Handle, + checkpoint_prepare_distributed_state, + checkpoint_restore_distributed_state, get_pp_group, get_tp_group, ) @@ -244,6 +246,12 @@ class Worker(WorkerBase): if tags is None or "kv_cache" in tags: self.model_runner.post_kv_cache_wake_up() + def checkpoint_prepare(self) -> None: + checkpoint_prepare_distributed_state() + + def checkpoint_restore(self) -> None: + checkpoint_restore_distributed_state() + def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager: if ( current_platform.is_cuda_alike() From 9e50e1037e268b4488bbd472bd5dfee4cc08bb75 Mon Sep 17 00:00:00 2001 From: aoshen02 <aoshen@inferact.ai> Date: Mon, 27 Jul 2026 03:09:13 +0800 Subject: [PATCH 070/185] [Bugfix][CuMem] Make KV-cache wake cleanup tag-safe (#49857) Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com> --- vllm/device_allocator/cumem.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index 7c4fedd34a3..2cb9805bae3 100644 --- a/vllm/device_allocator/cumem.py +++ b/vllm/device_allocator/cumem.py @@ -291,6 +291,9 @@ class CuMemAllocator: back to GPU memory. If None, all memory allocation will be loaded back to GPU memory. """ + gc.collect() + torch.accelerator.empty_cache() + for ptr, data in self.pointer_to_data.items(): if tags is None or data.tag in tags: handle = data.handle From 0934b267906f8cd9459f287b31647c3ed5c58e01 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" <khluu000@gmail.com> Date: Sun, 26 Jul 2026 13:07:32 -0700 Subject: [PATCH 071/185] [CI/Build] Refresh tags before building macOS wheel (#49901) Signed-off-by: khluu <khluu000@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com> --- .buildkite/scripts/build-macos-wheel.sh | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.buildkite/scripts/build-macos-wheel.sh b/.buildkite/scripts/build-macos-wheel.sh index 9e1719b793b..ac0e4eb1d76 100755 --- a/.buildkite/scripts/build-macos-wheel.sh +++ b/.buildkite/scripts/build-macos-wheel.sh @@ -7,6 +7,9 @@ set -euo pipefail +# The macmini queue uses persistent checkouts, so refresh tags for setuptools-scm. +git fetch --tags --force origin + # The Rust frontend build needs protoc. if ! command -v protoc >/dev/null 2>&1; then brew install protobuf From fdaa0d9e59238b6884f9515fa3245dea118edc66 Mon Sep 17 00:00:00 2001 From: Nick Hill <nickhill123@gmail.com> Date: Mon, 27 Jul 2026 01:38:44 +0100 Subject: [PATCH 072/185] [ModelRunner V2] Support encoder-only attention (#49331) Signed-off-by: Nick Hill <nickhill123@gmail.com> Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> Co-authored-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> --- .../models/language/pooling/test_embedding.py | 54 ++++++ vllm/v1/worker/gpu/attn_utils.py | 12 ++ vllm/v1/worker/gpu/block_table.py | 2 +- vllm/v1/worker/gpu/model_runner.py | 3 + vllm/v1/worker/gpu/model_states/__init__.py | 8 +- .../worker/gpu/model_states/encoder_only.py | 182 ++++++++++++++++++ vllm/v1/worker/gpu/model_states/interface.py | 10 + vllm/v1/worker/gpu/warmup.py | 9 +- 8 files changed, 276 insertions(+), 4 deletions(-) create mode 100644 vllm/v1/worker/gpu/model_states/encoder_only.py diff --git a/tests/models/language/pooling/test_embedding.py b/tests/models/language/pooling/test_embedding.py index e105195afe0..798aa890016 100644 --- a/tests/models/language/pooling/test_embedding.py +++ b/tests/models/language/pooling/test_embedding.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +import torch +from transformers import AutoModel from vllm.config import PoolerConfig @@ -87,3 +89,55 @@ def test_models( name_1="vllm", tol=1e-2, ) + + +@pytest.mark.parametrize( + "model", + [ + "BAAI/bge-base-en-v1.5", + "intfloat/multilingual-e5-small", + ], +) +@torch.inference_mode() +def test_encoder_only_model_runner_v2_attention( + hf_runner, + vllm_runner, + monkeypatch, + model: str, +) -> None: + prompts = [ + "short input", + "a longer input that exercises mixed sequence lengths", + ] + + with hf_runner(model, dtype="float", auto_cls=AutoModel) as hf_model: + hf_outputs = [] + for prompt in prompts: + inputs = hf_model.tokenizer(prompt, return_tensors="pt") + output = hf_model.model(**hf_model.wrap_device(inputs)) + embedding = torch.nn.functional.normalize( + output.last_hidden_state[0, -1].float(), dim=0 + ) + hf_outputs.append(embedding.cpu().tolist()) + + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + with vllm_runner( + model, + runner="pooling", + dtype="float", + max_model_len=64, + max_num_seqs=2, + gpu_memory_utilization=0.25, + pooler_config=PoolerConfig( + task="embed", seq_pooling_type="LAST", use_activation=True + ), + ) as vllm_model: + vllm_outputs = vllm_model.embed(prompts) + + check_embeddings_close( + embeddings_0_lst=hf_outputs, + embeddings_1_lst=vllm_outputs, + name_0="hf", + name_1="vllm", + tol=1e-2, + ) diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index b6bbe28e5f2..7dff4cbe4c9 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -46,6 +46,18 @@ class AttentionCGSupportInfo: min_cg_support: AttentionCGSupport = AttentionCGSupport.ALWAYS min_cg_attn_backend: str | None = None + def narrow( + self, support: AttentionCGSupport, backend: str | None + ) -> "AttentionCGSupportInfo": + """Return an info tightened by ``support`` if it is more restrictive. + + Lets attention groups built outside ``init_attn_backend`` (e.g. + encoder-only layers) contribute to the runner's cudagraph decision. + """ + if support.value < self.min_cg_support.value: + return AttentionCGSupportInfo(support, backend) + return self + def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]: kv_cache_spec: dict[str, KVCacheSpec] = {} diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index 22c4afc11bc..df5d2be629d 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -123,7 +123,7 @@ class BlockTables: if self.num_kv_cache_groups == 1: # Single group: write directly, skipping the per-write group lookup. self.block_tables[0].apply_write() - else: + elif self.num_kv_cache_groups > 1: # Multiple groups: apply all block tables with one fused kernel. assert self.fused_writer is not None self.fused_writer.apply( diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index fdedbfb86d0..64a4f55351a 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -459,6 +459,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.attn_groups, attn_cg_support, self.kernel_block_sizes = init_attn_backend( self.kv_cache_config, self.vllm_config, self.device ) + attn_cg_support = attn_cg_support.narrow( + *self.model_state.get_additional_cg_support() + ) self.block_tables = BlockTables( block_sizes=block_sizes, max_num_reqs=self.max_num_reqs, diff --git a/vllm/v1/worker/gpu/model_states/__init__.py b/vllm/v1/worker/gpu/model_states/__init__.py index dc52dc4ee57..2372a01dd30 100644 --- a/vllm/v1/worker/gpu/model_states/__init__.py +++ b/vllm/v1/worker/gpu/model_states/__init__.py @@ -4,7 +4,7 @@ import torch import torch.nn as nn from vllm.config import VllmConfig -from vllm.model_executor.layers.attention import CrossAttention +from vllm.model_executor.layers.attention import CrossAttention, EncoderOnlyAttention from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache @@ -27,6 +27,12 @@ def init_model_state( return EncoderDecoderModelState(vllm_config, model, encoder_cache, device) + # Encoder-only models (BERT/RoBERTa): non-causal self-attention, no KV cache. + if any(isinstance(m, EncoderOnlyAttention) for m in model.modules()): + from vllm.v1.worker.gpu.model_states.encoder_only import EncoderOnlyModelState + + return EncoderOnlyModelState(vllm_config, model, encoder_cache, device) + if vllm_config.model_config.is_hybrid: from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState diff --git a/vllm/v1/worker/gpu/model_states/encoder_only.py b/vllm/v1/worker/gpu/model_states/encoder_only.py new file mode 100644 index 00000000000..6595b9e434b --- /dev/null +++ b/vllm/v1/worker/gpu/model_states/encoder_only.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any, cast + +import torch +import torch.nn as nn + +from vllm.config import VllmConfig, get_layers_from_vllm_config +from vllm.config.compilation import CUDAGraphMode +from vllm.model_executor.layers.attention import Attention +from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE +from vllm.v1.attention.backend import ( + AttentionCGSupport, + AttentionType, + CommonAttentionMetadata, +) +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + EncoderOnlyAttentionSpec, + KVCacheConfig, +) +from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache +from vllm.v1.worker.gpu.model_states.default import DefaultModelState +from vllm.v1.worker.utils import AttentionGroup + + +class EncoderOnlyModelState(DefaultModelState): + """ModelState for encoder-only (BERT/RoBERTa) models. + + Encoder attention needs no KV cache: it runs full, bidirectional + self-attention over each request's tokens in a single forward. Such layers + return no ``get_kv_cache_spec`` and therefore never enter + ``kv_cache_config.kv_cache_groups``, so the KV-backed attention path never + builds their metadata. We build their (non-causal) metadata here, keeping + the normal KV-backed path untouched. + """ + + def __init__( + self, + vllm_config: VllmConfig, + model: nn.Module, + encoder_cache: EncoderCache | None, + device: torch.device, + ): + super().__init__(vllm_config, model, encoder_cache, device) + + cache_config = vllm_config.cache_config + if cache_config.cache_dtype == "auto": + kv_cache_dtype = self.dtype + else: + kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype] + + # Build an attention group (and its non-causal metadata builder) for the + # encoder-only layers, grouped by backend + query/KV head config. Models + # are typically uniform, yielding a single group. + attn_layers = get_layers_from_vllm_config(vllm_config, Attention) + groups: dict[tuple[tuple[str, str], int, int, int], AttentionGroup] = {} + for name, layer in attn_layers.items(): + if layer.attn_type != AttentionType.ENCODER_ONLY: + continue + # No KV cache is bound for these layers; give them an (empty) + # device tensor so the attention forward context is well-formed. + layer.kv_cache = torch.empty(0, dtype=kv_cache_dtype, device=device) + backend = layer.get_attn_backend() + # Metadata builders require a uniform query-head count per group. + key = ( + backend.full_cls_name(), + layer.num_heads, + layer.num_kv_heads, + layer.head_size, + ) + group = groups.get(key) + if group is None: + spec = EncoderOnlyAttentionSpec( + block_size=cache_config.block_size, + num_kv_heads=layer.num_kv_heads, + head_size=layer.head_size, + dtype=kv_cache_dtype, + ) + group = AttentionGroup(backend, [], spec, kv_cache_group_id=len(groups)) + groups[key] = group + group.layer_names.append(name) + + self.encoder_attn_groups = list(groups.values()) + for group in self.encoder_attn_groups: + group.create_metadata_builders(vllm_config, device) + + # Encoder attention reads neither the block table nor the slot mapping + # (full varlen self-attention over the batch's q/k/v), but the metadata + # builder expects both tensors to be present. + self._dummy_block_table = torch.zeros( + self.max_num_reqs, 1, dtype=torch.int32, device=device + ) + self._dummy_slot_mapping = torch.zeros( + self.max_num_tokens, dtype=torch.int64, device=device + ) + + def get_additional_cg_support(self) -> tuple[AttentionCGSupport, str | None]: + # Encoder groups are built here rather than in init_attn_backend, so + # their cudagraph support must be surfaced to the runner separately. + support = AttentionCGSupport.ALWAYS + backend: str | None = None + for group in self.encoder_attn_groups: + builder = group.get_metadata_builder(0) + cg_support = builder.get_cudagraph_support( + self.vllm_config, cast(AttentionSpec, group.kv_cache_spec) + ) + if cg_support.value < support.value: + support = cg_support + backend = group.backend.__name__ + return support, backend + + def prepare_attn( + self, + input_batch: InputBatch, + cudagraph_mode: CUDAGraphMode, + block_tables: tuple[torch.Tensor, ...], + slot_mappings: torch.Tensor, + attn_groups: list[list[AttentionGroup]], + kv_cache_config: KVCacheConfig, + for_capture: bool = False, + ) -> dict[str, Any]: + attn_metadata = super().prepare_attn( + input_batch, + cudagraph_mode, + block_tables, + slot_mappings, + attn_groups, + kv_cache_config, + for_capture, + ) + attn_metadata.update( + self._build_encoder_attn_metadata(input_batch, cudagraph_mode, for_capture) + ) + return attn_metadata + + def _build_encoder_attn_metadata( + self, + input_batch: InputBatch, + cudagraph_mode: CUDAGraphMode, + for_capture: bool, + ) -> dict[str, Any]: + if cudagraph_mode == CUDAGraphMode.FULL: + num_reqs = input_batch.num_reqs_after_padding + num_tokens = input_batch.num_tokens_after_padding + else: + num_reqs = input_batch.num_reqs + num_tokens = input_batch.num_tokens + max_query_len = int(input_batch.num_scheduled_tokens.max()) + if for_capture: + max_seq_len = self.max_model_len + else: + max_seq_len = int(input_batch.seq_lens_cpu_upper_bound[:num_reqs].max()) + + # The encoder builder forces ``causal=False`` regardless of this value. + common_attn_metadata = CommonAttentionMetadata( + query_start_loc=input_batch.query_start_loc, + query_start_loc_cpu=torch.from_numpy(input_batch.query_start_loc_np), + seq_lens=input_batch.seq_lens[:num_reqs], + num_reqs=num_reqs, + num_actual_tokens=num_tokens, + max_query_len=max_query_len, + max_seq_len=max_seq_len, + block_table_tensor=self._dummy_block_table[:num_reqs], + slot_mapping=self._dummy_slot_mapping[:num_tokens], + seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound[:num_reqs], + positions=input_batch.positions, + ) + + attn_metadata: dict[str, Any] = {} + for group in self.encoder_attn_groups: + builder = group.get_metadata_builder(0) + if for_capture: + metadata = builder.build_for_cudagraph_capture(common_attn_metadata) + else: + metadata = builder.build( + common_prefix_len=0, common_attn_metadata=common_attn_metadata + ) + for name in group.layer_names: + attn_metadata[name] = metadata + return attn_metadata diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index df86efa4a79..666d460c200 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -9,6 +9,7 @@ import torch.nn as nn from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode from vllm.tasks import GenerationTask +from vllm.v1.attention.backend import AttentionCGSupport from vllm.v1.core.sched.output import NewRequestData from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.input_batch import InputBatch @@ -95,6 +96,15 @@ class ModelState(ABC): def apply_staged_writes(self) -> None: return None + def get_additional_cg_support(self) -> tuple[AttentionCGSupport, str | None]: + """Cudagraph support of attention groups this ModelState builds outside + ``init_attn_backend`` (e.g. encoder-only layers). + + Returns the minimum support level and its backend name. The default of + ``ALWAYS`` imposes no extra constraint on the runner's cudagraph mode. + """ + return AttentionCGSupport.ALWAYS, None + def preprocess_state( self, input_batch: InputBatch, diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index f4047a0be8a..785188dc3c3 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -217,9 +217,14 @@ def warmup_kernels( model_runner.scheduler_config.max_num_seqs, model_runner.scheduler_config.max_num_batched_tokens // max(prompt_len, decode_query_len), - # Reserve block 0 (null block) and ensure we have enough blocks. - max(1, (model_runner.kv_cache_config.num_blocks - 1) // max_blocks_per_req), ) + if max_blocks_per_req > 0: + # Reserve block 0 (null block) and ensure we have enough blocks. + # Encoder-only models allocate no KV blocks, so this cap doesn't apply. + num_reqs = min( + num_reqs, + max(1, (model_runner.kv_cache_config.num_blocks - 1) // max_blocks_per_req), + ) req_ids = [f"_warmup_{i}_" for i in range(num_reqs)] From f0553889c04ae5bf3df48245302ba5ac19fe3432 Mon Sep 17 00:00:00 2001 From: Nick Iusiumbeli <nickuspro@gmail.com> Date: Mon, 27 Jul 2026 04:06:07 +0300 Subject: [PATCH 073/185] [Bugfix] Prevent NaN poisoning in xpu_mla_sparse for fully-masked index chunks (#48366) Signed-off-by: Nick Iusiumbeli <nickuspro@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com> --- .../kernels/attention/test_xpu_mla_sparse.py | 52 +++++++++++++++++++ vllm/v1/attention/ops/xpu_mla_sparse.py | 7 ++- 2 files changed, 57 insertions(+), 2 deletions(-) diff --git a/tests/kernels/attention/test_xpu_mla_sparse.py b/tests/kernels/attention/test_xpu_mla_sparse.py index 419644923ec..965c6f735f0 100644 --- a/tests/kernels/attention/test_xpu_mla_sparse.py +++ b/tests/kernels/attention/test_xpu_mla_sparse.py @@ -116,3 +116,55 @@ def test_bf16_triton_sparse_mla(device_str, dtype): assert torch.allclose(out, ref_out, atol=1e-2, rtol=1e-2) assert torch.allclose(max_logits, ref_max_logits, atol=1e-3, rtol=1e-3) assert torch.allclose(lse, ref_lse, atol=1e-3, rtol=1e-3) + + +@pytest.mark.parametrize("device_str", ["xpu"]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.skipif( + not torch.xpu.is_available(), + reason="XPU is required", +) +def test_bf16_triton_sparse_mla_masked_chunks(device_str, dtype): + """Rows whose leading BLOCK_N index entries are all masked must not NaN. + + Regression test: with an -inf running max, a fully-masked leading chunk + produced re_scale = exp2(-inf - -inf) = NaN, permanently poisoning the + accumulator even though valid keys followed in later chunks. + """ + device = torch.device(device_str) + s_q = 3 + s_kv = 256 + h_q = 64 + h_kv = 1 + d_qk = 576 + d_v = 512 + topk = 128 # 8 chunks of BLOCK_N=16 + + torch.random.manual_seed(1234) + + q = torch.randn((s_q, h_q, d_qk), dtype=dtype, device=device) + kv = torch.randn((s_kv, h_kv, d_qk), dtype=dtype, device=device) + indices = torch.full((s_q, h_kv, topk), -1, dtype=torch.int32, device=device) + # row 0: valid keys only in chunks 1-2 -> leading AND trailing masked chunks + indices[0, 0, 16:48] = torch.arange(32, dtype=torch.int32, device=device) + # row 1: fully valid + indices[1, 0, :] = torch.arange(topk, dtype=torch.int32, device=device) + # row 2: no valid key at all + + sm_scale = d_qk**-0.5 + + out, max_logits, lse = triton_bf16_mla_sparse_interface( + q, kv, indices, sm_scale, d_v + ) + assert out.isfinite().all() + + ref_out, _, ref_max_logits, ref_lse = reference_mla_sparse_prefill( + q, kv, indices, sm_scale, d_v + ) + assert torch.allclose(out[:2], ref_out[:2], atol=1e-2, rtol=1e-2) + assert torch.allclose(max_logits[:2], ref_max_logits[:2], atol=1e-3, rtol=1e-3) + assert torch.allclose(lse[:2], ref_lse[:2], atol=1e-3, rtol=1e-3) + # A row with no valid key yields zeros (the reference's convention); its + # lse/max_logits are large-negative finite rather than the reference's + # +inf/-inf placeholders, so only the output is compared here. + assert torch.allclose(out[2], torch.zeros_like(out[2])) diff --git a/vllm/v1/attention/ops/xpu_mla_sparse.py b/vllm/v1/attention/ops/xpu_mla_sparse.py index e73e5a2b28e..9c1a92e10f5 100644 --- a/vllm/v1/attention/ops/xpu_mla_sparse.py +++ b/vllm/v1/attention/ops/xpu_mla_sparse.py @@ -73,7 +73,10 @@ def _bf16_mla_sparse_kernel( q_buffer + off_qpe, mask=(mask_h[:, None]) & (mask_dpe[None, :]), other=0.0 ) - e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - float("inf") + # Use a finite sentinel rather than -inf: if every key in a chunk is + # masked (e.g. leading -1 padding in `indices`) an -inf running max gives + # re_scale = exp2(-inf - -inf) = NaN, permanently poisoning acc / e_sum. + e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - 1.0e30 e_sum = tl.zeros([BLOCK_H], dtype=tl.float32) acc = tl.zeros([BLOCK_H, BLOCK_DV], dtype=tl.float32) @@ -122,7 +125,7 @@ def _bf16_mla_sparse_kernel( # apply scaling qk *= sm_scale - qk = tl.where((mask_h[:, None]) & (mask_kv[None, :]), qk, -float("inf")) + qk = tl.where((mask_h[:, None]) & (mask_kv[None, :]), qk, -1.0e30) # load v mask_v_d = offs_dv < dim_v From 50aa83048219b70a3a68adf2fc8cd860ccc3e238 Mon Sep 17 00:00:00 2001 From: Nick Hill <nickhill123@gmail.com> Date: Mon, 27 Jul 2026 03:04:20 +0100 Subject: [PATCH 074/185] [BugFix][MRV2] Don't create dummy requests longer than `max_model_len` (#49751) Signed-off-by: Nick Hill <nickhill123@gmail.com> --- tests/v1/worker/test_gpu_input_batch_v2.py | 52 ++++++++++++++++++++++ vllm/v1/worker/gpu/input_batch.py | 13 ++++-- vllm/v1/worker/gpu/lora_utils.py | 6 ++- vllm/v1/worker/gpu/model_runner.py | 8 +++- 4 files changed, 72 insertions(+), 7 deletions(-) create mode 100644 tests/v1/worker/test_gpu_input_batch_v2.py diff --git a/tests/v1/worker/test_gpu_input_batch_v2.py b/tests/v1/worker/test_gpu_input_batch_v2.py new file mode 100644 index 00000000000..2c43cd72147 --- /dev/null +++ b/tests/v1/worker/test_gpu_input_batch_v2.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the V2 model runner's InputBatch (vllm.v1.worker.gpu.input_batch).""" + +import pytest +import torch + +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers + +DEVICE = "cuda" + + +@pytest.mark.parametrize( + "num_reqs,num_tokens", + [ + (256, 496), # remainder 240: previously gave the last request 241 tokens + (128, 512), # no remainder + (3, 8), + (1, 7), + ], +) +def test_make_dummy_distributes_remainder(num_reqs: int, num_tokens: int): + """No dummy request may exceed ceil(num_tokens / num_reqs) tokens. + + Dumping the remainder on a single request can produce a dummy request with + seq_len > max_model_len, which the block tables cannot back; attention + kernels running on the dummy batch during cudagraph capture then read + block-table entries out of bounds (https://github.com/vllm-project/vllm/pull/49364 + CI failure). + """ + buffers = InputBuffers( + max_num_reqs=num_reqs, max_num_tokens=num_tokens, device=torch.device(DEVICE) + ) + batch = InputBatch.make_dummy(num_reqs, num_tokens, buffers) + + max_per_req = -(-num_tokens // num_reqs) + assert batch.num_scheduled_tokens.sum() == num_tokens + assert batch.num_scheduled_tokens.max() == max_per_req + assert batch.num_scheduled_tokens.min() >= num_tokens // num_reqs + # Requests with an extra token are placed at the end of the batch. + assert (batch.num_scheduled_tokens[:-1] <= batch.num_scheduled_tokens[1:]).all() + + # seq_len == query_len for the dummy prefill-shaped batch, on GPU and CPU. + query_lens = batch.query_start_loc_np[1:] - batch.query_start_loc_np[:-1] + assert (query_lens == batch.num_scheduled_tokens).all() + assert torch.equal( + batch.seq_lens, torch.from_numpy(batch.num_scheduled_tokens).to(DEVICE) + ) + assert batch.query_start_loc_np[-1] == num_tokens + assert torch.equal( + batch.query_start_loc.cpu(), torch.from_numpy(batch.query_start_loc_np) + ) diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 64c1096dfbe..8c3a9cc030f 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -115,13 +115,18 @@ class InputBatch: expanded_idx_mapping = idx_mapping expanded_local_pos = torch.zeros(num_reqs, dtype=torch.int32, device=device) - num_scheduled_tokens = np.full(num_reqs, num_tokens // num_reqs, dtype=np.int32) - num_scheduled_tokens[-1] += num_tokens % num_reqs + # Distribute the remainder evenly so that no dummy request exceeds + # ceil(num_tokens / num_reqs) <= max_model_len tokens. + base_tokens = num_tokens // num_reqs + num_extra = num_tokens % num_reqs + num_scheduled_tokens = np.full(num_reqs, base_tokens, dtype=np.int32) + if num_extra > 0: + num_scheduled_tokens[-num_extra:] += 1 assert int(num_scheduled_tokens.sum()) == num_tokens # seq_len equals to query_len - input_buffers.seq_lens[:num_reqs] = num_tokens // num_reqs - input_buffers.seq_lens[num_reqs - 1] += num_tokens % num_reqs + input_buffers.seq_lens[: num_reqs - num_extra] = base_tokens + input_buffers.seq_lens[num_reqs - num_extra : num_reqs] = base_tokens + 1 # Pad for full CUDA graph mode. input_buffers.seq_lens[num_reqs:] = 0 seq_lens = input_buffers.seq_lens[:num_reqs] diff --git a/vllm/v1/worker/gpu/lora_utils.py b/vllm/v1/worker/gpu/lora_utils.py index fa281f6817b..6a00b35e4bc 100644 --- a/vllm/v1/worker/gpu/lora_utils.py +++ b/vllm/v1/worker/gpu/lora_utils.py @@ -59,8 +59,12 @@ def create_lora_capture_hook( return None def hook(num_active_loras: int, num_reqs: int, num_tokens: int) -> None: + # Match InputBatch.make_dummy: distribute the remainder evenly so no + # dummy request exceeds ceil(num_tokens / num_reqs) tokens. num_scheduled = np.full(num_reqs, num_tokens // num_reqs, dtype=np.int32) - num_scheduled[-1] += num_tokens % num_reqs + num_extra = num_tokens % num_reqs + if num_extra > 0: + num_scheduled[-num_extra:] += 1 with runner.maybe_select_dummy_loras( lora_config, num_scheduled, num_active_loras=num_active_loras ): diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 64a4f55351a..f0da049ce46 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -564,8 +564,12 @@ class GPUModelRunner(LoRAModelRunnerMixin): num_tokens = max(num_tokens, self.decode_query_len) num_reqs = num_tokens // self.decode_query_len assert num_tokens % self.decode_query_len == 0 - num_tokens_per_request = [num_tokens // num_reqs] * num_reqs - num_tokens_per_request[-1] += num_tokens % num_reqs + # Distribute the remainder evenly so no dummy request exceeds + # ceil(num_tokens / num_reqs) <= max_model_len tokens. + num_tokens_per_request = [ + num_tokens // num_reqs + (i >= num_reqs - num_tokens % num_reqs) + for i in range(num_reqs) + ] assert sum(num_tokens_per_request) == num_tokens num_scheduled_tokens = { From ffc4f08c8ee130d4ea6347c1bf31ffd4f8af28ab Mon Sep 17 00:00:00 2001 From: limeward <32970461+edwinlim0919@users.noreply.github.com> Date: Sun, 26 Jul 2026 19:10:00 -0700 Subject: [PATCH 075/185] [Core][KV-transfer] MoRIIO: heterogeneous TP<->DP prefill/decode read routing (#46116) Signed-off-by: Edwin Lim <edwin.lim@mangoboost.io> --- .../unit/test_moriio_routing_fairness.py | 210 ++++++++++++++ .../kv_connector/v1/moriio/moriio_common.py | 5 + .../v1/moriio/moriio_connector.py | 259 +++++++++++++++++- 3 files changed, 467 insertions(+), 7 deletions(-) create mode 100644 tests/v1/kv_connector/unit/test_moriio_routing_fairness.py diff --git a/tests/v1/kv_connector/unit/test_moriio_routing_fairness.py b/tests/v1/kv_connector/unit/test_moriio_routing_fairness.py new file mode 100644 index 00000000000..e38ec3f7e9b --- /dev/null +++ b/tests/v1/kv_connector/unit/test_moriio_routing_fairness.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Hardware-fair KV-read routing across MoRIIO heterogeneous P/D configs. + +Dependency-light (no GPU / ROCm / mori): builds bare ``MoRIIOConnectorWorker`` +instances via ``object.__new__`` and drives the REAL read-source routing +decision (``_resolve_read_source`` / ``_next_flex_tp_rank``) directly -- no +routing logic is re-implemented here, so a future change to those functions is +what these tests exercise. + +Scope -- the CONNECTOR (decode-side) read routing. The connector never selects +the prefill instance; the proxy does and hands the connector a ``remote_host`` + +``remote_dp_rank`` per request. The connector's only fairness lever is WHICH +prefill (dp, tp) rank each read targets, and the RFC's deployments collapse to +three real connector behaviours: + + symmetric TP (TP prefill + TP decode) -> decode tp_k reads prefill tp_k + flexible (TP prefill + DP decode) -> round-robin over prefill tp0..N-1 + owner DP (DP prefill, any decode) -> read the owner rank, tp0 + +These assume the proxy delivers a FAIR request stream (each prefill instance + +owner dp-rank an equal share) and verify the connector never re-introduces a +bottleneck. That proxy contract is checked separately against the real +``flat_interleaved_dp_route`` in the toy-proxy tests (PR #46115). + +A node runs one of two modes (8 GPUs each): + * TP8 -> (dp_size, tp_size) = (1, 8); MLA latent KV REPLICATED on all ranks. + * DP8EP -> (dp_size, tp_size) = (8, 1); KV PARTITIONED, one owner rank/request. +""" + +from collections import Counter +from dataclasses import dataclass +from unittest.mock import MagicMock + +import pytest + +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import ( + ReqMeta, + get_port_offset, +) +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import ( + MoRIIOConnectorWorker, +) + +MODE_DIMS = {"TP8": (1, 8), "DP8EP": (8, 1)} # (dp_size, tp_size), 8 GPUs/node + + +@dataclass(frozen=True) +class PDConfig: + name: str + p_mode: str + d_mode: str + + @property + def p_dp(self) -> int: + return MODE_DIMS[self.p_mode][0] + + @property + def p_tp(self) -> int: + return MODE_DIMS[self.p_mode][1] + + @property + def d_dp(self) -> int: + return MODE_DIMS[self.d_mode][0] + + @property + def d_tp(self) -> int: + return MODE_DIMS[self.d_mode][1] + + @property + def n_prefill_gpus(self) -> int: + return self.p_dp * self.p_tp + + +CONFIGS = [ + PDConfig("1P_TP8:1D_TP8", "TP8", "TP8"), + PDConfig("2P_TP8:1D_DP8EP", "TP8", "DP8EP"), + PDConfig("2P_TP8:2D_TP8", "TP8", "TP8"), + PDConfig("2P_DP8EP:3D_DP8EP", "DP8EP", "DP8EP"), + PDConfig("2P_DP8EP:4D_TP8", "DP8EP", "TP8"), +] +CONFIG_IDS = [c.name for c in CONFIGS] + + +def make_decode_worker(*, world_size: int, tp_rank: int, dp_rank: int): + w = object.__new__(MoRIIOConnectorWorker) + w.world_size = world_size + w.tp_rank = tp_rank + w.dp_rank = dp_rank + w.use_mla = True + return w + + +def make_meta(*, p_tp: int, p_dp: int, remote_dp_rank: int, host: str = "phost0"): + return ReqMeta( + transfer_id="t", + local_block_ids=[1], + remote_block_ids=[2], + remote_host=host, + remote_port=1234, + remote_handshake_port=6301, + remote_notify_port=61005, + remote_engine_id=f"{host}:6301", + tp_size=p_tp, + remote_dp_size=p_dp, + remote_dp_rank=remote_dp_rank, + ) + + +def build_decode_workers(cfg: PDConfig) -> list: + """Decode workers that issue reads for one prefill instance. A TP decode + instance reads from every tp rank; a DP decode instance from each dp-rank + worker. Reused across requests so per-worker round-robin state advances.""" + if cfg.d_tp > 1: + return [ + make_decode_worker(world_size=cfg.d_tp, tp_rank=r, dp_rank=0) + for r in range(cfg.d_tp) + ] + return [ + make_decode_worker(world_size=1, tp_rank=0, dp_rank=d) for d in range(cfg.d_dp) + ] + + +def prefill_target_multiset(cfg: PDConfig, rounds: int) -> Counter: + """Drive the REAL _resolve_read_source over a fair input stream; tally which + prefill GPU each read targets. The only modelled assumption is the proxy + contract -- each owner dp-rank delivered equally (``for owner_dp in + range(p_dp)``); no proxy algorithm is reproduced.""" + workers = build_decode_workers(cfg) + hits: Counter = Counter() + for _ in range(rounds): + for owner_dp in range(cfg.p_dp): + for worker in workers: + meta = make_meta(p_tp=cfg.p_tp, p_dp=cfg.p_dp, remote_dp_rank=owner_dp) + chosen_tp, _flexible = worker._resolve_read_source(meta) + target = get_port_offset(owner_dp, chosen_tp, cfg.p_tp) + assert 0 <= target < cfg.n_prefill_gpus + hits[target] += 1 + return hits + + +@pytest.mark.parametrize("cfg", CONFIGS, ids=CONFIG_IDS) +def test_kv_read_load_is_hardware_fair(cfg: PDConfig) -> None: + # rounds a multiple of p_tp so the round-robin closes on an exact cycle. + hits = prefill_target_multiset(cfg, rounds=16) + gpus = set(range(cfg.n_prefill_gpus)) + assert set(hits) == gpus, f"unused prefill GPUs: {sorted(gpus - set(hits))}" + assert max(hits.values()) == min(hits.values()), dict(sorted(hits.items())) + + +@pytest.mark.parametrize("cfg", CONFIGS, ids=CONFIG_IDS) +def test_flexible_gate_fires_only_for_tp_prefill_dp_decode(cfg: PDConfig) -> None: + flags = set() + for worker in build_decode_workers(cfg): + meta = make_meta(p_tp=cfg.p_tp, p_dp=cfg.p_dp, remote_dp_rank=0) + _chosen, flexible = worker._resolve_read_source(meta) + flags.add(flexible) + is_mirror = cfg.d_tp == 1 and cfg.p_dp == 1 and cfg.p_tp > 1 + assert flags == {is_mirror} + + +def test_symmetric_tp_is_a_bijection() -> None: + targets = [] + for tp_rank in range(8): + worker = make_decode_worker(world_size=8, tp_rank=tp_rank, dp_rank=0) + chosen_tp, flexible = worker._resolve_read_source( + make_meta(p_tp=8, p_dp=1, remote_dp_rank=0) + ) + assert not flexible + targets.append(chosen_tp) + assert sorted(targets) == list(range(8)) + + +def test_owner_dp_read_is_faithful_and_covers_every_rank() -> None: + worker = make_decode_worker(world_size=8, tp_rank=3, dp_rank=0) + targets = [] + for owner_dp in range(8): + chosen_tp, flexible = worker._resolve_read_source( + make_meta(p_tp=1, p_dp=8, remote_dp_rank=owner_dp) + ) + assert not flexible + assert chosen_tp == 0 # p_tp == 1, only tp0 exists + targets.append(get_port_offset(owner_dp, chosen_tp, 1)) + assert sorted(targets) == list(range(8)) + + +def test_flexible_round_robin_is_deterministic_uniform_and_staggered() -> None: + w0 = make_decode_worker(world_size=1, tp_rank=0, dp_rank=0) + seq = [w0._next_flex_tp_rank(8) for _ in range(64)] + assert seq[:8] == list(range(8)) + assert Counter(seq) == Counter({t: 8 for t in range(8)}) + + first_pick = [ + make_decode_worker(world_size=1, tp_rank=0, dp_rank=d)._next_flex_tp_rank(8) + for d in range(8) + ] + assert sorted(first_pick) == list(range(8)) + + +def test_read_blocks_for_req_threads_chosen_tp() -> None: + # The resolved (chosen_tp, flexible) must reach _read_blocks, which keys the + # session AND the notify port off that single value -- so a read and its + # completion notify address the same prefill rank. + worker = make_decode_worker(world_size=1, tp_rank=0, dp_rank=3) + worker._read_blocks = MagicMock() + worker._read_blocks_for_req("r", make_meta(p_tp=8, p_dp=1, remote_dp_rank=0)) + kw = worker._read_blocks.call_args.kwargs + assert kw["flexible"] is True + assert kw["chosen_tp"] == 3 # first flexible pick = dp_rank seed + assert get_port_offset(0, kw["chosen_tp"], 8) == 3 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py index 15585123e5c..45df75eeb3c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py @@ -422,6 +422,10 @@ class ReqMeta: remote_engine_id: str tp_size: int remote_dp_size: int + # Prefill DP rank that owns this request's KV (forwarded by the proxy). The + # read must target this rank's memory registration; the default 0 preserves + # the symmetric single-DP behaviour. + remote_dp_rank: int = 0 class MoRIIOConnectorMetadata(KVConnectorMetadata): @@ -475,6 +479,7 @@ class MoRIIOConnectorMetadata(KVConnectorMetadata): remote_notify_port=int(remote_notify_port), tp_size=kv_transfer_params.get("tp_size", 1), remote_dp_size=kv_transfer_params.get("remote_dp_size", 1), + remote_dp_rank=kv_transfer_params.get("remote_dp_rank", 0), ) if write_mode: self.reqs_to_save[request_id] = _req diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index ea41adf437c..c478ced2211 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -1014,6 +1014,8 @@ class MoRIIOConnectorWorker: self._handshake_futures: dict[EngineId, Future[set[str]]] = {} # Protects _handshake_futures and _remote_agents. self._handshake_lock = threading.RLock() + # Remote engines already covered by the eager pre-forward handshake. + self._eager_handshaked_engines: set[EngineId] = set() self.block_size = vllm_config.cache_config.block_size self.model_config = vllm_config.model_config @@ -1261,8 +1263,16 @@ class MoRIIOConnectorWorker: remote_tp_size: int, expected_engine_id: str, remote_dp_rank: int = 0, + remote_tp_rank: int | None = None, ) -> set[str]: - """Do a MoRIIO handshake with a remote instance.""" + """Do a MoRIIO handshake with a remote instance. + + remote_tp_rank: explicit remote TP index to dial. Flexible-read callers + pass the chosen prefill TP rank so the handshake, the (dp, tp) session + key and the notify port all address the SAME rank. None falls back to the + local-rank mapping _remote_tp_rank -- byte-identical for callers not yet + TP-aware. + """ start_time = time.perf_counter() @@ -1270,9 +1280,12 @@ class MoRIIOConnectorWorker: # a hack to keep us moving. We will switch when moving to etcd # or where we have a single ZMQ socket in the scheduler. - port_offset = get_port_offset( - remote_dp_rank, self._remote_tp_rank(remote_tp_size) + dial_tp_rank = ( + self._remote_tp_rank(remote_tp_size) + if remote_tp_rank is None + else int(remote_tp_rank) ) + port_offset = get_port_offset(remote_dp_rank, dial_tp_rank, remote_tp_size) path = make_zmq_path("tcp", host, port + port_offset) logger.debug("handshake Querying metadata on path: %s", path) @@ -1729,6 +1742,149 @@ class MoRIIOConnectorWorker: def get_engine_name_with_dp(self, engine_name, dp_rank): return f"{engine_name}_dp{dp_rank}" + def get_engine_name_with_dp_tp(self, engine_name, dp_rank, tp_rank): + # Per-(dp, tp) session key. The flexible mirror read keys sessions per + # (dp, tp) so one decode worker can hold a session to EACH prefill TP + # rank and spread reads across them; other configs keep the DP-only key. + return f"{engine_name}_dp{dp_rank}_tp{tp_rank}" + + def _eager_handshake_all_dp_ranks(self, metadata: MoRIIOConnectorMetadata) -> None: + """Handshake EVERY remote prefill DP rank BEFORE the decode forward pass, + identically across all local TP workers. + + Why this exists (the deadlock it prevents): with heterogeneous DP prefill + a decode TP worker reads KV from whichever prefill DP rank owns the + request, so across requests every worker must reach several prefill DP + ranks. The decode forward issues per-layer TP collectives (e.g. an + all-gather) that all local TP workers must enter together. If the + handshakes are left to fire lazily on the read path, the workers diverge: + a worker whose target rank is already cached races ahead into the forward + collective while a peer is still blocked in a handshake recv(). The first + worker then waits inside the collective for the stuck peer -> 600s NCCL + timeout / hang. This was observed directly with mixed TP<->DP configs. + + Fix: complete ALL prefill-DP-rank handshakes for every referenced remote + engine HERE, before any read enters the forward, so no worker is still + handshaking once its peers reach a collective. Fires ONCE per remote + engine (first contact), gated by _eager_handshaked_engines. The engine + set comes from scheduler-built metadata (identical on every TP worker), + so all workers run the same handshakes in the same order and reach the + all-reduce barrier below together. + + Failure handling: handshake exceptions are caught, never raised before + the collective (raising early would hang the peers still waiting for it). + Every worker reaches the all-reduce(MIN) vote; if ANY worker failed, ALL + raise the same error AFTER the collective, so the step fails fast and + uniformly in ~seconds instead of one rank hanging the forward for 600s. + """ + import torch.distributed as dist + + # Distinct remote engines referenced this step, in metadata (== + # scheduler) order so every TP worker iterates engines identically. + engines: dict[str, ReqMeta] = {} + for _req_id, meta in metadata.reqs_to_recv.items(): + remote_engine_id = ( + str(meta.remote_host) + ":" + str(meta.remote_handshake_port) + ) + engines.setdefault(remote_engine_id, meta) + + for remote_engine_id, meta in engines.items(): + if remote_engine_id in self._eager_handshaked_engines: + continue + + remote_dp_size = int(meta.remote_dp_size) + port = int(meta.remote_handshake_port) + tp_size = int(meta.tp_size) + + # Flexible mirror (TP prefill + MLA, world_size==1 decode): the read + # round-robins over prefill TP ranks, so pre-warm a session to EVERY + # (dp, tp) rank. Other configs pre-warm per DP rank (tp resolved by + # the fixed local-rank mapping) -- byte-identical to before. The + # mirror's decode is DP+EP, whose forward all-to-all is the collective + # that the eager barrier keeps everyone in step for. + flexible = ( + self.world_size == 1 + and self.use_mla + and remote_dp_size == 1 + and tp_size > 1 + ) + # (engine_id, dp_rank, tp_rank_or_None); tp_rank is None on the legacy + # path so _moriio_handshake falls back to its _remote_tp_rank mapping. + targets: list[tuple[Any, int, int | None]] + if flexible: + targets = [ + (self.get_engine_name_with_dp_tp(remote_engine_id, dp, tp), dp, tp) + for dp in range(remote_dp_size) + for tp in range(max(1, tp_size)) + ] + else: + targets = [ + (self.get_engine_name_with_dp(remote_engine_id, dp), dp, None) + for dp in range(remote_dp_size) + ] + + # Submit handshakes for every not-yet-known target UNDER the lock; do + # NOT hold it across the join or the collective (a stalled recv must + # not block another thread's lock acquisition). Gate on BOTH + # _remote_agents AND layer metadata: a rank with an agent entry but no + # layer metadata is half-handshaked and would KeyError at read time. + futures: list[tuple[str, Future[set[str]]]] = [] + with self._handshake_lock: + for eid, cur_dp_rank, cur_tp_rank in targets: + if ( + eid in self._remote_agents + and eid in self.layer_name_to_remote_kv_cache_metadata + ): + continue + fut = self._handshake_initiation_executor.submit( + self._moriio_handshake, + meta.remote_host, + port, + tp_size, + eid, + cur_dp_rank, + cur_tp_rank, + ) + futures.append((eid, fut)) + + # Join outside the lock. Bounded handshake errors are recorded here + # and reported after the all-reduce. + all_ok = True + results: dict[str, set[str]] = {} + for eid, fut in futures: + try: + results[eid] = fut.result() + except Exception: + logger.exception("Eager MoRIIO handshake failed for %s", eid) + all_ok = False + + with self._handshake_lock: + for eid, agents in results.items(): + self._remote_agents[eid] = agents + + logger.info( + "Eager MoRIIO handshake: engine=%s dp_size=%d new_ranks=%d " + "ok=%s tp_rank=%d", + remote_engine_id, + remote_dp_size, + len(futures), + all_ok, + self.tp_rank, + ) + # CPU all-reduce = TP-uniform success vote AND lockstep barrier: it + # blocks until every TP worker arrives, gives them the same verdict, + # and stays off the model compute stream. + vote = torch.tensor([1 if all_ok else 0], device="cpu", dtype=torch.int32) + dist.all_reduce(vote, group=self.tp_group.cpu_group, op=dist.ReduceOp.MIN) + if int(vote.item()) == 0: + raise HandshakeError( + f"Eager MoRIIO handshake failed for {remote_engine_id} on " + "at least one TP rank; failing this step fast to avoid a " + "TP collective hang" + ) + + self._eager_handshaked_engines.add(remote_engine_id) + def start_load_kv(self, metadata: MoRIIOConnectorMetadata): """ Start loading by triggering non-blocking moriio_xfer. @@ -1750,6 +1906,12 @@ class MoRIIOConnectorWorker: if self.mode == MoRIIOMode.WRITE: return + # Handshake every referenced remote prefill rank up front, before any + # read enters the forward pass. A lazy per-rank handshake on the read + # path lets TP workers diverge into a forward collective while a peer is + # still blocked handshaking -> NCCL hang (see below). + self._eager_handshake_all_dp_ranks(metadata) + wait_handshake_readd_req = False remote_engine_id = None @@ -1758,8 +1920,15 @@ class MoRIIOConnectorWorker: str(meta.remote_host) + ":" + str(meta.remote_handshake_port) ) meta.remote_engine_id = remote_engine_id + # The eager handshake above already covered every referenced engine + # (and keys the mirror per (dp, tp), which the DP-only dp0 probe below + # would miss). Only fall back to the lazy background handshake for an + # engine it did not cover. dp0_remote_engine_id = self.get_engine_name_with_dp(remote_engine_id, 0) - if dp0_remote_engine_id not in self._remote_agents: + if ( + remote_engine_id not in self._eager_handshaked_engines + and dp0_remote_engine_id not in self._remote_agents + ): # Initiate handshake with remote engine to exchange metadata. with self._handshake_lock: if remote_engine_id not in self._remote_agents: @@ -1809,12 +1978,53 @@ class MoRIIOConnectorWorker: self.save_kv_layer(metadata, layer_name, kv_layer, None) self._writer.seal_pending_transfers() + def _next_flex_tp_rank(self, remote_tp_size: int) -> int: + """Deterministic round-robin over prefill tp0..N-1 for the flexible read. + + Round-robin (not random): exactly uniform and testable, with the same + prefill-NIC balancing. Seeded from this decode rank's dp_rank so + concurrent decode DP ranks are phase-staggered -- at a given read index + distinct decode ranks target distinct prefill TP ranks. + """ + rr = getattr(self, "_flex_tp_rr", None) + if rr is None: + rr = int(getattr(self, "dp_rank", 0) or 0) + self._flex_tp_rr = rr + 1 + return rr % remote_tp_size + + def _resolve_read_source(self, meta: ReqMeta) -> tuple[int, bool]: + """Resolve (chosen_tp, flexible) for reading this request's KV. + + Flexible mirror (decode world_size==1 + MLA + pure-TP prefill): MLA + replicates the latent KV across the prefill TP ranks, so any is a valid + source; round-robin across them to spread RDMA/NIC load. Otherwise the + source TP rank is fixed by the local-rank mapping (_remote_tp_rank) -- + forward DP8EP->TP8 -> tp0; symmetric TP -> tp_rank -- byte-identical to + prior behaviour. chosen_tp is the single value threaded into the (dp, tp) + session key, the handshake dial and the notify port, so all three address + the SAME prefill rank (drift -> read one rank but notify another -> the + read rank's prefill buffer is never freed). + """ + remote_tp_size = int(meta.tp_size) + flexible = ( + self.world_size == 1 + and self.use_mla + and int(meta.remote_dp_size) == 1 + and remote_tp_size > 1 + ) + if flexible: + chosen_tp = self._next_flex_tp_rank(remote_tp_size) + else: + chosen_tp = self._remote_tp_rank(remote_tp_size) + return chosen_tp, flexible + def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): logger.debug( "Remote agent %s available, calling _read_blocks for req %s", meta.remote_engine_id, req_id, ) + chosen_tp, flexible = self._resolve_read_source(meta) self._read_blocks( request_id=req_id, transfer_id=meta.transfer_id, @@ -1824,6 +2034,9 @@ class MoRIIOConnectorWorker: remote_host=meta.remote_host, remote_notify_port=meta.remote_notify_port, remote_tp_size=meta.tp_size, + remote_dp_rank=meta.remote_dp_rank, + chosen_tp=chosen_tp, + flexible=flexible, ) def _write_blocks_for_req(self, req_id: ReqId, meta: ReqMeta, layer_name, kv_layer): @@ -1976,12 +2189,37 @@ class MoRIIOConnectorWorker: remote_host: str, remote_notify_port: int, remote_tp_size: int, + remote_dp_rank: int = 0, + chosen_tp: int | None = None, + flexible: bool = False, ) -> None: if self.mode == MoRIIOMode.WRITE: return - dp0_engine_id = self.get_engine_name_with_dp(dst_engine_id, 0) - sessions, remote_moriio_meta = self._get_built_session(dp0_engine_id) + # Read from the prefill rank that actually computed this request's KV + # (forwarded by the proxy). Hardcoding DP0 reads from a different rank's + # memory registration; per-rank num_blocks differ, so high block ids can + # overrun the wrong rank's region. + # + # eff_tp = the remote TP rank this read targets. The flexible mirror + # reads from a round-robin-chosen prefill TP rank and keys the session + # per (dp, tp); other configs use the fixed local-rank mapping (eff_tp == + # _remote_tp_rank), byte-identical to before. This key MUST match the one + # the eager handshake stored the session under. + eff_tp = ( + int(chosen_tp) + if chosen_tp is not None + else self._remote_tp_rank(remote_tp_size) + ) + if flexible: + remote_dp_engine_id = self.get_engine_name_with_dp_tp( + dst_engine_id, int(remote_dp_rank), eff_tp + ) + else: + remote_dp_engine_id = self.get_engine_name_with_dp( + dst_engine_id, int(remote_dp_rank) + ) + sessions, remote_moriio_meta = self._get_built_session(remote_dp_engine_id) # SQ-full backpressure deadline, shared across this request's layers. _sq_deadline = time.monotonic() + self.moriio_config.transfer_timeout @@ -2030,6 +2268,13 @@ class MoRIIOConnectorWorker: self._recving_transfers[request_id].append(transfer_status) self._recving_transfers_callback_addr[request_id] = ( remote_host, - str(remote_notify_port + self._remote_tp_rank(remote_tp_size)), + str( + remote_notify_port + + get_port_offset( + int(remote_dp_rank), + eff_tp, + remote_tp_size, + ) + ), transfer_id, ) From 439f336212227833e126526d3c5f3ef3968dfbf5 Mon Sep 17 00:00:00 2001 From: Nick Hill <nickhill123@gmail.com> Date: Mon, 27 Jul 2026 03:41:27 +0100 Subject: [PATCH 076/185] [Core] Fix gpu<->cpu syncs in MRV2 mamba_hybrid.py (#49736) Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Benjamin Chislett <bchislett@nvidia.com> --- vllm/v1/worker/gpu/model_states/mamba_hybrid.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index d04400afaea..03447c3cfad 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -101,12 +101,12 @@ class MambaHybridModelState(DefaultModelState): def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: super().add_request(req_index, new_req_data) # Must reset the speculative acceptance count in this idx which could be stale. - self.num_accepted_tokens_gpu[req_index] = 1 + self.num_accepted_tokens_gpu[req_index].fill_(1) if self._align_mode: # Seed the running state block from the resumed/prefilled position. - self._mamba_state_idx_gpu[req_index] = ( - new_req_data.num_computed_tokens - 1 - ) // self.cache_config.block_size + self._mamba_state_idx_gpu[req_index].fill_( + (new_req_data.num_computed_tokens - 1) // self.cache_config.block_size + ) def _get_mamba_group_info( self, kv_cache_config: KVCacheConfig From ac87549cbdff37f06c93f96f94005331772f2297 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas <akaratza@amd.com> Date: Sun, 26 Jul 2026 22:12:45 -0500 Subject: [PATCH 077/185] [CI][ROCm] Reduce V1 attention test runtime (#49916) Signed-off-by: Andreas Karatzas <akaratza@amd.com> --- .buildkite/test-amd.yaml | 26 ++----- tests/v1/attention/test_mla_backends.py | 91 ++++++++++++++++--------- 2 files changed, 66 insertions(+), 51 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index b67acbc40b3..ebd5d79cce0 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -364,22 +364,6 @@ steps: commands: - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness" -- label: V1 attention (H100-MI250) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/config/attention.py - - vllm/model_executor/layers/attention - - vllm/v1/attention - - tests/v1/attention - - vllm/_aiter_ops.py - - vllm/envs.py - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/attention - - label: V1 others (CPU) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] @@ -2694,11 +2678,12 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s v1/kv_connector/extract_hidden_states_integration -- label: V1 attention (H100-MI300) # TBD +- label: V1 attention (H100-MI300) %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 + parallelism: 2 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -2710,7 +2695,7 @@ steps: - vllm/envs.py - vllm/platforms/rocm.py commands: - - pytest -v -s v1/attention + - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: V1 Core + KV + Metrics # TBD timeout_in_minutes: 180 @@ -3676,11 +3661,12 @@ steps: #------------------------------------------------------------ mi355 · v1 -------------------------------------------------------------# -- label: V1 attention (B200-MI355) # TBD +- label: V1 attention (B200-MI355) %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false agent_pool: mi355_1 + parallelism: 2 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/config/attention.py @@ -3691,7 +3677,7 @@ steps: - vllm/envs.py - vllm/platforms/rocm.py commands: - - pytest -v -s v1/attention + - pytest -v -s v1/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: V1 Core + KV + Metrics # TBD timeout_in_minutes: 180 diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 7de61b32a7d..c520f22e7ee 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -40,6 +40,10 @@ from vllm.v1.attention.backends.mla.prefill import ( MLAPrefillBackendEnum, get_mla_prefill_backend, ) +from vllm.v1.attention.backends.mla.prefill.base import MLADimensions +from vllm.v1.attention.backends.mla.prefill.selector import ( + MLAPrefillSelectorConfig, +) from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported from vllm.v1.kv_cache_interface import ( @@ -145,14 +149,67 @@ def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch): torch.testing.assert_close(layer.W_UK_T, old_w_uk_t + 100) -# Filtered per-test via validate_configuration (capability/deps/dims). +# Validate parameter combinations during collection, before GPU fixtures run. PREFILL_BACKENDS_TO_TEST = [ + MLAPrefillBackendEnum.ROCM_AITER_FA, MLAPrefillBackendEnum.FLASH_ATTN, MLAPrefillBackendEnum.FLASHINFER, MLAPrefillBackendEnum.TRTLLM_RAGGED, MLAPrefillBackendEnum.TOKENSPEED_MLA, ] +MLA_DIMENSIONS_TO_TEST = [ + ("deepseek", 128, 128), + ("glm", 192, 256), +] + + +def _prefill_backend_dimension_params(): + device_capability = current_platform.get_device_capability() + params = [] + for prefill_backend in PREFILL_BACKENDS_TO_TEST: + for dimensions_id, qk_nope_head_dim, v_head_dim in MLA_DIMENSIONS_TO_TEST: + if device_capability is None: + invalid_reasons = ["device capability unavailable"] + else: + try: + invalid_reasons = ( + prefill_backend.get_class().validate_configuration( + device_capability, + MLAPrefillSelectorConfig( + dtype=torch.bfloat16, + mla_dimensions=MLADimensions( + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=64, + v_head_dim=v_head_dim, + ), + ), + ) + ) + except ImportError: + invalid_reasons = ["ImportError"] + + marks = [] + if invalid_reasons: + marks.append( + pytest.mark.skip( + reason=( + f"Prefill backend {prefill_backend.name} unavailable: " + f"{invalid_reasons}" + ) + ) + ) + params.append( + pytest.param( + prefill_backend, + qk_nope_head_dim, + v_head_dim, + id=f"{dimensions_id}-{prefill_backend}", + marks=marks, + ) + ) + return params + SPEC_DECODE_BACKENDS = [] for backend in BACKENDS_TO_TEST: @@ -1098,11 +1155,9 @@ def run_attention_backend( @pytest.mark.parametrize("tensor_parallel_size", [1, 4, 8, 16]) @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"]) @pytest.mark.parametrize(("q_scale", "k_scale"), [(1.0, 1.0), (2.0, 3.0)]) -@pytest.mark.parametrize("prefill_backend", PREFILL_BACKENDS_TO_TEST) @pytest.mark.parametrize( - ("qk_nope_head_dim", "v_head_dim"), - [(128, 128), (192, 256)], - ids=["deepseek", "glm"], + ("prefill_backend", "qk_nope_head_dim", "v_head_dim"), + _prefill_backend_dimension_params(), ) def test_backend_correctness( default_vllm_config, @@ -1153,32 +1208,6 @@ def test_backend_correctness( if not backends_to_test: pytest.skip(f"No backends support kv_cache_dtype={kv_cache_dtype}") - # Skip prefill backends that can't satisfy capability/deps/dimension constraints. - from vllm.v1.attention.backends.mla.prefill.base import MLADimensions - from vllm.v1.attention.backends.mla.prefill.selector import ( - MLAPrefillSelectorConfig, - ) - - try: - prefill_invalid_reasons = prefill_backend.get_class().validate_configuration( - current_platform.get_device_capability(), - MLAPrefillSelectorConfig( - dtype=torch.bfloat16, - mla_dimensions=MLADimensions( - qk_nope_head_dim=qk_nope_head_dim, - qk_rope_head_dim=64, - v_head_dim=v_head_dim, - ), - ), - ) - except ImportError: - prefill_invalid_reasons = ["ImportError"] - if prefill_invalid_reasons: - pytest.skip( - f"Prefill backend {prefill_backend.name} unavailable: " - f"{prefill_invalid_reasons}" - ) - batch_spec = BATCH_SPECS[batch_spec_name] is_spec_decode_test = batch_spec_name.startswith("spec_decode") unique_block_sizes = sorted(set(BACKEND_BLOCK_SIZES[b] for b in backends_to_test)) From 854c33f3801241574223644d25f80fbd735a44e5 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas <akaratza@amd.com> Date: Sun, 26 Jul 2026 22:19:55 -0500 Subject: [PATCH 078/185] [CI][ROCm] Keep global GPU memory cleanup opt-in (#49911) Signed-off-by: Andreas Karatzas <akaratza@amd.com> --- tests/conftest.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 2c5a907e557..6cb8e260c69 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1729,13 +1729,9 @@ def disable_deepgemm_ue8m0(monkeypatch): def _should_clean_gpu_memory_between_tests() -> bool: - setting = os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY") - if setting == "1": - return True - if setting == "0": - return False - # ROCm reclaims VRAM lazily; default to waiting between tests on ROCm CI. - return current_platform.is_rocm() + # This must stay opt-in: a function-scoped fixture cannot distinguish + # stale VRAM from allocations owned by longer-lived module/session fixtures. + return os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") == "1" @pytest.fixture(autouse=True) From bf4f633b4cb80e04ba28d2179efe1662d97a6fe0 Mon Sep 17 00:00:00 2001 From: Chaojun Zhang <chaojun.zhang@intel.com> Date: Mon, 27 Jul 2026 11:22:29 +0800 Subject: [PATCH 079/185] [XPU] Enable QK Norm + RoPE fusion pass on XPU (#49394) Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com> --- .buildkite/intel_jobs/misc_intel.yaml | 22 ++++++++++++++++++++++ vllm/compilation/passes/pass_manager.py | 4 ++-- vllm/platforms/xpu.py | 1 - 3 files changed, 24 insertions(+), 3 deletions(-) diff --git a/.buildkite/intel_jobs/misc_intel.yaml b/.buildkite/intel_jobs/misc_intel.yaml index 5ce12269fc4..07239e3b239 100644 --- a/.buildkite/intel_jobs/misc_intel.yaml +++ b/.buildkite/intel_jobs/misc_intel.yaml @@ -262,3 +262,25 @@ steps: pytest -v -s detokenizer && pytest -v -s -m "not cpu_test" ./multimodal && pytest -v -s utils_ --ignore=utils_/test_mem_utils.py' + +- label: Fusion Unit Tests + timeout_in_minutes: 30 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/compilation/ + - tests/compile/passes/test_qk_norm_rope_fusion.py + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pytest -v -s compile/passes/test_qk_norm_rope_fusion.py' \ No newline at end of file diff --git a/vllm/compilation/passes/pass_manager.py b/vllm/compilation/passes/pass_manager.py index 2fe863189c1..127be115c5d 100644 --- a/vllm/compilation/passes/pass_manager.py +++ b/vllm/compilation/passes/pass_manager.py @@ -30,19 +30,19 @@ if rocm_aiter_ops.is_enabled(): ) if current_platform.is_cuda_alike() or current_platform.is_xpu(): + from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass from .fusion.sequence_parallelism import SequenceParallelismPass + from .utility.split_coalescing import SplitCoalescingPass if current_platform.is_cuda_alike(): from .fusion.act_quant_fusion import ActivationQuantFusionPass from .fusion.attn_quant_fusion import AttnQuantFusionPass from .fusion.mla_attn_quant_fusion import MLAAttnQuantFusionPass from .fusion.mla_rope_kvcache_cat_fusion import MLARoPEKVCacheCatFusionPass - from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass from .fusion.qk_norm_rope_kvcache_fusion import QkNormRopeKvCacheFusionPass from .fusion.rms_quant_fusion import RMSNormQuantFusionPass from .fusion.rope_kvcache_fusion import RopeKVCacheFusionPass from .utility.scatter_split_replace import ScatterSplitReplacementPass - from .utility.split_coalescing import SplitCoalescingPass if current_platform.is_cuda(): from .fusion.allreduce_rms_fusion import AllReduceFusionPass diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index e673713330f..e9d1a905a4f 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -304,7 +304,6 @@ class XPUPlatform(Platform): "fuse_act_padding": "Activation + padding fusion", "fuse_rope_kvcache": "RoPE + KV cache fusion", "fuse_rope_kvcache_cat_mla": "RoPE + KV cache + MLA fusion", - "enable_qk_norm_rope_fusion": "QK Norm + RoPE fusion", } if compilation_config.mode != CompilationMode.NONE: for flag, feature_name in fusion_passes_to_disable.items(): From 8040ef242662af2f07614d78862c53338f046878 Mon Sep 17 00:00:00 2001 From: Walter Beller-Morales <walterbm@users.noreply.github.com> Date: Sun, 26 Jul 2026 23:30:37 -0400 Subject: [PATCH 080/185] [Frontend] expose stream_interval as req sampling param (#49754) Signed-off-by: walterbm <walter.beller.morales@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/engine/test_output_processor.py | 76 +++++++++++++++++++ .../openai/chat_completion/protocol.py | 11 +++ .../entrypoints/openai/completion/protocol.py | 11 +++ vllm/sampling_params.py | 13 ++++ vllm/v1/engine/output_processor.py | 3 + 5 files changed, 114 insertions(+) diff --git a/tests/v1/engine/test_output_processor.py b/tests/v1/engine/test_output_processor.py index 1919349790f..51dbb9c9895 100644 --- a/tests/v1/engine/test_output_processor.py +++ b/tests/v1/engine/test_output_processor.py @@ -141,6 +141,82 @@ def test_incremental_detokenization( assert not output_processor.has_unfinished_requests() +def test_request_stream_interval_raises_but_not_below_engine_default( + dummy_test_vectors, +): + """A per-request stream_interval can raise the interval above the engine + default but not below it (values under the default clamp up), without + altering the generated text.""" + engine_stream_interval = 5 + # Request 0 (below the default) clamps up to 5; request 1 raises it to 10. + request_stream_intervals = [1, 10] + output_processor = OutputProcessor( + dummy_test_vectors.tokenizer, + log_stats=False, + stream_interval=engine_stream_interval, + ) + + requests = [ + EngineCoreRequest( + request_id=f"request-{idx}-int", + external_req_id=f"request-{idx}", + prompt_token_ids=prompt_tokens, + mm_features=None, + arrival_time=0, + lora_request=None, + cache_salt=None, + data_parallel_rank=None, + sampling_params=SamplingParams( + skip_special_tokens=False, + spaces_between_special_tokens=False, + output_kind=RequestOutputKind.DELTA, + stop=[], + include_stop_str_in_output=False, + stream_interval=request_stream_intervals[idx], + ), + pooling_params=None, + ) + for idx, prompt_tokens in enumerate( + dummy_test_vectors.prompt_tokens[: len(request_stream_intervals)] + ) + ] + + num_requests = len(requests) + engine_core = MockEngineCore( + tokens_list=dummy_test_vectors.generation_tokens[:num_requests], + prompts_list=dummy_test_vectors.prompt_tokens[:num_requests], + request_ids=[req.request_id for req in requests], + ) + + for request, prompt in zip(requests, dummy_test_vectors.prompt_strings): + output_processor.add_request(request, prompt) + + gen_strings: dict[str, str] = {} + gen_tokens: dict[str, list[int]] = {} + while outputs := engine_core.get_outputs(): + for request_output in output_processor.process_outputs(outputs).request_outputs: + request_id = request_output.request_id + new_tokens = request_output.outputs[0].token_ids + if request_id not in gen_strings: + gen_strings[request_id] = request_output.outputs[0].text + gen_tokens[request_id] = list(new_tokens) + assert len(new_tokens) == 1, f"{len(new_tokens)=}" + continue + gen_strings[request_id] += request_output.outputs[0].text + gen_tokens[request_id].extend(new_tokens) + if not request_output.finished: + requested = request_stream_intervals[int(request_id.split("-")[1])] + interval = max(requested, engine_stream_interval) + assert len(new_tokens) == interval, f"{len(new_tokens)=}, {interval=}" + + for idx in range(num_requests): + request_id = f"request-{idx}" + assert gen_strings[request_id] == dummy_test_vectors.generation_strings[idx] + assert gen_tokens[request_id] == dummy_test_vectors.generation_tokens[idx] + + assert not output_processor.has_unfinished_requests() + + def _validate_logprobs( gen_tokens: dict[str, list[int]], gen_logprobs: dict[str, SampleLogprobs | None], diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 3b65faa5aa0..505bd01d4cf 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -476,6 +476,16 @@ class ChatCompletionRequest(OpenAIBaseModel): "can detect such behavior and terminate early, saving time and tokens.", ) + stream_interval: Annotated[int, Field(ge=1)] | None = Field( + default=None, + description=( + "Number of tokens to batch into each streamed chunk. Raises the " + "server's `--stream-interval` for this request. Values below the " + "server setting are clamped up to it. The first and last chunks " + "are always sent immediately. Ignored for non-streaming requests." + ), + ) + # --8<-- [end:chat-completion-extra-params] @model_validator(mode="before") @@ -690,6 +700,7 @@ class ChatCompletionRequest(OpenAIBaseModel): output_kind=( RequestOutputKind.DELTA if self.stream else RequestOutputKind.FINAL_ONLY ), + stream_interval=self.stream_interval, structured_outputs=self.extract_structured_outputs(), logit_bias=self.logit_bias, bad_words=self.bad_words, diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index 6a2bf47bb15..79ad9799e18 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -233,6 +233,16 @@ class CompletionRequest(OpenAIBaseModel): ), ) + stream_interval: Annotated[int, Field(ge=1)] | None = Field( + default=None, + description=( + "Number of tokens to batch into each streamed chunk. Raises the " + "server's `--stream-interval` for this request. Values below the " + "server setting are clamped up to it. The first and last chunks " + "are always sent immediately. Ignored for non-streaming requests." + ), + ) + # --8<-- [end:completion-extra-params] def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: @@ -365,6 +375,7 @@ class CompletionRequest(OpenAIBaseModel): output_kind=RequestOutputKind.DELTA if self.stream else RequestOutputKind.FINAL_ONLY, + stream_interval=self.stream_interval, structured_outputs=self.extract_structured_outputs(), logit_bias=self.logit_bias, allowed_token_ids=self.allowed_token_ids, diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 08580f6e8f6..25e36ceb568 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -299,6 +299,11 @@ class SamplingParams( include_stop_str_in_output: bool = False """Whether to include the stop strings in output text.""" output_kind: RequestOutputKind = RequestOutputKind.CUMULATIVE + stream_interval: int | None = None + """Number of newly generated tokens to batch into each streamed + `RequestOutput`. Raises the interval above the engine-level + `--stream-interval`. Values below engine setting are clamped up to it. + The first and final outputs are always emitted immediately.""" skip_clone: bool = False """Internal flag indicating that this SamplingParams instance is safe to reuse without cloning. When True, clone() will return self without @@ -377,6 +382,7 @@ class SamplingParams( skip_special_tokens: bool = True, spaces_between_special_tokens: bool = True, output_kind: RequestOutputKind = RequestOutputKind.CUMULATIVE, + stream_interval: int | None = None, structured_outputs: StructuredOutputsParams | None = None, logit_bias: dict[int, float] | dict[str, float] | None = None, allowed_token_ids: list[int] | None = None, @@ -439,6 +445,7 @@ class SamplingParams( skip_special_tokens=skip_special_tokens, spaces_between_special_tokens=spaces_between_special_tokens, output_kind=output_kind, + stream_interval=stream_interval, structured_outputs=structured_outputs, logit_bias=logit_bias, allowed_token_ids=allowed_token_ids, @@ -585,6 +592,12 @@ class SamplingParams( f"min_tokens must be less than or equal to " f"max_tokens={self.max_tokens}, got {self.min_tokens}." ) + if self.stream_interval is not None and self.stream_interval < 1: + raise VLLMValidationError( + f"stream_interval must be at least 1, got {self.stream_interval}.", + parameter="stream_interval", + value=self.stream_interval, + ) if self.logprobs is not None and self.logprobs != -1 and self.logprobs < 0: raise VLLMValidationError( f"logprobs must be non-negative or -1, got {self.logprobs}.", diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index 0d7d5fe18c9..c467703759e 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -224,6 +224,9 @@ class RequestState: if not sampling_params.detokenize: tokenizer = None output_kind = sampling_params.output_kind + if sampling_params.stream_interval is not None: + # clamp to the engine-level stream interval. + stream_interval = max(sampling_params.stream_interval, stream_interval) logprobs_processor = LogprobsProcessor.from_new_request( tokenizer=tokenizer, request=request, From da99ffcc13362ab446a22a9bc6ed36c4c14942e3 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas <akaratza@amd.com> Date: Sun, 26 Jul 2026 22:31:24 -0500 Subject: [PATCH 081/185] [ROCm][CI] Keep native datasets cache off shared NFS (#49516) Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> --- .buildkite/scripts/hardware_ci/run-amd-test.sh | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index 0f6bb2564c9..ed4a9ba3a2e 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -405,11 +405,14 @@ initialize_native_environment() { VLLM_CACHE_ROOT="${native_root}/cache/vllm" XDG_CACHE_HOME="${native_root}/cache/xdg" : "${HF_HOME:=/home/buildkite-agent/huggingface}" + # datasets uses POSIX locks that are unsupported by the shared HF NFS cache. + # Keep processed datasets job-local while retaining the persistent Hub cache. + HF_DATASETS_CACHE="${native_root}/cache/huggingface/datasets" : "${HF_HUB_DOWNLOAD_TIMEOUT:=300}" : "${HF_HUB_ETAG_TIMEOUT:=60}" export TMPDIR VLLM_RPC_BASE_PATH export TORCHINDUCTOR_CACHE_DIR TRITON_CACHE_DIR VLLM_CACHE_ROOT XDG_CACHE_HOME - export HF_HOME HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT + export HF_HOME HF_DATASETS_CACHE HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT export PYTORCH_ROCM_ARCH="" mkdir -p "${TMPDIR}" \ @@ -417,7 +420,8 @@ initialize_native_environment() { "${TRITON_CACHE_DIR}" \ "${VLLM_CACHE_ROOT}" \ "${XDG_CACHE_HOME}" \ - "${HF_HOME}" || return 1 + "${HF_HOME}" \ + "${HF_DATASETS_CACHE}" || return 1 echo "Native compile caches: VLLM_CACHE_ROOT=${VLLM_CACHE_ROOT} TORCHINDUCTOR_CACHE_DIR=${TORCHINDUCTOR_CACHE_DIR}" From 8de50e46d4a0a227db1888410cb6ac93058fb6cf Mon Sep 17 00:00:00 2001 From: Harjoth Khara <harjoth.khara@gmail.com> Date: Sun, 26 Jul 2026 20:53:28 -0700 Subject: [PATCH 082/185] [Docs] Document NVFP4 GEMM kernel selection and Marlin weight-only fallback (#49376) Signed-off-by: harjoth <harjoth.khara@gmail.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> --- docs/features/quantization/modelopt.md | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/docs/features/quantization/modelopt.md b/docs/features/quantization/modelopt.md index ad417bcb30a..7850dc75ccc 100644 --- a/docs/features/quantization/modelopt.md +++ b/docs/features/quantization/modelopt.md @@ -19,6 +19,20 @@ following `quantization.quant_algo` values: - `NVFP4`: ModelOpt NVFP4 checkpoints (use `quantization="modelopt_fp4"`). - `MXFP8`: ModelOpt MXFP8 checkpoints (use `quantization="modelopt_mxfp8"`). +!!! note + For NVFP4 checkpoints, vLLM selects a GEMM kernel automatically at load + time from the backends available on the current platform (CUTLASS, + FlashInfer, Marlin, and others). On GPUs without a supported native FP4 + GEMM kernel, vLLM falls back to weight-only (W4A16) execution via Marlin + and logs a warning; this may reduce throughput for compute-heavy + workloads. Use `--linear-backend` to override the automatic selection + (this replaces the deprecated `VLLM_NVFP4_GEMM_BACKEND` environment + variable). Values relevant to NVFP4 include `cutlass`, + `flashinfer_cutlass`, `flashinfer_trtllm`, `flashinfer_cudnn`, and + `marlin`; the full list is documented under `KernelConfig` on the + [Engine Arguments](../../configuration/engine_args.md) page and shown by + `vllm serve --help=KernelConfig`. + ## Quantizing HuggingFace Models with PTQ You can quantize HuggingFace models using the example scripts provided in the Model Optimizer repository. The primary script for LLM PTQ is typically found within the `examples/llm_ptq` directory. From ff6173997d54c5027971df8ecd1280f046a832b3 Mon Sep 17 00:00:00 2001 From: jcotant-inferact <joe@inferact.ai> Date: Sun, 26 Jul 2026 21:00:27 -0700 Subject: [PATCH 083/185] [CI] Add kimi and k3 auto-labeling rules (#49895) Signed-off-by: Joe Cotant <joe@inferact.ai> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Roger Wang <hey@rogerw.io> --- .github/mergify.yml | 25 +++++++++++++++++++++++++ .github/workflows/issue_autolabel.yml | 19 +++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/.github/mergify.yml b/.github/mergify.yml index 5615e78d0c2..faaf179ebdf 100644 --- a/.github/mergify.yml +++ b/.github/mergify.yml @@ -233,6 +233,31 @@ pull_request_rules: add: - gpt-oss +- name: label-kimi + description: Automatically apply kimi label + conditions: + - label != stale + - or: + - files~=(?i)kimi + - files~=(?i)moonshot + - title~=(?i)(?:kimi|moonshot) + actions: + label: + add: + - kimi + +- name: label-k3 + description: Automatically apply k3 label (launch triage; retire after ramp-down) + conditions: + - label != stale + - or: + - files~=(?i)kimi[-_]?k3 + - title~=(?i)(?:kimi[-\s]?k3|\bk3\b) + actions: + label: + add: + - k3 + - name: label-nvidia description: Automatically apply nvidia label conditions: diff --git a/.github/workflows/issue_autolabel.yml b/.github/workflows/issue_autolabel.yml index b758c967c1c..5e07d228b7f 100644 --- a/.github/workflows/issue_autolabel.yml +++ b/.github/workflows/issue_autolabel.yml @@ -130,6 +130,25 @@ jobs: }, ], }, + kimi: { + keywords: [ + { term: "Kimi", searchIn: "both" }, + { term: "Moonshot", searchIn: "both" }, + ], + substrings: [ + { term: "moonshotai/", searchIn: "both" }, + { term: "kimi", searchIn: "title" }, + ], + }, + k3: { + keywords: [ + { term: "Kimi K3", searchIn: "both" }, + { term: "K3", searchIn: "title" }, + ], + substrings: [ + { term: "moonshotai/kimi-k3", searchIn: "both" }, + ], + }, quantization: { keywords: [ { From 29fdeab2548fdb5cd61ec3613e5e72200c695ef0 Mon Sep 17 00:00:00 2001 From: xiangdong <40376367+zxd1997066@users.noreply.github.com> Date: Mon, 27 Jul 2026 12:40:22 +0800 Subject: [PATCH 084/185] [XPU][CI] Add more test cases in Intel GPU CI (#49422) Signed-off-by: zengxian <xiangdong.zeng@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com> --- .buildkite/intel_jobs/benchmarks_intel.yaml | 26 +++++++ .buildkite/intel_jobs/engine_intel.yaml | 76 +++++++++++++++++++ .../intel_jobs/model_executor_intel.yaml | 33 ++++++++ .../intel_jobs/model_runner_v2_intel.yaml | 56 +++++++++++++- .buildkite/intel_jobs/samplers_intel.yaml | 29 +++++++ 5 files changed, 219 insertions(+), 1 deletion(-) create mode 100644 .buildkite/intel_jobs/benchmarks_intel.yaml create mode 100644 .buildkite/intel_jobs/model_executor_intel.yaml create mode 100644 .buildkite/intel_jobs/samplers_intel.yaml diff --git a/.buildkite/intel_jobs/benchmarks_intel.yaml b/.buildkite/intel_jobs/benchmarks_intel.yaml new file mode 100644 index 00000000000..9ce373e4f42 --- /dev/null +++ b/.buildkite/intel_jobs/benchmarks_intel.yaml @@ -0,0 +1,26 @@ +group: Benchmarks +depends_on: + - image-build-xpu +steps: +- label: Benchmarks CLI Test + key: benchmarks-cli-test + timeout_in_minutes: 40 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/ + - tests/benchmarks/ + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pytest -v -s benchmarks/' diff --git a/.buildkite/intel_jobs/engine_intel.yaml b/.buildkite/intel_jobs/engine_intel.yaml index d1dc95b1d40..5f9357b2a77 100644 --- a/.buildkite/intel_jobs/engine_intel.yaml +++ b/.buildkite/intel_jobs/engine_intel.yaml @@ -2,6 +2,44 @@ group: Engine Intel depends_on: - image-build-xpu steps: +- label: Engine + key: engine + timeout_in_minutes: 40 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/compilation/ + - vllm/config/ + - vllm/engine/ + - vllm/entrypoints/logger.py + - vllm/envs.py + - vllm/logger.py + - vllm/logging_utils/ + - vllm/platforms/ + - vllm/sequence.py + - vllm/triton_utils/ + - vllm/utils/ + - tests/engine + - tests/test_sequence + - tests/test_config + - tests/test_logger + - tests/test_vllm_port + - tests/test_jit_monitor.py + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pytest -v -s engine/test_arg_utils.py test_sequence.py test_logger.py test_vllm_port.py test_jit_monitor.py' + - label: Engine (1 GPU) timeout_in_minutes: 30 device: intel_gpu @@ -23,3 +61,41 @@ steps: bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py' + +- label: V1 e2e (2 GPUs) + timeout_in_minutes: 30 + device: intel_gpu + agent_tags: + label: production + gpu: 2+ + mem: 16+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/compilation/ + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/envs.py + - vllm/forward_context.py + - vllm/inputs/ + - vllm/logger.py + - vllm/logging_utils/ + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/platforms/ + - vllm/sampling_params.py + - vllm/transformers_utils/ + - vllm/triton_utils/ + - vllm/utils/ + - vllm/v1/ + - tests/v1/e2e/spec_decode + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"' diff --git a/.buildkite/intel_jobs/model_executor_intel.yaml b/.buildkite/intel_jobs/model_executor_intel.yaml new file mode 100644 index 00000000000..14a853544ec --- /dev/null +++ b/.buildkite/intel_jobs/model_executor_intel.yaml @@ -0,0 +1,33 @@ +group: Model Executor Intel +depends_on: + - image-build-xpu +steps: +- label: Model Executor (Intel) + key: model-executor-intel + timeout_in_minutes: 45 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/engine/arg_utils.py + - vllm/config/model.py + - vllm/model_executor + - tests/model_executor + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'apt-get update && apt-get install -y curl libsodium23 && + pip3 install tensorizer==2.10.1 && + pip3 install runai-model-streamer[s3,gcs,azure]\>=0.15.7 && + export VLLM_WORKER_MULTIPROC_METHOD=spawn && + export PYTHONFAULTHANDLER=1 && + cd tests && + pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --deselect="tests/model_executor/model_loader/test_reload.py::test_kv_scale_reload"' diff --git a/.buildkite/intel_jobs/model_runner_v2_intel.yaml b/.buildkite/intel_jobs/model_runner_v2_intel.yaml index 0311b5dffb7..43ab989fbbd 100644 --- a/.buildkite/intel_jobs/model_runner_v2_intel.yaml +++ b/.buildkite/intel_jobs/model_runner_v2_intel.yaml @@ -8,7 +8,7 @@ steps: agent_tags: label: production gpu: 2+ - mem: 16+ + mem: 24+ no_plugin: true working_dir: "." env: @@ -28,7 +28,9 @@ steps: 'export VLLM_USE_V2_MODEL_RUNNER=1 && cd tests && pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics" && + pytest -v -s v1/e2e/general/test_context_length.py && ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram" && + pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0" && pytest -v -s v1/e2e/general/test_min_tokens.py' - label: Model Runner V2 Examples (Intel) @@ -60,3 +62,55 @@ steps: python3 basic/offline_inference/generate.py --model facebook/opt-125m && python3 generate/multimodal/vision_language_offline.py --seed 0 && python3 features/automatic_prefix_caching/prefix_caching_offline.py' + +- label: Model Runner V2 Distributed (2 GPUs) + timeout_in_minutes: 50 + device: intel_gpu + agent_tags: + label: production + gpu: 2+ + mem: 16+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/v1/worker/gpu/ + - vllm/v1/worker/gpu_worker.py + - tests/basic_correctness/test_basic_correctness.py + - tests/v1/distributed/test_async_llm_dp.py + - tests/v1/distributed/test_eagle_dp.py + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'export VLLM_USE_V2_MODEL_RUNNER=1 && + cd tests && + TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -m "distributed\(num_gpus=2\)" -k "not ray and not True"' + +- label: Model Runner V2 Spec Decode + timeout_in_minutes: 50 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/v1/worker/gpu/ + - vllm/v1/worker/gpu_worker.py + - tests/v1/spec_decode/test_max_len.py + - tests/v1/spec_decode/test_rejection_sampler_utils.py + - tests/v1/e2e/spec_decode/test_spec_decode.py + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'export VLLM_USE_V2_MODEL_RUNNER=1 && + cd tests && + pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py' diff --git a/.buildkite/intel_jobs/samplers_intel.yaml b/.buildkite/intel_jobs/samplers_intel.yaml new file mode 100644 index 00000000000..acd94803202 --- /dev/null +++ b/.buildkite/intel_jobs/samplers_intel.yaml @@ -0,0 +1,29 @@ +group: Samplers Intel +depends_on: + - image-build-xpu +steps: +- label: Samplers Test (FlashInfer) + key: samplers-test-flashinfer-intel + timeout_in_minutes: 40 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/model_executor/layers + - vllm/sampling_metadata.py + - tests/samplers + - tests/conftest.py + - vllm/entrypoints/generate/beam_search + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers' From 74d3b799e193756a8d42b59d223e38a39fc6be55 Mon Sep 17 00:00:00 2001 From: Nick Hill <nickhill123@gmail.com> Date: Sun, 26 Jul 2026 21:42:47 -0700 Subject: [PATCH 085/185] [Bugfix] Fix mHC block-M prenorm GEMM cross-row reduction carry-over (#49429) Signed-off-by: Nick Hill <nickhill123@gmail.com> --- .../model_executor/kernels/mhc/tilelang_kernels.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/vllm/model_executor/kernels/mhc/tilelang_kernels.py b/vllm/model_executor/kernels/mhc/tilelang_kernels.py index 50abac3e010..3afd7069aa3 100644 --- a/vllm/model_executor/kernels/mhc/tilelang_kernels.py +++ b/vllm/model_executor/kernels/mhc/tilelang_kernels.py @@ -856,21 +856,23 @@ def hc_prenorm_gemm_block_m_tilelang( T.sync_threads() if warp_id == 0: + reduced_acc = T.alloc_local((block_m,), T.float32) + reduced_sqr = T.alloc_local((block_m,), T.float32) + T.clear(reduced_acc) + T.clear(reduced_sqr) for i_m in T.unroll(block_m): token_idx = i_mt * block_m + i_m if token_idx < num_tokens: if lane < tile_n: - reduced_acc = T.alloc_var(T.float32, init=0.0) for i_w in T.unroll(num_warps): - reduced_acc += warp_acc[i_w, i_m, lane] + reduced_acc[i_m] += warp_acc[i_w, i_m, lane] out_idx = i_t * tile_n + lane if out_idx < n_out: - out[0, token_idx, out_idx] = reduced_acc + out[0, token_idx, out_idx] = reduced_acc[i_m] if lane == 0 and i_t == 0: - reduced_sqr = T.alloc_var(T.float32, init=0.0) for i_w in T.unroll(num_warps): - reduced_sqr += warp_sqr[i_w, i_m] - sqrsum[0, token_idx] = reduced_sqr + reduced_sqr[i_m] += warp_sqr[i_w, i_m] + sqrsum[0, token_idx] = reduced_sqr[i_m] if ENABLE_PDL: T.pdl_trigger() From 49f31d7cee425a6d38f8c5bc76877986daf832ed Mon Sep 17 00:00:00 2001 From: Andreas Karatzas <akaratza@amd.com> Date: Sun, 26 Jul 2026 23:59:14 -0500 Subject: [PATCH 086/185] [ROCm] Make vllm_c RMSNorm output contiguous (#49913) Signed-off-by: Andreas Karatzas <akaratza@amd.com> --- tests/kernels/ir/test_layernorm.py | 23 +++++++++++++++++++++++ vllm/kernels/vllm_c.py | 4 +++- 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/tests/kernels/ir/test_layernorm.py b/tests/kernels/ir/test_layernorm.py index 37ec1468a80..b25da988a7a 100644 --- a/tests/kernels/ir/test_layernorm.py +++ b/tests/kernels/ir/test_layernorm.py @@ -176,6 +176,29 @@ def test_vllm_c_rms_norm_accepts_nd_input(): assert_close(ir.ops.rms_norm, output, ref_output) +@pytest.mark.skipif( + not current_platform.is_rocm(), + reason="ROCm vllm_c RMSNorm needs a contiguous output for strided inputs", +) +def test_vllm_c_rms_norm_accepts_transposed_input(): + impl = ir.ops.rms_norm.impls["vllm_c"] + if not impl.supported: + pytest.skip("vllm_c impl not supported on this platform") + + x = torch.randn( + 1, 320, 120, dtype=torch.float16, device=current_platform.device_type + ).transpose(1, 2) + assert x.reshape(-1, x.shape[-1]).stride(-1) != 1 + weight = torch.randn(320, dtype=torch.float16, device=current_platform.device_type) + epsilon = 1e-5 + + output = impl.impl_fn(x, weight, epsilon) + ref_output = rms_norm_native(x, weight, epsilon) + + assert output.shape == x.shape + assert_close(ir.ops.rms_norm, output, ref_output) + + fused_add_rms_norm_native = ir.ops.fused_add_rms_norm.impls["native"].impl_fn diff --git a/vllm/kernels/vllm_c.py b/vllm/kernels/vllm_c.py index 6ae5d9939e3..0df4011f6de 100644 --- a/vllm/kernels/vllm_c.py +++ b/vllm/kernels/vllm_c.py @@ -32,7 +32,9 @@ def rms_norm( if IS_ROCM and (x.dim() > 2 or not x.is_contiguous()): original_shape = x.shape x = x.reshape(-1, original_shape[-1]) - output = torch.empty_like(x) + # empty_like preserves the strides of transposed inputs, but the + # libtorch-stable kernel requires a contiguous output tensor. + output = torch.empty(x.shape, device=x.device, dtype=x.dtype) torch.ops._C.rms_norm(output, x, weight, epsilon) return output.reshape(original_shape) From 53397fbfac4dcd421b9d467199068518b135bc34 Mon Sep 17 00:00:00 2001 From: Jason <31175216+wsyjh8@users.noreply.github.com> Date: Mon, 27 Jul 2026 01:07:00 -0400 Subject: [PATCH 087/185] [Bugfix][KV Offload][P2P] Fix EngineCore crash reconnecting to a reaped peer (#49823) Signed-off-by: Jason Yao <wsyjh8@gmail.com> Co-authored-by: Or Ozeri <oro@il.ibm.com> --- .../tiering/p2p/test_zmq_transport.py | 108 ++++++++++++++++++ .../v1/kv_offload/tiering/p2p/control/base.py | 5 + vllm/v1/kv_offload/tiering/p2p/control/zmq.py | 67 ++++++++--- 3 files changed, 167 insertions(+), 13 deletions(-) diff --git a/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py b/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py index 101d4291cfd..f814a04887c 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py +++ b/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py @@ -219,6 +219,7 @@ class TestZmqTransportConnectivity: # Pruning is synchronous within poll(). transport_a.poll() assert len(transport_a._connections) == 0 + assert new_conns[0]._sockets.dealer.closed finally: transport_a.close() transport_b.close() @@ -228,3 +229,110 @@ class TestZmqTransportConnectivity: transport, _ = _make_transport() transport.close() transport.close() # should not raise + + +class TestZmqReconnect: + """Reconnecting to a peer whose connection died (real ZMQ sockets). + + A session marks its connection dead while handling messages, which happens + after the transport's own sweep has run for that tick — so a dead + connection stays registered until the next poll(). Reconnecting in that + window must succeed, and the retired connection must release its sockets. + Real sockets are required: a mock reports every attribute as closed. + """ + + def test_close_after_mark_dead_releases_sockets(self): + """close() releases sockets even when mark_dead() ran first. + + mark_dead() must not set the flag close() guards on, or every peer + disconnect leaks a DEALER and a monitor socket. + """ + transport, _ = _make_transport() + try: + conn = transport.connect(f"127.0.0.1:{_free_port()}") + dealer, monitor = conn._sockets.dealer, conn._sockets.monitor + + conn.mark_dead() + assert not conn.alive + assert not dealer.closed + + conn.close() + assert dealer.closed + assert monitor.closed + finally: + transport.close() + + def test_connect_retires_dead_connection(self): + """connect() replaces a registered-but-dead connection.""" + transport, _ = _make_transport() + try: + # The peer port is never bound — only the peer id matters here. + peer_id = f"127.0.0.1:{_free_port()}" + dead = transport.connect(peer_id) + dead.mark_dead() + + conn = transport.connect(peer_id) + + assert conn is not dead + assert conn.alive + assert transport._connections[peer_id] is conn + assert dead._sockets.dealer.closed + finally: + transport.close() + + def test_repeated_reconnect_to_same_peer(self): + """A flapping peer stays reconnectable. + + Monitor endpoints are inproc addresses that libzmq releases + asynchronously, so deriving one from peer_id alone makes each + reconnect race the previous teardown and fail with EADDRINUSE. + """ + transport, _ = _make_transport() + try: + peer_id = f"127.0.0.1:{_free_port()}" + for _ in range(10): + conn = transport.connect(peer_id) + conn.mark_dead() + transport.poll() + assert conn._sockets.dealer.closed + + assert not transport._connections + finally: + transport.close() + + def test_inbound_message_survives_dead_registration(self): + """A reconnecting peer's first message is not dropped. + + poll() must retire connections killed by their session before routing + traffic, otherwise the message is enqueued into the dead connection and + discarded when it is swept — and a session announces itself only once. + Covers only that between-polls window: a peer dying while poll() runs, + or dying silently until the heartbeat expires, is out of scope. + """ + transport_a, port_a = _make_transport() + transport_b, _ = _make_transport() + + try: + conn_b = transport_b.connect(f"127.0.0.1:{port_a}") + conn_b.send({"type": "connect", "seq": 1}) + + inbound = _wait_for_inbound(transport_a)[0] + _wait_for_messages(transport_a, inbound, 1) + + inbound.mark_dead() + conn_b.send({"type": "connect", "seq": 2}) + + # Wait until the frame is readable on the ROUTER, so the message is + # known to have arrived rather than merely being slow. + poller = zmq.Poller() + poller.register(transport_a._router, zmq.POLLIN) + assert poller.poll(2000), "message never reached the ROUTER" + + new_conns = _wait_for_inbound(transport_a) + assert len(new_conns) == 1 + assert new_conns[0] is not inbound + msgs = _wait_for_messages(transport_a, new_conns[0], 1) + assert msgs == [{"type": "connect", "seq": 2}] + finally: + transport_a.close() + transport_b.close() diff --git a/vllm/v1/kv_offload/tiering/p2p/control/base.py b/vllm/v1/kv_offload/tiering/p2p/control/base.py index 6b4d4cfcb59..7016cfb162a 100644 --- a/vllm/v1/kv_offload/tiering/p2p/control/base.py +++ b/vllm/v1/kv_offload/tiering/p2p/control/base.py @@ -97,6 +97,7 @@ class ControlConnection(ABC): After this call, alive returns False. The session should stop using this connection and the transport will clean it up. + Resources are not released here — close() must still run. """ ... @@ -134,6 +135,10 @@ class ControlTransport(ABC): The connection's send queue is live immediately — messages sent before the remote peer's poll() will be buffered. + + A connection to peer_id that is registered but no longer alive is + retired and replaced; a live one is a duplicate and must not be + replaced. """ ... diff --git a/vllm/v1/kv_offload/tiering/p2p/control/zmq.py b/vllm/v1/kv_offload/tiering/p2p/control/zmq.py index 4e9069a471c..42dc2c69846 100644 --- a/vllm/v1/kv_offload/tiering/p2p/control/zmq.py +++ b/vllm/v1/kv_offload/tiering/p2p/control/zmq.py @@ -55,12 +55,16 @@ class ZmqConnection(ControlConnection): def __init__(self, peer_id: str, sockets: _Sockets) -> None: super().__init__(peer_id) self._sockets = sockets + # _dead: the peer is gone. _closed: the sockets have been released. + # Distinct, so mark_dead() cannot turn close() into a no-op and leak + # the DEALER and its monitor socket. + self._dead = False self._closed = False self._inbox: list[dict] = [] def send(self, msg: dict) -> None: """Send a msgpack-encoded message to this peer.""" - if self._closed: + if not self.alive: raise RuntimeError( f"ZmqConnection: send on closed connection to {self.peer_id}" ) @@ -77,12 +81,13 @@ class ZmqConnection(ControlConnection): @property def alive(self) -> bool: - return not self._closed + return not (self._dead or self._closed) def close(self) -> None: if self._closed: return self._closed = True + self._dead = True logger.info("ZmqConnection: closing connection to %s", self.peer_id) self._sockets.monitor.close() self._sockets.dealer.close() @@ -92,8 +97,12 @@ class ZmqConnection(ControlConnection): self._inbox.append(msg) def mark_dead(self) -> None: - """Mark connection as disconnected.""" - self._closed = True + """Mark connection as disconnected. + + Only flips liveness: the sockets stay open until close() releases + them, so the owner can still drain recv() before tearing down. + """ + self._dead = True @property def monitor_socket(self) -> zmq.Socket: @@ -114,6 +123,9 @@ class ZmqTransport(ControlTransport): self._connections: dict[str, ZmqConnection] = {} self._pending_inbound: list[tuple[str, dict]] = [] + # Monotonic suffix for inproc monitor endpoints — see + # _open_connection() for why peer_id alone is not enough. + self._monitor_seq = 0 self._zmq_ctx = zmq.Context() self._router: zmq.Socket = self._zmq_ctx.socket(zmq.ROUTER) @@ -127,10 +139,23 @@ class ZmqTransport(ControlTransport): # ------------------------------------------------------------------ def connect(self, peer_id: str) -> ZmqConnection: - """Open an outbound connection to a remote peer.""" - assert peer_id not in self._connections, ( - f"ZmqConnection to {peer_id} already exists" - ) + """Open an outbound connection to a remote peer. + + A dead connection can still be registered: its owning session may mark + it dead after this tick's sweep already ran, and poll() only + unregisters it on the next pass. Retire such an entry instead of + asserting, so a reconnect landing in that window succeeds. A live + entry is still a genuine duplicate. + """ + existing = self._connections.get(peer_id) + if existing is not None: + assert not existing.alive, f"ZmqConnection to {peer_id} already exists" + logger.info( + "ZmqTransport %s: retiring dead connection to %s before reconnect", + self._local_id, + peer_id, + ) + self._connections.pop(peer_id).close() logger.info( "ZmqTransport %s: opening OUTBOUND connection to %s", self._local_id, @@ -146,8 +171,18 @@ class ZmqTransport(ControlTransport): - Checks monitors for disconnections - Removes and closes dead connections """ + # Retire connections a session killed since the last poll() before + # routing traffic: otherwise a reconnecting peer's first message is + # enqueued into the dead connection and discarded along with it. + # This closes only that between-polls window. A peer that dies while + # poll() is running is not seen until the next _check_monitors(), and + # one that dies silently not until the ZMQ heartbeat expires; both + # still take a message into a doomed connection and are out of scope. + self._sweep_dead_connections() + self._recv_router() self._check_monitors() + self._sweep_dead_connections() # Create connections for new inbound peers new_connections: list[ControlConnection] | None = None @@ -166,10 +201,6 @@ class ZmqTransport(ControlTransport): conn.enqueue(msg) self._pending_inbound.clear() - # Remove dead connections - for pid in [p for p, c in self._connections.items() if not c.alive]: - self._connections.pop(pid).close() - return ( new_connections if new_connections is not None else _EMPTY_NEW_CONNECTIONS ) @@ -210,8 +241,13 @@ class ZmqTransport(ControlTransport): _apply_heartbeat(dealer) dealer.identity = self._local_id.encode() + # Unique per connection, not per peer: libzmq releases an inproc + # endpoint on its reaper thread after the DEALER's close() has already + # returned, so reusing the peer-derived address on a reconnect races + # that teardown and fails with EADDRINUSE. safe_id = peer_id.replace(":", "-").replace("/", "-") - monitor_addr = f"inproc://p2p-monitor-{safe_id}" + monitor_addr = f"inproc://p2p-monitor-{safe_id}-{self._monitor_seq}" + self._monitor_seq += 1 dealer.monitor(monitor_addr, zmq.EVENT_DISCONNECTED) monitor_sock = self._zmq_ctx.socket(zmq.PAIR) @@ -231,6 +267,11 @@ class ZmqTransport(ControlTransport): ) return conn + def _sweep_dead_connections(self) -> None: + """Unregister and release every connection that is no longer alive.""" + for pid in [p for p, c in self._connections.items() if not c.alive]: + self._connections.pop(pid).close() + def _recv_router(self) -> None: """Non-blocking: receive all pending messages from ROUTER.""" while True: From 5f89a03dcb52702a62644e15b93f766765d06b28 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas <akaratza@amd.com> Date: Mon, 27 Jul 2026 00:32:46 -0500 Subject: [PATCH 088/185] [CI] Explicitly tear down speculative decode runners (#49910) Signed-off-by: Andreas Karatzas <akaratza@amd.com> --- tests/v1/spec_decode/test_max_len.py | 76 ++++++++++++++++------------ 1 file changed, 43 insertions(+), 33 deletions(-) diff --git a/tests/v1/spec_decode/test_max_len.py b/tests/v1/spec_decode/test_max_len.py index 77c041d84a9..81f84241937 100644 --- a/tests/v1/spec_decode/test_max_len.py +++ b/tests/v1/spec_decode/test_max_len.py @@ -5,7 +5,7 @@ import pytest from tests.utils import get_attn_backend_list_based_on_platform -from vllm import LLM, SamplingParams +from vllm import SamplingParams from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig from vllm.platforms import current_platform from vllm.sampling_params import StructuredOutputsParams @@ -18,10 +18,12 @@ _PROMPTS = [ @pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10]) -def test_ngram_max_len(num_speculative_tokens: int): - llm = LLM( - model="facebook/opt-125m", +def test_ngram_max_len(num_speculative_tokens: int, vllm_runner): + with vllm_runner( + "facebook/opt-125m", + trust_remote_code=False, max_model_len=100, + enable_chunked_prefill=None, enforce_eager=True, # For faster initialization. speculative_config={ "method": "ngram", @@ -29,21 +31,26 @@ def test_ngram_max_len(num_speculative_tokens: int): "prompt_lookup_min": 3, "num_speculative_tokens": num_speculative_tokens, }, - ) - sampling_params = SamplingParams(max_tokens=100, ignore_eos=True) - llm.generate(_PROMPTS, sampling_params) + ) as runner: + sampling_params = SamplingParams(max_tokens=100, ignore_eos=True) + runner.llm.generate(_PROMPTS, sampling_params) @pytest.mark.parametrize("num_speculative_tokens", [1, 3, 10]) @pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform()) def test_eagle_max_len( - monkeypatch: pytest.MonkeyPatch, num_speculative_tokens: int, attn_backend: str + monkeypatch: pytest.MonkeyPatch, + num_speculative_tokens: int, + attn_backend: str, + vllm_runner, ): if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm(): monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") - llm = LLM( - model="meta-llama/Meta-Llama-3-8B-Instruct", + with vllm_runner( + "meta-llama/Meta-Llama-3-8B-Instruct", + trust_remote_code=False, + enable_chunked_prefill=None, enforce_eager=True, # For faster initialization. speculative_config={ "method": "eagle", @@ -53,31 +60,34 @@ def test_eagle_max_len( }, max_model_len=200, attention_config={"backend": attn_backend}, - ) - sampling_params = SamplingParams(max_tokens=200, ignore_eos=True) - outputs = llm.generate(_PROMPTS, sampling_params) - for o in outputs: - assert o.outputs[0].finish_reason == "length", ( - "This test is only meaningful if the output is truncated due to max length" - ) + ) as runner: + sampling_params = SamplingParams(max_tokens=200, ignore_eos=True) + outputs = runner.llm.generate(_PROMPTS, sampling_params) + for o in outputs: + assert o.outputs[0].finish_reason == "length", ( + "This test is only meaningful if the output is truncated " + "due to max length" + ) - sampling_params = SamplingParams( - max_tokens=200, - structured_outputs=StructuredOutputsParams(regex="^" + "a b c d e " * 15 + "$"), - ) - output = llm.generate(_PROMPTS, sampling_params) - for o in output: - assert o.prompt_token_ids is not None - assert ( - len(o.prompt_token_ids) - < 80 - < len(o.prompt_token_ids) + len(o.outputs[0].token_ids) - <= 200 - ), ( - "This test is only meaningful if the output " - "is longer than the eagle max length" + sampling_params = SamplingParams( + max_tokens=200, + structured_outputs=StructuredOutputsParams( + regex="^" + "a b c d e " * 15 + "$" + ), ) - assert o.outputs[0].text == "a b c d e " * 15 + output = runner.llm.generate(_PROMPTS, sampling_params) + for o in output: + assert o.prompt_token_ids is not None + assert ( + len(o.prompt_token_ids) + < 80 + < len(o.prompt_token_ids) + len(o.outputs[0].token_ids) + <= 200 + ), ( + "This test is only meaningful if the output " + "is longer than the eagle max length" + ) + assert o.outputs[0].text == "a b c d e " * 15 @pytest.mark.parametrize("spec_max_model_len", [80, 150]) From f19ee27e39b20564109c124c40d8544cc1df24cc Mon Sep 17 00:00:00 2001 From: Akash kaothalkar <61960177+Akashcodes732@users.noreply.github.com> Date: Mon, 27 Jul 2026 11:17:29 +0530 Subject: [PATCH 089/185] [Hardware][Power] Add FAST_EXP for Power (#49571) Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com> Co-authored-by: Akash kaothalkar <akash.kaothalkar@ibm.com> --- benchmarks/kernels/cpu/benchmark_cpu_attn.py | 2 +- csrc/cpu/cpu_arch_macros.h | 11 +++++++++++ csrc/cpu/cpu_types_vsx.hpp | 15 ++++++++++++--- 3 files changed, 24 insertions(+), 4 deletions(-) diff --git a/benchmarks/kernels/cpu/benchmark_cpu_attn.py b/benchmarks/kernels/cpu/benchmark_cpu_attn.py index 08afd693c33..cdfcb1d404a 100644 --- a/benchmarks/kernels/cpu/benchmark_cpu_attn.py +++ b/benchmarks/kernels/cpu/benchmark_cpu_attn.py @@ -154,7 +154,7 @@ def main( scale=scale, causal=True, alibi_slopes=None, - sliding_window=window_size, + sliding_window=window_size if sliding_window is not None else -1, block_table=block_tables, softcap=0, scheduler_metadata=metadata, diff --git a/csrc/cpu/cpu_arch_macros.h b/csrc/cpu/cpu_arch_macros.h index 53ae70497c0..cf3dd4b9616 100644 --- a/csrc/cpu/cpu_arch_macros.h +++ b/csrc/cpu/cpu_arch_macros.h @@ -172,4 +172,15 @@ #endif // __riscv_v +// Power VSX +#ifdef __powerpc__ + // FP32Vec16::exp() in cpu_types_vsx.hpp delegates to FP32Vec8::exp(), which + // implements a vectorised 5-term minimax polynomial using VSX intrinsics. + #define DEFINE_FAST_EXP \ + auto fast_exp = [&](const vec_op::FP32Vec16& vec) \ + __attribute__((always_inline)) { return vec.exp(); }; \ + auto fast_exp_f16 = fast_exp; + +#endif // __powerpc__ + #endif diff --git a/csrc/cpu/cpu_types_vsx.hpp b/csrc/cpu/cpu_types_vsx.hpp index 64fe961da22..42083bc3eb6 100644 --- a/csrc/cpu/cpu_types_vsx.hpp +++ b/csrc/cpu/cpu_types_vsx.hpp @@ -287,7 +287,7 @@ struct FP32Vec4 : public Vec<FP32Vec4> { explicit FP32Vec4(__vector float data) : reg(data) {} - explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {} + FP32Vec4(const FP32Vec4& data) : reg(data.reg) {} }; struct FP32Vec8 : public Vec<FP32Vec8> { @@ -316,7 +316,7 @@ struct FP32Vec8 : public Vec<FP32Vec8> { explicit FP32Vec8(f32x4x2_t data) : reg(data) {} - explicit FP32Vec8(const FP32Vec8& data) { + FP32Vec8(const FP32Vec8& data) { reg.val[0] = data.reg.val[0]; reg.val[1] = data.reg.val[1]; } @@ -593,7 +593,7 @@ struct FP32Vec16 : public Vec<FP32Vec16> { explicit FP32Vec16(bool, const float* ptr) : FP32Vec16(ptr) {} explicit FP32Vec16(f32x4x4_t data) : reg(data) {} - explicit FP32Vec16(const FP32Vec16& data) { + FP32Vec16(const FP32Vec16& data) { reg.val[0] = data.reg.val[0]; reg.val[1] = data.reg.val[1]; reg.val[2] = data.reg.val[2]; @@ -747,6 +747,15 @@ struct FP32Vec16 : public Vec<FP32Vec16> { vec_abs(reg.val[2]), vec_abs(reg.val[3])})); } + FP32Vec16 exp() const { + FP32Vec8 lo(f32x4x2_t{reg.val[0], reg.val[1]}); + FP32Vec8 hi(f32x4x2_t{reg.val[2], reg.val[3]}); + auto lo_e = lo.exp(); + auto hi_e = hi.exp(); + return FP32Vec16(f32x4x4_t{lo_e.reg.val[0], lo_e.reg.val[1], + hi_e.reg.val[0], hi_e.reg.val[1]}); + } + float reduce_max() { __vector float max01 = vec_max(reg.val[0], reg.val[1]); __vector float max23 = vec_max(reg.val[2], reg.val[3]); From c314af1abfddff7b6cce9af578be72c496c6a5e4 Mon Sep 17 00:00:00 2001 From: Fadi Arafeh <115173828+fadara01@users.noreply.github.com> Date: Mon, 27 Jul 2026 06:53:09 +0100 Subject: [PATCH 090/185] [CPU][Perf] INT8 Fused MoE Kernel for Arm CPUs (#48637) Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com> --- cmake/cpu_extension.cmake | 20 +- csrc/cpu/cpu_fused_moe.cpp | 192 +----- csrc/cpu/cpu_fused_moe_activations.hpp | 204 ++++++ csrc/cpu/cpu_fused_moe_int8.cpp | 647 ++++++++++++++++++ csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp | 3 + .../micro_gemm/cpu_micro_gemm_int8_neon.hpp | 424 ++++++++++++ csrc/cpu/micro_gemm/cpu_micro_gemm_neon.hpp | 22 +- csrc/cpu/torch_bindings.cpp | 34 +- tests/kernels/moe/test_cpu_fused_moe.py | 172 ++++- vllm/_custom_ops.py | 42 ++ .../layers/fused_moe/experts/cpu_moe.py | 182 ++++- .../layers/fused_moe/oracle/int8.py | 16 +- 12 files changed, 1742 insertions(+), 216 deletions(-) create mode 100644 csrc/cpu/cpu_fused_moe_activations.hpp create mode 100644 csrc/cpu/cpu_fused_moe_int8.cpp create mode 100644 csrc/cpu/micro_gemm/cpu_micro_gemm_int8_neon.hpp diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 4e7837df0d7..64df94e947d 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -15,6 +15,7 @@ endif() # set(ENABLE_X86_ISA $ENV{VLLM_CPU_X86}) set(ENABLE_ARM_BF16 $ENV{VLLM_CPU_ARM_BF16}) +set(ENABLE_ARM_I8MM $ENV{VLLM_CPU_ARM_I8MM}) set(ENABLE_RVV_BF16 $ENV{VLLM_CPU_RVV_BF16}) include_directories("${CMAKE_SOURCE_DIR}/csrc") @@ -96,12 +97,14 @@ if (MACOSX_FOUND AND CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64") set(ENABLE_NUMA OFF) check_sysctl(hw.optional.neon ASIMD_FOUND) check_sysctl(hw.optional.arm.FEAT_BF16 ARM_BF16_FOUND) + check_sysctl(hw.optional.arm.FEAT_I8MM ARM_I8MM_FOUND) else() find_isa(${CPUINFO} "Power11" POWER11_FOUND) find_isa(${CPUINFO} "POWER10" POWER10_FOUND) find_isa(${CPUINFO} "POWER9" POWER9_FOUND) find_isa(${CPUINFO} "asimd" ASIMD_FOUND) # Check for ARM NEON support find_isa(${CPUINFO} "bf16" ARM_BF16_FOUND) # Check for ARM BF16 support + find_isa(${CPUINFO} "i8mm" ARM_I8MM_FOUND) # Check for ARM I8MM support find_isa(${CPUINFO} "S390" S390_FOUND) find_isa(${CPUINFO} "zvfhmin" RVV_FP16_FOUND) # Check for RISC-V Vector FP16 support find_isa(${CPUINFO} "zvfbfmin" RVV_BF16_FOUND) # Check for RISC-V Vector BF16 support @@ -111,6 +114,11 @@ else() set(ARM_BF16_FOUND ON) message(STATUS "ARM BF16 support enabled via VLLM_CPU_ARM_BF16 environment variable") endif() + if (ENABLE_ARM_I8MM) + set(ARM_I8MM_FOUND ON) + message(STATUS + "ARM I8MM support enabled via VLLM_CPU_ARM_I8MM environment variable") + endif() # Some kernels (e.g. Bianbu on Spacemit X100) do not report zvfbfmin # in /proc/cpuinfo despite hardware support. VLLM_CPU_RVV_BF16=1 # overrides the detection result. @@ -166,6 +174,11 @@ elseif (ASIMD_FOUND) message(WARNING "BF16 functionality is not available") set(MARCH_FLAGS "-march=armv8.2-a+dotprod+fp16") endif() + if(ARM_I8MM_FOUND) + message(STATUS "I8MM extension detected") + string(APPEND MARCH_FLAGS "+i8mm") + add_compile_definitions(ARM_I8MM_SUPPORT) + endif() list(APPEND CXX_COMPILE_FLAGS ${MARCH_FLAGS}) elseif (S390_FOUND) message(STATUS "S390 detected") @@ -447,8 +460,13 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) "csrc/cpu/shm.cpp" "csrc/cpu/activation_lut_bf16.cpp" "csrc/cpu/cpu_tanhf_neon.hpp" - "csrc/cpu/cpu_fused_moe.cpp" ${VLLM_EXT_SRC}) + if (ARM_BF16_FOUND) + set(VLLM_EXT_SRC "csrc/cpu/cpu_fused_moe.cpp" ${VLLM_EXT_SRC}) + if (ARM_I8MM_FOUND) + set(VLLM_EXT_SRC "csrc/cpu/cpu_fused_moe_int8.cpp" ${VLLM_EXT_SRC}) + endif() + endif() endif() if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND) diff --git a/csrc/cpu/cpu_fused_moe.cpp b/csrc/cpu/cpu_fused_moe.cpp index 07b0aaf8688..b7bbef4c2d7 100644 --- a/csrc/cpu/cpu_fused_moe.cpp +++ b/csrc/cpu/cpu_fused_moe.cpp @@ -1,5 +1,6 @@ #include "cpu/cpu_types.hpp" #include "cpu/utils.hpp" +#include "cpu/cpu_fused_moe_activations.hpp" #include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp" #include "cpu/cpu_arch_macros.h" @@ -43,193 +44,9 @@ }() namespace { -enum class FusedMOEAct { - SiluAndMul, - SwigluOAIAndMul, - GeluAndMul, - GeluTanhAndMul, -}; -FusedMOEAct get_act_type(const std::string& act) { - if (act == "silu") { - return FusedMOEAct::SiluAndMul; - } else if (act == "swigluoai") { - return FusedMOEAct::SwigluOAIAndMul; - } else if (act == "gelu") { - return FusedMOEAct::GeluAndMul; - } else if (act == "gelu_tanh") { - return FusedMOEAct::GeluTanhAndMul; - } else { - TORCH_CHECK(false, "Invalid act type: " + act); - } -} - -template <typename scalar_t> -void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, - const int32_t m_size, const int32_t n_size, - const int32_t input_stride, - const int32_t output_stride) { - using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t; -#if !defined(__aarch64__) - // For GPT-OSS interleaved gate-up weights - alignas(64) static int32_t index[16] = {0, 2, 4, 6, 8, 10, 12, 14, - 16, 18, 20, 22, 24, 26, 28, 30}; - vec_op::INT32Vec16 index_vec(index); -#endif - vec_op::FP32Vec16 gate_up_max_vec(7.0); - vec_op::FP32Vec16 up_min_vec(-7.0); - vec_op::FP32Vec16 alpha_vec(1.702); - vec_op::FP32Vec16 one_vec(1.0); - - DEFINE_FAST_EXP - - for (int32_t m = 0; m < m_size; ++m) { - for (int32_t n = 0; n < n_size; n += 32) { - // Note: AdvSIMD does not support gather loads -#if defined(__aarch64__) - vec_op::FP32Vec16 gate_vec(vec_op::uninit); - vec_op::FP32Vec16 up_vec(vec_op::uninit); - vec_op::FP32Vec16::load_even_odd(input + n, gate_vec, up_vec); -#else - vec_op::FP32Vec16 gate_vec(input + n, index_vec); - vec_op::FP32Vec16 up_vec(input + n + 1, index_vec); -#endif - gate_vec = gate_vec.min(gate_up_max_vec); - up_vec = up_vec.clamp(up_min_vec, gate_up_max_vec); - auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec * alpha_vec)); - auto glu = gate_vec * sigmoid_vec; - auto gated_output_fp32 = (one_vec + up_vec) * glu; - scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32); - gated_output.save(output + n / 2); - } - input += input_stride; - output += output_stride; - } -} - -template <typename scalar_t> -void silu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, - const int32_t m_size, const int32_t n_size, - const int32_t input_stride, const int32_t output_stride) { - using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t; - const int32_t dim = n_size / 2; - float* __restrict__ gate = input; - float* __restrict__ up = input + dim; - vec_op::FP32Vec16 one_vec(1.0); - - DEFINE_FAST_EXP - - for (int32_t m = 0; m < m_size; ++m) { - for (int32_t n = 0; n < dim; n += 16) { - vec_op::FP32Vec16 gate_vec(gate + n); - vec_op::FP32Vec16 up_vec(up + n); - auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec)); - auto silu = gate_vec * sigmoid_vec; - auto gated_output_fp32 = up_vec * silu; - scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32); - gated_output.save(output + n); - } - gate += input_stride; - up += input_stride; - output += output_stride; - } -} - -template <typename scalar_t> -void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, - const int32_t m_size, const int32_t n_size, - const int32_t input_stride, const int32_t output_stride) { - using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t; - const int32_t dim = n_size / 2; - float* __restrict__ gate = input; - float* __restrict__ up = input + dim; - vec_op::FP32Vec16 one_vec(1.0); - vec_op::FP32Vec16 w1_vec(M_SQRT1_2); - vec_op::FP32Vec16 w2_vec(0.5); - alignas(64) float temp[16]; - - DEFINE_FAST_EXP - - for (int32_t m = 0; m < m_size; ++m) { - for (int32_t n = 0; n < dim; n += 16) { - vec_op::FP32Vec16 gate_vec(gate + n); - vec_op::FP32Vec16 up_vec(up + n); - auto er_input_vec = gate_vec * w1_vec; - - er_input_vec.save(temp); - for (int32_t i = 0; i < 16; ++i) { - temp[i] = std::erf(temp[i]); - } - vec_op::FP32Vec16 er_vec(temp); - auto gelu = gate_vec * w2_vec * (one_vec + er_vec); - auto gated_output_fp32 = up_vec * gelu; - scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32); - gated_output.save(output + n); - } - gate += input_stride; - up += input_stride; - output += output_stride; - } -} - -template <typename scalar_t> -void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, - const int32_t m_size, const int32_t n_size, - const int32_t input_stride, - const int32_t output_stride) { - using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t; - const int32_t dim = n_size / 2; - float* __restrict__ gate = input; - float* __restrict__ up = input + dim; - vec_op::FP32Vec16 one_vec(1.0); - vec_op::FP32Vec16 w1_vec(0.7978845608028654); - vec_op::FP32Vec16 w2_vec(0.5); - vec_op::FP32Vec16 w3_vec(0.044715); - - for (int32_t m = 0; m < m_size; ++m) { - for (int32_t n = 0; n < dim; n += 16) { - vec_op::FP32Vec16 gate_vec(gate + n); - vec_op::FP32Vec16 up_vec(up + n); - auto gate_pow3_vec = gate_vec * gate_vec * gate_vec; - auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec); - // Note: can't use fast_exp form because diffusiongemma will generate - // wrong results - auto tanh_vec = inner_vec.tanh(); - auto gelu_tanh = gate_vec * w2_vec * (one_vec + tanh_vec); - auto gated_output_fp32 = up_vec * gelu_tanh; - scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32); - gated_output.save(output + n); - } - gate += input_stride; - up += input_stride; - output += output_stride; - } -} - -template <typename scalar_t> -FORCE_INLINE void apply_gated_act(const FusedMOEAct act, - float* __restrict__ input, - scalar_t* __restrict__ output, - const int32_t m, const int32_t n, - const int32_t input_stride, - const int32_t output_stride) { - switch (act) { - case FusedMOEAct::SwigluOAIAndMul: - swigluoai_and_mul(input, output, m, n, input_stride, output_stride); - return; - case FusedMOEAct::SiluAndMul: - silu_and_mul(input, output, m, n, input_stride, output_stride); - return; - case FusedMOEAct::GeluAndMul: - gelu_and_mul(input, output, m, n, input_stride, output_stride); - return; - case FusedMOEAct::GeluTanhAndMul: - gelu_tanh_and_mul(input, output, m, n, input_stride, output_stride); - return; - default: - TORCH_CHECK(false, "Unsupported act type."); - } -} +using cpu_fused_moe_utils::apply_gated_act; +using cpu_fused_moe_utils::FusedMOEAct; template <typename scalar_t, typename gemm_t> void prepack_moe_weight_impl(scalar_t* __restrict__ weight_ptr, @@ -817,6 +634,7 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input, } } } + } // namespace void prepack_moe_weight( @@ -864,7 +682,7 @@ void cpu_fused_moe( const int32_t input_size_2 = w2.size(2); const int32_t output_size_2 = w2.size(1); const int32_t topk_num = topk_id.size(1); - const FusedMOEAct act_type = get_act_type(act); + const FusedMOEAct act_type = cpu_fused_moe_utils::get_act_type(act); cpu_utils::ISA isa_type = cpu_utils::get_isa(isa); TORCH_CHECK(!skip_weighted || topk_num == 1, "skip_weighted is only supported for topk=1 on CPU"); diff --git a/csrc/cpu/cpu_fused_moe_activations.hpp b/csrc/cpu/cpu_fused_moe_activations.hpp new file mode 100644 index 00000000000..31e8cb009b7 --- /dev/null +++ b/csrc/cpu/cpu_fused_moe_activations.hpp @@ -0,0 +1,204 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#ifndef CPU_FUSED_MOE_ACTIVATIONS_HPP +#define CPU_FUSED_MOE_ACTIVATIONS_HPP + +#include <cmath> +#include <cstdint> +#include <string> + +#include "cpu/cpu_arch_macros.h" +#include "cpu/utils.hpp" + +namespace cpu_fused_moe_utils { +enum class FusedMOEAct { + SiluAndMul, + SwigluOAIAndMul, + GeluAndMul, + GeluTanhAndMul, +}; + +inline FusedMOEAct get_act_type(const std::string& act) { + if (act == "silu") { + return FusedMOEAct::SiluAndMul; + } else if (act == "swigluoai") { + return FusedMOEAct::SwigluOAIAndMul; + } else if (act == "gelu") { + return FusedMOEAct::GeluAndMul; + } else if (act == "gelu_tanh") { + return FusedMOEAct::GeluTanhAndMul; + } else { + TORCH_CHECK(false, "Invalid act type: " + act); + } +} + +template <typename scalar_t> +void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, + const int32_t m_size, const int32_t n_size, + const int32_t input_stride, + const int32_t output_stride) { + using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t; +#if !defined(__aarch64__) + // For GPT-OSS interleaved gate-up weights + alignas(64) static int32_t index[16] = {0, 2, 4, 6, 8, 10, 12, 14, + 16, 18, 20, 22, 24, 26, 28, 30}; + vec_op::INT32Vec16 index_vec(index); +#endif + vec_op::FP32Vec16 gate_up_max_vec(7.0); + vec_op::FP32Vec16 up_min_vec(-7.0); + vec_op::FP32Vec16 alpha_vec(1.702); + vec_op::FP32Vec16 one_vec(1.0); + + DEFINE_FAST_EXP + + for (int32_t m = 0; m < m_size; ++m) { + for (int32_t n = 0; n < n_size; n += 32) { + // Note: AdvSIMD does not support gather loads +#if defined(__aarch64__) + vec_op::FP32Vec16 gate_vec(vec_op::uninit); + vec_op::FP32Vec16 up_vec(vec_op::uninit); + vec_op::FP32Vec16::load_even_odd(input + n, gate_vec, up_vec); +#else + vec_op::FP32Vec16 gate_vec(input + n, index_vec); + vec_op::FP32Vec16 up_vec(input + n + 1, index_vec); +#endif + gate_vec = gate_vec.min(gate_up_max_vec); + up_vec = up_vec.clamp(up_min_vec, gate_up_max_vec); + auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec * alpha_vec)); + auto glu = gate_vec * sigmoid_vec; + auto gated_output_fp32 = (one_vec + up_vec) * glu; + scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32); + gated_output.save(output + n / 2); + } + input += input_stride; + output += output_stride; + } +} + +template <typename scalar_t> +void silu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, + const int32_t m_size, const int32_t n_size, + const int32_t input_stride, const int32_t output_stride) { + using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t; + const int32_t dim = n_size / 2; + float* __restrict__ gate = input; + float* __restrict__ up = input + dim; + vec_op::FP32Vec16 one_vec(1.0); + + DEFINE_FAST_EXP + + for (int32_t m = 0; m < m_size; ++m) { + for (int32_t n = 0; n < dim; n += 16) { + vec_op::FP32Vec16 gate_vec(gate + n); + vec_op::FP32Vec16 up_vec(up + n); + auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec)); + auto silu = gate_vec * sigmoid_vec; + auto gated_output_fp32 = up_vec * silu; + scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32); + gated_output.save(output + n); + } + gate += input_stride; + up += input_stride; + output += output_stride; + } +} + +template <typename scalar_t> +void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, + const int32_t m_size, const int32_t n_size, + const int32_t input_stride, const int32_t output_stride) { + using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t; + const int32_t dim = n_size / 2; + float* __restrict__ gate = input; + float* __restrict__ up = input + dim; + vec_op::FP32Vec16 one_vec(1.0); + vec_op::FP32Vec16 w1_vec(M_SQRT1_2); + vec_op::FP32Vec16 w2_vec(0.5); + alignas(64) float temp[16]; + + DEFINE_FAST_EXP + + for (int32_t m = 0; m < m_size; ++m) { + for (int32_t n = 0; n < dim; n += 16) { + vec_op::FP32Vec16 gate_vec(gate + n); + vec_op::FP32Vec16 up_vec(up + n); + auto er_input_vec = gate_vec * w1_vec; + + er_input_vec.save(temp); + for (int32_t i = 0; i < 16; ++i) { + temp[i] = std::erf(temp[i]); + } + vec_op::FP32Vec16 er_vec(temp); + auto gelu = gate_vec * w2_vec * (one_vec + er_vec); + auto gated_output_fp32 = up_vec * gelu; + scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32); + gated_output.save(output + n); + } + gate += input_stride; + up += input_stride; + output += output_stride; + } +} + +template <typename scalar_t> +void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, + const int32_t m_size, const int32_t n_size, + const int32_t input_stride, + const int32_t output_stride) { + using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t; + const int32_t dim = n_size / 2; + float* __restrict__ gate = input; + float* __restrict__ up = input + dim; + vec_op::FP32Vec16 one_vec(1.0); + vec_op::FP32Vec16 w1_vec(0.7978845608028654); + vec_op::FP32Vec16 w2_vec(0.5); + vec_op::FP32Vec16 w3_vec(0.044715); + + for (int32_t m = 0; m < m_size; ++m) { + for (int32_t n = 0; n < dim; n += 16) { + vec_op::FP32Vec16 gate_vec(gate + n); + vec_op::FP32Vec16 up_vec(up + n); + auto gate_pow3_vec = gate_vec * gate_vec * gate_vec; + auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec); + // Note: can't use fast_exp form because diffusiongemma will generate + // wrong results + auto tanh_vec = inner_vec.tanh(); + auto gelu_tanh = gate_vec * w2_vec * (one_vec + tanh_vec); + auto gated_output_fp32 = up_vec * gelu_tanh; + scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32); + gated_output.save(output + n); + } + gate += input_stride; + up += input_stride; + output += output_stride; + } +} + +template <typename scalar_t> +FORCE_INLINE void apply_gated_act(const FusedMOEAct act, + float* __restrict__ input, + scalar_t* __restrict__ output, + const int32_t m, const int32_t n, + const int32_t input_stride, + const int32_t output_stride) { + switch (act) { + case FusedMOEAct::SwigluOAIAndMul: + swigluoai_and_mul(input, output, m, n, input_stride, output_stride); + return; + case FusedMOEAct::SiluAndMul: + silu_and_mul(input, output, m, n, input_stride, output_stride); + return; + case FusedMOEAct::GeluAndMul: + gelu_and_mul(input, output, m, n, input_stride, output_stride); + return; + case FusedMOEAct::GeluTanhAndMul: + gelu_tanh_and_mul(input, output, m, n, input_stride, output_stride); + return; + default: + TORCH_CHECK(false, "Unsupported act type."); + } +} +} // namespace cpu_fused_moe_utils + +#endif diff --git a/csrc/cpu/cpu_fused_moe_int8.cpp b/csrc/cpu/cpu_fused_moe_int8.cpp new file mode 100644 index 00000000000..649317b4786 --- /dev/null +++ b/csrc/cpu/cpu_fused_moe_int8.cpp @@ -0,0 +1,647 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#include "cpu/cpu_arch_macros.h" + +#include <algorithm> +#include <cstdint> +#include <cstring> +#include <optional> +#include <string> + +#include "cpu/cpu_fused_moe_activations.hpp" +#include "cpu/cpu_types.hpp" +#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp" +#include "cpu/utils.hpp" + +#if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) + #include "cpu/micro_gemm/cpu_micro_gemm_int8_neon.hpp" + #define NEON_DISPATCH(SCALAR_TYPE, ...) \ + case cpu_utils::ISA::NEON: { \ + using gemm_t = \ + cpu_micro_gemm::MicroGemmINT8<cpu_utils::ISA::NEON, SCALAR_TYPE>; \ + return __VA_ARGS__(); \ + } +#else + #define NEON_DISPATCH(SCALAR_TYPE, ...) case cpu_utils::ISA::NEON: +#endif + +#define CPU_INT8_ISA_DISPATCH_IMPL(ISA_TYPE, SCALAR_TYPE, ...) \ + [&] { \ + switch (ISA_TYPE) { \ + NEON_DISPATCH(SCALAR_TYPE, __VA_ARGS__) \ + default: { \ + TORCH_CHECK(false, "Invalid CPU ISA type."); \ + } \ + } \ + }() + +namespace { +using cpu_fused_moe_utils::apply_gated_act; +using cpu_fused_moe_utils::FusedMOEAct; + +template <typename gemm_t> +void prepack_moe_weight_int8_impl(const int8_t* __restrict__ weight_ptr, + int8_t* __restrict__ packed_weight_ptr, + const int32_t expert_num, + const int32_t output_size, + const int32_t input_size, + const int64_t expert_stride) { +#pragma omp parallel for + for (int32_t e_idx = 0; e_idx < expert_num; ++e_idx) { + gemm_t::pack_weight(weight_ptr + expert_stride * e_idx, + packed_weight_ptr + expert_stride * e_idx, output_size, + input_size); + } +} + +// INT8 MoE kernel, based on the original BF16 kernel in cpu_fused_moe.cpp +template <typename scalar_t, typename gemm_t> +void fused_moe_int8_impl( + scalar_t* __restrict__ output, const scalar_t* __restrict__ input, + const int8_t* __restrict__ w13, const int8_t* __restrict__ w2, + const float* __restrict__ w13_scales, const float* __restrict__ w2_scales, + scalar_t* __restrict__ w13_bias, scalar_t* __restrict__ w2_bias, + const float* __restrict__ topk_weights, const int32_t* __restrict__ topk_id, + const FusedMOEAct act_type, const int32_t token_num, + const int32_t expert_num, const int32_t topk_num, + const int32_t input_size_13, const int32_t output_size_13, + const int32_t input_size_2, const int32_t output_size_2, + const bool skip_weighted) { + using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t; + constexpr int32_t gemm_n_tile_size = gemm_t::NSize; + constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize; + constexpr int32_t min_w13_n_tile_size = 2 * gemm_n_tile_size; + + TORCH_CHECK_EQ(input_size_13 % gemm_t::K, 0); + TORCH_CHECK_EQ(input_size_2 % gemm_t::K, 0); + TORCH_CHECK_EQ(output_size_13 % min_w13_n_tile_size, 0); + TORCH_CHECK_EQ(output_size_2 % gemm_n_tile_size, 0); + TORCH_CHECK_EQ(output_size_13 / 2, input_size_2); + + const int32_t thread_num = cpu_utils::get_max_threads(); + const int32_t w13_input_buffer_size = cpu_utils::round_up<64>( + gemm_m_tile_size * input_size_13 * sizeof(int8_t)); + const int32_t w2_input_buffer_size = + cpu_utils::round_up<64>(gemm_m_tile_size * input_size_2 * sizeof(int8_t)); + + const int32_t w13_n_tile_size = [&]() { + const int64_t cache_size = cpu_utils::get_available_l2_size(); + const int32_t n_size_cache_limit = + (cache_size - w13_input_buffer_size) / + (gemm_m_tile_size * sizeof(float) + input_size_13 * sizeof(int8_t)); + const int32_t n_size_thread_limit = + output_size_13 / std::max(1, thread_num / topk_num); + const int32_t n_size = cpu_utils::round_down<min_w13_n_tile_size>( + std::min(n_size_cache_limit, n_size_thread_limit)); + return std::max(n_size, min_w13_n_tile_size); + }(); + + const int32_t w2_n_tile_size = [&]() { + const int64_t cache_size = cpu_utils::get_available_l2_size(); + const int32_t n_size_cache_limit = + (cache_size - w2_input_buffer_size) / (input_size_2 * sizeof(int8_t)); + const int32_t n_size_thread_limit = + output_size_2 / std::max(1, thread_num / topk_num); + const int32_t n_size = cpu_utils::round_down<gemm_n_tile_size>( + std::min(n_size_cache_limit, n_size_thread_limit)); + return std::max(n_size, gemm_n_tile_size); + }(); + + int32_t common_buffer_offset = 0; + const int32_t token_num_per_group_buffer_offset = common_buffer_offset; + common_buffer_offset += cpu_utils::round_up<64>(expert_num * sizeof(int32_t)); + const int32_t cu_token_num_per_group_buffer_offset = common_buffer_offset; + common_buffer_offset += + cpu_utils::round_up<64>((expert_num + 1) * sizeof(int32_t)); + const int32_t expanded_token_num = token_num * topk_num; + const int32_t expand_token_id_buffer_offset = common_buffer_offset; + common_buffer_offset += + cpu_utils::round_up<64>(expanded_token_num * sizeof(int32_t)); + const int32_t expand_token_id_index_buffer_offset = common_buffer_offset; + common_buffer_offset += + cpu_utils::round_up<64>(expanded_token_num * sizeof(int32_t)); + const int32_t input_quant_buffer_offset = common_buffer_offset; + common_buffer_offset += + cpu_utils::round_up<64>(token_num * input_size_13 * sizeof(int8_t)); + const int32_t input_scale_buffer_offset = common_buffer_offset; + common_buffer_offset += cpu_utils::round_up<64>(token_num * sizeof(float)); + const int32_t w13_gemm_output_buffer_offset = common_buffer_offset; + common_buffer_offset += cpu_utils::round_up<64>( + expanded_token_num * input_size_2 * sizeof(scalar_t)); + const int32_t w13_output_scale_buffer_offset = common_buffer_offset; + common_buffer_offset += + cpu_utils::round_up<64>(expanded_token_num * sizeof(float)); + const int32_t w2_gemm_output_buffer_offset = common_buffer_offset; + common_buffer_offset += cpu_utils::round_up<64>( + expanded_token_num * output_size_2 * sizeof(float)); + + int32_t gemm_thread_buffer_offset = 0; + const int32_t gemm_input_buffer_offset = gemm_thread_buffer_offset; + gemm_thread_buffer_offset += + std::max(w13_input_buffer_size, w2_input_buffer_size); + const int32_t gemm_output_buffer_offset = gemm_thread_buffer_offset; + gemm_thread_buffer_offset += cpu_utils::round_up<64>( + gemm_m_tile_size * std::max(w13_n_tile_size, w2_n_tile_size) * + sizeof(int32_t)); + + const int32_t ws_output_buffer_offset = 0; + const int32_t ws_thread_buffer_size = + cpu_utils::round_up<64>(output_size_2 * sizeof(float)); + const int32_t thread_buffer_size = + std::max(gemm_thread_buffer_offset, ws_thread_buffer_size); + const int32_t buffer_size = + common_buffer_offset + thread_buffer_size * thread_num; + cpu_utils::ScratchPadManager::get_scratchpad_manager()->realloc(buffer_size); + uint8_t* common_buffer_start = + cpu_utils::ScratchPadManager::get_scratchpad_manager() + ->get_data<uint8_t>(); + uint8_t* thread_buffer_start = common_buffer_start + common_buffer_offset; + + int32_t* __restrict__ token_num_per_group_buffer = reinterpret_cast<int32_t*>( + common_buffer_start + token_num_per_group_buffer_offset); + int32_t* __restrict__ cu_token_num_per_group_buffer = + reinterpret_cast<int32_t*>(common_buffer_start + + cu_token_num_per_group_buffer_offset); + int32_t* __restrict__ expand_token_id_buffer = reinterpret_cast<int32_t*>( + common_buffer_start + expand_token_id_buffer_offset); + int32_t* __restrict__ expand_token_id_index_buffer = + reinterpret_cast<int32_t*>(common_buffer_start + + expand_token_id_index_buffer_offset); + int8_t* __restrict__ input_quant_buffer = reinterpret_cast<int8_t*>( + common_buffer_start + input_quant_buffer_offset); + float* __restrict__ input_scale_buffer = + reinterpret_cast<float*>(common_buffer_start + input_scale_buffer_offset); + + std::memset(token_num_per_group_buffer, 0, expert_num * sizeof(int32_t)); + for (int32_t i = 0; i < expanded_token_num; ++i) { + ++token_num_per_group_buffer[topk_id[i]]; + } + + int32_t token_num_sum = 0; + cu_token_num_per_group_buffer[0] = 0; + int32_t* token_index_buffer = cu_token_num_per_group_buffer + 1; + for (int32_t i = 0; i < expert_num; ++i) { + token_index_buffer[i] = token_num_sum; + token_num_sum += token_num_per_group_buffer[i]; + } + + for (int32_t i = 0; i < token_num; ++i) { + const int32_t* curr_topk_id = topk_id + i * topk_num; + int32_t* curr_index_buffer = expand_token_id_index_buffer + i * topk_num; + for (int32_t j = 0; j < topk_num; ++j) { + const int32_t curr_expert_id = curr_topk_id[j]; + const int32_t curr_index = token_index_buffer[curr_expert_id]++; + expand_token_id_buffer[curr_index] = i; + curr_index_buffer[j] = curr_index; + } + } + +// quantize inputs +#pragma omp parallel for + for (int32_t token_idx = 0; token_idx < token_num; ++token_idx) { + gemm_t::quantize_row(input + token_idx * input_size_13, + input_quant_buffer + token_idx * input_size_13, + input_scale_buffer[token_idx], input_size_13); + } + + { + alignas(64) cpu_utils::Counter counter; + cpu_utils::Counter* counter_ptr = &counter; + +// w13 GEMM + act +#pragma omp parallel for schedule(static, 1) + for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) { + const int32_t task_num_per_expert = + (output_size_13 + w13_n_tile_size - 1) / w13_n_tile_size; + const int32_t task_num = task_num_per_expert * expert_num; + uint8_t* __restrict__ thread_buffer = + thread_buffer_start + thread_id * thread_buffer_size; + int8_t* __restrict__ gemm_input_buffer = + reinterpret_cast<int8_t*>(thread_buffer + gemm_input_buffer_offset); + float* __restrict__ gemm_output_buffer = + reinterpret_cast<float*>(thread_buffer + gemm_output_buffer_offset); + auto* __restrict__ w13_gemm_output_buffer = reinterpret_cast<scalar_t*>( + common_buffer_start + w13_gemm_output_buffer_offset); + gemm_t gemm; + + const int32_t w13_n_group_stride = + gemm_t::WeightOCGroupSize * input_size_13; + const int32_t w13_n_tile_stride = gemm_n_tile_size * input_size_13; + + for (;;) { + const int32_t task_id = counter_ptr->acquire_counter(); + if (task_id >= task_num) { + break; + } + const int32_t curr_expert_id = task_id / task_num_per_expert; + const int32_t curr_output_group_id = task_id % task_num_per_expert; + const int32_t curr_token_num = + token_num_per_group_buffer[curr_expert_id]; + if (curr_token_num == 0) { + continue; + } + + const int32_t actual_n_tile_size = + std::min(w13_n_tile_size, + output_size_13 - curr_output_group_id * w13_n_tile_size); + const int32_t* __restrict__ curr_expand_token_id_buffer = + expand_token_id_buffer + + cu_token_num_per_group_buffer[curr_expert_id]; + scalar_t* __restrict__ curr_w13_gemm_output_buffer = + w13_gemm_output_buffer + + cu_token_num_per_group_buffer[curr_expert_id] * input_size_2 + + curr_output_group_id * w13_n_tile_size / 2; + + const int8_t* w13_weight_ptr_0 = nullptr; + const int8_t* w13_weight_ptr_1 = nullptr; + const float* w13_scale_ptr_0 = nullptr; + const float* w13_scale_ptr_1 = nullptr; + scalar_t* w13_bias_ptr_0 = nullptr; + scalar_t* w13_bias_ptr_1 = nullptr; + if (act_type == FusedMOEAct::SwigluOAIAndMul) { + const int32_t output_offset = curr_output_group_id * w13_n_tile_size; + w13_weight_ptr_0 = w13 + + curr_expert_id * input_size_13 * output_size_13 + + output_offset * input_size_13; + w13_weight_ptr_1 = + w13_weight_ptr_0 + actual_n_tile_size / 2 * input_size_13; + w13_scale_ptr_0 = + w13_scales + curr_expert_id * output_size_13 + output_offset; + w13_scale_ptr_1 = w13_scale_ptr_0 + actual_n_tile_size / 2; + if (w13_bias != nullptr) { + w13_bias_ptr_0 = + w13_bias + curr_expert_id * output_size_13 + output_offset; + w13_bias_ptr_1 = w13_bias_ptr_0 + actual_n_tile_size / 2; + } + } else { + const int32_t output_offset = + curr_output_group_id * (w13_n_tile_size / 2); + w13_weight_ptr_0 = w13 + + curr_expert_id * input_size_13 * output_size_13 + + output_offset * input_size_13; + w13_weight_ptr_1 = + w13_weight_ptr_0 + output_size_13 / 2 * input_size_13; + w13_scale_ptr_0 = + w13_scales + curr_expert_id * output_size_13 + output_offset; + w13_scale_ptr_1 = w13_scale_ptr_0 + output_size_13 / 2; + if (w13_bias != nullptr) { + w13_bias_ptr_0 = + w13_bias + curr_expert_id * output_size_13 + output_offset; + w13_bias_ptr_1 = w13_bias_ptr_0 + output_size_13 / 2; + } + } + + for (int32_t token_idx = 0; token_idx < curr_token_num; + token_idx += gemm_m_tile_size) { + const int32_t actual_token_num = + std::min(gemm_m_tile_size, curr_token_num - token_idx); + const int8_t* input_rows[gemm_m_tile_size]; + alignas(64) float input_scales[gemm_m_tile_size]; + // gather and pack + for (int32_t i = 0; i < actual_token_num; ++i) { + const int32_t curr_token_id = curr_expand_token_id_buffer[i]; + input_rows[i] = input_quant_buffer + curr_token_id * input_size_13; + input_scales[i] = input_scale_buffer[curr_token_id]; + } + gemm_t::pack_input_from_rows(input_rows, gemm_input_buffer, + actual_token_num, input_size_13); + curr_expand_token_id_buffer += actual_token_num; + + const int8_t* w13_weight_ptr_0_iter = w13_weight_ptr_0; + const int8_t* w13_weight_ptr_1_iter = w13_weight_ptr_1; + const float* w13_scale_ptr_0_iter = w13_scale_ptr_0; + const float* w13_scale_ptr_1_iter = w13_scale_ptr_1; + scalar_t* w13_bias_ptr_0_iter = w13_bias_ptr_0; + scalar_t* w13_bias_ptr_1_iter = w13_bias_ptr_1; + float* w13_output_buffer_0_iter = gemm_output_buffer; + float* w13_output_buffer_1_iter = + gemm_output_buffer + actual_n_tile_size / 2; + + for (int32_t i = 0; i < actual_n_tile_size; + i += min_w13_n_tile_size) { + auto* output_0_int32 = + reinterpret_cast<int32_t*>(w13_output_buffer_0_iter); + gemm.gemm(gemm_input_buffer, w13_weight_ptr_0_iter, output_0_int32, + actual_token_num, input_size_13, w13_n_group_stride, + actual_n_tile_size); + gemm_t::dequantize_tile(output_0_int32, w13_output_buffer_0_iter, + input_scales, w13_scale_ptr_0_iter, + actual_token_num, gemm_n_tile_size, + actual_n_tile_size); + if (w13_bias != nullptr) { + cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>( + w13_output_buffer_0_iter, w13_output_buffer_0_iter, + w13_bias_ptr_0_iter, actual_token_num, actual_n_tile_size, + actual_n_tile_size); + w13_bias_ptr_0_iter += gemm_n_tile_size; + } + + auto* output_1_int32 = + reinterpret_cast<int32_t*>(w13_output_buffer_1_iter); + gemm.gemm(gemm_input_buffer, w13_weight_ptr_1_iter, output_1_int32, + actual_token_num, input_size_13, w13_n_group_stride, + actual_n_tile_size); + gemm_t::dequantize_tile(output_1_int32, w13_output_buffer_1_iter, + input_scales, w13_scale_ptr_1_iter, + actual_token_num, gemm_n_tile_size, + actual_n_tile_size); + if (w13_bias != nullptr) { + cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>( + w13_output_buffer_1_iter, w13_output_buffer_1_iter, + w13_bias_ptr_1_iter, actual_token_num, actual_n_tile_size, + actual_n_tile_size); + w13_bias_ptr_1_iter += gemm_n_tile_size; + } + + w13_weight_ptr_0_iter += w13_n_tile_stride; + w13_weight_ptr_1_iter += w13_n_tile_stride; + w13_scale_ptr_0_iter += gemm_n_tile_size; + w13_scale_ptr_1_iter += gemm_n_tile_size; + w13_output_buffer_0_iter += gemm_n_tile_size; + w13_output_buffer_1_iter += gemm_n_tile_size; + } + + apply_gated_act(act_type, gemm_output_buffer, + curr_w13_gemm_output_buffer, actual_token_num, + actual_n_tile_size, actual_n_tile_size, input_size_2); + curr_w13_gemm_output_buffer += gemm_m_tile_size * input_size_2; + } + } + } + } + + auto* __restrict__ w13_gemm_output_buffer = reinterpret_cast<scalar_t*>( + common_buffer_start + w13_gemm_output_buffer_offset); + float* __restrict__ w13_output_scale_buffer = reinterpret_cast<float*>( + common_buffer_start + w13_output_scale_buffer_offset); + +// quantize w2 inputs - in place +#pragma omp parallel for + for (int32_t token_idx = 0; token_idx < expanded_token_num; ++token_idx) { + scalar_t* input_row = w13_gemm_output_buffer + token_idx * input_size_2; + int8_t* output_row = reinterpret_cast<int8_t*>(input_row); + gemm_t::quantize_row(input_row, output_row, + w13_output_scale_buffer[token_idx], input_size_2); + } + + { + alignas(64) cpu_utils::Counter counter; + cpu_utils::Counter* counter_ptr = &counter; + +// w2 gemm +#pragma omp parallel for schedule(static, 1) + for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) { + const int32_t task_num_per_expert = + (output_size_2 + w2_n_tile_size - 1) / w2_n_tile_size; + const int32_t task_num = task_num_per_expert * expert_num; + uint8_t* __restrict__ thread_buffer = + thread_buffer_start + thread_id * thread_buffer_size; + int8_t* __restrict__ gemm_input_buffer = + reinterpret_cast<int8_t*>(thread_buffer + gemm_input_buffer_offset); + float* __restrict__ gemm_output_buffer = + reinterpret_cast<float*>(thread_buffer + gemm_output_buffer_offset); + float* __restrict__ w2_gemm_output_buffer = reinterpret_cast<float*>( + common_buffer_start + w2_gemm_output_buffer_offset); + gemm_t gemm; + + const int32_t w2_n_group_stride = + gemm_t::WeightOCGroupSize * input_size_2; + const int32_t w2_n_tile_stride = gemm_n_tile_size * input_size_2; + + for (;;) { + const int32_t task_id = counter_ptr->acquire_counter(); + if (task_id >= task_num) { + break; + } + const int32_t curr_expert_id = task_id / task_num_per_expert; + const int32_t curr_output_group_id = task_id % task_num_per_expert; + const int32_t curr_token_num = + token_num_per_group_buffer[curr_expert_id]; + if (curr_token_num == 0) { + continue; + } + + const int32_t actual_n_tile_size = + std::min(w2_n_tile_size, + output_size_2 - curr_output_group_id * w2_n_tile_size); + scalar_t* __restrict__ curr_w13_gemm_output_buffer = + w13_gemm_output_buffer + + cu_token_num_per_group_buffer[curr_expert_id] * input_size_2; + float* __restrict__ curr_w13_output_scale_buffer = + w13_output_scale_buffer + + cu_token_num_per_group_buffer[curr_expert_id]; + float* __restrict__ curr_w2_gemm_output_buffer = + w2_gemm_output_buffer + + cu_token_num_per_group_buffer[curr_expert_id] * output_size_2 + + curr_output_group_id * w2_n_tile_size; + const int8_t* __restrict__ w2_weight_ptr = + w2 + curr_expert_id * output_size_2 * input_size_2 + + curr_output_group_id * w2_n_tile_size * input_size_2; + const float* __restrict__ w2_scale_ptr = + w2_scales + curr_expert_id * output_size_2 + + curr_output_group_id * w2_n_tile_size; + scalar_t* w2_bias_ptr = nullptr; + if (w2_bias != nullptr) { + w2_bias_ptr = w2_bias + curr_expert_id * output_size_2 + + curr_output_group_id * w2_n_tile_size; + } + + for (int32_t token_idx = 0; token_idx < curr_token_num; + token_idx += gemm_m_tile_size) { + const int32_t actual_token_num = + std::min(gemm_m_tile_size, curr_token_num - token_idx); + const int8_t* input_rows[gemm_m_tile_size]; + alignas(64) float input_scales[gemm_m_tile_size]; + for (int32_t i = 0; i < actual_token_num; ++i) { + input_rows[i] = reinterpret_cast<const int8_t*>( + curr_w13_gemm_output_buffer + i * input_size_2); + input_scales[i] = curr_w13_output_scale_buffer[i]; + } + gemm_t::pack_input_from_rows(input_rows, gemm_input_buffer, + actual_token_num, input_size_2); + + const int8_t* w2_weight_ptr_iter = w2_weight_ptr; + const float* w2_scale_ptr_iter = w2_scale_ptr; + scalar_t* w2_bias_ptr_iter = w2_bias_ptr; + float* curr_w2_gemm_output_buffer_iter = curr_w2_gemm_output_buffer; + for (int32_t i = 0; i < actual_n_tile_size; i += gemm_n_tile_size) { + auto* output_int32 = reinterpret_cast<int32_t*>(gemm_output_buffer); + gemm.gemm(gemm_input_buffer, w2_weight_ptr_iter, output_int32, + actual_token_num, input_size_2, w2_n_group_stride, + gemm_n_tile_size); + gemm_t::dequantize_tile(output_int32, gemm_output_buffer, + input_scales, w2_scale_ptr_iter, + actual_token_num, gemm_n_tile_size, + gemm_n_tile_size); + if (w2_bias != nullptr) { + cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>( + gemm_output_buffer, gemm_output_buffer, w2_bias_ptr_iter, + actual_token_num, gemm_n_tile_size, gemm_n_tile_size); + w2_bias_ptr_iter += gemm_n_tile_size; + } + for (int32_t m_idx = 0; m_idx < actual_token_num; ++m_idx) { + std::memcpy( + curr_w2_gemm_output_buffer_iter + m_idx * output_size_2, + gemm_output_buffer + m_idx * gemm_n_tile_size, + gemm_n_tile_size * sizeof(float)); + } + + w2_weight_ptr_iter += w2_n_tile_stride; + w2_scale_ptr_iter += gemm_n_tile_size; + curr_w2_gemm_output_buffer_iter += gemm_n_tile_size; + } + + curr_w13_gemm_output_buffer += gemm_m_tile_size * input_size_2; + curr_w13_output_scale_buffer += gemm_m_tile_size; + curr_w2_gemm_output_buffer += gemm_m_tile_size * output_size_2; + } + } + } + } + + { + alignas(64) cpu_utils::Counter counter; + cpu_utils::Counter* counter_ptr = &counter; + +#pragma omp parallel for schedule(static, 1) + for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) { + uint8_t* __restrict__ thread_buffer = + thread_buffer_start + thread_id * thread_buffer_size; + float* __restrict__ ws_output_buffer = + reinterpret_cast<float*>(thread_buffer + ws_output_buffer_offset); + float* __restrict__ w2_gemm_output_buffer = reinterpret_cast<float*>( + common_buffer_start + w2_gemm_output_buffer_offset); + + for (;;) { + const int32_t token_id = counter_ptr->acquire_counter(); + if (token_id >= token_num) { + break; + } + int32_t* __restrict__ curr_expand_token_id_index_buffer = + expand_token_id_index_buffer + token_id * topk_num; + const float* __restrict__ curr_weight = + topk_weights + token_id * topk_num; + const float first_weight = skip_weighted ? 1.0f : curr_weight[0]; + scalar_t* __restrict__ curr_output_buffer = + output + token_id * output_size_2; + + if (topk_num > 1) { + int32_t w2_output_idx = curr_expand_token_id_index_buffer[0]; + float* w2_output_iter = + w2_gemm_output_buffer + w2_output_idx * output_size_2; + float* ws_output_buffer_iter = ws_output_buffer; + vec_op::FP32Vec16 weight_vec(first_weight); + for (int32_t i = 0; i < output_size_2; i += 16) { + vec_op::FP32Vec16 vec(w2_output_iter); + (vec * weight_vec).save(ws_output_buffer_iter); + w2_output_iter += 16; + ws_output_buffer_iter += 16; + } + + for (int32_t idx = 1; idx < topk_num - 1; ++idx) { + w2_output_idx = curr_expand_token_id_index_buffer[idx]; + w2_output_iter = + w2_gemm_output_buffer + w2_output_idx * output_size_2; + ws_output_buffer_iter = ws_output_buffer; + weight_vec = vec_op::FP32Vec16(curr_weight[idx]); + for (int32_t i = 0; i < output_size_2; i += 16) { + vec_op::FP32Vec16 vec(w2_output_iter); + vec_op::FP32Vec16 sum(ws_output_buffer_iter); + (sum + vec * weight_vec).save(ws_output_buffer_iter); + w2_output_iter += 16; + ws_output_buffer_iter += 16; + } + } + + const int32_t last_idx = topk_num - 1; + w2_output_idx = curr_expand_token_id_index_buffer[last_idx]; + w2_output_iter = + w2_gemm_output_buffer + w2_output_idx * output_size_2; + ws_output_buffer_iter = ws_output_buffer; + scalar_t* curr_output_buffer_iter = curr_output_buffer; + weight_vec = vec_op::FP32Vec16(curr_weight[last_idx]); + for (int32_t i = 0; i < output_size_2; i += 16) { + vec_op::FP32Vec16 vec(w2_output_iter); + vec_op::FP32Vec16 sum(ws_output_buffer_iter); + scalar_vec_t(sum + vec * weight_vec).save(curr_output_buffer_iter); + w2_output_iter += 16; + ws_output_buffer_iter += 16; + curr_output_buffer_iter += 16; + } + } else { + const int32_t w2_output_idx = curr_expand_token_id_index_buffer[0]; + float* w2_output_iter = + w2_gemm_output_buffer + w2_output_idx * output_size_2; + scalar_t* curr_output_buffer_iter = curr_output_buffer; + vec_op::FP32Vec16 weight_vec(first_weight); + for (int32_t i = 0; i < output_size_2; i += 16) { + vec_op::FP32Vec16 vec(w2_output_iter); + scalar_vec_t(vec * weight_vec).save(curr_output_buffer_iter); + w2_output_iter += 16; + curr_output_buffer_iter += 16; + } + } + } + } + } +} +} // namespace + +void prepack_moe_weight_int8( + const torch::Tensor& weight, // [expert_num, output_size, input_size] + torch::Tensor& packed_weight, const std::string& isa) { + TORCH_CHECK(weight.is_contiguous()); + const int32_t expert_num = weight.size(0); + const int32_t output_size = weight.size(1); + const int32_t input_size = weight.size(2); + const int64_t expert_stride = weight.stride(0); + const cpu_utils::ISA isa_type = cpu_utils::get_isa(isa); + TORCH_CHECK_EQ(output_size % 32, 0); + + CPU_INT8_ISA_DISPATCH_IMPL(isa_type, c10::BFloat16, [&]() { + TORCH_CHECK_EQ(input_size % gemm_t::K, 0); + prepack_moe_weight_int8_impl<gemm_t>( + weight.data_ptr<int8_t>(), packed_weight.data_ptr<int8_t>(), expert_num, + output_size, input_size, expert_stride); + }); +} + +void cpu_fused_moe_int8(torch::Tensor& output, const torch::Tensor& input, + const torch::Tensor& w13, const torch::Tensor& w2, + const torch::Tensor& w13_scale, + const torch::Tensor& w2_scale, + const std::optional<torch::Tensor>& w13_bias, + const std::optional<torch::Tensor>& w2_bias, + const torch::Tensor& topk_weights, + const torch::Tensor& topk_id, const bool skip_weighted, + const std::string& act, const std::string& isa) { + const int32_t token_num = input.size(0); + const int32_t input_size_13 = input.size(1); + const int64_t input_stride = input.stride(0); + TORCH_CHECK_EQ(input_stride, input_size_13); + const int32_t expert_num = w13.size(0); + const int32_t output_size_13 = w13.size(1); + const int32_t input_size_2 = w2.size(2); + const int32_t output_size_2 = w2.size(1); + const int32_t topk_num = topk_id.size(1); + const FusedMOEAct act_type = cpu_fused_moe_utils::get_act_type(act); + const cpu_utils::ISA isa_type = cpu_utils::get_isa(isa); + TORCH_CHECK(!skip_weighted || topk_num == 1, + "skip_weighted is only supported for topk=1 on CPU"); + + VLLM_DISPATCH_FLOATING_TYPES( + input.scalar_type(), "cpu_fused_moe_int8", [&]() { + CPU_INT8_ISA_DISPATCH_IMPL(isa_type, scalar_t, [&]() { + fused_moe_int8_impl<scalar_t, gemm_t>( + output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), + w13.data_ptr<int8_t>(), w2.data_ptr<int8_t>(), + w13_scale.data_ptr<float>(), w2_scale.data_ptr<float>(), + w13_bias.has_value() ? w13_bias->data_ptr<scalar_t>() : nullptr, + w2_bias.has_value() ? w2_bias->data_ptr<scalar_t>() : nullptr, + topk_weights.data_ptr<float>(), topk_id.data_ptr<int32_t>(), + act_type, token_num, expert_num, topk_num, input_size_13, + output_size_13, input_size_2, output_size_2, skip_weighted); + }); + }); +} diff --git a/csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp b/csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp index f0471f71470..75505176eee 100644 --- a/csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp +++ b/csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp @@ -31,6 +31,9 @@ class MicroGemm { } }; +template <cpu_utils::ISA isa, typename scalar_t> +class MicroGemmINT8; + template <int32_t n_size, typename scalar_t> FORCE_INLINE void default_epilogue(float* __restrict__ c_ptr, scalar_t* __restrict__ d_ptr, diff --git a/csrc/cpu/micro_gemm/cpu_micro_gemm_int8_neon.hpp b/csrc/cpu/micro_gemm/cpu_micro_gemm_int8_neon.hpp new file mode 100644 index 00000000000..2e194fd3903 --- /dev/null +++ b/csrc/cpu/micro_gemm/cpu_micro_gemm_int8_neon.hpp @@ -0,0 +1,424 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#ifndef CPU_MICRO_GEMM_INT8_NEON_HPP +#define CPU_MICRO_GEMM_INT8_NEON_HPP + +#include <algorithm> +#include <cstdint> + +#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp" + +#include <arm_bf16.h> +#include <arm_neon.h> +#include <c10/util/BFloat16.h> +#include <c10/util/Exception.h> +#include <c10/util/Half.h> + +namespace cpu_micro_gemm { + +namespace neon_smmla { + +constexpr int32_t K = 8; +constexpr int32_t Cols = 2; +constexpr int32_t TileSize = K * Cols; + +FORCE_INLINE float32x4x2_t load_as_f32(const float* input) { + float32x4x2_t result; + result.val[0] = vld1q_f32(input); + result.val[1] = vld1q_f32(input + 4); + return result; +} + +FORCE_INLINE float32x4x2_t load_as_f32(const c10::Half* input) { + const auto input_vec = vld1q_f16(reinterpret_cast<const float16_t*>(input)); + float32x4x2_t result; + result.val[0] = vcvt_f32_f16(vget_low_f16(input_vec)); + result.val[1] = vcvt_f32_f16(vget_high_f16(input_vec)); + return result; +} + +FORCE_INLINE float32x4x2_t load_as_f32(const c10::BFloat16* input) { + const auto input_vec = vld1q_bf16(reinterpret_cast<const bfloat16_t*>(input)); + float32x4x2_t result; + result.val[0] = vcvt_f32_bf16(vget_low_bf16(input_vec)); + result.val[1] = vcvt_f32_bf16(vget_high_bf16(input_vec)); + return result; +} + +FORCE_INLINE void store_acc_rowpair(const int32x4_t acc01, + const int32x4_t acc23, + const int32x4_t acc45, + const int32x4_t acc67, + int32_t* __restrict__ c_ptr, + const int64_t ldc, const int32_t m_rows) { + if (m_rows == 0) { + return; + } + + vst1q_s32(c_ptr, vcombine_s32(vget_low_s32(acc01), vget_low_s32(acc23))); + vst1q_s32(c_ptr + 4, vcombine_s32(vget_low_s32(acc45), vget_low_s32(acc67))); + + if (m_rows == 2) { + vst1q_s32(c_ptr + ldc, + vcombine_s32(vget_high_s32(acc01), vget_high_s32(acc23))); + vst1q_s32(c_ptr + ldc + 4, + vcombine_s32(vget_high_s32(acc45), vget_high_s32(acc67))); + } +} + +FORCE_INLINE void gemm_micro_smmla_8x8_packed_a( + const int8_t* __restrict__ a_packed, const int8_t* __restrict__ b_packed, + int32_t* __restrict__ c_ptr, const int32_t m, const int32_t k_size, + const int64_t ldc) { + const int32x4_t zero = vdupq_n_s32(0); + int32x4_t acc0101 = zero, acc0123 = zero, acc0145 = zero, acc0167 = zero; + int32x4_t acc2301 = zero, acc2323 = zero, acc2345 = zero, acc2367 = zero; + int32x4_t acc4501 = zero, acc4523 = zero, acc4545 = zero, acc4567 = zero; + int32x4_t acc6701 = zero, acc6723 = zero, acc6745 = zero, acc6767 = zero; + + const int8_t* __restrict__ a_tile = a_packed; + const int8_t* __restrict__ b_tile = b_packed; + +#pragma GCC unroll 8 + for (int32_t k_idx = 0; k_idx < k_size; k_idx += K) { + const int8x16_t a_tile01 = vld1q_s8(a_tile); + const int8x16_t a_tile23 = vld1q_s8(a_tile + TileSize); + const int8x16_t a_tile45 = vld1q_s8(a_tile + 2 * TileSize); + const int8x16_t a_tile67 = vld1q_s8(a_tile + 3 * TileSize); + + const int8x16_t b_tile01 = vld1q_s8(b_tile); + const int8x16_t b_tile23 = vld1q_s8(b_tile + TileSize); + const int8x16_t b_tile45 = vld1q_s8(b_tile + 2 * TileSize); + const int8x16_t b_tile67 = vld1q_s8(b_tile + 3 * TileSize); + + acc0101 = vmmlaq_s32(acc0101, a_tile01, b_tile01); + acc2301 = vmmlaq_s32(acc2301, a_tile23, b_tile01); + acc4501 = vmmlaq_s32(acc4501, a_tile45, b_tile01); + acc6701 = vmmlaq_s32(acc6701, a_tile67, b_tile01); + + acc0123 = vmmlaq_s32(acc0123, a_tile01, b_tile23); + acc2323 = vmmlaq_s32(acc2323, a_tile23, b_tile23); + acc4523 = vmmlaq_s32(acc4523, a_tile45, b_tile23); + acc6723 = vmmlaq_s32(acc6723, a_tile67, b_tile23); + + acc0145 = vmmlaq_s32(acc0145, a_tile01, b_tile45); + acc2345 = vmmlaq_s32(acc2345, a_tile23, b_tile45); + acc4545 = vmmlaq_s32(acc4545, a_tile45, b_tile45); + acc6745 = vmmlaq_s32(acc6745, a_tile67, b_tile45); + + acc0167 = vmmlaq_s32(acc0167, a_tile01, b_tile67); + acc2367 = vmmlaq_s32(acc2367, a_tile23, b_tile67); + acc4567 = vmmlaq_s32(acc4567, a_tile45, b_tile67); + acc6767 = vmmlaq_s32(acc6767, a_tile67, b_tile67); + + a_tile += 4 * TileSize; + b_tile += 4 * TileSize; + } + + store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, + std::min(2, m)); + store_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc, + std::min(2, std::max(0, m - 2))); + store_acc_rowpair(acc4501, acc4523, acc4545, acc4567, c_ptr + 4 * ldc, ldc, + std::min(2, std::max(0, m - 4))); + store_acc_rowpair(acc6701, acc6723, acc6745, acc6767, c_ptr + 6 * ldc, ldc, + std::min(2, std::max(0, m - 6))); +} + +FORCE_INLINE void gemm_micro_smmla_4x16_packed_a( + const int8_t* __restrict__ a_packed, const int8_t* __restrict__ b_packed, + int32_t* __restrict__ c_ptr, const int32_t m, const int32_t k_size, + const int64_t b_n_group_stride, const int64_t ldc) { + const int32_t m_rows_01 = std::min(2, m); + const int32_t m_rows_23 = std::min(2, std::max(0, m - 2)); + const int32x4_t zero = vdupq_n_s32(0); + + int32x4_t acc0101 = zero, acc0123 = zero, acc0145 = zero, acc0167 = zero; + int32x4_t acc2301 = zero, acc2323 = zero, acc2345 = zero, acc2367 = zero; + int32x4_t acc0189 = zero, acc011011 = zero, acc011213 = zero, + acc011415 = zero; + int32x4_t acc2389 = zero, acc231011 = zero, acc231213 = zero, + acc231415 = zero; + + const int8_t* __restrict__ a_tile = a_packed; + // note: b packs 8 panels contiguously, so we need 2 b_tile ptrs + // for the 4x16 microkernel + const int8_t* __restrict__ b_tile0 = b_packed; + const int8_t* __restrict__ b_tile1 = b_packed + b_n_group_stride; + +#pragma GCC unroll 8 + for (int32_t k_idx = 0; k_idx < k_size; k_idx += K) { + const int8x16_t a_tile01 = vld1q_s8(a_tile); + const int8x16_t a_tile23 = vld1q_s8(a_tile + TileSize); + const int8x16_t b_tile01 = vld1q_s8(b_tile0); + const int8x16_t b_tile23 = vld1q_s8(b_tile0 + TileSize); + const int8x16_t b_tile45 = vld1q_s8(b_tile0 + 2 * TileSize); + const int8x16_t b_tile67 = vld1q_s8(b_tile0 + 3 * TileSize); + const int8x16_t b_tile89 = vld1q_s8(b_tile1); + const int8x16_t b_tile1011 = vld1q_s8(b_tile1 + TileSize); + const int8x16_t b_tile1213 = vld1q_s8(b_tile1 + 2 * TileSize); + const int8x16_t b_tile1415 = vld1q_s8(b_tile1 + 3 * TileSize); + + acc0101 = vmmlaq_s32(acc0101, a_tile01, b_tile01); + acc2301 = vmmlaq_s32(acc2301, a_tile23, b_tile01); + acc0123 = vmmlaq_s32(acc0123, a_tile01, b_tile23); + acc2323 = vmmlaq_s32(acc2323, a_tile23, b_tile23); + + acc0145 = vmmlaq_s32(acc0145, a_tile01, b_tile45); + acc2345 = vmmlaq_s32(acc2345, a_tile23, b_tile45); + acc0167 = vmmlaq_s32(acc0167, a_tile01, b_tile67); + acc2367 = vmmlaq_s32(acc2367, a_tile23, b_tile67); + + acc0189 = vmmlaq_s32(acc0189, a_tile01, b_tile89); + acc2389 = vmmlaq_s32(acc2389, a_tile23, b_tile89); + acc011011 = vmmlaq_s32(acc011011, a_tile01, b_tile1011); + acc231011 = vmmlaq_s32(acc231011, a_tile23, b_tile1011); + + acc011213 = vmmlaq_s32(acc011213, a_tile01, b_tile1213); + acc231213 = vmmlaq_s32(acc231213, a_tile23, b_tile1213); + acc011415 = vmmlaq_s32(acc011415, a_tile01, b_tile1415); + acc231415 = vmmlaq_s32(acc231415, a_tile23, b_tile1415); + + a_tile += 2 * TileSize; + b_tile0 += 4 * TileSize; + b_tile1 += 4 * TileSize; + } + + // rows 0-1, columns 0-7 + store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, m_rows_01); + // rows 0-1, columns 8-15 + store_acc_rowpair(acc0189, acc011011, acc011213, acc011415, c_ptr + 8, ldc, + m_rows_01); + // rows 2-3, columns 0-7 + store_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc, + m_rows_23); + // rows 2-3, columns 8-15 + store_acc_rowpair(acc2389, acc231011, acc231213, acc231415, + c_ptr + 2 * ldc + 8, ldc, m_rows_23); +} + +} // namespace neon_smmla + +template <typename scalar_t> +class MicroGemmINT8<cpu_utils::ISA::NEON, scalar_t> { + public: + static constexpr int32_t K = neon_smmla::K; + static constexpr int32_t Mr = 8; + static constexpr int32_t Nr = 8; + static constexpr int32_t NrGemv = 16; + static constexpr int32_t MaxMSize = 8; + static constexpr int32_t NSize = 32; + static constexpr int32_t WeightOCGroupSize = Nr; + static_assert(MaxMSize % Mr == 0); + + static FORCE_INLINE void quantize_row(const scalar_t* input, int8_t* output, + float& scale, const int32_t size) { + TORCH_CHECK_EQ(size % K, 0); + float32x4_t max_vec = vdupq_n_f32(0.0f); + + for (int32_t i = 0; i < size; i += K) { + const float32x4x2_t input_vec = neon_smmla::load_as_f32(input + i); + max_vec = vmaxq_f32(max_vec, vabsq_f32(input_vec.val[0])); + max_vec = vmaxq_f32(max_vec, vabsq_f32(input_vec.val[1])); + } + + const float abs_max = std::max(vmaxvq_f32(max_vec), 1.0e-7f); + scale = abs_max / 127.0f; + const float32x4_t inv_scale_vec = vdupq_n_f32(127.0f / abs_max); + + for (int32_t i = 0; i < size; i += K) { + const float32x4x2_t input_vec = neon_smmla::load_as_f32(input + i); + const int32x4_t output_low = + vcvtnq_s32_f32(vmulq_f32(input_vec.val[0], inv_scale_vec)); + const int32x4_t output_high = + vcvtnq_s32_f32(vmulq_f32(input_vec.val[1], inv_scale_vec)); + const int16x8_t output_s16 = + vcombine_s16(vqmovn_s32(output_low), vqmovn_s32(output_high)); + vst1_s8(output + i, vqmovn_s16(output_s16)); + } + } + + // with current code, fusing this into the gemm micro kernel didn't move the + // needle + static FORCE_INLINE void dequantize_tile( + int32_t* input, float* output, const float* __restrict__ input_scales, + const float* __restrict__ weight_scales, const int32_t m, const int32_t n, + const int32_t stride) { + TORCH_CHECK_EQ(n % 4, 0); + for (int32_t m_idx = 0; m_idx < m; ++m_idx) { + const float32x4_t input_scale_vec = vdupq_n_f32(input_scales[m_idx]); + for (int32_t n_idx = 0; n_idx < n; n_idx += 4) { + const int32x4_t input_vec = vld1q_s32(input + m_idx * stride + n_idx); + const float32x4_t weight_scale_vec = vld1q_f32(weight_scales + n_idx); + const float32x4_t output_vec = + vmulq_f32(vcvtq_f32_s32(input_vec), + vmulq_f32(input_scale_vec, weight_scale_vec)); + vst1q_f32(output + m_idx * stride + n_idx, output_vec); + } + } + } + + // physical layout [ + // M / (8 or 4); Mr is 8 or 4 + // K / 8; K for smmla is 8 + // 4, ; 4 row-pairs for each 8 rows + // 2, ; row-pair is 2 rows + // 4 ; 4 elements per row + // ] + static void pack_input_from_rows(const int8_t* const* __restrict__ rows, + int8_t* __restrict__ a_packed, + const int32_t m, const int32_t k) { + TORCH_CHECK(m > 0 && m <= MaxMSize); + TORCH_CHECK(k % K == 0); + const int8x8_t zero = vdup_n_s8(0); + + for (int32_t row_base = 0; row_base < m; row_base += Mr) { + const int32_t panel_m = std::min(Mr, m - row_base); + const int8_t* const* panel_rows = rows + row_base; + int8_t* __restrict__ out = a_packed + row_base * k; + + // fast path for full 8-row panels (fast path for 4-row panels didn't move + // the needle) + if (panel_m == Mr) { + const int8_t* __restrict__ row0 = panel_rows[0]; + const int8_t* __restrict__ row1 = panel_rows[1]; + const int8_t* __restrict__ row2 = panel_rows[2]; + const int8_t* __restrict__ row3 = panel_rows[3]; + const int8_t* __restrict__ row4 = panel_rows[4]; + const int8_t* __restrict__ row5 = panel_rows[5]; + const int8_t* __restrict__ row6 = panel_rows[6]; + const int8_t* __restrict__ row7 = panel_rows[7]; + int32_t k_idx = 0; + for (; k_idx + 2 * K <= k; k_idx += 2 * K) { + int8_t* __restrict__ block0 = out; + int8_t* __restrict__ block1 = out + 4 * neon_smmla::TileSize; + + int8x16_t a0 = vld1q_s8(row0 + k_idx); + int8x16_t a1 = vld1q_s8(row1 + k_idx); + vst1q_s8(block0, vcombine_s8(vget_low_s8(a0), vget_low_s8(a1))); + vst1q_s8(block1, vcombine_s8(vget_high_s8(a0), vget_high_s8(a1))); + + a0 = vld1q_s8(row2 + k_idx); + a1 = vld1q_s8(row3 + k_idx); + vst1q_s8(block0 + neon_smmla::TileSize, + vcombine_s8(vget_low_s8(a0), vget_low_s8(a1))); + vst1q_s8(block1 + neon_smmla::TileSize, + vcombine_s8(vget_high_s8(a0), vget_high_s8(a1))); + + a0 = vld1q_s8(row4 + k_idx); + a1 = vld1q_s8(row5 + k_idx); + vst1q_s8(block0 + 2 * neon_smmla::TileSize, + vcombine_s8(vget_low_s8(a0), vget_low_s8(a1))); + vst1q_s8(block1 + 2 * neon_smmla::TileSize, + vcombine_s8(vget_high_s8(a0), vget_high_s8(a1))); + + a0 = vld1q_s8(row6 + k_idx); + a1 = vld1q_s8(row7 + k_idx); + vst1q_s8(block0 + 3 * neon_smmla::TileSize, + vcombine_s8(vget_low_s8(a0), vget_low_s8(a1))); + vst1q_s8(block1 + 3 * neon_smmla::TileSize, + vcombine_s8(vget_high_s8(a0), vget_high_s8(a1))); + + out += 8 * neon_smmla::TileSize; + } + + for (; k_idx < k; k_idx += K) { + int8x8_t a0 = vld1_s8(row0 + k_idx); + int8x8_t a1 = vld1_s8(row1 + k_idx); + vst1q_s8(out, vcombine_s8(a0, a1)); + + a0 = vld1_s8(row2 + k_idx); + a1 = vld1_s8(row3 + k_idx); + vst1q_s8(out + neon_smmla::TileSize, vcombine_s8(a0, a1)); + + a0 = vld1_s8(row4 + k_idx); + a1 = vld1_s8(row5 + k_idx); + vst1q_s8(out + 2 * neon_smmla::TileSize, vcombine_s8(a0, a1)); + + a0 = vld1_s8(row6 + k_idx); + a1 = vld1_s8(row7 + k_idx); + vst1q_s8(out + 3 * neon_smmla::TileSize, vcombine_s8(a0, a1)); + + out += 4 * neon_smmla::TileSize; + } + continue; + } + + const int32_t row_pairs = (panel_m <= 4) ? 2 : Mr / 2; + for (int32_t k_idx = 0; k_idx < k; k_idx += K) { + for (int32_t pair_idx = 0; pair_idx < row_pairs; ++pair_idx) { + const int32_t row_idx = pair_idx * 2; + const int8x8_t row0 = + (row_idx < panel_m) ? vld1_s8(panel_rows[row_idx] + k_idx) : zero; + const int8x8_t row1 = (row_idx + 1 < panel_m) + ? vld1_s8(panel_rows[row_idx + 1] + k_idx) + : zero; + vst1q_s8(out, vcombine_s8(row0, row1)); + out += neon_smmla::TileSize; + } + } + } + } + + // physical layout [ + // N / 8; Nr is 8 + // K / 8; K for smmla is 8 + // 4, ; 4 col-pairs for each 8 cols + // 2, ; col-pair is 2 cols + // 4 ; 4 elements per col + // ] + static void pack_weight(const int8_t* __restrict__ weight, + int8_t* __restrict__ packed_weight, + const int32_t output_size, const int32_t input_size) { + TORCH_CHECK(output_size % NSize == 0); + TORCH_CHECK(input_size % K == 0); + + for (int32_t o_idx = 0; o_idx < output_size; o_idx += Nr) { + int8_t* __restrict__ dst = packed_weight + o_idx * input_size; + for (int32_t k_idx = 0; k_idx < input_size; k_idx += K) { + for (int32_t pair_idx = 0; pair_idx < Nr; + pair_idx += neon_smmla::Cols) { + const int8_t* __restrict__ row0 = + weight + (o_idx + pair_idx) * input_size + k_idx; + const int8_t* __restrict__ row1 = row0 + input_size; + vst1q_s8(dst, vcombine_s8(vld1_s8(row0), vld1_s8(row1))); + dst += neon_smmla::TileSize; + } + } + } + } + + void gemm(const int8_t* __restrict__ a_packed, + const int8_t* __restrict__ b_packed, int32_t* __restrict__ c, + const int32_t m, const int32_t k, const int64_t b_n_group_stride, + const int64_t ldc) const { + TORCH_CHECK(m > 0 && m <= MaxMSize); + TORCH_CHECK(k % K == 0); + + for (int32_t n_idx = 0; n_idx < NSize; n_idx += NrGemv) { + const int8_t* __restrict__ b_panel = b_packed + n_idx * k; + + for (int32_t row_base = 0; row_base < m; row_base += Mr) { + const int32_t panel_m = std::min(Mr, m - row_base); + const int8_t* __restrict__ a_panel = a_packed + row_base * k; + int32_t* __restrict__ c_panel = c + row_base * ldc + n_idx; + + if (panel_m <= 4) { + neon_smmla::gemm_micro_smmla_4x16_packed_a( + a_panel, b_panel, c_panel, panel_m, k, b_n_group_stride, ldc); + } else { + neon_smmla::gemm_micro_smmla_8x8_packed_a(a_panel, b_panel, c_panel, + panel_m, k, ldc); + neon_smmla::gemm_micro_smmla_8x8_packed_a( + a_panel, b_panel + b_n_group_stride, c_panel + Nr, panel_m, k, + ldc); + } + } + } + } +}; + +} // namespace cpu_micro_gemm + +#endif diff --git a/csrc/cpu/micro_gemm/cpu_micro_gemm_neon.hpp b/csrc/cpu/micro_gemm/cpu_micro_gemm_neon.hpp index 7d4898852bb..b38337956f6 100644 --- a/csrc/cpu/micro_gemm/cpu_micro_gemm_neon.hpp +++ b/csrc/cpu/micro_gemm/cpu_micro_gemm_neon.hpp @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + #ifndef CPU_MICRO_GEMM_NEON_HPP #define CPU_MICRO_GEMM_NEON_HPP @@ -16,9 +19,6 @@ namespace { constexpr int32_t K = 4; constexpr int32_t Cols = 2; constexpr int32_t TileSize = K * Cols; -constexpr int32_t Mr = 8; -constexpr int32_t Nr = 8; -constexpr int32_t Nr_gemv = 16; // a = [a0, a1, a2, a3], b = [b0, b1, b2, b3] -> [a0, a1, b0, b1] FORCE_INLINE float32x4_t zip1_f32x4(const float32x4_t a, const float32x4_t b) { @@ -132,7 +132,7 @@ FORCE_INLINE void gemm_micro_bfmmla_8x8_packed_a( acc6767 = vbfmmlaq_f32(acc6767, a_tile67, b_tile67); a_tile += 4 * TileSize; - b_tile += Nr * K; + b_tile += 4 * TileSize; } store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, @@ -205,8 +205,8 @@ FORCE_INLINE void gemm_micro_bfmmla_4x16_packed_a( acc231415 = vbfmmlaq_f32(acc231415, a_tile23, b_tile1415); a_tile += 2 * TileSize; - b_tile0 += Nr * K; - b_tile1 += Nr * K; + b_tile0 += 4 * TileSize; + b_tile1 += 4 * TileSize; } store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, m_rows_01); @@ -223,6 +223,9 @@ FORCE_INLINE void gemm_micro_bfmmla_4x16_packed_a( template <typename scalar_t> class MicroGemm<cpu_utils::ISA::NEON, scalar_t> { public: + static constexpr int32_t Mr = 8; + static constexpr int32_t Nr = 8; + static constexpr int32_t NrGemv = 16; static constexpr int32_t MaxMSize = 8; static constexpr int32_t NSize = 32; static constexpr int32_t WeightOCGroupSize = Nr; @@ -246,6 +249,9 @@ class MicroGemm<cpu_utils::ISA::NEON, c10::BFloat16> { public: using scalar_t = c10::BFloat16; + static constexpr int32_t Mr = 8; + static constexpr int32_t Nr = 8; + static constexpr int32_t NrGemv = 16; static constexpr int32_t MaxMSize = 8; static constexpr int32_t NSize = 32; static constexpr int32_t WeightOCGroupSize = Nr; @@ -253,7 +259,7 @@ class MicroGemm<cpu_utils::ISA::NEON, c10::BFloat16> { public: // physical layout [ - // M / 8; Mr is 8 + // M / (8 or 4); Mr is 8 or 4 // K / 4; K for bfmmla is 4 // 4, ; 4 row-pairs for each 8 rows // 2, ; row-pair is 2 rows @@ -439,7 +445,7 @@ class MicroGemm<cpu_utils::ISA::NEON, c10::BFloat16> { (void)lda; // A is packed, so lda is not needed TORCH_CHECK_EQ(k % K, 0); - for (int32_t n_idx = 0; n_idx < NSize; n_idx += Nr_gemv) { + for (int32_t n_idx = 0; n_idx < NSize; n_idx += NrGemv) { const bfloat16_t* __restrict__ b_panel = reinterpret_cast<const bfloat16_t*>(b_ptr) + n_idx * k; diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 8b7d924dace..d5876202f15 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -207,6 +207,20 @@ void cpu_fused_moe(torch::Tensor& output, const torch::Tensor& input, const torch::Tensor& topk_id, const bool skip_weighted, const std::string& act, const std::string& isa); +void prepack_moe_weight_int8(const torch::Tensor& weight, + torch::Tensor& packed_weight, + const std::string& isa); + +void cpu_fused_moe_int8(torch::Tensor& output, const torch::Tensor& input, + const torch::Tensor& w13, const torch::Tensor& w2, + const torch::Tensor& w13_scale, + const torch::Tensor& w2_scale, + const std::optional<torch::Tensor>& w13_bias, + const std::optional<torch::Tensor>& w2_bias, + const torch::Tensor& topk_weights, + const torch::Tensor& topk_id, const bool skip_weighted, + const std::string& act, const std::string& isa); + void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc, const torch::Tensor positions, const torch::Tensor block_table, @@ -596,8 +610,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { #endif // fused moe -#if defined(__AVX512F__) || \ - (defined(__aarch64__) && !defined(__APPLE__) && defined(ARM_BF16_SUPPORT)) +#if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT) && !defined(__APPLE__)) ops.def( "prepack_moe_weight(Tensor weight, Tensor(a1!) packed_weight, str isa) " "-> ()"); @@ -608,7 +621,22 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "bool skip_weighted, " "str act, str isa) -> ()"); ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe); -#endif +#endif // #if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT) && + // !defined(__APPLE__)) +#if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) && \ + !defined(__APPLE__) + ops.def( + "prepack_moe_weight_int8(Tensor weight, Tensor(a1!) packed_weight, " + "str isa) -> ()"); + ops.impl("prepack_moe_weight_int8", torch::kCPU, &prepack_moe_weight_int8); + ops.def( + "cpu_fused_moe_int8(Tensor(a0!) output, Tensor input, Tensor w13, " + "Tensor w2, Tensor w13_scale, Tensor w2_scale, Tensor? w13_bias, " + "Tensor? w2_bias, Tensor topk_weights, Tensor topk_id, bool " + "skip_weighted, str act, str isa) -> ()"); + ops.impl("cpu_fused_moe_int8", torch::kCPU, &cpu_fused_moe_int8); +#endif // #if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) && + // !defined(__APPLE__) ops.def( "mla_decode_kvcache(" " Tensor! out, Tensor query, Tensor kv_cache," diff --git a/tests/kernels/moe/test_cpu_fused_moe.py b/tests/kernels/moe/test_cpu_fused_moe.py index 41ae9be5173..75e34e6766a 100644 --- a/tests/kernels/moe/test_cpu_fused_moe.py +++ b/tests/kernels/moe/test_cpu_fused_moe.py @@ -7,10 +7,14 @@ import torch from tests.kernels.allclose_default import get_default_atol, get_default_rtol from vllm._custom_ops import ( cpu_fused_moe, + cpu_fused_moe_int8, cpu_prepack_moe_weight, + cpu_prepack_moe_weight_int8, ) from vllm.model_executor.layers.fused_moe.activation import MoEActivation -from vllm.model_executor.layers.fused_moe.cpu_fused_moe import _CPU_MOE_ACT_FN +from vllm.model_executor.layers.fused_moe.cpu_fused_moe import ( + _CPU_MOE_ACT_FN, +) from vllm.platforms import CpuArchEnum, current_platform from vllm.utils.torch_utils import set_random_seed @@ -101,6 +105,90 @@ def ref_fused_moe( return final_out +def quantize_per_channel( + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Symmetrically quantize each weight output channel.""" + weight_fp32 = weight.float() + scale = weight_fp32.abs().amax(dim=-1).clamp_min(1e-12) / 127.0 + quantized = ( + (weight_fp32 / scale.unsqueeze(-1)).round().clamp(-127, 127).to(torch.int8) + ) + return quantized, scale + + +def quantize_per_token( + input: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Symmetrically quantize each input token.""" + input_fp32 = input.float() + scale = input_fp32.abs().amax(dim=-1, keepdim=True).clamp_min(1e-7) / 127.0 + quantized = (input_fp32 / scale).round().clamp(-127, 127).to(torch.int8) + return quantized, scale + + +def ref_fused_moe_int8( + input: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_bias: torch.Tensor | None, + w2_bias: torch.Tensor | None, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, +) -> torch.Tensor: + """Reference the two dynamically quantized INT8 GEMMs.""" + input_int8, input_scale = quantize_per_token(input) + expert_num = w13.size(0) + + counts = topk_ids.new_zeros((topk_ids.size(0), expert_num)) + counts.scatter_(1, topk_ids.to(torch.int64), 1) + tokens_per_expert = counts.sum(dim=0).tolist() + sorted_route_ids = topk_ids.view(-1).argsort() + sorted_token_ids = sorted_route_ids // topk_ids.size(1) + + outputs = [] + start_idx = 0 + for expert_idx, token_count in enumerate(tokens_per_expert): + end_idx = start_idx + token_count + if token_count == 0: + continue + + token_ids = sorted_token_ids[start_idx:end_idx] + gate_up = torch.matmul( + input_int8[token_ids].float(), + w13[expert_idx].float().T, + ) + gate_up *= input_scale[token_ids] * w13_scale[expert_idx] + if w13_bias is not None: + gate_up += w13_bias[expert_idx].float() + + intermediate = _CPU_MOE_ACT_FN[activation](gate_up).to(input.dtype) + intermediate_int8, intermediate_scale = quantize_per_token(intermediate) + output = torch.matmul( + intermediate_int8.float(), + w2[expert_idx].float().T, + ) + output *= intermediate_scale * w2_scale[expert_idx] + if w2_bias is not None: + output += w2_bias[expert_idx].float() + + outputs.append(output) + start_idx = end_idx + + sorted_output = torch.cat(outputs, dim=0) + routed_output = torch.empty_like(sorted_output) + routed_output[sorted_route_ids] = sorted_output + return ( + routed_output.view(*topk_ids.shape, -1) + .mul_(topk_weights.unsqueeze(-1)) + .sum(dim=1) + .to(input.dtype) + ) + + @pytest.mark.parametrize("batch_size", BATCH_SIZE) @pytest.mark.parametrize("expert_num", EXPERT_NUM) @pytest.mark.parametrize("hidden_size", HIDDEN_DIM) @@ -176,3 +264,85 @@ def test_cpu_fused_moe( torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol), f"{torch.max(torch.abs(output - ref_output))}", ) + + +@pytest.mark.skipif( + current_platform.get_cpu_architecture() != CpuArchEnum.ARM, + reason="Requires Arm CPU", +) +@pytest.mark.parametrize("batch_size", BATCH_SIZE) +@pytest.mark.parametrize("expert_num", EXPERT_NUM) +@pytest.mark.parametrize("hidden_size", HIDDEN_DIM) +@pytest.mark.parametrize("intermediate_size", INTERMEDIATE_DIM) +@pytest.mark.parametrize("use_bias", USE_BIAS) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("act", ACT) +@pytest.mark.parametrize("isa", ["neon"]) +def test_cpu_fused_moe_int8( + batch_size: int, + expert_num: int, + hidden_size: int, + intermediate_size: int, + use_bias: bool, + dtype: torch.dtype, + act: MoEActivation, + isa: str, +): + set_random_seed(0) + topk_num = max(expert_num // 2, 1) + up_dim = 2 * intermediate_size + + input = torch.randn((batch_size, hidden_size), dtype=dtype) / ( + 0.5 * hidden_size**0.5 + ) + w13_fp = torch.randn((expert_num, up_dim, hidden_size), dtype=dtype) / ( + 0.5 * hidden_size**0.5 + ) + w2_fp = torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype) / ( + 0.5 * intermediate_size**0.5 + ) + w13, w13_scale = quantize_per_channel(w13_fp) + w2, w2_scale = quantize_per_channel(w2_fp) + + w13_bias = None + w2_bias = None + if use_bias: + w13_bias = torch.randn((expert_num, up_dim), dtype=dtype) / (0.5 * up_dim**0.5) + w2_bias = torch.randn((expert_num, hidden_size), dtype=dtype) / ( + 0.5 * hidden_size**0.5 + ) + + router_logits = torch.randn((batch_size, expert_num), dtype=dtype) + score = torch.softmax(router_logits, dim=-1, dtype=torch.float32) + topk_weights, topk_ids = torch.topk(score, topk_num) + topk_ids = topk_ids.to(torch.int32) + + ref_output = ref_fused_moe_int8( + input, + w13, + w2, + w13_scale, + w2_scale, + w13_bias, + w2_bias, + topk_weights, + topk_ids, + act, + ) + packed_w13 = cpu_prepack_moe_weight_int8(w13, isa) + packed_w2 = cpu_prepack_moe_weight_int8(w2, isa) + output = cpu_fused_moe_int8( + input, + packed_w13, + packed_w2, + w13_scale, + w2_scale, + w13_bias, + w2_bias, + topk_weights, + topk_ids, + act.value, + isa, + ) + + torch.testing.assert_close(output, ref_output, atol=2e-2, rtol=2e-2) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index aa75c50a516..82044c87c1a 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3787,6 +3787,15 @@ def cpu_prepack_moe_weight( return output +def cpu_prepack_moe_weight_int8( + weight: torch.Tensor, + isa: str, +) -> torch.Tensor: + output = torch.empty_like(weight) + torch.ops._C.prepack_moe_weight_int8(weight, output, isa) + return output + + def cpu_fused_moe( input: torch.Tensor, w13: torch.Tensor, @@ -3816,6 +3825,39 @@ def cpu_fused_moe( return output +def cpu_fused_moe_int8( + input: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_bias: torch.Tensor | None, + w2_bias: torch.Tensor | None, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + act: str, + isa: str, + skip_weighted: bool = False, +) -> torch.Tensor: + output = torch.empty_like(input) + torch.ops._C.cpu_fused_moe_int8( + output, + input, + w13, + w2, + w13_scale, + w2_scale, + w13_bias, + w2_bias, + topk_weights, + topk_ids, + skip_weighted, + act, + isa, + ) + return output + + if hasattr(torch.ops._qutlass_C, "matmul_mxf4_bf16_tn"): @register_fake("_qutlass_C::matmul_mxf4_bf16_tn") diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py index 3ed1734cb91..bf6b00377fe 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py @@ -9,6 +9,8 @@ from vllm._custom_ops import ( CPUQuantAlgo, CPUQuantMethod, convert_weight_packed_scale_zp, + cpu_fused_moe_int8, + cpu_prepack_moe_weight_int8, fused_experts_cpu, ) from vllm.model_executor.layers.fused_moe.activation import MoEActivation @@ -27,7 +29,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kInt8StaticChannelSym, kMxfp4Static, ) -from vllm.platforms import current_platform +from vllm.platforms import CpuArchEnum, current_platform # =========================================================================== # FP8 W8A16 MoE @@ -556,7 +558,14 @@ def prepare_int8_moe_layer_for_cpu( w13: torch.Tensor, w2: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - """VNNI-prepack INT8 MoE weights for CPU kernel.""" + """Prepack INT8 MoE weights for the current CPU architecture.""" + # SMMLA packing for AArch64 + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: + return ( + cpu_prepack_moe_weight_int8(w13, "neon"), + cpu_prepack_moe_weight_int8(w2, "neon"), + ) + # VNNI packing for x86 packed_w13 = torch.ops._C.convert_weight_packed(w13) packed_w2 = torch.ops._C.convert_weight_packed(w2) return packed_w13, packed_w2 @@ -586,7 +595,10 @@ class CPUExpertsInt8(mk.FusedMoEExpertsMonolithic): @staticmethod def _supports_current_device() -> bool: - return current_platform.is_cpu() + return ( + current_platform.is_cpu() + and current_platform.get_cpu_architecture() == CpuArchEnum.X86 + ) @staticmethod def _supports_no_act_and_mul() -> bool: @@ -701,3 +713,167 @@ class CPUExpertsInt8(mk.FusedMoEExpertsMonolithic): None, # limit True, # is_vnni ) + + +class ArmCPUExpertsInt8(mk.FusedMoEExpertsMonolithic): + """Arm INT8 MoE with per-token activation and channelwise weight quantization.""" + + @property + def expects_unquantized_inputs(self) -> bool: + return True + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def is_supported_config( + cls: type[mk.FusedMoEExperts], + moe_config: FusedMoEConfig, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + activation_format: mk.FusedMoEActivationFormat, + ) -> tuple[bool, str | None]: + supported, reason = mk.FusedMoEExperts.is_supported_config( + cls, + moe_config, + weight_key, + activation_key, + activation_format, + ) + if not supported: + return supported, reason + if moe_config.in_dtype not in ( + torch.float32, + torch.float16, + torch.bfloat16, + ): + return False, "kernel requires float32, float16, or bfloat16 activations" + if moe_config.hidden_dim % 32 != 0: + return False, "kernel requires hidden dim divisible by 32" + if moe_config.intermediate_size_per_partition % 32 != 0: + return False, "kernel requires intermediate dim divisible by 32" + return True, None + + @staticmethod + def _supports_current_device() -> bool: + return ( + current_platform.is_cpu() + and current_platform.get_cpu_architecture() == CpuArchEnum.ARM + and hasattr(torch.ops._C, "cpu_fused_moe_int8") + ) + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation in ( + MoEActivation.SILU, + MoEActivation.SWIGLUOAI, + MoEActivation.GELU, + MoEActivation.GELU_TANH, + ) + + @staticmethod + def _supports_parallel_config( + moe_parallel_config: FusedMoEParallelConfig, + ) -> bool: + return not moe_parallel_config.use_ep + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return (weight_key, activation_key) == ( + kInt8StaticChannelSym, + kInt8DynamicTokenSym, + ) + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return routing_method in [ + RoutingMethodType.Default, + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ] + + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + return True + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + from vllm.model_executor.utils import replace_parameter + + w13, w2 = prepare_int8_moe_layer_for_cpu(layer.w13_weight, layer.w2_weight) + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + apply_router_weight_on_input: bool, + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + from vllm.model_executor.layers.fused_moe.cpu_fused_moe import ( + select_experts, + ) + + topk_weights, topk_ids = select_experts( + hidden_states=hidden_states, + router_logits=router_logits, + use_grouped_topk=num_expert_group is not None, + top_k=self.moe_config.experts_per_token, + renormalize=self.moe_config.routing_method + in ( + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ), + topk_group=topk_group, + num_expert_group=num_expert_group, + scoring_func="softmax", + routed_scaling_factor=( + routed_scaling_factor if routed_scaling_factor is not None else 1.0 + ), + e_score_correction_bias=e_score_correction_bias, + ) + + if apply_router_weight_on_input: + assert topk_ids.size(1) == 1 + hidden_states.mul_(topk_weights.to(hidden_states.dtype)) + + assert self.w1_scale is not None + assert self.w2_scale is not None + return cpu_fused_moe_int8( + hidden_states, + w1, + w2, + self.w1_scale, + self.w2_scale, + self.w1_bias, + self.w2_bias, + topk_weights, + topk_ids, + activation.value, + "neon", + skip_weighted=apply_router_weight_on_input, + ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/int8.py b/vllm/model_executor/layers/fused_moe/oracle/int8.py index 5a2b4c3a75b..e2fb3e9114f 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int8.py @@ -24,7 +24,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kInt8StaticChannelSym, ) from vllm.model_executor.utils import replace_parameter -from vllm.platforms import current_platform logger = init_logger(__name__) @@ -41,20 +40,12 @@ def _get_priority_backends( """ Get available backends in priority order based on platform and config. """ - _AVAILABLE_BACKENDS = [ + return [ Int8MoeBackend.TRITON, Int8MoeBackend.HUMMING, Int8MoeBackend.CPU, ] - def _move_to_front(backends: list[Int8MoeBackend], backend: Int8MoeBackend) -> None: - backends.insert(0, backends.pop(backends.index(backend))) - - if current_platform.is_cpu(): - _move_to_front(_AVAILABLE_BACKENDS, Int8MoeBackend.CPU) - - return _AVAILABLE_BACKENDS - def backend_to_kernel_cls( backend: Int8MoeBackend, @@ -78,14 +69,13 @@ def backend_to_kernel_cls( HummingGroupedExperts, HummingIndexedExperts, ] - elif backend == Int8MoeBackend.CPU: from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( + ArmCPUExpertsInt8, CPUExpertsInt8, ) - return [CPUExpertsInt8] - + return [ArmCPUExpertsInt8, CPUExpertsInt8] else: raise ValueError(f"Unknown Int8 MoE backend: {backend.value}") From 544cb724c888dda1eff06a02cb0baca783c3dd86 Mon Sep 17 00:00:00 2001 From: Tianmu Li <tianmu.li@intel.com> Date: Sun, 26 Jul 2026 23:20:14 -0700 Subject: [PATCH 091/185] [CPU][Spec Decode] Optimize GDN conv path for speculative decoding (#48577) Signed-off-by: Li, Tianmu <tianmu.li@intel.com> Co-authored-by: Codex <codex@openai.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com> --- csrc/cpu/sgl-kernels/conv.cpp | 185 +++++++- csrc/cpu/torch_bindings.cpp | 5 +- tests/kernels/mamba/cpu/test_cpu_gdn_ops.py | 436 +++++++++++++++++- vllm/_custom_ops.py | 3 +- .../layers/mamba/ops/cpu/gdn_attention.py | 260 ++++++----- 5 files changed, 767 insertions(+), 122 deletions(-) diff --git a/csrc/cpu/sgl-kernels/conv.cpp b/csrc/cpu/sgl-kernels/conv.cpp index b918aed8bff..8cac9ff1f01 100644 --- a/csrc/cpu/sgl-kernels/conv.cpp +++ b/csrc/cpu/sgl-kernels/conv.cpp @@ -451,6 +451,90 @@ void causal_conv1d_update_kernel_impl( }); } +template <typename scalar_t> +void causal_conv1d_update_multi_kernel_impl( + scalar_t* __restrict__ out, + const scalar_t* __restrict__ input, + scalar_t* __restrict__ conv_states, + const scalar_t* __restrict__ weight, + const scalar_t* __restrict__ bias, + const int32_t* __restrict__ num_accepted_tokens, + const int32_t* __restrict__ conv_indices, + bool silu_activation, + int64_t batch, + int64_t dim, + int64_t seqlen, + int64_t width, + int64_t state_len, + int64_t conv_state_slot_stride) { + constexpr int64_t BLOCK_N = block_size_n() * 2; + const int64_t NB = div_up(dim, BLOCK_N); + + AT_DISPATCH_BOOL2(bias != nullptr, has_bias, silu_activation, has_silu, [&] { + at::parallel_for(0, batch * NB, 0, [&](int64_t begin, int64_t end) { + int64_t bs{0}, nb{0}; + data_index_init(begin, bs, batch, nb, NB); + + for (int64_t i = begin; i < end; ++i) { + const int64_t nb_start = nb * BLOCK_N; + const int64_t nb_size = std::min(dim - nb_start, BLOCK_N); + const int32_t conv_state_index = conv_indices[bs]; + const int32_t history_offset = num_accepted_tokens[bs] - 1; + + switch (width << 4 | nb_size >> 4) { + case 0x42: + tinygemm_kernel<scalar_t, 4, 32, has_bias, has_silu>::apply( + input + bs * seqlen * dim + nb_start, + weight + nb_start * width, + out + bs * seqlen * dim + nb_start, + has_bias ? bias + nb_start : nullptr, + conv_states + conv_state_index * conv_state_slot_stride + + history_offset * dim + nb_start, + true, + seqlen, + dim, + true); + break; + case 0x44: + tinygemm_kernel<scalar_t, 4, 64, has_bias, has_silu>::apply( + input + bs * seqlen * dim + nb_start, + weight + nb_start * width, + out + bs * seqlen * dim + nb_start, + has_bias ? bias + nb_start : nullptr, + conv_states + conv_state_index * conv_state_slot_stride + + history_offset * dim + nb_start, + true, + seqlen, + dim, + true); + break; + default: + TORCH_CHECK(false, "Unexpected block size, ", width, " x ", nb_size); + } + + data_index_step(bs, batch, nb, NB); + } + }); + }); + + at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) { + for (int64_t bs = begin; bs < end; ++bs) { + const int32_t conv_state_index = conv_indices[bs]; + const int32_t num_accepted = num_accepted_tokens[bs]; + scalar_t* state = conv_states + conv_state_index * conv_state_slot_stride; + + std::memmove( + state, + state + num_accepted * dim, + (state_len - seqlen) * dim * sizeof(scalar_t)); + std::memcpy( + state + (state_len - seqlen) * dim, + input + bs * seqlen * dim, + seqlen * dim * sizeof(scalar_t)); + } + }); +} + } // anonymous namespace // from [dim, width] or [N, K] @@ -545,7 +629,7 @@ at::Tensor get_block_indices(const std::optional<at::Tensor>& offsets, int64_t n // query_start_loc: (batch + 1) int32 // cache_indices: (batch) int32 // has_initial_state: (batch) bool -// conv_states: (..., dim, width - 1) itype +// conv_states: (..., dim, state_len) itype, where state_len >= width - 1 // activation: either None or "silu" or "swish" // pad_slot_id: int // @@ -586,11 +670,14 @@ at::Tensor causal_conv1d_fwd_cpu( CHECK_EQ(conv_states_val.scalar_type(), scalar_type); CHECK_GE(padded_batch, batch); CHECK_EQ(conv_states_val.size(1), dim); - CHECK_EQ(conv_states_val.size(2), width - 1); + const int64_t state_len = conv_states_val.size(2); + CHECK_GE(state_len, width - 1); // adjust `conv_states` to be contiguous on `dim` // should happen only once if (conv_states_val.stride(-2) != 1) { + TORCH_CHECK(state_len == width - 1, + "causal_conv1d_fwd_cpu: wide conv_states must be contiguous on dim."); auto conv_states_copy = conv_states_val.clone(); conv_states_val.as_strided_({padded_batch, dim, width - 1}, {(width - 1) * dim, 1, dim}); conv_states_val.copy_(conv_states_copy); @@ -651,14 +738,14 @@ at::Tensor causal_conv1d_fwd_cpu( // API aligned with GPUs // -// x: (batch, dim) or (batch, dim, seqlen) +// x: (batch, dim) or (batch, seqlen, dim) // conv_state: (..., dim, state_len), where state_len >= width - 1 // weight: (dim, width) // bias: (dim,) -// cache_seqlens: (batch,), dtype int32. +// num_accepted_tokens: (batch,), dtype int32. // conv_state_indices: (batch,), dtype int32 // pad_slot_id: int -// out: (batch, dim) or (batch, dim, seqlen) +// out: (batch, dim) or (batch, seqlen, dim) // at::Tensor causal_conv1d_update_cpu( const at::Tensor& x, @@ -666,7 +753,7 @@ at::Tensor causal_conv1d_update_cpu( const at::Tensor& weight, const std::optional<at::Tensor>& bias, bool silu_activation, - const std::optional<at::Tensor>& cache_seqlens, + const std::optional<at::Tensor>& num_accepted_tokens, const std::optional<at::Tensor>& conv_state_indices, int64_t pad_slot_id, bool is_vnni) { @@ -674,13 +761,13 @@ at::Tensor causal_conv1d_update_cpu( CHECK_CONTIGUOUS(weight); auto packed_w = is_vnni ? weight : causal_conv1d_weight_pack(weight); - // TODO: add multi-token prediction support - TORCH_CHECK(x.dim() == 2, "causal_conv1d_update_cpu: expect x to be 2D tensor."); - TORCH_CHECK(!cache_seqlens.has_value(), "causal_conv1d_update_cpu: don't support cache_seqlens."); + TORCH_CHECK( + x.dim() == 2 || x.dim() == 3, + "causal_conv1d_update_cpu: expect x to be 2D or 3D tensor."); int64_t batch = x.size(0); - int64_t dim = x.size(1); - int64_t seqlen = 1; + int64_t dim = x.dim() == 2 ? x.size(1) : x.size(2); + int64_t seqlen = x.dim() == 2 ? 1 : x.size(1); int64_t width = weight.size(-1); const auto scalar_type = x.scalar_type(); @@ -690,10 +777,84 @@ at::Tensor causal_conv1d_update_cpu( CHECK_EQ(conv_states.scalar_type(), scalar_type); CHECK_EQ(conv_states.size(1), dim); - CHECK_EQ(conv_states.size(2), width - 1); + const int64_t state_len = conv_states.size(2); + CHECK_GE(state_len, width - 1); + + if (x.dim() == 3) { + TORCH_CHECK( + num_accepted_tokens.has_value(), + "causal_conv1d_update_cpu: num_accepted_tokens is required for 3D x."); + TORCH_CHECK( + conv_state_indices.has_value(), + "causal_conv1d_update_cpu: conv_state_indices is required for 3D x."); + CHECK_OPTIONAL_SHAPE_DTYPE(num_accepted_tokens, batch, at::kInt); + TORCH_CHECK( + width == 4, + "causal_conv1d_update_cpu: support only width of 4 for 3D x."); + TORCH_CHECK( + seqlen > 0, + "causal_conv1d_update_cpu: expect non-empty sequence for 3D x."); + TORCH_CHECK( + state_len >= seqlen, + "causal_conv1d_update_cpu: state_len must be >= seqlen for 3D x."); + TORCH_CHECK( + conv_states.stride(-2) == 1 && conv_states.stride(-1) == dim, + "causal_conv1d_update_cpu: 3D x requires SD conv_states layout."); + + const int32_t* accepted_counts = + num_accepted_tokens.value().data_ptr<int32_t>(); + const int32_t* indices = conv_state_indices.value().data_ptr<int32_t>(); + const int64_t num_slots = conv_states.size(0); + for (int64_t bs = 0; bs < batch; ++bs) { + const int32_t num_accepted = accepted_counts[bs]; + const int32_t conv_state_index = indices[bs]; + TORCH_CHECK( + conv_state_index != pad_slot_id, + "causal_conv1d_update_cpu: 3D x does not support pad slots."); + TORCH_CHECK( + conv_state_index >= 0 && conv_state_index < num_slots, + "causal_conv1d_update_cpu: conv_state_indices out of range."); + TORCH_CHECK( + num_accepted >= 1 && num_accepted <= seqlen, + "causal_conv1d_update_cpu: num_accepted_tokens must be in [1, " + "seqlen]."); + TORCH_CHECK( + num_accepted - 1 + width - 1 <= state_len, + "causal_conv1d_update_cpu: history window exceeds conv_states."); + } + + int64_t conv_state_slot_stride = conv_states.stride(0); + at::Tensor out = at::empty_like(x); + AT_DISPATCH_REDUCED_FLOATING_TYPES( + scalar_type, "causal_conv1d_update_multi_kernel_impl", [&] { + causal_conv1d_update_multi_kernel_impl<scalar_t>( + out.data_ptr<scalar_t>(), + x.data_ptr<scalar_t>(), + conv_states.data_ptr<scalar_t>(), + packed_w.data_ptr<scalar_t>(), + conditional_data_ptr<scalar_t>(bias), + accepted_counts, + indices, + silu_activation, + batch, + dim, + seqlen, + width, + state_len, + conv_state_slot_stride); + }); + return out; + } + + TORCH_CHECK( + !num_accepted_tokens.has_value(), + "causal_conv1d_update_cpu: num_accepted_tokens is only supported for 3D " + "x."); // adjust `conv_states` to be contiguous on `dim` if (conv_states.stride(-2) != 1) { + TORCH_CHECK(state_len == width - 1, + "causal_conv1d_update_cpu: wide conv_states must be contiguous on dim."); int64_t num_cache_lines = conv_states.size(0); auto conv_states_copy = conv_states.clone(); conv_states.as_strided_({num_cache_lines, dim, width - 1}, {(width - 1) * dim, 1, dim}); diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index d5876202f15..80b9409d991 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -147,7 +147,7 @@ at::Tensor causal_conv1d_fwd_cpu( at::Tensor causal_conv1d_update_cpu( const at::Tensor& x, const at::Tensor& conv_states, const at::Tensor& weight, const std::optional<at::Tensor>& bias, - bool silu_activation, const std::optional<at::Tensor>& cache_seqlens, + bool silu_activation, const std::optional<at::Tensor>& num_accepted_tokens, const std::optional<at::Tensor>& conv_state_indices, int64_t pad_slot_id, bool is_vnni); @@ -516,7 +516,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def( "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor " "weight, Tensor? bias, bool silu_activation," - "Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, " + "Tensor? num_accepted_tokens, Tensor? conv_state_indices, int " + "pad_slot_id, " "bool is_vnni) -> Tensor"); ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu); #endif diff --git a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py index 083e9ff22aa..7924b5a1acd 100644 --- a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py +++ b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py @@ -2,12 +2,14 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import functools +import types import pytest import torch import torch.nn.functional as F import vllm._custom_ops as ops +from vllm.model_executor.layers.mamba.ops.cpu import gdn_attention from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -417,6 +419,207 @@ def _conv_inputs(total_tokens: int): return x, weight, bias +def _sd_conv_states( + num_slots: int, state_len: int, dim: int = CONV_DIM +) -> torch.Tensor: + storage = torch.zeros(num_slots, state_len, dim, dtype=torch.bfloat16) + return storage.transpose(1, 2) + + +def _maybe_pack_conv_weight(weight: torch.Tensor, is_vnni: bool) -> torch.Tensor: + return ops.causal_conv1d_weight_pack(weight) if is_vnni else weight + + +@torch.inference_mode() +def test_spec_aware_mixed_routing_preserves_token_order( + monkeypatch: pytest.MonkeyPatch, +) -> None: + num_tokens = 4 + projection = torch.arange(num_tokens * 16, dtype=torch.float32).view(num_tokens, 16) + mixed_qkv, b, a = projection[:, :4], projection[:, 4:6], projection[:, 6:8] + assert all(not tensor.is_contiguous() for tensor in (mixed_qkv, b, a)) + + spec_indices = torch.tensor([0, 2]) + nonspec_indices = torch.tensor([1, 3]) + metadata = types.SimpleNamespace( + spec_sequence_masks=torch.ones(1, dtype=torch.bool), + spec_token_indx=spec_indices, + non_spec_token_indx=nonspec_indices, + num_prefills=1, + num_decodes=0, + ) + routed = [] + + def record(*args): + routed.append(args[2:5]) + return args[2] + + monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: True) + monkeypatch.setattr(gdn_attention, "_spec_forward", record) + monkeypatch.setattr(gdn_attention, "_spec_aware_nonspec_subset", record) + + layer = types.SimpleNamespace( + kv_cache=[torch.empty(1, 4, 6), torch.empty(1, 1, 1, 1)] + ) + core_attn_out = torch.empty_like(mixed_qkv) + gdn_attention._cpu_gdn_attention_spec_aware( + layer=layer, + attn_metadata_i=metadata, + mixed_qkv=mixed_qkv, + b=b, + a=a, + core_attn_out=core_attn_out, + width=CONV_KERNEL, + state_len=6, + ) + + expected_inputs = (mixed_qkv, b, a) + assert len(routed) == 2 + for actual_inputs, indices in zip(routed, (spec_indices, nonspec_indices)): + for actual, expected in zip(actual_inputs, expected_inputs): + assert actual.is_contiguous() + torch.testing.assert_close(actual, expected.index_select(0, indices)) + torch.testing.assert_close(core_attn_out, mixed_qkv) + + +@torch.inference_mode() +def test_spec_aware_nonspec_materializes_state_indices( + monkeypatch: pytest.MonkeyPatch, +) -> None: + block_table = torch.arange(8, dtype=torch.int32).view(2, 4) + state_indices = block_table[:, 0] + assert not state_indices.is_contiguous() + + metadata = types.SimpleNamespace( + non_spec_state_indices_tensor=state_indices, + non_spec_query_start_loc=torch.tensor([0, 2, 4], dtype=torch.int32), + num_decodes=0, + num_decode_tokens=0, + num_prefills=2, + num_prefill_tokens=4, + has_initial_state=torch.tensor([False, False]), + ) + + recorded_indices = None + + def causal_conv1d_fwd_cpu(**kwargs): + nonlocal recorded_indices + recorded_indices = kwargs["cache_indices"] + return kwargs["x"] + + def fused_gdn_gating_cpu(**kwargs): + return kwargs["a"], kwargs["b"] + + def chunk_gated_delta_rule_cpu(**kwargs): + out = torch.zeros(1, 4, 1, 1) + return out, kwargs["initial_state"] + + monkeypatch.setattr(torch.cpu, "_is_amx_tile_supported", lambda: True) + monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: False) + monkeypatch.setattr( + gdn_attention.ops, "causal_conv1d_fwd_cpu", causal_conv1d_fwd_cpu + ) + monkeypatch.setattr(gdn_attention.ops, "fused_gdn_gating_cpu", fused_gdn_gating_cpu) + monkeypatch.setattr( + gdn_attention.ops, + "chunk_gated_delta_rule_cpu", + chunk_gated_delta_rule_cpu, + ) + + layer = types.SimpleNamespace( + activation="silu", + conv1d=types.SimpleNamespace(weight=torch.empty(0), bias=None), + A_log=torch.empty(0), + dt_bias=torch.empty(0), + rearrange_mixed_qkv=lambda x: ( + x[:, :1].view(1, 4, 1, 1), + x[:, :1].view(1, 4, 1, 1), + x[:, :1].view(1, 4, 1, 1), + ), + ) + gdn_attention._spec_aware_nonspec( + layer=layer, + attn_metadata_i=metadata, + mixed_qkv=torch.zeros(4, 4), + b=torch.zeros(4, 1), + a=torch.zeros(4, 1), + core_attn_out=torch.zeros(4, 1, 1), + conv_buf=torch.zeros(8, 1, 6), + ssm_state=torch.zeros(8, 1, 1, 1), + width=4, + ) + + assert recorded_indices is not None + assert recorded_indices.is_contiguous() + torch.testing.assert_close( + recorded_indices, torch.tensor([0, 4], dtype=torch.int32) + ) + + +@torch.inference_mode() +def test_spec_forward_prepares_native_conv_metadata( + monkeypatch: pytest.MonkeyPatch, +) -> None: + block_table = torch.tensor([[0, 1], [4, 5]], dtype=torch.int32) + state_indices = block_table[:, 0] + accepted_counts = torch.tensor([1, 4], dtype=torch.int32) + assert not state_indices.is_contiguous() + metadata = types.SimpleNamespace( + num_spec_decodes=2, + spec_state_indices_tensor=block_table, + spec_query_start_loc=torch.tensor([0, 4, 8], dtype=torch.int32), + num_accepted_tokens=accepted_counts, + ) + forwarded_indices = None + forwarded_counts = None + + def causal_conv1d_update_cpu(**kwargs): + nonlocal forwarded_counts, forwarded_indices + forwarded_indices = kwargs["conv_state_indices"] + forwarded_counts = kwargs["num_accepted_tokens"] + return kwargs["x"] + + monkeypatch.setattr(torch.cpu, "_is_amx_tile_supported", lambda: True) + monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: False) + monkeypatch.setattr( + gdn_attention.ops, "causal_conv1d_update_cpu", causal_conv1d_update_cpu + ) + monkeypatch.setattr( + gdn_attention.ops, + "fused_sigmoid_gating_delta_rule_update_spec_cpu", + lambda **kwargs: kwargs["q"], + ) + + layer = types.SimpleNamespace( + activation="silu", + conv1d=types.SimpleNamespace(weight=torch.empty(1, CONV_KERNEL), bias=None), + A_log=None, + dt_bias=None, + rearrange_mixed_qkv=lambda x: (x.unsqueeze(0),) * 3, + ) + gdn_attention._spec_forward( + layer=layer, + attn_metadata_i=metadata, + mixed_qkv_spec=torch.zeros(8, 1, dtype=torch.bfloat16), + b_spec=torch.empty(0), + a_spec=torch.empty(0), + conv_buf=torch.empty(0), + ssm_state=torch.empty(0), + width=CONV_KERNEL, + state_len=0, + ) + + expected = ( + (forwarded_indices, torch.tensor([0, 4], dtype=torch.int32)), + (forwarded_counts, torch.tensor([1, 4], dtype=torch.int32)), + ) + for actual, reference in expected: + assert actual is not None + assert actual.is_contiguous() + assert actual.dtype == torch.int32 + torch.testing.assert_close(actual, reference) + + @pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS) @torch.inference_mode() def test_causal_conv1d_torch_two_call_split(total_tokens: int, split: int) -> None: @@ -475,7 +678,182 @@ def test_causal_conv1d_torch_two_call_split(total_tokens: int, split: int) -> No @pytest.mark.skipif( not torch.cpu._is_amx_tile_supported(), - reason="causal_conv1d_fwd_cpu requires AMX/AVX512", + reason="requires AMX support", +) +@torch.inference_mode() +def test_causal_conv1d_update_cpu_accepts_wide_state() -> None: + state_len = CONV_KERNEL - 1 + wide_state_len = state_len + 5 + batch_size = 3 + is_vnni = True + x, weight, bias = _conv_inputs(batch_size) + conv_state_indices = torch.tensor([2, 0, 1], dtype=torch.int32) + + narrow_state = _sd_conv_states(batch_size, state_len) + narrow_state.copy_( + tensor_cache(narrow_state.numel(), torch.bfloat16).view_as(narrow_state) + ) + wide_state = _sd_conv_states(batch_size, wide_state_len) + wide_state[:, :, :state_len].copy_(narrow_state) + wide_state[:, :, state_len:].fill_(7) + wide_tail = wide_state[:, :, state_len:].clone() + + conv_weight = _maybe_pack_conv_weight(weight, is_vnni) + out_narrow = ops.causal_conv1d_update_cpu( + x=x, + conv_states=narrow_state, + weight=conv_weight, + bias=bias, + silu_activation=True, + conv_state_indices=conv_state_indices, + is_vnni=is_vnni, + ) + out_wide = ops.causal_conv1d_update_cpu( + x=x, + conv_states=wide_state, + weight=conv_weight, + bias=bias, + silu_activation=True, + conv_state_indices=conv_state_indices, + is_vnni=is_vnni, + ) + + torch.testing.assert_close(out_wide, out_narrow, atol=1e-2, rtol=1e-2) + torch.testing.assert_close( + wide_state[:, :, :state_len], narrow_state, atol=0, rtol=0 + ) + torch.testing.assert_close(wide_state[:, :, state_len:], wide_tail, atol=0, rtol=0) + + +def _ref_causal_conv1d_update_cpu_multi( + x: torch.Tensor, + conv_states: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + silu_activation: bool, + conv_state_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor, +) -> torch.Tensor: + batch_size, seq_len, dim = x.shape + state_len = conv_states.size(2) + conv_out = torch.empty_like(x) + conv_weight = weight.unsqueeze(1) + + for i in range(batch_size): + slot = int(conv_state_indices[i].item()) + offset = int(num_accepted_tokens[i].item()) - 1 + state = conv_states[slot] + x_seq = x[i].transpose(0, 1).to(state.dtype) + prior = state[:, offset : offset + CONV_KERNEL - 1] + conv_in = torch.cat([prior, x_seq], dim=-1).unsqueeze(0) + out = F.conv1d(conv_in, conv_weight, bias, groups=dim)[0] + if silu_activation: + out = F.silu(out) + conv_out[i] = out.transpose(0, 1).to(conv_out.dtype) + keep = state[:, offset + 1 : offset + 1 + (state_len - seq_len)] + state.copy_(torch.cat([keep, x_seq], dim=-1)) + + return conv_out + + +@pytest.mark.skipif( + not torch.cpu._is_amx_tile_supported(), + reason="requires AMX support", +) +@pytest.mark.parametrize( + ("batch_size, seq_len, accepted_counts, has_bias, silu_activation, is_vnni"), + [ + (1, 1, [1], False, False, False), + (1, 1, [1], True, True, True), + (4, 4, [1, 2, 3, 4], False, True, False), + (4, 4, [4, 3, 2, 1], True, False, True), + (4, 16, [1, 5, 10, 16], False, False, True), + (4, 16, [16, 10, 5, 1], True, True, False), + ], +) +@torch.inference_mode() +def test_causal_conv1d_update_cpu_multi_token_matches_python( + batch_size: int, + seq_len: int, + accepted_counts: list[int], + has_bias: bool, + silu_activation: bool, + is_vnni: bool, +) -> None: + dim = 96 + state_len = seq_len + 2 + x = tensor_cache(batch_size * seq_len * dim, torch.bfloat16).view( + batch_size, seq_len, dim + ) + weight = tensor_cache(dim * CONV_KERNEL, torch.bfloat16).view(dim, CONV_KERNEL) + bias = tensor_cache(dim, torch.bfloat16) if has_bias else None + conv_state_indices = torch.arange(batch_size - 1, -1, -1, dtype=torch.int32) + num_accepted_tokens = torch.tensor(accepted_counts, dtype=torch.int32) + + conv_states_ref = _sd_conv_states(batch_size, state_len, dim) + conv_states_ref.copy_( + tensor_cache(conv_states_ref.numel(), torch.bfloat16).view_as(conv_states_ref) + ) + conv_states = conv_states_ref.clone() + + conv_weight = _maybe_pack_conv_weight(weight, is_vnni) + out = ops.causal_conv1d_update_cpu( + x=x, + conv_states=conv_states, + weight=conv_weight, + bias=bias, + silu_activation=silu_activation, + conv_state_indices=conv_state_indices, + is_vnni=is_vnni, + num_accepted_tokens=num_accepted_tokens, + ) + ref_out = _ref_causal_conv1d_update_cpu_multi( + x=x, + conv_states=conv_states_ref, + weight=weight, + bias=bias, + silu_activation=silu_activation, + conv_state_indices=conv_state_indices, + num_accepted_tokens=num_accepted_tokens, + ) + + torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2) + torch.testing.assert_close(conv_states, conv_states_ref, atol=0, rtol=0) + + +@pytest.mark.skipif( + not torch.cpu._is_amx_tile_supported(), + reason="requires AMX support", +) +@pytest.mark.parametrize("num_accepted", [0, 17]) +@torch.inference_mode() +def test_causal_conv1d_update_cpu_rejects_invalid_accepted_count( + num_accepted: int, +) -> None: + batch_size = 1 + seq_len = 16 + dim = 96 + state_len = seq_len + 2 + x = torch.zeros(batch_size, seq_len, dim, dtype=torch.bfloat16) + weight = torch.zeros(dim, CONV_KERNEL, dtype=torch.bfloat16) + conv_states = _sd_conv_states(batch_size, state_len, dim) + + with pytest.raises(RuntimeError, match="num_accepted_tokens must be in.*seqlen"): + ops.causal_conv1d_update_cpu( + x=x, + conv_states=conv_states, + weight=weight, + bias=None, + silu_activation=True, + conv_state_indices=torch.tensor([0], dtype=torch.int32), + is_vnni=False, + num_accepted_tokens=torch.tensor([num_accepted], dtype=torch.int32), + ) + + +@pytest.mark.skipif( + not torch.cpu._is_amx_tile_supported(), + reason="requires AMX support", ) @pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS) @torch.inference_mode() @@ -515,6 +893,62 @@ def test_causal_conv1d_fwd_cpu_two_call_split(total_tokens: int, split: int) -> torch.testing.assert_close(out_split, out_full, atol=1e-2, rtol=1e-2) +@pytest.mark.skipif( + not torch.cpu._is_amx_tile_supported(), + reason="requires AMX support", +) +@torch.inference_mode() +def test_causal_conv1d_fwd_cpu_accepts_wide_state() -> None: + state_len = CONV_KERNEL - 1 + wide_state_len = state_len + 5 + is_vnni = True + seq_lens = [CHUNK_SIZE - 1, CHUNK_SIZE + 5] + total_tokens = sum(seq_lens) + x, weight, bias = _conv_inputs(total_tokens) + query_start_loc = torch.tensor([0, seq_lens[0], total_tokens], dtype=torch.int32) + cache_indices = torch.tensor([2, 0], dtype=torch.int32) + has_initial_state = torch.tensor([True, False]) + + narrow_state = _sd_conv_states(3, state_len) + narrow_state.copy_( + tensor_cache(narrow_state.numel(), torch.bfloat16).view_as(narrow_state) + ) + wide_state = _sd_conv_states(3, wide_state_len) + wide_state[:, :, :state_len].copy_(narrow_state) + wide_state[:, :, state_len:].fill_(7) + wide_tail = wide_state[:, :, state_len:].clone() + + conv_weight = _maybe_pack_conv_weight(weight, is_vnni) + out_narrow = ops.causal_conv1d_fwd_cpu( + x=x.transpose(0, 1), + weight=conv_weight, + bias=bias, + conv_states=narrow_state, + query_start_loc=query_start_loc, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + silu_activation=True, + is_vnni=is_vnni, + ) + out_wide = ops.causal_conv1d_fwd_cpu( + x=x.transpose(0, 1), + weight=conv_weight, + bias=bias, + conv_states=wide_state, + query_start_loc=query_start_loc, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + silu_activation=True, + is_vnni=is_vnni, + ) + + torch.testing.assert_close(out_wide, out_narrow, atol=1e-2, rtol=1e-2) + torch.testing.assert_close( + wide_state[:, :, :state_len], narrow_state, atol=0, rtol=0 + ) + torch.testing.assert_close(wide_state[:, :, state_len:], wide_tail, atol=0, rtol=0) + + @torch.inference_mode() def test_batch_memcpy_cpu_fallback() -> None: """The ctypes batch_memcpy fallback (used when triton-cpu is absent) must diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 82044c87c1a..07a3a583f95 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3505,6 +3505,7 @@ def causal_conv1d_update_cpu( silu_activation: bool, conv_state_indices: torch.Tensor | None, is_vnni: bool, + num_accepted_tokens: torch.Tensor | None = None, ) -> torch.Tensor: return torch.ops._C.causal_conv1d_update_cpu( x, @@ -3512,7 +3513,7 @@ def causal_conv1d_update_cpu( weight, bias, silu_activation, - None, + num_accepted_tokens, conv_state_indices, -1, is_vnni, diff --git a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py index f93d2e18db6..b20561c8259 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py +++ b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py @@ -306,25 +306,21 @@ def _cpu_gdn_attention_spec_aware( width: int, state_len: int, ) -> None: - mixed_qkv = mixed_qkv.contiguous() - a = a.contiguous() - b = b.contiguous() - spec_sequence_masks = attn_metadata_i.spec_sequence_masks conv_buf = _conv_buffer_view(layer) # (num_slots, dim, state_len) ssm_state = _ssm_state_view(layer) if spec_sequence_masks is None: # No spec sequences in this batch (e.g. the prompt prefill step while - # speculative decoding is configured). Process as prefill/decode using - # torch conv (which only touches the first ``width-1`` columns of the - # wide buffer, leaving the rolling history untouched). + # speculative decoding is configured). Process as ordinary + # prefill/decode while touching only the first ``width-1`` columns of + # the wide buffer, leaving the rolling history untouched. _spec_aware_nonspec( layer, attn_metadata_i, - mixed_qkv, - b, - a, + mixed_qkv.contiguous(), + b.contiguous(), + a.contiguous(), core_attn_out, conv_buf, ssm_state, @@ -339,9 +335,9 @@ def _cpu_gdn_attention_spec_aware( num_decodes = attn_metadata_i.num_decodes if num_prefills == 0 and num_decodes == 0: - mixed_qkv_spec = mixed_qkv - b_spec = b - a_spec = a + mixed_qkv_spec = mixed_qkv.contiguous() + b_spec = b.contiguous() + a_spec = a.contiguous() spec_out_indx = None else: assert spec_token_indx is not None @@ -405,50 +401,71 @@ def _spec_forward( assert spec_qsl is not None assert num_accepted is not None - spec_qsl_cpu = spec_qsl[: num_spec_decodes + 1].to("cpu", torch.int64) - num_acc_cpu = num_accepted[:num_spec_decodes].to("cpu", torch.int64) + spec_qsl_cpu = spec_qsl[: num_spec_decodes + 1] seq_starts = spec_qsl_cpu[:-1] seq_lens = spec_qsl_cpu[1:] - spec_qsl_cpu[:-1] # ---- 1. Convolution (per-sequence rolling buffer) ---- - w2d = _unpacked_conv_weight(layer) # (dim, width) - dim = w2d.size(0) - w = w2d.unsqueeze(1) # (dim, 1, width) for F.conv1d depthwise + dim = mixed_qkv_spec.size(-1) bias = layer.conv1d.bias silu = layer.activation == "silu" - conv_out = torch.empty_like(mixed_qkv_spec) - col0 = spec_state_indices[:, 0].to("cpu", torch.int64) - for i in range(num_spec_decodes): - q_i = int(seq_lens[i].item()) - if q_i == 0: - continue - start = int(seq_starts[i].item()) - slot0 = int(col0[i].item()) - a_prev = int(num_acc_cpu[i].item()) - offset = a_prev - 1 - B = conv_buf[slot0] # (dim, state_len) - x_seq = mixed_qkv_spec[start : start + q_i].transpose(0, 1).to(B.dtype) - prior = B[:, offset : offset + (width - 1)] - conv_in = torch.cat([prior, x_seq], dim=-1).unsqueeze(0) # (1, dim, w-1+q) - out = F.conv1d(conv_in, w, bias, groups=dim)[0] # (dim, q_i) - if silu: - out = F.silu(out) - conv_out[start : start + q_i] = out.transpose(0, 1).to(conv_out.dtype) - # Roll the buffer: drop ``a_prev`` from the front, append the new - # draft tokens, keep total length == state_len. - keep = B[:, offset + 1 : offset + 1 + (state_len - q_i)] - new_B = torch.cat([keep, x_seq], dim=-1) - B.copy_(new_B) + can_use_native_conv = ( + torch.cpu._is_amx_tile_supported() + and not is_conv_state_dim_first() + and width == 4 + and num_spec_decodes > 0 + and bool(torch.all(seq_lens == seq_lens[0]).item()) + and int(seq_lens[0].item()) > 0 + ) + if can_use_native_conv: + q_i = int(seq_lens[0].item()) + conv_out = ops.causal_conv1d_update_cpu( + x=mixed_qkv_spec.view(num_spec_decodes, q_i, dim), + conv_states=conv_buf, + weight=layer.conv1d.weight, + bias=bias, + silu_activation=silu, + conv_state_indices=spec_state_indices[:num_spec_decodes, 0] + .to("cpu", torch.int32) + .contiguous(), + is_vnni=True, + num_accepted_tokens=num_accepted[:num_spec_decodes].to("cpu", torch.int32), + ).view_as(mixed_qkv_spec) + else: + w = _unpacked_conv_weight(layer).unsqueeze(1) + col0 = spec_state_indices[:num_spec_decodes, 0] + num_acc_cpu = num_accepted[:num_spec_decodes] + conv_out = torch.empty_like(mixed_qkv_spec) + for i in range(num_spec_decodes): + q_i = int(seq_lens[i].item()) + if q_i == 0: + continue + start = int(seq_starts[i].item()) + slot0 = int(col0[i].item()) + offset = int(num_acc_cpu[i].item()) - 1 + B = conv_buf[slot0] # (dim, state_len) + x_seq = mixed_qkv_spec[start : start + q_i].transpose(0, 1).to(B.dtype) + prior = B[:, offset : offset + (width - 1)] + conv_in = torch.cat([prior, x_seq], dim=-1).unsqueeze(0) + out = F.conv1d(conv_in, w, bias, groups=dim)[0] # (dim, q_i) + if silu: + out = F.silu(out) + conv_out[start : start + q_i] = out.transpose(0, 1).to(conv_out.dtype) + # Roll the buffer: drop the accepted history from the front, append + # the new draft tokens, keep total length == state_len. + keep = B[:, offset + 1 : offset + 1 + (state_len - q_i)] + new_B = torch.cat([keep, x_seq], dim=-1) + B.copy_(new_B) # ---- 2. Recurrent (multi-slot SSM state) ---- # Single fused kernel call: it runs the recurrence over each sequence's # draft tokens internally, resumes from slot ``num_accepted-1`` and stores # the state after token ``t`` into slot ``t`` (rollback for the next step). query, key, value = layer.rearrange_mixed_qkv(conv_out) - query = query.squeeze(0).contiguous() - key = key.squeeze(0).contiguous() - value = value.squeeze(0).contiguous() + query = query.squeeze(0) + key = key.squeeze(0) + value = value.squeeze(0) spec_idx = spec_state_indices[:num_spec_decodes].to(torch.int32).contiguous() num_acc = num_accepted[:num_spec_decodes].to(torch.int32).contiguous() cu = spec_qsl[: num_spec_decodes + 1].to(torch.int32).contiguous() @@ -458,8 +475,8 @@ def _spec_forward( q=query, k=key, v=value, - a=a_spec.contiguous(), - b=b_spec.contiguous(), + a=a_spec, + b=b_spec, initial_state_source=ssm_state, spec_state_indices=spec_idx, num_accepted_tokens=num_acc, @@ -469,15 +486,6 @@ def _spec_forward( return out_spec -def core_attn_out_like(layer, mixed_qkv_spec: torch.Tensor) -> torch.Tensor: - num_tokens = mixed_qkv_spec.size(0) - return torch.zeros( - (num_tokens, layer.num_v_heads // layer.tp_size, layer.head_v_dim), - dtype=mixed_qkv_spec.dtype, - device=mixed_qkv_spec.device, - ) - - def _spec_aware_nonspec( layer, attn_metadata_i: GDNAttentionMetadata, @@ -489,13 +497,19 @@ def _spec_aware_nonspec( ssm_state: torch.Tensor, width: int, ) -> None: - """Non-spec prefill/decode with a wide conv buffer (torch path).""" + """Non-spec prefill/decode with a wide conv buffer.""" state_indices_tensor = attn_metadata_i.non_spec_state_indices_tensor query_start_loc = attn_metadata_i.non_spec_query_start_loc assert state_indices_tensor is not None assert query_start_loc is not None + state_indices_tensor = state_indices_tensor.contiguous() - conv_weights = _unpacked_conv_weight(layer) + is_amx = torch.cpu._is_amx_tile_supported() + if is_amx and is_conv_state_dim_first(): + raise RuntimeError("AMX GDN attention requires `SD` conv_state layout.") + + if not is_amx: + conv_weights = _unpacked_conv_weight(layer) num_decodes = attn_metadata_i.num_decodes num_decode_tokens = attn_metadata_i.num_decode_tokens @@ -507,42 +521,48 @@ def _spec_aware_nonspec( decode_b = b[:num_decode_tokens] decode_a = a[:num_decode_tokens] decode_state_indices = state_indices_tensor[:num_decodes] - # Only the first ``width-1`` columns hold the real conv state. - if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: - conv_state_view = conv_buf[:, :, : width - 1] - decode_conv_state = conv_state_view[decode_state_indices].contiguous() - decode_mixed_qkv = causal_conv1d_update_torch( - x=decode_mixed_qkv.unsqueeze(-1), - conv_state=decode_conv_state, - weight=conv_weights, - bias=layer.conv1d.bias, - activation=layer.activation, - ).squeeze(-1) - conv_state_view[decode_state_indices] = decode_conv_state - else: - decode_mixed_qkv = causal_conv1d_update_cpu( + if is_amx: + decode_mixed_qkv = ops.causal_conv1d_update_cpu( x=decode_mixed_qkv, - conv_state=conv_buf[:, :, : width - 1], - weight=conv_weights, + conv_states=conv_buf, + weight=layer.conv1d.weight, bias=layer.conv1d.bias, - activation=layer.activation, + silu_activation=layer.activation == "silu", conv_state_indices=decode_state_indices, + is_vnni=True, ) + else: + # Only the first ``width-1`` columns hold the real conv state. + conv_state_view = conv_buf[:, :, : width - 1] + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: + decode_conv_state = conv_state_view[decode_state_indices].contiguous() + decode_mixed_qkv = causal_conv1d_update_torch( + x=decode_mixed_qkv.unsqueeze(-1), + conv_state=decode_conv_state, + weight=conv_weights, + bias=layer.conv1d.bias, + activation=layer.activation, + ).squeeze(-1) + conv_state_view[decode_state_indices] = decode_conv_state + else: + decode_mixed_qkv = causal_conv1d_update_cpu( + x=decode_mixed_qkv, + conv_state=conv_state_view, + weight=conv_weights, + bias=layer.conv1d.bias, + activation=layer.activation, + conv_state_indices=decode_state_indices, + ) query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv) - # rearrange_mixed_qkv can return views whose last dim is not - # contiguous; the fused CPU kernel requires a contiguous last dim. - query = query.contiguous() - key = key.contiguous() - value = value.contiguous() attn_out = ops.fused_sigmoid_gating_delta_rule_update_cpu( A_log=layer.A_log, dt_bias=layer.dt_bias, q=query, k=key, v=value, - a=decode_a.contiguous(), - b=decode_b.contiguous(), + a=decode_a, + b=decode_b, initial_state_source=ssm_state, initial_state_indices=decode_state_indices, cu_seqlens=query_start_loc[: num_decodes + 1], @@ -568,17 +588,29 @@ def _spec_aware_nonspec( prefill_has_initial_state = has_initial_state[ num_decodes : num_decodes + num_prefills ] - # ``causal_conv1d_torch`` only touches columns [:width-1] of the buffer. - prefill_mixed_qkv = causal_conv1d_torch( - x=prefill_mixed_qkv.transpose(0, 1), - weight=conv_weights, - bias=layer.conv1d.bias, - conv_states=conv_buf, - query_start_loc=prefill_query_start_loc, - cache_indices=prefill_state_indices, - has_initial_state=prefill_has_initial_state, - activation=layer.activation, - ).transpose(0, 1) + if is_amx: + prefill_mixed_qkv = ops.causal_conv1d_fwd_cpu( + x=prefill_mixed_qkv.transpose(0, 1), + weight=layer.conv1d.weight, + bias=layer.conv1d.bias, + conv_states=conv_buf, + query_start_loc=prefill_query_start_loc, + cache_indices=prefill_state_indices, + has_initial_state=prefill_has_initial_state, + silu_activation=layer.activation == "silu", + is_vnni=True, + ).transpose(0, 1) + else: + prefill_mixed_qkv = causal_conv1d_torch( + x=prefill_mixed_qkv.transpose(0, 1), + weight=conv_weights, + bias=layer.conv1d.bias, + conv_states=conv_buf, + query_start_loc=prefill_query_start_loc, + cache_indices=prefill_state_indices, + has_initial_state=prefill_has_initial_state, + activation=layer.activation, + ).transpose(0, 1) query, key, value = layer.rearrange_mixed_qkv(prefill_mixed_qkv) g, beta = ops.fused_gdn_gating_cpu( @@ -618,24 +650,41 @@ def _spec_aware_nonspec_subset( Returns outputs ordered like ``non_spec_token_indx``. """ - out = core_attn_out_like(layer, mixed_qkv) has_initial_state = attn_metadata_i.has_initial_state prefill_state_indices = attn_metadata_i.prefill_state_indices prefill_qsl = attn_metadata_i.prefill_query_start_loc assert prefill_state_indices is not None and prefill_qsl is not None assert has_initial_state is not None + prefill_state_indices = prefill_state_indices.contiguous() - conv_weights = _unpacked_conv_weight(layer) - conv_out = causal_conv1d_torch( - x=mixed_qkv.transpose(0, 1), - weight=conv_weights, - bias=layer.conv1d.bias, - conv_states=conv_buf, - query_start_loc=prefill_qsl, - cache_indices=prefill_state_indices, - has_initial_state=has_initial_state, - activation=layer.activation, - ).transpose(0, 1) + is_amx = torch.cpu._is_amx_tile_supported() + if is_amx and is_conv_state_dim_first(): + raise RuntimeError("AMX GDN attention requires `SD` conv_state layout.") + + if is_amx: + conv_out = ops.causal_conv1d_fwd_cpu( + x=mixed_qkv.transpose(0, 1), + weight=layer.conv1d.weight, + bias=layer.conv1d.bias, + conv_states=conv_buf, + query_start_loc=prefill_qsl, + cache_indices=prefill_state_indices, + has_initial_state=has_initial_state, + silu_activation=layer.activation == "silu", + is_vnni=True, + ).transpose(0, 1) + else: + conv_weights = _unpacked_conv_weight(layer) + conv_out = causal_conv1d_torch( + x=mixed_qkv.transpose(0, 1), + weight=conv_weights, + bias=layer.conv1d.bias, + conv_states=conv_buf, + query_start_loc=prefill_qsl, + cache_indices=prefill_state_indices, + has_initial_state=has_initial_state, + activation=layer.activation, + ).transpose(0, 1) query, key, value = layer.rearrange_mixed_qkv(conv_out) g, beta = ops.fused_gdn_gating_cpu( @@ -658,8 +707,7 @@ def _spec_aware_nonspec_subset( ssm_state[prefill_state_indices] = last_recurrent_state.to( ssm_state.dtype, copy=False ) - out[:] = attn_out.squeeze(0) - return out + return attn_out.squeeze(0) def cpu_gdn_attention_core_fake( From d74285661064541f12be14e49a16defa27de67c6 Mon Sep 17 00:00:00 2001 From: Dao007forever <daole@inferact.ai> Date: Sun, 26 Jul 2026 23:22:17 -0700 Subject: [PATCH 092/185] [3/N][Core][KV Connector] Support reliable partial-tail KV offload for sub-block prompts (#49502) Signed-off-by: Dao Le <Dao007forever@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../test_partial_prefix_cache_hits.py | 293 ++++++++++++++++++ .../unit/test_mooncake_store_coordinator.py | 55 ++++ .../unit/test_mooncake_store_hma_e2e.py | 247 ++++++++++++++- .../unit/test_mooncake_store_scheduler.py | 232 +++++++++++++- .../unit/test_mooncake_store_worker.py | 122 ++++++++ .../v1/mooncake/store/coordinator.py | 66 +++- .../kv_connector/v1/mooncake/store/data.py | 35 ++- .../v1/mooncake/store/scheduler.py | 55 +++- .../kv_connector/v1/mooncake/store/worker.py | 209 ++++++++++++- vllm/v1/core/kv_cache_manager.py | 63 +++- vllm/v1/core/sched/output.py | 6 + vllm/v1/core/sched/scheduler.py | 48 ++- vllm/v1/core/single_type_kv_cache_manager.py | 66 +++- 13 files changed, 1448 insertions(+), 49 deletions(-) diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index a42fc415011..70457d5af9e 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -223,6 +223,299 @@ def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue(): assert moved[0].block_hash_num_tokens == 6 +def test_take_partial_tail_offloads_returns_cow_target(): + """The connector offload hand-off exposes the mamba CoW *target* block Y + (the durable boundary state), not the overwritten source X, and only at + the CoW step.""" + hash_block_size = 2 + block_size = 2 * hash_block_size + kv_cache_config = KVCacheConfig( + num_blocks=24, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ) + + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None + + # Step A registered the partial tail but has not CoW'd yet: no offload. + assert manager.take_partial_tail_offloads() == {} + + partial_mamba_hash = req0.block_hashes[6 // hash_block_size - 1] + source_block = manager.block_pool.get_cached_block( + partial_mamba_hash, kv_cache_group_ids=[1] + ) + assert source_block is not None + source_block_id = source_block[0].block_id + + # Step B: the producer continues, triggering the CoW X->Y. + req0.num_computed_tokens = 6 + req0.append_output_token_ids([3]) + assert manager.allocate_slots(req0, 1) is not None + + offloads = manager.take_partial_tail_offloads() + assert list(offloads.keys()) == ["0"] + assert len(offloads["0"]) == 1 + group_id, block_id, boundary_tokens = offloads["0"][0] + assert group_id == 1 # the mamba group + assert boundary_tokens == 6 + copies, _ = manager.take_kv_cache_block_copies() + cow_copy = next(c for c in copies if c.src_block_id == source_block_id) + # The offload points at the durable CoW target Y, not the overwritten X. + assert block_id == cow_copy.dst_block_id + assert block_id != source_block_id + # Draining clears it. + assert manager.take_partial_tail_offloads() == {} + + # The hand-off pinned Y (its CoW retention is released after this step, + # and Y is off the request block table); freeing the request unpins it. + cow_block = manager.block_pool.blocks[block_id] + pinned_ref = cow_block.ref_cnt + assert pinned_ref >= 1 + manager.free(req0) + assert cow_block.ref_cnt == pinned_ref - 1 + + +def test_partial_tail_pin_survives_released_cow_retention(): + """If the CoW retention is released before the hand-off is drained + (immediate-free mode), the drain must rescue the cow block from the free + queue: a raw ref increment would leave a ref>0 block allocatable, and the + next allocation would pop it and assert.""" + hash_block_size = 2 + block_size = 2 * hash_block_size + kv_cache_config = KVCacheConfig( + num_blocks=24, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ) + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None + req0.num_computed_tokens = 6 + req0.append_output_token_ids([3]) + assert manager.allocate_slots(req0, 1) is not None + + # Retention released before the drain (defer_block_free=False ordering). + _copies, retained = manager.take_kv_cache_block_copies() + manager.block_pool.free_blocks(retained) + + offloads = manager.take_partial_tail_offloads() + ((_group_id, block_id, boundary_tokens),) = offloads["0"] + assert boundary_tokens == 6 + cow_block = manager.block_pool.blocks[block_id] + assert cow_block.ref_cnt == 1 + + # The pinned block is out of the free queue: draining every free block + # neither trips the allocator's ref_cnt assert nor hands it out. + new_blocks = manager.block_pool.get_new_blocks( + manager.block_pool.get_num_free_blocks() + ) + assert block_id not in {b.block_id for b in new_blocks} + + +def test_partial_tail_offload_dropped_when_request_freed_before_drain(): + """A hand-off recorded in the same scheduling pass as the request's death + must not be drained: its release hook has already run, so draining would + leak a pinned block.""" + hash_block_size = 2 + block_size = 2 * hash_block_size + kv_cache_config = KVCacheConfig( + num_blocks=24, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ) + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None + req0.num_computed_tokens = 6 + req0.append_output_token_ids([3]) + assert manager.allocate_slots(req0, 1) is not None + + # The request dies (preempt/abort) before the scheduler drains. + manager.block_pool.free_blocks(manager.pop_blocks_for_free(req0)) + assert manager.take_partial_tail_offloads() == {} + + +def test_take_partial_tail_offloads_empty_without_partial_tail(): + """A prompt ending on a block boundary registers no partial tail, so there + is nothing to offload.""" + hash_block_size = 2 + block_size = 2 * hash_block_size + kv_cache_config = KVCacheConfig( + num_blocks=24, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ) + + # 4-token prompt ends exactly on the mamba block boundary (block_size=4). + req0 = make_request("0", [0, 0, 1, 1], hash_block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) + assert manager.allocate_slots(req0, 4, num_computed, computed_blocks) is not None + assert manager.take_partial_tail_offloads() == {} + + req0.num_computed_tokens = 4 + req0.append_output_token_ids([2]) + assert manager.allocate_slots(req0, 1) is not None + assert manager.take_partial_tail_offloads() == {} + + +def test_truncate_computed_blocks_preserves_sparse_prefix_positions(): + """truncate_computed_blocks slices each group by its own block size, + keeps null placeholders in the retained prefix, and leaves the original + lookup result untouched (pure view, no refcount changes).""" + hash_block_size = 2 + kv_cache_config = KVCacheConfig( + num_blocks=24, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=2 * hash_block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ) + producer = make_request("producer", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + blocks, num_computed, _ = manager.get_computed_blocks(producer) + assert manager.allocate_slots(producer, 6, num_computed, blocks) is not None + manager.free(producer) + manager.new_step_starts() + + consumer = make_request( + "consumer", [0, 0, 1, 1, 2, 2, 3, 3], hash_block_size, sha256 + ) + blocks, num_computed, _ = manager.get_computed_blocks(consumer) + assert num_computed == 6 + assert [len(group) for group in blocks.blocks] == [3, 2] + assert blocks.blocks[1][0].is_null + + truncated = manager.truncate_computed_blocks(blocks, 4) + + assert [len(group) for group in truncated.blocks] == [2, 1] + assert truncated.blocks[1][0].is_null + assert [len(group) for group in blocks.blocks] == [3, 2] + + def test_hybrid_mamba_partial_tail_owner_continue_preserves_later_hit(): hash_block_size = 2 block_size = 2 * hash_block_size diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py index 6e003798c7a..6b56fd2f191 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py @@ -3,6 +3,8 @@ from math import lcm +import torch + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( # noqa: E501 ExternalCachedBlockPool, MooncakeStoreCoordinator, @@ -14,10 +16,20 @@ from vllm.v1.core.kv_cache_utils import BlockHash from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheGroupSpec, + MambaSpec, SlidingWindowSpec, ) +def _mamba_align(block_size=32): + return MambaSpec( + block_size=block_size, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + + def _make_coord(groups, hash_block_size, use_eagle=False, retention_interval=None): """Construct a coordinator using the natural LCM of group block sizes as the scheduler block size — mirrors ``resolve_kv_cache_block_sizes`` for @@ -196,6 +208,49 @@ def test_coordinator_group_block_size_double_hash(): assert hit % 32 == 0 +# ----- Fine-grained partial hits (full attention + mamba "align") ----- + + +def test_coordinator_fine_grained_partial_tail_hit(): + """K3 shape: FA + mamba-align, block_size=32 over hash_block_size=16. When + both groups have the sub-block boundary hash, the reconciled hit lands on + the hash boundary (48), not the block boundary (32).""" + groups = [ + KVCacheGroupSpec(["L0"], _full(32)), + KVCacheGroupSpec(["L1"], _mamba_align(32)), + ] + coord = _make_coord(groups, hash_block_size=16) + assert coord.enable_partial_hash_hits + hs = _hashes(4) # 4 hash units of 16 = 64 tokens; block 0 = [0,32), etc. + # Both groups: full block 0 (key = last sub-hash hs[1]) + partial boundary + # at token 48 (key = hs[2]). No hs[3] -> block 1 is not full. + exists = {(g, bytes(h)) for g in (0, 1) for h in (hs[1], hs[2])} + cmap = ExternalCachedBlockPool(16, exists) + _masks, hit = coord.find_longest_cache_hit( + hs, max_length=64, cached_block_pool=cmap + ) + assert hit == 48 + + +def test_coordinator_fine_grained_clips_when_one_group_missing_tail(): + """If only one group has the sub-block boundary, min-convergence clips the + reconciled hit back to the block boundary (32).""" + groups = [ + KVCacheGroupSpec(["L0"], _full(32)), + KVCacheGroupSpec(["L1"], _mamba_align(32)), + ] + coord = _make_coord(groups, hash_block_size=16) + hs = _hashes(4) + # Full block 0 for both; partial boundary hs[2] only for FA (group 0). + exists = {(g, bytes(hs[1])) for g in (0, 1)} + exists |= {(0, bytes(hs[2]))} + cmap = ExternalCachedBlockPool(16, exists) + _masks, hit = coord.find_longest_cache_hit( + hs, max_length=64, cached_block_pool=cmap + ) + assert hit == 32 + + # ----- store_mask ----- diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index 2ea776adf6f..6c8af7fb73c 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py @@ -32,6 +32,7 @@ from vllm.v1.kv_cache_interface import ( KVCacheConfig, KVCacheGroupSpec, KVCacheTensor, + MambaSpec, SlidingWindowSpec, ) @@ -333,14 +334,14 @@ def test_recv_skips_swa_blocks_before_window(): def test_chunked_token_database_hash_block_size_smaller_than_block_size(): """DSv4-style: hash_block_size=4, group block_size=16 — process_tokens - keys each 16-token chunk by its last fine hash, keeping the Mooncake key - at one digest instead of concatenating all 4 fine hashes.""" + keys each chunk by its ending fine hash, including a partial tail.""" md = KeyMetadata("m", 0, 0, 0, 0, group_id=3) db = ChunkedTokenDatabase(md, block_size=16, hash_block_size=4) db.set_kv_caches_base_addr([0]) db.set_block_len([512]) - # 8 fine-grained hashes (32 tokens at hash_block_size=4) → 2 group chunks. fine_hashes = [BlockHash(bytes([i + 1]) * 4) for i in range(8)] + + # 8 fine-grained hashes (32 tokens at hash_block_size=4) → 2 group chunks. out = list(db.process_tokens(token_len=32, block_hashes=fine_hashes)) assert len(out) == 2 assert out[0][0] == 0 and out[0][1] == 16 @@ -349,3 +350,243 @@ def test_chunked_token_database_hash_block_size_smaller_than_block_size(): # prior three. assert out[0][2].hex() == fine_hashes[3].hex() assert out[1][2].hex() == fine_hashes[7].hex() + + # Sub-block hit: emit the partial chunk under its ending fine hash. + out = list(db.process_tokens(token_len=12, block_hashes=fine_hashes[:3])) + assert [(s, e) for s, e, _ in out] == [(0, 12)] + assert out[0][2].hex() == fine_hashes[2].hex() + + # Cross-block hit: emit both the full chunk and its partial tail. + out = list(db.process_tokens(token_len=28, block_hashes=fine_hashes[:7])) + assert [(s, e) for s, e, _ in out] == [(0, 16), (16, 28)] + assert out[0][2].hex() == fine_hashes[3].hex() + assert out[1][2].hex() == fine_hashes[6].hex() + + +def test_sub_block_partial_tail_offload_reads_cow_block(): + """Sub-block prompt (the 900/128/1536 shape, scaled to 12/4/16): the + partial tail is offloaded for both groups under the boundary sub-hash. The + full-attention block is read from the request block table; the mamba block + is the core-provided CoW target, not block_ids.""" + full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) + mamba = MambaSpec( + block_size=16, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + groups = [ + KVCacheGroupSpec(["L0"], full), + KVCacheGroupSpec(["L1"], mamba), + ] + coord = MooncakeStoreCoordinator(groups, scheduler_block_size=16, hash_block_size=4) + assert coord.enable_partial_hash_hits + + class _RecordingStore(_DictStore): + def __init__(self): + super().__init__() + self.puts: dict[str, list[int]] = {} + + def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k): + for key, addr in zip(keys, addrs): + self.puts[key] = addr + return super().batch_put_from_multi_buffers(keys, addrs, sizes, *a, **k) + + store = _RecordingStore() + token_dbs = [] + for g_idx in range(2): + db = ChunkedTokenDatabase( + KeyMetadata("m", 0, 0, 0, 0, group_id=g_idx), + block_size=16, + hash_block_size=4, + ) + db.set_kv_caches_base_addr([g_idx * 10_000]) + db.set_block_len([512]) + token_dbs.append(db) + + send = KVCacheStoreSendingThread( + store=store, + coord=coord, + token_databases=token_dbs, + block_size=16, + tp_rank=0, + group_put_steps=[1, 1], + kv_role="kv_both", + ready_event=threading.Event(), + replicate_config=MagicMock(), + ) + + # The surrounding metadata may describe a longer resumed replay, but the + # handoff identifies the exact state boundary to persist. + hs = [BlockHash(bytes([i + 1]) * 4) for i in range(5)] + mamba_cow_block = 7 + req = ReqMeta( + req_id="r0", + token_len_chunk=0, + block_ids=([1], [2]), + block_hashes=hs, + can_save=True, + num_prompt_tokens=20, + partial_tail_offloads=[(1, mamba_cow_block, 12)], + ) + + send._maybe_offload_partial_tail(req) + + # boundary = 12 // 4 * 4 = 12 -> keyed by hs[12 // 4 - 1] = hs[2]. + partial_hash = hs[2] + fa_key = token_dbs[0].key_for(partial_hash) + mamba_key = token_dbs[1].key_for(partial_hash) + assert set(store.puts) == {fa_key, mamba_key} + # FA reads block_ids[0][0] = block 1: addr = base(0) + 1 * 512. + assert store.puts[fa_key] == [512] + # Mamba reads the CoW block 7, not block_ids[1][0]=2. + assert store.puts[mamba_key] == [10_000 + mamba_cow_block * 512] + + +def test_offload_syncs_event_before_put(): + """An offload-carrying meta synchronizes its CoW-fence event before the + store put reads the blocks, then completes in one pass and drains the + completion counter.""" + full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) + mamba = MambaSpec( + block_size=16, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + groups = [ + KVCacheGroupSpec(["L0"], full), + KVCacheGroupSpec(["L1"], mamba), + ] + coord = MooncakeStoreCoordinator(groups, scheduler_block_size=16, hash_block_size=4) + event = MagicMock() + + class _FencedStore(_DictStore): + def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k): + assert event.synchronize.called, "put must run after the event sync" + return super().batch_put_from_multi_buffers(keys, addrs, sizes, *a, **k) + + store = _FencedStore() + token_dbs = [] + for g_idx in range(2): + db = ChunkedTokenDatabase( + KeyMetadata("m", 0, 0, 0, 0, group_id=g_idx), + block_size=16, + hash_block_size=4, + ) + db.set_kv_caches_base_addr([g_idx * 10_000]) + db.set_block_len([512]) + token_dbs.append(db) + + send = KVCacheStoreSendingThread( + store=store, + coord=coord, + token_databases=token_dbs, + block_size=16, + tp_rank=0, + group_put_steps=[1, 1], + kv_role="kv_both", + ready_event=threading.Event(), + replicate_config=MagicMock(), + ) + + hs = [BlockHash(bytes([i + 1]) * 4) for i in range(3)] + req = ReqMeta( + req_id="r1", + token_len_chunk=0, + block_ids=([1], [2]), + block_hashes=hs, + can_save=True, + num_prompt_tokens=12, + partial_tail_offloads=[(1, 7, 12)], + ) + req.current_event = event + send.add_stored_request("r1") + + send.request_queue.put(req) + send._handle_request(send.request_queue.get()) + assert send.request_queue.qsize() == 0 + assert store._data + assert send.stored_requests["r1"] == 0 + event.synchronize.assert_called_once() + + +def test_sub_block_partial_tail_offload_covers_smaller_group_blocks(): + """The K3-shaped 900/128/1536 scenario scaled to 12/4/16, with a + full-attention group whose block (4) is smaller than the lcm (16): the + offload must persist every FA block up to the boundary — the normal save + floors to the lcm, so those blocks are otherwise never written and the + consumer's per-group lookup would miss. The mamba boundary block still + reads the core-provided CoW target.""" + full = FullAttentionSpec(block_size=4, num_kv_heads=8, head_size=64, dtype=None) + mamba = MambaSpec( + block_size=16, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + groups = [ + KVCacheGroupSpec(["L0"], full), + KVCacheGroupSpec(["L1"], mamba), + ] + coord = MooncakeStoreCoordinator(groups, scheduler_block_size=16, hash_block_size=4) + assert coord.enable_partial_hash_hits + + class _RecordingStore(_DictStore): + def __init__(self): + super().__init__() + self.puts: dict[str, list[int]] = {} + + def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k): + for key, addr in zip(keys, addrs): + self.puts[key] = addr + return super().batch_put_from_multi_buffers(keys, addrs, sizes, *a, **k) + + store = _RecordingStore() + token_dbs = [] + for g_idx, block_size in enumerate([4, 16]): + db = ChunkedTokenDatabase( + KeyMetadata("m", 0, 0, 0, 0, group_id=g_idx), + block_size=block_size, + hash_block_size=4, + ) + db.set_kv_caches_base_addr([g_idx * 10_000]) + db.set_block_len([512]) + token_dbs.append(db) + + send = KVCacheStoreSendingThread( + store=store, + coord=coord, + token_databases=token_dbs, + block_size=16, + tp_rank=0, + group_put_steps=[1, 1], + kv_role="kv_both", + ready_event=threading.Event(), + replicate_config=MagicMock(), + ) + + hs = [BlockHash(bytes([i + 1]) * 4) for i in range(3)] # 3 hash units = 12 tok + mamba_cow_block = 7 + req = ReqMeta( + req_id="r2", + token_len_chunk=0, + block_ids=([1, 2, 3], [4]), + block_hashes=hs, + can_save=True, + num_prompt_tokens=12, + partial_tail_offloads=[(1, mamba_cow_block, 12)], + ) + + send._maybe_offload_partial_tail(req) + + # FA (block 4): full blocks ending at 4, 8 and 12, keyed by their normal + # block-end hashes; mamba (block 16): the partial boundary block under + # the boundary sub-hash, read from the CoW target. + expected = { + token_dbs[0].key_for(hs[0]): [1 * 512], + token_dbs[0].key_for(hs[1]): [2 * 512], + token_dbs[0].key_for(hs[2]): [3 * 512], + token_dbs[1].key_for(hs[2]): [10_000 + mamba_cow_block * 512], + } + assert store.puts == expected diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py index e7a0cbc1a7f..2709b92e52e 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py @@ -13,11 +13,15 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler impor ) -def _make_bare_scheduler() -> MooncakeStoreScheduler: +def _make_bare_scheduler( + *, hash_block_size: int = 16, enable_partial_hash_hits: bool = False +) -> MooncakeStoreScheduler: scheduler = object.__new__(MooncakeStoreScheduler) scheduler.kv_role = "kv_both" scheduler.lookup_async = False scheduler._block_size = 16 + scheduler._hash_block_size = hash_block_size + scheduler.enable_partial_hash_hits = enable_partial_hash_hits scheduler.load_specs = {} scheduler._unfinished_request_ids = {"req-0"} scheduler._unfinished_requests = {} @@ -136,6 +140,7 @@ def test_preemption_resets_tracker_before_request_finished(): block_hashes=[b"h0", b"h1"], prefill_end_tokens=48, ) + scheduler._request_trackers["req-0"].has_pending_offload = True scheduler.build_connector_meta(_make_preemption_scheduler_output()) @@ -144,6 +149,7 @@ def test_preemption_resets_tracker_before_request_finished(): assert tracker.allocated_block_ids == () assert tracker.num_saved_tokens == 0 assert tracker.token_ids is None + assert tracker.has_pending_offload is False assert tracker.prefill_end_tokens == 0 request = SimpleNamespace(request_id="req-0") assert scheduler.request_finished(request, ([0, 1],)) == (False, None) @@ -536,3 +542,227 @@ def test_full_external_hit_with_full_local_hit_skips_load(): assert need_to_allocate == 0 assert load_async is False assert "req-0" not in scheduler.load_specs + + +def test_partial_hash_hit_block_aligned_local_loads_partial_tail(): + # Fine-grained on (hash=4, block=16): a block-aligned local hit can pull a + # sub-block remote hit (24 = a hash boundary inside block 1). Loads [16, 24). + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) + scheduler.load_async = True + scheduler.client = _StubLookupClient(hit_tokens=24) + + request = SimpleNamespace( + request_id="req-0", + num_tokens=32, + block_hashes=[b"h0", b"h1", b"h2", b"h3", b"h4", b"h5", b"h6", b"h7"], + ) + + need_to_allocate, load_async = scheduler.get_num_new_matched_tokens( + request, num_computed_tokens=16 + ) + + assert need_to_allocate == 8 + assert load_async is True + load_spec = scheduler.load_specs["req-0"] + assert load_spec.vllm_cached_tokens == 16 + assert load_spec.kvpool_cached_tokens == 24 + + +def test_partial_hash_hit_no_remote_gain_skips_load(): + # Core always presents a block-aligned local hit (it floors a sub-block + # tail before calling the connector). When the remote hit does not exceed + # that block-aligned local hit, nothing is loaded. + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) + scheduler.load_async = True + scheduler.client = _StubLookupClient(hit_tokens=16) + + request = SimpleNamespace( + request_id="req-0", + num_tokens=32, + block_hashes=[b"h0", b"h1", b"h2", b"h3", b"h4", b"h5", b"h6", b"h7"], + ) + + need_to_allocate, load_async = scheduler.get_num_new_matched_tokens( + request, num_computed_tokens=16 + ) + + assert need_to_allocate == 0 + assert load_async is False + assert "req-0" not in scheduler.load_specs + + +def test_sub_block_prompt_looks_up_with_fine_grained(): + # A prompt smaller than one block (12 < block 16). With fine-grained partial + # hits the sub-block prefix is worth looking up (floor is the hash unit 4, + # not a full block), so a remote partial hit is loaded. Pre-change the + # block-size floor returned (0, False) for such prompts. + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) + scheduler.load_async = True + scheduler.client = _StubLookupClient(hit_tokens=8) + + request = SimpleNamespace( + request_id="req-0", + num_tokens=12, + block_hashes=[b"h0", b"h1", b"h2"], + ) + + need_to_allocate, load_async = scheduler.get_num_new_matched_tokens( + request, num_computed_tokens=0 + ) + + assert need_to_allocate == 8 + assert load_async is True + assert scheduler.load_specs["req-0"].kvpool_cached_tokens == 8 + + +def test_sub_block_prompt_not_looked_up_without_fine_grained(): + # Without fine-grained partial hits, sub-block prompts still skip the lookup + # (there is no full block, and no sub-block key granularity). + scheduler = _make_bare_scheduler() + scheduler.client = _StubLookupClient(hit_tokens=8) + + request = SimpleNamespace( + request_id="req-0", + num_tokens=12, + block_hashes=[b"h0", b"h1", b"h2"], + ) + + need_to_allocate, load_async = scheduler.get_num_new_matched_tokens( + request, num_computed_tokens=0 + ) + + assert need_to_allocate == 0 + assert load_async is False + assert "req-0" not in scheduler.load_specs + + +def test_pending_partial_tail_emits_offload_only_reqmeta(): + # A sub-block prompt never produces a block-aligned save, so the partial- + # tail offload arriving this step is emitted as an offload-only ReqMeta + # (can_save=True so it takes the normal enqueue path, token_len_chunk=0 so + # the worker skips the normal save). Pending-offload state delays the free + # without advancing the normal-save watermark before the put succeeds. + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) + request = SimpleNamespace( + all_token_ids=list(range(12)), + block_hashes=[b"h0", b"h1", b"h2"], + num_output_placeholders=0, + num_prompt_tokens=12, + ) + scheduler._unfinished_requests["req-0"] = (request, ([0],)) + scheduler._request_trackers["req-0"] = RequestTracker( + req_id="req-0", + token_len=12, + allocated_block_ids=([0],), + num_saved_tokens=0, + token_ids=list(range(12)), + prefill_end_tokens=12, + ) + + out = SimpleNamespace( + finished_req_ids=set(), + preempted_req_ids=set(), + scheduled_new_reqs=[], + scheduled_cached_reqs=SimpleNamespace( + req_ids=[], + new_block_ids=[], + num_computed_tokens=[], + resumed_req_ids=set(), + ), + num_scheduled_tokens={}, + scheduled_spec_decode_tokens={}, + partial_tail_offloads={"req-0": [(1, 7, 12)]}, + ) + + meta = scheduler.build_connector_meta(out) + + assert len(meta.requests) == 1 + req_meta = meta.requests[0] + assert req_meta.req_id == "req-0" + assert req_meta.can_save is True + assert req_meta.token_len_chunk == 0 + assert req_meta.partial_tail_offloads == [(1, 7, 12)] + assert req_meta.num_prompt_tokens == 12 + assert req_meta.block_ids == ([0],) + tracker = scheduler._request_trackers["req-0"] + assert tracker.num_saved_tokens == 0 + assert tracker.has_pending_offload is True + request = SimpleNamespace(request_id="req-0") + assert scheduler.request_finished(request, ([0],)) == (True, None) + + +def test_resumed_partial_tail_uses_handoff_boundary(): + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) + request = SimpleNamespace( + all_token_ids=list(range(20)), + block_hashes=[b"h0", b"h1", b"h2", b"h3", b"h4"], + num_output_placeholders=0, + num_prompt_tokens=12, + ) + scheduler._unfinished_requests["req-0"] = (request, ([0, 1],)) + scheduler._request_trackers["req-0"] = RequestTracker( + req_id="req-0", + token_len=20, + allocated_block_ids=([0, 1],), + num_saved_tokens=0, + token_ids=list(range(20)), + # Resumption replays prompt + previously generated tokens. + prefill_end_tokens=20, + ) + + out = SimpleNamespace( + finished_req_ids=set(), + preempted_req_ids=set(), + scheduled_new_reqs=[], + scheduled_cached_reqs=SimpleNamespace( + req_ids=[], + new_block_ids=[], + num_computed_tokens=[], + resumed_req_ids=set(), + ), + num_scheduled_tokens={}, + scheduled_spec_decode_tokens={}, + partial_tail_offloads={"req-0": [(1, 7, 12)]}, + ) + + meta = scheduler.build_connector_meta(out) + + assert len(meta.requests) == 1 + assert meta.requests[0].partial_tail_offloads == [(1, 7, 12)] + # Ordinary metadata retains the full resumed prefill range. + assert meta.requests[0].num_prompt_tokens == 20 + tracker = scheduler._request_trackers["req-0"] + assert tracker.num_saved_tokens == 0 + assert tracker.has_pending_offload is True + + +def test_resumed_partial_tail_attached_to_save_keeps_handoff_boundary(): + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) + request = SimpleNamespace( + all_token_ids=list(range(48)), + block_hashes=[b"h0", b"h1", b"h2"], + num_output_placeholders=0, + num_prompt_tokens=36, + ) + scheduler._unfinished_requests["req-0"] = (request, ([0, 1],)) + scheduler._request_trackers["req-0"] = RequestTracker( + req_id="req-0", + token_len=44, + allocated_block_ids=([0, 1],), + num_saved_tokens=32, + token_ids=list(range(44)), + prefill_end_tokens=48, + ) + out = _make_scheduler_output(scheduled_spec_tokens=None) + out.partial_tail_offloads = {"req-0": [(0, 7, 36)]} + + meta = scheduler.build_connector_meta(out) + + assert len(meta.requests) == 1 + assert meta.requests[0].can_save is True + assert meta.requests[0].partial_tail_offloads == [(0, 7, 36)] + assert meta.requests[0].num_prompt_tokens == 48 + # Ordinary saving still covers the full resumed prefill range. + tracker = scheduler._request_trackers["req-0"] + assert tracker.num_saved_tokens == 48 + assert tracker.has_pending_offload is True diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index b9876b9827b..54553cd49b9 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -599,6 +599,88 @@ def test_store_sending_thread_retries_skipped_range_after_pressure(): assert store.batch_put_from_multi_buffers.call_args.args[0] == keys +def _make_partial_tail_send_thread(store): + coord = SimpleNamespace( + enable_partial_hash_hits=True, + hash_block_size=4, + lcm_block_size=16, + ) + db = ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0), + block_size=4, + hash_block_size=4, + ) + db.set_kv_caches_base_addr([0x1000]) + db.set_block_len([256]) + return _make_store_sending_thread( + store, + coord=coord, + token_databases=[db], + ) + + +def _make_partial_tail_req(block_ids: list[int]) -> ReqMeta: + return ReqMeta( + req_id="req-a", + token_len_chunk=0, + block_ids=(block_ids,), + block_hashes=[b"a0", b"a1", b"a2"], + can_save=True, + partial_tail_offloads=[(1, 7, 12)], + ) + + +def test_partial_tail_offload_skips_null_source_blocks(): + store = MagicMock() + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) + store.batch_put_from_multi_buffers.return_value = [256, 256] + thread = _make_partial_tail_send_thread(store) + + assert thread._maybe_offload_partial_tail(_make_partial_tail_req([0, 2, 3])) + + keys, addrs, _sizes, _replicate_config = ( + store.batch_put_from_multi_buffers.call_args.args + ) + assert keys == [ + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131", + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6132", + ] + assert addrs == [[0x1000 + 2 * 256], [0x1000 + 3 * 256]] + + +def test_partial_tail_offload_honors_active_pressure_gate(): + store = MagicMock() + thread = _make_partial_tail_send_thread(store) + thread._store_pressure_active = True + thread._skip_store_requests.add("req-a") + thread.add_stored_request("req-a") + + thread._handle_request(_make_partial_tail_req([1, 2, 3])) + + store.batch_is_exist.assert_not_called() + store.batch_put_from_multi_buffers.assert_not_called() + assert thread.stored_requests["req-a"] == 0 + + +def test_partial_tail_put_failure_activates_pressure_gate(): + store = MagicMock() + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) + store.batch_put_from_multi_buffers.return_value = [256, -200, 256] + thread = _make_partial_tail_send_thread(store) + thread.add_stored_request("req-a") + + thread._handle_request(_make_partial_tail_req([1, 2, 3])) + + assert thread._store_pressure_active is True + assert thread._skip_store_requests == {"req-a"} + assert thread._saved_offset.get("req-a", 0) == 0 + assert thread.stored_requests["req-a"] == 0 + + thread.add_stored_request("req-a") + thread._handle_request(_make_partial_tail_req([1, 2, 3])) + assert store.batch_put_from_multi_buffers.call_count == 1 + + def test_store_sending_thread_delta_start_rank_saves_second_local_chunk(): store = MagicMock() store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) @@ -1939,6 +2021,46 @@ def test_lookup_partial_prefix_returns_first_hit_length(): assert worker.lookup(48, [b"a0", b"a1", b"a2"]) == 32 +def test_lookup_partial_tail_uses_hash_alignment(): + """A stored sub-block tail can serve a request extending past it.""" + from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheGroupSpec, + MambaSpec, + ) + + worker = _make_bare_worker(block_size=16) + full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) + mamba = MambaSpec( + block_size=16, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + worker._kv_cache_groups = [ + KVCacheGroupSpec(["full"], full), + KVCacheGroupSpec(["mamba"], mamba), + ] + worker.hash_block_size = 4 + worker.token_dbs = [ + ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=group_id), + block_size=16, + hash_block_size=4, + ) + for group_id in range(2) + ] + worker.coord = mooncake_store_worker.MooncakeStoreCoordinator( + worker._kv_cache_groups, + scheduler_block_size=16, + hash_block_size=4, + ) + _refresh_group_tp_replication_factors(worker) + worker.store.batch_is_exist.return_value = [0, 0, 1, 0, 0, 1] + + assert worker.lookup(13, [b"h0", b"h1", b"h2"]) == 12 + + def test_lookup_full_hit_reuses_existing_boundary(): """A full hit is re-derived below the request end without another RPC.""" worker = _make_bare_worker(block_size=16) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py index db34079ef88..c70ddf95a2f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py @@ -21,6 +21,7 @@ from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheGroupSpec, KVCacheSpec, + MambaSpec, UniformTypeKVCacheSpecs, ) from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry @@ -83,6 +84,9 @@ class MooncakeStoreCoordinator: self.kv_cache_groups = kv_cache_groups self.hash_block_size = hash_block_size self.lcm_block_size = scheduler_block_size + self.enable_partial_hash_hits = partial_hash_hits_enabled( + kv_cache_groups, hash_block_size + ) self.use_eagle = use_eagle # Mirror vLLM core's KVCacheCoordinator.retention_interval. self.retention_interval = retention_interval @@ -94,7 +98,12 @@ class MooncakeStoreCoordinator: self._verify_and_split_kv_cache_groups() def align_lookup_length(self, length: int) -> int: - return length // self.lcm_block_size * self.lcm_block_size + alignment = ( + self.hash_block_size + if self.enable_partial_hash_hits + else self.lcm_block_size + ) + return length // alignment * alignment def _verify_and_split_kv_cache_groups(self) -> None: """Mirrors KVCacheCoordinator.verify_and_split_kv_cache_groups but @@ -226,9 +235,14 @@ class MooncakeStoreCoordinator: retention_interval: int | None, num_prompt_tokens: int | None, ) -> tuple[list[bool] | None, ...]: - assert aligned_token_len % self.lcm_block_size == 0, ( + mask_alignment = ( + self.hash_block_size + if self.enable_partial_hash_hits + else self.lcm_block_size + ) + assert aligned_token_len % mask_alignment == 0, ( f"aligned_token_len ({aligned_token_len}) must be a multiple of " - f"lcm_block_size ({self.lcm_block_size})" + f"{mask_alignment}" ) masks: list[list[bool] | None] = [] for g_idx, g in enumerate(self.kv_cache_groups): @@ -278,17 +292,21 @@ class MooncakeStoreCoordinator: one already removed by the lookup. """ eagle_indices = self.eagle_attn_group_indices if apply_eagle else set() + alignment_tokens = ( + self.hash_block_size + if self.enable_partial_hash_hits + else self.lcm_block_size + ) if len(self.attention_groups) == 1: spec, group_ids, manager_cls = self.attention_groups[0] - hashes = self.block_hashes_for_spec(block_hashes, spec) hit_blocks, hit_length = manager_cls.find_longest_cache_hit( - block_hashes=hashes, # type: ignore[arg-type] + block_hashes=block_hashes, # type: ignore[arg-type] max_length=max_length, kv_cache_group_ids=group_ids, block_pool=cast(BlockPool, cached_block_pool), kv_cache_spec=spec, drop_eagle_block=(0 in eagle_indices), - alignment_tokens=spec.block_size, + alignment_tokens=alignment_tokens, ) num_groups = len(self.kv_cache_groups) blocks_by_group: list[list[KVCacheBlock]] = [[] for _ in range(num_groups)] @@ -320,17 +338,23 @@ class MooncakeStoreCoordinator: drop_eagle_block = idx in eagle_indices and idx not in eagle_verified _max_length = curr_hit_length - if drop_eagle_block: - _max_length = min(curr_hit_length + spec.block_size, max_length) - hashes = self.block_hashes_for_spec(block_hashes, spec) + if drop_eagle_block and not isinstance(spec, MambaSpec): + eagle_margin = ( + self.hash_block_size + if self.enable_partial_hash_hits + and manager_cls.supports_fine_grained_hash_lookup + and spec.block_size > self.hash_block_size + else spec.block_size + ) + _max_length = min(curr_hit_length + eagle_margin, max_length) hit_blocks, _new_hit_length = manager_cls.find_longest_cache_hit( - block_hashes=hashes, # type: ignore[arg-type] + block_hashes=block_hashes, # type: ignore[arg-type] max_length=_max_length, kv_cache_group_ids=group_ids, block_pool=cast(BlockPool, cached_block_pool), kv_cache_spec=spec, drop_eagle_block=drop_eagle_block, - alignment_tokens=self.lcm_block_size, + alignment_tokens=alignment_tokens, ) if drop_eagle_block: eagle_verified.add(idx) @@ -348,10 +372,11 @@ class MooncakeStoreCoordinator: break # Truncate full-attention hit_blocks to final converged length; - # other specs already trim themselves inside their hit logic. + # other specs already trim themselves inside their hit logic. cdiv keeps + # the partial tail block when hit_length is not block-aligned. spec0, group_ids0, _ = self.attention_groups[0] if isinstance(spec0, FullAttentionSpec): - num_blocks = hit_length // spec0.block_size + num_blocks = cdiv(hit_length, spec0.block_size) for gid in group_ids0: full_blks = hit_blocks_by_group[gid] assert full_blks is not None @@ -368,3 +393,18 @@ def _unwrap_spec(spec: KVCacheSpec) -> KVCacheSpec: if isinstance(spec, UniformTypeKVCacheSpecs): return next(iter(spec.kv_cache_specs.values())) return spec + + +def partial_hash_hits_enabled( + kv_cache_groups: list[KVCacheGroupSpec], hash_block_size: int +) -> bool: + """Mirror of core's ``HybridKVCacheCoordinator.enable_partial_hash_hits`` + (its dcp == 1 clause holds: the connector rejects hybrid + DCP/PCP > 1). + Single copy on purpose — scheduler and coordinator must not disagree. + """ + return any( + isinstance(spec := _unwrap_spec(g.kv_cache_spec), MambaSpec) + and spec.mamba_cache_mode == "align" + and spec.block_size > hash_block_size + for g in kv_cache_groups + ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py index 57e6bd8de0b..6daa1e82ea6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py @@ -228,16 +228,28 @@ class ChunkedTokenDatabase: n = len(chunks) starts = np.fromiter((c[0] for c in chunks), dtype=np.int64, count=n) spans = np.fromiter((c[1] for c in chunks), dtype=np.int64, count=n) - starts - assert not (spans % self.block_size).any() + assert not (spans % self.hash_block_size).any() bids = np.fromiter( (block_ids[i] for i in (starts // self.block_size).tolist()), dtype=np.int64, count=n, ) addrs = base[None, :] + bids[:, None] * blen[None, :] - sizes = blen[None, :] * (spans // self.block_size)[:, None] + block_counts = (spans + self.block_size - 1) // self.block_size + sizes = blen[None, :] * block_counts[:, None] return addrs.tolist(), sizes.tolist(), bids.tolist() + def prepare_value_for_block(self, block_id: int) -> tuple[list[int], list[int]]: + """Return addresses and sizes for one physical block slot.""" + addr_list = [] + size_list = [] + length = len(self.block_len) + for index, base_addr in enumerate(self.kv_caches_base_addr): + addr = base_addr + block_id * self.block_len[index % length] + addr_list.append(addr) + size_list.append(self.block_len[index % length]) + return addr_list, size_list + def process_tokens( self, token_len: int, @@ -256,7 +268,8 @@ class ChunkedTokenDatabase: rank regardless of where the processed suffix begins. Args: - token_len: Total number of tokens. + token_len: Total number of tokens. Must be hash-block aligned and + covered by ``block_hashes`` when hashes are present. block_hashes: Block hashes computed at ``hash_block_size`` granularity. When ``block_size > hash_block_size`` each group's ``block_size`` chunk is keyed by its last sub-hash via ``chunk_hashes_for_block_size``. @@ -269,11 +282,10 @@ class ChunkedTokenDatabase: assert put_step > 0 if not block_hashes: return - chunk_hashes: Sequence[BlockHash] = chunk_hashes_for_block_size( - block_hashes, self.hash_block_size, self.block_size - ) + assert token_len % self.hash_block_size == 0 + assert token_len // self.hash_block_size <= len(block_hashes) start_chunk = max(0, cdiv(mask_num, self.block_size)) - max_chunks = min(len(chunk_hashes), cdiv(token_len, self.block_size)) + max_chunks = cdiv(token_len, self.block_size) if chunk_mask is not None: max_chunks = min(max_chunks, start_chunk + len(chunk_mask)) for chunk_id in range(start_chunk, max_chunks): @@ -281,9 +293,9 @@ class ChunkedTokenDatabase: continue if chunk_id % put_step != put_step_rank: continue - h = chunk_hashes[chunk_id] start_idx = chunk_id * self.block_size end_idx = min(start_idx + self.block_size, token_len) + h = block_hashes[end_idx // self.hash_block_size - 1] yield start_idx, end_idx, h @@ -306,6 +318,7 @@ class RequestTracker: allocated_block_ids: tuple[list[int], ...] num_saved_tokens: int = 0 token_ids: list[int] | None = None + has_pending_offload: bool = False # Snapshot of the prefill range length at tracker creation time. # For a fresh request this is len(prompt). For a resumed-from-preemption # request it includes previously-generated tokens, which are re-prefilled. @@ -316,6 +329,7 @@ class RequestTracker: self.allocated_block_ids = () self.num_saved_tokens = 0 self.token_ids = None + self.has_pending_offload = False self.prefill_end_tokens = 0 def update( @@ -352,6 +366,11 @@ class ReqMeta: token_ids: list[int] | None = None num_prompt_tokens: int | None = None + # Core-provided per-mamba-group + # (group_id, cow_block_id, boundary_tokens) for this request's partial tail. + # Present only on the producer's CoW step; drives the connector's offload + # (the FA group's block is derived from block_ids and boundary_tokens). + partial_tail_offloads: list[tuple[int, int, int]] | None = None @staticmethod def from_request_tracker( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py index 2fef49a3075..a14d02927f6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py @@ -11,6 +11,9 @@ from vllm.config import VllmConfig from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorMetadata, ) +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( # noqa: E501 + partial_hash_hits_enabled, +) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501 LoadSpec, MooncakeStoreConnectorMetadata, @@ -63,6 +66,9 @@ class MooncakeStoreScheduler: self._block_size, self._hash_block_size = resolve_kv_cache_block_sizes( kv_cache_config, vllm_config ) + self.enable_partial_hash_hits = partial_hash_hits_enabled( + kv_cache_config.kv_cache_groups, self._hash_block_size + ) # Per-request state self.load_specs: dict[str, LoadSpec] = {} # to be loaded @@ -80,7 +86,12 @@ class MooncakeStoreScheduler: Returns ``(None, False)`` when an async lookup is still in flight, signaling the scheduler to retry this request on a later step. """ - if request.num_tokens < self._block_size: + # Fine-grained hits may land on a hash boundary inside a block; without + # partial hits, prefixes shorter than one physical block are skipped. + align = ( + self._hash_block_size if self.enable_partial_hash_hits else self._block_size + ) + if request.num_tokens < align: return 0, False num_external_hit_tokens = self.client.lookup( @@ -338,6 +349,44 @@ class MooncakeStoreScheduler: if req_meta is not None: meta.add_request(req_meta) + # Flush partial-tail offloads in the step they arrive: the CoW copy is + # enqueued before the connector event records, so this step's event + # fences the cow block. Ride the request's save meta when present, else + # emit an offload-only ReqMeta (token_len_chunk=0 skips the normal + # save; can_save=True takes the normal enqueue path). + step_partial_tails = getattr(scheduler_output, "partial_tail_offloads", None) + if step_partial_tails and not force_skip_save: + pending = dict(step_partial_tails) + for req_meta in meta.requests: + if req_meta.can_save: + groups = pending.pop(req_meta.req_id, None) + if groups: + req_meta.partial_tail_offloads = groups + tracker = self._request_trackers.get(req_meta.req_id) + if tracker is not None: + tracker.has_pending_offload = True + for req_id, groups in pending.items(): + tracker = self._request_trackers.get(req_id) + req_tuple = self._unfinished_requests.get(req_id) + if tracker is None or req_tuple is None: + # Request finished/preempted within this step; its blocks + # are going away, so the offload is conservatively dropped. + logger.debug("Dropping partial-tail offload for request %s", req_id) + continue + assert len({boundary for _, _, boundary in groups}) == 1 + tracker.has_pending_offload = True + meta.add_request( + ReqMeta( + req_id=req_id, + token_len_chunk=0, + block_ids=tracker.allocated_block_ids, + block_hashes=req_tuple[0].block_hashes, + can_save=True, + num_prompt_tokens=tracker.prefill_end_tokens, + partial_tail_offloads=groups, + ) + ) + return meta def request_finished( @@ -352,7 +401,9 @@ class MooncakeStoreScheduler: # Missing tracker can happen when the request is aborted before the # connector observes the normal finished lifecycle or is preempted # before finishing. - if tracker is None or tracker.num_saved_tokens <= 0: + if tracker is None or ( + tracker.num_saved_tokens <= 0 and not tracker.has_pending_offload + ): return False, None total_blocks = sum(len(g) for g in block_ids) delay_free_blocks = total_blocks > 0 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index cc49a3569d6..0fd85eb57d8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -61,6 +61,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.protocol import from vllm.logger import init_logger from vllm.utils.math_utils import cdiv from vllm.utils.network_utils import get_ip, make_zmq_socket +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.core.kv_cache_utils import ( BlockHash, maybe_convert_block_hash, @@ -531,6 +532,164 @@ class KVCacheStoreSendingThread(KVTransferThread): self._skip_store_requests.clear() return True + def _maybe_offload_partial_tail(self, req_meta: ReqMeta) -> bool: + """Offload the request's sub-block partial tail (its last prompt hash + boundary) so a later request can hit the sub-block prefix. + + Covers every block from the normal save's lcm floor to the boundary: + the normal save floors to ``lcm_block_size``, so a smaller-block + group's full blocks in that gap are never persisted elsewhere, and + the consumer's lookup needs every group at every probed boundary. + Full blocks are keyed by their block-end hash, the partial boundary + block by the boundary sub-hash; the mamba "align" boundary block is + the core-provided CoW block. All keys are deduped against the store. + + Returns: + True when no put is needed or every put succeeds, False otherwise. + """ + if not self.coord.enable_partial_hash_hits or not req_meta.block_hashes: + return True + partial_tail_offloads = req_meta.partial_tail_offloads + if not partial_tail_offloads: + return True + hash_block_size = self.coord.hash_block_size + boundaries = {boundary for _, _, boundary in partial_tail_offloads} + if len(boundaries) != 1: + raise ValueError( + "Partial-tail offloads for one request must share a boundary" + ) + boundary = boundaries.pop() + if boundary == 0: + return True + if boundary // hash_block_size - 1 >= len(req_meta.block_hashes): + return True + mamba_offloads = { + group_id: block_id for group_id, block_id, _ in partial_tail_offloads + } + + keys: list[str] = [] + addrs: list[list[int]] = [] + sizes: list[list[int]] = [] + saved = self._saved_offset.get(req_meta.req_id, 0) + for g_idx, db in enumerate(self.token_databases): + group_blocks = req_meta.block_ids[g_idx] + # Distribute across ranks by the same rule as normal chunks. + put_step = self.group_put_steps[g_idx] + put_step_rank = (self.tp_rank + g_idx) % put_step + # Always include the boundary block: its sub-hash key is written + # only here, even if normal saves already advanced past it. + last_block = cdiv(boundary, db.block_size) - 1 + for block_idx in range( + min(saved // db.block_size, last_block), last_block + 1 + ): + if block_idx % put_step != put_step_rank: + continue + valid_end = min((block_idx + 1) * db.block_size, boundary) + key_hash = req_meta.block_hashes[valid_end // hash_block_size - 1] + if ( + g_idx in mamba_offloads + and valid_end == boundary + and boundary % db.block_size != 0 + ): + block_id = mamba_offloads[g_idx] + else: + if block_idx >= len(group_blocks): + continue + block_id = group_blocks[block_idx] + if block_id == NULL_BLOCK_ID: + logger.debug( + "Skipping unavailable partial-tail source block " + "(req=%s, group=%d, block=%d)", + req_meta.req_id, + g_idx, + block_idx, + ) + continue + addr, size = db.prepare_value_for_block(block_id) + keys.append(db.key_for(key_hash)) + addrs.append(addr) + sizes.append(size) + + if not keys: + return True + exists_start = time.perf_counter() + try: + exists = self.store.batch_is_exist(keys) + except Exception as e: + self._record_operation( + "save_exists", + exists_start, + len(keys), + status="error", + num_failed_keys=len(keys), + ) + logger.error( + "Failed to check partial-tail keys for request %s: %s", + req_meta.req_id, + e, + ) + return False + self._record_operation("save_exists", exists_start, len(keys)) + missing = [i for i, e in enumerate(exists) if e != 1] + if not missing: + return True + keys = [keys[i] for i in missing] + addrs = [addrs[i] for i in missing] + sizes = [sizes[i] for i in missing] + if req_meta.current_event is not None: + # Fence the CoW block copy enqueued earlier this step. + req_meta.current_event.synchronize() + batch_bytes = _sum_batch_bytes(sizes) + put_start = time.perf_counter() + try: + res = self.store.batch_put_from_multi_buffers( + keys, addrs, sizes, self.replicate_config + ) + except Exception as e: + self._record_operation( + "save_put", + put_start, + len(keys), + num_bytes=batch_bytes, + status="error", + num_failed_keys=len(keys), + ) + logger.error( + "Failed to put partial-tail keys for request %s: %s", + req_meta.req_id, + e, + ) + return False + + failed = [i for i, value in enumerate(res) if value < 0] + self._record_operation( + "save_put", + put_start, + len(keys), + num_bytes=batch_bytes, + status="partial_failure" if failed else "ok", + num_failed_keys=len(failed), + ) + if failed: + failed_codes = {res[i] for i in failed} + logger.warning( + "Partial-tail put failed for request %s: %d/%d keys failed (codes=%s)", + req_meta.req_id, + len(failed), + len(keys), + failed_codes, + ) + if MOONCAKE_NO_AVAILABLE_HANDLE in failed_codes: + self._mark_request_skipped_for_pressure(req_meta.req_id) + return False + + if self._clear_store_pressure(): + logger.info( + "Mooncake CPU/disk offloading pressure cleared after a " + "successful partial-tail batch" + ) + return True + def _handle_request(self, req_meta: ReqMeta): # Cache hits are always a multiple of ``lcm_block_size`` tokens, which # is also ``store_mask``'s precondition. @@ -548,9 +707,6 @@ class KVCacheStoreSendingThread(KVTransferThread): # so the scheduler can release the GPU blocks it pinned for this # request (via `delay_free_blocks`) even when the store path raises. try: - if token_len == 0: - return - if self._should_skip_request(req_id): logger.debug( "Skipping Mooncake store for request %s while CPU/disk " @@ -559,6 +715,16 @@ class KVCacheStoreSendingThread(KVTransferThread): ) return + # Offload the sub-block partial tail (independent of the normal + # block-aligned save, which may be skipped this step). + if req_meta.partial_tail_offloads is not None and not ( + self._maybe_offload_partial_tail(req_meta) + ): + return + + if token_len == 0: + return + # Resume from where this rank left off; only the new suffix is saved. save_start = self._saved_offset.get(req_id, 0) @@ -1385,7 +1551,7 @@ class MooncakeStoreWorker: self.recv_request_queue.put(request) assert self.load_async, "load_async must be True for better performance." - # Issue stores with CUDA event synchronization + # Issue stores with CUDA event synchronization. if self.kv_role in ["kv_producer", "kv_both"]: current_event = None for request in meta.requests: @@ -1501,21 +1667,32 @@ class MooncakeStoreWorker: # candidate_meta stores the (group, hash_bytes) for key slice. candidate_keys: list[str] = [] candidate_meta: list[tuple[int, bytes]] = [] - lookup_masks = self.coord.lookup_mask(token_len) + fine_grained = self.coord.enable_partial_hash_hits + lookup_masks = None if fine_grained else self.coord.lookup_mask(token_len) for g_idx, db in enumerate(self.token_dbs): spec_block_size = db.block_size - lookup_mask = lookup_masks[g_idx] key_prefixes = self._lookup_key_prefixes[g_idx] - group_hashes = self.coord.block_hashes_for_spec( - block_hashes, self._kv_cache_groups[g_idx].kv_cache_spec - ) - max_chunks = min(len(group_hashes), cdiv(token_len, spec_block_size)) - mask_limit = ( - max_chunks if lookup_mask is None else min(max_chunks, len(lookup_mask)) - ) - for chunk_id in range(mask_limit): - if lookup_mask is not None and not lookup_mask[chunk_id]: - continue + if fine_grained: + max_units = min(len(block_hashes), token_len // self.hash_block_size) + unit_ids: range | list[int] = range(max_units) + group_hashes: Sequence[BlockHash] = block_hashes + else: + lookup_mask = lookup_masks[g_idx] # type: ignore[index] + group_hashes = self.coord.block_hashes_for_spec( + block_hashes, self._kv_cache_groups[g_idx].kv_cache_spec + ) + max_chunks = min(len(group_hashes), cdiv(token_len, spec_block_size)) + mask_limit = ( + max_chunks + if lookup_mask is None + else min(max_chunks, len(lookup_mask)) + ) + unit_ids = [ + chunk_id + for chunk_id in range(mask_limit) + if lookup_mask is None or lookup_mask[chunk_id] + ] + for chunk_id in unit_ids: h = group_hashes[chunk_id] hash_hex = h.hex() for key_prefix in key_prefixes: diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index e93a1ce6400..fcb67726a51 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -186,6 +186,10 @@ class KVCacheManager: tuple(() for _ in range(self.num_kv_cache_groups)) ) + # Off-table cow blocks handed to a KV connector for partial-tail + # offload; pinned until the request's blocks are freed. + self._partial_tail_pins: dict[str, list[KVCacheBlock]] = {} + @property def usage(self) -> float: """Get the KV cache usage. @@ -568,6 +572,9 @@ class KVCacheManager: Args: request: The request to free the blocks. """ + pins = self._partial_tail_pins.pop(request.request_id, None) + if pins: + self.block_pool.free_blocks(pins) self.coordinator.free(request.request_id) def remove_skipped_blocks( @@ -600,7 +607,14 @@ class KVCacheManager: Returns: The request's blocks in allocation order. """ - return self.coordinator.pop_blocks_for_free(request.request_id) + blocks = self.coordinator.pop_blocks_for_free(request.request_id) + # Pins ride the same (possibly deferred) free as the request blocks. + # Preemption may release a pin under a still-queued offload — the same + # exposure normal saves of table blocks already have. + pins = self._partial_tail_pins.pop(request.request_id, None) + if pins: + blocks = pins + blocks + return blocks def evict_blocks(self, block_ids: set[int]) -> None: """evict blocks from the prefix cache by their block IDs. @@ -760,6 +774,25 @@ class KVCacheManager: # Only create new KVCacheBlocks for non-empty blocks return KVCacheBlocks(blocks) if any(blocks) else self.empty_kv_cache_blocks + def truncate_computed_blocks( + self, blocks: KVCacheBlocks, num_computed_tokens: int + ) -> KVCacheBlocks: + """Return a lookup-result view truncated at an aligned token endpoint. + + Pure slicing: refcounts are untouched and ``blocks`` is not mutated. + """ + truncated: list[list[KVCacheBlock]] = [] + for group_blocks, manager in zip( + blocks.blocks, + self.coordinator.single_type_managers, + strict=True, + ): + assert num_computed_tokens % manager.block_size == 0 + num_blocks = num_computed_tokens // manager.block_size + assert num_blocks <= len(group_blocks) + truncated.append(list(group_blocks[:num_blocks])) + return self.create_kv_cache_blocks(tuple(truncated)) + def take_new_block_ids(self) -> list[int]: """Drain and return new attention block IDs for zeroing.""" ids: list[int] = [] @@ -812,6 +845,34 @@ class KVCacheManager: retained_blocks = [block for pair in pending_copies for block in pair] return copies, retained_blocks + def take_partial_tail_offloads(self) -> dict[str, list[tuple[int, int, int]]]: + """Drain producer partial-tail offload hand-offs per request. + + Returns ``{request_id: [(group_id, block_id, boundary_tokens), ...]}`` + for the durable boundary blocks of producers' last-prompt-boundary + partial tails. Only mamba "align" groups contribute; empty otherwise. + A KV connector reads the referenced blocks and offloads them so a later + request can hit the sub-block prefix. + + Each handed-off block lives off the request block table, so it is + pinned here and unpinned when the request's blocks are freed — for a + producer with saved tokens, after the connector reports sends done. + """ + offloads: dict[str, list[tuple[int, int, int]]] = {} + for mgr in self.coordinator.single_type_managers: + for ( + req_id, + group_id, + block, + boundary_tokens, + ) in mgr.take_pending_partial_tail_offloads(): + self.block_pool.touch((block,)) + self._partial_tail_pins.setdefault(req_id, []).append(block) + offloads.setdefault(req_id, []).append( + (group_id, block.block_id, boundary_tokens) + ) + return offloads + def new_step_starts(self) -> None: """Notify the coordinator that a new step is starting.""" self.coordinator.new_step_starts() diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index dab86c015e0..8782061356f 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -258,6 +258,12 @@ class SchedulerOutput: # CoW copies to apply after zeroing new blocks and before forward. kv_cache_block_copies: list[KVCacheBlockCopy] | None = None + # Producer partial-tail offload hand-off for external KV connectors: + # {request_id: [(group_id, block_id, boundary_tokens), ...]} pointing at + # the durable boundary block of a producer's last-prompt-boundary partial + # tail (mamba "align" CoW target). None unless partial hash hits are active. + partial_tail_offloads: dict[str, list[tuple[int, int, int]]] | None = None + # Dynamic speculative decoding: optimal K chosen by scheduler. # Number of spec tokens to schedule for the next step. num_spec_tokens_to_schedule: int = 0 diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 7536f6b0606..fff8c15e224 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -740,9 +740,16 @@ class Scheduler(SchedulerInterface): # Get externally-cached tokens if using a KVConnector. if self.connector is not None: + # Present a block-aligned local hit to the connector so + # a strictly longer remote hit can supersede a local + # sub-block tail without racing its copy-on-write. + partial_tail = num_new_local_computed_tokens % self.block_size + block_aligned_local = ( + num_new_local_computed_tokens - partial_tail + ) ext_tokens, load_kv_async = ( self.connector.get_num_new_matched_tokens( - request, num_new_local_computed_tokens + request, block_aligned_local ) ) @@ -754,7 +761,27 @@ class Scheduler(SchedulerInterface): step_skipped_waiting.prepend_request(request) continue - num_external_computed_tokens = ext_tokens + if partial_tail and ext_tokens > partial_tail: + # Remote strictly exceeds the full local hit: drop the + # sub-block tail so no CoW is needed, and let the load + # cover it. Trim the partial block out of the local + # computed blocks so it is not adopted from the cache. + new_computed_blocks = ( + self.kv_cache_manager.truncate_computed_blocks( + new_computed_blocks, block_aligned_local + ) + ) + num_new_local_computed_tokens = block_aligned_local + num_external_computed_tokens = ext_tokens + elif partial_tail: + # Remote does not exceed the full local hit: keep the + # local sub-block tail and load nothing external. + num_external_computed_tokens = 0 + # Nothing to load remotely -> not an async-load step; + # clearing avoids the `load_kv_async` assert below. + load_kv_async = False + else: + num_external_computed_tokens = ext_tokens if hit_diverged and num_external_computed_tokens == 0: # No external tokens back the deeper local hit, so its @@ -1108,6 +1135,22 @@ class Scheduler(SchedulerInterface): self.prev_step_scheduled_req_ids.clear() self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys()) + # Producer partial-tail hand-off for external KV connectors. Drained + # before the CoW retentions are released below, so the pin lands while + # the cow block still holds a retention ref. Without a producer-side + # connector nothing consumes the hand-off, so skip the drain (and its + # pin); the manager drops stale entries when the request's blocks are + # popped for free. + pending_partial_tail_offloads = None + if ( + self.connector is not None + and self.vllm_config.kv_transfer_config is not None + and self.vllm_config.kv_transfer_config.is_kv_producer + ): + pending_partial_tail_offloads = ( + self.kv_cache_manager.take_partial_tail_offloads() or None + ) + kv_cache_block_copies, cow_retained_blocks = ( self.kv_cache_manager.take_kv_cache_block_copies() ) @@ -1153,6 +1196,7 @@ class Scheduler(SchedulerInterface): free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(), new_block_ids_to_zero=self._get_new_block_ids_to_zero(), kv_cache_block_copies=pending_kv_cache_block_copies, + partial_tail_offloads=pending_partial_tail_offloads, num_spec_tokens_to_schedule=num_spec_tokens_to_schedule, ec_manager_metadata=self.encoder_cache_manager.get_manager_metadata(), ) diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index f8578c68a38..75cf97f4f46 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -115,6 +115,15 @@ class SingleTypeKVCacheManager(ABC): # managers (full attention, mamba "align"); harmlessly empty elsewhere. self._partial_hit_reqs: dict[str, tuple[int, KVCacheBlock]] = {} self._pending_cow_copies: list[tuple[KVCacheBlock, KVCacheBlock]] = [] + # Partial-tail offload hand-off for external KV connectors: when a + # producer registers its last-prompt-boundary partial tail and the + # durable boundary block is not on the append-only request block table + # (mamba "align" CoW target), record the request, group, block, and + # exact token boundary so a connector can offload it under the right + # hash. Populated only by mamba "align". + self._pending_partial_tail_offloads: list[ + tuple[str, int, KVCacheBlock, int] + ] = [] @classmethod def _get_num_evictable_blocks(cls, blocks: Sequence[KVCacheBlock]): @@ -378,6 +387,21 @@ class SingleTypeKVCacheManager(ABC): self._pending_cow_copies = [] return pending_copies + def take_pending_partial_tail_offloads( + self, + ) -> list[tuple[str, int, KVCacheBlock, int]]: + """Drain producer partial-tail hand-offs. + + Entries are ``(req_id, group_id, block, boundary_tokens)``. + + Only mamba "align" populates this. The block lives off the request + block table, so the caller must pin it until the connector has read + it — nothing else keeps it alive once the CoW retention is released. + """ + pending = self._pending_partial_tail_offloads + self._pending_partial_tail_offloads = [] + return pending + def _apply_cow( self, request_id: str, @@ -693,7 +717,8 @@ class FullAttentionManager(SingleTypeKVCacheManager): alignment_tokens < block_size and block_size % alignment_tokens == 0 ) if fine_grained: - assert isinstance(block_hashes, list) + # list or lazy BlobBlockHashes view + assert isinstance(block_hashes, Sequence) full_block_hashes: BlockHashList = BlockHashListWithBlockSize( block_hashes, alignment_tokens, block_size ) @@ -716,7 +741,8 @@ class FullAttentionManager(SingleTypeKVCacheManager): # Phase 2 (fine-grained only): extend into the first non-full block by # probing its interior hash boundaries high-to-low (longest hit first). if fine_grained: - assert isinstance(block_hashes, list) + # list or lazy BlobBlockHashes view + assert isinstance(block_hashes, Sequence) scale_factor = block_size // alignment_tokens first_partial_idx = len(computed_blocks[0]) * scale_factor max_partial_idx = min( @@ -1244,6 +1270,11 @@ class MambaManager(SingleTypeKVCacheManager): self.last_state_block_idx: dict[str, int] = {} # The set of the requests that have been allocated blocks self._allocated_block_reqs: set[str] = set() + # Requests that registered their own last-prompt-boundary partial + # tail (producers). On the next step's CoW the boundary state moves + # into a private cow_block; we record that block for connector + # offload (see _pending_partial_tail_offloads). + self._producer_partial_tail_reqs: dict[str, int] = {} @classmethod def find_longest_cache_hit( @@ -1277,7 +1308,8 @@ class MambaManager(SingleTypeKVCacheManager): block_size = kv_cache_spec.block_size if alignment_tokens < block_size and block_size % alignment_tokens == 0: - assert isinstance(block_hashes, list) + # list or lazy BlobBlockHashes view + assert isinstance(block_hashes, Sequence) hash_block_size = alignment_tokens scale_factor = block_size // hash_block_size max_num_partial_units = min( @@ -1590,6 +1622,21 @@ class MambaManager(SingleTypeKVCacheManager): self.block_pool.move_block_hashes(source_block, cow_block) self._pending_cow_copies.append((source_block, cow_block)) source_block.ref_cnt += 1 + boundary_tokens = self._producer_partial_tail_reqs.pop( + request_id, None + ) + if boundary_tokens is not None: + # This CoW preserved a producer's own boundary + # state in cow_block; hand it to the connector for + # partial-tail offload once the copy has run. + self._pending_partial_tail_offloads.append( + ( + request_id, + self.kv_cache_group_id, + cow_block, + boundary_tokens, + ) + ) if cow_block.block_hash is not None: # The moved entry is only filled by this step's # copy, so defer same-step hits on it. @@ -1607,6 +1654,14 @@ class MambaManager(SingleTypeKVCacheManager): if self.mamba_cache_mode == "align": self._allocated_block_reqs.discard(request_id) self.last_state_block_idx.pop(request_id, None) + self._producer_partial_tail_reqs.pop(request_id, None) + # A hand-off whose request died in this same scheduling pass must + # not reach the connector: its unpin hook (free) has already run. + self._pending_partial_tail_offloads = [ + entry + for entry in self._pending_partial_tail_offloads + if entry[0] != request_id + ] return super().pop_blocks_for_free(request_id) def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: @@ -1681,6 +1736,11 @@ class MambaManager(SingleTypeKVCacheManager): if partial_hash is not None: self._partial_hit_reqs[request.request_id] = (block_idx, source_block) self.num_cached_block[request.request_id] = block_idx + # Producer of this partial tail: the boundary state currently lives + # in ``source_block`` but the next step's forward overwrites it. The + # upcoming CoW copies it into a durable cow_block; record the req so + # allocate_new_blocks hands that block to the connector for offload. + self._producer_partial_tail_reqs[request.request_id] = num_tokens return partial_hash From 5d07e268b11ae0240d3af2582516861dfe52e86a Mon Sep 17 00:00:00 2001 From: Zhenzhong Xu <zhenzhong.xu@intel.com> Date: Mon, 27 Jul 2026 14:26:31 +0800 Subject: [PATCH 093/185] [Quantization][INC]Add MXFP8 Linear Support (#47514) Signed-off-by: Zhenzhong1 <zhenzhong.xu@intel.com> Signed-off-by: Zhenzhong Xu <zhenzhong.xu@intel.com> Co-authored-by: Yi Liu <yi4.liu@intel.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/quantization/test_auto_round.py | 141 +++++++++++++++++- .../layers/quantization/inc/inc.py | 70 ++++++++- .../quantization/inc/schemes/__init__.py | 2 + .../quantization/inc/schemes/factory.py | 2 + .../inc/schemes/inc_mxfp8_linear.py | 84 +++++++++++ .../inc/schemes/inc_mxfp8_scheme.py | 31 ++++ .../inc/schemes/inc_wna16_linear.py | 1 + 7 files changed, 322 insertions(+), 9 deletions(-) create mode 100644 vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_linear.py create mode 100644 vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_scheme.py diff --git a/tests/quantization/test_auto_round.py b/tests/quantization/test_auto_round.py index 6df7211549b..732080fc967 100644 --- a/tests/quantization/test_auto_round.py +++ b/tests/quantization/test_auto_round.py @@ -18,9 +18,13 @@ from vllm.model_executor.layers.quantization.inc import INCConfig from vllm.model_executor.layers.quantization.inc.config_parser import INCLayerConfig from vllm.model_executor.layers.quantization.inc.inc_linear import INCLinearMethod from vllm.model_executor.layers.quantization.inc.schemes import ( + INCMxfp8Scheme, INCWna16Scheme, resolve_scheme, ) +from vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp8_linear import ( + INCMxfp8LinearScheme, +) from vllm.model_executor.layers.quantization.inc.schemes.inc_scheme import ( INCLinearScheme, ) @@ -52,12 +56,19 @@ MODELS = [ pytest.param( "Intel/Qwen3-8B-w2g64-for-ut", marks=pytest.mark.skipif( - not (current_platform.is_cuda() or current_platform.is_xpu()) - or current_platform.device_count() < 2, - reason="72B INT2 AutoRound model requires XPU with at least 2 devices.", + not (current_platform.is_xpu()), + reason="INC int2 on XPU requires the ARK backend.", ), id="auto_round:auto_gptq_int2_tp2", ), + pytest.param( + "INC4AI/Qwen3-8B-MXFP8-AR", + marks=pytest.mark.skipif( + not (current_platform.is_cuda() or current_platform.is_xpu()), + reason="MXFP8 AutoRound model only supports CUDA/XPU backend for now.", + ), + id="auto_round:llm_compressor_mxfp8", + ), ] MODEL_RUNNER_KWARGS = { @@ -66,6 +77,11 @@ MODEL_RUNNER_KWARGS = { "gpu_memory_utilization": 0.8, "max_model_len": 512, }, + "INC4AI/Qwen3-8B-MXFP8-AR": { + "block_size": 64, + "gpu_memory_utilization": 0.8, + "max_model_len": 512, + }, } @@ -273,6 +289,39 @@ def test_inc_config_parser_fused_module_requires_consistent_configs() -> None: config.config_parser.resolve(DummyLayer(), "layers.0.self_attn.qkv_proj") +def test_inc_mxfp8() -> None: + config = make_config( + weight_bits=8, + group_size=32, + sym=True, + packing_format="auto_round:llm_compressor", + data_type="mx_fp", + ) + + assert config.weight_bits == 8 + assert config.group_size == 32 + assert config.data_type == "mx_fp" + assert config.packing_format == "auto_round:llm_compressor" + + +def test_inc_config_rejects_invalid_mxfp8_activation_config() -> None: + with pytest.raises(AssertionError, match="act_dynamic=True"): + INCConfig.from_config( + { + "bits": 8, + "group_size": 32, + "sym": True, + "packing_format": "auto_round:llm_compressor", + "data_type": "mx_fp", + "act_bits": 8, + "act_data_type": "mx_fp", + "act_group_size": 32, + "act_sym": True, + "act_dynamic": False, + } + ) + + def test_inc_layer_config_mx_fp_helpers() -> None: layer_config = INCLayerConfig( bits=4, @@ -304,6 +353,22 @@ def test_inc_resolve_scheme_selects_wna16() -> None: assert isinstance(scheme, INCWna16Scheme) +def test_inc_resolve_scheme_selects_mxfp8() -> None: + layer_config = INCLayerConfig( + bits=8, + group_size=32, + sym=True, + packing_format="auto_round:llm_compressor", + backend="auto", + data_type="mx_fp", + quantized=True, + ) + + scheme = resolve_scheme(layer_config) + + assert isinstance(scheme, INCMxfp8Scheme) + + class DummyLinearScheme(INCLinearScheme): def __init__(self) -> None: self.calls: list[tuple] = [] @@ -323,6 +388,76 @@ class DummyLinearScheme(INCLinearScheme): return "applied" +def test_inc_mxfp8_linear_scheme_delegates_to_kernel(monkeypatch) -> None: + class DummyKernel: + def __init__(self) -> None: + self.calls: list[tuple] = [] + + def process_weights_after_loading(self, layer) -> None: + self.calls.append(("process", layer)) + + def apply_weights(self, layer, x, bias=None): + self.calls.append(("apply", layer, x, bias)) + return "applied" + + kernel = DummyKernel() + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp8_linear.init_mxfp8_linear_kernel", + lambda: kernel, + ) + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp8_linear.ModelWeightParameter", + lambda **kwargs: torch.nn.Parameter(kwargs["data"], requires_grad=False), + ) + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp8_linear.GroupQuantScaleParameter", + lambda **kwargs: torch.nn.Parameter(kwargs["data"], requires_grad=False), + ) + + scheme = INCMxfp8LinearScheme() + layer = torch.nn.Module() + + scheme.create_weights( + layer=layer, + input_size_per_partition=64, + output_partition_sizes=[48, 16], + input_size=64, + output_size=64, + params_dtype=torch.bfloat16, + weight_loader=lambda *args, **kwargs: None, + ) + + assert layer.weight.shape == (64, 64) + assert layer.weight.dtype == torch.float8_e4m3fn + assert layer.weight_scale.shape == (64, 2) + assert layer.weight_scale.dtype == torch.uint8 + + scheme.process_weights_after_loading(layer) + result = scheme.apply_weights(layer, torch.randn(1, 64), None) + + assert result == "applied" + assert [call[0] for call in kernel.calls] == ["process", "apply"] + + +def test_inc_mxfp8_linear_scheme_requires_block_32_input(monkeypatch) -> None: + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp8_linear.init_mxfp8_linear_kernel", + lambda: object(), + ) + scheme = INCMxfp8LinearScheme() + + with pytest.raises(ValueError, match="divisible by 32"): + scheme.create_weights( + layer=torch.nn.Module(), + input_size_per_partition=48, + output_partition_sizes=[32], + input_size=48, + output_size=32, + params_dtype=torch.bfloat16, + weight_loader=lambda *args, **kwargs: None, + ) + + def test_inc_linear_method_delegates() -> None: scheme = DummyLinearScheme() method = INCLinearMethod(scheme) diff --git a/vllm/model_executor/layers/quantization/inc/inc.py b/vllm/model_executor/layers/quantization/inc/inc.py index 86fa7cefcfc..ad8389a8b90 100644 --- a/vllm/model_executor/layers/quantization/inc/inc.py +++ b/vllm/model_executor/layers/quantization/inc/inc.py @@ -35,8 +35,12 @@ class INCConfig(QuantizationConfig): """ SUPPORTED_BITS = {2, 3, 4, 8} - SUPPORTED_DTYPES = {"int"} - SUPPORTED_FORMATS = {"auto_round:auto_gptq", "auto_round:auto_awq"} + SUPPORTED_DTYPES = {"int", "mx_fp"} + SUPPORTED_FORMATS = { + "auto_round:auto_gptq", + "auto_round:auto_awq", + "auto_round:llm_compressor", + } SUPPORTED_BACKENDS = { "auto", "gptq", @@ -45,6 +49,11 @@ class INCConfig(QuantizationConfig): "awq:marlin", "marlin", } + MXFP8_BITS = 8 + MXFP8_GROUP_SIZE = 32 + MXFP8_DATA_TYPE = "mx_fp" + MXFP8_PACKING_FORMAT = "auto_round:llm_compressor" + MXFP8_SUPPORTED_ACT_DTYPES = {"mx_fp", "mx_fp_rceil"} def __init__( self, @@ -65,8 +74,8 @@ class INCConfig(QuantizationConfig): ) if data_type not in self.SUPPORTED_DTYPES: raise ValueError( - f"Unsupported data_type: {data_type}," - f" currently only support {self.SUPPORTED_DTYPES}." + f"Unsupported data_type: {data_type}, " + f"currently only support {self.SUPPORTED_DTYPES}." ) if packing_format not in self.SUPPORTED_FORMATS: raise ValueError( @@ -75,7 +84,7 @@ class INCConfig(QuantizationConfig): ) if backend not in self.SUPPORTED_BACKENDS: raise ValueError( - f"Unsupported backend: {backend}, " + f"Unsupported backend: {backend}, " f"currently only support {self.SUPPORTED_BACKENDS}." ) @@ -94,12 +103,59 @@ class INCConfig(QuantizationConfig): self.pack_factor = Fraction(32, weight_bits) self.config_parser = INCConfigParser(self) + self._validate_supported_quantization() + def __repr__(self) -> str: return ( f"INCConfig(weight_bits={self.weight_bits}, " f"group_size={self.group_size}, sym={self.sym})" ) + def _validate_supported_quantization(self) -> None: + if self.data_type == self.MXFP8_DATA_TYPE: + assert self.weight_bits == self.MXFP8_BITS, ( + f"INC MXFP8 only supports bits=8, but found bits={self.weight_bits}." + ) + assert self.group_size == self.MXFP8_GROUP_SIZE, ( + "INC MXFP8 only supports group_size=32, " + f"but found group_size={self.group_size}." + ) + assert self.sym, "INC MXFP8 only supports symmetric weights." + assert self.packing_format == self.MXFP8_PACKING_FORMAT, ( + "INC MXFP8 only supports " + f"packing_format={self.MXFP8_PACKING_FORMAT!r}, " + f"but found {self.packing_format!r}." + ) + assert self.backend == "auto", ( + "INC MXFP8 only supports backend='auto', " + f"but found backend={self.backend!r}." + ) + elif self.packing_format == self.MXFP8_PACKING_FORMAT: + raise ValueError( + f"packing_format={self.MXFP8_PACKING_FORMAT!r} requires " + f"data_type={self.MXFP8_DATA_TYPE!r}." + ) + + def _validate_raw_config(self, config: dict[str, Any]) -> None: + if self.data_type != self.MXFP8_DATA_TYPE: + return + + expected_fields = { + "act_bits": self.MXFP8_BITS, + "act_data_type": self.MXFP8_DATA_TYPE, + "act_group_size": self.MXFP8_GROUP_SIZE, + "act_sym": True, + "act_dynamic": True, + "enable_quanted_input": False, + } + for field_name, expected_value in expected_fields.items(): + actual_value = self.get_from_keys_or(config, [field_name], expected_value) + assert actual_value == expected_value, ( + "INC MXFP8 only supports " + f"{field_name}={expected_value!r}, " + f"but found {field_name}={actual_value!r}." + ) + @classmethod def get_name(cls) -> QuantizationMethods: return "inc" @@ -118,7 +174,7 @@ class INCConfig(QuantizationConfig): @classmethod def from_config(cls, config: dict[str, Any]) -> "INCConfig": - return cls( + quant_config = cls( weight_bits=cls.get_from_keys(config, ["bits"]), group_size=cls.get_from_keys(config, ["group_size"]), sym=cls.get_from_keys(config, ["sym"]), @@ -132,6 +188,8 @@ class INCConfig(QuantizationConfig): data_type=cls.get_from_keys_or(config, ["data_type"], "int"), backend=cls.get_from_keys_or(config, ["backend", "vllm_backend"], "auto"), ) + quant_config._validate_raw_config(config) + return quant_config def get_layer_config(self, layer, layer_name: str): return self.config_parser.get_layer_config(layer, layer_name) diff --git a/vllm/model_executor/layers/quantization/inc/schemes/__init__.py b/vllm/model_executor/layers/quantization/inc/schemes/__init__.py index ea6c0a00d86..e4dc5cb5dca 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/__init__.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/__init__.py @@ -2,12 +2,14 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from .factory import resolve_scheme +from .inc_mxfp8_scheme import INCMxfp8Scheme from .inc_scheme import INCLinearScheme, INCScheme from .inc_wna16_scheme import INCWna16Scheme __all__ = [ "INCScheme", "INCLinearScheme", + "INCMxfp8Scheme", "INCWna16Scheme", "resolve_scheme", ] diff --git a/vllm/model_executor/layers/quantization/inc/schemes/factory.py b/vllm/model_executor/layers/quantization/inc/schemes/factory.py index 4ae85ed9a83..30ce420b505 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/factory.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/factory.py @@ -9,9 +9,11 @@ if TYPE_CHECKING: def resolve_scheme(layer_config: "INCLayerConfig") -> "INCScheme": + from .inc_mxfp8_scheme import INCMxfp8Scheme from .inc_wna16_scheme import INCWna16Scheme scheme_list: list[type[INCScheme]] = [ + INCMxfp8Scheme, INCWna16Scheme, ] diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_linear.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_linear.py new file mode 100644 index 00000000000..775b5f5bba6 --- /dev/null +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_linear.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +from vllm.model_executor.kernels.linear import init_mxfp8_linear_kernel +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + MXFP8_BLOCK_SIZE, + MXFP8_SCALE_DTYPE, + MXFP8_VALUE_DTYPE, +) +from vllm.model_executor.parameter import ( + GroupQuantScaleParameter, + ModelWeightParameter, +) + +from .inc_scheme import INCLinearScheme + + +class INCMxfp8LinearScheme(INCLinearScheme): + def __init__(self) -> None: + self.kernel = init_mxfp8_linear_kernel() + + @classmethod + def get_min_capability(cls) -> int: + return 75 + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + del input_size, output_size + if input_size_per_partition % MXFP8_BLOCK_SIZE != 0: + raise ValueError( + "INC MXFP8 requires input_size_per_partition " + f"({input_size_per_partition}) to be divisible by " + f"{MXFP8_BLOCK_SIZE}." + ) + + output_size_per_partition = sum(output_partition_sizes) + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.params_dtype = params_dtype + + weight = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition, + dtype=MXFP8_VALUE_DTYPE, + ), + input_dim=1, + output_dim=0, + weight_loader=extra_weight_attrs.get("weight_loader"), + ) + layer.register_parameter("weight", weight) + + weight_scale = GroupQuantScaleParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition // MXFP8_BLOCK_SIZE, + dtype=MXFP8_SCALE_DTYPE, + ), + input_dim=1, + output_dim=0, + weight_loader=extra_weight_attrs.get("weight_loader"), + ) + layer.register_parameter("weight_scale", weight_scale) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + self.kernel.process_weights_after_loading(layer) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.kernel.apply_weights(layer, x, bias) diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_scheme.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_scheme.py new file mode 100644 index 00000000000..2e251a8a542 --- /dev/null +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_scheme.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import TYPE_CHECKING + +from ..inc_linear import INCLinearMethod +from .inc_scheme import INCScheme + +if TYPE_CHECKING: + import torch + + from ..config_parser import INCLayerConfig + from ..inc import INCConfig + + +class INCMxfp8Scheme(INCScheme): + @staticmethod + def can_handle(layer_config: "INCLayerConfig") -> bool: + return layer_config.is_mxfp8 + + def get_linear_method( + self, + config: "INCConfig", + layer: "torch.nn.Module", + prefix: str, + layer_config: "INCLayerConfig", + ): + del config, layer, prefix, layer_config + from .inc_mxfp8_linear import INCMxfp8LinearScheme + + return INCLinearMethod(INCMxfp8LinearScheme()) diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py index dd6c2fa2eaa..5c99fd98b54 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py @@ -166,6 +166,7 @@ class INCXPULinearBase(INCLinearScheme): def __init__(self, layer_config: "INCLayerConfig") -> None: self.weight_bits = layer_config.bits self.group_size = layer_config.group_size + self.sym = layer_config.sym self.pack_factor = 32 // self.weight_bits self.is_awq_packed = layer_config.is_awq From e09900436ca2c531e980827cea2e6eaad625ae00 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas <akaratza@amd.com> Date: Mon, 27 Jul 2026 01:44:21 -0500 Subject: [PATCH 094/185] [CI][ROCm] Reduce kernel test runtime (#49915) Signed-off-by: Andreas Karatzas <akaratza@amd.com> --- .buildkite/test-amd.yaml | 7 +- .../core/test_fused_quant_layernorm.py | 126 +++++++++++++----- tests/kernels/core/test_vit_fp8_scaling.py | 5 + 3 files changed, 101 insertions(+), 37 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index ebd5d79cce0..6e5d115f19b 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1561,11 +1561,12 @@ steps: commands: - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: Kernels Core Operation Test # TBD +- label: Kernels Core Operation Test %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 + parallelism: 3 working_dir: "/vllm-workspace/tests" source_file_dependencies: - csrc/ @@ -1576,7 +1577,7 @@ steps: - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py + - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: Kernels KDA Test # TBD timeout_in_minutes: 180 @@ -2145,7 +2146,7 @@ steps: mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 - parallelism: 4 + parallelism: 8 optional: true working_dir: "/vllm-workspace/" source_file_dependencies: diff --git a/tests/kernels/core/test_fused_quant_layernorm.py b/tests/kernels/core/test_fused_quant_layernorm.py index 255833c48dc..02bdc90bb61 100644 --- a/tests/kernels/core/test_fused_quant_layernorm.py +++ b/tests/kernels/core/test_fused_quant_layernorm.py @@ -31,7 +31,7 @@ NUM_TOKENS_HIDDEN_SIZES = [ ADD_RESIDUAL = [False, True] SCALE_UBS = [True, False] -GROUP_SIZES = [None, [1, 64], [1, 128]] +GROUP_SIZES = [[1, 64], [1, 128]] TMA_ALIGNMENTS = [0, 4] SEEDS = [0] CUDA_DEVICES = [ @@ -40,6 +40,86 @@ CUDA_DEVICES = [ EPS = 1e-6 + +def _is_valid_config( + hidden_size: int, + has_scale_ub: bool, + quant_dtype: torch.dtype, + group_size: list[int] | None, + tma_alignment: int, +) -> bool: + if group_size is not None and hidden_size % group_size[1] != 0: + return False + if group_size is not None and has_scale_ub: + return False + if ( + group_size is None or quant_dtype != current_platform.fp8_dtype() + ) and tma_alignment != 0: + return False + if ( + group_size is not None + and tma_alignment != 0 + and hidden_size // group_size[1] % tma_alignment == 0 + ): + return False + return not (has_scale_ub and quant_dtype != current_platform.fp8_dtype()) + + +def _config_id( + num_tokens: int, + hidden_size: int, + has_scale_ub: bool, + quant_dtype: torch.dtype, + group_size: list[int] | None, + tma_alignment: int, +) -> str: + quant = str(quant_dtype).removeprefix("torch.") + group = "per-token" if group_size is None else f"group-{group_size[1]}" + return ( + f"{num_tokens}x{hidden_size}-{quant}-{group}-" + f"tma-{tma_alignment}-scale-ub-{has_scale_ub}" + ) + + +# Filter unsupported combinations during collection. Letting each case call +# pytest.skip still runs the global per-test teardown and dominates this suite. +RMS_NORM_CONFIGS = [ + pytest.param( + num_tokens, + hidden_size, + has_scale_ub, + quant_dtype, + group_size, + tma_alignment, + id=_config_id( + num_tokens, + hidden_size, + has_scale_ub, + quant_dtype, + group_size, + tma_alignment, + ), + ) + for ( + (num_tokens, hidden_size), + has_scale_ub, + quant_dtype, + (group_size, tma_alignment), + ) in itertools.product( + NUM_TOKENS_HIDDEN_SIZES, + SCALE_UBS, + QUANT_DTYPES, + [(None, 0), *itertools.product(GROUP_SIZES, TMA_ALIGNMENTS)], + ) + if _is_valid_config( + hidden_size, + has_scale_ub, + quant_dtype, + group_size, + tma_alignment, + ) +] + ## Helpers @@ -154,15 +234,19 @@ def ops_impl( ) -@pytest.mark.parametrize("num_tokens, hidden_size", NUM_TOKENS_HIDDEN_SIZES) -@pytest.mark.parametrize("add_residual", ADD_RESIDUAL) -@pytest.mark.parametrize("has_scale_ub", SCALE_UBS) -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("quant_dtype", QUANT_DTYPES) @pytest.mark.parametrize( - "group_size, tma_alignment", - [(None, 0), *itertools.product(GROUP_SIZES, TMA_ALIGNMENTS)], + ( + "num_tokens", + "hidden_size", + "has_scale_ub", + "quant_dtype", + "group_size", + "tma_alignment", + ), + RMS_NORM_CONFIGS, ) +@pytest.mark.parametrize("add_residual", ADD_RESIDUAL) +@pytest.mark.parametrize("dtype", DTYPES) @pytest.mark.parametrize("seed", SEEDS) @pytest.mark.parametrize("device", CUDA_DEVICES) @pytest.mark.parametrize("strided_input", [False, True]) @@ -185,32 +269,6 @@ def test_rms_norm( torch.set_default_device(device) torch.accelerator.set_device_index(device) - if group_size is not None and hidden_size % group_size[1] != 0: - # skip - pytest.skip("Skip non-divisible group sizes") - - if group_size is not None and has_scale_ub: - # blockwise baseline doesn't support scale_ub - pytest.skip("scale_ub not supported for blockwise/group quantization") - - if ( - group_size is None or quant_dtype != current_platform.fp8_dtype() - ) and tma_alignment != 0: - # TMA alignment is only supported for groupwise fp8 kernels - pytest.skip("tma alignment not supported for per-token or int8 quantization") - - if ( - group_size is not None - and tma_alignment != 0 - and hidden_size // group_size[1] % tma_alignment == 0 - ): - # Skip tests where TMA alignment doesn't create extra padding to save time - pytest.skip("Skip TMA alignment cases where no extra padding is added") - - if has_scale_ub and quant_dtype != current_platform.fp8_dtype(): - # skip - pytest.skip("scale_ub only supported for fp8 quantization") - layer = RMSNorm(hidden_size, EPS).to(dtype=dtype) # Make weights diff --git a/tests/kernels/core/test_vit_fp8_scaling.py b/tests/kernels/core/test_vit_fp8_scaling.py index a197439237f..f3b2dbfa56a 100644 --- a/tests/kernels/core/test_vit_fp8_scaling.py +++ b/tests/kernels/core/test_vit_fp8_scaling.py @@ -18,6 +18,11 @@ from vllm.utils.flashinfer import ( is_flashinfer_cudnn_fp8_prefill_attn_supported, ) +pytestmark = pytest.mark.skipif( + not is_flashinfer_cudnn_fp8_prefill_attn_supported(), + reason="FlashInfer cuDNN FP8 prefill attention not supported", +) + LAYER_0 = "visual.blocks.0.attn.attn" LAYER_1 = "visual.blocks.1.attn.attn" NUM_HEADS = 16 From fd9d2ede6fa05313e4a7482c5edb88bf8c565230 Mon Sep 17 00:00:00 2001 From: Bugen Zhao <i@bugenzhao.com> Date: Mon, 27 Jul 2026 14:50:08 +0800 Subject: [PATCH 095/185] [Rust Frontend] Keep `--max-model-len` engine-owned (#49944) Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com> --- rust/src/cmd/src/cli.rs | 6 --- rust/src/cmd/src/cli/tests.rs | 63 ++++++++++++++++++++++++++---- rust/src/managed-engine/src/cli.rs | 11 ++++-- 3 files changed, 63 insertions(+), 17 deletions(-) diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index 454a397ad05..3ae827d2868 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -148,11 +148,6 @@ pub struct SharedRuntimeArgs { #[arg(long)] #[serde(default)] pub language_model_only: bool, - /// Override the maximum model context length. When set, the frontend uses - /// this value instead of the model's `max_position_embeddings` from - /// `config.json`. - #[arg(long)] - pub max_model_len: Option<u32>, /// Maximum number of log probabilities to return when `logprobs` is /// specified in sampling parameters. `-1` means no cap. #[arg(long, value_parser = clap::value_parser!(i32).range(-1..), allow_negative_numbers = true)] @@ -664,7 +659,6 @@ impl ServeArgs { self.managed_engine.clone().into_config( self.runtime.model.clone(), - self.runtime.max_model_len, self.runtime.max_logprobs, profiler_config, reasoning_parser.as_deref(), diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index d8ec9792b7e..b0926fb24ac 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -58,9 +58,6 @@ fn serve_args_forward_python_flags_with_separator() { reasoning_parser: Auto, renderer: Auto, language_model_only: false, - max_model_len: Some( - 512, - ), max_logprobs: None, grpc_port: None, shutdown_timeout: 0, @@ -102,6 +99,9 @@ fn serve_args_forward_python_flags_with_separator() { handshake_port: None, data_parallel_size: 1, data_parallel_size_local: None, + max_model_len: Some( + "512", + ), python_args: [ "--dtype", "float16", @@ -756,7 +756,6 @@ fn frontend_args_accept_json() { reasoning_parser: None, renderer: Auto, language_model_only: false, - max_model_len: None, max_logprobs: None, grpc_port: None, shutdown_timeout: 0, @@ -823,11 +822,32 @@ fn frontend_args_json_applies_defaults() { assert_eq!(args.runtime.tool_call_parser, ParserSelection::None); assert_eq!(args.runtime.reasoning_parser, ParserSelection::None); assert_eq!(args.runtime.renderer, RendererSelection::Auto); - assert_eq!(args.runtime.max_model_len, None); assert_eq!(args.runtime.max_logprobs, None); assert_eq!(args.runtime.shutdown_timeout, 0); } +#[test] +fn frontend_args_json_ignores_engine_owned_max_model_len() { + let cli = Cli::try_parse_from([ + "vllm-rs", + "frontend", + "--listen-fd", + "3", + "--input-address", + "ipc:///tmp/input.sock", + "--output-address", + "ipc:///tmp/output.sock", + "--args-json", + r#"{"model_tag":"Qwen/Qwen3-0.6B","max_model_len":-1}"#, + ]) + .unwrap(); + + let Command::Frontend(args) = cli.command else { + panic!("expected frontend args"); + }; + assert_eq!(args.runtime.model, "Qwen/Qwen3-0.6B"); +} + #[test] fn frontend_args_json_accepts_supported_non_default_fields() { let cli = Cli::try_parse_from([ @@ -840,7 +860,7 @@ fn frontend_args_json_accepts_supported_non_default_fields() { "--output-address", "ipc:///tmp/output.sock", "--args-json", - r#"{"model_tag":"Qwen/Qwen3-0.6B","engine_ready_timeout_secs":42,"tool_call_parser":"hermes","reasoning_parser":"qwen3_thinking","tokenizer_mode":"deepseek_v32","language_model_only":true,"max_model_len":8192,"max_logprobs":-1,"shutdown_timeout":3}"#, + r#"{"model_tag":"Qwen/Qwen3-0.6B","engine_ready_timeout_secs":42,"tool_call_parser":"hermes","reasoning_parser":"qwen3_thinking","tokenizer_mode":"deepseek_v32","language_model_only":true,"max_logprobs":-1,"shutdown_timeout":3}"#, ]) .unwrap(); @@ -858,11 +878,38 @@ fn frontend_args_json_accepts_supported_non_default_fields() { ); assert_eq!(args.runtime.renderer, RendererSelection::DeepSeekV32); assert!(args.runtime.language_model_only); - assert_eq!(args.runtime.max_model_len, Some(8192)); assert_eq!(args.runtime.max_logprobs, Some(-1)); assert_eq!(args.runtime.shutdown_timeout, 3); } +#[test] +fn serve_args_forward_auto_max_model_len_to_managed_engine() { + let cli = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--max-model-len", + "auto", + ]) + .unwrap(); + + let Command::Serve(args) = cli.command else { + panic!("expected serve args"); + }; + assert_eq!(args.managed_engine.max_model_len.as_deref(), Some("auto")); + + let config = args.to_managed_engine_config(5555); + expect![[r#" + [ + "--max-model-len", + "auto", + "--reasoning-parser", + "qwen3", + ] + "#]] + .assert_debug_eq(&config.python_args); +} + #[test] fn serve_args_accept_none_reasoning_parser() { let cli = Cli::try_parse_from([ @@ -1278,7 +1325,6 @@ fn serve_args_accept_handshake_aliases() { reasoning_parser: Auto, renderer: Auto, language_model_only: false, - max_model_len: None, max_logprobs: None, grpc_port: None, shutdown_timeout: 0, @@ -1322,6 +1368,7 @@ fn serve_args_accept_handshake_aliases() { ), data_parallel_size: 4, data_parallel_size_local: None, + max_model_len: None, python_args: [], }, }, diff --git a/rust/src/managed-engine/src/cli.rs b/rust/src/managed-engine/src/cli.rs index 95d61f652c1..26a6e56e79c 100644 --- a/rust/src/managed-engine/src/cli.rs +++ b/rust/src/managed-engine/src/cli.rs @@ -39,6 +39,12 @@ pub struct ManagedEngineArgs { /// Number of data parallel replicas to run on this node. #[arg(long)] pub data_parallel_size_local: Option<usize>, + /// Maximum model context length forwarded to the managed Python engine. + /// + /// Rust leaves validation to Python so values such as `auto` and + /// human-readable integers retain their engine-owned semantics. + #[arg(long)] + pub max_model_len: Option<String>, /// Additional arguments forwarded to `python -m vllm.entrypoints.cli.main /// serve ...`. @@ -78,7 +84,6 @@ impl ManagedEngineArgs { pub fn into_config( self, model: String, - max_model_len: Option<u32>, max_logprobs: Option<i32>, profiler_config: Option<String>, reasoning_parser: Option<&str>, @@ -89,9 +94,9 @@ impl ManagedEngineArgs { ) -> ManagedEngineConfig { let mut python_args = self.python_args; // Manually forward some args to the Python engine. - if let Some(max_model_len) = max_model_len { + if let Some(max_model_len) = self.max_model_len { python_args.push("--max-model-len".to_string()); - python_args.push(max_model_len.to_string()); + python_args.push(max_model_len); } if let Some(max_logprobs) = max_logprobs { python_args.push("--max-logprobs".to_string()); From 8061dc26bd0b37a2ad58068b99d072be5a0e5f47 Mon Sep 17 00:00:00 2001 From: Xiaochang Wu <xiaochang.wu@intel.com> Date: Mon, 27 Jul 2026 15:02:48 +0800 Subject: [PATCH 096/185] [Bugfix] Normalize sparse MLA warmup compression ratios (#49392) Signed-off-by: Wu, Xiaochang <xiaochang.wu@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../test_indexer_deepseek_v4_slot_mapping.py | 24 ++++++++++++++++++- vllm/v1/attention/backends/mla/indexer.py | 2 +- 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 159bb8af3fb..c6f13b84408 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -1,15 +1,37 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest import torch from tests.v1.attention.utils import create_vllm_config from vllm.v1.attention.backend import CommonAttentionMetadata -from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadataBuilder +from vllm.v1.attention.backends.mla.indexer import ( + BuildPrefillChunkMetadataKernel, + DeepseekV32IndexerMetadataBuilder, +) from vllm.v1.kv_cache_interface import MLAAttentionSpec +def test_indexer_warmup_normalizes_zero_compress_ratios(): + config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=8), + model_config=SimpleNamespace( + hf_config=SimpleNamespace(compress_ratios=[0, 0, 4, 128, 0]) + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=1, + cp_kv_cache_interleave_size=1, + ), + ) + + keys = BuildPrefillChunkMetadataKernel().get_warmup_keys(config) + + assert {key.COMPRESS_RATIO for key in keys} == {1, 4, 128} + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_indexer_builder_deepseek_v4_compressed_slot_mapping_uses_storage_block_size(): """Regression test: DeepseekV4 compression path must compute slot_mapping from diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index fe8bfef5308..24ab48ddd05 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -328,7 +328,7 @@ class BuildPrefillChunkMetadataKernel( dcp_interleave = parallel_config.cp_kv_cache_interleave_size dcp_rank = get_dcp_group().rank_in_group if dcp_world > 1 else 0 compress_ratios = tuple( - int(ratio) + max(1, int(ratio)) for ratio in (getattr(hf_config, "compress_ratios", None) or (1,)) ) return self._trace_dispatch(self.dispatch)( From afc94523c965a61bd6869287c798b2a45f2c205e Mon Sep 17 00:00:00 2001 From: liuzhenwei <zhenweiliu@habana.ai> Date: Mon, 27 Jul 2026 15:28:02 +0800 Subject: [PATCH 097/185] [XPU][CI] Use platform device in InputBatch V2 test (#49939) Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com> --- tests/v1/worker/test_gpu_input_batch_v2.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/v1/worker/test_gpu_input_batch_v2.py b/tests/v1/worker/test_gpu_input_batch_v2.py index 2c43cd72147..f517abb2b14 100644 --- a/tests/v1/worker/test_gpu_input_batch_v2.py +++ b/tests/v1/worker/test_gpu_input_batch_v2.py @@ -5,9 +5,10 @@ import pytest import torch +from vllm.platforms import current_platform from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers -DEVICE = "cuda" +DEVICE = current_platform.device_type @pytest.mark.parametrize( From cbc3a872005d595493b546acf46f94a9c9f68cbb Mon Sep 17 00:00:00 2001 From: Andreas Karatzas <akaratza@amd.com> Date: Mon, 27 Jul 2026 02:49:11 -0500 Subject: [PATCH 098/185] [Tokenizer] Use HF config for HF tokenizers (#49907) Signed-off-by: Andreas Karatzas <akaratza@amd.com> --- vllm/tokenizers/registry.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py index 82df4339e85..84f3ccce553 100644 --- a/vllm/tokenizers/registry.py +++ b/vllm/tokenizers/registry.py @@ -205,17 +205,27 @@ def get_tokenizer( **kwargs, ) + if tokenizer_cls == TokenizerLike: + tokenizer_cls_ = TokenizerRegistry.load_tokenizer_cls(tokenizer_mode) + else: + tokenizer_cls_ = tokenizer_cls + # Ensure that, if the config were to come from vllm.transformers_utils.config, it is # registered with AutoConfig before the tokenizer is loaded. This is necessary since # tokenizer_cls_.from_pretrained will call AutoConfig.from_pretrained internally. # This may fail for paths that don't have a model config (e.g. LoRA adapters), # which is fine — those don't need custom config registration. + # HF-backed tokenizers must receive the HF config. In a dual-format Mistral + # repository, auto detection intentionally prefers params.json, but passing + # that generic config to AutoTokenizer can select the wrong tokenizer class. + config_format = "hf" if tokenizer_cls_ is CachedHfTokenizer else "auto" config = None with contextlib.suppress(ValueError, OSError): config = get_config( tokenizer_name, trust_remote_code=trust_remote_code, revision=revision, + config_format=config_format, ) # Some models have an incorrect tokenizer_class on the hub. @@ -229,10 +239,6 @@ def get_tokenizer( model_type, ) tokenizer_cls_ = TokenizersBackend - elif tokenizer_cls == TokenizerLike: - tokenizer_cls_ = TokenizerRegistry.load_tokenizer_cls(tokenizer_mode) - else: - tokenizer_cls_ = tokenizer_cls if config is not None and tokenizer_cls_ is CachedHfTokenizer: # AutoTokenizer otherwise reloads config.json internally. Reuse the From eb290ab673c8f3ff87648cd8e4cd60f50fe7b301 Mon Sep 17 00:00:00 2001 From: "Li, Jiang" <jiang1.li@intel.com> Date: Mon, 27 Jul 2026 16:32:23 +0800 Subject: [PATCH 099/185] [Bugfix][CPU] Zero-pad MoE intermediate size for grouped-gemm TP alignment (#49591) Signed-off-by: jiang1.li <jiang1.li@intel.com> Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com> --- tests/kernels/moe/test_cpu_fused_moe.py | 130 ++++++++++++++++++ .../layers/fused_moe/cpu_fused_moe.py | 95 +++++++++++-- 2 files changed, 211 insertions(+), 14 deletions(-) diff --git a/tests/kernels/moe/test_cpu_fused_moe.py b/tests/kernels/moe/test_cpu_fused_moe.py index 75e34e6766a..9e99a274fb3 100644 --- a/tests/kernels/moe/test_cpu_fused_moe.py +++ b/tests/kernels/moe/test_cpu_fused_moe.py @@ -14,6 +14,7 @@ from vllm._custom_ops import ( from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.cpu_fused_moe import ( _CPU_MOE_ACT_FN, + CPUFusedMOE, ) from vllm.platforms import CpuArchEnum, current_platform from vllm.utils.torch_utils import set_random_seed @@ -346,3 +347,132 @@ def test_cpu_fused_moe_int8( ) torch.testing.assert_close(output, ref_output, atol=2e-2, rtol=2e-2) + + +# moe_intermediate_size not a multiple of 32, e.g. what tensor-parallel +# sharding of moe_intermediate_size=704 at tp=4 produces (704 // 4 == 176). +UNALIGNED_INTERMEDIATE_DIM = 176 + + +class _StubMoELayer(torch.nn.Module): + """Minimal stand-in for the real MoE layer module, exposing just what + CPUFusedMOE reads/replaces (w13_weight, w2_weight, activation, and + optionally w13_bias/w2_bias).""" + + def __init__( + self, + w13_weight: torch.Tensor, + w2_weight: torch.Tensor, + activation: MoEActivation, + w13_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + ): + super().__init__() + self.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False) + self.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False) + self.activation = activation + if w13_bias is not None: + self.w13_bias = torch.nn.Parameter(w13_bias, requires_grad=False) + if w2_bias is not None: + self.w2_bias = torch.nn.Parameter(w2_bias, requires_grad=False) + + +@pytest.mark.parametrize("expert_num", EXPERT_NUM) +@pytest.mark.parametrize("hidden_size", HIDDEN_DIM) +@pytest.mark.parametrize("use_bias", USE_BIAS) +@pytest.mark.parametrize("dtype", DTYPE) +@pytest.mark.parametrize( + "act", + [MoEActivation.SILU, MoEActivation.GELU, MoEActivation.GELU_TANH], +) +def test_cpu_fused_moe_unaligned_intermediate_size( + default_vllm_config, + expert_num: int, + hidden_size: int, + use_bias: bool, + dtype: torch.dtype, + act: MoEActivation, +): + """An unaligned per-partition moe_intermediate_size must still hit the + grouped-gemm fast path via automatic zero-padding, with numerically + correct output, instead of silently falling back to the much slower + per-expert torch loop.""" + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: + pytest.skip("padding is only applied on the x86 AMX/vector kernels") + + set_random_seed(0) + batch_size = 64 + intermediate_size = UNALIGNED_INTERMEDIATE_DIM + topk_num = max(expert_num // 2, 1) + up_dim = 2 * intermediate_size + + input = torch.randn((batch_size, hidden_size), dtype=dtype) / ( + 0.5 * hidden_size**0.5 + ) + w13 = torch.randn((expert_num, up_dim, hidden_size), dtype=dtype) / ( + 0.5 * hidden_size**0.5 + ) + w2 = torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype) / ( + 0.5 * intermediate_size**0.5 + ) + router_logits = torch.randn((batch_size, expert_num), dtype=dtype) + w13_bias = None + w2_bias = None + if use_bias: + w13_bias = torch.randn((expert_num, up_dim), dtype=dtype) / (0.5 * up_dim**0.5) + w2_bias = torch.randn((expert_num, hidden_size), dtype=dtype) / ( + 0.5 * hidden_size**0.5 + ) + score = torch.softmax(router_logits, dim=-1, dtype=torch.float32) + topk_weight, topk_ids = torch.topk(score, topk_num) + topk_ids = topk_ids.to(torch.int32) + + ref_output = ref_fused_moe( + input, w13, w2, w13_bias, w2_bias, topk_weight, topk_ids, act + ) + + layer = _StubMoELayer( + w13.clone(), + w2.clone(), + act, + w13_bias.clone() if w13_bias is not None else None, + w2_bias.clone() if w2_bias is not None else None, + ) + + cpu_moe = CPUFusedMOE(layer) + assert cpu_moe.forward_method == cpu_moe.forward_grouped_gemm, ( + "expected the padded intermediate size to hit the grouped-gemm fast path" + ) + + output = cpu_moe.forward_method( + layer, input, topk_weight, topk_ids, act, expert_num, False + ) + + atol, rtol = get_default_atol(output), get_default_rtol(output) + torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol) + + +def test_cpu_fused_moe_unaligned_intermediate_size_swigluoai(default_vllm_config): + """swigluoai's interleaved gate/up layout isn't padded automatically. On + AMX-capable CPUs this must raise rather than silently falling back to + the (correct but much slower) per-expert torch loop; elsewhere it should + fall back exactly as before.""" + set_random_seed(0) + expert_num = 8 + hidden_size = 128 + intermediate_size = UNALIGNED_INTERMEDIATE_DIM + dtype = torch.bfloat16 + up_dim = 2 * intermediate_size + + layer = _StubMoELayer( + torch.randn((expert_num, up_dim, hidden_size), dtype=dtype), + torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype), + MoEActivation.SWIGLUOAI, + ) + + if torch.cpu._is_amx_tile_supported(): + with pytest.raises(RuntimeError): + CPUFusedMOE(layer) + else: + cpu_moe = CPUFusedMOE(layer) + assert cpu_moe.forward_method == cpu_moe.forward_torch diff --git a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py index 1a0acff058c..a957aca9eca 100644 --- a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py @@ -20,6 +20,11 @@ from vllm.platforms import CpuArchEnum, current_platform from vllm.utils.torch_utils import direct_register_custom_op _CPU_MOE_LAYER_CACHE = {} +# The CPU grouped-gemm MoE kernels (AMX and vector) tile the expert +# intermediate ("N") dimension in blocks of this size and have no tail/ +# remainder handling, so a shard is only eligible for the fast path when +# its per-partition intermediate size is a multiple of it. +_MOE_GROUPED_GEMM_N_TILE = 32 def _swigluoai_forward_native( @@ -230,6 +235,7 @@ class CPUFusedMOE: """CPU-based fused MoE implementation.""" def __init__(self, layer: torch.nn.Module) -> None: + self._pad_moe_intermediate_for_grouped_gemm(layer) use_grouped_gemm, isa = self.check_grouped_gemm(layer) self.isa = isa if use_grouped_gemm: @@ -284,6 +290,69 @@ class CPUFusedMOE: apply_router_weight_on_input, ) + def _pad_moe_intermediate_for_grouped_gemm(self, layer: torch.nn.Module) -> None: + """Zero-pad the per-partition MoE intermediate dim up to a multiple + of _MOE_GROUPED_GEMM_N_TILE, so the AMX/vector grouped-gemm kernels + can be used even when TP sharding (moe_intermediate_size // tp_size) + lands on an unaligned value (e.g. moe_intermediate_size=704 at tp=4 + -> 176). Only applies to the x86 (AMX/vec) kernels and half-split + gate/up activations; interleaved layouts (swigluoai) are left + untouched. + """ + if not hasattr(torch.ops._C, "prepack_moe_weight"): + return + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: + return + + intermediate_size = layer.w2_weight.size(2) + remainder = intermediate_size % _MOE_GROUPED_GEMM_N_TILE + if remainder == 0: + return + if layer.activation == MoEActivation.SWIGLUOAI: + return + + pad = _MOE_GROUPED_GEMM_N_TILE - remainder + padded_size = intermediate_size + pad + num_experts, _, hidden_size = layer.w13_weight.shape + + new_w13 = layer.w13_weight.new_zeros(num_experts, 2 * padded_size, hidden_size) + new_w13[:, :intermediate_size] = layer.w13_weight[:, :intermediate_size] + new_w13[:, padded_size : padded_size + intermediate_size] = layer.w13_weight[ + :, intermediate_size: + ] + replace_parameter(layer, "w13_weight", new_w13) + + new_w2 = layer.w2_weight.new_zeros(num_experts, hidden_size, padded_size) + new_w2[:, :, :intermediate_size] = layer.w2_weight + replace_parameter(layer, "w2_weight", new_w2) + + if hasattr(layer, "w13_bias"): + new_bias = layer.w13_bias.new_zeros(num_experts, 2 * padded_size) + new_bias[:, :intermediate_size] = layer.w13_bias[:, :intermediate_size] + new_bias[:, padded_size : padded_size + intermediate_size] = layer.w13_bias[ + :, intermediate_size: + ] + replace_parameter(layer, "w13_bias", new_bias) + + def _grouped_gemm_alignment_error(self, layer: torch.nn.Module) -> str: + # w2's input size is the per-partition MoE intermediate size (the + # dimension TP-sharding splits), and it's what most commonly breaks + # alignment, e.g. moe_intermediate_size=704 at tp=4 gives + # 704 // 4 == 176, which isn't a multiple of 32. + intermediate_size_per_partition = layer.w2_weight.size(2) + return ( + "CPU fused-MoE AMX grouped-gemm kernel cannot be used for a " + f"layer with w13 shape {tuple(layer.w13_weight.shape)} / w2 " + f"shape {tuple(layer.w2_weight.shape)}: the per-partition MoE " + f"intermediate size ({intermediate_size_per_partition}) is not " + f"a multiple of {_MOE_GROUPED_GEMM_N_TILE}, and automatic " + "zero-padding could not resolve this (typically because the " + "activation uses an interleaved gate/up layout, e.g. " + "swigluoai). vLLM refuses to silently fall back to the much " + "slower per-expert torch loop on AMX-capable CPUs; consider a " + "different --tensor-parallel-size." + ) + def check_grouped_gemm( self, layer: torch.nn.Module, @@ -297,20 +366,19 @@ class CPUFusedMOE: w2_input_size = layer.w2_weight.size(2) w2_output_size = layer.w2_weight.size(1) - if not (w13_output_size % 32 == 0 and w2_output_size % 32 == 0): - return False, "none" - supports_amx = torch.cpu._is_amx_tile_supported() - - if ( - supports_amx - and dtype == torch.bfloat16 - and w13_input_size % 32 == 0 - and w2_input_size % 32 == 0 - ): - return True, "amx" - if supports_amx: + if ( + dtype == torch.bfloat16 + and w13_output_size % 32 == 0 + and w2_output_size % 32 == 0 + and w13_input_size % 32 == 0 + and w2_input_size % 32 == 0 + ): + return True, "amx" + raise RuntimeError(self._grouped_gemm_alignment_error(layer)) + + if not (w13_output_size % 32 == 0 and w2_output_size % 32 == 0): return False, "none" supports_neon = current_platform.get_cpu_architecture() == CpuArchEnum.ARM @@ -321,8 +389,7 @@ class CPUFusedMOE: and w2_input_size % 4 == 0 ): return True, "neon" - else: - return False, "none" + return False, "none" return True, "vec" From 7f599d78546819948c32f2b23d913507bbb38875 Mon Sep 17 00:00:00 2001 From: haoyangli0109 <lihaoyang0109@gmail.com> Date: Mon, 27 Jul 2026 16:38:45 +0800 Subject: [PATCH 100/185] [communication] [bugfix] fix quickreduce acc error in cudagraph mode (#46913) Signed-off-by: Haoyang Li <lihaoyang0109@gmail.com> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> --- csrc/quickreduce/quick_reduce.h | 38 ++++-- tests/distributed/test_rocm_quick_reduce.py | 124 ++++++++++++++++++++ 2 files changed, 154 insertions(+), 8 deletions(-) diff --git a/csrc/quickreduce/quick_reduce.h b/csrc/quickreduce/quick_reduce.h index 7506329972b..449721b5713 100644 --- a/csrc/quickreduce/quick_reduce.h +++ b/csrc/quickreduce/quick_reduce.h @@ -22,17 +22,25 @@ template <typename AllReduceKernel, typename T> __global__ __quickreduce_launch_bounds_two_shot__ static void allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks, int rank, uint8_t** dbuffer_list, - uint32_t data_offset, uint32_t flag_color, + uint32_t data_offset, uint32_t* d_flag_counters, int64_t data_size_per_phase) { int block = blockIdx.x; int grid = gridDim.x; + // Load this block's counter from device memory and advance it on-device, + // so the color keeps changing across graph replays instead of being frozen. + uint32_t flag_color = d_flag_counters[blockIdx.x]; + while (block < num_blocks) { AllReduceKernel::run(A, B, N, block, rank, dbuffer_list, data_offset, flag_color, data_size_per_phase); block += grid; flag_color++; } + // All threads compute the same final value; one writer per block is enough. + if (threadIdx.x == 0 && threadIdx.y == 0) { + d_flag_counters[blockIdx.x] = flag_color; + } } #define TWOSHOT_DISPATCH(__codec) \ @@ -42,21 +50,21 @@ allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks, hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \ dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \ num_blocks, rank, dbuffer_list, data_offset, \ - flag_color, this->kMaxProblemSize); \ + d_flag_counters, this->kMaxProblemSize); \ } else if (world_size == 4) { \ using LineCodec = __codec<T, 4>; \ using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \ hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \ dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \ num_blocks, rank, dbuffer_list, data_offset, \ - flag_color, this->kMaxProblemSize); \ + d_flag_counters, this->kMaxProblemSize); \ } else if (world_size == 8) { \ using LineCodec = __codec<T, 8>; \ using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \ hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \ dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \ num_blocks, rank, dbuffer_list, data_offset, \ - flag_color, this->kMaxProblemSize); \ + d_flag_counters, this->kMaxProblemSize); \ } // INT3 only retains good performance on TP2 (world_size == 2). On TP4/TP8 @@ -69,7 +77,7 @@ allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks, hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \ dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \ num_blocks, rank, dbuffer_list, data_offset, \ - flag_color, this->kMaxProblemSize); \ + d_flag_counters, this->kMaxProblemSize); \ } else { \ throw std::runtime_error( \ "INT3 quick all-reduce is only supported for world_size == 2 " \ @@ -94,7 +102,7 @@ struct DeviceComms { static int constexpr kMaxWorldSize = 8; bool initialized = false; - uint32_t flag_color = 1; + uint32_t* d_flag_counters = nullptr; int world_size; int rank; @@ -128,6 +136,16 @@ struct DeviceComms { // Clear the flags buffer. HIP_CHECK(hipMemset(dbuffer, 0, flags_buffer_size)); + // One flag-color counter per block, advanced by the kernel. Start at 1 + // to stay clear of the flags buffer we just zeroed. + HIP_CHECK(hipMalloc(&d_flag_counters, kMaxNumBlocks * sizeof(uint32_t))); + { + std::vector<uint32_t> init_color(kMaxNumBlocks, 1u); + HIP_CHECK(hipMemcpy(d_flag_counters, init_color.data(), + kMaxNumBlocks * sizeof(uint32_t), + hipMemcpyHostToDevice)); + } + // Device-side list of IPC buffers. buffer_list.resize(world_size); HIP_CHECK(hipMalloc(&dbuffer_list, world_size * sizeof(uint8_t*))); @@ -144,6 +162,12 @@ struct DeviceComms { hipIpcMemHandle_t const get_handle() { return buffer_ipc_handle; } void destroy() { + // Allocated before `initialized` flips true, so free it on its own guard + // to avoid a leak if init fails partway through. + if (d_flag_counters) { + HIP_CHECK(hipFree(d_flag_counters)); + d_flag_counters = nullptr; + } if (initialized) { for (int i = 0; i < world_size; i++) { if (i != rank) { @@ -211,8 +235,6 @@ struct DeviceComms { break; } HIP_CHECK(cudaGetLastError()); - // Rotate the flag color. - flag_color += divceil(N, grid); } }; diff --git a/tests/distributed/test_rocm_quick_reduce.py b/tests/distributed/test_rocm_quick_reduce.py index a7e91f236a2..e422b032e25 100644 --- a/tests/distributed/test_rocm_quick_reduce.py +++ b/tests/distributed/test_rocm_quick_reduce.py @@ -23,6 +23,8 @@ from vllm.platforms import current_platform from vllm.transformers_utils.repo_utils import hf_api from vllm.utils.network_utils import get_open_port +from ..utils import multi_gpu_test + pytestmark = pytest.mark.skipif( not current_platform.is_rocm(), reason="ROCm-only quick-reduce tests", @@ -151,6 +153,116 @@ def _run_two_gpu_quick_allreduce_test( ) +CUDAGRAPH_WORLD_SIZE = 2 +CUDAGRAPH_ROUNDS = 10 +CUDAGRAPH_NUM_ELEMENTS = 1 << 21 # 2M fp16 = 4 MB, above the quick-reduce thresholds + + +def _quick_allreduce_cudagraph_worker( + rank: int, + world_size: int, + port: int, + quant_level: str, +): + # FP keeps the all-reduce bit-exact for small fp16 integers, so every + # replayed round can be checked exactly. + os.environ["VLLM_ROCM_QUICK_REDUCE_QUANTIZATION"] = quant_level + os.environ["VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16"] = "0" + _log(f"cudagraph worker start: rank={rank} quant={quant_level}") + + device = torch.device(f"cuda:{rank}") + torch.accelerator.set_device_index(device) + dist.init_process_group( + backend="gloo", + init_method=f"tcp://127.0.0.1:{port}", + rank=rank, + world_size=world_size, + ) + + qar = None + try: + from vllm.distributed.device_communicators.quick_all_reduce import ( + QuickAllReduce, + ) + + qar = QuickAllReduce(group=dist.GroupMember.WORLD, device=rank) + assert not qar.disabled + + N = CUDAGRAPH_NUM_ELEMENTS + inp = torch.empty(N, dtype=torch.float16, device=device) + out = torch.empty(N, dtype=torch.float16, device=device) + assert qar.should_quick_allreduce(inp) + + # Every rank contributes the same value v in a round, so the true + # cross-rank all-reduce sum is simply world_size * v. + def expected(v): + return float(world_size * v) + + if rank == 0: + print( + f"[repro] world_size={world_size} elems={N} regime={quant_level} fp16", + flush=True, + ) + + # Warmup, then capture a graph with EXACTLY ONE quick-reduce (isolated + # qr, so it is the sole writer of its flag slot -- the condition that + # triggers the stale-flag bug). + inp.fill_(1.0) + qar.quick_all_reduce(inp, out=out) + torch.accelerator.synchronize() + dist.barrier() + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + qar.quick_all_reduce(inp, out=out) + torch.accelerator.synchronize() + dist.barrier() + + for v in range(CUDAGRAPH_ROUNDS): + inp.fill_(float(v)) # in-place: same value on every rank + dist.barrier() + g.replay() + torch.accelerator.synchronize() + dist.barrier() + got = out.float() + expect = expected(v) + if rank == 0: + print(f"round {v}: got={got[:10]}, expected={expect}", flush=True) + mismatch = ~torch.isclose(got, torch.full_like(got, expect)) + num_mismatch = int(mismatch.sum().item()) + assert num_mismatch == 0, ( + f"rank={rank} round={v} expected={expect} " + f"mismatched {num_mismatch}/{got.numel()} elements; " + f"unique wrong values={torch.unique(got[mismatch])[:8].tolist()}" + ) + _log(f"cudagraph worker complete: rank={rank} rounds={CUDAGRAPH_ROUNDS}") + finally: + if qar is not None: + qar.close() + if dist.is_initialized(): + dist.destroy_process_group() + + +def _run_cudagraph_replay_test(*, world_size: int, quant_level: str): + _log(f"launch {world_size}-GPU cudagraph replay case: quant={quant_level}") + ctx = mp.get_context("spawn") + port = get_open_port() + procs = [] + + for rank in range(world_size): + proc = ctx.Process( + target=_quick_allreduce_cudagraph_worker, + args=(rank, world_size, port, quant_level), + ) + proc.start() + procs.append(proc) + + for proc in procs: + proc.join(timeout=120) + assert proc.exitcode == 0, f"worker exited with code {proc.exitcode}" + _log(f"finished {world_size}-GPU cudagraph replay case: quant={quant_level}") + + MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct" E2E_PREFILL_TOKENS = 1024 E2E_MAX_MODEL_LEN = 1536 @@ -729,6 +841,18 @@ def test_quick_allreduce_two_gpu_correctness(quant_level): ) +@multi_gpu_test(num_gpus=CUDAGRAPH_WORLD_SIZE) +def test_quick_allreduce_cudagraph_replay(): + # Regression test for the stale flag_color bug: a quick-reduce captured in a + # CUDA graph must return correct results on every replay, not the data from + # a previous round. + _log("cudagraph replay case") + _run_cudagraph_replay_test( + world_size=CUDAGRAPH_WORLD_SIZE, + quant_level="FP", + ) + + @pytest.mark.skipif( current_platform.device_count() < WORLD_SIZE, reason="requires 2 ROCm GPUs", From 394beb633b0b5b5d68aed3d2f2b5c1477e756d80 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas <akaratza@amd.com> Date: Mon, 27 Jul 2026 04:18:01 -0500 Subject: [PATCH 101/185] [Bugfix][ROCm] Use batch DMA for CPU KV cache loads (#49843) Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: OpenAI Codex <noreply@openai.com> --- tests/v1/kv_offload/cpu/test_gpu_worker.py | 14 ++++++++++++++ vllm/v1/kv_offload/cpu/gpu_worker.py | 4 ++-- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/tests/v1/kv_offload/cpu/test_gpu_worker.py b/tests/v1/kv_offload/cpu/test_gpu_worker.py index 4e5e94c2bcd..bfbb18251c6 100644 --- a/tests/v1/kv_offload/cpu/test_gpu_worker.py +++ b/tests/v1/kv_offload/cpu/test_gpu_worker.py @@ -7,6 +7,7 @@ import uuid import pytest import torch +from vllm import _custom_ops as ops from vllm.platforms import current_platform from vllm.utils.math_utils import round_up from vllm.utils.torch_utils import set_random_seed @@ -16,6 +17,7 @@ from vllm.v1.kv_offload.base import ( CanonicalKVCacheTensor, GPULoadStoreSpec, ) +from vllm.v1.kv_offload.cpu import gpu_worker from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec from vllm.v1.kv_offload.cpu.gpu_worker import CPUOffloadingWorker from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion @@ -32,6 +34,18 @@ NUM_MAPPINGS = [3] NUM_MAPPINGS_PER_GROUP = [2] +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific test") +def test_rocm_cpu_to_gpu_uses_dma(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(gpu_worker, "HAS_TRITON", True) + monkeypatch.setattr(gpu_worker.current_platform, "is_xpu", lambda: False) + monkeypatch.setattr(gpu_worker.current_platform, "is_rocm", lambda: True) + + refs = [[CanonicalKVCacheRef(tensor_idx=0, page_size_bytes=512)]] + assert gpu_worker._select_swap_blocks_fn(refs, gpu_to_cpu=False) is ( + ops.swap_blocks_batch + ) + + @pytest.mark.parametrize("gpu_to_cpu", [True, False]) @pytest.mark.parametrize("num_mappings", NUM_MAPPINGS) @pytest.mark.parametrize("gpu_page_size_bytes", GPU_PAGE_SIZES) diff --git a/vllm/v1/kv_offload/cpu/gpu_worker.py b/vllm/v1/kv_offload/cpu/gpu_worker.py index baf9a66719a..40b57ac867d 100644 --- a/vllm/v1/kv_offload/cpu/gpu_worker.py +++ b/vllm/v1/kv_offload/cpu/gpu_worker.py @@ -41,10 +41,10 @@ def _select_swap_blocks_fn( if gpu_to_cpu: return ops.swap_blocks_batch # Fall back to the C++ DMA path on platforms where Triton isn't usable - # (e.g. ROCm builds without Triton) or where GPU kernels cannot directly + # (e.g. ROCm host mappings) or where GPU kernels cannot directly # dereference CPU pointers (XPU lacks CUDA's unified virtual address space, # so the Triton kernel's tl.load(cpu_ptr) is invalid on XPU). - if not HAS_TRITON or current_platform.is_xpu(): + if not HAS_TRITON or current_platform.is_xpu() or current_platform.is_rocm(): return ops.swap_blocks_batch page_sizes = [r.page_size_bytes for g in kv_cache_groups_data_refs for r in g] # Triton wins only on small, 8-byte-aligned payloads. From 312ea82e758a27a333e4254cc7669ae1a5d7c69e Mon Sep 17 00:00:00 2001 From: Andreas Karatzas <akaratza@amd.com> Date: Mon, 27 Jul 2026 04:23:12 -0500 Subject: [PATCH 102/185] [CI][ROCm] Make hf-xet reconstruction safe on shared NFS (#49837) Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: OpenAI Codex <noreply@openai.com> --- .buildkite/scripts/hardware_ci/run-amd-test.sh | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index ed4a9ba3a2e..8b5afe5a9b0 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -387,6 +387,7 @@ initialize_native_environment() { local job_id="${BUILDKITE_JOB_ID:-${BUILDKITE_PARALLEL_JOB:-local}}" local job_id_suffix="" local native_root="" + local hf_fstype="" local hf_mount="" if [[ "$(id -u)" -ne 0 ]]; then @@ -436,6 +437,18 @@ initialize_native_environment() { return 1 fi fi + + if command -v findmnt >/dev/null 2>&1; then + hf_fstype=$(findmnt -n -T "${HF_HOME}" -o FSTYPE 2>/dev/null || true) + fi + if [[ "${hf_fstype}" == nfs || "${hf_fstype}" == nfs4 ]]; then + # Keep hf-xet state local and avoid vectored writes on shared NFS. + export HF_XET_CACHE="${native_root}/cache/hf-xet" + export HF_XET_HIGH_PERFORMANCE=0 + export HF_XET_RECONSTRUCTION_USE_VECTORED_WRITE=0 + mkdir -p "${HF_XET_CACHE}" || return 1 + echo "Configured hf-xet for shared ${hf_fstype} cache at ${HF_HOME}" + fi } run_native_preflight() { From bc3629b1c4760f7260d41a67e38ff460edbe22d7 Mon Sep 17 00:00:00 2001 From: fxmarty <9808326+fxmarty@users.noreply.github.com> Date: Mon, 27 Jul 2026 11:36:38 +0200 Subject: [PATCH 103/185] [ROCm][CI] Skip three torchao tests of gfx950 until `torchao==0.18` is released (#49732) Signed-off-by: Felix Marty <Felix.Marty@amd.com> Co-authored-by: Felix Marty <Felix.Marty@amd.com> Co-authored-by: Claude <noreply@anthropic.com> --- tests/quantization/test_torchao.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/quantization/test_torchao.py b/tests/quantization/test_torchao.py index 8efc6742a2d..a724803b9a1 100644 --- a/tests/quantization/test_torchao.py +++ b/tests/quantization/test_torchao.py @@ -1,17 +1,30 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import importlib.metadata import importlib.util import pytest import torch +from packaging import version from vllm.model_executor.model_loader import get_model_loader from vllm.platforms import current_platform +if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx950 +else: + + def on_gfx950() -> bool: + return False + + DEVICE_TYPE = current_platform.device_type DTYPE = ["bfloat16"] TORCHAO_AVAILABLE = importlib.util.find_spec("torchao") is not None +TORCHAO_VERSION_0_18_AVAILABLE = TORCHAO_AVAILABLE and version.parse( + importlib.metadata.version("torchao") +) >= version.parse("0.18.0") @pytest.mark.skipif( @@ -90,6 +103,10 @@ def test_opt_125m_awq_int4wo_model_loading_with_params(vllm_runner): @pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available") +@pytest.mark.skipif( + on_gfx950() and not TORCHAO_VERSION_0_18_AVAILABLE, + reason="requires torchao>=0.18.0 on gfx950", +) def test_online_quant_config_dict_json(vllm_runner, enable_pickle): """Testing online quantization, load_weights integration point, with config dict serialized to json string @@ -135,6 +152,10 @@ def test_online_quant_config_dict_json(vllm_runner, enable_pickle): @pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available") +@pytest.mark.skipif( + on_gfx950() and not TORCHAO_VERSION_0_18_AVAILABLE, + reason="requires torchao>=0.18.0 on gfx950", +) def test_online_quant_config_file(vllm_runner): """Testing on the fly quantization, load_weights integration point, with config file @@ -170,6 +191,10 @@ def test_online_quant_config_file(vllm_runner): @pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available") +@pytest.mark.skipif( + on_gfx950() and not TORCHAO_VERSION_0_18_AVAILABLE, + reason="requires torchao>=0.18.0 on gfx950", +) def test_reload_weights(): import json From 0906123953a5a253cc0cc4f26e548ca08ece102a Mon Sep 17 00:00:00 2001 From: TJian <tunjian.tan@embeddedllm.com> Date: Mon, 27 Jul 2026 18:05:46 +0800 Subject: [PATCH 104/185] [ROCm] [Model] Enable TML inkling (#48841) Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: OpenAI Codex <codex@openai.com> --- .../kernels/benchmark_inkling_qkvr_prep.py | 176 +++ tests/models/inkling/rocm/conftest.py | 17 + .../inkling/rocm/test_model_alignment.py | 59 + tests/models/inkling/rocm/test_mxfp4_load.py | 67 ++ .../models/inkling/rocm/test_rel_attention.py | 425 +++++++ .../rocm/test_rocm_mtp_input_fusion.py | 116 ++ .../inkling/rocm/test_sconv_cache_layout.py | 20 + vllm/models/inkling/__init__.py | 39 +- vllm/models/inkling/amd/__init__.py | 2 + vllm/models/inkling/amd/attention.py | 328 ++++++ vllm/models/inkling/amd/layernorm.py | 26 + vllm/models/inkling/amd/logits_processor.py | 129 ++ vllm/models/inkling/amd/mlp.py | 63 + vllm/models/inkling/amd/model.py | 669 +++++++++++ vllm/models/inkling/amd/moe.py | 693 +++++++++++ vllm/models/inkling/amd/mtp.py | 410 +++++++ vllm/models/inkling/amd/ops/__init__.py | 44 + .../inkling/amd/ops/fa4_rel_attention.py | 438 +++++++ vllm/models/inkling/amd/ops/fa4_warmup.py | 35 + vllm/models/inkling/amd/ops/gluon/__init__.py | 4 + .../amd/ops/gluon/rel_mha_decode_gfx950.py | 1037 +++++++++++++++++ .../amd/ops/gluon/rel_mha_extend_gfx950.py | 760 ++++++++++++ vllm/models/inkling/amd/ops/gluon/utils.py | 154 +++ vllm/models/inkling/amd/ops/lamport.py | 766 ++++++++++++ vllm/models/inkling/amd/ops/mm_towers.py | 190 +++ vllm/models/inkling/amd/ops/norm.py | 414 +++++++ vllm/models/inkling/amd/ops/qkvr_prep.py | 918 +++++++++++++++ .../inkling/amd/ops/rel_attention_decode.py | 403 +++++++ vllm/models/inkling/amd/ops/sconv.py | 290 +++++ vllm/models/inkling/amd/ops/silu_and_mul.py | 197 ++++ vllm/models/inkling/amd/sconv_swa_attn.py | 226 ++++ vllm/models/inkling/amd/short_conv.py | 98 ++ vllm/triton_utils/__init__.py | 9 + 33 files changed, 9213 insertions(+), 9 deletions(-) create mode 100644 benchmarks/kernels/benchmark_inkling_qkvr_prep.py create mode 100644 tests/models/inkling/rocm/conftest.py create mode 100644 tests/models/inkling/rocm/test_model_alignment.py create mode 100644 tests/models/inkling/rocm/test_mxfp4_load.py create mode 100644 tests/models/inkling/rocm/test_rel_attention.py create mode 100644 tests/models/inkling/rocm/test_rocm_mtp_input_fusion.py create mode 100644 tests/models/inkling/rocm/test_sconv_cache_layout.py create mode 100644 vllm/models/inkling/amd/__init__.py create mode 100644 vllm/models/inkling/amd/attention.py create mode 100644 vllm/models/inkling/amd/layernorm.py create mode 100644 vllm/models/inkling/amd/logits_processor.py create mode 100644 vllm/models/inkling/amd/mlp.py create mode 100644 vllm/models/inkling/amd/model.py create mode 100644 vllm/models/inkling/amd/moe.py create mode 100644 vllm/models/inkling/amd/mtp.py create mode 100644 vllm/models/inkling/amd/ops/__init__.py create mode 100644 vllm/models/inkling/amd/ops/fa4_rel_attention.py create mode 100644 vllm/models/inkling/amd/ops/fa4_warmup.py create mode 100644 vllm/models/inkling/amd/ops/gluon/__init__.py create mode 100644 vllm/models/inkling/amd/ops/gluon/rel_mha_decode_gfx950.py create mode 100644 vllm/models/inkling/amd/ops/gluon/rel_mha_extend_gfx950.py create mode 100644 vllm/models/inkling/amd/ops/gluon/utils.py create mode 100644 vllm/models/inkling/amd/ops/lamport.py create mode 100644 vllm/models/inkling/amd/ops/mm_towers.py create mode 100644 vllm/models/inkling/amd/ops/norm.py create mode 100644 vllm/models/inkling/amd/ops/qkvr_prep.py create mode 100644 vllm/models/inkling/amd/ops/rel_attention_decode.py create mode 100644 vllm/models/inkling/amd/ops/sconv.py create mode 100644 vllm/models/inkling/amd/ops/silu_and_mul.py create mode 100644 vllm/models/inkling/amd/sconv_swa_attn.py create mode 100644 vllm/models/inkling/amd/short_conv.py diff --git a/benchmarks/kernels/benchmark_inkling_qkvr_prep.py b/benchmarks/kernels/benchmark_inkling_qkvr_prep.py new file mode 100644 index 00000000000..8aa1233c50a --- /dev/null +++ b/benchmarks/kernels/benchmark_inkling_qkvr_prep.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import statistics + +import torch +from tabulate import tabulate + +from vllm.models.inkling.nvidia.ops import qkvr_prep +from vllm.utils.argparse_utils import FlexibleArgumentParser + + +def make_inputs(tokens: int, tp_size: int, is_local: bool): + torch.manual_seed(0) + num_q_heads = 64 // tp_size + num_kv_heads = (16 if is_local else 8) // tp_size + head_dim = 128 + d_rel = 16 + rel_extent = 512 if is_local else 1024 + page_size = 16 + num_blocks = (tokens + page_size - 1) // page_size + q_width = num_q_heads * head_dim + kv_width = num_kv_heads * head_dim + r_width = num_q_heads * d_rel + device = "cuda" + + qkvr = torch.randn( + tokens, + q_width + 2 * kv_width + r_width, + device=device, + dtype=torch.bfloat16, + ) + k_weight = torch.randn(kv_width, 4, device=device, dtype=torch.bfloat16) + v_weight = torch.randn_like(k_weight) + q_norm_weight = torch.randn(head_dim, device=device, dtype=torch.bfloat16) + k_norm_weight = torch.randn_like(q_norm_weight) + rel_proj = torch.randn(d_rel, rel_extent, device=device, dtype=torch.bfloat16) + conv_cache = torch.zeros( + num_blocks, + num_kv_heads, + page_size, + 2 * head_dim, + device=device, + dtype=torch.bfloat16, + ) + key_cache = torch.empty( + num_blocks, + page_size, + num_kv_heads, + head_dim, + device=device, + dtype=torch.bfloat16, + ) + value_cache = torch.empty_like(key_cache) + positions = torch.arange(tokens, device=device, dtype=torch.int64) + block_table = torch.arange(num_blocks, device=device, dtype=torch.int32)[None] + seq_idx = torch.zeros(tokens, device=device, dtype=torch.int32) + slots = torch.arange(tokens, device=device, dtype=torch.int64) + query_start = torch.zeros(tokens, device=device, dtype=torch.int32) + log_scaling = None + if not is_local: + effective_n = (positions + 1).to(torch.float32) + log_scaling = 1.0 + 0.1 * torch.log(torch.clamp(effective_n / 128000, min=1.0)) + return ( + qkvr, + k_weight, + v_weight, + q_norm_weight, + k_norm_weight, + rel_proj, + 1e-6, + num_q_heads, + num_kv_heads, + head_dim, + d_rel, + conv_cache, + key_cache, + value_cache, + positions, + block_table, + seq_idx, + slots, + query_start, + slots, + 0, + head_dim, + page_size, + log_scaling, + ) + + +def capture(implementation, inputs): + outputs = [] + + def run(): + outputs[:] = implementation.fused_qkvr_prep(*inputs) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(stream) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + torch.accelerator.synchronize() + return graph, outputs + + +def time_graph(graph: torch.cuda.CUDAGraph, warmup: int, repeats: int) -> float: + for _ in range(warmup): + graph.replay() + torch.accelerator.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(repeats): + graph.replay() + end.record() + end.synchronize() + return start.elapsed_time(end) * 1000 / repeats + + +def benchmark(inputs, args) -> float: + graph, _ = capture(qkvr_prep, inputs) + return statistics.median( + time_graph(graph, args.warmup, args.repeats) for _ in range(args.trials) + ) + + +@torch.inference_mode() +def main(args): + rows = [] + for tp_size in args.tp_sizes: + for tokens in args.tokens: + for is_local in (True, False): + triton_us = benchmark(make_inputs(tokens, tp_size, is_local), args) + rows.append( + [ + tp_size, + tokens, + "local" if is_local else "global", + triton_us, + ] + ) + + print("Inkling QKVR prep (CUDA graph, median latency)") + print( + tabulate( + rows, + headers=[ + "TP", + "tokens", + "scope", + "Triton (us)", + ], + floatfmt=("d", "d", "", ".2f"), + ) + ) + + +if __name__ == "__main__": + parser = FlexibleArgumentParser() + parser.add_argument( + "--tokens", + type=int, + nargs="+", + default=[1 << power for power in range(15)], + ) + parser.add_argument("--tp-sizes", type=int, nargs="+", default=[4, 8]) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--repeats", type=int, default=200) + parser.add_argument("--trials", type=int, default=5) + main(parser.parse_args()) diff --git a/tests/models/inkling/rocm/conftest.py b/tests/models/inkling/rocm/conftest.py new file mode 100644 index 00000000000..bbf1e7e12d5 --- /dev/null +++ b/tests/models/inkling/rocm/conftest.py @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.platforms import current_platform + + +@pytest.fixture(autouse=True) +def require_rocm_cdna4() -> None: + if not current_platform.is_rocm(): + pytest.skip("requires ROCm") + + from vllm.platforms.rocm import on_gfx950 + + if not on_gfx950(): + pytest.skip("requires CDNA4") diff --git a/tests/models/inkling/rocm/test_model_alignment.py b/tests/models/inkling/rocm/test_model_alignment.py new file mode 100644 index 00000000000..3bf3859de8b --- /dev/null +++ b/tests/models/inkling/rocm/test_model_alignment.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ROCm guards for shared Inkling model-definition contracts.""" + +from pathlib import Path + +import torch + +import vllm.models.inkling as inkling +from vllm.lora.utils import get_supported_lora_modules +from vllm.model_executor.models.interfaces import supports_lora +from vllm.models.inkling.amd import model as amd_model +from vllm.models.inkling.amd.mtp import InklingMTP + + +def test_backend_neutral_definitions_match_nvidia_reference() -> None: + inkling_dir = Path(amd_model.__file__).parents[1] + for filename in ("logits_processor.py", "mtp.py"): + assert (inkling_dir / "amd" / filename).read_bytes() == ( + inkling_dir / "nvidia" / filename + ).read_bytes() + + +def test_rocm_exports_its_aligned_mtp_implementation() -> None: + assert inkling.InklingMTP is InklingMTP + + +def test_rocm_model_exposes_the_upstream_lora_contract() -> None: + model_cls = amd_model._TmlForCausalLMBase + + assert supports_lora(model_cls) + assert model_cls.packed_modules_mapping == { + "qkvr": ["wq_du", "wk_dv", "wv_dv", "wr_du"], + "w13": ["w1", "w3"], + } + assert model_cls.embedding_modules == {"lm_head": "output_embeddings"} + + model = torch.nn.Module() + model.embedding_modules = model_cls.embedding_modules + supported = get_supported_lora_modules(model) + assert "embed_tokens" not in supported + assert "lm_head" in supported + + +def test_lightseek_bundled_adapter_weights_remain_opt_in() -> None: + mapper = amd_model._TmlForCausalLMBase.hf_to_vllm_mapper + weight = torch.empty(1) + name, mapped_weight = next( + iter( + mapper.apply([("language_model.layers.3.attn.wq_du.lora_A.weight", weight)]) + ) + ) + + assert name == "model.layers.3.attn.qkvr.lora_A.weight" + assert mapped_weight.shard_id == 0 + assert amd_model._is_peft_adapter_weight(name) + + lm_head_name = mapper.apply_list(["language_model.lm_head.lora_B.weight"]) + assert lm_head_name == ["lm_head.lora_B.weight"] diff --git a/tests/models/inkling/rocm/test_mxfp4_load.py b/tests/models/inkling/rocm/test_mxfp4_load.py new file mode 100644 index 00000000000..00ff2f4018f --- /dev/null +++ b/tests/models/inkling/rocm/test_mxfp4_load.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for Quark OCP MXFP4 weights loaded into AITER MoE layouts.""" + +from types import SimpleNamespace + +import torch +from torch import nn + +from vllm.models.inkling.amd.moe import InklingMoE + + +def _fake_moe() -> tuple[InklingMoE, SimpleNamespace]: + moe = InklingMoE.__new__(InklingMoE) + nn.Module.__init__(moe) + moe.n_routed_experts = 4 + routed = SimpleNamespace( + moe_config=SimpleNamespace( + moe_parallel_config=SimpleNamespace(tp_rank=1, tp_size=2) + ), + # Logical intermediate per rank is 3; AITER pads it to 4. + w13_weight=nn.Parameter(torch.zeros(4, 8, 2, dtype=torch.uint8), False), + w2_weight=nn.Parameter(torch.zeros(4, 5, 3, dtype=torch.uint8), False), + w13_weight_scale=nn.Parameter(torch.ones(4, 8, 2, dtype=torch.uint8), False), + w2_weight_scale=nn.Parameter(torch.ones(4, 5, 3, dtype=torch.uint8), False), + ) + moe.experts = SimpleNamespace(routed_experts=routed) + moe._local_expert_slots = ( # type: ignore[method-assign] + lambda: dict(enumerate(range(4))) + ) + return moe, routed + + +def test_loads_logical_tp_shards_and_preserves_aiter_padding(): + moe, routed = _fake_moe() + + w13 = torch.arange(4 * 12 * 2, dtype=torch.uint8).view(4, 12, 2) + w2 = torch.arange(4 * 5 * 4, dtype=torch.uint8).view(4, 5, 4) + moe.load_expert_weight("experts.w13_weight", w13) + moe.load_expert_weight("experts.w2_weight", w2) + + # TP rank 1 consumes the second six interleaved gate/up rows. AITER keeps + # one padding row in each destination half and one padding w2 column. + torch.testing.assert_close(routed.w13_weight[:, :3], w13[:, 6:12:2]) + torch.testing.assert_close(routed.w13_weight[:, 4:7], w13[:, 7:12:2]) + assert torch.count_nonzero(routed.w13_weight[:, 3]) == 0 + assert torch.count_nonzero(routed.w13_weight[:, 7]) == 0 + torch.testing.assert_close(routed.w2_weight[:, :, :2], w2[:, :, 2:4]) + assert torch.count_nonzero(routed.w2_weight[:, :, 2]) == 0 + + +def test_unflattens_quark_scales_and_preserves_scale_padding(): + moe, routed = _fake_moe() + + flat_w13_scale = torch.arange(4 * 12 * 2, dtype=torch.uint8).view(4 * 12, 2) + flat_w2_scale = torch.arange(4 * 5 * 4, dtype=torch.uint8).view(4 * 5, 4) + moe.load_expert_weight("experts.w13_weight_scale", flat_w13_scale) + moe.load_expert_weight("experts.w2_weight_scale", flat_w2_scale) + + w13_scale = flat_w13_scale.view(4, 12, 2) + w2_scale = flat_w2_scale.view(4, 5, 4) + torch.testing.assert_close(routed.w13_weight_scale[:, :3], w13_scale[:, 6:12:2]) + torch.testing.assert_close(routed.w13_weight_scale[:, 4:7], w13_scale[:, 7:12:2]) + assert torch.all(routed.w13_weight_scale[:, 3] == 1) + assert torch.all(routed.w13_weight_scale[:, 7] == 1) + torch.testing.assert_close(routed.w2_weight_scale[:, :, :2], w2_scale[:, :, 2:4]) + assert torch.all(routed.w2_weight_scale[:, :, 2] == 1) diff --git a/tests/models/inkling/rocm/test_rel_attention.py b/tests/models/inkling/rocm/test_rel_attention.py new file mode 100644 index 00000000000..c154faca34f --- /dev/null +++ b/tests/models/inkling/rocm/test_rel_attention.py @@ -0,0 +1,425 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness tests for Inkling's ROCm relative-bias paged attention.""" + +import pytest +import torch + +from vllm.models.inkling.amd.ops.fa4_rel_attention import ( + bucket_max_seqlen_q, + inkling_fa4_rel_attention, + use_gfx950_gluon_decode, + use_gfx950_gluon_extend, +) +from vllm.models.inkling.amd.ops.rel_attention_decode import ( + decode_split_count, + use_split_kv_decode, +) + +HEAD_DIM = 128 +DTYPE = torch.bfloat16 + + +def _reference( + q: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + rel_logits: torch.Tensor, + block_table: torch.Tensor, + q_lens: list[int], + kv_lens: list[int], + rel_extent: int, + window_left: int | None, +) -> torch.Tensor: + num_heads = q.shape[1] + num_kv_heads = key_cache.shape[2] + gqa_group = num_heads // num_kv_heads + out: list[torch.Tensor] = [] + q_start = 0 + + for req, (q_len, kv_len) in enumerate(zip(q_lens, kv_lens)): + page_size = key_cache.shape[1] + num_pages = (kv_len + page_size - 1) // page_size + page_ids = block_table[req, :num_pages].long() + k = key_cache[page_ids].reshape(-1, num_kv_heads, HEAD_DIM)[:kv_len] + v = value_cache[page_ids].reshape(-1, num_kv_heads, HEAD_DIM)[:kv_len] + k = k.float().repeat_interleave(gqa_group, dim=1) + v = v.float().repeat_interleave(gqa_group, dim=1) + + q_req = q[q_start : q_start + q_len].float() + rel_req = rel_logits[q_start : q_start + q_len].float() + scores = torch.einsum("qhd,khd->hqk", q_req, k) / HEAD_DIM + + q_pos = torch.arange(q_len, device=q.device)[:, None] + kv_len - q_len + k_pos = torch.arange(kv_len, device=q.device)[None, :] + distance = q_pos - k_pos + rel_idx = distance.clamp(0, rel_extent - 1) + bias = rel_req.permute(1, 0, 2).gather( + 2, rel_idx[None].expand(num_heads, -1, -1) + ) + in_rel_extent = (distance >= 0) & (distance < rel_extent) + scores += torch.where(in_rel_extent[None], bias, 0.0) + + masked = distance < 0 + if window_left is not None: + masked |= distance > window_left + scores.masked_fill_(masked[None], float("-inf")) + out.append(torch.einsum("hqk,khd->qhd", scores.softmax(-1), v)) + q_start += q_len + + return torch.cat(out).to(DTYPE) + + +def test_query_length_bucket(): + assert bucket_max_seqlen_q(1) == 1 + assert bucket_max_seqlen_q(17) == 32 + assert bucket_max_seqlen_q(33) == 64 + + +def test_split_kv_dispatch_thresholds(): + assert not use_split_kv_decode( + max_query_len=4, + max_kv_len=65536, + page_size=128, + window_left=-1, + ) + assert not use_split_kv_decode( + max_query_len=1, + max_kv_len=4096, + page_size=16, + window_left=-1, + ) + assert use_split_kv_decode( + max_query_len=1, + max_kv_len=8192, + page_size=16, + window_left=-1, + ) + assert use_split_kv_decode( + max_query_len=1, + max_kv_len=512, + page_size=128, + window_left=511, + ) + assert decode_split_count(512, 511) == 4 + assert decode_split_count(65536, -1) == 32 + + +def test_gfx950_gluon_dispatch_is_strict(monkeypatch: pytest.MonkeyPatch): + import vllm.models.inkling.amd.ops.fa4_rel_attention as rel_attention + + monkeypatch.setattr(rel_attention, "on_gfx950", lambda: True) + assert use_gfx950_gluon_decode(max_query_len=1, page_size=128, head_dim=128) + assert not use_gfx950_gluon_decode(max_query_len=1, page_size=16, head_dim=128) + assert not use_gfx950_gluon_decode(max_query_len=4, page_size=128, head_dim=128) + assert not use_gfx950_gluon_decode(max_query_len=1, page_size=128, head_dim=256) + + monkeypatch.setattr(rel_attention, "on_gfx950", lambda: False) + assert not use_gfx950_gluon_decode(max_query_len=1, page_size=128, head_dim=128) + + monkeypatch.setattr(rel_attention, "on_gfx950", lambda: True) + assert use_gfx950_gluon_extend( + max_query_len=4, + max_kv_len=8192, + page_size=128, + head_dim=128, + window_left=-1, + ) + assert not use_gfx950_gluon_extend( + max_query_len=4, + max_kv_len=8192, + page_size=128, + head_dim=128, + window_left=511, + ) + assert not use_gfx950_gluon_extend( + max_query_len=4, + max_kv_len=8192, + page_size=16, + head_dim=128, + window_left=-1, + ) + + +@pytest.mark.parametrize( + ("page_size", "window_left", "max_kv_len"), + [ + (16, None, None), + (16, 15, None), + (128, None, 8192), + ], +) +@torch.inference_mode() +def test_ragged_multi_page_relative_attention( + page_size: int, + window_left: int | None, + max_kv_len: int | None, +): + """Covers full/chunked prefill, decode, GQA, and the local window.""" + torch.manual_seed(19 + int(window_left is not None)) + device = "cuda" + q_lens = [17, 1] + kv_lens = [35, 80] + num_heads = 8 + num_kv_heads = 2 + rel_extent = 16 + max_pages = max((length + page_size - 1) // page_size for length in kv_lens) + + q = torch.randn(sum(q_lens), num_heads, HEAD_DIM, device=device) + q = torch.nn.functional.normalize(q.float(), dim=-1).to(DTYPE) + key_cache = torch.randn( + 1 + len(q_lens) * max_pages, + page_size, + num_kv_heads, + HEAD_DIM, + device=device, + ) + key_cache = torch.nn.functional.normalize(key_cache.float(), dim=-1).to(DTYPE) + value_cache = torch.randn_like(key_cache) + rel_logits = torch.randn( + sum(q_lens), num_heads, rel_extent, device=device, dtype=DTYPE + ) + + block_table = torch.stack( + [ + torch.arange(1 + req * max_pages, 1 + (req + 1) * max_pages) + for req in range(len(q_lens)) + ] + ).to(device=device, dtype=torch.int32) + cu_seqlens_q = torch.tensor( + [0, *torch.tensor(q_lens).cumsum(0).tolist()], + device=device, + dtype=torch.int32, + ) + cache_seqlens = torch.tensor(kv_lens, device=device, dtype=torch.int32) + window_size = (-1, -1) if window_left is None else (window_left, 0) + + preallocated = torch.empty_like(q) + actual = inkling_fa4_rel_attention( + q, + key_cache, + value_cache, + block_table=block_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=bucket_max_seqlen_q(max(q_lens)), + softmax_scale=1.0 / HEAD_DIM, + causal=True, + window_size=window_size, + rel_extent=rel_extent, + rel_logits=rel_logits, + max_kv_len=max_kv_len, + out=preallocated, + ) + assert actual.data_ptr() == preallocated.data_ptr() + + expected = _reference( + q, + key_cache, + value_cache, + rel_logits, + block_table, + q_lens, + kv_lens, + rel_extent, + window_left, + ) + torch.testing.assert_close(actual.float(), expected.float(), atol=3e-2, rtol=3e-2) + + +@pytest.mark.parametrize( + ("page_size", "kv_lens", "num_kv_heads", "rel_extent", "window_left"), + [ + (16, [257, 513], 1, 64, None), + (128, [257, 512], 2, 512, 511), + ], +) +@torch.inference_mode() +def test_split_kv_decode_matches_reference( + page_size: int, + kv_lens: list[int], + num_kv_heads: int, + rel_extent: int, + window_left: int | None, +): + """Covers page-16 split-KV and page-128 gfx950 Gluon decode.""" + torch.manual_seed(41 + page_size) + device = "cuda" + batch_size = len(kv_lens) + q_lens = [1] * batch_size + num_heads = 8 + max_pages = max((length + page_size - 1) // page_size for length in kv_lens) + + q = torch.randn(batch_size, num_heads, HEAD_DIM, device=device) + q = torch.nn.functional.normalize(q.float(), dim=-1).to(DTYPE) + key_cache = torch.randn( + batch_size * max_pages, + page_size, + num_kv_heads, + HEAD_DIM, + device=device, + ) + key_cache = torch.nn.functional.normalize(key_cache.float(), dim=-1).to(DTYPE) + value_cache = torch.randn_like(key_cache) + rel_logits = torch.randn( + batch_size, + num_heads, + rel_extent, + device=device, + dtype=DTYPE, + ) + block_table = torch.arange( + batch_size * max_pages, + device=device, + dtype=torch.int32, + ).view(batch_size, max_pages) + cache_seqlens = torch.tensor(kv_lens, device=device, dtype=torch.int32) + cu_seqlens_q = torch.arange( + batch_size + 1, + device=device, + dtype=torch.int32, + ) + window_size = (-1, -1) if window_left is None else (window_left, 0) + max_kv_len = 8192 if page_size == 16 else max(kv_lens) + + actual = inkling_fa4_rel_attention( + q, + key_cache, + value_cache, + block_table=block_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1, + softmax_scale=1.0 / HEAD_DIM, + causal=True, + window_size=window_size, + rel_extent=rel_extent, + rel_logits=rel_logits, + max_kv_len=max_kv_len, + ) + expected = _reference( + q, + key_cache, + value_cache, + rel_logits, + block_table, + q_lens, + kv_lens, + rel_extent, + window_left, + ) + torch.testing.assert_close(actual.float(), expected.float(), atol=3e-2, rtol=3e-2) + + +@pytest.mark.parametrize( + ( + "q_len", + "kv_len", + "num_kv_heads", + "rel_extent", + "window_left", + "gluon_enabled", + ), + [ + (1, 513, 2, 512, 511, True), + (1, 513, 2, 512, 511, False), + (4, 8192, 1, 1024, None, True), + ], +) +@torch.inference_mode() +def test_gfx950_page128_packed_kv_views_match_reference( + monkeypatch: pytest.MonkeyPatch, + q_len: int, + kv_len: int, + num_kv_heads: int, + rel_extent: int, + window_left: int | None, + gluon_enabled: bool, +): + """Cover both Gluon and split-KV against vLLM's packed KV allocation.""" + monkeypatch.setenv( + "INKLING_GFX950_GLUON", + "1" if gluon_enabled else "0", + ) + torch.manual_seed(71 + q_len) + device = "cuda" + page_size = 128 + batch_size = 2 + num_heads = 8 + q_lens = [q_len] * batch_size + kv_lens = [kv_len] * batch_size + pages_per_req = (kv_len + page_size - 1) // page_size + + q = torch.randn( + batch_size * q_len, + num_heads, + HEAD_DIM, + device=device, + dtype=DTYPE, + ) + q = torch.nn.functional.normalize(q.float(), dim=-1).to(DTYPE) + + # FlashAttentionBackend allocates logical [block, head, page, 2 * dim]. + # Inkling transposes to [block, page, head, 2 * dim] and splits K/V, + # producing non-contiguous views whose page/head strides include both. + packed_kv = torch.randn( + batch_size * pages_per_req, + num_kv_heads, + page_size, + 2 * HEAD_DIM, + device=device, + dtype=DTYPE, + ) + key_cache, value_cache = packed_kv.transpose(1, 2).split(HEAD_DIM, dim=-1) + assert not key_cache.is_contiguous() + assert key_cache.stride() == value_cache.stride() + + rel_logits = torch.randn( + batch_size * q_len, + num_heads, + rel_extent, + device=device, + dtype=DTYPE, + ) + block_table = torch.arange( + batch_size * pages_per_req, + device=device, + dtype=torch.int32, + ).view(batch_size, pages_per_req) + cache_seqlens = torch.tensor(kv_lens, device=device, dtype=torch.int32) + cu_seqlens_q = torch.arange( + 0, + batch_size * q_len + 1, + q_len, + device=device, + dtype=torch.int32, + ) + window_size = (-1, -1) if window_left is None else (window_left, 0) + + actual = inkling_fa4_rel_attention( + q, + key_cache, + value_cache, + block_table=block_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=q_len, + softmax_scale=1.0 / HEAD_DIM, + causal=True, + window_size=window_size, + rel_extent=rel_extent, + rel_logits=rel_logits, + max_kv_len=kv_len, + ) + expected = _reference( + q, + key_cache, + value_cache, + rel_logits, + block_table, + q_lens, + kv_lens, + rel_extent, + window_left, + ) + torch.testing.assert_close(actual.float(), expected.float(), atol=3e-2, rtol=3e-2) diff --git a/tests/models/inkling/rocm/test_rocm_mtp_input_fusion.py b/tests/models/inkling/rocm/test_rocm_mtp_input_fusion.py new file mode 100644 index 00000000000..948c3a7f2a8 --- /dev/null +++ b/tests/models/inkling/rocm/test_rocm_mtp_input_fusion.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ROCm numerical-equivalence tests for the fused MTP depth-layer input kernel. + +``embed_dual_rmsnorm_cat`` must match the unfused module sequence numerically: +each rmsnorm computes in fp32 and rounds to bf16 at the same points as the +vendored ``rmsnorm`` kernel (including the bf16 round-trip between the +chained backbone embed_norm and the depth embed_norm), and the fused row +gather matches ``F.embedding``. +""" + +from typing import cast + +import pytest +import torch + +from vllm.models.inkling.amd.ops.norm import ( + embed_dual_rmsnorm_cat, + embed_rmsnorm, + rmsnorm, +) + +EPS = 1e-6 +VOCAB = 4096 +BF16_ATOL = torch.finfo(torch.bfloat16).eps + + +def _assert_bf16_close(actual: torch.Tensor, expected: torch.Tensor) -> None: + # Different Triton reduction schedules can move a small number of outputs + # by one bf16 epsilon while remaining numerically equivalent. + torch.testing.assert_close(actual, expected, rtol=0, atol=BF16_ATOL) + + +def _ref(hidden, w_h, w_e, emb, w_pre=None): + if w_pre is not None: + emb = rmsnorm(emb, w_pre, EPS) + return torch.cat([rmsnorm(hidden, w_h, EPS), rmsnorm(emb, w_e, EPS)], dim=-1) + + +@pytest.mark.parametrize("n", [1536, 6144]) +@pytest.mark.parametrize("t", [0, 1, 7, 256]) +@pytest.mark.parametrize("ids_dtype", [torch.int32, torch.int64]) +def test_embed_dual_rmsnorm_cat(n: int, t: int, ids_dtype: torch.dtype) -> None: + torch.manual_seed(0) + dev = "cuda" + table = (torch.randn(VOCAB, n, device=dev) * 0.3).to(torch.bfloat16) + w_h = torch.randn(n, device=dev).to(torch.bfloat16) + w_e = (1 + 0.01 * torch.randn(n, device=dev)).to(torch.bfloat16) + w_pre = torch.randn(n, device=dev).to(torch.bfloat16) + ids = torch.randint(0, VOCAB, (t,), device=dev, dtype=ids_dtype) + hidden = (torch.randn(t, n, device=dev) * 2).to(torch.bfloat16) + emb = table[ids.long()] + + # Fused gather + chained backbone pre-norm (the decode draft-step path). + out = embed_dual_rmsnorm_cat( + hidden, + w_h, + w_e, + EPS, + input_ids=ids, + embed_table=table, + pre_norm_weight=w_pre, + ) + assert out.shape == (t, 2 * n) + _assert_bf16_close(out, _ref(hidden, w_h, w_e, emb, w_pre)) + + # Precomputed embeds, no pre-norm (draft prefill with target-merged MM + # embeddings, already backbone-normed). + out = embed_dual_rmsnorm_cat(hidden, w_h, w_e, EPS, embeds=emb) + _assert_bf16_close(out, _ref(hidden, w_h, w_e, emb)) + + # Fused gather, no pre-norm (use_embed_norm=False). + out = embed_dual_rmsnorm_cat( + hidden, w_h, w_e, EPS, input_ids=ids, embed_table=table + ) + _assert_bf16_close(out, _ref(hidden, w_h, w_e, emb)) + + +@pytest.mark.parametrize("n", [1536, 6144]) +@pytest.mark.parametrize("t", [0, 1, 7, 256]) +@pytest.mark.parametrize("ids_dtype", [torch.int32, torch.int64]) +def test_embed_rmsnorm(n: int, t: int, ids_dtype: torch.dtype) -> None: + torch.manual_seed(0) + dev = "cuda" + table = (torch.randn(VOCAB, n, device=dev) * 0.3).to(torch.bfloat16) + w = torch.randn(n, device=dev).to(torch.bfloat16) + ids = torch.randint(0, VOCAB, (t,), device=dev, dtype=ids_dtype) + ref_emb = table[ids.long()] + + # Gather + embed_norm (base model / MTP prefill embed path). + out = cast(torch.Tensor, embed_rmsnorm(ids, table, w, EPS)) + assert out.shape == (t, n) + _assert_bf16_close(out, rmsnorm(ref_emb, w, EPS) if t else ref_emb) + + # Pure gather (use_embed_norm=False / replicated module forward). + out = cast(torch.Tensor, embed_rmsnorm(ids, table, None, EPS)) + assert torch.equal(out, ref_emb) + + # Chained first-layer attn_norm (the target text-path forward): one launch + # emits both the residual and layer 0's normed attention input. + w_chain = (1 + 0.05 * torch.randn(n, device=dev)).to(torch.bfloat16) + res, attn_in = cast( + tuple[torch.Tensor, torch.Tensor], + embed_rmsnorm(ids, table, w, EPS, chain_weight=w_chain), + ) + ref_res = rmsnorm(ref_emb, w, EPS) if t else ref_emb + _assert_bf16_close(res, ref_res) + _assert_bf16_close(attn_in, rmsnorm(ref_res, w_chain, EPS) if t else ref_res) + + # Chained without embed_norm (use_embed_norm=False). + res, attn_in = cast( + tuple[torch.Tensor, torch.Tensor], + embed_rmsnorm(ids, table, None, EPS, chain_weight=w_chain), + ) + assert torch.equal(res, ref_emb) + _assert_bf16_close(attn_in, rmsnorm(ref_emb, w_chain, EPS) if t else ref_emb) diff --git a/tests/models/inkling/rocm/test_sconv_cache_layout.py b/tests/models/inkling/rocm/test_sconv_cache_layout.py new file mode 100644 index 00000000000..eeefb877d87 --- /dev/null +++ b/tests/models/inkling/rocm/test_sconv_cache_layout.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch import nn + +from vllm.models.inkling.amd.sconv_swa_attn import InklingConvState + + +def test_runtime_sconv_block_size_tracks_unified_cache_page(): + """The cache planner may enlarge W=4 blocks to match attention pages.""" + owner = InklingConvState.__new__(InklingConvState) + nn.Module.__init__(owner) + owner.block_size = 4 + + owner.kv_cache = torch.tensor([]) + assert owner.cache_block_size == 4 + + owner.kv_cache = torch.empty(2, 1, 32, 1024) + assert owner.cache_block_size == 32 diff --git a/vllm/models/inkling/__init__.py b/vllm/models/inkling/__init__.py index 32e58905e25..4726d68e480 100644 --- a/vllm/models/inkling/__init__.py +++ b/vllm/models/inkling/__init__.py @@ -2,12 +2,21 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import TYPE_CHECKING +from vllm.platforms import current_platform + if TYPE_CHECKING: - from .nvidia.model import ( - InklingForCausalLM, - InklingForConditionalGeneration, - ) - from .nvidia.mtp import InklingMTP + if current_platform.is_rocm(): + from .amd.model import ( + InklingForCausalLM, + InklingForConditionalGeneration, + ) + from .amd.mtp import InklingMTP as InklingMTP + else: + from .nvidia.model import ( + InklingForCausalLM, + InklingForConditionalGeneration, + ) + from .nvidia.mtp import InklingMTP as InklingMTP __all__ = [ "InklingForConditionalGeneration", @@ -18,11 +27,23 @@ __all__ = [ def __getattr__(name: str): if name == "InklingMTP": - from .nvidia import mtp + if current_platform.is_rocm(): + from .amd import mtp as amd_mtp + + return amd_mtp.InklingMTP + + from .nvidia import mtp as nvidia_mtp + + return nvidia_mtp.InklingMTP - return mtp.InklingMTP if name in __all__: - from .nvidia import model + if current_platform.is_rocm(): + from .amd import model as amd_model + + return getattr(amd_model, name) + + from .nvidia import model as nvidia_model + + return getattr(nvidia_model, name) - return getattr(model, name) raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/vllm/models/inkling/amd/__init__.py b/vllm/models/inkling/amd/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm/models/inkling/amd/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/inkling/amd/attention.py b/vllm/models/inkling/amd/attention.py new file mode 100644 index 00000000000..edbc64bd817 --- /dev/null +++ b/vllm/models/inkling/amd/attention.py @@ -0,0 +1,328 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +from typing import cast + +import torch +from torch import nn + +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.config import VllmConfig, get_current_vllm_config +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.utils.torch_utils import ( + canonicalize_singleton_dim_strides, + kv_cache_dtype_str_to_dtype, +) +from vllm.v1.attention.backend import AttentionBackend +from vllm.v1.attention.backends.flash_attn import ( + FlashAttentionBackend, + FlashAttentionMetadata, +) +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheSpec, + SlidingWindowSpec, +) + +from ..configs import InklingModelConfig +from .layernorm import InklingRMSNorm +from .ops.fa4_rel_attention import ( + bucket_max_seqlen_q, + inkling_fa4_num_splits, + inkling_fa4_rel_attention, +) +from .ops.fa4_warmup import InklingFA4WarmupConfig, register_fa4_warmup +from .ops.qkvr_prep import fused_qkvr_prep +from .sconv_swa_attn import _K, _V, InklingConvState, InklingSconvMetadata +from .short_conv import InklingShortConv + + +def compute_log_scaling_tau( + positions: torch.Tensor, n_floor: int, alpha: float +) -> torch.Tensor: + effective_n = (positions + 1).to(torch.float32) + return 1.0 + alpha * torch.log(torch.clamp(effective_n / float(n_floor), min=1.0)) + + +class RelLogitsProj(nn.Module): + """Project the per-head relative branch ``r`` to per-distance logits.""" + + def __init__(self, d_rel: int, rel_extent: int) -> None: + super().__init__() + self.d_rel = d_rel + self.rel_extent = rel_extent + self.proj = nn.Parameter(torch.empty(d_rel, rel_extent), requires_grad=False) + + def forward(self, r_out: torch.Tensor) -> torch.Tensor: + # r_out: (T, num_heads, d_rel) -> (T, num_heads, rel_extent) + return torch.einsum("thd,de->the", r_out, self.proj) + + +class InklingAttention(nn.Module, AttentionLayerBase): + def __init__( + self, + config: InklingModelConfig, + *, + num_heads: int, + num_kv_heads: int, + head_dim: int, + rel_extent: int, + local_extent: int, + is_local: bool, + prefix: str, + quant_config: QuantizationConfig | None = None, + conv_owner: InklingConvState, + ) -> None: + super().__init__() + self.prefix = prefix + self.is_local = is_local + self.hidden_size = config.hidden_size + self.head_dim = head_dim + self.d_rel = config.d_rel + self.log_scaling_n_floor = config.log_scaling_n_floor + self.log_scaling_alpha = config.log_scaling_alpha + # q/k are per-head RMS-normed (unit norm), so Inkling scales by 1/head_dim. + self.scaling = 1.0 / head_dim + + tp_size = get_tensor_model_parallel_world_size() + self.num_total_heads = num_heads + self.num_total_kv_heads = num_kv_heads + assert self.num_total_heads % tp_size == 0 + self.num_heads = self.num_total_heads // tp_size + if self.num_total_kv_heads >= tp_size: + assert self.num_total_kv_heads % tp_size == 0 + else: + assert tp_size % self.num_total_kv_heads == 0 + self.num_kv_heads = max(1, self.num_total_kv_heads // tp_size) + # When tp_size > num_kv_heads the K/V projections are padded up to + # tp_size heads so each rank gets at least one (GQA replication). + kv_total_for_sizing = max(self.num_total_kv_heads, tp_size) + + self.qkvr = MergedColumnParallelLinear( + input_size=config.hidden_size, + output_sizes=[ + head_dim * self.num_total_heads, + head_dim * kv_total_for_sizing, + head_dim * kv_total_for_sizing, + self.d_rel * self.num_total_heads, + ], + bias=config.q_bias, + quant_config=quant_config, + prefix=f"{prefix}.qkvr", + ) + self.wo_ud = RowParallelLinear( + input_size=head_dim * self.num_total_heads, + output_size=config.hidden_size, + bias=config.o_bias, + quant_config=quant_config, + # reduce_results=False: the partial output is all-reduced below + # (one-shot custom AR) so the attention-output sconv can run on the + # full hidden width fused with the residual add + rmsnorm. + reduce_results=False, + prefix=f"{prefix}.wo_ud", + ) + self.rel_extent = local_extent if is_local else rel_extent + self.local_extent = local_extent if is_local else None + self.rel_logits_proj = RelLogitsProj(self.d_rel, self.rel_extent) + self.q_norm = InklingRMSNorm(head_dim, eps=config.rms_norm_eps) + self.k_norm = InklingRMSNorm(head_dim, eps=config.rms_norm_eps) + + # Short convolution on the K/V streams (per-head-width, TP sharded), + # applied after the qkvr projection and before q/k norm. + kv_conv_dim = self.num_kv_heads * head_dim + self.conv_owner = conv_owner + self.k_sconv = InklingShortConv( + kv_conv_dim, config.sconv_kernel_size, owner=conv_owner, stream_idx=_K + ) + self.v_sconv = InklingShortConv( + kv_conv_dim, config.sconv_kernel_size, owner=conv_owner, stream_idx=_V + ) + + # FA4 left/right window; right=0 keeps it causal. local_extent-1 mirrors + # the source (sliding_window_size - 1). + self.window_size: tuple[int, int] = ( + (local_extent - 1, 0) if is_local else (-1, -1) + ) + # Static per-layer-type KV length bound for the split heuristic: local + # layers never see more than the sliding window. + vllm_config = get_current_vllm_config() + self._max_kv_len = ( + local_extent if is_local else vllm_config.model_config.max_model_len + ) + + # ---- KV-cache wiring (reuse FlashAttentionBackend for metadata) ---- + cache_config = vllm_config.cache_config + self.kv_cache_dtype = ( + cache_config.cache_dtype if cache_config is not None else "auto" + ) + self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype( + self.kv_cache_dtype, vllm_config.model_config + ) + self.register_buffer("k_scale", torch.ones((), dtype=torch.float32)) + self.register_buffer("v_scale", torch.ones((), dtype=torch.float32)) + + compilation_config = vllm_config.compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + self.kv_cache = torch.tensor([]) # replaced by bind_kv_cache + + register_fa4_warmup( + InklingFA4WarmupConfig( + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + rel_extent=self.rel_extent, + window_size=self.window_size, + is_local=self.is_local, + max_kv_len=self._max_kv_len, + dtype=vllm_config.model_config.dtype, + kv_dtype=self.kv_cache_torch_dtype, + block_size=vllm_config.cache_config.block_size, + max_num_reqs=vllm_config.scheduler_config.max_num_seqs, + max_num_batched_tokens=( + vllm_config.scheduler_config.max_num_batched_tokens + ), + ) + ) + + def get_attn_backend(self) -> type[AttentionBackend]: + return FlashAttentionBackend + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + block_size = vllm_config.cache_config.block_size + if self.is_local: + assert self.local_extent is not None + return SlidingWindowSpec( + block_size=block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_dim, + dtype=self.kv_cache_torch_dtype, + sliding_window=self.local_extent, + ) + return FullAttentionSpec( + block_size=block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_dim, + dtype=self.kv_cache_torch_dtype, + ) + + def _split_kv_cache(self) -> tuple[torch.Tensor, torch.Tensor]: + key_cache, value_cache = self.kv_cache.transpose(1, 2).split( + self.head_dim, dim=-1 + ) + return ( + canonicalize_singleton_dim_strides(key_cache), + canonicalize_singleton_dim_strides(value_cache), + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + log_scaling: torch.Tensor | None = None, + ) -> torch.Tensor: + num_tokens = hidden_states.shape[0] + qkvr, _ = self.qkvr(hidden_states) + + attn_metadata = get_forward_context().attn_metadata + attn_output = torch.empty( + (num_tokens, self.num_heads, self.head_dim), + dtype=qkvr.dtype, + device=qkvr.device, + ) + if not isinstance(attn_metadata, dict): + attn_output.zero_() + else: + conv_meta = attn_metadata[self.conv_owner.prefix] + md = attn_metadata[self.prefix] + assert isinstance(conv_meta, InklingSconvMetadata) + fa_md = cast(FlashAttentionMetadata, md) + assert self.kv_cache.numel() > 0 + assert self.conv_owner.kv_cache.numel() > 0 + # One launch: K/V sconv (conv-cache insert + conv + residual), + # Q/K per-head rmsnorm, and the attention KV-cache write. K/V are + # consumed via the KV cache; only normed q is materialized. + key_cache, value_cache = self._split_kv_cache() + off_k, _ = self.conv_owner.stream_ranges[_K] + off_v, _ = self.conv_owner.stream_ranges[_V] + q, rel_logits = fused_qkvr_prep( + qkvr, + self.k_sconv.weight.squeeze(1), + self.v_sconv.weight.squeeze(1), + self.q_norm.weight, + self.k_norm.weight, + self.rel_logits_proj.proj, + self.q_norm.variance_epsilon, + self.num_heads, + self.num_kv_heads, + self.head_dim, + self.d_rel, + self.conv_owner.kv_cache, + key_cache, + value_cache, + positions, + conv_meta.block_table, + conv_meta.seq_idx, + conv_meta.slot_mapping, + conv_meta.query_start, + fa_md.slot_mapping, + off_k, + off_v, + self.conv_owner.cache_block_size, + log_scaling if not self.is_local else None, + ) + q = q.view(num_tokens, self.num_heads, self.head_dim) + self._attention(q, rel_logits, attn_output) + + flat = attn_output.view(num_tokens, -1) + output, _ = self.wo_ud(flat) + return output + + @eager_break_during_capture + def _attention( + self, + q: torch.Tensor, + rel_logits: torch.Tensor, + output: torch.Tensor, + ) -> None: + attn_metadata = get_forward_context().attn_metadata + assert isinstance(attn_metadata, dict) + md = cast(FlashAttentionMetadata, attn_metadata[self.prefix]) + + nt = md.num_actual_tokens + key_cache, value_cache = self._split_kv_cache() + max_seqlen_q = bucket_max_seqlen_q(md.max_query_len) + num_splits = inkling_fa4_num_splits( + is_local=self.is_local, + batch_size=md.seq_lens.shape[0], + max_query_len=max_seqlen_q, + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + max_kv_len=md.max_seq_len, + ) + inkling_fa4_rel_attention( + q[:nt], + key_cache, + value_cache, + block_table=md.block_table, + cache_seqlens=md.seq_lens, + cu_seqlens_q=md.query_start_loc, + max_seqlen_q=max_seqlen_q, + softmax_scale=self.scaling, + causal=True, + window_size=self.window_size, + rel_extent=self.rel_extent, + rel_logits=rel_logits[:nt], + num_splits=num_splits, + max_kv_len=md.max_seq_len, + out=output[:nt], + ) diff --git a/vllm/models/inkling/amd/layernorm.py b/vllm/models/inkling/amd/layernorm.py new file mode 100644 index 00000000000..18f9951f65d --- /dev/null +++ b/vllm/models/inkling/amd/layernorm.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling RMSNorm (no bias, weight-scaled), backed by the vendored Triton kernel.""" + +from __future__ import annotations + +import torch +from torch import nn + +from .ops import rmsnorm + + +class InklingRMSNorm(nn.Module): + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + self.hidden_size = hidden_size + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if x.numel() == 0: + return x + original_shape = x.shape + x_2d = x.contiguous().view(-1, self.hidden_size) + y = rmsnorm(x_2d, self.weight, self.variance_epsilon) + return y.view(original_shape) diff --git a/vllm/models/inkling/amd/logits_processor.py b/vllm/models/inkling/amd/logits_processor.py new file mode 100644 index 00000000000..2e142c5961e --- /dev/null +++ b/vllm/models/inkling/amd/logits_processor.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling logits processor (muP + LoRA aware). + +Inkling divides the final logits by a muP width multiplier +(``logits_mup_width_multiplier``). This applies it two ways, depending on +whether an lm_head LoRA is attached: + +* No LoRA: fold ``1/mup`` into the lm_head GEMM alpha (fp32 epilogue) -- no + separate elementwise kernel, no extra rounding, no weight mutation. +* LoRA attached: the LoRA manager wraps this layer in + ``LogitsProcessorWithLoRA``, whose ``forward`` calls + ``type(base_layer).forward(self=wrapper)`` -- so this ``forward`` runs with + ``self`` bound to the wrapper. We detect that via ``base_layer`` and take the + LoRA path: run the wrapper's ``_get_logits`` (base logits + the lm_head LoRA + delta), then divide the full logits by the multiplier so the delta is scaled + too. muP thus composes with the LoRA delta, with the dispatch as the only + model-side branching. +""" + +from __future__ import annotations + +import torch + +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding + + +class InklingLogitsProcessor(LogitsProcessor): + """``LogitsProcessor`` that applies Inkling's muP logits width multiplier. + + Args: + vocab_size: Padded vocabulary size. + org_vocab_size: Unpadded vocabulary size (defaults to ``vocab_size``). + scale: Base logits scale (kept ``1.0`` for the served checkpoint). + logits_as_input: Whether the input is already logits. + soft_cap: Optional logit soft cap (``None`` for the served checkpoint). + logits_mup_width_multiplier: muP width divisor for the final logits; + ``None`` or ``0`` disables it. + """ + + def __init__( + self, + vocab_size: int, + org_vocab_size: int | None = None, + scale: float = 1.0, + logits_as_input: bool = False, + soft_cap: float | None = None, + logits_mup_width_multiplier: float | None = None, + ) -> None: + super().__init__( + vocab_size=vocab_size, + org_vocab_size=org_vocab_size, + scale=scale, + logits_as_input=logits_as_input, + soft_cap=soft_cap, + ) + self.logits_mup_width_multiplier = logits_mup_width_multiplier + self._logits_zero: torch.Tensor | None = None + + def forward( + self, + lm_head: VocabParallelEmbedding, + hidden_states: torch.Tensor, + embedding_bias: torch.Tensor | None = None, + ) -> torch.Tensor | None: + # ``base_layer`` exists only on the LogitsProcessorWithLoRA wrapper, + # which calls this forward with ``self`` bound to the wrapper. The + # wrapper is not an ``InklingLogitsProcessor`` instance, so dispatch + # ``_lora_forward`` explicitly through the base_layer's class (it + # provides ``_get_logits``/``logits_as_input``; only ``_lora_forward`` + # lives on this class). + if hasattr(self, "base_layer"): + return type(self.base_layer)._lora_forward( + self, lm_head, hidden_states, embedding_bias + ) + return self._base_forward(lm_head, hidden_states, embedding_bias) + + def _lora_forward( + self, + lm_head: VocabParallelEmbedding, + hidden_states: torch.Tensor, + embedding_bias: torch.Tensor | None = None, + ) -> torch.Tensor | None: + # ``self`` is the LogitsProcessorWithLoRA wrapper here: ``_get_logits`` + # returns the base logits plus the lm_head LoRA delta. Apply the muP + # divisor on the full logits so the LoRA delta is scaled too. + mup_multiplier = self.base_layer.logits_mup_width_multiplier + mup = 1.0 / mup_multiplier if mup_multiplier else None + if self.logits_as_input: + logits = hidden_states + else: + logits = self._get_logits(hidden_states, lm_head, embedding_bias) + # TODO: fuse this multiplication + if logits is not None and mup: + assert self.base_layer.soft_cap is None + assert self.base_layer.scale == 1.0 + logits = logits * mup + return logits + + def _base_forward( + self, + lm_head: VocabParallelEmbedding, + hidden_states: torch.Tensor, + embedding_bias: torch.Tensor | None = None, + ) -> torch.Tensor | None: + mup = self.logits_mup_width_multiplier + if not mup: + return super().forward(lm_head, hidden_states, embedding_bias) + # Fold the muP width divisor into the lm_head GEMM alpha (fp32 epilogue): + # no separate elementwise kernel, no bf16 rounding of scaled logits, and + # no weight mutation. Overfit to the served checkpoint: bf16 lm_head, no + # soft cap, unit logits scale. + assert self.soft_cap is None + assert self.scale == 1.0 + w = lm_head.weight + if self._logits_zero is None: + self._logits_zero = w.new_zeros(1) + logits = torch.addmm( + self._logits_zero, + hidden_states, + w.t(), + beta=0.0, + alpha=1.0 / mup, + ) + logits = self._gather_logits(logits) + if logits is not None: + logits = logits[..., : self.org_vocab_size] + return logits diff --git a/vllm/models/inkling/amd/mlp.py b/vllm/models/inkling/amd/mlp.py new file mode 100644 index 00000000000..51deb38e64e --- /dev/null +++ b/vllm/models/inkling/amd/mlp.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling dense SwiGLU MLP (also used as the MoE shared expert). + +The checkpoint stores the gate/up projection as a single fused, *interleaved* +weight (``[gate0, up0, gate1, up1, ...]``), so we use a plain +``ColumnParallelLinear`` whose contiguous row-sharding keeps each gate/up pair +together, and an interleaved SwiGLU activation. +""" + +from __future__ import annotations + +import torch +from torch import nn + +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.quantization import QuantizationConfig + + +class InklingDenseMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + *, + use_global_scale: bool = False, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = ColumnParallelLinear( + hidden_size, + 2 * intermediate_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.down_proj", + ) + if use_global_scale: + self.global_scale = nn.Parameter(torch.empty(1), requires_grad=False) + else: + self.global_scale = None + + def forward(self, x: torch.Tensor) -> torch.Tensor: + from .ops import silu_and_mul_triton + + gate_up, _ = self.gate_up_proj(x) + x = silu_and_mul_triton(gate_up) + x, _ = self.down_proj(x) + if self.global_scale is not None: + x.mul_(self.global_scale) + # TP-partial output: the layer's reduce-scatter fallback consumes it. + return x diff --git a/vllm/models/inkling/amd/model.py b/vllm/models/inkling/amd/model.py new file mode 100644 index 00000000000..56ffb1adcd3 --- /dev/null +++ b/vllm/models/inkling/amd/model.py @@ -0,0 +1,669 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling model implementation for AMD GPUs.""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +import regex as re +import torch +from torch import nn + +from vllm.config import VllmConfig +from vllm.distributed import ( + get_pp_group, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_gather, + tensor_model_parallel_reduce_scatter, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.models.interfaces import ( + MultiModalEmbeddings, + SupportsLoRA, + SupportsMultiModal, + SupportsPP, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + make_empty_intermediate_tensors_factory, + make_layers, + maybe_prefix, +) +from vllm.models.inkling.common.mm_preprocess import ( + InklingDummyInputsBuilder, + InklingMultiModalProcessor, + InklingProcessingInfo, + inkling_audio_enabled, + inkling_vision_enabled, +) +from vllm.models.inkling.common.towers import InklingAudio, InklingVision +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.sequence import IntermediateTensors + +from ..configs import InklingMMConfig, InklingModelConfig +from .attention import InklingAttention, compute_log_scaling_tau +from .layernorm import InklingRMSNorm +from .logits_processor import InklingLogitsProcessor +from .mlp import InklingDenseMLP +from .moe import InklingMoE +from .ops.norm import add_rmsnorm, embed_rmsnorm +from .sconv_swa_attn import _ATTN, _MLP, InklingConvState +from .short_conv import InklingShortConv + + +def _layer_id(name: str) -> int | None: + m = re.search(r"\.layers\.(\d+)\.", name) + return int(m.group(1)) if m else None + + +def _sconv_add_norm( + delta: torch.Tensor, + hidden: torch.Tensor, + sconv: InklingShortConv, + norm: InklingRMSNorm | None, + positions: torch.Tensor, +) -> tuple[torch.Tensor | None, torch.Tensor]: + """``h = hidden + sconv(TP-sum(delta)); y = rmsnorm(h)``. + + ROCm uses the portable collective path. The Lamport P2P implementation in + the NVIDIA model relies on CUDA GDC and CUDA-specific symmetric-memory + publication semantics which cannot be linked into a gfx950 HSACO.""" + norm_w = norm.weight if norm is not None else None + eps = norm.variance_epsilon if norm is not None else 0.0 + + # RCCL RS -> shard sconv -> AG -> fused add(+rmsnorm). + shard = tensor_model_parallel_reduce_scatter(delta, dim=-1) + shard = sconv(shard.contiguous(), positions) + full = tensor_model_parallel_all_gather(shard, dim=-1) + if norm is None: + return None, hidden + full + return add_rmsnorm(hidden, full, norm_w, eps) + + +class InklingDecoderLayer(nn.Module): + def __init__( + self, + config: InklingModelConfig, + layer_id: int, + is_local: bool, + quant_config: QuantizationConfig | None, + prefix: str, + force_dense_mlp: bool = False, + ) -> None: + super().__init__() + # Per-layer owner of the conv state as a paged SWA cache. The 4 sconv + # streams (K/V/attn/mlp) are packed head-major into one block and share + # it. Built first so the attention layer can wire its K/V sconv to it. + self.conv_state = InklingConvState( + num_kv_heads=( + config.swa_num_key_value_heads + if is_local + else config.num_key_value_heads + ), + head_dim=config.swa_head_dim if is_local else config.head_dim, + hidden_size=config.hidden_size, + kernel_size=config.sconv_kernel_size, + prefix=f"{prefix}.conv_state", + ) + self.attn_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.attn = InklingAttention( + config, + num_heads=( + config.swa_num_attention_heads + if is_local + else config.num_attention_heads + ), + num_kv_heads=( + config.swa_num_key_value_heads + if is_local + else config.num_key_value_heads + ), + head_dim=config.swa_head_dim if is_local else config.head_dim, + rel_extent=config.rel_extent, + local_extent=config.sliding_window_size, + is_local=is_local, + prefix=f"{prefix}.attn", + quant_config=quant_config, + conv_owner=self.conv_state, + ) + self.mlp_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if force_dense_mlp or layer_id < config.dense_mlp_idx: + self.mlp: nn.Module = InklingDenseMLP( + hidden_size=config.hidden_size, + intermediate_size=config.dense_intermediate_size, + use_global_scale=config.use_global_scale, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + else: + # InklingMoE decides per layer (from the checkpoint exclude list) + # whether the routed experts are NVFP4 or bf16; the shared sink + # experts are always bf16. + self.mlp = InklingMoE( + config, + prefix=f"{prefix}.mlp", + quant_config=quant_config, + ) + + # Short convolution on the attention-output and MLP-output residual + # streams, hidden-sharded: the sublayer outputs are reduce-scattered + # to [T, H/tp], the sconv runs on the shard, and an all-gather + # restores the full residual — all fused with the residual add + next + # rmsnorm via the Lamport P2P kernels for decode-sized batches. + tp_size = get_tensor_model_parallel_world_size() + sconv_dim = config.hidden_size // tp_size + self.attn_sconv = InklingShortConv( + sconv_dim, config.sconv_kernel_size, owner=self.conv_state, stream_idx=_ATTN + ) + self.mlp_sconv = InklingShortConv( + sconv_dim, config.sconv_kernel_size, owner=self.conv_state, stream_idx=_MLP + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + pending: tuple[torch.Tensor | None, InklingShortConv] | None = None, + defer_mlp_add: bool = False, + attn_in: torch.Tensor | None = None, + log_scaling: torch.Tensor | None = None, + ) -> ( + torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor | None, InklingShortConv]] + ): + # The previous sublayer's (pre-reduce, pre-sconv) delta is folded in + # fused with its RS/sconv/AG and this layer's pre-attention rmsnorm. + # A None delta means the partials sit in the NVLS symm buffer. + if pending is None: + if attn_in is None: + # First layer; on the text path attn_norm comes fused with + # the embedding gather (chain_weight in embed_rmsnorm). + attn_in = self.attn_norm(hidden_states) + else: + attn_in, hidden_states = _sconv_add_norm( + pending[0], hidden_states, pending[1], self.attn_norm, positions + ) + attn_output = self.attn(positions, attn_in, log_scaling) + mlp_in, hidden_states = _sconv_add_norm( + attn_output, hidden_states, self.attn_sconv, self.mlp_norm, positions + ) + mlp_output = self.mlp(mlp_in) + if defer_mlp_add: + # Caller folds mlp_output (pre-reduce, pre-sconv) into the next + # fused sconv+add+rmsnorm. + return hidden_states, (mlp_output, self.mlp_sconv) + # Standalone (MTP) tail: finish the sublayer without a norm. + return _sconv_add_norm( + mlp_output, hidden_states, self.mlp_sconv, None, positions + )[1] + + +class InklingReplicatedEmbedding(nn.Module): + """Full-vocab embedding table replicated on every TP rank. + + Trades the full table per rank (~2.3 GiB at V=201k / H=6144 bf16, vs a + 1/tp shard) for no masked lookup and no per-lookup TP all-reduce — one + all-reduce per MTP draft step plus one per verify pass — and keeps the + full table on-rank for the fused gather+norm kernels (``embed_rmsnorm``, + ``embed_dual_rmsnorm_cat``). Bit-exact vs vocab-parallel: the all-reduce + there only ever summed one real row against exact zeros. The LM head + stays vocab-sharded. + """ + + def __init__(self, num_embeddings: int, embedding_dim: int) -> None: + super().__init__() + self.weight = nn.Parameter( + torch.empty(num_embeddings, embedding_dim, dtype=torch.get_default_dtype()), + requires_grad=False, + ) + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + return embed_rmsnorm(input_ids, self.weight, None, 0.0) + + +class InklingModel(nn.Module): + def __init__( + self, + *, + config: InklingModelConfig, + quant_config: QuantizationConfig | None, + prefix: str, + ) -> None: + super().__init__() + self.config = config + self.embed_tokens = InklingReplicatedEmbedding( + config.padded_vocab_size, config.hidden_size + ) + self.embed_norm = ( + InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if config.use_embed_norm + else None + ) + local_ids = set(config.local_layer_ids) + + def get_layer(prefix: str) -> InklingDecoderLayer: + idx = _layer_id(prefix + ".") or int(prefix.split(".")[-1]) + return InklingDecoderLayer( + config, idx, idx in local_ids, quant_config, prefix + ) + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers" + ) + self.norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states"], config.hidden_size + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + # Row gather + embed_norm in one launch. + norm = self.embed_norm + return embed_rmsnorm( + input_ids, + self.embed_tokens.weight, + norm.weight if norm is not None else None, + norm.variance_epsilon if norm is not None else 0.0, + ) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + attn_in0: torch.Tensor | None = None + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + # embed_norm was already applied when producing inputs_embeds. + hidden_states = inputs_embeds + else: + # Gather + embed_norm + the first layer's attn_norm, one launch. + norm = self.embed_norm + hidden_states, attn_in0 = embed_rmsnorm( + input_ids, + self.embed_tokens.weight, + norm.weight if norm is not None else None, + self.config.rms_norm_eps, + chain_weight=self.layers[self.start_layer].attn_norm.weight, + ) + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + log_scaling = None + if self.config.log_scaling_n_floor is not None: + log_scaling = compute_log_scaling_tau( + positions, + self.config.log_scaling_n_floor, + self.config.log_scaling_alpha, + ) + + pending: tuple[torch.Tensor | None, InklingShortConv] | None = None + for layer in self.layers[self.start_layer : self.end_layer]: + hidden_states, pending = layer( + positions, + hidden_states, + pending=pending, + defer_mlp_add=True, + attn_in=attn_in0, + log_scaling=log_scaling, + ) + attn_in0 = None + + if not get_pp_group().is_last_rank: + if pending is not None: + hidden_states = _sconv_add_norm( + pending[0], hidden_states, pending[1], None, positions + )[1] + return IntermediateTensors({"hidden_states": hidden_states}) + if pending is not None: + # Final RS/sconv/AG + residual add fused with the final rmsnorm. + norm_out = _sconv_add_norm( + pending[0], hidden_states, pending[1], self.norm, positions + )[0] + assert norm_out is not None + return norm_out + return self.norm(hidden_states) + + +class _TmlForCausalLMBase(nn.Module, SupportsPP, SupportsLoRA): + """Shared text-backbone causal-LM scaffolding for both entry classes.""" + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_substr={ + ".w13_dn": ".gate_up_proj", + ".w2_md": ".down_proj", + }, + orig_to_new_stacked={ + ".attn.wq_du.": (".attn.qkvr.", 0), + ".attn.wk_dv.": (".attn.qkvr.", 1), + ".attn.wv_dv.": (".attn.qkvr.", 2), + ".attn.wr_du.": (".attn.qkvr.", 3), + }, + orig_to_new_prefix={ + "model.llm.layers.": "model.layers.", + "model.llm.embed_norm": "model.embed_norm", + "model.llm.embed": "model.embed_tokens", + "model.llm.norm": "model.norm", + "model.llm.unembed": "lm_head", + "language_model.layers.": "model.layers.", + "language_model.lm_head.": "lm_head.", + }, + orig_to_new_suffix={ + # NVFP4 scale + ".w13_weight.scale": ".w13_weight_scale", + ".w13_weight.scale2": ".w13_weight_scale_2", + ".w2_weight.scale": ".w2_weight_scale", + ".w2_weight.scale2": ".w2_weight_scale_2", + }, + ) + # Quark uses this mapping when resolving quantization exclusions for the + # four checkpoint attention projections fused into qkvr. + packed_modules_mapping = { + "qkvr": ["wq_du", "wk_dv", "wv_dv", "wr_du"], + "w13": ["w1", "w3"], + } + embedding_modules = { + "lm_head": "output_embeddings", + } + + def _build( + self, + vllm_config: VllmConfig, + text_config: InklingModelConfig, + prefix: str, + ) -> None: + quant_config = vllm_config.quant_config + self.config = text_config + # ROCm checkpoints use Quark OCP MXFP4. The global Quark config is + # passed into each routed MoE and its exclusion list keeps the rest of + # Inkling in bf16. + # Read by the MRV2 runner to publish per-request short-conv metadata. + # Short convolution is intrinsic to Inkling, so this is always set. + self.uses_sconv = True + self.model = InklingModel( + config=text_config, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "model"), + ) + self.lm_head = ParallelLMHead( + text_config.padded_vocab_size, + text_config.hidden_size, + org_num_embeddings=text_config.padded_vocab_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = InklingLogitsProcessor( + text_config.padded_vocab_size, + org_vocab_size=text_config.vocab_size, + soft_cap=text_config.final_logit_softcapping, + logits_mup_width_multiplier=text_config.logits_mup_width_multiplier, + ) + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.model.make_empty_intermediate_tensors + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors: + return self.model( + input_ids, + positions, + intermediate_tensors, + inputs_embeds, + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return self.logits_processor(self.lm_head, hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + return _load_inkling_weights(self, weights, self.config) + + +class InklingForCausalLM(_TmlForCausalLMBase): + """Text-only entry point (``inkling_model`` checkpoints).""" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self._build(vllm_config, vllm_config.model_config.hf_config, prefix) + + +@MULTIMODAL_REGISTRY.register_processor( + InklingMultiModalProcessor, + info=InklingProcessingInfo, + dummy_inputs=InklingDummyInputsBuilder, +) +class InklingForConditionalGeneration(_TmlForCausalLMBase, SupportsMultiModal): + """Top-level (multimodal) entry point. + + Builds the vision + audio towers on top of the shared text backbone. Inkling has + NO cross-modal fusion (the vision tower emits one token per patch, the audio + tower one token per frame), so generation reuses the inherited backbone + ``forward`` / ``compute_logits`` (the latter already applies muP) and this + class only adds multimodal embedding + merge. + """ + + hf_to_vllm_mapper = _TmlForCausalLMBase.hf_to_vllm_mapper | WeightsMapper( + orig_to_new_prefix={ + "model.audio.": "audio.", + "model.visual.": "visual.vision_encoder.", + }, + ) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality.startswith("image"): + return "<|content_image|>" + if modality.startswith("audio"): + return "<|content_audio_input|>" + raise ValueError("Only image or audio modality is supported") + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + config: InklingMMConfig = vllm_config.model_config.hf_config + + self.visual = ( + InklingVision(config.vision_config, prefix=maybe_prefix(prefix, "visual")) + if inkling_vision_enabled(config) + else None + ) + self.audio = ( + InklingAudio(config.audio_config, prefix=maybe_prefix(prefix, "audio")) + if inkling_audio_enabled(config) + else None + ) + + self._build(vllm_config, config.text_config, prefix) + + # -- multimodal embedding ------------------------------------------- + + def _process_image_input( + self, pixel_values: Any, num_patches: Any + ) -> tuple[torch.Tensor, ...]: + assert self.visual is not None + # pixel_values is a list (per item) of [P_i, 2, P, P, 3] tensors, + # or a single concatenated tensor. Normalize to a flat batch, run the + # tower once, then split back per item. + if isinstance(pixel_values, (list, tuple)): + if not pixel_values: + return () + sizes = [int(p.shape[0]) for p in pixel_values] + patches = torch.cat(list(pixel_values), dim=0) + else: + patches = pixel_values + sizes = self._sizes_from(num_patches, patches.shape[0]) + + patches = patches.to(device=self.visual.device, dtype=self.visual.dtype) + embeds = self.visual(patches) # [total_patches, D] + return tuple(embeds.split(sizes)) + + def _process_audio_input( + self, input_audio_features: Any, num_audio_tokens: Any + ) -> tuple[torch.Tensor, ...]: + assert self.audio is not None + if isinstance(input_audio_features, (list, tuple)): + if not input_audio_features: + return () + sizes = [int(d.shape[0]) for d in input_audio_features] + dmel = torch.cat(list(input_audio_features), dim=0) + else: + dmel = input_audio_features + sizes = self._sizes_from(num_audio_tokens, dmel.shape[0]) + + dmel = dmel.to(device=self.audio.device) + embeds = self.audio(dmel) # [total_frames, D] + return tuple(embeds.split(sizes)) + + @staticmethod + def _sizes_from(counts: Any, total: int) -> list[int]: + if counts is None: + return [total] + if isinstance(counts, torch.Tensor): + return [int(c) for c in counts.flatten().tolist()] + if isinstance(counts, (list, tuple)): + flat: list[int] = [] + for c in counts: + flat.append(int(c.item()) if isinstance(c, torch.Tensor) else int(c)) + return flat + return [int(counts)] + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + # Iterate modalities in a stable order so the returned per-item tensors + # line up with their appearance order; the positional merge in + # embed_input_ids handles actual placement. + pixel_values = kwargs.get("pixel_values") + num_patches = kwargs.get("num_patches") + input_audio_features = kwargs.get("input_audio_features") + num_audio_tokens = kwargs.get("num_audio_tokens") + + embeddings: tuple[torch.Tensor, ...] = () + if pixel_values is not None and self.visual is not None: + embeddings += self._process_image_input(pixel_values, num_patches) + if input_audio_features is not None and self.audio is not None: + embeddings += self._process_audio_input( + input_audio_features, num_audio_tokens + ) + return embeddings + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + # Override the base's 1-arg embed_input_ids: the runner calls this 3-arg + # signature for multimodal models. Text embeddings come from the shared + # backbone (which applies embed_norm); MM embeddings are scattered in. + from vllm.model_executor.models.utils import _merge_multimodal_embeddings + + # Placeholder ids use unused vocabulary slots and these positions are + # overwritten by MM embeds below. + inputs_embeds = self.model.embed_input_ids(input_ids) + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + assert is_multimodal is not None + return _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + def get_language_model(self) -> nn.Module: + # This class IS the causal LM (the towers are side branches), so the + # language model is self — callers expect a module exposing ``.model`` + # / ``.lm_head`` (e.g. the MTP/eagle loader shares embeddings via + # ``get_language_model().model.embed_tokens``). + return self + + +# =========================================================================== +# Weight loading +# =========================================================================== + + +_MOE_EXPERT_WEIGHT_RE = re.compile( + r"^(?P<mlp>.*\.mlp)\.(?P<rest>(?:shared_)?experts\..+)$" +) + + +def _is_peft_adapter_weight(name: str) -> bool: + return ".lora_A." in name or ".lora_B." in name + + +def _load_inkling_weights( + module: nn.Module, + weights: Iterable[tuple[str, torch.Tensor]], + config: InklingModelConfig, +) -> set[str]: + moe_modules = { + name: mod for name, mod in module.named_modules() if isinstance(mod, InklingMoE) + } + loaded: set[str] = set() + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + local_ids = set(config.local_layer_ids) + + def _iter_loadable_weights() -> Iterable[tuple[str, torch.Tensor]]: + for name, weight in module.hf_to_vllm_mapper.apply(weights): + # LightSeek's MXFP4 conversion bundles a quantized copy of the + # standalone PEFT adapter in the base safetensors index. These are + # not base-model parameters; adapters remain opt-in at serving. + if _is_peft_adapter_weight(name): + continue + shard_id = getattr(weight, "shard_id", None) + # Replicate K/V conv-free GQA heads when tp_size > num_kv_heads. + if ( + shard_id in (1, 2) + and name.endswith(".attn.qkvr.weight") + and weight.shape[0] > 0 + ): + lid = _layer_id(name) + if lid is not None: + is_local = lid in local_ids + n_kv = ( + config.swa_num_key_value_heads + if is_local + else config.num_key_value_heads + ) + head_dim = config.swa_head_dim if is_local else config.head_dim + if tp_size > n_kv and weight.shape[0] == n_kv * head_dim: + kv_idx = (tp_rank * n_kv) // tp_size + weight = weight.narrow(0, kv_idx * head_dim, head_dim) + weight.shard_id = shard_id + + # MoE expert tensors (fused stacked, routed + shared sink): translate + # the checkpoint layout to per-expert FusedMoE loads. + moe_match = _MOE_EXPERT_WEIGHT_RE.match(name) + if moe_match is not None and moe_match.group("mlp") in moe_modules: + moe = moe_modules[moe_match.group("mlp")] + for rel in moe.load_expert_weight(moe_match.group("rest"), weight): + loaded.add(f"{moe_match.group('mlp')}.{rel}") + continue + + yield name, weight + + loader = AutoWeightsLoader(module, skip_prefixes=["model.mtp."]) + loaded |= loader.load_weights(_iter_loadable_weights()) + + # Post-load MoE fixups (default input scales, zeroed EP-padding experts). + for moe_name, moe in moe_modules.items(): + for rel in moe.finalize_load(): + loaded.add(f"{moe_name}.{rel}") + return loaded + + +EntryClass = [InklingForCausalLM, InklingForConditionalGeneration] diff --git a/vllm/models/inkling/amd/moe.py b/vllm/models/inkling/amd/moe.py new file mode 100644 index 00000000000..90b00a3e5ae --- /dev/null +++ b/vllm/models/inkling/amd/moe.py @@ -0,0 +1,693 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling mixture-of-experts on vLLM's FusedMoE abstraction. + +Overfit to the served checkpoint: sigmoid gate (+ selection bias) top-k over +the routed experts, log-sigmoid renormalization over the k routed + S shared +"sink" logits, scaled by route_scale * global_scale. The routed top-k goes +through vLLM's FusedMoE (which handles TP/EP); the sink experts run in +:class:`InklingSinkExperts` -- replicated across EP ranks (every token +activates every sink) and always bf16 (the checkpoint excludes every +``shared_experts`` from quantization). + +MXFP4 routed experts reuse vLLM's Quark OCP MXFP4 fused-MoE method; excluded +(bf16) layers fall back to the unquantized method. The checkpoint's fused +stacked tensors (interleaved gate/up rows, ``.scale`` / ``.scale2`` / +``.input_amax`` aux tensors) are translated to the standard per-expert loads +in :meth:`InklingMoE.load_expert_weight`. +""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +import torch +from torch import nn +from torch.nn.parameter import Parameter + +import vllm.envs as envs +from vllm.config import get_current_vllm_config +from vllm.distributed import ( + get_dp_group, + get_pcp_group, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.model_executor.kernels.linear.cute_dsl import ll_bf16 +from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.utils import set_weight_attrs +from vllm.platforms import current_platform +from vllm.triton_utils import tl, tldevice, triton +from vllm.utils.multi_stream_utils import maybe_execute_in_parallel +from vllm.utils.torch_utils import aux_stream + +from ..configs import InklingModelConfig + +if TYPE_CHECKING: + from vllm.model_executor.layers.fused_moe.routed_experts import ( + RoutedExperts, + ) + +# --------------------------------------------------------------------------- +# Gate / expert selection +# --------------------------------------------------------------------------- + +_INKLING_LL_BF16_MAX_TOKENS = 64 +_MXFP4_INPUT_SCALE_DENOMINATOR = torch.finfo(torch.float8_e4m3fn).max * 6.0 + + +def _linear_with_fp32_out(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + leading = list(x.shape[:-1]) + flat = x.flatten(0, -2) + if ( + flat.shape[0] <= _INKLING_LL_BF16_MAX_TOKENS + and flat.dtype == torch.bfloat16 + and weight.dtype == torch.bfloat16 + and flat.is_cuda + and flat.is_contiguous() + and weight.is_contiguous() + and flat.shape[1] % 8 == 0 + and current_platform.has_device_capability(90) + and ll_bf16.is_available() + ): + out = ll_bf16.ll_bf16_gemm(flat, weight) + else: + out = torch.mm(flat, weight.T, out_dtype=torch.float32) + return out.view(*leading, weight.shape[0]) + + +@triton.jit(do_not_specialize=["T", "route_scale"]) +def _inkling_gate_select_kernel( + logits_ptr, # [T, G] fp32 gate logits (stride_logits_0 may include pad) + bias_ptr, # [R] fp32 selection bias (or 0 ptr if HAS_BIAS=False) + global_scale_ptr, # [1] fp32 (or unused if HAS_GSCALE=False) + ids_ptr, # [T, K + S] int32 out: selected expert ids + weights_ptr, # [T, K + S] fp32 out: renormalized weights + route_scale, + T, + G: tl.constexpr, # total gate experts (routed + shared) + stride_logits_0, + R: tl.constexpr, # routed experts + K: tl.constexpr, # top-k routed + S: tl.constexpr, # shared (sink) experts + HAS_BIAS: tl.constexpr, + HAS_GSCALE: tl.constexpr, + BLOCK_G: tl.constexpr, +): + pid = tl.program_id(0).to(tl.int64) + if pid >= T: + return + offs = tl.arange(0, BLOCK_G) + mask_r = offs < R + logits = tl.load( + logits_ptr + pid * stride_logits_0 + offs, + mask=offs < G, + other=float("-inf"), + ).to(tl.float32) + + # Selection scores: sigmoid(routed logits) (+ bias), non-routed lanes -inf. + sel = tl.where(mask_r, tl.sigmoid(logits), float("-inf")) + if HAS_BIAS: + bias = tl.load(bias_ptr + offs, mask=mask_r, other=0.0).to(tl.float32) + sel = tl.where(mask_r, sel + bias, float("-inf")) + + scale = route_scale + if HAS_GSCALE: + scale = scale * tl.load(global_scale_ptr).to(tl.float32) + + # Iterative top-K (K is small); argmax tie-breaks to the lowest index + # (stable ordering). + A: tl.constexpr = K + S + offs_a = tl.arange(0, A) + top_ids = tl.zeros([A], dtype=tl.int32) + active = tl.zeros([A], dtype=tl.float32) + for kk in tl.static_range(K): + idx = tl.argmax(sel, axis=0).to(tl.int32) + raw = tl.max(tl.where(offs == idx, logits, float("-inf")), axis=0) + top_ids = tl.where(offs_a == kk, idx, top_ids) + active = tl.where(offs_a == kk, raw, active) + sel = tl.where(offs == idx, float("-inf"), sel) + if S > 0: + # Shared sink logits sit at the tail of the gate output; their expert + # ids continue after the routed range (R + j). + for jj in tl.static_range(S): + raw = tl.max(tl.where(offs == R + jj, logits, float("-inf")), axis=0) + top_ids = tl.where(offs_a == K + jj, tl.full([], R + jj, tl.int32), top_ids) + active = tl.where(offs_a == K + jj, raw, active) + + # Log-sigmoid renormalization over the K + S active logits. + abs_l = tl.abs(active) + min_l = tl.minimum(active, 0.0) + log_probs = min_l - tldevice.log1p(tldevice.exp(-abs_l)) + max_lp = tl.max(log_probs, axis=0) + exp_shifted = tldevice.exp(log_probs - max_lp) + sum_exp = tl.sum(exp_shifted, axis=0) + weights = exp_shifted / sum_exp * scale + + tl.store(ids_ptr + pid * A + offs_a, top_ids) + tl.store(weights_ptr + pid * A + offs_a, weights) + + +def inkling_gate_select( + logits: torch.Tensor, # [T, >=G] fp32 (rows may carry GEMM padding) + n_gate_experts: int, + n_routed_experts: int, + topk: int, + n_shared_experts: int, + bias: torch.Tensor | None, + route_scale: float, + global_scale: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Sigmoid + bias + top-k + log-sigmoid renorm; returns (weights, ids).""" + assert logits.dtype == torch.float32 + tokens = logits.shape[0] + active = topk + n_shared_experts + topk_ids = torch.empty((tokens, active), dtype=torch.int32, device=logits.device) + topk_weights = torch.empty( + (tokens, active), dtype=torch.float32, device=logits.device + ) + if tokens == 0: + return topk_weights, topk_ids + _inkling_gate_select_kernel[(tokens,)]( + logits, + bias if bias is not None else logits, + global_scale if global_scale is not None else logits, + topk_ids, + topk_weights, + route_scale, + tokens, + n_gate_experts, + logits.stride(0), + n_routed_experts, + topk, + n_shared_experts, + HAS_BIAS=bias is not None, + HAS_GSCALE=global_scale is not None, + BLOCK_G=triton.next_power_of_2(n_gate_experts), + ) + return topk_weights, topk_ids + + +class InklingGate(nn.Module): + """Sigmoid gate with selection bias, log-sigmoid renorm after top-k, and + global scale (the served checkpoint's only configuration).""" + + def __init__( + self, + d_model: int, + n_routed_experts: int, + n_shared_experts: int, + experts_per_token: int, + route_scale: float, + *, + use_global_scale: bool = False, + use_gate_bias: bool = False, + ) -> None: + super().__init__() + self.n_routed_experts = n_routed_experts + self.n_shared_experts = n_shared_experts + self.n_total_experts = n_routed_experts + n_shared_experts + self.topk = experts_per_token + self.route_scale = route_scale + + padded_experts = self.n_total_experts + (-self.n_total_experts) % 8 + self.weight = Parameter( + torch.empty(padded_experts, d_model), requires_grad=False + ) + set_weight_attrs(self.weight, {"weight_loader": self._load_weight}) + self.global_scale: Parameter | None + if use_global_scale: + self.global_scale = Parameter( + torch.empty(1, dtype=torch.float32), requires_grad=False + ) + else: + self.global_scale = None + self.bias: Parameter | None + if use_gate_bias: + self.bias = Parameter( + torch.empty(n_routed_experts, dtype=torch.float32), + requires_grad=False, + ) + else: + self.bias = None + + @staticmethod + def _load_weight(param: Parameter, loaded_weight: torch.Tensor) -> None: + param.data.zero_() + param.data[: loaded_weight.shape[0]].copy_(loaded_weight) + + def compute_logits(self, x: torch.Tensor) -> torch.Tensor: + """fp32 gate logits [T, n_total_experts + pad] (pad columns are junk).""" + return _linear_with_fp32_out(x, self.weight) + + def select_experts( + self, gating_output: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Full selection: (weights, ids) of [T, K + S]. The first K entries + are the routed top-k; the S trailing entries are the sink gammas.""" + return inkling_gate_select( + gating_output, + self.n_total_experts, + self.n_routed_experts, + self.topk, + self.n_shared_experts, + self.bias, + self.route_scale, + self.global_scale, + ) + + +# --------------------------------------------------------------------------- +# MoE layer +# --------------------------------------------------------------------------- + + +def _inkling_moe_ep_size() -> int: + """EP size the FusedMoE layer will run with (mirrors + FusedMoEParallelConfig.make: experts shard over tp * dp * pcp when + expert parallelism is enabled).""" + parallel_config = get_current_vllm_config().parallel_config + if not parallel_config.enable_expert_parallel: + return 1 + world = ( + get_tensor_model_parallel_world_size() + * get_dp_group().world_size + * get_pcp_group().world_size + ) + return world if world > 1 else 1 + + +class InklingSinkExperts(nn.Module): + """Shared "sink" experts with per-token gammas, in bf16. + + Replicated across EP ranks (every token activates every sink, so + EP-sharding them would hotspot the owning rank) and TP-sharded on the + intermediate dim so the output remains a TP-partial sum like the routed + output. The sinks are always bf16 (the checkpoint excludes every + ``shared_experts`` from quantization): the experts concatenate into two + plain dense GEMMs with the fused sink epilogue between them. + """ + + def __init__( + self, n_experts: int, d_model: int, d_mlp: int, *, prefix: str = "" + ) -> None: + super().__init__() + self.n_experts = n_experts + tp_size = get_tensor_model_parallel_world_size() + self.tp_rank = get_tensor_model_parallel_rank() + intermediate_pp = d_mlp // tp_size + self.w13_weight = Parameter( + torch.empty(n_experts, 2 * intermediate_pp, d_model), + requires_grad=False, + ) + self.w2_weight = Parameter( + torch.empty(d_model, n_experts * intermediate_pp), + requires_grad=False, + ) + self._unit: torch.Tensor | None = None + + def load_weight(self, key: str, weight: torch.Tensor) -> list[str]: + """Load one checkpoint sink tensor (stacked over the S experts).""" + if key == "w13_weight": + if weight.shape != self.w13_weight.shape: + shard = self.w13_weight.shape[1] + weight = weight.narrow(1, self.tp_rank * shard, shard) + self.w13_weight.data.copy_(weight) + return [key] + + assert key == "w2_weight" + shard = self.w2_weight.shape[1] // self.n_experts + shard_start = 0 if weight.shape[2] == shard else self.tp_rank * shard + for expert_idx, expert_weight in enumerate(weight): + local_weight = expert_weight.narrow(1, shard_start, shard) + start = expert_idx * shard + self.w2_weight.data[:, start : start + shard].copy_(local_weight) + return [key] + + def forward(self, x: torch.Tensor, gammas: torch.Tensor) -> torch.Tensor: + """``sum_e gammas[:, e] * MLP_e(x)`` (TP-partial along d_mlp).""" + from .ops import sink_silu_mul_epilogue + + # One GEMM over the experts' stacked w13 (a view), fused epilogue, + # then one GEMM whose K-reduction over the K-concatenated w2 performs + # the expert sum. + if self._unit is None or self._unit.device != x.device: + self._unit = torch.ones( + self.n_experts, dtype=torch.float32, device=x.device + ) + raw = x @ self.w13_weight.view(-1, x.shape[-1]).T # (T, S*2F) + h = sink_silu_mul_epilogue( + raw, self._unit, gammas, self._unit, self.n_experts, x.dtype + ) + return h @ self.w2_weight.T # (T, D) + + +class InklingSinkExpertsLinear(nn.Module): + """LoRA-capable implementation of the Inkling sink experts.""" + + def __init__( + self, + n_experts: int, + d_model: int, + d_mlp: int, + *, + prefix: str = "", + ) -> None: + super().__init__() + from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + RowParallelLinear, + ) + + self.n_experts = n_experts + self.d_mlp = d_mlp + total = n_experts * d_mlp + self.w13 = MergedColumnParallelLinear( + input_size=d_model, + output_sizes=[total, total], + bias=False, + prefix=f"{prefix}.w13", + ) + self.w2 = RowParallelLinear( + input_size=total, + output_size=d_model, + bias=False, + reduce_results=False, + prefix=f"{prefix}.w2", + ) + self._w2_input_pp = self.w2.input_size_per_partition + self._col_expert: torch.Tensor | None = None + + def _gamma_expand(self, gammas: torch.Tensor) -> torch.Tensor: + if self._col_expert is None or self._col_expert.device != gammas.device: + local = self._w2_input_pp + start = get_tensor_model_parallel_rank() * local + cols = torch.arange(start, start + local, device=gammas.device) + self._col_expert = (cols // self.d_mlp).long() + return gammas[:, self._col_expert] + + def load_weight(self, key: str, weight: torch.Tensor) -> list[str]: + if key == "w13_weight": + d_model = weight.shape[-1] + gate = weight[:, 0::2, :].reshape(-1, d_model).contiguous() + up = weight[:, 1::2, :].reshape(-1, d_model).contiguous() + self.w13.weight_loader(self.w13.weight, gate, 0) + self.w13.weight_loader(self.w13.weight, up, 1) + return ["w13.weight"] + w = weight.permute(1, 0, 2).reshape(weight.shape[1], -1).contiguous() + self.w2.weight_loader(self.w2.weight, w) + return ["w2.weight"] + + def forward(self, x: torch.Tensor, gammas: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.w13(x) + gate, up = gate_up.chunk(2, dim=-1) + hidden_states = torch.nn.functional.silu(gate) * up + hidden_states = (hidden_states * self._gamma_expand(gammas)).to(x.dtype) + output, _ = self.w2(hidden_states) + return output + + +class InklingMoE(nn.Module): + def __init__( + self, + config: InklingModelConfig, + *, + prefix: str = "", + quant_config: QuantizationConfig | None = None, + ) -> None: + super().__init__() + # Overfit to the served checkpoint: sigmoid gate renormalized after + # top-k, shared sink experts, interleaved gate/up checkpoint rows. + assert config.gate_activation == "sigmoid" and config.norm_after_topk + assert config.n_shared_experts > 0 and config.shared_expert_sink + assert config.inference_moe_w13_interleaved + n_routed = config.n_routed_experts + n_shared = config.n_shared_experts + self.n_routed_experts = n_routed + self.gate = InklingGate( + d_model=config.hidden_size, + n_routed_experts=n_routed, + n_shared_experts=n_shared, + experts_per_token=config.num_experts_per_tok, + route_scale=config.route_scale, + use_global_scale=config.use_global_scale, + use_gate_bias=config.use_gate_bias, + ) + + # TRTLLM MoE kernels assume equal, contiguous per-rank expert slabs + # (local_expert_offset = ep_rank * local_num_experts), so pad the + # expert count to a multiple of the EP size. A no-op for the usual + # power-of-two EP sizes (n_routed is a power of two). + num_experts = n_routed + (-n_routed) % _inkling_moe_ep_size() + + # The released MXFP4 checkpoint keeps the first routed-expert layer + # in bf16 and lists its two expert weights explicitly in Quark's + # exclusion list. FusedMoE asks for a quant method at the module + # prefix, so an exact weight exclusion would otherwise be missed and + # the bf16 tensors would be loaded into MXFP4 parameters. + routed_quant_config = quant_config + if quant_config is not None: + excluded = set(getattr(quant_config, "quant_config", {}).get("exclude", ())) + expert_prefix = f"{prefix}.experts" + routed_weights = { + f"{expert_prefix}.w13_weight", + f"{expert_prefix}.w2_weight", + } + if routed_weights <= excluded: + routed_quant_config = None + + self.experts = FusedMoE( + num_experts=num_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + renormalize=False, + quant_config=routed_quant_config, + prefix=f"{prefix}.experts", + custom_routing_function=self._select_routed, + router_logits_dtype=torch.float32, + activation="silu", + ) + # The decoder layer reduce-scatters the MoE delta into the sconv + # stream itself (RS -> shard sconv -> AG); the runner must return the + # per-rank partial sum instead of all-reducing. + self.experts.moe_config.skip_final_all_reduce = True + + self._routed_sel: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None + sink_experts_cls = ( + InklingSinkExpertsLinear + if get_current_vllm_config().lora_config is not None + else InklingSinkExperts + ) + self.sink_experts = sink_experts_cls( + n_experts=n_shared, + d_model=config.hidden_size, + d_mlp=config.intermediate_size, + prefix=f"{prefix}.shared_experts", + ) + + # Sink chain overlaps the routed MoE call on the aux stream for + # decode-sized batches (same pattern as the runner's SharedExperts + # multi-stream overlap). The routed GEMM runs on the default stream and + # the sink chain on the aux stream, joined via these two events by + # ``maybe_execute_in_parallel``. + self._sink_stream: torch.cuda.Stream | None = aux_stream() + self._sink_events = (torch.cuda.Event(), torch.cuda.Event()) + + def _select_routed( + self, + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + """FusedMoE ``custom_routing_function``: the routed top-k slice of the + full (routed + sink) selection. + + forward() stashes its selection (keyed by logits identity) so the + gate select runs once per layer; the fallback covers paths where the + runner re-derives the logits (e.g. naive DP dispatch). + """ + del hidden_states, renormalize + assert topk == self.gate.topk + cached = self._routed_sel + self._routed_sel = None + if cached is not None and cached[0] is gating_output: + return cached[1], cached[2] + weights, ids = self.gate.select_experts(gating_output) + return weights[:, :topk].contiguous(), ids[:, :topk].contiguous() + + def forward(self, x: torch.Tensor) -> torch.Tensor | None: + router_logits = self.gate.compute_logits(x) + num_tokens = x.shape[0] + # One gate select per layer: the routed slice is stashed for the + # routing function inside the FusedMoE op; the sink gammas are the + # trailing columns. + k = self.gate.topk + weights, ids = self.gate.select_experts(router_logits) + self._routed_sel = ( + router_logits, + weights[:, :k].contiguous(), + ids[:, :k].contiguous(), + ) + gammas = weights[:, k:] + + out, sink_out = maybe_execute_in_parallel( + lambda: self.experts(hidden_states=x, router_logits=router_logits), + lambda: self.sink_experts(x, gammas), + self._sink_events[0], + self._sink_events[1], + self._sink_stream + if num_tokens <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD + else None, + ) + self._routed_sel = None + + return out.add_(sink_out) + + # -- weight loading ---------------------------------------------------- + + def _local_expert_slots(self) -> dict[int, int]: + """Global expert id -> local slot for this rank's expert partition.""" + manager = self.experts.routed_experts.expert_map_manager + if manager.expert_map is None: + return {g: g for g in range(manager.global_num_experts)} + emap = manager.expert_map.tolist() + return {g: slot for g, slot in enumerate(emap) if slot >= 0} + + def load_expert_weight(self, name: str, weight: torch.Tensor) -> list[str]: + """Load one checkpoint expert tensor. + + ``name`` is relative to the mlp module: ``experts.<t>`` (routed + stack) or ``shared_experts.shared_<t>`` (sink experts). Returns the + loaded param names (relative to this module). + """ + if name.startswith("shared_experts."): + key = name.split(".", 1)[1].replace("shared_", "", 1) + return [ + f"sink_experts.{p}" for p in self.sink_experts.load_weight(key, weight) + ] + + experts: RoutedExperts = self.experts.routed_experts + key = name.split(".", 1)[1] + + # original_shape is unused by the vLLM serving layout. + if key.endswith(".original_shape"): + return [] + if key.endswith(".input_amax"): + projection = "w13" if key.startswith("w13") else "w2" + amax = float(weight.max()) + assert math.isfinite(amax) and amax > 0, ( + f"bad {projection} input_amax: {amax}" + ) + input_scale = getattr(experts, f"{projection}_input_scale") + input_scale.data.fill_(amax / _MXFP4_INPUT_SCALE_DENOMINATOR) + return [f"experts.routed_experts.{projection}_input_scale"] + + # Quark's OCP MXFP4 converter stores block scales as a two-dimensional + # tensor with the expert and projection-row dimensions flattened: + # w13: [E * 2I, H / 32], w2: [E * H, I / 32]. + # Restore the expert dimension before applying the EP/TP slicing used + # for both packed weights and scales. + if key.endswith("_weight_scale") and weight.ndim == 2: + if weight.shape[0] % self.n_routed_experts != 0: + raise ValueError( + f"cannot unflatten {name} with shape {tuple(weight.shape)} " + f"over {self.n_routed_experts} experts" + ) + weight = weight.view( + self.n_routed_experts, + weight.shape[0] // self.n_routed_experts, + weight.shape[1], + ) + + param = getattr(experts, key) + slots = self._local_expert_slots() + gids = sorted(slots) + lids = [slots[g] for g in gids] + tp_rank = experts.moe_config.moe_parallel_config.tp_rank + tp_size = experts.moe_config.moe_parallel_config.tp_size + + if key.endswith("_scale_2"): + # Per-expert scalars, vectorized over the local experts. The + # fused w13 param carries one slot per gate/up half. + vals = weight[gids].float().to(param.device) + param.data[lids] = vals[:, None] if param.data.ndim == 2 else vals + elif key.startswith("w13"): + # Checkpoint w13 rows are interleaved [g0, u0, g1, u1, ...]; the + # fused param layout is [w1(gate); w3(up)]. The TP-local rows form + # one contiguous slab of the interleaved tensor, so upload just + # that slab (a single bounded synchronous H2D; pre-uploading whole + # untrimmed tensors pins the mmap pages of the entire checkpoint + # and OOMs the host) and de-interleave on device. + dst_half = param.shape[1] // 2 + if weight.shape[1] % (2 * tp_size) != 0: + raise ValueError( + f"cannot TP-shard {name} with shape {tuple(weight.shape)} " + f"over {tp_size} ranks" + ) + logical_half = weight.shape[1] // (2 * tp_size) + if logical_half > dst_half: + raise ValueError( + f"checkpoint shard for {name} has {logical_half} rows per " + f"gate/up half, but destination only has {dst_half}" + ) + for gid, lid in slots.items(): + slab = weight[gid].narrow( + 0, + tp_rank * 2 * logical_half, + 2 * logical_half, + ) + slab = slab.to(param.device) + param.data[lid, :logical_half].copy_(slab[0::2]) + param.data[ + lid, + dst_half : dst_half + logical_half, + ].copy_(slab[1::2]) + else: + # w2: shard the packed intermediate (last) dim. AITER rounds the + # destination intermediate width to 256, so derive the logical TP + # slice from the checkpoint and leave the destination tail at its + # initialized padding value (zero for weights, one for scales). + if weight.shape[2] % tp_size != 0: + raise ValueError( + f"cannot TP-shard {name} with shape {tuple(weight.shape)} " + f"over {tp_size} ranks" + ) + shard = weight.shape[2] // tp_size + if shard > param.shape[2]: + raise ValueError( + f"checkpoint shard for {name} has width {shard}, but " + f"destination only has {param.shape[2]}" + ) + for gid, lid in slots.items(): + param.data[lid, :, :shard].copy_( + weight[gid].narrow(1, tp_rank * shard, shard) + ) + return [f"experts.routed_experts.{key}"] + + def finalize_load(self) -> list[str]: + """Post-load fixups for zeroed padding experts.""" + experts = self.experts.routed_experts + out: list[str] = [] + # Zero the EP-alignment padding experts (if any) so their + # (never-routed) slots hold defined values. + slots = self._local_expert_slots() + for gid in range(self.n_routed_experts, experts.global_num_experts): + lid = slots.get(gid) + if lid is None: + continue + for pname in ( + "w13_weight", + "w2_weight", + "w13_weight_scale", + "w2_weight_scale", + "w13_weight_scale_2", + "w2_weight_scale_2", + ): + p = getattr(experts, pname, None) + if p is not None: + p.data[lid].zero_() + return out diff --git a/vllm/models/inkling/amd/mtp.py b/vllm/models/inkling/amd/mtp.py new file mode 100644 index 00000000000..34cde020ed3 --- /dev/null +++ b/vllm/models/inkling/amd/mtp.py @@ -0,0 +1,410 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling MTP (Multi-Token Prediction) draft model (NVIDIA). + +Implements the first MTP depth from the reference ``mtp_model.py`` shipped with +the checkpoint. It owns ``hidden_norm`` / ``embed_norm`` RMSNorms, a ``2H -> H`` +input projection, and a full Inkling transformer block with a dense bf16 MLP. + +The draft shares the target's token embedding table and LM head +(``load_eagle_model`` wires those references) and applies the backbone +``embed_norm`` on top: the depth layers were trained on the same normed +embeddings the backbone consumes (their own ``embed_norm`` weights are +near-identity trims, unlike the backbone's whitening ``embed_norm``). +""" + +from __future__ import annotations + +from collections.abc import Iterable + +import regex as re +import torch +from torch import nn + +from vllm.config import VllmConfig +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.model_loader.mtp_validation import ( + is_mtp_completeness_check_enabled, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.utils import maybe_prefix +from vllm.sequence import IntermediateTensors + +from ..configs import InklingModelConfig +from .layernorm import InklingRMSNorm +from .model import InklingDecoderLayer, InklingReplicatedEmbedding +from .ops.norm import embed_dual_rmsnorm_cat, embed_rmsnorm + +# Checkpoint attention projections (wq_du/wk_dv/wv_dv/wr_du) -> fused qkvr. +# Mirrors the backbone's hf_to_vllm_mapper.orig_to_new_stacked; kept as a +# local (pname, wname, shard) list since the MTP loader remaps by hand. +_ATTENTION_PARAMS_MAPPING = [ + ("qkvr", "wq_du", 0), + ("qkvr", "wk_dv", 1), + ("qkvr", "wv_dv", 2), + ("qkvr", "wr_du", 3), +] + + +def _mtp_depth_from_name(name: str) -> int | None: + m = re.search(r"\.mtp\.layers\.(\d+)\.", name) + return int(m.group(1)) if m else None + + +class InklingMTPDepthLayer(nn.Module): + """One MTP depth: norm both inputs, fuse (2H->H), run a Inkling block.""" + + def __init__(self, config: InklingModelConfig, prefix: str, is_local: bool) -> None: + super().__init__() + self.hidden_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.embed_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.input_proj = ReplicatedLinear( + config.hidden_size * 2, + config.hidden_size, + bias=False, + return_bias=False, + prefix=f"{prefix}.input_proj", + ) + # A force-dense-MLP bf16 block; ``is_local`` selects sliding-window vs + # full attention (the swa_* head config and sliding_window window) to + # match this depth's checkpoint transformer_block weights. + self.transformer_block = InklingDecoderLayer( + config, + layer_id=0, + is_local=is_local, + quant_config=None, + prefix=f"{prefix}.transformer_block", + force_dense_mlp=True, + ) + + def forward(self, combined: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + # ``combined`` is the fused-normed [rmsnorm(hidden) | embed_norm(emb)] + # input, built by InklingMultiTokenPredictor.fused_input_cat in one launch. + hidden = self.input_proj(combined) + # The short conv self-fetches its paged SWA-cache metadata from the + # forward context (via its conv_owner prefix); no conv_meta to thread. + return self.transformer_block(positions, hidden) + + +class InklingMultiTokenPredictor(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + assert vllm_config.speculative_config is not None + config: InklingModelConfig = ( + vllm_config.speculative_config.draft_model_config.hf_config + ) + self.config = config + if vllm_config.speculative_config.num_speculative_tokens != 1: + raise ValueError( + "Inkling MTP currently supports exactly one speculative token" + ) + self.chain_hidden_post_norm = config.chain_hidden_post_norm + local_ids = set(config.local_layer_ids) + self.layers = nn.ModuleDict( + {"0": InklingMTPDepthLayer(config, f"{prefix}.layers.0", 0 in local_ids)} + ) + self.chain_norm = ( + InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if self.chain_hidden_post_norm + else None + ) + # The target's raw token embedding (pre embed_norm), attached by + # load_eagle_model. Never materialized here: building our own + # replicated copy would transiently double the 2.3 GiB table. + self.embed_tokens: InklingReplicatedEmbedding = None # type: ignore[assignment] + # The depth layers consume the *backbone-normed* embedding + # (embed_norm(embed(ids))), not the raw one: mtp embed_norm weights + # are near-identity (trained on already-normalized inputs), and + # feeding raw embeddings drops MTP1 acceptance from ~0.85 to ~0.70. + # Weight loaded from the target's embed_norm.weight; gated like the + # target's InklingModel.embed_norm. + self.backbone_embed_norm = ( + InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if config.use_embed_norm + else None + ) + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: object | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + """Draft-prefill embedding: fused gather + backbone embed_norm, then + the target's tower embeddings scattered in unnormed (the backbone + convention — MM embeds are merged after embed_norm).""" + norm = self.backbone_embed_norm + embeds = embed_rmsnorm( + input_ids, + self.embed_tokens.weight, + norm.weight if norm is not None else None, + norm.variance_epsilon if norm is not None else 0.0, + ) + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: # type: ignore[arg-type] + return embeds + from vllm.model_executor.models.utils import _merge_multimodal_embeddings + + assert is_multimodal is not None + return _merge_multimodal_embeddings( + inputs_embeds=embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + def fused_input_cat( + self, + layer: InklingMTPDepthLayer, + previous_hidden: torch.Tensor, + input_ids: torch.Tensor, + inputs_embeds: torch.Tensor | None, + ) -> torch.Tensor: + """The depth layer's [rmsnorm(hidden) | embed_norm(embed)] input in one + launch: embedding row gather + the backbone embed_norm + the depth + embed_norm chain on one side, hidden_norm on the other, written + straight into the cat buffer.""" + hidden_w = layer.hidden_norm.weight + embed_w = layer.embed_norm.weight + eps = layer.hidden_norm.variance_epsilon + if inputs_embeds is not None: + # Draft prefill with target-merged MM embeddings (already + # backbone-normed via embed_input_ids); only the depth embed_norm + # remains. + return embed_dual_rmsnorm_cat( + previous_hidden, hidden_w, embed_w, eps, embeds=inputs_embeds + ) + return embed_dual_rmsnorm_cat( + previous_hidden, + hidden_w, + embed_w, + eps, + input_ids=input_ids, + embed_table=self.embed_tokens.weight, + pre_norm_weight=( + self.backbone_embed_norm.weight + if self.backbone_embed_norm is not None + else None + ), + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + # The draft's short conv is a paged SWA-cache layer (its conv_owner is + # auto-enumerated as a draft attention layer); its per-token metadata is + # built by the speculator's build_attn_metadata and read from the + # forward context, so nothing extra is threaded here. + if spec_step_idx != 0: + raise ValueError("Inkling MTP only supports spec_step_idx=0") + layer = self.layers["0"] + combined = self.fused_input_cat( + layer, previous_hidden_states, input_ids, inputs_embeds + ) + hidden = layer(combined, positions) + if self.chain_norm is not None: + hidden = self.chain_norm(hidden) + return hidden + + +class InklingMTP(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + assert vllm_config.speculative_config is not None + config: InklingModelConfig = ( + vllm_config.speculative_config.draft_model_config.hf_config + ) + self.config = config + self.model = InklingMultiTokenPredictor( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + # The target's (vocab-sharded) LM head, attached by load_eagle_model; + # never materialized here (same reasoning as model.embed_tokens). + self.lm_head: ParallelLMHead = None # type: ignore[assignment] + self.logits_processor = LogitsProcessor( + config.padded_vocab_size, + org_vocab_size=config.vocab_size, + soft_cap=config.final_logit_softcapping, + ) + self._logits_zero: torch.Tensor | None = None + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: object | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.model.embed_input_ids( + input_ids, multimodal_embeddings, is_multimodal=is_multimodal + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + hidden_states: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + return self.model( + input_ids, + positions, + hidden_states, + inputs_embeds, + spec_step_idx, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor | None: + # The MTP shares the base model's LM head, which is trained on + # ``hidden / mup``-scaled inputs, so apply the same mup scaling here + # for a matching logit scale — folded into the lm_head GEMM alpha + # (fp32 epilogue) like the target's compute_logits. (Argmax-invariant + # for greedy draft sampling, but it matters for the gumbel sampling + # distribution at temperature > 0.) + mup = self.config.logits_mup_width_multiplier + if not mup: + return self.logits_processor(self.lm_head, hidden_states) + assert self.logits_processor.soft_cap is None + assert self.logits_processor.scale == 1.0 + w = self.lm_head.weight + if self._logits_zero is None: + self._logits_zero = w.new_zeros(1) + logits = torch.addmm( + self._logits_zero, + hidden_states, + w.t(), + beta=0.0, + alpha=1.0 / mup, + ) + logits = self.logits_processor._gather_logits(logits) + if logits is not None: + logits = logits[..., : self.logits_processor.org_vocab_size] + return logits + + def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Greedy draft tokens via rank-local argmax + tiny (value, index) + reduction — no full-vocab logits all-gather. The muP divisor is a + positive scalar, so the argmax is invariant and the scaling is + skipped entirely.""" + return self.logits_processor.get_top_tokens(self.lm_head, hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + return _load_inkling_mtp_weights(self, weights) + + +def _load_inkling_mtp_weights( + module: InklingMTP, + weights: Iterable[tuple[str, torch.Tensor]], +) -> set[str]: + """Load ``model.mtp.*`` weights into the MTP module. + + Checkpoint keys look like ``model.mtp.chain_norm.weight`` and + ``model.mtp.layers.{i}.{...}``. The transformer block reuses the backbone + layer's fused-projection layout, so we apply the same qkvr / gate_up / down + remapping as ``_load_inkling_weights``. Token embedding and LM head are shared + (provided by ``load_eagle_model``) and are not present in mtp.safetensors. + """ + # Per-depth attention is full or sliding-window (config.local_layer_ids); + # each depth's qkvr MergedColumnParallelLinear is built with the matching + # (swa_)num_key_value_heads, and its weight_loader handles the TP sharding. + # The sconv SWA cache pins tp_size <= num_key_value_heads, so tp never + # exceeds a layer's kv-head count and no GQA K/V replication is needed here. + params = dict(module.named_parameters()) + loaded: set[str] = set() + + def _load(name: str, weight: torch.Tensor, shard_id: object = None) -> bool: + param = params.get(name) + if param is None: + return False + loader = getattr(param, "weight_loader", default_weight_loader) + if shard_id is None: + if loader is default_weight_loader or param.shape == weight.shape: + default_weight_loader(param, weight) + else: + loader(param, weight) + else: + loader(param, weight, shard_id) # type: ignore[call-arg] + loaded.add(name) + return True + + for name, weight in weights: + depth = _mtp_depth_from_name(name) + # Token embedding and LM head are never materialized on the draft + # (no params to load into); load_eagle_model attaches the target's. + if name in ("model.llm.embed.weight", "model.llm.unembed.weight"): + continue + # The backbone embed_norm, applied to the shared embedding before the + # depth layers (see InklingMultiTokenPredictor.embed_input_ids). The + # per-depth mtp.layers.{i}.embed_norm keys carry ".mtp." and are loaded + # below. Only the shared backbone key routes here. + if name == "model.llm.embed_norm.weight": + _load("model.backbone_embed_norm.weight", weight) + continue + # Only consume the MTP weights; everything else belongs to the target. + if ".mtp." not in name: + continue + # Only the first checkpoint depth is used for MTP=1. + if depth is not None and depth != 0: + continue + # model.mtp.chain_norm.weight -> model.chain_norm.weight + # model.mtp.layers.{i}.X -> model.layers.{i}.X + original_name = name + name = name.replace(".mtp.layers.", ".layers.").replace( + ".mtp.chain_norm.", ".chain_norm." + ) + + if ".chain_norm." in name and module.model.chain_norm is None: + raise ValueError( + "Inkling checkpoint contains chain_norm weights but " + "chain_hidden_post_norm is disabled." + ) + + # Fused attention qkvr (wq_du/wk_dv/wv_dv/wr_du -> qkvr). + matched = False + for pname, wname, shard in _ATTENTION_PARAMS_MAPPING: + if f".attn.{wname}." in name: + mapped_name = name.replace(f".{wname}.", f".{pname}.") + if not _load(mapped_name, weight, shard): + raise ValueError(f"Unexpected Inkling MTP weight: {original_name}") + matched = True + break + if matched: + continue + + # Dense MLP fused gate/up + down. + if ".mlp.w13_dn.weight" in name: + loaded_weight = _load(name.replace(".w13_dn.", ".gate_up_proj."), weight) + elif ".mlp.w2_md.weight" in name: + loaded_weight = _load(name.replace(".w2_md.", ".down_proj."), weight) + else: + if name.endswith(".bias") and name not in params: + continue + loaded_weight = _load(name, weight) + if not loaded_weight: + raise ValueError(f"Unexpected Inkling MTP weight: {original_name}") + required = { + name + for name in params + if name.startswith("model.layers.") or name.startswith("model.chain_norm.") + } + if (missing := sorted(required - loaded)) and is_mtp_completeness_check_enabled(): + raise ValueError( + "Inkling MTP checkpoint is missing required parameters: " + + ", ".join(missing) + ) + return loaded + + +EntryClass = [InklingMTP] diff --git a/vllm/models/inkling/amd/ops/__init__.py b/vllm/models/inkling/amd/ops/__init__.py new file mode 100644 index 00000000000..3c6ea5fe8f6 --- /dev/null +++ b/vllm/models/inkling/amd/ops/__init__.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling kernels (NVIDIA). + +``rmsnorm`` / ``sconv`` import eagerly. The SwiGLU kernels and the FA4 +relative-attention wrapper are exposed lazily to keep this package's +import path lightweight. +""" + +from typing import TYPE_CHECKING + +from .norm import add_rmsnorm, rmsnorm +from .sconv import fused_sconv + +_LAZY_EXPORTS = { + "silu_and_mul_triton": "silu_and_mul", + "sink_silu_mul_epilogue": "silu_and_mul", + "inkling_fa4_rel_attention": "fa4_rel_attention", + "inkling_rel_attention_split_kv_decode": "rel_attention_decode", +} + +if TYPE_CHECKING: + from .fa4_rel_attention import inkling_fa4_rel_attention # noqa: F401 + from .silu_and_mul import ( # noqa: F401 + silu_and_mul_triton, + sink_silu_mul_epilogue, + ) + +__all__ = [ + "add_rmsnorm", + "rmsnorm", + "fused_sconv", + *sorted(_LAZY_EXPORTS), +] + + +def __getattr__(name: str): + module = _LAZY_EXPORTS.get(name) + if module is not None: + import importlib + + mod = importlib.import_module(f".{module}", __name__) + return getattr(mod, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/vllm/models/inkling/amd/ops/fa4_rel_attention.py b/vllm/models/inkling/amd/ops/fa4_rel_attention.py new file mode 100644 index 00000000000..d448ac8e2c5 --- /dev/null +++ b/vllm/models/inkling/amd/ops/fa4_rel_attention.py @@ -0,0 +1,438 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ROCm paged attention with Inkling's query-dependent relative bias. + +The NVIDIA implementation uses the score-mod hook in tml-fa4. ROCm Flash +Attention and AITER do not expose an equivalent hook, so this module implements +the same operation directly in Triton. Query heads belonging to one KV head +are processed together and KV pages are gathered through vLLM's block table. +""" + +from __future__ import annotations + +import os +from typing import cast + +import torch + +from vllm.models.inkling.amd.ops.rel_attention_decode import ( + inkling_rel_attention_split_kv_decode, + use_split_kv_decode, +) +from vllm.platforms.rocm import on_gfx950 +from vllm.triton_utils import tl, triton + + +def bucket_max_seqlen_q(max_seqlen_q: int) -> int: + """Round the scheduling bound up to a power of two.""" + return 1 << max(0, max_seqlen_q - 1).bit_length() + + +def use_gfx950_gluon_decode( + *, max_query_len: int, page_size: int, head_dim: int +) -> bool: + """Use the vendored TokenSpeed CDNA4 decode kernel where it is supported.""" + return ( + os.getenv("INKLING_GFX950_GLUON", "1") == "1" + and on_gfx950() + and max_query_len == 1 + and page_size in (64, 128, 256) + and head_dim in (64, 128) + ) + + +def use_gfx950_gluon_extend( + *, + max_query_len: int, + max_kv_len: int, + page_size: int, + head_dim: int, + window_left: int, +) -> bool: + """Use Gluon only for the long full-attention extend regime it wins.""" + return ( + os.getenv("INKLING_GFX950_GLUON", "1") == "1" + and on_gfx950() + and max_query_len > 1 + and max_kv_len >= 8192 + and window_left < 0 + and page_size in (64, 128, 256) + and head_dim in (64, 128) + ) + + +def inkling_fa4_num_splits( + *, + is_local: bool, + batch_size: int, + max_query_len: int, + num_heads: int, + num_kv_heads: int, + max_kv_len: int, +) -> int: + """Keep the NVIDIA-facing split heuristic as API-compatible metadata. + + The ROCm Triton implementation performs online softmax in one program and + does not consume the result. Keeping this function unchanged avoids + platform-specific scheduling branches in :mod:`inkling.amd.attention`. + """ + if is_local: + return 1 + + q_rows = max_query_len * (num_heads // num_kv_heads) + q_tiles = (q_rows + 255) // 256 + base_ctas = batch_size * num_kv_heads * q_tiles + target_ctas = ( + 256 if q_tiles == 1 and batch_size == 1 else (128 if q_tiles == 1 else 64) + ) + max_splits = 128 + if q_tiles == 1 and batch_size == 1: + if num_kv_heads == 8: + max_splits = 16 + elif num_kv_heads == 4 or max_kv_len <= 8192: + max_splits = 32 + elif max_kv_len <= 65536: + max_splits = 64 + return max( + 1, + min(target_ctas // base_ctas, max_splits, (max_kv_len + 127) // 128), + ) + + +@triton.heuristics( + { + "BLOCK_D": lambda args: triton.next_power_of_2(args["head_dim"]), + "BLOCK_H": lambda args: triton.next_power_of_2(args["gqa_group_size"]), + "BLOCK_QH": lambda args: ( + args["BLOCK_Q"] * triton.next_power_of_2(args["gqa_group_size"]) + ), + } +) +@triton.jit(do_not_specialize_on_alignment=["cache_seqlens", "cu_seqlens_q"]) +def _inkling_rel_attention_kernel( + q_ptr, # [total_q, Hq, D] + k_ptr, # [blocks, page, Hkv, D] + v_ptr, # [blocks, page, Hkv, D] + rel_ptr, # [total_q, Hq, rel_extent] + out_ptr, # [total_q, Hq, D] + block_table_ptr, # [batch, max_pages] + cache_seqlens, + cu_seqlens_q, + gqa_group_size, + head_dim, + softmax_scale, + stride_q_t, + stride_q_h, + stride_q_d, + stride_k_b, + stride_k_p, + stride_k_h, + stride_k_d, + stride_v_b, + stride_v_p, + stride_v_h, + stride_v_d, + stride_r_t, + stride_r_h, + stride_r_e, + stride_o_t, + stride_o_h, + stride_o_d, + stride_bt_b, + page_size: tl.constexpr, + rel_extent: tl.constexpr, + window_left: tl.constexpr, + BLOCK_Q: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_QH: tl.constexpr, +): + # A program owns BLOCK_Q query positions and every query head sharing one + # KV head. This makes the QK/PV operations MFMA-friendly on CDNA. + pid_q = tl.program_id(0) + pid_kh = tl.program_id(1) + pid_b = tl.program_id(2) + + q_start = tl.load(cu_seqlens_q + pid_b) + q_len = tl.load(cu_seqlens_q + pid_b + 1) - q_start + q_block = pid_q * BLOCK_Q + if q_block >= q_len: + return + + kv_len = tl.load(cache_seqlens + pid_b) + prefix_len = kv_len - q_len + q_head_start = pid_kh * gqa_group_size + bt_row = block_table_ptr + pid_b * stride_bt_b + + q_block_ptr = tl.make_block_ptr( + base=q_ptr + q_start * stride_q_t + q_head_start * stride_q_h, + shape=(q_len, gqa_group_size, head_dim), + strides=(stride_q_t, stride_q_h, stride_q_d), + offsets=(q_block, 0, 0), + block_shape=(BLOCK_Q, BLOCK_H, BLOCK_D), + order=(2, 1, 0), + ) + q = tl.load(q_block_ptr, boundary_check=(0, 1, 2), padding_option="zero") + q = tl.reshape(q, (BLOCK_QH, BLOCK_D)) + + q_rows = q_block + tl.arange(0, BLOCK_Q) + q_abs = prefix_len + q_rows + q_valid = q_rows < q_len + off_d = tl.arange(0, BLOCK_D) + d_valid = off_d < head_dim + + m_i = tl.full((BLOCK_QH,), float("-inf"), dtype=tl.float32) + l_i = tl.zeros((BLOCK_QH,), dtype=tl.float32) + acc = tl.zeros((BLOCK_QH, BLOCK_D), dtype=tl.float32) + log2e: tl.constexpr = 1.4426950408889634 + + for k_start in range(0, kv_len, BLOCK_K): + k_pos = k_start + tl.arange(0, BLOCK_K) + k_valid = k_pos < kv_len + # Masked lanes still need an in-range page-table load. + safe_pos = tl.minimum(k_pos, tl.maximum(kv_len - 1, 0)) + page = tl.load(bt_row + safe_pos // page_size).to(tl.int64) + page_off = safe_pos % page_size + + k = tl.load( + k_ptr + + page[None, :] * stride_k_b + + page_off[None, :] * stride_k_p + + pid_kh * stride_k_h + + off_d[:, None] * stride_k_d, + mask=d_valid[:, None] & k_valid[None, :], + other=0.0, + ) + scores = tl.dot(q, k) * (softmax_scale * log2e) + + dist = q_abs[:, None] - k_pos[None, :] + score_valid = q_valid[:, None] & k_valid[None, :] & (dist >= 0) + if window_left >= 0: + score_valid &= dist <= window_left + + # rel_logits is query- and head-dependent. Expand [Q, K] distance + # indices across the GQA heads, then flatten to match QK's row order. + off_h = tl.arange(0, BLOCK_H) + rel_dist = dist[:, None, :] + rel_valid = ( + q_valid[:, None, None] + & (off_h[None, :, None] < gqa_group_size) + & (rel_dist >= 0) + & (rel_dist < rel_extent) + ) + safe_dist = tl.maximum(0, tl.minimum(rel_dist, rel_extent - 1)) + bias = tl.load( + rel_ptr + + (q_start + q_rows[:, None, None]) * stride_r_t + + (q_head_start + off_h[None, :, None]) * stride_r_h + + safe_dist * stride_r_e, + mask=rel_valid, + other=0.0, + ) + bias = tl.reshape(bias, (BLOCK_QH, BLOCK_K)) + scores += bias.to(tl.float32) * log2e + score_valid = tl.reshape( + score_valid[:, None, :] & (off_h[None, :, None] < gqa_group_size), + (BLOCK_QH, BLOCK_K), + ) + scores = tl.where(score_valid, scores, float("-inf")) + + m_ij = tl.maximum(m_i, tl.max(scores, axis=1)) + has_scores = tl.sum(score_valid.to(tl.int32), axis=1) > 0 + # A sliding-window query can have entire early KV tiles masked. Avoid + # the -inf - -inf NaN in online softmax until its first live tile. + alpha = tl.where(has_scores, tl.exp2(m_i - m_ij), 1.0) + p = tl.where( + score_valid, + tl.exp2(scores - m_ij[:, None]), + 0.0, + ) + l_i = l_i * alpha + tl.sum(p, axis=1) + acc *= alpha[:, None] + + v = tl.load( + v_ptr + + page[:, None] * stride_v_b + + page_off[:, None] * stride_v_p + + pid_kh * stride_v_h + + off_d[None, :] * stride_v_d, + mask=k_valid[:, None] & d_valid[None, :], + other=0.0, + ) + acc += tl.dot(p.to(v.dtype), v) + m_i = m_ij + + acc /= l_i[:, None] + acc = tl.reshape(acc, (BLOCK_Q, BLOCK_H, BLOCK_D)) + out_block_ptr = tl.make_block_ptr( + base=out_ptr + q_start * stride_o_t + q_head_start * stride_o_h, + shape=(q_len, gqa_group_size, head_dim), + strides=(stride_o_t, stride_o_h, stride_o_d), + offsets=(q_block, 0, 0), + block_shape=(BLOCK_Q, BLOCK_H, BLOCK_D), + order=(2, 1, 0), + ) + tl.store( + out_block_ptr, + acc.to(out_ptr.dtype.element_ty), + boundary_check=(0, 1, 2), + ) + + +@torch.no_grad() +def inkling_fa4_rel_attention( + q: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + *, + block_table: torch.Tensor, + cache_seqlens: torch.Tensor, + cu_seqlens_q: torch.Tensor, + max_seqlen_q: int, + softmax_scale: float, + causal: bool, + window_size: tuple[int, int], + rel_extent: int, + rel_logits: torch.Tensor, + num_splits: int = 32, + max_kv_len: int | None = None, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Paged varlen attention with Inkling's relative score modification.""" + del num_splits + if not causal or window_size[1] not in (0, -1): + raise NotImplementedError("Inkling ROCm attention requires causal masking") + if q.ndim != 3 or key_cache.ndim != 4 or value_cache.ndim != 4: + raise ValueError("expected q [T,H,D] and paged K/V [B,P,Hkv,D]") + if rel_logits.shape != (q.shape[0], q.shape[1], rel_extent): + raise ValueError( + f"relative logits have shape {tuple(rel_logits.shape)}, expected " + f"{(q.shape[0], q.shape[1], rel_extent)}" + ) + + num_kv_heads = key_cache.shape[2] + if q.shape[1] % num_kv_heads: + raise ValueError("query heads must be divisible by KV heads") + if out is None: + out = torch.empty_like(q) + gqa_group_size = q.shape[1] // num_kv_heads + batch = cu_seqlens_q.shape[0] - 1 + if max_kv_len is None: + max_kv_len = key_cache.shape[0] * key_cache.shape[1] + if use_gfx950_gluon_decode( + max_query_len=max_seqlen_q, + page_size=key_cache.shape[1], + head_dim=q.shape[2], + ): + # Import only after the architecture guard. The implementation uses + # gfx950-only CDNA4 Gluon layouts and async-copy operations. + from vllm.models.inkling.amd.ops.gluon.rel_mha_decode_gfx950 import ( + gluon_rel_mha_decode_gfx950, + ) + + return gluon_rel_mha_decode_gfx950( + q, + key_cache, + value_cache, + block_table, + cache_seqlens, + max_kv_len, + rel_logits, + cu_seqlens_q, + max_seqlen_q=max_seqlen_q, + window_left=window_size[0], + softmax_scale=softmax_scale, + out=out, + ) + if use_gfx950_gluon_extend( + max_query_len=max_seqlen_q, + max_kv_len=max_kv_len, + page_size=key_cache.shape[1], + head_dim=q.shape[2], + window_left=window_size[0], + ): + from vllm.models.inkling.amd.ops.gluon.rel_mha_extend_gfx950 import ( + gluon_rel_mha_extend_gfx950, + ) + + # The copied kernel retains TokenSpeed's cu_seqlens_kv parameter for + # API compatibility but does not consume it. + return cast( + torch.Tensor, + gluon_rel_mha_extend_gfx950( + q, + cu_seqlens_q, + cu_seqlens_q, + key_cache, + value_cache, + block_table, + cache_seqlens, + window_left=window_size[0], + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_kv_len, + rel_logits=rel_logits, + softmax_scale=softmax_scale, + out=out, + ), + ) + if use_split_kv_decode( + max_query_len=max_seqlen_q, + max_kv_len=max_kv_len, + page_size=key_cache.shape[1], + window_left=window_size[0], + ): + return inkling_rel_attention_split_kv_decode( + q, + key_cache, + value_cache, + block_table=block_table, + cache_seqlens=cache_seqlens, + softmax_scale=softmax_scale, + window_left=window_size[0], + rel_extent=rel_extent, + rel_logits=rel_logits, + max_kv_len=max_kv_len, + out=out, + ) + block_q = 1 if max_seqlen_q == 1 else 4 + grid = (triton.cdiv(max_seqlen_q, block_q), num_kv_heads, batch) + _inkling_rel_attention_kernel[grid]( + q, + key_cache, + value_cache, + rel_logits, + out, + block_table, + cache_seqlens, + cu_seqlens_q, + gqa_group_size, + q.shape[2], + softmax_scale, + q.stride(0), + q.stride(1), + q.stride(2), + key_cache.stride(0), + key_cache.stride(1), + key_cache.stride(2), + key_cache.stride(3), + value_cache.stride(0), + value_cache.stride(1), + value_cache.stride(2), + value_cache.stride(3), + rel_logits.stride(0), + rel_logits.stride(1), + rel_logits.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + block_table.stride(0), + page_size=key_cache.shape[1], + rel_extent=rel_extent, + window_left=window_size[0], + BLOCK_Q=block_q, + BLOCK_K=64, + num_warps=4, + num_stages=1, + ) + return out diff --git a/vllm/models/inkling/amd/ops/fa4_warmup.py b/vllm/models/inkling/amd/ops/fa4_warmup.py new file mode 100644 index 00000000000..998acdc0973 --- /dev/null +++ b/vllm/models/inkling/amd/ops/fa4_warmup.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compatibility shim for the ROCm relative-attention implementation. + +ROCm uses a Triton kernel which is compiled by vLLM's normal model warmup. The +NVIDIA path registers ahead-of-time CuTeDSL units; importing that provider on +ROCm would pull in CUDA-only tml-fa4 code. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class InklingFA4WarmupConfig: + num_heads: int + num_kv_heads: int + head_dim: int + rel_extent: int + window_size: tuple[int, int] + is_local: bool + max_kv_len: int + dtype: torch.dtype + kv_dtype: torch.dtype + block_size: int + max_num_reqs: int + max_num_batched_tokens: int + + +def register_fa4_warmup(config: InklingFA4WarmupConfig) -> None: + """Triton compilation is triggered by the ordinary vLLM warmup forward.""" + del config diff --git a/vllm/models/inkling/amd/ops/gluon/__init__.py b/vllm/models/inkling/amd/ops/gluon/__init__.py new file mode 100644 index 00000000000..54494063da9 --- /dev/null +++ b/vllm/models/inkling/amd/ops/gluon/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""GFX950-only Gluon relative-attention kernels for Inkling.""" diff --git a/vllm/models/inkling/amd/ops/gluon/rel_mha_decode_gfx950.py b/vllm/models/inkling/amd/ops/gluon/rel_mha_decode_gfx950.py new file mode 100644 index 00000000000..987508f91f2 --- /dev/null +++ b/vllm/models/inkling/amd/ops/gluon/rel_mha_decode_gfx950.py @@ -0,0 +1,1037 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""rel_mha decode Gluon kernel for AMD GFX950.""" + +import math +from typing import NamedTuple + +import torch + +from vllm.models.inkling.amd.ops.gluon.utils import ( + _INV_LN2_VALUE, + InputStrides, + PagedKVStrides, + max, + maximum, +) +from vllm.triton_utils import aggregate, gl, gluon + +cdna4 = gl.amd.cdna4 +async_copy = cdna4.async_copy +cdiv = gl.cdiv +_GFX950_SM_COUNT = 256 + + +# ===-----------------------------------------------------------------------===# +# Kernel Config +# ===-----------------------------------------------------------------------===# + + +@aggregate +class AttentionConfig: + SM_SCALE: gl.constexpr + PAGE_TABLE_STRIDE: gl.constexpr + PAGE_SIZE: gl.constexpr + NUM_KV_SPLITS: gl.constexpr + MAX_SEQLEN_Q: gl.constexpr + NUM_Q_HEADS: gl.constexpr + NUM_KV_HEADS: gl.constexpr + HEAD_DIM: gl.constexpr + BLOCK_M: gl.constexpr + BLOCK_N: gl.constexpr + IS_SLIDING: gl.constexpr + WINDOW_LEFT: gl.constexpr + REL_EXTENT: gl.constexpr + REL_BIAS_QK_SCALE: gl.constexpr + IS_FP8: gl.constexpr + GROUP_SIZE: gl.constexpr + NUM_GROUPS: gl.constexpr + q_strides: InputStrides + rel_strides: InputStrides + k_strides: PagedKVStrides + v_strides: PagedKVStrides + qk_layout: gl.constexpr + pv_layout: gl.constexpr + q_layout: gl.constexpr + k_layout: gl.constexpr + p_layout: gl.constexpr + v_layout: gl.constexpr + load_layout: gl.constexpr + store_layout: gl.constexpr + reduce_layout: gl.constexpr + k_smem_layout: gl.constexpr + v_smem_layout: gl.constexpr + + @gluon.constexpr_function + def __init__( + self, + SM_SCALE, + PAGE_TABLE_STRIDE, + PAGE_SIZE, + NUM_KV_SPLITS, + MAX_SEQLEN_Q, + NUM_Q_HEADS, + NUM_KV_HEADS, + HEAD_DIM, + BLOCK_M, + BLOCK_N, + IS_SLIDING, + WINDOW_LEFT, + REL_EXTENT, + REL_BIAS_QK_SCALE, + IS_FP8, + q_strides, + rel_strides, + k_strides, + v_strides, + ): + assert NUM_Q_HEADS % NUM_KV_HEADS == 0 + assert HEAD_DIM in (64, 128) + assert BLOCK_N == PAGE_SIZE + if IS_SLIDING: + assert WINDOW_LEFT >= 0 + else: + assert WINDOW_LEFT == -1 + + mfma_layout = gl.amd.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 32], + transposed=True, + warps_per_cta=[1, 1], + ) + qk_layout = mfma_layout + pv_layout = mfma_layout + # qk_kw is derived from a 128-bit load / dtype bitwidth. + # pv_kw is empirically tuned. + qk_kw = 16 if IS_FP8 else 8 + pv_kw = 8 if IS_FP8 else 4 + q_layout = gl.DotOperandLayout(0, qk_layout, k_width=qk_kw) + k_layout = gl.DotOperandLayout(1, qk_layout, k_width=qk_kw) + p_layout = gl.DotOperandLayout(0, pv_layout, k_width=pv_kw) + v_layout = gl.DotOperandLayout(1, pv_layout, k_width=pv_kw) + # Elements loaded per lane depend on the input dtype, matching qk_kw. + # load_threads is how many lanes span HEAD_DIM. + load_vec = 16 if IS_FP8 else 8 + load_threads = HEAD_DIM // load_vec + load_layout = gl.BlockedLayout( + [1, load_vec], [64 // load_threads, load_threads], [1, 1], [1, 0] + ) + # Output is always 16-bit, so a 128-bit store has 8 elements. + # store_threads is how many lanes span HEAD_DIM. + store_vec = 8 + store_threads = HEAD_DIM // store_vec + store_layout = gl.BlockedLayout( + [1, store_vec], [64 // store_threads, store_threads], [1, 1], [1, 0] + ) + reduce_layout = gl.BlockedLayout([1, HEAD_DIM // 64], [1, 64], [1, 1], [1, 0]) + # Padding interval is 64 lanes * load_vec elems. + pad_interval = 64 * load_vec + # Empirically tuned. + pad_k = 16 if IS_FP8 else 8 + pad_v = 32 + k_smem_layout = gl.PaddedSharedLayout.with_identity_for( + [[pad_interval, pad_k]], [BLOCK_N, HEAD_DIM], [1, 0] + ) + v_smem_layout = gl.PaddedSharedLayout.with_identity_for( + [[pad_interval, pad_v]], [BLOCK_N, HEAD_DIM], [1, 0] + ) + + self.SM_SCALE = gl.constexpr(SM_SCALE) + self.PAGE_TABLE_STRIDE = gl.constexpr(PAGE_TABLE_STRIDE) + self.PAGE_SIZE = gl.constexpr(PAGE_SIZE) + self.NUM_KV_SPLITS = gl.constexpr(NUM_KV_SPLITS) + self.MAX_SEQLEN_Q = gl.constexpr(MAX_SEQLEN_Q) + self.NUM_Q_HEADS = gl.constexpr(NUM_Q_HEADS) + self.NUM_KV_HEADS = gl.constexpr(NUM_KV_HEADS) + self.HEAD_DIM = gl.constexpr(HEAD_DIM) + self.BLOCK_M = gl.constexpr(BLOCK_M) + self.BLOCK_N = gl.constexpr(BLOCK_N) + self.IS_SLIDING = gl.constexpr(IS_SLIDING) + self.WINDOW_LEFT = gl.constexpr(WINDOW_LEFT) + self.REL_EXTENT = gl.constexpr(REL_EXTENT) + self.REL_BIAS_QK_SCALE = gl.constexpr(REL_BIAS_QK_SCALE) + self.IS_FP8 = gl.constexpr(IS_FP8) + self.GROUP_SIZE = gl.constexpr(NUM_Q_HEADS // NUM_KV_HEADS) + self.NUM_GROUPS = gl.constexpr((self.GROUP_SIZE + BLOCK_M - 1) // BLOCK_M) + self.q_strides = q_strides + self.rel_strides = rel_strides + self.k_strides = k_strides + self.v_strides = v_strides + self.qk_layout = gl.constexpr(qk_layout) + self.pv_layout = gl.constexpr(pv_layout) + self.q_layout = gl.constexpr(q_layout) + self.k_layout = gl.constexpr(k_layout) + self.p_layout = gl.constexpr(p_layout) + self.v_layout = gl.constexpr(v_layout) + self.load_layout = gl.constexpr(load_layout) + self.store_layout = gl.constexpr(store_layout) + self.reduce_layout = gl.constexpr(reduce_layout) + self.k_smem_layout = gl.constexpr(k_smem_layout) + self.v_smem_layout = gl.constexpr(v_smem_layout) + + +# ===-----------------------------------------------------------------------===# +# Kernel Program +# ===-----------------------------------------------------------------------===# + + +@aggregate +class AttentionProgram: + cfg: gl.constexpr + q_ptr: gl.tensor + rel_logits_ptr: gl.tensor + k_cache_ptr: gl.tensor + v_cache_ptr: gl.tensor + page_table_ptr: gl.tensor + cache_seqlens_ptr: gl.tensor + mid_o_ptr: gl.tensor + mid_lse_ptr: gl.tensor + q_index: gl.tensor + batch: gl.tensor + kv_head: gl.tensor + group_start: gl.tensor + split_id: gl.tensor + cache_len: gl.tensor + kv_start: gl.tensor + split_start: gl.tensor + split_end: gl.tensor + + @gluon.constexpr_function + def __init__( + self, + cfg, + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + cache_seqlens_ptr, + mid_o_ptr, + mid_lse_ptr, + q_index, + batch, + kv_head, + group_start, + split_id, + cache_len, + kv_start, + split_start, + split_end, + ): + self.cfg = gl.constexpr(cfg) + self.q_ptr = q_ptr + self.rel_logits_ptr = rel_logits_ptr + self.k_cache_ptr = k_cache_ptr + self.v_cache_ptr = v_cache_ptr + self.page_table_ptr = page_table_ptr + self.cache_seqlens_ptr = cache_seqlens_ptr + self.mid_o_ptr = mid_o_ptr + self.mid_lse_ptr = mid_lse_ptr + self.q_index = q_index + self.batch = batch + self.kv_head = kv_head + self.group_start = group_start + self.split_id = split_id + self.cache_len = cache_len + self.kv_start = kv_start + self.split_start = split_start + self.split_end = split_end + + @gluon.jit + def create( + cfg, + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + cache_seqlens_ptr, + mid_o_ptr, + mid_lse_ptr, + ): + # Gluon treats ``cfg`` as the constexpr factory argument, while mypy + # models the first parameter of this aggregate method as ``self``. + q_index = gl.program_id(0) + batch = q_index // cfg.MAX_SEQLEN_Q # type: ignore[attr-defined] + q_pos = q_index - batch * cfg.MAX_SEQLEN_Q # type: ignore[attr-defined] + head_block = gl.program_id(1) + kv_head = head_block // cfg.NUM_GROUPS # type: ignore[attr-defined] + group_block = head_block - kv_head * cfg.NUM_GROUPS # type: ignore[attr-defined] + group_start = group_block * cfg.BLOCK_M # type: ignore[attr-defined] + split_id = gl.program_id(2) + cache_len = gl.load(cache_seqlens_ptr + batch) + cache_len = cache_len - ( + cfg.MAX_SEQLEN_Q - 1 - q_pos # type: ignore[attr-defined] + ) + cache_len = maximum(cache_len, 0) + if cfg.IS_SLIDING: # type: ignore[attr-defined] + # WINDOW_LEFT is defined as exclusive (keys strictly to the left, not + # counting the current token), so the window is WINDOW_LEFT + 1 keys + # once the current token is included. E.g. with + # WINDOW_LEFT = 127 and cache_len = 500 (current token at index 499), + # kv_start = 500 - (127 + 1) = 372 keeps keys [372, 499] = 128 keys. + # Without the + 1, kv_start = 373 would drop the leftmost key. + kv_start = cache_len - min( + cache_len, + cfg.WINDOW_LEFT + 1, # type: ignore[attr-defined] + ) + else: + kv_start = cache_len - cache_len + first_page = kv_start // cfg.PAGE_SIZE # type: ignore[attr-defined] + end_page = cdiv(cache_len, cfg.PAGE_SIZE) # type: ignore[attr-defined] + num_pages = end_page - first_page + pages_per_split = cdiv( + num_pages, + cfg.NUM_KV_SPLITS, # type: ignore[attr-defined] + ) + split_start_page = first_page + split_id * pages_per_split + split_end_page = min(split_start_page + pages_per_split, end_page) + split_start = ( + split_start_page * cfg.PAGE_SIZE # type: ignore[attr-defined] + ) + split_end = min( + split_end_page * cfg.PAGE_SIZE, # type: ignore[attr-defined] + cache_len, + ) + return AttentionProgram( + gl.constexpr(cfg), + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + cache_seqlens_ptr, + mid_o_ptr, + mid_lse_ptr, + q_index, + batch, + kv_head, + group_start, + split_id, + cache_len, + kv_start, + split_start, + split_end, + ) + + @gluon.jit + def load_q(self): + cfg = self.cfg + offs_m = gl.arange(0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.q_layout)) + offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.q_layout)) + q_heads = self.kv_head * cfg.GROUP_SIZE + self.group_start + offs_m + valid = (self.group_start + offs_m) < cfg.GROUP_SIZE + offsets = cfg.q_strides.offsets(self.q_index, q_heads[:, None], offs_d[None, :]) + return cdna4.buffer_load(self.q_ptr, offsets, mask=valid[:, None], other=0.0) + + @gluon.jit + def init_state(self): + cfg = self.cfg + m_i = gl.full( + [cfg.BLOCK_M], + value=-float("inf"), + dtype=gl.float32, + layout=gl.SliceLayout(1, cfg.pv_layout), + ) + l_i = gl.full( + [cfg.BLOCK_M], + value=0.0, + dtype=gl.float32, + layout=gl.SliceLayout(1, cfg.pv_layout), + ) + acc = gl.zeros( + [cfg.BLOCK_M, cfg.HEAD_DIM], dtype=gl.float32, layout=cfg.pv_layout + ) + return m_i, l_i, acc + + @gluon.jit + def load_page(self, start_n): + cfg = self.cfg + page_index = start_n // cfg.PAGE_SIZE + valid = start_n < self.split_end + return gl.load( + self.page_table_ptr + self.batch * cfg.PAGE_TABLE_STRIDE + page_index, + mask=valid, + other=0, + ) + + @gluon.jit + def issue_load_k(self, physical_page, k_smem): + cfg = self.cfg + offs_n = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(1, cfg.load_layout)) + offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.load_layout)) + offsets = cfg.k_strides.offsets( + physical_page, + offs_n[:, None], + self.kv_head, + offs_d[None, :], + ) + # can't use buffer_load: paged KV offsets may exceed its 32-bit range. + async_copy.global_load_to_shared(k_smem, self.k_cache_ptr + offsets) + async_copy.commit_group() + + @gluon.jit + def issue_load_v(self, physical_page, v_smem): + cfg = self.cfg + offs_n = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(1, cfg.load_layout)) + offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.load_layout)) + offsets = cfg.v_strides.offsets( + physical_page, + offs_n[:, None], + self.kv_head, + offs_d[None, :], + ) + # can't use buffer_load: paged KV offsets may exceed its 32-bit range. + async_copy.global_load_to_shared(v_smem, self.v_cache_ptr + offsets) + async_copy.commit_group() + + @gluon.jit + def shared_load_k(self, k_smem): + return k_smem.permute([1, 0]).load(self.cfg.k_layout) + + @gluon.jit + def shared_load_v(self, v_smem): + return v_smem.load(self.cfg.v_layout) + + @gluon.jit + def compute_qk(self, q, k): + cfg = self.cfg + qk = gl.zeros( + [cfg.BLOCK_M, cfg.BLOCK_N], dtype=gl.float32, layout=cfg.qk_layout + ) + return cdna4.mfma(q, k, qk) + + @gluon.jit + def apply_rel_bias(self, qk, start_n): + cfg = self.cfg + offs_m = gl.arange(0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.qk_layout)) + offs_n = start_n + gl.arange( + 0, cfg.BLOCK_N, layout=gl.SliceLayout(0, cfg.qk_layout) + ) + q_heads = self.kv_head * cfg.GROUP_SIZE + self.group_start + offs_m + rel_dist = (self.cache_len - 1) - offs_n + rel_valid = (rel_dist >= 0) & (rel_dist < cfg.REL_EXTENT) + rel_idx = gl.where(rel_dist >= 0, rel_dist, 0) + rel_idx = gl.where(rel_idx < cfg.REL_EXTENT, rel_idx, cfg.REL_EXTENT - 1) + offsets = cfg.rel_strides.offsets( + self.q_index, q_heads[:, None], rel_idx[None, :] + ) + head_valid = (self.group_start + offs_m) < cfg.GROUP_SIZE + rel_bias = cdna4.buffer_load( + self.rel_logits_ptr, + offsets, + mask=head_valid[:, None] & rel_valid[None, :], + other=0.0, + ).to(gl.float32) + return qk + rel_bias * cfg.REL_BIAS_QK_SCALE + + @gluon.jit + def apply_kv_mask(self, qk, start_n): + cfg = self.cfg + offs_n = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(0, cfg.qk_layout)) + tokens = start_n + offs_n[None, :] + mask = (tokens >= self.kv_start) & (tokens < self.split_end) + return gl.where(mask, qk, -float("inf")) + + @gluon.jit + def softmax(self, qk, m_i, l_i, acc): + cfg = self.cfg + row_max = max(qk, axis=1) + row_max = gl.convert_layout(row_max, gl.SliceLayout(1, cfg.pv_layout)) + m_new = maximum(m_i, row_max) + m_new_scaled = m_new * cfg.SM_SCALE + qk_shifted = qk * cfg.SM_SCALE - m_new_scaled[:, None] + p = gl.exp2(qk_shifted) + m_diff = m_i * cfg.SM_SCALE - m_new_scaled + alpha = gl.exp2(m_diff) + l_ij = gl.sum(p, axis=1) + l_i = l_i * alpha + l_ij + acc = acc * alpha[:, None] + p = p.to(self.q_ptr.dtype.element_ty) + p = gl.convert_layout(p, cfg.p_layout) + return p, m_new, l_i, acc + + @gluon.jit + def compute_pv(self, p, v, acc): + return cdna4.mfma(p, v, acc) + + @gluon.jit + def store_split(self, acc, l_i, m_i): + cfg = self.cfg + offs_m = gl.arange(0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.store_layout)) + offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.store_layout)) + q_heads = self.kv_head * cfg.GROUP_SIZE + self.group_start + offs_m + valid = ((self.group_start + offs_m) < cfg.GROUP_SIZE) & ( + self.split_start < self.split_end + ) + acc = gl.convert_layout(acc, cfg.store_layout) + l_i = gl.convert_layout(l_i, gl.SliceLayout(1, cfg.store_layout)) + m_i = gl.convert_layout(m_i, gl.SliceLayout(1, cfg.store_layout)) + recip_l_i = 1.0 / l_i + part_o = acc * recip_l_i[:, None] + part_lse = m_i * cfg.SM_SCALE + gl.log2(l_i) + mid_o_offsets = ( + (self.q_index * cfg.NUM_Q_HEADS + q_heads[:, None]) * cfg.NUM_KV_SPLITS + + self.split_id + ) * cfg.HEAD_DIM + offs_d[None, :] + mid_lse_offsets = ( + self.q_index * cfg.NUM_Q_HEADS + q_heads + ) * cfg.NUM_KV_SPLITS + self.split_id + cdna4.buffer_store(part_o, self.mid_o_ptr, mid_o_offsets, mask=valid[:, None]) + cdna4.buffer_store(part_lse, self.mid_lse_ptr, mid_lse_offsets, mask=valid) + + @gluon.jit + def store_output(self, acc, l_i): + cfg = self.cfg + offs_m = gl.arange(0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.store_layout)) + offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.store_layout)) + q_heads = self.kv_head * cfg.GROUP_SIZE + self.group_start + offs_m + valid = (self.group_start + offs_m) < cfg.GROUP_SIZE + acc = gl.convert_layout(acc, cfg.store_layout) + l_i = gl.convert_layout(l_i, gl.SliceLayout(1, cfg.store_layout)) + output = acc * (1.0 / l_i)[:, None] + output = output.to(self.mid_o_ptr.dtype.element_ty) + offsets = (self.q_index * cfg.NUM_Q_HEADS + q_heads[:, None]) * cfg.HEAD_DIM + offsets += offs_d[None, :] + cdna4.buffer_store(output, self.mid_o_ptr, offsets, mask=valid[:, None]) + + +# ===-----------------------------------------------------------------------===# +# Entry Point +# ===-----------------------------------------------------------------------===# + + +@gluon.jit +def _rel_mha_decode_fp16( + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + cache_seqlens_ptr, + mid_o_ptr, + mid_lse_ptr, + Q_STRIDE_B: gl.constexpr, + Q_STRIDE_H: gl.constexpr, + Q_STRIDE_D: gl.constexpr, + K_STRIDE_B: gl.constexpr, + K_STRIDE_P: gl.constexpr, + K_STRIDE_H: gl.constexpr, + K_STRIDE_D: gl.constexpr, + V_STRIDE_B: gl.constexpr, + V_STRIDE_P: gl.constexpr, + V_STRIDE_H: gl.constexpr, + V_STRIDE_D: gl.constexpr, + SM_SCALE: gl.constexpr, + PAGE_TABLE_STRIDE: gl.constexpr, + PAGE_SIZE: gl.constexpr, + NUM_KV_SPLITS: gl.constexpr, + MAX_SEQLEN_Q: gl.constexpr, + NUM_Q_HEADS: gl.constexpr, + NUM_KV_HEADS: gl.constexpr, + HEAD_DIM: gl.constexpr, + BLOCK_M: gl.constexpr, + BLOCK_N: gl.constexpr, + IS_SLIDING: gl.constexpr, + WINDOW_LEFT: gl.constexpr, + REL_STRIDE_T: gl.constexpr, + REL_STRIDE_H: gl.constexpr, + REL_STRIDE_E: gl.constexpr, + REL_EXTENT: gl.constexpr, + REL_BIAS_QK_SCALE: gl.constexpr, + IS_FP8: gl.constexpr, +): + cfg = AttentionConfig( + SM_SCALE, + PAGE_TABLE_STRIDE, + PAGE_SIZE, + NUM_KV_SPLITS, + MAX_SEQLEN_Q, + NUM_Q_HEADS, + NUM_KV_HEADS, + HEAD_DIM, + BLOCK_M, + BLOCK_N, + IS_SLIDING, + WINDOW_LEFT, + REL_EXTENT, + REL_BIAS_QK_SCALE, + IS_FP8, + InputStrides(Q_STRIDE_B, Q_STRIDE_H, Q_STRIDE_D), + InputStrides(REL_STRIDE_T, REL_STRIDE_H, REL_STRIDE_E), + PagedKVStrides(K_STRIDE_B, K_STRIDE_P, K_STRIDE_H, K_STRIDE_D), + PagedKVStrides(V_STRIDE_B, V_STRIDE_P, V_STRIDE_H, V_STRIDE_D), + ) + program = AttentionProgram.create( + cfg, + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + cache_seqlens_ptr, + mid_o_ptr, + mid_lse_ptr, + ) + k_smem = gl.allocate_shared_memory( + k_cache_ptr.dtype.element_ty, [cfg.BLOCK_N, cfg.HEAD_DIM], cfg.k_smem_layout + ) + v_smem = gl.allocate_shared_memory( + v_cache_ptr.dtype.element_ty, [cfg.BLOCK_N, cfg.HEAD_DIM], cfg.v_smem_layout + ) + + q = program.load_q() + m_i, l_i, acc = program.init_state() + + physical_page = program.load_page(program.split_start) + + for start_n in range(program.split_start, program.split_end, cfg.BLOCK_N): + program.issue_load_k(physical_page, k_smem) + program.issue_load_v(physical_page, v_smem) + physical_page = program.load_page(start_n + cfg.BLOCK_N) + + async_copy.wait_group(1) + k = program.shared_load_k(k_smem) + qk = program.compute_qk(q, k) + qk = program.apply_rel_bias(qk, start_n) + qk = program.apply_kv_mask(qk, start_n) + p, m_i, l_i, acc = program.softmax(qk, m_i, l_i, acc) + + async_copy.wait_group(0) + v = program.shared_load_v(v_smem) + acc = program.compute_pv(p, v, acc) + + program.store_split(acc, l_i, m_i) + + +@gluon.jit +def _rel_mha_decode_sliding_fp16( + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + cache_seqlens_ptr, + out_ptr, + Q_STRIDE_B: gl.constexpr, + Q_STRIDE_H: gl.constexpr, + Q_STRIDE_D: gl.constexpr, + K_STRIDE_B: gl.constexpr, + K_STRIDE_P: gl.constexpr, + K_STRIDE_H: gl.constexpr, + K_STRIDE_D: gl.constexpr, + V_STRIDE_B: gl.constexpr, + V_STRIDE_P: gl.constexpr, + V_STRIDE_H: gl.constexpr, + V_STRIDE_D: gl.constexpr, + SM_SCALE: gl.constexpr, + PAGE_TABLE_STRIDE: gl.constexpr, + PAGE_SIZE: gl.constexpr, + MAX_SEQLEN_Q: gl.constexpr, + NUM_Q_HEADS: gl.constexpr, + NUM_KV_HEADS: gl.constexpr, + HEAD_DIM: gl.constexpr, + BLOCK_M: gl.constexpr, + BLOCK_N: gl.constexpr, + IS_SLIDING: gl.constexpr, + WINDOW_LEFT: gl.constexpr, + REL_STRIDE_T: gl.constexpr, + REL_STRIDE_H: gl.constexpr, + REL_STRIDE_E: gl.constexpr, + REL_EXTENT: gl.constexpr, + REL_BIAS_QK_SCALE: gl.constexpr, + IS_FP8: gl.constexpr, +): + cfg = AttentionConfig( + SM_SCALE, + PAGE_TABLE_STRIDE, + PAGE_SIZE, + 1, + MAX_SEQLEN_Q, + NUM_Q_HEADS, + NUM_KV_HEADS, + HEAD_DIM, + BLOCK_M, + BLOCK_N, + IS_SLIDING, + WINDOW_LEFT, + REL_EXTENT, + REL_BIAS_QK_SCALE, + IS_FP8, + InputStrides(Q_STRIDE_B, Q_STRIDE_H, Q_STRIDE_D), + InputStrides(REL_STRIDE_T, REL_STRIDE_H, REL_STRIDE_E), + PagedKVStrides(K_STRIDE_B, K_STRIDE_P, K_STRIDE_H, K_STRIDE_D), + PagedKVStrides(V_STRIDE_B, V_STRIDE_P, V_STRIDE_H, V_STRIDE_D), + ) + program = AttentionProgram.create( + cfg, + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + cache_seqlens_ptr, + out_ptr, + out_ptr, + ) + k_smem = gl.allocate_shared_memory( + k_cache_ptr.dtype.element_ty, [cfg.BLOCK_N, cfg.HEAD_DIM], cfg.k_smem_layout + ) + v_smem = gl.allocate_shared_memory( + v_cache_ptr.dtype.element_ty, [cfg.BLOCK_N, cfg.HEAD_DIM], cfg.v_smem_layout + ) + + q = program.load_q() + m_i, l_i, acc = program.init_state() + + for start_n in range(program.split_start, program.split_end, cfg.BLOCK_N): + physical_page = program.load_page(start_n) + program.issue_load_k(physical_page, k_smem) + program.issue_load_v(physical_page, v_smem) + async_copy.wait_group(1) + k = program.shared_load_k(k_smem) + qk = program.compute_qk(q, k) + qk = program.apply_rel_bias(qk, start_n) + qk = program.apply_kv_mask(qk, start_n) + p, m_i, l_i, acc = program.softmax(qk, m_i, l_i, acc) + + async_copy.wait_group(0) + v = program.shared_load_v(v_smem) + acc = program.compute_pv(p, v, acc) + + program.store_output(acc, l_i) + + +@gluon.jit +def _rel_mha_decode_reduce_fp16( + mid_o_ptr, + mid_lse_ptr, + out_ptr, + cache_seqlens_ptr, + SM_SCALE: gl.constexpr, + PAGE_TABLE_STRIDE: gl.constexpr, + NUM_KV_SPLITS: gl.constexpr, + MAX_SEQLEN_Q: gl.constexpr, + PAGE_SIZE: gl.constexpr, + NUM_Q_HEADS: gl.constexpr, + NUM_KV_HEADS: gl.constexpr, + HEAD_DIM: gl.constexpr, + BLOCK_M: gl.constexpr, + BLOCK_N: gl.constexpr, + IS_SLIDING: gl.constexpr, + WINDOW_LEFT: gl.constexpr, + IS_FP8: gl.constexpr, +): + cfg = AttentionConfig( + SM_SCALE, + PAGE_TABLE_STRIDE, + PAGE_SIZE, + NUM_KV_SPLITS, + MAX_SEQLEN_Q, + NUM_Q_HEADS, + NUM_KV_HEADS, + HEAD_DIM, + BLOCK_M, + BLOCK_N, + IS_SLIDING, + WINDOW_LEFT, + 1, # REL_EXTENT + 1.0, # REL_BIAS_QK_SCALE + IS_FP8, + InputStrides(1, 1, 1), + InputStrides(1, 1, 1), + PagedKVStrides(1, 1, 1, 1), + PagedKVStrides(1, 1, 1, 1), + ) + q_index = gl.program_id(0) + batch = q_index // MAX_SEQLEN_Q + q_pos = q_index - batch * MAX_SEQLEN_Q + q_head = gl.program_id(1) + cache_len = gl.load(cache_seqlens_ptr + batch) + cache_len = cache_len - (MAX_SEQLEN_Q - 1 - q_pos) + cache_len = maximum(cache_len, 0) + if cfg.IS_SLIDING: + kv_start = cache_len - min(cache_len, cfg.WINDOW_LEFT + 1) + else: + kv_start = cache_len - cache_len + first_page = kv_start // cfg.PAGE_SIZE + end_page = cdiv(cache_len, cfg.PAGE_SIZE) + num_pages = end_page - first_page + pages_per_split = cdiv(num_pages, cfg.NUM_KV_SPLITS) + + # SPLIT_TILE pads NUM_KV_SPLITS up to a power of 2. + SPLIT_TILE: gl.constexpr = 1 << (NUM_KV_SPLITS - 1).bit_length() + offs_s = gl.arange(0, SPLIT_TILE, layout=gl.SliceLayout(1, cfg.reduce_layout)) + offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.reduce_layout)) + # split_valid masks out empty splits and the power-of-2 padding tail. + split_start_page = first_page + offs_s * pages_per_split + split_end_page_raw = split_start_page + pages_per_split + split_end_page = gl.where( + split_end_page_raw < end_page, split_end_page_raw, end_page + ) + split_start_tok = split_start_page * cfg.PAGE_SIZE + split_end_raw = split_end_page * cfg.PAGE_SIZE + split_end_tok = gl.where(split_end_raw < cache_len, split_end_raw, cache_len) + split_valid = (split_start_tok < split_end_tok) & (offs_s < cfg.NUM_KV_SPLITS) + # Load every split's partial output and lse. + base = (q_index * cfg.NUM_Q_HEADS + q_head) * cfg.NUM_KV_SPLITS + offs_s + part_lse = gl.load(mid_lse_ptr + base, mask=split_valid, other=-float("inf")) + o_off = base[:, None] * cfg.HEAD_DIM + offs_d[None, :] + part_o = cdna4.buffer_load(mid_o_ptr, o_off, mask=split_valid[:, None], other=0.0) + + # Global softmax max over all splits. + m_i = max(part_lse, axis=0) + # Weighted sum of the split partials, normalized by the total softmax mass. + beta = gl.exp2(part_lse - m_i) + l_i = gl.sum(beta, axis=0) + acc = gl.sum(part_o * beta[:, None], axis=0) + + out_base = (q_index * cfg.NUM_Q_HEADS + q_head) * cfg.HEAD_DIM + output = acc * (1.0 / l_i) + output = output.to(out_ptr.dtype.element_ty) + cdna4.buffer_store(output, out_ptr, out_base + offs_d) + + +def _select_num_kv_splits( + *, + batch: int, + num_kv_heads: int, + num_groups: int, + num_pages: int, + sm_count: int, +) -> int: + """Pick num_kv_splits to balance occupancy against reduce overhead. + + The launch grid is (batch * num_kv_heads * num_groups) * num_kv_splits + work-groups. Too few splits under-fill the machine at low batch; too many + leave each split with a handful of pages, so the reduce kernel dominates. + + Return the smaller of two candidate counts: splits_for_occupancy (enough to + fill ~wave_target waves of CUs) and splits_for_pages (~min_pages_per_split + pages per split), with the pages candidate clamped to [min_page_splits, + max_page_splits] so a short context still splits without launching empty work + and a long one does not over-split where reduce cost outgrows the decode win. + """ + wave_target = 2 + min_pages_per_split = 2 + min_page_splits = 8 + max_page_splits = 32 + + base_ctas = batch * num_kv_heads * num_groups + target_ctas = sm_count * wave_target + splits_for_occupancy = (target_ctas + base_ctas - 1) // base_ctas + + splits_for_pages = num_pages // min_pages_per_split + min_page_splits = min(min_page_splits, num_pages) + if splits_for_pages < min_page_splits: + splits_for_pages = min_page_splits + if splits_for_pages > max_page_splits: + splits_for_pages = max_page_splits + return min(splits_for_occupancy, splits_for_pages) + + +class LaunchConfig(NamedTuple): + num_q_heads: int + num_kv_heads: int + num_groups: int + head_dim: int + page_size: int + num_kv_splits: int + block_m: int + block_n: int + sm_scale: float + rel_bias_qk_scale: float + is_sliding: bool + window_left: int + + +def get_config( + *, + q: torch.Tensor, + k_cache: torch.Tensor, + max_seqlen_k: int, + window_left: int, + softmax_scale: float | None, +) -> LaunchConfig: + head_dim = q.shape[2] + page_size = k_cache.shape[1] + block_m = 16 + block_n = page_size + group_size = q.shape[1] // k_cache.shape[2] + num_groups = math.ceil(group_size / block_m) + is_sliding = window_left >= 0 + window_left = window_left if is_sliding else -1 + if softmax_scale is None: + softmax_scale = 1.0 / math.sqrt(head_dim) + sm_scale = softmax_scale + effective_seqlen_k = ( + min(max_seqlen_k, window_left + 1) if is_sliding else max_seqlen_k + ) + num_pages = (effective_seqlen_k + page_size - 1) // page_size + if is_sliding: + num_kv_splits = 4 + else: + num_kv_splits = _select_num_kv_splits( + batch=q.shape[0], + num_kv_heads=k_cache.shape[2], + num_groups=num_groups, + num_pages=num_pages, + sm_count=_GFX950_SM_COUNT, + ) + return LaunchConfig( + num_q_heads=q.shape[1], + num_kv_heads=k_cache.shape[2], + num_groups=num_groups, + head_dim=head_dim, + page_size=page_size, + num_kv_splits=num_kv_splits, + block_m=block_m, + block_n=block_n, + sm_scale=sm_scale * _INV_LN2_VALUE, + rel_bias_qk_scale=1.0 / softmax_scale, + is_sliding=is_sliding, + window_left=window_left, + ) + + +def gluon_rel_mha_decode_gfx950( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + page_table: torch.Tensor, + cache_seqlens: torch.Tensor, + max_seqlen_k: int, + rel_logits: torch.Tensor, + cu_seqlens_q: torch.Tensor, + max_seqlen_q: int = 1, + window_left: int = -1, + softmax_scale: float | None = None, + q_scale: torch.Tensor | None = None, + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, + out: torch.Tensor | None = None, +) -> torch.Tensor: + total_q = q.shape[0] + + config = get_config( + q=q, + k_cache=k_cache, + max_seqlen_k=max_seqlen_k, + window_left=window_left, + softmax_scale=softmax_scale, + ) + + is_fp8 = q.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) + out_dtype = torch.bfloat16 if is_fp8 else q.dtype + if k_cache.shape[1] not in (64, 128, 256): + raise ValueError( + "gfx950 Gluon relative-attention decode requires page size " + f"64, 128, or 256; got {k_cache.shape[1]}" + ) + output = ( + torch.empty(q.shape, device=q.device, dtype=out_dtype) if out is None else out + ) + + # Always use split-k for full attention and for the small-batch sliding + # decode path. Sliding uses a fixed 8 splits, one page per split for the + # TP-4 local-attention shape with 512-token window and 64-token pages. + mid_o = torch.empty( + (total_q, config.num_q_heads, config.num_kv_splits, config.head_dim), + device=q.device, + dtype=torch.float32, + ) + mid_lse = torch.empty( + (total_q, config.num_q_heads, config.num_kv_splits), + device=q.device, + dtype=torch.float32, + ) + + grid = ( + total_q, + config.num_kv_heads * config.num_groups, + config.num_kv_splits, + ) + _rel_mha_decode_fp16[grid]( + q, + rel_logits, + k_cache, + v_cache, + page_table, + cache_seqlens, + mid_o, + mid_lse, + q.stride(0), + q.stride(1), + q.stride(2), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + k_cache.stride(3), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + v_cache.stride(3), + config.sm_scale, + page_table.stride(0), + config.page_size, + config.num_kv_splits, + max_seqlen_q, + config.num_q_heads, + config.num_kv_heads, + config.head_dim, + config.block_m, + config.block_n, + config.is_sliding, + config.window_left, + rel_logits.stride(0), + rel_logits.stride(1), + rel_logits.stride(2), + rel_logits.shape[2], + config.rel_bias_qk_scale, + is_fp8, + num_warps=1, + ) + + reduce_grid = (total_q, config.num_q_heads) + _rel_mha_decode_reduce_fp16[reduce_grid]( + mid_o, + mid_lse, + output, + cache_seqlens, + config.sm_scale, + page_table.stride(0), + config.num_kv_splits, + max_seqlen_q, + config.page_size, + config.num_q_heads, + config.num_kv_heads, + config.head_dim, + config.block_m, + config.block_n, + config.is_sliding, + config.window_left, + is_fp8, + num_warps=1, + ) + return output diff --git a/vllm/models/inkling/amd/ops/gluon/rel_mha_extend_gfx950.py b/vllm/models/inkling/amd/ops/gluon/rel_mha_extend_gfx950.py new file mode 100644 index 00000000000..224524e76b9 --- /dev/null +++ b/vllm/models/inkling/amd/ops/gluon/rel_mha_extend_gfx950.py @@ -0,0 +1,760 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""rel_mha extend Gluon kernel for AMD GFX950. + +This handles ragged, multi-token queries against a paged KV cache. The query +axis is tiled into the MFMA ``M`` dimension (prefill-style): ``BLOCK_M`` query +rows of a single q-head share each paged KV tile, so every KV tile is loaded +once and reused across all rows in the tile. The grid is +``(blocks_per_req, batch, n_heads)`` -- all host-known sizes -- and each program +self-locates its request from ``cu_seqlens_q`` / ``cache_seqlens`` in-kernel, so +the launch stays CUDA-graph static with no device->host sync. + +Visibility per query row depends on ``is_causal`` and the optional sliding +window: + +* ``is_causal=False``: every query token attends the full visible cache, i.e. + ``visible_kv = cache_seqlens[batch]``. +* ``is_causal=True``: query tokens are a causal suffix, so the ``i``-th query + token of a request (0-indexed) attends ``prefix + i + 1`` tokens, where + ``prefix = cache_seqlens[batch] - query_len[batch]``. + +Causal masking only touches the few KV tiles that reach the diagonal; the long +prefix is a mask-free fast path. Sinks and sliding windows (causal or not) are +applied per tile. This is the sole Gluon extend implementation. +""" + +import math + +import torch + +from vllm.models.inkling.amd.ops.gluon.utils import ( + _INV_LN2_VALUE, + _LN2, + InputStrides, + PagedKVStrides, + max, + maximum, +) +from vllm.triton_utils import aggregate, gl, gluon + +cdna4 = gl.amd.cdna4 +async_copy = cdna4.async_copy + + +# ===-----------------------------------------------------------------------===# +# Query-batched path +# +# This path tiles BLOCK_M query rows (of a single q-head) into the MFMA M +# dimension -- like the prefill kernel -- so each paged KV tile is loaded once +# and reused across all BLOCK_M rows. KV still comes from the paged cache +# (decode-style async page loads). Causal masking only touches the few KV tiles +# that reach the diagonal; the long prefix is a mask-free fast path. +# ===-----------------------------------------------------------------------===# + +_EXTEND_SHORT_Q_BLOCK_M = 64 +_EXTEND_SHORT_Q_NUM_WARPS = 2 +_EXTEND_LONG_Q_BLOCK_M = 128 +_EXTEND_LONG_Q_NUM_WARPS = 4 + + +def _select_extend_tile(max_seqlen_q: int, page_size: int) -> tuple[int, int, int]: + """Return (BLOCK_M, BLOCK_N, NUM_WARPS) for the given max query length. + + Queries that fit in a single short tile use it (least padding, most + occupancy); longer ones use the tall tile that covers more rows per + shared-KV pass. + """ + if max_seqlen_q <= _EXTEND_SHORT_Q_BLOCK_M: + return _EXTEND_SHORT_Q_BLOCK_M, page_size, _EXTEND_SHORT_Q_NUM_WARPS + return _EXTEND_LONG_Q_BLOCK_M, page_size, _EXTEND_LONG_Q_NUM_WARPS + + +@aggregate +class ExtendConfig: + N_HEADS: gl.constexpr + N_KV_HEADS: gl.constexpr + GROUP_SIZE: gl.constexpr + HEAD_DIM: gl.constexpr + SM_SCALE: gl.constexpr + BLOCK_M: gl.constexpr + BLOCK_N: gl.constexpr + NUM_WARPS: gl.constexpr + PAGE_SIZE: gl.constexpr + PAGE_TABLE_STRIDE: gl.constexpr + IS_CAUSAL: gl.constexpr + HAS_LSE: gl.constexpr + WINDOW_LEFT: gl.constexpr + REL_EXTENT: gl.constexpr + REL_BIAS_QK_SCALE: gl.constexpr + IS_FP8: gl.constexpr + q_strides: InputStrides + rel_strides: InputStrides + k_strides: PagedKVStrides + v_strides: PagedKVStrides + qk_layout: gl.constexpr + pv_layout: gl.constexpr + q_layout: gl.constexpr + k_layout: gl.constexpr + p_layout: gl.constexpr + v_layout: gl.constexpr + load_layout: gl.constexpr + store_layout: gl.constexpr + k_smem_layout: gl.constexpr + v_smem_layout: gl.constexpr + + @gluon.constexpr_function + def __init__( + self, + N_HEADS, + N_KV_HEADS, + HEAD_DIM, + SM_SCALE, + BLOCK_M, + BLOCK_N, + NUM_WARPS, + PAGE_SIZE, + PAGE_TABLE_STRIDE, + IS_CAUSAL, + HAS_LSE, + WINDOW_LEFT, + REL_EXTENT, + REL_BIAS_QK_SCALE, + IS_FP8, + q_strides, + rel_strides, + k_strides, + v_strides, + ): + assert HEAD_DIM in (64, 128) + assert BLOCK_N == PAGE_SIZE + + instr_shape = [32, 32, 16] + mfma_layout = gl.amd.AMDMFMALayout( + version=4, + instr_shape=instr_shape, + transposed=True, + warps_per_cta=[NUM_WARPS, 1], + ) + qk_layout = mfma_layout + pv_layout = mfma_layout + # Elements loaded per lane depend on the input dtype, matching qk_kw. + # load_threads is how many lanes span HEAD_DIM. + load_vec = 16 if IS_FP8 else 8 + load_threads = HEAD_DIM // load_vec + load_layout = gl.BlockedLayout( + [1, load_vec], [64 // load_threads, load_threads], [NUM_WARPS, 1], [1, 0] + ) + # Output is always 16-bit, so a 128-bit store has 8 elements. + # store_threads is how many lanes span HEAD_DIM. + store_vec = 8 + store_threads = HEAD_DIM // store_vec + store_layout = gl.BlockedLayout( + [1, store_vec], [64 // store_threads, store_threads], [NUM_WARPS, 1], [1, 0] + ) + # Padding interval is 64 lanes * load_vec elems. + pad_interval = 64 * load_vec + # Empirically tuned. + pad_k = 16 if IS_FP8 else 8 + pad_v = 16 if IS_FP8 else 32 + k_smem_layout = gl.PaddedSharedLayout.with_identity_for( + [[pad_interval, pad_k]], [BLOCK_N, HEAD_DIM], [1, 0] + ) + v_smem_layout = gl.PaddedSharedLayout.with_identity_for( + [[pad_interval, pad_v]], [BLOCK_N, HEAD_DIM], [1, 0] + ) + + self.N_HEADS = gl.constexpr(N_HEADS) + self.N_KV_HEADS = gl.constexpr(N_KV_HEADS) + self.GROUP_SIZE = gl.constexpr(N_HEADS // N_KV_HEADS) + self.HEAD_DIM = gl.constexpr(HEAD_DIM) + self.SM_SCALE = gl.constexpr(SM_SCALE) + self.BLOCK_M = gl.constexpr(BLOCK_M) + self.BLOCK_N = gl.constexpr(BLOCK_N) + self.NUM_WARPS = gl.constexpr(NUM_WARPS) + self.PAGE_SIZE = gl.constexpr(PAGE_SIZE) + self.PAGE_TABLE_STRIDE = gl.constexpr(PAGE_TABLE_STRIDE) + self.IS_CAUSAL = gl.constexpr(IS_CAUSAL) + self.HAS_LSE = gl.constexpr(HAS_LSE) + self.WINDOW_LEFT = gl.constexpr(WINDOW_LEFT) + self.REL_EXTENT = gl.constexpr(REL_EXTENT) + self.REL_BIAS_QK_SCALE = gl.constexpr(REL_BIAS_QK_SCALE) + self.IS_FP8 = gl.constexpr(IS_FP8) + self.q_strides = q_strides + self.rel_strides = rel_strides + self.k_strides = k_strides + self.v_strides = v_strides + self.qk_layout = gl.constexpr(qk_layout) + self.pv_layout = gl.constexpr(pv_layout) + # qk_kw is derived from a 128-bit load / dtype bitwidth. + # pv_kw is empirically tuned. + qk_kw = 16 if IS_FP8 else 8 + pv_kw = 16 if IS_FP8 else 4 + self.q_layout = gl.constexpr(gl.DotOperandLayout(0, qk_layout, k_width=qk_kw)) + self.k_layout = gl.constexpr(gl.DotOperandLayout(1, qk_layout, k_width=qk_kw)) + self.p_layout = gl.constexpr(gl.DotOperandLayout(0, pv_layout, k_width=pv_kw)) + self.v_layout = gl.constexpr(gl.DotOperandLayout(1, pv_layout, k_width=pv_kw)) + self.load_layout = gl.constexpr(load_layout) + self.store_layout = gl.constexpr(store_layout) + self.k_smem_layout = gl.constexpr(k_smem_layout) + self.v_smem_layout = gl.constexpr(v_smem_layout) + + +@aggregate +class ExtendProgram: + cfg: gl.constexpr + q_ptr: gl.tensor + rel_logits_ptr: gl.tensor + k_cache_ptr: gl.tensor + v_cache_ptr: gl.tensor + page_table_ptr: gl.tensor + output_ptr: gl.tensor + lse_ptr: gl.tensor + batch: gl.tensor + q_head: gl.tensor + kv_head: gl.tensor + q_start: gl.tensor + seq_base: gl.tensor + seq_len: gl.tensor + prefix: gl.tensor + cache_len: gl.tensor + + @gluon.constexpr_function + def __init__( + self, + cfg, + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + output_ptr, + lse_ptr, + batch, + q_head, + kv_head, + q_start, + seq_base, + seq_len, + prefix, + cache_len, + ): + self.cfg = gl.constexpr(cfg) + self.q_ptr = q_ptr + self.rel_logits_ptr = rel_logits_ptr + self.k_cache_ptr = k_cache_ptr + self.v_cache_ptr = v_cache_ptr + self.page_table_ptr = page_table_ptr + self.output_ptr = output_ptr + self.lse_ptr = lse_ptr + self.batch = batch + self.q_head = q_head + self.kv_head = kv_head + self.q_start = q_start + self.seq_base = seq_base + self.seq_len = seq_len + self.prefix = prefix + self.cache_len = cache_len + + @gluon.jit + def create( + cfg, + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + output_ptr, + lse_ptr, + cu_seqlens_q_ptr, + cache_seqlens_ptr, + ): + # Gluon treats ``cfg`` as the constexpr factory argument, while mypy + # models the first parameter of this aggregate method as ``self``. + block_in_req = gl.program_id(0) + batch = gl.program_id(1) + q_head = gl.program_id(2) + q_start = block_in_req * cfg.BLOCK_M # type: ignore[attr-defined] + seq_base = gl.load(cu_seqlens_q_ptr + batch) + seq_end = gl.load(cu_seqlens_q_ptr + batch + 1) + seq_len = seq_end - seq_base + cache_len = gl.load(cache_seqlens_ptr + batch) + prefix = cache_len - seq_len + kv_head = q_head // cfg.GROUP_SIZE # type: ignore[attr-defined] + return ExtendProgram( + gl.constexpr(cfg), + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + output_ptr, + lse_ptr, + batch, + q_head, + kv_head, + q_start, + seq_base, + seq_len, + prefix, + cache_len, + ) + + @gluon.jit + def load_q(self): + cfg = self.cfg + offs_m = self.q_start + gl.arange( + 0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.q_layout) + ) + offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.q_layout)) + row = self.seq_base + offs_m + offsets = cfg.q_strides.offsets(row[:, None], self.q_head, offs_d[None, :]) + mask = offs_m[:, None] < self.seq_len + return cdna4.buffer_load(self.q_ptr, offsets, mask=mask, other=0.0) + + @gluon.jit + def load_page(self, start_n): + cfg = self.cfg + page_index = start_n // cfg.PAGE_SIZE + valid = start_n < self.cache_len + return gl.load( + self.page_table_ptr + self.batch * cfg.PAGE_TABLE_STRIDE + page_index, + mask=valid, + other=0, + ) + + @gluon.jit + def issue_load_k(self, physical_page, k_smem): + cfg = self.cfg + offs_n = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(1, cfg.load_layout)) + offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.load_layout)) + offsets = cfg.k_strides.offsets( + physical_page, + offs_n[:, None], + self.kv_head, + offs_d[None, :], + ) + async_copy.global_load_to_shared(k_smem, self.k_cache_ptr + offsets) + async_copy.commit_group() + + @gluon.jit + def issue_load_v(self, physical_page, v_smem): + cfg = self.cfg + offs_n = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(1, cfg.load_layout)) + offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.load_layout)) + offsets = cfg.v_strides.offsets( + physical_page, + offs_n[:, None], + self.kv_head, + offs_d[None, :], + ) + async_copy.global_load_to_shared(v_smem, self.v_cache_ptr + offsets) + async_copy.commit_group() + + @gluon.jit + def shared_load_k(self, k_smem): + cfg = self.cfg + k_buffer = k_smem.permute([1, 0]) + return k_buffer.load(cfg.k_layout) + + @gluon.jit + def shared_load_v(self, v_smem): + cfg = self.cfg + return v_smem.load(cfg.v_layout) + + @gluon.jit + def compute_qk(self, q, k): + cfg = self.cfg + qk = gl.zeros( + [cfg.BLOCK_M, cfg.BLOCK_N], dtype=gl.float32, layout=cfg.qk_layout + ) + return cdna4.mfma(q, k, qk) + + @gluon.jit + def apply_rel_bias(self, qk, start_n): + cfg = self.cfg + offs_m = self.q_start + gl.arange( + 0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.qk_layout) + ) + offs_n_abs = start_n + gl.arange( + 0, cfg.BLOCK_N, layout=gl.SliceLayout(0, cfg.qk_layout) + ) + q_pos = self.prefix + offs_m + rel_dist = q_pos[:, None] - offs_n_abs[None, :] + rel_valid = (rel_dist >= 0) & (rel_dist < cfg.REL_EXTENT) + rel_idx = gl.where(rel_dist >= 0, rel_dist, 0) + rel_idx = gl.where(rel_idx < cfg.REL_EXTENT, rel_idx, cfg.REL_EXTENT - 1) + offsets = cfg.rel_strides.offsets( + self.seq_base + offs_m[:, None], self.q_head, rel_idx + ) + mask = (offs_m[:, None] < self.seq_len) & rel_valid + rel_bias = cdna4.buffer_load( + self.rel_logits_ptr, offsets, mask=mask, other=0.0 + ).to(gl.float32) + return qk + rel_bias * cfg.REL_BIAS_QK_SCALE + + @gluon.jit + def compute_pv(self, p, v, acc): + return cdna4.mfma(p, v, acc) + + @gluon.jit + def init_attention_state(self): + cfg = self.cfg + m_i = gl.full( + [cfg.BLOCK_M], + value=-float("inf"), + dtype=gl.float32, + layout=gl.SliceLayout(1, cfg.pv_layout), + ) + l_i = gl.full( + [cfg.BLOCK_M], + value=0, + dtype=gl.float32, + layout=gl.SliceLayout(1, cfg.pv_layout), + ) + acc = gl.zeros( + [cfg.BLOCK_M, cfg.HEAD_DIM], dtype=gl.float32, layout=cfg.pv_layout + ) + return m_i, l_i, acc + + @gluon.jit + def softmax(self, qk, m_i, l_i, acc): + cfg = self.cfg + # In sliding window case, some rows can see fully masked tiles before + # any valid KV. Guard the online softmax state so `-inf - -inf` does not + # produce NaNs. + HAS_INVALID: gl.constexpr = cfg.WINDOW_LEFT >= 0 + + row_max = max(qk, 1) + m_new = maximum(m_i, row_max) + m_new_scaled = m_new * cfg.SM_SCALE + if HAS_INVALID: + invalid = m_new == -float("inf") + m_new_scaled = gl.where(invalid, 0.0, m_new_scaled) + + qk_shifted = qk * cfg.SM_SCALE - m_new_scaled[:, None] + p = gl.exp2(qk_shifted) + m_diff = m_i * cfg.SM_SCALE - m_new_scaled + if HAS_INVALID: + m_diff = gl.where(invalid, 0.0, m_diff) + + alpha = gl.exp2(m_diff) + l_ij = gl.sum(p, axis=1) + l_i = l_i * alpha + l_ij + acc = acc * alpha[:, None] + p = p.to(self.q_ptr.dtype.element_ty) + p = gl.convert_layout(p, cfg.p_layout) + return p, m_new, l_i, acc + + @gluon.jit + def store_output(self, output): + cfg = self.cfg + offs_m = self.q_start + gl.arange( + 0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.store_layout) + ) + offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.store_layout)) + offsets = ( + ((self.seq_base + offs_m[:, None]) * cfg.N_HEADS + self.q_head) + * cfg.HEAD_DIM + + offs_d[None, :] + ).to(gl.int32) + mask = offs_m[:, None] < self.seq_len + output = output.to(self.output_ptr.dtype.element_ty) + cdna4.buffer_store(output, self.output_ptr, offsets, mask=mask) + + @gluon.jit + def store_lse(self, l_i, m_i): + cfg = self.cfg + if cfg.HAS_LSE: + offs_m = self.q_start + gl.arange( + 0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.pv_layout) + ) + offsets = ((self.seq_base + offs_m) * cfg.N_HEADS + self.q_head).to( + gl.int32 + ) + mask = offs_m < self.seq_len + lse_l_i = gl.where(l_i > 0.0, l_i, 1.0) + # Softmax runs in base-2 (exp2 hardware fast path), so m_i*SM_SCALE + + # log2(l_i) is the LSE in base-2 units. Convert to natural log (the + # public op contract / torch.logsumexp convention) by scaling by ln2. + lse = (m_i * cfg.SM_SCALE + gl.log2(lse_l_i)) * _LN2 + cdna4.buffer_store(lse, self.lse_ptr, offsets, mask=mask) + + +@gluon.jit +def _rel_mha_extend_fp16( + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + output_ptr, + lse_ptr, + cu_seqlens_q_ptr, + cache_seqlens_ptr, + Q_STRIDE_T: gl.constexpr, + Q_STRIDE_H: gl.constexpr, + Q_STRIDE_D: gl.constexpr, + K_STRIDE_B: gl.constexpr, + K_STRIDE_P: gl.constexpr, + K_STRIDE_H: gl.constexpr, + K_STRIDE_D: gl.constexpr, + V_STRIDE_B: gl.constexpr, + V_STRIDE_P: gl.constexpr, + V_STRIDE_H: gl.constexpr, + V_STRIDE_D: gl.constexpr, + SM_SCALE: gl.constexpr, + PAGE_TABLE_STRIDE: gl.constexpr, + PAGE_SIZE: gl.constexpr, + N_HEADS: gl.constexpr, + N_KV_HEADS: gl.constexpr, + HEAD_DIM: gl.constexpr, + BLOCK_M: gl.constexpr, + BLOCK_N: gl.constexpr, + NUM_WARPS: gl.constexpr, + IS_CAUSAL: gl.constexpr, + HAS_LSE: gl.constexpr, + WINDOW_LEFT: gl.constexpr, + REL_STRIDE_T: gl.constexpr, + REL_STRIDE_H: gl.constexpr, + REL_STRIDE_E: gl.constexpr, + REL_EXTENT: gl.constexpr, + REL_BIAS_QK_SCALE: gl.constexpr, + IS_FP8: gl.constexpr, +): + cfg = ExtendConfig( + N_HEADS, + N_KV_HEADS, + HEAD_DIM, + SM_SCALE, + BLOCK_M, + BLOCK_N, + NUM_WARPS, + PAGE_SIZE, + PAGE_TABLE_STRIDE, + IS_CAUSAL, + HAS_LSE, + WINDOW_LEFT, + REL_EXTENT, + REL_BIAS_QK_SCALE, + IS_FP8, + InputStrides(Q_STRIDE_T, Q_STRIDE_H, Q_STRIDE_D), + InputStrides(REL_STRIDE_T, REL_STRIDE_H, REL_STRIDE_E), + PagedKVStrides(K_STRIDE_B, K_STRIDE_P, K_STRIDE_H, K_STRIDE_D), + PagedKVStrides(V_STRIDE_B, V_STRIDE_P, V_STRIDE_H, V_STRIDE_D), + ) + program = ExtendProgram.create( + cfg, + q_ptr, + rel_logits_ptr, + k_cache_ptr, + v_cache_ptr, + page_table_ptr, + output_ptr, + lse_ptr, + cu_seqlens_q_ptr, + cache_seqlens_ptr, + ) + # Over-provisioned tile past this request's real query rows: nothing to do. + if program.q_start >= program.seq_len: + return + k_smem = gl.allocate_shared_memory( + k_cache_ptr.dtype.element_ty, [cfg.BLOCK_N, cfg.HEAD_DIM], cfg.k_smem_layout + ) + v_smem = gl.allocate_shared_memory( + v_cache_ptr.dtype.element_ty, [cfg.BLOCK_N, cfg.HEAD_DIM], cfg.v_smem_layout + ) + + q = program.load_q() + m_i, l_i, acc = program.init_attention_state() + + offs_m_q = program.q_start + gl.arange( + 0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.qk_layout) + ) + offs_n_q = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(0, cfg.qk_layout)) + diag_row = program.prefix + offs_m_q + + if IS_CAUSAL: + kv_end = min(program.cache_len, program.prefix + program.q_start + cfg.BLOCK_M) + else: + kv_end = program.cache_len + + # Sliding window (inclusive-left, matches flash-attn window_size=(W, 0)): + # skip KV tiles entirely below the window's lower edge. The tile's top query + # row sits at absolute position pos_top; its window opens at + # pos_top - WINDOW_LEFT (that key is visible -> W + 1 keys total). The min() + # form clamps to 0 without the shadowed builtin max(). + if cfg.WINDOW_LEFT >= 0: + pos_top = program.prefix + program.q_start + kv_start = pos_top - min(pos_top, cfg.WINDOW_LEFT) + kv_start = (kv_start // cfg.BLOCK_N) * cfg.BLOCK_N + else: + kv_start = 0 + + for start_n in range(kv_start, kv_end, cfg.BLOCK_N): + physical_page = program.load_page(start_n) + program.issue_load_k(physical_page, k_smem) + program.issue_load_v(physical_page, v_smem) + + async_copy.wait_group(1) + k = program.shared_load_k(k_smem) + qk = program.compute_qk(q, k) + qk = program.apply_rel_bias(qk, start_n) + + if cfg.WINDOW_LEFT >= 0: + # Window lower edge + cache bound always apply; causal upper edge only + # when IS_CAUSAL (independent layering keeps non-causal + window correct). + offs_n_abs = start_n + offs_n_q + mask = (offs_n_abs[None, :] >= diag_row[:, None] - cfg.WINDOW_LEFT) & ( + offs_n_abs[None, :] < program.cache_len + ) + if IS_CAUSAL: + mask &= offs_n_abs[None, :] <= diag_row[:, None] + qk = gl.where(mask, qk, -float("inf")) + elif IS_CAUSAL: + if start_n + cfg.BLOCK_N > program.prefix + program.q_start: + offs_n_abs = start_n + offs_n_q + mask = (offs_n_abs[None, :] <= diag_row[:, None]) & ( + offs_n_abs[None, :] < program.cache_len + ) + qk = gl.where(mask, qk, -float("inf")) + else: + if start_n + cfg.BLOCK_N > program.cache_len: + offs_n_abs = start_n + offs_n_q + qk = gl.where( + offs_n_abs[None, :] < program.cache_len, qk, -float("inf") + ) + + p, m_i, l_i, acc = program.softmax(qk, m_i, l_i, acc) + + async_copy.wait_group(0) + v = program.shared_load_v(v_smem) + acc = program.compute_pv(p, v, acc) + + denom = gl.where(l_i > 0.0, l_i, 1.0) + output = acc * (1.0 / denom)[:, None] + output = gl.convert_layout(output, cfg.store_layout) + program.store_output(output) + program.store_lse(l_i, m_i) + + +def gluon_rel_mha_extend_gfx950( + q: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + page_table: torch.Tensor, + cache_seqlens: torch.Tensor, + window_left: int = -1, + return_lse: bool = False, + max_seqlen_q: int = 1, + max_seqlen_k: int = 1, + rel_logits: torch.Tensor | None = None, + softmax_scale: float | None = None, + q_scale: torch.Tensor | None = None, + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, + out: torch.Tensor | None = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + assert rel_logits is not None + head_dim = q.shape[2] + n_heads = q.shape[1] + n_kv_heads = k_cache.shape[2] + page_size = k_cache.shape[1] + block_m, block_n, num_warps = _select_extend_tile(max_seqlen_q, page_size) + if softmax_scale is None: + softmax_scale = 1.0 / math.sqrt(head_dim) + sm_scale = softmax_scale * _INV_LN2_VALUE + rel_bias_qk_scale = 1.0 / softmax_scale + + # max_seqlen_q must be >= the true max query length; extra tiles early-exit. + batch = cu_seqlens_q.shape[0] - 1 + safe_max_q = max_seqlen_q if max_seqlen_q > 0 else 1 + blocks_per_req = (safe_max_q + block_m - 1) // block_m + cu_q_i32 = cu_seqlens_q.to(torch.int32).contiguous() + cache_i32 = cache_seqlens.to(torch.int32).contiguous() + + is_fp8 = q.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) + out_dtype = torch.bfloat16 if is_fp8 else q.dtype + if page_size not in (64, 128, 256): + raise ValueError( + "gfx950 Gluon relative-attention extend requires page size " + f"64, 128, or 256; got {page_size}" + ) + output = ( + torch.empty(q.shape, device=q.device, dtype=out_dtype) if out is None else out + ) + if return_lse: + lse = torch.empty((q.shape[0], n_heads), device=q.device, dtype=torch.float32) + lse_arg = lse + else: + lse = None + lse_arg = q + grid = (blocks_per_req, batch, n_heads) + _rel_mha_extend_fp16[grid]( + q, + rel_logits, + k_cache, + v_cache, + page_table, + output, + lse_arg, + cu_q_i32, + cache_i32, + q.stride(0), + q.stride(1), + q.stride(2), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + k_cache.stride(3), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + v_cache.stride(3), + sm_scale, + page_table.stride(0), + page_size, + n_heads, + n_kv_heads, + head_dim, + block_m, + block_n, + num_warps, + True, + return_lse, + window_left, + rel_logits.stride(0), + rel_logits.stride(1), + rel_logits.stride(2), + rel_logits.shape[2], + rel_bias_qk_scale, + is_fp8, + num_warps=num_warps, + ) + if return_lse: + assert lse is not None + return output, lse + return output diff --git a/vllm/models/inkling/amd/ops/gluon/utils.py b/vllm/models/inkling/amd/ops/gluon/utils.py new file mode 100644 index 00000000000..90f2e586bc6 --- /dev/null +++ b/vllm/models/inkling/amd/ops/gluon/utils.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +from vllm.triton_utils import aggregate, gl, gluon, tl + +_INV_LN2_VALUE = 1.4426950408889634 +_INV_LN2 = tl.constexpr(_INV_LN2_VALUE) +_LN2_VALUE = 0.6931471805599453 +_LN2 = tl.constexpr(_LN2_VALUE) +_PROPAGATE_NAN_ALL = gl.constexpr(tl.PropagateNan.ALL) + + +@gluon.jit +def maximum(a, b, propagate_nan: gl.constexpr = _PROPAGATE_NAN_ALL): + return gl.maximum(a, b, propagate_nan=propagate_nan) + + +@gluon.jit +def max(input, axis=None, keep_dims=False): + return gl.reduce(input, axis, maximum, keep_dims=keep_dims) + + +@gluon.constexpr_function +def attention_layouts(head_dim, block_n, is_fp8, dtype, num_warps, instr_shape): + mfma = gl.amd.AMDMFMALayout( + version=4, + instr_shape=instr_shape, + transposed=True, + warps_per_cta=[num_warps, 1], + ) + qk_layout = mfma + pv_layout = mfma + # qk_kw is derived from a 128-bit load / dtype bitwidth; pv_kw is tuned. + qk_kw = 16 if is_fp8 else 8 + pv_kw = 8 if is_fp8 else 4 + q_layout = gl.DotOperandLayout(0, qk_layout, k_width=qk_kw) + k_layout = gl.DotOperandLayout(1, qk_layout, k_width=qk_kw) + p_layout = gl.DotOperandLayout(0, pv_layout, k_width=pv_kw) + v_layout = gl.DotOperandLayout(1, pv_layout, k_width=pv_kw) + # load_vec = elems/lane (dtype-dependent, == qk_kw); load_threads span HEAD_DIM. + load_vec = 16 if is_fp8 else 8 + load_threads = head_dim // load_vec + load_layout = gl.BlockedLayout( + [1, load_vec], [64 // load_threads, load_threads], [num_warps, 1], [1, 0] + ) + # store_vec is always 16-bit (128 / 16 = 8) regardless of input dtype. + store_vec = 8 + store_threads = head_dim // store_vec + store_layout = gl.BlockedLayout( + [1, store_vec], [64 // store_threads, store_threads], [num_warps, 1], [1, 0] + ) + # Take only the built-in's padding, not its swizzle: the swizzle scatters + # head_dim across banks, making the LDS write stride non-constant, so it can't + # lower through async_copy's affine [1, 0] load. Padding-only is affine and + # DMA-legal. + # TODO(perf): to also use the swizzle, co-design a matched load layout so the + # DMA stays legal. + k_api = gl.amd.cdna4.compute_efficient_padded_shared_layout( + k_layout, [block_n, head_dim], dtype, is_k_contig=True + ) + v_api = gl.amd.cdna4.compute_efficient_padded_shared_layout( + v_layout, [block_n, head_dim], dtype, is_k_contig=False + ) + assert k_api is not None and v_api is not None, ( + "no CDNA4 padded shared layout for this operand/dtype" + ) + k_pairs = list(k_api.interval_padding_pairs) + v_pairs = list(v_api.interval_padding_pairs) + assert len(k_pairs) == 1 and len(v_pairs) == 1, ( + "expected a single interval padding pair from the built-in" + ) + k_smem_layout = gl.PaddedSharedLayout.with_identity_for( + [[int(k_pairs[0][0]), int(k_pairs[0][1])]], [block_n, head_dim], [1, 0] + ) + v_smem_layout = gl.PaddedSharedLayout.with_identity_for( + [[int(v_pairs[0][0]), int(v_pairs[0][1])]], [block_n, head_dim], [1, 0] + ) + return ( + qk_layout, + pv_layout, + q_layout, + k_layout, + p_layout, + v_layout, + load_layout, + store_layout, + k_smem_layout, + v_smem_layout, + ) + + +@aggregate +class InputStrides: + stride_t: gl.constexpr + stride_h: gl.constexpr + stride_d: gl.constexpr + + @gluon.constexpr_function + def __init__(self, stride_t, stride_h, stride_d): + self.stride_t = gl.constexpr(stride_t) + self.stride_h = gl.constexpr(stride_h) + self.stride_d = gl.constexpr(stride_d) + + @gluon.jit + def offsets(self, token, head, dim): + return (token * self.stride_t + head * self.stride_h + dim * self.stride_d).to( + gl.int32 + ) + + +@aggregate +class PagedKVStrides: + stride_b: gl.constexpr + stride_p: gl.constexpr + stride_h: gl.constexpr + stride_d: gl.constexpr + + @gluon.constexpr_function + def __init__(self, stride_b, stride_p, stride_h, stride_d): + self.stride_b = gl.constexpr(stride_b) + self.stride_p = gl.constexpr(stride_p) + self.stride_h = gl.constexpr(stride_h) + self.stride_d = gl.constexpr(stride_d) + + @gluon.jit + def offsets(self, page, token, head, dim): + # KV pools can exceed the 32-bit buffer-offset range. Keep the page + # arithmetic in int64, matching the original kernel's large-cache path. + return ( + page.to(gl.int64) * self.stride_b + + token.to(gl.int64) * self.stride_p + + head * self.stride_h + + dim * self.stride_d + ) diff --git a/vllm/models/inkling/amd/ops/lamport.py b/vllm/models/inkling/amd/ops/lamport.py new file mode 100644 index 00000000000..9705a9a3df3 --- /dev/null +++ b/vllm/models/inkling/amd/ops/lamport.py @@ -0,0 +1,766 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Deadlock-free fused RS + short-conv + AG + residual + RMSNorm. + +The public integration surface is ``LamportRSConv.rs_sconv_ag_add_norm``. + +Liveness +-------- +Large grids are deliberately split at the two communication dependencies: + + 1. ``_publish_input_kernel`` only publishes rank partials. + 2. ``_reduce_insert_kernel`` waits, reduces, and inserts the local shard. + 3. ``_sconv_publish_kernel`` only computes and publishes the local result. + 4. ``_gather_norm_kernel`` waits, gathers, and normalizes. + +Without PDL, CUDA stream order completes a producer before its consumer. With +PDL, every producer CTA posts all peer stores before triggering its dependent, +and the consumer executes ``gdc_wait`` before polling. A consumer therefore +waits only for stores from a producer that is already running (or complete) on +another GPU. It never waits for another block in its own grid. Consequently +spinning consumers cannot occupy resources needed by any producer, and the +wait-for graph has no cycle. The proof is independent of grid size, block +dispatch order, and occupancy. + +For one token, the first three phases use eight independent channel slices to +expose enough CTA parallelism for decode latency. Each slice has exclusive +ownership of its cache and Lamport columns; the gather/RMSNorm phase retains +one CTA per token so no cross-CTA reduction or completion counter is needed. + +Immediate buffer reuse (including replay of a captured CUDA graph) is also +safe. Rank R cannot republish input for call n+1 until its gather for n has +finished; that gather waited for owner O's output, which O publishes only +after consuming R's input for n. Likewise, R cannot republish output for n+1 +until its reduction for n+1 has observed destination D's input; D publishes +that input only after its gather consumed R's output for n. Thus every prior +read happens-before a same-slot rewrite. Three generations reduce incidental +coupling for ordinary launches, but correctness does not depend on rotation. + +The payload itself is the Lamport flag. A 32-bit store publishes two bf16s +atomically; 0x80008000 (two negative zeroes) denotes an empty pair. Real +negative zeroes are changed to positive zero before publication. Consumers +use volatile 32-bit loads and restore the sentinel after consuming a slot. +""" + +from __future__ import annotations + +import os + +import torch + +from vllm.distributed import get_tp_group +from vllm.distributed.parallel_state import in_the_same_node_as +from vllm.logger import init_logger +from vllm.triton_utils import tl, triton + +logger = init_logger(__name__) + +_MAX_TOKENS = 16384 +_EMPTY_PAIR = tl.constexpr(0x80008000) + + +@triton.jit +def _pack_bf16_pairs(values): + """Pack bf16 values into atomic u32 pairs and reserve negative zero.""" + lo, hi = tl.split(values.reshape([values.shape[0] // 2, 2])) + lo = lo.to(tl.uint16, bitcast=True) + hi = hi.to(tl.uint16, bitcast=True) + lo = tl.where(lo == 0x8000, 0, lo).to(tl.uint32) + hi = tl.where(hi == 0x8000, 0, hi).to(tl.uint32) + return lo | (hi << 16) + + +@triton.jit +def _unpack_bf16_pairs(values): + lo = (values & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True) + hi = (values >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True) + return tl.interleave(lo, hi) + + +@triton.jit +def _wait_pairs(ptr, offsets, mask): + values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True) + while tl.max(tl.where(mask & (values == _EMPTY_PAIR), 1, 0)) != 0: + values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True) + return values + + +@triton.jit +def _publish_input_kernel( + stage_ptr, + peer_ptrs, + peer_offset_u32, + stride_stage_t, + C: tl.constexpr, + CS: tl.constexpr, + CS_P2: tl.constexpr, + SPLITS: tl.constexpr, + RANK: tl.constexpr, + WORLD: tl.constexpr, + USE_PDL: tl.constexpr, + launch_pdl: tl.constexpr, +): + """Publish this rank's full partial row into every shard owner's slots.""" + token = tl.program_id(0).to(tl.int64) + split = tl.program_id(1).to(tl.int64) + CSS: tl.constexpr = CS // SPLITS + pair = tl.arange(0, CS_P2 // 2) + pair_mask = pair < CSS // 2 + elem = tl.arange(0, CS_P2) + elem_mask = elem < CSS + ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64)) + # Address generation is independent of the preceding stream kernel. The + # acquire remains before stage/peer loads, which may consume its writes. + if USE_PDL: + tl.extra.cuda.gdc_wait() + + for owner in tl.static_range(WORLD): + values = tl.load( + stage_ptr + token * stride_stage_t + owner * CS + split * CSS + elem, + mask=elem_mask, + other=0.0, + ) + packed = _pack_bf16_pairs(values) + base = tl.load(ptrs + owner).to(tl.pointer_type(tl.uint32)) + dst = (token * WORLD + RANK) * (CS // 2) + split * (CSS // 2) + pair + tl.store(base + peer_offset_u32 + dst, packed, mask=pair_mask) + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _reduce_insert_kernel( + input_peer_ptrs, + input_peer_offset_u32, + cache_ptr, + slot_ptr, + stride_cache_block, + stride_cache_head, + stride_cache_token, + stride_cache_dim, + block_size, + C: tl.constexpr, + CS: tl.constexpr, + CS_P2: tl.constexpr, + SPLITS: tl.constexpr, + HEAD_SIZE: tl.constexpr, + CACHE_OFFSET: tl.constexpr, + RANK: tl.constexpr, + WORLD: tl.constexpr, + USE_PDL: tl.constexpr, + launch_pdl: tl.constexpr, +): + """Consume all partials for this rank and insert the reduced cache row.""" + token = tl.program_id(0).to(tl.int64) + split = tl.program_id(1).to(tl.int64) + CSS: tl.constexpr = CS // SPLITS + source = tl.arange(0, WORLD) + pair = tl.arange(0, CS_P2 // 2) + pair_mask = pair < CSS // 2 + offsets = ( + (token * WORLD + source)[:, None] * (CS // 2) + + split * (CSS // 2) + + pair[None, :] + ) + mask = tl.full([WORLD], True, tl.int1)[:, None] & pair_mask[None, :] + input_ptrs = input_peer_ptrs.to(tl.pointer_type(tl.uint64)) + input_u32 = tl.load(input_ptrs + RANK).to(tl.pointer_type(tl.uint32)) + input_u32 += input_peer_offset_u32 + # Slot metadata and the cache destination do not depend on input publish. + slot = tl.load(slot_ptr + token) + valid = slot >= 0 + channel = tl.arange(0, CS_P2) + channel_mask = channel < CSS + global_channel = split * CSS + channel + head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1) + dim = CACHE_OFFSET + global_channel % HEAD_SIZE + safe_slot = tl.maximum(slot, 0).to(tl.int64) + dst = ( + cache_ptr + + (safe_slot // block_size) * stride_cache_block + + head * stride_cache_head + + (safe_slot % block_size) * stride_cache_token + + dim * stride_cache_dim + ) + if USE_PDL: + tl.extra.cuda.gdc_wait() + packed = _wait_pairs(input_u32, offsets, mask) + + lo = (packed & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True) + hi = (packed >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True) + reduced = tl.interleave( + tl.sum(lo.to(tl.float32), axis=0).to(tl.bfloat16), + tl.sum(hi.to(tl.float32), axis=0).to(tl.bfloat16), + ) + tl.store( + input_u32 + offsets, + tl.full([WORLD, CS_P2 // 2], _EMPTY_PAIR, tl.uint32), + mask=mask, + ) + + tl.store(dst, reduced, mask=valid & channel_mask) + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _sconv_publish_kernel( + peer_ptrs, + peer_offset_u32, + residual_ptr, + weight_ptr, + cache_ptr, + position_ptr, + sequence_ptr, + slot_ptr, + block_table_ptr, + stride_residual_t, + stride_cache_block, + stride_cache_head, + stride_cache_token, + stride_cache_dim, + stride_block_table_r, + max_blocks, + block_size, + C: tl.constexpr, + CS: tl.constexpr, + CS_P2: tl.constexpr, + SPLITS: tl.constexpr, + HEAD_SIZE: tl.constexpr, + CACHE_OFFSET: tl.constexpr, + RANK: tl.constexpr, + WORLD: tl.constexpr, + WINDOW: tl.constexpr, + USE_PDL: tl.constexpr, + launch_pdl: tl.constexpr, +): + """Compute this rank's shard, then publish it to every rank.""" + tl.static_assert(WINDOW == 4) + token = tl.program_id(0).to(tl.int64) + split = tl.program_id(1).to(tl.int64) + CSS: tl.constexpr = CS // SPLITS + channel = tl.arange(0, CS_P2) + channel_mask = channel < CSS + global_channel = split * CSS + channel + head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1) + dim = CACHE_OFFSET + global_channel % HEAD_SIZE + slot = tl.load(slot_ptr + token) + valid = slot >= 0 + position = tl.load(position_ptr + token) + sequence = tl.load(sequence_ptr + token) + safe_slot = tl.maximum(slot, 0).to(tl.int64) + own_ptr = ( + cache_ptr + + (safe_slot // block_size) * stride_cache_block + + head * stride_cache_head + + (safe_slot % block_size) * stride_cache_token + + dim * stride_cache_dim + ) + # These loads were made visible before reduce/insert could signal us and + # are independent of its cache store. Hoisting them gives PDL useful work + # to overlap while retaining the acquire before every cache read. + residual = tl.load( + residual_ptr + token * stride_residual_t + RANK * CS + global_channel, + mask=channel_mask, + other=0.0, + ) + weight0 = tl.load( + weight_ptr + global_channel * WINDOW, + mask=channel_mask, + other=0.0, + ).to(tl.float32) + weight1 = tl.load( + weight_ptr + global_channel * WINDOW + 1, + mask=channel_mask, + other=0.0, + ).to(tl.float32) + weight2 = tl.load( + weight_ptr + global_channel * WINDOW + 2, + mask=channel_mask, + other=0.0, + ).to(tl.float32) + weight3 = tl.load( + weight_ptr + global_channel * WINDOW + 3, + mask=channel_mask, + other=0.0, + ).to(tl.float32) + ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64)) + if USE_PDL: + tl.extra.cuda.gdc_wait() + current = tl.load(own_ptr, mask=valid & channel_mask, other=0.0) + + conv = tl.zeros([CS_P2], tl.float32) + for tap_idx in tl.static_range(WINDOW): + source_position = position - (WINDOW - 1) + tap_idx + take = valid & (source_position >= 0) + if tap_idx == WINDOW - 1: + value = tl.where(take, current.to(tl.float32), 0.0) + else: + safe_position = tl.maximum(source_position, 0) + logical_block = tl.minimum(safe_position // block_size, max_blocks - 1) + physical_block = tl.load( + block_table_ptr + sequence * stride_block_table_r + logical_block, + mask=take, + other=0, + ).to(tl.int64) + source_ptr = ( + cache_ptr + + physical_block * stride_cache_block + + head * stride_cache_head + + (safe_position % block_size) * stride_cache_token + + dim * stride_cache_dim + ) + cached = tl.load(source_ptr, mask=take & channel_mask, other=0.0) + value = tl.where(take, cached.to(tl.float32), 0.0) + if tap_idx == 0: + weight = weight0 + elif tap_idx == 1: + weight = weight1 + elif tap_idx == 2: + weight = weight2 + else: + weight = weight3 + conv += value * weight + + # Preserve both bf16 rounding points of the original sublayer. + short_conv_with_skip = (conv + current.to(tl.float32)).to(tl.bfloat16) + output = (residual.to(tl.float32) + short_conv_with_skip.to(tl.float32)).to( + tl.bfloat16 + ) + + pair = tl.arange(0, CS_P2 // 2) + pair_mask = pair < CSS // 2 + packed = _pack_bf16_pairs(output) + row_offset = token * (C // 2) + RANK * (CS // 2) + split * (CSS // 2) + pair + + for destination in tl.static_range(WORLD): + base = tl.load(ptrs + destination).to(tl.pointer_type(tl.uint32)) + tl.store(base + peer_offset_u32 + row_offset, packed, mask=pair_mask) + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _gather_norm_kernel( + output_peer_ptrs, + output_peer_offset_u32, + norm_weight_ptr, + normed_ptr, + residual_out_ptr, + eps, + stride_output_t, + C: tl.constexpr, + C_P2: tl.constexpr, + RANK: tl.constexpr, + HAS_NORM: tl.constexpr, + USE_PDL: tl.constexpr, + launch_pdl: tl.constexpr, +): + """Consume a complete row; one CTA owns all outputs for one token.""" + token = tl.program_id(0).to(tl.int64) + pair = tl.arange(0, C_P2 // 2) + pair_mask = pair < C // 2 + output_ptrs = output_peer_ptrs.to(tl.pointer_type(tl.uint64)) + output_u32 = tl.load(output_ptrs + RANK).to(tl.pointer_type(tl.uint32)) + output_u32 += output_peer_offset_u32 + offsets = token * (C // 2) + pair + channel = tl.arange(0, C_P2) + channel_mask = channel < C + # Gamma is independent of the preceding sconv publication. Keep the + # acquire immediately before polling the Lamport output slots. + if HAS_NORM: + weight = tl.load(norm_weight_ptr + channel, mask=channel_mask, other=0.0) + if USE_PDL: + tl.extra.cuda.gdc_wait() + packed = _wait_pairs(output_u32, offsets, pair_mask) + row = _unpack_bf16_pairs(packed) + tl.store( + residual_out_ptr + token * stride_output_t + channel, row, mask=channel_mask + ) + if HAS_NORM: + row_f32 = tl.where(channel_mask, row.to(tl.float32), 0.0) + inv_rms = tl.rsqrt(tl.sum(row_f32 * row_f32, axis=0) / C + eps) + tl.store( + normed_ptr + token * stride_output_t + channel, + (row_f32 * inv_rms * weight.to(tl.float32)).to(tl.bfloat16), + mask=channel_mask, + ) + tl.store( + output_u32 + offsets, + tl.full([C_P2 // 2], _EMPTY_PAIR, tl.uint32), + mask=pair_mask, + ) + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _validate_lamport_init_kernel( + input_peer_ptrs, + output_peer_ptrs, + bad_ptr, + num_pairs, + RANK: tl.constexpr, + BLOCK: tl.constexpr, +): + """Validate both complete local allocations through their fabric pointers.""" + offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < num_pairs + ptrs_in = input_peer_ptrs.to(tl.pointer_type(tl.uint64)) + ptrs_out = output_peer_ptrs.to(tl.pointer_type(tl.uint64)) + local_in = tl.load(ptrs_in + RANK).to(tl.pointer_type(tl.uint32)) + local_out = tl.load(ptrs_out + RANK).to(tl.pointer_type(tl.uint32)) + value_in = tl.load(local_in + offsets, mask=mask, other=_EMPTY_PAIR) + value_out = tl.load(local_out + offsets, mask=mask, other=_EMPTY_PAIR) + bad = tl.max( + tl.where(mask & ((value_in != _EMPTY_PAIR) | (value_out != _EMPTY_PAIR)), 1, 0) + ) + if bad != 0: + tl.atomic_max(bad_ptr, 1) + + +class LamportRSConv: + """Persistent symmetric buffers for one TP group.""" + + def _initialize_lamport_buffers(self) -> None: + """Arm every slot and collectively verify the exact sentinel bits. + + Do not replace this with a zero-fill: Lamport readiness distinguishes + the bf16 bit pattern for -0.0 (0x8000) from every published payload. + The validation is deliberately collective so one rank cannot enter the + first polling kernel while another rank still has an unarmed buffer. + """ + if not self.buf_in.is_contiguous() or not self.buf_out.is_contiguous(): + raise RuntimeError("Lamport symmetric buffers must be contiguous") + if self.buf_in.numel() % 2 or self.buf_out.numel() % 2: + raise RuntimeError("Lamport buffers must contain whole uint32 pairs") + + # Fill through int16 so each bf16 lane receives the exact -0.0 bits. + # This covers all three generations and all max-token slots, including + # slots that a smaller first invocation does not touch. + self.buf_in.view(torch.int16).fill_(-0x8000) + self.buf_out.view(torch.int16).fill_(-0x8000) + torch.accelerator.synchronize(self.device) + self.tp.barrier() + + # Scan the complete local allocations. Since every rank performs the + # scan and participates in MAX, success proves every symmetric backing + # allocation was armed before any rank is allowed to use generation 0. + bad = torch.logical_or( + self.buf_in.view(torch.int16).ne(-0x8000).any(), + self.buf_out.view(torch.int16).ne(-0x8000).any(), + ).to(dtype=torch.float32) + bad = self.tp.all_reduce(bad) + if int(bad.item()) != 0: + raise RuntimeError("Lamport sentinel initialization failed on a TP rank") + self.tp.barrier() + + def _initialize_mnnvl_buffers(self) -> None: + """Initialize and validate FlashInfer fabric-mapped Lamport storage.""" + self._mnnvl_input_handle.lamport_initialize(self.rank, torch.bfloat16) + self._mnnvl_output_handle.lamport_initialize(self.rank, torch.bfloat16) + torch.accelerator.synchronize(self.device) + self.tp.barrier() + + bad = torch.zeros((), dtype=torch.int32, device=self.device) + num_pairs = self.num_buffers * self.max_tokens * self.hidden_size // 2 + _validate_lamport_init_kernel[(triton.cdiv(num_pairs, 256),)]( + self.input_peer_ptrs, + self.output_peer_ptrs, + bad, + num_pairs, + RANK=self.rank, + BLOCK=256, + num_warps=4, + ) + bad = self.tp.all_reduce(bad.to(torch.float32)) + if int(bad.item()) != 0: + raise RuntimeError("MNNVL Lamport sentinel initialization failed") + self.tp.barrier() + + def __init__( + self, hidden_size: int, window_size: int, max_tokens: int = _MAX_TOKENS + ) -> None: + import torch.distributed._symmetric_memory as symm_mem + + tp = get_tp_group() + self.tp = tp + self.group = tp.device_group + self.world_size = tp.world_size + self.rank = tp.rank_in_group + self.device = torch.device(tp.device) + if self.world_size not in (2, 4, 8): + raise ValueError(f"TP world size must be 2, 4, or 8, got {self.world_size}") + if hidden_size % (2 * self.world_size) != 0: + raise ValueError("hidden size must produce an even shard on every rank") + if window_size != 4: + raise ValueError(f"short-conv window size must be 4, got {window_size}") + if max_tokens < 1 or max_tokens > _MAX_TOKENS: + raise ValueError(f"max_tokens must be in [1, {_MAX_TOKENS}]") + + is_cross_node = not all(in_the_same_node_as(tp.cpu_group)) + self.hidden_size = hidden_size + self.window_size = window_size + self.max_tokens = max_tokens + self.shard_size = hidden_size // self.world_size + # Three generations follow FlashInfer's Lamport layout. A generation + # is reused only after two intervening collective calls. + self.num_buffers = 3 + self.input_generation_bytes = max_tokens * hidden_size * 2 + self.output_generation_bytes = max_tokens * hidden_size * 2 + self.use_pdl = torch.cuda.get_device_capability(self.device)[0] >= 9 + if is_cross_node: + try: + from flashinfer.comm.mnnvl import ( + McastGPUBuffer, + TorchDistBackend, + is_mnnvl_fabric_supported, + ) + except ImportError as error: + raise RuntimeError( + "cross-node TP requires FlashInfer MNNVL support" + ) from error + + local_supported = int( + is_mnnvl_fabric_supported(torch.accelerator.current_device_index()) + ) + unsupported = torch.tensor( + 1 - local_supported, dtype=torch.float32, device=self.device + ) + unsupported = tp.all_reduce(unsupported) + if int(unsupported.item()) != 0: + raise RuntimeError("cross-node TP is supported only on MNNVL fabric") + + comm_backend = TorchDistBackend(self.group) + allocation_bytes = self.num_buffers * max_tokens * hidden_size * 2 + self._mnnvl_input_handle = McastGPUBuffer( + allocation_bytes, + self.world_size, + self.rank, + self.device, + comm_backend, + ) + self._mnnvl_output_handle = McastGPUBuffer( + allocation_bytes, + self.world_size, + self.rank, + self.device, + comm_backend, + ) + self.input_peer_ptrs = self._mnnvl_input_handle.get_buffer_ptrs_dev() + self.output_peer_ptrs = self._mnnvl_output_handle.get_buffer_ptrs_dev() + self._initialize_mnnvl_buffers() + logger.info("using FlashInfer fabric-mapped MNNVL Lamport buffers") + else: + self.buf_in = symm_mem.empty( + self.num_buffers, + max_tokens, + self.world_size, + self.shard_size, + dtype=torch.bfloat16, + device=self.device, + ) + self.buf_out = symm_mem.empty( + self.num_buffers, + max_tokens, + hidden_size, + dtype=torch.bfloat16, + device=self.device, + ) + group_name = self.group.group_name + input_handle = symm_mem.rendezvous(self.buf_in, group_name) + output_handle = symm_mem.rendezvous(self.buf_out, group_name) + self.input_peer_ptrs = input_handle.buffer_ptrs_dev + self.output_peer_ptrs = output_handle.buffer_ptrs_dev + self._input_handle = input_handle + self._output_handle = output_handle + self._initialize_lamport_buffers() + self.generation = 0 + + def usable(self, num_tokens: int) -> bool: + return 0 < num_tokens <= self.max_tokens + + def rs_sconv_ag_add_norm( + self, + input_tensor: torch.Tensor, + residual: torch.Tensor, + conv_weight: torch.Tensor, + norm_weight: torch.Tensor | None, + eps: float, + cache: torch.Tensor, + positions: torch.Tensor, + block_table: torch.Tensor, + seq_idx: torch.Tensor, + slot_mapping: torch.Tensor, + off_s: int, + ws: int, + block_size: int, + ) -> tuple[torch.Tensor | None, torch.Tensor]: + """Return ``(normed | None, new_residual)``, both shaped ``[T, 6144]``.""" + tokens, hidden_size = residual.shape + if not self.usable(tokens): + raise ValueError(f"num_tokens must be in [1, {self.max_tokens}]") + if hidden_size != self.hidden_size or residual.dtype != torch.bfloat16: + raise ValueError("residual must be bf16 [T, 6144]") + if ( + input_tensor.shape != residual.shape + or input_tensor.dtype != torch.bfloat16 + or input_tensor.stride(1) != 1 + ): + raise ValueError("input_tensor must be channel-contiguous bf16 [T, 6144]") + shard_size = hidden_size // self.world_size + if conv_weight.shape != (shard_size, self.window_size): + raise ValueError( + f"conv_weight must have shape [{shard_size}, {self.window_size}]" + ) + if ( + conv_weight.dtype != torch.bfloat16 + or conv_weight.stride(0) != self.window_size + ): + raise ValueError("conv_weight must be contiguous bf16") + if norm_weight is not None and ( + norm_weight.shape != (hidden_size,) or norm_weight.dtype != torch.bfloat16 + ): + raise ValueError("norm_weight must be bf16 [6144] or None") + if cache.dtype != torch.bfloat16 or cache.ndim != 4: + raise ValueError("cache must be a 4-D bf16 tensor") + if shard_size % ws != 0 or cache.shape[1] != shard_size // ws: + raise ValueError("cache head layout is inconsistent with ws") + if off_s < 0 or off_s + ws > cache.shape[3]: + raise ValueError("cache channel offset is out of bounds") + + index = self.generation + input_offset = index * self.input_generation_bytes // 4 + output_offset = index * self.output_generation_bytes // 4 + normed = torch.empty_like(residual) if norm_weight is not None else None + residual_out = torch.empty_like(residual) + phase_splits = 8 if tokens == 1 and shard_size % 8 == 0 else 1 + phase_tile_p2 = triton.next_power_of_2(shard_size // phase_splits) + phase_grid = (tokens, phase_splits) + # Wide CTAs win before the grid saturates; smaller CTAs reduce register + # pressure once high-throughput batches provide enough parallelism. + if 128 <= tokens <= 2048: + phase_warps = 16 + elif tokens > 2048: + phase_warps = 8 + else: + phase_warps = 4 + gather_warps = 4 if tokens >= 256 else 8 + _publish_input_kernel[phase_grid]( + input_tensor, + self.input_peer_ptrs, + input_offset, + input_tensor.stride(0), + C=hidden_size, + CS=shard_size, + CS_P2=phase_tile_p2, + SPLITS=phase_splits, + RANK=self.rank, + WORLD=self.world_size, + USE_PDL=self.use_pdl, + launch_pdl=self.use_pdl, + num_warps=phase_warps, + ) + _reduce_insert_kernel[phase_grid]( + self.input_peer_ptrs, + input_offset, + cache, + slot_mapping, + cache.stride(0), + cache.stride(1), + cache.stride(2), + cache.stride(3), + block_size, + C=hidden_size, + CS=shard_size, + CS_P2=phase_tile_p2, + SPLITS=phase_splits, + HEAD_SIZE=ws, + CACHE_OFFSET=off_s, + RANK=self.rank, + WORLD=self.world_size, + USE_PDL=self.use_pdl, + launch_pdl=self.use_pdl, + num_warps=phase_warps, + ) + _sconv_publish_kernel[phase_grid]( + self.output_peer_ptrs, + output_offset, + residual, + conv_weight, + cache, + positions, + seq_idx, + slot_mapping, + block_table, + residual.stride(0), + cache.stride(0), + cache.stride(1), + cache.stride(2), + cache.stride(3), + block_table.stride(0), + block_table.shape[1], + block_size, + C=hidden_size, + CS=shard_size, + CS_P2=phase_tile_p2, + SPLITS=phase_splits, + HEAD_SIZE=ws, + CACHE_OFFSET=off_s, + RANK=self.rank, + WORLD=self.world_size, + WINDOW=self.window_size, + USE_PDL=self.use_pdl, + launch_pdl=self.use_pdl, + num_warps=phase_warps, + ) + _gather_norm_kernel[(tokens,)]( + self.output_peer_ptrs, + output_offset, + norm_weight if norm_weight is not None else residual, + normed if normed is not None else residual_out, + residual_out, + eps, + residual.stride(0), + C=hidden_size, + C_P2=triton.next_power_of_2(hidden_size), + RANK=self.rank, + HAS_NORM=norm_weight is not None, + USE_PDL=self.use_pdl, + launch_pdl=self.use_pdl, + num_warps=gather_warps, + ) + self.generation = (index + 1) % self.num_buffers + return normed, residual_out + + +_STATE: LamportRSConv | None = None +_STATE_FAILED = False + + +def initialize_lamport_rs_conv( + hidden_size: int, window_size: int, max_num_batched_tokens: int +) -> None: + """Collectively initialize the TP-group state during model construction.""" + global _STATE, _STATE_FAILED + if _STATE is not None: + if _STATE.hidden_size != hidden_size or _STATE.window_size != window_size: + raise RuntimeError("all Lamport users must share hidden and window sizes") + return + if _STATE_FAILED or os.environ.get("LAMPORT_RS_SCONV", "1") == "0": + return + try: + max_tokens = min(_MAX_TOKENS, max_num_batched_tokens) + _STATE = LamportRSConv(hidden_size, window_size, max_tokens=max_tokens) + except Exception: + _STATE_FAILED = True + logger.exception("fused collective unavailable; use the NCCL fallback") + + +def get_lamport_rs_conv(hidden_size: int, window_size: int) -> LamportRSConv | None: + """Return the state initialized with the model, or ``None`` for fallback.""" + if _STATE is not None and ( + _STATE.hidden_size != hidden_size or _STATE.window_size != window_size + ): + raise RuntimeError("all Lamport users must share hidden and window sizes") + return _STATE diff --git a/vllm/models/inkling/amd/ops/mm_towers.py b/vllm/models/inkling/amd/ops/mm_towers.py new file mode 100644 index 00000000000..021cf5274bd --- /dev/null +++ b/vllm/models/inkling/amd/ops/mm_towers.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused CUDA kernels for the Inkling vision/audio towers. + +Both kernels keep the reference paths' fp32 accumulation and per-op bf16 +rounding points (native ``rms_norm`` / ``F.gelu``); outputs are frequently +bit-identical and otherwise differ by 1-2 bf16 ulps from reduction-order +(real-checkpoint-weight cosine vs reference > 0.9999998). +""" + +from __future__ import annotations + +import torch + +from vllm.triton_utils import tl, tldevice, triton + +from .norm import _get_num_warps_from_block_size + + +@triton.jit +def _dmel_embed_sum_norm_kernel( + idx_ptr, # [T, NB] int32 dMel bin indices (values in [0, VOCAB)) + w_ptr, # [NB * VOCAB, D] bf16 embedding table + norm_w_ptr, # [D] (unused if HAS_NORM=False) + out_ptr, # [T, D] bf16 + eps, + D, + stride_idx_t, + NB: tl.constexpr, + VOCAB: tl.constexpr, + D_P2: tl.constexpr, + HAS_NORM: tl.constexpr, +): + t = tl.program_id(0).to(tl.int64) + offs = tl.arange(0, D_P2) + mask = offs < D + # One embedding row per mel bin (bin b uses table rows [b*VOCAB, (b+1)*VOCAB)), + # summed in fp32 (matches torch's fp32-accumulated bf16 .sum()). + acc = tl.zeros([D_P2], dtype=tl.float32) + for b in tl.static_range(NB): + v = tl.load(idx_ptr + t * stride_idx_t + b) + row = (b * VOCAB + v).to(tl.int64) + acc += tl.load(w_ptr + row * D + offs, mask=mask, other=0.0).to(tl.float32) + h = acc.to(tl.bfloat16) + if HAS_NORM: + # Match ir.ops.rms_norm: fp32 variance/normalize, then a single-rounded + # bf16 multiply with the bf16 weight. + x32 = h.to(tl.float32) + var = tl.sum(tl.where(mask, x32 * x32, 0.0), axis=0) / D + xn = (x32 * tl.math.rsqrt(var + eps)).to(tl.bfloat16) + w = tl.load(norm_w_ptr + offs, mask=mask, other=0.0) + h = xn * w + tl.store(out_ptr + t * D + offs, h, mask=mask) + + +def dmel_embed_sum_norm( + dmel_idx: torch.Tensor, # [T, NB] int32 + weight: torch.Tensor, # [NB * VOCAB, D] bf16 + norm_weight: torch.Tensor | None, + eps: float, +) -> torch.Tensor: + """``rmsnorm(sum_b weight[b * VOCAB + idx[:, b]])`` in one launch (no + [T, NB, D] intermediate).""" + T, nb = dmel_idx.shape + D = weight.shape[1] + assert weight.shape[0] % nb == 0 + vocab = weight.shape[0] // nb + out = torch.empty((T, D), dtype=weight.dtype, device=weight.device) + if T == 0: + return out + d_p2 = triton.next_power_of_2(D) + _dmel_embed_sum_norm_kernel[(T,)]( + dmel_idx, + weight, + norm_weight if norm_weight is not None else weight, + out, + eps, + D, + dmel_idx.stride(0), + NB=nb, + VOCAB=vocab, + D_P2=d_p2, + HAS_NORM=norm_weight is not None, + # Swept on GB200: 8 warps beats the block-size heuristic's 16 by ~4% + # at large T (the 80-row gather chain is latency- not lane-bound). + num_warps=8, + ) + return out + + +@triton.jit +def _rmsnorm_gelu_kernel( + x_ptr, # [R, D] bf16 + w_ptr, # [D] + out_ptr, # [R, D] bf16 (or the folded layout when FOLD) + eps, + R, + D, + D_P2: tl.constexpr, + BLOCK_M: tl.constexpr, # rows per block (>1 for small D) + HAS_GELU: tl.constexpr, + FOLD: tl.constexpr, + # fold geometry: input rows index [N, T, H, W]; the store scatters each + # row to (out_row, slot) of fold_timespace_to_depth's output layout. + FT: tl.constexpr, + FH: tl.constexpr, + FW: tl.constexpr, + TF: tl.constexpr, + HF: tl.constexpr, +): + pid = tl.program_id(0).to(tl.int64) + rows = pid * BLOCK_M + tl.arange(0, BLOCK_M) + rmask = rows < R + offs = tl.arange(0, D_P2) + mask = rmask[:, None] & (offs < D)[None, :] + x32 = tl.load(x_ptr + rows[:, None] * D + offs[None, :], mask=mask, other=0.0).to( + tl.float32 + ) + var = tl.sum(x32 * x32, axis=1) / D + xn = (x32 * tl.math.rsqrt(var + eps)[:, None]).to(tl.bfloat16) + w = tl.load(w_ptr + offs, mask=offs < D, other=0.0) + h = xn * w[None, :] # bf16 multiply, matching ir.ops.rms_norm + if HAS_GELU: + # Exact (erf) GELU on the bf16-rounded norm output, fp32 math, + # matching F.gelu's opmath on a bf16 tensor. + g32 = h.to(tl.float32) + h = (0.5 * g32 * (1.0 + tldevice.erf(g32 * 0.7071067811865476))).to(tl.bfloat16) + if FOLD: + # Store directly into the next layer's folded layout (a pure + # permutation — replaces the separate fold copy pass). + t = (rows // (FH * FW)) % FT + hh = (rows // FW) % FH + ww = rows % FW + n = rows // (FT * FH * FW) + slot = ((t % TF) * HF + hh % HF) * HF + ww % HF + out_row = ((n * (FT // TF) + t // TF) * (FH // HF) + hh // HF) * ( + FW // HF + ) + ww // HF + base = (out_row * (TF * HF * HF) + slot) * D + tl.store(out_ptr + base[:, None] + offs[None, :], h, mask=mask) + else: + tl.store(out_ptr + rows[:, None] * D + offs[None, :], h, mask=mask) + + +def rmsnorm_gelu( + x: torch.Tensor, # [..., D] bf16 contiguous + weight: torch.Tensor, + eps: float, + gelu: bool = True, + fold: tuple[int, int] | None = None, # (t_fold, hw_fold) of the NEXT fold +) -> torch.Tensor: + """Fused ``gelu(rmsnorm(x))`` (or plain rmsnorm); multiple rows per block + when D is small. With ``fold``, x must be [N, T, H, W, D] and the output + comes back as ``fold_timespace_to_depth(result, *fold)``.""" + D = x.shape[-1] + flat = x.reshape(-1, D) + assert flat.stride(1) == 1 and flat.stride(0) == D + R = flat.shape[0] + if fold is None: + out = torch.empty_like(flat) + ft = fh = fw = tf = hf = 1 + out_shape = x.shape + else: + tf, hf = fold + N, ft, fh, fw, _ = x.shape + out_shape = (N, ft // tf, fh // hf, fw // hf, tf * hf * hf * D) + out = torch.empty(out_shape, dtype=x.dtype, device=x.device) + if R == 0: + return out.reshape(out_shape) + d_p2 = triton.next_power_of_2(D) + block_m = max(1, 4096 // d_p2) + _rmsnorm_gelu_kernel[(triton.cdiv(R, block_m),)]( + flat, + weight, + out, + eps, + R, + D, + D_P2=d_p2, + BLOCK_M=block_m, + HAS_GELU=gelu, + FOLD=fold is not None, + FT=ft, + FH=fh, + FW=fw, + TF=tf, + HF=hf, + num_warps=_get_num_warps_from_block_size(d_p2 * block_m), + ) + return out.reshape(out_shape) diff --git a/vllm/models/inkling/amd/ops/norm.py b/vllm/models/inkling/amd/ops/norm.py new file mode 100644 index 00000000000..d34c1c3259b --- /dev/null +++ b/vllm/models/inkling/amd/ops/norm.py @@ -0,0 +1,414 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +from functools import lru_cache + +import torch + +from vllm.triton_utils import tl, triton + +_MAX_FUSED_SIZE = 65536 + + +def _get_num_warps_from_block_size(block_size: int) -> int: + if block_size >= 32768: + return 32 + if block_size >= 8192: + return 16 + if block_size >= 2048: + return 8 + return 4 + + +def _largest_power_of_2(n: int) -> int: + assert n > 0, f"{n=}" + return 1 << (n.bit_length() - 1) + + +@lru_cache(maxsize=128) +def _get_grid_size_for_mem_bw_kernel(device: torch.device, factor: int = 8) -> int: + num_sms = torch.cuda.get_device_properties(device).multi_processor_count + return _largest_power_of_2(num_sms) * factor + + +@triton.jit +def _rmsnorm_fwd_kernel( + x_ptr, + weight_ptr, + y_ptr, + rstd_ptr, + eps, + x_stride_0, + y_stride_0, + n_cols, + block_size_n: tl.constexpr, +): + pid_m = tl.program_id(0).to(tl.int64) + offs_n = tl.arange(0, block_size_n) + mask_n = offs_n < n_cols + + weight = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + x = tl.load(x_ptr + pid_m * x_stride_0 + offs_n, mask=mask_n, other=0.0).to( + tl.float32 + ) + row_var = tl.sum(x * x, axis=0) / n_cols + rstd = tl.math.rsqrt(row_var + eps) + tl.store(rstd_ptr + pid_m, rstd) + + y = x * rstd * weight + tl.store(y_ptr + pid_m * y_stride_0 + offs_n, y, mask=mask_n) + + +@triton.jit(do_not_specialize=["n_rows"]) +def _rmsnorm_fwd_kernel_block_m( + x_ptr, + weight_ptr, + y_ptr, + rstd_ptr, + eps, + x_stride_0, + y_stride_0, + n_rows, + n_cols, + block_size_m: tl.constexpr, + block_size_n: tl.constexpr, +): + pid_m = tl.program_id(0).to(tl.int64) + offs_n = tl.arange(0, block_size_n) + mask_n = offs_n < n_cols + + num_blocks_m = tl.cdiv(n_rows, block_size_m) + blocks_per_pid = tl.cdiv(num_blocks_m, tl.num_programs(0)) + block_id_start = pid_m * blocks_per_pid + block_id_end = min(block_id_start + blocks_per_pid, num_blocks_m) + + weight = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + + for block_id in range(block_id_start, block_id_end): + offs_m = block_id * block_size_m + tl.arange(0, block_size_m) + mask_m = offs_m < n_rows + mask_mn = mask_m[:, None] & mask_n[None, :] + x = tl.load( + x_ptr + offs_m[:, None] * x_stride_0 + offs_n[None, :], + mask=mask_mn, + other=0.0, + ).to(tl.float32) + row_var = tl.sum(x * x, axis=1) / n_cols + rstd = tl.math.rsqrt(row_var + eps) + tl.store(rstd_ptr + offs_m, rstd, mask=mask_m) + + y = x * rstd[:, None] * weight + tl.store( + y_ptr + offs_m[:, None] * y_stride_0 + offs_n[None, :], + y, + mask=mask_mn, + ) + + +@triton.jit +def _add_rmsnorm_fwd_kernel( + res_ptr, # [T, N] residual (read) + delta_ptr, # [T, N] delta to add (read) + weight_ptr, + y_ptr, # [T, N] normed output + res_out_ptr, # [T, N] updated residual output + eps, + res_stride_0, + delta_stride_0, + y_stride_0, + res_out_stride_0, + n_cols, + block_size_n: tl.constexpr, +): + pid_m = tl.program_id(0).to(tl.int64) + offs_n = tl.arange(0, block_size_n) + mask_n = offs_n < n_cols + + weight = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + r = tl.load(res_ptr + pid_m * res_stride_0 + offs_n, mask=mask_n, other=0.0).to( + tl.float32 + ) + d = tl.load(delta_ptr + pid_m * delta_stride_0 + offs_n, mask=mask_n, other=0.0).to( + tl.float32 + ) + # Round the sum to the residual dtype first (matches the eager + # `residual + delta` then rmsnorm-on-bf16 sequence bit-for-bit). + s = (r + d).to(res_out_ptr.dtype.element_ty) + tl.store(res_out_ptr + pid_m * res_out_stride_0 + offs_n, s, mask=mask_n) + x = s.to(tl.float32) + row_var = tl.sum(x * x, axis=0) / n_cols + rstd = tl.math.rsqrt(row_var + eps) + y = x * rstd * weight + tl.store(y_ptr + pid_m * y_stride_0 + offs_n, y, mask=mask_n) + + +def add_rmsnorm( + residual: torch.Tensor, + delta: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Fused ``res = residual + delta; y = rmsnorm(res)``. + + Returns ``(y, res)``; both are fresh tensors (cudagraph-friendly, no + in-place update of the inputs). + """ + assert residual.ndim == 2 and delta.ndim == 2, (residual.shape, delta.shape) + n_rows, n_cols = residual.shape + assert weight.shape[0] == n_cols + y = torch.empty_like(residual) + res_out = torch.empty_like(residual) + if n_rows == 0: + return y, res_out + + block_size_n = triton.next_power_of_2(n_cols) + max_block_size_n = _MAX_FUSED_SIZE // residual.element_size() + if max_block_size_n < block_size_n: + raise RuntimeError(f"Large {n_cols=} is not supported") + num_warps = _get_num_warps_from_block_size(block_size_n) + _add_rmsnorm_fwd_kernel[(n_rows,)]( + residual, + delta, + weight, + y, + res_out, + eps, + residual.stride(0), + delta.stride(0), + y.stride(0), + res_out.stride(0), + n_cols, + block_size_n, + num_warps=num_warps, + ) + return y, res_out + + +@triton.jit +def _embed_rmsnorm_kernel( + ids_ptr, # [T] token ids + table_ptr, # [V, N] embedding table + weight_ptr, # [N] (HAS_NORM only) + chain_weight_ptr, # [N] (HAS_CHAIN only) + out_ptr, # [T, N] rmsnorm(table[ids], weight) + chain_out_ptr, # [T, N] rmsnorm(out, chain_weight) (HAS_CHAIN only) + eps, + table_stride_0, + n_cols, + block_size_n: tl.constexpr, + HAS_NORM: tl.constexpr, + HAS_CHAIN: tl.constexpr, +): + pid_m = tl.program_id(0).to(tl.int64) + offs_n = tl.arange(0, block_size_n) + mask_n = offs_n < n_cols + row = tl.load(ids_ptr + pid_m).to(tl.int64) + x = tl.load(table_ptr + row * table_stride_0 + offs_n, mask=mask_n, other=0.0) + if HAS_NORM: + xf = x.to(tl.float32) + rstd = tl.math.rsqrt(tl.sum(xf * xf, axis=0) / n_cols + eps) + w = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + # Round to the output dtype so the chained norm is bit-exact vs the + # unfused pair (which stores bf16 in between). + x = (xf * rstd * w).to(out_ptr.dtype.element_ty) + tl.store(out_ptr + pid_m * n_cols + offs_n, x, mask=mask_n) + if HAS_CHAIN: + xf = x.to(tl.float32) + rstd = tl.math.rsqrt(tl.sum(xf * xf, axis=0) / n_cols + eps) + w = tl.load(chain_weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + tl.store( + chain_out_ptr + pid_m * n_cols + offs_n, + (xf * rstd * w).to(chain_out_ptr.dtype.element_ty), + mask=mask_n, + ) + + +def embed_rmsnorm( + input_ids: torch.Tensor, + embed_table: torch.Tensor, + weight: torch.Tensor | None, + eps: float, + chain_weight: torch.Tensor | None = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Fused ``rmsnorm(embed_table[input_ids], weight)`` row gather + norm. + + Requires the full vocab on-rank (replicated or tp_size == 1). + ``weight=None`` skips the norm (``use_embed_norm=False``), leaving a pure + embedding-table gather. + ``chain_weight`` additionally emits ``rmsnorm(out, chain_weight)`` (the + first decoder layer's pre-attention norm) as a second output, still one + launch. Bit-exact vs the unfused module sequence.""" + ids = input_ids.view(-1) + (T,) = ids.shape + n = embed_table.shape[1] + out = torch.empty( + (*input_ids.shape, n), dtype=embed_table.dtype, device=embed_table.device + ) + chain_out = torch.empty_like(out) if chain_weight is not None else None + if T > 0: + block_size_n = triton.next_power_of_2(n) + _embed_rmsnorm_kernel[(T,)]( + ids, + embed_table, + weight if weight is not None else embed_table, + chain_weight if chain_weight is not None else embed_table, + out, + chain_out if chain_out is not None else out, + eps, + embed_table.stride(0), + n, + block_size_n, + HAS_NORM=weight is not None, + HAS_CHAIN=chain_weight is not None, + num_warps=_get_num_warps_from_block_size(block_size_n), + ) + if chain_out is not None: + return out, chain_out + return out + + +@triton.jit +def _embed_dual_rmsnorm_cat_kernel( + hidden_ptr, # [T, N] + emb_ptr, # [T, N] embeddings, or the [V, N] embedding table when GATHER + ids_ptr, # [T] token ids (GATHER only) + w_hidden_ptr, # [N] + w_pre_ptr, # [N] chained pre-norm on the embed side (HAS_PRE_NORM only) + w_embed_ptr, # [N] + out_ptr, # [T, 2N]: [rmsnorm(hidden) | rmsnorm(rmsnorm?(emb))] + eps, + hidden_stride_0, + emb_stride_0, + n_cols, + block_size_n: tl.constexpr, + GATHER: tl.constexpr, + HAS_PRE_NORM: tl.constexpr, +): + pid_m = tl.program_id(0).to(tl.int64) + which = tl.program_id(1) # 0 -> hidden into cols [0, N); 1 -> emb into [N, 2N) + offs_n = tl.arange(0, block_size_n) + mask_n = offs_n < n_cols + if which == 0: + x = tl.load( + hidden_ptr + pid_m * hidden_stride_0 + offs_n, mask=mask_n, other=0.0 + ).to(tl.float32) + w = tl.load(w_hidden_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + else: + row = tl.load(ids_ptr + pid_m).to(tl.int64) if GATHER else pid_m + x = tl.load(emb_ptr + row * emb_stride_0 + offs_n, mask=mask_n, other=0.0).to( + tl.float32 + ) + if HAS_PRE_NORM: + w_pre = tl.load(w_pre_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + rstd = tl.math.rsqrt(tl.sum(x * x, axis=0) / n_cols + eps) + # Round-trip through the output dtype so the chained norm is + # bit-exact vs the unfused pair (which stores bf16 in between). + x = (x * rstd * w_pre).to(out_ptr.dtype.element_ty).to(tl.float32) + w = tl.load(w_embed_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + rstd = tl.math.rsqrt(tl.sum(x * x, axis=0) / n_cols + eps) + tl.store( + out_ptr + pid_m * (2 * n_cols) + which * n_cols + offs_n, + (x * rstd * w).to(out_ptr.dtype.element_ty), + mask=mask_n, + ) + + +def embed_dual_rmsnorm_cat( + hidden: torch.Tensor, + hidden_weight: torch.Tensor, + embed_weight: torch.Tensor, + eps: float, + *, + embeds: torch.Tensor | None = None, + input_ids: torch.Tensor | None = None, + embed_table: torch.Tensor | None = None, + pre_norm_weight: torch.Tensor | None = None, +) -> torch.Tensor: + """The MTP depth-layer input in one launch: + ``cat([rmsnorm(hidden, w_h), rmsnorm(pre?(emb), w_e)], -1)``. + + The embed side is either a fused row gather ``embed_table[input_ids]`` + (draft decode steps) or precomputed ``embeds`` ([T, N], the target-merged + multimodal embeddings at draft prefill); ``pre_norm_weight`` chains the + backbone embed_norm in front of the depth embed_norm (bit-exact vs the + unfused sequence). The concat copies collapse into direct writes.""" + T, n = hidden.shape + if embeds is not None: + assert embeds.shape == hidden.shape + src, ids, src_stride = embeds, embeds, embeds.stride(0) + gather = False + else: + assert input_ids is not None and embed_table is not None + assert input_ids.shape == (T,) and embed_table.shape[1] == n + src, ids, src_stride = embed_table, input_ids, embed_table.stride(0) + gather = True + out = torch.empty((T, 2 * n), dtype=hidden.dtype, device=hidden.device) + if T == 0: + return out + block_size_n = triton.next_power_of_2(n) + _embed_dual_rmsnorm_cat_kernel[(T, 2)]( + hidden, + src, + ids, + hidden_weight, + pre_norm_weight if pre_norm_weight is not None else embed_weight, + embed_weight, + out, + eps, + hidden.stride(0), + src_stride, + n, + block_size_n, + GATHER=gather, + HAS_PRE_NORM=pre_norm_weight is not None, + num_warps=_get_num_warps_from_block_size(block_size_n), + ) + return out + + +def rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + assert x.ndim == 2, f"{x.shape=}" + assert weight.ndim == 1, f"{weight.shape=}" + n_rows, n_cols = x.shape + assert weight.shape[0] == n_cols, f"{weight.shape=} {x.shape=}" + y = torch.empty_like(x) + rstd = torch.empty((n_rows,), dtype=torch.float32, device=x.device) + + block_size_n = triton.next_power_of_2(n_cols) + max_block_size_n = _MAX_FUSED_SIZE // x.element_size() + if max_block_size_n < block_size_n: + raise RuntimeError(f"Large {n_cols=} is not supported") + block_size_m = max(1, 4096 // block_size_n) + num_warps = _get_num_warps_from_block_size(block_size_n) + + if block_size_m == 1: + _rmsnorm_fwd_kernel[(n_rows,)]( + x, + weight, + y, + rstd, + eps, + x.stride(0), + y.stride(0), + n_cols, + block_size_n, + num_warps=num_warps, + ) + else: + grid_size = _get_grid_size_for_mem_bw_kernel(x.device) + _rmsnorm_fwd_kernel_block_m[(grid_size,)]( + x, + weight, + y, + rstd, + eps, + x.stride(0), + y.stride(0), + n_rows, + n_cols, + block_size_m, + block_size_n, + num_warps=num_warps, + ) + return y diff --git a/vllm/models/inkling/amd/ops/qkvr_prep.py b/vllm/models/inkling/amd/ops/qkvr_prep.py new file mode 100644 index 00000000000..a910e72170d --- /dev/null +++ b/vllm/models/inkling/amd/ops/qkvr_prep.py @@ -0,0 +1,918 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import aux_stream + +LOW_BLOCK_M = 32 +LOW_BLOCK_N = 64 +LOW_NUM_WARPS = 4 +THROUGHPUT_BLOCK_M = 32 +THROUGHPUT_BLOCK_N = 128 +THROUGHPUT_GROUP_M = 2 +THROUGHPUT_NUM_WARPS = 4 +SMALL_TOKEN_THRESHOLD = 128 +SMALL_NUM_WARPS = 2 +Q_BLOCK_ROWS = 8 +Q_NUM_WARPS = 2 +KV_BLOCK_ROWS = 4 +KV_NUM_WARPS = 2 + + +@triton.jit(do_not_specialize=["rows"]) +def _rel_proj_low_latency_kernel( + qkvr_ptr, + rel_proj_ptr, + rel_out_ptr, + log_scaling_ptr, + rows, + stride_x_t, + R_OFFSET: tl.constexpr, + NUM_Q_HEADS: tl.constexpr, + REL_EXTENT: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + APPLY_LOG_SCALING: tl.constexpr, +): + row = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + col = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + inner = tl.arange(0, 16) + token = row // NUM_Q_HEADS + head = row % NUM_Q_HEADS + relative = tl.load( + qkvr_ptr + + token[:, None] * stride_x_t + + R_OFFSET + + head[:, None] * 16 + + inner[None, :], + mask=row[:, None] < rows, + other=0.0, + ) + projection = tl.load( + rel_proj_ptr + inner[:, None] * REL_EXTENT + col[None, :], + mask=col[None, :] < REL_EXTENT, + other=0.0, + ) + values = tl.dot(relative, projection, out_dtype=tl.float32).to( + rel_out_ptr.dtype.element_ty + ) + values = values.to(tl.float32) + if APPLY_LOG_SCALING: + values *= tl.load(log_scaling_ptr + token, mask=row < rows, other=1.0)[:, None] + tl.store( + rel_out_ptr + row[:, None] * REL_EXTENT + col[None, :], + values.to(rel_out_ptr.dtype.element_ty), + mask=(row[:, None] < rows) & (col[None, :] < REL_EXTENT), + ) + + +@triton.jit(do_not_specialize=["rows"]) +def _rel_proj_throughput_kernel( + qkvr_ptr, + rel_proj_ptr, + rel_out_ptr, + log_scaling_ptr, + rows, + stride_x_t, + R_OFFSET: tl.constexpr, + NUM_Q_HEADS: tl.constexpr, + REL_EXTENT: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + GROUP_M: tl.constexpr, + APPLY_LOG_SCALING: tl.constexpr, +): + row_group = tl.program_id(0) + col = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + inner = tl.arange(0, 16) + projection = tl.load( + rel_proj_ptr + inner[:, None] * REL_EXTENT + col[None, :], + mask=col[None, :] < REL_EXTENT, + other=0.0, + ) + row_offsets = tl.arange(0, BLOCK_M) + for group_offset in tl.static_range(GROUP_M): + row = (row_group * GROUP_M + group_offset) * BLOCK_M + row_offsets + token = row // NUM_Q_HEADS + head = row % NUM_Q_HEADS + relative = tl.load( + qkvr_ptr + + token[:, None] * stride_x_t + + R_OFFSET + + head[:, None] * 16 + + inner[None, :], + mask=row[:, None] < rows, + other=0.0, + ) + values = tl.dot(relative, projection, out_dtype=tl.float32).to( + rel_out_ptr.dtype.element_ty + ) + values = values.to(tl.float32) + if APPLY_LOG_SCALING: + values *= tl.load(log_scaling_ptr + token, mask=row < rows, other=1.0)[ + :, None + ] + tl.store( + rel_out_ptr + row[:, None] * REL_EXTENT + col[None, :], + values.to(rel_out_ptr.dtype.element_ty), + mask=(row[:, None] < rows) & (col[None, :] < REL_EXTENT), + ) + + +def use_rel_proj_throughput(rows: int, rel_extent: int) -> bool: + min_rows = 8192 if rel_extent == 512 else 2048 + return rows >= min_rows + + +def qkvr_rel_proj( + qkvr: torch.Tensor, + rel_proj: torch.Tensor, + rel_out: torch.Tensor, + log_scaling: torch.Tensor | None, + *, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + d_rel: int, +) -> None: + rows = qkvr.shape[0] * num_q_heads + rel_extent = rel_proj.shape[1] + assert d_rel == 16 and rel_proj.shape[0] == 16 + r_offset = num_q_heads * head_dim + 2 * num_kv_heads * head_dim + log_scaling_ptr = log_scaling if log_scaling is not None else qkvr + common = dict( + R_OFFSET=r_offset, + NUM_Q_HEADS=num_q_heads, + REL_EXTENT=rel_extent, + APPLY_LOG_SCALING=log_scaling is not None, + ) + + if use_rel_proj_throughput(rows, rel_extent): + grid = ( + triton.cdiv(rows, THROUGHPUT_BLOCK_M * THROUGHPUT_GROUP_M), + triton.cdiv(rel_extent, THROUGHPUT_BLOCK_N), + ) + _rel_proj_throughput_kernel[grid]( + qkvr, + rel_proj, + rel_out, + log_scaling_ptr, + rows, + qkvr.stride(0), + BLOCK_M=THROUGHPUT_BLOCK_M, + BLOCK_N=THROUGHPUT_BLOCK_N, + GROUP_M=THROUGHPUT_GROUP_M, + num_warps=THROUGHPUT_NUM_WARPS, + **common, + ) + return + + grid = ( + triton.cdiv(rows, LOW_BLOCK_M), + triton.cdiv(rel_extent, LOW_BLOCK_N), + ) + _rel_proj_low_latency_kernel[grid]( + qkvr, + rel_proj, + rel_out, + log_scaling_ptr, + rows, + qkvr.stride(0), + BLOCK_M=LOW_BLOCK_M, + BLOCK_N=LOW_BLOCK_N, + num_warps=LOW_NUM_WARPS, + **common, + ) + + +@triton.jit(do_not_specialize=["tokens", "stride_block_table_req", "max_blocks"]) +def _qkvr_qkv_kernel( + qkvr_ptr, + q_norm_weight_ptr, + q_out_ptr, + rel_proj_ptr, + rel_out_ptr, + k_weight_ptr, + v_weight_ptr, + k_norm_weight_ptr, + conv_cache_ptr, + key_cache_ptr, + value_cache_ptr, + positions_ptr, + seq_idx_ptr, + conv_slot_mapping_ptr, + conv_block_table_ptr, + query_start_ptr, + attention_slot_mapping_ptr, + log_scaling_ptr, + tokens, + eps, + stride_x_t, + stride_cc_block, + stride_cc_head, + stride_cc_token, + stride_cc_dim, + stride_kc_block, + stride_kc_token, + stride_kc_head, + stride_vc_block, + stride_vc_token, + stride_vc_head, + stride_block_table_req, + max_blocks, + conv_block_size, + attention_page_size, + Q_WIDTH: tl.constexpr, + KV_WIDTH: tl.constexpr, + NUM_Q_HEADS: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + WINDOW_SIZE: tl.constexpr, + OFF_K: tl.constexpr, + OFF_V: tl.constexpr, + APPLY_LOG_SCALING: tl.constexpr, + D_REL: tl.constexpr, + REL_EXTENT: tl.constexpr, + REL_EXTENT_PADDED: tl.constexpr, +): + block = tl.program_id(0) + num_q_rows = tokens * NUM_Q_HEADS + dims = tl.arange(0, HEAD_DIM) + + if block < num_q_rows: + row = block + token = row // NUM_Q_HEADS + head = row % NUM_Q_HEADS + values = tl.load( + qkvr_ptr + token * stride_x_t + head * HEAD_DIM + dims, + ).to(tl.float32) + weight = tl.load(q_norm_weight_ptr + dims).to(tl.float32) + rstd = tl.rsqrt(tl.sum(values * values, axis=0) / HEAD_DIM + eps) + normalized = values * rstd * weight + if APPLY_LOG_SCALING: + normalized = normalized.to(q_out_ptr.dtype.element_ty).to(tl.float32) + normalized *= tl.load(log_scaling_ptr + token) + tl.store( + q_out_ptr + row * HEAD_DIM + dims, + normalized.to(q_out_ptr.dtype.element_ty), + ) + rel_cols = tl.arange(0, REL_EXTENT_PADDED) + rel_mask = rel_cols < REL_EXTENT + projected = tl.zeros([REL_EXTENT_PADDED], dtype=tl.float32) + rel_offset = Q_WIDTH + 2 * KV_WIDTH + head * D_REL + for rel_dim in tl.static_range(D_REL): + rel_value = tl.load( + qkvr_ptr + token * stride_x_t + rel_offset + rel_dim + ).to(tl.float32) + proj = tl.load( + rel_proj_ptr + rel_dim * REL_EXTENT + rel_cols, + mask=rel_mask, + other=0.0, + ).to(tl.float32) + projected += rel_value * proj + projected = projected.to(rel_out_ptr.dtype.element_ty).to(tl.float32) + if APPLY_LOG_SCALING: + projected *= tl.load(log_scaling_ptr + token) + tl.store( + rel_out_ptr + row * REL_EXTENT + rel_cols, + projected.to(rel_out_ptr.dtype.element_ty), + mask=rel_mask, + ) + else: + row = block - num_q_rows + if row < tokens * NUM_KV_HEADS: + token = row // NUM_KV_HEADS + head = row % NUM_KV_HEADS + position = tl.load(positions_ptr + token) + request = tl.load(seq_idx_ptr + token) + conv_slot = tl.load(conv_slot_mapping_ptr + token) + query_start = tl.load(query_start_ptr + token) + attention_slot = tl.load(attention_slot_mapping_ptr + token) + valid = conv_slot >= 0 + + k_col = Q_WIDTH + head * HEAD_DIM + v_col = Q_WIDTH + KV_WIDTH + head * HEAD_DIM + k_value = tl.load(qkvr_ptr + token * stride_x_t + k_col + dims) + v_value = tl.load(qkvr_ptr + token * stride_x_t + v_col + dims) + + safe_slot = tl.maximum(conv_slot, 0) + cache_base = ( + conv_cache_ptr + + (safe_slot // conv_block_size) * stride_cc_block + + head * stride_cc_head + + (safe_slot % conv_block_size) * stride_cc_token + ) + tl.store( + cache_base + (OFF_K + dims) * stride_cc_dim, + k_value, + mask=valid, + ) + tl.store( + cache_base + (OFF_V + dims) * stride_cc_dim, + v_value, + mask=valid, + ) + + acc_k = tl.zeros([HEAD_DIM], dtype=tl.float32) + acc_v = tl.zeros([HEAD_DIM], dtype=tl.float32) + for tap in tl.static_range(WINDOW_SIZE): + source_position = position - (WINDOW_SIZE - 1) + tap + source_row = token - (WINDOW_SIZE - 1) + tap + in_window = valid & (source_position >= 0) + intra = in_window & (source_row >= query_start) + cached = in_window & (source_row < query_start) + safe_row = tl.maximum(source_row, 0) + source_k = tl.load( + qkvr_ptr + safe_row * stride_x_t + k_col + dims, + mask=intra, + other=0.0, + ).to(tl.float32) + source_v = tl.load( + qkvr_ptr + safe_row * stride_x_t + v_col + dims, + mask=intra, + other=0.0, + ).to(tl.float32) + safe_position = tl.maximum(source_position, 0) + logical_block = tl.minimum( + safe_position // conv_block_size, max_blocks - 1 + ) + physical_block = tl.load( + conv_block_table_ptr + + request * stride_block_table_req + + logical_block, + mask=cached, + other=0, + ).to(tl.int64) + tap_base = ( + conv_cache_ptr + + physical_block * stride_cc_block + + head * stride_cc_head + + (safe_position % conv_block_size) * stride_cc_token + ) + source_k += tl.load( + tap_base + (OFF_K + dims) * stride_cc_dim, + mask=cached, + other=0.0, + ).to(tl.float32) + source_v += tl.load( + tap_base + (OFF_V + dims) * stride_cc_dim, + mask=cached, + other=0.0, + ).to(tl.float32) + k_weight = tl.load( + k_weight_ptr + (head * HEAD_DIM + dims) * WINDOW_SIZE + tap + ).to(tl.float32) + v_weight = tl.load( + v_weight_ptr + (head * HEAD_DIM + dims) * WINDOW_SIZE + tap + ).to(tl.float32) + acc_k += source_k * k_weight + acc_v += source_v * v_weight + + k_rounded = (acc_k + k_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty) + v_rounded = (acc_v + v_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty) + k_float = k_rounded.to(tl.float32) + k_norm_weight = tl.load(k_norm_weight_ptr + dims).to(tl.float32) + rstd = tl.rsqrt(tl.sum(k_float * k_float, axis=0) / HEAD_DIM + eps) + k_normalized = (k_float * rstd * k_norm_weight).to( + qkvr_ptr.dtype.element_ty + ) + + safe_attention_slot = tl.maximum(attention_slot, 0) + attention_block = safe_attention_slot // attention_page_size + attention_offset = safe_attention_slot % attention_page_size + tl.store( + key_cache_ptr + + attention_block * stride_kc_block + + attention_offset * stride_kc_token + + head * stride_kc_head + + dims, + k_normalized, + mask=attention_slot >= 0, + ) + tl.store( + value_cache_ptr + + attention_block * stride_vc_block + + attention_offset * stride_vc_token + + head * stride_vc_head + + dims, + v_rounded, + mask=attention_slot >= 0, + ) + + +@triton.jit(do_not_specialize=["num_rows"]) +def _q_kernel( + qkvr_ptr, + q_norm_weight_ptr, + q_out_ptr, + log_scaling_ptr, + num_rows, + stride_x_t, + eps, + NUM_Q_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_ROWS: tl.constexpr, + APPLY_LOG_SCALING: tl.constexpr, +): + rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + dims = tl.arange(0, HEAD_DIM) + row_mask = rows < num_rows + tokens = rows // NUM_Q_HEADS + heads = rows % NUM_Q_HEADS + values = tl.load( + qkvr_ptr + + tokens[:, None] * stride_x_t + + heads[:, None] * HEAD_DIM + + dims[None, :], + mask=row_mask[:, None], + other=0.0, + ).to(tl.float32) + weight = tl.load(q_norm_weight_ptr + dims).to(tl.float32) + rstd = tl.rsqrt(tl.sum(values * values, axis=1) / HEAD_DIM + eps) + normalized = values * rstd[:, None] * weight[None, :] + if APPLY_LOG_SCALING: + normalized = normalized.to(q_out_ptr.dtype.element_ty).to(tl.float32) + tau = tl.load(log_scaling_ptr + tokens, mask=row_mask, other=1.0) + normalized *= tau[:, None] + tl.store( + q_out_ptr + rows[:, None] * HEAD_DIM + dims[None, :], + normalized.to(q_out_ptr.dtype.element_ty), + mask=row_mask[:, None], + ) + + +@triton.jit(do_not_specialize=["tokens", "stride_block_table_req", "max_blocks"]) +def _kv_kernel( + qkvr_ptr, + k_weight_ptr, + v_weight_ptr, + k_norm_weight_ptr, + conv_cache_ptr, + key_cache_ptr, + value_cache_ptr, + positions_ptr, + seq_idx_ptr, + conv_slot_mapping_ptr, + conv_block_table_ptr, + query_start_ptr, + attention_slot_mapping_ptr, + tokens, + eps, + stride_x_t, + stride_cc_block, + stride_cc_head, + stride_cc_token, + stride_cc_dim, + stride_kc_block, + stride_kc_token, + stride_kc_head, + stride_vc_block, + stride_vc_token, + stride_vc_head, + stride_block_table_req, + max_blocks, + conv_block_size, + attention_page_size, + Q_WIDTH: tl.constexpr, + KV_WIDTH: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + WINDOW_SIZE: tl.constexpr, + OFF_K: tl.constexpr, + OFF_V: tl.constexpr, + BLOCK_ROWS: tl.constexpr, +): + token = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + dims = tl.arange(0, HEAD_DIM) + row_mask = token < tokens + head_id = tl.program_id(1) + head = tl.full([BLOCK_ROWS], head_id, tl.int64) + position = tl.load(positions_ptr + token, mask=row_mask, other=0) + request = tl.load(seq_idx_ptr + token, mask=row_mask, other=0) + conv_slot = tl.load(conv_slot_mapping_ptr + token, mask=row_mask, other=-1) + query_start = tl.load(query_start_ptr + token, mask=row_mask, other=0) + attention_slot = tl.load( + attention_slot_mapping_ptr + token, mask=row_mask, other=-1 + ) + valid = row_mask & (conv_slot >= 0) + + k_col = Q_WIDTH + head * HEAD_DIM + v_col = Q_WIDTH + KV_WIDTH + head * HEAD_DIM + k_value = tl.load( + qkvr_ptr + token[:, None] * stride_x_t + k_col[:, None] + dims[None, :], + mask=row_mask[:, None], + other=0.0, + ) + v_value = tl.load( + qkvr_ptr + token[:, None] * stride_x_t + v_col[:, None] + dims[None, :], + mask=row_mask[:, None], + other=0.0, + ) + safe_slot = tl.maximum(conv_slot, 0) + cache_base = ( + conv_cache_ptr + + (safe_slot // conv_block_size) * stride_cc_block + + head * stride_cc_head + + (safe_slot % conv_block_size) * stride_cc_token + ) + tl.store( + cache_base[:, None] + (OFF_K + dims[None, :]) * stride_cc_dim, + k_value, + mask=valid[:, None], + ) + tl.store( + cache_base[:, None] + (OFF_V + dims[None, :]) * stride_cc_dim, + v_value, + mask=valid[:, None], + ) + + acc_k = tl.zeros([BLOCK_ROWS, HEAD_DIM], dtype=tl.float32) + acc_v = tl.zeros([BLOCK_ROWS, HEAD_DIM], dtype=tl.float32) + for tap in tl.static_range(WINDOW_SIZE): + source_position = position - (WINDOW_SIZE - 1) + tap + source_row = token - (WINDOW_SIZE - 1) + tap + in_window = valid & (source_position >= 0) + intra = in_window & (source_row >= query_start) + cached = in_window & (source_row < query_start) + safe_row = tl.maximum(source_row, 0) + source_k = tl.load( + qkvr_ptr + safe_row[:, None] * stride_x_t + k_col[:, None] + dims[None, :], + mask=intra[:, None], + other=0.0, + ).to(tl.float32) + source_v = tl.load( + qkvr_ptr + safe_row[:, None] * stride_x_t + v_col[:, None] + dims[None, :], + mask=intra[:, None], + other=0.0, + ).to(tl.float32) + safe_position = tl.maximum(source_position, 0) + logical_block = tl.minimum(safe_position // conv_block_size, max_blocks - 1) + physical_block = tl.load( + conv_block_table_ptr + request * stride_block_table_req + logical_block, + mask=cached, + other=0, + ).to(tl.int64) + tap_base = ( + conv_cache_ptr + + physical_block * stride_cc_block + + head * stride_cc_head + + (safe_position % conv_block_size) * stride_cc_token + ) + source_k += tl.load( + tap_base[:, None] + (OFF_K + dims[None, :]) * stride_cc_dim, + mask=cached[:, None], + other=0.0, + ).to(tl.float32) + source_v += tl.load( + tap_base[:, None] + (OFF_V + dims[None, :]) * stride_cc_dim, + mask=cached[:, None], + other=0.0, + ).to(tl.float32) + k_weight = tl.load( + k_weight_ptr + (head_id * HEAD_DIM + dims) * WINDOW_SIZE + tap + ).to(tl.float32) + v_weight = tl.load( + v_weight_ptr + (head_id * HEAD_DIM + dims) * WINDOW_SIZE + tap + ).to(tl.float32) + acc_k += source_k * k_weight[None, :] + acc_v += source_v * v_weight[None, :] + + k_rounded = (acc_k + k_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty) + v_rounded = (acc_v + v_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty) + k_float = k_rounded.to(tl.float32) + k_norm_weight = tl.load(k_norm_weight_ptr + dims).to(tl.float32) + rstd = tl.rsqrt(tl.sum(k_float * k_float, axis=1) / HEAD_DIM + eps) + k_normalized = (k_float * rstd[:, None] * k_norm_weight[None, :]).to( + qkvr_ptr.dtype.element_ty + ) + + safe_attention_slot = tl.maximum(attention_slot, 0) + attention_block = safe_attention_slot // attention_page_size + attention_offset = safe_attention_slot % attention_page_size + attention_mask = row_mask & (attention_slot >= 0) + tl.store( + key_cache_ptr + + attention_block[:, None] * stride_kc_block + + attention_offset[:, None] * stride_kc_token + + head[:, None] * stride_kc_head + + dims[None, :], + k_normalized, + mask=attention_mask[:, None], + ) + tl.store( + value_cache_ptr + + attention_block[:, None] * stride_vc_block + + attention_offset[:, None] * stride_vc_token + + head[:, None] * stride_vc_head + + dims[None, :], + v_rounded, + mask=attention_mask[:, None], + ) + + +def _run_tiled_q( + qkvr: torch.Tensor, + q_norm_weight: torch.Tensor, + q_out: torch.Tensor, + positions: torch.Tensor, + *, + eps: float, + num_q_heads: int, + head_dim: int, + log_scaling: torch.Tensor | None, +) -> None: + num_rows = qkvr.shape[0] * num_q_heads + _q_kernel[(triton.cdiv(num_rows, Q_BLOCK_ROWS),)]( + qkvr, + q_norm_weight, + q_out, + log_scaling if log_scaling is not None else positions, + num_rows, + qkvr.stride(0), + eps, + NUM_Q_HEADS=num_q_heads, + HEAD_DIM=head_dim, + BLOCK_ROWS=Q_BLOCK_ROWS, + APPLY_LOG_SCALING=log_scaling is not None, + num_warps=Q_NUM_WARPS, + ) + + +def _run_tiled_kv( + qkvr: torch.Tensor, + k_weight: torch.Tensor, + v_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + conv_cache: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + positions: torch.Tensor, + seq_idx: torch.Tensor, + conv_slot_mapping: torch.Tensor, + conv_block_table: torch.Tensor, + query_start: torch.Tensor, + attention_slot_mapping: torch.Tensor, + *, + eps: float, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + off_k: int, + off_v: int, + conv_block_size: int, +) -> None: + tokens = qkvr.shape[0] + _kv_kernel[(triton.cdiv(tokens, KV_BLOCK_ROWS), num_kv_heads)]( + qkvr, + k_weight, + v_weight, + k_norm_weight, + conv_cache, + key_cache, + value_cache, + positions, + seq_idx, + conv_slot_mapping, + conv_block_table, + query_start, + attention_slot_mapping, + tokens, + eps, + qkvr.stride(0), + conv_cache.stride(0), + conv_cache.stride(1), + conv_cache.stride(2), + conv_cache.stride(3), + key_cache.stride(0), + key_cache.stride(1), + key_cache.stride(2), + value_cache.stride(0), + value_cache.stride(1), + value_cache.stride(2), + conv_block_table.stride(0), + conv_block_table.shape[1], + conv_block_size, + key_cache.shape[1], + Q_WIDTH=num_q_heads * head_dim, + KV_WIDTH=num_kv_heads * head_dim, + NUM_KV_HEADS=num_kv_heads, + HEAD_DIM=head_dim, + WINDOW_SIZE=k_weight.shape[1], + OFF_K=off_k, + OFF_V=off_v, + BLOCK_ROWS=KV_BLOCK_ROWS, + num_warps=KV_NUM_WARPS, + ) + + +def _run_fused_small( + qkvr: torch.Tensor, + q_norm_weight: torch.Tensor, + q_out: torch.Tensor, + rel_proj: torch.Tensor, + rel_out: torch.Tensor, + k_weight: torch.Tensor, + v_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + conv_cache: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + positions: torch.Tensor, + seq_idx: torch.Tensor, + conv_slot_mapping: torch.Tensor, + conv_block_table: torch.Tensor, + query_start: torch.Tensor, + attention_slot_mapping: torch.Tensor, + *, + eps: float, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + off_k: int, + off_v: int, + conv_block_size: int, + log_scaling: torch.Tensor | None, +) -> None: + tokens = qkvr.shape[0] + + num_q_rows = tokens * num_q_heads + grid = (num_q_rows + tokens * num_kv_heads,) + _qkvr_qkv_kernel[grid]( + qkvr, + q_norm_weight, + q_out, + rel_proj, + rel_out, + k_weight, + v_weight, + k_norm_weight, + conv_cache, + key_cache, + value_cache, + positions, + seq_idx, + conv_slot_mapping, + conv_block_table, + query_start, + attention_slot_mapping, + log_scaling if log_scaling is not None else positions, + tokens, + eps, + qkvr.stride(0), + conv_cache.stride(0), + conv_cache.stride(1), + conv_cache.stride(2), + conv_cache.stride(3), + key_cache.stride(0), + key_cache.stride(1), + key_cache.stride(2), + value_cache.stride(0), + value_cache.stride(1), + value_cache.stride(2), + conv_block_table.stride(0), + conv_block_table.shape[1], + conv_block_size, + key_cache.shape[1], + Q_WIDTH=num_q_heads * head_dim, + KV_WIDTH=num_kv_heads * head_dim, + NUM_Q_HEADS=num_q_heads, + NUM_KV_HEADS=num_kv_heads, + HEAD_DIM=head_dim, + WINDOW_SIZE=k_weight.shape[1], + OFF_K=off_k, + OFF_V=off_v, + APPLY_LOG_SCALING=log_scaling is not None, + D_REL=16, + REL_EXTENT=rel_proj.shape[1], + REL_EXTENT_PADDED=triton.next_power_of_2(rel_proj.shape[1]), + num_warps=SMALL_NUM_WARPS, + ) + + +def fused_qkvr_prep( + qkvr: torch.Tensor, + k_weight: torch.Tensor, + v_weight: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + rel_proj: torch.Tensor, + eps: float, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + d_rel: int, + conv_cache: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + positions: torch.Tensor, + conv_block_table: torch.Tensor, + seq_idx: torch.Tensor, + conv_slot_mapping: torch.Tensor, + query_start: torch.Tensor, + attention_slot_mapping: torch.Tensor, + off_k: int, + off_v: int, + conv_block_size: int, + log_scaling: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + assert d_rel == 16 and rel_proj.shape[0] == 16 + assert head_dim == 128 + assert qkvr.is_contiguous() + assert k_weight.stride() == (k_weight.shape[1], 1) + assert v_weight.stride() == (v_weight.shape[1], 1) + assert rel_proj.stride() == (rel_proj.shape[1], 1) + assert conv_cache.stride(3) == 1 + assert key_cache.stride(3) == 1 and value_cache.stride(3) == 1 + tokens = qkvr.shape[0] + q_out = torch.empty( + (tokens, num_q_heads * head_dim), dtype=qkvr.dtype, device=qkvr.device + ) + rel_out = torch.empty( + (tokens, num_q_heads, rel_proj.shape[1]), + dtype=qkvr.dtype, + device=qkvr.device, + ) + if tokens == 0: + return q_out, rel_out + + if tokens < SMALL_TOKEN_THRESHOLD: + _run_fused_small( + qkvr, + q_norm_weight, + q_out, + rel_proj, + rel_out, + k_weight, + v_weight, + k_norm_weight, + conv_cache, + key_cache, + value_cache, + positions, + seq_idx, + conv_slot_mapping, + conv_block_table, + query_start, + attention_slot_mapping, + eps=eps, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + off_k=off_k, + off_v=off_v, + conv_block_size=conv_block_size, + log_scaling=log_scaling, + ) + return q_out, rel_out + + kv_stream = aux_stream() + assert kv_stream is not None + current_stream = torch.cuda.current_stream() + kv_stream.wait_stream(current_stream) + with torch.cuda.stream(kv_stream): + _run_tiled_kv( + qkvr, + k_weight, + v_weight, + k_norm_weight, + conv_cache, + key_cache, + value_cache, + positions, + seq_idx, + conv_slot_mapping, + conv_block_table, + query_start, + attention_slot_mapping, + eps=eps, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + off_k=off_k, + off_v=off_v, + conv_block_size=conv_block_size, + ) + _run_tiled_q( + qkvr, + q_norm_weight, + q_out, + positions, + eps=eps, + num_q_heads=num_q_heads, + head_dim=head_dim, + log_scaling=log_scaling, + ) + qkvr_rel_proj( + qkvr, + rel_proj, + rel_out, + log_scaling, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + d_rel=d_rel, + ) + current_stream.wait_stream(kv_stream) + return q_out, rel_out diff --git a/vllm/models/inkling/amd/ops/rel_attention_decode.py b/vllm/models/inkling/amd/ops/rel_attention_decode.py new file mode 100644 index 00000000000..810eb9cb5e6 --- /dev/null +++ b/vllm/models/inkling/amd/ops/rel_attention_decode.py @@ -0,0 +1,403 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +"""Split-KV decode for Inkling relative attention on ROCm. + +Adapted from LightSeek TokenSpeed's portable Triton relative-MHA decode. +""" + +from __future__ import annotations + +import os + +import torch + +from vllm.triton_utils import tl, triton + +_MIN_BLOCK_KV = tl.constexpr(32) + + +def decode_split_count(max_kv_len: int, window_left: int) -> int: + """Return the number of parallel KV partitions for decode.""" + if window_left >= 0: + effective_len = max(1, min(max_kv_len, window_left + 1)) + return min(4, max(1, triton.cdiv(effective_len, 128))) + return min(32, max(1, triton.cdiv(max(1, max_kv_len), 2048))) + + +def use_split_kv_decode( + *, + max_query_len: int, + max_kv_len: int, + page_size: int, + window_left: int, +) -> bool: + """Select split-KV only where it outperforms the single-pass kernel.""" + if os.getenv("INKLING_SPLIT_KV", "1") != "1": + return False + if max_query_len != 1: + return False + if page_size >= 64: + return True + if window_left >= 0: + return page_size >= 32 + return max_kv_len >= 8192 + + +@triton.jit +def _split_kv_stage1( + q_ptr, + rel_ptr, + k_ptr, + v_ptr, + block_table_ptr, + cache_seqlens, + mid_out_ptr, + mid_lse_ptr, + softmax_scale, + stride_q_t, + stride_q_h, + stride_k_b, + stride_k_p, + stride_k_h, + stride_k_d, + stride_v_b, + stride_v_p, + stride_v_h, + stride_v_d, + stride_r_t, + stride_r_h, + stride_r_e, + stride_mo_t, + stride_mo_h, + stride_mo_s, + stride_ml_t, + stride_ml_h, + stride_ml_s, + stride_bt_b: tl.constexpr, + page_size: tl.constexpr, + window_left: tl.constexpr, + gqa_group_size: tl.constexpr, + num_q_heads: tl.constexpr, + head_dim: tl.constexpr, + rel_extent: tl.constexpr, + max_kv_splits: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid_t = tl.program_id(0) + pid_h = tl.program_id(1) + pid_s = tl.program_id(2) + + head_blocks_per_kv = tl.cdiv(gqa_group_size, BLOCK_H) + kv_head = pid_h // head_blocks_per_kv + valid_block_h: tl.constexpr = min(BLOCK_H, gqa_group_size) + off_h = pid_h * valid_block_h + tl.arange(0, BLOCK_H) + h_valid = (off_h < (pid_h + 1) * valid_block_h) & (off_h < num_q_heads) + off_d = tl.arange(0, BLOCK_D) + d_valid = off_d < head_dim + + cache_len = tl.load(cache_seqlens + pid_t) + effective_len = ( + tl.minimum(cache_len, window_left + 1) if window_left >= 0 else cache_len + ) + kv_offset = cache_len - effective_len + split_len = ( + tl.cdiv(tl.cdiv(effective_len, max_kv_splits), _MIN_BLOCK_KV) * _MIN_BLOCK_KV + ) + split_start = split_len * pid_s + split_end = tl.minimum(split_start + split_len, effective_len) + + q_offsets = pid_t * stride_q_t + off_h[:, None] * stride_q_h + off_d[None, :] + row_max = tl.full((BLOCK_H,), float("-inf"), dtype=tl.float32) + row_sum = tl.zeros((BLOCK_H,), dtype=tl.float32) + acc = tl.zeros((BLOCK_H, BLOCK_D), dtype=tl.float32) + + if split_end > split_start: + q = tl.load( + q_ptr + q_offsets, + mask=h_valid[:, None] & d_valid[None, :], + other=0.0, + ) + for start_n in range(split_start, split_end, BLOCK_N): + off_n = start_n + tl.arange(0, BLOCK_N) + n_valid = off_n < split_end + token_idx = kv_offset + off_n + logical_page = token_idx // page_size + page_offset = token_idx % page_size + physical_page = tl.load( + block_table_ptr + pid_t * stride_bt_b + logical_page, + mask=n_valid, + other=0, + ) + k_offsets = ( + physical_page[None, :].to(tl.int64) * stride_k_b + + page_offset[None, :] * stride_k_p + + kv_head * stride_k_h + + off_d[:, None] * stride_k_d + ) + k = tl.load( + k_ptr + k_offsets, + mask=d_valid[:, None] & n_valid[None, :], + other=0.0, + ) + scores = tl.dot(q, k.to(q.dtype)) * softmax_scale + + rel_dist = cache_len - 1 - token_idx + rel_valid = (rel_dist >= 0) & (rel_dist < rel_extent) + rel_idx = tl.maximum(0, tl.minimum(rel_dist, rel_extent - 1)) + rel_offsets = ( + pid_t * stride_r_t + + off_h[:, None] * stride_r_h + + rel_idx[None, :] * stride_r_e + ) + rel_bias = tl.load( + rel_ptr + rel_offsets, + mask=h_valid[:, None] & rel_valid[None, :] & n_valid[None, :], + other=0.0, + ) + scores += rel_bias.to(tl.float32) + scores = tl.where( + h_valid[:, None] & n_valid[None, :], + scores, + float("-inf"), + ) + + v_offsets = ( + physical_page[:, None].to(tl.int64) * stride_v_b + + page_offset[:, None] * stride_v_p + + kv_head * stride_v_h + + off_d[None, :] * stride_v_d + ) + v = tl.load( + v_ptr + v_offsets, + mask=n_valid[:, None] & d_valid[None, :], + other=0.0, + ) + + next_max = tl.maximum(tl.max(scores, axis=1), row_max) + old_scale = tl.exp(row_max - next_max) + probs = tl.exp(scores - next_max[:, None]) + acc *= old_scale[:, None] + acc += tl.dot(probs.to(v.dtype), v) + row_sum = row_sum * old_scale + tl.sum(probs, axis=1) + row_max = next_max + + mid_out_offsets = ( + pid_t * stride_mo_t + + off_h[:, None] * stride_mo_h + + pid_s * stride_mo_s + + off_d[None, :] + ) + tl.store( + mid_out_ptr + mid_out_offsets, + acc / row_sum[:, None], + mask=h_valid[:, None] & d_valid[None, :], + ) + mid_lse_offsets = ( + pid_t * stride_ml_t + off_h * stride_ml_h + pid_s * stride_ml_s + ) + tl.store( + mid_lse_ptr + mid_lse_offsets, + row_max + tl.log(row_sum), + mask=h_valid, + ) + + +@triton.jit +def _split_kv_stage2( + mid_out_ptr, + mid_lse_ptr, + out_ptr, + cache_seqlens, + stride_mo_t, + stride_mo_h, + stride_mo_s, + stride_ml_t, + stride_ml_h, + stride_ml_s, + stride_o_t, + stride_o_h, + window_left: tl.constexpr, + head_dim: tl.constexpr, + max_kv_splits: tl.constexpr, + BLOCK_D: tl.constexpr, +): + pid_t = tl.program_id(0) + pid_h = tl.program_id(1) + + cache_len = tl.load(cache_seqlens + pid_t) + effective_len = ( + tl.minimum(cache_len, window_left + 1) if window_left >= 0 else cache_len + ) + split_len = ( + tl.cdiv(tl.cdiv(effective_len, max_kv_splits), _MIN_BLOCK_KV) * _MIN_BLOCK_KV + ) + + off_d = tl.arange(0, BLOCK_D) + d_valid = off_d < head_dim + row_max = -float("inf") + row_sum = 0.0 + acc = tl.zeros((BLOCK_D,), dtype=tl.float32) + + for split_id in range(max_kv_splits): + split_start = split_len * split_id + split_end = tl.minimum(split_start + split_len, effective_len) + if split_end > split_start: + value = tl.load( + mid_out_ptr + + pid_t * stride_mo_t + + pid_h * stride_mo_h + + split_id * stride_mo_s + + off_d, + mask=d_valid, + other=0.0, + ) + split_lse = tl.load( + mid_lse_ptr + + pid_t * stride_ml_t + + pid_h * stride_ml_h + + split_id * stride_ml_s + ) + next_max = tl.maximum(split_lse, row_max) + old_scale = tl.exp(row_max - next_max) + split_scale = tl.exp(split_lse - next_max) + acc = acc * old_scale + value * split_scale + row_sum = row_sum * old_scale + split_scale + row_max = next_max + + tl.store( + out_ptr + pid_t * stride_o_t + pid_h * stride_o_h + off_d, + acc / row_sum, + mask=d_valid, + ) + + +@torch.no_grad() +def inkling_rel_attention_split_kv_decode( + q: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + *, + block_table: torch.Tensor, + cache_seqlens: torch.Tensor, + softmax_scale: float, + window_left: int, + rel_extent: int, + rel_logits: torch.Tensor, + max_kv_len: int, + out: torch.Tensor, +) -> torch.Tensor: + """Run split-KV relative attention for single-token decode.""" + num_kv_heads = key_cache.shape[2] + gqa_group_size = q.shape[1] // num_kv_heads + max_kv_splits = decode_split_count(max_kv_len, window_left) + block_d = triton.next_power_of_2(q.shape[2]) + block_h = min(8, gqa_group_size) + + mid_out = torch.empty( + q.shape[0], + q.shape[1], + max_kv_splits, + q.shape[2], + dtype=torch.float32, + device=q.device, + ) + mid_lse = torch.empty( + q.shape[0], + q.shape[1], + max_kv_splits, + dtype=torch.float32, + device=q.device, + ) + stage1_grid = ( + q.shape[0], + triton.cdiv(q.shape[1], block_h), + max_kv_splits, + ) + _split_kv_stage1[stage1_grid]( + q, + rel_logits, + key_cache, + value_cache, + block_table, + cache_seqlens, + mid_out, + mid_lse, + softmax_scale, + q.stride(0), + q.stride(1), + key_cache.stride(0), + key_cache.stride(1), + key_cache.stride(2), + key_cache.stride(3), + value_cache.stride(0), + value_cache.stride(1), + value_cache.stride(2), + value_cache.stride(3), + rel_logits.stride(0), + rel_logits.stride(1), + rel_logits.stride(2), + mid_out.stride(0), + mid_out.stride(1), + mid_out.stride(2), + mid_lse.stride(0), + mid_lse.stride(1), + mid_lse.stride(2), + block_table.stride(0), + page_size=key_cache.shape[1], + window_left=window_left, + gqa_group_size=gqa_group_size, + num_q_heads=q.shape[1], + head_dim=q.shape[2], + rel_extent=rel_extent, + max_kv_splits=max_kv_splits, + BLOCK_D=block_d, + BLOCK_H=block_h, + BLOCK_N=key_cache.shape[1], + num_warps=4, + num_stages=1, + ) + stage2_grid = (q.shape[0], q.shape[1]) + _split_kv_stage2[stage2_grid]( + mid_out, + mid_lse, + out, + cache_seqlens, + mid_out.stride(0), + mid_out.stride(1), + mid_out.stride(2), + mid_lse.stride(0), + mid_lse.stride(1), + mid_lse.stride(2), + out.stride(0), + out.stride(1), + window_left=window_left, + head_dim=q.shape[2], + max_kv_splits=max_kv_splits, + BLOCK_D=block_d, + num_warps=4, + num_stages=2, + ) + return out diff --git a/vllm/models/inkling/amd/ops/sconv.py b/vllm/models/inkling/amd/ops/sconv.py new file mode 100644 index 00000000000..7ebb2a8a0a8 --- /dev/null +++ b/vllm/models/inkling/amd/ops/sconv.py @@ -0,0 +1,290 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling short-convolution kernels backed by a paged sliding-window state cache. + +Each layer's 4 conv streams (K, V, attn-output, mlp-output) share one paged KV +cache ``[num_blocks, H, N, D]`` (head-major; see ``sconv_swa_attn.py``). A stream +occupies the contiguous D-sub-range ``[off_s, off_s + ws)`` across all ``H`` +heads, so its flat per-token width is ``H * ws`` and that is the conv channel +dim. The cache stores the conv *input* at every absolute position. + +``fused_sconv`` is the single-launch path used by the model: per token it writes +the current input to its slot and convolves the ``W`` taps ending at its +absolute position. A tap landing inside the current forward is read from the +immutable input ``x`` (row ``src - pos + pid_t``); only pre-forward taps are read +from the paged cache (window position ``src`` -> physical block via +``block_table[req, src // N]``). The just-written slot is never read back this +step, so there is no write/read hazard within or across programs -- which is +why this needs no decode-vs-prefill split and is valid for prefill / decode / +spec alike. + +All kernels address the cache purely by ``(slot, absolute_position)`` and +allocate nothing inside the captured region; their grids depend only on the +token count (``fused_sconv`` on a fixed token/channel tiling), so the same +path replays correctly under eager, breakable PIECEWISE, and FULL cudagraphs +without any data-dependent shape or branch. +""" + +from __future__ import annotations + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _fused_sconv_kernel( + x_ptr, # [T, H*WS] head-major current-token inputs (also residual) + cache_ptr, # [num_blocks, H, N, D] paged (page-strided view) + weight_ptr, # [H*WS, W] + out_ptr, # [T, H*WS] + pos_ptr, # [T] int64 absolute position per token + seq_idx_ptr, # [T] int32 token -> batch request + slot_ptr, # [T] int64 flat slot (block*N + blk_off); < 0 => PAD + block_table_ptr, # [num_reqs, max_blocks] int32 block_table + qstart_ptr, # [T] int32 first x-row of the token's request + T, # num tokens + stride_x_t, + stride_c_blk, + stride_c_h, + stride_c_n, + stride_c_d, + stride_w_d, + stride_w_w, + stride_bt_r, + MAX_BLOCKS, + N, # block_size + W: tl.constexpr, + USE_SILU: tl.constexpr, + USE_RESIDUAL: tl.constexpr, + OFF_S: tl.constexpr, + WS: tl.constexpr, + H: tl.constexpr, + BT: tl.constexpr, + BLOCK_C: tl.constexpr, +): + # Each program owns a [BT tokens, BLOCK_C channels] tile. The flat channel + # index packs all H heads head-major (head = c // WS, in-stream offset = + # c % WS), so one program spans heads -- no per-head launch and no + # next_power_of_2(WS) lane waste. + pid_t = tl.program_id(0) + pid_c = tl.program_id(1) + toff = pid_t * BT + tl.arange(0, BT) # [BT] token rows + coff = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) # [BLOCK_C] flat channels + C = H * WS + t_mask = toff < T + c_mask = coff < C + head = tl.minimum(coff // WS, H - 1) # clamp keeps masked lanes in-buffer + cd = OFF_S + coff % WS # cache D-index of the channel's stream slot + + slot = tl.load(slot_ptr + toff, mask=t_mask, other=-1) # [BT] + valid = slot >= 0 + pos = tl.load(pos_ptr + toff, mask=t_mask, other=0) + req = tl.load(seq_idx_ptr + toff, mask=t_mask, other=0) + qstart = tl.load(qstart_ptr + toff, mask=t_mask, other=0) + + tc_mask = t_mask[:, None] & c_mask[None, :] + + # 1) Insert each token's input into its paged slot (skip PAD rows). + xv = tl.load(x_ptr + toff[:, None] * stride_x_t + coff[None, :], mask=tc_mask) + safe_slot = tl.maximum(slot, 0) + dst = ( + cache_ptr + + (safe_slot // N)[:, None] * stride_c_blk + + head[None, :] * stride_c_h + + (safe_slot % N)[:, None] * stride_c_n + + cd[None, :] * stride_c_d + ) + tl.store(dst, xv, mask=tc_mask & valid[:, None]) + + # 2) Convolve the W taps ending at each token's `pos`. Each tap is read from + # `x` if it falls inside this forward (row >= the request's first row), else + # from the paged cache. Exactly one source is unmasked per tap, so we sum both. + acc = tl.zeros([BT, BLOCK_C], dtype=tl.float32) + for iw in tl.static_range(W): + src = pos - (W - 1) + iw # [BT] absolute window position + row = toff - (W - 1) + iw # [BT] x-row of `src` (== src - pos + token) + in_win = valid & (src >= 0) + intra = in_win & (row >= qstart) + cached = in_win & (row < qstart) + # intra-forward tap: read the immutable input x (never the slot we just + # wrote), so there is no write/read hazard. + safe_row = tl.maximum(row, 0) + xt = tl.load( + x_ptr + safe_row[:, None] * stride_x_t + coff[None, :], + mask=c_mask[None, :] & intra[:, None], + other=0.0, + ).to(tl.float32) + # pre-forward tap: read from the paged cache via the block table. Clamp + # addressing terms; the load is masked off when out of window. + safe_src = tl.maximum(src, 0) + safe_lblk = tl.minimum(safe_src // N, MAX_BLOCKS - 1) + blk = tl.load( + block_table_ptr + req * stride_bt_r + safe_lblk, mask=cached, other=0 + ).to(tl.int64) + cbase = ( + cache_ptr + + blk[:, None] * stride_c_blk + + head[None, :] * stride_c_h + + (safe_src % N)[:, None] * stride_c_n + + cd[None, :] * stride_c_d + ) + cv = tl.load(cbase, mask=c_mask[None, :] & cached[:, None], other=0.0).to( + tl.float32 + ) + wv = tl.load( + weight_ptr + coff * stride_w_d + iw * stride_w_w, mask=c_mask, other=0.0 + ).to(tl.float32) + acc += (xt + cv) * wv[None, :] + + if USE_SILU: + acc = acc * tl.sigmoid(acc) + if USE_RESIDUAL: + acc += xv.to(tl.float32) + + tl.store( + out_ptr + toff[:, None] * stride_x_t + coff[None, :], + acc.to(out_ptr.dtype.element_ty), + mask=tc_mask, + ) + + +def fused_sconv( + x: torch.Tensor, # [T, H*ws] head-major current-token inputs + weight: torch.Tensor, # [H*ws, W] + cache: torch.Tensor, # [num_blocks, H, N, D] paged + positions: torch.Tensor, # [T] int64 absolute position per token + block_table: torch.Tensor, # [num_reqs, max_blocks] int32 + seq_idx: torch.Tensor, # [T] int32 token -> batch request + slot_mapping: torch.Tensor, # [T] int64 flat slot (PAD = -1 => skip) + query_start: torch.Tensor, # [T] int32 first x-row of the token's request + off_s: int, + ws: int, + block_size: int, + activation: str | None = None, + use_residual: bool = True, +) -> torch.Tensor: + """Single-launch insert + depthwise causal conv1d over the paged cache. + + Reads same-forward taps from ``x`` and pre-forward taps from the cache, so + it is race-free in one launch for prefill / decode / spec and cudagraph-safe + under eager / piecewise / full capture. + """ + T = x.shape[0] + out = torch.empty_like(x) + if T == 0: + return out + assert x.is_contiguous() + assert cache.stride(3) == 1, "cache D-dim must be contiguous" + H = cache.shape[1] + W = weight.shape[1] + C = H * ws # flat conv channel dim (all heads, head-major) + # Tile BT tokens x BLOCK_C channels per program: enough work per CTA to + # amortize the per-token addressing, while keeping the grid large for + # prefill. BLOCK_C spans heads so there is no per-head launch. A ~2K-element + # tile at 4 warps measured best on Blackwell; larger tiles spill registers. + BLOCK_C = min(triton.next_power_of_2(C), 256) + BT = 8 + grid = (triton.cdiv(T, BT), triton.cdiv(C, BLOCK_C)) + _fused_sconv_kernel[grid]( + x, + cache, + weight, + out, + positions, + seq_idx, + slot_mapping, + block_table, + query_start, + T, + x.stride(0), + cache.stride(0), + cache.stride(1), + cache.stride(2), + cache.stride(3), + weight.stride(0), + weight.stride(1), + block_table.stride(0), + block_table.shape[1], + block_size, + W=W, + USE_SILU=activation in ("silu", "swish"), + USE_RESIDUAL=use_residual, + OFF_S=off_s, + WS=ws, + H=H, + BT=BT, + BLOCK_C=BLOCK_C, + num_warps=4, + ) + return out + + +@triton.jit +def _seq_metadata_kernel( + qsl_ptr, # [num_reqs + 1] int32 cumulative query start rows + seq_idx_ptr, # [T] int32 out: token -> owning request + query_start_ptr, # [T] int32 out: first x-row of the token's request + num_reqs, + num_actual_tokens, + num_padded_tokens, + n_iters, # ceil(log2(num_reqs)): binary-search depth + BLOCK: tl.constexpr, +): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + tok = offs.to(tl.int32) + # Largest j in [0, num_reqs) with qsl[j] <= tok. + lo = tl.zeros([BLOCK], tl.int32) + hi = tl.full([BLOCK], num_reqs - 1, tl.int32) + for _ in range(n_iters): + mid = (lo + hi + 1) // 2 + below = tl.load(qsl_ptr + mid) <= tok + lo = tl.where(below, mid, lo) + hi = tl.where(below, hi, mid - 1) + actual = offs < num_actual_tokens + padded = offs < num_padded_tokens + query_start = tl.load(qsl_ptr + lo) + tl.store(seq_idx_ptr + offs, tl.where(actual, lo, 0), mask=padded) + tl.store( + query_start_ptr + offs, + tl.where(actual, query_start, 0), + mask=padded, + ) + + +def sconv_seq_metadata( + query_start_loc: torch.Tensor, + num_reqs: int, + num_actual_tokens: int, + seq_idx_out: torch.Tensor, + query_start_out: torch.Tensor, + num_padded_tokens: int | None = None, +) -> None: + """Fill static per-token seq_idx / query_start buffers in one launch. + + Replaces the arange + searchsorted + clamp + gather + 2x copy chain of the + sconv metadata build with a single kernel writing both persistent buffers. + Padded rows are filled with zero and must have ``slot_mapping == -1``. + """ + if num_padded_tokens is None: + num_padded_tokens = num_actual_tokens + if num_padded_tokens < num_actual_tokens: + raise ValueError("num_padded_tokens must cover all actual tokens") + if num_padded_tokens > seq_idx_out.shape[0]: + raise ValueError("seq_idx_out is too small for the padded token count") + if num_padded_tokens > query_start_out.shape[0]: + raise ValueError("query_start_out is too small for the padded token count") + + BLOCK = 256 + n_iters = (num_reqs - 1).bit_length() + grid = (triton.cdiv(num_padded_tokens, BLOCK),) + _seq_metadata_kernel[grid]( + query_start_loc, + seq_idx_out, + query_start_out, + num_reqs, + num_actual_tokens, + num_padded_tokens, + n_iters, + BLOCK=BLOCK, + ) diff --git a/vllm/models/inkling/amd/ops/silu_and_mul.py b/vllm/models/inkling/amd/ops/silu_and_mul.py new file mode 100644 index 00000000000..42ec25f6323 --- /dev/null +++ b/vllm/models/inkling/amd/ops/silu_and_mul.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""SwiGLU kernels for the Inkling MLP layers. + +``silu_and_mul_triton``: SiLU-and-mul over the checkpoint's interleaved +fused gate/up layout (dense MLP). ``sink_silu_mul_epilogue``: the sink-expert +variant with the per-expert dequant scale and per-token gamma fused in. +""" + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit(do_not_specialize=["M"]) +def _silu_and_mul_triton_kernel( + gateup_out_ptr, + down_inp_ptr, + M, + N: tl.constexpr, + GRID_SIZE: tl.constexpr, + NUM_STAGES: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + EVEN_N: tl.constexpr, + INT64_INDEX: tl.constexpr, +): + start_pid = tl.program_id(0) + if INT64_INDEX: + start_pid = start_pid.to(tl.int64) + M = M.to(tl.int64) + + NUM_BLOCKS_N: tl.constexpr = tl.cdiv(N, BLOCK_SIZE_N) + num_blocks_mn = tl.cdiv(M, BLOCK_SIZE_M) * NUM_BLOCKS_N + + for pid in tl.range(start_pid, num_blocks_mn, GRID_SIZE, num_stages=NUM_STAGES): + pid_m = pid // NUM_BLOCKS_N + pid_n = pid % NUM_BLOCKS_N + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + mask_m = offs_m < M + mask_n = offs_n < N + + # Interleaved fused gate/up: [g0, u0, g1, u1, ...]. + mask_offs_2n = pid_n * BLOCK_SIZE_N + tl.arange(0, 2 * BLOCK_SIZE_N) // 2 + tl.static_assert(BLOCK_SIZE_N % 8 == 0, f"{BLOCK_SIZE_N=}") + mask_2n = mask_offs_2n < N + mask_2n = tl.max_constancy(mask_2n, [16]) + + offs_2n = pid_n * 2 * BLOCK_SIZE_N + tl.arange(0, 2 * BLOCK_SIZE_N) + offs_m2n = offs_m[:, None] * N * 2 + offs_2n[None, :] + + if EVEN_N or pid_n * BLOCK_SIZE_N + BLOCK_SIZE_N <= N: + gateup_out = tl.load( + gateup_out_ptr + offs_m2n, mask=mask_m[:, None], other=0.0 + ) + else: + mask_m2n = mask_m[:, None] & mask_2n[None, :] + gateup_out = tl.load(gateup_out_ptr + offs_m2n, mask=mask_m2n, other=0.0) + + gate_out, up_out = tl.split( + tl.reshape(gateup_out, (BLOCK_SIZE_M, BLOCK_SIZE_N, 2)) + ) + gate_out = gate_out.to(tl.float32) + up_out = up_out.to(tl.float32) + + down_inp = gate_out * tl.sigmoid(gate_out) * up_out + + mask_mn = mask_m[:, None] if EVEN_N else mask_m[:, None] & mask_n[None, :] + offs_mn = offs_m[:, None] * N + offs_n[None, :] + tl.store(down_inp_ptr + offs_mn, down_inp, mask=mask_mn) + + +def silu_and_mul_triton(gateup_output: torch.Tensor) -> torch.Tensor: + """SiLU-and-mul for the interleaved fused gate/up layout. + + Adapted from ``inkling_kernels.activation.silu_and_mul_fwd`` (without MXFP). + """ + assert gateup_output.is_contiguous(), ( + f"{gateup_output.shape=} {gateup_output.stride()=}" + ) + assert gateup_output.ndim == 2, f"{gateup_output.shape=}" + + M = gateup_output.shape[0] + hidden_size = gateup_output.shape[1] + assert hidden_size % 2 == 0, f"{hidden_size=}" + N = hidden_size // 2 + + down_input = torch.empty( + (M, N), device=gateup_output.device, dtype=gateup_output.dtype + ) + if M == 0: + return down_input + + BLOCK_SIZE_N = max(8, min(256, triton.next_power_of_2(N))) + if M <= 1: + BLOCK_SIZE_M = 4 + elif M <= 256: + BLOCK_SIZE_M = 2 + elif M < 4096: + BLOCK_SIZE_M = 4 + else: + BLOCK_SIZE_M = 16 + BLOCK_SIZE_N = max(8, min(128, triton.next_power_of_2(N))) + max_grid_size = triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N) + num_sms = torch.cuda.get_device_properties( + gateup_output.device + ).multi_processor_count + grid_size = min(num_sms * 4, max_grid_size) + + _silu_and_mul_triton_kernel[(grid_size,)]( + gateup_out_ptr=gateup_output, + down_inp_ptr=down_input, + M=M, + N=N, + GRID_SIZE=grid_size, + NUM_STAGES=1, + BLOCK_SIZE_M=BLOCK_SIZE_M, + BLOCK_SIZE_N=BLOCK_SIZE_N, + EVEN_N=N % BLOCK_SIZE_N == 0, + INT64_INDEX=gateup_output.nbytes >= 2**31, + num_warps=8, + ) + + return down_input + + +@triton.jit(do_not_specialize=["T"]) +def _sink_epilogue_kernel( + raw_ptr, # [T, S * 2F] gemm1 output, interleaved g/u pairs per expert block + alpha_ptr, # [S] fp32 per-expert pre-SiLU dequant scale + gamma_ptr, # [T, S] fp32 per-token sink weights (may be strided) + ratio_ptr, # [S] fp32 per-expert post-SiLU scale (gemm2 alpha ratio) + out_ptr, # [T, S * F] output + T, + stride_raw_0, + stride_gamma_0, + F: tl.constexpr, + S: tl.constexpr, + BLOCK_F: tl.constexpr, +): + pid_t = tl.program_id(0).to(tl.int64) + pid_sf = tl.program_id(1) + if pid_t >= T: + return + s = pid_sf // (F // BLOCK_F) + offs_f = (pid_sf % (F // BLOCK_F)) * BLOCK_F + tl.arange(0, BLOCK_F) + + base = pid_t * stride_raw_0 + s * 2 * F + gate = tl.load(raw_ptr + base + 2 * offs_f).to(tl.float32) + up = tl.load(raw_ptr + base + 2 * offs_f + 1).to(tl.float32) + alpha = tl.load(alpha_ptr + s) + weight = tl.load(gamma_ptr + pid_t * stride_gamma_0 + s) * tl.load(ratio_ptr + s) + + gate *= alpha + up *= alpha + h = gate * tl.sigmoid(gate) * up * weight + tl.store(out_ptr + pid_t * (S * F) + s * F + offs_f, h) + + +def sink_silu_mul_epilogue( + raw: torch.Tensor, # [T, S * 2F] gemm1 output (interleaved gate/up rows) + alphas: torch.Tensor, # [S] fp32 + gammas: torch.Tensor, # [T, S] fp32 + ratios: torch.Tensor, # [S] fp32 + n_experts: int, + out_dtype: torch.dtype, +) -> torch.Tensor: + """Fused sink-expert epilogue: silu(g * a_e) * (u * a_e) * (gamma * r_e). + + One kernel replaces the per-expert dequant column scale, the SwiGLU, and + the per-token gamma multiply between the two sink GEMMs. + """ + tokens = raw.shape[0] + f = raw.shape[1] // (2 * n_experts) + out = torch.empty((tokens, n_experts * f), device=raw.device, dtype=out_dtype) + if tokens == 0: + return out + # raw may be a column-slice of a padded GEMM output (rows strided). + assert raw.stride(1) == 1 and gammas.stride(1) == 1 + # Largest power-of-two divisor of f (f = 768 -> 256), capped at 512. + block_f = min(512, f & (-f)) + _sink_epilogue_kernel[(tokens, n_experts * (f // block_f))]( + raw, + alphas, + gammas, + ratios, + out, + tokens, + raw.stride(0), + gammas.stride(0), + F=f, + S=n_experts, + BLOCK_F=block_f, + ) + return out diff --git a/vllm/models/inkling/amd/sconv_swa_attn.py b/vllm/models/inkling/amd/sconv_swa_attn.py new file mode 100644 index 00000000000..7b8958d556f --- /dev/null +++ b/vllm/models/inkling/amd/sconv_swa_attn.py @@ -0,0 +1,226 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling short-conv state managed as a sliding-window KV cache. + +Each decoder layer owns one ``InklingConvState`` (an ``AttentionLayerBase``) that +emits a single ``SlidingWindowSpec`` for the layer's 4 sconv streams (K, V, +attn-output, mlp-output), packed head-major into one block: + + H = num_kv_heads (per-rank), N = block_size = sconv_kernel_size, + D = head_dim(K) + head_dim(V) + hidden/H(attn) + hidden/H(mlp) + +``D`` is TP-invariant; per rank we store ``H/TP`` heads of width ``D``. The conv +reads/writes this cache out-of-band via a custom backend; the (smaller) conv page +is padded up to the uniform attention page by ``unify_kv_cache_spec_page_size``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import ClassVar + +import torch +from torch import nn + +from vllm.config import VllmConfig, get_current_vllm_config +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + AttentionMetadata, + AttentionMetadataBuilder, + CommonAttentionMetadata, +) +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec, SlidingWindowSpec + +from .ops.sconv import sconv_seq_metadata + +# Stream order within the per-head packed D (== contiguous sub-ranges). +_K, _V, _ATTN, _MLP = 0, 1, 2, 3 + + +@dataclass +class InklingSconvMetadata(AttentionMetadata): + block_table: torch.Tensor # [num_reqs, max_blocks] physical blocks per req + slot_mapping: torch.Tensor # [T] int64 flat slot of each token (-1 => skip) + seq_idx: torch.Tensor # [T] int32 token -> batch request + query_start: torch.Tensor # [T] int32 first x-row of each token's request + + +class InklingSconvMetadataBuilder(AttentionMetadataBuilder[InklingSconvMetadata]): + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + assert isinstance(kv_cache_spec, SlidingWindowSpec) + # Persistent per-token buffers for CUDA graph capture. + max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens + self.seq_idx_buffer = torch.empty( + max_num_tokens, dtype=torch.int32, device=device + ) + self.query_start_buffer = torch.empty( + max_num_tokens, dtype=torch.int32, device=device + ) + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> InklingSconvMetadata: + num_reqs = common_attn_metadata.num_reqs + num_actual_tokens = int(common_attn_metadata.query_start_loc_cpu[-1]) + num_padded_tokens = common_attn_metadata.slot_mapping.shape[0] + assert num_padded_tokens >= num_actual_tokens + + # Per-token seq_idx (owning request) and query_start (first x-row of + # that request; the fused kernel uses it to tell same-forward taps, + # read from x, from pre-forward taps, read from cache) in one launch. + sconv_seq_metadata( + common_attn_metadata.query_start_loc, + num_reqs, + num_actual_tokens, + self.seq_idx_buffer, + self.query_start_buffer, + num_padded_tokens, + ) + + return InklingSconvMetadata( + block_table=common_attn_metadata.block_table_tensor, + slot_mapping=common_attn_metadata.slot_mapping[:num_padded_tokens], + seq_idx=self.seq_idx_buffer[:num_padded_tokens], + query_start=self.query_start_buffer[:num_padded_tokens], + ) + + +class InklingSconvBackend(AttentionBackend): + """Custom dummy backend for the sconv sliding-window cache management.""" + + @staticmethod + def get_name() -> str: + return "INKLING_SCONV_SWA" + + @classmethod + def indexes_kv_by_block_stride(cls) -> bool: + # num_blocks is the outermost dim (HND, see get_kv_cache_shape), so the + # padded conv page is read through a strided view. + return True + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + # HND, num-blocks-first, head-major: [num_blocks, H, N, D]. + return (num_blocks, num_kv_heads, block_size, head_size) + + @staticmethod + def get_kv_cache_stride_order( + include_num_layers_dimension: bool = False, + ) -> tuple[int, ...]: + # Identity: physical layout == logical [num_blocks, H, N, D]. + if include_num_layers_dimension: + return (0, 1, 2, 3, 4) + return (0, 1, 2, 3) + + @staticmethod + def get_impl_cls(): + raise NotImplementedError( + "InklingSconvBackend has no attention impl; the conv runs out-of-band." + ) + + @staticmethod + def get_builder_cls() -> type[InklingSconvMetadataBuilder]: + return InklingSconvMetadataBuilder + + +class InklingConvState(nn.Module, AttentionLayerBase): + """Per-decoder-layer owner emitting one sliding-window conv-state spec.""" + + def __init__( + self, + *, + num_kv_heads: int, + head_dim: int, + hidden_size: int, + kernel_size: int, + prefix: str, + ) -> None: + super().__init__() + self.prefix = prefix + # Bound to the manager-allocated paged cache by bind_kv_cache; a + # placeholder until then. Read out-of-band by InklingShortConv. + self.kv_cache = torch.tensor([]) + tp_size = get_tensor_model_parallel_world_size() + # Guardrails for the conv-state layout below; only these are exercised. + # tp_size <= num_kv_heads keeps >=1 whole KV head per rank (no + # replication/clamping), so the per-head width stays TP-invariant. + assert tp_size <= num_kv_heads, ( + f"sconv SWA cache supports tp_size <= num_kv_heads ({num_kv_heads}), " + f"got {tp_size}" + ) + # Per-rank head count; D is TP-invariant (K/V heads and the hidden + # chunk both scale 1/TP together). The attn-/mlp-output sconv streams + # are hidden-sharded: each rank owns its H/tp chunk (the sublayer + # outputs are reduce-scattered / all-gathered around the conv). + self.num_kv_heads = num_kv_heads // tp_size + hidden_per_head = hidden_size // num_kv_heads + # Packed per-head width: K + V + attn-output chunk + mlp-output chunk, + # padded to a power of two so every layer's conv page is the same size + # and an exact multiple of the attention page (the page unifier then + # scales attention block sizes instead of padding). + raw_head_size = 2 * head_dim + 2 * hidden_per_head + self.head_size = 1 << (raw_head_size - 1).bit_length() + self.sliding_window = kernel_size + self.block_size = kernel_size + # Per-head D-sub-range (offset, width) for each stream. Streams share + # the cache; each writes/reads its own width across all H heads. + self.stream_ranges: tuple[tuple[int, int], ...] = ( + (0, head_dim), # _K + (head_dim, head_dim), # _V + (2 * head_dim, hidden_per_head), # _ATTN + (2 * head_dim + hidden_per_head, hidden_per_head), # _MLP + ) + vllm_config = get_current_vllm_config() + self._dtype = vllm_config.model_config.dtype + assert self._dtype == torch.bfloat16, ( + f"sconv SWA cache supports bfloat16 only, got {self._dtype}" + ) + # Register in the forward context so the runner enumerates this owner as + # an attention-like layer (get_kv_cache_spec / get_attn_backend). + compilation_config = vllm_config.compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + def forward(self): ... + + @property + def cache_block_size(self) -> int: + """Return the block size used by cache metadata and physical indexing.""" + if self.kv_cache.numel() > 0: + return self.kv_cache.shape[2] + return self.block_size + + def get_attn_backend(self) -> type[AttentionBackend]: + return InklingSconvBackend + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + return SlidingWindowSpec( + block_size=self.block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_size, + head_size_v=0, # all 4 streams packed into head_size + dtype=self._dtype, + sliding_window=self.sliding_window, + ) diff --git a/vllm/models/inkling/amd/short_conv.py b/vllm/models/inkling/amd/short_conv.py new file mode 100644 index 00000000000..3a4db8f8831 --- /dev/null +++ b/vllm/models/inkling/amd/short_conv.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling short convolution: depthwise causal conv1d (+ residual) over a paged +sliding-window conv-state cache. + +Each decoder layer owns one ``InklingConvState`` (``sconv_swa_attn.py``) holding the +manager-allocated paged cache for the layer's 4 sconv streams (K, V, attn-output, +mlp-output), packed head-major into one block. Each ``InklingShortConv`` is a +stateless weight + kernel launcher that, per forward (positions-addressed, the +same path for prefill / decode / mixed), inserts the current tokens' inputs +into their paged slot and convolves each token against the ``W`` taps ending +at its absolute position, reading pre-forward window positions out of the +paged cache via the block table. + +Per-forward metadata (``block_table`` / ``slot_mapping`` / ``seq_idx`` / +``query_start``) is built once by ``InklingSconvMetadataBuilder`` and published under +the owner's prefix in the forward context; the absolute ``positions`` are +threaded in from the model. The insert + conv run in a single ``fused_sconv`` +launch (same path for prefill / decode / mixed / spec). All inputs are +fixed-address persistent buffers and the grid is fixed, so the conv replays +correctly under eager, PIECEWISE, and FULL cudagraphs. +""" + +from __future__ import annotations + +import torch +from torch import nn +from torch.nn.parameter import Parameter + +from vllm.distributed import get_tensor_model_parallel_rank +from vllm.forward_context import get_forward_context +from vllm.model_executor.utils import set_weight_attrs + +from .ops import fused_sconv +from .sconv_swa_attn import InklingConvState, InklingSconvMetadata + + +class InklingShortConv(nn.Module): + def __init__( + self, dim: int, kernel_size: int, owner: InklingConvState, stream_idx: int + ) -> None: + super().__init__() + self.dim = dim + self.kernel_size = kernel_size + self.owner = owner + self.stream_idx = stream_idx + self.tp_rank = get_tensor_model_parallel_rank() + + # Depthwise conv weight; checkpoint stores (dim, 1, W). + self.weight = Parameter(torch.empty(dim, 1, kernel_size), requires_grad=False) + set_weight_attrs(self.weight, {"weight_loader": self.weight_loader}) + + def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor) -> None: + if loaded_weight.shape[0] != param.shape[0]: + shard = param.shape[0] + loaded_weight = loaded_weight.narrow(0, self.tp_rank * shard, shard) + param.data.copy_(loaded_weight) + + def forward(self, x: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + # x: (num_tokens, dim); positions: (num_tokens,) absolute positions. + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict): + # Memory-profiling / no metadata: identity (residual). + return x + m = attn_metadata.get(self.owner.prefix) + if m is None: + return x + assert isinstance(m, InklingSconvMetadata) + cache = self.owner.kv_cache + if cache.numel() == 0: + # Cache not yet bound (profiling before KV alloc): identity. + return x + + off_s, ws = self.owner.stream_ranges[self.stream_idx] + # The hybrid KV-cache planner can enlarge the logical conv block so + # that its physical page size matches the attention caches (for + # example, W=4 becomes 32 when --block-size=128). Metadata slot + # mappings are built with that enlarged size, so index the bound cache + # with its runtime token dimension rather than the kernel window size. + block_size = self.owner.cache_block_size + x = x.contiguous() + weight = self.weight.squeeze(1) # (dim, W) + + return fused_sconv( + x, + weight, + cache, + positions, + m.block_table, + m.seq_idx, + m.slot_mapping, + m.query_start, + off_s, + ws, + block_size, + activation=None, + use_residual=True, + ) diff --git a/vllm/triton_utils/__init__.py b/vllm/triton_utils/__init__.py index e20cb002cdc..5c9006c05c5 100644 --- a/vllm/triton_utils/__init__.py +++ b/vllm/triton_utils/__init__.py @@ -12,10 +12,16 @@ if TYPE_CHECKING or HAS_TRITON: import triton import triton.language as tl import triton.language.extra.libdevice as tldevice + from triton.experimental import gluon + from triton.experimental.gluon import language as gl + from triton.language.core import _aggregate as aggregate # noqa: E501 else: triton = TritonPlaceholder() tl = TritonLanguagePlaceholder() tldevice = TritonLanguagePlaceholder() + gluon = TritonLanguagePlaceholder() + gl = TritonLanguagePlaceholder() + aggregate = TritonLanguagePlaceholder() from vllm.triton_utils.tensor_descriptor import use_tensor_descriptor @@ -29,5 +35,8 @@ __all__ = [ "tldevice", "LOG2E", "LOGE2", + "gluon", + "gl", + "aggregate", "use_tensor_descriptor", ] From 30fbd055379ce4c5f26fcece6cfc90c5d8596f59 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas <akaratza@amd.com> Date: Mon, 27 Jul 2026 05:11:06 -0500 Subject: [PATCH 105/185] [ROCm] Use backend-default dot precision for ReplaySSM (#49909) Signed-off-by: Andreas Karatzas <akaratza@amd.com> --- .buildkite/test-amd.yaml | 29 ++++++++++--------- ...tive_state_update_replayssm_output_only.py | 25 +++++++++++----- 2 files changed, 33 insertions(+), 21 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 6e5d115f19b..7edc2931780 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -169,20 +169,6 @@ steps: - pip install helion==1.1.0 - pytest -v -s kernels/helion/ -- label: Kernels Mamba Test # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/mamba/ - - tests/kernels/mamba - - vllm/model_executor/layers/mamba/ops - - vllm/platforms/rocm.py - commands: - - pytest -v -s kernels/mamba - #------------------------------------------------------ mi250 · models / basic -------------------------------------------------------# - label: Basic Models Test (Other CPU) # TBD @@ -1595,6 +1581,21 @@ steps: commands: - pytest -v -s kernels/test_kda.py +- label: Kernels Mamba Test # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/mamba/ + - tests/kernels/mamba + - vllm/model_executor/layers/mamba/ops + - vllm/platforms/rocm.py + commands: + - pytest -v -s kernels/mamba + - label: Kernels MoE Test %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] diff --git a/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py b/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py index c6046b06a50..5b11bb7a08e 100644 --- a/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py +++ b/vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py @@ -6,14 +6,16 @@ import torch from vllm.model_executor.layers.mamba.ops.mamba_ssm import convert_rs_fp16x2, softplus from vllm.model_executor.layers.mamba.ops.replayssm_config import get_replayssm_config +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.utils import NULL_BLOCK_ID @triton.heuristics( { - "HAS_STATE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] - is not None + "HAS_STATE_BATCH_INDICES": lambda args: ( + args["state_batch_indices_ptr"] is not None + ) } ) @triton.heuristics( @@ -121,8 +123,9 @@ def _replayssm_output_only_precompute_kernel( @triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) @triton.heuristics( { - "HAS_STATE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] - is not None + "HAS_STATE_BATCH_INDICES": lambda args: ( + args["state_batch_indices_ptr"] is not None + ) } ) @triton.heuristics( @@ -210,6 +213,7 @@ def _replayssm_output_only_kernel( NF_NDS: tl.constexpr, FL_DSTATE_TILE: tl.constexpr, FL_NDS: tl.constexpr, + DOT_INPUT_PRECISION: tl.constexpr, USE_RS_ROUNDING: tl.constexpr, PHILOX_ROUNDS: tl.constexpr, # heuristic-computed @@ -409,12 +413,15 @@ def _replayssm_output_only_kernel( B_all_dot = tl.where( offs_k_dot[:, None] == write_pos, B_cur_tile[None, :], B_all_dot ) - # tf32x3 keeps fp32 parity with the elementwise baseline (plain tf32 on - # fp32 inputs drifts ~1e-2); bf16/fp16 inputs are unaffected by this flag. + # CUDA uses tf32x3 to keep fp32 parity with the elementwise + # baseline. Other platforms use their backend default; bf16/fp16 + # inputs are unaffected by this flag. B_scaled = (B_all_dot.to(tl.float32) * scale_dot[:, None]).to( x_ptr.dtype.element_ty ) - delta_state = tl.dot(x_all_ty, B_scaled, input_precision="tf32x3") + delta_state = tl.dot( + x_all_ty, B_scaled, input_precision=DOT_INPUT_PRECISION + ) state_ptrs = ( state_ptr + offs_m[:, None] * stride_state_dim @@ -574,6 +581,9 @@ def selective_state_update_replayssm_output_only( nf_nds = triton.cdiv(bs_dstate, nf_dstate_tile) fl_dstate_tile = max(16, min(fl_tile, bs_dstate)) fl_nds = triton.cdiv(bs_dstate, fl_dstate_tile) + # AMD Triton does not support tf32x3, so use its backend default. CUDA + # retains tf32x3 to preserve fp32 parity with the elementwise baseline. + dot_input_precision = None if current_platform.is_rocm() else "tf32x3" grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) z_strides = (z.stride(0), z.stride(1), z.stride(2)) if z is not None else (0, 0, 0) @@ -698,6 +708,7 @@ def selective_state_update_replayssm_output_only( nf_nds, fl_dstate_tile, fl_nds, + dot_input_precision, enable_stochastic_rounding, cache_philox_rounds, num_warps=num_warps, From 77cba0259f21dfc6f4f6298f54f3e43c9cc824c3 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer <ronen.schaffer@ibm.com> Date: Mon, 27 Jul 2026 13:29:57 +0300 Subject: [PATCH 106/185] [KV Offloading] Per-request tier filtering with TierFilter/TierMatcher (#48123) Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com> --- .../unit/offloading_connector/test_events.py | 24 ++- .../offloading_connector/test_scheduler.py | 7 +- tests/v1/kv_offload/cpu/test_manager.py | 4 +- tests/v1/kv_offload/tiering/test_fs_tier.py | 7 +- tests/v1/kv_offload/tiering/test_obj_tier.py | 4 +- .../tiering/test_tiering_offloading.py | 139 +++++++++++++++++- vllm/distributed/kv_events.py | 3 +- .../kv_connector/v1/offloading/events.py | 22 ++- .../kv_connector/v1/offloading/scheduler.py | 61 +++++++- vllm/v1/kv_offload/base.py | 55 +++++-- vllm/v1/kv_offload/cpu/manager.py | 4 +- vllm/v1/kv_offload/tiering/base.py | 8 +- vllm/v1/kv_offload/tiering/example/manager.py | 5 +- vllm/v1/kv_offload/tiering/fs/manager.py | 4 +- vllm/v1/kv_offload/tiering/manager.py | 2 + vllm/v1/kv_offload/tiering/obj/manager.py | 4 +- 16 files changed, 301 insertions(+), 52 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_events.py b/tests/v1/kv_connector/unit/offloading_connector/test_events.py index 5c172454b56..a6bab14e067 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_events.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_events.py @@ -9,8 +9,6 @@ from tests.v1.kv_connector.unit.utils import create_vllm_config from vllm.config import KVEventsConfig, KVTransferConfig from vllm.distributed.kv_events import ( MEDIUM_CPU, - MEDIUM_FS, - MEDIUM_OBJ, BlockRemoved, BlockStored, ) @@ -33,6 +31,7 @@ from vllm.v1.kv_cache_interface import ( ) from vllm.v1.kv_offload.base import ( Locality, + Medium, OffloadingEvent, OffloadingKVEventsConfig, OffloadKey, @@ -40,7 +39,7 @@ from vllm.v1.kv_offload.base import ( ) from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec -_CPU_MEDIUM = MEDIUM_CPU +_CPU_MEDIUM = Medium.CPU _FULL_ATTENTION_EVENT_SPEC = OffloadingEventGroupSpec( kv_cache_spec_kind=KVCacheSpecKind.FULL_ATTENTION.value, kv_cache_spec_sliding_window=None, @@ -135,7 +134,7 @@ def _record_lookup_chunks( def _stored_event( keys: list[OffloadKey], - medium: str = _CPU_MEDIUM, + medium: Medium = _CPU_MEDIUM, locality: Locality | None = None, ) -> OffloadingEvent: return OffloadingEvent( @@ -148,7 +147,7 @@ def _stored_event( def _removed_event( keys: list[OffloadKey], - medium: str = _CPU_MEDIUM, + medium: Medium = _CPU_MEDIUM, locality: Locality | None = None, ) -> OffloadingEvent: return OffloadingEvent( @@ -181,7 +180,7 @@ def test_take_events_forwards_locality_to_rich_store(): events = list( tracker.take_events( - [_stored_event([key], locality=Locality.LOCAL, medium=MEDIUM_FS)] + [_stored_event([key], locality=Locality.LOCAL, medium=Medium.STORAGE)] ) ) @@ -199,7 +198,7 @@ def test_take_events_forwards_locality_to_placeholder_store(): events = list( tracker.take_events( - [_stored_event([key], locality=Locality.REMOTE, medium=MEDIUM_FS)] + [_stored_event([key], locality=Locality.REMOTE, medium=Medium.STORAGE)] ) ) @@ -216,7 +215,7 @@ def test_take_events_forwards_locality_to_remove(): events = list( tracker.take_events( - [_removed_event([key], locality=Locality.LOCAL, medium=MEDIUM_FS)] + [_removed_event([key], locality=Locality.LOCAL, medium=Medium.STORAGE)] ) ) @@ -240,7 +239,7 @@ def test_take_events_publishes_routable_block_stored(): for i, event in enumerate(batch1): assert isinstance(event, BlockStored) - assert event.medium == _CPU_MEDIUM + assert event.medium == _CPU_MEDIUM.value assert event.block_hashes == [_wire_hash(_hash(i))] assert event.block_size == block_size assert event.token_ids == list( @@ -324,7 +323,7 @@ def test_lookup_promotion_factor_gt_1_store_and_remove(): assert len(removed) == 1 assert isinstance(removed[0], BlockRemoved) assert removed[0].block_hashes == expected_hashes - assert removed[0].medium == _CPU_MEDIUM + assert removed[0].medium == _CPU_MEDIUM.value assert removed[0].group_idx == 0 assert not tracker._pending_event_metadata @@ -444,12 +443,11 @@ def test_pending_cpu_removal_consumes_hit_backfill_until_next_hit(): ] -@pytest.mark.parametrize("medium", [MEDIUM_FS, MEDIUM_OBJ]) -def test_secondary_stored_event_does_not_mutate_cpu_metadata(medium: str): +def test_secondary_stored_event_does_not_mutate_cpu_metadata(): tracker, _, _, key = _lookup_chunk() expected_metadata = dict(tracker._pending_event_metadata) - stored = list(tracker.take_events([_stored_event([key], medium)])) + stored = list(tracker.take_events([_stored_event([key], Medium.STORAGE)])) assert stored[0].token_ids == [1, 2, 3, 4] assert tracker._pending_event_metadata == expected_metadata 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 d1463aebb4b..661078041f1 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -33,6 +33,7 @@ from vllm.v1.kv_cache_interface import ( ) from vllm.v1.kv_offload.base import ( LookupResult, + Medium, OffloadingEvent, OffloadingManager, OffloadPolicy, @@ -299,7 +300,7 @@ def test_abort_before_hit_uses_placeholder_then_later_hit_heals_removal( assert not tracker._pending_event_metadata - raw_events.append(OffloadingEvent(keys=[key], medium=MEDIUM_CPU, removed=False)) + raw_events.append(OffloadingEvent(keys=[key], medium=Medium.CPU, removed=False)) events = list(runner.connector_scheduler.take_events()) assert len(events) == 1 assert isinstance(events[0], BlockStored) @@ -320,7 +321,7 @@ def test_abort_before_hit_uses_placeholder_then_later_hit_heals_removal( ) assert key in tracker._pending_event_metadata - raw_events.append(OffloadingEvent(keys=[key], medium=MEDIUM_CPU, removed=True)) + raw_events.append(OffloadingEvent(keys=[key], medium=Medium.CPU, removed=True)) [event] = runner.connector_scheduler.take_events() assert isinstance(event, BlockRemoved) assert event.medium == MEDIUM_CPU @@ -355,7 +356,7 @@ def test_promotion_hit_precedes_stored_event_translation( raw_events: list[OffloadingEvent] = [] def lookup(key, req_context): - raw_events.append(OffloadingEvent(keys=[key], medium=MEDIUM_CPU, removed=False)) + raw_events.append(OffloadingEvent(keys=[key], medium=Medium.CPU, removed=False)) return LookupResult.HIT def take_raw_events(): diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index 97ed5168f74..6520a93fd1a 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -6,10 +6,10 @@ from dataclasses import dataclass import numpy as np import pytest -from vllm.distributed.kv_events import MEDIUM_CPU from vllm.v1.kv_offload.base import ( LoadStoreSpec, LookupResult, + Medium, OffloadingEvent, OffloadKey, PrepareStoreOutput, @@ -115,7 +115,7 @@ def verify_events( stores: list[set[OffloadKey]] = [] evictions: list[set[OffloadKey]] = [] for event in events: - assert event.medium == MEDIUM_CPU + assert event.medium == Medium.CPU if event.removed: evictions.append(set(event.keys)) else: diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index 589d7dda779..68b5512a3ec 100644 --- a/tests/v1/kv_offload/tiering/test_fs_tier.py +++ b/tests/v1/kv_offload/tiering/test_fs_tier.py @@ -18,10 +18,10 @@ import numpy as np import pytest import torch -from vllm.distributed.kv_events import MEDIUM_FS from vllm.v1.kv_offload.base import ( Locality, LookupResult, + Medium, OffloadingEvent, OffloadingKVEventsConfig, OffloadKey, @@ -514,8 +514,7 @@ def test_successful_store_emits_stored_event(fs_tier_with_events): events = list(tier.take_events()) assert len(events) == 1 assert events[0].keys == keys - # Literal medium pins the wire contract, not just the constant choice. - assert events[0].medium == "FS" + assert events[0].medium == Medium.STORAGE assert events[0].locality is Locality.LOCAL assert not events[0].removed # take_events drains the buffer. @@ -685,7 +684,7 @@ def test_cascade_store_emits_fs_event_through_tiering_manager(tmp_path): events.extend(manager.take_events()) time.sleep(0.01) - fs_events = [e for e in events if e.medium == MEDIUM_FS] + fs_events = [e for e in events if e.medium == Medium.STORAGE] assert len(fs_events) == 1 assert set(fs_events[0].keys) == set(keys) assert not fs_events[0].removed diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index 7500df8b4b9..fc30e1437a7 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -21,6 +21,7 @@ import torch from vllm.v1.kv_offload.base import ( Locality, LookupResult, + Medium, OffloadingKVEventsConfig, OffloadKey, ReqContext, @@ -473,8 +474,7 @@ class TestObjTierKVEvents: events = list(self.tier.take_events()) assert len(events) == 1 assert events[0].keys == keys - # Literal medium pins the wire contract, not just the constant choice. - assert events[0].medium == "OBJ" + assert events[0].medium == Medium.STORAGE assert events[0].locality is Locality.REMOTE assert not events[0].removed # take_events drains the buffer. diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index cf546b304cd..b19a270efb4 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -20,8 +20,13 @@ import torch from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( OffloadingConnectorStats, ) +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( + _parse_tier_filter, +) from vllm.v1.kv_offload.base import ( + Locality, LookupResult, + Medium, OffloadingCounterMetadata, OffloadingEvent, OffloadKey, @@ -29,6 +34,8 @@ from vllm.v1.kv_offload.base import ( ReqContext, RequestOffloadingContext, ScheduleEndContext, + TierFilter, + TierMatcher, make_offload_key, ) from vllm.v1.kv_offload.tiering.base import ( @@ -245,9 +252,9 @@ class TestTieringOffloadingManager: self.manager.on_new_request(req_context) def test_take_events_aggregates_tier_owned_events(self, manager_setup): - primary_event = OffloadingEvent(to_keys([1]), "CPU", removed=False) - secondary_event1 = OffloadingEvent(to_keys([2]), "tier-1", removed=False) - secondary_event2 = OffloadingEvent(to_keys([3]), "tier-2", removed=True) + primary_event = OffloadingEvent(to_keys([1]), Medium.CPU, removed=False) + secondary_event1 = OffloadingEvent(to_keys([2]), Medium.STORAGE, removed=False) + secondary_event2 = OffloadingEvent(to_keys([3]), Medium.STORAGE, removed=True) self.primary_tier.take_events = MagicMock(return_value=[primary_event]) self.secondary_tier1.take_events = MagicMock(return_value=[secondary_event1]) @@ -973,6 +980,56 @@ class TestTieringOffloadingManager: self.secondary_tier2.drain_jobs.assert_called_once() assert self.manager._transfer_jobs == {} + @pytest.mark.parametrize( + "load_tier_filter", + [ + TierFilter(matchers=(TierMatcher(medium=Medium.STORAGE),)), + TierFilter(matchers=()), + ], + ids=["non_matching_medium", "empty_no_load"], + ) + def test_tier_filter_skips_filtered_secondary( + self, manager_setup, load_tier_filter + ): + """Filter excluding secondary medium returns MISS from secondaries + even when they hold the block; primary is unaffected.""" + blocks = to_keys(range(2)) + # Put one block in primary, one only in secondary + self._start_request() + self.manager.prepare_store(blocks[:1], _CTX) + self.manager.complete_store(blocks[:1], _CTX, success=True) + self.secondary_tier1.blocks[blocks[1]] = True + + # Secondaries have medium=CPU, so load_tier_filter skips them. + self.secondary_tier1.lookup = MagicMock(wraps=self.secondary_tier1.lookup) + + ctx = ReqContext(req_id="r1", load_tier_filter=load_tier_filter) + assert self.manager.lookup(blocks[0], ctx) is LookupResult.HIT + assert self.manager.lookup(blocks[1], ctx) is LookupResult.MISS + self.secondary_tier1.lookup.assert_not_called() + + @pytest.mark.parametrize( + "load_tier_filter", + [ + TierFilter.ALL, + TierFilter(matchers=(TierMatcher(medium=Medium.CPU),)), + TierFilter(matchers=(TierMatcher(),)), + ], + ids=["all", "explicit_cpu", "unconstrained_matcher"], + ) + def test_tier_filter_allows_matching_secondary( + self, manager_setup, load_tier_filter + ): + """Filter that matches the secondary's medium allows lookup.""" + blocks = to_keys(range(1)) + self.secondary_tier1.blocks[blocks[0]] = True + + self.secondary_tier1.lookup = MagicMock(wraps=self.secondary_tier1.lookup) + + ctx = ReqContext(req_id="r2", load_tier_filter=load_tier_filter) + assert self.manager.lookup(blocks[0], ctx) is LookupResult.RETRY + self.secondary_tier1.lookup.assert_called() + class TestTieringOffloadingWithoutSecondaryTiers: """Test TieringOffloadingManager with no secondary tiers (backward compat).""" @@ -999,5 +1056,81 @@ class TestTieringOffloadingWithoutSecondaryTiers: assert count_hits(manager, blocks) == 3 +@pytest.mark.parametrize( + "raw,expected", + [ + ( + [{"medium": "storage"}], + TierFilter(matchers=(TierMatcher(medium=Medium.STORAGE),)), + ), + ( + [{"medium": "CPU"}], + TierFilter(matchers=(TierMatcher(medium=Medium.CPU),)), + ), + ( + [{}], + TierFilter(matchers=(TierMatcher(),)), + ), + ( + [{"medium": "storage", "locality": "local"}], + TierFilter( + matchers=(TierMatcher(medium=Medium.STORAGE, locality=Locality.LOCAL),) + ), + ), + ( + [{"medium": "cpu"}, {"medium": "storage"}], + TierFilter( + matchers=( + TierMatcher(medium=Medium.CPU), + TierMatcher(medium=Medium.STORAGE), + ) + ), + ), + ( + [], + TierFilter(matchers=()), + ), + ], + ids=[ + "medium_storage", + "medium_cpu_uppercase", + "unconstrained", + "with_locality", + "multiple_matchers", + "empty_list_deny_all", + ], +) +def test_parse_tier_filter_valid(raw, expected): + assert _parse_tier_filter(raw) == expected + + +@pytest.mark.parametrize( + "raw", + [ + "not a list", + [{"medium": "unknown"}], + [{"locality": "nowhere"}], + ], + ids=["non_list", "invalid_medium", "invalid_locality"], +) +def test_parse_tier_filter_invalid_returns_all(raw): + assert _parse_tier_filter(raw) is TierFilter.ALL + + +def test_parse_tier_filter_skips_bad_entries(): + result = _parse_tier_filter( + [ + {"medium": "storage"}, + "not a dict", + {"medium": "bogus"}, + {"medium": "cpu"}, + ] + ) + assert result.matchers == ( + TierMatcher(medium=Medium.STORAGE), + TierMatcher(medium=Medium.CPU), + ) + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/vllm/distributed/kv_events.py b/vllm/distributed/kv_events.py index be7e7363fb9..88cb3541f2b 100644 --- a/vllm/distributed/kv_events.py +++ b/vllm/distributed/kv_events.py @@ -44,8 +44,7 @@ class KVCacheEvent( MEDIUM_GPU = "GPU" MEDIUM_CPU = "CPU" -MEDIUM_FS = "FS" -MEDIUM_OBJ = "OBJ" +MEDIUM_STORAGE = "STORAGE" class BlockStored(KVCacheEvent): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py index 3c5a9956cac..67ada00e03f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py @@ -19,7 +19,13 @@ from collections.abc import Iterable from dataclasses import dataclass from typing import TYPE_CHECKING, Any, NamedTuple -from vllm.distributed.kv_events import BlockRemoved, BlockStored, KVCacheEvent +from vllm.distributed.kv_events import ( + MEDIUM_CPU, + MEDIUM_STORAGE, + BlockRemoved, + BlockStored, + KVCacheEvent, +) from vllm.logger import init_logger from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash from vllm.v1.kv_cache_interface import ( @@ -28,6 +34,7 @@ from vllm.v1.kv_cache_interface import ( get_kv_cache_spec_sliding_window, ) from vllm.v1.kv_offload.base import ( + Medium, OffloadingEvent, OffloadingKVEventsConfig, OffloadKey, @@ -43,6 +50,11 @@ if TYPE_CHECKING: logger = init_logger(__name__) +_MEDIUM_TO_EVENT_STR: dict[Medium, str] = { + Medium.CPU: MEDIUM_CPU, + Medium.STORAGE: MEDIUM_STORAGE, +} + class OffloadingEventGroupSpec(NamedTuple): kv_cache_spec_kind: str | None @@ -218,7 +230,7 @@ class OffloadingEventsTracker: def _placeholder_stored( self, key: OffloadKey, - medium: str, + medium: Medium, locality: str | None, ) -> BlockStored: return BlockStored( @@ -229,7 +241,7 @@ class OffloadingEventsTracker: token_ids=[], lora_id=None, block_size=0, - medium=medium, + medium=_MEDIUM_TO_EVENT_STR[medium], lora_name=None, group_idx=get_offload_group_idx(key), locality=locality, @@ -267,7 +279,7 @@ class OffloadingEventsTracker: token_ids=list(meta.token_ids), block_size=meta.block_size, lora_id=meta.lora_id, - medium=event.medium, + medium=_MEDIUM_TO_EVENT_STR[event.medium], lora_name=meta.lora_name, extra_keys=( list(meta.extra_keys) if meta.extra_keys is not None else None @@ -308,7 +320,7 @@ class OffloadingEventsTracker: for group_idx, hashes in by_group.items(): yield BlockRemoved( block_hashes=hashes, - medium=event.medium, + medium=_MEDIUM_TO_EVENT_STR[event.medium], group_idx=group_idx, locality=locality, ) 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 f809dca5ba4..fecce3d28fe 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -39,7 +39,9 @@ from vllm.v1.kv_cache_interface import ( ) from vllm.v1.kv_offload.base import ( GPULoadStoreSpec, + Locality, LookupResult, + Medium, OffloadingManager, OffloadingSpec, OffloadKey, @@ -47,6 +49,8 @@ from vllm.v1.kv_offload.base import ( ReqContext, RequestOffloadingContext, ScheduleEndContext, + TierFilter, + TierMatcher, make_offload_key, ) from vllm.v1.outputs import KVConnectorOutput @@ -54,6 +58,10 @@ from vllm.v1.request import Request, RequestStatus logger = init_logger(__name__) +KV_LOAD_TIERS_KEY = "kv_load_tiers" +MATCHER_MEDIUM_KEY = "medium" +MATCHER_LOCALITY_KEY = "locality" + @dataclass(slots=True) class TransferJobStatus: @@ -376,10 +384,61 @@ class RequestOffloadState: ) +def _parse_tier_filter(raw: Any) -> TierFilter: + """Parse raw kv_transfer_params tier matchers into a TierFilter.""" + if not isinstance(raw, list): + logger.warning( + "_parse_tier_filter: expected list, got %s; ignoring", + type(raw).__name__, + ) + return TierFilter.ALL + matchers: list[TierMatcher] = [] + for entry in raw: + if not isinstance(entry, dict): + logger.warning("_parse_tier_filter: entry is not a dict; skipping") + continue + medium: Medium | None = None + locality: Locality | None = None + raw_medium = entry.get(MATCHER_MEDIUM_KEY) + if raw_medium is not None: + try: + medium = Medium(raw_medium.upper()) + except (ValueError, AttributeError): + logger.warning( + "_parse_tier_filter: unknown medium %r; skipping entry", + raw_medium, + ) + continue + raw_locality = entry.get(MATCHER_LOCALITY_KEY) + if raw_locality is not None: + try: + locality = Locality(raw_locality.upper()) + except (ValueError, AttributeError): + logger.warning( + "_parse_tier_filter: unknown locality %r; skipping entry", + raw_locality, + ) + continue + matchers.append(TierMatcher(medium=medium, locality=locality)) + if not matchers: + if not raw: # input was [] — user explicitly wants nothing + return TierFilter(matchers=()) + # all entries were invalid — fall back to ALL + return TierFilter.ALL + return TierFilter(matchers=tuple(matchers)) + + def _create_req_context(req: Request) -> ReqContext: + params = req.kv_transfer_params + load_filter = TierFilter.ALL + if params: + raw = params.get(KV_LOAD_TIERS_KEY) + if raw is not None: + load_filter = _parse_tier_filter(raw) return ReqContext( req_id=req.request_id, - kv_transfer_params=req.kv_transfer_params, + kv_transfer_params=params, + load_tier_filter=load_filter, ) diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index cb90599cf81..6ff09f23d37 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -8,7 +8,7 @@ from abc import ABC, abstractmethod from collections.abc import Collection, Iterable, Sequence from dataclasses import dataclass, field from enum import Enum, auto -from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeVar +from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple, NewType, TypeVar import numpy as np import torch @@ -48,10 +48,54 @@ def get_offload_group_idx(key: OffloadKey) -> int: _T = TypeVar("_T") +class Medium(Enum): + """Storage medium of an offloading tier.""" + + CPU = "CPU" + STORAGE = "STORAGE" + + +class Locality(Enum): + """Locality of a tier's storage relative to the publishing instance.""" + + LOCAL = "LOCAL" + REMOTE = "REMOTE" + + +class TierMatcher(NamedTuple): + medium: Medium | None = None + locality: Locality | None = None + + def matches(self, medium: Medium | None, locality: Locality | None) -> bool: + medium_matches = self.medium is None or medium is None or self.medium == medium + locality_matches = ( + self.locality is None or locality is None or self.locality == locality + ) + return medium_matches and locality_matches + + +@dataclass(frozen=True) +class TierFilter: + """Per-request filter controlling which tiers participate.""" + + matchers: tuple[TierMatcher, ...] = () + + ALL: ClassVar["TierFilter"] + + def allows(self, medium: Medium | None, locality: Locality | None) -> bool: + if self is TierFilter.ALL: + return True + return any(m.matches(medium, locality) for m in self.matchers) + + +TierFilter.ALL = TierFilter(matchers=(TierMatcher(),)) + + @dataclass class ReqContext: req_id: str kv_transfer_params: dict[str, Any] | None = None + load_tier_filter: TierFilter = TierFilter.ALL # Per-request scratch space keyed by value type, so a tier can parse # kv_transfer_params once (in on_new_request) and read the result back # on later calls for the same request. @@ -110,17 +154,10 @@ class PrepareStoreOutput: evicted_keys: list[OffloadKey] -class Locality(Enum): - """Locality of a tier's storage relative to the publishing instance.""" - - LOCAL = "LOCAL" - REMOTE = "REMOTE" - - @dataclass class OffloadingEvent: keys: list[OffloadKey] - medium: str + medium: Medium # True if blocks are removed, False if stored removed: bool locality: Locality | None = None diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index 297c94aabb0..c2ec4170b8e 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -6,13 +6,13 @@ from typing import Literal from typing_extensions import override -from vllm.distributed.kv_events import MEDIUM_CPU from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( OffloadingConnectorStats, ) from vllm.v1.kv_offload.base import ( LoadStoreSpec, LookupResult, + Medium, OffloadingEvent, OffloadingManager, OffloadKey, @@ -52,7 +52,7 @@ class CPUOffloadingManager(OffloadingManager): store_threshold: int = 1, max_tracker_size: int = 64_000, ): - self.medium: str = MEDIUM_CPU + self.medium: Medium = Medium.CPU self._num_blocks: int = num_blocks self._num_allocated_blocks: int = 0 self._free_list: list[int] = [] diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index f5614171963..3577b4fa0ed 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -7,12 +7,14 @@ Abstract interfaces and data types for the secondary tiering layer. from abc import ABC, abstractmethod from collections.abc import Collection, Iterable from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, ClassVar import numpy as np from vllm.v1.kv_offload.base import ( + Locality, LookupResult, + Medium, OffloadingEvent, OffloadingMetricMetadata, OffloadKey, @@ -27,6 +29,7 @@ if TYPE_CHECKING: ) from vllm.v1.kv_offload.base import OffloadingSpec + # Type alias for job IDs used in async transfer tracking JobId = int @@ -108,6 +111,8 @@ class SecondaryTierManager(ABC): async jobs; get_finished_jobs() polls for completion. """ + medium: ClassVar[Medium | None] = None + def __init__( self, offloading_spec: "OffloadingSpec", @@ -124,6 +129,7 @@ class SecondaryTierManager(ABC): self._offloading_spec = offloading_spec self._primary_kv_view: memoryview = primary_kv_view self.tier_type = tier_type + self.locality: Locality | None = None @abstractmethod def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: diff --git a/vllm/v1/kv_offload/tiering/example/manager.py b/vllm/v1/kv_offload/tiering/example/manager.py index a9e4e4f689c..1a78a3feb71 100644 --- a/vllm/v1/kv_offload/tiering/example/manager.py +++ b/vllm/v1/kv_offload/tiering/example/manager.py @@ -11,12 +11,13 @@ TieringOffloadingManager without requiring actual storage or network backends. import logging from collections.abc import Iterable -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar from typing_extensions import override from vllm.v1.kv_offload.base import ( LookupResult, + Medium, OffloadKey, ReqContext, RequestOffloadingContext, @@ -42,6 +43,8 @@ class ExampleSecondaryTierManager(SecondaryTierManager): - Completes transfers immediately (synchronous) """ + medium: ClassVar[Medium] = Medium.CPU + def __init__( self, offloading_spec: "OffloadingSpec", diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index 4fde24d760e..4fe200629f3 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -30,11 +30,11 @@ except ImportError: from typing_extensions import override -from vllm.distributed.kv_events import MEDIUM_FS from vllm.logger import init_logger from vllm.v1.kv_offload.base import ( Locality, LookupResult, + Medium, OffloadingEvent, OffloadKey, ReqContext, @@ -100,7 +100,7 @@ class FileSystemTierManager(SecondaryTierManager): content. """ - medium: ClassVar[str] = MEDIUM_FS + medium: ClassVar[Medium] = Medium.STORAGE def __init__( self, diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 2f5b52e2220..c6738096135 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -324,6 +324,8 @@ class TieringOffloadingManager(OffloadingManager): for tier in self.secondary_tiers: if tier is exclude_tier: continue + if not req_context.load_tier_filter.allows(tier.medium, tier.locality): + continue result = tier.lookup(key, req_context) if result is LookupResult.HIT: promoted = self._initiate_promotion(tier, key, req_context) diff --git a/vllm/v1/kv_offload/tiering/obj/manager.py b/vllm/v1/kv_offload/tiering/obj/manager.py index f4af8c44c53..c7e0d4c4beb 100644 --- a/vllm/v1/kv_offload/tiering/obj/manager.py +++ b/vllm/v1/kv_offload/tiering/obj/manager.py @@ -7,13 +7,13 @@ import time from collections.abc import Iterable from typing import TYPE_CHECKING, ClassVar, NamedTuple -from vllm.distributed.kv_events import MEDIUM_OBJ from vllm.distributed.nixl_utils import NixlWrapper as nixl_agent from vllm.distributed.nixl_utils import nixl_agent_config from vllm.logger import init_logger from vllm.v1.kv_offload.base import ( Locality, LookupResult, + Medium, OffloadingEvent, OffloadKey, ReqContext, @@ -98,7 +98,7 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager): primary tier. Object keys are formed as ``{prefix}/{hash_shard}/{hash}.bin``. """ - medium: ClassVar[str] = MEDIUM_OBJ + medium: ClassVar[Medium] = Medium.STORAGE def __init__( self, From 92e8518d376c3c9a6979da4d56ee6263ef6af630 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Mon, 27 Jul 2026 12:51:16 +0100 Subject: [PATCH 107/185] Improve Transformers modelling backend `fx` tracer (#49957) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- .../models/transformers/fusers/test_linear.py | 7 +- .../transformers/fusers/test_rms_norm.py | 17 +- .../models/transformers/base.py | 9 +- .../models/transformers/fuser.py | 26 +- .../models/transformers/fusers/base.py | 25 +- .../models/transformers/fusers/glu.py | 12 +- .../models/transformers/fusers/qkv.py | 36 ++- .../models/transformers/fusers/rms_norm.py | 19 +- .../models/transformers/fx_utils.py | 284 +++++++++++++++--- 9 files changed, 307 insertions(+), 128 deletions(-) diff --git a/tests/models/transformers/fusers/test_linear.py b/tests/models/transformers/fusers/test_linear.py index 546dafb7921..eff78fa290d 100644 --- a/tests/models/transformers/fusers/test_linear.py +++ b/tests/models/transformers/fusers/test_linear.py @@ -368,12 +368,12 @@ def test_qkv_identifies_output_projection(): assert get_fuser(PerHeadQKNormAttention()).o_name == "o_proj" -def test_fuser_is_cached_per_class(): +def test_fuser_is_cached_per_class_and_structure(): with torch.device("meta"): fuser_a = get_fuser(GLUMLP()) fuser_b = get_fuser(GLUMLP()) assert fuser_a is fuser_b - assert GLUMLP in get_fuser.cache + assert any(key[0] is GLUMLP for key in get_fuser.cache) @pytest.mark.parametrize("cls", [NotAnMLP, UntraceableMLP]) @@ -460,8 +460,7 @@ def test_unfusable_modules_are_not_fused(cls, default_vllm_config): fuser = get_fuser(module) # Either no pattern matches the class, or this instance fails validation # (`recursive_replace` gates fusion and its weight mappings on `validate`) - model_config = default_vllm_config.model_config - assert fuser is None or not fuser.validate(module, model_config) + assert fuser is None or not fuser.validate(module, default_vllm_config) def test_act_and_mul_derived_from_module(default_vllm_config): diff --git a/tests/models/transformers/fusers/test_rms_norm.py b/tests/models/transformers/fusers/test_rms_norm.py index 6497b98c4ba..131897b1924 100644 --- a/tests/models/transformers/fusers/test_rms_norm.py +++ b/tests/models/transformers/fusers/test_rms_norm.py @@ -148,11 +148,13 @@ def test_rms_norm_builds_vllm_class(cls, expected, zero_centered, default_vllm_c # `default_vllm_config` supplies the config context the CustomOp needs; the # weightless path reads hidden size from the model config, so stub it. - model_config = SimpleNamespace(get_hidden_size=lambda: 16) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(get_hidden_size=lambda: 16) + ) with torch.device("meta"): module = cls() fuser = get_fuser(module) - built = fuser.fuse(module, "norm", model_config, None) + built = fuser.fuse(module, "norm", vllm_config) from vllm.model_executor.models.transformers.fusers.rms_norm import ( TPAwareNormMixin, ) @@ -176,8 +178,9 @@ def test_fused_rms_norm_op_default_eps(default_vllm_config): fuser = get_fuser(module) assert isinstance(fuser, RMSNormFuser) assert not fuser.zero_centered - model_config = SimpleNamespace(get_hidden_size=lambda: 16, dtype=torch.float32) - built = fuser.fuse(module, "norm", model_config, None) + mc = SimpleNamespace(get_hidden_size=lambda: 16, dtype=torch.float32) + vllm_config = SimpleNamespace(model_config=mc) + built = fuser.fuse(module, "norm", vllm_config) assert isinstance(built, VLLMRMSNorm) assert built.variance_epsilon == torch.finfo(torch.float32).eps @@ -185,11 +188,13 @@ def test_fused_rms_norm_op_default_eps(default_vllm_config): def test_eps_is_derived_per_instance(default_vllm_config): """Two instances of the same norm class with different eps must fuse to their own eps: the type-cached fuser holds only structure, not this value.""" - model_config = SimpleNamespace(get_hidden_size=lambda: 16) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(get_hidden_size=lambda: 16) + ) with torch.device("meta"): for eps in (1e-5, 1e-6): module = RMSNorm(16, eps=eps) - built = get_fuser(module).fuse(module, "norm", model_config, None) + built = get_fuser(module).fuse(module, "norm", vllm_config) assert built.variance_epsilon == eps diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py index 24cad1da20d..35e5c9d1a0f 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/vllm/model_executor/models/transformers/base.py @@ -104,6 +104,7 @@ class Base( super().__init__() logger.info("Using Transformers modeling backend.") + self.vllm_config = vllm_config self.config = vllm_config.model_config.hf_config self.text_config = self.config.get_text_config() self.cache_config = vllm_config.cache_config @@ -357,7 +358,7 @@ class Base( tip = get_feature_request_tip( self.model_config.model, self.model_config.trust_remote_code ) - logger.warning( + logger.warning_once( "%s does not define a pipeline parallel plan. The Transformers " "modeling backend will infer the split from the layers of %s in order " "of declaration and keep parameter-free modules on every rank. This " @@ -452,7 +453,7 @@ class Base( # Prefix the patterns because we always start from `self.model` tp_plan = {maybe_prefix("model", k): v for k, v in tp_plan.items()} # Detect fusable patterns once per module class (cached, so this is cheap) - fusers = Fusers(self.model, self.model_config) + fusers = Fusers(self.model, self.vllm_config) def register_fusion(fuser: BaseFuser, prefix: str): """Register a fused layer's mappings just before it is built.""" @@ -503,9 +504,7 @@ class Base( new_module = replace_conv_class(child_module) elif (fuser := fusers[child_module]) is not None: register_fusion(fuser, qual_name) - new_module = fuser.fuse( - child_module, qual_name, self.model_config, self.quant_config - ) + new_module = fuser.fuse(child_module, qual_name, self.vllm_config) logger.info_once(fuser.info(child_name)) _recursive_replace(new_module, prefix=qual_name) elif not isinstance(child_module, MoERunner): diff --git a/vllm/model_executor/models/transformers/fuser.py b/vllm/model_executor/models/transformers/fuser.py index 0aa7c419ceb..13d84a245d9 100644 --- a/vllm/model_executor/models/transformers/fuser.py +++ b/vllm/model_executor/models/transformers/fuser.py @@ -25,14 +25,19 @@ from vllm.model_executor.models.transformers.fusers import ( from vllm.model_executor.models.transformers.fx_utils import trace if TYPE_CHECKING: - from vllm.config.model import ModelConfig + from vllm.config import VllmConfig logger = init_logger(__name__) -@cached(cache={}, key=type) +def key(module: nn.Module) -> tuple: + """Cache key for `get_fuser`. Considers module type and its immediate children.""" + return (type(module), tuple(name for name, _ in module.named_children())) + + +@cached(cache={}, key=key) def get_fuser(module: nn.Module) -> BaseFuser | None: - """The fuser for `type(module)` (cached per class), or `None` if no match.""" + """The fuser for `module`'s class and shape (cached), or `None` if no match.""" # Projection fusions need >=2 sibling linears; the RMSNorm fusion needs a # leaf module (raw tensor math, no submodules). Nothing else can match, and # tracing is skipped for it. @@ -48,11 +53,14 @@ def get_fuser(module: nn.Module) -> BaseFuser | None: try: fuser.update_forward(module) except Exception as exc: - # An unrecognised source just means we cannot fuse here. logger.debug( - "Could not rewrite %s for fusion: %s", type(module), exc + "Attempted to fuse %s using %s but failed " + "to update its forward method: %s", + type(module), + fuser_cls.__name__, + exc, ) - return None + continue return fuser # A norm we could not match structurally is left unfused; flag likely misses. if module.__class__.__name__.endswith("RMSNorm"): @@ -67,12 +75,12 @@ def get_fuser(module: nn.Module) -> BaseFuser | None: class Fusers(UserDict): """Mapping from module class to fuser, for all fusable classes in a model.""" - def __init__(self, model: nn.Module, model_config: "ModelConfig"): - self.model_config = model_config + def __init__(self, model: nn.Module, vllm_config: "VllmConfig"): + self.vllm_config = vllm_config super().__init__({type(m): get_fuser(m) for m in model.modules()}) def __getitem__(self, m: nn.Module) -> BaseFuser | None: fuser = self.data.get(type(m)) - if fuser is not None and fuser.validate(m, self.model_config): + if fuser is not None and fuser.validate(m, self.vllm_config): return fuser return None diff --git a/vllm/model_executor/models/transformers/fusers/base.py b/vllm/model_executor/models/transformers/fusers/base.py index 54fb2d09165..abeff73e3ed 100644 --- a/vllm/model_executor/models/transformers/fusers/base.py +++ b/vllm/model_executor/models/transformers/fusers/base.py @@ -13,8 +13,7 @@ from torch import fx, nn from vllm.model_executor.models.utils import ShardId, maybe_prefix if TYPE_CHECKING: - from vllm.config.model import ModelConfig - from vllm.model_executor.layers.quantization import QuantizationConfig + from vllm.config import VllmConfig @dataclass @@ -36,16 +35,12 @@ class BaseFuser(ABC): """Match the pattern in `graph`, returning a fuser if found.""" @abstractmethod - def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool: + def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool: """Whether this fuser can be applied to this `module` instance.""" @abstractmethod def fuse( - self, - module: nn.Module, - prefix: str, - model_config: "ModelConfig", - quant_config: "QuantizationConfig", + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" ) -> nn.Module: """Apply the fusion to an already-validated `module`, returning the module to install in its place (mutated in place, or freshly built).""" @@ -123,24 +118,16 @@ class StackedFuser(BaseFuser): @abstractmethod def update_attrs( - self, - module: nn.Module, - prefix: str, - model_config: "ModelConfig", - quant_config: "QuantizationConfig", + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" ) -> None: """Replace `module`'s submodules with the merged module.""" def fuse( - self, - module: nn.Module, - prefix: str, - model_config: "ModelConfig", - quant_config: "QuantizationConfig", + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" ) -> nn.Module: """Fuse an already-validated `module` in place (see `Fusers.__getitem__`). Builds the merged submodule and binds the compiled forward.""" - self.update_attrs(module, prefix, model_config, quant_config) + self.update_attrs(module, prefix, vllm_config) module.forward = types.MethodType(self.fused_forward, module) return module diff --git a/vllm/model_executor/models/transformers/fusers/glu.py b/vllm/model_executor/models/transformers/fusers/glu.py index 951eb35777a..2da84addea2 100644 --- a/vllm/model_executor/models/transformers/fusers/glu.py +++ b/vllm/model_executor/models/transformers/fusers/glu.py @@ -33,8 +33,7 @@ from vllm.model_executor.models.transformers.utils import ( from vllm.model_executor.models.utils import ShardId, maybe_prefix if TYPE_CHECKING: - from vllm.config.model import ModelConfig - from vllm.model_executor.layers.quantization import QuantizationConfig + from vllm.config import VllmConfig logger = init_logger(__name__) @@ -168,7 +167,7 @@ class GLUFuser(StackedFuser): replace_expr(funcdef, muls[0], act_call) self.fused_forward = compile_forward(funcdef, fn) - def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool: + def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool: act = module.get_submodule(self.act_name) if self._get_act_and_mul_name(act) is None: logger.debug("No AndMul equivalent for %s; skipping fusion", type(act)) @@ -176,12 +175,9 @@ class GLUFuser(StackedFuser): return True def update_attrs( - self, - module: nn.Module, - prefix: str, - model_config: "ModelConfig", - quant_config: "QuantizationConfig", + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" ) -> None: + quant_config = vllm_config.quant_config act_fn = self._get_act_and_mul(module.get_submodule(self.act_name)) gate = module.get_submodule(self.gate_name) up = module.get_submodule(self.up_name) diff --git a/vllm/model_executor/models/transformers/fusers/qkv.py b/vllm/model_executor/models/transformers/fusers/qkv.py index 010a0018acc..b9fe9771fe1 100644 --- a/vllm/model_executor/models/transformers/fusers/qkv.py +++ b/vllm/model_executor/models/transformers/fusers/qkv.py @@ -17,6 +17,7 @@ from vllm.model_executor.models.transformers.fx_utils import ( is_linear, recover_forward, replace_expr, + returned_linear, single_self_call, ) from vllm.model_executor.models.transformers.utils import ( @@ -26,8 +27,7 @@ from vllm.model_executor.models.transformers.utils import ( from vllm.model_executor.models.utils import ShardId, maybe_prefix if TYPE_CHECKING: - from vllm.config.model import ModelConfig - from vllm.model_executor.layers.quantization import QuantizationConfig + from vllm.config import VllmConfig logger = init_logger(__name__) @@ -84,15 +84,16 @@ class QKVFuser(StackedFuser): return None q, k, v = qkv_nodes names = dict(q_name=q.target, k_name=k.target, v_name=v.target) - attn_width = module.get_submodule(q.target).out_features - candidates = [ - name - for name, child in module.named_children() - if isinstance(child, nn.Linear) - and name not in names.values() - and child.in_features == attn_width - ] - names["o_name"] = candidates[0] if len(candidates) == 1 else None + # o_proj produces the module's output. + o_name = returned_linear(graph, module) + # o_proj must be compatible with the q/k/v projections. + if o_name in names.values() or ( + o_name is not None + and module.get_submodule(o_name).in_features + != module.get_submodule(q.target).out_features + ): + o_name = None + names["o_name"] = o_name return cls(source_cls=type(module).__name__, **names) def update_forward(self, module: nn.Module) -> None: @@ -149,12 +150,12 @@ class QKVFuser(StackedFuser): replace_expr(funcdef, call, ast.Name(id=temp, ctx=ast.Load())) self.fused_forward = compile_forward(funcdef, fn) - def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool: + def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool: """Shapes must be compatible for a single merged, head-sharded GEMM.""" q = module.get_submodule(self.q_name) k = module.get_submodule(self.k_name) v = module.get_submodule(self.v_name) - head_size = model_config.get_head_size() + head_size = vllm_config.model_config.get_head_size() compatible = ( q.in_features == k.in_features == v.in_features and len({proj.bias is None for proj in (q, k, v)}) == 1 @@ -167,13 +168,10 @@ class QKVFuser(StackedFuser): return compatible def update_attrs( - self, - module: nn.Module, - prefix: str, - model_config: "ModelConfig", - quant_config: "QuantizationConfig", + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" ) -> None: - head_size = model_config.get_head_size() + quant_config = vllm_config.quant_config + head_size = vllm_config.model_config.get_head_size() q = module.get_submodule(self.q_name) k = module.get_submodule(self.k_name) merged = QKVParallelLinear( diff --git a/vllm/model_executor/models/transformers/fusers/rms_norm.py b/vllm/model_executor/models/transformers/fusers/rms_norm.py index 829fd9a541c..781b417177a 100644 --- a/vllm/model_executor/models/transformers/fusers/rms_norm.py +++ b/vllm/model_executor/models/transformers/fusers/rms_norm.py @@ -21,13 +21,13 @@ from vllm.model_executor.models.transformers.fx_utils import ( find_node, forward_input_count, is_op, + output_value, peel, trace, ) if TYPE_CHECKING: - from vllm.config.model import ModelConfig - from vllm.model_executor.layers.quantization import QuantizationConfig + from vllm.config import VllmConfig def _is_squared(node: object, x: fx.Node) -> bool: @@ -71,10 +71,8 @@ def _is_one_plus(node: object) -> bool: def _has_trailing_compute(graph: fx.Graph, node: fx.Node) -> bool: """Does the forward compute anything after `node` before returning?""" - output = find_node(graph, lambda n: n.op == "output") - if output is None or not output.args: - return False - return peel(output.args[0]) is not node + value = output_value(graph) + return value is not None and peel(value) is not node class TPAwareNormMixin(nn.Module): @@ -185,17 +183,14 @@ class RMSNormFuser(BaseFuser): return eps return None - def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool: + def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool: return True def fuse( - self, - module: nn.Module, - prefix: str, - model_config: "ModelConfig", - quant_config: "QuantizationConfig", + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" ) -> nn.Module: """Fuse the matched RMSNorm pattern into a vLLM fused RMSNorm CustomOp.""" + model_config = vllm_config.model_config weight = getattr(module, "weight", None) hidden_size = ( weight.size(0) if weight is not None else model_config.get_hidden_size() diff --git a/vllm/model_executor/models/transformers/fx_utils.py b/vllm/model_executor/models/transformers/fx_utils.py index 0e043941d8a..2dbc3a499de 100644 --- a/vllm/model_executor/models/transformers/fx_utils.py +++ b/vllm/model_executor/models/transformers/fx_utils.py @@ -9,10 +9,13 @@ stays live Python. `fusion.py` builds the concrete fusion patterns on top. """ import ast +import contextlib import inspect import operator import textwrap from collections.abc import Callable +from itertools import chain +from unittest import mock import torch from torch import fx, nn @@ -22,46 +25,64 @@ from vllm.logger import init_logger logger = init_logger(__name__) +_UNKNOWN = object() +"""Sentinel meta value for proxies whose concrete value could not be inferred. +Distinct from `None`, which is a valid concrete value (e.g. `attn_weights`).""" -def _infer_len(node: fx.Node) -> int | None: - """Concrete length of a proxy's value, inferred from its node chain. +_MODULE_CALL = nn.Module.__call__ +"""The unpatched `nn.Module.__call__`. During tracing fx patches it to record +`call_module` nodes; meta execution must call modules for real.""" - Lets tracing pass through the shape unpacks and `*`-splats (e.g. - `(*input_shape, -1, head_dim)`) that precede the patterns in HF attention. - """ - # `x.shape` has the rank of `x`, when known - if ( - node.op == "call_function" - and node.target is getattr - and node.args[1] == "shape" - and (rank := _rank(node.args[0])) is not None - ): - return rank - # Slices of known-length values - if node.op == "call_function" and node.target is operator.getitem: - src_len = _infer_len(node.args[0]) - index = node.args[1] - if src_len is not None and isinstance(index, slice): - return len(range(*index.indices(src_len))) + +def is_leaf_call(node: object) -> bool: + """Is node a call recorded by `_as_leaf_call` (e.g. an attention interface).""" + return isinstance(node, fx.Node) and node.meta.get("leaf_call", False) + + +def _reference_weight(module: nn.Module) -> torch.Tensor | None: + """A weight whose trailing dim is the module's hidden size. + + Linears and 2-D gate weights are `[out, hidden]`; norm weights are + `[hidden]`. Used to fabricate a placeholder input of matching size/dtype.""" + for child in module.modules(): + if isinstance(child, nn.Linear): + return child.weight + for param in module.parameters(): + if param.ndim in (1, 2): + return param return None -def _rank(node: fx.Node) -> int | None: - """The tensor rank of `node`'s value, if known.""" - # vLLM always feeds the model [1, seq_len, hidden_size] hidden states - if node.op == "placeholder" and node.target == "hidden_states": - return 3 - return None +class _MetaProxy(fx.Proxy): + """Proxy carrying the meta-tensor value of the traced expression. + Shape questions (`len`, iteration, `.shape` unpacks) are answered by + executing each op on the meta values, so PyTorch's meta kernels are the + single source of shape inference — no per-op rules.""" -class _SizedProxy(fx.Proxy): - """Proxy whose `len` is inferred from the graph (see `_infer_len`).""" + meta: object = _UNKNOWN def __len__(self) -> int: - length = _infer_len(self.node) - if length is None: - return super().__len__() - return length + if self.meta is not _UNKNOWN: + return len(self.meta) + return super().__len__() # type: ignore[misc] + + def __getattr__(self, k: str) -> "_MetaAttribute": + return _MetaAttribute(self, k) + + +class _MetaAttribute(_MetaProxy, fx.proxy.Attribute): + """Attribute proxy (e.g. `x.shape`) carrying its meta value. + + `Proxy.__getattr__` constructs `Attribute` directly, bypassing + `Tracer.proxy`, so the meta value must be grafted on here too.""" + + def __init__(self, root: fx.Proxy, attr: str): + super().__init__(root, attr) + root_meta = getattr(root, "meta", _UNKNOWN) + if root_meta is not _UNKNOWN: + with contextlib.suppress(Exception): + self.meta = getattr(root_meta, attr) class _AllLeafTracer(fx.Tracer): @@ -69,31 +90,159 @@ class _AllLeafTracer(fx.Tracer): Each child stays one `call_module` node, so matching sees the module's own forward structure (activations aren't decomposed into e.g. `sigmoid * x`). - `iter` traces through the leading shape unpacks (see `_infer_len`); anything - else untraceable ends the trace early and the partial graph is matched. + Every traced op is also executed on meta tensors (see `_MetaProxy`) so + shape unpacks and `*`-splats trace through; anything else untraceable ends + the trace early and the partial graph is matched. """ + varkw: str | None = None + """Name of the traced forward's `**kwargs` parameter, if any.""" + def is_leaf_module(self, m: nn.Module, module_qualified_name: str) -> bool: return True def proxy(self, node: fx.Node) -> fx.Proxy: - return _SizedProxy(node, self) + return _MetaProxy(node, self) + + def create_proxy(self, kind, target, args, kwargs, *extra, **extra_kwargs): + proxy = super().create_proxy(kind, target, args, kwargs, *extra, **extra_kwargs) + if isinstance(proxy, _MetaProxy) and proxy.meta is _UNKNOWN: + # A failure stays _UNKNOWN; only fatal if a shape question is asked. + with contextlib.suppress(Exception): + proxy.meta = self._infer_meta(kind, target, args, kwargs) + return proxy + + def _infer_meta(self, kind: str, target: object, args: tuple, kwargs: dict): + """Execute the op on meta tensors; PyTorch infers the output value.""" + if kind == "placeholder": + # vLLM always feeds the model [1, seq_len, hidden_size] hidden states. + weight = _reference_weight(self.root) + if str(target) == "hidden_states" and weight is not None: + return torch.empty( + 1, 8, weight.shape[-1], dtype=weight.dtype, device="meta" + ) + return _UNKNOWN + if kind == "get_attr": + value = operator.attrgetter(str(target))(self.root) + if isinstance(value, torch.Tensor): + value = torch.empty_like(value, device="meta") + return value + unknown = False + + def meta_of(arg: object) -> object: + nonlocal unknown + if isinstance(arg, fx.Proxy): + meta = getattr(arg, "meta", _UNKNOWN) + unknown = unknown or meta is _UNKNOWN + return meta + return arg + + meta_args = fx.node.map_aggregate(args, meta_of) + meta_kwargs = fx.node.map_aggregate(kwargs, meta_of) + if unknown: + return _UNKNOWN + if kind == "call_function": + return target(*meta_args, **meta_kwargs) + if kind == "call_method": + receiver, *rest = meta_args + return getattr(receiver, str(target))(*rest, **meta_kwargs) + if kind == "call_module": + # Run the child's forward with all its state on "meta", without + # mutating it (at match time params may be meta but buffers real). + # fx patches `nn.Module.__call__` while tracing; restore the real + # one so this execution is not itself recorded. + child = self.root.get_submodule(str(target)) + state = { + name: torch.empty_like(tensor, device="meta") + for name, tensor in chain( + child.named_parameters(), child.named_buffers() + ) + } + with mock.patch.object(nn.Module, "__call__", _MODULE_CALL): + return torch.func.functional_call(child, state, meta_args, meta_kwargs) + return _UNKNOWN + + def _is_varkw(self, node: object) -> bool: + return ( + isinstance(node, fx.Node) + and node.op == "placeholder" + and str(node.target).lstrip("*") == self.varkw + ) def iter(self, obj: fx.Proxy): - length = _infer_len(obj.node) - if length is None: + # Assume kwargs is always empty to simplify tracing. + node = obj.node + if self._is_varkw(node) or ( + node.op == "call_method" + and node.target == "keys" + and self._is_varkw(node.args[0]) + ): + return iter(()) + meta = getattr(obj, "meta", _UNKNOWN) + if meta is _UNKNOWN: return super().iter(obj) - return iter([obj[i] for i in range(length)]) + return iter([obj[i] for i in range(len(meta))]) + + +def _as_leaf_call(fn: Callable, length: int | None = None) -> Callable: + """Wrap any callable so tracing records it as one opaque `call_function` node. + + Lets the trace continue past untraceable bodies. Only the proxy arguments carry into + the node's dataflow; the rest are dropped rather than lifted into the graph. + `length` declares how many values the callable returns, so unpacking its result also + traces. Called without proxies (i.e. outside tracing), the wrapper is a passthrough. + """ + + def leaf(*args, **kwargs): + proxies = tuple(arg for arg in args if isinstance(arg, fx.Proxy)) + if not proxies: + return fn(*args, **kwargs) + proxy = proxies[0].tracer.create_proxy("call_function", fn, proxies, {}) + proxy.node.meta["leaf_call"] = True + if length is not None: + # The body never executes, so fabricate a value of the declared length. + proxy.meta = (_UNKNOWN,) * length + return proxy + + return leaf + + +def _leaf_attention_interfaces(): + """Patch `AttentionInterface.get_interface` so traced forwards see a leaf node. + + `vllm_attention_function` needs runtime context so it is untraceable. + Every interface returns `(attn_output, attn_weights)`.""" + from transformers.modeling_utils import AttentionInterface + + original = AttentionInterface.get_interface + + def get_interface(self, *args, **kwargs): + return _as_leaf_call(original(self, *args, **kwargs), length=2) + + return mock.patch.object(AttentionInterface, "get_interface", get_interface) def trace(module: nn.Module) -> fx.Graph | None: - """Trace `module.forward`, returning the partial graph on failure. - - The graph is only evidence for matching, and the patterns sit at the top of - their forwards, so a trace that fails partway can still be matched.""" + """Trace `module.forward`, returning the partial graph on failure.""" + parameters = forward_parameters(type(module)) + # vLLM never passes `past_key_values` so it is always the default value of `None`. + # Make this concrete to simplify tracing. + concrete_args = None + if "past_key_values" in parameters: + concrete_args = {"past_key_values": None} + # Get the name of the kwargs parameter passed to module.forward (usually "kwargs") tracer = _AllLeafTracer() + tracer.varkw = next( + ( + p.name + for p in parameters.values() + if p.kind is inspect.Parameter.VAR_KEYWORD + ), + None, + ) try: - return tracer.trace(module) + with _leaf_attention_interfaces(): + return tracer.trace(module, concrete_args=concrete_args) except Exception as exc: logger.debug("Could not fully trace %s: %s", type(module), exc) return getattr(tracer, "graph", None) @@ -129,20 +278,27 @@ def recover_forward(cls: type[nn.Module]) -> tuple[ast.FunctionDef, Callable]: return funcdef, fn +def forward_parameters(cls: type[nn.Module]) -> dict[str, inspect.Parameter]: + """`cls.forward`'s signature parameters, or empty if uninspectable.""" + try: + return dict(inspect.signature(cls.forward).parameters) + except (TypeError, ValueError): + return {} + + def forward_input_count(cls: type[nn.Module]) -> int: """The number of tensor inputs `cls.forward` declares, excluding `self` and any `*args`/`**kwargs`. Read from the signature, so it is independent of whether the trace completes (unlike counting placeholders).""" - try: - params = list(inspect.signature(cls.forward).parameters.values())[1:] - except (ValueError, TypeError): + params = list(forward_parameters(cls).values()) + if not params: return 1 # uninspectable: assume a single input and let matching decide fixed = ( inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY, ) - return sum(1 for p in params if p.kind in fixed) + return sum(1 for p in params[1:] if p.kind in fixed) def compile_forward(funcdef: ast.FunctionDef, fn: Callable) -> Callable: @@ -227,6 +383,42 @@ def find_node(graph: fx.Graph, predicate: Callable[[fx.Node], bool]) -> fx.Node return next((n for n in graph.nodes if predicate(n)), None) +def output_value(graph: fx.Graph) -> object | None: + """The value the graph's `output` node returns, if the trace reached one.""" + output = find_node(graph, lambda n: n.op == "output") + if output is None or not output.args: + return None + return output.args[0] + + +def upstream_linear(node: object, module: nn.Module) -> fx.Node | None: + """Nearest linear producing `node`, walking back through splits/reshapes. + + Never walks through a leaf call (e.g. an attention interface): its inputs + are what attention consumes, not what produced the value.""" + stack = [node] + seen: set[fx.Node] = set() + while stack: + current = stack.pop() + if not isinstance(current, fx.Node) or current in seen: + continue + seen.add(current) + if is_linear(current, module): + return current + if current.op in ("call_function", "call_method") and not is_leaf_call(current): + stack.extend(current.args) + return None + + +def returned_linear(graph: fx.Graph, module: nn.Module) -> str | None: + """Name of the Linear producing the graph's (first) output value.""" + value = output_value(graph) + if isinstance(value, (tuple, list)) and value: + value = value[0] + linear = upstream_linear(value, module) + return None if linear is None else str(linear.target) + + def is_linear(node: fx.Node, module: nn.Module) -> bool: """Is node `nn.Linear.__call__()`.""" return node.op == "call_module" and isinstance( From 81962bb6995eaebd1e49998c2a91c9e01e24da27 Mon Sep 17 00:00:00 2001 From: "rongfu.leng" <lenronfu@gmail.com> Date: Mon, 27 Jul 2026 20:12:23 +0800 Subject: [PATCH 108/185] [Bugfix]Reject invalid FlashInfer MNNVL workspaces (#49043) Signed-off-by: lengrongfu <lenronfu@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../device_communicators/flashinfer_all_reduce.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/vllm/distributed/device_communicators/flashinfer_all_reduce.py b/vllm/distributed/device_communicators/flashinfer_all_reduce.py index 8d3e8170924..4e4635035cc 100644 --- a/vllm/distributed/device_communicators/flashinfer_all_reduce.py +++ b/vllm/distributed/device_communicators/flashinfer_all_reduce.py @@ -67,6 +67,12 @@ def _create_workspace( comm_backend=comm_backend, group=group, ) + if backend == "mnnvl" and not getattr(workspace, "mc_ptr", 0): + workspace.destroy() + logger.warning_once( + "FlashInfer MNNVL multicast is unavailable on the current topology." + ) + return None except Exception as e: if "multicast" in str(e).lower(): logger.warning_once( From 96fa3f42c95acc75f40a8fddbc3a621c2355a30d Mon Sep 17 00:00:00 2001 From: neweyes <328719365@qq.com> Date: Mon, 27 Jul 2026 20:16:42 +0800 Subject: [PATCH 109/185] [Perf] Skip ll_bf16 router GEMM warmup for non-MoE models (#49659) Signed-off-by: neweyes <328719365@qq.com> --- vllm/model_executor/warmup/kernel_warmup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index e21ab112e73..fae88eb7f4d 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -122,7 +122,7 @@ def kernel_warmup(worker: "Worker"): elif has_flashinfer() and current_platform.has_device_capability(90): flashinfer_autotune(worker.model_runner) - if current_platform.has_device_capability(90): + if current_platform.has_device_capability(90) and worker.model_config.is_moe: _warmup_ll_bf16_router_gemm() # FlashInfer attention warmup From a89015c6df8eeb37a843b717c97a5be1355de83d Mon Sep 17 00:00:00 2001 From: liminfei-amd <91481003+liminfei-amd@users.noreply.github.com> Date: Mon, 27 Jul 2026 21:24:42 +0800 Subject: [PATCH 110/185] [Perf] Make merge attention context count a runtime argument (#48739) Signed-off-by: liminfei-amd <91481003+liminfei-amd@users.noreply.github.com> --- vllm/v1/attention/ops/triton_merge_attn_states.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/v1/attention/ops/triton_merge_attn_states.py b/vllm/v1/attention/ops/triton_merge_attn_states.py index ca06c2970b5..e8f90efce98 100644 --- a/vllm/v1/attention/ops/triton_merge_attn_states.py +++ b/vllm/v1/attention/ops/triton_merge_attn_states.py @@ -158,10 +158,10 @@ def merge_attn_states( prefix_head_stride, output_head_stride, output_scale, + prefill_tokens_with_context, head_size, padded_head_size, output_lse is not None, - prefill_tokens_with_context, output_scale is not None, ) @@ -177,10 +177,10 @@ def merge_attn_states_kernel( prefix_head_stride, output_head_stride, output_scale, # scale tensor or None + prefill_tokens_with_context, HEAD_SIZE: tl.constexpr, PADDED_HEAD_SIZE: tl.constexpr, OUTPUT_LSE: tl.constexpr, - prefill_tokens_with_context: tl.constexpr, USE_FP8: tl.constexpr, FP8_MIN: tl.constexpr = float8_info.min, FP8_MAX: tl.constexpr = float8_info.max, From dbccc5ae328d7b9168bc33b278a9125cfd89cc69 Mon Sep 17 00:00:00 2001 From: "Rui \"Garry\" Gao" <garrygaogg@gmail.com> Date: Mon, 27 Jul 2026 21:42:35 +0800 Subject: [PATCH 111/185] [Model] Enable EVS for Qwen3.5 (#48912) Signed-off-by: Rui "Garry" Gao <garrygaogg@gmail.com> --- vllm/model_executor/models/qwen3_5.py | 44 +++++++++++++++++++-------- 1 file changed, 32 insertions(+), 12 deletions(-) diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index bcd2576b74f..47076337531 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -53,6 +53,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.sequence import IntermediateTensors +from vllm.tokenizers.registry import cached_tokenizer_from_config from vllm.transformers_utils.configs.qwen3_5 import Qwen3_5Config, Qwen3_5TextConfig from vllm.transformers_utils.configs.qwen3_5_moe import ( Qwen3_5MoeConfig, @@ -397,8 +398,7 @@ class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLMBase, QwenNextMixtureOfExperts): dummy_inputs=Qwen3VLDummyInputsBuilder, ) class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid): - # Qwen3.5 does not support multimodal pruning (EVS). - supports_multimodal_pruning = False + supports_multimodal_pruning = True packed_modules_mapping = Qwen3VLForConditionalGeneration.packed_modules_mapping | { "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"], @@ -416,8 +416,21 @@ class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid) self.model_config = vllm_config.model_config self.multimodal_config = multimodal_config self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" - # Qwen3.5 does not support multimodal pruning (EVS). - self.is_multimodal_pruning_enabled = False + self.is_multimodal_pruning_enabled = ( + multimodal_config.is_multimodal_pruning_enabled() + ) + self.video_pruning_rate = self.multimodal_config.video_pruning_rate + self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config) + + # attributes needed by EVS-related functions inherited from Qwen3-VL + self.use_deepstack = hasattr(config.vision_config, "deepstack_visual_indexes") + self.deepstack_num_level = ( + len(config.vision_config.deepstack_visual_indexes) + if self.use_deepstack + else 0 + ) + self.visual_dim = config.vision_config.out_hidden_size + self.multiscale_dim = self.visual_dim * self.deepstack_num_level with self._mark_tower_model(vllm_config, {"image", "video"}): self.visual = Qwen3_VisionTransformer( @@ -462,12 +475,6 @@ class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid) return inputs_embeds - def recompute_mrope_positions(self, *args, **kwargs): - raise NotImplementedError( - "Qwen3.5 does not support multimodal pruning (EVS). " - "recompute_mrope_positions should never be called." - ) - def forward( self, input_ids: torch.Tensor, @@ -628,8 +635,21 @@ class Qwen3_5MoeForConditionalGeneration( self.model_config = vllm_config.model_config self.multimodal_config = multimodal_config self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" - # Qwen3.5 does not support multimodal pruning (EVS). - self.is_multimodal_pruning_enabled = False + self.is_multimodal_pruning_enabled = ( + multimodal_config.is_multimodal_pruning_enabled() + ) + self.video_pruning_rate = self.multimodal_config.video_pruning_rate + self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config) + + # attributes needed by EVS-related functions inherited from Qwen3-VL + self.use_deepstack = hasattr(config.vision_config, "deepstack_visual_indexes") + self.deepstack_num_level = ( + len(config.vision_config.deepstack_visual_indexes) + if self.use_deepstack + else 0 + ) + self.visual_dim = config.vision_config.out_hidden_size + self.multiscale_dim = self.visual_dim * self.deepstack_num_level with self._mark_tower_model(vllm_config, {"image", "video"}): self.visual = Qwen3_VisionTransformer( From 59a6b0411d1817c712728a17fa8c2e0cdb4cf1f9 Mon Sep 17 00:00:00 2001 From: Nick Hill <nickhill123@gmail.com> Date: Mon, 27 Jul 2026 07:00:46 -0700 Subject: [PATCH 112/185] [Core] Fix internal LB load-balancing (#49204) Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../engine-core-client/src/client/state.rs | 12 +-- tests/v1/engine/test_engine_core_client.py | 93 ++++++++++++++++--- vllm/v1/core/sched/interface.py | 4 + vllm/v1/core/sched/scheduler.py | 4 + vllm/v1/engine/coordinator.py | 69 ++++++++------ vllm/v1/engine/core.py | 26 +++++- vllm/v1/engine/core_client.py | 39 ++++++-- 7 files changed, 188 insertions(+), 59 deletions(-) diff --git a/rust/src/engine-core-client/src/client/state.rs b/rust/src/engine-core-client/src/client/state.rs index fad117ac31d..28e701ba732 100644 --- a/rust/src/engine-core-client/src/client/state.rs +++ b/rust/src/engine-core-client/src/client/state.rs @@ -74,17 +74,13 @@ impl EngineRoutingState { /// /// Scheduler stats can raise the load estimate above the frontend-local /// view, but they should not lower it below requests this frontend has - /// already admitted. Waiting requests still get the same extra penalty - /// as the original `waiting * 4 + running` score. + /// already admitted. fn routing_score(&self) -> usize { - const WAITING_WEIGHT: usize = 4; - let Some(stats) = self.last_scheduler_stats else { return self.inflight; }; - let scheduler_total = stats.running + stats.waiting; - self.inflight.max(scheduler_total) + stats.waiting * (WAITING_WEIGHT - 1) + self.inflight.max(stats.running + stats.waiting) } /// Replace the local routing view with a fresh real scheduler snapshot. @@ -750,7 +746,7 @@ mod tests { } #[test] - fn routing_score_keeps_extra_waiting_penalty() { + fn routing_score_counts_waiting_without_extra_penalty() { let state = EngineRoutingState { inflight: 1, last_scheduler_stats: Some(EngineLoadSnapshot { @@ -759,7 +755,7 @@ mod tests { }), }; - assert_eq!(state.routing_score(), 14); + assert_eq!(state.routing_score(), 5); } #[test] diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index 0b44b205cd4..0bdca7ada99 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -8,6 +8,7 @@ import os import signal import time import uuid +from collections import Counter from concurrent.futures import Future from dataclasses import dataclass from threading import Thread @@ -27,7 +28,11 @@ from vllm.platforms import current_platform from vllm.pooling_params import LateInteractionParams, PoolingParams from vllm.usage.usage_lib import UsageContext from vllm.utils.torch_utils import set_default_torch_num_threads -from vllm.v1.engine import EngineCoreReadyResponse, EngineCoreRequest +from vllm.v1.engine import ( + EngineCoreOutputs, + EngineCoreReadyResponse, + EngineCoreRequest, +) from vllm.v1.engine.core import EngineCore from vllm.v1.engine.core_client import ( AsyncMPClient, @@ -198,13 +203,19 @@ def _make_pooling_request( ) -def test_dplb_late_interaction_sticky_routing(): +def _make_dplb_client(num_engines: int = 3, client_count: int = 1) -> DPLBAsyncMPClient: client = object.__new__(DPLBAsyncMPClient) - client.client_count = 1 + client.client_count = client_count client.reqs_in_flight = {} - client.core_engines = [b"\x00\x00", b"\x01\x00", b"\x02\x00"] - client.lb_engines = [[0, 0], [0, 0], [0, 0]] + client.engine_inflight = Counter() + client.core_engines = [bytes([i, 0]) for i in range(num_engines)] + client.lb_engines = [[0, 0, 0.0] for _ in range(num_engines)] client.eng_start_index = 0 + return client + + +def test_dplb_late_interaction_sticky_routing(): + client = _make_dplb_client() query_key = "rerank-abc-query-0" query_request = _make_pooling_request( @@ -223,12 +234,8 @@ def test_dplb_late_interaction_sticky_routing(): def test_dplb_non_late_interaction_still_uses_lb(): - client = object.__new__(DPLBAsyncMPClient) - client.client_count = 1 - client.reqs_in_flight = {} - client.core_engines = [b"\x00\x00", b"\x01\x00", b"\x02\x00"] - client.lb_engines = [[2, 1], [0, 0], [1, 0]] - client.eng_start_index = 0 + client = _make_dplb_client() + client.lb_engines = [[2, 1, 0.0], [0, 0, 0.0], [1, 0, 0.0]] request = make_request(SamplingParams(max_tokens=1)) chosen_engine = client.get_core_engine_for_request(request) @@ -237,6 +244,70 @@ def test_dplb_non_late_interaction_still_uses_lb(): assert client.lb_engines[1][0] == 1 +def test_dplb_burst_round_robins_despite_snapshot_rebinds(): + """A stats snapshot rebind wipes the optimistic lb_engines increments; + the exact in-flight floor must keep a burst spreading round-robin.""" + client = _make_dplb_client(num_engines=4) + + for _ in range(4): + client.get_core_engine_for_request(make_request(SamplingParams(max_tokens=1))) + # Coordinator snapshot arrives, not yet reflecting the 4 routed requests. + client.lb_engines = [[0, 0, 0.0] for _ in range(4)] + for _ in range(4): + client.get_core_engine_for_request(make_request(SamplingParams(max_tokens=1))) + + assert sorted(client.engine_inflight.values()) == [2, 2, 2, 2] + + +def test_dplb_snapshot_backpressure_overrides_inflight(): + """An engine reported heavily loaded by the coordinator is avoided even + when this client has routed nothing to it.""" + client = _make_dplb_client(num_engines=2) + client.lb_engines = [[5, 10, 0.0], [0, 0, 0.0]] + + chosen = client.get_core_engine_for_request( + make_request(SamplingParams(max_tokens=1)) + ) + + assert chosen == client.core_engines[1] + + +def test_dplb_kv_pressure_amplifies_waiting_penalty(): + """A waiting queue on a KV-bound engine (slow drain) is penalized, while + the same queue with low KV usage is not (e.g. transient burst).""" + client = _make_dplb_client(num_engines=2) + # Engine 0 has a smaller total but is KV-bound with a queue. + client.lb_engines = [[5, 10, 1.0], [0, 20, 0.2]] + + chosen = client.get_core_engine_for_request( + make_request(SamplingParams(max_tokens=1)) + ) + assert chosen == client.core_engines[1] + + # Same counts without KV pressure: the smaller total wins. + client = _make_dplb_client(num_engines=2) + client.lb_engines = [[5, 10, 0.2], [0, 20, 0.2]] + + chosen = client.get_core_engine_for_request( + make_request(SamplingParams(max_tokens=1)) + ) + assert chosen == client.core_engines[0] + + +def test_dplb_finished_requests_release_inflight(): + client = _make_dplb_client(num_engines=2) + + req = make_request(SamplingParams(max_tokens=1)) + engine = client.get_core_engine_for_request(req) + assert client.engine_inflight[engine] == 1 + + outputs = EngineCoreOutputs(finished_requests={req.request_id}) + asyncio.run(DPLBAsyncMPClient.process_engine_outputs(client, outputs)) + + assert client.engine_inflight[engine] == 0 + assert req.request_id not in client.reqs_in_flight + + def test_apply_ready_response_syncs_block_size(): import msgspec diff --git a/vllm/v1/core/sched/interface.py b/vllm/v1/core/sched/interface.py index 4f13aa4b727..562946fb726 100644 --- a/vllm/v1/core/sched/interface.py +++ b/vllm/v1/core/sched/interface.py @@ -234,6 +234,10 @@ class SchedulerInterface(ABC): """Returns (num_running_reqs, num_waiting_reqs).""" raise NotImplementedError + def get_kv_cache_usage(self) -> float: + """Returns the fraction of the KV cache currently in use (0.0-1.0).""" + return 0.0 + @abstractmethod def make_stats(self) -> "SchedulerStats | None": """Make a SchedulerStats object for logging. diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index fff8c15e224..7fabf844a10 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -2154,6 +2154,10 @@ class Scheduler(SchedulerInterface): """Returns (num_running_reqs, num_waiting_reqs).""" return len(self.running), len(self.waiting) + len(self.skipped_waiting) + def get_kv_cache_usage(self) -> float: + """Returns the fraction of the KV cache currently in use (0.0-1.0).""" + return self.kv_cache_manager.usage + def add_request(self, request: Request) -> None: existing = self.requests.get(request.request_id) if existing is not None: diff --git a/vllm/v1/engine/coordinator.py b/vllm/v1/engine/coordinator.py index 151a07abe23..2f3b03636d7 100644 --- a/vllm/v1/engine/coordinator.py +++ b/vllm/v1/engine/coordinator.py @@ -139,7 +139,8 @@ class DPCoordinator: class EngineState: def __init__(self): - self.request_counts = [0, 0] # [waiting, running] + # [waiting, running, kv_cache_usage] + self.request_counts: list[int | float] = [0, 0, 0.0] class DPCoordinatorProc: @@ -202,7 +203,7 @@ class DPCoordinatorProc: stats_changed = False last_stats_step = -1 last_stats_wave = -1 - last_step_counts: list[list[int]] | None = None + last_step_counts: list[list[int | float]] | None = None with ( make_zmq_socket( @@ -260,8 +261,12 @@ class DPCoordinatorProc: wait_for = self.stats_update_interval_ms if stats_changed else 5000 # Wait at least 50ms to ensure we've received all stats for - # the current step. - min_timeout = 50 if last_step_counts is None else 0 + # the current step. Only applicable to lockstep (MoE) DP; + # non-lockstep engines have no synchronized step boundaries. + if self.enable_wave_coordination and last_step_counts is None: + min_timeout = 50 + else: + min_timeout = 0 events = poller.poll(timeout=max(min_timeout, wait_for - elapsed)) if not events: @@ -374,32 +379,38 @@ class DPCoordinatorProc: # 1. Updated request load stats - update our local # state with these. stats = self.engines[eng_index].request_counts - stats_step = scheduler_stats.step_counter - stats_wave = scheduler_stats.current_wave - if ( - stats_wave > last_stats_wave - or stats_wave == last_stats_wave - and stats_step > last_stats_step - ): - if stats_changed: - last_step_counts = self._get_engine_counts(do_copy=True) - last_stats_step = stats_step - last_stats_wave = stats_wave - elif stats_wave != last_stats_wave or ( - stats_step != last_stats_step - ): - logger.warning( - "Received stats for out-of-order " - "step (%d, %d) from engine %d (expected " - "> (%d, %d))", - stats_wave, - stats_step, - eng_index, - last_stats_wave, - last_stats_step, - ) + if self.enable_wave_coordination: + # Steps are synchronized across lockstep (MoE) DP + # ranks; snapshot counts at step boundaries. + stats_step = scheduler_stats.step_counter + stats_wave = scheduler_stats.current_wave + if ( + stats_wave > last_stats_wave + or stats_wave == last_stats_wave + and stats_step > last_stats_step + ): + if stats_changed: + last_step_counts = self._get_engine_counts( + do_copy=True + ) + last_stats_step = stats_step + last_stats_wave = stats_wave + elif stats_wave != last_stats_wave or ( + stats_step != last_stats_step + ): + logger.warning( + "Received stats for out-of-order " + "step (%d, %d) from engine %d (expected " + "> (%d, %d))", + stats_wave, + stats_step, + eng_index, + last_stats_wave, + last_stats_step, + ) stats[0] = scheduler_stats.num_waiting_reqs stats[1] = scheduler_stats.num_running_reqs + stats[2] = scheduler_stats.kv_cache_usage stats_changed = True # Wave coordination: handle wave completion and start notifications @@ -452,7 +463,7 @@ class DPCoordinatorProc: wave_encoded = msgspec.msgpack.encode((wave, exclude_engine_index)) socket.send_multipart((EngineCoreRequestType.START_DP_WAVE.value, wave_encoded)) - def _get_engine_counts(self, do_copy=False) -> list[list[int]]: + def _get_engine_counts(self, do_copy=False) -> list[list[int | float]]: """Return list of [waiting, running] count lists for each engine.""" if do_copy: return [copy.copy(e.request_counts) for e in self.engines] diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 8a62c8b2b1b..39393f64ae4 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -1057,6 +1057,7 @@ class EngineCoreProc(EngineCore): # Only publish request queue stats to coordinator for "internal" # and "hybrid" LB modes. self.publish_dp_lb_stats = internal_dp_balancing + self.last_counts = (0, 0) self.addresses = addresses self.process_input_queue_block = True @@ -1376,11 +1377,27 @@ class EngineCoreProc(EngineCore): while self._handle_shutdown(): # 1) Poll the input queue until there is work to do. self._process_input_queue() + # Publish request counts before and after GPU step to ensure freshness. + self._maybe_publish_request_counts() # 2) Step the engine core and return the outputs. self._process_engine_step() + self._maybe_publish_request_counts() raise SystemExit + def _maybe_publish_request_counts(self): + if not self.publish_dp_lb_stats: + return + + # Publish our request counts (if they've changed). + counts = self.scheduler.get_request_counts() + if counts != self.last_counts: + self.last_counts = counts + stats = SchedulerStats( + *counts, kv_cache_usage=self.scheduler.get_kv_cache_usage() + ) + self.output_queue.put_nowait((-1, EngineCoreOutputs(scheduler_stats=stats))) + def _process_input_queue(self): """Exits when an engine step needs to be performed.""" @@ -1890,7 +1907,6 @@ class DPEngineCoreProc(EngineCoreProc): # finished with DP peers every N steps. self.step_counter = 0 self.current_wave = 0 - self.last_counts = (0, 0) # Two-phase pause protocol state. When pending_pause is True, the # engine keeps stepping (dummy batches) while waiting for all DP @@ -2027,12 +2043,16 @@ class DPEngineCoreProc(EngineCoreProc): if not self.publish_dp_lb_stats: return - # Publish our request counts (if they've changed). + # Publish our request counts (if they've changed), stamped with the + # lockstep-synchronized step counter and wave number. counts = self.scheduler.get_request_counts() if counts != self.last_counts: self.last_counts = counts stats = SchedulerStats( - *counts, step_counter=self.step_counter, current_wave=self.current_wave + *counts, + kv_cache_usage=self.scheduler.get_kv_cache_usage(), + step_counter=self.step_counter, + current_wave=self.current_wave, ) self.output_queue.put_nowait((-1, EngineCoreOutputs(scheduler_stats=stats))) diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index f83e32096a9..0aa4b6f3312 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -7,7 +7,7 @@ import sys import uuid import weakref from abc import ABC, abstractmethod -from collections import defaultdict, deque +from collections import Counter, defaultdict, deque from collections.abc import Awaitable, Callable, Sequence from concurrent.futures import Future from dataclasses import dataclass @@ -1271,9 +1271,11 @@ class DPAsyncMPClient(AsyncMPClient): client_index, ) - # List of [waiting, running] pair per engine. + # List of [waiting, running, kv_cache_usage] per engine. # Used only by DPLBAsyncMPClient subclass. - self.lb_engines: list[list[int]] = [[0, 0] for _ in self.core_engines] + self.lb_engines: list[list[int | float]] = [ + [0, 0, 0.0] for _ in self.core_engines + ] self.eep_scaling_cache: ElasticScalingCache | None = None @@ -1351,7 +1353,7 @@ class DPAsyncMPClient(AsyncMPClient): ) if len(self.lb_engines) < new_engine_count: self.lb_engines = self.lb_engines + [ - [0, 0] + [0, 0, 0.0] for _ in range( new_engine_count - len(self.lb_engines) ) @@ -1445,6 +1447,9 @@ class DPLBAsyncMPClient(DPAsyncMPClient): # To route aborts to the correct engine. self.reqs_in_flight: dict[str, EngineIdentity] = {} + # Exact per-engine count of this client's unfinished requests. + self.engine_inflight: Counter[EngineIdentity] = Counter() + super().__init__( vllm_config, executor_class, @@ -1470,14 +1475,30 @@ class DPLBAsyncMPClient(DPAsyncMPClient): current_counts = self.lb_engines # TODO use P2C alg for larger DP sizes num_engines = len(current_counts) - min_score = sys.maxsize + min_score: float = sys.maxsize eng_index = 0 for i in range(num_engines): # Start from client_index to help with balancing when engines # are empty. idx = (self.eng_start_index + i) % num_engines - waiting, running = current_counts[idx] - score = waiting * 4 + running + waiting, running, kv_cache_usage = current_counts[idx] + # Estimate engine load as the greater of the coordinator's + # latest (waiting + running) snapshot and this client's own + # in-flight count (scaled by the number of clients). The + # in-flight floor is exact and can't be erased by a snapshot + # rebind, so a burst spreads round-robin even when snapshots + # race with routing decisions; the snapshot raises the score + # when other clients or stale requests load the engine. + inflight = self.engine_inflight[self.core_engines[idx]] + score: float = max(self.client_count * inflight, waiting + running) + if waiting: + # Waiting requests are penalized in proportion to KV cache + # pressure: a queue on a KV-bound engine drains slowly, so + # new requests should strongly prefer other engines. With + # low KV usage the queue is transient (e.g. mid-burst) and + # the penalty stays off, preserving exact round-robin. + # Ramps from 0 at <=50% usage to 3x waiting at 100%. + score += waiting * 6.0 * max(0.0, kv_cache_usage - 0.5) if score < min_score: min_score = score eng_index = idx @@ -1494,6 +1515,7 @@ class DPLBAsyncMPClient(DPAsyncMPClient): chosen_engine = self.core_engines[eng_index] # Record which engine is chosen for this request, to handle aborts. self.reqs_in_flight[request.request_id] = chosen_engine + self.engine_inflight[chosen_engine] += 1 return chosen_engine async def call_utility_async(self, method: str, *args) -> Any: @@ -1513,7 +1535,8 @@ class DPLBAsyncMPClient(DPAsyncMPClient): ): if outputs.finished_requests and self.reqs_in_flight: for req_id in outputs.finished_requests: - self.reqs_in_flight.pop(req_id, None) + if (engine := self.reqs_in_flight.pop(req_id, None)) is not None: + self.engine_inflight[engine] -= 1 @staticmethod async def eep_process_engine_core_notification( From 56c96b0d91f05140b61a2005c222d83e9ec042db Mon Sep 17 00:00:00 2001 From: "Roberto L. Castro" <38211239+LopezCastroRoberto@users.noreply.github.com> Date: Mon, 27 Jul 2026 16:25:37 +0200 Subject: [PATCH 113/185] [Perf] Tune LL BF16 Router GEMM (#48774) Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> --- .../kernels/linear/cute_dsl/ll_bf16.py | 70 ++++++++++++++++--- vllm/model_executor/warmup/kernel_warmup.py | 43 +++++++++--- 2 files changed, 93 insertions(+), 20 deletions(-) diff --git a/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py b/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py index 16bf9632965..b0471b56573 100644 --- a/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py +++ b/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py @@ -31,20 +31,71 @@ def is_available() -> bool: return _cutedsl_available +# Default configs _DEFAULT_DOTPROD_BS = 128 _DEFAULT_DOTPROD_MAX_M = 4 _DEFAULT_SPLITK_CONFIG = (6, 4) -_TUNED_DOTPROD_MAX_M: dict[tuple[int, int], int] = { - (7168, 256): 6, + +# ll_bf16 router shapes covered by warmup/tuning: +# (4096, 256) # DSV4-Flash +# (6144, 256) # GLM5.2 +# (7168, 256) # DSV3.2 +# (7168, 384) # DSV4-Pro + +# SM100f-specific tuned configs +_SM100F_TUNED_DOTPROD_BS: dict[tuple[int, int], dict[int, int]] = { + (6144, 256): {M: 256 for M in (1, 3, 4)}, } -_TUNED_CONFIGS: dict[tuple[int, int], dict[int, tuple[int, int]]] = { +_SM100F_TUNED_SPLITK_CONFIGS: dict[tuple[int, int], dict[int, tuple[int, int]]] = { + (4096, 256): { + **{M: (8, 5) for M in (5, 8)}, + 9: (8, 2), + }, + (7168, 256): {14: (8, 2)}, + (6144, 256): {M: (8, 2) for M in (9, 12, 16)}, + (7168, 384): {M: (7, 5) for M in (13, 16)}, +} + +# SM90-specific tuned configs +_SM90_TUNED_DOTPROD_BS: dict[tuple[int, int], dict[int, int]] = { + (4096, 256): {M: 256 for M in (1, 3)}, + (7168, 384): {M: 256 for M in (1, 2)}, +} +_SM90_TUNED_SPLITK_CONFIGS: dict[tuple[int, int], dict[int, tuple[int, int]]] = { + (4096, 256): { + **{M: (8, 2) for M in range(5, 8)}, + **{M: (8, 5) for M in range(10, 12)}, + **{M: (8, 2) for M in (13, 16)}, + **{M: (8, 5) for M in (14, 15)}, + }, + (7168, 256): {8: (6, 5)}, + (6144, 256): {M: (8, 2) for M in (9, 11)}, (7168, 384): { - 5: (4, 4), - **{M: (5, 4) for M in range(6, 17)}, + **{M: (8, 2) for M in (6, 8, 12)}, + **{M: (7, 5) for M in (7, 9, 10, 11, 13, 14, 15, 16)}, }, } +def _arch_tuned_configs() -> tuple[ + dict[tuple[int, int], dict[int, int]], + dict[tuple[int, int], dict[int, tuple[int, int]]], +]: + from vllm.platforms import current_platform + + if current_platform.is_device_capability_family(100): + return ( + _SM100F_TUNED_DOTPROD_BS, + _SM100F_TUNED_SPLITK_CONFIGS, + ) + if current_platform.is_device_capability(90): + return ( + _SM90_TUNED_DOTPROD_BS, + _SM90_TUNED_SPLITK_CONFIGS, + ) + return {}, {} + + _cute_ctx = None @@ -89,11 +140,12 @@ class LLBf16Gemm: self._splitk_cache: dict[tuple[int, int], Any] = {} def dispatch(self, *, M: int, K: int, N: int) -> CompileKey: - dotprod_max_m = _TUNED_DOTPROD_MAX_M.get((K, N), _DEFAULT_DOTPROD_MAX_M) - if dotprod_max_m >= M or K < 2048: - return self.CompileKey(backend="dotprod", M=M, K=K, bs=_DEFAULT_DOTPROD_BS) + tuned_bs, tuned_splitk = _arch_tuned_configs() + if M <= _DEFAULT_DOTPROD_MAX_M or K < 2048: + bs = tuned_bs.get((K, N), {}).get(M, _DEFAULT_DOTPROD_BS) + return self.CompileKey(backend="dotprod", M=M, K=K, bs=bs) - split_k, num_stages = _TUNED_CONFIGS.get((K, N), {}).get( + split_k, num_stages = tuned_splitk.get((K, N), {}).get( M, _DEFAULT_SPLITK_CONFIG ) return self.CompileKey(backend="splitk", split_k=split_k, num_stages=num_stages) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index fae88eb7f4d..e461dae0bb8 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -45,16 +45,30 @@ if TYPE_CHECKING: logger = init_logger(__name__) -_LL_BF16_WARMUP_MODEL_SHAPES: tuple[tuple[int, int], ...] = ( - (6144, 264), # Inkling - (7168, 256), # DSV3 - (7168, 384), # DSV4-Pro - (14400, 256), # DSV4-Flash -) _LL_BF16_WARMUP_M_RANGE = range(1, 17) -def _warmup_ll_bf16_router_gemm() -> None: +def _ll_bf16_router_shapes_from_model( + model: torch.nn.Module, +) -> tuple[tuple[int, int], ...]: + from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear + + shapes: set[tuple[int, int]] = set() + for module in model.modules(): + if not isinstance(module, GateLinear): + continue + weight = getattr(module, "weight", None) + if not isinstance(weight, torch.Tensor): + continue + if weight.dim() != 2 or weight.dtype != torch.bfloat16: + continue + n, k = weight.shape + if k % 8 == 0: + shapes.add((int(k), int(n))) + return tuple(sorted(shapes)) + + +def _warmup_ll_bf16_router_gemm(model: torch.nn.Module) -> None: from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ( is_available as is_ll_bf16_gemm_available, ) @@ -65,9 +79,16 @@ def _warmup_ll_bf16_router_gemm() -> None: if not is_ll_bf16_gemm_available(): return - logger.info("Warming up ll_bf16 router GEMM kernels.") + shapes = _ll_bf16_router_shapes_from_model(model) + if not shapes: + logger.info( + "Skipping ll_bf16 router GEMM warmup: no bf16 GateLinear shapes found." + ) + return + + logger.info("Warming up ll_bf16 router GEMM kernels for shapes: %s.", shapes) ll_bf16_gemm_kernel.warmup( - shapes=_LL_BF16_WARMUP_MODEL_SHAPES, + shapes=shapes, m_values=_LL_BF16_WARMUP_M_RANGE, ) @@ -122,8 +143,8 @@ def kernel_warmup(worker: "Worker"): elif has_flashinfer() and current_platform.has_device_capability(90): flashinfer_autotune(worker.model_runner) - if current_platform.has_device_capability(90) and worker.model_config.is_moe: - _warmup_ll_bf16_router_gemm() + if current_platform.has_device_capability(90): + _warmup_ll_bf16_router_gemm(worker.get_model()) # FlashInfer attention warmup # Only warmup if the model has FlashInfer attention groups From ef9975d021448b99a5408e8c78a4c4f6b63443c7 Mon Sep 17 00:00:00 2001 From: Guan-Ming Chiu <105915352+guan404ming@users.noreply.github.com> Date: Mon, 27 Jul 2026 22:37:54 +0800 Subject: [PATCH 114/185] [Bugfix] Reject pipeline parallelism for DiffusionGemma (#45828) Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com> --- vllm/model_executor/models/diffusion_gemma.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 1a5026d513f..1d457dd3958 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -61,7 +61,6 @@ from vllm.v1.worker.gpu.states import RequestState from .interfaces import ( SupportsMultiModal, - SupportsPP, SupportsQuant, ) @@ -143,7 +142,6 @@ class DiffusionGemmaForConditionalGeneration( nn.Module, SupportsMultiModal, SupportsQuant, - SupportsPP, ): """DiffusionGemma for vLLM. @@ -266,10 +264,6 @@ class DiffusionGemmaForConditionalGeneration( eps=getattr(text_config, "rms_norm_eps", 1e-6), ) - self.make_empty_intermediate_tensors = ( - self.model.make_empty_intermediate_tensors - ) - def compute_self_conditioning( self, inputs_embeds: torch.Tensor, From 27d7061ef62b3d853b4d335baf8ea9b3c56d9bf3 Mon Sep 17 00:00:00 2001 From: Umut Polat <52835619+umut-polat@users.noreply.github.com> Date: Mon, 27 Jul 2026 17:52:59 +0300 Subject: [PATCH 115/185] [Bugfix] Restore truncate_prompt_tokens for Jina rerank/score online (#49963) Signed-off-by: Umut Polat <52835619+umut-polat@users.noreply.github.com> --- .../test_jina_ranking_io_processor_unit.py | 59 +++++++++++++++++++ .../pooling/scoring/io_processor.py | 11 +++- 2 files changed, 69 insertions(+), 1 deletion(-) create mode 100644 tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py diff --git a/tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py b/tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py new file mode 100644 index 00000000000..c2fe62f609d --- /dev/null +++ b/tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for JinaRankingIOProcessor online request building.""" + +from unittest.mock import MagicMock + +import pytest + +from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor +from vllm.entrypoints.pooling.scoring.io_processor import JinaRankingIOProcessor +from vllm.entrypoints.pooling.scoring.protocol import RerankRequest +from vllm.entrypoints.pooling.scoring.typing import ScoringData + +pytestmark = pytest.mark.skip_global_cleanup + + +def test_online_forwards_truncate_prompt_tokens_to_proxy(monkeypatch): + """The proxy request handed to the base factory must carry + truncate_prompt_tokens/truncation_side from the real request. + + JinaRankingIOProcessor swaps ctx.request for a proxy + PoolingCompletionRequest before delegating to the base factory, which + reads truncation off ctx.request. Dropping the fields on the proxy + silently disables truncate_prompt_tokens for Jina rerank/score. + """ + proc = JinaRankingIOProcessor.__new__(JinaRankingIOProcessor) + proc.valid_inputs_online = MagicMock( + return_value=ScoringData(data_1=["query"], data_2=["doc"]) + ) + proc._get_token_limits = MagicMock(return_value=(0, 0)) + proc.ensure_str = MagicMock(side_effect=lambda data: list(data)) + proc.format_docs_prompts_func = MagicMock(return_value="formatted prompt") + + captured: dict[str, object] = {} + + def _spy_base(self, ctx): + captured["truncate_prompt_tokens"] = ctx.request.truncate_prompt_tokens + captured["truncation_side"] = ctx.request.truncation_side + return [] + + monkeypatch.setattr(PoolingIOProcessor, "get_request_factory_online", _spy_base) + + request = RerankRequest( + model="m", + query="query", + documents=["doc"], + truncate_prompt_tokens=512, + truncation_side="left", + ) + ctx = MagicMock() + ctx.request = request + ctx.prompt_extras = None + + proc.get_request_factory_online(ctx) + + assert captured["truncate_prompt_tokens"] == 512 + assert captured["truncation_side"] == "left" + # The real request is restored after delegating. + assert ctx.request is request diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index 8ebabca57cb..5dfe8a3a136 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -820,7 +820,16 @@ class JinaRankingIOProcessor(LateInteractionIOProcessor, JinaRankingIOProcessorM for q, d in zip(queries, docs) ] - ctx.request = PoolingCompletionRequest(task="token_embed", input=prompts) + # Forward truncation from the real request: the base factory reads + # these off ctx.request, so omitting them here silently drops + # truncate_prompt_tokens for Jina rerank/score (unlike the embed and + # bi/cross-encoder paths, which read them from the real request). + ctx.request = PoolingCompletionRequest( + task="token_embed", + input=prompts, + truncate_prompt_tokens=request.truncate_prompt_tokens, + truncation_side=request.truncation_side, + ) requests = PoolingIOProcessor.get_request_factory_online(self, ctx) ctx.request = request return requests From d2ca3002d93314a08bbddf9c6eb6ee78b1343407 Mon Sep 17 00:00:00 2001 From: Song Zhixin <szxfml@gmail.com> Date: Mon, 27 Jul 2026 23:31:36 +0800 Subject: [PATCH 116/185] [MRV2][Performance] Skip no-op FP32 logits materialization (#47711) Signed-off-by: jesse <szxfml@gmail.com> Signed-off-by: Song Zhixin <szxfml@gmail.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com> --- vllm/v1/worker/gpu/sample/sampler.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index e34e2acf377..3503b699f95 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -152,6 +152,9 @@ class Sampler: expanded_local_pos: torch.Tensor, skip_top_k_top_p: bool = False, ) -> torch.Tensor: + if not self._requires_logits_processing(idx_mapping_np): + return logits + # Copy logits to a new FP32 tensor. logits = torch.empty_like(logits, dtype=torch.float32).copy_(logits) @@ -194,6 +197,24 @@ class Sampler: logits, expanded_idx_mapping, idx_mapping_np ) + def _requires_logits_processing(self, idx_mapping_np: np.ndarray) -> bool: + if np.any(self.logit_bias_state.use_logit_bias[idx_mapping_np]): + return True + if np.any(self.penalties_state.use_penalty[idx_mapping_np]): + return True + if np.any(self.bad_words_state.num_bad_words.np[idx_mapping_np] > 0): + return True + + states = self.sampling_states + temperatures = states.temperature.np[idx_mapping_np] + if np.any((temperatures != 0.0) & (temperatures != 1.0)): + return True + if np.any(states.min_p.np[idx_mapping_np] != 0.0): + return True + if np.any(states.top_k.np[idx_mapping_np] != states.vocab_size): + return True + return bool(np.any(states.top_p.np[idx_mapping_np] != 1.0)) + def sample( self, logits: torch.Tensor, From 04502deca2f66b874d0fceae446c668c987413c0 Mon Sep 17 00:00:00 2001 From: Guan-Ming Chiu <105915352+guan404ming@users.noreply.github.com> Date: Mon, 27 Jul 2026 23:57:55 +0800 Subject: [PATCH 117/185] [Perf] Hash videos by source bytes (#49607) Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com> Signed-off-by: Guan-Ming Chiu <105915352+guan404ming@users.noreply.github.com> Co-authored-by: Isotr0py <2037008807@qq.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/multimodal/test_hasher.py | 25 +++++++++++++++++++++++++ vllm/multimodal/hasher.py | 6 ++++++ vllm/multimodal/inputs.py | 5 ++++- vllm/multimodal/media/base.py | 14 ++++++++++++-- vllm/multimodal/media/connector.py | 6 +++--- vllm/multimodal/media/video.py | 19 ++++++++++++------- vllm/multimodal/parse.py | 30 ++++++++++++++++++++++++------ vllm/multimodal/utils.py | 10 ++++++++-- 8 files changed, 94 insertions(+), 21 deletions(-) diff --git a/tests/multimodal/test_hasher.py b/tests/multimodal/test_hasher.py index fdedcaea27c..f1ddc378753 100644 --- a/tests/multimodal/test_hasher.py +++ b/tests/multimodal/test_hasher.py @@ -9,6 +9,8 @@ import torch from PIL import Image, ImageDraw from vllm.multimodal.hasher import MultiModalHasher +from vllm.multimodal.media.base import MediaWithBytes +from vllm.multimodal.parse import MultiModalDataParser pytestmark = pytest.mark.cpu_test @@ -82,6 +84,29 @@ def test_hash_collision_array_shape(): assert hasher.hash_kwargs(data=arr1) != hasher.hash_kwargs(data=arr2) +def test_hash_collision_video_num_frames(): + source = b"x" * 100 + + def item_for_hash(num_frames: int): + frames: np.ndarray = np.zeros((num_frames, 8, 8, 3), dtype=np.uint8) + metadata = { + "total_num_frames": 16, + "fps": 2.0, + "duration": 8.0, + "video_backend": "opencv", + "frames_indices": list(range(num_frames)), + "do_sample_frames": False, + } + video = MediaWithBytes((frames, metadata), source) + items = MultiModalDataParser()._parse_video_data([video]) + return items.get_all_items_for_hash()[0] + + hasher = MultiModalHasher + assert hasher.hash_kwargs(video=item_for_hash(2)) != hasher.hash_kwargs( + video=item_for_hash(4) + ) + + def test_hash_non_contiguous_array(): arr = np.arange(24).reshape(4, 6).T assert not arr.flags.c_contiguous diff --git a/vllm/multimodal/hasher.py b/vllm/multimodal/hasher.py index 6caf9c11427..ba9bcd0ea81 100644 --- a/vllm/multimodal/hasher.py +++ b/vllm/multimodal/hasher.py @@ -83,6 +83,12 @@ class MultiModalHasher: return cls.iter_item_to_bytes("image", obj.original_bytes) + if isinstance(obj, MediaWithBytes) and isinstance(obj.media, np.ndarray): + frames = obj.media + if frames.nbytes < len(obj.original_bytes): + return cls.iter_item_to_bytes("video", frames) + return cls.iter_item_to_bytes("video", obj.original_bytes) + if isinstance(obj, torch.Tensor): tensor_obj: torch.Tensor = obj.cpu() tensor_dtype = tensor_obj.dtype diff --git a/vllm/multimodal/inputs.py b/vllm/multimodal/inputs.py index c55bbe4623c..71f17a8648b 100644 --- a/vllm/multimodal/inputs.py +++ b/vllm/multimodal/inputs.py @@ -66,7 +66,10 @@ these are directly passed to the model without HF processing. """ VideoItem: TypeAlias = Union[ - HfVideoItem, "torch.Tensor", tuple[HfVideoItem, dict[str, Any]] + HfVideoItem, + "torch.Tensor", + tuple[HfVideoItem, dict[str, Any]], + MediaWithBytes[tuple[HfVideoItem, dict[str, Any]]], ] """ A `transformers.video_utils.VideoInput` representing a single video item. diff --git a/vllm/multimodal/media/base.py b/vllm/multimodal/media/base.py index 91e7a494717..3c072a09a54 100644 --- a/vllm/multimodal/media/base.py +++ b/vllm/multimodal/media/base.py @@ -2,9 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from abc import ABC, abstractmethod +from collections.abc import Iterable, Iterator from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Generic, TypeVar +from typing import Any, Generic, TypeVar, cast import numpy as np @@ -22,7 +23,8 @@ class MediaWithBytes(Generic[_T]): The wrapper delegates attribute access to the underlying media object, making it behave transparently like the wrapped type (e.g., PIL.Image). - NOTE: Currently, this wrapper is used only for the image modality. + NOTE: Currently, this wrapper is used only for the image and video + modalities. """ media: _T @@ -32,6 +34,14 @@ class MediaWithBytes(Generic[_T]): """Allow np.array(obj) to return np.array(obj.media).""" return np.array(self.media, *args, **kwargs) + def __iter__(self) -> Iterator[Any]: + """Allow unpacking obj to unpack obj.media (e.g. video tuples).""" + return iter(cast(Iterable[Any], self.media)) + + def __getitem__(self, index: Any) -> Any: + """Allow obj[i] to index obj.media (e.g. video tuples).""" + return cast(Any, self.media)[index] + def __getstate__(self): return self.__dict__.copy() diff --git a/vllm/multimodal/media/connector.py b/vllm/multimodal/media/connector.py index 656f537c5b1..fed41657e98 100644 --- a/vllm/multimodal/media/connector.py +++ b/vllm/multimodal/media/connector.py @@ -29,7 +29,7 @@ from vllm.multimodal.video import get_video_loader_backend_for_processor from vllm.utils.registry import ExtensionManager from .audio import AudioEmbeddingMediaIO, AudioMediaIO -from .base import MediaIO +from .base import MediaIO, MediaWithBytes from .image import ImageEmbeddingMediaIO, ImageMediaIO from .video import VideoMediaIO @@ -536,7 +536,7 @@ class MediaConnector: *, image_mode: str | None = "RGB", video_processor: str | None = None, - ) -> tuple[npt.NDArray, dict[str, Any]]: + ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: """ Load video from an HTTP or base64 data URL. """ @@ -562,7 +562,7 @@ class MediaConnector: *, image_mode: str | None = "RGB", video_processor: str | None = None, - ) -> tuple[npt.NDArray, dict[str, Any]]: + ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: """ Asynchronously load video from an HTTP or base64 data URL. diff --git a/vllm/multimodal/media/video.py b/vllm/multimodal/media/video.py index 978cf06c23a..124dcd7f7e5 100644 --- a/vllm/multimodal/media/video.py +++ b/vllm/multimodal/media/video.py @@ -13,13 +13,13 @@ from vllm import envs from vllm.logger import init_logger from ..video import VIDEO_LOADER_REGISTRY -from .base import MediaIO +from .base import MediaIO, MediaWithBytes from .image import ImageMediaIO logger = init_logger(__name__) -class VideoMediaIO(MediaIO[tuple[npt.NDArray, dict[str, Any]]]): +class VideoMediaIO(MediaIO[MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]]): """Configuration values can be user-provided either by --media-io-kwargs or by the runtime API field "media_io_kwargs". Ensure proper validation and error handling. @@ -91,14 +91,17 @@ class VideoMediaIO(MediaIO[tuple[npt.NDArray, dict[str, Any]]]): self.kwargs = kwargs self.video_loader = VIDEO_LOADER_REGISTRY.load(video_loader_backend) - def load_bytes(self, data: bytes) -> tuple[npt.NDArray, dict[str, Any]]: - return self.video_loader.load_bytes( + def load_bytes( + self, data: bytes + ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: + video = self.video_loader.load_bytes( data, num_frames=self.num_frames, **self.kwargs ) + return MediaWithBytes(video, data) def load_base64( self, media_type: str, data: str - ) -> tuple[npt.NDArray, dict[str, Any]]: + ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: if media_type.lower() == "video/jpeg": load_frame = partial( self.image_io.load_base64, @@ -160,11 +163,13 @@ class VideoMediaIO(MediaIO[tuple[npt.NDArray, dict[str, Any]]]): "frames_indices": frames_indices, "do_sample_frames": self.kwargs.get("do_sample_frames", False), } - return frames, metadata + return MediaWithBytes((frames, metadata), data.encode()) return self.load_bytes(pybase64.b64decode(data)) - def load_file(self, filepath: Path) -> tuple[npt.NDArray, dict[str, Any]]: + def load_file( + self, filepath: Path + ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: with filepath.open("rb") as f: data = f.read() diff --git a/vllm/multimodal/parse.py b/vllm/multimodal/parse.py index 9d1e005da7f..729a5bad18a 100644 --- a/vllm/multimodal/parse.py +++ b/vllm/multimodal/parse.py @@ -366,6 +366,20 @@ class VideoProcessorItems(ProcessorBatchItems[HfVideoItem | None]): self.metadata = metadata + def _unwrap(self, item: Any) -> Any: + if isinstance(item, tuple): + frames, metadata = item + return super()._unwrap(frames), metadata + return super()._unwrap(item) + + def get_item_for_hash(self, index: int) -> Any: + item = self.data[index] + if isinstance(item, MediaWithBytes) and isinstance(self.metadata, list): + metadata = self.metadata[index] + if metadata is not None: + return item, metadata + return item + def get_num_frames(self, item_idx: int) -> int: video = self.get(item_idx) if video is None: @@ -552,7 +566,10 @@ class MultiModalDataParser: def _get_video_with_metadata( self, video: VideoItem, - ) -> tuple[np.ndarray, dict[str, Any] | None]: + ) -> tuple[np.ndarray | MediaWithBytes[np.ndarray], dict[str, Any] | None]: + if isinstance(video, MediaWithBytes): + new_video, metadata = self._get_video_with_metadata(video.media) + return MediaWithBytes(new_video, video.original_bytes), metadata if isinstance(video, tuple): return video if isinstance(video, list): @@ -653,7 +670,11 @@ class MultiModalDataParser: else: data_items = data # type: ignore[assignment] - new_videos = list[tuple[np.ndarray, dict[str, Any] | None]]() + new_videos = list[ + np.ndarray + | MediaWithBytes[np.ndarray] + | tuple[np.ndarray | MediaWithBytes[np.ndarray], dict[str, Any]] + ]() metadata_lst: list[dict[str, Any] | None] = [] for data_item in data_items: video, metadata = self._get_video_with_metadata(data_item) @@ -664,12 +685,9 @@ class MultiModalDataParser: "Please check your video input in `multi_modal_data`" ) new_videos.append((video, metadata)) - metadata_lst.append(metadata) else: new_videos.append(video) - - if not self.video_needs_metadata: - metadata = None + metadata_lst.append(metadata) return VideoProcessorItems(new_videos, metadata=metadata_lst) diff --git a/vllm/multimodal/utils.py b/vllm/multimodal/utils.py index 401e6e517bd..eae274a81e9 100644 --- a/vllm/multimodal/utils.py +++ b/vllm/multimodal/utils.py @@ -24,7 +24,13 @@ from .inputs import ( MultiModalKwargsItem, MultiModalSharedField, ) -from .media import AudioMediaIO, ImageMediaIO, MediaConnector, VideoMediaIO +from .media import ( + AudioMediaIO, + ImageMediaIO, + MediaConnector, + MediaWithBytes, + VideoMediaIO, +) if TYPE_CHECKING: import torch.types @@ -332,7 +338,7 @@ def fetch_image( def fetch_video( video_url: str, video_io_kwargs: dict[str, Any] | None = None, -) -> tuple[npt.NDArray, dict[str, Any]]: +) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: """ Args: video_url: URL of the video file to fetch. From 3f47a8384d90a6e29b05b18d84942f8d7c895def Mon Sep 17 00:00:00 2001 From: yzong-rh <yzong@redhat.com> Date: Mon, 27 Jul 2026 12:12:14 -0400 Subject: [PATCH 118/185] [Bugfix] Fix VLLM_ENFORCE_STRICT_TOOL_CALLING mutation in tests (#49846) Signed-off-by: Yifan Zong <yzong@redhat.com> --- tests/parser/test_include_reasoning.py | 28 +++++---------------- tests/parser/test_parse.py | 34 +++++++++----------------- tests/parser/test_streaming.py | 8 ++++++ 3 files changed, 26 insertions(+), 44 deletions(-) diff --git a/tests/parser/test_include_reasoning.py b/tests/parser/test_include_reasoning.py index 3d1893577ef..a75b3517356 100644 --- a/tests/parser/test_include_reasoning.py +++ b/tests/parser/test_include_reasoning.py @@ -7,31 +7,15 @@ streaming (parse_delta), and ParsableContext.append_output() paths. """ import json -import os import pytest -_STRICT_TOOL_CALLING_ENV = "VLLM_ENFORCE_STRICT_TOOL_CALLING" -_STRICT_TOOL_CALLING_ENV_VALUE = os.environ.get(_STRICT_TOOL_CALLING_ENV) -os.environ[_STRICT_TOOL_CALLING_ENV] = "0" - -from vllm.entrypoints.openai.chat_completion.protocol import ( # noqa: E402 - ChatCompletionRequest, -) -from vllm.entrypoints.openai.engine.protocol import DeltaMessage # noqa: E402 -from vllm.entrypoints.openai.responses.protocol import ResponsesRequest # noqa: E402 -from vllm.parser.abstract_parser import DelegatingParser # noqa: E402 -from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser # noqa: E402 -from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser # noqa: E402 - - -@pytest.fixture(scope="module", autouse=True) -def restore_strict_tool_calling_env(): - yield - if _STRICT_TOOL_CALLING_ENV_VALUE is None: - os.environ.pop(_STRICT_TOOL_CALLING_ENV, None) - else: - os.environ[_STRICT_TOOL_CALLING_ENV] = _STRICT_TOOL_CALLING_ENV_VALUE +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.engine.protocol import DeltaMessage +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.parser.abstract_parser import DelegatingParser +from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser +from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser class ThinkReasoningParser(BaseThinkingReasoningParser): diff --git a/tests/parser/test_parse.py b/tests/parser/test_parse.py index 2a34ac7eea7..2379b8a4c2b 100644 --- a/tests/parser/test_parse.py +++ b/tests/parser/test_parse.py @@ -2,34 +2,24 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json -import os from types import SimpleNamespace import pytest -_STRICT_TOOL_CALLING_ENV = "VLLM_ENFORCE_STRICT_TOOL_CALLING" -_STRICT_TOOL_CALLING_ENV_VALUE = os.environ.get(_STRICT_TOOL_CALLING_ENV) -os.environ[_STRICT_TOOL_CALLING_ENV] = "0" - -from vllm.entrypoints.openai.chat_completion.protocol import ( # noqa: E402 - ChatCompletionRequest, -) -from vllm.entrypoints.openai.responses.protocol import ResponsesRequest # noqa: E402 -from vllm.parser.abstract_parser import DelegatingParser # noqa: E402 -from vllm.parser.utils import count_history_tool_calls # noqa: E402 -from vllm.reasoning.basic_parsers import ( # noqa: E402 - BaseThinkingReasoningParser, -) -from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser # noqa: E402 +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.parser.abstract_parser import DelegatingParser +from vllm.parser.utils import count_history_tool_calls +from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser +from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser -@pytest.fixture(scope="module", autouse=True) -def restore_strict_tool_calling_env(): - yield - if _STRICT_TOOL_CALLING_ENV_VALUE is None: - os.environ.pop(_STRICT_TOOL_CALLING_ENV, None) - else: - os.environ[_STRICT_TOOL_CALLING_ENV] = _STRICT_TOOL_CALLING_ENV_VALUE +@pytest.fixture(autouse=True) +def enable_hermes_required_named_parsing(monkeypatch): + # With VLLM_ENFORCE_STRICT_TOOL_CALLING (default on), Hermes sets + # supports_required_and_named=False and uses structural-tag guided tool + # calling. Force it True to test the non-guided JSON required/named parse path. + monkeypatch.setattr(Hermes2ProToolParser, "supports_required_and_named", True) class ThinkReasoningParser(BaseThinkingReasoningParser): diff --git a/tests/parser/test_streaming.py b/tests/parser/test_streaming.py index 1e3cca57d6d..619aa0e491a 100644 --- a/tests/parser/test_streaming.py +++ b/tests/parser/test_streaming.py @@ -14,6 +14,14 @@ from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser +@pytest.fixture(autouse=True) +def enable_hermes_required_named_parsing(monkeypatch): + # With VLLM_ENFORCE_STRICT_TOOL_CALLING (default on), Hermes sets + # supports_required_and_named=False and uses structural-tag guided tool + # calling. Force it True to test the non-guided JSON required/named parse path. + monkeypatch.setattr(Hermes2ProToolParser, "supports_required_and_named", True) + + class ThinkReasoningParser(BaseThinkingReasoningParser): @property def start_token(self) -> str: From 2b465b2c42e6f7d37fbbc67956dc9741e832dc29 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= <nicolo.lucchesi@mistral.ai> Date: Mon, 27 Jul 2026 18:51:37 +0200 Subject: [PATCH 119/185] [Misc][PD] Nixl cleanup `get_backend_aware_kv_block_len` and `virtually_split_kv_in_blocks` (#49988) Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai> --- .../kv_connector/unit/test_nixl_connector.py | 12 -- .../kv_connector/v1/nixl/base_worker.py | 117 +++++++----------- 2 files changed, 44 insertions(+), 85 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 8bb554a0433..d58f117fa31 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -1115,18 +1115,6 @@ class TestNixlHandshake: block_lens=[remote_block_len], ) - assert worker.get_backend_aware_kv_block_len(0, mamba_view=False) == ( - local_block_len - ) - assert ( - worker.get_backend_aware_kv_block_len(0, first_split=True, mamba_view=True) - == worker._mamba_ssm_size[0] - ) - assert ( - worker.get_backend_aware_kv_block_len(0, first_split=False, mamba_view=True) - == worker._mamba_ssm_size[1] - ) - assert worker._build_fa_remote(plan, meta, block_size_ratio=1).tolist() == [ [0x1000 + local_block_len, local_block_len, 0] ] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 671f5272dbf..d4e516d6df7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -96,7 +96,6 @@ class NixlBaseConnectorWorker: physical_blocks_per_logical: int, ) -> np.ndarray: """Compute NIXL descriptor IDs for given block IDs.""" - num_fa_regions = self.num_regions num_ssm_regions = 0 if self._has_mamba: assert self._conv_decomp is not None @@ -108,7 +107,7 @@ class NixlBaseConnectorWorker: num_blocks = dst_num_blocks if block_size_ratio is not None: num_blocks = int(num_blocks * block_size_ratio) - num_fa_descs = num_fa_regions * num_blocks + num_fa_descs = self.num_regions * num_blocks # All-attention fast path: single vectorized broadcast. if num_ssm_regions == 0: @@ -119,7 +118,7 @@ class NixlBaseConnectorWorker: # always differ (different areas). Therefore we can just flatten the # block_ids and compute the descs ids for all groups at once. block_arr = np.concatenate(block_ids)[None, :] - region_ids = np.arange(num_fa_regions)[:, None] + region_ids = np.arange(self.num_regions)[:, None] return (region_ids * num_blocks + block_arr).flatten() # Compute desc ids per group using the right stride: FA descs have @@ -130,7 +129,7 @@ class NixlBaseConnectorWorker: for i, group in enumerate(block_ids): group_arr = np.asarray(group) if _is_attention_spec(self._group_spec_types[i]): - fa_region_ids = np.arange(num_fa_regions)[:, None] + fa_region_ids = np.arange(self.num_regions)[:, None] all_descs.append( (fa_region_ids * num_blocks + group_arr[None, :]).flatten() ) @@ -210,7 +209,7 @@ class NixlBaseConnectorWorker: def _fa_desc_replicated(self, num_fa_descs: int) -> list[bool]: """Per-FA-descriptor replicate flag, in _build_fa_local emission order - (region-major; K then optional V per region). Length ``num_fa_descs``. + (region-major; one desc per block, with K/V packed). Length ``num_fa_descs``. """ assert self.transfer_topo is not None n_regions = len(self.block_len_per_layer) @@ -1082,21 +1081,15 @@ class NixlBaseConnectorWorker: # With hybrid allocator, layers can share a kv cache tensor seen_base_addresses = [] - # Note(tms): I modified this from the original region setup code. - # K and V are now in different regions. Advantage is that we can - # elegantly support MLA and any cases where the K and V tensors - # are non-contiguous (it's not locally guaranteed that they will be) - # Disadvantage is that the encoded NixlAgentMetadata is now larger - # (roughly 8KB vs 5KB). - # Conversely for FlashInfer, K and V are registered in the same region - # to better exploit the memory layout (ie num_blocks is the first dim). + # K and V are packed into the content dim, so each attention layer is a + # single NIXL region whose block transfers as one unit. Mamba layers instead + # register separate conv/ssm sub-regions (see `_build_mamba_local`). tensor_size_bytes = None for layer_name, cache in xfer_buffers.items(): - # NOTE (NickLucche) Hybrid SSM models assume a layout that is similar to - # that of FI, with block laid out as in `get_backend_aware_kv_block_len`. - # However, physical page_size may differ when kernel requires a specific - # block size. This leads to SSM and FA layers having different num_blocks. + # NOTE (NickLucche) Hybrid SSM mamba/FA physical page_size may differ when + # kernel requires a specific block size. This leads to SSM and FA layers + # having different num_blocks. # `_physical_blocks_per_logical_kv_block` ratio is used to adjust for this. layer_spec = self._layer_specs.get(layer_name) if layer_spec is None: @@ -1107,7 +1100,7 @@ class NixlBaseConnectorWorker: ) continue if isinstance(layer_spec, UniformTypeKVCacheSpecs): - # MLA DSv32 Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs + # DSA Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs layer_spec = layer_spec.kv_cache_specs[layer_name] # `layer_spec.page_size_bytes` only accounts for logical page_size, that is # the page_size assuming constant `self._logical_num_blocks`. @@ -1270,7 +1263,32 @@ class NixlBaseConnectorWorker: block_size_ratio: int, ) -> np.ndarray: """Build desc regions (conv sub-projections + ssm) per layer for - local mamba blocks with DS conv layout, as an Nx3 uint64 array.""" + local mamba blocks with DS conv layout, as an Nx3 uint64 array. + + A Mamba block interleaves conv and SSM state, which crucially differ in + size, so the two are indexed as separate sub-regions. Attention blocks + instead pack K and V into the content dim and transfer as a single unit. + Reference diagram: + KVCacheTensor (Shared) + / \\ + / \\ + / \\ + Attention (FlashInfer) View Mamba View + | | + | | + +-------------------+ +-------------------+ + | KVCacheTensor | | KVCacheTensor | + | | | | + |<----- page ------>| |<----- page ------->| + | size | | size | + | Key 0 | Val 0 | |Conv 0 | SSM 0 | + | Key 1 | Val 1 | |Conv 1 | SSM 1 | + | ... | ... | | ... | ... | + | Key N-2 | Val N-2 | |Conv N-2| SSM N-2 | + | Key N-1 | Val N-1 | |Conv N-1| SSM N-1 | + +-------------------+ +--------------------+ + |1st_split-2nd_split| |1st_split-2nd_split | + """ assert block_size_ratio == 1, ( "Mamba 3-read transfer with block_size_ratio != 1 is not tested. " f"Got block_size_ratio={block_size_ratio}." @@ -1362,15 +1380,11 @@ class NixlBaseConnectorWorker: block_arange = np.arange(num_blocks, dtype=np.uint64) parts: list[np.ndarray] = [] for i, base_addr in enumerate(base_addresses): - kv_block_len = ( - self.get_backend_aware_kv_block_len( - layer_idx=i, first_split=True, mamba_view=False - ) - // block_size_ratio - ) - page_stride = self.block_len_per_layer[i] // block_size_ratio - addrs = base_addr + block_arange * page_stride - parts.append(self._stack_descs(addrs, kv_block_len, device_id)) + # K/V are packed into the content dim, so the whole block transfers + # as one unit: desc length equals the block stride. + block_len = self.block_len_per_layer[i] // block_size_ratio + addrs = base_addr + block_arange * block_len + parts.append(self._stack_descs(addrs, block_len, device_id)) return np.concatenate(parts) def _build_fa_remote( @@ -1397,9 +1411,7 @@ class NixlBaseConnectorWorker: for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr): replicated = self._is_region_replicated(i) # Read our whole local region size from remote.. - local_block_len = self.get_backend_aware_kv_block_len( - layer_idx=i, first_split=True, mamba_view=False - ) + local_block_len = self.block_len_per_layer[i] remote_kv_block_len = local_block_len // block_size_ratio if block_size_ratio > 1: # ..using remote kv_block_len as transfer unit @@ -1521,8 +1533,7 @@ class NixlBaseConnectorWorker: ) return self._remote_agents[engine_id][(0, remote_tp_rank)] - # Compare physical regions, not self.num_regions (doubled by - # FlashInfer's virtual K/V split). + # Number of physical regions registered locally (one per layer/tensor). num_local_regions = len(self.block_len_per_layer) if ( self.pp_size > 1 @@ -2301,46 +2312,6 @@ class NixlBaseConnectorWorker: remote_block_ids[i] = remote_group[:num_blocks] return local_block_ids, remote_block_ids - def get_backend_aware_kv_block_len( - self, layer_idx: int, first_split: bool = True, mamba_view: bool = False - ) -> int: - """ - Get the block length for one K/V element (K and V have the same size). - - For FA and other backends, this is equal to the length of the whole - block, as K and V are in separate regions. - For FlashInfer, this is half the length of the whole block, as K and V - share the same region. - Similarly, for SSM-based models, state and conv are interleaved, but crucially - the their size differs. - Reference diagram: - KVCacheTensor (Shared) - / \\ - / \\ - / \\ - Attention (FlashInfer) View Mamba View - | | - | | - +-------------------+ +-------------------+ - | KVCacheTensor | | KVCacheTensor | - | | | | - |<----- page ------>| |<----- page ------->| - | size | | size | - | Key 0 | Val 0 | |Conv 0 | SSM 0 | - | Key 1 | Val 1 | |Conv 1 | SSM 1 | - | ... | ... | | ... | ... | - | Key N-2 | Val N-2 | |Conv N-2| SSM N-2 | - | Key N-1 | Val N-1 | |Conv N-1| SSM N-1 | - +-------------------+ +--------------------+ - |1st_split-2nd_split| |1st_split-2nd_split | - """ - assert self.transfer_topo is not None - if self.transfer_topo.virtually_split_kv_in_blocks and mamba_view: - block_len = self._mamba_ssm_size[not first_split] - else: - block_len = self.block_len_per_layer[layer_idx] - return block_len - def get_kv_connector_stats(self) -> KVConnectorStats | None: """ Get the KV transfer stats for the connector. From e3c2fc3b3ca3f41466126cd2aa0b228eb8465844 Mon Sep 17 00:00:00 2001 From: Connor Carpenter <connorc@nvidia.com> Date: Mon, 27 Jul 2026 09:53:27 -0700 Subject: [PATCH 120/185] [Rust Frontend][gRPC] Add server and model discovery (#49491) Signed-off-by: Connor Carpenter <connorc@nvidia.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> --- rust/proto/control.proto | 53 +++++ .../{vllm_grpc.proto => inference.proto} | 16 +- rust/src/chat/src/lib.rs | 27 +++ rust/src/engine-core-client/src/client.rs | 11 + .../src/engine-core-client/src/mock_engine.rs | 7 + .../src/protocol/handshake.rs | 15 ++ .../src/tests/python_compat.py | 14 ++ rust/src/server/build.rs | 8 +- rust/src/server/src/grpc/control.rs | 106 ++++++++++ rust/src/server/src/grpc/health.rs | 8 +- rust/src/server/src/grpc/inference.rs | 156 ++++++++++++++ rust/src/server/src/grpc/mod.rs | 194 +---------------- rust/src/server/src/grpc/tests.rs | 195 ++++++++++++++---- rust/src/server/src/lib.rs | 10 +- rust/src/server/src/middleware/offload.rs | 8 +- tests/v1/engine/test_engine_core_client.py | 7 + vllm/v1/engine/__init__.py | 7 + vllm/v1/engine/core.py | 42 ++-- 18 files changed, 615 insertions(+), 269 deletions(-) create mode 100644 rust/proto/control.proto rename rust/proto/{vllm_grpc.proto => inference.proto} (93%) create mode 100644 rust/src/server/src/grpc/control.rs create mode 100644 rust/src/server/src/grpc/inference.rs diff --git a/rust/proto/control.proto b/rust/proto/control.proto new file mode 100644 index 00000000000..7b858cb3e20 --- /dev/null +++ b/rust/proto/control.proto @@ -0,0 +1,53 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +syntax = "proto3"; +package vllm; + +service Control { + rpc GetServerInfo (GetServerInfoRequest) returns (ServerInfo) {} + rpc GetModelInfo (GetModelInfoRequest) returns (ModelInfo) {} + rpc Abort (AbortRequest) returns (AbortResponse) {} +} + +message GetServerInfoRequest {} + +message ServerInfo { + string engine_version = 1; + string api_version = 2; + string instance_id = 3; + ParallelismInfo parallelism = 4; + uint32 max_model_len = 5; + uint32 kv_block_size = 6; + uint64 total_kv_blocks = 7; + uint64 max_running_requests = 8; + uint64 max_batched_tokens = 9; +} + +message ParallelismInfo { + uint32 tensor_parallel_size = 1; + uint32 pipeline_parallel_size = 2; + uint32 data_parallel_size = 3; + uint32 data_parallel_rank = 4; + uint32 decode_context_parallel_size = 5; +} + +message GetModelInfoRequest {} + +message ModelInfo { + string model_id = 1; + string served_model_name = 2; + repeated string served_model_aliases = 3; + + bool supports_text_input = 20; + bool supports_token_ids_input = 21; + bool supports_multimodal = 23; + string reasoning_parser = 24; + string tool_call_parser = 25; +} + +message AbortRequest { + repeated string request_ids = 1; +} + +message AbortResponse {} diff --git a/rust/proto/vllm_grpc.proto b/rust/proto/inference.proto similarity index 93% rename from rust/proto/vllm_grpc.proto rename to rust/proto/inference.proto index c10da4818e3..b1c08ae76e6 100644 --- a/rust/proto/vllm_grpc.proto +++ b/rust/proto/inference.proto @@ -7,17 +7,13 @@ package vllm; import "google/protobuf/struct.proto"; -service Generate { +service Inference { // Generates text given a prompt rpc Generate (GenerateRequest) returns (GenerateResponse) {} // Generates text given a prompt, streaming the outputs rpc GenerateStream (GenerateRequest) returns (stream GenerateResponse) {} } -service Control { - rpc Abort (AbortRequest) returns (AbortResponse) {} -} - // ====================================================================================== // Generate Request // ====================================================================================== @@ -204,13 +200,3 @@ message CandidateTokenInfo { message TokenIds { repeated uint32 ids = 1; } - -// ====================================================================================== -// Control -// ====================================================================================== - -message AbortRequest { - repeated string request_ids = 1; -} - -message AbortResponse {} diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index e584257eaa4..9c38c40be4e 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -255,6 +255,33 @@ impl ChatLlm { self.text.engine_core_client() } + /// Whether the loaded backend has a registered multimodal processor. + pub fn supports_multimodal(&self) -> bool { + self.processor.backend.multimodal_model_info().is_some() + } + + /// Effective tool-call parser name for this model, if parsing is enabled. + pub fn tool_call_parser_name(&self) -> Option<&str> { + match &self.tool_call_parser { + ParserSelection::Auto => { + ToolParserFactory::global().resolve_name_for_model(self.model_id()) + } + ParserSelection::None => None, + ParserSelection::Explicit(name) => Some(name), + } + } + + /// Effective reasoning parser name for this model, if parsing is enabled. + pub fn reasoning_parser_name(&self) -> Option<&str> { + match &self.reasoning_parser { + ParserSelection::Auto => { + ReasoningParserFactory::global().resolve_name_for_model(self.model_id()) + } + ParserSelection::None => None, + ParserSelection::Explicit(name) => Some(name), + } + } + /// Render, tokenize, and submit one chat request. pub async fn chat(&self, request: ChatRequest) -> Result<ChatEventStream> { let (text_request, output_processor) = self diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index 9ea7b703671..705cb003e95 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -394,6 +394,17 @@ impl EngineCoreClient { self.engines.iter().map(|engine| &engine.ready_response).collect() } + /// Return the first engine's ready response. + /// + /// Per-engine fields such as `data_parallel_rank` should be read through + /// [`ready_responses`](Self::ready_responses). + pub fn ready_response(&self) -> &EngineCoreReadyResponse { + &self + .engines + .first() + .expect("engine core client requires at least one engine") + .ready_response + } /// Return the engine-reported effective model dtype. pub fn model_dtype(&self) -> ModelDtype { self.engines diff --git a/rust/src/engine-core-client/src/mock_engine.rs b/rust/src/engine-core-client/src/mock_engine.rs index e8f277b128d..5540ae8bc11 100644 --- a/rust/src/engine-core-client/src/mock_engine.rs +++ b/rust/src/engine-core-client/src/mock_engine.rs @@ -57,6 +57,13 @@ pub fn default_ready_response() -> EngineCoreReadyResponse { vllm_version: "test-vllm-version".to_string(), world_size: 1, data_parallel_size: 1, + tensor_parallel_size: 1, + pipeline_parallel_size: 1, + decode_context_parallel_size: 1, + data_parallel_rank: 0, + max_num_seqs: 256, + max_num_batched_tokens: 8192, + instance_id: "test-instance".to_string(), kv_cache_size_tokens: None, kv_cache_max_concurrency: None, } diff --git a/rust/src/engine-core-client/src/protocol/handshake.rs b/rust/src/engine-core-client/src/protocol/handshake.rs index c8545017a96..125a82b39d2 100644 --- a/rust/src/engine-core-client/src/protocol/handshake.rs +++ b/rust/src/engine-core-client/src/protocol/handshake.rs @@ -52,6 +52,21 @@ pub struct EngineCoreReadyResponse { pub world_size: u64, /// Data parallelism size from the parallel config. pub data_parallel_size: u64, + // Required discovery metadata; EngineCore and client versions must match. + /// Tensor-parallel size of this engine. + pub tensor_parallel_size: u32, + /// Pipeline-parallel size of this engine. + pub pipeline_parallel_size: u32, + /// Decode-context-parallel size of this engine. + pub decode_context_parallel_size: u32, + /// This engine's data-parallel rank. + pub data_parallel_rank: u32, + /// Scheduler cap on concurrently running sequences. + pub max_num_seqs: u64, + /// Scheduler cap on batched tokens per step. + pub max_num_batched_tokens: u64, + /// Unique identifier for this server instance. + pub instance_id: String, /// Total KV cache capacity in tokens, if reported. pub kv_cache_size_tokens: Option<u64>, /// Maximum achievable request concurrency given the KV cache, if reported. 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 0005ff20f92..ed16f2bd8b3 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -363,6 +363,13 @@ class EngineCoreReadyResponse: vllm_version: str world_size: int data_parallel_size: int + tensor_parallel_size: int + pipeline_parallel_size: int + decode_context_parallel_size: int + data_parallel_rank: int + max_num_seqs: int + max_num_batched_tokens: int + instance_id: str kv_cache_size_tokens: int | None = None kv_cache_max_concurrency: float | None = None @@ -376,6 +383,13 @@ ready_response = EngineCoreReadyResponse( vllm_version="0.0.0", data_parallel_size=1, world_size=1, + tensor_parallel_size=1, + pipeline_parallel_size=1, + decode_context_parallel_size=1, + data_parallel_rank=0, + max_num_seqs=256, + max_num_batched_tokens=8192, + instance_id="test-instance", ) print(msgspec.msgpack.encode(request).hex()) diff --git a/rust/src/server/build.rs b/rust/src/server/build.rs index c20ff1c86b8..585c3c70b99 100644 --- a/rust/src/server/build.rs +++ b/rust/src/server/build.rs @@ -9,7 +9,13 @@ fn main() -> Result<(), Box<dyn std::error::Error>> { .build_server(true) .build_client(true) .protoc_arg("--experimental_allow_proto3_optional") // be compatible with old compilers - .compile_protos(&[format!("{proto_dir}/vllm_grpc.proto")], &[proto_dir])?; + .compile_protos( + &[ + format!("{proto_dir}/control.proto"), + format!("{proto_dir}/inference.proto"), + ], + &[proto_dir], + )?; Ok(()) } diff --git a/rust/src/server/src/grpc/control.rs b/rust/src/server/src/grpc/control.rs new file mode 100644 index 00000000000..e989e6abaed --- /dev/null +++ b/rust/src/server/src/grpc/control.rs @@ -0,0 +1,106 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::sync::Arc; + +use thiserror_ext::AsReport as _; +use tonic::{Request, Response, Status}; +use vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse; + +use super::{ControlServer, pb}; +use crate::state::AppState; + +pub(crate) type ControlGrpcService = ControlServer<ControlServiceImpl>; + +/// gRPC control service backed by the shared application state. +pub struct ControlServiceImpl { + state: Arc<AppState>, +} + +impl ControlServiceImpl { + pub fn new(state: Arc<AppState>) -> Self { + Self { state } + } + + fn ready(&self) -> &EngineCoreReadyResponse { + self.state.engine_core_client().ready_response() + } + + fn parallelism_info(&self) -> pb::ParallelismInfo { + let ready = self.ready(); + pb::ParallelismInfo { + tensor_parallel_size: ready.tensor_parallel_size, + pipeline_parallel_size: ready.pipeline_parallel_size, + data_parallel_size: ready.data_parallel_size.min(u64::from(u32::MAX)) as u32, + data_parallel_rank: ready.data_parallel_rank, + decode_context_parallel_size: ready.decode_context_parallel_size, + } + } +} + +const GRPC_API_VERSION: &str = "vllm"; + +#[tonic::async_trait] +impl pb::control_server::Control for ControlServiceImpl { + async fn get_server_info( + &self, + _request: Request<pb::GetServerInfoRequest>, + ) -> Result<Response<pb::ServerInfo>, Status> { + let ready = self.ready(); + Ok(Response::new(pb::ServerInfo { + engine_version: ready.vllm_version.clone(), + api_version: GRPC_API_VERSION.to_string(), + instance_id: ready.instance_id.clone(), + parallelism: Some(self.parallelism_info()), + max_model_len: self.state.engine_core_client().max_model_len(), + kv_block_size: ready.block_size.min(u64::from(u32::MAX)) as u32, + total_kv_blocks: self.state.engine_core_client().total_num_gpu_blocks(), + max_running_requests: ready.max_num_seqs, + max_batched_tokens: ready.max_num_batched_tokens, + })) + } + + async fn get_model_info( + &self, + _request: Request<pb::GetModelInfoRequest>, + ) -> Result<Response<pb::ModelInfo>, Status> { + let served = self.state.served_model_names(); + Ok(Response::new(pb::ModelInfo { + model_id: self.state.chat.text().model_id().to_string(), + served_model_name: self.state.primary_model_name().to_string(), + served_model_aliases: served.iter().skip(1).cloned().collect(), + // GenerateRequest accepts both prompt representations. + supports_text_input: true, + supports_token_ids_input: true, + supports_multimodal: self.state.chat.supports_multimodal(), + reasoning_parser: self + .state + .chat + .reasoning_parser_name() + .unwrap_or_default() + .to_string(), + tool_call_parser: self + .state + .chat + .tool_call_parser_name() + .unwrap_or_default() + .to_string(), + })) + } + + async fn abort( + &self, + request: Request<pb::AbortRequest>, + ) -> Result<Response<pb::AbortResponse>, Status> { + let request_ids = request.into_inner().request_ids; + if request_ids.is_empty() { + return Ok(Response::new(pb::AbortResponse {})); + } + self.state + .chat + .abort(&request_ids) + .await + .map_err(|error| Status::internal(error.to_report_string()))?; + Ok(Response::new(pb::AbortResponse {})) + } +} diff --git a/rust/src/server/src/grpc/health.rs b/rust/src/server/src/grpc/health.rs index 554dfae675f..cd35aec972e 100644 --- a/rust/src/server/src/grpc/health.rs +++ b/rust/src/server/src/grpc/health.rs @@ -8,14 +8,14 @@ use tonic_health::ServingStatus; use tonic_health::server::HealthReporter; use tracing::{info, warn}; -use super::{ControlGrpcService, GenerateGrpcService}; +use super::{ControlGrpcService, InferenceGrpcService}; pub(crate) async fn monitor_health( mut health_reporter: HealthReporter, mut engine_health: watch::Receiver<bool>, shutdown: CancellationToken, ) { - let generate_service = GenerateGrpcService::NAME; + let inference_service = InferenceGrpcService::NAME; let control_service = ControlGrpcService::NAME; let status = ServingStatus::NotServing; let health_event_first = tokio::select! { @@ -45,7 +45,7 @@ pub(crate) async fn monitor_health( } }; - health_reporter.set_not_serving::<GenerateGrpcService>().await; + health_reporter.set_not_serving::<InferenceGrpcService>().await; health_reporter.set_not_serving::<ControlGrpcService>().await; // Both gRPC services use the same engine client, so overall server health // mirrors their shared engine health. @@ -59,7 +59,7 @@ pub(crate) async fn monitor_health( ); } - health_reporter.clear_service_status(generate_service).await; + health_reporter.clear_service_status(inference_service).await; health_reporter.clear_service_status(control_service).await; health_reporter.clear_service_status("").await; } diff --git a/rust/src/server/src/grpc/inference.rs b/rust/src/server/src/grpc/inference.rs new file mode 100644 index 00000000000..56fa40924cf --- /dev/null +++ b/rust/src/server/src/grpc/inference.rs @@ -0,0 +1,156 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::pin::Pin; +use std::sync::Arc; + +use futures::{Stream, StreamExt as _}; +use thiserror_ext::AsReport as _; +use tokio::sync::mpsc; +use tokio_stream::wrappers::ReceiverStream; +use tonic::{Request, Response, Status}; +use tracing::info; +use vllm_text::{DecodedTextEvent, TextOutputStreamExt as _}; + +use super::convert::{self, ResponseOpts}; +use super::{InferenceServer, pb}; +use crate::state::AppState; + +pub(crate) type InferenceGrpcService = InferenceServer<InferenceServiceImpl>; + +/// gRPC inference service backed by the shared application state. +pub struct InferenceServiceImpl { + state: Arc<AppState>, +} + +impl InferenceServiceImpl { + pub fn new(state: Arc<AppState>) -> Self { + Self { state } + } +} + +#[tonic::async_trait] +impl pb::inference_server::Inference for InferenceServiceImpl { + type GenerateStreamStream = + Pin<Box<dyn Stream<Item = Result<pb::GenerateResponse, Status>> + Send>>; + + /// Unary generate: collect all output and return a single response. + async fn generate( + &self, + request: Request<pb::GenerateRequest>, + ) -> Result<Response<pb::GenerateResponse>, Status> { + let proto_req = request.into_inner(); + let response_opts = ResponseOpts::from_proto(proto_req.response.as_ref()); + let text_request = + convert::to_text_request(proto_req, false, self.state.served_model_names())?; + + let request_id = text_request.request_id.clone(); + info!(%request_id, "grpc generate (unary)"); + + let stream = self.state.chat.text().generate(text_request).await; + let stream = stream.map_err(text_error_to_status)?; + + let collected = stream.collect_output().await.map_err(text_error_to_status)?; + + // Build the single aggregated response. + let prompt_info = convert::to_prompt_info( + &collected.prompt_token_ids, + collected.prompt_logprobs.as_ref(), + &response_opts, + ); + + let finish_info = vllm_text::Finished { + usage: collected.usage, + finish_reason: collected.finish_reason, + kv_transfer_params: collected.kv_transfer_params, + ec_transfer_params: collected.ec_transfer_params, + }; + + let outputs = convert::to_sequence_output( + &collected.text, + &collected.token_ids, + collected.logprobs.as_ref(), + Some(&finish_info), + &response_opts, + ); + + Ok(Response::new(pb::GenerateResponse { + prompt_info: Some(prompt_info), + outputs: Some(outputs), + })) + } + + /// Streaming generate: yield incremental responses as tokens are produced. + async fn generate_stream( + &self, + request: Request<pb::GenerateRequest>, + ) -> Result<Response<Self::GenerateStreamStream>, Status> { + let proto_req = request.into_inner(); + let response_opts = ResponseOpts::from_proto(proto_req.response.as_ref()); + let text_request = + convert::to_text_request(proto_req, true, self.state.served_model_names())?; + + let request_id = text_request.request_id.clone(); + info!(%request_id, "grpc generate (stream)"); + + let stream = self.state.chat.text().generate(text_request).await; + let stream = stream.map_err(text_error_to_status)?; + + let (tx, rx) = mpsc::channel(32); + + tokio::spawn(async move { + futures::pin_mut!(stream); + while let Some(event) = stream.next().await { + let response = match event { + Err(e) => Err(text_error_to_status(e)), + Ok(DecodedTextEvent::Start { + prompt_token_ids, + prompt_logprobs, + }) => { + let prompt_info = convert::to_prompt_info( + &prompt_token_ids, + prompt_logprobs.as_ref(), + &response_opts, + ); + Ok(pb::GenerateResponse { + prompt_info: Some(prompt_info), + outputs: None, + }) + } + Ok(DecodedTextEvent::TextDelta { + delta, + token_ids, + logprobs, + finished, + }) => Ok(pb::GenerateResponse { + prompt_info: None, + outputs: Some(convert::to_sequence_output( + &delta, + &token_ids, + logprobs.as_ref(), + finished.as_ref(), + &response_opts, + )), + }), + }; + + if tx.send(response).await.is_err() { + // Client disconnected. + break; + } + } + }); + + let response_stream = ReceiverStream::new(rx); + Ok(Response::new(Box::pin(response_stream))) + } +} + +fn text_error_to_status(error: vllm_text::Error) -> Status { + let message = error.to_report_string(); + if error.is_request_validation_error() { + Status::invalid_argument(message) + } else { + Status::internal(message) + } +} diff --git a/rust/src/server/src/grpc/mod.rs b/rust/src/server/src/grpc/mod.rs index f5d2f347bb3..06c7c9259eb 100644 --- a/rust/src/server/src/grpc/mod.rs +++ b/rust/src/server/src/grpc/mod.rs @@ -1,203 +1,25 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -//! gRPC Generate service backed by the shared [`vllm_text::TextLlm`] facade. +//! gRPC services backed by the shared application state. +mod control; mod convert; mod health; - -use std::pin::Pin; -use std::sync::Arc; - -use futures::{Stream, StreamExt as _}; -use thiserror_ext::AsReport as _; -use tokio::sync::mpsc; -use tokio_stream::wrappers::ReceiverStream; -use tonic::{Request, Response, Status}; -use tracing::info; -use vllm_text::{DecodedTextEvent, TextOutputStreamExt as _}; - -use self::convert::ResponseOpts; -use crate::state::AppState; +mod inference; /// Generated protobuf/gRPC types for the `vllm` package. pub mod pb { tonic::include_proto!("vllm"); } +pub(crate) use control::ControlGrpcService; +pub use control::ControlServiceImpl; pub(crate) use health::monitor_health; +pub(crate) use inference::InferenceGrpcService; +pub use inference::InferenceServiceImpl; pub use pb::control_server::ControlServer; -pub use pb::generate_server::GenerateServer; - -pub(crate) type ControlGrpcService = ControlServer<ControlServiceImpl>; -pub(crate) type GenerateGrpcService = GenerateServer<GenerateServiceImpl>; +pub use pb::inference_server::InferenceServer; #[cfg(test)] mod tests; - -/// gRPC Generate service implementation backed by the shared application state. -pub struct GenerateServiceImpl { - state: Arc<AppState>, -} - -impl GenerateServiceImpl { - pub fn new(state: Arc<AppState>) -> Self { - Self { state } - } -} - -/// gRPC control service backed by the shared application state. -pub struct ControlServiceImpl { - state: Arc<AppState>, -} - -impl ControlServiceImpl { - pub fn new(state: Arc<AppState>) -> Self { - Self { state } - } -} - -#[tonic::async_trait] -impl pb::control_server::Control for ControlServiceImpl { - async fn abort( - &self, - request: Request<pb::AbortRequest>, - ) -> Result<Response<pb::AbortResponse>, Status> { - let request_ids = request.into_inner().request_ids; - if request_ids.is_empty() { - return Ok(Response::new(pb::AbortResponse {})); - } - self.state - .chat - .abort(&request_ids) - .await - .map_err(|error| Status::internal(error.to_report_string()))?; - Ok(Response::new(pb::AbortResponse {})) - } -} - -#[tonic::async_trait] -impl pb::generate_server::Generate for GenerateServiceImpl { - type GenerateStreamStream = - Pin<Box<dyn Stream<Item = Result<pb::GenerateResponse, Status>> + Send>>; - - /// Unary generate: collect all output and return a single response. - async fn generate( - &self, - request: Request<pb::GenerateRequest>, - ) -> Result<Response<pb::GenerateResponse>, Status> { - let proto_req = request.into_inner(); - let response_opts = ResponseOpts::from_proto(proto_req.response.as_ref()); - let text_request = - convert::to_text_request(proto_req, false, self.state.served_model_names())?; - - let request_id = text_request.request_id.clone(); - info!(%request_id, "grpc generate (unary)"); - - let stream = self.state.chat.text().generate(text_request).await; - let stream = stream.map_err(text_error_to_status)?; - - let collected = stream.collect_output().await.map_err(text_error_to_status)?; - - // Build the single aggregated response. - let prompt_info = convert::to_prompt_info( - &collected.prompt_token_ids, - collected.prompt_logprobs.as_ref(), - &response_opts, - ); - - let finish_info = vllm_text::Finished { - usage: collected.usage, - finish_reason: collected.finish_reason, - kv_transfer_params: collected.kv_transfer_params, - ec_transfer_params: collected.ec_transfer_params, - }; - - let outputs = convert::to_sequence_output( - &collected.text, - &collected.token_ids, - collected.logprobs.as_ref(), - Some(&finish_info), - &response_opts, - ); - - Ok(Response::new(pb::GenerateResponse { - prompt_info: Some(prompt_info), - outputs: Some(outputs), - })) - } - - /// Streaming generate: yield incremental responses as tokens are produced. - async fn generate_stream( - &self, - request: Request<pb::GenerateRequest>, - ) -> Result<Response<Self::GenerateStreamStream>, Status> { - let proto_req = request.into_inner(); - let response_opts = ResponseOpts::from_proto(proto_req.response.as_ref()); - let text_request = - convert::to_text_request(proto_req, true, self.state.served_model_names())?; - - let request_id = text_request.request_id.clone(); - info!(%request_id, "grpc generate (stream)"); - - let stream = self.state.chat.text().generate(text_request).await; - let stream = stream.map_err(text_error_to_status)?; - - let (tx, rx) = mpsc::channel(32); - - tokio::spawn(async move { - futures::pin_mut!(stream); - while let Some(event) = stream.next().await { - let response = match event { - Err(e) => Err(text_error_to_status(e)), - Ok(DecodedTextEvent::Start { - prompt_token_ids, - prompt_logprobs, - }) => { - let prompt_info = convert::to_prompt_info( - &prompt_token_ids, - prompt_logprobs.as_ref(), - &response_opts, - ); - Ok(pb::GenerateResponse { - prompt_info: Some(prompt_info), - outputs: None, - }) - } - Ok(DecodedTextEvent::TextDelta { - delta, - token_ids, - logprobs, - finished, - }) => Ok(pb::GenerateResponse { - prompt_info: None, - outputs: Some(convert::to_sequence_output( - &delta, - &token_ids, - logprobs.as_ref(), - finished.as_ref(), - &response_opts, - )), - }), - }; - - if tx.send(response).await.is_err() { - // Client disconnected. - break; - } - } - }); - - let response_stream = ReceiverStream::new(rx); - Ok(Response::new(Box::pin(response_stream))) - } -} - -fn text_error_to_status(error: vllm_text::Error) -> Status { - let message = error.to_report_string(); - if error.is_request_validation_error() { - Status::invalid_argument(message) - } else { - Status::internal(message) - } -} diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 4cd67f4a4b3..75da670ec3b 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -25,12 +25,18 @@ use vllm_chat::{ ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor, DynChatOutputProcessor, DynChatRenderer, NewChatOutputProcessorOptions, RenderedPrompt, }; +use vllm_engine_core_client::mock_engine::{ + DEFAULT_MOCK_BLOCK_SIZE, DEFAULT_MOCK_MAX_MODEL_LEN, DEFAULT_MOCK_NUM_GPU_BLOCKS, + default_ready_response, +}; use vllm_engine_core_client::protocol::output::{ EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, RequestBatchOutputs, }; use vllm_engine_core_client::protocol::request::EngineCoreRequest; -use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task}; -use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId}; +use vllm_engine_core_client::test_utils::{ + IpcNamespace, spawn_mock_engine_task, spawn_mock_engine_task_with_ready, +}; +use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId, TransportMode}; use vllm_llm::Llm; use vllm_text::tokenizer::DynTokenizer; use vllm_text::{Prompt, TextBackend}; @@ -39,8 +45,8 @@ use zeromq::prelude::{SocketRecv, SocketSend}; use zeromq::{DealerSocket, PushSocket, ZmqMessage}; use super::pb::control_client::ControlClient; -use super::pb::generate_client::GenerateClient; -use super::{ControlServer, ControlServiceImpl, GenerateServer, GenerateServiceImpl, pb}; +use super::pb::inference_client::InferenceClient; +use super::{ControlServer, ControlServiceImpl, InferenceServer, InferenceServiceImpl, pb}; use crate::listener::{Listener, MaybeTlsListener}; use crate::state::AppState; use crate::tls; @@ -202,7 +208,7 @@ async fn setup_grpc_service( engine_id: impl Into<EngineId>, output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>, ) -> ( - GenerateServer<GenerateServiceImpl>, + InferenceServer<InferenceServiceImpl>, ControlServer<ControlServiceImpl>, tokio::sync::watch::Receiver<bool>, MockEngineTask, @@ -246,7 +252,7 @@ async fn setup_grpc_service( ); let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat)); ( - GenerateServer::new(GenerateServiceImpl::new(state.clone())), + InferenceServer::new(InferenceServiceImpl::new(state.clone())), ControlServer::new(ControlServiceImpl::new(state)), engine_health, engine_task, @@ -259,30 +265,30 @@ async fn grpc_test_server( engine_id: impl Into<EngineId>, output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>, ) -> ( - GenerateClient<tonic::transport::Channel>, + InferenceClient<tonic::transport::Channel>, tokio::task::JoinHandle<()>, MockEngineTask, ) { - let (generate_service, control_service, engine_health, engine_task) = + let (inference_service, control_service, engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await; let (channel, server_task) = start_grpc_test_server( - generate_service, + inference_service, control_service, engine_health, tokio_util::sync::CancellationToken::new(), ) .await; - (GenerateClient::new(channel), server_task, engine_task) + (InferenceClient::new(channel), server_task, engine_task) } async fn start_grpc_test_server( - generate_service: GenerateServer<GenerateServiceImpl>, + inference_service: InferenceServer<InferenceServiceImpl>, control_service: ControlServer<ControlServiceImpl>, engine_health: tokio::sync::watch::Receiver<bool>, shutdown: tokio_util::sync::CancellationToken, ) -> (Channel, tokio::task::JoinHandle<()>) { let (health_reporter, health_service) = health_reporter(); - health_reporter.set_serving::<GenerateServer<GenerateServiceImpl>>().await; + health_reporter.set_serving::<InferenceServer<InferenceServiceImpl>>().await; health_reporter.set_serving::<ControlServer<ControlServiceImpl>>().await; let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener"); @@ -293,7 +299,7 @@ async fn start_grpc_test_server( let server = TonicServer::builder() .add_service(health_service) .add_service(control_service) - .add_service(generate_service) + .add_service(inference_service) .serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned()); let health_monitor = super::monitor_health(health_reporter, engine_health, shutdown.clone()); @@ -323,7 +329,7 @@ async fn grpc_tls_test_server( certs: &TestCerts, cert_reqs: i32, ) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) { - let (generate_service, control_service, _engine_health, engine_task) = + let (inference_service, control_service, _engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await; let context = tls::build_grpc_server_config(&server_tls(certs, cert_reqs)) .expect("build grpc tls config"); @@ -335,7 +341,7 @@ async fn grpc_tls_test_server( let incoming = MaybeTlsListener::tls(Listener::Tcp(listener), context); TonicServer::builder() .add_service(control_service) - .add_service(generate_service) + .add_service(inference_service) .serve_with_incoming(incoming) .await .expect("grpc tls server"); @@ -351,7 +357,7 @@ async fn grpc_tls_client( certs: &TestCerts, addr: &str, identity: Option<&str>, -) -> Result<GenerateClient<Channel>, tonic::transport::Error> { +) -> Result<InferenceClient<Channel>, tonic::transport::Error> { let ca = certs.path("ca.pem"); let identity = identity.map(|name| { ( @@ -388,7 +394,7 @@ async fn grpc_tls_client( .expect("grpc endpoint") .connect_with_connector(connector) .await?; - Ok(GenerateClient::new(channel)) + Ok(InferenceClient::new(channel)) } /// Complete a raw TLS handshake against the gRPC port (offering ALPN `h2`) for @@ -415,7 +421,7 @@ async fn grpc_server_with_keepalive( engine_id: impl Into<EngineId>, keepalive: Option<Duration>, ) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) { - let (generate_service, control_service, _engine_health, engine_task) = + let (inference_service, control_service, _engine_health, engine_task) = setup_grpc_service(engine_id, default_stream_output_specs()).await; let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener"); @@ -432,7 +438,7 @@ async fn grpc_server_with_keepalive( let incoming = MaybeTlsListener::plain(Listener::Tcp(listener)); builder .add_service(control_service) - .add_service(generate_service) + .add_service(inference_service) .serve_with_incoming(incoming) .await .expect("grpc server"); @@ -1083,20 +1089,20 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn control_abort_resolves_external_id_and_empty_is_noop() { - let (generate_service, control_service, engine_health, engine_task) = + let (inference_service, control_service, engine_health, engine_task) = setup_grpc_service(b"engine-grpc-abort-active", vec![(vec![b'h' as u32], None)]).await; let (channel, server_task) = start_grpc_test_server( - generate_service, + inference_service, control_service, engine_health, tokio_util::sync::CancellationToken::new(), ) .await; - let mut generate_client = GenerateClient::new(channel.clone()); + let mut inference_client = InferenceClient::new(channel.clone()); let mut control_client = ControlClient::new(channel); let request_id = "test-abort-active"; - let mut stream = generate_client + let mut stream = inference_client .generate_stream(pb::GenerateRequest { request_id: request_id.to_string(), model: "test-model".to_string(), @@ -1171,14 +1177,127 @@ async fn control_abort_resolves_external_id_and_empty_is_noop() { server_task.abort(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn control_reports_server_and_model_info() { + let (generate_service, control_service, engine_health, _engine_task) = + setup_grpc_service(b"engine-grpc-info", default_stream_output_specs()).await; + let (channel, server_task) = start_grpc_test_server( + generate_service, + control_service, + engine_health, + tokio_util::sync::CancellationToken::new(), + ) + .await; + let mut client = ControlClient::new(channel); + + let server = client + .get_server_info(pb::GetServerInfoRequest {}) + .await + .expect("get server info") + .into_inner(); + assert_eq!(server.engine_version, "test-vllm-version"); + 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); + assert_eq!(server.kv_block_size, DEFAULT_MOCK_BLOCK_SIZE as u32); + assert_eq!(server.total_kv_blocks, DEFAULT_MOCK_NUM_GPU_BLOCKS); + assert_eq!(server.max_running_requests, 256); + assert_eq!(server.max_batched_tokens, 8_192); + let parallelism = server.parallelism.expect("parallelism metadata"); + assert_eq!(parallelism.tensor_parallel_size, 1); + assert_eq!(parallelism.pipeline_parallel_size, 1); + assert_eq!(parallelism.data_parallel_size, 1); + assert_eq!(parallelism.data_parallel_rank, 0); + assert_eq!(parallelism.decode_context_parallel_size, 1); + + let model = client + .get_model_info(pb::GetModelInfoRequest {}) + .await + .expect("get model info") + .into_inner(); + assert_eq!(model.model_id, "test-model"); + assert_eq!(model.served_model_name, "test-model"); + assert!(model.served_model_aliases.is_empty()); + assert!(model.supports_text_input); + assert!(model.supports_token_ids_input); + assert!(!model.supports_multimodal); + assert!(model.reasoning_parser.is_empty()); + assert!(model.tool_call_parser.is_empty()); + + server_task.abort(); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn control_aggregates_multi_engine_capacity() { + let ipc = IpcNamespace::new().expect("create ipc namespace"); + let handshake_address = ipc.handshake_endpoint(); + + let mut ready_0 = default_ready_response(); + ready_0.max_model_len = 8_192; + ready_0.num_gpu_blocks = 10; + ready_0.data_parallel_size = 2; + + let mut ready_1 = default_ready_response(); + ready_1.max_model_len = 4_096; + ready_1.num_gpu_blocks = 20; + ready_1.data_parallel_size = 2; + ready_1.data_parallel_rank = 1; + + let engine_tasks = [ready_0, ready_1].map(|ready| { + let engine_id = EngineId::from_engine_index(ready.data_parallel_rank); + MockEngineTask::new(spawn_mock_engine_task_with_ready( + handshake_address.clone(), + engine_id, + ready, + |_, _| boxed_test_future(async {}), + )) + }); + + let client = EngineCoreClient::connect(EngineCoreClientConfig { + transport_mode: TransportMode::HandshakeOwner { + handshake_address, + advertised_host: "127.0.0.1".to_string(), + engine_count: 2, + ready_timeout: Duration::from_secs(2), + local_input_address: Some(ipc.input_endpoint()), + local_output_address: Some(ipc.output_endpoint()), + }, + coordinator_mode: None, + model_name: "test-model".to_string(), + client_index: 0, + }) + .await + .expect("connect multi-engine client"); + let chat = ChatLlm::from_shared_backend( + Llm::new(client), + Arc::new(FakeTextBackend) as Arc<dyn ChatTextBackend>, + ); + let service = ControlServiceImpl::new(Arc::new(AppState::new( + vec!["test-model".to_string()], + chat, + ))); + + let server = pb::control_server::Control::get_server_info( + &service, + tonic::Request::new(pb::GetServerInfoRequest {}), + ) + .await + .expect("get server info") + .into_inner(); + assert_eq!(server.max_model_len, 4_096); + assert_eq!(server.total_kv_blocks, 30); + + drop(engine_tasks); +} #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy() { - let (generate_service, control_service, _connected_engine_health, _engine_task) = + let (inference_service, control_service, _connected_engine_health, _engine_task) = setup_grpc_service(b"engine-grpc-health-failure", default_stream_output_specs()).await; let (engine_health_tx, engine_health) = tokio::sync::watch::channel(true); let (channel, server_task) = start_grpc_test_server( - generate_service, + inference_service, control_service, engine_health, tokio_util::sync::CancellationToken::new(), @@ -1187,7 +1306,7 @@ async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy() let mut health_client = HealthClient::new(channel); let mut health_streams = Vec::new(); - for service in ["vllm.Generate", "vllm.Control", ""] { + for service in ["vllm.Inference", "vllm.Control", ""] { let service_label = if service.is_empty() { "overall" } else { @@ -1242,14 +1361,14 @@ async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy() #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn grpc_health_watch_closes_on_graceful_shutdown() { - let (generate_service, control_service, engine_health, _engine_task) = setup_grpc_service( + let (inference_service, control_service, engine_health, _engine_task) = setup_grpc_service( b"engine-grpc-health-shutdown", default_stream_output_specs(), ) .await; let shutdown = tokio_util::sync::CancellationToken::new(); let (channel, server_task) = start_grpc_test_server( - generate_service, + inference_service, control_service, engine_health, shutdown.clone(), @@ -1258,43 +1377,43 @@ async fn grpc_health_watch_closes_on_graceful_shutdown() { let mut health_client = HealthClient::new(channel); let mut stream = health_client .watch(HealthCheckRequest { - service: "vllm.Generate".to_string(), + service: "vllm.Inference".to_string(), }) .await - .expect("start health watch for vllm.Generate") + .expect("start health watch for vllm.Inference") .into_inner(); let initial = stream .message() .await - .expect("read initial health status for vllm.Generate") + .expect("read initial health status for vllm.Inference") .expect("health watch ended before its initial status"); assert_eq!( initial.status, HealthServingStatus::Serving as i32, - "unexpected initial health status for vllm.Generate" + "unexpected initial health status for vllm.Inference" ); shutdown.cancel(); let update = tokio::time::timeout(Duration::from_secs(2), stream.message()) .await - .expect("timed out waiting for shutdown health update for vllm.Generate") - .expect("failed to read shutdown health update for vllm.Generate") + .expect("timed out waiting for shutdown health update for vllm.Inference") + .expect("failed to read shutdown health update for vllm.Inference") .expect("health watch ended before its shutdown update"); assert_eq!( update.status, HealthServingStatus::NotServing as i32, - "unexpected shutdown health status for vllm.Generate" + "unexpected shutdown health status for vllm.Inference" ); let stream_end = tokio::time::timeout(Duration::from_secs(2), stream.message()) .await - .expect("timed out waiting for vllm.Generate health watch to close") - .expect("failed while closing vllm.Generate health watch"); + .expect("timed out waiting for vllm.Inference health watch to close") + .expect("failed while closing vllm.Inference health watch"); assert!( stream_end.is_none(), - "vllm.Generate health watch remained open" + "vllm.Inference health watch remained open" ); tokio::time::timeout(Duration::from_secs(2), server_task) diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index 1f3468fe46d..470c424cdb8 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -187,7 +187,7 @@ where let model = state.primary_model_name().to_owned(); let app = extend_router(build_router(state.clone())); - // Optionally bind the gRPC Generate server on a separate port. Bind + // Optionally bind the gRPC Inference server on a separate port. Bind // synchronously here so bind errors (port in use, permission denied, ...) // surface before serving rather than being deferred until shutdown. let grpc_setup = if let Some(grpc_port) = config.grpc_port { @@ -206,19 +206,19 @@ where .context("invalid gRPC TLS configuration")?; let (health_reporter, health_service) = health_reporter(); let engine_health = state.engine_core_client().subscribe_health(); - health_reporter.set_serving::<grpc::GenerateGrpcService>().await; + health_reporter.set_serving::<grpc::InferenceGrpcService>().await; health_reporter.set_serving::<grpc::ControlGrpcService>().await; let control_service = grpc::ControlGrpcService::new(grpc::ControlServiceImpl::new(state.clone())); - let generate_service = - grpc::GenerateGrpcService::new(grpc::GenerateServiceImpl::new(state.clone())); + let inference_service = + grpc::InferenceGrpcService::new(grpc::InferenceServiceImpl::new(state.clone())); let svc = TonicServer::builder() .http2_keepalive_interval(Some(GRPC_KEEPALIVE_INTERVAL)) .http2_keepalive_timeout(Some(GRPC_KEEPALIVE_TIMEOUT)) .layer(middleware::request_runtime_layer(state.clone())) .add_service(health_service) .add_service(control_service) - .add_service(generate_service); + .add_service(inference_service); info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server"); Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health)) } else { diff --git a/rust/src/server/src/middleware/offload.rs b/rust/src/server/src/middleware/offload.rs index edc6560bf20..c310ab442fc 100644 --- a/rust/src/server/src/middleware/offload.rs +++ b/rust/src/server/src/middleware/offload.rs @@ -30,8 +30,8 @@ const OFFLOADED_PATHS: &[&str] = &[ "/detokenize", "/inference/v1/generate", // gRPC routes: - "/vllm.Generate/Generate", - "/vllm.Generate/GenerateStream", + "/vllm.Inference/Generate", + "/vllm.Inference/GenerateStream", ]; /// Return a Tower layer that runs selected data-plane requests on the request runtime, @@ -124,8 +124,8 @@ mod tests { assert!(should_offload("/tokenize")); assert!(should_offload("/detokenize")); assert!(should_offload("/inference/v1/generate")); - assert!(should_offload("/vllm.Generate/Generate")); - assert!(should_offload("/vllm.Generate/GenerateStream")); + assert!(should_offload("/vllm.Inference/Generate")); + assert!(should_offload("/vllm.Inference/GenerateStream")); } #[test] diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index 0bdca7ada99..64adf7a8b3c 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -328,6 +328,13 @@ def test_apply_ready_response_syncs_block_size(): vllm_version="test", world_size=1, data_parallel_size=1, + tensor_parallel_size=1, + pipeline_parallel_size=1, + decode_context_parallel_size=1, + data_parallel_rank=0, + max_num_seqs=256, + max_num_batched_tokens=8192, + instance_id="test-instance", ) ) client._apply_ready_response(payload) diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index 4ac27be5068..e80be0e45d7 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -82,6 +82,13 @@ class EngineCoreReadyResponse: vllm_version: str world_size: int data_parallel_size: int + tensor_parallel_size: int + pipeline_parallel_size: int + decode_context_parallel_size: int + data_parallel_rank: int + max_num_seqs: int + max_num_batched_tokens: int + instance_id: str # KV cache capacity (None for encoder-only/attention-free models). kv_cache_size_tokens: int | None = None kv_cache_max_concurrency: float | None = None diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 39393f64ae4..9817c474343 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -1613,6 +1613,31 @@ class EngineCoreProc(EngineCore): "to send. Please report this issue." ) + def _make_ready_response(self) -> EngineCoreReadyResponse: + parallel_config = self.vllm_config.parallel_config + scheduler_config = self.vllm_config.scheduler_config + return EngineCoreReadyResponse( + max_model_len=self.vllm_config.model_config.max_model_len, + num_gpu_blocks=self.vllm_config.cache_config.num_gpu_blocks or 0, + block_size=self.vllm_config.cache_config.block_size, + dp_stats_address=self.frontend_stats_publish_address, + dtype=str(self.vllm_config.model_config.dtype).removeprefix("torch."), + vllm_version=VLLM_VERSION, + world_size=self.vllm_config.parallel_config.world_size, + data_parallel_size=parallel_config.data_parallel_size, + kv_cache_size_tokens=self.vllm_config.cache_config.kv_cache_size_tokens, + kv_cache_max_concurrency=( + self.vllm_config.cache_config.kv_cache_max_concurrency + ), + tensor_parallel_size=parallel_config.tensor_parallel_size, + pipeline_parallel_size=parallel_config.pipeline_parallel_size, + decode_context_parallel_size=parallel_config.decode_context_parallel_size, + data_parallel_rank=self.engine_index, + max_num_seqs=scheduler_config.max_num_seqs, + max_num_batched_tokens=scheduler_config.max_num_batched_tokens, + instance_id=self.vllm_config.instance_id, + ) + def process_input_sockets( self, input_addresses: list[str], @@ -1654,22 +1679,7 @@ class EngineCoreProc(EngineCore): # Register sockets with poller. poller = zmq.Poller() - ready_response = EngineCoreReadyResponse( - max_model_len=self.vllm_config.model_config.max_model_len, - num_gpu_blocks=self.vllm_config.cache_config.num_gpu_blocks or 0, - block_size=self.vllm_config.cache_config.block_size, - dp_stats_address=self.frontend_stats_publish_address, - dtype=str(self.vllm_config.model_config.dtype).removeprefix("torch."), - vllm_version=VLLM_VERSION, - world_size=self.vllm_config.parallel_config.world_size, - data_parallel_size=self.vllm_config.parallel_config.data_parallel_size, - kv_cache_size_tokens=( - self.vllm_config.cache_config.kv_cache_size_tokens - ), - kv_cache_max_concurrency=( - self.vllm_config.cache_config.kv_cache_max_concurrency - ), - ) + ready_response = self._make_ready_response() ready_payload = msgspec.msgpack.encode(ready_response) for input_socket in input_sockets: # Send initial message to each input socket - this is required From 15d65f86694ac647b139179f88c9d7aea0820108 Mon Sep 17 00:00:00 2001 From: TobyJBell <51262478+TobyB1702@users.noreply.github.com> Date: Mon, 27 Jul 2026 13:06:46 -0400 Subject: [PATCH 121/185] [Bugfix] Changed speech to text chunk timestamp to cumulative approach (#41131) Signed-off-by: Toby Bell <toby.bell1702@hotmail.co.uk> --- docs/contributing/model/transcription.md | 2 +- .../test_speech_to_text_cancellation.py | 12 +++- .../test_chunk_timestamp_offset.py | 70 +++++++++++++++++++ .../test_transcription_inter_chunk_spacing.py | 4 +- .../speech_to_text/base/serving.py | 24 +++++-- 5 files changed, 101 insertions(+), 11 deletions(-) create mode 100644 tests/entrypoints/speech_to_text/transcription/test_chunk_timestamp_offset.py diff --git a/docs/contributing/model/transcription.md b/docs/contributing/model/transcription.md index b076ef84a46..f82bafa44d5 100644 --- a/docs/contributing/model/transcription.md +++ b/docs/contributing/model/transcription.md @@ -195,7 +195,7 @@ Provide a fast duration→token estimate to improve streaming usage statistics: The API server takes care of basic audio I/O and optional chunking before building prompts: - Resampling: Input audio is resampled to `SpeechToTextConfig.sample_rate` using `AudioResampler`. -- Chunking: If `SpeechToTextConfig.allow_audio_chunking` is True and the duration exceeds `max_audio_clip_s`, the server splits the audio into overlapping chunks and generates a prompt per chunk. Overlap is controlled by `overlap_chunk_second`. +- Chunking: If `SpeechToTextConfig.allow_audio_chunking` is True and the duration exceeds `max_audio_clip_s`, the server splits the audio into chunks and generates a prompt per chunk. There is no overlap between chunks, overlap_chunk_second controls the size of the search window used to find the split point. - Energy-aware splitting: When `min_energy_split_window_size` is set, the server finds low-energy regions to minimize cutting within words. Relevant server logic: diff --git a/tests/entrypoints/speech_to_text/test_speech_to_text_cancellation.py b/tests/entrypoints/speech_to_text/test_speech_to_text_cancellation.py index 040fc1a48ff..45772708d92 100644 --- a/tests/entrypoints/speech_to_text/test_speech_to_text_cancellation.py +++ b/tests/entrypoints/speech_to_text/test_speech_to_text_cancellation.py @@ -53,7 +53,13 @@ async def test_non_streaming_cancel_aborts_engine_requests( server.asr_config = SimpleNamespace(max_audio_clip_s=30) server._check_model = AsyncMock(return_value=None) server._maybe_get_adapters = Mock(return_value=None) - server._preprocess_speech_to_text = AsyncMock(return_value=(engine_inputs, 40.0)) + server._preprocess_speech_to_text = AsyncMock( + return_value=( + engine_inputs, + 40.0, + [30.0 * i for i in range(len(engine_inputs))], + ) + ) server._log_inputs = Mock() request = SimpleNamespace( @@ -122,7 +128,9 @@ async def test_non_streaming_cancel_advances_all_chunk_generators(): server.asr_config = SimpleNamespace(max_audio_clip_s=30) server._check_model = AsyncMock(return_value=None) server._maybe_get_adapters = Mock(return_value=None) - server._preprocess_speech_to_text = AsyncMock(return_value=(engine_inputs, 90.0)) + server._preprocess_speech_to_text = AsyncMock( + return_value=(engine_inputs, 90.0, [0.0, 29.5, 29.5 + 29.7]) + ) server._log_inputs = Mock() request = SimpleNamespace( diff --git a/tests/entrypoints/speech_to_text/transcription/test_chunk_timestamp_offset.py b/tests/entrypoints/speech_to_text/transcription/test_chunk_timestamp_offset.py new file mode 100644 index 00000000000..22231e0d85c --- /dev/null +++ b/tests/entrypoints/speech_to_text/transcription/test_chunk_timestamp_offset.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression test for chunk timestamp offset drift in _preprocess_speech_to_text.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import numpy as np +import pytest + +from vllm.config.speech_to_text import SpeechToTextConfig +from vllm.entrypoints.speech_to_text.base.serving import SpeechToTextBaseServing + +SR = 16_000 +_PATCH = "vllm.entrypoints.speech_to_text.base.serving" + + +@pytest.mark.asyncio +async def test_chunk_offsets_are_cumulative_not_nominal(): + """ + chunk_start_offsets must be cumulative actual chunk lengths, not the old approach of + 'idx * max_audio_clip_s'. When split_audio places a boundary before the + nominal 30 s mark, the old formula drifts; the fixed formula stays exact. + """ + # Chunks shorter than exactly 30 s, as split_audio produces when a quiet + # region falls inside the 1 s overlap window before the nominal boundary. + chunk_lengths = [int(29.5 * SR), int(29.7 * SR), int(5.0 * SR)] + chunks = [np.zeros(n, dtype=np.float32) for n in chunk_lengths] + + duration = sum(chunk_lengths) / SR + + expected_offsets = [0.0, 29.5, 29.5 + 29.7] # cumulative seconds + wrong_offsets = [0.0, 30.0, 60.0] # what the old bug produced + + serving = SpeechToTextBaseServing.__new__(SpeechToTextBaseServing) + serving._decode_and_chunk_speech_async = AsyncMock(return_value=(chunks, duration)) + serving.asr_config = SpeechToTextConfig( + sample_rate=float(SR), + max_audio_clip_s=30, + overlap_chunk_second=1, + min_energy_split_window_size=1600, + ) + serving.max_audio_filesize_mb = 100.0 + serving.model_cls = MagicMock() + serving.model_cls.validate_language.side_effect = lambda lang: lang + serving.model_cls.supports_explicit_language_detection = False + serving.model_cls.get_generation_prompt.return_value = {} + serving.model_config = MagicMock() + serving.task_type = "transcribe" + serving.renderer = MagicMock() + serving.renderer.render_cmpl_async = AsyncMock( + return_value=[MagicMock()] * len(chunks) + ) + + request = MagicMock() + request.language = "en" + request.to_language = None + request.response_format = "json" + request.build_stt_params.return_value = MagicMock() + + with patch(f"{_PATCH}.parse_model_prompt", return_value=MagicMock()): + _, _, offsets = await serving._preprocess_speech_to_text( + request=request, + audio_data=b"\x00", + request_id="test", + ) + + assert offsets == pytest.approx(expected_offsets, abs=1e-6) + assert offsets != pytest.approx(wrong_offsets, abs=1e-6) diff --git a/tests/entrypoints/speech_to_text/transcription/test_transcription_inter_chunk_spacing.py b/tests/entrypoints/speech_to_text/transcription/test_transcription_inter_chunk_spacing.py index 7e51e5be44a..8c99429af1e 100644 --- a/tests/entrypoints/speech_to_text/transcription/test_transcription_inter_chunk_spacing.py +++ b/tests/entrypoints/speech_to_text/transcription/test_transcription_inter_chunk_spacing.py @@ -342,7 +342,9 @@ async def test_create_transcription_non_streaming_joins_chunks_by_language(): models.lora_requests = {} models.is_base_model.return_value = True - preprocess_mock = AsyncMock(return_value=([MagicMock(), MagicMock()], 1.0)) + preprocess_mock = AsyncMock( + return_value=([MagicMock(), MagicMock()], 1.0, [0.0, 29.5]) + ) with ( patch( diff --git a/vllm/entrypoints/speech_to_text/base/serving.py b/vllm/entrypoints/speech_to_text/base/serving.py index 7c703f2535c..c646f533837 100644 --- a/vllm/entrypoints/speech_to_text/base/serving.py +++ b/vllm/entrypoints/speech_to_text/base/serving.py @@ -258,7 +258,7 @@ class SpeechToTextBaseServing(GenerateBaseServing): request: SpeechToTextRequest, audio_data: bytes, request_id: str, - ) -> tuple[list[EngineInput], float]: + ) -> tuple[list[EngineInput], float, list[float]]: # Validate request request.language = self.model_cls.validate_language(request.language) request.to_language = ( @@ -277,6 +277,13 @@ class SpeechToTextBaseServing(GenerateBaseServing): # Run cpu intensive preprocess step in a separate thread pool executor. chunks, duration = await self._decode_and_chunk_speech_async(audio_data) + chunk_start_offsets: list[float] = [0.0] + + for chunk in chunks[:-1]: + chunk_start_offsets.append( + chunk_start_offsets[-1] + chunk.shape[-1] / self.asr_config.sample_rate + ) + if request.language is None and getattr( self.model_cls, "supports_explicit_language_detection", False ): @@ -306,7 +313,7 @@ class SpeechToTextBaseServing(GenerateBaseServing): engine_inputs = await self.renderer.render_cmpl_async(parsed_prompts) - return engine_inputs, duration + return engine_inputs, duration, chunk_start_offsets def _preprocess_verbose_prompt(self, prompt: EncoderDecoderDictPrompt): dec_prompt = prompt["decoder_prompt"] @@ -391,7 +398,7 @@ class SpeechToTextBaseServing(GenerateBaseServing): SpeechToTextSegment, segment_class( id=len(segments), - seek=start_time, + seek=int(start_time), start=start_time + BASE_OFFSET * start_timestamp, end=start_time + BASE_OFFSET * end_timestamp, temperature=request.temperature, @@ -467,7 +474,11 @@ class SpeechToTextBaseServing(GenerateBaseServing): lora_request = self._maybe_get_adapters(request) - engine_inputs, duration_s = await self._preprocess_speech_to_text( + ( + engine_inputs, + duration_s, + chunk_start_offsets, + ) = await self._preprocess_speech_to_text( request=request, audio_data=audio_data, request_id=request_id, @@ -587,11 +598,10 @@ class SpeechToTextBaseServing(GenerateBaseServing): assert len(list_result_generator) == 1, ( "`max_audio_clip_s` is set to None, audio cannot be chunked" ) + assert len(chunk_start_offsets) == len(list_result_generator) result_generator = merge_async_iterators(*list_result_generator) async for idx, op in result_generator: - start_time = ( - float(idx * chunk_size_in_s) if chunk_size_in_s is not None else 0.0 - ) + start_time = chunk_start_offsets[idx] if request.response_format == "verbose_json": assert op.outputs[0].logprobs segments: list[SpeechToTextSegment] = self._get_verbose_segments( From 8112b6c9972067012e08ce4f0dbf1d1d9906dbe3 Mon Sep 17 00:00:00 2001 From: Nick Hill <nickhill123@gmail.com> Date: Mon, 27 Jul 2026 10:25:17 -0700 Subject: [PATCH 122/185] [MRV2] Always build attn metadata at capture time (#49364) (#49995) Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai> --- vllm/v1/worker/gpu/cudagraph_utils.py | 51 ++++++++++++------- .../autoregressive/cudagraph_utils.py | 5 +- 2 files changed, 35 insertions(+), 21 deletions(-) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 26f331fe843..3a174cba80d 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -498,10 +498,7 @@ class ModelCudaGraphManager(CudaGraphManager): block_tables, attn_groups, kv_cache_config, - skip_attn=( - desc.cg_mode == CUDAGraphMode.PIECEWISE - and not self.use_breakable_cg - ), + full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL, ) # Capture with dummy rows marked as padding. @@ -510,7 +507,6 @@ class ModelCudaGraphManager(CudaGraphManager): def forward_fn(cg_mode: CUDAGraphMode) -> None: batch_descriptor = None if cg_mode == CUDAGraphMode.PIECEWISE: - assert (attn_metadata is not None) == self.use_breakable_cg batch_descriptor = BatchDescriptor( num_tokens=num_tokens, has_lora=has_lora, @@ -593,7 +589,7 @@ def prepare_inputs_to_capture( block_tables: BlockTables, attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, - skip_attn: bool = False, + full_cudagraph: bool, ) -> AttentionState: input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers) input_block_tables = block_tables.get_dummy_block_tables(num_reqs) @@ -614,15 +610,36 @@ def prepare_inputs_to_capture( ) input_batch.dcp_local_seq_lens = input_buffers.dcp_local_seq_lens[:num_reqs] - attn_metadata = None - if not skip_attn: - attn_metadata = model_state.prepare_attn( - input_batch, - CUDAGraphMode.NONE, - input_block_tables, - slot_mappings, - attn_groups, - kv_cache_config, - for_capture=True, - ) + # NOTE(woosuk): Attention metadata is required not just by standard attention + # kernels, but also by specialized attention-like operations (e.g., Inkling's sconv, + # DSV4 compressor), which maintain their own states and require special metadata + # such as block tables. + # During CUDA graph capture: + # - For FULL CUDA graphs: We set for_capture=True so that both attention and + # attention-like ops produce capturable metadata compatible with CUDA graphs. + # - For PIECEWISE CUDA graphs: We still build attention metadata, but set + # for_capture=False. This is because: + # * Attention-like ops (such as sconv or DSV4 compressor) may not be used as + # breakpoints in PIECEWISE CUDA graphs, so we must generate their attention + # metadata so they can execute and be captured during graph capture. + # * Standard attention ops that are treated as breakpoints will be executed + # eagerly at capture time (not included in the graph itself), and for these, + # setting for_capture=False is essential. Some attention backends + # (like linear attention) cannot generate capturable metadata for prefill, + # so for_capture=False ensures they execute without issue. + # * We assume that attention-like operations intended for capture will still + # produce capturable metadata, even when for_capture=False. While this + # assumption is brittle, it currently works in practice. + # In summary: We always generate attention metadata for both FULL and PIECEWISE + # CUDA graphs, setting for_capture=True for FULL graphs, and for_capture=False + # for PIECEWISE graphs, to ensure correct execution and capture. + attn_metadata = model_state.prepare_attn( + input_batch, + CUDAGraphMode.NONE, + input_block_tables, + slot_mappings, + attn_groups, + kv_cache_config, + for_capture=full_cudagraph, + ) return AttentionState(attn_metadata, slot_mappings_by_layer) diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py index 19919043c83..ef3b6e2ed53 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py @@ -56,10 +56,7 @@ class SpeculatorCudaGraphManager(CudaGraphManager): block_tables, attn_groups, kv_cache_config, - skip_attn=( - desc.cg_mode == CUDAGraphMode.PIECEWISE - and not self.use_breakable_cg - ), + full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL, ) return lambda cg_mode: forward_fn( From bf2b45b5d6a991f47702d4d87c81bb0ef8619a65 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni <mbonanni@redhat.com> Date: Mon, 27 Jul 2026 14:25:50 -0400 Subject: [PATCH 123/185] [Attention] Integrate FlashAttention 4 SM100 headdim 256 support (#42669) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- benchmarks/attention_benchmarks/benchmark.py | 4 ++ docs/design/attention_backends.md | 4 +- .../v1/attention/test_mla_prefill_selector.py | 14 +++---- vllm/v1/attention/backends/fa_utils.py | 22 ++++++++--- vllm/v1/attention/backends/flash_attn.py | 1 + .../backends/mla/prefill/flash_attn.py | 37 +++++++++---------- .../backends/mla/prefill/selector.py | 18 ++++++++- 7 files changed, 66 insertions(+), 34 deletions(-) diff --git a/benchmarks/attention_benchmarks/benchmark.py b/benchmarks/attention_benchmarks/benchmark.py index 2fdd5ce88dc..a4941b4a18f 100644 --- a/benchmarks/attention_benchmarks/benchmark.py +++ b/benchmarks/attention_benchmarks/benchmark.py @@ -1358,6 +1358,10 @@ def main(): profile_memory=args.profile_memory, warmup_ms=args.warmup_ms, prefill_backend=pb, + kv_lora_rank=args.kv_lora_rank, + qk_nope_head_dim=args.qk_nope_head_dim, + qk_rope_head_dim=args.qk_rope_head_dim, + v_head_dim=args.v_head_dim, ) result = run_benchmark(config) diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index 755ce8961e2..c4cedd9e4e1 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -153,7 +153,9 @@ hardware and configuration. > **‡** Automatic selection tries FlashAttention first. On Blackwell > (SM100), the fallback order is TRT-LLM Ragged, FlashInfer, then -> TokenSpeed MLA. On other GPUs, only FlashAttention is considered. +> TokenSpeed MLA; for (qk_nope_head_dim=192, qk_rope_head_dim=64, +> v_head_dim=256) TRT-LLM Ragged is tried before FlashAttention. On other +> GPUs, only FlashAttention is considered. ### Decode Backends diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py index d82397591f7..f5985e7bc8e 100644 --- a/tests/v1/attention/test_mla_prefill_selector.py +++ b/tests/v1/attention/test_mla_prefill_selector.py @@ -163,7 +163,7 @@ class TestGetMLAPrefillBackend: class TestAutoSelectMLAPrefillBackend: """Tests for fallback and error paths in auto-selection.""" - def test_blackwell_glm_dimensions_fall_back_to_trtllm(self): + def test_blackwell_glm_dimensions_use_trtllm(self): capability = DeviceCapability(major=10, minor=0) selector_config = MLAPrefillSelectorConfig( dtype=torch.bfloat16, @@ -184,11 +184,6 @@ class TestAutoSelectMLAPrefillBackend: with ( patch("vllm.platforms.current_platform") as mock_platform, patch.object(flash_attn_cls, "is_available", return_value=True), - patch( - "vllm.v1.attention.backends.mla.prefill.flash_attn." - "get_flash_attn_version", - return_value=4, - ), patch.object(trtllm_cls, "validate_configuration", return_value=[]), ): # Force the non-ROCm priority on the Blackwell. @@ -300,7 +295,12 @@ class TestROCmAiterFAPrefillSelection: with patch("vllm.platforms.current_platform") as mock_platform: mock_platform.is_rocm.return_value = True priorities = _get_mla_prefill_backend_priorities( - DeviceCapability(major=9, minor=5) + DeviceCapability(major=9, minor=5), + MLADimensions( + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + ), ) assert priorities == [ diff --git a/vllm/v1/attention/backends/fa_utils.py b/vllm/v1/attention/backends/fa_utils.py index 718580dcf69..9299af09a0a 100644 --- a/vllm/v1/attention/backends/fa_utils.py +++ b/vllm/v1/attention/backends/fa_utils.py @@ -134,6 +134,7 @@ def get_flash_attn_version( head_size: int | None = None, head_size_v: int | None = None, has_sinks: bool = False, + requires_local_attention: bool = False, ) -> int | None: if current_platform.is_xpu(): return 2 @@ -230,17 +231,28 @@ def get_flash_attn_version( ) fa_version = 2 + if ( + fa_version == 4 + and device_capability.major >= 10 + and head_size == 256 + and requires_local_attention + ): + logger.warning_once( + "FA4 on Blackwell does not support local attention with " + "head_size=256, defaulting to FA version 2." + ) + fa_version = 2 + # FA4 on SM100 (Blackwell) has TMEM capacity limits that restrict - # supported head dimensions. - # See: https://github.com/Dao-AILab/flash-attention/issues/1959 - # Exception: hdim 192 is supported for MLA's diff-headdim case - # (qk=192, v=128), added upstream in commits 1a15733e/1b36ab19. + # supported head dimensions to ≤128, with exceptions for 256 and 192/128 (MLA + # prefill). Development of symmetric 192, 384, and 512 support is being tracked + # in https://github.com/Dao-AILab/flash-attention/issues/2456 if ( fa_version == 4 and device_capability.major >= 10 and head_size is not None and head_size > 128 - and head_size != 192 + and not (head_size == 256 or (head_size == 192 and head_size_v == 128)) ): logger.warning_once( "FA4 on Blackwell does not support head_size=%d due to TMEM " diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 13bf435a542..416377deabc 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -782,6 +782,7 @@ class FlashAttentionImpl(AttentionImpl): self.attn_type = attn_type self.vllm_flash_attn_version = get_flash_attn_version( requires_alibi=alibi_slopes is not None, + requires_local_attention=sliding_window is not None, head_size=head_size, has_sinks=sinks is not None, ) diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py index c5fdbaebc02..d073f9b8283 100644 --- a/vllm/v1/attention/backends/mla/prefill/flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py @@ -288,26 +288,23 @@ class FlashAttnPrefillBackend(MLAPrefillBackend): @classmethod def supports_mla_dimensions(cls, mla_dimensions: MLADimensions) -> bool: - dims_deepseek = MLADimensions( - qk_nope_head_dim=128, - qk_rope_head_dim=64, - v_head_dim=128, - ) - dims_glm = MLADimensions( - qk_nope_head_dim=192, - qk_rope_head_dim=64, - v_head_dim=256, - ) - dims_mistral_s4 = MLADimensions( - qk_nope_head_dim=64, - qk_rope_head_dim=64, - v_head_dim=128, - ) - fa_version = get_flash_attn_version() - if fa_version == 4: - return mla_dimensions in [dims_deepseek, dims_mistral_s4] - else: - return mla_dimensions in [dims_deepseek, dims_glm, dims_mistral_s4] + return mla_dimensions in [ + MLADimensions( + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + ), + MLADimensions( + qk_nope_head_dim=192, + qk_rope_head_dim=64, + v_head_dim=256, + ), + MLADimensions( + qk_nope_head_dim=64, + qk_rope_head_dim=64, + v_head_dim=128, + ), + ] def __init__( self, diff --git a/vllm/v1/attention/backends/mla/prefill/selector.py b/vllm/v1/attention/backends/mla/prefill/selector.py index e0d54eee101..c313c901772 100644 --- a/vllm/v1/attention/backends/mla/prefill/selector.py +++ b/vllm/v1/attention/backends/mla/prefill/selector.py @@ -47,11 +47,13 @@ class MLAPrefillSelectorConfig(NamedTuple): def _get_mla_prefill_backend_priorities( device_capability: DeviceCapability, + mla_dimensions: MLADimensions, ) -> list[MLAPrefillBackendEnum]: """Get MLA prefill backend priorities based on device capability. Args: device_capability: The device's compute capability. + mla_dimensions: The model's MLA head dimensions. Returns: List of backends in priority order (highest priority first). @@ -65,6 +67,17 @@ def _get_mla_prefill_backend_priorities( ] if device_capability.major == 10: # Blackwell + if mla_dimensions == MLADimensions( + qk_nope_head_dim=192, + qk_rope_head_dim=64, + v_head_dim=256, + ): + return [ + MLAPrefillBackendEnum.TRTLLM_RAGGED, + MLAPrefillBackendEnum.FLASH_ATTN, + MLAPrefillBackendEnum.FLASHINFER, + MLAPrefillBackendEnum.TOKENSPEED_MLA, + ] return [ MLAPrefillBackendEnum.FLASH_ATTN, MLAPrefillBackendEnum.TRTLLM_RAGGED, @@ -157,7 +170,10 @@ def _auto_select_mla_prefill_backend( Returns: The selected prefill backend class. """ - priorities = _get_mla_prefill_backend_priorities(device_capability) + priorities = _get_mla_prefill_backend_priorities( + device_capability, + selector_config.mla_dimensions, + ) all_invalid_reasons: dict[str, list[str]] = {} for backend_enum in priorities: From 99de48e98fe9570d52b733a545d27661929c170c Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Mon, 27 Jul 2026 19:32:32 +0100 Subject: [PATCH 124/185] Fix MLA padding and grouped topk routing in the Transformers modelling backend (#49982) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> --- tests/models/transformers/fusers/test_moe.py | 83 +++++++++++++++++- .../models/transformers/__init__.py | 13 ++- .../models/transformers/fusers/moe.py | 77 ++++++++++++----- .../model_executor/models/transformers/moe.py | 86 ++++++++++++++++--- .../models/transformers/utils.py | 7 ++ 5 files changed, 227 insertions(+), 39 deletions(-) diff --git a/tests/models/transformers/fusers/test_moe.py b/tests/models/transformers/fusers/test_moe.py index 04eadac3f78..7de4589eef9 100644 --- a/tests/models/transformers/fusers/test_moe.py +++ b/tests/models/transformers/fusers/test_moe.py @@ -29,8 +29,48 @@ class TopKRouter(nn.Module): return logits, value, index +class ScaledRouter(TopKRouter): + """Greedy router scaling its top-k weights (DeepSeek `routed_scaling_factor`).""" + + def forward(self, hidden_states): + logits = F.linear(hidden_states, self.weight) + scores = F.softmax(logits, dim=-1) + value, index = torch.topk(scores, self.top_k, dim=-1) + value = value * 16.0 + return logits, value, index + + +class Fp32Router(TopKRouter): + """Router that computes its logits in fp32 (DeepSeek/GLM style).""" + + def forward(self, hidden_states): + logits = F.linear( + hidden_states.type(torch.float32), self.weight.type(torch.float32) + ) + scores = F.softmax(logits, dim=-1) + value, index = torch.topk(scores, self.top_k, dim=-1) + return logits, value.to(hidden_states.dtype), index + + +class GroupedRouter(TopKRouter): + """Group-limited router (DeepSeek `group_limited_greedy`), scaled weights.""" + + def forward(self, hidden_states): + logits = F.linear(hidden_states, self.weight) + scores = F.softmax(logits, dim=-1) + group_scores = scores.view(-1, 4, 2).max(dim=-1).values + group_idx = torch.topk(group_scores, k=2, dim=-1)[1] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_idx, 1) + score_mask = group_mask.unsqueeze(-1).expand(-1, 4, 2).reshape(-1, 8) + scores = scores.masked_fill(~score_mask.bool(), 0.0) + value, index = torch.topk(scores, self.top_k, dim=-1) + value = value * 16.0 + return logits, value, index + + class CorrectionRouter(nn.Module): - """Grouped router with a score-correction bias buffer (DeepSeek-V3) -> declined.""" + """Router with a score-correction bias buffer (DeepSeek-V3 noaux) -> matched.""" def __init__(self, num_experts=8, hidden=16): super().__init__() @@ -228,6 +268,44 @@ def test_moe_fuser_detects_router(sigmoid): assert fuser.shared_name is None and fuser.shared_gate_name is None +def test_moe_fuser_matches_scaled_router(): + """Weight scaling after the top-k (DeepSeek style) does not break matching.""" + with torch.device("meta"): + block = MoEBlock(ScaledRouter) + assert isinstance(MoEBlockFuser.match(block, "experts"), MoEBlockFuser) + + +def test_moe_fuser_matches_grouped_router(): + """Group masking between the score and the routing top-k still matches: the + scorer is anchored on the routing (last) top-k, not the group one.""" + with torch.device("meta"): + block = MoEBlock(GroupedRouter) + fuser = MoEBlockFuser.match(block, "experts") + assert isinstance(fuser, MoEBlockFuser) + assert fuser.scoring_func == "softmax" + + +def test_moe_fuser_matches_correction_router(): + """A score-correction bias buffer (DeepSeek-V3 noaux) is allowed; the + rebuilt gate carries it in fp32 for FusedMoE's biased routers.""" + with torch.device("meta"): + block = MoEBlock(CorrectionRouter) + fuser = MoEBlockFuser.match(block, "experts") + assert isinstance(fuser, MoEBlockFuser) + assert fuser.scoring_func == "sigmoid" + + +def test_moe_fuser_reads_router_dtype_from_the_gate(): + """A router that computes its logits in fp32 must keep routing in fp32 when + rebuilt, even though no config field names the dtype. The cast back to the + activation dtype after the top-k must not be mistaken for the routing dtype.""" + with torch.device("meta"): + fp32 = MoEBlockFuser.match(MoEBlock(Fp32Router), "experts") + default = MoEBlockFuser.match(MoEBlock(TopKRouter), "experts") + assert fp32.router_dtype == torch.float32 + assert default.router_dtype is None + + def test_moe_fuser_detects_shared_experts(): with torch.device("meta"): block = MoEBlockShared() @@ -269,8 +347,7 @@ def test_moe_fuser_detects_non_glu_shared_expert(): @pytest.mark.parametrize( "block_cls", [ - lambda: MoEBlock(CorrectionRouter), # score-correction buffer (grouped) - lambda: MoEBlock(BiasedRouter), # router not weight-only (extra param) + lambda: MoEBlock(BiasedRouter), # router with an unrecognized extra param MoEBlockTuple, # tuple-returning block (e.g. gpt-oss) MoEBlockTupleVar, # tuple returned via a name binding, not a literal MoEBlockUnaccounted, # weight-bearing child outside the fused dataflow diff --git a/vllm/model_executor/models/transformers/__init__.py b/vllm/model_executor/models/transformers/__init__.py index 78a12876e66..9dbfe4b5031 100644 --- a/vllm/model_executor/models/transformers/__init__.py +++ b/vllm/model_executor/models/transformers/__init__.py @@ -18,6 +18,7 @@ from typing import TYPE_CHECKING +import torch.nn.functional as F from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS from vllm.model_executor.models.transformers.base import Base @@ -59,9 +60,19 @@ def vllm_attention_forward( if scaling is not None: self_attn.impl.scale = float(scaling) hidden = query.shape[-2] + head_dim_qk = query.shape[-1] + head_dim_v = value.shape[-1] query, key, value = (x.transpose(1, 2) for x in (query, key, value)) query, key, value = (x.reshape(hidden, -1) for x in (query, key, value)) - return self_attn.forward(query, key, value), None + # Pad `value` up to the query/key head size when they differ (expanded MLA). + if head_dim_v != head_dim_qk: + value = F.pad(value.view(-1, head_dim_v), (0, head_dim_qk - head_dim_v)) + value = value.reshape(hidden, -1) + attn_output = self_attn.forward(query, key, value) + if head_dim_v != head_dim_qk: + attn_output = attn_output.view(-1, head_dim_qk)[..., :head_dim_v] + attn_output = attn_output.reshape(hidden, -1) + return attn_output, None ALL_ATTENTION_FUNCTIONS["vllm"] = vllm_attention_forward diff --git a/vllm/model_executor/models/transformers/fusers/moe.py b/vllm/model_executor/models/transformers/fusers/moe.py index 6a3c7e85d6e..af3a101a0ce 100644 --- a/vllm/model_executor/models/transformers/fusers/moe.py +++ b/vllm/model_executor/models/transformers/fusers/moe.py @@ -6,14 +6,14 @@ import ast import inspect import textwrap import types -from collections.abc import Iterator +from collections.abc import Iterable, Iterator from dataclasses import dataclass -from itertools import chain import torch from torch import fx, nn from vllm.distributed import tensor_model_parallel_all_gather +from vllm.model_executor.layers.fused_moe import GateLinear from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.models.transformers.fx_utils import ( find_node, @@ -21,14 +21,10 @@ from vllm.model_executor.models.transformers.fx_utils import ( peel, trace, ) +from vllm.model_executor.models.transformers.utils import named_state from vllm.model_executor.models.utils import maybe_prefix, sequence_parallel_chunk -def named_state(module: nn.Module) -> Iterator[tuple[str, torch.Tensor]]: - """`module`'s own state (i.e. named parameters and buffers).""" - return chain(module.named_parameters(), module.named_buffers()) - - def _own_returns(node: ast.AST) -> Iterator[ast.Return]: """`return` statements in `node`'s own scope, not in nested functions.""" stack = list(ast.iter_child_nodes(node)) @@ -79,6 +75,25 @@ def _is_scalar_gate(module: nn.Module) -> bool: ) +def _forced_dtype(nodes: Iterable[fx.Node]) -> torch.dtype | None: + """The floating dtype `nodes` cast to, if any. + + Computations that must run in higher precision say so in their forward (e.g. + `hidden_states.type(torch.float32)`), so the dtype is readable from the graph + even when the config does not name it.""" + for node in nodes: + if node.op not in ("call_method", "call_function"): + continue + name = str(node.target).rsplit(".", 1)[-1] + if name == "float": + return torch.float32 + if name in ("to", "type"): + for arg in (*node.args[1:], *node.kwargs.values()): + if isinstance(arg, torch.dtype) and arg.is_floating_point: + return arg + return None + + def _reaches(node: fx.Node, key: str) -> set[fx.Node]: """Returns the set of nodes reachable from `node` by following `key` edges.""" seen: set[fx.Node] = set() @@ -131,18 +146,24 @@ class MoEBlockFuser: scoring_func: str shared_name: str | None shared_gate_name: str | None + router_dtype: torch.dtype | None = None @staticmethod - def _match_router(gate: nn.Module) -> str | None: - """Matches `topk(score(linear(x)))`, `score` being `softmax`/`sigmoid`.""" - if [name for name, _ in named_state(gate)] != ["weight"]: + def _match_router(gate: nn.Module) -> tuple[str, torch.dtype | None] | None: + """Matches `topk(score(linear(x)))`, `score` being `softmax`/`sigmoid`. + + Returns the scoring function and the dtype the router computes in.""" + state = {name for name, _ in named_state(gate)} + if "weight" not in state or state - {"weight", "e_score_correction_bias"}: return None graph = trace(gate) if graph is None: return None - topk = find_node(graph, lambda n: is_op(n, "topk")) - if topk is None: + # The routing top-k is the last one; any earlier one scores expert groups. + topks = [node for node in graph.nodes if is_op(node, "topk")] + if not topks: return None + topk = topks[-1] # Exactly one scoring op upstream of the top-k, fed (transitively) by a linear. scorers = [ n @@ -152,9 +173,11 @@ class MoEBlockFuser: if len(scorers) != 1: return None scorer = scorers[0] - if not any(is_op(n, "linear") for n in _reaches(scorer, "all_input_nodes")): + logits_cone = _reaches(scorer, "all_input_nodes") + if not any(is_op(n, "linear") for n in logits_cone): return None - return "softmax" if is_op(scorer, "softmax") else "sigmoid" + scoring_func = "softmax" if is_op(scorer, "softmax") else "sigmoid" + return scoring_func, _forced_dtype(logits_cone) @staticmethod def _match_shared_experts( @@ -198,10 +221,14 @@ class MoEBlockFuser: if _returns_tuple(type(moe_block)): return None # Router: the child that scores + top-k selects. - gate_name = scoring_func = None + gate_name = scoring_func = router_dtype = None for name, child in moe_block.named_children(): - if name != experts_name and (func := cls._match_router(child)) is not None: - gate_name, scoring_func = name, func + if ( + name != experts_name + and (router := cls._match_router(child)) is not None + ): + gate_name = name + scoring_func, router_dtype = router break if gate_name is None or scoring_func is None: return None @@ -229,17 +256,23 @@ class MoEBlockFuser: for name, child in moe_block.named_children(): if name not in accounted and next(named_state(child), None) is not None: return None - return cls(gate_name, scoring_func, shared_name, shared_gate_name) + return cls(gate_name, scoring_func, shared_name, shared_gate_name, router_dtype) - def gate(self, moe_block: nn.Module, prefix: str) -> ReplicatedLinear: - """Rebuild the HF gate as a `ReplicatedLinear` for vLLM's fused MoE.""" - num_experts, hidden_size = getattr(moe_block, self.gate_name).weight.shape - gate = ReplicatedLinear( + def gate( + self, moe_block: nn.Module, prefix: str, out_dtype: torch.dtype | None = None + ) -> GateLinear: + """Rebuild the HF gate as a `GateLinear` for vLLM's fused MoE.""" + hf_gate = getattr(moe_block, self.gate_name) + num_experts, hidden_size = hf_gate.weight.shape + gate = GateLinear( hidden_size, num_experts, bias=False, + out_dtype=out_dtype or self.router_dtype, prefix=maybe_prefix(prefix, self.gate_name), ) + if (bias := getattr(hf_gate, "e_score_correction_bias", None)) is not None: + gate.register_buffer("e_score_correction_bias", bias) setattr(moe_block, self.gate_name, gate) return gate diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index d1f5dba0373..ce964036973 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/vllm/model_executor/models/transformers/moe.py @@ -23,6 +23,7 @@ from typing import TYPE_CHECKING, Any import torch import torch.nn as nn +from vllm._aiter_ops import rocm_aiter_ops from vllm.config.utils import getattr_iter from vllm.distributed import get_dp_group, get_ep_group from vllm.forward_context import ForwardContext, get_forward_context @@ -30,9 +31,10 @@ from vllm.logger import init_logger from vllm.model_executor.custom_op import PluggableLayer from vllm.model_executor.layers.fused_moe import FusedMoE, MoERunner, RoutedExperts from vllm.model_executor.models.interfaces import MixtureOfExperts +from vllm.model_executor.models.transformers.fuser import get_fuser from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser from vllm.model_executor.models.utils import maybe_prefix -from vllm.utils.torch_utils import direct_register_custom_op +from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, direct_register_custom_op from .utils import log_replacement @@ -176,17 +178,31 @@ class MoEMixin(MixtureOfExperts): 0, ) - # Unused kwargs since we use custom_routing_function: - # - `scoring_func` and `e_score_correction_bias` only used for grouped - # topk routing inside vLLM and are non-trivial to infer - # and hard code `use_grouped_topk=False` - # - `renormalize` passed anyway because it's easy to infer - # - `num_expert_group` and `topk_group` used for inferring expert - # placement strategy in FusedMoE - # - `apply_router_weight_on_input` is already applied in Transformers + # Common kwargs renormalize = getattr(text_config, "norm_topk_prob", top_k > 1) + + # Routed scaling factor kwargs + routed_scaling_factor = getattr(text_config, "routed_scaling_factor", 1.0) + # aiter applies routed_scaling_factor internally + apply_routed_scale_to_output = not rocm_aiter_ops.is_fused_moe_enabled() + routed_scaling_factor_kwargs = dict( + routed_scaling_factor=routed_scaling_factor, + apply_routed_scale_to_output=apply_routed_scale_to_output, + ) + + # Dtype the router computes in, if it is not the activation dtype. + config_router_dtype = getattr(text_config, "moe_router_dtype", None) + config_router_dtype = STR_DTYPE_TO_TORCH_DTYPE.get(config_router_dtype) + + # Grouped topk routing kwargs num_expert_group = getattr(text_config, "n_group", None) topk_group = getattr(text_config, "topk_group", None) + use_grouped_topk = num_expert_group is not None and topk_group is not None + grouped_topk_routing_kwargs = dict( + use_grouped_topk=use_grouped_topk, + num_expert_group=num_expert_group, + topk_group=topk_group, + ) # MoE activation function activation = "silu" @@ -211,6 +227,9 @@ class MoEMixin(MixtureOfExperts): self.num_shared_experts = num_shared_experts self.num_redundant_experts = num_redundant_experts + # Down projections of shared experts consumed by FusedMoE + shared_down_projs: list[tuple[nn.Module, str]] = [] + # Recursively fuse MoE layers def _recursive_replace(module: nn.Module, prefix: str): for child_name, child_module in module.named_children(): @@ -249,7 +268,6 @@ class MoEMixin(MixtureOfExperts): hidden_size=hidden_size, intermediate_size=intermediate_size, renormalize=renormalize, - use_grouped_topk=False, quant_config=self.quant_config, prefix=qual_name, activation=activation, @@ -259,17 +277,55 @@ class MoEMixin(MixtureOfExperts): routed_experts_cls=TransformersRoutedExperts, ) fuser = MoEBlockFuser.match(moe_block, experts_name) - if self.num_expert_groups <= 1 and fuser is not None: + # _maybe_apply_routed_scale_to_output edge case. Transformers + # decoder layers do not compensate for dividing by scaling factor. + reaches_fp16_trick = ( + routed_scaling_factor != 1.0 + and apply_routed_scale_to_output + and self.model_config.dtype == torch.float16 + and fuser is not None + and fuser.shared_name is not None + ) + if reaches_fp16_trick: + logger.warning_once( + "%s could be fused but routing it in vLLM would apply " + "`routed_scaling_factor` by dividing the shared expert " + "output, which only fp16 overflow protection expects the " + "decoder layer to compensate for. Falling back to routing " + "in Transformers; run in bfloat16 to fuse it.", + moe_block_cls, + ) + if fuser is not None and not reaches_fp16_trick: # MoE block forward is fully replaced. # gate/router and shared expert (if any) runs in FusedMoE. + shared_experts = fuser.shared_experts(moe_block, prefix) + # Store shared experts for later down projection adjustment + if shared_experts is not None: + hf_shared = shared_experts.shared_experts + glu_fuser = get_fuser(hf_shared) + down_name = getattr(glu_fuser, "down_name", None) + if down_name is not None: + shared_down_projs.append((hf_shared, down_name)) + # Prefer config, otherwise read it from fuser. + router_dtype = config_router_dtype or fuser.router_dtype + gate = fuser.gate(moe_block, prefix, router_dtype) kwargs |= dict( scoring_func=fuser.scoring_func, is_sequence_parallel=( self.parallel_config.use_sequence_parallel_moe ), - gate=fuser.gate(moe_block, prefix), - shared_experts=fuser.shared_experts(moe_block, prefix), + gate=gate, + shared_experts=shared_experts, ) + if router_dtype is not None: + kwargs["router_logits_dtype"] = router_dtype + if use_grouped_topk: + kwargs |= grouped_topk_routing_kwargs + if routed_scaling_factor != 1.0: + kwargs |= routed_scaling_factor_kwargs + bias = getattr(gate, "e_score_correction_bias", None) + if bias is not None: + kwargs["e_score_correction_bias"] = bias fuser.rewrite_forward(moe_block) routed = "gate + experts" if fuser.shared_name: @@ -333,3 +389,7 @@ class MoEMixin(MixtureOfExperts): self.num_moe_layers = len(self.moe_layers) # Continue with the replacement of layers in Base super().recursive_replace() + # GLUFuser likely fused shared_experts. The down projection after the GLU + # normally immediately reduces but we want FusedMoE to handle the reduction. + for hf_shared, down_name in shared_down_projs: + hf_shared.get_submodule(down_name).reduce_results = False diff --git a/vllm/model_executor/models/transformers/utils.py b/vllm/model_executor/models/transformers/utils.py index 4d9b01ce393..9000ace9051 100644 --- a/vllm/model_executor/models/transformers/utils.py +++ b/vllm/model_executor/models/transformers/utils.py @@ -16,7 +16,9 @@ # limitations under the License. """Transformers modeling backend utilities.""" +from collections.abc import Iterator from contextlib import contextmanager +from itertools import chain from pathlib import Path from typing import TYPE_CHECKING, Literal @@ -209,6 +211,11 @@ def recursive_replace_linear( _recursive_replace(model, prefix=prefix) +def named_state(module: nn.Module) -> Iterator[tuple[str, torch.Tensor]]: + """`module`'s own state (i.e. named parameters and buffers).""" + return chain(module.named_parameters(), module.named_buffers()) + + def log_replacement(name: str, old_module: nn.Module, new_module: nn.Module): logger.debug("%s: %s -> %s", name, old_module, new_module) From ed13deb37622cf9f603596743642807561c7691d Mon Sep 17 00:00:00 2001 From: oops-oom <liubin8905@vip.qq.com> Date: Tue, 28 Jul 2026 02:35:57 +0800 Subject: [PATCH 125/185] [Bugfix][CPU] Fall back to torch for unaligned swigluoai on NEON/vec MoE (#49985) Signed-off-by: oops-oom <73481342@qq.com> Co-authored-by: oops-oom <73481342@qq.com> Co-authored-by: Claude <noreply@anthropic.com> --- vllm/model_executor/layers/fused_moe/cpu_fused_moe.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py index a957aca9eca..c89e5ddc304 100644 --- a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py @@ -381,6 +381,12 @@ class CPUFusedMOE: if not (w13_output_size % 32 == 0 and w2_output_size % 32 == 0): return False, "none" + if ( + layer.activation == MoEActivation.SWIGLUOAI + and w2_input_size % _MOE_GROUPED_GEMM_N_TILE != 0 + ): + return False, "none" + supports_neon = current_platform.get_cpu_architecture() == CpuArchEnum.ARM if supports_neon: if ( From fd10e8946d7b27d1f3070d4f58febd8eac1c3a4b Mon Sep 17 00:00:00 2001 From: Rishi Puri <riship@nvidia.com> Date: Mon, 27 Jul 2026 16:03:58 -0300 Subject: [PATCH 126/185] [Test] Regression test for hybrid-Mamba eagle cache-peek in Mooncake connector (#43559) (#48361) Signed-off-by: Rishi Puri <riship@nvidia.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> --- .../unit/test_mooncake_store_coordinator.py | 36 +++++++++++++++++++ .../v1/mooncake/store/coordinator.py | 3 ++ 2 files changed, 39 insertions(+) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py index 6b56fd2f191..8b9d1e46b3c 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py @@ -531,3 +531,39 @@ def test_load_mask_without_eagle_unchanged(): assert hit == 64 masks = coord.load_mask(hs, token_len=hit) assert masks[0] == [True, True, True, True] + + +def _mamba(block_size=16): + return MambaSpec( + block_size=block_size, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + + +def test_lookup_with_eagle_hybrid_full_plus_mamba_no_overrun(): + """Full+Mamba with eagle must not overrun the attention-verified hit. + + ``MambaManager`` ignores ``drop_eagle_block`` (a Mamba block at position + p IS the recurrent state after (p + 1) * block_size tokens; there is + nothing to recompute), so granting the Mamba group the one-block eagle + peek margin lets it match one block PAST the eagle-pruned full-attention + hit, and adopting that length resumes the recurrent state ahead of the + verified token prefix (#43559). Gating the margin on ``not + isinstance(spec, MambaSpec)`` pins the hit to the attention-verified 48. + """ + groups = [ + KVCacheGroupSpec(["L0"], _full(16)), + KVCacheGroupSpec(["L1"], _mamba(16)), + ] + coord = _make_coord(groups, hash_block_size=16, use_eagle=True) + hs = _hashes(4) + exists = {(g, bytes(h)) for g in (0, 1) for h in hs} + cmap = ExternalCachedBlockPool(16, exists) + _masks, hit = coord.find_longest_cache_hit( + hs, max_length=64, cached_block_pool=cmap + ) + # FullAttn matches 4 blocks, eagle pops 1 -> 48 verified tokens. The + # Mamba group must serve its state@48 snapshot, not peek to state@64. + assert hit == 48 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py index c70ddf95a2f..2fc71ee4883 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py @@ -338,6 +338,9 @@ class MooncakeStoreCoordinator: drop_eagle_block = idx in eagle_indices and idx not in eagle_verified _max_length = curr_hit_length + # No eagle peek margin for a recurrent (Mamba) group: its finder + # never drops a block, so a widened bound would match past the + # attention-verified hit and resume from speculative state (#43559). if drop_eagle_block and not isinstance(spec, MambaSpec): eagle_margin = ( self.hash_block_size From 831d3848f16e21e652bf1c377d010c2b5fe76f37 Mon Sep 17 00:00:00 2001 From: Andrea Tassi <andrea.tassi@gmail.com> Date: Mon, 27 Jul 2026 20:48:35 +0100 Subject: [PATCH 127/185] [Core] Fail fast when /dev/shm is too small for the shm ring buffer (#48879) Signed-off-by: Dr Andrea Tassi <andrea@verticular.uk> Co-authored-by: Dr Andrea Tassi <andrea@verticular.uk> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> --- tests/distributed/test_shm_broadcast.py | 44 ++++++++++++++++++- .../device_communicators/shm_broadcast.py | 29 ++++++++++++ 2 files changed, 72 insertions(+), 1 deletion(-) diff --git a/tests/distributed/test_shm_broadcast.py b/tests/distributed/test_shm_broadcast.py index 33affb2a396..17957924051 100644 --- a/tests/distributed/test_shm_broadcast.py +++ b/tests/distributed/test_shm_broadcast.py @@ -4,6 +4,7 @@ import random import threading import time +from types import SimpleNamespace from unittest import mock import multiprocess as mp @@ -11,7 +12,12 @@ import numpy as np import pytest import torch.distributed as dist -from vllm.distributed.device_communicators.shm_broadcast import MessageQueue +from vllm.distributed.device_communicators import shm_broadcast +from vllm.distributed.device_communicators.shm_broadcast import ( + MessageQueue, + ShmRingBuffer, + check_shm_free_space, +) from vllm.distributed.utils import StatelessProcessGroup from vllm.utils.network_utils import get_open_port from vllm.utils.system_utils import update_environment_variables @@ -522,3 +528,39 @@ def test_warning_logs(caplog_vllm): # Clean up when done writer.shutdown() reader.shutdown() + + +def _fake_disk_usage(free_bytes: int): + return SimpleNamespace(total=free_bytes, used=0, free=free_bytes) + + +def test_check_shm_free_space_raises_when_insufficient(tmp_path): + with ( + mock.patch.object( + shm_broadcast.shutil, "disk_usage", return_value=_fake_disk_usage(32 << 20) + ), + pytest.raises(RuntimeError, match="Insufficient space"), + ): + check_shm_free_space(240 << 20, shm_path=str(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) + ): + 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")) + + +def test_shm_ring_buffer_creation_checks_free_space(): + with ( + mock.patch.object( + shm_broadcast.shutil, "disk_usage", return_value=_fake_disk_usage(1 << 20) + ), + mock.patch.object(shm_broadcast.os.path, "isdir", return_value=True), + pytest.raises(RuntimeError, match="Insufficient space"), + ): + ShmRingBuffer(n_reader=1, max_chunk_bytes=24 * 1024 * 1024, max_chunks=10) diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index 6b9dd4068b9..afabdf18c80 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import functools +import os import pickle +import shutil import sys import threading import time @@ -218,6 +220,32 @@ class SpinCondition: self.local_notify_socket.send(b"\x00") +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. + + Args: + required_bytes: Size of the shared-memory segment to be created. + shm_path: Mount point backing POSIX shared memory; skipped if absent. + + Raises: + RuntimeError: If ``required_bytes`` exceeds the free 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)." + ) + + class ShmRingBuffer: def __init__( self, @@ -288,6 +316,7 @@ class ShmRingBuffer: if name is None: # we are creating a buffer self.is_creator = True + check_shm_free_space(self.total_bytes_of_buffer) self.shared_memory = shared_memory.SharedMemory( create=True, size=self.total_bytes_of_buffer ) From b2f9e4caa49425d93667e01be5c9ad4c45bf81df Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Mon, 27 Jul 2026 15:56:52 -0400 Subject: [PATCH 128/185] [DSv4 Perf] Adaptive topk width, 1.0% E2E throughput improvement (#50004) Signed-off-by: yewentao256 <zhyanwentao@126.com> --- .../kernels/attention/test_flashmla_sparse.py | 37 +++++++++++++++++++ vllm/models/deepseek_v4/sparse_mla.py | 26 +++++++++---- 2 files changed, 56 insertions(+), 7 deletions(-) diff --git a/tests/kernels/attention/test_flashmla_sparse.py b/tests/kernels/attention/test_flashmla_sparse.py index 010c4479766..fa5602fdf96 100644 --- a/tests/kernels/attention/test_flashmla_sparse.py +++ b/tests/kernels/attention/test_flashmla_sparse.py @@ -4,6 +4,43 @@ import pytest import torch +def test_deepseek_v4_c128a_dynamic_topk_packed_buffers(): + from vllm.models.deepseek_v4.sparse_mla import build_c128a_topk_metadata + + device = torch.device("cuda") + capacity_width = 256 + active_width = 128 + global_decode_buffer = torch.empty( + (2, capacity_width), dtype=torch.int32, device=device + ) + decode_lens_buffer = torch.empty(2, dtype=torch.int32, device=device) + prefill_buffer = torch.empty((2, capacity_width), dtype=torch.int32, device=device) + + global_decode, decode_lens, prefill_local = build_c128a_topk_metadata( + positions=torch.tensor([255, 511], dtype=torch.int64, device=device), + compress_ratio=128, + num_decode_tokens=1, + token_to_req_indices=torch.tensor([0, 0], dtype=torch.int32, device=device), + block_table=torch.tensor([[3]], dtype=torch.int32, device=device), + block_size=capacity_width, + slot_mapping=torch.tensor([0, 1], dtype=torch.int64, device=device), + global_decode_buffer=global_decode_buffer, + decode_lens_buffer=decode_lens_buffer, + prefill_buffer=prefill_buffer, + max_compressed_tokens=active_width, + ) + + assert global_decode.shape == (1, active_width) + assert prefill_local.shape == (1, active_width) + assert global_decode.stride() == (active_width, 1) + assert prefill_local.stride() == (active_width, 1) + assert global_decode[0, :2].cpu().tolist() == [768, 769] + assert decode_lens.cpu().tolist() == [2] + assert prefill_local[0, :4].cpu().tolist() == list(range(4)) + assert torch.all(global_decode[0, 2:] == -1) + assert torch.all(prefill_local[0, 4:] == -1) + + def test_sparse_flashmla_metadata_smoke(): import vllm.v1.attention.ops.flashmla as fm diff --git a/vllm/models/deepseek_v4/sparse_mla.py b/vllm/models/deepseek_v4/sparse_mla.py index 4523d1875eb..4ac27bd59a8 100644 --- a/vllm/models/deepseek_v4/sparse_mla.py +++ b/vllm/models/deepseek_v4/sparse_mla.py @@ -261,6 +261,13 @@ class DeepseekV4FlashMLAMetadataBuilder( assert cm.positions is not None, ( "positions is required for C128A metadata build" ) + active_topk_width = min( + max( + triton.next_power_of_2(max(cm.max_seq_len // self.compress_ratio, 1)), + _C128A_TOPK_ALIGNMENT, + ), + self.c128a_max_compressed, + ) block_size = self.kv_cache_spec.block_size // self.compress_ratio global_decode, decode_lens, prefill_local = build_c128a_topk_metadata( cm.positions[:num_total], @@ -273,7 +280,7 @@ class DeepseekV4FlashMLAMetadataBuilder( self.c128a_global_decode_buffer, self.c128a_decode_lens_buffer, self.c128a_prefill_buffer, - max_compressed_tokens=self.c128a_max_compressed, + max_compressed_tokens=active_topk_width, ) result: dict[str, torch.Tensor | None] = {} @@ -305,25 +312,30 @@ def build_c128a_topk_metadata( Decode tokens: position → block_table lookup → global slot ids + topk_lens. Prefill tokens: position → local indices [0, ..., n-1, -1, ...]. - Writes into pre-allocated buffers for CUDA graph address stability. - Returns slices of the buffers. + Writes into packed views of pre-allocated buffers for CUDA graph stability. """ num_tokens = positions.shape[0] num_prefill_tokens = num_tokens - num_decode_tokens - global_decode = global_decode_buffer[:num_decode_tokens] + # view(-1) as 1-d array and then expanded to + # [num_decode_tokens, max_compressed_tokens] + global_decode = global_decode_buffer.view(-1)[ + : num_decode_tokens * max_compressed_tokens + ].view(num_decode_tokens, max_compressed_tokens) decode_lens = decode_lens_buffer[:num_decode_tokens] - prefill_local = prefill_buffer[:num_prefill_tokens] + prefill_local = prefill_buffer.view(-1)[ + : num_prefill_tokens * max_compressed_tokens + ].view(num_prefill_tokens, max_compressed_tokens) if num_tokens == 0: return global_decode, decode_lens, prefill_local _build_c128a_topk_metadata_kernel[(num_tokens,)]( global_decode_buffer, - global_decode_buffer.stride(0), + max_compressed_tokens, decode_lens_buffer, prefill_buffer, - prefill_buffer.stride(0), + max_compressed_tokens, positions, compress_ratio, max_compressed_tokens, From b5bcb3ce881e1d324ff7f6176ef27606558dbd74 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Mon, 27 Jul 2026 15:58:26 -0400 Subject: [PATCH 129/185] [Refactor] Remove dead code in multiple files (#49745) Signed-off-by: yewentao256 <zhyanwentao@126.com> --- .../linear/cute_dsl/_ll_bf16_dotprod.py | 4 - .../fused_moe/prepare_finalize/nixl_ep.py | 5 -- .../model_executor/layers/quantization/fp8.py | 10 --- .../layers/quantization/modelopt.py | 87 ------------------- .../utils/nvfp4_emulation_utils.py | 13 --- vllm/model_executor/models/funaudiochat.py | 19 ---- .../models/hyperclovax_vision.py | 32 ------- vllm/model_executor/models/idefics3.py | 24 ----- .../model_executor/models/llava_onevision2.py | 4 - .../models/qwen3_omni_moe_thinker.py | 46 ---------- vllm/model_executor/models/radio.py | 48 ---------- .../hunyuan_a13b_reasoning_parser.py | 12 --- vllm/tokenizers/mistral.py | 14 --- vllm/tool_parsers/gemma4_utils.py | 7 -- vllm/tool_parsers/streaming.py | 6 -- vllm/transformers_utils/configs/hunyuan_vl.py | 41 --------- .../transformers_utils/processors/minicpmo.py | 13 --- .../processors/minimax_m3.py | 80 ----------------- vllm/utils/nvtx_pytorch_hooks.py | 3 - .../v1/attention/ops/rocm_aiter_mla_sparse.py | 31 ------- vllm/v1/worker/gpu/model_runner.py | 4 - vllm/v1/worker/tpu_input_batch.py | 16 ---- 22 files changed, 519 deletions(-) diff --git a/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py b/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py index f6931ceb1c8..704fa7da4b4 100644 --- a/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py +++ b/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py @@ -307,7 +307,3 @@ class LLBf16Dotprod: ) if const_expr(self.use_pdl): cute.arch.griddepcontrol_launch_dependents() - - -def make_host_bf16(k_val: int, bs: int = 128): - return LLBf16Dotprod(k=k_val, bs=bs) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py index 89571278c6e..fd7d45655dd 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py @@ -160,11 +160,6 @@ class NixlEPPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): return topk_ids return self.global_to_physical[topk_ids] - def _map_local_to_global_ids(self, expert_topk_ids: torch.Tensor) -> torch.Tensor: - if self.local_expert_global_ids is None: - return expert_topk_ids - return self.local_expert_global_ids[expert_topk_ids] - def _do_quant( self, x: torch.Tensor | tuple[torch.Tensor, torch.Tensor], diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 626fc83cdff..49dd3a6d224 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -254,16 +254,6 @@ class CopyNumelCounter(TorchDispatchMode): return out -def _copy_missing_attrs(old: torch.Tensor, new: torch.Tensor) -> None: - """Copies any attrs present in `old` but not in `new` to `new`""" - new_attrs = set(dir(new)) - attrs_to_set = {} - for attr in dir(old): - if attr not in new_attrs: - attrs_to_set[attr] = getattr(old, attr) - set_weight_attrs(new, attrs_to_set) - - class Fp8LinearMethod(LinearMethodBase): """Linear method for FP8. Supports loading FP8 checkpoints with static weight scale and diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 7df8178ca71..e11097b7dea 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -58,9 +58,6 @@ from vllm.model_executor.layers.quantization.base_config import ( QuantizeMethodBase, ) from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod -from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( - swap_w13_to_w31, -) from vllm.model_executor.layers.quantization.utils.fp8_utils import ( process_fp8_input_tensor_strategy_moe, process_fp8_weight_channel_strategy, @@ -2014,90 +2011,6 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase): f"Expected {name} dtype {expected_dtype}, got {actual}." ) - def _shuffle_weights_for_trtllm(self, layer: torch.nn.Module) -> None: - """Shuffle weights and scales into FlashInfer TRTLLM MXFP8 layout.""" - from flashinfer import ( - reorder_rows_for_gated_act_gemm, - shuffle_matrix_a, - shuffle_matrix_sf_a, - ) - - epilogue_tile_m = 128 - num_experts = layer.w13_weight.shape[0] - is_gated = self.moe.is_act_and_mul - intermediate_size_factor = 2 if is_gated else 1 - - w13_weight = layer.w13_weight.data - w13_scale = layer.w13_weight_scale.data - if is_gated: - # FI TRTLLM gated kernels use W31 ordering. Model checkpoints store - # gated projection as W13, so convert once before shuffling. - w13_weight = swap_w13_to_w31(w13_weight) - w13_scale = swap_w13_to_w31(w13_scale) - - w13_weight_shuffled = [] - w2_weight_shuffled = [] - w13_scale_shuffled = [] - w2_scale_shuffled = [] - for i in range(num_experts): - w13_i = w13_weight[i].reshape( - intermediate_size_factor * layer.intermediate_size_per_partition, -1 - ) - w13_sf_i = w13_scale[i].reshape( - intermediate_size_factor * layer.intermediate_size_per_partition, -1 - ) - if is_gated: - # Reorder rows for gated activation layout expected by TRTLLM. - w13_i = reorder_rows_for_gated_act_gemm(w13_i.clone()) - w13_sf_i = reorder_rows_for_gated_act_gemm(w13_sf_i.clone()) - - w13_shuffled_i = shuffle_matrix_a(w13_i.view(torch.uint8), epilogue_tile_m) - w2_shuffled_i = shuffle_matrix_a( - layer.w2_weight.data[i].view(torch.uint8), epilogue_tile_m - ) - w13_weight_shuffled.append( - w13_shuffled_i.contiguous().view(MXFP8_VALUE_DTYPE) - ) - w2_weight_shuffled.append( - w2_shuffled_i.contiguous().view(MXFP8_VALUE_DTYPE) - ) - w13_sf_shuffled_i = shuffle_matrix_sf_a( - w13_sf_i.view(torch.uint8).reshape( - intermediate_size_factor * layer.intermediate_size_per_partition, - -1, - ), - epilogue_tile_m, - ) - w2_sf_shuffled_i = shuffle_matrix_sf_a( - layer.w2_weight_scale.data[i] - .view(torch.uint8) - .reshape(layer.hidden_size, -1), - epilogue_tile_m, - ) - w13_scale_shuffled.append( - w13_sf_shuffled_i.contiguous().view(MXFP8_SCALE_DTYPE) - ) - w2_scale_shuffled.append( - w2_sf_shuffled_i.contiguous().view(MXFP8_SCALE_DTYPE) - ) - - replace_parameter( - layer, "w13_weight", torch.stack(w13_weight_shuffled).contiguous() - ) - replace_parameter( - layer, "w2_weight", torch.stack(w2_weight_shuffled).contiguous() - ) - replace_parameter( - layer, - "w13_weight_scale", - torch.stack(w13_scale_shuffled).contiguous(), - ) - replace_parameter( - layer, - "w2_weight_scale", - torch.stack(w2_scale_shuffled).contiguous(), - ) - def _dequant_mxfp8_weights_to_bf16(self, layer: RoutedExperts) -> None: """One-time MXFP8->BF16 weight dequant for the emulation path. diff --git a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py b/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py index ad6b272371e..abb2043d7c4 100644 --- a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py +++ b/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py @@ -116,19 +116,6 @@ def _dequantize_nvfp4_kernel( tl.store(output_ptr + out_indices, result, mask=block_mask[:, None]) -@triton.jit -def _e2m1_lookup(magnitude): - """Lookup E2M1 float value from 3-bit magnitude.""" - result = tl.where(magnitude == 1, 0.5, 0.0) - result = tl.where(magnitude == 2, 1.0, result) - result = tl.where(magnitude == 3, 1.5, result) - result = tl.where(magnitude == 4, 2.0, result) - result = tl.where(magnitude == 5, 3.0, result) - result = tl.where(magnitude == 6, 4.0, result) - result = tl.where(magnitude == 7, 6.0, result) - return result - - @triton.jit def _round_to_fp4(x): """Round float values to the nearest E2M1 representable value. diff --git a/vllm/model_executor/models/funaudiochat.py b/vllm/model_executor/models/funaudiochat.py index 72b12e26b5c..2d3334f2751 100644 --- a/vllm/model_executor/models/funaudiochat.py +++ b/vllm/model_executor/models/funaudiochat.py @@ -257,25 +257,6 @@ class FunAudioChatAudioEncoder(nn.Module): def dtype(self) -> torch.dtype: return self.conv1.weight.dtype - def _prepare_attention_mask( - self, inputs_tensor: torch.Tensor, cu_seqlens: torch.Tensor - ) -> torch.Tensor | None: - if getattr(self.config, "_attn_implementation", "eager") == "flash_attention_2": - return None - - seq_length = inputs_tensor.shape[0] - attention_mask = torch.full( - (1, 1, seq_length, seq_length), - torch.finfo(inputs_tensor.dtype).min, - device=inputs_tensor.device, - dtype=inputs_tensor.dtype, - ) - for i in range(1, len(cu_seqlens)): - start = int(cu_seqlens[i - 1].item()) - end = int(cu_seqlens[i].item()) - attention_mask[..., start:end, start:end] = 0 - return attention_mask - def forward( self, input_features: torch.Tensor, diff --git a/vllm/model_executor/models/hyperclovax_vision.py b/vllm/model_executor/models/hyperclovax_vision.py index 53923d88438..593be5d88f6 100644 --- a/vllm/model_executor/models/hyperclovax_vision.py +++ b/vllm/model_executor/models/hyperclovax_vision.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # copied from : https://github.com/huggingface/transformers import ast -from collections import defaultdict from collections.abc import Iterable, Mapping, Sequence from functools import partial from itertools import accumulate @@ -891,37 +890,6 @@ class HCXVisionForCausalLM(nn.Module, SupportsMultiModal, SupportsPP): for i in range(len(feats_per_video)) ) - def _prepare_multimodal_kwargs(self, **kwargs: object): - output = defaultdict(list) - for k, v in kwargs.items(): - if len(v) < 1 or len(v[0]) < 1: - continue # if empty batch of empty sample - - new_k, is_video = k, False - if not k.endswith("_images") and not k.endswith("_videos"): - pass - else: - new_k, is_video = k.split("_")[:-1], k.split("_")[-1] - new_k = "_".join(new_k) - is_video = is_video == "videos" - - for _sample_idx, _v in enumerate(v): # batch -> sample - if new_k not in ["pixel_values"]: - if len(output[new_k]) < _sample_idx + 1: - output[new_k].append(list()) - _v = _v.detach().cpu().numpy().tolist() - output[new_k][_sample_idx] += _v - elif isinstance(_v, torch.Tensor): - if len(output[new_k]) < _sample_idx + 1: - output[new_k].append(list()) - output["is_videos"].append(list()) - _v = list(torch.unbind(_v, dim=0)) - output[new_k][_sample_idx] += _v - output["is_videos"][_sample_idx] += [ - is_video, - ] * len(_v) - return dict(output) - def compute_logits( self, hidden_states: torch.Tensor, diff --git a/vllm/model_executor/models/idefics3.py b/vllm/model_executor/models/idefics3.py index ad94719241c..4b3cbcabd2f 100644 --- a/vllm/model_executor/models/idefics3.py +++ b/vllm/model_executor/models/idefics3.py @@ -138,30 +138,6 @@ class Idefics3ProcessingInfo(BaseProcessingInfo): return height, width - def _get_resize_output_image_size( - self, - *, - image_width: int, - image_height: int, - resolution_max_side: int, - ) -> tuple[int, int]: - hf_processor = self.get_hf_processor() - image_processor: Idefics3ImageProcessor = hf_processor.image_processor - max_image_size = image_processor.size["longest_edge"] - if resolution_max_side > max_image_size: - raise ValueError( - "`resolution_max_side` cannot be larger than `max_image_size`" - ) - - height, width = image_height, image_width - - # Find the output size, when rescaling the longest edge to max_len and - # preserving the aspect ratio - height, width = self._resize_output_size( - height=height, width=width, max_len=resolution_max_side - ) - return height, width - def _get_image_feature_grid_size( self, *, diff --git a/vllm/model_executor/models/llava_onevision2.py b/vllm/model_executor/models/llava_onevision2.py index e17398cc8b7..58179ec00de 100644 --- a/vllm/model_executor/models/llava_onevision2.py +++ b/vllm/model_executor/models/llava_onevision2.py @@ -603,10 +603,6 @@ _OV2_FRAME_FACTOR = 2 _OV2_FPS_MIN_FRAMES = 4 -def _round_by_factor(n: float, factor: int) -> int: - return round(n / factor) * factor - - def _ceil_by_factor(n: float, factor: int) -> int: import math as _math diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index 32a622567ce..0afc6734afa 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -1530,52 +1530,6 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( result_placeholders["audio"] = audio_placeholders return result_placeholders - def _get_raw_input_ids( - self, - token_ids: list[int], - use_audio_in_video: bool = False, - ) -> list[int]: - tokenizer = self.info.get_tokenizer() - vision_bos_token = tokenizer.encode(tokenizer.vision_bos_token)[0] - vision_eos_token = tokenizer.encode(tokenizer.vision_eos_token)[0] - audio_bos_token = tokenizer.encode(tokenizer.audio_bos_token)[0] - audio_eos_token = tokenizer.encode(tokenizer.audio_eos_token)[0] - audio_token = tokenizer.encode("<|audio_pad|>")[0] - image_token = tokenizer.encode("<|image_pad|>")[0] - video_token = tokenizer.encode("<|video_pad|>")[0] - - result = token_ids[:] - if use_audio_in_video: - while True: - start = None - for i in range(len(result) - 1): - if result[i : i + 2] == [vision_bos_token, audio_bos_token]: - start = i - break - if start is not None: - end = None - for i in range(start + 2, len(result) - 1): - if result[i : i + 2] == [audio_eos_token, vision_eos_token]: - end = i - break - if end is not None: - result = ( - result[:start] - + [vision_bos_token, video_token, vision_eos_token] - + result[end + 2 :] - ) - else: - break - - for mm_token in [audio_token, image_token, video_token]: - compressed = [] - for x in result: - if x != mm_token or (not compressed or compressed[-1] != mm_token): - compressed.append(x) - result = compressed - - return result - class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMixin): def _process_audio_input( diff --git a/vllm/model_executor/models/radio.py b/vllm/model_executor/models/radio.py index 7ec320c5348..213083392df 100644 --- a/vllm/model_executor/models/radio.py +++ b/vllm/model_executor/models/radio.py @@ -331,54 +331,6 @@ class ViTPatchGenerator(nn.Module): def num_skip(self): return self.num_cls_tokens + self.num_registers - def _load_embed(self, src_embed: torch.Tensor, targ_embed: nn.Parameter): - if src_embed.shape != targ_embed.shape: - src_size = int(math.sqrt(src_embed.shape[1])) - - assert src_size**2 == src_embed.shape[1], ( - "Unable to interpolate non-square embedding" - ) - - src_embed = rearrange( - src_embed, "b (h w) c -> b c h w", h=src_size, w=src_size - ) - src_embed = F.interpolate( - src_embed, - size=(self.num_rows, self.num_cols), - mode="bicubic", - align_corners=True, - antialias=False, - ) - src_embed = rearrange(src_embed, "b c h w -> b (h w) c") - targ_embed.data.copy_(src_embed) - - def _load_projection( - self, src_proj_weight: torch.Tensor, targ_proj_weight: torch.Tensor - ): - if src_proj_weight.shape != targ_proj_weight.shape: - src_patch_size = int(math.sqrt(src_proj_weight.shape[1] // 3)) - - assert (src_patch_size**2) * 3 == src_proj_weight.shape[1], ( - "Unable to interpolate non-square patch size" - ) - - src_proj_weight = rearrange( - src_proj_weight, - "b (c h w) -> b c h w", - c=3, - h=src_patch_size, - w=src_patch_size, - ) - src_proj_weight = F.interpolate( - src_proj_weight, - size=(self.patch_size, self.patch_size), - mode="bicubic", - align_corners=True, - antialias=False, - ) - src_proj_weight = rearrange(src_proj_weight, "b c h w -> b (c h w)") - targ_proj_weight.data.copy_(src_proj_weight) - def embed_patches(self, x: torch.Tensor) -> torch.Tensor: patches = self.im_to_patches(x) patches = self.embedder(patches) diff --git a/vllm/reasoning/hunyuan_a13b_reasoning_parser.py b/vllm/reasoning/hunyuan_a13b_reasoning_parser.py index 257dc0f9540..1da1938e72b 100644 --- a/vllm/reasoning/hunyuan_a13b_reasoning_parser.py +++ b/vllm/reasoning/hunyuan_a13b_reasoning_parser.py @@ -130,18 +130,6 @@ class HunyuanA13BReasoningParser(ReasoningParser): return None, model_output - def _is_strict_increasing_subsequence( - self, subsequence: Sequence[int], sequence: Sequence[int] - ) -> bool: - if not subsequence: - return False - - sub_idx = 0 - for num in sequence: - if sub_idx < len(subsequence) and num == subsequence[sub_idx]: - sub_idx += 1 - return sub_idx == len(subsequence) - def extract_reasoning_streaming( self, previous_text: str, diff --git a/vllm/tokenizers/mistral.py b/vllm/tokenizers/mistral.py index 1164f7c41a7..3d8bf37eacb 100644 --- a/vllm/tokenizers/mistral.py +++ b/vllm/tokenizers/mistral.py @@ -45,20 +45,6 @@ if TYPE_CHECKING: logger = init_logger(__name__) -def _pop_unallowed_keys_and_warn( - dictionary: dict[str, Any], allowed_keys: set[str], err_dict_name: str -): - keys = list(dictionary.keys()) - for key in keys: - if key not in allowed_keys: - dictionary.pop(key) - logger.warning_once( - f"'{key=}' is not supported by mistral-common " - f"for {err_dict_name}. It has been popped from the " - "object." - ) - - def maybe_serialize_tool_calls(request: "MistralChatCompletionRequest"): # SEE: https://github.com/vllm-project/vllm/pull/9951 # Credits go to: @gcalmettes diff --git a/vllm/tool_parsers/gemma4_utils.py b/vllm/tool_parsers/gemma4_utils.py index a72e16ea56f..d9aad254c72 100644 --- a/vllm/tool_parsers/gemma4_utils.py +++ b/vllm/tool_parsers/gemma4_utils.py @@ -37,15 +37,8 @@ do not need a transformers dependency for output parsing. import regex as re -# Tool call delimiter tokens as they appear in decoded text. -# Standard format: <|tool_call>call:name{args}<tool_call|> -_TOOL_CALL_START_TAG = "<|tool_call>" -_TOOL_CALL_END_TAG = "<tool_call|>" _TOOL_RESPONSE_START_TAG = "<|tool_response>" -# Gemma4 escape token as it appears in decoded text. -_ESCAPE_TOKEN = '<|"|>' - def _parse_tool_arguments(args_str: str) -> dict[str, str]: """Parse tool call arguments from the Gemma4 compact format. diff --git a/vllm/tool_parsers/streaming.py b/vllm/tool_parsers/streaming.py index 5ee7c6f6c29..c3b8a11ba32 100644 --- a/vllm/tool_parsers/streaming.py +++ b/vllm/tool_parsers/streaming.py @@ -47,12 +47,6 @@ def _bracket_level_state( return level, in_string, escaped -def _bracket_level(s: str, opening: str = "{", closing: str = "}") -> int: - """Calculate the current level of nested brackets in a string.""" - level, _, _ = _bracket_level_state(s, opening, closing) - return level - - def filter_delta_text( delta_text: str, previous_text: str, diff --git a/vllm/transformers_utils/configs/hunyuan_vl.py b/vllm/transformers_utils/configs/hunyuan_vl.py index a826ed9b515..548dcfefcaf 100644 --- a/vllm/transformers_utils/configs/hunyuan_vl.py +++ b/vllm/transformers_utils/configs/hunyuan_vl.py @@ -194,7 +194,6 @@ class HunYuanVLTextConfig(PretrainedConfig): self.use_cache = use_cache self.rope_theta = rope_theta self.rope_scaling = rope_scaling - # self._rope_scaling_validation() # TODO: Need validation? self.attention_bias = attention_bias self.attention_dropout = attention_dropout @@ -206,46 +205,6 @@ class HunYuanVLTextConfig(PretrainedConfig): **kwargs, ) - def _rope_scaling_validation(self): - """ - Validate the `rope_scaling` configuration. - """ - if self.rope_scaling is None: - return - - if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2: - raise ValueError( - "`rope_scaling` must be a dictionary with with two fields, `type` and " - f"`factor` or `type` and `alpha`, got {self.rope_scaling}" - ) - rope_scaling_type = self.rope_scaling.get("type", None) - rope_scaling_factor = self.rope_scaling.get("factor", None) - rope_scaling_alpha = self.rope_scaling.get("alpha", None) - if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]: - raise ValueError( - "`rope_scaling`'s type field must be one of ['linear', 'dynamic'], " - f"got {rope_scaling_type}" - ) - if rope_scaling_factor is None and rope_scaling_alpha is None: - raise ValueError( - "`rope_scaling`'s factor or alpha field must be have one, " - "got both of none" - ) - if rope_scaling_factor is not None and ( - not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0 - ): - raise ValueError( - "`rope_scaling`'s factor field must be a float > 1.0, " - f"got {rope_scaling_factor}" - ) - if rope_scaling_alpha is not None and ( - not isinstance(rope_scaling_alpha, float) or rope_scaling_alpha <= 1.0 - ): - raise ValueError( - "`rope_scaling`'s alpha field must be a float > 1.0, " - f"got {rope_scaling_alpha}" - ) - class HunYuanVLConfig(PretrainedConfig): model_type = "hunyuan_vl" diff --git a/vllm/transformers_utils/processors/minicpmo.py b/vllm/transformers_utils/processors/minicpmo.py index 899e0402ba5..d9ac841cd86 100644 --- a/vllm/transformers_utils/processors/minicpmo.py +++ b/vllm/transformers_utils/processors/minicpmo.py @@ -72,19 +72,6 @@ class MiniCPMOProcessor(ProcessorMixin): self.version = getattr(image_processor, "version", None) self.pool_step = pool_step - def _safe_get_token_id(self, attr_name, default_token_str): - """Get token ID safely, with fallback to default.""" - val = getattr(self.tokenizer, attr_name, None) - if val is None: - val = self.tokenizer.convert_tokens_to_ids(default_token_str) - if val is None: - return -1 - return val - - def _safe_get_token_str(self, attr_name, default_token_str): - """Get token string safely, with fallback to default.""" - return getattr(self.tokenizer, attr_name, default_token_str) - def __call__( self, text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput], diff --git a/vllm/transformers_utils/processors/minimax_m3.py b/vllm/transformers_utils/processors/minimax_m3.py index 13dbce5368f..60fa7ed536b 100644 --- a/vllm/transformers_utils/processors/minimax_m3.py +++ b/vllm/transformers_utils/processors/minimax_m3.py @@ -15,7 +15,6 @@ feeds decoded frames to the processor. import math -import regex as re import torch from torchvision.transforms import InterpolationMode from transformers import AutoTokenizer, BatchFeature @@ -554,85 +553,6 @@ class MiniMaxVLProcessor(ProcessorMixin): self.VISION_END_TOKEN ) - def _prune_video_tokens( - self, - input_text: str, - video_segments: list[int], - video_token: str, - ) -> str: - """Prune video tokens by temporal_patch_size (e.g., 2:1). - - Expects the prompt to carry exactly sum(video_segments) video tokens - — i.e. one token per *sampled* frame — then drops tokens. - """ - # If no videos or temporal_patch_size <= 1, no pruning needed - if not video_segments or self.video_processor.temporal_patch_size <= 1: - return input_text - - # Split while keeping delimiters - special_tokens = [video_token] - pattern = "|".join(map(re.escape, special_tokens)) - parts = re.split(f"({pattern})", input_text) - - def is_timestamp(text: str) -> bool: - """Check if text ends with timestamp format like ']<]0.0 seconds[>['""" - return ( - text.endswith("seconds[>[") - or text.endswith("seconds[>[ ") - or text.endswith("seconds [>[") - or text.endswith("seconds [>[ ") - ) - - def extract_timestamp(text: str) -> str: - """Extract timestamp text from the end, starting from ']<]'""" - start_index = text.rfind("]<]") - if start_index == -1: - raise ValueError(f"Failed to extract timestamp: {text}") - return text[start_index:] - - # Build new text with pruned video tokens - final_parts = [] - current_seg_idx = 0 # Which video segment we're in - frame_in_seg = 0 # Frame index within current segment - last_timestamp_len = 0 # Length of timestamp to potentially remove - - for part in parts: - if part == video_token: - if current_seg_idx < len(video_segments): - if frame_in_seg % self.video_processor.temporal_patch_size == 0: - # Keep this video token - final_parts.append(part) - frame_in_seg += 1 - if frame_in_seg >= video_segments[current_seg_idx]: - current_seg_idx += 1 - frame_in_seg = 0 - last_timestamp_len = 0 - else: - # Skip this video token - frame_in_seg += 1 - if frame_in_seg >= video_segments[current_seg_idx]: - current_seg_idx += 1 - frame_in_seg = 0 - # Remove the timestamp that was already appended - if last_timestamp_len > 0: - assert len(final_parts) > 0 - final_parts[-1] = final_parts[-1][:-last_timestamp_len] - last_timestamp_len = 0 - else: - # No more video segments, keep as is - final_parts.append(part) - last_timestamp_len = 0 - else: - # Text part - final_parts.append(part) - # Check if this text ends with a timestamp - if is_timestamp(part): - last_timestamp_len = len(extract_timestamp(part)) - else: - last_timestamp_len = 0 - - return "".join(final_parts) - def __call__( self, images=None, diff --git a/vllm/utils/nvtx_pytorch_hooks.py b/vllm/utils/nvtx_pytorch_hooks.py index 39e2a9a136e..25f76edbfaa 100644 --- a/vllm/utils/nvtx_pytorch_hooks.py +++ b/vllm/utils/nvtx_pytorch_hooks.py @@ -231,9 +231,6 @@ class PytHooks: super().__init__() self.module_to_name_map = {} - def _process_layer_params(self, module_obj): - return process_layer_params(module_obj) - def module_fwd_hook(self, module_obj, in_tensor, out_tensor): """Callback function that ends the NVTX marker. Records the module name and tensor information. diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 96f7693430a..28ec8d4b592 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -592,37 +592,6 @@ def rocm_fp8_mqa_logits( return fp8_mqa_logits_torch(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke) -def _topk_indices_torch( - logits: torch.Tensor, - topk_tokens: int, - row_starts: torch.Tensor | None = None, -) -> torch.Tensor: - k = min(topk_tokens, logits.shape[-1]) - values, indices = torch.topk(logits, k=k, dim=-1) - indices = indices.to(torch.int32) - indices = torch.where( - values == float("-inf"), - torch.full_like(indices, -1, dtype=torch.int32), - indices, - ) - if row_starts is not None: - # Match the CUDA top_k_per_row_prefill contract: indices are local to - # each row's valid [row_start, row_end) range, not columns in the - # concatenated chunk logits matrix. - starts = row_starts.to(dtype=torch.int32).view(-1, 1) - indices = torch.where(indices < 0, indices, indices - starts) - if k == topk_tokens: - return indices - padded = torch.full( - (logits.shape[0], topk_tokens), - -1, - dtype=torch.int32, - device=logits.device, - ) - padded[:, :k] = indices - return padded - - def rocm_aiter_sparse_attn_indexer_fake( hidden_states: torch.Tensor, k_cache_prefix: LayerNameType, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index f0da049ce46..9182ada4394 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -731,10 +731,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): if self.encoder_cache is not None: self.encoder_cache.reset_encoder_cache() - def _get_num_input_tokens(self, num_scheduled_tokens: int) -> int: - # SP is not supported yet. - return num_scheduled_tokens - def profile_cudagraph_memory(self) -> int: # NOTE(woosuk): It is TBD whether we keep this API or not. return 0 diff --git a/vllm/v1/worker/tpu_input_batch.py b/vllm/v1/worker/tpu_input_batch.py index 1396c8ad9d5..245ac3c9028 100644 --- a/vllm/v1/worker/tpu_input_batch.py +++ b/vllm/v1/worker/tpu_input_batch.py @@ -495,22 +495,6 @@ class InputBatch: del self._req_ids[self.num_reqs :] del self.req_output_token_ids[self.num_reqs :] - def _make_prompt_token_ids_tensor(self) -> torch.Tensor: - max_prompt_len = self.num_prompt_tokens[: self.num_reqs].max() - prompt_token_ids_cpu_tensor = torch.empty( - (self.num_reqs, max_prompt_len), - device="cpu", - dtype=torch.int64, - pin_memory=self.pin_memory, - ) - prompt_token_ids = prompt_token_ids_cpu_tensor.numpy() - prompt_token_ids[:] = self.token_ids_cpu[: self.num_reqs, :max_prompt_len] - # Use the value of vocab_size as a pad since we don't have a - # token_id of this value. - for i in range(self.num_reqs): - prompt_token_ids[i, self.num_prompt_tokens[i] :] = self.vocab_size - return prompt_token_ids_cpu_tensor.to(device=self.device, non_blocking=True) - def make_lora_inputs( self, num_scheduled_tokens: np.ndarray, num_sampled_tokens: np.ndarray ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]: From 1053e248f02f453390fbaadbecd2b94beff2fbb4 Mon Sep 17 00:00:00 2001 From: amd-sourjya <sourroy@amd.com> Date: Mon, 27 Jul 2026 14:01:34 -0700 Subject: [PATCH 130/185] [ROCm][Quantization][5/N] Refactor quark_moe w8a8-int8 w/ oracle (#46765) Signed-off-by: amd-sourjya <amd-sourjya@users.noreply.github.com> Co-authored-by: amd-sourjya <amd-sourjya@users.noreply.github.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> --- .../configs/Qwen1.5-MoE-A2.7B-Chat-INT8.yaml | 5 + .../configs/models-mi3xx-fp8-and-mixed.txt | 1 + tests/quantization/test_int8_moe_oracle.py | 105 ++++++++++++++++ tests/quantization/test_quark.py | 2 +- .../layers/fused_moe/experts/triton_moe.py | 18 ++- .../layers/fused_moe/oracle/int8.py | 4 +- .../layers/quantization/quark/quark_moe.py | 114 ++++++++++++++++-- .../layers/quantization/utils/quant_utils.py | 2 + 8 files changed, 234 insertions(+), 17 deletions(-) create mode 100644 tests/evals/gsm8k/configs/Qwen1.5-MoE-A2.7B-Chat-INT8.yaml create mode 100644 tests/quantization/test_int8_moe_oracle.py diff --git a/tests/evals/gsm8k/configs/Qwen1.5-MoE-A2.7B-Chat-INT8.yaml b/tests/evals/gsm8k/configs/Qwen1.5-MoE-A2.7B-Chat-INT8.yaml new file mode 100644 index 00000000000..20e3948cce6 --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen1.5-MoE-A2.7B-Chat-INT8.yaml @@ -0,0 +1,5 @@ +model_name: "amd/Qwen1.5-MoE-A2.7B-Chat-w-int8-a-int8-sym" +accuracy_threshold: 0.50 +num_questions: 1319 +num_fewshot: 5 +server_args: "--enforce-eager --max-model-len 4096" diff --git a/tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt b/tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt index bcd00044bc0..66d4224665a 100644 --- a/tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt +++ b/tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt @@ -1,6 +1,7 @@ Qwen3-0.6B-FP8.yaml Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml Qwen1.5-MoE-W4A16-CT.yaml +Qwen1.5-MoE-A2.7B-Chat-INT8.yaml DeepSeek-V2-Lite-Instruct-FP8.yaml Qwen3-Next-FP8-EP2_MI355.yaml Qwen3-30B-A3B-Thinking-2507-FP8.yaml diff --git a/tests/quantization/test_int8_moe_oracle.py b/tests/quantization/test_int8_moe_oracle.py new file mode 100644 index 00000000000..2eb6fa92a54 --- /dev/null +++ b/tests/quantization/test_int8_moe_oracle.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Tests for INT8 (W8A8) fused-MoE oracle backend selection. + +These exercise ``select_int8_moe_backend`` only (no kernels are launched), so +they run on any platform where the Triton INT8 MoE kernel is available — CUDA +(SM >= 7.5) or ROCm — not just gfx950. +""" + +import pytest +import torch + +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, +) +from vllm.model_executor.layers.fused_moe.oracle.int8 import ( + Int8MoeBackend, + select_int8_moe_backend, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kInt8DynamicTensorSym, + kInt8DynamicTokenSym, + kInt8StaticChannelSym, + kInt8StaticTensorSym, +) +from vllm.platforms import current_platform + +# The Triton int8_w8a8 fused-MoE kernel is available on CUDA (Turing+) and on +# ROCm CDNA GPUs. Gate on that rather than on a specific arch. +INT8_MOE_SUPPORTED = ( + current_platform.is_cuda() and current_platform.has_device_capability((7, 5)) +) or current_platform.is_rocm() + +requires_int8_moe = pytest.mark.skipif( + not INT8_MOE_SUPPORTED, + reason="Requires a GPU with Triton INT8 MoE support (CUDA SM>=7.5 or ROCm)", +) + + +def _make_int8_moe_config(moe_backend: str = "auto") -> FusedMoEConfig: + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + + return FusedMoEConfig( + num_experts=8, + experts_per_token=2, + hidden_dim=256, + intermediate_size=256, + num_local_experts=8, + num_logical_experts=8, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device="cuda", + routing_method=RoutingMethodType.Renormalize, + moe_backend=moe_backend, + ) + + +@requires_int8_moe +@pytest.mark.parametrize( + "weight_key,activation_key", + [ + # per-channel weight + dynamic per-token activation + (kInt8StaticChannelSym, kInt8DynamicTokenSym), + # per-tensor weight + dynamic per-tensor activation + (kInt8StaticTensorSym, kInt8DynamicTensorSym), + ], +) +def test_int8_dynamic_schemes_dispatch_to_triton(weight_key, activation_key): + """Both dynamic-activation INT8 MoE schemes (per-channel + per-tensor + weights) select the Triton backend.""" + config = _make_int8_moe_config() + backend, experts_cls = select_int8_moe_backend( + config, weight_key=weight_key, activation_key=activation_key + ) + assert backend == Int8MoeBackend.TRITON + assert experts_cls is not None + + +@requires_int8_moe +def test_int8_explicit_moe_backend_triton(): + """An explicit --moe-backend triton selects the Triton INT8 backend.""" + config = _make_int8_moe_config(moe_backend="triton") + backend, experts_cls = select_int8_moe_backend( + config, + weight_key=kInt8StaticChannelSym, + activation_key=kInt8DynamicTokenSym, + ) + assert backend == Int8MoeBackend.TRITON + assert experts_cls is not None + + +@requires_int8_moe +def test_int8_unsupported_moe_backend_raises(): + """An unsupported --moe-backend for INT8 MoE raises a clear error.""" + config = _make_int8_moe_config(moe_backend="cutlass") + with pytest.raises(ValueError, match="not supported for Int8 MoE"): + select_int8_moe_backend( + config, + weight_key=kInt8StaticChannelSym, + activation_key=kInt8DynamicTokenSym, + ) diff --git a/tests/quantization/test_quark.py b/tests/quantization/test_quark.py index b31b8580e34..009559c22a9 100644 --- a/tests/quantization/test_quark.py +++ b/tests/quantization/test_quark.py @@ -150,7 +150,7 @@ def test_quark_int8_w_per_tensor_a_per_tensor(vllm_runner, tp): @pytest.mark.parametrize("tp", [1]) def test_quark_int8_w8a8_moe(vllm_runner, tp): """Test W8A8 INT8 MoE quantization with a tiny Qwen3 MoE model.""" - model_path = "nameistoken/tiny-qwen3-moe-w8a8-int8-quark" + model_path = "amd/tiny-qwen3-moe-w8a8-int8" with vllm_runner( model_path, enforce_eager=True, diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py index 8610e6b1a70..3656e90c946 100644 --- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py @@ -45,9 +45,11 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8StaticTensorSym, kInt4Static, kInt4Static32, + kInt8DynamicTensorSym, kInt8DynamicTokenSym, kInt8Static, kInt8StaticChannelSym, + kInt8StaticTensorSym, ) from vllm.platforms import current_platform from vllm.triton_utils import tl @@ -103,15 +105,25 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular): weight_key: QuantKey | None, activation_key: QuantKey | None, ) -> bool: - # INT8 requires at least 7.5 (Turing). + # INT8 requires at least 7.5 (Turing) on CUDA. ROCm CDNA GPUs + # (e.g. MI2xx/MI3xx/gfx950) provide native INT8 matrix-core support and + # the Triton int8_w8a8 fused MoE kernel handles them. device_supports_int8 = ( current_platform.is_cuda() and current_platform.has_device_capability((7, 5)) - ) + ) or current_platform.is_rocm() supported: list[tuple[QuantKey | None, QuantKey | None]] = [(None, None)] if device_supports_int8: - supported.append((kInt8StaticChannelSym, kInt8DynamicTokenSym)) + # Activations are consumed as float and quantized to int8 + # dynamically inside the kernel, so only dynamic-activation int8 + # schemes are supported (static-activation int8 is not). + supported += [ + # per-channel weight + dynamic per-token activation + (kInt8StaticChannelSym, kInt8DynamicTokenSym), + # per-tensor weight + dynamic per-tensor activation + (kInt8StaticTensorSym, kInt8DynamicTensorSym), + ] if current_platform.supports_fp8(): supported += [ (kFp8Static128BlockSym, kFp8Dynamic128Sym), diff --git a/vllm/model_executor/layers/fused_moe/oracle/int8.py b/vllm/model_executor/layers/fused_moe/oracle/int8.py index e2fb3e9114f..db484b9b31b 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int8.py @@ -166,7 +166,9 @@ def select_int8_moe_backend( logger.debug_once(_make_log_unsupported(backend, reason)) raise NotImplementedError( - "No Int8 MoE backend supports the deployment configuration." + "No Int8 MoE backend supports the deployment configuration " + f"(weight_key={weight_key}, activation_key={activation_key}). " + "Set `VLLM_LOGGING_LEVEL=DEBUG` to see per-backend unsupported reasons." ) diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 15023d7ca39..c5c9c6988f8 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -33,6 +33,13 @@ from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( make_fp8_moe_quant_config, select_fp8_moe_backend, ) +from vllm.model_executor.layers.fused_moe.oracle.int8 import ( + Int8MoeBackend, + convert_to_int8_moe_kernel_format, + make_int8_moe_kernel, + make_int8_moe_quant_config, + select_int8_moe_backend, +) from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( TRITON_BACKENDS, Mxfp4MoeBackend, @@ -59,6 +66,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8DynamicTokenSym, kFp8StaticChannelSym, kFp8StaticTensorSym, + kInt8DynamicTensorSym, + kInt8DynamicTokenSym, + kInt8StaticChannelSym, + kInt8StaticTensorSym, kMxfp4Dynamic, kNvfp4Dynamic, kNvfp4Static, @@ -502,6 +513,35 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod): self.weight_qscheme = self.weight_quant.get("qscheme", "per_tensor") self.static_input_scales = not self.input_quant.get("is_dynamic", False) + self.moe_quant_config: FusedMoEQuantConfig | None = None + self.moe_kernel: mk.FusedMoEKernel | None = None + self.int8_backend: Int8MoeBackend | None = None + self.experts_cls: type[mk.FusedMoEExperts] | None = None + + # Dynamic-activation INT8 MoE goes through the oracle + modular kernel. + # The modular TritonExperts kernel consumes float activations and + # quantizes them to int8 itself, so it cannot apply a loaded static + # activation scale (this matches CompressedTensorsW8A8Int8MoEMethod). + # TODO: Static-activation INT8 therefore stays on the legacy fused_experts + # path (see apply()) for now, preserving pre-refactor behavior. + # Needs to be migrated to expert backend. + if not self.static_input_scales: + # Map the Quark weight scheme to oracle quant keys. Per-channel + # weights pair with dynamic per-token activations; per-tensor + # weights with dynamic per-tensor activations. + if self.weight_qscheme == "per_channel": + weight_key = kInt8StaticChannelSym + activation_key = kInt8DynamicTokenSym + else: + weight_key = kInt8StaticTensorSym + activation_key = kInt8DynamicTensorSym + + self.int8_backend, self.experts_cls = select_int8_moe_backend( + config=moe, + weight_key=weight_key, + activation_key=activation_key, + ) + def create_weights( self, layer: torch.nn.Module, @@ -563,7 +603,7 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod): set_weight_attrs(w13_weight_scale, extra_weight_attrs) set_weight_attrs(w2_weight_scale, extra_weight_attrs) else: - # per-tensor: one scalar per expert + # per-tensor: one scalar per expert (two for the fused w1/w3) w13_weight_scale = torch.nn.Parameter( torch.ones(num_experts, 2, dtype=torch.float32), requires_grad=False, @@ -582,6 +622,8 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod): # INPUT_SCALES if self.static_input_scales: + # Static activations: the per-expert scales are loaded from the + # checkpoint (used by the legacy fused_experts path). w13_input_scale = torch.nn.Parameter( torch.ones(num_experts, dtype=torch.float32), requires_grad=False, @@ -596,6 +638,7 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod): layer.register_parameter("w2_input_scale", w2_input_scale) set_weight_attrs(w2_input_scale, extra_weight_attrs) else: + # Dynamic activations are quantized in-kernel (no stored scale). layer.w13_input_scale = None layer.w2_input_scale = None @@ -673,7 +716,8 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod): if hasattr(layer, attr): delattr(layer, attr) - # For static input scales, collapse per-expert scales to single max + # For static input scales, collapse the per-expert scales to a single + # value (the legacy fused_experts path expects one scale per layer). if self.static_input_scales: if layer.w13_input_scale is None or layer.w2_input_scale is None: raise ValueError( @@ -709,7 +753,8 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod): ), ) - # For per-tensor weights, merge w1/w3 scales into single per-expert + # For per-tensor weights, merge the w1/w3 scales into a single + # per-expert scale (dequant -> requant at the max scale). if self.weight_qscheme == "per_tensor": assert layer.w13_weight_scale is not None shard_size = layer.intermediate_size_per_partition @@ -734,10 +779,42 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod): max_w13_scales, requires_grad=False ) + # Dynamic activations run through the oracle's modular kernel; static + # activations use the legacy fused_experts path in apply(). + if not self.static_input_scales: + assert self.int8_backend is not None + assert self.experts_cls is not None + w13, w2 = convert_to_int8_moe_kernel_format( + int8_backend=self.int8_backend, + w13=layer.w13_weight, + w2=layer.w2_weight, + layer=layer, + w13_scale=layer.w13_weight_scale, + ) + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + assert self.moe_quant_config is not None + + if not self.static_input_scales: + assert self.int8_backend is not None + assert self.experts_cls is not None + self.moe_kernel = make_int8_moe_kernel( + int8_backend=self.int8_backend, + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + experts_cls=self.experts_cls, + routing_tables=layer._expert_routing_tables(), + layer=layer, + ) + def get_fused_moe_quant_config( self, layer: torch.nn.Module ) -> FusedMoEQuantConfig | None: - if self.weight_qscheme == "per_channel" and not self.static_input_scales: + # Static-activation INT8 has no oracle backend (it uses the legacy + # fused_experts path); build its config directly. + if self.int8_backend is None: return int8_w8a8_moe_quant_config( w1_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, @@ -745,21 +822,18 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod): a2_scale=layer.w2_input_scale, w1_bias=getattr(layer, "w13_bias", None), w2_bias=getattr(layer, "w2_bias", None), - per_act_token_quant=True, + per_act_token_quant=False, ) - is_dynamic = not self.static_input_scales - is_per_channel = self.weight_qscheme == "per_channel" - return FusedMoEQuantConfig.make( - torch.int8, + return make_int8_moe_quant_config( + int8_backend=self.int8_backend, w1_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, a1_scale=layer.w13_input_scale, a2_scale=layer.w2_input_scale, w1_bias=getattr(layer, "w13_bias", None), w2_bias=getattr(layer, "w2_bias", None), - per_act_token_quant=is_dynamic, - per_out_ch_quant=is_per_channel, - block_shape=None, + per_act_token_quant=(self.weight_qscheme == "per_channel"), + layer=layer, ) def apply( @@ -771,6 +845,22 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod): shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: + if self.moe_kernel is not None: + return self.moe_kernel.apply( + hidden_states=x, + w1=layer.w13_weight, + w2=layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + expert_map=layer.expert_map, + shared_experts_input=shared_experts_input, + ) + + # Static-activation INT8 MoE: legacy monolithic path (the modular kernel + # quantizes activations dynamically and cannot apply a loaded scale). from vllm.model_executor.layers.fused_moe import fused_experts return fused_experts( diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py index 705da43c373..453e08e6120 100644 --- a/vllm/model_executor/layers/quantization/utils/quant_utils.py +++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py @@ -190,6 +190,8 @@ kInt4Static32Asym = QuantKey( kInt8StaticChannelSym = QuantKey(torch.int8, kStaticChannelScale, symmetric=True) kInt8DynamicTokenSym = QuantKey(torch.int8, kDynamicTokenScale, symmetric=True) +kInt8StaticTensorSym = QuantKey(torch.int8, kStaticTensorScale, symmetric=True) +kInt8DynamicTensorSym = QuantKey(torch.int8, kDynamicTensorScale, symmetric=True) # INT4 W4A8 quantization keys From 99115fcdcde9736d4097ea15e7bbab368c405edf Mon Sep 17 00:00:00 2001 From: Andreas Karatzas <akaratza@amd.com> Date: Mon, 27 Jul 2026 16:29:08 -0500 Subject: [PATCH 131/185] [CI] Initialize DeepEP FP8 test weights (#49912) Signed-off-by: Andreas Karatzas <akaratza@amd.com> --- tests/kernels/moe/test_deepep_moe.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/kernels/moe/test_deepep_moe.py b/tests/kernels/moe/test_deepep_moe.py index 8d12e2888d0..933718ef333 100644 --- a/tests/kernels/moe/test_deepep_moe.py +++ b/tests/kernels/moe/test_deepep_moe.py @@ -66,8 +66,14 @@ def make_weights( # per-out-channel weight quantization assert dtype == current_platform.fp8_dtype() - w1 = torch.empty((e, 2 * n, k), device="cuda", dtype=torch.float16) - w2 = torch.empty((e, k, n), device="cuda", dtype=torch.float16) + # Keep FP8 inputs finite and bounded. torch.empty made this test depend on + # allocator contents, while larger values exceed its fixed W8A8 tolerance. + w1 = torch.randn( + (e, 2 * n, k), device=current_platform.device_type, dtype=torch.float16 + ).div_(100) + w2 = torch.randn( + (e, k, n), device=current_platform.device_type, dtype=torch.float16 + ).div_(100) n_b_scales = 2 * n k_b_scales = k From 53f6dd5c6f7725df4e5ac9441569860023100870 Mon Sep 17 00:00:00 2001 From: fxmarty-amd <felmarty@amd.com> Date: Mon, 27 Jul 2026 23:47:07 +0200 Subject: [PATCH 132/185] [CI][ROCm] Fix `test_ocp_mx_wikitext_correctness` reference value (#49690) Signed-off-by: Felix Marty <Felix.Marty@amd.com> --- tests/quantization/test_quark.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/quantization/test_quark.py b/tests/quantization/test_quark.py index 009559c22a9..d2f9dcb8af1 100644 --- a/tests/quantization/test_quark.py +++ b/tests/quantization/test_quark.py @@ -279,7 +279,7 @@ WIKITEXT_ACCURACY_CONFIGS = [ excepted_value=10.6, ), AccuracyTestConfig( - model_name="fxmarty/qwen_1.5-moe-a2.7b-mxfp4", excepted_value=12.4 + model_name="fxmarty/qwen_1.5-moe-a2.7b-mxfp4", excepted_value=12.45 ), ] From 28158b2fc3640fa18063b85b3357f41ea9919d7d Mon Sep 17 00:00:00 2001 From: Colin Z <59755453+ColinZ22@users.noreply.github.com> Date: Mon, 27 Jul 2026 15:38:22 -0700 Subject: [PATCH 133/185] [ROCm] [BugFix] Fix Quark GLM-5.2 Checkpoint inference: indexer wk per-channel FP8 dequant + missing sparse-MLA metadata fields (#48886) Signed-off-by: Colin Zeng <Colin.Zeng@amd.com> Signed-off-by: ColinZ22 <Colin.Zeng@amd.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com> Co-authored-by: fanxingran <xingran.fan@amd.com> --- ...est_rocm_aiter_mla_sparse_metadata_sync.py | 75 +++++++++++++++++++ .../layers/attention/mla_attention.py | 45 ++++++----- vllm/model_executor/models/deepseek_v2.py | 9 ++- vllm/v1/attention/backend.py | 4 + .../backends/mla/rocm_aiter_mla_sparse.py | 17 +++++ 5 files changed, 129 insertions(+), 21 deletions(-) diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py index cc9ac6b8d71..aa603626414 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py @@ -88,6 +88,81 @@ def _make_common_metadata(): ) +def _make_mixed_common_metadata(): + # req0: query_len 1 (decode), req1: query_len 4 (prefill) -> decode-first + query_start_loc = torch.tensor([0, 1, 5], dtype=torch.int32, device="cpu") + seq_lens = torch.tensor([16, 10], dtype=torch.int32, device="cpu") + return CommonAttentionMetadata( + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc, + seq_lens=seq_lens, + _seq_lens_cpu=seq_lens, + num_reqs=2, + num_actual_tokens=5, + max_query_len=4, + max_seq_len=16, + block_table_tensor=torch.arange(16, dtype=torch.int32, device="cpu").view(2, 8), + slot_mapping=torch.arange(5, dtype=torch.int64, device="cpu"), + ) + + +def _patch_build_deps(monkeypatch, events=None): + """Stub the aiter kernel, triton helper and CUDA sync so ``build()`` runs + on CPU.""" + + def fake_generate_sparse_seqlen_triton( + query_lens, seq_lens, cu_query_lens, topk_token, num_tokens, max_query_len + ): + return torch.zeros(num_tokens, dtype=torch.int32, device="cpu") + + fake_aiter = _FakeAiter("aiter") + fake_aiter.get_mla_metadata_v1 = Mock(side_effect=lambda *a, **k: None) + monkeypatch.setitem(sys.modules, "aiter", fake_aiter) + monkeypatch.setattr( + sparse_mod, "generate_sparse_seqlen_triton", fake_generate_sparse_seqlen_triton + ) + monkeypatch.setattr( + sparse_mod.torch.cuda, + "current_stream", + lambda device=None: SimpleNamespace( + synchronize=lambda: events.append("sync") if events is not None else None + ), + ) + + +def test_build_populates_decode_only_split_fields(monkeypatch): + """Decode-only batch: all reqs count as decodes, prefill fields default.""" + builder = _make_builder() + _patch_build_deps(monkeypatch) + + md = builder.build( + common_prefix_len=0, common_attn_metadata=_make_common_metadata() + ) + + assert md.num_decodes == 2 + assert md.num_prefills == 0 + assert md.num_decode_tokens == 2 + assert md.prefill_max_seq_len == 0 + assert md.prefill is None + + +def test_build_populates_mixed_split_fields(monkeypatch): + """Mixed decode+prefill batch: split is reported, prefill fields stay + default because this impl always runs the MQA path.""" + builder = _make_builder() + _patch_build_deps(monkeypatch) + + md = builder.build( + common_prefix_len=0, common_attn_metadata=_make_mixed_common_metadata() + ) + + assert md.num_decodes == 1 + assert md.num_prefills == 1 + assert md.num_decode_tokens == 1 + assert md.prefill_max_seq_len == 0 + assert md.prefill is None + + def test_sparse_persistent_metadata_syncs_only_after_recompute(monkeypatch): builder = _make_builder() common_metadata = _make_common_metadata() diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 884c4ca4e68..1dea276d2d2 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -515,29 +515,36 @@ class MLAAttention(nn.Module, AttentionLayerBase): compilation_config.static_forward_context[prefix] = self self.prefill_backend: MLAPrefillBackend | None - try: - prefill_backend_cls = get_mla_prefill_backend(vllm_config) - except ValueError: - if ( - not self.impl.is_sparse - or vllm_config.attention_config.mla_prefill_backend is not None - ): - raise + if self.impl.is_sparse and not self.impl.supports_dense_mha_prefill: logger.warning_once( - "No MLA prefill backend supports this model; sparse MLA will use the " - "top-k MQA path only (no dense-MHA prefill)." + "Sparse MLA impl has no dense-MHA prefill path; using the top-k " + "MQA path only." ) self.prefill_backend = None else: - self.prefill_backend = prefill_backend_cls( - num_heads=self.num_heads, - scale=self.scale, - kv_lora_rank=self.kv_lora_rank, - qk_nope_head_dim=self.qk_nope_head_dim, - qk_rope_head_dim=self.qk_rope_head_dim, - v_head_dim=self.v_head_dim, - vllm_config=vllm_config, - ) + try: + prefill_backend_cls = get_mla_prefill_backend(vllm_config) + except ValueError: + if ( + not self.impl.is_sparse + or vllm_config.attention_config.mla_prefill_backend is not None + ): + raise + logger.warning_once( + "No MLA prefill backend supports this model; sparse MLA will " + "use the top-k MQA path only (no dense-MHA prefill)." + ) + self.prefill_backend = None + else: + self.prefill_backend = prefill_backend_cls( + num_heads=self.num_heads, + scale=self.scale, + kv_lora_rank=self.kv_lora_rank, + qk_nope_head_dim=self.qk_nope_head_dim, + qk_rope_head_dim=self.qk_rope_head_dim, + v_head_dim=self.v_head_dim, + vllm_config=vllm_config, + ) self.kv_cache = torch.tensor([]) diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 2e30fd4aa2f..4b92e351caa 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -851,11 +851,16 @@ def _try_load_fp8_indexer_wk( # We have both weight and scale: dequantize FP8 to BF16. weight_fp8, scale_inv = entry["weight"], entry["scale"] del buf[layer_prefix] - block_size = weight_fp8.shape[1] // scale_inv.shape[1] + if scale_inv.ndim == 1: + # Per-channel scale: one scale per row of [out, in] + group_shape = GroupShape(1, weight_fp8.shape[1]) + else: + block_size = weight_fp8.shape[1] // scale_inv.shape[1] + group_shape = GroupShape(block_size, block_size) weight_bf16 = scaled_dequantize( weight_fp8, scale_inv, - group_shape=GroupShape(block_size, block_size), + group_shape=group_shape, out_dtype=torch.bfloat16, ) diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index b479aab5816..b5d8962fd40 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -803,6 +803,10 @@ class AttentionImplBase(ABC, Generic[T]): # route between the dense-MHA prefill and sparse-MQA paths. is_sparse: ClassVar[bool] = False + # Whether this impl provides a dense-MHA prefill path (forward_mha). Sparse + # impls without one run the top-k MQA path for all requests. + supports_dense_mha_prefill: ClassVar[bool] = True + # Required attributes that all impls should have num_heads: int head_size: int diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index 2bbaaad7654..d5db42cb7f9 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -30,6 +30,7 @@ from vllm.v1.attention.backend import ( from vllm.v1.attention.backends.mla.rocm_aiter_mla import ( AiterMLAHelper, ) +from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.kv_cache_interface import AttentionSpec from vllm.v1.worker.workspace import current_workspace_manager @@ -333,6 +334,14 @@ class ROCMAiterMLASparseMetadata(AttentionMetadata): block_size: int = 1 topk_tokens: int = 2048 + # Fields read by the shared MLA forward. This impl has no dense-MHA prefill + # path (supports_dense_mha_prefill=False), so it always runs the MQA path; + num_decodes: int = 0 + num_prefills: int = 0 + num_decode_tokens: int = 0 + prefill_max_seq_len: int = 0 + prefill: object = None + # Persistent MLA metadata (only populated when persistent mode is enabled, # i.e. when the aiter sparse decode kernel supports work-stealing splits). work_meta_data: torch.Tensor | None = None @@ -488,6 +497,10 @@ class ROCMAiterMLASparseMetadataBuilder( fast_build: bool = False, ) -> ROCMAiterMLASparseMetadata: num_tokens = common_attn_metadata.num_actual_tokens + (num_decodes, num_prefills, num_decode_tokens, _) = split_decodes_and_prefills( + common_attn_metadata, + decode_threshold=self.reorder_batch_threshold or 1, + ) starts = np.asarray(common_attn_metadata.query_start_loc_cpu, dtype=np.int32) seg_lengths = np.diff(starts) req_id_per_token = np.repeat( @@ -604,6 +617,9 @@ class ROCMAiterMLASparseMetadataBuilder( block_size=self.kv_cache_spec.block_size, attn_out_dtype=self.model_dtype, topk_tokens=self.topk_tokens, + num_decodes=num_decodes, + num_prefills=num_prefills, + num_decode_tokens=num_decode_tokens, qo_indptr=qo_indptr, paged_kv_last_page_len=paged_kv_last_page_len, paged_kv_indices=paged_kv_indices, @@ -649,6 +665,7 @@ def reference_mla_sparse_prefill( class ROCMAiterMLASparseImpl(MLAAttentionImpl[ROCMAiterMLASparseMetadata]): is_sparse = True + supports_dense_mha_prefill = False def __init__( self, From ebcef33766f8bcce6344f2bafa245d47097d3c31 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mi=C5=82osz=20Grunwald?= <milosz.grunwald@intel.com> Date: Tue, 28 Jul 2026 00:51:17 +0200 Subject: [PATCH 134/185] Fix MQA with tensor parallelism on transformers modeling backend (#49987) Signed-off-by: microslaw <milosz.grunwald@intel.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- .../models/transformers/fusers/test_linear.py | 163 +++++++++++++++++- .../models/transformers/fuser.py | 7 +- .../models/transformers/fusers/__init__.py | 9 +- .../models/transformers/fusers/base.py | 79 +++++---- .../models/transformers/fusers/packed_qkv.py | 163 ++++++++++++++++++ .../models/transformers/fusers/qkv.py | 7 +- .../models/transformers/fx_utils.py | 12 +- 7 files changed, 396 insertions(+), 44 deletions(-) create mode 100644 vllm/model_executor/models/transformers/fusers/packed_qkv.py diff --git a/tests/models/transformers/fusers/test_linear.py b/tests/models/transformers/fusers/test_linear.py index eff78fa290d..842a1f54794 100644 --- a/tests/models/transformers/fusers/test_linear.py +++ b/tests/models/transformers/fusers/test_linear.py @@ -11,7 +11,11 @@ import torch.nn as nn import torch.nn.functional as F from vllm.model_executor.models.transformers.fuser import get_fuser -from vllm.model_executor.models.transformers.fusers import GLUFuser, QKVFuser +from vllm.model_executor.models.transformers.fusers import ( + GLUFuser, + PackedQKVFuser, + QKVFuser, +) class SiluAndMulStub(nn.Module): @@ -203,6 +207,103 @@ class PerHeadQKNormAttention(FakeAttention): return self.o_proj((q + k + v).flatten(-2)), None +class ResidDropoutAttention(FakeAttention): + """GPT-style dropout after `o_proj` -> the output projection is still found.""" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.resid_dropout = nn.Dropout(0.0) + + def forward( + self, hidden_states, attention_mask=None, past_key_values=None, **kwargs + ): + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, None + ) + attn_output, _ = attention_interface( + self, q, k, v, attention_mask, scaling=self.scaling, **kwargs + ) + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + return self.resid_dropout(self.o_proj(attn_output)), None + + +class PackedQKVAttention(nn.Module): + """GPTBigCode-style: one packed projection split into q/k/v in the forward.""" + + is_causal = True + + def __init__( + self, + hidden: int = 32, + head_dim: int = 8, + heads: int = 4, + kv_heads: int = 1, + bias: bool = False, + layer_idx: int = 0, + ): + super().__init__() + self.config = SimpleNamespace(_attn_implementation="vllm") + self.layer_idx = layer_idx + self.head_dim = head_dim + self.scaling = head_dim**-0.5 + self.embed_dim = heads * head_dim + self.kv_dim = kv_heads * head_dim + self.c_attn = nn.Linear(hidden, self.embed_dim + 2 * self.kv_dim, bias=bias) + self.c_proj = nn.Linear(self.embed_dim, hidden, bias=bias) + self.resid_dropout = nn.Dropout(0.0) + + def forward( + self, hidden_states, attention_mask=None, past_key_values=None, **kwargs + ): + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + + input_shape = hidden_states.shape[:-1] + q, k, v = ( + self.c_attn(hidden_states) + .unsqueeze(1) + .split((self.embed_dim, self.kv_dim, self.kv_dim), dim=3) + ) + q = q.view(*input_shape, -1, self.head_dim).transpose(1, 2) + if past_key_values is not None: + k, v = past_key_values.update(k, v, self.layer_idx) + attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, None + ) + attn_output, attn_weights = attention_interface( + self, q, k, v, attention_mask, scaling=self.scaling, **kwargs + ) + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + return self.resid_dropout(self.c_proj(attn_output)), attn_weights + + +class PerHeadSplitAttention(nn.Module): + """A packed projection reshaped and split *per head* -> not a q/k/v split.""" + + def __init__(self, hidden: int = 32, head_dim: int = 8, heads: int = 4): + super().__init__() + self.head_dim = head_dim + self.heads = heads + self.c_attn = nn.Linear(hidden, 3 * heads * head_dim) + self.c_proj = nn.Linear(heads * head_dim, hidden) + + def forward(self, hidden_states): + shape = (*hidden_states.shape[:2], self.heads, 3 * self.head_dim) + q, k, v = ( + self.c_attn(hidden_states) + .view(shape) + .transpose(1, 2) + .split((self.head_dim, self.head_dim, self.head_dim), dim=3) + ) + return self.c_proj((q + k + v).transpose(1, 2).flatten(-2)) + + class FakeSelfAttn(nn.Module): """Stand-in for the vLLM `Attention` looked up in `attention_instances`.""" @@ -215,6 +316,14 @@ class FakeSelfAttn(nn.Module): return q + 2 * k + 3 * v +class FakeMQASelfAttn(FakeSelfAttn): + """Stand-in for grouped/multi-query layouts, where `k`/`v` are narrower.""" + + def forward(self, q, k, v): + groups = q.shape[-1] // k.shape[-1] + return q + (2 * k + 3 * v).repeat(1, groups) + + @pytest.fixture(autouse=True) def _clear_fuser_cache(): get_fuser.cache_clear() @@ -267,6 +376,15 @@ def _apply_qkv_fuser_with_stubs(module: nn.Module, fuser: QKVFuser): return module +def _apply_packed_qkv_fuser_with_stubs(module: nn.Module, fuser: PackedQKVFuser): + """Apply a fuser at `tp_size == 1`, where the rewritten split is unchanged.""" + qkv = module.get_submodule(fuser.qkv_name) + qkv.output_sizes = [fuser.q_size, fuser.kv_size, fuser.kv_size] + qkv.tp_size = 1 + module.forward = MethodType(fuser.fused_forward, module) + return module + + @pytest.mark.parametrize("mlp_cls", [GLUMLP, ReversedGLUMLP]) @pytest.mark.parametrize("bias", [False, True]) def test_detects_and_rewrites_glu(mlp_cls, bias): @@ -366,6 +484,49 @@ def test_qkv_identifies_output_projection(): # Norm children (q_norm/k_norm) must not disturb o_proj identification. assert get_fuser(QKNormAttention()).o_name == "o_proj" assert get_fuser(PerHeadQKNormAttention()).o_name == "o_proj" + # A module between o_proj and the return is transparent. + assert get_fuser(ResidDropoutAttention()).o_name == "o_proj" + + +@pytest.mark.parametrize("kv_heads", [1, 2]) +def test_detects_and_rewrites_packed_qkv(kv_heads): + """A single projection split into q/k/v must be re-sharded, not merged. + + Only the split sizes change: `QKVParallelLinear` loads the packed + checkpoint weight as-is, and shards q by heads while replicating k/v.""" + with torch.device("meta"): + meta = PackedQKVAttention(kv_heads=kv_heads) + fuser = get_fuser(meta) + assert isinstance(fuser, PackedQKVFuser) + assert (fuser.qkv_name, fuser.o_name) == ("c_attn", "c_proj") + assert (fuser.q_size, fuser.kv_size) == (32, 8 * kv_heads) + + # The hard-coded widths become the per-rank widths of the sharded linear + names = fuser.fused_forward.__code__.co_names + assert "output_sizes" in names and "tp_size" in names + assert "kv_dim" not in names and "embed_dim" not in names + + # Numerics: the rewritten forward must match the original on a real instance + real = PackedQKVAttention(kv_heads=kv_heads, layer_idx=3) + for p in real.parameters(): + nn.init.normal_(p, std=0.05) + x = torch.randn(1, 5, 32) + attention_instances = {3: FakeMQASelfAttn()} + expected, _ = real(x, attention_instances=attention_instances) + fused = _apply_packed_qkv_fuser_with_stubs(real, fuser) + + # Fusion is in place: the module keeps its class and other attributes + assert fused is real and type(fused) is PackedQKVAttention + assert fused.layer_idx == 3 and fused.is_causal + out, _ = fused(x, attention_instances=attention_instances) + torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5) + + +def test_per_head_split_is_not_packed_qkv(): + """The split must consume the whole projection, else its sizes are head + widths and re-sharding by them would be wrong.""" + with torch.device("meta"): + assert get_fuser(PerHeadSplitAttention()) is None def test_fuser_is_cached_per_class_and_structure(): diff --git a/vllm/model_executor/models/transformers/fuser.py b/vllm/model_executor/models/transformers/fuser.py index 13d84a245d9..cbcb24e2a95 100644 --- a/vllm/model_executor/models/transformers/fuser.py +++ b/vllm/model_executor/models/transformers/fuser.py @@ -18,9 +18,10 @@ from vllm.logger import init_logger from vllm.model_executor.models.transformers.fusers import ( BaseFuser, GLUFuser, + PackedQKVFuser, QKVFuser, + RewriteFuser, RMSNormFuser, - StackedFuser, ) from vllm.model_executor.models.transformers.fx_utils import trace @@ -47,9 +48,9 @@ def get_fuser(module: nn.Module) -> BaseFuser | None: return None if (graph := trace(module)) is None: return None - for fuser_cls in (GLUFuser, QKVFuser, RMSNormFuser): + for fuser_cls in (GLUFuser, QKVFuser, PackedQKVFuser, RMSNormFuser): if (fuser := fuser_cls.match(graph, module)) is not None: - if isinstance(fuser, StackedFuser): + if isinstance(fuser, RewriteFuser): try: fuser.update_forward(module) except Exception as exc: diff --git a/vllm/model_executor/models/transformers/fusers/__init__.py b/vllm/model_executor/models/transformers/fusers/__init__.py index 58910b0ecc3..ed4da5155cf 100644 --- a/vllm/model_executor/models/transformers/fusers/__init__.py +++ b/vllm/model_executor/models/transformers/fusers/__init__.py @@ -2,17 +2,24 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Concrete fusers for the Transformers modeling backend.""" -from vllm.model_executor.models.transformers.fusers.base import BaseFuser, StackedFuser +from vllm.model_executor.models.transformers.fusers.base import ( + BaseFuser, + RewriteFuser, + StackedFuser, +) from vllm.model_executor.models.transformers.fusers.glu import GLUFuser from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser +from vllm.model_executor.models.transformers.fusers.packed_qkv import PackedQKVFuser from vllm.model_executor.models.transformers.fusers.qkv import QKVFuser from vllm.model_executor.models.transformers.fusers.rms_norm import RMSNormFuser __all__ = [ "BaseFuser", + "RewriteFuser", "StackedFuser", "GLUFuser", "MoEBlockFuser", + "PackedQKVFuser", "QKVFuser", "RMSNormFuser", ] diff --git a/vllm/model_executor/models/transformers/fusers/base.py b/vllm/model_executor/models/transformers/fusers/base.py index abeff73e3ed..920e63101ac 100644 --- a/vllm/model_executor/models/transformers/fusers/base.py +++ b/vllm/model_executor/models/transformers/fusers/base.py @@ -57,27 +57,61 @@ class BaseFuser(ABC): return {} +def local_output_sizes(merged_name: str) -> str: + """Source for the per-rank widths of the merged linear `self.<merged_name>`.""" + merged = f"self.{merged_name}" + return f"[s // {merged}.tp_size for s in {merged}.output_sizes]" + + @dataclass -class StackedFuser(BaseFuser): - """A fuser that merges sibling projections into one stacked linear and - rewrites the forward to call it. +class RewriteFuser(BaseFuser): + """A fuser that rewrites the module's forward and rebinds it. - `match` and `update_forward` analyse the class once; `fuse` builds the merged - submodule and binds the compiled forward on an instance in place, so it keeps - its class and any attribute the fusion does not consume. + `match` and `update_forward` analyse the class once; `fuse` swaps the + submodules and binds the compiled forward on an instance in place, so it + keeps its class and any attribute the fusion does not consume. """ - merged_name: ClassVar[str] - """Attribute name of the merged module created by `update_attrs`.""" - merged_cls: ClassVar[str] - """Name of the vLLM class the merged projection becomes (for logging).""" - source_cls: str """Class of the HF module the fused projections belonged to (for logging).""" fused_forward: Callable = field(init=False, repr=False) """The compiled rewritten forward, set by `update_forward`.""" + @abstractmethod + def update_forward(self, module: nn.Module) -> None: + """Rewrite and compile `type(module)`'s forward source. + + Raises if the source does not admit the rewrite (fusion is then skipped). + """ + + @abstractmethod + def update_attrs( + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" + ) -> None: + """Replace `module`'s submodules with their vLLM equivalents.""" + + def fuse( + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" + ) -> nn.Module: + """Fuse an already-validated `module` in place (see `Fusers.__getitem__`). + + Builds the merged submodule and binds the compiled forward.""" + self.update_attrs(module, prefix, vllm_config) + module.forward = types.MethodType(self.fused_forward, module) + return module + + +@dataclass +class StackedFuser(RewriteFuser): + """A fuser that merges sibling projections into one stacked linear and + rewrites the forward to call it.""" + + merged_name: ClassVar[str] + """Attribute name of the merged module created by `update_attrs`.""" + merged_cls: ClassVar[str] + """Name of the vLLM class the merged projection becomes (for logging).""" + def info(self, name: str) -> str: sources = " + ".join(shard for shard, _ in self.shards) return ( @@ -108,26 +142,3 @@ class StackedFuser(BaseFuser): """`{merged_name: [projection names]}` so quantization can unpack the fused layer into its per-shard configs.""" return {self.merged_name: [name for name, _ in self.shards]} - - @abstractmethod - def update_forward(self, module: nn.Module) -> None: - """Rewrite and compile `type(module)`'s forward source. - - Raises if the source does not admit the rewrite (fusion is then skipped). - """ - - @abstractmethod - def update_attrs( - self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" - ) -> None: - """Replace `module`'s submodules with the merged module.""" - - def fuse( - self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" - ) -> nn.Module: - """Fuse an already-validated `module` in place (see `Fusers.__getitem__`). - - Builds the merged submodule and binds the compiled forward.""" - self.update_attrs(module, prefix, vllm_config) - module.forward = types.MethodType(self.fused_forward, module) - return module diff --git a/vllm/model_executor/models/transformers/fusers/packed_qkv.py b/vllm/model_executor/models/transformers/fusers/packed_qkv.py new file mode 100644 index 00000000000..ae5612fa816 --- /dev/null +++ b/vllm/model_executor/models/transformers/fusers/packed_qkv.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Packed-QKV fuser: `c_attn(x).split((q, kv, kv))` -> a `QKVParallelLinear`.""" + +import ast +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from torch import fx, nn + +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import QKVParallelLinear +from vllm.model_executor.models.transformers.fusers.base import ( + RewriteFuser, + local_output_sizes, +) +from vllm.model_executor.models.transformers.fx_utils import ( + compile_forward, + is_method, + recover_forward, + returned_linear, + upstream_linear, +) +from vllm.model_executor.models.transformers.utils import ( + log_replacement, + replace_linear_class, +) +from vllm.model_executor.models.utils import maybe_prefix + +if TYPE_CHECKING: + from vllm.config import VllmConfig + +logger = init_logger(__name__) + + +@dataclass +class PackedQKVFuser(RewriteFuser): + """Fuser for attention with q, k and v packed into one projection.""" + + qkv_name: str + o_name: str | None + q_size: int + kv_size: int + + def info(self, name: str) -> str: + return ( + f"Fused: {self.qkv_name} ({name}: {self.source_cls}) -> QKVParallelLinear" + ) + + @staticmethod + def _packed_sizes(node: fx.Node) -> tuple[int, int] | None: + """`(q, kv)` from a `split((q, kv, kv), ...)` call, if it is one.""" + if not is_method(node, "split") or len(node.args) < 2: + return None + sizes = node.args[1] + if not isinstance(sizes, (tuple, list)) or len(sizes) != 3: + return None + if not all(isinstance(size, int) for size in sizes): + return None + q_size, k_size, v_size = sizes + if k_size != v_size or q_size < k_size: + return None + return q_size, k_size + + @classmethod + def match(cls, graph: fx.Graph, module: nn.Module) -> "PackedQKVFuser | None": + for node in graph.nodes: + if (sizes := cls._packed_sizes(node)) is None: + continue + q_size, kv_size = sizes + qkv_node = upstream_linear(node.args[0], module) + if qkv_node is None: + continue + qkv_name = str(qkv_node.target) + # The split must consume the whole projection. + if module.get_submodule(qkv_name).out_features != q_size + 2 * kv_size: + continue + # o_proj produces the module's output and consumes the query width. + o_name = returned_linear(graph, module) + if o_name == qkv_name or ( + o_name is not None + and module.get_submodule(o_name).in_features != q_size + ): + o_name = None + return cls( + source_cls=type(module).__name__, + qkv_name=qkv_name, + o_name=o_name, + q_size=q_size, + kv_size=kv_size, + ) + return None + + def _split_call(self, funcdef: ast.FunctionDef) -> ast.Call: + """The unique `self.<qkv_name>(...)....split((a, b, c), ...)` call.""" + calls = [ + node + for node in ast.walk(funcdef) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "split" + and node.args + and isinstance(node.args[0], (ast.Tuple, ast.List)) + and len(node.args[0].elts) == 3 + and any( + isinstance(inner, ast.Attribute) and inner.attr == self.qkv_name + for inner in ast.walk(node.func.value) + ) + ] + if len(calls) != 1: + raise ValueError(f"{self.qkv_name} has {len(calls)} three-way splits") + return calls[0] + + def update_forward(self, module: nn.Module) -> None: + """Rewrite the split sizes to the sharded projection's per-rank widths.""" + funcdef, fn = recover_forward(type(module)) + split = self._split_call(funcdef) + # (q, kv, kv) -> [s // qkv.tp_size for s in qkv.output_sizes] + sections = local_output_sizes(self.qkv_name) + split.args[0] = ast.parse(sections, mode="eval").body + self.fused_forward = compile_forward(funcdef, fn) + + def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool: + """Shapes must be compatible with a head-sharded packed GEMM.""" + head_size = vllm_config.model_config.get_head_size() + qkv = module.get_submodule(self.qkv_name) + compatible = ( + self.q_size % head_size == 0 + and self.kv_size % head_size == 0 + and qkv.out_features == self.q_size + 2 * self.kv_size + ) + if not compatible: + logger.debug("%s is not compatible with packed QKV fusion", type(module)) + return compatible + + def update_attrs( + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" + ) -> None: + quant_config = vllm_config.quant_config + head_size = vllm_config.model_config.get_head_size() + qkv_prefix = maybe_prefix(prefix, self.qkv_name) + qkv = module.get_submodule(self.qkv_name) + merged = QKVParallelLinear( + hidden_size=qkv.in_features, + head_size=head_size, + total_num_heads=self.q_size // head_size, + total_num_kv_heads=self.kv_size // head_size, + bias=qkv.bias is not None, + quant_config=quant_config, + prefix=qkv_prefix, + return_bias=False, + ) + setattr(module, self.qkv_name, merged) + log_replacement(qkv_prefix, qkv, merged) + # If there is an output projection, we know it must be rowwise. + if self.o_name is not None: + o_prefix = maybe_prefix(prefix, self.o_name) + o_proj = module.get_submodule(self.o_name) + new_o = replace_linear_class( + o_proj, "rowwise", quant_config, prefix=o_prefix + ) + setattr(module, self.o_name, new_o) + log_replacement(o_prefix, o_proj, new_o) diff --git a/vllm/model_executor/models/transformers/fusers/qkv.py b/vllm/model_executor/models/transformers/fusers/qkv.py index b9fe9771fe1..370c0cdf4fa 100644 --- a/vllm/model_executor/models/transformers/fusers/qkv.py +++ b/vllm/model_executor/models/transformers/fusers/qkv.py @@ -10,7 +10,10 @@ from torch import fx, nn from vllm.logger import init_logger from vllm.model_executor.layers.linear import QKVParallelLinear -from vllm.model_executor.models.transformers.fusers.base import StackedFuser +from vllm.model_executor.models.transformers.fusers.base import ( + StackedFuser, + local_output_sizes, +) from vllm.model_executor.models.transformers.fx_utils import ( compile_forward, innermost_block, @@ -134,7 +137,7 @@ class QKVFuser(StackedFuser): if names & set(temps): raise ValueError("fused temporaries would shadow existing names") merged = f"self.{self.merged_name}" - sections = f"[s // {merged}.tp_size for s in {merged}.output_sizes]" + sections = local_output_sizes(self.merged_name) template = f"{', '.join(temps)} = {merged}(__arg__).split({sections}, -1)" assign = ast.parse(template).body[0] arg = next( diff --git a/vllm/model_executor/models/transformers/fx_utils.py b/vllm/model_executor/models/transformers/fx_utils.py index 2dbc3a499de..d665bcfe8cf 100644 --- a/vllm/model_executor/models/transformers/fx_utils.py +++ b/vllm/model_executor/models/transformers/fx_utils.py @@ -394,8 +394,10 @@ def output_value(graph: fx.Graph) -> object | None: def upstream_linear(node: object, module: nn.Module) -> fx.Node | None: """Nearest linear producing `node`, walking back through splits/reshapes. - Never walks through a leaf call (e.g. an attention interface): its inputs - are what attention consumes, not what produced the value.""" + Non-linear submodules are transparent too (e.g. the dropout GPT-style + attentions apply after their output projection). Never walks through a leaf + call (e.g. an attention interface): its inputs are what attention consumes, + not what produced the value.""" stack = [node] seen: set[fx.Node] = set() while stack: @@ -405,7 +407,11 @@ def upstream_linear(node: object, module: nn.Module) -> fx.Node | None: seen.add(current) if is_linear(current, module): return current - if current.op in ("call_function", "call_method") and not is_leaf_call(current): + if current.op in ( + "call_function", + "call_method", + "call_module", + ) and not is_leaf_call(current): stack.extend(current.args) return None From 1e34a135390b4fa3370cf55885b442667c3b8ba6 Mon Sep 17 00:00:00 2001 From: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Date: Tue, 28 Jul 2026 01:56:44 +0300 Subject: [PATCH 135/185] Fix Humming non-gated MoE (#49096) --- tests/kernels/moe/test_moe.py | 96 +++++++++++++++++++ .../fused_moe/experts/fused_humming_moe.py | 14 ++- .../quantization/utils/humming_utils.py | 8 +- 3 files changed, 111 insertions(+), 7 deletions(-) diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index 7fc0e777e60..9c43aa97409 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -1165,6 +1165,102 @@ def test_fused_marlin_moe_non_gated( torch.testing.assert_close(marlin_output, torch_output, atol=1e-1, rtol=0) +@pytest.mark.parametrize( + "activation", + [ + MoEActivation.SILU, + MoEActivation.RELU2_NO_MUL, + ], + ids=["gated", "non_gated"], +) +def test_humming_gated_non_gated_shape_contract(activation: MoEActivation): + pytest.importorskip("humming") + from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( + HummingIndexedExperts, + ) + from vllm.model_executor.layers.quantization.utils import humming_utils + from vllm.utils import humming + + top_k, num_experts = 6, 12 + hidden_size, intermediate_size = 2688, 1856 + gate_up_size = intermediate_size * 2 if activation.is_gated else intermediate_size + num_w13_stacks = 2 if activation.is_gated else 1 + moe_config = make_dummy_moe_config( + num_experts=num_experts, + experts_per_token=top_k, + hidden_dim=hidden_size, + intermediate_size=intermediate_size, + activation=activation, + ) + + layer = torch.nn.Module() + layer.moe_config = moe_config + layer.params_dtype = torch.bfloat16 + + weight_schema = humming.ModeloptNvfp4WeightSchema() + for sublayer_name, shape_n, shape_k, stack_size in ( + ("w13", gate_up_size, hidden_size, num_w13_stacks), + ("w2", hidden_size, intermediate_size, 1), + ): + tensor_attrs = weight_schema.get_tensors_attrs( + shape_n=shape_n, + shape_k=shape_k, + param_dtype=layer.params_dtype, + num_experts=num_experts, + stack_size=stack_size, + ) + for tensor_name, attrs in tensor_attrs.items(): + layer.register_parameter( + f"{sublayer_name}_{tensor_name}", + Parameter( + torch.ones( + attrs["shape"], + dtype=attrs["dtype"], + device="cuda", + ), + requires_grad=False, + ), + ) + + humming_utils.convert_to_humming_moe_kernel_format( + layer, + weight_schema=weight_schema, + input_schema=humming.HummingInputSchema(a_dtype=humming.dtypes.bfloat16), + ) + + w13_meta, w2_meta = (layer.humming_metas[name] for name in ("w13", "w2")) + for meta in (w13_meta, w2_meta): + assert meta.a_dtype == humming.dtypes.bfloat16 + assert meta.b_dtype == humming.dtypes.float4e2m1 + + assert w13_meta.shape_n - w13_meta.pad_shape_n == gate_up_size + assert w2_meta.shape_k - w2_meta.pad_shape_k == intermediate_size + + layer.local_num_experts = layer.global_num_experts = num_experts + layer.hidden_size = hidden_size + layer.intermediate_size_per_partition = intermediate_size + quant_config = humming_utils.get_humming_moe_quant_config(layer) + experts = HummingIndexedExperts( + layer, + moe_config, + quant_config, + ) + + buffer_metas, _ = experts.get_buffer_metas( + M=1, + topk=top_k, + activation=moe_config.activation, + ) + assert buffer_metas["gate_up_output"]["shape"][-1] == gate_up_size + assert buffer_metas["activation_output"]["shape"][-1] == intermediate_size + assert experts.moe_problem_size( + a1=torch.empty(1, hidden_size), + w1=torch.empty(num_experts, 1), + w2=torch.empty(num_experts, 1), + topk_ids=torch.empty(1, top_k, dtype=torch.long), + ) == (num_experts, 1, intermediate_size, hidden_size, top_k) + + @pytest.mark.parametrize("ep_size", [1, 2]) def test_moe_align_block_size_opcheck(ep_size): num_experts = 4 diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py index 5f112380cf9..4cc11317338 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py @@ -289,7 +289,14 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): assert a1.size(0) == num_experts num_tokens = a1.size(1) - return meta1.num_experts, num_tokens, meta1.shape_n // 2, meta1.shape_k, top_k + return ( + meta1.num_experts, + num_tokens, + # Logical intermediate width for both gated and non-gated activations + self.layer.intermediate_size_per_partition, + meta1.shape_k, + top_k, + ) def get_buffer_metas(self, M: int, topk: int, activation: MoEActivation): from vllm.utils.humming import GemmType as HummingGemmType @@ -327,7 +334,8 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): real_shape_m = M * topk output_shape = (M, K) - down_input_size = N if activation.is_gated else (N * 2) + gate_up_size = N * (2 if activation.is_gated else 1) + down_input_size = N a_dtype = self.layer.humming_metas["w13"].a_dtype c_dtype = self.layer.humming_metas["w13"].c_dtype num_bits = a_dtype.num_bits @@ -347,7 +355,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): "dtype": torch_dtype_map[a_dtype], }, "gate_up_output": { - "shape": (real_shape_m, N * 2), + "shape": (real_shape_m, gate_up_size), "dtype": torch_dtype_map[c_dtype], }, "activation_output": { diff --git a/vllm/model_executor/layers/quantization/utils/humming_utils.py b/vllm/model_executor/layers/quantization/utils/humming_utils.py index e08e80ec4bb..15c6a178b05 100644 --- a/vllm/model_executor/layers/quantization/utils/humming_utils.py +++ b/vllm/model_executor/layers/quantization/utils/humming_utils.py @@ -788,7 +788,7 @@ def _convert_sublayer_to_humming( shape_k_stacks = [shape_k] shape_n_stacks = [shape_n] - if sublayer_name == "w13": + if sublayer_name == "w13" and layer.moe_config.activation.is_gated: shape_n_stacks = [shape_n // 2] * 2 converted_weight_schema, converted_tensors = weight_schema.convert_humming( @@ -991,15 +991,15 @@ def convert_to_humming_moe_kernel_format( # Build sublayer configs from layer properties if not provided if sublayer_configs is None: is_gated = layer.moe_config.activation.is_gated + intermediate_size = layer.moe_config.intermediate_size_per_partition sublayer_configs = { "w13": { - "shape_n": layer.moe_config.intermediate_size_per_partition * 2, + "shape_n": intermediate_size * (2 if is_gated else 1), "shape_k": layer.moe_config.hidden_dim, }, "w2": { "shape_n": layer.moe_config.hidden_dim, - "shape_k": layer.moe_config.intermediate_size_per_partition - * (1 if is_gated else 2), + "shape_k": intermediate_size, }, } From 272abd5f486967f1fb9db7ca7504f8c34235ef50 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Mon, 27 Jul 2026 16:15:54 -0700 Subject: [PATCH 136/185] [Tests][Spec Decode] Add gemma4 MTP acceptance rates test (#47920) Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai> --- .buildkite/test_areas/spec_decode.yaml | 12 ++ tests/v1/e2e/spec_decode/test_spec_decode.py | 115 +++++++++++------- vllm/v1/spec_decode/gemma4.py | 24 ++++ .../spec_decode/autoregressive/speculator.py | 17 ++- 4 files changed, 120 insertions(+), 48 deletions(-) diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index 977562ee97e..c63aaa18d8b 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -188,3 +188,15 @@ steps: - tests/v1/e2e/spec_decode/test_mtp_parallel_load.py commands: - pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py + +- label: Spec Decode Acceptance Rates Nightly + key: spec-decode-acceptance-rates-nightly + timeout_in_minutes: 60 + device: h200_35gb + optional: true + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - tests/v1/e2e/spec_decode/ + commands: + - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "acceptance_rates" diff --git a/tests/v1/e2e/spec_decode/test_spec_decode.py b/tests/v1/e2e/spec_decode/test_spec_decode.py index 01bec69640b..17c48721f80 100644 --- a/tests/v1/e2e/spec_decode/test_spec_decode.py +++ b/tests/v1/e2e/spec_decode/test_spec_decode.py @@ -1393,51 +1393,81 @@ def load_and_process_dataset(data_name: str): return dataset -@pytest.fixture -def dflash_config(): - target_model = "Qwen/Qwen3-8B" - draft_model = "z-lab/Qwen3-8B-DFlash-b16" - - return dict( - model=target_model, - trust_remote_code=True, - speculative_config={ - "method": "dflash", - "model": draft_model, - "num_speculative_tokens": 16, - "max_model_len": 32768, - }, - max_model_len=32768, - max_num_seqs=128, - gpu_memory_utilization=0.85, - enforce_eager=False, - disable_log_stats=False, - ) - - +@pytest.mark.parametrize( + ["spec_config", "expected_acceptance_lengths", "chat_template_kwargs"], + [ + pytest.param( + dict( + model="Qwen/Qwen3-8B", + trust_remote_code=True, + speculative_config={ + "method": "dflash", + "model": "z-lab/Qwen3-8B-DFlash-b16", + "num_speculative_tokens": 16, + "max_model_len": 32768, + }, + max_model_len=32768, + max_num_seqs=128, + gpu_memory_utilization=0.85, + enforce_eager=False, + disable_log_stats=False, + ), + # All scores from Table 1 in https://arxiv.org/pdf/2602.06036 + { + "mt-bench": 4.24, + "humaneval": 6.50, + # runs with a subset of prompts so extra wide tol here + "gsm8k": 6.54 * 0.975, + }, + {"enable_thinking": False}, + id="dflash", + ), + pytest.param( + dict( + model="google/gemma-4-E4B-it", + trust_remote_code=True, + speculative_config={ + "method": "mtp", + "model": "google/gemma-4-E4B-it-assistant", + "num_speculative_tokens": 2, + "max_model_len": 32768, + }, + max_model_len=32768, + # Skip multimodal profiling; this is a text-only eval. + limit_mm_per_prompt={"image": 0, "audio": 0}, + disable_log_stats=False, + ), + { + "mt-bench": 2.28, + "humaneval": 2.68, + # runs with a subset of prompts so extra wide tol here + "gsm8k": 2.67 * 0.975, + }, + {}, + id="gemma4", + ), + ], +) @pytest.mark.parametrize("use_mrv2", [False, True]) -def test_dflash_acceptance_rates( - monkeypatch: pytest.MonkeyPatch, use_mrv2: bool, dflash_config +def test_acceptance_rates( + monkeypatch: pytest.MonkeyPatch, + spec_config: dict[str, Any], + expected_acceptance_lengths: dict[str, float], + chat_template_kwargs: dict[str, Any], + use_mrv2: bool, ): """ - E2E test for DFlash (block diffusion) speculative decoding. - Runs acceptance rate validation on GSM8k, MT-Bench, and HumanEval - comparing against baseline results from the paper (Table 1). - See https://github.com/z-lab/dflash/blob/main/benchmark_sglang.py for methodology. + E2E acceptance-rate validation for speculative decoding. + + Drives one or more datasets (keyed in ``expected_acceptance_lengths``) + through the spec decode engine and asserts the mean acceptance length + stays within tolerance of the reference figure for each dataset. """ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_mrv2 else "0") - - spec_llm = LLM(**dflash_config) + spec_llm = LLM(**spec_config) max_prompts_per_dataset = 200 # mt-bench has 80, humaneval has 164, truncates gsm8k - # All scores from Table 1 in https://arxiv.org/pdf/2602.06036 - expected_acceptance_lengths = { - "mt-bench": 4.24, - "humaneval": 6.50, - "gsm8k": 6.54 * 0.975, # runs with a subset of prompts so extra wide tol here - } - tokenizer = spec_llm.get_tokenizer() for dataset_name, expected_len in expected_acceptance_lengths.items(): dataset = load_and_process_dataset(dataset_name) @@ -1452,10 +1482,11 @@ def test_dflash_acceptance_rates( [{"role": "user", "content": user_content}], tokenize=False, add_generation_prompt=True, - enable_thinking=False, + **chat_template_kwargs, ) - # Temp=0, MaxTokens=2048 from the paper + # Greedy (temp=0) so acceptance length is deterministic and comparable + # across runs. spec_llm.generate( [prompt_text], SamplingParams(temperature=0, max_tokens=2048), @@ -1467,17 +1498,17 @@ def test_dflash_acceptance_rates( acceptance_lengths.append(acceptance_len) mean_acceptance_length = sum(acceptance_lengths) / len(acceptance_lengths) - # Fairly tight tolerance of 95% against the paper's figures, + # Fairly tight tolerance of 95% against the reference figures, # watching for regressions. Can be relaxed if test is flaky but be sure to # check for genuine issues such as #40727. expected_len = expected_len * 0.95 print( - f"DFlash acceptance_len for {dataset_name}: {mean_acceptance_length:.2f}" + f"acceptance_len for {dataset_name}: {mean_acceptance_length:.2f}" f" (expected at least {expected_len:.2f})" ) assert mean_acceptance_length >= expected_len, ( - f"DFlash acceptance_len for {dataset_name} is below expected threshold:" + f"acceptance_len for {dataset_name} is below expected threshold: " f"{mean_acceptance_length:.2f} < {expected_len:.2f}" ) diff --git a/vllm/v1/spec_decode/gemma4.py b/vllm/v1/spec_decode/gemma4.py index 7f67ae9f499..07cfaa2f4db 100644 --- a/vllm/v1/spec_decode/gemma4.py +++ b/vllm/v1/spec_decode/gemma4.py @@ -14,6 +14,7 @@ import torch import torch.nn as nn from vllm.config import VllmConfig, get_layers_from_vllm_config, replace +from vllm.distributed.parallel_state import get_pp_group from vllm.logger import init_logger from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.v1.attention.backend import CommonAttentionMetadata @@ -165,6 +166,29 @@ class Gemma4Proposer(SpecDecodeBaseProposer): ) return base + def _maybe_share_embeddings(self, target_language_model: nn.Module) -> None: + """Gemma4 MTP requires dim-mismatched embedding sharing. + + The draft checkpoint's embed_tokens is a draft-dim placeholder + (tied to lm_head so load_weights populates both); the model + expects it to be replaced by the target's backbone-dim embedding, + so the base class's embedding-dim equality guard must not apply. + """ + if get_pp_group().world_size != 1: + return + inner_model = getattr(target_language_model, "model", None) + target_embed_tokens = getattr(inner_model, "embed_tokens", None) + if target_embed_tokens is None: + raise AttributeError( + "Target model does not have an 'embed_tokens' attribute" + ) + del self.model.model.embed_tokens + self.model.model.embed_tokens = target_embed_tokens + logger.info( + "Gemma4 MTP: sharing target model's backbone-dim embed_tokens " + "with the draft model." + ) + def _maybe_share_lm_head(self, target_language_model: nn.Module) -> None: """Gemma4 MTP always keeps its own draft-dim lm_head. diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index b890574971c..9035ea38b23 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -446,10 +446,16 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): ) last_hidden_states = last_hidden_states[:num_reqs] + sample_positions = positions + if not self.advance_draft_positions: + # The forward pass holds positions fixed (Q-only, shared target KV), + # but Gumbel sampling still needs the absolute draft position. + sample_positions = positions + self.current_draft_step + # Sample the draft tokens. draft_tokens = self.sample_draft( last_hidden_states, - positions, + sample_positions, idx_mapping, self.temperature, self.seeds, @@ -630,6 +636,9 @@ def _prepare_decode_inputs_kernel( draft_token = tl.load(draft_tokens_ptr + req_idx * draft_tokens_stride) tl.store(input_ids_ptr + req_idx, draft_token) + target_seq_len = tl.load(target_seq_lens_ptr + req_idx) + num_rejected = tl.load(num_rejected_ptr + req_idx) + seq_len = target_seq_len - num_rejected if ADVANCE_DRAFT_POSITIONS: # Compute position and seq_lens. # NOTE(woosuk): To prevent out-of-range access, we clamp these values @@ -637,12 +646,8 @@ def _prepare_decode_inputs_kernel( position = tl.load(positions_ptr + req_idx) position = tl.minimum(position + 1, max_model_len - 1) tl.store(positions_ptr + req_idx, position) - - target_seq_len = tl.load(target_seq_lens_ptr + req_idx) - num_rejected = tl.load(num_rejected_ptr + req_idx) - seq_len = target_seq_len - num_rejected seq_len = tl.minimum(seq_len + 1, max_model_len) - tl.store(seq_lens_ptr + req_idx, seq_len) + tl.store(seq_lens_ptr + req_idx, seq_len) def prepare_decode_inputs( From 60417b4b744c371453eddcf5c8fa0f184418c957 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Date: Mon, 27 Jul 2026 20:03:26 -0400 Subject: [PATCH 137/185] [Core][PCP] Select MRV2 when PCP is enabled (#50034) Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> --- vllm/config/vllm.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index a7f185cad40..82ca8ea26a3 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -557,6 +557,10 @@ class VllmConfig: if use_v2_model_runner is not None: return use_v2_model_runner + # PCP runtime support is implemented only by the V2 model runner. + if self.parallel_config.prefill_context_parallel_size > 1: + return True + # DSpark is implemented only by the V2 GPU model runner, and DeepSeek-V4 # is not otherwise a default-V2 architecture, so force V2 for it. If V2 # is unsupported for the rest of the config, _validate_v2_model_runner @@ -1460,6 +1464,11 @@ class VllmConfig: if self.use_v2_model_runner: self._validate_v2_model_runner() + elif self.parallel_config.prefill_context_parallel_size > 1: + raise ValueError( + "Prefill context parallelism requires Model Runner V2. " + "Remove VLLM_USE_V2_MODEL_RUNNER=0." + ) # Re-compute compile ranges after platform-specific config updates # (e.g., XPU may lower max_num_batched_tokens when MLA is enabled) From 60b3d39cd36c53a698040edbf51406d3febc97a7 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon <woosuk@inferact.ai> Date: Mon, 27 Jul 2026 17:06:28 -0700 Subject: [PATCH 138/185] [Docs] Remove experimental warning for EP (#50057) Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> --- docs/serving/expert_parallel_deployment.md | 3 --- 1 file changed, 3 deletions(-) diff --git a/docs/serving/expert_parallel_deployment.md b/docs/serving/expert_parallel_deployment.md index b348c5dd965..d1dff32fc94 100644 --- a/docs/serving/expert_parallel_deployment.md +++ b/docs/serving/expert_parallel_deployment.md @@ -31,9 +31,6 @@ vLLM provides multiple communication backends for EP. Use `--all2all-backend` to ## Single Node Deployment -!!! warning - EP is an experimental feature. Argument names and default values may change in the future. - ### Configuration Enable EP by setting the `--enable-expert-parallel` flag. The EP size is automatically calculated as: From 1206891822ca8befe421879a4230ef42a3fc93be Mon Sep 17 00:00:00 2001 From: avininjamay8 <Avinash.Paul@amd.com> Date: Tue, 28 Jul 2026 06:56:24 +0530 Subject: [PATCH 139/185] [ROCm][KVConnector][MoRI-IO] Fix WRITE-mode remote-TP rank collapse (#46332 follow-up) (#47764) Signed-off-by: avininjamay8 <avininjamay8@users.noreply.github.com> Signed-off-by: avininjamay8 <Avinash.Paul@amd.com> Co-authored-by: avininjamay8 <avininjamay8@users.noreply.github.com> Co-authored-by: avininjamay8 <avpaul@amd.com> --- .../kv_connector/v1/moriio/moriio_common.py | 10 +++++++++- .../kv_connector/v1/moriio/moriio_connector.py | 7 ++++++- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py index 45df75eeb3c..811bedb0841 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py @@ -477,7 +477,15 @@ class MoRIIOConnectorMetadata(KVConnectorMetadata): remote_port=int(remote_handshake_port), remote_handshake_port=int(remote_handshake_port), remote_notify_port=int(remote_notify_port), - tp_size=kv_transfer_params.get("tp_size", 1), + # Remote peer TP degree (used as remote_tp_size downstream). The + # proxy advertises it under "remote_tp_size"; #46332 read "tp_size" + # which is absent on WRITE producer requests -> defaulted to 1 -> + # rank collapse. Read the right key; 0 == unknown (== homogeneous). + tp_size=int( + kv_transfer_params.get("remote_tp_size") + or kv_transfer_params.get("tp_size") + or 0 + ), remote_dp_size=kv_transfer_params.get("remote_dp_size", 1), remote_dp_rank=kv_transfer_params.get("remote_dp_rank", 0), ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index c478ced2211..2a4c028e15b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -1351,6 +1351,9 @@ class MoRIIOConnectorWorker: return {remote_agent_name} def _remote_tp_rank(self, remote_tp_size: int) -> int: + # 0/unknown remote TP == homogeneous (avoids collapsing all ranks to 0). + if remote_tp_size == 0: + remote_tp_size = self.world_size return get_moriio_remote_tp_rank(self.tp_rank, self.world_size, remote_tp_size) def _background_moriio_handshake( @@ -2145,7 +2148,9 @@ class MoRIIOConnectorWorker: validate_moriio_heterogeneous_tp_kv_heads( local_tp_size=self.world_size, remote_tp_size=( - remote_tp_size if remote_tp_size is not None else self.world_size + remote_tp_size + if remote_tp_size and remote_tp_size > 0 + else self.world_size ), total_num_kv_heads=self.model_config.get_total_num_kv_heads(), is_mla=self._is_mla_cache_layer(layer_name), From 02b6ecf07cf6a3e6dd395f73af5d1904312165cf Mon Sep 17 00:00:00 2001 From: frida-andersson <fanderss@amd.com> Date: Tue, 28 Jul 2026 04:12:20 +0200 Subject: [PATCH 140/185] [ROCm][DSv3.2] Eliminate per-decode FillFunctor launches in sparse-MLA hot loop (#44527) Signed-off-by: Frida Andersson <fanderss@amd.com> --- vllm/v1/attention/ops/rocm_aiter_mla_sparse.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 28ec8d4b592..6a28324208f 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -429,7 +429,6 @@ def rocm_fp8_paged_mqa_logits( (out_logits,) = current_workspace_manager().get_simultaneous( ((batch_size * next_n, max_model_len), torch.float32), ) - out_logits.fill_(float("-inf")) deepgemm_fp8_paged_mqa_logits( q_fp8, kv_cache_fp8, @@ -715,7 +714,6 @@ def rocm_aiter_sparse_attn_indexer( scale_fmt, ) - topk_indices_buffer[: hidden_states.shape[0]] = -1 if has_prefill: prefill_metadata = layer_attn_metadata.prefill assert prefill_metadata is not None From e68bfc28281eb3023a859bd1b2c512acfb9ff924 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas <akaratza@amd.com> Date: Mon, 27 Jul 2026 21:22:53 -0500 Subject: [PATCH 141/185] [CI][ROCm] Soft-fail Python-only installation mirror (#50041) Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> --- .buildkite/test_areas/misc.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index cadbaed5b05..1a53a92961f 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -295,6 +295,7 @@ steps: amd: device: mi250_1 timeout_in_minutes: 55 + soft_fail: true depends_on: - image-build-amd source_file_dependencies: From 73af7a362ac18725814cd161a802e9c428123850 Mon Sep 17 00:00:00 2001 From: Yan Ma <yan.ma@intel.com> Date: Tue, 28 Jul 2026 10:34:33 +0800 Subject: [PATCH 142/185] [XPU] Add online fp8 quantization test (#44513) Signed-off-by: Yan Ma <yan.ma@intel.com> --- .buildkite/intel_jobs/test-intel.yaml | 5 +++-- .../models/multimodal/processing/test_common.py | 1 + tests/quantization/test_online.py | 5 ++++- tests/quantization/utils.py | 8 ++++---- vllm/platforms/xpu.py | 16 ++++++++++++++++ 5 files changed, 28 insertions(+), 7 deletions(-) diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index 5c717c9201c..3fadb07f391 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -145,7 +145,8 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && - pytest -v -s quantization/test_auto_round.py' + pytest -v -s quantization/test_auto_round.py && + pytest -v -s quantization/test_online.py' - label: "XPU compressed tensors FP8 test" depends_on: - image-build-xpu @@ -168,4 +169,4 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && - pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8' \ No newline at end of file + pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8' diff --git a/tests/models/multimodal/processing/test_common.py b/tests/models/multimodal/processing/test_common.py index 9b586be4797..e9cb4ab23c1 100644 --- a/tests/models/multimodal/processing/test_common.py +++ b/tests/models/multimodal/processing/test_common.py @@ -88,6 +88,7 @@ _XPU_EXCLUDED_MODEL_IDS = { "baidu/Unlimited-OCR", "mistralai/Mistral-Large-3-675B-Instruct-2512-NVFP4", "Qwen/Qwen2.5-Omni-7B-AWQ", + "thinkingmachines/Inkling-NVFP4", } diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index 3c21441ed65..d977cff96d0 100644 --- a/tests/quantization/test_online.py +++ b/tests/quantization/test_online.py @@ -89,6 +89,9 @@ def test_online_quantization( if use_rocm_aiter: monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + if current_platform.is_xpu() and quant_scheme == "fp8_per_block": + pytest.skip("Skip test for online fp8_per_block on XPU platform.") + # `LLM.apply_model` requires pickling a function. monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") @@ -121,7 +124,7 @@ def test_online_quantization( if moe is not None: assert isinstance(moe._quant_method, expected_moe_cls) - if current_platform.is_cuda(): + if current_platform.is_cuda() or current_platform.is_xpu(): assert o_proj.weight.dtype == torch.float8_e4m3fn elif current_platform.is_rocm(): assert o_proj.weight.dtype == current_platform.fp8_dtype() diff --git a/tests/quantization/utils.py b/tests/quantization/utils.py index 8ab9c310dca..5e34199f942 100644 --- a/tests/quantization/utils.py +++ b/tests/quantization/utils.py @@ -10,15 +10,15 @@ from vllm.platforms import current_platform def is_quant_method_supported(quant_method: str) -> bool: - # Currently, all quantization methods require Nvidia or AMD GPUs - if not (current_platform.is_cuda() or current_platform.is_rocm()): + # Currently, quantization tests only run GPUs + if current_platform.is_cpu(): return False - try: current_platform.verify_quantization(quant_method) except ValueError: return False - + if current_platform.is_xpu(): + return True capability = current_platform.get_device_capability() assert capability is not None diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index e9d1a905a4f..b9fb8487b72 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -110,6 +110,22 @@ class XPUPlatform(Platform): ray_device_key: str = "GPU" dist_backend: str = "xccl" # xccl only device_control_env_var: str = "ZE_AFFINITY_MASK" + supported_quantization: list[str] = [ + "awq", + "gptq", + "auto_awq", + "auto_gptq", + "inc", + "fp8", + "mxfp4", + "mxfp8", + "fp8_per_tensor", + "fp8_per_block", + "online", + "gpt_oss_mxfp4", + "modelopt", + "compressed-tensors", + ] @classmethod def import_kernels(cls) -> None: From 7aea73d83d6064449ff8147899de161e6eb68a20 Mon Sep 17 00:00:00 2001 From: fxmarty-amd <felmarty@amd.com> Date: Tue, 28 Jul 2026 04:35:57 +0200 Subject: [PATCH 143/185] [ROCm][Quark][6/N] Use MXFP4 linear kernel abstraction for `aiter` backend (#49348) Signed-off-by: Felix Marty <Felix.Marty@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> --- .../evals/gsm8k/configs/Qwen3-1.7B-MXFP4.yaml | 5 + .../test_mxfp4_kernel_selection.py | 130 +++++++++++++ tests/quantization/test_quark.py | 42 +++++ .../model_executor/kernels/linear/__init__.py | 8 + .../kernels/linear/mxfp4/aiter.py | 177 ++++++++++++++++++ .../quark/schemes/quark_ocp_mx.py | 174 +---------------- 6 files changed, 370 insertions(+), 166 deletions(-) create mode 100644 tests/evals/gsm8k/configs/Qwen3-1.7B-MXFP4.yaml create mode 100644 tests/kernels/quantization/test_mxfp4_kernel_selection.py create mode 100644 vllm/model_executor/kernels/linear/mxfp4/aiter.py diff --git a/tests/evals/gsm8k/configs/Qwen3-1.7B-MXFP4.yaml b/tests/evals/gsm8k/configs/Qwen3-1.7B-MXFP4.yaml new file mode 100644 index 00000000000..b0249f41b60 --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3-1.7B-MXFP4.yaml @@ -0,0 +1,5 @@ +model_name: "amd-quark/Qwen3-1.7B-MXFP4" +accuracy_threshold: 0.27 +num_questions: 1319 +num_fewshot: 5 +server_args: "--enforce-eager --max-model-len 4096" diff --git a/tests/kernels/quantization/test_mxfp4_kernel_selection.py b/tests/kernels/quantization/test_mxfp4_kernel_selection.py new file mode 100644 index 00000000000..20c171a87ce --- /dev/null +++ b/tests/kernels/quantization/test_mxfp4_kernel_selection.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for MXFP4 linear kernel selection logic (CPU-only) + +Run `pytest tests/kernels/quantization/test_mxfp4_kernel_selection.py`. +""" + +from unittest.mock import patch + +import pytest +import torch + +from vllm.model_executor.kernels.linear import ( + AiterMxfp4LinearKernel, + MxFp4LinearKernel, + MxFp4LinearLayerConfig, + init_mxfp4_linear_kernel, + register_linear_kernel, +) +from vllm.platforms import PlatformEnum + +pytestmark = pytest.mark.cpu_test + + +def test_can_implement_is_abstract(): + """Test that can_implement()/is_supported() are properly defined.""" + assert hasattr(MxFp4LinearKernel, "can_implement") + assert hasattr(MxFp4LinearKernel, "is_supported") + + +def test_aiter_kernel_is_supported_requires_native_mx_support(): + """AiterMxfp4LinearKernel must not be selected on platforms without + native MX compute, even if AITER itself is importable.""" + with patch( + "vllm.model_executor.kernels.linear.mxfp4.aiter.current_platform.supports_mx", + return_value=False, + ): + is_supported, reason = AiterMxfp4LinearKernel.is_supported() + assert not is_supported + assert reason + + +class OOTMxFp4LinearKernel(MxFp4LinearKernel): + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + return True, None + + @classmethod + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + pass + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + pass + + +@patch("vllm.model_executor.kernels.linear.current_platform") +def test_init_mxfp4_linear_kernel_dispatches_to_registered_kernel(platform_mock): + """init_mxfp4_linear_kernel should select a registered kernel that + reports itself as supported, and construct it with a fresh config.""" + platform_mock._enum = PlatformEnum.OOT + register_linear_kernel(OOTMxFp4LinearKernel, PlatformEnum.OOT, "mxfp4") + + kernel = init_mxfp4_linear_kernel() + + assert isinstance(kernel, OOTMxFp4LinearKernel) + assert kernel.config == MxFp4LinearLayerConfig() + + +class UnsupportedMxFp4LinearKernel(MxFp4LinearKernel): + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + return False, "never supported" + + @classmethod + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + pass + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + pass + + +@patch("vllm.model_executor.kernels.linear.current_platform") +def test_init_mxfp4_linear_kernel_raises_when_no_kernel_matches(platform_mock): + platform_mock._enum = PlatformEnum.UNSPECIFIED + register_linear_kernel( + UnsupportedMxFp4LinearKernel, PlatformEnum.UNSPECIFIED, "mxfp4" + ) + + with pytest.raises(ValueError, match="Failed to find a kernel"): + init_mxfp4_linear_kernel() + + +@patch("vllm.model_executor.kernels.linear.mxfp4.aiter.is_aiter_found_and_supported") +@patch("vllm.model_executor.kernels.linear.mxfp4.aiter.current_platform") +@patch("vllm.model_executor.kernels.linear.current_platform") +def test_init_mxfp4_linear_kernel_raises_on_rocm_without_aiter( + linear_platform_mock, aiter_platform_mock, is_aiter_found_and_supported_mock +): + """On ROCm, the only registered MXFP4 linear kernel is AITER-based. + If AITER is not found/supported, no kernel should be selected.""" + linear_platform_mock._enum = PlatformEnum.ROCM + aiter_platform_mock.supports_mx.return_value = True + is_aiter_found_and_supported_mock.return_value = False + + with pytest.raises( + ValueError, + match="(?s)Failed to find a kernel.*" + "AITER not found or not supported on the current platform", + ): + init_mxfp4_linear_kernel() diff --git a/tests/quantization/test_quark.py b/tests/quantization/test_quark.py index d2f9dcb8af1..8440db10719 100644 --- a/tests/quantization/test_quark.py +++ b/tests/quantization/test_quark.py @@ -17,6 +17,7 @@ import pytest import torch from packaging import version +from vllm._aiter_ops import is_aiter_found_and_supported from vllm.model_executor.layers.quantization.quark.quark import ( # noqa: E501 QuarkLinearMethod, QuarkW8A8Fp8, @@ -26,6 +27,9 @@ from vllm.model_executor.layers.quantization.quark.quark_moe import ( # noqa: E QuarkW4A8Fp8MoEMethod, QuarkW8A8Int8MoEMethod, ) +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + quant_dequant_mxfp4, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( is_layer_skipped, ) @@ -52,6 +56,8 @@ QUARK_MXFP4_AVAILABLE = find_spec("quark") is not None and version.parse( importlib.metadata.version("amd-quark") ) >= version.parse(QUARK_MXFP4_MIN_VERSION) +AITER_AVAILABLE = is_aiter_found_and_supported() + DEVICE_TYPE = current_platform.device_type if QUARK_MXFP4_AVAILABLE: @@ -487,6 +493,42 @@ def test_mxfp4_dequant_kernel_match_quark( assert torch.equal(out_hip, out_torch) +@pytest.mark.skipif( + not QUARK_MXFP4_AVAILABLE, + reason=f"amd-quark>={QUARK_MXFP4_MIN_VERSION} is not available", +) +@pytest.mark.skipif( + not AITER_AVAILABLE, + reason="AITER is not found or not supported on the current platform", +) +@pytest.mark.parametrize("float_dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("scalings", [[2.3, 0.03, 7.3, 0.1, 0.004, 17.3, 1e4, 1e-4]]) +def test_mxfp4_dynamic_quant_match_quark( + float_dtype: torch.dtype, scalings: list[float] +): + """`AiterMxfp4LinearKernel` quantizes weights dynamically through AITER's + `dynamic_mxfp4_quant`, while the emulation path quantizes/dequantizes + through Quark's `qdq_mxfp4`. Check that both agree on the same input. + """ + from aiter.ops.triton.quant import dynamic_mxfp4_quant + + torch.manual_seed(0) + + hidden_size = 32 * 64 + inp = (torch.rand(48, hidden_size, dtype=float_dtype, device=DEVICE_TYPE) - 0.5) * 2 + for i in range(hidden_size // 32): + inp[:, i * 32 : (i + 1) * 32] = ( + inp[:, i * 32 : (i + 1) * 32] * scalings[i % len(scalings)] + ) + + x_q, x_s = dynamic_mxfp4_quant(inp) + out_dynamic_quant = dq_mxfp4_torch(x_q, x_s, float_dtype) + + out_quark_qdq = quant_dequant_mxfp4(inp) + + assert torch.equal(out_dynamic_quant, out_quark_qdq) + + # Unit tests for ``is_layer_skipped`` fused-name handling. FUSED_MAPPING = { diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index fcc50ffcb0e..39ed4309cad 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -74,6 +74,9 @@ from vllm.model_executor.kernels.linear.mxfp4 import ( MxFp4LinearKernel, MxFp4LinearLayerConfig, ) +from vllm.model_executor.kernels.linear.mxfp4.aiter import ( + AiterMxfp4LinearKernel, +) from vllm.model_executor.kernels.linear.mxfp4.flashinfer import ( FlashInferMxFp4LinearKernel, ) @@ -274,6 +277,7 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = { AiterFp8BlockScaledMMKernel, AiterPerTokenFp8ScaledMMLinearKernel, AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, + AiterMxfp4LinearKernel, }, "machete": { MacheteLinearKernel, @@ -469,6 +473,9 @@ _POSSIBLE_MXFP4_KERNELS: dict[PlatformEnum, list[type[MxFp4LinearKernel]]] = { MarlinMxFp4LinearKernel, HummingMxFp4LinearKernel, ], + PlatformEnum.ROCM: [ + AiterMxfp4LinearKernel, + ], PlatformEnum.XPU: [ XPUMxFp4LinearKernel, ], @@ -1079,6 +1086,7 @@ __all__ = [ "init_mxfp4_linear_kernel", "MxFp4LinearKernel", "MxFp4LinearLayerConfig", + "AiterMxfp4LinearKernel", "FlashInferMxFp4LinearKernel", "MarlinMxFp4LinearKernel", "FlashInferCutedslMxfp8LinearKernel", diff --git a/vllm/model_executor/kernels/linear/mxfp4/aiter.py b/vllm/model_executor/kernels/linear/mxfp4/aiter.py new file mode 100644 index 00000000000..54f8098b9a5 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp4/aiter.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch.nn.parameter import Parameter + +from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops +from vllm.platforms import current_platform + +from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig + +# NOTE: Do not import aiter at module scope. Importing aiter eagerly initializes HIP +# which can force the engine core to spawn instead of fork. +# is_aiter_found_and_supported() checks platform + arch + library availability via +# find_spec/amdsmi, so it stays HIP-free. +# Actual aiter imports are deferred to the functions/methods that need them, +# where HIP initialization is expected. +if is_aiter_found_and_supported(): + from vllm.utils.torch_utils import direct_register_custom_op + + def gemm_with_dynamic_quant( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + rocm_use_aiter_fp4_asm_gemm: bool = False, + out_dtype: torch.dtype | None = torch.bfloat16, + x_scales: torch.Tensor | None = None, + ) -> torch.Tensor: + from aiter.ops.triton.gemm_afp4wfp4 import ( + gemm_afp4wfp4, + gemm_afp4wfp4_preshuffled_weight_scales, + ) + from aiter.ops.triton.quant import dynamic_mxfp4_quant + + if rocm_use_aiter_fp4_asm_gemm: + from aiter import gemm_a4w4, per_1x32_f4_quant_hip + + M = x.shape[0] + N = weight.shape[0] + K = weight.shape[1] + if rocm_use_aiter_fp4_asm_gemm: + if M <= 64 and rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K): + if x_scales is None: + # use hip quant kernel for performance + if M >= 32: + x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=True) + else: + x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=False) + else: + x_q = x + x_s = x_scales + + if M >= 32: + x_s = x_s.view(torch.uint8).view(x_s.shape[0] // 32, -1) + else: + x_s = x_s[:M, ...].view(torch.uint8) + + y = torch.empty(M, N, device=x_q.device, dtype=out_dtype) + gemm_afp4wfp4_preshuffled_weight_scales( + x_q.view(torch.uint8), + weight.view(torch.uint8).view(weight.shape[0] // 16, -1), + x_s, + weight_scale.view(torch.uint8).view( + weight_scale.shape[0] // 32, -1 + ), + out_dtype, + y, + ) + else: + if x_scales is None: + # use hip quant kernel for performance + x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=True) + else: + x_q = x + x_s = x_scales + + y = gemm_a4w4( + x_q, + weight.view(x_q.dtype), + x_s, + weight_scale.view(x_s.dtype), + dtype=out_dtype, + bpreshuffle=True, + ) + return y[:M] + else: + if x_scales is None: + x_q, x_s = dynamic_mxfp4_quant(x) + else: + x_q = x + x_s = x_scales + y = torch.empty( + x_q.shape[0], weight.shape[0], device=x_q.device, dtype=out_dtype + ) + + gemm_afp4wfp4(x_q, weight, x_s, weight_scale.T, out_dtype, y) + return y + + def gemm_with_dynamic_quant_fake( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + x_scales: torch.Tensor = None, + rocm_use_aiter_fp4_asm_gemm: bool = False, + out_dtype: torch.dtype | None = torch.bfloat16, + ) -> torch.Tensor: + return torch.empty( + (*x.shape[:-1], weight.shape[0]), dtype=out_dtype, device=x.device + ) + + direct_register_custom_op( + op_name="gemm_with_dynamic_quant", + op_func=gemm_with_dynamic_quant, + mutates_args=[], + fake_impl=gemm_with_dynamic_quant_fake, + dispatch_key=current_platform.dispatch_key, + ) + + +class AiterMxfp4LinearKernel(MxFp4LinearKernel): + """AITER-based native MXFP4 GEMM kernel for ROCm.""" + + def __init__(self, config: MxFp4LinearLayerConfig) -> None: + super().__init__(config) + self.use_asm_gemm = rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled() + self.out_dtype = torch.get_default_dtype() + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.supports_mx(): + return False, "current platform does not support native MXFP4 computation" + if is_aiter_found_and_supported(): + return True, None + return False, "AITER not found or not supported on the current platform" + + @classmethod + def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + if self.use_asm_gemm: + from aiter.ops.shuffle import shuffle_weight + + weight_scale = layer.weight_scale.data + sm, sn = weight_scale.shape + weight_scale = weight_scale.view(sm // 32, 2, 16, sn // 8, 2, 4, 1) + weight_scale = weight_scale.permute(0, 3, 5, 2, 4, 1, 6).contiguous() + weight_scale = weight_scale.view(sm, sn) + layer.weight_scale = Parameter(weight_scale, requires_grad=False) + + layer.weight = Parameter( + shuffle_weight(layer.weight.data, layout=(16, 16)), + requires_grad=False, + ) + else: + layer.weight_scale = Parameter( + layer.weight_scale.data.T.contiguous(), requires_grad=False + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + y = torch.ops.vllm.gemm_with_dynamic_quant( + x, + layer.weight, + layer.weight_scale, + self.use_asm_gemm, + self.out_dtype, + ) + if bias is not None: + y = y + bias + return y diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py index ea63d1bcf12..2d42babeb38 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py @@ -9,8 +9,8 @@ from typing import Any import torch import torch.nn.functional as F -from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops from vllm.logger import init_logger +from vllm.model_executor.kernels.linear import init_mxfp4_linear_kernel from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( dequant_mxfp4, quant_dequant_mxfp4, @@ -36,120 +36,6 @@ from .quark_scheme import QuarkScheme logger = init_logger(__name__) -# NOTE: Do not import aiter at module scope. Importing aiter eagerly initializes HIP -# which can force the engine core to spawn instead of fork. -# is_aiter_found_and_supported() checks platform + arch + library availability via -# find_spec/amdsmi, so it stays HIP-free. -# Actual aiter imports are deferred to the functions/methods that need them, -# where HIP initialization is expected. -if is_aiter_found_and_supported(): - from vllm.utils.torch_utils import direct_register_custom_op - - def gemm_with_dynamic_quant( - x: torch.Tensor, - weight: torch.Tensor, - weight_scale: torch.Tensor, - rocm_use_aiter_fp4_asm_gemm: bool = False, - out_dtype: torch.dtype | None = torch.bfloat16, - x_scales: torch.Tensor | None = None, - ) -> torch.Tensor: - from aiter.ops.triton.gemm_afp4wfp4 import ( - gemm_afp4wfp4, - gemm_afp4wfp4_preshuffled_weight_scales, - ) - from aiter.ops.triton.quant import dynamic_mxfp4_quant - - if rocm_use_aiter_fp4_asm_gemm: - from aiter import gemm_a4w4, per_1x32_f4_quant_hip - - M = x.shape[0] - N = weight.shape[0] - K = weight.shape[1] - if rocm_use_aiter_fp4_asm_gemm: - if M <= 64 and rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K): - if x_scales is None: - # use hip quant kernel for performance - if M >= 32: - x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=True) - else: - x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=False) - else: - x_q = x - x_s = x_scales - - if M >= 32: - x_s = x_s.view(torch.uint8).view(x_s.shape[0] // 32, -1) - else: - x_s = x_s[:M, ...].view(torch.uint8) - - y = torch.empty(M, N, device=x_q.device, dtype=out_dtype) - gemm_afp4wfp4_preshuffled_weight_scales( - x_q.view(torch.uint8), - weight.view(torch.uint8).view(weight.shape[0] // 16, -1), - x_s, - weight_scale.view(torch.uint8).view( - weight_scale.shape[0] // 32, -1 - ), - out_dtype, - y, - ) - else: - if x_scales is None: - # use hip quant kernel for performance - x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=True) - else: - x_q = x - x_s = x_scales - - y = gemm_a4w4( - x_q, - weight.view(x_q.dtype), - x_s, - weight_scale.view(x_s.dtype), - dtype=out_dtype, - bpreshuffle=True, - ) - return y[:M] - else: - if x_scales is None: - x_q, x_s = dynamic_mxfp4_quant(x) - else: - x_q = x - x_s = x_scales - y = torch.empty( - x_q.shape[0], weight.shape[0], device=x_q.device, dtype=out_dtype - ) - - gemm_afp4wfp4(x_q, weight, x_s, weight_scale.T, out_dtype, y) - return y - - def gemm_with_dynamic_quant_fake( - x: torch.Tensor, - weight: torch.Tensor, - weight_scale: torch.Tensor, - x_scales: torch.Tensor = None, - rocm_use_aiter_fp4_asm_gemm: bool = False, - out_dtype: torch.dtype | None = torch.bfloat16, - ) -> torch.Tensor: - return torch.empty( - (*x.shape[:-1], weight.shape[0]), dtype=out_dtype, device=x.device - ) - - direct_register_custom_op( - op_name="gemm_with_dynamic_quant", - op_func=gemm_with_dynamic_quant, - mutates_args=[], - fake_impl=gemm_with_dynamic_quant_fake, - dispatch_key=current_platform.dispatch_key, - ) -elif current_platform.is_rocm(): - logger.warning( - "AITER is not found or not supported on the current platform, " - "QuarkOCP_MX will fall back to emulation." - "Native MXFP4/MXFP6 acceleration will not be available." - ) - - class QuarkOCP_MX(QuarkScheme): def __init__( self, @@ -157,8 +43,6 @@ class QuarkOCP_MX(QuarkScheme): input_quant_spec: dict[str, Any] | None, dynamic_mxfp4_quant: bool = False, ): - self.out_dtype = torch.get_default_dtype() - self.qscheme = "per_group" self.weight_quant_spec = weight_quant_spec self.input_quant_spec = input_quant_spec self.dynamic_mxfp4_quant = dynamic_mxfp4_quant @@ -211,17 +95,10 @@ class QuarkOCP_MX(QuarkScheme): self.input_dtype != "mxfp4" or self.weight_dtype != "mxfp4" ) - self.rocm_use_aiter_fp4_asm_gemm = ( - rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled() - ) - - if not self.emulate and not is_aiter_found_and_supported(): - # Currently need AITER kernels if not emulating - raise NotImplementedError( - f"{self.__class__.__name__} requires AITER to be installed " - "for non-emulation mode! Please refer to " - "https://github.com/ROCm/aiter for installation details." - ) + # TODO: Move emulation code path as a kernel, and always + # use init_mxfp4_linear_kernel. + if not self.emulate: + self.ocp_mx_linear = init_mxfp4_linear_kernel() if not current_platform.supports_mx(): logger.warning_once( @@ -268,7 +145,7 @@ class QuarkOCP_MX(QuarkScheme): from aiter.ops.triton.quant import dynamic_mxfp4_quant w_q, w_s = dynamic_mxfp4_quant(layer.weight) - layer.weight_scale = torch.nn.Parameter(w_s.T.contiguous(), requires_grad=False) + layer.weight_scale = torch.nn.Parameter(w_s, requires_grad=False) layer.weight = torch.nn.Parameter(w_q, requires_grad=False) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: @@ -284,31 +161,7 @@ class QuarkOCP_MX(QuarkScheme): else: if self.dynamic_mxfp4_quant: self.process_dynamic_mxfp4_weights_after_loading(layer) - elif self.rocm_use_aiter_fp4_asm_gemm: - from aiter.ops.shuffle import shuffle_weight - - # shuffle weight scale - weight_scale_shuffle = layer.weight_scale.data - sm, sn = weight_scale_shuffle.shape - weight_scale_shuffle = weight_scale_shuffle.view( - sm // 32, 2, 16, sn // 8, 2, 4, 1 - ) - weight_scale_shuffle = weight_scale_shuffle.permute( - 0, 3, 5, 2, 4, 1, 6 - ).contiguous() - weight_scale_shuffle = weight_scale_shuffle.view(sm, sn) - layer.weight_scale = torch.nn.Parameter( - weight_scale_shuffle, requires_grad=False - ) - - # shuffle weight - weight_shuffle = layer.weight.data - weight_shuffle = shuffle_weight(weight_shuffle, layout=(16, 16)) - layer.weight = torch.nn.Parameter(weight_shuffle, requires_grad=False) - else: - layer.weight_scale = torch.nn.Parameter( - layer.weight_scale.data.T.contiguous(), requires_grad=False - ) + self.ocp_mx_linear.process_weights_after_loading(layer) def create_weights( self, @@ -375,15 +228,4 @@ class QuarkOCP_MX(QuarkScheme): dq_w = self.dequant_func(layer.weight, layer.weight_scale, x.dtype) qdq_x = self.quant_dequant_func(x) return F.linear(qdq_x, dq_w, bias) - y = torch.ops.vllm.gemm_with_dynamic_quant( - x, - layer.weight, - layer.weight_scale, - self.rocm_use_aiter_fp4_asm_gemm, - self.out_dtype, - ) - # gemm_with_dynamic_quant has no bias argument; add it here so the - # native path matches F.linear (e.g. qkv_proj with qkv_bias=True). - if bias is not None: - y = y + bias - return y + return self.ocp_mx_linear.apply_weights(layer, x, bias) From 60915c972ccf0fa5a1993fef15a8bb667dbaa0f6 Mon Sep 17 00:00:00 2001 From: nvbfalk <bfalk@nvidia.com> Date: Tue, 28 Jul 2026 05:20:17 +0200 Subject: [PATCH 144/185] [Feature] Add VidCom2 video token pruning (#47750) Signed-off-by: Benedikt Falk <bfalk@nvidia.com> --- docs/design/cuda_graphs_multimodal.md | 2 +- docs/features/multimodal_inputs.md | 25 +++ tests/multimodal/test_vidcom2.py | 144 ++++++++++++++++++ vllm/config/model.py | 16 ++ vllm/config/multimodal.py | 21 ++- vllm/engine/arg_utils.py | 6 + vllm/model_executor/models/interfaces.py | 8 + vllm/model_executor/models/interns1_pro.py | 5 +- .../model_executor/models/nano_nemotron_vl.py | 8 +- vllm/model_executor/models/qwen2_5_vl.py | 12 +- vllm/model_executor/models/qwen3_vl.py | 85 ++++++++--- vllm/model_executor/models/qwen3_vl_moe.py | 5 +- vllm/model_executor/models/registry.py | 4 + vllm/multimodal/video_prune/__init__.py | 2 + vllm/multimodal/{ => video_prune}/evs.py | 0 vllm/multimodal/video_prune/vidcom2.py | 124 +++++++++++++++ .../processors/nano_nemotron_vl.py | 2 +- 17 files changed, 421 insertions(+), 48 deletions(-) create mode 100644 tests/multimodal/test_vidcom2.py create mode 100644 vllm/multimodal/video_prune/__init__.py rename vllm/multimodal/{ => video_prune}/evs.py (100%) create mode 100644 vllm/multimodal/video_prune/vidcom2.py diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index 7502186ae46..04fee73af60 100644 --- a/docs/design/cuda_graphs_multimodal.md +++ b/docs/design/cuda_graphs_multimodal.md @@ -101,7 +101,7 @@ When `mm_encoder_tp_mode="data"`, the manager distributes images across TP ranks Following <https://github.com/vllm-project/vllm/pull/35963> (ViT full CUDA graph support for image inference), <https://github.com/vllm-project/vllm/pull/38061> extends the encoder CUDA graph framework to support video inference for Qwen3-VL. Previously, the CUDA graph capture/replay path only handled image inputs (`pixel_values` + `image_grid_thw`). Video inputs use different keys (`pixel_values_videos` + `video_grid_thw`) and require larger `cu_seqlens` buffers because each video item contributes multiple frames (`T` attention sequences). This PR generalizes the protocol and manager to handle both modalities through a single shared graph manager. !!! note - Video CUDA graphs are automatically disabled when EVS (Efficient Video Sampling) pruning is enabled, since EVS makes the token count data-dependent and incompatible with CUDA graph capture. + Video CUDA graphs are automatically disabled when video token pruning (EVS or VidCom2) is enabled, since pruning makes the token count data-dependent and incompatible with CUDA graph capture. Mixed inputs (image+video) per prompt are also supported now. diff --git a/docs/features/multimodal_inputs.md b/docs/features/multimodal_inputs.md index 7f6c1aee760..4cb624030e9 100644 --- a/docs/features/multimodal_inputs.md +++ b/docs/features/multimodal_inputs.md @@ -350,6 +350,31 @@ Instead of NumPy arrays, you can also pass `'torch.Tensor'` instances, as shown Full example: [examples/generate/multimodal/vision_language_offline.py](../../examples/generate/multimodal/vision_language_offline.py) +#### Video Token Pruning + +For supported models, vLLM can prune video tokens after the vision encoder to +reduce prefill time and KV cache usage, at some cost in accuracy. Set +`--video-pruning-rate <q>` to prune the fraction `q` of video tokens from each +video, and `--video-pruning-method` to choose the training-free algorithm: + +- **`evs`** (Efficient Video Sampling, default): drops the tokens with the + lowest temporal dissimilarity to the previous frame. The first frame is + always fully retained. +- **`vidcom2`** (Video Compression Commander): scores tokens by similarity to + video-level and frame-level feature centers and gives distinctive frames a + larger share of the budget. At least one token per frame is retained. + +```bash +vllm serve Qwen/Qwen3-VL-8B-Instruct \ + --video-pruning-rate 0.75 --video-pruning-method vidcom2 +``` + +!!! note + `evs` is supported by all models implementing multimodal pruning; + `vidcom2` is currently supported by Qwen3-VL only. Unsupported combinations + are rejected at startup. Enabling video pruning also disables encoder CUDA + graphs, since the retained token count becomes data-dependent. + ### Audio Inputs You can pass a tuple `(array, sampling_rate)` to the `'audio'` field of the multi-modal dictionary. diff --git a/tests/multimodal/test_vidcom2.py b/tests/multimodal/test_vidcom2.py new file mode 100644 index 00000000000..4a62bbd87cf --- /dev/null +++ b/tests/multimodal/test_vidcom2.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.multimodal.video_prune.vidcom2 import ( + compute_retained_tokens_count, + compute_retention_mask, +) + + +def _fake_video_embeds( + num_frames: int, + rows: int, + cols: int, + hidden: int = 64, + seed: int = 0, +) -> torch.Tensor: + """Deterministic fake ViT output with a distinct mean per frame.""" + g = torch.Generator().manual_seed(seed) + frames = [] + for f in range(num_frames): + base = torch.randn(hidden, generator=g) * (0.1 + 0.05 * f) + frames.append( + base[None, :].expand(rows * cols, hidden) + + 0.01 * torch.randn(rows * cols, hidden, generator=g) + ) + return torch.cat(frames, dim=0) + + +@pytest.mark.parametrize("q", [0.25, 0.5, 0.75, 0.9]) +@pytest.mark.parametrize("num_frames", [1, 4, 16]) +def test_mask_shape_and_dtype(q: float, num_frames: int) -> None: + merge = 2 + rows, cols = 6, 8 + embeds = _fake_video_embeds(num_frames, rows, cols) + mask = compute_retention_mask( + embeds, + (num_frames, rows * merge, cols * merge), + spatial_merge_size=merge, + q=q, + ) + assert mask.dtype == torch.bool + assert mask.shape == (num_frames * rows * cols,) + + +def test_retained_count_floors_at_one_token_per_frame() -> None: + """The global minimum is one token per frame (not a full first frame).""" + assert ( + compute_retained_tokens_count(tokens_per_frame=48, num_frames=4, q=0.999) == 4 + ) + assert ( + compute_retained_tokens_count(tokens_per_frame=48, num_frames=4, q=0.0) + == 48 * 4 + ) + + +@pytest.mark.parametrize("q", [0.25, 0.5, 0.75, 0.9]) +@pytest.mark.parametrize("num_frames", [1, 4, 16]) +def test_total_retained_matches_target(q: float, num_frames: int) -> None: + """Mask total must equal the placeholder-sizing helper.""" + merge = 2 + rows, cols = 6, 8 + tpf = rows * cols + embeds = _fake_video_embeds(num_frames, rows, cols) + mask = compute_retention_mask( + embeds, + (num_frames, rows * merge, cols * merge), + spatial_merge_size=merge, + q=q, + ) + expected = compute_retained_tokens_count( + tokens_per_frame=tpf, num_frames=num_frames, q=q + ) + assert int(mask.sum().item()) == expected + + +def test_per_frame_min_one_when_budget_allows() -> None: + """No frame is fully dropped when the budget allows.""" + merge = 2 + rows, cols = 6, 8 + num_frames = 8 + embeds = _fake_video_embeds(num_frames, rows, cols) + mask = compute_retention_mask( + embeds, + (num_frames, rows * merge, cols * merge), + spatial_merge_size=merge, + q=0.25, + ) + per_frame = mask.view(num_frames, rows * cols).sum(dim=1) + assert (per_frame >= 1).all(), f"zero-token frame detected: {per_frame.tolist()}" + + +def test_dynamic_per_frame_budget() -> None: + """A distinctive frame gets more retained tokens than bland ones.""" + merge = 2 + rows, cols = 6, 8 + tpf = rows * cols + hidden = 64 + torch.manual_seed(0) + bland = 0.01 * torch.randn(tpf, hidden) + frames = [torch.randn(tpf, hidden) * 1.0] + for _ in range(7): + frames.append(bland + 0.001 * torch.randn(tpf, hidden)) + embeds = torch.cat(frames, dim=0) + mask = compute_retention_mask( + embeds, + (8, rows * merge, cols * merge), + spatial_merge_size=merge, + q=0.5, + ) + per_frame = mask.view(8, tpf).sum(dim=1) + assert per_frame[0].item() > per_frame[1:].float().mean().item() + + +def test_empty_input_safe() -> None: + embeds = torch.zeros(0, 32) + mask = compute_retention_mask(embeds, (0, 0, 0), spatial_merge_size=2, q=0.25) + assert mask.numel() == 0 + + +@pytest.mark.parametrize("q", [0.0, 0.25, 0.5, 0.75]) +def test_first_frame_not_privileged(q: float) -> None: + """A bland first frame is not force-retained (unlike EVS).""" + merge = 2 + rows, cols = 6, 8 + tpf = rows * cols + torch.manual_seed(1) + bland = 0.01 * torch.randn(tpf, 64) + frames = [bland] + for f in range(7): + frames.append(torch.randn(tpf, 64) * (1.0 + 0.1 * f)) + embeds = torch.cat(frames, dim=0) + mask = compute_retention_mask( + embeds, + (8, rows * merge, cols * merge), + spatial_merge_size=merge, + q=q, + ) + per_frame = mask.view(8, tpf).sum(dim=1) + assert per_frame[0].item() <= tpf + if q > 0.0: + assert per_frame[0].item() < int(mask.sum().item()) diff --git a/vllm/config/model.py b/vllm/config/model.py index d64bd57e87d..c2aa650f88d 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -376,6 +376,7 @@ class ModelConfig: interleave_mm_strings: InitVar[bool | None] = None skip_mm_profiling: InitVar[bool | None] = None video_pruning_rate: InitVar[float | None] = None + video_pruning_method: InitVar[str | None] = None mm_tensor_ipc: InitVar[MMTensorIPC] = None mm_ipc_gpu_memory_gb: InitVar[float | None] = None @@ -504,6 +505,7 @@ class ModelConfig: interleave_mm_strings: bool | None, skip_mm_profiling: bool | None, video_pruning_rate: float | None, + video_pruning_method: str | None, mm_tensor_ipc: MMTensorIPC, mm_ipc_gpu_memory_gb: float | None, ) -> None: @@ -735,6 +737,7 @@ class ModelConfig: interleave_mm_strings=interleave_mm_strings, skip_mm_profiling=skip_mm_profiling, video_pruning_rate=video_pruning_rate, + video_pruning_method=video_pruning_method, mm_tensor_ipc=mm_tensor_ipc, mm_ipc_gpu_memory_gb=mm_ipc_gpu_memory_gb, ) @@ -745,6 +748,19 @@ class ModelConfig: self.multimodal_config = MultiModalConfig(**mm_config_kwargs) # type: ignore[arg-type] + pruning_spec = self.multimodal_config.get_video_pruning_spec() + supported_pruning = self._model_info.supported_video_pruning_methods + if ( + pruning_spec is not None + and supported_pruning + and pruning_spec[0] not in supported_pruning + ): + raise ValueError( + f"Video pruning method '{pruning_spec[0]}' is not " + f"supported by {self._model_info.architecture} " + f"(supported methods: {supported_pruning})." + ) + if ( self.renderer_num_workers > 1 and self.multimodal_config.mm_processor_cache_gb > 0 diff --git a/vllm/config/multimodal.py b/vllm/config/multimodal.py index 865615dff79..ab5cf50fc3a 100644 --- a/vllm/config/multimodal.py +++ b/vllm/config/multimodal.py @@ -61,6 +61,7 @@ class MultiModalDummyOptionsBuiltins(TypedDict, total=False): MMEncoderTPMode = Literal["weights", "data"] MMCacheType = Literal["shm", "lru"] +VideoPruningMethod = Literal["evs", "vidcom2"] MMTensorIPC = Literal["direct_rpc", "torch_shm"] MMDummyOptions: TypeAlias = dict[str, BaseDummyOptions] """ @@ -189,9 +190,14 @@ class MultiModalConfig: estimating the peak memory usage of the activation of multimodal encoder and embedding cache.""" video_pruning_rate: float | None = Field(default=None, ge=0.0, lt=1.0) - """Sets pruning rate for video pruning via Efficient Video Sampling. - Value sits in range [0;1) and determines fraction of media tokens - from each video to be pruned. + """Fraction of video tokens to prune from each video. Value sits in range + [0;1); pruning is enabled when it is greater than 0. The pruning algorithm + is selected by `video_pruning_method`. + """ + video_pruning_method: VideoPruningMethod = "evs" + """Video token pruning algorithm applied when `video_pruning_rate` > 0: + - "evs": Efficient Video Sampling. + - "vidcom2": Video Compression Commander. """ mm_tensor_ipc: MMTensorIPC = "direct_rpc" """IPC (inter-process communication) method for multimodal tensors. @@ -360,4 +366,11 @@ class MultiModalConfig: ) def is_multimodal_pruning_enabled(self): - return self.video_pruning_rate is not None and self.video_pruning_rate > 0 + return self.get_video_pruning_spec() is not None + + def get_video_pruning_spec(self) -> tuple[VideoPruningMethod, float] | None: + """Return `(method, rate)` when video pruning is enabled, else None. + `rate` is the fraction of video tokens to prune.""" + if self.video_pruning_rate is not None and self.video_pruning_rate > 0: + return (self.video_pruning_method, float(self.video_pruning_rate)) + return None diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index b7c5f746545..fc9349a0bea 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -586,6 +586,7 @@ class EngineArgs: renderer_num_workers: int = 1 skip_mm_profiling: bool = MultiModalConfig.skip_mm_profiling video_pruning_rate: float | None = MultiModalConfig.video_pruning_rate + video_pruning_method: str = MultiModalConfig.video_pruning_method mm_tensor_ipc: MMTensorIPC = MultiModalConfig.mm_tensor_ipc mm_ipc_gpu_memory_gb: float = MultiModalConfig.mm_ipc_gpu_memory_gb # LoRA fields @@ -1333,6 +1334,10 @@ class EngineArgs: multimodal_group.add_argument( "--video-pruning-rate", **multimodal_kwargs["video_pruning_rate"] ) + multimodal_group.add_argument( + "--video-pruning-method", + **multimodal_kwargs["video_pruning_method"], + ) multimodal_group.add_argument( "--mm-tensor-ipc", **multimodal_kwargs["mm_tensor_ipc"] ) @@ -1715,6 +1720,7 @@ class EngineArgs: override_attention_dtype=self.override_attention_dtype, logits_processors=self.logits_processors, video_pruning_rate=self.video_pruning_rate, + video_pruning_method=self.video_pruning_method, mm_tensor_ipc=self.mm_tensor_ipc, mm_ipc_gpu_memory_gb=self.mm_ipc_gpu_memory_gb, io_processor_plugin=self.io_processor_plugin, diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index df9c18ce534..71229439f66 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -41,6 +41,7 @@ if TYPE_CHECKING: SpeechToTextParams, VllmConfig, ) + from vllm.config.multimodal import VideoPruningMethod from vllm.inputs import PromptType, TokensPrompt from vllm.lora.model_manager import LoRAModelManager from vllm.model_executor.layers.fused_moe import MoERunner @@ -424,6 +425,13 @@ class SupportsMultiModalPruning(Protocol): supports_multimodal_pruning: ClassVar[Literal[True]] = True + supported_video_pruning_methods: ClassVar[tuple["VideoPruningMethod", ...]] = ( + "evs", + ) + """Video pruning methods (as reported by + `MultiModalConfig.get_video_pruning_spec`) implemented by this model. + Models supporting methods beyond EVS should override this.""" + def recompute_mrope_positions( self, input_ids: list[int] | torch.Tensor, diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index c04b4729454..ef7d3d58343 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -570,10 +570,7 @@ class InternS1ProForConditionalGeneration( self.config = config self.multimodal_config = multimodal_config self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" - self.video_pruning_rate = multimodal_config.video_pruning_rate - self.is_multimodal_pruning_enabled = ( - multimodal_config.is_multimodal_pruning_enabled() - ) + self._init_video_pruning(multimodal_config) with self._mark_tower_model(vllm_config, {"image", "video"}): self.visual = Qwen3_VisionTransformer( diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index 64667503d57..5b4233b0781 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -42,10 +42,6 @@ from vllm.model_executor.models.utils import ( maybe_prefix, ) from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.evs import ( - compute_retained_tokens_count, - compute_retention_mask, -) from vllm.multimodal.inputs import ( AudioItem, BatchedTensorInputs, @@ -74,6 +70,10 @@ from vllm.multimodal.processing.processor import ( PromptReplacement, PromptUpdate, ) +from vllm.multimodal.video_prune.evs import ( + compute_retained_tokens_count, + compute_retention_mask, +) from vllm.renderers import TokenizeParams from vllm.sequence import IntermediateTensors from vllm.tokenizers import cached_tokenizer_from_config diff --git a/vllm/model_executor/models/qwen2_5_vl.py b/vllm/model_executor/models/qwen2_5_vl.py index c987e07b43d..7957ca805ff 100644 --- a/vllm/model_executor/models/qwen2_5_vl.py +++ b/vllm/model_executor/models/qwen2_5_vl.py @@ -67,12 +67,6 @@ from vllm.model_executor.layers.rotary_embedding.common import ( ) from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.evs import ( - compute_mrope_for_media, - compute_retained_tokens_count, - compute_retention_mask, - recompute_mrope_positions, -) from vllm.multimodal.inputs import ( MultiModalFeatureSpec, MultiModalFieldConfig, @@ -80,6 +74,12 @@ from vllm.multimodal.inputs import ( ) from vllm.multimodal.parse import MultiModalDataItems from vllm.multimodal.processing import PromptReplacement, PromptUpdate +from vllm.multimodal.video_prune.evs import ( + compute_mrope_for_media, + compute_retained_tokens_count, + compute_retention_mask, + recompute_mrope_positions, +) from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.tensor_schema import TensorSchema, TensorShape diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index f86560e5f4e..baf75fa2dd5 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -50,7 +50,12 @@ from transformers.video_utils import VideoMetadata from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions +from vllm.config.multimodal import ( + BaseDummyOptions, + MultiModalConfig, + VideoDummyOptions, + VideoPruningMethod, +) from vllm.distributed import get_pp_group, parallel_state from vllm.inputs import MultiModalDataDict from vllm.logger import init_logger @@ -69,12 +74,6 @@ from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.evs import ( - compute_mrope_for_media, - compute_retained_tokens_count, - compute_retention_mask, - recompute_mrope_positions, -) from vllm.multimodal.inputs import ( MultiModalFeatureSpec, MultiModalFieldConfig, @@ -92,6 +91,18 @@ from vllm.multimodal.processing import ( PromptUpdate, PromptUpdateDetails, ) +from vllm.multimodal.video_prune.evs import ( + compute_mrope_for_media, + compute_retained_tokens_count, + compute_retention_mask, + recompute_mrope_positions, +) +from vllm.multimodal.video_prune.vidcom2 import ( + compute_retained_tokens_count as vidcom2_compute_retained_tokens_count, +) +from vllm.multimodal.video_prune.vidcom2 import ( + compute_retention_mask as vidcom2_compute_retention_mask, +) from vllm.sequence import IntermediateTensors from vllm.tokenizers.protocol import TokenizerLike from vllm.tokenizers.registry import cached_tokenizer_from_config @@ -1256,7 +1267,7 @@ class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo]) hf_config = self.info.get_hf_config() tokenizer = self.info.get_tokenizer() merge_size = hf_config.vision_config.spatial_merge_size - video_pruning_rate = self.info.ctx.get_mm_config().video_pruning_rate + pruning_spec = self.info.ctx.get_mm_config().get_video_pruning_spec() vision_start_token_id = hf_config.vision_start_token_id vision_end_token_id = hf_config.vision_end_token_id video_token_id = hf_config.video_token_id @@ -1339,11 +1350,18 @@ class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo]) merge_size**2 ) - if video_pruning_rate is not None and video_pruning_rate > 0.0: - num_tokens = compute_retained_tokens_count( + # Apply video pruning (EVS or VidCom2) if enabled. + if pruning_spec is not None: + method, prune_q = pruning_spec + count_fn = ( + vidcom2_compute_retained_tokens_count + if method == "vidcom2" + else compute_retained_tokens_count + ) + num_tokens = count_fn( tokens_per_frame=tokens_per_frame_base, num_frames=num_frames, - q=video_pruning_rate, + q=prune_q, ) tokens_per_frame = [num_tokens] + [0] * (num_frames - 1) select_token_id = False @@ -1459,16 +1477,22 @@ class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo]) f"video length ({grid_thw[0]})." ) - # Compute tokens per frame, with EVS support + # Compute tokens per frame, with EVS / VidCom2 support num_frames = int(grid_thw[0]) tokens_per_frame_base = int(grid_thw[1:].prod()) // merge_length - video_pruning_rate = self.info.ctx.get_mm_config().video_pruning_rate - if video_pruning_rate is not None and video_pruning_rate > 0.0: - num_tokens = compute_retained_tokens_count( + pruning_spec = self.info.ctx.get_mm_config().get_video_pruning_spec() + if pruning_spec is not None: + method, prune_q = pruning_spec + count_fn = ( + vidcom2_compute_retained_tokens_count + if method == "vidcom2" + else compute_retained_tokens_count + ) + num_tokens = count_fn( tokens_per_frame=tokens_per_frame_base, num_frames=num_frames, - q=video_pruning_rate, + q=prune_q, ) tokens_per_frame = [num_tokens] + [0] * (num_frames - 1) select_token_id = False @@ -1701,6 +1725,8 @@ class Qwen3VLForConditionalGeneration( supports_encoder_tp_data = True + supported_video_pruning_methods = ("evs", "vidcom2") + # To ensure correct weight loading and mapping. hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ @@ -1719,6 +1745,17 @@ class Qwen3VLForConditionalGeneration( raise ValueError("Only image or video modality is supported") + def _init_video_pruning(self, multimodal_config: MultiModalConfig) -> None: + pruning_spec = multimodal_config.get_video_pruning_spec() + if pruning_spec is None: + self.video_pruning_method: VideoPruningMethod | None = None + self.video_pruning_rate = multimodal_config.video_pruning_rate + else: + self.video_pruning_method, self.video_pruning_rate = pruning_spec + self.is_multimodal_pruning_enabled = ( + multimodal_config.is_multimodal_pruning_enabled() + ) + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"): super().__init__() config: Qwen3VLConfig = vllm_config.model_config.hf_config @@ -1730,10 +1767,7 @@ class Qwen3VLForConditionalGeneration( self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config) self.multimodal_config = multimodal_config self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" - self.video_pruning_rate = multimodal_config.video_pruning_rate - self.is_multimodal_pruning_enabled = ( - multimodal_config.is_multimodal_pruning_enabled() - ) + self._init_video_pruning(multimodal_config) self.use_deepstack = hasattr(config.vision_config, "deepstack_visual_indexes") self.deepstack_num_level = ( @@ -1848,7 +1882,7 @@ class Qwen3VLForConditionalGeneration( EncoderCudaGraphConfig, ) - # When EVS pruning is enabled, embed_multimodal post-processes both + # When video pruning is enabled, embed_multimodal post-processes both # image and video embeddings (mrope positions are appended for image, # prune+append for video). The encoder CUDA graph path bypasses that # post-process, producing inconsistent embedding formats vs eager. So @@ -2291,9 +2325,12 @@ class Qwen3VLForConditionalGeneration( t, h, w = size if self.is_multimodal_pruning_enabled: - # For each video, compute retention mask using EVS. - # retention_mask: [11424]. - retention_mask = compute_retention_mask( + # Compute the retention mask for each video (EVS or VidCom2). + if self.video_pruning_method == "vidcom2": + mask_fn = vidcom2_compute_retention_mask + else: + mask_fn = compute_retention_mask + retention_mask = mask_fn( emb, size, spatial_merge_size=self.visual.spatial_merge_size, diff --git a/vllm/model_executor/models/qwen3_vl_moe.py b/vllm/model_executor/models/qwen3_vl_moe.py index 4413e1213bb..f1409d23399 100644 --- a/vllm/model_executor/models/qwen3_vl_moe.py +++ b/vllm/model_executor/models/qwen3_vl_moe.py @@ -220,10 +220,7 @@ class Qwen3VLMoeForConditionalGeneration( self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config) self.multimodal_config = multimodal_config self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" - self.video_pruning_rate = multimodal_config.video_pruning_rate - self.is_multimodal_pruning_enabled = ( - multimodal_config.is_multimodal_pruning_enabled() - ) + self._init_video_pruning(multimodal_config) self.use_deepstack = hasattr(config.vision_config, "deepstack_visual_indexes") self.deepstack_num_level = ( diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 1dcacc7936e..697f55b3727 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -797,6 +797,7 @@ class _ModelInfo: supports_replayssm: bool supports_transcription: bool supports_transcription_only: bool + supported_video_pruning_methods: tuple[str, ...] @staticmethod def from_model_cls(model: type[nn.Module]) -> "_ModelInfo": @@ -827,6 +828,9 @@ class _ModelInfo: supports_transcription(model) and model.supports_transcription_only ), has_noops=has_noops(model), + supported_video_pruning_methods=getattr( + model, "supported_video_pruning_methods", () + ), ) diff --git a/vllm/multimodal/video_prune/__init__.py b/vllm/multimodal/video_prune/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm/multimodal/video_prune/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/multimodal/evs.py b/vllm/multimodal/video_prune/evs.py similarity index 100% rename from vllm/multimodal/evs.py rename to vllm/multimodal/video_prune/evs.py diff --git a/vllm/multimodal/video_prune/vidcom2.py b/vllm/multimodal/video_prune/vidcom2.py new file mode 100644 index 00000000000..830e47bc055 --- /dev/null +++ b/vllm/multimodal/video_prune/vidcom2.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# VidCom2 (Video Compression Commander) video token pruning. +# Liu et al., EMNLP 2025 — https://arxiv.org/abs/2505.14454 +# Adapted from the reference implementation: +# https://github.com/xuyang-liu16/VidCom2 (Apache-2.0, +# Copyright (c) 2025 the VidCom2 authors). + +import torch +import torch.nn.functional as F + +# Multi-scale Gaussian bandwidths from the reference implementation. +_ALPHAS: tuple[float, ...] = tuple(2.0**k for k in range(-3, 2)) +_LOW_VAR_CHANNEL_RATIO: float = 0.5 +_SOFTMAX_TEMPERATURE: float = 0.01 + + +def compute_retained_tokens_count( + tokens_per_frame: int, num_frames: int, q: float +) -> int: + """Number of video tokens retained after VidCom2 pruning. + + The target is `(1 - q) * total_tokens`, i.e. a retention ratio of + `1 - q` averaged across frames. Because the per-frame budget is floored + at one token, the global minimum is `num_frames` (one token per frame). + """ + total_tokens = tokens_per_frame * num_frames + base_num = int(total_tokens * (1.0 - q)) + return max(num_frames, min(base_num, total_tokens)) + + +def compute_retention_mask( + video_embeds: torch.Tensor, + video_size_thw: torch.LongTensor | tuple[int, int, int], + spatial_merge_size: int, + q: float, +) -> torch.Tensor: + """Compute the VidCom2 retention mask for a single video. + + Args: + video_embeds: `(T*H*W/merge^2, hidden_size)` post-ViT token features. + video_size_thw: `(T, H, W)` grid dimensions. + spatial_merge_size: ViT spatial merge factor (e.g. 2). + q: Pruning fraction in `[0, 1)`; retention ratio is `1 - q`. + + Returns: + Flat bool tensor of shape `(T*H*W/merge^2,)`, True for retained + tokens. The True count equals `compute_retained_tokens_count` so + placeholders sized at prompt-processing time match exactly. + """ + T, H, W = map(int, video_size_thw) + rows = H // spatial_merge_size + cols = W // spatial_merge_size + tokens_per_frame = rows * cols + total_tokens = T * tokens_per_frame + + device = video_embeds.device + if tokens_per_frame == 0 or total_tokens == 0: + return torch.ones(0, dtype=torch.bool, device=device) + + target_retained = compute_retained_tokens_count( + tokens_per_frame=tokens_per_frame, num_frames=T, q=q + ) + target_retained = min(target_retained, total_tokens) + + # 1. Score in the lowest-variance half of channels. + variances = video_embeds.var(dim=0, unbiased=False) + k_channels = max(1, int(video_embeds.size(-1) * _LOW_VAR_CHANNEL_RATIO)) + _, low_var_idx = torch.topk(variances, k=k_channels, largest=False) + sel = video_embeds.index_select(-1, low_var_idx) + + # 2. Multi-scale Gaussian similarity to video and per-frame centers. + frames = sel.view(T, tokens_per_frame, sel.size(-1)) + frames = F.normalize(frames, dim=-1) + vid_center = frames.mean(dim=(0, 1), keepdim=True) # (1, 1, C) + frame_center = frames.mean(dim=1, keepdim=True) # (T, 1, C) + v_score = _multi_scale_gaussian(frames, vid_center) + f_score = _multi_scale_gaussian(frames, frame_center) + # Higher similarity = more redundant; lowest-similarity tokens are kept. + similarity = v_score + f_score # (T, tpf) + + # 3. Per-frame dynamic budget: distinctive frames get a larger share. + base = 1.0 - q + frame_scores = -v_score.mean(dim=-1) # (T,) + probs = F.softmax((frame_scores - frame_scores.max()) / _SOFTMAX_TEMPERATURE, dim=0) + scales = (base * (1.0 + probs - probs.mean())).clamp(max=1.0) + ks = (scales * tokens_per_frame).round().long().clamp(min=1, max=tokens_per_frame) + + # 4. Retain the smallest-similarity tokens per frame. + mask_2d = torch.zeros(T, tokens_per_frame, dtype=torch.bool, device=device) + for i in range(T): + k_i = int(ks[i].item()) + if k_i <= 0: + continue + _, idx = torch.topk(similarity[i], k=k_i, largest=False, sorted=False) + mask_2d[i].scatter_(0, idx, True) + + # 5. Reconcile rounding/clamp drift to the exact target count by score. + flat_mask = mask_2d.view(-1) + flat_sim = similarity.view(-1) + current = int(flat_mask.sum().item()) + if current > target_retained: + drop_n = current - target_retained + retained_idx = flat_mask.nonzero(as_tuple=False).squeeze(-1) + retained_sim = flat_sim[retained_idx] + _, worst = torch.topk(retained_sim, k=drop_n, largest=True, sorted=False) + flat_mask[retained_idx[worst]] = False + elif current < target_retained: + add_n = target_retained - current + available_idx = (~flat_mask).nonzero(as_tuple=False).squeeze(-1) + if available_idx.numel() > 0: + available_sim = flat_sim[available_idx] + add_n = min(add_n, available_idx.numel()) + _, best = torch.topk(available_sim, k=add_n, largest=False, sorted=False) + flat_mask[available_idx[best]] = True + + return flat_mask + + +def _multi_scale_gaussian(x: torch.Tensor, center: torch.Tensor) -> torch.Tensor: + """Sum Gaussian kernels over `_ALPHAS`; `(T, N, C) -> (T, N)` scores.""" + dist_sq = ((x - center) ** 2).sum(dim=-1) + return sum(torch.exp(-dist_sq / (2.0 * a)) for a in _ALPHAS) diff --git a/vllm/transformers_utils/processors/nano_nemotron_vl.py b/vllm/transformers_utils/processors/nano_nemotron_vl.py index d48a29d6b43..028d207a25a 100644 --- a/vllm/transformers_utils/processors/nano_nemotron_vl.py +++ b/vllm/transformers_utils/processors/nano_nemotron_vl.py @@ -23,9 +23,9 @@ from PIL import Image from transformers import BatchFeature, PretrainedConfig, TensorType from vllm.model_executor.models.parakeet import ParakeetExtractor -from vllm.multimodal.evs import compute_retained_tokens_count from vllm.multimodal.inputs import AudioItem from vllm.multimodal.processing.processor import PromptUpdateDetails +from vllm.multimodal.video_prune.evs import compute_retained_tokens_count from vllm.tokenizers.hf import HfTokenizer from .internvl import calculate_internvl_targets, get_internvl_target_ratios From d18ed2304a2703e3211fc384a58607e754f5b723 Mon Sep 17 00:00:00 2001 From: Varun Sundar Rabindranath <varunsundar08@gmail.com> Date: Mon, 27 Jul 2026 23:21:53 -0400 Subject: [PATCH 145/185] [KV-offload][FS] : Batch store/load_block in C (#49152) Signed-off-by: <> Co-authored-by: Varun Sundar Rabindranath <varun-sundar-rabindranath@h100-01.nemg-001.lab.rdu2.dc.redhat.com> --- csrc/fs_io.cpp | 284 +++++++++++++++++++- tests/v1/kv_offload/tiering/test_fs_tier.py | 99 ++++--- vllm/v1/kv_offload/tiering/fs/io.py | 82 +++++- vllm/v1/kv_offload/tiering/fs/manager.py | 45 ++-- 4 files changed, 452 insertions(+), 58 deletions(-) diff --git a/csrc/fs_io.cpp b/csrc/fs_io.cpp index fdf3e614e64..6b7c8d7fe1e 100644 --- a/csrc/fs_io.cpp +++ b/csrc/fs_io.cpp @@ -3,19 +3,155 @@ #include <Python.h> +#include <errno.h> +#include <fcntl.h> #include <unistd.h> +#include <filesystem> +#include <string> #include <vector> +#if defined(O_DIRECT) +constexpr int kODirectFlag = O_DIRECT; +#else +constexpr int kODirectFlag = 0; +#endif + extern "C" { -static void _batch_lookup(const std::vector<const char*>& paths, +namespace { + +// Returns 0 on success, or the std::error_code's POSIX-compatible value on +// failure, mirroring the errno convention used by the syscalls below. +inline int ensure_parent_dirs(const std::string& path) { + const auto parent = std::filesystem::path(path).parent_path(); + if (parent.empty()) { + return 0; + } + std::error_code ec; + std::filesystem::create_directories(parent, ec); + return ec ? ec.value() : 0; +} + +// Core single-block store: src/size are raw pointer + byte count. Returns 0 +// on success, or the errno of the failing step on failure -- captured +// before any subsequent cleanup call can overwrite it. On failure, the temp +// file is removed. +inline int _store_block(const char* tmp_path, const char* dest_path, + const char* src, size_t size, bool use_o_direct) { + if (access(dest_path, F_OK) == 0) { + return 0; // Already present. + } + + if (const int err = ensure_parent_dirs(dest_path); err != 0) { + return err; + } + + const int o_direct_flag = use_o_direct ? kODirectFlag : 0; + const int fd = open( + tmp_path, O_CREAT | O_EXCL | O_WRONLY | O_TRUNC | o_direct_flag, 0644); + if (fd < 0) { + return errno; + } + + const ssize_t written = write(fd, src, size); + if (written < 0 || static_cast<size_t>(written) != size) { + const int err = written < 0 ? errno : EIO; + close(fd); // Best-effort cleanup; the real error is already captured. + unlink(tmp_path); + return err; + } + + if (close(fd) != 0) { + const int err = errno; + unlink(tmp_path); + return err; + } + + if (rename(tmp_path, dest_path) != 0) { + const int err = errno; + unlink(tmp_path); + return err; + } + + return 0; +} + +// Core single-block load: dst/size are raw pointer + byte count. Returns 0 +// on success, or the errno of the failing step on failure. On failure, +// the source file is removed since a partially-read block should not be reused. +inline int _load_block(const char* source_path, char* dst, size_t size, + bool use_o_direct) { + const int o_direct_flag = use_o_direct ? kODirectFlag : 0; + const int fd = open(source_path, O_RDONLY | o_direct_flag, 0); + if (fd < 0) { + const int err = errno; + unlink(source_path); + return err; + } + + const ssize_t bytes_read = read(fd, dst, size); + if (bytes_read < 0 || static_cast<size_t>(bytes_read) != size) { + const int err = bytes_read < 0 ? errno : EIO; + close(fd); + unlink(source_path); + return err; + } + + if (close(fd) != 0) { + const int err = errno; + unlink(source_path); + return err; + } + + return 0; +} + +inline void _batch_lookup(const std::vector<const char*>& paths, std::vector<int>& exists_flags) { for (size_t i = 0; i < paths.size(); i++) { exists_flags[i] = (access(paths[i], F_OK) == 0) ? 1 : 0; } } +// Helper: extract a list[str] of length n into a vector<const char*>. +// Returns false and sets a Python exception on error. +inline bool extract_str_list(PyObject* list, Py_ssize_t n, + std::vector<const char*>& out) { + for (Py_ssize_t i = 0; i < n; i++) { + out[i] = PyUnicode_AsUTF8AndSize(PyList_GetItem(list, i), nullptr); + if (out[i] == nullptr) { + return false; + } + } + return true; +} + +// Helper: extract a Py_buffer per element of a list[bytes-like] of length n. +// On success, `out` holds n acquired buffers (caller must PyBuffer_Release +// each). On failure, any buffers already acquired are released before +// returning false, and a Python exception is set. +inline bool extract_buffer_list(PyObject* list, Py_ssize_t n, int flags, + std::vector<Py_buffer>& out) { + for (Py_ssize_t i = 0; i < n; i++) { + if (PyObject_GetBuffer(PyList_GetItem(list, i), &out[i], flags) != 0) { + for (Py_ssize_t j = 0; j < i; j++) { + PyBuffer_Release(&out[j]); + } + return false; + } + } + return true; +} + +inline void release_buffer_list(std::vector<Py_buffer>& buffers) { + for (auto& buf : buffers) { + PyBuffer_Release(&buf); + } +} + +} // namespace + /// @brief Check file existence for a batch of paths. /// @param paths list[str] – absolute paths to check. /// @return list[bool] – True if the corresponding path exists, False otherwise. @@ -51,11 +187,157 @@ static PyObject* batch_lookup(PyObject* /*self*/, PyObject* args) { return result; } +/// @brief Store a batch of blocks, each from its own buffer, to disk. +/// @param tmp_paths list[str] – one temp path per block. +/// @param dest_paths list[str] – one destination path per block. +/// @param buffers list[bytes-like] – one source buffer per block. +/// @param use_o_direct bool – whether to open files with O_DIRECT +/// (default True). Ignored where O_DIRECT is unsupported +/// by the platform. +/// @note Releases the GIL for the entire batch. Raises on first error. +static PyObject* batch_store_block(PyObject* /*self*/, PyObject* args) { + PyObject* tmp_paths_obj = nullptr; + PyObject* dest_paths_obj = nullptr; + PyObject* buffers_obj = nullptr; + int use_o_direct = 1; + + if (!PyArg_ParseTuple(args, "O!O!O!|p", &PyList_Type, &tmp_paths_obj, + &PyList_Type, &dest_paths_obj, &PyList_Type, + &buffers_obj, &use_o_direct)) { + return nullptr; + } + + const Py_ssize_t n = PyList_Size(tmp_paths_obj); + if (PyList_Size(dest_paths_obj) != n || PyList_Size(buffers_obj) != n) { + PyErr_SetString( + PyExc_ValueError, + "tmp_paths, dest_paths and buffers must have the same length"); + return nullptr; + } + + std::vector<const char*> tmp_paths(n); + std::vector<const char*> dest_paths(n); + + if (!extract_str_list(tmp_paths_obj, n, tmp_paths)) return nullptr; + if (!extract_str_list(dest_paths_obj, n, dest_paths)) return nullptr; + + std::vector<Py_buffer> buffers(n); + if (!extract_buffer_list(buffers_obj, n, PyBUF_SIMPLE, buffers)) { + return nullptr; + } + + Py_ssize_t failed_index = -1; + int failure_errno = 0; + + { + Py_BEGIN_ALLOW_THREADS for (Py_ssize_t i = 0; i < n; i++) { + const char* buf = static_cast<const char*>(buffers[i].buf); + const int err = + _store_block(tmp_paths[i], dest_paths[i], buf, + static_cast<size_t>(buffers[i].len), use_o_direct); + if (err != 0) { + failed_index = i; + failure_errno = err; + break; + } + } + Py_END_ALLOW_THREADS + } + + release_buffer_list(buffers); + + if (failed_index >= 0) { + // PyErr_SetFromErrnoWithFilename() reads the errno to format exception. + errno = failure_errno; + return PyErr_SetFromErrnoWithFilename(PyExc_OSError, + dest_paths[failed_index]); + } + + Py_RETURN_NONE; +} + +/// @brief Load a batch of blocks from disk, each into its own buffer. +/// @param source_paths list[str] – one source path per block. +/// @param buffers list[writable bytes-like] – one destination buffer +/// per block. +/// @param use_o_direct bool – whether to open files with O_DIRECT +/// (default True). Ignored where O_DIRECT is unsupported +/// by the platform. +/// @note Releases the GIL for the entire batch. Raises on first error. +static PyObject* batch_load_block(PyObject* /*self*/, PyObject* args) { + PyObject* source_paths_obj = nullptr; + PyObject* buffers_obj = nullptr; + int use_o_direct = 1; + + if (!PyArg_ParseTuple(args, "O!O!|p", &PyList_Type, &source_paths_obj, + &PyList_Type, &buffers_obj, &use_o_direct)) { + return nullptr; + } + + const Py_ssize_t n = PyList_Size(source_paths_obj); + if (PyList_Size(buffers_obj) != n) { + PyErr_SetString(PyExc_ValueError, + "source_paths and buffers must have the same length"); + return nullptr; + } + + std::vector<const char*> source_paths(n); + if (!extract_str_list(source_paths_obj, n, source_paths)) return nullptr; + + std::vector<Py_buffer> buffers(n); + if (!extract_buffer_list(buffers_obj, n, PyBUF_WRITABLE, buffers)) { + return nullptr; + } + + Py_ssize_t failed_index = -1; + int failure_errno = 0; + + { + Py_BEGIN_ALLOW_THREADS for (Py_ssize_t i = 0; i < n; i++) { + char* buf = static_cast<char*>(buffers[i].buf); + const int err = + _load_block(source_paths[i], buf, static_cast<size_t>(buffers[i].len), + use_o_direct); + if (err != 0) { + failed_index = i; + failure_errno = err; + break; + } + } + Py_END_ALLOW_THREADS + } + + release_buffer_list(buffers); + + if (failed_index >= 0) { + // PyErr_SetFromErrnoWithFilename() reads the errno to format exception. + errno = failure_errno; + return PyErr_SetFromErrnoWithFilename(PyExc_OSError, + source_paths[failed_index]); + } + + Py_RETURN_NONE; +} + static PyMethodDef fs_io_C_methods[] = { {"batch_lookup", batch_lookup, METH_VARARGS, "batch_lookup(paths: list[str]) -> list[bool]\n" "\n" "Check file existence for a batch of paths."}, + {"batch_store_block", batch_store_block, METH_VARARGS, + "batch_store_block(tmp_paths: list[str], dest_paths: list[str],\n" + " buffers: list[bytes-like],\n" + " use_o_direct: bool = True) -> None\n" + "\n" + "Store a batch of blocks, each from its own buffer, to disk. Raises on " + "first error."}, + {"batch_load_block", batch_load_block, METH_VARARGS, + "batch_load_block(source_paths: list[str],\n" + " buffers: list[writable bytes-like],\n" + " use_o_direct: bool = True) -> None\n" + "\n" + "Load a batch of blocks from disk into corresponding buffers. " + "Raises on first error."}, {nullptr, nullptr, 0, nullptr}, }; diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index 68b5512a3ec..66533e3676b 100644 --- a/tests/v1/kv_offload/tiering/test_fs_tier.py +++ b/tests/v1/kv_offload/tiering/test_fs_tier.py @@ -46,6 +46,7 @@ from vllm.v1.kv_offload.tiering.fs.thread_pool import DualQueueThreadPool # Helpers # --------------------------------------------------------------------------- +_NUM_BLOCKS = 8 _BLOCK_ELEMENTS = 128 * mmap.PAGESIZE # 2MB per block for pagesize 4096. _DTYPE: torch.dtype = torch.float32 _CTX = ReqContext(req_id="test") @@ -162,7 +163,7 @@ def _page_aligned_rand_tensor( @pytest.fixture def fs_tier(tmp_path): - tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS) + tensor = _page_aligned_zero_tensor(_NUM_BLOCKS, _BLOCK_ELEMENTS) mock_view = memoryview(tensor.numpy()) tier = FileSystemTierManager( offloading_spec=_MOCK_OFFLOADING_SPEC, @@ -178,7 +179,7 @@ def fs_tier(tmp_path): @pytest.fixture def fs_tier_with_events(tmp_path): - tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS) + tensor = _page_aligned_zero_tensor(_NUM_BLOCKS, _BLOCK_ELEMENTS) mock_view = memoryview(tensor.numpy()) tier = FileSystemTierManager( offloading_spec=_make_offloading_spec(enable_kv_cache_events=True), @@ -347,33 +348,43 @@ def test_shutdown_discards_pending_tasks(fs_tier): assert all(not t.is_alive() for t in tier._pool._threads) -def test_store_load_data_integrity(fs_tier): - """Data written by store must be exactly recovered by load.""" +@pytest.mark.parametrize("batch_size", [0, 1, 2, 5]) +@pytest.mark.parametrize("use_c_ext", [True, False]) +def test_store_load_data_integrity(fs_tier, monkeypatch, use_c_ext, batch_size): + """Data written by store must be exactly recovered by load, for batches + of any size -- including the empty batch.""" + import vllm.v1.kv_offload.tiering.fs.io as io_mod + + if use_c_ext and not io_mod._HAS_FSIO_C: + pytest.skip("fs_io_C extension not built") + monkeypatch.setattr(io_mod, "_HAS_FSIO_C", use_c_ext) + tier, tensor = fs_tier # Populate tensor with random data - tensor[:] = _page_aligned_rand_tensor(4, _BLOCK_ELEMENTS) + tensor[:] = _page_aligned_rand_tensor(_NUM_BLOCKS, _BLOCK_ELEMENTS) - # Store first 2 blocks - num_store = 2 - expected = tensor[:num_store].clone() + keys = [key(i) for i in range(batch_size)] + store_block_ids = list(range(batch_size)) + load_block_ids = list(range(_NUM_BLOCKS - batch_size, _NUM_BLOCKS)) + expected = tensor[:batch_size].clone() - store_ids = list(range(num_store)) - keys = [key(i) for i in range(num_store)] + tier.submit_store(make_job(1, keys, store_block_ids)) + store_results = drain(tier) + assert len(store_results) == 1 + assert store_results[0].success + assert all(os.path.exists(tier.file_mapper.get_file_name(k)) for k in keys) - tier.submit_store(make_job(1, keys, store_ids)) - results = drain(tier) - assert all(r.success for r in results) + # reset tensor to prove data is read from disk + tensor[:] = 0.0 - # Overwrite source blocks to prove data is read from disk - tensor[:num_store] = 0.0 + # Load into a range disjoint by index from the store ids, to also + # exercise loading a block into a different id than it was stored from. + tier.submit_load(make_job(2, keys, load_block_ids, is_promotion=True)) + load_results = drain(tier) + assert len(load_results) == 1 + assert load_results[0].success - # Load into last 2 blocks - load_ids = [2, 3] - tier.submit_load(make_job(2, keys, load_ids, is_promotion=True)) - results = drain(tier) - assert all(r.success for r in results) - - for i, bid in enumerate(load_ids): + for i, bid in enumerate(load_block_ids): assert torch.allclose(tensor[bid], expected[i]), ( f"Block {bid} data mismatch after store+load" ) @@ -499,6 +510,30 @@ def test_batch_lookup_dispatch(fs_tier, monkeypatch, use_c_ext): assert results == [LookupResult.HIT, LookupResult.MISS] +@pytest.mark.parametrize("use_c_ext", [True, False]) +def test_out_of_bounds_block_id_smoke(fs_tier, monkeypatch, use_c_ext): + """Smoke test: a block id beyond the primary tensor's block count must + fail the job, for both the C extension and the Python fallback.""" + import vllm.v1.kv_offload.tiering.fs.io as io_mod + + if use_c_ext and not io_mod._HAS_FSIO_C: + pytest.skip("fs_io_C extension not built") + monkeypatch.setattr(io_mod, "_HAS_FSIO_C", use_c_ext) + + tier, tensor = fs_tier + out_of_bounds_bid = tensor.shape[0] # one past the last valid block + + tier.submit_store(make_job(1, [key(1)], [out_of_bounds_bid])) + store_results = drain(tier) + assert len(store_results) == 1 + assert not store_results[0].success + + tier.submit_load(make_job(2, [key(1)], [out_of_bounds_bid], is_promotion=True)) + load_results = drain(tier) + assert len(load_results) == 1 + assert not load_results[0].success + + # --------------------------------------------------------------------------- # KV events # --------------------------------------------------------------------------- @@ -571,14 +606,14 @@ def test_mixed_job_results_emit_event_only_for_successful_job( tier = fs_tier_with_events failing_path = tier.file_mapper.get_file_name(key(1)) - original_store_block = mgr_mod.store_block + original_batch_store_block = mgr_mod.batch_store_block - def flaky_store_block(dest_path, *args, **kwargs): - if dest_path == failing_path: + def flaky_batch_store_block(paths, *args, **kwargs): + if failing_path in paths: raise OSError("injected store failure") - return original_store_block(dest_path, *args, **kwargs) + return original_batch_store_block(paths, *args, **kwargs) - monkeypatch.setattr(mgr_mod, "store_block", flaky_store_block) + monkeypatch.setattr(mgr_mod, "batch_store_block", flaky_batch_store_block) tier.submit_store(make_job(1, [key(1)], [0])) tier.submit_store(make_job(2, [key(2)], [1])) @@ -599,14 +634,14 @@ def test_partially_failed_store_emits_no_event(fs_tier_with_events, monkeypatch) tier = fs_tier_with_events failing_path = tier.file_mapper.get_file_name(key(2)) - original_store_block = mgr_mod.store_block + original_batch_store_block = mgr_mod.batch_store_block - def flaky_store_block(dest_path, *args, **kwargs): - if dest_path == failing_path: + def flaky_batch_store_block(paths, *args, **kwargs): + if failing_path in paths: raise OSError("injected store failure") - return original_store_block(dest_path, *args, **kwargs) + return original_batch_store_block(paths, *args, **kwargs) - monkeypatch.setattr(mgr_mod, "store_block", flaky_store_block) + monkeypatch.setattr(mgr_mod, "batch_store_block", flaky_batch_store_block) tier.submit_store(make_job(1, [key(1), key(2)], [0, 1])) results = drain(tier) diff --git a/vllm/v1/kv_offload/tiering/fs/io.py b/vllm/v1/kv_offload/tiering/fs/io.py index 7785591fd2f..7c70fbae9f1 100644 --- a/vllm/v1/kv_offload/tiering/fs/io.py +++ b/vllm/v1/kv_offload/tiering/fs/io.py @@ -8,6 +8,18 @@ import os import random import threading +try: + from vllm.fs_io_C import ( # pyright: ignore[reportMissingImports] + batch_load_block as batch_load_block_C, + ) + from vllm.fs_io_C import ( + batch_store_block as batch_store_block_C, + ) + + _HAS_FSIO_C = True +except ImportError: + _HAS_FSIO_C = False + logger = logging.getLogger(__name__) # O_DIRECT is Linux-specific and not available on macOS @@ -59,7 +71,23 @@ def _ensure_dirs(path: str) -> None: os.makedirs(os.path.dirname(path), exist_ok=True) -def store_block( +def _validate_offsets(view: memoryview, offsets: list[int], block_size: int) -> None: + """Raise if any block would read/write past the bounds of `view`. + + Without this, an out-of-range offset silently clips to a shorter (or + empty) slice instead of failing, since memoryview slicing follows + Python's slice-clamping semantics rather than raising. + """ + total_len = len(view.cast("B")) + for offset in offsets: + if offset < 0 or offset + block_size > total_len: + raise ValueError( + f"block offset {offset} (block_size {block_size}) is out of " + f"bounds for a buffer of size {total_len}" + ) + + +def _store_block( dest_path: str, buffer: memoryview, offset: int, @@ -104,7 +132,7 @@ def store_block( raise -def load_block( +def _load_block( source_path: str, view: memoryview, offset: int, @@ -117,6 +145,7 @@ def load_block( fd: int | None = None view_slice = view.cast("B")[offset : offset + block_size] o_direct = O_DIRECT if use_o_direct else 0 + try: fd = os.open(source_path, os.O_RDONLY | o_direct) bytes_read = os.readv(fd, [view_slice]) @@ -133,3 +162,52 @@ def load_block( finally: if fd is not None: os.close(fd) + + +def batch_store_block( + paths: list[str], + view: memoryview, + offsets: list[int], + block_size: int, + use_o_direct: bool = True, +) -> None: + """ + Store a batch of KV blocks from a shared buffer to disk in one call. + + Each block buffer[offsets[i] : offsets[i]+block_size] is written atomically + to dest_paths[i] via a temp-file rename. Raises on first error. + """ + _validate_offsets(view, offsets, block_size) + + if _HAS_FSIO_C: + view_B = view.cast("B") + view_slices = [view_B[x : x + block_size] for x in offsets] + tmp_paths = [p + _get_tmp_suffix() for p in paths] + return batch_store_block_C(tmp_paths, paths, view_slices, use_o_direct) + else: + for path, offset in zip(paths, offsets): + _store_block(path, view, offset, block_size, use_o_direct) + + +def batch_load_block( + paths: list[str], + view: memoryview, + offsets: list[int], + block_size: int, + use_o_direct: bool = True, +) -> None: + """ + Load a batch of KV blocks from disk into a shared buffer in one call. + + Block i is read from source_paths[i] into view[offsets[i] : offsets[i]+block_size]. + Raises on first error and removes the offending file. + """ + _validate_offsets(view, offsets, block_size) + + if _HAS_FSIO_C: + view_B = view.cast("B") + view_slices = [view_B[x : x + block_size] for x in offsets] + return batch_load_block_C(paths, view_slices, use_o_direct) + else: + for path, offset in zip(paths, offsets): + _load_block(path, view, offset, block_size, use_o_direct) diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index 4fe200629f3..1c8d8b26676 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -49,7 +49,11 @@ from vllm.v1.kv_offload.tiering.base import ( ScheduleEndContext, SecondaryTierManager, ) -from vllm.v1.kv_offload.tiering.fs.io import load_block, probe_o_direct, store_block +from vllm.v1.kv_offload.tiering.fs.io import ( + batch_load_block, + batch_store_block, + probe_o_direct, +) from vllm.v1.kv_offload.tiering.fs.thread_pool import DualQueueThreadPool if TYPE_CHECKING: @@ -203,33 +207,28 @@ class FileSystemTierManager(SecondaryTierManager): def submit_store(self, job_metadata: JobMetadata) -> None: if self.events is not None: self._store_job_keys[job_metadata.job_id] = list(job_metadata.keys) - tasks = ( - functools.partial( - store_block, - self.file_mapper.get_file_name(key), - self._primary_kv_view, - int(bid) * self._block_size, - self._block_size, - self._use_o_direct, - ) - for key, bid in zip(job_metadata.keys, job_metadata.block_ids) + task = functools.partial( + batch_store_block, + [self.file_mapper.get_file_name(key) for key in job_metadata.keys], + self._primary_kv_view, + [int(bid) * self._block_size for bid in job_metadata.block_ids], + self._block_size, + self._use_o_direct, ) - self._pool.enqueue_store(job_metadata.job_id, len(job_metadata.keys), tasks) + self._pool.enqueue_store(job_metadata.job_id, 1, [task]) @override def submit_load(self, job_metadata: JobMetadata) -> None: - tasks = ( - functools.partial( - load_block, - self.file_mapper.get_file_name(key), - self._primary_kv_view, - int(bid) * self._block_size, - self._block_size, - self._use_o_direct, - ) - for key, bid in zip(job_metadata.keys, job_metadata.block_ids) + task = functools.partial( + batch_load_block, + [self.file_mapper.get_file_name(key) for key in job_metadata.keys], + self._primary_kv_view, + [int(bid) * self._block_size for bid in job_metadata.block_ids], + self._block_size, + self._use_o_direct, ) - self._pool.enqueue_load(job_metadata.job_id, len(job_metadata.keys), tasks) + + self._pool.enqueue_load(job_metadata.job_id, 1, [task]) @override def get_finished_jobs(self) -> Iterable[JobResult]: From 33fe71a4d3316c847ccd44a35a2055a361148600 Mon Sep 17 00:00:00 2001 From: fxmarty <9808326+fxmarty@users.noreply.github.com> Date: Tue, 28 Jul 2026 05:31:23 +0200 Subject: [PATCH 146/185] [AMD] Revert `Mxfp4MoeBackend.TRITON_UNFUSED` fallback (#46491) Signed-off-by: Felix Marty <Felix.Marty@amd.com> Co-authored-by: Felix Marty <Felix.Marty@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> --- tests/kernels/moe/test_ocp_mx_moe.py | 48 +++++++++++++++++++ tests/quantization/test_gfx950_moe.py | 15 ------ .../layers/fused_moe/oracle/mxfp4.py | 35 ++++++-------- 3 files changed, 63 insertions(+), 35 deletions(-) diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index 2f819c09aaa..7e8d4e3028a 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -1558,3 +1558,51 @@ def test_mxfp4_emulation_rounds_up_to_block_size( # The block-scale buffer (dim // OCP_MX_BLOCK_SIZE) must not floor-truncate. assert rounded_hidden % OCP_MX_BLOCK_SIZE == 0 assert rounded_intermediate % OCP_MX_BLOCK_SIZE == 0 + + +def test_select_mxfp4_moe_backend_raises_with_unsupported_reasons( + monkeypatch: pytest.MonkeyPatch, +): + """ + select_mxfp4_moe_backend() must raise NotImplementedError, with the + collected per-backend unsupported reasons in the message, when no + backend supports the requested deployment configuration. + """ + import vllm.model_executor.layers.fused_moe.oracle.mxfp4 as mxfp4_oracle + from vllm.model_executor.layers.fused_moe import FusedMoEConfig + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEParallelConfig, + RoutingMethodType, + ) + + class UnsupportedExperts: + @staticmethod + def is_supported_config( + cls, moe_config, weight_key, activation_key, activation_format + ): + return False, f"unsupported reason for {cls.__name__}" + + monkeypatch.setattr( + mxfp4_oracle, "backend_to_kernel_cls", lambda backend: [UnsupportedExperts] + ) + monkeypatch.setattr(mxfp4_oracle, "_user_moe_activation_override", lambda: None) + monkeypatch.setattr(current_platform, "is_xpu", lambda: False) + monkeypatch.setattr(current_platform, "is_cpu", lambda: False) + + moe_config = FusedMoEConfig( + num_experts=8, + experts_per_token=2, + hidden_dim=256, + intermediate_size=256, + num_local_experts=8, + num_logical_experts=8, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device="cpu", + routing_method=RoutingMethodType.Renormalize, + ) + + with pytest.raises(NotImplementedError, match="Unsupported reasons"): + mxfp4_oracle.select_mxfp4_moe_backend(moe_config) diff --git a/tests/quantization/test_gfx950_moe.py b/tests/quantization/test_gfx950_moe.py index 0efcc8a3c62..c8d34bb0ab5 100644 --- a/tests/quantization/test_gfx950_moe.py +++ b/tests/quantization/test_gfx950_moe.py @@ -79,21 +79,6 @@ def test_w4a4_dispatches_to_aiter(mxfp4_oracle_config): assert experts_cls is not None -@pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)") -@pytest.mark.skipif( - ROCM_AITER_AVAILABLE, - reason="Test requires AITER disabled (unset VLLM_ROCM_USE_AITER)", -) -def test_w4a4_falls_back_to_triton_unfused_without_aiter(mxfp4_oracle_config): - """Without AITER and no --moe-backend, ROCm falls back to TRITON_UNFUSED.""" - config = _make_w4a4_moe_config() - backend, experts_cls = select_mxfp4_moe_backend( - config, activation_key=kMxfp4Dynamic - ) - assert backend == Mxfp4MoeBackend.TRITON_UNFUSED - assert experts_cls is not None - - @pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)") def test_w4a4_dispatches_to_emulation_with_moe_backend(mxfp4_oracle_config): """With --moe-backend emulation, W4A4 selects EMULATION.""" diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 921e8f114d5..07cff8fa22c 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -502,6 +502,7 @@ def select_mxfp4_moe_backend( _get_priority_backends_for_gpt_oss(), requested_activation_key ) + unsupported_reasons = [] for backend in AVAILABLE_BACKENDS: # Use requested_activation_key if provided, otherwise use backend default act_key = ( @@ -518,6 +519,7 @@ def select_mxfp4_moe_backend( return backend, k_cls else: logger.debug_once(_make_log_unsupported(backend, reason)) + unsupported_reasons.append((backend, reason)) if current_platform.is_xpu(): backend = Mxfp4MoeBackend.XPU @@ -541,26 +543,19 @@ def select_mxfp4_moe_backend( activation_format, ) - if current_platform.is_rocm(): - backend = Mxfp4MoeBackend.TRITON_UNFUSED - logger.info_once(_make_log_backend(backend)) - return _return_or_raise( - Mxfp4MoeBackend.TRITON_UNFUSED, - config, - kMxfp4Static, - None, - activation_format, - ) - - if current_platform.is_cuda(): - raise NotImplementedError( - "No MXFP4 MoE backend supports the deployment configuration. " - f"weight_key=kMxfp4Static, activation_key={activation_key}. " - "Native backends require specific hardware. " - "Set `VLLM_LOGGING_LEVEL=DEBUG` to see detailed unsupported reasons. " - ) - - return Mxfp4MoeBackend.NONE, None + unsupported_log = "; ".join( + [ + f"backend: {backend.value}, reason: {reason}" + for backend, reason in unsupported_reasons + ] + ) + raise NotImplementedError( + "No MXFP4 MoE backend supports the deployment configuration. " + f"weight_key=kMxfp4Static, activation_key={activation_key}. " + f"Candidate backends were: " + f"{[backend.value for backend in AVAILABLE_BACKENDS]}. " + f"Unsupported reasons: {unsupported_log}. " + ) def select_deepseek_v4_mxfp4_moe_backend( From fbb1ef68030991a14291231bf8877063b57e1ada Mon Sep 17 00:00:00 2001 From: Colin Z <59755453+ColinZ22@users.noreply.github.com> Date: Mon, 27 Jul 2026 20:32:42 -0700 Subject: [PATCH 147/185] [Bugfix] Fix DeepseekV4FP8 Quark MXFP4 crash on list-valued weight (#49634) Signed-off-by: Colin Zeng <Colin.Zeng@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> --- vllm/models/deepseek_v4/quant_config.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/vllm/models/deepseek_v4/quant_config.py b/vllm/models/deepseek_v4/quant_config.py index 89cf695baf0..293d71f2f41 100644 --- a/vllm/models/deepseek_v4/quant_config.py +++ b/vllm/models/deepseek_v4/quant_config.py @@ -120,7 +120,11 @@ class DeepseekV4FP8Config(Fp8Config): @staticmethod def _is_quark_mxfp4_ocp(hf_quant_cfg: dict) -> bool: """True for AMD-Quark exports whose global scheme is MXFP4.""" - weight = (hf_quant_cfg.get("global_quant_config") or {}).get("weight") or {} + weight = (hf_quant_cfg.get("global_quant_config") or {}).get("weight") + # A non-dict weight (e.g. a list of multiple specs) means not an OCP + # MXFP4 scheme (e.g. NVFP4 with 2-level scale). + if not isinstance(weight, dict): + return False return ( weight.get("dtype") == "fp4" and weight.get("qscheme") == "per_group" From a8f296083f43e236f908d4d898ccfb0b108db1f6 Mon Sep 17 00:00:00 2001 From: Chang Guo <changg@nvidia.com> Date: Mon, 27 Jul 2026 21:10:20 -0700 Subject: [PATCH 148/185] [KV Offload] Make compact secondary identity TP-independent (#49858) Signed-off-by: Change72 <changg@nvidia.com> Co-authored-by: GPT-5.6 Sol <noreply@openai.com> Co-authored-by: Or Ozeri <oro@il.ibm.com> --- tests/v1/kv_offload/test_file_mapper.py | 105 +++++++++++++++++++ tests/v1/kv_offload/tiering/test_fs_tier.py | 66 +++++++++++- tests/v1/kv_offload/tiering/test_obj_tier.py | 55 +++++++++- vllm/v1/kv_offload/file_mapper.py | 11 +- 4 files changed, 226 insertions(+), 11 deletions(-) diff --git a/tests/v1/kv_offload/test_file_mapper.py b/tests/v1/kv_offload/test_file_mapper.py index c2c4427e184..027f01133b2 100644 --- a/tests/v1/kv_offload/test_file_mapper.py +++ b/tests/v1/kv_offload/test_file_mapper.py @@ -51,6 +51,7 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper: data_parallel_index=0, is_parallelism_agnostic=kwargs.get("is_parallelism_agnostic", False), ), + replicated_layout=kwargs.get("replicated_layout", False), ) spec = MagicMock(spec=OffloadingSpec) spec.config = config @@ -205,3 +206,107 @@ def test_parallel_agnostic_separates_persistent_layouts(): assert agnostic.base_path != specific.base_path assert "parallel_agnostic" not in agnostic.fields assert specific.fields["parallel_agnostic"] is False + + +# --------------------------------------------------------------------------- +# replicated_layout: OR'd into parallel-agnostic identity for compact rows +# --------------------------------------------------------------------------- + + +def test_replicated_layout_collapses_parallel_identity(): + shared = dict( + model_name="mla-model", + groups=((16, "mla_layer"),), + replicated_layout=True, + parallel_agnostic=True, + ) + tp2 = make_mapper_from_offloading_spec(tp_size=2, world_size=2, rank=1, **shared) + tp4 = make_mapper_from_offloading_spec(tp_size=4, world_size=4, rank=3, **shared) + + assert tp2.base_path == tp4.base_path + for fm in (tp2, tp4): + assert fm.fields["tp_size"] == 1 + assert fm.fields["pp_size"] == 1 + assert fm.fields["pcp_size"] == 1 + assert fm.fields["dcp_size"] == 1 + assert fm.rank == 0 + assert "parallel_agnostic" not in fm.fields + assert fm.fields["replicated_layout"] is True + assert fm.get_file_name(make_offload_key(b"\x01" * 8, 0)).startswith( + f"{fm.base_path}_r0/" + ) + + +def test_replicated_layout_requires_caller_opt_in(): + fm = make_mapper_from_offloading_spec( + tp_size=2, + world_size=2, + rank=1, + replicated_layout=True, + parallel_agnostic=False, + ) + assert fm.fields["tp_size"] == 2 + assert fm.rank == 1 + assert fm.fields["parallel_agnostic"] is False + assert "replicated_layout" not in fm.fields + baseline = make_mapper_from_offloading_spec( + tp_size=2, + world_size=2, + rank=1, + replicated_layout=False, + parallel_agnostic=False, + ) + assert fm.base_path == baseline.base_path + + +def test_non_replicated_keeps_parallel_identity(): + fm = make_mapper_from_offloading_spec( + tp_size=4, + world_size=4, + rank=2, + replicated_layout=False, + is_parallelism_agnostic=False, + parallel_agnostic=True, + ) + assert fm.fields["tp_size"] == 4 + assert fm.rank == 2 + assert fm.fields["parallel_agnostic"] is False + assert fm.get_file_name(make_offload_key(b"\x02" * 8, 0)).startswith( + f"{fm.base_path}_r2/" + ) + + +def test_replicated_and_parallelism_agnostic_separate_layouts(): + shared = dict( + model_name="shared-model", + groups=((16, "layer0"),), + tp_size=2, + world_size=2, + rank=1, + parallel_agnostic=True, + ) + via_agnostic = make_mapper_from_offloading_spec( + is_parallelism_agnostic=True, + replicated_layout=False, + **shared, + ) + via_replicated = make_mapper_from_offloading_spec( + is_parallelism_agnostic=False, + replicated_layout=True, + **shared, + ) + assert via_agnostic.base_path != via_replicated.base_path + assert "replicated_layout" not in via_agnostic.fields + assert via_replicated.fields["replicated_layout"] is True + + +def test_replicated_layout_run_config_tp_invariant(): + shared = dict( + model_name="mla-model", + groups=((16, "mla_layer"),), + replicated_layout=True, + parallel_agnostic=True, + ) + tp2 = make_mapper_from_offloading_spec(tp_size=2, world_size=2, rank=0, **shared) + tp4 = make_mapper_from_offloading_spec(tp_size=4, world_size=4, rank=2, **shared) + assert tp2.get_run_config() == tp4.get_run_config() diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index 66533e3676b..2959ac1aa03 100644 --- a/tests/v1/kv_offload/tiering/test_fs_tier.py +++ b/tests/v1/kv_offload/tiering/test_fs_tier.py @@ -52,8 +52,18 @@ _DTYPE: torch.dtype = torch.float32 _CTX = ReqContext(req_id="test") -def _make_offloading_spec(enable_kv_cache_events: bool) -> MagicMock: +def _make_offloading_spec( + enable_kv_cache_events: bool = False, + *, + tp_size: int = 1, + rank: int = 0, + world_size: int | None = None, + replicated_layout: bool = False, + is_parallelism_agnostic: bool = False, +) -> MagicMock: """Mock spec with an explicit global KV events flag.""" + if world_size is None: + world_size = tp_size spec = MagicMock() spec.config = OffloadingConfig( groups=(), @@ -64,15 +74,16 @@ def _make_offloading_spec(enable_kv_cache_events: bool) -> MagicMock: model=OffloadingModelConfig(name="test-model", dtype="float32"), cache=OffloadingCacheConfig(tokens_per_hash=16, blocks_per_chunk=1), parallel=OffloadingParallelConfig( - rank=0, - world_size=1, - tp_size=1, + rank=rank, + world_size=world_size, + tp_size=tp_size, pp_size=1, pcp_size=1, dcp_size=1, data_parallel_index=0, - is_parallelism_agnostic=False, + is_parallelism_agnostic=is_parallelism_agnostic, ), + replicated_layout=replicated_layout, ) spec.blocks_per_chunk = 1 spec.kv_events_config = OffloadingKVEventsConfig( @@ -725,3 +736,48 @@ def test_cascade_store_emits_fs_event_through_tiering_manager(tmp_path): assert not fs_events[0].removed finally: tier.shutdown() + + +def test_fs_tier_cross_tp_round_trip(tmp_path): + """TP=2 replicated writer and TP=4 reader share namespace and bytes.""" + root = str(tmp_path) + writer_tensor = _page_aligned_rand_tensor(4, _BLOCK_ELEMENTS) + expected = writer_tensor[0].clone() + writer = FileSystemTierManager( + offloading_spec=_make_offloading_spec( + tp_size=2, world_size=2, rank=0, replicated_layout=True + ), + primary_kv_view=memoryview(writer_tensor.numpy()), + tier_type="fs", + root_dir=root, + n_read_threads=2, + n_write_threads=2, + ) + try: + writer.submit_store(make_job(1, [key(7)], [0])) + assert all(r.success for r in drain(writer)) + writer_base = writer.file_mapper.base_path + writer_path = writer.file_mapper.get_file_name(key(7)) + finally: + writer.shutdown() + + reader_tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS) + reader = FileSystemTierManager( + offloading_spec=_make_offloading_spec( + tp_size=4, world_size=4, rank=3, replicated_layout=True + ), + primary_kv_view=memoryview(reader_tensor.numpy()), + tier_type="fs", + root_dir=root, + n_read_threads=2, + n_write_threads=2, + ) + try: + assert reader.file_mapper.base_path == writer_base + assert reader.file_mapper.get_file_name(key(7)) == writer_path + assert lookup_and_wait(reader, [key(7)]) == [LookupResult.HIT] + reader.submit_load(make_job(2, [key(7)], [1], is_promotion=True)) + assert all(r.success for r in drain(reader)) + assert torch.allclose(reader_tensor[1], expected) + finally: + reader.shutdown() diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index fc30e1437a7..82ba183ea5f 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -43,7 +43,17 @@ from vllm.v1.kv_offload.tiering.obj.manager import ObjectStoreSecondaryTierManag # --------------------------------------------------------------------------- -def _make_offloading_config(enable_kv_cache_events: bool) -> OffloadingConfig: +def _make_offloading_config( + enable_kv_cache_events: bool, + *, + tp_size: int = 1, + rank: int = 0, + world_size: int | None = None, + replicated_layout: bool = False, + is_parallelism_agnostic: bool = False, +) -> OffloadingConfig: + if world_size is None: + world_size = tp_size return OffloadingConfig( groups=(), worker_kv_bytes_per_block=0, @@ -53,15 +63,16 @@ def _make_offloading_config(enable_kv_cache_events: bool) -> OffloadingConfig: model=OffloadingModelConfig(name="test/model", dtype="float16"), cache=OffloadingCacheConfig(tokens_per_hash=16, blocks_per_chunk=1), parallel=OffloadingParallelConfig( - rank=0, - world_size=1, - tp_size=1, + rank=rank, + world_size=world_size, + tp_size=tp_size, pp_size=1, pcp_size=1, dcp_size=1, data_parallel_index=0, - is_parallelism_agnostic=False, + is_parallelism_agnostic=is_parallelism_agnostic, ), + replicated_layout=replicated_layout, ) @@ -617,3 +628,37 @@ class TestObjStoreConfig: params = cfg.to_nixl_params() assert params["ca_bundle"] == "/path/to/ca.pem" assert "access_key" not in params + + +def test_obj_tier_replicated_layout_collapses_mapper_identity(): + """TP=2 and TP=4 replicated configs share the obj FileMapper namespace.""" + tp2_spec = SimpleNamespace( + config=_make_offloading_config( + False, tp_size=2, world_size=2, rank=1, replicated_layout=True + ), + kv_events_config=OffloadingKVEventsConfig( + enable_kv_cache_events=False, + self_describing_kv_events=False, + ), + ) + tp4_spec = SimpleNamespace( + config=_make_offloading_config( + False, tp_size=4, world_size=4, rank=3, replicated_layout=True + ), + kv_events_config=OffloadingKVEventsConfig( + enable_kv_cache_events=False, + self_describing_kv_events=False, + ), + ) + tp2_tier, _ = _make_tier(offloading_spec=tp2_spec) + tp4_tier, _ = _make_tier(offloading_spec=tp4_spec) + try: + assert tp2_tier._file_mapper.base_path == tp4_tier._file_mapper.base_path + assert tp2_tier._file_mapper.rank == 0 + assert tp4_tier._file_mapper.rank == 0 + assert tp2_tier._file_mapper.get_run_config() == ( + tp4_tier._file_mapper.get_run_config() + ) + finally: + tp2_tier.shutdown() + tp4_tier.shutdown() diff --git a/vllm/v1/kv_offload/file_mapper.py b/vllm/v1/kv_offload/file_mapper.py index 8e8c19d53d6..4b12dba913d 100644 --- a/vllm/v1/kv_offload/file_mapper.py +++ b/vllm/v1/kv_offload/file_mapper.py @@ -35,6 +35,7 @@ class FileMapper: kv_cache_groups: list[dict] | None = None, inference_engine: str = "vllm", parallel_agnostic: bool = False, + replicated_layout: bool = False, ): """ Initialize the file mapper. Each worker constructs its own, but @@ -60,6 +61,10 @@ class FileMapper: } if not parallel_agnostic: self.fields["parallel_agnostic"] = False + # Only written when True so existing deployments' hashed fields are + # unchanged (False is the historical default and must not appear). + if replicated_layout: + self.fields["replicated_layout"] = True self.base_path: str = self._compute_base_path(root_dir, self.fields) @classmethod @@ -92,7 +97,11 @@ class FileMapper: rank=parallel.rank, dtype=config.model.dtype, kv_cache_groups=kv_cache_groups, - parallel_agnostic=(parallel_agnostic and parallel.is_parallelism_agnostic), + parallel_agnostic=( + parallel_agnostic + and (parallel.is_parallelism_agnostic or config.replicated_layout) + ), + replicated_layout=(parallel_agnostic and config.replicated_layout), ) def get_file_name(self, key: OffloadKey) -> str: From 52c3c4a42fd13b62ba985b9ceb9b9969964ee83e Mon Sep 17 00:00:00 2001 From: MINJUN GIL <alswnsrlf12@naver.com> Date: Tue, 28 Jul 2026 13:10:52 +0900 Subject: [PATCH 149/185] [Bugfix][KV Offload][OBJ] Preserve job completion during cleanup (#49947) Signed-off-by: MINJUN GIL <alswnsrlf12@naver.com> Co-authored-by: OpenAI Codex <codex@openai.com> --- tests/v1/kv_offload/tiering/test_obj_tier.py | 111 ++++++++++++++++++- vllm/v1/kv_offload/tiering/obj/manager.py | 29 ++++- 2 files changed, 133 insertions(+), 7 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index 82ba183ea5f..661438dce63 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -35,6 +35,10 @@ from vllm.v1.kv_offload.config import ( OffloadingParallelConfig, ) from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult +from vllm.v1.kv_offload.tiering.manager import ( + CPUPrimaryTierOffloadingManager, + TieringOffloadingManager, +) from vllm.v1.kv_offload.tiering.obj.config import ObjStoreConfig from vllm.v1.kv_offload.tiering.obj.manager import ObjectStoreSecondaryTierManager @@ -220,12 +224,14 @@ def _make_events_spec(enable_kv_cache_events: bool) -> SimpleNamespace: def _make_tier( num_blocks: int = 4, offloading_spec: SimpleNamespace = _OFFLOADING_SPEC, + primary_kv_view: memoryview | None = None, **tier_kwargs, ) -> tuple[ObjectStoreSecondaryTierManager, MockNixlAgent]: """Create a tier backed by a fresh MockNixlAgent.""" mock_agent = MockNixlAgent() - tensor = torch.zeros((num_blocks, _BLOCK_ELEMENTS), dtype=_DTYPE) - view = memoryview(tensor.numpy()) + if primary_kv_view is None: + tensor = torch.zeros((num_blocks, _BLOCK_ELEMENTS), dtype=_DTYPE) + primary_kv_view = memoryview(tensor.numpy()) with ( patch("vllm.v1.kv_offload.tiering.obj.manager.nixl_agent_config"), patch( @@ -235,7 +241,7 @@ def _make_tier( ): tier = ObjectStoreSecondaryTierManager( offloading_spec=offloading_spec, - primary_kv_view=view, + primary_kv_view=primary_kv_view, tier_type="obj", store_config=_STORE_CONFIG, prefix=_RUN_PREFIX, @@ -449,6 +455,105 @@ class TestMockObjTierFailures: assert not by_id[1].success assert by_id[2].success + def test_release_xfer_failure_retries_without_losing_result(self, monkeypatch): + tier, agent = _make_tier(num_blocks=4) + agent.check_xfer_state = MagicMock(side_effect=RuntimeError("poll failed")) + release_xfer = MagicMock( + side_effect=[RuntimeError("transfer is still active"), None] + ) + monkeypatch.setattr(agent, "release_xfer_handle", release_xfer) + + tier.submit_store(make_job(1, [key(1)], [0])) + + # The transfer handle could not be released safely, so the job must + # remain tracked and must not be finalized yet. + assert list(tier.get_finished_jobs()) == [] + assert 1 in tier._transfers + + # Cleanup is retried without polling again or changing the failure + # verdict. The completion is then returned exactly once. + results = list(tier.get_finished_jobs()) + assert len(results) == 1 + assert results[0].job_id == 1 + assert not results[0].success + assert agent.check_xfer_state.call_count == 2 + assert release_xfer.call_count == 2 + assert not tier._transfers + assert list(tier.get_finished_jobs()) == [] + + @pytest.mark.parametrize( + "cleanup_method", ["release_dlist_handle", "deregister_memory"] + ) + def test_post_transfer_cleanup_failure_does_not_lose_result( + self, monkeypatch, cleanup_method + ): + tier, agent = _make_tier(num_blocks=4) + monkeypatch.setattr( + agent, + cleanup_method, + MagicMock(side_effect=RuntimeError("cleanup failed")), + ) + + tier.submit_store(make_job(1, [key(1)], [0])) + results = list(tier.get_finished_jobs()) + + assert len(results) == 1 + assert results[0].job_id == 1 + assert results[0].success + assert not tier._transfers + assert list(tier.get_finished_jobs()) == [] + + def test_xfer_cleanup_retry_finalizes_parent_job_and_primary_pin(self, monkeypatch): + num_blocks = 4 + tensor = torch.zeros((num_blocks, _BLOCK_ELEMENTS), dtype=_DTYPE) + primary_kv_view = memoryview(tensor.numpy()) + mmap_region = MagicMock() + mmap_region.create_kv_memoryview.return_value = primary_kv_view + primary_tier = CPUPrimaryTierOffloadingManager( + num_blocks=num_blocks, mmap_region=mmap_region + ) + obj_tier, agent = _make_tier( + num_blocks=num_blocks, primary_kv_view=primary_kv_view + ) + manager = TieringOffloadingManager( + primary_tier=primary_tier, secondary_tiers=[obj_tier] + ) + + keys = [key(1)] + primary_result = primary_tier.prepare_store(keys, _CTX) + assert primary_result is not None + primary_tier.complete_store(keys, _CTX, success=True) + job = manager.create_store_job(keys, _CTX) + obj_tier.submit_store(job) + + block = primary_tier._policy.get(keys[0]) + assert block is not None + assert block.ref_cnt == 1 + assert len(manager._transfer_jobs) == 1 + + agent.check_xfer_state = MagicMock(side_effect=RuntimeError("poll failed")) + release_xfer = MagicMock( + side_effect=[RuntimeError("transfer is still active"), None] + ) + monkeypatch.setattr(agent, "release_xfer_handle", release_xfer) + schedule_context = ScheduleEndContext(new_req_ids=[], preempted_req_ids=()) + + manager.on_schedule_end(schedule_context) + + assert len(manager._transfer_jobs) == 1 + assert block.ref_cnt == 1 + assert len(obj_tier._transfers) == 1 + assert manager.has_pending_work() + + manager.on_schedule_end(schedule_context) + + assert manager._transfer_jobs == {} + assert block.ref_cnt == 0 + assert obj_tier._transfers == {} + assert not manager.has_pending_work() + assert agent.check_xfer_state.call_count == 2 + assert release_xfer.call_count == 2 + class TestMockObjTierShutdown: def test_shutdown_clears_in_flight_transfers(self): diff --git a/vllm/v1/kv_offload/tiering/obj/manager.py b/vllm/v1/kv_offload/tiering/obj/manager.py index c7e0d4c4beb..2dfc2d30fa2 100644 --- a/vllm/v1/kv_offload/tiering/obj/manager.py +++ b/vllm/v1/kv_offload/tiering/obj/manager.py @@ -307,15 +307,36 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager): else: if state == NIXL_PROC: continue - elif state == NIXL_DONE: + if state == NIXL_DONE: success = True else: success = False logger.warning("transfer failed job=%d state=%s", job_id, state) + + try: + self._agent.release_xfer_handle(entry.xfer_handle) + except Exception as exc: + # Keep the entry until NIXL confirms that the transfer handle + # can be released. The transfer may still access primary-tier + # memory, so publishing its result would allow unsafe reuse. + logger.warning("release_xfer_handle failed for job %d: %s", job_id, exc) + continue + + # Once the transfer handle is released, these remaining cleanup + # failures must not suppress the job completion. They can leak + # NIXL metadata, but cannot leave an active data transfer behind. + try: + self._agent.release_dlist_handle(entry.obj_handle) + except Exception as exc: + logger.warning( + "release_dlist_handle failed for job %d: %s", job_id, exc + ) + try: + self._agent.deregister_memory(entry.files_desc) + except Exception as exc: + logger.warning("deregister_memory failed for job %d: %s", job_id, exc) + del self._transfers[job_id] - self._agent.release_xfer_handle(entry.xfer_handle) - self._agent.release_dlist_handle(entry.obj_handle) - self._agent.deregister_memory(entry.files_desc) self._pending_results.append(JobResult(job_id=job_id, success=success)) def get_finished_jobs(self) -> Iterable[JobResult]: From d223c900d85224c02f2162ee2c757a769e99f519 Mon Sep 17 00:00:00 2001 From: Nick Hill <nickhill123@gmail.com> Date: Mon, 27 Jul 2026 21:36:45 -0700 Subject: [PATCH 150/185] [Bugfix] Only pad transformers backend `value` when it is narrower (#50060) Signed-off-by: Nick Hill <nickhill123@gmail.com> --- vllm/model_executor/models/transformers/__init__.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/models/transformers/__init__.py b/vllm/model_executor/models/transformers/__init__.py index 9dbfe4b5031..ff4bb6cd433 100644 --- a/vllm/model_executor/models/transformers/__init__.py +++ b/vllm/model_executor/models/transformers/__init__.py @@ -64,12 +64,15 @@ def vllm_attention_forward( head_dim_v = value.shape[-1] query, key, value = (x.transpose(1, 2) for x in (query, key, value)) query, key, value = (x.reshape(hidden, -1) for x in (query, key, value)) - # Pad `value` up to the query/key head size when they differ (expanded MLA). - if head_dim_v != head_dim_qk: + # Pad `value` up to the query/key head size when it is smaller (expanded + # MLA). A larger last dim just means `value` isn't split per head, e.g. + # packed grouped/multi-query projections, and needs no padding. + pad_value = head_dim_v < head_dim_qk + if pad_value: value = F.pad(value.view(-1, head_dim_v), (0, head_dim_qk - head_dim_v)) value = value.reshape(hidden, -1) attn_output = self_attn.forward(query, key, value) - if head_dim_v != head_dim_qk: + if pad_value: attn_output = attn_output.view(-1, head_dim_qk)[..., :head_dim_v] attn_output = attn_output.reshape(hidden, -1) return attn_output, None From 74587939b17b4eb3ba281ec9ca43f93aa230a3cb Mon Sep 17 00:00:00 2001 From: Ayushman Singh <40520701+ayush1399@users.noreply.github.com> Date: Tue, 28 Jul 2026 01:01:56 -0400 Subject: [PATCH 151/185] [Build] Fix CUDA arch detection producing kernel-less builds on SM121 (#49904) --- CMakeLists.txt | 13 +++++++++---- cmake/utils.cmake | 11 ++++++----- tests/test_cmake_utils.py | 25 +++++++++++++++++++++++++ 3 files changed, 40 insertions(+), 9 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 9bd73f18e70..3a4e23ad0b3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -219,10 +219,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") # the set of architectures we want to compile for and remove the from the # CMAKE_CUDA_FLAGS so that they are not applied globally. # - # `+PTX` in TORCH_CUDA_ARCH_LIST is not preserved here. It is emitted by torch - # as `code=compute_*`, while extract_unique_cuda_archs_ascending() records only - # `arch=compute_*`. If a kernel really needs PTX, add `+PTX` to that kernel's - # component-specific arch list below. + # `+PTX` in TORCH_CUDA_ARCH_LIST is not preserved here. If a kernel really + # needs PTX, add `+PTX` to that kernel's component-specific arch list below. # clear_cuda_arches(CUDA_ARCH_FLAGS) extract_unique_cuda_archs_ascending(CUDA_ARCHS "${CUDA_ARCH_FLAGS}") @@ -232,6 +230,13 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") cuda_archs_loose_intersection(CUDA_ARCHS "${CUDA_SUPPORTED_ARCHS}" "${CUDA_ARCHS}") message(STATUS "CUDA supported target architectures: ${CUDA_ARCHS}") + if(NOT CUDA_ARCHS) + message(FATAL_ERROR + "No supported CUDA architectures; the build would produce a binary " + "with no usable kernels. Detected gencode flags: ${CUDA_ARCH_FLAGS}; " + "supported: ${CUDA_SUPPORTED_ARCHS}. " + "Set TORCH_CUDA_ARCH_LIST for your GPU (e.g. 12.0).") + endif() else() # # For other GPU targets override the GPU architectures detected by cmake/torch diff --git a/cmake/utils.cmake b/cmake/utils.cmake index 14a94eebb22..bbae89c1f57 100644 --- a/cmake/utils.cmake +++ b/cmake/utils.cmake @@ -241,14 +241,15 @@ endmacro() # `<major>.<minor>`, dedupes them and then sorts them in ascending order and # stores them in `OUT_ARCHES`. # -# Example: -# CUDA_ARCH_FLAGS="-gencode arch=compute_75,code=sm_75;...;-gencode arch=compute_90a,code=sm_90a" -# extract_unique_cuda_archs_ascending(OUT_ARCHES CUDA_ARCH_FLAGS) -# OUT_ARCHES="7.5;...;9.0" +# Prefer `code=sm_*`; fall back to `arch=compute_*` for PTX-only flags. +# This handles mismatches such as `arch=compute_20,code=sm_121`. function(extract_unique_cuda_archs_ascending OUT_ARCHES CUDA_ARCH_FLAGS) set(_CUDA_ARCHES) foreach(_ARCH ${CUDA_ARCH_FLAGS}) - string(REGEX MATCH "arch=compute_\([0-9]+[af]?\)" _COMPUTE ${_ARCH}) + string(REGEX MATCH "code=sm_\([0-9]+[af]?\)" _COMPUTE ${_ARCH}) + if (NOT _COMPUTE) + string(REGEX MATCH "arch=compute_\([0-9]+[af]?\)" _COMPUTE ${_ARCH}) + endif() if (_COMPUTE) set(_COMPUTE ${CMAKE_MATCH_1}) endif() diff --git a/tests/test_cmake_utils.py b/tests/test_cmake_utils.py index 227ec231eb2..d0673bc462e 100644 --- a/tests/test_cmake_utils.py +++ b/tests/test_cmake_utils.py @@ -21,3 +21,28 @@ endif() ) subprocess.run(["cmake", "-P", script], check=True) + + +def test_extract_archs_prefers_sass_target_over_corrupted_virtual_arch( + tmp_path: Path, +): + """torch's autodetection can emit a bogus arch=compute_* half (e.g. + capability 12.1 corrupted to arch=compute_20,code=sm_121); the SASS + target must win, while PTX-only entries keep the virtual arch.""" + repo_root = Path(__file__).parents[1] + script = tmp_path / "test_extract_archs.cmake" + script.write_text( + f""" +cmake_minimum_required(VERSION 3.26) +include("{repo_root / "cmake" / "utils.cmake"}") +extract_unique_cuda_archs_ascending(actual + "-gencode arch=compute_20,code=sm_121;\ +-gencode arch=compute_80,code=sm_80;\ +-gencode arch=compute_80,code=compute_80") +if(NOT "${{actual}}" STREQUAL "8.0;12.1") + message(FATAL_ERROR "Expected '8.0;12.1', got '${{actual}}'") +endif() +""" + ) + + subprocess.run(["cmake", "-P", script], check=True) From f472ab0a4c0987d2badc2f2aeecd38fee4180a4c Mon Sep 17 00:00:00 2001 From: afriedri <afriedri@amd.com> Date: Tue, 28 Jul 2026 00:53:46 -0500 Subject: [PATCH 152/185] Remove triton per group quant [ROCm] [Bugfix] (#49621) Signed-off-by: Andy Friedrich <afriedri@amd.com> --- .../distributed/test_fusion_all_reduce.py | 31 ++++------------- tests/compile/passes/test_fusion.py | 2 -- .../passes/test_silu_mul_quant_fusion.py | 11 +----- .../passes/fusion/allreduce_rms_fusion.py | 3 +- .../layers/quantization/input_quant_fp8.py | 5 --- .../layers/quantization/utils/fp8_utils.py | 34 ------------------- 6 files changed, 8 insertions(+), 78 deletions(-) diff --git a/tests/compile/passes/distributed/test_fusion_all_reduce.py b/tests/compile/passes/distributed/test_fusion_all_reduce.py index 1aac4b2bec4..e9c8d0deaa7 100644 --- a/tests/compile/passes/distributed/test_fusion_all_reduce.py +++ b/tests/compile/passes/distributed/test_fusion_all_reduce.py @@ -272,12 +272,10 @@ class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module): token_num=16, eps=1e-6, dtype: torch.dtype = torch.bfloat16, - use_triton_quant: bool = False, ): super().__init__() self.hidden_size = hidden_size self.eps = eps - self.use_triton_quant = use_triton_quant assert hidden_size % self.quant_group_size == 0, ( f"hidden_size ({hidden_size}) must be a multiple of " f"quant_group_size ({self.quant_group_size}) for per-group FP8 quant" @@ -289,10 +287,6 @@ class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module): ] def _group_quant(self, rms: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - if self.use_triton_quant: - return torch.ops.vllm.triton_per_token_group_quant_fp8( - rms, self.quant_group_size - ) return torch.ops.vllm.rocm_aiter_group_fp8_quant.default( rms, self.quant_group_size ) @@ -339,11 +333,7 @@ class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module): def ops_in_model_before(self): return [ torch.ops.vllm.all_reduce.default, - ( - torch.ops.vllm.triton_per_token_group_quant_fp8.default - if self.use_triton_quant - else torch.ops.vllm.rocm_aiter_group_fp8_quant.default - ), + torch.ops.vllm.rocm_aiter_group_fp8_quant.default, ] def ops_in_model_after(self): @@ -646,7 +636,6 @@ def all_reduce_fusion_pass_on_test_model( @multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize("use_triton_quant", [True, False]) @pytest.mark.parametrize("batch_size", [8]) @pytest.mark.parametrize("seq_len", [8]) @pytest.mark.parametrize("hidden_size", [128]) @@ -663,7 +652,6 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace( hidden_size: int, dtype: torch.dtype, enable_rms_norm_custom_op: bool, - use_triton_quant: bool, monkeypatch: pytest.MonkeyPatch, ): """Sibling of ``test_all_reduce_fusion_pass_replace`` for the new @@ -676,9 +664,9 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace( * ``AiterAllreduceFusedAddRMSNormGroupQuantFP8Pattern`` (with-residual, single ``rms`` consumer) * ``AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern`` (with- - residual, DSv3.2 indexer fan-out; parametrized over both - ``triton_per_token_group_quant_fp8`` and ``rocm_aiter_group_fp8_quant`` - producers). + residual, DSv3.2 indexer fan-out; parametrized over + ``rocm_aiter_group_fp8_quant`` + producer). """ with monkeypatch.context() as m: m.setenv("VLLM_ROCM_USE_AITER", "1") @@ -703,7 +691,6 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace( hidden_size, dtype, enable_rms_norm_custom_op, - use_triton_quant, monkeypatch, ), nprocs=nprocs, @@ -721,7 +708,6 @@ def rocm_aiter_group_quant_fusion_pass_on_test_model( hidden_size: int, dtype: torch.dtype, enable_rms_norm_custom_op: bool, - use_triton_quant: bool, monkeypatch: pytest.MonkeyPatch, ): set_random_seed(0) @@ -749,10 +735,7 @@ def rocm_aiter_group_quant_fusion_pass_on_test_model( custom_ops = [] if enable_rms_norm_custom_op: custom_ops.append("+rms_norm") - # ``triton_per_token_group_quant_fp8`` is emitted by ``QuantFP8.forward_hip`` - # only when QuantFP8 is enabled as a custom op (and ``use_triton=True`` at - # the call site). The patterns in this PR are robust to both Triton and - # rocm_aiter forms; we always enable +quant_fp8 so the matcher's example + # We always enable +quant_fp8 so the matcher's example # trace finds the same form the test model uses. custom_ops.append("+quant_fp8") @@ -783,9 +766,7 @@ def rocm_aiter_group_quant_fusion_pass_on_test_model( ) token_num = batch_size * seq_len - model = test_model_cls( - hidden_size, token_num, dtype=dtype, use_triton_quant=use_triton_quant - ) + model = test_model_cls(hidden_size, token_num, dtype=dtype) hidden_states = torch.randn((token_num, hidden_size), requires_grad=False) diff --git a/tests/compile/passes/test_fusion.py b/tests/compile/passes/test_fusion.py index 92d1902b2c2..591b014d9e2 100644 --- a/tests/compile/passes/test_fusion.py +++ b/tests/compile/passes/test_fusion.py @@ -195,8 +195,6 @@ class TestModel(torch.nn.Module): # Blockwise path if self.use_aiter_fusion and self.use_aiter_quant_op: return [rocm_aiter_ops.get_group_quant_op()] - if self.use_aiter_fusion: - return [torch.ops.vllm.triton_per_token_group_quant_fp8.default] else: if self.use_aiter_quant_op: return [rocm_aiter_ops.get_per_token_quant_op()] diff --git a/tests/compile/passes/test_silu_mul_quant_fusion.py b/tests/compile/passes/test_silu_mul_quant_fusion.py index bc134ed427a..7d291cc5044 100644 --- a/tests/compile/passes/test_silu_mul_quant_fusion.py +++ b/tests/compile/passes/test_silu_mul_quant_fusion.py @@ -158,13 +158,6 @@ class TestSiluMulGroupFp8QuantModel(torch.nn.Module): input_dtype=dtype, ) - if not current_platform.is_fp8_fnuz(): - kernel = self.w8a8_block_fp8_linear.kernel - orig_quant = kernel.quant_fp8 - kernel.quant_fp8 = lambda *a, use_triton=False, **kw: orig_quant( - *a, use_triton=True, **kw - ) - self.enable_silu_mul_custom_op = self.silu_and_mul.enabled() def forward(self, x): @@ -175,9 +168,7 @@ class TestSiluMulGroupFp8QuantModel(torch.nn.Module): def ops_in_model_before(self): return [ SILU_MUL_OP if self.enable_silu_mul_custom_op else torch.ops.aten.mul, - rocm_aiter_ops.get_group_quant_op() - if current_platform.is_fp8_fnuz() - else torch.ops.vllm.triton_per_token_group_quant_fp8.default, + rocm_aiter_ops.get_group_quant_op(), ] def ops_in_model_after(self): diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index f7ff7df66cd..1722b524eeb 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -1416,8 +1416,7 @@ class AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern( The trailing FP8 group-quant is matched via ``MatcherQuantFP8`` (consistent with the sibling patterns above), which traces both ``QuantFP8.forward_hip`` and ``forward_native`` paths and so matches whichever op the call site - lowers to (``vllm.triton_per_token_group_quant_fp8`` or - ``vllm.rocm_aiter_group_fp8_quant``). + lowers to (``vllm.rocm_aiter_group_fp8_quant``). """ def __init__( diff --git a/vllm/model_executor/layers/quantization/input_quant_fp8.py b/vllm/model_executor/layers/quantization/input_quant_fp8.py index e8810919c20..2eb34630aa6 100644 --- a/vllm/model_executor/layers/quantization/input_quant_fp8.py +++ b/vllm/model_executor/layers/quantization/input_quant_fp8.py @@ -139,11 +139,6 @@ class QuantFP8(CustomOp): scale_ub: torch.Tensor | None = None, use_triton: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - if self.is_group_quant and use_triton: - assert scale is None, "Dynamic group quantization does not use scale" - - return torch.ops.vllm.triton_per_token_group_quant_fp8(x, self.group_size) - use_aiter_quant = self.use_aiter and scale_ub is None and x.is_contiguous() use_aiter_per_tensor_quant = ( use_aiter_quant and self.group_shape.is_per_tensor() diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index 83e56a4567b..2e4fbdf4c64 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -34,7 +34,6 @@ from vllm.utils.deep_gemm import ( transform_sf_into_required_layout, ) from vllm.utils.platform_utils import get_device_name_as_file_name -from vllm.utils.torch_utils import direct_register_custom_op logger = init_logger(__name__) @@ -45,39 +44,6 @@ def is_fp8(x: torch.dtype | torch.Tensor) -> bool: return x == torch.float8_e4m3fn or x == torch.float8_e4m3fnuz -def _triton_per_token_group_quant_fp8_impl( - x: torch.Tensor, - group_size: int, -) -> tuple[torch.Tensor, torch.Tensor]: - return per_token_group_quant_fp8( - x, group_size, column_major_scales=False, use_ue8m0=False - ) - - -def _triton_per_token_group_quant_fp8_fake( - x: torch.Tensor, - group_size: int, -) -> tuple[torch.Tensor, torch.Tensor]: - M, N = x.shape - x_fp8 = torch.empty((M, N), dtype=current_platform.fp8_dtype(), device=x.device) - out_bs = torch.empty( - ( - M, - (N + group_size - 1) // group_size, - ), - dtype=torch.float32, - device=x.device, - ) - return x_fp8, out_bs - - -direct_register_custom_op( - "triton_per_token_group_quant_fp8", - _triton_per_token_group_quant_fp8_impl, - fake_impl=_triton_per_token_group_quant_fp8_fake, -) - - def input_to_float8( x: torch.Tensor, dtype: torch.dtype | None = None ) -> tuple[torch.Tensor, torch.Tensor]: From 90245f4190a35593a625e4bc349485c39c774d39 Mon Sep 17 00:00:00 2001 From: "Li, Jiang" <jiang1.li@intel.com> Date: Tue, 28 Jul 2026 14:09:48 +0800 Subject: [PATCH 153/185] [Bugfix] Fix multi-modal support on CPU MRV2 (#50073) Signed-off-by: jiang1.li <jiang1.li@intel.com> --- vllm/multimodal/inputs.py | 2 ++ vllm/v1/worker/cpu/shm.py | 12 ++++++++++++ vllm/v1/worker/gpu/model_states/encoder_decoder.py | 5 ++++- 3 files changed, 18 insertions(+), 1 deletion(-) diff --git a/vllm/multimodal/inputs.py b/vllm/multimodal/inputs.py index 71f17a8648b..dc96c366ad2 100644 --- a/vllm/multimodal/inputs.py +++ b/vllm/multimodal/inputs.py @@ -457,6 +457,8 @@ class BaseMultiModalField(ABC): device = "cpu" if pin_memory and self.keep_on_cpu: pin_memory = False + if device == "cpu" or device == torch.device("cpu"): + pin_memory = False batch = [elem.data for elem in elems] out = self._reduce_data(batch, pin_memory=pin_memory) diff --git a/vllm/v1/worker/cpu/shm.py b/vllm/v1/worker/cpu/shm.py index bd1f96c71ed..970e8a2d414 100644 --- a/vllm/v1/worker/cpu/shm.py +++ b/vllm/v1/worker/cpu/shm.py @@ -24,12 +24,17 @@ def fake_pin_memory(self: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tens class _EventPlaceholder: def __init__(self, *args, **kwargs) -> None: self.record = noop + self.wait = noop self.synchronize = noop class _StreamPlaceholder: def __init__(self, *args, **kwargs) -> None: self.wait_stream = noop + self.wait_event = noop + self.record_event = noop + self.synchronize = noop + self.query = lambda: True self.device = torch.device("cpu") def __enter__(self, *args, **kwargs): @@ -55,6 +60,7 @@ torch.cuda.current_stream = lambda *args, **kwargs: _StreamPlaceholder() torch.accelerator.synchronize = noop torch.accelerator.empty_cache = noop torch.Tensor.pin_memory = fake_pin_memory +torch.Tensor.record_stream = noop torch.accelerator.get_memory_info = get_memory_info # Patch vLLM torch utils @@ -80,3 +86,9 @@ import vllm.v1.worker.gpu.buffer_utils as gpu_buffer_utils import vllm.v1.worker.cpu.buffer_utils as cpu_buffer_utils gpu_buffer_utils.UvaBuffer = cpu_buffer_utils.UvaBuffer + +# Patch Triton +from vllm.triton_utils import HAS_TRITON, tl + +if HAS_TRITON: + tl.debug_barrier = noop diff --git a/vllm/v1/worker/gpu/model_states/encoder_decoder.py b/vllm/v1/worker/gpu/model_states/encoder_decoder.py index f759c0b1e15..618984c97f3 100644 --- a/vllm/v1/worker/gpu/model_states/encoder_decoder.py +++ b/vllm/v1/worker/gpu/model_states/encoder_decoder.py @@ -9,6 +9,7 @@ import torch.nn as nn from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode +from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.kv_cache_interface import CrossAttentionSpec, KVCacheConfig from vllm.v1.worker.gpu.attn_utils import build_attn_metadata from vllm.v1.worker.gpu.input_batch import InputBatch @@ -157,7 +158,9 @@ class EncoderDecoderModelState(ModelState): for_capture: bool, num_reqs: int, ) -> dict[int, tuple[torch.Tensor, np.ndarray]]: - encoder_seq_lens = torch.zeros(num_reqs, dtype=torch.int32, pin_memory=True) + encoder_seq_lens = torch.zeros( + num_reqs, dtype=torch.int32, pin_memory=PIN_MEMORY + ) encoder_seq_lens_np = encoder_seq_lens.numpy() if not for_capture: # During normal execution, use actual encoder lengths. From 9069a57139bc075733ffb82c228d9f33945e524b Mon Sep 17 00:00:00 2001 From: Shuolei Wang <46160365+ShuoleiWang@users.noreply.github.com> Date: Tue, 28 Jul 2026 14:35:29 +0800 Subject: [PATCH 154/185] [Core][Frontend] Add weight version tagging for RL rollouts (#49040) Signed-off-by: Shuolei Wang <shuoleiwang123@gmail.com> Signed-off-by: Shuolei Wang <948904026@qq.com> --- docs/serving/offline_inference.md | 2 ++ docs/serving/online_serving/README.md | 2 ++ docs/training/async_rl.md | 3 +- docs/training/weight_transfer/README.md | 6 ++-- tests/distributed/test_weight_transfer.py | 9 +++++- .../entrypoints/openai/test_openai_schema.py | 1 + .../test_weight_transfer_llm.py | 12 +++++++- vllm/distributed/weight_transfer/base.py | 2 +- vllm/distributed/weight_transfer/clients.py | 13 ++++++-- vllm/engine/protocol.py | 12 ++++++-- vllm/entrypoints/llm.py | 14 +++++++-- vllm/entrypoints/serve/dev/rlhf/api_router.py | 24 +++++++++++++-- vllm/v1/engine/async_llm.py | 14 +++++++-- vllm/v1/engine/core.py | 9 ++++++ vllm/v1/engine/core_client.py | 30 +++++++++++++++++++ vllm/v1/engine/llm_engine.py | 7 +++++ 16 files changed, 142 insertions(+), 18 deletions(-) diff --git a/docs/serving/offline_inference.md b/docs/serving/offline_inference.md index 4512f4a0720..9a71612f262 100644 --- a/docs/serving/offline_inference.md +++ b/docs/serving/offline_inference.md @@ -65,6 +65,8 @@ For further details on Weight Transfer, please refer to [this page](../training/ - `LLM.start_weight_update` - Starts a new weight update cycle. - `LLM.update_weights` - Updates the model weights. - `LLM.finish_weight_update` - Finishes the current weight update cycle. +- `LLM.update_weight_version` - Sets the weight version without updating model weights. +- `LLM.get_weight_version` - Returns the latest committed weight version. ## Additional APIs diff --git a/docs/serving/online_serving/README.md b/docs/serving/online_serving/README.md index f7914bc0582..90ff7a3e3d8 100644 --- a/docs/serving/online_serving/README.md +++ b/docs/serving/online_serving/README.md @@ -179,6 +179,8 @@ For further details on Weight Transfer, please refer to [this page](../../traini - `/start_weight_update` - Prepares the inference engine for a weight update. - `/update_weights` - Update model weights (can alter model behavior) - `/finish_weight_update` - Finalizes the weight update +- `/update_weight_version` - Set the weight version without updating model weights +- `/weight_info` - Get the latest committed weight version - `/get_world_size` - Get distributed world size ### Collective RPC diff --git a/docs/training/async_rl.md b/docs/training/async_rl.md index e655f9c39ff..9e75a24eaa1 100644 --- a/docs/training/async_rl.md +++ b/docs/training/async_rl.md @@ -38,11 +38,12 @@ Resumes the scheduler after a pause. Any requests frozen with `mode="keep"` will ### HTTP Endpoints -When using the vLLM HTTP server, the same functionality is available via: +With `VLLM_SERVER_DEV_MODE=1`, the vLLM HTTP server exposes the same functionality via: - `POST /pause?mode=keep` - Pause generation - `POST /resume` - Resume generation - `POST /abort_requests` - Abort in-flight requests without pausing the scheduler (send `{}` to abort all, or `{"request_ids": [...]}`) +- `GET /weight_info` - Return the latest committed `weight_version` !!! note "Data Parallelism" When using data parallelism with vLLM's **internal load balancer** (i.e. `data_parallel_backend="ray"`), pause and resume are handled automatically across all DP ranks -- a single call is sufficient. When using an **external load balancer** (i.e. multiple independent vLLM instances behind a proxy), you must send pause and resume requests to **every** engine instance individually before and after the weight update. diff --git a/docs/training/weight_transfer/README.md b/docs/training/weight_transfer/README.md index 7579e5fd4d0..b8d39763181 100644 --- a/docs/training/weight_transfer/README.md +++ b/docs/training/weight_transfer/README.md @@ -53,7 +53,9 @@ When running vLLM as an HTTP server, the following endpoints are available for w | `/init_weight_transfer_engine` | POST | Initialize the weight transfer engine with backend-specific info | | `/start_weight_update` | POST | Start a weight update | | `/update_weights` | POST | Transfer a batch of weights with backend-specific metadata | -| `/finish_weight_update` | POST | Finish the weight update and run post-processing | +| `/finish_weight_update` | POST | Finish the update and optionally commit its `weight_version` | +| `/update_weight_version` | POST | Update `weight_version` without changing model weights | +| `/weight_info` | GET | Get the latest committed weight version | | `/pause` | POST | Pause generation before weight sync to handle inflight requests | | `/resume` | POST | Resume generation after weight sync | | `/get_world_size` | GET | Get the number of inference workers (useful for NCCL world size calculation) | @@ -79,7 +81,7 @@ EngineClass.trainer_send_weights( ) # 4. Finish weight update on inference side -llm.finish_weight_update() +llm.finish_weight_update(weight_version="step-42") ``` See the [NCCL](nccl.md) and [IPC](ipc.md) pages for backend-specific trainer APIs and full examples. diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index b79aa1974d1..eeeceb95998 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -1247,7 +1247,7 @@ class RecordingClient: self.order.append("update") self.last_update_info = update_info - def finish_weight_update(self) -> None: + def finish_weight_update(self, weight_version: str | None = None) -> None: self.order.append("finish") @@ -1303,6 +1303,10 @@ class TestTrainerClients: assert isinstance(update_req, WeightTransferUpdateRequest) assert update_req.update_info == {"names": ["w"]} + client.finish_weight_update("step-42") + handle.finish_weight_update.remote.assert_called_once_with() + handle.update_weight_version.remote.assert_called_once_with("step-42") + def test_http_client_pickles_ipc_handles_for_json(self, monkeypatch): """HTTP update_weights must encode raw ipc_handles as a base64 pickle.""" captured = {} @@ -1334,6 +1338,9 @@ class TestTrainerClients: client.update_weights(update_info) assert captured["json"]["update_info"] == update_info + client.finish_weight_update("step-42") + assert captured["json"] == {"weight_version": "step-42"} + class TestModuleSource: """`ModuleSource` metadata vs. materialized iteration (dense, no GPU).""" diff --git a/tests/entrypoints/openai/test_openai_schema.py b/tests/entrypoints/openai/test_openai_schema.py index 2985c539518..6d3fc2f4474 100644 --- a/tests/entrypoints/openai/test_openai_schema.py +++ b/tests/entrypoints/openai/test_openai_schema.py @@ -148,6 +148,7 @@ def test_openapi_stateless(case: schemathesis.Case): "/start_draft_weight_update", "/update_weights", "/finish_weight_update", + "/update_weight_version", ): return diff --git a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py b/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py index 9088b3c5e8d..31d562e5ccf 100644 --- a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py +++ b/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py @@ -234,6 +234,7 @@ def test_update_weights_calls_engine(): assert shapes == test_shapes llm.finish_weight_update() + assert llm.get_weight_version() == "default" @create_new_process_for_each_test() @@ -259,6 +260,8 @@ def test_full_weight_transfer_flow(): weight_transfer_config=WeightTransferConfig(backend="nccl"), ) + assert llm.get_weight_version() == "default" + # Step 1: Initialize weight transfer engine llm.init_weight_transfer_engine( WeightTransferInitRequest(init_info={"test_param": "flow_test"}) @@ -278,8 +281,15 @@ def test_full_weight_transfer_flow(): ) ) + assert llm.get_weight_version() == "default" + # Step 4: Finish weight update - llm.finish_weight_update() + llm.finish_weight_update("step-42") + + assert llm.get_weight_version() == "step-42" + + llm.update_weight_version("manual-version") + assert llm.get_weight_version() == "manual-version" # Verify the full flow completed def check_flow(self): diff --git a/vllm/distributed/weight_transfer/base.py b/vllm/distributed/weight_transfer/base.py index 2e377e29253..adddf41ff4e 100644 --- a/vllm/distributed/weight_transfer/base.py +++ b/vllm/distributed/weight_transfer/base.py @@ -370,7 +370,7 @@ class VLLMWeightSyncClient(Protocol): def update_weights(self, update_info: dict[str, Any]) -> None: ... - def finish_weight_update(self) -> None: ... + def finish_weight_update(self, weight_version: str | None = None) -> None: ... class TrainerWeightTransferEngine(ABC, Generic[TConfig, TInitInfo]): diff --git a/vllm/distributed/weight_transfer/clients.py b/vllm/distributed/weight_transfer/clients.py index 4f54a6e291e..12dd0c9eacc 100644 --- a/vllm/distributed/weight_transfer/clients.py +++ b/vllm/distributed/weight_transfer/clients.py @@ -77,8 +77,11 @@ class HTTPVLLMWeightSyncClient: "update_weights", {"update_info": _json_safe_update_info(update_info)} ) - def finish_weight_update(self) -> None: - self._post("finish_weight_update") + def finish_weight_update(self, weight_version: str | None = None) -> None: + json = ( + {"weight_version": weight_version} if weight_version is not None else None + ) + self._post("finish_weight_update", json) class RayVLLMWeightSyncClient: @@ -108,7 +111,11 @@ class RayVLLMWeightSyncClient: request = WeightTransferUpdateRequest(update_info=update_info) ray.get([h.update_weights.remote(request) for h in self.handles]) - def finish_weight_update(self) -> None: + def finish_weight_update(self, weight_version: str | None = None) -> None: import ray ray.get([h.finish_weight_update.remote() for h in self.handles]) + if weight_version is not None: + ray.get( + [h.update_weight_version.remote(weight_version) for h in self.handles] + ) diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index ef3be178ac8..5a9b9f96d2c 100644 --- a/vllm/engine/protocol.py +++ b/vllm/engine/protocol.py @@ -267,6 +267,14 @@ class EngineClient(ABC): """Batched weight update for RL training.""" raise NotImplementedError - async def finish_weight_update(self) -> None: - """Finish the current weight update.""" + async def finish_weight_update(self, weight_version: str | None = None) -> None: + """Finish the weight update and set its version if provided.""" + raise NotImplementedError + + async def update_weight_version(self, new_version: str) -> None: + """Set the weight version without updating weights.""" + raise NotImplementedError + + async def get_weight_version(self) -> str: + """Return the latest committed weight version.""" raise NotImplementedError diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index b3205728e49..4274819d988 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -885,9 +885,19 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin): "update_weights", kwargs={"update_info": update_info_dict} ) - def finish_weight_update(self) -> None: - """Finish the current weight update.""" + def finish_weight_update(self, weight_version: str | None = None) -> None: + """Finish the weight update and set its version if provided.""" self.llm_engine.collective_rpc("finish_weight_update") + if weight_version is not None: + self.llm_engine.set_weight_version(weight_version) + + def update_weight_version(self, new_version: str) -> None: + """Set the weight version without updating weights.""" + self.llm_engine.set_weight_version(new_version) + + def get_weight_version(self) -> str: + """Return the latest committed weight version.""" + return self.llm_engine.get_weight_version() def __repr__(self) -> str: """Return a transformers-style hierarchical view of the model.""" diff --git a/vllm/entrypoints/serve/dev/rlhf/api_router.py b/vllm/entrypoints/serve/dev/rlhf/api_router.py index 8a2494a59df..392fcf56747 100644 --- a/vllm/entrypoints/serve/dev/rlhf/api_router.py +++ b/vllm/entrypoints/serve/dev/rlhf/api_router.py @@ -5,7 +5,7 @@ import json from http import HTTPStatus from typing import Annotated -from fastapi import APIRouter, FastAPI, HTTPException, Query, Request +from fastapi import APIRouter, Body, FastAPI, HTTPException, Query, Request from fastapi.responses import JSONResponse from vllm.distributed.weight_transfer.base import ( @@ -203,11 +203,29 @@ async def update_weights(raw_request: Request): @router.post("/finish_weight_update") -async def finish_weight_update(raw_request: Request): - await engine_client(raw_request).finish_weight_update() +async def finish_weight_update( + raw_request: Request, + weight_version: Annotated[str | None, Body(embed=True)] = None, +): + await engine_client(raw_request).finish_weight_update(weight_version) return JSONResponse(content={"message": "Weight update finished"}) +@router.post("/update_weight_version") +async def update_weight_version( + raw_request: Request, + new_version: Annotated[str, Body(embed=True)], +): + await engine_client(raw_request).update_weight_version(new_version) + return JSONResponse(content={"success": True, "new_version": new_version}) + + +@router.get("/weight_info") +async def weight_info(raw_request: Request): + weight_version = await engine_client(raw_request).get_weight_version() + return JSONResponse(content={"weight_version": weight_version}) + + @router.get("/get_world_size") async def get_world_size( raw_request: Request, diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 922a8aa5982..5c2e01cf44b 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -1106,6 +1106,16 @@ class AsyncLLM(EngineClient): "update_weights", kwargs={"update_info": request.update_info} ) - async def finish_weight_update(self) -> None: - """Finish the current weight update.""" + async def finish_weight_update(self, weight_version: str | None = None) -> None: + """Finish the weight update and set its version if provided.""" await self.collective_rpc("finish_weight_update") + if weight_version is not None: + await self.update_weight_version(weight_version) + + async def update_weight_version(self, new_version: str) -> None: + """Set the weight version without updating weights.""" + await self.engine_core.set_weight_version_async(new_version) + + async def get_weight_version(self) -> str: + """Return the latest committed weight version.""" + return await self.engine_core.get_weight_version_async() diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 9817c474343..9917f810b5b 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -125,6 +125,8 @@ class EngineCore: ) self.log_stats = log_stats + # Opaque weight version supplied by the caller. + self._weight_version = "default" # Setup Model. self.model_executor = executor_class(vllm_config) @@ -956,6 +958,13 @@ class EngineCore: ) -> list[_R]: return self.model_executor.collective_rpc(method, timeout, args, kwargs) + def set_weight_version(self, weight_version: str) -> None: + self._weight_version = weight_version + + def get_weight_version(self) -> str: + """Return the latest committed weight version.""" + return self._weight_version + def preprocess_add_request(self, request: EngineCoreRequest) -> tuple[Request, int]: """Preprocess the request. diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index 0aa4b6f3312..a6c232b2ab7 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -176,9 +176,21 @@ class EngineCoreClient(ABC): def execute_dummy_batch(self) -> None: raise NotImplementedError + def set_weight_version(self, weight_version: str) -> None: + raise NotImplementedError + + def get_weight_version(self) -> str: + raise NotImplementedError + async def execute_dummy_batch_async(self) -> None: raise NotImplementedError + async def set_weight_version_async(self, weight_version: str) -> None: + raise NotImplementedError + + async def get_weight_version_async(self) -> str: + raise NotImplementedError + def abort_requests(self, request_ids: list[str]) -> None: raise NotImplementedError @@ -351,6 +363,12 @@ class InprocClient(EngineCoreClient): def execute_dummy_batch(self) -> None: self.engine_core.execute_dummy_batch() + def set_weight_version(self, weight_version: str) -> None: + self.engine_core.set_weight_version(weight_version) + + def get_weight_version(self) -> str: + return self.engine_core.get_weight_version() + def add_lora(self, lora_request: LoRARequest) -> bool: return self.engine_core.add_lora(lora_request) @@ -947,6 +965,12 @@ class SyncMPClient(MPClient): def execute_dummy_batch(self) -> None: self.call_utility("execute_dummy_batch") + def set_weight_version(self, weight_version: str) -> None: + self.call_utility("set_weight_version", weight_version) + + def get_weight_version(self) -> str: + return self.call_utility("get_weight_version") + def collective_rpc( self, method: str | Callable[..., _R], @@ -1199,6 +1223,12 @@ class AsyncMPClient(MPClient): async def execute_dummy_batch_async(self) -> None: await self.call_utility_async("execute_dummy_batch") + async def set_weight_version_async(self, weight_version: str) -> None: + await self.call_utility_async("set_weight_version", weight_version) + + async def get_weight_version_async(self) -> str: + return await self.call_utility_async("get_weight_version") + async def add_lora_async(self, lora_request: LoRARequest) -> bool: return await self.call_utility_async("add_lora", lora_request) diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py index ff86a1dffd9..17e40630859 100644 --- a/vllm/v1/engine/llm_engine.py +++ b/vllm/v1/engine/llm_engine.py @@ -425,6 +425,13 @@ class LLMEngine: ) -> list[_R]: return self.engine_core.collective_rpc(method, timeout, args, kwargs) + def set_weight_version(self, weight_version: str) -> None: + self.engine_core.set_weight_version(weight_version) + + def get_weight_version(self) -> str: + """Return the latest committed weight version.""" + return self.engine_core.get_weight_version() + def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]: return self.collective_rpc("apply_model", args=(func,)) From 03a2d033673d8804b001bae41f30059273e1a669 Mon Sep 17 00:00:00 2001 From: Chauncey <chaunceyjiang@gmail.com> Date: Tue, 28 Jul 2026 14:38:20 +0800 Subject: [PATCH 155/185] [Bugfix] Respect cgroup memory limits on all platforms (#49966) Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> --- vllm/model_executor/model_loader/weight_utils.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index db161a58988..e1897a77f10 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -694,12 +694,10 @@ def _get_checkpoints_size_bytes(files: list[str]) -> int: def _get_available_ram_bytes() -> int: - """Return available RAM, honoring cgroup limits on ROCm.""" + """Return available RAM, honoring cgroup limits.""" import psutil host_available = psutil.virtual_memory().available - if not current_platform.is_rocm(): - return host_available from vllm.utils.cpu_resource_utils import get_cgroup_memory_limit From b09688a6e77f0c061c08c67aeed56ad669c617dc Mon Sep 17 00:00:00 2001 From: aoshen02 <aoshen@inferact.ai> Date: Tue, 28 Jul 2026 15:16:21 +0800 Subject: [PATCH 156/185] [Bugfix][Spec Decode] Preserve draft buffers across level-2 sleep (#49774) Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com> Signed-off-by: vx120 <893600387@qq.com> Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: vx120 <893600387@qq.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Roger Wang <hey@rogerw.io> --- vllm/v1/worker/gpu_worker.py | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 556b1e6c7d9..bbab98dbd7b 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -154,7 +154,7 @@ class Worker(WorkerBase): self.worker_sentinel = WorkerSentinel(worker=self) # Buffers saved before sleep self._sleep_saved_buffers: dict[str, torch.Tensor] = {} - self._sleep_rebuild_draft_metadata_buffers = False + self._sleep_saved_draft_buffers: dict[str, torch.Tensor] = {} # Weight transfer engine is created in `load_model` once the model # is available, since the engine needs a reference to the model. @@ -200,10 +200,10 @@ class Worker(WorkerBase): name: buffer.cpu().clone() for name, buffer in model.named_buffers() } draft = self.get_draft_model() - inner = getattr(draft, "model", None) if draft is not None else None - self._sleep_rebuild_draft_metadata_buffers = inner is not None and hasattr( - inner, "_build_fused_kv_buffers" - ) + if draft is not None: + self._sleep_saved_draft_buffers = { + name: buffer.cpu().clone() for name, buffer in draft.named_buffers() + } self._get_sleep_mode_backend().suspend(level) @@ -228,20 +228,21 @@ class Worker(WorkerBase): self._get_sleep_mode_backend().resume(tags) # Restore the buffers after level 2 sleep - if len(self._sleep_saved_buffers): + wake_weights = tags is None or "weights" in tags + if wake_weights and len(self._sleep_saved_buffers): model = self.model_runner.model for name, buffer in model.named_buffers(): if name in self._sleep_saved_buffers: buffer.data.copy_(self._sleep_saved_buffers[name].data) self._sleep_saved_buffers = {} - if self._sleep_rebuild_draft_metadata_buffers: + if wake_weights and len(self._sleep_saved_draft_buffers): draft = self.get_draft_model() if draft is not None: - inner = getattr(draft, "model", None) - if inner is not None and hasattr(inner, "_build_fused_kv_buffers"): - inner._build_fused_kv_buffers() - self._sleep_rebuild_draft_metadata_buffers = False + for name, buffer in draft.named_buffers(): + if name in self._sleep_saved_draft_buffers: + buffer.data.copy_(self._sleep_saved_draft_buffers[name].data) + self._sleep_saved_draft_buffers = {} if tags is None or "kv_cache" in tags: self.model_runner.post_kv_cache_wake_up() From 99b57a4823d8fe22d3e249b40efa5d0f6fdfa2b1 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas <akaratza@amd.com> Date: Tue, 28 Jul 2026 02:18:04 -0500 Subject: [PATCH 157/185] [CI][ROCm] Soft fail LoRA mirror (#50086) Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: OpenAI Codex <codex@openai.com> --- .buildkite/test_areas/lora.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/.buildkite/test_areas/lora.yaml b/.buildkite/test_areas/lora.yaml index a79196ffbd8..214d4a12bf7 100644 --- a/.buildkite/test_areas/lora.yaml +++ b/.buildkite/test_areas/lora.yaml @@ -16,6 +16,7 @@ steps: amd: dind: false device: mi300_1 + soft_fail: true working_dir: "/vllm-workspace/tests" timeout_in_minutes: 85 source_file_dependencies: From 61ac368021033dea54448d36ed98c7426e82cd18 Mon Sep 17 00:00:00 2001 From: Thien Tran <gau.nernst@yahoo.com.sg> Date: Tue, 28 Jul 2026 15:22:09 +0800 Subject: [PATCH 158/185] [Kimi-K3] Add AttnRes kernels (#50090) Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> --- .buildkite/test_areas/models_basic.yaml | 12 + CMakeLists.txt | 22 + .../kimi_k3/attn_res_kernel.cu | 954 ++++++++++++++++++ csrc/libtorch_stable/ops.h | 11 + csrc/libtorch_stable/torch_bindings.cpp | 11 + tests/models/kimi_k3/test_amd_attn_res.py | 102 ++ tests/models/kimi_k3/test_attn_res.py | 193 ++++ vllm/_custom_ops.py | 27 + vllm/models/kimi_k3/amd/__init__.py | 2 + vllm/models/kimi_k3/amd/ops/__init__.py | 0 vllm/models/kimi_k3/amd/ops/attn_res.py | 132 +++ vllm/models/kimi_k3/nvidia/__init__.py | 2 + vllm/models/kimi_k3/nvidia/ops/__init__.py | 6 + vllm/models/kimi_k3/nvidia/ops/attn_res.py | 245 +++++ 14 files changed, 1719 insertions(+) create mode 100644 csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu create mode 100644 tests/models/kimi_k3/test_amd_attn_res.py create mode 100644 tests/models/kimi_k3/test_attn_res.py create mode 100644 vllm/models/kimi_k3/amd/__init__.py create mode 100644 vllm/models/kimi_k3/amd/ops/__init__.py create mode 100644 vllm/models/kimi_k3/amd/ops/attn_res.py create mode 100644 vllm/models/kimi_k3/nvidia/__init__.py create mode 100644 vllm/models/kimi_k3/nvidia/ops/__init__.py create mode 100644 vllm/models/kimi_k3/nvidia/ops/attn_res.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index af90308d6c3..a7c7aa9022d 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -61,6 +61,18 @@ steps: # FA4 kernel tests require SM100; the suite skips them elsewhere. - pytest -v -s models/inkling +- label: Kimi K3 Unit Tests (B200) + key: kimi-k3-unit-tests-b200 + timeout_in_minutes: 40 + device: b200-k8s + source_file_dependencies: + - vllm/models/kimi_k3/ + - csrc/libtorch_stable/kimi_k3/ + - tests/models/kimi_k3/ + commands: + # The native NVIDIA AttnRes kernel requires the SM100 family. + - pytest -v -s models/kimi_k3 + - label: Basic Models Test (Other CPU) # 5min key: basic-models-test-other-cpu depends_on: diff --git a/CMakeLists.txt b/CMakeLists.txt index 3a4e23ad0b3..cdda81ea46e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1079,6 +1079,24 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") set(MLA_ARCHS) endif() + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(KIMI_K3_ATTN_RES_ARCHS + "10.0f" "${CUDA_ARCHS}") + endif() + if(KIMI_K3_ATTN_RES_ARCHS) + set(KIMI_K3_ATTN_RES_SRC + "csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu") + set_gencode_flags_for_srcs( + SRCS "${KIMI_K3_ATTN_RES_SRC}" + CUDA_ARCHS "${KIMI_K3_ATTN_RES_ARCHS}") + set_property(SOURCE ${KIMI_K3_ATTN_RES_SRC} APPEND PROPERTY + COMPILE_OPTIONS + "$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr;--expt-extended-lambda;--use_fast_math>") + list(APPEND VLLM_STABLE_EXT_SRC "${KIMI_K3_ATTN_RES_SRC}") + message(STATUS + "Building Kimi K3 AttnRes for archs: ${KIMI_K3_ATTN_RES_ARCHS}") + endif() + # Hadacore kernels cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}") if(HADACORE_ARCHS) @@ -1120,6 +1138,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") target_compile_definitions(_C_stable_libtorch PRIVATE VLLM_ENABLE_COOPERATIVE_TOPK=1) endif() + if(KIMI_K3_ATTN_RES_ARCHS) + target_compile_definitions(_C_stable_libtorch PRIVATE + VLLM_ENABLE_KIMI_K3_ATTN_RES=1) + endif() # Needed by CUTLASS kernels target_compile_definitions(_C_stable_libtorch PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1) diff --git a/csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu b/csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu new file mode 100644 index 00000000000..eb4dcf6bf5a --- /dev/null +++ b/csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu @@ -0,0 +1,954 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + */ + +// Production AttnRes forward for Blackwell (SM100). +// +// Warp-specialized online softmax + residual + RMSNorm: +// - 1 producer warp issues cp.async.bulk row loads into shared memory. +// - 8 consumer warps compute reductions and output. +// - Q=res_weight*rms_weight remains in registers across persistent tokens. +// - V rows are converted once and cached as FP32 in TMEM between passes. +// +// Integration contract: Kimi K3 H=7168, 1<=num_blocks<=8, and token-major +// block residual storage. + +#include "../torch_utils.h" + +#include <cfloat> +#include <cstdint> +#include <cstdio> +#include <cuda_runtime.h> +#include <type_traits> + +using bf16_t = __nv_bfloat16; + +namespace sm100 { +namespace fwd_prod_v2 { + +constexpr int K_TILE = 1024; +constexpr int N_CHUNK_DEFAULT = 4; +constexpr int CHUNK_DEPTH = 2; +constexpr int BLK = 288; // 1 producer warp + 8 consumer warps +constexpr int CONSUMER_THREADS = BLK - 32; // 256 +constexpr int CONSUMER_WARPS = CONSUMER_THREADS / 32; +constexpr int CONSUMER_GROUPS = 2; // two 128-thread consumer groups +constexpr int CONSUMER_THREADS_PER_GROUP = CONSUMER_THREADS / CONSUMER_GROUPS; +constexpr int FIRST_USER_NAMED_BARRIER = 8; + +__device__ __forceinline__ const bf16_t* residual_addr( + const bf16_t* block_res, const bf16_t* layer_res, int source, int N, + int token, int block_stride_m, int block_stride_r, int H) { + if (source < N - 1) { + return block_res + static_cast<long long>(token) * block_stride_m + + source * block_stride_r; + } + return layer_res + static_cast<long long>(token) * H; +} + +__device__ __forceinline__ uint32_t elect_one_sync() { + uint32_t pred = 0; + uint32_t laneid = 0; + asm volatile( + "{\n" + ".reg .b32 %%rx;\n" + ".reg .pred %%px;\n" + " elect.sync %%rx|%%px, %2;\n" + "@%%px mov.s32 %1, 1;\n" + " mov.s32 %0, %%rx;\n" + "}\n" + : "+r"(laneid), "+r"(pred) + : "r"(0xffffffff)); + return pred; +} + +__device__ __forceinline__ void mbarrier_init(uint64_t& barrier, + int thread_count) { + uint32_t const barrier_addr = + static_cast<uint32_t>(__cvta_generic_to_shared(&barrier)); + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(barrier_addr), + "r"(thread_count)); +} + +__device__ __forceinline__ void mbarrier_expect_tx(uint64_t& barrier, + uint32_t bytes) { + uint32_t const barrier_addr = + static_cast<uint32_t>(__cvta_generic_to_shared(&barrier)); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"( + barrier_addr), + "r"(bytes)); +} + +__device__ __forceinline__ void mbarrier_wait(uint64_t& barrier, int phase) { + uint32_t const barrier_addr = + static_cast<uint32_t>(__cvta_generic_to_shared(&barrier)); + asm volatile( + "{\n" + ".reg .pred p;\n" + "WAIT:\n" + "mbarrier.try_wait.parity.shared::cta.b64 p, [%0], %1;\n" + "@p bra DONE;\n" + "bra WAIT;\n" + "DONE:\n" + "}\n" ::"r"(barrier_addr), + "r"(phase)); +} + +__device__ __forceinline__ void mbarrier_arrive(uint64_t& barrier) { + uint32_t const barrier_addr = + static_cast<uint32_t>(__cvta_generic_to_shared(&barrier)); + asm volatile( + "{\n" + ".reg .b64 state;\n" + "mbarrier.arrive.shared::cta.b64 state, [%0];\n" + "}\n" ::"r"(barrier_addr)); +} + +__device__ __forceinline__ void fence_mbarrier_init() { + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); +} + +__device__ __forceinline__ void named_barrier_sync(uint32_t num_threads, + uint32_t user_barrier_id) { + asm volatile( + "bar.sync %0, %1;" ::"r"(user_barrier_id + FIRST_USER_NAMED_BARRIER), + "r"(num_threads) + : "memory"); +} + +__device__ __forceinline__ void tmem_allocate(int num_columns, uint32_t* dst) { + uint32_t const dst_addr = + static_cast<uint32_t>(__cvta_generic_to_shared(dst)); + asm volatile( + "tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"( + dst_addr), + "r"(num_columns)); +} + +__device__ __forceinline__ void tmem_free(uint32_t tmem_ptr, int num_columns) { + asm volatile( + "tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(tmem_ptr), + "r"(num_columns)); +} + +__device__ __forceinline__ void tmem_release_allocation_lock() { + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); +} + +__device__ __forceinline__ void tmem_store_wait() { + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); +} + +template <int N, typename T> +__device__ __forceinline__ void tmem_load(uint32_t src_addr, T* dst) { + uint32_t* values = reinterpret_cast<uint32_t*>(dst); + if constexpr (N == 8) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x8.b32" + "{%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n" + : "=r"(values[0]), "=r"(values[1]), "=r"(values[2]), "=r"(values[3]), + "=r"(values[4]), "=r"(values[5]), "=r"(values[6]), "=r"(values[7]) + : "r"(src_addr)); + } else { + static_assert(N == 4, "AttnRes TMEM helpers support x4 and x8"); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + "{%0, %1, %2, %3}, [%4];\n" + : "=r"(values[0]), "=r"(values[1]), "=r"(values[2]), "=r"(values[3]) + : "r"(src_addr)); + } +} + +template <int N, typename T> +__device__ __forceinline__ void tmem_store(uint32_t dst_addr, T* src) { + uint32_t* values = reinterpret_cast<uint32_t*>(src); + if constexpr (N == 8) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x8.b32" + "[%8], {%0, %1, %2, %3, %4, %5, %6, %7};\n" ::"r"(values[0]), + "r"(values[1]), "r"(values[2]), "r"(values[3]), "r"(values[4]), + "r"(values[5]), "r"(values[6]), "r"(values[7]), "r"(dst_addr)); + } else { + static_assert(N == 4, "AttnRes TMEM helpers support x4 and x8"); + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + "[%4], {%0, %1, %2, %3};\n" ::"r"(values[0]), + "r"(values[1]), "r"(values[2]), "r"(values[3]), "r"(dst_addr)); + } +} + +__device__ __forceinline__ float2 float2_add(const float2& a, const float2& b) { + float2 result; + asm volatile("add.rn.f32x2 %0, %1, %2;\n" + : "=l"(reinterpret_cast<uint64_t&>(result)) + : "l"(reinterpret_cast<uint64_t const&>(a)), + "l"(reinterpret_cast<uint64_t const&>(b))); + return result; +} + +__device__ __forceinline__ float2 float2_mul(const float2& a, const float2& b) { + float2 result; + asm volatile("mul.f32x2 %0, %1, %2;\n" + : "=l"(reinterpret_cast<uint64_t&>(result)) + : "l"(reinterpret_cast<uint64_t const&>(a)), + "l"(reinterpret_cast<uint64_t const&>(b))); + return result; +} + +__device__ __forceinline__ float2 float2_fma(const float2& a, const float2& b, + const float2& c) { + float2 result; + asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n" + : "=l"(reinterpret_cast<uint64_t&>(result)) + : "l"(reinterpret_cast<uint64_t const&>(a)), + "l"(reinterpret_cast<uint64_t const&>(b)), + "l"(reinterpret_cast<uint64_t const&>(c))); + return result; +} + +template <int NC> +struct FwdSmemPlan { + alignas(16) uint64_t bar_ready[CHUNK_DEPTH]; + alignas(16) uint64_t bar_consumed[CHUNK_DEPTH]; + alignas(16) uint64_t bar_output_norm_ready; + alignas(16) float2 ws_stats[CONSUMER_WARPS][NC]; + uint32_t tmem_base; +}; + +__device__ __forceinline__ void cp_async_bulk(void* smem_dst, + const void* gmem_src, int bytes, + uint64_t& mbar) { + uint32_t const s = static_cast<uint32_t>(__cvta_generic_to_shared(smem_dst)); + uint32_t const m = static_cast<uint32_t>(__cvta_generic_to_shared(&mbar)); + asm volatile( + "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], " + "[%1], %2, [%3];\n" ::"r"(s), + "l"(gmem_src), "r"(bytes), "r"(m) + : "memory"); +} + +template <int H, int NC = N_CHUNK_DEFAULT, bool RELEASE_TMEM = false, + bool HAS_DELTA = false, bool HAS_OUTPUT_NORM = false, + bool OUTPUT_NORM_IN_SMEM = false> +__global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel( + const bf16_t* __restrict__ block_res, bf16_t* __restrict__ layer_res, + const bf16_t* __restrict__ delta, const bf16_t* __restrict__ res_w, + const bf16_t* __restrict__ rms_w, bf16_t* __restrict__ output, int N, int T, + int B, int block_stride_m, int block_stride_r, float rms_eps, + const bf16_t* __restrict__ output_norm_weight, float output_norm_eps) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1100 + constexpr float LOG2_E = 1.4426950408889634f; + constexpr int N_CHUNK = NC; + // The two-source specialization only consumes half of the TMEM columns. + constexpr int TMEM_COLS_ALLOC = NC == 2 ? 128 : 256; + constexpr int NUM_BUFS = CHUNK_DEPTH * NC; + constexpr int NHT = H / K_TILE; + constexpr int SLICES_PER_GROUP = + (NHT + CONSUMER_GROUPS - 1) / CONSUMER_GROUPS; + constexpr int VEC = 8; + constexpr int ACC_PER_THREAD = H == 7168 ? 28 : SLICES_PER_GROUP * VEC; + constexpr int TMEM_V_COLS_PER_GROUP = SLICES_PER_GROUP * N_CHUNK * VEC; + constexpr int TMEM_V_COLS_TOTAL = CONSUMER_GROUPS * TMEM_V_COLS_PER_GROUP; + static_assert(TMEM_V_COLS_TOTAL <= TMEM_COLS_ALLOC); + static_assert(H >= 4096 && H <= 8192); + static_assert(H % K_TILE == 0); + + const int tid = threadIdx.x; + const int wid = tid >> 5; + const int lane = tid & 31; + const int TB = T * B; + const int num_ctas = gridDim.x; + const int num_chunks = (N + N_CHUNK - 1) / N_CHUNK; + + const int comp_wid = wid - 1; + const int comp_tid = tid - 32; + const int group = (comp_wid >= 4) ? 1 : 0; + const int ct_in_group = + (comp_tid >= 0) ? (comp_tid & (CONSUMER_THREADS_PER_GROUP - 1)) : -1; + const int k_local = ct_in_group * VEC; + + constexpr size_t V_BYTES = (size_t)NUM_BUFS * H * sizeof(bf16_t); + constexpr size_t DELTA_BYTES = + HAS_DELTA ? (size_t)CHUNK_DEPTH * H * sizeof(bf16_t) : 0; + constexpr size_t OUTPUT_NORM_BYTES = + OUTPUT_NORM_IN_SMEM ? (size_t)H * sizeof(bf16_t) : 0; + extern __shared__ __align__(16) char smem_raw[]; + bf16_t* v_bufs = reinterpret_cast<bf16_t*>(smem_raw); // [NUM_BUFS][H] + bf16_t* delta_bufs = reinterpret_cast<bf16_t*>(smem_raw + V_BYTES); + bf16_t* output_norm_buf = + reinterpret_cast<bf16_t*>(smem_raw + V_BYTES + DELTA_BYTES); + FwdSmemPlan<NC>& plan = *reinterpret_cast<FwdSmemPlan<NC>*>( + smem_raw + V_BYTES + DELTA_BYTES + OUTPUT_NORM_BYTES); + + auto slot_of = [](long long gci, int n) { + return (int)(gci % CHUNK_DEPTH) * N_CHUNK + n; + }; + auto phase_of = [](long long gci) { return (int)((gci / CHUNK_DEPTH) & 1); }; + auto buf_ptr = [&](int slot) -> bf16_t* { return v_bufs + slot * H; }; + auto delta_buf_ptr = [&](int chunk_slot) -> bf16_t* { + return delta_bufs + chunk_slot * H; + }; + + if (wid == 0 && elect_one_sync()) { + #pragma unroll + for (int i = 0; i < CHUNK_DEPTH; i++) { + mbarrier_init(plan.bar_ready[i], 1); + mbarrier_init(plan.bar_consumed[i], CONSUMER_WARPS); + } + if constexpr (OUTPUT_NORM_IN_SMEM) { + mbarrier_init(plan.bar_output_norm_ready, 1); + } + fence_mbarrier_init(); + } + + // gdc wait BEFORE tmem alloc + cudaGridDependencySynchronize(); + + if (wid == 1) { + tmem_allocate(TMEM_COLS_ALLOC, &plan.tmem_base); + if constexpr (RELEASE_TMEM) { + tmem_release_allocation_lock(); + } + } + __syncthreads(); + + if constexpr (OUTPUT_NORM_IN_SMEM) { + if (wid == 0 && elect_one_sync()) { + mbarrier_expect_tx(plan.bar_output_norm_ready, H * (int)sizeof(bf16_t)); + cp_async_bulk(output_norm_buf, output_norm_weight, H * sizeof(bf16_t), + plan.bar_output_norm_ready); + } + } + + const uint32_t my_v_tmem = + comp_tid >= 0 ? plan.tmem_base + group * TMEM_V_COLS_PER_GROUP : 0; + float q_cache[ACC_PER_THREAD]; + if (comp_tid >= 0) { + #pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) { + if constexpr (H == 7168) { + if (si == SLICES_PER_GROUP - 1) { + int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4; + #pragma unroll + for (int j = 0; j < 4; j++) { + int h = h_base + j; + q_cache[si * VEC + j] = + __bfloat162float(rms_w[h]) * __bfloat162float(res_w[h]); + } + continue; + } + } + int dt = si * CONSUMER_GROUPS + group; + if (dt >= NHT) continue; + int h_base = dt * K_TILE + k_local; + #pragma unroll + for (int j = 0; j < VEC; j++) { + int h = h_base + j; + q_cache[si * VEC + j] = + __bfloat162float(rms_w[h]) * __bfloat162float(res_w[h]); + } + } + } + + if (wid == 0) { + if (elect_one_sync()) { + long long gci = 0; + for (int tb = blockIdx.x; tb < TB; tb += num_ctas) { + const int t = tb / B; + for (int ci = 0; ci < num_chunks; ci++, gci++) { + int ns = ci * N_CHUNK; + int an = min(N_CHUNK, N - ns); + int chunk_slot = (int)(gci % CHUNK_DEPTH); + int pc = phase_of(gci); + mbarrier_wait(plan.bar_consumed[chunk_slot], pc ^ 1); + int transaction_bytes = an * H * (int)sizeof(bf16_t); + if constexpr (HAS_DELTA) { + int prefix_n = N - 1 - ns; + if (prefix_n >= 0 && prefix_n < an) { + transaction_bytes += H * (int)sizeof(bf16_t); + } + } + mbarrier_expect_tx(plan.bar_ready[chunk_slot], transaction_bytes); + #pragma unroll + for (int n = 0; n < N_CHUNK; n++) { + if (n >= an) continue; + int slot = slot_of(gci, n); + const bf16_t* src = + residual_addr(block_res, layer_res, ns + n, N, t, + block_stride_m, block_stride_r, H); + cp_async_bulk(buf_ptr(slot), src, H * sizeof(bf16_t), + plan.bar_ready[chunk_slot]); + } + if constexpr (HAS_DELTA) { + int prefix_n = N - 1 - ns; + if (prefix_n >= 0 && prefix_n < an) { + cp_async_bulk(delta_buf_ptr(chunk_slot), + delta + (long long)tb * H, H * sizeof(bf16_t), + plan.bar_ready[chunk_slot]); + } + } + } + } + } + } else { + float acc32[ACC_PER_THREAD] = {}; + float eps_cache; + asm volatile("mov.b32 %0, %1;" : "=f"(eps_cache) : "f"(rms_eps)); + + long long gci = 0; + for (int tb = blockIdx.x; tb < TB; tb += num_ctas) { + float m_running = -FLT_MAX; + float s_running = 0.f; + #pragma unroll + for (int i = 0; i < ACC_PER_THREAD; i++) { + acc32[i] = 0.f; + } + + for (int ci = 0; ci < num_chunks; ci++, gci++) { + int ns = ci * N_CHUNK; + int an = min(N_CHUNK, N - ns); + int chunk_slot = (int)(gci % CHUNK_DEPTH); + int pr = phase_of(gci); + mbarrier_wait(plan.bar_ready[chunk_slot], pr); + + float2 sq_local[N_CHUNK] = {}; + float2 dot_local[N_CHUNK] = {}; + + auto pass_A_body = [&](auto AN_TOK) { + constexpr int AN = decltype(AN_TOK)::value; + #pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) { + if constexpr (H == 7168) { + if (si == SLICES_PER_GROUP - 1) { + int h_base = + 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4; + const float* qv = &q_cache[si * VEC]; + #pragma unroll + for (int n = 0; n < AN; n++) { + int slot = slot_of(gci, n); + int2 vp = + *reinterpret_cast<const int2*>(buf_ptr(slot) + h_base); + auto* v2 = reinterpret_cast<__nv_bfloat162*>(&vp); + if constexpr (HAS_DELTA) { + int prefix_n = N - 1 - ns; + if (n == prefix_n) { + const bf16_t* delta_ptr = + delta_buf_ptr(chunk_slot) + h_base; + #pragma unroll + for (int j = 0; j < 2; j++) { + auto delta2 = *reinterpret_cast<const __nv_bfloat162*>( + delta_ptr + 2 * j); + v2[j] = __hadd2(v2[j], delta2); + } + *reinterpret_cast<int2*>(layer_res + (long long)tb * H + + h_base) = vp; + } + } + float2 f[2] = {__bfloat1622float2(v2[0]), + __bfloat1622float2(v2[1])}; + tmem_store<4>(my_v_tmem + (si * N_CHUNK + n) * VEC, f); + sq_local[n] = float2_fma(f[0], f[0], sq_local[n]); + sq_local[n] = float2_fma(f[1], f[1], sq_local[n]); + dot_local[n] = + float2_fma(f[0], make_float2(qv[0], qv[1]), dot_local[n]); + dot_local[n] = + float2_fma(f[1], make_float2(qv[2], qv[3]), dot_local[n]); + } + continue; + } + } + int dt = si * CONSUMER_GROUPS + group; + if (dt >= NHT) continue; + int h_base = dt * K_TILE + k_local; + const float* qv = &q_cache[si * VEC]; + + #pragma unroll + for (int n = 0; n < AN; n++) { + int slot = slot_of(gci, n); + int4 vp = *reinterpret_cast<const int4*>(buf_ptr(slot) + h_base); + auto* v2 = reinterpret_cast<__nv_bfloat162*>(&vp); + if constexpr (HAS_DELTA) { + int prefix_n = N - 1 - ns; + if (n == prefix_n) { + const bf16_t* delta_ptr = delta_buf_ptr(chunk_slot) + h_base; + #pragma unroll + for (int j = 0; j < VEC / 2; j++) { + auto delta2 = *reinterpret_cast<const __nv_bfloat162*>( + delta_ptr + 2 * j); + v2[j] = __hadd2(v2[j], delta2); + } + *reinterpret_cast<int4*>(layer_res + (long long)tb * H + + h_base) = vp; + } + } + float2 f[4] = { + __bfloat1622float2(v2[0]), __bfloat1622float2(v2[1]), + __bfloat1622float2(v2[2]), __bfloat1622float2(v2[3])}; + tmem_store<VEC>(my_v_tmem + (si * N_CHUNK + n) * VEC, f); + #pragma unroll + for (int j = 0; j < VEC / 2; j++) { + sq_local[n] = float2_fma(f[j], f[j], sq_local[n]); + dot_local[n] = float2_fma( + f[j], make_float2(qv[2 * j], qv[2 * j + 1]), dot_local[n]); + } + } + } + }; + if constexpr (NC == 4) { + switch (an) { + case 4: + pass_A_body(std::integral_constant<int, 4>{}); + break; + case 3: + pass_A_body(std::integral_constant<int, 3>{}); + break; + case 2: + pass_A_body(std::integral_constant<int, 2>{}); + break; + case 1: + pass_A_body(std::integral_constant<int, 1>{}); + break; + default: + __builtin_unreachable(); + } + } else if constexpr (NC == 3) { + switch (an) { + case 3: + pass_A_body(std::integral_constant<int, 3>{}); + break; + case 2: + pass_A_body(std::integral_constant<int, 2>{}); + break; + case 1: + pass_A_body(std::integral_constant<int, 1>{}); + break; + default: + __builtin_unreachable(); + } + } else { + static_assert(NC == 2); + switch (an) { + case 2: + pass_A_body(std::integral_constant<int, 2>{}); + break; + case 1: + pass_A_body(std::integral_constant<int, 1>{}); + break; + default: + __builtin_unreachable(); + } + } + if (lane == 0) { + mbarrier_arrive(plan.bar_consumed[chunk_slot]); + } + tmem_store_wait(); + + float2 reduce_pair[N_CHUNK]; + #pragma unroll + for (int n = 0; n < N_CHUNK; n++) { + reduce_pair[n] = make_float2(sq_local[n].x + sq_local[n].y, + dot_local[n].x + dot_local[n].y); + } + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + #pragma unroll + for (int n = 0; n < N_CHUNK; n++) { + uint64_t packed = reinterpret_cast<uint64_t&>(reduce_pair[n]); + packed = __shfl_xor_sync(0xffffffff, packed, offset); + float2 other = reinterpret_cast<float2&>(packed); + reduce_pair[n] = float2_add(reduce_pair[n], other); + } + } + if (lane == 0) { + #pragma unroll + for (int n = 0; n < N_CHUNK; n++) { + plan.ws_stats[comp_wid][n] = reduce_pair[n]; + } + } + named_barrier_sync(CONSUMER_THREADS, 0); + + float local_rsig = 0.f; + float local_logit = 0.f; + int stat_n = lane / CONSUMER_WARPS; + int stat_w = lane % CONSUMER_WARPS; + float2 totals = {}; + if (stat_n < N_CHUNK) { + totals = plan.ws_stats[stat_w][stat_n]; + } + #pragma unroll + for (int offset = CONSUMER_WARPS / 2; offset > 0; offset >>= 1) { + totals.x += + __shfl_down_sync(0xffffffff, totals.x, offset, CONSUMER_WARPS); + totals.y += + __shfl_down_sync(0xffffffff, totals.y, offset, CONSUMER_WARPS); + } + if (stat_n < N_CHUNK && stat_w == 0) { + local_rsig = rsqrtf(totals.x / H + eps_cache); + local_logit = totals.y * local_rsig; + } + float logit_n[N_CHUNK]; + #pragma unroll + for (int n = 0; n < N_CHUNK; n++) { + logit_n[n] = __shfl_sync(0xffffffff, local_logit, n * CONSUMER_WARPS); + } + + float m_chunk = -FLT_MAX; + #pragma unroll + for (int n = 0; n < N_CHUNK; n++) { + if (n < an) m_chunk = fmaxf(m_chunk, logit_n[n]); + } + float m_new = fmaxf(m_running, m_chunk); + float corr = exp2f((m_running - m_new) * LOG2_E); + float w_n[N_CHUNK] = {}; + float w_sum = 0.f; + #pragma unroll + for (int n = 0; n < N_CHUNK; n++) { + if (n < an) { + w_n[n] = exp2f((logit_n[n] - m_new) * LOG2_E); + w_sum += w_n[n]; + } + } + + auto pass_B_body = [&](auto AN_TOK) { + constexpr int AN = decltype(AN_TOK)::value; + #pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) { + if constexpr (H == 7168) { + if (si == SLICES_PER_GROUP - 1) { + float2 corr2 = make_float2(corr, corr); + float2 a[2]; + #pragma unroll + for (int j = 0; j < 2; j++) { + float2 old = make_float2(acc32[si * VEC + 2 * j], + acc32[si * VEC + 2 * j + 1]); + a[j] = float2_mul(old, corr2); + } + float2 f_cache[AN][2]; + #pragma unroll + for (int n = 0; n < AN; n++) { + tmem_load<4>(my_v_tmem + (si * N_CHUNK + n) * VEC, + f_cache[n]); + } + #pragma unroll + for (int n = 0; n < AN; n++) { + float2 wn = make_float2(w_n[n], w_n[n]); + #pragma unroll + for (int j = 0; j < 2; j++) { + a[j] = float2_fma(wn, f_cache[n][j], a[j]); + } + } + #pragma unroll + for (int j = 0; j < 2; j++) { + acc32[si * VEC + 2 * j] = a[j].x; + acc32[si * VEC + 2 * j + 1] = a[j].y; + } + continue; + } + } + int dt = si * CONSUMER_GROUPS + group; + if (dt >= NHT) continue; + float2 corr2 = make_float2(corr, corr); + float2 a[VEC / 2]; + #pragma unroll + for (int j = 0; j < VEC / 2; j++) { + float2 old = make_float2(acc32[si * VEC + 2 * j], + acc32[si * VEC + 2 * j + 1]); + a[j] = float2_mul(old, corr2); + } + float2 f_cache[AN][VEC / 2]; + #pragma unroll + for (int n = 0; n < AN; n++) { + tmem_load<VEC>(my_v_tmem + (si * N_CHUNK + n) * VEC, f_cache[n]); + } + #pragma unroll + for (int n = 0; n < AN; n++) { + float2 wn = make_float2(w_n[n], w_n[n]); + #pragma unroll + for (int j = 0; j < VEC / 2; j++) { + a[j] = float2_fma(wn, f_cache[n][j], a[j]); + } + } + #pragma unroll + for (int j = 0; j < VEC / 2; j++) { + acc32[si * VEC + 2 * j] = a[j].x; + acc32[si * VEC + 2 * j + 1] = a[j].y; + } + } + }; + if constexpr (NC == 4) { + switch (an) { + case 4: + pass_B_body(std::integral_constant<int, 4>{}); + break; + case 3: + pass_B_body(std::integral_constant<int, 3>{}); + break; + case 2: + pass_B_body(std::integral_constant<int, 2>{}); + break; + case 1: + pass_B_body(std::integral_constant<int, 1>{}); + break; + default: + __builtin_unreachable(); + } + } else if constexpr (NC == 3) { + switch (an) { + case 3: + pass_B_body(std::integral_constant<int, 3>{}); + break; + case 2: + pass_B_body(std::integral_constant<int, 2>{}); + break; + case 1: + pass_B_body(std::integral_constant<int, 1>{}); + break; + default: + __builtin_unreachable(); + } + } else { + static_assert(NC == 2); + switch (an) { + case 2: + pass_B_body(std::integral_constant<int, 2>{}); + break; + case 1: + pass_B_body(std::integral_constant<int, 1>{}); + break; + default: + __builtin_unreachable(); + } + } + + s_running = s_running * corr + w_sum; + m_running = m_new; + } + + float inv_s = 1.f / s_running; + bf16_t* out_ptr = output + (long long)tb * H; + float2 output_sq_pair = {}; + // When output RMSNorm is fused, the softmax denominator cancels: + // (acc / s) * rsqrt(mean((acc / s)^2) + eps) + // = acc * rsqrt(mean(acc^2) + eps * s^2). + #pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) { + if constexpr (H == 7168) { + if (si == SLICES_PER_GROUP - 1) { + int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4; + uint2 packed; + auto* ov2 = reinterpret_cast<__nv_bfloat162*>(&packed); + float2 inv2 = make_float2(inv_s, inv_s); + #pragma unroll + for (int j = 0; j < 2; j++) { + float2 old = make_float2(acc32[si * VEC + 2 * j], + acc32[si * VEC + 2 * j + 1]); + if constexpr (HAS_OUTPUT_NORM) { + output_sq_pair = float2_fma(old, old, output_sq_pair); + } else { + float2 mixed = float2_mul(old, inv2); + ov2[j] = __float22bfloat162_rn(mixed); + } + } + if constexpr (!HAS_OUTPUT_NORM) { + *reinterpret_cast<uint2*>(out_ptr + h_base) = packed; + } + continue; + } + } + int dt = si * CONSUMER_GROUPS + group; + if (dt >= NHT) continue; + int h_base = dt * K_TILE + k_local; + uint4 packed; + auto* ov2 = reinterpret_cast<__nv_bfloat162*>(&packed); + float2 inv2 = make_float2(inv_s, inv_s); + #pragma unroll + for (int j = 0; j < VEC / 2; j++) { + float2 old = + make_float2(acc32[si * VEC + 2 * j], acc32[si * VEC + 2 * j + 1]); + if constexpr (HAS_OUTPUT_NORM) { + output_sq_pair = float2_fma(old, old, output_sq_pair); + } else { + float2 mixed = float2_mul(old, inv2); + ov2[j] = __float22bfloat162_rn(mixed); + } + } + if constexpr (!HAS_OUTPUT_NORM) { + *reinterpret_cast<uint4*>(out_ptr + h_base) = packed; + } + } + + if constexpr (HAS_OUTPUT_NORM) { + if constexpr (OUTPUT_NORM_IN_SMEM) { + // The immutable weight copy is acquired once, at its first use. + if (tb == blockIdx.x) { + mbarrier_wait(plan.bar_output_norm_ready, 0); + } + } + float output_sq = output_sq_pair.x + output_sq_pair.y; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + output_sq += __shfl_xor_sync(0xffffffff, output_sq, offset); + } + if (lane == 0) { + plan.ws_stats[comp_wid][0] = make_float2(output_sq, 0.f); + } + named_barrier_sync(CONSUMER_THREADS, 0); + float total_sq = lane < CONSUMER_WARPS ? plan.ws_stats[lane][0].x : 0.f; + #pragma unroll + for (int offset = CONSUMER_WARPS / 2; offset > 0; offset >>= 1) { + total_sq += + __shfl_down_sync(0xffffffff, total_sq, offset, CONSUMER_WARPS); + } + if (lane == 0) { + total_sq = + rsqrtf(total_sq / H + output_norm_eps * s_running * s_running); + } + float output_rsigma = __shfl_sync(0xffffffff, total_sq, 0); + #pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) { + if constexpr (H == 7168) { + if (si == SLICES_PER_GROUP - 1) { + int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4; + uint2 packed; + auto* values = reinterpret_cast<bf16_t*>(&packed); + #pragma unroll + for (int j = 0; j < 4; j++) { + const bf16_t* weight_ptr = + OUTPUT_NORM_IN_SMEM ? output_norm_buf : output_norm_weight; + float weight = __bfloat162float(weight_ptr[h_base + j]); + values[j] = __float2bfloat16(acc32[si * VEC + j] * + output_rsigma * weight); + } + *reinterpret_cast<uint2*>(out_ptr + h_base) = packed; + continue; + } + } + int dt = si * CONSUMER_GROUPS + group; + if (dt >= NHT) continue; + int h_base = dt * K_TILE + k_local; + uint4 packed; + auto* values = reinterpret_cast<bf16_t*>(&packed); + #pragma unroll + for (int j = 0; j < VEC; j++) { + const bf16_t* weight_ptr = + OUTPUT_NORM_IN_SMEM ? output_norm_buf : output_norm_weight; + float weight = __bfloat162float(weight_ptr[h_base + j]); + values[j] = + __float2bfloat16(acc32[si * VEC + j] * output_rsigma * weight); + } + *reinterpret_cast<uint4*>(out_ptr + h_base) = packed; + } + } + } + } + + cudaTriggerProgrammaticLaunchCompletion(); + __syncthreads(); + if (wid == 1) { + tmem_free(plan.tmem_base, TMEM_COLS_ALLOC); + } +#else + if (threadIdx.x == 0) { + printf("attn_res_fwd_online_v2_kernel requires sm_10x\n"); + } +#endif +} + +template <int H, int NC = N_CHUNK_DEFAULT, bool RELEASE_TMEM = false, + bool HAS_DELTA = false, bool HAS_OUTPUT_NORM = false, + bool OUTPUT_NORM_IN_SMEM = false> +static void launch_fwd(const bf16_t* block_residual, bf16_t* layer_residual, + const bf16_t* delta, const bf16_t* res_weight, + const bf16_t* rms_weight, bf16_t* output, int N, int T, + int B, float rms_eps, int num_sm, cudaStream_t stream, + const bf16_t* output_norm_weight = nullptr, + float output_norm_eps = 0.f, int block_stride_m = 0, + int block_stride_r = 0) { + constexpr size_t smem_size = + ((size_t)CHUNK_DEPTH * (NC + (HAS_DELTA ? 1 : 0)) * H * sizeof(bf16_t) + + (OUTPUT_NORM_IN_SMEM ? (size_t)H * sizeof(bf16_t) : 0) + + sizeof(FwdSmemPlan<NC>) + 15) & + ~size_t(15); + auto kernel = + &attn_res_fwd_online_v2_kernel<H, NC, RELEASE_TMEM, HAS_DELTA, + HAS_OUTPUT_NORM, OUTPUT_NORM_IN_SMEM>; + static bool attrs_set = false; + if (!attrs_set) { + if (smem_size > 48 * 1024) { + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + } + attrs_set = true; + } + int grid = RELEASE_TMEM ? num_sm * 2 : num_sm; + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = BLK; + config.dynamicSmemBytes = smem_size; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attrs; + config.numAttrs = 1; + cudaLaunchKernelEx(&config, kernel, block_residual, layer_residual, delta, + res_weight, rms_weight, output, N, T, B, block_stride_m, + block_stride_r, rms_eps, output_norm_weight, + output_norm_eps); +} + +} // namespace fwd_prod_v2 +} // namespace sm100 + +void kimi_k3_attn_res(torch::stable::Tensor& prefix, + torch::stable::Tensor const& delta, + torch::stable::Tensor const& blocks, + torch::stable::Tensor const& norm_weight, + torch::stable::Tensor const& qk_weight, + torch::stable::Tensor const& output_norm_weight, + torch::stable::Tensor& output, int64_t num_blocks, + double eps, double output_norm_eps) { + int const num_tokens = static_cast<int>(prefix.size(0)); + int const device = prefix.get_device_index(); + torch::stable::accelerator::DeviceGuard const device_guard(device); + cudaDeviceProp const* properties = get_device_prop(); + STD_TORCH_CHECK(properties->major == 10, + "Kimi K3 AttnRes requires the SM100 family"); + + using namespace sm100::fwd_prod_v2; + // Two-source chunks and two resident CTAs are beneficial once setup is + // amortized by the long, full eight-block prefill workload. + if (num_blocks == 8 && num_tokens >= 4096) { + launch_fwd<7168, 2, true, true, true, true>( + static_cast<bf16_t const*>(blocks.data_ptr()), + static_cast<bf16_t*>(prefix.data_ptr()), + static_cast<bf16_t const*>(delta.data_ptr()), + static_cast<bf16_t const*>(qk_weight.data_ptr()), + static_cast<bf16_t const*>(norm_weight.data_ptr()), + static_cast<bf16_t*>(output.data_ptr()), + static_cast<int>(num_blocks) + 1, num_tokens, 1, + static_cast<float>(eps), properties->multiProcessorCount, + get_current_cuda_stream(device), + static_cast<bf16_t const*>(output_norm_weight.data_ptr()), + static_cast<float>(output_norm_eps), static_cast<int>(blocks.stride(0)), + static_cast<int>(blocks.stride(1))); + } else { + launch_fwd<7168, 4, false, true, true, true>( + static_cast<bf16_t const*>(blocks.data_ptr()), + static_cast<bf16_t*>(prefix.data_ptr()), + static_cast<bf16_t const*>(delta.data_ptr()), + static_cast<bf16_t const*>(qk_weight.data_ptr()), + static_cast<bf16_t const*>(norm_weight.data_ptr()), + static_cast<bf16_t*>(output.data_ptr()), + static_cast<int>(num_blocks) + 1, num_tokens, 1, + static_cast<float>(eps), properties->multiProcessorCount, + get_current_cuda_stream(device), + static_cast<bf16_t const*>(output_norm_weight.data_ptr()), + static_cast<float>(output_norm_eps), static_cast<int>(blocks.stride(0)), + static_cast<int>(blocks.stride(1))); + } + cudaError_t const error = cudaGetLastError(); + STD_TORCH_CHECK( + error == cudaSuccess, + "Kimi K3 AttnRes kernel launch failed: ", cudaGetErrorString(error)); +} diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 3834bea5857..5a9c91d4563 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -315,6 +315,17 @@ void fused_minimax_m3_qknorm_rope_kv_insert( std::optional<torch::stable::Tensor> index_q_out, const std::string& kv_cache_dtype, bool skip_index_branch); +#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES +void kimi_k3_attn_res(torch::stable::Tensor& prefix, + torch::stable::Tensor const& delta, + torch::stable::Tensor const& blocks, + torch::stable::Tensor const& norm_weight, + torch::stable::Tensor const& qk_weight, + torch::stable::Tensor const& output_norm_weight, + torch::stable::Tensor& output, int64_t num_blocks, + double eps, double output_norm_eps); +#endif + // Sampler kernels (shared CUDA/ROCm) void apply_repetition_penalties_( torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 0a475d02c6f..c364e211474 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -468,6 +468,14 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "int block_size, Tensor!? q_out, Tensor!? index_q_out, " "str kv_cache_dtype, bool skip_index_branch=False) -> ()"); +#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES + ops.def( + "kimi_k3_attn_res(" + "Tensor! prefix, Tensor delta, Tensor blocks, Tensor norm_weight, " + "Tensor qk_weight, Tensor output_norm_weight, Tensor! output, " + "int num_blocks, float eps, float output_norm_eps) -> ()"); +#endif + // Apply repetition penalties to logits in-place. ops.def( "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, " @@ -693,6 +701,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { #endif ops.impl("fused_minimax_m3_qknorm_rope_kv_insert", TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert)); +#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES + ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res)); +#endif // Sampler kernels (shared CUDA/ROCm) ops.impl("apply_repetition_penalties_", diff --git a/tests/models/kimi_k3/test_amd_attn_res.py b/tests/models/kimi_k3/test_amd_attn_res.py new file mode 100644 index 00000000000..f6dfaa422b3 --- /dev/null +++ b/tests/models/kimi_k3/test_amd_attn_res.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch.nn.functional as F + +from vllm.models.kimi_k3.amd.ops.attn_res import attn_res +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), + reason="AMD AttnRes requires ROCm", +) + + +def _randn_with_row_padding(*shape: int, padding: int = 0) -> torch.Tensor: + storage = torch.randn( + *shape[:-1], + shape[-1] + padding, + device="cuda", + dtype=torch.bfloat16, + ) + return storage[..., : shape[-1]] + + +def _reference( + prefix: torch.Tensor, + blocks: torch.Tensor, + norm_weight: torch.Tensor, + qk_weight: torch.Tensor, + num_blocks: int, + eps: float, +) -> torch.Tensor: + hidden_size = prefix.shape[-1] + values = torch.cat((blocks[:, :num_blocks], prefix.unsqueeze(1)), dim=1) + keys = F.rms_norm(values, (hidden_size,), norm_weight, eps) + probs = (keys @ qk_weight).softmax(dim=-1) + return torch.matmul(probs.unsqueeze(1), values).squeeze(1) + + +@pytest.mark.parametrize( + ( + "num_tokens", + "num_blocks", + "block_capacity", + "hidden_size", + "row_padding", + ), + [ + pytest.param(0, 3, 5, 128, 0, id="empty"), + pytest.param(1, 1, 2, 128, 0, id="decode-single"), + pytest.param(17, 4, 6, 1024, 7, id="decode-padded"), + pytest.param(320, 8, 10, 7168, 0, id="prefill-full"), + ], +) +def test_amd_attn_res_matches_reference( + num_tokens: int, + num_blocks: int, + block_capacity: int, + hidden_size: int, + row_padding: int, +) -> None: + eps = 1e-5 + prefix = _randn_with_row_padding(num_tokens, hidden_size, padding=row_padding) + blocks = _randn_with_row_padding( + num_tokens, + block_capacity, + hidden_size, + padding=row_padding, + ) + norm_weight = 1 + 0.1 * torch.randn( + hidden_size, device="cuda", dtype=torch.bfloat16 + ) + qk_weight = ( + torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16) / hidden_size**0.5 + ) + expected = _reference( + prefix, + blocks, + norm_weight, + qk_weight, + num_blocks, + eps, + ) + original_prefix = prefix.clone() + original_blocks = blocks.clone() + + actual = attn_res( + prefix, + blocks, + norm_weight, + qk_weight, + num_blocks, + eps, + ) + + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + torch.testing.assert_close(prefix, original_prefix, atol=0, rtol=0) + torch.testing.assert_close(blocks, original_blocks, atol=0, rtol=0) + assert actual.shape == prefix.shape + assert actual.is_contiguous() diff --git a/tests/models/kimi_k3/test_attn_res.py b/tests/models/kimi_k3/test_attn_res.py new file mode 100644 index 00000000000..69c9213647a --- /dev/null +++ b/tests/models/kimi_k3/test_attn_res.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch.nn.functional as F + +from vllm.models.kimi_k3.nvidia.ops import attn_res +from vllm.platforms import current_platform + +HIDDEN_SIZE = 7168 +MAX_BLOCKS = 8 +EPS = 1e-5 + + +def _randn_with_row_padding(*shape: int, padding: int = 0) -> torch.Tensor: + storage = torch.randn( + *shape[:-1], + shape[-1] + padding, + device="cuda", + dtype=torch.bfloat16, + ) + return storage[..., : shape[-1]] + + +def _reference( + prefix: torch.Tensor, + delta: torch.Tensor | None, + blocks: torch.Tensor, + norm_weight: torch.Tensor, + qk_weight: torch.Tensor, + output_norm_weight: torch.Tensor | None, + num_blocks: int, +) -> tuple[torch.Tensor, torch.Tensor]: + if delta is not None: + prefix = prefix + delta + values = torch.cat((blocks[:, :num_blocks], prefix.unsqueeze(1)), dim=1) + keys = F.rms_norm(values, (HIDDEN_SIZE,), norm_weight, EPS) + probs = (keys @ qk_weight).softmax(dim=-1) + output = torch.matmul(probs.unsqueeze(1), values).squeeze(1) + if output_norm_weight is not None: + output = F.rms_norm(output, (HIDDEN_SIZE,), output_norm_weight, EPS) + return output, prefix + + +@pytest.mark.parametrize( + ( + "num_tokens", + "num_blocks", + "row_padding", + "write_block", + "has_delta", + "backend", + ), + [ + pytest.param(1, 0, 0, True, False, "triton", id="triton-empty"), + pytest.param(1, 0, 0, True, True, "triton", id="triton-empty-add"), + pytest.param(17, 5, 7, True, False, "triton", id="triton-write"), + pytest.param(17, 5, 7, True, True, "triton", id="triton-write-add"), + pytest.param(3, 8, 0, False, False, "triton", id="triton-full"), + pytest.param(3, 8, 0, False, True, "triton", id="triton-full-add"), + pytest.param(320, 1, 0, False, True, "nvidia", id="nvidia-1"), + pytest.param(320, 4, 0, False, True, "nvidia", id="nvidia-4"), + pytest.param(320, 8, 0, False, True, "nvidia", id="nvidia-8"), + ], +) +def test_attn_res( + num_tokens: int, + num_blocks: int, + row_padding: int, + write_block: bool, + has_delta: bool, + backend: str, +): + if backend == "nvidia" and not current_platform.is_device_capability_family(100): + pytest.skip("NVIDIA AttnRes requires the SM100 family") + + prefix = _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=row_padding) + delta = ( + _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=row_padding) + if has_delta + else None + ) + blocks = _randn_with_row_padding( + num_tokens, MAX_BLOCKS, HIDDEN_SIZE, padding=row_padding + ) + norm_weight = 1 + 0.1 * torch.randn( + HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16 + ) + qk_weight = ( + torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5 + ) + output_norm_weight = 1 + 0.1 * torch.randn( + HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16 + ) + original_blocks = blocks.clone() + expected, expected_prefix = _reference( + prefix.clone(), + delta, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + num_blocks, + ) + block_write_idx = num_blocks if write_block else -1 + + actual = attn_res( + prefix, + delta, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + num_blocks, + block_write_idx, + EPS, + EPS, + ) + + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + torch.testing.assert_close(prefix, expected_prefix, atol=0, rtol=0) + if write_block: + original_blocks[:, block_write_idx].copy_(expected_prefix) + torch.testing.assert_close(blocks, original_blocks, atol=0, rtol=0) + assert actual.is_contiguous() + + +@pytest.mark.parametrize("num_blocks", range(MAX_BLOCKS + 1)) +def test_attn_res_block_counts(num_blocks: int): + prefix = torch.randn(1, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + blocks = torch.randn( + 1, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16 + ) + norm_weight = torch.ones(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + qk_weight = ( + torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5 + ) + output_norm_weight = torch.ones_like(norm_weight) + expected, _ = _reference( + prefix.clone(), + None, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + num_blocks, + ) + + actual = attn_res( + prefix, + None, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + num_blocks, + -1, + EPS, + EPS, + ) + + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + + +def test_attn_res_without_output_norm(): + prefix = torch.randn(7, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + delta = torch.randn_like(prefix) + blocks = torch.randn( + 7, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16 + ) + norm_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + qk_weight = ( + torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5 + ) + expected, _ = _reference( + prefix.clone(), delta, blocks, norm_weight, qk_weight, None, MAX_BLOCKS + ) + + actual = attn_res( + prefix, + delta, + blocks, + norm_weight, + qk_weight, + None, + MAX_BLOCKS, + -1, + EPS, + 0.0, + ) + + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 07a3a583f95..b1d2f1344d6 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2746,6 +2746,33 @@ def concat_and_cache_mla( ) +def kimi_k3_attn_res( + prefix: torch.Tensor, + delta: torch.Tensor, + blocks: torch.Tensor, + norm_weight: torch.Tensor, + qk_weight: torch.Tensor, + output_norm_weight: torch.Tensor, + num_blocks: int, + eps: float, + output_norm_eps: float, +) -> torch.Tensor: + output = torch.empty_like(prefix) + torch.ops._C.kimi_k3_attn_res( + prefix, + delta, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + output, + num_blocks, + eps, + output_norm_eps, + ) + return output + + def concat_and_cache_mla_rope_fused( positions: torch.Tensor, q_pe: torch.Tensor, diff --git a/vllm/models/kimi_k3/amd/__init__.py b/vllm/models/kimi_k3/amd/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm/models/kimi_k3/amd/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/amd/ops/__init__.py b/vllm/models/kimi_k3/amd/ops/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm/models/kimi_k3/amd/ops/attn_res.py b/vllm/models/kimi_k3/amd/ops/attn_res.py new file mode 100644 index 00000000000..c00a3422ca1 --- /dev/null +++ b/vllm/models/kimi_k3/amd/ops/attn_res.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code adapted from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _attn_res_kernel( + prefix_ptr, + blocks_ptr, + norm_weight_ptr, + qk_weight_ptr, + output_ptr, + stride_prefix_m: tl.constexpr, + stride_block_m: tl.constexpr, + stride_block_r: tl.constexpr, + stride_output_m: tl.constexpr, + num_blocks: tl.constexpr, + hidden_size: tl.constexpr, + eps: tl.constexpr, + BLOCK_L: tl.constexpr, + BLOCK_D: tl.constexpr, +): + row_idx = tl.program_id(0).to(tl.int64) + d_offsets = tl.max_contiguous(tl.arange(0, BLOCK_D), BLOCK_D) + d_mask = d_offsets < hidden_size + + prefix = tl.load( + prefix_ptr + row_idx * stride_prefix_m + d_offsets, + mask=d_mask, + other=0.0, + ).to(tl.float32) + input_qk_weight = tl.load(norm_weight_ptr + d_offsets, mask=d_mask, other=0.0).to( + tl.float32 + ) * tl.load(qk_weight_ptr + d_offsets, mask=d_mask, other=0.0).to(tl.float32) + + max_logit = tl.full((), -float("inf"), tl.float32) + denominator = tl.zeros((), tl.float32) + mixed = tl.zeros((BLOCK_D,), tl.float32) + num_sources = num_blocks + 1 + + for source_tile in range(tl.cdiv(num_sources, BLOCK_L)): + source_offsets = source_tile * BLOCK_L + tl.arange(0, BLOCK_L) + source_mask = source_offsets < num_sources + is_prefix = source_offsets == num_blocks + block_ptrs = ( + blocks_ptr + + row_idx * stride_block_m + + source_offsets[:, None] * stride_block_r + + d_offsets[None, :] + ) + block_values = tl.load( + block_ptrs, + mask=(source_mask[:, None] & ~is_prefix[:, None] & d_mask[None, :]), + other=0.0, + eviction_policy="evict_first", + ).to(tl.float32) + values = tl.where(is_prefix[:, None], prefix[None, :], block_values) + reciprocal_std = tl.rsqrt( + tl.sum(values * values, axis=1) * (1.0 / hidden_size) + eps + ) + logits = tl.sum(values * input_qk_weight[None, :], axis=1) * reciprocal_std + scores = tl.where(source_mask, logits, -float("inf")) + + new_max_logit = tl.maximum(max_logit, tl.max(scores, axis=0)) + old_scale = tl.exp(max_logit - new_max_logit) + block_scales = tl.exp(scores - new_max_logit) + denominator = denominator * old_scale + tl.sum(block_scales, axis=0) + mixed = mixed * old_scale + tl.sum(block_scales[:, None] * values, axis=0) + max_logit = new_max_logit + + output = mixed / denominator + tl.store( + output_ptr + row_idx * stride_output_m + d_offsets, + output, + mask=d_mask, + ) + + +def attn_res( + prefix: torch.Tensor, + blocks: torch.Tensor, + norm_weight: torch.Tensor, + qk_weight: torch.Tensor, + num_blocks: int, + eps: float, +) -> torch.Tensor: + num_tokens, hidden_size = prefix.shape + assert 0 < num_blocks <= blocks.shape[1] + assert blocks.shape[0] == num_tokens + assert norm_weight.numel() == hidden_size + assert qk_weight.numel() == hidden_size + assert prefix.stride(-1) == 1 + assert blocks.stride(-1) == 1 + assert norm_weight.stride(-1) == 1 + assert qk_weight.stride(-1) == 1 + + output = prefix.new_empty(prefix.shape) + if num_tokens == 0: + return output + + if num_tokens >= 256 or num_blocks <= 1: + block_l, num_warps = 1, 4 + else: + block_l, num_warps = 4, 8 + _attn_res_kernel[(num_tokens,)]( + prefix, + blocks, + norm_weight, + qk_weight, + output, + prefix.stride(0), + blocks.stride(0), + blocks.stride(1), + output.stride(0), + num_blocks, + hidden_size, + eps, + BLOCK_L=block_l, + BLOCK_D=triton.next_power_of_2(hidden_size), + num_warps=num_warps, + num_stages=2, + ) + return output diff --git a/vllm/models/kimi_k3/nvidia/__init__.py b/vllm/models/kimi_k3/nvidia/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/nvidia/ops/__init__.py b/vllm/models/kimi_k3/nvidia/ops/__init__.py new file mode 100644 index 00000000000..bbaee887ba2 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .attn_res import attn_res + +__all__ = ["attn_res"] diff --git a/vllm/models/kimi_k3/nvidia/ops/attn_res.py b/vllm/models/kimi_k3/nvidia/ops/attn_res.py new file mode 100644 index 00000000000..01078d6c01d --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/attn_res.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code adapted from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li + + +import torch + +from vllm import _custom_ops as ops +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + + +# Consumed by kimi_k3_triton_warmup.py during kernel_warmup(). +def get_attn_res_triton_warmup_profiles( + max_blocks: int, +) -> tuple[tuple[int, bool, int, bool], ...]: + """Return the small-batch profiles that bypass the native kernel.""" + profiles = [ + (num_blocks, False, -1, True) for num_blocks in range(2, max_blocks + 1) + ] + profiles.extend( + (block_write_idx, True, block_write_idx, True) + for block_write_idx in range(2, max_blocks) + ) + profiles.append((max_blocks, True, -1, False)) + return tuple(profiles) + + +@triton.jit +def _attn_res_kernel( + prefix_ptr, + delta_ptr, + blocks_ptr, + norm_weight_ptr, + qk_weight_ptr, + output_norm_weight_ptr, + output_ptr, + stride_prefix_m: tl.constexpr, + stride_delta_m: tl.constexpr, + stride_block_m: tl.constexpr, + stride_block_r: tl.constexpr, + stride_output_m: tl.constexpr, + num_blocks: tl.constexpr, + hidden_size: tl.constexpr, + block_write_idx: tl.constexpr, + eps: tl.constexpr, + output_norm_eps: tl.constexpr, + HAS_DELTA: tl.constexpr, + WRITE_BLOCK: tl.constexpr, + APPLY_OUTPUT_NORM: tl.constexpr, + BLOCK_L: tl.constexpr, + BLOCK_D: tl.constexpr, + launch_pdl: tl.constexpr, +): + row_idx = tl.program_id(0).to(tl.int64) + d_offsets = tl.max_contiguous(tl.arange(0, BLOCK_D), BLOCK_D) + d_mask = d_offsets < hidden_size + + if launch_pdl: + tl.extra.cuda.gdc_wait() + + updated_prefix = tl.load( + prefix_ptr + row_idx * stride_prefix_m + d_offsets, + mask=d_mask, + other=0.0, + ).to(tl.float32) + if HAS_DELTA: + delta = tl.load( + delta_ptr + row_idx * stride_delta_m + d_offsets, + mask=d_mask, + other=0.0, + ).to(tl.float32) + updated_prefix += delta + # Match the BF16 prefix-add result before using it as a residual source. + updated_prefix = updated_prefix.to(prefix_ptr.dtype.element_ty).to(tl.float32) + tl.store( + prefix_ptr + row_idx * stride_prefix_m + d_offsets, + updated_prefix, + mask=d_mask, + ) + if WRITE_BLOCK: + tl.store( + blocks_ptr + + row_idx * stride_block_m + + block_write_idx * stride_block_r + + d_offsets, + updated_prefix, + mask=d_mask, + ) + # With only the prefix source, the AttnRes softmax is exactly one. + if num_blocks == 0: + mixed = updated_prefix + else: + # Reloading avoids keeping the full prefix vector live across the loop. + if HAS_DELTA: + tl.debug_barrier() + input_qk_weight = tl.load( + norm_weight_ptr + d_offsets, mask=d_mask, other=0.0 + ).to(tl.float32) * tl.load( + qk_weight_ptr + d_offsets, mask=d_mask, other=0.0 + ).to(tl.float32) + max_logit = tl.full((), -float("inf"), tl.float32) + denominator = tl.zeros((), tl.float32) + mixed = tl.zeros((BLOCK_D,), tl.float32) + + num_sources = num_blocks + 1 + for source_tile in range(tl.cdiv(num_sources, BLOCK_L)): + source_offsets = source_tile * BLOCK_L + tl.arange(0, BLOCK_L) + source_mask = source_offsets < num_sources + is_prefix = source_offsets == num_blocks + block_ptrs = ( + blocks_ptr + + row_idx * stride_block_m + + source_offsets[:, None] * stride_block_r + + d_offsets[None, :] + ) + prefix_ptrs = ( + prefix_ptr + + row_idx * stride_prefix_m + + source_offsets[:, None] * 0 + + d_offsets[None, :] + ) + value_ptrs = tl.where(is_prefix[:, None], prefix_ptrs, block_ptrs) + values = tl.load( + value_ptrs, + mask=source_mask[:, None] & d_mask[None, :], + other=0.0, + eviction_policy="evict_first", + ).to(tl.float32) + reciprocal_std = tl.rsqrt( + tl.sum(values * values, axis=1) * (1.0 / hidden_size) + eps + ) + logits = tl.sum(values * input_qk_weight[None, :], axis=1) * reciprocal_std + scores = tl.where(source_mask, logits, -float("inf")) + + new_max_logit = tl.maximum(max_logit, tl.max(scores, axis=0)) + old_scale = tl.exp(max_logit - new_max_logit) + block_scales = tl.exp(scores - new_max_logit) + denominator = denominator * old_scale + tl.sum(block_scales, axis=0) + mixed = mixed * old_scale + tl.sum(block_scales[:, None] * values, axis=0) + max_logit = new_max_logit + + mixed /= denominator + output = mixed + + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + + if APPLY_OUTPUT_NORM: + output_reciprocal_std = tl.rsqrt( + tl.sum(tl.where(d_mask, mixed * mixed, 0.0), axis=0) * (1.0 / hidden_size) + + output_norm_eps + ) + output_norm_weight = tl.load( + output_norm_weight_ptr + d_offsets, mask=d_mask, other=0.0 + ).to(tl.float32) + output = mixed * output_reciprocal_std * output_norm_weight + tl.store( + output_ptr + row_idx * stride_output_m + d_offsets, + output, + mask=d_mask, + ) + + +def attn_res( + prefix: torch.Tensor, + delta: torch.Tensor | None, + blocks: torch.Tensor, + norm_weight: torch.Tensor, + qk_weight: torch.Tensor, + output_norm_weight: torch.Tensor | None, + num_blocks: int, + block_write_idx: int, + eps: float, + output_norm_eps: float, +) -> torch.Tensor: + num_tokens, hidden_size = prefix.shape + assert prefix.stride(-1) == 1 + assert delta is None or delta.stride(-1) == 1 + assert blocks.stride(-1) == 1 + assert norm_weight.stride(-1) == 1 + assert qk_weight.stride(-1) == 1 + assert output_norm_weight is None or output_norm_weight.stride(-1) == 1 + # The in-tree NVIDIA kernel covers the common fused-add + output-norm path; + # Triton handles block boundaries and final pre-norm output. + if ( + hidden_size == 7168 + and delta is not None + and output_norm_weight is not None + and num_blocks > 0 + and block_write_idx < 0 + and current_platform.is_device_capability_family(100) + ): + return ops.kimi_k3_attn_res( + prefix, + delta, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + num_blocks, + eps, + output_norm_eps, + ) + output = prefix.new_empty(prefix.shape) + # Tuned on GB300: source tiling helps decode, while one-source tiles scale + # better for prefill. + # Keep get_attn_res_triton_warmup_profiles in sync with these fallbacks. + if num_tokens >= 256 or num_blocks <= 1: + block_l, num_warps = 1, 4 + else: + block_l, num_warps = 4, 8 + _attn_res_kernel[(num_tokens,)]( + prefix, + delta, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + output, + prefix.stride(0), + 0 if delta is None else delta.stride(0), + blocks.stride(0), + blocks.stride(1), + output.stride(0), + num_blocks, + hidden_size, + block_write_idx, + eps, + output_norm_eps, + HAS_DELTA=delta is not None, + WRITE_BLOCK=block_write_idx >= 0, + APPLY_OUTPUT_NORM=output_norm_weight is not None, + BLOCK_L=block_l, + BLOCK_D=triton.next_power_of_2(hidden_size), + num_warps=num_warps, + num_stages=2, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return output From 5ed3faa43ddf075f24482396b634edc33e047a40 Mon Sep 17 00:00:00 2001 From: Bugen Zhao <i@bugenzhao.com> Date: Tue, 28 Jul 2026 15:46:51 +0800 Subject: [PATCH 159/185] [Rust Frontend] Add ordinary-text tokenizer encoding (#49992) Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com> --- rust/src/chat/src/renderer/inkling/tests.rs | 4 + rust/src/parser/benches/utils/adapter.rs | 4 + rust/src/parser/src/unified/inkling.rs | 8 + rust/src/text/src/backend/hf/mod.rs | 4 + rust/src/tokenizer/src/hf.rs | 290 +++++++++++++++++++- rust/src/tokenizer/src/incremental.rs | 12 + rust/src/tokenizer/src/lib.rs | 4 + rust/src/tokenizer/src/tekken.rs | 54 ++++ rust/src/tokenizer/src/test_utils.rs | 30 ++ rust/src/tokenizer/src/tiktoken.rs | 40 +++ 10 files changed, 448 insertions(+), 2 deletions(-) diff --git a/rust/src/chat/src/renderer/inkling/tests.rs b/rust/src/chat/src/renderer/inkling/tests.rs index d0d55e1be2d..57e08ac0158 100644 --- a/rust/src/chat/src/renderer/inkling/tests.rs +++ b/rust/src/chat/src/renderer/inkling/tests.rs @@ -31,6 +31,10 @@ impl Tokenizer for FixtureTokenizer { Ok(text.bytes().map(u32::from).collect()) } + fn encode_ordinary(&self, text: &str) -> vllm_tokenizer::Result<Vec<u32>> { + self.encode(text, false) + } + fn decode( &self, token_ids: &[u32], diff --git a/rust/src/parser/benches/utils/adapter.rs b/rust/src/parser/benches/utils/adapter.rs index 243f19e3ce8..9cba325116c 100644 --- a/rust/src/parser/benches/utils/adapter.rs +++ b/rust/src/parser/benches/utils/adapter.rs @@ -19,6 +19,10 @@ impl Tokenizer for BenchTokenizer { Ok(text.chars().map(|_| u32::MAX).collect()) } + fn encode_ordinary(&self, text: &str) -> vllm_tokenizer::Result<Vec<u32>> { + self.encode(text, false) + } + fn decode( &self, token_ids: &[u32], diff --git a/rust/src/parser/src/unified/inkling.rs b/rust/src/parser/src/unified/inkling.rs index eb78321d126..1ce69d8bfb4 100644 --- a/rust/src/parser/src/unified/inkling.rs +++ b/rust/src/parser/src/unified/inkling.rs @@ -414,6 +414,10 @@ mod tests { Ok(text.chars().map(u32::from).collect()) } + fn encode_ordinary(&self, text: &str) -> vllm_tokenizer::Result<Vec<u32>> { + self.encode(text, false) + } + fn decode( &self, token_ids: &[u32], @@ -733,6 +737,10 @@ mod tests { Ok(vec![]) } + fn encode_ordinary(&self, text: &str) -> vllm_tokenizer::Result<Vec<u32>> { + self.encode(text, false) + } + fn decode( &self, _token_ids: &[u32], diff --git a/rust/src/text/src/backend/hf/mod.rs b/rust/src/text/src/backend/hf/mod.rs index 49ae5dbd6b9..4fc7a18753a 100644 --- a/rust/src/text/src/backend/hf/mod.rs +++ b/rust/src/text/src/backend/hf/mod.rs @@ -177,6 +177,10 @@ mod tests { Ok(vec![]) } + fn encode_ordinary(&self, text: &str) -> vllm_tokenizer::Result<Vec<u32>> { + self.encode(text, false) + } + fn decode( &self, _token_ids: &[u32], diff --git a/rust/src/tokenizer/src/hf.rs b/rust/src/tokenizer/src/hf.rs index 08ec5a22d6b..eb527294a5b 100644 --- a/rust/src/tokenizer/src/hf.rs +++ b/rust/src/tokenizer/src/hf.rs @@ -1,13 +1,20 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +use std::borrow::Cow; use std::path::Path; -use std::sync::Arc; +use std::sync::{Arc, LazyLock}; use fastokens::Tokenizer as FastokensTokenizer; use fastokens::decoders::Decoder as FastokensDecoder; +use fastokens::pre_tokenized::{ + PreTokenizedString as FastokensPreTokenizedString, Split as FastokensSplit, +}; +use fastokens::{PreTokenizer as FastokensPreTokenizer, Split as FastokensSplitPreTokenizer}; use thiserror_ext::AsReport as _; -use tokenizers::Tokenizer as HfTokenizer; +use tokenizers::{ + AddedVocabulary, Model as _, OffsetType, PreTokenizer as _, Tokenizer as HfTokenizer, +}; use tracing::{info, warn}; use crate::byte_level_decode::decode_byte_level; @@ -16,6 +23,8 @@ use crate::{Result, Tokenizer}; mod added_tokens; +static EMPTY_HF_ADDED_VOCABULARY: LazyLock<AddedVocabulary> = LazyLock::new(AddedVocabulary::new); + enum Backend { Hf(Box<HfTokenizer>), Fastokens(Box<FastokensTokenizer>), @@ -53,6 +62,85 @@ fn decode_fastokens_byte_level( Ok(decode_byte_level(tokens)) } +fn encode_hf_ordinary(tokenizer: &HfTokenizer, text: &str) -> tokenizers::Result<Vec<u32>> { + let mut pretokenized = + EMPTY_HF_ADDED_VOCABULARY.extract_and_normalize(tokenizer.get_normalizer(), text); + + if let Some(pre_tokenizer) = tokenizer.get_pre_tokenizer() { + pre_tokenizer.pre_tokenize(&mut pretokenized)?; + } + pretokenized.tokenize(|normalized| tokenizer.get_model().tokenize(normalized.get()))?; + let encoding = pretokenized.into_encoding(None, 0, OffsetType::Byte)?; + let encoding = tokenizer.post_process(encoding, None, false)?; + Ok(encoding.get_ids().to_vec()) +} + +fn fastokens_fused_split(tokenizer: &FastokensTokenizer) -> Option<&FastokensSplitPreTokenizer> { + // Keep this predicate aligned with fastokens::Tokenizer::detect_fused_byte_level. + let FastokensPreTokenizer::Sequence(steps) = tokenizer.pre_tokenizer()? else { + return None; + }; + let [ + FastokensPreTokenizer::Split(split), + FastokensPreTokenizer::ByteLevel(byte_level), + ] = steps.as_slice() + else { + return None; + }; + byte_level.is_bulk_only().then_some(split) +} + +fn fastokens_pre_tokenized_ordinary( + tokenizer: &FastokensTokenizer, + text: &str, +) -> FastokensPreTokenizedString { + // This is fastokens::Tokenizer::build_pre_tokenized with added_tokens = None. + let normalized = tokenizer + .normalizer() + .map_or(Cow::Borrowed(text), |normalizer| normalizer.normalize(text)); + match normalized { + Cow::Borrowed(_) => FastokensPreTokenizedString::from_text(text), + Cow::Owned(text) => { + let len = text.len(); + FastokensPreTokenizedString::new( + text, + vec![FastokensSplit { + range: 0..len, + token_id: None, + }], + ) + } + } +} + +fn encode_fastokens_ordinary( + tokenizer: &FastokensTokenizer, + text: &str, +) -> std::result::Result<Vec<u32>, fastokens::Error> { + if text.is_empty() { + return Ok(Vec::new()); + } + + let mut pretokenized = fastokens_pre_tokenized_ordinary(tokenizer, text); + let ids = if let Some(split) = fastokens_fused_split(tokenizer) { + split.pre_tokenize(&mut pretokenized)?; + pretokenized + .tokenize_batched(|buffer, splits, output| { + tokenizer.model().tokenize_batch_fused(buffer, splits, output) + }) + .map_err(fastokens::Error::Model)? + } else { + if let Some(pre_tokenizer) = tokenizer.pre_tokenizer() { + pre_tokenizer.pre_tokenize(&mut pretokenized)?; + } + pretokenized + .tokenize(|text, output| tokenizer.model().tokenize_into(text, output)) + .map_err(fastokens::Error::Model)? + }; + + Ok(tokenizer.post_process(ids, false)) +} + /// Tokenizer from `tokenizer.json` in HuggingFace format. /// /// This tries to load with `fastokens` first for better performance, then falls @@ -156,6 +244,17 @@ impl Tokenizer for HuggingFaceTokenizer { } } + fn encode_ordinary(&self, text: &str) -> Result<Vec<u32>> { + match &self.backend { + Backend::Hf(tokenizer) => encode_hf_ordinary(tokenizer, text) + .map_err(|error| tokenizer_error!("encoding failed: {}", error.as_report())), + Backend::Fastokens(tokenizer) | Backend::FastokensByteLevel(tokenizer) => { + encode_fastokens_ordinary(tokenizer, text) + .map_err(|error| tokenizer_error!("encoding failed: {}", error.as_report())) + } + } + } + fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result<String> { match &self.backend { Backend::Hf(t) => t @@ -200,12 +299,19 @@ impl Tokenizer for HuggingFaceTokenizer { #[cfg(test)] mod tests { + use std::path::{Path, PathBuf}; + + use serde_json::{Value, json}; use tempfile::tempdir; use tokenizers::models::bpe::BPE; + use tokenizers::pre_tokenizers::byte_level::ByteLevel; use tokenizers::{AddedToken, Tokenizer as HfTokenizer}; use super::{HuggingFaceTokenizer, Tokenizer}; + const REGULAR_TOKEN: &str = "<|regular|>"; + const SPECIAL_TOKEN: &str = "<|special|>"; + fn tiny_bpe_tokenizer() -> HfTokenizer { let vocab = [ ("<unk>".to_string(), 0), @@ -232,6 +338,186 @@ mod tests { HfTokenizer::new(model) } + fn ordinary_test_tokenizer_json(fused: bool, with_added_tokens: bool) -> Value { + let mut alphabet: Vec<char> = ByteLevel::alphabet().into_iter().collect(); + alphabet.sort_unstable(); + let vocab = alphabet + .into_iter() + .enumerate() + .map(|(id, token)| (token.to_string(), json!(id))) + .collect::<serde_json::Map<_, _>>(); + + let pre_tokenizer = if fused { + json!({ + "type": "Sequence", + "pretokenizers": [ + { + "type": "Split", + "pattern": {"Regex": "\\S+|\\s+"}, + "behavior": "Isolated", + "invert": false + }, + { + "type": "ByteLevel", + "add_prefix_space": false, + "trim_offsets": true, + "use_regex": false + } + ] + }) + } else { + json!({ + "type": "ByteLevel", + "add_prefix_space": false, + "trim_offsets": true, + "use_regex": true + }) + }; + let added_tokens = with_added_tokens.then(|| { + json!([ + { + "id": 256, + "content": REGULAR_TOKEN, + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": true, + "special": false + }, + { + "id": 257, + "content": SPECIAL_TOKEN, + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + } + ]) + }); + + json!({ + "version": "1.0", + "truncation": { + "direction": "Right", + "max_length": 24, + "strategy": "LongestFirst", + "stride": 0 + }, + "padding": null, + "added_tokens": added_tokens.unwrap_or_else(|| json!([])), + "normalizer": {"type": "NFC"}, + "pre_tokenizer": pre_tokenizer, + "post_processor": { + "type": "ByteLevel", + "add_prefix_space": false, + "trim_offsets": true, + "use_regex": true + }, + "decoder": { + "type": "ByteLevel", + "add_prefix_space": false, + "trim_offsets": true, + "use_regex": true + }, + "model": { + "type": "BPE", + "dropout": null, + "unk_token": null, + "continuing_subword_prefix": null, + "end_of_word_suffix": null, + "fuse_unk": false, + "byte_fallback": false, + "ignore_merges": false, + "vocab": vocab, + "merges": [] + } + }) + } + + fn write_tokenizer_json(dir: &Path, name: &str, value: &Value) -> PathBuf { + let path = dir.join(name); + std::fs::write( + &path, + serde_json::to_vec(value).expect("serialize tokenizer"), + ) + .expect("write tokenizer"); + path + } + + fn assert_ordinary_matches_added_empty( + constructor: fn(&Path) -> crate::Result<HuggingFaceTokenizer>, + fused: bool, + ) { + let dir = tempdir().expect("create temp dir"); + let added_path = write_tokenizer_json( + dir.path(), + "with-added.json", + &ordinary_test_tokenizer_json(fused, true), + ); + let empty_path = write_tokenizer_json( + dir.path(), + "added-empty.json", + &ordinary_test_tokenizer_json(fused, false), + ); + let tokenizer = constructor(&added_path).expect("load tokenizer with added tokens"); + let added_empty = constructor(&empty_path).expect("load tokenizer with empty added tokens"); + + if let super::Backend::Fastokens(inner) | super::Backend::FastokensByteLevel(inner) = + &tokenizer.backend + { + assert_eq!(super::fastokens_fused_split(inner).is_some(), fused); + } + + assert_eq!( + tokenizer.encode(REGULAR_TOKEN, false).unwrap(), + vec![tokenizer.token_to_id(REGULAR_TOKEN).unwrap()] + ); + assert_eq!( + tokenizer.encode(SPECIAL_TOKEN, false).unwrap(), + vec![tokenizer.token_to_id(SPECIAL_TOKEN).unwrap()] + ); + + for text in [ + "", + "hello", + "Cafe\u{301}", + REGULAR_TOKEN, + SPECIAL_TOKEN, + "hello <|regular|> Cafe\u{301} <|special|> tail", + ] { + assert_eq!( + tokenizer.encode_ordinary(text).unwrap(), + added_empty.encode(text, false).unwrap(), + "fused={fused}, text={text:?}", + ); + } + if matches!(&tokenizer.backend, super::Backend::Hf(_)) { + assert_eq!( + tokenizer + .encode_ordinary("hello <|regular|> Cafe\u{301} <|special|> tail") + .unwrap() + .len(), + 24, + "HF post-processing must retain configured truncation", + ); + } + } + + #[test] + fn hf_ordinary_matches_original_encode_with_added_empty() { + for fused in [false, true] { + assert_ordinary_matches_added_empty(HuggingFaceTokenizer::new_hf, fused); + } + } + + #[test] + fn fastokens_ordinary_matches_original_encode_with_added_empty() { + for fused in [false, true] { + assert_ordinary_matches_added_empty(HuggingFaceTokenizer::new_fastokens, fused); + } + } + #[test] fn hf_constructor_resolves_added_token_ids() { let mut tokenizer = tiny_bpe_tokenizer(); diff --git a/rust/src/tokenizer/src/incremental.rs b/rust/src/tokenizer/src/incremental.rs index 5e470ae4b00..f608fa874b8 100644 --- a/rust/src/tokenizer/src/incremental.rs +++ b/rust/src/tokenizer/src/incremental.rs @@ -199,6 +199,10 @@ mod tests { unreachable!() } + fn encode_ordinary(&self, _text: &str) -> Result<Vec<u32>> { + unreachable!() + } + fn decode(&self, token_ids: &[u32], _skip_special_tokens: bool) -> Result<String> { let bytes = token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>(); Ok(String::from_utf8_lossy(&bytes).into_owned()) @@ -273,6 +277,10 @@ mod tests { unreachable!() } + fn encode_ordinary(&self, _text: &str) -> Result<Vec<u32>> { + unreachable!() + } + fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result<String> { let mut text = String::new(); for &token_id in token_ids { @@ -410,6 +418,10 @@ mod tests { unreachable!() } + fn encode_ordinary(&self, _text: &str) -> Result<Vec<u32>> { + unreachable!() + } + fn decode(&self, token_ids: &[u32], _skip_special_tokens: bool) -> Result<String> { match token_ids { [1] => Ok("abc".into()), diff --git a/rust/src/tokenizer/src/lib.rs b/rust/src/tokenizer/src/lib.rs index 6c9fcd3fdea..0f8c7dc16e5 100644 --- a/rust/src/tokenizer/src/lib.rs +++ b/rust/src/tokenizer/src/lib.rs @@ -25,6 +25,10 @@ pub trait Tokenizer: Send + Sync { /// Encode one prompt string into token IDs. fn encode(&self, text: &str, add_special_tokens: bool) -> Result<Vec<u32>>; + /// Equivalent to `encode(text, false)`, except that every added, + /// special, and control-token matcher is bypassed. + fn encode_ordinary(&self, text: &str) -> Result<Vec<u32>>; + /// Decode one token sequence into text. fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result<String>; diff --git a/rust/src/tokenizer/src/tekken.rs b/rust/src/tokenizer/src/tekken.rs index 20b6c26ffc8..5f9342e4a5e 100644 --- a/rust/src/tokenizer/src/tekken.rs +++ b/rust/src/tokenizer/src/tekken.rs @@ -35,6 +35,12 @@ impl Tokenizer for TekkenTokenizer { .map_err(|error| tokenizer_error!("encoding failed: {error}")) } + fn encode_ordinary(&self, text: &str) -> Result<Vec<u32>> { + self.inner + .encode(text, false, false) + .map_err(|error| tokenizer_error!("encoding failed: {error}")) + } + fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result<String> { let policy = if skip_special_tokens { tekken::SpecialTokenPolicy::Ignore @@ -67,3 +73,51 @@ impl Tokenizer for TekkenTokenizer { self.inner.is_special_token(token_id) } } + +#[cfg(test)] +mod tests { + use base64::Engine as _; + use tekken::config::TokenizerVersion; + use tekken::{SpecialTokenInfo, TokenInfo}; + + use super::*; + + fn test_tokenizer() -> TekkenTokenizer { + let vocab = (0_u8..=255) + .map(|byte| TokenInfo { + rank: byte as usize, + token_bytes: base64::engine::general_purpose::STANDARD.encode([byte]), + token_str: None, + }) + .collect(); + let special_tokens = vec![SpecialTokenInfo { + rank: 0, + token_str: "<control>".to_string(), + is_control: true, + }]; + let inner = Tekkenizer::new( + vocab, + &special_tokens, + r"(?s).", + 257, + 1, + TokenizerVersion::V3, + None, + ) + .expect("build Tekken tokenizer"); + TekkenTokenizer { inner } + } + + #[test] + fn ordinary_matches_tekkens_empty_special_encoding() { + let tokenizer = test_tokenizer(); + let text = "user <control> text"; + let control_id = tokenizer.token_to_id("<control>").unwrap(); + let ordinary_ids = tokenizer.encode_ordinary(text).unwrap(); + + assert_eq!(control_id, 0); + assert_eq!(ordinary_ids, tokenizer.encode(text, false).unwrap()); + assert!(!ordinary_ids.contains(&control_id)); + assert_eq!(tokenizer.decode(&ordinary_ids, false).unwrap(), text); + } +} diff --git a/rust/src/tokenizer/src/test_utils.rs b/rust/src/tokenizer/src/test_utils.rs index 36d1f3d18aa..6fdb6e02721 100644 --- a/rust/src/tokenizer/src/test_utils.rs +++ b/rust/src/tokenizer/src/test_utils.rs @@ -208,6 +208,10 @@ impl Tokenizer for TestTokenizer { Ok(ids) } + fn encode_ordinary(&self, text: &str) -> Result<Vec<u32>> { + Ok(text.as_bytes().iter().copied().map(u32::from).collect()) + } + fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result<String> { let mut output = String::new(); let mut pending_bytes = Vec::new(); @@ -374,6 +378,32 @@ mod tests { assert!(!tokenizer.is_special_id(0xF002)); } + #[test] + fn ordinary_encoding_bypasses_all_configured_tokens() { + let tokenizer = TestTokenizer::new() + .with_bos_token("<bos>", 256) + .with_special_token("<control>", 257) + .with_regular_token("<visible>", 258); + let ordinary_text = "user <control> and <visible>"; + + assert_eq!(tokenizer.encode("<control>", false).unwrap(), vec![257]); + assert_eq!(tokenizer.encode("<visible>", false).unwrap(), vec![258]); + assert_eq!( + tokenizer.encode_ordinary(ordinary_text).unwrap(), + ordinary_text.as_bytes().iter().copied().map(u32::from).collect::<Vec<_>>() + ); + + let mut segmented = tokenizer.encode("<control>", false).unwrap(); + segmented.extend(tokenizer.encode_ordinary(ordinary_text).unwrap()); + segmented.extend(tokenizer.encode("<visible>", false).unwrap()); + assert_eq!(segmented.first(), Some(&257)); + assert_eq!(segmented.last(), Some(&258)); + assert_eq!( + tokenizer.decode(&segmented, false).unwrap(), + format!("<control>{ordinary_text}<visible>") + ); + } + #[test] #[should_panic(expected = "configured test token id 255 overlaps byte fallback range 0..=255")] fn configured_token_id_must_stay_outside_byte_range() { diff --git a/rust/src/tokenizer/src/tiktoken.rs b/rust/src/tokenizer/src/tiktoken.rs index d5ca119b808..0f355566b94 100644 --- a/rust/src/tokenizer/src/tiktoken.rs +++ b/rust/src/tokenizer/src/tiktoken.rs @@ -462,6 +462,13 @@ impl Tokenizer for TiktokenTokenizer { }) } + fn encode_ordinary(&self, text: &str) -> Result<Vec<u32>> { + Ok(match &self.backend { + Backend::Riptoken(backend) => backend.inner.encode_ordinary(text), + Backend::TiktokenRs(backend) => backend.inner.encode_ordinary(text), + }) + } + fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result<String> { // Filter passes: // @@ -752,6 +759,39 @@ mod tests { } } + #[test] + fn tiktoken_ordinary_bypasses_every_registered_added_token() { + let dir = tempfile::tempdir().expect("create temp dir"); + let bpe_path = write_synthetic_bpe_file(dir.path()); + fs::write( + dir.path().join("tokenizer_config.json"), + r#"{ + "added_tokens_decoder": { + "257": { "content": "<|im_end|>", "special": true }, + "258": { "content": "<|tool_call_begin|>", "special": false } + } + }"#, + ) + .expect("write tokenizer_config.json"); + fs::write(dir.path().join("config.json"), r#"{"vocab_size": 260}"#) + .expect("write config.json"); + + let input = "<|im_end|><|tool_call_begin|><|reserved_token_259|>"; + let expected: Vec<u32> = input.as_bytes().iter().copied().map(u32::from).collect(); + for backend in explicit_backends(&bpe_path) { + assert_eq!(backend.encode("<|im_end|>", false).unwrap(), vec![257]); + assert_eq!( + backend.encode("<|tool_call_begin|>", false).unwrap(), + vec![258] + ); + assert_eq!( + backend.encode("<|reserved_token_259|>", false).unwrap(), + vec![259] + ); + assert_eq!(backend.encode_ordinary(input).unwrap(), expected); + } + } + /// `vocab_size` may live under `text_config` for composite (e.g. /// multimodal) configs. #[test] From 88402a41c4ab272ebbbd33f4a77fbbac0431cbb9 Mon Sep 17 00:00:00 2001 From: Liangliang Ma <liangliang.ma@intel.com> Date: Tue, 28 Jul 2026 16:26:49 +0800 Subject: [PATCH 160/185] [Test] Skip ROCm AITER MLA prefill tests on non-ROCm platforms (#49945) Signed-off-by: Liangliang-Ma <liangliang.ma@intel.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/attention/test_mla_prefill_selector.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py index f5985e7bc8e..e6d9f939ea5 100644 --- a/tests/v1/attention/test_mla_prefill_selector.py +++ b/tests/v1/attention/test_mla_prefill_selector.py @@ -8,6 +8,7 @@ import pytest import torch from vllm.config import AttentionConfig, ModelConfig, VllmConfig +from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability from vllm.v1.attention.backends.mla.prefill.base import MLADimensions from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum @@ -287,6 +288,11 @@ class TestBackendValidation: assert invalid_reasons == [] +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="Imports vllm.platforms.rocm, whose module init requires a CUDA or " + "ROCm torch build; not importable on XPU/CPU/TPU.", +) class TestROCmAiterFAPrefillSelection: """Tests for the ROCm AITER FlashAttention MLA prefill backend.""" From 247470f23a4c8191a296c1993c4e86727ffb5191 Mon Sep 17 00:00:00 2001 From: Chris Leonard <chleonar@redhat.com> Date: Tue, 28 Jul 2026 05:37:01 -0400 Subject: [PATCH 161/185] [CI] Add PyTorch stable ABI audit check (#48164) Signed-off-by: Chris Leonard <chleonar@redhat.com> Co-authored-by: Shengqi Chen <harry-chen@outlook.com> --- .buildkite/check-torch-abi.py | 102 +++++++++++++++++++++++++++ .buildkite/ci_config.yaml | 1 + .buildkite/test_areas/torch_abi.yaml | 14 ++++ requirements/test/cpu.txt | 6 ++ requirements/test/cuda.in | 2 +- requirements/test/cuda.txt | 6 ++ 6 files changed, 130 insertions(+), 1 deletion(-) create mode 100644 .buildkite/check-torch-abi.py create mode 100644 .buildkite/test_areas/torch_abi.yaml diff --git a/.buildkite/check-torch-abi.py b/.buildkite/check-torch-abi.py new file mode 100644 index 00000000000..493952c33ec --- /dev/null +++ b/.buildkite/check-torch-abi.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Audit vLLM compiled libraries for PyTorch stable ABI compliance.""" + +import fnmatch +import sys +from pathlib import Path + +from torch_abi_audit import inspect_package +from torch_abi_audit.report import ExtensionReport, PackageReport + +# Temporary allowlist of extensions not yet on the stable ABI. +# Shrink and remove over time. +ALLOWED_UNSTABLE_LIBRARIES: tuple[str, ...] = ( + "vllm_flash_attn/_vllm_fa2_C.abi3.so", + "vllm_flash_attn/_vllm_fa3_C.abi3.so", + "third_party/deep_gemm/_C*.so", +) + + +def _relative_path(lib: ExtensionReport, package_root: Path) -> str: + try: + return lib.path.relative_to(package_root).as_posix() + except ValueError: + return lib.path.name + + +def _is_torch_unstable(lib: ExtensionReport) -> bool: + return lib.error is None and lib.torch.uses_torch and not lib.torch.stable + + +def _matches_allowlist(rel_path: str, patterns: tuple[str, ...]) -> bool: + return any(fnmatch.fnmatch(rel_path, pattern) for pattern in patterns) + + +def _iter_libs(report: PackageReport) -> tuple[ExtensionReport, ...]: + return (*report.extensions, *report.bundled_libs) + + +def _collect_unstable(report: PackageReport) -> list[str]: + return sorted( + _relative_path(lib, report.root) + for lib in _iter_libs(report) + if _is_torch_unstable(lib) + ) + + +def _find_stale_allowlist_entries( + report: PackageReport, patterns: tuple[str, ...] +) -> list[str]: + """Allowlist patterns that match a built library which is no longer unstable.""" + stale: list[str] = [] + for pattern in patterns: + for lib in _iter_libs(report): + if lib.error is not None: + continue + if not fnmatch.fnmatch(_relative_path(lib, report.root), pattern): + continue + if not _is_torch_unstable(lib): + stale.append(pattern) + break + return stale + + +def check_torch_abi( + package: str = "vllm", + patterns: tuple[str, ...] = ALLOWED_UNSTABLE_LIBRARIES, +) -> int: + report = inspect_package(package) + if report.error: + print(f"error: failed to inspect {package!r}: {report.error}", file=sys.stderr) + return 2 + + unstable = _collect_unstable(report) + unexpected = [ + rel_path for rel_path in unstable if not _matches_allowlist(rel_path, patterns) + ] + stale = _find_stale_allowlist_entries(report, patterns) + + if unexpected or stale: + if unexpected: + print( + "Not allowed: torch-unstable libraries outside " + f"ALLOWED_UNSTABLE_LIBRARIES: {', '.join(unexpected)}", + file=sys.stderr, + ) + if stale: + print( + "Not allowed: stale ALLOWED_UNSTABLE_LIBRARIES entries: " + f"{', '.join(stale)}", + file=sys.stderr, + ) + return 1 + + print("Torch stable ABI check passed.") + return 0 + + +if __name__ == "__main__": + print(">>> Auditing vLLM extension modules for PyTorch stable ABI compliance") + sys.exit(check_torch_abi()) diff --git a/.buildkite/ci_config.yaml b/.buildkite/ci_config.yaml index 21ffa1b9b8d..9e1e46db67e 100644 --- a/.buildkite/ci_config.yaml +++ b/.buildkite/ci_config.yaml @@ -14,6 +14,7 @@ run_all_patterns: - "setup.py" - "csrc/" - "cmake/" + - ".buildkite/check-torch-abi.py" run_all_exclude_patterns: - "docker/Dockerfile." - "csrc/cpu/" diff --git a/.buildkite/test_areas/torch_abi.yaml b/.buildkite/test_areas/torch_abi.yaml new file mode 100644 index 00000000000..eaef3551664 --- /dev/null +++ b/.buildkite/test_areas/torch_abi.yaml @@ -0,0 +1,14 @@ +group: Torch ABI +depends_on: + - image-build +steps: +- label: Torch Stable ABI Audit + key: torch-stable-abi-audit + timeout_in_minutes: 5 + source_file_dependencies: + - .buildkite/check-torch-abi.py + - csrc/ + - cmake/ + - setup.py + commands: + - python3 /vllm-workspace/.buildkite/check-torch-abi.py diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt index 3cb251a308b..923b54bb14e 100644 --- a/requirements/test/cpu.txt +++ b/requirements/test/cpu.txt @@ -1,5 +1,7 @@ # This file was autogenerated by uv via the following command: # uv pip compile requirements/test/cuda.in -o requirements/test/cpu.txt --index-strategy unsafe-best-match --torch-backend cpu --python-platform x86_64-manylinux_2_28 --python-version 3.12 +abi3info==2025.11.29 + # via torch-abi-audit absl-py==2.1.0 # via rouge-score accelerate==1.13.0 @@ -763,6 +765,8 @@ pycparser==2.22 # via cffi pycryptodomex==3.22.0 # via blobfile +pycxxfilt==0.1.0 + # via torch-abi-audit pydantic==2.12.0 # via # -r requirements/test/../common.txt @@ -1127,6 +1131,8 @@ torch==2.13.0+cpu # vector-quantize-pytorch # vocos # xgrammar +torch-abi-audit==0.0.1 + # via -r requirements/test/cuda.in torchaudio==2.11.0+cpu # via # -r requirements/test/cuda.in diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index b33257250e7..377224eac36 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -45,7 +45,7 @@ schemathesis>=4.0.0 # Required for openai schema test. # quantization bitsandbytes==0.49.2 buildkite-test-collector==0.1.9 - +torch-abi-audit # CI check for PyTorch stable ABI compliance genai_perf>=0.0.8 tritonclient>=2.51.0 diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 6490502cdda..6c155d89bd2 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -1,5 +1,7 @@ # This file was autogenerated by uv via the following command: # uv pip compile requirements/test/cuda.in -c requirements/cuda.txt -o requirements/test/cuda.txt --index-strategy unsafe-best-match --torch-backend cu130 --python-platform x86_64-manylinux_2_28 --python-version 3.12 +abi3info==2025.11.29 + # via torch-abi-audit absl-py==2.1.0 # via rouge-score accelerate==1.13.0 @@ -850,6 +852,8 @@ pycparser==2.22 # via cffi pycryptodomex==3.22.0 # via blobfile +pycxxfilt==0.1.0 + # via torch-abi-audit pydantic==2.12.0 # via # -c requirements/common.txt @@ -1225,6 +1229,8 @@ torch==2.13.0+cu130 # vector-quantize-pytorch # vocos # xgrammar +torch-abi-audit==0.0.1 + # via -r requirements/test/cuda.in torchaudio==2.11.0+cu130 # via # -c requirements/cuda.txt From 25ace8fe5df07fc13f4aef5a89db391f326e60ee Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu <riverclouds.zhu@qq.com> Date: Tue, 28 Jul 2026 18:04:36 +0800 Subject: [PATCH 162/185] [CI] Increase Qwen3.5 MTP GSM8K generation length (#49881) Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../evals/gsm8k/configs/Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/evals/gsm8k/configs/Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml b/tests/evals/gsm8k/configs/Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml index d247515a0f0..921365ae686 100644 --- a/tests/evals/gsm8k/configs/Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml +++ b/tests/evals/gsm8k/configs/Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml @@ -3,8 +3,9 @@ accuracy_threshold: 0.88 tolerance: 0.03 num_questions: 1319 num_fewshot: 5 +max_tokens: 12000 server_args: >- - --max-model-len 4096 + --max-model-len 16384 --data-parallel-size 2 --enable-expert-parallel --max-num-seqs 384 From bf9f23003c67698b973a674b54e4ca77aae1adff Mon Sep 17 00:00:00 2001 From: Reid <61492567+reidliu41@users.noreply.github.com> Date: Tue, 28 Jul 2026 18:53:18 +0800 Subject: [PATCH 163/185] [Rust Frontend] Fix finish reason for named tool choices (#49496) Signed-off-by: reidliu41 <reid201711@gmail.com> --- .../src/routes/openai/chat_completions.rs | 16 ++-- .../routes/openai/chat_completions/convert.rs | 6 ++ rust/src/server/src/routes/tests.rs | 83 ++++++++++++++++--- 3 files changed, 88 insertions(+), 17 deletions(-) diff --git a/rust/src/server/src/routes/openai/chat_completions.rs b/rust/src/server/src/routes/openai/chat_completions.rs index bd86a06136a..8a1e9c4383a 100644 --- a/rust/src/server/src/routes/openai/chat_completions.rs +++ b/rust/src/server/src/routes/openai/chat_completions.rs @@ -131,6 +131,7 @@ async fn collect_chat_completion( echo, return_token_ids, return_tokens_as_token_ids, + is_named_tool_choice, }: ResponseOptions, ) -> Result<ChatCompletionResponse, ApiError> { let collected = stream.collect_message().await.map_err(|error| { @@ -157,7 +158,9 @@ async fn collect_chat_completion( // When reasoning is hidden, omit them rather than leaking hidden reasoning // tokens through per-token metadata. let include_output_metadata = include_reasoning || reasoning.is_none(); - let finish_reason = chat_finish_reason_to_openai(&finish_reason, saw_tool_calls)?.to_string(); + let finish_reason = + chat_finish_reason_to_openai(&finish_reason, saw_tool_calls && !is_named_tool_choice)? + .to_string(); let tool_calls = message .tool_calls() .map(|call| ToolCall { @@ -254,6 +257,7 @@ async fn chat_completion_chunk_stream( echo, return_token_ids, return_tokens_as_token_ids, + is_named_tool_choice, }: ResponseOptions, mut y: TryYielder<ChatCompletionStreamResponse, ApiError>, ) -> Result<(), ApiError> { @@ -454,7 +458,7 @@ async fn chat_completion_chunk_stream( &response_model, created, finish_reason, - saw_tool_calls, + saw_tool_calls && !is_named_tool_choice, ) { Ok(chunk) => yield_chunk!(chunk), Err(error) => { @@ -787,10 +791,10 @@ fn final_chunk( response_model: &str, created: u64, finish_reason: FinishReason, - saw_tool_calls: bool, + use_tool_calls_finish_reason: bool, ) -> Result<ChatCompletionStreamResponse, ApiError> { let stop_reason = finish_reason.as_stop_reason().map(stop_reason_to_json); - let finish_reason = chat_finish_reason_to_openai(&finish_reason, saw_tool_calls)?; + let finish_reason = chat_finish_reason_to_openai(&finish_reason, use_tool_calls_finish_reason)?; debug!( finish_reason = %finish_reason, @@ -809,10 +813,10 @@ fn final_chunk( fn chat_finish_reason_to_openai( finish_reason: &FinishReason, - saw_tool_calls: bool, + use_tool_calls_finish_reason: bool, ) -> Result<&'static str, ApiError> { match finish_reason { - FinishReason::Stop(_) if saw_tool_calls => Ok("tool_calls"), + FinishReason::Stop(_) if use_tool_calls_finish_reason => Ok("tool_calls"), FinishReason::Stop(_) => Ok("stop"), FinishReason::Length => Ok("length"), FinishReason::Abort => Ok("abort"), diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index 5b2fa3c19ed..d75f3fab162 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -54,6 +54,8 @@ pub(super) struct ResponseOptions { pub return_token_ids: bool, /// Whether to format logprob tokens as `token_id:{id}`. pub return_tokens_as_token_ids: bool, + /// Whether the request forces one named function tool. + pub is_named_tool_choice: bool, } /// Validate and lower one OpenAI chat completion request into the internal chat @@ -98,6 +100,7 @@ pub(super) fn prepare_chat_request( .and_then(|options| options.continuous_usage_stats) .unwrap_or(false); let requested_logprobs = request.logprobs; + let is_named_tool_choice = matches!(&request.tool_choice, Some(ToolChoice::Function { .. })); // Auto-enable prompt logprobs for non-streaming echo, matching Python vLLM's // behavior. @@ -180,6 +183,7 @@ pub(super) fn prepare_chat_request( echo, return_token_ids: request.return_token_ids.unwrap_or(false), return_tokens_as_token_ids: request.return_tokens_as_token_ids.unwrap_or(false), + is_named_tool_choice, }, chat_request, }) @@ -1068,6 +1072,7 @@ mod tests { .expect("request is valid"); assert_eq!(prepared.chat_request.tool_choice, ChatToolChoice::Required); + assert!(!prepared.options.is_named_tool_choice); } #[test] @@ -1107,6 +1112,7 @@ mod tests { name: "get_weather".to_string(), } ); + assert!(prepared.options.is_named_tool_choice); } #[test] diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 64e904b562b..1c228f99f1f 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -153,6 +153,20 @@ fn default_stream_output_specs() -> Vec<(Vec<u32>, Option<EngineCoreFinishReason ] } +fn weather_tool_call_output_specs() -> Vec<(Vec<u32>, Option<EngineCoreFinishReason>)> { + vec![ + (bytes_to_token_ids(b"<think>Need tool.</think>"), None), + ( + bytes_to_token_ids(b"<tool_call>\n{\"name\":\"get_weather\", "), + None, + ), + ( + bytes_to_token_ids(b"\"arguments\":{\"city\":\"Paris\"}}\n</tool_call>"), + Some(EngineCoreFinishReason::Stop), + ), + ] +} + fn assert_adapter_a_lora_request(request: &EngineCoreRequest) { let lora = request.lora_request.as_ref().expect("lora request"); assert_eq!(lora.lora_name, "adapter-a"); @@ -4554,17 +4568,7 @@ async fn include_reasoning_false_suppresses_non_stream_output_metadata() { async fn tool_calls_are_mapped_to_tool_call_sse_chunks() { let (app, engine_task) = test_app_with_backend_and_stream_output_specs( Arc::new(FakeChatBackend::with_model_id("Qwen/Qwen3-0.6B")), - vec![ - (bytes_to_token_ids(b"<think>Need tool.</think>"), None), - ( - bytes_to_token_ids(b"<tool_call>\n{\"name\":\"get_weather\", "), - None, - ), - ( - bytes_to_token_ids(b"\"arguments\":{\"city\":\"Paris\"}}\n</tool_call>"), - Some(EngineCoreFinishReason::Stop), - ), - ], + weather_tool_call_output_specs(), ) .await; @@ -4613,6 +4617,63 @@ async fn tool_calls_are_mapped_to_tool_call_sse_chunks() { assert!(text.contains("\"finish_reason\":\"tool_calls\""), "{text}"); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn named_tool_choice_uses_stop_finish_reason() { + for stream in [false, true] { + let (app, engine_task) = test_app_with_backend_and_stream_output_specs( + Arc::new(FakeChatBackend::with_model_id("Qwen/Qwen3-0.6B")), + weather_tool_call_output_specs(), + ) + .await; + + let response = app + .clone() + .call( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "stream": stream, + "messages": [{"role": "user", "content": "hello"}], + "tools": [{ + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}} + } + } + }], + "tool_choice": { + "type": "function", + "function": {"name": "get_weather"} + } + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + engine_task.await.expect("mock engine task"); + let text = String::from_utf8(body.to_vec()).expect("utf8 body"); + + assert!(text.contains("\"tool_calls\":"), "{text}"); + assert!(text.contains("\"name\":\"get_weather\""), "{text}"); + assert!(text.contains("\"finish_reason\":\"stop\""), "{text}"); + assert!(!text.contains("\"finish_reason\":\"tool_calls\""), "{text}"); + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn tool_call_sse_chunks_can_carry_logprobs() { From 912d6b619de0f2a44df74704973cff90fa0fb4e9 Mon Sep 17 00:00:00 2001 From: Reid <61492567+reidliu41@users.noreply.github.com> Date: Tue, 28 Jul 2026 20:02:00 +0800 Subject: [PATCH 164/185] [Rust Frontend] Align sampling validation with Python (#47494) Signed-off-by: reidliu41 <reid201711@gmail.com> --- rust/src/server/src/error.rs | 16 ++++ rust/src/server/src/grpc/tests.rs | 29 +++++++ rust/src/text/src/error.rs | 4 + rust/src/text/src/lib.rs | 2 +- rust/src/text/src/lower.rs | 119 +++++++++++++++++++++++++++- rust/src/text/src/lower/sampling.rs | 92 +++++++++++++++++++++ 6 files changed, 260 insertions(+), 2 deletions(-) create mode 100644 rust/src/text/src/lower/sampling.rs diff --git a/rust/src/server/src/error.rs b/rust/src/server/src/error.rs index fbd13b77baf..3a8c0ae5c92 100644 --- a/rust/src/server/src/error.rs +++ b/rust/src/server/src/error.rs @@ -141,6 +141,22 @@ mod tests { assert!(response.error.message.contains("max_tokens=4")); } + #[test] + fn sampling_params_validation_maps_to_invalid_request() { + let api_error = text_submit_error( + "failed to submit completion request", + vllm_text::Error::SamplingParams(vllm_text::SamplingParamsError::OutOfRange { + parameter: "top_p", + value: 0.0, + expected: "(0, 1]", + }), + ); + assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST); + let response = api_error.to_error_response(); + assert_eq!(response.error.error_type, "invalid_request_error"); + assert!(response.error.message.contains("top_p")); + } + #[test] fn chat_wrapped_prompt_too_long_maps_to_invalid_request() { let error = vllm_chat::Error::Text(vllm_text::Error::PromptTooLong { diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 75da670ec3b..808f3422758 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -648,6 +648,35 @@ async fn unary_generate_min_tokens_above_max_tokens_returns_invalid_argument() { server_task.abort(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn unary_generate_invalid_sampling_params_returns_invalid_argument() { + let (mut client, server_task, _engine_task) = grpc_test_server( + b"engine-grpc-invalid-sampling", + default_stream_output_specs(), + ) + .await; + + let status = client + .generate(pb::GenerateRequest { + request_id: "test-invalid-sampling".to_string(), + model: "test-model".to_string(), + prompt: Some(pb::generate_request::Prompt::Text("hi".to_string())), + sampling: Some(pb::RandomSampling { + top_p: 2.0, + ..Default::default() + }), + ..Default::default() + }) + .await + .expect_err("should fail when top_p is out of range"); + + assert_eq!(status.code(), tonic::Code::InvalidArgument); + assert!(status.message().contains("top_p")); + + server_task.abort(); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn streaming_generate_yields_incremental_responses() { diff --git a/rust/src/text/src/error.rs b/rust/src/text/src/error.rs index e72e68196c3..6b385fad68d 100644 --- a/rust/src/text/src/error.rs +++ b/rust/src/text/src/error.rs @@ -6,6 +6,7 @@ use vllm_engine_core_client::Error as EngineCoreError; use vllm_llm::Error as LlmError; pub use crate::lower::logprobs::LogprobsError; +pub use crate::lower::sampling::SamplingParamsError; pub use crate::lower::token_ids::TokenIdsError; #[derive(Debug, Error)] @@ -23,6 +24,8 @@ pub enum Error { Logprobs(#[from] LogprobsError), #[error(transparent)] TokenIds(#[from] TokenIdsError), + #[error(transparent)] + SamplingParams(#[from] SamplingParamsError), #[error( "`min_tokens` must be less than or equal to `max_tokens`, \ got min_tokens={min_tokens}, max_tokens={max_tokens}" @@ -50,6 +53,7 @@ impl Error { | Self::EmptyPromptTokenIds { .. } | Self::Logprobs(_) | Self::TokenIds(_) + | Self::SamplingParams(_) | Self::MinTokensExceedsMaxTokens { .. } | Self::InvalidThinkingTokenBudget | Self::InvalidRepetitionDetection { .. } diff --git a/rust/src/text/src/lib.rs b/rust/src/text/src/lib.rs index b00155999e8..646e4bac9f3 100644 --- a/rust/src/text/src/lib.rs +++ b/rust/src/text/src/lib.rs @@ -10,7 +10,7 @@ use std::mem::take; pub use backend::{DynTextBackend, SamplingHints, SamplingLimits, TextBackend}; -pub use error::{Error, LogprobsError, Result, TokenIdsError}; +pub use error::{Error, LogprobsError, Result, SamplingParamsError, TokenIdsError}; use futures::Stream; pub use lower::{ PreparedTextRequest, lower_sampling_params, lower_text_request, resolve_max_tokens, diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index bd43a1d141c..aa54595acda 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -4,9 +4,11 @@ use std::collections::BTreeSet; pub(crate) mod logprobs; +pub(crate) mod sampling; pub(crate) mod token_ids; use logprobs::validate_logprobs; +use sampling::validate_resolved_sampling_params; use token_ids::{validate_prompt_token_ids, validate_vocab_range}; use vllm_engine_core_client::protocol::sampling::{ EngineCoreSamplingParams, RepetitionDetectionParams, @@ -186,6 +188,7 @@ pub fn lower_sampling_params( skip_reading_prefix_cache, extra_args: vllm_xargs, }; + validate_resolved_sampling_params(¶ms)?; validate_vocab_range(¶ms, &sampling_limits)?; Ok(params) } @@ -319,7 +322,7 @@ mod tests { use super::*; use crate::backend::hf::HfTextBackend; use crate::backend::{SamplingHints, TextBackend as _}; - use crate::error::{LogprobsError, TokenIdsError}; + use crate::error::{LogprobsError, SamplingParamsError, TokenIdsError}; use crate::request::{Prompt, TextRequest}; fn stub_tokenizer() -> TestTokenizer { @@ -482,6 +485,120 @@ mod tests { assert!(message.contains("min_count=1")); } + #[test] + fn lower_sampling_params_rejects_invalid_sampling_ranges() { + let cases = [ + ( + "temperature", + SamplingParams { + temperature: Some(5.0), + ..SamplingParams::default() + }, + ), + ( + "top_p", + SamplingParams { + top_p: Some(0.0), + ..SamplingParams::default() + }, + ), + ( + "min_p", + SamplingParams { + min_p: Some(2.0), + ..SamplingParams::default() + }, + ), + ( + "repetition_penalty", + SamplingParams { + repetition_penalty: Some(0.0), + ..SamplingParams::default() + }, + ), + ( + "frequency_penalty", + SamplingParams { + frequency_penalty: Some(100.0), + ..SamplingParams::default() + }, + ), + ( + "presence_penalty", + SamplingParams { + presence_penalty: Some(100.0), + ..SamplingParams::default() + }, + ), + ]; + + for (expected_parameter, sampling_params) in cases { + let error = + lower_sampling_params_with_limits(sampling_params, sample_sampling_limits()) + .unwrap_err(); + + assert!( + matches!( + error, + Error::SamplingParams(SamplingParamsError::OutOfRange { + parameter, + .. + }) if parameter == expected_parameter + ), + "{expected_parameter} should be rejected" + ); + } + } + + #[test] + fn lower_sampling_params_rejects_non_finite_sampling_values() { + for (expected_parameter, sampling_params) in [ + ( + "temperature", + SamplingParams { + temperature: Some(f32::INFINITY), + ..SamplingParams::default() + }, + ), + ( + "repetition_penalty", + SamplingParams { + repetition_penalty: Some(f32::NAN), + ..SamplingParams::default() + }, + ), + ] { + let error = + lower_sampling_params_with_limits(sampling_params, sample_sampling_limits()) + .unwrap_err(); + + assert!( + matches!( + error, + Error::SamplingParams(SamplingParamsError::NotFinite { + parameter, + .. + }) if parameter == expected_parameter + ), + "{expected_parameter} should reject non-finite values" + ); + } + } + + #[test] + fn lower_sampling_params_accepts_python_compatible_repetition_penalty_above_two() { + let params = lower_sampling_params_with_limits( + SamplingParams { + repetition_penalty: Some(2.5), + ..SamplingParams::default() + }, + sample_sampling_limits(), + ) + .unwrap(); + + assert_eq!(params.repetition_penalty, 2.5); + } + #[test] fn lower_text_request_applies_python_style_eos_hints() { let prepared = lower_text_request( diff --git a/rust/src/text/src/lower/sampling.rs b/rust/src/text/src/lower/sampling.rs new file mode 100644 index 00000000000..edcdc4b0492 --- /dev/null +++ b/rust/src/text/src/lower/sampling.rs @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use thiserror::Error; +use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; + +#[derive(Debug, Error, PartialEq)] +pub enum SamplingParamsError { + #[error("{parameter} must be a finite number, got {value}")] + NotFinite { parameter: &'static str, value: f32 }, + #[error("{parameter} must be in {expected}, got {value}")] + OutOfRange { + parameter: &'static str, + value: f32, + expected: &'static str, + }, +} + +fn validate_frequency_penalty(value: f32) -> Result<(), SamplingParamsError> { + validate_closed_range("frequency_penalty", value, -2.0, 2.0, "[-2, 2]") +} + +fn validate_presence_penalty(value: f32) -> Result<(), SamplingParamsError> { + validate_closed_range("presence_penalty", value, -2.0, 2.0, "[-2, 2]") +} + +fn validate_temperature(value: f32) -> Result<(), SamplingParamsError> { + validate_finite("temperature", value)?; + validate_closed_range("temperature", value, 0.0, 2.0, "[0, 2]") +} + +fn validate_top_p(value: f32) -> Result<(), SamplingParamsError> { + if value > 0.0 && value <= 1.0 { + return Ok(()); + } + Err(SamplingParamsError::OutOfRange { + parameter: "top_p", + value, + expected: "(0, 1]", + }) +} + +fn validate_min_p(value: f32) -> Result<(), SamplingParamsError> { + validate_closed_range("min_p", value, 0.0, 1.0, "[0, 1]") +} + +fn validate_repetition_penalty(value: f32) -> Result<(), SamplingParamsError> { + validate_finite("repetition_penalty", value)?; + if value > 0.0 { + return Ok(()); + } + Err(SamplingParamsError::OutOfRange { + parameter: "repetition_penalty", + value, + expected: "(0, inf)", + }) +} + +pub(crate) fn validate_resolved_sampling_params( + params: &EngineCoreSamplingParams, +) -> Result<(), SamplingParamsError> { + validate_temperature(params.temperature)?; + validate_top_p(params.top_p)?; + validate_min_p(params.min_p)?; + validate_frequency_penalty(params.frequency_penalty)?; + validate_presence_penalty(params.presence_penalty)?; + validate_repetition_penalty(params.repetition_penalty) +} + +fn validate_finite(parameter: &'static str, value: f32) -> Result<(), SamplingParamsError> { + if value.is_finite() { + return Ok(()); + } + Err(SamplingParamsError::NotFinite { parameter, value }) +} + +fn validate_closed_range( + parameter: &'static str, + value: f32, + min: f32, + max: f32, + expected: &'static str, +) -> Result<(), SamplingParamsError> { + if value >= min && value <= max { + return Ok(()); + } + Err(SamplingParamsError::OutOfRange { + parameter, + value, + expected, + }) +} From d2bfc6fe20343c638840c8867c29f7365fe23378 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" <khluu000@gmail.com> Date: Tue, 28 Jul 2026 05:09:19 -0700 Subject: [PATCH 165/185] [Build] Fix DeepEP CUDA driver stub linking (#50103) Signed-off-by: khluu <khluu000@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com> --- tools/ep_kernels/install_python_libraries.sh | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tools/ep_kernels/install_python_libraries.sh b/tools/ep_kernels/install_python_libraries.sh index 739f031c9ef..5f5a597baca 100755 --- a/tools/ep_kernels/install_python_libraries.sh +++ b/tools/ep_kernels/install_python_libraries.sh @@ -197,6 +197,21 @@ do_build() { #endif' csrc/kernels/backend/symmetric.hpp fi + if [[ "$name" == "DeepEP" ]]; then + # DeepEP links against the CUDA driver API in driverless build images. + local cuda_driver_stub + local cuda_driver_stub_dir + cuda_driver_stub=$( + find -H "$CUDA_HOME" -path "*/stubs/libcuda.so" -print -quit + ) + if [[ -z "$cuda_driver_stub" ]]; then + echo "CUDA driver stub not found under $CUDA_HOME" >&2 + exit 1 + fi + cuda_driver_stub_dir=$(dirname "$cuda_driver_stub") + export LIBRARY_PATH="${cuda_driver_stub_dir}${LIBRARY_PATH:+:$LIBRARY_PATH}" + fi + if [ "$MODE" = "install" ]; then echo "Installing $name into environment" eval "$extra_env" uv pip install --no-build-isolation -vvv . From 35efdf6b34f1ef76c23c8980f6e8a3b73cab50f1 Mon Sep 17 00:00:00 2001 From: Itay Alroy <75032521+itayalroy@users.noreply.github.com> Date: Tue, 28 Jul 2026 15:19:13 +0300 Subject: [PATCH 166/185] [Elastic EP] Async preparation (#47288) Signed-off-by: Itay Alroy <ialroy@nvidia.com> --- .buildkite/test_areas/expert_parallelism.yaml | 1 + tests/distributed/test_elastic_ep.py | 166 +++++-- vllm/config/parallel.py | 9 + .../base_device_communicator.py | 16 +- .../device_communicators/cpu_communicator.py | 5 +- .../device_communicators/cuda_communicator.py | 2 + .../device_communicators/xpu_communicator.py | 5 +- .../distributed/elastic_ep/elastic_execute.py | 181 +++---- vllm/distributed/elastic_ep/elastic_state.py | 460 ++++++------------ vllm/distributed/elastic_ep/standby_state.py | 3 +- vllm/distributed/parallel_state.py | 14 +- vllm/distributed/stateless_coordinator.py | 2 + .../serve/elastic_ep/api_router.py | 9 +- vllm/model_executor/warmup/kernel_warmup.py | 31 +- vllm/v1/engine/__init__.py | 2 - vllm/v1/engine/async_llm.py | 45 +- vllm/v1/engine/coordinator.py | 3 + vllm/v1/engine/core.py | 75 ++- vllm/v1/engine/core_client.py | 310 ++++++------ vllm/v1/engine/utils.py | 31 +- vllm/v1/executor/abstract.py | 8 +- 21 files changed, 701 insertions(+), 677 deletions(-) diff --git a/.buildkite/test_areas/expert_parallelism.yaml b/.buildkite/test_areas/expert_parallelism.yaml index a3b46b58285..1d1609d46b8 100644 --- a/.buildkite/test_areas/expert_parallelism.yaml +++ b/.buildkite/test_areas/expert_parallelism.yaml @@ -52,4 +52,5 @@ steps: - vllm/compilation/ - tests/distributed/ commands: + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - pytest -v -s distributed/test_elastic_ep.py diff --git a/tests/distributed/test_elastic_ep.py b/tests/distributed/test_elastic_ep.py index 4ce7497598a..01c254c2d03 100644 --- a/tests/distributed/test_elastic_ep.py +++ b/tests/distributed/test_elastic_ep.py @@ -3,7 +3,9 @@ import os import subprocess +import threading import time +from concurrent.futures import ThreadPoolExecutor import pytest import requests @@ -40,6 +42,106 @@ def _send_scale_command(server: RemoteOpenAIServer, new_dp_size: int) -> bool: return False +def _traffic_loop( + server: RemoteOpenAIServer, + dp_rank: int | None, + ready: threading.Barrier, + stop: threading.Event, + finished: threading.Event, + is_probe: bool = False, +) -> list[tuple[float, float, int | None]]: + url = server.url_for("is_scaling_elastic_ep" if is_probe else "v1/completions") + payload = {"model": MODEL_NAME, "prompt": "Hello", "max_tokens": 4} + headers = None if dp_rank is None else {"X-data-parallel-rank": str(dp_rank)} + request_payload = None if is_probe else payload + responses = [] + is_ready = False + while not stop.is_set(): + request_start = time.perf_counter() + try: + response = requests.post( + url, json=request_payload, headers=headers, timeout=120 + ) + status_code = response.status_code + except requests.exceptions.RequestException: + status_code = None + responses.append((request_start, time.perf_counter(), status_code)) + if status_code == 200: + if not is_ready: + ready.wait(timeout=120) + is_ready = True + if finished.is_set(): + return responses + time.sleep(0.05) + return responses + + +def _downtime(responses: list[tuple[float, float, int | None]]) -> float: + rejected = [end for _, end, status in responses if status == 503] + if not rejected: + return 0 + recovered = next( + end for _, end, status in responses if status == 200 and end > rejected[-1] + ) + return recovered - rejected[0] + + +def _scale_with_traffic( + server: RemoteOpenAIServer, + source_dp_size: int, + new_dp_size: int, + traffic_mode: str, +) -> None: + traffic_clients: list[int | None] = [] + if traffic_mode == "light": + traffic_clients = [0] + elif traffic_mode == "heavy": + traffic_clients = [None] * source_dp_size + clients = [(None, True)] + [(rank, False) for rank in traffic_clients] + ready = threading.Barrier(len(clients) + 1) + stop = threading.Event() + finished = threading.Event() + + with ThreadPoolExecutor(max_workers=len(clients)) as executor: + futures = [ + executor.submit( + _traffic_loop, server, rank, ready, stop, finished, is_probe + ) + for rank, is_probe in clients + ] + try: + ready.wait(timeout=120) + start_time = time.perf_counter() + assert _send_scale_command(server, new_dp_size) + scale_seconds = time.perf_counter() - start_time + finished.set() + probe_result, *results = [future.result(timeout=120) for future in futures] + finally: + stop.set() + + bad_statuses = { + status + for responses in [probe_result, *results] + for _, _, status in responses + if status not in (200, 503) + } + assert not bad_statuses, f"traffic got unexpected statuses {bad_statuses}" + probe_503 = [start for start, _, status in probe_result if status == 503] + assert probe_503, "Scaling probe did not observe commit" + assert not results or any( + status == 200 and start_time <= request_start and request_end < probe_503[0] + for responses in results + for request_start, request_end, status in responses + ), "No request completed successfully during preparation" + + print( + f"[Elastic EP timing][{source_dp_size}->{new_dp_size}]" + f"[traffic={traffic_mode}] " + f"scale_seconds={scale_seconds:.3f} " + f"downtime_seconds={_downtime(probe_result):.3f}" + ) + + def _run_gsm8k_eval(server: RemoteOpenAIServer, stage: str) -> float: assert server.port is not None result = evaluate_gsm8k( @@ -59,7 +161,7 @@ def _run_gsm8k_eval(server: RemoteOpenAIServer, stage: str) -> float: return accuracy -def _base_serve_args(use_async_eplb: bool = False) -> list[str]: +def _base_serve_args(dp_size: int = 2, enforce_eager: bool = False) -> list[str]: args = [ "--trust-remote-code", "--tensor-parallel-size", @@ -78,57 +180,65 @@ def _base_serve_args(use_async_eplb: bool = False) -> list[str]: "--eplb-config.num_redundant_experts", "0", "--eplb-config.use_async", - "true" if use_async_eplb else "false", + "true", "--eplb-config.step_interval", - "10", + "300", "--eplb-config.window_size", "5", "--data-parallel-backend", "ray", "--data-parallel-size", - "2", + str(dp_size), "--api-server-count", "1", + "--disable-access-log-for-endpoints", + "/is_scaling_elastic_ep", ] leader_address = os.environ.get("LEADER_ADDRESS") if leader_address: args.extend(["--data-parallel-address", leader_address]) + if enforce_eager: + args.append("--enforce-eager") return args @pytest.mark.parametrize( - "use_async_eplb", [False, True], ids=["sync_eplb", "async_eplb"] + ("enforce_eager", "traffic_mode"), + [ + pytest.param(True, "none", id="enforce_eager_none"), + pytest.param(True, "light", id="enforce_eager_light"), + pytest.param(True, "heavy", id="enforce_eager_heavy"), + pytest.param(False, "heavy", id="cuda_graphs_heavy"), + ], ) @multi_gpu_test(num_gpus=4) -def test_elastic_ep_scaling(use_async_eplb: bool): - if use_async_eplb: - from vllm.distributed.eplb.eplb_communicator import has_nixl +def test_elastic_ep_scaling(enforce_eager: bool, traffic_mode: str): + from vllm.distributed.eplb.eplb_communicator import has_nixl - if not has_nixl(): - pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)") + if not has_nixl(): + pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)") - vllm_serve_args = _base_serve_args(use_async_eplb) + initial_dp_size = int(os.getenv("VLLM_TEST_ELASTIC_EP_INITIAL_DP", "2")) + target_dp_size = int(os.getenv("VLLM_TEST_ELASTIC_EP_TARGET_DP", "4")) + assert target_dp_size > initial_dp_size + vllm_serve_args = _base_serve_args(initial_dp_size, enforce_eager) with RemoteOpenAIServer( MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200 ) as server: - initial_accuracy = _run_gsm8k_eval(server, "Initial (2 GPUs)") - - assert _send_scale_command(server, 4) - time.sleep(10) - scale_up_accuracy = _run_gsm8k_eval(server, "After scale up (4 GPUs)") + initial_accuracy = _run_gsm8k_eval(server, "Initial") + _scale_with_traffic(server, initial_dp_size, target_dp_size, traffic_mode) + scale_up_accuracy = _run_gsm8k_eval(server, "After scale up") assert scale_up_accuracy >= initial_accuracy - ACCURACY_TOL, ( f"Scale up accuracy {scale_up_accuracy:.3f} dropped more than " f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}" ) - assert _send_scale_command(server, 2) - time.sleep(5) - scale_down_accuracy = _run_gsm8k_eval(server, "After scale down (2 GPUs)") - + _scale_with_traffic(server, target_dp_size, initial_dp_size, traffic_mode) + scale_down_accuracy = _run_gsm8k_eval(server, "After scale down") assert scale_down_accuracy >= initial_accuracy - ACCURACY_TOL, ( f"Scale down accuracy {scale_down_accuracy:.3f} dropped more than " f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}" @@ -147,24 +257,20 @@ def test_elastic_ep_scaling(use_async_eplb: bool): print(f" Tolerance: {ACCURACY_TOL:.3f}") -@pytest.mark.parametrize( - "use_async_eplb", [False, True], ids=["sync_eplb", "async_eplb"] -) @multi_gpu_test(num_gpus=4) -def test_elastic_ep_scaling_uneven(use_async_eplb: bool): +def test_elastic_ep_scaling_uneven(): """Test scale up with uneven worker distribution. This tests the case where num_new_workers % old_dp_size != 0, specifically 2 -> 3 where remainder = 1 % 2 = 1. This exercises the remainder handling in sender-receiver pairing. """ - if use_async_eplb: - from vllm.distributed.eplb.eplb_communicator import has_nixl + from vllm.distributed.eplb.eplb_communicator import has_nixl - if not has_nixl(): - pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)") + if not has_nixl(): + pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)") - vllm_serve_args = _base_serve_args(use_async_eplb) + vllm_serve_args = _base_serve_args() with RemoteOpenAIServer( MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200 @@ -174,7 +280,6 @@ def test_elastic_ep_scaling_uneven(use_async_eplb: bool): # Scale 2 -> 3: This has remainder = 1 % 2 = 1 # Tests uneven sender-receiver pairing assert _send_scale_command(server, 3) - time.sleep(10) scale_up_accuracy = _run_gsm8k_eval(server, "After scale up (3 GPUs)") assert scale_up_accuracy >= initial_accuracy - ACCURACY_TOL, ( @@ -184,7 +289,6 @@ def test_elastic_ep_scaling_uneven(use_async_eplb: bool): # Scale back down to 2 assert _send_scale_command(server, 2) - time.sleep(5) scale_down_accuracy = _run_gsm8k_eval(server, "After scale down (2 GPUs)") assert scale_down_accuracy >= initial_accuracy - ACCURACY_TOL, ( diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 949eb298a17..5ebc410f3c6 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -686,6 +686,14 @@ class ParallelConfig: and self.data_parallel_size > 1 ) + @property + def use_all2all(self) -> bool: + return ( + self.data_parallel_size > 1 + or self.use_sequence_parallel_moe + or (self.enable_expert_parallel and self.prefill_context_parallel_size > 1) + ) + @property def use_batched_dp_moe(self) -> bool: return ( @@ -786,6 +794,7 @@ class ParallelConfig: "data_parallel_master_ip", "data_parallel_master_port", "_data_parallel_master_port_list", + "_coord_store_port", "data_parallel_rpc_port", "rank", "master_addr", diff --git a/vllm/distributed/device_communicators/base_device_communicator.py b/vllm/distributed/device_communicators/base_device_communicator.py index 45438a54691..73fd1331f5c 100644 --- a/vllm/distributed/device_communicators/base_device_communicator.py +++ b/vllm/distributed/device_communicators/base_device_communicator.py @@ -175,6 +175,7 @@ class DeviceCommunicatorBase: unique_name: str = "", global_ranks: list[int] | None = None, global_world_size: int | None = None, + use_all2all: bool = False, ): self.device = device or torch.device("cpu") self.cpu_group = cpu_group @@ -204,26 +205,15 @@ class DeviceCommunicatorBase: self.global_world_size = dist.get_world_size() self.rank_in_group = dist.get_group_rank(self.cpu_group, self.global_rank) - use_ep = False all2all_backend = None from vllm.config import get_current_vllm_config_or_none config = get_current_vllm_config_or_none() if config is not None: - # initialize the all2all manager for DP or sequence-parallel EP. - parallel_config = config.parallel_config - use_ep = ( - parallel_config.data_parallel_size > 1 - or parallel_config.use_sequence_parallel_moe - or ( - parallel_config.enable_expert_parallel - and parallel_config.prefill_context_parallel_size > 1 - ) - ) - all2all_backend = parallel_config.all2all_backend + all2all_backend = config.parallel_config.all2all_backend self.is_ep_communicator = unique_name.split(":")[0] == "ep" - self.use_all2all = self.is_ep_communicator and use_ep + self.use_all2all = self.is_ep_communicator and use_all2all self.all2all_backend = all2all_backend self.all2all_manager: All2AllManagerBase | None = None diff --git a/vllm/distributed/device_communicators/cpu_communicator.py b/vllm/distributed/device_communicators/cpu_communicator.py index 9ec4b72f80d..8ea12d9255a 100644 --- a/vllm/distributed/device_communicators/cpu_communicator.py +++ b/vllm/distributed/device_communicators/cpu_communicator.py @@ -24,8 +24,11 @@ class CpuCommunicator(DeviceCommunicatorBase): device: torch.device | None = None, device_group: ProcessGroup | None = None, unique_name: str = "", + use_all2all: bool = False, ): - super().__init__(cpu_group, device, device_group, unique_name) + super().__init__( + cpu_group, device, device_group, unique_name, use_all2all=use_all2all + ) self.dist_module = torch.distributed if ( diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index fccc6ba60c3..23e37ca830e 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -36,6 +36,7 @@ class CudaCommunicator(DeviceCommunicatorBase): global_ranks: list[int] | None = None, global_world_size: int | None = None, tcp_store_group: StatelessProcessGroup | None = None, + use_all2all: bool = False, ): super().__init__( cpu_group, @@ -44,6 +45,7 @@ class CudaCommunicator(DeviceCommunicatorBase): unique_name, global_ranks, global_world_size, + use_all2all=use_all2all, ) if "tp" not in unique_name: # custom allreduce or torch symm mem can be used only by tp diff --git a/vllm/distributed/device_communicators/xpu_communicator.py b/vllm/distributed/device_communicators/xpu_communicator.py index 1b6ce9e8aae..7ca132824ec 100644 --- a/vllm/distributed/device_communicators/xpu_communicator.py +++ b/vllm/distributed/device_communicators/xpu_communicator.py @@ -20,8 +20,11 @@ class XpuCommunicator(DeviceCommunicatorBase): device: torch.device | None = None, device_group: ProcessGroup | None = None, unique_name: str = "", + use_all2all: bool = False, ): - super().__init__(cpu_group, device, device_group, unique_name) + super().__init__( + cpu_group, device, device_group, unique_name, use_all2all=use_all2all + ) self.ca_comm: None = None if self.use_all2all: if self.all2all_backend in ("naive", "allgather_reducescatter"): diff --git a/vllm/distributed/elastic_ep/elastic_execute.py b/vllm/distributed/elastic_ep/elastic_execute.py index b0c3740f57e..cea7fcb2f01 100644 --- a/vllm/distributed/elastic_ep/elastic_execute.py +++ b/vllm/distributed/elastic_ep/elastic_execute.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import copy import gc import weakref from collections.abc import Iterable, Sequence +from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import replace from typing import TYPE_CHECKING @@ -43,6 +43,7 @@ from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig from vllm.model_executor.layers.fused_moe.eep_reconfigure import ( make_eep_staged_quant_method, ) +from vllm.model_executor.warmup.kernel_warmup import kernel_warmup from vllm.utils import is_moe_layer from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper @@ -145,6 +146,10 @@ class ElasticEPScalingExecutor: self.worker_ref = weakref.ref(worker) self.reconfig_request = None self._staged_moe_quant_methods: dict[nn.Module, FusedMoEMethodBase] = {} + self._async_executor = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="ElasticEPAsync" + ) + self._async_future: Future[None] | None = None @property def worker(self): @@ -159,59 +164,68 @@ class ElasticEPScalingExecutor: raise ValueError(f"Unknown execute method: {execute_method}") return method(*args, **kwargs) - def _set_eplb_suppressed(self, suppressed: bool) -> None: - self.worker.model_runner.eep_eplb_suppressed = suppressed - ep_group = get_standby_ep_group() or get_ep_group() - if ep_group.rank == 0: - logger.info( - "[Elastic EP] EPLB %s elastic scaling transition", - "disabled during" if suppressed else "re-enabled after", - ) + def start_async(self, execute_method: str, *args, **kwargs) -> str: + if self._async_future is not None: + raise RuntimeError("Another Elastic EP async method is active") + if args and isinstance(args[0], ReconfigureDistributedRequest): + self.reconfig_request = args[0] + dp_rank = self.worker.vllm_config.parallel_config.data_parallel_rank + done_key = f"eep_async/{execute_method}/{dp_rank}/{self.worker.rank}" + self._async_future = self._async_executor.submit( + self._run_async, execute_method, *args, **kwargs + ) + self._async_future.add_done_callback(lambda _: self._mark_async_done(done_key)) + return done_key + + def _run_async(self, execute_method: str, *args, **kwargs) -> None: + from vllm.platforms import current_platform + + self.worker.vllm_config.enable_trace_function_call_for_thread() + assert hasattr(self.worker, "device") + current_platform.set_device(self.worker.device) + with set_current_vllm_config(self.worker.vllm_config): + self.execute(execute_method, *args, **kwargs) + + def _mark_async_done(self, done_key: str) -> None: + from vllm.distributed.utils import get_cached_tcp_store_client + + assert self.reconfig_request is not None + get_cached_tcp_store_client( + self.reconfig_request.new_data_parallel_master_ip, + self.reconfig_request.coord_store_port, + ).set(done_key, b"1") + + def clear_async(self) -> None: + future = self._async_future + if future is None: + raise RuntimeError("No Elastic EP async method is active") + if not future.done(): + raise RuntimeError("Elastic EP async method is not done") + self._async_future = None + future.result() def load_model(self) -> None: - ( - expanded_physical_to_logical, - num_logical_experts, - old_num_physical_experts, - ) = self.receive_expert_mapping() - num_physical_experts = expanded_physical_to_logical.shape[1] - self.worker.parallel_config.eplb_config.num_redundant_experts = ( - num_physical_experts - num_logical_experts - ) self.worker.load_model(load_dummy_weights=True) - self.worker.model_runner.setup_eplb_from_mapping( - expanded_physical_to_logical, old_num_physical_experts - ) - self._set_eplb_suppressed(True) def create_standby_groups( - self, reconfig_request: ReconfigureDistributedRequest + self, reconfig_request: ReconfigureDistributedRequest, use_all2all: bool ) -> None: self.reconfig_request = reconfig_request new_dp_size = reconfig_request.new_data_parallel_size old_dp_size = get_dp_group().world_size - world_size = self.worker.vllm_config.parallel_config.world_size + parallel_config = self.worker.vllm_config.parallel_config + world_size = parallel_config.world_size new_world_size_across_dp = world_size * new_dp_size - updated_config = copy.copy(self.worker.vllm_config) - updated_config.parallel_config = copy.deepcopy( - self.worker.vllm_config.parallel_config + create_standby_groups( + new_dp_size=new_dp_size, + new_world_size_across_dp=new_world_size_across_dp, + master_ip=reconfig_request.new_data_parallel_master_ip, + coord_store_port=reconfig_request.coord_store_port, + use_all2all=use_all2all, + enable_eplb=parallel_config.enable_eplb, ) - updated_config.parallel_config.data_parallel_size = new_dp_size - with set_current_vllm_config(updated_config): - create_standby_groups( - new_dp_size=new_dp_size, - new_world_size_across_dp=new_world_size_across_dp, - master_ip=reconfig_request.new_data_parallel_master_ip, - coord_store_port=reconfig_request.coord_store_port, - enable_eplb=updated_config.parallel_config.enable_eplb, - ) - if new_dp_size > old_dp_size: - self._set_eplb_suppressed(True) - eplb_state = self.worker.model_runner.eplb_state - if eplb_state is not None: - eplb_state.drain_async() - elif new_dp_size < old_dp_size: - self._stage_standby_moe_quant_methods() + if new_dp_size < old_dp_size: + self.stage_standby_moe_quant_methods() def transfer_weights(self, old_dp_size: int, new_dp_size: int) -> None: standby_dp_group = get_standby_dp_group() @@ -265,6 +279,7 @@ class ElasticEPScalingExecutor: model_config = self.worker.model_runner.model_config eplb_state = self.worker.model_runner.eplb_state assert eplb_state is not None + eplb_state.drain_async() eplb_model_state = eplb_state.model_states[model_config.compute_hash()] physical_to_logical = eplb_model_state.physical_to_logical_map num_physical_experts = physical_to_logical.shape[1] @@ -278,10 +293,6 @@ class ElasticEPScalingExecutor: src_rank=0, device=self.worker.device, ) - # New workers enter load_model after receiving the expert mapping. - # Stage replacement MoE kernels before returning to the state machine - # so existing ranks can participate in collective EP comm creation. - self._stage_standby_moe_quant_methods() def _make_eep_moe_config(self, module, dp_group, ep_group): parallel_config = self.worker.vllm_config.parallel_config @@ -300,7 +311,7 @@ class ElasticEPScalingExecutor: moe_parallel_config=moe_parallel_config, ) - def _stage_standby_moe_quant_methods(self) -> None: + def stage_standby_moe_quant_methods(self) -> None: standby_dp_group = get_standby_dp_group() standby_ep_group = get_standby_ep_group() model = self.worker.model_runner.get_model() @@ -500,26 +511,6 @@ class ElasticEPScalingExecutor: compilation_counter.stock_torch_compile_count += 1 self.worker.model_runner.model.compile(fullgraph=True, backend=backend) - multi_block_table = self.worker.model_runner.input_batch.block_table - saved_block_tables: list[tuple[torch.Tensor, torch.Tensor]] = [] - for bt in multi_block_table.block_tables: - saved_block_tables.append( - (bt.block_table.gpu.clone(), bt.block_table.cpu.clone()) - ) - multi_block_table.clear() - - unlock_workspace() - self.worker.compile_or_warm_up_model() - lock_workspace() - - for bt, (saved_gpu, saved_cpu) in zip( - multi_block_table.block_tables, saved_block_tables - ): - bt.block_table.gpu.copy_(saved_gpu) - bt.block_table.cpu.copy_(saved_cpu) - if new_dp_size < old_dp_size: - self._set_eplb_suppressed(False) - def _perform_eplb_reshuffle( self, rank_mapping: dict[int, int] | None = None ) -> None: @@ -553,12 +544,25 @@ class ElasticEPScalingExecutor: if get_ep_group().rank == 0: logger.info("[Elastic EP] Expert resharding completed") - def perform_eplb_reshuffle(self) -> None: + def commit_scale_up(self, is_existing_worker: bool) -> None: + if is_existing_worker: + self.broadcast_expert_mapping() + self.switch_and_prepare() + else: + mapping, _, num_valid_experts = self.receive_expert_mapping() + self.worker.model_runner.setup_eplb_from_mapping(mapping, num_valid_experts) self._perform_eplb_reshuffle() - self._set_eplb_suppressed(False) + self.warm_and_capture() + + def commit_scale_down(self, new_dp_size: int, removing: bool) -> None: + self.perform_scale_down_eplb_reshuffle(new_dp_size) + if removing: + self.switch_and_remove() + else: + self.switch_and_prepare() + self.warm_and_capture() def perform_scale_down_eplb_reshuffle(self, new_dp_size: int) -> None: - self._set_eplb_suppressed(True) eplb_state = self.worker.model_runner.eplb_state if eplb_state is not None: eplb_state.drain_async() @@ -599,12 +603,17 @@ class ElasticEPScalingExecutor: ) model = self.worker.model_runner.get_model() + expert_weights = [ + module.get_expert_weights() + for module in model.modules() + if is_moe_layer(module) + ] batch_transfer_weights( model=model, is_sender=False, peer_rank=sender_rank, dp_group=dp_group, - expert_weights=model.expert_weights, + expert_weights=expert_weights, ) torch.accelerator.synchronize() @@ -643,14 +652,17 @@ class ElasticEPScalingExecutor: with set_current_vllm_config(self.worker.vllm_config): prepare_communication_buffer_for_model(self.worker.model_runner.get_model()) - def rewarm_workspace(self) -> None: + def warmup_local_kernels(self) -> None: + with set_current_vllm_config(self.worker.vllm_config): + kernel_warmup(self.worker, process_local_only=True) + + def warm_and_capture(self) -> None: # Must run on every DP sibling in lockstep: _dummy_run calls # coordinate_batch_across_dp whenever data_parallel_size > 1 # (gpu_model_runner.py:3663), which deadlocks if any rank skips it. - # Save and clear block tables so profile_run/compile_or_warm_up_model - # don't write dummy slot mappings into real KV-cache blocks (mirrors - # switch_and_prepare's pattern). + # Save and clear block tables so the dummy MoE forward doesn't + # write dummy slot mappings into real KV-cache blocks. multi_block_table = self.worker.model_runner.input_batch.block_table saved_block_tables: list[tuple[torch.Tensor, torch.Tensor]] = [] for bt in multi_block_table.block_tables: @@ -660,19 +672,16 @@ class ElasticEPScalingExecutor: multi_block_table.clear() # _ensure_workspace_size allocates a fresh tensor on grow, leaving - # captured CUDA graphs with stale data pointers; drop graphs before - # re-warm so captures realign with the resized buffer. + # any captured CUDA graph with a stale data pointer; drop graphs + # before re-warm so captures realign with the resized buffer. self._release_cuda_graphs() unlock_workspace() - # Grow the MoE workspace at max_num_tokens. - # compile_or_warm_up_model alone only exercises cudagraph-capture - # sizes (≤64 tokens for this test) and leaves the workspace at - # ~10-14 MB; the post-all-to-all per-rank token count under real - # post-reshuffle routing needs hundreds of MB. Use _dummy_run - # directly (rather than profile_run) with skip_eplb=True so dummy - # routing doesn't pollute the just-rebalanced EPLB stats — same - # convention compile_or_warm_up_model itself uses. + # Grow the MoE workspace at max_num_tokens. compile_or_warm_up_model + # alone only exercises cudagraph-capture sizes and can leave the + # workspace too small for post-reshuffle routing. Use _dummy_run + # directly with skip_eplb=True so dummy routing doesn't pollute the + # just-rebalanced EPLB stats. runner = self.worker.model_runner runner._dummy_run(runner.max_num_tokens, is_profile=True, skip_eplb=True) self.worker.compile_or_warm_up_model() diff --git a/vllm/distributed/elastic_ep/elastic_state.py b/vllm/distributed/elastic_ep/elastic_state.py index 256efe46a4a..f33fd90e863 100644 --- a/vllm/distributed/elastic_ep/elastic_state.py +++ b/vllm/distributed/elastic_ep/elastic_state.py @@ -1,18 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import enum -import time import weakref -from datetime import timedelta -from typing import TYPE_CHECKING, Literal, TypeAlias +from concurrent.futures import Future, ThreadPoolExecutor +from typing import TYPE_CHECKING, Any, Literal, TypeAlias import torch.distributed from vllm.config import ParallelConfig from vllm.distributed import ( - sched_yield, stateless_destroy_torch_distributed_process_group, ) +from vllm.distributed.utils import get_cached_tcp_store_client from vllm.logger import init_logger from vllm.v1.engine import ( EEPNotificationType, @@ -31,35 +30,29 @@ WorkerType = Literal["existing", "new", "removing"] class ScaleUpExistingEngineState(enum.IntEnum): - WAIT_NEW_CORE_ENGINES_INIT = 0 - CREATE_STANDBY_GROUPS = 1 - TRANSFER_EXPERT_MAPPING = 2 - WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT = 3 - TRANSFER_WEIGHTS = 4 - SYNC_KV_CACHE_MEMORY_SIZE = 5 - SWITCH_AND_PREPARE = 6 - EPLB_RESHUFFLE = 7 - COMPLETE = 8 + CREATE_STANDBY_GROUPS = 0 + STAGE_QUANT_METHODS = 1 + TRANSFER_WEIGHTS = 2 + SYNC_KV_CACHE_MEMORY_SIZE = 3 + COMMIT_SCALE_UP = 4 # Blocks forward passes. + COMPLETE = 5 class ScaleUpNewEngineState(enum.IntEnum): PRE_KV_INIT = 0 PREPARE = 1 - EPLB_RESHUFFLE = 2 - COMPLETE = 3 + COMPLETE = 2 class ScaleDownRemainingEngineState(enum.IntEnum): PREPARE = 0 - EPLB_RESHUFFLE = 1 - SWITCH_AND_PREPARE = 2 - COMPLETE = 3 + COMMIT_SCALE_DOWN = 1 # Blocks forward passes. + COMPLETE = 2 class ScaleDownRemovingEngineState(enum.IntEnum): PREPARE = 0 - EPLB_RESHUFFLE = 1 - COMPLETE = 2 + COMPLETE = 1 EngineState: TypeAlias = ( @@ -70,15 +63,6 @@ EngineState: TypeAlias = ( ) -class _BarrierTimeoutError(RuntimeError): - """ - Exception raised for timeout - in the first stage of our two-staged - TCPStore based barrier to synchronize the - execution of all engines in the DP group. - """ - - class ElasticEPScalingState: def __init__( self, @@ -94,20 +78,24 @@ class ElasticEPScalingState: self.engine_core_ref = weakref.ref(engine_core) self.vllm_config = vllm_config self.old_dp_group = self.engine_core.dp_group if worker_type != "new" else None - self.old_dp_store = self.engine_core.dp_store if worker_type != "new" else None self.new_parallel_config: ParallelConfig = new_parallel_config self.new_dp_group = self.engine_core.dp_group if worker_type == "new" else None self.new_dp_store = self.engine_core.dp_store if worker_type == "new" else None self.worker_type = worker_type self.scale_type = scale_type self.reconfig_request = reconfig_request - + self.commit_requested = False + self._prepare_executor = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="ElasticEPPrepare" + ) + self._prepare_future: Future[Any] | None = None + self._new_dp_sync: tuple[object, Any] | None = None self.state: EngineState if scale_type == "scale_up": self.state = ( ScaleUpNewEngineState.PRE_KV_INIT if worker_type == "new" - else ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_INIT + else ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS ) else: self.state = ( @@ -130,6 +118,31 @@ class ElasticEPScalingState: raise RuntimeError("Engine core has been garbage collected") return engine_core + def _collective_rpc(self, *args, **kwargs): + return self.model_executor.collective_rpc(*args, **kwargs) + + def _execute_async(self, execute_method: str, *args) -> bool: + if self._prepare_future is None: + done_keys = self._collective_rpc( + "elastic_ep_execute", + args=("start_async", execute_method, *args), + ) + assert self.reconfig_request is not None + coord_store = get_cached_tcp_store_client( + self.reconfig_request.new_data_parallel_master_ip, + self.reconfig_request.coord_store_port, + ) + self._prepare_future = self._prepare_executor.submit( + coord_store.wait, done_keys + ) + if not self._prepare_future.done(): + return False + + self._prepare_future.result() + self._collective_rpc("elastic_ep_execute", args=("clear_async",)) + self._prepare_future = None + return True + def progress(self) -> bool: if self.scale_type == "scale_up": return ( @@ -149,157 +162,43 @@ class ElasticEPScalingState: assert self.progress() assert self.state == ScaleUpNewEngineState.PREPARE - def _execute_tcp_store_barrier( - self, dp_store, group_rank, group_size, barrier_id, timeout=None - ): - arrival_key = f"arrival_{barrier_id}_{group_rank}" - dp_store.set(arrival_key, b"1") - - start_time = time.time() - processes_arrived: set[int] = set() - - while len(processes_arrived) < group_size: - if ( - timeout is not None - and time.time() - start_time > timeout.total_seconds() - ): - raise _BarrierTimeoutError( - f"Barrier timed out after {timeout.total_seconds()} seconds" - ) - - for i in range(group_size): - if i in processes_arrived: - continue - - key = f"arrival_{barrier_id}_{i}" - present = dp_store.check([key]) - if present: - processes_arrived.add(i) - - if len(processes_arrived) < group_size: - sched_yield() - - def _staged_barrier(self, use_new_group: bool, barrier_name: str) -> bool: - """ - Execute a two-staged barrier to synchronize all engines in the DP group. - - Some DP EngineCores may receive the reconfiguration notifications - later than others, and already proceed to engine step (model forward) - in the busy loop. - In this case, EngineCores that already proceed to reconfiguration - should skip reconfiguration and execute model forward for one more - step, so in the next step, all EngineCores will be synchronized. - We use a two-staged barrier to achieve this. The first time each - EngineCore executes the barrier, if a timeout is reached before the - barrier completes, that means some EngineCores have already entered - engine step. The EngineCores that timed out will then proceed to - engine step, and will synchronize with the other EngineCores in the - next step with a barrier without timeout. - """ - dp_group = self.new_dp_group if use_new_group else self.old_dp_group - dp_store = self.new_dp_store if use_new_group else self.old_dp_store - assert dp_group is not None and dp_store is not None - - group_rank = dp_group.rank() - group_size = dp_group.size() - barrier_id = f"eep_barrier_{barrier_name}" - sync_key = f"{barrier_id}_sync" - - # TODO(yongji): figure out appropriate timeout for the barrier - timeout = None if dp_store.check([sync_key]) else timedelta(seconds=5) - - try: - self._execute_tcp_store_barrier( - dp_store, group_rank, group_size, barrier_id, timeout=timeout - ) - torch.distributed.barrier(dp_group) - if group_rank == 0: - dp_store.delete_key(sync_key) - for i in range(group_size): - dp_store.delete_key(f"arrival_{barrier_id}_{i}") - return True - except _BarrierTimeoutError as e: - if timeout is None: - raise RuntimeError("Unexpected timeout encountered") from e - dp_store.compare_set(sync_key, "", b"1") - return False - def _progress_existing_engine(self) -> bool: state = self.state - assert self.old_dp_group is not None and self.old_dp_store is not None + assert self.old_dp_group is not None - if state == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_INIT: - return False - - elif state == ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS: - # NOTE(yongji): wait for all existing workers to receive the request - if ( - int(self.old_dp_store.get("eep_barrier_engine_count")) - < self.old_dp_group.size() - ): + if state == ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS: + if not self._create_standby_groups(): return False - if not self._staged_barrier( - use_new_group=False, barrier_name="create_standby_groups" - ): - return False - if self.old_dp_group.rank() == 0: - self.old_dp_store.delete_key("eep_barrier_engine_count") - self._create_standby_groups() - self.state = ScaleUpExistingEngineState.TRANSFER_EXPERT_MAPPING + self.state = ScaleUpExistingEngineState.STAGE_QUANT_METHODS return True - elif state == ScaleUpExistingEngineState.TRANSFER_EXPERT_MAPPING: - self._transfer_expert_mapping() - self.state = ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT + elif state == ScaleUpExistingEngineState.STAGE_QUANT_METHODS: + if not self._execute_async("stage_standby_moe_quant_methods"): + return False + self.state = ScaleUpExistingEngineState.TRANSFER_WEIGHTS return True - elif state == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT: - return False - elif state == ScaleUpExistingEngineState.TRANSFER_WEIGHTS: - if ( - int(self.old_dp_store.get("eep_barrier_engine_count")) - < self.old_dp_group.size() - ): + if not self._transfer_weights(): return False - if not self._staged_barrier( - use_new_group=False, barrier_name="transfer_weights" - ): - return False - if self.old_dp_group.rank() == 0: - self.old_dp_store.delete_key("eep_barrier_engine_count") - self._transfer_weights() self.state = ScaleUpExistingEngineState.SYNC_KV_CACHE_MEMORY_SIZE return True elif state == ScaleUpExistingEngineState.SYNC_KV_CACHE_MEMORY_SIZE: - self._sync_kv_cache_memory_size() - self.state = ScaleUpExistingEngineState.SWITCH_AND_PREPARE + if not self._sync_kv_cache_memory_size(): + return False + self.state = ScaleUpExistingEngineState.COMMIT_SCALE_UP + self._mark_ready_for_switch() return True - elif state == ScaleUpExistingEngineState.SWITCH_AND_PREPARE: - self._switch_and_prepare() - self.state = ScaleUpExistingEngineState.EPLB_RESHUFFLE - assert self.new_dp_store is not None - self.new_dp_store.add("eep_barrier_engine_count", 1) - return True - - elif state == ScaleUpExistingEngineState.EPLB_RESHUFFLE: - assert self.new_dp_group is not None and self.new_dp_store is not None - if ( - int(self.new_dp_store.get("eep_barrier_engine_count")) - < self.new_dp_group.size() - ): + elif state == ScaleUpExistingEngineState.COMMIT_SCALE_UP: + if not self.commit_requested: return False - if not self._staged_barrier( - use_new_group=True, barrier_name="eplb_reshuffle" - ): - return False - if self.new_dp_group.rank() == 0: - self.new_dp_store.delete_key("eep_barrier_engine_count") - self._eplb_reshuffle() + self._commit_new_dp_group() + self._collective_rpc("elastic_ep_execute", args=("commit_scale_up", True)) self.state = ScaleUpExistingEngineState.COMPLETE self._update_parallel_config() + self._send_reconfigure_finished() return True else: @@ -311,22 +210,17 @@ class ElasticEPScalingState: assert self.new_dp_group is not None and self.new_dp_store is not None if state == ScaleUpNewEngineState.PRE_KV_INIT: - self.engine_core._eep_send_engine_core_notification( - EEPNotificationType.NEW_CORE_ENGINES_WEIGHTS_INIT_READY - ) - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("receive_weights",) - ) + self._collective_rpc("elastic_ep_execute", args=("receive_weights",)) self.engine_core.available_gpu_memory_for_kv_cache = ( ParallelConfig.sync_kv_cache_memory_size(self.new_dp_group, -1) ) - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("prepare_new_worker",) - ) + self._collective_rpc("elastic_ep_execute", args=("prepare_new_worker",)) self.state = ScaleUpNewEngineState.PREPARE return True elif state == ScaleUpNewEngineState.PREPARE: + self._collective_rpc("elastic_ep_execute", args=("warmup_local_kernels",)) + self._mark_ready_for_switch() tensor = torch.tensor([0, 0, 0], dtype=torch.int32, device="cpu") torch.distributed.all_reduce( tensor, @@ -337,22 +231,7 @@ class ElasticEPScalingState: self.engine_core.engines_running = bool(data[0]) self.engine_core.current_wave = int(data[1]) self.engine_core.step_counter = int(data[2]) - self.state = ScaleUpNewEngineState.EPLB_RESHUFFLE - self.new_dp_store.add("eep_barrier_engine_count", 1) - return True - - elif state == ScaleUpNewEngineState.EPLB_RESHUFFLE: - if ( - int(self.new_dp_store.get("eep_barrier_engine_count")) - < self.new_dp_group.size() - ): - return False - if not self._staged_barrier( - use_new_group=True, barrier_name="eplb_reshuffle" - ): - return False - assert self.new_dp_group.rank() > 0 - self._eplb_reshuffle() + self._collective_rpc("elastic_ep_execute", args=("commit_scale_up", False)) self.state = ScaleUpNewEngineState.COMPLETE return True @@ -362,38 +241,23 @@ class ElasticEPScalingState: def _progress_remaining_engine(self) -> bool: state = self.state - assert self.old_dp_group is not None and self.old_dp_store is not None + assert self.old_dp_group is not None if state == ScaleDownRemainingEngineState.PREPARE: - self.state = ScaleDownRemainingEngineState.EPLB_RESHUFFLE - self.old_dp_store.add("eep_barrier_engine_count", 1) - return True + if self._create_standby_groups(): + self.state = ScaleDownRemainingEngineState.COMMIT_SCALE_DOWN + self._mark_ready_for_switch() + return True + return False - elif state == ScaleDownRemainingEngineState.EPLB_RESHUFFLE: - if ( - int(self.old_dp_store.get("eep_barrier_engine_count")) - < self.old_dp_group.size() - ): + elif state == ScaleDownRemainingEngineState.COMMIT_SCALE_DOWN: + if not self.commit_requested: return False - if not self._staged_barrier( - use_new_group=False, barrier_name="eplb_reshuffle" - ): - return False - if self.old_dp_group.rank() == 0: - self.old_dp_store.delete_key("eep_barrier_engine_count") - self._eplb_reshuffle_before_scale_down() - self.state = ScaleDownRemainingEngineState.SWITCH_AND_PREPARE - # NOTE(yongji): currently, after EPLB reshuffle - # that redistributes experts to remaining workers, workers - # to be removed will immediately initiate shutdown; - # existing workers can no longer execute forward steps using - # the old setup. In the future, we may keep - # the removing workers alive a bit longer, - # e.g., to drain in-batch requests. - self._create_standby_groups() - self._switch_and_prepare() + self._commit_scale_down(removing=False) + self._commit_new_dp_group() self._update_parallel_config() self.state = ScaleDownRemainingEngineState.COMPLETE + self._send_reconfigure_finished() return True else: @@ -402,26 +266,11 @@ class ElasticEPScalingState: def _progress_removing_engine(self) -> bool: state = self.state - assert self.old_dp_group is not None and self.old_dp_store is not None + assert self.old_dp_group is not None if state == ScaleDownRemovingEngineState.PREPARE: - self.state = ScaleDownRemovingEngineState.EPLB_RESHUFFLE - self.old_dp_store.add("eep_barrier_engine_count", 1) - return True - - if state == ScaleDownRemovingEngineState.EPLB_RESHUFFLE: - if ( - int(self.old_dp_store.get("eep_barrier_engine_count")) - < self.old_dp_group.size() - ): - return False - if not self._staged_barrier( - use_new_group=False, barrier_name="eplb_reshuffle" - ): - return False assert self.old_dp_group.rank() > 0 - self._eplb_reshuffle_before_scale_down() - self._switch_and_remove() + self._commit_scale_down(removing=True) self.state = ScaleDownRemovingEngineState.COMPLETE self.engine_core._eep_send_engine_core_notification( EEPNotificationType.SHUTDOWN_COMPLETE @@ -432,22 +281,22 @@ class ElasticEPScalingState: assert self.state == ScaleDownRemovingEngineState.COMPLETE return True - def handle_notification(self, notification_type: EEPNotificationType): - assert self.worker_type != "new" - assert self.old_dp_store is not None - if ( - notification_type == EEPNotificationType.NEW_CORE_ENGINES_INIT_READY - and self.state == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_INIT - ): - self.old_dp_store.add("eep_barrier_engine_count", 1) - self.state = ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS - elif ( - notification_type == EEPNotificationType.NEW_CORE_ENGINES_WEIGHTS_INIT_READY - and self.state - == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT - ): - self.old_dp_store.add("eep_barrier_engine_count", 1) - self.state = ScaleUpExistingEngineState.TRANSFER_WEIGHTS + def is_ready_for_switch(self) -> bool: + return self.worker_type == "existing" and ( + self.state is ScaleUpExistingEngineState.COMMIT_SCALE_UP + or self.state is ScaleDownRemainingEngineState.COMMIT_SCALE_DOWN + ) + + @property + def ready_key(self) -> str: + return f"eep_ready/{self.engine_core.dp_rank}" + + def _mark_ready_for_switch(self) -> None: + parallel_config = self.new_parallel_config + get_cached_tcp_store_client( + parallel_config.data_parallel_master_ip, + parallel_config._coord_store_port, + ).set(self.ready_key, b"1") def is_complete(self) -> bool: if self.scale_type == "scale_up": @@ -462,50 +311,78 @@ class ElasticEPScalingState: else self.state == ScaleDownRemainingEngineState.COMPLETE ) - def _create_standby_groups(self): + def _init_new_dp_group(self) -> tuple[Any, Any]: + return self.new_parallel_config.stateless_init_dp_group(return_store=True) + + def _ensure_new_dp_group(self) -> bool: + if self.new_dp_group is not None: + return True + + if self._prepare_future is None: + self._prepare_future = self._prepare_executor.submit( + self._init_new_dp_group + ) + if not self._prepare_future.done(): + return False + + self.new_dp_group, self.new_dp_store = self._prepare_future.result() + self._prepare_future = None + return True + + def _create_standby_groups(self) -> bool: assert self.old_dp_group is not None - self.new_dp_group, self.new_dp_store = ( - self.new_parallel_config.stateless_init_dp_group(return_store=True) - ) - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("create_standby_groups", self.reconfig_request) - ) + if not self._ensure_new_dp_group(): + return False + if not self._execute_async( + "create_standby_groups", + self.reconfig_request, + self.new_parallel_config.use_all2all, + ): + return False if self.old_dp_group.rank() == 0: logger.info("[Elastic EP] Created standby communication groups") + return True - def _transfer_weights(self): + def _transfer_weights(self) -> bool: assert self.reconfig_request is not None and self.old_dp_group is not None old_dp_size = self.old_dp_group.size() new_dp_size = self.reconfig_request.new_data_parallel_size - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("transfer_weights", old_dp_size, new_dp_size) - ) + if not self._execute_async("transfer_weights", old_dp_size, new_dp_size): + return False if self.old_dp_group.rank() == 0: logger.info("[Elastic EP] Transferred weights to new workers") + return True - def _transfer_expert_mapping(self): - assert self.old_dp_group is not None - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("broadcast_expert_mapping",) - ) - if self.old_dp_group.rank() == 0: - logger.info("[Elastic EP] Broadcasted expert mapping to new workers") - - def _sync_kv_cache_memory_size(self): + def _sync_kv_cache_memory_size(self) -> bool: assert self.engine_core.available_gpu_memory_for_kv_cache > 0 assert self.new_dp_group is not None and self.old_dp_group is not None - ParallelConfig.sync_kv_cache_memory_size( - self.new_dp_group, - self.engine_core.available_gpu_memory_for_kv_cache, - ) + + if self._new_dp_sync is None: + tensor = torch.tensor( + [self.engine_core.available_gpu_memory_for_kv_cache], + dtype=torch.int64, + device="cpu", + ) + work = torch.distributed.all_reduce( + tensor, + op=torch.distributed.ReduceOp.MIN, + group=self.new_dp_group, + async_op=True, + ) + self._new_dp_sync = (tensor, work) + return False + + _, work = self._new_dp_sync + if not work.is_completed(): + return False + work.wait() + self._new_dp_sync = None if self.old_dp_group.rank() == 0: logger.info("[Elastic EP] Synced KV cache memory size to new workers") + return True - def _switch_and_prepare(self): - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("switch_and_prepare",) - ) + def _commit_new_dp_group(self): old_dp_group = self.old_dp_group stateless_destroy_torch_distributed_process_group(old_dp_group) assert self.new_dp_group is not None @@ -529,41 +406,28 @@ class ElasticEPScalingState: self.engine_core.current_wave = int(data[1]) self.engine_core.step_counter = int(data[2]) if new_dp_group.rank() == 0: + logger.info("[Elastic EP] Switched to new setup") + + def _send_reconfigure_finished(self): + assert self.new_dp_group is not None + if self.new_dp_group.rank() == 0: self.engine_core._eep_send_engine_core_notification( EEPNotificationType.RECONFIGURE_FINISHED ) - logger.info("[Elastic EP] Switched to new setup") - def _eplb_reshuffle(self): - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("perform_eplb_reshuffle",) - ) - # Reshuffle changes per-rank token routing; the locked MoE workspace - # may now be too small. Rewarm covers both new and existing engines. - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("rewarm_workspace",) - ) - assert self.new_dp_group is not None - if self.new_dp_group.rank() == 0: - logger.info("[Elastic EP] EPLB reshuffle completed") - - def _eplb_reshuffle_before_scale_down(self): + def _commit_scale_down(self, removing: bool): assert self.reconfig_request is not None and self.old_dp_group is not None - self.model_executor.collective_rpc( + self._collective_rpc( "elastic_ep_execute", args=( - "perform_scale_down_eplb_reshuffle", + "commit_scale_down", self.reconfig_request.new_data_parallel_size, + removing, ), ) if self.old_dp_group.rank() == 0: logger.info("[Elastic EP] EPLB reshuffle completed") - def _switch_and_remove(self): - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("switch_and_remove",) - ) - def _update_parallel_config(self): assert self.reconfig_request is not None reconfig_request = self.reconfig_request diff --git a/vllm/distributed/elastic_ep/standby_state.py b/vllm/distributed/elastic_ep/standby_state.py index 846793a955f..1892f3e7942 100644 --- a/vllm/distributed/elastic_ep/standby_state.py +++ b/vllm/distributed/elastic_ep/standby_state.py @@ -39,6 +39,7 @@ def create_standby_groups( new_world_size_across_dp: int, master_ip: str, coord_store_port: int, + use_all2all: bool, enable_eplb: bool = True, backend: str | None = None, ) -> None: @@ -86,7 +87,7 @@ def create_standby_groups( ) standby_ep_ranks = [x.tolist() for x in standby_ep_ranks] _STANDBY_EP = _init_stateless_group( - standby_ep_ranks, "ep", master_ip, backend, coord_store=coord_store + standby_ep_ranks, "ep", master_ip, backend, coord_store, use_all2all=use_all2all ) if enable_eplb: diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 4284a609d67..a90e8acbcad 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -414,6 +414,7 @@ class GroupCoordinator: use_device_communicator: bool, # whether to use device communicator use_message_queue_broadcaster: bool = False, group_name: str | None = None, + use_all2all: bool = False, ): group_name = group_name or "anonymous" self.unique_name = _get_unique_name(group_name) @@ -508,6 +509,7 @@ class GroupCoordinator: device=self.device, device_group=self.device_group, unique_name=self.unique_name, + use_all2all=use_all2all, ) from vllm.distributed.device_communicators.shm_broadcast import MessageQueue @@ -1321,6 +1323,7 @@ def init_model_parallel_group( use_message_queue_broadcaster: bool = False, group_name: str | None = None, use_device_communicator: bool = True, + use_all2all: bool = False, ) -> GroupCoordinator: return GroupCoordinator( group_ranks=group_ranks, @@ -1329,6 +1332,7 @@ def init_model_parallel_group( use_device_communicator=use_device_communicator, use_message_queue_broadcaster=use_message_queue_broadcaster, group_name=group_name, + use_all2all=use_all2all, ) @@ -1339,6 +1343,7 @@ def _init_stateless_group( backend: str, coord_store: Store, use_device_communicator: bool = True, + use_all2all: bool = False, ) -> "StatelessGroupCoordinator": """Create a StatelessGroupCoordinator with the given parameters.""" from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator @@ -1354,6 +1359,7 @@ def _init_stateless_group( coord_store=coord_store, global_rank=world.rank, global_world_size=world.world_size, + use_all2all=use_all2all, ) @@ -1924,6 +1930,7 @@ def initialize_model_parallel( .unbind(0) ) group_ranks = [x.tolist() for x in group_ranks] + use_all2all = parallel_config.use_all2all if enable_elastic_ep: _EP = _init_stateless_group( group_ranks, @@ -1931,10 +1938,15 @@ def initialize_model_parallel( parallel_config.data_parallel_master_ip, backend, coord_store=coord_store, + use_all2all=use_all2all, ) else: _EP = init_model_parallel_group( - group_ranks, get_world_group().local_rank, backend, group_name="ep" + group_ranks, + get_world_group().local_rank, + backend, + group_name="ep", + use_all2all=use_all2all, ) # Create EPLB group with the same ranks as EP if EPLB is enabled. diff --git a/vllm/distributed/stateless_coordinator.py b/vllm/distributed/stateless_coordinator.py index 5f4597d07cb..38c74a97c55 100644 --- a/vllm/distributed/stateless_coordinator.py +++ b/vllm/distributed/stateless_coordinator.py @@ -79,6 +79,7 @@ class StatelessGroupCoordinator(GroupCoordinator): host: str = "127.0.0.1", global_rank: int = 0, global_world_size: int = 1, + use_all2all: bool = False, ): group_name = group_name or "anonymous" self.unique_name = _get_unique_name(group_name) @@ -191,6 +192,7 @@ class StatelessGroupCoordinator(GroupCoordinator): global_ranks=self.ranks, global_world_size=global_world_size, tcp_store_group=self.tcp_store_group, + use_all2all=use_all2all, ) self.mq_broadcaster = None diff --git a/vllm/entrypoints/serve/elastic_ep/api_router.py b/vllm/entrypoints/serve/elastic_ep/api_router.py index e711a257ddd..02a24250905 100644 --- a/vllm/entrypoints/serve/elastic_ep/api_router.py +++ b/vllm/entrypoints/serve/elastic_ep/api_router.py @@ -12,10 +12,7 @@ from vllm.engine.protocol import EngineClient from vllm.entrypoints.openai.engine.protocol import ( ErrorResponse, ) -from vllm.entrypoints.serve.elastic_ep.middleware import ( - get_scaling_elastic_ep, - set_scaling_elastic_ep, -) +from vllm.entrypoints.serve.elastic_ep.middleware import get_scaling_elastic_ep from vllm.entrypoints.serve.utils.api_utils import validate_json_request from vllm.logger import init_logger @@ -64,8 +61,6 @@ async def scale_elastic_ep(raw_request: Request): status_code=400, detail="drain_timeout must be a positive integer" ) - # Set scaling flag to prevent new requests - set_scaling_elastic_ep(True) client = engine_client(raw_request) try: await client.scale_elastic_ep(new_data_parallel_size, drain_timeout) @@ -83,8 +78,6 @@ async def scale_elastic_ep(raw_request: Request): except Exception as e: logger.error("Scale failed: %s", e) raise HTTPException(status_code=500, detail="Scale failed") from e - finally: - set_scaling_elastic_ep(False) @router.post("/is_scaling_elastic_ep") diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index e461dae0bb8..b2c989b295c 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -93,7 +93,7 @@ def _warmup_ll_bf16_router_gemm(model: torch.nn.Module) -> None: ) -def kernel_warmup(worker: "Worker"): +def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): from vllm.model_executor.warmup.minimax_m3_msa_warmup import ( minimax_m3_msa_warmup, ) @@ -118,6 +118,22 @@ def kernel_warmup(worker: "Worker"): ) # Run next so input-prep kernels JIT against pristine runner state. + if worker.vllm_config.kernel_config.enable_jit_warmup: + fa4_cutedsl_warmup(worker) + sparse_mla_triton_warmup(worker) + + if current_platform.has_device_capability(90): + _warmup_ll_bf16_router_gemm(worker.get_model()) + + if worker.vllm_config.kernel_config.enable_cutedsl_warmup: + # TODO(roberto): Remove after registered CuTeDSL warmups are migrated + # to the shared JIT warmup infrastructure. + # https://github.com/vllm-project/vllm/pull/47451 + cutedsl_warmup() + + if process_local_only: + return + flashinfer_sparse_mla_decode_autotune_warmup(worker) deepseek_v4_sparse_mla_attention_warmup(worker) @@ -143,9 +159,6 @@ def kernel_warmup(worker: "Worker"): elif has_flashinfer() and current_platform.has_device_capability(90): flashinfer_autotune(worker.model_runner) - if current_platform.has_device_capability(90): - _warmup_ll_bf16_router_gemm(worker.get_model()) - # FlashInfer attention warmup # Only warmup if the model has FlashInfer attention groups # and is not a pooling model @@ -178,16 +191,6 @@ def kernel_warmup(worker: "Worker"): create_mixed_batch=True, ) - if worker.vllm_config.kernel_config.enable_cutedsl_warmup: - # TODO(roberto): Remove after registered CuTeDSL warmups are migrated - # to the shared JIT warmup infrastructure. - # https://github.com/vllm-project/vllm/pull/47451 - cutedsl_warmup() - - if worker.vllm_config.kernel_config.enable_jit_warmup: - fa4_cutedsl_warmup(worker) - sparse_mla_triton_warmup(worker) - def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None: if envs.VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS is not None: diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index e80be0e45d7..83033ecf81d 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -35,8 +35,6 @@ FT_STATUS_CALL_ID = -2 class EEPNotificationType(enum.Enum): - NEW_CORE_ENGINES_INIT_READY = "NEW_CORE_ENGINES_INIT_READY" - NEW_CORE_ENGINES_WEIGHTS_INIT_READY = "NEW_CORE_ENGINES_WEIGHTS_INIT_READY" RECONFIGURE_FINISHED = "RECONFIGURE_FINISHED" SHUTDOWN_COMPLETE = "SHUTDOWN_COMPLETE" diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 5c2e01cf44b..e8d33c961dc 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -109,6 +109,7 @@ class AsyncLLM(EngineClient): maybe_register_config_serialize_by_value() self.vllm_config = vllm_config + self._elastic_ep_lock = asyncio.Lock() self.model_config = vllm_config.model_config self.observability_config = vllm_config.observability_config @@ -998,17 +999,27 @@ class AsyncLLM(EngineClient): "waiting for requests to drain." ) + async def _drain_requests_for_elastic_ep(self, drain_timeout: int) -> None: + try: + logger.info( + "VLLM_ELASTIC_EP_DRAIN_REQUESTS is set, " + "waiting for requests to drain before scaling" + ) + await self.wait_for_requests_to_drain(drain_timeout) + except BaseException: + set_scaling_elastic_ep(False) + raise + async def scale_elastic_ep( self, new_data_parallel_size: int, drain_timeout: int = 300 ): - """ - Scale up or down the data parallel size by adding or removing - engine cores. - Args: - new_data_parallel_size: The new number of data parallel workers - drain_timeout: - Maximum time to wait for requests to drain (seconds) - """ + """Scale the elastic EP data parallel size.""" + async with self._elastic_ep_lock: + await self._scale_elastic_ep(new_data_parallel_size, drain_timeout) + + async def _scale_elastic_ep( + self, new_data_parallel_size: int, drain_timeout: int + ) -> None: old_data_parallel_size = self.vllm_config.parallel_config.data_parallel_size if old_data_parallel_size == new_data_parallel_size: logger.info( @@ -1017,12 +1028,7 @@ class AsyncLLM(EngineClient): ) return - if envs.VLLM_ELASTIC_EP_DRAIN_REQUESTS: - logger.info( - "VLLM_ELASTIC_EP_DRAIN_REQUESTS is set, " - "waiting for requests to drain before scaling" - ) - await self.wait_for_requests_to_drain(drain_timeout) + await self.engine_core.prepare_elastic_ep(new_data_parallel_size) # recreate stat loggers if new_data_parallel_size > old_data_parallel_size and self.log_stats: @@ -1042,11 +1048,12 @@ class AsyncLLM(EngineClient): self.logger_manager.log_engine_initialized() set_scaling_elastic_ep(True) - try: - await self.engine_core.scale_elastic_ep(new_data_parallel_size) - self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size - finally: - set_scaling_elastic_ep(False) + if envs.VLLM_ELASTIC_EP_DRAIN_REQUESTS: + await self._drain_requests_for_elastic_ep(drain_timeout) + + await self.engine_core.commit_elastic_ep() + self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size + set_scaling_elastic_ep(False) async def handle_fault( self, fault_tolerance_request: FaultToleranceRequest diff --git a/vllm/v1/engine/coordinator.py b/vllm/v1/engine/coordinator.py index 2f3b03636d7..d7f05cffc8a 100644 --- a/vllm/v1/engine/coordinator.py +++ b/vllm/v1/engine/coordinator.py @@ -376,6 +376,9 @@ class DPCoordinatorProc: eng_index = outputs.engine_index scheduler_stats = outputs.scheduler_stats if scheduler_stats: + # Elastic EP stats may arrive while the engine list changes. + if eng_index >= len(self.engines): + continue # 1. Updated request load stats - update our local # state with these. stats = self.engines[eng_index].request_counts diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 9917f810b5b..ecac92f5fe0 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -327,8 +327,9 @@ class EngineCore: vllm_config.validate_block_size() - # Initialize kv cache and warmup the execution self.model_executor.initialize_from_config(kv_cache_configs) + if not envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: + self.model_executor.compile_or_warm_up_model() elapsed = time.time() - start compile_time = vllm_config.compilation_config.compilation_time @@ -993,9 +994,7 @@ class EngineCore: raise NotImplementedError def _eep_send_engine_core_notification( - self, - notification_type: EEPNotificationType, - vllm_config: VllmConfig | None = None, + self, notification_type: EEPNotificationType ): raise NotImplementedError @@ -1070,11 +1069,6 @@ class EngineCoreProc(EngineCore): self.addresses = addresses self.process_input_queue_block = True - if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: - self._eep_send_engine_core_notification( - EEPNotificationType.NEW_CORE_ENGINES_INIT_READY, - vllm_config=vllm_config, - ) self._init_data_parallel(vllm_config) super().__init__( @@ -2096,12 +2090,16 @@ class DPEngineCoreProc(EngineCoreProc): self._maybe_publish_request_counts() if self.eep_scaling_state is not None: - _ = self.eep_scaling_state.progress() - if self.eep_scaling_state.is_complete(): - if self.eep_scaling_state.worker_type == "removing": + state = self.eep_scaling_state + if state.commit_requested or not state.is_ready_for_switch(): + state.progress() + if state.is_complete(): + if state.worker_type == "removing": raise SystemExit self.process_input_queue_block = True self.eep_scaling_state = None + elif not state.commit_requested and state.is_ready_for_switch(): + self.process_input_queue_block = True executed = self._process_engine_step() self._maybe_publish_request_counts() @@ -2171,7 +2169,7 @@ class DPEngineCoreProc(EngineCoreProc): def reinitialize_distributed( self, reconfig_request: ReconfigureDistributedRequest - ) -> None: + ) -> str: from copy import deepcopy from vllm.distributed.elastic_ep.elastic_state import ElasticEPScalingState @@ -2203,7 +2201,10 @@ class DPEngineCoreProc(EngineCoreProc): == ReconfigureRankType.SHUTDOWN_CURRENT_RANK ) - self.eep_scaling_state = ElasticEPScalingState( + if self.eep_scaling_state is not None: + raise RuntimeError("Elastic EP reconfiguration is already active") + + state = ElasticEPScalingState( model_executor=self.model_executor, engine_core=self, vllm_config=self.vllm_config, @@ -2212,30 +2213,34 @@ class DPEngineCoreProc(EngineCoreProc): scale_type="scale_down" if is_scale_down else "scale_up", reconfig_request=reconfig_request, ) + self.eep_scaling_state = state + self.process_input_queue_block = False logger.info( "[Elastic EP] Received reconfiguration request and starting scaling up/down" ) + return state.ready_key + + def commit_prepared_elastic_ep(self) -> None: + state = self.eep_scaling_state + if state is None or state.commit_requested or not state.is_ready_for_switch(): + raise RuntimeError("No prepared Elastic EP reconfiguration is ready") + state.commit_requested = True + self.process_input_queue_block = False + logger.info("[Elastic EP] Committing prepared reconfiguration") def _eep_send_engine_core_notification( - self, - notification_type: EEPNotificationType, - vllm_config: VllmConfig | None = None, + self, notification_type: EEPNotificationType ): """ Send notifications to EngineCoreClient, which can then forward the notifications to other engine core processes. It is used for: - 1) In scale up: new core engines to notify existing core engines - that they are ready; - 2) In scale down: removing core engines to notify EngineCoreClient + 1) In scale down: removing core engines to notify EngineCoreClient so EngineCoreClient can release their ray placement groups; - 3) Both scale up/down: to notify EngineCoreClient that existing + 2) Both scale up/down: to notify EngineCoreClient that existing core engines have already switched to the new parallel setup. """ - if vllm_config is None: - dp_rank = self.vllm_config.parallel_config.data_parallel_rank - else: - dp_rank = vllm_config.parallel_config.data_parallel_rank + dp_rank = self.vllm_config.parallel_config.data_parallel_rank notification_data = (notification_type.value, dp_rank) outputs = EngineCoreOutputs( utility_output=UtilityOutput( @@ -2257,22 +2262,11 @@ class DPEngineCoreProc(EngineCoreProc): ): socket.send_multipart(encoder.encode(outputs)) - def eep_handle_engine_core_notification( - self, notification_type: str | EEPNotificationType - ): - """ - Handle notification received from EngineCoreClient - (forwarded from new core engines). - """ - assert self.eep_scaling_state is not None - if isinstance(notification_type, str): - notification_type = EEPNotificationType(notification_type) - self.eep_scaling_state.handle_notification(notification_type) - def _eep_scale_up_before_kv_init(self): from vllm.distributed.elastic_ep.elastic_state import ElasticEPScalingState - self.eep_scaling_state = ElasticEPScalingState( + self.ignore_start_dp_wave = True + state = ElasticEPScalingState( model_executor=self.model_executor, engine_core=self, vllm_config=self.vllm_config, @@ -2281,7 +2275,10 @@ class DPEngineCoreProc(EngineCoreProc): scale_type="scale_up", reconfig_request=None, ) - self.eep_scaling_state.run_pre_kv_init_states() + if self.eep_scaling_state is not None: + raise RuntimeError("Elastic EP reconfiguration is already active") + self.eep_scaling_state = state + state.run_pre_kv_init_states() self.process_input_queue_block = False diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index a6c232b2ab7..9460fdf48f7 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -225,7 +225,10 @@ class EngineCoreClient(ABC): running state.""" raise NotImplementedError - async def scale_elastic_ep(self, new_data_parallel_size: int) -> None: + async def commit_elastic_ep(self) -> None: + raise NotImplementedError + + async def prepare_elastic_ep(self, new_data_parallel_size: int) -> None: raise NotImplementedError async def get_output_async(self) -> EngineCoreOutputs: @@ -1490,6 +1493,7 @@ class DPLBAsyncMPClient(DPAsyncMPClient): ) assert len(self.core_engines) > 1 + self._prepared_elastic_ep: tuple[int, int] | None = None self.eng_start_index = ( len(self.core_engines) * self.client_index @@ -1603,31 +1607,16 @@ class DPLBAsyncMPClient(DPAsyncMPClient): if len(cache.pending_notifications[notification_type]) >= abs( cache.num_new_core_engines ): - if notification_type == EEPNotificationType.SHUTDOWN_COMPLETE: - assert isinstance(self.resources.engine_manager, CoreEngineActorManager) - assert cache.num_new_core_engines < 0 - old_dp_size = len(cache.existing_core_engines) - new_dp_size = old_dp_size + cache.num_new_core_engines - self.resources.engine_manager.scale_down_elastic_ep( - old_dp_size, new_dp_size - ) - else: - await asyncio.gather( - *[ - self._call_utility_async( - "eep_handle_engine_core_notification", - notification_type, - engine=engine, - ) - for engine in cache.existing_core_engines - ] - ) - cache.pending_notifications[notification_type] = set() - if notification_type in [ - EEPNotificationType.SHUTDOWN_COMPLETE, - EEPNotificationType.NEW_CORE_ENGINES_WEIGHTS_INIT_READY, - ]: - self.eep_scaling_cache = None + engine_manager = self.resources.engine_manager + assert isinstance(engine_manager, CoreEngineActorManager) + assert cache.num_new_core_engines < 0 + old_dp_size = len(cache.existing_core_engines) + new_dp_size = old_dp_size + cache.num_new_core_engines + engine_manager.scale_down_elastic_ep(old_dp_size, new_dp_size) + self.vllm_config.parallel_config.data_parallel_size_local = len( + engine_manager.local_engine_actors + ) + self.eep_scaling_cache = None async def abort_requests_async(self, request_ids: list[str]) -> None: if not request_ids or self.resources.engine_dead: @@ -1651,31 +1640,46 @@ class DPLBAsyncMPClient(DPAsyncMPClient): ) -> None: await self._send_input(EngineCoreRequestType.ABORT, request_ids, engine) - async def scale_elastic_ep(self, new_data_parallel_size: int) -> None: - """Scale elastic EP data parallel size""" + async def commit_elastic_ep(self) -> None: + """Commit prepared elastic EP scaling.""" + prepared = self._prepared_elastic_ep + if prepared is None: + raise RuntimeError("Elastic EP scaling has not been prepared") + new_data_parallel_size, num_redundant_experts = prepared cur_data_parallel_size = len(self.core_engines) - - assert new_data_parallel_size != cur_data_parallel_size, ( - f"new_data_parallel_size {new_data_parallel_size} must be " - f"different from cur_data_parallel_size {cur_data_parallel_size}" + if new_data_parallel_size > cur_data_parallel_size: + await self._commit_scale_up_elastic_ep(new_data_parallel_size) + else: + await self._commit_scale_down_elastic_ep(new_data_parallel_size) + self.vllm_config.parallel_config.eplb_config.num_redundant_experts = ( + num_redundant_experts ) + self._prepared_elastic_ep = None + async def prepare_elastic_ep(self, new_data_parallel_size: int) -> None: + """Prepare elastic EP scaling without routing requests to new engines.""" + if (prepared := self._prepared_elastic_ep) is not None: + if prepared[0] == new_data_parallel_size: + return + raise RuntimeError("Elastic EP scaling is already prepared") + cur_data_parallel_size = len(self.core_engines) assert self.vllm_config.parallel_config.data_parallel_backend == "ray", ( "Only ray DP backend supports scaling elastic EP" ) - - scale_up = new_data_parallel_size > cur_data_parallel_size - - if scale_up: - await self._scale_up_elastic_ep( - cur_data_parallel_size, new_data_parallel_size - ) + parallel_config = self.vllm_config.parallel_config + num_experts = self.vllm_config.model_config.get_num_experts() + num_redundant_experts = ( + num_experts + parallel_config.eplb_config.num_redundant_experts + ) * new_data_parallel_size // cur_data_parallel_size - num_experts + if new_data_parallel_size < cur_data_parallel_size: + await self._prepare_scale_down_elastic_ep(new_data_parallel_size) else: - await self._scale_down_elastic_ep( - cur_data_parallel_size, new_data_parallel_size + await self._prepare_scale_up_elastic_ep( + new_data_parallel_size, num_redundant_experts ) + self._prepared_elastic_ep = new_data_parallel_size, num_redundant_experts - async def _eep_wait_for_setup_switch_complete(self) -> None: + def _eep_wait_for_setup_switch_complete(self) -> asyncio.Future: """ Wait for core engines to switch to the new setup. @@ -1687,9 +1691,26 @@ class DPLBAsyncMPClient(DPAsyncMPClient): future = asyncio.get_running_loop().create_future() self.utility_results[EEP_NOTIFICATION_CALL_ID] = future self._ensure_output_queue_task() - await future + return future - def _setup_elastic_ep_reconfig_bootstrap(self) -> tuple[str, int]: + def _wait_for_new_engine_ready(self, new_core_engines: list[bytes]) -> None: + new_engine_identities = set(new_core_engines) + sync_input_socket = zmq.Socket.shadow(self.input_socket) + while new_engine_identities: + if not sync_input_socket.poll(timeout=VLLM_ENGINE_READY_TIMEOUT_S * 1000): + raise TimeoutError( + f"Timed out waiting for new engine core processes to " + f"start. Waited " + f"{VLLM_ENGINE_READY_TIMEOUT_S}s (configured by " + f"VLLM_ENGINE_READY_TIMEOUT_S). To increase the " + f"timeout, set the environment variable: " + f"VLLM_ENGINE_READY_TIMEOUT_S=<seconds>" + ) + identity, payload = sync_input_socket.recv_multipart() + new_engine_identities.discard(identity) + self._apply_ready_response(payload) + + def _setup_elastic_ep_reconfig_bootstrap(self) -> None: from vllm.distributed.utils import create_tcp_store from vllm.utils.network_utils import get_open_ports_list @@ -1709,36 +1730,36 @@ class DPLBAsyncMPClient(DPAsyncMPClient): ) parallel_config._coord_store_port = store.port self._coord_store = store - return ip, store.port - async def _scale_up_elastic_ep( - self, cur_data_parallel_size: int, new_data_parallel_size: int - ) -> None: - """Scale up the data parallel size by creating new engine cores - and reconfiguring existing ones.""" - cur_data_parallel_size = len(self.core_engines) - - self.eep_scaling_cache = ElasticScalingCache( - existing_core_engines=self.core_engines.copy(), - num_new_core_engines=new_data_parallel_size - cur_data_parallel_size, - pending_notifications=dict(), + def _make_reconfig_request( + self, + new_data_parallel_size: int, + rank_type: ReconfigureRankType = ReconfigureRankType.KEEP_CURRENT_RANK, + ) -> ReconfigureDistributedRequest: + parallel_config = self.vllm_config.parallel_config + return ReconfigureDistributedRequest( + new_data_parallel_size=new_data_parallel_size, + new_data_parallel_rank=rank_type, + new_data_parallel_rank_local=ReconfigureRankType.KEEP_CURRENT_RANK, + new_data_parallel_master_ip=parallel_config.data_parallel_master_ip, + new_data_parallel_master_port=parallel_config.data_parallel_master_port, + new_data_parallel_master_port_list=parallel_config._data_parallel_master_port_list, + coord_store_port=parallel_config._coord_store_port, ) - parallel_config = self.vllm_config.parallel_config - ip, coord_store_port = self._setup_elastic_ep_reconfig_bootstrap() + async def _prepare_scale_up_elastic_ep( + self, + new_data_parallel_size: int, + num_redundant_experts: int, + ) -> None: + """Prepare scale up by creating new engine cores and reconfiguring + existing ones.""" + self._setup_elastic_ep_reconfig_bootstrap() # Phase 1: Send reconfig messages to existing engines reconfig_futures = [] for engine in self.core_engines: - reconfig_request = ReconfigureDistributedRequest( - new_data_parallel_size=new_data_parallel_size, - new_data_parallel_rank=ReconfigureRankType.KEEP_CURRENT_RANK, - new_data_parallel_rank_local=ReconfigureRankType.KEEP_CURRENT_RANK, - new_data_parallel_master_ip=ip, - new_data_parallel_master_port=parallel_config.data_parallel_master_port, - new_data_parallel_master_port_list=parallel_config._data_parallel_master_port_list, - coord_store_port=coord_store_port, - ) + reconfig_request = self._make_reconfig_request(new_data_parallel_size) coro = self._call_utility_async( "reinitialize_distributed", reconfig_request, engine=engine ) @@ -1746,51 +1767,54 @@ class DPLBAsyncMPClient(DPAsyncMPClient): # Phase 2: Create new engines assert isinstance(self.resources.engine_manager, CoreEngineActorManager) - parallel_config.eplb_config.num_redundant_experts = 0 start_new_worker_future = asyncio.to_thread( self.resources.engine_manager.scale_up_elastic_ep, self.vllm_config, new_data_parallel_size, + num_redundant_experts, ) - wait_future = self._eep_wait_for_setup_switch_complete() # Phase 3: Wait for new engines to be created # and reconfig messages to be received await asyncio.gather(start_new_worker_future, *reconfig_futures) + ready_keys = [future.result() for future in reconfig_futures] + ready_keys.extend( + f"eep_ready/{rank}" + for rank in range(len(self.core_engines), new_data_parallel_size) + ) + await asyncio.to_thread(self._coord_store.wait, ready_keys) logger.info("[Elastic EP] Successfully started new engines") - # Create new CoreEngine objects for the new engines - new_engine_identities = set() - for i in range(cur_data_parallel_size, new_data_parallel_size): - new_engine = i.to_bytes(2, "little") - self.core_engines.append(new_engine) - # NOTE(yongji): we don't update lb_engines here, - # we let run_engine_stats_update_task to update it. - new_engine_identities.add(new_engine) + async def _commit_scale_up_elastic_ep(self, new_data_parallel_size: int) -> None: + new_core_engines = [ + rank.to_bytes(2, "little") + for rank in range(len(self.core_engines), new_data_parallel_size) + ] - # Wait for ready messages from new engines on the input socket - sync_input_socket = zmq.Socket.shadow(self.input_socket) - while new_engine_identities: - if not sync_input_socket.poll( - timeout=VLLM_ENGINE_READY_TIMEOUT_S * 1000 # convert to ms - ): - raise TimeoutError( - f"Timed out waiting for new engine core processes to " - f"start. Waited " - f"{VLLM_ENGINE_READY_TIMEOUT_S}s (configured by " - f"VLLM_ENGINE_READY_TIMEOUT_S). To increase the " - f"timeout, set the environment variable: " - f"VLLM_ENGINE_READY_TIMEOUT_S=<seconds>" - ) - identity, payload = sync_input_socket.recv_multipart() - new_engine_identities.discard(identity) - self._apply_ready_response(payload) + await self.pause_scheduler_async(mode="keep", clear_cache=False) + wait_future = self._eep_wait_for_setup_switch_complete() + finish_futures = [ + asyncio.create_task( + self._call_utility_async("commit_prepared_elastic_ep", engine=engine) + ) + for engine in self.core_engines + ] + try: + await asyncio.gather(*finish_futures) + await wait_future + self._wait_for_new_engine_ready(new_core_engines) + except Exception: + wait_future.cancel() + raise - # NOTE(yongji): Before we schedule any requests on the new workers, - # we should wait for them to switch to the new setup. - await wait_future + self.core_engines.extend(new_core_engines) # Update the parallel config - self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size + parallel_config = self.vllm_config.parallel_config + parallel_config.data_parallel_size = new_data_parallel_size + if isinstance(self.resources.engine_manager, CoreEngineActorManager): + parallel_config.data_parallel_size_local = len( + self.resources.engine_manager.local_engine_actors + ) # Notify coordinator about scale up through existing # stats_update_task connection self._ensure_stats_update_task() @@ -1803,10 +1827,23 @@ class DPLBAsyncMPClient(DPAsyncMPClient): "[Elastic EP] Scale up completed, new data parallel size: %s", new_data_parallel_size, ) + await self.resume_scheduler_async() - async def _scale_down_elastic_ep( - self, cur_data_parallel_size: int, new_data_parallel_size: int - ) -> None: + async def _prepare_scale_down_elastic_ep(self, new_data_parallel_size: int) -> None: + self._setup_elastic_ep_reconfig_bootstrap() + + reconfig_futures = [] + for engine in self.core_engines[:new_data_parallel_size]: + reconfig_request = self._make_reconfig_request(new_data_parallel_size) + coro = self._call_utility_async( + "reinitialize_distributed", reconfig_request, engine=engine + ) + reconfig_futures.append(asyncio.create_task(coro)) + + ready_keys = await asyncio.gather(*reconfig_futures) + await asyncio.to_thread(self._coord_store.wait, ready_keys) + + async def _commit_scale_down_elastic_ep(self, new_data_parallel_size: int) -> None: """Scale down the data parallel size by shutting down and reconfiguring existing engine cores.""" cur_data_parallel_size = len(self.core_engines) @@ -1817,50 +1854,51 @@ class DPLBAsyncMPClient(DPAsyncMPClient): pending_notifications=dict(), ) - parallel_config = self.vllm_config.parallel_config - ip, coord_store_port = self._setup_elastic_ep_reconfig_bootstrap() - + old_core_engines = self.core_engines + # NOTE(yongji): Immediately stop sending requests to the removing engines. + self.core_engines = old_core_engines[:new_data_parallel_size] + self.lb_engines = self.lb_engines[:new_data_parallel_size] removed_dp_size = cur_data_parallel_size - new_data_parallel_size + pause_modes = ["keep"] * new_data_parallel_size + ["abort"] * removed_dp_size + pause_futures = [ + self._call_utility_async("pause_scheduler", mode, False, engine=engine) + for mode, engine in zip(pause_modes, old_core_engines) + ] + await asyncio.gather(*pause_futures) assert isinstance(self.resources.engine_manager, CoreEngineActorManager) self.resources.engine_manager.remove_run_refs_for_scale_down(removed_dp_size) + wait_future = self._eep_wait_for_setup_switch_complete() reconfig_futures = [] - for cur_dp_rank, engine in enumerate(self.core_engines): - reconfig_request = ReconfigureDistributedRequest( - new_data_parallel_size=new_data_parallel_size, - new_data_parallel_rank=ReconfigureRankType.KEEP_CURRENT_RANK, - new_data_parallel_rank_local=ReconfigureRankType.KEEP_CURRENT_RANK, - new_data_parallel_master_ip=ip, - new_data_parallel_master_port=parallel_config.data_parallel_master_port, - new_data_parallel_master_port_list=parallel_config._data_parallel_master_port_list, - coord_store_port=coord_store_port, - ) - if cur_dp_rank >= new_data_parallel_size: - reconfig_request.new_data_parallel_rank = ( - ReconfigureRankType.SHUTDOWN_CURRENT_RANK + for cur_dp_rank, engine in enumerate(old_core_engines): + if cur_dp_rank < new_data_parallel_size: + coro = self._call_utility_async( + "commit_prepared_elastic_ep", engine=engine + ) + else: + reconfig_request = self._make_reconfig_request( + new_data_parallel_size, + ReconfigureRankType.SHUTDOWN_CURRENT_RANK, + ) + coro = self._call_utility_async( + "reinitialize_distributed", reconfig_request, engine=engine ) - coro = self._call_utility_async( - "reinitialize_distributed", reconfig_request, engine=engine - ) reconfig_futures.append(asyncio.create_task(coro)) - # NOTE(yongji): Immediately stop sending requests to the removing engines. - self.core_engines = self.core_engines[:new_data_parallel_size] - self.lb_engines = self.lb_engines[:new_data_parallel_size] - wait_future = self._eep_wait_for_setup_switch_complete() + try: + await asyncio.gather(*reconfig_futures) - await asyncio.gather(*reconfig_futures) + self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size + self._ensure_stats_update_task() + scale_down_marker = msgspec.msgpack.encode( + ("SCALE_ELASTIC_EP", new_data_parallel_size) + ) + await self.first_req_send_socket.send(scale_down_marker) + await wait_future + await self.resume_scheduler_async() + except Exception: + wait_future.cancel() + raise - self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size - self._ensure_stats_update_task() - scale_down_marker = msgspec.msgpack.encode( - ("SCALE_ELASTIC_EP", new_data_parallel_size) - ) - await self.first_req_send_socket.send(scale_down_marker) - - # NOTE(yongji): Unlike scaling up, - # here we don't actually need to wait for the setup switch to complete. - # We may want to remove it in the future. - await wait_future logger.info( "[Elastic EP] Scale down completed, new data parallel size: %s", new_data_parallel_size, diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index db1896b0946..9b3bea0db9c 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -821,7 +821,10 @@ class CoreEngineActorManager: return placement_groups, local_dp_ranks def scale_up_elastic_ep( - self, cur_vllm_config: VllmConfig, new_data_parallel_size: int + self, + cur_vllm_config: VllmConfig, + new_data_parallel_size: int, + num_redundant_experts: int, ) -> None: import copy @@ -864,6 +867,9 @@ class CoreEngineActorManager: if new_data_parallel_size > 1: _apply_dp_identity_suffix(dp_vllm_config, rank) dp_vllm_config.parallel_config.data_parallel_size = new_data_parallel_size + dp_vllm_config.parallel_config.eplb_config.num_redundant_experts = ( + num_redundant_experts + ) dp_vllm_config.parallel_config.placement_group = pg # Check if this placement group is on the head node @@ -906,39 +912,18 @@ class CoreEngineActorManager: self.created_placement_groups.append(pg) self.placement_group_is_local.append(local_client) - ray.get( - [ - actor.wait_for_init.remote() - for actor in ( - self.local_engine_actors[-new_local_engines:] - if new_local_engines > 0 - else [] - ) - + self.remote_engine_actors[ - -(len(placement_groups) - new_local_engines) : - ] - ] - ) - actors = ( self.local_engine_actors[-new_local_engines:] if new_local_engines > 0 else [] ) + self.remote_engine_actors[-(len(placement_groups) - new_local_engines) :] + ray.get([actor.wait_for_init.remote() for actor in actors]) for actor in actors: ref = actor.run.remote() self.run_refs.append(ref) self.actor_run_ref_dict[actor] = ref - cur_vllm_config.parallel_config.data_parallel_size = new_data_parallel_size - # Update old_vllm_config with new data_parallel_size_local if any new - # local engines were added - if new_local_engines > 0: - cur_vllm_config.parallel_config.data_parallel_size_local += ( - new_local_engines - ) - def scale_down_elastic_ep( self, cur_data_parallel_size: int, new_data_parallel_size: int ) -> None: diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py index 4063844d469..404acd50de9 100644 --- a/vllm/v1/executor/abstract.py +++ b/vllm/v1/executor/abstract.py @@ -116,11 +116,11 @@ class Executor(ABC): raise NotImplementedError def initialize_from_config(self, kv_cache_configs: list[KVCacheConfig]) -> None: - """ - Initialize the KV caches and begin the model execution loop of the - underlying workers. - """ + """Initialize the KV caches on the underlying workers.""" self.collective_rpc("initialize_from_config", args=(kv_cache_configs,)) + + def compile_or_warm_up_model(self) -> None: + """Compile/warm up the model and capture cudagraphs on workers.""" compilation_times: list[CompilationTimes] = self.collective_rpc( "compile_or_warm_up_model" ) From 948107acf7ef8813b8ec94fff7c5ab62d4aba9ae Mon Sep 17 00:00:00 2001 From: Xin He <xin3.he@intel.com> Date: Tue, 28 Jul 2026 20:37:55 +0800 Subject: [PATCH 167/185] [Bugfix] Enhance extra_config handling for layer name suffix matching (#48589) Signed-off-by: Xin He <xin3.he@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com> --- tests/quantization/test_auto_round.py | 49 +++++++++++++++++++ .../layers/quantization/inc/config_parser.py | 8 +++ 2 files changed, 57 insertions(+) diff --git a/tests/quantization/test_auto_round.py b/tests/quantization/test_auto_round.py index 732080fc967..59c6a1e326b 100644 --- a/tests/quantization/test_auto_round.py +++ b/tests/quantization/test_auto_round.py @@ -243,6 +243,27 @@ def test_inc_config_parser_parallel_lm_head_defaults_to_unquantized() -> None: assert layer_config.bits == 16 +def test_inc_config_parser_suffix_match_for_lm_head() -> None: + """Short extra_config key should match fully-qualified lm_head layer name.""" + layer = object.__new__(ParallelLMHead) + config = make_config( + extra_config={ + "lm_head": { + "bits": 4, + "group_size": 128, + "sym": True, + } + } + ) + + layer_config = config.config_parser.resolve(layer, "model.language_model.lm_head") + + assert layer_config.quantized is True + assert layer_config.bits == 4 + assert layer_config.group_size == 128 + assert layer_config.sym is True + + def test_inc_config_parser_fused_moe_requires_consistent_configs() -> None: config = make_config( extra_config={ @@ -790,6 +811,34 @@ def test_inc_get_quant_method_linear_uses_resolved_scheme(monkeypatch) -> None: assert method is sentinel +def test_inc_get_quant_method_lm_head_uses_suffix_match(monkeypatch) -> None: + """lm_head extra_config should apply to fully-qualified prefix.""" + config = make_config( + extra_config={ + "lm_head": { + "bits": 4, + "group_size": 128, + "sym": True, + } + } + ) + layer = object.__new__(ParallelLMHead) + sentinel = object() + + class DummyScheme: + def get_linear_method(self, _config, _layer, _prefix, _layer_config): + return sentinel + + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.inc.schemes.factory.resolve_scheme", + lambda _layer_config: DummyScheme(), + ) + + method = config.get_quant_method(layer, "model.language_model.lm_head") + + assert method is sentinel + + def test_inc_get_quant_method_moe_uses_resolved_scheme(monkeypatch) -> None: config = make_config() layer = object.__new__(RoutedExperts) diff --git a/vllm/model_executor/layers/quantization/inc/config_parser.py b/vllm/model_executor/layers/quantization/inc/config_parser.py index 603b80b7cd0..6e94cad2cc6 100644 --- a/vllm/model_executor/layers/quantization/inc/config_parser.py +++ b/vllm/model_executor/layers/quantization/inc/config_parser.py @@ -142,6 +142,14 @@ class INCConfigParser: if self._config.extra_config and layer_name in self._config.extra_config: return get_config(layer_name) + # Suffix match: handle cases where extra_config keys use short names + # (e.g. "lm_head") but the layer_name is fully qualified + # (e.g. "model.language_model.lm_head") due to model nesting. + if self._config.extra_config: + for cfg_key in self._config.extra_config: + if layer_name.endswith(f".{cfg_key}"): + return get_config(cfg_key) + quantized = not isinstance(layer, ParallelLMHead) if self._config.block_name_to_quantize: quantized = any( From 98e91a9600eb75b2de14ef27f13b10088d1a1279 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= <nicolo.lucchesi@mistral.ai> Date: Tue, 28 Jul 2026 14:57:12 +0200 Subject: [PATCH 168/185] [PD][NixlPush] Skip extra `add_remote_agent` step in D->P handshake (#49345) Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai> --- .../kv_transfer/kv_connector/v1/nixl/base_worker.py | 3 ++- .../kv_transfer/kv_connector/v1/nixl/push_worker.py | 6 ++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index d4e516d6df7..9063308ce2e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -700,8 +700,9 @@ class NixlBaseConnectorWorker: ) setup_agent_time = time.perf_counter() logger.debug( - "NIXL handshake: add agent took: %s", + "NIXL handshake: add agent took: %s (notif_agents_only=%s)", setup_agent_time - got_metadata_time, + notif_agents_only, ) remote_ranks = (remote_pp_rank, remote_rank) remote_rank_to_agent_name[remote_ranks] = remote_agent_name diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 43b4f893907..8cb710e7e58 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -299,8 +299,10 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): reg_data["remote_port"], reg_data["remote_tp_size"], pp_size=remote_pp_size, - # D never addresses P memory in push mode; just load P's agents. - notif_agents_only=remote_pp_size > 1, + # D only ever sends PUSH_REG notifs to P and never reads or writes + # P's memory in push mode, so it never needs the transfer + # descriptors set up by the full add_remote_agent path. + notif_agents_only=True, ) if fut is None: self._do_send_reg_notif(req_id, reg_data) From 9b9fc4039c25a6e4fe0ae97361b62edd74b8b47e Mon Sep 17 00:00:00 2001 From: liangel-02 <liangel@meta.com> Date: Tue, 28 Jul 2026 07:26:55 -0600 Subject: [PATCH 169/185] add epilogue hook to flex attention (#45841) Signed-off-by: Angel Li <liangel@meta.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Michael Goin <mgoin64@gmail.com> --- vllm/v1/attention/backends/flex_attention.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/vllm/v1/attention/backends/flex_attention.py b/vllm/v1/attention/backends/flex_attention.py index e7fe5852194..83144751aeb 100644 --- a/vllm/v1/attention/backends/flex_attention.py +++ b/vllm/v1/attention/backends/flex_attention.py @@ -12,6 +12,7 @@ import torch import torch._dynamo.decorators import torch.nn.functional as F from torch.nn.attention.flex_attention import ( + AuxRequest, BlockMask, _mask_mod_signature, _score_mod_signature, @@ -1228,6 +1229,9 @@ class FlexAttentionImpl(AttentionImpl): if block_n is not None: self.block_n = block_n + # Optional post-attention epilogue transform + self.out_transform = kwargs.get("out_transform") + @staticmethod def view_as_4d(tensor: torch.Tensor) -> torch.Tensor: """View a 3d tensor as 4D.""" @@ -1392,8 +1396,13 @@ class FlexAttentionImpl(AttentionImpl): self.scale, enable_gqa=enable_gqa, kernel_options=kernel_options, + return_aux=AuxRequest(lse=True) if self.out_transform is not None else None, ) + if self.out_transform is not None: + out, aux = out + out = self.out_transform(out, aux.lse) + # Flex doesn't have an out variant today, rely on epilogue fusion out = out.permute(0, 2, 1, 3).squeeze(0) output[:num_actual_tokens, :, :].copy_(out) From 601fa9a74e8cd2ae0a4fa127d68b91bbf363a24e Mon Sep 17 00:00:00 2001 From: Nick Hill <nickhill123@gmail.com> Date: Tue, 28 Jul 2026 06:45:02 -0700 Subject: [PATCH 170/185] [KV Connector] Support NIXL heterogeneous P/D block sizes for hybrid models (#49612) Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> --- .../kv_connector/unit/test_nixl_connector.py | 10 +- .../unit/test_nixl_connector_hma.py | 114 ++++ .../unit/test_nixl_desc_geometry.py | 619 ++++++++++++++++++ tests/v1/kv_connector/unit/test_tp_mapping.py | 54 ++ .../kv_connector/v1/nixl/base_worker.py | 271 ++++++-- .../kv_connector/v1/nixl/pull_worker.py | 35 +- .../kv_connector/v1/nixl/push_worker.py | 22 +- 7 files changed, 1014 insertions(+), 111 deletions(-) create mode 100644 tests/v1/kv_connector/unit/test_nixl_desc_geometry.py diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index d58f117fa31..c474077ab22 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -479,8 +479,9 @@ class FakeNixlConnectorWorker(NixlConnectorWorker): super().__init__(*args, kv_cache_config=kv_cache_config, **kwargs) self._hand_shake_latency = hand_shake_latency self.kv_cache_layout = kv_cache_layout - # Mock register_kv_caches attribute needed for tests that do not call it. + # Mock register_kv_caches attributes needed for tests that do not call it. self.src_xfer_handles_by_block_size = {self.block_size: 1} + self.src_blocks_data = np.empty((0, 3), dtype=np.uint64) test_shape = self.attn_backends[0].get_kv_cache_shape( num_blocks=1, block_size=16, num_kv_heads=1, head_size=1 ) @@ -765,8 +766,9 @@ class TestNixlHandshake: assert remote_info.remote_tp_size == remote_tp_size assert -tp_ratio == worker.transfer_topo.tp_ratio(remote_tp_size) # ensure src_xfer_handles_by_tp_ratio is populated with tpratio chunks - assert -tp_ratio in worker.src_xfer_handles_by_tp_ratio - assert len(worker.src_xfer_handles_by_tp_ratio[-tp_ratio]) == tp_ratio + split_key = (-tp_ratio, worker.block_size) + assert split_key in worker.src_xfer_handles_by_tp_ratio + assert len(worker.src_xfer_handles_by_tp_ratio[split_key]) == tp_ratio assert remote_engine_id in worker.dst_xfer_side_handles assert set(worker.dst_xfer_side_handles[remote_engine_id].keys()) == set( range(tp_ratio) @@ -2091,7 +2093,7 @@ def test_shutdown_cleans_up_resources(default_vllm_config, dist_init): # Mock register_kv_cache which registers local handle worker.src_xfer_handles_by_block_size = {worker.block_size: 455} # P TP = 2 * D TP case, we should register 2 local handles - worker.src_xfer_handles_by_tp_ratio = {-2: [456, 457]} + worker.src_xfer_handles_by_tp_ratio = {(-2, 16): [456, 457]} worker.dst_xfer_side_handles = {"engine1": {0: 789}} worker._remote_agents = {"engine1": {(0, 0): "agent1"}} # _cleanup_remote_engine (called by shutdown) also clears these: diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 4945942ba3a..1f7a62d2c9a 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -708,6 +708,120 @@ def test_get_block_descs_ids_kernel_block_mismatch(): assert list(result) == expected, f"Expected {expected}, got {list(result)}" +@pytest.mark.cpu_test +def test_get_block_descs_ids_hetero_block_size_hybrid(): + """With a block-size ratio, FA desc ids are ratio-expanded while SSM + desc ids keep the unexpanded logical stride (state blocks are never + sub-split).""" + from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec + + worker = _make_mock_worker_for_desc_ids( + num_regions=2, + has_mamba=True, + group_spec_types=(FullAttentionSpec, MambaSpec), + block_len_per_layer=[100], + ) + + ratio = 4 + # FA ids are already remote-granularity (expanded) sub-block ids. + fa_sub_blocks = [3, 5] + ssm_blocks = [1] + result = worker._compute_desc_ids( + block_ids=(fa_sub_blocks, ssm_blocks), + dst_num_blocks=100, + block_size_ratio=ratio, + physical_blocks_per_logical=1, + ) + + # FA regions have 100*4 entries each; SSM regions (4 per layer) start at + # 2*400 and stride by the unexpanded 100 logical blocks. + expected = [3, 5, 403, 405, 801, 901, 1001, 1101] + assert list(result) == expected, f"Expected {expected}, got {list(result)}" + + +def _bind_worker_method(worker, name): + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + + method = getattr(NixlConnectorWorker, name) + setattr(worker, name, method.__get__(worker, NixlConnectorWorker)) + + +@pytest.mark.cpu_test +def test_map_block_ids_for_block_size_ratio_hybrid(): + """Attention groups expand to remote granularity and clip to the remote + coverage; mamba state blocks pass through 1:1.""" + from unittest.mock import MagicMock + + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec + + worker = MagicMock(spec=NixlConnectorWorker) + worker._group_spec_types = (FullAttentionSpec, MambaSpec) + _bind_worker_method(worker, "get_mapped_blocks") + _bind_worker_method(worker, "_map_block_ids_for_block_size_ratio") + + local, remote = worker._map_block_ids_for_block_size_ratio( + [[1, 2, 3], [7]], + [list(range(30, 40)), [42]], + 4, + ) + # [1, 2, 3] expand to sub-blocks [4..15], clipped to the 10 remote blocks. + assert local == [list(range(4, 14)), [7]] + assert remote == [list(range(30, 40)), [42]] + + # Attention-only full prefix hit: empty local list is preserved. + worker._group_spec_types = (FullAttentionSpec,) + local, remote = worker._map_block_ids_for_block_size_ratio([[]], [[30, 31]], 4) + assert local == [] + + +@pytest.mark.cpu_test +def test_post_process_zeroes_untransferred_tail(): + """The untransferred sub-blocks of the last local block are zeroed on + receive; mamba state caches are untouched by the attention permute.""" + from unittest.mock import MagicMock + + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec + + ratio = 4 + block_tokens = 8 # 2 tokens per remote sub-block + + worker = MagicMock(spec=NixlConnectorWorker) + worker._group_spec_types = (FullAttentionSpec, MambaSpec) + worker.transfer_topo = MagicMock() + worker.device_type = "cpu" + worker.enable_permute_local_kv = False + attn_cache = torch.ones(6, block_tokens, 2, 4) + mamba_cache = torch.ones(6, 16) + worker.device_kv_caches = {"attn.0": attn_cache, "mamba.0": mamba_cache} + fa_group = MagicMock(layer_names=["attn.0"]) + ssm_group = MagicMock(layer_names=["mamba.0"]) + worker.kv_cache_config = MagicMock(kv_cache_groups=[fa_group, ssm_group]) + # The cached property filters mamba layers out of the permuted caches. + attn_caches = NixlConnectorWorker._attention_kv_caches.func(worker) + assert len(attn_caches) == 1 and attn_caches[0] is attn_cache + worker._attention_kv_caches = attn_caches + _bind_worker_method(worker, "post_process_device_kv_on_receive") + + # Request occupies blocks [2, 3]; only 6 of 8 sub-blocks were received. + worker.post_process_device_kv_on_receive(ratio, [([2, 3], 6)]) + + # Block 2 fully covered; block 3 covered for 2 sub-blocks (4 tokens). + assert torch.all(attn_cache[2] == 1) + assert torch.all(attn_cache[3, :4] == 1) + assert torch.all(attn_cache[3, 4:] == 0) + # Untouched blocks and the mamba cache keep their content. + assert torch.all(attn_cache[4] == 1) + assert torch.all(mamba_cache == 1) + + @pytest.mark.cpu_test def test_nixl_metadata_hybrid_ssm_block_ids(): """Test NixlConnectorMetadata correctly stores block IDs for FA + SSM diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py new file mode 100644 index 00000000000..e2fc41a0229 --- /dev/null +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -0,0 +1,619 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end NIXL descriptor geometry invariants for hybrid MLA+SSM models +under heterogeneous P/D block geometry (TP-sharded KDA-style state, so +the mamba-aligned logical block size differs between P and D while the +kernel-granularity pages stay equal). + +The invariant under test: every LOCAL byte range a request's READ transfers +into must lie within that request's own blocks. A violation means an +incoming transfer can overwrite a co-resident request's KV or mamba state +mid-decode (silent corruption of an unrelated request). +""" + +from unittest.mock import patch + +import numpy as np +import pytest +import torch + +from .utils import create_vllm_config + + +class _RecordingNixl: + """Minimal NIXL wrapper stand-in that records descriptor lists and + prepared transfers so tests can resolve desc ids to byte ranges.""" + + def __init__(self, *args, **kwargs): + self.dlists: dict[int, np.ndarray] = {} + self.xfers: list[tuple] = [] + self._next_handle = 1 + + def get_reg_descs(self, caches_data, mem_type): + return caches_data + + def register_memory(self, descs, backends=None): + pass + + def deregister_memory(self, descs): + pass + + def get_agent_metadata(self): + return b"agent-meta" + + def get_xfer_descs(self, blocks_data, mem_type): + return blocks_data + + def prep_xfer_dlist(self, agent, descs): + handle = self._next_handle + self._next_handle += 1 + self.dlists[handle] = np.asarray(descs, dtype=np.uint64).reshape(-1, 3) + return handle + + def add_remote_agent(self, metadata): + return "remote-agent" + + def make_prepped_xfer( + self, op, local_handle, local_ids, remote_handle, remote_ids, notif_msg=None + ): + handle = self._next_handle + self._next_handle += 1 + self.xfers.append( + ( + op, + local_handle, + np.asarray(local_ids), + remote_handle, + np.asarray(remote_ids), + ) + ) + return handle + + def transfer(self, handle): + pass + + def check_xfer_state(self, handle): + return "DONE" + + def get_xfer_telemetry(self, handle): + from types import SimpleNamespace + + return SimpleNamespace( + xferDuration=1.0, postDuration=1.0, totalBytes=1, descCount=1 + ) + + def release_xfer_handle(self, handle): + pass + + def release_dlist_handle(self, handle): + pass + + def send_notif(self, agent, notif_msg=None): + pass + + def get_new_notifs(self): + return {} + + def remove_remote_agent(self, agent): + pass + + +def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blocks): + """Build a real pull worker with a hybrid MLA + 2xKDA HMA layout.""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl import ( + base_worker as bw, + ) + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum + from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupSpec, + KVCacheTensor, + MambaSpec, + MLAAttentionSpec, + ) + + mla_spec = MLAAttentionSpec( + block_size=local_block_size, + num_kv_heads=1, + head_size=6, + dtype=torch.float16, + ) + unified_page = mla_spec.page_size_bytes + kda_spec = MambaSpec( + block_size=local_block_size, + shapes=((8, 3), (1, 4, 4)), + dtypes=(torch.float16, torch.float32), + page_size_padded=unified_page, + mamba_type=MambaAttentionBackendEnum.GDN_ATTN, + ) + kv_cache_config = KVCacheConfig( + num_blocks=num_logical_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=num_logical_blocks * unified_page, + shared_by=[f"mla.{i}", f"kda_a.{i}", f"kda_b.{i}"], + ) + for i in range(2) + ], + kv_cache_groups=[ + KVCacheGroupSpec(["mla.0", "mla.1"], mla_spec), + KVCacheGroupSpec(["kda_a.0", "kda_a.1"], kda_spec), + KVCacheGroupSpec(["kda_b.0", "kda_b.1"], kda_spec), + ], + ) + + vllm_config = create_vllm_config(block_size=local_block_size) + vllm_config.cache_config.enable_prefix_caching = False + # kv_buffer_device defaults to the *real* platform's device type, which on + # a CPU-only test host would make this a host-buffer worker: host xfer + # buffers are per-layer, so the HMA shared-tensor regions this test builds + # would not be deduplicated. Pin it to the faked device type. + vllm_config.kv_transfer_config.kv_buffer_device = "cuda" + + from unittest.mock import MagicMock + + fake_backend = MagicMock() + fake_backend.get_supported_kernel_block_sizes.return_value = [kernel_block_size] + fake_backend.get_name.return_value = "FLASHMLA" + fake_backend.full_cls_name.return_value = "fake.FLASHMLA" + fake_platform = MagicMock() + fake_platform.device_type = "cuda" + fake_platform.get_nixl_memory_type.return_value = "VRAM" + + from vllm.config import set_current_vllm_config + + with ( + patch.object(bw, "NixlWrapper", _RecordingNixl), + patch.object(bw, "get_tensor_model_parallel_rank", return_value=0), + patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1), + patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]), + patch.object(bw, "current_platform", fake_platform), + patch( + "vllm.model_executor.layers.mamba.mamba_utils.get_conv_state_layout", + return_value="DS", + ), + set_current_vllm_config(vllm_config), + ): + worker = NixlConnectorWorker(vllm_config, "local-engine", kv_cache_config) + worker.use_mla = True + + # Attention caches are kernel-block granular on dim 0, as the + # receive post-process assumes. + ppl = local_block_size // kernel_block_size + tensors = [ + torch.zeros( + num_logical_blocks * ppl, unified_page // ppl, dtype=torch.uint8 + ) + for _ in range(2) + ] + worker.register_kv_caches( + { + "kda_a.0": tensors[0], + "mla.0": tensors[0], + "kda_b.0": tensors[0], + "kda_a.1": tensors[1], + "mla.1": tensors[1], + "kda_b.1": tensors[1], + } + ) + # Keep tensors alive alongside the worker; flat views for byte checks. + worker._test_tensors = [t.view(-1) for t in tensors] + worker._test_tensors_2d = tensors + worker._test_unified_page = unified_page + return worker + + +def _make_remote_meta( + worker, + remote_block_size, + remote_kernel_block_size, + remote_num_logical, + remote_ssm_sizes, +): + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + NixlAgentMetadata, + ) + + remote_ppl = remote_block_size // remote_kernel_block_size + # Kernel-granularity pages are TP-independent for MLA hybrids and must + # match the local ones for the handshake to pass, scaled down by the + # block-size ratio when the remote's kernel block is smaller. + block_size_ratio = worker.block_size // remote_kernel_block_size + kernel_page = worker.block_len_per_layer[0] // block_size_ratio + return NixlAgentMetadata( + engine_id="remote-engine", + agent_metadata=b"remote-agent-meta", + device_id=0, + kv_caches_base_addr=[0x10_000_000, 0x20_000_000], + num_blocks=remote_num_logical * remote_ppl, + block_lens=[kernel_page, kernel_page], + kv_cache_layout=worker.kv_cache_layout, + block_size=remote_kernel_block_size, + ssm_sizes=remote_ssm_sizes, + attn_backend_name=worker.backend_name, + physical_blocks_per_logical_kv_block=remote_ppl, + ) + + +def _owned_byte_ranges(worker, group_logical_ids): + """Byte ranges owned by a request: for each HMA region tensor, every + logical block id of every group maps to one unified page.""" + unified_page = worker._test_unified_page + bases = [t.data_ptr() for t in worker._test_tensors] + owned = [] + for base in bases: + for ids in group_logical_ids: + for b in ids: + owned.append((base + b * unified_page, base + (b + 1) * unified_page)) + return owned + + +def _assert_local_writes_within(worker, owned_ranges): + nixl = worker.nixl_wrapper + assert nixl.xfers, "no transfers were posted" + violations = [] + total_descs = 0 + for op, local_handle, local_ids, _, remote_ids in nixl.xfers: + assert len(local_ids) == len(remote_ids) + desc_arr = nixl.dlists[local_handle] + for i in local_ids: + addr, length, _dev = desc_arr[int(i)] + addr, length = int(addr), int(length) + total_descs += 1 + if not any(lo <= addr and addr + length <= hi for lo, hi in owned_ranges): + violations.append((int(i), hex(addr), length)) + assert not violations, ( + f"{len(violations)}/{total_descs} local descriptors write outside " + f"the request's own blocks: {violations[:10]}" + ) + return total_descs + + +@pytest.mark.cpu_test +def test_hetero_ppl_multi_read_writes_stay_within_request_blocks(): + """MLA-hybrid hetero geometry: local (D, TP1) logical blocks of 12 tokens + (kernel 4, ppl=3) vs remote (P, TP2) logical blocks of 8 tokens (ppl=2), + equal kernel pages, tp_ratio=-2 multi-read with replicated MLA and + TP-sharded KDA state. Every local descriptor of the request's reads must + stay within its own blocks.""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + NixlConnectorMetadata, + ) + + worker = _make_mla_hybrid_worker( + local_block_size=12, kernel_block_size=4, num_logical_blocks=8 + ) + assert worker._physical_blocks_per_logical_kv_block == 3 + + meta_r = _make_remote_meta( + worker, + remote_block_size=8, + remote_kernel_block_size=4, + remote_num_logical=12, + remote_ssm_sizes=(24, 32), + ) + for rank in (0, 1): + worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=2) + + # Request B: 17 matched tokens. Local: 2 logical blocks (24 tok + # capacity); remote: 16 prefilled tokens -> 2 remote logical blocks. + # Sparse, non-contiguous ids so neighbor blocks exist on all sides. + local_ids = ([2, 5], [1], [7]) + remote_ids = [[1, 4], [5], [2]] + + metadata = NixlConnectorMetadata() + metadata.add_new_req_to_recv( + request_id="req-b", + local_block_ids=local_ids, + kv_transfer_params={ + "remote_block_ids": remote_ids, + "remote_engine_id": "remote-engine", + "remote_request_id": "prefill-req-b", + "remote_host": "localhost", + "remote_port": 1234, + "tp_size": 2, + }, + ) + meta = metadata.reqs_to_recv["req-b"] + meta.local_physical_block_ids = worker._logical_to_kernel_block_ids( + meta.local_block_ids, worker._physical_blocks_per_logical_kv_block + ) + worker._recving_metadata["req-b"] = meta + + worker._read_blocks_for_req("req-b", meta) + + owned = _owned_byte_ranges(worker, local_ids) + total = _assert_local_writes_within(worker, owned) + # Multi-read: rank 0 carries the replicated MLA + its SSM shard, + # rank 1 carries only its SSM shard. + assert len(worker.nixl_wrapper.xfers) == 2 + assert total > 0 + + +def _resolve( + desc_arr, + idx, + bases, + region_size, + unified_page, + desc_page, + logical_ids_attn, + block_tokens, +): + """Resolve a desc id to (region, kind, token_start) where kind is 'attn' + (desc-page sized, sub-block-aligned, in the request's attention blocks) + or 'mamba'. token_start is the request-relative token offset, so local + and remote are comparable even when their kernel blocks differ in size.""" + addr, length, _ = (int(x) for x in desc_arr[int(idx)]) + for region, base in enumerate(bases): + off = addr - base + if 0 <= off < region_size: + b = off // unified_page + rem = off % unified_page + if length == desc_page and rem % desc_page == 0 and b in logical_ids_attn: + pos = logical_ids_attn.index(b) + tokens_per_desc = block_tokens * desc_page // unified_page + sub = rem // desc_page + return (region, "attn", pos * block_tokens + sub * tokens_per_desc) + return (region, "mamba", None) + raise AssertionError(f"desc {idx} addr {addr:#x} not in any region") + + +def _run_hetero_case( + local_block, kernel, remote_block, num_tokens, tp_size=2, remote_kernel=None +): + """Full pull-path run for one geometry; returns pairing records. + + ``remote_kernel`` defaults to the local kernel block size; a smaller + value additionally exercises block_size_ratio > 1. + """ + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + NixlConnectorMetadata, + ) + + remote_kernel = remote_kernel or kernel + block_size_ratio = kernel // remote_kernel + remote_ppl = remote_block // remote_kernel + matched = num_tokens - 1 # mamba N-1 rule + n_local = -(-num_tokens // local_block) + n_remote = -(-matched // remote_block) + + worker = _make_mla_hybrid_worker( + local_block_size=local_block, + kernel_block_size=kernel, + num_logical_blocks=max(2 * n_local + 4, 8), + ) + # Local KDA state pages are (48, 64) bytes; the remote holds 1/tp_size + # shards of each. + meta_r = _make_remote_meta( + worker, + remote_block_size=remote_block, + remote_kernel_block_size=remote_kernel, + remote_num_logical=max(2 * n_remote + 4, 8), + remote_ssm_sizes=(48 // tp_size, 64 // tp_size), + ) + for rank in range(tp_size): + worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=tp_size) + + # Sparse ids so neighbors exist between the request's blocks. + local_attn = [2 * i + 1 for i in range(n_local)] + remote_attn = [2 * i + 2 for i in range(n_remote)] + local_ids = (local_attn, [0], [2 * n_local + 2]) + remote_ids = [remote_attn, [1], [0]] + + metadata = NixlConnectorMetadata() + metadata.add_new_req_to_recv( + request_id="req-b", + local_block_ids=local_ids, + kv_transfer_params={ + "remote_block_ids": remote_ids, + "remote_engine_id": "remote-engine", + "remote_request_id": "prefill-req-b", + "remote_host": "localhost", + "remote_port": 1234, + "tp_size": tp_size, + }, + ) + meta = metadata.reqs_to_recv["req-b"] + meta.local_physical_block_ids = worker._logical_to_kernel_block_ids( + meta.local_block_ids, worker._physical_blocks_per_logical_kv_block + ) + worker._recving_metadata["req-b"] = meta + + # Sentinel-fill the local KV so untouched bytes are detectable. + for t in worker._test_tensors: + t.fill_(0xAA) + + worker._read_blocks_for_req("req-b", meta) + + # Invariant 1: all local writes within the request's own blocks. + owned = _owned_byte_ranges(worker, local_ids) + _assert_local_writes_within(worker, owned) + + # Invariant 2: local<->remote attention pairs are token-aligned. + nixl = worker.nixl_wrapper + local_bases = [t.data_ptr() for t in worker._test_tensors] + remote_bases = [0x10_000_000, 0x20_000_000] + local_unified = worker._test_unified_page + remote_unified = (local_unified // local_block) * remote_block + # With block_size_ratio > 1 the local page is split into ratio sub-descs, + # each the size of a whole remote kernel page. + desc_page = worker.block_len_per_layer[0] // block_size_ratio + meta_r_num_blocks_bytes = (meta_r.num_blocks // remote_ppl) * remote_unified + covered_tokens = set() + for op, lh, lids, rh, rids in nixl.xfers: + larr, rarr = nixl.dlists[lh], nixl.dlists[rh] + for li, ri in zip(lids, rids): + lreg, lkind, ltok = _resolve( + larr, + li, + local_bases, + len(worker._test_tensors[0]), + local_unified, + desc_page, + local_attn, + local_block, + ) + rreg, rkind, rtok = _resolve( + rarr, + ri, + remote_bases, + meta_r_num_blocks_bytes, + remote_unified, + desc_page, + remote_attn, + remote_block, + ) + assert lkind == rkind, ( + f"pair kind mismatch: local {lkind} vs remote {rkind} " + f"(local desc {li}, remote desc {ri})" + ) + assert lreg == rreg, ( + f"region mismatch: local {lreg} vs remote {rreg} for " + f"tokens {ltok} vs {rtok}" + ) + if lkind == "attn": + assert ltok == rtok, ( + f"TOKEN MISALIGNMENT: local sub-block holds tokens " + f"[{ltok}..) but receives remote tokens [{rtok}..) " + f"(geometry local_block={local_block}, " + f"remote_block={remote_block}, N={num_tokens})" + ) + covered_tokens.add(ltok) + + # Invariant 3: full coverage of the matched tokens, at the finest + # transfer granularity (the remote kernel block). + needed = {t for t in range(0, matched - matched % remote_kernel, remote_kernel)} + missing = needed - covered_tokens + assert not missing, ( + f"tokens never transferred: {sorted(missing)[:8]} " + f"(geometry local_block={local_block}, remote_block={remote_block}, " + f"N={num_tokens}, matched={matched})" + ) + + # Invariant 4: no stale bytes after receive completion. The scheduler + # excludes the blocks covering the matched tokens from alloc-time KV + # zeroing (the zeroing would race the RDMA write), so every byte of + # those blocks must be either written by the transfer or zeroed by the + # receive post-process. Stale bytes surface as mid-response garbage + # once decode grows into the untransferred tail. + for op, lh, lids, rh, rids in nixl.xfers: + larr = nixl.dlists[lh] + for li in lids: + addr, length, _ = (int(x) for x in larr[int(li)]) + for t in worker._test_tensors: + off = addr - t.data_ptr() + if 0 <= off < t.numel(): + t[off : off + length] = 0 # simulate the RDMA write + break + done_sending, done_recving = worker.get_finished() + assert "req-b" in done_recving + n_excluded = -(-matched // local_block) + stale = [] + for b in local_attn[:n_excluded]: + for region, t in enumerate(worker._test_tensors): + page = t[b * local_unified : (b + 1) * local_unified] + n_stale = int((page == 0xAA).sum()) + if n_stale: + stale.append((region, b, n_stale)) + assert not stale, ( + f"stale (unzeroed, untransferred) bytes in matched-range attention " + f"blocks (region, block, bytes): {stale} " + f"(geometry local_block={local_block}, remote_block={remote_block}, " + f"N={num_tokens}, matched={matched})" + ) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "local_block,remote_block", + [ + (12, 8), # ppl 3 vs 2 + (36, 8), # ppl 9 vs 2 (large ppl asymmetry, scaled) + (24, 4), # ppl 6 vs 1 + (16, 24), # remote larger than local (D_TP > P_TP direction) + ], +) +@pytest.mark.parametrize("num_tokens", list(range(2, 40))) +def test_hetero_ppl_token_alignment_sweep(local_block, remote_block, num_tokens): + """Sweep prompt lengths across block-boundary residues for several + hetero-ppl geometries; assert neighbor-safety, token alignment, and + coverage of every transferred kernel block.""" + _run_hetero_case( + local_block, kernel=4, remote_block=remote_block, num_tokens=num_tokens + ) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "num_tokens", + # Residues around the remote kernel block (4), the local kernel block + # (8), the remote logical block (8) and the local logical block (24). + [2, 5, 8, 9, 13, 16, 17, 21, 24, 25, 29, 32, 33, 41, 48, 49], +) +def test_hetero_ppl_with_block_size_ratio(num_tokens): + """Both hetero regimes at once: kernel blocks differ (local 8 / remote + 4, block_size_ratio=2) *and* physical_blocks_per_logical differs (3 vs + 2). The transfer is clipped at remote sub-block granularity by the + pairing and front-trimmed by _apply_prefix_caching, so the + untransferred tail can span both a partial block and whole blocks — + the case each of the two former zeroing paths handled only half of.""" + _run_hetero_case( + local_block=24, + kernel=8, + remote_block=8, + remote_kernel=4, + num_tokens=num_tokens, + ) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "num_tokens", + # Residues around every geometric boundary: kernel block (64), remote + # logical block (768), local logical block (5760), plus odd offsets. + [ + 2, + 63, + 64, + 65, + 127, + 128, + 300, + 640, + 767, + 768, + 769, + 831, + 832, + 1000, + 1535, + 1536, + 1537, + 2303, + 2304, + 2305, + 3001, + 5759, + 5760, + 5761, + 5824, + 6528, + 6529, + ], +) +def test_mla_hybrid_large_ppl_geometry(num_tokens): + """KimiLinear-scale MLA-hybrid geometry (TP8 prefill -> TP1 decode): + decode (local) logical block 5760 / kernel 64 (ppl=90), prefill + (remote) logical block 768 (ppl=12), tp_ratio=-8 multi-read with + replicated MLA and 8-way TP-sharded KDA state.""" + _run_hetero_case( + local_block=5760, + kernel=64, + remote_block=768, + num_tokens=num_tokens, + tp_size=8, + ) diff --git a/tests/v1/kv_connector/unit/test_tp_mapping.py b/tests/v1/kv_connector/unit/test_tp_mapping.py index 7c735098230..d72bbc9ec2d 100644 --- a/tests/v1/kv_connector/unit/test_tp_mapping.py +++ b/tests/v1/kv_connector/unit/test_tp_mapping.py @@ -161,3 +161,57 @@ class TestMambaPlanSplitHandles: # FA: chunk=200//1=200, slot=0 (skip_fa) → (1000, 200, 0), (2000, 200, 0) # SSM: chunk=400//2=200, idx=1 → (3200, 200, 0) assert splits[1] == [(1000, 200, 0), (2000, 200, 0), (3200, 200, 0)] + + def test_hetero_block_size_splits(self): + """With a block-size ratio, single-source FA sub-block descs pass + through whole; SSM descs are unexpanded and split per source.""" + plan = TPMapping( + source_ranks_per_group=((0,), (0, 1)), + all_source_ranks=(0, 1), + rank_to_attention_slot={0: 0, 1: 0}, + rank_offset_factor=0, + ) + + worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec)) + # 2 FA blocks x ratio 2 sub-blocks + 1 SSM desc (never expanded). + src_blocks_data = np.array( + [ + (1000, 100, 0), + (1100, 100, 0), + (2000, 100, 0), + (2100, 100, 0), + (3000, 400, 0), + ], + dtype=np.uint64, + ) + + splits = list(worker._build_local_splits_from_plan(plan, src_blocks_data, 4, 2)) + + assert len(splits) == 2 + fa_passthrough = [ + (1000, 100, 0), + (1100, 100, 0), + (2000, 100, 0), + (2100, 100, 0), + ] + assert splits[0] == fa_passthrough + [(3000, 200, 0)] + assert splits[1] == fa_passthrough + [(3200, 200, 0)] + + def test_hetero_block_size_head_sharded_asserts(self): + """Head-sharded FA reads (multiple FA sources) are incompatible with + a block-size mismatch and must fail loudly.""" + plan = TPMapping( + source_ranks_per_group=((0, 1), (0, 1)), + all_source_ranks=(0, 1), + rank_to_attention_slot={0: 0, 1: 1}, + rank_offset_factor=0, + ) + + worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec)) + src_blocks_data = np.array( + [(1000, 100, 0), (1100, 100, 0), (3000, 400, 0)], + dtype=np.uint64, + ) + + with pytest.raises(AssertionError, match="Head-sharded"): + list(worker._build_local_splits_from_plan(plan, src_blocks_data, 2, 2)) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 9063308ce2e..480ac8937df 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -12,6 +12,7 @@ import uuid from collections import defaultdict from collections.abc import Iterator from concurrent.futures import Future, ThreadPoolExecutor +from functools import cached_property from typing import TYPE_CHECKING, Any, cast import msgspec @@ -67,6 +68,7 @@ from vllm.distributed.parallel_state import ( from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.network_utils import make_zmq_path +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backends.utils import get_kv_cache_layout from vllm.v1.kv_cache_interface import ( FullAttentionSpec, @@ -122,9 +124,11 @@ class NixlBaseConnectorWorker: return (region_ids * num_blocks + block_arr).flatten() # Compute desc ids per group using the right stride: FA descs have - # num_blocks entries per region (kernel granularity), SSM descs have - # logical_blocks entries per region (no kernel splitting). - logical_blocks = num_blocks // physical_blocks_per_logical + # num_blocks entries per region (kernel granularity, expanded by + # block_size_ratio for heterogeneous block sizes), SSM descs have + # logical_blocks entries per region (no kernel splitting, and never + # ratio-expanded since state blocks are indivisible). + logical_blocks = dst_num_blocks // physical_blocks_per_logical all_descs: list[np.ndarray] = [] for i, group in enumerate(block_ids): group_arr = np.asarray(group) @@ -161,6 +165,7 @@ class NixlBaseConnectorWorker: plan: TPMapping, src_blocks_data: np.ndarray, num_fa_descs: int, + block_size_ratio: int = 1, ) -> Iterator[list[tuple[int, int, int]]]: """Build split handle data for P_TP > D_TP scenario. @@ -187,6 +192,11 @@ class NixlBaseConnectorWorker: # Per-FA-descriptor replicate flag, in _build_fa_local emission order. fa_desc_replicated = self._fa_desc_replicated(num_fa_descs) + + assert block_size_ratio == 1 or fa_num_splits == 1 or all(fa_desc_replicated), ( + "Head-sharded attention reads with P_TP > D_TP and heterogeneous " + "block sizes are not supported" + ) src_blocks_list = src_blocks_data.tolist() for p_idx, p_rank in enumerate(plan.all_source_ranks): @@ -442,9 +452,12 @@ class NixlBaseConnectorWorker: # nixl_prepped_dlist_handle. self.src_xfer_handles_by_block_size: dict[int, int] = {} + # Local descriptor arrays per remote block size (block_size_ratio>1), + # kept for building per-tp-ratio splits at the same granularity. + self.src_blocks_data_by_block_size: dict[int, np.ndarray] = {} # Populated dynamically during handshake based on remote configuration. - # Keep track of regions at different tp_ratio values. tp_ratio->handles - self.src_xfer_handles_by_tp_ratio: dict[int, list[int]] = {} + # Per-source split handles, keyed by (tp_ratio, remote_block_size). + self.src_xfer_handles_by_tp_ratio: dict[tuple[int, int], list[int]] = {} # Map of engine_id -> {tp_rank: nixl_prepped_dlist_handle (int)}. self.dst_xfer_side_handles = defaultdict[EngineId, dict[int, int]](dict) @@ -1258,11 +1271,7 @@ class NixlBaseConnectorWorker: agent_metadata_bytes=encoder.encode(agent_metadata), ) - def _build_mamba_local( - self, - base_addresses: list[int], - block_size_ratio: int, - ) -> np.ndarray: + def _build_mamba_local(self, base_addresses: list[int]) -> np.ndarray: """Build desc regions (conv sub-projections + ssm) per layer for local mamba blocks with DS conv layout, as an Nx3 uint64 array. @@ -1289,16 +1298,17 @@ class NixlBaseConnectorWorker: | Key N-1 | Val N-1 | |Conv N-1| SSM N-1 | +-------------------+ +--------------------+ |1st_split-2nd_split| |1st_split-2nd_split | + + Mamba state blocks are indivisible (not token-extent data), so the + descriptors always use the local page geometry regardless of any + attention block-size ratio; their desc ids are likewise never + ratio-expanded (see _compute_desc_ids). """ - assert block_size_ratio == 1, ( - "Mamba 3-read transfer with block_size_ratio != 1 is not tested. " - f"Got block_size_ratio={block_size_ratio}." - ) assert base_addresses, "Local KV cache base addresses must not be empty." assert self._conv_decomp is not None conv_offsets = self._conv_decomp.local_conv_offsets conv_size, ssm_size = self._mamba_ssm_size - num_blocks = self._logical_num_blocks * block_size_ratio + num_blocks = self._logical_num_blocks physical_per_logical = self._physical_blocks_per_logical_kv_block device_id = self.device_id block_arange = np.arange(num_blocks, dtype=np.uint64) @@ -1307,9 +1317,7 @@ class NixlBaseConnectorWorker: for i, base_addr in enumerate(base_addresses): # Jump one page_size, but ssm page_size may be bigger when kernel # locks block size to a specific value (physical_per_logical scale). - page_stride = ( - self.block_len_per_layer[i] // block_size_ratio * physical_per_logical - ) + page_stride = self.block_len_per_layer[i] * physical_per_logical blk_addrs = base_addr + block_arange * page_stride for off, sz in conv_offsets: parts.append(self._stack_descs(blk_addrs + off, sz, device_id)) @@ -1459,7 +1467,7 @@ class NixlBaseConnectorWorker: self.device_id, ) if self._has_mamba: - assert self.num_descs == len(blocks_data) + assert self.num_descs * block_size_ratio == len(blocks_data) # TODO (ZhanqiuHu): For homogeneous TP (tp_ratio == 1), the 3-descs split # is unnecessary — a single conv desc per block suffices. Consider # adding a fast path that falls back to the standard 2-region @@ -1467,7 +1475,7 @@ class NixlBaseConnectorWorker: # remote has been seen. Currently we always register 4 regions # because local descs are created before knowing the remote TP. logger.debug("Registering local Mamba descriptors (4 regions/layer)") - mamba = self._build_mamba_local(local_base_addresses, block_size_ratio) + mamba = self._build_mamba_local(local_base_addresses) blocks_data = np.concatenate([blocks_data, mamba]) descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type) @@ -1607,29 +1615,44 @@ class NixlBaseConnectorWorker: plan = self.tp_mappings[engine_id] + ### (Optional) Register a local handler at the remote engine's block + ### granularity (remote/prefill blocks smaller than local). + remote_block_size = nixl_agent_meta.block_size + src_blocks_data = self.src_blocks_data + if block_size_ratio > 1: + if remote_block_size not in self.src_xfer_handles_by_block_size: + handle, blocks_data = self.register_local_xfer_handler( + remote_block_size + ) + self.src_xfer_handles_by_block_size[remote_block_size] = handle + self.src_blocks_data_by_block_size[remote_block_size] = blocks_data + src_blocks_data = self.src_blocks_data_by_block_size[remote_block_size] + ### (Optional) Register local agent memory regions. MLA is not split. + split_key = (tp_ratio, remote_block_size) if ( tp_ratio < 0 and (not self.use_mla or len(plan.all_source_ranks) > 1) - and tp_ratio not in self.src_xfer_handles_by_tp_ratio + and split_key not in self.src_xfer_handles_by_tp_ratio ): # Remote tp_size > local tp_size: read from multiple remote ranks. # Logically "split" own regions into per-source chunks. Hybrid # MLA+SSM also needs this path: MLA is replicated and read once, # while the SSM state is sharded across every remote TP rank. - # We only do this once per remote tp_size (replica-friendly). - self.src_xfer_handles_by_tp_ratio[tp_ratio] = [] + # We only do this once per remote (tp_size, block_size). + self.src_xfer_handles_by_tp_ratio[split_key] = [] for handle_data in self._build_local_splits_from_plan( plan, - self.src_blocks_data, - self.num_descs, + src_blocks_data, + self.num_descs * block_size_ratio, + block_size_ratio, ): descs = self.nixl_wrapper.get_xfer_descs( handle_data, self.nixl_memory_type ) handle = self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs) - self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle) + self.src_xfer_handles_by_tp_ratio[split_key].append(handle) ### Register remote agent memory regions # With homogeneous TP, D pulls the whole kv cache from corresponding rank. With @@ -1665,13 +1688,6 @@ class NixlBaseConnectorWorker: self.nixl_wrapper.prep_xfer_dlist(remote_agent_name, descs) ) - if block_size_ratio > 1: - # when prefill with smaller block_size, we need to init a - # new handler with same block_len to match - self.src_xfer_handles_by_block_size[nixl_agent_meta.block_size] = ( - self.register_local_xfer_handler(nixl_agent_meta.block_size)[0] - ) - return remote_agent_name def _validate_remote_agent_handshake( @@ -1716,9 +1732,13 @@ class NixlBaseConnectorWorker: "Disable prefix caching with --no-enable-prefix-caching." ) - if self._is_hma_required: - assert block_size_ratio == 1, ( - "HMA does not support different remote block size yet" + if block_size_ratio != 1: + # Heterogeneous block sizes transfer at remote-block granularity; + # the untransferred tail of the last local attention block is + # zeroed in the receive post-process, and mamba state pages + # transfer 1:1 (never sub-split). + assert not self.use_host_buffer, ( + "Heterogeneous block sizes are not supported with host buffer" ) kv_cache_layout = ( self.kv_cache_layout @@ -1875,27 +1895,57 @@ class NixlBaseConnectorWorker: "d2h", ) + @cached_property + def _attention_kv_caches(self) -> list[torch.Tensor]: + """Device KV caches of attention layers (mamba states excluded), + as consumed by the receive post-process.""" + assert self.device_kv_caches, ( + "_attention_kv_caches accessed before register_kv_caches" + ) + mamba_layers = { + name + for g, group in enumerate(self.kv_cache_config.kv_cache_groups) + if _is_ssm_spec(self._group_spec_types[g]) + for name in group.layer_names + } + kv_caches = self.device_kv_caches + return [cache for name, cache in kv_caches.items() if name not in mamba_layers] + def post_process_device_kv_on_receive( self, block_size_ratio: int, - block_ids_list: list[list[int]], + block_ids_list: list[tuple[list[int], int]], + convert: bool = True, ): """ Post process device kv cache after receiving from remote. - 3 types of post processing supported: + 3 types of conversion supported (``convert``): * kv_cache_postprocess_layout => convert from HND to NHD * kv_cache_postprocess_blksize => convert from small block size to large block size * kv_cache_postprocess_blksize_and_layout => convert from small block size to large block size and convert from HND to NHD + The transfer only covers ``covered_sub_blocks`` remote-sized + sub-blocks of each request's local attention blocks; the rest was + clipped, either by remote-block pairing (block-size ratio) or by the + hetero-ppl front trim in ``_apply_prefix_caching``. Those blocks were + excluded from the scheduler's alloc-time KV zeroing (which would race + the RDMA write), so everything past the covered range is zeroed here. + Stale bytes would otherwise surface as garbage or NaNs once decode + grows into the untransferred tail. """ if len(self.device_kv_caches) == 0: return assert block_size_ratio >= 1, "Only nP < nD supported currently." assert self.transfer_topo is not None - if self.enable_permute_local_kv and block_size_ratio > 1: + if not convert: + logger.debug( + "Post-processing device kv cache on receive by zeroing " + "untransferred blocks." + ) + elif self.enable_permute_local_kv and block_size_ratio > 1: logger.debug( "Post-processing device kv cache on receive by converting " "block_size with %sx bigger and permuting layout from HND" @@ -1914,18 +1964,45 @@ class NixlBaseConnectorWorker: block_size_ratio, ) - for block_ids in block_ids_list: - indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long) + attn_caches = self._attention_kv_caches + device = attn_caches[0].device + for block_ids, covered_sub_blocks in block_ids_list: + # Blocks the transfer didn't write: the token tail of the last + # partially covered block, then everything beyond it. + covered_blocks, sub_blocks_in_last = divmod( + covered_sub_blocks, block_size_ratio + ) + first_stale = covered_blocks + (1 if sub_blocks_in_last else 0) + has_stale = first_stale < len(block_ids) + indices = None + if convert or has_stale: + indices = async_tensor_h2d(block_ids, device, torch.long) - for cache in self.device_kv_caches.values(): - if self.enable_permute_local_kv and block_size_ratio > 1: - kv_postprocess_blksize_and_layout_on_receive( - cache, indices, block_size_ratio - ) - elif self.enable_permute_local_kv: - kv_postprocess_layout_on_receive(cache, indices) - else: - kv_postprocess_blksize_on_receive(cache, indices, block_size_ratio) + if convert: + for cache in attn_caches: + if self.enable_permute_local_kv and block_size_ratio > 1: + kv_postprocess_blksize_and_layout_on_receive( + cache, indices, block_size_ratio + ) + elif self.enable_permute_local_kv: + kv_postprocess_layout_on_receive(cache, indices) + else: + kv_postprocess_blksize_on_receive( + cache, indices, block_size_ratio + ) + + if sub_blocks_in_last: + last_block_id = block_ids[covered_blocks] + for cache in attn_caches: + # Both post-processed layouts leave tokens on dim 1. + sub_block_tokens = cache.shape[1] // block_size_ratio + zero_from = sub_blocks_in_last * sub_block_tokens + cache[last_block_id, zero_from:].zero_() + if has_stale: + assert indices is not None + stale_ids = indices[first_stale:] + for cache in attn_caches: + cache.index_fill_(0, stale_ids, 0) def post_process_device_kv_on_receive_heterogeneous_attn( self, block_ids: list[int] @@ -1995,18 +2072,33 @@ class NixlBaseConnectorWorker: if self.use_host_buffer: self.sync_recved_kv_to_device(req_id, meta) - # post processing for heteroblocksize + # Post processing for heteroblocksize/layout, and for blocks the + # transfer clipped. The latter happens either at remote-block + # granularity (block_size_ratio > 1) or at kernel-block + # granularity, when equal kernel pages meet differing logical + # block sizes and _apply_prefix_caching front-trims to the + # minimum count (hybrid heterogeneous TP). remote_info = self.transfer_topo.get_engine_info(meta.remote.engine_id) block_size_ratio = self.transfer_topo.block_size_ratio( remote_info.remote_block_size ) - if not self.use_mla and ( - block_size_ratio > 1 or self.enable_permute_local_kv - ): - assert not self._is_hma_required - block_ids_for_blocksize_post_process[block_size_ratio].append( - meta.local_physical_block_ids[0] - ) + hetero_ppl = ( + remote_info.remote_physical_blocks_per_logical + != self._physical_blocks_per_logical_kv_block + ) + if block_size_ratio > 1 or self.enable_permute_local_kv or hetero_ppl: + for g, local_group in enumerate(meta.local_physical_block_ids): + if not local_group or _is_ssm_spec(self._group_spec_types[g]): + continue + # Number of remote-sized sub-blocks the transfer covered; + # everything past this was clipped and must be zeroed. + covered_sub_blocks = min( + len(local_group) * block_size_ratio, + len(meta.remote.block_ids[g]), + ) + block_ids_for_blocksize_post_process[block_size_ratio].append( + (local_group, covered_sub_blocks) + ) # post processing for heterogeneous attention if self.enable_heterogeneous_attn_post_process: block_ids_for_heterogeneous_attn_post_process.append( @@ -2016,7 +2108,14 @@ class NixlBaseConnectorWorker: block_size_ratio, block_ids_list, ) in block_ids_for_blocksize_post_process.items(): - self.post_process_device_kv_on_receive(block_size_ratio, block_ids_list) + # MLA never needs the block-size/layout conversion, but its + # clipped blocks still need zeroing. + convert = not self.use_mla and ( + block_size_ratio > 1 or self.enable_permute_local_kv + ) + self.post_process_device_kv_on_receive( + block_size_ratio, block_ids_list, convert + ) for block_ids in block_ids_for_heterogeneous_attn_post_process: self.post_process_device_kv_on_receive_heterogeneous_attn(block_ids) @@ -2206,6 +2305,45 @@ class NixlBaseConnectorWorker: return mapped_2d.flatten().astype(np.int64) + def _map_block_ids_for_block_size_ratio( + self, + local_block_ids: BlockIds, + remote_block_ids: BlockIds, + block_size_ratio: int, + ) -> tuple[BlockIds, BlockIds]: + """Map attention-group block ids to remote-block granularity. + + Each local attention block is split into ``block_size_ratio`` + sub-blocks paired 1:1 with remote blocks. Sub-blocks beyond the + remote list — the untransferred tail of the last local block — are + clipped here and zeroed in the receive post-process. Mamba state + blocks are indivisible and transfer 1:1, unexpanded. + + ex: remote (prefill) block ids with block_size 4: + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + Local (decode) block ids with block_size 16: [1, 2, 3] expand to + [4, 5, ..., 15], then clip to the first 10 to pair 1:1 with remote. + """ + mapped_local: list[list[int]] = [] + mapped_remote: list[list[int]] = [] + for i, remote_group in enumerate(remote_block_ids): + local_group = local_block_ids[i] if local_block_ids else [] + if _is_ssm_spec(self._group_spec_types[i]): + mapped_local.append(list(local_group)) + mapped_remote.append(list(remote_group)) + continue + mapped = self.get_mapped_blocks( + np.asarray(local_group), block_size_ratio + ).tolist() + if len(mapped) > len(remote_group): + mapped = mapped[: len(remote_group)] + mapped_local.append(mapped) + mapped_remote.append(list(remote_group)) + if not any(mapped_local): + # Full prefix cache hit is indicated with an empty list. + return [], mapped_remote + return mapped_local, mapped_remote + def _logical_to_kernel_block_ids(self, block_ids: BlockIds, ratio: int) -> BlockIds: """ Convert block ids to kernel physical block ids. @@ -2300,13 +2438,18 @@ class NixlBaseConnectorWorker: remote_block_ids[i] = remote_group[-num_local_blocks:] else: # TODO Handle prefix caching with different block_sizes - max_padding = max( - self._physical_blocks_per_logical_kv_block, - remote_physical_per_logical, + # Allocation rounding legitimately leaves up to + # ppl - 1 trailing dead kernel blocks per side (plus one + # extra local block for the recomputed final token), so + # the counts may differ by up to the sum of the two + # ratios; anything larger indicates mismatched lists. + max_padding = ( + self._physical_blocks_per_logical_kv_block + + remote_physical_per_logical ) - assert abs(num_local_blocks - num_remote_blocks) < max_padding, ( + assert abs(num_local_blocks - num_remote_blocks) <= max_padding, ( f"Group {i}: |{num_local_blocks} - " - f"{num_remote_blocks}| >= {max_padding}" + f"{num_remote_blocks}| > {max_padding}" ) num_blocks = min(num_local_blocks, num_remote_blocks) local_block_ids[i] = local_block_ids[i][:num_blocks] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index 40d6851769c..0c66596c4b7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -5,8 +5,6 @@ import time from typing import TYPE_CHECKING -import numpy as np - from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import ( NixlBaseConnectorWorker, ) @@ -178,10 +176,10 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): ) # Get side handles. if tp_ratio < 0 and (not self.use_mla or len(read_specs) > 1): - assert remote_block_size == self.block_size # Remote tp_size > local tp_size: we must perform multiple # reads. Get the memory chunk onto which we will write to. - local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[tp_ratio][i] + split_key = (tp_ratio, remote_block_size) + local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[split_key][i] else: # Single read from remote, we write to the whole memory region. # Also handle remote block size different from local block size. @@ -235,30 +233,11 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): remote_info.remote_block_size ) if block_size_ratio > 1: - # TODO (NickLucche) assume HMA is off. Change to handle multiple KV groups. - assert not self._is_hma_required - local_block_ids0 = local_block_ids[0] if local_block_ids else [] - remote_block_ids0 = remote_block_ids[0] - local_block_ids_mapped = self.get_mapped_blocks( - np.asarray(local_block_ids0), block_size_ratio - ).tolist() - if len(local_block_ids_mapped) > len(remote_block_ids0): - # NOTE: - # get_mapped_blocks will always expand block_ids for n times. - # ex: - # prefill block_ids with block_size as 4: - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] - # Local decode block_ids with block_size as 16: [1, 2, 3] - # expanded decode block_ids with get_mapped_blocks from [1, 2, 3] to - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] - # Then we clip local to align with prefill - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] to - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] - local_block_ids_mapped = local_block_ids_mapped[ - : len(remote_block_ids0) - ] - local_block_ids = [local_block_ids_mapped] if local_block_ids_mapped else [] - remote_block_ids = [remote_block_ids0] + local_block_ids, remote_block_ids = ( + self._map_block_ids_for_block_size_ratio( + local_block_ids, remote_block_ids, block_size_ratio + ) + ) # NOTE(rob): having the staging blocks be on the READER side is # not going to work well (since we will have to call rearrange tensors). # after we detect the txn is complete (which means we cannot make the diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 8cb710e7e58..6bc48963c58 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -39,7 +39,6 @@ from concurrent.futures import Future from typing import TYPE_CHECKING, Any import msgspec -import numpy as np from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import ( @@ -553,8 +552,8 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): req_id, ) if tp_ratio < 0 and not self.use_mla: - assert remote_block_size == self.block_size - local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[tp_ratio][i] + split_key = (tp_ratio, remote_block_size) + local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[split_key][i] else: local_xfer_side_handle = self.src_xfer_handles_by_block_size[ remote_block_size @@ -606,18 +605,11 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): remote_info.remote_block_size ) if block_size_ratio > 1: - assert not self._is_hma_required - local_block_ids0 = local_block_ids[0] if local_block_ids else [] - remote_block_ids0 = remote_block_ids[0] - local_block_ids_mapped = self.get_mapped_blocks( - np.asarray(local_block_ids0), block_size_ratio - ).tolist() - if len(local_block_ids_mapped) > len(remote_block_ids0): - local_block_ids_mapped = local_block_ids_mapped[ - : len(remote_block_ids0) - ] - local_block_ids = [local_block_ids_mapped] if local_block_ids_mapped else [] - remote_block_ids = [remote_block_ids0] + local_block_ids, remote_block_ids = ( + self._map_block_ids_for_block_size_ratio( + local_block_ids, remote_block_ids, block_size_ratio + ) + ) notif_id = f"{remote_request_id}:{self.world_size}".encode() From 62d8db7c05af8b9ef3655cf13d68416dbe3185d8 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Tue, 28 Jul 2026 14:47:04 +0100 Subject: [PATCH 171/185] [Bugfix] Add missing `vllm/models/kimi_k3/__init__.py` (#50131) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- vllm/models/kimi_k3/__init__.py | 2 ++ vllm/models/kimi_k3/amd/ops/__init__.py | 2 ++ 2 files changed, 4 insertions(+) create mode 100644 vllm/models/kimi_k3/__init__.py diff --git a/vllm/models/kimi_k3/__init__.py b/vllm/models/kimi_k3/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm/models/kimi_k3/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/amd/ops/__init__.py b/vllm/models/kimi_k3/amd/ops/__init__.py index e69de29bb2d..208f01a7cb5 100644 --- a/vllm/models/kimi_k3/amd/ops/__init__.py +++ b/vllm/models/kimi_k3/amd/ops/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project From 94100b5915d449aed0685cc5c6fc8949fcf5fe40 Mon Sep 17 00:00:00 2001 From: Nick Hill <nickhill123@gmail.com> Date: Tue, 28 Jul 2026 07:02:31 -0700 Subject: [PATCH 172/185] [CI] Wire untethered test files into CI jobs (#49340) Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --- .buildkite/test_areas/cuda.yaml | 1 + .buildkite/test_areas/disaggregated.yaml | 16 +++++++++ .buildkite/test_areas/engine.yaml | 2 ++ .buildkite/test_areas/kernels.yaml | 36 +++++++++++++++++++ .buildkite/test_areas/misc.yaml | 10 ++++-- .buildkite/test_areas/models_basic.yaml | 3 +- .buildkite/test_areas/spec_decode.yaml | 2 ++ ..._fused_minimax_m3_qknorm_rope_kv_insert.py | 21 +++++++---- .../nixl_integration/run_edge_case_test.sh | 10 +++--- 9 files changed, 88 insertions(+), 13 deletions(-) diff --git a/.buildkite/test_areas/cuda.yaml b/.buildkite/test_areas/cuda.yaml index 927b5bd27f2..431ce07af4d 100644 --- a/.buildkite/test_areas/cuda.yaml +++ b/.buildkite/test_areas/cuda.yaml @@ -16,6 +16,7 @@ steps: commands: - pytest -v -s cuda/test_cuda_context.py - pytest -v -s cuda/test_platform_no_cuda_init.py + - pytest -v -s cuda/test_cuda_compatibility_path.py - label: Cudagraph device: h200_35gb diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index a3342e362ed..f1a89b39682 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -131,6 +131,22 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +- label: NixlConnector PD edge case test (2 GPUs) + key: nixlconnector-pd-edge-cases-2-gpus + timeout_in_minutes: 40 + working_dir: "/vllm-workspace/tests" + num_devices: 2 + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/v1/core/sched/ + - tests/v1/kv_connector/nixl_integration/ + env: + PREFILL_GPU_ID: "0" + DECODE_GPU_ID: "1" + commands: + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh + - bash v1/kv_connector/nixl_integration/run_edge_case_test.sh + - label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs) key: hybrid-ssm-nixlconnector-pd-prefix-cache-2-gpus timeout_in_minutes: 25 diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index ed593c4aba2..ce4fc590eec 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -40,9 +40,11 @@ steps: source_file_dependencies: - vllm/v1/engine/ - tests/v1/engine/ + - tests/v1/test_tensor_ipc_queue.py commands: - pytest -v -s v1/engine/test_preprocess_error_handling.py - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py + - pytest -v -s v1/test_tensor_ipc_queue.py mirror: amd: device: mi250_1 diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index e1951685f60..938f8690551 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -61,9 +61,45 @@ steps: source_file_dependencies: - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu - vllm/models/deepseek_v4/common/ops/ + - vllm/models/deepseek_v4/nvidia/ - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py + - tests/models/test_deepseek_v4_mega_moe.py commands: - pytest -v -s kernels/test_fused_deepseek_v4_*.py + - pytest -v -s models/test_deepseek_v4_mega_moe.py + +# Catch-all for test files at the tests/kernels root. This job collects +# the whole root so new files are wired by default. +# Files with dedicated jobs elsewhere in this file are excluded via --ignore +# (test_kda, test_bf16x3_router_gemm_cutedsl and test_ll_bf16_gemm run in +# their own jobs / Kernels (B200)). +- label: Kernels Root Misc Test (B200) + key: kernels-root-misc-test-b200 + timeout_in_minutes: 45 + device: b200-k8s + source_file_dependencies: + - csrc/ + - vllm/ + - tests/kernels/ + commands: + - pytest -v -s kernels/ + --ignore=kernels/attention + --ignore=kernels/core + --ignore=kernels/helion + --ignore=kernels/ir + --ignore=kernels/mamba + --ignore=kernels/moe + --ignore=kernels/quantization + --ignore=kernels/test_concat_mla_q.py + --ignore=kernels/test_fused_qk_norm_rope_gate.py + --ignore=kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py + --ignore=kernels/test_top_k_per_row.py + --ignore=kernels/test_kda.py + --ignore=kernels/test_bf16x3_router_gemm_cutedsl.py + --ignore=kernels/test_ll_bf16_gemm.py + --ignore=kernels/test_shuffle_rows.py + # BROKEN on main, pending kernel fixes (B200): + # test_shuffle_rows.py (1: test_shuffle_rows_edge_cases) - label: Kernels Attention Test %N key: kernels-attention-test diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 1a53a92961f..763fbcbfec2 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -148,6 +148,7 @@ steps: - pytest -v -s -m 'cpu_test' v1/core - pytest -v -s v1/structured_output - pytest -v -s v1/test_serial_utils.py + - pytest -v -s v1/test_kv_cache_spec_registry.py - pytest -v -s v1/cudagraph/test_cudagraph_manager.py - pytest -v -s -m 'cpu_test' v1/kv_connector/unit - pytest -v -s -m 'cpu_test' v1/metrics @@ -265,6 +266,7 @@ steps: - vllm/utils/ - vllm/v1/ - tests/v1/tracing + - tests/tracing/ commands: - "pip install \ 'opentelemetry-sdk>=1.26.0' \ @@ -272,6 +274,7 @@ steps: 'opentelemetry-exporter-otlp>=1.26.0' \ 'opentelemetry-semantic-conventions-ai>=0.4.1'" - pytest -v -s v1/tracing + - pytest -v -s tracing mirror: amd: dind: false @@ -425,7 +428,7 @@ steps: - label: Batch Invariance (B200) key: batch-invariance-b200 - timeout_in_minutes: 35 + timeout_in_minutes: 45 device: b200-k8s source_file_dependencies: - vllm/v1/attention @@ -440,7 +443,10 @@ steps: - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN] - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py - pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py - + - pytest -v -s v1/determinism/test_matmul_batch_invariant.py + - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py + - pytest -v -s v1/determinism/test_online_batch_invariance.py + - label: Acceptance Length Test (Large Models) # optional device: h200_35gb key: acceptance-length-test-large-models diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index a7c7aa9022d..24bfff5d756 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -82,7 +82,8 @@ steps: - vllm/ - tests/models/test_utils.py - tests/models/test_vision.py + - tests/models/test_adapters.py - tests/models/transformers/fusers/ device: cpu-small commands: - - pytest -v -s models/test_utils.py models/test_vision.py models/transformers/fusers/ + - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/transformers/fusers/ diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index c63aaa18d8b..7cde7124cdc 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -90,8 +90,10 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - tests/v1/e2e/spec_decode/ + - tests/spec_decode/ commands: - pytest -v -s v1/e2e/spec_decode -k "ngram or suffix" + - python3 spec_decode/test_custom_proposer.py mirror: amd: dind: false diff --git a/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py b/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py index 626b06290e0..9c4a996438e 100644 --- a/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py +++ b/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py @@ -147,8 +147,11 @@ def test_dense_norm_rope(num_tokens, num_heads, num_kv_heads): eps, ).view(num_tokens, kvsz) - torch.testing.assert_close(q_out, q_ref, rtol=1e-2, atol=1e-2) - torch.testing.assert_close(k_out, k_ref, rtol=1e-2, atol=1e-2) + # The fused kernel keeps an fp32 intermediate across norm->rope, while the + # reference materializes bf16 after the norm (the unfused boundary), so + # rounding-boundary elements can differ by ~1 bf16 ulp. + torch.testing.assert_close(q_out, q_ref, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(k_out, k_ref, rtol=2e-2, atol=2e-2) # V is untouched. torch.testing.assert_close(v_out, v_in, rtol=0, atol=0) @@ -255,8 +258,11 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype): ik_orig.view(num_tokens, 1, HEAD_DIM), ik_w, positions, cos_sin, eps ).view(num_tokens, HEAD_DIM) - torch.testing.assert_close(q_out, q_ref, rtol=1e-2, atol=1e-2) - torch.testing.assert_close(k_out, k_ref, rtol=1e-2, atol=1e-2) + # The fused kernel keeps an fp32 intermediate across norm->rope, while the + # reference materializes bf16 after the norm (the unfused boundary), so + # rounding-boundary elements can differ by ~1 bf16 ulp. + torch.testing.assert_close(q_out, q_ref, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(k_out, k_ref, rtol=2e-2, atol=2e-2) torch.testing.assert_close(index_q, iq_ref, rtol=1e-2, atol=1e-2) torch.testing.assert_close(index_k, ik_ref, rtol=1e-2, atol=1e-2) @@ -376,8 +382,11 @@ def test_sparse_skip_index_branch(num_tokens, block_size, kv_cache_dtype): eps, ).view(num_tokens, kvsz) - torch.testing.assert_close(q_out, q_ref, rtol=1e-2, atol=1e-2) - torch.testing.assert_close(k_out, k_ref, rtol=1e-2, atol=1e-2) + # The fused kernel keeps an fp32 intermediate across norm->rope, while the + # reference materializes bf16 after the norm (the unfused boundary), so + # rounding-boundary elements can differ by ~1 bf16 ulp. + torch.testing.assert_close(q_out, q_ref, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(k_out, k_ref, rtol=2e-2, atol=2e-2) torch.testing.assert_close(v_out, v_in, rtol=0, atol=0) torch.testing.assert_close(index_q_out, index_q_in, rtol=0, atol=0) torch.testing.assert_close(index_k_out, index_k_in, rtol=0, atol=0) diff --git a/tests/v1/kv_connector/nixl_integration/run_edge_case_test.sh b/tests/v1/kv_connector/nixl_integration/run_edge_case_test.sh index 9d8e4df8c53..c3240ab5c17 100755 --- a/tests/v1/kv_connector/nixl_integration/run_edge_case_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_edge_case_test.sh @@ -3,8 +3,8 @@ set -xe # Parse command line arguments KV_BUFFER_DEVICE="cuda" # Default to cuda -PREFILL_GPU_ID=4 # Default GPU IDs -DECODE_GPU_ID=5 +PREFILL_GPU_ID="${PREFILL_GPU_ID:-4}" # Default GPU IDs +DECODE_GPU_ID="${DECODE_GPU_ID:-5}" while [[ $# -gt 0 ]]; do case $1 in --kv_buffer_device) @@ -70,6 +70,7 @@ run_tests_for_model() { --port $PREFILL_PORT \ --enforce-eager \ --gpu-memory-utilization 0.2 \ + --max-model-len 8192 \ --kv-transfer-config '$KV_CONFIG'" FULL_CMD="$BASE_CMD" @@ -84,6 +85,7 @@ run_tests_for_model() { --port $DECODE_PORT \ --enforce-eager \ --gpu-memory-utilization 0.2 \ + --max-model-len 8192 \ --kv-transfer-config '$KV_CONFIG'" FULL_CMD="$BASE_CMD" @@ -98,7 +100,7 @@ run_tests_for_model() { # Build the command for the proxy server with all the hosts and ports PROXY_PORT=8192 - PROXY_CMD="python ${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py --port $PROXY_PORT" + PROXY_CMD="python3 ${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py --port $PROXY_PORT" PROXY_CMD+=" --prefiller-ports ${PREFILL_PORT}" PROXY_CMD+=" --decoder-ports ${DECODE_PORT}" # Start the proxy server @@ -110,7 +112,7 @@ run_tests_for_model() { # Run lm eval for this model echo "Running tests for $model_name" - PREFILL_PORT=$PREFILL_PORT DECODE_PORT=$DECODE_PORT PROXY_PORT=$PROXY_PORT python -m pytest -s -v "${GIT_ROOT}"/tests/v1/kv_connector/nixl_integration/test_edge_cases.py + PREFILL_PORT=$PREFILL_PORT DECODE_PORT=$DECODE_PORT PROXY_PORT=$PROXY_PORT python3 -m pytest -s -v "${GIT_ROOT}"/tests/v1/kv_connector/nixl_integration/test_edge_cases.py # Clean up before running next model cleanup_instances From b6cbba8bc893c61e412a205533aafbee1ae6be31 Mon Sep 17 00:00:00 2001 From: oops-oom <liubin8905@vip.qq.com> Date: Tue, 28 Jul 2026 22:24:02 +0800 Subject: [PATCH 173/185] [Bugfix][Kernel] Fix batch invariance in RMSNorm kernels by pinning block size (#48391) Signed-off-by: oops-oom <73481342@qq.com> Signed-off-by: oops-oom <liubin8905@vip.qq.com> Co-authored-by: oops-oom <73481342@qq.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Shengqi Chen <harry-chen@outlook.com> --- .buildkite/test_areas/misc.yaml | 14 +- csrc/libtorch_stable/layernorm_kernels.cu | 14 +- .../layernorm_quant_kernels.cu | 9 +- ...fused_layernorm_dynamic_per_token_quant.cu | 5 +- tests/v1/determinism/test_batch_invariance.py | 18 +++ .../test_rms_norm_batch_invariant.py | 150 +++++++++++++++++- 6 files changed, 189 insertions(+), 21 deletions(-) diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 763fbcbfec2..caa56c21b37 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -398,7 +398,7 @@ steps: - label: Batch Invariance (A100) key: batch-invariance-a100 - timeout_in_minutes: 40 + timeout_in_minutes: 60 device: a100 source_file_dependencies: - vllm/v1/attention @@ -408,11 +408,11 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pip install pytest-timeout pytest-forked - pytest -v -s v1/determinism/test_batch_invariance.py - - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA] + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA - label: Batch Invariance (H100) key: batch-invariance-h100 - timeout_in_minutes: 40 + timeout_in_minutes: 60 device: h100 source_file_dependencies: - vllm/v1/attention @@ -423,8 +423,8 @@ steps: - pip install pytest-timeout pytest-forked - pytest -v -s v1/determinism/test_batch_invariance.py - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py - - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA] - - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN] + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN - label: Batch Invariance (B200) key: batch-invariance-b200 @@ -439,8 +439,8 @@ steps: - pip install pytest-timeout pytest-forked - pytest -v -s v1/determinism/test_batch_invariance.py - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py - - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA] - - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN] + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py - pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py - pytest -v -s v1/determinism/test_matmul_batch_invariant.py diff --git a/csrc/libtorch_stable/layernorm_kernels.cu b/csrc/libtorch_stable/layernorm_kernels.cu index 878b44df936..7a1051d2c00 100644 --- a/csrc/libtorch_stable/layernorm_kernels.cu +++ b/csrc/libtorch_stable/layernorm_kernels.cu @@ -249,7 +249,9 @@ void rms_norm(torch::stable::Tensor& out, // [..., hidden_size] int64_t input_shape_d3 = (num_dims >= 4) ? input.size(-3) : 0; // For large num_tokens, use smaller blocks to increase SM concurrency. - const int max_block_size = (num_tokens < 256) ? 1024 : 256; + const bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); + const int max_block_size = + batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256); dim3 grid(num_tokens); const torch::stable::accelerator::DeviceGuard device_guard( input.get_device_index()); @@ -325,8 +327,13 @@ void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size] /* This kernel is memory-latency bound in many scenarios. When num_tokens is large, a smaller block size allows for increased block occupancy on CUs and better latency - hiding on global mem ops. */ - const int max_block_size = (num_tokens < 256) ? 1024 : 256; + hiding on global mem ops. In batch-invariant mode the block size must + not depend on num_tokens, otherwise the same token would use a different + reduction width (and thus a different floating-point summation order) + across batches; lock it to 1024 to keep results bit-exact. */ + const bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); + const int max_block_size = + batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256); dim3 block(std::min(hidden_size, max_block_size)); const torch::stable::accelerator::DeviceGuard device_guard( input.get_device_index()); @@ -337,7 +344,6 @@ void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size] auto res_ptr = reinterpret_cast<std::uintptr_t>(residual.data_ptr()); bool offsets_are_multiple_of_vector_width = hidden_size % vector_width == 0 && input_stride % vector_width == 0; - bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); const bool has_weight = weight.has_value(); if (has_weight) { auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight->data_ptr()); diff --git a/csrc/libtorch_stable/layernorm_quant_kernels.cu b/csrc/libtorch_stable/layernorm_quant_kernels.cu index f3bf8882e77..f43be531de0 100644 --- a/csrc/libtorch_stable/layernorm_quant_kernels.cu +++ b/csrc/libtorch_stable/layernorm_quant_kernels.cu @@ -215,7 +215,9 @@ void rms_norm_static_fp8_quant( int num_tokens = input.numel() / hidden_size; // For large num_tokens, use smaller blocks to increase SM concurrency. - const int max_block_size = (num_tokens < 256) ? 1024 : 256; + const bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); + const int max_block_size = + batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256); dim3 grid(num_tokens); const torch::stable::accelerator::DeviceGuard device_guard( input.get_device_index()); @@ -279,7 +281,9 @@ void fused_add_rms_norm_static_fp8_quant( When num_tokens is large, a smaller block size allows for increased block occupancy on CUs and better latency hiding on global mem ops. */ - const int max_block_size = (num_tokens < 256) ? 1024 : 256; + const bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); + const int max_block_size = + batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256); dim3 block(std::min(hidden_size, max_block_size)); const torch::stable::accelerator::DeviceGuard device_guard( input.get_device_index()); @@ -296,7 +300,6 @@ void fused_add_rms_norm_static_fp8_quant( auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight.data_ptr()); bool ptrs_are_aligned = inp_ptr % 16 == 0 && res_ptr % 16 == 0 && wt_ptr % 16 == 0; - bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); if (ptrs_are_aligned && hidden_size % 8 == 0 && input_stride % 8 == 0 && !batch_invariant_launch) { LAUNCH_FUSED_ADD_RMS_NORM(8); diff --git a/csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu b/csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu index 2152e64dc96..56dd4703872 100644 --- a/csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu +++ b/csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu @@ -2,6 +2,7 @@ #include "../../torch_utils.h" #include "../../dispatch_utils.h" +#include "../../../core/batch_invariant.hpp" #include "layernorm_utils.cuh" #include "quant_conversions.cuh" @@ -231,7 +232,9 @@ void rms_norm_per_block_quant_dispatch( auto num_tokens = input.numel() / hidden_size; dim3 grid(num_tokens); - const int max_block_size = (num_tokens <= 256) ? 512 : 256; + const bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); + const int max_block_size = + batch_invariant_launch ? 512 : ((num_tokens <= 256) ? 512 : 256); dim3 block(std::min(hidden_size, max_block_size)); const torch::stable::accelerator::DeviceGuard device_guard( input.get_device_index()); diff --git a/tests/v1/determinism/test_batch_invariance.py b/tests/v1/determinism/test_batch_invariance.py index b2706ed89b7..37fd5cba6a5 100644 --- a/tests/v1/determinism/test_batch_invariance.py +++ b/tests/v1/determinism/test_batch_invariance.py @@ -27,8 +27,10 @@ from vllm.platforms import current_platform "backend", BACKENDS, ) +@pytest.mark.parametrize("rms_norm_impl", ["default", "vllm_c"]) def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( backend, + rms_norm_impl, ): """ Ensures that the same request (the 'needle' prompt) yields identical output @@ -60,6 +62,16 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( random.seed(seed) attention_config = {"backend": backend} + # Force the C++ RMSNorm implementation so we actually exercise the + # num_tokens-dependent block-size branches. + kernel_config = None + if rms_norm_impl == "vllm_c": + kernel_config = { + "ir_op_priority": { + "rms_norm": ["vllm_c"], + "fused_add_rms_norm": ["vllm_c"], + } + } # Allow overrides from environment (useful for CI tuning) # "facebook/opt-125m" is too small, doesn't reliably test determinism model = TEST_MODEL @@ -96,6 +108,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( gpu_memory_utilization=gpu_mem_util, max_model_len=max_model_len, attention_config=attention_config, + kernel_config=kernel_config, ) # Baseline generation for the needle prompt alone. @@ -923,11 +936,15 @@ def LLM_with_max_seqs( gpu_memory_utilization: float, max_model_len: int, attention_config: dict | None = None, + kernel_config: dict | None = None, ) -> LLM: """ Helper to construct an LLM with a specific max_num_seqs (batch-size limit) using the high-level v1 LLM API, while constraining memory usage. """ + extra_kwargs: dict = {} + if kernel_config is not None: + extra_kwargs["kernel_config"] = kernel_config return LLM( model=model, max_num_seqs=max_num_seqs, @@ -939,4 +956,5 @@ def LLM_with_max_seqs( attention_config=attention_config, # Enable for MOE models # enable_expert_parallel=True, + **extra_kwargs, ) diff --git a/tests/v1/determinism/test_rms_norm_batch_invariant.py b/tests/v1/determinism/test_rms_norm_batch_invariant.py index dfd08351277..232a43b1f98 100644 --- a/tests/v1/determinism/test_rms_norm_batch_invariant.py +++ b/tests/v1/determinism/test_rms_norm_batch_invariant.py @@ -28,16 +28,18 @@ def _rms_norm_reference( @skip_if_not_cuda -@pytest.mark.parametrize("batch_size", [1, 4, 16, 64]) +@pytest.mark.parametrize("batch_size", [1, 4, 64, 300]) @pytest.mark.parametrize("hidden_size", [512, 2048, 4096, 8192]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("eps", [1e-6, 1e-5]) +@pytest.mark.parametrize("seed", list(range(4))) def test_rms_norm_batch_invariant_vs_reference( default_vllm_config, batch_size: int, hidden_size: int, dtype: torch.dtype, eps: float, + seed: int, ): """ Compare batch-invariant Triton RMS norm against a PyTorch reference. @@ -48,7 +50,7 @@ def test_rms_norm_batch_invariant_vs_reference( device = torch.device(DEVICE_TYPE) # Create test input and weight - torch.manual_seed(42) + torch.manual_seed(seed) input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device) weight = torch.randn(hidden_size, dtype=dtype, device=device) @@ -71,7 +73,7 @@ def test_rms_norm_batch_invariant_vs_reference( atol=atol, msg=f"RMS norm mismatch for batch_size={batch_size}, " f"hidden_size={hidden_size}, " - f"dtype={dtype}, eps={eps}", + f"dtype={dtype}, eps={eps}, seed={seed}", ) @@ -79,17 +81,21 @@ def test_rms_norm_batch_invariant_vs_reference( @pytest.mark.parametrize("hidden_size", [512, 4096]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("eps", [1e-6]) +@pytest.mark.parametrize("n_extra", [3, 299]) +@pytest.mark.parametrize("seed", list(range(16))) def test_fused_add_rms_norm_batch_invariant_residual_path( hidden_size: int, dtype: torch.dtype, eps: float, + n_extra: int, + seed: int, ): """ Test the batch-invariant fused residual-add + RMSNorm helper directly. """ device = torch.device(DEVICE_TYPE) - torch.manual_seed(42) + torch.manual_seed(seed) x_single = torch.randn(1, hidden_size, dtype=dtype, device=device) residual_single = torch.randn(1, hidden_size, dtype=dtype, device=device) weight = torch.randn(hidden_size, dtype=dtype, device=device) @@ -97,14 +103,14 @@ def test_fused_add_rms_norm_batch_invariant_residual_path( x_batch = torch.cat( [ x_single, - torch.randn(3, hidden_size, dtype=dtype, device=device), + torch.randn(n_extra, hidden_size, dtype=dtype, device=device), ], dim=0, ) residual_batch = torch.cat( [ residual_single, - torch.randn(3, hidden_size, dtype=dtype, device=device), + torch.randn(n_extra, hidden_size, dtype=dtype, device=device), ], dim=0, ) @@ -168,6 +174,138 @@ def test_fused_add_rms_norm_batch_invariant_residual_path( ) +FP8_DTYPE = current_platform.fp8_dtype() + +# The large launch (num_tokens=300 >= 256) drops an un-pinned kernel to block +# 256, while the small launch (255 rows) stays under the threshold and keeps the +# larger block (1024, or 512 for per-block quant). Under the pin the two launches +# use the same block, so the shared first 255 rows must match bit-for-bit; 255 is +# the most rows a single small launch can hold (< 256, and <= 256 for per-block). +_LARGE_TOKENS = 300 +_SMALL_TOKENS = 255 + + +def _assert_rows_bit_identical(small, large, msg): + if small.dtype == FP8_DTYPE: + assert torch.equal(small.view(torch.uint8), large.view(torch.uint8)), msg + else: + torch.testing.assert_close(small, large, rtol=0.0, atol=0.0, msg=msg) + + +@skip_if_not_cuda +@pytest.mark.parametrize("hidden_size", [512, 4096]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("seed", list(range(4))) +def test_rms_norm_batch_invariant_nonresidual_kernel( + hidden_size: int, dtype: torch.dtype, seed: int +): + """C++ ``rms_norm`` (no residual) must be batch invariant across the block + threshold. Reached in compiled mode with ``ir_op_priority.rms_norm=["vllm_c"]`` + (default priority is ``native``/inductor codegen when compiling). + """ + import vllm._custom_ops as ops + + device = torch.device(DEVICE_TYPE) + torch.manual_seed(seed) + rows = torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device) + weight = torch.randn(hidden_size, dtype=dtype, device=device) + + def rms_norm(x): + out = torch.empty_like(x) + ops.rms_norm(out, x, weight, 1e-6) + return out + + large = rms_norm(rows.clone()) + small = rms_norm(rows[:_SMALL_TOKENS].clone()) + _assert_rows_bit_identical( + small, + large[:_SMALL_TOKENS], + "rms_norm output depends on num_tokens (block size)", + ) + + +@skip_if_not_cuda +@pytest.mark.parametrize("hidden_size", [512, 4096]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("seed", list(range(4))) +@pytest.mark.parametrize("add_residual", [False, True]) +def test_rms_norm_static_fp8_quant_batch_invariant( + hidden_size: int, dtype: torch.dtype, seed: int, add_residual: bool +): + """C++ static per-tensor fp8-quant RMSNorm must be batch invariant across + the block threshold. Covers ``rms_norm_static_fp8_quant`` and, with + ``add_residual``, ``fused_add_rms_norm_static_fp8_quant`` (the compiled fp8 + path where ``RMSNormQuantFusionPass`` rewrites norm + quant into them). + """ + device = torch.device(DEVICE_TYPE) + torch.manual_seed(seed) + rows = torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device) + residual = ( + torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device) + if add_residual + else None + ) + weight = torch.randn(hidden_size, dtype=dtype, device=device) + quant_scale = torch.tensor(1.0, dtype=torch.float32, device=device) + + def quant(x, res): + out = torch.empty_like(x, dtype=FP8_DTYPE) + if add_residual: + torch.ops._C.fused_add_rms_norm_static_fp8_quant( + out, x, res, weight, quant_scale, 1e-6 + ) + else: + torch.ops._C.rms_norm_static_fp8_quant(out, x, weight, quant_scale, 1e-6) + return out + + large = quant(rows.clone(), residual.clone() if residual is not None else None) + small = quant( + rows[:_SMALL_TOKENS].clone(), + residual[:_SMALL_TOKENS].clone() if residual is not None else None, + ) + _assert_rows_bit_identical( + small, + large[:_SMALL_TOKENS], + "static-fp8-quant RMSNorm output depends on num_tokens (block size)", + ) + + +@skip_if_not_cuda +@pytest.mark.parametrize("hidden_size", [512, 4096]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("seed", list(range(4))) +def test_rms_norm_per_block_quant_batch_invariant( + hidden_size: int, dtype: torch.dtype, seed: int +): + """C++ ``rms_norm_per_block_quant`` must be batch invariant across the + block threshold (compiled fp8 block-quant path; block pinned to 512).""" + import vllm._custom_ops as ops + + device = torch.device(DEVICE_TYPE) + torch.manual_seed(seed) + rows = torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device) + weight = torch.randn(hidden_size, dtype=dtype, device=device) + group_size = [1, 128] + + def per_block_quant(x): + return ops.rms_norm_per_block_quant(x, weight, 1e-6, FP8_DTYPE, group_size) + + out_large, scale_large = per_block_quant(rows.clone()) + out_small, scale_small = per_block_quant(rows[:_SMALL_TOKENS].clone()) + _assert_rows_bit_identical( + out_small, + out_large[:_SMALL_TOKENS], + "rms_norm_per_block_quant output depends on num_tokens (block size)", + ) + torch.testing.assert_close( + scale_small, + scale_large[:_SMALL_TOKENS], + rtol=0.0, + atol=0.0, + msg="rms_norm_per_block_quant scales depend on num_tokens (block size)", + ) + + @skip_if_not_cuda @pytest.mark.parametrize("batch_size", [1, 16, 128]) @pytest.mark.parametrize("seq_len", [1, 32, 512]) From 1e81853afc8701cb50b649e94a23c977da8c1ed0 Mon Sep 17 00:00:00 2001 From: Jonguk Cheong <jdal3031@snu.ac.kr> Date: Tue, 28 Jul 2026 23:41:55 +0900 Subject: [PATCH 174/185] [Bugfix][KV Offload] Keep Mamba block span unscaled under DCP (#49964) Signed-off-by: Jonguk Cheong <jdal3031@snu.ac.kr> Co-authored-by: OpenAI Codex <codex@openai.com> --- .../unit/offloading_connector/test_config.py | 56 +++++++++++++++++++ .../kv_connector/v1/offloading/config.py | 12 +++- 2 files changed, 66 insertions(+), 2 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index fc426ff318d..5a66b463e30 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -8,10 +8,15 @@ from unittest.mock import MagicMock, patch import pytest import torch +from tests.v1.kv_connector.unit.offloading_connector.utils import MockOffloadingSpec from vllm.config import KVTransferConfig, ParallelConfig, VllmConfig from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import ( build_offloading_config, ) +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( + SchedulerOffloadConfig, + is_store_reachable_swa_chunk, +) from vllm.platforms import current_platform from vllm.v1.kv_cache_interface import ( FullAttentionSpec, @@ -179,6 +184,25 @@ def _make_hybrid_kv_cache_config() -> KVCacheConfig: ) +def _make_mamba_hybrid_kv_cache_config() -> KVCacheConfig: + return KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(["full_layer"], _full_attention_spec()), + KVCacheGroupSpec( + ["mamba_layer"], + MambaSpec( + block_size=16, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + + def _parallelism_agnostic(kv_cache_groups: list[KVCacheGroupSpec]) -> bool: config = _make_vllm_config() kv_cache_config = KVCacheConfig( @@ -267,6 +291,38 @@ def test_prefill_context_parallelism_does_not_scale_group_blocks(): assert offloading_config.cache.blocks_per_chunk == 4 +def test_dcp_scales_attention_but_not_mamba_group_blocks(): + config = _make_vllm_config(tensor_parallel_size=2, decode_context_parallel_size=2) + config.speculative_config = None + + offloading_config = build_offloading_config( + config, _make_mamba_hybrid_kv_cache_config() + ) + + assert tuple(group.tokens_per_block for group in offloading_config.groups) == ( + 32, + 16, + ) + scheduler_config = SchedulerOffloadConfig.from_spec( + MockOffloadingSpec(offloading_config), + config, + _make_mamba_hybrid_kv_cache_config(), + ) + mamba_group = scheduler_config.kv_group_configs[1] + assert mamba_group.alignment_chunk_count == 2 + assert [ + chunk_idx + for chunk_idx in range(4) + if is_store_reachable_swa_chunk( + chunk_idx, + 4, + mamba_group.alignment_chunk_count, + mamba_group.sliding_window_size_in_chunks, + mamba_group.is_eagle_group, + ) + ] == [1, 3] + + def test_preserves_data_parallel_index(): config = _make_vllm_config() config.parallel_config.data_parallel_index = 2 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index b86ebf96bb6..b9837bcb4b0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -5,7 +5,11 @@ from typing import TYPE_CHECKING from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes -from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + FullAttentionSpec, + MLAAttentionSpec, +) from vllm.v1.kv_offload.config import ( OffloadingCacheConfig, OffloadingConfig, @@ -40,7 +44,11 @@ def build_offloading_config( OffloadingGroupConfig( tokens_per_block=( group.kv_cache_spec.block_size - * parallel_config.decode_context_parallel_size + * ( + parallel_config.decode_context_parallel_size + if isinstance(group.kv_cache_spec, AttentionSpec) + else 1 + ) ), layer_names=tuple(group.layer_names), ) From 0d0504b54c73119ed643c80b4ed56ef3cf80e209 Mon Sep 17 00:00:00 2001 From: Nick Hill <nickhill123@gmail.com> Date: Tue, 28 Jul 2026 07:58:02 -0700 Subject: [PATCH 175/185] [Core] Warm up runner-owned Triton kernels before the first request (#49903) --- tests/v1/worker/test_kv_block_zeroer.py | 66 +++++++++- vllm/model_executor/warmup/kernel_warmup.py | 16 ++- .../warmup/qwen_triton_warmup.py | 114 ----------------- .../warmup/v1_block_table_warmup.py | 42 ++---- vllm/v1/worker/gpu/warmup.py | 121 ++++++++++++------ vllm/v1/worker/mamba_utils.py | 6 +- vllm/v1/worker/utils.py | 7 +- 7 files changed, 183 insertions(+), 189 deletions(-) diff --git a/tests/v1/worker/test_kv_block_zeroer.py b/tests/v1/worker/test_kv_block_zeroer.py index b212e3ae17b..17aa1bf38d4 100644 --- a/tests/v1/worker/test_kv_block_zeroer.py +++ b/tests/v1/worker/test_kv_block_zeroer.py @@ -4,7 +4,7 @@ import pytest import torch -from vllm.v1.worker.utils import KVBlockZeroer +from vllm.v1.worker.utils import KVBlockZeroer, _zero_kv_blocks_kernel @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -86,3 +86,67 @@ def test_non_uniform_page_sizes(): assert torch.all(storage[1] == 0) assert torch.all(storage[2] == 0) assert torch.all(storage[3] == 1) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_warmup_compiles_every_n_blocks_specialization(): + """After warmup, no launch should trigger a first-request JIT compile. + + ``n_blocks`` is ``do_not_specialize``, so a single warmup launch must + cover every block count. + """ + device = torch.device("cuda") + num_blocks = 64 + page_size_el = 4 + storage = torch.ones((num_blocks, page_size_el), dtype=torch.int32, device=device) + + zeroer = KVBlockZeroer.__new__(KVBlockZeroer) + zeroer.device = device + zeroer._meta = ( + torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device), + torch.tensor([page_size_el], dtype=torch.int64, device=device), + 1, # max_chunks + page_size_el, # blk_size + 1, # n_segs + ) + + def compiled_variants() -> set: + return { + key + for caches in _zero_kv_blocks_kernel.device_caches.values() + for key in caches[0] + } + + zeroer.warmup(num_blocks) + torch.accelerator.synchronize() + warmed = compiled_variants() + assert warmed + + for n_blocks in (1, 2, 3, 16, 32): + zeroer.zero_block_ids(list(range(n_blocks))) + torch.accelerator.synchronize() + + assert compiled_variants() == warmed + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_warmup_respects_available_block_count(): + """An empty KV cache must not be warmed with out-of-range block IDs.""" + device = torch.device("cuda") + page_size_el = 4 + storage = torch.ones((1, page_size_el), dtype=torch.int32, device=device) + + zeroer = KVBlockZeroer.__new__(KVBlockZeroer) + zeroer.device = device + zeroer._meta = ( + torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device), + torch.tensor([page_size_el], dtype=torch.int64, device=device), + 1, + page_size_el, + 1, + ) + + zeroer.warmup(0) + torch.accelerator.synchronize() + + assert torch.all(storage == 1) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index b2c989b295c..afe47f39ced 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -98,12 +98,16 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): minimax_m3_msa_warmup, ) - # Pooling models do not use the generation slot-mapping path. - if not worker.use_v2_model_runner and not worker.model_runner.is_pooling_model: - warm_v1_block_table_kernels( - getattr(worker.model_runner, "device", torch.device("cuda")), - worker.scheduler_config.max_num_batched_tokens, - ) + if not worker.use_v2_model_runner: + # Pooling models do not use the generation slot-mapping path. + if not worker.model_runner.is_pooling_model: + warm_v1_block_table_kernels(worker.model_runner) + # The KV-block zeroing kernel is driven by the scheduler's + # `new_block_ids_to_zero`, so no dummy run ever reaches it. + zeroer = getattr(worker.model_runner, "_kv_block_zeroer", None) + if zeroer is not None: + zeroer.warmup(worker.model_runner.kv_cache_config.num_blocks) + qwen_triton_warmup(worker.model_runner, worker.vllm_config.model_config) # DSv4 mHC TileLang kernels (hc_pre/hc_post/hc_head_op) run every decoder diff --git a/vllm/model_executor/warmup/qwen_triton_warmup.py b/vllm/model_executor/warmup/qwen_triton_warmup.py index 8cbfa539b7e..06a016a5a06 100644 --- a/vllm/model_executor/warmup/qwen_triton_warmup.py +++ b/vllm/model_executor/warmup/qwen_triton_warmup.py @@ -24,24 +24,10 @@ _QWEN_MODEL_TYPES = frozenset( } ) -_ZERO_KV_N_BLOCKS = (1, 2) - -_SLOT_MAPPING_KV_BLOCK_SIZE = 16 -_SLOT_MAPPING_CP_KV_CACHE_INTERLEAVE_SIZE = 1 -_SLOT_MAPPING_BLOCK_TABLE_STRIDES = (1, 3) - # Covers L=1 constexpr, non-divisible runtime L, and divisible runtime L. _FLA_POST_CONV_WARMUP_LENGTHS = (1, 2, 16) -@dataclass(frozen=True) -class _ZeroKvWarmupConfig: - seg_page_sizes: torch.Tensor - max_chunks: int - block_size: int - n_segs: int - - @dataclass(frozen=True) class _QwenGDNWarmupConfig: h: int @@ -148,96 +134,6 @@ def _qwen_gdn_warmup_config( return None -def _get_kv_block_zeroer(runner: object) -> object | None: - zeroer = getattr(runner, "kv_block_zeroer", None) - if zeroer is None: - zeroer = getattr(runner, "_kv_block_zeroer", None) - return zeroer - - -def _zero_kv_warmup_config(runner: object) -> _ZeroKvWarmupConfig | None: - zeroer = _get_kv_block_zeroer(runner) - meta = getattr(zeroer, "_meta", None) - if meta is None: - return None - - _, seg_page_sizes, max_chunks, block_size, n_segs = meta - return _ZeroKvWarmupConfig( - seg_page_sizes=seg_page_sizes, - max_chunks=int(max_chunks), - block_size=int(block_size), - n_segs=int(n_segs), - ) - - -def _warm_zero_kv_blocks_with_runner_zeroer(runner: object) -> bool: - zeroer = _get_kv_block_zeroer(runner) - zero_block_ids = getattr(zeroer, "zero_block_ids", None) - if not callable(zero_block_ids): - return False - - for n_blocks in _ZERO_KV_N_BLOCKS: - zero_block_ids(list(range(n_blocks))) - return True - - -def _warm_zero_kv_blocks_kernel( - device: torch.device, config: _ZeroKvWarmupConfig -) -> None: - from vllm.v1.worker.utils import _zero_kv_blocks_kernel - - max_n_blocks = max(_ZERO_KV_N_BLOCKS) - max_page_size = int(config.seg_page_sizes.max().item()) - scratch = torch.empty( - max_n_blocks * max_page_size, - dtype=torch.int32, - device=device, - ) - seg_addrs = torch.tensor( - [scratch.data_ptr()] * config.n_segs, - dtype=torch.uint64, - device=device, - ) - - for n_blocks in _ZERO_KV_N_BLOCKS: - block_ids = torch.arange(n_blocks, dtype=torch.int64, device=device) - grid = (n_blocks * config.n_segs * config.max_chunks,) - _zero_kv_blocks_kernel[grid]( - seg_addrs, - config.seg_page_sizes, - block_ids, - n_blocks, - N_SEGS=config.n_segs, - MAX_CHUNKS=config.max_chunks, - BLOCK_SIZE=config.block_size, - ) - - -def _warm_compute_slot_mapping_kernel(device: torch.device) -> None: - from vllm.v1.worker.block_table import BlockTable - - # num_tokens/max_num_tokens are do_not_specialize; keep the launch tiny. - num_tokens = 1 - query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) - positions = torch.arange(num_tokens, dtype=torch.int64, device=device) - - for block_table_stride in _SLOT_MAPPING_BLOCK_TABLE_STRIDES: - # Use BlockTable so the JIT key matches the production slot-mapping call. - block_table = BlockTable( - block_size=_SLOT_MAPPING_KV_BLOCK_SIZE, - max_num_reqs=1, - max_num_blocks_per_req=block_table_stride, - max_num_batched_tokens=num_tokens, - pin_memory=False, - device=device, - kernel_block_size=_SLOT_MAPPING_KV_BLOCK_SIZE, - cp_kv_cache_interleave_size=_SLOT_MAPPING_CP_KV_CACHE_INTERLEAVE_SIZE, - ) - block_table.add_row(list(range(block_table_stride)), 0) - block_table.commit_block_table(num_reqs=1) - block_table.compute_slot_mapping(1, query_start_loc, positions) - - def _warm_causal_conv1d_fwd_kernel( device: torch.device, config: _QwenGDNWarmupConfig ) -> None: @@ -373,16 +269,6 @@ def qwen_triton_warmup( device = getattr(runner, "device", torch.device("cuda")) logger.info("Warming up Qwen Triton kernels for model_type=%s.", model_type) - zero_config = _zero_kv_warmup_config(runner) - warmed_zeroer = _warm_zero_kv_blocks_with_runner_zeroer(runner) - if zero_config is not None: - _warm_zero_kv_blocks_kernel(device, zero_config) - elif not warmed_zeroer: - logger.info("Skipping Qwen zero-kv warmup: no KVBlockZeroer metadata.") - - _warm_compute_slot_mapping_kernel(device) - _synchronize_device(device) - compilation_config = getattr(runner, "compilation_config", None) static_forward_context = getattr(compilation_config, "static_forward_context", None) gdn_config = _qwen_gdn_warmup_config(static_forward_context) diff --git a/vllm/model_executor/warmup/v1_block_table_warmup.py b/vllm/model_executor/warmup/v1_block_table_warmup.py index 8d2328432eb..d49e1ba7cc8 100644 --- a/vllm/model_executor/warmup/v1_block_table_warmup.py +++ b/vllm/model_executor/warmup/v1_block_table_warmup.py @@ -2,42 +2,28 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Warm up v1 block-table Triton kernels.""" +from typing import TYPE_CHECKING + import torch +if TYPE_CHECKING: + from vllm.v1.worker.gpu_model_runner import GPUModelRunner + _SLOT_MAPPING_WARMUP_TOKENS = 8 -_SLOT_MAPPING_WARMUP_BLOCK_SIZES = (3, 16) -_SLOT_MAPPING_WARMUP_CP_KV_CACHE_INTERLEAVE_SIZE = 1 -def warm_v1_block_table_kernels( - device: torch.device, - max_tokens: int, -) -> None: - from vllm.v1.worker.block_table import BlockTable +def warm_v1_block_table_kernels(runner: "GPUModelRunner") -> None: + """JIT-compile ``_compute_slot_mapping_kernel`` for the real block tables.""" - num_tokens = max(0, min(_SLOT_MAPPING_WARMUP_TOKENS, max_tokens)) + device = runner.device + block_table = runner.input_batch.block_table + num_tokens = min( + _SLOT_MAPPING_WARMUP_TOKENS, + runner.scheduler_config.max_num_batched_tokens, + ) if num_tokens <= 0: return query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) positions = torch.arange(num_tokens, dtype=torch.int64, device=device) - for block_size in _SLOT_MAPPING_WARMUP_BLOCK_SIZES: - max_num_blocks_per_req = max( - 1, (max(num_tokens, max_tokens) + block_size - 1) // block_size - ) - max_num_blocks_per_req = ((max_num_blocks_per_req + 15) // 16) * 16 - block_table = BlockTable( - block_size=block_size, - max_num_reqs=1, - max_num_blocks_per_req=max_num_blocks_per_req, - max_num_batched_tokens=max(num_tokens, max_tokens), - pin_memory=False, - device=device, - kernel_block_size=block_size, - cp_kv_cache_interleave_size=( - _SLOT_MAPPING_WARMUP_CP_KV_CACHE_INTERLEAVE_SIZE - ), - ) - block_table.add_row(list(range(max_num_blocks_per_req)), 0) - block_table.commit_block_table(1) - block_table.compute_slot_mapping(1, query_start_loc, positions) + block_table.compute_slot_mapping(1, query_start_loc, positions) diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 785188dc3c3..260ca7a5f17 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -157,13 +157,10 @@ def warmup_kernels( worker_execute_model: Callable[[SchedulerOutput], Any], worker_sample_tokens: Callable[[GrammarOutput | None], Any], ) -> None: - """Run two execute_model + sample_tokens iterations to JIT compile - triton kernels. We must call the provided worker's execute_model for - pipeline parallel coordination. + """Run scheduler-realistic prefill and decode steps to JIT compile kernels. - The first iteration simulates a prefill with requests of - decode_query_len + 1 prompt tokens each. The second iteration simulates - a decode step with all requests generating decode_query_len tokens. + We must call the provided worker's execute_model for pipeline parallel + coordination. """ num_spec_steps = model_runner.num_speculative_steps decode_query_len = model_runner.decode_query_len @@ -172,8 +169,13 @@ def warmup_kernels( # a uniform decode batch. prompt_len = decode_query_len + 1 prompt_token_ids = list(range(prompt_len)) - # After prefill, decode generates decode_query_len tokens. - decode_len = prompt_len + decode_query_len + # Upper bound on the decode steps built in `decode_steps` below. + num_decode_steps = 1 + if not model_runner.is_pooling_model: + num_decode_steps = 5 if num_spec_steps > 0 else 3 + # Size the block allocation for the worst case: every request advancing + # decode_query_len tokens on every decode step. + decode_len = prompt_len + num_decode_steps * decode_query_len kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups num_kv_cache_groups = len(kv_cache_groups) @@ -208,9 +210,6 @@ def warmup_kernels( kv_cache_specs = [g.kv_cache_spec for g in kv_cache_groups] prefill_block_counts = [_warmup_block_count(prompt_len, s) for s in kv_cache_specs] decode_block_counts = [_warmup_block_count(decode_len, s) for s in kv_cache_specs] - decode_block_deltas = [ - d - p for d, p in zip(decode_block_counts, prefill_block_counts) - ] max_blocks_per_req = sum(decode_block_counts) num_reqs = min( @@ -243,6 +242,11 @@ def warmup_kernels( nonlocal next_block_id return list(range(next_block_id, next_block_id := next_block_id + num_blocks)) + # The KV-block zeroing kernel is driven by the scheduler's + # new_block_ids_to_zero, so none of the steps below reach it. + if model_runner.kv_block_zeroer is not None: + model_runner.kv_block_zeroer.warmup(model_runner.kv_cache_config.num_blocks) + # Step 1: Prefill all requests with 1 + decode_query_len prompt tokens each. new_reqs = [ NewRequestData.from_request( @@ -287,33 +291,78 @@ def warmup_kernels( worker_sample_tokens(grammar_output) - # Step 2: Decode all requests with decode_query_len tokens each. - cached_req_data = CachedRequestData.make_empty() - cached_req_data.req_ids = list(req_ids) - cached_req_data.num_computed_tokens = [prompt_len] * num_reqs - cached_req_data.num_output_tokens = [1] * num_reqs - new_block = any(decode_block_deltas) - cached_req_data.new_block_ids = [ - tuple(_alloc_blocks(n) for n in decode_block_deltas) if new_block else None - for _ in range(num_reqs) + # Per-request state carried across the decode steps. + req_computed = [prompt_len] * num_reqs + req_blocks = [list(prefill_block_counts) for _ in range(num_reqs)] + + def _run_decode_step(indices: list[int], spec_flags: list[bool]) -> None: + """Decode `indices`, spec-decoding the ones flagged in `spec_flags`.""" + cached_req_data = CachedRequestData.make_empty() + cached_req_data.req_ids = [req_ids[i] for i in indices] + cached_req_data.num_computed_tokens = [req_computed[i] for i in indices] + cached_req_data.num_output_tokens = [1] * len(indices) + cached_req_data.new_block_ids = [] + + step_num_scheduled_tokens: dict[str, int] = {} + step_spec_tokens: dict[str, list[int]] = {} + for i, use_spec in zip(indices, spec_flags): + num_tokens = decode_query_len if use_spec else 1 + after = req_computed[i] + num_tokens + deltas = [ + _warmup_block_count(after, spec) - held + for spec, held in zip(kv_cache_specs, req_blocks[i]) + ] + cached_req_data.new_block_ids.append( + tuple(_alloc_blocks(n) for n in deltas) if any(deltas) else None + ) + req_blocks[i] = [ + held + delta for held, delta in zip(req_blocks[i], deltas) + ] + step_num_scheduled_tokens[req_ids[i]] = num_tokens + if use_spec: + step_spec_tokens[req_ids[i]] = [0] * num_spec_steps + + decode_output = SchedulerOutput.make_empty() + decode_output.scheduled_cached_reqs = cached_req_data + decode_output.num_scheduled_tokens = step_num_scheduled_tokens + decode_output.scheduled_spec_decode_tokens = step_spec_tokens + decode_output.total_num_scheduled_tokens = sum( + step_num_scheduled_tokens.values() + ) + decode_output.num_common_prefix_blocks = [0] * num_kv_cache_groups + + worker_execute_model(decode_output) + worker_sample_tokens(None) + + for i, use_spec in zip(indices, spec_flags): + req_computed[i] += decode_query_len if use_spec else 1 + + all_indices = list(range(num_reqs)) + use_spec_decode = num_spec_steps > 0 + + # Decode steps to warm, as (request indices, per-request spec flag). + # Under spec decoding the scheduler drops requests the drafter proposed + # nothing for, so warm each batch shape with and without draft tokens. + decode_steps: list[tuple[list[int], list[bool]]] = [ + (all_indices, [use_spec_decode] * num_reqs), ] + if num_reqs >= 2: + # Mixed spec / non-spec: GDN and KDA reclassify the non-spec decode + # as a prefill and split the batch into spec/non-spec token indices. + decode_steps.append(([0, 1], [use_spec_decode, False])) + if use_spec_decode: + # Exercise the model paths that split a batch by whether each + # request received draft tokens. + decode_steps.append(([0, 1], [False, False])) + if num_reqs > 1: + decode_steps.append(([0], [use_spec_decode])) + if use_spec_decode: + decode_steps.append(([0], [False])) + elif use_spec_decode: + decode_steps.append(([0], [False])) - decode_output = SchedulerOutput.make_empty() - decode_output.scheduled_cached_reqs = cached_req_data - decode_output.num_scheduled_tokens = { - req_id: decode_query_len for req_id in req_ids - } - if num_spec_steps > 0: - decode_output.scheduled_spec_decode_tokens = { - req_id: [0] * num_spec_steps for req_id in req_ids - } - decode_output.total_num_scheduled_tokens = sum( - decode_output.num_scheduled_tokens.values() - ) - decode_output.num_common_prefix_blocks = [0] * num_kv_cache_groups - - worker_execute_model(decode_output) - worker_sample_tokens(None) + for step_indices, step_spec_flags in decode_steps: + _run_decode_step(step_indices, step_spec_flags) # Clean up - process finish_req_ids. cleanup_output = SchedulerOutput.make_empty() diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 7f611a2eeff..36b72efe35b 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -150,7 +150,7 @@ def _copy_mamba_state_block( tl.store(tail_dst + tail_off, tail_data, mask=tail_mask) -@triton.jit +@triton.jit(do_not_specialize=["num_reqs"]) def postprocess_mamba_fused_kernel( # Decision inputs (per-request) num_accepted_tokens_ptr, @@ -280,7 +280,7 @@ def postprocess_mamba_fused_kernel( ) -@triton.jit +@triton.jit(do_not_specialize=["num_reqs"]) def preprocess_mamba_align_fused_kernel( idx_mapping_ptr, state_idx_ptr, @@ -326,7 +326,7 @@ def preprocess_mamba_align_fused_kernel( tl.store(num_accepted_tokens_ptr + req_indices, 1, mask=mask & should_reset) -@triton.jit +@triton.jit(do_not_specialize=["num_reqs"]) def precopy_mamba_align_fused_kernel( # Per-request-slot inputs (indexed by req_idx via idx_mapping), produced by # the V2 fused align preprocess kernel for the current step: diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index d00974046fc..4afb941c12a 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -40,7 +40,7 @@ from vllm.v1.kv_cache_interface import ( logger = init_logger(__name__) -@triton.jit +@triton.jit(do_not_specialize=["n_blocks"]) def _zero_kv_blocks_kernel( seg_addrs_ptr, seg_page_sizes_ptr, @@ -206,6 +206,11 @@ class KVBlockZeroer: BLOCK_SIZE=blk_size, ) + def warmup(self, num_kv_blocks: int) -> None: + """JIT-compile the zeroing kernel before the first real request.""" + if num_kv_blocks > 0: + self.zero_block_ids([0]) + @dataclass class AttentionGroup: From 30217b0e809ef045b3b4aa7c011d6c6d2b0b0362 Mon Sep 17 00:00:00 2001 From: Itay Etelis <92247226+Etelis@users.noreply.github.com> Date: Tue, 28 Jul 2026 18:01:55 +0300 Subject: [PATCH 176/185] [Bugfix][KV Offload][P2P] Scope serve state to fetch rounds (#49877) Signed-off-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Itay Etelis <Itay.etelis@gmail.com> --- .../v1/kv_offload/tiering/p2p/test_manager.py | 13 +- .../kv_offload/tiering/p2p/test_sessions.py | 116 ++++- .../kv_offload/tiering/p2p/session/client.py | 293 +++++------ .../tiering/p2p/session/protocol.py | 48 +- .../kv_offload/tiering/p2p/session/server.py | 474 ++++++++++-------- .../kv_offload/tiering/p2p/session/session.py | 16 +- 6 files changed, 570 insertions(+), 390 deletions(-) diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index cd6e8f382a5..c8a9d7f3ea3 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -1612,16 +1612,17 @@ class TestBindHostPortDefaults: monkeypatch.setattr( manager_module, "NixlTransport", - lambda agent_name, *a, **k: calls.update(nixl_name=agent_name) - or SimpleNamespace(), + lambda agent_name, *a, **k: ( + calls.update(nixl_name=agent_name) or SimpleNamespace() + ), ) monkeypatch.setattr( manager_module, "ZmqTransport", - lambda local_id, host, port, *a, **k: calls.update( - zmq_id=local_id, zmq_host=host, zmq_port=port - ) - or SimpleNamespace(), + lambda local_id, host, port, *a, **k: ( + calls.update(zmq_id=local_id, zmq_host=host, zmq_port=port) + or SimpleNamespace() + ), ) spec = SimpleNamespace( blocks_per_chunk=1, diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py index cab4ce6705a..20e4bad2186 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py +++ b/tests/v1/kv_offload/tiering/p2p/test_sessions.py @@ -49,6 +49,7 @@ from vllm.v1.kv_offload.tiering.p2p.session.protocol import ( from vllm.v1.kv_offload.tiering.p2p.session.server import ( _CANCEL_DRAIN_TIMEOUT_S, _InflightXfer, + _OutboundRequestState, ) from vllm.v1.kv_offload.tiering.p2p.session.session import ( _MAX_CONSECUTIVE_DISPATCH_ERRORS, @@ -315,10 +316,26 @@ def _activate( # either a missing entry or a None field. These helpers paper over that. +def _client_load(session: P2PSession, kv_request_id: str): + """The single in-flight load of a kv_request_id (loads are per-round).""" + loads = session._client._requests[kv_request_id].loads + assert len(loads) == 1 + return next(iter(loads.values())) + + def _srv_outbound(session: P2PSession, kv_request_id: str): - """Outbound serve state for a kv_request_id, or None (idle / GC'd).""" + """Serve-side round for a kv_request_id, or None (idle / GC'd). + + Rounds are keyed by wire round_seq; surfaces the demanded round when + a fetch has bound one, else any parked supply round. + """ st = session._server._requests.get(kv_request_id) - return st.outbound if st is not None else None + if st is None or not st.outbound: + return None + for rnd in st.outbound.values(): + if rnd.demand_received: + return rnd + return next(iter(st.outbound.values())) def _srv_lookups(session: P2PSession) -> list: @@ -330,8 +347,10 @@ def _srv_lookups(session: P2PSession) -> list: def _srv_abort_started(session: P2PSession, kv_request_id: str) -> float | None: """Pending-abort start time for a kv_request_id, or None.""" - st = session._server._requests.get(kv_request_id) - return st.abort_started_at if st is not None else None + for (kv, _), started in session._server._pending_aborts.items(): + if kv == kv_request_id: + return started + return None def _srv_inflight_count(session: P2PSession, kv_request_id: str) -> int: @@ -479,6 +498,7 @@ class TestClientFlows: conn.enqueue( { TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-1", TransferDoneMsg.SUCCESS: True, } @@ -495,6 +515,7 @@ class TestClientFlows: conn.enqueue( { TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-1", TransferDoneMsg.SUCCESS: False, } @@ -538,6 +559,7 @@ class TestClientFlows: conn.enqueue( { TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-1", TransferDoneMsg.SUCCESS: True, } @@ -552,7 +574,7 @@ class TestClientFlows: session.request_blocks( job_id=1, kv_request_id="req-1", keys=[b"k"], block_ids=[0] ) - session._client._requests["req-1"].load.submitted_at = time.monotonic() - 60.0 + _client_load(session, "req-1").submitted_at = time.monotonic() - 60.0 session.poll() abort = conn._sent[-1] assert abort[TYPE_KEY] == AbortFetchMsg.TYPE @@ -569,7 +591,7 @@ class TestClientFlows: job_id=7, kv_request_id="req-7", keys=[b"k"], block_ids=[0] ) # 1) Trip the load timeout to send AbortFetch and stamp aborted_at. - session._client._requests["req-7"].load.submitted_at = ( + _client_load(session, "req-7").submitted_at = ( time.monotonic() - _LOAD_TIMEOUT_S - 1.0 ) loads = session.poll().loads @@ -579,11 +601,11 @@ class TestClientFlows: and m[AbortFetchMsg.KV_REQUEST_ID] == "req-7" for m in conn._sent ) - assert session._client._requests["req-7"].load.aborted_at is not None + assert _client_load(session, "req-7").aborted_at is not None # 2) Now backdate aborted_at past the abort-ack timeout. No ack ever # arrived from the peer. - session._client._requests["req-7"].load.aborted_at = ( + _client_load(session, "req-7").aborted_at = ( time.monotonic() - _ABORT_ACK_TIMEOUT_S - 1.0 ) loads = session.poll().loads @@ -599,17 +621,18 @@ class TestClientFlows: session.request_blocks( job_id=8, kv_request_id="req-8", keys=[b"k"], block_ids=[0] ) - session._client._requests["req-8"].load.submitted_at = ( + _client_load(session, "req-8").submitted_at = ( time.monotonic() - _LOAD_TIMEOUT_S - 1.0 ) # First poll: AbortFetch goes out. session.poll() - assert session._client._requests["req-8"].load.aborted_at is not None + assert _client_load(session, "req-8").aborted_at is not None # Peer acks the abort. conn.enqueue( { TYPE_KEY: AbortAckMsg.TYPE, + AbortAckMsg.ROUND_SEQ: 0, AbortAckMsg.KV_REQUEST_ID: "req-8", } ) @@ -915,6 +938,7 @@ class TestLookupFlow: conn.enqueue( { TYPE_KEY: LookupMsg.TYPE, + LookupMsg.ROUND_SEQ: 0, LookupMsg.KV_REQUEST_ID: "req-1", LookupMsg.KEYS: [b"hX", b"hY", b"hZ"], } @@ -951,6 +975,7 @@ def _send_lookup(conn: FakeConnection, kv_request_id: str, keys: list[bytes]): conn.enqueue( { TYPE_KEY: LookupMsg.TYPE, + LookupMsg.ROUND_SEQ: 0, LookupMsg.KV_REQUEST_ID: kv_request_id, LookupMsg.KEYS: list(keys), } @@ -1201,6 +1226,7 @@ class TestServerLookupHandling: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [], FetchMsg.BLOCK_INDEXES: [], @@ -1240,6 +1266,7 @@ class TestServerLookupHandling: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"hA", b"hB"], FetchMsg.BLOCK_INDEXES: [20, 21], @@ -1277,6 +1304,7 @@ class TestServerFlows: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], @@ -1295,6 +1323,7 @@ class TestServerFlows: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], @@ -1313,6 +1342,7 @@ class TestServerFlows: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], @@ -1331,6 +1361,7 @@ class TestServerFlows: conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) @@ -1349,13 +1380,19 @@ class TestServerFlows: tid = 42 session._server._inflight_add( tid, - _InflightXfer(kv_request_id="req-1", block_count=1, job_ids={1}), + _InflightXfer( + kv_request_id="req-1", + block_count=1, + job_ids={1}, + round=_OutboundRequestState(inflight=1), + ), ) transport._cancel_still_inflight.add(tid) conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) @@ -1376,13 +1413,19 @@ class TestServerFlows: tid = 42 session._server._inflight_add( tid, - _InflightXfer(kv_request_id="req-1", block_count=1, job_ids={1}), + _InflightXfer( + kv_request_id="req-1", + block_count=1, + job_ids={1}, + round=_OutboundRequestState(inflight=1), + ), ) transport._cancel_still_inflight.add(tid) conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) @@ -1409,20 +1452,26 @@ class TestServerFlows: tid = 42 session._server._inflight_add( tid, - _InflightXfer(kv_request_id="req-1", block_count=1, job_ids={1}), + _InflightXfer( + kv_request_id="req-1", + block_count=1, + job_ids={1}, + round=_OutboundRequestState(inflight=1), + ), ) transport._cancel_still_inflight.add(tid) conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) session.poll() assert _srv_abort_started(session, "req-1") is not None # Backdate past the drain deadline. - session._server._requests["req-1"].abort_started_at = ( + session._server._pending_aborts[("req-1", 0)] = ( time.monotonic() - _CANCEL_DRAIN_TIMEOUT_S - 1.0 ) # Even if the transport still claims it can't cancel, the @@ -1445,13 +1494,19 @@ class TestServerFlows: tid = 42 session._server._inflight_add( tid, - _InflightXfer(kv_request_id="req-1", block_count=1, job_ids={1}), + _InflightXfer( + kv_request_id="req-1", + block_count=1, + job_ids={1}, + round=_OutboundRequestState(inflight=1), + ), ) transport._cancel_still_inflight.add(tid) conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) @@ -1463,6 +1518,7 @@ class TestServerFlows: conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) @@ -1497,6 +1553,7 @@ class TestServerFlows: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], @@ -1529,6 +1586,7 @@ class TestServerFlows: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], @@ -1571,6 +1629,7 @@ class TestFinishRequestServerSide: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], @@ -1596,6 +1655,7 @@ class TestFinishRequestServerSide: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], @@ -1631,6 +1691,7 @@ class TestFinishRequestServerSide: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], @@ -1665,6 +1726,7 @@ class TestFinishRequestServerSide: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], @@ -1686,6 +1748,7 @@ class TestFinishRequestServerSide: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-2", FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], @@ -1719,6 +1782,7 @@ class TestFinishRequestServerSide: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"demand"], FetchMsg.BLOCK_INDEXES: [5], @@ -1753,6 +1817,7 @@ class TestFinishRequestServerSide: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], @@ -1790,6 +1855,7 @@ class TestFinishRequestServerSide: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], @@ -1824,6 +1890,7 @@ class TestFinishRequestServerSide: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1", b"k2", b"k3"], FetchMsg.BLOCK_INDEXES: [10, 11, 12], @@ -1885,6 +1952,7 @@ class TestFinishRequestServerSide: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], @@ -1951,6 +2019,7 @@ class TestBidirectional: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-srv", FetchMsg.KEYS: [b"served"], FetchMsg.BLOCK_INDEXES: [7], @@ -1981,6 +2050,7 @@ class TestBidirectional: conn.enqueue( { TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-cli", TransferDoneMsg.SUCCESS: True, } @@ -2168,6 +2238,7 @@ class TestAdversarial: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-bad", FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [1], @@ -2216,6 +2287,7 @@ class TestDispatchErrorHandling: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-bad", FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [1], @@ -2246,6 +2318,7 @@ class TestDispatchErrorHandling: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [0], @@ -2269,6 +2342,7 @@ class TestDispatchErrorHandling: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [0], @@ -2296,6 +2370,7 @@ class TestDispatchErrorHandling: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [0], @@ -2341,6 +2416,7 @@ class TestInflightPerReqInvariant: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: kv_id, FetchMsg.KEYS: keys, FetchMsg.BLOCK_INDEXES: indexes, @@ -2394,7 +2470,12 @@ class TestInflightPerReqInvariant: tid = kv_id_idx * 10 + j session._server._inflight_add( tid, - _InflightXfer(kv_request_id=kv_id, block_count=1, job_ids={tid}), + _InflightXfer( + kv_request_id=kv_id, + block_count=1, + job_ids={tid}, + round=_OutboundRequestState(inflight=1), + ), ) assert _srv_total_inflight(session) == len(session._server._inflight) assert session._server._has_inflight_for("req-0") @@ -2488,6 +2569,7 @@ class TestFetchMsgValidation: def _valid_msg(self) -> dict: return { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [0, 1], @@ -2513,6 +2595,7 @@ class TestTransferDoneMsgValidation: def test_valid_message_passes(self): msg = { TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-1", TransferDoneMsg.SUCCESS: True, } @@ -2521,6 +2604,7 @@ class TestTransferDoneMsgValidation: def test_success_wrong_type(self): msg = { TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-1", TransferDoneMsg.SUCCESS: 1, } diff --git a/vllm/v1/kv_offload/tiering/p2p/session/client.py b/vllm/v1/kv_offload/tiering/p2p/session/client.py index ed91f7f9b00..07a4620780e 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/client.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/client.py @@ -11,7 +11,6 @@ callback injected by the coordinator (which gates on ConnectAck). from __future__ import annotations -import enum import time from collections.abc import Callable, Sequence from dataclasses import dataclass, field @@ -35,26 +34,12 @@ _LOAD_TIMEOUT_S = 30.0 _ABORT_ACK_TIMEOUT_S = 10.0 -class ClientPhase(enum.Enum): - """Lifecycle of a request's client-side lookup/fetch signalling. - - Advances monotonically. Only ``finish`` reads it, to decide - whether a terminal empty FetchMsg is owed to release the peer's - lookup state (owed only from ``PROBING``: a LookupMsg went out but no - FetchMsg has since closed the peer's lookup phase). - """ - - REGISTERED = enum.auto() # keys in probes/unsent, nothing sent yet - PROBING = enum.auto() # LookupMsg flushed, awaiting responses - FETCH_SENT = enum.auto() # FetchMsg sent (real or terminal empty) - - @dataclass class _InboundLoadState: """Client-role state for a single in-flight load request. - Lives on ``_ClientRequestState.load`` for the duration of a fetch; - the owning kv_request_id is the dict key, so it isn't stored here. + Lives in ``_ClientRequestState.loads`` keyed by round_seq for the + duration of a fetch; the owning kv_request_id is the outer dict key. """ job_id: int # opaque ID assigned by the manager to this load request @@ -68,7 +53,7 @@ class _ClientRequestState: One entry per kv_request_id we're driving. Lookup-phase fields are used only by symmetric P2P (``do_p2p_fetch``); PD-only loads leave - ``probes``/``unsent`` empty and drive just ``phase`` and ``load``. An + ``probes``/``unsent`` empty and drive just ``phase`` and ``loads``. An entry is dropped once every field is idle — see ``ClientRole._maybe_prune``. """ @@ -82,11 +67,23 @@ class _ClientRequestState: # OffloadKeys registered but not yet flushed onto the wire. Drained and # cleared by the next flush_pending_lookups. unsent: list[OffloadKey] = field(default_factory=list) + # Current lookup round. LookupMsgs carry it, each fetch closes it and + # advances it, so every round's supply/demand/completion is isolated + # on the wire. PD clients never probe and stay on round 0. + round_seq: int = 0 + # This id ran the symmetric lookup phase (register_lookup); a fetch + # with keys then requires every key to be a confirmed probe. PD + # loads never probe. + probed: bool = False - # Monotonic lookup/fetch signalling phase; see ``ClientPhase``. - phase: ClientPhase = ClientPhase.REGISTERED - # Set while a fetch is in flight; cleared on completion/abort/timeout. - load: _InboundLoadState | None = None + # The peer holds lookup state no FetchMsg has closed: a LookupMsg + # was flushed since the last fetch. finish owes a terminal empty + # FetchMsg while set, so the peer releases parked supply. + peer_lookup_open: bool = False + # In-flight loads keyed by the round their fetch carried. The + # scheduler submits loads incrementally as chunks resolve, so several + # can be in flight at once; TransferDone/AbortAck match by round. + loads: dict[int, _InboundLoadState] = field(default_factory=dict) class LoadResult(NamedTuple): @@ -133,11 +130,10 @@ class ClientRole: # _serve_pending. Populated by register_lookup, drained by # flush_pending_lookups, and discarded on finish/close. self._flush_pending: set[str] = set() - # kv_request_ids with a fetch in flight (``st.load is not None``) — - # the work-list collect_results walks for timeouts, and the - # has_active_loads predicate, instead of scanning every request. - # Kept in exact sync with ``st.load``: armed in request_blocks, - # discarded wherever load is cleared, and cleared on close. + # kv_request_ids with at least one fetch in flight — the work-list + # collect_results walks for timeouts, and the has_active_loads + # predicate, instead of scanning every request. Kept in exact sync + # with ``st.loads``. self._active_loads: set[str] = set() self._completed_loads: list[LoadResult] = [] @@ -156,17 +152,21 @@ class ClientRole: def _maybe_prune(self, kv_request_id: str) -> None: """Drop the entry once it holds no live load or lookup state. - The sticky ``phase`` is only read by ``finish``. A probe - clears when its fetch is issued (``request_blocks``) or when the - request finishes (``finish``/``close``); in the former case - ``load`` is set and keeps the entry alive, in the latter the phase - is no longer needed — so dropping on emptiness never loses a phase - still in use. + ``peer_lookup_open`` is only read by ``finish``, and every path + that clears the last probe (fetch / finish / close) also settles + it, so dropping on emptiness never loses a flag still in use. """ st = self._requests.get(kv_request_id) - if st is not None and st.load is None and not st.probes and not st.unsent: + if st is not None and not st.loads and not st.probes and not st.unsent: del self._requests[kv_request_id] + def _on_load_terminal(self, kv_request_id: str, st: _ClientRequestState) -> None: + """Wind down id-level state once no load remains in flight.""" + if st.loads: + return + self._active_loads.discard(kv_request_id) + self._maybe_prune(kv_request_id) + @property def has_active_loads(self) -> bool: """True if any kv_request_id has a fetch in flight.""" @@ -184,7 +184,12 @@ class ClientRole: block_ids: Sequence[int], send_ready: bool, ) -> None: - """Register a load request and send the FetchMsg.""" + """Send the FetchMsg closing the current lookup round. + + The scheduler may submit several loads per kv_request_id as its + matched prefix resolves incrementally; each fetch carries the + round it closes so the loads stay independent on the wire. + """ logger.debug( "P2PSession %s: request_blocks job_id=%d kv_request_id=%s " "blocks=%d ready=%s", @@ -195,77 +200,74 @@ class ClientRole: send_ready, ) st = self._get_or_create_request(kv_request_id) - st.load = _InboundLoadState( + round_seq = st.round_seq + st.round_seq += 1 + st.loads[round_seq] = _InboundLoadState( job_id=job_id, submitted_at=time.monotonic(), ) self._active_loads.add(kv_request_id) - st.phase = ClientPhase.FETCH_SENT + st.peer_lookup_open = False self._send( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: kv_request_id, FetchMsg.KEYS: list(keys), FetchMsg.BLOCK_INDEXES: [int(idx) for idx in block_ids], + FetchMsg.ROUND_SEQ: round_seq, } ) - # Issuing the fetch ends this request's lookup phase, so drop all - # probe state. Once the peer serves this fetch both sides unpin, so - # the producer may evict the block; a stale cached True would - # otherwise let a re-scheduled lookup() return HIT without - # re-probing, pointing at a block the producer no longer holds. - # Clearing forces a fresh LookupMsg on re-schedule so the producer - # answers from current state. For a symmetric-P2P request (probes - # populated) every fetched block was a confirmed HIT; a PD-only load - # never probes, so probes is empty and the clear is a no-op. - if st.probes: + # Issuing the fetch closes this lookup round, so drop all probe + # state. Once the peer serves the fetch both sides unpin, so a + # stale cached True would let a later lookup() return HIT for a + # block the producer may have evicted; clearing forces a fresh + # probe under the next round. + if st.probed and keys: + assert st.probes, ( + f"symmetric fetch for {kv_request_id} has keys but no probes" + ) assert all(st.probes.get(key) is True for key in keys) st.probes.clear() def finish(self, kv_request_id: str) -> None: - """Finish a request: abort any in-flight load and release lookup state. + """Finish a request: abort in-flight loads and release lookup state. - Called from the session's ``finish_request``. The two branches are - mutually exclusive: ``load`` is set only by ``request_blocks``, which - also advances ``phase`` to ``FETCH_SENT``, and nothing moves it back - to ``PROBING`` — so a fetch in flight never coexists with the - ``PROBING`` phase. - - - Fetch in flight (``load`` set, phase ``FETCH_SENT``): send an - AbortFetchMsg unless the load is already aborting, then drop it. - - Outstanding lookups (``PROBING``): a LookupMsg was flushed but no - FetchMsg has closed the peer's lookup phase. Every FetchMsg the - server receives in p2p mode is its "request finished" signal (it - releases lookup state and fires ``cb.finish_request``); when the - client's lookups all missed no FetchMsg is otherwise sent, so emit - a terminal empty one purely to trigger those semantics. In - ``REGISTERED`` the peer never received a LookupMsg and in - ``FETCH_SENT`` a FetchMsg already closed the phase, so neither owes - a terminal FetchMsg. + Called from the session's ``finish_request``. Sends an + AbortFetchMsg per load not already aborting, and — independently — + the terminal empty FetchMsg when the peer still holds lookup + state no fetch has closed (its "request finished" signal: it + releases lookup state, drains parked supply, and fires + ``cb.finish_request``). A later round's supply can be parked + while an earlier round's load is still in flight, so both can be + owed at once. Then drop all probe/lookup state and prune the entry. """ st = self._requests.get(kv_request_id) if st is None: return - if st.load is not None: - if st.load.aborted_at is None: + if st.loads: + for round_seq, load in st.loads.items(): + if load.aborted_at is not None: + continue self._send( { TYPE_KEY: AbortFetchMsg.TYPE, AbortFetchMsg.KV_REQUEST_ID: kv_request_id, + AbortFetchMsg.ROUND_SEQ: round_seq, } ) - st.load = None + st.loads.clear() self._active_loads.discard(kv_request_id) - elif st.phase is ClientPhase.PROBING: - st.phase = ClientPhase.FETCH_SENT + if st.peer_lookup_open: + st.peer_lookup_open = False self._send( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: kv_request_id, FetchMsg.KEYS: [], FetchMsg.BLOCK_INDEXES: [], + FetchMsg.ROUND_SEQ: st.round_seq, } ) st.probes.clear() @@ -273,20 +275,21 @@ class ClientRole: self._flush_pending.discard(kv_request_id) self._maybe_prune(kv_request_id) - def on_transfer_done(self, kv_request_id: str, success: bool) -> None: + def on_transfer_done( + self, kv_request_id: str, success: bool, round_seq: int + ) -> None: """Handle a TransferDoneMsg from the peer.""" st = self._requests.get(kv_request_id) - if st is not None and st.load is not None: + load = st.loads.pop(round_seq, None) if st is not None else None + if st is not None and load is not None: self._completed_loads.append( LoadResult( - job_id=st.load.job_id, + job_id=load.job_id, kv_request_id=kv_request_id, success=success, ) ) - st.load = None - self._active_loads.discard(kv_request_id) - self._maybe_prune(kv_request_id) + self._on_load_terminal(kv_request_id, st) else: # No matching in-flight load: either a duplicate # transfer_done from the peer (protocol violation) or a @@ -295,41 +298,44 @@ class ClientRole: # so we can't tell — log so it's findable. logger.warning( "P2PSession %s: transfer_done for unknown kv_request_id=%s " - "(duplicate from peer, or raced with local cancel/timeout)", + "round=%s (duplicate from peer, or raced with local " + "cancel/timeout)", self._peer_id, kv_request_id, + round_seq, ) - def on_abort_ack(self, kv_request_id: str) -> None: + def on_abort_ack(self, kv_request_id: str, round_seq: int) -> None: """Handle an AbortAckMsg from the peer.""" st = self._requests.get(kv_request_id) - if st is not None and st.load is not None: + load = st.loads.pop(round_seq, None) if st is not None else None + if st is not None and load is not None: logger.warning( "P2PSession %s: load request %s (job_id=%d) timed out; " "load job completed with failure. If this recurs, ensure " "PYTHONHASHSEED is set to the same value on all nodes.", self._peer_id, kv_request_id, - st.load.job_id, + load.job_id, ) self._completed_loads.append( LoadResult( - job_id=st.load.job_id, + job_id=load.job_id, kv_request_id=kv_request_id, success=False, ) ) - st.load = None - self._active_loads.discard(kv_request_id) - self._maybe_prune(kv_request_id) + self._on_load_terminal(kv_request_id, st) else: # See on_transfer_done: same ambiguity (duplicate ack # vs. raced with local cancel/timeout that already popped). logger.warning( "P2PSession %s: abort_ack for unknown kv_request_id=%s " - "(duplicate from peer, or raced with local cancel/timeout)", + "round=%s (duplicate from peer, or raced with local " + "cancel/timeout)", self._peer_id, kv_request_id, + round_seq, ) # ------------------------------------------------------------------ @@ -345,17 +351,18 @@ class ClientRole: - Once a LookupRespMsg has resolved the entry: returns the cached bool result on every call without popping it. - A resolved entry is retained until its fetch is issued - (``request_blocks`` pops it) or the request finishes + A resolved entry is retained until a fetch closes the round + (``request_blocks`` clears all probes) or the request finishes (``finish`` clears all entries for the id). A request's block set can be re-probed across steps, so popping on read would make a repeat probe of an already-resolved key look brand-new and re-queue it, emitting a redundant LookupMsg for an answer we already hold. Keeping the entry until fetch makes repeat probes - free; clearing it at fetch forces a fresh probe if the request is - re-scheduled, since the block is unpinned once served. + free; clearing at fetch forces a fresh probe under the next + round, since the block is unpinned once served. """ st = self._get_or_create_request(kv_request_id) + st.probed = True okey = OffloadKey(key) if okey in st.probes: return st.probes[okey] @@ -378,14 +385,11 @@ class ClientRole: ``on_schedule_end()``. A request's block set may be discovered across several scheduler steps, so more than one LookupMsg can go out per kv_request_id — one per step that registered new - keys. register_lookup() de-dups in-flight and already-resolved - (req_id, key) pairs, so each LookupMsg carries only the keys - first probed in that step. The peer's lookup phase for the id is - still closed by exactly one FetchMsg, which the client contract - guarantees is sent after every lookup for the id has resolved - (see request_blocks / finish). Send-gating is handled by - the injected ``_send`` callback (queues until ConnectAckMsg if - needed). + keys, all tagged with the current round. register_lookup() + de-dups in-flight and already-resolved (req_id, key) pairs, so + each LookupMsg carries only the keys first probed in that step. + Send-gating is handled by the injected ``_send`` callback + (queues until ConnectAckMsg if needed). Only requests that registered new keys since the last flush are visited — the ``_flush_pending`` work-list avoids scanning every @@ -395,13 +399,9 @@ class ClientRole: st = self._requests.get(req_id) if st is None or not st.unsent: continue - # Record that the peer now holds lookup state for this id so - # finish knows a terminal empty FetchMsg may be owed. - # Only promote from REGISTERED: once a fetch has gone out - # (FETCH_SENT) a later LookupMsg must not regress the phase, as - # no terminal FetchMsg is owed for an already-fetched request. - if st.phase is ClientPhase.REGISTERED: - st.phase = ClientPhase.PROBING + # The peer now holds lookup state for this id; finish owes a + # terminal empty FetchMsg until a fetch closes it. + st.peer_lookup_open = True logger.debug( "P2P LOOKUP client %s: SEND LookupMsg kv_request_id=%s keys=%d", self._peer_id, @@ -413,6 +413,7 @@ class ClientRole: TYPE_KEY: LookupMsg.TYPE, LookupMsg.KV_REQUEST_ID: req_id, LookupMsg.KEYS: list(st.unsent), + LookupMsg.ROUND_SEQ: st.round_seq, } ) st.unsent = [] @@ -451,52 +452,56 @@ class ClientRole: def collect_results(self) -> list[LoadResult]: """Walk load timeouts and drain completed loads. - Active requests past ``_LOAD_TIMEOUT_S`` get an AbortFetchMsg - sent and enter the aborting phase. Aborting requests past - ``_ABORT_ACK_TIMEOUT_S`` are surfaced as failed loads. + Loads past ``_LOAD_TIMEOUT_S`` get an AbortFetchMsg sent and + enter the aborting phase. Aborting loads past + ``_ABORT_ACK_TIMEOUT_S`` are surfaced as failed. Lookups have no timeout: an unanswered probe stays None (RETRY) until finish_request clears it — see ``_ClientRequestState.probes``. """ now = time.monotonic() - to_remove: list[str] = [] + to_remove: list[tuple[str, int]] = [] for req_id in self._active_loads: st = self._requests[req_id] - assert st.load is not None - load = st.load - if load.aborted_at is None: - if now - load.submitted_at >= _LOAD_TIMEOUT_S: - load.aborted_at = now - logger.warning( - "P2PSession %s: %s timed out, sending abort", - self._peer_id, - req_id, - ) - self._send( - { - TYPE_KEY: AbortFetchMsg.TYPE, - AbortFetchMsg.KV_REQUEST_ID: req_id, - } - ) - else: - if now - load.aborted_at >= _ABORT_ACK_TIMEOUT_S: - to_remove.append(req_id) - self._completed_loads.append( - LoadResult( - job_id=load.job_id, - kv_request_id=req_id, - success=False, + assert st.loads + for round_seq, load in st.loads.items(): + if load.aborted_at is None: + if now - load.submitted_at >= _LOAD_TIMEOUT_S: + load.aborted_at = now + logger.warning( + "P2PSession %s: %s round=%s timed out, sending abort", + self._peer_id, + req_id, + round_seq, ) - ) - logger.warning( - "P2PSession %s: abort_ack timed out for kv_request_id=%s", - self._peer_id, - req_id, - ) - for req_id in to_remove: - self._requests[req_id].load = None - self._active_loads.discard(req_id) - self._maybe_prune(req_id) + self._send( + { + TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.KV_REQUEST_ID: req_id, + AbortFetchMsg.ROUND_SEQ: round_seq, + } + ) + else: + if now - load.aborted_at >= _ABORT_ACK_TIMEOUT_S: + to_remove.append((req_id, round_seq)) + self._completed_loads.append( + LoadResult( + job_id=load.job_id, + kv_request_id=req_id, + success=False, + ) + ) + logger.warning( + "P2PSession %s: abort_ack timed out for " + "kv_request_id=%s round=%s", + self._peer_id, + req_id, + round_seq, + ) + for req_id, round_seq in to_remove: + st = self._requests[req_id] + st.loads.pop(round_seq) + self._on_load_terminal(req_id, st) results = self._completed_loads self._completed_loads = [] @@ -510,12 +515,12 @@ class ClientRole: forever on an answer that can never arrive). See ``ClientCloseResult``. """ failed_jobs = [ - st.load.job_id for st in self._requests.values() if st.load is not None + load.job_id for st in self._requests.values() for load in st.loads.values() ] failed_req_ids = [ req_id for req_id, st in self._requests.items() - if st.load is not None or any(hit is None for hit in st.probes.values()) + if st.loads or any(hit is None for hit in st.probes.values()) ] self._requests.clear() self._flush_pending.clear() diff --git a/vllm/v1/kv_offload/tiering/p2p/session/protocol.py b/vllm/v1/kv_offload/tiering/p2p/session/protocol.py index 8988c7e79f3..33ae5bf2f79 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/protocol.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/protocol.py @@ -26,14 +26,12 @@ Block Transfer Flow (happy path) 1. Client sends FetchMsg with a kv_request_id and lists of block keys + remote indexes where it wants the data written. - In p2p mode FetchMsg is also the server-side "request finished" - signal for the id: no further ``cb.create_store_job`` will fire - (parked LookupMsg batches are popped, so pending-key resolution - cannot promote a HIT after this point), all server-side lookup - state for the id is released, and ``cb.finish_request`` fires on - each dropped batch. The client emits exactly one FetchMsg per - lookup-touched request, including an empty one when no blocks - end up being fetched. + A request may run several lookup→fetch rounds; symmetric-P2P + messages carry ROUND_SEQ so each round's supply, demand, and + completion stay isolated. The terminal empty FetchMsg is the + server-side "request finished" signal for the id: parked + LookupMsg batches are popped and ``cb.finish_request`` fires on + each. 2. Server matches requested blocks against locally stored blocks: - Blocks already available are transferred immediately via RDMA. - Blocks not yet available are recorded as "demanded" and @@ -178,33 +176,32 @@ class DisconnectMsg: class FetchMsg: - """Client → Server: request blocks by key and close the lookup phase. + """Client → Server: request blocks for one lookup round. - In p2p mode FetchMsg is also the server-side "request finished" - signal for ``kv_request_id``: on receipt the server (a) fires no - further ``cb.create_store_job`` for this id — parked LookupMsg - batches are popped, so ``_resolve_pending_lookups`` cannot promote - a HIT_PENDING / RETRY key into a fresh pin after this point — and - (b) calls ``cb.finish_request(batch.ctx)`` on each dropped batch - so the TieringManager can release per-batch bookkeeping. In the - all-miss case the client emits an empty FetchMsg (``KEYS`` - and ``BLOCK_INDEXES`` both empty) purely to fire this signal. + A non-empty fetch closes only its round. The terminal empty FetchMsg + (``KEYS`` and ``BLOCK_INDEXES`` both empty) is the "request + finished" signal: the server pops parked LookupMsg batches, calls + ``cb.finish_request`` on each, and drains any leftover supply. Fields: KV_REQUEST_ID: Identifies this block transfer request. KEYS: List of block keys (OffloadKey bytes). May be empty. BLOCK_INDEXES: List of remote block indexes (same length as KEYS). + ROUND_SEQ: Lookup round this fetch closes. PD clients never probe + and stay on their single round 0. """ TYPE = "fetch" KV_REQUEST_ID = "kv_request_id" KEYS = "keys" BLOCK_INDEXES = "block_indexes" + ROUND_SEQ = "round_seq" @staticmethod def validate(msg: dict) -> None: """Raise ValueError if any field has an invalid type or value.""" _require(msg, FetchMsg.KV_REQUEST_ID, str) + _require_non_neg_int(msg, FetchMsg.ROUND_SEQ) _require_list(msg, FetchMsg.KEYS) _require_list(msg, FetchMsg.BLOCK_INDEXES) keys = msg[FetchMsg.KEYS] @@ -229,17 +226,21 @@ class LookupMsg: Fields: KV_REQUEST_ID: Identifies this lookup transaction. KEYS: List of block keys (OffloadKey bytes) to probe. + ROUND_SEQ: Lookup round these probes belong to; pinned supply is + parked under it for that round's fetch. """ TYPE = "lookup" KV_REQUEST_ID = "kv_request_id" KEYS = "keys" + ROUND_SEQ = "round_seq" @staticmethod def validate(msg: dict) -> None: """Raise ValueError if any field has an invalid type or value.""" _require(msg, LookupMsg.KV_REQUEST_ID, str) _require_list(msg, LookupMsg.KEYS) + _require_non_neg_int(msg, LookupMsg.ROUND_SEQ) class LookupRespMsg: @@ -284,17 +285,22 @@ class TransferDoneMsg: Fields: KV_REQUEST_ID: The request that completed. SUCCESS: Whether the transfer completed successfully. + ROUND_SEQ: The fetch round that completed. Several loads can be + in flight per id (the scheduler submits loads incrementally), + so completions are matched by round. """ TYPE = "transfer_done" KV_REQUEST_ID = "kv_request_id" SUCCESS = "success" + ROUND_SEQ = "round_seq" @staticmethod def validate(msg: dict) -> None: """Raise ValueError if any field has an invalid type or value.""" _require(msg, TransferDoneMsg.KV_REQUEST_ID, str) _require(msg, TransferDoneMsg.SUCCESS, bool) + _require_non_neg_int(msg, TransferDoneMsg.ROUND_SEQ) class AbortFetchMsg: @@ -302,15 +308,18 @@ class AbortFetchMsg: Fields: KV_REQUEST_ID: The request to cancel. + ROUND_SEQ: The fetch round to cancel. """ TYPE = "abort_fetch" KV_REQUEST_ID = "kv_request_id" + ROUND_SEQ = "round_seq" @staticmethod def validate(msg: dict) -> None: """Raise ValueError if any field has an invalid type or value.""" _require(msg, AbortFetchMsg.KV_REQUEST_ID, str) + _require_non_neg_int(msg, AbortFetchMsg.ROUND_SEQ) class AbortAckMsg: @@ -318,12 +327,15 @@ class AbortAckMsg: Fields: KV_REQUEST_ID: The request that was cancelled. + ROUND_SEQ: The round that was cancelled; echoes AbortFetchMsg. """ TYPE = "abort_ack" KV_REQUEST_ID = "kv_request_id" + ROUND_SEQ = "round_seq" @staticmethod def validate(msg: dict) -> None: """Raise ValueError if any field has an invalid type or value.""" _require(msg, AbortAckMsg.KV_REQUEST_ID, str) + _require_non_neg_int(msg, AbortAckMsg.ROUND_SEQ) diff --git a/vllm/v1/kv_offload/tiering/p2p/session/server.py b/vllm/v1/kv_offload/tiering/p2p/session/server.py index 3709ecd9c53..8558acdf919 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/server.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/server.py @@ -53,15 +53,6 @@ class StoreResult(NamedTuple): success: bool -class _InflightXfer(NamedTuple): - """Metadata for a single inflight RDMA transfer, keyed by transfer_id.""" - - kv_request_id: str - block_count: int - # The set of store job IDs that contributed blocks to this transfer. - job_ids: set[int] - - class _MatchResult(NamedTuple): """Result of block matching: pairs ready for transfer.""" @@ -73,12 +64,17 @@ class _MatchResult(NamedTuple): @dataclass class _OutboundRequestState: - """Server-role state for a single peer fetch request. + """Server-role state for a single fetch round of a peer request. - The owning ``kv_request_id`` is the ``ServerRole._requests`` dict key - and is not duplicated on the value. + A kv_request_id may run several lookup→fetch rounds; rounds live in + ``_ServerRequestState.outbound`` keyed by the wire ``round_seq``, so + terminals touch only their own round. """ + # Supply came from inbound lookup pins (symmetric): no late + # submit_store can arrive, so unmatched fetch demand fails fast. PD + # rounds park demand for stores instead. + lookup_supplied: bool = False demand_received: bool = False available: dict[OffloadKey, tuple[int, int]] = field( default_factory=dict @@ -88,7 +84,8 @@ class _OutboundRequestState: ) # key → remote_block_idx: blocks peer wants, awaiting supply remaining: int = 0 # blocks that need to be transferred to client finishing: bool = False # Signal finish request ASAP - # Job IDs that submit_store'd blocks for this request and have not + inflight: int = 0 # transfers submitted for this round, not yet polled + # Job IDs that submit_store'd blocks for this round and have not # yet emitted a StoreResult. The terminal-finalize helper drains # this set; poll-done and poll-failed discard entries as their # StoreResults fire. @@ -145,6 +142,21 @@ class _OutboundRequestState: ) +@dataclass +class _InflightXfer: + """Metadata for a single inflight RDMA transfer, keyed by transfer_id.""" + + kv_request_id: str + block_count: int + # The set of store job IDs that contributed blocks to this transfer. + job_ids: set[int] + # Round this transfer serves and its key in ``st.outbound``; + # remaining/finalize apply only while the round is still registered. + # Dummy default for test-seeded entries. + round: _OutboundRequestState = field(default_factory=_OutboundRequestState) + round_key: int = 0 + + @dataclass class _ActiveLookup: """In-flight state for one inbound LookupMsg. @@ -159,6 +171,8 @@ class _ActiveLookup: lookup_id: int kv_request_id: str ctx: ReqContext + # Wire round these probes belong to; pins park under it. + round_seq: int = 0 # Keys from the inbound LookupMsg, preserved in wire order so # the aggregated response goes back in the same order. keys: list[OffloadKey] = field(default_factory=list) @@ -185,6 +199,7 @@ class _PendingLookup(NamedTuple): keys: list[OffloadKey] enqueued_at: float + round_seq: int = 0 @dataclass @@ -198,17 +213,15 @@ class _ServerRequestState: is idle — see ``ServerRole._maybe_prune``. """ - # Outbound serve state (PD + symmetric producer side). None until the - # first add_stored_blocks / on_fetch; reset to None on finalize/abort. - outbound: _OutboundRequestState | None = None + # Fetch rounds keyed by wire round_seq. A round is created by its + # first supply or its fetch and removed at its terminal (finalize / + # failure / abort). + outbound: dict[int, _OutboundRequestState] = field(default_factory=dict) # Raw inbound LookupMsgs not yet processed against the ParentManager. pending_lookups: list[_PendingLookup] = field(default_factory=list) # Per-LookupMsg state parked with HIT_PENDING / RETRY keys, keyed by # the (globally unique) lookup_id and re-polled each serve. lookups: dict[int, _ActiveLookup] = field(default_factory=dict) - # Start time of a pending abort drain (``time.monotonic``); None when - # no abort is in progress. - abort_started_at: float | None = None # Transfer ids in ``ServerRole._inflight`` for this id. Kept in sync # via _inflight_add / _inflight_pop so a non-empty set is an exact # "has any inflight transfer" predicate and the abort drain can @@ -256,9 +269,9 @@ class ServerRole: # ``parent.on_request_finished`` in ``serve_external_requests``. self._finished_lookup_ctxs: list[ReqContext] = [] self._lookup_id_counter: int = 0 - # kv_request_ids with a parked abort awaiting drain — work-list so - # drain_pending_aborts doesn't scan every request each poll tick. - self._parked_aborts: set[str] = set() + # Parked aborts awaiting drain, keyed by (kv_request_id, round) + # with the abort start time. + self._pending_aborts: dict[tuple[str, int], float] = {} # ------------------------------------------------------------------ # State helpers @@ -277,11 +290,11 @@ class ServerRole: st = self._requests.get(kv_request_id) if ( st is not None - and st.outbound is None + and not st.outbound and not st.inflight_tids and not st.lookups and not st.pending_lookups - and st.abort_started_at is None + and not any(kv == kv_request_id for kv, _ in self._pending_aborts) ): del self._requests[kv_request_id] @@ -295,87 +308,100 @@ class ServerRole: keys: Sequence[OffloadKey], block_ids: Sequence[int], job_id: JobId, + round_seq: int = 0, + *, + from_lookup: bool = False, ) -> None: - """New blocks stored locally — match against pending fetch demand.""" + """New blocks stored locally — match within their fetch round. + + Lookup pins carry the round they were probed under; PD + submit_store batches share PD's single round 0. + """ self._store_jobs[job_id] = time.monotonic() st = self._get_or_create_request(kv_request_id) - if st.outbound is None: - st.outbound = _OutboundRequestState() - result = st.outbound.add_stored_blocks(keys, block_ids, job_id) - if result.local_idxs and st.outbound.demand_received: - self._submit_transfer(kv_request_id, result) + rnd = st.outbound.get(round_seq) + if rnd is None: + rnd = st.outbound[round_seq] = _OutboundRequestState() + if from_lookup: + rnd.lookup_supplied = True + result = rnd.add_stored_blocks(keys, block_ids, job_id) + if result.local_idxs and rnd.demand_received: + self._submit_transfer(kv_request_id, result, rnd, round_seq) def on_fetch( self, kv_request_id: str, keys: Sequence[OffloadKey], block_indexes: Sequence[int], + round_seq: int = 0, ) -> None: """Handle a FetchMsg from the peer. - In p2p mode FetchMsg is the server-side "request finished" - signal for ``kv_request_id``. Three consequences flow from that: - - - No further ``parent.create_store_job`` will fire for this id: - the request's parked ``lookups`` are popped here (via - ``_finish_inbound_lookups``) before the next - ``serve_external_requests`` runs ``_resolve_pending_lookups``, - so any HIT_PENDING / RETRY key that would otherwise later - promote to HIT and pin a slot is dropped instead. Any raw - not-yet-processed LookupMsg for this id is dropped too. - - All server-side lookup state for the id is cleaned up (the - ``lookups`` entries themselves). - - The synthetic ctx is queued for ``parent.on_request_finished`` - (fired by the next ``serve_external_requests``) so the - TieringManager can release per-lookup bookkeeping. - - The client contract guarantees exactly one FetchMsg per - lookup-touched request — including an empty one when no - blocks end up being fetched. - - Raises ``ValueError`` on a duplicate fetch for the same - ``kv_request_id``; the coordinator's dispatch loop turns that - into a protocol-error disconnect. + A non-empty fetch binds and closes its round, leaving lookup + state alone (the next round's LookupMsg may already be in + flight). The terminal empty fetch closes the id: parked lookups + are popped and every remaining round drained. A second fetch for + a round already holding demand raises ValueError + (protocol-error disconnect). """ logger.debug( - "P2PSession %s: fetch RECEIVED kv_request_id=%s blocks=%d", + "P2PSession %s: fetch RECEIVED kv_request_id=%s round=%s blocks=%d", self._peer_id, kv_request_id, + round_seq, len(keys), ) st = self._requests.get(kv_request_id) - existing = st.outbound if st is not None else None + existing = st.outbound.get(round_seq) if st is not None else None if existing is not None and existing.demand_received: - # A second fetch for the same kv_request_id would overwrite - # `remaining` and leak inflight bookkeeping. Treat as a - # protocol violation. - raise ValueError(f"duplicate fetch for kv_request_id={kv_request_id}") + raise ValueError( + f"duplicate fetch for kv_request_id={kv_request_id} round={round_seq}" + ) st = self._get_or_create_request(kv_request_id) - if st.outbound is None: - st.outbound = _OutboundRequestState() - req = st.outbound + req = st.outbound.get(round_seq) + if req is None: + req = st.outbound[round_seq] = _OutboundRequestState() result = req.add_fetch_demand(keys, block_indexes) + if not keys: + # Terminal empty fetch: close the lookup phase and drain + # every round with no TransferDoneMsg (nothing waits on it). + self._finish_inbound_lookups(kv_request_id) + for key in list(st.outbound): + self._finalize_outbound(kv_request_id, key, send_done=False) + return + if req.lookup_supplied and req.demanded: + # A symmetric round's supply always precedes its fetch, so + # unmatched demand is unservable — fail now, not at the load + # timeout. PD rounds keep parking demand for stores that + # arrive later. + logger.warning( + "P2PSession %s: fetch kv_request_id=%s round=%s demanded %d " + "blocks but %d have no pinned supply; failing fetch " + "immediately", + self._peer_id, + kv_request_id, + round_seq, + len(keys), + len(req.demanded), + ) + self._finalize_outbound(kv_request_id, round_seq, success=False) + return if result.local_idxs: - self._submit_transfer(kv_request_id, result) - # Close the peer's request as far as the server's lookup phase - # is concerned: pop parked lookups so no further - # ``parent.create_store_job`` fires for this id, and queue their - # ctxs for ``parent.on_request_finished``. Done before the - # finalize path below so bookkeeping releases in-order. - self._finish_inbound_lookups(kv_request_id) + self._submit_transfer(kv_request_id, result, req, round_seq) # Prefiller-first mode: finish_request may have run before # fetch arrived. If so, finalize once we know what was # demanded — fully satisfied → success, else early-fail. - if req.finishing and not self._has_inflight_for(kv_request_id): - self._finalize_outbound(kv_request_id) + if req.finishing and req.inflight == 0: + self._finalize_outbound(kv_request_id, round_seq) - def on_abort_fetch(self, kv_request_id: str) -> None: - """Handle an AbortFetchMsg from the peer.""" + def on_abort_fetch(self, kv_request_id: str, round_seq: int = 0) -> None: + """Handle an AbortFetchMsg from the peer, cancelling one round.""" # Abort for an unknown id may be a benign race/duplicate or a # real protocol violation; we don't track completed ids, so warn. st = self._requests.get(kv_request_id) - has_outbound = st is not None and st.outbound is not None - if not has_outbound and not self._has_inflight_for(kv_request_id): + if (st is None or not st.outbound) and not self._has_inflight_for( + kv_request_id + ): logger.warning( "P2PSession %s: abort_fetch for unknown kv_request_id=%s " "(no outbound or inflight state); benign race or stale", @@ -385,16 +411,15 @@ class ServerRole: # Idempotent: receiving AbortFetchMsg again before we've sent the # ack just triggers another drain attempt without resetting the # deadline. - st = self._get_or_create_request(kv_request_id) - if st.abort_started_at is None: - st.abort_started_at = time.monotonic() - self._parked_aborts.add(kv_request_id) - self._drain_abort(kv_request_id) + self._get_or_create_request(kv_request_id) + self._pending_aborts.setdefault((kv_request_id, round_seq), time.monotonic()) + self._drain_abort(kv_request_id, round_seq) def on_lookup( self, kv_request_id: str, keys: Sequence[OffloadKey], + round_seq: int = 0, ) -> None: """Enqueue a LookupMsg from a symmetric-P2P consumer. @@ -406,13 +431,18 @@ class ServerRole: parent calls are valid. """ logger.debug( - "P2P LOOKUP server %s: RECV LookupMsg kv_request_id=%s keys=%d", + "P2P LOOKUP server %s: RECV LookupMsg kv_request_id=%s round=%s keys=%d", self._peer_id, kv_request_id, + round_seq, len(keys), ) self._get_or_create_request(kv_request_id).pending_lookups.append( - _PendingLookup(keys=list(keys), enqueued_at=time.monotonic()) + _PendingLookup( + keys=list(keys), + enqueued_at=time.monotonic(), + round_seq=round_seq, + ) ) self._serve_pending.add(kv_request_id) @@ -434,7 +464,7 @@ class ServerRole: st.pending_lookups = [] for pl in pending: self._process_inbound_lookup( - kv_request_id, pl.keys, pl.enqueued_at, parent + kv_request_id, pl.keys, pl.enqueued_at, pl.round_seq, parent ) self._resolve_pending_lookups(kv_request_id, parent) st = self._requests.get(kv_request_id) @@ -481,9 +511,7 @@ class ServerRole: else: lookup.pending.add(h) if new_hits: - self._pin_and_register_hits( - lookup.kv_request_id, new_hits, lookup.ctx, parent - ) + self._pin_and_register_hits(lookup, new_hits, parent) return new_hits def _process_inbound_lookup( @@ -491,6 +519,7 @@ class ServerRole: kv_request_id: str, keys: list[OffloadKey], enqueued_at: float, + round_seq: int, parent: ParentManager, ) -> None: """Resolve one enqueued LookupMsg against ``parent``. @@ -515,6 +544,7 @@ class ServerRole: lookup_id=lookup_id, kv_request_id=kv_request_id, ctx=ctx, + round_seq=round_seq, keys=list(keys), deadline=enqueued_at + _LOOKUP_PENDING_TIMEOUT_S, ) @@ -547,25 +577,26 @@ class ServerRole: def _pin_and_register_hits( self, - kv_request_id: str, + lookup: _ActiveLookup, keys: list[OffloadKey], - ctx: ReqContext, parent: ParentManager, ) -> None: - """Pin primary slots for HIT keys and feed them into the - existing ``add_stored_blocks`` matching path. + """Pin primary slots for HIT keys and park them as the lookup's + round supply via ``add_stored_blocks``. Caller has already confirmed every key is HIT (single-threaded scheduler ⇒ no eviction race), so the JobMetadata returned by ``parent.create_store_job`` carries parallel ``keys``/``block_ids`` of length ``len(keys)``. """ - meta = parent.create_store_job(keys, ctx) + meta = parent.create_store_job(keys, lookup.ctx) self.add_stored_blocks( - kv_request_id, + lookup.kv_request_id, list(meta.keys), list(meta.block_ids), meta.job_id, + round_seq=lookup.round_seq, + from_lookup=True, ) def _resolve_pending_lookups( @@ -650,9 +681,8 @@ class ServerRole: Called on the two events that mean "no more lookup traffic for ``kv_request_id`` is expected on this session": the terminal - FetchMsg from the peer (client contract: exactly one FetchMsg - per lookup-touched request, even if empty), and a local - ``finish``. Whichever fires second is a no-op. + empty FetchMsg from the peer and a local ``finish``. Whichever + fires second is a no-op. """ st = self._requests.get(kv_request_id) if st is None: @@ -688,20 +718,15 @@ class ServerRole: self._finish_inbound_lookups(kv_request_id) st = self._requests.get(kv_request_id) - req = st.outbound if st is not None else None - if req is None: + if st is None: return - req.finishing = True - if not req.demand_received: - return - if self._has_inflight_for(kv_request_id): - return - # Remaining > 0 here: if it had hit 0, the poll-done success - # branch would have already cleared outbound and we'd have - # returned at `req is None` above. Helper derives success from - # remaining and emits StoreResult(success=False) for any - # leftover pending jobs. - self._finalize_outbound(kv_request_id) + for key, req in list(st.outbound.items()): + req.finishing = True + if not req.demand_received or req.inflight: + # No demand yet (prefiller-first): on_fetch finalizes via + # `finishing`. Inflight: the last completion finalizes. + continue + self._finalize_outbound(kv_request_id, key) def collect_results(self) -> list[StoreResult]: """Drain timeouts, deferred results, and transport completions. @@ -741,20 +766,24 @@ class ServerRole: ) continue results.extend(self._settle_xfer_jobs(xfer, success=True)) + rnd = xfer.round st = self._requests.get(xfer.kv_request_id) - req = st.outbound if st is not None else None - if req is not None and req.demand_received: - req.remaining -= xfer.block_count - assert req.remaining >= 0, ( + if st is not None and st.outbound.get(xfer.round_key) is rnd: + rnd.remaining -= xfer.block_count + assert rnd.remaining >= 0, ( f"remaining went negative for kv_request_id={xfer.kv_request_id}" ) - if req.remaining == 0: - self._finalize_outbound(xfer.kv_request_id, success=True) - elif req.finishing and not self._has_inflight_for(xfer.kv_request_id): - self._finalize_outbound(xfer.kv_request_id, success=False) + if rnd.remaining == 0: + self._finalize_outbound( + xfer.kv_request_id, xfer.round_key, success=True + ) + elif rnd.finishing and rnd.inflight == 0: + self._finalize_outbound( + xfer.kv_request_id, xfer.round_key, success=False + ) self._maybe_prune(xfer.kv_request_id) - failed_kv_request_ids: set[str] | None = None + failed_rounds: list[tuple[str, _OutboundRequestState]] | None = None for tid in poll_result.failed: xfer = self._inflight_pop(tid) if xfer is None: @@ -767,35 +796,36 @@ class ServerRole: tid, ) continue - if failed_kv_request_ids is None: - failed_kv_request_ids = set() - failed_kv_request_ids.add(xfer.kv_request_id) results.extend(self._settle_xfer_jobs(xfer, success=False)) + rnd = xfer.round st = self._requests.get(xfer.kv_request_id) - req = st.outbound if st is not None else None - if st is not None: - st.outbound = None - if req is not None and req.demand_received: + if st is not None and st.outbound.get(xfer.round_key) is rnd: + del st.outbound[xfer.round_key] + if failed_rounds is None: + failed_rounds = [] + failed_rounds.append((xfer.kv_request_id, rnd)) self._send( { TYPE_KEY: TransferDoneMsg.TYPE, TransferDoneMsg.KV_REQUEST_ID: xfer.kv_request_id, TransferDoneMsg.SUCCESS: False, + TransferDoneMsg.ROUND_SEQ: xfer.round_key, } ) self._maybe_prune(xfer.kv_request_id) - # Cancel other inflight for the same failed kv_request_ids - if failed_kv_request_ids: - ids_to_cancel = [ - tid - for tid, xfer in self._inflight.items() - if xfer.kv_request_id in failed_kv_request_ids - ] - for tid in ids_to_cancel: - self._inflight_pop(tid) - self._transport.cancel(ids_to_cancel) - for kv_request_id in failed_kv_request_ids: + # Cancel each failed round's other inflight and fail its + # remaining store jobs — nothing else will settle them. + if failed_rounds: + for kv_request_id, rnd in failed_rounds: + ids_to_cancel = [ + tid for tid, x in self._inflight.items() if x.round is rnd + ] + for tid in ids_to_cancel: + self._inflight_pop(tid) + if ids_to_cancel: + self._transport.cancel(ids_to_cancel) + results.extend(self._fail_round_jobs(rnd)) self._maybe_prune(kv_request_id) return results @@ -811,8 +841,8 @@ class ServerRole: def drain_pending_aborts(self) -> None: """Re-attempt every parked abort once per poll tick.""" - for kv_request_id in list(self._parked_aborts): - self._drain_abort(kv_request_id) + for kv_request_id, round_seq in list(self._pending_aborts): + self._drain_abort(kv_request_id, round_seq) def close(self) -> tuple[list[int], list[ReqContext]]: """Tear down. Cancels inflight. @@ -839,7 +869,7 @@ class ServerRole: failed_serves.extend(self._finished_lookup_ctxs) self._requests.clear() self._serve_pending.clear() - self._parked_aborts.clear() + self._pending_aborts.clear() self._finished_lookup_ctxs.clear() return failed_stores, failed_serves @@ -870,6 +900,8 @@ class ServerRole: xfer = self._inflight.pop(tid, None) if xfer is None: return None + xfer.round.inflight -= 1 + assert xfer.round.inflight >= 0 st = self._requests.get(xfer.kv_request_id) if st is not None: st.inflight_tids.discard(tid) @@ -880,80 +912,108 @@ class ServerRole: ) -> list[StoreResult]: """Emit StoreResults for a completed transfer's store jobs. - Pops each attached job from ``_store_jobs`` and clears it from the - request's pending set. A job already popped (via timeout, cancel, - etc.) is skipped so we never double-emit a contradictory result. + Pops each attached job from ``_store_jobs`` and clears it from + its round's pending set. A job already popped (via timeout, + cancel, etc.) is skipped so we never double-emit a contradictory + result. """ results: list[StoreResult] = [] - st = self._requests.get(xfer.kv_request_id) - req = st.outbound if st is not None else None for job_id in xfer.job_ids: if self._store_jobs.pop(job_id, None) is None: continue results.append(StoreResult(job_id=job_id, success=success)) - if req is not None: - req.pending_job_ids.discard(job_id) + xfer.round.pending_job_ids.discard(job_id) return results # ------------------------------------------------------------------ # Internal — finalize / abort drain # ------------------------------------------------------------------ + def _fail_round_jobs(self, rnd: _OutboundRequestState) -> list[StoreResult]: + """Fail a terminated round's still-pending store jobs (idempotent).""" + results: list[StoreResult] = [] + for job_id in rnd.pending_job_ids: + if self._store_jobs.pop(job_id, None) is None: + continue + results.append(StoreResult(job_id=job_id, success=False)) + rnd.pending_job_ids.clear() + return results + def _finalize_outbound( self, kv_request_id: str, + round_key: int, success: bool | None = None, + send_done: bool = True, ) -> None: - """Pop the outbound state and emit terminal results. - - Called when no further work will happen for this kv_request_id - on the server side: either request_finish has fired and there - are no inflight transfers, or the last inflight just completed - while finishing. + """Pop one round and emit its terminal results. If ``success`` is None, derive it from ``req.remaining == 0``. - The same flag is used for both the peer's TransferDoneMsg and - the StoreResult(s) emitted for any leftover pending job_ids. + ``send_done=False`` skips the TransferDoneMsg (terminal empty + fetch). Other rounds of the id are untouched. """ st = self._requests[kv_request_id] - assert st.outbound is not None - req = st.outbound - st.outbound = None + req = st.outbound.pop(round_key) if success is None: - success = req.remaining == 0 - for job_id in req.pending_job_ids: - self._store_jobs.pop(job_id, None) - self._pending_store_results.append( - StoreResult(job_id=job_id, success=success) - ) - self._send( - { - TYPE_KEY: TransferDoneMsg.TYPE, - TransferDoneMsg.KV_REQUEST_ID: kv_request_id, - TransferDoneMsg.SUCCESS: success, - } + success = req.demand_received and req.remaining == 0 + settled = self._fail_round_jobs(req) if not success else None + if settled is not None: + self._pending_store_results.extend(settled) + else: + for job_id in req.pending_job_ids: + if self._store_jobs.pop(job_id, None) is None: + continue + self._pending_store_results.append( + StoreResult(job_id=job_id, success=True) + ) + req.pending_job_ids.clear() + logger.debug( + "P2PSession %s: finalize kv_request_id=%s round=%s success=%s " + "remaining=%d leftover_available=%d send_done=%s", + self._peer_id, + kv_request_id, + round_key, + success, + req.remaining, + len(req.available), + send_done, ) + if send_done and req.demand_received: + self._send( + { + TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.KV_REQUEST_ID: kv_request_id, + TransferDoneMsg.SUCCESS: success, + TransferDoneMsg.ROUND_SEQ: round_key, + } + ) self._maybe_prune(kv_request_id) - def _drain_abort(self, kv_request_id: str) -> None: + def _drain_abort(self, kv_request_id: str, round_seq: int) -> None: """One drain attempt for a pending abort. - Stops accepting more blocks for ``kv_request_id``, then asks the - transport to cancel any matching inflight transfers in - ``mode="wait"``. Sends ``AbortAckMsg`` once nothing remains - inflight, or after ``_CANCEL_DRAIN_TIMEOUT_S`` falls back to - ``mode="immediate"`` and acks anyway. + Detaches the aborted round, then asks the transport to cancel its + inflight transfers in ``mode="wait"``. Sends ``AbortAckMsg`` once + nothing remains inflight, or after ``_CANCEL_DRAIN_TIMEOUT_S`` + falls back to ``mode="immediate"`` and acks anyway. """ st = self._requests[kv_request_id] - st.outbound = None - ids = list(st.inflight_tids) + rnd = st.outbound.pop(round_seq, None) + if rnd is not None: + # Its transfers are being cancelled; fail its jobs now + # instead of leaking them to the store timeout. + self._pending_store_results.extend(self._fail_round_jobs(rnd)) + ids = [ + tid + for tid, x in self._inflight.items() + if x.kv_request_id == kv_request_id and x.round_key == round_seq + ] if not ids: - self._finalize_abort(kv_request_id) + self._finalize_abort(kv_request_id, round_seq) return - assert st.abort_started_at is not None - expired = time.monotonic() - st.abort_started_at >= _CANCEL_DRAIN_TIMEOUT_S - if expired: + started_at = self._pending_aborts[(kv_request_id, round_seq)] + if time.monotonic() - started_at >= _CANCEL_DRAIN_TIMEOUT_S: for tid in ids: self._inflight_pop(tid) self._transport.cancel(ids, mode="immediate") @@ -964,7 +1024,7 @@ class ServerRole: kv_request_id, len(ids), ) - self._finalize_abort(kv_request_id) + self._finalize_abort(kv_request_id, round_seq) return still = self._transport.cancel(ids, mode="wait") @@ -977,16 +1037,15 @@ class ServerRole: if tid not in still_set: self._inflight_pop(tid) if not still: - self._finalize_abort(kv_request_id) + self._finalize_abort(kv_request_id, round_seq) - def _finalize_abort(self, kv_request_id: str) -> None: - st = self._requests[kv_request_id] - st.abort_started_at = None - self._parked_aborts.discard(kv_request_id) + def _finalize_abort(self, kv_request_id: str, round_seq: int) -> None: + self._pending_aborts.pop((kv_request_id, round_seq), None) self._send( { TYPE_KEY: AbortAckMsg.TYPE, AbortAckMsg.KV_REQUEST_ID: kv_request_id, + AbortAckMsg.ROUND_SEQ: round_seq, } ) self._maybe_prune(kv_request_id) @@ -995,7 +1054,13 @@ class ServerRole: # Internal — transfers and store-job timeouts # ------------------------------------------------------------------ - def _submit_transfer(self, kv_request_id: str, result: _MatchResult) -> None: + def _submit_transfer( + self, + kv_request_id: str, + result: _MatchResult, + rnd: _OutboundRequestState, + round_key: int, + ) -> None: logger.debug( "P2PSession %s: NIXL write_blocks CALL kv_request_id=%s " "local_idxs=%d remote_idxs=%d", @@ -1016,12 +1081,15 @@ class ServerRole: transfer_id, len(result.local_idxs), ) + rnd.inflight += 1 self._inflight_add( transfer_id, _InflightXfer( kv_request_id=kv_request_id, block_count=len(result.local_idxs), job_ids=result.job_ids, + round=rnd, + round_key=round_key, ), ) else: @@ -1031,22 +1099,24 @@ class ServerRole: kv_request_id, len(result.local_idxs), ) - # The matched blocks were popped from req.demanded / - # req.available, but no inflight will satisfy them, so - # remaining will never reach 0 on its own. Mark the - # request as finishing so the existing terminal paths - # clean up: if other inflight is in flight, the last one - # to drain will fire _finalize_outbound(success=False) - # via the elif branch in collect_results. If - # nothing else is in flight, finalize now so the peer - # and the local store jobs don't wait for finish_request - # or for _STORE_TIMEOUT_S / _LOAD_TIMEOUT_S. + # The matched blocks were popped from rnd.demanded / + # rnd.available, but no inflight will satisfy them, so + # remaining will never reach 0 on its own. Mark the round + # as finishing so the existing terminal paths clean up: if + # other transfers of this round are in flight, the last one + # to drain will fire _finalize_outbound(success=False) via + # the elif branch in collect_results. If nothing else is in + # flight, finalize now so the peer and the local store jobs + # don't wait for finish_request or for _STORE_TIMEOUT_S / + # _LOAD_TIMEOUT_S. + rnd.finishing = True st = self._requests.get(kv_request_id) - req = st.outbound if st is not None else None - if req is not None: - req.finishing = True - if not self._has_inflight_for(kv_request_id): - self._finalize_outbound(kv_request_id, success=False) + if ( + st is not None + and st.outbound.get(round_key) is rnd + and rnd.inflight == 0 + ): + self._finalize_outbound(kv_request_id, round_key, success=False) def _timeout_pending_store_jobs(self) -> list[StoreResult]: if not self._store_jobs: diff --git a/vllm/v1/kv_offload/tiering/p2p/session/session.py b/vllm/v1/kv_offload/tiering/p2p/session/session.py index 7d19913b0a2..4bd780159a4 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/session.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/session.py @@ -374,26 +374,34 @@ class P2PSession: for bh in msg[FetchMsg.KEYS] ] block_indexes = msg[FetchMsg.BLOCK_INDEXES] + round_seq = msg[FetchMsg.ROUND_SEQ] # Run the server-role state machine inline as today — # add_fetch_demand records demand against any blocks we've # already seen in `available`. Report the kv_request_id so # the manager (after poll() returns) can replay any parked # submit_store batches; their add_stored_blocks calls hit # the demand recorded here and submit transfers immediately. - self._server.on_fetch(kv_request_id, keys, block_indexes) + self._server.on_fetch(kv_request_id, keys, block_indexes, round_seq) self._new_fetch_ids.append(kv_request_id) elif msg_type == AbortFetchMsg.TYPE: AbortFetchMsg.validate(msg) - self._server.on_abort_fetch(msg[AbortFetchMsg.KV_REQUEST_ID]) + self._server.on_abort_fetch( + msg[AbortFetchMsg.KV_REQUEST_ID], + msg[AbortFetchMsg.ROUND_SEQ], + ) elif msg_type == TransferDoneMsg.TYPE: TransferDoneMsg.validate(msg) self._client.on_transfer_done( msg[TransferDoneMsg.KV_REQUEST_ID], msg[TransferDoneMsg.SUCCESS], + msg[TransferDoneMsg.ROUND_SEQ], ) elif msg_type == AbortAckMsg.TYPE: AbortAckMsg.validate(msg) - self._client.on_abort_ack(msg[AbortAckMsg.KV_REQUEST_ID]) + self._client.on_abort_ack( + msg[AbortAckMsg.KV_REQUEST_ID], + msg[AbortAckMsg.ROUND_SEQ], + ) elif msg_type == LookupMsg.TYPE: LookupMsg.validate(msg) kv_request_id = msg[LookupMsg.KV_REQUEST_ID] @@ -401,7 +409,7 @@ class P2PSession: OffloadKey(bh if isinstance(bh, bytes) else bytes(bh)) for bh in msg[LookupMsg.KEYS] ] - self._server.on_lookup(kv_request_id, keys) + self._server.on_lookup(kv_request_id, keys, msg[LookupMsg.ROUND_SEQ]) elif msg_type == LookupRespMsg.TYPE: LookupRespMsg.validate(msg) kv_request_id = msg[LookupRespMsg.KV_REQUEST_ID] From 4fb483ca86566d163a886416cdd822cf716b1df5 Mon Sep 17 00:00:00 2001 From: IBRAHIM IBRAHIM <66755652+Ibrahim2595@users.noreply.github.com> Date: Tue, 28 Jul 2026 11:47:56 -0400 Subject: [PATCH 177/185] [Docs] Expand llm-d integration page (#45432) Signed-off-by: ibrahimibrahim <ibib2595@gmail.com> Co-authored-by: ibrahimibrahim <ibib2595@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --- docs/deployment/integrations/llm-d.md | 36 +++++++++++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/docs/deployment/integrations/llm-d.md b/docs/deployment/integrations/llm-d.md index 6060b98f642..7d261eb910b 100644 --- a/docs/deployment/integrations/llm-d.md +++ b/docs/deployment/integrations/llm-d.md @@ -1,5 +1,37 @@ # llm-d -vLLM can be deployed with [llm-d](https://github.com/llm-d/llm-d), a Kubernetes-native distributed inference serving stack providing well-lit paths for anyone to serve large generative AI models at scale. It helps achieve the fastest "time to state-of-the-art (SOTA) performance" for key OSS models across most hardware accelerators and infrastructure providers. +[llm-d](https://llm-d.ai/) is a Kubernetes-native distributed inference framework for serving large language models at scale, with vLLM as its primary inference engine. llm-d coordinates a fleet of vLLM instances across a cluster so that performance holds up under real production traffic, achieving the fastest "time to state-of-the-art (SOTA) performance" for key OSS models across most hardware accelerators. -You can use vLLM with llm-d directly by following [the official guides](https://llm-d.ai/docs/guides) or via [KServe's LLMInferenceService](https://kserve.github.io/website/docs/model-serving/generative-inference/llmisvc/llmisvc-overview). +It is a [CNCF Sandbox project](https://www.cncf.io/blog/2026/03/24/welcome-llm-d-to-the-cncf-evolving-kubernetes-into-sota-ai-infrastructure/) founded by Red Hat, Google Cloud, IBM Research, CoreWeave, and NVIDIA. + +## What llm-d adds to vLLM + +A single vLLM server is fast, but at scale the picture changes: across many replicas, cache locality breaks under round-robin load balancing, long prompts inflate time-to-first-token, and accelerators sit underused. llm-d adds the cluster-level layer that vLLM does not aim to provide on its own: + +- **[Prefix-aware routing](https://llm-d.ai/docs/guides/precise-prefix-cache-aware).** Instead of round-robin, llm-d reads vLLM's KV-cache events and routes each request to the replica that already holds its prefix, reusing cache instead of recomputing it. +- **[Distributed KV-cache management](https://llm-d.ai/docs/guides#advanced-kv-cache-management).** A global index tracks which token blocks live on which replica, and [tiered offloading](https://llm-d.ai/docs/guides/tiered-prefix-cache) spills cache to CPU memory or local SSD, extending the working set beyond accelerator HBM. +- **[Prefill/decode disaggregation](https://llm-d.ai/docs/guides/pd-disaggregation).** Prompt processing and token generation run on separate vLLM workers, with KV-cache moved over the vLLM [NIXL connector](https://docs.vllm.ai/en/latest/features/nixl_connector_usage/), lowering TTFT and steadying per-token latency on long prompts. +- **[Wide expert-parallelism](https://llm-d.ai/docs/guides/wide-expert-parallelism).** Serve large Mixture-of-Experts models such as DeepSeek-R1 and GPT-OSS across nodes with combined data and expert parallelism, for more KV-cache capacity and throughput. +- **SLO-aware [autoscaling](https://llm-d.ai/docs/guides/workload-autoscaling) and [flow control](https://llm-d.ai/docs/guides/flow-control).** Scale vLLM pools on real inference signals (queue depth, true demand) rather than raw GPU utilization, with multi-tenant fairness and priority dispatch. + +These are composable. Most teams start by adding prefix-aware routing over an existing vLLM pool, then layer in the rest as specific bottlenecks appear. + +## Performance + +Representative benchmarked results across accelerators: + +- **3x higher output throughput** and **2x faster TTFT** from prefix-aware routing vs round-robin (Llama 3.1 70B, AMD MI300X) +- **Up to 70% higher tokens/sec** from prefill/decode disaggregation (GPT-OSS, NVIDIA B200) +- **13.9x throughput** from hierarchical KV offloading at high concurrency vs GPU-only (NVIDIA H100) + +See the [full list](https://github.com/llm-d/llm-d#performance-highlights) and reproducible benchmarks on [Prism](https://prism.llm-d.ai/). + +## Get started + +1. Deploy the [Optimized Baseline](https://llm-d.ai/docs/guides/optimized-baseline) with the [Quickstart](https://llm-d.ai/docs/getting-started/quickstart). It stands up an intelligent router over a vLLM pool on Kubernetes in a tested configuration. +2. Browse the [well-lit path guides](https://llm-d.ai/docs/guides), each a tested recipe for one of the capabilities above, and add the optimization that fits your workload. +3. Read the [Introduction](https://llm-d.ai/docs/getting-started) and [Architecture overview](https://llm-d.ai/docs/architecture) to see how the pieces wrap your vLLM deployment. + +You can also deploy vLLM with llm-d via [KServe's LLMInferenceService](https://kserve.github.io/website/docs/model-serving/generative-inference/llmisvc/llmisvc-overview). + +Questions and contributions are welcome on [GitHub](https://github.com/llm-d/llm-d) and [Slack](https://llm-d.ai/slack). From 6453fc0b8cb50668dceaf5839033b0a2a8fa7f7f Mon Sep 17 00:00:00 2001 From: Nick Hill <nickhill123@gmail.com> Date: Tue, 28 Jul 2026 09:09:28 -0700 Subject: [PATCH 178/185] [Bugfix] Don't reuse engine core payload buffer while zmq is sending it (#50053) Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> --- tests/v1/test_serial_utils.py | 136 ++++++++++++++++++++++++++++++++++ vllm/envs.py | 2 +- vllm/v1/engine/core.py | 41 +++++++--- vllm/v1/engine/core_client.py | 58 +++------------ 4 files changed, 179 insertions(+), 58 deletions(-) diff --git a/tests/v1/test_serial_utils.py b/tests/v1/test_serial_utils.py index 4ed8724e60f..9f7761c22b5 100644 --- a/tests/v1/test_serial_utils.py +++ b/tests/v1/test_serial_utils.py @@ -423,3 +423,139 @@ def test_multiple_senders_single_receiver_ipc(): assert torch.allclose(decoded.prompt_embeds, original_tensor), ( f"Value mismatch for sender {sender_idx} msg {msg_idx}" ) + + +def _logprobs_outputs(num_reqs: int, num_prompt_tokens: int): + """An EngineCoreOutputs carrying prompt logprobs, as the engine core sends + it: many requests, each with per-token tensors small enough that pyzmq + copies their frames, while the accumulated payload frame is large enough + that pyzmq sends it zero-copy.""" + from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs + from vllm.v1.outputs import LogprobsTensors + + outputs = [] + for req in range(num_reqs): + num_tokens = num_prompt_tokens + req % 4 + outputs.append( + EngineCoreOutput( + request_id=f"req-{req:08d}", + new_token_ids=[req], + new_prompt_logprobs_tensors=LogprobsTensors( + logprob_token_ids=torch.arange( + num_tokens * 2, dtype=torch.int64 + ).view(num_tokens, 2), + logprobs=torch.zeros(num_tokens, 2, dtype=torch.float32), + selected_token_ranks=torch.zeros(num_tokens, dtype=torch.int32), + ), + ) + ) + return EngineCoreOutputs(outputs=outputs) + + +def test_payload_buffer_reuse_does_not_corrupt_in_flight_messages(): + """The engine core recycles the msgpack payload buffer across messages + (`MsgpackEncoder.encode_into`). It may only do so once zmq has finished + sending that buffer, otherwise a newer payload is delivered alongside the + older message's zero-copy tensor frames. + + `Socket.send_multipart(track=True)` cannot be used to detect this: it + returns a tracker for the last frame only, and pyzmq copies frames below + `zmq.COPY_THRESHOLD` and reports them as already-sent. + """ + import zmq + + from vllm.v1.engine import EngineCoreOutputs + from vllm.v1.engine.core import EngineCoreProc + + num_msgs = 100 + encoder = MsgpackEncoder() + decoder = MsgpackDecoder(EngineCoreOutputs) + # Enough requests that the payload frame is zero-copied rather than copied + # by pyzmq, which is what makes early reuse observable. + messages = [_logprobs_outputs(300, 24 + i % 8) for i in range(num_msgs)] + assert len(encoder.encode(messages[0])[0]) >= zmq.COPY_THRESHOLD + + reuse_buffers: list[bytearray] = [] + pending: list[tuple[zmq.MessageTracker, bytearray]] = [] + with zmq.Context() as ctx: + push = ctx.socket(zmq.PUSH) + push.bind("inproc://test-payload-reuse") + pull = ctx.socket(zmq.PULL) + pull.connect("inproc://test-payload-reuse") + + for outputs in messages: + while pending and pending[0][0].done: + reuse_buffers.append(pending.pop(0)[1]) + buffer = reuse_buffers.pop() if reuse_buffers else bytearray() + buffers = encoder.encode_into(outputs, buffer) + tracker = EngineCoreProc._send_msg_tracking_payload(push, buffers) + if tracker.done: + reuse_buffers.append(buffer) + else: + pending.append((tracker, buffer)) + + for i, sent in enumerate(messages): + received = decoder.decode(pull.recv_multipart(copy=False)) + assert len(received.outputs) == len(sent.outputs), f"message {i}" + for expected, actual in zip(sent.outputs, received.outputs): + sent_ids = expected.new_prompt_logprobs_tensors.logprob_token_ids + got_ids = actual.new_prompt_logprobs_tensors.logprob_token_ids + assert actual.request_id == expected.request_id, f"message {i}" + assert torch.equal(got_ids, sent_ids), ( + f"message {i} request {actual.request_id}: corrupted " + f"prompt logprobs, {got_ids.shape} vs {sent_ids.shape}" + ) + push.close(linger=0) + pull.close(linger=0) + + +def test_zero_copy_frames_survive_without_caller_side_references(): + """Callers don't need to retain the encoded object until zmq has sent it: + for a zero-copy frame, zmq holds its own reference to the backing buffer. + + The engine core clients rely on this when sending requests that carry + tensors (e.g. prompt embeds) without tracking the messages. + + What makes that safe is that `tensor_data()` hands zmq a memoryview which + transitively references the source tensor, so refcounting - not timing - + keeps the memory from being freed and reused underneath zmq. + """ + import gc + + import zmq + + from vllm.v1.utils import tensor_data + + num_elems = 100_000 # comfortably over zmq.COPY_THRESHOLD + expected = torch.arange(num_elems, dtype=torch.int64) + encoder = MsgpackEncoder() + decoder = MsgpackDecoder(RequestWithTensor) + + # The buffer handed to zmq must keep the tensor's storage alive by itself. + holder = tensor_data(expected).obj + while getattr(holder, "base", None) is not None: + holder = holder.base + assert isinstance(holder, torch.Tensor) + assert holder.data_ptr() == expected.data_ptr() + + with zmq.Context() as ctx: + push = ctx.socket(zmq.PUSH) + push.bind("inproc://test-zero-copy-lifetime") + pull = ctx.socket(zmq.PULL) + pull.connect("inproc://test-zero-copy-lifetime") + + request = RequestWithTensor(prompt_embeds=expected.clone(), data="req") + buffers = encoder.encode(request) + assert max(len(buf) for buf in buffers) >= zmq.COPY_THRESHOLD + push.send_multipart(buffers, copy=False) + + # Drop every reference the sender holds, then churn the allocator. + del request, buffers + gc.collect() + torch.arange(num_elems * 4, dtype=torch.int64) + + decoded = decoder.decode(pull.recv_multipart(copy=False)) + assert decoded.prompt_embeds is not None + assert torch.equal(decoded.prompt_embeds, expected) + push.close(linger=0) + pull.close(linger=0) diff --git a/vllm/envs.py b/vllm/envs.py index fb54619c748..f7e01c33275 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -1547,7 +1547,7 @@ environment_variables: dict[str, Callable[[], Any]] = { # tensors above will instead be sent via a separate message. # While the sending side still actually copies the tensor # in all cases, on the receiving side, tensors above this - # limit will actually be zero-copy decoded. + # limit will actually be zero-copy decoded. The unit is bytes. "VLLM_MSGPACK_ZERO_COPY_THRESHOLD": lambda: int( os.getenv("VLLM_MSGPACK_ZERO_COPY_THRESHOLD", "256") ), diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index ecac92f5fe0..66135273002 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -7,7 +7,7 @@ import signal import threading import time from collections import defaultdict, deque -from collections.abc import Callable, Generator +from collections.abc import Callable, Generator, Sequence from concurrent.futures import Future from contextlib import ExitStack, contextmanager from enum import IntEnum @@ -87,7 +87,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig, get_kv_cache_spec_kind from vllm.v1.metrics.stats import SchedulerIterationDetails, SchedulerStats from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus -from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder +from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder, bytestr from vllm.v1.structured_output import StructuredOutputManager from vllm.v1.utils import compute_iteration_details from vllm.version import __version__ as VLLM_VERSION @@ -1748,10 +1748,11 @@ class EngineCoreProc(EngineCore): encoder = MsgpackEncoder() # Send buffers to reuse. reuse_buffers: list[bytearray] = [] - # Keep references to outputs and buffers until zmq is finished - # with them (outputs may contain tensors/np arrays whose - # backing buffers were extracted for zero-copy send). - pending = deque[tuple[zmq.MessageTracker, Any, bytearray]]() + # Payload buffers that can't be reused yet because zmq may still be + # sending them. + # Buffers of the zero-copy tensor/ndarray frames don't need tracking + # here: zmq itself holds a reference to each until it's done with it. + pending = deque[tuple[zmq.MessageTracker, bytearray]]() # We must set linger to ensure the ENGINE_CORE_DEAD # message is sent prior to closing the socket. @@ -1792,20 +1793,38 @@ class EngineCoreProc(EngineCore): # Reclaim buffers that zmq is finished with. while pending and pending[-1][0].done: - reuse_buffers.append(pending.pop()[2]) + reclaimed = pending.pop()[1] + if len(reuse_buffers) < max_reuse_bufs: + reuse_buffers.append(reclaimed) buffer = reuse_buffers.pop() if reuse_buffers else bytearray() buffers = encoder.encode_into(outputs, buffer) - tracker = sockets[client_index].send_multipart( - buffers, copy=False, track=True + tracker = self._send_msg_tracking_payload( + sockets[client_index], buffers ) if not tracker.done: - ref = outputs if len(buffers) > 1 else None - pending.appendleft((tracker, ref, buffer)) + pending.appendleft((tracker, buffer)) elif len(reuse_buffers) < max_reuse_bufs: # Limit the number of buffers to reuse. reuse_buffers.append(buffer) + @staticmethod + def _send_msg_tracking_payload( + socket: zmq.Socket, buffers: Sequence[bytestr] + ) -> zmq.MessageTracker: + """Send `buffers` as a zero-copy multipart message, returning a tracker + for the *first* frame. + + Used instead of `Socket.send_multipart()` because we reuse the buffer + passed to `MsgpackEncoder.encode_into()`: `send_multipart()` returns a + tracker for the last frame only. + """ + more_flag = zmq.SNDMORE if len(buffers) > 1 else 0 + tracker = socket.send(buffers[0], more_flag, copy=False, track=True) + if more_flag: + socket.send_multipart(buffers[1:], copy=False) + return tracker + def _handle_request_preproc_error(self, request: EngineCoreRequest) -> None: """Log and return a request-scoped error response for exceptions raised from the add request preprocessing in the input socket processing thread. diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index 9460fdf48f7..febaa10ce61 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -7,7 +7,7 @@ import sys import uuid import weakref from abc import ABC, abstractmethod -from collections import Counter, defaultdict, deque +from collections import Counter, defaultdict from collections.abc import Awaitable, Callable, Sequence from concurrent.futures import Future from dataclasses import dataclass @@ -671,11 +671,6 @@ class MPClient(EngineCoreClient): self.core_engine: EngineIdentity = self.core_engines[0] self.utility_results: dict[int, AnyFuture] = {} - # Request objects which may contain pytorch-allocated tensors - # that we need to keep references to until zmq is done with the - # underlying data. - self.pending_messages = deque[tuple[zmq.MessageTracker, Any]]() - # Start monitoring engine core processes for unexpected failures self.start_engine_core_monitor() @@ -707,14 +702,6 @@ class MPClient(EngineCoreClient): if self.resources.engine_dead: raise EngineDeadError() - def add_pending_message(self, tracker: zmq.MessageTracker, msg: Any): - if not tracker.done: - self.pending_messages.appendleft((tracker, msg)) - - def free_pending_messages(self): - while self.pending_messages and self.pending_messages[-1][0].done: - self.pending_messages.pop() - def dp_engines_running(self) -> bool: return self.engines_running @@ -896,17 +883,12 @@ class SyncMPClient(MPClient): def _send_input(self, request_type: EngineCoreRequestType, request: Any): self.ensure_alive() - self.free_pending_messages() # (Identity, RequestType, SerializedRequest) msg = (self.core_engine, request_type.value, *self.encoder.encode(request)) - - if len(msg) <= 3: - # No auxiliary buffers => no tensor backing buffers in request. - self.input_socket.send_multipart(msg, copy=False) - return - - tracker = self.input_socket.send_multipart(msg, copy=False, track=True) - self.add_pending_message(tracker, request) + # Any zero-copy tensor/ndarray frames are kept alive by zmq itself + # until it's finished sending them (there is a ref chain from the underlying + # memoryview back to the original owning tensor/ndarray). + self.input_socket.send_multipart(msg, copy=False) def call_utility(self, method: str, *args) -> Any: call_id = uuid.uuid1().int >> 64 @@ -1129,32 +1111,16 @@ class AsyncMPClient(MPClient): engine = self.core_engine message = (request_type.value, *self.encoder.encode(request)) - return self._send_input_message(message, engine, request) + return self._send_input_message(message, engine) def _send_input_message( - self, message: tuple[bytestr, ...], engine: EngineIdentity, objects: Any + self, message: tuple[bytestr, ...], engine: EngineIdentity ) -> Awaitable[Any]: - """ - objects is a reference to retain until zmq is finished with the - buffers, in case they were extracted from tensors in the request. - """ self.ensure_alive() - self.free_pending_messages() - - msg = (engine,) + message - if not objects or len(msg) <= 3: - # No auxiliary buffers => no tensor backing buffers in request. - return self.input_socket.send_multipart(msg, copy=False) - - future: asyncio.Future[zmq.MessageTracker] - future = self.input_socket.send_multipart(msg, copy=False, track=True) - - def add_pending(f: asyncio.Future[zmq.MessageTracker]): - with contextlib.suppress(BaseException): - self.add_pending_message(f.result(), objects) - - future.add_done_callback(add_pending) - return future + # Any zero-copy tensor/ndarray frames are kept alive by zmq itself + # until it's finished sending them (there is a ref chain from the underlying + # memoryview back to the original owning tensor/ndarray). + return self.input_socket.send_multipart((engine,) + message, copy=False) async def call_utility_async(self, method: str, *args) -> Any: return await self._call_utility_async(method, *args, engine=self.core_engine) @@ -1169,7 +1135,7 @@ class AsyncMPClient(MPClient): EngineCoreRequestType.UTILITY.value, *self.encoder.encode((self.client_index, call_id, method, args)), ) - await self._send_input_message(message, engine, args) + await self._send_input_message(message, engine) self._ensure_output_queue_task() return await future From ba702e978e3bc6af3a601cee10fefdeb49e7e8b5 Mon Sep 17 00:00:00 2001 From: Yiliu Dong <1098822169@qq.com> Date: Wed, 29 Jul 2026 00:17:22 +0800 Subject: [PATCH 179/185] [Attention] Skip sparse indexer scoring for dense short prefills (#48407) Signed-off-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com> Co-authored-by: OpenAI Codex <codex@openai.com> --- .../layers/test_mla_short_prefill_indexer.py | 165 ++++++++++++++++++ .../layers/attention/mla_attention.py | 13 +- .../layers/attention/sparse_mla_attention.py | 4 + vllm/model_executor/layers/mla.py | 23 ++- .../layers/sparse_attn_indexer.py | 28 ++- vllm/model_executor/models/deepseek_v2.py | 3 + vllm/models/deepseek_v32/nvidia/attention.py | 6 +- 7 files changed, 230 insertions(+), 12 deletions(-) create mode 100644 tests/model_executor/layers/test_mla_short_prefill_indexer.py diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py new file mode 100644 index 00000000000..6e1e10e8b45 --- /dev/null +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +import vllm.model_executor.layers.sparse_attn_indexer as sparse_indexer +from vllm.config import CUDAGraphMode +from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadata + +INDEXER_LAYER = "model.layers.0.self_attn.indexer.k_cache" +MLA_LAYER = "model.layers.0.self_attn.attn" + + +def make_indexer_metadata( + *, + num_decodes: int = 0, + num_decode_tokens: int = 0, + num_prefills: int = 1, + num_prefill_tokens: int = 1, + slot_mapping: torch.Tensor | None = None, +) -> DeepseekV32IndexerMetadata: + if slot_mapping is None: + slot_mapping = torch.zeros(num_prefill_tokens, dtype=torch.long) + return DeepseekV32IndexerMetadata( + seq_lens=torch.empty(0, dtype=torch.int32), + max_seq_len=2048, + slot_mapping=slot_mapping, + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + prefill=SimpleNamespace(chunks=[]) if num_prefills else None, + ) + + +def make_mla_metadata(*, use_dense_mha: bool = True, num_decode_tokens: int = 0): + return SimpleNamespace( + num_decode_tokens=num_decode_tokens, + prefill=SimpleNamespace(use_dense_mha=use_dense_mha), + ) + + +@pytest.mark.parametrize( + "batch_kind", + ["short", "threshold_mismatch", "force_mqa", "mla_decode", "capture", "full"], +) +def test_short_prefill_updates_k_cache_before_scoring_decision( + monkeypatch: pytest.MonkeyPatch, + batch_kind: str, +): + slot_mapping = torch.tensor([63, 64, 127, 128, -1]) + mla_num_decode_tokens = 1 if batch_kind == "mla_decode" else 0 + runtime_mode = ( + CUDAGraphMode.FULL if batch_kind == "full" else CUDAGraphMode.PIECEWISE + ) + should_skip = batch_kind in ("short", "threshold_mismatch") + num_decodes = int(batch_kind == "threshold_mismatch") + num_decode_tokens = 3 if batch_kind == "threshold_mismatch" else 0 + num_prefills = 0 if batch_kind == "threshold_mismatch" else 2 + num_prefill_tokens = 0 if batch_kind == "threshold_mismatch" else 5 + if batch_kind == "threshold_mismatch": + # With MTP=3 the indexer threshold is four. A main MLA backend whose + # threshold is one (for example FlashMLA under DCP) still routes this + # three-token extend through dense prefill attention. + slot_mapping = slot_mapping[:3] + indexer_metadata = make_indexer_metadata( + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + slot_mapping=slot_mapping, + ) + if indexer_metadata.num_decodes: + indexer_metadata.decode = object() + mla_metadata = make_mla_metadata( + use_dense_mha=batch_kind != "force_mqa", + num_decode_tokens=mla_num_decode_tokens, + ) + + observed: dict[str, object] = {} + + monkeypatch.setattr( + sparse_indexer, + "get_forward_context", + lambda: SimpleNamespace( + attn_metadata={ + INDEXER_LAYER: indexer_metadata, + MLA_LAYER: mla_metadata, + }, + cudagraph_runtime_mode=runtime_mode, + ), + ) + monkeypatch.setattr( + sparse_indexer.current_platform, "fp8_dtype", lambda: torch.float16 + ) + monkeypatch.setattr( + torch.cuda, + "is_current_stream_capturing", + lambda: batch_kind == "capture", + ) + + def record_cache_update(k, kv_cache, slots, block_size, scale_fmt): + observed.update(k=k.clone(), slots=slots) + + monkeypatch.setattr( + sparse_indexer.ops, "indexer_k_quant_and_cache", record_cache_update + ) + + class ScoringReached(Exception): + pass + + def scoring_trigger(): + if should_skip: + pytest.fail("short dense-MHA prefill must not enter indexer scoring") + raise ScoringReached + + def scoring_decode(*args): + raise ScoringReached + + monkeypatch.setattr(sparse_indexer, "current_workspace_manager", scoring_trigger) + monkeypatch.setattr( + sparse_indexer, + "kv_cache_as_quant_view", + scoring_decode, + ) + + hidden_states = torch.full((7, 1), float("inf")) + k = torch.arange(28, dtype=torch.float32).reshape(7, 4) + topk_indices = torch.full((7, 2048), 17, dtype=torch.int32) + + def run_indexer(): + return sparse_indexer.sparse_attn_indexer( + hidden_states, + INDEXER_LAYER, + torch.empty(1), + torch.full((7, 1), float("inf")), + None, + k, + torch.full((7, 1), float("inf")), + 128, + "ue8m0", + 2048, + 4, + 4096, + 4096, + topk_indices, + False, + False, + MLA_LAYER, + ) + + if should_skip: + assert run_indexer() is topk_indices + assert torch.all(topk_indices == 17) + else: + with pytest.raises(ScoringReached): + run_indexer() + assert torch.all(topk_indices == -1) + + # K cache is always updated before the scoring decision. + torch.testing.assert_close(observed["k"], k[: slot_mapping.numel()]) + assert observed["slots"] is slot_mapping diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 1dea276d2d2..16fb961e563 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -769,13 +769,8 @@ class MLAAttention(nn.Module, AttentionLayerBase): num_mha_tokens = q.size(0) - num_mqa_tokens if self.impl.is_sparse and num_mha_tokens > 0: - prefill_max_seq_len = attn_metadata.prefill_max_seq_len # type: ignore[attr-defined] - use_mha = ( - self.prefill_backend is not None - and prefill_max_seq_len <= attn_metadata.topk_tokens # type: ignore[attr-defined] - and not self._vllm_config.attention_config.sparse_mla_force_mqa - ) - if not use_mha: + prefill_metadata = getattr(attn_metadata, "prefill", None) + if not getattr(prefill_metadata, "use_dense_mha", False): num_mqa_tokens = q.size(0) num_mha_tokens = 0 @@ -1409,6 +1404,10 @@ class MLACommonPrefillMetadata: q_data_type: torch.dtype | None = None output_dtype: torch.dtype | None = None prefill_backend: MLAPrefillBackend | None = None + # Whether the prefill suffix is routed through dense MHA. + # Indexer scoring may be skipped only for a pure-prefill batch, + # since decode tokens still consume top-k indices. + use_dense_mha: bool = False @dataclass diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 19cad7986bf..1463f9fb35b 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -203,6 +203,10 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]): q_data_type=self.model_config.dtype, output_dtype=self.model_config.dtype, prefill_backend=self._prefill_backend, + use_dense_mha=( + prefill_max_seq_len <= self.topk_tokens + and not self.vllm_config.attention_config.sparse_mla_force_mqa + ), ) self._prefill_backend.prepare_metadata(prefill) diff --git a/vllm/model_executor/layers/mla.py b/vllm/model_executor/layers/mla.py index ca4a20874a0..ac956461415 100644 --- a/vllm/model_executor/layers/mla.py +++ b/vllm/model_executor/layers/mla.py @@ -8,6 +8,7 @@ from vllm.config import CacheConfig from vllm.model_executor.custom_op import PluggableLayer from vllm.model_executor.layers.attention import MLAAttention from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.platforms import current_platform @dataclass @@ -65,6 +66,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer): quant_config: QuantizationConfig | None = None, prefix: str = "", skip_topk: bool = False, + allow_short_prefill_indexer_scoring_skip: bool = False, ) -> None: super().__init__() self.hidden_size = hidden_size @@ -119,7 +121,26 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer): indexer=self.indexer, topk_indices_buffer=mla_modules.topk_indices_buffer, ) - + indexer_op = getattr(self.indexer, "indexer_op", None) + if indexer_op is not None and hasattr( + indexer_op, "dense_mha_metadata_layer_name" + ): + enable_short_prefill_scoring_skip = ( + allow_short_prefill_indexer_scoring_skip + and not self.skip_topk + and not getattr(indexer_op, "use_pcp", False) + and current_platform.is_cuda() + ) + # The indexer and main MLA use independent decode thresholds and + # may classify the same short extend differently. Bind the main + # MLA layer name so the eager indexer op can check whether the + # batch's top-k indices will be consumed. + # PCP is excluded because indexer cache/scoring ownership differs + # across ranks and the no-consumer invariant has not been + # established there. + indexer_op.dense_mha_metadata_layer_name = ( + self.mla_attn.layer_name if enable_short_prefill_scoring_skip else "" + ) self.prefix = prefix def forward( diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 5b8e2bf008e..9a671be5563 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -8,7 +8,7 @@ import vllm.envs as envs from vllm import _custom_ops as ops from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.breakable_cudagraph import eager_break_during_capture -from vllm.config import get_current_vllm_config +from vllm.config import CUDAGraphMode, get_current_vllm_config from vllm.distributed import get_dcp_group, get_pcp_group from vllm.forward_context import get_forward_context from vllm.logger import init_logger @@ -310,6 +310,7 @@ def sparse_attn_indexer( topk_indices_buffer: torch.Tensor, skip_k_cache_insert: bool, use_pcp: bool, + dense_mha_metadata_layer_name: LayerNameType, use_fp4_cache: bool = False, dcp_rank: int = 0, dcp_world_size: int = 1, @@ -317,7 +318,8 @@ def sparse_attn_indexer( skip_topk_buffer_clear: bool = False, ) -> torch.Tensor: # careful! this will be None in dummy run - attn_metadata = get_forward_context().attn_metadata + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata fp8_dtype = current_platform.fp8_dtype() k_cache_prefix = _resolve_layer_name(k_cache_prefix) @@ -357,6 +359,7 @@ def sparse_attn_indexer( topk_indices_buffer, skip_k_cache_insert, use_pcp, + dense_mha_metadata_layer_name, use_fp4_cache, ) attn_metadata_narrowed = attn_metadata[k_cache_prefix] @@ -402,6 +405,24 @@ def sparse_attn_indexer( scale_fmt, ) + # The indexer and main MLA may classify the same short extend differently + # because they use independent decode thresholds. Only the main MLA route + # can determine whether the top-k indices will be consumed. + if forward_context.cudagraph_runtime_mode != CUDAGraphMode.FULL: + dense_mha_layer = _resolve_layer_name(dense_mha_metadata_layer_name) + if dense_mha_layer: + mla_metadata = attn_metadata.get(dense_mha_layer) + prefill_metadata = getattr(mla_metadata, "prefill", None) + if ( + getattr(prefill_metadata, "use_dense_mha", False) + and getattr(mla_metadata, "num_decode_tokens", -1) == 0 + and not torch.cuda.is_current_stream_capturing() + ): + # Deliberately leave the buffer untouched. Dense MHA does not + # consume top-k indices for this batch; clearing it would be + # unnecessary work. + return topk_indices_buffer + # The buffer must be pre-filled with -1 (the "no token" sentinel) before the # top-k kernels scatter valid indices into it. On the fused deepseek_v32 # nvidia path, _fused_norm_rope_kernel already cleared the same @@ -684,6 +705,7 @@ def sparse_attn_indexer_fake( topk_indices_buffer: torch.Tensor | None, skip_k_cache_insert: bool, use_pcp: bool, + dense_mha_metadata_layer_name: LayerNameType, use_fp4_cache: bool = False, dcp_rank: int = 0, dcp_world_size: int = 1, @@ -739,6 +761,7 @@ class SparseAttnIndexer(CustomOp): self.topk_indices_buffer = topk_indices_buffer self.skip_k_cache_insert = skip_k_cache_insert self.use_fp4_cache = use_fp4_cache + self.dense_mha_metadata_layer_name = "" # DCP scalars are constant for the run; resolve them here (config is set # during model construction) and pass them into the custom op, rather # than threading them through per-step metadata. @@ -800,6 +823,7 @@ class SparseAttnIndexer(CustomOp): self.topk_indices_buffer, self.skip_k_cache_insert, self.use_pcp, + _encode_layer_name(self.dense_mha_metadata_layer_name), self.use_fp4_cache, self.dcp_rank, self.dcp_world_size, diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 4b92e351caa..bf67e040a15 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -1178,6 +1178,9 @@ class DeepseekV2MLAAttention(nn.Module): # the V1 proposer. A frozen True would leave the draft reading a # never-written topk buffer. skip_topk=_skip_topk and not is_mtp_layer, + # Do not skip scoring for MTP layers: their top-k buffer may be + # reused by later draft iterations through index sharing. + allow_short_prefill_indexer_scoring_skip=not is_mtp_layer, ) def forward( diff --git a/vllm/models/deepseek_v32/nvidia/attention.py b/vllm/models/deepseek_v32/nvidia/attention.py index dcf955ad59b..0e604da87ba 100644 --- a/vllm/models/deepseek_v32/nvidia/attention.py +++ b/vllm/models/deepseek_v32/nvidia/attention.py @@ -494,8 +494,10 @@ class DeepseekV32Attention(MLAAttention): self.indexer.max_model_len, self.indexer.max_total_seq_len, self.topk_indices_buffer, - True, # skip_k_cache_insert - False, # use_fp4_cache + skip_k_cache_insert=True, + use_pcp=False, + dense_mha_metadata_layer_name="", + use_fp4_cache=False, # fused_norm_rope already cleared the topk buffer this forward. skip_topk_buffer_clear=True, ) From 01661cc57f48ce95c639efce7c88e6dd37349007 Mon Sep 17 00:00:00 2001 From: Bugen Zhao <i@bugenzhao.com> Date: Wed, 29 Jul 2026 00:21:12 +0800 Subject: [PATCH 180/185] [Rust][Benchmark] Make `vllm bench serve` Rust delegation opt-in (#50081) Signed-off-by: Bugen Zhao <i@bugenzhao.com> --- .../entrypoints/openai/test_dp_supervisor.py | 1 + tests/test_envs.py | 9 +++ vllm/entrypoints/cli/benchmark/main.py | 20 ++++++ vllm/entrypoints/cli/benchmark/serve.py | 62 +------------------ vllm/entrypoints/cli/main.py | 2 + vllm/entrypoints/cli/serve.py | 14 +++-- vllm/entrypoints/openai/dp_supervisor.py | 2 +- vllm/envs.py | 27 ++++---- 8 files changed, 61 insertions(+), 76 deletions(-) diff --git a/tests/entrypoints/openai/test_dp_supervisor.py b/tests/entrypoints/openai/test_dp_supervisor.py index 576e7ef16df..df80053deb5 100644 --- a/tests/entrypoints/openai/test_dp_supervisor.py +++ b/tests/entrypoints/openai/test_dp_supervisor.py @@ -201,6 +201,7 @@ def test_run_vllm_dp_server_uses_rust_frontend_when_enabled(monkeypatch): monkeypatch.setattr(dp_sup.os, "setpgrp", lambda: None) monkeypatch.setattr(dp_sup, "set_process_title", lambda *_args: None) monkeypatch.setattr(dp_sup, "decorate_logs", lambda *_args: None) + monkeypatch.setattr(dp_sup.envs, "VLLM_USE_RUST_FRONTEND", True, raising=False) monkeypatch.setattr( dp_sup.envs, "VLLM_RUST_FRONTEND_PATH", diff --git a/tests/test_envs.py b/tests/test_envs.py index 56c04dd6f2e..5917c28fab2 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -145,6 +145,15 @@ def test_precompiled_install_flags_are_orthogonal() -> None: assert environment_variables["VLLM_USE_PRECOMPILED_RUST"]() is True +def test_rust_bench_auto_path_missing_fails_fast() -> None: + with ( + patch.dict(os.environ, {"VLLM_USE_RUST_BENCH": "1"}, clear=True), + patch("vllm.envs.os.path.isfile", return_value=False), + pytest.raises(FileNotFoundError, match="vllm-rs binary was not found"), + ): + environment_variables["VLLM_RUST_FRONTEND_PATH"]() + + class TestEnvWithChoices: """Test cases for env_with_choices function.""" diff --git a/vllm/entrypoints/cli/benchmark/main.py b/vllm/entrypoints/cli/benchmark/main.py index 1afac64b148..9ea49987091 100644 --- a/vllm/entrypoints/cli/benchmark/main.py +++ b/vllm/entrypoints/cli/benchmark/main.py @@ -2,18 +2,38 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import argparse +import os import sys import typing +from vllm import envs from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase from vllm.entrypoints.cli.types import CLISubcommand from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG +from vllm.logger import init_logger if typing.TYPE_CHECKING: from vllm.utils.argparse_utils import FlexibleArgumentParser else: FlexibleArgumentParser = argparse.ArgumentParser +logger = init_logger(__name__) + + +def maybe_exec_rust_bench() -> None: + if sys.argv[1:3] != ["bench", "serve"] or not envs.VLLM_USE_RUST_BENCH: + return + + rust_cli = envs.VLLM_RUST_FRONTEND_PATH + if rust_cli is None: + raise RuntimeError( + "VLLM_USE_RUST_BENCH=1 requires VLLM_RUST_FRONTEND_PATH " + "to resolve to the vllm-rs binary." + ) + + logger.info("Delegating `vllm bench serve` to Rust binary at %s.", rust_cli) + os.execv(rust_cli, [rust_cli, "bench", "serve", *sys.argv[3:]]) + def _import_bench_subcommand_modules() -> None: # Imported lazily so `BenchmarkSubcommandBase` subclasses register only diff --git a/vllm/entrypoints/cli/benchmark/serve.py b/vllm/entrypoints/cli/benchmark/serve.py index 41a65273ba8..188afd6c703 100644 --- a/vllm/entrypoints/cli/benchmark/serve.py +++ b/vllm/entrypoints/cli/benchmark/serve.py @@ -1,68 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import argparse -import os -import sys -from pathlib import Path -from vllm.benchmarks.serve import add_cli_args -from vllm.benchmarks.serve import main as python_main +from vllm.benchmarks.serve import add_cli_args, main from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase -from vllm.logger import init_logger from vllm.utils.argparse_utils import FlexibleArgumentParser -logger = init_logger(__name__) -_RUST_CLI_PATH = Path(__file__).resolve().parents[3] / "vllm-rs" -_RUST_SUPPORTED_DATASETS = frozenset( - { - "custom", - "hf", - "prefix_repetition", - "random", - "random-mm", - "random-rerank", - "sharegpt", - "sonnet", - "speed_bench", - } -) -_RUST_SUPPORTED_BACKENDS = frozenset( - { - "openai", - "openai-chat", - "openai-embeddings", - "openai-embeddings-chat", - "vllm", - "vllm-pooling", - "vllm-rerank", - } -) - - -def _rust_unsupported_reason(args: argparse.Namespace) -> str | None: - if args.dataset_name not in _RUST_SUPPORTED_DATASETS: - return f"dataset {args.dataset_name!r} is not supported by the Rust benchmark" - if args.backend not in _RUST_SUPPORTED_BACKENDS: - return f"backend {args.backend!r} is not supported by the Rust benchmark" - return None - - -def _maybe_exec_rust_bench(args: argparse.Namespace) -> None: - if reason := _rust_unsupported_reason(args): - logger.info("Using Python benchmark: %s.", reason) - return - - if not _RUST_CLI_PATH.is_file(): - logger.warning( - "Rust benchmark binary not found at %s; falling back to Python.", - _RUST_CLI_PATH, - ) - return - - rust_cli = str(_RUST_CLI_PATH) - logger.info("Delegating `vllm bench serve` to Rust binary at %s.", rust_cli) - os.execv(rust_cli, [rust_cli, "bench", "serve", *sys.argv[3:]]) - class BenchmarkServingSubcommand(BenchmarkSubcommandBase): """The `serve` subcommand for `vllm bench`.""" @@ -76,5 +19,4 @@ class BenchmarkServingSubcommand(BenchmarkSubcommandBase): @staticmethod def cmd(args: argparse.Namespace) -> None: - _maybe_exec_rust_bench(args) - python_main(args) + main(args) diff --git a/vllm/entrypoints/cli/main.py b/vllm/entrypoints/cli/main.py index fe0b339b3ed..3dc69dd3ad2 100644 --- a/vllm/entrypoints/cli/main.py +++ b/vllm/entrypoints/cli/main.py @@ -54,6 +54,8 @@ def main(): logger.info("Delegating entrypoint handling to vllm-omni") omni_main() else: + vllm.entrypoints.cli.benchmark.main.maybe_exec_rust_bench() + # For 'vllm bench *': use CPU instead of UnspecifiedPlatform by default if len(sys.argv) > 1 and sys.argv[1] == "bench": logger.debug( diff --git a/vllm/entrypoints/cli/serve.py b/vllm/entrypoints/cli/serve.py index d5e9b2bc874..08cb79f2081 100644 --- a/vllm/entrypoints/cli/serve.py +++ b/vllm/entrypoints/cli/serve.py @@ -58,6 +58,10 @@ class ServeSubcommand(CLISubcommand): uvloop.run(serve_grpc(args)) return + rust_frontend_path = ( + envs.VLLM_RUST_FRONTEND_PATH if envs.VLLM_USE_RUST_FRONTEND else None + ) + if args.headless: if args.api_server_count is not None and args.api_server_count > 0: raise ValueError( @@ -103,7 +107,7 @@ class ServeSubcommand(CLISubcommand): # - Hybrid LB: Use local DP size (internal LB for local ranks only) # - Internal LB: Use full DP size if args.api_server_count is None: - if is_multi_port or is_external_lb or envs.VLLM_RUST_FRONTEND_PATH: + if is_multi_port or is_external_lb or rust_frontend_path: args.api_server_count = 1 elif is_hybrid_lb: args.api_server_count = args.data_parallel_size_local or 1 @@ -120,7 +124,7 @@ class ServeSubcommand(CLISubcommand): "Defaulting api_server_count to data_parallel_size (%d).", args.api_server_count, ) - elif envs.VLLM_RUST_FRONTEND_PATH and args.api_server_count > 1: + elif rust_frontend_path and args.api_server_count > 1: logger.warning( "Ignoring --api-server-count=%d when using rust front-end process", args.api_server_count, @@ -140,7 +144,7 @@ class ServeSubcommand(CLISubcommand): run_dp_supervisor(args) elif args.api_server_count < 1: run_headless(args) - elif args.api_server_count > 1 or envs.VLLM_RUST_FRONTEND_PATH: + elif args.api_server_count > 1 or rust_frontend_path: run_multi_api_server(args) else: # Single API server (this process). @@ -256,7 +260,9 @@ def run_headless(args: argparse.Namespace): def run_multi_api_server(args: argparse.Namespace): assert not args.headless - rust_frontend_path = envs.VLLM_RUST_FRONTEND_PATH + rust_frontend_path = ( + envs.VLLM_RUST_FRONTEND_PATH if envs.VLLM_USE_RUST_FRONTEND else None + ) num_api_servers: int = args.api_server_count assert num_api_servers > 0 diff --git a/vllm/entrypoints/openai/dp_supervisor.py b/vllm/entrypoints/openai/dp_supervisor.py index d669ec4d1d5..8ce6233c1ba 100644 --- a/vllm/entrypoints/openai/dp_supervisor.py +++ b/vllm/entrypoints/openai/dp_supervisor.py @@ -257,7 +257,7 @@ def _run_vllm_dp_server(child_args: argparse.Namespace) -> None: name = f"APIServer_DP{child_args.data_parallel_rank}" set_process_title(name) decorate_logs(name) - if envs.VLLM_RUST_FRONTEND_PATH: + if envs.VLLM_USE_RUST_FRONTEND and envs.VLLM_RUST_FRONTEND_PATH: _run_rust_vllm_dp_server(child_args) else: _run_python_vllm_dp_server(child_args) diff --git a/vllm/envs.py b/vllm/envs.py index f7e01c33275..84ec5a8af85 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -153,6 +153,7 @@ if TYPE_CHECKING: K_SCALE_CONSTANT: int = 200 V_SCALE_CONSTANT: int = 100 VLLM_USE_RUST_FRONTEND: bool = False + VLLM_USE_RUST_BENCH: bool = False VLLM_RUST_FRONTEND_PATH: str | None = "auto" VLLM_SERVER_DEV_MODE: bool = False VLLM_V1_OUTPUT_PROC_CHUNK_SIZE: int = 128 @@ -548,22 +549,24 @@ def _deprecated_triton_attn_use_td() -> None: return None -def _resolve_rust_frontend_path() -> str | None: - """Resolve the Rust frontend binary path. +def _resolve_rust_cli_path() -> str | None: + """Resolve the vllm-rs binary path. - Returns None if VLLM_USE_RUST_FRONTEND is not enabled. + Returns None unless VLLM_USE_RUST_FRONTEND or VLLM_USE_RUST_BENCH is enabled. When enabled, resolves VLLM_RUST_FRONTEND_PATH ("auto" by default) to the actual binary path. """ - use_rust = bool(int(os.environ.get("VLLM_USE_RUST_FRONTEND", "0"))) + use_rust = bool(int(os.environ.get("VLLM_USE_RUST_FRONTEND", "0"))) or bool( + int(os.environ.get("VLLM_USE_RUST_BENCH", "0")) + ) raw = os.environ.get("VLLM_RUST_FRONTEND_PATH", "auto") if not use_rust: if os.environ.get("VLLM_RUST_FRONTEND_PATH") is not None: logger.warning( - "VLLM_RUST_FRONTEND_PATH is set but VLLM_USE_RUST_FRONTEND " - "is not enabled. The Rust frontend will not be used. " - "Set VLLM_USE_RUST_FRONTEND=1 to enable it." + "VLLM_RUST_FRONTEND_PATH is set without enabling " + "VLLM_USE_RUST_FRONTEND or VLLM_USE_RUST_BENCH. " + "Set one of them to 1 to use the vllm-rs binary." ) return None @@ -1340,10 +1343,12 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_USE_RUST_FRONTEND": lambda: bool( int(os.getenv("VLLM_USE_RUST_FRONTEND", "0")) ), - # Path to the Rust frontend binary. Defaults to "auto" which discovers - # the binary installed with the vllm package. Only used when - # VLLM_USE_RUST_FRONTEND=1. - "VLLM_RUST_FRONTEND_PATH": lambda: _resolve_rust_frontend_path(), + # If set, use the packaged Rust client for `vllm bench serve`. + "VLLM_USE_RUST_BENCH": lambda: bool(int(os.getenv("VLLM_USE_RUST_BENCH", "0"))), + # Path to the vllm-rs binary. Defaults to "auto" which discovers the + # binary installed with the vllm package. Used when VLLM_USE_RUST_FRONTEND=1 + # or VLLM_USE_RUST_BENCH=1. + "VLLM_RUST_FRONTEND_PATH": lambda: _resolve_rust_cli_path(), # If set, vllm will run in development mode, which will enable # some additional endpoints for developing and debugging, # e.g. `/reset_prefix_cache` From 4f56321d7ec25dd041c7bf0aa47b4eadd075a7e5 Mon Sep 17 00:00:00 2001 From: jiacao-amd <jiahui.cao@amd.com> Date: Tue, 28 Jul 2026 09:34:16 -0700 Subject: [PATCH 181/185] [ROCm] Cache fp32 upcast of static e8m0 weight scale in AITER scaled_mm (#47773) Signed-off-by: jiacao-amd <jiahui.cao@amd.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com> --- .../kernels/linear/scaled_mm/aiter.py | 24 ++++++++++++------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py index 1b39491ab34..da8f69fa97b 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py @@ -9,6 +9,9 @@ from vllm._aiter_ops import ( rocm_aiter_ops, ) from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + _upcast_e8m0_to_fp32, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, ) @@ -16,6 +19,7 @@ from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform from .BlockScaledMMLinearKernel import ( + FP8BlockParams, Fp8BlockScaledMMLinearKernel, ) from .cutlass import CutlassInt8ScaledMMLinearKernel @@ -375,6 +379,17 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): and rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k) ) + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + super().process_weights_after_loading(layer) + + params = FP8BlockParams.from_layer(layer) + if params.weight_scale_inv is not None: + ws, attr = params.weight_scale_inv, params.WEIGHT_SCALE_INV + else: + ws, attr = params.weight_scale, params.WEIGHT_SCALE + if ws is not None and ws.dtype == torch.float8_e8m0fnu: + replace_parameter(layer, attr, _upcast_e8m0_to_fp32(ws).contiguous()) + @classmethod def is_supported(cls, compute_capability=None): return ( @@ -406,19 +421,12 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): Bs: torch.Tensor, ) -> torch.Tensor: if As.dtype != Bs.dtype: - from vllm.model_executor.layers.quantization.utils.fp8_utils import ( - _upcast_e8m0_to_fp32, - ) - if As.dtype == torch.float8_e8m0fnu: As = _upcast_e8m0_to_fp32(As).contiguous() else: As = As.to(torch.float32) - if Bs.dtype == torch.float8_e8m0fnu: - Bs = _upcast_e8m0_to_fp32(Bs).contiguous() - else: - Bs = Bs.to(torch.float32) + Bs = Bs.to(torch.float32) out_dtype = self.config.out_dtype if self.use_triton: From 05a08148631ab2d852c1984ddb3abebad8039a29 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas <akaratza@amd.com> Date: Tue, 28 Jul 2026 11:50:10 -0500 Subject: [PATCH 182/185] [ROCm] Fix and optimize GPT-J-style MRoPE (#49906) Signed-off-by: Andreas Karatzas <akaratza@amd.com> --- tests/kernels/core/test_mrope.py | 10 +- .../layers/rotary_embedding/mrope.py | 149 +++++++++++------- 2 files changed, 100 insertions(+), 59 deletions(-) diff --git a/tests/kernels/core/test_mrope.py b/tests/kernels/core/test_mrope.py index 29051b4a00c..6f64fabfe9b 100644 --- a/tests/kernels/core/test_mrope.py +++ b/tests/kernels/core/test_mrope.py @@ -38,6 +38,7 @@ def generate_test_data( class MRoPETestInfo(NamedTuple): model_name: str + is_neox_style: bool = True # https://github.com/pytorch/pytorch/blob/main/torch/testing/_comparison.py#L1317 atol: float = 1e-2 rtol: float = 1.6e-2 @@ -45,7 +46,10 @@ class MRoPETestInfo(NamedTuple): MODELS_TO_TEST = [ - MRoPETestInfo(model_name="zai-org/GLM-4.1V-9B-Thinking"), + MRoPETestInfo( + model_name="zai-org/GLM-4.1V-9B-Thinking", + is_neox_style=False, + ), MRoPETestInfo(model_name="Qwen/Qwen2-VL-7B-Instruct"), MRoPETestInfo(model_name="Qwen/Qwen2-VL-72B-Instruct"), MRoPETestInfo(model_name="Qwen/Qwen2.5-VL-72B-Instruct"), @@ -92,7 +96,7 @@ def test_mrope( if hasattr(config, "head_dim") else config.hidden_size // total_num_heads ) - is_neox_style = True + is_neox_style = model_info.is_neox_style max_position = config.max_position_embeddings @@ -162,7 +166,7 @@ def test_mrope_torch_compile_tracing( if hasattr(config, "head_dim") else config.hidden_size // total_num_heads ) - is_neox_style = True + is_neox_style = model_info.is_neox_style max_position = config.max_position_embeddings mrope_helper_class = get_rope( diff --git a/vllm/model_executor/layers/rotary_embedding/mrope.py b/vllm/model_executor/layers/rotary_embedding/mrope.py index 3c946dd130c..29ce9e5000d 100644 --- a/vllm/model_executor/layers/rotary_embedding/mrope.py +++ b/vllm/model_executor/layers/rotary_embedding/mrope.py @@ -5,6 +5,7 @@ import numpy as np import torch +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from .base import RotaryEmbeddingBase @@ -24,16 +25,17 @@ def _triton_mrope_forward( rd: tl.constexpr, pad_n_qh: tl.constexpr, pad_n_kh: tl.constexpr, - pad_hd: tl.constexpr, + pad_rd: tl.constexpr, mrope_section_t: tl.constexpr, mrope_section_h: tl.constexpr, mrope_section_w: tl.constexpr, is_interleaved: tl.constexpr, + is_neox_style: tl.constexpr, ): # Adapted from # https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/qwen2vl_mrope.py # This version supports flatten input tensors from vllm - # and supports cos and sin cache with shape (3, num_tokens, head_dim // 2) + # and supports cos and sin cache with shape (3, num_tokens, rotary_dim // 2) # instead of (3, bsz, seq_len, head_dim), also supports interleaved rotary pid = tl.program_id(0) # locate start address @@ -44,9 +46,9 @@ def _triton_mrope_forward( # get the cos(mθ_{i...d/2}) and sin(mθ_{i...d/2}) for token position # m of this program instance # #################################################################### - # Note: cos and sin now have shape (3, num_tokens, head_dim // 2) + # Note: cos and sin now have shape (3, num_tokens, rotary_dim // 2) - # Updated stride calculation for half head_dim + # Updated stride calculation for half rotary_dim half_rd = rd // 2 t_cos = cos + pid * half_rd h_cos = t_cos + num_tokens * half_rd @@ -55,12 +57,17 @@ def _triton_mrope_forward( h_sin = t_sin + num_tokens * half_rd w_sin = h_sin + num_tokens * half_rd - # Updated offsets for half head_dim - cos_offsets = tl.arange(0, pad_hd // 2) + # Updated offsets for half rotary_dim + cos_offsets = tl.arange(0, pad_rd // 2) if is_interleaved: - h_mask = ((cos_offsets % 3) == 1) & (cos_offsets <= 3 * mrope_section_h) - w_mask = ((cos_offsets % 3) == 2) & (cos_offsets <= 3 * mrope_section_w) - t_mask = ~(h_mask | w_mask) + valid_mask = cos_offsets < half_rd + h_mask = ( + valid_mask & ((cos_offsets % 3) == 1) & (cos_offsets <= 3 * mrope_section_h) + ) + w_mask = ( + valid_mask & ((cos_offsets % 3) == 2) & (cos_offsets <= 3 * mrope_section_w) + ) + t_mask = valid_mask & ~(h_mask | w_mask) else: t_end = mrope_section_t h_end = t_end + mrope_section_h @@ -79,55 +86,74 @@ def _triton_mrope_forward( sin_row = t_sin_row + h_sin_row + w_sin_row # #################################################################### - # Load the left and right half of q and k for the current - # program instance (i.e. for the current token) separately + # Load the two values in each rotary pair for the current token. + # NeoX pairs the first and second halves, while GPT-J pairs + # adjacent values. # #################################################################### - # left half of the head - first_half_q_offsets = ( - tl.arange(0, pad_n_qh)[:, None] * hd + tl.arange(0, pad_hd // 2)[None, :] - ) - first_half_k_offsets = ( - tl.arange(0, pad_n_kh)[:, None] * hd + tl.arange(0, pad_hd // 2)[None, :] - ) - first_q_mask = (tl.arange(0, pad_n_qh)[:, None] < n_qh) & ( - tl.arange(0, pad_hd // 2)[None, :] < rd // 2 - ) - first_k_mask = (tl.arange(0, pad_n_kh)[:, None] < n_kh) & ( - tl.arange(0, pad_hd // 2)[None, :] < rd // 2 - ) + if is_neox_style: + rotary_offsets = tl.arange(0, pad_rd // 2) + first_q_offsets = tl.arange(0, pad_n_qh)[:, None] * hd + rotary_offsets[None, :] + first_k_offsets = tl.arange(0, pad_n_kh)[:, None] * hd + rotary_offsets[None, :] + first_q_mask = (tl.arange(0, pad_n_qh)[:, None] < n_qh) & ( + rotary_offsets[None, :] < rd // 2 + ) + first_k_mask = (tl.arange(0, pad_n_kh)[:, None] < n_kh) & ( + rotary_offsets[None, :] < rd // 2 + ) - q_tile_1 = tl.load(q_ptr + first_half_q_offsets, mask=first_q_mask, other=0).to( - sin_row.dtype - ) - k_tile_1 = tl.load(k_ptr + first_half_k_offsets, mask=first_k_mask, other=0).to( - sin_row.dtype - ) + q_tile_1 = tl.load(q_ptr + first_q_offsets, mask=first_q_mask, other=0).to( + sin_row.dtype + ) + k_tile_1 = tl.load(k_ptr + first_k_offsets, mask=first_k_mask, other=0).to( + sin_row.dtype + ) - # right half of the head - second_half_q_offsets = first_half_q_offsets + (rd // 2) - second_half_k_offsets = first_half_k_offsets + (rd // 2) - second_q_mask = first_q_mask - second_k_mask = first_k_mask + second_q_offsets = first_q_offsets + (rd // 2) + second_k_offsets = first_k_offsets + (rd // 2) + q_tile_2 = tl.load(q_ptr + second_q_offsets, mask=first_q_mask, other=0).to( + sin_row.dtype + ) + k_tile_2 = tl.load(k_ptr + second_k_offsets, mask=first_k_mask, other=0).to( + sin_row.dtype + ) - q_tile_2 = tl.load(q_ptr + second_half_q_offsets, mask=second_q_mask, other=0).to( - sin_row.dtype - ) - k_tile_2 = tl.load(k_ptr + second_half_k_offsets, mask=second_k_mask, other=0).to( - sin_row.dtype - ) + new_q_tile_1 = q_tile_1 * cos_row - q_tile_2 * sin_row + tl.store(q_ptr + first_q_offsets, new_q_tile_1, mask=first_q_mask) + new_q_tile_2 = q_tile_2 * cos_row + q_tile_1 * sin_row + tl.store(q_ptr + second_q_offsets, new_q_tile_2, mask=first_q_mask) - # y = [x1, x2] * [cos, cos] + [-x2, x1] * [sin, sin] - # Since cos and sin are now half-size, - # we use the same cos_row and sin_row for both halves - new_q_tile_1 = q_tile_1 * cos_row - q_tile_2 * sin_row - tl.store(q_ptr + first_half_q_offsets, new_q_tile_1, mask=first_q_mask) - new_q_tile_2 = q_tile_2 * cos_row + q_tile_1 * sin_row - tl.store(q_ptr + second_half_q_offsets, new_q_tile_2, mask=second_q_mask) + new_k_tile_1 = k_tile_1 * cos_row - k_tile_2 * sin_row + tl.store(k_ptr + first_k_offsets, new_k_tile_1, mask=first_k_mask) + new_k_tile_2 = k_tile_2 * cos_row + k_tile_1 * sin_row + tl.store(k_ptr + second_k_offsets, new_k_tile_2, mask=first_k_mask) + else: + # Load and store adjacent rotary pairs contiguously. Using stride-two + # even/odd offsets makes Triton emit scalar 16-bit memory operations on + # AMD, while split/interleave only rearranges values in registers. + rotary_offsets = tl.arange(0, pad_rd) + q_offsets = tl.arange(0, pad_n_qh)[:, None] * hd + rotary_offsets[None, :] + k_offsets = tl.arange(0, pad_n_kh)[:, None] * hd + rotary_offsets[None, :] + q_mask = (tl.arange(0, pad_n_qh)[:, None] < n_qh) & ( + rotary_offsets[None, :] < rd + ) + k_mask = (tl.arange(0, pad_n_kh)[:, None] < n_kh) & ( + rotary_offsets[None, :] < rd + ) - new_k_tile_1 = k_tile_1 * cos_row - k_tile_2 * sin_row - tl.store(k_ptr + first_half_k_offsets, new_k_tile_1, mask=first_k_mask) - new_k_tile_2 = k_tile_2 * cos_row + k_tile_1 * sin_row - tl.store(k_ptr + second_half_k_offsets, new_k_tile_2, mask=second_k_mask) + q_tile = tl.load(q_ptr + q_offsets, mask=q_mask, other=0).to(sin_row.dtype) + k_tile = tl.load(k_ptr + k_offsets, mask=k_mask, other=0).to(sin_row.dtype) + q_tile_1, q_tile_2 = tl.split(tl.reshape(q_tile, (pad_n_qh, pad_rd // 2, 2))) + k_tile_1, k_tile_2 = tl.split(tl.reshape(k_tile, (pad_n_kh, pad_rd // 2, 2))) + + new_q_tile_1 = q_tile_1 * cos_row - q_tile_2 * sin_row + new_q_tile_2 = q_tile_2 * cos_row + q_tile_1 * sin_row + new_q_tile = tl.interleave(new_q_tile_1, new_q_tile_2) + tl.store(q_ptr + q_offsets, new_q_tile, mask=q_mask) + + new_k_tile_1 = k_tile_1 * cos_row - k_tile_2 * sin_row + new_k_tile_2 = k_tile_2 * cos_row + k_tile_1 * sin_row + new_k_tile = tl.interleave(new_k_tile_1, new_k_tile_2) + tl.store(k_ptr + k_offsets, new_k_tile, mask=k_mask) def triton_mrope( @@ -139,23 +165,26 @@ def triton_mrope( head_size: int, rotary_dim: int, mrope_interleaved: bool, + is_neox_style: bool, ) -> tuple[torch.Tensor, torch.Tensor]: """Qwen2VL mrope kernel. Args: q: [num_tokens, num_heads * head_size] k: [num_tokens, num_kv_heads * head_size] - cos: [3, num_tokens, head_size //2 ] + cos: [3, num_tokens, rotary_dim // 2] (T/H/W positions with multimodal inputs) - sin: [3, num_tokens, head_size //2 ] + sin: [3, num_tokens, rotary_dim // 2] (T/H/W positions with multimodal inputs) mrope_section: [t, h, w] head_size: int + is_neox_style: Whether rotary pairs use split-half (NeoX) or + adjacent (GPT-J) layout. """ n_row, n_q_head_head_dim = q.shape n_q_head = n_q_head_head_dim // head_size n_kv_head = k.shape[1] // head_size - pad_hd = triton.next_power_of_2(head_size) + pad_rd = triton.next_power_of_2(rotary_dim) pad_n_q_head = triton.next_power_of_2(n_q_head) pad_n_kv_head = triton.next_power_of_2(n_kv_head) @@ -166,6 +195,11 @@ def triton_mrope( cos = cos.contiguous() sin = sin.contiguous() + # Small adjacent-pair tiles perform best with one wave per program on + # ROCm. Keep the existing launch shape for larger rotary dimensions, + # NeoX, and other backends. + use_single_wave = current_platform.is_rocm() and not is_neox_style and pad_rd <= 64 + num_warps = 1 if use_single_wave else 4 _triton_mrope_forward[(n_row,)]( q, k, @@ -178,11 +212,13 @@ def triton_mrope( rotary_dim, pad_n_q_head, pad_n_kv_head, - pad_hd, + pad_rd, mrope_section[0], mrope_section[1], mrope_section[2], mrope_interleaved, + is_neox_style, + num_warps=num_warps, ) return q, k @@ -349,6 +385,7 @@ class MRotaryEmbedding(RotaryEmbeddingBase): self.head_size, self.rotary_dim, self.mrope_interleaved, + self.is_neox_style, ) return q.reshape(query_shape), k.reshape(key_shape) From 8a7b3c299053efbca2669081ddf50a46b6a9d149 Mon Sep 17 00:00:00 2001 From: Brian Dellabetta <brian-dellabetta@users.noreply.github.com> Date: Tue, 28 Jul 2026 13:03:08 -0400 Subject: [PATCH 183/185] [compressed-tensors] update `find_matched_target` order to prioritize fused name matches over class match (#49483) Signed-off-by: Brian Dellabetta <bdellabe@redhat.com> --- .../layers/quantization/compressed_tensors/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/utils.py b/vllm/model_executor/layers/quantization/compressed_tensors/utils.py index afb899cd6d7..872771af9ab 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/utils.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/utils.py @@ -145,8 +145,8 @@ def find_matched_target( matched_target = ( _find_first_match(layer_name, targets) - or _find_first_match(module.__class__.__name__, targets, True) or _match_fused_layer(layer_name, targets, fused_mapping) + or _find_first_match(module.__class__.__name__, targets, True) ) return matched_target From 1db989bbf14b9896c1306fd722ced9a9b1238465 Mon Sep 17 00:00:00 2001 From: labAxiaoming <34019940+labAxiaoming@users.noreply.github.com> Date: Wed, 29 Jul 2026 01:57:31 +0800 Subject: [PATCH 184/185] [Bugfix][Multimodal] Fix video temporal padding estimates (#49030) Signed-off-by: xiaoming <1259730330@qq.com> --- .../multimodal/processing/test_glm4_1v.py | 25 +++++++++++++++++++ vllm/model_executor/models/glm4_1v.py | 4 +-- vllm/model_executor/models/kanana_v.py | 4 +-- vllm/model_executor/models/keye.py | 2 +- .../model_executor/models/llava_onevision2.py | 2 +- vllm/model_executor/models/mimo_v2_omni.py | 2 +- vllm/model_executor/models/qwen2_vl.py | 4 +-- 7 files changed, 34 insertions(+), 9 deletions(-) diff --git a/tests/models/multimodal/processing/test_glm4_1v.py b/tests/models/multimodal/processing/test_glm4_1v.py index 8a777832826..e45e741c5b4 100644 --- a/tests/models/multimodal/processing/test_glm4_1v.py +++ b/tests/models/multimodal/processing/test_glm4_1v.py @@ -63,6 +63,31 @@ def test_encoder_cudagraph_uses_model_video_frame_limit(): assert Glm4vForConditionalGeneration.get_max_frames_per_video(model) == 600 +@pytest.mark.parametrize( + ("temporal_patch_size", "expected_grid_t"), + [(2, 9), (4, 5), (8, 3)], +) +def test_vision_info_rounds_up_temporal_frames( + temporal_patch_size: int, + expected_grid_t: int, +): + info = Mock(spec=Glm4vProcessingInfo) + vision_config = info.get_hf_config.return_value.vision_config + vision_config.patch_size = 14 + vision_config.spatial_merge_size = 2 + vision_config.temporal_patch_size = temporal_patch_size + + _, num_vision_tokens = Glm4vProcessingInfo._get_vision_info( + info, + image_width=28, + image_height=28, + num_frames=17, + do_resize=False, + ) + + assert num_vision_tokens == expected_grid_t + + @pytest.mark.parametrize("model_id", ["zai-org/GLM-4.1V-9B-Thinking"]) @pytest.mark.parametrize("expected_toks_per_frame", [299]) @pytest.mark.parametrize( diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py index 787e53e8df2..19389d84463 100644 --- a/vllm/model_executor/models/glm4_1v.py +++ b/vllm/model_executor/models/glm4_1v.py @@ -1082,8 +1082,8 @@ class Glm4vProcessingInfo(BaseProcessingInfo): preprocessed_size = ImageSize(width=image_width, height=image_height) # NOTE: Frames are padded to be divisible by `temporal_patch_size` - # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/qwen2_vl/image_processing_qwen2_vl.py#L294 - padded_num_frames = num_frames + num_frames % temporal_patch_size + # https://github.com/huggingface/transformers/blob/v5.13.0/src/transformers/models/qwen2_vl/video_processing_qwen2_vl.py#L249-L252 + padded_num_frames = num_frames + (-num_frames % temporal_patch_size) grid_t = max(padded_num_frames // temporal_patch_size, 1) grid_h = preprocessed_size.height // patch_size diff --git a/vllm/model_executor/models/kanana_v.py b/vllm/model_executor/models/kanana_v.py index 125d7e71c7b..b1a5f78b1b3 100644 --- a/vllm/model_executor/models/kanana_v.py +++ b/vllm/model_executor/models/kanana_v.py @@ -409,8 +409,8 @@ class KananaVProcessingInfo(BaseProcessingInfo): preprocessed_size = ImageSize(width=image_width, height=image_height) # NOTE: Frames are padded to be divisible by `temporal_patch_size` - # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/qwen2_vl/image_processing_qwen2_vl.py#L294 - padded_num_frames = num_frames + num_frames % temporal_patch_size + # https://github.com/huggingface/transformers/blob/v5.13.0/src/transformers/models/qwen2_vl/video_processing_qwen2_vl.py#L249-L252 + padded_num_frames = num_frames + (-num_frames % temporal_patch_size) grid_t = max(padded_num_frames // temporal_patch_size, 1) grid_h = preprocessed_size.height // patch_size diff --git a/vllm/model_executor/models/keye.py b/vllm/model_executor/models/keye.py index dd1fb892ad1..c3d69836a79 100644 --- a/vllm/model_executor/models/keye.py +++ b/vllm/model_executor/models/keye.py @@ -983,7 +983,7 @@ class KeyeProcessingInfo(BaseProcessingInfo): else: preprocessed_size = ImageSize(width=image_width, height=image_height) - padded_num_frames = num_frames + num_frames % temporal_patch_size + padded_num_frames = num_frames + (-num_frames % temporal_patch_size) grid_t = max(padded_num_frames // temporal_patch_size, 1) grid_h = preprocessed_size.height // patch_size diff --git a/vllm/model_executor/models/llava_onevision2.py b/vllm/model_executor/models/llava_onevision2.py index 58179ec00de..9dc11d493a0 100644 --- a/vllm/model_executor/models/llava_onevision2.py +++ b/vllm/model_executor/models/llava_onevision2.py @@ -1345,7 +1345,7 @@ class LlavaOnevision2ProcessingInfo(BaseProcessingInfo): preprocessed = ImageSize(width=rw, height=rh) else: preprocessed = ImageSize(width=image_width, height=image_height) - padded_frames = num_frames + num_frames % temporal_patch_size + padded_frames = num_frames + (-num_frames % temporal_patch_size) grid_t = max(padded_frames // temporal_patch_size, 1) grid_h = preprocessed.height // patch_size grid_w = preprocessed.width // patch_size diff --git a/vllm/model_executor/models/mimo_v2_omni.py b/vllm/model_executor/models/mimo_v2_omni.py index d0d9589ae1d..747cb0e88b2 100644 --- a/vllm/model_executor/models/mimo_v2_omni.py +++ b/vllm/model_executor/models/mimo_v2_omni.py @@ -715,7 +715,7 @@ class MiMoV2OmniProcessingInfo(BaseProcessingInfo): effective_frames = num_frames * tokens_per_second else: effective_frames = num_frames - padded_num_frames = effective_frames + effective_frames % temporal_patch_size + padded_num_frames = effective_frames + (-effective_frames % temporal_patch_size) grid_t = max(padded_num_frames // temporal_patch_size, 1) grid_h = preprocessed_size.height // patch_size grid_w = preprocessed_size.width // patch_size diff --git a/vllm/model_executor/models/qwen2_vl.py b/vllm/model_executor/models/qwen2_vl.py index 539f141cbaa..e2e9f245248 100644 --- a/vllm/model_executor/models/qwen2_vl.py +++ b/vllm/model_executor/models/qwen2_vl.py @@ -898,8 +898,8 @@ class Qwen2VLProcessingInfo(BaseProcessingInfo): preprocessed_size = ImageSize(width=image_width, height=image_height) # NOTE: Frames are padded to be divisible by `temporal_patch_size` - # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/qwen2_vl/image_processing_qwen2_vl.py#L294 - padded_num_frames = num_frames + num_frames % temporal_patch_size + # https://github.com/huggingface/transformers/blob/v5.13.0/src/transformers/models/qwen2_vl/video_processing_qwen2_vl.py#L249-L252 + padded_num_frames = num_frames + (-num_frames % temporal_patch_size) grid_t = max(padded_num_frames // temporal_patch_size, 1) grid_h = preprocessed_size.height // patch_size From 6c7e679f048dc6123caecc3985150766e455ff22 Mon Sep 17 00:00:00 2001 From: Shanshan Shen <467638484@qq.com> Date: Wed, 29 Jul 2026 02:08:50 +0800 Subject: [PATCH 185/185] [ROCm][Bugfix] Sanitize AITER paged-MQA logits before sparse top-k for DeepSeek-V4 (#49714) Signed-off-by: shen-shanshan <467638484@qq.com> --- .../attention/test_rocm_triton_attn_dsv4.py | 64 +++++++++++++++++++ .../v1/attention/ops/rocm_aiter_mla_sparse.py | 1 + 2 files changed, 65 insertions(+) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 77e068ab171..6fe2a3e7758 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest import torch @@ -196,6 +198,68 @@ def _ragged_from_rows( ) +@torch.inference_mode() +def test_paged_mqa_logits_do_not_contain_nan(monkeypatch) -> None: + from vllm._aiter_ops import rocm_aiter_ops + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + device = torch.device("cuda") + + class FakeWorkspaceManager: + def get_simultaneous(self, *shapes_and_dtypes): + return [ + torch.empty(shape, dtype=dtype, device=device) + for shape, dtype in shapes_and_dtypes + ] + + def fake_paged_mqa_logits( + q_fp8, + kv_cache_fp8, + weights, + out_logits, + context_lens, + block_tables, + max_seq_len, + **kwargs, + ): + del ( + q_fp8, + kv_cache_fp8, + weights, + context_lens, + block_tables, + max_seq_len, + kwargs, + ) + out_logits.fill_(float("nan")) + + monkeypatch.setattr(mod, "_ON_GFX942", False) + monkeypatch.setattr(mod, "_ON_GFX950", True) + monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True) + monkeypatch.setattr( + mod, + "paged_mqa_logits_module", + lambda: SimpleNamespace(deepgemm_fp8_paged_mqa_logits=fake_paged_mqa_logits), + ) + monkeypatch.setattr( + mod, "current_workspace_manager", lambda: FakeWorkspaceManager() + ) + + q_fp8 = torch.empty((1, 1, 1, 1), dtype=torch.uint8, device=device) + kv_cache_fp8 = torch.empty((1, 1, 1, 5), dtype=torch.uint8, device=device) + logits = mod.rocm_fp8_paged_mqa_logits( + q_fp8, + kv_cache_fp8, + torch.empty((1, 1), dtype=torch.float32, device=device), + torch.ones(1, dtype=torch.int32, device=device), + torch.zeros((1, 1), dtype=torch.int32, device=device), + torch.empty(0, dtype=torch.int32, device=device), + 1, + ) + + assert not torch.isnan(logits).any() + + @torch.inference_mode() def test_compute_global_topk_ragged_indices_and_indptr() -> None: from vllm.models.deepseek_v4.amd.rocm import ( diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 6a28324208f..63bad3cfad0 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -442,6 +442,7 @@ def rocm_fp8_paged_mqa_logits( KVBlockSize=block_size, WavePerEU=2, ) + out_logits.nan_to_num_(float("-inf")) return out_logits deepgemm_fp8_paged_mqa_logits_stage1 = ( aiter_paged_mqa_logits_module.deepgemm_fp8_paged_mqa_logits_stage1