batched encode with padded videos

Author: hengtaoguo

src/maxtext/layers/embeddings.py

@@ -1401,44 +1401,115 @@ def __call__(
       height: int,
       width: int,
       token_mask: Array | None = None,
-      valid_grid: tuple[int, int, int] | None = None,
+      valid_grid: Array | None = None,
   ) -> Array:
     """Apply rotary position embeddings directly to inputs (Q or K tensors).
 
     Args:
-      inputs: Input tensor of shape [B, T*H*W, N, head_dim] (batch, sequence, heads, head_dim)
+      inputs: Input tensor of shape [B, S, N, head_dim] (batch, sequence, heads, head_dim)
              where T=num_frames, H=height, W=width (all static)
       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.
+      valid_grid: Optional actual video grid with shape (batch, 3) (dynamic)
 
     Returns:
-      Rotated inputs with same shape [B, T*H*W, N, head_dim]
+      Rotated inputs with same shape [B, S, 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]
-      sin_emb = sin_emb[None, :, None, :]
-    elif len(inputs.shape) == 3:
-      # For [S, N, H] case
-      cos_emb = cos_emb[:, None, :]  # [S, 1, H]
-      sin_emb = sin_emb[:, None, :]
+    is_3d = inputs.ndim == 3
+    if is_3d:
+      inputs = inputs[None, :, :, :]
+
+    batch_size = inputs.shape[0]
+    max_patches = num_frames * height * width
+
+    if valid_grid is None:
+      valid_grid = jnp.tile(
+          jnp.array([[num_frames, height, width]], dtype=jnp.int32),
+          (batch_size, 1),
+      )
+
+    # Generate coordinates in block-based order
+    row, col = generate_block_coords(valid_grid, max_patches, self.spatial_merge_size)
+
+    # Compute frequencies dynamically using coordinates
+    max_hw = max(height, width)
+    freq_table = self._compute_freq_table(max_hw)  # [max_hw, head_dim//4]
+
+    row_freqs = freq_table[row]  # [B, max_patches, head_dim//4]
+    col_freqs = freq_table[col]  # [B, max_patches, head_dim//4]
+
+    # Concatenate row and column frequencies
+    embeddings = jnp.concatenate([row_freqs, col_freqs], axis=-1)  # [B, max_patches, head_dim//2]
+
+    # Double the embeddings to match head_dim
+    embeddings = jnp.concatenate([embeddings, embeddings], axis=-1)  # [B, max_patches, head_dim]
+
+    cos_emb = jnp.cos(embeddings)
+    sin_emb = jnp.sin(embeddings)
+
+    # Mask out invalid (padded) tokens by setting cos to 1 and sin to 0
+    if token_mask is not None:
+      is_valid = token_mask[:, :, None]  # [B, max_patches, 1]
+      cos_emb = jnp.where(is_valid, cos_emb, 1.0)
+      sin_emb = jnp.where(is_valid, sin_emb, 0.0)
+
+    if self.cast_as_fprop_dtype:
+      cos_emb = cos_emb.astype(self.fprop_dtype)
+      sin_emb = sin_emb.astype(self.fprop_dtype)
+
+    # Reshape for broadcasting to inputs [B, S, N, H]
+    cos_emb = cos_emb[:, :, None, :]  # [B, S, 1, H]
+    sin_emb = sin_emb[:, :, None, :]
 
     rotated = inputs * cos_emb + self._rotate_half(inputs) * sin_emb
 
+    if is_3d:
+      rotated = rotated[0]
+
     return rotated
 
 
+def generate_block_coords(video_grid_thw: Array, max_patches: int, merge_size: int) -> tuple[Array, Array]:
+  """Generate row and col coordinates in block-based order for padded video.
+
+  Args:
+    video_grid_thw: Actual video grid with shape (batch, 3), in Qwen grid units.
+    max_patches: Maximum number of patches (static).
+    merge_size: Spatial merge block size (static).
+
+  Returns:
+    Tuple of (row, col) each of shape [batch, max_patches]
+  """
+  V_T = video_grid_thw[:, 0:1]  # [B, 1]
+  V_H = video_grid_thw[:, 1:2]
+  V_W = video_grid_thw[:, 2:3]
+
+  merged_w = V_W // merge_size
+  V_len = V_T * V_H * V_W
+
+  idx = jnp.arange(max_patches, dtype=jnp.int32)[None, :]  # [1, max_patches]
+  is_valid = idx < V_len  # [B, max_patches]
+
+  stride_lh = V_H * V_W
+  safe_stride_lh = jnp.maximum(stride_lh, 1)
+  s_idx = idx % safe_stride_lh
+
+  intra_col = s_idx % merge_size
+  intra_row = (s_idx // merge_size) % merge_size
+
+  safe_merged_w = jnp.maximum(merged_w, 1)
+  block_elements = merge_size * merge_size
+  block_col = (s_idx // block_elements) % safe_merged_w
+  block_row = s_idx // (safe_merged_w * block_elements)
+
+  row = jnp.where(is_valid, block_row * merge_size + intra_row, 0)
+  col = jnp.where(is_valid, block_col * merge_size + intra_col, 0)
+
+  return row, col
+
+
 def qwen3omnimoe_vision_pos_embed_interpolate_as_linen(
     *,
     num_position_embeddings: int,
@@ -1596,46 +1667,83 @@ def _interpolate_single(self, t: int, h: int, w: int) -> tuple[Array, Array]:
 
     return indices, weights
 
-  def __call__(self, num_frames: int, height: int, width: int) -> Array:
+  def __call__(
+      self,
+      num_frames: int,
+      height: int,
+      width: int,
+      video_grid_thw: Array | None = None,
+      attention_mask: Array | None = None,
+  ) -> Array:
     """Interpolate positional embeddings for given static grid dimensions.
 
     Args:
       num_frames: Number of temporal frames (static)
       height: Height in patches (static)
       width: Width in patches (static)
+      video_grid_thw: Optional actual video grid with shape (batch, 3) (dynamic)
+      attention_mask: Optional attention mask with shape (batch, max_patches) (dynamic)
 
     Returns:
-      Interpolated positional embeddings of shape [num_frames * height * width, hidden_size]
+      Interpolated positional embeddings of shape [batch, num_frames * height * width, hidden_size]
     """
-    # Get interpolation indices and weights
-    indices, weights = self._interpolate_single(num_frames, height, width)  # [4, h*w], [4, h*w]
+    if video_grid_thw is not None:
+      batch_size = video_grid_thw.shape[0]
+    elif attention_mask is not None:
+      batch_size = attention_mask.shape[0]
+    else:
+      batch_size = 1
 
-    # Lookup embeddings for all 4 corners
-    corner_embeds = self.pos_embed.value[indices]  # [4, h*w, hidden_size]
+    max_patches = num_frames * height * width
 
-    # Apply bilinear weights and sum
-    weighted_embeds = corner_embeds * weights[:, :, None]  # [4, h*w, hidden_size]
-    interpolated = jnp.sum(weighted_embeds, axis=0)  # [h*w, hidden_size]
+    if video_grid_thw is None:
+      video_grid_thw = jnp.tile(
+          jnp.array([[num_frames, height, width]], dtype=jnp.int32),
+          (batch_size, 1),
+      )
+    if attention_mask is None:
+      attention_mask = jnp.ones((batch_size, max_patches), dtype=jnp.int32)
 
-    # Repeat for temporal frames
-    if num_frames > 1:
-      interpolated = jnp.tile(interpolated, (num_frames, 1))  # [t*h*w, hidden_size]
+    # Generate coordinates in block-based order
+    row, col = generate_block_coords(video_grid_thw, max_patches, self.spatial_merge_size)
 
-    # Apply spatial merge permutation
-    # Reshape to [t, h, w, hidden_size] then permute for block-based processing
-    merge_size = self.spatial_merge_size
-    merged_h = height // merge_size
-    merged_w = width // merge_size
+    # Normalize coordinates to [0, 1] range based on valid grid sizes
+    V_H = video_grid_thw[:, 1:2]
+    V_W = video_grid_thw[:, 2:3]
+    row_norm = row / jnp.maximum(V_H - 1, 1)
+    col_norm = col / jnp.maximum(V_W - 1, 1)
+
+    # Interpolate from N x N grid
+    N = self.num_grid_per_side
+    table = self.pos_embed.value.reshape(N, N, self.hidden_size)
+
+    y = row_norm * (N - 1)
+    x = col_norm * (N - 1)
+
+    y0 = jnp.floor(y).astype(jnp.int32)
+    x0 = jnp.floor(x).astype(jnp.int32)
+    y1 = jnp.minimum(y0 + 1, N - 1)
+    x1 = jnp.minimum(x0 + 1, N - 1)
+
+    dy = y - y0
+    dx = x - x0
 
-    # Reshape: [t*h*w, hidden_size] -> [t, h, w, hidden_size]
-    interpolated = interpolated.reshape(num_frames, height, width, self.hidden_size)
+    # Gather 4 corners
+    embed_00 = table[y0, x0]
+    embed_01 = table[y0, x1]
+    embed_10 = table[y1, x0]
+    embed_11 = table[y1, x1]
+
+    # Apply bilinear weights
+    dy = dy[:, :, None]
+    dx = dx[:, :, None]
+
+    interpolated = (
+        (1.0 - dy) * (1.0 - dx) * embed_00 + (1.0 - dy) * dx * embed_01 + dy * (1.0 - dx) * embed_10 + dy * dx * embed_11
+    )
 
-    # Permute for spatial merging: [t, merged_h, merge_size, merged_w, merge_size, hidden_size]
-    interpolated = interpolated.reshape(num_frames, merged_h, merge_size, merged_w, merge_size, self.hidden_size)
-    # -> [t, merged_h, merged_w, merge_size, merge_size, hidden_size]
-    interpolated = jnp.transpose(interpolated, (0, 1, 3, 2, 4, 5))
-    # Flatten back to [t*merged_h*merged_w*merge_size*merge_size, hidden_size]
-    interpolated = interpolated.reshape(-1, self.hidden_size)
+    # Mask out invalid (padded) tokens by setting their embeddings to zero
+    interpolated = interpolated * attention_mask[:, :, None]
 
     if self.cast_as_fprop_dtype:
       interpolated = interpolated.astype(self.fprop_dtype)

src/maxtext/models/qwen3.py

@@ -2136,22 +2136,17 @@ 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, :, :]
+    pos = self.pos_embed_interpolate(
+        num_frames,
+        height,
+        width,
+        video_grid_thw=video_grid_thw,
+        attention_mask=attention_mask,
+    )
     x = x + pos
+    valid_grid = video_grid_thw
 
     h_traj = []
     for i in range(self.depth):

tests/unit/qwen3_omni_layers_test.py

@@ -645,6 +645,9 @@ def _copy_weights_and_test(self, num_frames, height, width):
     pos_embed_jax = self.jax_model(num_frames, height, width)
     pos_embed_torch = self.torch_encoder.fast_pos_embed_interpolate(grid_thw_torch)
 
+    if pos_embed_jax.ndim == 3 and pos_embed_torch.ndim == 2:
+      pos_embed_jax = pos_embed_jax[0]
+
     assert_all_close_jax_torch(pos_embed_jax, pos_embed_torch, rtol=1e-2, atol=1e-2)
 
   def test_pos_embed_interpolate_matches_torch(self):
@@ -762,7 +765,7 @@ def test_padded_video_valid_outputs_match_unpadded(self):
 
     encoder = JaxQwen3OmniMoeVisionEncoder(config=config, mesh=self.mesh, rngs=nnx.Rngs(42))
     raw_output, _ = encoder(raw_video)
-    padded_output, _ = encoder(padded_video, video_mask=video_mask, video_grid_thw=raw_grid)
+    padded_output, _ = encoder(padded_video, video_mask=video_mask, video_grid_thw=jnp.array([raw_grid], dtype=jnp.int32))
     patch_mask = np.asarray(video_mask).reshape(1, -1, temporal_patch_size * patch_size * patch_size).max(-1).astype(bool)
 
     np.testing.assert_allclose(
@@ -772,6 +775,99 @@ def test_padded_video_valid_outputs_match_unpadded(self):
         atol=1e-4,
     )
 
+  def test_padded_video_valid_outputs_match_unpadded_batched(self):
+    """Padded tokens do not change valid outputs anywhere in the ViT for batched inputs."""
+    raw_grid_1 = (2, 2, 2)
+    raw_grid_2 = (1, 4, 4)
+    max_grid = (3, 4, 4)
+    config = pyconfig.initialize(
+        ["", base_config_path],
+        model_name="qwen3-omni-30b-a3b",
+        attention="dot_product",
+        attention_type="full",
+        dtype="float32",
+        dtype_mm="float32",
+        weight_dtype="float32",
+        override_model_config=True,
+        attention_for_vit="dot_product",
+        hidden_size_for_vit=16,
+        num_attention_heads_for_vit=2,
+        intermediate_size_for_vit=32,
+        num_hidden_layers_for_vit=1,
+        deepstack_visual_indexes_for_vit=[],
+        video_max_grid_t=max_grid[0],
+        video_max_grid_h=max_grid[1],
+        video_max_grid_w=max_grid[2],
+    )
+    patch_size = config.patch_size_for_vit
+    temporal_patch_size = config.temporal_patch_size_for_vit
+
+    # Generate raw videos
+    raw_shape_1 = (
+        1,
+        config.num_channels_for_vit,
+        raw_grid_1[0] * temporal_patch_size,
+        raw_grid_1[1] * patch_size,
+        raw_grid_1[2] * patch_size,
+    )
+    raw_shape_2 = (
+        1,
+        config.num_channels_for_vit,
+        raw_grid_2[0] * temporal_patch_size,
+        raw_grid_2[1] * patch_size,
+        raw_grid_2[2] * patch_size,
+    )
+    raw_video_1, _ = create_random_jax_torch(*raw_shape_1)
+    raw_video_2, _ = create_random_jax_torch(*raw_shape_2)
+
+    # Pad them individually
+    padded_video_1, _, video_mask_1 = maybe_pad_video_values_to_max_grid(
+        np.asarray(raw_video_1), np.asarray([raw_grid_1]), config
+    )
+    padded_video_2, _, video_mask_2 = maybe_pad_video_values_to_max_grid(
+        np.asarray(raw_video_2), np.asarray([raw_grid_2]), config
+    )
+
+    # Batch them
+    batched_padded_video = jnp.concatenate([padded_video_1, padded_video_2], axis=0)
+    batched_video_mask = jnp.concatenate([video_mask_1, video_mask_2], axis=0)
+    batched_grid_thw = jnp.array([raw_grid_1, raw_grid_2], dtype=jnp.int32)
+
+    encoder = JaxQwen3OmniMoeVisionEncoder(config=config, mesh=self.mesh, rngs=nnx.Rngs(42))
+
+    # Run individual unpadded
+    raw_output_1, _ = encoder(raw_video_1)
+    raw_output_2, _ = encoder(raw_video_2)
+
+    # Run batched padded
+    batched_padded_output, _ = encoder(
+        batched_padded_video,
+        video_mask=batched_video_mask,
+        video_grid_thw=batched_grid_thw,
+    )
+
+    # Verify equivalence for batch item 0
+    patch_mask_1 = (
+        np.asarray(video_mask_1).reshape(1, -1, temporal_patch_size * patch_size * patch_size).max(-1).astype(bool)
+    )
+    np.testing.assert_allclose(
+        np.asarray(raw_output_1).reshape(-1, config.hidden_size_for_vit),
+        np.asarray(batched_padded_output[0])[np.asarray(patch_mask_1)[0]],
+        rtol=1e-4,
+        atol=1e-4,
+    )
+
+    # Verify equivalence for batch item 1
+    patch_mask_2 = (
+        np.asarray(video_mask_2).reshape(1, -1, temporal_patch_size * patch_size * patch_size).max(-1).astype(bool)
+    )
+    np.testing.assert_allclose(
+        np.asarray(raw_output_2).reshape(-1, config.hidden_size_for_vit),
+        np.asarray(batched_padded_output[1])[np.asarray(patch_mask_2)[0]],
+        rtol=1e-4,
+        atol=1e-4,
+    )
+
 
 class TestDeepstackProcess(unittest.TestCase):
   """Tests for deepstack_process.