apps/FastDeploy
1
0
Fork 0
mirror of https://github.com/PaddlePaddle/FastDeploy.git synced 2025-12-24 13:28:13 +08:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

[Feature] support flash_mask_attention backend (#5134)

Browse Source 
* [Feature] suppert flash_mask_attention backend

* fix unittest

* clean code
This commit is contained in:
lizhenyun01
2025-11-28 10:12:16 +08:00
committed by GitHub
parent b935101008
commit aba4fc657f
13 changed files with 542 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

143
custom_ops/gpu_ops/append_attn/gqa_rope_write_cache.cu
View File

@@ -24,6 +24,8 @@ __global__ void GQAVariableLengthRotarySplitKernel(
const T *qkv,
const float *cos_emb,
const float *sin_emb,
const float *q_norm_weight,
const float *k_norm_weight,
const int *batch_id_per_token,
const int *cu_seqlens_q,
const int *seq_lens,
@@ -38,37 +40,46 @@ __global__ void GQAVariableLengthRotarySplitKernel(
const int kv_num_head,
const int seq_len,
const int last_dim,
const bool rope_3d) {
const bool rope_3d,
const float rms_norm_eps) {
using LoadT = AlignedVector<T, VecSize>;
constexpr int HalfVecSize = VecSize / 2;
using LoadEmbT = AlignedVector<float, HalfVecSize>;
using LoadFloat = AlignedVector<float, VecSize>;
LoadT src_vec;
LoadEmbT cos_emb_vec;
LoadEmbT sin_emb_vec;
int64_t global_thread_idx = blockDim.x * blockIdx.x + threadIdx.x;
LoadFloat tmp_vec;
LoadFloat q_norm_vec, k_norm_vec;
int64_t global_warp_idx = blockDim.y * blockIdx.x + threadIdx.y;
int64_t all_warp_num = gridDim.x * blockDim.y;
const int half_lastdim = last_dim / 2;
const int offset = (q_num_head + kv_num_head * 2) * last_dim;
for (int64_t linear_index = global_thread_idx * VecSize,
step = gridDim.x * blockDim.x * VecSize;
linear_index < elem_cnt;
linear_index += step) {
const int token_idx = linear_index / offset;
const int ori_bi = batch_id_per_token[token_idx];
const int offset =
(q_num_head + kv_num_head * 2) * last_dim; // for all q,k,v
const int all_head_num = elem_cnt / last_dim;
for (int gloabl_hi = global_warp_idx; gloabl_hi < all_head_num;
gloabl_hi += all_warp_num) {
int64_t linear_index =
gloabl_hi * last_dim + threadIdx.x * VecSize; // 全局index
const int token_idx =
linear_index / offset; // token id(第几个token,不分qkv)
const int ori_bi = batch_id_per_token[token_idx]; // 第几个batch
if (seq_lens[ori_bi] == 0) continue;
const int bias = linear_index % offset;
const int hi = bias / last_dim;
const int h_bias = bias % last_dim;
const int ori_seq_id =
(token_idx - cu_seqlens_q[ori_bi]) + seq_lens_decoder[ori_bi];
const int kv_write_idx = cu_seqlens_k[ori_bi] + ori_seq_id;
const int64_t emb_idx = ori_seq_id * half_lastdim + h_bias / 2;
int64_t new_emb_idx =
rope_3d ? emb_idx + ori_bi * last_dim * seq_len : emb_idx;
(token_idx - cu_seqlens_q[ori_bi]) +
seq_lens_decoder
[ori_bi]; // 在当前seq中的id(拼接了seq到一个batch的情况下有效)
const int64_t emb_idx =
ori_seq_id * half_lastdim + h_bias / 2; // embedding的id
const int64_t base_idx =
token_idx * (q_num_head + 2 * kv_num_head) * last_dim + hi * last_dim +
h_bias;
Load<T, VecSize>(&qkv[base_idx], &src_vec);
const int kv_write_idx = cu_seqlens_k[ori_bi] + ori_seq_id;
int64_t base_split_idx;
T *out_p = nullptr;
if (hi < q_num_head) {
@@ -84,21 +95,67 @@ __global__ void GQAVariableLengthRotarySplitKernel(
base_split_idx = kv_write_idx * kv_num_head * last_dim +
(hi - q_num_head - kv_num_head) * last_dim + h_bias;
}
Load<T, VecSize>(&qkv[base_idx], &src_vec);
// do rope
if (hi < q_num_head + kv_num_head) {
Load<float, HalfVecSize>(&cos_emb[new_emb_idx], &cos_emb_vec);
Load<float, HalfVecSize>(&sin_emb[new_emb_idx], &sin_emb_vec);
// TODO check this correct or not
int64_t new_emb_idx =
rope_3d ? emb_idx + ori_bi * last_dim * seq_len : emb_idx;
float thread_m2 = 0.0f;
float warp_m2 = 0.0f;
if (q_norm_weight && k_norm_weight) {
if (hi < q_num_head + kv_num_head) { // only q and k need rope
Load<float, HalfVecSize>(&cos_emb[new_emb_idx], &cos_emb_vec);
Load<float, HalfVecSize>(&sin_emb[new_emb_idx], &sin_emb_vec);
#pragma unroll
for (int i = 0; i < HalfVecSize; i++) {
const float input_left = static_cast<float>(src_vec[2 * i]);
const float input_right = static_cast<float>(src_vec[2 * i + 1]);
const float cos_tmp = cos_emb_vec[i];
const float sin_tmp = sin_emb_vec[i];
src_vec[2 * i] =
static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
src_vec[2 * i + 1] =
static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
for (int i = 0; i < HalfVecSize; i++) {
const float input_left = static_cast<float>(src_vec[2 * i]);
const float input_right = static_cast<float>(src_vec[2 * i + 1]);
const float cos_tmp = cos_emb_vec[i];
const float sin_tmp = sin_emb_vec[i];
float tmp1 = input_left * cos_tmp - input_right * sin_tmp;
float tmp2 = input_right * cos_tmp + input_left * sin_tmp;
tmp_vec[2 * i] = tmp1;
tmp_vec[2 * i + 1] = tmp2;
thread_m2 += tmp1 * tmp1 + tmp2 * tmp2;
}
}
WelfordWarpAllReduce<float, 32>(thread_m2, &warp_m2); // 单个head的标准差
if (hi < q_num_head + kv_num_head) { // only q and k need norm
float row_variance = max(warp_m2 / last_dim, 0.0f);
float row_inv_var = Rsqrt(row_variance + rms_norm_eps);
if (hi < q_num_head) {
Load<float, VecSize>(&q_norm_weight[threadIdx.x * VecSize],
&q_norm_vec);
#pragma unroll
for (int i = 0; i < VecSize; i++) {
src_vec[i] =
static_cast<T>(tmp_vec[i] * row_inv_var * q_norm_vec[i]);
}
} else {
Load<float, VecSize>(&k_norm_weight[threadIdx.x * VecSize],
&k_norm_vec);
for (int i = 0; i < VecSize; i++) {
src_vec[i] =
static_cast<T>(tmp_vec[i] * row_inv_var * k_norm_vec[i]);
}
}
}
} else {
if (hi < q_num_head + kv_num_head) {
Load<float, HalfVecSize>(&cos_emb[new_emb_idx], &cos_emb_vec);
Load<float, HalfVecSize>(&sin_emb[new_emb_idx], &sin_emb_vec);
#pragma unroll
for (int i = 0; i < HalfVecSize; i++) {
const float input_left = static_cast<float>(src_vec[2 * i]);
const float input_right = static_cast<float>(src_vec[2 * i + 1]);
const float cos_tmp = cos_emb_vec[i];
const float sin_tmp = sin_emb_vec[i];
src_vec[2 * i] =
static_cast<T>(input_left * cos_tmp - input_right * sin_tmp);
src_vec[2 * i + 1] =
static_cast<T>(input_right * cos_tmp + input_left * sin_tmp);
}
}
}
Store<T, VecSize>(src_vec, &qkv_out[base_idx]);
@@ -114,6 +171,8 @@ void gqa_rotary_qk_split_variable(
T *v,
const T *qkv_input,
const float *rotary_emb, // [2, 1, 1, seq_len, dim_head / 2]
const float *q_norm_weight,
const float *k_norm_weight,
const int *batch_id_per_token,
const int *seq_lens_encoder,
const int *seq_lens_decoder,
@@ -126,24 +185,31 @@ void gqa_rotary_qk_split_variable(
const int input_output_len,
const int dim_head,
const bool rope_3d,
const float rms_norm_eps,
const cudaStream_t &stream) {
assert(dim_head == 128 && "dim_head must be 128");
int64_t elem_nums = token_num * (num_heads + 2 * kv_num_heads) * dim_head;
constexpr int PackSize = 16 / sizeof(T);
constexpr int HEAD_DIM = 128;
constexpr int PackSize = HEAD_DIM / kWarpSize;
const int pack_num = elem_nums / PackSize;
const int blocksize = 128;
int grid_size = 1;
GetNumBlocks<128>(pack_num, &grid_size);
dim3 block_size(kWarpSize, blocksize / kWarpSize);
const float *cos_emb = rotary_emb;
const float *sin_emb = rotary_emb + input_output_len * dim_head / 2;
launchWithPdlWhenEnabled(GQAVariableLengthRotarySplitKernel<T, PackSize>,
grid_size,
blocksize,
block_size,
0,
stream,
qkv_input,
cos_emb,
sin_emb,
q_norm_weight,
k_norm_weight,
batch_id_per_token,
cu_seqlens_q,
seq_lens_encoder,
@@ -158,7 +224,8 @@ void gqa_rotary_qk_split_variable(
kv_num_heads,
seq_len,
dim_head,
rope_3d);
rope_3d,
rms_norm_eps);
}
template <typename T,
@@ -1054,6 +1121,8 @@ std::vector<paddle::Tensor> GQARopeWriteCacheKernel(
const paddle::Tensor &cache_batch_ids,
const paddle::Tensor &cache_tile_ids,
const paddle::Tensor &cache_num_blocks,
const paddle::optional<paddle::Tensor> &q_norm_weight,
const paddle::optional<paddle::Tensor> &k_norm_weight,
const paddle::optional<paddle::Tensor> &cache_k_quant_scales,
const paddle::optional<paddle::Tensor> &cache_v_quant_scales,
const paddle::optional<paddle::Tensor> &cache_k_dequant_scales,
@@ -1063,6 +1132,7 @@ std::vector<paddle::Tensor> GQARopeWriteCacheKernel(
const paddle::optional<paddle::Tensor> &kv_signal_data,
const int kv_token_num,
const int max_seq_len,
const float rms_norm_eps,
const std::string &cache_quant_type,
const bool rope_3d) {
typedef PDTraits<paddle::DataType::BFLOAT16> traits_;
@@ -1113,6 +1183,8 @@ std::vector<paddle::Tensor> GQARopeWriteCacheKernel(
v.data<data_t>(),
qkv.data<data_t>(),
rotary_embs.data<float>(),
q_norm_weight ? q_norm_weight.get().data<float>() : nullptr,
k_norm_weight ? k_norm_weight.get().data<float>() : nullptr,
batch_id_per_token.data<int>(),
seq_lens_encoder.data<int>(),
seq_lens_decoder.data<int>(),
@@ -1125,6 +1197,7 @@ std::vector<paddle::Tensor> GQARopeWriteCacheKernel(
rope_3d ? rotary_embs.dims()[3] : rotary_embs.dims()[2],
head_dim,
rope_3d,
rms_norm_eps,
stream);
if (token_num < kv_token_num) {
@@ -1259,6 +1332,8 @@ PD_BUILD_STATIC_OP(gqa_rope_write_cache)
"cache_batch_ids",
"cache_tile_ids_per_batch",
"cache_num_blocks",
paddle::Optional("q_norm_weight"),
paddle::Optional("k_norm_weight"),
paddle::Optional("cache_k_quant_scales"),
paddle::Optional("cache_v_quant_scales"),
paddle::Optional("cache_k_dequant_scales"),
@@ -1271,5 +1346,7 @@ PD_BUILD_STATIC_OP(gqa_rope_write_cache)
{"value_cache", "value_cache_out"}})
.Attrs({"kv_token_num: int",
"max_seq_len: int",
"cache_quant_type: std::string"})
"rms_norm_eps: float",
"cache_quant_type: std::string",
"rope_3d: bool"})
.SetKernelFn(PD_KERNEL(GQARopeWriteCacheKernel));

3
custom_ops/gpu_ops/cpp_extensions.cc
View File

@@ -178,6 +178,8 @@ std::vector<paddle::Tensor> GQARopeWriteCacheKernel(
const paddle::Tensor& cache_batch_ids,
const paddle::Tensor& cache_tile_ids,
const paddle::Tensor& cache_num_blocks,
const paddle::optional<paddle::Tensor>& q_norm_weight,
const paddle::optional<paddle::Tensor>& k_norm_weight,
const paddle::optional<paddle::Tensor>& cache_k_quant_scales,
const paddle::optional<paddle::Tensor>& cache_v_quant_scales,
const paddle::optional<paddle::Tensor>& cache_k_dequant_scales,
@@ -187,6 +189,7 @@ std::vector<paddle::Tensor> GQARopeWriteCacheKernel(
const paddle::optional<paddle::Tensor>& kv_signal_data,
const int kv_token_num,
const int max_seq_len,
const float rms_norm_eps,
const std::string& cache_quant_type,
const bool rope_3d);

29
custom_ops/gpu_ops/flash_mask_attn/flash_mask_attn.cu
View File

@@ -46,12 +46,11 @@ void DispatchFlashAttentionMask(const paddle::Tensor& q_input,
const int kv_head_num,
const int head_dim,
const int max_seq_len,
const int max_enc_len_this_time,
const int max_dec_len_this_time) {
const int q_token_num,
const int k_token_num) {
constexpr int kBlockM = 128;
constexpr int kBlockN = 128;
const int batch_size = seq_len_encoder.dims()[0];
const int batch_size = cu_seq_k.dims()[0] - 1;
Flash_mask_params params;
memset(&params, 0, sizeof(Flash_mask_params));
@@ -63,8 +62,8 @@ void DispatchFlashAttentionMask(const paddle::Tensor& q_input,
params.seq_len_encoder = const_cast<int*>(seq_len_encoder.data<int>());
params.head_num = head_num;
params.kv_head_num = kv_head_num;
params.max_seq_len_q = max_enc_len_this_time;
params.max_seq_len_k = max_enc_len_this_time + max_dec_len_this_time;
params.q_token_num = q_token_num;
params.k_token_num = k_token_num;
params.batch_size = batch_size;
params.gqa_group_size = head_num / kv_head_num;
constexpr float kLog2e = 1.4426950408889634074;
@@ -132,8 +131,8 @@ void FlashAttentionMask(const paddle::Tensor& q_input,
const int kv_head_num,
const int head_dim,
const int max_seq_len,
const int max_enc_len_this_time,
const int max_dec_len_this_time) {
const int q_token_num,
const int k_token_num) {
if (q_input.dtype() == paddle::DataType::FLOAT16) {
using T = phi::dtype::float16;
DispatchFlashAttentionMask<T>(q_input,
@@ -148,8 +147,8 @@ void FlashAttentionMask(const paddle::Tensor& q_input,
kv_head_num,
head_dim,
max_seq_len,
max_enc_len_this_time,
max_dec_len_this_time);
q_token_num,
k_token_num);
} else if (q_input.dtype() == paddle::DataType::BFLOAT16) {
using T = phi::dtype::bfloat16;
DispatchFlashAttentionMask<T>(q_input,
@@ -164,12 +163,12 @@ void FlashAttentionMask(const paddle::Tensor& q_input,
kv_head_num,
head_dim,
max_seq_len,
max_enc_len_this_time,
max_dec_len_this_time);
q_token_num,
k_token_num);
}
}
PD_BUILD_STATIC_OP(flash_attention_mask)
PD_BUILD_STATIC_OP(flash_mask_attention)
.Inputs({"q_input",
"k_input",
"v_input",
@@ -182,8 +181,8 @@ PD_BUILD_STATIC_OP(flash_attention_mask)
"kv_head_num: int",
"head_dim: int",
"max_seq_len: int",
"max_enc_len_this_time: int",
"max_dec_len_this_time: int"})
"q_token_num: int",
"k_token_num: int"})
.Outputs({"out"})
.SetInplaceMap({{"attn_out", "out"}})
.SetKernelFn(PD_KERNEL(FlashAttentionMask));

30
custom_ops/gpu_ops/flash_mask_attn/flash_mask_attn_kernel.hpp
View File

@@ -59,19 +59,26 @@ __global__ void __launch_bounds__(Ktraits::kNWarps *cutlass::NumThreadsPerWarp,
auto &shared_storage =
*reinterpret_cast<typename Ktraits::SharedStorage *>(shared_memory);
__align__(16) __shared__ int mask[kBlockM];
__align__(16) __shared__ int mask_end[kBlockM];
__align__(16) __shared__ int mask_start[kBlockM];
const int m_block = blockIdx.x;
const int bidh = blockIdx.y;
const int bidb = blockIdx.z;
if (data_params.seq_len_encoder[bidb] <= 0) {
return;
}
if constexpr (NeedMask) {
const int *mask_this_batch =
data_params.mask + data_params.cu_seq_q[bidb] + m_block * kBlockM;
const int2 *mask_this_batch =
reinterpret_cast<int2 *>(data_params.mask) +
(data_params.cu_seq_q[bidb] + m_block * kBlockM);
for (int i = threadIdx.x; i < kBlockM;
i += Ktraits::kNWarps * cutlass::NumThreadsPerWarp) {
mask[i] = mask_this_batch[i];
int2 mask_value = mask_this_batch[i];
mask_start[i] = mask_value.x;
mask_end[i] = mask_value.y;
}
}
@@ -119,7 +126,7 @@ __global__ void __launch_bounds__(Ktraits::kNWarps *cutlass::NumThreadsPerWarp,
const int n_block_max =
NeedMask
? cute::ceil_div(mask[min(kBlockM - 1, real_seq - 1)], kBlockN)
? cute::ceil_div(mask_end[min(kBlockM - 1, real_seq - 1)], kBlockN)
: min(cute::ceil_div((m_block + 1) * kBlockM + seq_len_k - seq_len_q,
kBlockN),
cute::ceil_div(seq_len_k, kBlockN));
@@ -170,7 +177,7 @@ __global__ void __launch_bounds__(Ktraits::kNWarps *cutlass::NumThreadsPerWarp,
smem_pipe_read_v,
tOrO,
softmax,
mask,
mask_end,
n_block_max,
threadIdx.x - NumCopyThreads,
m_block,
@@ -207,18 +214,15 @@ void run_flash_mask(Flash_mask_params &params, cudaStream_t stream) {
typename CollectiveMainloop::Params mainloop_params =
CollectiveMainloop::to_underlying_arguments(
{static_cast<Element const *>(params.q_ptr),
get_gmem_layout<kHeadDim>(params.max_seq_len_q * params.batch_size,
params.head_num),
get_gmem_layout<kHeadDim>(params.q_token_num, params.head_num),
static_cast<Element const *>(params.k_ptr),
get_gmem_layout<kHeadDim>(params.max_seq_len_k * params.batch_size,
params.kv_head_num),
get_gmem_layout<kHeadDim>(params.k_token_num, params.kv_head_num),
static_cast<Element const *>(params.v_ptr),
get_gmem_layout<kHeadDim>(params.max_seq_len_k * params.batch_size,
params.kv_head_num),
get_gmem_layout<kHeadDim>(params.k_token_num, params.kv_head_num),
params.scale_softmax_log2});
int num_blocks_m =
cutlass::ceil_div(params.max_seq_len_q, Kernel_traits::kBlockM);
cutlass::ceil_div(params.q_token_num, Kernel_traits::kBlockM);
num_blocks_m = cutlass::ceil_div(num_blocks_m, size<0>(ClusterShape{})) *
size<0>(ClusterShape{});

4
custom_ops/gpu_ops/flash_mask_attn/kernel_traits.h
View File

@@ -35,8 +35,8 @@ struct Flash_mask_params {
int *seq_len_encoder;
int head_num;
int kv_head_num;
int max_seq_len_q;
int max_seq_len_k;
int q_token_num;
int k_token_num;
int batch_size;
int gqa_group_size;
float scale_softmax_log2;