[jzakrzew]

Overview

Fix partial prefill of multimodal pooling requests when chunked prefill is disabled. Bidirectional MEAN-pooling models can crash under async encoder-cache pressure or an explicit long-prefill threshold.

Claims

• Defer the whole request when encoder limits would force partial prefill, releasing acquired cache references and restoring the compute budget.
• Ignore long_prefill_token_threshold when chunked prefill is disabled.
• Add scheduler and real-model regression coverage for both triggers.

Validation

Reproduced with real weights for nvidia/llama-nemotron-embed-vl-1b-v2, using Model Runner V2. Both cases queue two distinct images with chunked prefill disabled:

• Encoder-cache pressure: async scheduling, a 384-token budget/cache, and 256 tokens per image. The second request fits the next step's token budget while the first image still occupies the encoder cache.
• Prefill threshold: synchronous scheduling with long_prefill_token_threshold=128, which incorrectly splits the queued prompts.

With the unchanged upstream scheduler (c7a7607519), both cases raise RuntimeError: partial prefill is not supported with MEAN pooling. With this fix, both complete and every batched embedding matches its individually processed reference with cosine similarity ≥0.999.

Focused regressions, 2 passed each:

python -m pytest \
  tests/v1/core/test_scheduler.py \
  -k encoder_cache_pressure_respects_disabled_chunked_prefill -v

python -m pytest \
  tests/models/multimodal/pooling/test_llama_nemotron_vl.py \
  -k test_image_embeddings_without_chunked_prefill -v

Details

The scheduler checks the token budget before scheduling encoder inputs. Encoder-cache exhaustion can subsequently shorten the request to the text before an image, violating disabled chunked prefill. Async scheduling exposes this when an earlier batch still holds encoder-cache references. An explicit long-prefill threshold provides a separate path to partial prefill under synchronous scheduling.

The encoder check lives in _try_schedule_encoder_inputs. The threshold is applied only when chunked prefill is enabled.

Open-PR searches found no equivalent fix. Related work: #63033 adds a PaliGemma-specific whole-prompt guard; #61733 rejects permanently oversized prompts at admission. Neither covers these two triggers.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[njhill]

Thanks @jzakrzew

[njhill]

/ci run

[github-actions]

✅ Triggered Buildkite CI #99993 for commit 2a9aa6d34873.

ℹ️ This PR is 151 commits behind upstream main. Merged upstream main into this branch as 2a9aa6d34873 before starting CI. Pull the branch before pushing more changes.