| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Inference-only Moondream3 model implementation.""" |
| 4 | |
| 5 | from collections.abc import Iterable, Mapping |
| 6 | from dataclasses import dataclass |
| 7 | from functools import cached_property |
| 8 | from itertools import islice |
| 9 | |
| 10 | import torch |
| 11 | import torch.nn as nn |
| 12 | import torch.nn.functional as F |
| 13 | from transformers import BatchFeature |
| 14 | |
| 15 | from vllm.config import VllmConfig |
| 16 | from vllm.config.multimodal import MultiModalDummyOptions |
| 17 | from vllm.distributed import ( |
| 18 | get_pp_group, |
| 19 | get_tensor_model_parallel_rank, |
| 20 | get_tensor_model_parallel_world_size, |
| 21 | tensor_model_parallel_all_gather, |
| 22 | tensor_model_parallel_all_reduce, |
| 23 | ) |
| 24 | from vllm.inputs import MultiModalDataDict |
| 25 | from vllm.logger import init_logger |
| 26 | from vllm.model_executor.layers.activation import get_act_fn |
| 27 | from vllm.model_executor.layers.attention import Attention |
| 28 | from vllm.model_executor.layers.attention.mm_encoder_attention import ( |
| 29 | MMEncoderAttention, |
| 30 | ) |
| 31 | from vllm.model_executor.layers.fused_moe import ( |
| 32 | GateLinear, |
| 33 | MoEActivation, |
| 34 | fused_experts, |
| 35 | ) |
| 36 | from vllm.model_executor.layers.fused_moe.config import biased_moe_quant_config |
| 37 | from vllm.model_executor.layers.linear import ( |
| 38 | ColumnParallelLinear, |
| 39 | QKVParallelLinear, |
| 40 | RowParallelLinear, |
| 41 | ) |
| 42 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 43 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 44 | from vllm.model_executor.layers.rotary_embedding import get_rope |
| 45 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 46 | ParallelLMHead, |
| 47 | VocabParallelEmbedding, |
| 48 | ) |
| 49 | from vllm.model_executor.model_loader.weight_utils import default_weight_loader |
| 50 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 51 | from vllm.multimodal.inputs import ( |
| 52 | MultiModalFieldConfig, |
| 53 | MultiModalKwargsItem, |
| 54 | MultiModalKwargsItems, |
| 55 | ) |
| 56 | from vllm.multimodal.parse import ImageSize, MultiModalDataItems |
| 57 | from vllm.multimodal.processing import ( |
| 58 | BaseDummyInputsBuilder, |
| 59 | BaseMultiModalProcessor, |
| 60 | BaseProcessingInfo, |
| 61 | PromptReplacement, |
| 62 | PromptUpdate, |
| 63 | PromptUpdateDetails, |
| 64 | cached_encode, |
| 65 | ) |
| 66 | from vllm.sequence import IntermediateTensors |
| 67 | from vllm.transformers_utils.configs.moondream3 import ( |
| 68 | Moondream3Config, |
| 69 | Moondream3TextConfig, |
| 70 | Moondream3VisionConfig, |
| 71 | ) |
| 72 | from vllm.transformers_utils.processors.moondream3 import Moondream3Processor |
| 73 | |
| 74 | from .interfaces import ( |
| 75 | MultiModalEmbeddings, |
| 76 | SupportsMultiModal, |
| 77 | SupportsPP, |
| 78 | ) |
| 79 | from .utils import ( |
| 80 | extract_layer_index, |
| 81 | make_empty_intermediate_tensors_factory, |
| 82 | make_layers, |
| 83 | maybe_prefix, |
| 84 | ) |
| 85 | |
| 86 | logger = init_logger(__name__) |
| 87 | |
| 88 | |
| 89 | # ============================================================================ |
| 90 | # Image Processing Utilities |
| 91 | # ============================================================================ |
| 92 | |
| 93 | |
| 94 | def reconstruct_from_crops( |
| 95 | crops: torch.Tensor, |
| 96 | tiling: tuple[int, int], |
| 97 | overlap_margin: int, |
| 98 | patch_size: int = 14, |
| 99 | ) -> torch.Tensor: |
| 100 | """Reconstruct features from overlapping crops.""" |
| 101 | tiling_h, tiling_w = tiling |
| 102 | crop_height, crop_width = crops[0].shape[:2] |
| 103 | margin_pixels = overlap_margin * patch_size |
| 104 | |
| 105 | output_h = (crop_height - 2 * margin_pixels) * tiling_h + 2 * margin_pixels |
| 106 | output_w = (crop_width - 2 * margin_pixels) * tiling_w + 2 * margin_pixels |
| 107 | |
| 108 | reconstructed = torch.zeros( |
| 109 | (output_h, output_w, crops[0].shape[2]), |
| 110 | device=crops[0].device, |
| 111 | dtype=crops[0].dtype, |
| 112 | ) |
| 113 | |
| 114 | for i, crop in enumerate(crops): |
| 115 | tile_y = i // tiling_w |
| 116 | tile_x = i % tiling_w |
| 117 | |
| 118 | x_start = 0 if tile_x == 0 else margin_pixels |
| 119 | x_end = crop_width if tile_x == tiling_w - 1 else crop_width - margin_pixels |
| 120 | y_start = 0 if tile_y == 0 else margin_pixels |
| 121 | y_end = crop_height if tile_y == tiling_h - 1 else crop_height - margin_pixels |
| 122 | |
| 123 | out_x = tile_x * (crop_width - 2 * margin_pixels) |
| 124 | out_y = tile_y * (crop_height - 2 * margin_pixels) |
| 125 | |
| 126 | reconstructed[ |
| 127 | out_y + y_start : out_y + y_end, out_x + x_start : out_x + x_end |
| 128 | ] = crop[y_start:y_end, x_start:x_end] |
| 129 | |
| 130 | return reconstructed |
| 131 | |
| 132 | |
| 133 | # ============================================================================ |
| 134 | # Vision Encoder Components |
| 135 | # ============================================================================ |
| 136 | |
| 137 | |
| 138 | class Moondream3VisionMLP(nn.Module): |
| 139 | """MLP for vision encoder blocks.""" |
| 140 | |
| 141 | def __init__( |
| 142 | self, |
| 143 | hidden_size: int, |
| 144 | intermediate_size: int, |
| 145 | quant_config: QuantizationConfig | None = None, |
| 146 | prefix: str = "", |
| 147 | ): |
| 148 | super().__init__() |
| 149 | self.fc1 = ColumnParallelLinear( |
| 150 | hidden_size, |
| 151 | intermediate_size, |
| 152 | bias=True, |
| 153 | quant_config=quant_config, |
| 154 | prefix=f"{prefix}.fc1", |
| 155 | ) |
| 156 | self.act = get_act_fn("gelu_pytorch_tanh") |
| 157 | self.fc2 = RowParallelLinear( |
| 158 | intermediate_size, |
| 159 | hidden_size, |
| 160 | bias=True, |
| 161 | quant_config=quant_config, |
| 162 | prefix=f"{prefix}.fc2", |
| 163 | ) |
| 164 | |
| 165 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 166 | x, _ = self.fc1(x) |
| 167 | x = self.act(x) |
| 168 | x, _ = self.fc2(x) |
| 169 | return x |
| 170 | |
| 171 | |
| 172 | class Moondream3VisionAttention(nn.Module): |
| 173 | """Self-attention for vision encoder (bidirectional).""" |
| 174 | |
| 175 | def __init__( |
| 176 | self, |
| 177 | hidden_size: int, |
| 178 | num_heads: int, |
| 179 | quant_config: QuantizationConfig | None = None, |
| 180 | prefix: str = "", |
| 181 | ): |
| 182 | super().__init__() |
| 183 | self.hidden_size = hidden_size |
| 184 | self.num_heads = num_heads |
| 185 | self.head_dim = hidden_size // num_heads |
| 186 | |
| 187 | self.qkv_proj = QKVParallelLinear( |
| 188 | hidden_size=hidden_size, |
| 189 | head_size=self.head_dim, |
| 190 | total_num_heads=num_heads, |
| 191 | bias=True, |
| 192 | quant_config=quant_config, |
| 193 | prefix=f"{prefix}.qkv_proj", |
| 194 | ) |
| 195 | self.out_proj = RowParallelLinear( |
| 196 | input_size=hidden_size, |
| 197 | output_size=hidden_size, |
| 198 | bias=True, |
| 199 | quant_config=quant_config, |
| 200 | prefix=f"{prefix}.out_proj", |
| 201 | ) |
| 202 | |
| 203 | tp_size = get_tensor_model_parallel_world_size() |
| 204 | self.num_heads_per_partition = num_heads // tp_size |
| 205 | |
| 206 | self.attn = MMEncoderAttention( |
| 207 | num_heads=self.num_heads_per_partition, |
| 208 | head_size=self.head_dim, |
| 209 | scale=self.head_dim**-0.5, |
| 210 | prefix=f"{prefix}.attn", |
| 211 | ) |
| 212 | |
| 213 | def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 214 | qkv, _ = self.qkv_proj(hidden_states) |
| 215 | q, k, v = qkv.chunk(3, dim=-1) |
| 216 | out = self.attn(q, k, v) |
| 217 | out, _ = self.out_proj(out) |
| 218 | return out |
| 219 | |
| 220 | |
| 221 | class Moondream3VisionBlock(nn.Module): |
| 222 | """Transformer block for vision encoder.""" |
| 223 | |
| 224 | def __init__( |
| 225 | self, |
| 226 | hidden_size: int, |
| 227 | intermediate_size: int, |
| 228 | num_heads: int, |
| 229 | quant_config: QuantizationConfig | None = None, |
| 230 | prefix: str = "", |
| 231 | ): |
| 232 | super().__init__() |
| 233 | self.ln1 = nn.LayerNorm(hidden_size, eps=1e-5) |
| 234 | self.attn = Moondream3VisionAttention( |
| 235 | hidden_size=hidden_size, |
| 236 | num_heads=num_heads, |
| 237 | quant_config=quant_config, |
| 238 | prefix=f"{prefix}.attn", |
| 239 | ) |
| 240 | self.ln2 = nn.LayerNorm(hidden_size, eps=1e-5) |
| 241 | self.mlp = Moondream3VisionMLP( |
| 242 | hidden_size=hidden_size, |
| 243 | intermediate_size=intermediate_size, |
| 244 | quant_config=quant_config, |
| 245 | prefix=f"{prefix}.mlp", |
| 246 | ) |
| 247 | |
| 248 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 249 | x = x + self.attn(self.ln1(x)) |
| 250 | x = x + self.mlp(self.ln2(x)) |
| 251 | return x |
| 252 | |
| 253 | |
| 254 | class Moondream3VisionEncoder(nn.Module): |
| 255 | """Vision encoder (SigLIP-style ViT).""" |
| 256 | |
| 257 | def __init__( |
| 258 | self, |
| 259 | config: Moondream3VisionConfig, |
| 260 | quant_config: QuantizationConfig | None = None, |
| 261 | prefix: str = "", |
| 262 | ): |
| 263 | super().__init__() |
| 264 | self.config = config |
| 265 | |
| 266 | # Patch embedding |
| 267 | self.patch_emb = nn.Linear( |
| 268 | config.enc_patch_size * config.enc_patch_size * 3, |
| 269 | config.enc_dim, |
| 270 | bias=True, |
| 271 | ) |
| 272 | |
| 273 | # Position embeddings (27x27 = 729 patches for 378x378 / 14) |
| 274 | num_patches = (config.crop_size // config.enc_patch_size) ** 2 |
| 275 | self.pos_emb = nn.Parameter(torch.zeros(1, num_patches, config.enc_dim)) |
| 276 | |
| 277 | # Transformer blocks |
| 278 | self.blocks = nn.ModuleList( |
| 279 | [ |
| 280 | Moondream3VisionBlock( |
| 281 | hidden_size=config.enc_dim, |
| 282 | intermediate_size=config.enc_ff_dim, |
| 283 | num_heads=config.enc_n_heads, |
| 284 | quant_config=quant_config, |
| 285 | prefix=f"{prefix}.blocks.{i}", |
| 286 | ) |
| 287 | for i in range(config.enc_n_layers) |
| 288 | ] |
| 289 | ) |
| 290 | |
| 291 | self.post_ln = nn.LayerNorm(config.enc_dim, eps=1e-5) |
| 292 | |
| 293 | def create_patches(self, images: torch.Tensor) -> torch.Tensor: |
| 294 | """Convert images to patch embeddings. |
| 295 | |
| 296 | Args: |
| 297 | images: (batch, channels, height, width) |
| 298 | |
| 299 | Returns: |
| 300 | patches: (batch, num_patches, patch_dim) |
| 301 | |
| 302 | """ |
| 303 | patch_size = self.config.enc_patch_size |
| 304 | batch, channels, height, width = images.shape |
| 305 | patches_h = height // patch_size |
| 306 | patches_w = width // patch_size |
| 307 | |
| 308 | # Unfold into patches |
| 309 | patches = images.unfold(2, patch_size, patch_size).unfold( |
| 310 | 3, patch_size, patch_size |
| 311 | ) |
| 312 | # (batch, channels, patches_h, patches_w, patch_size, patch_size) |
| 313 | patches = patches.permute(0, 2, 3, 1, 4, 5).contiguous() |
| 314 | # (batch, patches_h, patches_w, channels, patch_size, patch_size) |
| 315 | patches = patches.view(batch, patches_h * patches_w, -1) |
| 316 | # (batch, num_patches, channels * patch_size * patch_size) |
| 317 | |
| 318 | return patches |
| 319 | |
| 320 | def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: |
| 321 | """Encode images. |
| 322 | |
| 323 | Args: |
| 324 | pixel_values: (batch, channels, height, width) |
| 325 | |
| 326 | Returns: |
| 327 | features: (batch, num_patches, hidden_size) |
| 328 | |
| 329 | """ |
| 330 | # Create patches and embed |
| 331 | patches = self.create_patches(pixel_values) |
| 332 | x = self.patch_emb(patches) |
| 333 | |
| 334 | # Add position embeddings |
| 335 | x = x + self.pos_emb |
| 336 | |
| 337 | # Apply transformer blocks |
| 338 | for block in self.blocks: |
| 339 | x = block(x) |
| 340 | |
| 341 | # Final layer norm |
| 342 | x = self.post_ln(x) |
| 343 | |
| 344 | return x |
| 345 | |
| 346 | |
| 347 | class Moondream3VisionProjection(nn.Module): |
| 348 | """Projects vision features to text embedding dimension.""" |
| 349 | |
| 350 | def __init__( |
| 351 | self, |
| 352 | input_dim: int, |
| 353 | inner_dim: int, |
| 354 | output_dim: int, |
| 355 | quant_config: QuantizationConfig | None = None, |
| 356 | prefix: str = "", |
| 357 | ): |
| 358 | super().__init__() |
| 359 | # Input is concatenated global and local features (2 * input_dim) |
| 360 | self.fc1 = ColumnParallelLinear( |
| 361 | input_dim * 2, |
| 362 | inner_dim, |
| 363 | bias=True, |
| 364 | quant_config=quant_config, |
| 365 | prefix=f"{prefix}.fc1", |
| 366 | ) |
| 367 | self.act = get_act_fn("gelu_pytorch_tanh") |
| 368 | self.fc2 = RowParallelLinear( |
| 369 | inner_dim, |
| 370 | output_dim, |
| 371 | bias=True, |
| 372 | quant_config=quant_config, |
| 373 | prefix=f"{prefix}.fc2", |
| 374 | ) |
| 375 | |
| 376 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 377 | x, _ = self.fc1(x) |
| 378 | x = self.act(x) |
| 379 | x, _ = self.fc2(x) |
| 380 | return x |
| 381 | |
| 382 | |
| 383 | # ============================================================================ |
| 384 | # Text Decoder Components |
| 385 | # ============================================================================ |
| 386 | |
| 387 | |
| 388 | class Moondream3TextMLP(nn.Module): |
| 389 | """Standard MLP for non-MoE layers (layers 0-3).""" |
| 390 | |
| 391 | def __init__( |
| 392 | self, |
| 393 | hidden_size: int, |
| 394 | intermediate_size: int, |
| 395 | quant_config: QuantizationConfig | None = None, |
| 396 | prefix: str = "", |
| 397 | ): |
| 398 | super().__init__() |
| 399 | self.fc1 = ColumnParallelLinear( |
| 400 | hidden_size, |
| 401 | intermediate_size, |
| 402 | bias=True, |
| 403 | quant_config=quant_config, |
| 404 | prefix=f"{prefix}.fc1", |
| 405 | ) |
| 406 | self.act = get_act_fn("gelu_pytorch_tanh") |
| 407 | self.fc2 = RowParallelLinear( |
| 408 | intermediate_size, |
| 409 | hidden_size, |
| 410 | bias=True, |
| 411 | quant_config=quant_config, |
| 412 | prefix=f"{prefix}.fc2", |
| 413 | ) |
| 414 | |
| 415 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 416 | x, _ = self.fc1(x) |
| 417 | x = self.act(x) |
| 418 | x, _ = self.fc2(x) |
| 419 | return x |
| 420 | |
| 421 | |
| 422 | class Moondream3TextMoE(nn.Module): |
| 423 | """Mixture of Experts layer for layers 4+ with expert parallelism. |
| 424 | |
| 425 | Moondream3 uses a custom GeGLU activation: gelu(h) * (g + 1) |
| 426 | where fc1 outputs [gate, up] and the activation is gelu(gate) * (up + 1). |
| 427 | |
| 428 | Uses expert parallelism where each GPU stores num_experts/tp_size experts. |
| 429 | Routing and communication handled via all-to-all or replicated computation. |
| 430 | |
| 431 | Checkpoint format: |
| 432 | - fc1.weight: [num_experts, expert_inner_dim * 2, hidden_size] (gate+up) |
| 433 | - fc2.weight: [num_experts, hidden_size, expert_inner_dim] (down) |
| 434 | - router.weight: [num_experts, hidden_size] |
| 435 | - router.bias: [num_experts] |
| 436 | """ |
| 437 | |
| 438 | def __init__( |
| 439 | self, |
| 440 | hidden_size: int, |
| 441 | expert_inner_dim: int, |
| 442 | num_experts: int, |
| 443 | experts_per_token: int, |
| 444 | quant_config: QuantizationConfig | None = None, |
| 445 | prefix: str = "", |
| 446 | ): |
| 447 | super().__init__() |
| 448 | self.hidden_size = hidden_size |
| 449 | self.expert_inner_dim = expert_inner_dim |
| 450 | self.num_experts = num_experts |
| 451 | self.experts_per_token = experts_per_token |
| 452 | |
| 453 | # Expert parallelism: each GPU stores a subset of experts |
| 454 | self.tp_size = get_tensor_model_parallel_world_size() |
| 455 | self.experts_per_rank = num_experts // self.tp_size |
| 456 | self.num_local_experts = self.experts_per_rank |
| 457 | |
| 458 | # Router (gate) - use ReplicatedLinear for compatibility |
| 459 | self.gate = GateLinear( |
| 460 | hidden_size, |
| 461 | num_experts, |
| 462 | bias=True, |
| 463 | prefix=f"{prefix}.gate", |
| 464 | ) |
| 465 | |
| 466 | # Local expert weights (only store experts_per_rank experts) |
| 467 | # fc1: [experts_per_rank, expert_inner_dim * 2, hidden_size] |
| 468 | # fc2: [experts_per_rank, hidden_size, expert_inner_dim] |
| 469 | self.fc1_weight = nn.Parameter( |
| 470 | torch.empty(self.num_local_experts, expert_inner_dim * 2, hidden_size) |
| 471 | ) |
| 472 | self.fc2_weight = nn.Parameter( |
| 473 | torch.empty(self.num_local_experts, hidden_size, expert_inner_dim) |
| 474 | ) |
| 475 | self._use_fused_moe = True |
| 476 | |
| 477 | local_expert_start = get_tensor_model_parallel_rank() * self.experts_per_rank |
| 478 | expert_map = torch.full((num_experts,), -1, dtype=torch.int32) |
| 479 | expert_map[local_expert_start : local_expert_start + self.num_local_experts] = ( |
| 480 | torch.arange(self.num_local_experts, dtype=torch.int32) |
| 481 | ) |
| 482 | self.register_buffer("_expert_map", expert_map, persistent=False) |
| 483 | |
| 484 | # Preserve Moondream3's exact GeGLU variant (gelu(h) * (g + 1)) by |
| 485 | # adding +1 bias to the second half of the fused fc1 activations. |
| 486 | fused_w1_bias = torch.zeros( |
| 487 | self.num_local_experts, |
| 488 | expert_inner_dim * 2, |
| 489 | dtype=torch.float32, |
| 490 | ) |
| 491 | fused_w1_bias[:, expert_inner_dim:] = 1.0 |
| 492 | self.register_buffer("_fused_w1_bias", fused_w1_bias, persistent=False) |
| 493 | |
| 494 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 495 | """Forward pass with expert parallelism and custom GeGLU activation.""" |
| 496 | # Get router logits and compute top-k |
| 497 | router_logits, _ = self.gate(x) # [num_tokens, num_experts] |
| 498 | topk_logits, topk_ids = torch.topk( |
| 499 | router_logits, self.experts_per_token, dim=-1 |
| 500 | ) |
| 501 | # Softmax over selected experts |
| 502 | topk_weights = F.softmax(topk_logits, dim=-1, dtype=torch.float32).to(x.dtype) |
| 503 | |
| 504 | out = None |
| 505 | if self._use_fused_moe and x.is_cuda: |
| 506 | # Only the expert computation may fall back. The all-reduce below |
| 507 | # is deliberately outside this `try`: if it were inside and raised |
| 508 | # on a single rank, that rank would run the fallback loop and issue |
| 509 | # a *second* all-reduce while its peers issued one, desyncing the |
| 510 | # TP group and hanging it for good. |
| 511 | try: |
| 512 | out = fused_experts( |
| 513 | hidden_states=x.contiguous(), |
| 514 | w1=self.fc1_weight, |
| 515 | w2=self.fc2_weight, |
| 516 | topk_weights=topk_weights.contiguous(), |
| 517 | topk_ids=topk_ids.contiguous(), |
| 518 | activation=MoEActivation.GELU, |
| 519 | global_num_experts=self.num_experts, |
| 520 | expert_map=self._expert_map, |
| 521 | quant_config=biased_moe_quant_config(self._fused_w1_bias, None), |
| 522 | ) |
| 523 | except (NotImplementedError, RuntimeError) as exc: |
| 524 | self._use_fused_moe = False |
| 525 | logger.warning_once( |
| 526 | "Disabling fused Moondream3 MoE path and falling back to " |
| 527 | "the Python expert loop: %s", |
| 528 | str(exc), |
| 529 | ) |
| 530 | |
| 531 | if out is None: |
| 532 | out = self._expert_loop(x, topk_weights, topk_ids) |
| 533 | |
| 534 | # Combine each rank's partial expert outputs. Reached exactly once per |
| 535 | # forward, on every rank, whichever expert path ran above. |
| 536 | return tensor_model_parallel_all_reduce(out) |
| 537 | |
| 538 | def _expert_loop( |
| 539 | self, |
| 540 | x: torch.Tensor, |
| 541 | topk_weights: torch.Tensor, |
| 542 | topk_ids: torch.Tensor, |
| 543 | ) -> torch.Tensor: |
| 544 | """Fallback for environments where fused kernels are unavailable. |
| 545 | |
| 546 | Returns this rank's partial output; the caller performs the all-reduce. |
| 547 | """ |
| 548 | tp_rank = get_tensor_model_parallel_rank() |
| 549 | # Compute local expert range |
| 550 | local_expert_start = tp_rank * self.experts_per_rank |
| 551 | |
| 552 | out = x.new_zeros(x.shape) |
| 553 | |
| 554 | for local_expert_idx in range(self.num_local_experts): |
| 555 | global_expert_id = local_expert_start + local_expert_idx |
| 556 | |
| 557 | # Find tokens assigned to this expert |
| 558 | token_pos, which_k = (topk_ids == global_expert_id).nonzero(as_tuple=True) |
| 559 | if token_pos.numel() == 0: |
| 560 | continue |
| 561 | |
| 562 | # Get tokens and their routing weights |
| 563 | x_tok = x.index_select(0, token_pos) # [n_tokens, hidden_size] |
| 564 | gate_tok = topk_weights[token_pos, which_k] # [n_tokens] |
| 565 | |
| 566 | # fc1: [expert_inner_dim * 2, hidden_size] |
| 567 | # h_full: [n_tokens, expert_inner_dim * 2] |
| 568 | h_full = F.linear(x_tok, self.fc1_weight[local_expert_idx]) |
| 569 | |
| 570 | # GeGLU with (g + 1): h, g = split; output = gelu(h) * (g + 1) |
| 571 | # HF MoE uses exact GELU (not tanh approximation). |
| 572 | h, g = h_full.chunk(2, dim=-1) # Each [n_tokens, expert_inner_dim] |
| 573 | h = F.gelu(h) * (g + 1.0) |
| 574 | |
| 575 | # fc2: [hidden_size, expert_inner_dim] |
| 576 | # y: [n_tokens, hidden_size] |
| 577 | y = F.linear(h, self.fc2_weight[local_expert_idx]) |
| 578 | |
| 579 | # Apply routing weight |
| 580 | y = y * gate_tok.unsqueeze(-1) |
| 581 | |
| 582 | # Accumulate output |
| 583 | out.index_add_(0, token_pos, y) |
| 584 | |
| 585 | return out |
| 586 | |
| 587 | |
| 588 | class Moondream3Attention(nn.Module): |
| 589 | """Decoder attention with RoPE and tau scaling. |
| 590 | |
| 591 | Moondream3 uses a tau attention mechanism that scales Q and V |
| 592 | based on both token content and position. |
| 593 | """ |
| 594 | |
| 595 | def __init__( |
| 596 | self, |
| 597 | config: Moondream3TextConfig, |
| 598 | layer_idx: int, |
| 599 | cache_config=None, |
| 600 | quant_config: QuantizationConfig | None = None, |
| 601 | prefix: str = "", |
| 602 | ): |
| 603 | super().__init__() |
| 604 | self.hidden_size = config.dim |
| 605 | self.num_heads = config.n_heads |
| 606 | self.num_kv_heads = config.n_kv_heads |
| 607 | self.head_dim = config.dim // config.n_heads |
| 608 | |
| 609 | tp_size = get_tensor_model_parallel_world_size() |
| 610 | self.num_heads_per_partition = self.num_heads // tp_size |
| 611 | self.num_kv_heads_per_partition = max(1, self.num_kv_heads // tp_size) |
| 612 | |
| 613 | self.qkv_proj = QKVParallelLinear( |
| 614 | hidden_size=self.hidden_size, |
| 615 | head_size=self.head_dim, |
| 616 | total_num_heads=self.num_heads, |
| 617 | total_num_kv_heads=self.num_kv_heads, |
| 618 | bias=True, |
| 619 | quant_config=quant_config, |
| 620 | prefix=f"{prefix}.qkv_proj", |
| 621 | ) |
| 622 | |
| 623 | self.out_proj = RowParallelLinear( |
| 624 | input_size=self.hidden_size, |
| 625 | output_size=self.hidden_size, |
| 626 | bias=True, |
| 627 | quant_config=quant_config, |
| 628 | prefix=f"{prefix}.out_proj", |
| 629 | ) |
| 630 | |
| 631 | # Moondream uses 32-dim rotation out of 64-dim head (partial_rotary_factor=0.5) |
| 632 | # HF Moondream uses non-interleaved RoPE (split by half) |
| 633 | # In vLLM, is_neox_style=True means split by half (GPT-NeoX style) |
| 634 | rope_parameters = { |
| 635 | "rope_theta": config.rope_theta, |
| 636 | "partial_rotary_factor": 32 / self.head_dim, # 32/64 = 0.5 |
| 637 | } |
| 638 | self.rotary_emb = get_rope( |
| 639 | head_size=self.head_dim, |
| 640 | max_position=config.max_context, |
| 641 | rope_parameters=rope_parameters, |
| 642 | is_neox_style=True, # Moondream uses split-by-half (GPT-NeoX) style |
| 643 | ) |
| 644 | |
| 645 | self.scaling = self.head_dim**-0.5 |
| 646 | self.attn = Attention( |
| 647 | num_heads=self.num_heads_per_partition, |
| 648 | head_size=self.head_dim, |
| 649 | scale=self.scaling, |
| 650 | num_kv_heads=self.num_kv_heads_per_partition, |
| 651 | cache_config=cache_config, |
| 652 | quant_config=quant_config, |
| 653 | prefix=f"{prefix}.attn", |
| 654 | ) |
| 655 | |
| 656 | # Tau scaling parameters for position-dependent attention |
| 657 | # These are learned during training to modulate attention based on position |
| 658 | # tau_wq and tau_wv need full qkv_dim for correct computation |
| 659 | # Only heads are partitioned, qkv dimension is kept full for all-gather |
| 660 | qkv_dim = self.hidden_size * 3 # Q + K + V dimension (full) |
| 661 | self.tau_alpha = nn.Parameter(torch.zeros(self.num_heads_per_partition)) |
| 662 | self.tau_wq = nn.Parameter(torch.zeros(self.num_heads_per_partition, qkv_dim)) |
| 663 | self.tau_wv = nn.Parameter(torch.zeros(self.num_heads_per_partition, qkv_dim)) |
| 664 | self.tp_size = tp_size |
| 665 | |
| 666 | def forward( |
| 667 | self, |
| 668 | positions: torch.Tensor, |
| 669 | hidden_states: torch.Tensor, |
| 670 | ) -> torch.Tensor: |
| 671 | qkv, _ = self.qkv_proj(hidden_states) |
| 672 | |
| 673 | q, k, v = qkv.split( |
| 674 | [ |
| 675 | self.num_heads_per_partition * self.head_dim, |
| 676 | self.num_kv_heads_per_partition * self.head_dim, |
| 677 | self.num_kv_heads_per_partition * self.head_dim, |
| 678 | ], |
| 679 | dim=-1, |
| 680 | ) |
| 681 | |
| 682 | # Apply tau scaling to Q and V |
| 683 | # Tau scaling has two components: |
| 684 | # 1. Token-based: tok_q = tanh(gelu(qkv) @ tau_wq.T) |
| 685 | # 2. Position-based: tau_pos = 1 + (sigmoid(alpha * log(pos+1)) - 0.5) |
| 686 | # Final: tau = tok + tau_pos |
| 687 | # |
| 688 | # For TP, tau weights are sharded by head, but qkv_dim is kept full |
| 689 | |
| 690 | # Get full qkv for tau computation |
| 691 | # With TP, reconstruct qkv in correct layout [q_full, k_full, v_full] |
| 692 | # (all-gather would produce [q_0, k_0, v_0, q_1, k_1, v_1] - wrong) |
| 693 | if self.tp_size > 1: |
| 694 | # All-gather once, then reconstruct [q_full, k_full, v_full]. |
| 695 | qkv_full_sharded = tensor_model_parallel_all_gather(qkv.contiguous()) |
| 696 | q_local_dim = q.shape[-1] |
| 697 | kv_local_dim = k.shape[-1] |
| 698 | qkv_full_sharded = qkv_full_sharded.view( |
| 699 | qkv.shape[0], |
| 700 | self.tp_size, |
| 701 | q_local_dim + 2 * kv_local_dim, |
| 702 | ) |
| 703 | q_full = qkv_full_sharded[:, :, :q_local_dim].reshape(qkv.shape[0], -1) |
| 704 | k_full = qkv_full_sharded[ |
| 705 | :, :, q_local_dim : q_local_dim + kv_local_dim |
| 706 | ].reshape(qkv.shape[0], -1) |
| 707 | v_full = qkv_full_sharded[:, :, q_local_dim + kv_local_dim :].reshape( |
| 708 | qkv.shape[0], -1 |
| 709 | ) |
| 710 | qkv_full = torch.cat([q_full, k_full, v_full], dim=-1).contiguous() |
| 711 | else: |
| 712 | qkv_full = qkv |
| 713 | |
| 714 | # Compute tau scaling factors matching HF implementation exactly: |
| 715 | # tok_feat = gelu(qkv) |
| 716 | # tok_q = tanh(tok_feat @ tau_wq.T) # [num_tokens, num_heads] |
| 717 | # tau_pos = 1 + (sigmoid(alpha * log(pos+1)) - 0.5) # [num_heads, num_tokens] |
| 718 | # tau = (tok_q.T + tau_pos).T # [num_tokens, num_heads] |
| 719 | num_tokens = qkv_full.shape[0] |
| 720 | orig_dtype = q.dtype |
| 721 | |
| 722 | # Token-based component |
| 723 | tok_feat = F.gelu(qkv_full) # Apply GELU activation |
| 724 | tok_q = torch.tanh(tok_feat @ self.tau_wq.t()) # [N, H_per_partition] |
| 725 | tok_v = torch.tanh(tok_feat @ self.tau_wv.t()) # [N, H_per_partition] |
| 726 | |
| 727 | # Position-based component |
| 728 | # tau_pos = 1 + (sigmoid(alpha * log(pos+1)) - 0.5) |
| 729 | # positions is [num_tokens], need to compute for each head |
| 730 | # tau_alpha: [num_heads_per_partition] |
| 731 | pos_float = (positions.to(orig_dtype) + 1.0).clamp(min=1e-6) |
| 732 | pos_log = pos_float.log() # [num_tokens] |
| 733 | # alpha[:, None] * pos_log[None, :] -> [num_heads, num_tokens] |
| 734 | tau_pos = 1.0 + ( |
| 735 | torch.sigmoid(self.tau_alpha[:, None] * pos_log[None, :]) - 0.5 |
| 736 | ) # [H_per_partition, N] |
| 737 | |
| 738 | # Combine token and position components |
| 739 | tau_q = (tok_q + tau_pos.t()).to(orig_dtype) # [N, H_per_partition] |
| 740 | tau_v = (tok_v + tau_pos.t()).to(orig_dtype) # [N, H_per_partition] |
| 741 | |
| 742 | # Reshape q and v to apply per-head tau scaling |
| 743 | q = q.view(num_tokens, self.num_heads_per_partition, self.head_dim) |
| 744 | v = v.view(num_tokens, self.num_kv_heads_per_partition, self.head_dim) |
| 745 | |
| 746 | # Apply tau scaling |
| 747 | q = q * tau_q.unsqueeze(-1) |
| 748 | v = v * tau_v[:, : self.num_kv_heads_per_partition].unsqueeze(-1) |
| 749 | |
| 750 | # Reshape back |
| 751 | q = q.view(num_tokens, -1) |
| 752 | v = v.view(num_tokens, -1) |
| 753 | |
| 754 | q, k = self.rotary_emb(positions, q, k) |
| 755 | |
| 756 | attn_output = self.attn(q, k, v) |
| 757 | |
| 758 | output, _ = self.out_proj(attn_output) |
| 759 | return output |
| 760 | |
| 761 | |
| 762 | class Moondream3DecoderLayer(nn.Module): |
| 763 | """Decoder layer with attention + MLP/MoE.""" |
| 764 | |
| 765 | def __init__( |
| 766 | self, |
| 767 | config: Moondream3TextConfig, |
| 768 | cache_config=None, |
| 769 | quant_config: QuantizationConfig | None = None, |
| 770 | prefix: str = "", |
| 771 | ): |
| 772 | super().__init__() |
| 773 | layer_idx = extract_layer_index(prefix) |
| 774 | self.layer_idx = layer_idx |
| 775 | |
| 776 | self.ln = nn.LayerNorm(config.dim, eps=1e-5, bias=True) |
| 777 | |
| 778 | self.attn = Moondream3Attention( |
| 779 | config=config, |
| 780 | layer_idx=layer_idx, |
| 781 | cache_config=cache_config, |
| 782 | quant_config=quant_config, |
| 783 | prefix=f"{prefix}.attn", |
| 784 | ) |
| 785 | |
| 786 | # Use MoE for layers >= moe_start_layer, standard MLP otherwise |
| 787 | if layer_idx >= config.moe_start_layer: |
| 788 | self.mlp = Moondream3TextMoE( |
| 789 | hidden_size=config.dim, |
| 790 | expert_inner_dim=config.moe_expert_inner_dim, |
| 791 | num_experts=config.moe_num_experts, |
| 792 | experts_per_token=config.moe_experts_per_token, |
| 793 | quant_config=quant_config, |
| 794 | prefix=f"{prefix}.mlp", |
| 795 | ) |
| 796 | else: |
| 797 | self.mlp = Moondream3TextMLP( |
| 798 | hidden_size=config.dim, |
| 799 | intermediate_size=config.ff_dim, |
| 800 | quant_config=quant_config, |
| 801 | prefix=f"{prefix}.mlp", |
| 802 | ) |
| 803 | |
| 804 | def forward( |
| 805 | self, |
| 806 | positions: torch.Tensor, |
| 807 | hidden_states: torch.Tensor, |
| 808 | ) -> torch.Tensor: |
| 809 | # Pre-norm architecture |
| 810 | normed = self.ln(hidden_states) |
| 811 | attn_out = self.attn(positions, normed) |
| 812 | mlp_out = self.mlp(normed) |
| 813 | hidden_states = hidden_states + attn_out + mlp_out |
| 814 | return hidden_states |
| 815 | |
| 816 | |
| 817 | class Moondream3TextModel(nn.Module): |
| 818 | """Text decoder model.""" |
| 819 | |
| 820 | def __init__( |
| 821 | self, |
| 822 | config: Moondream3TextConfig, |
| 823 | cache_config=None, |
| 824 | quant_config: QuantizationConfig | None = None, |
| 825 | prefix: str = "", |
| 826 | ): |
| 827 | super().__init__() |
| 828 | self.config = config |
| 829 | |
| 830 | self.wte = VocabParallelEmbedding( |
| 831 | config.vocab_size, |
| 832 | config.dim, |
| 833 | prefix=f"{prefix}.wte", |
| 834 | ) |
| 835 | |
| 836 | blocks_prefix = maybe_prefix(prefix, "blocks") |
| 837 | self.start_layer, self.end_layer, self.blocks = make_layers( |
| 838 | config.n_layers, |
| 839 | lambda prefix: Moondream3DecoderLayer( |
| 840 | config=config, |
| 841 | cache_config=cache_config, |
| 842 | quant_config=quant_config, |
| 843 | prefix=prefix, |
| 844 | ), |
| 845 | prefix=blocks_prefix, |
| 846 | ) |
| 847 | |
| 848 | self.post_ln = nn.LayerNorm(config.dim, eps=1e-5, bias=True) |
| 849 | self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( |
| 850 | ["hidden_states"], config.dim |
| 851 | ) |
| 852 | |
| 853 | def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 854 | return self.wte(input_ids) |
| 855 | |
| 856 | def forward( |
| 857 | self, |
| 858 | input_ids: torch.Tensor | None, |
| 859 | positions: torch.Tensor, |
| 860 | intermediate_tensors: IntermediateTensors | None = None, |
| 861 | inputs_embeds: torch.Tensor | None = None, |
| 862 | ) -> torch.Tensor | IntermediateTensors: |
| 863 | pp_group = get_pp_group() |
| 864 | if pp_group.is_first_rank: |
| 865 | if inputs_embeds is not None: |
| 866 | hidden_states = inputs_embeds |
| 867 | else: |
| 868 | assert input_ids is not None |
| 869 | hidden_states = self.embed_input_ids(input_ids) |
| 870 | else: |
| 871 | assert intermediate_tensors is not None |
| 872 | hidden_states = intermediate_tensors["hidden_states"] |
| 873 | |
| 874 | for i, layer in enumerate( |
| 875 | islice(self.blocks, self.start_layer, self.end_layer) |
| 876 | ): |
| 877 | hidden_states = layer(positions, hidden_states) |
| 878 | |
| 879 | if not pp_group.is_last_rank: |
| 880 | return IntermediateTensors({"hidden_states": hidden_states}) |
| 881 | |
| 882 | hidden_states = self.post_ln(hidden_states) |
| 883 | return hidden_states |
| 884 | |
| 885 | |
| 886 | @dataclass(frozen=True) |
| 887 | class Moondream3ImageInput: |
| 888 | """Container holding per-image inputs for embedding.""" |
| 889 | |
| 890 | pixel_values: torch.Tensor |
| 891 | tiling: tuple[int, int] | None |
| 892 | |
| 893 | |
| 894 | # ============================================================================ |
| 895 | # Multimodal Processing |
| 896 | # ============================================================================ |
| 897 | |
| 898 | |
| 899 | class Moondream3ProcessingInfo(BaseProcessingInfo): |
| 900 | """Processing info for Moondream3.""" |
| 901 | |
| 902 | def get_hf_config(self): |
| 903 | return self.ctx.get_hf_config() |
| 904 | |
| 905 | def get_hf_processor(self, **kwargs: object): |
| 906 | return self.ctx.get_hf_processor(Moondream3Processor, **kwargs) |
| 907 | |
| 908 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 909 | return {"image": 1} |
| 910 | |
| 911 | def get_num_image_tokens( |
| 912 | self, |
| 913 | *, |
| 914 | image_width: int, |
| 915 | image_height: int, |
| 916 | ) -> int: |
| 917 | # HF pre-fills BOS together with the fixed 27x27 vision grid under |
| 918 | # the same bidirectional prefix mask: 1 BOS + 729 image embeddings. |
| 919 | return 730 |
| 920 | |
| 921 | def get_image_size_with_most_features(self) -> ImageSize: |
| 922 | return ImageSize(width=378, height=378) |
| 923 | |
| 924 | def get_max_image_tokens(self) -> int: |
| 925 | return 730 |
| 926 | |
| 927 | def get_mm_max_tokens_per_item( |
| 928 | self, |
| 929 | seq_len: int, |
| 930 | mm_counts: Mapping[str, int], |
| 931 | ) -> Mapping[str, int]: |
| 932 | return {"image": self.get_max_image_tokens()} |
| 933 | |
| 934 | |
| 935 | class Moondream3DummyInputsBuilder(BaseDummyInputsBuilder[Moondream3ProcessingInfo]): |
| 936 | """Dummy inputs builder for profiling.""" |
| 937 | |
| 938 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 939 | return ( |
| 940 | "<|endoftext|><image><|md_reserved_0|>query<|md_reserved_1|>" |
| 941 | "What is this image?<|md_reserved_2|>" |
| 942 | ) |
| 943 | |
| 944 | def get_dummy_mm_data( |
| 945 | self, |
| 946 | seq_len: int, |
| 947 | mm_counts: Mapping[str, int], |
| 948 | mm_options: MultiModalDummyOptions | None = None, |
| 949 | mm_processor_kwargs: Mapping[str, object] | None = None, |
| 950 | ) -> MultiModalDataDict: |
| 951 | return { |
| 952 | "image": self._get_dummy_images( |
| 953 | width=378, |
| 954 | height=378, |
| 955 | num_images=mm_counts.get("image", 0), |
| 956 | ) |
| 957 | } |
| 958 | |
| 959 | |
| 960 | class Moondream3MultiModalProcessor(BaseMultiModalProcessor[Moondream3ProcessingInfo]): |
| 961 | """Multimodal processor for Moondream3.""" |
| 962 | |
| 963 | image_placeholder: str = "<image>" |
| 964 | bos_image_placeholder: str = "<|endoftext|><image>" |
| 965 | |
| 966 | @cached_property |
| 967 | def bos_image_placeholder_tokens(self) -> list[int]: |
| 968 | tokenizer = self.info.get_tokenizer() |
| 969 | token_ids = cached_encode( |
| 970 | tokenizer, |
| 971 | self.bos_image_placeholder, |
| 972 | add_special_tokens=False, |
| 973 | ) |
| 974 | if len(token_ids) < 2: |
| 975 | raise ValueError( |
| 976 | "Tokenizer could not encode Moondream3 BOS/image placeholder " |
| 977 | f"{self.bos_image_placeholder!r}." |
| 978 | ) |
| 979 | return token_ids |
| 980 | |
| 981 | def _get_mm_fields_config( |
| 982 | self, |
| 983 | hf_inputs: BatchFeature, |
| 984 | hf_processor_mm_kwargs: Mapping[str, object], |
| 985 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 986 | return { |
| 987 | "pixel_values": MultiModalFieldConfig.batched("image"), |
| 988 | "tilings": MultiModalFieldConfig.batched("image", keep_on_cpu=True), |
| 989 | } |
| 990 | |
| 991 | def _get_prompt_updates( |
| 992 | self, |
| 993 | mm_items: MultiModalDataItems, |
| 994 | hf_processor_mm_kwargs: Mapping[str, object], |
| 995 | out_mm_kwargs: MultiModalKwargsItems, |
| 996 | ) -> list[PromptUpdate]: |
| 997 | image_size = self.info.get_image_size_with_most_features() |
| 998 | num_image_tokens = self.info.get_num_image_tokens( |
| 999 | image_width=image_size.width, |
| 1000 | image_height=image_size.height, |
| 1001 | ) |
| 1002 | placeholder_tokens = self.bos_image_placeholder_tokens |
| 1003 | bos_token = placeholder_tokens[0] |
| 1004 | image_token = placeholder_tokens[-1] |
| 1005 | return [ |
| 1006 | PromptReplacement( |
| 1007 | modality="image", |
| 1008 | target=placeholder_tokens, |
| 1009 | replacement=PromptUpdateDetails( |
| 1010 | full=[bos_token] + [image_token] * (num_image_tokens - 1), |
| 1011 | ), |
| 1012 | ), |
| 1013 | ] |
| 1014 | |
| 1015 | |
| 1016 | # ============================================================================ |
| 1017 | # Main Model |
| 1018 | # ============================================================================ |
| 1019 | |
| 1020 | |
| 1021 | @MULTIMODAL_REGISTRY.register_processor( |
| 1022 | Moondream3MultiModalProcessor, |
| 1023 | info=Moondream3ProcessingInfo, |
| 1024 | dummy_inputs=Moondream3DummyInputsBuilder, |
| 1025 | ) |
| 1026 | class Moondream3ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP): |
| 1027 | """Moondream3 multimodal model for causal language modeling. |
| 1028 | |
| 1029 | vLLM supports the standard autoregressive Moondream3 query and caption |
| 1030 | prompt formats. The region-module point/detect skills require custom |
| 1031 | coordinate decoding and are intentionally not exposed here. |
| 1032 | """ |
| 1033 | |
| 1034 | supports_multimodal = True |
| 1035 | packed_modules_mapping = { |
| 1036 | "qkv_proj": ["q_proj", "k_proj", "v_proj"], |
| 1037 | } |
| 1038 | |
| 1039 | def __init__( |
| 1040 | self, |
| 1041 | *, |
| 1042 | vllm_config: VllmConfig, |
| 1043 | prefix: str = "", |
| 1044 | ): |
| 1045 | super().__init__() |
| 1046 | |
| 1047 | hf_config = vllm_config.model_config.hf_config |
| 1048 | quant_config = vllm_config.quant_config |
| 1049 | cache_config = vllm_config.cache_config |
| 1050 | |
| 1051 | # Reuse the transformers_utils config implementation. |
| 1052 | if isinstance(hf_config, Moondream3Config): |
| 1053 | self.config = hf_config |
| 1054 | else: |
| 1055 | config_dict = hf_config.config if hasattr(hf_config, "config") else {} |
| 1056 | self.config = Moondream3Config(config=config_dict) |
| 1057 | |
| 1058 | with self._mark_tower_model(vllm_config, "image"): |
| 1059 | # Vision encoder |
| 1060 | self.vision = Moondream3VisionEncoder( |
| 1061 | config=self.config.vision_config, |
| 1062 | quant_config=quant_config, |
| 1063 | prefix=maybe_prefix(prefix, "vision"), |
| 1064 | ) |
| 1065 | |
| 1066 | # Vision projection |
| 1067 | self.vision_proj = Moondream3VisionProjection( |
| 1068 | input_dim=self.config.vision_config.enc_dim, |
| 1069 | inner_dim=self.config.vision_config.proj_inner_dim, |
| 1070 | output_dim=self.config.text_config.dim, |
| 1071 | quant_config=quant_config, |
| 1072 | prefix=maybe_prefix(prefix, "vision_proj"), |
| 1073 | ) |
| 1074 | |
| 1075 | with self._mark_language_model(vllm_config): |
| 1076 | # Text decoder |
| 1077 | self.text = Moondream3TextModel( |
| 1078 | config=self.config.text_config, |
| 1079 | cache_config=cache_config, |
| 1080 | quant_config=quant_config, |
| 1081 | prefix=maybe_prefix(prefix, "text"), |
| 1082 | ) |
| 1083 | |
| 1084 | # LM head (with bias - Moondream3 has lm_head bias) |
| 1085 | self.lm_head = ParallelLMHead( |
| 1086 | self.config.text_config.vocab_size, |
| 1087 | self.config.text_config.dim, |
| 1088 | bias=True, |
| 1089 | quant_config=quant_config, |
| 1090 | prefix=maybe_prefix(prefix, "lm_head"), |
| 1091 | ) |
| 1092 | |
| 1093 | self.logits_processor = LogitsProcessor(self.config.text_config.vocab_size) |
| 1094 | self.make_empty_intermediate_tensors = self.text.make_empty_intermediate_tensors |
| 1095 | self._answer_id = getattr( |
| 1096 | self.config, |
| 1097 | "answer_token_id", |
| 1098 | getattr(hf_config, "answer_token_id", 3), |
| 1099 | ) |
| 1100 | |
| 1101 | @classmethod |
| 1102 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 1103 | if modality == "image": |
| 1104 | return "<image>" |
| 1105 | return None |
| 1106 | |
| 1107 | def get_language_model(self) -> nn.Module: |
| 1108 | return self.text |
| 1109 | |
| 1110 | def get_mm_lora_token_counts( |
| 1111 | self, |
| 1112 | *, |
| 1113 | modality: str, |
| 1114 | mm_kwargs: MultiModalKwargsItem | None, |
| 1115 | num_mm_embeds: int, |
| 1116 | ) -> tuple[int, int | None]: |
| 1117 | del modality, mm_kwargs |
| 1118 | return num_mm_embeds, num_mm_embeds |
| 1119 | |
| 1120 | def _split_pixel_values( |
| 1121 | self, |
| 1122 | pixel_values: object, |
| 1123 | ) -> list[torch.Tensor]: |
| 1124 | # The processor should standardize image inputs into: |
| 1125 | # - torch.Tensor [num_images, num_crops, C, H, W], or |
| 1126 | # - list[torch.Tensor[num_crops, C, H, W]] for ragged crops. |
| 1127 | if isinstance(pixel_values, torch.Tensor): |
| 1128 | if pixel_values.dim() != 5: |
| 1129 | raise ValueError( |
| 1130 | "Expected `pixel_values` tensor with shape " |
| 1131 | "[num_images, num_crops, C, H, W], got " |
| 1132 | f"{tuple(pixel_values.shape)}." |
| 1133 | ) |
| 1134 | return [pv.contiguous() for pv in pixel_values] |
| 1135 | |
| 1136 | if isinstance(pixel_values, (list, tuple)): |
| 1137 | tensors: list[torch.Tensor] = [] |
| 1138 | for value in pixel_values: |
| 1139 | if not isinstance(value, torch.Tensor): |
| 1140 | raise TypeError( |
| 1141 | "Expected each `pixel_values` element to be a tensor, " |
| 1142 | f"got {type(value)!r}." |
| 1143 | ) |
| 1144 | if value.dim() != 4: |
| 1145 | raise ValueError( |
| 1146 | f"Unsupported pixel_values element shape {tuple(value.shape)}." |
| 1147 | ) |
| 1148 | tensors.append(value.contiguous()) |
| 1149 | return tensors |
| 1150 | |
| 1151 | raise TypeError( |
| 1152 | "pixel_values must be a tensor or a sequence of tensors, " |
| 1153 | f"got {type(pixel_values)!r}." |
| 1154 | ) |
| 1155 | |
| 1156 | def _split_tilings( |
| 1157 | self, |
| 1158 | tilings: object, |
| 1159 | expected: int, |
| 1160 | ) -> list[tuple[int, int] | None]: |
| 1161 | if tilings is None: |
| 1162 | return [None] * expected |
| 1163 | |
| 1164 | if isinstance(tilings, torch.Tensor): |
| 1165 | if tilings.dim() != 2 or tilings.shape[1] != 2: |
| 1166 | raise ValueError( |
| 1167 | "Expected `tilings` tensor with shape [num_images, 2], got " |
| 1168 | f"{tuple(tilings.shape)}." |
| 1169 | ) |
| 1170 | tiling_items = tilings.tolist() |
| 1171 | elif isinstance(tilings, (list, tuple)): |
| 1172 | tiling_items = list(tilings) |
| 1173 | else: |
| 1174 | raise TypeError( |
| 1175 | "tilings must be None, a tensor or a sequence of tuples, " |
| 1176 | f"got {type(tilings)!r}." |
| 1177 | ) |
| 1178 | |
| 1179 | if len(tiling_items) != expected: |
| 1180 | raise ValueError( |
| 1181 | "Mismatch between the number of pixel_values entries " |
| 1182 | f"({expected}) and tilings ({len(tiling_items)})." |
| 1183 | ) |
| 1184 | |
| 1185 | normalized: list[tuple[int, int] | None] = [] |
| 1186 | for tiling in tiling_items: |
| 1187 | if tiling is None: |
| 1188 | normalized.append(None) |
| 1189 | continue |
| 1190 | if isinstance(tiling, torch.Tensor): |
| 1191 | tiling = tiling.tolist() |
| 1192 | if isinstance(tiling, (list, tuple)) and len(tiling) == 2: |
| 1193 | normalized.append((int(tiling[0]), int(tiling[1]))) |
| 1194 | else: |
| 1195 | raise ValueError( |
| 1196 | f"Each tiling entry must be a pair of integers, got {tiling!r}." |
| 1197 | ) |
| 1198 | return normalized |
| 1199 | |
| 1200 | def _parse_image_inputs(self, **kwargs: object) -> list[Moondream3ImageInput]: |
| 1201 | pixel_values = kwargs.get("pixel_values") |
| 1202 | if pixel_values is None: |
| 1203 | return [] |
| 1204 | |
| 1205 | pixel_values_list = self._split_pixel_values(pixel_values) |
| 1206 | tilings_list = self._split_tilings( |
| 1207 | kwargs.get("tilings"), len(pixel_values_list) |
| 1208 | ) |
| 1209 | |
| 1210 | image_inputs: list[Moondream3ImageInput] = [] |
| 1211 | for value, tiling in zip(pixel_values_list, tilings_list): |
| 1212 | if value.dim() != 4: |
| 1213 | raise ValueError( |
| 1214 | f"Expected 4D tensor for crops, got {tuple(value.shape)}." |
| 1215 | ) |
| 1216 | image_inputs.append(Moondream3ImageInput(pixel_values=value, tiling=tiling)) |
| 1217 | return image_inputs |
| 1218 | |
| 1219 | def _encode_image_input(self, image_input: Moondream3ImageInput) -> torch.Tensor: |
| 1220 | pixel_values = image_input.pixel_values |
| 1221 | if pixel_values.dim() != 4: |
| 1222 | raise ValueError( |
| 1223 | f"Expected 4D tensor for crops, got {tuple(pixel_values.shape)}." |
| 1224 | ) |
| 1225 | |
| 1226 | device = self.vision.patch_emb.weight.device |
| 1227 | dtype = self.vision.patch_emb.weight.dtype |
| 1228 | pixel_values = pixel_values.to(device=device, dtype=dtype) |
| 1229 | |
| 1230 | features = self.vision(pixel_values) |
| 1231 | |
| 1232 | # Grid size = crop_size / patch_size (e.g., 378 / 14 = 27) |
| 1233 | grid_size = ( |
| 1234 | self.config.vision_config.crop_size |
| 1235 | // self.config.vision_config.enc_patch_size |
| 1236 | ) |
| 1237 | enc_dim = self.config.vision_config.enc_dim |
| 1238 | global_features = features[0] |
| 1239 | |
| 1240 | if features.shape[0] > 1: |
| 1241 | if image_input.tiling is None: |
| 1242 | raise ValueError( |
| 1243 | "Missing tiling metadata for multi-crop Moondream image." |
| 1244 | ) |
| 1245 | local = features[1:].contiguous().view(-1, grid_size, grid_size, enc_dim) |
| 1246 | reconstructed = reconstruct_from_crops( |
| 1247 | local, |
| 1248 | image_input.tiling, |
| 1249 | overlap_margin=self.config.vision_config.overlap_margin, |
| 1250 | patch_size=1, |
| 1251 | ) |
| 1252 | else: |
| 1253 | reconstructed = global_features.view(grid_size, grid_size, enc_dim) |
| 1254 | |
| 1255 | recon = reconstructed.permute(2, 0, 1).contiguous() |
| 1256 | # Mirror HF reference behavior: reconstructed local features are pooled |
| 1257 | # to enc_n_layers x enc_n_layers. For moondream3-preview this is 27x27. |
| 1258 | pooled_size = self.config.vision_config.enc_n_layers |
| 1259 | if pooled_size != grid_size: |
| 1260 | logger.warning_once( |
| 1261 | "Moondream3 pooled_size (%d) differs from crop grid (%d). " |
| 1262 | "Using enc_n_layers to match HF reference behavior.", |
| 1263 | pooled_size, |
| 1264 | grid_size, |
| 1265 | ) |
| 1266 | recon = F.adaptive_avg_pool2d(recon, output_size=(pooled_size, pooled_size)) |
| 1267 | recon = recon.permute(1, 2, 0).contiguous().view(-1, enc_dim) |
| 1268 | |
| 1269 | combined = torch.cat([global_features, recon], dim=-1).unsqueeze(0) |
| 1270 | projected = self.vision_proj(combined).squeeze(0) |
| 1271 | |
| 1272 | # Note: Vision embeddings are already synchronized across TP ranks |
| 1273 | # because the vision projection uses RowParallelLinear which performs |
| 1274 | # all-reduce internally, ensuring identical outputs on all ranks. |
| 1275 | |
| 1276 | return projected |
| 1277 | |
| 1278 | def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: |
| 1279 | """Generate the HF image prefix: BOS embedding + 729 image embeddings.""" |
| 1280 | image_inputs = self._parse_image_inputs(**kwargs) |
| 1281 | if not image_inputs: |
| 1282 | return [] |
| 1283 | |
| 1284 | device = self.vision.patch_emb.weight.device |
| 1285 | bos_ids = torch.tensor([self.config.bos_token_id], device=device) |
| 1286 | bos_embedding = self.text.embed_input_ids(bos_ids) |
| 1287 | |
| 1288 | embeddings: list[torch.Tensor] = [] |
| 1289 | for image_input in image_inputs: |
| 1290 | image_embeddings = self._encode_image_input(image_input) |
| 1291 | embeddings.append( |
| 1292 | torch.cat([bos_embedding.to(image_embeddings.dtype), image_embeddings]) |
| 1293 | ) |
| 1294 | return embeddings |
| 1295 | |
| 1296 | def forward( |
| 1297 | self, |
| 1298 | input_ids: torch.Tensor | None, |
| 1299 | positions: torch.Tensor, |
| 1300 | intermediate_tensors: IntermediateTensors | None = None, |
| 1301 | inputs_embeds: torch.Tensor | None = None, |
| 1302 | **kwargs, |
| 1303 | ) -> torch.Tensor | IntermediateTensors: |
| 1304 | hidden_states = self.text( |
| 1305 | input_ids=input_ids, |
| 1306 | positions=positions, |
| 1307 | intermediate_tensors=intermediate_tensors, |
| 1308 | inputs_embeds=inputs_embeds, |
| 1309 | ) |
| 1310 | return hidden_states |
| 1311 | |
| 1312 | def compute_logits( |
| 1313 | self, |
| 1314 | hidden_states: torch.Tensor, |
| 1315 | ) -> torch.Tensor | None: |
| 1316 | logits = self.logits_processor(self.lm_head, hidden_states) |
| 1317 | if logits is not None: |
| 1318 | logits[:, self._answer_id] = float("-inf") |
| 1319 | return logits |
| 1320 | |
| 1321 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 1322 | """Load weights with remapping from HuggingFace format.""" |
| 1323 | params_dict = dict(self.named_parameters()) |
| 1324 | loaded_params: set[str] = set() |
| 1325 | |
| 1326 | # Get expert intermediate size for fc1 splitting |
| 1327 | |
| 1328 | for name, loaded_weight in weights: |
| 1329 | # Map from HF naming to vLLM naming |
| 1330 | # model.vision.* -> vision.* |
| 1331 | # model.text.* -> text.* |
| 1332 | if name.startswith("model."): |
| 1333 | name = name[6:] # Remove "model." prefix |
| 1334 | |
| 1335 | # Specific name mappings |
| 1336 | # Vision projection: vision.proj_mlp.fc1 -> vision_proj.fc1 |
| 1337 | name = name.replace("vision.proj_mlp.", "vision_proj.") |
| 1338 | |
| 1339 | # Text embedding: text.wte (no suffix) -> text.wte.weight |
| 1340 | if name == "text.wte": |
| 1341 | name = "text.wte.weight" |
| 1342 | |
| 1343 | # LM head: text.lm_head -> lm_head |
| 1344 | name = name.replace("text.lm_head.", "lm_head.") |
| 1345 | |
| 1346 | # Attention mapping |
| 1347 | name = name.replace(".attn.qkv.", ".attn.qkv_proj.") |
| 1348 | name = name.replace(".attn.proj.", ".attn.out_proj.") |
| 1349 | |
| 1350 | # Tau attention scaling weights |
| 1351 | # HF format: .attn.tau.alpha -> .attn.tau_alpha |
| 1352 | name = name.replace(".attn.tau.alpha", ".attn.tau_alpha") |
| 1353 | name = name.replace(".attn.tau.wq", ".attn.tau_wq") |
| 1354 | name = name.replace(".attn.tau.wv", ".attn.tau_wv") |
| 1355 | |
| 1356 | # MoE router mapping: mlp.router -> mlp.gate |
| 1357 | name = name.replace(".mlp.router.", ".mlp.gate.") |
| 1358 | |
| 1359 | # Handle MoE expert weights for layers 4+ with expert parallelism |
| 1360 | # fc1.weight: [n_experts, expert_inner_dim * 2, hidden_size] (gate+up) |
| 1361 | # fc2.weight: [n_experts, hidden_size, expert_inner_dim] (down) |
| 1362 | # Each GPU stores n_experts/tp_size experts |
| 1363 | # Note: Only 3D weights are MoE, 2D weights are standard MLP |
| 1364 | if ".mlp.fc1.weight" in name and loaded_weight.dim() == 3: |
| 1365 | from vllm.distributed import get_tensor_model_parallel_rank |
| 1366 | |
| 1367 | tp_size = get_tensor_model_parallel_world_size() |
| 1368 | tp_rank = get_tensor_model_parallel_rank() |
| 1369 | num_experts = loaded_weight.shape[0] |
| 1370 | experts_per_rank = num_experts // tp_size |
| 1371 | expert_start = tp_rank * experts_per_rank |
| 1372 | expert_end = expert_start + experts_per_rank |
| 1373 | # Shard by expert dimension |
| 1374 | loaded_weight = loaded_weight[expert_start:expert_end].contiguous() |
| 1375 | # Map to our custom MoE format: mlp.fc1_weight |
| 1376 | name = name.replace(".mlp.fc1.weight", ".mlp.fc1_weight") |
| 1377 | |
| 1378 | if ".mlp.fc2.weight" in name and loaded_weight.dim() == 3: |
| 1379 | from vllm.distributed import get_tensor_model_parallel_rank |
| 1380 | |
| 1381 | tp_size = get_tensor_model_parallel_world_size() |
| 1382 | tp_rank = get_tensor_model_parallel_rank() |
| 1383 | num_experts = loaded_weight.shape[0] |
| 1384 | experts_per_rank = num_experts // tp_size |
| 1385 | expert_start = tp_rank * experts_per_rank |
| 1386 | expert_end = expert_start + experts_per_rank |
| 1387 | # Shard by expert dimension |
| 1388 | loaded_weight = loaded_weight[expert_start:expert_end].contiguous() |
| 1389 | # Map to our custom MoE format: mlp.fc2_weight |
| 1390 | name = name.replace(".mlp.fc2.weight", ".mlp.fc2_weight") |
| 1391 | |
| 1392 | # Handle tau weights with tensor parallelism |
| 1393 | # tau_alpha: [num_heads] -> [num_heads/tp] |
| 1394 | # tau_wq: [num_heads, qkv_dim] -> [num_heads/tp, qkv_dim/tp] |
| 1395 | # tau_wv: [num_heads, qkv_dim] -> [num_heads/tp, qkv_dim/tp] |
| 1396 | if ".tau_alpha" in name: |
| 1397 | from vllm.distributed import get_tensor_model_parallel_rank |
| 1398 | |
| 1399 | tp_size = get_tensor_model_parallel_world_size() |
| 1400 | tp_rank = get_tensor_model_parallel_rank() |
| 1401 | num_heads = loaded_weight.shape[0] |
| 1402 | heads_per_partition = num_heads // tp_size |
| 1403 | start = tp_rank * heads_per_partition |
| 1404 | end = start + heads_per_partition |
| 1405 | loaded_weight = loaded_weight[start:end].contiguous() |
| 1406 | |
| 1407 | if ".tau_wq" in name or ".tau_wv" in name: |
| 1408 | from vllm.distributed import get_tensor_model_parallel_rank |
| 1409 | |
| 1410 | tp_size = get_tensor_model_parallel_world_size() |
| 1411 | tp_rank = get_tensor_model_parallel_rank() |
| 1412 | num_heads, qkv_dim = loaded_weight.shape |
| 1413 | heads_per_partition = num_heads // tp_size |
| 1414 | # Only shard by head dimension, keep full qkv_dim for all-gather |
| 1415 | head_start = tp_rank * heads_per_partition |
| 1416 | head_end = head_start + heads_per_partition |
| 1417 | loaded_weight = loaded_weight[head_start:head_end, :].contiguous() |
| 1418 | |
| 1419 | if name in params_dict: |
| 1420 | param = params_dict[name] |
| 1421 | weight_loader = getattr(param, "weight_loader", default_weight_loader) |
| 1422 | weight_loader(param, loaded_weight) |
| 1423 | loaded_params.add(name) |
| 1424 | |
| 1425 | return loaded_params |