@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7198940b39

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 10:00 UTC

[Model] Add Moondream3 model support(only query and caption skills) (#32325)

Signed-off-by: Dong Wang <dongw2019@gmail.com>

Dong W committed Apr 30, 2026 at 19:06 UTC 7198940b395e74337101e7b62437c55a7e1d4284

19 files changed +3238 -9

docs/features/multimodal_inputs.md

+61

@@ -215,6 +215,67 @@ When loading RGBA images (images with transparency), vLLM converts them to RGB f

215 - This setting only affects RGBA images with transparency; RGB images are unchanged

216 - If not specified, the default white background `(255, 255, 255)` is used for backward compatibility

217

218 +#### Moondream3 Prompt Recipes { #moondream3-prompt-recipes }

219 +

220 +`Moondream3ForCausalLM` supports two task-specific prompt formats:

221 +

222 +- `query`: ask a question about the image.

223 +- `caption`: generate a caption for the image.

224 +

225 +```python

226 +from vllm import LLM, SamplingParams

227 +from vllm.assets.image import ImageAsset

228 +

229 +llm = LLM(

230 + model="moondream/moondream3-preview",

231 + tokenizer="moondream/starmie-v1",

232 + trust_remote_code=True,

233 + max_model_len=2048,

234 + limit_mm_per_prompt={"image": 1},

235 +)

236 +

237 +image = ImageAsset("stop_sign").pil_image

238 +

239 +

240 +def make_query_prompt(question: str) -> str:

241 + return (

243 + f"{question}<|md_reserved_2|>"

244 + )

245 +

246 +

247 +def make_caption_prompt(length: str = "normal") -> str:

248 + return (

249 + "<|endoftext|><image><|md_reserved_0|>"

250 + f"describe<|md_reserved_1|>{length}<|md_reserved_2|>"

251 + )

252 +

253 +

254 +query_out = llm.generate(

255 + {

256 + "prompt": make_query_prompt("What is shown in this image?"),

257 + "multi_modal_data": {"image": image},

258 + },

259 + SamplingParams(max_tokens=64, temperature=0),

260 +)[0].outputs[0].text

261 +

262 +caption_out = llm.generate(

263 + {

264 + "prompt": make_caption_prompt(),

265 + "multi_modal_data": {"image": image},

266 + },

267 + SamplingParams(max_tokens=100, temperature=0),

268 +)[0].outputs[0].text

269 +

270 +print("query:", query_out)

271 +print("caption:", caption_out)

272 +```

273 +

274 +!!! note

275 + The native Moondream3 model also has `detect` and `point` skills. Those

276 + require custom coordinate decoding and are not exposed by this vLLM

277 + implementation.

278 +

279 ### Video Inputs

280

281 You can pass a list of NumPy arrays directly to the `'video'` field of the multi-modal dictionary

docs/models/supported_models.md

@@ -599,6 +599,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen

599 | `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I+ | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ |

600 | `MolmoForCausalLM` | Molmo | T + I+ | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ |

601 | `Molmo2ForConditionalGeneration` | Molmo2 | T + I+ / V | `allenai/Molmo2-4B`, `allenai/Molmo2-8B`, `allenai/Molmo2-O-7B` | ✅︎ | ✅︎ |

602 +| `Moondream3ForCausalLM` | Moondream3 | T + I | `moondream/moondream3-preview` | | ✅︎ |

603 | `MusicFlamingoForConditionalGeneration` | MusicFlamingo | T + A | `nvidia/music-flamingo-2601-hf`, `nvidia/music-flamingo-think-2601-hf` | ✅︎ | ✅︎ |

604 | `NVLM_D_Model` | NVLM-D 1.0 | T + I+ | `nvidia/NVLM-D-72B`, etc. | | ✅︎ |

605 | `OpenCUAForConditionalGeneration` | OpenCUA-7B | T + IE+ | `xlangai/OpenCUA-7B` | ✅︎ | ✅︎ |

@@ -661,6 +662,12 @@ Some models are supported only via the [Transformers modeling backend](#transfor

662 !!! note

663 For `InternVLChatModel`, only InternVL2.5 with Qwen2.5 text backbone (`OpenGVLab/InternVL2.5-1B` etc.), InternVL3 and InternVL3.5 have video inputs support currently.

664

665 +!!! note

666 + `Moondream3ForCausalLM` uses task-specific prompt templates for `query`

667 + and `caption`. The native `detect` and `point` skills require custom

668 + coordinate decoding and are not exposed by this vLLM implementation.

669 + See [Moondream3 prompt recipes](../features/multimodal_inputs.md#moondream3-prompt-recipes).

670 +

671 !!! note

672 To use `TIGER-Lab/Mantis-8B-siglip-llama3`, you have to pass `--hf_overrides '{"architectures": ["MantisForConditionalGeneration"]}'` when running vLLM.

673

tests/conftest.py

+31 -8

@@ -371,6 +371,8 @@ class HfRunner:

371 is_cross_encoder: bool = False,

372 skip_tokenizer_init: bool = False,

373 auto_cls: type[_BaseAutoModelClass] = AutoModelForCausalLM,

374 + tokenizer_name: str | None = None,

375 + processor: Any | None = None,

376 # Set this to avoid hanging issue

377 default_torch_num_threads: int | None = None,

378 ) -> None:

@@ -391,6 +393,8 @@ class HfRunner:

393 is_cross_encoder=is_cross_encoder,

394 skip_tokenizer_init=skip_tokenizer_init,

395 auto_cls=auto_cls,

396 + tokenizer_name=tokenizer_name,

397 + processor=processor,

398 )

399

400 def _init(

@@ -405,6 +409,8 @@ class HfRunner:

409 is_cross_encoder: bool = False,

410 skip_tokenizer_init: bool = False,

411 auto_cls: type[_BaseAutoModelClass] = AutoModelForCausalLM,

412 + tokenizer_name: str | None = None,

413 + processor: Any | None = None,

414 ) -> None:

415 model_name = maybe_model_redirect(model_name)

416 self.model_name = model_name

@@ -484,20 +490,27 @@ class HfRunner:

490 if not skip_tokenizer_init:

491 self.tokenizer: "PreTrainedTokenizer | PreTrainedTokenizerFast" = (

492 AutoTokenizer.from_pretrained(

487 - model_name,

493 + tokenizer_name or model_name,

494 trust_remote_code=trust_remote_code,

495 )

496 )

497

492 - # don't put this import at the top level

493 - # it will call torch.accelerator.device_count()

494 - from transformers import AutoProcessor

498 + if processor is not None:

499 + self.processor = processor

500 + else:

501 + # don't put this import at the top level

502 + # it will call torch.accelerator.device_count()

503 + from transformers import AutoProcessor

504

496 - self.processor = AutoProcessor.from_pretrained(

497 - model_name,

498 - trust_remote_code=trust_remote_code,

499 - )

505 + self.processor = AutoProcessor.from_pretrained(

506 + model_name,

507 + trust_remote_code=trust_remote_code,

508 + )

509 if skip_tokenizer_init:

510 + if self.processor is None:

511 + raise ValueError(

512 + "skip_tokenizer_init=True requires processor initialization."

513 + )

514 self.tokenizer = self.processor.tokenizer

515

516 def get_inputs(

@@ -520,6 +533,12 @@ class HfRunner:

533 all_inputs: list[BatchFeature | BatchEncoding | dict[str, torch.Tensor]] = []

534 for i, prompt in enumerate(prompts):

535 if isinstance(prompt, str):

536 + if self.processor is None:

537 + raise RuntimeError(

538 + "HfRunner.processor is not initialized. "

539 + "Pass processor=... to HfRunner or set "

540 + "hf_model.processor before generation."

541 + )

542 # Create a copy to avoid modifying the original dict

543 processor_kwargs = (

544 tokenization_kwargs.copy()

@@ -617,6 +636,10 @@ class HfRunner:

636 use_cache=True,

637 **kwargs,

638 )

639 + if self.processor is None:

640 + raise RuntimeError(

641 + "HfRunner.processor is not initialized; cannot decode output."

642 + )

643 output_str = self.processor.batch_decode(

644 output_ids,

645 skip_special_tokens=True,

tests/models/multimodal/generation/test_common.py

+26

@@ -805,6 +805,32 @@ VLM_TEST_SETTINGS = {

805 max_num_seqs=2,

806 patch_hf_runner=model_utils.molmo_patch_hf_runner,

807 ),

808 + "moondream3": VLMTestInfo(

809 + models=["moondream/moondream3-preview"],

810 + test_type=VLMTestType.IMAGE,

811 + prompt_formatter=identity,

812 + img_idx_to_prompt=lambda idx: "<|endoftext|><image>",

813 + # Common-image coverage here targets query/caption. The native

814 + # detect/point skills are not exposed by vLLM.

815 + single_image_prompts=IMAGE_ASSETS.prompts(

816 + {

818 + "cherry_blossom": (

820 + ),

821 + }

822 + ),

823 + max_model_len=4096,

824 + max_num_seqs=2,

825 + dtype="bfloat16",

826 + hf_processor=model_utils.moondream3_processor,

827 + patch_hf_runner=model_utils.moondream3_patch_hf_runner,

828 + # Single size factor to avoid GPU OOM when running multiple test

829 + # cases sequentially (9B MoE model uses ~18 GiB per instance).

830 + image_size_factors=[(1.0,)],

831 + # Moondream3 is 9B params with MoE, needs significant GPU memory

832 + marks=[large_gpu_mark(min_gb=48)],

833 + ),

834 "ovis1_6-gemma2": VLMTestInfo(

835 models=["AIDC-AI/Ovis1.6-Gemma2-9B"],

836 test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),

tests/models/multimodal/generation/test_moondream3.py new

+176

@@ -0,0 +1,176 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Generation tests for Moondream3 query and caption support."""

4 +

5 +import pytest

6 +

7 +from tests.models.registry import HF_EXAMPLE_MODELS

8 +from vllm.platforms import current_platform

9 +

10 +from ....conftest import IMAGE_ASSETS, ImageTestAssets

11 +from ....utils import large_gpu_mark, multi_gpu_test

12 +

13 +MOONDREAM3_MODEL_ID = "moondream/moondream3-preview"

14 +MOONDREAM3_TOKENIZER = "moondream/starmie-v1"

15 +

16 +HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts(

17 + {

20 + }

21 +)

22 +

23 +

24 +def make_query_prompt(question: str) -> str:

25 + """Create a direct-answer query prompt for Moondream3."""

26 + return (

28 + f"{question}<|md_reserved_2|>"

29 + )

30 +

31 +

32 +def make_caption_prompt(length: str = "normal") -> str:

33 + """Create a caption prompt for Moondream3."""

34 + return (

35 + "<|endoftext|><image><|md_reserved_0|>"

36 + f"describe<|md_reserved_1|>{length}<|md_reserved_2|>"

37 + )

38 +

39 +

40 +@multi_gpu_test(num_gpus=2)

41 +@large_gpu_mark(min_gb=80)

42 +def test_tensor_parallel(image_assets: ImageTestAssets):

43 + import gc

44 +

45 + from vllm import LLM, SamplingParams

46 + from vllm.distributed.parallel_state import destroy_model_parallel

47 +

48 + destroy_model_parallel()

49 + gc.collect()

50 + current_platform.empty_cache()

51 +

52 + llm = LLM(

53 + model=MOONDREAM3_MODEL_ID,

54 + tokenizer=MOONDREAM3_TOKENIZER,

55 + trust_remote_code=True,

56 + dtype="bfloat16",

57 + tensor_parallel_size=2,

58 + max_model_len=1024,

59 + enforce_eager=True,

60 + limit_mm_per_prompt={"image": 1},

61 + gpu_memory_utilization=0.45,

62 + )

63 +

64 + image = image_assets[0].pil_image

65 + prompt = make_query_prompt("What color is the stop sign?")

66 +

67 + try:

68 + outputs = llm.generate(

69 + {"prompt": prompt, "multi_modal_data": {"image": image}},

70 + SamplingParams(max_tokens=20, temperature=0),

71 + )

72 +

73 + assert len(outputs) > 0

74 + assert outputs[0].outputs[0].text is not None

75 + finally:

76 + del llm

77 + gc.collect()

78 + current_platform.empty_cache()

79 +

80 +

81 +@pytest.fixture(scope="module")

82 +def llm():

83 + model_info = HF_EXAMPLE_MODELS.get_hf_info("Moondream3ForCausalLM")

84 + model_info.check_transformers_version(on_fail="skip")

85 +

86 + from vllm import LLM

87 +

88 + try:

89 + return LLM(

90 + model=MOONDREAM3_MODEL_ID,

91 + tokenizer=MOONDREAM3_TOKENIZER,

92 + trust_remote_code=True,

93 + dtype="bfloat16",

94 + max_model_len=2048,

95 + enforce_eager=True,

96 + limit_mm_per_prompt={"image": 1},

97 + gpu_memory_utilization=0.45,

98 + )

99 + except Exception as exc:

100 + pytest.skip(f"Failed to load {MOONDREAM3_MODEL_ID}: {exc}")

101 +

102 +

103 +@large_gpu_mark(min_gb=48)

104 +def test_model_loading(llm):

105 + assert llm is not None

106 +

107 +

108 +@large_gpu_mark(min_gb=48)

109 +def test_query_skill(llm, image_assets: ImageTestAssets):

110 + from vllm import SamplingParams

111 +

112 + image = image_assets[0].pil_image

113 + prompt = make_query_prompt("What color is the stop sign?")

114 +

115 + outputs = llm.generate(

116 + {"prompt": prompt, "multi_modal_data": {"image": image}},

117 + SamplingParams(max_tokens=50, temperature=0),

118 + )

119 +

120 + output_text = outputs[0].outputs[0].text

121 + assert output_text is not None

122 + assert len(output_text) > 0

123 +

124 +

125 +@large_gpu_mark(min_gb=48)

126 +def test_caption_skill(llm, image_assets: ImageTestAssets):

127 + from vllm import SamplingParams

128 +

129 + image = image_assets[1].pil_image

130 + prompt = make_caption_prompt()

131 +

132 + outputs = llm.generate(

133 + {"prompt": prompt, "multi_modal_data": {"image": image}},

134 + SamplingParams(max_tokens=100, temperature=0),

135 + )

136 +

137 + output_text = outputs[0].outputs[0].text

138 + assert output_text is not None

139 + assert len(output_text) > 0

140 +

141 +

142 +@large_gpu_mark(min_gb=48)

143 +def test_batched_inference(llm, image_assets: ImageTestAssets):

144 + from vllm import SamplingParams

145 +

146 + images = [asset.pil_image for asset in image_assets]

147 + prompts = [

148 + {"prompt": prompt, "multi_modal_data": {"image": img}}

149 + for img, prompt in zip(images, HF_IMAGE_PROMPTS)

150 + ]

151 +

152 + outputs = llm.generate(prompts, SamplingParams(max_tokens=50, temperature=0))

153 +

154 + assert len(outputs) == len(images)

155 + for output in outputs:

156 + assert output.outputs[0].text is not None

157 + assert len(output.outputs[0].text) > 0

158 +

159 +

160 +@pytest.mark.parametrize("asset_name", ["stop_sign", "cherry_blossom"])

161 +@large_gpu_mark(min_gb=48)

162 +def test_image_assets(llm, image_assets: ImageTestAssets, asset_name: str):

163 + from vllm import SamplingParams

164 +

165 + asset_idx = 0 if asset_name == "stop_sign" else 1

166 + image = image_assets[asset_idx].pil_image

167 + prompt = HF_IMAGE_PROMPTS[asset_idx]

168 +

169 + outputs = llm.generate(

170 + {"prompt": prompt, "multi_modal_data": {"image": image}},

171 + SamplingParams(max_tokens=50, temperature=0),

172 + )

173 +

174 + output_text = outputs[0].outputs[0].text

175 + assert output_text is not None

176 + assert len(output_text) > 0

tests/models/multimodal/generation/vlm_utils/core.py

+12 -1

@@ -38,6 +38,7 @@ def run_test(

38 limit_mm_per_prompt: dict[str, int],

39 vllm_runner_kwargs: dict[str, Any] | None,

40 hf_model_kwargs: dict[str, Any] | None,

41 + hf_processor: Callable[[str], Any] | None,

42 patch_hf_runner: Callable[[HfRunner], HfRunner] | None,

43 runner: RunnerOption = "auto",

44 distributed_executor_backend: str | None = None,

@@ -116,8 +117,18 @@ def run_test(

117 )

118 vllm_outputs_per_mm.append(vllm_output)

119

120 + hf_runner_kwargs: dict[str, Any] = {}

121 + if model_info.tokenizer:

122 + hf_runner_kwargs["tokenizer_name"] = model_info.tokenizer

123 + if hf_processor is not None:

124 + hf_runner_kwargs["processor"] = hf_processor(model)

125 +

126 hf_model = hf_runner(

120 - model, dtype=dtype, auto_cls=auto_cls, model_kwargs=hf_model_kwargs

127 + model,

128 + dtype=dtype,

129 + auto_cls=auto_cls,

130 + model_kwargs=hf_model_kwargs,

131 + **hf_runner_kwargs,

132 )

133

134 # Some models need to patch things like the model processor, e.g., internvl

tests/models/multimodal/generation/vlm_utils/model_utils.py

+218

@@ -1336,3 +1336,221 @@ def voxtral_patch_hf_runner(hf_model: "HfRunner") -> "HfRunner":

1336 hf_model.get_inputs = patched_get_inputs # type: ignore[method-assign, assignment]

1337 hf_model.model.generate = patched_generate # type: ignore[method-assign]

1338 return hf_model

1339 +

1340 +

1341 +def moondream3_processor(model: str):

1342 + from vllm.transformers_utils.processors.moondream3 import Moondream3Processor

1343 +

1344 + return Moondream3Processor.from_pretrained(model, trust_remote_code=True)

1345 +

1346 +

1347 +def moondream3_patch_hf_runner(hf_model: HfRunner) -> HfRunner:

1348 + """Patch HfRunner for Moondream3."""

1349 + moondream_processor = hf_model.processor

1350 +

1351 + def processor(*args, text="", images=None, **kwargs):

1352 + if images is None:

1353 + return moondream_processor(text=text, **kwargs)

1354 +

1355 + images_list = [images] if isinstance(images, Image) else images

1356 + return moondream_processor(images=images_list, text=text, **kwargs)

1357 +

1358 + hf_model.processor = processor

1359 +

1360 + # Expose the LM head for logprob extraction.

1361 + hf_model.model.get_output_embeddings = lambda: hf_model.model.model.text.lm_head

1362 +

1363 + native_model = hf_model.model.model # MoondreamModel instance

1364 +

1365 + from torch.nn import functional as F

1366 +

1367 + from vllm.model_executor.models.moondream3 import reconstruct_from_crops

1368 +

1369 + # Resolve the placeholder tokens from the tokenizer instead of hard-coding.

1370 + image_placeholder_ids = moondream_processor.tokenizer.encode(

1371 + "<image>", add_special_tokens=False

1372 + )

1373 +

1374 + def _normalize_tiling(tilings):

1375 + """Extract (h, w) tuple from various tiling container formats."""

1376 + tiling = tilings

1377 + if isinstance(tiling, torch.Tensor):

1378 + tiling = tuple(tiling.squeeze().tolist())

1379 + elif isinstance(tiling, (list, tuple)):

1380 + t0 = tiling[0]

1381 + if isinstance(t0, torch.Tensor):

1382 + tiling = tuple(t0.tolist())

1383 + elif isinstance(t0, (list, tuple)):

1384 + tiling = tuple(t0)

1385 + return tiling

1386 +

1387 + def _encode_vision(pixel_values, tilings):

1388 + """Run preprocessed crops through vision encoder + projection."""

1389 + device = native_model.device

1390 + dtype = native_model.vision.pos_emb.dtype

1391 + config = native_model.config

1392 +

1393 + pv = pixel_values

1394 + while pv.dim() > 4:

1395 + pv = pv.squeeze(0)

1396 + pv = pv.to(device=device, dtype=dtype)

1397 +

1398 + features = native_model._vis_enc(pv)

1399 + grid_size = config.vision.crop_size // config.vision.enc_patch_size

1400 + global_feat = features[0]

1401 +

1402 + if features.shape[0] > 1 and tilings is not None:

1403 + tiling = _normalize_tiling(tilings)

1404 + local = features[1:].view(-1, grid_size, grid_size, config.vision.enc_dim)

1405 + reconstructed = reconstruct_from_crops(

1406 + local,

1407 + tiling,

1408 + config.vision.overlap_margin,

1409 + patch_size=1,

1410 + )

1411 + else:

1412 + reconstructed = global_feat.view(

1413 + grid_size, grid_size, config.vision.enc_dim

1414 + )

1415 +

1416 + return native_model._vis_proj(global_feat, reconstructed)

1417 +

1418 + def _find_subsequence(seq, subseq):

1419 + """Find start index of subseq in seq, or None."""

1420 + n = len(subseq)

1421 + for i in range(len(seq) - n + 1):

1422 + if seq[i : i + n] == subseq:

1423 + return i

1424 + return None

1425 +

1426 + def _generate(

1427 + self,

1428 + input_ids=None,

1429 + pixel_values=None,

1430 + tilings=None,

1431 + attention_mask=None,

1432 + **kwargs,

1433 + ):

1434 + max_new_tokens = kwargs.get("max_new_tokens", 128)

1435 + return_dict = kwargs.get("return_dict_in_generate", False)

1436 + output_hs = kwargs.get("output_hidden_states", False)

1437 +

1438 + if pixel_values is None:

1439 + sequences = input_ids

1440 + if return_dict:

1441 + return types.SimpleNamespace(

1442 + sequences=sequences,

1443 + hidden_states=() if output_hs else None,

1444 + )

1445 + return sequences

1446 +

1447 + # Processor may return lists; extract the single element.

1448 + if isinstance(pixel_values, (list, tuple)):

1449 + pixel_values = pixel_values[0]

1450 + if (

1451 + isinstance(tilings, (list, tuple))

1452 + and tilings

1453 + and not isinstance(tilings[0], int)

1454 + ):

1455 + tilings = tilings[0]

1456 +

1457 + hf_model.model._setup_caches()

1458 + native_model.use_flex_decoding = False

1459 +

1460 + device = native_model.device

1461 + config = native_model.config

1462 +

1463 + with torch.inference_mode():

1464 + for block in native_model.text.blocks:

1465 + block.kv_cache.k_cache.zero_()

1466 + block.kv_cache.v_cache.zero_()

1467 +

1468 + img_emb = _encode_vision(pixel_values, tilings)

1469 +

1470 + bos_emb = F.embedding(

1471 + torch.tensor([[config.tokenizer.bos_id]], device=device),

1472 + native_model.text.wte,

1473 + )

1474 + img_input = torch.cat([bos_emb, img_emb.unsqueeze(0)], dim=1)

1475 + prefix_len = img_input.size(1)

1476 +

1477 + mask = native_model.attn_mask[:, :, :prefix_len, :]

1478 + pos_ids = torch.arange(prefix_len, dtype=torch.long, device=device)

1479 + native_model._prefill(img_input, mask, pos_ids, None)

1480 +

1481 + ids = input_ids.squeeze(0).tolist()

1482 + img_start = _find_subsequence(ids, image_placeholder_ids)

1483 +

1484 + if img_start is None:

1485 + sequences = input_ids

1486 + if return_dict:

1487 + return types.SimpleNamespace(

1488 + sequences=sequences,

1489 + hidden_states=() if output_hs else None,

1490 + )

1491 + return sequences

1492 +

1493 + prompt_tokens = ids[img_start + len(image_placeholder_ids) :]

1494 +

1495 + if not prompt_tokens:

1496 + sequences = input_ids

1497 + if return_dict:

1498 + return types.SimpleNamespace(

1499 + sequences=sequences,

1500 + hidden_states=() if output_hs else None,

1501 + )

1502 + return sequences

1503 +

1504 + prompt_tensor = torch.tensor([prompt_tokens], device=device)

1505 + prompt_emb = F.embedding(prompt_tensor, native_model.text.wte)

1506 + prompt_len = prompt_emb.size(1)

1507 +

1508 + mask = native_model.attn_mask[:, :, prefix_len : prefix_len + prompt_len, :]

1509 + pos_ids = torch.arange(

1510 + prefix_len,

1511 + prefix_len + prompt_len,

1512 + dtype=torch.long,

1513 + device=device,

1514 + )

1515 + hidden = native_model._prefill(prompt_emb, mask, pos_ids, None)

1516 + pos = prefix_len + prompt_len

1517 +

1518 + hidden_last = native_model.text.post_ln(hidden[:, -1:, :])

1519 + logits = native_model.text.lm_head(hidden_last.squeeze(1))

1520 +

1521 + generated = []

1522 + all_hidden_states = []

1523 + # Record the hidden state that predicted each generated token.

1524 + prev_hs = hidden_last

1525 + for _ in range(max_new_tokens):

1526 + next_token = logits.argmax(dim=-1).item()

1527 + if next_token == 0:

1528 + break

1529 + generated.append(next_token)

1530 + if output_hs:

1531 + all_hidden_states.append((prev_hs,))

1532 +

1533 + next_emb = F.embedding(

1534 + torch.tensor([[next_token]], device=device),

1535 + native_model.text.wte,

1536 + )

1537 + mask = native_model.attn_mask[:, :, pos : pos + 1, :]

1538 + pos_ids_step = torch.tensor([pos], dtype=torch.long, device=device)

1539 + hidden = native_model._prefill(next_emb, mask, pos_ids_step, None)

1540 + hidden_last = native_model.text.post_ln(hidden[:, -1:, :])

1541 + prev_hs = hidden_last

1542 + logits = native_model.text.lm_head(hidden_last.squeeze(1))

1543 + pos += 1

1544 +

1545 + result_ids = ids + generated

1546 + sequences = torch.tensor([result_ids], device=device)

1547 +

1548 + if return_dict:

1549 + return types.SimpleNamespace(

1550 + sequences=sequences,

1551 + hidden_states=tuple(all_hidden_states) if output_hs else None,

1552 + )

1553 + return sequences

1554 +

1555 + hf_model.model.generate = types.MethodType(_generate, hf_model.model)

1556 + return hf_model

tests/models/multimodal/generation/vlm_utils/types.py

@@ -133,6 +133,7 @@ class VLMTestInfo(NamedTuple):

133

134 # Exposed options for HF runner

135 hf_model_kwargs: dict[str, Any] | None = None

136 + hf_processor: Callable[[str], Any] | None = None

137 # Indicates we should explicitly pass the EOS from the tokenizer

138 use_tokenizer_eos: bool = False

139 auto_cls: type[_BaseAutoModelClass] = AutoModelForCausalLM

@@ -196,6 +197,7 @@ class VLMTestInfo(NamedTuple):

197 "comparator": self.comparator,

198 "get_stop_token_ids": self.get_stop_token_ids,

199 "hf_model_kwargs": self.hf_model_kwargs,

200 + "hf_processor": self.hf_processor,

201 "stop_str": self.stop_str,

202 "patch_hf_runner": self.patch_hf_runner,

203 }

tests/models/multimodal/processing/test_moondream3.py new

+553

@@ -0,0 +1,553 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Tests for Moondream3 multimodal processing.

4 +

5 +Includes:

6 +- Processor creation and application tests

7 +- Image tokenization and placeholder expansion tests

8 +- Tiling and cropping logic tests (CPU-based)

9 +- Pixel normalization tests

10 +"""

11 +

12 +import numpy as np

13 +import pytest

14 +import torch

15 +

16 +from vllm.multimodal import MULTIMODAL_REGISTRY

17 +

18 +from ....conftest import ImageTestAssets

19 +from ...utils import build_model_context

20 +

21 +MOONDREAM3_MODEL_ID = "moondream/moondream3-preview"

22 +# Expected multimodal prefix: BOS + 729 image tokens.

23 +EXPECTED_IMAGE_TOKENS = 730

24 +# Vision encoder constants

25 +CROP_SIZE = 378

26 +PATCH_SIZE = 14

27 +MAX_CROPS = 12

28 +

29 +

30 +@pytest.mark.parametrize("model_id", [MOONDREAM3_MODEL_ID])

31 +def test_processor_creation(model_id: str):

32 + """Test that Moondream3 processor can be created."""

33 + ctx = build_model_context(

34 + model_id,

35 + limit_mm_per_prompt={"image": 1},

36 + )

37 + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)

38 + assert processor is not None

39 +

40 +

41 +@pytest.mark.parametrize("model_id", [MOONDREAM3_MODEL_ID])

42 +def test_processor_apply(

43 + image_assets: ImageTestAssets,

44 + model_id: str,

45 +):

46 + """Test that Moondream3 processor can process inputs.

47 +

48 + NOTE: The prompt includes the leading BOS token because Moondream3

49 + pre-fills BOS and image embeddings together.

50 + """

51 + ctx = build_model_context(

52 + model_id,

53 + limit_mm_per_prompt={"image": 1},

54 + )

55 + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)

56 +

58 + mm_data = {"image": [image_assets[0].pil_image]}

59 +

60 + processed_inputs = processor(

61 + prompt,

62 + mm_items=processor.info.parse_mm_data(mm_data),

63 + hf_processor_mm_kwargs={},

64 + )

65 +

66 + assert "prompt_token_ids" in processed_inputs

67 + image_placeholders = processed_inputs["mm_placeholders"]["image"]

68 + assert len(image_placeholders) == 1

69 + assert image_placeholders[0].length == EXPECTED_IMAGE_TOKENS

70 +

71 +

72 +@pytest.mark.parametrize("model_id", [MOONDREAM3_MODEL_ID])

73 +def test_processor_pixel_values(

74 + image_assets: ImageTestAssets,

75 + model_id: str,

76 +):

77 + """Test that pixel values are correctly produced."""

78 + ctx = build_model_context(

79 + model_id,

80 + limit_mm_per_prompt={"image": 1},

81 + )

82 + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)

83 +

85 + mm_data = {"image": [image_assets[0].pil_image]}

86 +

87 + processed_inputs = processor(

88 + prompt,

89 + mm_items=processor.info.parse_mm_data(mm_data),

90 + hf_processor_mm_kwargs={},

91 + )

92 +

93 + # Check mm_kwargs contains pixel_values

94 + mm_kwargs = processed_inputs.get("mm_kwargs")

95 + assert mm_kwargs is not None

96 + mm_data_result = mm_kwargs.get_data()

97 + assert "pixel_values" in mm_data_result

98 +

99 + # Verify pixel_values shape

100 + pixel_values = mm_data_result["pixel_values"]

101 + assert pixel_values.dim() == 5 # [batch, num_crops, C, H, W]

102 + assert pixel_values.shape[2] == 3 # RGB channels

103 + assert pixel_values.shape[3] == 378 # crop height

104 + assert pixel_values.shape[4] == 378 # crop width

105 +

106 +

107 +@pytest.mark.parametrize("model_id", [MOONDREAM3_MODEL_ID])

108 +def test_processor_image_token_expansion(

109 + image_assets: ImageTestAssets,

110 + model_id: str,

111 +):

112 + """Test that <image> placeholder is expanded to correct number of tokens."""

113 + ctx = build_model_context(

114 + model_id,

115 + limit_mm_per_prompt={"image": 1},

116 + )

117 + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)

118 +

120 + mm_data = {"image": [image_assets[0].pil_image]}

121 +

122 + processed_inputs = processor(

123 + prompt,

124 + mm_items=processor.info.parse_mm_data(mm_data),

125 + hf_processor_mm_kwargs={},

126 + )

127 + image_placeholders = processed_inputs["mm_placeholders"]["image"]

128 + assert len(image_placeholders) == 1

129 + assert image_placeholders[0].length == EXPECTED_IMAGE_TOKENS

130 +

131 +

132 +@pytest.mark.parametrize("model_id", [MOONDREAM3_MODEL_ID])

133 +def test_multi_crop_tiling(

134 + model_id: str,

135 +):

136 + """Test that large images produce correct multi-crop tiling."""

137 + from PIL import Image

138 +

139 + from vllm.transformers_utils.processors.moondream3 import Moondream3Processor

140 +

141 + processor = Moondream3Processor.from_pretrained(model_id, trust_remote_code=True)

142 +

143 + # Create a large image that requires multiple crops

144 + large_image = Image.new("RGB", (1000, 1000), color="blue")

145 + pixel_values, tiling = processor.preprocess_image(large_image)

146 +

147 + # Large images should produce more than 1x1 tiling

148 + assert tiling[0] >= 1 and tiling[1] >= 1

149 + # Check that we have global crop + local crops

150 + expected_crops = tiling[0] * tiling[1] + 1

151 + assert pixel_values.shape[0] == expected_crops

152 +

153 +

154 +@pytest.mark.parametrize(

155 + "image_size",

156 + [

157 + (500, 500),

158 + (800, 600),

159 + (1920, 1080),

160 + ],

161 +)

162 +@pytest.mark.parametrize("model_id", [MOONDREAM3_MODEL_ID])

163 +def test_tiling_various_sizes(

164 + image_size: tuple[int, int],

165 + model_id: str,

166 +):

167 + """Test tiling with various image sizes."""

168 + from PIL import Image

169 +

170 + from vllm.transformers_utils.processors.moondream3 import Moondream3Processor

171 +

172 + processor = Moondream3Processor.from_pretrained(model_id, trust_remote_code=True)

173 +

174 + width, height = image_size

175 + image = Image.new("RGB", (width, height), color="red")

176 + pixel_values, tiling = processor.preprocess_image(image)

177 +

178 + # Basic shape checks

179 + assert pixel_values.dim() == 4 # [num_crops, C, H, W]

180 + assert pixel_values.shape[1] == 3 # RGB

181 + assert pixel_values.shape[2] == 378 # crop height

182 + assert pixel_values.shape[3] == 378 # crop width

183 +

184 + # Tiling should respect max_crops (12)

185 + assert tiling[0] * tiling[1] <= 12

186 +

187 +

188 +@pytest.mark.parametrize("model_id", [MOONDREAM3_MODEL_ID])

189 +def test_pixel_normalization(

190 + model_id: str,

191 +):

192 + """Test that pixel values are normalized to [-1, 1] range."""

193 + from PIL import Image

194 +

195 + from vllm.transformers_utils.processors.moondream3 import Moondream3Processor

196 +

197 + processor = Moondream3Processor.from_pretrained(model_id, trust_remote_code=True)

198 +

199 + # Create test image

200 + image = Image.new("RGB", (378, 378), color="green")

201 + pixel_values, _ = processor.preprocess_image(image)

202 +

203 + # Normalization: (x - 0.5) / 0.5 = 2*x - 1

204 + # For input [0, 1], output should be [-1, 1]

205 + assert pixel_values.min() >= -1.0

206 + assert pixel_values.max() <= 1.0

207 +

208 +

209 +@pytest.mark.parametrize("model_id", [MOONDREAM3_MODEL_ID])

210 +def test_chat_template_with_image(

211 + image_assets: ImageTestAssets,

212 + model_id: str,

213 +):

214 + """Test that chat template correctly formats BOS + image + prompt."""

215 + ctx = build_model_context(

216 + model_id,

217 + limit_mm_per_prompt={"image": 1},

218 + )

219 + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)

220 + tokenizer = ctx.tokenizer

221 +

222 + # Use the chat template format

224 + mm_data = {"image": [image_assets[0].pil_image]}

225 +

226 + processed_inputs = processor(

227 + prompt,

228 + mm_items=processor.info.parse_mm_data(mm_data),

229 + hf_processor_mm_kwargs={},

230 + )

231 + token_ids = processed_inputs["prompt_token_ids"]

232 +

233 + # BOS token (<|endoftext|>) should be token ID 0

234 + bos_token_id = tokenizer.encode("<|endoftext|>", add_special_tokens=False)[0]

235 + assert bos_token_id == 0

236 +

237 + # First token should be BOS

238 + assert token_ids[0] == bos_token_id

239 +

240 +

241 +@pytest.mark.parametrize(

242 + "content",

243 + [

244 + pytest.param(

245 + [

246 + {

247 + "type": "image_url",

248 + "image_url": {"url": "https://example.invalid/image.png"},

249 + },

250 + {"type": "text", "text": "What is in this image?"},

251 + ],

252 + id="image-first",

253 + ),

254 + pytest.param(

255 + [

256 + {"type": "text", "text": "What is in this image?"},

257 + {

258 + "type": "image_url",

259 + "image_url": {"url": "https://example.invalid/image.png"},

260 + },

261 + ],

262 + id="text-first",

263 + ),

264 + ],

265 +)

266 +@pytest.mark.parametrize("model_id", [MOONDREAM3_MODEL_ID])

267 +def test_chat_template_content_list_uses_moondream_image_prefix(

268 + image_assets: ImageTestAssets,

269 + content: list[dict[str, object]],

270 + model_id: str,

271 +):

272 + ctx = build_model_context(

273 + model_id,

274 + limit_mm_per_prompt={"image": 1},

275 + )

276 + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)

277 + hf_processor = processor.info.get_hf_processor()

278 +

279 + prompt = hf_processor.tokenizer.apply_chat_template(

280 + [{"role": "user", "content": content}],

281 + chat_template=hf_processor.chat_template,

282 + tokenize=False,

283 + )

284 +

285 + expected_prompt = (

287 + "What is in this image?<|md_reserved_2|>"

288 + )

289 + assert prompt == expected_prompt

290 +

291 + processed_inputs = processor(

292 + prompt,

293 + mm_items=processor.info.parse_mm_data({"image": [image_assets[0].pil_image]}),

294 + hf_processor_mm_kwargs={},

295 + )

296 + image_placeholders = processed_inputs["mm_placeholders"]["image"]

297 + assert len(image_placeholders) == 1

298 + assert image_placeholders[0].length == EXPECTED_IMAGE_TOKENS

299 +

300 +

301 +@pytest.mark.parametrize("model_id", [MOONDREAM3_MODEL_ID])

302 +def test_bos_token_always_first(

303 + image_assets: ImageTestAssets,

304 + model_id: str,

305 +):

306 + """Test that BOS token (ID 0) is always at position 0."""

307 + ctx = build_model_context(

308 + model_id,

309 + limit_mm_per_prompt={"image": 1},

310 + )

311 + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)

312 +

313 + # Start with BOS token explicitly

315 + mm_data = {"image": [image_assets[0].pil_image]}

316 +

317 + processed_inputs = processor(

318 + prompt,

319 + mm_items=processor.info.parse_mm_data(mm_data),

320 + hf_processor_mm_kwargs={},

321 + )

322 + token_ids = processed_inputs["prompt_token_ids"]

323 +

324 + # Token ID 0 (<|endoftext|>) should be the first token

325 + assert token_ids[0] == 0, (

326 + f"Expected BOS token (0) at position 0, got {token_ids[0]}"

327 + )

328 +

329 +

330 +@pytest.mark.parametrize("model_id", [MOONDREAM3_MODEL_ID])

331 +def test_processor_with_small_image(

332 + model_id: str,

333 +):

334 + """Test processor with image smaller than crop size."""

335 + from PIL import Image

336 +

337 + from vllm.transformers_utils.processors.moondream3 import Moondream3Processor

338 +

339 + processor = Moondream3Processor.from_pretrained(model_id, trust_remote_code=True)

340 +

341 + # Small image (smaller than crop size)

342 + small_image = Image.new("RGB", (100, 100), color="yellow")

343 + pixel_values, tiling = processor.preprocess_image(small_image)

344 +

345 + # Small images should use 1x1 tiling

346 + assert tiling == (1, 1)

347 + # Should have 2 crops (global + 1 local)

348 + assert pixel_values.shape[0] == 2

349 +

350 +

351 +@pytest.mark.parametrize(

352 + "image_kind",

353 + [

354 + pytest.param("numpy_hwc", id="numpy-hwc"),

355 + pytest.param("numpy_chw", id="numpy-chw"),

356 + pytest.param("torch_chw", id="torch-chw"),

357 + ],

358 +)

359 +@pytest.mark.parametrize("model_id", [MOONDREAM3_MODEL_ID])

360 +def test_preprocess_image_accepts_non_pil_inputs(

361 + image_assets: ImageTestAssets,

362 + image_kind: str,

363 + model_id: str,

364 +):

365 + from vllm.transformers_utils.processors.moondream3 import Moondream3Processor

366 +

367 + processor = Moondream3Processor.from_pretrained(model_id, trust_remote_code=True)

368 + pil_image = image_assets[0].pil_image.convert("RGB")

369 + hwc_array = np.asarray(pil_image)

370 + expected_pixel_values, expected_tiling = processor.preprocess_image(pil_image)

371 +

372 + if image_kind == "numpy_hwc":

373 + image = hwc_array

374 + elif image_kind == "numpy_chw":

375 + image = np.transpose(hwc_array, (2, 0, 1))

376 + else:

377 + image = torch.from_numpy(np.transpose(hwc_array, (2, 0, 1)).copy())

378 +

379 + pixel_values, tiling = processor.preprocess_image(image)

380 +

381 + assert tiling == expected_tiling

382 + assert pixel_values.shape == expected_pixel_values.shape

383 + assert pixel_values.dtype == torch.bfloat16

384 + assert torch.equal(pixel_values, expected_pixel_values)

385 +

386 +

387 +@pytest.mark.parametrize("image_kind", ["numpy_chw", "torch_chw"])

388 +@pytest.mark.parametrize("model_id", [MOONDREAM3_MODEL_ID])

389 +def test_processor_apply_accepts_non_pil_image_inputs(

390 + image_assets: ImageTestAssets,

391 + image_kind: str,

392 + model_id: str,

393 +):

394 + ctx = build_model_context(

395 + model_id,

396 + limit_mm_per_prompt={"image": 1},

397 + )

398 + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)

399 +

401 + hwc_array = np.asarray(image_assets[0].pil_image.convert("RGB"))

402 + chw_array = np.transpose(hwc_array, (2, 0, 1)).copy()

403 + image = chw_array if image_kind == "numpy_chw" else torch.from_numpy(chw_array)

404 +

405 + processed_inputs = processor(

406 + prompt,

407 + mm_items=processor.info.parse_mm_data({"image": [image]}),

408 + hf_processor_mm_kwargs={},

409 + )

410 +

411 + image_placeholders = processed_inputs["mm_placeholders"]["image"]

412 + assert len(image_placeholders) == 1

413 + assert image_placeholders[0].length == EXPECTED_IMAGE_TOKENS

414 +

415 + mm_kwargs = processed_inputs["mm_kwargs"].get_data()

416 + assert mm_kwargs["pixel_values"].shape[2:] == (3, 378, 378)

417 +

418 +

419 +class TestMoondream3TilingLogic:

420 + """CPU-based tests for Moondream3 tiling selection logic.

421 +

422 + These tests validate the select_tiling() function which determines

423 + how images are divided into crops for the vision encoder.

424 + """

425 +

426 + def test_small_image_no_tiling(self):

427 + """Small images should use 1x1 tiling."""

428 + from vllm.transformers_utils.processors.moondream3 import select_tiling

429 +

430 + tiling = select_tiling(

431 + height=300, width=300, crop_size=CROP_SIZE, max_crops=MAX_CROPS

432 + )

433 + assert tiling == (1, 1)

434 +

435 + def test_exact_crop_size(self):

436 + """Image exactly at crop size should use 1x1."""

437 + from vllm.transformers_utils.processors.moondream3 import select_tiling

438 +

439 + tiling = select_tiling(

440 + height=CROP_SIZE, width=CROP_SIZE, crop_size=CROP_SIZE, max_crops=MAX_CROPS

441 + )

442 + assert tiling == (1, 1)

443 +

444 + def test_large_square_image(self):

445 + """Large square image should use multiple tiles."""

446 + from vllm.transformers_utils.processors.moondream3 import select_tiling

447 +

448 + tiling = select_tiling(

449 + height=800, width=800, crop_size=CROP_SIZE, max_crops=MAX_CROPS

450 + )

451 + h_tiles, w_tiles = tiling

452 + assert h_tiles >= 2

453 + assert w_tiles >= 2

454 + assert h_tiles * w_tiles <= MAX_CROPS

455 +

456 + def test_wide_image(self):

457 + """Wide image should have more width tiles."""

458 + from vllm.transformers_utils.processors.moondream3 import select_tiling

459 +

460 + tiling = select_tiling(

461 + height=400, width=1200, crop_size=CROP_SIZE, max_crops=MAX_CROPS

462 + )

463 + h_tiles, w_tiles = tiling

464 + assert w_tiles >= h_tiles

465 +

466 + def test_tall_image(self):

467 + """Tall image should have more height tiles."""

468 + from vllm.transformers_utils.processors.moondream3 import select_tiling

469 +

470 + tiling = select_tiling(

471 + height=1200, width=400, crop_size=CROP_SIZE, max_crops=MAX_CROPS

472 + )

473 + h_tiles, w_tiles = tiling

474 + assert h_tiles >= w_tiles

475 +

476 + def test_respects_max_crops(self):

477 + """Tiling should not exceed max_crops."""

478 + from vllm.transformers_utils.processors.moondream3 import select_tiling

479 +

480 + tiling = select_tiling(

481 + height=2000, width=2000, crop_size=CROP_SIZE, max_crops=4

482 + )

483 + h_tiles, w_tiles = tiling

484 + assert h_tiles * w_tiles <= 4

485 +

486 +

487 +class TestMoondream3VisionShapes:

488 + """CPU-based tests for vision encoder expected shapes.

489 +

490 + These tests verify the mathematical relationships between

491 + crop size, patch size, and token counts.

492 + """

493 +

494 + def test_expected_patch_count(self):

495 + """Test 378/14 = 27 patches per side, 729 total."""

496 + patches_per_side = CROP_SIZE // PATCH_SIZE

497 + total_patches = patches_per_side**2

498 +

499 + assert patches_per_side == 27

500 + assert total_patches == EXPECTED_IMAGE_TOKENS - 1

501 +

502 + def test_patch_embedding_input_dim(self):

503 + """Test patch embedding input dimension."""

504 + channels = 3

505 + input_dim = PATCH_SIZE * PATCH_SIZE * channels

506 +

507 + assert input_dim == 14 * 14 * 3

508 + assert input_dim == 588

509 +

510 +

511 +class TestMoondream3TauAttention:

512 + """CPU-based tests for tau attention scaling components.

513 +

514 + These tests validate the tau attention formula used in Moondream3:

515 + - Token-based: tok_q = tanh(gelu(qkv) @ tau_wq.T)

516 + - Position-based: tau_pos = 1 + (sigmoid(alpha * log(pos+1)) - 0.5)

517 + """

518 +

519 + def test_tau_position_range(self):

520 + """Test tau position scaling produces values in valid range."""

521 + num_heads = 32

522 + seq_len = 100

523 +

524 + tau_alpha = torch.randn(num_heads)

525 + positions = torch.arange(seq_len)

526 +

527 + pos_float = (positions.float() + 1.0).clamp(min=1e-6)

528 + pos_log = pos_float.log()

529 + tau_pos = 1.0 + (torch.sigmoid(tau_alpha[:, None] * pos_log[None, :]) - 0.5)

530 +

531 + assert tau_pos.shape == (num_heads, seq_len)

532 + # tau_pos should be between 0.5 and 1.5

533 + assert tau_pos.min() >= 0.5

534 + assert tau_pos.max() <= 1.5

535 +

536 + def test_tau_token_output_range(self):

537 + """Test tau token scaling output is bounded by tanh."""

538 + import torch.nn.functional as F

539 +

540 + seq_len = 100

541 + qkv_dim = 6144 # 2048 * 3

542 + num_heads = 32

543 +

544 + qkv = torch.randn(seq_len, qkv_dim)

545 + tau_wq = torch.randn(num_heads, qkv_dim)

546 +

547 + tok_feat = F.gelu(qkv)

548 + tok_q = torch.tanh(tok_feat @ tau_wq.t())

549 +

550 + assert tok_q.shape == (seq_len, num_heads)

551 + # tanh output is bounded by [-1, 1]

552 + assert tok_q.min() >= -1.0

553 + assert tok_q.max() <= 1.0

tests/models/registry.py

+10

@@ -1122,6 +1122,16 @@ _MULTIMODAL_EXAMPLE_MODELS = {

1122 extras={"olmo": "allenai/Molmo-7B-O-0924"},

1123 trust_remote_code=True,

1124 ),

1125 + "Moondream3ForCausalLM": _HfExamplesInfo(

1126 + "moondream/moondream3-preview",

1127 + tokenizer="moondream/starmie-v1",

1128 + trust_remote_code=True,

1129 + ),

1130 + "HfMoondream": _HfExamplesInfo(

1131 + "moondream/moondream3-preview",

1132 + tokenizer="moondream/starmie-v1",

1133 + trust_remote_code=True,

1134 + ),

1135 "Molmo2ForConditionalGeneration": _HfExamplesInfo(

1136 "allenai/Molmo2-8B",

1137 extras={"olmo": "allenai/Molmo2-O-7B"},

tests/models/utils.py

+14

@@ -523,6 +523,17 @@ def dummy_hf_overrides(

523

524 text_config.update(update_dict)

525

526 + # Update n_layers and moe configs for Moondream3 model

527 + if model_arch in ("Moondream3ForCausalLM", "HfMoondream"):

528 + text_config.update(

529 + {

530 + "n_layers": num_hidden_layers,

531 + "moe_num_experts": num_experts,

532 + "moe_experts_per_token": 2,

533 + "moe_start_layer": num_hidden_layers,

534 + }

535 + )

536 +

537 if hasattr(hf_config, "vision_config"):

538 hf_config.vision_config.update(

539 {

@@ -531,6 +542,9 @@ def dummy_hf_overrides(

542 }

543 )

544

545 + if model_arch in ("Moondream3ForCausalLM", "HfMoondream"):

546 + hf_config.vision_config.update({"enc_n_layers": 1})

547 +

548 # e.g.: ibm-granite/granite-speech-3.3-2b

549 if hasattr(hf_config, "encoder_config"):

550 hf_config.encoder_config.update(

vllm/model_executor/models/moondream3.py new

+1423

@@ -0,0 +1,1423 @@

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 BaseDummyOptions

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 MoEActivation, fused_experts

32 +from vllm.model_executor.layers.fused_moe.config import biased_moe_quant_config

33 +from vllm.model_executor.layers.linear import (

34 + ColumnParallelLinear,

35 + QKVParallelLinear,

36 + ReplicatedLinear,

37 + RowParallelLinear,

38 +)

39 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

40 +from vllm.model_executor.layers.quantization import QuantizationConfig

41 +from vllm.model_executor.layers.rotary_embedding import get_rope

42 +from vllm.model_executor.layers.vocab_parallel_embedding import (

43 + ParallelLMHead,

44 + VocabParallelEmbedding,

45 +)

46 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader

47 +from vllm.multimodal import MULTIMODAL_REGISTRY

48 +from vllm.multimodal.inputs import (

49 + MultiModalFieldConfig,

50 + MultiModalKwargsItems,

51 +)

52 +from vllm.multimodal.parse import ImageSize, MultiModalDataItems

53 +from vllm.multimodal.processing import (

54 + BaseDummyInputsBuilder,

55 + BaseMultiModalProcessor,

56 + BaseProcessingInfo,

57 + PromptReplacement,

58 + PromptUpdate,

59 + PromptUpdateDetails,

60 +)

61 +from vllm.sequence import IntermediateTensors

62 +from vllm.transformers_utils.configs.moondream3 import (

63 + Moondream3Config,

64 + Moondream3TextConfig,

65 + Moondream3VisionConfig,

66 +)

67 +from vllm.transformers_utils.processors.moondream3 import Moondream3Processor

68 +

69 +from .interfaces import (

70 + MultiModalEmbeddings,

71 + SupportsMultiModal,

72 + SupportsPP,

73 +)

74 +from .utils import (

75 + extract_layer_index,

76 + make_empty_intermediate_tensors_factory,

77 + make_layers,

78 + maybe_prefix,

79 +)

80 +

81 +logger = init_logger(__name__)

82 +

83 +

84 +# ============================================================================

85 +# Image Processing Utilities

86 +# ============================================================================

87 +

88 +

89 +def reconstruct_from_crops(

90 + crops: torch.Tensor,

91 + tiling: tuple[int, int],

92 + overlap_margin: int,

93 + patch_size: int = 14,

94 +) -> torch.Tensor:

95 + """Reconstruct features from overlapping crops."""

96 + tiling_h, tiling_w = tiling

97 + crop_height, crop_width = crops[0].shape[:2]

98 + margin_pixels = overlap_margin * patch_size

99 +

100 + output_h = (crop_height - 2 * margin_pixels) * tiling_h + 2 * margin_pixels

101 + output_w = (crop_width - 2 * margin_pixels) * tiling_w + 2 * margin_pixels

102 +

103 + reconstructed = torch.zeros(

104 + (output_h, output_w, crops[0].shape[2]),

105 + device=crops[0].device,

106 + dtype=crops[0].dtype,

107 + )

108 +

109 + for i, crop in enumerate(crops):

110 + tile_y = i // tiling_w

111 + tile_x = i % tiling_w

112 +

113 + x_start = 0 if tile_x == 0 else margin_pixels

114 + x_end = crop_width if tile_x == tiling_w - 1 else crop_width - margin_pixels

115 + y_start = 0 if tile_y == 0 else margin_pixels

116 + y_end = crop_height if tile_y == tiling_h - 1 else crop_height - margin_pixels

117 +

118 + out_x = tile_x * (crop_width - 2 * margin_pixels)

119 + out_y = tile_y * (crop_height - 2 * margin_pixels)

120 +

121 + reconstructed[

122 + out_y + y_start : out_y + y_end, out_x + x_start : out_x + x_end

123 + ] = crop[y_start:y_end, x_start:x_end]

124 +

125 + return reconstructed

126 +

127 +

128 +# ============================================================================

129 +# Vision Encoder Components

130 +# ============================================================================

131 +

132 +

133 +class Moondream3VisionMLP(nn.Module):

134 + """MLP for vision encoder blocks."""

135 +

136 + def __init__(

137 + self,

138 + hidden_size: int,

139 + intermediate_size: int,

140 + quant_config: QuantizationConfig | None = None,

141 + prefix: str = "",

142 + ):

143 + super().__init__()

144 + self.fc1 = ColumnParallelLinear(

145 + hidden_size,

146 + intermediate_size,

147 + bias=True,

148 + quant_config=quant_config,

149 + prefix=f"{prefix}.fc1",

150 + )

151 + self.act = get_act_fn("gelu_pytorch_tanh")

152 + self.fc2 = RowParallelLinear(

153 + intermediate_size,

154 + hidden_size,

155 + bias=True,

156 + quant_config=quant_config,

157 + prefix=f"{prefix}.fc2",

158 + )

159 +

160 + def forward(self, x: torch.Tensor) -> torch.Tensor:

161 + x, _ = self.fc1(x)

162 + x = self.act(x)

163 + x, _ = self.fc2(x)

164 + return x

165 +

166 +

167 +class Moondream3VisionAttention(nn.Module):

168 + """Self-attention for vision encoder (bidirectional)."""

169 +

170 + def __init__(

171 + self,

172 + hidden_size: int,

173 + num_heads: int,

174 + quant_config: QuantizationConfig | None = None,

175 + prefix: str = "",

176 + ):

177 + super().__init__()

178 + self.hidden_size = hidden_size

179 + self.num_heads = num_heads

180 + self.head_dim = hidden_size // num_heads

181 +

182 + self.qkv_proj = QKVParallelLinear(

183 + hidden_size=hidden_size,

184 + head_size=self.head_dim,

185 + total_num_heads=num_heads,

186 + bias=True,

187 + quant_config=quant_config,

188 + prefix=f"{prefix}.qkv_proj",

189 + )

190 + self.out_proj = RowParallelLinear(

191 + input_size=hidden_size,

192 + output_size=hidden_size,

193 + bias=True,

194 + quant_config=quant_config,

195 + prefix=f"{prefix}.out_proj",

196 + )

197 +

198 + tp_size = get_tensor_model_parallel_world_size()

199 + self.num_heads_per_partition = num_heads // tp_size

200 +

201 + self.attn = MMEncoderAttention(

202 + num_heads=self.num_heads_per_partition,

203 + head_size=self.head_dim,

204 + scale=self.head_dim**-0.5,

205 + prefix=f"{prefix}.attn",

206 + )

207 +

208 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

209 + qkv, _ = self.qkv_proj(hidden_states)

210 + q, k, v = qkv.chunk(3, dim=-1)

211 + out = self.attn(q, k, v)

212 + out, _ = self.out_proj(out)

213 + return out

214 +

215 +

216 +class Moondream3VisionBlock(nn.Module):

217 + """Transformer block for vision encoder."""

218 +

219 + def __init__(

220 + self,

221 + hidden_size: int,

222 + intermediate_size: int,

223 + num_heads: int,

224 + quant_config: QuantizationConfig | None = None,

225 + prefix: str = "",

226 + ):

227 + super().__init__()

228 + self.ln1 = nn.LayerNorm(hidden_size, eps=1e-5)

229 + self.attn = Moondream3VisionAttention(

230 + hidden_size=hidden_size,

231 + num_heads=num_heads,

232 + quant_config=quant_config,

233 + prefix=f"{prefix}.attn",

234 + )

235 + self.ln2 = nn.LayerNorm(hidden_size, eps=1e-5)

236 + self.mlp = Moondream3VisionMLP(

237 + hidden_size=hidden_size,

238 + intermediate_size=intermediate_size,

239 + quant_config=quant_config,

240 + prefix=f"{prefix}.mlp",

241 + )

242 +

243 + def forward(self, x: torch.Tensor) -> torch.Tensor:

244 + x = x + self.attn(self.ln1(x))

245 + x = x + self.mlp(self.ln2(x))

246 + return x

247 +

248 +

249 +class Moondream3VisionEncoder(nn.Module):

250 + """Vision encoder (SigLIP-style ViT)."""

251 +

252 + def __init__(

253 + self,

254 + config: Moondream3VisionConfig,

255 + quant_config: QuantizationConfig | None = None,

256 + prefix: str = "",

257 + ):

258 + super().__init__()

259 + self.config = config

260 +

261 + # Patch embedding

262 + self.patch_emb = nn.Linear(

263 + config.enc_patch_size * config.enc_patch_size * 3,

264 + config.enc_dim,

265 + bias=True,

266 + )

267 +

268 + # Position embeddings (27x27 = 729 patches for 378x378 / 14)

269 + num_patches = (config.crop_size // config.enc_patch_size) ** 2

270 + self.pos_emb = nn.Parameter(torch.zeros(1, num_patches, config.enc_dim))

271 +

272 + # Transformer blocks

273 + self.blocks = nn.ModuleList(

274 + [

275 + Moondream3VisionBlock(

276 + hidden_size=config.enc_dim,

277 + intermediate_size=config.enc_ff_dim,

278 + num_heads=config.enc_n_heads,

279 + quant_config=quant_config,

280 + prefix=f"{prefix}.blocks.{i}",

281 + )

282 + for i in range(config.enc_n_layers)

283 + ]

284 + )

285 +

286 + self.post_ln = nn.LayerNorm(config.enc_dim, eps=1e-5)

287 +

288 + def create_patches(self, images: torch.Tensor) -> torch.Tensor:

289 + """Convert images to patch embeddings.

290 +

291 + Args:

292 + images: (batch, channels, height, width)

293 +

294 + Returns:

295 + patches: (batch, num_patches, patch_dim)

296 + """

297 + patch_size = self.config.enc_patch_size

298 + batch, channels, height, width = images.shape

299 + patches_h = height // patch_size

300 + patches_w = width // patch_size

301 +

302 + # Unfold into patches

303 + patches = images.unfold(2, patch_size, patch_size).unfold(

304 + 3, patch_size, patch_size

305 + )

306 + # (batch, channels, patches_h, patches_w, patch_size, patch_size)

307 + patches = patches.permute(0, 2, 3, 1, 4, 5).contiguous()

308 + # (batch, patches_h, patches_w, channels, patch_size, patch_size)

309 + patches = patches.view(batch, patches_h * patches_w, -1)

310 + # (batch, num_patches, channels * patch_size * patch_size)

311 +

312 + return patches

313 +

314 + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:

315 + """Encode images.

316 +

317 + Args:

318 + pixel_values: (batch, channels, height, width)

319 +

320 + Returns:

321 + features: (batch, num_patches, hidden_size)

322 + """

323 + # Create patches and embed

324 + patches = self.create_patches(pixel_values)

325 + x = self.patch_emb(patches)

326 +

327 + # Add position embeddings

328 + x = x + self.pos_emb

329 +

330 + # Apply transformer blocks

331 + for block in self.blocks:

332 + x = block(x)

333 +

334 + # Final layer norm

335 + x = self.post_ln(x)

336 +

337 + return x

338 +

339 +

340 +class Moondream3VisionProjection(nn.Module):

341 + """Projects vision features to text embedding dimension."""

342 +

343 + def __init__(

344 + self,

345 + input_dim: int,

346 + inner_dim: int,

347 + output_dim: int,

348 + quant_config: QuantizationConfig | None = None,

349 + prefix: str = "",

350 + ):

351 + super().__init__()

352 + # Input is concatenated global and local features (2 * input_dim)

353 + self.fc1 = ColumnParallelLinear(

354 + input_dim * 2,

355 + inner_dim,

356 + bias=True,

357 + quant_config=quant_config,

358 + prefix=f"{prefix}.fc1",

359 + )

360 + self.act = get_act_fn("gelu_pytorch_tanh")

361 + self.fc2 = RowParallelLinear(

362 + inner_dim,

363 + output_dim,

364 + bias=True,

365 + quant_config=quant_config,

366 + prefix=f"{prefix}.fc2",

367 + )

368 +

369 + def forward(self, x: torch.Tensor) -> torch.Tensor:

370 + x, _ = self.fc1(x)

371 + x = self.act(x)

372 + x, _ = self.fc2(x)

373 + return x

374 +

375 +

376 +# ============================================================================

377 +# Text Decoder Components

378 +# ============================================================================

379 +

380 +

381 +class Moondream3TextMLP(nn.Module):

382 + """Standard MLP for non-MoE layers (layers 0-3)."""

383 +

384 + def __init__(

385 + self,

386 + hidden_size: int,

387 + intermediate_size: int,

388 + quant_config: QuantizationConfig | None = None,

389 + prefix: str = "",

390 + ):

391 + super().__init__()

392 + self.fc1 = ColumnParallelLinear(

393 + hidden_size,

394 + intermediate_size,

395 + bias=True,

396 + quant_config=quant_config,

397 + prefix=f"{prefix}.fc1",

398 + )

399 + self.act = get_act_fn("gelu_pytorch_tanh")

400 + self.fc2 = RowParallelLinear(

401 + intermediate_size,

402 + hidden_size,

403 + bias=True,

404 + quant_config=quant_config,

405 + prefix=f"{prefix}.fc2",

406 + )

407 +

408 + def forward(self, x: torch.Tensor) -> torch.Tensor:

409 + x, _ = self.fc1(x)

410 + x = self.act(x)

411 + x, _ = self.fc2(x)

412 + return x

413 +

414 +

415 +class Moondream3TextMoE(nn.Module):

416 + """Mixture of Experts layer for layers 4+ with expert parallelism.

417 +

418 + Moondream3 uses a custom GeGLU activation: gelu(h) * (g + 1)

419 + where fc1 outputs [gate, up] and the activation is gelu(gate) * (up + 1).

420 +

421 + Uses expert parallelism where each GPU stores num_experts/tp_size experts.

422 + Routing and communication handled via all-to-all or replicated computation.

423 +

424 + Checkpoint format:

425 + - fc1.weight: [num_experts, expert_inner_dim * 2, hidden_size] (gate+up)

426 + - fc2.weight: [num_experts, hidden_size, expert_inner_dim] (down)

427 + - router.weight: [num_experts, hidden_size]

428 + - router.bias: [num_experts]

429 + """

430 +

431 + def __init__(

432 + self,

433 + hidden_size: int,

434 + expert_inner_dim: int,

435 + num_experts: int,

436 + experts_per_token: int,

437 + quant_config: QuantizationConfig | None = None,

438 + prefix: str = "",

439 + ):

440 + super().__init__()

441 + self.hidden_size = hidden_size

442 + self.expert_inner_dim = expert_inner_dim

443 + self.num_experts = num_experts

444 + self.experts_per_token = experts_per_token

445 +

446 + # Expert parallelism: each GPU stores a subset of experts

447 + self.tp_size = get_tensor_model_parallel_world_size()

448 + self.experts_per_rank = num_experts // self.tp_size

449 + self.num_local_experts = self.experts_per_rank

450 +

451 + # Router (gate) - use ReplicatedLinear for compatibility

452 + self.gate = ReplicatedLinear(

453 + hidden_size,

454 + num_experts,

455 + bias=True,

456 + quant_config=None,

457 + prefix=f"{prefix}.gate",

458 + )

459 +

460 + # Local expert weights (only store experts_per_rank experts)

461 + # fc1: [experts_per_rank, expert_inner_dim * 2, hidden_size]

462 + # fc2: [experts_per_rank, hidden_size, expert_inner_dim]

463 + self.fc1_weight = nn.Parameter(

464 + torch.empty(self.num_local_experts, expert_inner_dim * 2, hidden_size)

465 + )

466 + self.fc2_weight = nn.Parameter(

467 + torch.empty(self.num_local_experts, hidden_size, expert_inner_dim)

468 + )

469 + self._use_fused_moe = True

470 +

471 + local_expert_start = get_tensor_model_parallel_rank() * self.experts_per_rank

472 + expert_map = torch.full((num_experts,), -1, dtype=torch.int32)

473 + expert_map[local_expert_start : local_expert_start + self.num_local_experts] = (

474 + torch.arange(self.num_local_experts, dtype=torch.int32)

475 + )

476 + self.register_buffer("_expert_map", expert_map, persistent=False)

477 +

478 + # Preserve Moondream3's exact GeGLU variant (gelu(h) * (g + 1)) by

479 + # adding +1 bias to the second half of the fused fc1 activations.

480 + fused_w1_bias = torch.zeros(

481 + self.num_local_experts,

482 + expert_inner_dim * 2,

483 + dtype=torch.float32,

484 + )

485 + fused_w1_bias[:, expert_inner_dim:] = 1.0

486 + self.register_buffer("_fused_w1_bias", fused_w1_bias, persistent=False)

487 +

488 + def forward(self, x: torch.Tensor) -> torch.Tensor:

489 + """Forward pass with expert parallelism and custom GeGLU activation."""

490 +

491 + # Get router logits and compute top-k

492 + router_logits, _ = self.gate(x) # [num_tokens, num_experts]

493 + topk_logits, topk_ids = torch.topk(

494 + router_logits, self.experts_per_token, dim=-1

495 + )

496 + # Softmax over selected experts

497 + topk_weights = F.softmax(topk_logits, dim=-1, dtype=torch.float32).to(x.dtype)

498 +

499 + if self._use_fused_moe and x.is_cuda:

500 + try:

501 + out = fused_experts(

502 + hidden_states=x.contiguous(),

503 + w1=self.fc1_weight,

504 + w2=self.fc2_weight,

505 + topk_weights=topk_weights.contiguous(),

506 + topk_ids=topk_ids.contiguous(),

507 + activation=MoEActivation.GELU,

508 + global_num_experts=self.num_experts,

509 + expert_map=self._expert_map,

510 + quant_config=biased_moe_quant_config(self._fused_w1_bias, None),

511 + )

512 + out = tensor_model_parallel_all_reduce(out)

513 + return out

514 + except (NotImplementedError, RuntimeError) as exc:

515 + self._use_fused_moe = False

516 + logger.warning_once(

517 + "Disabling fused Moondream3 MoE path and falling back to "

518 + "the Python expert loop: %s",

519 + str(exc),

520 + )

521 +

522 + tp_rank = get_tensor_model_parallel_rank()

523 + # Compute local expert range

524 + local_expert_start = tp_rank * self.experts_per_rank

525 +

526 + # Fallback path for environments where fused kernels are unavailable.

527 + out = x.new_zeros(x.shape)

528 +

529 + for local_expert_idx in range(self.num_local_experts):

530 + global_expert_id = local_expert_start + local_expert_idx

531 +

532 + # Find tokens assigned to this expert

533 + token_pos, which_k = (topk_ids == global_expert_id).nonzero(as_tuple=True)

534 + if token_pos.numel() == 0:

535 + continue

536 +

537 + # Get tokens and their routing weights

538 + x_tok = x.index_select(0, token_pos) # [n_tokens, hidden_size]

539 + gate_tok = topk_weights[token_pos, which_k] # [n_tokens]

540 +

541 + # fc1: [expert_inner_dim * 2, hidden_size]

542 + # h_full: [n_tokens, expert_inner_dim * 2]

543 + h_full = F.linear(x_tok, self.fc1_weight[local_expert_idx])

544 +

545 + # GeGLU with (g + 1): h, g = split; output = gelu(h) * (g + 1)

546 + # HF MoE uses exact GELU (not tanh approximation).

547 + h, g = h_full.chunk(2, dim=-1) # Each [n_tokens, expert_inner_dim]

548 + h = F.gelu(h) * (g + 1.0)

549 +

550 + # fc2: [hidden_size, expert_inner_dim]

551 + # y: [n_tokens, hidden_size]

552 + y = F.linear(h, self.fc2_weight[local_expert_idx])

553 +

554 + # Apply routing weight

555 + y = y * gate_tok.unsqueeze(-1)

556 +

557 + # Accumulate output

558 + out.index_add_(0, token_pos, y)

559 +

560 + # All-reduce to combine results from all experts across GPUs

561 + out = tensor_model_parallel_all_reduce(out)

562 +

563 + return out

564 +

565 +

566 +class Moondream3Attention(nn.Module):

567 + """Decoder attention with RoPE and tau scaling.

568 +

569 + Moondream3 uses a tau attention mechanism that scales Q and V

570 + based on both token content and position.

571 + """

572 +

573 + def __init__(

574 + self,

575 + config: Moondream3TextConfig,

576 + layer_idx: int,

577 + cache_config=None,

578 + quant_config: QuantizationConfig | None = None,

579 + prefix: str = "",

580 + ):

581 + super().__init__()

582 + self.hidden_size = config.dim

583 + self.num_heads = config.n_heads

584 + self.num_kv_heads = config.n_kv_heads

585 + self.head_dim = config.dim // config.n_heads

586 +

587 + tp_size = get_tensor_model_parallel_world_size()

588 + self.num_heads_per_partition = self.num_heads // tp_size

589 + self.num_kv_heads_per_partition = max(1, self.num_kv_heads // tp_size)

590 +

591 + self.qkv_proj = QKVParallelLinear(

592 + hidden_size=self.hidden_size,

593 + head_size=self.head_dim,

594 + total_num_heads=self.num_heads,

595 + total_num_kv_heads=self.num_kv_heads,

596 + bias=True,

597 + quant_config=quant_config,

598 + prefix=f"{prefix}.qkv_proj",

599 + )

600 +

601 + self.out_proj = RowParallelLinear(

602 + input_size=self.hidden_size,

603 + output_size=self.hidden_size,

604 + bias=True,

605 + quant_config=quant_config,

606 + prefix=f"{prefix}.out_proj",

607 + )

608 +

609 + # Moondream uses 32-dim rotation out of 64-dim head (partial_rotary_factor=0.5)

610 + # HF Moondream uses non-interleaved RoPE (split by half)

611 + # In vLLM, is_neox_style=True means split by half (GPT-NeoX style)

612 + rope_parameters = {

613 + "rope_theta": config.rope_theta,

614 + "partial_rotary_factor": 32 / self.head_dim, # 32/64 = 0.5

615 + }

616 + self.rotary_emb = get_rope(

617 + head_size=self.head_dim,

618 + max_position=config.max_context,

619 + rope_parameters=rope_parameters,

620 + is_neox_style=True, # Moondream uses split-by-half (GPT-NeoX) style

621 + )

622 +

623 + self.scaling = self.head_dim**-0.5

624 + self.attn = Attention(

625 + num_heads=self.num_heads_per_partition,

626 + head_size=self.head_dim,

627 + scale=self.scaling,

628 + num_kv_heads=self.num_kv_heads_per_partition,

629 + cache_config=cache_config,

630 + quant_config=quant_config,

631 + prefix=f"{prefix}.attn",

632 + )

633 +

634 + # Tau scaling parameters for position-dependent attention

635 + # These are learned during training to modulate attention based on position

636 + # tau_wq and tau_wv need full qkv_dim for correct computation

637 + # Only heads are partitioned, qkv dimension is kept full for all-gather

638 + qkv_dim = self.hidden_size * 3 # Q + K + V dimension (full)

639 + self.tau_alpha = nn.Parameter(torch.zeros(self.num_heads_per_partition))

640 + self.tau_wq = nn.Parameter(torch.zeros(self.num_heads_per_partition, qkv_dim))

641 + self.tau_wv = nn.Parameter(torch.zeros(self.num_heads_per_partition, qkv_dim))

642 + self.tp_size = tp_size

643 +

644 + def forward(

645 + self,

646 + positions: torch.Tensor,

647 + hidden_states: torch.Tensor,

648 + ) -> torch.Tensor:

649 + qkv, _ = self.qkv_proj(hidden_states)

650 +

651 + q, k, v = qkv.split(

652 + [

653 + self.num_heads_per_partition * self.head_dim,

654 + self.num_kv_heads_per_partition * self.head_dim,

655 + self.num_kv_heads_per_partition * self.head_dim,

656 + ],

657 + dim=-1,

658 + )

659 +

660 + # Apply tau scaling to Q and V

661 + # Tau scaling has two components:

662 + # 1. Token-based: tok_q = tanh(gelu(qkv) @ tau_wq.T)

663 + # 2. Position-based: tau_pos = 1 + (sigmoid(alpha * log(pos+1)) - 0.5)

664 + # Final: tau = tok + tau_pos

665 + #

666 + # For TP, tau weights are sharded by head, but qkv_dim is kept full

667 +

668 + # Get full qkv for tau computation

669 + # With TP, reconstruct qkv in correct layout [q_full, k_full, v_full]

670 + # (all-gather would produce [q_0, k_0, v_0, q_1, k_1, v_1] - wrong)

671 + if self.tp_size > 1:

672 + # All-gather once, then reconstruct [q_full, k_full, v_full].

673 + qkv_full_sharded = tensor_model_parallel_all_gather(qkv.contiguous())

674 + q_local_dim = q.shape[-1]

675 + kv_local_dim = k.shape[-1]

676 + qkv_full_sharded = qkv_full_sharded.view(

677 + qkv.shape[0],

678 + self.tp_size,

679 + q_local_dim + 2 * kv_local_dim,

680 + )

681 + q_full = qkv_full_sharded[:, :, :q_local_dim].reshape(qkv.shape[0], -1)

682 + k_full = qkv_full_sharded[

683 + :, :, q_local_dim : q_local_dim + kv_local_dim

684 + ].reshape(qkv.shape[0], -1)

685 + v_full = qkv_full_sharded[:, :, q_local_dim + kv_local_dim :].reshape(

686 + qkv.shape[0], -1

687 + )

688 + qkv_full = torch.cat([q_full, k_full, v_full], dim=-1).contiguous()

689 + else:

690 + qkv_full = qkv

691 +

692 + # Compute tau scaling factors matching HF implementation exactly:

693 + # tok_feat = gelu(qkv)

694 + # tok_q = tanh(tok_feat @ tau_wq.T) # [num_tokens, num_heads]

695 + # tau_pos = 1 + (sigmoid(alpha * log(pos+1)) - 0.5) # [num_heads, num_tokens]

696 + # tau = (tok_q.T + tau_pos).T # [num_tokens, num_heads]

697 + num_tokens = qkv_full.shape[0]

698 + orig_dtype = q.dtype

699 +

700 + # Token-based component

701 + tok_feat = F.gelu(qkv_full) # Apply GELU activation

702 + tok_q = torch.tanh(tok_feat @ self.tau_wq.t()) # [N, H_per_partition]

703 + tok_v = torch.tanh(tok_feat @ self.tau_wv.t()) # [N, H_per_partition]

704 +

705 + # Position-based component

706 + # tau_pos = 1 + (sigmoid(alpha * log(pos+1)) - 0.5)

707 + # positions is [num_tokens], need to compute for each head

708 + # tau_alpha: [num_heads_per_partition]

709 + pos_float = (positions.to(orig_dtype) + 1.0).clamp(min=1e-6)

710 + pos_log = pos_float.log() # [num_tokens]

711 + # alpha[:, None] * pos_log[None, :] -> [num_heads, num_tokens]

712 + tau_pos = 1.0 + (

713 + torch.sigmoid(self.tau_alpha[:, None] * pos_log[None, :]) - 0.5

714 + ) # [H_per_partition, N]

715 +

716 + # Combine token and position components

717 + tau_q = (tok_q + tau_pos.t()).to(orig_dtype) # [N, H_per_partition]

718 + tau_v = (tok_v + tau_pos.t()).to(orig_dtype) # [N, H_per_partition]

719 +

720 + # Reshape q and v to apply per-head tau scaling

721 + q = q.view(num_tokens, self.num_heads_per_partition, self.head_dim)

722 + v = v.view(num_tokens, self.num_kv_heads_per_partition, self.head_dim)

723 +

724 + # Apply tau scaling

725 + q = q * tau_q.unsqueeze(-1)

726 + v = v * tau_v[:, : self.num_kv_heads_per_partition].unsqueeze(-1)

727 +

728 + # Reshape back

729 + q = q.view(num_tokens, -1)

730 + v = v.view(num_tokens, -1)

731 +

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

733 +

734 + attn_output = self.attn(q, k, v)

735 +

736 + output, _ = self.out_proj(attn_output)

737 + return output

738 +

739 +

740 +class Moondream3DecoderLayer(nn.Module):

741 + """Decoder layer with attention + MLP/MoE."""

742 +

743 + def __init__(

744 + self,

745 + config: Moondream3TextConfig,

746 + cache_config=None,

747 + quant_config: QuantizationConfig | None = None,

748 + prefix: str = "",

749 + ):

750 + super().__init__()

751 + layer_idx = extract_layer_index(prefix)

752 + self.layer_idx = layer_idx

753 +

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

755 +

756 + self.attn = Moondream3Attention(

757 + config=config,

758 + layer_idx=layer_idx,

759 + cache_config=cache_config,

760 + quant_config=quant_config,

761 + prefix=f"{prefix}.attn",

762 + )

763 +

764 + # Use MoE for layers >= moe_start_layer, standard MLP otherwise

765 + if layer_idx >= config.moe_start_layer:

766 + self.mlp = Moondream3TextMoE(

767 + hidden_size=config.dim,

768 + expert_inner_dim=config.moe_expert_inner_dim,

769 + num_experts=config.moe_num_experts,

770 + experts_per_token=config.moe_experts_per_token,

771 + quant_config=quant_config,

772 + prefix=f"{prefix}.mlp",

773 + )

774 + else:

775 + self.mlp = Moondream3TextMLP(

776 + hidden_size=config.dim,

777 + intermediate_size=config.ff_dim,

778 + quant_config=quant_config,

779 + prefix=f"{prefix}.mlp",

780 + )

781 +

782 + def forward(

783 + self,

784 + positions: torch.Tensor,

785 + hidden_states: torch.Tensor,

786 + ) -> torch.Tensor:

787 + # Pre-norm architecture

788 + normed = self.ln(hidden_states)

789 + attn_out = self.attn(positions, normed)

790 + mlp_out = self.mlp(normed)

791 + hidden_states = hidden_states + attn_out + mlp_out

792 + return hidden_states

793 +

794 +

795 +class Moondream3TextModel(nn.Module):

796 + """Text decoder model."""

797 +

798 + def __init__(

799 + self,

800 + config: Moondream3TextConfig,

801 + cache_config=None,

802 + quant_config: QuantizationConfig | None = None,

803 + prefix: str = "",

804 + ):

805 + super().__init__()

806 + self.config = config

807 +

808 + self.wte = VocabParallelEmbedding(

809 + config.vocab_size,

810 + config.dim,

811 + prefix=f"{prefix}.wte",

812 + )

813 +

814 + blocks_prefix = maybe_prefix(prefix, "blocks")

815 + self.start_layer, self.end_layer, self.blocks = make_layers(

816 + config.n_layers,

817 + lambda prefix: Moondream3DecoderLayer(

818 + config=config,

819 + cache_config=cache_config,

820 + quant_config=quant_config,

821 + prefix=prefix,

822 + ),

823 + prefix=blocks_prefix,

824 + )

825 +

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

827 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

828 + ["hidden_states"], config.dim

829 + )

830 +

831 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

832 + return self.wte(input_ids)

833 +

834 + def forward(

835 + self,

836 + input_ids: torch.Tensor | None,

837 + positions: torch.Tensor,

838 + intermediate_tensors: IntermediateTensors | None = None,

839 + inputs_embeds: torch.Tensor | None = None,

840 + ) -> torch.Tensor | IntermediateTensors:

841 + pp_group = get_pp_group()

842 + if pp_group.is_first_rank:

843 + if inputs_embeds is not None:

844 + hidden_states = inputs_embeds

845 + else:

846 + assert input_ids is not None

847 + hidden_states = self.embed_input_ids(input_ids)

848 + else:

849 + assert intermediate_tensors is not None

850 + hidden_states = intermediate_tensors["hidden_states"]

851 +

852 + for i, layer in enumerate(

853 + islice(self.blocks, self.start_layer, self.end_layer)

854 + ):

855 + hidden_states = layer(positions, hidden_states)

856 +

857 + if not pp_group.is_last_rank:

858 + return IntermediateTensors({"hidden_states": hidden_states})

859 +

860 + hidden_states = self.post_ln(hidden_states)

861 + return hidden_states

862 +

863 +

864 +@dataclass(frozen=True)

865 +class Moondream3ImageInput:

866 + """Container holding per-image inputs for embedding."""

867 +

868 + pixel_values: torch.Tensor

869 + tiling: tuple[int, int] | None

870 +

871 +

872 +# ============================================================================

873 +# Multimodal Processing

874 +# ============================================================================

875 +

876 +

877 +class Moondream3ProcessingInfo(BaseProcessingInfo):

878 + """Processing info for Moondream3."""

879 +

880 + def get_hf_config(self):

881 + return self.ctx.get_hf_config()

882 +

883 + def get_hf_processor(self, **kwargs: object):

884 + return self.ctx.get_hf_processor(Moondream3Processor, **kwargs)

885 +

886 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

887 + return {"image": 1}

888 +

889 + def get_num_image_tokens(

890 + self,

891 + *,

892 + image_width: int,

893 + image_height: int,

894 + ) -> int:

895 + # HF pre-fills BOS together with the fixed 27x27 vision grid under

896 + # the same bidirectional prefix mask: 1 BOS + 729 image embeddings.

897 + return 730

898 +

899 + def get_image_size_with_most_features(self) -> ImageSize:

900 + return ImageSize(width=378, height=378)

901 +

902 + def get_max_image_tokens(self) -> int:

903 + return 730

904 +

905 + def get_mm_max_tokens_per_item(

906 + self,

907 + seq_len: int,

908 + mm_counts: Mapping[str, int],

909 + ) -> Mapping[str, int]:

910 + return {"image": self.get_max_image_tokens()}

911 +

912 +

913 +class Moondream3DummyInputsBuilder(BaseDummyInputsBuilder[Moondream3ProcessingInfo]):

914 + """Dummy inputs builder for profiling."""

915 +

916 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:

917 + return (

919 + "What is this image?<|md_reserved_2|>"

920 + )

921 +

922 + def get_dummy_mm_data(

923 + self,

924 + seq_len: int,

925 + mm_counts: Mapping[str, int],

926 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

927 + mm_processor_kwargs: Mapping[str, object] | None = None,

928 + ) -> MultiModalDataDict:

929 + num_images = mm_counts.get("image", 0)

930 + return {

931 + "image": self._get_dummy_images(

932 + width=378,

933 + height=378,

934 + num_images=num_images,

935 + )

936 + }

937 +

938 +

939 +class Moondream3MultiModalProcessor(BaseMultiModalProcessor[Moondream3ProcessingInfo]):

940 + """Multimodal processor for Moondream3."""

941 +

942 + image_placeholder: str = "<image>"

943 + bos_image_placeholder: str = "<|endoftext|><image>"

944 +

945 + def _call_hf_processor(

946 + self,

947 + prompt: str,

948 + mm_data: Mapping[str, object],

949 + mm_kwargs: Mapping[str, object],

950 + tok_kwargs: Mapping[str, object],

951 + ) -> BatchFeature:

952 + # Moondream3's processor handles images directly rather than exposing a

953 + # separate `image_processor`, so keep the cache path on text+MM calls.

954 + return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs)

955 +

956 + @cached_property

957 + def bos_image_placeholder_tokens(self) -> list[int]:

958 + tokenizer = self.info.get_tokenizer()

959 + token_ids = tokenizer.encode(

960 + self.bos_image_placeholder,

961 + add_special_tokens=False,

962 + )

963 + if len(token_ids) < 2:

964 + raise ValueError(

965 + "Tokenizer could not encode Moondream3 BOS/image placeholder "

966 + f"{self.bos_image_placeholder!r}."

967 + )

968 + return token_ids

969 +

970 + def _get_mm_fields_config(

971 + self,

972 + hf_inputs: BatchFeature,

973 + hf_processor_mm_kwargs: Mapping[str, object],

974 + ) -> Mapping[str, MultiModalFieldConfig]:

975 + return {

976 + "pixel_values": MultiModalFieldConfig.batched("image"),

977 + "tilings": MultiModalFieldConfig.batched("image", keep_on_cpu=True),

978 + }

979 +

980 + def _hf_processor_applies_updates(

981 + self,

982 + prompt_text: str,

983 + mm_items: MultiModalDataItems,

984 + hf_processor_mm_kwargs: Mapping[str, object],

985 + tokenization_kwargs: Mapping[str, object],

986 + ) -> bool:

987 + # Moondream3 HF processor does NOT expand placeholder tokens.

988 + # vLLM expands BOS + <image> so the whole HF image prefix is marked

989 + # bidirectional by the multimodal prefix-LM mask.

990 + return False

991 +

992 + def _get_prompt_updates(

993 + self,

994 + mm_items: MultiModalDataItems,

995 + hf_processor_mm_kwargs: Mapping[str, object],

996 + out_mm_kwargs: MultiModalKwargsItems,

997 + ) -> list[PromptUpdate]:

998 + image_size = self.info.get_image_size_with_most_features()

999 + num_image_tokens = self.info.get_num_image_tokens(

1000 + image_width=image_size.width,

1001 + image_height=image_size.height,

1002 + )

1003 + placeholder_tokens = self.bos_image_placeholder_tokens

1004 + bos_token = placeholder_tokens[0]

1005 + image_token = placeholder_tokens[-1]

1006 + return [

1007 + PromptReplacement(

1008 + modality="image",

1009 + target=placeholder_tokens,

1010 + replacement=PromptUpdateDetails(

1011 + full=[bos_token] + [image_token] * (num_image_tokens - 1),

1012 + ),

1013 + ),

1014 + ]

1015 +

1016 +

1017 +# ============================================================================

1018 +# Main Model

1019 +# ============================================================================

1020 +

1021 +

1022 +@MULTIMODAL_REGISTRY.register_processor(

1023 + Moondream3MultiModalProcessor,

1024 + info=Moondream3ProcessingInfo,

1025 + dummy_inputs=Moondream3DummyInputsBuilder,

1026 +)

1027 +class Moondream3ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

1028 + """Moondream3 multimodal model for causal language modeling.

1029 +

1030 + vLLM supports the standard autoregressive Moondream3 query and caption

1031 + prompt formats. The region-module point/detect skills require custom

1032 + coordinate decoding and are intentionally not exposed here.

1033 + """

1034 +

1035 + supports_multimodal = True

1036 + packed_modules_mapping = {

1037 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

1038 + }

1039 +

1040 + def __init__(

1041 + self,

1042 + *,

1043 + vllm_config: VllmConfig,

1044 + prefix: str = "",

1045 + ):

1046 + super().__init__()

1047 +

1048 + hf_config = vllm_config.model_config.hf_config

1049 + quant_config = vllm_config.quant_config

1050 + cache_config = vllm_config.cache_config

1051 +

1052 + # Reuse the transformers_utils config implementation.

1053 + if isinstance(hf_config, Moondream3Config):

1054 + self.config = hf_config

1055 + else:

1056 + config_dict = hf_config.config if hasattr(hf_config, "config") else {}

1057 + self.config = Moondream3Config(config=config_dict)

1058 +

1059 + with self._mark_tower_model(vllm_config, "image"):

1060 + # Vision encoder

1061 + self.vision = Moondream3VisionEncoder(

1062 + config=self.config.vision_config,

1063 + quant_config=quant_config,

1064 + prefix=maybe_prefix(prefix, "vision"),

1065 + )

1066 +

1067 + # Vision projection

1068 + self.vision_proj = Moondream3VisionProjection(

1069 + input_dim=self.config.vision_config.enc_dim,

1070 + inner_dim=self.config.vision_config.proj_inner_dim,

1071 + output_dim=self.config.text_config.dim,

1072 + quant_config=quant_config,

1073 + prefix=maybe_prefix(prefix, "vision_proj"),

1074 + )

1075 +

1076 + with self._mark_language_model(vllm_config):

1077 + # Text decoder

1078 + self.text = Moondream3TextModel(

1079 + config=self.config.text_config,

1080 + cache_config=cache_config,

1081 + quant_config=quant_config,

1082 + prefix=maybe_prefix(prefix, "text"),

1083 + )

1084 +

1085 + # LM head (with bias - Moondream3 has lm_head bias)

1086 + self.lm_head = ParallelLMHead(

1087 + self.config.text_config.vocab_size,

1088 + self.config.text_config.dim,

1089 + bias=True,

1090 + quant_config=quant_config,

1091 + prefix=maybe_prefix(prefix, "lm_head"),

1092 + )

1093 +

1094 + self.logits_processor = LogitsProcessor(self.config.text_config.vocab_size)

1095 + self.make_empty_intermediate_tensors = self.text.make_empty_intermediate_tensors

1096 + self._answer_id = getattr(

1097 + self.config,

1098 + "answer_token_id",

1099 + getattr(hf_config, "answer_token_id", 3),

1100 + )

1101 +

1102 + @classmethod

1103 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

1104 + if modality == "image":

1105 + return "<image>"

1106 + return None

1107 +

1108 + def get_language_model(self) -> nn.Module:

1109 + return self.text

1110 +

1111 + def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:

1112 + return num_image_tokens

1113 +

1114 + def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:

1115 + return num_vision_tokens

1116 +

1117 + def _split_pixel_values(

1118 + self,

1119 + pixel_values: object,

1120 + ) -> list[torch.Tensor]:

1121 + # The processor should standardize image inputs into:

1122 + # - torch.Tensor [num_images, num_crops, C, H, W], or

1123 + # - list[torch.Tensor[num_crops, C, H, W]] for ragged crops.

1124 + if isinstance(pixel_values, torch.Tensor):

1125 + if pixel_values.dim() != 5:

1126 + raise ValueError(

1127 + "Expected `pixel_values` tensor with shape "

1128 + "[num_images, num_crops, C, H, W], got "

1129 + f"{tuple(pixel_values.shape)}."

1130 + )

1131 + return [pv.contiguous() for pv in pixel_values]

1132 +

1133 + if isinstance(pixel_values, (list, tuple)):

1134 + tensors: list[torch.Tensor] = []

1135 + for value in pixel_values:

1136 + if not isinstance(value, torch.Tensor):

1137 + raise TypeError(

1138 + "Expected each `pixel_values` element to be a tensor, "

1139 + f"got {type(value)!r}."

1140 + )

1141 + if value.dim() != 4:

1142 + raise ValueError(

1143 + f"Unsupported pixel_values element shape {tuple(value.shape)}."

1144 + )

1145 + tensors.append(value.contiguous())

1146 + return tensors

1147 +

1148 + raise TypeError(

1149 + "pixel_values must be a tensor or a sequence of tensors, "

1150 + f"got {type(pixel_values)!r}."

1151 + )

1152 +

1153 + def _split_tilings(

1154 + self,

1155 + tilings: object,

1156 + expected: int,

1157 + ) -> list[tuple[int, int] | None]:

1158 + if tilings is None:

1159 + return [None] * expected

1160 +

1161 + if isinstance(tilings, torch.Tensor):

1162 + if tilings.dim() != 2 or tilings.shape[1] != 2:

1163 + raise ValueError(

1164 + "Expected `tilings` tensor with shape [num_images, 2], got "

1165 + f"{tuple(tilings.shape)}."

1166 + )

1167 + tiling_items = tilings.tolist()

1168 + elif isinstance(tilings, (list, tuple)):

1169 + tiling_items = list(tilings)

1170 + else:

1171 + raise TypeError(

1172 + "tilings must be None, a tensor or a sequence of tuples, "

1173 + f"got {type(tilings)!r}."

1174 + )

1175 +

1176 + if len(tiling_items) != expected:

1177 + raise ValueError(

1178 + "Mismatch between the number of pixel_values entries "

1179 + f"({expected}) and tilings ({len(tiling_items)})."

1180 + )

1181 +

1182 + normalized: list[tuple[int, int] | None] = []

1183 + for tiling in tiling_items:

1184 + if tiling is None:

1185 + normalized.append(None)

1186 + continue

1187 + if isinstance(tiling, torch.Tensor):

1188 + tiling = tiling.tolist()

1189 + if isinstance(tiling, (list, tuple)) and len(tiling) == 2:

1190 + normalized.append((int(tiling[0]), int(tiling[1])))

1191 + else:

1192 + raise ValueError(

1193 + f"Each tiling entry must be a pair of integers, got {tiling!r}."

1194 + )

1195 + return normalized

1196 +

1197 + def _parse_image_inputs(self, **kwargs: object) -> list[Moondream3ImageInput]:

1198 + pixel_values = kwargs.get("pixel_values")

1199 + if pixel_values is None:

1200 + return []

1201 +

1202 + pixel_values_list = self._split_pixel_values(pixel_values)

1203 + tilings_list = self._split_tilings(

1204 + kwargs.get("tilings"), len(pixel_values_list)

1205 + )

1206 +

1207 + image_inputs: list[Moondream3ImageInput] = []

1208 + for value, tiling in zip(pixel_values_list, tilings_list):

1209 + if value.dim() != 4:

1210 + raise ValueError(

1211 + f"Expected 4D tensor for crops, got {tuple(value.shape)}."

1212 + )

1213 + image_inputs.append(Moondream3ImageInput(pixel_values=value, tiling=tiling))

1214 + return image_inputs

1215 +

1216 + def _encode_image_input(self, image_input: Moondream3ImageInput) -> torch.Tensor:

1217 + pixel_values = image_input.pixel_values

1218 + if pixel_values.dim() != 4:

1219 + raise ValueError(

1220 + f"Expected 4D tensor for crops, got {tuple(pixel_values.shape)}."

1221 + )

1222 +

1223 + device = self.vision.patch_emb.weight.device

1224 + dtype = self.vision.patch_emb.weight.dtype

1225 + pixel_values = pixel_values.to(device=device, dtype=dtype)

1226 +

1227 + features = self.vision(pixel_values)

1228 +

1229 + # Grid size = crop_size / patch_size (e.g., 378 / 14 = 27)

1230 + grid_size = (

1231 + self.config.vision_config.crop_size

1232 + // self.config.vision_config.enc_patch_size

1233 + )

1234 + enc_dim = self.config.vision_config.enc_dim

1235 + global_features = features[0]

1236 +

1237 + if features.shape[0] > 1:

1238 + if image_input.tiling is None:

1239 + raise ValueError(

1240 + "Missing tiling metadata for multi-crop Moondream image."

1241 + )

1242 + local = features[1:].contiguous().view(-1, grid_size, grid_size, enc_dim)

1243 + reconstructed = reconstruct_from_crops(

1244 + local,

1245 + image_input.tiling,

1246 + overlap_margin=self.config.vision_config.overlap_margin,

1247 + patch_size=1,

1248 + )

1249 + else:

1250 + reconstructed = global_features.view(grid_size, grid_size, enc_dim)

1251 +

1252 + recon = reconstructed.permute(2, 0, 1).contiguous()

1253 + # Mirror HF reference behavior: reconstructed local features are pooled

1254 + # to enc_n_layers x enc_n_layers. For moondream3-preview this is 27x27.

1255 + pooled_size = self.config.vision_config.enc_n_layers

1256 + if pooled_size != grid_size:

1257 + logger.warning_once(

1258 + "Moondream3 pooled_size (%d) differs from crop grid (%d). "

1259 + "Using enc_n_layers to match HF reference behavior.",

1260 + pooled_size,

1261 + grid_size,

1262 + )

1263 + recon = F.adaptive_avg_pool2d(recon, output_size=(pooled_size, pooled_size))

1264 + recon = recon.permute(1, 2, 0).contiguous().view(-1, enc_dim)

1265 +

1266 + combined = torch.cat([global_features, recon], dim=-1).unsqueeze(0)

1267 + projected = self.vision_proj(combined).squeeze(0)

1268 +

1269 + # Note: Vision embeddings are already synchronized across TP ranks

1270 + # because the vision projection uses RowParallelLinear which performs

1271 + # all-reduce internally, ensuring identical outputs on all ranks.

1272 +

1273 + return projected

1274 +

1275 + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:

1276 + """Generate the HF image prefix: BOS embedding + 729 image embeddings."""

1277 + image_inputs = self._parse_image_inputs(**kwargs)

1278 + if not image_inputs:

1279 + return []

1280 +

1281 + device = self.vision.patch_emb.weight.device

1282 + bos_ids = torch.tensor([self.config.bos_token_id], device=device)

1283 + bos_embedding = self.text.embed_input_ids(bos_ids)

1284 +

1285 + embeddings: list[torch.Tensor] = []

1286 + for image_input in image_inputs:

1287 + image_embeddings = self._encode_image_input(image_input)

1288 + embeddings.append(

1289 + torch.cat([bos_embedding.to(image_embeddings.dtype), image_embeddings])

1290 + )

1291 + return embeddings

1292 +

1293 + def forward(

1294 + self,

1295 + input_ids: torch.Tensor | None,

1296 + positions: torch.Tensor,

1297 + intermediate_tensors: IntermediateTensors | None = None,

1298 + inputs_embeds: torch.Tensor | None = None,

1299 + **kwargs,

1300 + ) -> torch.Tensor | IntermediateTensors:

1301 + hidden_states = self.text(

1302 + input_ids=input_ids,

1303 + positions=positions,

1304 + intermediate_tensors=intermediate_tensors,

1305 + inputs_embeds=inputs_embeds,

1306 + )

1307 + return hidden_states

1308 +

1309 + def compute_logits(

1310 + self,

1311 + hidden_states: torch.Tensor,

1312 + ) -> torch.Tensor | None:

1313 + logits = self.logits_processor(self.lm_head, hidden_states)

1314 + if logits is not None:

1315 + logits[:, self._answer_id] = float("-inf")

1316 + return logits

1317 +

1318 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

1319 + """Load weights with remapping from HuggingFace format."""

1320 +

1321 + params_dict = dict(self.named_parameters())

1322 + loaded_params: set[str] = set()

1323 +

1324 + # Get expert intermediate size for fc1 splitting

1325 +

1326 + for name, loaded_weight in weights:

1327 + # Map from HF naming to vLLM naming

1328 + # model.vision.* -> vision.*

1329 + # model.text.* -> text.*

1330 + if name.startswith("model."):

1331 + name = name[6:] # Remove "model." prefix

1332 +

1333 + # Specific name mappings

1334 + # Vision projection: vision.proj_mlp.fc1 -> vision_proj.fc1

1335 + name = name.replace("vision.proj_mlp.", "vision_proj.")

1336 +

1337 + # Text embedding: text.wte (no suffix) -> text.wte.weight

1338 + if name == "text.wte":

1339 + name = "text.wte.weight"

1340 +

1341 + # LM head: text.lm_head -> lm_head

1342 + name = name.replace("text.lm_head.", "lm_head.")

1343 +

1344 + # Attention mapping

1345 + name = name.replace(".attn.qkv.", ".attn.qkv_proj.")

1346 + name = name.replace(".attn.proj.", ".attn.out_proj.")

1347 +

1348 + # Tau attention scaling weights

1349 + # HF format: .attn.tau.alpha -> .attn.tau_alpha

1350 + name = name.replace(".attn.tau.alpha", ".attn.tau_alpha")

1351 + name = name.replace(".attn.tau.wq", ".attn.tau_wq")

1352 + name = name.replace(".attn.tau.wv", ".attn.tau_wv")

1353 +

1354 + # MoE router mapping: mlp.router -> mlp.gate

1355 + name = name.replace(".mlp.router.", ".mlp.gate.")

1356 +

1357 + # Handle MoE expert weights for layers 4+ with expert parallelism

1358 + # fc1.weight: [n_experts, expert_inner_dim * 2, hidden_size] (gate+up)

1359 + # fc2.weight: [n_experts, hidden_size, expert_inner_dim] (down)

1360 + # Each GPU stores n_experts/tp_size experts

1361 + # Note: Only 3D weights are MoE, 2D weights are standard MLP

1362 + if ".mlp.fc1.weight" in name and loaded_weight.dim() == 3:

1363 + from vllm.distributed import get_tensor_model_parallel_rank

1364 +

1365 + tp_size = get_tensor_model_parallel_world_size()

1366 + tp_rank = get_tensor_model_parallel_rank()

1367 + num_experts = loaded_weight.shape[0]

1368 + experts_per_rank = num_experts // tp_size

1369 + expert_start = tp_rank * experts_per_rank

1370 + expert_end = expert_start + experts_per_rank

1371 + # Shard by expert dimension

1372 + loaded_weight = loaded_weight[expert_start:expert_end].contiguous()

1373 + # Map to our custom MoE format: mlp.fc1_weight

1374 + name = name.replace(".mlp.fc1.weight", ".mlp.fc1_weight")

1375 +

1376 + if ".mlp.fc2.weight" in name and loaded_weight.dim() == 3:

1377 + from vllm.distributed import get_tensor_model_parallel_rank

1378 +

1379 + tp_size = get_tensor_model_parallel_world_size()

1380 + tp_rank = get_tensor_model_parallel_rank()

1381 + num_experts = loaded_weight.shape[0]

1382 + experts_per_rank = num_experts // tp_size

1383 + expert_start = tp_rank * experts_per_rank

1384 + expert_end = expert_start + experts_per_rank

1385 + # Shard by expert dimension

1386 + loaded_weight = loaded_weight[expert_start:expert_end].contiguous()

1387 + # Map to our custom MoE format: mlp.fc2_weight

1388 + name = name.replace(".mlp.fc2.weight", ".mlp.fc2_weight")

1389 +

1390 + # Handle tau weights with tensor parallelism

1391 + # tau_alpha: [num_heads] -> [num_heads/tp]

1392 + # tau_wq: [num_heads, qkv_dim] -> [num_heads/tp, qkv_dim/tp]

1393 + # tau_wv: [num_heads, qkv_dim] -> [num_heads/tp, qkv_dim/tp]

1394 + if ".tau_alpha" in name:

1395 + from vllm.distributed import get_tensor_model_parallel_rank

1396 +

1397 + tp_size = get_tensor_model_parallel_world_size()

1398 + tp_rank = get_tensor_model_parallel_rank()

1399 + num_heads = loaded_weight.shape[0]

1400 + heads_per_partition = num_heads // tp_size

1401 + start = tp_rank * heads_per_partition

1402 + end = start + heads_per_partition

1403 + loaded_weight = loaded_weight[start:end].contiguous()

1404 +

1405 + if ".tau_wq" in name or ".tau_wv" in name:

1406 + from vllm.distributed import get_tensor_model_parallel_rank

1407 +

1408 + tp_size = get_tensor_model_parallel_world_size()

1409 + tp_rank = get_tensor_model_parallel_rank()

1410 + num_heads, qkv_dim = loaded_weight.shape

1411 + heads_per_partition = num_heads // tp_size

1412 + # Only shard by head dimension, keep full qkv_dim for all-gather

1413 + head_start = tp_rank * heads_per_partition

1414 + head_end = head_start + heads_per_partition

1415 + loaded_weight = loaded_weight[head_start:head_end, :].contiguous()

1416 +

1417 + if name in params_dict:

1418 + param = params_dict[name]

1419 + weight_loader = getattr(param, "weight_loader", default_weight_loader)

1420 + weight_loader(param, loaded_weight)

1421 + loaded_params.add(name)

1422 +

1423 + return loaded_params

vllm/model_executor/models/registry.py

@@ -484,6 +484,8 @@ _MULTIMODAL_MODELS = {

484 ),

485 "MolmoForCausalLM": ("molmo", "MolmoForCausalLM"),

486 "Molmo2ForConditionalGeneration": ("molmo2", "Molmo2ForConditionalGeneration"),

487 + "Moondream3ForCausalLM": ("moondream3", "Moondream3ForCausalLM"),

488 + "HfMoondream": ("moondream3", "Moondream3ForCausalLM"),

489 "NemotronH_Nano_VL_V2": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"),

490 "NemotronH_Nano_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"),

491 "NemotronH_Super_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"),

vllm/transformers_utils/config.py

@@ -114,6 +114,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(

114 mlp_speculator="MLPSpeculatorConfig",

115 medusa="MedusaConfig",

116 midashenglm="MiDashengLMConfig",

117 + moondream3="Moondream3Config",

118 eagle="EAGLEConfig",

119 speculators="SpeculatorsConfig",

120 nemotron="NemotronConfig",

vllm/transformers_utils/configs/__init__.py

@@ -50,6 +50,9 @@ _CLASS_TO_MODULE: dict[str, str] = {

50 "MedusaConfig": "vllm.transformers_utils.configs.medusa",

51 "MiDashengLMConfig": "vllm.transformers_utils.configs.midashenglm",

52 "MLPSpeculatorConfig": "vllm.transformers_utils.configs.mlp_speculator",

53 + "Moondream3Config": "vllm.transformers_utils.configs.moondream3",

54 + "Moondream3TextConfig": "vllm.transformers_utils.configs.moondream3",

55 + "Moondream3VisionConfig": "vllm.transformers_utils.configs.moondream3",

56 "MoonViTConfig": "vllm.transformers_utils.configs.moonvit",

57 "KimiLinearConfig": "vllm.transformers_utils.configs.kimi_linear",

58 "KimiVLConfig": "vllm.transformers_utils.configs.kimi_vl",

@@ -111,6 +114,9 @@ __all__ = [

114 "MedusaConfig",

115 "MiDashengLMConfig",

116 "MLPSpeculatorConfig",

117 + "Moondream3Config",

118 + "Moondream3TextConfig",

119 + "Moondream3VisionConfig",

120 "MoonViTConfig",

121 "KimiLinearConfig",

122 "KimiVLConfig",

vllm/transformers_utils/configs/moondream3.py new

+152

@@ -0,0 +1,152 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Configuration for Moondream3 model."""

4 +

5 +from transformers import PretrainedConfig

6 +

7 +

8 +class Moondream3VisionConfig(PretrainedConfig):

9 + """Vision encoder configuration for Moondream3."""

10 +

11 + model_type = "moondream3_vision"

12 +

13 + def __init__(

14 + self,

15 + enc_dim: int = 1152,

16 + enc_patch_size: int = 14,

17 + enc_n_layers: int = 27,

18 + enc_ff_dim: int = 4304,

19 + enc_n_heads: int = 16,

20 + proj_inner_dim: int = 8192,

21 + crop_size: int = 378,

22 + max_crops: int = 12,

23 + overlap_margin: int = 4,

24 + **kwargs,

25 + ):

26 + super().__init__(**kwargs)

27 + self.enc_dim = enc_dim

28 + self.enc_patch_size = enc_patch_size

29 + self.enc_n_layers = enc_n_layers

30 + self.enc_ff_dim = enc_ff_dim

31 + self.enc_n_heads = enc_n_heads

32 + self.proj_inner_dim = proj_inner_dim

33 + self.crop_size = crop_size

34 + self.max_crops = max_crops

35 + self.overlap_margin = overlap_margin

36 +

37 + # Standard HuggingFace attributes for vision config

38 + self.hidden_size = enc_dim

39 + self.num_attention_heads = enc_n_heads

40 + self.num_hidden_layers = enc_n_layers

41 + self.intermediate_size = enc_ff_dim

42 + self.patch_size = enc_patch_size

43 + self.image_size = crop_size

44 +

45 +

46 +class Moondream3TextConfig(PretrainedConfig):

47 + """Text decoder configuration for Moondream3."""

48 +

49 + model_type = "moondream3_text"

50 +

51 + def __init__(

52 + self,

53 + dim: int = 2048,

54 + ff_dim: int = 8192,

55 + n_layers: int = 24,

56 + vocab_size: int = 51200,

57 + max_context: int = 4096,

58 + n_heads: int = 32,

59 + n_kv_heads: int = 32,

60 + prefix_attn: int = 730,

61 + rope_theta: float = 1500000.0,

62 + moe: dict | None = None,

63 + **kwargs,

64 + ):

65 + super().__init__(**kwargs)

66 +

67 + # Store original moondream3 config names

68 + self.dim = dim

69 + self.ff_dim = ff_dim

70 + self.n_layers = n_layers

71 + self.n_heads = n_heads

72 + self.n_kv_heads = n_kv_heads

73 + self.prefix_attn = prefix_attn

74 + self.max_context = max_context

75 + self.rope_theta = rope_theta

76 +

77 + # MoE config

78 + moe = moe or {}

79 + self.moe_start_layer = moe.get("start_layer", 4)

80 + self.moe_num_experts = moe.get("n_experts", 64)

81 + self.moe_experts_per_token = moe.get("n_experts_per_tok", 8)

82 + self.moe_expert_inner_dim = moe.get("expert_inner_dim", 1024)

83 +

84 + # Standard HuggingFace attributes (required by vLLM)

85 + self.hidden_size = dim

86 + self.num_attention_heads = n_heads

87 + self.num_key_value_heads = n_kv_heads

88 + self.num_hidden_layers = n_layers

89 + self.intermediate_size = ff_dim

90 + self.vocab_size = vocab_size

91 + self.max_position_embeddings = max_context

92 +

93 + # Moondream3 uses token 0 (<|endoftext|>) as both BOS and EOS.

94 + # Token 3 (<|md_reserved_2|>) is an answer delimiter that the model

95 + # implementation suppresses during generation.

96 + self.bos_token_id = 0

97 + self.eos_token_id = 0

98 +

99 + # MoE standard attributes

100 + self.num_local_experts = self.moe_num_experts

101 + self.num_experts_per_tok = self.moe_experts_per_token

102 +

103 +

104 +class Moondream3Config(PretrainedConfig):

105 + """Combined configuration for Moondream3 multimodal model."""

106 +

107 + model_type = "moondream3"

108 + is_composition = True

109 +

110 + def __init__(

111 + self,

112 + config: dict | None = None,

113 + **kwargs,

114 + ):

115 + config = config or {}

116 +

117 + # Parse text config

118 + text_config = config.get("text", {})

119 + self.text_config: Moondream3TextConfig = Moondream3TextConfig(**text_config)

120 +

121 + # Parse vision config

122 + vision_config = config.get("vision", {})

123 + self.vision_config = Moondream3VisionConfig(**vision_config)

124 +

125 + # Store the original config dict for model access

126 + self.config = config

127 + tokenizer_config = config.get("tokenizer", {})

128 + self.answer_token_id = tokenizer_config.get("answer_id", 3)

129 +

130 + super().__init__(**kwargs)

131 +

132 + # Expose key attributes at top level for vLLM compatibility

133 + self.hidden_size = self.text_config.hidden_size

134 + self.num_attention_heads = self.text_config.num_attention_heads

135 + self.num_key_value_heads = self.text_config.num_key_value_heads

136 + self.num_hidden_layers = self.text_config.num_hidden_layers

137 + self.vocab_size = self.text_config.vocab_size

138 + self.intermediate_size = self.text_config.intermediate_size

139 +

140 + # Moondream3 uses token 0 (<|endoftext|>) as both BOS and EOS.

141 + # Token 3 (<|md_reserved_2|>) is an answer delimiter that the model

142 + # implementation suppresses during generation.

143 + self.bos_token_id = 0

144 + self.eos_token_id = 0

145 +

146 + def get_text_config(self, decoder: bool = False) -> "Moondream3TextConfig":

147 + """Return the text config for vLLM's text_config detection.

148 +

149 + Args:

150 + decoder: Ignored. Only used for encoder-decoder models.

151 + """

152 + return self.text_config

vllm/transformers_utils/model_arch_config_convertor.py

@@ -267,6 +267,7 @@ class ModelArchConfigConvertorBase:

267 "bagel",

268 "gemma3",

269 "molmo2",

270 + "moondream3",

271 "paligemma",

272 "umm",

273 )

vllm/transformers_utils/processors/__init__.py

@@ -23,6 +23,7 @@ __all__ = [

23 "H2OVLProcessor",

24 "HunYuanVLProcessor",

25 "HunYuanVLImageProcessor",

26 + "Moondream3Processor",

27 "InternVLProcessor",

28 "IsaacProcessor",

29 "KimiAudioProcessor",

@@ -61,6 +62,7 @@ _CLASS_TO_MODULE: dict[str, str] = {

62 "MiMoOmniProcessor": "vllm.transformers_utils.processors.mimo_v2_omni",

63 "MistralCommonPixtralProcessor": "vllm.transformers_utils.processors.pixtral",

64 "MistralCommonVoxtralProcessor": "vllm.transformers_utils.processors.voxtral",

65 + "Moondream3Processor": "vllm.transformers_utils.processors.moondream3",

66 "NanoNemotronVLProcessor": "vllm.transformers_utils.processors.nano_nemotron_vl",

67 "NemotronVLProcessor": "vllm.transformers_utils.processors.nemotron_vl",

68 "LlamaNemotronVLEmbedProcessor": "vllm.transformers_utils.processors.nemotron_vl",

vllm/transformers_utils/processors/moondream3.py new

+541

@@ -0,0 +1,541 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Custom processor for Moondream3 model."""

4 +

5 +import math

6 +

7 +import numpy as np

8 +import torch

9 +from PIL import Image

10 +from transformers import AutoProcessor, BatchFeature

11 +from transformers.image_utils import ImageInput

12 +from transformers.processing_utils import ProcessingKwargs, ProcessorMixin, Unpack

13 +from transformers.tokenization_utils_base import (

14 + PreTokenizedInput,

15 + PreTrainedTokenizerBase,

16 + TextInput,

17 +)

18 +

19 +from vllm.multimodal.image import convert_image_mode

20 +

21 +__all__ = ["Moondream3Processor"]

22 +

23 +

24 +class Moondream3ProcessorKwargs(ProcessingKwargs, total=False): # type: ignore[call-arg]

25 + _defaults = {

26 + "text_kwargs": {

27 + "padding": False,

28 + },

29 + "images_kwargs": {

30 + "max_crops": 12,

31 + "overlap_margin": 4,

32 + "crop_size": 378,

33 + "patch_size": 14,

34 + "convert_to_rgb": True,

35 + "return_tensors": "pt",

36 + },

37 + }

38 +

39 +

40 +def select_tiling(

41 + height: int, width: int, crop_size: int, max_crops: int

42 +) -> tuple[int, int]:

43 + """Determine the optimal number of tiles to cover an image."""

44 + if height <= crop_size or width <= crop_size:

45 + return (1, 1)

46 +

47 + min_h = math.ceil(height / crop_size)

48 + min_w = math.ceil(width / crop_size)

49 +

50 + if min_h * min_w > max_crops:

51 + ratio = math.sqrt(max_crops / (min_h * min_w))

52 + return (max(1, math.floor(min_h * ratio)), max(1, math.floor(min_w * ratio)))

53 +

54 + h_tiles = math.floor(math.sqrt(max_crops * height / width))

55 + w_tiles = math.floor(math.sqrt(max_crops * width / height))

56 +

57 + h_tiles = max(h_tiles, min_h)

58 + w_tiles = max(w_tiles, min_w)

59 +

60 + if h_tiles * w_tiles > max_crops:

61 + if w_tiles > h_tiles:

62 + w_tiles = math.floor(max_crops / h_tiles)

63 + else:

64 + h_tiles = math.floor(max_crops / w_tiles)

65 +

66 + return (max(1, h_tiles), max(1, w_tiles))

67 +

68 +

69 +class Moondream3Processor(ProcessorMixin):

70 + """

71 + Constructs a Moondream3 processor which handles image preprocessing

72 + and tokenization for the Moondream3 multimodal model.

73 +

74 + Args:

75 + tokenizer: The tokenizer to use for text processing.

76 + chat_template: Optional chat template string.

77 + crop_size: Size of each image crop.

78 + max_crops: Maximum number of crops per image.

79 + overlap_margin: Margin for overlapping crops in patches.

80 + patch_size: Size of each patch.

81 + """

82 +

83 + attributes = ["tokenizer"]

84 + valid_kwargs = [

85 + "chat_template",

86 + "crop_size",

87 + "max_crops",

88 + "overlap_margin",

89 + "patch_size",

90 + ]

91 +

92 + tokenizer_class = "AutoTokenizer"

93 + # Use separate tokenizer repo

94 + _tokenizer_repo = "moondream/starmie-v1"

95 +

96 + # Default chat template for Moondream3

97 + # Moondream uses special tokens for prompting:

98 + # - Token 0 (<|endoftext|>): BOS token (ALWAYS present at position 0)

99 + # - Token 1 (<|md_reserved_0|>): Start of instruction

100 + # - Token 2 (<|md_reserved_1|>): Separator before question

101 + # - Token 3 (<|md_reserved_2|>): End of question / start of answer

102 + #

103 + # Task routing based on text prefix:

104 + # "caption [short|normal|long]" → describe<|md_reserved_1|>{length}

105 + # "describe [short|normal|long]" → describe<|md_reserved_1|>{length}

106 + # otherwise → query<|md_reserved_1|><text>

107 + #

108 + # Format with image:

109 + # <|endoftext|><image><|md_reserved_0|>{task}<|md_reserved_1|>{q}<|md_reserved_2|>

110 + # Format without image:

111 + # <|endoftext|><|md_reserved_0|>{task}<|md_reserved_1|>{q}<|md_reserved_2|>

112 + _default_chat_template = (

113 + "{% for message in messages %}"

114 + "{% if message['role'] == 'user' %}"

115 + "{% if message['content'] is string %}"

116 + # Simple string content (with image assumed) - route by prefix

117 + "<|endoftext|><image><|md_reserved_0|>"

118 + "{% if message['content'] == 'caption' %}"

119 + "describe<|md_reserved_1|>normal<|md_reserved_2|>"

120 + "{% elif message['content'].startswith('caption ') %}"

121 + "describe<|md_reserved_1|>{{ message['content'][8:] }}<|md_reserved_2|>"

122 + "{% elif message['content'] == 'describe' %}"

123 + "describe<|md_reserved_1|>normal<|md_reserved_2|>"

124 + "{% elif message['content'].startswith('describe ') %}"

125 + "describe<|md_reserved_1|>{{ message['content'][9:] }}<|md_reserved_2|>"

126 + "{% else %}"

127 + "query<|md_reserved_1|>{{ message['content'] }}<|md_reserved_2|>"

128 + "{% endif %}"

129 + "{% else %}"

130 + # List content - build Moondream's image prefix independently of

131 + # OpenAI-style content part order, then render the text task.

132 + "<|endoftext|>"

133 + "{% for content in message['content'] %}"

134 + "{% if content['type'] in ['image', 'image_url', 'input_image', 'image_pil'] %}" # noqa: E501

135 + "<image>"

136 + "{% endif %}"

137 + "{% endfor %}"

138 + "{% for content in message['content'] %}"

139 + "{% if content['type'] == 'text' %}"

140 + "<|md_reserved_0|>"

141 + "{% if content['text'] == 'caption' %}"

142 + "describe<|md_reserved_1|>normal<|md_reserved_2|>"

143 + "{% elif content['text'].startswith('caption ') %}"

144 + "describe<|md_reserved_1|>{{ content['text'][8:] }}<|md_reserved_2|>"

145 + "{% elif content['text'] == 'describe' %}"

146 + "describe<|md_reserved_1|>normal<|md_reserved_2|>"

147 + "{% elif content['text'].startswith('describe ') %}"

148 + "describe<|md_reserved_1|>{{ content['text'][9:] }}<|md_reserved_2|>"

149 + "{% else %}"

150 + "query<|md_reserved_1|>{{ content['text'] }}<|md_reserved_2|>"

151 + "{% endif %}"

152 + "{% endif %}"

153 + "{% endfor %}"

154 + "{% endif %}"

155 + "{% elif message['role'] == 'assistant' %}"

156 + "{{ message['content'] }}"

157 + "{% endif %}"

158 + "{% endfor %}"

159 + )

160 +

161 + def __init__(

162 + self,

163 + tokenizer: PreTrainedTokenizerBase | None = None,

164 + chat_template: str | None = None,

165 + crop_size: int = 378,

166 + max_crops: int = 12,

167 + overlap_margin: int = 4,

168 + patch_size: int = 14,

169 + **kwargs,

170 + ):

171 + self.image_token = "<image>"

172 + self.crop_size = crop_size

173 + self.max_crops = max_crops

174 + self.overlap_margin = overlap_margin

175 + self.patch_size = patch_size

176 +

177 + # Number of patches per crop (27x27 = 729 for 378/14)

178 + self.patches_per_crop = (crop_size // patch_size) ** 2

179 +

180 + # Use default chat template if none provided

181 + if chat_template is None:

182 + chat_template = self._default_chat_template

183 +

184 + super().__init__(tokenizer, chat_template=chat_template)

185 +

186 + @classmethod

187 + def from_pretrained(

188 + cls,

189 + pretrained_model_name_or_path,

190 + **kwargs,

191 + ):

192 + """

193 + Load the processor, using a separate tokenizer repo.

194 +

195 + The moondream3 model uses a custom tokenizer from 'moondream/starmie-v1'

196 + instead of having tokenizer files in the model repo.

197 + """

198 + from transformers import AutoTokenizer, PreTrainedTokenizerFast

199 + from transformers.utils import cached_file

200 +

201 + tokenizer = kwargs.pop("tokenizer", None)

202 +

203 + tokenizer_kwargs = {

204 + "trust_remote_code": kwargs.get("trust_remote_code", False),

205 + }

206 + for key in (

207 + "cache_dir",

208 + "force_download",

209 + "local_files_only",

210 + "revision",

211 + "subfolder",

212 + "token",

213 + "use_fast",

214 + ):

215 + if key in kwargs:

216 + tokenizer_kwargs[key] = kwargs[key]

217 +

218 + cached_file_kwargs = {

219 + key: tokenizer_kwargs[key]

220 + for key in (

221 + "cache_dir",

222 + "force_download",

223 + "local_files_only",

224 + "revision",

225 + "subfolder",

226 + "token",

227 + )

228 + if key in tokenizer_kwargs

229 + }

230 +

231 + def load_tokenizer(repo_or_path):

232 + try:

233 + return AutoTokenizer.from_pretrained(repo_or_path, **tokenizer_kwargs)

234 + except Exception:

235 + tokenizer_file = cached_file(

236 + repo_or_path,

237 + "tokenizer.json",

238 + **cached_file_kwargs,

239 + )

240 + return PreTrainedTokenizerFast(

241 + tokenizer_file=tokenizer_file,

242 + clean_up_tokenization_spaces=False,

243 + )

244 +

245 + if isinstance(tokenizer, str):

246 + tokenizer = load_tokenizer(tokenizer)

247 +

248 + if tokenizer is None:

249 + # Prefer model-local tokenizer files first. If unavailable, fall

250 + # back to moondream's dedicated tokenizer repository.

251 + try:

252 + tokenizer = load_tokenizer(pretrained_model_name_or_path)

253 + except Exception:

254 + tokenizer = load_tokenizer(cls._tokenizer_repo)

255 +

256 + # Configure special tokens for Moondream3

257 + # BOS and EOS are both token 0 (<|endoftext|>), matching the native

258 + # config (TokenizerConfig.bos_id=0, eos_id=0). This is standard for

259 + # GPT-2 style models where <|endoftext|> signals both start and end.

260 + # Token 1 (<|md_reserved_0|>) is a template delimiter, NOT the EOS.

261 + tokenizer.bos_token = "<|endoftext|>"

262 + tokenizer.bos_token_id = 0

263 + tokenizer.eos_token = "<|endoftext|>"

264 + tokenizer.eos_token_id = 0

265 +

266 + # Extract processor-specific kwargs

267 + crop_size = kwargs.pop("crop_size", 378)

268 + max_crops = kwargs.pop("max_crops", 12)

269 + overlap_margin = kwargs.pop("overlap_margin", 4)

270 + patch_size = kwargs.pop("patch_size", 14)

271 + chat_template = kwargs.pop("chat_template", None)

272 +

273 + # Set default chat template on tokenizer if not already set

274 + if chat_template is None:

275 + chat_template = cls._default_chat_template

276 + if tokenizer.chat_template is None:

277 + tokenizer.chat_template = chat_template

278 +

279 + return cls(

280 + tokenizer=tokenizer,

281 + chat_template=chat_template,

282 + crop_size=crop_size,

283 + max_crops=max_crops,

284 + overlap_margin=overlap_margin,

285 + patch_size=patch_size,

286 + )

287 +

288 + def __call__(

289 + self,

290 + images: ImageInput = None,

291 + text: TextInput

292 + | PreTokenizedInput

293 + | list[TextInput]

294 + | list[PreTokenizedInput] = None,

295 + **kwargs: Unpack[Moondream3ProcessorKwargs],

296 + ) -> BatchFeature:

297 + """

298 + Process images and text for Moondream3 model.

299 +

300 + Args:

301 + images: Input images (PIL Image, numpy array, or list thereof).

302 + text: Input text or list of texts.

303 + **kwargs: Additional processing arguments.

304 +

305 + Returns:

306 + BatchFeature with processed inputs.

307 + """

308 + output_kwargs = self._merge_kwargs(

309 + Moondream3ProcessorKwargs,

310 + tokenizer_init_kwargs=self.tokenizer.init_kwargs,

311 + **kwargs,

312 + )

313 +

314 + # Process images

315 + image_features = {}

316 + if images is not None:

317 + processed_images = []

318 + tilings = []

319 +

320 + images_list = images if isinstance(images, list) else [images]

321 + for image in images_list:

322 + pixel_values, tiling = self.preprocess_image(

323 + image, **output_kwargs["images_kwargs"]

324 + )

325 + processed_images.append(pixel_values)

326 + tilings.append(tiling)

327 +

328 + if processed_images:

329 + image_features["pixel_values"] = processed_images

330 + image_features["tilings"] = tilings

331 +

332 + # Process text

333 + if text is not None:

334 + if not isinstance(text, list):

335 + text = [text]

336 +

337 + # Get text kwargs, remove keys we set ourselves

338 + text_kwargs = output_kwargs.get("text_kwargs", {}).copy()

339 + text_kwargs.pop("return_tensors", None)

340 + text_kwargs.pop("add_special_tokens", None)

341 +

342 + # Tokenize text

343 + tokenized = self.tokenizer(

344 + text,

345 + add_special_tokens=True,

346 + return_tensors="pt",

347 + **text_kwargs,

348 + )

349 +

350 + output = BatchFeature(data=dict(tokenized))

351 +

352 + # Add image features

353 + if image_features:

354 + output["pixel_values"] = image_features["pixel_values"]

355 + output["tilings"] = image_features["tilings"]

356 +

357 + return output

358 +

359 + # If only images were provided

360 + return BatchFeature(data=image_features)

361 +

362 + @staticmethod

363 + def _image_array_to_uint8(array: np.ndarray) -> np.ndarray:

364 + if array.dtype == np.uint8:

365 + return np.ascontiguousarray(array)

366 +

367 + if array.dtype == np.bool_:

368 + return np.ascontiguousarray(array.astype(np.uint8) * 255)

369 +

370 + if np.issubdtype(array.dtype, np.floating):

371 + array = np.nan_to_num(array, nan=0.0, posinf=255.0, neginf=0.0)

372 + if array.size > 0 and array.max() <= 1.0:

373 + array = array * 255.0

374 + array = np.rint(array)

375 +

376 + return np.ascontiguousarray(np.clip(array, 0, 255).astype(np.uint8))

377 +

378 + @staticmethod

379 + def _to_pil_image(image: ImageInput) -> Image.Image:

380 + if isinstance(image, Image.Image):

381 + return image

382 +

383 + if isinstance(image, torch.Tensor):

384 + tensor = image.detach().cpu()

385 + if tensor.dtype == torch.bfloat16:

386 + tensor = tensor.to(torch.float32)

387 + image_array = tensor.numpy()

388 + elif isinstance(image, np.ndarray):

389 + image_array = image

390 + else:

391 + raise TypeError(

392 + "Moondream3 images must be PIL images, numpy arrays, "

393 + f"or torch tensors, got {type(image)!r}."

394 + )

395 +

396 + if image_array.ndim == 2:

397 + image_array = Moondream3Processor._image_array_to_uint8(image_array)

398 + return Image.fromarray(image_array)

399 +

400 + if image_array.ndim != 3:

401 + raise ValueError(

402 + "Moondream3 image arrays must have 2 or 3 dimensions, "

403 + f"got shape {image_array.shape}."

404 + )

405 +

406 + channel_dims = (1, 3, 4)

407 + if image_array.shape[-1] not in channel_dims:

408 + if image_array.shape[0] not in channel_dims:

409 + raise ValueError(

410 + "Moondream3 image arrays must be HWC or CHW with 1, 3, "

411 + f"or 4 channels, got shape {image_array.shape}."

412 + )

413 + image_array = np.transpose(image_array, (1, 2, 0))

414 +

415 + image_array = Moondream3Processor._image_array_to_uint8(image_array)

416 + if image_array.shape[-1] == 1:

417 + image_array = image_array[..., 0]

418 +

419 + return Image.fromarray(image_array)

420 +

421 + def preprocess_image(

422 + self,

423 + image: ImageInput,

424 + max_crops: int = 12,

425 + overlap_margin: int = 4,

426 + crop_size: int = 378,

427 + patch_size: int = 14,

428 + convert_to_rgb: bool = True,

429 + return_tensors: str = "pt",

430 + ) -> tuple[torch.Tensor, tuple[int, int]]:

431 + """

432 + Preprocess an image using overlap-and-resize cropping strategy.

433 +

434 + Args:

435 + image: Input PIL image, numpy array, or torch tensor.

436 + max_crops: Maximum number of crops.

437 + overlap_margin: Margin for overlapping in patches.

438 + crop_size: Size of each crop.

439 + patch_size: Size of each patch.

440 + convert_to_rgb: Whether to convert to RGB.

441 + return_tensors: Return type ("pt" for PyTorch).

442 +

443 + Returns:

444 + Tuple of (pixel_values tensor, tiling tuple).

445 + """

446 + image = self._to_pil_image(image)

447 + if convert_to_rgb:

448 + image = convert_image_mode(image, "RGB")

449 +

450 + # Convert to numpy array

451 + image_array = np.array(image)

452 + original_h, original_w = image_array.shape[:2]

453 +

454 + margin_pixels = patch_size * overlap_margin

455 + total_margin_pixels = margin_pixels * 2

456 +

457 + crop_patches = crop_size // patch_size

458 + crop_window_patches = crop_patches - (2 * overlap_margin)

459 + crop_window_size = crop_window_patches * patch_size

460 +

461 + tiling = select_tiling(

462 + original_h - total_margin_pixels,

463 + original_w - total_margin_pixels,

464 + crop_window_size,

465 + max_crops,

466 + )

467 +

468 + n_crops = tiling[0] * tiling[1] + 1

469 + crops = np.zeros((n_crops, crop_size, crop_size, 3), dtype=np.uint8)

470 +

471 + target_size = (

472 + tiling[0] * crop_window_size + total_margin_pixels,

473 + tiling[1] * crop_window_size + total_margin_pixels,

474 + )

475 +

476 + # Resize image

477 + pil_img = Image.fromarray(image_array)

478 + resized = pil_img.resize(

479 + (int(target_size[1]), int(target_size[0])),

480 + resample=Image.Resampling.LANCZOS,

481 + )

482 + resized_array = np.asarray(resized)

483 +

484 + # Create global crop

485 + global_pil = pil_img.resize(

486 + (crop_size, crop_size), resample=Image.Resampling.LANCZOS

487 + )

488 + crops[0] = np.asarray(global_pil)

489 +

490 + # Create local crops

491 + for i in range(tiling[0]):

492 + for j in range(tiling[1]):

493 + y0 = i * crop_window_size

494 + x0 = j * crop_window_size

495 + y_end = min(y0 + crop_size, resized_array.shape[0])

496 + x_end = min(x0 + crop_size, resized_array.shape[1])

497 +

498 + crop_region = resized_array[y0:y_end, x0:x_end]

499 + crop_idx = 1 + i * tiling[1] + j

500 + h_slice = slice(None, crop_region.shape[0])

501 + w_slice = slice(None, crop_region.shape[1])

502 + crops[crop_idx, h_slice, w_slice] = crop_region

503 +

504 + # Convert to tensor: (n_crops, H, W, C) -> (n_crops, C, H, W)

505 + pixel_values = np.transpose(crops, (0, 3, 1, 2))

506 +

507 + if return_tensors == "pt":

508 + # Match HF reference preprocessing exactly: convert uint8 crops to

509 + # bfloat16 before in-place normalization.

510 + pixel_values = (

511 + torch.from_numpy(pixel_values)

512 + .to(dtype=torch.bfloat16)

513 + .div_(255.0)

514 + .sub_(0.5)

515 + .div_(0.5)

516 + )

517 + else:

518 + pixel_values = pixel_values.astype(np.float32) / 255.0

519 + pixel_values = (pixel_values - 0.5) / 0.5

520 +

521 + return pixel_values, tiling

522 +

523 + def get_num_image_tokens(self) -> int:

524 + """Return the number of image tokens (729 = 27x27 patches)."""

525 + return self.patches_per_crop

526 +

527 + def batch_decode(self, *args, **kwargs):

528 + """Forward to tokenizer's batch_decode."""

529 + return self.tokenizer.batch_decode(*args, **kwargs)

530 +

531 + def decode(self, *args, **kwargs):

532 + """Forward to tokenizer's decode."""

533 + return self.tokenizer.decode(*args, **kwargs)

534 +

535 + @property

536 + def model_input_names(self):

537 + tokenizer_input_names = self.tokenizer.model_input_names

538 + return tokenizer_input_names + ["pixel_values", "tilings"]

539 +

540 +

541 +AutoProcessor.register("Moondream3Processor", Moondream3Processor)