From 44354be5930f0d89ec94ec0be7aded4926ef4741 Mon Sep 17 00:00:00 2001 From: "Yuliang Feng (i26389)" Date: Mon, 29 Jun 2026 18:07:09 +0800 Subject: [PATCH] new oj problem description --- .../xpuoj_problem/meta.json | 6 +- .../xpuoj_problem/spj.py | 207 ++++++++++++++++++ .../xpuoj_problem/testcase_config.py | 52 +++++ ...樼洰鎻忚堪.md => 00_题目描述.md} | 0 ...﹀畾.cuda.md => 01_接口约定.cuda.md} | 0 ...ilelang.md => 01_接口约定.tilelang.md} | 0 ....triton.md => 01_接口约定.triton.md} | 0 ...撳叆鏍煎紡.md => 02_输入格式.md} | 0 ...撳嚭鏍煎紡.md => 03_输出格式.md} | 0 .../zh_CN/{04_鏍蜂緥.md => 04_样例.md} | 0 .../xpuoj_problem/zh_CN/_title.txt | 2 +- 11 files changed, 263 insertions(+), 4 deletions(-) create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/spj.py rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/{00_棰樼洰鎻忚堪.md => 00_题目描述.md} (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/{01_鎺ュ彛绾﹀畾.cuda.md => 01_接口约定.cuda.md} (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/{01_鎺ュ彛绾﹀畾.tilelang.md => 01_接口约定.tilelang.md} (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/{01_鎺ュ彛绾﹀畾.triton.md => 01_接口约定.triton.md} (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/{02_杈撳叆鏍煎紡.md => 02_输入格式.md} (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/{03_杈撳嚭鏍煎紡.md => 03_输出格式.md} (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/{04_鏍蜂緥.md => 04_样例.md} (100%) diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/meta.json b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/meta.json index 712ba01..09ee0f2 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/meta.json +++ b/基于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": [] -} +} \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/spj.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/spj.py new file mode 100644 index 0000000..ddf9a37 --- /dev/null +++ b/基于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 ').") + 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)")) diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/testcase_config.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/testcase_config.py index 7bb1f5b..2006e77 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/testcase_config.py +++ b/基于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 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/00_棰樼洰鎻忚堪.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/00_题目描述.md similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/00_棰樼洰鎻忚堪.md rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/00_题目描述.md diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/01_鎺ュ彛绾﹀畾.cuda.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/01_接口约定.cuda.md similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/01_鎺ュ彛绾﹀畾.cuda.md rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/01_接口约定.cuda.md diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/01_鎺ュ彛绾﹀畾.tilelang.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/01_接口约定.tilelang.md similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/01_鎺ュ彛绾﹀畾.tilelang.md rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/01_接口约定.tilelang.md diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/01_鎺ュ彛绾﹀畾.triton.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/01_接口约定.triton.md similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/01_鎺ュ彛绾﹀畾.triton.md rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/01_接口约定.triton.md diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/02_杈撳叆鏍煎紡.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/02_输入格式.md similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/02_杈撳叆鏍煎紡.md rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/02_输入格式.md diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/03_杈撳嚭鏍煎紡.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/03_输出格式.md similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/03_杈撳嚭鏍煎紡.md rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/03_输出格式.md diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/04_鏍蜂緥.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/04_样例.md similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/04_鏍蜂緥.md rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/04_样例.md diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/_title.txt b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/_title.txt index 1359d35..ff8cc61 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/_title.txt +++ b/基于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 \ No newline at end of file