[diffusion] Enable breakable CUDA graph for SANA (H200 1024px e2e -26%, bit-exact) (#33989)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Fable 5
parent
12de7fb1f6
commit
6c7498113f
@@ -136,7 +136,6 @@ from sglang.multimodal_gen.runtime.utils.torch_compile import (
|
|||||||
maybe_enable_inductor_compute_comm_overlap,
|
maybe_enable_inductor_compute_comm_overlap,
|
||||||
resolve_torch_compile_mode,
|
resolve_torch_compile_mode,
|
||||||
)
|
)
|
||||||
from sglang.multimodal_gen.utils import dict_to_3d_list
|
|
||||||
|
|
||||||
logger = init_logger(__name__)
|
logger = init_logger(__name__)
|
||||||
|
|
||||||
@@ -978,7 +977,6 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
|
|||||||
{
|
{
|
||||||
# TODO: make sure on-device
|
# TODO: make sure on-device
|
||||||
"encoder_hidden_states_image": image_embeds,
|
"encoder_hidden_states_image": image_embeds,
|
||||||
"mask_strategy": dict_to_3d_list(None, t_max=50, l_max=60, h_max=24),
|
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -20,7 +20,6 @@ from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler
|
|||||||
from sglang.multimodal_gen.runtime.utils.precision import (
|
from sglang.multimodal_gen.runtime.utils.precision import (
|
||||||
autocast_context as precision_autocast_context,
|
autocast_context as precision_autocast_context,
|
||||||
)
|
)
|
||||||
from sglang.multimodal_gen.utils import dict_to_3d_list
|
|
||||||
|
|
||||||
logger = init_logger(__name__)
|
logger = init_logger(__name__)
|
||||||
|
|
||||||
@@ -91,7 +90,6 @@ class DmdDenoisingStage(DenoisingStage):
|
|||||||
self.transformer.forward,
|
self.transformer.forward,
|
||||||
{
|
{
|
||||||
"encoder_hidden_states_image": image_embeds,
|
"encoder_hidden_states_image": image_embeds,
|
||||||
"mask_strategy": dict_to_3d_list(None, t_max=50, l_max=60, h_max=24),
|
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -133,6 +133,8 @@ DEFAULT_BCG_TEXT_BUCKETS = (64, 128, 256, 512, 1024)
|
|||||||
BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS = frozenset(
|
BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS = frozenset(
|
||||||
{
|
{
|
||||||
"comfy-org/ideogram-4",
|
"comfy-org/ideogram-4",
|
||||||
|
"efficient-large-model/sana1.5_1.6b_1024px_diffusers",
|
||||||
|
"sana1.5_1.6b_1024px_diffusers",
|
||||||
"fal/ideogram-v4-fast",
|
"fal/ideogram-v4-fast",
|
||||||
"fal/ideogram-v4-instant",
|
"fal/ideogram-v4-instant",
|
||||||
"glm-image",
|
"glm-image",
|
||||||
@@ -166,6 +168,7 @@ BREAKABLE_CUDA_GRAPH_SUPPORTED_PIPELINE_CONFIGS = frozenset(
|
|||||||
"LTX2PipelineConfig",
|
"LTX2PipelineConfig",
|
||||||
"MiniMaxH3PipelineConfig",
|
"MiniMaxH3PipelineConfig",
|
||||||
"QwenImagePipelineConfig",
|
"QwenImagePipelineConfig",
|
||||||
|
"SanaPipelineConfig",
|
||||||
"ZImagePipelineConfig",
|
"ZImagePipelineConfig",
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
@@ -560,7 +563,7 @@ class ServerArgs(DisaggServerArgsMixin):
|
|||||||
|
|
||||||
logger.warning(
|
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, MiniMax-H3, "
|
||||||
"Qwen/Qwen-Image, Qwen/Qwen-Image-2512, "
|
"Qwen/Qwen-Image, Qwen/Qwen-Image-2512, SANA1.5, "
|
||||||
"Tongyi-MAI/Z-Image/Z-Image-Turbo, and zai-org/GLM-Image are "
|
"Tongyi-MAI/Z-Image/Z-Image-Turbo, and zai-org/GLM-Image are "
|
||||||
"currently supported.",
|
"currently supported.",
|
||||||
pipeline_config_name,
|
pipeline_config_name,
|
||||||
|
|||||||
Reference in New Issue
Block a user