@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 894ebb27f5

Add Cosmos3 Edge Reasoner model (#48291)

Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com> Co-authored-by: Bartosz Stefaniak <bstefaniak@nvidia.com>

adsridhar committed Jul 14, 2026 at 01:14 UTC 894ebb27f5c7d7f8758266163ffd503132e432d2
13 files changed +1474 -27
docs/models/supported_models.md
+1
@@ -514,6 +514,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
514 | `CheersForConditionalGeneration` | Cheers | T + I | `ai9stars/Cheers` | | ✅︎ |
515 | `Cohere2VisionForConditionalGeneration` | Command A Vision, Command-A+ | T + I<sup>+</sup> | `CohereLabs/command-a-vision-07-2025`, `CohereLabs/command-a-plus-05-2026`, etc. | | ✅︎ |
516 | `Cosmos3ForConditionalGeneration` | Cosmos3 (understanding tower) | T + I<sup>E+</sup> + V<sup>E+</sup> | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | | ✅︎ |
517 +| `Cosmos3EdgeForConditionalGeneration` | Cosmos3-Edge (understanding tower) | T + I<sup>E+</sup> + V<sup>E+</sup> | `nvidia/Cosmos3-Edge` | | ✅︎ |
518 | `DeepseekVLV2ForCausalLM` | DeepSeek-VL2 | T + I<sup>+</sup> | `deepseek-ai/deepseek-vl2-tiny`, `deepseek-ai/deepseek-vl2-small`, `deepseek-ai/deepseek-vl2`, etc. | | ✅︎ |
519 | `DeepseekOCRForCausalLM` | DeepSeek-OCR | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR`, etc. | ✅︎ | ✅︎ |
520 | `DeepseekOCR2ForCausalLM` | DeepSeek-OCR-2 | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR-2`, etc. | ✅︎ | ✅︎ |
tests/models/multimodal/processing/test_cosmos3_edge.py new
+101
@@ -0,0 +1,101 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import os
5 +
6 +import pytest
7 +
8 +from vllm.assets.video import VideoAsset
9 +from vllm.config import ModelConfig
10 +from vllm.multimodal import MULTIMODAL_REGISTRY
11 +
12 +from ....conftest import ImageTestAssets
13 +from ...utils import build_model_context
14 +
15 +MODEL_ID = "nvidia/Cosmos3-Edge"
16 +IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>"
17 +VIDEO_PLACEHOLDER = "<|vision_start|><|video_pad|><|vision_end|>"
18 +LOCAL_MODEL_PATH = os.getenv("COSMOS3_EDGE_MODEL_PATH")
19 +
20 +
21 +@pytest.fixture(scope="module")
22 +def processor():
23 + if LOCAL_MODEL_PATH:
24 + model_config = ModelConfig(
25 + LOCAL_MODEL_PATH,
26 + tokenizer=LOCAL_MODEL_PATH,
27 + max_model_len=4096,
28 + limit_mm_per_prompt={"image": 2, "video": 1},
29 + )
30 + else:
31 + ctx = build_model_context(
32 + MODEL_ID,
33 + limit_mm_per_prompt={"image": 2, "video": 1},
34 + )
35 + model_config = ctx.model_config
36 +
37 + return MULTIMODAL_REGISTRY.create_processor(model_config)
38 +
39 +
40 +def _assert_image_outputs(processor, processed, num_images: int) -> None:
41 + mm_data = processed["mm_kwargs"].get_data()
42 + grid_thw = mm_data["image_grid_thw"]
43 + pixel_values = mm_data["pixel_values"]
44 +
45 + assert grid_thw.shape == (num_images, 3)
46 + assert pixel_values.shape[0] == int(grid_thw.prod(dim=-1).sum())
47 +
48 + merge_size = processor.info.get_hf_config().vision_config.spatial_merge_size
49 + expected_tokens = (grid_thw.prod(dim=-1) // merge_size**2).tolist()
50 + image_placeholders = processed["mm_placeholders"]["image"]
51 +
52 + assert len(image_placeholders) == num_images
53 + assert [placeholder.length for placeholder in image_placeholders] == (
54 + expected_tokens
55 + )
56 +
57 + image_token_id = processor.info.get_hf_processor().image_token_id
58 + assert processed["prompt_token_ids"].count(image_token_id) == sum(expected_tokens)
59 +
60 +
61 +def _assert_video_outputs(processor, processed) -> None:
62 + mm_data = processed["mm_kwargs"].get_data()
63 + grid_thw = mm_data["video_grid_thw"]
64 + pixel_values = mm_data["pixel_values_videos"]
65 +
66 + assert grid_thw.shape == (1, 3)
67 + assert pixel_values.shape[0] == int(grid_thw.prod())
68 + assert len(processed["mm_placeholders"]["video"]) == 1
69 +
70 + merge_size = processor.info.get_hf_config().vision_config.spatial_merge_size
71 + expected_tokens = int(grid_thw.prod()) // merge_size**2
72 + video_token_id = processor.info.get_hf_config().video_token_id
73 + assert processed["prompt_token_ids"].count(video_token_id) == expected_tokens
74 +
75 +
76 +@pytest.mark.parametrize("num_images", [1, 2])
77 +def test_process_images(
78 + processor,
79 + image_assets: ImageTestAssets,
80 + num_images: int,
81 +) -> None:
82 + images = [asset.pil_image for asset in image_assets[:num_images]]
83 + processed = processor(
84 + IMAGE_PLACEHOLDER * num_images,
85 + mm_items=processor.info.parse_mm_data({"image": images}),
86 + hf_processor_mm_kwargs={},
87 + )
88 +
89 + _assert_image_outputs(processor, processed, num_images)
90 +
91 +
92 +def test_process_video(processor) -> None:
93 + video_asset = VideoAsset(name="baby_reading", num_frames=8)
94 + video = (video_asset.np_ndarrays, video_asset.metadata)
95 + processed = processor(
96 + VIDEO_PLACEHOLDER,
97 + mm_items=processor.info.parse_mm_data({"video": [video]}),
98 + hf_processor_mm_kwargs={},
99 + )
100 +
101 + _assert_video_outputs(processor, processed)
tests/models/multimodal/test_mapping.py
+71
@@ -74,6 +74,77 @@ def test_cosmos3_new_checkpoint_weights_mapper():
74 )
75
76
77 +def test_cosmos3_edge_checkpoint_weights_mapper():
78 + from vllm.model_executor.models.cosmos3_edge import (
79 + Cosmos3EdgeForConditionalGeneration,
80 + )
81 +
82 + mapper = Cosmos3EdgeForConditionalGeneration.hf_to_vllm_mapper
83 +
84 + assert mapper.apply_list(
85 + [
86 + "embed_tokens.weight",
87 + "norm.weight",
88 + "layers.0.input_layernorm.weight",
89 + "layers.0.self_attn.to_q.weight",
90 + "layers.0.self_attn.to_k.weight",
91 + "layers.0.self_attn.to_v.weight",
92 + "layers.0.self_attn.to_out.weight",
93 + "layers.0.post_attention_layernorm.weight",
94 + "layers.0.mlp.up_proj.weight",
95 + "layers.0.mlp.down_proj.weight",
96 + "layers.27.input_layernorm.weight",
97 + "layers.27.self_attn.to_q.weight",
98 + "layers.27.post_attention_layernorm.weight",
99 + "layers.27.mlp.down_proj.weight",
100 + "model.visual.embeddings.patch_embedding.weight",
101 + "model.visual.encoder.layers.0.self_attn.q_proj.weight",
102 + "model.projector.linear_fc1.weight",
103 + "lm_head.weight",
104 + ]
105 + ) == [
106 + "language_model.model.embed_tokens.weight",
107 + "language_model.model.norm_f.weight",
108 + "language_model.model.layers.0.norm.weight",
109 + "language_model.model.layers.0.mixer.qkv_proj.weight",
110 + "language_model.model.layers.0.mixer.qkv_proj.weight",
111 + "language_model.model.layers.0.mixer.qkv_proj.weight",
112 + "language_model.model.layers.0.mixer.o_proj.weight",
113 + "language_model.model.layers.1.norm.weight",
114 + "language_model.model.layers.1.mixer.up_proj.weight",
115 + "language_model.model.layers.1.mixer.down_proj.weight",
116 + "language_model.model.layers.54.norm.weight",
117 + "language_model.model.layers.54.mixer.qkv_proj.weight",
118 + "language_model.model.layers.55.norm.weight",
119 + "language_model.model.layers.55.mixer.down_proj.weight",
120 + "visual.encoder.embeddings.patch_embedding.weight",
121 + "visual.encoder.encoder.layers.0.self_attn.qkv_proj.weight",
122 + "visual.projector.linear_fc1.weight",
123 + "language_model.lm_head.weight",
124 + ]
125 +
126 + assert (
127 + mapper.apply_list(
128 + [
129 + "layers.0.self_attn.add_q_proj.weight",
130 + "layers.0.self_attn.add_k_proj.weight",
131 + "layers.0.self_attn.add_v_proj.weight",
132 + "layers.0.self_attn.to_add_out.weight",
133 + "layers.0.self_attn.norm_added_q.weight",
134 + "layers.0.self_attn.norm_added_k.weight",
135 + "layers.0.self_attn.q_proj_moe_gen.weight",
136 + "layers.0.mlp_moe_gen.up_proj.weight",
137 + "norm_moe_gen.weight",
138 + "proj_in.weight",
139 + "proj_out.weight",
140 + "audio_modality_embed",
141 + "action_modality_embed",
142 + ]
143 + )
144 + == []
145 + )
146 +
147 +
148 def create_repo_dummy_weights(repo: str) -> Iterable[tuple[str, torch.Tensor]]:
149 """Create weights from safetensors checkpoint metadata"""
150 metadata = try_get_safetensors_metadata(repo)
tests/models/registry.py
+5
@@ -811,6 +811,11 @@ _MULTIMODAL_EXAMPLE_MODELS = {
811 max_model_len=4096,
812 min_transformers_version="4.57",
813 ),
814 + "Cosmos3EdgeForConditionalGeneration": _HfExamplesInfo(
815 + "nvidia/Cosmos3-Edge",
816 + max_model_len=4096,
817 + is_available_online=False,
818 + ),
819 "DeepseekVLV2ForCausalLM": _HfExamplesInfo(
820 "deepseek-ai/deepseek-vl2-tiny",
821 extras={"fork": "Isotr0py/deepseek-vl2-tiny"},
vllm/model_executor/models/cosmos3.py
+4 -21
@@ -3,8 +3,6 @@
3
4 import regex
5
6 -from vllm.config import VllmConfig
7 -from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
6 from vllm.model_executor.models.qwen3_vl import Qwen3VLForConditionalGeneration
7 from vllm.model_executor.models.utils import WeightsMapper
8
@@ -52,22 +50,7 @@ class Cosmos3ForConditionalGeneration(Qwen3VLForConditionalGeneration):
50 },
51 )
52
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 - ]
53 + # Cosmos3 unified Diffusers checkpoints store reasoner weights across
54 + # transformer/ and vision_encoder/. Match both while excluding VAE and
55 + # sound_tokenizer weights; the mapper drops generation-only tensors.
56 + allow_patterns_overrides = ["[tv]*er/*.safetensors"]
vllm/model_executor/models/cosmos3_edge.py new
+751
@@ -0,0 +1,751 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc import Iterable
5 +
6 +import torch
7 +import torch.nn as nn
8 +from transformers import ProcessorMixin
9 +
10 +from vllm.config import VllmConfig
11 +from vllm.model_executor.layers.layernorm import RMSNorm
12 +from vllm.model_executor.layers.linear import (
13 + ColumnParallelLinear,
14 + RowParallelLinear,
15 +)
16 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
17 +from vllm.model_executor.layers.rotary_embedding import get_rope
18 +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
19 +from vllm.model_executor.models.module_mapping import MultiModelKeys
20 +from vllm.multimodal import MULTIMODAL_REGISTRY
21 +from vllm.multimodal.inputs import MultiModalFeatureSpec
22 +from vllm.sequence import IntermediateTensors
23 +from vllm.transformers_utils.configs.nemotron_h import NemotronHConfig
24 +
25 +from .interfaces import (
26 + MultiModalEmbeddings,
27 + SupportsMRoPE,
28 + SupportsMultiModal,
29 + SupportsPP,
30 +)
31 +from .lfm2_siglip2 import Siglip2VisionTransformer
32 +from .nemotron_h import (
33 + NemotronHAttention,
34 + NemotronHMLPDecoderLayer,
35 + NemotronHModel,
36 +)
37 +from .qwen2_5_vl import (
38 + Qwen2_5_VLImageEmbeddingInputs,
39 + Qwen2_5_VLImageInputs,
40 + Qwen2_5_VLImagePixelInputs,
41 + Qwen2_5_VLVideoEmbeddingInputs,
42 + Qwen2_5_VLVideoInputs,
43 + Qwen2_5_VLVideoPixelInputs,
44 +)
45 +from .qwen3_vl import (
46 + Qwen3VLDummyInputsBuilder,
47 + Qwen3VLForConditionalGeneration,
48 + Qwen3VLMultiModalProcessor,
49 + Qwen3VLProcessingInfo,
50 +)
51 +from .utils import (
52 + AutoWeightsLoader,
53 + WeightsMapper,
54 + maybe_prefix,
55 +)
56 +from .vision import (
57 + is_vit_use_data_parallel,
58 + run_dp_sharded_mrope_vision_model,
59 +)
60 +
61 +
62 +class Cosmos3EdgeVisionEncoder(Siglip2VisionTransformer):
63 + """Adapts Cosmos (T, H, W) metadata to vLLM packed SigLIP2."""
64 +
65 + @property
66 + def dtype(self) -> torch.dtype:
67 + return self.embeddings.patch_embedding.weight.dtype
68 +
69 + def forward(
70 + self,
71 + pixel_values: torch.Tensor,
72 + grid_thw: torch.Tensor,
73 + ) -> torch.Tensor:
74 + # SigLIP2 attends to each frame independently, so expand every THW
75 + # entry into T separate HW attention sequences.
76 + grid_thw_cpu = grid_thw.to(device="cpu")
77 + spatial_shapes = torch.repeat_interleave(
78 + grid_thw_cpu[:, 1:],
79 + grid_thw_cpu[:, 0],
80 + dim=0,
81 + )
82 + lengths_cpu = spatial_shapes.prod(dim=-1).to(torch.int32)
83 + lengths = lengths_cpu.to(
84 + device=pixel_values.device,
85 + non_blocking=True,
86 + )
87 +
88 + cu_seqlens = torch.zeros(
89 + lengths.numel() + 1,
90 + dtype=torch.int32,
91 + device=pixel_values.device,
92 + )
93 + cu_seqlens[1:] = lengths.cumsum(dim=0)
94 +
95 + max_seqlen = lengths_cpu.max().reshape(1)
96 +
97 + return super().forward(
98 + pixel_values_packed=pixel_values,
99 + spatial_shapes=spatial_shapes,
100 + cu_seqlens=cu_seqlens,
101 + max_seqlen=max_seqlen,
102 + )[0]
103 +
104 +
105 +def patch_merging_by_param(
106 + image_embeds: torch.Tensor,
107 + grid_thw: torch.Tensor,
108 + merge_size: int = 2,
109 +) -> torch.Tensor:
110 + """Merge each spatial ``merge_size`` block into the channel dimension."""
111 + merged_embeds: list[torch.Tensor] = []
112 + current_idx = 0
113 + hidden_size = image_embeds.shape[-1]
114 +
115 + for grid in grid_thw:
116 + t, h, w = grid.tolist()
117 + if h % merge_size != 0 or w % merge_size != 0:
118 + raise ValueError(
119 + "Grid height and width must be divisible by merge_size: "
120 + f"got grid={(t, h, w)} and merge_size={merge_size}."
121 + )
122 +
123 + num_patches = t * h * w
124 + media_embeds = image_embeds[current_idx : current_idx + num_patches]
125 + if media_embeds.shape[0] != num_patches:
126 + raise ValueError(
127 + "image_embeds contains fewer patches than grid_thw describes."
128 + )
129 + current_idx += num_patches
130 +
131 + media_embeds = media_embeds.view(t, h, w, hidden_size)
132 + media_embeds = media_embeds.view(
133 + t,
134 + h // merge_size,
135 + merge_size,
136 + w // merge_size,
137 + merge_size,
138 + hidden_size,
139 + )
140 + media_embeds = media_embeds.permute(0, 1, 3, 2, 4, 5).contiguous()
141 + merged_embeds.append(
142 + media_embeds.view(-1, merge_size * merge_size * hidden_size)
143 + )
144 +
145 + if current_idx != image_embeds.shape[0]:
146 + raise ValueError(
147 + "image_embeds contains more patches than grid_thw describes: "
148 + f"got {image_embeds.shape[0]}, expected {current_idx}."
149 + )
150 + if not merged_embeds:
151 + return image_embeds.new_empty((0, merge_size * merge_size * hidden_size))
152 + return torch.cat(merged_embeds, dim=0)
153 +
154 +
155 +class Cosmos3EdgePatchMerger(nn.Module):
156 + """
157 + Projector: LayerNorm -> Linear -> GELU -> Linear
158 +
159 + Reads config from projector_config (not vision_config).
160 + input_hidden_size * spatial_merge_size² -> merger_intermediate_size
161 + -> out_hidden_size
162 + """
163 +
164 + def __init__(
165 + self,
166 + input_hidden_size: int,
167 + out_hidden_size: int,
168 + merger_intermediate_size: int,
169 + spatial_merge_size: int = 2,
170 + use_postshuffle_norm: bool = False,
171 + quant_config=None,
172 + prefix: str = "",
173 + ) -> None:
174 + super().__init__()
175 + use_data_parallel = is_vit_use_data_parallel()
176 + self.spatial_merge_size = spatial_merge_size
177 + self.hidden_size = input_hidden_size * (spatial_merge_size**2)
178 + self.input_hidden_size = input_hidden_size
179 + self.out_hidden_size = out_hidden_size
180 + self.use_postshuffle_norm = use_postshuffle_norm
181 +
182 + norm_dim = self.hidden_size if use_postshuffle_norm else input_hidden_size
183 + self.norm = nn.LayerNorm(norm_dim, eps=1e-6)
184 +
185 + self.linear_fc1 = ColumnParallelLinear(
186 + self.hidden_size,
187 + merger_intermediate_size,
188 + bias=True,
189 + quant_config=quant_config,
190 + prefix=f"{prefix}.linear_fc1",
191 + disable_tp=use_data_parallel,
192 + )
193 + self.act_fn = nn.GELU()
194 + self.linear_fc2 = RowParallelLinear(
195 + merger_intermediate_size,
196 + out_hidden_size,
197 + bias=True,
198 + quant_config=quant_config,
199 + prefix=f"{prefix}.linear_fc2",
200 + disable_tp=use_data_parallel,
201 + )
202 +
203 + def forward(self, x: torch.Tensor) -> torch.Tensor:
204 + if self.use_postshuffle_norm:
205 + x = self.norm(x.view(-1, self.hidden_size))
206 + else:
207 + x = self.norm(x).view(-1, self.hidden_size)
208 +
209 + x, _ = self.linear_fc1(x)
210 + x = self.act_fn(x)
211 + x, _ = self.linear_fc2(x)
212 + return x
213 +
214 +
215 +class Cosmos3EdgeVisionModel(nn.Module):
216 + """Complete Cosmos vision tower returning language-model embeddings."""
217 +
218 + def __init__(
219 + self,
220 + vision_config,
221 + projector_config,
222 + quant_config=None,
223 + prefix: str = "",
224 + ) -> None:
225 + super().__init__()
226 + self.spatial_merge_size = projector_config.spatial_merge_size
227 + self.out_hidden_size = projector_config.out_hidden_size
228 +
229 + self.encoder = Cosmos3EdgeVisionEncoder(
230 + vision_config,
231 + quant_config=quant_config,
232 + prefix=maybe_prefix(prefix, "encoder"),
233 + )
234 + self.projector = Cosmos3EdgePatchMerger(
235 + input_hidden_size=projector_config.input_hidden_size,
236 + out_hidden_size=projector_config.out_hidden_size,
237 + merger_intermediate_size=projector_config.merger_intermediate_size,
238 + spatial_merge_size=self.spatial_merge_size,
239 + use_postshuffle_norm=getattr(
240 + projector_config, "use_postshuffle_norm", False
241 + ),
242 + quant_config=quant_config,
243 + prefix=maybe_prefix(prefix, "projector"),
244 + )
245 +
246 + @property
247 + def dtype(self) -> torch.dtype:
248 + return self.encoder.dtype
249 +
250 + def forward(
251 + self,
252 + pixel_values: torch.Tensor,
253 + grid_thw: torch.Tensor | list[list[int]],
254 + ) -> torch.Tensor:
255 + grid_thw = torch.as_tensor(grid_thw, dtype=torch.int64, device="cpu")
256 + image_embeds = self.encoder(pixel_values.type(self.dtype), grid_thw=grid_thw)
257 + image_embeds = patch_merging_by_param(
258 + image_embeds,
259 + grid_thw,
260 + merge_size=self.spatial_merge_size,
261 + )
262 + image_embeds = image_embeds.view(
263 + -1,
264 + self.spatial_merge_size**2,
265 + self.projector.input_hidden_size,
266 + )
267 + return self.projector(image_embeds)
268 +
269 +
270 +class Cosmos3EdgeProcessingInfo(Qwen3VLProcessingInfo):
271 + def get_hf_config(self):
272 + return self.ctx.get_hf_config()
273 +
274 + def get_hf_processor(self, **kwargs: object) -> ProcessorMixin:
275 + return self.ctx.get_hf_processor(**kwargs)
276 +
277 +
278 +class Cosmos3EdgeMultiModalProcessor(Qwen3VLMultiModalProcessor):
279 + pass
280 +
281 +
282 +class Cosmos3EdgeDummyInputsBuilder(Qwen3VLDummyInputsBuilder):
283 + pass
284 +
285 +
286 +class Cosmos3EdgeAttention(NemotronHAttention):
287 + """Nemotron-H attention with interleaved multimodal RoPE."""
288 +
289 + def __init__(
290 + self,
291 + config: NemotronHConfig,
292 + layer_idx: int,
293 + model_config=None,
294 + cache_config=None,
295 + quant_config=None,
296 + prefix: str = "",
297 + ) -> None:
298 + super().__init__(
299 + config,
300 + layer_idx,
301 + model_config,
302 + cache_config,
303 + quant_config,
304 + prefix,
305 + )
306 + self.rotary_emb = get_rope(
307 + self.head_dim,
308 + max_position=config.max_position_embeddings,
309 + rope_parameters=config.rope_parameters,
310 + )
311 +
312 + def forward(
313 + self,
314 + positions: torch.Tensor,
315 + hidden_states: torch.Tensor,
316 + **kwargs,
317 + ) -> torch.Tensor:
318 + qkv, _ = self.qkv_proj(hidden_states)
319 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
320 + q, k = self.rotary_emb(positions, q, k)
321 + attn_output = self.attn(q, k, v)
322 + output, _ = self.o_proj(attn_output)
323 + return output
324 +
325 +
326 +class Cosmos3EdgeAttentionDecoderLayer(nn.Module):
327 + """Pre-norm attention layer for the Cosmos3 Edge dense text model."""
328 +
329 + def __init__(
330 + self,
331 + config: NemotronHConfig,
332 + layer_idx: int,
333 + model_config=None,
334 + cache_config=None,
335 + quant_config=None,
336 + parallel_config=None,
337 + prefix: str = "",
338 + ) -> None:
339 + super().__init__()
340 + get_layer_config = getattr(config, "get_nemotron_h_config_for_layer", None)
341 + layer_config = get_layer_config(layer_idx) if get_layer_config else config
342 + self.mixer = Cosmos3EdgeAttention(
343 + layer_config,
344 + layer_idx,
345 + model_config,
346 + cache_config,
347 + quant_config,
348 + prefix=f"{prefix}.mixer",
349 + )
350 + self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
351 +
352 + def forward(
353 + self,
354 + positions: torch.Tensor,
355 + hidden_states: torch.Tensor,
356 + residual: torch.Tensor | None,
357 + **kwargs,
358 + ) -> tuple[torch.Tensor, torch.Tensor]:
359 + if residual is None:
360 + residual = hidden_states
361 + hidden_states = self.norm(hidden_states)
362 + else:
363 + hidden_states, residual = self.norm(hidden_states, residual)
364 + hidden_states = self.mixer(positions=positions, hidden_states=hidden_states)
365 + return hidden_states, residual
366 +
367 +
368 +class Cosmos3EdgeTextModel(NemotronHModel):
369 + """Nemotron-H backbone with Cosmos-owned checkpoint name mapping."""
370 +
371 + # The enclosing Cosmos model maps and packs checkpoint weights. Disabling
372 + # Nemotron-H's checkpoint-specific hook prevents those names from being
373 + # remapped a second time when AutoWeightsLoader descends into this module.
374 + load_weights = None # type: ignore[assignment]
375 +
376 +
377 +class Cosmos3EdgeForCausalLM(nn.Module):
378 + """Minimal CausalLM wrapper for the Cosmos3 Edge dense text model."""
379 +
380 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
381 + super().__init__()
382 + hf_config = vllm_config.model_config.hf_config
383 + text_config = getattr(hf_config, "text_config", hf_config)
384 + vllm_config = vllm_config.with_hf_config(text_config)
385 +
386 + self.vllm_config = vllm_config
387 + self.model_config = vllm_config.model_config
388 + self.quant_config = vllm_config.quant_config
389 + self.config = text_config
390 + pattern = text_config.hybrid_override_pattern
391 + unsupported = set(pattern) - {"*", "-"}
392 + if unsupported:
393 + raise ValueError(
394 + "Cosmos3 Edge only supports attention (`*`) and MLP (`-`) "
395 + f"layers, but found: {sorted(unsupported)}"
396 + )
397 + if len(pattern) != text_config.num_hidden_layers:
398 + raise ValueError(
399 + "hybrid_override_pattern must contain one entry per layer: "
400 + f"expected {text_config.num_hidden_layers}, "
401 + f"got {len(pattern)}"
402 + )
403 +
404 + self.model = Cosmos3EdgeTextModel(
405 + vllm_config=vllm_config,
406 + prefix=maybe_prefix(prefix, "model"),
407 + decoder_layer_types={
408 + "*": Cosmos3EdgeAttentionDecoderLayer,
409 + "-": NemotronHMLPDecoderLayer,
410 + },
411 + )
412 + self.lm_head = ParallelLMHead(
413 + text_config.vocab_size,
414 + text_config.hidden_size,
415 + quant_config=self.quant_config,
416 + prefix=maybe_prefix(prefix, "lm_head"),
417 + )
418 + self.logits_processor = LogitsProcessor(text_config.vocab_size)
419 + self.make_empty_intermediate_tensors = (
420 + self.model.make_empty_intermediate_tensors
421 + )
422 +
423 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
424 + return self.model.embed_input_ids(input_ids)
425 +
426 + def forward(
427 + self,
428 + input_ids: torch.Tensor | None,
429 + positions: torch.Tensor,
430 + intermediate_tensors: IntermediateTensors | None = None,
431 + inputs_embeds: torch.Tensor | None = None,
432 + **kwargs,
433 + ):
434 + return self.model(
435 + input_ids=input_ids,
436 + positions=positions,
437 + intermediate_tensors=intermediate_tensors,
438 + inputs_embeds=inputs_embeds,
439 + )
440 +
441 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
442 + return self.logits_processor(self.lm_head, hidden_states)
443 +
444 +
445 +def _cosmos3_edge_diffusers_prefix_map() -> dict[str, str]:
446 + """Map Diffusers blocks to interleaved Nemotron-H layers.
447 +
448 + The checkpoint stores attention and MLP in one ``layers.N`` block, while
449 + the Nemotron-H model represents them as separate decoder layers. Therefore,
450 + checkpoint block ``N`` maps its attention to layer ``2N`` and its MLP to
451 + layer ``2N + 1``.
452 + """
453 + mappings = {
454 + "embed_tokens.": "language_model.model.embed_tokens.",
455 + "norm.": "language_model.model.norm_f.",
456 + }
457 + for physical_idx in range(28):
458 + attention_idx = 2 * physical_idx
459 + mlp_idx = attention_idx + 1
460 + source_prefix = f"layers.{physical_idx}"
461 + attention_prefix = f"language_model.model.layers.{attention_idx}"
462 + mlp_prefix = f"language_model.model.layers.{mlp_idx}"
463 + mappings.update(
464 + {
465 + f"{source_prefix}.input_layernorm.": f"{attention_prefix}.norm.",
466 + f"{source_prefix}.self_attn.qkv_proj.": (
467 + f"{attention_prefix}.mixer.qkv_proj."
468 + ),
469 + f"{source_prefix}.self_attn.o_proj.": (
470 + f"{attention_prefix}.mixer.o_proj."
471 + ),
472 + f"{source_prefix}.post_attention_layernorm.": (f"{mlp_prefix}.norm."),
473 + f"{source_prefix}.mlp.": f"{mlp_prefix}.mixer.",
474 + }
475 + )
476 + return mappings
477 +
478 +
479 +@MULTIMODAL_REGISTRY.register_processor(
480 + Cosmos3EdgeMultiModalProcessor,
481 + info=Cosmos3EdgeProcessingInfo,
482 + dummy_inputs=Cosmos3EdgeDummyInputsBuilder,
483 +)
484 +class Cosmos3EdgeForConditionalGeneration(
485 + nn.Module,
486 + SupportsMultiModal,
487 + SupportsPP,
488 + SupportsMRoPE,
489 +):
490 + """
491 + Cosmos3 Edge model with a SigLIP2 vision encoder.
492 +
493 + Architecture:
494 + - self.visual: SigLIP2 encoder + patch merger + projector
495 + - self.language_model: Cosmos3EdgeForCausalLM (pure attention + RoPE)
496 + """
497 +
498 + hf_to_vllm_mapper = WeightsMapper(
499 + orig_to_new_stacked={
500 + ".self_attn.q_proj": (".self_attn.qkv_proj", "q"),
501 + ".self_attn.k_proj": (".self_attn.qkv_proj", "k"),
502 + ".self_attn.v_proj": (".self_attn.qkv_proj", "v"),
503 + },
504 + orig_to_new_prefix={
505 + **_cosmos3_edge_diffusers_prefix_map(),
506 + "proj_in.": None,
507 + "proj_out.": None,
508 + "time_embedder.": None,
509 + "action_proj_in.": None,
510 + "action_proj_out.": None,
511 + "audio_modality_embed": None,
512 + "action_modality_embed": None,
513 + "model.visual.": "visual.encoder.",
514 + "model.projector.": "visual.projector.",
515 + "lm_head.": "language_model.lm_head.",
516 + "model.language_model.": "language_model.model.",
517 + },
518 + orig_to_new_substr={
519 + "_moe_gen": None,
520 + ".add_q_proj.": None,
521 + ".add_k_proj.": None,
522 + ".add_v_proj.": None,
523 + ".to_add_out.": None,
524 + ".norm_added_q.": None,
525 + ".norm_added_k.": None,
526 + ".self_attn.to_q.": ".self_attn.q_proj.",
527 + ".self_attn.to_k.": ".self_attn.k_proj.",
528 + ".self_attn.to_v.": ".self_attn.v_proj.",
529 + ".self_attn.to_out.": ".self_attn.o_proj.",
530 + "language_model.embeddings": "language_model.embed_tokens",
531 + },
532 + )
533 +
534 + packed_modules_mapping = {
535 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],
536 + }
537 +
538 + # Cosmos3 Edge unified Diffusers checkpoints store reasoner weights across
539 + # transformer/ and vision_encoder/. Match both while excluding VAE weights.
540 + allow_patterns_overrides = ["[tv]*er/*.safetensors"]
541 +
542 + supports_encoder_tp_data = True
543 +
544 + @classmethod
545 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:
546 + if modality.startswith("image"):
547 + return "<|vision_start|><|image_pad|><|vision_end|>"
548 + if modality.startswith("video"):
549 + return "<|vision_start|><|video_pad|><|vision_end|>"
550 + raise ValueError(f"Unsupported modality: {modality}")
551 +
552 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
553 + super().__init__()
554 +
555 + config = vllm_config.model_config.hf_config
556 + quant_config = vllm_config.quant_config
557 + multimodal_config = vllm_config.model_config.multimodal_config
558 +
559 + self.config = config
560 + self.multimodal_config = multimodal_config
561 + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
562 +
563 + with self._mark_tower_model(vllm_config, {"image", "video"}):
564 + self.visual = Cosmos3EdgeVisionModel(
565 + vision_config=config.vision_config,
566 + projector_config=config.projector_config,
567 + quant_config=quant_config,
568 + prefix=maybe_prefix(prefix, "visual"),
569 + )
570 +
571 + with self._mark_language_model(vllm_config):
572 + self.language_model = Cosmos3EdgeForCausalLM(
573 + vllm_config=vllm_config,
574 + prefix=maybe_prefix(prefix, "language_model"),
575 + )
576 +
577 + self.make_empty_intermediate_tensors = (
578 + self.language_model.make_empty_intermediate_tensors
579 + )
580 +
581 + def _get_image_features(
582 + self,
583 + pixel_values: torch.Tensor,
584 + grid_thw: torch.Tensor,
585 + ) -> tuple[torch.Tensor, ...]:
586 + """Run the complete vision tower and split by media item."""
587 + image_embeds = self.visual(pixel_values, grid_thw=grid_thw)
588 + sizes = (grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist()
589 + return image_embeds.split(sizes)
590 +
591 + def _parse_and_validate_image_input(
592 + self, **kwargs: object
593 + ) -> Qwen2_5_VLImageInputs | None:
594 + pixel_values = kwargs.pop("pixel_values", None)
595 + image_embeds = kwargs.pop("image_embeds", None)
596 + image_grid_thw = kwargs.pop("image_grid_thw", None)
597 +
598 + if pixel_values is None and image_embeds is None:
599 + return None
600 +
601 + if pixel_values is not None:
602 + return Qwen2_5_VLImagePixelInputs(
603 + type="pixel_values",
604 + pixel_values=pixel_values,
605 + image_grid_thw=image_grid_thw,
606 + )
607 +
608 + if image_embeds is not None:
609 + return Qwen2_5_VLImageEmbeddingInputs(
610 + type="image_embeds",
611 + image_embeds=image_embeds,
612 + image_grid_thw=image_grid_thw,
613 + )
614 +
615 + def _parse_and_validate_video_input(
616 + self, **kwargs: object
617 + ) -> Qwen2_5_VLVideoInputs | None:
618 + pixel_values_videos = kwargs.pop("pixel_values_videos", None)
619 + video_embeds = kwargs.pop("video_embeds", None)
620 + video_grid_thw = kwargs.pop("video_grid_thw", None)
621 + second_per_grid_ts = kwargs.pop("second_per_grid_ts", None)
622 +
623 + if pixel_values_videos is None and video_embeds is None:
624 + return None
625 +
626 + if pixel_values_videos is not None:
627 + return Qwen2_5_VLVideoPixelInputs(
628 + type="pixel_values_videos",
629 + pixel_values_videos=pixel_values_videos,
630 + video_grid_thw=video_grid_thw,
631 + second_per_grid_ts=second_per_grid_ts,
632 + )
633 +
634 + if video_embeds is not None:
635 + return Qwen2_5_VLVideoEmbeddingInputs(
636 + type="video_embeds",
637 + video_embeds=video_embeds,
638 + video_grid_thw=video_grid_thw,
639 + )
640 +
641 + def _process_image_input(
642 + self, image_input: Qwen2_5_VLImageInputs
643 + ) -> tuple[torch.Tensor, ...]:
644 + grid_thw = image_input["image_grid_thw"]
645 + assert grid_thw.ndim == 2
646 +
647 + if image_input["type"] == "image_embeds":
648 + image_embeds = image_input["image_embeds"].type(self.visual.dtype)
649 + merge_size = self.visual.spatial_merge_size
650 + sizes = (grid_thw.prod(-1) // merge_size // merge_size).tolist()
651 + return image_embeds.split(sizes)
652 +
653 + pixel_values = image_input["pixel_values"].type(self.visual.dtype)
654 + if self.use_data_parallel:
655 + return run_dp_sharded_mrope_vision_model(
656 + self.visual, pixel_values, grid_thw.tolist(), rope_type="rope_3d"
657 + )
658 + return self._get_image_features(pixel_values, grid_thw)
659 +
660 + def _process_video_input(
661 + self, video_input: Qwen2_5_VLVideoInputs
662 + ) -> tuple[torch.Tensor, ...]:
663 + grid_thw = video_input["video_grid_thw"]
664 + assert grid_thw.ndim == 2
665 +
666 + if video_input["type"] == "video_embeds":
667 + video_embeds = video_input["video_embeds"].type(self.visual.dtype)
668 + merge_size = self.visual.spatial_merge_size
669 + sizes = (grid_thw.prod(-1) // merge_size // merge_size).tolist()
670 + return video_embeds.split(sizes)
671 +
672 + pixel_values_videos = video_input["pixel_values_videos"].type(self.visual.dtype)
673 + if self.use_data_parallel:
674 + return run_dp_sharded_mrope_vision_model(
675 + self.visual, pixel_values_videos, grid_thw.tolist(), rope_type="rope_3d"
676 + )
677 + return self._get_image_features(pixel_values_videos, grid_thw)
678 +
679 + def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
680 + modalities = {}
681 + image_input = self._parse_and_validate_image_input(**kwargs)
682 + if image_input is not None:
683 + modalities["image"] = image_input
684 + video_input = self._parse_and_validate_video_input(**kwargs)
685 + if video_input is not None:
686 + modalities["video"] = video_input
687 + return modalities
688 +
689 + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None:
690 + mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
691 + if not mm_input_by_modality:
692 + return None
693 +
694 + multimodal_embeddings: tuple[torch.Tensor, ...] = ()
695 +
696 + for modality in mm_input_by_modality:
697 + multimodal_input = mm_input_by_modality[modality]
698 + if modality == "image":
699 + image_embeddings = self._process_image_input(multimodal_input)
700 + multimodal_embeddings += tuple(image_embeddings)
701 + if modality == "video":
702 + video_embeddings = self._process_video_input(multimodal_input)
703 + multimodal_embeddings += tuple(video_embeddings)
704 +
705 + return multimodal_embeddings
706 +
707 + def forward(
708 + self,
709 + input_ids: torch.Tensor | None,
710 + positions: torch.Tensor,
711 + intermediate_tensors: IntermediateTensors | None = None,
712 + inputs_embeds: torch.Tensor | None = None,
713 + **kwargs: object,
714 + ) -> torch.Tensor | IntermediateTensors:
715 + if intermediate_tensors is not None:
716 + inputs_embeds = None
717 +
718 + return self.language_model.model(
719 + input_ids=input_ids,
720 + positions=positions,
721 + intermediate_tensors=intermediate_tensors,
722 + inputs_embeds=inputs_embeds,
723 + )
724 +
725 + def compute_logits(
726 + self,
727 + hidden_states: torch.Tensor,
728 + ) -> torch.Tensor | None:
729 + return self.language_model.compute_logits(hidden_states)
730 +
731 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
732 + loader = AutoWeightsLoader(self)
733 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
734 +
735 + def get_mm_mapping(self) -> MultiModelKeys:
736 + return MultiModelKeys.from_string_field(
737 + language_model="language_model",
738 + connector="visual.projector",
739 + tower_model="visual.",
740 + )
741 +
742 + def get_mrope_input_positions(
743 + self,
744 + input_tokens: list[int],
745 + mm_features: list[MultiModalFeatureSpec],
746 + ) -> tuple[torch.Tensor, int]:
747 + return Qwen3VLForConditionalGeneration._get_mrope_input_positions(
748 + input_tokens=input_tokens,
749 + mm_features=mm_features,
750 + config=self.config,
751 + )
vllm/model_executor/models/nemotron_h.py
+24 -6
@@ -19,7 +19,7 @@
19 """Inference-only NemotronH model."""
20
21 import typing
22 -from collections.abc import Callable, Iterable
22 +from collections.abc import Callable, Iterable, Mapping
23 from itertools import islice
24
25 import torch
@@ -540,9 +540,22 @@ ALL_DECODER_LAYER_TYPES = {
540 }
541
542
543 -@support_torch_compile
543 +@support_torch_compile(
544 + dynamic_arg_dims={
545 + "input_ids": 0,
546 + "positions": -1,
547 + "intermediate_tensors": 0,
548 + "inputs_embeds": 0,
549 + }
550 +)
551 class NemotronHModel(nn.Module, EagleModelMixin):
545 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
552 + def __init__(
553 + self,
554 + *,
555 + vllm_config: VllmConfig,
556 + prefix: str = "",
557 + decoder_layer_types: Mapping[str, type[nn.Module]] | None = None,
558 + ):
559 super().__init__()
560
561 config: NemotronHConfig = vllm_config.model_config.hf_config
@@ -562,11 +575,16 @@ class NemotronHModel(nn.Module, EagleModelMixin):
575
576 self.has_moe = "E" in config.hybrid_override_pattern
577
578 + layer_types = decoder_layer_types or ALL_DECODER_LAYER_TYPES
579 +
580 def get_layer(prefix: str):
581 layer_idx = int(prefix.rsplit(".", 1)[1])
567 - layer_class = ALL_DECODER_LAYER_TYPES[
568 - config.hybrid_override_pattern[layer_idx]
569 - ]
582 + layer_type = config.hybrid_override_pattern[layer_idx]
583 + if layer_type not in layer_types:
584 + raise ValueError(
585 + f"Unsupported layer type {layer_type!r} at layer {layer_idx}"
586 + )
587 + layer_class = layer_types[layer_type]
588 return layer_class(
589 config=config,
590 layer_idx=layer_idx,
vllm/model_executor/models/registry.py
+4
@@ -349,6 +349,10 @@ _MULTIMODAL_MODELS = {
349 "Cohere2VisionForConditionalGeneration",
350 ),
351 "Cosmos3ForConditionalGeneration": ("cosmos3", "Cosmos3ForConditionalGeneration"),
352 + "Cosmos3EdgeForConditionalGeneration": (
353 + "cosmos3_edge",
354 + "Cosmos3EdgeForConditionalGeneration",
355 + ),
356 "DeepseekVLV2ForCausalLM": ("deepseek_vl2", "DeepseekVLV2ForCausalLM"),
357 "DeepseekOCRForCausalLM": ("deepseek_ocr", "DeepseekOCRForCausalLM"),
358 "DeepseekOCR2ForCausalLM": ("deepseek_ocr2", "DeepseekOCR2ForCausalLM"),
vllm/transformers_utils/config.py
+1
@@ -81,6 +81,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
81 ops_colqwen3="OpsColQwen3Config",
82 qwen3_vl_nemotron_embed="Qwen3VLNemotronEmbedConfig",
83 cosmos3_omni="Cosmos3Config",
84 + cosmos3_edge="Cosmos3EdgeConfig",
85 diffusion_gemma="DiffusionGemmaConfig",
86 deepseek_vl_v2="DeepseekVLV2Config",
87 deepseek_v32="DeepseekV3Config",
vllm/transformers_utils/configs/__init__.py
+8
@@ -27,6 +27,10 @@ _CLASS_TO_MODULE: dict[str, str] = {
27 "OpsColQwen3Config": "vllm.transformers_utils.configs.colqwen3",
28 "Qwen3VLNemotronEmbedConfig": "vllm.transformers_utils.configs.colqwen3",
29 "Cosmos3Config": "vllm.transformers_utils.configs.cosmos3",
30 + "Cosmos3EdgeConfig": "vllm.transformers_utils.configs.cosmos3_edge",
31 + "Cosmos3EdgeProjectorConfig": "vllm.transformers_utils.configs.cosmos3_edge",
32 + "Cosmos3EdgeTextConfig": "vllm.transformers_utils.configs.cosmos3_edge",
33 + "Cosmos3EdgeVisionConfig": "vllm.transformers_utils.configs.cosmos3_edge",
34 "DiffusionGemmaConfig": "vllm.transformers_utils.configs.diffusion_gemma",
35 "DiffusionGemmaTextConfig": "vllm.transformers_utils.configs.diffusion_gemma",
36 "DeepseekVLV2Config": "vllm.transformers_utils.configs.deepseek_vl2",
@@ -107,6 +111,10 @@ __all__ = [
111 "OpsColQwen3Config",
112 "Qwen3VLNemotronEmbedConfig",
113 "Cosmos3Config",
114 + "Cosmos3EdgeConfig",
115 + "Cosmos3EdgeProjectorConfig",
116 + "Cosmos3EdgeTextConfig",
117 + "Cosmos3EdgeVisionConfig",
118 "DiffusionGemmaConfig",
119 "DiffusionGemmaTextConfig",
120 "DeepseekVLV2Config",
vllm/transformers_utils/configs/cosmos3_edge.py new
+146
@@ -0,0 +1,146 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from transformers.configuration_utils import PretrainedConfig
5 +from transformers.models.siglip2.configuration_siglip2 import Siglip2VisionConfig
6 +
7 +from .nemotron_h import NemotronHConfig
8 +
9 +
10 +def _normalize_mrope_parameters(config: NemotronHConfig) -> None:
11 + """Normalize the checkpoint's mRoPE configuration for vLLM."""
12 + rope_parameters = dict(getattr(config, "rope_parameters", None) or {})
13 + mrope_section = rope_parameters.get("mrope_section")
14 + if mrope_section is None:
15 + return
16 +
17 + rope_parameters["mrope_section"] = list(mrope_section)
18 + rope_parameters.setdefault("mrope_interleaved", True)
19 + config.rope_parameters = rope_parameters
20 +
21 +
22 +class Cosmos3EdgeTextConfig(NemotronHConfig):
23 + """Dense Cosmos3 Edge text config backed by Nemotron-H layers."""
24 +
25 + model_type = "cosmos3_edge_text"
26 + has_no_defaults_at_init = True
27 + ignore_keys_at_rope_validation = {"mrope_section", "mrope_interleaved"}
28 +
29 + def __init__(
30 + self,
31 + num_hidden_layers: int = 28,
32 + hidden_act: str = "relu2",
33 + rms_norm_eps: float = 1e-5,
34 + **kwargs,
35 + ) -> None:
36 + super().__init__(
37 + num_hidden_layers=2 * num_hidden_layers,
38 + hybrid_override_pattern="*-" * num_hidden_layers,
39 + mlp_hidden_act=hidden_act,
40 + layer_norm_epsilon=rms_norm_eps,
41 + **kwargs,
42 + )
43 +
44 + def to_dict(self) -> dict:
45 + config_dict = super().to_dict()
46 + config_dict["num_hidden_layers"] = self.num_hidden_layers // 2
47 + config_dict["hidden_act"] = config_dict.pop("mlp_hidden_act")
48 + config_dict["rms_norm_eps"] = config_dict.pop("layer_norm_epsilon")
49 + config_dict.pop("hybrid_override_pattern", None)
50 + return config_dict
51 +
52 +
53 +class Cosmos3EdgeProjectorConfig(PretrainedConfig):
54 + model_type = "cosmos3_edge_projector"
55 +
56 + def __init__(
57 + self,
58 + input_hidden_size: int = 1024,
59 + use_postshuffle_norm: bool = False,
60 + spatial_merge_size: int = 2,
61 + merger_intermediate_size: int = 11520,
62 + out_hidden_size: int = 1024,
63 + **kwargs,
64 + ) -> None:
65 + super().__init__(**kwargs)
66 + self.input_hidden_size = input_hidden_size
67 + self.use_postshuffle_norm = use_postshuffle_norm
68 + self.spatial_merge_size = spatial_merge_size
69 + self.merger_intermediate_size = merger_intermediate_size
70 + self.out_hidden_size = out_hidden_size
71 +
72 +
73 +class Cosmos3EdgeVisionConfig(Siglip2VisionConfig):
74 + model_type = "cosmos3_edge_vision"
75 +
76 +
77 +class Cosmos3EdgeConfig(PretrainedConfig):
78 + model_type = "cosmos3_edge"
79 + keys_to_ignore_at_inference = ["past_key_values"]
80 + has_no_defaults_at_init = True
81 + sub_configs = {
82 + "vision_config": Cosmos3EdgeVisionConfig,
83 + "projector_config": Cosmos3EdgeProjectorConfig,
84 + "text_config": Cosmos3EdgeTextConfig,
85 + }
86 +
87 + def __init__(
88 + self,
89 + text_config: Cosmos3EdgeTextConfig | dict | None = None,
90 + vision_config: Cosmos3EdgeVisionConfig | dict | None = None,
91 + projector_config: Cosmos3EdgeProjectorConfig | dict | None = None,
92 + image_token_id: int = 19,
93 + video_token_id: int = 18,
94 + vision_start_token_id: int = 20,
95 + vision_end_token_id: int = 21,
96 + tie_word_embeddings: bool = False,
97 + **kwargs,
98 + ) -> None:
99 + if text_config is None:
100 + text_config = {}
101 + if vision_config is None:
102 + vision_config = {}
103 + if projector_config is None:
104 + projector_config = {}
105 +
106 + self.text_config = (
107 + Cosmos3EdgeTextConfig(**text_config)
108 + if isinstance(text_config, dict)
109 + else text_config
110 + )
111 + _normalize_mrope_parameters(self.text_config)
112 +
113 + self.vision_config = (
114 + Cosmos3EdgeVisionConfig(**vision_config)
115 + if isinstance(vision_config, dict)
116 + else vision_config
117 + )
118 + self.projector_config = (
119 + Cosmos3EdgeProjectorConfig(**projector_config)
120 + if isinstance(projector_config, dict)
121 + else projector_config
122 + )
123 +
124 + # Qwen3-VL processing and M-RoPE read these attributes from the
125 + # vision config, while this checkpoint stores them in projector config.
126 + self.vision_config.spatial_merge_size = self.projector_config.spatial_merge_size
127 + self.vision_config.temporal_patch_size = 1
128 + self.vision_config.out_hidden_size = self.projector_config.out_hidden_size
129 +
130 + self.image_token_id = image_token_id
131 + self.video_token_id = video_token_id
132 + self.vision_start_token_id = vision_start_token_id
133 + self.vision_end_token_id = vision_end_token_id
134 +
135 + super().__init__(
136 + tie_word_embeddings=tie_word_embeddings,
137 + **kwargs,
138 + )
139 +
140 +
141 +__all__ = [
142 + "Cosmos3EdgeConfig",
143 + "Cosmos3EdgeProjectorConfig",
144 + "Cosmos3EdgeTextConfig",
145 + "Cosmos3EdgeVisionConfig",
146 +]
vllm/transformers_utils/processors/__init__.py
+2
@@ -14,6 +14,7 @@ __all__ = [
14 "BagelProcessor",
15 "CheersProcessor",
16 "CohereASRProcessor",
17 + "Cosmos3EdgeProcessor",
18 "DeepseekVLV2Processor",
19 "FireRedASR2Processor",
20 "FireRedLIDProcessor",
@@ -49,6 +50,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
50 "BagelProcessor": "vllm.transformers_utils.processors.bagel",
51 "CheersProcessor": "vllm.transformers_utils.processors.cheers",
52 "CohereASRProcessor": "vllm.transformers_utils.processors.cohere_asr",
53 + "Cosmos3EdgeProcessor": "vllm.transformers_utils.processors.cosmos3_edge",
54 "DeepseekVLV2Processor": "vllm.transformers_utils.processors.deepseek_vl2",
55 "FireRedASR2Processor": "vllm.transformers_utils.processors.fireredasr2",
56 "FireRedLIDProcessor": "vllm.transformers_utils.processors.fireredlid",
vllm/transformers_utils/processors/cosmos3_edge.py new
+356
@@ -0,0 +1,356 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +
5 +import math
6 +
7 +import torch
8 +from torchvision.transforms import InterpolationMode
9 +from transformers import AutoTokenizer
10 +from transformers.feature_extraction_utils import BatchFeature
11 +from transformers.image_utils import ChannelDimension, PILImageResampling, SizeDict
12 +from transformers.models.qwen3_vl.processing_qwen3_vl import Qwen3VLProcessor
13 +from transformers.models.qwen3_vl.video_processing_qwen3_vl import (
14 + Qwen3VLVideoProcessor,
15 + Qwen3VLVideoProcessorInitKwargs,
16 + get_image_size,
17 + smart_resize,
18 +)
19 +from transformers.models.siglip2.image_processing_siglip2 import (
20 + Siglip2ImageProcessor,
21 + convert_image_to_patches,
22 +)
23 +from transformers.processing_utils import Unpack
24 +from transformers.utils import TensorType
25 +from transformers.video_utils import group_videos_by_shape, reorder_videos
26 +
27 +
28 +def round_by_factor(number: int, factor: int) -> int:
29 + """Returns the closest integer to 'number' that is divisible by 'factor'."""
30 + return round(number / factor) * factor
31 +
32 +
33 +def ceil_by_factor(number: int, factor: int) -> int:
34 + """Return the smallest multiple of ``factor`` greater than ``number``."""
35 + return math.ceil(number / factor) * factor
36 +
37 +
38 +def floor_by_factor(number: int, factor: int) -> int:
39 + """Return the largest multiple of ``factor`` less than ``number``."""
40 + return math.floor(number / factor) * factor
41 +
42 +
43 +class Cosmos3EdgeImagesKwargs( # type: ignore[call-arg]
44 + Siglip2ImageProcessor.valid_kwargs,
45 + total=False,
46 +):
47 + # global setting for all images, can be overridden by per-image kwargs
48 + max_pixels: int | None
49 + min_pixels: int | None
50 +
51 + per_image_kwargs: list[dict] | None
52 +
53 +
54 +class Cosmos3EdgeImageProcessor(Siglip2ImageProcessor):
55 + resample = PILImageResampling.BICUBIC
56 + valid_kwargs = Cosmos3EdgeImagesKwargs
57 +
58 + def __init__(self, **kwargs: Unpack[Cosmos3EdgeImagesKwargs]):
59 + super().__init__(**kwargs)
60 +
61 + def _resize_image(
62 + self,
63 + image: torch.Tensor,
64 + max_ratio=200,
65 + resample: PILImageResampling | InterpolationMode | int | None = None,
66 + **kwargs,
67 + ) -> torch.Tensor:
68 + """
69 + Rescales the image so that the following conditions are met:
70 +
71 + 1. Both dimensions (height and width) are divisible by 'factor'.
72 + 2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
73 + 3. The aspect ratio of the image is maintained as closely as possible.
74 + """
75 + image_min_pixels = kwargs.get("min_pixels") or self.size.get(
76 + "shortest_edge", None
77 + )
78 + image_max_pixels = kwargs.get("max_pixels") or self.size.get(
79 + "longest_edge", None
80 + )
81 + assert image_min_pixels is not None and image_max_pixels is not None, (
82 + "When do_resize is True, min_pixels and max_pixels must be provided."
83 + )
84 + assert image_max_pixels >= image_min_pixels, (
85 + "The max_pixels of image must be greater than or equal to min_pixels."
86 + )
87 +
88 + _, height, width = image.shape
89 + aspect_ratio = max(height, width) / min(height, width)
90 + if aspect_ratio > max_ratio:
91 + raise ValueError(
92 + f"absolute aspect ratio must be smaller than {max_ratio}, "
93 + f"got {aspect_ratio}"
94 + )
95 + factor = self.merge_size * self.patch_size
96 + h_bar = max(factor, round_by_factor(height, factor))
97 + w_bar = max(factor, round_by_factor(width, factor))
98 + if h_bar * w_bar > image_max_pixels:
99 + beta = math.sqrt((height * width) / image_max_pixels)
100 + h_bar = floor_by_factor(height / beta, factor)
101 + w_bar = floor_by_factor(width / beta, factor)
102 + elif h_bar * w_bar < image_min_pixels:
103 + beta = math.sqrt(image_min_pixels / (height * width))
104 + h_bar = ceil_by_factor(height * beta, factor)
105 + w_bar = ceil_by_factor(width * beta, factor)
106 +
107 + image = self.resize(
108 + image,
109 + size=SizeDict(height=h_bar, width=w_bar),
110 + resample=resample,
111 + )
112 + return image
113 +
114 + def _preprocess(
115 + self,
116 + images: list["torch.Tensor"],
117 + do_resize: bool,
118 + patch_size: int,
119 + max_num_patches: int,
120 + resample: PILImageResampling | InterpolationMode | int | None,
121 + do_rescale: bool,
122 + rescale_factor: float,
123 + do_normalize: bool,
124 + image_mean: float | list[float] | None,
125 + image_std: float | list[float] | None,
126 + return_tensors: str | TensorType | None,
127 + **kwargs,
128 + ) -> BatchFeature:
129 + per_image_kwargs = kwargs.pop("per_image_kwargs", None)
130 + pixel_values = []
131 + spatial_shapes = []
132 + for idx, image in enumerate(images):
133 + # merge global kwargs with per-image overrides (if provided)
134 + image_kwargs = kwargs.copy()
135 + if (
136 + per_image_kwargs is not None
137 + and idx < len(per_image_kwargs)
138 + and per_image_kwargs[idx]
139 + ):
140 + image_kwargs.update(per_image_kwargs[idx])
141 + if do_resize:
142 + image = self._resize_image(
143 + image,
144 + max_ratio=200,
145 + resample=resample,
146 + **image_kwargs,
147 + )
148 +
149 + image = self.rescale_and_normalize(
150 + image, do_rescale, rescale_factor, do_normalize, image_mean, image_std
151 + )
152 +
153 + # (channels, height, width) -> (patches, patch_size**2 * channels)
154 + patches = convert_image_to_patches(image, patch_size)
155 +
156 + num_patches_height = image.shape[1] // patch_size
157 + num_patches_width = image.shape[2] // patch_size
158 +
159 + spatial_shapes.append((num_patches_height, num_patches_width))
160 + pixel_values.append(patches)
161 +
162 + pixel_values = torch.cat(pixel_values, dim=0)
163 +
164 + spatial_shapes = torch.tensor(spatial_shapes)
165 + t_dim = torch.ones(
166 + (spatial_shapes.shape[0], 1),
167 + dtype=spatial_shapes.dtype,
168 + device=spatial_shapes.device,
169 + )
170 + image_grid_thw = torch.cat([t_dim, spatial_shapes], dim=1)
171 +
172 + batch_feature = BatchFeature(
173 + data={
174 + "pixel_values": pixel_values,
175 + "image_grid_thw": image_grid_thw,
176 + },
177 + tensor_type=return_tensors,
178 + )
179 + return batch_feature
180 +
181 +
182 +class Cosmos3EdgeVideoProcessor(Qwen3VLVideoProcessor):
183 + def __init__(self, **kwargs: Unpack[Qwen3VLVideoProcessorInitKwargs]):
184 + super().__init__(**kwargs)
185 +
186 + def _preprocess(
187 + self,
188 + videos: list[torch.Tensor],
189 + do_convert_rgb: bool = True,
190 + do_resize: bool = True,
191 + size: SizeDict | None = None,
192 + resample: PILImageResampling
193 + | InterpolationMode
194 + | int
195 + | None = PILImageResampling.BICUBIC,
196 + do_rescale: bool = True,
197 + rescale_factor: float = 1 / 255.0,
198 + do_normalize: bool = True,
199 + image_mean: float | list[float] | None = None,
200 + image_std: float | list[float] | None = None,
201 + patch_size: int | None = None,
202 + return_tensors: str | TensorType | None = None,
203 + **kwargs,
204 + ):
205 + assert patch_size is not None
206 +
207 + merge_size = self.merge_size
208 + grouped_videos, grouped_videos_index = group_videos_by_shape(videos)
209 + resized_videos_grouped = {}
210 +
211 + for shape, stacked_videos in grouped_videos.items():
212 + B, T, C, H, W = stacked_videos.shape
213 + num_frames, height, width = T, H, W
214 + if do_resize:
215 + assert size is not None
216 + # 1. Both dimensions (height and width) are divisible by 'factor'.
217 + # 2. Total pixels are within the configured min/max range.
218 + resized_height, resized_width = smart_resize(
219 + num_frames=num_frames,
220 + height=height,
221 + width=width,
222 + temporal_factor=1,
223 + factor=patch_size * merge_size,
224 + min_pixels=size.shortest_edge,
225 + max_pixels=size.longest_edge,
226 + )
227 + stacked_videos = stacked_videos.view(B * T, C, H, W)
228 + stacked_videos = self.resize(
229 + stacked_videos,
230 + size=SizeDict(height=resized_height, width=resized_width),
231 + resample=resample,
232 + )
233 + stacked_videos = stacked_videos.view(
234 + B, T, C, resized_height, resized_width
235 + )
236 + resized_videos_grouped[shape] = stacked_videos
237 + resized_videos = reorder_videos(resized_videos_grouped, grouped_videos_index)
238 +
239 + # Group videos by size for further processing
240 + # Resize may return videos with different sizes.
241 + grouped_videos, grouped_videos_index = group_videos_by_shape(resized_videos)
242 + processed_videos_grouped = {}
243 + processed_grids = {}
244 + for shape, stacked_videos in grouped_videos.items():
245 + resized_height, resized_width = get_image_size(
246 + stacked_videos[0], channel_dim=ChannelDimension.FIRST
247 + )
248 + # Fused rescale and normalize
249 + stacked_videos = self.rescale_and_normalize(
250 + stacked_videos,
251 + do_rescale,
252 + rescale_factor,
253 + do_normalize,
254 + image_mean,
255 + image_std,
256 + )
257 + patches = stacked_videos
258 +
259 + batch_size, grid_t, channel = patches.shape[:3]
260 + grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
261 +
262 + # [1, 16, 3, 320, 608] -> [1, 16, 3, 20, 16, 38, 16]
263 + patches = patches.view(
264 + batch_size, # 0
265 + grid_t, # 1
266 + channel, # 2
267 + grid_h, # 3
268 + patch_size, # 4
269 + grid_w, # 5
270 + patch_size, # 6
271 + )
272 + # [1, 16, 3, 20, 16, 38, 16] -> [1, 16, 20, 38, 16, 16, 3]
273 + # Result: [batch, grid_t, grid_h, grid_w, patch_h, patch_w, channel]
274 + patches = patches.permute(0, 1, 3, 5, 4, 6, 2)
275 +
276 + # [1, 16, 20, 38, 16, 16, 3] -> [1, 5120, 768]
277 + flatten_patches = patches.reshape(
278 + batch_size,
279 + grid_t * grid_h * grid_w,
280 + patch_size * patch_size * channel,
281 + )
282 +
283 + processed_videos_grouped[shape] = flatten_patches
284 + processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size
285 +
286 + processed_videos = reorder_videos(
287 + processed_videos_grouped, grouped_videos_index
288 + )
289 + processed_grids = reorder_videos(processed_grids, grouped_videos_index)
290 + pixel_values_videos = torch.cat(processed_videos, dim=0)
291 + video_grid_thw = torch.tensor(processed_grids)
292 + data = {
293 + "pixel_values_videos": pixel_values_videos,
294 + "video_grid_thw": video_grid_thw,
295 + }
296 +
297 + return BatchFeature(data=data, tensor_type=return_tensors)
298 +
299 +
300 +class Cosmos3EdgeProcessor(Qwen3VLProcessor):
301 + """Construct a processor for Cosmos3 Edge multimodal inputs.
302 +
303 + This processor combines :class:`Cosmos3EdgeImageProcessor`,
304 + :class:`Cosmos3EdgeVideoProcessor`, and a tokenizer into a single
305 + interface. It inherits the text, image, and video input handling and the
306 + decoding methods from :class:`Qwen3VLProcessor`, while using the Cosmos3
307 + Edge-specific image and video preprocessing implementations.
308 +
309 + Args:
310 + image_processor: A :class:`Cosmos3EdgeImageProcessor` used to resize,
311 + normalize, and patch image inputs.
312 + tokenizer: The tokenizer used to encode and decode text.
313 + video_processor: A :class:`Cosmos3EdgeVideoProcessor` used to resize,
314 + normalize, and patch video inputs.
315 + chat_template: An optional Jinja template used to render chat messages
316 + into a tokenizable string.
317 + """
318 +
319 + attributes = ["image_processor", "tokenizer", "video_processor"]
320 +
321 + # Transformers' auto loaders cannot resolve these vLLM-native image and
322 + # video processors without remote code or global auto-class registration.
323 + @classmethod
324 + def _get_arguments_from_pretrained(
325 + cls,
326 + pretrained_model_name_or_path,
327 + processor_dict=None,
328 + **kwargs,
329 + ):
330 + component_kwargs = dict(kwargs)
331 + component_kwargs.pop("trust_remote_code", None)
332 +
333 + image_processor = Cosmos3EdgeImageProcessor.from_pretrained(
334 + pretrained_model_name_or_path,
335 + **component_kwargs,
336 + )
337 + tokenizer = AutoTokenizer.from_pretrained(
338 + pretrained_model_name_or_path,
339 + trust_remote_code=False,
340 + **component_kwargs,
341 + )
342 + video_processor = Cosmos3EdgeVideoProcessor.from_pretrained(
343 + pretrained_model_name_or_path,
344 + **component_kwargs,
345 + )
346 +
347 + return [
348 + image_processor,
349 + tokenizer,
350 + video_processor,
351 + ]
352 +
353 +
354 +__all__ = [
355 + "Cosmos3EdgeProcessor",
356 +]