vllm.model_executor.models.gemma4_mm ¶

class Gemma4AudioInputs(TensorSchema):
    """
    Dimensions:
        - bn: Batch size * number of audios
        - s: Sequence length (MEL spectrogram frames)
        - f: Number of features (MEL bins)
    """

    type: Literal["audio"] = "audio"
    input_features_padded: Annotated[
        torch.Tensor, TensorShape("bn", "s", "f", dynamic_dims={"s"})
    ]
    input_features_mask: Annotated[
        torch.Tensor, TensorShape("bn", "s", dynamic_dims={"s"})
    ]

@MULTIMODAL_REGISTRY.register_processor(
    Gemma4MultiModalProcessor,
    info=Gemma4ProcessingInfo,
    dummy_inputs=Gemma4DummyInputsBuilder,
)
class Gemma4ForConditionalGeneration(
    nn.Module,
    SupportsMultiModal,
    SupportsPP,
    SupportsLoRA,
    SupportsEagle3,
):
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

    # Maps checkpoint prefixes to vLLM module paths.
    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_prefix={
            "model.embed_audio.": "embed_audio.",
            "model.embed_vision.": "embed_vision.",
            "model.language_model.": "language_model.model.",
            "model.vision_tower.": "vision_tower.",
            "model.audio_tower.": "audio_tower.",
            "lm_head.": "language_model.lm_head.",
            "model": "language_model.model",
        }
    )

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

        # ---- Vision tower (shared by image and video) ----
        with self._mark_tower_model(vllm_config, {"image", "video"}):
            self.vision_tower = AutoModel.from_config(config=config.vision_config)
            self.embed_vision = Gemma4MultimodalEmbedder(
                config.vision_config, config.text_config
            )

        # ---- Audio tower (variants with audio_config) ----
        if config.audio_config is not None:
            with self._mark_tower_model(vllm_config, "audio"):
                self.audio_tower = AutoModel.from_config(config=config.audio_config)
                # AutoModel.from_config does NOT call post_init(),
                # which is needed to initialize buffers that are absent
                # from the checkpoint (e.g. inv_timescales for relative
                # position embeddings, softcap, gradient_clipping).
                self.audio_tower.post_init()
                self.embed_audio = Gemma4MultimodalEmbedder(
                    config.audio_config, config.text_config
                )
        else:
            self.audio_tower = None
            self.embed_audio = None

        # ---- Language model (vLLM optimised) ----
        with self._mark_language_model(vllm_config):
            self.language_model: Gemma4ForCausalLM = init_vllm_registered_model(
                vllm_config=vllm_config,
                hf_config=config.text_config,
                prefix=maybe_prefix(prefix, "language_model"),
                architectures=["Gemma4ForCausalLM"],
            )

            # Pre-allocate PLE buffer for CUDA graph compatibility.
            # Some variants have hidden_size_per_layer_input=None (no PLE).
            ple_dim = config.text_config.hidden_size_per_layer_input
            if ple_dim is not None:
                self.per_layer_embeddings = torch.zeros(
                    vllm_config.scheduler_config.max_num_batched_tokens,
                    config.text_config.num_hidden_layers,
                    ple_dim,
                    device=(self.language_model.model.embed_tokens.weight.device),
                    dtype=(self.language_model.model.embed_tokens.weight.dtype),
                )
            else:
                self.per_layer_embeddings = None

        self.make_empty_intermediate_tensors = (
            self.language_model.make_empty_intermediate_tensors
        )

        # --- Precompute full-attention layer indices for bidi clearing ---
        self._full_attn_layer_idxs: frozenset[int] = frozenset()
        text_config = config.text_config
        if getattr(text_config, "use_bidirectional_attention", None) == "vision":
            layer_types = getattr(text_config, "layer_types", None)
            if layer_types:
                self._full_attn_layer_idxs = frozenset(
                    i for i, lt in enumerate(layer_types) if lt != "sliding_attention"
                )

        # --- MixtureOfExperts delegation to language_model ---
        self.expert_weights = self.language_model.expert_weights
        self.moe_layers = self.language_model.moe_layers
        self.num_moe_layers = self.language_model.num_moe_layers
        self.num_logical_experts = self.language_model.num_logical_experts
        self.num_physical_experts = self.language_model.num_physical_experts
        self.num_local_physical_experts = self.language_model.num_local_physical_experts
        self.num_routed_experts = self.language_model.num_routed_experts
        self.num_expert_groups = self.language_model.num_expert_groups
        self.num_shared_experts = self.language_model.num_shared_experts
        self.num_redundant_experts = self.language_model.num_redundant_experts

    # ------------------------------------------------------------------ #
    # Input parsing
    # ------------------------------------------------------------------ #

    def _parse_and_validate_image_input(
        self, **kwargs: object
    ) -> Gemma4ImageInputs | None:
        pixel_values = kwargs.pop("pixel_values", None)
        pixel_position_ids = kwargs.pop("pixel_position_ids", None)
        image_embeds = kwargs.pop("image_embeds", None)
        assert image_embeds is None, "Gemma4 does not support image_embeds."
        if pixel_values is None:
            return None
        return Gemma4ImagePixelInputs(
            pixel_values=pixel_values,
            pixel_position_ids=pixel_position_ids,
        )

    def _parse_and_validate_audio_input(
        self, **kwargs: object
    ) -> Gemma4AudioInputs | None:
        input_features_padded = kwargs.pop("input_features_padded", None)
        if input_features_padded is None:
            return None
        input_features_mask = kwargs.pop("input_features_mask", None)
        if input_features_mask is None:
            return None
        return Gemma4AudioInputs(
            input_features_padded=input_features_padded,
            input_features_mask=input_features_mask,
        )

    def _parse_and_validate_video_input(
        self, **kwargs: object
    ) -> dict[str, torch.Tensor] | None:
        pixel_values_videos = kwargs.pop("pixel_values_videos", None)
        pixel_position_ids_videos = kwargs.pop("pixel_position_ids_videos", None)
        video_frame_counts = kwargs.pop("video_frame_counts", None)
        if pixel_values_videos is None:
            return None
        return {
            "pixel_values_videos": pixel_values_videos,
            "pixel_position_ids_videos": pixel_position_ids_videos,
            "video_frame_counts": video_frame_counts,
        }

    def _parse_and_validate_multimodal_inputs(
        self, **kwargs: object
    ) -> dict[str, Gemma4ImageInputs | Gemma4AudioInputs | Gemma4VideoInputs | None]:
        mm_input_by_modality = {}
        for input_key in list(kwargs):
            if (
                input_key in ("pixel_values", "image_embeds")
                and "image" not in mm_input_by_modality
            ):
                mm_input_by_modality["image"] = self._parse_and_validate_image_input(
                    **kwargs
                )
            if (
                input_key == "pixel_values_videos"
                and "video" not in mm_input_by_modality
            ):
                mm_input_by_modality["video"] = self._parse_and_validate_video_input(
                    **kwargs
                )
            if (
                input_key == "input_features_padded"
                and "audio" not in mm_input_by_modality
            ):
                mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(
                    **kwargs
                )
        return mm_input_by_modality

    @staticmethod
    def _encoder_chunk(
        patches_per_item: int,
        free_bytes: int,
        total_bytes: int,
        position_embedding_size: int,
    ) -> int:
        """Max chunk size whose F.one_hot transient fits in the budget.

        The dominant transient inside HF's ``Gemma4VisionPatchEmbedder.
        _position_embeddings`` is
        ``F.one_hot(clamped_positions, num_classes=position_embedding_size)``
        with shape ``(chunk, patches, 2, position_embedding_size)``,
        int64, plus its simultaneous cast to the position embedding
        table dtype. That, not the encoder residual stream, sets peak
        memory.
        """
        if patches_per_item <= 0:
            return 1
        # Half of currently-free, capped at 10% of total so we leave room
        # for the rest of profile_run / the subsequent encoder + pooler.
        budget = min(free_bytes // 2, total_bytes // 10)
        if budget <= 0:
            return 1
        # F.one_hot allocates (chunk, patches, 2, pos_emb_size) int64
        # (the inner 2 is the (x, y) coordinate axis, 8 is sizeof(int64)).
        # Outer 2x covers the int64 buffer and its concurrent bf16 cast
        # plus the matmul output that live alongside it at peak.
        cost = patches_per_item * 4 * position_embedding_size * 8
        return max(1, budget // cost) if cost > 0 else 1

    # ------------------------------------------------------------------ #
    # Image processing
    # ------------------------------------------------------------------ #

    def _process_image_input(
        self,
        image_input: Gemma4ImageInputs,
    ) -> list[torch.Tensor]:
        """Batch-encode images through the vision tower.

        Groups images by patch count (resolution bucket) so each
        encoder call processes a uniform-shape batch with no
        cross-resolution padding.  Pooling and projection are then
        applied over a single concatenated tensor for all images.
        """
        pixel_values = image_input["pixel_values"]
        pixel_position_ids = image_input["pixel_position_ids"]

        vt = self.vision_tower
        vision_cfg = self.config.vision_config
        pooling_k2 = vision_cfg.pooling_kernel_size**2

        # Concurrent requests with different image resolutions may
        # arrive as a list of per-image tensors, while same-resolution
        # batches may arrive as a stacked tensor.
        buckets: dict[int, list[tuple[int, torch.Tensor, torch.Tensor]]] = {}
        total_images = (
            len(pixel_values)
            if isinstance(pixel_values, list)
            else pixel_values.shape[0]
        )

        for idx in range(total_images):
            pv = pixel_values[idx]
            pp = pixel_position_ids[idx]
            buckets.setdefault(pv.shape[0], []).append((idx, pv, pp))

        # Encode each resolution bucket in memory-safe chunks. Re-read
        # free memory per bucket because the previous bucket's encoder
        # pass has already allocated activations we should account for.
        last_hidden_states_map: dict[int, torch.Tensor] = {}
        for patches, items in buckets.items():
            free, total = current_platform.mem_get_info()
            max_batch_size = min(
                len(items),
                self._encoder_chunk(
                    patches, free, total, vision_cfg.position_embedding_size
                ),
            )

            for chunk_idx in range(0, len(items), max_batch_size):
                chunk_items = items[chunk_idx : chunk_idx + max_batch_size]

                pv_tensor = torch.cat(
                    [item[1].unsqueeze(0) for item in chunk_items], dim=0
                )
                pp_tensor = torch.cat(
                    [item[2].unsqueeze(0) for item in chunk_items], dim=0
                )
                pad_tensor = (pp_tensor == -1).all(dim=-1)

                inputs_embeds = vt.patch_embedder(pv_tensor, pp_tensor, pad_tensor)
                encoder_outputs = vt.encoder(
                    inputs_embeds=inputs_embeds,
                    attention_mask=~pad_tensor,
                    pixel_position_ids=pp_tensor,
                )
                hidden_states = encoder_outputs.last_hidden_state

                for i, (orig_idx, _, _) in enumerate(chunk_items):
                    last_hidden_states_map[orig_idx] = hidden_states[i]

        # Pool per image to strip padding and reduce spatial resolution.
        all_valid_states: list[torch.Tensor] = [None] * total_images  # type: ignore[list-item]
        valid_lens = [0] * total_images

        for orig_idx in range(total_images):
            chunk_hidden = last_hidden_states_map[orig_idx]
            output_length = chunk_hidden.shape[0] // pooling_k2

            single_hidden = chunk_hidden.unsqueeze(0)
            single_pos_ids = pixel_position_ids[orig_idx].unsqueeze(0)
            padding_positions = (single_pos_ids == -1).all(dim=-1)

            pooled_states, valid_mask = vt.pooler(
                hidden_states=single_hidden,
                pixel_position_ids=single_pos_ids,
                padding_positions=padding_positions,
                output_length=output_length,
            )
            valid_states = pooled_states[valid_mask]

            if getattr(vt.config, "standardize", False):
                valid_states = (valid_states - vt.std_bias) * vt.std_scale

            all_valid_states[orig_idx] = valid_states
            valid_lens[orig_idx] = valid_states.shape[0]

        target_dtype = self.embed_vision.embedding_projection.weight.dtype

        # Project all images in a single batched call.
        flat_valid_states = torch.cat(all_valid_states, dim=0).to(target_dtype)
        flat_proj_embs = self.embed_vision(
            inputs_embeds=flat_valid_states.unsqueeze(0)
        ).squeeze(0)

        # Split back into per-image tensors (slicing returns views).
        per_image_embeddings: list[torch.Tensor] = []
        offset = 0
        for length in valid_lens:
            per_image_embeddings.append(flat_proj_embs[offset : offset + length])
            offset += length

        return per_image_embeddings

    # ------------------------------------------------------------------ #
    # Video processing (frames through vision tower)
    # ------------------------------------------------------------------ #

    def _process_video_input(
        self,
        video_input: dict[str, torch.Tensor],
    ) -> list[torch.Tensor]:
        """Batch-encode video frames through the vision tower.

        Gemma4 has no separate video tower; video frames are images at
        lower resolution (max_soft_tokens=70).  All frames across all
        videos in the batch are encoded together in chunks, then pooled
        and projected in a single batched call.

        Returns one concatenated embedding tensor per video (not per
        frame), matching the flat_from_sizes grouping that vLLM expects
        for embed_multimodal.
        """
        pixel_values = video_input["pixel_values_videos"]
        pixel_position_ids = video_input["pixel_position_ids_videos"]
        frame_counts = video_input["video_frame_counts"]

        vt = self.vision_tower
        vision_cfg = self.config.vision_config
        pooling_k2 = vision_cfg.pooling_kernel_size**2
        target_dtype = self.embed_vision.embedding_projection.weight.dtype

        if isinstance(frame_counts, torch.Tensor):
            fc_list = frame_counts.tolist()
        else:
            fc_list = list(frame_counts)

        total_frames = pixel_values.shape[0]
        free, total = current_platform.mem_get_info()
        max_batch_size = min(
            total_frames,
            self._encoder_chunk(
                pixel_values.shape[1],
                free,
                total,
                vision_cfg.position_embedding_size,
            ),
        )

        padding_positions = (pixel_position_ids == -1).all(dim=-1)

        # Encode frames in chunks bounded by _encoder_chunk.
        last_hidden_states_list: list[torch.Tensor] = []
        for i in range(0, total_frames, max_batch_size):
            pv_chunk = pixel_values[i : i + max_batch_size]
            pp_chunk = pixel_position_ids[i : i + max_batch_size]
            pad_chunk = padding_positions[i : i + max_batch_size]

            inputs_embeds = vt.patch_embedder(pv_chunk, pp_chunk, pad_chunk)
            encoder_outputs = vt.encoder(
                inputs_embeds=inputs_embeds,
                attention_mask=~pad_chunk,
                pixel_position_ids=pp_chunk,
            )
            last_hidden_states_list.append(encoder_outputs.last_hidden_state)

        last_hidden_states = torch.cat(last_hidden_states_list, dim=0)

        # Pool per frame to strip padding and reduce spatial resolution.
        output_length = pixel_values.shape[1] // pooling_k2
        all_frame_valid_states: list[torch.Tensor] = []
        frame_valid_lens: list[int] = []

        for i in range(total_frames):
            single_hidden = last_hidden_states[i].unsqueeze(0)
            single_pos_ids = pixel_position_ids[i].unsqueeze(0)
            single_pad_pos = padding_positions[i].unsqueeze(0)

            pooled_states, valid_mask = vt.pooler(
                hidden_states=single_hidden,
                pixel_position_ids=single_pos_ids,
                padding_positions=single_pad_pos,
                output_length=output_length,
            )
            valid_states = pooled_states[valid_mask]

            if getattr(vt.config, "standardize", False):
                valid_states = (valid_states - vt.std_bias) * vt.std_scale

            all_frame_valid_states.append(valid_states)
            frame_valid_lens.append(valid_states.shape[0])

        # Project all frames in a single batched call.
        flat_valid_states = torch.cat(all_frame_valid_states, dim=0).to(target_dtype)
        flat_proj_embs = self.embed_vision(
            inputs_embeds=flat_valid_states.unsqueeze(0)
        ).squeeze(0)

        # Regroup into per-video tensors (slicing returns views).
        per_video_embeddings: list[torch.Tensor] = []
        frame_idx = 0
        offset = 0
        for count in fc_list:
            video_tokens = sum(frame_valid_lens[frame_idx : frame_idx + count])
            per_video_embeddings.append(flat_proj_embs[offset : offset + video_tokens])
            offset += video_tokens
            frame_idx += count

        return per_video_embeddings

    # ------------------------------------------------------------------ #
    # Audio processing
    # ------------------------------------------------------------------ #

    def _process_audio_input(
        self,
        audio_input: Gemma4AudioInputs,
    ) -> list[torch.Tensor]:
        input_features = audio_input["input_features_padded"].squeeze(1)
        input_features_mask = audio_input["input_features_mask"].squeeze(1)

        # Run audio tower — mask uses standard HF convention
        # (True=valid, False=padding).
        audio_outputs = self.audio_tower(input_features, input_features_mask)
        if isinstance(audio_outputs, tuple):
            audio_encodings, audio_mask = audio_outputs
        else:
            audio_encodings = audio_outputs.last_hidden_state
            audio_mask = audio_outputs.attention_mask

        # Project into LM embedding space.
        audio_features = self.embed_audio(inputs_embeds=audio_encodings)

        # Strip padding per-batch element: only keep real (non-padding)
        # tokens. audio_mask is True for valid positions (HF convention).
        per_audio = []
        for enc, mask in zip(audio_features, audio_mask, strict=True):
            per_audio.append(enc[mask])  # [num_real, hidden_size]

        return per_audio

    # ------------------------------------------------------------------ #
    # MultiModalEmbeddings interface
    # ------------------------------------------------------------------ #

    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
        mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
        multimodal_embeddings: list[torch.Tensor] = []

        for modality, multimodal_input in mm_input_by_modality.items():
            if multimodal_input is None:
                continue
            if modality == "image":
                multimodal_embeddings.extend(
                    self._process_image_input(multimodal_input)
                )
            elif modality == "video":
                multimodal_embeddings.extend(
                    self._process_video_input(multimodal_input)
                )
            elif modality == "audio":
                multimodal_embeddings.extend(
                    self._process_audio_input(multimodal_input)
                )

        return multimodal_embeddings

    def embed_input_ids(
        self,
        input_ids: torch.Tensor,
        multimodal_embeddings: MultiModalEmbeddings | None = None,
        *,
        is_multimodal: torch.Tensor | None = None,
    ) -> torch.Tensor:
        # Cache per-layer embeddings (PLE) for the language model's
        # forward pass.  During profiling embed_input_ids is not called,
        # so the pre-allocated zeros are used instead.
        if self.per_layer_embeddings is not None:
            # Mask multimodal tokens (image/audio) to 0 for PLE
            # computation (using token_type_ids == 0 as text_mask).
            # Replicate this: map image token positions to token 0.
            if is_multimodal is not None:
                ple_input_ids = torch.where(
                    is_multimodal.to(input_ids.device, non_blocking=True),
                    torch.zeros_like(input_ids),
                    input_ids,
                )
            else:
                ple_input_ids = input_ids

            per_layer_inputs = self.language_model.model.get_per_layer_inputs(
                ple_input_ids
            )
            if per_layer_inputs is not None:
                per_layer_inputs = per_layer_inputs.reshape(
                    -1,
                    self.config.text_config.num_hidden_layers,
                    self.config.text_config.hidden_size_per_layer_input,
                )
                self.per_layer_embeddings[: per_layer_inputs.shape[0]].copy_(
                    per_layer_inputs
                )

        if multimodal_embeddings is None or is_multimodal is None:
            return super().embed_input_ids(input_ids)

        return super().embed_input_ids(
            input_ids,
            multimodal_embeddings=multimodal_embeddings,
            is_multimodal=is_multimodal,
        )

    # ------------------------------------------------------------------ #
    # Forward
    # ------------------------------------------------------------------ #

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

        # Select the pre-cached PLEs for this batch (None when PLE
        # is disabled for variants without PLE).
        per_layer_inputs = (
            self.per_layer_embeddings[: inputs_embeds.shape[0]]
            if self.per_layer_embeddings is not None and inputs_embeds is not None
            else None
        )

        # Gemma4 bidi: clear mm_prefix_range for full_attention layers.
        # Must run here (outside @support_torch_compile boundary) because
        # _run_decoder_layers is inside a compiled graph where Python
        # side effects are eliminated.
        self._clear_mm_prefix_for_full_attn_layers()

        hidden_states = self.language_model.model(
            input_ids,
            positions,
            per_layer_inputs=per_layer_inputs,
            intermediate_tensors=intermediate_tensors,
            inputs_embeds=inputs_embeds,
            **kwargs,
        )

        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor | None:
        return self.language_model.compute_logits(hidden_states)

    # ------------------------------------------------------------------ #
    # Bidirectional attention helpers
    # ------------------------------------------------------------------ #

    def _clear_mm_prefix_for_full_attn_layers(self) -> None:
        """Clear mm_prefix_range for non-sliding layers.

        Gemma4 with use_bidirectional_attention='vision' applies
        bidirectional attention only to sliding_attention layers.
        Full attention layers use plain causal masking.

        Uses _full_attn_layer_idxs (precomputed in __init__) for O(1)
        lookup instead of per-call regex parsing.
        """
        if not self._full_attn_layer_idxs:
            return

        from vllm.forward_context import get_forward_context

        attn_metadata = get_forward_context().attn_metadata
        if attn_metadata is None:
            return

        def _process(metadata_dict: dict) -> None:
            for layer_name, metadata in metadata_dict.items():
                if ".layers." not in layer_name:
                    continue
                try:
                    layer_idx = int(layer_name.split(".layers.")[1].split(".")[0])
                except (ValueError, IndexError):
                    continue
                if layer_idx in self._full_attn_layer_idxs:
                    if hasattr(metadata, "mm_prefix_range"):
                        metadata.mm_prefix_range = None
                    if hasattr(metadata, "mm_prefix_range_tensor"):
                        metadata.mm_prefix_range_tensor = None

        if isinstance(attn_metadata, list):
            for ub_metadata in attn_metadata:
                _process(ub_metadata)
        elif isinstance(attn_metadata, dict):
            _process(attn_metadata)

    # ------------------------------------------------------------------ #
    # Weight loading
    # ------------------------------------------------------------------ #

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        # Some checkpoints have vestigial embed_vision.embedding and
        # embed_audio.embedding weights from the Gemma3n architecture
        # that are not used by Gemma4's MultimodalEmbedder (which only
        # has embedding_projection + embedding_post_projection_norm).
        ignore_prefixes = [
            "embed_vision.embedding.",
            "embed_audio.embedding.",
        ]
        # Models without audio tower should skip
        # audio weights entirely.
        if self.audio_tower is None:
            ignore_prefixes.extend(
                [
                    "audio_tower.",
                    "embed_audio.",
                ]
            )
        loader = AutoWeightsLoader(
            self,
            ignore_unexpected_prefixes=ignore_prefixes,
        )
        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

    # ------------------------------------------------------------------ #
    # LoRA / multimodal mapping
    # ------------------------------------------------------------------ #

    def get_mm_mapping(self) -> MultiModelKeys:
        """Get the module prefix mapping for multimodal models."""
        connectors = ["embed_vision"]
        tower_models = ["vision_tower"]
        if self.audio_tower is not None:
            connectors.append("embed_audio")
            tower_models.append("audio_tower")

        return MultiModelKeys.from_string_field(
            language_model="language_model",
            connector=connectors,
            tower_model=tower_models,
        )

    @classmethod
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
        if modality == "image":
            return "<image_soft_token>"
        if modality == "audio":
            return "<audio_soft_token>"
        if modality == "video":
            return "<|video|>"
        raise ValueError(f"Unsupported modality: {modality}")