diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 45ce3fc70..5017535f7 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -745,11 +745,23 @@ class _NormScaleShift(CustomOp): def forward_npu( self, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor ) -> torch.Tensor: - from sgl_kernel_npu.norm.scale_shift import fused_scale_shift + hidden_size = x.shape[-1] + x_numel = x.numel() - normalized = self.norm(x) - modulated = fused_scale_shift(normalized, scale, shift) - return modulated.to(x.dtype) + if scale.numel() in (1, hidden_size) and shift.numel() in ( + 1, + hidden_size, + x_numel, + ): + from sgl_kernel_npu.norm.scale_shift import fused_scale_shift + + normalized = self.norm(x) + modulated = fused_scale_shift( + normalized, scale.contiguous(), shift.contiguous() + ) + return modulated.to(x.dtype) + + return self.forward_native(x, shift, scale) class LayerNormScaleShift(_NormScaleShift): diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_helios.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_helios.py index bdb1adec6..a881d1a9b 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_helios.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_helios.py @@ -14,6 +14,8 @@ from dataclasses import dataclass import numpy as np import torch +from sglang.multimodal_gen.runtime.platforms import current_platform + @dataclass class HeliosSchedulerOutput: @@ -255,6 +257,11 @@ class HeliosScheduler: self.timesteps = torch.from_numpy(timesteps).to(device=device) self.sigmas = torch.cat([sigmas, torch.zeros(1)]).to(device=device) + if current_platform.is_npu(): + # self.sigmas is float64 (np.linspace default); Ascend aclnnExpm1 does + # not support float64 (DT_DOUBLE) and crashes the UniPC step's expm1. + # Pin fp32 on NPU; remove once aclnnExpm1 supports float64. + self.sigmas = self.sigmas.to(torch.float32) self._step_index = None self.reset_scheduler_history()