align platform build config with working submission

qwen3_6_scripts/paged_attn.py

@@ -289,7 +289,22 @@ class PagedAttention:
     ) -> torch.Tensor:
         # NOTE: The Triton context_attention_fwd kernel hangs on Iluvatar
         # BI-V100 hardware (same class of issue as cudnnFlashAttnForward).
-        # Use a pure-PyTorch fallback that reads the paged KV cache directly.
+        # 单次 pad-方阵原生 flash(verified) when possible; on any anomaly or
+        # when disabled via env, fall back to the pure-PyTorch path.
+        import os as _os
+        if _os.environ.get("PREFIX_FLASH", "1") != "0":
+            try:
+                return PagedAttention._forward_prefix_flash(
+                    query, key, value,
+                    key_cache, value_cache,
+                    block_tables, query_start_loc,
+                    seq_lens_tensor, context_lens,
+                )
+            except Exception as _e:
+                import sys as _sys, traceback as _tb
+                print(f"[prefix_flash FALLBACK] {type(_e).__name__}: {_e}",
+                      file=_sys.stderr, flush=True)
+                _tb.print_exc(file=_sys.stderr)
         return PagedAttention._forward_prefix_pytorch(
             query, key, value,
             key_cache, value_cache,
@@ -297,6 +312,107 @@ class PagedAttention:
             seq_lens_tensor, context_lens,
         )
 
+    @staticmethod
+    def _forward_prefix_flash(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        key_cache: torch.Tensor,
+        value_cache: torch.Tensor,
+        block_tables: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        seq_lens_tensor: torch.Tensor,
+        context_lens: torch.Tensor,
+    ) -> torch.Tensor:
+        """单次 pad-方阵原生 flash 的 prefix/chunked-prefill attention。
+
+        对每个 sequence:
+          1. gather 历史 KV(paged cache) → 连续张量,拼上当前 chunk 的 KV
+             → K/V = [ctx_len + q_len, kv_h, d]。
+          2. 把 query 前面 pad ctx_len 个零行 → [ctx_len+q_len, q_h, d](方阵)。
+          3. 单次 flash_attn_varlen_func(causal=True) 方阵 —— causal 在方阵下语义正确,
+             query 第 (ctx_len+j) 行只看到 key [0, ctx_len+j], 正是 prefix 语义。
+          4. 丢弃前 ctx_len 行(pad 产生的垃圾输出),返回后 q_len 行。
+
+        比两段式(flash+lse归merge)简单且快 1.3~3.5x —— 无需 lse 重算(ixformer 推理版
+        flash 不暴露 softmax_lse, 重算 lse 占两段式 77% 成本)。方阵 causal 已验证精确
+        (verify_single_correct.py, max_abs≤2e-3, GQA 4Q/1KV, ctx≤80K)。
+
+        显存兜底:总长 (ctx+q) 超阈值时抛异常, 由 forward_prefix 回退纯 PyTorch。
+        pad 浪费兜底:ctx_len 远大于 q_len(缓存命中态)时 pad 方阵浪费大, 抛异常回退
+        纯 PyTorch(O(q·k) 在大 ctx 反而更省)。只有冷 prefill 的 chunked(ctx≲q_len)吃单次 pad。
+        """
+        from ixformer.functions.flash_attn_lib import flash_attn_varlen_func
+
+        # pad 方阵总长上限。方阵 flash 是 O(tk^2), ctx 远大于 q_len 时 pad 浪费大;
+        # 超阈值交给纯 PyTorch(O(q·k), 大 ctx 反而相对省)。
+        _FLASH_TOTAL_LIMIT = 131072
+        # pad 浪费比例上限: ctx_len > q_len * R 时该 seq 退回纯 PyTorch。
+        # 冷 prefill 的后续 chunk: ctx/q ≤ ~4 (65K/16K); 缓存命中: ctx/q ≫ 10 → 退回。
+        _FLASH_PAD_RATIO = 4
+
+        batch_size   = seq_lens_tensor.shape[0]
+        num_q_heads  = query.shape[1]
+        num_kv_heads = key_cache.shape[1]
+        head_dim     = query.shape[2]
+        block_size   = value_cache.shape[3]
+        scale        = head_dim ** -0.5
+        orig_dtype   = query.dtype
+        dev          = query.device
+        output       = torch.empty_like(query)
+
+        for i in range(batch_size):
+            ctx_len = int(context_lens[i].item())
+            q_start = int(query_start_loc[i].item())
+            q_end   = int(query_start_loc[i + 1].item())
+            q_len   = q_end - q_start
+            total_k = ctx_len + q_len
+
+            if total_k > _FLASH_TOTAL_LIMIT:
+                raise RuntimeError(
+                    f"total_k={total_k} > flash limit {_FLASH_TOTAL_LIMIT}; "
+                    "fall back to pytorch for memory safety")
+            if ctx_len > q_len * _FLASH_PAD_RATIO:
+                raise RuntimeError(
+                    f"pad ratio too high (ctx={ctx_len} q={q_len}); "
+                    "fall back to pytorch (single-pad wasteful for cache-hit)")
+
+            q_i = query[q_start:q_end]                 # [q_len, q_h,  d]
+            k_i = key  [q_start:q_end]                 # [q_len, kv_h, d]
+            v_i = value[q_start:q_end]
+
+            if ctx_len > 0:
+                num_ctx_blocks = (ctx_len + block_size - 1) // block_size
+                if num_ctx_blocks > block_tables.shape[1]:
+                    # 与 pytorch 版一致的 undersized 保护(prefix_cache_hit bug)
+                    num_ctx_blocks = block_tables.shape[1]
+                blk_ids = block_tables[i, :num_ctx_blocks]
+                # gather K/V: paged 布局 → 连续 [ctx_len, kv_h, d]
+                k_ctx = (key_cache[blk_ids]
+                         .permute(0, 3, 1, 2, 4)
+                         .contiguous()
+                         .view(-1, num_kv_heads, head_dim))[:ctx_len]
+                v_ctx = (value_cache[blk_ids]
+                         .permute(0, 3, 1, 2)
+                         .contiguous()
+                         .view(-1, num_kv_heads, head_dim))[:ctx_len]
+                k_full = torch.cat([k_ctx, k_i], dim=0).contiguous()   # [tk, kv_h, d]
+                v_full = torch.cat([v_ctx, v_i], dim=0).contiguous()
+                # query pad 到方阵: 前 ctx_len 行零(输出丢弃)
+                q_pad = torch.zeros(total_k, num_q_heads, head_dim,
+                                    device=dev, dtype=q_i.dtype)
+                q_pad[ctx_len:] = q_i
+            else:
+                k_full, v_full, q_pad = k_i.contiguous(), v_i.contiguous(), q_i.contiguous()
+
+            cu = torch.tensor([0, total_k], dtype=torch.int32, device=dev)
+            o = flash_attn_varlen_func(
+                q_pad, k_full, v_full, cu, cu, total_k, total_k,
+                0.0, scale, causal=True)               # [tk, q_h, d]
+            output[q_start:q_end] = o[ctx_len:].to(orig_dtype)
+
+        return output
+
     @staticmethod
     def _forward_prefix_pytorch(
         query: torch.Tensor,