[muse_glimmer] Change vision encoder forward internally - #4130
Merged
Conversation
meta-cla
Bot
added
the
CLA Signed
Summary: Convert MuseGlimmerVisionEncoder.forward from the internal packed (1, total_tokens, D) path to a padded (N, P, D) batch, matching qwen3_5 and kimi_k2_7. Per-image block-diagonal attention is expressed via the batch dim (get_vision_block_mask_mod); positional embeddings, complex 2D RoPE, and the sparse-window permutation + window ids are built as padded (N, P, ...) tensors via per-unique-(h,w) grouped helpers. The sparse-window permutation is applied per row through _VisionTokenPermute.forward (which now dispatches on index.ndim: 1-D packed advanced-index or 2-D row-wise gather), keeping each window contiguous so the FlexAttention mask stays block-sparse. Windows are un-permuted and pixel-shuffle downsample loops per image, then valid tokens are flattened to (total_output_tokens, output_dim). Public forward contract and model.py are unchanged. Dead packed-path helpers (_embed_image, _finalize_image, _block_diag_mask) and the unused create_attention_mask import are removed. Test Plan: Verified numerical parity against the pre-refactor packed forward with a golden oracle (max abs diff ~5e-7, well within 1e-4), plus padded pos-emb / 2D-RoPE / sparse-perm+win_id builder parity, token_permute dispatch, block-mask semantics, TP sharding-config compatibility, model-level encoder->adapter forward, and sparse-vs-global block-sparsity. Verification harness kept local (not landed).
jinsooihm
force-pushed
the
muse-glimmer-encoder-refactor
branch
from
August 12, 2026 22:26
842c78c to
5910ed3
Compare
jinsooihm
marked this pull request as ready for review
August 12, 2026 22:36
jinsooihm
requested review from
fegin,
tianyu-l,
wconstab and
wwwjn
as code owners
August 12, 2026 22:36
tianyu-l
approved these changes
Aug 13, 2026
fegin
approved these changes
Aug 13, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
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.
Summary:
Although
MuseGlimmerVisionEncodertakes in the same padded(N, P, D)(num images, max patches, latent dimension) formatted input as Qwen3.5 and Kimi2.7, internally, it ran in a packed(1, total_tokens, D)format.This PR changes the internal behavior to keep
(N, P, ...)shape throughout.This improves consistency with other models and would make spmd_types enablement easier after #3895 and #4079.
Test Plan:
Verified numerical parity against the pre-refactor packed forward (max abs diff ~5e-7, well within 1e-4).