[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>
This commit is contained in:
Артем Савкин
2026-03-17 21:23:09 +03:00
committed by GitHub
parent d35fea1b2b
commit 09f5097fe4
2 changed files with 8 additions and 2 deletions

View File

@@ -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

View File

@@ -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)