vllm.model_executor.models.muse_glimmer

Inference-only MuseGlimmer multimodal model for vLLM.

Native port of the MuseGlimmer text decoder (MuseGlimmerForCausalLM). The text stack is a Gemma2 derivative with the following MuseGlimmer-specific deltas, each of which is handled explicitly here:

• SiLU-gated MLP (hidden_activation="silu"), not Gemma's gelu-tanh.
• Scaleless RMSNorm on the token embeddings (no sqrt(hidden) scaling).
• Per-layer sandwich RMSNorms with a baked +1 weight offset (x * (1 + w)), matching Gemma, but with distinct eps for the pre/post norms (rms_norm_eps vs post_norm_eps).
• QK-norm (weightless, fp32) applied before RoPE, followed by a query pre-scale of qk_scale_factor / sqrt(head_dim).
• A per-head sigmoid attention output gate.
• iRoPE layout: NoPE layers use full attention, RoPE layers use sliding window attention. RoPE is applied NEOX-style (is_neox_style=True): the HF converter (convert_muse_glimmer_weights_to_hf.py, 20260806+) permutes q/k into the half-split (NEOX) layout via _permute_for_rope so they pair with rotate_half — matching the reference's interleaved rotation on the native (unpermuted) weights. Serving the permuted HF weights with is_neox_style=False scrambles q/k and causes token-repetition collapse.
• Final logits are pre-scaled by output_multiplier and then tanh soft-capped at final_logit_softcapping.
• Untied lm_head.

The vision path supports variable-resolution images and temporally patched videos. It mirrors the checkpoint's native vision encoder, including sparse block attention, 2-D RoPE, pixel-shuffle downsampling, and the two-layer adapter/projection stack.

Classes:

• MuseGlimmerForCausalLM –
• MuseGlimmerImagePixelInputs –

  Dimensions:

• MuseGlimmerRMSNorm –

  RMSNorm mirroring HF MuseGlimmer exactly (fp32 compute, cast at the end).

• MuseGlimmerVideoPixelInputs –

  Dimensions:

MuseGlimmerForCausalLM

Bases: Module, SupportsLoRA, SupportsMultiModal, SupportsPP, SupportsEagle3

Methods:

• get_mm_mapping –

  Get the module prefix in multimodal models

Source code in vllm/model_executor/models/muse_glimmer.py

@MULTIMODAL_REGISTRY.register_processor(
    MuseGlimmerMultiModalProcessor,
    info=MuseGlimmerProcessingInfo,
    dummy_inputs=MuseGlimmerDummyInputsBuilder,
)
class MuseGlimmerForCausalLM(
    nn.Module, SupportsLoRA, SupportsMultiModal, SupportsPP, SupportsEagle3
):
    # Weight-name normalization. Two checkpoint conventions are supported:
    #
    #   * HF MuseGlimmer export (``convert_muse_glimmer_weights_to_hf.py``): the
    #     multimodal ``MuseGlimmerConfig`` prefixes the language model with
    #     ``model.language_model.``; the per-layer sandwich norms are already
    #     named ``input_layernorm`` / ``post_attention_layernorm`` /
    #     ``pre_feedforward_layernorm`` / ``post_feedforward_layernorm``.
    #
    #   * Legacy HF export (an earlier checkpoint convention): uses ``model.``
    #     and a different sandwich-norm naming where ``post_attn_norm`` is the
    #     true post-attention norm and ``post_attention_layernorm`` is actually
    #     the pre-feedforward norm. We remap those to MuseGlimmer's names.
    #
    # CONVENTION DISAMBIGUATION (critical): the two checkpoint families use
    # DIFFERENT sandwich-norm names, and they must not be conflated:
    #
    #   * Canonical MuseGlimmer export (current
    #     ``convert_muse_glimmer_weights_to_hf.py`` — what partners ship):
    #     keys are ``model.language_model.layers.N.*`` and
    #     the norms are ALREADY named ``input_layernorm`` /
    #     ``post_attention_layernorm`` / ``pre_feedforward_layernorm`` /
    #     ``post_feedforward_layernorm``. No norm rename needed — pass through.
    #
    #   * Legacy HF export (an earlier checkpoint convention): keys are
    #     ``model.layers.N.*`` and the sandwich norms are
    #     named ``input_layernorm`` / ``post_attention_layernorm`` (this one is
    #     actually the PRE-feedforward norm) / ``post_attn_norm`` (the true
    #     post-attention norm) / ``post_ffn_norm``. These must be remapped.
    #
    # The unambiguous discriminator is the PREFIX: legacy keys start with
    # ``model.layers.`` while canonical keys start with
    # ``model.language_model.layers.``. ``orig_to_new_regex`` runs BEFORE the
    # prefix strip (see WeightsMapper._map_name_with_shard), so we anchor the
    # legacy renames on ``^model\.layers\.`` — they fire ONLY on legacy keys and
    # leave canonical/partner checkpoints untouched. Rule order within the regex
    # dict matters: the ``post_attention_layernorm`` -> ``pre_feedforward_...``
    # rule must precede the ``post_attn_norm`` -> ``post_attention_layernorm``
    # rule so the latter's output is not re-captured by the former (regex rules
    # apply as a single forward pass).
    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_substr={
            "model.vision_tower.patch_embedder.position_embedding_table.weight": (
                "model.vision_tower.positional_embedding_vlm"
            ),
            "model.vision_tower.layers.": "model.vision_tower.transformer.",
            ".norm1.": ".ln_1.",
            ".norm2.": ".ln_2.",
            ".attn.proj.": ".attn.o_proj.",
            ".mlp.fc1.": ".mlp.c_fc.",
            ".mlp.fc2.": ".mlp.c_proj.",
            ".self_attn.gate_proj": ".self_attn.output_gate_proj",
        },
        orig_to_new_prefix={
            "model.rotary_emb.": None,
            "model.language_model.": "model.",
            "language_model.": "model.",
            "model.vision_tower.patch_embedder.patch_embedding.": (
                "model.vision_tower.conv1_linear."
            ),
            "model.vision_tower.": "vision_encoder.",
            "vision_tower.": "vision_encoder.",
            "model.vision_encoder.": "vision_encoder.",
            "model.vision_adapter.fc1.": "model.vision_adapter.c_fc.",
            "model.vision_adapter.fc2.": "model.vision_adapter.c_proj.",
            "model.vision_adapter.": "vision_adapter.",
            "model.vision_projection.": "vision_projection.",
            "model.perception_emb_norm.": "perception_emb_norm.",
        },
        orig_to_new_stacked={
            ".q_proj": (".qkv_proj", "q"),
            ".k_proj": (".qkv_proj", "k"),
            ".v_proj": (".qkv_proj", "v"),
            ".gate_proj": (".gate_up_proj", 0),
            ".up_proj": (".gate_up_proj", 1),
        },
    )

    packed_modules_mapping = {
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
        "gate_up_proj": ["gate_proj", "up_proj"],
    }

    @classmethod
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
        if modality.startswith("image"):
            return IMAGE_TOKEN
        if modality.startswith("video"):
            return VIDEO_TOKEN
        raise ValueError(f"Unsupported MuseGlimmer modality: {modality}")

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        config = vllm_config.model_config.hf_config
        text_config = _text_config(config)
        quant_config = vllm_config.quant_config
        self.config = config
        self.text_config = text_config
        self.quant_config = quant_config
        self.has_vision = _muse_glimmer_has_vision(config)

        with self._mark_language_model(vllm_config):
            self.model = MuseGlimmerModel(
                vllm_config=vllm_config,
                prefix=maybe_prefix(prefix, "model"),
            )

        self.vision_encoder: MuseGlimmerVisionEncoder | None
        self.vision_adapter: MuseGlimmerVisionAdapter | None
        self.perception_emb_norm: MuseGlimmerRMSNorm | None
        if self.has_vision:
            with self._mark_tower_model(vllm_config, {"image", "video"}):
                self.vision_encoder = MuseGlimmerVisionEncoder(
                    config,
                    prefix=maybe_prefix(prefix, "vision_encoder"),
                )
                self.vision_adapter = MuseGlimmerVisionAdapter(config)
                self.vision_projection = nn.Linear(
                    config.projector_hidden_size,
                    text_config.hidden_size,
                    bias=False,
                )
                self.perception_emb_norm = MuseGlimmerRMSNorm(
                    eps=text_config.rms_norm_eps, with_scale=False
                )
        else:
            self.vision_encoder = None
            self.vision_adapter = None
            self.vision_projection = None
            self.perception_emb_norm = None

        self.lm_head = ParallelLMHead(
            text_config.vocab_size,
            text_config.hidden_size,
            quant_config=quant_config,
            prefix=maybe_prefix(prefix, "lm_head"),
        )
        self.output_multiplier = text_config.output_multiplier
        self.final_logit_softcapping = text_config.final_logit_softcapping
        self.logits_processor = LogitsProcessor(text_config.vocab_size)
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors
        )
        self.configure_mm_token_handling(
            text_config.vocab_size,
            [
                getattr(config, "image_token_id", 200092),
                getattr(config, "video_token_id", 200091),
            ],
        )

    def _parse_and_validate_image_input(
        self,
        **kwargs: object,
    ) -> MuseGlimmerImagePixelInputs | None:
        pixel_values = kwargs.pop("pixel_values", None)
        if pixel_values is None:
            return None
        return MuseGlimmerImagePixelInputs(
            type="pixel_values",
            pixel_values=pixel_values,
            image_grid_thw=kwargs.pop("image_grid_thw", None),
        )

    def _parse_and_validate_video_input(
        self,
        **kwargs: object,
    ) -> MuseGlimmerVideoPixelInputs | None:
        pixel_values_videos = kwargs.pop("pixel_values_videos", None)
        if pixel_values_videos is None:
            return None
        return MuseGlimmerVideoPixelInputs(
            type="pixel_values_videos",
            pixel_values_videos=pixel_values_videos,
            video_grid_thw=kwargs.pop("video_grid_thw", None),
        )

    def _encode_pixel_values(
        self,
        pixel_values: torch.Tensor,
        grid_thw: torch.Tensor,
    ) -> tuple[torch.Tensor, ...]:
        if (
            self.vision_encoder is None
            or self.vision_adapter is None
            or self.vision_projection is None
            or self.perception_emb_norm is None
        ):
            raise ValueError("This MuseGlimmer checkpoint has no vision tower")
        features = self.vision_encoder(pixel_values, grid_thw)
        features = self.vision_adapter(features)
        features = self.vision_projection(features)
        features = self.perception_emb_norm(features)
        merge_length = self.vision_encoder.merge_size**2
        return features.split((grid_thw.prod(-1) // merge_length).tolist())

    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
        image_input = self._parse_and_validate_image_input(**kwargs)
        video_input = self._parse_and_validate_video_input(**kwargs)

        embeddings: list[torch.Tensor] = []
        for key in kwargs:
            if key == "pixel_values" and image_input is not None:
                embeddings.extend(
                    self._encode_pixel_values(
                        image_input["pixel_values"], image_input["image_grid_thw"]
                    )
                )
            elif key == "pixel_values_videos" and video_input is not None:
                embeddings.extend(
                    self._encode_pixel_values(
                        video_input["pixel_values_videos"],
                        video_input["video_grid_thw"],
                    )
                )
        return embeddings

    def forward(
        self,
        input_ids: torch.Tensor | None,
        positions: torch.Tensor,
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
        **kwargs: object,
    ) -> torch.Tensor | IntermediateTensors:
        return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)

    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
        logits = self.logits_processor(self.lm_head, hidden_states)
        if logits is None:
            return None
        logits = logits * self.output_multiplier
        if self.final_logit_softcapping is not None:
            cap = self.final_logit_softcapping
            logits = cap * torch.tanh(logits / cap)
        return logits

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

    def get_mm_mapping(self) -> MultiModelKeys:
        """Get the module prefix in multimodal models"""
        return MultiModelKeys.from_string_field(
            language_model="model",
            connector=["vision_adapter.", "vision_projection."],
            tower_model="vision_encoder.",
        )

get_mm_mapping()

Get the module prefix in multimodal models

Source code in vllm/model_executor/models/muse_glimmer.py

def get_mm_mapping(self) -> MultiModelKeys:
    """Get the module prefix in multimodal models"""
    return MultiModelKeys.from_string_field(
        language_model="model",
        connector=["vision_adapter.", "vision_projection."],
        tower_model="vision_encoder.",
    )

MuseGlimmerImagePixelInputs

Bases: TensorSchema

Dimensions: - np: Total number of patches over all images - cps: Temporal patch size * channels * patch size * patch size - ni: Number of images

Source code in vllm/model_executor/models/muse_glimmer.py

class MuseGlimmerImagePixelInputs(TensorSchema):
    """Dimensions:
    - np: Total number of patches over all images
    - cps: Temporal patch size * channels * patch size * patch size
    - ni: Number of images
    """

    type: Literal["pixel_values"]
    pixel_values: Annotated[torch.Tensor, TensorShape("np", "cps")]
    image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]

MuseGlimmerRMSNorm

Bases: Module

RMSNorm mirroring HF MuseGlimmer exactly (fp32 compute, cast at the end).

normed = _norm(x.float()) * (w.float() + weight_offset) cast back to the input dtype. When with_scale is False the layer is weightless (used for QK-norm and the token-embedding norm).

Source code in vllm/model_executor/models/muse_glimmer.py

class MuseGlimmerRMSNorm(nn.Module):
    """RMSNorm mirroring HF MuseGlimmer exactly (fp32 compute, cast at the end).

    ``normed = _norm(x.float()) * (w.float() + weight_offset)`` cast back to the
    input dtype. When ``with_scale`` is False the layer is weightless (used for
    QK-norm and the token-embedding norm).
    """

    def __init__(
        self,
        dim: int | None = None,
        eps: float = 1e-6,
        with_scale: bool = True,
        weight_offset: int = 0,
    ) -> None:
        super().__init__()
        self.eps = eps
        self.with_scale = with_scale
        self.weight_offset = weight_offset
        if with_scale:
            assert dim is not None
            self.weight = nn.Parameter(torch.zeros(dim))
        else:
            self.register_parameter("weight", None)

    def _norm(self, x: torch.Tensor) -> torch.Tensor:
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        out = self._norm(hidden_states.float())
        if self.with_scale:
            out = out * (self.weight.float() + self.weight_offset)
        return out.type_as(hidden_states)

MuseGlimmerVideoPixelInputs

Bases: TensorSchema

Dimensions: - np: Total number of patches over all videos - ctps: Temporal patch size * channels * patch size * patch size - nv: Number of videos

Source code in vllm/model_executor/models/muse_glimmer.py

class MuseGlimmerVideoPixelInputs(TensorSchema):
    """Dimensions:
    - np: Total number of patches over all videos
    - ctps: Temporal patch size * channels * patch size * patch size
    - nv: Number of videos
    """

    type: Literal["pixel_values_videos"]
    pixel_values_videos: Annotated[torch.Tensor, TensorShape("np", "ctps")]
    video_grid_thw: Annotated[torch.Tensor, TensorShape("nv", 3)]

_text_config(config)

MuseGlimmer checkpoints may nest the text config under text_config (multimodal MuseGlimmerConfig) or expose it directly (MuseGlimmerTextConfig).

Source code in vllm/model_executor/models/muse_glimmer.py

def _text_config(config):
    """MuseGlimmer checkpoints may nest the text config under ``text_config``
    (multimodal ``MuseGlimmerConfig``) or expose it directly
    (``MuseGlimmerTextConfig``)."""
    return getattr(config, "text_config", config)