align platform build config with working submission

qwen3_6_scripts/patch_xformers_sdpa_seq.py

@@ -29,10 +29,6 @@ flash attention kernel（ixformer / cudnnFlashAttnForward）。
   关闭选项），原意是防止 profiling OOM。但 _run_sdpa_fallback 已通过 Q-tiling
   解决了该问题，chunked prefill 反而会把推理路径从 _run_sdpa_fallback 切换到
   _forward_prefix_pytorch，属于不必要的行为变更，因此一并禁用该自动逻辑。
-  同时解除 EngineArgs.create_*_config 中强制 enforce_eager=True 和
-  disable_custom_all_reduce=True 的硬编码，让 Iluvatar 环境可以实际验证
-  CUDA Graph / custom all-reduce；需要回退时仍可通过命令行显式传
-  --enforce-eager --disable-custom-all-reduce。
 
 Deploy:
   python3 modified_scripts/patch_xformers_sdpa_seq.py
@@ -95,13 +91,6 @@ _ARG_NEW_BLOCK = """\
                           # handles long-context memory without chunked prefill\
 """
 
-_ARG_FORCE_EAGER_OLD = "            enforce_eager=True,"
-_ARG_FORCE_EAGER_NEW = "            enforce_eager=self.enforce_eager,"
-_ARG_FORCE_ALLREDUCE_OLD = "            disable_custom_all_reduce=True,"
-_ARG_FORCE_ALLREDUCE_NEW = (
-    "            disable_custom_all_reduce=self.disable_custom_all_reduce,"
-)
-
 FALLBACK_METHOD = '''
     def _run_sdpa_fallback(
         self,
@@ -225,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)\
 """
 
@@ -286,26 +291,6 @@ def patch_arg_utils(path):
     else:
         print("  [warn] target block not found — check arg_utils.py version")
 
-    if _ARG_FORCE_EAGER_NEW in content:
-        print("  [skip] enforce_eager already respects CLI")
-    elif _ARG_FORCE_EAGER_OLD in content:
-        content = content.replace(_ARG_FORCE_EAGER_OLD,
-                                  _ARG_FORCE_EAGER_NEW, 1)
-        print("  [ok]   enforce_eager now respects CLI")
-        changed = True
-    else:
-        print("  [warn] enforce_eager assignment not found")
-
-    if _ARG_FORCE_ALLREDUCE_NEW in content:
-        print("  [skip] custom all-reduce already respects CLI")
-    elif _ARG_FORCE_ALLREDUCE_OLD in content:
-        content = content.replace(_ARG_FORCE_ALLREDUCE_OLD,
-                                  _ARG_FORCE_ALLREDUCE_NEW, 1)
-        print("  [ok]   custom all-reduce now respects CLI")
-        changed = True
-    else:
-        print("  [warn] custom all-reduce assignment not found")
-
     if changed:
         with open(path, "w") as f:
             f.write(content)