From 3cb29f67479206fe551c4b3dfaa0a184f627adac Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 15 Jun 2026 00:34:10 +0800 Subject: [PATCH] [diffusion] feat: use regional torch.compile (compile_repeated_blocks) for DiT of diffusers backend (#28193) --- .../runtime/pipelines/diffusers_pipeline.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py index 3da874a64..a2b623a5b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py @@ -638,10 +638,18 @@ class DiffusersPipeline(ComposedPipelineBase): if hasattr(pipe, comp): try: component = getattr(pipe, comp) - # TODO(DefTruth): Add support for 'compile_repeated_blocks' for 'transformer' - # modules which can significantly reduce compilation time for large models - # with repeated blocks. - if isinstance(component, torch.nn.Module) and hasattr( + repeated_blocks = getattr(component, "_repeated_blocks", None) + if ( + isinstance(component, torch.nn.Module) + and repeated_blocks + and hasattr(component, "compile_repeated_blocks") + ): + # Regional compilation: compile a single instance of each + # repeated transformer block and let inductor's cache reuse + # it for all repeats, instead of compiling the whole DiT as + # one graph + component.compile_repeated_blocks() + elif isinstance(component, torch.nn.Module) and hasattr( component, "compile" ): # Prefer in-place compilation if supported. According to PyTorch documentation: