From 96bfd2476c40bc575d87fd22c8508ece7c199614 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Sat, 22 Aug 2026 12:57:09 +0800 Subject: [PATCH] [diffusion] Enable SANA-Video breakable CUDA graphs (#35729) --- docs/docs/sglang-diffusion/api/cli.mdx | 13 ++++++ .../existing-fast-paths.md | 16 +++++-- .../sglang-diffusion-performance/SKILL.md | 4 +- .../model_padders/sana_video.py | 45 +++++++++++++++++++ .../breakable_cuda_graph/prompt_padding.py | 1 + .../runtime/server_args/server_args.py | 5 ++- .../test/unit/test_diffusion_bcg_padding.py | 27 +++++++++++ 7 files changed, 104 insertions(+), 7 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/sana_video.py diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index 316bfcca9..6ce48c555 100644 --- a/docs/docs/sglang-diffusion/api/cli.mdx +++ b/docs/docs/sglang-diffusion/api/cli.mdx @@ -234,6 +234,19 @@ sglang serve --model-path meituan-longcat/LongCat-Image \ --port 30010 ``` +SANA-Video supports the same path for fixed-resolution serving. Its default +text stage emits a fixed 300-token prompt shape, so the runtime reuses one +graph across prompt lengths without padding it to a generic text bucket: + +```bash Command +sglang serve \ + --model-path Efficient-Large-Model/SANA-Video_2B_480p_diffusers \ + --enable-breakable-cuda-graph \ + --warmup-resolutions 832x480 \ + --enable-torch-compile false \ + --port 30010 +``` + ### Component Residency Use `--component-residency COMPONENT=MODE` to assign one runtime residency mode to each native pipeline component: diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md index e9ca2c6ca..ec7b8a74e 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md @@ -306,10 +306,15 @@ framework-specific optimization workflow. joint-stream `cat`/split inside each single block. Do not misclassify those as a missed existing packed path; they are model-local structural opportunities that need their own weight-loader and parity coverage. -- SANA-Video already packs self QKV and cross KV. Its conv/modulation formulas - mirror SANA, but it does not yet call SANA's bit-exact bias-SiLU, bias-GLU, - residual-gate, LayerNorm-modulation, or one-time contiguous-layout helpers. - Reuse or extract those helpers before authoring a video-only kernel. +- SANA-Video already packs self QKV and cross KV. For fixed 832x480 serving, + its default 300-token prompt shape can reuse one breakable CUDA graph without + generic text-bucket padding. An H200 81-frame, 8-step run measured + 920.6--925.3 ms/step eager versus 797.8--798.9 ms/step with BCG, with + bit-exact final videos; reserved peak memory increased by about 3.4 GB. + Its conv/modulation formulas mirror SANA, but it does not yet call SANA's + bit-exact bias-SiLU, bias-GLU, residual-gate, LayerNorm-modulation, or + one-time contiguous-layout helpers. Reuse or extract those helpers before + authoring a video-only kernel. - LingBot Video MoE's router implements sigmoid+bias grouped top-k in `multimodal_gen/runtime/layers/moe.py`. Check parameter and output-order compatibility with `srt/layers/moe/topk.py::biased_grouped_topk` before @@ -381,6 +386,9 @@ framework-specific optimization workflow. Keep eager as the baseline because the gain is hardware-dependent; an H200 50-step, three-prompt run measured 177.0--177.3 ms/step eager versus 173.1--173.3 ms/step with BCG, with bit-exact final images. +- SANA-Video uses this runner directly at declared 832x480 resolutions. Its + default text pipeline always emits 300 prompt slots, so one graph covers + different raw prompt lengths without padding cross-attention to 512 slots. - Dual-stream diffusion models: `use_dual_stream = True` in models such as `hunyuan3d.py` is an existing overlap family. - Workflow rule: if a hotspot is communication-heavy, rule out these in-repo overlap families before proposing a brand new overlap design. 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 b023f21b2..7f130b88a 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 @@ -37,7 +37,7 @@ These options are intended to preserve output quality. In practice, some paths ( |---|---|---|---|---| | **Performance Mode** | `--performance-mode auto\|speed\|memory\|manual` (`--mode` alias) | Applies model-aware residency, FSDP/CFG, and compile defaults without overriding explicit flags. `auto` is the safe default; `speed` favors GPU residency; `memory` favors offload; `manual` leaves performance args explicit. | Fastest way to establish a sensible deployment baseline | `speed` may OOM and enables `torch.compile` only when the model deployment config allows it. Explicit offload/FSDP/parallelism/compile flags win. Use `manual` for controlled A/B benchmarks. | | **torch.compile** | `--enable-torch-compile` | Applies `torch.compile` to the DiT forward pass. Treat it as a measured comparator, not an assumed upgrade. | Model- and shape-dependent; recent B300 coverage found eager or valid BCG faster or within 1% for every valid compile control | First request is slow and some models time out or drift numerically. Keep eager as the ground truth, use a warmup watchdog, and validate the target model. See the [H200/B300 survey](https://github.com/BBuf/how-to-optim-algorithm-in-cuda/issues/21). | -| **Breakable CUDA Graph** | `--enable-breakable-cuda-graph` plus optional `--warmup-resolutions ` and `--bcg-text-buckets ...` | Captures fixed-resolution DiT segments while leaving attention/collectives eager, reducing launch overhead on supported pipelines. | Large on launch-bound paths; merged SANA and LTX-2 cases show material e2e gains | Mutually exclusive with `torch.compile` and Cache-DiT; BCG takes priority. The model's default resolution is captured automatically; declare every additional production resolution. Current support is model-specific (Ideogram4, LTX-2/2.3, LongCat-Image, MiniMax-H3, Qwen-Image, SANA1.5, Z-Image, GLM-Image); benchmark before keeping it. | +| **Breakable CUDA Graph** | `--enable-breakable-cuda-graph` plus optional `--warmup-resolutions ` and `--bcg-text-buckets ...` | Captures fixed-resolution DiT segments while leaving attention/collectives eager, reducing launch overhead on supported pipelines. | Large on launch-bound paths; merged SANA and LTX-2 cases show material e2e gains | Mutually exclusive with `torch.compile` and Cache-DiT; BCG takes priority. The model's default resolution is captured automatically; declare every additional production resolution. Current support is model-specific (Ideogram4, LTX-2/2.3, LongCat-Image, MiniMax-H3, Qwen-Image, SANA1.5, SANA-Video, Z-Image, GLM-Image); benchmark before keeping it. | | **Warmup** | `--warmup-mode request` | Runs dummy forward passes to warm up CUDA caches, JIT, and `torch.compile`. Eliminates cold-start penalty. | Removes first-request latency spike | Adds startup time. Without `--warmup-resolutions`, warmup happens on first request. | | **Warmup Resolutions** | `--warmup-resolutions 256x256 720x720` | Pre-compiles and warms up specific resolutions at server startup (instead of lazily on first request). | Faster first request per resolution | Each resolution adds to startup time. Serving mode only; useful when you know your target resolutions in advance. | | **Multi-GPU (SP)** | `--num-gpus N --ulysses-degree N` | Sequence parallelism across GPUs. Shards sequence tokens (not frames) to minimize padding. | Near-linear scaling with N GPUs | Requires NCCL; inter-GPU bandwidth matters. `ulysses_degree * ring_degree = sp_degree`. For Wan2.2 video, start by benchmarking pure Ulysses before assuming a mixed Ulysses/Ring layout is fastest. | @@ -371,7 +371,7 @@ Use these as first commands to benchmark, not as universal winners. | 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. | | LongCat-Image | 1024x1024, 50 steps, guidance 4.5, 1 GPU | `--performance-mode manual --enable-prompt-rewrite false` for a DiT-only eager baseline; compare `--enable-breakable-cuda-graph --warmup-resolutions 1024x1024 --enable-torch-compile false` for fixed-resolution serving | Prompt rewriting is enabled by the model defaults and runs a Qwen2.5-VL component. Disable it for kernel A/B, then keep a separate end-to-end recipe with rewriting enabled. LongCat always sends a 512-token prompt body to the DiT, so BCG reuses one signature across prompt lengths without a custom text bucket. | -| SANA-Video | 832x480, 17 frames, 8 steps for CI-sized profiling; 81 frames, 50 steps for release quality | `--performance-mode manual` and eager first | Self QKV and cross KV are already packed. Check SANA's shared bit-exact conv/modulation fast paths and one-time contiguous layout before adding a new kernel. | +| SANA-Video | 832x480, 17 frames, 8 steps for CI-sized profiling; 81 frames, 50 steps for release quality | `--performance-mode manual` and eager first; compare `--enable-breakable-cuda-graph --warmup-resolutions 832x480 --enable-torch-compile false` for fixed-resolution serving | Self QKV and cross KV are already packed. The default 300-token prompt shape reuses one BCG signature without a custom text bucket. Check SANA's shared bit-exact conv/modulation fast paths and one-time contiguous layout before adding a new kernel. | | LTX-2 / LTX-2.3 | 768x512 or HQ 1920x1088, 121 frames | `--pipeline-class-name LTX2TwoStagePipeline --enable-torch-compile --warmup-mode request`; HQ uses `LTX2TwoStageHQPipeline` | Use benchmark/profile presets for nightly alignment, one-stage, high-resolution stress, and HQ. Device mode choices are `original` and `resident`; `resident` is fastest but uses more VRAM. `snapshot` is a deprecated alias for `original`, so do not use it in new commands. | | LTX-2.5 | One-stage distilled: 960x544, 121 frames, 8 steps; two-stage: 1920x1088 | `--pipeline-class-name LTX2Pipeline --performance-mode manual`; add `--use-diffusion-decoder` only for the decoder A/B | Benchmark the DiT and optional diffusion decoder as separate stages. Confirm NATTEN `na3d` is active before comparing decoder latency; a FlexAttention fallback is a different backend. Distilled weights run unguided. | | HunyuanVideo | 848x480 or 720p class video | `--text-encoder-cpu-offload --pin-cpu-memory --enable-torch-compile --warmup-mode request` | Check VAE decode separately. GroupNorm+SiLU is default-eligible in mainline when wrapper guards pass; use `bench_group_norm_silu.py` when VAE residual blocks are hot. | diff --git a/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/sana_video.py b/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/sana_video.py new file mode 100644 index 000000000..5b5894842 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/sana_video.py @@ -0,0 +1,45 @@ +# Copyright 2023-2026 SGLang Team +# Licensed under the Apache License, Version 2.0 +# ============================================================================== +"""SANA-Video breakable CUDA graph (BCG) prompt handling.""" + +from __future__ import annotations + +from typing import Any + +import torch + +from sglang.multimodal_gen.runtime.breakable_cuda_graph import ( + prompt_padding as bcg_utils, +) + +_DEFAULT_PROMPT_LENGTH = 300 + + +def is_sana_video_transformer(current_model: Any, call_kwargs: dict) -> bool: + encoder_hidden_states = bcg_utils.first_tensor( + call_kwargs.get("encoder_hidden_states") + ) + return ( + bcg_utils.transformer_class_name_matches(current_model, "sanavideo") + and torch.is_tensor(encoder_hidden_states) + and encoder_hidden_states.dim() >= 2 + and encoder_hidden_states.shape[1] == _DEFAULT_PROMPT_LENGTH + ) + + +def keep_sana_video_prompt_shape( + call_kwargs: dict, current_model: Any, buckets: tuple[int, ...] +) -> dict: + """Keep the pipeline's fixed 300-token prompt shape unchanged. + + The SANA-Video text stage pads both CFG branches to 300 tokens by default. + Expanding that fixed shape to the generic 512/1024 buckets only increases + cross-attention work and captures unused graph signatures. + """ + return call_kwargs + + +bcg_utils.register_prompt_padder( + is_sana_video_transformer, keep_sana_video_prompt_shape +) diff --git a/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/prompt_padding.py b/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/prompt_padding.py index c83140dce..b9d510553 100644 --- a/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/prompt_padding.py +++ b/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/prompt_padding.py @@ -304,5 +304,6 @@ def _ensure_model_padders_registered() -> None: longcat_image, minimax_h3, qwen_image, + sana_video, zimage, ) diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index 451062f19..d444537a1 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -146,7 +146,9 @@ BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS = frozenset( { "comfy-org/ideogram-4", "efficient-large-model/sana1.5_1.6b_1024px_diffusers", + "efficient-large-model/sana-video_2b_480p_diffusers", "sana1.5_1.6b_1024px_diffusers", + "sana-video_2b_480p_diffusers", "fal/ideogram-v4-fast", "fal/ideogram-v4-instant", "glm-image", @@ -186,6 +188,7 @@ BREAKABLE_CUDA_GRAPH_SUPPORTED_PIPELINE_CONFIGS = frozenset( "MiniMaxH3PipelineConfig", "QwenImagePipelineConfig", "SanaPipelineConfig", + "SanaVideoPipelineConfig", "ZImagePipelineConfig", } ) @@ -646,7 +649,7 @@ class ServerArgs(DisaggServerArgsMixin): logger.warning( "[Diffusion BCG] disabled for %s: only Ideogram-4, " "Lightricks/LTX-2, LongCat-Image, MiniMax-H3, " - "Qwen/Qwen-Image, Qwen/Qwen-Image-2512, SANA1.5, " + "Qwen/Qwen-Image, Qwen/Qwen-Image-2512, SANA1.5, SANA-Video, " "Tongyi-MAI/Z-Image/Z-Image-Turbo, and zai-org/GLM-Image are " "currently supported.", pipeline_config_name, diff --git a/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py b/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py index 7f2f39938..8dfb7e69b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py +++ b/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py @@ -54,6 +54,10 @@ class ZImageTransformer2DModel(torch.nn.Module): return torch.zeros(pos_ids.shape[0], 8, device=pos_ids.device) +class SanaVideoTransformer3DModel(torch.nn.Module): + pass + + def _fake_cache_dit_batch(*, is_warmup: bool) -> SimpleNamespace: return SimpleNamespace( is_warmup=is_warmup, @@ -69,6 +73,7 @@ class TestDiffusionBCGPadding(unittest.TestCase): self.longcat_model = LongCatImageTransformer2DModel() self.minimax_h3_model = MiniMaxH3DiTModel() self.zimage_model = ZImageTransformer2DModel() + self.sana_video_model = SanaVideoTransformer3DModel() self.other_model = OtherTransformer2DModel() def _patch_buckets(self, *buckets: int): @@ -436,6 +441,28 @@ class TestDiffusionBCGPadding(unittest.TestCase): "LTX23PipelineConfig", BREAKABLE_CUDA_GRAPH_SUPPORTED_PIPELINE_CONFIGS ) + def test_sana_video_keeps_its_fixed_prompt_shape(self): + kwargs = { + "encoder_hidden_states": torch.zeros(1, 300, 2304), + "encoder_attention_mask": torch.ones(1, 300, dtype=torch.long), + } + with self._patch_buckets(64, 128, 256, 512, 1024): + out = self.stage._bcg_pad_prompt_kwargs( + kwargs, current_model=self.sana_video_model + ) + + self.assertIs(out, kwargs) + self.assertEqual(out["encoder_hidden_states"].shape, (1, 300, 2304)) + self.assertEqual(out["encoder_attention_mask"].shape, (1, 300)) + self.assertIn( + "efficient-large-model/sana-video_2b_480p_diffusers", + BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS, + ) + self.assertIn( + "SanaVideoPipelineConfig", + BREAKABLE_CUDA_GRAPH_SUPPORTED_PIPELINE_CONFIGS, + ) + def test_dynamic_varlen_mask_meta_rebuilds_once_per_replay_token(self): builder = DynamicVarlenMaskMeta() mask = torch.tensor([[True, True, False, False]])