[diffusion] Enable LongCat breakable CUDA graphs (#35724)
This commit is contained in:
Xiaoyu Zhang
@@ -207,6 +207,18 @@ For supported native pipelines, set `SGLANG_CACHE_DIT_ENABLED=true` to enable Ca
|
||||
|
||||
For supported image pipelines, breakable CUDA graph can be enabled with `--enable-breakable-cuda-graph`, but you must declare every served resolution in `--warmup-resolutions` so warmup captures matching graph signatures.
|
||||
|
||||
LongCat-Image is supported by this path. Its DiT always consumes a fixed
|
||||
512-token prompt body, so different prompt lengths reuse the same graph without
|
||||
model-specific padding. For a 1024x1024 deployment, use:
|
||||
|
||||
```bash Command
|
||||
sglang serve --model-path meituan-longcat/LongCat-Image \
|
||||
--enable-breakable-cuda-graph \
|
||||
--warmup-resolutions 1024x1024 \
|
||||
--enable-torch-compile false \
|
||||
--port 30010
|
||||
```
|
||||
|
||||
### Component Residency
|
||||
|
||||
Use `--component-residency COMPONENT=MODE` to assign one runtime residency mode to each native pipeline component:
|
||||
|
||||
+11
-1
@@ -297,7 +297,12 @@ framework-specific optimization workflow.
|
||||
|
||||
**Recent Model Audit Boundaries**
|
||||
|
||||
- LongCat-Image currently has split image/text QKV projections and performs
|
||||
- LongCat-Image supports breakable CUDA graph at fixed, captured resolutions.
|
||||
Its DiT always receives a 512-token prompt body, so different raw prompt
|
||||
lengths reuse the same graph signature without padding. A model-specific
|
||||
pass-through padder prevents the generic buckets from expanding this fixed
|
||||
shape into unused graph signatures.
|
||||
The model still has split image/text QKV projections and performs
|
||||
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.
|
||||
@@ -371,6 +376,11 @@ framework-specific optimization workflow.
|
||||
every additional served resolution in `--warmup-resolutions`, and use
|
||||
`--bcg-text-buckets` for prompt signatures. Check this path before proposing
|
||||
a second graph-capture mechanism for launch-bound traces.
|
||||
- LongCat-Image uses this generic runner directly: one 1024x1024 capture covers
|
||||
short and long prompts because text conditioning is fixed at 512 tokens.
|
||||
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.
|
||||
- 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, 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, 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. |
|
||||
@@ -370,7 +370,7 @@ Use these as first commands to benchmark, not as universal winners.
|
||||
| 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. |
|
||||
| LongCat-Image | 1024x1024, 50 steps, guidance 4.5, 1 GPU | `--performance-mode manual --enable-prompt-rewrite false` for a DiT-only eager baseline | 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-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. |
|
||||
| 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. |
|
||||
|
||||
+33
@@ -0,0 +1,33 @@
|
||||
# Copyright 2023-2026 SGLang Team
|
||||
# Licensed under the Apache License, Version 2.0
|
||||
# ==============================================================================
|
||||
"""LongCat-Image breakable CUDA graph (BCG) prompt handling."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from sglang.multimodal_gen.runtime.breakable_cuda_graph import (
|
||||
prompt_padding as bcg_utils,
|
||||
)
|
||||
|
||||
|
||||
def is_longcat_image_transformer(current_model: Any, call_kwargs: dict) -> bool:
|
||||
return (
|
||||
bcg_utils.transformer_class_name_matches(current_model, "longcatimage")
|
||||
and "encoder_hidden_states" in call_kwargs
|
||||
and "txt_ids" in call_kwargs
|
||||
)
|
||||
|
||||
|
||||
def keep_longcat_prompt_shape(
|
||||
call_kwargs: dict, _current_model: Any, _buckets: tuple[int, ...]
|
||||
) -> dict:
|
||||
# LongCat always supplies the complete 512-token prompt body to its DiT.
|
||||
# Generic bucket padding would only create larger, unused graph signatures.
|
||||
return call_kwargs
|
||||
|
||||
|
||||
bcg_utils.register_prompt_padder(
|
||||
is_longcat_image_transformer, keep_longcat_prompt_shape
|
||||
)
|
||||
@@ -301,6 +301,7 @@ def _ensure_model_padders_registered() -> None:
|
||||
_model_padders_registered = True
|
||||
from sglang.multimodal_gen.runtime.breakable_cuda_graph.model_padders import ( # noqa: F401
|
||||
ideogram,
|
||||
longcat_image,
|
||||
minimax_h3,
|
||||
qwen_image,
|
||||
zimage,
|
||||
|
||||
@@ -159,6 +159,7 @@ BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS = frozenset(
|
||||
"ideogram-ai/ideogram-4-nf4",
|
||||
"lightricks/ltx-2",
|
||||
"lightricks/ltx-2.3",
|
||||
"meituan-longcat/longcat-image",
|
||||
"ltx-2",
|
||||
"ltx-2.3",
|
||||
"minimax-h3",
|
||||
@@ -181,6 +182,7 @@ BREAKABLE_CUDA_GRAPH_SUPPORTED_PIPELINE_CONFIGS = frozenset(
|
||||
"Ideogram4PipelineConfig",
|
||||
"LTX2PipelineConfig",
|
||||
"LTX23PipelineConfig",
|
||||
"LongCatImagePipelineConfig",
|
||||
"MiniMaxH3PipelineConfig",
|
||||
"QwenImagePipelineConfig",
|
||||
"SanaPipelineConfig",
|
||||
@@ -642,7 +644,8 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
return
|
||||
|
||||
logger.warning(
|
||||
"[Diffusion BCG] disabled for %s: only Ideogram-4, Lightricks/LTX-2, MiniMax-H3, "
|
||||
"[Diffusion BCG] disabled for %s: only Ideogram-4, "
|
||||
"Lightricks/LTX-2, LongCat-Image, MiniMax-H3, "
|
||||
"Qwen/Qwen-Image, Qwen/Qwen-Image-2512, SANA1.5, "
|
||||
"Tongyi-MAI/Z-Image/Z-Image-Turbo, and zai-org/GLM-Image are "
|
||||
"currently supported.",
|
||||
|
||||
@@ -45,6 +45,10 @@ class MiniMaxH3DiTModel(torch.nn.Module):
|
||||
pass
|
||||
|
||||
|
||||
class LongCatImageTransformer2DModel(torch.nn.Module):
|
||||
pass
|
||||
|
||||
|
||||
class ZImageTransformer2DModel(torch.nn.Module):
|
||||
def rotary_emb(self, pos_ids):
|
||||
return torch.zeros(pos_ids.shape[0], 8, device=pos_ids.device)
|
||||
@@ -62,6 +66,7 @@ class TestDiffusionBCGPadding(unittest.TestCase):
|
||||
self.stage = DenoisingStage.__new__(DenoisingStage)
|
||||
self.qwen_model = QwenImageTransformer2DModel()
|
||||
self.ideogram_model = Ideogram4Transformer2DModel()
|
||||
self.longcat_model = LongCatImageTransformer2DModel()
|
||||
self.minimax_h3_model = MiniMaxH3DiTModel()
|
||||
self.zimage_model = ZImageTransformer2DModel()
|
||||
self.other_model = OtherTransformer2DModel()
|
||||
@@ -133,6 +138,26 @@ class TestDiffusionBCGPadding(unittest.TestCase):
|
||||
self.assertEqual(_attn_mask_meta_local_pad({"local_pad": 7}), 7)
|
||||
self.assertEqual(_attn_mask_meta_local_pad(DynamicVarlenMaskMeta()), 0)
|
||||
|
||||
def test_longcat_keeps_its_fixed_512_token_prompt_shape(self):
|
||||
kwargs = {
|
||||
"hidden_states": torch.zeros(1, 4096, 64),
|
||||
"timestep": torch.zeros(1),
|
||||
"encoder_hidden_states": torch.zeros(1, 512, 3584),
|
||||
"encoder_attention_mask": [torch.ones(1, 512, dtype=torch.long)],
|
||||
"encoder_hidden_states_mask": [torch.ones(1, 512, dtype=torch.bool)],
|
||||
"txt_ids": torch.zeros(512, 3),
|
||||
"img_ids": torch.zeros(4096, 3),
|
||||
}
|
||||
|
||||
with self._patch_buckets(64, 128, 256, 512, 1024):
|
||||
out = self.stage._bcg_pad_prompt_kwargs(
|
||||
kwargs, current_model=self.longcat_model
|
||||
)
|
||||
|
||||
self.assertIs(out, kwargs)
|
||||
self.assertEqual(out["encoder_hidden_states"].shape, (1, 512, 3584))
|
||||
self.assertEqual(out["txt_ids"].shape, (512, 3))
|
||||
|
||||
def test_qwen_default_bucket_preserves_mask(self):
|
||||
def kwargs(valid_len: int):
|
||||
mask = torch.zeros(1, 64, dtype=torch.bool)
|
||||
@@ -380,6 +405,7 @@ class TestDiffusionBCGPadding(unittest.TestCase):
|
||||
|
||||
def test_image_generation_models_are_registered_as_bcg_supported(self):
|
||||
for model_id in (
|
||||
"meituan-longcat/longcat-image",
|
||||
"qwen/qwen-image",
|
||||
"qwen/qwen-image-2512",
|
||||
"tongyi-mai/z-image",
|
||||
@@ -390,6 +416,7 @@ class TestDiffusionBCGPadding(unittest.TestCase):
|
||||
|
||||
for config_name in (
|
||||
"GlmImagePipelineConfig",
|
||||
"LongCatImagePipelineConfig",
|
||||
"QwenImagePipelineConfig",
|
||||
"ZImagePipelineConfig",
|
||||
):
|
||||
|
||||
Reference in New Issue
Block a user