[diffusion] Keep Cosmos3 Nano resident on 96 GB GPUs (#36641)

docs/cookbook/diffusion/Cosmos/Cosmos3.mdx

@@ -62,13 +62,14 @@ sglang serve \
   --num-gpus 1
 ```
 
-With `--performance-mode auto`, a Cosmos3 checkpoint keeps its DiT and VAE
+With `--performance-mode auto`, Cosmos3 Nano keeps its DiT and VAE resident
-resident when every selected GPU has at least 120 GiB available at startup.
+when every selected GPU has at least 90 GiB available at startup. Other
-Below that threshold, auto mode retains the conservative DiT
+Cosmos3 checkpoints use a 120 GiB threshold. Below the applicable threshold,
-component-offload policy. Cosmos3 runs one DiT per pipeline, so component
+auto mode retains the conservative DiT component-offload policy. Cosmos3 runs
-offload above the threshold only pays to copy the weights out to host memory
+one DiT per pipeline, so component offload above the threshold only pays to
-and back on every request. Serve `Cosmos3-Super` across multiple GPUs as
+copy the weights out to host memory and back on every request. Serve
-shown below so each rank holds a shard of the weights.
+`Cosmos3-Super` across multiple GPUs as shown below so each rank holds a shard
+of the weights.
 
 For `Cosmos3-Super`, split the model across multiple GPUs:
 

python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py

@@ -24,6 +24,20 @@ from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config impo
 )
 
 COSMOS3_EDGE_BACKBONE_TYPE = "cosmos3_edge_nemotron_dense"
+COSMOS3_NANO_ARCH_SIGNATURE = (4096, 36, 32)
+COSMOS3_NANO_KEEP_RESIDENT_MIN_AVAILABLE_GB = 90
+COSMOS3_DEFAULT_KEEP_RESIDENT_MIN_AVAILABLE_GB = 120
+
+
+@functools.lru_cache(maxsize=None)
+def _transformer_config(model_path: str) -> dict:
+    from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
+        get_diffusers_component_config,
+    )
+
+    return get_diffusers_component_config(
+        component_path=os.path.join(model_path, "transformer")
+    )
 
 
 @functools.lru_cache(maxsize=None)
@@ -34,19 +48,25 @@ def is_edge_checkpoint(model_path: str) -> bool:
     is available before the weights are on device (e.g. when resolving sampling
     defaults in the client process).
     """
-    from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
+    config = _transformer_config(model_path)
-        get_diffusers_component_config,
-    )
-
-    config = get_diffusers_component_config(
-        component_path=os.path.join(model_path, "transformer")
-    )
     return (
         config.get("backbone_type") == COSMOS3_EDGE_BACKBONE_TYPE
         or config.get("hidden_act") == "relu2"
     )
 
 
+@functools.lru_cache(maxsize=None)
+def is_nano_checkpoint(model_path: str) -> bool:
+    """Whether the checkpoint uses the Nano transformer architecture."""
+    config = _transformer_config(model_path)
+    signature = (
+        config.get("hidden_size"),
+        config.get("num_hidden_layers"),
+        config.get("num_attention_heads"),
+    )
+    return signature == COSMOS3_NANO_ARCH_SIGNATURE
+
+
 @functools.lru_cache(maxsize=None)
 def _distilled_sampler_config(model_path: str) -> dict | None:
     """The fixed-step sampler config for a distilled checkpoint, else ``None``.
@@ -125,6 +145,7 @@ class Cosmos3Config(PipelineConfig):
     # Pre-computed once in update_config_from_dict from the resolved model_path.
     # None until that point (e.g. in unit-test mocks that never call update_config_from_dict).
     is_edge: bool | None = None
+    is_nano: bool | None = None
     distilled_sigmas: list[float] | None = None
 
     def __post_init__(self):
@@ -145,6 +166,7 @@ class Cosmos3Config(PipelineConfig):
         if self.model_path:
             self.distilled_sigmas = get_distilled_sigmas(self.model_path)
             self.is_edge = is_edge_checkpoint(self.model_path)
+            self.is_nano = is_nano_checkpoint(self.model_path)
             if self.distilled_sigmas is not None:
                 self.scheduler_class_override = None
 
@@ -171,7 +193,17 @@ class Cosmos3Config(PipelineConfig):
 
     def get_model_deployment_config(self) -> ModelDeploymentConfig:
         # Keep the DiT and VAE resident when the GPUs have the headroom.
+        is_nano = self.is_nano
+        if is_nano is None:
+            # Directly constructed configs in callers/tests have not resolved
+            # checkpoint metadata yet; registered model IDs remain unambiguous.
+            is_nano = "cosmos3-nano" in (self.model_path or "").lower()
+        threshold_gb = (
+            COSMOS3_NANO_KEEP_RESIDENT_MIN_AVAILABLE_GB
+            if is_nano
+            else COSMOS3_DEFAULT_KEEP_RESIDENT_MIN_AVAILABLE_GB
+        )
         return ModelDeploymentConfig(
-            keep_resident_min_available_gb=120,
+            keep_resident_min_available_gb=threshold_gb,
             keep_resident_components=("dit", "vae"),
         )

python/sglang/multimodal_gen/test/unit/test_cosmos3.py

@@ -14,7 +14,10 @@ from PIL import Image
 from sglang.multimodal_gen.configs.models.dits.cosmos3video import (
     _build_cosmos3_param_names_mapping,
 )
-from sglang.multimodal_gen.configs.pipeline_configs.cosmos3 import Cosmos3Config
+from sglang.multimodal_gen.configs.pipeline_configs.cosmos3 import (
+    Cosmos3Config,
+    is_nano_checkpoint,
+)
 from sglang.multimodal_gen.configs.sample.cosmos3 import (
     COSMOS3_EDGE_SUPPORTED_RESOLUTIONS,
     Cosmos3SamplingParams,
@@ -841,6 +844,44 @@ class TestCosmos3ActionEndpoint(unittest.TestCase):
 class TestCosmos3ModelResolution(unittest.TestCase):
     """Verify Cosmos3 checkpoints resolve to the native SGLang pipeline."""
 
+    def test_nano_architecture_detection_does_not_depend_on_model_path(self):
+        cases = (
+            (
+                {
+                    "hidden_size": 4096,
+                    "num_hidden_layers": 36,
+                    "num_attention_heads": 32,
+                },
+                True,
+            ),
+            (
+                {
+                    "hidden_size": 5120,
+                    "num_hidden_layers": 64,
+                    "num_attention_heads": 64,
+                },
+                False,
+            ),
+            (
+                {
+                    "hidden_size": 2048,
+                    "num_hidden_layers": 28,
+                    "num_attention_heads": 16,
+                },
+                False,
+            ),
+        )
+        for index, (transformer_config, expected) in enumerate(cases):
+            with self.subTest(transformer_config=transformer_config):
+                is_nano_checkpoint.cache_clear()
+                with mock.patch(
+                    "sglang.multimodal_gen.configs.pipeline_configs.cosmos3._transformer_config",
+                    return_value=transformer_config,
+                ):
+                    self.assertEqual(
+                        is_nano_checkpoint(f"/models/checkpoint-{index}"), expected
+                    )
+
     def test_hf_checkpoint_uses_registered_native_pipeline_config(self):
         for model_path in (
             "nvidia/Cosmos3-Nano",

python/sglang/multimodal_gen/test/unit/test_server_args.py

@@ -1440,6 +1440,12 @@ class TestOffloadDefaults(unittest.TestCase):
         cosmos3_deployment = Cosmos3Config(
             model_path="nvidia/Cosmos3-Nano"
         ).get_model_deployment_config()
+        cosmos3_super_deployment = Cosmos3Config(
+            model_path="nvidia/Cosmos3-Super"
+        ).get_model_deployment_config()
+        local_cosmos3_nano_deployment = Cosmos3Config(
+            model_path="/models/custom-checkpoint", is_nano=True
+        ).get_model_deployment_config()
         wan_deployment = WanT2V480PConfig().get_model_deployment_config()
         mova_deployment = MOVAPipelineConfig().get_model_deployment_config()
         zimage_deployment = ZImagePipelineConfig().get_model_deployment_config()
@@ -1452,8 +1458,12 @@ class TestOffloadDefaults(unittest.TestCase):
         self.assertIsNone(qwen_deployment.fsdp_auto_min_available_memory_gb)
         self.assertEqual(qwen_deployment.dit_layerwise_offload_modes, ())
 
-        self.assertEqual(cosmos3_deployment.keep_resident_min_available_gb, 120)
+        self.assertEqual(cosmos3_deployment.keep_resident_min_available_gb, 90)
         self.assertEqual(cosmos3_deployment.keep_resident_components, ("dit", "vae"))
+        self.assertEqual(cosmos3_super_deployment.keep_resident_min_available_gb, 120)
+        self.assertEqual(
+            local_cosmos3_nano_deployment.keep_resident_min_available_gb, 90
+        )
 
         self.assertIsNone(wan_deployment.fsdp_auto_min_available_memory_gb)
         self.assertEqual(wan_deployment.dit_layerwise_offload_modes, ("memory",))
@@ -1930,7 +1940,7 @@ class TestOffloadDefaults(unittest.TestCase):
     def test_auto_cosmos3_keeps_dit_resident_on_high_memory_gpu(self):
         args = self._from_dict_with_pipeline_config(
             Cosmos3Config(model_path="nvidia/Cosmos3-Nano"),
-            available_memory_gb=139,
+            available_memory_gb=95,
             kwargs={
                 "model_path": "nvidia/Cosmos3-Nano",
                 "performance_mode": "auto",
@@ -1947,7 +1957,7 @@ class TestOffloadDefaults(unittest.TestCase):
     def test_auto_cosmos3_offloads_dit_below_resident_threshold(self):
         args = self._from_dict_with_pipeline_config(
             Cosmos3Config(model_path="nvidia/Cosmos3-Nano"),
-            available_memory_gb=100,
+            available_memory_gb=85,
             kwargs={
                 "model_path": "nvidia/Cosmos3-Nano",
                 "performance_mode": "auto",