[diffusion] Enable SANA-Video breakable CUDA graphs (#35729)
This commit is contained in:
Xiaoyu Zhang
+12
-4
@@ -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.
|
||||
|
||||
|
||||
@@ -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 <WxH...>` 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 <WxH...>` 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. |
|
||||
|
||||
@@ -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
|
||||
)
|
||||
@@ -304,5 +304,6 @@ def _ensure_model_padders_registered() -> None:
|
||||
longcat_image,
|
||||
minimax_h3,
|
||||
qwen_image,
|
||||
sana_video,
|
||||
zimage,
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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]])
|
||||
|
||||
Reference in New Issue
Block a user