Decode with padded video masks for Qwen3 models

[hengtaoguo]

Description

E2E decode now works for padded/resized video input:

• Pass video_masks, video_grid_thw through MaxEngine and the vision encoder.
• Convert pixel masks into post-projector token masks.
• Merge only valid video embeddings into decoder placeholders.
• Use video placeholder masks for deepstack features.
• Preserve existing tile-mask behavior for other vision models.
• Add focused mask reduction and embedding merge tests.

Tests

Original video grid [5, 30, 28], tensor size [10, 480, 448]:

python -m maxtext.inference.decode maxtext/configs/base.yml tokenizer_type=huggingface model_name=qwen3-vl-2b tokenizer_path=Qwen/Qwen3-VL-2B-Instruct load_parameters_path=gs://hengtaoguo-maxtext-logs/checkpoints/qwen3-vl-2b/unscanned/2026-07-01-11-55/0/items per_device_batch_size=1 run_name=ht_test scan_layers=false use_multimodal=true prompt=\'What\ is\ the\ classification\ of\ the\ single\ exhibit\ in\ this\ video\?\' video_path=\'tests/assets/test_video.mp4\' max_prefill_predict_length=1240 max_target_length=1280 ici_tensor_parallelism=4 override_model_config=true attention=\'dot_product\' hf_access_token=xxx

# Logs, video tensor shape [1, 3, 10, 480, 448]
DEBUG: vision encoder video input shape=(Array(1, dtype=int32), Array(3, dtype=int32), Array(10, dtype=int32), Array(480, dtype=int32),Array(448, dtype=int32))
` -> `The exhibit is a dinosaur, specifically a large, multi-headed creature. It is likely a model or replica of a prehistoric animal, possibly a type of theropod dinosaur, given its size and features`

Padded video [5, 30, 28] -> [32, 32, 32], tensor size [64, 512, 512]. video_max_grid_t/h/w=32, attention_for_vit="autoselected"

python -m maxtext.inference.decode maxtext/configs/base.yml tokenizer_type=huggingface model_name=qwen3-vl-2b tokenizer_path=Qwen/Qwen3-VL-2B-Instruct load_parameters_path=gs://hengtaoguo-maxtext-logs/checkpoints/qwen3-vl-2b/unscanned/2026-07-01-11-55/0/items per_device_batch_size=1 run_name=ht_test scan_layers=false use_multimodal=true prompt=\'What\ is\ the\ classification\ of\ the\ single\ exhibit\ in\ this\ video\?\' video_path=\'tests/assets/test_video.mp4\' max_prefill_predict_length=1240 max_target_length=1280 ici_tensor_parallelism=4 override_model_config=true video_max_grid_t=32 video_max_grid_h=32 video_max_grid_w=32 attention=\'dot_product\' attention_for_vit=\'autoselected\' hf_access_token=xxx

# Logs, video tensor shape [1, 3, 64, 512, 512], fixed canvas
DEBUG: vision encoder video input shape=(Array(1, dtype=int32), Array(3, dtype=int32), Array(64, dtype=int32), Array(512, dtype=int32),Array(512, dtype=int32))
` -> `The exhibit is a dinosaur, specifically a large, multi-headed creature. It is likely a model or replica of a prehistoric animal, possibly a type of theropod dinosaur, and is displayed in a`

The output string is basically identical, and drifted after ~35 tokens generation. This could be caused by small numerical differences in ViT attention kernel using dot_product vs autoselected.

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.

[codecov]

Codecov Report

❌ Patch coverage is 6.25000% with 30 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/multimodal/processor_qwen3_omni.py	0.00%	9 Missing ⚠️
src/maxtext/multimodal/utils.py	0.00%	8 Missing ⚠️
src/maxtext/multimodal/processor.py	14.28%	6 Missing ⚠️
src/maxtext/layers/decoders.py	0.00%	3 Missing ⚠️
src/maxtext/layers/nnx_decoders.py	33.33%	1 Missing and 1 partial ⚠️
src/maxtext/models/models.py	0.00%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

🤖 Hi @hengtaoguo, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This Pull Request successfully introduces end-to-end decode support for padded and resized video inputs in Qwen3* multimodal models. It elegantly manages the propagation of video masks, handles post-projector token-mask downsampling, and correctly configures positional embedding interpolation under JIT compilation. The implementation is highly professional, clean, and backed by robust unit tests.

🔍 General Feedback

• Excellent Test Coverage: Great addition of comprehensive unit tests in multimodal_utils_test.py and qwen3_omni_layers_test.py verifying mask downsampling and token-level embedding merging.
• Flax and NNX Uniformity: Outstanding consistency in integrating the new multimodal features across both the Flax/Linen and NNX model execution paths.
• Future-proofing VisionEncoder: In src/maxtext/layers/encoders.py (which is not modified in this PR), VisionEncoder.__call__ currently only propagates video_grid_thw to the underlying encoder if input_masks is not None. Consider decoupling this check in a future PR to support unpadded videos that require grid information but do not use padding masks.

[github-actions]

🤖 I'm sorry @hengtaoguo, but I was unable to process your request. Please see the logs for more details.

[github-actions]

🤖 I'm sorry @hengtaoguo, but I was unable to process your request. Please see the logs for more details.