[model] add cosmos3 reasoner to llm only inference (#33572)

Signed-off-by: joeltg <joel@reflection.ai> Signed-off-by: Joe Rowell <joe@poolside.ai> Co-authored-by: Dawid Majchrowski <dmajchrowski@nvidia.com> Co-authored-by: Kedi Wu <kediw@nvidia.com> Co-authored-by: Kedi Wu <31940276+kediwu0331@users.noreply.github.com> Co-authored-by: Joel Gustafson <joelgustafson@protonmail.com>
2026-09-04 22:11:58 +08:00
parent 19b46863f3
commit 4349538c02
22 changed files with 2698 additions and 37 deletions
@@ -35,7 +35,10 @@ from sglang.srt.model_loader.loader import (
    ModelOptModelLoader,
    get_model_loader,
 )
-from sglang.srt.model_loader.weight_utils import get_quant_config
+from sglang.srt.model_loader.weight_utils import (
+    _modelopt_quant_section,
+    get_quant_config,
+)
 from sglang.srt.models.minimax_m3 import MiniMaxM3SparseForCausalLM
 from sglang.srt.models.utils import WeightsMapper
 from sglang.srt.utils import get_device
@@ -925,6 +928,50 @@ class TestModelOptMixedPrecisionConfig(CustomTestCase):

        self.assertEqual(result["quant_method"], "modelopt_mixed")

+    def test_flat_hf_quant_config_without_quantization_key(self):
+        """Diffusion/unified ModelOpt exports use a flat hf_quant_config.json.
+
+        Regression for Cosmos3-style checkpoints that put quant_algo at the top
+        level (no nested ``quantization`` key).
+        """
+        model_config = ModelConfig.__new__(ModelConfig)
+
+        result = model_config._parse_modelopt_quant_config(
+            {
+                "quant_method": "modelopt",
+                "quant_algo": "FP8",
+                "quant_type": "FP8_FP8",
+                "ignore": ["lm_head", "visual*"],
+            }
+        )
+
+        self.assertEqual(result["quant_method"], "modelopt_fp8")
+        self.assertEqual(result["quant_algo"], "FP8")
+
+    def test_hf_quant_config_missing_quant_algo_returns_none(self):
+        model_config = ModelConfig.__new__(ModelConfig)
+        self.assertIsNone(
+            model_config._parse_modelopt_quant_config(
+                {"quant_method": "modelopt", "producer": {"name": "modelopt"}}
+            )
+        )
+
+    def test_modelopt_quant_section_supports_nested_and_flat(self):
+        nested = {"quantization": {"quant_algo": "FP8", "exclude_modules": ["lm_head"]}}
+        self.assertEqual(
+            _modelopt_quant_section(nested)["quant_algo"],
+            "FP8",
+        )
+
+        flat = {
+            "quant_method": "modelopt",
+            "quant_algo": "FP8",
+            "ignore": ["lm_head"],
+            "producer": {"name": "modelopt"},
+        }
+        self.assertIs(_modelopt_quant_section(flat), flat)
+        self.assertEqual(_modelopt_quant_section(flat)["quant_algo"], "FP8")
+
    def test_mixed_precision_override_does_not_hijack_w4afp8(self):
        self.assertIsNone(
            ModelOptMixedPrecisionConfig.override_quantization_method(
@@ -379,13 +379,12 @@ class TestPrefetchDispatch(CustomTestCase):
        return DefaultModelLoader(load_config)

    def _make_source(self):
-        # model_config=None skips maybe_add_mtp_safetensors.
-        return SimpleNamespace(
+        # model_config=None skips maybe_add_mtp_safetensors. A real Source
+        # (not a stand-in) so new fields with defaults are picked up.
+        return DefaultModelLoader.Source(
            model_or_path="/dummy",
            revision=None,
            fall_back_to_pt=False,
-            model_config=None,
-            prefix="",
        )

    def _server_args(self, prefetch, disable_mmap=False, drop_cache=False):
@@ -21,6 +21,7 @@ def _write_index(folder, weight_map):

 def _touch(folder, name):
    path = os.path.join(folder, name)
+    os.makedirs(os.path.dirname(path), exist_ok=True)
    open(path, "w").close()
    return path

@@ -72,6 +73,52 @@ class TestFilterDuplicateSafetensorsFiles(CustomTestCase):
        )
        self.assertEqual(sorted(result), sorted([shard1, shard2]))

+    def test_missing_shard_outside_allow_patterns_is_ignored(self):
+        # Cosmos3-style checkpoints use one root index for multiple subfolder
+        # weight sources. Loading the transformer source should not require the
+        # vision encoder shard to already be present; the secondary source
+        # downloads and loads it separately.
+        _write_index(
+            self.folder,
+            {
+                "llm": "transformer/diffusion_pytorch_model.safetensors",
+                "vit": "vision_encoder/model.safetensors",
+            },
+        )
+        transformer = _touch(
+            self.folder, "transformer/diffusion_pytorch_model.safetensors"
+        )
+
+        result = filter_duplicate_safetensors_files(
+            hf_weights_files=[transformer],
+            hf_folder=self.folder,
+            index_file=INDEX_NAME,
+            allow_patterns=["transformer/*.safetensors"],
+        )
+        self.assertEqual(result, [transformer])
+
+    def test_missing_shard_inside_allow_patterns_raises(self):
+        _write_index(
+            self.folder,
+            {
+                "llm1": "transformer/model-00001-of-00002.safetensors",
+                "llm2": "transformer/model-00002-of-00002.safetensors",
+                "vit": "vision_encoder/model.safetensors",
+            },
+        )
+        transformer = _touch(
+            self.folder, "transformer/model-00001-of-00002.safetensors"
+        )
+
+        with self.assertRaises(RuntimeError) as cm:
+            filter_duplicate_safetensors_files(
+                hf_weights_files=[transformer],
+                hf_folder=self.folder,
+                index_file=INDEX_NAME,
+                allow_patterns=["transformer/*.safetensors"],
+            )
+        self.assertIn("model-00002-of-00002.safetensors", str(cm.exception))
+
    def test_single_file_model_no_index_returns_unchanged(self):
        # No index on disk (single-file / dummy / object-storage): early return.
        single = _touch(self.folder, "model.safetensors")
@@ -0,0 +1,310 @@
+"""Unit tests for the Cosmos3 reasoner (understanding tower).
+
+Covers the three pieces of load-time logic that make the diffusers-layout
+Cosmos3 checkpoint loadable through the Qwen3-VL inference stack:
+
+1. ``Cosmos3ForConditionalGeneration.hf_to_sglang_mapper`` - renames the
+   understanding-tower keys into the nested Qwen3-VL checkpoint form and drops
+   the generation-tower weights.
+2. ``Cosmos3Config`` - reuses the Qwen3-VL schema under ``model_type
+   "cosmos3_omni"`` and is registered with ``AutoConfig``.
+3. ``DefaultModelLoader`` ``allow_patterns_overrides`` - globs weights from the
+   ``transformer/`` and ``vision_encoder/`` subfolders rather than the repo root.
+
+All of this is pure CPU logic (no server / engine launch).
+"""
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=15, suite="base-a-test-cpu")
+
+import os
+import tempfile
+import unittest
+
+import safetensors.torch
+import torch
+
+from sglang.srt.configs import Cosmos3Config
+from sglang.srt.configs.load_config import LoadConfig, LoadFormat
+from sglang.srt.model_loader.loader import DefaultModelLoader
+from sglang.srt.models.cosmos3 import Cosmos3ForConditionalGeneration
+from sglang.srt.runtime_context import get_context
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestCosmos3WeightsMapper(CustomTestCase):
+    """Validate the understanding-tower rename + generation-tower drop rules."""
+
+    def setUp(self):
+        self.mapper = Cosmos3ForConditionalGeneration.hf_to_sglang_mapper
+
+    def test_understanding_tower_is_renamed(self):
+        # Flat Cosmos3 keys -> nested Qwen3-VL checkpoint keys. The parent
+        # Qwen3VLForConditionalGeneration.load_weights then strips the
+        # `model.language_model.`/`model.visual.` prefixes and fuses q/k/v.
+        inputs = [
+            "layers.0.self_attn.to_q.weight",
+            "layers.0.self_attn.to_k.weight",
+            "layers.0.self_attn.to_v.weight",
+            "layers.0.self_attn.to_out.weight",
+            "layers.0.self_attn.norm_q.weight",
+            "layers.0.self_attn.norm_k.weight",
+            "layers.0.mlp.gate_proj.weight",
+            "layers.0.mlp.up_proj.weight",
+            "layers.0.mlp.down_proj.weight",
+            "layers.0.input_layernorm.weight",
+            "embed_tokens.weight",
+            "norm.weight",
+            "lm_head.weight",
+        ]
+        expected = [
+            "model.language_model.layers.0.self_attn.q_proj.weight",
+            "model.language_model.layers.0.self_attn.k_proj.weight",
+            "model.language_model.layers.0.self_attn.v_proj.weight",
+            "model.language_model.layers.0.self_attn.o_proj.weight",
+            "model.language_model.layers.0.self_attn.q_norm.weight",
+            "model.language_model.layers.0.self_attn.k_norm.weight",
+            "model.language_model.layers.0.mlp.gate_proj.weight",
+            "model.language_model.layers.0.mlp.up_proj.weight",
+            "model.language_model.layers.0.mlp.down_proj.weight",
+            "model.language_model.layers.0.input_layernorm.weight",
+            "model.language_model.embed_tokens.weight",
+            "model.language_model.norm.weight",
+            "lm_head.weight",
+        ]
+        self.assertEqual(self.mapper.apply_list(inputs), expected)
+
+    def test_vision_encoder_is_prefixed(self):
+        inputs = [
+            "blocks.0.attn.qkv.weight",
+            "merger.norm.weight",
+            "patch_embed.proj.weight",
+            "pos_embed.weight",
+            "deepstack_merger_list.0.norm.weight",
+        ]
+        expected = [
+            "model.visual.blocks.0.attn.qkv.weight",
+            "model.visual.merger.norm.weight",
+            "model.visual.patch_embed.proj.weight",
+            "model.visual.pos_embed.weight",
+            "model.visual.deepstack_merger_list.0.norm.weight",
+        ]
+        self.assertEqual(self.mapper.apply_list(inputs), expected)
+
+    def test_generation_tower_is_dropped(self):
+        dropped = [
+            "layers.0.self_attn.add_q_proj.weight",
+            "layers.0.self_attn.add_k_proj.weight",
+            "layers.0.self_attn.add_v_proj.weight",
+            "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.q_proj_moe_gen.weight",
+            "layers.0.mlp_moe_gen.gate_up_proj.weight",
+            "norm_moe_gen.weight",
+            "proj_in.weight",
+            "proj_out.weight",
+            "time_embedder.linear_1.weight",
+            "audio_proj_in.weight",
+            "audio_proj_out.weight",
+            "action_proj_in.weight",
+            "action_proj_out.weight",
+            "audio_modality_embed",
+            "action_modality_embed",
+        ]
+        self.assertEqual(self.mapper.apply_list(dropped), [])
+
+    def test_modelopt_quantizer_buffers_are_dropped(self):
+        # FP8 Cosmos3 exports keep ModelOpt TensorQuantizer state alongside the
+        # inference weight_scale/input_scale tensors. SGLang uses the latter;
+        # transformers restores the former via ModelOpt HF checkpointing.
+        dropped = [
+            "layers.0.mlp.gate_proj.input_quantizer._amax",
+            "layers.0.mlp.gate_proj.weight_quantizer._amax",
+            "layers.0.mlp.gate_proj.weight_quantizer._scale",
+            "layers.0.self_attn.to_q.input_quantizer._amax",
+            "layers.0.self_attn.to_q.weight_quantizer._amax",
+            "layers.0.self_attn.to_q.weight_quantizer._scale",
+        ]
+        self.assertEqual(self.mapper.apply_list(dropped), [])
+
+    def test_fp8_inference_scales_are_kept(self):
+        inputs = [
+            "layers.0.mlp.gate_proj.weight",
+            "layers.0.mlp.gate_proj.weight_scale",
+            "layers.0.mlp.gate_proj.input_scale",
+            "layers.0.self_attn.to_q.weight",
+            "layers.0.self_attn.to_q.weight_scale",
+            "layers.0.self_attn.to_q.input_scale",
+        ]
+        expected = [
+            "model.language_model.layers.0.mlp.gate_proj.weight",
+            "model.language_model.layers.0.mlp.gate_proj.weight_scale",
+            "model.language_model.layers.0.mlp.gate_proj.input_scale",
+            "model.language_model.layers.0.self_attn.q_proj.weight",
+            "model.language_model.layers.0.self_attn.q_proj.weight_scale",
+            "model.language_model.layers.0.self_attn.q_proj.input_scale",
+        ]
+        self.assertEqual(self.mapper.apply_list(inputs), expected)
+
+    def test_moe_gen_substring_wins_over_norm_prefix(self):
+        # `norm_moe_gen.weight` must be dropped (generation), not routed to the
+        # final `norm.` -> language-model norm.
+        self.assertEqual(self.mapper.apply_list(["norm_moe_gen.weight"]), [])
+        self.assertEqual(
+            self.mapper.apply_list(["norm.weight"]),
+            ["model.language_model.norm.weight"],
+        )
+
+
+class TestCosmos3Config(CustomTestCase):
+    def test_model_type(self):
+        self.assertEqual(Cosmos3Config.model_type, "cosmos3_omni")
+
+    def test_subconfigs_are_objects(self):
+        # The Qwen3-VL inference stack reads sub-configs as objects, e.g.
+        # `config.vision_config.hidden_size` and
+        # `config.vision_config.deepstack_visual_indexes`. Some transformers
+        # versions leave sub-configs as raw dicts after construction, which would
+        # raise `'dict' object has no attribute 'hidden_size'` at model init.
+        # Cosmos3Config coerces them into config objects, so assert that here.
+        cfg = Cosmos3Config(
+            text_config={
+                "hidden_size": 128,
+                "num_hidden_layers": 2,
+                "num_attention_heads": 4,
+            },
+            vision_config={"depth": 3, "hidden_size": 64},
+        )
+        self.assertNotIsInstance(cfg.text_config, dict)
+        self.assertNotIsInstance(cfg.vision_config, dict)
+        self.assertEqual(cfg.text_config.hidden_size, 128)
+        self.assertEqual(cfg.text_config.num_hidden_layers, 2)
+        self.assertEqual(cfg.text_config.num_attention_heads, 4)
+        self.assertEqual(cfg.vision_config.hidden_size, 64)
+        self.assertEqual(cfg.vision_config.depth, 3)
+
+    def test_registered_with_autoconfig(self):
+        # Importing common runs the AutoConfig registration side effects.
+        # Newer transformers ship a native Cosmos3OmniConfig owning this model
+        # type, and sglang's registration deliberately yields to it — so assert
+        # the contract the Qwen3-VL stack needs from whichever class resolves,
+        # not one implementation class.
+        from transformers import AutoConfig
+
+        import sglang.srt.utils.hf_transformers.common  # noqa: F401
+
+        cfg = AutoConfig.for_model(
+            "cosmos3_omni",
+            text_config={"hidden_size": 128, "num_hidden_layers": 2},
+            vision_config={"hidden_size": 64},
+        )
+        self.assertEqual(cfg.model_type, "cosmos3_omni")
+        # Sub-configs must resolve to attribute-accessible objects, not dicts:
+        # the inference stack reads e.g. `config.vision_config.hidden_size`.
+        self.assertEqual(cfg.text_config.hidden_size, 128)
+        self.assertEqual(cfg.vision_config.hidden_size, 64)
+
+
+class TestCosmos3MropeIndex(CustomTestCase):
+    """Cosmos3 reuses the Qwen3-VL mrope path.
+
+    The multimodal processor calls ``MRotaryEmbedding.get_rope_index`` with the
+    config's ``model_type``. Because Cosmos3 declares its own ``cosmos3_omni``
+    type (for AutoConfig resolution), the mrope dispatch must recognize it as a
+    Qwen3-VL-family model or it raises ``RuntimeError: Unimplemented model type:
+    cosmos3_omni``. This guards that regression.
+    """
+
+    def _get_rope_index(self, model_type):
+        from sglang.srt.layers.rotary_embedding.mrope import MRotaryEmbedding
+
+        # A single image: grid [t=1, h=2, w=2] with spatial_merge_size 2 expands
+        # to exactly one image placeholder token after the vision-start token.
+        input_ids = torch.tensor([[1, 99, 100, 2]], dtype=torch.long)
+        image_grid_thw = torch.tensor([[1, 2, 2]], dtype=torch.long)
+        return MRotaryEmbedding.get_rope_index(
+            spatial_merge_size=2,
+            image_token_id=100,
+            video_token_id=101,
+            vision_start_token_id=99,
+            model_type=model_type,
+            input_ids=input_ids,
+            image_grid_thw=image_grid_thw,
+        )
+
+    def test_cosmos3_omni_is_supported(self):
+        # Must not raise "Unimplemented model type".
+        positions, delta = self._get_rope_index("cosmos3_omni")
+        self.assertEqual(positions.shape[0], 3)
+
+    def test_cosmos3_matches_qwen3_vl(self):
+        # Cosmos3 must produce identical rope indices to the qwen3_vl path it
+        # reuses, so behavior can't silently diverge.
+        pos_cosmos, delta_cosmos = self._get_rope_index("cosmos3_omni")
+        pos_qwen, delta_qwen = self._get_rope_index("qwen3_vl")
+        self.assertTrue(torch.equal(pos_cosmos, pos_qwen))
+        self.assertTrue(torch.equal(delta_cosmos, delta_qwen))
+
+
+class TestAllowPatternsOverrides(CustomTestCase):
+    """Validate DefaultModelLoader subfolder globbing for diffusers layouts."""
+
+    def setUp(self):
+        # Publish a default config so _prepare_weights reads real bags; the
+        # default model_checksum=None skips checksum verification.
+        override = get_context().override_server_args()
+        override.install()
+        self.addCleanup(override.restore)
+        self.loader = DefaultModelLoader(LoadConfig(load_format=LoadFormat.AUTO))
+
+    def _make_checkpoint(self, root):
+        os.makedirs(os.path.join(root, "transformer"))
+        os.makedirs(os.path.join(root, "vision_encoder"))
+        # A root-level file that must be ignored when an override is given.
+        safetensors.torch.save_file(
+            {"root": torch.zeros(2)}, os.path.join(root, "model.safetensors")
+        )
+        safetensors.torch.save_file(
+            {"llm": torch.zeros(2)},
+            os.path.join(root, "transformer", "diffusion_pytorch_model.safetensors"),
+        )
+        safetensors.torch.save_file(
+            {"vit": torch.zeros(2)},
+            os.path.join(root, "vision_encoder", "model.safetensors"),
+        )
+
+    def test_override_selects_transformer_subfolder(self):
+        with tempfile.TemporaryDirectory() as root:
+            self._make_checkpoint(root)
+            _, files, use_safetensors = self.loader._prepare_weights(
+                root, None, True, ["transformer/*.safetensors"]
+            )
+            self.assertTrue(use_safetensors)
+            self.assertEqual(len(files), 1)
+            self.assertTrue(
+                files[0].endswith("transformer/diffusion_pytorch_model.safetensors")
+            )
+
+    def test_override_selects_vision_subfolder(self):
+        with tempfile.TemporaryDirectory() as root:
+            self._make_checkpoint(root)
+            _, files, _ = self.loader._prepare_weights(
+                root, None, True, ["vision_encoder/*.safetensors"]
+            )
+            self.assertEqual(len(files), 1)
+            self.assertTrue(files[0].endswith("vision_encoder/model.safetensors"))
+
+    def test_no_override_globs_repo_root(self):
+        with tempfile.TemporaryDirectory() as root:
+            self._make_checkpoint(root)
+            _, files, _ = self.loader._prepare_weights(root, None, True)
+            # Without an override only the root-level file is discovered.
+            self.assertEqual(
+                [os.path.basename(f) for f in files], ["model.safetensors"]
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()
@@ -0,0 +1,442 @@
+# Copyright 2023-2025 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# CPU coverage for Cosmos3-Edge checkpoint mapping and video prompt preparation.
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+
+import asyncio
+import unittest
+
+import numpy as np
+import torch
+
+from sglang.srt.configs.cosmos3 import Cosmos3EdgeConfig
+from sglang.srt.layers.rotary_embedding import MRotaryEmbedding
+from sglang.srt.managers.schedule_batch import (
+    Modality,
+    MultimodalDataItem,
+    MultimodalInputFormat,
+)
+from sglang.srt.models.cosmos3_edge import Cosmos3EdgeForConditionalGeneration
+from sglang.srt.multimodal.processors.base_processor import (
+    BaseMultiModalProcessorOutput,
+    MultimodalSpecialTokens,
+)
+from sglang.srt.multimodal.processors.cosmos3_edge import (
+    Cosmos3EdgeProcessor,
+    _smart_resize,
+)
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestCosmos3EdgeConfig(CustomTestCase):
+    def test_checkpoint_without_pad_token_has_explicit_none(self):
+        config = Cosmos3EdgeConfig(text_config={"eos_token_id": 11})
+
+        self.assertTrue(hasattr(config.text_config, "pad_token_id"))
+        self.assertIsNone(config.text_config.pad_token_id)
+        self.assertEqual(config.text_config.eos_token_id, 11)
+        self.assertFalse(config.text_config.tie_word_embeddings)
+        self.assertTrue(hasattr(config, "pad_token_id"))
+        self.assertIsNone(config.pad_token_id)
+        self.assertEqual(config.eos_token_id, 11)
+
+
+class TestCosmos3EdgeWeightsMapper(CustomTestCase):
+    def setUp(self):
+        self.mapper = Cosmos3EdgeForConditionalGeneration.hf_to_sglang_mapper
+
+    def test_text_tower_is_renamed_for_arcee(self):
+        inputs = [
+            "embed_tokens.weight",
+            "layers.0.input_layernorm.weight",
+            "layers.0.self_attn.to_q.weight",
+            "layers.0.self_attn.to_k.weight",
+            "layers.0.self_attn.to_v.weight",
+            "layers.0.self_attn.to_out.weight",
+            "layers.0.mlp.up_proj.weight",
+            "layers.0.mlp.down_proj.weight",
+            "norm.weight",
+            "lm_head.weight",
+        ]
+        expected = [
+            "model.embed_tokens.weight",
+            "model.layers.0.input_layernorm.weight",
+            "model.layers.0.self_attn.q_proj.weight",
+            "model.layers.0.self_attn.k_proj.weight",
+            "model.layers.0.self_attn.v_proj.weight",
+            "model.layers.0.self_attn.o_proj.weight",
+            "model.layers.0.mlp.up_proj.weight",
+            "model.layers.0.mlp.down_proj.weight",
+            "model.norm.weight",
+            "lm_head.weight",
+        ]
+        self.assertEqual(self.mapper.apply_list(inputs), expected)
+
+    def test_generation_and_routed_vision_weights_are_dropped(self):
+        dropped = [
+            "layers.0.self_attn.k_norm_und_for_gen.weight",
+            "layers.0.self_attn.add_q_proj.weight",
+            "layers.0.self_attn.to_add_out.weight",
+            "layers.0.mlp_moe_gen.up_proj.weight",
+            "proj_in.weight",
+            "time_embedder.linear_1.weight",
+            "model.visual.encoder.layers.0.self_attn.q_proj.weight",
+            "model.projector.linear_fc1.weight",
+        ]
+        self.assertEqual(self.mapper.apply_list(dropped), [])
+
+
+class TestCosmos3EdgeVideoSampling(CustomTestCase):
+    def _processor(self, video_config=None):
+        processor = object.__new__(Cosmos3EdgeProcessor)
+        processor.video_config = video_config or {}
+        return processor
+
+    def test_default_sampling_matches_qwen3_vl(self):
+        processor = self._processor()
+        indices = processor._select_frame_indices(total_frames=300, video_fps=30.0)
+        expected = np.linspace(0, 299, num=20).round().astype(np.int64).tolist()
+        self.assertEqual(indices, expected)
+
+    def test_default_sampling_clamps_to_frame_limits(self):
+        processor = self._processor()
+        self.assertEqual(
+            len(processor._select_frame_indices(total_frames=30, video_fps=30.0)),
+            4,
+        )
+        self.assertEqual(
+            len(processor._select_frame_indices(total_frames=15_000, video_fps=10.0)),
+            768,
+        )
+
+    def test_num_frames_and_legacy_nframes(self):
+        num_frames = self._processor({"num_frames": 5})
+        nframes = self._processor({"nframes": 5})
+        expected = [0, 25, 50, 74, 99]
+        self.assertEqual(num_frames._select_frame_indices(100, 30.0), expected)
+        self.assertEqual(nframes._select_frame_indices(100, 30.0), expected)
+
+    def test_explicit_frame_count_and_fps_are_mutually_exclusive(self):
+        processor = self._processor({"num_frames": 5, "fps": 2.0})
+        with self.assertRaisesRegex(ValueError, "Specify only one"):
+            processor._select_frame_indices(100, 30.0)
+
+    def test_missing_source_fps_uses_24_fps(self):
+        processor = self._processor()
+        self.assertEqual(
+            len(processor._select_frame_indices(total_frames=240, video_fps=0.0)),
+            20,
+        )
+
+
+class TestCosmos3EdgeResize(CustomTestCase):
+    def test_video_max_pixels_is_a_total_temporal_budget(self):
+        height, width = _smart_resize(
+            1024,
+            1024,
+            factor=32,
+            min_pixels=4096,
+            max_pixels=4 * 1024 * 1024,
+            num_frames=16,
+        )
+        self.assertEqual((height, width), (512, 512))
+        self.assertLessEqual(16 * height * width, 4 * 1024 * 1024)
+
+    def test_image_resize_keeps_single_frame_semantics(self):
+        height, width = _smart_resize(
+            1024,
+            1024,
+            factor=32,
+            min_pixels=4096,
+            max_pixels=512 * 512,
+        )
+        self.assertEqual((height, width), (512, 512))
+
+
+class _TimestampTokenizer:
+    def encode(self, text, add_special_tokens=False):
+        del add_special_tokens
+        timestamp = float(text.removeprefix("<").split()[0])
+        return [1000 + int(timestamp * 10)]
+
+
+class TestCosmos3EdgePromptExpansion(CustomTestCase):
+    def _processor(self):
+        processor = object.__new__(Cosmos3EdgeProcessor)
+        processor.IMAGE_TOKEN_ID = 19
+        processor.VIDEO_TOKEN_ID = 18
+        processor.IM_START_TOKEN_ID = 20
+        processor.IM_END_TOKEN_ID = 21
+        processor._spatial_merge_size = 2
+        processor._tokenizer = _TimestampTokenizer()
+        return processor
+
+    def test_video_placeholder_expands_once_per_frame(self):
+        processor = self._processor()
+        prompt = [7, 20, 18, 21, 8]
+        video_grid_thw = torch.tensor([[3, 4, 4]], dtype=torch.long)
+        timestamps = [[0.0, 0.5, 1.0]]
+
+        input_ids, offsets, modalities = processor._build_input_ids(
+            prompt,
+            img_grid_thw=None,
+            video_grid_thw=video_grid_thw,
+            video_timestamps=timestamps,
+        )
+
+        expected = [7]
+        expected_offsets = []
+        for timestamp_id in (1000, 1005, 1010):
+            expected.extend([timestamp_id, 20])
+            offset_start = len(expected)
+            expected.extend([18] * 4)
+            expected_offsets.append((offset_start, len(expected) - 1))
+            expected.append(21)
+        expected.append(8)
+
+        self.assertEqual(input_ids, expected)
+        self.assertEqual(offsets, [expected_offsets])
+        self.assertEqual(modalities, [Modality.VIDEO])
+
+    def test_expanded_video_mrope_matches_qwen3_vl(self):
+        processor = self._processor()
+        video_grid_thw = torch.tensor([[3, 4, 4]], dtype=torch.long)
+        input_ids, _, _ = processor._build_input_ids(
+            [7, 20, 18, 21, 8],
+            img_grid_thw=None,
+            video_grid_thw=video_grid_thw,
+            video_timestamps=[[0.0, 0.5, 1.0]],
+        )
+        input_ids = torch.tensor([input_ids], dtype=torch.long)
+
+        kwargs = dict(
+            spatial_merge_size=2,
+            image_token_id=19,
+            video_token_id=18,
+            vision_start_token_id=20,
+            input_ids=input_ids,
+            video_grid_thw=video_grid_thw,
+        )
+        edge_positions, edge_delta = MRotaryEmbedding.get_rope_index(
+            model_type="cosmos3_edge", **kwargs
+        )
+        qwen_positions, qwen_delta = MRotaryEmbedding.get_rope_index(
+            model_type="qwen3_vl", **kwargs
+        )
+
+        self.assertEqual(edge_positions.shape, (3, 1, len(input_ids[0])))
+        self.assertTrue(torch.equal(edge_positions, qwen_positions))
+        self.assertTrue(torch.equal(edge_delta, qwen_delta))
+
+
+def _processed_input_processor():
+    processor = object.__new__(Cosmos3EdgeProcessor)
+    processor.IMAGE_TOKEN_ID = 19
+    processor.VIDEO_TOKEN_ID = 18
+    processor.IM_START_TOKEN_ID = 20
+    processor.IM_END_TOKEN_ID = 21
+    processor.vision_start_token_id = 20
+    processor.model_type = "cosmos3_edge"
+    processor._spatial_merge_size = 2
+    processor._tokenizer = _TimestampTokenizer()
+    processor._processor = processor._tokenizer
+    processor.mm_tokens = MultimodalSpecialTokens(
+        image_token="<image>",
+        video_token="<video>",
+        image_token_id=19,
+        video_token_id=18,
+    )
+    processor.mm_processor_executor = None
+    processor.use_cuda_ipc = False
+    processor.precompute_hash_before_cpu_transfer = False
+    processor.ATTR_NAME_TO_MODALITY = {
+        "pixel_values": Modality.IMAGE,
+        "image_grid_thw": Modality.IMAGE,
+        "pixel_values_videos": Modality.VIDEO,
+        "video_grid_thw": Modality.VIDEO,
+    }
+    processor.FEATURE_NAMES = ["pixel_values", "pixel_values_videos"]
+    return processor
+
+
+class TestCosmos3EdgeProcessedInputs(CustomTestCase):
+    def test_image_processor_output_keeps_pixels_for_vision_tower(self):
+        processor = _processed_input_processor()
+        input_ids = [7, 20, 19, 19, 19, 19, 21, 8]
+        grid = torch.tensor([[1, 4, 4]], dtype=torch.long)
+        pixels = torch.randn(16, 768)
+        processor_data = {
+            "format": "processor_output",
+            "input_ids": input_ids,
+            "pixel_values": pixels,
+            "image_grid_thw": grid,
+        }
+        base_output = BaseMultiModalProcessorOutput(
+            input_text="",
+            input_ids=input_ids,
+            images=[processor_data],
+        )
+
+        output = asyncio.run(processor._process_preprocessed_mm_data(base_output))
+
+        self.assertEqual(output.input_ids, input_ids)
+        self.assertEqual(len(output.mm_items), 1)
+        item = output.mm_items[0]
+        self.assertEqual(item.format, MultimodalInputFormat.PROCESSOR_OUTPUT)
+        self.assertIs(item.feature, pixels)
+        self.assertEqual(item.offsets, [(2, 5)])
+        self.assertEqual(output.mrope_positions.shape, (3, len(input_ids)))
+
+    def test_image_precomputed_embedding_is_preserved(self):
+        processor = _processed_input_processor()
+        input_ids = [7, 20, 19, 19, 19, 19, 21, 8]
+        grid = torch.tensor([[1, 4, 4]], dtype=torch.long)
+        embeddings = torch.randn(4, 2048)
+        processor_data = {
+            "format": "precomputed_embedding",
+            "input_ids": input_ids,
+            "feature": embeddings,
+            "image_grid_thw": grid,
+        }
+        base_output = BaseMultiModalProcessorOutput(
+            input_text="",
+            input_ids=input_ids,
+            images=[processor_data],
+        )
+
+        output = asyncio.run(processor._process_preprocessed_mm_data(base_output))
+
+        item = output.mm_items[0]
+        self.assertEqual(item.format, MultimodalInputFormat.PRECOMPUTED_EMBEDDING)
+        self.assertIs(item.feature, embeddings)
+        self.assertEqual(item.offsets, [(2, 5)])
+        self.assertEqual(output.mrope_positions.shape, (3, len(input_ids)))
+
+    def test_video_processor_output_retains_per_frame_offsets(self):
+        processor = _processed_input_processor()
+        grid = torch.tensor([[3, 4, 4]], dtype=torch.long)
+        input_ids, offsets, _ = processor._build_input_ids(
+            [7, 20, 18, 21, 8],
+            img_grid_thw=None,
+            video_grid_thw=grid,
+            video_timestamps=[[0.0, 0.5, 1.0]],
+        )
+        pixels = torch.randn(48, 768)
+        processor_data = {
+            "format": "processor_output",
+            "input_ids": input_ids,
+            "pixel_values_videos": pixels,
+            "video_grid_thw": grid,
+        }
+        base_output = BaseMultiModalProcessorOutput(
+            input_text="",
+            input_ids=input_ids,
+            videos=[processor_data],
+        )
+
+        output = asyncio.run(processor._process_preprocessed_mm_data(base_output))
+
+        self.assertEqual(len(output.mm_items), 1)
+        item = output.mm_items[0]
+        self.assertEqual(item.format, MultimodalInputFormat.PROCESSOR_OUTPUT)
+        self.assertTrue(torch.equal(item.feature, pixels))
+        self.assertEqual(item.offsets, offsets[0])
+        self.assertEqual(output.mrope_positions.shape, (3, len(input_ids)))
+
+    def test_video_precomputed_embedding_is_preserved(self):
+        processor = _processed_input_processor()
+        grid = torch.tensor([[3, 4, 4]], dtype=torch.long)
+        input_ids, offsets, _ = processor._build_input_ids(
+            [7, 20, 18, 21, 8],
+            img_grid_thw=None,
+            video_grid_thw=grid,
+            video_timestamps=[[0.0, 0.5, 1.0]],
+        )
+        embeddings = torch.randn(12, 2048)
+        processor_data = {
+            "format": "precomputed_embedding",
+            "input_ids": input_ids,
+            "feature": embeddings,
+            "video_grid_thw": grid,
+        }
+        base_output = BaseMultiModalProcessorOutput(
+            input_text="",
+            input_ids=input_ids,
+            videos=[processor_data],
+        )
+
+        output = asyncio.run(processor._process_preprocessed_mm_data(base_output))
+
+        self.assertEqual(len(output.mm_items), 1)
+        item = output.mm_items[0]
+        self.assertEqual(item.format, MultimodalInputFormat.PRECOMPUTED_EMBEDDING)
+        self.assertIs(item.feature, embeddings)
+        self.assertEqual(item.offsets, offsets[0])
+        self.assertEqual(output.mrope_positions.shape, (3, len(input_ids)))
+
+
+class TestCosmos3EdgePrecomputedVisionFeatures(CustomTestCase):
+    def _model(self):
+        model = Cosmos3EdgeForConditionalGeneration.__new__(
+            Cosmos3EdgeForConditionalGeneration
+        )
+        torch.nn.Module.__init__(model)
+        model.language_model_only = False
+        return model
+
+    def test_image_and_video_embeddings_bypass_vision_tower(self):
+        model = self._model()
+        image_embeddings = torch.randn(4, 2048)
+        video_embeddings = torch.randn(12, 2048)
+        image_item = MultimodalDataItem(
+            modality=Modality.IMAGE,
+            feature=image_embeddings,
+            format=MultimodalInputFormat.PRECOMPUTED_EMBEDDING,
+        )
+        video_item = MultimodalDataItem(
+            modality=Modality.VIDEO,
+            feature=video_embeddings,
+            format=MultimodalInputFormat.PRECOMPUTED_EMBEDDING,
+        )
+
+        self.assertTrue(
+            torch.equal(model.get_image_feature([image_item]), image_embeddings)
+        )
+        self.assertTrue(
+            torch.equal(model.get_video_feature([video_item]), video_embeddings)
+        )
+
+    def test_mixed_features_are_rejected(self):
+        model = self._model()
+        items = [
+            MultimodalDataItem(
+                modality=Modality.IMAGE,
+                feature=torch.randn(4, 2048),
+                format=MultimodalInputFormat.PRECOMPUTED_EMBEDDING,
+            ),
+            MultimodalDataItem(
+                modality=Modality.IMAGE,
+                feature=torch.randn(16, 768),
+            ),
+        ]
+
+        with self.assertRaisesRegex(ValueError, "cannot mix"):
+            model.get_image_feature(items)
+
+
+if __name__ == "__main__":
+    unittest.main()