[Diffusion] [NPU] enable Helios on npu (#29011)

Co-authored-by: ronnie_zheng <zl19940307@163.com>

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

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