[Bugfix] Removed the NPU-specific code path in _run_local_attention (#597)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>

vllm_omni/diffusion/attention/layer.py

@@ -17,7 +17,6 @@ from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention
 from vllm_omni.diffusion.attention.selector import get_attn_backend
 from vllm_omni.diffusion.data import get_current_omni_diffusion_config
 from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
-from vllm_omni.utils.platform_utils import is_npu
 
 logger = init_logger(__name__)
 
@@ -112,19 +111,6 @@ class Attention(nn.Module):
             )
             self.backend_pref = "sdpa"
 
-        if is_npu():
-            return self.attention(
-                query,
-                key,
-                value,
-                num_heads=query.shape[-2],
-                input_layout="BSND",
-                scale=self.softmax_scale,
-                softmax_lse_flag=True,
-                pre_tokens=65535,
-                next_tokens=65535,
-            )[0]
-
         # Fallback to standard attention
         return self.attention.forward(query, key, value, attn_metadata)
 

vllm_omni/diffusion/attention/parallel/ulysses.py

@@ -192,6 +192,8 @@ class UlyssesParallelAttention:
             # 2. Process Joint part: AllGather on Heads
             # Input: (B, JointLen, H/P, D). Output: (B, JointLen, H, D).
             # AllGather along dim 2.
+            # Ensure tensor is contiguous for all_gather (slicing may create non-contiguous views)
+            output_joint = output_joint.contiguous()
             gathered_joint = [torch.zeros_like(output_joint) for _ in range(dist.get_world_size(ctx.ulysses_pg))]
             dist.all_gather(gathered_joint, output_joint, group=ctx.ulysses_pg)
             output_joint = torch.cat(gathered_joint, dim=2)