[NPU][Diffusion] fix sp modulate for qwen-image-edit (#20974)

Co-authored-by: 高鑫 <gaoxin@gaoxindeMacBook-Pro.local>

python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py

@@ -47,12 +47,16 @@ from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
     apply_flashinfer_rope_qk_inplace,
 )
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
-from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.platforms import (
+    AttentionBackendEnum,
+    current_platform,
+)
 from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)  # pylint: disable=invalid-name
 
+
 try:
     from nunchaku.models.attention import NunchakuFeedForward  # type: ignore[import]
 except Exception:
@@ -826,6 +830,9 @@ class QwenImageTransformerBlock(nn.Module):
 
         shift, scale, gate = mod_params.chunk(3, dim=-1)
         if index is not None:
+            # ROCm currently fails to compile the select01 Triton kernel, so
+            # keep using the torch.where fallback there.
+            if x.is_cuda and not current_platform.is_hip():
                 actual_batch = x.shape[0]
                 shift0, shift1 = (
                     shift[:actual_batch],
@@ -835,7 +842,10 @@ class QwenImageTransformerBlock(nn.Module):
                     scale[:actual_batch],
                     scale[actual_batch : 2 * actual_batch],
                 )
-            gate0, gate1 = gate[:actual_batch], gate[actual_batch : 2 * actual_batch]
+                gate0, gate1 = (
+                    gate[:actual_batch],
+                    gate[actual_batch : 2 * actual_batch],
+                )
                 if not x.is_contiguous():
                     x = x.contiguous()
                 if not index.is_contiguous():
@@ -878,6 +888,28 @@ class QwenImageTransformerBlock(nn.Module):
                         eps=norm_module.eps,
                     )
                     return x, gate_result
+            else:
+                actual_batch = x.shape[0]
+                shift0, shift1 = (
+                    shift[:actual_batch],
+                    shift[actual_batch : 2 * actual_batch],
+                )
+                scale0, scale1 = (
+                    scale[:actual_batch],
+                    scale[actual_batch : 2 * actual_batch],
+                )
+                gate0, gate1 = (
+                    gate[:actual_batch],
+                    gate[actual_batch : 2 * actual_batch],
+                )
+                index = index.to(dtype=torch.bool).unsqueeze(-1)
+                shift_result = torch.where(
+                    index, shift1.unsqueeze(1), shift0.unsqueeze(1)
+                )
+                scale_result = torch.where(
+                    index, scale1.unsqueeze(1), scale0.unsqueeze(1)
+                )
+                gate_result = torch.where(index, gate1.unsqueeze(1), gate0.unsqueeze(1))
         else:
             shift_result = shift.unsqueeze(1)
             scale_result = scale.unsqueeze(1)
@@ -1127,8 +1159,8 @@ class QwenImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
             first_size = sample[0][0] * sample[0][1] * sample[0][2]
             total_size = sum(s[0] * s[1] * s[2] for s in sample)
             if sp_world_size > 1:
-                first_local_size = _local_seq_len(first_size)
-                tail_local_size = _local_seq_len(total_size - first_size)
+                first_local_size = _local_seq_len(first_size, sp_world_size)
+                tail_local_size = _local_seq_len(total_size - first_size, sp_world_size)
                 idx = torch.cat(
                     [
                         torch.zeros(first_local_size, device=device, dtype=torch.int),