forked from metax-maca/op_optimization
Merge branch 'master' of https://gitlink.org.cn/metax-maca/op_optimization
This commit is contained in:
commit
cab590db95
|
|
@ -260,7 +260,10 @@
|
|||
[GitHub - MetaX-MACA/flashattn · GitHub](https://github.com/MetaX-MACA/flashattn)
|
||||
|
||||
> 链接内容可供用于学习 API、算子实现思路、benchmark 方法和优化策略。选手仍需根据 XPU-OJ 题包接口**自行实现**可提交的 `run_kernel(...)`
|
||||
---
|
||||
|
||||
**mctlass 组件仓库**: [GitHub - MetaX-MACA/mcTlass · GitHub](https://github.com/MetaX-MACA/mcTlass)
|
||||
|
||||
> 沐曦版 CUTLASS 组件，本题包要求核心矩阵计算必须基于此实现。仓库包含头文件和示例代码，可参考其接口设计和使用方式。
|
||||
|
||||
## 6. 项目实践:FlashAttention KV-Cache Benchmark & OJ 评测
|
||||
|
||||
|
|
@ -303,8 +306,13 @@ python -c "import torch; print(f'PyTorch {torch.__version__}')"
|
|||
python -c "import flash_attn; print(f'flash-attn {flash_attn.__version__}')"
|
||||
python -c "import einops; print('einops OK')"
|
||||
|
||||
# 检查 mctlass 头文件是否可用（在镜像环境中应已预装）
|
||||
ls $MACA_PATH/include/mctlass/ 2>/dev/null || echo "mctlass headers not found in default path"
|
||||
|
||||
|
||||
```
|
||||
|
||||
|
||||
**预期结果:**
|
||||
|
||||
* `mx-smi` 显示沐曦 GPU 信息
|
||||
|
|
@ -326,6 +334,10 @@ python -c "import einops; print('einops OK')"
|
|||
|
||||
|
||||
|
||||
* mctlass头文件可用
|
||||
|
||||
|
||||
|
||||
|
||||
**常见问题:**
|
||||
|
||||
|
|
@ -615,6 +627,9 @@ extern "C" void run_kernel(
|
|||
|
||||
```
|
||||
|
||||
> ⚠️ **重要提交限制**：
|
||||
> 在使用 CUDA MACA 提交 FlashAttention KV cache decode 算子时，`QKᵀ`、`PV` 等 attention 核心矩阵计算应使用沐曦提供的 `mctlass` 组件或其基础计算原语实现，不应完全采用手写 CUDA 循环替代 `mctlass` 完成核心矩阵乘法逻辑。可用的头文件位于：https://github.com/MetaX-MACA/flashattn/tree/2.25.2/csrc/mctlass/include/mctlass ，可以通过 `#include "mctlass/mctlass.h"` 的方式导入。
|
||||
|
||||
#### 6.9.2 参数说明
|
||||
|
||||
| 参数 | 说明 |
|
||||
|
|
@ -754,13 +769,42 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
|
||||
|
||||
|
||||
为方便参赛者先跑通完整流程,这里直接提供一份完整的冒烟代码，使用语言为`CUDA Maca`,可直接复制粘贴到右侧编辑器,用于验证提交链路是否正常:
|
||||
> ⚠️ **再次强调（关于冒烟代码）**：
|
||||
> 为方便参赛者先跑通完整提交链路，下文提供的冒烟代码已使用 `mctlass` 组件实现核心计算逻辑。请注意，该代码仅用于验证接口签名和平台环境是否正常。在确认链路畅通后，后续的算子优化迭代仍需确保使用 `mctlass`（沐曦版 CUTLASS）实现，否则最终评测将受限或无法通过。
|
||||
|
||||
```C++
|
||||
#include <stdint.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <math.h>
|
||||
|
||||
#include "mctlass/mctlass.h"
|
||||
#include "mctlass/bfloat16.h"
|
||||
#include "mctlass/numeric_conversion.h"
|
||||
|
||||
static_assert(sizeof(mctlass::bfloat16_t) == sizeof(__nv_bfloat16),
|
||||
"mctlass::bfloat16_t and __nv_bfloat16 must both be 16-bit values");
|
||||
|
||||
union Bf16Bits {
|
||||
__nv_bfloat16 nv;
|
||||
uint16_t bits;
|
||||
};
|
||||
|
||||
MCTLASS_DEVICE float mctlass_bf16_to_float(__nv_bfloat16 x) {
|
||||
Bf16Bits u;
|
||||
u.nv = x;
|
||||
mctlass::bfloat16_t mx = mctlass::bfloat16_t::bitcast(u.bits);
|
||||
mctlass::NumericConverter<float, mctlass::bfloat16_t> convert;
|
||||
return convert(mx);
|
||||
}
|
||||
|
||||
MCTLASS_DEVICE __nv_bfloat16 mctlass_float_to_bf16(float x) {
|
||||
mctlass::NumericConverter<mctlass::bfloat16_t, float> convert;
|
||||
mctlass::bfloat16_t mx = convert(x);
|
||||
Bf16Bits u;
|
||||
u.bits = mx.raw();
|
||||
return u.nv;
|
||||
}
|
||||
|
||||
__global__ void paged_kv_decode_kernel(
|
||||
const __nv_bfloat16* __restrict__ q,
|
||||
const __nv_bfloat16* __restrict__ k_cache_paged,
|
||||
|
|
@ -783,18 +827,14 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
int ty = threadIdx.y;
|
||||
int num_warps = blockDim.y;
|
||||
|
||||
// 当前 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) {
|
||||
|
|
@ -802,38 +842,36 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
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);
|
||||
output[(int64_t)b * num_heads * headdim + (int64_t)h * headdim + i] =
|
||||
mctlass_float_to_bf16(0.0f);
|
||||
}
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// 第 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]);
|
||||
q_smem[i] = mctlass_bf16_to_float(q[q_idx]);
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// 独立维持 FlashAttention Softmax 局部状态
|
||||
float m_warp = -1e20f;
|
||||
float d_warp = 0.0f;
|
||||
float acc[32]; // 支持 max_headdim = 32*32=1024 (安全冗余)
|
||||
float acc[32];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 32; ++i) {
|
||||
acc[i] = 0.0f;
|
||||
}
|
||||
|
||||
float scale = 1.0f / sqrtf((float)headdim);
|
||||
int kv_h = h / (num_heads / num_heads_k); // MQA 或 GQA 支持映射
|
||||
int kv_h = h / (num_heads / num_heads_k);
|
||||
|
||||
// 每个 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;
|
||||
|
|
@ -843,30 +881,26 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
+ (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]);
|
||||
float k_val = mctlass_bf16_to_float(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;
|
||||
|
|
@ -874,13 +908,12 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
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]);
|
||||
float v_val = mctlass_bf16_to_float(v_cache_paged[v_base + i]);
|
||||
acc[step] = acc[step] * exp_old + exp_val * v_val;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Warp 将局部运算结果写入共享内存
|
||||
if (tx == 0) {
|
||||
smem_m[ty] = m_warp;
|
||||
smem_d[ty] = d_warp;
|
||||
|
|
@ -893,19 +926,17 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
}
|
||||
__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) {
|
||||
|
|
@ -915,7 +946,7 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
}
|
||||
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);
|
||||
output[out_idx] = mctlass_float_to_bf16(out_val);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -938,15 +969,11 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
int64_t num_blocks,
|
||||
int64_t causal)
|
||||
{
|
||||
// 配置 Kernel 计算资源与 Layout
|
||||
int num_warps = 8;
|
||||
dim3 block(32, num_warps);
|
||||
dim3 grid(batch_size, num_heads);
|
||||
|
||||
// 动态计算需要的共享内存大小 (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>>>(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,187 @@
|
|||
"""FlashAttention KV Cache Decode in TileLang."""
|
||||
|
||||
import tilelang
|
||||
import tilelang.language as T
|
||||
from tilelang import jit
|
||||
|
||||
NUM_SPLITS = 4
|
||||
real_kernel = None
|
||||
|
||||
@jit(
|
||||
pass_configs={
|
||||
tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: False,
|
||||
},
|
||||
)
|
||||
def build_kernel(
|
||||
batch_size,
|
||||
num_heads,
|
||||
num_heads_k,
|
||||
headdim,
|
||||
page_block_size,
|
||||
num_blocks,
|
||||
causal,
|
||||
):
|
||||
blocks_per_batch = num_blocks // batch_size
|
||||
assert blocks_per_batch % NUM_SPLITS == 0, (
|
||||
f"blocks_per_batch={blocks_per_batch} must be divisible by NUM_SPLITS={NUM_SPLITS}"
|
||||
)
|
||||
blocks_per_split = blocks_per_batch // NUM_SPLITS
|
||||
|
||||
BLOCK_M = 1
|
||||
BLOCK_N = page_block_size
|
||||
scale = (1.0 / headdim) ** 0.5 * 1.44269504 # log2(e)
|
||||
dtype = "bfloat16"
|
||||
accum_dtype = "float32"
|
||||
|
||||
# Use a large-negative-finite sentinel instead of -inf to avoid
|
||||
# (-inf) - (-inf) = NaN when an entire split is masked out.
|
||||
NEG_INF_SAFE = -1e30
|
||||
|
||||
@T.prim_func
|
||||
def kernel(
|
||||
Q: T.Tensor([batch_size, 1, num_heads, headdim], dtype),
|
||||
K: T.Tensor([num_blocks, page_block_size, num_heads_k, headdim], dtype),
|
||||
V: T.Tensor([num_blocks, page_block_size, num_heads_k, headdim], dtype),
|
||||
Output: T.Tensor([batch_size, 1, num_heads, headdim], dtype),
|
||||
cache_seqlens: T.Tensor([batch_size], "int32"),
|
||||
block_table: T.Tensor([batch_size, blocks_per_batch], "int32"),
|
||||
):
|
||||
# float32 workspace — avoids BF16StorageLegalize var-remap bug
|
||||
glse = T.alloc_global([batch_size, num_heads, NUM_SPLITS], accum_dtype)
|
||||
Output_partial = T.alloc_global(
|
||||
[batch_size, 1, num_heads, NUM_SPLITS, headdim], accum_dtype
|
||||
)
|
||||
|
||||
# ============= Stage 1: split kernel =============
|
||||
with T.Kernel(NUM_SPLITS, num_heads, batch_size, threads=128) as (bs, bh, bz):
|
||||
Q_shared = T.alloc_shared([BLOCK_M, headdim], dtype)
|
||||
K_shared = T.alloc_shared([BLOCK_N, headdim], dtype)
|
||||
V_shared = T.alloc_shared([BLOCK_N, headdim], dtype)
|
||||
acc_s = T.alloc_fragment([BLOCK_M, BLOCK_N], accum_dtype)
|
||||
acc_o = T.alloc_fragment([BLOCK_M, headdim], accum_dtype)
|
||||
scores_max = T.alloc_fragment([BLOCK_M], accum_dtype)
|
||||
scores_max_prev = T.alloc_fragment([BLOCK_M], accum_dtype)
|
||||
scores_scale = T.alloc_fragment([BLOCK_M], accum_dtype)
|
||||
scores_sum = T.alloc_fragment([BLOCK_M], accum_dtype)
|
||||
logsum = T.alloc_fragment([BLOCK_M], accum_dtype)
|
||||
|
||||
T.copy(Q[bz, 0, bh, :], Q_shared)
|
||||
|
||||
kv_seqlen = cache_seqlens[bz]
|
||||
split_k_start = bs * blocks_per_split
|
||||
|
||||
T.fill(acc_o, 0)
|
||||
T.fill(logsum, 0)
|
||||
# KEY FIX: use -1e30 instead of -inf to avoid (-inf)-(-inf)=NaN
|
||||
T.fill(scores_max, NEG_INF_SAFE)
|
||||
|
||||
for k in T.Pipelined(blocks_per_split, num_stages=2):
|
||||
global_k = split_k_start + k
|
||||
physical_block = block_table[bz, global_k]
|
||||
tok_offset = global_k * page_block_size
|
||||
|
||||
# ----- Q @ K^T (hand-written, M=1, masked) -----
|
||||
T.copy(K[physical_block, 0:BLOCK_N, bh, :], K_shared)
|
||||
T.fill(acc_s, 0)
|
||||
for j in T.Parallel(BLOCK_N):
|
||||
if tok_offset + j < kv_seqlen:
|
||||
for d in T.serial(headdim):
|
||||
acc_s[0, j] = acc_s[0, j] + Q_shared[0, d] * K_shared[j, d]
|
||||
else:
|
||||
acc_s[0, j] = -T.infinity(accum_dtype)
|
||||
|
||||
# ----- online softmax -----
|
||||
T.copy(scores_max, scores_max_prev)
|
||||
# KEY FIX: use -1e30 instead of -inf here too
|
||||
T.fill(scores_max, NEG_INF_SAFE)
|
||||
T.reduce_max(acc_s, scores_max, dim=1, clear=False)
|
||||
scores_max[0] = T.max(scores_max[0], scores_max_prev[0])
|
||||
# (prev - cur) is now (finite - finite) = 0 when both are sentinel,
|
||||
# never (-inf - (-inf)) = NaN
|
||||
scores_scale[0] = T.exp2((scores_max_prev[0] - scores_max[0]) * scale)
|
||||
for j in T.Parallel(BLOCK_N):
|
||||
acc_s[0, j] = T.exp2((acc_s[0, j] - scores_max[0]) * scale)
|
||||
T.reduce_sum(acc_s, scores_sum, dim=1)
|
||||
logsum[0] = logsum[0] * scores_scale[0] + scores_sum[0]
|
||||
for d in T.Parallel(headdim):
|
||||
acc_o[0, d] = acc_o[0, d] * scores_scale[0]
|
||||
|
||||
# ----- P @ V (hand-written, fp32 accum) -----
|
||||
T.copy(V[physical_block, 0:BLOCK_N, bh, :], V_shared)
|
||||
for d in T.Parallel(headdim):
|
||||
for j in T.serial(BLOCK_N):
|
||||
acc_o[0, d] = acc_o[0, d] + acc_s[0, j] * V_shared[j, d]
|
||||
|
||||
# ----- final normalise & write partial state -----
|
||||
# KEY FIX: add epsilon to avoid 0/0 = NaN when split is all-masked
|
||||
safe_logsum = logsum[0] + 1e-30
|
||||
for d in T.Parallel(headdim):
|
||||
acc_o[0, d] = acc_o[0, d] / safe_logsum
|
||||
|
||||
lse_local = T.alloc_fragment([1], accum_dtype)
|
||||
lse_local[0] = T.log2(safe_logsum) + scores_max[0] * scale
|
||||
glse[bz, bh, bs] = lse_local[0]
|
||||
|
||||
for d in T.Parallel(headdim):
|
||||
Output_partial[bz, 0, bh, bs, d] = acc_o[0, d]
|
||||
|
||||
# ============= Stage 2: combine kernel =============
|
||||
with T.Kernel(num_heads, batch_size, threads=128) as (bh, bz):
|
||||
lse_local = T.alloc_fragment([NUM_SPLITS], accum_dtype)
|
||||
for s in T.serial(NUM_SPLITS):
|
||||
lse_local[s] = glse[bz, bh, s]
|
||||
|
||||
lse_max = T.alloc_fragment([1], accum_dtype)
|
||||
lse_max[0] = -T.infinity(accum_dtype)
|
||||
for s in T.serial(NUM_SPLITS):
|
||||
lse_max[0] = T.max(lse_max[0], lse_local[s])
|
||||
|
||||
lse_logsum = T.alloc_fragment([1], accum_dtype)
|
||||
lse_logsum[0] = 0
|
||||
for s in T.serial(NUM_SPLITS):
|
||||
lse_logsum[0] = lse_logsum[0] + T.exp2(lse_local[s] - lse_max[0])
|
||||
lse_logsum[0] = T.log2(lse_logsum[0]) + lse_max[0]
|
||||
|
||||
o_accum = T.alloc_fragment([headdim], accum_dtype)
|
||||
T.fill(o_accum, 0)
|
||||
for s in T.serial(NUM_SPLITS):
|
||||
s_scale = T.exp2(lse_local[s] - lse_logsum[0])
|
||||
for d in T.Parallel(headdim):
|
||||
o_accum[d] = o_accum[d] + Output_partial[bz, 0, bh, s, d] * s_scale
|
||||
|
||||
for d in T.Parallel(headdim):
|
||||
Output[bz, 0, bh, d] = T.Cast(dtype, o_accum[d])
|
||||
|
||||
return kernel
|
||||
|
||||
|
||||
def run_kernel(
|
||||
q,
|
||||
k_cache_paged,
|
||||
v_cache_paged,
|
||||
output,
|
||||
cache_seqlens,
|
||||
block_table,
|
||||
batch_size,
|
||||
seqlen_k,
|
||||
seqlen_q,
|
||||
num_heads,
|
||||
num_heads_k,
|
||||
headdim,
|
||||
page_block_size,
|
||||
num_blocks,
|
||||
causal,
|
||||
):
|
||||
global real_kernel
|
||||
|
||||
B = int(batch_size)
|
||||
H = int(num_heads)
|
||||
HK = int(num_heads_k)
|
||||
D = int(headdim)
|
||||
PBS = int(page_block_size)
|
||||
NB = int(num_blocks)
|
||||
|
||||
if real_kernel is None:
|
||||
real_kernel = build_kernel(B, H, HK, D, PBS, NB, int(causal))
|
||||
|
||||
real_kernel(q, k_cache_paged, v_cache_paged, output, cache_seqlens, block_table)
|
||||
|
|
@ -0,0 +1,106 @@
|
|||
import triton
|
||||
import triton.language as tl
|
||||
import torch
|
||||
|
||||
@triton.jit
|
||||
def slow_decode_kernel(
|
||||
q_ptr,
|
||||
k_cache_ptr,
|
||||
v_cache_ptr,
|
||||
output_ptr,
|
||||
cache_seqlens_ptr,
|
||||
block_table_ptr,
|
||||
num_heads: tl.constexpr,
|
||||
num_heads_k: tl.constexpr,
|
||||
headdim: tl.constexpr,
|
||||
page_block_size: tl.constexpr,
|
||||
blocks_per_batch,
|
||||
):
|
||||
# 维度索引
|
||||
pid_b = tl.program_id(0) # Batch index
|
||||
pid_h = tl.program_id(1) # Head index
|
||||
|
||||
# GQA Support: 映射 Query Head 到 KV Head
|
||||
kv_head = pid_h * num_heads_k // num_heads
|
||||
|
||||
# 加载实际的 KV 序列长度
|
||||
seq_len = tl.load(cache_seqlens_ptr + pid_b).to(tl.int32)
|
||||
|
||||
# 维度偏移量 [0, 1, ..., headdim-1]
|
||||
offs_d = tl.arange(0, headdim)
|
||||
|
||||
# Online Softmax 累加器
|
||||
acc = tl.zeros([headdim], dtype=tl.float32)
|
||||
l_i = 0.0
|
||||
m_i = float('-inf')
|
||||
scale = 1.0 / tl.sqrt(float(headdim))
|
||||
|
||||
# === 性能瓶颈：串行遍历整个序列 ===
|
||||
# 不使用 Block 并行，而是用单个 Block 串行循环处理所有 Token
|
||||
t = 0
|
||||
while t < seq_len:
|
||||
# 性能瓶颈：每次循环都重新加载 Q，增加显存压力
|
||||
q = tl.load(q_ptr + pid_b * num_heads * headdim + pid_h * headdim + offs_d).to(tl.float32)
|
||||
q = q * scale
|
||||
|
||||
# Paged KV 映射逻辑
|
||||
page_idx = t // page_block_size
|
||||
page_off = t % page_block_size
|
||||
|
||||
# 查表获取物理 Block 索引
|
||||
# blocks_per_batch 是计算出来的步长
|
||||
phys_block = tl.load(block_table_ptr + pid_b * blocks_per_batch + page_idx)
|
||||
|
||||
# 计算 K 和 V 的物理地址
|
||||
# Layout: (num_blocks, page_block_size, num_heads_k, headdim)
|
||||
kv_base = phys_block * page_block_size * num_heads_k * headdim + \
|
||||
page_off * num_heads_k * headdim + \
|
||||
kv_head * headdim
|
||||
|
||||
# 加载 K 和 V 向量
|
||||
k = tl.load(k_cache_ptr + kv_base + offs_d).to(tl.float32)
|
||||
v = tl.load(v_cache_ptr + kv_base + offs_d).to(tl.float32)
|
||||
|
||||
# Attention 计算
|
||||
s = tl.sum(q * k) # 点积
|
||||
|
||||
# Online Softmax 更新
|
||||
m_new = tl.maximum(m_i, s)
|
||||
p = tl.exp(s - m_new)
|
||||
alpha = tl.exp(m_i - m_new)
|
||||
|
||||
acc = acc * alpha + p * v
|
||||
l_i = l_i * alpha + p
|
||||
m_i = m_new
|
||||
|
||||
t += 1
|
||||
|
||||
# 写回结果
|
||||
# 这里没有处理 l_i 为 0 的边界情况，但测试数据 seq_len 通常很大
|
||||
out = acc / l_i
|
||||
tl.store(output_ptr + pid_b * num_heads * headdim + pid_h * headdim + offs_d, out)
|
||||
|
||||
def run_kernel(
|
||||
q, k_cache_paged, v_cache_paged, output,
|
||||
cache_seqlens, block_table,
|
||||
batch_size, seqlen_k, seqlen_q, num_heads, num_heads_k, headdim,
|
||||
page_block_size, num_blocks, causal,
|
||||
):
|
||||
# 计算每个 batch 对应的 block_table 行宽
|
||||
blocks_per_batch = num_blocks // batch_size
|
||||
|
||||
# 启动配置：每个 Head 一个 Block
|
||||
# 总 Block 数 = batch_size * num_heads (最大 128个)，并行度极低
|
||||
grid = (batch_size, num_heads)
|
||||
|
||||
slow_decode_kernel[grid](
|
||||
q, k_cache_paged, v_cache_paged, output,
|
||||
cache_seqlens, block_table,
|
||||
num_heads=num_heads,
|
||||
num_heads_k=num_heads_k,
|
||||
headdim=headdim,
|
||||
page_block_size=page_block_size,
|
||||
blocks_per_batch=blocks_per_batch,
|
||||
num_warps=1, # 性能瓶颈：仅使用 1 个 warp，限制计算吞吐
|
||||
num_stages=1, # 性能瓶颈：禁用流水线并行
|
||||
)
|
||||
Loading…
Reference in New Issue