Merge pull request #4232 from AI-Hypercomputer:hengtaoguo-resize

PiperOrigin-RevId: 936353929

Author: Google-ML-Automation

src/maxtext/configs/types.py

@@ -1884,6 +1884,7 @@ class MultimodalGeneral(BaseModel):
   """General configuration for Multimodal models."""
 
   use_multimodal: bool = Field(False, description="Enable multimodal capabilities.")
+  attention_for_vit: str = Field("dot_product", description="The attention algorithm to use for vision encoder.")
   freeze_vision_encoder_params: bool = Field(True, description="Freeze the parameters of the vision encoder.")
   freeze_audio_encoder_params: bool = Field(True, description="Freeze the parameters of the audio encoder.")
   use_audio: bool = Field(False, description="Enable audio encoder for multimodal models.")

src/maxtext/models/gemma3.py

@@ -449,7 +449,7 @@ def __init__(
         dtype=self.config.dtype_mm,
         weight_dtype=self.config.weight_dtype,
         mesh=self.mesh,
-        attention_kernel="dot_product",
+        attention_kernel=self.config.attention_for_vit,
         inputs_q_shape=(self.config.per_device_batch_size, self.seq_len, self.config.hidden_size_for_vit),
         inputs_kv_shape=(self.config.per_device_batch_size, self.seq_len, self.config.hidden_size_for_vit),
         dropout_rate=0,

src/maxtext/models/gemma4_vision.py

@@ -472,7 +472,7 @@ def __init__(
         head_dim=config.hidden_size_for_vit // config.num_attention_heads_for_vit,
         max_target_length=seq_len,
         mesh=mesh,
-        attention_kernel="dot_product",
+        attention_kernel=config.attention_for_vit,
         inputs_q_shape=dummy_shape,
         inputs_kv_shape=dummy_shape,
         dtype=config.dtype_mm,

src/maxtext/models/llama4.py

@@ -642,7 +642,7 @@ def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs = None):
         num_kv_heads=self.config.num_attention_heads_for_vit,
         head_dim=self.config.hidden_size_for_vit // self.config.num_attention_heads_for_vit,
         max_target_length=(self.config.image_size_for_vit // self.config.patch_size_for_vit) ** 2 + 1,
-        attention_kernel="dot_product",
+        attention_kernel=self.config.attention_for_vit,
         inputs_q_shape=self.hidden_states_shape,
         inputs_kv_shape=self.hidden_states_shape,
         float32_qk_product=self.config.float32_qk_product,

src/maxtext/models/qwen3.py

@@ -1906,7 +1906,7 @@ def __init__(self, config: Config, *, mesh=None, rngs: nnx.Rngs = None):
         num_kv_heads=self.config.num_attention_heads_for_vit,
         head_dim=head_dim,
         max_target_length=self.config.num_position_embeddings_for_vit,
-        attention_kernel="dot_product",
+        attention_kernel=self.config.attention_for_vit,
         inputs_q_shape=(1, 1, self.config.hidden_size_for_vit),
         inputs_kv_shape=(1, 1, self.config.hidden_size_for_vit),
         float32_qk_product=self.config.float32_qk_product,