From 09f5097fe428c2c57afcafb4a5ef817c27a72eb2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=90=D1=80=D1=82=D0=B5=D0=BC=20=D0=A1=D0=B0=D0=B2=D0=BA?= =?UTF-8?q?=D0=B8=D0=BD?= <58187114+OrangeRedeng@users.noreply.github.com> Date: Tue, 17 Mar 2026 21:23:09 +0300 Subject: [PATCH] [NPU] [Bugfix] [diffusion] Fix NZ performance bug for diffusion models (#20684) Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- python/sglang/multimodal_gen/runtime/loader/fsdp_load.py | 3 ++- python/sglang/srt/hardware_backend/npu/utils.py | 7 ++++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index 2b2630fb1..dcfbb6eac 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py @@ -22,6 +22,7 @@ from torch.distributed.fsdp import ( ) from torch.nn.modules.module import _IncompatibleKeys +from sglang.multimodal_gen.runtime.layers.linear import UnquantizedLinearMethod from sglang.multimodal_gen.runtime.loader.utils import ( get_param_names_mapping, hf_to_custom_state_dict, @@ -145,7 +146,7 @@ def maybe_load_fsdp_model( if quant_method is not None and hasattr( quant_method, "process_weights_after_loading" ): - if _is_npu: + if _is_npu and not isinstance(quant_method, UnquantizedLinearMethod): # Activate the NZ format for storing weights, # which is a specific optimization for Ascend NPU torch.npu.config.allow_internal_format = True diff --git a/python/sglang/srt/hardware_backend/npu/utils.py b/python/sglang/srt/hardware_backend/npu/utils.py index 816370111..a78de6130 100644 --- a/python/sglang/srt/hardware_backend/npu/utils.py +++ b/python/sglang/srt/hardware_backend/npu/utils.py @@ -124,7 +124,12 @@ def npu_format_cast( return tensor if tensor.device == torch.device("cpu"): - return torch.ops.npu.npu_format_cast(tensor.npu(), acl_format.value).cpu() + logger.warning_once( + "Warning: The conversion from 'ND' to 'NZ' does not work on the CPU. " + "Please disable offloading, otherwise the performance will be " + "significantly reduced." + ) + return tensor else: return torch.ops.npu.npu_format_cast(tensor, acl_format.value)