修复教程格式，更新内容 #45
|
|
@ -190,8 +190,10 @@
|
|||
|
||||
3. 创建完成后,进入算力容器,点击"工具-lab"即可打开 JupyterLab 终端开始项目创作。
|
||||
|
||||
|
||||
|
||||
\*\*说明:\*\*由于本次使用的是预装的专属镜像,环境中已经默认安装并配置好了 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 信息
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
* Python 版本 = 3.8
|
||||
* Python 版本 = 3.12
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
* `torch.cuda.is_available()` 返回 `True`
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
* 所有依赖版本符合要求
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
**常见问题:**
|
||||
|
||||
|
|
@ -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。
|
||||
|
||||
|
||||
|
||||
|
||||
---
|
||||
|
||||
### 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.点击进入题目详情页,查看题目描述、接口约定、数据范围和提交入口。
|
||||
|
||||
|
||||
|
||||
### 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)
|
|||
|
||||
|
||||
|
||||
为方便参赛者先跑通完整流程,这里直接提供一份完整的冒烟代码,可直接复制粘贴到右侧编辑器,用于验证提交链路是否正常:
|
||||
为方便参赛者先跑通完整流程,这里直接提供一份完整的冒烟代码，使用语言为`CUDA Maca`,可直接复制粘贴到右侧编辑器,用于验证提交链路是否正常:
|
||||
|
||||
```python
|
||||
```C++
|
||||
#include <stdint.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
|
||||
#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<float>(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<<<grid, block>>>(
|
||||
// 动态计算需要的共享内存大小 (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<<<grid, block, smem_size>>>(
|
||||
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. 查看结果
|
||||
|
||||
提交详情会显示状态、得分、时间、内存、编译信息以及各测试点结果。
|
||||
提交详情会显示状态、总得分、时间、内存、编译信息以及各测试点结果。
|
||||
|
||||
|
||||
|
||||
|
||||
**单测试点查看：**
|
||||
|
||||
|
||||
**榜单查看**
|
||||
|
||||
|
||||
**优化思路： **
|
||||
- **算子融合**：将矩阵乘、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();优化 <<<grid, block>>> 参数以提升并行度。
|
||||
**原因:** 发散分支内的 \_\_syncthreads() 导致内核死锁、误加主机端同步指令或并行度划分错误导致串行。 **解决:**确保同步指令在所有线程必经路径上;移除主机端多余的 cudaDeviceSynchronize();优化 <<<grid, block>>> 参数以提升并行度。
|
||||
|
||||
## 9. 下一步学习建议
|
||||
|
||||
|
|
|
|||
|
|
@ -1,330 +0,0 @@
|
|||
# guide
|
||||
|
||||
# FlashAttention KV Cache Decode - 参赛指南
|
||||
|
||||
## 一、登录 XPU-OJ 平台
|
||||
|
||||
打开浏览器，访问：\*\*https://xpuoj.com/\*\*
|
||||
|
||||
1. 等待组委会统一发放 XPU-OJ 账号;
|
||||
|
||||
2. 使用分配的用户名和初始密码登录平台
|
||||
|
||||
|
||||
|
||||
|
||||
点击 **"登录"** 进入平台。
|
||||
|
||||
## 二、进入比赛
|
||||
|
||||
### 2.1 点击导航栏"比赛"
|
||||
|
||||
登录成功后,会跳转到 XPUOJ 平台首页。
|
||||
|
||||
找到页面顶部的导航栏,点击 **"比赛"** 标签,进入比赛列表页。
|
||||
|
||||
|
||||
|
||||
### 2.2 找到目标比赛
|
||||
|
||||
在比赛列表中,根据比赛状态(全部 / 未开始 / 进行中 / 已结束)找到目标比赛。
|
||||
|
||||
|
||||
|
||||
## 三、进入题目
|
||||
|
||||
### 3.1 在题目列表中找到目标题目
|
||||
|
||||
进入比赛后,页面会展示该比赛的**题目列表**。每道题都有编号(1 ~ N)和标题。
|
||||
|
||||
在列表中找到 **FlashAttention KV Cache Decode** 这道题(前缀为 `FlashAttention`),点击标题即可进入题目详情页。
|
||||
|
||||
|
||||
|
||||
### 3.2 查看题目要求
|
||||
|
||||
进入题目详情页后,可以看到以下几个区域:
|
||||
|
||||
* **左侧**:题目描述、接口约定、参数说明
|
||||
|
||||
* **右侧**:代码编辑器,用于编写并提交代码
|
||||
|
||||
|
||||
请仔细阅读左侧的 **题目描述** 和 **接口约定**,重点关注:
|
||||
|
||||
* 入口函数名(本题为 `run_kernel`)
|
||||
|
||||
* 必传的参数列表及其类型、顺序
|
||||
|
||||
* 编译/运行环境(语言选择,目标硬件 C500)
|
||||
|
||||
|
||||
### 3.3 编写并提交代码
|
||||
|
||||
在右侧的代码编辑器中,按照题目要求填入完整代码(可点击右上角"重置"恢复初始模板)。
|
||||
|
||||
#### 最小正确性代码(可先复制跑通)
|
||||
|
||||
为方便参赛者先跑通完整流程,这里提供一份 **最小正确性代码**,可直接复制粘贴到右侧编辑器,用于验证提交链路是否正常:
|
||||
|
||||
```cpp
|
||||
|
||||
#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
|
||||
|
||||
);
|
||||
|
||||
}
|
||||
```
|
||||
|
||||
完成后:
|
||||
|
||||
1. 在编辑器下方 **语言** 选项中,根据代码实际情况选择对应语言(可选 `Triton` / `CUDA Maca` / `MXMACA C++` / `TileLang` 等)
|
||||
|
||||
2. 选择 **目标硬件** 为 `C500`
|
||||
|
||||
3. 点击右上角 **"提交"** 按钮
|
||||
|
||||
|
||||
|
||||
|
||||
## 四、查看提交结果
|
||||
|
||||
点击 **"提交"** 后,系统会自动跳转到提交记录页,展示本次提交的详细信息。
|
||||
|
||||
页面顶部会显示一行汇总信息,包括:
|
||||
|
||||
* **状态** —— 评测结果(如 `Accepted` 表示通过)
|
||||
|
||||
* **分数** —— 本次提交获得的分数
|
||||
|
||||
* **题目** —— 对应的题目名称
|
||||
|
||||
* **用时** —— 程序运行耗时
|
||||
|
||||
* **内存** —— 占用内存大小
|
||||
|
||||
* **答案** —— 提交所用的语言/硬件
|
||||
|
||||
* **提交时间** —— 提交的时刻
|
||||
|
||||
|
||||
|
||||
|
||||
页面下方会展开 **编译信息** 和 **各测试点**(样例、测试点 #1 ~ #N)的结果,逐个显示:
|
||||
|
||||
|
||||
Loading…
Reference in New Issue