forked from metax-maca/op_optimization
Merge pull request '更新XPU-OJ教程并修复合并冲突' (#25) from xiao-ke/op_optimization:master into master
This commit is contained in:
commit
59a4fb8f8c
|
|
@ -11,6 +11,8 @@
|
|||
3. 跑通 BatchDecode、BatchPrefill、MLA 等典型算子的 API 调用示例；
|
||||
4. 完成各算子在不同参数配置下的 Benchmark 测试；
|
||||
5. 输出各算子的 Baseline 性能结果记录表，为后续算子优化提供对比基准。
|
||||
6. 理解 XPU-OJ 评测`run_kernel`接口与精度要求。
|
||||
7. 理解 Baseline 与 XPU-OJ 评测题包之间的关系，能够根据题包接口实现一个最小正确版 run_kernel。
|
||||
|
||||
|
||||
## 三、适用对象
|
||||
|
|
@ -39,9 +41,9 @@
|
|||
* 步骤1：获取算力券[https://developer.metax-tech.com/activities/6](https://developer.metax-tech.com/activities/6)
|
||||
|
||||
* 首次登录需要先进行注册（使用邮箱或者手机号进行注册）
|
||||
|
||||
|
||||
* 登录成功后进行第二步-邮箱验证，填入自己的邮箱。
|
||||
|
||||
|
||||
* 第三步，提交申请。
|
||||
|
||||
* 获得兑换码
|
||||
|
|
@ -124,7 +126,9 @@ prefill = 并行处理用户输入，decode = 逐个生成回答 token
|
|||
|
||||
## 六、项目实践--FlashInfer-Baseline
|
||||
|
||||
### Step 1：检查运行环境
|
||||
## 本地项目运行
|
||||
|
||||
#### Step 1：检查运行环境
|
||||
|
||||
**目标：** 进入terminal确认当前环境满足本模块运行要求。
|
||||
|
||||
|
|
@ -140,7 +144,7 @@ prefill = 并行处理用户输入，decode = 逐个生成回答 token
|
|||
| `No GPUs are available` | 检查 MXMACA 驱动是否正确安装 |
|
||||
| `ModuleNotFoundError: No module named 'flashinfer'` | 执行 `pip install flashinfer` 安装 |
|
||||
|
||||
### Step 2：进入项目目录
|
||||
#### Step 2：进入项目目录
|
||||
|
||||
**目标：** 进入本模块所需的源码目录[flashinfer_baseline](https://gitlink.org.cn/metax-maca/op_optimization/tree/master/%E5%9F%BA%E4%BA%8EAI%20Agent%E5%BC%80%E5%8F%91%E8%8C%83%E5%BC%8F%E7%9A%84%E5%9B%BD%E4%BA%A7GPU%E5%A4%A7%E6%A8%A1%E5%9E%8B%E6%8E%A8%E7%90%86%E7%AE%97%E5%AD%90%E5%BA%93%E4%BC%98%E5%8C%96%2Fbaselines%2Fflashinfer_baseline)。
|
||||
|
||||
|
|
@ -154,12 +158,12 @@ prefill = 并行处理用户输入，decode = 逐个生成回答 token
|
|||
|
||||
在仓库目录 `op_optimization/基于AI Agent开发范式的国产GPU大模型推理算子库优化` 下，找到 `flashinfer_baseline` 文件夹。可以将 `flashinfer_baseline` 整个目录复制到工作目录 `data/` 下。
|
||||
|
||||
**下一步操作：** 切换到 FlashInfer Baseline 项目目录。
|
||||
**下一步操作：** 切换到 FlashInfer_Baseline 项目目录。
|
||||
|
||||
**命令示例：**
|
||||
|
||||
```bash
|
||||
cd /data/flashinfer_baseline
|
||||
cd baselines\flashinfer_baseline\FlashInfer_Baseline
|
||||
ls -la
|
||||
```
|
||||
|
||||
|
|
@ -176,7 +180,7 @@ README.md
|
|||
|
||||
```
|
||||
|
||||
### Step 3：验证项目脚本
|
||||
#### Step 3：验证项目脚本
|
||||
|
||||
**目标：** 确认所有基准测试脚本可正常执行。
|
||||
|
||||
|
|
@ -205,7 +209,7 @@ python -c "from bench_common import setup_workspace, get_csv_path; print('脚本
|
|||
|
||||
```
|
||||
|
||||
### Step 4：运行单算子 Benchmark
|
||||
#### Step 4：运行单算子 Benchmark
|
||||
|
||||
**目标：** 执行基准测试，获取 Baseline 性能数据。
|
||||
|
||||
|
|
@ -237,7 +241,7 @@ Results saved to BatchDecodeWithPagedKVCacheWrapper_20260525_xxxxxx.csv
|
|||
| `out of memory` | 减小 batch\_size 或 seq\_len 参数 |
|
||||
| 运行时间过长 | 脚本会自动调整重复次数，耐心等待 |
|
||||
|
||||
### Step 5：查看测试结果
|
||||
#### Step 5：查看测试结果
|
||||
|
||||
**目标：** 查看并分析 Benchmark 输出结果。
|
||||
|
||||
|
|
@ -272,7 +276,226 @@ if csv_files:
|
|||
|
||||
```
|
||||
|
||||
### **XPU-OJ 评测与提交**
|
||||
|
||||
#### Step 1：访问XPUOJ平台(https://xpuoj.com/)进行登录
|
||||
|
||||
- 完成 baseline benchmark 后，要到统一评测工具XPU-OJ完成提交，进行评测。
|
||||
|
||||
- 等待组委会统一发放 XPU-OJ 账号
|
||||
- 使用分配的用户名和初始密码登录平台
|
||||
|
||||
|
||||
|
||||
#### Step 2：进入比赛页面
|
||||
|
||||
- 登录后进入对应赛题页面
|
||||
- 选择对应的赛题
|
||||
|
||||
|
||||
|
||||
#### Step 3：选择对应的算子题目进行评测
|
||||
|
||||
|
||||
|
||||
#### Step 4：阅读题目要求，接口约定，编写符合`run_kernel`接口规范的CUDA算子
|
||||
|
||||
以下代码仅用于说明接口结构，不代表最优实现，也不作为评分参考：
|
||||
|
||||
```python
|
||||
FlashInfer Ragged Prefill示例代码
|
||||
#include <stdint.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
|
||||
#ifndef max
|
||||
#define max(a,b) ((a) > (b) ? (a) : (b))
|
||||
#endif
|
||||
|
||||
#ifndef min
|
||||
#define min(a,b) ((a) < (b) ? (a) : (b))
|
||||
#endif
|
||||
|
||||
template <int D_QK, int D_VO, int Br, int Bc>
|
||||
__global__ void ragged_prefill_kernel(
|
||||
const __nv_bfloat16* __restrict__ q,
|
||||
const __nv_bfloat16* __restrict__ k,
|
||||
const __nv_bfloat16* __restrict__ v,
|
||||
__nv_bfloat16* __restrict__ output,
|
||||
const int32_t* __restrict__ qo_indptr,
|
||||
const int32_t* __restrict__ kv_indptr,
|
||||
int64_t seq_len,
|
||||
int64_t num_qo_heads,
|
||||
int64_t num_kv_heads,
|
||||
float scale,
|
||||
int64_t causal
|
||||
) {
|
||||
int b = blockIdx.x;
|
||||
int hq = blockIdx.y;
|
||||
int hkv = hq / (num_qo_heads / num_kv_heads);
|
||||
|
||||
int ty = threadIdx.y; // 0~31, Block 内的 query 索引 (对应 Br)
|
||||
int tx = threadIdx.x; // 0~31, Warp 内的线程 (lane)
|
||||
|
||||
// 使用 int64_t 防止长序列+大Batch组合下的整数溢出
|
||||
int64_t q_idx = (int64_t)blockIdx.z * Br + ty;
|
||||
bool is_valid = q_idx < seq_len;
|
||||
|
||||
int64_t cur_qo_start = qo_indptr[b];
|
||||
int64_t cur_kv_start = kv_indptr[b];
|
||||
|
||||
float r_Q[D_QK / 32];
|
||||
if (is_valid) {
|
||||
// 显式使用 int64_t 计算全局偏移量，防止超过 2GB 显存边界
|
||||
int64_t q_offset_base = (cur_qo_start + q_idx) * num_qo_heads * D_QK + hq * D_QK;
|
||||
#pragma unroll
|
||||
for (int d = 0; d < D_QK / 32; ++d) {
|
||||
r_Q[d] = __bfloat162float(q[q_offset_base + tx + d * 32]);
|
||||
}
|
||||
}
|
||||
|
||||
float r_O[D_VO / 32];
|
||||
#pragma unroll
|
||||
for (int d = 0; d < D_VO / 32; ++d) r_O[d] = 0.0f;
|
||||
|
||||
float m = -1e20f;
|
||||
float l = 0.0f;
|
||||
|
||||
extern __shared__ __nv_bfloat16 smem[];
|
||||
__nv_bfloat16* shared_K = smem;
|
||||
__nv_bfloat16* shared_V = smem + Bc * D_QK;
|
||||
|
||||
int tid = ty * 32 + tx;
|
||||
int total_threads = Br * 32;
|
||||
|
||||
int64_t block_max_q = min((int64_t)seq_len - 1, (int64_t)(blockIdx.z * Br + Br - 1));
|
||||
|
||||
int total_k_ldg = (Bc * D_QK) / 8;
|
||||
int total_v_ldg = (Bc * D_VO) / 8;
|
||||
|
||||
for (int64_t kv_start = 0; kv_start <= block_max_q; kv_start += Bc) {
|
||||
|
||||
// 协同加载 K
|
||||
for (int i = tid; i < total_k_ldg; i += total_threads) {
|
||||
int token_idx = i / (D_QK / 8);
|
||||
int d_idx = i % (D_QK / 8);
|
||||
int64_t global_kv_idx = kv_start + token_idx;
|
||||
|
||||
if (global_kv_idx < seq_len) {
|
||||
// 同样强制为 int64_t
|
||||
int64_t k_offset = (cur_kv_start + global_kv_idx) * num_kv_heads * D_QK + hkv * D_QK + d_idx * 8;
|
||||
((float4*)shared_K)[i] = *(reinterpret_cast<const float4*>(&k[k_offset]));
|
||||
} else {
|
||||
((float4*)shared_K)[i] = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
|
||||
}
|
||||
}
|
||||
|
||||
// 协同加载 V
|
||||
for (int i = tid; i < total_v_ldg; i += total_threads) {
|
||||
int token_idx = i / (D_VO / 8);
|
||||
int d_idx = i % (D_VO / 8);
|
||||
int64_t global_kv_idx = kv_start + token_idx;
|
||||
|
||||
if (global_kv_idx < seq_len) {
|
||||
int64_t v_offset = (cur_kv_start + global_kv_idx) * num_kv_heads * D_VO + hkv * D_VO + d_idx * 8;
|
||||
((float4*)shared_V)[i] = *(reinterpret_cast<const float4*>(&v[v_offset]));
|
||||
} else {
|
||||
((float4*)shared_V)[i] = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (is_valid) {
|
||||
int64_t kv_end = causal ? min(kv_start + Bc, q_idx + 1) : min(kv_start + Bc, seq_len);
|
||||
|
||||
for (int j = 0; j < kv_end - kv_start; ++j) {
|
||||
float score = 0.0f;
|
||||
#pragma unroll
|
||||
for (int d = 0; d < D_QK / 32; ++d) {
|
||||
score += r_Q[d] * __bfloat162float(shared_K[j * D_QK + tx + d * 32]);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int offset = 16; offset > 0; offset /= 2) {
|
||||
score += __shfl_down_sync(0xffffffff, score, offset);
|
||||
}
|
||||
score = __shfl_sync(0xffffffff, score, 0);
|
||||
score *= scale;
|
||||
|
||||
float m_prev = m;
|
||||
m = max(m, score);
|
||||
float P = __expf(score - m);
|
||||
l = l * __expf(m_prev - m) + P;
|
||||
|
||||
#pragma unroll
|
||||
for (int d = 0; d < D_VO / 32; ++d) {
|
||||
r_O[d] = r_O[d] * __expf(m_prev - m) + P * __bfloat162float(shared_V[j * D_VO + tx + d * 32]);
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (is_valid) {
|
||||
int64_t out_offset_base = (cur_qo_start + q_idx) * num_qo_heads * D_VO + hq * D_VO;
|
||||
#pragma unroll
|
||||
for (int d = 0; d < D_VO / 32; ++d) {
|
||||
float out_val = r_O[d] / l;
|
||||
output[out_offset_base + tx + d * 32] = __float2bfloat16(out_val);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
extern "C" void run_kernel(
|
||||
const __nv_bfloat16* q,
|
||||
const __nv_bfloat16* k,
|
||||
const __nv_bfloat16* v,
|
||||
__nv_bfloat16* output,
|
||||
const int32_t* qo_indptr,
|
||||
const int32_t* kv_indptr,
|
||||
int64_t batch_size,
|
||||
int64_t seq_len,
|
||||
int64_t num_qo_heads,
|
||||
int64_t num_kv_heads,
|
||||
int64_t head_dim_qk,
|
||||
int64_t head_dim_vo,
|
||||
int64_t causal
|
||||
) {
|
||||
float scale = 1.0f / sqrtf((float)head_dim_qk);
|
||||
const int Br = 32;
|
||||
|
||||
dim3 grid(batch_size, num_qo_heads, (seq_len + Br - 1) / Br);
|
||||
dim3 block(32, Br);
|
||||
|
||||
// 根据维度配置绝对安全的共享内存策略
|
||||
if (head_dim_qk == 128 && head_dim_vo == 128) {
|
||||
const int Bc = 64;
|
||||
int smem = Bc * 128 * 2 + Bc * 128 * 2; // 32 KB
|
||||
cudaFuncSetAttribute(ragged_prefill_kernel<128, 128, Br, Bc>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
|
||||
ragged_prefill_kernel<128, 128, Br, Bc><<<grid, block, smem>>>(
|
||||
q, k, v, output, qo_indptr, kv_indptr, seq_len, num_qo_heads, num_kv_heads, scale, causal);
|
||||
} else if (head_dim_qk == 192 && head_dim_vo == 128) {
|
||||
const int Bc = 64;
|
||||
int smem = Bc * 192 * 2 + Bc * 128 * 2; // 40 KB
|
||||
cudaFuncSetAttribute(ragged_prefill_kernel<192, 128, Br, Bc>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
|
||||
ragged_prefill_kernel<192, 128, Br, Bc><<<grid, block, smem>>>(
|
||||
q, k, v, output, qo_indptr, kv_indptr, seq_len, num_qo_heads, num_kv_heads, scale, causal);
|
||||
} else if (head_dim_qk == 256 && head_dim_vo == 256) {
|
||||
// 关键降级：256维度下，将 Bc 从 64 降为 32。
|
||||
// 将所需的共享内存从 64KB 压回到绝对安全的 32KB，防止在部分环境触发 Launch Failure (OOM/Crash)。
|
||||
const int Bc = 32;
|
||||
int smem = Bc * 256 * 2 + Bc * 256 * 2; // 32 KB
|
||||
cudaFuncSetAttribute(ragged_prefill_kernel<256, 256, Br, Bc>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
|
||||
ragged_prefill_kernel<256, 256, Br, Bc><<<grid, block, smem>>>(
|
||||
q, k, v, output, qo_indptr, kv_indptr, seq_len, num_qo_heads, num_kv_heads, scale, causal);
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
#### Step 5：提交代码开始评测
|
||||
|
||||
|
||||
|
||||
## 七、常见问题
|
||||
|
||||
|
|
@ -353,4 +576,8 @@ mv *.csv results/
|
|||
| 优化项 | 改动内容 | Baseline | 优化后 | 提升比例 |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| 例：调整 block\_size | 16 → 32 | xx ms | xx ms | xx% |
|
||||
| 完成优化后，再次运行本模块的 Benchmark 脚本，对比前后性能变化。 | | | | |
|
||||
| 完成优化后，再次运行本模块的 Benchmark 脚本，对比前后性能变化。 | | | | |
|
||||
|
||||
### 6. 使用多语言完成算子优化加速
|
||||
|
||||
可以使用Triton和TileLang语言实现`run_kernel`接口完成算子优化，对比不同的语言对于性能加速的影响
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
@ -0,0 +1,16 @@
|
|||
{
|
||||
"id": 193,
|
||||
"displayId": 20001,
|
||||
"type": "Traditional",
|
||||
"isPublic": false,
|
||||
"locales": [
|
||||
"zh_CN"
|
||||
],
|
||||
"samples": [
|
||||
{
|
||||
"inputData": "1\n",
|
||||
"outputData": ""
|
||||
}
|
||||
],
|
||||
"problemTagIds": []
|
||||
}
|
||||
|
|
@ -0,0 +1,307 @@
|
|||
from __future__ import annotations
|
||||
|
||||
|
||||
HEAD_DIM_CONFIGS = [(128, 128)]
|
||||
BATCH_SIZES = [1, 4, 16]
|
||||
SEQ_LENS = [1024, 4096, 8192, 16384]
|
||||
NUM_QO_HEADS = 32
|
||||
NUM_KV_HEADS = 4
|
||||
CAUSAL = 1
|
||||
|
||||
|
||||
def _build_cases():
|
||||
cases = []
|
||||
for head_dim_qk, head_dim_vo in HEAD_DIM_CONFIGS:
|
||||
for batch_size in BATCH_SIZES:
|
||||
for seq_len in SEQ_LENS:
|
||||
cases.append(
|
||||
(
|
||||
batch_size,
|
||||
seq_len,
|
||||
NUM_QO_HEADS,
|
||||
NUM_KV_HEADS,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
CAUSAL,
|
||||
)
|
||||
)
|
||||
return cases
|
||||
|
||||
|
||||
TESTCASES = _build_cases()
|
||||
|
||||
|
||||
def getNumOfTestcases() -> int:
|
||||
return len(TESTCASES)
|
||||
|
||||
|
||||
try:
|
||||
from pathlib import Path
|
||||
from typing import List, Tuple, Union
|
||||
import sys
|
||||
|
||||
import torch
|
||||
|
||||
KernelArg = Union[torch.Tensor, int, float]
|
||||
CURRENT_CASE = None
|
||||
|
||||
def _ensure_flashinfer_importable():
|
||||
try:
|
||||
import flashinfer # noqa: F401
|
||||
|
||||
return
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
here = Path(__file__).resolve()
|
||||
for parent in here.parents:
|
||||
candidate = parent / "McFlashInfer"
|
||||
if (candidate / "flashinfer").is_dir():
|
||||
sys.path.insert(0, str(candidate))
|
||||
return
|
||||
|
||||
def _get_testcase_index() -> int:
|
||||
try:
|
||||
raw = input().strip()
|
||||
except EOFError:
|
||||
return 0
|
||||
if raw == "":
|
||||
return 0
|
||||
try:
|
||||
testcase_id = int(raw.split()[0])
|
||||
except ValueError:
|
||||
return 0
|
||||
if 1 <= testcase_id <= len(TESTCASES):
|
||||
return testcase_id - 1
|
||||
if 0 <= testcase_id < len(TESTCASES):
|
||||
return testcase_id
|
||||
return 0
|
||||
|
||||
def _compute_reps(batch_size: int, seq_len: int, head_dim: int, base_reps: int = 100) -> int:
|
||||
workload = batch_size * seq_len * head_dim
|
||||
if workload < 1e5:
|
||||
return base_reps
|
||||
if workload < 1e6:
|
||||
return base_reps // 2
|
||||
if workload < 1e7:
|
||||
return base_reps // 4
|
||||
if workload < 1e8:
|
||||
return base_reps // 8
|
||||
if workload < 1e9:
|
||||
return base_reps // 16
|
||||
return base_reps // 32
|
||||
|
||||
def getTestCaseSize() -> Tuple[List[Tuple[int, ...]], Tuple[int, int]]:
|
||||
testcase_id = _get_testcase_index()
|
||||
global CURRENT_CASE
|
||||
batch_size, seq_len, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, causal = TESTCASES[testcase_id]
|
||||
CURRENT_CASE = (
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
20260610 + testcase_id,
|
||||
)
|
||||
qo_len = batch_size * seq_len
|
||||
kv_len = batch_size * seq_len
|
||||
warmup = 3
|
||||
iters = max(1, _compute_reps(batch_size, seq_len, head_dim_qk + head_dim_vo))
|
||||
return [
|
||||
(qo_len, num_qo_heads, head_dim_qk),
|
||||
(kv_len, num_kv_heads, head_dim_qk),
|
||||
(kv_len, num_kv_heads, head_dim_vo),
|
||||
(qo_len, num_qo_heads, head_dim_vo),
|
||||
(batch_size + 1,),
|
||||
(batch_size + 1,),
|
||||
(), (), (), (), (), (), (),
|
||||
], (warmup, iters)
|
||||
|
||||
def genTestCase(testcase_sizes, device: str = "cuda") -> List[KernelArg]:
|
||||
del testcase_sizes
|
||||
batch_size, seq_len, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, causal, seed = CURRENT_CASE
|
||||
gen = torch.Generator(device=device)
|
||||
gen.manual_seed(seed)
|
||||
dtype = torch.bfloat16
|
||||
qo_len = batch_size * seq_len
|
||||
q = torch.rand(
|
||||
qo_len,
|
||||
num_qo_heads,
|
||||
head_dim_qk,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
generator=gen,
|
||||
).contiguous()
|
||||
kv_len = batch_size * seq_len
|
||||
k = torch.rand(
|
||||
kv_len,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
generator=gen,
|
||||
).contiguous()
|
||||
v = torch.rand(
|
||||
kv_len,
|
||||
num_kv_heads,
|
||||
head_dim_vo,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
generator=gen,
|
||||
).contiguous()
|
||||
output = torch.empty(
|
||||
qo_len,
|
||||
num_qo_heads,
|
||||
head_dim_vo,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
qo_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device=device) * seq_len
|
||||
kv_indptr = qo_indptr.clone()
|
||||
return [
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
output,
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
]
|
||||
|
||||
def baseline(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
output,
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
):
|
||||
_ensure_flashinfer_importable()
|
||||
import flashinfer
|
||||
|
||||
workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device)
|
||||
wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper(
|
||||
workspace_buffer,
|
||||
kv_layout="NHD",
|
||||
backend="auto",
|
||||
)
|
||||
wrapper.plan(
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
int(num_qo_heads),
|
||||
int(num_kv_heads),
|
||||
int(head_dim_qk),
|
||||
int(head_dim_vo),
|
||||
causal=bool(causal),
|
||||
q_data_type=torch.bfloat16,
|
||||
kv_data_type=torch.bfloat16,
|
||||
)
|
||||
wrapper.run(q, k, v, out=output)
|
||||
return [
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
output,
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
]
|
||||
|
||||
def check(
|
||||
testcase_sizes,
|
||||
original_input_tensors,
|
||||
target_kernel_input_tensors,
|
||||
baseline_input_tensors,
|
||||
rtol=1e-2,
|
||||
atol=1e-2,
|
||||
) -> bool:
|
||||
del testcase_sizes, original_input_tensors
|
||||
output_t = target_kernel_input_tensors[3]
|
||||
output_ref = baseline_input_tensors[3]
|
||||
if output_t.shape != output_ref.shape:
|
||||
print(f"[FAIL] shape mismatch: target {output_t.shape}, ref {output_ref.shape}", file=sys.stderr)
|
||||
return False
|
||||
if output_t.dtype != output_ref.dtype:
|
||||
print(f"[FAIL] dtype mismatch: target {output_t.dtype}, ref {output_ref.dtype}", file=sys.stderr)
|
||||
return False
|
||||
if not torch.allclose(output_t.float(), output_ref.float(), rtol=rtol, atol=atol):
|
||||
diff = (output_t.float() - output_ref.float()).abs()
|
||||
print(
|
||||
f"[FAIL] allclose failed: max_abs_diff={float(diff.max().item()):.6f}, "
|
||||
f"mean_abs_diff={float(diff.mean().item()):.6f} (rtol={rtol}, atol={atol})",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return False
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
INPUT_CLASS = [
|
||||
"INPUT",
|
||||
"INPUT",
|
||||
"INPUT",
|
||||
"OUTPUT",
|
||||
"INPUT",
|
||||
"INPUT",
|
||||
"INPUT",
|
||||
"INPUT",
|
||||
"INPUT",
|
||||
"INPUT",
|
||||
"INPUT",
|
||||
"INPUT",
|
||||
"INPUT",
|
||||
]
|
||||
|
||||
|
||||
def getWorkload(testcase_sizes) -> dict:
|
||||
raw_sizes = testcase_sizes[0] if isinstance(testcase_sizes, tuple) and len(testcase_sizes) == 2 else testcase_sizes
|
||||
q_shape, k_shape, v_shape, output_shape, qo_indptr_shape, kv_indptr_shape = raw_sizes[:6]
|
||||
qo_len, num_qo_heads, head_dim_qk = q_shape
|
||||
kv_len, num_kv_heads, k_dim = k_shape
|
||||
v_len, v_heads, head_dim_vo = v_shape
|
||||
assert k_dim == head_dim_qk
|
||||
assert v_len == kv_len
|
||||
assert v_heads == num_kv_heads
|
||||
assert qo_len == kv_len
|
||||
assert output_shape == (qo_len, num_qo_heads, head_dim_vo)
|
||||
assert qo_indptr_shape == kv_indptr_shape
|
||||
batch_size = qo_indptr_shape[0] - 1
|
||||
seq_len = kv_len // batch_size
|
||||
flops = batch_size * seq_len * seq_len * num_qo_heads * (head_dim_qk + head_dim_vo)
|
||||
memory_bytes = (
|
||||
qo_len * num_qo_heads * head_dim_qk * 2
|
||||
+ kv_len * num_kv_heads * head_dim_qk * 2
|
||||
+ kv_len * num_kv_heads * head_dim_vo * 2
|
||||
+ qo_len * num_qo_heads * head_dim_vo * 2
|
||||
+ (batch_size + 1) * 4 * 2
|
||||
)
|
||||
return {
|
||||
"flops": flops,
|
||||
"memory_bytes": memory_bytes,
|
||||
"dtype": "bf16",
|
||||
}
|
||||
|
||||
|
||||
DESIGNED_VRAM_SIZE = 48
|
||||
|
|
@ -0,0 +1,23 @@
|
|||
---
|
||||
sectionTitle: "题目描述"
|
||||
type: "Text"
|
||||
---
|
||||
你需要实现 FlashInfer ragged KV cache prefill 的CUDA C++前向算子。
|
||||
|
||||
本题输入采用 FlashInfer `BatchPrefillWithRaggedKVCacheWrapper` 的 ragged `NHD` 布局。每个 batch 中有 `seq_len` 个 query token，KV cache 中也有 `seq_len` 个 token：
|
||||
|
||||
其中 query heads 采用 GQA 布局：`num_qo_heads` 个 query/output heads 共享 `num_kv_heads` 个 KV heads，`G = num_qo_heads / num_kv_heads`。
|
||||
|
||||
评测程序会调用你提交代码中的 `run_kernel` 函数。你需要根据 `qo_indptr` 和 `kv_indptr` 读取 ragged Q/K/V，并将结果写入 `output`。
|
||||
|
||||
baseline 使用 FlashInfer ragged prefill 的 Python API：
|
||||
|
||||
```python
|
||||
wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper(workspace, kv_layout="NHD", backend="auto")
|
||||
wrapper.plan(qo_indptr, kv_indptr, num_qo_heads, num_kv_heads,
|
||||
head_dim_qk, head_dim_vo, causal=True,
|
||||
q_data_type=torch.bfloat16, kv_data_type=torch.bfloat16)
|
||||
wrapper.run(q, k, v, out=output)
|
||||
```
|
||||
|
||||
如何提交代码详见[评测指南](/d/2)。
|
||||
|
|
@ -0,0 +1,41 @@
|
|||
---
|
||||
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,
|
||||
const __nv_bfloat16* v,
|
||||
__nv_bfloat16* output,
|
||||
const int32_t* qo_indptr,
|
||||
const int32_t* kv_indptr,
|
||||
int64_t batch_size,
|
||||
int64_t seq_len,
|
||||
int64_t num_qo_heads,
|
||||
int64_t num_kv_heads,
|
||||
int64_t head_dim_qk,
|
||||
int64_t head_dim_vo,
|
||||
int64_t causal
|
||||
);
|
||||
```
|
||||
|
||||
### 参数说明
|
||||
|
||||
* `q`：query tensor，shape `(batch_size * seq_len, num_qo_heads, head_dim_qk)`，连续 `bf16`
|
||||
* `k`：key tensor，shape `(batch_size * seq_len, num_kv_heads, head_dim_qk)`，连续 `bf16`
|
||||
* `v`：value tensor，shape `(batch_size * seq_len, num_kv_heads, head_dim_vo)`，连续 `bf16`
|
||||
* `output`：输出缓冲区，shape `(batch_size * seq_len, num_qo_heads, head_dim_vo)`，连续 `bf16`
|
||||
* `qo_indptr`：query/output ragged indptr，shape `(batch_size + 1)`，连续 `int32`
|
||||
* `kv_indptr`：KV ragged indptr，shape `(batch_size + 1)`，连续 `int32`
|
||||
* `causal`：是否启用 causal mask，评测中固定为 `1`
|
||||
|
||||
本题测试中 `qo_indptr[b + 1] - qo_indptr[b] == seq_len`，`kv_indptr[b + 1] - kv_indptr[b] == seq_len`。
|
||||
|
||||
`run_kernel` 内部需要自行计算合适的 launch 配置并启动 CUDA kernel。为保证计时准确，不建议在 `run_kernel` 内部做 `cudaDeviceSynchronize()` 或显式同步。
|
||||
|
|
@ -0,0 +1,26 @@
|
|||
---
|
||||
sectionTitle: "输入格式"
|
||||
type: "Text"
|
||||
---
|
||||
## 数据范围
|
||||
|
||||
- 数据类型：`q/kv_data/output` 均为 `bfloat16`
|
||||
- KV layout：`NHD`
|
||||
- `num_qo_heads = 32`
|
||||
- `num_kv_heads = 4`
|
||||
- `page_block_size = 16`
|
||||
- `causal = 0`
|
||||
- `head_dim` 固定为 `128`
|
||||
- `batch_size` 取值为 `1, 4, 16`
|
||||
- `seq_len` 取值为 `1024, 4096, 8192, 16384`
|
||||
|
||||
测试点顺序与 `McFlashInfer/benchmarks/bench_batch_prefill_paged.py` 中的 cases 一致，即：
|
||||
|
||||
```python
|
||||
for head_dim in [128]:
|
||||
for batch_size in [1, 4, 16]:
|
||||
for seq_len in [1024, 4096, 8192, 16384]:
|
||||
...
|
||||
```
|
||||
|
||||
输出参与 `torch.allclose` 校验，容差为 `rtol=1e-2, atol=1e-2`。
|
||||
|
|
@ -0,0 +1,28 @@
|
|||
pytorch参考实现：
|
||||
|
||||
```python
|
||||
def baseline(q, kv_data, output, qo_indptr, kv_indptr, kv_indices, last_page_len,
|
||||
batch_size, seq_len, num_qo_heads, num_kv_heads,
|
||||
head_dim, page_block_size, causal):
|
||||
workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device)
|
||||
wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(
|
||||
workspace_buffer,
|
||||
kv_layout="NHD",
|
||||
backend="auto",
|
||||
)
|
||||
wrapper.plan(
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
last_page_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
page_block_size,
|
||||
causal=bool(causal),
|
||||
q_data_type=torch.bfloat16,
|
||||
kv_data_type=torch.bfloat16,
|
||||
)
|
||||
wrapper.run(q, kv_data, out=output)
|
||||
```
|
||||
|
||||
|
|
@ -0,0 +1 @@
|
|||
FlashInfer Ragged Prefill
|
||||
|
|
@ -0,0 +1,16 @@
|
|||
{
|
||||
"id": 194,
|
||||
"displayId": 20002,
|
||||
"type": "Traditional",
|
||||
"isPublic": false,
|
||||
"locales": [
|
||||
"zh_CN"
|
||||
],
|
||||
"samples": [
|
||||
{
|
||||
"inputData": "1\n",
|
||||
"outputData": ""
|
||||
}
|
||||
],
|
||||
"problemTagIds": []
|
||||
}
|
||||
|
|
@ -0,0 +1,273 @@
|
|||
from __future__ import annotations
|
||||
|
||||
|
||||
HEAD_DIMS = [128]
|
||||
BATCH_SIZES = [1, 4, 16]
|
||||
SEQ_LENS = [1024, 4096, 8192, 16384]
|
||||
NUM_QO_HEADS = 32
|
||||
PAGE_BLOCK_SIZE = 16
|
||||
CAUSAL = 0
|
||||
|
||||
|
||||
def _build_cases():
|
||||
cases = []
|
||||
for head_dim in HEAD_DIMS:
|
||||
for batch_size in BATCH_SIZES:
|
||||
for seq_len in SEQ_LENS:
|
||||
num_kv_heads = 8 if head_dim == 64 else 4
|
||||
cases.append(
|
||||
(
|
||||
batch_size,
|
||||
seq_len,
|
||||
NUM_QO_HEADS,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
PAGE_BLOCK_SIZE,
|
||||
CAUSAL,
|
||||
)
|
||||
)
|
||||
return cases
|
||||
|
||||
|
||||
TESTCASES = _build_cases()
|
||||
|
||||
|
||||
def getNumOfTestcases() -> int:
|
||||
return len(TESTCASES)
|
||||
|
||||
|
||||
try:
|
||||
from pathlib import Path
|
||||
from typing import List, Tuple, Union
|
||||
import sys
|
||||
|
||||
import torch
|
||||
|
||||
KernelArg = Union[torch.Tensor, int, float]
|
||||
CURRENT_CASE = None
|
||||
|
||||
def _ensure_flashinfer_importable():
|
||||
try:
|
||||
import flashinfer # noqa: F401
|
||||
|
||||
return
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
here = Path(__file__).resolve()
|
||||
for parent in here.parents:
|
||||
candidate = parent / "McFlashInfer"
|
||||
if (candidate / "flashinfer").is_dir():
|
||||
sys.path.insert(0, str(candidate))
|
||||
return
|
||||
|
||||
def _get_testcase_index() -> int:
|
||||
try:
|
||||
raw = input().strip()
|
||||
except EOFError:
|
||||
return 0
|
||||
if raw == "":
|
||||
return 0
|
||||
try:
|
||||
testcase_id = int(raw.split()[0])
|
||||
except ValueError:
|
||||
return 0
|
||||
if 1 <= testcase_id <= len(TESTCASES):
|
||||
return testcase_id - 1
|
||||
if 0 <= testcase_id < len(TESTCASES):
|
||||
return testcase_id
|
||||
return 0
|
||||
|
||||
def _compute_reps(batch_size: int, seq_len: int, head_dim: int, base_reps: int = 100) -> int:
|
||||
workload = batch_size * seq_len * head_dim
|
||||
if workload < 1e5:
|
||||
return base_reps
|
||||
if workload < 1e6:
|
||||
return base_reps // 2
|
||||
if workload < 1e7:
|
||||
return base_reps // 4
|
||||
if workload < 1e8:
|
||||
return base_reps // 8
|
||||
if workload < 1e9:
|
||||
return base_reps // 16
|
||||
return base_reps // 32
|
||||
|
||||
def _setup_paged_kv_indptr(batch_size: int, seq_len: int, page_block_size: int, device: str):
|
||||
seq_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, device=device)
|
||||
seq_lens_blocks = torch.div(seq_lens + page_block_size - 1, page_block_size, rounding_mode="floor")
|
||||
kv_indptr = torch.empty((batch_size + 1,), dtype=torch.int32, device=device)
|
||||
kv_indptr[0] = 0
|
||||
kv_indptr[1:] = torch.cumsum(seq_lens_blocks, dim=0)
|
||||
num_blocks = int(kv_indptr[-1].item())
|
||||
last_page_len = (seq_lens - 1) % page_block_size + 1
|
||||
return kv_indptr, last_page_len, num_blocks
|
||||
|
||||
def getTestCaseSize() -> Tuple[List[Tuple[int, ...]], Tuple[int, int]]:
|
||||
testcase_id = _get_testcase_index()
|
||||
global CURRENT_CASE
|
||||
batch_size, seq_len, num_qo_heads, num_kv_heads, head_dim, page_block_size, causal = TESTCASES[testcase_id]
|
||||
num_blocks = batch_size * ((seq_len + page_block_size - 1) // page_block_size)
|
||||
qo_len = batch_size * seq_len
|
||||
CURRENT_CASE = (
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
page_block_size,
|
||||
causal,
|
||||
num_blocks,
|
||||
20260620 + testcase_id,
|
||||
)
|
||||
warmup = 3
|
||||
iters = max(1, _compute_reps(batch_size, seq_len, head_dim))
|
||||
return [
|
||||
(qo_len, num_qo_heads, head_dim),
|
||||
(num_blocks, 2, page_block_size, num_kv_heads, head_dim),
|
||||
(qo_len, num_qo_heads, head_dim),
|
||||
(batch_size + 1,),
|
||||
(batch_size + 1,),
|
||||
(num_blocks,),
|
||||
(batch_size,),
|
||||
(), (), (), (), (), (), (),
|
||||
], (warmup, iters)
|
||||
|
||||
def genTestCase(testcase_sizes, device: str = "cuda") -> List[KernelArg]:
|
||||
del testcase_sizes
|
||||
(
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
page_block_size,
|
||||
causal,
|
||||
num_blocks,
|
||||
seed,
|
||||
) = CURRENT_CASE
|
||||
gen = torch.Generator(device=device)
|
||||
gen.manual_seed(seed)
|
||||
dtype = torch.bfloat16
|
||||
qo_len = batch_size * seq_len
|
||||
q = torch.rand(qo_len, num_qo_heads, head_dim, dtype=dtype, device=device, generator=gen).contiguous()
|
||||
kv_data = torch.randn(
|
||||
num_blocks,
|
||||
2,
|
||||
page_block_size,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
generator=gen,
|
||||
).contiguous()
|
||||
output = torch.empty(qo_len, num_qo_heads, head_dim, dtype=dtype, device=device)
|
||||
qo_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device=device) * seq_len
|
||||
kv_indptr, last_page_len, num_blocks_check = _setup_paged_kv_indptr(
|
||||
batch_size,
|
||||
seq_len,
|
||||
page_block_size,
|
||||
device,
|
||||
)
|
||||
assert num_blocks_check == num_blocks
|
||||
kv_indices = torch.arange(num_blocks, dtype=torch.int32, device=device)
|
||||
return [
|
||||
q,
|
||||
kv_data,
|
||||
output,
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
last_page_len,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
page_block_size,
|
||||
causal,
|
||||
]
|
||||
|
||||
def baseline(
|
||||
q,
|
||||
kv_data,
|
||||
output,
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
last_page_len,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
page_block_size,
|
||||
causal,
|
||||
):
|
||||
_ensure_flashinfer_importable()
|
||||
import flashinfer
|
||||
|
||||
workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device)
|
||||
wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(
|
||||
workspace_buffer,
|
||||
kv_layout="NHD",
|
||||
backend="auto",
|
||||
)
|
||||
wrapper.plan(
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
last_page_len,
|
||||
int(num_qo_heads),
|
||||
int(num_kv_heads),
|
||||
int(head_dim),
|
||||
int(page_block_size),
|
||||
causal=bool(causal),
|
||||
q_data_type=torch.bfloat16,
|
||||
kv_data_type=torch.bfloat16,
|
||||
)
|
||||
wrapper.run(q, kv_data, out=output)
|
||||
return [
|
||||
q,
|
||||
kv_data,
|
||||
output,
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
last_page_len,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
page_block_size,
|
||||
causal,
|
||||
]
|
||||
|
||||
def check(
|
||||
testcase_sizes,
|
||||
original_input_tensors,
|
||||
target_kernel_input_tensors,
|
||||
baseline_input_tensors,
|
||||
rtol=1e-2,
|
||||
atol=1e-2,
|
||||
) -> bool:
|
||||
del testcase_sizes, original_input_tensors
|
||||
output_t = target_kernel_input_tensors[2]
|
||||
output_ref = baseline_input_tensors[2]
|
||||
if output_t.shape != output_ref.shape:
|
||||
print(f"[FAIL] shape mismatch: target {output_t.shape}, ref {output_ref.shape}", file=sys.stderr)
|
||||
return False
|
||||
if output_t.dtype != output_ref.dtype:
|
||||
print(f"[FAIL] dtype mismatch: target {output_t.dtype}, ref {output_ref.dtype}", file=sys.stderr)
|
||||
return False
|
||||
if not torch.allclose(output_t.float(), output_ref.float(), rtol=rtol, atol=atol):
|
||||
diff = (output_t.float() - output_ref.float()).abs()
|
||||
print(
|
||||
f"[FAIL] allclose failed: max_abs_diff={float(diff.max().item()):.6f}, "
|
||||
f"mean_abs_diff={float(diff.mean().item()):.6f} (rtol={rtol}, atol={atol})",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return False
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
|
|
@ -0,0 +1,22 @@
|
|||
---
|
||||
sectionTitle: "题目描述"
|
||||
type: "Text"
|
||||
---
|
||||
你需要实现 FlashInfer paged KV cache prefill 的 CUDA C++ 前向算子。
|
||||
|
||||
本题输入采用 FlashInfer `BatchPrefillWithPagedKVCacheWrapper` 的 paged `NHD` 布局。每个 batch 中有 `seq_len` 个 query token，KV cache 也有 `seq_len` 个 token，并按 page 存储。
|
||||
|
||||
评测程序会调用你提交代码中的 `run_kernel` 函数。你需要根据 `qo_indptr`、`kv_indptr`、`kv_indices` 和 `last_page_len` 读取 paged KV cache，并将结果写入 `output`。
|
||||
|
||||
baseline 使用 FlashInfer paged prefill 的 Python API：
|
||||
|
||||
```python
|
||||
wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(workspace, kv_layout="NHD", backend="auto")
|
||||
wrapper.plan(qo_indptr, kv_indptr, kv_indices, last_page_len,
|
||||
num_qo_heads, num_kv_heads, head_dim, page_block_size,
|
||||
causal=bool(causal),
|
||||
q_data_type=torch.bfloat16, kv_data_type=torch.bfloat16)
|
||||
wrapper.run(q, kv_data, out=output)
|
||||
```
|
||||
|
||||
如何提交代码详见[评测指南](/d/2)。
|
||||
|
|
@ -0,0 +1,42 @@
|
|||
---
|
||||
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* kv_data,
|
||||
__nv_bfloat16* output,
|
||||
const int32_t* qo_indptr,
|
||||
const int32_t* kv_indptr,
|
||||
const int32_t* kv_indices,
|
||||
const int32_t* last_page_len,
|
||||
int64_t batch_size,
|
||||
int64_t seq_len,
|
||||
int64_t num_qo_heads,
|
||||
int64_t num_kv_heads,
|
||||
int64_t head_dim,
|
||||
int64_t page_block_size,
|
||||
int64_t causal
|
||||
);
|
||||
```
|
||||
|
||||
### 参数说明
|
||||
|
||||
* `q`：query tensor，shape `(batch_size * seq_len, num_qo_heads, head_dim)`，连续 `bf16`
|
||||
* `kv_data`：paged KV cache，shape `(num_blocks, 2, page_block_size, num_kv_heads, head_dim)`，连续 `bf16`，其中 `kv_data[:, 0]` 为 key，`kv_data[:, 1]` 为 value
|
||||
* `output`：输出缓冲区，shape `(batch_size * seq_len, num_qo_heads, head_dim)`，连续 `bf16`
|
||||
* `qo_indptr`：query/output indptr，shape `(batch_size + 1)`，连续 `int32`
|
||||
* `kv_indptr`：paged KV indptr，shape `(batch_size + 1)`，连续 `int32`
|
||||
* `kv_indices`：page index，shape `(num_blocks)`，连续 `int32`
|
||||
* `last_page_len`：每个 batch 最后一个 page 的有效 token 数，shape `(batch_size)`，连续 `int32`
|
||||
* `page_block_size`：page size，评测中固定为 `16`
|
||||
* `causal`：是否启用 causal mask，本题按 benchmark case 固定为 `0`
|
||||
|
||||
`run_kernel` 内部需要自行计算合适的 launch 配置并启动 CUDA kernel。为保证计时准确，不建议在 `run_kernel` 内部做 `cudaDeviceSynchronize()` 或显式同步。
|
||||
|
|
@ -0,0 +1,26 @@
|
|||
---
|
||||
sectionTitle: "输入格式"
|
||||
type: "Text"
|
||||
---
|
||||
## 数据范围
|
||||
|
||||
- 数据类型：`q/kv_data/output` 均为 `bfloat16`
|
||||
- KV layout：`NHD`
|
||||
- `num_qo_heads = 32`
|
||||
- `num_kv_heads = 4`
|
||||
- `page_block_size = 16`
|
||||
- `causal = 0`
|
||||
- `head_dim` 固定为 `128`
|
||||
- `batch_size` 取值为 `1, 4, 16`
|
||||
- `seq_len` 取值为 `1024, 4096, 8192, 16384`
|
||||
|
||||
测试点顺序与 `McFlashInfer/benchmarks/bench_batch_prefill_paged.py` 中的 cases 一致，即：
|
||||
|
||||
```python
|
||||
for head_dim in [128]:
|
||||
for batch_size in [1, 4, 16]:
|
||||
for seq_len in [1024, 4096, 8192, 16384]:
|
||||
...
|
||||
```
|
||||
|
||||
输出参与 `torch.allclose` 校验，容差为 `rtol=1e-2, atol=1e-2`。
|
||||
|
|
@ -0,0 +1,31 @@
|
|||
---
|
||||
sectionTitle: "输出格式"
|
||||
type: "Text"
|
||||
---
|
||||
PyTorch 参考实现
|
||||
|
||||
```python
|
||||
def baseline(q, kv_data, output, qo_indptr, kv_indptr, kv_indices, last_page_len,
|
||||
batch_size, seq_len, num_qo_heads, num_kv_heads,
|
||||
head_dim, page_block_size, causal):
|
||||
workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device)
|
||||
wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(
|
||||
workspace_buffer,
|
||||
kv_layout="NHD",
|
||||
backend="auto",
|
||||
)
|
||||
wrapper.plan(
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
last_page_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
page_block_size,
|
||||
causal=bool(causal),
|
||||
q_data_type=torch.bfloat16,
|
||||
kv_data_type=torch.bfloat16,
|
||||
)
|
||||
wrapper.run(q, kv_data, out=output)
|
||||
```
|
||||
|
|
@ -0,0 +1 @@
|
|||
FlashInfer Paged Prefill
|
||||
|
|
@ -0,0 +1,16 @@
|
|||
{
|
||||
"id": 195,
|
||||
"displayId": 20003,
|
||||
"type": "Traditional",
|
||||
"isPublic": false,
|
||||
"locales": [
|
||||
"zh_CN"
|
||||
],
|
||||
"samples": [
|
||||
{
|
||||
"inputData": "1\n",
|
||||
"outputData": ""
|
||||
}
|
||||
],
|
||||
"problemTagIds": []
|
||||
}
|
||||
|
|
@ -0,0 +1,245 @@
|
|||
from __future__ import annotations
|
||||
|
||||
|
||||
HEAD_DIM_CKV = 512
|
||||
HEAD_DIM_KPE = 64
|
||||
NUM_HEADS_LIST = [64, 128]
|
||||
BATCH_SIZES = [1, 4, 16]
|
||||
SEQ_LENS = [1024, 4096, 8192, 16384]
|
||||
PAGE_SIZE = 1
|
||||
CAUSAL = 0
|
||||
|
||||
|
||||
def _build_cases():
|
||||
cases = []
|
||||
for num_heads in NUM_HEADS_LIST:
|
||||
for batch_size in BATCH_SIZES:
|
||||
for seq_len in SEQ_LENS:
|
||||
cases.append(
|
||||
(
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_heads,
|
||||
HEAD_DIM_CKV,
|
||||
HEAD_DIM_KPE,
|
||||
PAGE_SIZE,
|
||||
CAUSAL,
|
||||
)
|
||||
)
|
||||
return cases
|
||||
|
||||
|
||||
TESTCASES = _build_cases()
|
||||
|
||||
|
||||
def getNumOfTestcases() -> int:
|
||||
return len(TESTCASES)
|
||||
|
||||
|
||||
try:
|
||||
from pathlib import Path
|
||||
from typing import List, Tuple, Union
|
||||
import sys
|
||||
|
||||
import torch
|
||||
|
||||
KernelArg = Union[torch.Tensor, int, float]
|
||||
CURRENT_CASE = None
|
||||
|
||||
def _ensure_flashinfer_importable():
|
||||
try:
|
||||
import flashinfer # noqa: F401
|
||||
|
||||
return
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
here = Path(__file__).resolve()
|
||||
for parent in here.parents:
|
||||
candidate = parent / "McFlashInfer"
|
||||
if (candidate / "flashinfer").is_dir():
|
||||
sys.path.insert(0, str(candidate))
|
||||
return
|
||||
|
||||
def _get_testcase_index() -> int:
|
||||
try:
|
||||
raw = input().strip()
|
||||
except EOFError:
|
||||
return 0
|
||||
if raw == "":
|
||||
return 0
|
||||
try:
|
||||
testcase_id = int(raw.split()[0])
|
||||
except ValueError:
|
||||
return 0
|
||||
if 1 <= testcase_id <= len(TESTCASES):
|
||||
return testcase_id - 1
|
||||
if 0 <= testcase_id < len(TESTCASES):
|
||||
return testcase_id
|
||||
return 0
|
||||
|
||||
def _compute_reps(batch_size: int, seq_len: int, head_dim: int, base_reps: int = 100) -> int:
|
||||
workload = batch_size * seq_len * head_dim
|
||||
if workload < 1e5:
|
||||
return base_reps
|
||||
if workload < 1e6:
|
||||
return base_reps // 2
|
||||
if workload < 1e7:
|
||||
return base_reps // 4
|
||||
if workload < 1e8:
|
||||
return base_reps // 8
|
||||
if workload < 1e9:
|
||||
return base_reps // 16
|
||||
return base_reps // 32
|
||||
|
||||
def getTestCaseSize() -> Tuple[List[Tuple[int, ...]], Tuple[int, int]]:
|
||||
testcase_id = _get_testcase_index()
|
||||
global CURRENT_CASE
|
||||
batch_size, seq_len, num_heads, head_dim_ckv, head_dim_kpe, page_size, causal = TESTCASES[testcase_id]
|
||||
CURRENT_CASE = (
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim_ckv,
|
||||
head_dim_kpe,
|
||||
page_size,
|
||||
causal,
|
||||
20260630 + testcase_id,
|
||||
)
|
||||
warmup = 3
|
||||
iters = max(1, _compute_reps(batch_size, seq_len, head_dim_ckv + head_dim_kpe))
|
||||
kv_len = batch_size * seq_len
|
||||
return [
|
||||
(batch_size, num_heads, head_dim_ckv),
|
||||
(batch_size, num_heads, head_dim_kpe),
|
||||
(kv_len, 1, head_dim_ckv),
|
||||
(kv_len, 1, head_dim_kpe),
|
||||
(batch_size, num_heads, head_dim_ckv),
|
||||
(batch_size + 1,),
|
||||
(batch_size + 1,),
|
||||
(kv_len,),
|
||||
(batch_size,),
|
||||
(), (), (), (), (), (), (),
|
||||
], (warmup, iters)
|
||||
|
||||
def genTestCase(testcase_sizes, device: str = "cuda") -> List[KernelArg]:
|
||||
del testcase_sizes
|
||||
batch_size, seq_len, num_heads, head_dim_ckv, head_dim_kpe, page_size, causal, seed = CURRENT_CASE
|
||||
gen = torch.Generator(device=device)
|
||||
gen.manual_seed(seed)
|
||||
dtype = torch.bfloat16
|
||||
kv_len = batch_size * seq_len
|
||||
q_nope = torch.randn(batch_size, num_heads, head_dim_ckv, dtype=dtype, device=device, generator=gen).contiguous()
|
||||
q_pe = torch.zeros(batch_size, num_heads, head_dim_kpe, dtype=dtype, device=device)
|
||||
ckv = torch.randn(kv_len, 1, head_dim_ckv, dtype=dtype, device=device, generator=gen).contiguous()
|
||||
kpe = torch.zeros(kv_len, 1, head_dim_kpe, dtype=dtype, device=device)
|
||||
output = torch.empty(batch_size, num_heads, head_dim_ckv, dtype=dtype, device=device)
|
||||
q_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device=device)
|
||||
kv_indptr = q_indptr * seq_len
|
||||
kv_indices = torch.arange(0, kv_len, dtype=torch.int32, device=device)
|
||||
kv_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, device=device)
|
||||
return [
|
||||
q_nope,
|
||||
q_pe,
|
||||
ckv,
|
||||
kpe,
|
||||
output,
|
||||
q_indptr,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
kv_lens,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim_ckv,
|
||||
head_dim_kpe,
|
||||
page_size,
|
||||
causal,
|
||||
]
|
||||
|
||||
def baseline(
|
||||
q_nope,
|
||||
q_pe,
|
||||
ckv,
|
||||
kpe,
|
||||
output,
|
||||
q_indptr,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
kv_lens,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim_ckv,
|
||||
head_dim_kpe,
|
||||
page_size,
|
||||
causal,
|
||||
):
|
||||
_ensure_flashinfer_importable()
|
||||
import flashinfer
|
||||
|
||||
sm_scale = 1.0 / ((int(head_dim_ckv) + int(head_dim_kpe)) ** 0.5)
|
||||
workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q_nope.device)
|
||||
wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(workspace_buffer, backend="auto")
|
||||
wrapper.plan(
|
||||
q_indptr,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
kv_lens,
|
||||
int(num_heads),
|
||||
int(head_dim_ckv),
|
||||
int(head_dim_kpe),
|
||||
int(page_size),
|
||||
bool(causal),
|
||||
sm_scale,
|
||||
q_nope.dtype,
|
||||
ckv.dtype,
|
||||
)
|
||||
wrapper.run(q_nope, q_pe, ckv, kpe, out=output, return_lse=False)
|
||||
return [
|
||||
q_nope,
|
||||
q_pe,
|
||||
ckv,
|
||||
kpe,
|
||||
output,
|
||||
q_indptr,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
kv_lens,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim_ckv,
|
||||
head_dim_kpe,
|
||||
page_size,
|
||||
causal,
|
||||
]
|
||||
|
||||
def check(
|
||||
testcase_sizes,
|
||||
original_input_tensors,
|
||||
target_kernel_input_tensors,
|
||||
baseline_input_tensors,
|
||||
rtol=1e-2,
|
||||
atol=1e-2,
|
||||
) -> bool:
|
||||
del testcase_sizes, original_input_tensors
|
||||
output_t = target_kernel_input_tensors[4]
|
||||
output_ref = baseline_input_tensors[4]
|
||||
if output_t.shape != output_ref.shape:
|
||||
print(f"[FAIL] shape mismatch: target {output_t.shape}, ref {output_ref.shape}", file=sys.stderr)
|
||||
return False
|
||||
if output_t.dtype != output_ref.dtype:
|
||||
print(f"[FAIL] dtype mismatch: target {output_t.dtype}, ref {output_ref.dtype}", file=sys.stderr)
|
||||
return False
|
||||
if not torch.allclose(output_t.float(), output_ref.float(), rtol=rtol, atol=atol):
|
||||
diff = (output_t.float() - output_ref.float()).abs()
|
||||
print(
|
||||
f"[FAIL] allclose failed: max_abs_diff={float(diff.max().item()):.6f}, "
|
||||
f"mean_abs_diff={float(diff.mean().item()):.6f} (rtol={rtol}, atol={atol})",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return False
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
|
|
@ -0,0 +1,22 @@
|
|||
---
|
||||
sectionTitle: "题目描述"
|
||||
type: "Text"
|
||||
---
|
||||
你需要实现 FlashInfer MLA paged attention 的 CUDA C++ 前向算子。
|
||||
|
||||
本题使用 `BatchMLAPagedAttentionWrapper` 的 DeepSeek MLA decode 配置：`q_nope` 表示不含 RoPE 的 query，`q_pe` 表示 RoPE 部分的 query，`ckv` 与 `kpe` 分别表示 compressed KV cache 与 RoPE KV cache。MLA page size 固定为 `1`。
|
||||
|
||||
评测程序会调用你提交代码中的 `run_kernel` 函数。你需要根据 `q_indptr`、`kv_indptr`、`kv_indices` 和 `kv_lens` 读取 cache，并将结果写入 `output`。
|
||||
|
||||
baseline 使用 FlashInfer MLA 的 Python API：
|
||||
|
||||
```python
|
||||
wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(workspace, backend="auto")
|
||||
wrapper.plan(q_indptr, kv_indptr, kv_indices, kv_lens,
|
||||
num_heads, head_dim_ckv, head_dim_kpe,
|
||||
page_size, False, sm_scale,
|
||||
q_nope.dtype, ckv.dtype)
|
||||
wrapper.run(q_nope, q_pe, ckv, kpe, out=output, return_lse=False)
|
||||
```
|
||||
|
||||
如何提交代码详见[评测指南](/d/2)。
|
||||
|
|
@ -0,0 +1,46 @@
|
|||
---
|
||||
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_nope,
|
||||
const __nv_bfloat16* q_pe,
|
||||
const __nv_bfloat16* ckv,
|
||||
const __nv_bfloat16* kpe,
|
||||
__nv_bfloat16* output,
|
||||
const int32_t* q_indptr,
|
||||
const int32_t* kv_indptr,
|
||||
const int32_t* kv_indices,
|
||||
const int32_t* kv_lens,
|
||||
int64_t batch_size,
|
||||
int64_t seq_len,
|
||||
int64_t num_heads,
|
||||
int64_t head_dim_ckv,
|
||||
int64_t head_dim_kpe,
|
||||
int64_t page_size,
|
||||
int64_t causal
|
||||
);
|
||||
```
|
||||
|
||||
### 参数说明
|
||||
|
||||
* `q_nope`：query 的 compressed/nope 部分，shape `(batch_size, num_heads, head_dim_ckv)`，连续 `bf16`
|
||||
* `q_pe`：query 的 RoPE 部分，shape `(batch_size, num_heads, head_dim_kpe)`，连续 `bf16`
|
||||
* `ckv`：compressed KV cache，shape `(batch_size * seq_len, 1, head_dim_ckv)`，连续 `bf16`
|
||||
* `kpe`：RoPE KV cache，shape `(batch_size * seq_len, 1, head_dim_kpe)`，连续 `bf16`
|
||||
* `output`：输出缓冲区，shape `(batch_size, num_heads, head_dim_ckv)`，连续 `bf16`
|
||||
* `q_indptr`：decode query indptr，shape `(batch_size + 1)`，内容为 `[0, 1, ..., batch_size]`
|
||||
* `kv_indptr`：KV indptr，shape `(batch_size + 1)`，每段长度为 `seq_len`
|
||||
* `kv_indices`：page index，shape `(batch_size * seq_len)`，连续 `int32`
|
||||
* `kv_lens`：每个 batch 的 KV 长度，shape `(batch_size)`，连续 `int32`
|
||||
* `page_size`：评测中固定为 `1`
|
||||
* `causal`：评测中固定为 `0`
|
||||
|
||||
`run_kernel` 内部需要自行计算合适的 launch 配置并启动 CUDA kernel。为保证计时准确，不建议在 `run_kernel` 内部做 `cudaDeviceSynchronize()` 或显式同步。
|
||||
|
|
@ -0,0 +1,26 @@
|
|||
---
|
||||
sectionTitle: "输入格式"
|
||||
type: "Text"
|
||||
---
|
||||
## 数据范围
|
||||
|
||||
- 数据类型：输入输出均为 `bfloat16`
|
||||
- `head_dim_ckv = 512`
|
||||
- `head_dim_kpe = 64`
|
||||
- `page_size = 1`
|
||||
- `causal = 0`
|
||||
- `num_heads` 取值为 `64, 128`
|
||||
- `batch_size` 取值为 `1, 4, 16`
|
||||
- `seq_len` 取值为 `1024, 4096, 8192, 16384`
|
||||
- `sm_scale = 1 / sqrt(head_dim_ckv + head_dim_kpe)`
|
||||
|
||||
测试点顺序与 `McFlashInfer/benchmarks/bench_batch_mla.py` 中的 cases 一致，即：
|
||||
|
||||
```python
|
||||
for num_heads in [64, 128]:
|
||||
for batch_size in [1, 4, 16]:
|
||||
for seq_len in [1024, 4096, 8192, 16384]:
|
||||
...
|
||||
```
|
||||
|
||||
输出参与 `torch.allclose` 校验，容差为 `rtol=1e-2, atol=1e-2`。
|
||||
|
|
@ -0,0 +1,30 @@
|
|||
---
|
||||
sectionTitle: "输出格式"
|
||||
type: "Text"
|
||||
---
|
||||
PyTorch 参考实现
|
||||
|
||||
```python
|
||||
def baseline(q_nope, q_pe, ckv, kpe, output,
|
||||
q_indptr, kv_indptr, kv_indices, kv_lens,
|
||||
batch_size, seq_len, num_heads,
|
||||
head_dim_ckv, head_dim_kpe, page_size, causal):
|
||||
sm_scale = 1.0 / ((head_dim_ckv + head_dim_kpe) ** 0.5)
|
||||
workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q_nope.device)
|
||||
wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(workspace_buffer, backend="auto")
|
||||
wrapper.plan(
|
||||
q_indptr,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
kv_lens,
|
||||
num_heads,
|
||||
head_dim_ckv,
|
||||
head_dim_kpe,
|
||||
page_size,
|
||||
bool(causal),
|
||||
sm_scale,
|
||||
q_nope.dtype,
|
||||
ckv.dtype,
|
||||
)
|
||||
wrapper.run(q_nope, q_pe, ckv, kpe, out=output, return_lse=False)
|
||||
```
|
||||
|
|
@ -0,0 +1 @@
|
|||
FlashInfer MLA Paged Attention
|
||||
|
|
@ -0,0 +1,16 @@
|
|||
{
|
||||
"id": 196,
|
||||
"displayId": 20004,
|
||||
"type": "Traditional",
|
||||
"isPublic": false,
|
||||
"locales": [
|
||||
"zh_CN"
|
||||
],
|
||||
"samples": [
|
||||
{
|
||||
"inputData": "1\n",
|
||||
"outputData": ""
|
||||
}
|
||||
],
|
||||
"problemTagIds": []
|
||||
}
|
||||
|
|
@ -0,0 +1,248 @@
|
|||
from __future__ import annotations
|
||||
|
||||
|
||||
BATCH_SIZES = [1, 2, 4, 8, 16]
|
||||
SEQ_LENS_KV = [512, 1024, 2048, 4096, 8192, 16384]
|
||||
HEAD_DIMS = [128]
|
||||
NUM_QO_HEADS = 32
|
||||
PAGE_BLOCK_SIZE = 16
|
||||
|
||||
|
||||
def _build_cases():
|
||||
cases = []
|
||||
for batch_size in BATCH_SIZES:
|
||||
for seq_len_kv in SEQ_LENS_KV:
|
||||
for head_dim in HEAD_DIMS:
|
||||
num_kv_heads = 8 if head_dim == 64 else 4
|
||||
cases.append(
|
||||
(
|
||||
batch_size,
|
||||
seq_len_kv,
|
||||
NUM_QO_HEADS,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
PAGE_BLOCK_SIZE,
|
||||
)
|
||||
)
|
||||
return cases
|
||||
|
||||
|
||||
TESTCASES = _build_cases()
|
||||
|
||||
|
||||
def getNumOfTestcases() -> int:
|
||||
return len(TESTCASES)
|
||||
|
||||
|
||||
try:
|
||||
from pathlib import Path
|
||||
from typing import List, Tuple, Union
|
||||
import sys
|
||||
|
||||
import torch
|
||||
|
||||
KernelArg = Union[torch.Tensor, int, float]
|
||||
CURRENT_CASE = None
|
||||
|
||||
def _ensure_flashinfer_importable():
|
||||
try:
|
||||
import flashinfer # noqa: F401
|
||||
|
||||
return
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
here = Path(__file__).resolve()
|
||||
for parent in here.parents:
|
||||
candidate = parent / "McFlashInfer"
|
||||
if (candidate / "flashinfer").is_dir():
|
||||
sys.path.insert(0, str(candidate))
|
||||
return
|
||||
|
||||
def _get_testcase_index() -> int:
|
||||
try:
|
||||
raw = input().strip()
|
||||
except EOFError:
|
||||
return 0
|
||||
if raw == "":
|
||||
return 0
|
||||
try:
|
||||
testcase_id = int(raw.split()[0])
|
||||
except ValueError:
|
||||
return 0
|
||||
if 1 <= testcase_id <= len(TESTCASES):
|
||||
return testcase_id - 1
|
||||
if 0 <= testcase_id < len(TESTCASES):
|
||||
return testcase_id
|
||||
return 0
|
||||
|
||||
def _compute_reps(batch_size: int, seq_len: int, head_dim: int, base_reps: int = 100) -> int:
|
||||
workload = batch_size * seq_len * head_dim
|
||||
if workload < 1e5:
|
||||
return base_reps
|
||||
if workload < 1e6:
|
||||
return base_reps // 2
|
||||
if workload < 1e7:
|
||||
return base_reps // 4
|
||||
if workload < 1e8:
|
||||
return base_reps // 8
|
||||
if workload < 1e9:
|
||||
return base_reps // 16
|
||||
return base_reps // 32
|
||||
|
||||
def _setup_paged_kv_indptr(batch_size: int, seq_len: int, page_block_size: int, device: str):
|
||||
seq_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, device=device)
|
||||
seq_lens_blocks = torch.div(seq_lens + page_block_size - 1, page_block_size, rounding_mode="floor")
|
||||
kv_indptr = torch.empty((batch_size + 1,), dtype=torch.int32, device=device)
|
||||
kv_indptr[0] = 0
|
||||
kv_indptr[1:] = torch.cumsum(seq_lens_blocks, dim=0)
|
||||
num_blocks = int(kv_indptr[-1].item())
|
||||
last_page_len = (seq_lens - 1) % page_block_size + 1
|
||||
return kv_indptr, last_page_len, num_blocks
|
||||
|
||||
def getTestCaseSize() -> Tuple[List[Tuple[int, ...]], Tuple[int, int]]:
|
||||
testcase_id = _get_testcase_index()
|
||||
global CURRENT_CASE
|
||||
batch_size, seq_len_kv, num_qo_heads, num_kv_heads, head_dim, page_block_size = TESTCASES[testcase_id]
|
||||
num_blocks = batch_size * ((seq_len_kv + page_block_size - 1) // page_block_size)
|
||||
CURRENT_CASE = (
|
||||
batch_size,
|
||||
seq_len_kv,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
page_block_size,
|
||||
num_blocks,
|
||||
20260710 + testcase_id,
|
||||
)
|
||||
warmup = 3
|
||||
iters = max(1, _compute_reps(batch_size, seq_len_kv, head_dim))
|
||||
return [
|
||||
(batch_size, num_qo_heads, head_dim),
|
||||
(num_blocks, 2, page_block_size, num_kv_heads, head_dim),
|
||||
(batch_size, num_qo_heads, head_dim),
|
||||
(batch_size + 1,),
|
||||
(num_blocks,),
|
||||
(batch_size,),
|
||||
(), (), (), (), (), (),
|
||||
], (warmup, iters)
|
||||
|
||||
def genTestCase(testcase_sizes, device: str = "cuda") -> List[KernelArg]:
|
||||
del testcase_sizes
|
||||
batch_size, seq_len_kv, num_qo_heads, num_kv_heads, head_dim, page_block_size, num_blocks, seed = CURRENT_CASE
|
||||
gen = torch.Generator(device=device)
|
||||
gen.manual_seed(seed)
|
||||
dtype = torch.bfloat16
|
||||
q = torch.rand(batch_size, num_qo_heads, head_dim, dtype=dtype, device=device, generator=gen).contiguous()
|
||||
kv_data = torch.randn(
|
||||
num_blocks,
|
||||
2,
|
||||
page_block_size,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
generator=gen,
|
||||
).contiguous()
|
||||
output = torch.empty(batch_size, num_qo_heads, head_dim, dtype=dtype, device=device)
|
||||
kv_indptr, last_page_len, num_blocks_check = _setup_paged_kv_indptr(
|
||||
batch_size,
|
||||
seq_len_kv,
|
||||
page_block_size,
|
||||
device,
|
||||
)
|
||||
assert num_blocks_check == num_blocks
|
||||
kv_indices = torch.arange(num_blocks, dtype=torch.int32, device=device)
|
||||
return [
|
||||
q,
|
||||
kv_data,
|
||||
output,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
last_page_len,
|
||||
batch_size,
|
||||
seq_len_kv,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
page_block_size,
|
||||
]
|
||||
|
||||
def baseline(
|
||||
q,
|
||||
kv_data,
|
||||
output,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
last_page_len,
|
||||
batch_size,
|
||||
seq_len_kv,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
page_block_size,
|
||||
):
|
||||
_ensure_flashinfer_importable()
|
||||
import flashinfer
|
||||
|
||||
workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device)
|
||||
wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
|
||||
workspace_buffer,
|
||||
kv_layout="NHD",
|
||||
use_tensor_cores=True,
|
||||
)
|
||||
wrapper.plan(
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
last_page_len,
|
||||
int(num_qo_heads),
|
||||
int(num_kv_heads),
|
||||
int(head_dim),
|
||||
int(page_block_size),
|
||||
data_type=torch.bfloat16,
|
||||
q_data_type=torch.bfloat16,
|
||||
)
|
||||
wrapper.run(q, kv_data, out=output)
|
||||
return [
|
||||
q,
|
||||
kv_data,
|
||||
output,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
last_page_len,
|
||||
batch_size,
|
||||
seq_len_kv,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
page_block_size,
|
||||
]
|
||||
|
||||
def check(
|
||||
testcase_sizes,
|
||||
original_input_tensors,
|
||||
target_kernel_input_tensors,
|
||||
baseline_input_tensors,
|
||||
rtol=1e-2,
|
||||
atol=1e-2,
|
||||
) -> bool:
|
||||
del testcase_sizes, original_input_tensors
|
||||
output_t = target_kernel_input_tensors[2]
|
||||
output_ref = baseline_input_tensors[2]
|
||||
if output_t.shape != output_ref.shape:
|
||||
print(f"[FAIL] shape mismatch: target {output_t.shape}, ref {output_ref.shape}", file=sys.stderr)
|
||||
return False
|
||||
if output_t.dtype != output_ref.dtype:
|
||||
print(f"[FAIL] dtype mismatch: target {output_t.dtype}, ref {output_ref.dtype}", file=sys.stderr)
|
||||
return False
|
||||
if not torch.allclose(output_t.float(), output_ref.float(), rtol=rtol, atol=atol):
|
||||
diff = (output_t.float() - output_ref.float()).abs()
|
||||
print(
|
||||
f"[FAIL] allclose failed: max_abs_diff={float(diff.max().item()):.6f}, "
|
||||
f"mean_abs_diff={float(diff.mean().item()):.6f} (rtol={rtol}, atol={atol})",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return False
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
|
|
@ -0,0 +1,22 @@
|
|||
---
|
||||
sectionTitle: "题目描述"
|
||||
type: "Text"
|
||||
---
|
||||
你需要实现 FlashInfer paged KV cache decode 的 CUDA C++ 前向算子。
|
||||
|
||||
本题输入采用 FlashInfer `BatchDecodeWithPagedKVCacheWrapper` 的 paged `NHD` 布局。每个 batch 只有 1 个 query token，KV cache 长度为 `seq_len_kv`。
|
||||
|
||||
评测程序会调用你提交代码中的 `run_kernel` 函数。你需要根据 `kv_indptr`、`kv_indices` 和 `last_page_len` 读取 paged KV cache，并将结果写入 `output`。
|
||||
|
||||
baseline 使用 FlashInfer paged decode 的 Python API：
|
||||
|
||||
```python
|
||||
wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
|
||||
workspace, kv_layout="NHD", use_tensor_cores=True)
|
||||
wrapper.plan(kv_indptr, kv_indices, last_page_len,
|
||||
num_qo_heads, num_kv_heads, head_dim, page_block_size,
|
||||
data_type=torch.bfloat16, q_data_type=torch.bfloat16)
|
||||
wrapper.run(q, kv_data, out=output)
|
||||
```
|
||||
|
||||
如何提交代码详见[评测指南](/d/2)。
|
||||
|
|
@ -0,0 +1,38 @@
|
|||
---
|
||||
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* kv_data,
|
||||
__nv_bfloat16* output,
|
||||
const int32_t* kv_indptr,
|
||||
const int32_t* kv_indices,
|
||||
const int32_t* last_page_len,
|
||||
int64_t batch_size,
|
||||
int64_t seq_len_kv,
|
||||
int64_t num_qo_heads,
|
||||
int64_t num_kv_heads,
|
||||
int64_t head_dim,
|
||||
int64_t page_block_size
|
||||
);
|
||||
```
|
||||
|
||||
### 参数说明
|
||||
|
||||
* `q`：decode query tensor，shape `(batch_size, num_qo_heads, head_dim)`，连续 `bf16`
|
||||
* `kv_data`：paged KV cache，shape `(num_blocks, 2, page_block_size, num_kv_heads, head_dim)`，连续 `bf16`，其中 `kv_data[:, 0]` 为 key，`kv_data[:, 1]` 为 value
|
||||
* `output`：输出缓冲区，shape `(batch_size, num_qo_heads, head_dim)`，连续 `bf16`
|
||||
* `kv_indptr`：paged KV indptr，shape `(batch_size + 1)`，连续 `int32`
|
||||
* `kv_indices`：page index，shape `(num_blocks)`，连续 `int32`
|
||||
* `last_page_len`：每个 batch 最后一个 page 的有效 token 数，shape `(batch_size)`，连续 `int32`
|
||||
* `page_block_size`：page size，评测中固定为 `16`
|
||||
|
||||
`run_kernel` 内部需要自行计算合适的 launch 配置并启动 CUDA kernel。为保证计时准确，不建议在 `run_kernel` 内部做 `cudaDeviceSynchronize()` 或显式同步。
|
||||
|
|
@ -0,0 +1,26 @@
|
|||
---
|
||||
sectionTitle: "输入格式"
|
||||
type: "Text"
|
||||
---
|
||||
## 数据范围
|
||||
|
||||
- 数据类型：`q/kv_data/output` 均为 `bfloat16`
|
||||
- KV layout：`NHD`
|
||||
- `seq_len_q = 1`
|
||||
- `num_qo_heads = 32`
|
||||
- `page_block_size = 16`
|
||||
- `head_dim` 固定为 `128`
|
||||
- `num_kv_heads = 4`
|
||||
- `batch_size` 取值为 `1, 2, 4, 8, 16`
|
||||
- `seq_len_kv` 取值为 `512, 1024, 2048, 4096, 8192, 16384`
|
||||
|
||||
测试点顺序与 `McFlashInfer/benchmarks/bench_batch_decode.py` 中的 cases 一致，即：
|
||||
|
||||
```python
|
||||
for batch_size in [1, 2, 4, 8, 16]:
|
||||
for seq_len_kv in [512, 1024, 2048, 4096, 8192, 16384]:
|
||||
for head_dim in [128]:
|
||||
...
|
||||
```
|
||||
|
||||
输出参与 `torch.allclose` 校验，容差为 `rtol=1e-2, atol=1e-2`。
|
||||
|
|
@ -0,0 +1,29 @@
|
|||
---
|
||||
sectionTitle: "输出格式"
|
||||
type: "Text"
|
||||
---
|
||||
PyTorch 参考实现
|
||||
|
||||
```python
|
||||
def baseline(q, kv_data, output, kv_indptr, kv_indices, last_page_len,
|
||||
batch_size, seq_len_kv, num_qo_heads,
|
||||
num_kv_heads, head_dim, page_block_size):
|
||||
workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device)
|
||||
wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
|
||||
workspace_buffer,
|
||||
kv_layout="NHD",
|
||||
use_tensor_cores=True,
|
||||
)
|
||||
wrapper.plan(
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
last_page_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
page_block_size,
|
||||
data_type=torch.bfloat16,
|
||||
q_data_type=torch.bfloat16,
|
||||
)
|
||||
wrapper.run(q, kv_data, out=output)
|
||||
```
|
||||
|
|
@ -0,0 +1 @@
|
|||
FlashInfer Paged Decode
|
||||
Loading…
Reference in New Issue