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

[amitkumar307d]

Description

📊 Technical Analysis

This commit addresses several critical memory inefficiencies, OOMs, and runtime stalls when performing large-scale video inference with Qwen architectures on TPUs.

---

1. 🧮 Memory Inefficiency & OOMs in Multimodal Forward Pass

• ❌ The Issue: For large inputs like video, the pixel_values tensor was being fully replicated across all TPU devices. This led to severe HBM/Host Out-of-Memory (OOM) errors. Furthermore, input sequence lengths were often not divisible by the Tensor Parallelism (tp_size) factor, making efficient sharding impossible without crashing downstream spatial downsampling operations.
• 🧠 Approach: Instead of replicating huge video tensors, shard them across the Model (TP) axis. To do this safely, dynamically pad the sequence length to a multiple of tp_size and valid factoring dimensions.
• ✅ The Fix:
  • Padded pixel_values and grid_thw to be multiples of tp_size.
  • Ensured padding factors into valid (t, h, w) dimensions (where $h \ge 2, w \ge 2$) to satisfy downstream kernel constraints.
  • Applied explicit sharding using PartitionSpec("model", None).
  • Stripped padding post-forward pass using original lengths to maintain exact correctness.

🔍 Code (tpu_inference/models/vllm/vllm_model_wrapper.py)

# Pad pixel_values sequence length to multiple of tp_size to avoid full replication OOM
seq_len = pixels.shape[0]
if seq_len % tp_size != 0:
    pad_seq = tp_size - (seq_len % tp_size)
    # ... logic to find valid dummy_grid_thw [t, h, w] ...
    kwargs[pixel_key] = torch.cat([pixels, pad_pixels_seq], dim=0)

# Shard pixel values over model (TP) axis if divisible
if param_name in ("pixel_values", "pixel_values_videos"):
    spec = PartitionSpec("model", None) if arr.shape[0] % tp_size == 0 else PartitionSpec()

---

2. 💥 VMEM OOM in Vision Tower Attention

• ❌ The Issue: Native PyTorch scaled_dot_product_attention (SDPA) consumed excessive VMEM ($>40\text{MB}$) for long video sequences in the Vision Tower, triggering hard TPU VMEM crashes.
• 🧠 Approach: Intercept the standard SDPA call and transparently redirect execution to a memory-efficient Flash Attention kernel (encoder_only_flash_attention) based on dynamic memory heuristics.
• ✅ The Fix:
  • Globally patched F.scaled_dot_product_attention for Qwen3 VL.
  • Implemented _chunked_sdpa which estimates native VMEM usage ($B \times H \times S \times S \times 2\text{ bytes}$).
  • If estimated usage exceeds 40MB, it redirects to the highly optimized encoder_only_flash_attention kernel from tpu-inference.

🔍 Code (tpu_inference/models/vllm/experimental/qwen3_vl_patcher.py)

# Global Patch
F.scaled_dot_product_attention = _chunked_sdpa

def _chunked_sdpa(...):
    # ... heuristics ...
    if attn_mask is None and dropout_p == 0.0 and estimated_vmem_native > 40 * 1024 * 1024:
        # Fallback to custom memory-efficient kernel
        out_jax = encoder_only_flash_attention(q_jax, k_jax, v_jax, seq_lens, ...)
        return torch_view(out_jax)

---

3. ⏳ Stalls Due to Missing JIT Support & Precompilation

• ❌ The Issue: Newer Qwen 3.5 architectures were not registered as "Jittable", missing out on JAX optimizations. Additionally, the vision precompiler only generated graphs for static images ($t=1$), causing major runtime compilation stalls ("JIT lag") during the first video request.
• 🧠 Approach: Whitelist newer architectures and expand the precompiler heat-up routine to simulate video frame batches.
• ✅ The Fix:
  • Added Qwen3_5MoeForConditionalGeneration and Qwen3_5ForConditionalGeneration to JITTABLE_ARCHS.
  • Expanded maybe_precompile_vision_encoder_fn to iterate over common video frame counts [1, 2, 4, 8, 16] to fully warm up the JIT cache before requests arrive.

🔍 Code (tpu_inference/models/vllm/experimental/vision_tower_jit.py)

JITTABLE_ARCHS = {
    Qwen3OmniMoeThinkerForConditionalGeneration,
    Qwen3_5MoeForConditionalGeneration,
    Qwen3_5ForConditionalGeneration,
}

# Video Precompilation Warming Up
frame_counts = [1, 2, 4, 8, 16]
for t_val in frame_counts:
    # ... warmup execution ...

---

4. 🗄️ Missing Multimodal Cache Integration

• ❌ The Issue: Redundant computation and data transfers occurred for multimodal features because the TPU worker and engine cores were not actively retrieving features from the mm_receiver_cache.
• 🧠 Approach: Activate and plumb the receiver cache directly into the core execution paths.
• ✅ The Fix: Initialized and invoked mm_receiver_cache in TPUWorker and DisaggEngineCoreProc before running model execution to reuse cached features.

🔍 Code (tpu_inference/worker/tpu_worker.py)

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)

---

Tests

Model server command

python3 -m vllm.entrypoints.openai.api_server \
  --host=0.0.0.0 \
  --port=8000 \
  --download-dir=/data \
  --kv-cache-dtype=fp8 \
  --max-model-len=4096 \
  --gpu-memory-utilization=0.9 \
  --tensor-parallel-size=8 \
  --max-num-seqs=128 \
  --max-num-batched-tokens=1024 \
  --no-enable-prefix-caching \
  --enable-expert-parallel \
  --async-scheduling \
  --model=gs://amitmkumar-bucket/gcloud/Qwen3.5-397B-A17B-FP8 \
  --served-model-name=Qwen/Qwen3.5-397B-A17B-FP8 \
  --load-format=auto \
  --enable-auto-tool-choice \
  --reasoning-parser=qwen3 \
  --limit-mm-per-prompt='{"image": 0, "video": 1}' \
  --tool-call-parser=qwen3_coder \
  --block-size=256 \
  --additional-config='{"compilation_sizes": [512, 1024, 2048, 4096],"sharding": {"sharding_strategy": {"enable_dp_attention": true}}}' || true

Benchmark command

vllm bench serve \
  --model Qwen/Qwen3.5-397B-A17B-FP8 \
  --backend openai-chat \
  --endpoint /v1/chat/completions \
  --dataset-name random-mm \
  --base-url http://qwen-3-5-397b-tpu-9k-video-svc:8000 \
  --num-prompts 640 \
  --num-warmups 640
  --max-concurrency 128 \
  --random-input-len 500 \
  --random-output-len 1024 \
  --random-mm-base-items-per-request 1 \
  --random-mm-limit-mm-per-prompt '{"image": 0, "video": 1}' \
  --random-mm-bucket-config '{"(256, 256, 16)": 1.0}' \
  --request-rate inf \
  --percentile-metrics ttft,tpot,itl,e2el \
  --ignore-eos 

Performance metrics

============ Serving Benchmark Result ============
Successful requests:                     640       
Failed requests:                         0         
Maximum request concurrency:             128       
Benchmark duration (s):                  202.84    
Total input tokens:                      420098    
Total generated tokens:                  655360    
Request throughput (req/s):              3.16      
Output token throughput (tok/s):         3230.92   
Peak output token throughput (tok/s):    4480.00   
Peak concurrent requests:                192.00    
Total token throughput (tok/s):          5302.00   
---------------Time to First Token----------------
Mean TTFT (ms):                          2609.97   
Median TTFT (ms):                        1219.77   
P99 TTFT (ms):                           11138.57  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          37.05     
Median TPOT (ms):                        37.37     
P99 TPOT (ms):                           41.15     
---------------Inter-token Latency----------------
Mean ITL (ms):                           38.05     
Median ITL (ms):                         32.10     
P99 ITL (ms):                            96.49     
----------------End-to-end Latency----------------
Mean E2EL (ms):                          40507.93  
Median E2EL (ms):                        39579.70  
P99 E2EL (ms):                           45238.58  
==================================================

Checklist

Before submitting this PR, please make sure:
-[X] I have performed a self-review of my code.
-[X] I have necessary comments in my code, particularly in hard-to-understand areas.
-[X] I have made or will make corresponding changes to any relevant documentation.

[QiliangCui]

Thanks — the problems this targets (host/HBM replication OOM on large video inputs, VMEM OOM in vision-tower SDPA, cold video compiles) are real, and the pixel-shard + strip-after-forward design and the GridTHW pytree/slice changes look sound. I reviewed this against the vLLM commit this repo pins (.buildkite/vllm_lkg.version = 0ba2aa35a), and there are three things that need to be addressed before this can merge, plus inline comments.

1. The mm receiver-cache integration (item 4 in the description) doesn't function at the pinned vLLM. engine_receiver_cache_from_config is called with a ModelConfig where it takes the full VllmConfig → AttributeError at init of both disagg engine classes (any model, only under PREFILL_SLICES, which PR CI doesn't run). The application loops in add_request / execute_model check mm_inputs / mm_positions, which don't exist on the pinned request/scheduler types (only mm_features does), and _apply_mm_cache only exists on vLLM's WorkerWrapperBase, not on TPUWorker — so every branch is dead code. Details inline.

2. No unit tests. ~400 lines of production change across 5 modules with no test coverage — this repo requires unit tests for tpu_inference/** changes. The padding factorization, output-strip logic, _chunked_sdpa routing, and the GridTHW pytree round-trip are all cheaply unit-testable, and several of the inline findings would have been caught by them.

3. No accuracy validation. The PR validation is a throughput benchmark only. The padding/strip changes sit on the numerics path; please include an accuracy run (e.g. the CI accuracy harness or lm-eval on a VL benchmark) for at least one image and one video workload. Note the green CI here doesn't exercise the new seq-pad branch (its inputs are 4-aligned at small TP) nor the PREFILL_SLICES path.

Given the four changes are independent (pixel sharding+padding / SDPA flash redirect / precompile expansion / mm cache), splitting them into separate PRs would make each easy to test and land — the sharding+padding and SDPA pieces are closest to ready.

[QiliangCui]

[blocker] engine_receiver_cache_from_config takes the full VllmConfig, not a ModelConfig: at the pinned vLLM, MultiModalRegistry._get_cache_type immediately does model_config = vllm_config.model_config (vllm multimodal/registry.py), so passing vllm_config.model_config here raises AttributeError as soon as this class is constructed — for any model, multimodal or not. Same issue at line 744. PR CI doesn't catch it because these classes are only instantiated under PREFILL_SLICES disagg mode.

vLLM's own EngineCore.__init__ calls it with the full config: engine_receiver_cache_from_config(vllm_config).

[QiliangCui]

[blocker] This block can never execute at the vLLM commit this repo pins: EngineCoreRequest has neither mm_inputs nor mm_positions — multimodal data lives in mm_features (vllm/v1/engine/__init__.py). So the first branch is always falsy and getattr(request, "mm_positions", None) is always None. (The comment is also inverted — mm_features is the newer schema, mm_inputs the older one.)

If engine-side feature caching is the goal, mirror vLLM's EngineCore.preprocess_add_request:

if self.mm_receiver_cache is not None and request.mm_features:
    request.mm_features = self.mm_receiver_cache.get_and_update_features(request.mm_features)

Also note the sibling DisaggEngineCore now creates the cache in __init__ but its add_request never applies it.

[QiliangCui]

[blocker] This whole block is dead code at the pinned vLLM:

• _apply_mm_cache is defined on vLLM's WorkerWrapperBase; TPUWorker subclasses the abstract WorkerBase, so getattr(self, "_apply_mm_cache", None) is always None here. When the executor path goes through WorkerWrapperBase, its execute_model already applies the cache before delegating (vllm/v1/worker/worker_base.py).
• In the elif, NewRequestData has no mm_inputs and no mm_positions — only mm_features (vllm/v1/core/sched/output.py) — so neither branch can fire.

Net effect: nothing runs. Either rewrite against req.mm_features (as WorkerWrapperBase._apply_mm_cache does) or drop this.

[QiliangCui]

[minor] Two issues: (1) worker_receiver_cache_from_config only returns a cache for mm_processor_cache_type == "shm", and vLLM deliberately refuses to build it without the executor-provided shared_worker_lock (it raises in WorkerWrapperBase.init_worker) — a fresh private multiprocessing.Lock() can't provide cross-process mutual exclusion for the shm object store. (2) This is now the third place a worker receiver cache is created (vLLM's WorkerWrapperBase, disagg_executor.py:82). Plumb the shared lock through the executor or drop this. Also move the function-level multiprocessing / MULTIMODAL_REGISTRY imports to module top.

[QiliangCui]

[major] This search only enforces h >= 2, w >= 2, but the vision tower also needs h and w divisible by spatial_merge_size: e.g. pad_seq=6 yields (1, 2, 3), and Qwen3-VL's rot_pos_ids does reshape(h//merge, merge, w//merge, merge) → ValueError (6 elements into (1,2,1,2)).

Today this is unreachable only because the qwen processors emit h/w as multiples of merge_size, so patch counts are ≡ 0 mod 4 and pad_seq always lands on a multiple of 4 (the h=2, w=2 case) — but nothing asserts that invariant, and anything that violates it crashes deep in the vision tower.

Simpler and safe: pad to a multiple of lcm(tp_size, merge_size**2) and use dummy_grid_thw = [pad_seq // merge_size**2, merge_size, merge_size]; the nested search loop then goes away entirely.

[QiliangCui]

[nit] _chunked_sdpa doesn't chunk — it redirects to encoder_only_flash_attention. Rename to reflect behavior (e.g. _flash_attn_sdpa), and lift 40 * 1024 * 1024 into a named constant. Side note: the GQA bail-out comment says the kernel requires num_heads == num_kv_heads, but the kernel signature accepts distinct num_kv_heads — the fallback is fine, the comment overstates the constraint.

[QiliangCui]

[major] Startup-cost / coverage concerns with this expansion:

• Fan-out per patch bucket goes 1 → 12 compilations for the listed model types, and 18 for everything else. Qwen3-Omni's model_type (qwen3_omni_moe) is not in the tuple, so the currently-supported omni path now gets video_grid_thw and bare-grid_thw dummy calls it never received before, on top of the multiplier — and the VISION_MIN_SHIFT default change (4 → 6) silently drops the 16/32-patch buckets it used to warm.
• The grid args are jit static_argnames, so a precompiled (t, 2^a, 2^b) grid only matches media whose processed grid is exactly that. Real-world resolutions still compile at first request; the warm-up helps pow-2 shapes (like the 256x256x16 benchmark), but "fully warm up the JIT cache" overstates it.

Suggest scoping frame_counts to the video-capable archs that need it and keeping the previous defaults (min_shift=4, single image-grid call) for omni and image-only deployments.

[QiliangCui]

[major] use_dlpack=False was chosen explicitly when this conversion was introduced (#2396); flipping it to True (+ .contiguous()) changes the tensor-conversion path for every multimodal model, and is unrelated to the padding/sharding this PR is about. Was False load-bearing (dtype/layout issues)? If the flip is intentional, add a comment with the rationale and validation; otherwise keep False here and make the dlpack switch its own change.

[QiliangCui]

[minor] self.vllm_config is always set in __init__, so the hasattr chain and this tp_size = 8 fallback are dead — and wrong if they ever fired (a TP=2/TP=4 deployment would pad/shard as if TP=8). Also, divisibility for PartitionSpec("model", None) should be checked against the actual mesh axis size (self.mesh.shape["model"]), not tensor_parallel_size — they aren't guaranteed equal under all sharding configs. The same block is duplicated at ~705; extract one helper.

[QiliangCui]

[minor] GridTHW is imported at the top of this file — use isinstance(grid, GridTHW) instead of the type(grid).__name__ == "GridTHW" string check (4 occurrences), and this list(grid) if ... else grid.tolist() if ... else list(grid) ternary appears twice — worth a small _grid_to_list(grid) helper.

[lk-chen]

these were intentionally removed in https://github.com/vllm-project/tpu-inference/pull/2878/changes

[lk-chen]

if you add a space around this line

tpu-inference/.buildkite/benchmark/cases/daily/daily_qwen3_5_397b_fp8_tpu7x_8.json

Line 176 in 068cd11

"__COMMENT__": "Shardy disabled due to b/517297815#comment9",

, CI will trigger perf. benchmark, current throughput is 13.81 req/s so you have a sense how much your change impacts

[amitkumar307d]

@lk-chen we can definitely leverage the benchmarking capability. One thing to note is the client args are different than what is needed. We need a video input with text and 1k output tokens.

[amitkumar307d]

Split this PR into 3 PRs -
#3442
#3443
#3444