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 2, 13:30 UTC
Add Cosmos3 Reasoner model (#43356)
Signed-off-by: Maciej Bala <mbala@nvidia.com> Signed-off-by: MaciejBalaNV <mbala@nvidia.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Isotr0py <2037008807@qq.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Roger Wang <hey@rogerw.io>
MaciejBalaNV committed
May 28, 2026 at 18:43 UTC
9aa131f9444649b3036d811d5bdb26c75f8d4157
10 files changed
+170
-1
docs/models/supported_models.md
+1
@@ -551,6 +551,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
551
| `ChameleonForConditionalGeneration` | Chameleon | T + I | `facebook/chameleon-7b`, etc. | | ✅︎ |
552
| `CheersForConditionalGeneration` | Cheers | T + I | `ai9stars/Cheers` | | ✅︎ |
553
| `Cohere2VisionForConditionalGeneration` | Command A Vision | T + I<sup>+</sup> | `CohereLabs/command-a-vision-07-2025`, etc. | | ✅︎ |
554
+| `Cosmos3ForConditionalGeneration` | Cosmos3 (understanding tower) | T + I<sup>E+</sup> + V<sup>E+</sup> | `nvidia/Cosmos3-Nano` | | ✅︎ |
555
| `DeepseekVLV2ForCausalLM` | DeepSeek-VL2 | T + I<sup>+</sup> | `deepseek-ai/deepseek-vl2-tiny`, `deepseek-ai/deepseek-vl2-small`, `deepseek-ai/deepseek-vl2`, etc. | | ✅︎ |
556
| `DeepseekOCRForCausalLM` | DeepSeek-OCR | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR`, etc. | ✅︎ | ✅︎ |
557
| `DeepseekOCR2ForCausalLM` | DeepSeek-OCR-2 | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR-2`, etc. | ✅︎ | ✅︎ |
tests/models/multimodal/generation/test_common.py
+20
@@ -383,6 +383,26 @@ VLM_TEST_SETTINGS = {
383
max_tokens=8,
384
dtype="bfloat16",
385
),
386
+ "cosmos3": VLMTestInfo(
387
+ models=["nvidia/Cosmos3-Nano"],
388
+ test_type=(
389
+ VLMTestType.IMAGE,
390
+ VLMTestType.MULTI_IMAGE,
391
+ VLMTestType.VIDEO,
392
+ ),
393
+ enforce_eager=False,
394
+ needs_video_metadata=True,
395
+ prompt_formatter=lambda img_prompt: f"<|im_start|>User\n{img_prompt}<|im_end|>\n<|im_start|>assistant\n", # noqa: E501
396
+ img_idx_to_prompt=lambda idx: "<|vision_start|><|image_pad|><|vision_end|>", # noqa: E501
397
+ video_idx_to_prompt=lambda idx: "<|vision_start|><|video_pad|><|vision_end|>", # noqa: E501
398
+ max_model_len=4096,
399
+ max_num_seqs=2,
400
+ num_logprobs=20,
401
+ auto_cls=AutoModelForImageTextToText,
402
+ vllm_output_post_proc=model_utils.qwen2_vllm_to_hf_output,
403
+ patch_hf_runner=model_utils.qwen3_vl_patch_hf_runner,
404
+ image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
405
+ ),
406
"deepseek_vl_v2": VLMTestInfo(
407
models=["Isotr0py/deepseek-vl2-tiny"], # model repo using dynamic module
408
test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),
tests/models/multimodal/test_mapping.py
+58
@@ -16,6 +16,64 @@ from vllm.transformers_utils.config import try_get_safetensors_metadata
16
from ..registry import _MULTIMODAL_EXAMPLE_MODELS, HF_EXAMPLE_MODELS
17
18
19
+def test_cosmos3_new_checkpoint_weights_mapper():
20
+ from vllm.model_executor.models.cosmos3 import Cosmos3ForConditionalGeneration
21
+
22
+ mapper = Cosmos3ForConditionalGeneration.hf_to_vllm_mapper
23
+
24
+ assert mapper.apply_list(
25
+ [
26
+ "layers.0.self_attn.to_q.weight",
27
+ "layers.0.self_attn.to_k.weight",
28
+ "layers.0.self_attn.to_v.weight",
29
+ "layers.0.self_attn.to_out.weight",
30
+ "layers.0.self_attn.norm_q.weight",
31
+ "layers.0.self_attn.norm_k.weight",
32
+ "embed_tokens.weight",
33
+ "norm.weight",
34
+ "lm_head.weight",
35
+ "blocks.0.attn.qkv.weight",
36
+ ]
37
+ ) == [
38
+ "language_model.model.layers.0.self_attn.q_proj.weight",
39
+ "language_model.model.layers.0.self_attn.k_proj.weight",
40
+ "language_model.model.layers.0.self_attn.v_proj.weight",
41
+ "language_model.model.layers.0.self_attn.o_proj.weight",
42
+ "language_model.model.layers.0.self_attn.q_norm.weight",
43
+ "language_model.model.layers.0.self_attn.k_norm.weight",
44
+ "language_model.model.embed_tokens.weight",
45
+ "language_model.model.norm.weight",
46
+ "language_model.lm_head.weight",
47
+ "visual.blocks.0.attn.qkv.weight",
48
+ ]
49
+
50
+ assert (
51
+ mapper.apply_list(
52
+ [
53
+ "layers.0.self_attn.add_q_proj.weight",
54
+ "layers.0.self_attn.add_k_proj.weight",
55
+ "layers.0.self_attn.add_v_proj.weight",
56
+ "layers.0.self_attn.to_add_out.weight",
57
+ "layers.0.self_attn.norm_added_q.weight",
58
+ "layers.0.self_attn.norm_added_k.weight",
59
+ "layers.0.self_attn.q_proj_moe_gen.weight",
60
+ "layers.0.mlp_moe_gen.gate_up_proj.weight",
61
+ "norm_moe_gen.weight",
62
+ "proj_in.weight",
63
+ "proj_out.weight",
64
+ "time_embedder.linear_1.weight",
65
+ "audio_proj_in.weight",
66
+ "audio_proj_out.weight",
67
+ "action_proj_in.weight",
68
+ "action_proj_out.weight",
69
+ "audio_modality_embed",
70
+ "action_modality_embed",
71
+ ]
72
+ )
73
+ == []
74
+ )
75
+
76
+
77
def create_repo_dummy_weights(repo: str) -> Iterable[tuple[str, torch.Tensor]]:
78
"""Create weights from safetensors checkpoint metadata"""
79
metadata = try_get_safetensors_metadata(repo)
tests/models/registry.py
+6
@@ -846,6 +846,12 @@ _MULTIMODAL_EXAMPLE_MODELS = {
846
"Cohere2VisionForConditionalGeneration": _HfExamplesInfo(
847
"CohereLabs/command-a-vision-07-2025"
848
),
849
+ "Cosmos3ForConditionalGeneration": _HfExamplesInfo(
850
+ "nvidia/Cosmos3-Nano",
851
+ max_model_len=4096,
852
+ min_transformers_version="4.57",
853
+ is_available_online=False,
854
+ ),
855
"DeepseekVLV2ForCausalLM": _HfExamplesInfo(
856
"deepseek-ai/deepseek-vl2-tiny",
857
extras={"fork": "Isotr0py/deepseek-vl2-tiny"},
vllm/model_executor/model_loader/default_loader.py
+1
-1
@@ -177,7 +177,7 @@ class DefaultModelLoader(BaseModelLoader):
177
for pattern in allow_patterns:
178
hf_weights_files += glob.glob(os.path.join(hf_folder, pattern))
179
if len(hf_weights_files) > 0:
180
- if pattern == "*.safetensors":
180
+ if pattern.endswith(".safetensors"):
181
use_safetensors = True
182
break
183
vllm/model_executor/models/cosmos3.py
new
+73
@@ -0,0 +1,73 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import regex
5
+
6
+from vllm.config import VllmConfig
7
+from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
8
+from vllm.model_executor.models.qwen3_vl import Qwen3VLForConditionalGeneration
9
+from vllm.model_executor.models.utils import WeightsMapper
10
+
11
+
12
+class Cosmos3ForConditionalGeneration(Qwen3VLForConditionalGeneration):
13
+ # Cosmos3 unified checkpoints store a Qwen3-VL understanding tower
14
+ # alongside a generation tower in a flat key layout. This mapper drops
15
+ # the generation tower weights and rewrites the understanding tower keys
16
+ # into the nested form expected by Qwen3VLForConditionalGeneration.
17
+ hf_to_vllm_mapper = WeightsMapper(
18
+ orig_to_new_regex={
19
+ regex.compile(
20
+ r"^(layers\.|embed_tokens\.|norm\.)(.+)$"
21
+ ): r"language_model.model.\1\2",
22
+ regex.compile(
23
+ r"^(blocks\.|merger\.|patch_embed\.|pos_embed\.|deepstack_merger_list\.)"
24
+ ): r"visual.\1",
25
+ regex.compile(r"^audio_modality_embed(?:\..*)?$"): None,
26
+ regex.compile(r"^action_modality_embed(?:\..*)?$"): None,
27
+ },
28
+ orig_to_new_substr={
29
+ "_moe_gen": None,
30
+ ".add_q_proj.": None,
31
+ ".add_k_proj.": None,
32
+ ".add_v_proj.": None,
33
+ ".to_add_out.": None,
34
+ ".norm_added_q.": None,
35
+ ".norm_added_k.": None,
36
+ ".to_q.": ".q_proj.",
37
+ ".to_k.": ".k_proj.",
38
+ ".to_v.": ".v_proj.",
39
+ ".to_out.": ".o_proj.",
40
+ ".norm_q.": ".q_norm.",
41
+ ".norm_k.": ".k_norm.",
42
+ },
43
+ orig_to_new_prefix={
44
+ "proj_in.": None,
45
+ "proj_out.": None,
46
+ "time_embedder.": None,
47
+ "audio_proj_in.": None,
48
+ "audio_proj_out.": None,
49
+ "action_proj_in.": None,
50
+ "action_proj_out.": None,
51
+ "lm_head.": "language_model.lm_head.",
52
+ },
53
+ )
54
+
55
+ allow_patterns_overrides = ["transformer/*.safetensors"]
56
+
57
+ """
58
+ Cosmos3 checkpoint separates transformer weights and vision_encoder weights
59
+ into separate directories, as it's in diffusers checkpoint format.
60
+ Using secondary_weights here to load all necessary weights for
61
+ the Reasoner-only part.
62
+ """
63
+
64
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
65
+ super().__init__(vllm_config=vllm_config, prefix=prefix)
66
+ self.secondary_weights = [
67
+ DefaultModelLoader.Source(
68
+ model_or_path=vllm_config.model_config.model,
69
+ revision=vllm_config.model_config.revision,
70
+ prefix="",
71
+ allow_patterns_overrides=["vision_encoder/*.safetensors"],
72
+ ),
73
+ ]
vllm/model_executor/models/registry.py
+1
@@ -369,6 +369,7 @@ _MULTIMODAL_MODELS = {
369
"cohere2_vision",
370
"Cohere2VisionForConditionalGeneration",
371
),
372
+ "Cosmos3ForConditionalGeneration": ("cosmos3", "Cosmos3ForConditionalGeneration"),
373
"DeepseekVLV2ForCausalLM": ("deepseek_vl2", "DeepseekVLV2ForCausalLM"),
374
"DeepseekOCRForCausalLM": ("deepseek_ocr", "DeepseekOCRForCausalLM"),
375
"DeepseekOCR2ForCausalLM": ("deepseek_ocr2", "DeepseekOCR2ForCausalLM"),
vllm/transformers_utils/config.py
+1
@@ -94,6 +94,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
94
colqwen3="ColQwen3Config",
95
ops_colqwen3="OpsColQwen3Config",
96
qwen3_vl_nemotron_embed="Qwen3VLNemotronEmbedConfig",
97
+ cosmos3_omni="Cosmos3Config",
98
deepseek_vl_v2="DeepseekVLV2Config",
99
deepseek_v32="DeepseekV3Config",
100
deepseek_v4="DeepseekV4Config",
vllm/transformers_utils/configs/__init__.py
+2
@@ -25,6 +25,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
25
"ColQwen3Config": "vllm.transformers_utils.configs.colqwen3",
26
"OpsColQwen3Config": "vllm.transformers_utils.configs.colqwen3",
27
"Qwen3VLNemotronEmbedConfig": "vllm.transformers_utils.configs.colqwen3",
28
+ "Cosmos3Config": "vllm.transformers_utils.configs.cosmos3",
29
"DeepseekVLV2Config": "vllm.transformers_utils.configs.deepseek_vl2",
30
"DeepseekV4Config": "vllm.transformers_utils.configs.deepseek_v4",
31
"DotsOCRConfig": "vllm.transformers_utils.configs.dotsocr",
@@ -95,6 +96,7 @@ __all__ = [
96
"ColQwen3Config",
97
"OpsColQwen3Config",
98
"Qwen3VLNemotronEmbedConfig",
99
+ "Cosmos3Config",
100
"DeepseekVLV2Config",
101
"DeepseekV3Config",
102
"DeepseekV4Config",
vllm/transformers_utils/configs/cosmos3.py
new
+7
@@ -0,0 +1,7 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig
4
+
5
+
6
+class Cosmos3Config(Qwen3VLConfig):
7
+ model_type = "cosmos3_omni"