From f302a83b62bd620da2861f1e809c34447699d10c Mon Sep 17 00:00:00 2001 From: xiaoke-shu <2732928868@qq.com> Date: Tue, 16 Jun 2026 21:57:22 +0800 Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0XPU-OJ=E6=95=99=E7=A8=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../FlashInfer关键算子迁移与优化.md | 247 +++++++++++++- ...ithPagedKVCacheWrapper_20260525_093143.csv | 0 ...APagedAttentionWrapper_20260525_074200.csv | 0 ...ithPagedKVCacheWrapper_20260526_150417.csv | 0 ...thRaggedKVCacheWrapper_20260526_145454.csv | 0 ...ithPagedKVCacheWrapper_20260525_093143.csv | Bin ...APagedAttentionWrapper_20260525_074200.csv | Bin ...ithPagedKVCacheWrapper_20260526_150417.csv | Bin ...thRaggedKVCacheWrapper_20260526_145454.csv | Bin .../__MACOSX/._bench_batch_decode.py | Bin .../__MACOSX/._bench_batch_mla.py | Bin .../__MACOSX/._bench_batch_prefill_paged.py | Bin .../__MACOSX/._bench_batch_prefill_ragged.py | Bin .../__MACOSX/._bench_common.py | Bin .../bench_batch_decode.py | 0 .../bench_batch_mla.py | 0 .../bench_batch_prefill_paged.py | 0 .../bench_batch_prefill_ragged.py | 0 .../{ => FlashInfer_Baseline}/bench_common.py | 0 .../xpuoj_problem/__MACOSX/._problem_20001 | Bin 0 -> 163 bytes .../xpuoj_problem/__MACOSX/._problem_20002 | Bin 0 -> 163 bytes .../xpuoj_problem/__MACOSX/._problem_20003 | Bin 0 -> 163 bytes .../xpuoj_problem/__MACOSX/._problem_20004 | Bin 0 -> 163 bytes .../xpuoj_problem/__MACOSX/._problem_20005 | Bin 0 -> 163 bytes .../__MACOSX/problem_20001/._meta.json | Bin 0 -> 163 bytes .../problem_20001/._testcase_config.py | Bin 0 -> 163 bytes .../__MACOSX/problem_20001/._zh_CN | Bin 0 -> 163 bytes .../zh_CN/._00_棰樼洰鎻忚堪.md | Bin 0 -> 163 bytes .../zh_CN/._01_鎺ュ彛绾﹀畾.cuda.md | Bin 0 -> 163 bytes .../zh_CN/._02_杈撳叆鏍煎紡.md | Bin 0 -> 163 bytes .../zh_CN/._03_杈撳嚭鏍煎紡.md | Bin 0 -> 163 bytes .../problem_20001/zh_CN/._04_鏍蜂緥.md | Bin 0 -> 163 bytes .../__MACOSX/problem_20001/zh_CN/.__title.txt | Bin 0 -> 163 bytes .../__MACOSX/problem_20002/._meta.json | Bin 0 -> 163 bytes .../problem_20002/._testcase_config.py | Bin 0 -> 163 bytes .../__MACOSX/problem_20002/._zh_CN | Bin 0 -> 163 bytes .../zh_CN/._00_棰樼洰鎻忚堪.md | Bin 0 -> 163 bytes .../zh_CN/._01_鎺ュ彛绾﹀畾.cuda.md | Bin 0 -> 163 bytes .../zh_CN/._02_杈撳叆鏍煎紡.md | Bin 0 -> 163 bytes .../zh_CN/._03_杈撳嚭鏍煎紡.md | Bin 0 -> 163 bytes .../problem_20002/zh_CN/._04_鏍蜂緥.md | Bin 0 -> 163 bytes .../__MACOSX/problem_20002/zh_CN/.__title.txt | Bin 0 -> 163 bytes .../__MACOSX/problem_20003/._meta.json | Bin 0 -> 163 bytes .../problem_20003/._testcase_config.py | Bin 0 -> 163 bytes .../__MACOSX/problem_20003/._zh_CN | Bin 0 -> 163 bytes .../zh_CN/._00_棰樼洰鎻忚堪.md | Bin 0 -> 163 bytes .../zh_CN/._01_鎺ュ彛绾﹀畾.cuda.md | Bin 0 -> 163 bytes .../zh_CN/._02_杈撳叆鏍煎紡.md | Bin 0 -> 163 bytes .../zh_CN/._03_杈撳嚭鏍煎紡.md | Bin 0 -> 163 bytes .../problem_20003/zh_CN/._04_鏍蜂緥.md | Bin 0 -> 163 bytes .../__MACOSX/problem_20003/zh_CN/.__title.txt | Bin 0 -> 163 bytes .../__MACOSX/problem_20004/._meta.json | Bin 0 -> 163 bytes .../problem_20004/._testcase_config.py | Bin 0 -> 163 bytes .../__MACOSX/problem_20004/._zh_CN | Bin 0 -> 163 bytes .../zh_CN/._00_棰樼洰鎻忚堪.md | Bin 0 -> 163 bytes .../zh_CN/._01_鎺ュ彛绾﹀畾.cuda.md | Bin 0 -> 163 bytes .../zh_CN/._02_杈撳叆鏍煎紡.md | Bin 0 -> 163 bytes .../zh_CN/._03_杈撳嚭鏍煎紡.md | Bin 0 -> 163 bytes .../problem_20004/zh_CN/._04_鏍蜂緥.md | Bin 0 -> 163 bytes .../__MACOSX/problem_20004/zh_CN/.__title.txt | Bin 0 -> 163 bytes .../__MACOSX/problem_20005/._meta.json | Bin 0 -> 163 bytes .../problem_20005/._testcase_config.py | Bin 0 -> 163 bytes .../__MACOSX/problem_20005/._zh_CN | Bin 0 -> 163 bytes .../zh_CN/._00_棰樼洰鎻忚堪.md | Bin 0 -> 163 bytes .../zh_CN/._01_鎺ュ彛绾﹀畾.cuda.md | Bin 0 -> 163 bytes .../zh_CN/._02_杈撳叆鏍煎紡.md | Bin 0 -> 163 bytes .../zh_CN/._03_杈撳嚭鏍煎紡.md | Bin 0 -> 163 bytes .../problem_20005/zh_CN/._04_鏍蜂緥.md | Bin 0 -> 163 bytes .../__MACOSX/problem_20005/zh_CN/.__title.txt | Bin 0 -> 163 bytes .../xpuoj_problem/problem_20001/meta.json | 16 + .../problem_20001/testcase_config.py | 307 ++++++++++++++++++ .../problem_20001/zh_CN/00_题目描述.md | 23 ++ .../problem_20001/zh_CN/01_接口约定.md | 41 +++ .../problem_20001/zh_CN/02_数据范围.md | 26 ++ .../problem_20001/zh_CN/03_参考实现.md | 28 ++ .../problem_20001/zh_CN/_title.txt | 1 + .../xpuoj_problem/problem_20002/meta.json | 16 + .../problem_20002/testcase_config.py | 273 ++++++++++++++++ .../problem_20002/zh_CN/00_题目描述.md | 22 ++ .../problem_20002/zh_CN/01_接口约定.md | 42 +++ .../problem_20002/zh_CN/02_数据范围.md | 26 ++ .../problem_20002/zh_CN/03_参考实现.md | 31 ++ .../problem_20002/zh_CN/_title.txt | 1 + .../xpuoj_problem/problem_20003/meta.json | 16 + .../problem_20003/testcase_config.py | 245 ++++++++++++++ .../problem_20003/zh_CN/00_题目描述.md | 22 ++ .../problem_20003/zh_CN/01_接口约定.md | 46 +++ .../problem_20003/zh_CN/02_数据范围.md | 26 ++ .../problem_20003/zh_CN/03_参考实现.md | 30 ++ .../problem_20003/zh_CN/_title.txt | 1 + .../xpuoj_problem/problem_20004/meta.json | 16 + .../problem_20004/testcase_config.py | 248 ++++++++++++++ .../problem_20004/zh_CN/00_题目描述.md | 22 ++ .../problem_20004/zh_CN/01_接口约定.md | 38 +++ .../problem_20004/zh_CN/02_数据范围.md | 26 ++ .../problem_20004/zh_CN/03_参考实现.md | 29 ++ .../problem_20004/zh_CN/_title.txt | 1 + 97 files changed, 1856 insertions(+), 10 deletions(-) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/{ => FlashInfer_Baseline}/BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/{ => FlashInfer_Baseline}/BatchMLAPagedAttentionWrapper_20260525_074200.csv (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/{ => FlashInfer_Baseline}/BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/{ => FlashInfer_Baseline}/BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/{ => FlashInfer_Baseline}/__MACOSX/._BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/{ => FlashInfer_Baseline}/__MACOSX/._BatchMLAPagedAttentionWrapper_20260525_074200.csv (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/{ => FlashInfer_Baseline}/__MACOSX/._BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/{ => FlashInfer_Baseline}/__MACOSX/._BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/{ => FlashInfer_Baseline}/__MACOSX/._bench_batch_decode.py (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/{ => FlashInfer_Baseline}/__MACOSX/._bench_batch_mla.py (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/{ => FlashInfer_Baseline}/__MACOSX/._bench_batch_prefill_paged.py (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/{ => FlashInfer_Baseline}/__MACOSX/._bench_batch_prefill_ragged.py (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/{ => FlashInfer_Baseline}/__MACOSX/._bench_common.py (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/{ => FlashInfer_Baseline}/bench_batch_decode.py (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/{ => FlashInfer_Baseline}/bench_batch_mla.py (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/{ => FlashInfer_Baseline}/bench_batch_prefill_paged.py (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/{ => FlashInfer_Baseline}/bench_batch_prefill_ragged.py (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/{ => FlashInfer_Baseline}/bench_common.py (100%) create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/._problem_20001 create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/._problem_20002 create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/._problem_20003 create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/._problem_20004 create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/._problem_20005 create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20001/._meta.json create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20001/._testcase_config.py create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20001/._zh_CN create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20001/zh_CN/._00_棰樼洰鎻忚堪.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20001/zh_CN/._01_鎺ュ彛绾﹀畾.cuda.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20001/zh_CN/._02_杈撳叆鏍煎紡.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20001/zh_CN/._03_杈撳嚭鏍煎紡.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20001/zh_CN/._04_鏍蜂緥.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20001/zh_CN/.__title.txt create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20002/._meta.json create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20002/._testcase_config.py create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20002/._zh_CN create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20002/zh_CN/._00_棰樼洰鎻忚堪.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20002/zh_CN/._01_鎺ュ彛绾﹀畾.cuda.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20002/zh_CN/._02_杈撳叆鏍煎紡.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20002/zh_CN/._03_杈撳嚭鏍煎紡.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20002/zh_CN/._04_鏍蜂緥.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20002/zh_CN/.__title.txt create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20003/._meta.json create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20003/._testcase_config.py create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20003/._zh_CN create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20003/zh_CN/._00_棰樼洰鎻忚堪.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20003/zh_CN/._01_鎺ュ彛绾﹀畾.cuda.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20003/zh_CN/._02_杈撳叆鏍煎紡.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20003/zh_CN/._03_杈撳嚭鏍煎紡.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20003/zh_CN/._04_鏍蜂緥.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20003/zh_CN/.__title.txt create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20004/._meta.json create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20004/._testcase_config.py create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20004/._zh_CN create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20004/zh_CN/._00_棰樼洰鎻忚堪.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20004/zh_CN/._01_鎺ュ彛绾﹀畾.cuda.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20004/zh_CN/._02_杈撳叆鏍煎紡.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20004/zh_CN/._03_杈撳嚭鏍煎紡.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20004/zh_CN/._04_鏍蜂緥.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20004/zh_CN/.__title.txt create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20005/._meta.json create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20005/._testcase_config.py create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20005/._zh_CN create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20005/zh_CN/._00_棰樼洰鎻忚堪.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20005/zh_CN/._01_鎺ュ彛绾﹀畾.cuda.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20005/zh_CN/._02_杈撳叆鏍煎紡.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20005/zh_CN/._03_杈撳嚭鏍煎紡.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20005/zh_CN/._04_鏍蜂緥.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/problem_20005/zh_CN/.__title.txt create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/meta.json create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/testcase_config.py create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/00_题目描述.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/01_接口约定.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/02_数据范围.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/03_参考实现.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/_title.txt create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/meta.json create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/testcase_config.py create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/00_题目描述.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/01_接口约定.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/02_数据范围.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/03_参考实现.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/_title.txt create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/meta.json create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/testcase_config.py create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/00_题目描述.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/01_接口约定.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/02_数据范围.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/03_参考实现.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/_title.txt create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/meta.json create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/testcase_config.py create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/00_题目描述.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/01_接口约定.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/02_数据范围.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/03_参考实现.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/_title.txt diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer关键算子迁移与优化.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer关键算子迁移与优化.md index 1d1e509..36de497 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer关键算子迁移与优化.md +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer关键算子迁移与优化.md @@ -11,6 +11,8 @@ 3. 跑通 BatchDecode、BatchPrefill、MLA 等典型算子的 API 调用示例； 4. 完成各算子在不同参数配置下的 Benchmark 测试； 5. 输出各算子的 Baseline 性能结果记录表，为后续算子优化提供对比基准。 +6. 理解 XPU-OJ 评测`run_kernel`接口与精度要求。 +7. 理解 Baseline 与 XPU-OJ 评测题包之间的关系，能够根据题包接口实现一个最小正确版 run_kernel。 ## 三、适用对象 @@ -39,9 +41,9 @@ * 步骤1：获取算力券[https://developer.metax-tech.com/activities/6](https://developer.metax-tech.com/activities/6) * 首次登录需要先进行注册（使用邮箱或者手机号进行注册） - + * 登录成功后进行第二步-邮箱验证，填入自己的邮箱。 - + * 第三步，提交申请。 * 获得兑换码 @@ -124,7 +126,9 @@ prefill = 并行处理用户输入，decode = 逐个生成回答 token ## 六、项目实践--FlashInfer-Baseline -### Step 1：检查运行环境 +## 本地项目运行 + +#### Step 1：检查运行环境 **目标：** 进入terminal确认当前环境满足本模块运行要求。 @@ -140,7 +144,7 @@ prefill = 并行处理用户输入，decode = 逐个生成回答 token | `No GPUs are available` | 检查 MXMACA 驱动是否正确安装 | | `ModuleNotFoundError: No module named 'flashinfer'` | 执行 `pip install flashinfer` 安装 | -### Step 2：进入项目目录 +#### Step 2：进入项目目录 **目标：** 进入本模块所需的源码目录[flashinfer_baseline](https://gitlink.org.cn/metax-maca/op_optimization/tree/master/%E5%9F%BA%E4%BA%8EAI%20Agent%E5%BC%80%E5%8F%91%E8%8C%83%E5%BC%8F%E7%9A%84%E5%9B%BD%E4%BA%A7GPU%E5%A4%A7%E6%A8%A1%E5%9E%8B%E6%8E%A8%E7%90%86%E7%AE%97%E5%AD%90%E5%BA%93%E4%BC%98%E5%8C%96%2Fbaselines%2Fflashinfer_baseline)。 @@ -154,12 +158,12 @@ prefill = 并行处理用户输入，decode = 逐个生成回答 token 在仓库目录 `op_optimization/基于AI Agent开发范式的国产GPU大模型推理算子库优化` 下，找到 `flashinfer_baseline` 文件夹。可以将 `flashinfer_baseline` 整个目录复制到工作目录 `data/` 下。 -**下一步操作：** 切换到 FlashInfer Baseline 项目目录。 +**下一步操作：** 切换到 FlashInfer_Baseline 项目目录。 **命令示例：** ```bash -cd /data/flashinfer_baseline +cd baselines\flashinfer_baseline\FlashInfer_Baseline ls -la ``` @@ -176,7 +180,7 @@ README.md ``` -### Step 3：验证项目脚本 +#### Step 3：验证项目脚本 **目标：** 确认所有基准测试脚本可正常执行。 @@ -205,7 +209,7 @@ python -c "from bench_common import setup_workspace, get_csv_path; print('脚本 ``` -### Step 4：运行单算子 Benchmark +#### Step 4：运行单算子 Benchmark **目标：** 执行基准测试，获取 Baseline 性能数据。 @@ -237,7 +241,7 @@ Results saved to BatchDecodeWithPagedKVCacheWrapper_20260525_xxxxxx.csv | `out of memory` | 减小 batch\_size 或 seq\_len 参数 | | 运行时间过长 | 脚本会自动调整重复次数，耐心等待 | -### Step 5：查看测试结果 +#### Step 5：查看测试结果 **目标：** 查看并分析 Benchmark 输出结果。 @@ -272,7 +276,226 @@ if csv_files: ``` +### **XPU-OJ 评测与提交** +#### Step 1：访问XPUOJ平台(https://xpuoj.com/)进行登录 + +- 完成 baseline benchmark 后，要到统一评测工具XPU-OJ完成提交，进行评测。 + +- 等待组委会统一发放 XPU-OJ 账号 +- 使用分配的用户名和初始密码登录平台 + +![image 20260616152804368](https://origin.picgo.net/2026/06/16/image-20260616152804368aa0d44b72c4f9572.png) + +#### Step 2：进入比赛页面 + +- 登录后进入对应赛题页面 +- 选择对应的赛题 + +![image 20260616152953345](https://origin.picgo.net/2026/06/16/image-20260616152953345986ce39fda69da55.png) + +#### Step 3：选择对应的算子题目进行评测 + +![image 20260616153142383](https://origin.picgo.net/2026/06/16/image-202606161531423833b0f0428edf2e35e.png) + +#### Step 4：阅读题目要求，接口约定，编写符合`run_kernel`接口规范的CUDA算子 + + 以下代码仅用于说明接口结构，不代表最优实现，也不作为评分参考： + +```python +FlashInfer Ragged Prefill示例代码 +#include +#include +#include +#include + +#ifndef max +#define max(a,b) ((a) > (b) ? (a) : (b)) +#endif + +#ifndef min +#define min(a,b) ((a) < (b) ? (a) : (b)) +#endif + +template +__global__ void ragged_prefill_kernel( + const __nv_bfloat16* __restrict__ q, + const __nv_bfloat16* __restrict__ k, + const __nv_bfloat16* __restrict__ v, + __nv_bfloat16* __restrict__ output, + const int32_t* __restrict__ qo_indptr, + const int32_t* __restrict__ kv_indptr, + int64_t seq_len, + int64_t num_qo_heads, + int64_t num_kv_heads, + float scale, + int64_t causal +) { + int b = blockIdx.x; + int hq = blockIdx.y; + int hkv = hq / (num_qo_heads / num_kv_heads); + + int ty = threadIdx.y; // 0~31, Block 内的 query 索引 (对应 Br) + int tx = threadIdx.x; // 0~31, Warp 内的线程 (lane) + + // 使用 int64_t 防止长序列+大Batch组合下的整数溢出 + int64_t q_idx = (int64_t)blockIdx.z * Br + ty; + bool is_valid = q_idx < seq_len; + + int64_t cur_qo_start = qo_indptr[b]; + int64_t cur_kv_start = kv_indptr[b]; + + float r_Q[D_QK / 32]; + if (is_valid) { + // 显式使用 int64_t 计算全局偏移量，防止超过 2GB 显存边界 + int64_t q_offset_base = (cur_qo_start + q_idx) * num_qo_heads * D_QK + hq * D_QK; + #pragma unroll + for (int d = 0; d < D_QK / 32; ++d) { + r_Q[d] = __bfloat162float(q[q_offset_base + tx + d * 32]); + } + } + + float r_O[D_VO / 32]; + #pragma unroll + for (int d = 0; d < D_VO / 32; ++d) r_O[d] = 0.0f; + + float m = -1e20f; + float l = 0.0f; + + extern __shared__ __nv_bfloat16 smem[]; + __nv_bfloat16* shared_K = smem; + __nv_bfloat16* shared_V = smem + Bc * D_QK; + + int tid = ty * 32 + tx; + int total_threads = Br * 32; + + int64_t block_max_q = min((int64_t)seq_len - 1, (int64_t)(blockIdx.z * Br + Br - 1)); + + int total_k_ldg = (Bc * D_QK) / 8; + int total_v_ldg = (Bc * D_VO) / 8; + + for (int64_t kv_start = 0; kv_start <= block_max_q; kv_start += Bc) { + + // 协同加载 K + for (int i = tid; i < total_k_ldg; i += total_threads) { + int token_idx = i / (D_QK / 8); + int d_idx = i % (D_QK / 8); + int64_t global_kv_idx = kv_start + token_idx; + + if (global_kv_idx < seq_len) { + // 同样强制为 int64_t + int64_t k_offset = (cur_kv_start + global_kv_idx) * num_kv_heads * D_QK + hkv * D_QK + d_idx * 8; + ((float4*)shared_K)[i] = *(reinterpret_cast(&k[k_offset])); + } else { + ((float4*)shared_K)[i] = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + } + } + + // 协同加载 V + for (int i = tid; i < total_v_ldg; i += total_threads) { + int token_idx = i / (D_VO / 8); + int d_idx = i % (D_VO / 8); + int64_t global_kv_idx = kv_start + token_idx; + + if (global_kv_idx < seq_len) { + int64_t v_offset = (cur_kv_start + global_kv_idx) * num_kv_heads * D_VO + hkv * D_VO + d_idx * 8; + ((float4*)shared_V)[i] = *(reinterpret_cast(&v[v_offset])); + } else { + ((float4*)shared_V)[i] = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + } + } + __syncthreads(); + + if (is_valid) { + int64_t kv_end = causal ? min(kv_start + Bc, q_idx + 1) : min(kv_start + Bc, seq_len); + + for (int j = 0; j < kv_end - kv_start; ++j) { + float score = 0.0f; + #pragma unroll + for (int d = 0; d < D_QK / 32; ++d) { + score += r_Q[d] * __bfloat162float(shared_K[j * D_QK + tx + d * 32]); + } + + #pragma unroll + for (int offset = 16; offset > 0; offset /= 2) { + score += __shfl_down_sync(0xffffffff, score, offset); + } + score = __shfl_sync(0xffffffff, score, 0); + score *= scale; + + float m_prev = m; + m = max(m, score); + float P = __expf(score - m); + l = l * __expf(m_prev - m) + P; + + #pragma unroll + for (int d = 0; d < D_VO / 32; ++d) { + r_O[d] = r_O[d] * __expf(m_prev - m) + P * __bfloat162float(shared_V[j * D_VO + tx + d * 32]); + } + } + } + __syncthreads(); + } + + if (is_valid) { + int64_t out_offset_base = (cur_qo_start + q_idx) * num_qo_heads * D_VO + hq * D_VO; + #pragma unroll + for (int d = 0; d < D_VO / 32; ++d) { + float out_val = r_O[d] / l; + output[out_offset_base + tx + d * 32] = __float2bfloat16(out_val); + } + } +} + +extern "C" void run_kernel( + const __nv_bfloat16* q, + const __nv_bfloat16* k, + const __nv_bfloat16* v, + __nv_bfloat16* output, + const int32_t* qo_indptr, + const int32_t* kv_indptr, + int64_t batch_size, + int64_t seq_len, + int64_t num_qo_heads, + int64_t num_kv_heads, + int64_t head_dim_qk, + int64_t head_dim_vo, + int64_t causal +) { + float scale = 1.0f / sqrtf((float)head_dim_qk); + const int Br = 32; + + dim3 grid(batch_size, num_qo_heads, (seq_len + Br - 1) / Br); + dim3 block(32, Br); + + // 根据维度配置绝对安全的共享内存策略 + if (head_dim_qk == 128 && head_dim_vo == 128) { + const int Bc = 64; + int smem = Bc * 128 * 2 + Bc * 128 * 2; // 32 KB + cudaFuncSetAttribute(ragged_prefill_kernel<128, 128, Br, Bc>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + ragged_prefill_kernel<128, 128, Br, Bc><<>>( + q, k, v, output, qo_indptr, kv_indptr, seq_len, num_qo_heads, num_kv_heads, scale, causal); + } else if (head_dim_qk == 192 && head_dim_vo == 128) { + const int Bc = 64; + int smem = Bc * 192 * 2 + Bc * 128 * 2; // 40 KB + cudaFuncSetAttribute(ragged_prefill_kernel<192, 128, Br, Bc>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + ragged_prefill_kernel<192, 128, Br, Bc><<>>( + q, k, v, output, qo_indptr, kv_indptr, seq_len, num_qo_heads, num_kv_heads, scale, causal); + } else if (head_dim_qk == 256 && head_dim_vo == 256) { + // 关键降级：256维度下，将 Bc 从 64 降为 32。 + // 将所需的共享内存从 64KB 压回到绝对安全的 32KB，防止在部分环境触发 Launch Failure (OOM/Crash)。 + const int Bc = 32; + int smem = Bc * 256 * 2 + Bc * 256 * 2; // 32 KB + cudaFuncSetAttribute(ragged_prefill_kernel<256, 256, Br, Bc>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + ragged_prefill_kernel<256, 256, Br, Bc><<>>( + q, k, v, output, qo_indptr, kv_indptr, seq_len, num_qo_heads, num_kv_heads, scale, causal); + } +} +``` + +#### Step 5：提交代码开始评测 + +![image 20260616155023443](https://origin.picgo.net/2026/06/16/image-202606161550234433e54d703d581f858.png) ## 七、常见问题 @@ -353,4 +576,8 @@ mv *.csv results/ | 优化项 | 改动内容 | Baseline | 优化后 | 提升比例 | | --- | --- | --- | --- | --- | | 例：调整 block\_size | 16 → 32 | xx ms | xx ms | xx% | -| 完成优化后，再次运行本模块的 Benchmark 脚本，对比前后性能变化。 | | | | | \ No newline at end of file +| 完成优化后，再次运行本模块的 Benchmark 脚本，对比前后性能变化。 | | | | | + +### 6. 使用多语言完成算子优化加速 + +可以使用Triton和TileLang语言实现`run_kernel`接口完成算子优化，对比不同的语言对于性能加速的影响 \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/BatchMLAPagedAttentionWrapper_20260525_074200.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/BatchMLAPagedAttentionWrapper_20260525_074200.csv similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/BatchMLAPagedAttentionWrapper_20260525_074200.csv rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/BatchMLAPagedAttentionWrapper_20260525_074200.csv diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/__MACOSX/._BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/__MACOSX/._BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._BatchMLAPagedAttentionWrapper_20260525_074200.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/__MACOSX/._BatchMLAPagedAttentionWrapper_20260525_074200.csv similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._BatchMLAPagedAttentionWrapper_20260525_074200.csv rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/__MACOSX/._BatchMLAPagedAttentionWrapper_20260525_074200.csv diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/__MACOSX/._BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/__MACOSX/._BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/__MACOSX/._BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/__MACOSX/._BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_batch_decode.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/__MACOSX/._bench_batch_decode.py similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_batch_decode.py rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/__MACOSX/._bench_batch_decode.py diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_batch_mla.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/__MACOSX/._bench_batch_mla.py similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_batch_mla.py rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/__MACOSX/._bench_batch_mla.py diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_batch_prefill_paged.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/__MACOSX/._bench_batch_prefill_paged.py similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_batch_prefill_paged.py rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/__MACOSX/._bench_batch_prefill_paged.py diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_batch_prefill_ragged.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/__MACOSX/._bench_batch_prefill_ragged.py similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_batch_prefill_ragged.py rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/__MACOSX/._bench_batch_prefill_ragged.py diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_common.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/__MACOSX/._bench_common.py similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_common.py rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/__MACOSX/._bench_common.py diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_decode.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/bench_batch_decode.py similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_decode.py rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/bench_batch_decode.py diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_mla.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/bench_batch_mla.py similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_mla.py rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/bench_batch_mla.py diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_prefill_paged.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/bench_batch_prefill_paged.py similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_prefill_paged.py rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/bench_batch_prefill_paged.py diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_prefill_ragged.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/bench_batch_prefill_ragged.py similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_prefill_ragged.py rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/bench_batch_prefill_ragged.py diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_common.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/bench_common.py similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_common.py rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/FlashInfer_Baseline/bench_common.py diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/._problem_20001 b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/__MACOSX/._problem_20001 new file mode 100644 index 0000000000000000000000000000000000000000..e02b168933473a6c317e19406aaeeb1b918f8c63 GIT binary patch literal 163 zcmZQz6=P>$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K$Vqox1Ojhs@R)|o50+1L3ClDI}aUl?c_=|y<2;dkJ5(HHS(lG;wxzV&S oBE&_L^K int: + return len(TESTCASES) + + +try: + from pathlib import Path + from typing import List, Tuple, Union + import sys + + import torch + + KernelArg = Union[torch.Tensor, int, float] + CURRENT_CASE = None + + def _ensure_flashinfer_importable(): + try: + import flashinfer # noqa: F401 + + return + except ImportError: + pass + + here = Path(__file__).resolve() + for parent in here.parents: + candidate = parent / "McFlashInfer" + if (candidate / "flashinfer").is_dir(): + sys.path.insert(0, str(candidate)) + return + + def _get_testcase_index() -> int: + try: + raw = input().strip() + except EOFError: + return 0 + if raw == "": + return 0 + try: + testcase_id = int(raw.split()[0]) + except ValueError: + return 0 + if 1 <= testcase_id <= len(TESTCASES): + return testcase_id - 1 + if 0 <= testcase_id < len(TESTCASES): + return testcase_id + return 0 + + def _compute_reps(batch_size: int, seq_len: int, head_dim: int, base_reps: int = 100) -> int: + workload = batch_size * seq_len * head_dim + if workload < 1e5: + return base_reps + if workload < 1e6: + return base_reps // 2 + if workload < 1e7: + return base_reps // 4 + if workload < 1e8: + return base_reps // 8 + if workload < 1e9: + return base_reps // 16 + return base_reps // 32 + + def getTestCaseSize() -> Tuple[List[Tuple[int, ...]], Tuple[int, int]]: + testcase_id = _get_testcase_index() + global CURRENT_CASE + batch_size, seq_len, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, causal = TESTCASES[testcase_id] + CURRENT_CASE = ( + batch_size, + seq_len, + num_qo_heads, + num_kv_heads, + head_dim_qk, + head_dim_vo, + causal, + 20260610 + testcase_id, + ) + qo_len = batch_size * seq_len + kv_len = batch_size * seq_len + warmup = 3 + iters = max(1, _compute_reps(batch_size, seq_len, head_dim_qk + head_dim_vo)) + return [ + (qo_len, num_qo_heads, head_dim_qk), + (kv_len, num_kv_heads, head_dim_qk), + (kv_len, num_kv_heads, head_dim_vo), + (qo_len, num_qo_heads, head_dim_vo), + (batch_size + 1,), + (batch_size + 1,), + (), (), (), (), (), (), (), + ], (warmup, iters) + + def genTestCase(testcase_sizes, device: str = "cuda") -> List[KernelArg]: + del testcase_sizes + batch_size, seq_len, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, causal, seed = CURRENT_CASE + gen = torch.Generator(device=device) + gen.manual_seed(seed) + dtype = torch.bfloat16 + qo_len = batch_size * seq_len + q = torch.rand( + qo_len, + num_qo_heads, + head_dim_qk, + dtype=dtype, + device=device, + generator=gen, + ).contiguous() + kv_len = batch_size * seq_len + k = torch.rand( + kv_len, + num_kv_heads, + head_dim_qk, + dtype=dtype, + device=device, + generator=gen, + ).contiguous() + v = torch.rand( + kv_len, + num_kv_heads, + head_dim_vo, + dtype=dtype, + device=device, + generator=gen, + ).contiguous() + output = torch.empty( + qo_len, + num_qo_heads, + head_dim_vo, + dtype=dtype, + device=device, + ) + qo_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device=device) * seq_len + kv_indptr = qo_indptr.clone() + return [ + q, + k, + v, + output, + qo_indptr, + kv_indptr, + batch_size, + seq_len, + num_qo_heads, + num_kv_heads, + head_dim_qk, + head_dim_vo, + causal, + ] + + def baseline( + q, + k, + v, + output, + qo_indptr, + kv_indptr, + batch_size, + seq_len, + num_qo_heads, + num_kv_heads, + head_dim_qk, + head_dim_vo, + causal, + ): + _ensure_flashinfer_importable() + import flashinfer + + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device) + wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper( + workspace_buffer, + kv_layout="NHD", + backend="auto", + ) + wrapper.plan( + qo_indptr, + kv_indptr, + int(num_qo_heads), + int(num_kv_heads), + int(head_dim_qk), + int(head_dim_vo), + causal=bool(causal), + q_data_type=torch.bfloat16, + kv_data_type=torch.bfloat16, + ) + wrapper.run(q, k, v, out=output) + return [ + q, + k, + v, + output, + qo_indptr, + kv_indptr, + batch_size, + seq_len, + num_qo_heads, + num_kv_heads, + head_dim_qk, + head_dim_vo, + causal, + ] + + def check( + testcase_sizes, + original_input_tensors, + target_kernel_input_tensors, + baseline_input_tensors, + rtol=1e-2, + atol=1e-2, + ) -> bool: + del testcase_sizes, original_input_tensors + output_t = target_kernel_input_tensors[3] + output_ref = baseline_input_tensors[3] + if output_t.shape != output_ref.shape: + print(f"[FAIL] shape mismatch: target {output_t.shape}, ref {output_ref.shape}", file=sys.stderr) + return False + if output_t.dtype != output_ref.dtype: + print(f"[FAIL] dtype mismatch: target {output_t.dtype}, ref {output_ref.dtype}", file=sys.stderr) + return False + if not torch.allclose(output_t.float(), output_ref.float(), rtol=rtol, atol=atol): + diff = (output_t.float() - output_ref.float()).abs() + print( + f"[FAIL] allclose failed: max_abs_diff={float(diff.max().item()):.6f}, " + f"mean_abs_diff={float(diff.mean().item()):.6f} (rtol={rtol}, atol={atol})", + file=sys.stderr, + ) + return False + return True +except Exception: + pass + + +INPUT_CLASS = [ + "INPUT", + "INPUT", + "INPUT", + "OUTPUT", + "INPUT", + "INPUT", + "INPUT", + "INPUT", + "INPUT", + "INPUT", + "INPUT", + "INPUT", + "INPUT", +] + + +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, qo_indptr_shape, kv_indptr_shape = raw_sizes[:6] + qo_len, num_qo_heads, head_dim_qk = q_shape + kv_len, num_kv_heads, k_dim = k_shape + v_len, v_heads, head_dim_vo = v_shape + assert k_dim == head_dim_qk + assert v_len == kv_len + assert v_heads == num_kv_heads + assert qo_len == kv_len + assert output_shape == (qo_len, num_qo_heads, head_dim_vo) + assert qo_indptr_shape == kv_indptr_shape + batch_size = qo_indptr_shape[0] - 1 + seq_len = kv_len // batch_size + flops = batch_size * seq_len * seq_len * num_qo_heads * (head_dim_qk + head_dim_vo) + memory_bytes = ( + qo_len * num_qo_heads * head_dim_qk * 2 + + kv_len * num_kv_heads * head_dim_qk * 2 + + kv_len * num_kv_heads * head_dim_vo * 2 + + qo_len * num_qo_heads * head_dim_vo * 2 + + (batch_size + 1) * 4 * 2 + ) + return { + "flops": flops, + "memory_bytes": memory_bytes, + "dtype": "bf16", + } + + +DESIGNED_VRAM_SIZE = 48 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/00_题目描述.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/00_题目描述.md new file mode 100644 index 0000000..4f97990 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/00_题目描述.md @@ -0,0 +1,23 @@ +--- +sectionTitle: "题目描述" +type: "Text" +--- +你需要实现 FlashInfer ragged KV cache prefill 的CUDA C++前向算子。 + +本题输入采用 FlashInfer `BatchPrefillWithRaggedKVCacheWrapper` 的 ragged `NHD` 布局。每个 batch 中有 `seq_len` 个 query token，KV cache 中也有 `seq_len` 个 token： + +其中 query heads 采用 GQA 布局：`num_qo_heads` 个 query/output heads 共享 `num_kv_heads` 个 KV heads，`G = num_qo_heads / num_kv_heads`。 + +评测程序会调用你提交代码中的 `run_kernel` 函数。你需要根据 `qo_indptr` 和 `kv_indptr` 读取 ragged Q/K/V，并将结果写入 `output`。 + +baseline 使用 FlashInfer ragged prefill 的 Python API： + +```python +wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper(workspace, kv_layout="NHD", backend="auto") +wrapper.plan(qo_indptr, kv_indptr, num_qo_heads, num_kv_heads, + head_dim_qk, head_dim_vo, causal=True, + q_data_type=torch.bfloat16, kv_data_type=torch.bfloat16) +wrapper.run(q, k, v, out=output) +``` + +如何提交代码详见[评测指南](/d/2)。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/01_接口约定.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/01_接口约定.md new file mode 100644 index 0000000..06f1703 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/01_接口约定.md @@ -0,0 +1,41 @@ +--- +sectionTitle: "接口约定" +type: "codeSample" +lang: "cuda" +--- +你必须在提交的 CUDA 源码中提供如下 **C 符号**，函数名、参数类型、顺序必须完全一致，并使用 `extern "C"` 防止 name mangling： + +```cpp +#include +#include + +extern "C" void run_kernel( + const __nv_bfloat16* q, + const __nv_bfloat16* k, + const __nv_bfloat16* v, + __nv_bfloat16* output, + const int32_t* qo_indptr, + const int32_t* kv_indptr, + int64_t batch_size, + int64_t seq_len, + int64_t num_qo_heads, + int64_t num_kv_heads, + int64_t head_dim_qk, + int64_t head_dim_vo, + int64_t causal +); +``` + +### 参数说明 + +* `q`：query tensor，shape `(batch_size * seq_len, num_qo_heads, head_dim_qk)`，连续 `bf16` +* `k`：key tensor，shape `(batch_size * seq_len, num_kv_heads, head_dim_qk)`，连续 `bf16` +* `v`：value tensor，shape `(batch_size * seq_len, num_kv_heads, head_dim_vo)`，连续 `bf16` +* `output`：输出缓冲区，shape `(batch_size * seq_len, num_qo_heads, head_dim_vo)`，连续 `bf16` +* `qo_indptr`：query/output ragged indptr，shape `(batch_size + 1)`，连续 `int32` +* `kv_indptr`：KV ragged indptr，shape `(batch_size + 1)`，连续 `int32` +* `causal`：是否启用 causal mask，评测中固定为 `1` + +本题测试中 `qo_indptr[b + 1] - qo_indptr[b] == seq_len`，`kv_indptr[b + 1] - kv_indptr[b] == seq_len`。 + +`run_kernel` 内部需要自行计算合适的 launch 配置并启动 CUDA kernel。为保证计时准确，不建议在 `run_kernel` 内部做 `cudaDeviceSynchronize()` 或显式同步。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/02_数据范围.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/02_数据范围.md new file mode 100644 index 0000000..1e57f2e --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/02_数据范围.md @@ -0,0 +1,26 @@ +--- +sectionTitle: "输入格式" +type: "Text" +--- +## 数据范围 + +- 数据类型：`q/kv_data/output` 均为 `bfloat16` +- KV layout：`NHD` +- `num_qo_heads = 32` +- `num_kv_heads = 4` +- `page_block_size = 16` +- `causal = 0` +- `head_dim` 固定为 `128` +- `batch_size` 取值为 `1, 4, 16` +- `seq_len` 取值为 `1024, 4096, 8192, 16384` + +测试点顺序与 `McFlashInfer/benchmarks/bench_batch_prefill_paged.py` 中的 cases 一致，即： + +```python +for head_dim in [128]: + for batch_size in [1, 4, 16]: + for seq_len in [1024, 4096, 8192, 16384]: + ... +``` + +输出参与 `torch.allclose` 校验，容差为 `rtol=1e-2, atol=1e-2`。 \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/03_参考实现.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/03_参考实现.md new file mode 100644 index 0000000..b4417ae --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/03_参考实现.md @@ -0,0 +1,28 @@ +pytorch参考实现： + +```python +def baseline(q, kv_data, output, qo_indptr, kv_indptr, kv_indices, last_page_len, + batch_size, seq_len, num_qo_heads, num_kv_heads, + head_dim, page_block_size, causal): + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device) + wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( + workspace_buffer, + kv_layout="NHD", + backend="auto", + ) + wrapper.plan( + qo_indptr, + kv_indptr, + kv_indices, + last_page_len, + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, + causal=bool(causal), + q_data_type=torch.bfloat16, + kv_data_type=torch.bfloat16, + ) + wrapper.run(q, kv_data, out=output) +``` + diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/_title.txt b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/_title.txt new file mode 100644 index 0000000..1c5167b --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/_title.txt @@ -0,0 +1 @@ +FlashInfer Ragged Prefill diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/meta.json b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/meta.json new file mode 100644 index 0000000..cd4684c --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/meta.json @@ -0,0 +1,16 @@ +{ + "id": 194, + "displayId": 20002, + "type": "Traditional", + "isPublic": false, + "locales": [ + "zh_CN" + ], + "samples": [ + { + "inputData": "1\n", + "outputData": "" + } + ], + "problemTagIds": [] +} diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/testcase_config.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/testcase_config.py new file mode 100644 index 0000000..4066882 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/testcase_config.py @@ -0,0 +1,273 @@ +from __future__ import annotations + + +HEAD_DIMS = [128] +BATCH_SIZES = [1, 4, 16] +SEQ_LENS = [1024, 4096, 8192, 16384] +NUM_QO_HEADS = 32 +PAGE_BLOCK_SIZE = 16 +CAUSAL = 0 + + +def _build_cases(): + cases = [] + for head_dim in HEAD_DIMS: + for batch_size in BATCH_SIZES: + for seq_len in SEQ_LENS: + num_kv_heads = 8 if head_dim == 64 else 4 + cases.append( + ( + batch_size, + seq_len, + NUM_QO_HEADS, + num_kv_heads, + head_dim, + PAGE_BLOCK_SIZE, + CAUSAL, + ) + ) + return cases + + +TESTCASES = _build_cases() + + +def getNumOfTestcases() -> int: + return len(TESTCASES) + + +try: + from pathlib import Path + from typing import List, Tuple, Union + import sys + + import torch + + KernelArg = Union[torch.Tensor, int, float] + CURRENT_CASE = None + + def _ensure_flashinfer_importable(): + try: + import flashinfer # noqa: F401 + + return + except ImportError: + pass + + here = Path(__file__).resolve() + for parent in here.parents: + candidate = parent / "McFlashInfer" + if (candidate / "flashinfer").is_dir(): + sys.path.insert(0, str(candidate)) + return + + def _get_testcase_index() -> int: + try: + raw = input().strip() + except EOFError: + return 0 + if raw == "": + return 0 + try: + testcase_id = int(raw.split()[0]) + except ValueError: + return 0 + if 1 <= testcase_id <= len(TESTCASES): + return testcase_id - 1 + if 0 <= testcase_id < len(TESTCASES): + return testcase_id + return 0 + + def _compute_reps(batch_size: int, seq_len: int, head_dim: int, base_reps: int = 100) -> int: + workload = batch_size * seq_len * head_dim + if workload < 1e5: + return base_reps + if workload < 1e6: + return base_reps // 2 + if workload < 1e7: + return base_reps // 4 + if workload < 1e8: + return base_reps // 8 + if workload < 1e9: + return base_reps // 16 + return base_reps // 32 + + def _setup_paged_kv_indptr(batch_size: int, seq_len: int, page_block_size: int, device: str): + seq_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, device=device) + seq_lens_blocks = torch.div(seq_lens + page_block_size - 1, page_block_size, rounding_mode="floor") + kv_indptr = torch.empty((batch_size + 1,), dtype=torch.int32, device=device) + kv_indptr[0] = 0 + kv_indptr[1:] = torch.cumsum(seq_lens_blocks, dim=0) + num_blocks = int(kv_indptr[-1].item()) + last_page_len = (seq_lens - 1) % page_block_size + 1 + return kv_indptr, last_page_len, num_blocks + + def getTestCaseSize() -> Tuple[List[Tuple[int, ...]], Tuple[int, int]]: + testcase_id = _get_testcase_index() + global CURRENT_CASE + batch_size, seq_len, num_qo_heads, num_kv_heads, head_dim, page_block_size, causal = TESTCASES[testcase_id] + num_blocks = batch_size * ((seq_len + page_block_size - 1) // page_block_size) + qo_len = batch_size * seq_len + CURRENT_CASE = ( + batch_size, + seq_len, + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, + causal, + num_blocks, + 20260620 + testcase_id, + ) + warmup = 3 + iters = max(1, _compute_reps(batch_size, seq_len, head_dim)) + return [ + (qo_len, num_qo_heads, head_dim), + (num_blocks, 2, page_block_size, num_kv_heads, head_dim), + (qo_len, num_qo_heads, head_dim), + (batch_size + 1,), + (batch_size + 1,), + (num_blocks,), + (batch_size,), + (), (), (), (), (), (), (), + ], (warmup, iters) + + def genTestCase(testcase_sizes, device: str = "cuda") -> List[KernelArg]: + del testcase_sizes + ( + batch_size, + seq_len, + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, + causal, + num_blocks, + seed, + ) = CURRENT_CASE + gen = torch.Generator(device=device) + gen.manual_seed(seed) + dtype = torch.bfloat16 + qo_len = batch_size * seq_len + q = torch.rand(qo_len, num_qo_heads, head_dim, dtype=dtype, device=device, generator=gen).contiguous() + kv_data = torch.randn( + num_blocks, + 2, + page_block_size, + num_kv_heads, + head_dim, + dtype=dtype, + device=device, + generator=gen, + ).contiguous() + output = torch.empty(qo_len, num_qo_heads, head_dim, dtype=dtype, device=device) + qo_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device=device) * seq_len + kv_indptr, last_page_len, num_blocks_check = _setup_paged_kv_indptr( + batch_size, + seq_len, + page_block_size, + device, + ) + assert num_blocks_check == num_blocks + kv_indices = torch.arange(num_blocks, dtype=torch.int32, device=device) + return [ + q, + kv_data, + output, + qo_indptr, + kv_indptr, + kv_indices, + last_page_len, + batch_size, + seq_len, + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, + causal, + ] + + def baseline( + q, + kv_data, + output, + qo_indptr, + kv_indptr, + kv_indices, + last_page_len, + batch_size, + seq_len, + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, + causal, + ): + _ensure_flashinfer_importable() + import flashinfer + + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device) + wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( + workspace_buffer, + kv_layout="NHD", + backend="auto", + ) + wrapper.plan( + qo_indptr, + kv_indptr, + kv_indices, + last_page_len, + int(num_qo_heads), + int(num_kv_heads), + int(head_dim), + int(page_block_size), + causal=bool(causal), + q_data_type=torch.bfloat16, + kv_data_type=torch.bfloat16, + ) + wrapper.run(q, kv_data, out=output) + return [ + q, + kv_data, + output, + qo_indptr, + kv_indptr, + kv_indices, + last_page_len, + batch_size, + seq_len, + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, + causal, + ] + + def check( + testcase_sizes, + original_input_tensors, + target_kernel_input_tensors, + baseline_input_tensors, + rtol=1e-2, + atol=1e-2, + ) -> bool: + del testcase_sizes, original_input_tensors + output_t = target_kernel_input_tensors[2] + output_ref = baseline_input_tensors[2] + if output_t.shape != output_ref.shape: + print(f"[FAIL] shape mismatch: target {output_t.shape}, ref {output_ref.shape}", file=sys.stderr) + return False + if output_t.dtype != output_ref.dtype: + print(f"[FAIL] dtype mismatch: target {output_t.dtype}, ref {output_ref.dtype}", file=sys.stderr) + return False + if not torch.allclose(output_t.float(), output_ref.float(), rtol=rtol, atol=atol): + diff = (output_t.float() - output_ref.float()).abs() + print( + f"[FAIL] allclose failed: max_abs_diff={float(diff.max().item()):.6f}, " + f"mean_abs_diff={float(diff.mean().item()):.6f} (rtol={rtol}, atol={atol})", + file=sys.stderr, + ) + return False + return True +except Exception: + pass diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/00_题目描述.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/00_题目描述.md new file mode 100644 index 0000000..3e12eca --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/00_题目描述.md @@ -0,0 +1,22 @@ +--- +sectionTitle: "题目描述" +type: "Text" +--- +你需要实现 FlashInfer paged KV cache prefill 的 CUDA C++ 前向算子。 + +本题输入采用 FlashInfer `BatchPrefillWithPagedKVCacheWrapper` 的 paged `NHD` 布局。每个 batch 中有 `seq_len` 个 query token，KV cache 也有 `seq_len` 个 token，并按 page 存储。 + +评测程序会调用你提交代码中的 `run_kernel` 函数。你需要根据 `qo_indptr`、`kv_indptr`、`kv_indices` 和 `last_page_len` 读取 paged KV cache，并将结果写入 `output`。 + +baseline 使用 FlashInfer paged prefill 的 Python API： + +```python +wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(workspace, kv_layout="NHD", backend="auto") +wrapper.plan(qo_indptr, kv_indptr, kv_indices, last_page_len, + num_qo_heads, num_kv_heads, head_dim, page_block_size, + causal=bool(causal), + q_data_type=torch.bfloat16, kv_data_type=torch.bfloat16) +wrapper.run(q, kv_data, out=output) +``` + +如何提交代码详见[评测指南](/d/2)。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/01_接口约定.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/01_接口约定.md new file mode 100644 index 0000000..48fbe06 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/01_接口约定.md @@ -0,0 +1,42 @@ +--- +sectionTitle: "接口约定" +type: "codeSample" +lang: "cuda" +--- +你必须在提交的 CUDA 源码中提供如下 **C 符号**，函数名、参数类型、顺序必须完全一致，并使用 `extern "C"` 防止 name mangling： + +```cpp +#include +#include + +extern "C" void run_kernel( + const __nv_bfloat16* q, + const __nv_bfloat16* kv_data, + __nv_bfloat16* output, + const int32_t* qo_indptr, + const int32_t* kv_indptr, + const int32_t* kv_indices, + const int32_t* last_page_len, + int64_t batch_size, + int64_t seq_len, + int64_t num_qo_heads, + int64_t num_kv_heads, + int64_t head_dim, + int64_t page_block_size, + int64_t causal +); +``` + +### 参数说明 + +* `q`：query tensor，shape `(batch_size * seq_len, num_qo_heads, head_dim)`，连续 `bf16` +* `kv_data`：paged KV cache，shape `(num_blocks, 2, page_block_size, num_kv_heads, head_dim)`，连续 `bf16`，其中 `kv_data[:, 0]` 为 key，`kv_data[:, 1]` 为 value +* `output`：输出缓冲区，shape `(batch_size * seq_len, num_qo_heads, head_dim)`，连续 `bf16` +* `qo_indptr`：query/output indptr，shape `(batch_size + 1)`，连续 `int32` +* `kv_indptr`：paged KV indptr，shape `(batch_size + 1)`，连续 `int32` +* `kv_indices`：page index，shape `(num_blocks)`，连续 `int32` +* `last_page_len`：每个 batch 最后一个 page 的有效 token 数，shape `(batch_size)`，连续 `int32` +* `page_block_size`：page size，评测中固定为 `16` +* `causal`：是否启用 causal mask，本题按 benchmark case 固定为 `0` + +`run_kernel` 内部需要自行计算合适的 launch 配置并启动 CUDA kernel。为保证计时准确，不建议在 `run_kernel` 内部做 `cudaDeviceSynchronize()` 或显式同步。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/02_数据范围.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/02_数据范围.md new file mode 100644 index 0000000..1e57f2e --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/02_数据范围.md @@ -0,0 +1,26 @@ +--- +sectionTitle: "输入格式" +type: "Text" +--- +## 数据范围 + +- 数据类型：`q/kv_data/output` 均为 `bfloat16` +- KV layout：`NHD` +- `num_qo_heads = 32` +- `num_kv_heads = 4` +- `page_block_size = 16` +- `causal = 0` +- `head_dim` 固定为 `128` +- `batch_size` 取值为 `1, 4, 16` +- `seq_len` 取值为 `1024, 4096, 8192, 16384` + +测试点顺序与 `McFlashInfer/benchmarks/bench_batch_prefill_paged.py` 中的 cases 一致，即： + +```python +for head_dim in [128]: + for batch_size in [1, 4, 16]: + for seq_len in [1024, 4096, 8192, 16384]: + ... +``` + +输出参与 `torch.allclose` 校验，容差为 `rtol=1e-2, atol=1e-2`。 \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/03_参考实现.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/03_参考实现.md new file mode 100644 index 0000000..a55a635 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/03_参考实现.md @@ -0,0 +1,31 @@ +--- +sectionTitle: "输出格式" +type: "Text" +--- +PyTorch 参考实现 + +```python +def baseline(q, kv_data, output, qo_indptr, kv_indptr, kv_indices, last_page_len, + batch_size, seq_len, num_qo_heads, num_kv_heads, + head_dim, page_block_size, causal): + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device) + wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( + workspace_buffer, + kv_layout="NHD", + backend="auto", + ) + wrapper.plan( + qo_indptr, + kv_indptr, + kv_indices, + last_page_len, + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, + causal=bool(causal), + q_data_type=torch.bfloat16, + kv_data_type=torch.bfloat16, + ) + wrapper.run(q, kv_data, out=output) +``` \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/_title.txt b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/_title.txt new file mode 100644 index 0000000..34e4801 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/_title.txt @@ -0,0 +1 @@ +FlashInfer Paged Prefill diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/meta.json b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/meta.json new file mode 100644 index 0000000..f9593e3 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/meta.json @@ -0,0 +1,16 @@ +{ + "id": 195, + "displayId": 20003, + "type": "Traditional", + "isPublic": false, + "locales": [ + "zh_CN" + ], + "samples": [ + { + "inputData": "1\n", + "outputData": "" + } + ], + "problemTagIds": [] +} diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/testcase_config.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/testcase_config.py new file mode 100644 index 0000000..ee3ae71 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/testcase_config.py @@ -0,0 +1,245 @@ +from __future__ import annotations + + +HEAD_DIM_CKV = 512 +HEAD_DIM_KPE = 64 +NUM_HEADS_LIST = [64, 128] +BATCH_SIZES = [1, 4, 16] +SEQ_LENS = [1024, 4096, 8192, 16384] +PAGE_SIZE = 1 +CAUSAL = 0 + + +def _build_cases(): + cases = [] + for num_heads in NUM_HEADS_LIST: + for batch_size in BATCH_SIZES: + for seq_len in SEQ_LENS: + cases.append( + ( + batch_size, + seq_len, + num_heads, + HEAD_DIM_CKV, + HEAD_DIM_KPE, + PAGE_SIZE, + CAUSAL, + ) + ) + return cases + + +TESTCASES = _build_cases() + + +def getNumOfTestcases() -> int: + return len(TESTCASES) + + +try: + from pathlib import Path + from typing import List, Tuple, Union + import sys + + import torch + + KernelArg = Union[torch.Tensor, int, float] + CURRENT_CASE = None + + def _ensure_flashinfer_importable(): + try: + import flashinfer # noqa: F401 + + return + except ImportError: + pass + + here = Path(__file__).resolve() + for parent in here.parents: + candidate = parent / "McFlashInfer" + if (candidate / "flashinfer").is_dir(): + sys.path.insert(0, str(candidate)) + return + + def _get_testcase_index() -> int: + try: + raw = input().strip() + except EOFError: + return 0 + if raw == "": + return 0 + try: + testcase_id = int(raw.split()[0]) + except ValueError: + return 0 + if 1 <= testcase_id <= len(TESTCASES): + return testcase_id - 1 + if 0 <= testcase_id < len(TESTCASES): + return testcase_id + return 0 + + def _compute_reps(batch_size: int, seq_len: int, head_dim: int, base_reps: int = 100) -> int: + workload = batch_size * seq_len * head_dim + if workload < 1e5: + return base_reps + if workload < 1e6: + return base_reps // 2 + if workload < 1e7: + return base_reps // 4 + if workload < 1e8: + return base_reps // 8 + if workload < 1e9: + return base_reps // 16 + return base_reps // 32 + + def getTestCaseSize() -> Tuple[List[Tuple[int, ...]], Tuple[int, int]]: + testcase_id = _get_testcase_index() + global CURRENT_CASE + batch_size, seq_len, num_heads, head_dim_ckv, head_dim_kpe, page_size, causal = TESTCASES[testcase_id] + CURRENT_CASE = ( + batch_size, + seq_len, + num_heads, + head_dim_ckv, + head_dim_kpe, + page_size, + causal, + 20260630 + testcase_id, + ) + warmup = 3 + iters = max(1, _compute_reps(batch_size, seq_len, head_dim_ckv + head_dim_kpe)) + kv_len = batch_size * seq_len + return [ + (batch_size, num_heads, head_dim_ckv), + (batch_size, num_heads, head_dim_kpe), + (kv_len, 1, head_dim_ckv), + (kv_len, 1, head_dim_kpe), + (batch_size, num_heads, head_dim_ckv), + (batch_size + 1,), + (batch_size + 1,), + (kv_len,), + (batch_size,), + (), (), (), (), (), (), (), + ], (warmup, iters) + + def genTestCase(testcase_sizes, device: str = "cuda") -> List[KernelArg]: + del testcase_sizes + batch_size, seq_len, num_heads, head_dim_ckv, head_dim_kpe, page_size, causal, seed = CURRENT_CASE + gen = torch.Generator(device=device) + gen.manual_seed(seed) + dtype = torch.bfloat16 + kv_len = batch_size * seq_len + q_nope = torch.randn(batch_size, num_heads, head_dim_ckv, dtype=dtype, device=device, generator=gen).contiguous() + q_pe = torch.zeros(batch_size, num_heads, head_dim_kpe, dtype=dtype, device=device) + ckv = torch.randn(kv_len, 1, head_dim_ckv, dtype=dtype, device=device, generator=gen).contiguous() + kpe = torch.zeros(kv_len, 1, head_dim_kpe, dtype=dtype, device=device) + output = torch.empty(batch_size, num_heads, head_dim_ckv, dtype=dtype, device=device) + q_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device=device) + kv_indptr = q_indptr * seq_len + kv_indices = torch.arange(0, kv_len, dtype=torch.int32, device=device) + kv_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, device=device) + return [ + q_nope, + q_pe, + ckv, + kpe, + output, + q_indptr, + kv_indptr, + kv_indices, + kv_lens, + batch_size, + seq_len, + num_heads, + head_dim_ckv, + head_dim_kpe, + page_size, + causal, + ] + + def baseline( + q_nope, + q_pe, + ckv, + kpe, + output, + q_indptr, + kv_indptr, + kv_indices, + kv_lens, + batch_size, + seq_len, + num_heads, + head_dim_ckv, + head_dim_kpe, + page_size, + causal, + ): + _ensure_flashinfer_importable() + import flashinfer + + sm_scale = 1.0 / ((int(head_dim_ckv) + int(head_dim_kpe)) ** 0.5) + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q_nope.device) + wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(workspace_buffer, backend="auto") + wrapper.plan( + q_indptr, + kv_indptr, + kv_indices, + kv_lens, + int(num_heads), + int(head_dim_ckv), + int(head_dim_kpe), + int(page_size), + bool(causal), + sm_scale, + q_nope.dtype, + ckv.dtype, + ) + wrapper.run(q_nope, q_pe, ckv, kpe, out=output, return_lse=False) + return [ + q_nope, + q_pe, + ckv, + kpe, + output, + q_indptr, + kv_indptr, + kv_indices, + kv_lens, + batch_size, + seq_len, + num_heads, + head_dim_ckv, + head_dim_kpe, + page_size, + causal, + ] + + def check( + testcase_sizes, + original_input_tensors, + target_kernel_input_tensors, + baseline_input_tensors, + rtol=1e-2, + atol=1e-2, + ) -> bool: + del testcase_sizes, original_input_tensors + output_t = target_kernel_input_tensors[4] + output_ref = baseline_input_tensors[4] + if output_t.shape != output_ref.shape: + print(f"[FAIL] shape mismatch: target {output_t.shape}, ref {output_ref.shape}", file=sys.stderr) + return False + if output_t.dtype != output_ref.dtype: + print(f"[FAIL] dtype mismatch: target {output_t.dtype}, ref {output_ref.dtype}", file=sys.stderr) + return False + if not torch.allclose(output_t.float(), output_ref.float(), rtol=rtol, atol=atol): + diff = (output_t.float() - output_ref.float()).abs() + print( + f"[FAIL] allclose failed: max_abs_diff={float(diff.max().item()):.6f}, " + f"mean_abs_diff={float(diff.mean().item()):.6f} (rtol={rtol}, atol={atol})", + file=sys.stderr, + ) + return False + return True +except Exception: + pass diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/00_题目描述.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/00_题目描述.md new file mode 100644 index 0000000..c3bf938 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/00_题目描述.md @@ -0,0 +1,22 @@ +--- +sectionTitle: "题目描述" +type: "Text" +--- +你需要实现 FlashInfer MLA paged attention 的 CUDA C++ 前向算子。 + +本题使用 `BatchMLAPagedAttentionWrapper` 的 DeepSeek MLA decode 配置：`q_nope` 表示不含 RoPE 的 query，`q_pe` 表示 RoPE 部分的 query，`ckv` 与 `kpe` 分别表示 compressed KV cache 与 RoPE KV cache。MLA page size 固定为 `1`。 + +评测程序会调用你提交代码中的 `run_kernel` 函数。你需要根据 `q_indptr`、`kv_indptr`、`kv_indices` 和 `kv_lens` 读取 cache，并将结果写入 `output`。 + +baseline 使用 FlashInfer MLA 的 Python API： + +```python +wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(workspace, backend="auto") +wrapper.plan(q_indptr, kv_indptr, kv_indices, kv_lens, + num_heads, head_dim_ckv, head_dim_kpe, + page_size, False, sm_scale, + q_nope.dtype, ckv.dtype) +wrapper.run(q_nope, q_pe, ckv, kpe, out=output, return_lse=False) +``` + +如何提交代码详见[评测指南](/d/2)。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/01_接口约定.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/01_接口约定.md new file mode 100644 index 0000000..7cd07bb --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/01_接口约定.md @@ -0,0 +1,46 @@ +--- +sectionTitle: "接口约定" +type: "codeSample" +lang: "cuda" +--- +你必须在提交的 CUDA 源码中提供如下 **C 符号**，函数名、参数类型、顺序必须完全一致，并使用 `extern "C"` 防止 name mangling： + +```cpp +#include +#include + +extern "C" void run_kernel( + const __nv_bfloat16* q_nope, + const __nv_bfloat16* q_pe, + const __nv_bfloat16* ckv, + const __nv_bfloat16* kpe, + __nv_bfloat16* output, + const int32_t* q_indptr, + const int32_t* kv_indptr, + const int32_t* kv_indices, + const int32_t* kv_lens, + int64_t batch_size, + int64_t seq_len, + int64_t num_heads, + int64_t head_dim_ckv, + int64_t head_dim_kpe, + int64_t page_size, + int64_t causal +); +``` + +### 参数说明 + +* `q_nope`：query 的 compressed/nope 部分，shape `(batch_size, num_heads, head_dim_ckv)`，连续 `bf16` +* `q_pe`：query 的 RoPE 部分，shape `(batch_size, num_heads, head_dim_kpe)`，连续 `bf16` +* `ckv`：compressed KV cache，shape `(batch_size * seq_len, 1, head_dim_ckv)`，连续 `bf16` +* `kpe`：RoPE KV cache，shape `(batch_size * seq_len, 1, head_dim_kpe)`，连续 `bf16` +* `output`：输出缓冲区，shape `(batch_size, num_heads, head_dim_ckv)`，连续 `bf16` +* `q_indptr`：decode query indptr，shape `(batch_size + 1)`，内容为 `[0, 1, ..., batch_size]` +* `kv_indptr`：KV indptr，shape `(batch_size + 1)`，每段长度为 `seq_len` +* `kv_indices`：page index，shape `(batch_size * seq_len)`，连续 `int32` +* `kv_lens`：每个 batch 的 KV 长度，shape `(batch_size)`，连续 `int32` +* `page_size`：评测中固定为 `1` +* `causal`：评测中固定为 `0` + +`run_kernel` 内部需要自行计算合适的 launch 配置并启动 CUDA kernel。为保证计时准确，不建议在 `run_kernel` 内部做 `cudaDeviceSynchronize()` 或显式同步。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/02_数据范围.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/02_数据范围.md new file mode 100644 index 0000000..8919075 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/02_数据范围.md @@ -0,0 +1,26 @@ +--- +sectionTitle: "输入格式" +type: "Text" +--- +## 数据范围 + +- 数据类型：输入输出均为 `bfloat16` +- `head_dim_ckv = 512` +- `head_dim_kpe = 64` +- `page_size = 1` +- `causal = 0` +- `num_heads` 取值为 `64, 128` +- `batch_size` 取值为 `1, 4, 16` +- `seq_len` 取值为 `1024, 4096, 8192, 16384` +- `sm_scale = 1 / sqrt(head_dim_ckv + head_dim_kpe)` + +测试点顺序与 `McFlashInfer/benchmarks/bench_batch_mla.py` 中的 cases 一致，即： + +```python +for num_heads in [64, 128]: + for batch_size in [1, 4, 16]: + for seq_len in [1024, 4096, 8192, 16384]: + ... +``` + +输出参与 `torch.allclose` 校验，容差为 `rtol=1e-2, atol=1e-2`。 \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/03_参考实现.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/03_参考实现.md new file mode 100644 index 0000000..23285e8 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/03_参考实现.md @@ -0,0 +1,30 @@ +--- +sectionTitle: "输出格式" +type: "Text" +--- +PyTorch 参考实现 + +```python +def baseline(q_nope, q_pe, ckv, kpe, output, + q_indptr, kv_indptr, kv_indices, kv_lens, + batch_size, seq_len, num_heads, + head_dim_ckv, head_dim_kpe, page_size, causal): + sm_scale = 1.0 / ((head_dim_ckv + head_dim_kpe) ** 0.5) + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q_nope.device) + wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(workspace_buffer, backend="auto") + wrapper.plan( + q_indptr, + kv_indptr, + kv_indices, + kv_lens, + num_heads, + head_dim_ckv, + head_dim_kpe, + page_size, + bool(causal), + sm_scale, + q_nope.dtype, + ckv.dtype, + ) + wrapper.run(q_nope, q_pe, ckv, kpe, out=output, return_lse=False) +``` \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/_title.txt b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/_title.txt new file mode 100644 index 0000000..9f19bce --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/_title.txt @@ -0,0 +1 @@ +FlashInfer MLA Paged Attention diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/meta.json b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/meta.json new file mode 100644 index 0000000..035e3ac --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/meta.json @@ -0,0 +1,16 @@ +{ + "id": 196, + "displayId": 20004, + "type": "Traditional", + "isPublic": false, + "locales": [ + "zh_CN" + ], + "samples": [ + { + "inputData": "1\n", + "outputData": "" + } + ], + "problemTagIds": [] +} diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/testcase_config.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/testcase_config.py new file mode 100644 index 0000000..3273aa3 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/testcase_config.py @@ -0,0 +1,248 @@ +from __future__ import annotations + + +BATCH_SIZES = [1, 2, 4, 8, 16] +SEQ_LENS_KV = [512, 1024, 2048, 4096, 8192, 16384] +HEAD_DIMS = [128] +NUM_QO_HEADS = 32 +PAGE_BLOCK_SIZE = 16 + + +def _build_cases(): + cases = [] + for batch_size in BATCH_SIZES: + for seq_len_kv in SEQ_LENS_KV: + for head_dim in HEAD_DIMS: + num_kv_heads = 8 if head_dim == 64 else 4 + cases.append( + ( + batch_size, + seq_len_kv, + NUM_QO_HEADS, + num_kv_heads, + head_dim, + PAGE_BLOCK_SIZE, + ) + ) + return cases + + +TESTCASES = _build_cases() + + +def getNumOfTestcases() -> int: + return len(TESTCASES) + + +try: + from pathlib import Path + from typing import List, Tuple, Union + import sys + + import torch + + KernelArg = Union[torch.Tensor, int, float] + CURRENT_CASE = None + + def _ensure_flashinfer_importable(): + try: + import flashinfer # noqa: F401 + + return + except ImportError: + pass + + here = Path(__file__).resolve() + for parent in here.parents: + candidate = parent / "McFlashInfer" + if (candidate / "flashinfer").is_dir(): + sys.path.insert(0, str(candidate)) + return + + def _get_testcase_index() -> int: + try: + raw = input().strip() + except EOFError: + return 0 + if raw == "": + return 0 + try: + testcase_id = int(raw.split()[0]) + except ValueError: + return 0 + if 1 <= testcase_id <= len(TESTCASES): + return testcase_id - 1 + if 0 <= testcase_id < len(TESTCASES): + return testcase_id + return 0 + + def _compute_reps(batch_size: int, seq_len: int, head_dim: int, base_reps: int = 100) -> int: + workload = batch_size * seq_len * head_dim + if workload < 1e5: + return base_reps + if workload < 1e6: + return base_reps // 2 + if workload < 1e7: + return base_reps // 4 + if workload < 1e8: + return base_reps // 8 + if workload < 1e9: + return base_reps // 16 + return base_reps // 32 + + def _setup_paged_kv_indptr(batch_size: int, seq_len: int, page_block_size: int, device: str): + seq_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, device=device) + seq_lens_blocks = torch.div(seq_lens + page_block_size - 1, page_block_size, rounding_mode="floor") + kv_indptr = torch.empty((batch_size + 1,), dtype=torch.int32, device=device) + kv_indptr[0] = 0 + kv_indptr[1:] = torch.cumsum(seq_lens_blocks, dim=0) + num_blocks = int(kv_indptr[-1].item()) + last_page_len = (seq_lens - 1) % page_block_size + 1 + return kv_indptr, last_page_len, num_blocks + + def getTestCaseSize() -> Tuple[List[Tuple[int, ...]], Tuple[int, int]]: + testcase_id = _get_testcase_index() + global CURRENT_CASE + batch_size, seq_len_kv, num_qo_heads, num_kv_heads, head_dim, page_block_size = TESTCASES[testcase_id] + num_blocks = batch_size * ((seq_len_kv + page_block_size - 1) // page_block_size) + CURRENT_CASE = ( + batch_size, + seq_len_kv, + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, + num_blocks, + 20260710 + testcase_id, + ) + warmup = 3 + iters = max(1, _compute_reps(batch_size, seq_len_kv, head_dim)) + return [ + (batch_size, num_qo_heads, head_dim), + (num_blocks, 2, page_block_size, num_kv_heads, head_dim), + (batch_size, num_qo_heads, head_dim), + (batch_size + 1,), + (num_blocks,), + (batch_size,), + (), (), (), (), (), (), + ], (warmup, iters) + + def genTestCase(testcase_sizes, device: str = "cuda") -> List[KernelArg]: + del testcase_sizes + batch_size, seq_len_kv, num_qo_heads, num_kv_heads, head_dim, page_block_size, num_blocks, seed = CURRENT_CASE + gen = torch.Generator(device=device) + gen.manual_seed(seed) + dtype = torch.bfloat16 + q = torch.rand(batch_size, num_qo_heads, head_dim, dtype=dtype, device=device, generator=gen).contiguous() + kv_data = torch.randn( + num_blocks, + 2, + page_block_size, + num_kv_heads, + head_dim, + dtype=dtype, + device=device, + generator=gen, + ).contiguous() + output = torch.empty(batch_size, num_qo_heads, head_dim, dtype=dtype, device=device) + kv_indptr, last_page_len, num_blocks_check = _setup_paged_kv_indptr( + batch_size, + seq_len_kv, + page_block_size, + device, + ) + assert num_blocks_check == num_blocks + kv_indices = torch.arange(num_blocks, dtype=torch.int32, device=device) + return [ + q, + kv_data, + output, + kv_indptr, + kv_indices, + last_page_len, + batch_size, + seq_len_kv, + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, + ] + + def baseline( + q, + kv_data, + output, + kv_indptr, + kv_indices, + last_page_len, + batch_size, + seq_len_kv, + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, + ): + _ensure_flashinfer_importable() + import flashinfer + + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device) + wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper( + workspace_buffer, + kv_layout="NHD", + use_tensor_cores=True, + ) + wrapper.plan( + kv_indptr, + kv_indices, + last_page_len, + int(num_qo_heads), + int(num_kv_heads), + int(head_dim), + int(page_block_size), + data_type=torch.bfloat16, + q_data_type=torch.bfloat16, + ) + wrapper.run(q, kv_data, out=output) + return [ + q, + kv_data, + output, + kv_indptr, + kv_indices, + last_page_len, + batch_size, + seq_len_kv, + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, + ] + + def check( + testcase_sizes, + original_input_tensors, + target_kernel_input_tensors, + baseline_input_tensors, + rtol=1e-2, + atol=1e-2, + ) -> bool: + del testcase_sizes, original_input_tensors + output_t = target_kernel_input_tensors[2] + output_ref = baseline_input_tensors[2] + if output_t.shape != output_ref.shape: + print(f"[FAIL] shape mismatch: target {output_t.shape}, ref {output_ref.shape}", file=sys.stderr) + return False + if output_t.dtype != output_ref.dtype: + print(f"[FAIL] dtype mismatch: target {output_t.dtype}, ref {output_ref.dtype}", file=sys.stderr) + return False + if not torch.allclose(output_t.float(), output_ref.float(), rtol=rtol, atol=atol): + diff = (output_t.float() - output_ref.float()).abs() + print( + f"[FAIL] allclose failed: max_abs_diff={float(diff.max().item()):.6f}, " + f"mean_abs_diff={float(diff.mean().item()):.6f} (rtol={rtol}, atol={atol})", + file=sys.stderr, + ) + return False + return True +except Exception: + pass diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/00_题目描述.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/00_题目描述.md new file mode 100644 index 0000000..74508b0 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/00_题目描述.md @@ -0,0 +1,22 @@ +--- +sectionTitle: "题目描述" +type: "Text" +--- +你需要实现 FlashInfer paged KV cache decode 的 CUDA C++ 前向算子。 + +本题输入采用 FlashInfer `BatchDecodeWithPagedKVCacheWrapper` 的 paged `NHD` 布局。每个 batch 只有 1 个 query token，KV cache 长度为 `seq_len_kv`。 + +评测程序会调用你提交代码中的 `run_kernel` 函数。你需要根据 `kv_indptr`、`kv_indices` 和 `last_page_len` 读取 paged KV cache，并将结果写入 `output`。 + +baseline 使用 FlashInfer paged decode 的 Python API： + +```python +wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper( + workspace, kv_layout="NHD", use_tensor_cores=True) +wrapper.plan(kv_indptr, kv_indices, last_page_len, + num_qo_heads, num_kv_heads, head_dim, page_block_size, + data_type=torch.bfloat16, q_data_type=torch.bfloat16) +wrapper.run(q, kv_data, out=output) +``` + +如何提交代码详见[评测指南](/d/2)。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/01_接口约定.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/01_接口约定.md new file mode 100644 index 0000000..c498b62 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/01_接口约定.md @@ -0,0 +1,38 @@ +--- +sectionTitle: "接口约定" +type: "codeSample" +lang: "cuda" +--- +你必须在提交的 CUDA 源码中提供如下 **C 符号**，函数名、参数类型、顺序必须完全一致，并使用 `extern "C"` 防止 name mangling： + +```cpp +#include +#include + +extern "C" void run_kernel( + const __nv_bfloat16* q, + const __nv_bfloat16* kv_data, + __nv_bfloat16* output, + const int32_t* kv_indptr, + const int32_t* kv_indices, + const int32_t* last_page_len, + int64_t batch_size, + int64_t seq_len_kv, + int64_t num_qo_heads, + int64_t num_kv_heads, + int64_t head_dim, + int64_t page_block_size +); +``` + +### 参数说明 + +* `q`：decode query tensor，shape `(batch_size, num_qo_heads, head_dim)`，连续 `bf16` +* `kv_data`：paged KV cache，shape `(num_blocks, 2, page_block_size, num_kv_heads, head_dim)`，连续 `bf16`，其中 `kv_data[:, 0]` 为 key，`kv_data[:, 1]` 为 value +* `output`：输出缓冲区，shape `(batch_size, num_qo_heads, head_dim)`，连续 `bf16` +* `kv_indptr`：paged KV indptr，shape `(batch_size + 1)`，连续 `int32` +* `kv_indices`：page index，shape `(num_blocks)`，连续 `int32` +* `last_page_len`：每个 batch 最后一个 page 的有效 token 数，shape `(batch_size)`，连续 `int32` +* `page_block_size`：page size，评测中固定为 `16` + +`run_kernel` 内部需要自行计算合适的 launch 配置并启动 CUDA kernel。为保证计时准确，不建议在 `run_kernel` 内部做 `cudaDeviceSynchronize()` 或显式同步。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/02_数据范围.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/02_数据范围.md new file mode 100644 index 0000000..9fe7bee --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/02_数据范围.md @@ -0,0 +1,26 @@ +--- +sectionTitle: "输入格式" +type: "Text" +--- +## 数据范围 + +- 数据类型：`q/kv_data/output` 均为 `bfloat16` +- KV layout：`NHD` +- `seq_len_q = 1` +- `num_qo_heads = 32` +- `page_block_size = 16` +- `head_dim` 固定为 `128` +- `num_kv_heads = 4` +- `batch_size` 取值为 `1, 2, 4, 8, 16` +- `seq_len_kv` 取值为 `512, 1024, 2048, 4096, 8192, 16384` + +测试点顺序与 `McFlashInfer/benchmarks/bench_batch_decode.py` 中的 cases 一致，即： + +```python +for batch_size in [1, 2, 4, 8, 16]: + for seq_len_kv in [512, 1024, 2048, 4096, 8192, 16384]: + for head_dim in [128]: + ... +``` + +输出参与 `torch.allclose` 校验，容差为 `rtol=1e-2, atol=1e-2`。 \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/03_参考实现.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/03_参考实现.md new file mode 100644 index 0000000..5d17935 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/03_参考实现.md @@ -0,0 +1,29 @@ +--- +sectionTitle: "输出格式" +type: "Text" +--- +PyTorch 参考实现 + +```python +def baseline(q, kv_data, output, kv_indptr, kv_indices, last_page_len, + batch_size, seq_len_kv, num_qo_heads, + num_kv_heads, head_dim, page_block_size): + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device) + wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper( + workspace_buffer, + kv_layout="NHD", + use_tensor_cores=True, + ) + wrapper.plan( + kv_indptr, + kv_indices, + last_page_len, + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, + data_type=torch.bfloat16, + q_data_type=torch.bfloat16, + ) + wrapper.run(q, kv_data, out=output) +``` \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/_title.txt b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/_title.txt new file mode 100644 index 0000000..552b751 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/_title.txt @@ -0,0 +1 @@ +FlashInfer Paged Decode