From 88fd4036a80b0d807c9f37c139ec0e5e50e86b7c Mon Sep 17 00:00:00 2001 From: quic-xiyushi Date: Tue, 2 Jun 2026 02:35:31 -0700 Subject: [PATCH 1/3] Add user_vision_size in VLM's get_specializations for chunked embedding in vLLM v1 (#996) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit vLLM v1 introduces chunked prefill with encoder cache, where the model runner processes requests in scheduled token windows. For each window: - Overlap detection determines which multimodal embeddings are required - Only the relevant sub-tensors are gathered and used From a QEfficient perspective, this is analogous to: - Identifying the image indices for the current chunk, and - Gathering the corresponding vision embeddings accordingly Additionally, v1 eliminates the strict prefill/decode distinction, allowing: - Prefill chunks of one request to be interleaved with decode steps of other requests However, if we reuse the full vision embedding for every prefill chunk, which is what currently QEfficient supports, this: - Breaks the intended vLLM v1 design - Introduces unnecessary overhead due to repeated large set_buffer calls This PR proposes adding user_vision_size to get_specializations for VLM models, with the following scope: ✅ Apply to all current VLM models ❌ Exclude mllama due to it's cross attention ❌ Exclude molmo (not yet supported in vLLM on QAIC, not sure about how the team is going to support it) For models onboarded by QEfficient in the future, it should also follow the same way to enable user_vision_size. This change enables: ✅ In vLLM v1, align vision embedding size with prefill sequence length and enable efficient chunked prefill for multimodal inputs ✅ Better alignment with vLLM v1 scheduling model ✅ Reduced overhead from repeated large buffer updates ✅ Easier support for multi-resolution input and multiple images per request 🚫 Not intended for 1.21 release ✅ Target release: 1.22 This draft PR is currently based on the 1.21.6 branch because multi-resolution and multi-frame support for Qwen2.5-VL / Qwen3-VL is only available in that branch, and these features are prerequisites for this change Once multi-resolution and multi-frame support are available in the QEfficient main branch, this PR will be rebased and migrated to the main branch. --------- Signed-off-by: quic-xiyushi Signed-off-by: sanising Co-authored-by: sanising --- .../models/gemma3/modeling_gemma3.py | 33 ++++++++++++++----- .../models/internvl/modeling_internvl.py | 8 ++++- .../models/llama4/modeling_llama4.py | 8 ++++- .../models/llava/modeling_llava.py | 8 ++++- .../models/llava_next/modeling_llava_next.py | 8 ++++- .../models/mistral3/modeling_mistral3.py | 12 +++++-- .../models/qwen2_5_vl/modeling_qwen2_5_vl.py | 15 +++++---- .../models/qwen3_vl/modeling_qwen3_vl.py | 15 +++++---- .../qwen3_vl_moe/modeling_qwen3_vl_moe.py | 15 +++++---- 9 files changed, 86 insertions(+), 36 deletions(-) diff --git a/QEfficient/transformers/models/gemma3/modeling_gemma3.py b/QEfficient/transformers/models/gemma3/modeling_gemma3.py index a3e9257a73..96bf6cd1a0 100644 --- a/QEfficient/transformers/models/gemma3/modeling_gemma3.py +++ b/QEfficient/transformers/models/gemma3/modeling_gemma3.py @@ -795,7 +795,18 @@ def get_specializations( elif img_size is None: img_size = 896 # FIXME based on gemma3 Image size logger.warning("Setting img_size to be 336, as it was neither passed nor found in vision_config") - mm_tokens_per_image = getattr(self.config, "mm_tokens_per_image", 256) + 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: + _mm_tokens_per_image = getattr(self.config, "mm_tokens_per_image", None) + if _mm_tokens_per_image: + logger.warning_once("mm_tokens_per_image is deprecated and will be removed in the next release.") + vision_size = _mm_tokens_per_image + else: + vision_size = 256 vision = [ { @@ -816,7 +827,7 @@ def get_specializations( "comp_ctx_lengths": comp_ctx_lengths_prefill[i], "sliding_window": self.language_model.config.sliding_window, "img_size": img_size, - "mm_tokens_per_image": mm_tokens_per_image, + "vision_size": vision_size, "vision_batch_size": batch_size, } if continuous_batching: @@ -835,7 +846,7 @@ def get_specializations( "comp_ctx_lengths": comp_ctx_lengths_decode[i], "sliding_window": self.language_model.config.sliding_window, "img_size": img_size, - "mm_tokens_per_image": mm_tokens_per_image, + "vision_size": vision_size, "vision_batch_size": batch_size, } if continuous_batching: @@ -851,7 +862,7 @@ def get_specializations( "ctx_len": ctx_len, "sliding_window": self.language_model.config.sliding_window, "img_size": img_size, - "mm_tokens_per_image": mm_tokens_per_image, + "vision_size": vision_size, "vision_batch_size": batch_size, } if continuous_batching: @@ -867,7 +878,7 @@ def get_specializations( "ctx_len": ctx_len, "sliding_window": self.language_model.config.sliding_window, "img_size": img_size, - "mm_tokens_per_image": mm_tokens_per_image, + "vision_size": vision_size, "vision_batch_size": batch_size, } if continuous_batching: @@ -893,7 +904,7 @@ def get_onnx_dynamic_axes( lang_dynamic_axes = {} lang_dynamic_axes["input_ids"] = {0: "batch_size", 1: "seq_len"} lang_dynamic_axes["position_ids"] = {0: "batch_size", 1: "seq_len"} - lang_dynamic_axes["vision_embeds"] = {0: "vision_batch_size", 1: "mm_tokens_per_image"} + lang_dynamic_axes["vision_embeds"] = {0: "vision_batch_size", 1: "vision_size"} if continuous_batching: lang_dynamic_axes["batch_index"] = {0: "batch_size"} vision_dynamic_axes["pixel_values"] = {0: "batch_size", 2: "img_size", 3: "img_size"} @@ -984,13 +995,19 @@ def get_dummy_inputs( else: img_size = 896 - mm_tokens_per_image = getattr(self.config, "mm_tokens_per_image", 256) + _mm_tokens_per_image = getattr(self.config, "mm_tokens_per_image", None) + if _mm_tokens_per_image: + logger.warning_once("mm_tokens_per_image is deprecated and will be removed in the next release.") + vision_size = _mm_tokens_per_image + else: + vision_size = 256 + # Define shapes inputs_shapes = {} inputs_shapes["input_ids"] = (constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE, prefill_seq_len) inputs_shapes["vision_embeds"] = ( 1, # constants.INTERN_NUM_PATCHES, - mm_tokens_per_image, # constants.INTERN_FEATURE_SIZE, + vision_size, # constants.INTERN_FEATURE_SIZE, self.language_model.config.hidden_size, # 5120 ) inputs_shapes["position_ids"] = ( diff --git a/QEfficient/transformers/models/internvl/modeling_internvl.py b/QEfficient/transformers/models/internvl/modeling_internvl.py index 563c42e256..1496e972b0 100644 --- a/QEfficient/transformers/models/internvl/modeling_internvl.py +++ b/QEfficient/transformers/models/internvl/modeling_internvl.py @@ -134,7 +134,13 @@ def get_specializations( raise NotImplementedError("Image Size other than 448 is not supported for Intern models yet.") per_patch_embed_size = (img_size // self.config.vision_config.patch_size * self.config.downsample_ratio) ** 2 - vision_size = int(batch_size * num_patches * per_patch_embed_size) + 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 = int(batch_size * num_patches * per_patch_embed_size) vision = [ { "batch_size": batch_size, diff --git a/QEfficient/transformers/models/llama4/modeling_llama4.py b/QEfficient/transformers/models/llama4/modeling_llama4.py index 2cf5dbb2e9..3dcd69e79e 100644 --- a/QEfficient/transformers/models/llama4/modeling_llama4.py +++ b/QEfficient/transformers/models/llama4/modeling_llama4.py @@ -1003,7 +1003,13 @@ def get_specializations( * (img_size // self.config.vision_config.patch_size) // downsample_ratio ) - vision_size = num_features_per_tile * max_num_tiles + 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 = num_features_per_tile * max_num_tiles vision = [ { diff --git a/QEfficient/transformers/models/llava/modeling_llava.py b/QEfficient/transformers/models/llava/modeling_llava.py index 88bb5e1027..3cbe77a95b 100644 --- a/QEfficient/transformers/models/llava/modeling_llava.py +++ b/QEfficient/transformers/models/llava/modeling_llava.py @@ -255,7 +255,13 @@ def get_specializations( logger.warning("Setting img_size to be 336, as it was neither passed nor found in vision_config") if img_size != 336 and kv_offload: raise NotImplementedError("Image Size other than 336 is not supported for Llava models yet.") - vision_size = (img_size // self.config.vision_config.patch_size) ** 2 + 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 = (img_size // self.config.vision_config.patch_size) ** 2 vision = [ { "batch_size": batch_size, diff --git a/QEfficient/transformers/models/llava_next/modeling_llava_next.py b/QEfficient/transformers/models/llava_next/modeling_llava_next.py index 342269ce50..9d7881b012 100755 --- a/QEfficient/transformers/models/llava_next/modeling_llava_next.py +++ b/QEfficient/transformers/models/llava_next/modeling_llava_next.py @@ -330,7 +330,13 @@ def get_specializations( logger.warning("Setting img_size to be 384, as it was neither passed nor found in vision_config") if img_size != constants.GRANITEVISION_IMG_SIZE and kv_offload: logger.warning("Image Size other than 384 is not supported for LlavaNext models yet.") - vision_size = constants.GRANITEVISION_FEATURE_SIZE + 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 = constants.GRANITEVISION_FEATURE_SIZE vision = [ { "batch_size": batch_size, diff --git a/QEfficient/transformers/models/mistral3/modeling_mistral3.py b/QEfficient/transformers/models/mistral3/modeling_mistral3.py index 628d1dee2c..d8cbdb17b6 100644 --- a/QEfficient/transformers/models/mistral3/modeling_mistral3.py +++ b/QEfficient/transformers/models/mistral3/modeling_mistral3.py @@ -443,9 +443,15 @@ def get_specializations( ctx_len = ctx_len if ctx_len else constants.INTERN_CTX_LEN patch_size = self.config.vision_config.patch_size kernel_size = self.config.spatial_merge_size - vision_size = ( - ((img_size // patch_size) * (img_size // patch_size)) * (batch_size) // (kernel_size * kernel_size) - ) + 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 = ( + ((img_size // patch_size) * (img_size // patch_size)) * (batch_size) // (kernel_size * kernel_size) + ) vision = [ { diff --git a/QEfficient/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py b/QEfficient/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py index fd4f33e759..ecd3704e6f 100644 --- a/QEfficient/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py +++ b/QEfficient/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py @@ -972,13 +972,14 @@ def get_specializations( "resolution." ) else: - assert vision_size * f < user_vision_size, ( - f"Computed vision_size of {vision_size * f} tokens " - f"(vision_size={vision_size}, num_frames={f}) for image resolution " - f"(width={w}, height={h}) cannot exceed the provided " - f"vision_size={user_vision_size}. Please adjust the image resolution or " - "increase the vision_size." - ) + if vision_size * f > user_vision_size: + logger.warning_once( + f"Computed vision_size of {vision_size * f} tokens " + f"(vision_size={vision_size}, num_frames={f}) for image resolution " + f"(width={w}, height={h}) exceed the provided " + f"vision_size={user_vision_size}. " + f"Vision embedding need to be chunked during prefill." + ) vision.append( { diff --git a/QEfficient/transformers/models/qwen3_vl/modeling_qwen3_vl.py b/QEfficient/transformers/models/qwen3_vl/modeling_qwen3_vl.py index d83dc0a5d5..c022b52652 100644 --- a/QEfficient/transformers/models/qwen3_vl/modeling_qwen3_vl.py +++ b/QEfficient/transformers/models/qwen3_vl/modeling_qwen3_vl.py @@ -1002,13 +1002,14 @@ def get_specializations( "resolution." ) else: - assert vision_size * f < user_vision_size, ( - f"Computed vision_size of {vision_size * f} tokens " - f"(vision_size={vision_size}, num_frames={f}) for image resolution " - f"(width={w}, height={h}) cannot exceed the provided " - f"vision_size={user_vision_size}. Please adjust the image resolution or " - "increase the vision_size." - ) + if vision_size * f > user_vision_size: + logger.warning_once( + f"Computed vision_size of {vision_size * f} tokens " + f"(vision_size={vision_size}, num_frames={f}) for image resolution " + f"(width={w}, height={h}) exceed the provided " + f"vision_size={user_vision_size}. " + f"Vision embedding need to be chunked during prefill." + ) vision.append( { diff --git a/QEfficient/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py b/QEfficient/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py index 446c84df30..ef4d378d87 100644 --- a/QEfficient/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py +++ b/QEfficient/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py @@ -1144,13 +1144,14 @@ def get_specializations( "resolution." ) else: - assert vision_size * f < user_vision_size, ( - f"Computed vision_size of {vision_size * f} tokens " - f"(vision_size={vision_size}, num_frames={f}) for image resolution " - f"(width={w}, height={h}) cannot exceed the provided " - f"vision_size={user_vision_size}. Please adjust the image resolution or " - "increase the vision_size." - ) + if vision_size * f > user_vision_size: + logger.warning_once( + f"Computed vision_size of {vision_size * f} tokens " + f"(vision_size={vision_size}, num_frames={f}) for image resolution " + f"(width={w}, height={h}) exceed the provided " + f"vision_size={user_vision_size}. " + f"Vision embedding need to be chunked during prefill." + ) vision.append( { From 2c6bab1a9116820b2de2b40a590a7c819c532597 Mon Sep 17 00:00:00 2001 From: sanising Date: Tue, 9 Jun 2026 18:09:09 -0700 Subject: [PATCH 2/3] Add user_vision_size in VLM's get_specializations for chunked embedding in vLLM v1 Signed-off-by: sanising --- .../models/qwen3_5/modeling_qwen3_5.py | 15 ++++++++------- .../models/qwen3_5_moe/modeling_qwen3_5_moe.py | 15 ++++++++------- 2 files changed, 16 insertions(+), 14 deletions(-) diff --git a/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py b/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py index 7a66be2958..4e8b8da712 100644 --- a/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -1626,13 +1626,14 @@ def get_specializations( "resolution." ) else: - assert vision_size * f <= user_vision_size, ( - f"Computed vision_size of {vision_size * f} tokens " - f"(vision_size={vision_size}, num_frames={f}) for image resolution " - f"(width={w}, height={h}) cannot exceed the provided " - f"vision_size={user_vision_size}. Please adjust the image resolution or " - "increase the vision_size." - ) + if vision_size * f > user_vision_size: + logger.warning_once( + f"Computed vision_size of {vision_size * f} tokens " + f"(vision_size={vision_size}, num_frames={f}) for image resolution " + f"(width={w}, height={h}) exceed the provided " + f"vision_size={user_vision_size}. " + f"Vision embedding need to be chunked during prefill." + ) vision.append( { diff --git a/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py index b93a08c5a8..7d65186b65 100644 --- a/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -1816,13 +1816,14 @@ def get_specializations( "resolution." ) else: - assert vision_size * f <= user_vision_size, ( - f"Computed vision_size of {vision_size * f} tokens " - f"(vision_size={vision_size}, num_frames={f}) for image resolution " - f"(width={w}, height={h}) cannot exceed the provided " - f"vision_size={user_vision_size}. Please adjust the image resolution or " - "increase the vision_size." - ) + if vision_size * f > user_vision_size: + logger.warning_once( + f"Computed vision_size of {vision_size * f} tokens " + f"(vision_size={vision_size}, num_frames={f}) for image resolution " + f"(width={w}, height={h}) exceed the provided " + f"vision_size={user_vision_size}. " + f"Vision embedding need to be chunked during prefill." + ) vision.append( { From e968d6c885c9488eb89da62faef90754fabc706a Mon Sep 17 00:00:00 2001 From: sanising Date: Tue, 9 Jun 2026 18:11:22 -0700 Subject: [PATCH 3/3] Fix typos Signed-off-by: sanising --- .../transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py | 4 ++-- QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py | 4 ++-- .../transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py | 4 ++-- QEfficient/transformers/models/qwen3_vl/modeling_qwen3_vl.py | 4 ++-- .../transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py | 4 ++-- 5 files changed, 10 insertions(+), 10 deletions(-) diff --git a/QEfficient/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py b/QEfficient/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py index ecd3704e6f..2959a8d0de 100644 --- a/QEfficient/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py +++ b/QEfficient/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py @@ -976,9 +976,9 @@ def get_specializations( logger.warning_once( f"Computed vision_size of {vision_size * f} tokens " f"(vision_size={vision_size}, num_frames={f}) for image resolution " - f"(width={w}, height={h}) exceed the provided " + f"(width={w}, height={h}) exceeds the provided " f"vision_size={user_vision_size}. " - f"Vision embedding need to be chunked during prefill." + f"Vision embedding needs to be chunked during prefill." ) vision.append( diff --git a/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py b/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py index 4e8b8da712..fd67dbb805 100644 --- a/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -1630,9 +1630,9 @@ def get_specializations( logger.warning_once( f"Computed vision_size of {vision_size * f} tokens " f"(vision_size={vision_size}, num_frames={f}) for image resolution " - f"(width={w}, height={h}) exceed the provided " + f"(width={w}, height={h}) exceeds the provided " f"vision_size={user_vision_size}. " - f"Vision embedding need to be chunked during prefill." + f"Vision embedding needs to be chunked during prefill." ) vision.append( diff --git a/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py index 7d65186b65..a429ebc15a 100644 --- a/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/QEfficient/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -1820,9 +1820,9 @@ def get_specializations( logger.warning_once( f"Computed vision_size of {vision_size * f} tokens " f"(vision_size={vision_size}, num_frames={f}) for image resolution " - f"(width={w}, height={h}) exceed the provided " + f"(width={w}, height={h}) exceeds the provided " f"vision_size={user_vision_size}. " - f"Vision embedding need to be chunked during prefill." + f"Vision embedding needs to be chunked during prefill." ) vision.append( diff --git a/QEfficient/transformers/models/qwen3_vl/modeling_qwen3_vl.py b/QEfficient/transformers/models/qwen3_vl/modeling_qwen3_vl.py index c022b52652..08d90dbaf8 100644 --- a/QEfficient/transformers/models/qwen3_vl/modeling_qwen3_vl.py +++ b/QEfficient/transformers/models/qwen3_vl/modeling_qwen3_vl.py @@ -1006,9 +1006,9 @@ def get_specializations( logger.warning_once( f"Computed vision_size of {vision_size * f} tokens " f"(vision_size={vision_size}, num_frames={f}) for image resolution " - f"(width={w}, height={h}) exceed the provided " + f"(width={w}, height={h}) exceeds the provided " f"vision_size={user_vision_size}. " - f"Vision embedding need to be chunked during prefill." + f"Vision embedding needs to be chunked during prefill." ) vision.append( diff --git a/QEfficient/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py b/QEfficient/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py index ef4d378d87..ac6ab7913d 100644 --- a/QEfficient/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py +++ b/QEfficient/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py @@ -1148,9 +1148,9 @@ def get_specializations( logger.warning_once( f"Computed vision_size of {vision_size * f} tokens " f"(vision_size={vision_size}, num_frames={f}) for image resolution " - f"(width={w}, height={h}) exceed the provided " + f"(width={w}, height={h}) exceeds the provided " f"vision_size={user_vision_size}. " - f"Vision embedding need to be chunked during prefill." + f"Vision embedding needs to be chunked during prefill." ) vision.append(