[CI] Fix "test_awq_load[gemma4-moe-*]" failure

[haosdent]

Purpose

Fix nightly Quantized Models Test #67339 failures of test_awq_load[gemma4-moe-*] introduced by #43169.

That PR sized the batched vision-encoder chunk by 0.05 * total_memory with a cost model that counted only the encoder residual stream. On a 22 GiB L4 with a 26B AWQ model loaded (~3 GiB free), the heuristic admitted chunk ≈ 53 and OOMed inside Gemma4VisionPatchEmbedder._position_embeddings allocating a 2.88 GiB F.one_hot(num_classes=position_embedding_size) int64 buffer, which is the actual dominant transient and ~4x larger than the old cost term.

This PR sizes the chunk by currently-free GPU memory (min(free // 2, total // 10)) with F.one_hot as the dominant per-patch cost, hoists torch.cuda.mem_get_info() out of the per-bucket loop, and extracts the math into a pure static _encoder_chunk for GPU-free unit testing. Batching speedup is preserved on roomy GPUs (chunk = 92 on 80 GiB A100 with 60 GiB free vs. 17 on the failing L4).

Test Plan

Added tests/models/multimodal/test_gemma4_mm.py with 4 unit tests against _encoder_chunk (no GPU required): tight-budget allocation fits in free memory, roomy GPU keeps batching, zero-patches safe, zero-free falls back to 1.

Test Result

Unit tests (no GPU):

$ pytest -v tests/models/multimodal/test_gemma4_mm.py
4 passed in 0.88s

Integration tests (the original failing parametrizations):

Platform	Test	Result
A30 (24 GiB)	test_awq_load[gemma4-moe-standard-awq-dot-suffix]	passed in 237s
A30 (24 GiB)	test_awq_load[gemma4-moe-compressed-tensors-underscore-suffix]	passed in 85s
A100 (80 GiB, CUDA_VISIBLE_DEVICES=0)	both gemma4-moe-* parametrizations	2 passed in 416s

profile_run reaches the Encoder cache will be initialized with a budget of 8192 tokens, and profiled with 3 video items of the maximum feature size. log line, the exact point that OOMed in CI, and proceeds without OutOfMemoryError on both SKUs.

[gemini-code-assist]

Code Review

This pull request refactors the Gemma4 vision encoder's batching logic to prevent out-of-memory (OOM) errors by dynamically sizing chunks based on currently free GPU memory. It replaces the static budget calculation with a more accurate model of transient memory costs, specifically targeting the F.one_hot allocation. Feedback highlights that the new _encoder_chunk method lacks a safety check for zero-cost scenarios, which could lead to a division by zero. Additionally, the direct use of torch.cuda.mem_get_info() in both image and video processing paths breaks portability for non-CUDA backends and should be abstracted through vLLM's platform layer.

[haosdent]

@lucianommartins @Isotr0py, may you help review this? This is the follow-up to #43169. Thank you in advance!

[mergify]

Hi @haosdent, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

Hi @haosdent, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

Hi @haosdent, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[Isotr0py]

The pre-commit failures fixed in #43378

[haosdent]

Likely it also fixes Multi-Modal Models (Extended Generation 3)

Let me monitor if https://buildkite.com/vllm/ci/builds/67561/list?sid=019e4d84-f66a-472b-9bdb-73d6a1b9193b&tab=output pass
Update: Pass