diff --git a/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx b/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx index cf982620e..51686bbdd 100644 --- a/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx +++ b/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: diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md index 3d88bc932..9dcaff378 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md +++ b/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. | diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py b/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py index e1b634f2c..6ad73fa18 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py +++ b/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"), 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 03b5c6a86..6f954eb30 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/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)