Merge pull request #4425 from AI-Hypercomputer:hengtaoguo-pad

PiperOrigin-RevId: 947231694

Author: Google-ML-Automation

src/maxtext/layers/attentions.py

@@ -929,8 +929,12 @@ def apply_rotary_embedding(
       num_frames = rope_kwargs.get("num_frames")
       height = rope_kwargs.get("height")
       width = rope_kwargs.get("width")
+      token_mask = rope_kwargs.get("token_mask")
+      valid_grid = rope_kwargs.get("valid_grid")
       # Type cast required: Omni rotary embedding uses different __call__ parameters than other embeddings.
-      return cast(Qwen3OmniMoeVisionRotaryEmbedding, self.rotary_embedding)(inputs, num_frames, height, width)
+      return cast(Qwen3OmniMoeVisionRotaryEmbedding, self.rotary_embedding)(
+          inputs, num_frames, height, width, token_mask=token_mask, valid_grid=valid_grid
+      )
     else:
       return self.rotary_embedding(inputs, inputs_positions)
 

src/maxtext/layers/embeddings.py

@@ -1394,7 +1394,15 @@ def _rotate_half(self, x: Array) -> Array:
     x2 = x[..., x.shape[-1] // 2 :]
     return jnp.concatenate([-x2, x1], axis=-1)
 
-  def __call__(self, inputs: Array, num_frames: int, height: int, width: int) -> Array:
+  def __call__(
+      self,
+      inputs: Array,
+      num_frames: int,
+      height: int,
+      width: int,
+      token_mask: Array | None = None,
+      valid_grid: tuple[int, int, int] | None = None,
+  ) -> Array:
     """Apply rotary position embeddings directly to inputs (Q or K tensors).
 
     Args:
@@ -1403,11 +1411,20 @@ def __call__(self, inputs: Array, num_frames: int, height: int, width: int) -> A
       num_frames: Number of temporal frames (static)
       height: Height in patches (static)
       width: Width in patches (static)
+      token_mask: Optional mask identifying valid tokens in the padded sequence.
+      valid_grid: Optional unpadded `(frames, height, width)` grid used to compute
+        RoPE for valid tokens.
 
     Returns:
       Rotated inputs with same shape [B, T*H*W, N, head_dim]
     """
     cos_emb, sin_emb = self.compute_cos_sin(num_frames, height, width)
+    if token_mask is not None and valid_grid is not None:
+      valid_cos, valid_sin = self.compute_cos_sin(*valid_grid)
+      valid_len = math.prod(valid_grid)
+      valid_indices = jnp.nonzero(token_mask[0], size=valid_len)[0]
+      cos_emb = jnp.ones_like(cos_emb).at[valid_indices].set(valid_cos)
+      sin_emb = jnp.zeros_like(sin_emb).at[valid_indices].set(valid_sin)
 
     if len(inputs.shape) == 4:
       cos_emb = cos_emb[None, :, None, :]  # [1, S, 1, H]

src/maxtext/layers/encoders.py

@@ -93,10 +93,15 @@ def _setup_vision_encoder_layers(self):
     else:
       raise ValueError(f"No VisionEncoder implemented for {self.config.model_name} yet")
 
-  def __call__(self, input_images, deterministic=False):
+  def __call__(self, input_images, input_masks=None, video_grid_thw=None, deterministic=False):
     # vision encoder output, frozen params in many cases
     encoder = getattr(self, self.encoder_name)
-    encoder_output = encoder(input_images, deterministic=deterministic)
+    if self.config.model_name.startswith("qwen3") and input_masks is not None:
+      encoder_output = encoder(
+          input_images, video_mask=input_masks, video_grid_thw=video_grid_thw, deterministic=deterministic
+      )
+    else:
+      encoder_output = encoder(input_images, deterministic=deterministic)
     deep_feats = None
     if isinstance(encoder_output, tuple):
       embeddings = encoder_output[0]

src/maxtext/models/qwen3.py

@@ -1885,8 +1885,8 @@ def __call__(self, hidden_states: Array, video_mask: Array | None = None) -> tup
 
     attention_mask = None
     if video_mask is not None:
-      patch_mask = video_mask[:, 0, :: self.temporal_patch_size, :: self.patch_size, :: self.patch_size]
-      attention_mask = patch_mask.reshape(video_mask.shape[0], -1).astype(jnp.int32)
+      mask_patch_elements = self.temporal_patch_size * self.patch_size * self.patch_size
+      attention_mask = video_mask.reshape(video_mask.shape[0], -1, mask_patch_elements).max(axis=-1).astype(jnp.int32)
 
     return hidden_states, attention_mask
 
@@ -1941,6 +1941,8 @@ def __call__(
       num_frames: int,
       height: int,
       width: int,
+      attention_mask: Array | None = None,
+      valid_grid: tuple[int, int, int] | None = None,
       deterministic: bool = True,
   ) -> Array:
     """
@@ -1949,6 +1951,8 @@ def __call__(
         num_frames: Number of temporal frames (static)
         height: Height in patches (static)
         width: Width in patches (static)
+        attention_mask: Optional mask identifying valid tokens in the padded sequence.
+        valid_grid: Optional unpadded `(frames, height, width)` grid used for vision RoPE.
         deterministic: Whether to use deterministic mode (disable dropout)
 
     Returns:
@@ -1959,11 +1963,14 @@ def __call__(
         "num_frames": num_frames,
         "height": height,
         "width": width,
+        "token_mask": attention_mask,
+        "valid_grid": valid_grid,
     }
     output, _ = self.attn(
         inputs_q=hidden_states,
         inputs_kv=hidden_states,
         deterministic=deterministic,
+        decoder_segment_ids=attention_mask,
         rope_kwargs=rope_kwargs,
     )
 
@@ -2016,6 +2023,8 @@ def __call__(
       num_frames: int,
       height: int,
       width: int,
+      attention_mask: Array | None = None,
+      valid_grid: tuple[int, int, int] | None = None,
   ) -> Array:
     """
     Args:
@@ -2027,7 +2036,14 @@ def __call__(
     Returns:
         Output tensor of shape (batch, T*H*W, hidden_size)
     """
-    x = x + self.attn(self.ln1(x), num_frames=num_frames, height=height, width=width)
+    x = x + self.attn(
+        self.ln1(x),
+        num_frames=num_frames,
+        height=height,
+        width=width,
+        attention_mask=attention_mask,
+        valid_grid=valid_grid,
+    )
     y = self.ln2(x)
     y = self.mlp(y)
     y = jax.nn.gelu(y)
@@ -2088,6 +2104,8 @@ def __init__(self, config: Config, *, mesh=None, rngs: nnx.Rngs = None):
   def __call__(
       self,
       hidden_states: Array,
+      video_mask: Array | None = None,
+      video_grid_thw: Array | tuple[int, int, int] | None = None,
       deterministic: bool = True,
   ):
     """
@@ -2104,6 +2122,12 @@ def __call__(
     num_frames = num_frames // self.config.temporal_patch_size_for_vit
     height = height // self.config.patch_size_for_vit
     width = width // self.config.patch_size_for_vit
+    attention_mask = None
+    if video_mask is not None:
+      mask_patch_elements = (
+          self.config.temporal_patch_size_for_vit * self.config.patch_size_for_vit * self.config.patch_size_for_vit
+      )
+      attention_mask = video_mask.reshape(batch_size, -1, mask_patch_elements).max(axis=-1).astype(jnp.int32)
     hidden_states = hidden_states.reshape(
         -1,
         self.config.num_channels_for_vit,
@@ -2114,7 +2138,19 @@ def __call__(
 
     x, _ = self.patch_embed(hidden_states)
     x = x.reshape(batch_size, -1, self.config.hidden_size_for_vit)
+    valid_grid = None
+    if attention_mask is not None and video_grid_thw is None:
+      raise ValueError("video_grid_thw is required when video_mask is provided.")
+    if attention_mask is not None and batch_size != 1:
+      raise ValueError("Padded Qwen3-Omni vision encoding currently supports batch size one.")
+    if video_grid_thw is not None:
+      grid = video_grid_thw[0] if getattr(video_grid_thw, "ndim", 1) == 2 else video_grid_thw
+      valid_grid = tuple(int(dim) for dim in grid)
     pos = self.pos_embed_interpolate(num_frames, height, width)
+    if attention_mask is not None and valid_grid is not None:
+      valid_pos = self.pos_embed_interpolate(*valid_grid)
+      valid_indices = jnp.nonzero(attention_mask[0], size=math.prod(valid_grid))[0]
+      pos = jnp.zeros_like(pos).at[valid_indices].set(valid_pos)
 
     pos = pos[jnp.newaxis, :, :]
     x = x + pos
@@ -2123,7 +2159,14 @@ def __call__(
     for i in range(self.depth):
       block_name = f"blocks_{i}"
       blk = getattr(self, block_name)
-      x = blk(x, num_frames=num_frames, height=height, width=width)
+      x = blk(
+          x,
+          num_frames=num_frames,
+          height=height,
+          width=width,
+          attention_mask=attention_mask,
+          valid_grid=valid_grid,
+      )
       h_traj.append(x)
 
     deep_feats = []

src/maxtext/multimodal/processor_qwen3_omni.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Qwen3-Omni-specific preprocessing utilities for multimodal features. 
+"""Qwen3-Omni-specific preprocessing utilities for multimodal features.
 
 Original implementation from HuggingFace: Qwen/Qwen3-Omni-30B-A3B-Instruct.
 """
@@ -177,9 +177,6 @@ def maybe_pad_video_values_to_max_grid(
 
   temporal_patch_size = config.temporal_patch_size_for_vit
   patch_size = config.patch_size_for_vit
-  valid_t_px = actual_t * temporal_patch_size
-  valid_h_px = actual_h * patch_size
-  valid_w_px = actual_w * patch_size
   max_t_px = max_t * temporal_patch_size
   max_h_px = max_h * patch_size
   max_w_px = max_w * patch_size
@@ -188,12 +185,15 @@ def maybe_pad_video_values_to_max_grid(
       (video_values.shape[0], video_values.shape[1], max_t_px, max_h_px, max_w_px),
       dtype=video_values.dtype,
   )
-  padded_video_values[:, :, :valid_t_px, :valid_h_px, :valid_w_px] = video_values[
-      :, :, :valid_t_px, :valid_h_px, :valid_w_px
-  ]
+  patch_elements = video_values.shape[1] * temporal_patch_size * patch_size * patch_size
+  valid_patches = actual_t * actual_h * actual_w
+  padded_video_values.reshape((video_values.shape[0], -1, patch_elements))[:, :valid_patches] = video_values.reshape(
+      (video_values.shape[0], valid_patches, patch_elements)
+  )
 
   video_mask = np.zeros((video_values.shape[0], 1, max_t_px, max_h_px, max_w_px), dtype=np.int32)
-  video_mask[:, :, :valid_t_px, :valid_h_px, :valid_w_px] = 1
+  mask_patch_elements = temporal_patch_size * patch_size * patch_size
+  video_mask.reshape((video_values.shape[0], -1, mask_patch_elements))[:, :valid_patches] = 1
 
   return padded_video_values, video_grid_thw, video_mask
 
@@ -1136,7 +1136,9 @@ def get_rope_index(
       # Process modality-specific content
       # Audio Only
       if min_ed == ed_audio_start:
-        audio_len = _get_feat_extract_output_lengths(audio_lengths[audio_idx]).item()  # pyrefly: ignore[unsupported-operation]
+        audio_len = _get_feat_extract_output_lengths(
+            audio_lengths[audio_idx]
+        ).item()  # pyrefly: ignore[unsupported-operation]
         audio_pos = np.arange(audio_len).reshape(1, -1).repeat(3, axis=0) + st_idx
         llm_pos_ids_list.append(audio_pos)
 
@@ -1151,10 +1153,14 @@ def get_rope_index(
         grid_ws = image_grid_thw[:, 2]  # pyrefly: ignore[unsupported-operation]
         t_index = np.arange(grid_t, dtype=np.float32) * 1 * position_id_per_seconds
 
-        image_pos = get_llm_pos_ids_for_vision(st_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws)  # pyrefly: ignore[bad-argument-type]
+        image_pos = get_llm_pos_ids_for_vision(
+            st_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+        )  # pyrefly: ignore[bad-argument-type]
         llm_pos_ids_list.append(image_pos)
 
-        image_len = int(np.prod(image_grid_thw[image_idx]).item() // (spatial_merge_size**2))  # pyrefly: ignore[unsupported-operation]
+        image_len = int(
+            np.prod(image_grid_thw[image_idx]).item() // (spatial_merge_size**2)
+        )  # pyrefly: ignore[unsupported-operation]
         st += int(text_len + bos_len + image_len + eos_len)
         image_idx += 1
         remain_images -= 1
@@ -1164,27 +1170,39 @@ def get_rope_index(
         grid_t = video_grid_thw[video_idx, 0].item()  # pyrefly: ignore[unsupported-operation]
         grid_hs = video_grid_thw[:, 1]  # pyrefly: ignore[unsupported-operation]
         grid_ws = video_grid_thw[:, 2]  # pyrefly: ignore[unsupported-operation]
-        t_index = np.arange(grid_t, dtype=np.float32) * second_per_grids[video_idx].item() * position_id_per_seconds  # pyrefly: ignore[unsupported-operation]
+        t_index = (
+            np.arange(grid_t, dtype=np.float32) * second_per_grids[video_idx].item() * position_id_per_seconds
+        )  # pyrefly: ignore[unsupported-operation]
 
-        video_pos = get_llm_pos_ids_for_vision(st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws)  # pyrefly: ignore[bad-argument-type]
+        video_pos = get_llm_pos_ids_for_vision(
+            st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+        )  # pyrefly: ignore[bad-argument-type]
         llm_pos_ids_list.append(video_pos)
 
-        video_len = int(np.prod(video_grid_thw[video_idx]).item() // (spatial_merge_size**2))  # pyrefly: ignore[unsupported-operation]
+        video_len = int(
+            np.prod(video_grid_thw[video_idx]).item() // (spatial_merge_size**2)
+        )  # pyrefly: ignore[unsupported-operation]
         st += int(text_len + bos_len + video_len + eos_len)
         video_idx += 1
         remain_videos -= 1
 
       # Audio in Video (interleaved)
       elif min_ed == ed_vision_start and ed_vision_start + 1 == ed_audio_start:
-        audio_len = _get_feat_extract_output_lengths(audio_lengths[audio_idx]).item()  # pyrefly: ignore[unsupported-operation]
+        audio_len = _get_feat_extract_output_lengths(
+            audio_lengths[audio_idx]
+        ).item()  # pyrefly: ignore[unsupported-operation]
         audio_llm_pos_ids = np.arange(audio_len).reshape(1, -1).repeat(3, axis=0) + st_idx
 
         grid_t = video_grid_thw[video_idx, 0].item()  # pyrefly: ignore[unsupported-operation]
         grid_hs = video_grid_thw[:, 1]  # pyrefly: ignore[unsupported-operation]
         grid_ws = video_grid_thw[:, 2]  # pyrefly: ignore[unsupported-operation]
-        t_index = np.arange(grid_t, dtype=np.float32) * second_per_grids[video_idx].item() * position_id_per_seconds  # pyrefly: ignore[unsupported-operation]
+        t_index = (
+            np.arange(grid_t, dtype=np.float32) * second_per_grids[video_idx].item() * position_id_per_seconds
+        )  # pyrefly: ignore[unsupported-operation]
 
-        video_llm_pos_ids = get_llm_pos_ids_for_vision(st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws)  # pyrefly: ignore[bad-argument-type]
+        video_llm_pos_ids = get_llm_pos_ids_for_vision(
+            st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
+        )  # pyrefly: ignore[bad-argument-type]
 
         # Interleave audio and video based on temporal ordering
         video_data_index = 0
@@ -1203,7 +1221,9 @@ def get_rope_index(
         if audio_data_index < audio_llm_pos_ids.shape[1]:
           llm_pos_ids_list.append(audio_llm_pos_ids[:, audio_data_index:])
 
-        video_len = int(np.prod(video_grid_thw[video_idx]).item() // (spatial_merge_size**2))  # pyrefly: ignore[unsupported-operation]
+        video_len = int(
+            np.prod(video_grid_thw[video_idx]).item() // (spatial_merge_size**2)
+        )  # pyrefly: ignore[unsupported-operation]
         st += int(text_len + bos_len + audio_len + video_len + eos_len)
 
         audio_idx += 1