minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-09-09 14:09:21 +08:00
committed by GitHub
parent 5998e9321c
commit a27be5ff62
4 changed files with 101 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/server_args/server_args.py
+4 -1
View File
@@ -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
+25
View File
@@ -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
+47
View File
@@ -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