# SPDX-License-Identifier: Apache-2.0 # Copyright 2026 The Oblivion Authors. """The fused Oblivion model for vLLM v1 -- multimodal, single route. Oblivion subclasses the full native multimodal ``MuseGlimmerForCausalLM`` and splices two donor macroblocks into the text decoder: a Qwen3.5 block after muse layer 35 and a Gemma4 block after layer 47. Both donors are ALWAYS applied at the canonical gate (``CANONICAL_GATE``) -- there is no per-request routing or gating. Everything multimodal (vision tower, image/video embedding merge, the ``@MULTIMODAL_REGISTRY.register_processor``) is inherited unchanged, so image prompts run through Glimmer's vision path. * ``super().__init__`` builds the vision tower, the native ``self.model`` (``MuseGlimmerModel``, 52 text layers at ``model.layers.0..51``), ``lm_head``, and multimodal token handling. * Donor decoder layers are first-class vLLM modules built under the live fused ``vllm_config`` (so each Attention / GatedDeltaNet self-registers KV / mamba state), held in ``self.qwen_layers`` / ``self.gemma_layers`` (ModuleDicts keyed by source index, prefixes ``qwen_layers.{24..39}`` / ``gemma_layers.{48..59}``). * ``forward`` runs the text loop with two seams: muse 0..35 -> qwen seam -> muse 36..47 -> gemma seam -> muse 48..51 -> norm. It also captures auxiliary hidden states when spec decoding is active (for the EAGLE3 drafter). The Qwen donor is recurrent (Gated DeltaNet), so the fused arch is hybrid: it implements ``IsHybrid`` + ``HasInnerState`` (launch ``--mamba-cache-mode=align``). """ from __future__ import annotations from collections.abc import Iterable from typing import Any, ClassVar, Literal import torch from torch import nn from vllm.config import VllmConfig from vllm.model_executor.models.gemma4 import Gemma4DecoderLayer from vllm.model_executor.models.interfaces import HasInnerState, IsHybrid from vllm.model_executor.models.muse_glimmer import ( MuseGlimmerDummyInputsBuilder, MuseGlimmerForCausalLM, MuseGlimmerMultiModalProcessor, MuseGlimmerProcessingInfo, ) from vllm.model_executor.models.qwen3_5 import Qwen3_5DecoderLayer from vllm.model_executor.models.qwen3_next import Qwen3NextForCausalLM from vllm.model_executor.models.utils import AutoWeightsLoader from vllm.multimodal import MULTIMODAL_REGISTRY from .config import ( CANONICAL_GATE, GEMMA_BOUNDARY_MUSE_LAYER, HOST_HIDDEN, QWEN_BOUNDARY_MUSE_LAYER, ) from .donor_config import ( GEMMA_LAYER_INDICES, QWEN_LAYER_INDICES, clone_vllm_config_for_donor, load_donor_text_config, load_manifest, ) from .layers import OblivionDonorBranch, OblivionSeam, OblivionStitch from .weights import load_oblivion_weights _MUSE_TEXT_LAYERS = 52 _DONOR_SKIP_PREFIXES = ("qwen_layers.", "gemma_layers.", "qwen_seam.", "gemma_seam.") @MULTIMODAL_REGISTRY.register_processor( MuseGlimmerMultiModalProcessor, info=MuseGlimmerProcessingInfo, dummy_inputs=MuseGlimmerDummyInputsBuilder, ) class OblivionForCausalLM(MuseGlimmerForCausalLM, IsHybrid, HasInnerState): """Fused Muse (multimodal) + Qwen3.5 + Gemma4, one model, both donors always on.""" is_hybrid: ClassVar[Literal[True]] = True def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: # (a) Build the full native multimodal Muse model: vision tower, the # native text model self.model (52 layers), lm_head, mm-token handling. super().__init__(vllm_config=vllm_config, prefix=prefix) # Spec-decode (EAGLE3) navigates the target as ``get_language_model().model`` # and reads ``.embed_tokens`` off it. On multimodal MuseGlimmer, # ``get_language_model()`` returns the inner ``MuseGlimmerModel`` (which # already holds ``embed_tokens``/``layers``), so give it a self-referential # ``.model`` alias. object.__setattr__ keeps it OUT of the submodule graph. if not hasattr(self.model, "model"): object.__setattr__(self.model, "model", self.model) # vLLM's EAGLE3 proposer reads ``target_model.config.image_token_index`` # for multimodal targets; MuseGlimmer exposes it as ``image_token_id``. if not hasattr(self.config, "image_token_index"): self.config.image_token_index = getattr(self.config, "image_token_id", 200092) self.vllm_config = vllm_config cache_config = vllm_config.cache_config quant_config = vllm_config.quant_config max_rms = float(getattr(self.config, "max_relative_rms", 1.0)) # (b) Donor decoder layers, built under the live fused vllm_config so each # donor Attention / GDN self-registers KV/mamba state. Unique # single-integer prefixes (source index) never collide with # model.layers.{i}. qwen_text_config = load_donor_text_config(vllm_config, "qwen") gemma_text_config = load_donor_text_config(vllm_config, "gemma") self._qwen_donor_config = qwen_text_config qwen_vllm_config = clone_vllm_config_for_donor(vllm_config, qwen_text_config) qwen_manifest = load_manifest(vllm_config, "qwen") qwen_layer_types = qwen_manifest["layer_types"] # per local index self.qwen_layers = nn.ModuleDict() for local, src in enumerate(QWEN_LAYER_INDICES): self.qwen_layers[str(src)] = Qwen3_5DecoderLayer( qwen_vllm_config, layer_type=qwen_layer_types[local], prefix=f"qwen_layers.{src}", ) self.gemma_layers = nn.ModuleDict() for src in GEMMA_LAYER_INDICES: # Gemma4DecoderLayer derives sliding/full geometry + per-layer RoPE from # config.layer_types[extract_layer_index(prefix)] == src. self.gemma_layers[str(src)] = Gemma4DecoderLayer( gemma_text_config, cache_config, quant_config, prefix=f"gemma_layers.{src}", ) # (c) Seams: own the two stitches, reference the donor layers by order. self.qwen_seam = self._build_seam( "qwen", HOST_HIDDEN, qwen_text_config.hidden_size, list(self.qwen_layers.values()), max_rms, ) self.gemma_seam = self._build_seam( "gemma", HOST_HIDDEN, gemma_text_config.hidden_size, list(self.gemma_layers.values()), max_rms, ) # -- seam construction --------------------------------------------------- def _build_seam( self, donor: str, host_hidden: int, donor_hidden: int, donor_layers: list[nn.Module], max_rms: float, ) -> OblivionSeam: seam_prefix = f"{donor}_seam" input_stitch = OblivionStitch( host_hidden, donor_hidden, prefix=f"{seam_prefix}.input_stitch" ) output_stitch = OblivionStitch( donor_hidden, host_hidden, prefix=f"{seam_prefix}.output_stitch" ) branch = OblivionDonorBranch( donor, input_stitch, nn.ModuleList(donor_layers), output_stitch, (), max_relative_rms=max_rms, ) return OblivionSeam(branch, donor, prefix=seam_prefix) # -- forward ------------------------------------------------------------- def forward( self, input_ids: torch.Tensor | None, positions: torch.Tensor, intermediate_tensors: Any = None, inputs_embeds: torch.Tensor | None = None, **kwargs: Any, ) -> torch.Tensor | Any: m = self.model # Merged text+image embeddings arrive as inputs_embeds from the runner. if inputs_embeds is not None: hidden = inputs_embeds else: hidden = m.embed_input_ids(input_ids) residual: torch.Tensor | None = None layers = m.layers # Both donors always on: a constant per-token gate at the canonical value. gate = hidden.new_full((hidden.shape[0],), CANONICAL_GATE) # EAGLE3 auxiliary hidden-state capture: when spec decoding is active vLLM # sets aux_hidden_state_layers on the inner model and expects forward to # return (hidden, aux). Capture the true post-layer hidden at each requested # muse-layer index; layers after a seam expose donor-influenced states. aux_layers = getattr(m, "aux_hidden_state_layers", None) or () aux: list[torch.Tensor] = [] # muse 0..35 for i in range(QWEN_BOUNDARY_MUSE_LAYER + 1): hidden, residual = layers[i](positions, hidden, residual) if aux_layers and (i + 1) in aux_layers: aux.append(hidden) hidden = self.qwen_seam(hidden, positions, gate) residual = None # muse 36..47 for i in range(QWEN_BOUNDARY_MUSE_LAYER + 1, GEMMA_BOUNDARY_MUSE_LAYER + 1): hidden, residual = layers[i](positions, hidden, residual) if aux_layers and (i + 1) in aux_layers: aux.append(hidden) hidden = self.gemma_seam(hidden, positions, gate) residual = None # muse 48..51 for i in range(GEMMA_BOUNDARY_MUSE_LAYER + 1, _MUSE_TEXT_LAYERS): hidden, residual = layers[i](positions, hidden, residual) if aux_layers and (i + 1) in aux_layers: aux.append(hidden) normed = m.norm(hidden) if aux: return normed, aux return normed # -- weight loading ------------------------------------------------------ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Muse (vision tower + language + lm_head) loads via the inherited # AutoWeightsLoader path with muse's own mapper; the donor submodules load # from the co-located donor artifacts in the self-contained checkpoint dir. loaded = AutoWeightsLoader(self, skip_prefixes=list(_DONOR_SKIP_PREFIXES)) loaded_names = loaded.load_weights(weights, mapper=self.hf_to_vllm_mapper) loaded_names |= load_oblivion_weights(self) return loaded_names # -- IsHybrid / HasInnerState (mamba state from the Qwen GDN donor) ------- @classmethod def get_mamba_state_shape_from_config( cls, vllm_config: VllmConfig ) -> tuple[tuple[int, ...], ...]: donor = clone_vllm_config_for_donor( vllm_config, load_donor_text_config(vllm_config, "qwen") ) return Qwen3NextForCausalLM.get_mamba_state_shape_from_config(donor) @classmethod def get_mamba_state_dtype_from_config( cls, vllm_config: VllmConfig ) -> tuple[torch.dtype, ...]: donor = clone_vllm_config_for_donor( vllm_config, load_donor_text_config(vllm_config, "qwen") ) return Qwen3NextForCausalLM.get_mamba_state_dtype_from_config(donor) @classmethod def get_mamba_state_copy_func(cls) -> Any: return Qwen3NextForCausalLM.get_mamba_state_copy_func()