From b2eb0fa51ef01058155a0b61f9651bc5d59d5b49 Mon Sep 17 00:00:00 2001
From: Xiaoyu Zhang <1182563586@qq.com>
Date: Mon, 24 Aug 2026 08:51:58 +0800
Subject: [PATCH] [diffusion] Keep Cosmos3 Nano resident on high-memory GPUs
 (#36000)

---
 docs/cookbook/diffusion/Cosmos/Cosmos3.mdx    |  6 +++
 .../configs/pipeline_configs/cosmos3.py       | 12 +++++
 .../test/unit/test_server_args.py             | 50 +++++++++++++++++++
 3 files changed, 68 insertions(+)

diff --git a/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx b/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx
index 6fd48027f..cf982620e 100644
--- a/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx
+++ b/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx
@@ -62,6 +62,12 @@ sglang serve \
   --num-gpus 1
 ```
 
+With `--performance-mode auto`, a Cosmos3 Nano checkpoint keeps its DiT and
+VAE resident when every selected GPU has at least 120 GiB available at
+startup. Below that threshold, auto mode retains the conservative DiT
+component-offload policy. This high-memory override is limited to Nano;
+Cosmos3 Super checkpoints keep their existing multi-GPU placement defaults.
+
 For `Cosmos3-Super`, split the model across multiple GPUs:
 
 ```bash Command
diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py b/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py
index 8065e2503..e1b634f2c 100644
--- a/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py
+++ b/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py
@@ -19,6 +19,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
     ModelTaskType,
     PipelineConfig,
 )
+from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import (
+    ModelDeploymentConfig,
+)
 
 COSMOS3_EDGE_BACKBONE_TYPE = "cosmos3_edge_nemotron_dense"
 
@@ -165,3 +168,12 @@ class Cosmos3Config(PipelineConfig):
         # and action-capable checkpoints. The loaded transformer validates that
         # an action head is actually present when an action request is submitted.
         return True
+
+    def get_model_deployment_config(self) -> ModelDeploymentConfig:
+        if "cosmos3-nano" not in self.model_path.lower():
+            return ModelDeploymentConfig()
+
+        return ModelDeploymentConfig(
+            keep_resident_min_available_gb=120,
+            keep_resident_components=("dit", "vae"),
+        )
diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py
index 6a73359c9..9fd2aa7d6 100644
--- a/python/sglang/multimodal_gen/test/unit/test_server_args.py
+++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py
@@ -16,6 +16,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
     ModelTaskType,
     PipelineConfig,
 )
+from sglang.multimodal_gen.configs.pipeline_configs.cosmos3 import Cosmos3Config
 from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import FastHunyuanConfig
 from sglang.multimodal_gen.configs.pipeline_configs.lingbot_world import (
     LingBotWorldCausalDMDConfig,
@@ -1383,6 +1384,9 @@ class TestOffloadDefaults(unittest.TestCase):
 
     def test_pipeline_configs_declare_auto_tune_hints(self):
         qwen_deployment = QwenImagePipelineConfig().get_model_deployment_config()
+        cosmos3_deployment = Cosmos3Config(
+            model_path="nvidia/Cosmos3-Nano"
+        ).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()
@@ -1394,6 +1398,9 @@ 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_components, ("dit", "vae"))
+
         self.assertIsNone(wan_deployment.fsdp_auto_min_available_memory_gb)
         self.assertEqual(wan_deployment.dit_layerwise_offload_modes, ("memory",))
         self.assertEqual(wan_deployment.keep_resident_min_available_gb, 60)
@@ -1842,6 +1849,49 @@ class TestOffloadDefaults(unittest.TestCase):
             ["text_encoder", "image_encoder"],
         )
 
+    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,
+            kwargs={
+                "model_path": "nvidia/Cosmos3-Nano",
+                "performance_mode": "auto",
+            },
+        )
+
+        self.assertFalse(args.dit_cpu_offload)
+        self.assertFalse(args.vae_cpu_offload)
+        self.assertEqual(
+            args.layerwise_offload_components,
+            ["text_encoder", "image_encoder"],
+        )
+
+    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,
+            kwargs={
+                "model_path": "nvidia/Cosmos3-Nano",
+                "performance_mode": "auto",
+            },
+        )
+
+        self.assertTrue(args.dit_cpu_offload)
+        self.assertFalse(args.vae_cpu_offload)
+
+    def test_auto_cosmos3_super_keeps_default_offload_policy(self):
+        args = self._from_dict_with_pipeline_config(
+            Cosmos3Config(model_path="nvidia/Cosmos3-Super"),
+            available_memory_gb=139,
+            kwargs={
+                "model_path": "nvidia/Cosmos3-Super",
+                "performance_mode": "auto",
+            },
+        )
+
+        self.assertTrue(args.dit_cpu_offload)
+        self.assertFalse(args.vae_cpu_offload)
+
     def test_memory_sana_wm_layerwise_offload_adds_dit(self):
         args = self._from_dict_with_pipeline_config(
             SanaWMPipelineConfig(),