From 436763b0d372c9362ad739274f827cc1c1292710 Mon Sep 17 00:00:00 2001 From: quic-gthiruko Date: Fri, 3 Jul 2026 08:11:50 +0530 Subject: [PATCH] fix: use vision_size for Gemma4 specializations to support chunked embedding in vLLM v1 (#1135) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary To keep Gemma4 consistent with the other VLMs, use `vision_size` instead of `mm_tokens_per_image` for the number of vision tokens when building specializations. This allows chunked embedding in vLLM v1. Follows the same pattern introduced in PR #1059 (quic-sanising) for gemma3, llava, internvl, qwen2_5_vl, etc. ## Changes - Pop `vision_size` from `compiler_options` (same as all other VLMs) - If user provides `vision_size`, validate it (`vision_size < ctx_len`) and use it as the vision token count in specializations - Otherwise fall back to `mm_tokens_per_image` (default 256 for Gemma4) - Replace `mm_tokens_per_image` with `vision_size` in `build_lang_prefill_spec` and `build_lang_decode_spec` ## Testing Tested with vLLM flow for the following models: - google/gemma-4-E2B-it & google/gemma-4-E4B-it → 2 QPC (Dense) - google/gemma-4-26B-A4B-it → 3 QPC (MoE) Signed-off-by: Gokkulnath Thirukonda Srinivasan --------- Signed-off-by: Gokkulnath Thirukonda Srinivasan --- .../transformers/models/gemma4/modeling_gemma4.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/QEfficient/transformers/models/gemma4/modeling_gemma4.py b/QEfficient/transformers/models/gemma4/modeling_gemma4.py index 37a8852326..62c04f1584 100755 --- a/QEfficient/transformers/models/gemma4/modeling_gemma4.py +++ b/QEfficient/transformers/models/gemma4/modeling_gemma4.py @@ -1179,7 +1179,13 @@ def get_specializations( prefill_seq_len = prefill_seq_len if prefill_seq_len else 32 ctx_len = ctx_len if ctx_len else constants.INTERN_CTX_LEN max_patches = self._get_vision_max_patches() - mm_tokens_per_image = self._get_mm_tokens_per_image() + user_vision_size = compiler_options.pop("vision_size", None) + if user_vision_size: + if user_vision_size >= ctx_len: + raise ValueError("vision_size must be less than ctx_len") + vision_size = user_vision_size + else: + vision_size = self._get_mm_tokens_per_image() vision = [{"batch_size": batch_size, "max_patches": max_patches}] @@ -1190,7 +1196,7 @@ def build_lang_prefill_spec(comp_ctx_lengths: Optional[int] = None): "ctx_len": ctx_len, "sliding_window": self.model.language_model.config.sliding_window, "vision_batch_size": batch_size, - "vision_tokens": mm_tokens_per_image, + "vision_size": vision_size, } if comp_ctx_lengths is not None: spec["comp_ctx_lengths"] = comp_ctx_lengths @@ -1209,7 +1215,7 @@ def build_lang_decode_spec(comp_ctx_lengths: Optional[int] = None): "ctx_len": ctx_len, "sliding_window": self.model.language_model.config.sliding_window, "vision_batch_size": batch_size, - "vision_tokens": mm_tokens_per_image, + "vision_size": vision_size, } if comp_ctx_lengths is not None: spec["comp_ctx_lengths"] = comp_ctx_lengths @@ -1237,7 +1243,7 @@ def get_onnx_dynamic_axes( } lang_dynamic_axes = { "input_ids": {0: "batch_size", 1: "seq_len"}, - "vision_embeds": {0: "vision_batch_size", 1: "vision_tokens"}, + "vision_embeds": {0: "vision_batch_size", 1: "vision_size"}, "position_ids": {0: "batch_size", 1: "seq_len"}, "mm_token_type_ids": {0: "batch_size", 1: "seq_len"}, }