fix perforence

qwen3_6_scripts/patch_xformers_sdpa_seq.py

@@ -214,18 +214,34 @@ NEW_XFORMER_BLOCK = """\
             query = query.unsqueeze(0)
             key = key.unsqueeze(0)
             value = value.unsqueeze(0)
-            if self.head_size > 128:
+            if self.head_size == 256:
+                # head_dim=256: ixformer flash_attn supports it natively
+                # Use ixinfer_flash_attn_unpad directly (bypasses FwOp check)
+                from ixformer.functions.flash_attn_lib import ixinfer_flash_attn_unpad as _ixf_unpad
+                total_q = query.shape[0] * query.shape[1]
+                total_k = key.shape[0] * key.shape[1]
+                q_flat = query.reshape(total_q, self.num_heads, self.head_size)
+                k_flat = key.reshape(total_k, self.num_kv_heads, self.head_size)
+                v_flat = value.reshape(total_k, self.num_kv_heads, self.head_size)
+                if isinstance(attn_bias[0], BlockDiagonalCausalMask):
+                    cu_seqlens_q = attn_bias[0].q_seqinfo.seqstart.int().to(query.device)
+                    cu_seqlens_k = attn_bias[0].k_seqinfo.seqstart.int().to(query.device)
+                    max_seqlen_q = attn_bias[0].q_seqinfo.max_seqlen
+                    max_seqlen_k = attn_bias[0].k_seqinfo.max_seqlen
+                else:
+                    batch_size = query.shape[0]
+                    seqlen = query.shape[1]
+                    cu_seqlens_q = torch.arange(0, batch_size + 1, device=query.device, dtype=torch.int32) * seqlen
+                    cu_seqlens_k = cu_seqlens_q
+                    max_seqlen_q = seqlen
+                    max_seqlen_k = seqlen
+                out = _ixf_unpad(q_flat, k_flat, v_flat,
+                                 cu_seqlens_q, cu_seqlens_k,
+                                 max_seqlen_q, max_seqlen_k,
+                                 True, self.scale, out=None)
+                out = out.view(query.shape[0], query.shape[1], self.num_heads, self.head_size)
+            elif self.head_size not in (32, 64, 128, 256):
                 out = self._run_sdpa_fallback(query, key, value, attn_metadata)
-            else:
-                out = xops.memory_efficient_attention_forward(
-                    query,
-                    key,
-                    value,
-                    attn_bias=attn_bias[0],
-                    p=0.0,
-                    scale=self.scale,
-                    op=self.attn_op,
-                )
             return out.view_as(original_query)\
 """