vllm.model_executor.models.moondream3 ¶

@dataclass
class DetectPointState:
    """Per-request state for detect/point generation.

    Tracks the current step in the state machine, pending embedding
    information for the next decode step, the object currently being
    constructed, and the accumulated results.
    """

    mode: Literal["detect", "point"]
    step: Literal["decode_x_or_stop", "decode_y", "decode_size"] = "decode_x_or_stop"

    # Embedding info for the *next* decode step (set in compute_logits,
    # consumed in forward to replace the token embedding).
    pending_embed_type: str | None = None  # "coord" or "size"
    pending_embed_coord: float | None = None
    pending_embed_w: float | None = None
    pending_embed_h: float | None = None

    # Current object being constructed.
    current_x: float | None = None
    current_y: float | None = None

    # Accumulated results.
    objects: list[dict] = field(default_factory=list)
    finished: bool = False
    max_objects: int = _DEFAULT_MAX_OBJECTS

class DetectPointStateManager:
    """Manages per-request detect/point state for the model."""

    def __init__(
        self,
        coord_id: int = 5,
        size_id: int = 6,
        eos_id: int = 0,
    ) -> None:
        self._states: dict[str, DetectPointState] = {}
        self.coord_id = coord_id
        self.size_id = size_id
        self.eos_id = eos_id

    def register_request(
        self,
        req_id: str,
        mode: Literal["detect", "point"],
        max_objects: int = _DEFAULT_MAX_OBJECTS,
    ) -> None:
        self._states[req_id] = DetectPointState(mode=mode, max_objects=max_objects)

    def get_state(self, req_id: str) -> DetectPointState | None:
        return self._states.get(req_id)

    def has_active_requests(self) -> bool:
        return bool(self._states)

    def remove_request(self, req_id: str) -> None:
        self._states.pop(req_id, None)

    def get_json_result(self, req_id: str) -> str | None:
        """Serialize the accumulated objects for a finished request."""
        state = self._states.get(req_id)
        if state is None:
            return None
        if state.mode == "detect":
            return json.dumps({"objects": state.objects})
        return json.dumps({"points": state.objects})

    def update_after_sample(self, req_id: str, sampled_token_id: int) -> None:
        """Transition state after a token is sampled."""
        state = self._states.get(req_id)
        if state is None:
            return

        tok = sampled_token_id

        if state.step == "decode_x_or_stop":
            if tok == self.eos_id:
                state.finished = True
            else:
                # coord token sampled → move to decode_y
                state.step = "decode_y"

        elif state.step == "decode_y":
            if state.mode == "point":
                # For point: y was decoded, object is complete.
                state.objects.append({"x": state.current_x, "y": state.current_y})
                state.step = "decode_x_or_stop"
            else:
                # For detect: y decoded, need size next.
                state.step = "decode_size"

        elif state.step == "decode_size":
            # Size decoded, compute bbox from center + size.
            x, y = state.current_x, state.current_y
            w, h = state.pending_embed_w, state.pending_embed_h
            state.objects.append(
                {
                    "x_min": x - w / 2,
                    "y_min": y - h / 2,
                    "x_max": x + w / 2,
                    "y_max": y + h / 2,
                }
            )
            state.step = "decode_x_or_stop"

        # Check max objects limit.
        if len(state.objects) >= state.max_objects:
            state.finished = True

class Moondream3Attention(nn.Module):
    """Decoder attention with RoPE and tau scaling.

    Moondream3 uses a tau attention mechanism that scales Q and V
    based on both token content and position.
    """

    def __init__(
        self,
        config: Moondream3TextConfig,
        layer_idx: int,
        cache_config=None,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ):
        super().__init__()
        self.hidden_size = config.dim
        self.num_heads = config.n_heads
        self.num_kv_heads = config.n_kv_heads
        self.head_dim = config.dim // config.n_heads

        tp_size = get_tensor_model_parallel_world_size()
        self.num_heads_per_partition = self.num_heads // tp_size
        self.num_kv_heads_per_partition = max(1, self.num_kv_heads // tp_size)

        self.qkv_proj = QKVParallelLinear(
            hidden_size=self.hidden_size,
            head_size=self.head_dim,
            total_num_heads=self.num_heads,
            total_num_kv_heads=self.num_kv_heads,
            bias=True,
            quant_config=quant_config,
            prefix=f"{prefix}.qkv",
        )

        self.out_proj = RowParallelLinear(
            input_size=self.hidden_size,
            output_size=self.hidden_size,
            bias=True,
            quant_config=quant_config,
            prefix=f"{prefix}.proj",
        )

        # Moondream uses 32-dim rotation out of 64-dim head (partial_rotary_factor=0.5)
        # HF Moondream uses non-interleaved RoPE (split by half)
        # In vLLM, is_neox_style=True means split by half (GPT-NeoX style)
        rope_parameters = {
            "rope_theta": config.rope_theta,
            "partial_rotary_factor": 32 / self.head_dim,  # 32/64 = 0.5
        }
        self.rotary_emb = get_rope(
            head_size=self.head_dim,
            max_position=config.max_context,
            rope_parameters=rope_parameters,
            is_neox_style=True,  # Moondream uses split-by-half (GPT-NeoX) style
        )

        self.scaling = self.head_dim**-0.5
        self.attn = Attention(
            num_heads=self.num_heads_per_partition,
            head_size=self.head_dim,
            scale=self.scaling,
            num_kv_heads=self.num_kv_heads_per_partition,
            cache_config=cache_config,
            quant_config=quant_config,
            prefix=f"{prefix}.attn",
        )

        # Tau scaling parameters for position-dependent attention
        # These are learned during training to modulate attention based on position
        # tau_wq and tau_wv need full qkv_dim for correct computation
        # Only heads are partitioned, qkv dimension is kept full for all-gather
        qkv_dim = self.hidden_size * 3  # Q + K + V dimension (full)
        self.tau_alpha = nn.Parameter(torch.zeros(self.num_heads_per_partition))
        self.tau_wq = nn.Parameter(torch.zeros(self.num_heads_per_partition, qkv_dim))
        self.tau_wv = nn.Parameter(torch.zeros(self.num_heads_per_partition, qkv_dim))
        self.tp_size = tp_size

        # Prefix-LM attention length: BOS (1) + vision patches (729) = 730
        self._prefix_attn_len = config.prefix_attn  # 730
        if self._prefix_attn_len > config.max_context:
            raise ValueError(
                "prefix_attn must be <= max_context, "
                f"got {self._prefix_attn_len} > {config.max_context}."
            )
        # Build once and slice per prefill call to avoid allocating an N x N
        # mask on every forward pass.
        prefill_mask = torch.tril(
            torch.ones(
                1,
                1,
                config.max_context,
                config.max_context,
                dtype=torch.bool,
            )
        )
        prefill_mask[:, :, : self._prefix_attn_len, : self._prefix_attn_len] = True
        self.register_buffer("_prefill_mask", prefill_mask, persistent=False)

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)

        q, k, v = qkv.split(
            [
                self.num_heads_per_partition * self.head_dim,
                self.num_kv_heads_per_partition * self.head_dim,
                self.num_kv_heads_per_partition * self.head_dim,
            ],
            dim=-1,
        )

        # Apply tau scaling to Q and V
        # Tau scaling has two components:
        # 1. Token-based: tok_q = tanh(gelu(qkv) @ tau_wq.T)
        # 2. Position-based: tau_pos = 1 + (sigmoid(alpha * log(pos+1)) - 0.5)
        # Final: tau = tok + tau_pos
        #
        # For TP, tau weights are sharded by head, but qkv_dim is kept full

        # Get full qkv for tau computation
        # With TP, reconstruct qkv in correct layout [q_full, k_full, v_full]
        # (all-gather would produce [q_0, k_0, v_0, q_1, k_1, v_1] - wrong)
        if self.tp_size > 1:
            from vllm.distributed import tensor_model_parallel_all_gather

            # All-gather once, then reconstruct [q_full, k_full, v_full].
            qkv_full_sharded = tensor_model_parallel_all_gather(qkv.contiguous())
            q_local_dim = q.shape[-1]
            kv_local_dim = k.shape[-1]
            qkv_full_sharded = qkv_full_sharded.view(
                qkv.shape[0],
                self.tp_size,
                q_local_dim + 2 * kv_local_dim,
            )
            q_full = qkv_full_sharded[:, :, :q_local_dim].reshape(qkv.shape[0], -1)
            k_full = qkv_full_sharded[
                :, :, q_local_dim : q_local_dim + kv_local_dim
            ].reshape(qkv.shape[0], -1)
            v_full = qkv_full_sharded[:, :, q_local_dim + kv_local_dim :].reshape(
                qkv.shape[0], -1
            )
            qkv_full = torch.cat([q_full, k_full, v_full], dim=-1).contiguous()
        else:
            qkv_full = qkv

        # Compute tau scaling factors matching HF implementation exactly:
        # tok_feat = gelu(qkv)
        # tok_q = tanh(tok_feat @ tau_wq.T)  # [num_tokens, num_heads]
        # tau_pos = 1 + (sigmoid(alpha * log(pos+1)) - 0.5)  # [num_heads, num_tokens]
        # tau = (tok_q.T + tau_pos).T  # [num_tokens, num_heads]
        num_tokens = qkv_full.shape[0]
        orig_dtype = q.dtype

        # Token-based component
        tok_feat = F.gelu(qkv_full)  # Apply GELU activation
        tok_q = torch.tanh(tok_feat @ self.tau_wq.t())  # [N, H_per_partition]
        tok_v = torch.tanh(tok_feat @ self.tau_wv.t())  # [N, H_per_partition]

        # Position-based component
        # tau_pos = 1 + (sigmoid(alpha * log(pos+1)) - 0.5)
        # positions is [num_tokens], need to compute for each head
        # tau_alpha: [num_heads_per_partition]
        pos_float = (positions.to(orig_dtype) + 1.0).clamp(min=1e-6)
        pos_log = pos_float.log()  # [num_tokens]
        # alpha[:, None] * pos_log[None, :] -> [num_heads, num_tokens]
        tau_pos = 1.0 + (
            torch.sigmoid(self.tau_alpha[:, None] * pos_log[None, :]) - 0.5
        )  # [H_per_partition, N]

        # Combine token and position components
        tau_q = (tok_q + tau_pos.t()).to(orig_dtype)  # [N, H_per_partition]
        tau_v = (tok_v + tau_pos.t()).to(orig_dtype)  # [N, H_per_partition]

        # Reshape q and v to apply per-head tau scaling
        q = q.view(num_tokens, self.num_heads_per_partition, self.head_dim)
        v = v.view(num_tokens, self.num_kv_heads_per_partition, self.head_dim)

        # Apply tau scaling
        q = q * tau_q.unsqueeze(-1)
        v = v * tau_v[:, : self.num_kv_heads_per_partition].unsqueeze(-1)

        # Reshape back
        q = q.view(num_tokens, -1)
        v = v.view(num_tokens, -1)

        q, k = self.rotary_emb(positions, q, k)

        # ---- SDPA prefill override ----
        # During prefill (num_tokens >= prefix_attn_len), use PyTorch SDPA
        # with an explicit bidirectional mask for the prefix-LM region.
        # This produces hidden states that closely match the HF reference
        # (which also uses SDPA for prefill).  Decode tokens and short
        # sequences fall through to the configured attention backend (e.g.
        # FLEX_ATTENTION) for KV-cache-aware decoding.
        H = self.num_heads_per_partition
        Hkv = self.num_kv_heads_per_partition
        D = self.head_dim
        P = self._prefix_attn_len  # 730

        if num_tokens > 1 and num_tokens >= P:
            q_4d = q.view(num_tokens, H, D).transpose(0, 1).unsqueeze(0)
            k_4d = k.view(num_tokens, Hkv, D).transpose(0, 1).unsqueeze(0)
            v_4d = v.view(num_tokens, Hkv, D).transpose(0, 1).unsqueeze(0)

            # Causal mask with bidirectional prefix region. Reuse the prebuilt
            # mask when possible; some profiling paths can exceed max_context.
            if num_tokens <= self._prefill_mask.shape[-1]:
                bool_mask = self._prefill_mask[:, :, :num_tokens, :num_tokens]
            else:
                bool_mask = torch.tril(
                    torch.ones(
                        1,
                        1,
                        num_tokens,
                        num_tokens,
                        dtype=torch.bool,
                        device=q.device,
                    )
                )
                bool_mask[:, :, :P, :P] = True

            attn_output = F.scaled_dot_product_attention(
                q_4d,
                k_4d,
                v_4d,
                attn_mask=bool_mask,
                scale=self.scaling,
            )
            attn_output = (
                attn_output.squeeze(0)
                .transpose(0, 1)
                .contiguous()
                .view(num_tokens, H * D)
            )

            # Still call self.attn to populate the KV cache
            _ = self.attn(q, k, v)
        else:
            attn_output = self.attn(q, k, v)

        output, _ = self.out_proj(attn_output)
        return output

@dataclass
class Moondream3Config:
    """Combined configuration for Moondream3 model."""

    text: Moondream3TextConfig
    vision: Moondream3VisionConfig
    region: Moondream3RegionConfig

    @classmethod
    def from_dict(cls, d: dict) -> "Moondream3Config":
        return cls(
            text=Moondream3TextConfig.from_dict(d),
            vision=Moondream3VisionConfig.from_dict(d),
            region=Moondream3RegionConfig.from_dict(d),
        )

class Moondream3DecoderLayer(nn.Module):
    """Decoder layer with attention + MLP/MoE."""

    def __init__(
        self,
        config: Moondream3TextConfig,
        cache_config=None,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ):
        super().__init__()
        layer_idx = extract_layer_index(prefix)
        self.layer_idx = layer_idx

        self.ln = nn.LayerNorm(config.dim, eps=1e-5, bias=True)

        self.attn = Moondream3Attention(
            config=config,
            layer_idx=layer_idx,
            cache_config=cache_config,
            quant_config=quant_config,
            prefix=f"{prefix}.attn",
        )

        # Use MoE for layers >= moe_start_layer, standard MLP otherwise
        if layer_idx >= config.moe_start_layer:
            self.mlp = Moondream3TextMoE(
                hidden_size=config.dim,
                expert_inner_dim=config.moe_expert_inner_dim,
                num_experts=config.moe_num_experts,
                experts_per_token=config.moe_experts_per_token,
                quant_config=quant_config,
                prefix=f"{prefix}.mlp",
            )
        else:
            self.mlp = Moondream3TextMLP(
                hidden_size=config.dim,
                intermediate_size=config.ff_dim,
                quant_config=quant_config,
                prefix=f"{prefix}.mlp",
            )

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        # Pre-norm architecture
        normed = self.ln(hidden_states)
        attn_out = self.attn(positions, normed)
        mlp_out = self.mlp(normed)
        hidden_states = hidden_states + attn_out + mlp_out
        return hidden_states

class Moondream3DummyInputsBuilder(BaseDummyInputsBuilder[Moondream3ProcessingInfo]):
    """Dummy inputs builder for profiling."""

    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
        return (
            "<|endoftext|><image><|md_reserved_0|>query<|md_reserved_1|>"
            "What is this image?<|md_reserved_2|>"
        )

    def get_dummy_mm_data(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
        mm_options: Mapping[str, BaseDummyOptions] | None = None,
        mm_processor_kwargs: Mapping[str, object] | None = None,
    ) -> MultiModalDataDict:
        num_images = mm_counts.get("image", 0)
        return {
            "image": self._get_dummy_images(
                width=378,
                height=378,
                num_images=num_images,
            )
        }

@MULTIMODAL_REGISTRY.register_processor(
    Moondream3MultiModalProcessor,
    info=Moondream3ProcessingInfo,
    dummy_inputs=Moondream3DummyInputsBuilder,
)
class Moondream3ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
    """Moondream3 multimodal model for causal language modeling.

    Moondream3 has four capabilities:

    - **query**: Visual QA.
    - **caption**: Image description.
    - **detect**: Object detection (bounding boxes).
    - **point**: Object pointing (x, y coordinates).

    All four capabilities are supported. Query and caption use standard
    autoregressive generation. Detect and point use a custom state machine
    that intercepts ``compute_logits`` and ``forward`` to decode
    coordinates from hidden states and feed Fourier-encoded coordinate
    embeddings back as the next input.

    Detect/point mode is activated by setting
    ``SamplingParams(extra_args={"moondream3_task": "detect"})``
    (or ``"point"``).
    """

    supports_multimodal = True
    packed_modules_mapping = {
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
    }

    def __init__(
        self,
        *,
        vllm_config: VllmConfig,
        prefix: str = "",
    ):
        super().__init__()

        hf_config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        cache_config = vllm_config.cache_config

        # Parse config from HuggingFace config
        config_dict = hf_config.config if hasattr(hf_config, "config") else {}

        self.config = Moondream3Config.from_dict(config_dict)

        with self._mark_tower_model(vllm_config, "image"):
            # Vision encoder
            self.vision = Moondream3VisionEncoder(
                config=self.config.vision,
                quant_config=quant_config,
                prefix=maybe_prefix(prefix, "vision"),
            )

            # Vision projection
            self.vision_proj = Moondream3VisionProjection(
                input_dim=self.config.vision.enc_dim,
                inner_dim=self.config.vision.proj_inner_dim,
                output_dim=self.config.text.dim,
                quant_config=quant_config,
                prefix=maybe_prefix(prefix, "vision_proj"),
            )

        with self._mark_language_model(vllm_config):
            # Text decoder
            self.text = Moondream3TextModel(
                config=self.config.text,
                cache_config=cache_config,
                quant_config=quant_config,
                prefix=maybe_prefix(prefix, "text"),
            )

            # LM head (with bias - Moondream3 has lm_head bias)
            self.lm_head = ParallelLMHead(
                self.config.text.vocab_size,
                self.config.text.dim,
                bias=True,
                quant_config=quant_config,
                prefix=maybe_prefix(prefix, "lm_head"),
            )

        # Region module for point/detect coordinate encoding/decoding.
        self.region = Moondream3RegionModule(
            config=self.config.region,
            prefix=maybe_prefix(prefix, "region"),
        )
        self.logits_processor = LogitsProcessor(self.config.text.vocab_size)
        self.make_empty_intermediate_tensors = self.text.make_empty_intermediate_tensors

        # Detect/point token IDs (from config, with defaults).
        self._coord_id = getattr(hf_config, "coord_token_id", 5)
        self._size_id = getattr(hf_config, "size_token_id", 6)
        self._eos_id = getattr(hf_config, "region_eos_token_id", 0)

        # Detect/point state management and per-step scratch buffers.
        # These buffers are populated by model hooks invoked by the v1 runner.
        self.detect_point_manager = DetectPointStateManager(
            coord_id=self._coord_id,
            size_id=self._size_id,
            eos_id=self._eos_id,
        )
        self._dp_row_states: list[DetectPointState | None] | None = None
        self._dp_embed_data: dict[int, dict] | None = None

    def on_new_request(
        self,
        *,
        req_id: str,
        sampling_params: object | None,
    ) -> None:
        """Register detect/point requests from per-request extra args."""
        if sampling_params is None:
            return
        extra = getattr(sampling_params, "extra_args", None) or {}
        dp_task = extra.get("moondream3_task")
        if dp_task == "detect":
            mode: Literal["detect", "point"] = "detect"
        elif dp_task == "point":
            mode = "point"
        else:
            return

        raw_max = extra.get("moondream3_max_objects", _DEFAULT_MAX_OBJECTS)
        try:
            max_obj = int(raw_max)
        except (TypeError, ValueError):
            raise ValueError(
                "moondream3_max_objects must be an integer, "
                f"got {type(raw_max).__name__}"
            ) from None
        if max_obj < 1:
            raise ValueError(f"moondream3_max_objects must be >= 1, got {max_obj}")
        self.detect_point_manager.register_request(req_id, mode, max_obj)

    def on_requests_finished(self, req_ids: Iterable[str]) -> None:
        for req_id in req_ids:
            self.detect_point_manager.remove_request(req_id)

    def on_before_model_forward(
        self,
        *,
        req_ids: list[str],
        logits_indices: torch.Tensor,
        device: torch.device,
    ) -> None:
        """Prepare pending coordinate/size embed replacements for forward."""
        dp_mgr = self.detect_point_manager
        if not dp_mgr.has_active_requests():
            self._dp_embed_data = None
            return

        pp = get_pp_group()
        if pp.world_size > 1:
            self._sync_dp_pending_embeds(req_ids=req_ids, device=device, pp=pp)

        dp_embed_data: dict[int, dict[str, Any]] = {}
        for i, rid in enumerate(req_ids):
            dp_st = dp_mgr.get_state(rid)
            if dp_st is None or dp_st.pending_embed_type is None:
                continue
            pos = int(logits_indices[i].item())
            if dp_st.pending_embed_type == "coord":
                dp_embed_data[pos] = {
                    "type": "coord",
                    "value": dp_st.pending_embed_coord,
                }
            elif dp_st.pending_embed_type == "size":
                dp_embed_data[pos] = {
                    "type": "size",
                    "w": dp_st.pending_embed_w,
                    "h": dp_st.pending_embed_h,
                }
        self._dp_embed_data = dp_embed_data or None

    def on_before_compute_logits(self, *, req_ids: list[str]) -> None:
        dp_mgr = self.detect_point_manager
        if not dp_mgr.has_active_requests():
            self._dp_row_states = None
            return
        self._dp_row_states = [dp_mgr.get_state(req_id) for req_id in req_ids]

    def on_after_sample(
        self,
        *,
        req_ids: list[str],
        sampled_token_ids: torch.Tensor,
    ) -> None:
        dp_mgr = self.detect_point_manager
        if not dp_mgr.has_active_requests():
            return
        for i, req_id in enumerate(req_ids):
            if dp_mgr.get_state(req_id) is None:
                continue
            tok = int(sampled_token_ids[i, 0].item())
            dp_mgr.update_after_sample(req_id, tok)

    def get_per_request_extra_output(
        self,
        *,
        req_ids: list[str],
    ) -> dict[str, dict[str, Any]] | None:
        """Return per-request final text overrides for active requests."""
        dp_mgr = self.detect_point_manager
        if not get_pp_group().is_last_rank or not dp_mgr.has_active_requests():
            return None

        per_request_extra: dict[str, dict[str, Any]] = {}
        for req_id in req_ids:
            state = dp_mgr.get_state(req_id)
            if state is None:
                continue
            json_str = dp_mgr.get_json_result(req_id)
            if json_str is not None:
                per_request_extra[req_id] = {"text_override": json_str}
        return per_request_extra or None

    def _sync_dp_pending_embeds(
        self,
        *,
        req_ids: list[str],
        device: torch.device,
        pp: Any,
    ) -> None:
        """Broadcast detect/point pending embed data across PP ranks."""
        num_reqs = len(req_ids)
        sync = torch.zeros(
            num_reqs,
            3,
            device=device,
            dtype=torch.float32,
        )

        dp_mgr = self.detect_point_manager
        if pp.is_last_rank:
            for i, req_id in enumerate(req_ids):
                st = dp_mgr.get_state(req_id)
                if st is None or st.pending_embed_type is None:
                    continue
                if st.pending_embed_type == "coord":
                    sync[i, 0] = 1.0
                    sync[i, 1] = st.pending_embed_coord or 0.0
                elif st.pending_embed_type == "size":
                    sync[i, 0] = 2.0
                    sync[i, 1] = st.pending_embed_w or 0.0
                    sync[i, 2] = st.pending_embed_h or 0.0

        torch.distributed.broadcast(
            sync,
            src=pp.last_rank,
            group=pp.device_group,
        )

        if not pp.is_last_rank:
            for i, req_id in enumerate(req_ids):
                st = dp_mgr.get_state(req_id)
                if st is None:
                    continue
                state_type = int(sync[i, 0].item())
                if state_type == 0:
                    st.pending_embed_type = None
                elif state_type == 1:
                    st.pending_embed_type = "coord"
                    st.pending_embed_coord = sync[i, 1].item()
                elif state_type == 2:
                    st.pending_embed_type = "size"
                    st.pending_embed_w = sync[i, 1].item()
                    st.pending_embed_h = sync[i, 2].item()

    @classmethod
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
        if modality == "image":
            return "<image>"
        return None

    def get_language_model(self) -> nn.Module:
        return self.text

    def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
        return num_image_tokens

    def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
        return num_vision_tokens

    def _split_pixel_values(
        self,
        pixel_values: object,
    ) -> list[torch.Tensor]:
        if isinstance(pixel_values, torch.Tensor):
            if pixel_values.dim() == 5:
                return [pv.contiguous() for pv in pixel_values]
            if pixel_values.dim() == 4:
                return [pixel_values.contiguous()]
            if pixel_values.dim() == 3:
                return [pixel_values.unsqueeze(0).contiguous()]
            raise ValueError(
                f"Unsupported pixel_values shape {tuple(pixel_values.shape)}."
            )

        if isinstance(pixel_values, (list, tuple)):
            tensors: list[torch.Tensor] = []
            for value in pixel_values:
                tensor_value = (
                    value if isinstance(value, torch.Tensor) else torch.as_tensor(value)
                )
                if tensor_value.dim() == 3:
                    tensor_value = tensor_value.unsqueeze(0)
                elif tensor_value.dim() != 4:
                    raise ValueError(
                        f"Unsupported pixel_values element shape "
                        f"{tuple(tensor_value.shape)}."
                    )
                tensors.append(tensor_value.contiguous())
            return tensors

        raise TypeError(
            "pixel_values must be a tensor or a sequence of tensors, "
            f"got {type(pixel_values)!r}."
        )

    def _split_tilings(
        self,
        tilings: object,
        expected: int,
    ) -> list[tuple[int, int] | None]:
        if tilings is None:
            return [None] * expected

        if isinstance(tilings, torch.Tensor):
            tiling_items = tilings.tolist()
        elif isinstance(tilings, (list, tuple)):
            tiling_items = list(tilings)
        else:
            raise TypeError(
                "tilings must be None, a tensor or a sequence of tuples, "
                f"got {type(tilings)!r}."
            )

        if len(tiling_items) != expected:
            raise ValueError(
                "Mismatch between the number of pixel_values entries "
                f"({expected}) and tilings ({len(tiling_items)})."
            )

        normalized: list[tuple[int, int] | None] = []
        for tiling in tiling_items:
            if tiling is None:
                normalized.append(None)
                continue
            if isinstance(tiling, torch.Tensor):
                tiling = tiling.tolist()
            if isinstance(tiling, (list, tuple)) and len(tiling) == 2:
                normalized.append((int(tiling[0]), int(tiling[1])))
            else:
                raise ValueError(
                    f"Each tiling entry must be a pair of integers, got {tiling!r}."
                )
        return normalized

    def _parse_image_inputs(self, **kwargs: object) -> list[Moondream3ImageInput]:
        pixel_values = kwargs.get("pixel_values")
        if pixel_values is None:
            return []

        pixel_values_list = self._split_pixel_values(pixel_values)
        tilings_list = self._split_tilings(
            kwargs.get("tilings"), len(pixel_values_list)
        )

        image_inputs: list[Moondream3ImageInput] = []
        for value, tiling in zip(pixel_values_list, tilings_list):
            if value.dim() != 4:
                raise ValueError(
                    f"Expected 4D tensor for crops, got {tuple(value.shape)}."
                )
            image_inputs.append(Moondream3ImageInput(pixel_values=value, tiling=tiling))
        return image_inputs

    def _encode_image_input(self, image_input: Moondream3ImageInput) -> torch.Tensor:
        pixel_values = image_input.pixel_values
        if pixel_values.dim() != 4:
            raise ValueError(
                f"Expected 4D tensor for crops, got {tuple(pixel_values.shape)}."
            )

        device = self.vision.patch_emb.weight.device
        dtype = self.vision.patch_emb.weight.dtype
        pixel_values = pixel_values.to(device=device, dtype=dtype)

        features = self.vision(pixel_values)

        # Grid size = crop_size / patch_size (e.g., 378 / 14 = 27)
        grid_size = self.config.vision.crop_size // self.config.vision.enc_patch_size
        enc_dim = self.config.vision.enc_dim
        global_features = features[0]

        if features.shape[0] > 1:
            if image_input.tiling is None:
                raise ValueError(
                    "Missing tiling metadata for multi-crop Moondream image."
                )
            local = features[1:].contiguous().view(-1, grid_size, grid_size, enc_dim)
            reconstructed = reconstruct_from_crops(
                local,
                image_input.tiling,
                overlap_margin=self.config.vision.overlap_margin,
                patch_size=1,
            )
        else:
            reconstructed = global_features.view(grid_size, grid_size, enc_dim)

        recon = reconstructed.permute(2, 0, 1).contiguous()
        # Mirror HF reference behavior: reconstructed local features are pooled
        # to enc_n_layers x enc_n_layers. For moondream3-preview this is 27x27.
        pooled_size = self.config.vision.enc_n_layers
        if pooled_size != grid_size:
            logger.warning_once(
                "Moondream3 pooled_size (%d) differs from crop grid (%d). "
                "Using enc_n_layers to match HF reference behavior.",
                pooled_size,
                grid_size,
            )
        recon = F.adaptive_avg_pool2d(recon, output_size=(pooled_size, pooled_size))
        recon = recon.permute(1, 2, 0).contiguous().view(-1, enc_dim)

        combined = torch.cat([global_features, recon], dim=-1).unsqueeze(0)
        projected = self.vision_proj(combined).squeeze(0)

        # Note: Vision embeddings are already synchronized across TP ranks
        # because the vision projection uses RowParallelLinear which performs
        # all-reduce internally, ensuring identical outputs on all ranks.

        return projected

    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
        """Generate 729 vision embeddings per image (27x27 patches)."""
        image_inputs = self._parse_image_inputs(**kwargs)
        if not image_inputs:
            return []

        return [self._encode_image_input(image_input) for image_input in image_inputs]

    def forward(
        self,
        input_ids: torch.Tensor | None,
        positions: torch.Tensor,
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
        **kwargs,
    ) -> torch.Tensor | IntermediateTensors:
        # Replace embeddings for detect/point decode positions with
        # Fourier-encoded coordinate/size embeddings from the region module.
        dp_embed_data = self._dp_embed_data
        if dp_embed_data and inputs_embeds is not None:
            for pos, info in dp_embed_data.items():
                if pos >= inputs_embeds.shape[0]:
                    continue
                if info["type"] == "coord":
                    coord_t = torch.tensor(
                        [[info["value"]]],
                        device=inputs_embeds.device,
                        dtype=inputs_embeds.dtype,
                    )
                    emb = self.region.encode_coordinate(coord_t)  # [1, dim]
                    inputs_embeds[pos] = emb.squeeze(0)
                elif info["type"] == "size":
                    emb = self.region.encode_size(
                        info["w"],
                        info["h"],
                        device=inputs_embeds.device,
                        dtype=inputs_embeds.dtype,
                    )  # [1, dim]
                    inputs_embeds[pos] = emb.squeeze(0)
            # Clear after use.
            self._dp_embed_data = None

        hidden_states = self.text(
            input_ids=input_ids,
            positions=positions,
            intermediate_tensors=intermediate_tensors,
            inputs_embeds=inputs_embeds,
        )
        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor | None:
        row_states = self._dp_row_states
        if row_states is None or not any(s is not None for s in row_states):
            # No detect/point requests — standard path.
            return self.logits_processor(self.lm_head, hidden_states)

        # Compute standard logits for ALL rows (needed for text requests
        # and for the continue/stop sparse check in decode_x_or_stop).
        logits = self.logits_processor(self.lm_head, hidden_states)
        if logits is None:
            return logits

        for i, state in enumerate(row_states):
            if state is None:
                continue
            h = hidden_states[i : i + 1]  # [1, dim]

            if state.step == "decode_x_or_stop":
                # Check continue/stop via sparse lm_head logits.
                coord_score = logits[i, self._coord_id].item()
                eos_score = logits[i, self._eos_id].item()

                if state.finished or eos_score > coord_score:
                    # Model wants to stop — force EOS.
                    logits[i] = float("-inf")
                    logits[i, self._eos_id] = 0.0
                else:
                    # Decode x coordinate from the *same* hidden state.
                    coord_logits = self.region.decode_coordinate(h)
                    x_bin = torch.argmax(coord_logits, dim=-1)
                    x_val = x_bin.float().item() / coord_logits.shape[-1]
                    state.current_x = x_val
                    state.pending_embed_type = "coord"
                    state.pending_embed_coord = x_val
                    # Force COORD_ID token.
                    logits[i] = float("-inf")
                    logits[i, self._coord_id] = 0.0

            elif state.step == "decode_y":
                coord_logits = self.region.decode_coordinate(h)
                y_bin = torch.argmax(coord_logits, dim=-1)
                y_val = y_bin.float().item() / coord_logits.shape[-1]
                state.current_y = y_val
                state.pending_embed_type = "coord"
                state.pending_embed_coord = y_val
                logits[i] = float("-inf")
                logits[i, self._coord_id] = 0.0

            elif state.step == "decode_size":
                w, h_val = self.region.decode_size(h)
                state.pending_embed_type = "size"
                state.pending_embed_w = w
                state.pending_embed_h = h_val
                logits[i] = float("-inf")
                logits[i, self._size_id] = 0.0

        return logits

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        """Load weights with remapping from HuggingFace format."""

        params_dict = dict(self.named_parameters())
        loaded_params: set[str] = set()

        # Get expert intermediate size for fc1 splitting

        for name, loaded_weight in weights:
            # Map from HF naming to vLLM naming
            # model.vision.* -> vision.*
            # model.text.* -> text.*
            if name.startswith("model."):
                name = name[6:]  # Remove "model." prefix

            # Specific name mappings
            # Vision projection: vision.proj_mlp.fc1 -> vision_proj.fc1
            name = name.replace("vision.proj_mlp.", "vision_proj.")

            # Text embedding: text.wte (no suffix) -> text.wte.weight
            if name == "text.wte":
                name = "text.wte.weight"

            # LM head: text.lm_head -> lm_head
            name = name.replace("text.lm_head.", "lm_head.")

            # Attention mapping
            name = name.replace(".attn.qkv.", ".attn.qkv_proj.")
            name = name.replace(".attn.proj.", ".attn.out_proj.")

            # Tau attention scaling weights
            # HF format: .attn.tau.alpha -> .attn.tau_alpha
            name = name.replace(".attn.tau.alpha", ".attn.tau_alpha")
            name = name.replace(".attn.tau.wq", ".attn.tau_wq")
            name = name.replace(".attn.tau.wv", ".attn.tau_wv")

            # MoE router mapping: mlp.router -> mlp.gate
            name = name.replace(".mlp.router.", ".mlp.gate.")

            # Handle MoE expert weights for layers 4+ with expert parallelism
            # fc1.weight: [n_experts, expert_inner_dim * 2, hidden_size] (gate+up)
            # fc2.weight: [n_experts, hidden_size, expert_inner_dim] (down)
            # Each GPU stores n_experts/tp_size experts
            # Note: Only 3D weights are MoE, 2D weights are standard MLP
            if ".mlp.fc1.weight" in name and loaded_weight.dim() == 3:
                from vllm.distributed import get_tensor_model_parallel_rank

                tp_size = get_tensor_model_parallel_world_size()
                tp_rank = get_tensor_model_parallel_rank()
                num_experts = loaded_weight.shape[0]
                experts_per_rank = num_experts // tp_size
                expert_start = tp_rank * experts_per_rank
                expert_end = expert_start + experts_per_rank
                # Shard by expert dimension
                loaded_weight = loaded_weight[expert_start:expert_end].contiguous()
                # Map to our custom MoE format: mlp.fc1_weight
                name = name.replace(".mlp.fc1.weight", ".mlp.fc1_weight")

            if ".mlp.fc2.weight" in name and loaded_weight.dim() == 3:
                from vllm.distributed import get_tensor_model_parallel_rank

                tp_size = get_tensor_model_parallel_world_size()
                tp_rank = get_tensor_model_parallel_rank()
                num_experts = loaded_weight.shape[0]
                experts_per_rank = num_experts // tp_size
                expert_start = tp_rank * experts_per_rank
                expert_end = expert_start + experts_per_rank
                # Shard by expert dimension
                loaded_weight = loaded_weight[expert_start:expert_end].contiguous()
                # Map to our custom MoE format: mlp.fc2_weight
                name = name.replace(".mlp.fc2.weight", ".mlp.fc2_weight")

            # Handle tau weights with tensor parallelism
            # tau_alpha: [num_heads] -> [num_heads/tp]
            # tau_wq: [num_heads, qkv_dim] -> [num_heads/tp, qkv_dim/tp]
            # tau_wv: [num_heads, qkv_dim] -> [num_heads/tp, qkv_dim/tp]
            if ".tau_alpha" in name:
                from vllm.distributed import get_tensor_model_parallel_rank

                tp_size = get_tensor_model_parallel_world_size()
                tp_rank = get_tensor_model_parallel_rank()
                num_heads = loaded_weight.shape[0]
                heads_per_partition = num_heads // tp_size
                start = tp_rank * heads_per_partition
                end = start + heads_per_partition
                loaded_weight = loaded_weight[start:end].contiguous()

            if ".tau_wq" in name or ".tau_wv" in name:
                from vllm.distributed import get_tensor_model_parallel_rank

                tp_size = get_tensor_model_parallel_world_size()
                tp_rank = get_tensor_model_parallel_rank()
                num_heads, qkv_dim = loaded_weight.shape
                heads_per_partition = num_heads // tp_size
                # Only shard by head dimension, keep full qkv_dim for all-gather
                head_start = tp_rank * heads_per_partition
                head_end = head_start + heads_per_partition
                loaded_weight = loaded_weight[head_start:head_end, :].contiguous()

            if name in params_dict:
                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                weight_loader(param, loaded_weight)
                loaded_params.add(name)

        return loaded_params