[diffusion] Keep the Cosmos3 Super DiT resident on high-memory GPUs (#36375)

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

docs/cookbook/diffusion/Cosmos/Cosmos3.mdx

@@ -62,11 +62,13 @@ 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.
+With `--performance-mode auto`, a Cosmos3 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. Cosmos3 runs one DiT per pipeline, so component
+offload above the threshold only pays to copy the weights out to host memory
+and back on every request. Serve `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/.claude/skills/sglang-diffusion-performance/SKILL.md

@@ -384,7 +384,7 @@ Use these as first commands to benchmark, not as universal winners.
 | Wan2.2 A14B T2V/I2V | 1280x720, 81 frames | Nightly: `--num-gpus 4 --enable-cfg-parallel --ulysses-degree 2 --text-encoder-cpu-offload --pin-cpu-memory` | For lowest latency, also benchmark pure Ulysses on the same GPUs. |
 | Wan2.2 TI2V 5B | 1280x720, 81 frames, 1 GPU | `--enable-torch-compile --warmup-mode request` | Keep the input image and motion prompt fixed when comparing sparse attention or Cache-DiT. |
 | Wan2.1 / FastWan / TurboWan variants | 480p or 720p video, family defaults | `--enable-torch-compile --warmup-mode request`; add `--ulysses-degree` / CFG parallel only after measuring | Current registry includes Wan2.1, FastWan2.1, FastWan2.2 TI2V, TurboWan2.1, TurboWan2.2 I2V, and Wan2.1-Fun InP. Use the compatibility matrix and benchmark presets before choosing topology. |
-| Cosmos3 Nano / Super | T2I: 1024x1024 with `--num-frames 1`; T2V/I2V: 480p/720p video | Start with `--performance-mode auto --warmup-mode request`; use `SGLANG_DISABLE_COSMOS3_GUARDRAILS=1` only for benchmark isolation, and compare compile separately | One checkpoint serves T2I/T2V/I2V. Mode is request-driven: `num_frames == 1` means T2I, `--image-path` means I2V. On GPUs with at least 120 GiB available, auto mode keeps Cosmos3 Nano's DiT and VAE resident; a 1xH200 832x480x9f, 4-step eager ABBA reduced e2e from 1.576 to 0.428 seconds with exact output parity. The override is Nano-only; keep Super on its conservative multi-GPU policy. |
+| Cosmos3 Nano / Super | T2I: 1024x1024 with `--num-frames 1`; T2V/I2V: 480p/720p video | Start with `--performance-mode auto --warmup-mode request`; use `SGLANG_DISABLE_COSMOS3_GUARDRAILS=1` only for benchmark isolation, and compare compile separately | One checkpoint serves T2I/T2V/I2V. Mode is request-driven: `num_frames == 1` means T2I, `--image-path` means I2V. On GPUs with at least 120 GiB available, auto mode keeps the Cosmos3 DiT and VAE resident for every checkpoint in the family; a 1xH200 832x480x9f, 4-step eager ABBA reduced e2e from 1.576 to 0.428 seconds with exact output parity. Cosmos3 runs one DiT per pipeline, so component offload above that threshold only buys a DiT copy out to host memory and back per request -- it cost Cosmos3-Super 720p 81f T2V ~4s of ~115s on 2xH200. |
 | Cosmos3 Edge / distilled Super | Edge T2I: 640x640, 35 steps, 1 GPU; distilled Super T2I: 640x640, fixed 4-step schedule, 4 GPUs | Start eager with `--performance-mode manual`; use `SGLANG_DISABLE_COSMOS3_GUARDRAILS=1` only for benchmark isolation | Edge is trained for 256p/480p shapes. Distilled checkpoints own their sigma schedule and force guidance 1.0; do not override steps or flow shift. Do not retry the closed experimental Cosmos BCG path without a new lifecycle design. |
 | Ideogram 4 FP8/NVFP4 | 1024x1024, native preset defaults | `--enable-torch-compile --warmup-mode request` | Do not set `--num-inference-steps` or `--guidance-scale` directly unless you also update the Ideogram preset; sampling params derive them from `preset`. |
 | ERNIE-Image / GLM-Image / SANA / SD3 | 1024-class image, family defaults | `--enable-torch-compile --warmup-mode request`; disable offload only after checking VRAM | Treat these as current native image families. Start with benchmark/profile presets for ERNIE, GLM, and SANA; use registry/config defaults for SD3 unless you add a new preset. |

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

@@ -170,9 +170,7 @@ class Cosmos3Config(PipelineConfig):
         return True
 
     def get_model_deployment_config(self) -> ModelDeploymentConfig:
-        if "cosmos3-nano" not in self.model_path.lower():
-            return ModelDeploymentConfig()
-
+        # Keep the DiT and VAE resident when the GPUs have the headroom.
         return ModelDeploymentConfig(
             keep_resident_min_available_gb=120,
             keep_resident_components=("dit", "vae"),

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

@@ -1945,7 +1945,9 @@ class TestOffloadDefaults(unittest.TestCase):
         self.assertTrue(args.dit_cpu_offload)
         self.assertFalse(args.vae_cpu_offload)
 
-    def test_auto_cosmos3_super_keeps_default_offload_policy(self):
+    def test_auto_cosmos3_super_keeps_dit_resident_on_high_memory_gpu(self):
+        # Super is a single-DiT pipeline like Nano, so above the threshold the
+        # component-offload round trip is pure per-request copy cost.
         args = self._from_dict_with_pipeline_config(
             Cosmos3Config(model_path="nvidia/Cosmos3-Super"),
             available_memory_gb=139,
@@ -1955,6 +1957,19 @@ class TestOffloadDefaults(unittest.TestCase):
             },
         )
 
+        self.assertFalse(args.dit_cpu_offload)
+        self.assertFalse(args.vae_cpu_offload)
+
+    def test_auto_cosmos3_super_offloads_dit_below_resident_threshold(self):
+        args = self._from_dict_with_pipeline_config(
+            Cosmos3Config(model_path="nvidia/Cosmos3-Super"),
+            available_memory_gb=100,
+            kwargs={
+                "model_path": "nvidia/Cosmos3-Super",
+                "performance_mode": "auto",
+            },
+        )
+
         self.assertTrue(args.dit_cpu_offload)
         self.assertFalse(args.vae_cpu_offload)