vllm.models.inkling.amd.model ¶

class InklingForCausalLM(_TmlForCausalLMBase):
    """Text-only entry point (``inkling_model`` checkpoints)."""

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
        super().__init__()
        self._build(vllm_config, vllm_config.model_config.hf_config, prefix)

@MULTIMODAL_REGISTRY.register_processor(
    InklingMultiModalProcessor,
    info=InklingProcessingInfo,
    dummy_inputs=InklingDummyInputsBuilder,
)
class InklingForConditionalGeneration(_TmlForCausalLMBase, SupportsMultiModal):
    """Top-level (multimodal) entry point.

    Builds the vision + audio towers on top of the shared text backbone. Inkling has
    NO cross-modal fusion (the vision tower emits one token per patch, the audio
    tower one token per frame), so generation reuses the inherited backbone
    ``forward`` / ``compute_logits`` (the latter already applies muP) and this
    class only adds multimodal embedding + merge.
    """

    hf_to_vllm_mapper = _TmlForCausalLMBase.hf_to_vllm_mapper | WeightsMapper(
        orig_to_new_prefix={
            "model.audio.": "audio.",
            "model.visual.": "visual.vision_encoder.",
        },
    )

    @classmethod
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
        if modality.startswith("image"):
            return "<|content_image|>"
        if modality.startswith("audio"):
            return "<|content_audio_input|>"
        raise ValueError("Only image or audio modality is supported")

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
        super().__init__()
        config: InklingMMConfig = vllm_config.model_config.hf_config

        self.visual = (
            InklingVision(config.vision_config, prefix=maybe_prefix(prefix, "visual"))
            if inkling_vision_enabled(config)
            else None
        )
        self.audio = (
            InklingAudio(config.audio_config, prefix=maybe_prefix(prefix, "audio"))
            if inkling_audio_enabled(config)
            else None
        )

        self._build(vllm_config, config.text_config, prefix)

    # -- multimodal embedding -------------------------------------------

    def _process_image_input(
        self, pixel_values: Any, num_patches: Any
    ) -> tuple[torch.Tensor, ...]:
        assert self.visual is not None
        # pixel_values is a list (per item) of [P_i, 2, P, P, 3] tensors,
        # or a single concatenated tensor. Normalize to a flat batch, run the
        # tower once, then split back per item.
        if isinstance(pixel_values, (list, tuple)):
            if not pixel_values:
                return ()
            sizes = [int(p.shape[0]) for p in pixel_values]
            patches = torch.cat(list(pixel_values), dim=0)
        else:
            patches = pixel_values
            sizes = self._sizes_from(num_patches, patches.shape[0])

        patches = patches.to(device=self.visual.device, dtype=self.visual.dtype)
        embeds = self.visual(patches)  # [total_patches, D]
        return tuple(embeds.split(sizes))

    def _process_audio_input(
        self, input_audio_features: Any, num_audio_tokens: Any
    ) -> tuple[torch.Tensor, ...]:
        assert self.audio is not None
        if isinstance(input_audio_features, (list, tuple)):
            if not input_audio_features:
                return ()
            sizes = [int(d.shape[0]) for d in input_audio_features]
            dmel = torch.cat(list(input_audio_features), dim=0)
        else:
            dmel = input_audio_features
            sizes = self._sizes_from(num_audio_tokens, dmel.shape[0])

        dmel = dmel.to(device=self.audio.device)
        embeds = self.audio(dmel)  # [total_frames, D]
        return tuple(embeds.split(sizes))

    @staticmethod
    def _sizes_from(counts: Any, total: int) -> list[int]:
        if counts is None:
            return [total]
        if isinstance(counts, torch.Tensor):
            return [int(c) for c in counts.flatten().tolist()]
        if isinstance(counts, (list, tuple)):
            flat: list[int] = []
            for c in counts:
                flat.append(int(c.item()) if isinstance(c, torch.Tensor) else int(c))
            return flat
        return [int(counts)]

    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
        # Iterate modalities in a stable order so the returned per-item tensors
        # line up with their appearance order; the positional merge in
        # embed_input_ids handles actual placement.
        pixel_values = kwargs.get("pixel_values")
        num_patches = kwargs.get("num_patches")
        input_audio_features = kwargs.get("input_audio_features")
        num_audio_tokens = kwargs.get("num_audio_tokens")

        embeddings: tuple[torch.Tensor, ...] = ()
        if pixel_values is not None and self.visual is not None:
            embeddings += self._process_image_input(pixel_values, num_patches)
        if input_audio_features is not None and self.audio is not None:
            embeddings += self._process_audio_input(
                input_audio_features, num_audio_tokens
            )
        return embeddings

    def embed_input_ids(
        self,
        input_ids: torch.Tensor,
        multimodal_embeddings: MultiModalEmbeddings | None = None,
        *,
        is_multimodal: torch.Tensor | None = None,
    ) -> torch.Tensor:
        # Override the base's 1-arg embed_input_ids: the runner calls this 3-arg
        # signature for multimodal models. Text embeddings come from the shared
        # backbone (which applies embed_norm); MM embeddings are scattered in.
        from vllm.model_executor.models.utils import _merge_multimodal_embeddings

        # Placeholder ids use unused vocabulary slots and these positions are
        # overwritten by MM embeds below.
        inputs_embeds = self.model.embed_input_ids(input_ids)
        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
            return inputs_embeds
        assert is_multimodal is not None
        return _merge_multimodal_embeddings(
            inputs_embeds=inputs_embeds,
            multimodal_embeddings=multimodal_embeddings,
            is_multimodal=is_multimodal,
        )

    def get_language_model(self) -> nn.Module:
        # This class IS the causal LM (the towers are side branches), so the
        # language model is self — callers expect a module exposing ``.model``
        # / ``.lm_head`` (e.g. the MTP/eagle loader shares embeddings via
        # ``get_language_model().model.embed_tokens``).
        return self

class InklingReplicatedEmbedding(nn.Module):
    """Full-vocab embedding table replicated on every TP rank.

    Trades the full table per rank (~2.3 GiB at V=201k / H=6144 bf16, vs a
    1/tp shard) for no masked lookup and no per-lookup TP all-reduce — one
    all-reduce per MTP draft step plus one per verify pass — and keeps the
    full table on-rank for the fused gather+norm kernels (``embed_rmsnorm``,
    ``embed_dual_rmsnorm_cat``). Bit-exact vs vocab-parallel: the all-reduce
    there only ever summed one real row against exact zeros. The LM head
    stays vocab-sharded.
    """

    def __init__(self, num_embeddings: int, embedding_dim: int) -> None:
        super().__init__()
        self.weight = nn.Parameter(
            torch.empty(num_embeddings, embedding_dim, dtype=torch.get_default_dtype()),
            requires_grad=False,
        )

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        return embed_rmsnorm(input_ids, self.weight, None, 0.0)

class _TmlForCausalLMBase(nn.Module, SupportsPP, SupportsLoRA):
    """Shared text-backbone causal-LM scaffolding for both entry classes."""

    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_substr={
            ".w13_dn": ".gate_up_proj",
            ".w2_md": ".down_proj",
        },
        orig_to_new_stacked={
            ".attn.wq_du.": (".attn.qkvr.", 0),
            ".attn.wk_dv.": (".attn.qkvr.", 1),
            ".attn.wv_dv.": (".attn.qkvr.", 2),
            ".attn.wr_du.": (".attn.qkvr.", 3),
        },
        orig_to_new_prefix={
            "model.llm.layers.": "model.layers.",
            "model.llm.embed_norm": "model.embed_norm",
            "model.llm.embed": "model.embed_tokens",
            "model.llm.norm": "model.norm",
            "model.llm.unembed": "lm_head",
            "language_model.layers.": "model.layers.",
            "language_model.lm_head.": "lm_head.",
        },
        orig_to_new_suffix={
            # NVFP4 scale
            ".w13_weight.scale": ".w13_weight_scale",
            ".w13_weight.scale2": ".w13_weight_scale_2",
            ".w2_weight.scale": ".w2_weight_scale",
            ".w2_weight.scale2": ".w2_weight_scale_2",
        },
    )
    # Quark uses this mapping when resolving quantization exclusions for the
    # four checkpoint attention projections fused into qkvr.
    packed_modules_mapping = {
        "qkvr": ["wq_du", "wk_dv", "wv_dv", "wr_du"],
        "w13": ["w1", "w3"],
    }
    embedding_modules = {
        "lm_head": "output_embeddings",
    }

    def _build(
        self,
        vllm_config: VllmConfig,
        text_config: InklingModelConfig,
        prefix: str,
    ) -> None:
        quant_config = vllm_config.quant_config
        self.config = text_config
        # ROCm checkpoints use Quark OCP MXFP4. The global Quark config is
        # passed into each routed MoE and its exclusion list keeps the rest of
        # Inkling in bf16.
        # Read by the MRV2 runner to publish per-request short-conv metadata.
        # Short convolution is intrinsic to Inkling, so this is always set.
        self.uses_sconv = True
        self.model = InklingModel(
            config=text_config,
            quant_config=quant_config,
            prefix=maybe_prefix(prefix, "model"),
        )
        self.lm_head = ParallelLMHead(
            text_config.padded_vocab_size,
            text_config.hidden_size,
            org_num_embeddings=text_config.padded_vocab_size,
            quant_config=quant_config,
            prefix=maybe_prefix(prefix, "lm_head"),
        )
        self.logits_processor = InklingLogitsProcessor(
            text_config.padded_vocab_size,
            org_vocab_size=text_config.vocab_size,
            soft_cap=text_config.final_logit_softcapping,
            logits_mup_width_multiplier=text_config.logits_mup_width_multiplier,
        )
        self.make_empty_intermediate_tensors = (  # type: ignore[method-assign]
            self.model.make_empty_intermediate_tensors
        )

    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)

    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:
        return self.logits_processor(self.lm_head, hidden_states)

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        return _load_inkling_weights(self, weights, self.config)