[bugfix] Fix Cosmos3 Edge checkpoint weights filtering, video loading, prompt expansion (#49190)

Signed-off-by: Bartosz Stefaniak <[email protected]>

tests/models/multimodal/processing/test_cosmos3_edge.py

@@ -70,7 +70,19 @@ def _assert_video_outputs(processor, processed) -> None:
     merge_size = processor.info.get_hf_config().vision_config.spatial_merge_size
     expected_tokens = int(grid_thw.prod()) // merge_size**2
     video_token_id = processor.info.get_hf_config().video_token_id
-    assert processed["prompt_token_ids"].count(video_token_id) == expected_tokens
+    prompt_token_ids = processed["prompt_token_ids"]
+    assert prompt_token_ids.count(video_token_id) == expected_tokens
+
+    hf_processor = processor.info.get_hf_processor()
+    expected_frame_wrappers = int(grid_thw[:, 0].sum())
+    assert (
+        prompt_token_ids.count(hf_processor.vision_start_token_id)
+        == expected_frame_wrappers
+    )
+    assert (
+        prompt_token_ids.count(hf_processor.vision_end_token_id)
+        == expected_frame_wrappers
+    )
 
 
 @pytest.mark.parametrize("num_images", [1, 2])

tests/models/multimodal/test_mapping.py

@@ -170,6 +170,7 @@ def test_cosmos3_edge_checkpoint_weights_mapper():
                 "layers.0.self_attn.to_add_out.weight",
                 "layers.0.self_attn.norm_added_q.weight",
                 "layers.0.self_attn.norm_added_k.weight",
+                "layers.0.self_attn.k_norm_und_for_gen.weight",
                 "layers.0.self_attn.q_proj_moe_gen.weight",
                 "layers.0.mlp_moe_gen.up_proj.weight",
                 "norm_moe_gen.weight",

tests/multimodal/test_video.py

@@ -304,6 +304,13 @@ def test_pynvvideocodec_decoder_slot_retains_simple_decoder():
 # ============================================================================
 
 
+def test_cosmos3_edge_uses_qwen3_vl_video_backend():
+    backend = get_video_loader_backend_for_processor("Cosmos3EdgeVideoProcessor")
+
+    assert backend == "qwen3_vl"
+    assert isinstance(VIDEO_LOADER_REGISTRY.load(backend), Qwen3VLVideoBackend)
+
+
 @pytest.mark.parametrize(
     "model_repo, expected_loader_cls, hf_sample_kwargs",
     [

vllm/model_executor/models/cosmos3_edge.py

@@ -276,7 +276,13 @@ class Cosmos3EdgeProcessingInfo(Qwen3VLProcessingInfo):
 
 
 class Cosmos3EdgeMultiModalProcessor(Qwen3VLMultiModalProcessor):
-    pass
+    @staticmethod
+    def _expands_only_video_token(_hf_processor: ProcessorMixin) -> bool:
+        # Cosmos renders each video as a full
+        # <|vision_start|><|video_pad|><|vision_end|> placeholder, and its
+        # reference processor replaces that entire triplet with timestamped,
+        # per-frame vision sequences.
+        return False
 
 
 class Cosmos3EdgeDummyInputsBuilder(Qwen3VLDummyInputsBuilder):
@@ -517,6 +523,7 @@ class Cosmos3EdgeForConditionalGeneration(
         },
         orig_to_new_substr={
             "_moe_gen": None,
+            "k_norm_und_for_gen": None,
             ".add_q_proj.": None,
             ".add_k_proj.": None,
             ".add_v_proj.": None,

vllm/multimodal/video.py

@@ -1177,7 +1177,7 @@ class PyNvVideoCodecVideoBackend(VideoBackend):
 
 @VIDEO_LOADER_REGISTRY.register(
     "qwen3_vl",
-    video_processor="Qwen3VLVideoProcessor",
+    video_processor=("Qwen3VLVideoProcessor", "Cosmos3EdgeVideoProcessor"),
 )
 class Qwen3VLVideoBackend(VideoBackend):
     @classmethod