Update vLLM model wrapper, Qwen3 VL patcher, and vision tower JIT for…

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()

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

tpu_inference/models/vllm/experimental/qwen3_vl_patcher.py

@@ -45,19 +45,78 @@
 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
 from vllm.sequence import IntermediateTensors
 
 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."""

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