diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashAttention关键算子迁移与优化.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashAttention关键算子迁移与优化.md index 850452f..1aa5524 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashAttention关键算子迁移与优化.md +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashAttention关键算子迁移与优化.md @@ -190,8 +190,10 @@ 3. 创建完成后,进入算力容器,点击"工具-lab"即可打开 JupyterLab 终端开始项目创作。 +![image](https://origin.picgo.net/2026/06/23/image86811f30a57ae7b8.png) -\*\*说明:\*\*由于本次使用的是预装的专属镜像,环境中已经默认安装并配置好了 PyTorch、FlashAttention、einops 等所有依赖包。因此在启动实例后,无需再进行繁琐的依赖库版本验证即可直接进入测试环节。 + +**说明:** 由于本次使用的是预装的专属镜像,环境中已经默认安装并配置好了 PyTorch、FlashAttention、einops 等所有依赖包。因此在启动实例后,无需再进行繁琐的依赖库版本验证即可直接进入测试环节。 ### 4.2 代码准备 @@ -208,50 +210,50 @@ ### 5.1 关键术语 -* **KV-Cache:**缓存历史 Token 的 Key/Value 向量,避免 Transformer 推理时重复计算。 +* **KV-Cache:** 缓存历史 Token 的 Key/Value 向量,避免 Transformer 推理时重复计算。 -* **Paged KV-Cache:**将 KV-Cache 分页管理,减少显存碎片,提高利用率。 +* **Paged KV-Cache:** 将 KV-Cache 分页管理,减少显存碎片,提高利用率。 -* **Batch Size:**一次处理的样本数,越大并行度越高,但显存占用越大。 +* **Batch Size:** 一次处理的样本数,越大并行度越高,但显存占用越大。 -* **seq\_len\_kv:**KV-Cache 中已缓存的历史 Token 数量。 +* **seq\_len\_kv:** KV-Cache 中已缓存的历史 Token 数量。 -* **headdim:**注意力头维度,常见为 64/128/256。 +* **headdim:** 注意力头维度,常见为 64/128/256。 ### 5.2 核心知识 -* **正确性测试 :**验证算子输出结果的数学精度是否与标准实现一致,这是绝对底线。 +* **正确性测试 :** 验证算子输出结果的数学精度是否与标准实现一致,这是绝对底线。 -* **性能测试 :**在正确的前提下,测算速度与吞吐。通过建立性能基线(baseline),才能量化后续每次代码修改带来的真实收益(加速比)。 +* **性能测试 :** 在正确的前提下,测算速度与吞吐。通过建立性能基线(baseline),才能量化后续每次代码修改带来的真实收益(加速比)。 -* **Benchmark (基准测试):**在固定条件下反复运行同一任务,获取可重复的性能指标,用于建立基线、量化优化效果和定位瓶颈。 +* **Benchmark (基准测试):** 在固定条件下反复运行同一任务,获取可重复的性能指标,用于建立基线、量化优化效果和定位瓶颈。 * **XPU-OJ**:比赛官方在线评测平台,最终评测会调用参赛者提交代码中的 `run_kernel`。 ### 5.3 关键指标 -* **Kernel 执行时间:**GPU 核函数运行耗时(ms),使用 GPU 端同步计时获得。 +* **Kernel 执行时间:** GPU 核函数运行耗时(ms),使用 GPU 端同步计时获得。 -* **有效带宽:**数据传输量 (GB) ÷ Kernel 时间 (s),越接近理论峰值说明显存带宽利用越充分。 +* **有效带宽:** 数据传输量 (GB) ÷ Kernel 时间 (s),越接近理论峰值说明显存带宽利用越充分。 ### 5.4 其他要点 -* **Warmup:**预热若干次(不记录),使 GPU 进入稳定状态。 +* **Warmup:** 预热若干次(不记录),使 GPU 进入稳定状态。 -* **Repeat:**正式运行多次,取平均值或中位数以消除波动。 +* **Repeat:** 正式运行多次,取平均值或中位数以消除波动。 -* **同步:**调用 torch.cuda.synchronize() 确保精确计时。 +* **同步:** 调用 torch.cuda.synchronize() 确保精确计时。 -* **数据类型:**本教程使用 bfloat16,在精度和性能取得平衡。 +* **数据类型:** 本教程使用 bfloat16,在精度和性能取得平衡。 -* **显存占用估算:**KV-Cache ≈ batch × seq\_len\_kv × num\_heads\_k × headdim × 2(K+V) × 字节数。 +* **显存占用估算:** KV-Cache ≈ batch × seq\_len\_kv × num\_heads\_k × headdim × 2(K+V) × 字节数。 -* **OOM 应对:**减小 batch/seq\_len\_kv、使用更小 dtype 或释放中间变量。 +* **OOM 应对:** 减小 batch/seq\_len\_kv、使用更小 dtype 或释放中间变量。 -* **Tensor Core:**现代 GPU(含沐曦 C500)的矩阵乘法专用单元,要求维度对齐为 8 或 16 的倍数。 +* **Tensor Core:** 现代 GPU(含沐曦 C500)的矩阵乘法专用单元,要求维度对齐为 8 或 16 的倍数。 --- @@ -304,22 +306,22 @@ python -c "import einops; print('einops OK')" * `mx-smi` 显示沐曦 GPU 信息 -![016119b5 960b 478d 9119 af27b9f3c727](https://origin.picgo.net/2026/06/18/016119b5-960b-478d-9119-af27b9f3c727ac0af039c38a9f5e.png) +![image](https://origin.picgo.net/2026/06/23/image4591670d6c9f207d.png) -* Python 版本 = 3.8 +* Python 版本 = 3.12 -![568886e4 85a7 41c6 a94c 89a09913e0f7](https://origin.picgo.net/2026/06/18/568886e4-85a7-41c6-a94c-89a09913e0f712d9e076005617b3.png) +![image](https://origin.picgo.net/2026/06/23/image5f5fb78c43a0c220.png) * `torch.cuda.is_available()` 返回 `True` -![d1c97030 d40f 416f 89b3 b81caff88966](https://origin.picgo.net/2026/06/18/d1c97030-d40f-416f-89b3-b81caff88966e271eefaf7ce35bc.png) +![image](https://origin.picgo.net/2026/06/23/image10a748b701eed6f6.png) * 所有依赖版本符合要求 -![352b41ee 8a87 405e 991b 1c3714bdeff0](https://origin.picgo.net/2026/06/18/352b41ee-8a87-405e-991b-1c3714bdeff0e5fcd7b9b87d7b59.png) +![image](https://origin.picgo.net/2026/06/23/image62099cfafe955cac.png) **常见问题:** @@ -341,38 +343,22 @@ python -c "import einops; print('einops OK')" git clone https://gitlink.org.cn/metax-maca/op_optimization.git ``` -2. 准备flashattn\_baseline +2. 准备基准测试 - 在仓库目录 `op_optimization/基于AI Agent开发范式的国产GPU大模型推理算子库优化` 下,找到 `flashattn_baseline` 文件夹。可以将 `flashattn_baseline` 整个目录复制到工作目录 `data/` 下。 - - **下一步操作:** 切换到 Flashattn\_Baseline 项目目录。 + 在仓库目录 `op_optimization/基于AI Agent开发范式的国产GPU大模型推理算子库优化` 下,找到 `operator_task_package` 文件夹。切换到 flashattn_task_package 项目目录。 **命令示例:** ```bash - cd flashattn_baseline/Flashattn_Baselinels -la + cd op_optimization/基于AI\ Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/ + ls ``` **预期结果:** ```Plain - total xx - drwxr-xr-x 2 root root 4096 Jun 1 09:00 __MACOSX - - - rw-r--r-- 1 root root 5232 Jun 1 09:00 benchmark_kvcache.py - - - - rw-r--r-- 1 root root 1440 Jun 1 09:00 benchmark_kvcache_20260526_150953.csv - - ... - + baseline结果实例 benchmark_kvcache.py ``` - -3. 将Flashattn\_Baseline文件加入JupyterLab。 - - -![3eb4d4fe e948 405e 8d11 6301b4d8f16e](https://origin.picgo.net/2026/06/18/3eb4d4fe-e948-405e-8d11-6301b4d8f16e342580914be51bbf.png) - --- ### 6.4 步骤 3:配置基准测试参数 @@ -423,7 +409,6 @@ causal = True **命令示例:** ```Bash -cd flashattn_baseline python benchmark_kvcache.py ``` @@ -573,7 +558,7 @@ benchmark 脚本 用于理解目标算子的调用方式、输入输出 shape > 注意:每个子题的接口参数、数据范围和精度要求可能不同,正式要求以对应 XPU-OJ 题包为准。本节以 **FlashAttention KV Cache Decode** 为例,演示从 baseline benchmark 到 XPU-OJ 提交的完整流程。 -* **对应 benchmark 脚本**:`flashattn_baseline/baseline/benchmark_kvcache.py` +* **对应 benchmark 脚本**:`op_optimization/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmarkbenchmark_kvcache.py` * **对应 FlashAttention API**:`flash_attn_with_kvcache`(paged KV cache 版本) @@ -587,7 +572,7 @@ benchmark 脚本 用于理解目标算子的调用方式、输入输出 shape 完成 benchmark 后,需要注意:benchmark 脚本 主要用于建立性能基线,**不是最终提交物**。最终评测以 XPU-OJ 题包为准,评测程序会调用选手提交代码中的 `run_kernel`,并将输出结果与 OJ 参考结果进行比较。 -下面以 **FlashAttention KV Cache Decode** 题为例,题包目录中通常包含以下文件: +下面以 **FlashAttention KV Cache Decode** 题为例,题包目录（op_optimization/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/xpuoj_problem）中通常包含以下文件: * `zh_CN/00_题目描述.md`:说明需要实现的算子功能; @@ -693,13 +678,15 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) 5.点击进入题目详情页,查看题目描述、接口约定、数据范围和提交入口。 +![image](https://origin.picgo.net/2026/06/23/image123512a0b877cd31.png) + ### 6.11 步骤 10:在Agent的帮助下提交 OJ 冒烟代码 -**目标:**完成一次最小提交,确认 OJ 提交链路、语言环境和 `run_kernel(...)` 接口可用。 +**目标:** 完成一次最小提交,确认 OJ 提交链路、语言环境和 `run_kernel(...)` 接口可用。 **操作:** -1. 在语言下拉框中选择本题支持的提交语言,例如 `MXMACA C++`、`TileLang` 、 `Triton`; +1. 在语言下拉框中选择本题支持的提交语言,例如 `CUDA Maca`、`TileLang` 、 `Triton`，本次提供的冒烟代码使用的是`CUDA Maca`; 2. 将实现了题目要求接口的代码复制到提交框中; @@ -763,17 +750,14 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) ![095af1c8 aa7e 44e3 ad0e 620206aaba1e](https://origin.picgo.net/2026/06/18/095af1c8-aa7e-44e3-ad0e-620206aaba1e23a43b2eb7b0ead3.png) - 为方便参赛者先跑通完整流程,这里直接提供一份完整的冒烟代码,可直接复制粘贴到右侧编辑器,用于验证提交链路是否正常: + 为方便参赛者先跑通完整流程,这里直接提供一份完整的冒烟代码，使用语言为`CUDA Maca`,可直接复制粘贴到右侧编辑器,用于验证提交链路是否正常: - ```python + ```C++ #include #include - #include + #include - #define PAGE_SIZE 16 - #define HEAD_DIM 128 - - __global__ void paged_attention_kernel( + __global__ void paged_kv_decode_kernel( const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k_cache_paged, const __nv_bfloat16* __restrict__ v_cache_paged, @@ -781,119 +765,156 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) const int32_t* __restrict__ cache_seqlens, const int32_t* __restrict__ 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 blocks_per_batch) + int64_t max_num_blocks_per_seq) { - 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 b = blockIdx.x; + int h = blockIdx.y; + int tx = threadIdx.x; + int ty = threadIdx.y; + int num_warps = blockDim.y; - int seqlen = cache_seqlens[batch_idx]; - if (seqlen <= 0) { - // 无有效 KV,输出 0 - int64_t out_base = ((batch_idx * seqlen_q + 0) * num_heads + head_idx) * headdim; - for (int i = threadIdx.x; i < headdim; i += blockDim.x) - output[out_base + i] = __float2bfloat16(0.0f); + // 当前 batch 的有效 KV token 长度 + int seqlen = cache_seqlens[b]; + + // 动态共享内存布局 + // q_smem [headdim] -> smem_m [num_warps] -> smem_d [num_warps] -> smem_acc [num_warps * headdim] + extern __shared__ char smem[]; + float* q_smem = (float*)smem; + float* smem_m = (float*)(q_smem + headdim); + float* smem_d = (float*)(smem_m + num_warps); + float* smem_acc = (float*)(smem_d + num_warps); + + // 计算 headdim 需要分几个 float (每 32 元素处理一次) + int num_iters = (headdim + 31) / 32; + + if (seqlen == 0) { + if (ty == 0) { + for (int step = 0; step < num_iters; ++step) { + int i = step * 32 + tx; + if (i < headdim) { + output[(int64_t)b * num_heads * headdim + (int64_t)h * headdim + i] = __float2bfloat16(0.0f); + } + } + } return; } - 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]); - - // 全局 softmax 状态(每个线程维护自己维度的累加器) - float max_val = -1e38f; - float sum_exp = 0.0f; - float out_acc = 0.0f; - float scale = rsqrtf(static_cast(headdim)); - - // 共享内存布局: - // K_tile[PAGE_SIZE][HEAD_DIM] (bf16) - // V_tile[PAGE_SIZE][HEAD_DIM] (bf16) - // partial_scores[PAGE_SIZE][HEAD_DIM] (float, 用于归约点积) - __shared__ __nv_bfloat16 K_tile[PAGE_SIZE][HEAD_DIM]; - __shared__ __nv_bfloat16 V_tile[PAGE_SIZE][HEAD_DIM]; - __shared__ float partial_scores[PAGE_SIZE][HEAD_DIM]; - - int total_pages = (seqlen + PAGE_SIZE - 1) / PAGE_SIZE; - - for (int page = 0; page < total_pages; ++page) { - int physical_block = block_table[batch_idx * blocks_per_batch + page]; - int tokens_this_page = min(seqlen - page * PAGE_SIZE, PAGE_SIZE); - - // 1. 将当前 page 的 K 和 V 从全局显存加载到共享内存 - // 每个线程负责加载所有 token 的同一个 head 维度 - const int64_t kv_stride = num_heads_k * headdim; // 每个 (block, offset) 的 stride - for (int j = 0; j < tokens_this_page; ++j) { - int64_t offset = (physical_block * PAGE_SIZE + j) * kv_stride + head_idx * headdim + tid; - K_tile[j][tid] = k_cache_paged[offset]; - V_tile[j][tid] = v_cache_paged[offset]; - } - __syncthreads(); - - // 2. 计算该 page 内每个 token 与 Q 的部分点积,存入 partial_scores - for (int j = 0; j < tokens_this_page; ++j) { - float k_val = __bfloat162float(K_tile[j][tid]); - partial_scores[j][tid] = q_val * k_val; - } - __syncthreads(); - - // 3. 对 partial_scores 做 tree reduction,得到每个 token 的完整点积 - #pragma unroll - for (int stride = HEAD_DIM / 2; stride > 0; stride >>= 1) { - if (tid < stride) { - #pragma unroll - for (int j = 0; j < tokens_this_page; ++j) { - partial_scores[j][tid] += partial_scores[j][tid + stride]; - } + // 第 0 个 warp 将 Query 读取至 Shared Memory + if (ty == 0) { + for (int step = 0; step < num_iters; ++step) { + int i = step * 32 + tx; + if (i < headdim) { + int64_t q_idx = (int64_t)b * num_heads * headdim + (int64_t)h * headdim + i; + q_smem[i] = __bfloat162float(q[q_idx]); } - __syncthreads(); } + } + __syncthreads(); - // 4. 在线 safe softmax 更新 + V 累加 - // 4.1 找出本 page 内点积的最大值,结合全局 max 得到 new_max - float local_max = -1e38f; - #pragma unroll - for (int j = 0; j < tokens_this_page; ++j) { - local_max = fmaxf(local_max, partial_scores[j][0]); - } - float new_max = fmaxf(max_val, local_max); - - // 4.2 用旧的 max 对全局状态进行重缩放 - float rescale = expf(max_val - new_max); - sum_exp *= rescale; - out_acc *= rescale; - max_val = new_max; - - // 4.3 累加 V,并更新 sum_exp - #pragma unroll - for (int j = 0; j < tokens_this_page; ++j) { - float score = partial_scores[j][0] * scale; - float weight = expf(score - new_max); - sum_exp += weight; - - float v_val = __bfloat162float(V_tile[j][tid]); - out_acc += weight * v_val; - } - - __syncthreads(); // 准备下一个 page 的共享内存加载 + // 独立维持 FlashAttention Softmax 局部状态 + float m_warp = -1e20f; + float d_warp = 0.0f; + float acc[32]; // 支持 max_headdim = 32*32=1024 (安全冗余) + #pragma unroll + for (int i = 0; i < 32; ++i) { + acc[i] = 0.0f; } - // 5. 最终归一化并写回 - if (sum_exp > 0.0f) { - out_acc /= sum_exp; - } else { - out_acc = 0.0f; + float scale = 1.0f / sqrtf((float)headdim); + int kv_h = h / (num_heads / num_heads_k); // MQA 或 GQA 支持映射 + + // 每个 Warp 以跨步的方式 (stride = num_warps) 并行消化长 Token + for (int t = ty; t < seqlen; t += num_warps) { + int page_idx = t / page_block_size; + int page_offset = t % page_block_size; + int block_id = block_table[b * max_num_blocks_per_seq + page_idx]; + + int64_t k_base = (int64_t)block_id * (page_block_size * num_heads_k * headdim) + + (int64_t)page_offset * (num_heads_k * headdim) + + (int64_t)kv_h * headdim; + + // 计算当前 token (K) 与 Q 的内积 + float score = 0.0f; + for (int step = 0; step < num_iters; ++step) { + int i = step * 32 + tx; + if (i < headdim) { + float k_val = __bfloat162float(k_cache_paged[k_base + i]); + score += q_smem[i] * k_val; + } + } + score *= scale; + + // Warp 级归约累加 + for (int mask = 16; mask > 0; mask /= 2) { + score += __shfl_xor_sync(0xffffffff, score, mask); + } + + // Online Softmax: 维护最大值与底数和 + float m_old = m_warp; + m_warp = fmaxf(m_warp, score); + float exp_val = expf(score - m_warp); + float exp_old = expf(m_old - m_warp); + d_warp = d_warp * exp_old + exp_val; + + // 同步 V 计算与更新 + int64_t v_base = (int64_t)block_id * (page_block_size * num_heads_k * headdim) + + (int64_t)page_offset * (num_heads_k * headdim) + + (int64_t)kv_h * headdim; + + for (int step = 0; step < num_iters; ++step) { + int i = step * 32 + tx; + if (i < headdim) { + float v_val = __bfloat162float(v_cache_paged[v_base + i]); + acc[step] = acc[step] * exp_old + exp_val * v_val; + } + } } - int64_t out_offset = ((batch_idx * seqlen_q + 0) * num_heads + head_idx) * headdim + tid; - output[out_offset] = __float2bfloat16(out_acc); + // Warp 将局部运算结果写入共享内存 + if (tx == 0) { + smem_m[ty] = m_warp; + smem_d[ty] = d_warp; + } + for (int step = 0; step < num_iters; ++step) { + int i = step * 32 + tx; + if (i < headdim) { + smem_acc[ty * headdim + i] = acc[step]; + } + } + __syncthreads(); + + // Block 层级化简：Warp 0 收敛所有 Warp 的局部统计信息并写入 output + if (ty == 0) { + float global_m = -1e20f; + for (int w = 0; w < num_warps; ++w) { + global_m = fmaxf(global_m, smem_m[w]); + } + + float global_d = 0.0f; + for (int w = 0; w < num_warps; ++w) { + global_d += smem_d[w] * expf(smem_m[w] - global_m); + } + + // 缩放加权合并所有局部的 acc V 值，写入最后结果 + for (int step = 0; step < num_iters; ++step) { + int i = step * 32 + tx; + if (i < headdim) { + float global_acc = 0.0f; + for (int w = 0; w < num_warps; ++w) { + global_acc += smem_acc[w * headdim + i] * expf(smem_m[w] - global_m); + } + float out_val = global_acc / global_d; + int64_t out_idx = (int64_t)b * num_heads * headdim + (int64_t)h * headdim + i; + output[out_idx] = __float2bfloat16(out_val); + } + } + } } extern "C" void run_kernel( @@ -913,15 +934,23 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) 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); + // 配置 Kernel 计算资源与 Layout + int num_warps = 8; + dim3 block(32, num_warps); + dim3 grid(batch_size, num_heads); - paged_attention_kernel<<>>( + // 动态计算需要的共享内存大小 (Q + Metadata + 归约缓存 V) + size_t smem_size = (headdim + num_warps * 2 + num_warps * headdim) * sizeof(float); + + // 基于约定，block_table shape 为 (batch_size, num_blocks / batch_size) + int64_t max_num_blocks_per_seq = num_blocks / batch_size; + + paged_kv_decode_kernel<<>>( 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 + batch_size, seqlen_k, seqlen_q, + num_heads, num_heads_k, headdim, + page_block_size, max_num_blocks_per_seq ); } ``` @@ -952,9 +981,24 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) 6. 查看结果 - 提交详情会显示状态、得分、时间、内存、编译信息以及各测试点结果。 + 提交详情会显示状态、总得分、时间、内存、编译信息以及各测试点结果。 - ![cd119ff85c024846b91d69b9ca9a1e00.png](https://origin.picgo.net/2026/06/17/cd119ff85c024846b91d69b9ca9a1e006456033e499779c6.png) +![image](https://origin.picgo.net/2026/06/23/imageba30dbcea3e96f38.png) + +**单测试点查看：** +![image](https://origin.picgo.net/2026/06/23/image073f38d8acbf295c.png) + +**榜单查看** +![image](https://origin.picgo.net/2026/06/23/image4035d30ecf0356f4.png) + +**优化思路： ** +- **算子融合**：将矩阵乘、Scale、Softmax 等独立步骤合并为单个内核，避免中间结果在全局显存（HBM）中的反复读写。 +- **并行策略调整**：在 Decode 阶段采用 Split-K（如 Flash-Decoding），将长 KV 序列分块分配给多个线程块并行计算，解决单 Query 并行度极低的问题。 +- **在线 Softmax**：引入局部最大值和局部求和动态缩放历史累加器，实现单遍计算，完美契合算子融合需求并保证数值稳定。 +- **显存访问合并**：确保同一线程束内相邻线程访问相邻的显存地址，按连续维度（如 HeadDim）进行向量化加载，最大化利用显存带宽。 +- **利用内存层级**：将频繁更新的标量存入寄存器，需块内复用的数据存入共享内存，仅在全局显存放输入输出，缩短数据访问延迟。 +- **软件流水线**：通过异步拷贝等技术，在当前分块计算的同时预取下一分块数据，将内存加载延迟隐藏在计算时间中。 +- **减少冗余计算与分支**：将循环不变量（如 Q 向量的加载和 Scale 操作）提取到循环外部；在处理变长序列时，使用高效的掩码指令替代复杂的 `if-else` 分支。 ## 7. 从性能基线到参赛作品的路径回顾 @@ -995,19 +1039,19 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) ### 8.5 Q5: 评测状态显示 `Compile Error` 或提示 `Undefined reference to run_kernel` -**原因:**C++函数名被修饰(Name Mangling)或参数类型/顺序与接口约定不符。 **解决:**在 run\_kernel 前添加 extern "C",并严格逐字核对参数的类型和修饰符。 +**原因:** C++函数名被修饰(Name Mangling)或参数类型/顺序与接口约定不符。 **解决:**在 run\_kernel 前添加 extern "C",并严格逐字核对参数的类型和修饰符。 ### 8.6 Q6: 评测状态显示 `Wrong Answer`，提示 torch.allclose 校验失败 -**原因:**bfloat16精度截断溢出、线程同步缺失或无效Token(Padding区域)处理错误。 **解决:**累加和 Softmax 强制转为 float32 计算;检查 \_\_syncthreads() 逻辑;增加 seqlen 的边界判空。 +**原因:** bfloat16精度截断溢出、线程同步缺失或无效Token(Padding区域)处理错误。 **解决:**累加和 Softmax 强制转为 float32 计算;检查 \_\_syncthreads() 逻辑;增加 seqlen 的边界判空。 ### 8.7 Q7: 评测状态显示 `Runtime Error`（非法内存访问/段错误） -**原因:**Paged KV 地址映射索引错误、尾部 Page 越界读取,或线程块维度超限。 **解决:**仔细核对物理块寻址公式;增加当前 Page 有效 Token 数量的越界判断;检查单 Block 线程数配置。 +**原因:** Paged KV 地址映射索引错误、尾部 Page 越界读取,或线程块维度超限。 **解决:**仔细核对物理块寻址公式;增加当前 Page 有效 Token 数量的越界判断;检查单 Block 线程数配置。 ### 8.8 Q8: 评测状态显示 `Time Limit Exceeded`（评测超时） -**原因:**发散分支内的 \_\_syncthreads() 导致内核死锁、误加主机端同步指令或并行度划分错误导致串行。 **解决:**确保同步指令在所有线程必经路径上;移除主机端多余的 cudaDeviceSynchronize();优化 <<>> 参数以提升并行度。 +**原因:** 发散分支内的 \_\_syncthreads() 导致内核死锁、误加主机端同步指令或并行度划分错误导致串行。 **解决:**确保同步指令在所有线程必经路径上;移除主机端多余的 cudaDeviceSynchronize();优化 <<>> 参数以提升并行度。 ## 9. 下一步学习建议 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/starter/guide.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/starter/guide.md deleted file mode 100644 index 4fc288c..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/starter/guide.md +++ /dev/null @@ -1,330 +0,0 @@ -# guide - -# FlashAttention KV Cache Decode - 参赛指南 - -## 一、登录 XPU-OJ 平台 - -打开浏览器，访问：\*\*https://xpuoj.com/\*\* - -1. 等待组委会统一发放 XPU-OJ 账号; - -2. 使用分配的用户名和初始密码登录平台 - - -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/4jKqm0bXGBkPvnw1/img/a3af510d-2b4a-4fa0-8e0f-5137ec2cc2fe.png) - -点击 **"登录"** 进入平台。 - -## 二、进入比赛 - -### 2.1 点击导航栏"比赛" - -登录成功后,会跳转到 XPUOJ 平台首页。 - -找到页面顶部的导航栏,点击 **"比赛"** 标签,进入比赛列表页。 - -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/4jKqm0bXGBkPvnw1/img/d3c48202-7e3f-48a1-8a9e-e96192e51256.png) - -### 2.2 找到目标比赛 - -在比赛列表中,根据比赛状态(全部 / 未开始 / 进行中 / 已结束)找到目标比赛。 - -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/4jKqm0bXGBkPvnw1/img/bff83a88-21ed-4766-9450-6db58a83df43.png) - -## 三、进入题目 - -### 3.1 在题目列表中找到目标题目 - -进入比赛后,页面会展示该比赛的**题目列表**。每道题都有编号(1 ~ N)和标题。 - -在列表中找到 **FlashAttention KV Cache Decode** 这道题(前缀为 `FlashAttention`),点击标题即可进入题目详情页。 - -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/4jKqm0bXGBkPvnw1/img/cc304e39-753f-443f-9267-6b14c3ec99da.png) - -### 3.2 查看题目要求 - -进入题目详情页后,可以看到以下几个区域: - -* **左侧**:题目描述、接口约定、参数说明 - -* **右侧**:代码编辑器,用于编写并提交代码 - - -请仔细阅读左侧的 **题目描述** 和 **接口约定**,重点关注: - -* 入口函数名(本题为 `run_kernel`) - -* 必传的参数列表及其类型、顺序 - -* 编译/运行环境(语言选择,目标硬件 C500) - - -### 3.3 编写并提交代码 - -在右侧的代码编辑器中,按照题目要求填入完整代码(可点击右上角"重置"恢复初始模板)。 - -#### 最小正确性代码(可先复制跑通) - -为方便参赛者先跑通完整流程,这里提供一份 **最小正确性代码**,可直接复制粘贴到右侧编辑器,用于验证提交链路是否正常: - -```cpp - -#include - -#include - -#include - - -#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(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<<>>( - - 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 - - ); - -} -``` - -完成后: - -1. 在编辑器下方 **语言** 选项中,根据代码实际情况选择对应语言(可选 `Triton` / `CUDA Maca` / `MXMACA C++` / `TileLang` 等) - -2. 选择 **目标硬件** 为 `C500` - -3. 点击右上角 **"提交"** 按钮 - - -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/4jKqm0bXGBkPvnw1/img/d4dba895-a1fd-414c-832d-86111734378b.png) - -## 四、查看提交结果 - -点击 **"提交"** 后,系统会自动跳转到提交记录页,展示本次提交的详细信息。 - -页面顶部会显示一行汇总信息,包括: - -* **状态** —— 评测结果(如 `Accepted` 表示通过) - -* **分数** —— 本次提交获得的分数 - -* **题目** —— 对应的题目名称 - -* **用时** —— 程序运行耗时 - -* **内存** —— 占用内存大小 - -* **答案** —— 提交所用的语言/硬件 - -* **提交时间** —— 提交的时刻 - - -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/4jKqm0bXGBkPvnw1/img/14eedaa5-8072-4247-a85b-5a12b284d23c.png) - -页面下方会展开 **编译信息** 和 **各测试点**(样例、测试点 #1 ~ #N)的结果,逐个显示: - -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/4jKqm0bXGBkPvnw1/img/f1a3c988-c26a-4c30-b32b-f4e10d7270e2.png) \ No newline at end of file