diff --git a/tpu_inference/core/core_tpu.py b/tpu_inference/core/core_tpu.py index 162f7e9aa6..771fe02126 100644 --- a/tpu_inference/core/core_tpu.py +++ b/tpu_inference/core/core_tpu.py @@ -16,6 +16,7 @@ # ====================================================================================== from vllm.config import VllmConfig from vllm.logger import init_logger +from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.tasks import POOLING_TASKS, SupportedTask from vllm.v1.core.kv_cache_utils import (get_request_block_hasher, init_none_hash) @@ -513,7 +514,8 @@ def __init__( self.step_fn = (self.step if self.batch_queue is None else self.step_with_batch_queue) - self.mm_receiver_cache = None + self.mm_receiver_cache = MULTIMODAL_REGISTRY.engine_receiver_cache_from_config( + vllm_config.model_config) self._orchestrator = _DisaggOrchestrator( config=vllm_config, output_queue=self.output_queue, @@ -739,7 +741,8 @@ def executor_fail_callback(): self.request_block_hasher = get_request_block_hasher( block_size, caching_hash_fn) - self.mm_receiver_cache = None + self.mm_receiver_cache = MULTIMODAL_REGISTRY.engine_receiver_cache_from_config( + vllm_config.model_config) self._orchestrator = _DisaggOrchestrator( config=vllm_config, output_queue=self.output_queue, @@ -754,6 +757,18 @@ def add_request(self, request: EngineCoreRequest, request_wave: int = 0): raise TypeError( f"request_id must be a string, got {type(request.request_id)}") + if self.mm_receiver_cache is not None and getattr( + request, "mm_inputs", None): + request.mm_inputs = self.mm_receiver_cache.get_and_update_features( + request.mm_inputs) + elif self.mm_receiver_cache is not None and getattr( + request, "mm_positions", None) and hasattr( + request, "mm_features"): + # Older vllm might use mm_features, checking both to be safe + if request.mm_features: + request.mm_features = self.mm_receiver_cache.get_and_update_features( + request.mm_features) + if pooling_params := request.pooling_params: supported_pooling_tasks = [ task for task in self.get_supported_tasks() diff --git a/tpu_inference/envs.py b/tpu_inference/envs.py index da14eed861..c2ac7d9315 100644 --- a/tpu_inference/envs.py +++ b/tpu_inference/envs.py @@ -64,6 +64,8 @@ MOE_APPROX_TOPK: bool = False MOE_APPROX_TOPK_RECALL_TARGET: float | None = None VLLM_TPU_PATCH_MM_EMBEDDINGS: bool = False + VISION_PRECOMPILE_FRAMES: list[int] = [] + VISION_MIN_SHIFT: int = 6 ENABLE_RS_KERNEL: bool = False NUM_PRECOMPILE_WORKERS: int = 1 DP_SCHED_BATCH_PREFILL: bool = False @@ -392,6 +394,10 @@ def _get_int_list_env() -> list[int]: env_bool("DISABLE_WEIGHT_REQUANTIZATION", default=False), "VLLM_TPU_PATCH_MM_EMBEDDINGS": env_bool("VLLM_TPU_PATCH_MM_EMBEDDINGS", default=False), + "VISION_PRECOMPILE_FRAMES": + env_int_list("VISION_PRECOMPILE_FRAMES"), + "VISION_MIN_SHIFT": + lambda: int(os.getenv("VISION_MIN_SHIFT", "6")), "DISABLE_MLA_Q_ACTIVATION_QUANTIZATION": env_bool("DISABLE_MLA_Q_ACTIVATION_QUANTIZATION", default=False), # Enable hierarchical reduce-scatter kernel for MoE diff --git a/tpu_inference/models/vllm/experimental/qwen3_vl_patcher.py b/tpu_inference/models/vllm/experimental/qwen3_vl_patcher.py index d53ebe9c10..eeacd2cfe9 100644 --- a/tpu_inference/models/vllm/experimental/qwen3_vl_patcher.py +++ b/tpu_inference/models/vllm/experimental/qwen3_vl_patcher.py @@ -45,8 +45,10 @@ import jax.numpy as jnp import torch import torch.nn as nn +import torch.nn.functional as F import vllm.model_executor.models.qwen3_vl as qwen3_vl_mod import vllm.model_executor.models.utils as vllm_utils +from jax.experimental.pallas import tpu as pltpu from torchax.interop import jax_view, torch_view from vllm.model_executor.models.qwen3_vl import Qwen3VLForConditionalGeneration from vllm.multimodal import NestedTensors @@ -54,10 +56,67 @@ from tpu_inference.distributed.jax_parallel_state import \ get_pp_group as jax_get_pp_group +from tpu_inference.kernels.flash_attention.kernel import \ + encoder_only_flash_attention from tpu_inference.logger import init_logger +from tpu_inference.utils import align_to logger = init_logger(__name__) +_orig_sdpa = F.scaled_dot_product_attention + + +def _chunked_sdpa(query, + key, + value, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + scale=None): + if query.ndim == 4: + B, H, S, D = query.shape + + # 1. Fallback for GQA (our custom kernel requires num_heads == num_kv_heads) + if query.shape != key.shape: + return _orig_sdpa(query, key, value, attn_mask, dropout_p, + is_causal, scale) + + # 2. Heuristic for native SDPA memory usage: B * H * S * S * element_size bytes + # We account for hardware alignment (e.g. 128 lanes) to prevent VMEM OOM. + lane_count = pltpu.get_tpu_info().num_lanes + S_padded = align_to(S, lane_count) + estimated_vmem_native = B * H * S_padded * S_padded * query.element_size( + ) + + # 3. Only use chunked flash attention if we will OOM native SDPA (>40MB) + if attn_mask is None and dropout_p == 0.0 and estimated_vmem_native > 40 * 1024 * 1024: + q_jax = jax_view(query) + k_jax = jax_view(key) + v_jax = jax_view(value) + + # [B, H, S, D] -> [B, S, H, D] -> [B*S, H, D] + q_jax = jnp.reshape(jnp.swapaxes(q_jax, 1, 2), (B * S, H, D)) + k_jax = jnp.reshape(jnp.swapaxes(k_jax, 1, 2), (B * S, H, D)) + v_jax = jnp.reshape(jnp.swapaxes(v_jax, 1, 2), (B * S, H, D)) + + seq_lens = jnp.full((B, ), S, dtype=jnp.int32) + if scale is None: + scale = 1.0 / (D**0.5) + + out_jax = encoder_only_flash_attention(q_jax, + k_jax, + v_jax, + seq_lens, + causal=is_causal, + sm_scale=scale) + + # [B*S, H, D] -> [B, S, H, D] -> [B, H, S, D] + out_jax = jnp.swapaxes(jnp.reshape(out_jax, (B, S, H, D)), 1, 2) + return torch_view(out_jax) + + return _orig_sdpa(query, key, value, attn_mask, dropout_p, is_causal, + scale) + def _patched_set_deepstack(vllm_model, deepstack_input_embeds): """Intercepts Deepstack embeddings to store them in a JAX-friendly cached tensors (`_deepstack_tensors`). @@ -223,6 +282,12 @@ def _patched_embed_input_ids(vllm_model, orig_embed_input_ids, *args, # 1. Get the base text embeddings from the native model. inputs_embeds = orig_embed_input_ids(*args, **kwargs) + # 1.5. Bypass deepstack packing during autoregressive decode or text-only prefill. + # In vllm_model_wrapper, mm_embeds is passed as args[1] during prefill. If len(args) == 1, + # it means mm_embeds is None (decode phase or text-only). We also check size(0) == 1 just in case. + if len(args) == 1 or inputs_embeds.size(0) == 1: + return inputs_embeds + # 2. Check if there are any deepstack features to pack. # Read from _deepstack_tensors (our JAX-compatible cache) rather than # vllm_model.deepstack_input_embeds (the pre-allocated placeholder buffer that @@ -428,6 +493,9 @@ def apply_qwen3_vl_patches(vllm_model): "Disabled dynamo for Qwen3LLMModel; JAX JIT handles outer compilation" ) + # 8. Patch PyTorch SDPA globally to catch Qwen Vision Tower calls and prevent VMEM OOM + F.scaled_dot_product_attention = _chunked_sdpa + def is_qwen3_vl(vllm_model) -> bool: """Check if the given vLLM model is of architecture Qwen3VLForConditionalGeneration.""" diff --git a/tpu_inference/models/vllm/experimental/vision_tower_jit.py b/tpu_inference/models/vllm/experimental/vision_tower_jit.py index 6b4a2e86c1..f8aba5df8b 100644 --- a/tpu_inference/models/vllm/experimental/vision_tower_jit.py +++ b/tpu_inference/models/vllm/experimental/vision_tower_jit.py @@ -24,9 +24,12 @@ from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import \ Qwen3OmniMoeConfig from vllm.config import VllmConfig +from vllm.model_executor.models.qwen3_5 import ( + Qwen3_5ForConditionalGeneration, Qwen3_5MoeForConditionalGeneration) from vllm.model_executor.models.qwen3_omni_moe_thinker import \ Qwen3OmniMoeThinkerForConditionalGeneration +from tpu_inference import envs from tpu_inference.logger import init_logger from tpu_inference.utils import to_jax_dtype @@ -35,6 +38,8 @@ # Architectures whose embed_multimodal function is safe to wrap with jax.jit. JITTABLE_ARCHS = { Qwen3OmniMoeThinkerForConditionalGeneration, + Qwen3_5MoeForConditionalGeneration, + Qwen3_5ForConditionalGeneration, } @@ -82,6 +87,7 @@ def maybe_jit_embed_multimodal_func(embed_multimodal_func_jax: Callable, return embed_multimodal_func_jax +@jax.tree_util.register_pytree_node_class class GridTHW(tuple): """Tensor-like wrapper for image/video grid_thw arguments. @@ -103,6 +109,12 @@ def _nested_to_tuple(v): flat: tuple = _nested_to_tuple(values) return super().__new__(cls, flat) + def __getitem__(self, key): + val = super().__getitem__(key) + if isinstance(key, slice): + return type(self)(val) + return val + # ---- tensor-like API expected by _process_image_input ---- @property @@ -124,6 +136,13 @@ def prod(self, dim=-1): def __repr__(self): return f"GridTHW({tuple(self)})" + def tree_flatten(self): + return (), tuple(self) + + @classmethod + def tree_unflatten(cls, aux_data, children): + return cls(aux_data) + def maybe_precompile_vision_encoder_fn( params: Any, embed_multimodal_fn: Optional[Callable], vllm_model, @@ -151,7 +170,12 @@ def maybe_precompile_vision_encoder_fn( spatial_merge_unit = vc.spatial_merge_size**2 max_patches = (vllm_config.scheduler_config.max_num_batched_tokens // spatial_merge_unit) - min_shift = 4 # 1 << 4 = 16 patches minimum + min_shift = envs.VISION_MIN_SHIFT + if min_shift < 6: + logger.warning( + f"VISION_MIN_SHIFT is set to {min_shift} (< 6). " + "This may cause JAX divisibility errors on TPUs with 8+ devices " + "if spatial merging is active.") max_shift = max(min_shift, (max(max_patches, 1) - 1).bit_length()) num_patches_paddings = [1 << i for i in range(min_shift, max_shift + 1)] @@ -165,20 +189,49 @@ def precompile_fn(run_compilation_fn: Callable) -> None: h = 1 << (k // 2) w = 1 << (k - k // 2) - dummy_pixel_values = jnp.ones((num_patches, patch_input_dim), - dtype=jax_dtype) - dummy_image_grid_thw = GridTHW([(1, h, w)]) - - run_compilation_fn( - f"vllm embed_multimodal {dummy_image_grid_thw}", - embed_multimodal_fn, - params, - call_kwargs={ - "pixel_values": dummy_pixel_values, - "image_grid_thw": dummy_image_grid_thw, - }, - num_patches=num_patches, - ) + # By default, we precompile for common small frame counts to balance startup time. + # Users can override this via the VISION_PRECOMPILE_FRAMES environment variable + # (e.g., VISION_PRECOMPILE_FRAMES="1,2,4,8,16,64") to support specific video lengths. + # + # ⚠️ WARNING: Adding more frames or larger buckets here will significantly increase + # server startup time (XLA compilation) and can cause Host CPU OOMs during boot. + frame_counts = [1, 2, 4, 8, 16] + if envs.VISION_PRECOMPILE_FRAMES: + frame_counts = envs.VISION_PRECOMPILE_FRAMES + logger.info( + f"Using custom vision precompile frames: {frame_counts}") + + for t_val in frame_counts: + # Limit batch sizes to prevent astronomical compilation time and host OOMs. + # If users submit larger batches, they will incur a one-time compilation cost at runtime. + batch_sizes = [1, 2] if t_val == 1 else [1] + for b in batch_sizes: + dummy_pixel_values = jnp.ones( + (b * t_val * num_patches, patch_input_dim), + dtype=jax_dtype) + dummy_image_grid_thw = GridTHW([(t_val, h, w)] * b) + + model_type = getattr(vllm_config.model_config.hf_config, + "model_type", "") + if model_type in ("qwen2_vl", "qwen2_5_vl", "qwen", + "qwen3_5_moe", "qwen3_5", "qwen3_vl"): + grid_keys = ("image_grid_thw", "video_grid_thw") + else: + grid_keys = ("image_grid_thw", "video_grid_thw", + "grid_thw") + + for grid_key in grid_keys: + pixel_key = "pixel_values_videos" if grid_key == "video_grid_thw" else "pixel_values" + run_compilation_fn( + f"vllm embed_multimodal {grid_key}={dummy_image_grid_thw}", + embed_multimodal_fn, + params, + call_kwargs={ + pixel_key: dummy_pixel_values, + grid_key: dummy_image_grid_thw, + }, + num_patches=num_patches, + ) return precompile_fn @@ -199,4 +252,10 @@ def maybe_prepare_for_jit(kwargs: dict, vllm_model) -> dict: elif k == "audio_feature_lengths" and isinstance(v, torch.Tensor): kwargs[k] = tuple(v.tolist()) + elif k == "timestamps": + if isinstance(v, list): + kwargs[k] = torch.tensor(v, dtype=torch.float32) + elif isinstance(v, (float, int)): + kwargs[k] = torch.tensor([v], dtype=torch.float32) + return kwargs diff --git a/tpu_inference/models/vllm/vllm_model_wrapper.py b/tpu_inference/models/vllm/vllm_model_wrapper.py index cb71896a23..42900efb8f 100644 --- a/tpu_inference/models/vllm/vllm_model_wrapper.py +++ b/tpu_inference/models/vllm/vllm_model_wrapper.py @@ -63,8 +63,8 @@ from tpu_inference.models.vllm.experimental.model_patcher import ( apply_model_specific_patches, patch_mm_model) from tpu_inference.models.vllm.experimental.vision_tower_jit import ( - maybe_jit_embed_multimodal_func, maybe_precompile_vision_encoder_fn, - maybe_prepare_for_jit) + GridTHW, maybe_jit_embed_multimodal_func, + maybe_precompile_vision_encoder_fn, maybe_prepare_for_jit) from tpu_inference.models.vllm.vllm_model_wrapper_context import ( get_vllm_model_wrapper_context, set_vllm_model_wrapper_context) from tpu_inference.runner.lora_utils import replace_lora_metadata @@ -505,8 +505,16 @@ def embed_multimodal_func_jax( params_and_buffers: Any, **kwargs, ) -> Any: + # Convert JAX arrays back to Torch Tensors for the Torch model proxy. + # Static metadata like GridTHW should be passed as is (it's already converted to static). call_kwargs = { - k: jax.tree.map(torch_view, v) + k: + jax.tree.map( + lambda x: torch_view(x) if isinstance(x, jax.Array) else x, + v, + is_leaf=lambda x: isinstance(x, (GridTHW, tuple)) and + not isinstance(x, jax.Array), + ) for k, v in kwargs.items() } @@ -541,23 +549,223 @@ def embed_multimodal_func_torch(params_and_buffers: Any, kwargs = maybe_prepare_for_jit(kwargs, self.model.vllm_model) - def move(v: torch.Tensor) -> torch.Tensor: - if not isinstance(v, torch.Tensor): - logger.warning(f"Expect torch.Tensor, got {type(v)}") - return v - return t2j(v, use_dlpack=False) + is_video = "video_grid_thw" in kwargs + grid_key = "video_grid_thw" if is_video else "image_grid_thw" + pixel_key = "pixel_values_videos" if is_video else "pixel_values" + + padded_anything = False + if grid_key in kwargs and pixel_key in kwargs: + grid = kwargs[grid_key] + pixels = kwargs[pixel_key] + b = len(grid) + original_batch_len = b + original_pixels_len = pixels.shape[0] + padded_anything = False + + # Limit valid batch sizes for vision tower to avoid multi-hour XLA compilations. + valid_batch_sizes = [1, 2, 4, 8, 16, 32, 64] + if b not in valid_batch_sizes: + target_b = next( + (x for x in valid_batch_sizes if x >= b), None) + if target_b is None: + target_b = 1 << (b - 1).bit_length() + + if target_b > b: + pad_b = target_b - b + padded_anything = True + + # 1. Homogeneous padding to maximize cache hits + grid_list = list(grid) if type( + grid).__name__ == "GridTHW" else grid.tolist( + ) if isinstance(grid, + torch.Tensor) else list(grid) + + # Find the maximum patch count across the batch to use for padding + max_grid = grid_list[-1] # fallback + max_patches = 0 + for g in grid_list: + p = int(g[0] * g[1] * g[2]) + if p > max_patches: + max_patches = p + max_grid = g + + kwargs[grid_key] = type(grid)( + list(grid) + [max_grid] * pad_b) if type( + grid + ).__name__ == "GridTHW" else torch.tensor( + grid_list + [max_grid] * pad_b, + device=grid.device) + + pad_pixels = torch.zeros( + (pad_b * max_patches, *pixels.shape[1:]), + dtype=pixels.dtype, + device=pixels.device) + kwargs[pixel_key] = torch.cat([pixels, pad_pixels], + dim=0) + + if "second_per_grid_ts" in kwargs: + ts = kwargs["second_per_grid_ts"] + if ts is not None and isinstance( + ts, torch.Tensor): + pad_ts = torch.full((pad_b, ), + ts[-1].item(), + dtype=ts.dtype, + device=ts.device) + kwargs["second_per_grid_ts"] = torch.cat( + [ts, pad_ts], dim=0) + + if "timestamps" in kwargs: + ts_vals = kwargs["timestamps"] + if ts_vals is not None and isinstance( + ts_vals, torch.Tensor): + pad_ts_vals = ts_vals[-1].unsqueeze( + 0).repeat(pad_b, + *([1] * (ts_vals.ndim - 1))) + kwargs["timestamps"] = torch.cat( + [ts_vals, pad_ts_vals], dim=0) + + # 2. Pad pixel_values sequence length to multiple of tp_size to avoid full replication OOM + # and pad grid_thw to match the sequence length + if hasattr(self, "vllm_config") and hasattr( + self.vllm_config, "parallel_config"): + tp_size = self.vllm_config.parallel_config.tensor_parallel_size + else: + tp_size = 8 + + pixels = kwargs[pixel_key] + seq_len = pixels.shape[0] + if seq_len % tp_size != 0: + padded_anything = True + pad_seq = tp_size - (seq_len % tp_size) + + # Find a pad_seq that can be factored into t, h, w where h>=2 and w>=2 + # to prevent spatial_merge_size downsampling from crashing on 0. + dummy_grid_thw = None + while dummy_grid_thw is None: + for h in range(2, pad_seq + 1): + if dummy_grid_thw is not None: + break + for w in range(2, (pad_seq // h) + 1): + if pad_seq % (h * w) == 0: + dummy_grid_thw = [ + pad_seq // (h * w), h, w + ] + break + if dummy_grid_thw is None: + pad_seq += tp_size + + pad_pixels_seq = torch.zeros( + (pad_seq, *pixels.shape[1:]), + dtype=pixels.dtype, + device=pixels.device) + kwargs[pixel_key] = torch.cat([pixels, pad_pixels_seq], + dim=0) + + # Add a dummy grid entry for the sequence padding so the length matches + grid = kwargs[grid_key] + grid_list = list(grid) if type( + grid).__name__ == "GridTHW" else grid.tolist( + ) if isinstance(grid, torch.Tensor) else list(grid) + dummy_grid = type( + grid_list[-1])(dummy_grid_thw) if isinstance( + grid_list[-1], (list, tuple)) else pad_seq + kwargs[grid_key] = type(grid)( + grid_list + [dummy_grid]) if type( + grid).__name__ == "GridTHW" else torch.tensor( + grid_list + [dummy_grid], + device=grid.device) + + def make_move(param_name: str): + + def move(v: Any) -> Any: + if isinstance( + v, + (int, float, str, bool, jax.Array + )) or v is None or type(v).__name__ == "GridTHW": + return v + if isinstance(v, (tuple, list)): + return type(v)(move(x) for x in v) + if not isinstance(v, torch.Tensor): + logger.warning( + f"Expect torch.Tensor, got {type(v)}") + return v + arr = t2j(v.contiguous(), use_dlpack=True) + if hasattr(self, "mesh") and getattr( + self.mesh, "devices", None) is not None: + from jax.sharding import (NamedSharding, + PartitionSpec) + + # Shard pixel values over model (TP) axis if divisible + if param_name in ("pixel_values", + "pixel_values_videos"): + if hasattr(self, "vllm_config") and hasattr( + self.vllm_config, "parallel_config"): + tp_size = self.vllm_config.parallel_config.tensor_parallel_size + else: + tp_size = 8 + if arr.shape[0] % tp_size == 0: + spec = PartitionSpec("model", None) + else: + spec = PartitionSpec() + else: + spec = PartitionSpec() + + arr = jax.device_put( + arr, NamedSharding(self.mesh, spec)) + return arr + + return move # Ensure all tensors are moved into accelerator so the # computation with weights can work properly. call_kwargs = { - k: jax.tree.map(move, v) + k: + jax.tree.map( + make_move(k), + v, + is_leaf=lambda x: type(x).__name__ == "GridTHW" or + isinstance(x, (tuple, torch.Tensor, jax.Array))) for k, v in kwargs.items() } - return maybe_jit_embed_multimodal_func( + out = maybe_jit_embed_multimodal_func( embed_multimodal_func_jax, self.model.vllm_model)(params_and_buffers, **call_kwargs) + # 3. Strip both batch padding and sequence padding using the original lengths + if padded_anything: + if isinstance(out, (list, tuple)): + # Case A: List/Tuple of items -> Strip batch padding directly + out = out[:original_batch_len] + elif hasattr(out, "shape"): # jax.Array or torch.Tensor + if len(out.shape) == 3: + # Case B: 3D Batched Tensor (Batch, SeqLen, Hidden) + # Strip batch padding on Dimension 0 + out = out[:original_batch_len, ...] + elif len(out.shape) == 2: + # Case C: 2D Flattened/Packed Tensor (TotalTokens, Hidden) + # Dynamic merge factor calculation + visual = getattr(self.model.vllm_model, "visual", + None) + if visual is None and hasattr( + self.model.vllm_model, "model"): + visual = getattr(self.model.vllm_model.model, + "vision_tower", None) + + merge_size = getattr(visual, "spatial_merge_size", + 1) if visual else 1 + merge_factor = merge_size**2 + + # Just grab how many patches existed before padding, divided by merge factor + original_patch_len = original_pixels_len // merge_factor + out = out[:original_patch_len, :] + else: + logger.warning( + f"Unexpected multimodal output shape: {out.shape}. Stripping skipped." + ) + + return out + return embed_multimodal_func_torch def wrap_embed_input_ids_func(self): diff --git a/tpu_inference/worker/tpu_worker.py b/tpu_inference/worker/tpu_worker.py index 69cbe675cc..ab83795756 100644 --- a/tpu_inference/worker/tpu_worker.py +++ b/tpu_inference/worker/tpu_worker.py @@ -262,6 +262,12 @@ def __init__( # step_counter is used to calculate uuid to transfer intermediate tensors. self.step_counter = 0 + import multiprocessing + + from vllm.multimodal import MULTIMODAL_REGISTRY + self.mm_receiver_cache = MULTIMODAL_REGISTRY.worker_receiver_cache_from_config( + self.vllm_config, multiprocessing.Lock()) + def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks: int) -> None: self.cache_config.num_gpu_blocks = num_gpu_blocks @@ -567,6 +573,18 @@ def execute_model( self, scheduler_output: SchedulerOutput, ) -> Optional[ModelRunnerOutput]: + if getattr(self, "_apply_mm_cache", None) is not None: + self._apply_mm_cache(scheduler_output) + elif self.mm_receiver_cache is not None: + for req in scheduler_output.scheduled_new_reqs: + if getattr(req, "mm_inputs", None): + req.mm_inputs = self.mm_receiver_cache.get_and_update_features( + req.mm_inputs) + elif getattr(req, "mm_positions", None) and hasattr( + req, "mm_features") and req.mm_features: + req.mm_features = self.mm_receiver_cache.get_and_update_features( + req.mm_features) + # NOTE: This method intentionally returns a concrete vLLM type, which # violates the pure abstract contract of the base class. This is a # deliberate, temporary compromise for the same reasons outlined in