[Bugfix][Gemma 4] Clamp soft-token estimate to max_soft_tokens

[hnt2601]

Summary

For extreme-aspect-ratio images (e.g. 3x900), the prompt-side Gemma4ProcessingInfo._compute_num_soft_tokens returned more soft tokens than the HF Gemma 4 image processor's vision tower actually emits (which is capped at max_soft_tokens). The mismatch crashed _merge_multimodal_embeddings mid-forward and propagated ValueError out of EngineCore.step() with:

Attempted to assign 280 multimodal tokens to 289 placeholders

This PR fixes the bug by clamping the prompt-side estimator to max_soft_tokens so the placeholder count always matches what the encoder will emit, and adds a small typed-exception layer in the Gemma 4 model so a future engine-layer change can classify this failure mode without parsing free-form ValueError messages.

Root cause

Gemma4ProcessingInfo._compute_num_soft_tokens computes target H/W via max(unit, floor(dim * scale / unit) * unit). For very thin or very tall images the max(unit, …) floor clamps one dimension up to unit, while the other scales freely. After dividing by patch_size**2 * pooling_kernel_size**2, the result can exceed max_soft_tokens.

Concrete repro for image_height=3, image_width=900, max_soft_tokens=280, patch_size=14, pooling_kernel_size=2:

• scale ≈ 9.02
• target_h = 28 (floor would give 0; max(unit=28, …) lifts it)
• target_w = 8092
• num_patches = (28 // 14) * (8092 // 14) = 1156
• soft_tokens = 1156 // 4 = 289

The HF Gemma 4 image processor caps its vision-tower output at max_soft_tokens = 280, so the prompt has 289 image placeholders but only 280 embeddings to fill them.

Changes

1. Clamp at the prompt-side estimator (the fix)

Gemma4ProcessingInfo._compute_num_soft_tokens now returns min(num_patches // (pooling_kernel_size**2), max_soft_tokens). Strict tightening of an existing upper bound: any caller that previously got a value ≤ max_soft_tokens is unaffected; any caller that got a value > max_soft_tokens was already producing a hard crash downstream.

2. Typed exception for the count mismatch (defense in depth)

Adds Gemma4MultimodalPlaceholderMismatch(ValueError) and a small _count_multimodal_embedding_rows helper to gemma4_mm.py, and an explicit count check at the top of Gemma4ForConditionalGeneration.embed_input_ids:

expected = int(is_multimodal.sum().item())
actual = _count_multimodal_embedding_rows(multimodal_embeddings)
if actual != expected:
    raise Gemma4MultimodalPlaceholderMismatch(
        actual=actual, expected=expected
    )

The clamp above means this branch should not fire under normal operation. It exists so that any future regression reintroducing a count mismatch raises a typed ValueError subclass with structured actual / expected attributes — instead of the generic ValueError that _merge_multimodal_embeddings raises after a failing inputs_embeds[is_multimodal] = mm_embeds_flat index-put. Subclassing ValueError keeps existing except ValueError handlers working unchanged.

Note: the exception is not yet caught at the engine layer — that hardening is a separate follow-up. Until then, the failure mode for a count mismatch is the same as before, but the offending case is identifiable without string matching on the error message.

Why this is not duplicating an existing PR

Open PRs touching gemma4 / multimodal were checked:

• [MM][Gemma4] Support max_soft_tokens="auto" to avoid wasting the vision budget on small images #40599 — [MM][Gemma4] Support max_soft_tokens="auto". Adds
  an auto-mode that selects max_soft_tokens per image; orthogonal
  to the prompt-side estimator over-shooting whatever cap is set.
• [Bugfix][Gemma4] Fix vision fp16 overflow causing <pad> output #40347 — [Bugfix][Gemma4] Fix vision fp16 overflow causing <pad> output. Different bug (numerical overflow), different layer.
• gemma4: batch multimodal vision processing #40464 — gemma4: batch multimodal vision processing.
  Performance work; does not touch the soft-token arithmetic.
• Fix gemma4 core: quantized inference + gguf loading + fused moe gelu_tanh #40281 — Fix gemma4 core: quantized inference + gguf loading + fused moe gelu_tanh. Unrelated to multimodal.
• Fix RMSNorm hidden_size validation crash for weightless norms #39073, [Gemma4] Allow per-layer attention backend selection for heterogeneou… #38891 — RMSNorm validation, per-layer attention
  backend. Unrelated.

No open PR addresses the placeholder/encoder soft-token count
mismatch.

Test plan

New tests in tests/models/multimodal/processing/test_gemma4.py
run without loading the real Gemma 4 weights (uses MagicMock for
get_hf_config()):

.venv/bin/python -m pytest tests/models/multimodal/processing/test_gemma4.py -v

Test	Pins
test_compute_num_soft_tokens_does_not_exceed_max_soft_tokens (4 cases)	the clamp — production repro (900,3,280), swapped (3,900,280), video-frame budget (900,3,70), high cap (4000,2,1120)
test_gemma4_multimodal_placeholder_mismatch_is_value_error	exception subclasses ValueError (preserves except ValueError call sites)
test_gemma4_multimodal_placeholder_mismatch_carries_counts	actual / expected attributes + message contains both counts
test_gemma4_multimodal_placeholder_mismatch_requires_kwargs	constructor is keyword-only
test_count_multimodal_embedding_rows (4 cases)	helper handles tensor, list, tuple, and empty inputs

All pass locally on Python 3.12 in a uv venv.

Backwards compatibility

• min(…, max_soft_tokens) is a strict tightening of an existing upper bound; no public API surface changes.
• Gemma4MultimodalPlaceholderMismatch is a new public name in vllm.model_executor.models.gemma4_mm. Subclasses ValueError, so any caller using except ValueError continues to catch it.
• _count_multimodal_embedding_rows is module-private (leading underscore).

[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.

[gemini-code-assist]

Code Review

This pull request introduces a typed MultimodalPlaceholderMismatch exception to handle client-input errors in multimodal models, preventing engine crashes and allowing the OpenAI serving layer to return HTTP 400 instead of 500. The changes include updates to EngineCore to catch these errors during the step phase, propagation of error metadata through the output pipeline, and comprehensive regression tests. Feedback was provided to ensure type consistency in EngineCoreOutputs by converting a list of request IDs to a set to match the expected schema.

[hnt2601]

@Mergifyio update

[mergify]

update

✅ Branch has been successfully updated