AddingZero
/
op_optimization
forked from metax-maca/op_optimization
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

更新Flashinfer问答文档

This commit is contained in:
Ke Xiao (i26293) 2026-06-18 17:54:48 +08:00
parent 0a1db9a6a6
commit adf3096c01
2 changed files with 163 additions and 133 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -150,6 +150,8 @@ python -c "import torch; print(f'GPU available: {torch.cuda.is_available()}'); p
# 检查依赖版本
python -c "import torch; print(f'PyTorch {torch.__version__}')"
python -c "import einops; print('einops OK')"
# 安装必要依赖
pip install pandas
```
**预期结果：**
@ -172,7 +174,6 @@ python -c "import einops; print('einops OK')"
![](https://origin.picgo.net/2026/06/18/-2026-06-18-164201----5fed48dda82423727.png)
**常见问题：**
@ -204,7 +205,7 @@ python -c "import einops; print('einops OK')"
3. 切换到项目目录 `FlashInfer_Baseline`
```bash
cd data/operator_task_package/flashinfer_task_package/FlashInfer_Baseline
ls -al
ls
```
**预期结果：**
@ -585,7 +586,7 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
```
以上代码仅用于说明接口结构，不代表最优实现，也不作为评分参考。
4. 点击提交，等待评测结果返回；
4. 点击提交，等待评测结果返回
![image 20260616155023443](https://origin.picgo.net/2026/06/16/image-202606161550234433e54d703d581f858.png)

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

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