Batched Video Encoding and Token Mask Propagation for Qwen3*

[hengtaoguo]

Description

This PR enables JIT-compatible, batched video encoding and correct token mask propagation for Qwen3* under mixed-length batches (vision_encoder + merge_mm_embeddings).

How we handled batched padding & recovery:

1. Coordinate-Based Embeddings: Refactored PosEmbedInterpolate and RotaryEmbedding to dynamically generate Row/Col coordinates from a dynamic video_grid_thw JAX array. This allows processing mixed-length padded videos in a batch without triggering JIT recompilations.
2. Decoder Merge Recovery: Post-projector, patch-level masks are downsampled to token-level masks. In merge_mm_embeddings, we sort the projected video tokens (putting valid ones first) and use a JAX-safe index-0 overwrite trick to insert only valid video embeddings into the text sequence while discarding padded ones.
3. Inference Compatibility: Restored support for static tuple-based grids to avoid breaking the existing MaxEngine decode pipeline.

Tests

• Added test_padded_video_projected_outputs_match_unpadded_batched in qwen3_omni_layers_test.py. It verifies that the entire pipeline—from raw video inputs, through the encoder and projector, downsampling the masks, to merging into the text sequence—produces outputs for batched padded items that are identical to independent unpadded runs.

python -m pytest tests/unit/qwen3_omni_layers_test.py::TestQwen3OmniMoeVisionEncoderEndToEnd::test_padded_video_projected_outputs_match_unpadded_batched -vv -s

• Test Case Configurations:
  • Max Grid: [3, 4, 4] (T, H, W in patches).
  • Video 1 Grid: [2, 2, 2] (valid). Video 2 Grid: [1, 4, 4] (valid).
  • Patch size: 14, temporal patch size: 2. Channels: 3.
  • ViT Dim: 16. Projector Dim: 24. Text Seq Len: 20.
• Input & Padding Stage:
  • Videos are padded to max grid and batched.
  • batched_padded_video: [2, 3, 6, 56, 56] (B, C, T_px, H_px, W_px).
  • batched_video_mask: [2, 1, 6, 56, 56] (pixel mask).
  • batched_grid_thw: [[2, 2, 2], [1, 4, 4]] (dynamic grids).
• Vision Encoder (ViT):
  • Pixels embedded to patches: [2, 48, 16] (B, Max_Patches = $3 \times 4 \times 4 = 48$, ViT_Dim=16).
• Projector (Patch Merger):
  • 2x2 spatial blocks merged: [2, 12, 24] (B, Max_Blocks = $48/4 = 12$, Projector_Dim=24).
• Mask Downsampling:
  • Downsampled token masks: [2, 12].
  • Video 1: [1, 1, 0, ..., 0] (2 valid). Video 2: [1, 1, 1, 1, 0, ..., 0] (4 valid).
• Decoder Merging Stage:
  • batched_text_embeddings: [2, 20, 24] (B, Seq_Len=20, Dim=24).
  • batched_placeholder_mask: [2, 20] (Video 1 has 2 placeholders at 5:7; Video 2 has 4 at 5:9).
  • batched_merged: [2, 20, 24]. Valid tokens inserted, padded tokens discarded at index 0.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[codecov]

Codecov Report

❌ Patch coverage is 89.01099% with 10 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/embeddings.py	88.76%	5 Missing and 5 partials ⚠️

📢 Thoughts on this report? Let us know!

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[aireenmei]

Thanks for the work! I like the test