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..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 @@ -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}) exceeds the provided " + f"vision_size={user_vision_size}. " + 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 7a66be2958..fd67dbb805 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}) exceeds the provided " + f"vision_size={user_vision_size}. " + 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 b93a08c5a8..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 @@ -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}) exceeds the provided " + f"vision_size={user_vision_size}. " + 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 d83dc0a5d5..08d90dbaf8 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}) exceeds the provided " + f"vision_size={user_vision_size}. " + 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 446c84df30..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 @@ -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}) exceeds the provided " + f"vision_size={user_vision_size}. " + f"Vision embedding needs to be chunked during prefill." + ) vision.append( {