[bugfix] Fix Cosmos3 Edge checkpoint weights filtering, video loading, prompt expansion (#49190)
Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>
bastefaniak committed
Jul 21, 2026 at 11:18 UTC
4d30c510ce36e692f385f92fa69a05939ac44a52
5 files changed
+30
-3
tests/models/multimodal/processing/test_cosmos3_edge.py
+13
-1
@@ -70,7 +70,19 @@ def _assert_video_outputs(processor, processed) -> None:
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
73
+ prompt_token_ids = processed["prompt_token_ids"]
74
+ assert prompt_token_ids.count(video_token_id) == expected_tokens
75
+
76
+ hf_processor = processor.info.get_hf_processor()
77
+ expected_frame_wrappers = int(grid_thw[:, 0].sum())
78
+ assert (
79
+ prompt_token_ids.count(hf_processor.vision_start_token_id)
80
+ == expected_frame_wrappers
81
+ )
82
+ assert (
83
+ prompt_token_ids.count(hf_processor.vision_end_token_id)
84
+ == expected_frame_wrappers
85
+ )
86
87
88
@pytest.mark.parametrize("num_images", [1, 2])
tests/models/multimodal/test_mapping.py
+1
@@ -170,6 +170,7 @@ def test_cosmos3_edge_checkpoint_weights_mapper():
170
"layers.0.self_attn.to_add_out.weight",
171
"layers.0.self_attn.norm_added_q.weight",
172
"layers.0.self_attn.norm_added_k.weight",
173
+ "layers.0.self_attn.k_norm_und_for_gen.weight",
174
"layers.0.self_attn.q_proj_moe_gen.weight",
175
"layers.0.mlp_moe_gen.up_proj.weight",
176
"norm_moe_gen.weight",
tests/multimodal/test_video.py
+7
@@ -304,6 +304,13 @@ def test_pynvvideocodec_decoder_slot_retains_simple_decoder():
304
# ============================================================================
305
306
307
+def test_cosmos3_edge_uses_qwen3_vl_video_backend():
308
+ backend = get_video_loader_backend_for_processor("Cosmos3EdgeVideoProcessor")
309
+
310
+ assert backend == "qwen3_vl"
311
+ assert isinstance(VIDEO_LOADER_REGISTRY.load(backend), Qwen3VLVideoBackend)
312
+
313
+
314
@pytest.mark.parametrize(
315
"model_repo, expected_loader_cls, hf_sample_kwargs",
316
[
vllm/model_executor/models/cosmos3_edge.py
+8
-1
@@ -276,7 +276,13 @@ class Cosmos3EdgeProcessingInfo(Qwen3VLProcessingInfo):
276
277
278
class Cosmos3EdgeMultiModalProcessor(Qwen3VLMultiModalProcessor):
279
- pass
279
+ @staticmethod
280
+ def _expands_only_video_token(_hf_processor: ProcessorMixin) -> bool:
281
+ # Cosmos renders each video as a full
282
+ # <|vision_start|><|video_pad|><|vision_end|> placeholder, and its
283
+ # reference processor replaces that entire triplet with timestamped,
284
+ # per-frame vision sequences.
285
+ return False
286
287
288
class Cosmos3EdgeDummyInputsBuilder(Qwen3VLDummyInputsBuilder):
@@ -517,6 +523,7 @@ class Cosmos3EdgeForConditionalGeneration(
523
},
524
orig_to_new_substr={
525
"_moe_gen": None,
526
+ "k_norm_und_for_gen": None,
527
".add_q_proj.": None,
528
".add_k_proj.": None,
529
".add_v_proj.": None,
vllm/multimodal/video.py
+1
-1
@@ -1177,7 +1177,7 @@ class PyNvVideoCodecVideoBackend(VideoBackend):
1177
1178
@VIDEO_LOADER_REGISTRY.register(
1179
"qwen3_vl",
1180
- video_processor="Qwen3VLVideoProcessor",
1180
+ video_processor=("Qwen3VLVideoProcessor", "Cosmos3EdgeVideoProcessor"),
1181
)
1182
class Qwen3VLVideoBackend(VideoBackend):
1183
@classmethod