Add Llama4VisionModel for multimodal decoding #1809
+127
−46
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
This PR implements a complete class
Llama4VisionModel
, by integrating all Llama4 basic vision components. It allows Llama4 multimodal decode by describing an input image. Joint by @hengtaoguo and @aireenmei .Llama4VisionModel
, which converts image tiles (batch_size, num_tiles, C, H, W) to feature activations (batch_size, num_tiles, num_patches, vision_output_dim_for_vit). ThenLlama4MultiModalProjector
projects it to (batch_size, num_tiles, num_patches, base_emb_dim). Example: (8, 5, 3, 336, 336) -> (8, 5, 144, 4096) -> (8, 5, 144, 5120)get_dummy_image_shape_for_init()
to create desired dummy images for different models, for jit init purpose.Tests
Tested full multimodal decode on v5p-16 cluster with this command, and workload with screenshot.

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