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

[diffusion] feat: use regional torch.compile (compile_repeated_blocks) for DiT of diffusers backend (#28193)

Browse Source
This commit is contained in:
Mick
2026-06-15 00:34:10 +08:00
committed by GitHub
parent ec36dde580
commit 3cb29f6747
1 changed files with 12 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py
+12 -4
View File
@@ -638,10 +638,18 @@ class DiffusersPipeline(ComposedPipelineBase):
if hasattr(pipe, comp): if hasattr(pipe, comp):
try: try:
component = getattr(pipe, comp) component = getattr(pipe, comp)
# TODO(DefTruth): Add support for 'compile_repeated_blocks' for 'transformer' repeated_blocks = getattr(component, "_repeated_blocks", None)
# modules which can significantly reduce compilation time for large models if (
# with repeated blocks. isinstance(component, torch.nn.Module)
if isinstance(component, torch.nn.Module) and hasattr( 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" component, "compile"
): ):
# Prefer in-place compilation if supported. According to PyTorch documentation: # Prefer in-place compilation if supported. According to PyTorch documentation: