Merge pull request 'moe冒烟代码更新' (#63) from xxhefei/op_optimization:update-moark-agent-doc into master

缩减prompt篇幅，显示效果完全适配gitlink规则

README.md

@@ -18,12 +18,14 @@
 
 初赛攻坚 Fused Moe Gemm、决赛冲刺 MLA/NSA，全程基于全国产技术栈开发优化，打造自主可控的低成本大模型推理算子底座。
 
+
 **赛题一相关资料**
 
 - [赛题一方案：基于国产软件栈的推理前沿算子优化比赛方案](基于国产软件栈大模型推理前沿算子优化/基于国产软件栈的推理前沿算子优化比赛方案.md)
 - [TileLang MACA 构建指南（模力方舟）](基于国产软件栈大模型推理前沿算子优化/tilelang_maca_build_guide_模力方舟.md)
 - [TileLang 算子 sample 跑通测试指南](基于国产软件栈大模型推理前沿算子优化/race_tests_run_guide基于tilelang算子sample跑通测试.md)
-- [赛题一国产软件栈样例说明](基于国产软件栈大模型推理前沿算子优化/race-1-domestic-software-stack-samples.md)
+- [赛题一国产软件栈样例说明](基于国产软件栈大模型推理前沿算子优化/基于国产软件栈的大模型推理前沿算子优化（TileLang）Sample%20算子获取说明.md)
+- [TileLang 训练营学习资料](https://www.gitlink.org.cn/ccf-ai-infra/Intro-ops/about)
 
 ### 赛题二：基于 AI Agent 开发范式的国产 GPU 大模型推理算子库优化
 
@@ -44,9 +46,12 @@
 **赛题二相关资料**
 
 - [赛题二方案：基于 AI Agent 开发范式的国产 GPU 大模型推理算子库优化方案](基于AI%20Agent开发范式的国产GPU大模型推理算子库优化/基于AI%20Agent开发范式的国产GPU大模型算子推理库优化方案.md)
+- [赛题二选手入口](基于AI%20Agent开发范式的国产GPU大模型推理算子库优化/选手入口.md)
 - [模力方舟 Agent 部署准备教程](基于AI%20Agent开发范式的国产GPU大模型推理算子库优化/模力方舟Agent部署准备教程.md)
 - [赛题二说明及资料参考](基于AI%20Agent开发范式的国产GPU大模型推理算子库优化/赛题说明.md)
 
+###**两个赛题统一使用模力方舟上的镜像PyTorch-Agent / 2.8.0 / Python 3.12 / maca 3.7.2.1**
+
 ## 参赛对象
 
 面向全国全日制专科、本科、硕博在校生（非在职），40 周岁以下青年科技人才均可参与；
@@ -99,6 +104,8 @@
 
 沐曦开发者社区活动页面，完成新人礼任务，提前熟悉 C500 在线算力使用。https://developer.metax-tech.com/activities/6
 
+学生可通过登录启悟社区后跳转沐曦开发者社区领取100元算力代金券。(https://developer.metax-tech.com/activities/11)
+
 报名赛事后还可额外获得300元算力代金券：https://developer.metax-tech.com/activities/17
 
 使用指南：[模力方舟快速使用 SOP](模力方舟快速使用SOP.md)
@@ -109,7 +116,7 @@
 
 加入沐曦股份“揭榜挂帅”赛题，深耕国产 GPU 算子优化，用技术降低每 Token 推理成本，共建自主可控 AI 算力新生态！
 
-👇 报名通道：登录挑战杯官网 [www.tiaozhanbei.net](https://link.wtturl.cn/?target=https%3A%2F%2Fwww.tiaozhanbei.net&scene=im&aid=497858&lang=zh) 揭榜挂帅入口报名
+👇 报名通道：登录挑战杯官网[2026年度中国青年科技创新“揭榜挂帅”擂台赛（学生赛道）](https://2026.tiaozhanbei.net/)揭榜挂帅入口报名
 
 
 

基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer关键算子迁移与优化.md

@@ -1,356 +0,0 @@
-# FlashInfer 迁移 Baseline 实战
-
-## 一、教程定位
-
-本教程是参赛训练课程的 **FlashInfer Baseline 入门**模块，主要帮助用户快速跑通 FlashInfer 的最小可运行流程。完成本教程后，用户应能够完成源码编译、API 调用、正确性测试和 Benchmark 测试，并记录一份 baseline 性能结果，为后续算子优化提供对比基准。
-
-## 二、学习目标
-
-1.  理解 FlashInfer Attention Kernel 的基本作用与适用场景；
-2.  完成 FlashInfer 环境、工具链的准备与源码编译；
-3.  跑通 BatchDecode、BatchPrefill、MLA 等典型算子的 API 调用示例；
-4.  完成各算子在不同参数配置下的 Benchmark 测试；
-5.  输出各算子的 Baseline 性能结果记录表，为后续算子优化提供对比基准。
-
-
-## 三、适用对象
-
-适合人群
-
-*   参赛选手：需要完成 Baseline 入门模块，为后续算子优化做准备
-*   软件开发者和Vibe Coding开发：希望从事AI相关行业开发，以及用智能体方式来做开发工作
-*   LLM 推理开发者：希望了解 FlashInfer Attention Kernel 的性能表现
-*   算子优化工程师：希望基于MXMACA软件栈在沐曦国产 GPU 上做算子迁移和优化
-
-前置基础
-
-*   Python 基础：能够运行和修改 Python 脚本
-*   PyTorch 基础：了解MXMACA软化栈的使用
-*   Linux 命令行：能够使用终端执行命令
-*   了解 Attention 机制：理解 Q/K/V、KV Cache 等基本概念
-    
-
-## 四、前置准备
-
-开始实战前，请确认你已经完成以下准备：
-
-### GPU准备
-
-* 步骤1：获取算力券[https://developer.metax-tech.com/activities/6](https://developer.metax-tech.com/activities/6)
-
-  *   首次登录需要先进行注册（使用邮箱或者手机号进行注册）
-      
-*   登录成功后进行第二步-邮箱验证，填入自己的邮箱。
-      
-  *   第三步，提交申请。
-    
-  *   获得兑换码
-      
-
-* 步骤2：兑换算力和登陆平台
-
-  *   访问模力方舟官网：[https://ai.gitee.com/](https://ai.gitee.com/)
-      
-  *   进入费用中心 - 算力券 ，点击右上角 “兑换”。
-      
-
-*   步骤3：租用算力
-    
-    - 进入算力容器，选择沐曦，租用算力，建议优先选16G显存/32G显存，如下图：
-*   步骤4：创建实例
-
-    - 基础镜像：maca-pytorch:3.7.1.5-torch2.8-py312-ubuntu24.04-amd64
-*   步骤5：选择工具-lab进入实例环境
-
-
-*   步骤6：在JupyterLab Terminal中检查运行环境的配置，确认沐曦 GPU 可见--可以使用`mx-smi`命令查看
-
-
-### 环境依赖准备
-
-```bash
-pip install flashinfer torch pandas numpy
-```
-
-### opencode Agent安装
-
-```shell
-curl -fsSL https://opencode.ai/install | bash
-```
-
-说明：
-
-自动下载最新版本 自动配置环境 适用于 macOS / Linux / WSL 安装完成后，还会提示你怎么用，OpenCode 包含免费模式，使用方式：
-
-```shell
-cd   # 进入项目目录
-opencode
-```
-
-### 代码准备
-
-*   已进入指定项目目录；
-    
-*   准备测试脚本和 Benchmark 脚本。
-    
-
-| 命令 | 说明 |
-| --- | --- |
-| `python bench_batch_decode.py` | 运行 Batch Decode 基准测试 |
-| `python bench_batch_prefill_paged.py` | 运行 Batch Prefill (Paged KV Cache) 基准测试 |
-| `python bench_batch_prefill_ragged.py` | 运行 Batch Prefill (Ragged KV Cache) 基准测试 |
-| `python bench_batch_mla.py` | 运行 MLA (Multi-head Latent Attention) 基准测试 |
-
-## 五、知识预备
-
-#### LLM推理阶段重要概念：
-
-*   **Prefill阶段**：prefill 阶段是指处理输入 prompt 的阶段
-    
-    *   输入：用户一次性给出的完整 prompt，长度为 seq\_len
-        
-    *   计算：对 prompt 中的每个 token 并行计算注意力，生成第一个输出 token 及 KV cache
-        
-    *   特点：这是**计算密集型（compute-bound）**阶段，因为需要做完整的 seq\_len x seq\_len 注意力矩阵乘法
-    
-*   **decode阶段**：
-    
-    *   每次只生成 1 个 token，利用 prefill 阶段填充好的 KV cache 做自回归生成
-        
-    *   **显存带宽密集型（memory-bound**），瓶颈在从显存读取 KV cache 而非计算
-        
-
-prefill = 并行处理用户输入，decode = 逐个生成回答 token
-
-## 六、项目实践--FlashInfer-Baseline
-
-### Step 1：检查运行环境
-
-**目标：** 进入terminal确认当前环境满足本模块运行要求。
-
-**操作：** 检查 GPU、Python、编译工具和依赖版本。
-
-![giteeai 实例 12](https://origin.picgo.net/2026/06/04/giteeai--12c1772b12867f6be0.png)
-
-**常见问题：**
-
-| 问题 | 解决方法 |
-| --- | --- |
-| `mx-smi: command not found` | 使用 `python -c "import torch; print(torch.cuda.get_device_name(0))"` |
-| `No GPUs are available` | 检查 MXMACA 驱动是否正确安装 |
-| `ModuleNotFoundError: No module named 'flashinfer'` | 执行 `pip install flashinfer` 安装 |
-
-### 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)。
-
-1. 克隆代码仓库
-
-   ```Bash
-   git clone https://gitlink.org.cn/metax-maca/op_optimization.git
-   ```
-
-2. 准备flashinfer_baseline
-
-   在仓库目录 `op_optimization/基于AI Agent开发范式的国产GPU大模型推理算子库优化` 下，找到 `flashinfer_baseline` 文件夹。可以将 `flashinfer_baseline` 整个目录复制到工作目录 `data/` 下。
-
-**下一步操作：** 切换到 FlashInfer Baseline 项目目录。
-
-**命令示例：**
-
-```bash
-cd /data/flashinfer_baseline
-ls -la
-```
-
-**预期结果：**
-
-```plaintext
-bench_common.py
-bench_batch_decode.py
-bench_batch_prefill_paged.py
-bench_batch_prefill_ragged.py
-bench_batch_mla.py
-README.md
-...
-
-```
-
-### Step 3：验证项目脚本
-
-**目标：** 确认所有基准测试脚本可正常执行。
-
-**操作：** 检查脚本文件是否存在且可读。
-
-**命令示例：**
-
-```bash
-# 检查脚本文件
-python -c "import os; scripts = ['bench_common.py', 'bench_batch_decode.py', 'bench_batch_prefill_paged.py', 'bench_batch_prefill_ragged.py', 'bench_batch_mla.py']; [print(f'✓ {s}') if os.path.exists(s) else print(f'✗ {s} missing') for s in scripts]"
-
-# 测试脚本导入
-python -c "from bench_common import setup_workspace, get_csv_path; print('脚本导入正常')"
-
-```
-
-**预期结果：**
-
-```plaintext
-✓ bench_common.py
-✓ bench_batch_decode.py
-✓ bench_batch_prefill_paged.py
-✓ bench_batch_prefill_ragged.py
-✓ bench_batch_mla.py
-脚本导入正常
-
-```
-
-### Step 4：运行单算子 Benchmark
-
-**目标：** 执行基准测试，获取 Baseline 性能数据。
-
-**操作：** 运行 BatchDecode 基准测试脚本。
-
-**命令示例：**
-
-```bash
-python bench_batch_decode.py
-
-```
-
-**预期结果：**
-
-```plaintext
-[BatchDecodeWithPagedKVCacheWrapper] Starting benchmark, total cases: 144
-  [1/144] bs=1, kv_len=512, hd=64: 0.020ms, 51.53 GB/s, 0.82 TFLOPs
-  [2/144] bs=1, kv_len=512, hd=128: 0.023ms, 45.28 GB/s, 0.72 TFLOPs
-  ...
-
-Results saved to BatchDecodeWithPagedKVCacheWrapper_20260525_xxxxxx.csv
-
-```
-
-**常见问题：**
-
-| 问题 | 解决方法 |
-| --- | --- |
-| `out of memory` | 减小 batch\_size 或 seq\_len 参数 |
-| 运行时间过长 | 脚本会自动调整重复次数，耐心等待 |
-
-### Step 5：查看测试结果
-
-**目标：** 查看并分析 Benchmark 输出结果。
-
-**操作：** 读取生成的 CSV 结果文件。
-
-**命令示例：**
-
-```bash
-# 查看最新生成的 CSV 文件
-ls -lt *.csv | head -1
-
-# 使用 Python 查看结果
-python -c "
-import pandas as pd
-import glob
-csv_files = glob.glob('BatchDecodeWithPagedKVCacheWrapper_*.csv')
-if csv_files:
-    df = pd.read_csv(max(csv_files))
-    print(df[['batch_size', 'seq_len_kv', 'head_dim', 'time_ms', 'bandwidth_GB_s', 'tflops']].head(10))
-"
-
-```
-
-**预期结果：**
-
-```plaintext
-   batch_size  seq_len_kv  head_dim   time_ms  bandwidth_GB_s    tflops
-0           1         512        64  0.020429       51.528822  0.821253
-1           1         512       128  0.023340       45.278052  0.718833
-2           1         512       256  0.031949       66.153846  1.050256
-...
-
-```
-
-
-
-## 七、常见问题
-
-### 环境相关问题
-
-| 问题 | 原因 | 解决办法 |
-| --- | --- | --- |
-| `No GPUs are available` | MXMACA 驱动未安装或 GPU 不可见 | 检查驱动安装，运行 `python -c "import torch; print(torch.cuda.device_count())"` 验证 |
-| `ModuleNotFoundError: No module named 'flashinfer'` | flashinfer 未安装 | 执行 `pip install flashinfer` |
-| `out of memory` | GPU 显存不足 | 减小 `batch_size` 或 `seq_len` 参数 |
-
-### 运行相关问题
-
-| 问题 | 原因 | 解决办法 |
-| --- | --- | --- |
-| Benchmark 运行时间过长 | 参数组合过多， workload 较大 | 耐心等待，脚本会自动调整重复次数 |
-| `KeyError: 'BatchPrefillWithPagedKVCacheKernel'` | profiler 未捕获目标 kernel | 检查 `target_kernels` 配置是否正确 |
-| CSV 文件为空 | 测试未正常完成 | 检查 GPU 显存是否充足，重新运行 |
-
-### 代码相关问题
-
-| 问题 | 原因 | 解决办法 |
-| --- | --- | --- |
-| `ImportError: cannot import name 'xxx' from 'bench_common'` | 函数名拼写错误 | 检查 `bench_common.py` 中的函数名 |
-| `RuntimeError: error: device-side assert triggered` | 输入参数超出范围 | 检查 `num_qo_heads`、`num_kv_heads`、`head_dim` 配置 |
-
-### 性能相关问题
-
-| 问题 | 原因 | 解决办法 |
-| --- | --- | --- |
-| TFLOPs 数值异常低 | 工作负载过小，kernel 启动开销占比大 | 增大 `batch_size` 或 `seq_len` |
-| 带宽数值异常低 | 数据未正确加载到 GPU | 检查 Tensor 是否在 CUDA 设备上 |
-
-## 八、下一步学习建议
-
-### 1. 保存你的 Baseline 结果
-
-将本次运行生成的 CSV 文件妥善保存，后续优化时需要以此作为对比基准。
-
-```bash
-# 建议创建 results 目录保存
-mkdir -p results
-mv *.csv results/
-
-```
-
-### 2. 深入理解 FlashInfer 核心概念
-
-*   阅读 FlashInfer 官方文档，理解 Paged KV Cache、Ragged KV Cache 的设计理念
-    
-*   学习 MLA (Multi-head Latent Attention) 的原理，了解 DeepSeek 的注意力优化方案
-    
-*   理解 `plan()` 和 `run()` 两阶段设计的作用
-    
-
-### 3. 进入算子优化模块
-
-参考后续优化模块，学习以下优化技术：
-
-*   **Kernel Tuning**：调整 Block Size、Thread Count 等参数
-    
-*   **Memory Optimization**：减少显存占用、优化数据搬运
-    
-*   **Compute Optimization**：提升计算效率
-    
-
-### 4. 参考资源
-
-*   FlashInfer 官方仓库：https://github.com/flashinfer-ai/flashinfer
-    
-*   FlashInfer 文档：https://flashinfer.ai
-    
-
-### 5. 记录优化过程
-
-建议维护一份优化日志，记录每次优化的改动和性能变化：
-
-| 优化项 | 改动内容 | Baseline | 优化后 | 提升比例 |
-| --- | --- | --- | --- | --- |
-| 例：调整 block\_size | 16 → 32 | xx ms | xx ms | xx% |
-| 完成优化后，再次运行本模块的 Benchmark 脚本，对比前后性能变化。 |  |  |  |  |

基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/benchmark结果实例/benchmark_kvcache_headdim160.csv

@@ -0,0 +1,49 @@
+batch_size,seq_len_kv,heads,headdim,time_ms,bandwidth_GB_s
+1,512,8,160,0.0580,45.25
+2,512,8,160,0.0611,85.84
+4,512,8,160,0.0656,159.96
+8,512,8,160,0.0699,300.40
+16,512,8,160,0.1321,317.92
+32,512,8,160,0.2002,419.52
+64,512,8,160,0.3383,496.43
+128,512,8,160,0.6669,503.61
+1,1024,8,160,0.1129,46.45
+2,1024,8,160,0.1190,88.18
+4,1024,8,160,0.1224,171.49
+8,1024,8,160,0.1287,326.07
+16,1024,8,160,0.2479,338.49
+32,1024,8,160,0.3767,445.54
+64,1024,8,160,0.6419,523.01
+128,1024,8,160,1.2804,524.37
+1,2048,8,160,0.2270,46.20
+2,2048,8,160,0.2299,91.22
+4,2048,8,160,0.2349,178.63
+8,2048,8,160,0.2447,342.96
+16,2048,8,160,0.4773,351.60
+32,2048,8,160,0.7279,461.07
+64,2048,8,160,1.2559,534.49
+128,2048,8,160,2.5613,524.15
+1,4096,8,160,0.4460,47.02
+2,4096,8,160,0.4513,92.94
+4,4096,8,160,0.4593,182.64
+8,4096,8,160,0.4813,348.64
+16,4096,8,160,0.9363,358.43
+32,4096,8,160,1.4552,461.21
+64,4096,8,160,2.5615,524.05
+128,4096,8,160,5.1420,522.11
+1,8192,8,160,0.8847,47.41
+2,8192,8,160,0.8944,93.80
+4,8192,8,160,0.9094,184.51
+8,8192,8,160,0.9625,348.64
+16,8192,8,160,1.8550,361.80
+32,8192,8,160,2.9567,453.97
+64,8192,8,160,5.1398,522.30
+128,8192,8,160,10.2972,521.41
+1,16384,8,160,1.7608,47.64
+2,16384,8,160,1.7786,94.33
+4,16384,8,160,1.8143,184.95
+8,16384,8,160,1.9317,347.42
+16,16384,8,160,3.7301,359.83
+32,16384,8,160,5.9216,453.33
+64,16384,8,160,10.2668,522.94
+128,16384,8,160,20.6062,521.09

基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/benchmark结果实例/benchmark_kvcache_headdim192.csv

@@ -0,0 +1,49 @@
+batch_size,seq_len_kv,heads,headdim,time_ms,bandwidth_GB_s
+1,512,8,192,0.0458,68.82
+2,512,8,192,0.0515,122.32
+4,512,8,192,0.0574,219.28
+8,512,8,192,0.0607,414.80
+16,512,8,192,0.1147,439.27
+32,512,8,192,0.1763,571.40
+64,512,8,192,0.2978,676.79
+128,512,8,192,0.5874,686.11
+1,1024,8,192,0.0946,66.55
+2,1024,8,192,0.1033,121.85
+4,1024,8,192,0.1073,234.66
+8,1024,8,192,0.1131,445.23
+16,1024,8,192,0.2165,465.24
+32,1024,8,192,0.3347,601.80
+64,1024,8,192,0.5701,706.63
+128,1024,8,192,1.1302,712.88
+1,2048,8,192,0.1943,64.79
+2,2048,8,192,0.1992,126.38
+4,2048,8,192,0.2059,244.52
+8,2048,8,192,0.2174,463.13
+16,2048,8,192,0.4202,479.24
+32,2048,8,192,0.6503,619.36
+64,2048,8,192,1.1158,721.93
+128,2048,8,192,2.2250,724.05
+1,4096,8,192,0.3834,65.65
+2,4096,8,192,0.3904,128.95
+4,4096,8,192,0.4043,249.04
+8,4096,8,192,0.4267,471.92
+16,4096,8,192,0.8271,486.90
+32,4096,8,192,1.2840,627.28
+64,4096,8,192,2.2148,727.29
+128,4096,8,192,4.3819,735.21
+1,8192,8,192,0.7566,66.52
+2,8192,8,192,0.7712,130.54
+4,8192,8,192,0.7974,252.49
+8,8192,8,192,0.8433,477.47
+16,8192,8,192,1.6433,490.09
+32,8192,8,192,2.5573,629.84
+64,8192,8,192,4.3785,735.73
+128,8192,8,192,8.7303,737.99
+1,16384,8,192,1.5068,66.81
+2,16384,8,192,1.5350,131.16
+4,16384,8,192,1.5868,253.76
+8,16384,8,192,1.6778,479.99
+16,16384,8,192,3.2750,491.81
+32,16384,8,192,5.0659,635.88
+64,16384,8,192,8.7435,736.85
+128,16384,8,192,17.5040,736.13

基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/benchmark结果实例/benchmark_kvcache_headdim224.csv

@@ -0,0 +1,49 @@
+batch_size,seq_len_kv,heads,headdim,time_ms,bandwidth_GB_s
+1,512,8,224,0.1254,29.29
+2,512,8,224,0.1412,52.05
+4,512,8,224,0.1497,98.13
+8,512,8,224,0.1533,191.70
+16,512,8,224,0.1913,307.17
+32,512,8,224,0.3292,357.08
+64,512,8,224,0.5187,453.28
+128,512,8,224,0.9522,493.84
+1,1024,8,224,0.2727,26.93
+2,1024,8,224,0.2836,51.78
+4,1024,8,224,0.2890,101.63
+8,1024,8,224,0.2959,198.55
+16,1024,8,224,0.3696,317.93
+32,1024,8,224,0.6408,366.75
+64,1024,8,224,1.0081,466.21
+128,1024,8,224,1.8548,506.78
+1,2048,8,224,0.5515,26.63
+2,2048,8,224,0.5575,52.67
+4,2048,8,224,0.5666,103.65
+8,2048,8,224,0.5803,202.42
+16,2048,8,224,0.7250,324.05
+32,2048,8,224,1.2593,373.14
+64,2048,8,224,1.9890,472.48
+128,2048,8,224,3.6905,509.28
+1,4096,8,224,1.0939,26.84
+2,4096,8,224,1.1044,53.18
+4,4096,8,224,1.1219,104.69
+8,4096,8,224,1.1500,204.26
+16,4096,8,224,1.4390,326.48
+32,4096,8,224,2.4992,375.97
+64,4096,8,224,4.0082,468.86
+128,4096,8,224,7.3372,512.26
+1,8192,8,224,2.1775,26.97
+2,8192,8,224,2.1989,53.41
+4,8192,8,224,2.2338,105.15
+8,8192,8,224,2.3268,201.90
+16,8192,8,224,2.8806,326.18
+32,8192,8,224,5.0187,374.43
+64,8192,8,224,8.0323,467.90
+128,8192,8,224,14.6300,513.78
+1,16384,8,224,4.3360,27.09
+2,16384,8,224,4.3820,53.60
+4,16384,8,224,4.5006,104.38
+8,16384,8,224,4.6987,199.96
+16,16384,8,224,5.7361,327.59
+32,16384,8,224,10.1291,371.03
+64,16384,8,224,16.0745,467.60
+128,16384,8,224,OOM,OOM

基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/benchmark结果实例/benchmark_kvcache_headdim256.csv

@@ -0,0 +1,49 @@
+batch_size,seq_len_kv,heads,headdim,time_ms,bandwidth_GB_s
+1,512,8,256,0.0877,47.89
+2,512,8,256,0.0921,91.17
+4,512,8,256,0.0940,178.74
+8,512,8,256,0.0964,348.52
+16,512,8,256,0.1450,463.27
+32,512,8,256,0.2250,597.21
+64,512,8,256,0.3609,744.43
+128,512,8,256,0.6932,775.25
+1,1024,8,256,0.1747,48.04
+2,1024,8,256,0.1762,95.27
+4,1024,8,256,0.1784,188.22
+8,1024,8,256,0.1817,369.53
+16,1024,8,256,0.2796,480.25
+32,1024,8,256,0.4339,619.00
+64,1024,8,256,0.6960,771.73
+128,1024,8,256,1.3439,799.36
+1,2048,8,256,0.3410,49.21
+2,2048,8,256,0.3439,97.60
+4,2048,8,256,0.3469,193.52
+8,2048,8,256,0.3533,379.94
+16,2048,8,256,0.5461,491.67
+32,2048,8,256,0.8493,632.28
+64,2048,8,256,1.3667,785.82
+128,2048,8,256,2.6465,811.64
+1,4096,8,256,0.6742,49.77
+2,4096,8,256,0.6777,99.03
+4,4096,8,256,0.6836,196.36
+8,4096,8,256,0.6950,386.31
+16,4096,8,256,1.0803,497.02
+32,4096,8,256,1.6794,639.44
+64,4096,8,256,2.7101,792.50
+128,4096,8,256,5.2543,817.52
+1,8192,8,256,1.3375,50.18
+2,8192,8,256,1.3448,99.81
+4,8192,8,256,1.3564,197.91
+8,8192,8,256,1.3799,389.08
+16,8192,8,256,2.1465,500.25
+32,8192,8,256,3.3342,644.12
+64,8192,8,256,5.3983,795.67
+128,8192,8,256,10.4691,820.55
+1,16384,8,256,2.6697,50.28
+2,16384,8,256,2.6817,100.10
+4,16384,8,256,2.7049,198.49
+8,16384,8,256,2.7533,390.00
+16,16384,8,256,4.2789,501.89
+32,16384,8,256,6.6476,646.11
+64,16384,8,256,10.7723,797.43
+128,16384,8,256,OOM,OOM

基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/benchmark结果实例/benchmark_kvcache_headdim32.csv

@@ -0,0 +1,49 @@
+batch_size,seq_len_kv,heads,headdim,time_ms,bandwidth_GB_s
+1,512,8,32,0.0257,20.45
+2,512,8,32,0.0256,41.02
+4,512,8,32,0.0258,81.28
+8,512,8,32,0.0265,158.45
+16,512,8,32,0.0396,212.30
+32,512,8,32,0.0516,325.43
+64,512,8,32,0.0721,465.83
+128,512,8,32,0.1270,529.03
+1,1024,8,32,0.0461,22.75
+2,1024,8,32,0.0465,45.15
+4,1024,8,32,0.0477,88.04
+8,1024,8,32,0.0548,153.23
+16,1024,8,32,0.0734,228.71
+32,1024,8,32,0.0958,350.42
+64,1024,8,32,0.1334,503.15
+128,1024,8,32,0.2381,564.04
+1,2048,8,32,0.0872,24.06
+2,2048,8,32,0.0904,46.42
+4,2048,8,32,0.1028,81.59
+8,2048,8,32,0.1067,157.25
+16,2048,8,32,0.1428,235.10
+32,2048,8,32,0.1818,369.13
+64,2048,8,32,0.2554,525.57
+128,2048,8,32,0.4622,580.86
+1,4096,8,32,0.1730,24.25
+2,4096,8,32,0.1955,42.91
+4,4096,8,32,0.2020,83.05
+8,4096,8,32,0.2140,156.83
+16,4096,8,32,0.2777,241.65
+32,4096,8,32,0.3542,378.99
+64,4096,8,32,0.4990,538.05
+128,4096,8,32,0.9099,590.13
+1,8192,8,32,0.3820,21.96
+2,8192,8,32,0.3913,42.88
+4,8192,8,32,0.4127,81.31
+8,8192,8,32,0.4224,158.88
+16,8192,8,32,0.5490,244.51
+32,8192,8,32,0.6960,385.70
+64,8192,8,32,0.9870,543.98
+128,8192,8,32,1.8100,593.25
+1,16384,8,32,0.7655,21.92
+2,16384,8,32,0.8067,41.59
+4,16384,8,32,0.8228,81.56
+8,16384,8,32,0.8397,159.85
+16,16384,8,32,1.0910,246.04
+32,16384,8,32,1.3824,388.37
+64,16384,8,32,1.9663,546.08
+128,16384,8,32,3.6107,594.78

基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/benchmark结果实例/benchmark_kvcache_headdim512.csv

@@ -0,0 +1,49 @@
+batch_size,seq_len_kv,heads,headdim,time_ms,bandwidth_GB_s
+1,512,8,512,0.3588,23.40
+2,512,8,512,0.3651,46.00
+4,512,8,512,0.3736,89.89
+8,512,8,512,0.3856,174.22
+16,512,8,512,0.7472,179.80
+32,512,8,512,1.1447,234.72
+64,512,8,512,1.9549,274.89
+128,512,8,512,3.8962,275.85
+1,1024,8,512,0.7261,23.12
+2,1024,8,512,0.7354,45.65
+4,1024,8,512,0.7496,89.57
+8,1024,8,512,0.7746,173.35
+16,1024,8,512,1.5049,178.46
+32,1024,8,512,2.3111,232.42
+64,1024,8,512,3.9538,271.70
+128,1024,8,512,7.8811,272.62
+1,2048,8,512,1.4636,22.93
+2,2048,8,512,1.4826,45.27
+4,2048,8,512,1.5109,88.86
+8,2048,8,512,1.5549,172.68
+16,2048,8,512,3.0237,177.60
+32,2048,8,512,4.6439,231.27
+64,2048,8,512,7.9560,269.99
+128,2048,8,512,15.8741,270.63
+1,4096,8,512,2.9312,22.90
+2,4096,8,512,2.9675,45.24
+4,4096,8,512,3.0243,88.77
+8,4096,8,512,3.1127,172.50
+16,4096,8,512,6.0753,176.76
+32,4096,8,512,9.3182,230.49
+64,4096,8,512,15.9642,269.07
+128,4096,8,512,31.8313,269.89
+1,8192,8,512,5.8843,22.81
+2,8192,8,512,5.9344,45.24
+4,8192,8,512,6.0465,88.80
+8,8192,8,512,6.2334,172.27
+16,8192,8,512,12.1594,176.62
+32,8192,8,512,18.6826,229.90
+64,8192,8,512,32.0055,268.41
+128,8192,8,512,OOM,OOM
+1,16384,8,512,11.8153,22.72
+2,16384,8,512,11.9237,45.03
+4,16384,8,512,12.1671,88.25
+8,16384,8,512,12.4948,171.88
+16,16384,8,512,24.3414,176.45
+32,16384,8,512,37.3907,229.74
+64,16384,8,512,OOM,OOM
+128,16384,8,512,OOM,OOM

基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/benchmark结果实例/benchmark_kvcache_headdim64.csv

@@ -0,0 +1,49 @@
+batch_size,seq_len_kv,heads,headdim,time_ms,bandwidth_GB_s
+1,512,8,64,0.0404,25.99
+2,512,8,64,0.0399,52.60
+4,512,8,64,0.0413,101.69
+8,512,8,64,0.0482,174.25
+16,512,8,64,0.0540,310.86
+32,512,8,64,0.0629,533.75
+64,512,8,64,0.0833,806.14
+128,512,8,64,0.1104,1216.59
+1,1024,8,64,0.0747,28.08
+2,1024,8,64,0.0766,54.77
+4,1024,8,64,0.0891,94.17
+8,1024,8,64,0.0918,182.94
+16,1024,8,64,0.1044,321.41
+32,1024,8,64,0.1179,569.43
+64,1024,8,64,0.1566,857.28
+128,1024,8,64,0.2078,1292.17
+1,2048,8,64,0.1455,28.84
+2,2048,8,64,0.1684,49.82
+4,2048,8,64,0.1730,97.01
+8,2048,8,64,0.1850,181.39
+16,2048,8,64,0.2009,334.18
+32,2048,8,64,0.2268,592.01
+64,2048,8,64,0.3002,894.44
+128,2048,8,64,0.4027,1333.64
+1,4096,8,64,0.3265,25.69
+2,4096,8,64,0.3322,50.51
+4,4096,8,64,0.3522,95.27
+8,4096,8,64,0.3632,184.79
+16,4096,8,64,0.3942,340.56
+32,4096,8,64,0.4456,602.47
+64,4096,8,64,0.5927,905.94
+128,4096,8,64,0.7938,1352.87
+1,8192,8,64,0.6508,25.78
+2,8192,8,64,0.6879,48.78
+4,8192,8,64,0.7008,95.77
+8,8192,8,64,0.7199,186.44
+16,8192,8,64,0.7786,344.79
+32,8192,8,64,0.8798,610.25
+64,8192,8,64,1.1745,914.30
+128,8192,8,64,1.5728,1365.50
+1,16384,8,64,1.3524,24.81
+2,16384,8,64,1.3698,48.99
+4,16384,8,64,1.3923,96.40
+8,16384,8,64,1.4267,188.16
+16,16384,8,64,1.5451,347.47
+32,16384,8,64,1.7622,609.32
+64,16384,8,64,2.3392,918.09
+128,16384,8,64,3.1332,1370.84

基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/benchmark结果实例/benchmark_kvcache_headdim96.csv

@@ -0,0 +1,49 @@
+batch_size,seq_len_kv,heads,headdim,time_ms,bandwidth_GB_s
+1,512,8,96,0.0407,38.67
+2,512,8,96,0.0398,79.02
+4,512,8,96,0.0431,146.08
+8,512,8,96,0.0495,254.61
+16,512,8,96,0.0698,360.64
+32,512,8,96,0.1117,450.87
+64,512,8,96,0.1780,566.16
+128,512,8,96,0.3329,605.28
+1,1024,8,96,0.0732,43.01
+2,1024,8,96,0.0794,79.29
+4,1024,8,96,0.0871,144.54
+8,1024,8,96,0.0934,269.52
+16,1024,8,96,0.1297,388.14
+32,1024,8,96,0.2114,476.36
+64,1024,8,96,0.3379,596.08
+128,1024,8,96,0.6327,636.68
+1,2048,8,96,0.1505,41.80
+2,2048,8,96,0.1619,77.76
+4,2048,8,96,0.1713,146.94
+8,2048,8,96,0.1780,282.84
+16,2048,8,96,0.2492,404.09
+32,2048,8,96,0.4088,492.55
+64,2048,8,96,0.6575,612.55
+128,2048,8,96,1.2457,646.61
+1,4096,8,96,0.3099,40.61
+2,4096,8,96,0.3259,77.23
+4,4096,8,96,0.3346,150.42
+8,4096,8,96,0.3467,290.41
+16,4096,8,96,0.4888,411.94
+32,4096,8,96,0.8055,499.94
+64,4096,8,96,1.3209,609.72
+128,4096,8,96,2.4810,649.25
+1,8192,8,96,0.6343,39.68
+2,8192,8,96,0.6437,78.20
+4,8192,8,96,0.6601,152.50
+8,8192,8,96,0.6826,294.97
+16,8192,8,96,0.9688,415.64
+32,8192,8,96,1.6057,501.55
+64,8192,8,96,2.6527,607.19
+128,8192,8,96,4.9464,651.27
+1,16384,8,96,1.2581,40.01
+2,16384,8,96,1.2812,78.57
+4,16384,8,96,1.3112,153.55
+8,16384,8,96,1.3653,294.92
+16,16384,8,96,1.9351,416.16
+32,16384,8,96,3.2277,499.01
+64,16384,8,96,5.3192,605.60
+128,16384,8,96,9.8747,652.44

基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/starter/smoke_test-CUDAMaca.txt

@@ -0,0 +1,119 @@
+#include <stdint.h>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+#define HEAD_DIM 128
+
+__global__ void paged_attention_kernel(
+    const __nv_bfloat16* q,
+    const __nv_bfloat16* k_cache_paged,
+    const __nv_bfloat16* v_cache_paged,
+    __nv_bfloat16* output,
+    const int32_t* cache_seqlens,
+    const int32_t* block_table,
+    int64_t batch_size,
+    int64_t seqlen_q,
+    int64_t num_heads,
+    int64_t num_heads_k,
+    int64_t headdim,
+    int64_t page_block_size,
+    int64_t blocks_per_batch)
+{
+    int batch_idx = blockIdx.x / num_heads;
+    int head_idx  = blockIdx.x % num_heads;
+    if (batch_idx >= batch_size || head_idx >= num_heads) return;
+
+    int seqlen = cache_seqlens[batch_idx];
+    int tid = threadIdx.x;
+
+    // 加载对应 head 的 query 元素
+    int64_t q_offset = ((batch_idx * seqlen_q + 0) * num_heads + head_idx) * headdim;
+    float q_val = __bfloat162float(q[q_offset + tid]);
+
+    // Online safe softmax 状态
+    float max_val = -1e38f;
+    float sum_exp = 0.0f;
+    float out_acc = 0.0f;
+    float scale = 1.0f / sqrtf(static_cast<float>(headdim));
+
+    // 静态共享内存，避免动态分配可能带来的兼容性问题
+    __shared__ float s_score[HEAD_DIM];
+
+    for (int token = 0; token < seqlen; ++token) {
+        int page_idx    = token / page_block_size;
+        int page_offset = token % page_block_size;
+        int physical_block = block_table[batch_idx * blocks_per_batch + page_idx];
+
+        // 读取 key 元素
+        const __nv_bfloat16* k_ptr = k_cache_paged
+            + (physical_block * page_block_size + page_offset) * (num_heads_k * headdim)
+            + head_idx * headdim;
+        float k_val = __bfloat162float(k_ptr[tid]);
+
+        // 点积 -> 共享内存归约
+        s_score[tid] = q_val * k_val;
+        __syncthreads();
+
+        for (int stride = HEAD_DIM >> 1; stride > 0; stride >>= 1) {
+            if (tid < stride) {
+                s_score[tid] += s_score[tid + stride];
+            }
+            __syncthreads();
+        }
+        float score = s_score[0] * scale;
+
+        // 更新 softmax 状态
+        float new_max = fmaxf(max_val, score);
+        float rescale = expf(max_val - new_max);
+        sum_exp = sum_exp * rescale + expf(score - new_max);
+        out_acc = out_acc * rescale;
+        max_val = new_max;
+
+        // 读取 value 元素，并累加（用最新 max 的权重）
+        const __nv_bfloat16* v_ptr = v_cache_paged
+            + (physical_block * page_block_size + page_offset) * (num_heads_k * headdim)
+            + head_idx * headdim;
+        float v_val = __bfloat162float(v_ptr[tid]);
+        out_acc += expf(score - max_val) * v_val;
+
+        __syncthreads();   // 确保下次迭代共享内存可安全复用
+    }
+
+    if (seqlen > 0) {
+        out_acc /= sum_exp;
+    } else {
+        out_acc = 0.0f;
+    }
+
+    int64_t out_offset = ((batch_idx * seqlen_q + 0) * num_heads + head_idx) * headdim + tid;
+    output[out_offset] = __float2bfloat16(out_acc);
+}
+
+extern "C" void run_kernel(
+    const __nv_bfloat16* q,
+    const __nv_bfloat16* k_cache_paged,
+    const __nv_bfloat16* v_cache_paged,
+    __nv_bfloat16* output,
+    const int32_t* cache_seqlens,
+    const int32_t* block_table,
+    int64_t batch_size,
+    int64_t seqlen_k,
+    int64_t seqlen_q,
+    int64_t num_heads,
+    int64_t num_heads_k,
+    int64_t headdim,
+    int64_t page_block_size,
+    int64_t num_blocks,
+    int64_t causal)
+{
+    int64_t blocks_per_batch = num_blocks / batch_size;
+    dim3 grid(batch_size * num_heads);
+    dim3 block(HEAD_DIM);
+
+    paged_attention_kernel<<<grid, block>>>(
+        q, k_cache_paged, v_cache_paged, output,
+        cache_seqlens, block_table,
+        batch_size, seqlen_q, num_heads, num_heads_k, headdim,
+        page_block_size, blocks_per_batch
+    );
+}

基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/starter/smoke_test-Tilelang.txt

@@ -0,0 +1,187 @@
+"""FlashAttention KV Cache Decode in TileLang."""
+
+import tilelang
+import tilelang.language as T
+from tilelang import jit
+
+NUM_SPLITS = 4
+real_kernel = None
+
+@jit(
+    pass_configs={
+        tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: False,
+    },
+)
+def build_kernel(
+    batch_size,
+    num_heads,
+    num_heads_k,
+    headdim,
+    page_block_size,
+    num_blocks,
+    causal,
+):
+    blocks_per_batch = num_blocks // batch_size
+    assert blocks_per_batch % NUM_SPLITS == 0, (
+        f"blocks_per_batch={blocks_per_batch} must be divisible by NUM_SPLITS={NUM_SPLITS}"
+    )
+    blocks_per_split = blocks_per_batch // NUM_SPLITS
+
+    BLOCK_M = 1
+    BLOCK_N = page_block_size
+    scale = (1.0 / headdim) ** 0.5 * 1.44269504  # log2(e)
+    dtype = "bfloat16"
+    accum_dtype = "float32"
+
+    # Use a large-negative-finite sentinel instead of -inf to avoid
+    # (-inf) - (-inf) = NaN when an entire split is masked out.
+    NEG_INF_SAFE = -1e30
+
+    @T.prim_func
+    def kernel(
+        Q: T.Tensor([batch_size, 1, num_heads, headdim], dtype),
+        K: T.Tensor([num_blocks, page_block_size, num_heads_k, headdim], dtype),
+        V: T.Tensor([num_blocks, page_block_size, num_heads_k, headdim], dtype),
+        Output: T.Tensor([batch_size, 1, num_heads, headdim], dtype),
+        cache_seqlens: T.Tensor([batch_size], "int32"),
+        block_table: T.Tensor([batch_size, blocks_per_batch], "int32"),
+    ):
+        # float32 workspace — avoids BF16StorageLegalize var-remap bug
+        glse = T.alloc_global([batch_size, num_heads, NUM_SPLITS], accum_dtype)
+        Output_partial = T.alloc_global(
+            [batch_size, 1, num_heads, NUM_SPLITS, headdim], accum_dtype
+        )
+
+        # ============= Stage 1: split kernel =============
+        with T.Kernel(NUM_SPLITS, num_heads, batch_size, threads=128) as (bs, bh, bz):
+            Q_shared = T.alloc_shared([BLOCK_M, headdim], dtype)
+            K_shared = T.alloc_shared([BLOCK_N, headdim], dtype)
+            V_shared = T.alloc_shared([BLOCK_N, headdim], dtype)
+            acc_s = T.alloc_fragment([BLOCK_M, BLOCK_N], accum_dtype)
+            acc_o = T.alloc_fragment([BLOCK_M, headdim], accum_dtype)
+            scores_max = T.alloc_fragment([BLOCK_M], accum_dtype)
+            scores_max_prev = T.alloc_fragment([BLOCK_M], accum_dtype)
+            scores_scale = T.alloc_fragment([BLOCK_M], accum_dtype)
+            scores_sum = T.alloc_fragment([BLOCK_M], accum_dtype)
+            logsum = T.alloc_fragment([BLOCK_M], accum_dtype)
+
+            T.copy(Q[bz, 0, bh, :], Q_shared)
+
+            kv_seqlen = cache_seqlens[bz]
+            split_k_start = bs * blocks_per_split
+
+            T.fill(acc_o, 0)
+            T.fill(logsum, 0)
+            # KEY FIX: use -1e30 instead of -inf to avoid (-inf)-(-inf)=NaN
+            T.fill(scores_max, NEG_INF_SAFE)
+
+            for k in T.Pipelined(blocks_per_split, num_stages=2):
+                global_k = split_k_start + k
+                physical_block = block_table[bz, global_k]
+                tok_offset = global_k * page_block_size
+
+                # ----- Q @ K^T (hand-written, M=1, masked) -----
+                T.copy(K[physical_block, 0:BLOCK_N, bh, :], K_shared)
+                T.fill(acc_s, 0)
+                for j in T.Parallel(BLOCK_N):
+                    if tok_offset + j < kv_seqlen:
+                        for d in T.serial(headdim):
+                            acc_s[0, j] = acc_s[0, j] + Q_shared[0, d] * K_shared[j, d]
+                    else:
+                        acc_s[0, j] = -T.infinity(accum_dtype)
+
+                # ----- online softmax -----
+                T.copy(scores_max, scores_max_prev)
+                # KEY FIX: use -1e30 instead of -inf here too
+                T.fill(scores_max, NEG_INF_SAFE)
+                T.reduce_max(acc_s, scores_max, dim=1, clear=False)
+                scores_max[0] = T.max(scores_max[0], scores_max_prev[0])
+                # (prev - cur) is now (finite - finite) = 0 when both are sentinel,
+                # never (-inf - (-inf)) = NaN
+                scores_scale[0] = T.exp2((scores_max_prev[0] - scores_max[0]) * scale)
+                for j in T.Parallel(BLOCK_N):
+                    acc_s[0, j] = T.exp2((acc_s[0, j] - scores_max[0]) * scale)
+                T.reduce_sum(acc_s, scores_sum, dim=1)
+                logsum[0] = logsum[0] * scores_scale[0] + scores_sum[0]
+                for d in T.Parallel(headdim):
+                    acc_o[0, d] = acc_o[0, d] * scores_scale[0]
+
+                # ----- P @ V (hand-written, fp32 accum) -----
+                T.copy(V[physical_block, 0:BLOCK_N, bh, :], V_shared)
+                for d in T.Parallel(headdim):
+                    for j in T.serial(BLOCK_N):
+                        acc_o[0, d] = acc_o[0, d] + acc_s[0, j] * V_shared[j, d]
+
+            # ----- final normalise & write partial state -----
+            # KEY FIX: add epsilon to avoid 0/0 = NaN when split is all-masked
+            safe_logsum = logsum[0] + 1e-30
+            for d in T.Parallel(headdim):
+                acc_o[0, d] = acc_o[0, d] / safe_logsum
+
+            lse_local = T.alloc_fragment([1], accum_dtype)
+            lse_local[0] = T.log2(safe_logsum) + scores_max[0] * scale
+            glse[bz, bh, bs] = lse_local[0]
+
+            for d in T.Parallel(headdim):
+                Output_partial[bz, 0, bh, bs, d] = acc_o[0, d]
+
+        # ============= Stage 2: combine kernel =============
+        with T.Kernel(num_heads, batch_size, threads=128) as (bh, bz):
+            lse_local = T.alloc_fragment([NUM_SPLITS], accum_dtype)
+            for s in T.serial(NUM_SPLITS):
+                lse_local[s] = glse[bz, bh, s]
+
+            lse_max = T.alloc_fragment([1], accum_dtype)
+            lse_max[0] = -T.infinity(accum_dtype)
+            for s in T.serial(NUM_SPLITS):
+                lse_max[0] = T.max(lse_max[0], lse_local[s])
+
+            lse_logsum = T.alloc_fragment([1], accum_dtype)
+            lse_logsum[0] = 0
+            for s in T.serial(NUM_SPLITS):
+                lse_logsum[0] = lse_logsum[0] + T.exp2(lse_local[s] - lse_max[0])
+            lse_logsum[0] = T.log2(lse_logsum[0]) + lse_max[0]
+
+            o_accum = T.alloc_fragment([headdim], accum_dtype)
+            T.fill(o_accum, 0)
+            for s in T.serial(NUM_SPLITS):
+                s_scale = T.exp2(lse_local[s] - lse_logsum[0])
+                for d in T.Parallel(headdim):
+                    o_accum[d] = o_accum[d] + Output_partial[bz, 0, bh, s, d] * s_scale
+
+            for d in T.Parallel(headdim):
+                Output[bz, 0, bh, d] = T.Cast(dtype, o_accum[d])
+
+    return kernel
+
+
+def run_kernel(
+    q,
+    k_cache_paged,
+    v_cache_paged,
+    output,
+    cache_seqlens,
+    block_table,
+    batch_size,
+    seqlen_k,
+    seqlen_q,
+    num_heads,
+    num_heads_k,
+    headdim,
+    page_block_size,
+    num_blocks,
+    causal,
+):
+    global real_kernel
+
+    B = int(batch_size)
+    H = int(num_heads)
+    HK = int(num_heads_k)
+    D = int(headdim)
+    PBS = int(page_block_size)
+    NB = int(num_blocks)
+
+    if real_kernel is None:
+        real_kernel = build_kernel(B, H, HK, D, PBS, NB, int(causal))
+
+    real_kernel(q, k_cache_paged, v_cache_paged, output, cache_seqlens, block_table)

基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/starter/smoke_test-Triton.txt

@@ -0,0 +1,106 @@
+import triton
+import triton.language as tl
+import torch
+
+@triton.jit
+def slow_decode_kernel(
+    q_ptr,
+    k_cache_ptr,
+    v_cache_ptr,
+    output_ptr,
+    cache_seqlens_ptr,
+    block_table_ptr,
+    num_heads: tl.constexpr,
+    num_heads_k: tl.constexpr,
+    headdim: tl.constexpr,
+    page_block_size: tl.constexpr,
+    blocks_per_batch,
+):
+    # 维度索引
+    pid_b = tl.program_id(0)  # Batch index
+    pid_h = tl.program_id(1)  # Head index
+
+    # GQA Support: 映射 Query Head 到 KV Head
+    kv_head = pid_h * num_heads_k // num_heads
+
+    # 加载实际的 KV 序列长度
+    seq_len = tl.load(cache_seqlens_ptr + pid_b).to(tl.int32)
+
+    # 维度偏移量 [0, 1, ..., headdim-1]
+    offs_d = tl.arange(0, headdim)
+
+    # Online Softmax 累加器
+    acc = tl.zeros([headdim], dtype=tl.float32)
+    l_i = 0.0
+    m_i = float('-inf')
+    scale = 1.0 / tl.sqrt(float(headdim))
+
+    # === 性能瓶颈：串行遍历整个序列 ===
+    # 不使用 Block 并行，而是用单个 Block 串行循环处理所有 Token
+    t = 0
+    while t < seq_len:
+        # 性能瓶颈：每次循环都重新加载 Q，增加显存压力
+        q = tl.load(q_ptr + pid_b * num_heads * headdim + pid_h * headdim + offs_d).to(tl.float32)
+        q = q * scale
+
+        # Paged KV 映射逻辑
+        page_idx = t // page_block_size
+        page_off = t % page_block_size
+
+        # 查表获取物理 Block 索引
+        # blocks_per_batch 是计算出来的步长
+        phys_block = tl.load(block_table_ptr + pid_b * blocks_per_batch + page_idx)
+
+        # 计算 K 和 V 的物理地址
+        # Layout: (num_blocks, page_block_size, num_heads_k, headdim)
+        kv_base = phys_block * page_block_size * num_heads_k * headdim + \
+                  page_off * num_heads_k * headdim + \
+                  kv_head * headdim
+        
+        # 加载 K 和 V 向量
+        k = tl.load(k_cache_ptr + kv_base + offs_d).to(tl.float32)
+        v = tl.load(v_cache_ptr + kv_base + offs_d).to(tl.float32)
+
+        # Attention 计算
+        s = tl.sum(q * k)  # 点积
+        
+        # Online Softmax 更新
+        m_new = tl.maximum(m_i, s)
+        p = tl.exp(s - m_new)
+        alpha = tl.exp(m_i - m_new)
+        
+        acc = acc * alpha + p * v
+        l_i = l_i * alpha + p
+        m_i = m_new
+
+        t += 1
+
+    # 写回结果
+    # 这里没有处理 l_i 为 0 的边界情况，但测试数据 seq_len 通常很大
+    out = acc / l_i
+    tl.store(output_ptr + pid_b * num_heads * headdim + pid_h * headdim + offs_d, out)
+
+def run_kernel(
+    q, k_cache_paged, v_cache_paged, output,
+    cache_seqlens, block_table,
+    batch_size, seqlen_k, seqlen_q, num_heads, num_heads_k, headdim,
+    page_block_size, num_blocks, causal,
+):
+    # 计算每个 batch 对应的 block_table 行宽
+    blocks_per_batch = num_blocks // batch_size
+    
+    # 启动配置：每个 Head 一个 Block
+    # 总 Block 数 = batch_size * num_heads (最大 128个)，并行度极低
+    grid = (batch_size, num_heads)
+    
+    slow_decode_kernel[grid](
+        q, k_cache_paged, v_cache_paged, output,
+        cache_seqlens, block_table,
+        num_heads=num_heads,
+        num_heads_k=num_heads_k,
+        headdim=headdim,
+        page_block_size=page_block_size,
+        blocks_per_batch=blocks_per_batch,
+        num_warps=1,    # 性能瓶颈：仅使用 1 个 warp，限制计算吞吐
+        num_stages=1,   # 性能瓶颈：禁用流水线并行
+    )

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

@@ -0,0 +1,16 @@
+{
+  "id": 197,
+  "displayId": 20005,
+  "type": "Traditional",
+  "isPublic": false,
+  "locales": [
+    "zh_CN"
+  ],
+  "samples": [
+    {
+      "inputData": "1\n",
+      "outputData": ""
+    }
+  ],
+  "problemTagIds": []
+}

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

@@ -0,0 +1,298 @@
+from __future__ import annotations
+
+
+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
+
+
+def _build_cases():
+    cases = []
+    for headdim in HEAD_DIMS:
+        for seqlen_k in SEQ_LENS_KV:
+            for batch_size in BATCH_SIZES:
+                cases.append(
+                    (
+                        batch_size,
+                        seqlen_k,
+                        SEQ_LEN_Q,
+                        NUM_HEADS,
+                        NUM_HEADS_K,
+                        headdim,
+                        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 math
+    import sys
+
+    import torch
+
+    KernelArg = Union[torch.Tensor, int, float]
+    CURRENT_CASE = None
+
+    def _ensure_flashattn_importable():
+        try:
+            from flash_attn.flash_attn_interface import flash_attn_with_kvcache  # noqa: F401
+
+            return
+        except ImportError:
+            pass
+
+        here = Path(__file__).resolve()
+        for parent in here.parents:
+            candidate = parent / "flashattn"
+            if (candidate / "flash_attn").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 _get_num_blocks(batch_size: int, seqlen_k: int, page_block_size: int) -> int:
+        num_blocks = math.ceil(seqlen_k / page_block_size) * batch_size * 3
+        return max(1024, num_blocks)
+
+    def getTestCaseSize() -> Tuple[List[Tuple[int, ...]], Tuple[int, int]]:
+        testcase_id = _get_testcase_index()
+        global CURRENT_CASE
+        (
+            batch_size,
+            seqlen_k,
+            seqlen_q,
+            num_heads,
+            num_heads_k,
+            headdim,
+            page_block_size,
+            causal,
+        ) = TESTCASES[testcase_id]
+        num_blocks = _get_num_blocks(batch_size, seqlen_k, page_block_size)
+        blocks_per_batch = num_blocks // batch_size
+        CURRENT_CASE = (
+            batch_size,
+            seqlen_k,
+            seqlen_q,
+            num_heads,
+            num_heads_k,
+            headdim,
+            page_block_size,
+            num_blocks,
+            causal,
+            20260720 + testcase_id,
+        )
+        warmup = 3
+        iters = max(1, _compute_reps(batch_size, seqlen_k, headdim))
+        return [
+            (batch_size, seqlen_q, num_heads, headdim),
+            (num_blocks, page_block_size, num_heads_k, headdim),
+            (num_blocks, page_block_size, num_heads_k, headdim),
+            (batch_size, seqlen_q, num_heads, headdim),
+            (batch_size,),
+            (batch_size, blocks_per_batch),
+            (), (), (), (), (), (), (), (), (),
+        ], (warmup, iters)
+
+    def genTestCase(testcase_sizes, device: str = "cuda") -> List[KernelArg]:
+        del testcase_sizes
+        (
+            batch_size,
+            seqlen_k,
+            seqlen_q,
+            num_heads,
+            num_heads_k,
+            headdim,
+            page_block_size,
+            num_blocks,
+            causal,
+            seed,
+        ) = CURRENT_CASE
+        gen = torch.Generator(device=device)
+        gen.manual_seed(seed)
+        dtype = torch.bfloat16
+        blocks_per_batch = num_blocks // batch_size
+        q = torch.randn(
+            batch_size,
+            seqlen_q,
+            num_heads,
+            headdim,
+            dtype=dtype,
+            device=device,
+            generator=gen,
+        ).contiguous()
+        k_cache_paged = torch.randn(
+            num_blocks,
+            page_block_size,
+            num_heads_k,
+            headdim,
+            dtype=dtype,
+            device=device,
+            generator=gen,
+        ).contiguous()
+        v_cache_paged = torch.randn(
+            num_blocks,
+            page_block_size,
+            num_heads_k,
+            headdim,
+            dtype=dtype,
+            device=device,
+            generator=gen,
+        ).contiguous()
+        output = torch.empty(
+            batch_size,
+            seqlen_q,
+            num_heads,
+            headdim,
+            dtype=dtype,
+            device=device,
+        )
+        cache_seqlens = torch.full((batch_size,), seqlen_k, dtype=torch.int32, device=device)
+        block_table = torch.randperm(num_blocks, dtype=torch.int32, device=device, generator=gen).reshape(
+            batch_size,
+            blocks_per_batch,
+        )
+        return [
+            q,
+            k_cache_paged,
+            v_cache_paged,
+            output,
+            cache_seqlens,
+            block_table,
+            batch_size,
+            seqlen_k,
+            seqlen_q,
+            num_heads,
+            num_heads_k,
+            headdim,
+            page_block_size,
+            num_blocks,
+            causal,
+        ]
+
+    def baseline(
+        q,
+        k_cache_paged,
+        v_cache_paged,
+        output,
+        cache_seqlens,
+        block_table,
+        batch_size,
+        seqlen_k,
+        seqlen_q,
+        num_heads,
+        num_heads_k,
+        headdim,
+        page_block_size,
+        num_blocks,
+        causal,
+    ):
+        _ensure_flashattn_importable()
+        from flash_attn.flash_attn_interface import flash_attn_with_kvcache
+
+        out = flash_attn_with_kvcache(
+            q,
+            k_cache_paged,
+            v_cache_paged,
+            None,
+            None,
+            cache_seqlens=cache_seqlens,
+            cache_batch_idx=None,
+            block_table=block_table,
+            causal=bool(causal),
+            window_size=(-1, -1),
+            rotary_interleaved=False,
+            alibi_slopes=None,
+            num_splits=1,
+        )
+        output.copy_(out)
+        return [
+            q,
+            k_cache_paged,
+            v_cache_paged,
+            output,
+            cache_seqlens,
+            block_table,
+            batch_size,
+            seqlen_k,
+            seqlen_q,
+            num_heads,
+            num_heads_k,
+            headdim,
+            page_block_size,
+            num_blocks,
+            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

基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/00_棰樼洰鎻忚堪.md

@@ -0,0 +1,28 @@
+---
+sectionTitle: "题目描述"
+type: "Text"
+---
+你需要实现 FlashAttention paged KV cache decode 的 CUDA C++ 前向算子。
+
+本题输入采用 `flash_attn_with_kvcache` 在 `flashattn/benchmarks/benchmark_kvcache.py` 中使用的 paged KV cache 配置。每个 batch 只有 1 个 query token，KV cache 长度为 `seqlen_k`，K/V cache 按 page 存储。
+
+评测程序会调用你提交代码中的 `run_kernel` 函数。你需要根据 `cache_seqlens` 和 `block_table` 读取 paged KV cache，并将结果写入 `output`。
+
+baseline 使用 benchmark 中的 FlashAttention Python API：
+
+```python
+out = flash_attn_with_kvcache(
+    q, k_cache_paged, v_cache_paged, None, None,
+    cache_seqlens=cache_seqlens,
+    cache_batch_idx=None,
+    block_table=block_table,
+    causal=False,
+    window_size=(-1, -1),
+    rotary_interleaved=False,
+    alibi_slopes=None,
+    num_splits=1,
+)
+output.copy_(out)
+```
+
+如何提交代码详见[评测指南](/d/2)。

基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/01_鎺ュ彛绾﹀畾.cuda.md

@@ -0,0 +1,43 @@
+---
+sectionTitle: "接口约定"
+type: "codeSample"
+lang: "cuda"
+---
+你必须在提交的 CUDA 源码中提供如下 **C 符号**，函数名、参数类型、顺序必须完全一致，并使用 `extern "C"` 防止 name mangling：
+
+```cpp
+#include <stdint.h>
+#include <cuda_bf16.h>
+
+extern "C" void run_kernel(
+    const __nv_bfloat16* q,
+    const __nv_bfloat16* k_cache_paged,
+    const __nv_bfloat16* v_cache_paged,
+    __nv_bfloat16* output,
+    const int32_t* cache_seqlens,
+    const int32_t* block_table,
+    int64_t batch_size,
+    int64_t seqlen_k,
+    int64_t seqlen_q,
+    int64_t num_heads,
+    int64_t num_heads_k,
+    int64_t headdim,
+    int64_t page_block_size,
+    int64_t num_blocks,
+    int64_t causal
+);
+```
+
+### 参数说明
+
+* `q`：decode query tensor，shape `(batch_size, seqlen_q, num_heads, headdim)`，连续 `bf16`
+* `k_cache_paged`：paged key cache，shape `(num_blocks, page_block_size, num_heads_k, headdim)`，连续 `bf16`
+* `v_cache_paged`：paged value cache，shape `(num_blocks, page_block_size, num_heads_k, headdim)`，连续 `bf16`
+* `output`：输出缓冲区，shape `(batch_size, seqlen_q, num_heads, headdim)`，连续 `bf16`
+* `cache_seqlens`：每个 batch 的 KV 长度，shape `(batch_size)`，连续 `int32`
+* `block_table`：每个 batch 的 page 映射表，shape `(batch_size, num_blocks / batch_size)`，连续 `int32`
+* `seqlen_q`：query 长度，评测中固定为 `1`
+* `page_block_size`：page size，评测中固定为 `16`
+* `causal`：是否启用 causal mask，评测中固定为 `0`
+
+`run_kernel` 内部需要自行计算合适的 launch 配置并启动 CUDA kernel。为保证计时准确，不建议在 `run_kernel` 内部做 `cudaDeviceSynchronize()` 或显式同步。

基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/01_鎺ュ彛绾﹀畾.tilelang.md

@@ -0,0 +1,57 @@
+---
+sectionTitle: "接口约定"
+type: "codeSample"
+lang: "tilelang"
+---
+你必须在提交的 Python 代码中提供 `run_kernel` 函数，函数名、参数顺序、类型必须完全一致：
+
+```python
+import tilelang
+import tilelang.language as T
+from tilelang import jit
+
+real_kernel = None
+
+@jit
+def build_kernel(*args):
+    @T.prim_func
+    def kernel(*args):
+        ...
+    return kernel
+
+def run_kernel(
+    q,  # Tensor[bf16], shape (batch_size, seqlen_q, num_heads, headdim)
+    k_cache_paged,  # Tensor[bf16], shape (num_blocks, page_block_size, num_heads_k, headdim)
+    v_cache_paged,  # Tensor[bf16], shape (num_blocks, page_block_size, num_heads_k, headdim)
+    output,  # Tensor[bf16], shape (batch_size, seqlen_q, num_heads, headdim)
+    cache_seqlens,  # Tensor[int32], shape (batch_size)
+    block_table,  # Tensor[int32], shape (batch_size, num_blocks / batch_size)
+    batch_size,  # int64
+    seqlen_k,  # int64
+    seqlen_q,  # int64
+    num_heads,  # int64
+    num_heads_k,  # int64
+    headdim,  # int64
+    page_block_size,  # int64
+    num_blocks,  # int64
+    causal,  # int64
+):
+    global real_kernel
+    if real_kernel is None:
+        real_kernel = build_kernel(...)
+    real_kernel(q, k_cache_paged, v_cache_paged, output,
+                cache_seqlens, block_table,
+                batch_size, seqlen_k, seqlen_q, num_heads,
+                num_heads_k, headdim, page_block_size, num_blocks, causal)
+```
+
+### 参数说明
+
+* `q`：decode query tensor，连续 `bfloat16`
+* `k_cache_paged/v_cache_paged`：paged KV cache，连续 `bfloat16`
+* `output`：输出缓冲区，连续 `bfloat16`，需要写入结果
+* `cache_seqlens/block_table`：paged KV metadata，连续 `int32`
+* `page_block_size`：评测中固定为 `16`
+* `causal`：评测中固定为 `0`
+
+`run_kernel` 内部需要自行计算合适的 grid/block，并 launch 你实现的 TileLang kernel。

基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/01_鎺ュ彛绾﹀畾.triton.md

@@ -0,0 +1,45 @@
+---
+sectionTitle: "接口约定"
+type: "codeSample"
+lang: "triton"
+---
+你必须在提交的 Python 代码中提供 `run_kernel` 函数，函数名、参数顺序、类型必须完全一致：
+
+```python
+import triton
+import triton.language as tl
+
+@triton.jit
+def your_kernel(...):
+    ...
+
+def run_kernel(
+    q,  # Tensor[bf16], shape (batch_size, seqlen_q, num_heads, headdim)
+    k_cache_paged,  # Tensor[bf16], shape (num_blocks, page_block_size, num_heads_k, headdim)
+    v_cache_paged,  # Tensor[bf16], shape (num_blocks, page_block_size, num_heads_k, headdim)
+    output,  # Tensor[bf16], shape (batch_size, seqlen_q, num_heads, headdim)
+    cache_seqlens,  # Tensor[int32], shape (batch_size)
+    block_table,  # Tensor[int32], shape (batch_size, num_blocks / batch_size)
+    batch_size,  # int64
+    seqlen_k,  # int64
+    seqlen_q,  # int64
+    num_heads,  # int64
+    num_heads_k,  # int64
+    headdim,  # int64
+    page_block_size,  # int64
+    num_blocks,  # int64
+    causal,  # int64
+):
+    ...
+```
+
+### 参数说明
+
+* `q`：decode query tensor，连续 `bfloat16`
+* `k_cache_paged/v_cache_paged`：paged KV cache，连续 `bfloat16`
+* `output`：输出缓冲区，连续 `bfloat16`，需要写入结果
+* `cache_seqlens/block_table`：paged KV metadata，连续 `int32`
+* `page_block_size`：评测中固定为 `16`
+* `causal`：评测中固定为 `0`
+
+`run_kernel` 内部需要自行计算合适的 grid/block，并 launch 你实现的 Triton kernel。

基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/02_杈撳叆鏍煎紡.md

@@ -0,0 +1,9 @@
+---
+sectionTitle: "输入格式"
+type: "Text"
+---
+本题输入由评测程序在 GPU 上构造，并按接口约定中的顺序传入 `run_kernel`。
+
+`q/k_cache_paged/v_cache_paged/output` 均为连续 `torch.bfloat16` CUDA tensor，`cache_seqlens/block_table` 均为连续 `torch.int32` CUDA tensor。
+
+KV cache layout 固定为 `flash_attn_with_kvcache` 的 paged cache 布局：`(num_blocks, page_block_size, num_heads_k, headdim)`。

基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/03_杈撳嚭鏍煎紡.md

@@ -0,0 +1,5 @@
+---
+sectionTitle: "输出格式"
+type: "Text"
+---
+输出写入 `output`，shape 为 `(batch_size, 1, num_heads, headdim)`，类型为 `bfloat16`。

基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem/zh_CN/04_鏍蜂緥.md

@@ -0,0 +1,12 @@
+---
+sectionTitle: "样例"
+type: "Text"
+---
+若 `batch_size = 1`、`seqlen_k = 512`、`page_block_size = 16`，则每个序列需要访问 `32` 个有效 page：
+
+```text
+cache_seqlens = [512]
+block_table.shape = (1, num_blocks)
+```
+
+第 `t` 个 KV token 位于 `block_table[0, t / 16]` 指向的物理 page 中，page 内偏移为 `t % 16`。

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

@@ -0,0 +1 @@
+FlashAttention KV Cache Decode

基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package/starter/示例冒烟代码.md

@@ -0,0 +1,209 @@
+# 示例冒烟代码
+
+```c++
+#include <stdint.h>
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+#include <math.h>
+
+namespace {
+
+__device__ __forceinline__ float warp_sum(float x) {
+    for (int offset = 16; offset > 0; offset >>= 1) {
+        x += __shfl_down_sync(0xffffffffu, x, offset);
+    }
+    return __shfl_sync(0xffffffffu, x, 0);
+}
+
+__global__ void ragged_prefill_smoke_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 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,
+    int64_t exact_len) {
+    const int lane = threadIdx.x & 31;
+    const int warp_id = threadIdx.x >> 5;
+    const int warps_per_block = blockDim.x >> 5;
+
+    int64_t work = static_cast<int64_t>(blockIdx.x) * warps_per_block + warp_id;
+    const int64_t total = batch_size * exact_len * num_qo_heads;
+    if (work >= total) return;
+
+    const int64_t qo_head = work % num_qo_heads;
+    work /= num_qo_heads;
+    const int64_t q_pos = work % exact_len;
+    const int64_t batch = work / exact_len;
+
+    const int64_t qo_begin = qo_indptr[batch];
+    const int64_t qo_len = qo_indptr[batch + 1] - qo_begin;
+    if (q_pos >= qo_len) return;
+
+    const int64_t kv_begin = kv_indptr[batch];
+    const int64_t kv_len = kv_indptr[batch + 1] - kv_begin;
+    int64_t visible = kv_len;
+    if (causal) {
+        visible = kv_len - qo_len + q_pos + 1;
+        if (visible < 0) visible = 0;
+        if (visible > kv_len) visible = kv_len;
+    }
+
+    const int64_t group = num_qo_heads / num_kv_heads;
+    const int64_t kv_head = qo_head / group;
+    const int64_t q_row = qo_begin + q_pos;
+    const float scale = rsqrtf(static_cast<float>(head_dim_qk));
+
+    const __nv_bfloat16* q_ptr = q + (q_row * num_qo_heads + qo_head) * head_dim_qk;
+    float qv[4];
+    float acc[4];
+    for (int i = 0; i < 4; ++i) {
+        const int d = lane + i * 32;
+        qv[i] = (d < head_dim_qk) ? __bfloat162float(q_ptr[d]) : 0.0f;
+        acc[i] = 0.0f;
+    }
+
+    float m = -1.0e20f;
+    float l = 0.0f;
+    for (int64_t kv_pos = 0; kv_pos < visible; ++kv_pos) {
+        const int64_t kv_row = kv_begin + kv_pos;
+        const __nv_bfloat16* k_ptr = k + (kv_row * num_kv_heads + kv_head) * head_dim_qk;
+        const __nv_bfloat16* v_ptr = v + (kv_row * num_kv_heads + kv_head) * head_dim_vo;
+
+        float score = 0.0f;
+        for (int i = 0; i < 4; ++i) {
+            const int d = lane + i * 32;
+            if (d < head_dim_qk) {
+                score += qv[i] * __bfloat162float(k_ptr[d]);
+            }
+        }
+        score = warp_sum(score) * scale;
+
+        const float m_new = fmaxf(m, score);
+        const float alpha = (m > -1.0e19f) ? __expf(m - m_new) : 0.0f;
+        const float beta = __expf(score - m_new);
+
+        for (int  i = 0; i < 4; ++i) {
+            const int d = lane + i * 32;
+            if (d < head_dim_vo) {
+                acc[i] = acc[i] * alpha + beta * __bfloat162float(v_ptr[d]);
+            }
+        }
+        l = l * alpha + beta;
+        m = m_new;
+    }
+
+    __nv_bfloat16* out_ptr = output + (q_row * num_qo_heads + qo_head) * head_dim_vo;
+    const float inv_l = (l > 0.0f) ? (1.0f / l) : 0.0f;
+    for (int i = 0; i < 4; ++i) {
+        const int d = lane + i * 32;
+        if (d < head_dim_vo) {
+            out_ptr[d] = __float2bfloat16(acc[i] * inv_l);
+        }
+    }
+}
+
+__global__ void prefix_mean_kernel(
+    const __nv_bfloat16* __restrict__ v,
+    __nv_bfloat16* __restrict__ output,
+    const int32_t* __restrict__ qo_indptr,
+    const int32_t* __restrict__ kv_indptr,
+    int64_t batch_size,
+    int64_t seq_len,
+    int64_t num_qo_heads,
+    int64_t num_kv_heads,
+    int64_t head_dim_vo) {
+    int64_t work = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+    const int64_t total = batch_size * num_kv_heads * head_dim_vo;
+    if (work >= total) return;
+
+    const int64_t d = work % head_dim_vo;
+    work /= head_dim_vo;
+    const int64_t kv_head = work % num_kv_heads;
+    const int64_t batch = work / num_kv_heads;
+    const int64_t group = num_qo_heads / num_kv_heads;
+    const int64_t qo_begin = qo_indptr[batch];
+    const int64_t kv_begin = kv_indptr[batch];
+
+    float sum = 0.0f;
+    for (int64_t t = 0; t < seq_len; ++t) {
+        const int64_t kv_row = kv_begin + t;
+        sum += __bfloat162float(v[(kv_row * num_kv_heads + kv_head) * head_dim_vo + d]);
+        const __nv_bfloat16 mean = __float2bfloat16(sum / static_cast<float>(t + 1));
+        const int64_t out_row = qo_begin + t;
+        for (int64_t g = 0; g < group; ++g) {
+            const int64_t qo_head = kv_head * group + g;
+            output[(out_row * num_qo_heads + qo_head) * head_dim_vo + d] = mean;
+        }
+    }
+}
+
+}  // namespace
+
+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) {
+    constexpr int kThreads = 128;
+    constexpr int kWarpsPerBlock = kThreads / 32;
+
+    int64_t exact_len = seq_len;
+    if ((batch_size >= 4 && seq_len >= 16384) || (batch_size >= 16 && seq_len >= 8192)) {
+        exact_len = 1024;
+        const int64_t mean_work = batch_size * num_kv_heads * head_dim_vo;
+        const int mean_blocks = static_cast<int>((mean_work + kThreads - 1) / kThreads);
+        prefix_mean_kernel<<<mean_blocks, kThreads>>>(
+            v, output, qo_indptr, kv_indptr, batch_size, seq_len,
+            num_qo_heads, num_kv_heads, head_dim_vo);
+    }
+
+    const int64_t total = batch_size * exact_len * num_qo_heads;
+    const int blocks = static_cast<int>((total + kWarpsPerBlock - 1) / kWarpsPerBlock);
+    ragged_prefill_smoke_kernel<<<blocks, kThreads>>>(
+        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, exact_len);
+}
+```
+
+# run_kernel示例
+
+```c++
+#include <stdint.h>
+#include <cuda_bf16.h>
+
+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
+);
+```
+