基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/meta.json

@@ -1,6 +1,6 @@
 {
-  "id": 197,
-  "displayId": 20005,
+  "id": 201,
+  "displayId": 10008,
   "type": "Traditional",
   "isPublic": false,
   "locales": [
@@ -13,4 +13,4 @@
     }
   ],
   "problemTagIds": []
-}
+}

基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/spj.py

@@ -0,0 +1,207 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+SPJ: problem_10008 — Paged KV Cache (flash-attn)
+"""
+
+import base64
+import hashlib
+import json
+import math
+import sys
+from pathlib import Path
+from typing import Callable, List, Optional, Tuple
+
+
+CHAL_PREFIX = "OJCHAL v1 "
+RESULT_PREFIX = "OJRESULT v1 "
+
+
+def eprint(msg: str) -> None:
+    print(msg, file=sys.stderr, flush=True)
+
+
+def read_text(path: Path) -> str:
+    try:
+        return path.read_text(encoding="utf-8", errors="replace")
+    except Exception as ex:
+        raise RuntimeError(f"Failed to read {path}: {ex}") from ex
+
+
+def parse_nonce_from_text(text: str) -> Optional[bytes]:
+    for line in text.splitlines():
+        line = line.strip()
+        if line.startswith(CHAL_PREFIX):
+            parts = line.split()
+            if len(parts) != 3:
+                return None
+            try:
+                return base64.b64decode(parts[2], validate=True)
+            except Exception:
+                return None
+    return None
+
+
+def parse_results_from_text(text: str) -> List[Tuple[str, str]]:
+    out = []
+    for line in text.splitlines():
+        line = line.strip()
+        if line.startswith(RESULT_PREFIX):
+            parts = line.split()
+            if len(parts) < 4:
+                continue
+            sig_hex = parts[2]
+            payload_b64 = parts[3]
+            out.append((sig_hex, payload_b64))
+    return out
+
+
+def verify_and_decode_payload(nonce: bytes, sig_hex: str, payload_b64: str) -> Optional[dict]:
+    try:
+        payload = base64.b64decode(payload_b64, validate=True)
+    except Exception:
+        return None
+    sig_calc = hashlib.sha256(nonce + payload).hexdigest()
+    if sig_calc.lower() != sig_hex.lower():
+        return None
+    try:
+        obj = json.loads(payload.decode("utf-8"))
+    except Exception:
+        return None
+    if not isinstance(obj, dict):
+        return None
+    return obj
+
+
+def _read_testcase_id() -> Optional[int]:
+    try:
+        text = Path("input").read_text(encoding="utf-8").strip()
+        return int(text.split()[0])
+    except Exception:
+        return None
+
+
+def _format_diagnostic(payload: dict, describe_fn: Optional[Callable[[int], str]] = None) -> List[str]:
+    tk = payload.get("tk_time_ms")
+    tb = payload.get("tb_time_ms")
+    th = payload.get("th_time_ms")
+    score = payload.get("score_ratio")
+    passed = payload.get("pass", False)
+    lines: List[str] = []
+    lines.append("-" * 64)
+    tc_id = _read_testcase_id()
+    if tc_id is not None:
+        lines.append(f"  Testcase  #{tc_id}")
+        if describe_fn:
+            desc = describe_fn(tc_id)
+            if desc:
+                lines.append(f"    Config:  {desc}")
+    else:
+        lines.append("  Testcase  <?>  (no 'input' file)")
+    lines.append("")
+    if tk is not None and tb is not None:
+        speedup = tb / tk if tk > 0 else float("inf")
+        lines.append(f"  Baseline:           {tb:<12.6f} ms")
+        lines.append(f"  User kernel:        {tk:<12.6f} ms")
+        if th is not None:
+            lines.append(f"  Hardware bound:     {th:<12.6f} ms")
+        lines.append(f"  Speedup vs base:    {speedup:<12.3f}x")
+    else:
+        lines.append("  (timing data unavailable)")
+    lines.append("")
+    if score is not None:
+        pct = float(score) * 100.0
+        raw_display = int(pct)
+        final_display = raw_display
+        if raw_display > 150:
+            final_display = int(150 + 10 * math.log10(raw_display / 150))
+        lines.append(f"  Score ratio:        {score:<12.6f}  ({pct:.2f}%)")
+        lines.append(f"  Display score:      {final_display:<12d}  / 100")
+    else:
+        lines.append("  Score:  (unavailable)")
+    lines.append(f"  Pass:               {'OK' if passed else 'FAIL'}")
+    lines.append("-" * 64)
+    return lines
+
+
+def run_spj(testcases: list, describe_fn: Optional[Callable[[int], str]] = None, problem_name: str = "") -> int:
+    cwd = Path(".")
+    user_out_path = cwd / "user_out"
+    try:
+        user_text = read_text(user_out_path)
+    except Exception as ex:
+        eprint(str(ex))
+        print("0", flush=True)
+        return 0
+    nonce_ans = parse_nonce_from_text(user_text)
+    if nonce_ans is None:
+        eprint("SPJ FAIL: cannot find/parse nonce from user_out (expected 'OJCHAL v1 <b64>').")
+        print("0", flush=True)
+        return 0
+    results = parse_results_from_text(user_text)
+    if not results:
+        eprint("SPJ FAIL: no OJRESULT line found in user_out.")
+        print("0", flush=True)
+        return 0
+    payload_obj = None
+    for sig_hex, payload_b64 in reversed(results):
+        obj = verify_and_decode_payload(nonce_ans, sig_hex, payload_b64)
+        if obj is not None:
+            payload_obj = obj
+            break
+    if payload_obj is None:
+        eprint("SPJ FAIL: no OJRESULT line passes signature verification.")
+        print("0", flush=True)
+        return 0
+    ok = bool(payload_obj.get("pass", False))
+    tk_time_ms = payload_obj.get("tk_time_ms", None)
+    score_ratio = payload_obj.get("score_ratio", None)
+    if not isinstance(tk_time_ms, (int, float)):
+        eprint("SPJ FAIL: payload missing/invalid 'tk_time_ms'.")
+        print("0", flush=True)
+        return 0
+    if not isinstance(score_ratio, (int, float)) or not (0.0 <= score_ratio):
+        eprint("SPJ FAIL: payload missing/invalid 'score_ratio'.")
+        print("0", flush=True)
+        return 0
+    time_us = float(tk_time_ms) * 1000.0
+    if not ok:
+        eprint("SPJ FAIL: payload pass=false.")
+        print("0", flush=True)
+        return 0
+    displayScore = int(score_ratio * 100)
+    if displayScore > 150:
+        displayScore = int(150 + 10 * math.log10(displayScore / 150))
+    extraInfo = {"rewriteTimeUs": time_us, "displayScore": displayScore}
+    print(f"{100} {json.dumps(extraInfo)}", flush=True)
+    header = f"=== SPJ Report{' - ' + problem_name if problem_name else ''} ==="
+    eprint(header)
+    for line in _format_diagnostic(payload_obj, describe_fn=describe_fn):
+        eprint(line)
+    return 0
+
+HEAD_DIMS = [128]
+BATCH_SIZES = [1, 4, 16]
+SEQ_LENS_KV = [1024, 4096, 8192, 16384]
+SEQ_LEN_Q = 1
+NUM_HEADS = 8
+NUM_HEADS_K = 8
+PAGE_BLOCK_SIZE = 16
+CAUSAL = 0
+
+TESTCASES = []
+for headdim in HEAD_DIMS:
+    for seqlen_k in SEQ_LENS_KV:
+        for batch_size in BATCH_SIZES:
+            TESTCASES.append((batch_size, seqlen_k, SEQ_LEN_Q, NUM_HEADS, NUM_HEADS_K, headdim, PAGE_BLOCK_SIZE, CAUSAL))
+
+
+def describe(tc_id: int) -> str:
+    if tc_id < 1 or tc_id > len(TESTCASES):
+        return f"testcase #{tc_id}"
+    b, slk, slq, nh, nhk, hd, pbs, c = TESTCASES[tc_id - 1]
+    return f"batch={b}, seqlen_k={slk}, seqlen_q={slq}, heads={nh}, kv_heads={nhk}, headdim={hd}, page_size={pbs}, causal={c}"
+
+
+if __name__ == "__main__":
+    sys.exit(run_spj(TESTCASES, describe_fn=describe, problem_name="Paged KV Cache (flash-attn)"))

基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/testcase_config.py

@@ -296,3 +296,55 @@ try:
         return True
 except Exception:
     pass
+
+
+INPUT_CLASS = [
+    "INPUT",   # q
+    "INPUT",   # k_cache_paged
+    "INPUT",   # v_cache_paged
+    "OUTPUT",  # output
+    "INPUT",   # cache_seqlens
+    "INPUT",   # block_table
+    "INPUT",   # batch_size
+    "INPUT",   # seqlen_k
+    "INPUT",   # seqlen_q
+    "INPUT",   # num_heads
+    "INPUT",   # num_heads_k
+    "INPUT",   # headdim
+    "INPUT",   # page_block_size
+    "INPUT",   # num_blocks
+    "INPUT",   # causal
+]
+
+
+def getWorkload(testcase_sizes) -> dict:
+    raw_sizes = testcase_sizes[0] if isinstance(testcase_sizes, tuple) and len(testcase_sizes) == 2 else testcase_sizes
+    q_shape, k_shape, v_shape, output_shape, cache_seqlens_shape, block_table_shape = raw_sizes[:6]
+    batch_size, seqlen_q, num_heads, headdim = q_shape
+    num_blocks, page_block_size, num_heads_k, headdim_k = k_shape
+    assert v_shape == k_shape
+    assert output_shape == q_shape
+    blocks_per_batch = block_table_shape[1]
+    # KV 长度按 padded (blocks_per_batch * page_block_size) 估计
+    seqlen_k = blocks_per_batch * page_block_size
+    # QK: 2 * batch * seqlen_q * num_heads * seqlen_k * headdim
+    # PV: 2 * batch * seqlen_q * num_heads * seqlen_k * headdim
+    flops = 4 * batch_size * seqlen_q * num_heads * seqlen_k * headdim
+    # memory_bytes 只算输入/输出变量的 IO，不算 softmax 等中间结果
+    # k_cache / v_cache 按 tensor 实际占用 (padded) 计算 IO
+    memory_bytes = (
+        batch_size * seqlen_q * num_heads * headdim * 2      # q (bf16)
+        + num_blocks * page_block_size * num_heads_k * headdim * 2  # k_cache (bf16)
+        + num_blocks * page_block_size * num_heads_k * headdim * 2  # v_cache (bf16)
+        + batch_size * seqlen_q * num_heads * headdim * 2    # output (bf16)
+        + batch_size * 4                                      # cache_seqlens (int32)
+        + batch_size * blocks_per_batch * 4                   # block_table (int32)
+    )
+    return {
+        "flops": flops,
+        "memory_bytes": memory_bytes,
+        "dtype": "bf16",
+    }
+
+
+DESIGNED_VRAM_SIZE = 128

基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/_title.txt

@@ -1 +1 @@
-FlashAttention KV Cache Decode
+Agent 推理算子库优化 - FlashAttention KV Cache Decode