From cac3e70cd4f522eaab9d28323285d5e31238f907 Mon Sep 17 00:00:00 2001 From: ap9272 Date: Thu, 9 Jul 2026 16:46:22 -0700 Subject: [PATCH 0001/1526] Correct model layer aliasing for Bert style models (#43896) --- vllm/model_executor/models/modernbert.py | 8 +++++++- vllm/model_executor/warmup/deep_gemm_warmup.py | 5 ++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/models/modernbert.py b/vllm/model_executor/models/modernbert.py index 8195f61b05f2..d182fa071594 100644 --- a/vllm/model_executor/models/modernbert.py +++ b/vllm/model_executor/models/modernbert.py @@ -237,7 +237,11 @@ def forward( @default_pooling_type(seq_pooling_type="CLS") class ModernBertModel(nn.Module): hf_to_vllm_mapper = WeightsMapper( - orig_to_new_prefix={"layers.": "encoder_layer.layers."} + orig_to_new_prefix={ + "model.layers.": "encoder_layer.layers.", + "layers.": "encoder_layer.layers.", + "model.": "", + } ) def __init__( @@ -266,6 +270,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: for name, loaded_weight in weights: if name.endswith(".bias") and name not in params_dict: continue + if name not in params_dict: + continue param = params_dict[name] weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) diff --git a/vllm/model_executor/warmup/deep_gemm_warmup.py b/vllm/model_executor/warmup/deep_gemm_warmup.py index cfff491ab2bd..2c9182e86194 100644 --- a/vllm/model_executor/warmup/deep_gemm_warmup.py +++ b/vllm/model_executor/warmup/deep_gemm_warmup.py @@ -138,9 +138,6 @@ def _fp8_linear_may_use_deep_gemm(module: torch.nn.Module) -> bool: Return True if the input module/layer could be processed with DeepGEMM. """ - # FIXME: this logic is brittle and incorrect - since we - # could use DeepGEMM with for than just Fp8LinearMethod - block_size = get_mk_alignment_for_contiguous_layout()[0] if not ( isinstance(module, LinearBase) and isinstance(module.quant_method, Fp8LinearMethod) @@ -156,6 +153,8 @@ def _fp8_linear_may_use_deep_gemm(module: torch.nn.Module) -> bool: ): return False + block_size = get_mk_alignment_for_contiguous_layout()[0] + w, _, block_sizes = _extract_data_from_linear_base_module(module) return ( block_sizes == get_mk_alignment_for_contiguous_layout() From f1a5adddb815610db46cae81f2a1d5a4609bf99d Mon Sep 17 00:00:00 2001 From: gnovack Date: Thu, 9 Jul 2026 16:52:52 -0700 Subject: [PATCH 0002/1526] update marlin M size for EP (#48144) Signed-off-by: gnovack --- .../layers/fused_moe/experts/marlin_moe.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py index 867f71b9bf64..17c166961397 100644 --- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Fused MoE utilities for GPTQ.""" +import math from collections.abc import Callable import torch @@ -312,6 +313,12 @@ def fused_marlin_moe( assert num_bits in [4, 8] assert topk_weights.dtype == torch.float32 + if global_num_experts == -1: + global_num_experts = E + else: + # Set M to estimated valid tokens per rank + M = math.ceil(M * E / global_num_experts) + # M block size selection logic # TODO: tune this further for specific models for block_size_m in [8, 16, 32, 48, 64]: @@ -321,8 +328,6 @@ def fused_marlin_moe( if input_dtype is not None and input_dtype.itemsize == 1: block_size_m = max(block_size_m, 16) - if global_num_experts == -1: - global_num_experts = E sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( topk_ids, block_size_m, From a0f6d767e42acf09f94e7af2bacca7f4c268c264 Mon Sep 17 00:00:00 2001 From: peizhang56 Date: Thu, 9 Jul 2026 17:20:15 -0700 Subject: [PATCH 0003/1526] [ROCm][CI] Move remaining engine/samplers AMD steps to mi325_1 (#48169) Signed-off-by: pei.zhang Co-authored-by: Claude --- .buildkite/test_areas/engine.yaml | 4 ++-- .buildkite/test_areas/samplers.yaml | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index 9edd9343dedb..7bb605465720 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -60,7 +60,7 @@ steps: - pytest -v -s v1/e2e/general/test_async_scheduling.py mirror: amd: - device: mi250_1 + device: mi325_1 timeout_in_minutes: 60 depends_on: - image-build-amd @@ -76,7 +76,7 @@ steps: - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py mirror: amd: - device: mi250_1 + device: mi325_1 timeout_in_minutes: 35 depends_on: - image-build-amd diff --git a/.buildkite/test_areas/samplers.yaml b/.buildkite/test_areas/samplers.yaml index 6ec6f8efd351..5abc16889434 100644 --- a/.buildkite/test_areas/samplers.yaml +++ b/.buildkite/test_areas/samplers.yaml @@ -19,7 +19,7 @@ steps: - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers mirror: amd: - device: mi250_1 + device: mi325_1 depends_on: - image-build-amd commands: From feb384ada2a1da1981ccb713ea6bd3f90c44f7b4 Mon Sep 17 00:00:00 2001 From: Augusto Yao Date: Fri, 10 Jul 2026 10:03:00 +0800 Subject: [PATCH 0004/1526] [bugfix] bge-m3-sparse-plugin mismatch requests (#48112) Signed-off-by: augusto.yjh --- .../sparse_embeddings_processor.py | 52 ++----------------- 1 file changed, 4 insertions(+), 48 deletions(-) diff --git a/tests/plugins/bge_m3_sparse_plugin/bge_m3_sparse_processor/sparse_embeddings_processor.py b/tests/plugins/bge_m3_sparse_plugin/bge_m3_sparse_processor/sparse_embeddings_processor.py index 77f8884fef23..a845dae77b48 100644 --- a/tests/plugins/bge_m3_sparse_plugin/bge_m3_sparse_processor/sparse_embeddings_processor.py +++ b/tests/plugins/bge_m3_sparse_plugin/bge_m3_sparse_processor/sparse_embeddings_processor.py @@ -3,9 +3,8 @@ from collections.abc import Sequence -from vllm.config import ModelConfig, PoolerConfig, VllmConfig +from vllm.config import PoolerConfig, VllmConfig from vllm.entrypoints.openai.engine.protocol import UsageInfo -from vllm.entrypoints.pooling.base.protocol import EmbedRequestMixin from vllm.inputs import PromptType from vllm.outputs import PoolingRequestOutput from vllm.plugins.io_processors.interface import IOProcessor @@ -14,7 +13,6 @@ from vllm.tokenizers.detokenizer_utils import convert_ids_list_to_tokens from .types import ( - EMBED_TASKS, SparseEmbeddingCompletionRequestMixin, SparseEmbeddingResponse, SparseEmbeddingResponseData, @@ -38,7 +36,6 @@ def __init__(self, vllm_config: VllmConfig, renderer: BaseRenderer): continue self.default_pooling_params[param] = getattr(pooler_config, param) self.embed_dimensions = vllm_config.model_config.embedding_size - self.embed_request_queue: list[EmbedRequestMixin] = [] def __repr__(self) -> str: return ( @@ -56,44 +53,9 @@ def merge_pooling_params( # refer to PoolingCompletionRequest.to_pooling_params # set and verify pooling params params.skip_reading_prefix_cache = True - - raw_embed_request = self.embed_request_queue.pop(0) - if raw_embed_request.embed_task not in EMBED_TASKS: - raise ValueError( - f"Unsupported task {raw_embed_request}, " - f"Supported tasks are {EMBED_TASKS}" - ) params.task = "embed&token_classify" - params.use_activation = raw_embed_request.use_activation - if params.use_activation is None: - params.use_activation = True - - params.dimensions = raw_embed_request.dimensions - - model_config: ModelConfig = self.vllm_config.model_config - for param in self.default_pooling_params: - if getattr(params, param, None) is None: - setattr(params, param, self.default_pooling_params[param]) - - if params.dimensions is not None: - if not model_config.is_matryoshka: - raise ValueError( - f'Model "{model_config.served_model_name}" does not ' - f"support matryoshka representation, " - f"changing output dimensions will lead to poor results." - ) - - mds = model_config.matryoshka_dimensions - if mds is not None: - if params.dimensions not in mds: - raise ValueError( - f"Model {model_config.served_model_name!r} " - f"only supports {str(mds)} matryoshka dimensions, " - f"use other output dimensions will " - f"lead to poor results." - ) - elif params.dimensions < 1: - raise ValueError("Dimensions must be greater than 0") + params.use_activation = True + params.dimensions = self.embed_dimensions return params def parse_request( @@ -113,10 +75,8 @@ def pre_process( if request_id is not None: assert request_id not in self.online_requests, "request_id duplicated" self.online_requests[request_id] = prompt - self.embed_request_queue.extend(prompt.to_embed_requests_online()) else: self.offline_requests.append(prompt) - self.embed_request_queue.extend(prompt.to_embed_requests_offline()) return prompt.input def _get_sparse_embedding_request(self, request_id: str | None = None): @@ -157,11 +117,7 @@ def post_process( raw_request = self._get_sparse_embedding_request(request_id) has_dense_embed = raw_request.embed_task in ["dense", "dense&sparse"] has_sparse_embed = raw_request.embed_task in ["sparse", "dense&sparse"] - embed_dimensions = ( - self.embed_dimensions - if raw_request.dimensions is None - else raw_request.dimensions - ) + embed_dimensions = self.embed_dimensions for idx in range(len(model_output)): mo = model_output[idx] sparse_embedding_dict: dict[int, float] = {} From 88e5e2c57be8ce6e25510c1249352a23b8a85ec4 Mon Sep 17 00:00:00 2001 From: peizhang56 Date: Thu, 9 Jul 2026 19:14:38 -0700 Subject: [PATCH 0005/1526] [CI/Build][AMD] Fix ROCm OOM in eagle_correctness_heavy by reserving CUDA graph memory (#47366) Signed-off-by: pei.zhang Co-authored-by: Claude Co-authored-by: Andreas Karatzas --- tests/v1/e2e/spec_decode/test_spec_decode.py | 7 +++++++ vllm/v1/worker/gpu_worker.py | 12 +++++++----- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/tests/v1/e2e/spec_decode/test_spec_decode.py b/tests/v1/e2e/spec_decode/test_spec_decode.py index 2eff02ea6ed6..96d6684594cc 100644 --- a/tests/v1/e2e/spec_decode/test_spec_decode.py +++ b/tests/v1/e2e/spec_decode/test_spec_decode.py @@ -17,6 +17,7 @@ multi_gpu_marks, multi_gpu_only, single_gpu_only, + wait_for_rocm_memory_to_settle, ) from vllm import LLM, SamplingParams from vllm.assets.base import VLLM_S3_BUCKET_URL @@ -448,6 +449,9 @@ def _run_eagle_correctness( del ref_llm torch.accelerator.empty_cache() cleanup_dist_env_and_memory() + # ROCm frees VRAM lazily; wait so the spec engine started right after + # does not OOM on its startup memory guard. + wait_for_rocm_memory_to_settle() spec_llm = LLM( model=model_name, @@ -485,6 +489,9 @@ def _run_eagle_correctness( del spec_llm torch.accelerator.empty_cache() cleanup_dist_env_and_memory() + # ROCm frees VRAM lazily; wait so the next parametrization's engine does + # not OOM on its startup memory guard. + wait_for_rocm_memory_to_settle() @single_gpu_only diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 182476e25336..1cd7e2598a03 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -478,11 +478,14 @@ def determine_available_memory(self) -> int: ) # Profile CUDA graph memory if graphs will be captured. - # Skip on ROCm/HIP/XPU as graph pool handles and get_memory_info - # behave differently and can produce incorrect/negative estimates. + # ROCm is included: #44825 moved the profiler to + # torch.accelerator.get_memory_info (reliable on ROCm, as used by + # the AMD-CI mem tests), and graph_pool_handle resolves to the same + # torch.cuda handle the live capture path already uses on ROCm. + # XPU stays excluded (see #39977). cudagraph_memory_estimate = 0 if ( - current_platform.is_cuda() + current_platform.is_cuda_alike() and self.vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE ): @@ -498,8 +501,7 @@ def determine_available_memory(self) -> int: + profile_result.weights_memory ) - # On ROCm, cudagraph_memory_estimate is always 0 so this is a no-op. - # On CUDA, respect the opt-in flag as originally designed. + # Respect the opt-in flag as originally designed. cudagraph_memory_estimate_applied = ( cudagraph_memory_estimate if envs.VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS From 2d814a00820daec7082599bea75ae1d0959a346c Mon Sep 17 00:00:00 2001 From: Chang Guo Date: Thu, 9 Jul 2026 20:17:23 -0700 Subject: [PATCH 0006/1526] [kv_offload] Emit tier-owned BlockStored events from FS/OBJ secondary tiers (#47923) Signed-off-by: Change72 Co-authored-by: Claude Fable 5 Co-authored-by: Or Ozeri --- docs/features/kv_offloading_usage.md | 28 +++ tests/v1/kv_offload/tiering/test_fs_tier.py | 196 +++++++++++++++++++ tests/v1/kv_offload/tiering/test_obj_tier.py | 108 +++++++++- vllm/distributed/kv_events.py | 2 + vllm/v1/kv_offload/tiering/fs/manager.py | 53 ++++- vllm/v1/kv_offload/tiering/obj/manager.py | 55 +++++- 6 files changed, 433 insertions(+), 9 deletions(-) diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index cff65753d99a..8a48d11be73a 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -81,6 +81,8 @@ vllm serve \ Each entry in `secondary_tiers` is a dict with a required `type` field plus tier-specific fields. +The filesystem and object-store tiers can publish hash-only `BlockStored` KV events for blocks they successfully store, tagged with a stable per-tier `medium` (`FS` for the filesystem tier, `OBJ` for the object-store tier). Set `enable_kv_events: true` in the tier's entry to opt in; events are published only when KV cache events are also enabled globally via `--kv-events-config`. + ### Filesystem (FS) The filesystem tier (`type: "fs"`) writes blocks to a directory on local storage. @@ -91,6 +93,7 @@ The filesystem tier (`type: "fs"`) writes blocks to a directory on local storage | `root_dir` | yes | — | Base directory; vLLM creates subdirectories beneath it (see [On-Disk Layout](#on-disk-layout)). | | `n_read_threads` | no | `16` | Read-priority I/O threads (load path). | | `n_write_threads` | no | `16` | Write-priority I/O threads (store path). | +| `enable_kv_events` | no | `false` | Publish `BlockStored` KV events (medium `FS`) for successfully stored blocks. Requires KV cache events to be enabled globally. | Each thread group prefers its own queue but pulls from the other when its primary queue is empty, so a write-heavy or read-heavy burst won't leave the off-priority queue waiting. Size the totals to your storage's effective concurrency. @@ -120,6 +123,31 @@ To enable KV cache sharing between multiple vLLM instances using the same `root_ PYTHONHASHSEED=0 vllm serve ... ``` +### Object Store (OBJ) + +The object-store tier (`type: "obj"`) offloads blocks to an S3-compatible object store through the NIXL OBJ backend. + +| Key | Required | Default | Notes | +| --- | --- | --- | --- | +| `type` | yes | — | Must be `obj`. | +| `store_config` | yes | — | Object store connection parameters (see below). | +| `prefix` | no | `""` | Key prefix prepended to all object keys. | +| `io_threads` | no | `4` | Number of NIXL OBJ backend I/O threads. | +| `enable_kv_events` | no | `false` | Publish `BlockStored` KV events (medium `OBJ`) for successfully stored blocks. Requires KV cache events to be enabled globally. | + +`store_config` fields: + +| Key | Required | Default | Notes | +| --- | --- | --- | --- | +| `bucket` | yes | — | Bucket name. | +| `endpoint_override` | yes | — | Object store endpoint host; the URL scheme is set separately via `scheme`. | +| `scheme` | no | `http` | `http` or `https`. | +| `access_key`, `secret_key`, `session_token` | no | `""` | Explicit credentials. When left empty, the NIXL OBJ plugin falls back to the AWS SDK default credential provider chain (IAM roles, environment variables, credential files), which enables workload-identity auth on Kubernetes. | +| `region` | no | `""` | Bucket region, if the endpoint requires one. | +| `ca_bundle` | no | `""` | CA bundle path for TLS verification. | + +Object keys follow the same run-configuration digest scheme as the filesystem tier (see [On-Disk Layout](#on-disk-layout)) and are stored under the optional `prefix`. The [Cross-Process Sharing](#cross-process-sharing) requirement (`PYTHONHASHSEED`) applies to shared buckets as well, so instances sharing a bucket produce identical keys for identical content. At startup the tier probes object store connectivity and fails fast with a configuration error if the bucket is unreachable. + ### P2P (Including P/D) The P2P tier (`type: "p2p"`) shares completed KV blocks between vLLM instances over RDMA via NIXL. Each instance binds a control socket on `host:port` and exchanges blocks directly with peers — no shared filesystem required. diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index 680a9584787a..c90a4dad5e79 100644 --- a/tests/v1/kv_offload/tiering/test_fs_tier.py +++ b/tests/v1/kv_offload/tiering/test_fs_tier.py @@ -18,8 +18,10 @@ import pytest import torch +from vllm.distributed.kv_events import MEDIUM_FS from vllm.v1.kv_offload.base import ( LookupResult, + OffloadingEvent, OffloadKey, ReqContext, ScheduleEndContext, @@ -58,6 +60,16 @@ _MOCK_OFFLOADING_SPEC.block_size_factor = 1 +def _make_offloading_spec(enable_kv_cache_events: bool) -> MagicMock: + """Mock spec with an explicit global KV events flag.""" + spec = MagicMock() + spec.vllm_config = _MOCK_VLLM_CONFIG + spec.kv_cache_config = _MOCK_KV_CACHE_CONFIG + spec.block_size_factor = 1 + spec.kv_events_config.enable_kv_cache_events = enable_kv_cache_events + return spec + + def key(n: int) -> OffloadKey: return make_offload_key(n.to_bytes(8, "big"), 0) @@ -164,6 +176,23 @@ def fs_tier(tmp_path): tier.shutdown() +@pytest.fixture +def fs_tier_with_events(tmp_path): + tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS) + mock_view = memoryview(tensor.numpy()) + tier = FileSystemTierManager( + offloading_spec=_make_offloading_spec(enable_kv_cache_events=True), + primary_kv_view=mock_view, + tier_type="fs", + root_dir=str(tmp_path), + n_read_threads=4, + n_write_threads=4, + enable_kv_events=True, + ) + yield tier + tier.shutdown() + + # --------------------------------------------------------------------------- # Tests # --------------------------------------------------------------------------- @@ -396,3 +425,170 @@ def test_batch_lookup_dispatch(fs_tier, monkeypatch, use_c_ext): results = lookup_and_wait(tier, [key(1), key(2)]) assert results == [LookupResult.HIT, LookupResult.MISS] + + +# --------------------------------------------------------------------------- +# KV events +# --------------------------------------------------------------------------- + + +def test_successful_store_emits_stored_event(fs_tier_with_events): + """A completed store job emits one stored event with the job's keys.""" + tier = fs_tier_with_events + keys = [key(1), key(2)] + tier.submit_store(make_job(1, keys, [0, 1])) + assert all(r.success for r in drain(tier)) + + 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 not events[0].removed + # take_events drains the buffer. + assert list(tier.take_events()) == [] + + +def test_load_job_emits_no_event(fs_tier_with_events): + tier = fs_tier_with_events + tier.submit_store(make_job(1, [key(1)], [0])) + results = drain(tier) + assert len(results) == 1 + assert results[0].success + list(tier.take_events()) + + tier.submit_load(make_job(2, [key(1)], [1], is_promotion=True)) + results = drain(tier) + assert len(results) == 1 + assert results[0].success + assert list(tier.take_events()) == [] + + +def test_mixed_job_results_emit_event_only_for_successful_job( + fs_tier_with_events, monkeypatch +): + """With a failed and a successful store job in flight, exactly one event + is emitted and its keys belong to the successful job.""" + import vllm.v1.kv_offload.tiering.fs.manager as mgr_mod + + tier = fs_tier_with_events + failing_path = tier.file_mapper.get_file_name(key(1)) + original_store_block = mgr_mod.store_block + + def flaky_store_block(dest_path, *args, **kwargs): + if dest_path == failing_path: + raise OSError("injected store failure") + return original_store_block(dest_path, *args, **kwargs) + + monkeypatch.setattr(mgr_mod, "store_block", flaky_store_block) + + tier.submit_store(make_job(1, [key(1)], [0])) + tier.submit_store(make_job(2, [key(2)], [1])) + results = drain(tier) + assert len(results) == 2 + by_id = {r.job_id: r for r in results} + assert not by_id[1].success + assert by_id[2].success + + events = list(tier.take_events()) + assert len(events) == 1 + assert events[0].keys == [key(2)] + + +def test_partially_failed_store_emits_no_event(fs_tier_with_events, monkeypatch): + """A store job with any failed block emits no event for the whole job.""" + import vllm.v1.kv_offload.tiering.fs.manager as mgr_mod + + tier = fs_tier_with_events + failing_path = tier.file_mapper.get_file_name(key(2)) + original_store_block = mgr_mod.store_block + + def flaky_store_block(dest_path, *args, **kwargs): + if dest_path == failing_path: + raise OSError("injected store failure") + return original_store_block(dest_path, *args, **kwargs) + + monkeypatch.setattr(mgr_mod, "store_block", flaky_store_block) + + tier.submit_store(make_job(1, [key(1), key(2)], [0, 1])) + results = drain(tier) + assert len(results) == 1 + assert not results[0].success + assert list(tier.take_events()) == [] + assert tier._store_job_keys == {} + + +def test_events_disabled_by_default(fs_tier): + tier, _ = fs_tier + tier.submit_store(make_job(1, [key(1)], [0])) + results = drain(tier) + assert len(results) == 1 + assert results[0].success + assert tier.events is None + assert tier._store_job_keys == {} + assert list(tier.take_events()) == [] + + +def test_events_require_global_kv_events_flag(tmp_path): + """Tier-level opt-in alone is not enough; the global flag gates events.""" + tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS) + tier = FileSystemTierManager( + offloading_spec=_make_offloading_spec(enable_kv_cache_events=False), + primary_kv_view=memoryview(tensor.numpy()), + tier_type="fs", + root_dir=str(tmp_path), + enable_kv_events=True, + ) + try: + assert tier.events is None + tier.submit_store(make_job(1, [key(1)], [0])) + results = drain(tier) + assert len(results) == 1 + assert results[0].success + assert list(tier.take_events()) == [] + assert tier._store_job_keys == {} + finally: + tier.shutdown() + + +def test_cascade_store_emits_fs_event_through_tiering_manager(tmp_path): + """A GPU->CPU->fs cascade surfaces the tier-owned FS stored event via the + TieringOffloadingManager's aggregated take_events().""" + from vllm.v1.kv_offload.tiering.manager import ( + CPUPrimaryTierOffloadingManager, + TieringOffloadingManager, + ) + + tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS) + view = memoryview(tensor.numpy()) + mock_region = MagicMock() + mock_region.create_kv_memoryview.return_value = view + primary = CPUPrimaryTierOffloadingManager(num_blocks=4, mmap_region=mock_region) + tier = FileSystemTierManager( + offloading_spec=_make_offloading_spec(enable_kv_cache_events=True), + primary_kv_view=primary.get_kv_memoryview(), + tier_type="fs", + root_dir=str(tmp_path), + enable_kv_events=True, + ) + manager = TieringOffloadingManager(primary_tier=primary, secondary_tiers=[tier]) + try: + keys = [key(1), key(2)] + manager.on_new_request(_CTX) + assert manager.prepare_store(keys, _CTX) is not None + manager.complete_store(keys, _CTX) # cascades to the fs tier + + events: list[OffloadingEvent] = [] + ctx = ScheduleEndContext(new_req_ids=[], preempted_req_ids=()) + deadline = time.monotonic() + 5.0 + while time.monotonic() < deadline and not events: + manager.on_schedule_end(ctx) + events.extend(manager.take_events()) + time.sleep(0.01) + + fs_events = [e for e in events if e.medium == MEDIUM_FS] + assert len(fs_events) == 1 + assert set(fs_events[0].keys) == set(keys) + assert not fs_events[0].removed + finally: + tier.shutdown() diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index f429d6cf62d6..37687adce0cb 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -179,8 +179,19 @@ def _query_memory(self, queries, mem_type, agent_name): # --------------------------------------------------------------------------- +def _make_events_spec(enable_kv_cache_events: bool) -> SimpleNamespace: + """Offloading spec stub with an explicit global KV events flag.""" + return SimpleNamespace( + vllm_config=_make_vllm_config(), + kv_cache_config=SimpleNamespace(kv_cache_groups=[]), + kv_events_config=SimpleNamespace(enable_kv_cache_events=enable_kv_cache_events), + ) + + def _make_tier( num_blocks: int = 4, + offloading_spec: SimpleNamespace = _OFFLOADING_SPEC, + **tier_kwargs, ) -> tuple[ObjectStoreSecondaryTierManager, MockNixlAgent]: """Create a tier backed by a fresh MockNixlAgent.""" mock_agent = MockNixlAgent() @@ -194,11 +205,12 @@ def _make_tier( ), ): tier = ObjectStoreSecondaryTierManager( - offloading_spec=_OFFLOADING_SPEC, + offloading_spec=offloading_spec, primary_kv_view=view, tier_type="obj", store_config=_STORE_CONFIG, prefix=_RUN_PREFIX, + **tier_kwargs, ) return tier, mock_agent @@ -421,6 +433,100 @@ def test_shutdown_idempotent(self): tier.shutdown() # must not raise +class TestObjTierKVEvents: + def setup_method(self): + self.tier, self.agent = _make_tier( + offloading_spec=_make_events_spec(enable_kv_cache_events=True), + enable_kv_events=True, + ) + + def test_successful_store_emits_stored_event(self): + """A completed store transfer emits one stored event with the job's keys.""" + keys = [key(1), key(2)] + self.tier.submit_store(make_job(1, keys, [0, 1])) + assert all(r.success for r in drain(self.tier)) + + 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 not events[0].removed + # take_events drains the buffer. + assert list(self.tier.take_events()) == [] + + def test_mixed_job_results_emit_event_only_for_successful_job(self): + """With a failed and a successful store job resolving in the same + poll, exactly one event is emitted and its keys belong to the + successful job.""" + original = self.agent.check_xfer_state + self.agent.check_xfer_state = lambda h: "ERR" if h._id == 0 else original(h) + self.tier.submit_store(make_job(1, [key(1)], [0])) # handle 0: fails + self.tier.submit_store(make_job(2, [key(2)], [1])) # handle 1: succeeds + results = drain(self.tier) + by_id = {r.job_id: r for r in results} + assert not by_id[1].success + assert by_id[2].success + + events = list(self.tier.take_events()) + assert len(events) == 1 + assert events[0].keys == [key(2)] + assert self.tier._store_job_keys == {} + + def test_load_job_emits_no_event(self): + self.tier.submit_store(make_job(1, [key(1)], [0])) + results = drain(self.tier) + assert len(results) == 1 + assert results[0].success + list(self.tier.take_events()) + + self.tier.submit_load(make_job(2, [key(1)], [0])) + results = drain(self.tier) + assert len(results) == 1 + assert results[0].success + assert list(self.tier.take_events()) == [] + + def test_failed_transfer_emits_no_event(self): + self.agent.check_xfer_state = lambda h: "ERR" + self.tier.submit_store(make_job(1, [key(1)], [0])) + results = drain(self.tier) + assert not results[0].success + assert list(self.tier.take_events()) == [] + assert self.tier._store_job_keys == {} + + def test_submission_failure_emits_no_event(self): + self.agent.make_prepped_xfer = lambda *a, **k: None + self.tier.submit_store(make_job(1, [key(1)], [0])) + results = list(self.tier.get_finished_jobs()) + assert not results[0].success + assert list(self.tier.take_events()) == [] + assert self.tier._store_job_keys == {} + + def test_events_disabled_by_default(self): + tier, _ = _make_tier() + tier.submit_store(make_job(1, [key(1)], [0])) + results = drain(tier) + assert len(results) == 1 + assert results[0].success + assert tier.events is None + assert tier._store_job_keys == {} + assert list(tier.take_events()) == [] + + def test_events_require_global_kv_events_flag(self): + """Tier-level opt-in alone is not enough; the global flag gates events.""" + tier, _ = _make_tier( + offloading_spec=_make_events_spec(enable_kv_cache_events=False), + enable_kv_events=True, + ) + tier.submit_store(make_job(1, [key(1)], [0])) + results = drain(tier) + assert len(results) == 1 + assert results[0].success + assert tier.events is None + assert tier._store_job_keys == {} + assert list(tier.take_events()) == [] + + class TestObjStoreConfig: def test_explicit_credentials_included(self): cfg = ObjStoreConfig( diff --git a/vllm/distributed/kv_events.py b/vllm/distributed/kv_events.py index a7d83bb378f4..c2faf34095de 100644 --- a/vllm/distributed/kv_events.py +++ b/vllm/distributed/kv_events.py @@ -44,6 +44,8 @@ class KVCacheEvent( MEDIUM_GPU = "GPU" MEDIUM_CPU = "CPU" +MEDIUM_FS = "FS" +MEDIUM_OBJ = "OBJ" class BlockStored(KVCacheEvent): diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index a38ac8eb4044..d8d17002856a 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -19,7 +19,7 @@ import json import os from collections.abc import Iterable -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar try: from vllm.fs_io_C import batch_lookup as batch_lookup_C @@ -30,11 +30,18 @@ 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 LookupResult, OffloadKey, ReqContext +from vllm.v1.kv_offload.base import ( + LookupResult, + OffloadingEvent, + OffloadKey, + ReqContext, +) from vllm.v1.kv_offload.file_mapper import FileMapper from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager from vllm.v1.kv_offload.tiering.base import ( + JobId, JobMetadata, JobResult, RequestOffloadingContext, @@ -92,6 +99,8 @@ class FileSystemTierManager(SecondaryTierManager): content. """ + medium: ClassVar[str] = MEDIUM_FS + def __init__( self, offloading_spec: "OffloadingSpec", @@ -100,6 +109,7 @@ def __init__( root_dir: str, n_read_threads: int = 16, n_write_threads: int = 16, + enable_kv_events: bool = False, ): """ Args: @@ -110,9 +120,26 @@ def __init__( root_dir: Root directory for block files. n_read_threads: Number of read-priority I/O threads. n_write_threads: Number of write-priority I/O threads. + enable_kv_events: Emit BlockStored KV events for blocks + successfully stored to this tier. Effective only when KV + cache events are enabled globally (kv_events_config). """ super().__init__(offloading_spec, primary_kv_view, tier_type) + self.events: list[OffloadingEvent] | None = None + if enable_kv_events: + if offloading_spec.kv_events_config.enable_kv_cache_events: + self.events = [] + else: + logger.warning( + "enable_kv_events is set on secondary tier '%s' but KV " + "cache events are disabled globally; the tier will not " + "emit events.", + tier_type, + ) + # Keys of in-flight store jobs, tracked only when events are enabled. + self._store_job_keys: dict[JobId, list[OffloadKey]] = {} + # Extract block size from primary view assert primary_kv_view.strides is not None, ( "primary_kv_view.strides cannot be None" @@ -157,6 +184,8 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: @override 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, @@ -188,10 +217,22 @@ def get_finished_jobs(self) -> Iterable[JobResult]: """ Collect completed jobs from the finished-jobs queue. """ - return ( - JobResult(job_id=job_id, success=success) - for job_id, success in self._pool.get_finished() - ) + results = [] + for job_id, success in self._pool.get_finished(): + if self.events is not None: + keys = self._store_job_keys.pop(job_id, None) + if success and keys: + self.events.append( + OffloadingEvent(keys=keys, medium=self.medium, removed=False) + ) + results.append(JobResult(job_id=job_id, success=success)) + return results + + @override + def take_events(self) -> Iterable[OffloadingEvent]: + if self.events is not None: + yield from self.events + self.events.clear() @override def drain_jobs(self) -> None: diff --git a/vllm/v1/kv_offload/tiering/obj/manager.py b/vllm/v1/kv_offload/tiering/obj/manager.py index 6060370ea067..954f86ed162f 100644 --- a/vllm/v1/kv_offload/tiering/obj/manager.py +++ b/vllm/v1/kv_offload/tiering/obj/manager.py @@ -5,15 +5,22 @@ import ctypes import time from collections.abc import Iterable -from typing import TYPE_CHECKING, NamedTuple +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 LookupResult, OffloadKey, ReqContext +from vllm.v1.kv_offload.base import ( + LookupResult, + OffloadingEvent, + OffloadKey, + ReqContext, +) from vllm.v1.kv_offload.file_mapper import FileMapper from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager from vllm.v1.kv_offload.tiering.base import ( + JobId, JobMetadata, JobResult, RequestOffloadingContext, @@ -90,6 +97,8 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager): primary tier. Object keys are formed as ``{prefix}/{hash_shard}/{hash}.bin``. """ + medium: ClassVar[str] = MEDIUM_OBJ + def __init__( self, offloading_spec: "OffloadingSpec", @@ -98,8 +107,36 @@ def __init__( store_config: dict, prefix: str = "", io_threads: int = 4, + enable_kv_events: bool = False, ): + """ + Args: + offloading_spec: Offloading configuration. + primary_kv_view: Memoryview of the primary tier's CPU KV cache. + tier_type: Tier type identifier, set by SecondaryTierFactory. + store_config: Object store connection parameters (see ObjStoreConfig). + prefix: Key prefix prepended to all object keys. + io_threads: Number of NIXL I/O threads. + enable_kv_events: Emit BlockStored KV events for blocks + successfully stored to this tier. Effective only when KV + cache events are enabled globally (kv_events_config). + """ super().__init__(offloading_spec, primary_kv_view, tier_type) + + self.events: list[OffloadingEvent] | None = None + if enable_kv_events: + if offloading_spec.kv_events_config.enable_kv_cache_events: + self.events = [] + else: + logger.warning( + "enable_kv_events is set on secondary tier '%s' but KV " + "cache events are disabled globally; the tier will not " + "emit events.", + tier_type, + ) + # Keys of in-flight store jobs, tracked only when events are enabled. + self._store_job_keys: dict[JobId, list[OffloadKey]] = {} + agent_config = nixl_agent_config(backends=[]) self._agent = nixl_agent("ObjAgent", agent_config) obj_config = ObjStoreConfig(**store_config) @@ -231,6 +268,8 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: return LookupResult.HIT if result else LookupResult.MISS 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) obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys) self._submit_transfer( job_metadata.job_id, job_metadata.block_ids, obj_keys, NIXL_WRITE @@ -279,8 +318,20 @@ def get_finished_jobs(self) -> Iterable[JobResult]: self._poll_active_transfers() results = self._pending_results self._pending_results = [] + if self.events is not None: + for result in results: + keys = self._store_job_keys.pop(result.job_id, None) + if result.success and keys: + self.events.append( + OffloadingEvent(keys=keys, medium=self.medium, removed=False) + ) return results + def take_events(self) -> Iterable[OffloadingEvent]: + if self.events is not None: + yield from self.events + self.events.clear() + def drain_jobs(self) -> None: """Block until every submitted transfer has completed or failed. From 95ed0feaa5cd7fb16d72c53ce04950aaf07c4698 Mon Sep 17 00:00:00 2001 From: Yan Xu Date: Fri, 10 Jul 2026 12:34:45 +0800 Subject: [PATCH 0007/1526] DCP supports hybrid attention (#40996) Signed-off-by: YanXu Signed-off-by: Jingyi Yang Co-authored-by: Jingyi Yang --- tests/distributed/test_context_parallel.py | 8 + tests/distributed/test_pynccl.py | 47 ++++ .../models/language/generation/test_hybrid.py | 14 +- .../generation/test_vit_cudagraph.py | 1 + tests/test_config.py | 8 +- .../test_gpu_model_runner_streaming.py | 1 + tests/v1/worker/test_cp_utils.py | 45 ++++ tests/v1/worker/test_gpu_input_batch.py | 6 + tests/v1/worker/test_gpu_model_runner.py | 3 + vllm/config/model.py | 7 +- .../device_communicators/cuda_communicator.py | 23 +- vllm/engine/arg_utils.py | 6 +- .../layers/attention_layer_base.py | 1 + vllm/model_executor/layers/mamba/abstract.py | 1 + vllm/utils/cpu_triton_utils.py | 14 +- vllm/v1/attention/backend.py | 2 + vllm/v1/attention/backends/flash_attn.py | 221 ++++++++++++---- vllm/v1/core/kv_cache_coordinator.py | 49 +++- vllm/v1/core/kv_cache_utils.py | 16 +- vllm/v1/core/single_type_kv_cache_manager.py | 4 + vllm/v1/kv_cache_interface.py | 34 +++ vllm/v1/worker/block_table.py | 77 ++++-- vllm/v1/worker/cp_utils.py | 237 ++++++++++++++++++ vllm/v1/worker/gpu_input_batch.py | 7 +- vllm/v1/worker/gpu_model_runner.py | 77 ++++-- vllm/v1/worker/tpu_input_batch.py | 3 +- 26 files changed, 785 insertions(+), 127 deletions(-) create mode 100644 tests/v1/worker/test_cp_utils.py diff --git a/tests/distributed/test_context_parallel.py b/tests/distributed/test_context_parallel.py index 484d29c5b536..007948aeba82 100644 --- a/tests/distributed/test_context_parallel.py +++ b/tests/distributed/test_context_parallel.py @@ -33,6 +33,7 @@ # [LANGUAGE GENERATION] "deepseek-ai/DeepSeek-V2-Lite-Chat", "Qwen/Qwen2.5-1.5B-Instruct", + "Qwen/Qwen3.5-0.8B", # hybrid attention model ] # GSM8K eval configuration @@ -46,6 +47,7 @@ "deepseek-ai/DeepSeek-V2-Lite-Chat": 0.64, # .buildkite/lm-eval-harness/configs/Qwen2.5-1.5B-Instruct.yaml "Qwen/Qwen2.5-1.5B-Instruct": 0.52, + "Qwen/Qwen3.5-0.8B": 0.33, } @@ -151,6 +153,12 @@ def iter_params(self, model_id: str): cp_kv_cache_interleave_size=16, attn_backend="FLASHINFER" ), ], + "Qwen/Qwen3.5-0.8B": [ + CPTestSettings.detailed( + cp_kv_cache_interleave_size=16, + attn_backend="FLASH_ATTN", + ), + ], } diff --git a/tests/distributed/test_pynccl.py b/tests/distributed/test_pynccl.py index d7b04f68091d..62eba4843a0b 100644 --- a/tests/distributed/test_pynccl.py +++ b/tests/distributed/test_pynccl.py @@ -16,6 +16,7 @@ from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary from vllm.distributed.parallel_state import ( ensure_model_parallel_initialized, + get_tp_group, get_world_group, graph_capture, init_distributed_environment, @@ -199,6 +200,52 @@ def test_pynccl_all_gather(): distributed_run(all_gather_worker_fn, 2) +@worker_fn_wrapper +def cuda_communicator_all_gather_dim_worker_fn(): + with ensure_current_vllm_config(): + ensure_model_parallel_initialized(2, 1) + + tp_group = get_tp_group() + comm = tp_group.device_communicator + assert comm is not None + + rank = tp_group.rank_in_group + world_size = tp_group.world_size + device = tp_group.device + + shape = (2, 3, 4) + num_elems = 1 + for size in shape: + num_elems *= size + + for dim in (1, -1): + tensor = ( + torch.arange(num_elems, dtype=torch.float32, device=device).reshape(shape) + + rank * num_elems + ) + expected = torch.cat( + [ + torch.arange(num_elems, dtype=torch.float32, device=device).reshape( + shape + ) + + r * num_elems + for r in range(world_size) + ], + dim=dim, + ) + + result = comm.all_gather(tensor, dim=dim) + torch.accelerator.synchronize() + torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-8) + + +@pytest.mark.skipif( + torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs to run the test." +) +def test_cuda_communicator_all_gather_dim_not_zero(): + distributed_run(cuda_communicator_all_gather_dim_worker_fn, 2) + + @worker_fn_wrapper def all_gatherv_worker_fn(): pynccl_comm = PyNcclCommunicator( diff --git a/tests/models/language/generation/test_hybrid.py b/tests/models/language/generation/test_hybrid.py index 0f19c1038ec7..f06998e07f62 100644 --- a/tests/models/language/generation/test_hybrid.py +++ b/tests/models/language/generation/test_hybrid.py @@ -43,6 +43,11 @@ "tiny-random/qwen3-next-moe", ] +HYBRID_MODELS_REQUIRING_CHUNKED_PREFILL = { + "LiquidAI/LFM2-1.2B", + "tiny-random/qwen3-next-moe", +} + FULL_CUDA_GRAPH_MODELS = [ "ai21labs/Jamba-tiny-dev", "pfnet/plamo-2-1b", @@ -92,8 +97,15 @@ def test_models( example_prompts, max_tokens, num_logprobs ) + extra_kwargs = {} + if model in HYBRID_MODELS_REQUIRING_CHUNKED_PREFILL: + extra_kwargs["enable_chunked_prefill"] = True + with vllm_runner( - model, max_num_seqs=MAX_NUM_SEQS, attention_backend=ATTN_BACKEND + model, + max_num_seqs=MAX_NUM_SEQS, + attention_backend=ATTN_BACKEND, + **extra_kwargs, ) as vllm_model: vllm_outputs = vllm_model.generate_greedy_logprobs( example_prompts, max_tokens, num_logprobs diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index 954bbdbb9b83..387046f3e742 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/tests/models/multimodal/generation/test_vit_cudagraph.py @@ -171,6 +171,7 @@ def gemma3_chat_template(content: str) -> str: "Describe this video in one sentence." ), needs_video_metadata=True, + vllm_runner_kwargs={"enable_chunked_prefill": True}, marks=[pytest.mark.core_model], ), "internvl": VitCudagraphTestConfig( diff --git a/tests/test_config.py b/tests/test_config.py index 1e93b610da56..6f2be35ae86f 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1121,14 +1121,14 @@ def test_is_chunked_prefill_supported( ( "Qwen/Qwen3-Next-80B-A3B-Instruct", "hybrid", - False, - "Hybrid models do not support prefix caching since the feature is still experimental.", # noqa: E501 + True, + "Generative hybrid models support prefix caching.", # noqa: E501 ), ( "ibm-granite/granite-4.0-h-small", "hybrid", - False, - "Hybrid models do not support prefix caching since the feature is still experimental.", # noqa: E501 + True, + "Generative hybrid models support prefix caching.", # noqa: E501 ), ( "state-spaces/mamba-130m-hf", diff --git a/tests/v1/streaming_input/test_gpu_model_runner_streaming.py b/tests/v1/streaming_input/test_gpu_model_runner_streaming.py index 9b130e570f66..fd619610b767 100644 --- a/tests/v1/streaming_input/test_gpu_model_runner_streaming.py +++ b/tests/v1/streaming_input/test_gpu_model_runner_streaming.py @@ -38,6 +38,7 @@ def mock_model_runner_with_input_batch(): vocab_size=32000, block_sizes=[16], kernel_block_sizes=[16], + max_num_blocks_per_req=[64], logitsprocs=None, is_pooling_model=False, ) diff --git a/tests/v1/worker/test_cp_utils.py b/tests/v1/worker/test_cp_utils.py new file mode 100644 index 000000000000..b38ae4d0636c --- /dev/null +++ b/tests/v1/worker/test_cp_utils.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest +import torch + +from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens +from vllm.v1.worker.cp_utils import should_skip_dcp_context_attention + + +def test_skip_gate_only_for_zero_context(): + assert should_skip_dcp_context_attention(torch.zeros(3, dtype=torch.int32)) + assert not should_skip_dcp_context_attention( + torch.tensor([0, 5, 0], dtype=torch.int32) + ) + + +@pytest.mark.parametrize( + "dcp_world_size,interleave_size,context_len", + [(2, 16, 10), (4, 16, 10), (8, 16, 10), (4, 1, 2)], +) +def test_skip_gate_rank_invariant_with_divergent_local_context( + dcp_world_size: int, interleave_size: int, context_len: int +): + """Contexts shorter than a full interleave round land entirely on a + subset of DCP ranks, so the per-rank local context lengths diverge: + some ranks hold zero local context while others hold all of it. Ranks + with zero local context must still take the collective (non-skip) path, + otherwise the query all-gather in _forward_with_dcp deadlocks across + ranks. The skip gate must therefore depend only on the rank-invariant + global context lengths, never on get_dcp_local_seq_lens output. + """ + context_kv_lens = torch.tensor([context_len], dtype=torch.int32) + local_maxes = [ + int( + get_dcp_local_seq_lens( + context_kv_lens, dcp_world_size, rank, interleave_size + ).max() + ) + for rank in range(dcp_world_size) + ] + # Precondition: the local view diverges across ranks. + assert 0 in local_maxes + assert max(local_maxes) > 0 + # The batch still has context globally, so no rank may skip. + assert not should_skip_dcp_context_attention(context_kv_lens) diff --git a/tests/v1/worker/test_gpu_input_batch.py b/tests/v1/worker/test_gpu_input_batch.py index 4d0a1698a70b..e7aa2e03e0e4 100644 --- a/tests/v1/worker/test_gpu_input_batch.py +++ b/tests/v1/worker/test_gpu_input_batch.py @@ -238,6 +238,7 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int): vocab_size=1024, block_sizes=[1], kernel_block_sizes=[1], + max_num_blocks_per_req=[1024], ) reqs: list[CachedRequestState] = [] req_id_reqs = {} @@ -332,6 +333,7 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis vocab_size=1024, block_sizes=[1], kernel_block_sizes=[1], + max_num_blocks_per_req=[1024], ) ref_input_batch: InputBatch = InputBatch( max_num_reqs=batch_size, @@ -341,6 +343,7 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis vocab_size=1024, block_sizes=[1], kernel_block_sizes=[1], + max_num_blocks_per_req=[1024], ) reqs: list[CachedRequestState] = [] @@ -409,6 +412,7 @@ def test_pooling_prompt_lens_not_aliased(device: str): vocab_size=VOCAB_SIZE, block_sizes=[16], kernel_block_sizes=[16], + max_num_blocks_per_req=[64], is_pooling_model=True, ) @@ -444,6 +448,7 @@ def test_placeholder_spec_token_ids_written_verbatim(): vocab_size=VOCAB_SIZE, block_sizes=[16], kernel_block_sizes=[16], + max_num_blocks_per_req=[1], ) req = CachedRequestState( req_id="req", @@ -491,6 +496,7 @@ def test_pooling_metadata_token_id_buffers( vocab_size=VOCAB_SIZE, block_sizes=[16], kernel_block_sizes=[16], + max_num_blocks_per_req=[64], is_pooling_model=True, ) req = _construct_pooling_request(0, PoolingParams(**pooling_params)) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 7490cbb27be5..6e5b1aa5b882 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -89,6 +89,7 @@ def initialize_kv_cache(runner: GPUModelRunner): kernel_block_sizes=[ kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size ], + max_num_blocks_per_req=[NUM_BLOCKS], ) runner.initialize_attn_backend(kv_cache_config) @@ -1397,6 +1398,7 @@ def test_input_batch_with_kernel_block_sizes(): vocab_size=vocab_size, block_sizes=block_sizes, kernel_block_sizes=kernel_block_sizes, + max_num_blocks_per_req=[16, 8], ) # Verify that block tables were created with kernel block sizes @@ -1457,6 +1459,7 @@ def test_hybrid_cache_integration(default_vllm_config, dist_init): vocab_size=runner.model_config.get_vocab_size(), block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size], kernel_block_sizes=[16], + max_num_blocks_per_req=[NUM_BLOCKS], ) # Use kernel block size runner.initialize_attn_backend(kv_cache_config) diff --git a/vllm/config/model.py b/vllm/config/model.py index 6e3fef0dcca7..b5f4e69e031f 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1881,11 +1881,8 @@ def is_prefix_caching_supported(self) -> bool: else: # for generative models if attn_type == "hybrid": - logger.debug( - "Hybrid models do not support prefix caching since the feature " - "is still experimental." - ) - return False + logger.debug("Generative hybrid models support prefix caching.") + return True elif attn_type == "attention_free": logger.debug( "Attention free models do not support prefix caching since the " diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 555fd0ec9489..a196a54f5091 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -341,13 +341,30 @@ def all_reduce(self, input_): def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: # Route uniform dim-0 all-gathers through NVLS symmetric memory when # enabled (mirrors reduce_scatter); otherwise fall back to the - # base-class ring all-gather. Sequence parallelism's gather-before-GEMM - # uses dim=0 with tp-aligned (uniform) shards. + # PyNccl/base-class all-gather. Sequence parallelism's + # gather-before-GEMM uses dim=0 with tp-aligned (uniform) shards. if dim < 0: dim += input_.dim() if dim == 0 and should_nccl_symm_mem_ag_rs(): return self._all_gather_symm_mem(input_.contiguous()) - return super().all_gather(input_, dim) + + pynccl_comm = self.pynccl_comm + if pynccl_comm is None or pynccl_comm.disabled: + return super().all_gather(input_, dim) + + input_size = input_.size() + output_size = (input_size[0] * self.world_size,) + input_size[1:] + output_tensor = torch.empty( + output_size, dtype=input_.dtype, device=input_.device + ) + pynccl_comm.all_gather(output_tensor, input_.contiguous()) + output_tensor = output_tensor.reshape((self.world_size,) + input_size) + output_tensor = output_tensor.movedim(0, dim) + return output_tensor.reshape( + input_size[:dim] + + (self.world_size * input_size[dim],) + + input_size[dim + 1 :] + ) def reduce_scatter(self, input_: torch.Tensor, dim: int = -1): world_size = self.world_size diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index d80ecb9c8555..c602eaf8f2ed 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -2493,7 +2493,11 @@ def _set_default_chunked_prefill_and_prefix_caching_args( self, model_config: ModelConfig ) -> None: default_chunked_prefill = model_config.is_chunked_prefill_supported - default_prefix_caching = model_config.is_prefix_caching_supported + # Hybrid models support prefix caching but keep it opt-in for now + # while the feature matures. + default_prefix_caching = ( + model_config.is_prefix_caching_supported and not model_config.is_hybrid + ) if self.enable_chunked_prefill is None: self.enable_chunked_prefill = default_chunked_prefill diff --git a/vllm/model_executor/layers/attention_layer_base.py b/vllm/model_executor/layers/attention_layer_base.py index 97395b641497..53994114558e 100644 --- a/vllm/model_executor/layers/attention_layer_base.py +++ b/vllm/model_executor/layers/attention_layer_base.py @@ -19,6 +19,7 @@ class AttentionLayerBase(ABC): """ impl: "AttentionImpl" + supports_dcp: bool = True @abstractmethod def get_attn_backend(self) -> type[AttentionBackend]: diff --git a/vllm/model_executor/layers/mamba/abstract.py b/vllm/model_executor/layers/mamba/abstract.py index 8bbb21d7bc90..62fd64bd0dea 100644 --- a/vllm/model_executor/layers/mamba/abstract.py +++ b/vllm/model_executor/layers/mamba/abstract.py @@ -22,6 +22,7 @@ class MambaBase(AttentionLayerBase): # Contains the KV cache (mamba state) for the layer # in the shape specified by `self.get_state_shape`. kv_cache: tuple[torch.Tensor, ...] + supports_dcp: bool = False @abstractmethod def get_state_shape(self) -> Iterable[tuple[int, ...]]: diff --git a/vllm/utils/cpu_triton_utils.py b/vllm/utils/cpu_triton_utils.py index 3b5012d01751..9ce33b835707 100644 --- a/vllm/utils/cpu_triton_utils.py +++ b/vllm/utils/cpu_triton_utils.py @@ -29,13 +29,17 @@ def _compute_slot_mapping_kernel_impl( block_table_stride: int, # max_num_blocks_per_req block_size: int, slot_mapping: torch.Tensor, # [max_num_tokens], int64 - TOTAL_CP_WORLD_SIZE: int, - TOTAL_CP_RANK: int, - CP_KV_CACHE_INTERLEAVE_SIZE: int, - PAD_ID: int, - BLOCK_SIZE: int, + KV_CACHE_BLOCK_SIZE: int | None = None, + BLOCKS_PER_KV_BLOCK: int = 1, + TOTAL_CP_WORLD_SIZE: int = 1, + TOTAL_CP_RANK: int = 0, + CP_KV_CACHE_INTERLEAVE_SIZE: int = 1, + PAD_ID: int = -1, + BLOCK_SIZE: int = 1024, ) -> None: assert TOTAL_CP_WORLD_SIZE == 1, "Context Parallelism is not supported on CPU." + if BLOCKS_PER_KV_BLOCK != 1: + assert block_size * BLOCKS_PER_KV_BLOCK == KV_CACHE_BLOCK_SIZE torch.ops._C.compute_slot_mapping_kernel_impl( query_start_loc, positions, diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index bfecb3c952e1..ec07023f80ef 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -796,6 +796,8 @@ class AttentionImplBase(ABC, Generic[T]): # Whether the attention impl supports Prefill Context Parallelism. supports_pcp: bool = False + # Whether the attention impl supports Decode Context Parallelism. + supports_dcp: bool = True # Whether the attention impl(or ops) supports MTP # when cp_kv_cache_interleave_size > 1 supports_mtp_with_cp_non_trivial_interleave_size: bool = False diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index db5b0dda367e..d83d2f4d810a 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -56,10 +56,14 @@ AttentionMetadataBuilder, CommonAttentionMetadata, ) -from vllm.v1.attention.backends.utils import ( - get_kv_cache_layout, -) +from vllm.v1.attention.backends.utils import get_kv_cache_layout from vllm.v1.kv_cache_interface import AttentionSpec +from vllm.v1.worker.cp_utils import ( + run_split_fa2_dcp_context_attention, + should_skip_dcp_context_attention, + should_split_fa2_dcp_context_attention, + split_dcp_context_queries, +) logger = init_logger(__name__) @@ -245,6 +249,13 @@ class FlashAttentionMetadata: max_dcp_context_kv_len: int | None = None dcp_context_kv_lens: torch.Tensor | None = None + # Split counts for FA2 DCP context attention. num_prefill_* tracks + # context-bearing extend rows; pure prefills do not attend to DCP context. + num_decode_reqs: int = 0 + num_prefill_reqs: int = 0 + num_decode_tokens: int = 0 + num_prefill_tokens: int = 0 + # Optional aot scheduling scheduler_metadata: torch.Tensor | None = None prefix_scheduler_metadata: torch.Tensor | None = None @@ -513,6 +524,10 @@ def schedule( use_cascade = common_prefix_len > 0 max_dcp_context_kv_len = 0 dcp_context_kv_lens = None + num_decode_reqs = 0 + num_prefill_reqs = 0 + num_decode_tokens = 0 + num_prefill_tokens = 0 cu_prefix_query_lens = None prefix_kv_lens = None @@ -532,23 +547,54 @@ def schedule( self._dcp_context_kv_lens[num_reqs:] = 0 dcp_context_kv_lens = self._dcp_context_kv_lens[:num_reqs] + skip_dcp_context_attention = False + if common_attn_metadata.seq_lens_cpu_upper_bound is not None: + query_lens_cpu = ( + common_attn_metadata.query_start_loc_cpu[1 : num_reqs + 1] + - common_attn_metadata.query_start_loc_cpu[:num_reqs] + ) + context_kv_lens_cpu = ( + common_attn_metadata.seq_lens_cpu_upper_bound[:num_reqs] + - query_lens_cpu + ) + skip_dcp_context_attention = should_skip_dcp_context_attention( + context_kv_lens_cpu + ) + + if max_query_len > 1: + ( + num_decode_reqs, + num_prefill_reqs, + num_decode_tokens, + num_prefill_tokens, + ) = split_dcp_context_queries( + common_attn_metadata.query_start_loc_cpu, + common_attn_metadata.seq_lens_cpu_upper_bound, + max_query_len, + num_actual_tokens, + ) + # After DCP distribution, the maximum number of tokens for any rank is # ceil(L / (N * I)) * I, where L is max_seq_len, N is dcp_world_size, # and I is cp_kv_cache_interleave_size. # This eliminates GPU->CPU sync while minimizing workspace over-allocation. - num_partitions = self.dcp_world_size * self.cp_kv_cache_interleave_size - max_dcp_context_kv_len = ( - (max_seq_len + num_partitions - 1) // num_partitions - ) * self.cp_kv_cache_interleave_size - - scheduler_metadata = schedule( - batch_size=num_reqs, - cu_query_lens=query_start_loc, - max_query_len=max_query_len, - seqlens=dcp_context_kv_lens, - max_seq_len=max_dcp_context_kv_len, - causal=False, - ) + if skip_dcp_context_attention: + max_dcp_context_kv_len = 0 + scheduler_metadata = None + else: + num_partitions = self.dcp_world_size * self.cp_kv_cache_interleave_size + max_dcp_context_kv_len = ( + (max_seq_len + num_partitions - 1) // num_partitions + ) * self.cp_kv_cache_interleave_size + + scheduler_metadata = schedule( + batch_size=num_reqs, + cu_query_lens=query_start_loc, + max_query_len=max_query_len, + seqlens=dcp_context_kv_lens, + max_seq_len=max_dcp_context_kv_len, + causal=False, + ) elif use_cascade: cu_prefix_query_lens = torch.tensor( [0, num_actual_tokens], dtype=torch.int32, device=self.device @@ -614,6 +660,10 @@ def schedule( slot_mapping=slot_mapping, max_dcp_context_kv_len=max_dcp_context_kv_len, dcp_context_kv_lens=dcp_context_kv_lens, + num_decode_reqs=num_decode_reqs, + num_prefill_reqs=num_prefill_reqs, + num_decode_tokens=num_decode_tokens, + num_prefill_tokens=num_prefill_tokens, use_cascade=use_cascade, common_prefix_len=common_prefix_len, scheduler_metadata=scheduler_metadata, @@ -743,8 +793,12 @@ def __init__( self.dcp_combine = dcp_a2a_lse_reduce if dcp_a2a else cp_lse_ag_out_rs self._dcp_dtype: torch.dtype | None = None + self._dcp_max_num_tokens: int = 0 if vllm_config is not None and self.dcp_world_size > 1: self._dcp_dtype = vllm_config.model_config.dtype + self._dcp_max_num_tokens = ( + vllm_config.scheduler_config.max_num_batched_tokens + ) def forward( self, @@ -1063,40 +1117,120 @@ def _forward_with_dcp( block_table = attn_metadata.block_table query = query.contiguous() + if attn_metadata.max_dcp_context_kv_len == 0: + flash_attn_varlen_func( + q=query, + k=key, + v=value, + out=output, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=max_seqlen_q, + cu_seqlens_k=cu_seqlens_q, + max_seqlen_k=max_seqlen_q, + softmax_scale=self.scale, + causal=attn_metadata.causal, + alibi_slopes=self.alibi_slopes, + window_size=list(self.sliding_window) + if self.sliding_window is not None + else None, + softcap=self.logits_soft_cap, + return_softmax_lse=True, + fa_version=self.vllm_flash_attn_version, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + num_splits=attn_metadata.max_num_splits, + ) + return output + query_across_dcp = get_dcp_group().all_gather(query, dim=1) sliding_window_size = ( list(self.sliding_window) if self.sliding_window is not None else None ) n = query_across_dcp.shape[0] - (dcp_context_out,) = current_workspace_manager().get_simultaneous( + num_reqs = cu_seqlens_q.shape[0] - 1 + num_decodes = attn_metadata.num_decode_reqs + num_context_prefills = attn_metadata.num_prefill_reqs + num_decode_tokens = attn_metadata.num_decode_tokens + num_context_prefill_tokens = attn_metadata.num_prefill_tokens + split_dcp_context = should_split_fa2_dcp_context_attention( + self.vllm_flash_attn_version, + max_seqlen_q, + num_reqs, + num_decodes, + num_context_prefills, + ) + dcp_context_out_tokens = max(n, self._dcp_max_num_tokens) + dcp_context_out_spec = ( ( - (n, self.num_heads * self.dcp_world_size, self.head_size), - self._dcp_dtype, + dcp_context_out_tokens, + self.num_heads * self.dcp_world_size, + self.head_size, ), + self._dcp_dtype, ) - context_attn_out, context_lse = flash_attn_varlen_func( - q=query_across_dcp, - k=key_cache, - v=value_cache, - out=dcp_context_out, - cu_seqlens_q=cu_seqlens_q, - max_seqlen_q=max_seqlen_q, - seqused_k=attn_metadata.dcp_context_kv_lens, - max_seqlen_k=attn_metadata.max_dcp_context_kv_len, - softmax_scale=self.scale, - causal=False, - alibi_slopes=self.alibi_slopes, - window_size=sliding_window_size, - block_table=block_table, - softcap=self.logits_soft_cap, - return_softmax_lse=True, - scheduler_metadata=attn_metadata.scheduler_metadata, - fa_version=self.vllm_flash_attn_version, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, - num_splits=attn_metadata.max_num_splits, + (dcp_context_out_workspace,) = current_workspace_manager().get_simultaneous( + dcp_context_out_spec, ) + dcp_context_out = dcp_context_out_workspace[:n] + + if split_dcp_context: + # TODO: Remove this DCP + FA2 mixed decode/prefill workaround once + # FA4 supports this Qwen3.5 shape. + assert attn_metadata.dcp_context_kv_lens is not None + assert attn_metadata.max_dcp_context_kv_len is not None + assert self.vllm_flash_attn_version is not None + context_attn_out, context_lse = run_split_fa2_dcp_context_attention( + flash_attn_varlen_func, + query_across_dcp, + key_cache, + value_cache, + dcp_context_out, + cu_seqlens_q, + max_seqlen_q, + attn_metadata.dcp_context_kv_lens, + attn_metadata.max_dcp_context_kv_len, + self.scale, + self.alibi_slopes, + sliding_window_size, + block_table, + self.logits_soft_cap, + self.vllm_flash_attn_version, + q_descale, + k_descale, + v_descale, + attn_metadata.max_num_splits, + self.num_heads, + self.dcp_world_size, + num_decodes, + num_context_prefills, + num_decode_tokens, + num_context_prefill_tokens, + ) + else: + context_attn_out, context_lse = flash_attn_varlen_func( + q=query_across_dcp, + k=key_cache, + v=value_cache, + out=dcp_context_out, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=max_seqlen_q, + seqused_k=attn_metadata.dcp_context_kv_lens, + max_seqlen_k=attn_metadata.max_dcp_context_kv_len, + softmax_scale=self.scale, + causal=False, + alibi_slopes=self.alibi_slopes, + window_size=sliding_window_size, + block_table=block_table, + softcap=self.logits_soft_cap, + return_softmax_lse=True, + scheduler_metadata=attn_metadata.scheduler_metadata, + fa_version=self.vllm_flash_attn_version, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + num_splits=attn_metadata.max_num_splits, + ) # FA returns LSE in shape [ H, B ] but DCP combine wants [ B, H ] context_attn_out_cor, context_lse_cor = self.dcp_combine( context_attn_out, @@ -1106,14 +1240,11 @@ def _forward_with_dcp( ) context_lse_cor = context_lse_cor.transpose(0, 1).contiguous() - (dcp_query_out,) = current_workspace_manager().get_simultaneous( - ((query.shape[0], self.num_heads, self.head_size), self._dcp_dtype), - ) query_attn_out, query_lse = flash_attn_varlen_func( q=query, k=key, v=value, - out=dcp_query_out, + out=output, cu_seqlens_q=cu_seqlens_q, max_seqlen_q=max_seqlen_q, cu_seqlens_k=cu_seqlens_q, diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 4756136f03cf..2e978451b885 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -550,12 +550,28 @@ def __init__( # different KV cache groups have different block sizes, the actual block size # can be a multiple of hash_block_size. self.hash_block_size = hash_block_size + self.dcp_world_size = dcp_world_size + group_block_sizes = [ + manager.block_size for manager in self.single_type_managers + ] assert all( - g.kv_cache_spec.block_size % hash_block_size == 0 - for g in kv_cache_config.kv_cache_groups - ), "block_size must be divisible by hash_block_size" - assert dcp_world_size == 1, "DCP not support hybrid attn now." + block_size % hash_block_size == 0 for block_size in group_block_sizes + ), ( + "Each KV cache group's real block_size must be divisible by " + f"hash_block_size. block_sizes={group_block_sizes}, " + f"hash_block_size={hash_block_size}" + ) assert pcp_world_size == 1, "PCP not support hybrid attn now." + if dcp_world_size > 1: + # DCP shards full-attention KV across ranks and replicates Mamba + # state; other spec types (e.g. sliding window) have no DCP-aware + # handling yet, so reject them explicitly. + for g in kv_cache_config.kv_cache_groups: + assert isinstance(g.kv_cache_spec, (FullAttentionSpec, MambaSpec)), ( + "DCP with hybrid KV cache layouts only supports " + "full-attention and Mamba groups, got: " + f"{type(g.kv_cache_spec).__name__}." + ) self.verify_and_split_kv_cache_groups() def verify_and_split_kv_cache_groups(self) -> None: @@ -651,11 +667,11 @@ def find_longest_cache_hit( - The number of tokens of the longest cache hit. """ - def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList: - if kv_cache_spec.block_size == self.hash_block_size: + def _get_block_hashes(block_size: int) -> BlockHashList: + if block_size == self.hash_block_size: return block_hashes return BlockHashListWithBlockSize( - block_hashes, self.hash_block_size, kv_cache_spec.block_size + block_hashes, self.hash_block_size, block_size ) num_groups = len(self.kv_cache_config.kv_cache_groups) @@ -680,13 +696,14 @@ def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList: for idx, (spec, group_ids, manager_cls, use_eagle) in enumerate( self.attention_groups ): + group_block_size = self.single_type_managers[group_ids[0]].block_size cached_blocks = hit_blocks_by_group[group_ids[0]] if isinstance(spec, FullAttentionSpec) and cached_blocks is not None: # Full attention is downward-closed: we only need to look # up cached blocks once; on subsequent iterations just trim # to the (reduced) current hit length. curr_hit_length = ( - curr_hit_length // spec.block_size * spec.block_size + curr_hit_length // group_block_size * group_block_size ) continue @@ -696,18 +713,23 @@ def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList: if drop_eagle_block: # Eagle needs to match one more block and then pop the last. _max_length = min( - curr_hit_length + spec.block_size, max_cache_hit_length + curr_hit_length + group_block_size, max_cache_hit_length ) hit_blocks = manager_cls.find_longest_cache_hit( - block_hashes=_get_block_hashes(spec), + block_hashes=_get_block_hashes(group_block_size), max_length=_max_length, kv_cache_group_ids=group_ids, block_pool=self.block_pool, kv_cache_spec=spec, drop_eagle_block=drop_eagle_block, alignment_tokens=self.scheduler_block_size, + dcp_world_size=( + self.dcp_world_size + if isinstance(spec, FullAttentionSpec) + else 1 + ), ) - _new_hit_length = len(hit_blocks[0]) * spec.block_size + _new_hit_length = len(hit_blocks[0]) * group_block_size if drop_eagle_block: eagle_verified.add(idx) elif _new_hit_length < curr_hit_length: @@ -728,7 +750,10 @@ def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList: # Truncate full attention blocks to final hit_length (if present) first_group = self.attention_groups[0] if isinstance(first_group.spec, FullAttentionSpec): - num_blocks = hit_length // first_group.spec.block_size + group_block_size = self.single_type_managers[ + first_group.group_ids[0] + ].block_size + num_blocks = hit_length // group_block_size for group_id in first_group.group_ids: if (blks := hit_blocks_by_group[group_id]) is not None: del blks[num_blocks:] diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index aa42f90bb90e..6e9530e381bf 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -627,7 +627,8 @@ def resolve_kv_cache_block_sizes( - ``scheduler_block_size`` is the token-alignment invariant used by the scheduler (e.g. for ``num_computed_tokens`` rounding). Single group: ``cache_config.block_size * dcp * pcp``. Multiple groups: LCM of every - group's block size — context parallelism is not supported here. + group's effective block size. Attention groups are scaled by DCP/PCP; + Mamba groups keep their full per-rank state and are not scaled. - ``hash_block_size`` is the granularity at which ``Request.block_hashes`` is computed. Single group: equals scheduler block size. Multiple groups: ``cache_config.hash_block_size`` override if set, else the GCD of group @@ -645,13 +646,12 @@ def resolve_kv_cache_block_sizes( bs = cache_config.block_size * dcp * pcp return bs, bs - if dcp != 1 or pcp != 1: - raise ValueError( - "Hybrid KV cache groups with multiple block sizes do not " - "support context parallelism (dcp_world_size/pcp_world_size > 1)." - ) - - group_block_sizes = [g.kv_cache_spec.block_size for g in groups] + group_block_sizes = [ + g.kv_cache_spec.block_size * dcp * pcp + if isinstance(g.kv_cache_spec, AttentionSpec) + else g.kv_cache_spec.block_size + for g in groups + ] scheduler_block_size = math.lcm(*group_block_sizes) # Block hashes are only consumed by prefix caching and KV connectors diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 87c3f8feb725..324d0016cd0b 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -1031,6 +1031,10 @@ def __init__( self, kv_cache_spec: MambaSpec, block_pool: BlockPool, **kwargs ) -> None: super().__init__(kv_cache_spec, block_pool, **kwargs) + # Mamba layers use TP instead of DCP, so each rank holds the full + # recurrent state. Undo the DCP/PCP block_size scaling that the base + # class applies for attention groups whose KV cache is partitioned. + self.block_size = kv_cache_spec.block_size self.cached_blocks_this_step: set[BlockHashWithGroupId] = set() self.mamba_cache_mode = kv_cache_spec.mamba_cache_mode self.num_speculative_blocks: int = kv_cache_spec.num_speculative_blocks diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 4204c31be58d..b0409b16dee2 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -128,6 +128,18 @@ def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: """ raise NotImplementedError + def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: + """ + The number of block table entries needed per request, i.e. the row + length of the worker-side block table for this cache group. + + Args: + vllm_config: The vllm config. + max_len: The maximum sequence length to size for, including the + encoder length for encoder-decoder models. + """ + return cdiv(max_len, self.block_size) + def copy_with_new_block_size(self, block_size: int) -> Self: """ Create a new KVCacheSpec from self but replacing the block size. @@ -201,6 +213,16 @@ def real_page_size_bytes(self) -> int: * get_dtype_size(self.dtype) ) + def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: + # Attention KV is token-interleaved across DCP/PCP ranks, so each rank + # only stores max_len // (dcp * pcp) tokens per request. + parallel_config = vllm_config.parallel_config + total_cp_size = ( + parallel_config.decode_context_parallel_size + * parallel_config.prefill_context_parallel_size + ) + return cdiv(max_len, self.block_size * total_cp_size) + @dataclass(frozen=True, kw_only=True) class FullAttentionSpec(AttentionSpec): @@ -699,6 +721,18 @@ def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: else: return self.page_size_bytes * (1 + self.num_speculative_blocks) + def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: + # Mamba state is replicated across DCP/PCP ranks, never sharded, so + # no CP scaling applies. + if vllm_config.cache_config.mamba_cache_mode == "align": + # Block table rows are position-indexed over the full sequence + # even though only 2 + num_speculative_blocks state blocks are + # resident at a time (earlier states are nulled out by + # remove_skipped_blocks), so the row length must cover max_len + # rather than max_memory_usage_bytes. + return cdiv(max_len, self.block_size) + self.num_speculative_blocks + return cdiv(self.max_memory_usage_bytes(vllm_config), self.page_size_bytes) + def is_uniform_with_collection( self, kv_cache_specs: dict[str, KVCacheSpec] ) -> bool: diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py index d9c041ba0b89..d40887879fcb 100644 --- a/vllm/v1/worker/block_table.py +++ b/vllm/v1/worker/block_table.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from enum import Enum + import numpy as np import torch @@ -10,11 +12,15 @@ from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm.v1.utils import CpuGpuBuffer -from vllm.v1.worker.cp_utils import get_total_cp_world_size logger = init_logger(__name__) +class SlotMappingMode(Enum): + TOKEN_TO_KV_SLOT = "token_to_kv_slot" + NONE = "none" + + class BlockTable: def __init__( self, @@ -26,6 +32,7 @@ def __init__( device: torch.device, kernel_block_size: int, cp_kv_cache_interleave_size: int, + slot_mapping_mode: SlotMappingMode = SlotMappingMode.TOKEN_TO_KV_SLOT, ): """ Args: @@ -38,11 +45,15 @@ def __init__( kernel_block_size: The block_size of underlying attention kernel. Will be the same as `block_size` if `block_size` is supported by the attention kernel. + slot_mapping_mode: How this cache group maps scheduled tokens to + cache slots. Mamba-like state caches do not use token slot + mappings and should use SlotMappingMode.NONE. """ self.max_num_reqs = max_num_reqs self.max_num_batched_tokens = max_num_batched_tokens self.pin_memory = pin_memory self.device = device + self.kv_cache_block_size = block_size if kernel_block_size == block_size: # Standard case: allocation and computation use same block size @@ -98,6 +109,7 @@ def __init__( self.dcp_world_size = 1 self.dcp_rank = 0 self.cp_kv_cache_interleave_size = cp_kv_cache_interleave_size + self.slot_mapping_mode = slot_mapping_mode def append_row( self, @@ -145,6 +157,12 @@ def compute_slot_mapping( positions: torch.Tensor, ) -> None: num_tokens = positions.shape[0] + if self.slot_mapping_mode == SlotMappingMode.NONE: + # Mamba/GDN groups consume the block table as recurrent state + # indices and do not use per-token slot mappings. + return + assert self.slot_mapping_mode == SlotMappingMode.TOKEN_TO_KV_SLOT + total_cp_world_size = self.pcp_world_size * self.dcp_world_size total_cp_rank = self.pcp_rank * self.dcp_world_size + self.dcp_rank _compute_slot_mapping_kernel[(num_reqs + 1,)]( @@ -156,6 +174,8 @@ def compute_slot_mapping( self.block_table.gpu.stride(0), self.block_size, self.slot_mapping.gpu, + KV_CACHE_BLOCK_SIZE=self.kv_cache_block_size, + BLOCKS_PER_KV_BLOCK=self.blocks_per_kv_block, TOTAL_CP_WORLD_SIZE=total_cp_world_size, TOTAL_CP_RANK=total_cp_rank, CP_KV_CACHE_INTERLEAVE_SIZE=self.cp_kv_cache_interleave_size, @@ -226,30 +246,27 @@ class MultiGroupBlockTable: def __init__( self, max_num_reqs: int, - max_model_len: int, max_num_batched_tokens: int, pin_memory: bool, device: torch.device, block_sizes: list[int], kernel_block_sizes: list[int], - max_num_blocks: list[int] | None = None, + max_num_blocks: list[int], cp_kv_cache_interleave_size: int = 1, + slot_mapping_modes: list[SlotMappingMode] | None = None, ) -> None: if len(kernel_block_sizes) != len(block_sizes): raise ValueError( f"kernel_block_sizes length ({len(kernel_block_sizes)}) " f"must match block_sizes length ({len(block_sizes)})" ) - if max_num_blocks is None: - # Note(hc): each dcp rank only store - # (max_model_len//dcp_world_size) tokens in kvcache, - # so the block_size which used for calc max_num_blocks_per_req - # must be multiplied by dcp_world_size. - total_cp_world_size = get_total_cp_world_size() - max_num_blocks = [ - cdiv(max_model_len, block_size * total_cp_world_size) - for block_size in block_sizes - ] + if slot_mapping_modes is None: + slot_mapping_modes = [SlotMappingMode.TOKEN_TO_KV_SLOT] * len(block_sizes) + if len(slot_mapping_modes) != len(block_sizes): + raise ValueError( + f"slot_mapping_modes length ({len(slot_mapping_modes)}) " + f"must match block_sizes length ({len(block_sizes)})" + ) if len(max_num_blocks) != len(block_sizes): raise ValueError( @@ -274,9 +291,15 @@ def __init__( device, kernel_block_size, cp_kv_cache_interleave_size, + slot_mapping_mode=slot_mapping_mode, ) - for block_size, kernel_block_size, max_num_blocks_per_req in zip( - block_sizes, kernel_block_sizes, max_num_blocks + for ( + block_size, + kernel_block_size, + max_num_blocks_per_req, + slot_mapping_mode, + ) in zip( + block_sizes, kernel_block_sizes, max_num_blocks, slot_mapping_modes ) ] @@ -332,6 +355,8 @@ def _compute_slot_mapping_kernel( block_table_stride, # max_num_blocks_per_req block_size, slot_mapping_ptr, # [max_num_tokens], int64 + KV_CACHE_BLOCK_SIZE: tl.constexpr, + BLOCKS_PER_KV_BLOCK: tl.constexpr, TOTAL_CP_WORLD_SIZE: tl.constexpr, TOTAL_CP_RANK: tl.constexpr, CP_KV_CACHE_INTERLEAVE_SIZE: tl.constexpr, @@ -354,18 +379,14 @@ def _compute_slot_mapping_kernel( start_idx = tl.load(query_start_loc_ptr + req_idx).to(tl.int64) end_idx = tl.load(query_start_loc_ptr + req_idx + 1).to(tl.int64) - virtual_block_size = block_size * TOTAL_CP_WORLD_SIZE + virtual_block_size = KV_CACHE_BLOCK_SIZE * TOTAL_CP_WORLD_SIZE row_offset = req_idx * block_table_stride for i in range(start_idx, end_idx, BLOCK_SIZE): offsets = i + tl.arange(0, BLOCK_SIZE) mask = offsets < end_idx pos = tl.load(positions_ptr + offsets, mask=mask, other=0) - block_indices = pos // virtual_block_size - block_numbers = tl.load(block_table_ptr + row_offset + block_indices).to( - tl.int64 - ) - - virtual_block_offsets = pos - block_indices * virtual_block_size + virtual_block_indices = pos // virtual_block_size + virtual_block_offsets = pos - virtual_block_indices * virtual_block_size is_local = ( virtual_block_offsets // CP_KV_CACHE_INTERLEAVE_SIZE ) % TOTAL_CP_WORLD_SIZE == TOTAL_CP_RANK @@ -375,6 +396,16 @@ def _compute_slot_mapping_kernel( virtual_block_offsets % CP_KV_CACHE_INTERLEAVE_SIZE ) - slot_ids = block_numbers * block_size + local_block_offsets + block_indices = ( + virtual_block_indices * BLOCKS_PER_KV_BLOCK + + local_block_offsets // block_size + ) + block_numbers = tl.load( + block_table_ptr + row_offset + block_indices, + mask=mask & is_local, + other=0, + ).to(tl.int64) + slot_offsets = local_block_offsets % block_size + slot_ids = block_numbers * block_size + slot_offsets slot_ids = tl.where(is_local, slot_ids, PAD_ID) tl.store(slot_mapping_ptr + offsets, slot_ids, mask=mask) diff --git a/vllm/v1/worker/cp_utils.py b/vllm/v1/worker/cp_utils.py index 05cca52fc0db..11edd86a5db0 100644 --- a/vllm/v1/worker/cp_utils.py +++ b/vllm/v1/worker/cp_utils.py @@ -1,15 +1,23 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace from typing import TYPE_CHECKING, Any, cast +import torch + from vllm.config import VllmConfig, get_layers_from_vllm_config from vllm.distributed import get_dcp_group, get_pcp_group +from vllm.logger import init_logger +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.attention.backends.utils import split_decodes_prefills_and_extends if TYPE_CHECKING: from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase else: AttentionLayerBase = object +logger = init_logger(__name__) + def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None: pcp_size = vllm_config.parallel_config.prefill_context_parallel_size @@ -56,3 +64,232 @@ def get_total_cp_world_size(): # DCP might not be initialized in testing dcp_world_size = 1 return dcp_world_size * pcp_world_size + + +def get_dcp_dummy_context_len( + dcp_world_size: int, + cp_kv_cache_interleave_size: int, + has_kv_cache_config: bool, + create_mixed_batch: bool, + is_graph_capturing: bool, + uniform_decode: bool, +) -> int: + if ( + dcp_world_size <= 1 + or not has_kv_cache_config + or not (create_mixed_batch or (is_graph_capturing and uniform_decode)) + ): + return 0 + return dcp_world_size * cp_kv_cache_interleave_size + + +def prepare_dcp_dummy_context_metadata( + *, + input_batch: Any, + kv_cache_config: Any, + query_pos: Any, + positions: torch.Tensor, + query_start_loc: Any, + num_reqs: int, + num_tokens_unpadded: int, + dcp_dummy_context_len: int, +) -> None: + """Populate valid fake KV metadata for DCP CUDA graph warmup/capture.""" + if dcp_dummy_context_len == 0: + return + + # DCP graph warmup may exercise context attention, so block-table entries + # must point at allocated KV blocks. + assert kv_cache_config is not None + max_valid_block_id = kv_cache_config.num_blocks - 1 + assert max_valid_block_id > 0 + for blk_table in input_batch.block_table.block_tables: + max_row_blocks = ( + blk_table.max_num_blocks_per_req // blk_table.blocks_per_kv_block + ) + block_ids = [ + (block_idx % max_valid_block_id) + 1 for block_idx in range(max_row_blocks) + ] + for req_idx in range(num_reqs): + blk_table.add_row(block_ids, req_idx) + blk_table.commit_block_table(num_reqs) + + query_pos.copy_to_gpu(num_tokens_unpadded) + positions[:num_tokens_unpadded] = ( + query_pos.gpu[:num_tokens_unpadded] + dcp_dummy_context_len + ) + input_batch.block_table.compute_slot_mapping( + num_reqs, + query_start_loc.gpu[: num_reqs + 1], + positions[:num_tokens_unpadded], + ) + + +def should_skip_dcp_context_attention(context_kv_lens_cpu: torch.Tensor) -> bool: + """Whether DCP context attention can be skipped for this batch. + + Must be computed from rank-invariant inputs only (the global context + lengths, NOT this rank's local share from get_dcp_local_seq_lens): the + non-skip path in _forward_with_dcp issues DCP collectives (query + all-gather + LSE combine), so every DCP rank must take the same branch. + A rank can hold zero local context tokens while other ranks still hold + context for the same batch. + """ + return int(context_kv_lens_cpu.max().item()) == 0 + + +def split_dcp_context_queries( + query_start_loc: torch.Tensor, + seq_lens_cpu_upper_bound: torch.Tensor | None, + max_query_len: int, + num_actual_tokens: int, +) -> tuple[int, int, int, int]: + """Split reordered DCP context queries into decode and extend regions.""" + num_reqs = query_start_loc.shape[0] - 1 + if max_query_len <= 1: + return num_reqs, 0, num_actual_tokens, 0 + if seq_lens_cpu_upper_bound is None: + return 0, num_reqs, 0, num_actual_tokens + + common_attn_metadata = cast( + CommonAttentionMetadata, + SimpleNamespace( + max_query_len=max_query_len, + num_reqs=num_reqs, + num_actual_tokens=num_actual_tokens, + query_start_loc_cpu=query_start_loc, + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, + is_prefilling=None, + ), + ) + ( + num_decodes, + num_extends, + _num_prefills, + num_decode_tokens, + num_extend_tokens, + _num_prefill_tokens, + ) = split_decodes_prefills_and_extends(common_attn_metadata) + return num_decodes, num_extends, num_decode_tokens, num_extend_tokens + + +def should_split_fa2_dcp_context_attention( + fa_version: int | None, + max_query_len: int, + num_reqs: int, + num_decode_reqs: int, + num_context_prefill_reqs: int, +) -> bool: + num_prefills = num_reqs - num_decode_reqs + # TODO: Remove this FA2-only DCP compatibility path once FA4 supports + # the Qwen3.5 head_size=256 shape on Blackwell and can be used here. + # FA2 paged-varlen context attention can fail for DCP mixed batches when + # decode rows, context-bearing extend rows, and zero-context pure prefill + # rows are submitted together. + return ( + fa_version == 2 + and max_query_len > 1 + and num_prefills > 0 + and (num_decode_reqs > 0 or num_context_prefill_reqs < num_prefills) + ) + + +def run_split_fa2_dcp_context_attention( + flash_attn_varlen_func: Any, + query_across_dcp: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + dcp_context_out: torch.Tensor, + cu_seqlens_q: torch.Tensor, + max_seqlen_q: int, + dcp_context_kv_lens: torch.Tensor, + max_dcp_context_kv_len: int, + softmax_scale: float, + alibi_slopes: torch.Tensor | None, + sliding_window_size: list[int] | None, + block_table: torch.Tensor, + softcap: float, + fa_version: int, + q_descale: torch.Tensor | None, + k_descale: torch.Tensor | None, + v_descale: torch.Tensor | None, + max_num_splits: int, + num_heads: int, + dcp_world_size: int, + num_decode_reqs: int, + num_context_prefill_reqs: int, + num_decode_tokens: int, + num_context_prefill_tokens: int, +) -> tuple[torch.Tensor, torch.Tensor]: + dcp_context_out.zero_() + context_lse = torch.full( + (num_heads * dcp_world_size, query_across_dcp.shape[0]), + -torch.inf, + dtype=torch.float32, + device=query_across_dcp.device, + ) + + if num_decode_tokens > 0: + _, decode_context_lse = flash_attn_varlen_func( + q=query_across_dcp[:num_decode_tokens], + k=key_cache, + v=value_cache, + out=dcp_context_out[:num_decode_tokens], + cu_seqlens_q=cu_seqlens_q[: num_decode_reqs + 1], + max_seqlen_q=1, + seqused_k=dcp_context_kv_lens[:num_decode_reqs], + max_seqlen_k=max_dcp_context_kv_len, + softmax_scale=softmax_scale, + causal=False, + alibi_slopes=alibi_slopes, + window_size=sliding_window_size, + block_table=block_table[:num_decode_reqs], + softcap=softcap, + return_softmax_lse=True, + scheduler_metadata=None, + fa_version=fa_version, + q_descale=q_descale[:num_decode_reqs] if q_descale is not None else None, + k_descale=k_descale[:num_decode_reqs] if k_descale is not None else None, + v_descale=v_descale[:num_decode_reqs] if v_descale is not None else None, + num_splits=max_num_splits, + ) + context_lse[:, :num_decode_tokens] = decode_context_lse + + if num_context_prefill_tokens > 0: + prefill_start = num_decode_tokens + prefill_end = prefill_start + num_context_prefill_tokens + prefill_query_start_loc = ( + cu_seqlens_q[ + num_decode_reqs : num_decode_reqs + num_context_prefill_reqs + 1 + ] + - num_decode_tokens + ) + prefill_req_slice = slice( + num_decode_reqs, num_decode_reqs + num_context_prefill_reqs + ) + _, prefill_context_lse = flash_attn_varlen_func( + q=query_across_dcp[prefill_start:prefill_end], + k=key_cache, + v=value_cache, + out=dcp_context_out[prefill_start:prefill_end], + cu_seqlens_q=prefill_query_start_loc, + max_seqlen_q=max_seqlen_q, + seqused_k=dcp_context_kv_lens[prefill_req_slice], + max_seqlen_k=max_dcp_context_kv_len, + softmax_scale=softmax_scale, + causal=False, + alibi_slopes=alibi_slopes, + window_size=sliding_window_size, + block_table=block_table[prefill_req_slice], + softcap=softcap, + return_softmax_lse=True, + scheduler_metadata=None, + fa_version=fa_version, + q_descale=q_descale[prefill_req_slice] if q_descale is not None else None, + k_descale=k_descale[prefill_req_slice] if k_descale is not None else None, + v_descale=v_descale[prefill_req_slice] if v_descale is not None else None, + num_splits=max_num_splits, + ) + context_lse[:, prefill_start:prefill_end] = prefill_context_lse + + return dcp_context_out, context_lse diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py index 28d1a04b780c..bb4fb3fe453c 100644 --- a/vllm/v1/worker/gpu_input_batch.py +++ b/vllm/v1/worker/gpu_input_batch.py @@ -28,7 +28,7 @@ maybe_create_thinking_budget_state_holder, ) from vllm.v1.utils import copy_slice -from vllm.v1.worker.block_table import MultiGroupBlockTable +from vllm.v1.worker.block_table import MultiGroupBlockTable, SlotMappingMode @dataclass @@ -99,13 +99,14 @@ def __init__( vocab_size: int, block_sizes: list[int], # The block_size of each kv cache group kernel_block_sizes: list[int], - max_num_blocks_per_req: list[int] | None = None, + max_num_blocks_per_req: list[int], logitsprocs: LogitsProcessors | None = None, logitsprocs_need_output_token_ids: bool = False, num_spec_tokens: int = 0, is_pooling_model: bool = False, cp_kv_cache_interleave_size: int = 1, reasoning_config: ReasoningConfig | None = None, + slot_mapping_modes: list[SlotMappingMode] | None = None, ): self.thinking_budget_state_holder = maybe_create_thinking_budget_state_holder( reasoning_config, @@ -171,7 +172,6 @@ def __init__( # Block table. self.block_table = MultiGroupBlockTable( max_num_reqs=max_num_reqs, - max_model_len=max_model_len, max_num_batched_tokens=max_num_batched_tokens, pin_memory=PIN_MEMORY, device=device, @@ -179,6 +179,7 @@ def __init__( kernel_block_sizes=kernel_block_sizes, max_num_blocks=max_num_blocks_per_req, cp_kv_cache_interleave_size=cp_kv_cache_interleave_size, + slot_mapping_modes=slot_mapping_modes, ) # Sampling-related. diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 38500ab0514d..efb302b90f47 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -153,10 +153,12 @@ KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, + KVCacheSpecKind, KVQuantMode, MambaSpec, SlidingWindowSpec, UniformTypeKVCacheSpecs, + get_kv_cache_spec_kind, ) from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry from vllm.v1.outputs import ( @@ -201,9 +203,11 @@ from vllm.v1.structured_output.utils import apply_grammar_bitmask from vllm.v1.utils import CpuGpuBuffer, record_function_or_nullcontext from vllm.v1.worker import mamba_utils +from vllm.v1.worker.block_table import SlotMappingMode from vllm.v1.worker.cp_utils import ( check_attention_cp_compatibility, - get_total_cp_world_size, + get_dcp_dummy_context_len, + prepare_dcp_dummy_context_metadata, ) from vllm.v1.worker.dp_utils import coordinate_batch_across_dp from vllm.v1.worker.ec_connector_model_runner_mixin import ECConnectorModelRunnerMixin @@ -681,8 +685,13 @@ def __init__( placeholder_block_size = ( self.cache_config.block_size or CacheConfig.DEFAULT_BLOCK_SIZE ) + placeholder_max_num_blocks = cdiv( + max(self.max_model_len, self.max_encoder_len), placeholder_block_size + ) self._init_block_sizes = [placeholder_block_size] self._init_kernel_block_sizes = [placeholder_block_size] + self._init_max_num_blocks = [placeholder_max_num_blocks] + self._init_slot_mapping_modes = [SlotMappingMode.TOKEN_TO_KV_SLOT] self.input_batch = InputBatch( max_num_reqs=self.max_num_reqs, # We need to use the encoder length for encoder-decoder @@ -693,6 +702,7 @@ def __init__( vocab_size=self.model_config.get_vocab_size(), block_sizes=[placeholder_block_size], kernel_block_sizes=[placeholder_block_size], + max_num_blocks_per_req=[placeholder_max_num_blocks], num_spec_tokens=self.num_spec_tokens, logitsprocs=build_logitsprocs( self.vllm_config, @@ -5858,6 +5868,14 @@ def _dummy_run( num_reqs_padded = ( batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs ) + dcp_dummy_context_len = get_dcp_dummy_context_len( + self.dcp_world_size, + self.parallel_config.cp_kv_cache_interleave_size, + hasattr(self, "kv_cache_config"), + create_mixed_batch, + is_graph_capturing, + uniform_decode, + ) ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices( should_ubatch, num_scheduled_tokens, @@ -5900,10 +5918,19 @@ def _dummy_run( # In the mixed batch mode (used for FI warmup), we use # shorter sequence lengths to run faster. # TODO(luka) better system for describing dummy batches - seq_lens = torch.tensor( # type: ignore[assignment] - [1] * num_decode_tokens + [num_prefill_tokens + 1], - dtype=torch.int, - ) + if dcp_dummy_context_len > 0: + seq_lens = torch.tensor( # type: ignore[assignment] + [1 + dcp_dummy_context_len] * num_decode_tokens + + [num_prefill_tokens + dcp_dummy_context_len], + dtype=torch.int, + ) + else: + seq_lens = torch.tensor( # type: ignore[assignment] + [1] * num_decode_tokens + [num_prefill_tokens + 1], + dtype=torch.int, + ) + elif dcp_dummy_context_len > 0: + seq_lens = max_query_len + dcp_dummy_context_len # type: ignore[assignment] else: seq_lens = max_query_len # type: ignore[assignment] self.optimistic_seq_lens_cpu[:num_reqs] = seq_lens @@ -5919,6 +5946,17 @@ def _dummy_run( ) self.query_start_loc.copy_to_gpu() + prepare_dcp_dummy_context_metadata( + input_batch=self.input_batch, + kv_cache_config=getattr(self, "kv_cache_config", None), + query_pos=self.query_pos, + positions=self.positions, + query_start_loc=self.query_start_loc, + num_reqs=num_reqs, + num_tokens_unpadded=num_tokens_unpadded, + dcp_dummy_context_len=dcp_dummy_context_len, + ) + # Sync block table CPU->GPU so cleared rows from # remove_request() are visible to the attention metadata # builder. Without this, stale block IDs from finished @@ -7029,29 +7067,34 @@ def may_reinitialize_input_batch( """ block_sizes = [] max_num_blocks = [] + slot_mapping_modes = [] max_model_len = max(self.max_model_len, self.max_encoder_len) for kv_cache_group in kv_cache_config.kv_cache_groups: - if isinstance(kv_cache_group.kv_cache_spec, EncoderOnlyAttentionSpec): + kv_cache_spec = kv_cache_group.kv_cache_spec + kv_cache_spec_kind = get_kv_cache_spec_kind(kv_cache_spec) + if kv_cache_spec_kind == KVCacheSpecKind.ENCODER_ONLY_ATTENTION: continue - block_size = kv_cache_group.kv_cache_spec.block_size + block_size = kv_cache_spec.block_size block_sizes.append(block_size) - max_num_blocks_per_req = cdiv( - max_model_len, block_size * get_total_cp_world_size() - ) - if isinstance(kv_cache_group.kv_cache_spec, MambaSpec): - max_num_blocks_per_req = ( - max_num_blocks_per_req - if self.cache_config.enable_prefix_caching - else 1 - ) + kv_cache_group.kv_cache_spec.num_speculative_blocks + if kv_cache_spec_kind == KVCacheSpecKind.MAMBA: + slot_mapping_modes.append(SlotMappingMode.NONE) + else: + slot_mapping_modes.append(SlotMappingMode.TOKEN_TO_KV_SLOT) + max_num_blocks_per_req = kv_cache_spec.max_num_blocks_per_req( + self.vllm_config, max_model_len + ) max_num_blocks.append(max_num_blocks_per_req) if ( block_sizes != self._init_block_sizes or kernel_block_sizes != self._init_kernel_block_sizes + or max_num_blocks != self._init_max_num_blocks + or slot_mapping_modes != self._init_slot_mapping_modes ): self._init_block_sizes = block_sizes self._init_kernel_block_sizes = kernel_block_sizes + self._init_max_num_blocks = max_num_blocks + self._init_slot_mapping_modes = slot_mapping_modes self.input_batch = InputBatch( max_num_reqs=self.max_num_reqs, max_model_len=max_model_len, @@ -7065,7 +7108,9 @@ def may_reinitialize_input_batch( logitsprocs=self.input_batch.logitsprocs, logitsprocs_need_output_token_ids=self.input_batch.logitsprocs_need_output_token_ids, is_pooling_model=self.is_pooling_model, + cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size, reasoning_config=self.vllm_config.reasoning_config, + slot_mapping_modes=slot_mapping_modes, ) assert self._init_block_sizes == block_sizes, ( diff --git a/vllm/v1/worker/tpu_input_batch.py b/vllm/v1/worker/tpu_input_batch.py index 3758a73ee496..1396c8ad9d5a 100644 --- a/vllm/v1/worker/tpu_input_batch.py +++ b/vllm/v1/worker/tpu_input_batch.py @@ -29,6 +29,7 @@ def __init__( vocab_size: int, block_sizes: list[int], # The block_size of each kv cache group kernel_block_sizes: list[int], + max_num_blocks_per_req: list[int], ): self.max_num_reqs = max_num_reqs self.max_model_len = max_model_len @@ -64,12 +65,12 @@ def __init__( # Block table. self.block_table = MultiGroupBlockTable( max_num_reqs=max_num_reqs, - max_model_len=max_model_len, max_num_batched_tokens=max_num_batched_tokens, pin_memory=pin_memory, device=device, block_sizes=block_sizes, kernel_block_sizes=kernel_block_sizes, + max_num_blocks=max_num_blocks_per_req, ) # Sampling-related. From e5588e49bc2642670116664a7fc4096e27adb179 Mon Sep 17 00:00:00 2001 From: GongLei-HW <1327185943@qq.com> Date: Fri, 10 Jul 2026 13:46:54 +0800 Subject: [PATCH 0008/1526] [Core][KV events] Report prefix-cache-reused blocks in full report mode (#45261) Signed-off-by: Lei Gong Co-authored-by: Lei Gong Co-authored-by: Claude --- tests/v1/core/test_prefix_caching.py | 119 +++++++++++++++++++++++- vllm/v1/core/block_pool.py | 131 ++++++++++++++++++++++----- vllm/v1/core/kv_cache_manager.py | 21 +++++ vllm/v1/core/kv_cache_utils.py | 19 ++++ vllm/v1/request.py | 5 + 5 files changed, 270 insertions(+), 25 deletions(-) diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 59260a499ef4..03d61ec299f1 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -12,7 +12,12 @@ import vllm.v1.core.kv_cache_manager as kv_cache_manager import vllm.v1.core.kv_cache_utils as kv_cache_utils -from vllm.distributed.kv_events import AllBlocksCleared, BlockRemoved, BlockStored +from vllm.distributed.kv_events import ( + MEDIUM_GPU, + AllBlocksCleared, + BlockRemoved, + BlockStored, +) from vllm.lora.request import LoRARequest from vllm.multimodal.inputs import ( MultiModalFeatureSpec, @@ -2454,6 +2459,118 @@ def test_block_removed_event_group_idx(group_id: int): assert event.group_idx == group_id +def test_emit_cached_block_events(): + """emit_cached_block_events emits one BlockStored for already-cached + (reused) prefix blocks, carrying the correct group_idx / + parent_block_hash / token_ids, and without mutating block state.""" + block_size = 4 + num_cached_blocks = 3 + kv_cache_group_id = 1 + num_tokens = block_size * 4 # 4 full blocks; reuse the first 3 + + pool = BlockPool( + num_gpu_blocks=8, + enable_caching=True, + hash_block_size=block_size, + enable_kv_cache_events=True, + ) + + req = make_request( + "req_emit_cached", + prompt_token_ids=list(range(num_tokens)), + block_size=block_size, + hash_fn=sha256, + ) + assert len(req.block_hashes) >= num_cached_blocks + + # Snapshot block state to prove emit_cached_block_events does not mutate it. + free_before = pool.get_num_free_blocks() + assert len(pool.cached_block_hash_to_block) == 0 + + pool.emit_cached_block_events( + request=req, + num_cached_blocks=num_cached_blocks, + block_size=block_size, + kv_cache_group_id=kv_cache_group_id, + ) + + # No block-state mutation: nothing allocated, nothing inserted into the + # prefix-cache map. + assert pool.get_num_free_blocks() == free_before + assert len(pool.cached_block_hash_to_block) == 0 + + events = pool.take_events() + assert len(events) == 1 + event = events[0] + assert isinstance(event, BlockStored) + + expected_hashes = [ + kv_cache_utils.maybe_convert_block_hash(req.block_hashes[i]) + for i in range(num_cached_blocks) + ] + assert event.block_hashes == expected_hashes + # Reused blocks start from block 0, so there is no parent block hash. + assert event.parent_block_hash is None + assert event.token_ids == list(req.all_token_ids[: num_cached_blocks * block_size]) + assert event.group_idx == kv_cache_group_id + assert event.block_size == block_size + assert event.medium == MEDIUM_GPU + assert event.lora_id is None + assert event.lora_name is None + + +def test_emit_cached_block_events_disabled(): + """No events are emitted when enable_kv_cache_events is False.""" + block_size = 4 + pool = BlockPool( + num_gpu_blocks=8, + enable_caching=True, + hash_block_size=block_size, + enable_kv_cache_events=False, + ) + req = make_request( + "req_emit_disabled", + prompt_token_ids=list(range(block_size * 4)), + block_size=block_size, + hash_fn=sha256, + ) + + pool.emit_cached_block_events( + request=req, + num_cached_blocks=3, + block_size=block_size, + kv_cache_group_id=0, + ) + + assert pool.take_events() == [] + + +def test_emit_cached_block_events_zero_cached(): + """No events are emitted when num_cached_blocks == 0.""" + block_size = 4 + pool = BlockPool( + num_gpu_blocks=8, + enable_caching=True, + hash_block_size=block_size, + enable_kv_cache_events=True, + ) + req = make_request( + "req_emit_zero", + prompt_token_ids=list(range(block_size * 4)), + block_size=block_size, + hash_fn=sha256, + ) + + pool.emit_cached_block_events( + request=req, + num_cached_blocks=0, + block_size=block_size, + kv_cache_group_id=0, + ) + + assert pool.take_events() == [] + + def test_eagle_enabled_removes_last_block(): """Verify Eagle does NOT remove blocks when request length is divisible by block size.""" diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py index 81ac05f36589..bc8f2d87adb9 100644 --- a/vllm/v1/core/block_pool.py +++ b/vllm/v1/core/block_pool.py @@ -14,8 +14,6 @@ from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector from vllm.v1.core.kv_cache_utils import ( BlockHash, - BlockHashList, - BlockHashListWithBlockSize, BlockHashWithGroupId, ExternalBlockHash, FreeKVCacheBlockQueue, @@ -25,6 +23,7 @@ get_group_id, make_block_hash_with_group_id, maybe_convert_block_hash, + resolve_block_hashes, ) from vllm.v1.request import Request @@ -261,17 +260,7 @@ def cache_full_blocks( return new_full_blocks = blocks[num_cached_blocks:num_full_blocks] assert block_mask is None or len(block_mask) == len(new_full_blocks) - if block_size == self.hash_block_size: - # Common case. - block_hashes: BlockHashList = request.block_hashes - else: - # block_size is a multiple of hash_block_size. This happens when - # different KV cache groups have different block sizes. - assert block_size % self.hash_block_size == 0 - block_hashes = BlockHashListWithBlockSize( - request.block_hashes, self.hash_block_size, block_size - ) - assert len(block_hashes) >= num_full_blocks + block_hashes = resolve_block_hashes(request, self.hash_block_size, block_size) new_block_hashes = block_hashes[num_cached_blocks:] new_hashes: list[ExternalBlockHash] | None = ( @@ -338,23 +327,117 @@ def cache_full_blocks( extra_keys_list.append(extra_keys) self.kv_event_queue.append( - BlockStored( + self._build_block_stored_event( + request, block_hashes=new_hashes, parent_block_hash=parent_block_hash, - token_ids=request.all_token_ids[start_token_idx:end_token_idx], + start_token_idx=start_token_idx, + end_token_idx=end_token_idx, block_size=block_size, - lora_id=request.lora_request.adapter_id - if request.lora_request - else None, - medium=MEDIUM_GPU, - lora_name=request.lora_request.name - if request.lora_request - else None, - extra_keys=extra_keys_list if extra_keys_list else None, - group_idx=kv_cache_group_id, + kv_cache_group_id=kv_cache_group_id, + extra_keys_list=extra_keys_list, ) ) + def _build_block_stored_event( + self, + request: Request, + block_hashes: list[ExternalBlockHash] | None, + parent_block_hash: ExternalBlockHash | None, + start_token_idx: int, + end_token_idx: int, + block_size: int, + kv_cache_group_id: int, + extra_keys_list: list[tuple[Any, ...] | None], + ) -> BlockStored: + """Build a ``BlockStored`` KV event for ``request``. + + Shared by ``cache_full_blocks`` (newly cached blocks) and + ``emit_cached_block_events`` (prefix-cache-reused blocks) so both emit + identical event shapes for downstream consumers. + """ + return BlockStored( + block_hashes=block_hashes, + parent_block_hash=parent_block_hash, + token_ids=request.all_token_ids[start_token_idx:end_token_idx], + block_size=block_size, + lora_id=request.lora_request.adapter_id if request.lora_request else None, + medium=MEDIUM_GPU, + lora_name=request.lora_request.name if request.lora_request else None, + extra_keys=extra_keys_list if extra_keys_list else None, + group_idx=kv_cache_group_id, + ) + + def emit_cached_block_events( + self, + request: Request, + num_cached_blocks: int, + block_size: int, + kv_cache_group_id: int, + ) -> None: + """Generate BlockStored events for blocks reused from prefix cache. + + Unlike cache_full_blocks(), this does NOT modify block state — + the blocks are already cached. It only generates events so that + external consumers (e.g. gateway) can learn about reused blocks. + + Args: + request: The request whose prefix cache blocks were reused. + num_cached_blocks: Number of blocks that were cache hits. + block_size: Number of tokens per block. + kv_cache_group_id: The KV cache group ID. + """ + if not self.enable_kv_cache_events or num_cached_blocks == 0: + return + + block_hashes = resolve_block_hashes(request, self.hash_block_size, block_size) + + # Collect external hashes and extra_keys for cached blocks. + cached_hashes: list[ExternalBlockHash] = [] + extra_keys_list: list[tuple[Any, ...] | None] = [] + curr_mm_idx = 0 + for i in range(num_cached_blocks): + block_start = i * block_size + block_end = block_start + block_size + cached_hashes.append(maybe_convert_block_hash(block_hashes[i])) + extra_keys, curr_mm_idx = generate_block_hash_extra_keys( + request, block_start, block_end, curr_mm_idx + ) + extra_keys_list.append(extra_keys) + + if not cached_hashes: + return + + # Prefix-cache hits always form a contiguous prefix starting at block 0, + # so the first (and thus the whole group's) parent block hash is None. + parent_block_hash: ExternalBlockHash | None = None + start_token_idx = 0 + end_token_idx = num_cached_blocks * block_size + + logger.debug( + "EmitCachedBlock event: block_size=%d, " + "num_cached_blocks=%d, parent_block_hash=%s, " + "token_ids_len=%d, group_idx=%s", + block_size, + num_cached_blocks, + parent_block_hash, + len(request.all_token_ids[start_token_idx:end_token_idx]), + kv_cache_group_id, + ) + + self.kv_event_queue.append( + self._build_block_stored_event( + request, + block_hashes=cached_hashes, + parent_block_hash=parent_block_hash, + start_token_idx=start_token_idx, + end_token_idx=end_token_idx, + block_size=block_size, + kv_cache_group_id=kv_cache_group_id, + extra_keys_list=extra_keys_list, + ) + ) + def cache_partial_block( self, request: Request, diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 4b62915edeff..4d5d613866ad 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -136,6 +136,7 @@ def __init__( max_in_flight_tokens = max_model_len self.enable_caching = enable_caching + self.enable_kv_cache_events = enable_kv_cache_events self.use_eagle = use_eagle self.log_stats = log_stats self.metrics_collector = metrics_collector @@ -235,6 +236,26 @@ def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int]: ) ) + # When kv_cache_report_mode is "full", emit BlockStored events + # for the reused prefix cache blocks so that external consumers + # (e.g. gateway) can learn about them. + if ( + num_new_computed_tokens > 0 + and self.enable_kv_cache_events + and getattr(request, "kv_cache_report_mode", "incremental") == "full" + ): + for group_idx, group_blocks in enumerate(computed_blocks): + num_blocks = len(group_blocks) + if num_blocks > 0: + group = self.kv_cache_config.kv_cache_groups[group_idx] + block_size = group.kv_cache_spec.block_size + self.block_pool.emit_cached_block_events( + request, + num_blocks, + block_size, + group_idx, + ) + if self.log_stats: assert self.prefix_cache_stats is not None self.prefix_cache_stats.record( diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 6e9530e381bf..27af83c9a516 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -2249,3 +2249,22 @@ def _get_value_at(self, idx: int) -> BlockHash: BlockHashList = list[BlockHash] | BlockHashListWithBlockSize + + +def resolve_block_hashes( + request: Request, + hash_block_size: int, + block_size: int, +) -> BlockHashList: + """Resolve the block-hash view for ``request`` at ``block_size``. + + When ``block_size`` equals ``hash_block_size``, reuse the request's + precomputed ``block_hashes`` directly; otherwise recalculate at + ``block_size`` granularity (``block_size`` must be a multiple of + ``hash_block_size``, which happens when KV cache groups differ in + block size). + """ + if block_size == hash_block_size: + return request.block_hashes + assert block_size % hash_block_size == 0 + return BlockHashListWithBlockSize(request.block_hashes, hash_block_size, block_size) diff --git a/vllm/v1/request.py b/vllm/v1/request.py index e9946a7f76b8..058d498d621d 100644 --- a/vllm/v1/request.py +++ b/vllm/v1/request.py @@ -115,6 +115,11 @@ def __init__( self.kv_transfer_params = sampling_params.extra_args.get( "kv_transfer_params" ) + self.kv_cache_report_mode = sampling_params.extra_args.get( + "kv_cache_report_mode", "incremental" + ) + else: + self.kv_cache_report_mode = "incremental" else: raise ValueError("sampling_params and pooling_params can't both be unset") From 5715fde12c1e28eccd08a2394a114339b12a96c1 Mon Sep 17 00:00:00 2001 From: alberto Date: Fri, 10 Jul 2026 08:34:02 +0100 Subject: [PATCH 0009/1526] [Feature][Parser] Support include_reasoning param for non-Harmony models (#44301) Signed-off-by: Alberto Perdomo Co-authored-by: Chauncey --- docs/features/reasoning_outputs.md | 63 +++ .../chat_completion/test_include_reasoning.py | 158 ++++++ tests/parser/engine/conftest.py | 1 + tests/parser/engine/test_nemotron_v3.py | 1 + tests/parser/engine/test_parser_engine.py | 1 + tests/parser/test_include_reasoning.py | 456 ++++++++++++++++++ .../openai/chat_completion/serving.py | 50 +- vllm/entrypoints/openai/responses/context.py | 6 +- vllm/entrypoints/openai/responses/protocol.py | 9 + vllm/entrypoints/openai/responses/serving.py | 7 + vllm/parser/abstract_parser.py | 19 +- vllm/parser/engine/parser_engine.py | 10 +- vllm/parser/harmony.py | 10 + 13 files changed, 765 insertions(+), 26 deletions(-) create mode 100644 tests/entrypoints/openai/chat_completion/test_include_reasoning.py create mode 100644 tests/parser/test_include_reasoning.py diff --git a/docs/features/reasoning_outputs.md b/docs/features/reasoning_outputs.md index 50a58b8b3e37..9495c4672ce0 100644 --- a/docs/features/reasoning_outputs.md +++ b/docs/features/reasoning_outputs.md @@ -351,6 +351,69 @@ print(response.choices[0].message.reasoning) print(response.choices[0].message.content) ``` +## Suppressing Reasoning Output + +You can suppress reasoning content from API responses using the `include_reasoning` parameter. When set to `false`, reasoning tokens are still generated (so model quality is unaffected) but excluded from the response. This reduces network traffic without changing inference behavior. + +The parameter is supported in both the Chat Completions API and the Responses API, for streaming and non-streaming requests. + +When `include_reasoning=false`, vLLM also suppresses per-token metadata (logprobs and token IDs) to prevent leaking reasoning content through decoded token text in logprob entries or raw token IDs. + +### Chat Completions API + +```python +from openai import OpenAI + +client = OpenAI(base_url="http://localhost:8000/v1", api_key="EMPTY") +model = client.models.list().data[0].id + +# Reasoning is included by default (include_reasoning=True) +response = client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": "What is 15 * 37?"}], + extra_body={"include_reasoning": False}, +) + +msg = response.choices[0].message +assert msg.content # Content is still present +assert not getattr(msg, "reasoning", None) # Reasoning is suppressed +``` + +Streaming works the same way, reasoning deltas are omitted from chunks: + +```python +stream = client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": "What is 15 * 37?"}], + stream=True, + extra_body={"include_reasoning": False}, +) + +for chunk in stream: + delta = chunk.choices[0].delta + # delta.reasoning will always be None + if delta.content: + print(delta.content, end="", flush=True) +``` + +### Responses API + +```python +from openai import OpenAI + +client = OpenAI(base_url="http://localhost:8000/v1", api_key="EMPTY") + +response = client.responses.create( + model=client.models.list().data[0].id, + input="What is 15 * 37?", + include_reasoning=False, +) + +# No "reasoning" items in output +types = [item.type for item in response.output] +assert "reasoning" not in types +``` + ## Limitations - The reasoning content is only available for online serving's chat completion endpoint (`/v1/chat/completions`), Anthropic Messages API (`/v1/messages`) and the Responses API (`/v1/responses`). diff --git a/tests/entrypoints/openai/chat_completion/test_include_reasoning.py b/tests/entrypoints/openai/chat_completion/test_include_reasoning.py new file mode 100644 index 000000000000..50261cfff804 --- /dev/null +++ b/tests/entrypoints/openai/chat_completion/test_include_reasoning.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""E2E tests for ``include_reasoning`` with non-Harmony reasoning models. + +Verifies that reasoning content is included by default and suppressed +when ``include_reasoning=False``, for both streaming and non-streaming +Chat Completions. +""" + +import openai +import pytest +import pytest_asyncio + +from tests.utils import RemoteOpenAIServer + +MODEL_NAME = "Qwen/Qwen3-0.6B" +MESSAGES = [{"role": "user", "content": "What is 1+1? Be concise."}] + + +@pytest.fixture(scope="module") +def server(): + args = [ + "--reasoning-parser", + "qwen3", + "--max-model-len", + "2048", + "--enforce-eager", + "--gpu-memory-utilization", + "0.4", + ] + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: + yield remote_server + + +@pytest_asyncio.fixture +async def client(server): + async with server.get_async_client() as async_client: + yield async_client + + +@pytest.mark.asyncio +async def test_include_reasoning_true_non_streaming(client: openai.AsyncOpenAI): + """Default: reasoning content appears in non-streaming response.""" + response = await client.chat.completions.create( + model=MODEL_NAME, + messages=MESSAGES, + max_tokens=200, + extra_body={"include_reasoning": True}, + ) + + msg = response.choices[0].message + reasoning = getattr(msg, "reasoning", None) or getattr( + msg, "reasoning_content", None + ) + assert reasoning, "Expected reasoning content when include_reasoning=True" + assert msg.content, "Expected content in response" + + +@pytest.mark.asyncio +async def test_include_reasoning_false_non_streaming(client: openai.AsyncOpenAI): + """Reasoning content is suppressed when include_reasoning=False.""" + response = await client.chat.completions.create( + model=MODEL_NAME, + messages=MESSAGES, + max_tokens=200, + extra_body={"include_reasoning": False}, + ) + + msg = response.choices[0].message + reasoning = getattr(msg, "reasoning", None) or getattr( + msg, "reasoning_content", None + ) + assert not reasoning, ( + f"Expected no reasoning when include_reasoning=False, got: {reasoning}" + ) + assert msg.content, "Expected content in response even without reasoning" + + +@pytest.mark.asyncio +async def test_include_reasoning_true_streaming(client: openai.AsyncOpenAI): + """Default: reasoning deltas appear in streaming response.""" + stream = await client.chat.completions.create( + model=MODEL_NAME, + messages=MESSAGES, + max_tokens=200, + stream=True, + extra_body={"include_reasoning": True}, + ) + + reasoning_parts = [] + content_parts = [] + async for chunk in stream: + delta = chunk.choices[0].delta if chunk.choices else None + if delta: + r = getattr(delta, "reasoning", None) or getattr( + delta, "reasoning_content", None + ) + if r: + reasoning_parts.append(r) + if delta.content: + content_parts.append(delta.content) + + reasoning_text = "".join(reasoning_parts) + content_text = "".join(content_parts) + + assert reasoning_text, "Expected reasoning deltas when include_reasoning=True" + assert content_text, "Expected content deltas in streaming response" + + +@pytest.mark.asyncio +async def test_include_reasoning_false_streaming(client: openai.AsyncOpenAI): + """Reasoning deltas are suppressed in streaming when include_reasoning=False.""" + stream = await client.chat.completions.create( + model=MODEL_NAME, + messages=MESSAGES, + max_tokens=200, + stream=True, + extra_body={"include_reasoning": False}, + ) + + reasoning_parts = [] + content_parts = [] + async for chunk in stream: + delta = chunk.choices[0].delta if chunk.choices else None + if delta: + r = getattr(delta, "reasoning", None) or getattr( + delta, "reasoning_content", None + ) + if r: + reasoning_parts.append(r) + if delta.content: + content_parts.append(delta.content) + + reasoning_text = "".join(reasoning_parts) + content_text = "".join(content_parts) + + assert not reasoning_text, ( + f"Expected no reasoning deltas when include_reasoning=False, " + f"got: {reasoning_text[:100]}" + ) + assert content_text, "Expected content deltas even without reasoning" + + +@pytest.mark.asyncio +async def test_default_includes_reasoning(client: openai.AsyncOpenAI): + """Without specifying include_reasoning, reasoning appears (default=True).""" + response = await client.chat.completions.create( + model=MODEL_NAME, + messages=MESSAGES, + max_tokens=200, + ) + + msg = response.choices[0].message + reasoning = getattr(msg, "reasoning", None) or getattr( + msg, "reasoning_content", None + ) + assert reasoning, "Expected reasoning content by default" diff --git a/tests/parser/engine/conftest.py b/tests/parser/engine/conftest.py index 157234c852aa..2522d146c895 100644 --- a/tests/parser/engine/conftest.py +++ b/tests/parser/engine/conftest.py @@ -49,4 +49,5 @@ def mock_request(): req = MagicMock(spec=ChatCompletionRequest) req.tools = [] req.tool_choice = "auto" + req.include_reasoning = True return req diff --git a/tests/parser/engine/test_nemotron_v3.py b/tests/parser/engine/test_nemotron_v3.py index 6aedcd1513b2..a7f166f9f3a9 100644 --- a/tests/parser/engine/test_nemotron_v3.py +++ b/tests/parser/engine/test_nemotron_v3.py @@ -43,6 +43,7 @@ def _make_request(**chat_template_kwargs): request = MagicMock(spec=ChatCompletionRequest) request.tools = [] request.tool_choice = "auto" + request.include_reasoning = True request.chat_template_kwargs = chat_template_kwargs or None return request diff --git a/tests/parser/engine/test_parser_engine.py b/tests/parser/engine/test_parser_engine.py index 635f51196a50..36258668215b 100644 --- a/tests/parser/engine/test_parser_engine.py +++ b/tests/parser/engine/test_parser_engine.py @@ -896,6 +896,7 @@ def _make_delegating_request(): req = MagicMock(spec=ChatCompletionRequest) req.tools = [] req.tool_choice = "auto" + req.include_reasoning = True return req diff --git a/tests/parser/test_include_reasoning.py b/tests/parser/test_include_reasoning.py new file mode 100644 index 000000000000..3d1893577ef0 --- /dev/null +++ b/tests/parser/test_include_reasoning.py @@ -0,0 +1,456 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for include_reasoning suppression in the unified Parser interface. + +Covers non-streaming (parser.parse() + build_response_output_items), +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 + + +class ThinkReasoningParser(BaseThinkingReasoningParser): + @property + def start_token(self) -> str: + return "" + + @property + def end_token(self) -> str: + return "" + + +MODEL_OUTPUT_REASONING_AND_CONTENT = ( + "let me think about thisThe answer is 42." +) + +MODEL_OUTPUT_REASONING_AND_TOOL = ( + "I need to call a tool" + '\n{"name": "get_weather", ' + '"arguments": {"city": "Dallas"}}\n' +) + +MODEL_OUTPUT_CONTENT_ONLY = "The answer is 42." + + +@pytest.fixture(scope="module") +def tokenizer(): + from vllm.tokenizers import get_tokenizer + + return get_tokenizer("Qwen/Qwen3-32B") + + +def make_responses_request(**kwargs) -> ResponsesRequest: + defaults = dict(model="test-model", input="test input") + defaults.update(kwargs) + return ResponsesRequest(**defaults) + + +def make_chat_request(**kwargs) -> ChatCompletionRequest: + defaults = dict( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + ) + defaults.update(kwargs) + return ChatCompletionRequest(**defaults) + + +def make_parser(tokenizer, reasoning=False, tool=False): + class TestParser(DelegatingParser): + reasoning_parser_cls = ThinkReasoningParser if reasoning else None + tool_parser_cls = Hermes2ProToolParser if tool else None + + return TestParser(tokenizer) + + +# ── Non-streaming: parser.parse() + build_response_output_items ────── + + +def parse_and_build(parser, request, model_output, enable_auto_tools=False): + """Mirrors the non-streaming path in _make_response_output_items / + ParsableContext.append_output(): parse → suppress reasoning → build items. + """ + from vllm.entrypoints.openai.responses.utils import ( + build_response_output_items, + ) + + reasoning, content, tool_calls = parser.parse( + model_output, request, enable_auto_tools=enable_auto_tools + ) + if not request.include_reasoning: + reasoning = None + return build_response_output_items( + reasoning=reasoning, + content=content, + tool_calls=tool_calls, + ) + + +class TestNonStreamingIncludeReasoning: + def test_include_reasoning_true_has_reasoning_item(self, tokenizer): + """Default: reasoning items appear in output.""" + parser = make_parser(tokenizer, reasoning=True) + request = make_responses_request(include_reasoning=True) + + outputs = parse_and_build(parser, request, MODEL_OUTPUT_REASONING_AND_CONTENT) + + types = [o.type for o in outputs] + assert "reasoning" in types + assert "message" in types + + def test_include_reasoning_false_no_reasoning_item(self, tokenizer): + """Reasoning item is suppressed when include_reasoning=False.""" + parser = make_parser(tokenizer, reasoning=True) + request = make_responses_request(include_reasoning=False) + + outputs = parse_and_build(parser, request, MODEL_OUTPUT_REASONING_AND_CONTENT) + + types = [o.type for o in outputs] + assert "reasoning" not in types + assert "message" in types + assert outputs[0].content[0].text == "The answer is 42." + + def test_include_reasoning_false_content_preserved(self, tokenizer): + """Content is extracted correctly even when reasoning is suppressed.""" + parser = make_parser(tokenizer, reasoning=True) + request = make_responses_request(include_reasoning=False) + + outputs = parse_and_build(parser, request, MODEL_OUTPUT_REASONING_AND_CONTENT) + + message = next(o for o in outputs if o.type == "message") + assert message.content[0].text == "The answer is 42." + + def test_include_reasoning_false_tool_calls_preserved(self, tokenizer): + """Tool calls still work when reasoning is suppressed.""" + parser = make_parser(tokenizer, reasoning=True, tool=True) + request = make_responses_request( + include_reasoning=False, + tools=[ + { + "type": "function", + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + } + ], + ) + + outputs = parse_and_build( + parser, + request, + MODEL_OUTPUT_REASONING_AND_TOOL, + enable_auto_tools=True, + ) + + types = [o.type for o in outputs] + assert "reasoning" not in types + assert "function_call" in types + fc = next(o for o in outputs if o.type == "function_call") + assert fc.name == "get_weather" + assert json.loads(fc.arguments) == {"city": "Dallas"} + + def test_no_reasoning_parser_include_false_is_noop(self, tokenizer): + """include_reasoning=False is harmless when no reasoning parser.""" + parser = make_parser(tokenizer, reasoning=False) + request = make_responses_request(include_reasoning=False) + + outputs = parse_and_build(parser, request, MODEL_OUTPUT_CONTENT_ONLY) + + assert len(outputs) == 1 + assert outputs[0].type == "message" + assert outputs[0].content[0].text == MODEL_OUTPUT_CONTENT_ONLY + + def test_default_include_reasoning_is_true(self, tokenizer): + """ResponsesRequest defaults to include_reasoning=True.""" + request = make_responses_request() + assert request.include_reasoning is True + + def test_include_reasoning_false_suppresses_all_reasoning(self, tokenizer): + """Reasoning is suppressed regardless of request type.""" + parser = make_parser(tokenizer, reasoning=True) + request = make_responses_request(include_reasoning=False) + + outputs = parse_and_build(parser, request, MODEL_OUTPUT_REASONING_AND_CONTENT) + + assert all(o.type != "reasoning" for o in outputs) + + +# ── Streaming: parse_delta ─────────────────────────────────────────── + + +def stream_text(parser, tokenizer, text, request, prompt_token_ids=None): + token_ids = tokenizer.encode(text, add_special_tokens=False) + results: list[DeltaMessage | None] = [] + for i, tid in enumerate(token_ids): + delta_text = tokenizer.decode([tid]) + is_last = i == len(token_ids) - 1 + result = parser.parse_delta( + delta_text, + [tid], + request, + prompt_token_ids=prompt_token_ids, + finished=is_last, + ) + prompt_token_ids = None + results.append(result) + return results + + +def collect_fields(results): + all_reasoning = "".join(r.reasoning for r in results if r and r.reasoning) + all_content = "".join(r.content for r in results if r and r.content) + all_tool_calls = [tc for r in results if r and r.tool_calls for tc in r.tool_calls] + return all_reasoning, all_content, all_tool_calls + + +class TestParseDeltaIncludeReasoning: + def test_streaming_include_true_emits_reasoning(self, tokenizer): + """With include_reasoning=True, reasoning deltas are emitted.""" + parser = make_parser(tokenizer, reasoning=True) + request = make_responses_request(include_reasoning=True) + + results = stream_text( + parser, + tokenizer, + MODEL_OUTPUT_REASONING_AND_CONTENT, + request, + prompt_token_ids=[], + ) + reasoning, content, _ = collect_fields(results) + + assert "let me think about this" in reasoning + assert "42" in content + + def test_streaming_include_false_suppresses_reasoning(self, tokenizer): + """With include_reasoning=False, no reasoning deltas are emitted.""" + parser = make_parser(tokenizer, reasoning=True) + request = make_responses_request(include_reasoning=False) + + results = stream_text( + parser, + tokenizer, + MODEL_OUTPUT_REASONING_AND_CONTENT, + request, + prompt_token_ids=[], + ) + reasoning, content, _ = collect_fields(results) + + assert reasoning == "" + assert "42" in content + + def test_streaming_include_false_content_still_works(self, tokenizer): + """Content is correctly extracted in streaming even with suppression.""" + parser = make_parser(tokenizer, reasoning=True) + request = make_responses_request(include_reasoning=False) + + results = stream_text( + parser, + tokenizer, + MODEL_OUTPUT_REASONING_AND_CONTENT, + request, + prompt_token_ids=[], + ) + _, content, _ = collect_fields(results) + + assert "The answer is 42" in content + + def test_streaming_include_false_tool_calls_preserved(self, tokenizer): + """Tool calls stream correctly when reasoning is suppressed.""" + parser = make_parser(tokenizer, reasoning=True, tool=True) + request = make_responses_request( + include_reasoning=False, + tools=[ + { + "type": "function", + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + } + ], + ) + + results = stream_text( + parser, + tokenizer, + MODEL_OUTPUT_REASONING_AND_TOOL, + request, + prompt_token_ids=[], + ) + reasoning, content, tool_calls = collect_fields(results) + + assert reasoning == "" + assert len(tool_calls) > 0 + assert tool_calls[0].function.name == "get_weather" + tool_args = "".join( + tc.function.arguments for tc in tool_calls if tc.function.arguments + ) + assert json.loads(tool_args) == {"city": "Dallas"} + + def test_streaming_no_reasoning_parser_include_false(self, tokenizer): + """No crash when reasoning parser absent and include_reasoning=False.""" + parser = make_parser(tokenizer, reasoning=False) + request = make_responses_request(include_reasoning=False) + + results = stream_text( + parser, + tokenizer, + MODEL_OUTPUT_CONTENT_ONLY, + request, + prompt_token_ids=[], + ) + reasoning, content, _ = collect_fields(results) + + assert reasoning == "" + assert "42" in content + + def test_streaming_chat_completion_include_false(self, tokenizer): + """parse_delta also respects ChatCompletionRequest.include_reasoning.""" + parser = make_parser(tokenizer, reasoning=True) + request = make_chat_request(include_reasoning=False) + + results = stream_text( + parser, + tokenizer, + MODEL_OUTPUT_REASONING_AND_CONTENT, + request, + prompt_token_ids=[], + ) + reasoning, content, _ = collect_fields(results) + + assert reasoning == "" + assert "42" in content + + def test_streaming_reasoning_only_deltas_become_none(self, tokenizer): + """Deltas that carry only reasoning become None (not empty).""" + parser = make_parser(tokenizer, reasoning=True) + request = make_responses_request(include_reasoning=False) + + results = stream_text( + parser, + tokenizer, + MODEL_OUTPUT_REASONING_AND_CONTENT, + request, + prompt_token_ids=[], + ) + + for r in results: + if r is not None: + assert r.reasoning is None + + +# ── ParsableContext.append_output() ─────────────────────────────────── + + +class TestParsableContextIncludeReasoning: + def _make_context(self, tokenizer, request): + from vllm.entrypoints.openai.responses.context import ParsableContext + + class TestParser(DelegatingParser): + reasoning_parser_cls = ThinkReasoningParser + tool_parser_cls = None + + return ParsableContext( + tokenizer=tokenizer, + parser_cls=TestParser, + response_messages=[], + request=request, + available_tools=None, + chat_template=None, + chat_template_content_format="auto", + ) + + def test_process_include_false_suppresses_reasoning(self, tokenizer): + """ParsableContext.process() suppresses reasoning items.""" + from vllm.outputs import CompletionOutput, RequestOutput + + request = make_responses_request(include_reasoning=False) + ctx = self._make_context(tokenizer, request) + + output = RequestOutput( + request_id="test", + prompt=None, + prompt_token_ids=[], + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text=MODEL_OUTPUT_REASONING_AND_CONTENT, + token_ids=tokenizer.encode( + MODEL_OUTPUT_REASONING_AND_CONTENT, + add_special_tokens=False, + ), + cumulative_logprob=None, + logprobs=None, + finish_reason="stop", + ) + ], + finished=True, + ) + + ctx.append_output(output) + + types = [getattr(m, "type", None) for m in ctx.response_messages] + assert "reasoning" not in types + assert "message" in types + + def test_process_include_true_has_reasoning(self, tokenizer): + """ParsableContext.process() includes reasoning by default.""" + from vllm.outputs import CompletionOutput, RequestOutput + + request = make_responses_request(include_reasoning=True) + ctx = self._make_context(tokenizer, request) + + output = RequestOutput( + request_id="test", + prompt=None, + prompt_token_ids=[], + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text=MODEL_OUTPUT_REASONING_AND_CONTENT, + token_ids=tokenizer.encode( + MODEL_OUTPUT_REASONING_AND_CONTENT, + add_special_tokens=False, + ), + cumulative_logprob=None, + logprobs=None, + finish_reason="stop", + ) + ], + finished=True, + ) + + ctx.append_output(output) + + types = [getattr(m, "type", None) for m in ctx.response_messages] + assert "reasoning" in types + assert "message" in types diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index caa3f724da45..eddcf014afee 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -600,19 +600,8 @@ async def chat_completion_stream_generator( prompt_token_ids=res.prompt_token_ids, finished=output.finish_reason is not None, ) - if delta_message is not None: - if delta_message.tool_calls: - tools_streamed[i] = True - - if ( - delta_message.reasoning - and not request.include_reasoning - ): - delta_message.reasoning = None - if not ( - delta_message.content or delta_message.tool_calls - ): - delta_message = None + if delta_message is not None and delta_message.tool_calls: + tools_streamed[i] = True # handle streaming just a content delta (no parsers) else: @@ -627,13 +616,22 @@ async def chat_completion_stream_generator( # "control token" for tool calls or the parser otherwise # wasn't ready to send a token, then # get the next token without streaming a chunk + # When reasoning is hidden, suppress per-token + # metadata (logprobs, token_ids) on every chunk to + # prevent leaking reasoning tokens through decoded + # token text in logprob entries or raw token IDs. + hide_stream_metadata = ( + not request.include_reasoning and parser is not None + ) + if hide_stream_metadata: + logprobs = None + if delta_message is None: # NOTE: If return_token_ids is enabled, we still need to # send a chunk with token_ids even if delta_message is None # to ensure all tokens are included in the response - if ( - output.finish_reason is None - and not request.return_token_ids + if output.finish_reason is None and ( + not request.return_token_ids or hide_stream_metadata ): continue delta_message = DeltaMessage() @@ -666,6 +664,10 @@ async def chat_completion_stream_generator( delta=True, ) + include_token_ids = ( + request.return_token_ids and not hide_stream_metadata + ) + if output.finish_reason is None: # Send token-by-token response for each request.n choice_data = ChatCompletionResponseStreamChoice( @@ -674,9 +676,7 @@ async def chat_completion_stream_generator( logprobs=logprobs, finish_reason=None, token_ids=( - as_list(output.token_ids) - if request.return_token_ids - else None + as_list(output.token_ids) if include_token_ids else None ), ) @@ -704,9 +704,7 @@ async def chat_completion_stream_generator( finish_reason=finish_reason_, stop_reason=output.stop_reason, token_ids=( - as_list(output.token_ids) - if request.return_token_ids - else None + as_list(output.token_ids) if include_token_ids else None ), ) @@ -883,12 +881,16 @@ async def chat_completion_full_generator( enable_auto_tools=self.enable_auto_tools, model_output_token_ids=token_ids, ) + suppress_metadata = not request.include_reasoning and parser is not None if not request.include_reasoning: reasoning = None + if suppress_metadata: + logprobs = None else: reasoning = None content = output.text tool_calls = [] + suppress_metadata = False auto_tools_called = False is_named_tool_choice = ( @@ -982,7 +984,9 @@ async def chat_completion_full_generator( else "stop", stop_reason=output.stop_reason, token_ids=( - as_list(output.token_ids) if request.return_token_ids else None + as_list(output.token_ids) + if request.return_token_ids and not suppress_metadata + else None ), routed_experts=routed_experts_b64, ) diff --git a/vllm/entrypoints/openai/responses/context.py b/vllm/entrypoints/openai/responses/context.py index ed989e2ba9f4..e1d3f4cf7ef4 100644 --- a/vllm/entrypoints/openai/responses/context.py +++ b/vllm/entrypoints/openai/responses/context.py @@ -310,7 +310,9 @@ def __init__( self.finish_reason: str | None = None self.enable_auto_tools = enable_auto_tools - self.response_parser = response_parser + self.response_parser = response_parser or ( + parser_cls(tokenizer, request.tools) if parser_cls is not None else None + ) self.parser_cls = parser_cls self.request = request @@ -344,6 +346,8 @@ def append_output(self, output: RequestOutput) -> None: enable_auto_tools=self.enable_auto_tools, model_output_token_ids=completion.token_ids, ) + if not self.request.include_reasoning: + reasoning = None self.response_messages.extend( build_response_output_items( reasoning=reasoning, diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index ba8bc5a40f1b..423068c71ca8 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -161,6 +161,15 @@ class ResponsesRequest(OpenAIBaseModel): previous_response_id: str | None = None prompt: ResponsePrompt | None = None reasoning: Reasoning | None = None + include_reasoning: bool = Field( + default=True, + description=( + "Whether to include reasoning content in the response. " + "When false, reasoning tokens are still generated but " + "excluded from the output. This reduces network traffic " + "without affecting model inference." + ), + ) service_tier: Literal["auto", "default", "flex", "scale", "priority"] = "auto" store: bool | None = True stream: bool | None = False diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index 40a520127920..3926284d701f 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -1068,6 +1068,9 @@ def _make_response_output_items( enable_auto_tools=self.enable_auto_tools, model_output_token_ids=final_output.token_ids, ) + if not request.include_reasoning: + reasoning = None + logprobs = None return build_response_output_items( reasoning=reasoning, content=content, @@ -1342,11 +1345,15 @@ async def _process_simple_streaming_events( ) -> AsyncGenerator[StreamingResponsesResponse, None]: processor = SimpleStreamingEventProcessor(tools=request.tools) + hide_stream_metadata = not request.include_reasoning and self.parser is not None + def _get_logprobs( output: CompletionOutput, ) -> list[response_text_delta_event.Logprob]: if not request.is_include_output_logprobs(): return [] + if hide_stream_metadata: + return [] return self._create_stream_response_logprobs( token_ids=output.token_ids, logprobs=output.logprobs, diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index 639aec0fa992..b643cbfc1c66 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -447,7 +447,7 @@ def _extract_tool_calls( tool_calls = list[FunctionCall]() if is_named_tool_choice and supports_required_and_named: - if content is None: + if content is None or (isinstance(content, str) and not content.strip()): return [], None function_name = self._get_function_name(request) tool_calls.append( @@ -496,6 +496,14 @@ def _extract_tool_calls( content = None else: # No tool calls. + # For required/named tool choice (when falling back to auto + # parsing), if content is empty or whitespace-only, return + # empty list with None content. + if (is_required_tool_choice or is_named_tool_choice) and ( + content is None + or (isinstance(content, str) and not content.strip()) + ): + return [], None return None, content return tool_calls, content @@ -917,6 +925,15 @@ def parse_delta( delta_message = self.finalize_generation(delta_message, request, state) delta_message = self._flush_engine_parsers(delta_message) + # Suppress reasoning deltas if not requested + if delta_message and not request.include_reasoning: + delta_message.reasoning = None + + # If only reasoning was in the message (no content, no tool_calls) + # skip emitting entirely + if not delta_message.content and not delta_message.tool_calls: + delta_message = None + return delta_message def _flush_engine_parsers( diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index 6f21cbca768a..fc2d653779ed 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -443,7 +443,15 @@ def parse_delta( if finished: events.extend(self._engine.finish()) result = self._events_to_delta(events, finished=finished) - return self._strip_trailing_reasoning(result) + result = self._strip_trailing_reasoning(result) + + # Suppress reasoning deltas if not requested + if result and not request.include_reasoning: + result.reasoning = None + if not result.content and not result.tool_calls: + result = None + + return result def _strip_trailing_reasoning( self, diff --git a/vllm/parser/harmony.py b/vllm/parser/harmony.py index 5043ca191f3e..1442246f139c 100644 --- a/vllm/parser/harmony.py +++ b/vllm/parser/harmony.py @@ -293,6 +293,16 @@ def parse_delta( delta_message.reasoning = combined_reasoning if tool_messages: delta_message.tool_calls = tool_messages + + # Suppress reasoning deltas if not requested + if delta_message and not request.include_reasoning: + delta_message.reasoning = None + + # If only reasoning was in the message (no content, no tool_calls) + # skip emitting entirely + if not delta_message.content and not delta_message.tool_calls: + return None + return delta_message def process_chunk(self, token_ids: Sequence[int]) -> ChunkResult: From 300e33797fac8a949a0ed89fab2633de96e65fa8 Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Fri, 10 Jul 2026 15:37:51 +0800 Subject: [PATCH 0010/1526] [Perf] fuse more rmsnorm and all-reduce in qwen3.5 (#46998) Signed-off-by: zjy0516 --- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 36 +++++++++---------- vllm/model_executor/models/qwen3_next.py | 16 +++------ 2 files changed, 21 insertions(+), 31 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 06bfe5c5de27..6e10b9a9932b 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -844,17 +844,14 @@ def rearrange_mixed_qkv(self, mixed_qkv): def forward( self, hidden_states: torch.Tensor, - output: torch.Tensor, - ): - self._forward_method(hidden_states, output) + ) -> torch.Tensor: + return self._forward_method(hidden_states) def _output_projection( self, core_attn_out: torch.Tensor, z: torch.Tensor, - output: torch.Tensor, - num_tokens: int, - ): + ) -> torch.Tensor: """Part 3: RMSNormGated + output linear projection. The RMSNormGated + quant sequence is eligible for fusion @@ -866,13 +863,13 @@ def _output_projection( core_attn_out = self.norm(core_attn_out, z) core_attn_out = core_attn_out.reshape(z_shape_og) core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d) - output[:num_tokens], _ = self.out_proj(core_attn_out) + output, _ = self.out_proj(core_attn_out) + return output def forward_hip( self, hidden_states: torch.Tensor, - output: torch.Tensor, - ): + ) -> torch.Tensor: """ROCm forward using AITER Triton fused projection+attention when available, otherwise falling back to the generic CUDA path.""" if GDN_AITER_TRITON_AVAILABLE: @@ -901,15 +898,14 @@ def forward_hip( use_aiter=True, ) - self._output_projection(core_attn_out, z, output, num_tokens) + return self._output_projection(core_attn_out, z) else: - self.forward_cuda(hidden_states, output) + return self.forward_cuda(hidden_states) def forward_cuda( self, hidden_states: torch.Tensor, - output: torch.Tensor, - ): + ) -> torch.Tensor: """ Forward pass with three parts: 1. Input projection @@ -964,13 +960,12 @@ def forward_cuda( # ============================================================ # Part 3: Output Projection # ============================================================ - self._output_projection(core_attn_out, z, output, num_tokens) + return self._output_projection(core_attn_out, z) def forward_xpu( self, hidden_states: torch.Tensor, - output: torch.Tensor, - ): + ) -> torch.Tensor: """ Forward pass with three parts: 1. Input projection @@ -1013,13 +1008,13 @@ def forward_xpu( core_attn_out = self.norm(core_attn_out, z) core_attn_out = core_attn_out.reshape(z_shape_og) core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d) - output[:num_tokens], _ = self.out_proj(core_attn_out) + out, _ = self.out_proj(core_attn_out) + return out def forward_cpu( self, hidden_states: torch.Tensor, - output: torch.Tensor, - ): + ) -> torch.Tensor: assert not hasattr(self, "in_proj_qkv"), "lora isn't supported on CPU." mixed_qkvz, _ = self.in_proj_qkvz(hidden_states) @@ -1063,7 +1058,8 @@ def forward_cpu( core_attn_out = self.norm(core_attn_out, z) core_attn_out = core_attn_out.reshape(z_shape_og) core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d) - output[:num_tokens], _ = self.out_proj(core_attn_out) + out, _ = self.out_proj(core_attn_out) + return out def _warmup_prefill_kernels(self, qkv_or_qkvz: torch.Tensor, v_dim: int) -> None: """Warm up GDN prefill kernels during V1 profiling. diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index ef320dd526b7..9a18cd4aad7b 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -381,15 +381,15 @@ def _project_qkv_gate( def forward( self, positions: torch.Tensor, - output: torch.Tensor, hidden_states: torch.Tensor, - ): + ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) q, k, v, gate = self._project_qkv_gate(qkv, positions) attn_output = self.attn(q, k, v) if gate is not None: attn_output = attn_output * torch.sigmoid(gate) - output[:], _ = self.o_proj(attn_output) + output, _ = self.o_proj(attn_output) + return output class Qwen3NextDecoderLayer(nn.Module): @@ -484,21 +484,15 @@ def forward( else: hidden_states, residual = self.input_layernorm(hidden_states, residual) - self_attention_output = torch.empty_like(hidden_states) if self.layer_type == "linear_attention": - self.linear_attn( - hidden_states=hidden_states, - output=self_attention_output, - ) + hidden_states = self.linear_attn(hidden_states=hidden_states) elif self.layer_type == "full_attention": - self.self_attn( + hidden_states = self.self_attn( hidden_states=hidden_states, - output=self_attention_output, positions=positions, ) else: raise ValueError("Invalid layer_type") - hidden_states = self_attention_output if self.layer_scale: if len(hidden_states.shape) == 2: From 28eaf05d5690a0c0d20e0999e3aa03a81102dfab Mon Sep 17 00:00:00 2001 From: Chaojun Zhang Date: Fri, 10 Jul 2026 15:40:51 +0800 Subject: [PATCH 0011/1526] [XPU] Enable v1/sample tests on XPU CI (#44472) Signed-off-by: Chaojun Zhang --- .buildkite/intel_jobs/misc_intel.yaml | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/.buildkite/intel_jobs/misc_intel.yaml b/.buildkite/intel_jobs/misc_intel.yaml index 656df8791b9e..55c9b9d9b424 100644 --- a/.buildkite/intel_jobs/misc_intel.yaml +++ b/.buildkite/intel_jobs/misc_intel.yaml @@ -72,9 +72,7 @@ steps: pytest -v -s v1/test_oracle.py && pytest -v -s v1/test_request.py && pytest -v -s v1/test_outputs.py && - pytest -v -s v1/sample/test_topk_topp_sampler.py && - pytest -v -s v1/sample/test_logprobs.py && - pytest -v -s v1/sample/test_logprobs_e2e.py' + pytest -v -s v1/sample' - label: Basic Models Tests (Initialization) timeout_in_minutes: 60 From 433f291195ded3ca8d278bc78da9280c5d4e5329 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Fri, 10 Jul 2026 00:53:16 -0700 Subject: [PATCH 0012/1526] [CI] Right-size test-area timeouts from nightly durations (#48186) Signed-off-by: khluu Co-authored-by: Claude Opus 4.8 --- .buildkite/test_areas/attention.yaml | 6 +-- .buildkite/test_areas/basic_correctness.yaml | 4 +- .buildkite/test_areas/benchmarks.yaml | 4 +- .buildkite/test_areas/compile.yaml | 22 +++++------ .buildkite/test_areas/cuda.yaml | 4 +- .buildkite/test_areas/disaggregated.yaml | 24 ++++++------ .buildkite/test_areas/distributed.yaml | 20 +++++----- .buildkite/test_areas/docker.yaml | 2 +- .buildkite/test_areas/e2e_integration.yaml | 10 ++--- .buildkite/test_areas/engine.yaml | 22 +++++------ .buildkite/test_areas/entrypoints.yaml | 20 +++++----- .buildkite/test_areas/expert_parallelism.yaml | 6 +-- .buildkite/test_areas/kernels.yaml | 38 +++++++++---------- .buildkite/test_areas/lm_eval.yaml | 38 +++++++++---------- .buildkite/test_areas/lora.yaml | 6 +-- .buildkite/test_areas/misc.yaml | 30 +++++++-------- .buildkite/test_areas/model_executor.yaml | 2 +- .buildkite/test_areas/model_runner_v2.yaml | 8 ++-- .buildkite/test_areas/models_basic.yaml | 8 ++-- .buildkite/test_areas/models_distributed.yaml | 2 +- .buildkite/test_areas/models_language.yaml | 16 ++++---- .buildkite/test_areas/models_multimodal.yaml | 14 +++---- .buildkite/test_areas/plugins.yaml | 2 +- .buildkite/test_areas/pytorch.yaml | 10 ++--- .buildkite/test_areas/quantization.yaml | 8 ++-- .buildkite/test_areas/rust_frontend.yaml | 10 ++--- .../test_areas/rust_frontend_cargo.yaml | 4 +- .buildkite/test_areas/samplers.yaml | 2 +- .buildkite/test_areas/spec_decode.yaml | 22 +++++------ .buildkite/test_areas/weight_loading.yaml | 2 +- 30 files changed, 183 insertions(+), 183 deletions(-) diff --git a/.buildkite/test_areas/attention.yaml b/.buildkite/test_areas/attention.yaml index 01e43b501497..9c662fec1afc 100644 --- a/.buildkite/test_areas/attention.yaml +++ b/.buildkite/test_areas/attention.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: V1 attention (H100-MI300) key: v1-attention-h100-mi300 - timeout_in_minutes: 30 + timeout_in_minutes: 85 device: h100 source_file_dependencies: - vllm/config/attention.py @@ -16,7 +16,7 @@ steps: mirror: amd: device: mi325_1 - timeout_in_minutes: 70 + timeout_in_minutes: 95 depends_on: - image-build-amd source_file_dependencies: @@ -30,7 +30,7 @@ steps: - label: V1 attention (B200) key: v1-attention-b200 - timeout_in_minutes: 30 + timeout_in_minutes: 80 device: b200-k8s source_file_dependencies: - vllm/config/attention.py diff --git a/.buildkite/test_areas/basic_correctness.yaml b/.buildkite/test_areas/basic_correctness.yaml index d7173b6438d8..1b92babb3ce7 100644 --- a/.buildkite/test_areas/basic_correctness.yaml +++ b/.buildkite/test_areas/basic_correctness.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Basic Correctness key: basic-correctness - timeout_in_minutes: 30 + timeout_in_minutes: 45 device: h200_18gb source_file_dependencies: - vllm/ @@ -19,6 +19,6 @@ steps: mirror: amd: device: mi325_1 - timeout_in_minutes: 50 + timeout_in_minutes: 70 depends_on: - image-build-amd diff --git a/.buildkite/test_areas/benchmarks.yaml b/.buildkite/test_areas/benchmarks.yaml index 622ebd44f409..7c26a8d6e982 100644 --- a/.buildkite/test_areas/benchmarks.yaml +++ b/.buildkite/test_areas/benchmarks.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Benchmarks CLI Test key: benchmarks-cli-test - timeout_in_minutes: 20 + timeout_in_minutes: 30 device: h200_18gb source_file_dependencies: - vllm/ @@ -23,7 +23,7 @@ steps: num_gpus: 2 optional: true working_dir: "/vllm-workspace/" - timeout_in_minutes: 10 + timeout_in_minutes: 20 source_file_dependencies: - benchmarks/attention_benchmarks/ - vllm/v1/attention/ diff --git a/.buildkite/test_areas/compile.yaml b/.buildkite/test_areas/compile.yaml index 01248738d519..eba37f657d90 100644 --- a/.buildkite/test_areas/compile.yaml +++ b/.buildkite/test_areas/compile.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Sequence Parallel Correctness Tests (2 GPUs) key: sequence-parallel-correctness-tests-2-gpus - timeout_in_minutes: 50 + timeout_in_minutes: 80 working_dir: "/vllm-workspace/" num_devices: 2 source_file_dependencies: @@ -19,7 +19,7 @@ steps: - label: Sequence Parallel Correctness Tests (2xH100) key: sequence-parallel-correctness-tests-2xh100 - timeout_in_minutes: 50 + timeout_in_minutes: 75 working_dir: "/vllm-workspace/" device: h100 optional: true @@ -30,7 +30,7 @@ steps: - label: AsyncTP Correctness Tests (2xH100) key: asynctp-correctness-tests-2xh100 - timeout_in_minutes: 50 + timeout_in_minutes: 30 working_dir: "/vllm-workspace/" device: h100 optional: true @@ -41,7 +41,7 @@ steps: - label: AsyncTP Correctness Tests (B200) key: asynctp-correctness-tests-b200 - timeout_in_minutes: 50 + timeout_in_minutes: 30 working_dir: "/vllm-workspace/" device: b200-k8s optional: true @@ -52,7 +52,7 @@ steps: - label: Distributed Compile Unit Tests (2xH100) key: distributed-compile-unit-tests-2xh100 - timeout_in_minutes: 20 + timeout_in_minutes: 45 working_dir: "/vllm-workspace/" device: h100 num_devices: 2 @@ -66,7 +66,7 @@ steps: - label: Fusion and Compile Unit Tests (2xB200) key: fusion-and-compile-unit-tests-2xb200 - timeout_in_minutes: 20 + timeout_in_minutes: 30 working_dir: "/vllm-workspace/" device: b200-k8s source_file_dependencies: @@ -96,7 +96,7 @@ steps: - label: Fusion E2E Quick (H100) key: fusion-e2e-quick-h100 - timeout_in_minutes: 15 + timeout_in_minutes: 25 working_dir: "/vllm-workspace/" device: h100 num_devices: 1 @@ -115,7 +115,7 @@ steps: - label: Fusion E2E Config Sweep (H100) key: fusion-e2e-config-sweep-h100 - timeout_in_minutes: 30 + timeout_in_minutes: 25 working_dir: "/vllm-workspace/" device: h100 num_devices: 1 @@ -149,7 +149,7 @@ steps: - label: Fusion E2E TP2 Quick (H100) key: fusion-e2e-tp2-quick-h100 - timeout_in_minutes: 20 + timeout_in_minutes: 35 working_dir: "/vllm-workspace/" device: h100 num_devices: 2 @@ -167,7 +167,7 @@ steps: - label: Fusion E2E TP2 AR-RMS Config Sweep (H100) key: fusion-e2e-tp2-ar-rms-config-sweep-h100 - timeout_in_minutes: 40 + timeout_in_minutes: 30 working_dir: "/vllm-workspace/" device: h100 num_devices: 2 @@ -207,7 +207,7 @@ steps: - label: Fusion E2E TP2 (B200) key: fusion-e2e-tp2-b200 - timeout_in_minutes: 20 + timeout_in_minutes: 45 working_dir: "/vllm-workspace/" device: b200-k8s num_devices: 2 diff --git a/.buildkite/test_areas/cuda.yaml b/.buildkite/test_areas/cuda.yaml index 956c76cf05f5..99e1949fef78 100644 --- a/.buildkite/test_areas/cuda.yaml +++ b/.buildkite/test_areas/cuda.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Platform Tests key: platform-tests - timeout_in_minutes: 15 + timeout_in_minutes: 20 device: h200_18gb source_file_dependencies: - vllm/envs.py @@ -19,7 +19,7 @@ steps: - label: Cudagraph key: cudagraph - timeout_in_minutes: 20 + timeout_in_minutes: 30 source_file_dependencies: - tests/v1/cudagraph - vllm/v1/cudagraph_dispatcher.py diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index 4cab698f322b..38020d91ad00 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Distributed NixlConnector PD accuracy (4 GPUs) key: distributed-nixlconnector-pd-accuracy-4-gpus - timeout_in_minutes: 30 + timeout_in_minutes: 55 working_dir: "/vllm-workspace/tests" num_devices: 4 source_file_dependencies: @@ -16,7 +16,7 @@ steps: mirror: amd: device: mi300_4 - timeout_in_minutes: 110 + timeout_in_minutes: 85 depends_on: - image-build-amd source_file_dependencies: @@ -29,7 +29,7 @@ steps: - label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs) key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus - timeout_in_minutes: 30 + timeout_in_minutes: 55 working_dir: "/vllm-workspace/tests" num_devices: 4 source_file_dependencies: @@ -66,7 +66,7 @@ steps: mirror: amd: device: mi300_4 - timeout_in_minutes: 50 + timeout_in_minutes: 60 depends_on: - image-build-amd source_file_dependencies: @@ -79,7 +79,7 @@ steps: - label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs) key: crosslayer-kv-layout-distributed-nixlconnector-pd-accuracy-tests-4-gpus - timeout_in_minutes: 30 + timeout_in_minutes: 55 working_dir: "/vllm-workspace/tests" num_devices: 4 source_file_dependencies: @@ -91,7 +91,7 @@ steps: mirror: amd: device: mi300_4 - timeout_in_minutes: 110 + timeout_in_minutes: 85 depends_on: - image-build-amd source_file_dependencies: @@ -104,7 +104,7 @@ steps: - label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs) key: hybrid-ssm-nixlconnector-pd-accuracy-tests-4-gpus - timeout_in_minutes: 25 + timeout_in_minutes: 60 working_dir: "/vllm-workspace/tests" num_devices: 4 source_file_dependencies: @@ -116,7 +116,7 @@ steps: mirror: amd: device: mi300_4 - timeout_in_minutes: 60 + timeout_in_minutes: 80 depends_on: - image-build-amd source_file_dependencies: @@ -143,7 +143,7 @@ steps: - label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs) key: multiconnector-nixl-offloading-pd-accuracy-2-gpus - timeout_in_minutes: 30 + timeout_in_minutes: 40 working_dir: "/vllm-workspace/tests" num_devices: 2 source_file_dependencies: @@ -158,7 +158,7 @@ steps: - label: NixlConnector PD + Spec Decode acceptance (2 GPUs) key: nixlconnector-pd-spec-decode-acceptance-2-gpus - timeout_in_minutes: 30 + timeout_in_minutes: 45 device: a100 working_dir: "/vllm-workspace/tests" num_devices: 2 @@ -172,7 +172,7 @@ steps: mirror: amd: device: mi300_2 - timeout_in_minutes: 60 + timeout_in_minutes: 70 depends_on: - image-build-amd source_file_dependencies: @@ -186,7 +186,7 @@ steps: - label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs) key: multiconnector-nixl-offloading-pd-edge-cases-2-gpus - timeout_in_minutes: 30 + timeout_in_minutes: 25 working_dir: "/vllm-workspace/tests" num_devices: 2 source_file_dependencies: diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 0b81cdb9d116..b519132c6f1f 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Distributed Comm Ops key: distributed-comm-ops - timeout_in_minutes: 20 + timeout_in_minutes: 25 working_dir: "/vllm-workspace/tests" num_devices: 2 source_file_dependencies: @@ -18,7 +18,7 @@ steps: - label: Distributed DP Tests (2 GPUs) key: distributed-dp-tests-2-gpus - timeout_in_minutes: 20 + timeout_in_minutes: 35 working_dir: "/vllm-workspace/tests" num_devices: 2 source_file_dependencies: @@ -55,7 +55,7 @@ steps: - label: Distributed Compile + RPC Tests (2 GPUs) key: distributed-compile-rpc-tests-2-gpus - timeout_in_minutes: 20 + timeout_in_minutes: 65 working_dir: "/vllm-workspace/tests" num_devices: 2 source_file_dependencies: @@ -78,7 +78,7 @@ steps: - label: Distributed Torchrun + Shutdown Tests (2 GPUs) key: distributed-torchrun-shutdown-tests-2-gpus - timeout_in_minutes: 20 + timeout_in_minutes: 30 working_dir: "/vllm-workspace/tests" num_devices: 2 source_file_dependencies: @@ -133,7 +133,7 @@ steps: - label: Distributed DP Tests (4 GPUs) key: distributed-dp-tests-4-gpus - timeout_in_minutes: 30 + timeout_in_minutes: 45 working_dir: "/vllm-workspace/tests" num_devices: 4 source_file_dependencies: @@ -154,7 +154,7 @@ steps: - label: Distributed Compile + Comm (4 GPUs) key: distributed-compile-comm-4-gpus - timeout_in_minutes: 30 + timeout_in_minutes: 70 working_dir: "/vllm-workspace/tests" num_devices: 4 source_file_dependencies: @@ -176,7 +176,7 @@ steps: - label: Distributed Tests (8xH100) key: distributed-tests-8xh100 - timeout_in_minutes: 10 + timeout_in_minutes: 20 device: h100 num_devices: 8 working_dir: "/vllm-workspace/tests" @@ -212,7 +212,7 @@ steps: - label: Distributed Tests (2xH100-2xMI300) key: distributed-tests-2xh100-2xmi300 - timeout_in_minutes: 15 + timeout_in_minutes: 30 device: h100 optional: true working_dir: "/vllm-workspace/" @@ -259,7 +259,7 @@ steps: - label: Pipeline + Context Parallelism (4 GPUs) key: pipeline-context-parallelism-4-gpus - timeout_in_minutes: 60 + timeout_in_minutes: 55 working_dir: "/vllm-workspace/tests" num_devices: 4 source_file_dependencies: @@ -274,7 +274,7 @@ steps: - label: RayExecutorV2 (4 GPUs) key: rayexecutorv2-4-gpus - timeout_in_minutes: 60 + timeout_in_minutes: 45 working_dir: "/vllm-workspace/tests" num_devices: 4 source_file_dependencies: diff --git a/.buildkite/test_areas/docker.yaml b/.buildkite/test_areas/docker.yaml index 9bf96221abe0..9f0562ca3bcf 100644 --- a/.buildkite/test_areas/docker.yaml +++ b/.buildkite/test_areas/docker.yaml @@ -3,7 +3,7 @@ depends_on: - image-build-cpu steps: - label: Docker Build Metadata - timeout_in_minutes: 10 + timeout_in_minutes: 20 device: cpu-small source_file_dependencies: - .buildkite/release-pipeline.yaml diff --git a/.buildkite/test_areas/e2e_integration.yaml b/.buildkite/test_areas/e2e_integration.yaml index 3f87e3958d0b..6655ae781e8d 100644 --- a/.buildkite/test_areas/e2e_integration.yaml +++ b/.buildkite/test_areas/e2e_integration.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: DeepSeek V2-Lite Sync EPLB Accuracy (4xH100) key: deepseek-v2-lite-sync-eplb-accuracy-4xh100 - timeout_in_minutes: 60 + timeout_in_minutes: 25 device: h100 optional: true num_devices: 4 @@ -14,7 +14,7 @@ steps: - label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (4xH100) key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-4xh100 - timeout_in_minutes: 60 + timeout_in_minutes: 25 device: h100 optional: true num_devices: 4 @@ -24,7 +24,7 @@ steps: - label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (2xB200) key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-2xb200 - timeout_in_minutes: 60 + timeout_in_minutes: 20 device: b200-k8s optional: true num_devices: 2 @@ -34,7 +34,7 @@ steps: - label: Qwen3-30B-A3B-FP8 DP4 Async EPLB Accuracy key: qwen3-30b-a3b-fp8-dp4-async-eplb-accuracy - timeout_in_minutes: 60 + timeout_in_minutes: 25 device: h100 optional: true num_devices: 4 @@ -44,7 +44,7 @@ steps: - label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100) key: deepseek-v2-lite-prefetch-offload-accuracy-h100 - timeout_in_minutes: 60 + timeout_in_minutes: 20 device: h100 optional: true num_devices: 1 diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index 7bb605465720..1c50ab7a1132 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Engine key: engine - timeout_in_minutes: 15 + timeout_in_minutes: 30 device: h200_18gb source_file_dependencies: - vllm/compilation/ @@ -29,13 +29,13 @@ steps: mirror: amd: device: mi325_1 - timeout_in_minutes: 60 + timeout_in_minutes: 50 depends_on: - image-build-amd - label: Engine (1 GPU) key: engine-1-gpu - timeout_in_minutes: 30 + timeout_in_minutes: 45 source_file_dependencies: - vllm/v1/engine/ - tests/v1/engine/ @@ -45,13 +45,13 @@ steps: mirror: amd: device: mi325_1 - timeout_in_minutes: 40 + timeout_in_minutes: 55 depends_on: - image-build-amd - label: e2e Scheduling (1 GPU) key: e2e-scheduling-1-gpu - timeout_in_minutes: 30 + timeout_in_minutes: 35 device: h200_18gb source_file_dependencies: - vllm/v1/ @@ -61,14 +61,14 @@ steps: mirror: amd: device: mi325_1 - timeout_in_minutes: 60 + timeout_in_minutes: 70 depends_on: - image-build-amd - label: e2e Core (1 GPU) device: h200_35gb key: e2e-core-1-gpu - timeout_in_minutes: 30 + timeout_in_minutes: 40 source_file_dependencies: - vllm/v1/ - tests/v1/e2e/general/ @@ -77,7 +77,7 @@ steps: mirror: amd: device: mi325_1 - timeout_in_minutes: 35 + timeout_in_minutes: 60 depends_on: - image-build-amd source_file_dependencies: @@ -87,7 +87,7 @@ steps: - label: V1 e2e (2 GPUs) key: v1-e2e-2-gpus - timeout_in_minutes: 60 # TODO: Fix timeout after we have more confidence in the test stability + timeout_in_minutes: 25 # TODO: Fix timeout after we have more confidence in the test stability optional: true num_devices: 2 source_file_dependencies: @@ -120,7 +120,7 @@ steps: - label: V1 e2e (4 GPUs) key: v1-e2e-4-gpus - timeout_in_minutes: 60 # TODO: Fix timeout after we have more confidence in the test stability + timeout_in_minutes: 20 # TODO: Fix timeout after we have more confidence in the test stability optional: true num_devices: 4 source_file_dependencies: @@ -148,7 +148,7 @@ steps: - label: V1 e2e (4xH100) key: v1-e2e-4xh100 - timeout_in_minutes: 60 + timeout_in_minutes: 35 device: h100 num_devices: 4 optional: true diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index 5ef88d4b97b5..2db4c5ad5a21 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Entrypoints Unit Tests key: entrypoints-unit-tests - timeout_in_minutes: 10 + timeout_in_minutes: 25 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/entrypoints @@ -16,7 +16,7 @@ steps: - label: Entrypoints Integration (LLM) key: entrypoints-integration-llm - timeout_in_minutes: 40 + timeout_in_minutes: 60 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -37,7 +37,7 @@ steps: - label: Entrypoints Integration (API Server) key: entrypoints-integration-api-server device: h200_35gb - timeout_in_minutes: 130 + timeout_in_minutes: 50 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -56,7 +56,7 @@ steps: - label: Entrypoints Integration (API Server OpenAI - Part 1) key: entrypoints-integration-api-server-openai-part-1 - timeout_in_minutes: 50 + timeout_in_minutes: 45 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -68,13 +68,13 @@ steps: mirror: amd: device: mi325_1 - timeout_in_minutes: 80 + timeout_in_minutes: 65 depends_on: - image-build-amd - label: Entrypoints Integration (API Server OpenAI - Part 2) key: entrypoints-integration-api-server-openai-part-2 - timeout_in_minutes: 50 + timeout_in_minutes: 45 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -109,7 +109,7 @@ steps: mirror: amd: device: mi325_1 - timeout_in_minutes: 60 + timeout_in_minutes: 65 depends_on: - image-build-amd @@ -126,7 +126,7 @@ steps: - label: Entrypoints Integration (Speech to Text) device: h200_35gb key: entrypoints-integration-speech_to_text - timeout_in_minutes: 50 + timeout_in_minutes: 45 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -138,7 +138,7 @@ steps: - label: Entrypoints Integration (Multimodal) device: h200_35gb key: entrypoints-integration-multimodal - timeout_in_minutes: 50 + timeout_in_minutes: 45 working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -160,7 +160,7 @@ steps: - label: OpenAI API Correctness key: openai-api-correctness - timeout_in_minutes: 30 + timeout_in_minutes: 20 device: h200_18gb source_file_dependencies: - csrc/ diff --git a/.buildkite/test_areas/expert_parallelism.yaml b/.buildkite/test_areas/expert_parallelism.yaml index ccb3054f2e9c..d02ffeb77485 100644 --- a/.buildkite/test_areas/expert_parallelism.yaml +++ b/.buildkite/test_areas/expert_parallelism.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: EPLB Algorithm key: eplb-algorithm - timeout_in_minutes: 15 + timeout_in_minutes: 20 device: h200_18gb working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -27,7 +27,7 @@ steps: - label: EPLB Execution # 17min key: eplb-execution - timeout_in_minutes: 27 + timeout_in_minutes: 25 working_dir: "/vllm-workspace/tests" num_devices: 4 source_file_dependencies: @@ -39,7 +39,7 @@ steps: - label: Elastic EP Scaling Test key: elastic-ep-scaling-test - timeout_in_minutes: 20 + timeout_in_minutes: 30 device: h100 working_dir: "/vllm-workspace/tests" num_devices: 4 diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 10c132da095a..d9e350f5093a 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: vLLM IR Tests key: vllm-ir-tests - timeout_in_minutes: 10 + timeout_in_minutes: 35 device: h200_18gb working_dir: "/vllm-workspace/" source_file_dependencies: @@ -16,7 +16,7 @@ steps: - label: Kernels Core Operation Test key: kernels-core-operation-test - timeout_in_minutes: 75 + timeout_in_minutes: 120 source_file_dependencies: - csrc/ - tests/kernels/core @@ -27,7 +27,7 @@ steps: - label: Kernels MiniMax Reduce RMS Test (2 GPUs) key: kernels-minimax-reduce-rms-test-2-gpus - timeout_in_minutes: 15 + timeout_in_minutes: 20 num_devices: 2 device: h100 source_file_dependencies: @@ -41,7 +41,7 @@ steps: - label: Deepseek V4 Kernel Test (H100) key: deepseek-v4-kernel-test-h100 - timeout_in_minutes: 15 + timeout_in_minutes: 30 device: h100 source_file_dependencies: - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -54,7 +54,7 @@ steps: - label: Deepseek V4 Kernel Test (B200) key: deepseek-v4-kernel-test-b200 - timeout_in_minutes: 15 + timeout_in_minutes: 20 device: b200-k8s source_file_dependencies: - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -65,7 +65,7 @@ steps: - label: Kernels Attention Test %N key: kernels-attention-test - timeout_in_minutes: 35 + timeout_in_minutes: 65 source_file_dependencies: - csrc/attention/ - vllm/v1/attention @@ -79,7 +79,7 @@ steps: mirror: amd: device: mi325_1 - timeout_in_minutes: 55 + timeout_in_minutes: 90 depends_on: - image-build-amd source_file_dependencies: @@ -106,7 +106,7 @@ steps: - label: Kernels Quantization Test %N key: kernels-quantization-test - timeout_in_minutes: 90 + timeout_in_minutes: 60 source_file_dependencies: - csrc/quantization/ - vllm/model_executor/layers/quantization @@ -131,7 +131,7 @@ steps: - label: Kernels MoE Test %N key: kernels-moe-test - timeout_in_minutes: 25 + timeout_in_minutes: 50 source_file_dependencies: - csrc/quantization/cutlass_w8a8/moe/ - csrc/moe/ @@ -147,7 +147,7 @@ steps: mirror: amd: device: mi325_1 - timeout_in_minutes: 50 + timeout_in_minutes: 65 source_file_dependencies: - csrc/quantization/cutlass_w8a8/moe/ - csrc/moe/ @@ -163,7 +163,7 @@ steps: - label: Kernels Mamba Test key: kernels-mamba-test - timeout_in_minutes: 45 + timeout_in_minutes: 40 source_file_dependencies: - csrc/mamba/ - tests/kernels/mamba @@ -172,7 +172,7 @@ steps: - pytest -v -s kernels/mamba - label: Kernels KDA Test - timeout_in_minutes: 20 + timeout_in_minutes: 25 device: h200_18gb source_file_dependencies: - vllm/model_executor/layers/fla/ops/kda.py @@ -184,7 +184,7 @@ steps: - label: Kernels DeepGEMM Test (H100) key: kernels-deepgemm-test-h100 - timeout_in_minutes: 45 + timeout_in_minutes: 35 device: h100 num_devices: 1 source_file_dependencies: @@ -211,7 +211,7 @@ steps: - label: Kernels (B200) key: kernels-b200 - timeout_in_minutes: 30 + timeout_in_minutes: 80 working_dir: "/vllm-workspace/" device: b200-k8s # optional: true @@ -264,7 +264,7 @@ steps: - label: Kernels Helion Test key: kernels-helion-test - timeout_in_minutes: 30 + timeout_in_minutes: 115 device: h100 source_file_dependencies: - vllm/utils/import_utils.py @@ -276,7 +276,7 @@ steps: - label: Kernels FP8 MoE Test (1xH100) key: kernels-fp8-moe-test-1xh100 - timeout_in_minutes: 90 + timeout_in_minutes: 40 device: h100 num_devices: 1 optional: true @@ -293,7 +293,7 @@ steps: - label: Kernels FP8 MoE Test (2xH100) key: kernels-fp8-moe-test-2xh100 - timeout_in_minutes: 90 + timeout_in_minutes: 45 device: h100 num_devices: 2 optional: true @@ -303,7 +303,7 @@ steps: - label: Kernels Fp4 MoE Test (B200) key: kernels-fp4-moe-test-b200 - timeout_in_minutes: 60 + timeout_in_minutes: 25 device: b200-k8s num_devices: 1 optional: true @@ -316,7 +316,7 @@ steps: - label: Kernels FusedMoE Layer Test (2 H100s) key: kernels-fusedmoe-layer-test-2-h100s - timeout_in_minutes: 90 + timeout_in_minutes: 30 device: h100 num_devices: 2 source_file_dependencies: diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index feba8f26eb32..1f7d27875367 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -5,7 +5,7 @@ steps: - label: LM Eval Small Models device: h200_35gb key: lm-eval-small-models - timeout_in_minutes: 75 + timeout_in_minutes: 45 source_file_dependencies: - csrc/ - vllm/model_executor/layers/quantization @@ -56,7 +56,7 @@ steps: - label: LM Eval Small Models (1xB200) key: lm-eval-small-models-1xb200 - timeout_in_minutes: 120 + timeout_in_minutes: 50 device: b200-k8s optional: true source_file_dependencies: @@ -80,7 +80,7 @@ steps: - label: LM Eval Large Models EP (2xB200) key: lm-eval-large-models-ep-2xb200 - timeout_in_minutes: 120 + timeout_in_minutes: 60 device: b200-k8s optional: true num_devices: 2 @@ -92,7 +92,7 @@ steps: - label: LM Eval Qwen3.5 Models (2xB200) key: lm-eval-qwen3-5-models-2xb200 - timeout_in_minutes: 120 + timeout_in_minutes: 45 device: b200-k8s optional: true num_devices: 2 @@ -109,7 +109,7 @@ steps: - label: LM Eval Large Models (8xH200) key: lm-eval-large-models-8xh200 - timeout_in_minutes: 60 + timeout_in_minutes: 50 device: h200 optional: true num_devices: 8 @@ -118,7 +118,7 @@ steps: mirror: amd: device: mi300_8 - timeout_in_minutes: 180 + timeout_in_minutes: 60 depends_on: - image-build-amd commands: @@ -152,7 +152,7 @@ steps: - label: LM Eval Humming f16 (A100 - TEMPORARY) key: lm-eval-humming-f16-a100 - timeout_in_minutes: 120 + timeout_in_minutes: 75 device: a100 optional: true num_devices: 1 @@ -167,7 +167,7 @@ steps: - label: LM Eval Humming Act int8 (A100 - TEMPORARY) key: lm-eval-humming-act-a100 - timeout_in_minutes: 120 + timeout_in_minutes: 45 device: a100 optional: true num_devices: 1 @@ -182,7 +182,7 @@ steps: - label: LM Eval Humming f16 (H100 - TEMPORARY) key: lm-eval-humming-f16-h100 - timeout_in_minutes: 120 + timeout_in_minutes: 70 device: h100 optional: true num_devices: 1 @@ -197,7 +197,7 @@ steps: - label: LM Eval Humming Act fp8/int8 (H100 - TEMPORARY) key: lm-eval-humming-act-h100 - timeout_in_minutes: 120 + timeout_in_minutes: 70 device: h100 optional: true num_devices: 1 @@ -213,7 +213,7 @@ steps: - label: LM Eval Humming f16 (B200 - TEMPORARY) key: lm-eval-humming-f16-b200 - timeout_in_minutes: 120 + timeout_in_minutes: 50 device: b200-k8s optional: true num_devices: 1 @@ -228,7 +228,7 @@ steps: - label: LM Eval Humming Act fp8/int8 (B200 - TEMPORARY) key: lm-eval-humming-act-b200 - timeout_in_minutes: 120 + timeout_in_minutes: 50 device: b200-k8s optional: true num_devices: 1 @@ -244,7 +244,7 @@ steps: - label: LM Eval TurboQuant KV Cache key: lm-eval-turboquant-kv-cache - timeout_in_minutes: 75 + timeout_in_minutes: 55 device: h200_18gb source_file_dependencies: - vllm/model_executor/layers/quantization/turboquant/ @@ -256,7 +256,7 @@ steps: - label: GPQA Eval (GPT-OSS) (2xH100) key: gpqa-eval-gpt-oss-2xh100 - timeout_in_minutes: 120 + timeout_in_minutes: 35 device: h100 optional: true num_devices: 2 @@ -270,7 +270,7 @@ steps: - label: GPQA Eval (GPT-OSS) (2xB200) key: gpqa-eval-gpt-oss-2xb200 - timeout_in_minutes: 120 + timeout_in_minutes: 30 device: b200-k8s optional: true num_devices: 2 @@ -284,7 +284,7 @@ steps: - label: GPQA Eval (GPT-OSS) (DGX Spark) key: gpqa-eval-gpt-oss-spark - timeout_in_minutes: 120 + timeout_in_minutes: 35 device: dgx-spark optional: true num_devices: 1 @@ -313,7 +313,7 @@ steps: - label: LM Eval KV-Offload (2xH100) key: kv-offload-medium - timeout_in_minutes: 60 + timeout_in_minutes: 30 device: h100 num_devices: 2 source_file_dependencies: @@ -327,7 +327,7 @@ steps: - label: LM Eval KV-Offload (4xH100) key: kv-offload-large - timeout_in_minutes: 60 + timeout_in_minutes: 40 device: h100 num_devices: 4 source_file_dependencies: @@ -341,7 +341,7 @@ steps: - label: MRCR Eval Small Models device: h200_35gb - timeout_in_minutes: 30 + timeout_in_minutes: 25 source_file_dependencies: - tests/evals/mrcr/ commands: diff --git a/.buildkite/test_areas/lora.yaml b/.buildkite/test_areas/lora.yaml index bd437c52265f..46a3710ea808 100644 --- a/.buildkite/test_areas/lora.yaml +++ b/.buildkite/test_areas/lora.yaml @@ -5,7 +5,7 @@ steps: - label: LoRA %N device: h200_35gb key: lora - timeout_in_minutes: 30 + timeout_in_minutes: 40 source_file_dependencies: - vllm/lora - tests/lora @@ -16,7 +16,7 @@ steps: amd: device: mi325_1 working_dir: "/vllm-workspace/tests" - timeout_in_minutes: 60 + timeout_in_minutes: 65 source_file_dependencies: - vllm/lora - tests/lora @@ -27,7 +27,7 @@ steps: - label: LoRA TP (Distributed) key: lora-tp-distributed - timeout_in_minutes: 30 + timeout_in_minutes: 60 num_devices: 4 source_file_dependencies: - vllm/lora diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index fd6ef2e61bae..13840c8db2fc 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -5,7 +5,7 @@ steps: - label: V1 Spec Decode device: h200_35gb key: v1-spec-decode - timeout_in_minutes: 30 + timeout_in_minutes: 40 source_file_dependencies: - vllm/config/ - vllm/distributed/ @@ -24,13 +24,13 @@ steps: mirror: amd: device: mi300_1 - timeout_in_minutes: 65 + timeout_in_minutes: 75 depends_on: - image-build-amd - label: V1 Sample + Logits key: v1-sample-logits - timeout_in_minutes: 30 + timeout_in_minutes: 45 device: h200_18gb source_file_dependencies: - vllm/config/ @@ -64,7 +64,7 @@ steps: - label: V1 Core + KV + Metrics key: v1-core-kv-metrics - timeout_in_minutes: 30 + timeout_in_minutes: 60 source_file_dependencies: - vllm/config/ - vllm/distributed/ @@ -108,7 +108,7 @@ steps: mirror: amd: device: mi325_1 - timeout_in_minutes: 60 + timeout_in_minutes: 75 depends_on: - image-build-amd @@ -172,7 +172,7 @@ steps: - label: Regression key: regression - timeout_in_minutes: 20 + timeout_in_minutes: 30 device: h200_18gb source_file_dependencies: - vllm/config/ @@ -195,7 +195,7 @@ steps: - label: Examples device: h200_35gb key: examples - timeout_in_minutes: 45 + timeout_in_minutes: 40 working_dir: "/vllm-workspace/examples" source_file_dependencies: - vllm/entrypoints @@ -237,7 +237,7 @@ steps: - label: Metrics, Tracing (2 GPUs) key: metrics-tracing-2-gpus - timeout_in_minutes: 20 + timeout_in_minutes: 25 num_devices: 2 source_file_dependencies: - vllm/config/ @@ -281,7 +281,7 @@ steps: mirror: amd: device: mi325_1 - timeout_in_minutes: 20 + timeout_in_minutes: 45 depends_on: - image-build-amd source_file_dependencies: @@ -292,7 +292,7 @@ steps: - label: Async Engine, Inputs, Utils, Worker device: h200_35gb key: async-engine-inputs-utils-worker - timeout_in_minutes: 50 + timeout_in_minutes: 25 source_file_dependencies: - vllm/assets/ - vllm/config/ @@ -319,7 +319,7 @@ steps: key: async-engine-inputs-utils-worker-config-cpu depends_on: - image-build-cpu - timeout_in_minutes: 30 + timeout_in_minutes: 65 source_file_dependencies: - vllm/assets/ - vllm/config/ @@ -381,7 +381,7 @@ steps: - label: Batch Invariance (A100) key: batch-invariance-a100 - timeout_in_minutes: 30 + timeout_in_minutes: 40 device: a100 source_file_dependencies: - vllm/v1/attention @@ -395,7 +395,7 @@ steps: - label: Batch Invariance (H100) key: batch-invariance-h100 - timeout_in_minutes: 30 + timeout_in_minutes: 40 device: h100 source_file_dependencies: - vllm/v1/attention @@ -411,7 +411,7 @@ steps: - label: Batch Invariance (B200) key: batch-invariance-b200 - timeout_in_minutes: 30 + timeout_in_minutes: 35 device: b200-k8s source_file_dependencies: - vllm/v1/attention @@ -430,7 +430,7 @@ steps: - label: Acceptance Length Test (Large Models) # optional device: h200_35gb key: acceptance-length-test-large-models - timeout_in_minutes: 25 + timeout_in_minutes: 20 gpu: h100 optional: true num_gpus: 1 diff --git a/.buildkite/test_areas/model_executor.yaml b/.buildkite/test_areas/model_executor.yaml index aaf85b4f2753..4280e600df4b 100644 --- a/.buildkite/test_areas/model_executor.yaml +++ b/.buildkite/test_areas/model_executor.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Model Executor key: model-executor - timeout_in_minutes: 35 + timeout_in_minutes: 45 source_file_dependencies: - vllm/engine/arg_utils.py - vllm/config/model.py diff --git a/.buildkite/test_areas/model_runner_v2.yaml b/.buildkite/test_areas/model_runner_v2.yaml index dbb35df80be9..3601aeee1172 100644 --- a/.buildkite/test_areas/model_runner_v2.yaml +++ b/.buildkite/test_areas/model_runner_v2.yaml @@ -5,7 +5,7 @@ steps: - label: Model Runner V2 Core Tests device: h200_35gb key: model-runner-v2-core-tests - timeout_in_minutes: 45 + timeout_in_minutes: 35 source_file_dependencies: - vllm/v1/worker/gpu/ - vllm/v1/worker/gpu_worker.py @@ -27,7 +27,7 @@ steps: - label: Model Runner V2 Examples device: h200_35gb key: model-runner-v2-examples - timeout_in_minutes: 45 + timeout_in_minutes: 35 working_dir: "/vllm-workspace/examples" source_file_dependencies: - vllm/v1/worker/gpu/ @@ -63,7 +63,7 @@ steps: - label: Model Runner V2 Distributed (2 GPUs) key: model-runner-v2-distributed-2-gpus - timeout_in_minutes: 45 + timeout_in_minutes: 30 working_dir: "/vllm-workspace/tests" num_devices: 2 source_file_dependencies: @@ -84,7 +84,7 @@ steps: - label: Model Runner V2 Pipeline Parallelism (4 GPUs) key: model-runner-v2-pipeline-parallelism-4-gpus - timeout_in_minutes: 60 + timeout_in_minutes: 50 working_dir: "/vllm-workspace/tests" num_devices: 4 source_file_dependencies: diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 3a113f1982a1..0fcd4f410487 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Basic Models Tests (Initialization) key: basic-models-tests-initialization - timeout_in_minutes: 45 + timeout_in_minutes: 25 device: h200_18gb source_file_dependencies: - vllm/ @@ -17,7 +17,7 @@ steps: - label: Basic Models Tests (Extra Initialization) %N device: h200_35gb key: basic-models-tests-extra-initialization - timeout_in_minutes: 45 + timeout_in_minutes: 100 source_file_dependencies: - vllm/model_executor/models/ - tests/models/test_initialization.py @@ -32,7 +32,7 @@ steps: - label: Basic Models Tests (Other) device: h200_35gb key: basic-models-tests-other - timeout_in_minutes: 45 + timeout_in_minutes: 35 source_file_dependencies: - vllm/ - tests/models/test_terratorch.py @@ -50,7 +50,7 @@ steps: key: basic-models-test-other-cpu depends_on: - image-build-cpu - timeout_in_minutes: 10 + timeout_in_minutes: 20 source_file_dependencies: - vllm/ - tests/models/test_utils.py diff --git a/.buildkite/test_areas/models_distributed.yaml b/.buildkite/test_areas/models_distributed.yaml index a3ee7666ed09..c1ec5eb00ae5 100644 --- a/.buildkite/test_areas/models_distributed.yaml +++ b/.buildkite/test_areas/models_distributed.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Distributed Model Tests (2 GPUs) key: distributed-model-tests-2-gpus - timeout_in_minutes: 50 + timeout_in_minutes: 60 working_dir: "/vllm-workspace/tests" num_devices: 2 source_file_dependencies: diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index 2c163bc80491..d89d3fabaecf 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Language Models Tests (Standard) key: language-models-tests-standard - timeout_in_minutes: 25 + timeout_in_minutes: 30 device: h200_18gb source_file_dependencies: - vllm/ @@ -21,7 +21,7 @@ steps: - label: Language Models Tests (Extra Standard) %N key: language-models-tests-extra-standard - timeout_in_minutes: 45 + timeout_in_minutes: 40 source_file_dependencies: - vllm/model_executor/models/ - tests/models/language/pooling/test_embedding.py @@ -52,7 +52,7 @@ steps: - label: Language Models Tests (Hybrid) %N key: language-models-tests-hybrid - timeout_in_minutes: 75 + timeout_in_minutes: 65 source_file_dependencies: - vllm/ - tests/models/language/generation @@ -67,7 +67,7 @@ steps: mirror: amd: device: mi325_1 - timeout_in_minutes: 90 + timeout_in_minutes: 70 depends_on: - image-build-amd commands: @@ -78,7 +78,7 @@ steps: - label: Language Models Test (Extended Generation) # 80min device: h200_35gb key: language-models-test-extended-generation - timeout_in_minutes: 110 + timeout_in_minutes: 65 optional: true source_file_dependencies: - vllm/ @@ -92,7 +92,7 @@ steps: - label: Language Models Test (PPL) key: language-models-test-ppl - timeout_in_minutes: 110 + timeout_in_minutes: 30 device: h200_18gb optional: true source_file_dependencies: @@ -104,7 +104,7 @@ steps: - label: Language Models Test (Extended Pooling) # 36min device: h200_35gb key: language-models-test-extended-pooling - timeout_in_minutes: 50 + timeout_in_minutes: 70 optional: true source_file_dependencies: - vllm/ @@ -120,7 +120,7 @@ steps: - label: Language Models Test (MTEB) key: language-models-test-mteb - timeout_in_minutes: 110 + timeout_in_minutes: 45 device: h200_18gb optional: true source_file_dependencies: diff --git a/.buildkite/test_areas/models_multimodal.yaml b/.buildkite/test_areas/models_multimodal.yaml index a721efa6067e..6720483c25f2 100644 --- a/.buildkite/test_areas/models_multimodal.yaml +++ b/.buildkite/test_areas/models_multimodal.yaml @@ -20,7 +20,7 @@ steps: - label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" key: multi-modal-models-standard-2-qwen3-gemma - timeout_in_minutes: 45 + timeout_in_minutes: 50 device: h200_18gb source_file_dependencies: - vllm/ @@ -38,7 +38,7 @@ steps: - label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl" device: h200_35gb key: multi-modal-models-standard-3-llava-qwen2-vl - timeout_in_minutes: 45 + timeout_in_minutes: 40 source_file_dependencies: - vllm/ - tests/models/multimodal @@ -54,7 +54,7 @@ steps: - label: "Multi-Modal Models (Standard) 4: other + whisper" device: h200_35gb key: multi-modal-models-standard-4-other-whisper - timeout_in_minutes: 45 + timeout_in_minutes: 50 source_file_dependencies: - vllm/ - tests/models/multimodal @@ -73,7 +73,7 @@ steps: key: multi-modal-processor-cpu depends_on: - image-build-cpu - timeout_in_minutes: 60 + timeout_in_minutes: 125 source_file_dependencies: - vllm/ - tests/models/multimodal @@ -84,7 +84,7 @@ steps: - label: Multi-Modal Processor # 44min key: multi-modal-processor - timeout_in_minutes: 60 + timeout_in_minutes: 65 device: h200_18gb source_file_dependencies: - vllm/ @@ -96,7 +96,7 @@ steps: - label: Multi-Modal Accuracy Eval (Small Models) # 50min device: h200_35gb key: multi-modal-accuracy-eval-small-models - timeout_in_minutes: 70 + timeout_in_minutes: 30 working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" source_file_dependencies: - vllm/multimodal/ @@ -164,7 +164,7 @@ steps: mirror: amd: device: mi325_1 - timeout_in_minutes: 60 + timeout_in_minutes: 75 depends_on: - image-build-amd source_file_dependencies: diff --git a/.buildkite/test_areas/plugins.yaml b/.buildkite/test_areas/plugins.yaml index 5effe17513d4..e0bb67b1bb47 100644 --- a/.buildkite/test_areas/plugins.yaml +++ b/.buildkite/test_areas/plugins.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Plugin Tests (2 GPUs) key: plugin-tests-2-gpus - timeout_in_minutes: 60 + timeout_in_minutes: 35 working_dir: "/vllm-workspace/tests" num_devices: 2 source_file_dependencies: diff --git a/.buildkite/test_areas/pytorch.yaml b/.buildkite/test_areas/pytorch.yaml index 5c3060582aa8..72cb00696644 100644 --- a/.buildkite/test_areas/pytorch.yaml +++ b/.buildkite/test_areas/pytorch.yaml @@ -5,7 +5,7 @@ steps: - label: PyTorch Compilation Unit Tests device: h200_35gb key: pytorch-compilation-unit-tests - timeout_in_minutes: 10 + timeout_in_minutes: 90 source_file_dependencies: - vllm/__init__.py - vllm/_aiter_ops.py @@ -78,7 +78,7 @@ steps: - label: PyTorch Compilation Passes Unit Tests key: pytorch-compilation-passes-unit-tests - timeout_in_minutes: 20 + timeout_in_minutes: 45 source_file_dependencies: - vllm/__init__.py - vllm/_aiter_ops.py @@ -110,13 +110,13 @@ steps: mirror: amd: device: mi300_1 - timeout_in_minutes: 180 + timeout_in_minutes: 65 depends_on: - image-build-amd - label: PyTorch Fullgraph Smoke Test key: pytorch-fullgraph-smoke-test - timeout_in_minutes: 35 + timeout_in_minutes: 60 source_file_dependencies: - vllm/__init__.py - vllm/_aiter_ops.py @@ -152,7 +152,7 @@ steps: - label: PyTorch Fullgraph key: pytorch-fullgraph - timeout_in_minutes: 30 + timeout_in_minutes: 40 device: h200_18gb source_file_dependencies: - vllm/__init__.py diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index a92ee24f4aac..ce3e58e501b9 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Quantization key: quantization - timeout_in_minutes: 90 + timeout_in_minutes: 60 source_file_dependencies: - csrc/ - vllm/model_executor/layers/quantization @@ -23,7 +23,7 @@ steps: - label: Quantized Fusions key: quantized-fusions - timeout_in_minutes: 30 + timeout_in_minutes: 20 source_file_dependencies: - tests/fusion - vllm/model_executor/layers/fusion @@ -35,7 +35,7 @@ steps: - label: Quantized MoE Test (B200) key: quantized-moe-test-b200 - timeout_in_minutes: 60 + timeout_in_minutes: 120 working_dir: "/vllm-workspace/" device: b200-k8s source_file_dependencies: @@ -53,7 +53,7 @@ steps: - label: Quantized Models Test key: quantized-models-test - timeout_in_minutes: 60 + timeout_in_minutes: 50 source_file_dependencies: - vllm/model_executor/layers/quantization - tests/models/quantization diff --git a/.buildkite/test_areas/rust_frontend.yaml b/.buildkite/test_areas/rust_frontend.yaml index 1dfe912aa89e..9e5e09c3ec30 100644 --- a/.buildkite/test_areas/rust_frontend.yaml +++ b/.buildkite/test_areas/rust_frontend.yaml @@ -3,7 +3,7 @@ depends_on: - image-build steps: - label: Rust Frontend OpenAI Coverage - timeout_in_minutes: 90 + timeout_in_minutes: 30 device: h200_18gb working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -38,7 +38,7 @@ steps: - pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server" - label: Rust Frontend Serve/Admin Coverage - timeout_in_minutes: 60 + timeout_in_minutes: 25 device: h200_18gb working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -67,7 +67,7 @@ steps: - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info" - label: Rust Frontend Core Correctness - timeout_in_minutes: 30 + timeout_in_minutes: 20 device: h200_18gb working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -81,7 +81,7 @@ steps: - pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine - label: Rust Frontend Tool Use - timeout_in_minutes: 60 + timeout_in_minutes: 25 working_dir: "/vllm-workspace/tests" source_file_dependencies: - rust/ @@ -95,7 +95,7 @@ steps: - pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice" - label: Rust Frontend Distributed - timeout_in_minutes: 30 + timeout_in_minutes: 25 num_devices: 4 working_dir: "/vllm-workspace/tests" source_file_dependencies: diff --git a/.buildkite/test_areas/rust_frontend_cargo.yaml b/.buildkite/test_areas/rust_frontend_cargo.yaml index 06f9eb9c245d..21d4c2ac2192 100644 --- a/.buildkite/test_areas/rust_frontend_cargo.yaml +++ b/.buildkite/test_areas/rust_frontend_cargo.yaml @@ -4,7 +4,7 @@ steps: - label: Rust Frontend Cargo Style + Clippy key: rust-frontend-cargo-style-clippy depends_on: [] - timeout_in_minutes: 30 + timeout_in_minutes: 20 device: cpu-medium no_plugin: true source_file_dependencies: @@ -18,7 +18,7 @@ steps: - label: Rust Frontend Cargo Tests key: rust-frontend-cargo-tests depends_on: [] - timeout_in_minutes: 30 + timeout_in_minutes: 20 device: cpu-medium no_plugin: true source_file_dependencies: diff --git a/.buildkite/test_areas/samplers.yaml b/.buildkite/test_areas/samplers.yaml index 5abc16889434..2e7cd4a623e6 100644 --- a/.buildkite/test_areas/samplers.yaml +++ b/.buildkite/test_areas/samplers.yaml @@ -5,7 +5,7 @@ steps: - label: Samplers Test device: h200_35gb key: samplers-test - timeout_in_minutes: 75 + timeout_in_minutes: 40 source_file_dependencies: - vllm/model_executor/layers - vllm/sampling_metadata.py diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index 671638f6f642..096c324bb8e2 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Spec Decode Eagle key: spec-decode-eagle - timeout_in_minutes: 30 + timeout_in_minutes: 25 device: h200_18gb source_file_dependencies: - vllm/v1/spec_decode/ @@ -15,7 +15,7 @@ steps: mirror: amd: device: mi325_1 - timeout_in_minutes: 45 + timeout_in_minutes: 60 depends_on: - image-build-amd source_file_dependencies: @@ -29,7 +29,7 @@ steps: - label: Spec Decode Eagle Nightly B200 key: spec-decode-eagle-nightly-b200 - timeout_in_minutes: 30 + timeout_in_minutes: 25 device: b200-k8s optional: true source_file_dependencies: @@ -41,7 +41,7 @@ steps: - label: Spec Decode Speculators + MTP key: spec-decode-speculators-mtp - timeout_in_minutes: 30 + timeout_in_minutes: 20 device: h200_18gb source_file_dependencies: - vllm/v1/spec_decode/ @@ -82,7 +82,7 @@ steps: - label: Spec Decode Ngram + Suffix key: spec-decode-ngram-suffix - timeout_in_minutes: 30 + timeout_in_minutes: 20 device: h200_18gb source_file_dependencies: - vllm/v1/spec_decode/ @@ -93,7 +93,7 @@ steps: mirror: amd: device: mi325_1 - timeout_in_minutes: 65 + timeout_in_minutes: 55 # TODO(akaratza): Test after Torch >= 2.12 bump soft_fail: true depends_on: @@ -109,7 +109,7 @@ steps: - label: Spec Decode Draft Model key: spec-decode-draft-model - timeout_in_minutes: 30 + timeout_in_minutes: 45 device: h200_18gb source_file_dependencies: - vllm/v1/spec_decode/ @@ -120,7 +120,7 @@ steps: mirror: amd: device: mi325_1 - timeout_in_minutes: 50 + timeout_in_minutes: 55 depends_on: - image-build-amd source_file_dependencies: @@ -134,7 +134,7 @@ steps: - label: Spec Decode Draft Model Nightly B200 key: spec-decode-draft-model-nightly-b200 - timeout_in_minutes: 30 + timeout_in_minutes: 40 device: b200-k8s optional: true source_file_dependencies: @@ -146,7 +146,7 @@ steps: - label: Speculators Correctness key: speculators-correctness - timeout_in_minutes: 60 + timeout_in_minutes: 30 device: h100 optional: true num_devices: 1 @@ -159,7 +159,7 @@ steps: - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test - label: Spec Decode MTP hybrid (B200) - timeout_in_minutes: 30 + timeout_in_minutes: 20 device: b200-k8s optional: true source_file_dependencies: diff --git a/.buildkite/test_areas/weight_loading.yaml b/.buildkite/test_areas/weight_loading.yaml index 9d7bd0bce91c..eeb24a49d8d3 100644 --- a/.buildkite/test_areas/weight_loading.yaml +++ b/.buildkite/test_areas/weight_loading.yaml @@ -4,7 +4,7 @@ depends_on: steps: - label: Weight Loading Multiple GPU # 33min key: weight-loading-multiple-gpu - timeout_in_minutes: 45 + timeout_in_minutes: 50 working_dir: "/vllm-workspace/tests" num_devices: 2 optional: true From 216ee58780a85470971ac9be02d0fd2eabc942b7 Mon Sep 17 00:00:00 2001 From: wenjun liu Date: Fri, 10 Jul 2026 15:56:00 +0800 Subject: [PATCH 0013/1526] Add XPU nightly and release image publishing to DockerHub (#48126) Signed-off-by: wenjun.liu Signed-off-by: jun,du Co-authored-by: jun,du Co-authored-by: Kunshang Ji --- .buildkite/release-pipeline.yaml | 18 ++++++++++++++++ .buildkite/scripts/publish-release-images.sh | 16 ++++++++++++++ .../scripts/xpu/push-nightly-builds-xpu.sh | 21 +++++++++++++++++++ 3 files changed, 55 insertions(+) create mode 100644 .buildkite/scripts/xpu/push-nightly-builds-xpu.sh diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index b1fe875bc2d8..1d3a1d4a5348 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -848,6 +848,23 @@ steps: DOCKER_BUILDKIT: "1" S3_BUCKET: "vllm-wheels" + - label: "Publish nightly XPU image to DockerHub" + depends_on: + - create-manifest-xpu + if: build.env("NIGHTLY") == "1" + agents: + queue: small_cpu_queue_release + commands: + - "bash .buildkite/scripts/xpu/push-nightly-builds-xpu.sh" + - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-xpu" + plugins: + - docker-login#v3.0.0: + username: vllmbot + password-env: DOCKERHUB_TOKEN + env: + DOCKER_BUILDKIT: "1" + DOCKERHUB_USERNAME: "vllmbot" + - label: "Publish nightly ROCm image to DockerHub" depends_on: - build-rocm-release-image @@ -878,6 +895,7 @@ steps: - create-multi-arch-manifest-cuda-12-9 - create-multi-arch-manifest-ubuntu2404 - create-multi-arch-manifest-cuda-12-9-ubuntu2404 + - create-manifest-xpu - build-rocm-release-image - input-release-version # Wait for CPU builds if their block steps were unblocked, so publish diff --git a/.buildkite/scripts/publish-release-images.sh b/.buildkite/scripts/publish-release-images.sh index ec319aa76006..91b5c3ace1be 100755 --- a/.buildkite/scripts/publish-release-images.sh +++ b/.buildkite/scripts/publish-release-images.sh @@ -130,6 +130,22 @@ docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm docker push vllm/vllm-openai-rocm:latest-base docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base +# ---- XPU ---- + +docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu + +docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu vllm/vllm-openai-xpu:latest-x86_64 +docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 +docker push vllm/vllm-openai-xpu:latest-x86_64 +docker push vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 + +docker manifest rm vllm/vllm-openai-xpu:latest || true +docker manifest rm vllm/vllm-openai-xpu:v${RELEASE_VERSION} || true +docker manifest create vllm/vllm-openai-xpu:latest vllm/vllm-openai-xpu:latest-x86_64 --amend +docker manifest create vllm/vllm-openai-xpu:v${RELEASE_VERSION} vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 --amend +docker manifest push vllm/vllm-openai-xpu:latest +docker manifest push vllm/vllm-openai-xpu:v${RELEASE_VERSION} + # ---- CPU ---- # CPU images are behind separate block steps and may not have been built. # All-or-nothing: inspect both arches first, then either publish everything diff --git a/.buildkite/scripts/xpu/push-nightly-builds-xpu.sh b/.buildkite/scripts/xpu/push-nightly-builds-xpu.sh new file mode 100644 index 000000000000..ca50f2543bba --- /dev/null +++ b/.buildkite/scripts/xpu/push-nightly-builds-xpu.sh @@ -0,0 +1,21 @@ +#!/bin/bash + +set -ex + +ORIG_TAG_NAME="$BUILDKITE_COMMIT" +REPO="vllm/vllm-openai-xpu" + +echo "Pushing original XPU tag ${ORIG_TAG_NAME}-xpu to nightly tags in ${REPO}" + +aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7 +docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:"$ORIG_TAG_NAME"-x86_64-xpu + +docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:"$ORIG_TAG_NAME"-x86_64-xpu ${REPO}:nightly-x86_64 +docker push ${REPO}:nightly-x86_64 + +docker manifest rm ${REPO}:nightly || true +docker manifest rm ${REPO}:nightly-"$BUILDKITE_COMMIT" || true +docker manifest create ${REPO}:nightly ${REPO}:nightly-x86_64 --amend +docker manifest create ${REPO}:nightly-"$BUILDKITE_COMMIT" ${REPO}:nightly-x86_64 --amend +docker manifest push ${REPO}:nightly +docker manifest push ${REPO}:nightly-"$BUILDKITE_COMMIT" \ No newline at end of file From 074bdd0d9961fb7f81c9311e00c00cfd5aff77a3 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Fri, 10 Jul 2026 16:15:33 +0800 Subject: [PATCH 0014/1526] [Rust Frontend] Integrate MM video support (#47959) Signed-off-by: Bugen Zhao --- rust/Cargo.lock | 4 +- rust/Cargo.toml | 4 +- rust/src/chat/Cargo.toml | 1 + rust/src/chat/src/backend/hf.rs | 17 +- rust/src/chat/src/error.rs | 32 +- rust/src/chat/src/multimodal.rs | 894 ++++++++---------- rust/src/chat/src/multimodal/expand.rs | 446 +++++++++ rust/src/chat/src/multimodal/image.rs | 141 +++ rust/src/chat/src/multimodal/tensor.rs | 46 +- rust/src/chat/src/multimodal/video.rs | 316 +++++++ rust/src/chat/src/renderer/hf/mod.rs | 91 +- rust/src/chat/src/request.rs | 37 +- .../routes/openai/chat_completions/convert.rs | 21 +- .../server/src/routes/openai/utils/types.rs | 6 +- rust/src/server/src/routes/tests.rs | 12 +- rust/src/text/src/backend/hf/model_files.rs | 19 + 16 files changed, 1571 insertions(+), 516 deletions(-) create mode 100644 rust/src/chat/src/multimodal/expand.rs create mode 100644 rust/src/chat/src/multimodal/image.rs create mode 100644 rust/src/chat/src/multimodal/video.rs diff --git a/rust/Cargo.lock b/rust/Cargo.lock index cc055505dc4e..d155e07edadd 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -2167,7 +2167,7 @@ checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" [[package]] name = "llm-multimodal" version = "1.7.1" -source = "git+https://github.com/smg-project/llm-multimodal?rev=7d74582aeaf0e4086a44964382655d22f1af0686#7d74582aeaf0e4086a44964382655d22f1af0686" +source = "git+https://github.com/smg-project/llm-multimodal?rev=c8a29dcc755139fdc26185f400ea48c6d6d48273#c8a29dcc755139fdc26185f400ea48c6d6d48273" dependencies = [ "anyhow", "base64 0.22.1", @@ -2473,6 +2473,7 @@ dependencies = [ "portable-atomic", "portable-atomic-util", "rawpointer", + "serde", ] [[package]] @@ -5093,6 +5094,7 @@ dependencies = [ "llm-multimodal", "minijinja", "minijinja-contrib", + "ndarray 0.17.2", "oss-harmony", "paste", "reqwest", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 435350c07117..db9764672f9e 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -53,12 +53,12 @@ hyper-util = { version = "0.1.20", features = [ indexmap = "2.13.0" itertools = "0.14.0" libc = "0.2.177" -llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "7d74582aeaf0e4086a44964382655d22f1af0686" } +llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "c8a29dcc755139fdc26185f400ea48c6d6d48273" } mimalloc = "0.1.52" minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] } minijinja-contrib = { version = "2.0", features = ["pycompat"] } native-tls-vendored = { package = "native-tls", version = "0.2.18", features = ["vendored"] } -ndarray = { version = "0.16.1", features = ["serde"] } +ndarray = { version = "0.17", features = ["serde"] } openai-harmony = { package = "oss-harmony", git = "https://github.com/oss-harmony/harmony", tag = "v0.0.11", default-features = false } openai-protocol = "1.6.0" openssl = "0.10" diff --git a/rust/src/chat/Cargo.toml b/rust/src/chat/Cargo.toml index 95ce5ff2e42c..00a6e223b775 100644 --- a/rust/src/chat/Cargo.toml +++ b/rust/src/chat/Cargo.toml @@ -42,6 +42,7 @@ anyhow.workspace = true bytes.workspace = true clap.workspace = true expect-test.workspace = true +ndarray.workspace = true paste.workspace = true rmp-serde.workspace = true serial_test.workspace = true diff --git a/rust/src/chat/src/backend/hf.rs b/rust/src/chat/src/backend/hf.rs index fdfe8620b20a..47c0bfb78c82 100644 --- a/rust/src/chat/src/backend/hf.rs +++ b/rust/src/chat/src/backend/hf.rs @@ -10,7 +10,7 @@ use crate::backend::{ NewChatOutputProcessorOptions, }; use crate::error::Result; -use crate::multimodal::MultimodalModelInfo; +use crate::multimodal::{MultimodalConfigFiles, MultimodalModelInfo}; use crate::output::{ DefaultChatOutputProcessor, HarmonyChatOutputProcessor, validate_harmony_parser_overrides, }; @@ -46,8 +46,12 @@ impl HfChatBackend { MultimodalModelInfo::from_paths( model_id.clone(), (!model_type.is_empty()).then_some(model_type.to_string()), - files.config_path.as_deref(), - files.preprocessor_config_path.as_deref(), + MultimodalConfigFiles { + config: files.config_path.as_deref(), + preprocessor_config: files.preprocessor_config_path.as_deref(), + video_preprocessor_config: files.video_preprocessor_config_path.as_deref(), + processor_config: files.processor_config_path.as_deref(), + }, tokenizer.clone(), )? }; @@ -139,8 +143,11 @@ pub(super) async fn load_model_backends( fn resolve_multimodal_render_info( info: Option<&MultimodalModelInfo>, ) -> Option { + use llm_multimodal::Modality; + info.map(|info| MultimodalRenderInfo { - placeholder_token: info.placeholder_token().to_string(), + image_token: info.placeholder_token(Modality::Image).map(str::to_string), + video_token: info.placeholder_token(Modality::Video).map(str::to_string), }) } @@ -192,6 +199,8 @@ mod tests { tokenizer_config_path: Some(tokenizer_config_path), generation_config_path: None, preprocessor_config_path: None, + video_preprocessor_config_path: None, + processor_config_path: None, chat_template_path: None, config_path: Some(config_path), } diff --git a/rust/src/chat/src/error.rs b/rust/src/chat/src/error.rs index da2396c2198d..8e5d4ef84f89 100644 --- a/rust/src/chat/src/error.rs +++ b/rust/src/chat/src/error.rs @@ -1,5 +1,5 @@ use thiserror::Error; -use thiserror_ext::Macro; +use thiserror_ext::{AsReport as _, Macro}; type BoxedError = Box; @@ -18,6 +18,8 @@ pub enum Error { UnsupportedMultimodalRenderer, #[error("unsupported multimodal content: {0}")] UnsupportedMultimodalContent(&'static str), + #[error("`{modality}` input is not supported by this model")] + UnsupportedModality { modality: String }, #[error("multimodal preprocessing error: {0}")] Multimodal(#[message] String), #[error("{kind} parsing is not available for model `{model_id}`")] @@ -80,11 +82,39 @@ impl Error { match self { Self::PromptTooLong { .. } => true, Self::Text(error) => error.is_request_validation_error(), + Self::UnsupportedMultimodalRenderer + | Self::UnsupportedMultimodalContent(_) + | Self::UnsupportedModality { .. } => true, + _ => false, } } } +impl From for Error { + fn from(error: llm_multimodal::MediaConnectorError) -> Self { + Self::Multimodal(error.to_report_string()) + } +} + +impl From for Error { + fn from(error: llm_multimodal::MultiModalError) -> Self { + Self::Multimodal(error.to_report_string()) + } +} + +impl From for Error { + fn from(error: llm_multimodal::TransformError) -> Self { + Self::Multimodal(error.to_report_string()) + } +} + +impl From for Error { + fn from(error: llm_multimodal::registry::ModelRegistryError) -> Self { + Self::Multimodal(error.to_report_string()) + } +} + /// Format the available-parser suffix used in user-facing error messages. fn available_parser_hint(available_names: &[String]) -> String { if available_names.is_empty() { diff --git a/rust/src/chat/src/multimodal.rs b/rust/src/chat/src/multimodal.rs index 7d950d581da5..422696ae7107 100644 --- a/rust/src/chat/src/multimodal.rs +++ b/rust/src/chat/src/multimodal.rs @@ -1,8 +1,8 @@ -//! Chat-layer multimodal image preparation. +//! Chat-layer multimodal media preparation. //! -//! This module owns the narrow image-only multimodal path for chat requests: -//! it extracts image parts from structured chat messages, fetches and -//! preprocesses them through `llm-multimodal`, expands rendered prompt +//! This module owns the multimodal path for chat requests: it extracts media +//! parts from structured chat messages, fetches and preprocesses them through +//! `llm-multimodal` one modality at a time, expands rendered prompt //! placeholders after tokenization, and builds the engine-facing //! `MmFeatures` payload. //! @@ -16,19 +16,15 @@ use std::sync::{Arc, LazyLock}; use itertools::izip; use llm_multimodal::{ - AsyncMultiModalTracker, FieldLayout, MediaConnector, MediaConnectorConfig, MediaContentPart, - Modality, ModelMetadata, ModelProcessorSpec, ModelRegistry, PreProcessorConfig, - PreprocessedEncoderInputs as PreprocessedImages, PromptReplacement, Tokenizer as TokenResolver, - TrackedMedia, VisionPreProcessor as ImagePreProcessor, - VisionProcessorRegistry as ImageProcessorRegistry, + AsyncMultiModalTracker, FieldLayout, ImageFrame, MediaConnector, MediaConnectorConfig, + MediaContentPart, Modality, ModelMetadata, ModelProcessorSpec, ModelRegistry, + PreProcessorConfig, PreprocessedEncoderInputs, PromptReplacement, Tokenizer as TokenResolver, + TrackedMedia, VideoClip, VisionPreProcessor, VisionProcessorRegistry, }; +use thiserror_ext::AsReport as _; use tracing::warn; use vllm_engine_core_client::protocol::dtype::ModelDtype; -use vllm_engine_core_client::protocol::multimodal::{ - MmBatchedField, MmFeatureSpec, MmFeatures, MmField, MmFieldElem, MmFlatField, MmKwargsItem, - MmSharedField, MmSlice, PlaceholderRange, SliceSpec, -}; -use vllm_engine_core_client::protocol::tensor::WireTensor; +use vllm_engine_core_client::protocol::multimodal::{MmFeatureSpec, MmFeatures, MmKwargsItem}; use vllm_text::Prompt; use vllm_text::tokenizer::{DynTokenizer, Tokenizer}; @@ -36,14 +32,20 @@ use crate::error::{Error, Result, bail_multimodal, multimodal}; use crate::renderer::RenderedPrompt; use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest}; +mod expand; +mod image; mod tensor; +mod video; + +use self::expand::expand_prompt_token_ids; /// Resolved multimodal support for one loaded model. #[derive(Clone)] pub struct MultimodalModelInfo { context: MultimodalModelContext, spec: ResolvedMultimodalSpec, - image_processor: ResolvedImageProcessor, + image: Option, + video: Option, media_connector: Arc, } @@ -75,91 +77,171 @@ impl MultimodalModelContext { REGISTRY.lookup(&self.metadata()) } - /// Resolve a static image preprocessor for one loaded model. - fn resolve_image_processor(&self) -> Option<&'static dyn ImagePreProcessor> { - static REGISTRY: LazyLock = - LazyLock::new(ImageProcessorRegistry::with_defaults); + /// Resolve a static vision preprocessor for one loaded model. + /// + /// The vision preprocessor serves both the image and video modalities. + fn resolve_vision_processor(&self) -> Option<&'static dyn VisionPreProcessor> { + static REGISTRY: LazyLock = + LazyLock::new(VisionProcessorRegistry::with_defaults); REGISTRY.find(&self.model_id, self.model_type.as_deref()) } } -/// Static model-specific prompt and tensor-layout behavior. +/// Static model-specific tensor-layout behavior shared across modalities. #[derive(Clone)] struct ResolvedMultimodalSpec { raw: &'static dyn ModelProcessorSpec, - placeholder_token: String, - placeholder_marker_token_id: u32, - placeholder_embed_token_id: u32, field_layouts: HashMap, keep_on_cpu_keys: HashSet, } impl ResolvedMultimodalSpec { - fn new(raw: &'static dyn ModelProcessorSpec, context: &MultimodalModelContext) -> Result { - let metadata = context.metadata(); - let placeholder_token = - raw.placeholder_token(&metadata).map_err(|error| multimodal!("{error}"))?; - // This is the rendered prompt marker, so resolve it from the token - // string itself. Do not use `ModelProcessorSpec::placeholder_token_id()`: - // for some specs that ID is the replacement vision/patch token, - // not necessarily the token ID of `placeholder_token`. - let placeholder_marker_token_id = - context.tokenizer().token_to_id(&placeholder_token).ok_or_else(|| { - multimodal!( - "placeholder token `{placeholder_token}` is not in the tokenizer vocabulary" - ) - })?; - let placeholder_embed_token_id = - raw.placeholder_token_id(&metadata).map_err(|error| multimodal!("{error}"))? as u32; - - Ok(Self { + fn new(raw: &'static dyn ModelProcessorSpec) -> Self { + Self { raw, - placeholder_token, - placeholder_marker_token_id, - placeholder_embed_token_id, field_layouts: raw.field_layouts(), keep_on_cpu_keys: raw.keep_on_cpu_keys().into_iter().collect(), - }) + } } - fn prompt_replacements( + fn prompt_replacements_for( &self, context: &MultimodalModelContext, - preprocessed: &PreprocessedImages, + preprocessed: &PreprocessedEncoderInputs, + modality: Modality, ) -> Result> { - self.raw - .prompt_replacements(&context.metadata(), preprocessed) - .map_err(|error| multimodal!("{error}")) + Ok(self.raw.prompt_replacements_for(&context.metadata(), preprocessed, modality)?) } } -/// Static image preprocessor plus its loaded config. +/// Resolved placeholder tokens for one modality. #[derive(Clone)] -struct ResolvedImageProcessor { - raw: &'static dyn ImagePreProcessor, +struct ResolvedPlaceholder { + token: String, + /// The token ID emitted for `token` in the rendered prompt. + marker_token_id: u32, + /// The model-declared embed token ID marked in `is_embed` masks. + embed_token_id: u32, +} + +impl ResolvedPlaceholder { + fn resolve( + raw: &'static dyn ModelProcessorSpec, + context: &MultimodalModelContext, + modality: Modality, + ) -> Result { + let metadata = context.metadata(); + let token = raw.placeholder_token_for(&metadata, modality)?; + // This is the rendered prompt marker, so resolve it from the token + // string itself. Do not use `ModelProcessorSpec::placeholder_token_id_for()`: + // for some specs that ID is the replacement vision/patch token, + // not necessarily the token ID of the placeholder token. + let marker_token_id = context.tokenizer().token_to_id(&token).ok_or_else(|| { + multimodal!("placeholder token `{token}` is not in the tokenizer vocabulary") + })?; + let embed_token_id = raw.placeholder_token_id_for(&metadata, modality)? as u32; + + Ok(Self { + token, + marker_token_id, + embed_token_id, + }) + } +} + +/// Static per-modality vision preprocessor plus its loaded config and +/// resolved placeholder tokens. +#[derive(Clone)] +struct ModalitySupport { + placeholder: ResolvedPlaceholder, + processor: &'static dyn VisionPreProcessor, config: PreProcessorConfig, } -/// Request-scoped fetched media, kept together with tracker UUID metadata. -struct FetchedImageMedia { - frames: Vec>, - uuids: Vec>, +/// Model-repo config file locations consumed by multimodal support. +#[derive(Debug, Default, Clone, Copy)] +pub struct MultimodalConfigFiles<'a> { + pub config: Option<&'a Path>, + pub preprocessor_config: Option<&'a Path>, + /// Video-specific preprocessor config (`video_preprocessor_config.json`). + pub video_preprocessor_config: Option<&'a Path>, + /// Combined processor config (`processor_config.json`), whose modality + /// sections are fallback preprocessor config sources. + pub processor_config: Option<&'a Path>, +} + +/// Load a modality's dedicated preprocessor config, falling back to its section +/// in the combined processor config. +fn load_preprocessor_config( + dedicated_path: Option<&Path>, + dedicated_name: &str, + processor_config_path: Option<&Path>, + processor_section: &str, +) -> Result> { + if let Some(path) = dedicated_path { + let text = fs::read_to_string(path) + .map_err(|error| multimodal!("failed to read {dedicated_name}: {error}"))?; + let config = PreProcessorConfig::from_json(&text) + .map_err(|error| multimodal!("failed to parse {dedicated_name}: {error}"))?; + return Ok(Some(config)); + } + + let Some(path) = processor_config_path else { + return Ok(None); + }; + let text = fs::read_to_string(path) + .map_err(|error| multimodal!("failed to read processor_config.json: {error}"))?; + let value: serde_json::Value = serde_json::from_str(&text) + .map_err(|error| multimodal!("failed to parse processor_config.json: {error}"))?; + let Some(processor) = value.get(processor_section) else { + return Ok(None); + }; + let config = PreProcessorConfig::from_value(processor.clone()).map_err(|error| { + multimodal!("failed to parse {processor_section} from processor_config.json: {error}") + })?; + Ok(Some(config)) +} + +/// Request-scoped fetched media, split per modality with tracker UUID +/// metadata preserved in request order. +struct FetchedMedia { + images: Vec>, + image_uuids: Vec>, + videos: Vec>, + video_uuids: Vec>, +} + +/// One modality's preprocessed output, ready for the shared expansion and +/// feature-assembly tail. +struct PreparedMedia { + modality: Modality, + placeholder: ResolvedPlaceholder, + /// One replacement per media item, in request order. + replacements: Vec, + /// One entry per media item, aligned with `replacements`. + items: Vec, +} + +/// One media item's complete engine kwargs plus identity metadata. +struct PreparedItem { + data: MmKwargsItem, + hash: String, + uuid: Option, } impl MultimodalModelInfo { /// Load and resolve multimodal support from model files. /// - /// Returns `Ok(Some(_))` only when both the model spec and image processor - /// are registered. File read/parse failures are real errors; unsupported - /// model families are logged and returned as `Ok(None)`. + /// Returns `Ok(Some(_))` only when the model spec is registered and at + /// least one modality resolves. File read/parse failures are real errors; + /// unsupported model families are logged and returned as `Ok(None)`. pub fn from_paths( model_id: String, model_type: Option, - config_path: Option<&Path>, - preprocessor_config_path: Option<&Path>, + files: MultimodalConfigFiles<'_>, tokenizer: DynTokenizer, ) -> Result> { - let config = match config_path { + let config = match files.config { Some(path) => { let text = fs::read_to_string(path) .map_err(|error| multimodal!("failed to read config.json: {error}"))?; @@ -168,17 +250,20 @@ impl MultimodalModelInfo { } None => serde_json::Value::Object(Default::default()), }; - let preprocessor_config = match preprocessor_config_path { - Some(path) => { - let text = fs::read_to_string(path).map_err(|error| { - multimodal!("failed to read preprocessor_config.json: {error}") - })?; - PreProcessorConfig::from_json(&text).map_err(|error| { - multimodal!("failed to parse preprocessor_config.json: {error}") - })? - } - None => PreProcessorConfig::default(), - }; + let image_preprocessor_config = load_preprocessor_config( + files.preprocessor_config, + "preprocessor_config.json", + files.processor_config, + "image_processor", + )? + .unwrap_or_default(); + let video_preprocessor_config = load_preprocessor_config( + files.video_preprocessor_config, + "video_preprocessor_config.json", + files.processor_config, + "video_processor", + )? + .unwrap_or_else(|| image_preprocessor_config.clone()); let context = MultimodalModelContext { model_id, @@ -187,7 +272,21 @@ impl MultimodalModelInfo { tokenizer: TokenizerResolver(tokenizer), }; - let Some(spec) = context.resolve_model_spec() else { + Self::from_loaded( + context, + image_preprocessor_config, + video_preprocessor_config, + ) + } + + /// Resolve multimodal support from an assembled context and parsed + /// preprocessor configs. + fn from_loaded( + context: MultimodalModelContext, + image_preprocessor_config: PreProcessorConfig, + video_preprocessor_config: PreProcessorConfig, + ) -> Result> { + let Some(raw_spec) = context.resolve_model_spec() else { warn!( model_id = context.model_id, model_type = context.model_type, @@ -195,47 +294,99 @@ impl MultimodalModelInfo { ); return Ok(None); }; - let spec = ResolvedMultimodalSpec::new(spec, &context)?; - let Some(image_processor) = context.resolve_image_processor() else { + let Some(processor) = context.resolve_vision_processor() else { warn!( model_id = context.model_id, model_type = context.model_type, - "image processor is not registered; disabling multimodal support for this model" + "vision processor is not registered; disabling multimodal support for this model" ); return Ok(None); }; - let media_connector = Arc::new( - MediaConnector::new(reqwest::Client::new(), MediaConnectorConfig::default()) - .map_err(|error| multimodal!("{error}"))?, - ); + // Warn and disable the modality if the placeholder resolution fails. + let resolve_placeholder = + |modality: Modality| match ResolvedPlaceholder::resolve(raw_spec, &context, modality) { + Ok(placeholder) => Some(placeholder), + Err(error) => { + warn!( + model_id = context.model_id, + %modality, + error = %error.as_report(), + "placeholder tokens did not resolve; disabling this modality for this model" + ); + None + } + }; + + let image = resolve_placeholder(Modality::Image).map(|placeholder| ModalitySupport { + placeholder, + processor, + config: image_preprocessor_config, + }); + + let video = resolve_placeholder(Modality::Video).and_then(|placeholder| { + // Placeholder expansion attributes markers to modalities by token + // ID, so a marker shared with the image modality is ambiguous. + let image_marker = image.as_ref().map(|image| image.placeholder.marker_token_id); + if image_marker == Some(placeholder.marker_token_id) { + warn!( + model_id = context.model_id, + token = placeholder.token, + "video placeholder token collides with the image placeholder; disabling video support for this model" + ); + None + } else { + Some(ModalitySupport { + placeholder, + processor, + config: video_preprocessor_config, + }) + } + }); + + if image.is_none() && video.is_none() { + warn!( + model_id = context.model_id, + model_type = context.model_type, + "no multimodal modality resolved; disabling multimodal support for this model" + ); + return Ok(None); + } + + let media_connector = Arc::new(MediaConnector::new( + reqwest::Client::new(), + MediaConnectorConfig::default(), + )?); Ok(Some(Self { context, - spec, - image_processor: ResolvedImageProcessor { - raw: image_processor, - config: preprocessor_config, - }, + spec: ResolvedMultimodalSpec::new(raw_spec), + image, + video, media_connector, })) } - /// Return the template-visible placeholder token for this model. + /// Return the template-visible placeholder token for one modality, when + /// this model supports it. /// - /// The HF renderer uses this token while flattening image content in string - /// content format. - pub fn placeholder_token(&self) -> &str { - &self.spec.placeholder_token + /// The HF renderer uses these tokens while flattening media content in + /// string content format. + pub fn placeholder_token(&self, modality: Modality) -> Option<&str> { + match modality { + Modality::Image => self.image.as_ref()?.placeholder.token.as_str().into(), + Modality::Video => self.video.as_ref()?.placeholder.token.as_str().into(), + _ => None, + } } } /// Finalize a rendered chat prompt into text-generation input. /// /// Text-only requests pass through unchanged as `Prompt::Text`. Multimodal -/// requests are tokenized in chat, their image placeholders are expanded, and -/// preprocessed image features are attached for engine-core transport. +/// requests are tokenized in chat, their media placeholders are expanded, and +/// preprocessed media features are attached for engine-core transport. pub(crate) async fn finalize_rendered_prompt( request: &ChatRequest, rendered: RenderedPrompt, @@ -260,7 +411,7 @@ pub(crate) async fn finalize_rendered_prompt( Ok((Prompt::TokenIds(prompt_token_ids), Some(prepared))) } -/// Extract image media parts from chat messages in message/content order. +/// Extract media parts from chat messages in message/content order. /// /// Assistant history is skipped because generated assistant blocks are already /// represented as text for prompt rendering in this crate. @@ -289,6 +440,12 @@ fn extract_media_parts(request: &ChatRequest) -> Result> { detail: *detail, uuid: uuid.clone(), }), + ChatContentPart::VideoUrl { video_url, uuid } => { + all_parts.push(MediaContentPart::VideoUrl { + url: video_url.clone(), + uuid: uuid.clone(), + }) + } } } } @@ -296,8 +453,8 @@ fn extract_media_parts(request: &ChatRequest) -> Result> { } impl MultimodalModelInfo { - /// Run media fetch, image preprocessing, prompt expansion, and feature - /// build. + /// Run media fetch, per-modality preprocessing, prompt expansion, and + /// feature build. /// /// `prompt_token_ids` is mutated in place because placeholder expansion /// changes both the final prompt and the offsets recorded in @@ -313,12 +470,47 @@ impl MultimodalModelInfo { } let media_parts_len = media_parts.len(); - let fetched = self.fetch_images(media_parts).await?; - let preprocessed = self.preprocess_images(&fetched.frames).await?; - let replacements = self.spec.prompt_replacements(&self.context, &preprocessed)?; - let ranges = self.expand_prompt_tokens(prompt_token_ids, replacements)?; + // TODO: enforce per-modality item-count limits, aligned with the + // engine's `--limit-mm-per-prompt` semantics. + let fetched = self.fetch_media(media_parts).await?; + + let mut prepared = Vec::new(); + if !fetched.images.is_empty() { + prepared + .push(self.prepare_images(fetched.images, fetched.image_uuids, model_dtype).await?); + } + if !fetched.videos.is_empty() { + prepared + .push(self.prepare_videos(fetched.videos, fetched.video_uuids, model_dtype).await?); + } + + let mut ranges = expand_prompt_token_ids(prompt_token_ids, &prepared)?; + + let mut features = Vec::with_capacity(media_parts_len); + for media in prepared { + let media_ranges = ranges.remove(&media.modality).unwrap_or_default(); + if media_ranges.len() != media.items.len() { + bail_multimodal!( + "number of expanded `{}` placeholders {} does not match number of media items {}", + media.modality, + media_ranges.len(), + media.items.len() + ); + } + for (item, range) in izip!(media.items, media_ranges) { + features.push(MmFeatureSpec { + data: Some(item.data), + modality: media.modality.to_string(), + identifier: item.uuid.unwrap_or_else(|| item.hash.clone()), + mm_position: range, + mm_hash: Some(item.hash), + }); + } + } + // Mirror the Python frontend (`argsort_mm_positions`): features are + // ordered by their placeholder position in the prompt. + features.sort_by_key(|feature| feature.mm_position.offset); - let features = self.build_features(preprocessed, fetched, ranges, model_dtype)?; if features.len() != media_parts_len { bail_multimodal!( "number of built multimodal features {} does not match number of media parts {}", @@ -329,219 +521,51 @@ impl MultimodalModelInfo { Ok(features) } - /// Fetch all image parts and preserve their request-order UUID metadata. - async fn fetch_images(&self, media_parts: Vec) -> Result { + /// Fetch all media parts and split them per modality, preserving their + /// request-order UUID metadata. + async fn fetch_media(&self, media_parts: Vec) -> Result { let mut tracker = AsyncMultiModalTracker::new(Arc::clone(&self.media_connector)); for part in media_parts { - tracker.push_part(part).map_err(|error| multimodal!("{error}"))?; + tracker.push_part(part)?; } - let tracker_output = tracker.finalize().await.map_err(|error| multimodal!("{error}"))?; - let images = tracker_output.data.get(&Modality::Image).cloned().unwrap_or_default(); - let uuids = tracker_output.uuids.get(&Modality::Image).cloned().unwrap_or_default(); + let mut tracker_output = tracker.finalize().await?; - let frames = images + let images = tracker_output + .data + .remove(&Modality::Image) + .unwrap_or_default() .into_iter() .map(|media| match media { TrackedMedia::Image(frame) => Ok(frame), - _ => Err(Error::UnsupportedMultimodalContent("non-image")), + _ => Err(multimodal!( + "tracker returned non-image media for the image modality" + )), }) .collect::>>()?; + let image_uuids = tracker_output.uuids.remove(&Modality::Image).unwrap_or_default(); - Ok(FetchedImageMedia { frames, uuids }) - } + let videos = tracker_output + .data + .remove(&Modality::Video) + .unwrap_or_default() + .into_iter() + .map(|media| match media { + TrackedMedia::Video(clip) => Ok(clip), + _ => Err(multimodal!( + "tracker returned non-video media for the video modality" + )), + }) + .collect::>>()?; + let video_uuids = tracker_output.uuids.remove(&Modality::Video).unwrap_or_default(); - /// Preprocess fetched image frames with the model's resolved image - /// processor. - /// - /// The processor work is CPU-heavy relative to request wiring, so it runs - /// in a blocking task and returns owned tensors ready for wire - /// conversion. - async fn preprocess_images( - &self, - image_frames: &[Arc], - ) -> Result { - let config = self.image_processor.config.clone(); - let processor = self.image_processor.raw; - let images = image_frames.iter().map(|frame| frame.data().clone()).collect::>(); - - // TODO: is it still necessary given that we've already in a dedicated runtime? - tokio::task::spawn_blocking(move || { - processor.preprocess(&images, &config).map_err(|error| multimodal!("{error}")) + Ok(FetchedMedia { + images, + image_uuids, + videos, + video_uuids, }) - .await - .map_err(|error| multimodal!("image preprocessing task failed: {error}"))? - } - - /// Replace rendered placeholder markers with model-specific replacement - /// tokens. - /// - /// Replacements are consumed in order, matching the original media-part - /// order. The returned ranges point into the already-expanded prompt. - fn expand_prompt_tokens( - &self, - prompt_token_ids: &mut Vec, - replacements: Vec, - ) -> Result> { - expand_prompt_token_ids( - prompt_token_ids, - replacements, - self.spec.placeholder_marker_token_id, - self.spec.placeholder_embed_token_id, - &self.spec.placeholder_token, - ) } - - /// Convert preprocessed image tensors into engine-core multimodal features. - /// - /// One `MmFeatureSpec` is produced per image. Tensor fields are - /// sliced according to the model spec's field layout declarations. - fn build_features( - &self, - preprocessed: PreprocessedImages, - images: FetchedImageMedia, - ranges: Vec, - model_dtype: ModelDtype, - ) -> Result { - let len = images.frames.len(); - let tensors = tensor::collect_tensors(preprocessed, model_dtype)?; - - let mut features = Vec::with_capacity(images.frames.len()); - for (index, (frame, uuid, range)) in izip!(images.frames, images.uuids, ranges).enumerate() - { - let mut data = MmKwargsItem::new(); - for (key, tensor) in &tensors { - let keep_on_cpu = self.spec.keep_on_cpu_keys.contains(key); - let (value, field) = match self.spec.field_layouts.get(key) { - Some(FieldLayout::Batched) => ( - tensor.batched_value_at(index)?, - MmField::Batched(MmBatchedField { keep_on_cpu }), - ), - Some(FieldLayout::Flat { sizes_key }) => { - let sizes = tensors.get(sizes_key).ok_or_else(|| { - multimodal!("flat tensor sizes key `{sizes_key}` is missing") - })?; - let (start, end) = tensor::flat_range_for_index(sizes, sizes_key, index)?; - ( - tensor.flat_value_range(start, end)?, - MmField::Flat(MmFlatField { - slices: vec![MmSlice::Slice(SliceSpec { - start: Some(0), - stop: Some((end - start) as isize), - step: None, - })], - dim: 0, - keep_on_cpu, - }), - ) - } - None => ( - tensor.clone(), - MmField::Shared(MmSharedField { - batch_size: len, - keep_on_cpu, - }), - ), - }; - - data.insert( - key.clone(), - MmFieldElem { - data: Some(value.try_into()?), - field, - }, - ); - } - - let hash = frame.hash.clone(); - features.push(MmFeatureSpec { - data: Some(data), - modality: "image".to_string(), - identifier: uuid.unwrap_or_else(|| hash.clone()), - mm_position: range, - mm_hash: Some(hash), - }); - } - - Ok(features) - } -} - -fn expand_prompt_token_ids( - prompt_token_ids: &mut Vec, - replacements: Vec, - placeholder_marker_token_id: u32, - placeholder_embed_token_id: u32, - placeholder_token: &str, -) -> Result> { - if replacements.is_empty() { - return Ok(Vec::new()); - } - - let replacement_growth = replacements.iter().fold(0usize, |total, replacement| { - total.saturating_add(replacement.tokens.len().saturating_sub(1)) - }); - let mut expanded = - Vec::with_capacity(prompt_token_ids.len().saturating_add(replacement_growth)); - let mut ranges = Vec::with_capacity(replacements.len()); - let mut cursor = 0usize; - - for replacement in replacements { - if replacement.modality != Modality::Image { - bail_multimodal!( - "unsupported prompt replacement modality `{}`", - replacement.modality - ); - } - - let offset = find_next_token(prompt_token_ids, placeholder_marker_token_id, cursor) - .ok_or_else(|| { - multimodal!( - "placeholder token `{placeholder_token}` was not found in tokenized prompt" - ) - })?; - - if replacement.tokens.is_empty() { - bail_multimodal!("placeholder token `{placeholder_token}` expanded to no tokens"); - } - - let replacement_len = replacement.tokens.len(); - let is_embed = { - let mask = replacement - .tokens - .iter() - .map(|&token| token as u32 == placeholder_embed_token_id) - .collect::>(); - WireTensor::from_bool(vec![replacement_len], mask).map_err(Error::Multimodal)? - }; - - expanded.extend_from_slice(&prompt_token_ids[cursor..offset]); - let expanded_offset = expanded.len(); - expanded.extend(replacement.tokens.into_iter().map(|token| token as u32)); - ranges.push(PlaceholderRange { - offset: expanded_offset, - length: replacement_len, - is_embed: Some(is_embed), - }); - cursor = offset + 1; - } - - expanded.extend_from_slice(&prompt_token_ids[cursor..]); - *prompt_token_ids = expanded; - - Ok(ranges) -} - -/// Find `needle` in `haystack`, starting at `start`. -/// -/// This is intentionally order-preserving rather than a global replace: each -/// image consumes the next placeholder occurrence. -fn find_next_token(haystack: &[u32], needle: u32, start: usize) -> Option { - haystack - .get(start..)? - .iter() - .position(|token| *token == needle) - .map(|offset| start + offset) } /// Adapter from the frontend tokenizer trait to `llm-multimodal`. @@ -566,18 +590,19 @@ impl TokenResolver for TokenizerResolver { mod tests { use std::sync::Arc; - use llm_multimodal::TokenId; - use vllm_engine_core_client::protocol::tensor::WireArrayData; use vllm_tokenizer::test_utils::TestTokenizer; use super::*; - const LLAMA4_IMAGE_START_ID: u32 = 200088; - const LLAMA4_IMAGE_END_ID: u32 = 200089; - const LLAMA4_IMAGE_ID: u32 = 200090; - const LLAMA4_PATCH_ID: u32 = 200092; - const LLAMA4_TILE_X_SEPARATOR_ID: u32 = 200093; - const LLAMA4_TILE_Y_SEPARATOR_ID: u32 = 200094; + pub(super) const LLAMA4_IMAGE_START_ID: u32 = 200088; + pub(super) const LLAMA4_IMAGE_END_ID: u32 = 200089; + pub(super) const LLAMA4_IMAGE_ID: u32 = 200090; + pub(super) const LLAMA4_PATCH_ID: u32 = 200092; + pub(super) const LLAMA4_TILE_X_SEPARATOR_ID: u32 = 200093; + pub(super) const LLAMA4_TILE_Y_SEPARATOR_ID: u32 = 200094; + + pub(super) const QWEN3_IMAGE_PAD_ID: u32 = 151655; + pub(super) const QWEN3_VIDEO_PAD_ID: u32 = 151656; fn llama4_tokenizer() -> TestTokenizer { TestTokenizer::new() @@ -589,33 +614,31 @@ mod tests { .with_regular_token("<|tile_y_separator|>", LLAMA4_TILE_Y_SEPARATOR_ID) } - fn test_info(model_type: &str, config: serde_json::Value) -> MultimodalModelInfo { + pub(super) fn qwen3_vl_tokenizer() -> TestTokenizer { + TestTokenizer::new() + .with_regular_token("<|image_pad|>", QWEN3_IMAGE_PAD_ID) + .with_regular_token("<|video_pad|>", QWEN3_VIDEO_PAD_ID) + } + + fn test_info( + model_type: &str, + config: serde_json::Value, + tokenizer: TestTokenizer, + ) -> MultimodalModelInfo { let context = MultimodalModelContext { model_id: format!("{model_type}-test"), model_type: Some(model_type.to_string()), config, - tokenizer: TokenizerResolver(Arc::new(llama4_tokenizer())), + tokenizer: TokenizerResolver(Arc::new(tokenizer)), }; - let spec = context - .resolve_model_spec() - .unwrap_or_else(|| panic!("{model_type} spec should match")); - let spec = ResolvedMultimodalSpec::new(spec, &context).unwrap(); - let raw_image_processor = context - .resolve_image_processor() - .unwrap_or_else(|| panic!("{model_type} image processor should match")); - let media_connector = Arc::new( - MediaConnector::new(reqwest::Client::new(), MediaConnectorConfig::default()).unwrap(), - ); - MultimodalModelInfo { + MultimodalModelInfo::from_loaded( context, - spec, - image_processor: ResolvedImageProcessor { - raw: raw_image_processor, - config: PreProcessorConfig::default(), - }, - media_connector, - } + PreProcessorConfig::default(), + PreProcessorConfig::default(), + ) + .unwrap() + .unwrap_or_else(|| panic!("{model_type} multimodal support should resolve")) } fn llama4_info() -> MultimodalModelInfo { @@ -624,173 +647,96 @@ mod tests { "image_token_index": LLAMA4_PATCH_ID, "vision_config": {"image_size": 336, "patch_size": 14} }); - test_info("llama4", config) + test_info("llama4", config, llama4_tokenizer()) } - fn llama4_single_tile_replacement() -> PromptReplacement { - PromptReplacement::sequence( - Modality::Image, - "<|image|>", - vec![ - LLAMA4_IMAGE_START_ID as TokenId, - LLAMA4_IMAGE_ID as TokenId, - LLAMA4_PATCH_ID as TokenId, - LLAMA4_PATCH_ID as TokenId, - LLAMA4_IMAGE_END_ID as TokenId, - ], - ) + pub(super) fn qwen3_vl_info() -> MultimodalModelInfo { + let config = serde_json::json!({ + "model_type": "qwen3_vl", + "image_token_id": QWEN3_IMAGE_PAD_ID, + "video_token_id": QWEN3_VIDEO_PAD_ID, + "vision_start_token_id": 151652, + "vision_end_token_id": 151653, + "vision_config": {"patch_size": 16} + }); + test_info("qwen3_vl", config, qwen3_vl_tokenizer()) } - fn llama4_multi_tile_replacement() -> PromptReplacement { - PromptReplacement::sequence( - Modality::Image, - "<|image|>", - vec![ - LLAMA4_IMAGE_START_ID as TokenId, - LLAMA4_PATCH_ID as TokenId, - LLAMA4_TILE_X_SEPARATOR_ID as TokenId, - LLAMA4_PATCH_ID as TokenId, - LLAMA4_TILE_Y_SEPARATOR_ID as TokenId, - LLAMA4_IMAGE_ID as TokenId, - LLAMA4_PATCH_ID as TokenId, - LLAMA4_IMAGE_END_ID as TokenId, - ], + #[test] + fn from_paths_resolves_image_config_from_processor_config() { + let dir = tempfile::tempdir().unwrap(); + let config_path = dir.path().join("config.json"); + std::fs::write( + &config_path, + serde_json::json!({ + "model_type": "qwen3_vl", + "image_token_id": QWEN3_IMAGE_PAD_ID, + }) + .to_string(), ) - } + .unwrap(); + let processor_config_path = dir.path().join("processor_config.json"); + std::fs::write( + &processor_config_path, + r#"{"image_processor":{"size":{"shortest_edge":64}}}"#, + ) + .unwrap(); + + let info = MultimodalModelInfo::from_paths( + "qwen3-vl-test".to_string(), + Some("qwen3_vl".to_string()), + MultimodalConfigFiles { + config: Some(&config_path), + processor_config: Some(&processor_config_path), + ..Default::default() + }, + Arc::new(qwen3_vl_tokenizer()), + ) + .unwrap() + .unwrap(); - fn assert_bool_mask(range: &PlaceholderRange, expected: &[bool]) { - let tensor = range.is_embed.as_ref().expect("is_embed mask"); - assert_eq!(tensor.dtype, "bool"); - assert_eq!(tensor.shape, vec![expected.len()]); - assert_eq!( - tensor.data, - WireArrayData::RawView(expected.iter().map(|value| u8::from(*value)).collect()) - ); + assert_eq!(info.image.unwrap().config.get_shortest_edge(), Some(64)); } #[test] - fn expand_prompt_tokens_marks_only_llama4_patch_tokens_as_embed() { - let info = llama4_info(); - let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; - let replacements = vec![llama4_multi_tile_replacement()]; - - let ranges = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap(); + fn qwen3_vl_resolves_image_and_video_support() { + let info = qwen3_vl_info(); assert_eq!( - prompt_token_ids, - vec![ - 1, - LLAMA4_IMAGE_START_ID, - LLAMA4_PATCH_ID, - LLAMA4_TILE_X_SEPARATOR_ID, - LLAMA4_PATCH_ID, - LLAMA4_TILE_Y_SEPARATOR_ID, - LLAMA4_IMAGE_ID, - LLAMA4_PATCH_ID, - LLAMA4_IMAGE_END_ID, - 2, - ] + info.placeholder_token(Modality::Image), + Some("<|image_pad|>") ); - assert_eq!(ranges[0].offset, 1); - assert_eq!(ranges[0].length, 8); - assert_bool_mask( - &ranges[0], - &[false, true, false, true, false, false, true, false], + assert_eq!( + info.placeholder_token(Modality::Video), + Some("<|video_pad|>") + ); + assert_ne!( + info.image.as_ref().unwrap().placeholder.marker_token_id, + info.video.as_ref().unwrap().placeholder.marker_token_id, ); } #[test] - fn expand_prompt_tokens_errors_when_placeholder_missing() { - let info = llama4_info(); - let mut prompt_token_ids = vec![1, 2, 3]; - let replacements = vec![llama4_single_tile_replacement()]; - - let error = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap_err(); - - assert!(matches!(error, Error::Multimodal(message) if message.contains("not found"))); - } - - #[test] - fn expand_prompt_tokens_ignores_empty_replacements() { - let info = llama4_info(); - let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; - let original_prompt_token_ids = prompt_token_ids.clone(); - - let ranges = info.expand_prompt_tokens(&mut prompt_token_ids, Vec::new()).unwrap(); - - assert!(ranges.is_empty()); - assert_eq!(prompt_token_ids, original_prompt_token_ids); - } - - #[test] - fn expand_prompt_tokens_leaves_prompt_unchanged_when_later_placeholder_missing() { - let info = llama4_info(); - let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; - let original_prompt_token_ids = prompt_token_ids.clone(); - let replacements = vec![ - llama4_single_tile_replacement(), - llama4_single_tile_replacement(), - ]; - - let error = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap_err(); - - assert!(matches!(error, Error::Multimodal(message) if message.contains("not found"))); - assert_eq!(prompt_token_ids, original_prompt_token_ids); - } + fn qwen3_vl_without_video_token_id_disables_video_support_only() { + let config = serde_json::json!({ + "model_type": "qwen3_vl", + "image_token_id": QWEN3_IMAGE_PAD_ID, + "vision_config": {"patch_size": 16} + }); + let info = test_info("qwen3_vl", config, qwen3_vl_tokenizer()); - #[test] - fn expand_prompt_tokens_errors_when_replacement_is_empty() { - let info = llama4_info(); - let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; - let original_prompt_token_ids = prompt_token_ids.clone(); - let replacements = vec![PromptReplacement::sequence( - Modality::Image, - "<|image|>", - Vec::new(), - )]; - - let error = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap_err(); - - assert!( - matches!(error, Error::Multimodal(message) if message.contains("expanded to no tokens")) + assert_eq!( + info.placeholder_token(Modality::Image), + Some("<|image_pad|>") ); - assert_eq!(prompt_token_ids, original_prompt_token_ids); + assert_eq!(info.placeholder_token(Modality::Video), None); } #[test] - fn expand_prompt_tokens_skips_llama4_image_marker_inside_replacement() { + fn llama4_resolves_image_support_only() { let info = llama4_info(); - let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2, LLAMA4_IMAGE_ID, 3]; - let replacements = vec![ - llama4_single_tile_replacement(), - llama4_single_tile_replacement(), - ]; - - let ranges = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap(); - assert_eq!( - prompt_token_ids, - vec![ - 1, - LLAMA4_IMAGE_START_ID, - LLAMA4_IMAGE_ID, - LLAMA4_PATCH_ID, - LLAMA4_PATCH_ID, - LLAMA4_IMAGE_END_ID, - 2, - LLAMA4_IMAGE_START_ID, - LLAMA4_IMAGE_ID, - LLAMA4_PATCH_ID, - LLAMA4_PATCH_ID, - LLAMA4_IMAGE_END_ID, - 3, - ] - ); - assert_eq!(ranges[0].offset, 1); - assert_eq!(ranges[0].length, 5); - assert_bool_mask(&ranges[0], &[false, false, true, true, false]); - assert_eq!(ranges[1].offset, 7); - assert_eq!(ranges[1].length, 5); - assert_bool_mask(&ranges[1], &[false, false, true, true, false]); + assert_eq!(info.placeholder_token(Modality::Image), Some("<|image|>")); + assert_eq!(info.placeholder_token(Modality::Video), None); } } diff --git a/rust/src/chat/src/multimodal/expand.rs b/rust/src/chat/src/multimodal/expand.rs new file mode 100644 index 000000000000..36a2a9bdb775 --- /dev/null +++ b/rust/src/chat/src/multimodal/expand.rs @@ -0,0 +1,446 @@ +//! Prompt placeholder expansion shared across modalities. + +use std::collections::{HashMap, VecDeque}; + +use llm_multimodal::{Modality, PromptReplacement}; +use vllm_engine_core_client::protocol::multimodal::PlaceholderRange; +use vllm_engine_core_client::protocol::tensor::WireTensor; + +use super::PreparedMedia; +use crate::error::{Error, Result, bail_multimodal}; + +/// One modality's queue of pending placeholder replacements for prompt +/// expansion. +struct ExpansionLane<'a> { + modality: Modality, + marker_token_id: u32, + embed_token_id: u32, + placeholder_token: String, + replacements: VecDeque<&'a PromptReplacement>, +} + +impl<'a> ExpansionLane<'a> { + fn from_prepared(media: &'a PreparedMedia) -> Option { + if media.replacements.is_empty() { + return None; + } + + Some(Self { + modality: media.modality, + marker_token_id: media.placeholder.marker_token_id, + embed_token_id: media.placeholder.embed_token_id, + placeholder_token: media.placeholder.token.clone(), + replacements: media.replacements.iter().collect(), + }) + } +} + +/// Replace rendered placeholder markers with model-specific replacement +/// tokens across all modalities in one left-to-right pass. +/// +/// Each prepared modality consumes its own marker occurrences in order, +/// matching the original media-part order within that modality; markers of +/// different modalities may interleave freely. +/// +/// The returned ranges point into the already-expanded prompt, grouped per +/// modality in item order. +pub(super) fn expand_prompt_token_ids( + prompt_token_ids: &mut Vec, + prepared: &[PreparedMedia], +) -> Result>> { + let mut lanes = prepared.iter().filter_map(ExpansionLane::from_prepared).collect::>(); + if lanes.is_empty() { + return Ok(HashMap::new()); + } + + let replacement_growth = lanes + .iter() + .flat_map(|lane| lane.replacements.iter()) + .fold(0usize, |total, replacement| { + total.saturating_add(replacement.tokens.len().saturating_sub(1)) + }); + let expanded_len = prompt_token_ids.len().saturating_add(replacement_growth); + + let mut expanded = Vec::with_capacity(expanded_len); + let mut ranges = HashMap::>::new(); + + for &token in prompt_token_ids.iter() { + let lane = lanes + .iter_mut() + .find(|lane| lane.marker_token_id == token && !lane.replacements.is_empty()); + let Some(lane) = lane else { + expanded.push(token); + continue; + }; + + let replacement = lane.replacements.pop_front().expect("lane queue is non-empty"); + debug_assert_eq!(replacement.modality, lane.modality); + if replacement.tokens.is_empty() { + bail_multimodal!( + "placeholder token `{}` expanded to no tokens", + lane.placeholder_token + ); + } + + let replacement_len = replacement.tokens.len(); + let is_embed = { + let mask = replacement + .tokens + .iter() + .map(|&token| token as u32 == lane.embed_token_id) + .collect::>(); + WireTensor::from_bool(vec![replacement_len], mask).map_err(Error::Multimodal)? + }; + + let expanded_offset = expanded.len(); + expanded.extend(replacement.tokens.iter().map(|&token| token as u32)); + ranges.entry(lane.modality).or_default().push(PlaceholderRange { + offset: expanded_offset, + length: replacement_len, + is_embed: Some(is_embed), + }); + } + + for lane in &lanes { + if !lane.replacements.is_empty() { + bail_multimodal!( + "placeholder token `{}` was not found in tokenized prompt for {} remaining `{}` item(s)", + lane.placeholder_token, + lane.replacements.len(), + lane.modality + ); + } + } + + *prompt_token_ids = expanded; + + Ok(ranges) +} + +#[cfg(test)] +mod tests { + use llm_multimodal::TokenId; + use vllm_engine_core_client::protocol::tensor::WireArrayData; + + use super::super::tests::{ + LLAMA4_IMAGE_END_ID, LLAMA4_IMAGE_ID, LLAMA4_IMAGE_START_ID, LLAMA4_PATCH_ID, + LLAMA4_TILE_X_SEPARATOR_ID, LLAMA4_TILE_Y_SEPARATOR_ID, QWEN3_IMAGE_PAD_ID, + QWEN3_VIDEO_PAD_ID, + }; + use super::super::{PreparedMedia, ResolvedPlaceholder}; + use super::*; + + /// Build prepared media directly from placeholder token IDs. + fn prepared_media( + modality: Modality, + placeholder_token: &str, + marker_token_id: u32, + embed_token_id: u32, + replacements: Vec, + ) -> PreparedMedia { + PreparedMedia { + modality, + placeholder: ResolvedPlaceholder { + token: placeholder_token.to_string(), + marker_token_id, + embed_token_id, + }, + replacements, + items: Vec::new(), + } + } + + /// Llama4 image prepared media: the `<|image|>` marker expands to + /// sequences whose embed positions are the `<|patch|>` tokens. + fn llama4_prepared(replacements: Vec) -> PreparedMedia { + prepared_media( + Modality::Image, + "<|image|>", + LLAMA4_IMAGE_ID, + LLAMA4_PATCH_ID, + replacements, + ) + } + + fn qwen3_image_prepared(replacements: Vec) -> PreparedMedia { + prepared_media( + Modality::Image, + "<|image_pad|>", + QWEN3_IMAGE_PAD_ID, + QWEN3_IMAGE_PAD_ID, + replacements, + ) + } + + fn qwen3_video_prepared(replacements: Vec) -> PreparedMedia { + prepared_media( + Modality::Video, + "<|video_pad|>", + QWEN3_VIDEO_PAD_ID, + QWEN3_VIDEO_PAD_ID, + replacements, + ) + } + + fn llama4_single_tile_replacement() -> PromptReplacement { + PromptReplacement::sequence( + Modality::Image, + "<|image|>", + vec![ + LLAMA4_IMAGE_START_ID as TokenId, + LLAMA4_IMAGE_ID as TokenId, + LLAMA4_PATCH_ID as TokenId, + LLAMA4_PATCH_ID as TokenId, + LLAMA4_IMAGE_END_ID as TokenId, + ], + ) + } + + fn llama4_multi_tile_replacement() -> PromptReplacement { + PromptReplacement::sequence( + Modality::Image, + "<|image|>", + vec![ + LLAMA4_IMAGE_START_ID as TokenId, + LLAMA4_PATCH_ID as TokenId, + LLAMA4_TILE_X_SEPARATOR_ID as TokenId, + LLAMA4_PATCH_ID as TokenId, + LLAMA4_TILE_Y_SEPARATOR_ID as TokenId, + LLAMA4_IMAGE_ID as TokenId, + LLAMA4_PATCH_ID as TokenId, + LLAMA4_IMAGE_END_ID as TokenId, + ], + ) + } + + fn assert_bool_mask(range: &PlaceholderRange, expected: &[bool]) { + let tensor = range.is_embed.as_ref().expect("is_embed mask"); + assert_eq!(tensor.dtype, "bool"); + assert_eq!(tensor.shape, vec![expected.len()]); + assert_eq!( + tensor.data, + WireArrayData::RawView(expected.iter().map(|value| u8::from(*value)).collect()) + ); + } + + #[test] + fn expand_prompt_tokens_marks_only_llama4_patch_tokens_as_embed() { + let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; + let prepared = vec![llama4_prepared(vec![llama4_multi_tile_replacement()])]; + + let ranges = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap(); + let ranges = &ranges[&Modality::Image]; + + assert_eq!( + prompt_token_ids, + vec![ + 1, + LLAMA4_IMAGE_START_ID, + LLAMA4_PATCH_ID, + LLAMA4_TILE_X_SEPARATOR_ID, + LLAMA4_PATCH_ID, + LLAMA4_TILE_Y_SEPARATOR_ID, + LLAMA4_IMAGE_ID, + LLAMA4_PATCH_ID, + LLAMA4_IMAGE_END_ID, + 2, + ] + ); + assert_eq!(ranges[0].offset, 1); + assert_eq!(ranges[0].length, 8); + assert_bool_mask( + &ranges[0], + &[false, true, false, true, false, false, true, false], + ); + } + + #[test] + fn expand_prompt_tokens_errors_when_placeholder_missing() { + let mut prompt_token_ids = vec![1, 2, 3]; + let prepared = vec![llama4_prepared(vec![llama4_single_tile_replacement()])]; + + let error = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap_err(); + + assert!(matches!(error, Error::Multimodal(message) if message.contains("not found"))); + } + + #[test] + fn expand_prompt_tokens_ignores_empty_replacements() { + let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; + let original_prompt_token_ids = prompt_token_ids.clone(); + let prepared = vec![llama4_prepared(Vec::new())]; + + let ranges = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap(); + + assert!(ranges.is_empty()); + assert_eq!(prompt_token_ids, original_prompt_token_ids); + } + + #[test] + fn expand_prompt_tokens_leaves_prompt_unchanged_when_later_placeholder_missing() { + let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; + let original_prompt_token_ids = prompt_token_ids.clone(); + let prepared = vec![llama4_prepared(vec![ + llama4_single_tile_replacement(), + llama4_single_tile_replacement(), + ])]; + + let error = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap_err(); + + assert!(matches!(error, Error::Multimodal(message) if message.contains("not found"))); + assert_eq!(prompt_token_ids, original_prompt_token_ids); + } + + #[test] + fn expand_prompt_tokens_errors_when_replacement_is_empty() { + let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2]; + let original_prompt_token_ids = prompt_token_ids.clone(); + let prepared = vec![llama4_prepared(vec![PromptReplacement::sequence( + Modality::Image, + "<|image|>", + Vec::new(), + )])]; + + let error = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap_err(); + + assert!( + matches!(error, Error::Multimodal(message) if message.contains("expanded to no tokens")) + ); + assert_eq!(prompt_token_ids, original_prompt_token_ids); + } + + #[test] + fn expand_prompt_tokens_skips_llama4_image_marker_inside_replacement() { + let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2, LLAMA4_IMAGE_ID, 3]; + let prepared = vec![llama4_prepared(vec![ + llama4_single_tile_replacement(), + llama4_single_tile_replacement(), + ])]; + + let ranges = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap(); + let ranges = &ranges[&Modality::Image]; + + assert_eq!( + prompt_token_ids, + vec![ + 1, + LLAMA4_IMAGE_START_ID, + LLAMA4_IMAGE_ID, + LLAMA4_PATCH_ID, + LLAMA4_PATCH_ID, + LLAMA4_IMAGE_END_ID, + 2, + LLAMA4_IMAGE_START_ID, + LLAMA4_IMAGE_ID, + LLAMA4_PATCH_ID, + LLAMA4_PATCH_ID, + LLAMA4_IMAGE_END_ID, + 3, + ] + ); + assert_eq!(ranges[0].offset, 1); + assert_eq!(ranges[0].length, 5); + assert_bool_mask(&ranges[0], &[false, false, true, true, false]); + assert_eq!(ranges[1].offset, 7); + assert_eq!(ranges[1].length, 5); + assert_bool_mask(&ranges[1], &[false, false, true, true, false]); + } + + #[test] + fn expand_prompt_tokens_interleaves_image_and_video_prepared_media() { + let mut prompt_token_ids = vec![ + 1, + QWEN3_IMAGE_PAD_ID, + 2, + QWEN3_VIDEO_PAD_ID, + 3, + QWEN3_IMAGE_PAD_ID, + 4, + ]; + let prepared = vec![ + qwen3_image_prepared(vec![ + PromptReplacement::repeated( + Modality::Image, + "<|image_pad|>", + QWEN3_IMAGE_PAD_ID as TokenId, + 2, + ), + PromptReplacement::repeated( + Modality::Image, + "<|image_pad|>", + QWEN3_IMAGE_PAD_ID as TokenId, + 3, + ), + ]), + qwen3_video_prepared(vec![PromptReplacement::repeated( + Modality::Video, + "<|video_pad|>", + QWEN3_VIDEO_PAD_ID as TokenId, + 4, + )]), + ]; + + let ranges = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap(); + + assert_eq!( + prompt_token_ids, + vec![ + 1, + QWEN3_IMAGE_PAD_ID, + QWEN3_IMAGE_PAD_ID, + 2, + QWEN3_VIDEO_PAD_ID, + QWEN3_VIDEO_PAD_ID, + QWEN3_VIDEO_PAD_ID, + QWEN3_VIDEO_PAD_ID, + 3, + QWEN3_IMAGE_PAD_ID, + QWEN3_IMAGE_PAD_ID, + QWEN3_IMAGE_PAD_ID, + 4, + ] + ); + + let image_ranges = &ranges[&Modality::Image]; + assert_eq!(image_ranges[0].offset, 1); + assert_eq!(image_ranges[0].length, 2); + assert_bool_mask(&image_ranges[0], &[true, true]); + assert_eq!(image_ranges[1].offset, 9); + assert_eq!(image_ranges[1].length, 3); + assert_bool_mask(&image_ranges[1], &[true, true, true]); + + let video_ranges = &ranges[&Modality::Video]; + assert_eq!(video_ranges[0].offset, 4); + assert_eq!(video_ranges[0].length, 4); + assert_bool_mask(&video_ranges[0], &[true, true, true, true]); + } + + #[test] + fn expand_prompt_tokens_error_names_modality_with_leftover_replacements() { + let mut prompt_token_ids = vec![1, QWEN3_IMAGE_PAD_ID, 2]; + let original_prompt_token_ids = prompt_token_ids.clone(); + let prepared = vec![ + qwen3_image_prepared(vec![PromptReplacement::repeated( + Modality::Image, + "<|image_pad|>", + QWEN3_IMAGE_PAD_ID as TokenId, + 2, + )]), + qwen3_video_prepared(vec![PromptReplacement::repeated( + Modality::Video, + "<|video_pad|>", + QWEN3_VIDEO_PAD_ID as TokenId, + 4, + )]), + ]; + + let error = expand_prompt_token_ids(&mut prompt_token_ids, &prepared).unwrap_err(); + + assert!(matches!( + error, + Error::Multimodal(message) + if message.contains("<|video_pad|>") && message.contains("`video`") + )); + assert_eq!(prompt_token_ids, original_prompt_token_ids); + } +} diff --git a/rust/src/chat/src/multimodal/image.rs b/rust/src/chat/src/multimodal/image.rs new file mode 100644 index 000000000000..f71f1ee35261 --- /dev/null +++ b/rust/src/chat/src/multimodal/image.rs @@ -0,0 +1,141 @@ +//! Image-modality preparation: batch preprocessing and per-item feature +//! build. + +use std::sync::Arc; + +use itertools::izip; +use llm_multimodal::{FieldLayout, ImageFrame, Modality, PreprocessedEncoderInputs}; +use vllm_engine_core_client::protocol::dtype::ModelDtype; +use vllm_engine_core_client::protocol::multimodal::{ + MmBatchedField, MmField, MmFieldElem, MmFlatField, MmKwargsItem, MmSharedField, MmSlice, + SliceSpec, +}; + +use super::{ModalitySupport, MultimodalModelInfo, PreparedItem, PreparedMedia, tensor}; +use crate::error::{Error, Result, bail_multimodal, multimodal}; + +impl MultimodalModelInfo { + /// Preprocess all fetched image frames as one batch and build per-item + /// features. + pub(super) async fn prepare_images( + &self, + frames: Vec>, + uuids: Vec>, + model_dtype: ModelDtype, + ) -> Result { + let support = self.image.as_ref().ok_or_else(|| Error::UnsupportedModality { + modality: Modality::Image.to_string(), + })?; + let preprocessed = self.preprocess_images(support, &frames).await?; + let replacements = + self.spec + .prompt_replacements_for(&self.context, &preprocessed, Modality::Image)?; + if replacements.len() != frames.len() { + bail_multimodal!( + "number of image prompt replacements {} does not match number of images {}", + replacements.len(), + frames.len() + ); + } + let items = self.build_image_items(preprocessed, &frames, uuids, model_dtype)?; + + Ok(PreparedMedia { + modality: Modality::Image, + placeholder: support.placeholder.clone(), + replacements, + items, + }) + } + + /// Preprocess fetched image frames with the model's resolved vision + /// processor. + /// + /// The processor work is CPU-heavy relative to request wiring, so it runs + /// in a blocking task and returns owned tensors ready for wire + /// conversion. + async fn preprocess_images( + &self, + support: &ModalitySupport, + image_frames: &[Arc], + ) -> Result { + let config = support.config.clone(); + let processor = support.processor; + let images = image_frames.iter().map(|frame| frame.data().clone()).collect::>(); + + // TODO: is it still necessary given that we've already in a dedicated runtime? + tokio::task::spawn_blocking(move || Ok(processor.preprocess(&images, &config)?)) + .await + .map_err(|error| multimodal!("image preprocessing task failed: {error}"))? + } + + /// Convert one batch of preprocessed image tensors into per-item engine + /// kwargs. + /// + /// Tensor fields are sliced per item according to the model spec's field + /// layout declarations. + fn build_image_items( + &self, + preprocessed: PreprocessedEncoderInputs, + frames: &[Arc], + uuids: Vec>, + model_dtype: ModelDtype, + ) -> Result> { + let len = frames.len(); + let tensors = tensor::collect_tensors(preprocessed, "pixel_values", model_dtype)?; + + let mut items = Vec::with_capacity(len); + for (index, (frame, uuid)) in izip!(frames, uuids).enumerate() { + let mut data = MmKwargsItem::new(); + for (key, tensor) in &tensors { + let keep_on_cpu = self.spec.keep_on_cpu_keys.contains(key); + let (value, field) = match self.spec.field_layouts.get(key) { + Some(FieldLayout::Batched) => ( + tensor.batched_value_at(index)?, + MmField::Batched(MmBatchedField { keep_on_cpu }), + ), + Some(FieldLayout::Flat { sizes_key }) => { + let sizes = tensors.get(sizes_key).ok_or_else(|| { + multimodal!("flat tensor sizes key `{sizes_key}` is missing") + })?; + let (start, end) = tensor::flat_range_for_index(sizes, sizes_key, index)?; + ( + tensor.flat_value_range(start, end)?, + MmField::Flat(MmFlatField { + slices: vec![MmSlice::Slice(SliceSpec { + start: Some(0), + stop: Some((end - start) as isize), + step: None, + })], + dim: 0, + keep_on_cpu, + }), + ) + } + None => ( + tensor.clone(), + MmField::Shared(MmSharedField { + batch_size: len, + keep_on_cpu, + }), + ), + }; + + data.insert( + key.clone(), + MmFieldElem { + data: Some(value.try_into()?), + field, + }, + ); + } + + items.push(PreparedItem { + data, + hash: frame.hash.clone(), + uuid, + }); + } + + Ok(items) + } +} diff --git a/rust/src/chat/src/multimodal/tensor.rs b/rust/src/chat/src/multimodal/tensor.rs index 95259f1a93fe..e26022aea038 100644 --- a/rust/src/chat/src/multimodal/tensor.rs +++ b/rust/src/chat/src/multimodal/tensor.rs @@ -1,7 +1,7 @@ use std::collections::HashMap; use half::{bf16, f16}; -use llm_multimodal::{ModelSpecificValue, PreprocessedEncoderInputs as PreprocessedImages}; +use llm_multimodal::{ModelSpecificValue, PreprocessedEncoderInputs}; use vllm_engine_core_client::protocol::dtype::ModelDtype; use vllm_engine_core_client::protocol::multimodal::MmKwargValue as ProtocolKwargValue; use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireTensor}; @@ -25,25 +25,31 @@ pub(super) enum KwargValue { Passthrough(ProtocolKwargValue), } -/// Collect `pixel_values` and model-specific outputs into one tensor map. +/// Collect the primary encoder input and model-specific outputs into one +/// tensor map. +/// +/// `primary_key` names the encoder-input tensor as the model's forward kwargs +/// expect it (e.g. `pixel_values` for images, `pixel_values_videos` for +/// videos). pub(super) fn collect_tensors( - preprocessed: PreprocessedImages, + preprocessed: PreprocessedEncoderInputs, + primary_key: &str, float_dtype: ModelDtype, ) -> Result> { - let PreprocessedImages { + let PreprocessedEncoderInputs { encoder_input, model_specific, .. } = preprocessed; - let pixel_values = { + let primary_value = { let shape = encoder_input.shape().to_vec(); let data = encoder_input.into_iter().collect(); KwargValue::from_f32_tensor(data, shape, float_dtype)? }; let mut tensors = HashMap::new(); - tensors.insert("pixel_values".to_string(), pixel_values); + tensors.insert(primary_key.to_string(), primary_value); for (key, value) in model_specific { tensors.insert(key, KwargValue::from_model_specific(value, float_dtype)?); } @@ -124,10 +130,22 @@ impl TryFrom for ProtocolKwargValue { } impl KwargValue { - /// Extract one image from a batched tensor field. + /// First-axis length for tensor values; `None` for passthrough kwargs. + pub(super) fn first_dim(&self) -> Option { + match self { + Self::F32Tensor { shape, .. } + | Self::F16Tensor { shape, .. } + | Self::Bf16Tensor { shape, .. } + | Self::I64Tensor { shape, .. } + | Self::U32Tensor { shape, .. } => shape.first().copied(), + Self::Passthrough(_) => None, + } + } + + /// Extract one media item from a batched tensor field. /// - /// Batched fields use their first axis as image index and drop that axis in - /// the per-feature value, matching vLLM's batched-field semantics. + /// Batched fields use their first axis as media-item index and drop that + /// axis in the per-feature value, matching vLLM's batched-field semantics. pub(super) fn batched_value_at(&self, index: usize) -> Result { match self { Self::F32Tensor { data, shape } => { @@ -154,9 +172,9 @@ impl KwargValue { } } - /// Extract one image's variable-length range from a flat tensor field. + /// Extract one media item's variable-length range from a flat tensor field. /// - /// Flat fields keep the first axis as the sliced length for this image. + /// Flat fields keep the first axis as the sliced length for this item. pub(super) fn flat_value_range(&self, start: usize, end: usize) -> Result { match self { Self::F32Tensor { data, shape } => { @@ -184,10 +202,10 @@ impl KwargValue { } } -/// Compute the first-axis range for one image in a flat tensor. +/// Compute the first-axis range for one media item in a flat tensor. /// /// `sizes_key` names a companion tensor whose entries are cumulative slice -/// sizes per image. +/// sizes per media item. pub(super) fn flat_range_for_index( sizes: &KwargValue, sizes_key: &str, @@ -195,7 +213,7 @@ pub(super) fn flat_range_for_index( ) -> Result<(usize, usize)> { let sizes = tensor_as_usize_vec(sizes)?; let size = *sizes.get(index).ok_or_else(|| { - multimodal!("flat tensor sizes key `{sizes_key}` has no entry for image {index}") + multimodal!("flat tensor sizes key `{sizes_key}` has no entry for media item {index}") })?; let start = sizes[..index].iter().sum::(); Ok((start, start + size)) diff --git a/rust/src/chat/src/multimodal/video.rs b/rust/src/chat/src/multimodal/video.rs new file mode 100644 index 000000000000..482074abde65 --- /dev/null +++ b/rust/src/chat/src/multimodal/video.rs @@ -0,0 +1,316 @@ +//! Video-modality preparation: per-clip preprocessing, config resolution, +//! and per-item feature build. + +use std::sync::Arc; + +use itertools::izip; +use llm_multimodal::{FieldLayout, Modality, PreprocessedEncoderInputs, VideoClip}; +use thiserror_ext::AsReport as _; +use tracing::warn; +use vllm_engine_core_client::protocol::dtype::ModelDtype; +use vllm_engine_core_client::protocol::multimodal::{ + MmBatchedField, MmField, MmFieldElem, MmFlatField, MmKwargsItem, MmSharedField, MmSlice, + SliceSpec, +}; + +use super::{ModalitySupport, MultimodalModelInfo, PreparedItem, PreparedMedia, tensor}; +use crate::error::{Error, Result, bail_multimodal, multimodal}; + +/// Forward-kwargs name of the primary video encoder input. +/// +/// Video-capable vLLM models read `pixel_values_videos` alongside +/// `video_grid_thw`, mirroring the HF processor output naming. +const VIDEO_PRIMARY_KEY: &str = "pixel_values_videos"; + +impl MultimodalModelInfo { + /// Preprocess fetched video clips one at a time and build per-item + /// features. + /// + /// Unlike images, each clip runs through the preprocessor independently + /// (a batch of one), so its tensors are complete per item and need no + /// cross-item slicing. + pub(super) async fn prepare_videos( + &self, + clips: Vec>, + uuids: Vec>, + model_dtype: ModelDtype, + ) -> Result { + let support = self.video.as_ref().ok_or_else(|| Error::UnsupportedModality { + modality: Modality::Video.to_string(), + })?; + let mut replacements = Vec::with_capacity(clips.len()); + let mut items = Vec::with_capacity(clips.len()); + + for (clip, uuid) in izip!(&clips, uuids) { + let preprocessed = self.preprocess_video_clip(support, Arc::clone(clip)).await?; + let mut clip_replacements = + self.spec + .prompt_replacements_for(&self.context, &preprocessed, Modality::Video)?; + if clip_replacements.len() != 1 { + bail_multimodal!( + "expected exactly one prompt replacement per video clip, got {}", + clip_replacements.len() + ); + } + replacements.push(clip_replacements.pop().unwrap()); + items.push(self.build_video_item( + preprocessed, + clip.hash.clone(), + uuid, + model_dtype, + )?); + } + + Ok(PreparedMedia { + modality: Modality::Video, + placeholder: support.placeholder.clone(), + replacements, + items, + }) + } + + /// Preprocess one decoded video clip with the model's resolved vision + /// processor. + async fn preprocess_video_clip( + &self, + support: &ModalitySupport, + clip: Arc, + ) -> Result { + let config = support.config.clone(); + let processor = support.processor; + + tokio::task::spawn_blocking(move || { + // Prefer the borrowed-RGB fast path, which avoids materializing a + // `DynamicImage` per sampled frame after media decode. + if let Some(rgb_video) = clip.rgb_video() { + match rgb_video.frame_refs() { + Ok(frame_refs) => match processor.preprocess_video_rgb(&frame_refs, &config) { + Ok(preprocessed) => return Ok(preprocessed), + Err(error) => warn!( + error = %error.as_report(), + "RGB video preprocessing fast path failed; falling back to materialized frames" + ), + }, + Err(error) => warn!( + error, + "RGB video frame refs are invalid; falling back to materialized frames" + ), + } + } + + let frames = clip.materialized_frames().map_err(|error| multimodal!("{error}"))?; + Ok(processor.preprocess_video(&frames, &config)?) + }) + .await + .map_err(|error| multimodal!("video preprocessing task failed: {error}"))? + } + + /// Convert one preprocessed video clip into engine kwargs. + /// + /// The clip is a batch of one, so no per-item slicing is required: the + /// primary tensor ships as a full-range flat field (the engine re-batches + /// flat fields by concatenating along the declared dim, matching vLLM's + /// `flat_from_sizes` treatment of video patches), and batched metadata + /// tensors drop their singleton batch axis. + fn build_video_item( + &self, + preprocessed: PreprocessedEncoderInputs, + hash: String, + uuid: Option, + model_dtype: ModelDtype, + ) -> Result { + let tensors = tensor::collect_tensors(preprocessed, VIDEO_PRIMARY_KEY, model_dtype)?; + + let mut data = MmKwargsItem::new(); + for (key, tensor) in tensors { + let keep_on_cpu = self.spec.keep_on_cpu_keys.contains(&key); + let (value, field) = if key == VIDEO_PRIMARY_KEY { + let len = tensor + .first_dim() + .ok_or_else(|| multimodal!("video encoder input `{key}` is not a tensor"))?; + ( + tensor, + MmField::Flat(MmFlatField { + slices: vec![MmSlice::Slice(SliceSpec { + start: Some(0), + stop: Some(len as isize), + step: None, + })], + dim: 0, + keep_on_cpu, + }), + ) + } else if matches!( + self.spec.field_layouts.get(&key), + Some(FieldLayout::Batched) + ) { + ( + tensor.batched_value_at(0)?, + MmField::Batched(MmBatchedField { keep_on_cpu }), + ) + } else { + ( + tensor, + MmField::Shared(MmSharedField { + batch_size: 1, + keep_on_cpu, + }), + ) + }; + + data.insert( + key, + MmFieldElem { + data: Some(value.try_into()?), + field, + }, + ); + } + + Ok(PreparedItem { data, hash, uuid }) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + + use llm_multimodal::ModelSpecificValue; + use ndarray::ArrayD; + use vllm_engine_core_client::protocol::multimodal::MmKwargValue; + + use super::super::tests::{ + QWEN3_IMAGE_PAD_ID, QWEN3_VIDEO_PAD_ID, qwen3_vl_info, qwen3_vl_tokenizer, + }; + use super::super::{MultimodalConfigFiles, MultimodalModelInfo}; + use super::*; + + #[test] + fn from_paths_resolves_video_config_from_dedicated_file_or_processor_config() { + let dir = tempfile::tempdir().unwrap(); + let config_path = dir.path().join("config.json"); + std::fs::write( + &config_path, + serde_json::json!({ + "model_type": "qwen3_vl", + "image_token_id": QWEN3_IMAGE_PAD_ID, + "video_token_id": QWEN3_VIDEO_PAD_ID, + }) + .to_string(), + ) + .unwrap(); + + let info_for = |files: MultimodalConfigFiles<'_>| { + MultimodalModelInfo::from_paths( + "qwen3-vl-test".to_string(), + Some("qwen3_vl".to_string()), + files, + Arc::new(qwen3_vl_tokenizer()), + ) + }; + + // Dedicated video preprocessor config file. + let video_config_path = dir.path().join("video_preprocessor_config.json"); + std::fs::write(&video_config_path, r#"{"size":{"shortest_edge":128}}"#).unwrap(); + let info = info_for(MultimodalConfigFiles { + config: Some(&config_path), + video_preprocessor_config: Some(&video_config_path), + ..Default::default() + }) + .unwrap() + .unwrap(); + assert!(info.video.is_some()); + + // `video_processor` section of the combined processor config. + let processor_config_path = dir.path().join("processor_config.json"); + std::fs::write( + &processor_config_path, + r#"{"video_processor":{"size":{"shortest_edge":128}}}"#, + ) + .unwrap(); + let info = info_for(MultimodalConfigFiles { + config: Some(&config_path), + processor_config: Some(&processor_config_path), + ..Default::default() + }) + .unwrap() + .unwrap(); + assert!(info.video.is_some()); + + // Neither source: video support still resolves on the image config. + let info = info_for(MultimodalConfigFiles { + config: Some(&config_path), + ..Default::default() + }) + .unwrap() + .unwrap(); + assert!(info.video.is_some()); + + // Malformed dedicated file is a real error, not a silent fallback. + std::fs::write(&video_config_path, r#"{"size""#).unwrap(); + let error = match info_for(MultimodalConfigFiles { + config: Some(&config_path), + video_preprocessor_config: Some(&video_config_path), + ..Default::default() + }) { + Err(error) => error, + Ok(_) => panic!("malformed video preprocessor config should fail"), + }; + assert!(matches!( + error, + Error::Multimodal(message) + if message.contains("failed to parse video_preprocessor_config.json") + )); + } + + #[test] + fn build_video_item_names_primary_tensor_and_layouts() { + let info = qwen3_vl_info(); + // One clip flattened to 6 patches with 4 features each. + let preprocessed = PreprocessedEncoderInputs { + encoder_input: ArrayD::zeros(vec![6, 4]), + feature_token_counts: vec![6], + item_sizes: vec![(32, 32)], + model_specific: HashMap::from([ + ( + "video_grid_thw".to_string(), + ModelSpecificValue::int_2d(vec![1, 2, 3], 1, 3), + ), + ( + "patches_per_video".to_string(), + ModelSpecificValue::int_1d(vec![6]), + ), + ]), + }; + + let item = info + .build_video_item( + preprocessed, + "".to_string(), + None, + ModelDtype::Float32, + ) + .unwrap(); + + let primary = &item.data[VIDEO_PRIMARY_KEY]; + assert!(matches!( + &primary.field, + MmField::Flat(MmFlatField { slices, dim: 0, .. }) + if matches!( + slices.as_slice(), + [MmSlice::Slice(SliceSpec { start: Some(0), stop: Some(6), step: None })] + ) + )); + + // Batched metadata drops its singleton batch axis per item. + let grid = &item.data["video_grid_thw"]; + assert!(matches!(&grid.field, MmField::Batched(_))); + let MmKwargValue::Tensor(grid_tensor) = grid.data.as_ref().unwrap() else { + panic!("expected tensor value for video_grid_thw"); + }; + assert_eq!(grid_tensor.shape, vec![3]); + + assert_eq!(item.hash, ""); + } +} diff --git a/rust/src/chat/src/renderer/hf/mod.rs b/rust/src/chat/src/renderer/hf/mod.rs index d9031d73a4bf..1c87e75337f7 100644 --- a/rust/src/chat/src/renderer/hf/mod.rs +++ b/rust/src/chat/src/renderer/hf/mod.rs @@ -31,9 +31,14 @@ pub use template::{load_chat_template, resolve_chat_template}; pub use self::format::ChatTemplateContentFormatOption; -#[derive(Debug, Clone)] +/// Template-visible placeholder tokens per supported modality. +/// +/// A `None` token means the loaded model does not support that modality, and +/// content parts of that modality are rejected during rendering. +#[derive(Debug, Clone, Default)] pub struct MultimodalRenderInfo { - pub placeholder_token: String, + pub image_token: Option, + pub video_token: Option, } /// Hugging Face chat-template renderer backed by the local Jinja chat-template @@ -254,6 +259,7 @@ enum TemplateContent { enum TemplateContentPart { Text { text: String }, Image, + Video, } #[derive(Debug, Serialize)] @@ -417,9 +423,17 @@ fn to_template_openai_content( } // All multimodal contents are normalized to `{ "type": }`. ChatContentPart::ImageUrl { .. } => { - multimodal.ok_or(Error::UnsupportedMultimodalContent("image_url"))?; + multimodal + .and_then(|multimodal| multimodal.image_token.as_ref()) + .ok_or(Error::UnsupportedMultimodalContent("image_url"))?; Ok(TemplateContentPart::Image) } + ChatContentPart::VideoUrl { .. } => { + multimodal + .and_then(|multimodal| multimodal.video_token.as_ref()) + .ok_or(Error::UnsupportedMultimodalContent("video_url"))?; + Ok(TemplateContentPart::Video) + } }) .collect(), } @@ -437,9 +451,16 @@ fn to_template_string_content( match part { ChatContentPart::Text { text } => out.push_str(text), ChatContentPart::ImageUrl { .. } => { - let multimodal = - multimodal.ok_or(Error::UnsupportedMultimodalContent("image_url"))?; - out.push_str(&multimodal.placeholder_token); + let image_token = multimodal + .and_then(|multimodal| multimodal.image_token.as_ref()) + .ok_or(Error::UnsupportedMultimodalContent("image_url"))?; + out.push_str(image_token); + } + ChatContentPart::VideoUrl { .. } => { + let video_token = multimodal + .and_then(|multimodal| multimodal.video_token.as_ref()) + .ok_or(Error::UnsupportedMultimodalContent("video_url"))?; + out.push_str(video_token); } } } @@ -468,7 +489,7 @@ fn append_continue_final_message_tag(message: &mut TemplateMessage) -> Result parts.iter_mut().rev().find_map(|part| match part { TemplateContentPart::Text { text } => Some(text), - TemplateContentPart::Image => None, + TemplateContentPart::Image | TemplateContentPart::Video => None, }), }; let text = text.ok_or_else(|| { @@ -577,7 +598,8 @@ mod tests { ) -> Result { HfChatRenderer::new(Some(template.to_string()), HashMap::new(), content_format)? .with_multimodal(Some(MultimodalRenderInfo { - placeholder_token: "".to_string(), + image_token: Some("".to_string()), + video_token: Some("