diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index 7be087b22..3ad5428ee 100644 --- a/docs/docs/sglang-diffusion/api/cli.mdx +++ b/docs/docs/sglang-diffusion/api/cli.mdx @@ -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: 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 64077e801..e9ca2c6ca 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 @@ -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. 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 1d1e21dd4..b023f21b2 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, 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, 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. | diff --git a/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/longcat_image.py b/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/longcat_image.py new file mode 100644 index 000000000..e97255990 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/longcat_image.py @@ -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 +) 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 d0ea0d442..c83140dce 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 @@ -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, 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 28dbcf149..b6709d27c 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -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.", 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 689fbecf9..7f2f39938 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 @@ -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", ):