From 21349939a17a3073b6b4ba48f87e002596236713 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Sun, 12 Jul 2026 07:19:51 +0000 Subject: [PATCH 01/49] [sync] [ROCm] Revert Part of `[ROCm] Fix pooling startup workspace lock` #47912 (#48154) Upstream-vLLM: 766469a4c460043ae52cda19b1c52f0dc87e555c Co-authored-by: Micah Williamson --- .sync/vllm-sha | 2 +- aphrodite/v1/worker/gpu_worker.py | 22 +++------------------- 2 files changed, 4 insertions(+), 20 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 8828e4021a..c86fd21b19 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -ea0fa34f4992c09fc5aa4ae9a12e67e03125289c +766469a4c460043ae52cda19b1c52f0dc87e555c diff --git a/aphrodite/v1/worker/gpu_worker.py b/aphrodite/v1/worker/gpu_worker.py index ec0033aca5..3d65cd17cc 100644 --- a/aphrodite/v1/worker/gpu_worker.py +++ b/aphrodite/v1/worker/gpu_worker.py @@ -686,11 +686,6 @@ def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None: @instrument(span_name="Warmup (GPU)") def compile_or_warm_up_model(self) -> CompilationTimes: warmup_sizes: list[int] = [] - cg_capture_sizes: list[int] = [] - - if self.aphrodite_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE: - cg_sizes = self.aphrodite_config.compilation_config.cudagraph_capture_sizes - cg_capture_sizes = [] if cg_sizes is None else cg_sizes if self.aphrodite_config.compilation_config.mode == CompilationMode.APHRODITE_COMPILE: # warm up sizes that are not in cudagraph capture sizes, @@ -698,8 +693,11 @@ def compile_or_warm_up_model(self) -> CompilationTimes: # e.g. for the max-num-batched token size in chunked prefill. compile_sizes = self.aphrodite_config.compilation_config.compile_sizes warmup_sizes = compile_sizes.copy() if compile_sizes is not None else [] # type: ignore[assignment] + cg_capture_sizes: list[int] = [] if self.aphrodite_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE: + cg_sizes = self.aphrodite_config.compilation_config.cudagraph_capture_sizes + cg_capture_sizes = [] if cg_sizes is None else cg_sizes warmup_sizes = [x for x in warmup_sizes if x not in cg_capture_sizes] compile_ranges = self.aphrodite_config.compilation_config.get_compile_ranges() @@ -712,20 +710,6 @@ def compile_or_warm_up_model(self) -> CompilationTimes: if not any(x in compile_range for x in all_sizes): warmup_sizes.append(compile_range.end) - # TODO(LucasWilkinson, akaratza): Remove when MRV1 is deprecated - if ( - current_platform.is_rocm() - and not self.use_v2_model_runner - and self.aphrodite_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE - and get_pp_group().is_last_rank - ): - max_num_reqs = min( - self.scheduler_config.max_num_seqs, - self.scheduler_config.max_num_batched_tokens, - ) - if max_num_reqs not in cg_capture_sizes and max_num_reqs not in warmup_sizes: - warmup_sizes.append(max_num_reqs) - # We skip EPLB here since we don't want to record dummy metrics for size in sorted(warmup_sizes, reverse=True): logger.info("Compile and warming up model for size %d", size) From 4805c2c8fbb105b9a100769f307ddae8afb3389b Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Sun, 12 Jul 2026 07:22:03 +0000 Subject: [PATCH 02/49] [sync] Correct model layer aliasing for Bert style models (#43896) Upstream-vLLM: cac3e70cd4f522eaab9d28323285d5e31238f907 Co-authored-by: ap9272 --- .sync/vllm-sha | 2 +- aphrodite/model_executor/models/modernbert.py | 10 +++++++++- aphrodite/model_executor/warmup/deep_gemm_warmup.py | 5 ++--- 3 files changed, 12 insertions(+), 5 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index c86fd21b19..536aca7e03 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -766469a4c460043ae52cda19b1c52f0dc87e555c +cac3e70cd4f522eaab9d28323285d5e31238f907 diff --git a/aphrodite/model_executor/models/modernbert.py b/aphrodite/model_executor/models/modernbert.py index 3de1b685d9..69ce41f598 100644 --- a/aphrodite/model_executor/models/modernbert.py +++ b/aphrodite/model_executor/models/modernbert.py @@ -220,7 +220,13 @@ def forward( @support_torch_compile @default_pooling_type(seq_pooling_type="CLS") class ModernBertModel(nn.Module): - hf_to_aphrodite_mapper = WeightsMapper(orig_to_new_prefix={"layers.": "encoder_layer.layers."}) + hf_to_aphrodite_mapper = WeightsMapper( + orig_to_new_prefix={ + "model.layers.": "encoder_layer.layers.", + "layers.": "encoder_layer.layers.", + "model.": "", + } + ) def __init__( self, @@ -244,6 +250,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/aphrodite/model_executor/warmup/deep_gemm_warmup.py b/aphrodite/model_executor/warmup/deep_gemm_warmup.py index 8659e85d82..68bb2e2e36 100644 --- a/aphrodite/model_executor/warmup/deep_gemm_warmup.py +++ b/aphrodite/model_executor/warmup/deep_gemm_warmup.py @@ -128,9 +128,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) @@ -146,6 +143,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 9d8fee74f739747bd6e3f349354701ff1741e50d Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Sun, 12 Jul 2026 07:23:06 +0000 Subject: [PATCH 03/49] [sync] update marlin M size for EP (#48144) Upstream-vLLM: f1a5adddb815610db46cae81f2a1d5a4609bf99d Co-authored-by: gnovack --- .sync/vllm-sha | 2 +- .../layers/fused_moe/experts/marlin_moe.py | 9 +++++++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 536aca7e03..fd770268fb 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -cac3e70cd4f522eaab9d28323285d5e31238f907 +f1a5adddb815610db46cae81f2a1d5a4609bf99d diff --git a/aphrodite/model_executor/layers/fused_moe/experts/marlin_moe.py b/aphrodite/model_executor/layers/fused_moe/experts/marlin_moe.py index 9af65fa95d..91254167f4 100644 --- a/aphrodite/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/aphrodite/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 @@ -306,6 +307,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]: @@ -315,8 +322,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 da0f73a5354cb86a826e987f55c22f44aacbdc21 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Sun, 12 Jul 2026 07:25:44 +0000 Subject: [PATCH 04/49] [sync] [bugfix] bge-m3-sparse-plugin mismatch requests (#48112) Upstream-vLLM: feb384ada2a1da1981ccb713ea6bd3f90c44f7b4 Co-authored-by: Augusto Yao --- .sync/vllm-sha | 2 +- .../sparse_embeddings_processor.py | 45 ++----------------- 2 files changed, 5 insertions(+), 42 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index fd770268fb..a32afdc462 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -f1a5adddb815610db46cae81f2a1d5a4609bf99d +feb384ada2a1da1981ccb713ea6bd3f90c44f7b4 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 62ae56f1a4..61ebcf95c6 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 aphrodite.config import AphroditeConfig, ModelConfig, PoolerConfig +from aphrodite.config import AphroditeConfig, PoolerConfig from aphrodite.entrypoints.openai.engine.protocol import UsageInfo -from aphrodite.entrypoints.pooling.base.protocol import EmbedRequestMixin from aphrodite.inputs import PromptType from aphrodite.outputs import PoolingRequestOutput from aphrodite.plugins.io_processors.interface import IOProcessor @@ -14,7 +13,6 @@ from aphrodite.tokenizers.detokenizer_utils import convert_ids_list_to_tokens from .types import ( - EMBED_TASKS, SparseEmbeddingCompletionRequestMixin, SparseEmbeddingResponse, SparseEmbeddingResponseData, @@ -36,7 +34,6 @@ def __init__(self, aphrodite_config: AphroditeConfig, renderer: BaseRenderer): continue self.default_pooling_params[param] = getattr(pooler_config, param) self.embed_dimensions = aphrodite_config.model_config.embedding_size - self.embed_request_queue: list[EmbedRequestMixin] = [] def __repr__(self) -> str: return ( @@ -54,41 +51,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}, 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.aphrodite_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(self, request_data: object) -> SparseEmbeddingCompletionRequestMixin: @@ -106,10 +71,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): @@ -144,7 +107,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 1d71d3f74e90ae92fa01cfba4c2f0a13dfce8227 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Sun, 12 Jul 2026 07:27:17 +0000 Subject: [PATCH 05/49] [sync] [CI/Build][AMD] Fix ROCm OOM in eagle_correctness_heavy by reserving CUDA graph memory (#47366) Upstream-vLLM: 88e5e2c57be8ce6e25510c1249352a23b8a85ec4 Co-authored-by: peizhang56 --- .sync/vllm-sha | 2 +- aphrodite/v1/worker/gpu_worker.py | 12 +++++++----- tests/v1/e2e/spec_decode/test_spec_decode.py | 7 +++++++ 3 files changed, 15 insertions(+), 6 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index a32afdc462..d6677214a6 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -feb384ada2a1da1981ccb713ea6bd3f90c44f7b4 +88e5e2c57be8ce6e25510c1249352a23b8a85ec4 diff --git a/aphrodite/v1/worker/gpu_worker.py b/aphrodite/v1/worker/gpu_worker.py index 3d65cd17cc..48237a82ef 100644 --- a/aphrodite/v1/worker/gpu_worker.py +++ b/aphrodite/v1/worker/gpu_worker.py @@ -454,11 +454,14 @@ def determine_available_memory(self) -> int: profile_torch_peak = torch.accelerator.memory_stats(self.device).get("allocated_bytes.all.peak", 0) # 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.aphrodite_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE ): cudagraph_memory_estimate = self.model_runner.profile_cudagraph_memory() @@ -469,8 +472,7 @@ def determine_available_memory(self) -> int: profile_result.non_torch_increase + profile_result.torch_peak_increase + 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.APHRODITE_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS else 0 ) diff --git a/tests/v1/e2e/spec_decode/test_spec_decode.py b/tests/v1/e2e/spec_decode/test_spec_decode.py index 5417f3cf0e..3169ace7e2 100644 --- a/tests/v1/e2e/spec_decode/test_spec_decode.py +++ b/tests/v1/e2e/spec_decode/test_spec_decode.py @@ -26,6 +26,7 @@ multi_gpu_marks, multi_gpu_only, single_gpu_only, + wait_for_rocm_memory_to_settle, ) MTP_SIMILARITY_RATE = 0.8 @@ -418,6 +419,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, @@ -453,6 +457,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 From bd0646d0ff680318cb85da99ddac321f2f4cb907 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Sun, 12 Jul 2026 07:31:43 +0000 Subject: [PATCH 06/49] [sync] [kv_offload] Emit tier-owned BlockStored events from FS/OBJ secondary tiers (#47923) Upstream-vLLM: 2d814a00820daec7082599bea75ae1d0959a346c Co-authored-by: Chang Guo --- .sync/vllm-sha | 2 +- aphrodite/distributed/kv_events.py | 2 + aphrodite/v1/kv_offload/tiering/fs/manager.py | 48 ++++- .../v1/kv_offload/tiering/obj/manager.py | 53 ++++- tests/v1/kv_offload/tiering/test_fs_tier.py | 194 ++++++++++++++++++ tests/v1/kv_offload/tiering/test_obj_tier.py | 108 +++++++++- 6 files changed, 400 insertions(+), 7 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index d6677214a6..79dad54462 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -88e5e2c57be8ce6e25510c1249352a23b8a85ec4 +2d814a00820daec7082599bea75ae1d0959a346c diff --git a/aphrodite/distributed/kv_events.py b/aphrodite/distributed/kv_events.py index 1a287715cc..e454ce6dc3 100644 --- a/aphrodite/distributed/kv_events.py +++ b/aphrodite/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/aphrodite/v1/kv_offload/tiering/fs/manager.py b/aphrodite/v1/kv_offload/tiering/fs/manager.py index 594cdc04d4..a35129a2fa 100644 --- a/aphrodite/v1/kv_offload/tiering/fs/manager.py +++ b/aphrodite/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 aphrodite.fs_io_C import batch_lookup as batch_lookup_C @@ -30,11 +30,18 @@ from typing_extensions import override +from aphrodite.distributed.kv_events import MEDIUM_FS from aphrodite.logger import init_logger -from aphrodite.v1.kv_offload.base import LookupResult, OffloadKey, ReqContext +from aphrodite.v1.kv_offload.base import ( + LookupResult, + OffloadingEvent, + OffloadKey, + ReqContext, +) from aphrodite.v1.kv_offload.file_mapper import FileMapper from aphrodite.v1.kv_offload.tiering.async_lookup import AsyncLookupManager from aphrodite.v1.kv_offload.tiering.base import ( + JobId, JobMetadata, JobResult, RequestOffloadingContext, @@ -90,6 +97,8 @@ class FileSystemTierManager(SecondaryTierManager): content. """ + medium: ClassVar[str] = MEDIUM_FS + def __init__( self, offloading_spec: "OffloadingSpec", @@ -98,6 +107,7 @@ def __init__( root_dir: str, n_read_threads: int = 16, n_write_threads: int = 16, + enable_kv_events: bool = False, ): """ Args: @@ -108,9 +118,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" self._block_size: int = primary_kv_view.strides[0] @@ -151,6 +178,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, @@ -182,7 +211,20 @@ 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/aphrodite/v1/kv_offload/tiering/obj/manager.py b/aphrodite/v1/kv_offload/tiering/obj/manager.py index 4bbc44e49d..3d59cbbf07 100644 --- a/aphrodite/v1/kv_offload/tiering/obj/manager.py +++ b/aphrodite/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 aphrodite.distributed.kv_events import MEDIUM_OBJ from aphrodite.distributed.nixl_utils import NixlWrapper as nixl_agent from aphrodite.distributed.nixl_utils import nixl_agent_config from aphrodite.logger import init_logger -from aphrodite.v1.kv_offload.base import LookupResult, OffloadKey, ReqContext +from aphrodite.v1.kv_offload.base import ( + LookupResult, + OffloadingEvent, + OffloadKey, + ReqContext, +) from aphrodite.v1.kv_offload.file_mapper import FileMapper from aphrodite.v1.kv_offload.tiering.async_lookup import AsyncLookupManager from aphrodite.v1.kv_offload.tiering.base import ( + JobId, JobMetadata, JobResult, RequestOffloadingContext, @@ -88,6 +95,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", @@ -96,8 +105,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) @@ -217,6 +254,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) @@ -261,8 +300,18 @@ 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. diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index 38b38c52a7..3b9f6cc595 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 aphrodite.distributed.kv_events import MEDIUM_FS from aphrodite.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.aphrodite_config = _MOCK_APHRODITE_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) @@ -160,6 +172,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 # --------------------------------------------------------------------------- @@ -388,3 +417,168 @@ 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 aphrodite.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 aphrodite.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 aphrodite.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) + + 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 0c35103666..b60c9e8e9a 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -66,6 +66,15 @@ def _make_aphrodite_config(): _CTX = ReqContext(req_id="test-req") +def _make_events_spec(enable_kv_cache_events: bool) -> SimpleNamespace: + """Offloading spec stub with an explicit global KV events flag.""" + return SimpleNamespace( + aphrodite_config=_make_aphrodite_config(), + kv_cache_config=SimpleNamespace(kv_cache_groups=[]), + kv_events_config=SimpleNamespace(enable_kv_cache_events=enable_kv_cache_events), + ) + + def key(n: int) -> OffloadKey: return make_offload_key(n.to_bytes(8, "big"), 0) @@ -181,6 +190,8 @@ def _query_memory(self, queries, mem_type, agent_name): 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 @@ -417,6 +429,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( From e24724120a8f06cf7b9846343009242a486f1985 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Sun, 12 Jul 2026 07:35:02 +0000 Subject: [PATCH 07/49] [sync] DCP supports hybrid attention (#40996) Upstream-vLLM: 95ed0feaa5cd7fb16d72c53ce04950aaf07c4698 Co-authored-by: Yan Xu --- .sync/vllm-sha | 2 +- aphrodite/v1/worker/cp_utils.py | 6 +++ tests/distributed/test_context_parallel.py | 8 ++++ tests/distributed/test_pynccl.py | 39 +++++++++++++++++++ .../models/language/generation/test_hybrid.py | 16 +++++++- .../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 | 39 +++++++++++++++++++ tests/v1/worker/test_gpu_input_batch.py | 3 ++ 10 files changed, 117 insertions(+), 6 deletions(-) create mode 100644 tests/v1/worker/test_cp_utils.py diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 79dad54462..0a799371f6 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -2d814a00820daec7082599bea75ae1d0959a346c +95ed0feaa5cd7fb16d72c53ce04950aaf07c4698 diff --git a/aphrodite/v1/worker/cp_utils.py b/aphrodite/v1/worker/cp_utils.py index fec81dcb88..c30c8fe96a 100644 --- a/aphrodite/v1/worker/cp_utils.py +++ b/aphrodite/v1/worker/cp_utils.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import TYPE_CHECKING, Any, cast +import torch + from aphrodite.config import AphroditeConfig, get_layers_from_aphrodite_config from aphrodite.distributed import get_dcp_group, get_pcp_group @@ -55,3 +57,7 @@ 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 should_skip_dcp_context_attention(context_kv_lens: torch.Tensor) -> bool: + return bool(context_kv_lens.max().item() == 0) diff --git a/tests/distributed/test_context_parallel.py b/tests/distributed/test_context_parallel.py index ee1a957f2b..bb638d6908 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, } @@ -147,6 +149,12 @@ def iter_params(self, model_id: str): CPTestSettings.detailed(cp_kv_cache_interleave_size=16, attn_backend="FLASH_ATTN"), CPTestSettings.detailed(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 16500d6768..425b68279a 100644 --- a/tests/distributed/test_pynccl.py +++ b/tests/distributed/test_pynccl.py @@ -15,6 +15,7 @@ from aphrodite.distributed.device_communicators.pynccl_wrapper import NCCLLibrary from aphrodite.distributed.parallel_state import ( ensure_model_parallel_initialized, + get_tp_group, get_world_group, graph_capture, init_distributed_environment, @@ -178,6 +179,44 @@ 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_aphrodite_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(get_world_group().cpu_group, device=get_world_group().device) diff --git a/tests/models/language/generation/test_hybrid.py b/tests/models/language/generation/test_hybrid.py index 02d7b3fff3..f031d05051 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", @@ -90,7 +95,16 @@ def test_models( with hf_runner(model) as hf_model: hf_outputs = hf_model.generate_greedy_logprobs_limit(example_prompts, max_tokens, num_logprobs) - with aphrodite_runner(model, max_num_seqs=MAX_NUM_SEQS, attention_backend=ATTN_BACKEND) as aphrodite_model: + extra_kwargs = {} + if model in HYBRID_MODELS_REQUIRING_CHUNKED_PREFILL: + extra_kwargs["enable_chunked_prefill"] = True + + with aphrodite_runner( + model, + max_num_seqs=MAX_NUM_SEQS, + attention_backend=ATTN_BACKEND, + **extra_kwargs, + ) as aphrodite_model: aphrodite_outputs = aphrodite_model.generate_greedy_logprobs(example_prompts, max_tokens, num_logprobs) check_logprobs_close( diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index e11960eec7..49e11a6f3b 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/tests/models/multimodal/generation/test_vit_cudagraph.py @@ -144,6 +144,7 @@ def gemma3_chat_template(content: str) -> str: "<|vision_start|><|video_pad|><|vision_end|>Describe this video in one sentence." ), needs_video_metadata=True, + aphrodite_runner_kwargs={"enable_chunked_prefill": True}, marks=[pytest.mark.core_model], ), "qwen3_5": VitCudagraphTestConfig( diff --git a/tests/test_config.py b/tests/test_config.py index 2fbcb49c19..180857aaf1 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1082,14 +1082,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 0e06c43f8f..b8818c99ea 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 0000000000..21e87914ba --- /dev/null +++ b/tests/v1/worker/test_cp_utils.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest +import torch + +from aphrodite.v1.attention.backends.utils import get_dcp_local_seq_lens +from aphrodite.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 ab34da4582..4ad0d8aa8f 100644 --- a/tests/v1/worker/test_gpu_input_batch.py +++ b/tests/v1/worker/test_gpu_input_batch.py @@ -204,6 +204,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 = {} @@ -286,6 +287,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, @@ -295,6 +297,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] = [] From c93f4779929dd16e81328deb436e5c55f859181b Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Sun, 12 Jul 2026 07:38:31 +0000 Subject: [PATCH 08/49] [sync] [Core][KV events] Report prefix-cache-reused blocks in full report mode (#45261) Upstream-vLLM: e5588e49bc2642670116664a7fc4096e27adb179 Co-authored-by: GongLei-HW <1327185943@qq.com> --- .sync/vllm-sha | 2 +- aphrodite/v1/core/block_pool.py | 88 ++++++++++++++++----- aphrodite/v1/core/kv_cache_manager.py | 18 +++++ aphrodite/v1/core/kv_cache_utils.py | 17 +++++ aphrodite/v1/request.py | 2 + tests/v1/core/test_prefix_caching.py | 105 +++++++++++++++++++++++++- 6 files changed, 213 insertions(+), 19 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 0a799371f6..9283039c5e 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -95ed0feaa5cd7fb16d72c53ce04950aaf07c4698 +e5588e49bc2642670116664a7fc4096e27adb179 diff --git a/aphrodite/v1/core/block_pool.py b/aphrodite/v1/core/block_pool.py index 3b353662ce..089fa5d092 100644 --- a/aphrodite/v1/core/block_pool.py +++ b/aphrodite/v1/core/block_pool.py @@ -14,8 +14,6 @@ from aphrodite.v1.core.kv_cache_metrics import KVCacheMetricsCollector from aphrodite.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 aphrodite.v1.request import Request @@ -251,14 +250,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) + block_hashes = resolve_block_hashes(request, self.hash_block_size, block_size) assert len(block_hashes) >= num_full_blocks new_block_hashes = block_hashes[num_cached_blocks:] @@ -315,19 +307,81 @@ 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``.""" + 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 prefix-cache-reused blocks.""" + 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) + assert len(block_hashes) >= num_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) + + self.kv_event_queue.append( + self._build_block_stored_event( + request, + block_hashes=cached_hashes, + parent_block_hash=None, + start_token_idx=0, + end_token_idx=num_cached_blocks * block_size, + 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/aphrodite/v1/core/kv_cache_manager.py b/aphrodite/v1/core/kv_cache_manager.py index 2e2500de90..f3778ab040 100644 --- a/aphrodite/v1/core/kv_cache_manager.py +++ b/aphrodite/v1/core/kv_cache_manager.py @@ -127,6 +127,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 @@ -222,6 +223,23 @@ def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int]: request.block_hashes, max_cache_hit_length ) + 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): + if len(group_blocks) == 0: + continue + 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, + len(group_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/aphrodite/v1/core/kv_cache_utils.py b/aphrodite/v1/core/kv_cache_utils.py index 30461a6e22..3180b82cad 100644 --- a/aphrodite/v1/core/kv_cache_utils.py +++ b/aphrodite/v1/core/kv_cache_utils.py @@ -2075,3 +2075,20 @@ 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, provide a view at the + larger cache block granularity used by that KV cache group. + """ + 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/aphrodite/v1/request.py b/aphrodite/v1/request.py index 6a10a64b0b..6754a054f3 100644 --- a/aphrodite/v1/request.py +++ b/aphrodite/v1/request.py @@ -96,6 +96,7 @@ def __init__( # P/D: Connector-specific KV transfer parameters. self.kv_transfer_params: dict[str, Any] | None = None + self.kv_cache_report_mode = "incremental" if pooling_params is not None: # Pooling models. @@ -109,6 +110,7 @@ def __init__( if sampling_params.extra_args is not None: 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: raise ValueError("sampling_params and pooling_params can't both be unset") diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 33bb5003fb..71550dc0c5 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -12,7 +12,12 @@ import aphrodite.v1.core.kv_cache_manager as kv_cache_manager import aphrodite.v1.core.kv_cache_utils as kv_cache_utils -from aphrodite.distributed.kv_events import AllBlocksCleared, BlockRemoved, BlockStored +from aphrodite.distributed.kv_events import ( + MEDIUM_GPU, + AllBlocksCleared, + BlockRemoved, + BlockStored, +) from aphrodite.lora.request import LoRARequest from aphrodite.multimodal.inputs import ( MultiModalFeatureSpec, @@ -2295,6 +2300,104 @@ def test_block_removed_event_group_idx(group_id: int): assert event.group_idx == group_id +def test_emit_cached_block_events(): + block_size = 4 + num_cached_blocks = 3 + kv_cache_group_id = 1 + num_tokens = 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_cached", + prompt_token_ids=list(range(num_tokens)), + block_size=block_size, + hash_fn=sha256, + ) + 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, + ) + + 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 + 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(): + 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(): + 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.""" From 6082897bddd2c76123751e62d219e6e303bed6c1 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Sun, 12 Jul 2026 07:45:45 +0000 Subject: [PATCH 09/49] [sync] [Feature][Parser] Support include_reasoning param for non-Harmony models (#44301) Upstream-vLLM: 5715fde12c1e28eccd08a2394a114339b12a96c1 Co-authored-by: alberto --- .sync/vllm-sha | 2 +- .../openai/chat_completion/serving.py | 31 +- .../entrypoints/openai/responses/context.py | 6 +- .../entrypoints/openai/responses/protocol.py | 9 + .../entrypoints/openai/responses/serving.py | 6 + aphrodite/parser/abstract_parser.py | 18 +- aphrodite/parser/engine/parser_engine.py | 10 +- aphrodite/parser/harmony.py | 10 + aphrodite/v1/core/block_pool.py | 4 +- .../chat_completion/test_include_reasoning.py | 143 ++++++ 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 | 423 ++++++++++++++++++ 14 files changed, 646 insertions(+), 19 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/.sync/vllm-sha b/.sync/vllm-sha index 9283039c5e..35b51e9f59 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -e5588e49bc2642670116664a7fc4096e27adb179 +5715fde12c1e28eccd08a2394a114339b12a96c1 diff --git a/aphrodite/entrypoints/openai/chat_completion/serving.py b/aphrodite/entrypoints/openai/chat_completion/serving.py index 1cafaa1421..bcc9c95f71 100644 --- a/aphrodite/entrypoints/openai/chat_completion/serving.py +++ b/aphrodite/entrypoints/openai/chat_completion/serving.py @@ -564,14 +564,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: @@ -586,11 +580,19 @@ 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() @@ -622,6 +624,8 @@ 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( @@ -629,7 +633,7 @@ async def chat_completion_stream_generator( delta=delta_message, logprobs=logprobs, finish_reason=None, - token_ids=(as_list(output.token_ids) if request.return_token_ids else None), + token_ids=(as_list(output.token_ids) if include_token_ids else None), ) # if the model is finished generating @@ -653,7 +657,7 @@ async def chat_completion_stream_generator( logprobs=logprobs, finish_reason=finish_reason_, stop_reason=output.stop_reason, - token_ids=(as_list(output.token_ids) if request.return_token_ids else None), + token_ids=(as_list(output.token_ids) if include_token_ids else None), ) finish_reason_sent[i] = True @@ -817,12 +821,15 @@ 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 if not request.include_reasoning: reasoning = None + logprobs = None else: reasoning = None content = output.text tool_calls = [] + suppress_metadata = False auto_tools_called = False is_named_tool_choice = ( @@ -906,7 +913,7 @@ async def chat_completion_full_generator( if output.finish_reason else "stop", stop_reason=output.stop_reason, - token_ids=(as_list(output.token_ids) if request.return_token_ids else None), + token_ids=(as_list(output.token_ids) if request.return_token_ids and not suppress_metadata else None), routed_experts=routed_experts_b64, ) choice_data = maybe_filter_parallel_tool_calls(choice_data, request) diff --git a/aphrodite/entrypoints/openai/responses/context.py b/aphrodite/entrypoints/openai/responses/context.py index d63068f407..460bafbd24 100644 --- a/aphrodite/entrypoints/openai/responses/context.py +++ b/aphrodite/entrypoints/openai/responses/context.py @@ -305,7 +305,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 @@ -339,6 +341,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/aphrodite/entrypoints/openai/responses/protocol.py b/aphrodite/entrypoints/openai/responses/protocol.py index 7f2c526066..cb12b369d1 100644 --- a/aphrodite/entrypoints/openai/responses/protocol.py +++ b/aphrodite/entrypoints/openai/responses/protocol.py @@ -159,6 +159,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/aphrodite/entrypoints/openai/responses/serving.py b/aphrodite/entrypoints/openai/responses/serving.py index a3f0d66fdf..d2b294eeeb 100644 --- a/aphrodite/entrypoints/openai/responses/serving.py +++ b/aphrodite/entrypoints/openai/responses/serving.py @@ -997,6 +997,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, @@ -1248,12 +1251,15 @@ async def _process_simple_streaming_events( _increment_sequence_number_and_return: Callable[[StreamingResponsesResponse], StreamingResponsesResponse], ) -> 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/aphrodite/parser/abstract_parser.py b/aphrodite/parser/abstract_parser.py index 98da5632bb..5f1cef71cd 100644 --- a/aphrodite/parser/abstract_parser.py +++ b/aphrodite/parser/abstract_parser.py @@ -431,7 +431,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( @@ -478,6 +478,13 @@ 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 @@ -867,6 +874,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(self, delta_message: DeltaMessage | None) -> DeltaMessage | None: diff --git a/aphrodite/parser/engine/parser_engine.py b/aphrodite/parser/engine/parser_engine.py index 4d0bae9bd1..138115655b 100644 --- a/aphrodite/parser/engine/parser_engine.py +++ b/aphrodite/parser/engine/parser_engine.py @@ -429,7 +429,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/aphrodite/parser/harmony.py b/aphrodite/parser/harmony.py index 3b296b28ff..ddbaec3848 100644 --- a/aphrodite/parser/harmony.py +++ b/aphrodite/parser/harmony.py @@ -270,6 +270,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: diff --git a/aphrodite/v1/core/block_pool.py b/aphrodite/v1/core/block_pool.py index 089fa5d092..e5123f7c86 100644 --- a/aphrodite/v1/core/block_pool.py +++ b/aphrodite/v1/core/block_pool.py @@ -364,9 +364,7 @@ def emit_cached_block_events( 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, curr_mm_idx = generate_block_hash_extra_keys(request, block_start, block_end, curr_mm_idx) extra_keys_list.append(extra_keys) self.kv_event_queue.append( 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 0000000000..f576716a0b --- /dev/null +++ b/tests/entrypoints/openai/chat_completion/test_include_reasoning.py @@ -0,0 +1,143 @@ +# 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, 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 bc0e5e4a1f..3975d4e7e3 100644 --- a/tests/parser/engine/conftest.py +++ b/tests/parser/engine/conftest.py @@ -47,4 +47,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 8518313933..4244cdcff7 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 58233167a5..e6ba2f6a9f 100644 --- a/tests/parser/engine/test_parser_engine.py +++ b/tests/parser/engine/test_parser_engine.py @@ -879,6 +879,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 0000000000..efe1105e9b --- /dev/null +++ b/tests/parser/test_include_reasoning.py @@ -0,0 +1,423 @@ +# 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 aphrodite.entrypoints.openai.chat_completion.protocol import ( # noqa: E402 + ChatCompletionRequest, +) +from aphrodite.entrypoints.openai.engine.protocol import DeltaMessage # noqa: E402 +from aphrodite.entrypoints.openai.responses.protocol import ResponsesRequest # noqa: E402 +from aphrodite.parser.abstract_parser import DelegatingParser # noqa: E402 +from aphrodite.reasoning.basic_parsers import BaseThinkingReasoningParser # noqa: E402 +from aphrodite.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 aphrodite.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) + + +def parse_and_build(parser, request, model_output, enable_auto_tools=False): + """Mirror the non-streaming parse, suppression, and build-items path.""" + from aphrodite.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): + 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): + 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): + 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): + 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): + 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): + request = make_responses_request() + assert request.include_reasoning is True + + def test_include_reasoning_false_suppresses_all_reasoning(self, tokenizer): + 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) + + +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): + 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): + 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): + 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): + 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): + 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): + 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): + 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 + + +class TestParsableContextIncludeReasoning: + def _make_context(self, tokenizer, request): + from aphrodite.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): + from aphrodite.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): + from aphrodite.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 From 15a0fc67b7a3324b0641007c28afaaa314345115 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Sun, 12 Jul 2026 07:47:59 +0000 Subject: [PATCH 10/49] [sync] [Perf] fuse more rmsnorm and all-reduce in qwen3.5 (#46998) Upstream-vLLM: 300e33797fac8a949a0ed89fab2633de96e65fa8 Co-authored-by: Jiangyun Zhu --- .sync/vllm-sha | 2 +- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 36 +++++++++---------- aphrodite/model_executor/models/qwen3_next.py | 16 +++------ 3 files changed, 22 insertions(+), 32 deletions(-) diff --git a/.sync/vllm-sha b/.sync/vllm-sha index 35b51e9f59..f9130962e0 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -5715fde12c1e28eccd08a2394a114339b12a96c1 +300e33797fac8a949a0ed89fab2633de96e65fa8 diff --git a/aphrodite/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/aphrodite/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 8672010a6f..6b4ff589ba 100644 --- a/aphrodite/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/aphrodite/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -820,17 +820,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 @@ -842,13 +839,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: @@ -877,15 +874,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 @@ -936,13 +932,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 @@ -985,13 +980,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) @@ -1031,7 +1026,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/aphrodite/model_executor/models/qwen3_next.py b/aphrodite/model_executor/models/qwen3_next.py index 8035a487cf..9b2e644308 100644 --- a/aphrodite/model_executor/models/qwen3_next.py +++ b/aphrodite/model_executor/models/qwen3_next.py @@ -356,15 +356,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): @@ -452,21 +452,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 39e9ecd1a87b167597a681b98b168cd286d29d24 Mon Sep 17 00:00:00 2001 From: AlpinDale Date: Sun, 12 Jul 2026 07:57:00 +0000 Subject: [PATCH 11/49] [sync] [Rust Frontend] Integrate MM video support (#47959) Upstream-vLLM: 074bdd0d9961fb7f81c9311e00c00cfd5aff77a3 Co-authored-by: Bugen Zhao --- .sync/vllm-sha | 2 +- 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 + 17 files changed, 1571 insertions(+), 518 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/.sync/vllm-sha b/.sync/vllm-sha index f9130962e0..2f063f1f93 100644 --- a/.sync/vllm-sha +++ b/.sync/vllm-sha @@ -1 +1 @@ -300e33797fac8a949a0ed89fab2633de96e65fa8 +074bdd0d9961fb7f81c9311e00c00cfd5aff77a3 diff --git a/rust/Cargo.lock b/rust/Cargo.lock index d950956d82..b5ea0b8c6f 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -2483,7 +2483,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", @@ -2774,8 +2774,8 @@ dependencies = [ "portable-atomic", "portable-atomic-util", "rawpointer", + "serde", ] - [[package]] name = "ndarray" version = "0.17.2" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 7868165792..08352acefb 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 fcca07346b..0425042dba 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 1c3aebf111..b94a446f70 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 3ab1fa47a5..6e2662b1b7 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 5e25f45d68..70014099c0 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 aphrodite_engine_core_client::protocol::dtype::ModelDtype; -use aphrodite_engine_core_client::protocol::multimodal::{ - MmBatchedField, MmFeatureSpec, MmFeatures, MmField, MmFieldElem, MmFlatField, MmKwargsItem, - MmSharedField, MmSlice, PlaceholderRange, SliceSpec, -}; -use aphrodite_engine_core_client::protocol::tensor::WireTensor; +use aphrodite_engine_core_client::protocol::multimodal::{MmFeatureSpec, MmFeatures, MmKwargsItem}; use aphrodite_text::Prompt; use aphrodite_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 aphrodite_engine_core_client::protocol::tensor::WireArrayData; use aphrodite_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 0000000000..3ea20db3b9 --- /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 aphrodite_engine_core_client::protocol::multimodal::PlaceholderRange; +use aphrodite_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 aphrodite_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 0000000000..53741b6534 --- /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 aphrodite_engine_core_client::protocol::dtype::ModelDtype; +use aphrodite_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 6c9bab8ad1..dbed8e1512 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 aphrodite_engine_core_client::protocol::dtype::ModelDtype; use aphrodite_engine_core_client::protocol::multimodal::MmKwargValue as ProtocolKwargValue; use aphrodite_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 Aphrodite's batched-field semantics. + /// Batched fields use their first axis as media-item index and drop that + /// axis in the per-feature value, matching Aphrodite'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 0000000000..f33c4675b3 --- /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 aphrodite_engine_core_client::protocol::dtype::ModelDtype; +use aphrodite_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 aphrodite_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 3e252ac0a1..ab694ec543 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("