diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 818563ab4..e7fa4c66f 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/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) diff --git a/vllm_omni/diffusion/attention/parallel/ulysses.py b/vllm_omni/diffusion/attention/parallel/ulysses.py index 1a368d5d7..6ea2d38fd 100644 --- a/vllm_omni/diffusion/attention/parallel/ulysses.py +++ b/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)