[Diffusion] Enable lossless BCG for FLUX.1-dev (#38591)

docs/cookbook/diffusion/FLUX/FLUX.mdx

@@ -49,6 +49,31 @@ See [Performance Optimization](/docs/sglang-diffusion/performance-optimization)
 - `--ulysses-degree`: The degree of DeepSpeed-Ulysses-style SP in USP
 - `--ring-degree`: The degree of ring attention-style SP in USP
 
+### 3.3 Breakable CUDA graph for FLUX.1-dev
+
+For repeated `quality=lossless` requests, FLUX.1-dev supports breakable CUDA graph (BCG) execution with the native backend. BCG captures DiT segments during startup and replays them for the warmed resolutions. It can reduce recurring host launch overhead without enabling `torch.compile`.
+
+The following configuration was validated on two NVIDIA H200 GPUs with BF16 weights, PyTorch 2.13, and CUDA 13.0. Set `HF_TOKEN` to a Hugging Face token with access to the checkpoint before downloading it.
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1 sglang generate \
+  --model-path black-forest-labs/FLUX.1-dev \
+  --backend sglang \
+  --num-gpus 2 --tp-size 2 \
+  --component-residency dit=resident \
+  --enable-torch-compile false \
+  --enable-breakable-cuda-graph \
+  --warmup-resolutions 1024x1024 \
+  --quality lossless \
+  --width 1024 --height 1024 \
+  --num-inference-steps 50 --guidance-scale 3.5 --seed 42 \
+  --prompt "A futuristic cyberpunk city at night, neon lights reflecting on wet streets"
+```
+
+Confirm `[Diffusion BCG] captured` in the log, then compare warmed request latency with eager execution on the same GPUs. Capture time and graph memory are additional startup costs. Add other served resolutions to `--warmup-resolutions`; a request with an uncaptured signature falls back to eager.
+
+FLUX.1-dev uses a fixed 512-token T5 conditioning sequence, so changing `--bcg-text-buckets` does not create additional prompt-length graphs. Its request-gated DiT fusions at `quality=high` and `extra-high` cannot be combined with BCG: the runtime rejects those requests because the captured graph uses the lossless branches. FLUX.2 and quantized transformer overrides require separate validation.
+
 ## 4. API Usage
 
 For complete API documentation, please refer to the [official API usage guide](/docs/sglang-diffusion/api/openai_api).

python/sglang/multimodal_gen/runtime/server_args/server_args.py

@@ -166,6 +166,7 @@ DEFAULT_BCG_TEXT_BUCKETS = (64, 128, 256, 512, 1024)
 
 BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS = frozenset(
     {
+        "black-forest-labs/flux.1-dev",
         "comfy-org/ideogram-4",
         "efficient-large-model/sana1.5_1.6b_1024px_diffusers",
         "efficient-large-model/sana-video_2b_480p_diffusers",
@@ -173,6 +174,7 @@ BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS = frozenset(
         "sana-video_2b_480p_diffusers",
         "fal/ideogram-v4-fast",
         "fal/ideogram-v4-instant",
+        "flux.1-dev",
         "glm-image",
         "ideogram-4",
         "ideogram-4-fp8",
@@ -204,6 +206,7 @@ BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS = frozenset(
 
 BREAKABLE_CUDA_GRAPH_SUPPORTED_PIPELINE_CONFIGS = frozenset(
     {
+        "FluxPipelineConfig",
         "GlmImagePipelineConfig",
         "Ideogram4PipelineConfig",
         "JoyEchoPipelineConfig",
@@ -742,7 +745,7 @@ class ServerArgs(DisaggServerArgsMixin):
             return
 
         logger.warning(
-            "[Diffusion BCG] disabled for %s: only Ideogram-4, "
+            "[Diffusion BCG] disabled for %s: only FLUX.1-dev, Ideogram-4, "
             "jdopensource/JoyAI-Echo, Lightricks/LTX-2, LongCat-Image, "
             "MiniMax-H3, Qwen/Qwen-Image, Qwen/Qwen-Image-2512, SANA1.5, "
             "SANA-Video, Tongyi-MAI/Z-Image/Z-Image-Turbo, and "

python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py

@@ -37,6 +37,10 @@ class OtherTransformer2DModel(torch.nn.Module):
     pass
 
 
+class FluxTransformer2DModel(torch.nn.Module):
+    pass
+
+
 class Ideogram4Transformer2DModel(torch.nn.Module):
     pass
 
@@ -132,6 +136,7 @@ class TestDiffusionBCGPadding(unittest.TestCase):
         self.zimage_model = ZImageTransformer2DModel()
         self.sana_video_model = SanaVideoTransformer3DModel()
         self.other_model = OtherTransformer2DModel()
+        self.flux_model = FluxTransformer2DModel()
 
     def _patch_buckets(self, *buckets: int):
         resolved = tuple(sorted({b for b in buckets if b > 0}))
@@ -520,6 +525,26 @@ class TestDiffusionBCGPadding(unittest.TestCase):
             BREAKABLE_CUDA_GRAPH_SUPPORTED_PIPELINE_CONFIGS,
         )
 
+    def test_flux_unmasked_conditioning_keeps_one_signature_for_all_buckets(self):
+        # FLUX attends to all 512 T5 tokens, including tokenizer padding.
+        # Adding unmasked tokens would change attention and generated pixels.
+        kwargs = {
+            "hidden_states": torch.zeros(1, 4096, 64, dtype=torch.bfloat16),
+            "encoder_hidden_states": torch.ones(1, 512, 4096, dtype=torch.bfloat16),
+            "pooled_projections": torch.ones(1, 768, dtype=torch.bfloat16),
+            "timestep": torch.zeros(1),
+            "guidance": torch.full((1,), 3.5, dtype=torch.bfloat16),
+            "freqs_cis": (torch.zeros(4608, 64), torch.ones(4608, 64)),
+        }
+        signature = _signature_kwargs(kwargs)
+        for bucket in (64, 128, 256, 512, 1024):
+            with self.subTest(bucket=bucket):
+                out = self.stage._bcg_pad_prompt_kwargs(
+                    kwargs, current_model=self.flux_model, force_bucket=bucket
+                )
+                self.assertIs(out, kwargs)
+                self.assertEqual(_signature_kwargs(out), signature)
+
     def test_dynamic_varlen_mask_meta_rebuilds_once_per_replay_token(self):
         builder = DynamicVarlenMaskMeta()
         mask = torch.tensor([[True, True, False, False]])

python/sglang/multimodal_gen/test/unit/test_server_args.py

@@ -17,6 +17,10 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
     PipelineConfig,
 )
 from sglang.multimodal_gen.configs.pipeline_configs.cosmos3 import Cosmos3Config
+from sglang.multimodal_gen.configs.pipeline_configs.flux import (
+    Flux2PipelineConfig,
+    FluxPipelineConfig,
+)
 from sglang.multimodal_gen.configs.pipeline_configs.helios import (
     HeliosDistilledConfig,
 )
@@ -60,6 +64,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.wan import (
     WanT2V720PConfig,
 )
 from sglang.multimodal_gen.configs.pipeline_configs.zimage import ZImagePipelineConfig
+from sglang.multimodal_gen.configs.sample.flux import FluxSamplingParams
 from sglang.multimodal_gen.registry import (
     _get_config_info,
     get_non_diffusers_pipeline_name,
@@ -820,6 +825,48 @@ class TestWarmupModeNormalization(unittest.TestCase):
         sa.bcg_text_buckets = None
         sa._validate_breakable_cuda_graph()  # must not raise
 
+    def test_flux_bcg_resolves_hub_and_local_checkpoint_warmup(self):
+        for model_path, model_id in (
+            ("black-forest-labs/FLUX.1-dev", None),
+            ("/models/FLUX.1-dev", None),
+            ("/cache/models--black-forest-labs--FLUX.1-dev/snapshots/revision", None),
+            ("/models/pinned-checkpoint", "black-forest-labs/FLUX.1-dev"),
+        ):
+            with self.subTest(model_path=model_path, model_id=model_id):
+                sa = ServerArgs.__new__(ServerArgs)
+                sa.model_path = model_path
+                sa.model_id = model_id
+                sa.pipeline_class_name = "FluxPipeline"
+                sa.pipeline_config = FluxPipelineConfig()
+                sa.enable_breakable_cuda_graph = True
+                # Resolve native sampling defaults without loading checkpoint
+                # metadata for the synthetic local paths in this unit test.
+                with patch(
+                    "sglang.multimodal_gen.runtime.warmup_request_builder."
+                    "get_model_sampling_defaults",
+                    return_value=FluxSamplingParams(),
+                ):
+                    sa._adjust_breakable_cuda_graph_support()
+                sa._adjust_warmup()
+
+                self.assertTrue(sa.enable_breakable_cuda_graph)
+                self.assertEqual(sa.warmup_resolutions, ["1024x1024"])
+                self.assertEqual(sa.warmup_mode, "server")
+
+    def test_flux_bcg_requires_both_supported_checkpoint_and_pipeline(self):
+        for model_path, config in (
+            ("black-forest-labs/FLUX.2-dev", Flux2PipelineConfig()),
+            ("black-forest-labs/FLUX.1-schnell", FluxPipelineConfig()),
+            ("black-forest-labs/FLUX.1-dev", Flux2PipelineConfig()),
+        ):
+            with self.subTest(model_path=model_path, config=type(config).__name__):
+                sa = ServerArgs.__new__(ServerArgs)
+                sa.model_path = model_path
+                sa.pipeline_config = config
+                sa.enable_breakable_cuda_graph = True
+                sa._adjust_breakable_cuda_graph_support()
+                self.assertFalse(sa.enable_breakable_cuda_graph)
+
     def test_disagg_role_disables_server_warmup(self):
         from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType