chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[NPU][diffusion] npu support enable_torch_compile for torchair backend on diffusion models (#20687)

Browse Source
This commit is contained in:
Yaochen Han
2026-03-19 03:40:35 +08:00
committed by GitHub
parent 39e83d1401
commit c7a71740a5
2 changed files with 45 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
View File

@@ -79,6 +79,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler
from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler
from sglang.multimodal_gen.utils import dict_to_3d_list, masks_like
from sglang.srt.utils.common import get_compiler_backend
logger = init_logger(__name__)
@@ -137,16 +138,28 @@ class DenoisingStage(PipelineStage):
module, nn.Module
):
return
try:
import torch._inductor.config as _inductor_cfg
compile_kwargs: dict[str, Any] = {"fullgraph": False, "dynamic": None}
if current_platform.is_npu():
backend = get_compiler_backend()
compile_kwargs["backend"] = backend
compile_kwargs["dynamic"] = False
logger.info("Compiling transformer with torchair backend on NPU")
else:
try:
import torch._inductor.config as _inductor_cfg
_inductor_cfg.reorder_for_compute_comm_overlap = True
except ImportError:
pass
mode = os.environ.get(
"SGLANG_TORCH_COMPILE_MODE", "max-autotune-no-cudagraphs"
)
compile_kwargs["mode"] = mode
logger.info(f"Compiling transformer with mode: {mode}")
_inductor_cfg.reorder_for_compute_comm_overlap = True
except ImportError:
pass
mode = os.environ.get("SGLANG_TORCH_COMPILE_MODE", "max-autotune-no-cudagraphs")
logger.info(f"Compiling transformer with mode: {mode}")
# TODO(triple-mu): support customized fullgraph and dynamic in the future
module.compile(mode=mode, fullgraph=False, dynamic=None)
module.compile(**compile_kwargs)
def _maybe_enable_cache_dit(
self, num_inference_steps: int | tuple[int, int], batch: Req

32
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py
View File

@@ -67,6 +67,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler
from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler
from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
from sglang.srt.utils.common import get_compiler_backend
logger = init_logger(__name__)
@@ -208,16 +209,31 @@ class MOVADenoisingStage(PipelineStage):
"""
if not server_args.enable_torch_compile or not isinstance(module, nn.Module):
return
try:
import torch._inductor.config as _inductor_cfg
compile_kwargs: dict[str, object] = {"fullgraph": False, "dynamic": None}
if current_platform.is_npu():
backend = get_compiler_backend()
compile_kwargs["backend"] = backend
compile_kwargs["dynamic"] = False
logger.info(
"Compiling %s with torchair backend on NPU",
module.__class__.__name__,
)
else:
try:
import torch._inductor.config as _inductor_cfg
_inductor_cfg.reorder_for_compute_comm_overlap = True
except ImportError:
pass
mode = os.environ.get(
"SGLANG_TORCH_COMPILE_MODE", "max-autotune-no-cudagraphs"
)
compile_kwargs["mode"] = mode
logger.info("Compiling %s with mode: %s", module.__class__.__name__, mode)
_inductor_cfg.reorder_for_compute_comm_overlap = True
except ImportError:
pass
mode = os.environ.get("SGLANG_TORCH_COMPILE_MODE", "max-autotune-no-cudagraphs")
logger.info("Compiling %s with mode: %s", module.__class__.__name__, mode)
# TODO(triple-mu): support customized fullgraph and dynamic in the future
module.compile(mode=mode, fullgraph=False, dynamic=None)
module.compile(**compile_kwargs)
def _maybe_compile_dits(self, server_args: ServerArgs):
if self._torch_compiled or not server_args.enable_torch_compile: