wehub/paddlepaddle--paddle
1
0
Fork 0
Code Issues Pull Requests Actions 2 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
paddlepaddle--paddle/paddle/phi/kernels/gpu/flash_attn_v3_kernel.cu
T
wehub-resource-sync e25996e7db chore: import upstream snapshot with attribution
2026-07-13 12:40:42 +08:00

2425 lines
101 KiB
Plaintext
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "paddle/phi/kernels/flash_attn_kernel.h"
#include <cstddef>
#include "glog/logging.h" // For VLOG()
#include "paddle/common/enforce.h"
#include "paddle/common/errors.h"
#include "paddle/common/flags.h"
#include "paddle/phi/common/data_type.h"
#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/core/platform/device_context.h"
#include "paddle/phi/core/tensor_utils.h"
#include "paddle/phi/core/utils/data_type.h"
#include "paddle/phi/kernels/empty_kernel.h"
#include "paddle/phi/kernels/funcs/elementwise_base.h"
#include "paddle/phi/kernels/slice_kernel.h"
#include "paddle/utils/none.h"
#include "paddle/phi/kernels/gpu/flash_attn_utils.h"
#include "paddle/phi/kernels/gpu/flash_attn_v3_utils.h"
#include "paddle/phi/kernels/gpu/flash_attn_v3_kernel.h"
namespace phi {
template <typename T, typename Context>
void FlashAttnV3BaseKernel(
const Context &dev_ctx,
const DenseTensor &q,
const DenseTensor &k,
const DenseTensor &v,
const optional<DenseTensor>
&k_new_, // (b, s_k_new, h_k, d) or (total_k_new, h_k, d) if there is
// cu_seqlens_k_new
const optional<DenseTensor>
&v_new_, // (b, s_k_new, h_k, dv) or (total_k_new, h_k, dv) if there is
// cu_seqlens_k_new
const optional<DenseTensor> &q_v_, // (b, s_q, h, dv) or (total_q_new, h,
// dv) if there is cu_seqlens_q
const optional<DenseTensor>
&out_, // (b, s_q, h, dv) or (total_q, h, dv) if there is cu_seqlens_q
const optional<DenseTensor> &cu_seqlens_q_, // b+1
const optional<DenseTensor> &cu_seqlens_k_, // b+1
const optional<DenseTensor> &cu_seqlens_k_new_, // b+1
const optional<DenseTensor>
&seqused_q_, // b. If given, only this many elements of each batch
// element's queries and outputs are used.
const optional<DenseTensor>
&seqused_k_, // b. If given, only this many elements of each batch
// element's keys are used.
const optional<DenseTensor> &page_table_, // (b_k, max_num_pages_per_seq)
const optional<DenseTensor>
&kv_batch_idx_, // b. indices to index into the KV cache
const optional<DenseTensor> &leftpad_k_, // b
const optional<DenseTensor> &rotary_cos_, // seqlen_ro x (rotary_dim / 2)
const optional<DenseTensor> &rotary_sin_, // seqlen_ro x (rotary_dim / 2)
const optional<DenseTensor> &q_descale_, // (b, h_k), not (b, h)
const optional<DenseTensor> &k_descale_, // (b, h_k)
const optional<DenseTensor> &v_descale_, // (b, h_k)
const optional<DenseTensor> &scheduler_metadata_, // (b + 1)
const int
max_seqlen_q_, // if max_seqlen_q_ is set to 0, it indicates that it is
// uninitialized and should not be referenced
// TODO(tridao): check if we need max_seqlen_k
const int
max_seqlen_k_, // if max_seqlen_q_ is set to 0, it indicates that it is
// uninitialized and should not be referenced
const float softmax_scale,
bool is_causal,
int window_size_left,
int window_size_right,
const float softcap,
const bool is_rotary_interleaved, // if true, rotary combines indices 0 &
// 1, else indices 0 & rotary_dim / 2
int num_splits,
const bool manual_set_pack_gqa,
const bool
pack_gqa_, // the pack_gqa_ will be used only if manual_set_pack_gqa is
// set to True; otherwise, the internal heuristic
// get_pack_gqa() from fa3 will decide whether to pack gqa
const int sm_margin,
DenseTensor *out,
DenseTensor *softmax_lse,
DenseTensor *out_accum,
DenseTensor *softmax_lse_accum) {
#ifdef PADDLE_WITH_FLASHATTN_V3
// TODO(umiswing): support ampere
int device_id = dev_ctx.GetPlace().GetDeviceId();
auto dprops = paddle::platform::GetDeviceProperties(device_id);
const bool is_sm90 = dprops.major == 9 && dprops.minor == 0;
PADDLE_ENFORCE_EQ(is_sm90,
true,
common::errors::Unavailable(
"FlashAttention-3 only supports Hopper GPUs."));
auto q_type = q.dtype();
PADDLE_ENFORCE_EQ(
(q_type == DataType::FLOAT16 || q_type == DataType::BFLOAT16 ||
q_type == DataType::FLOAT8_E4M3FN),
true,
common::errors::InvalidArgument(
"FlashAttention-3 only supports fp16, bf16, and fp8_e4m3 data type"));
PADDLE_ENFORCE_EQ(k.dtype(),
q_type,
common::errors::InvalidArgument(
"query and key must have the same dtype"));
PADDLE_ENFORCE_EQ(v.dtype(),
q_type,
common::errors::InvalidArgument(
"query and value must have the same dtype"));
CHECK_DEVICE(q);
CHECK_DEVICE(k);
CHECK_DEVICE(v);
PADDLE_ENFORCE_EQ(q.strides()[q.strides().size() - 1],
1,
common::errors::InvalidArgument(
"Input tensor must have contiguous last dimension"));
PADDLE_ENFORCE_EQ(k.strides()[k.strides().size() - 1],
1,
common::errors::InvalidArgument(
"Input tensor must have contiguous last dimension"));
PADDLE_ENFORCE_EQ(v.strides()[v.strides().size() - 1],
1,
common::errors::InvalidArgument(
"Input tensor must have contiguous last dimension"));
DenseTensor page_table;
// const bool paged_KV = page_table_.has_value();
// umiswing: this is stupid but idk how to use optional
const bool paged_KV = page_table_.is_initialized();
if (paged_KV) {
page_table = page_table_.get();
CHECK_DEVICE(page_table);
PADDLE_ENFORCE_EQ(page_table.dtype(),
DataType::INT32,
common::errors::InvalidArgument(
"page_table must have dtype paddle.int32"));
PADDLE_ENFORCE_EQ(page_table.strides()[page_table.strides().size() - 1],
1,
common::errors::InvalidArgument(
"page_table must have contiguous last dimension"));
}
// TODO(umiswing): support cusum
DenseTensor cu_seqlens_q;
// bool const is_varlen_q = cu_seqlens_q_.has_value();
// TODO(umiswing): this is stupid, must fix it (after understand
// optional)
const bool is_varlen_q = cu_seqlens_q_.is_initialized();
if (is_varlen_q) {
cu_seqlens_q = cu_seqlens_q_.get();
CHECK_DEVICE(cu_seqlens_q);
CHECK_CONTIGUOUS(cu_seqlens_q);
PADDLE_ENFORCE_EQ(cu_seqlens_q.dtype(),
DataType::INT32,
common::errors::InvalidArgument(
"cu_seqlens_q must have dtype paddle.int32"));
PADDLE_ENFORCE_NE(
max_seqlen_q_,
0,
common::errors::InvalidArgument(
"max_seqlen_q must be provided if cu_seqlens_q is provided"));
}
DenseTensor cu_seqlens_k;
const bool is_varlen_k = cu_seqlens_k_.is_initialized();
if (is_varlen_k) {
cu_seqlens_k = cu_seqlens_k_.get();
CHECK_DEVICE(cu_seqlens_k);
CHECK_CONTIGUOUS(cu_seqlens_k);
PADDLE_ENFORCE_EQ(cu_seqlens_k.dtype(),
DataType::INT32,
common::errors::InvalidArgument(
"cu_seqlens_k must have dtype paddle.int32"));
PADDLE_ENFORCE_NE(
max_seqlen_k_,
0,
common::errors::InvalidArgument(
"max_seqlen_k must be provided if cu_seqlens_k is provided"));
PADDLE_ENFORCE_EQ(
!paged_KV,
true,
common::errors::InvalidArgument(
"If cu_seqlens_k is passed in, then page table is not supported"));
PADDLE_ENFORCE_EQ(
!kv_batch_idx_,
true,
common::errors::InvalidArgument(
"If cu_seqlens_k is passed in, then page table is not supported"));
}
auto const sizes = q.dims();
const int batch_size = !is_varlen_q ? sizes[0] : cu_seqlens_q.dims()[0] - 1;
int seqlen_q = !is_varlen_q ? sizes[1] : max_seqlen_q_;
int total_q = !is_varlen_q ? batch_size * sizes[1] : sizes[0];
int64_t num_heads = q.dims()[q.dims().size() - 2];
int64_t const head_size = q.dims()[q.dims().size() - 1];
int const head_size_v = v.dims()[v.dims().size() - 1];
int const max_num_pages_per_seq = !paged_KV ? 0 : page_table.dims()[1];
int const num_pages = !paged_KV ? 0 : k.dims()[0];
int const page_size = !paged_KV ? 1 : k.dims()[1];
int const seqlen_k =
!is_varlen_k
? (!paged_KV ? k.dims()[1] : max_num_pages_per_seq * page_size)
: max_seqlen_k_;
int const total_k = !is_varlen_k ? batch_size * k.dims()[1] : k.dims()[0];
int const num_heads_k = k.dims()[k.dims().size() - 2];
int const batch_size_k =
!paged_KV ? (!is_varlen_k ? k.dims()[0] : cu_seqlens_k.dims()[0] - 1)
: page_table.dims()[0];
if (!kv_batch_idx_.is_initialized()) {
PADDLE_ENFORCE_EQ(batch_size,
batch_size_k,
common::errors::InvalidArgument(
"batch_size must be equal to batch_size_k"));
}
int const max_headdim = get_max_headdim();
PADDLE_ENFORCE_LE(
head_size,
max_headdim,
common::errors::InvalidArgument(
"FlashAttention forward only supports head dimension at most %d",
max_headdim));
PADDLE_ENFORCE_EQ(
num_heads % num_heads_k,
0,
common::errors::InvalidArgument(
"Number of heads in key/value must divide number of heads in query"));
if (head_size_v != head_size) {
PADDLE_ENFORCE_EQ(
((head_size > 128 && head_size <= 192 && head_size_v > 96 &&
head_size_v <= 128) ||
(head_size <= 64 && head_size_v <= 512)),
true,
common::errors::InvalidArgument(
"If V headdim is different from Q/K dim, we only support "
"Q/K headdim in (128, 192] and V headdim in (96, 128], "
"or (Q/K <= 64 and V <= 512)."));
PADDLE_ENFORCE_EQ(dprops.major,
9,
common::errors::InvalidArgument(
"Only Hopper supports different V headdim"));
if (head_size_v > 256) {
PADDLE_ENFORCE_EQ(
(q_type == DataType::FLOAT16 || q_type == DataType::BFLOAT16),
true,
common::errors::InvalidArgument(
"HeaddimV > 256 requires fp16 and bf16 data type"));
}
}
// This needs to go before kBlockM & kBlockN since we rely on the correct
// window_size and is_causal to set kBlockM
// TODO(tridao): check this
if (window_size_left >= seqlen_k - 1) {
window_size_left = -1;
}
if (window_size_right >= seqlen_q - 1) {
window_size_right = -1;
}
// causal=true is the same as causal=false in this case
if (seqlen_q == 1 && window_size_left == -1 && window_size_right == -1) {
// Special case of hdim 128 where we want causal to have kBlockN=128, better
// for pagedKV and TMA
if ((head_size <= 64 || head_size > 128) || !paged_KV) {
is_causal = false;
}
}
if (is_causal) {
window_size_right = 0;
}
// There's a case where is_causal=false, window_size=(-1, 0). Then
// set_params_fprop will set params.is_causal=true. If we don't have is_causal
// here matching params.is_causal, we might get the wrong kBlockM.
is_causal = window_size_left < 0 && window_size_right == 0;
if (!is_varlen_q) {
CHECK_SHAPE(q, batch_size, seqlen_q, num_heads, head_size);
} else {
CHECK_SHAPE(q, total_q, num_heads, head_size);
CHECK_SHAPE(cu_seqlens_q, batch_size + 1);
}
if (!paged_KV) {
if (!is_varlen_k) {
CHECK_SHAPE(k, batch_size_k, seqlen_k, num_heads_k, head_size);
CHECK_SHAPE(v, batch_size_k, seqlen_k, num_heads_k, head_size_v);
} else {
CHECK_SHAPE(k, total_k, num_heads_k, head_size);
CHECK_SHAPE(v, total_k, num_heads_k, head_size_v);
CHECK_SHAPE(cu_seqlens_k, batch_size + 1);
}
} else {
CHECK_SHAPE(k, num_pages, page_size, num_heads_k, head_size);
CHECK_SHAPE(v, num_pages, page_size, num_heads_k, head_size_v);
CHECK_SHAPE(page_table, batch_size_k, max_num_pages_per_seq);
}
if (seqused_q_.is_initialized()) {
auto seqused_q = seqused_q_.get();
PADDLE_ENFORCE_EQ(
seqused_q.dtype(),
DataType::INT32,
common::errors::InvalidArgument("seqused_q must have dtype int32"));
CHECK_DEVICE(seqused_q);
CHECK_CONTIGUOUS(seqused_q);
CHECK_SHAPE(seqused_q, batch_size);
}
if (seqused_k_.is_initialized()) {
auto seqused_k = seqused_k_.get();
PADDLE_ENFORCE_EQ(
seqused_k.dtype(),
DataType::INT32,
common::errors::InvalidArgument("seqused_k must have dtype int32"));
CHECK_DEVICE(seqused_k);
CHECK_CONTIGUOUS(seqused_k);
CHECK_SHAPE(seqused_k, batch_size);
}
if (leftpad_k_.is_initialized()) {
auto leftpad_k = leftpad_k_.get();
PADDLE_ENFORCE_EQ(
leftpad_k.dtype(),
DataType::INT32,
common::errors::InvalidArgument("leftpad_k must have dtype int32"));
CHECK_DEVICE(leftpad_k);
CHECK_CONTIGUOUS(leftpad_k);
CHECK_SHAPE(leftpad_k, batch_size);
}
// This is what we will template on
bool const is_varlen =
is_varlen_q || is_varlen_k || seqused_q_.is_initialized() ||
seqused_k_.is_initialized() || leftpad_k_.is_initialized();
#ifdef FLASHATTENTION_DISABLE_VARLEN
PADDLE_ENFORCE_EQ(!is_varlen,
true,
common::errors::Unavailable(
"This flash attention build does not support varlen."));
#endif
int const alignment = q_type == DataType::FLOAT8_E4M3FN ? 16 : 8;
PADDLE_ENFORCE_EQ(head_size % alignment,
0,
common::errors::InvalidArgument(
"head_size should be a multiple of %d", alignment));
PADDLE_ENFORCE_EQ(head_size_v % alignment,
0,
common::errors::InvalidArgument(
"head_size_v should be a multiple of %d", alignment));
auto out_type =
q_type == DataType::FLOAT8_E4M3FN ? DataType::BFLOAT16 : q_type;
if (out_.is_initialized()) {
*out = out_.get();
PADDLE_ENFORCE_EQ(
out->dtype(),
out_type,
common::errors::InvalidArgument(
"For FP16/BF16 input, output must have the same dtype as "
"inputs. For FP8 input, output must have dtype BF16"));
CHECK_DEVICE((*out));
PADDLE_ENFORCE_EQ(out->strides()[out->strides().size() - 1],
1,
common::errors::InvalidArgument(
"Output tensor must have contiguous last dimension"));
if (!is_varlen_q) {
CHECK_SHAPE((*out), batch_size, seqlen_q, num_heads, head_size_v);
} else {
CHECK_SHAPE((*out), total_q, num_heads, head_size_v);
}
} else {
if (!is_varlen_q) {
out->Resize({batch_size, seqlen_q, num_heads, head_size_v});
} else {
out->Resize({total_q, num_heads, head_size_v});
}
if (q_type == DataType::FLOAT8_E4M3FN) {
dev_ctx.template Alloc<phi::bfloat16>(out);
} else {
// umiswing: assuming T is Input Type
dev_ctx.template Alloc<T>(out);
}
}
auto round_multiple = [](int x, int m) { return (x + m - 1) / m * m; };
int const head_size_rounded = round_up_headdim(head_size);
int const head_size_v_rounded = round_up_headdim(head_size_v);
int const seqlen_q_rounded = round_multiple(seqlen_q, 128);
int const seqlen_k_rounded = round_multiple(seqlen_k, 128);
if (!is_varlen_q) {
softmax_lse->Resize({batch_size, num_heads, seqlen_q});
} else {
softmax_lse->Resize({num_heads, total_q});
}
dev_ctx.template Alloc<float>(softmax_lse);
Flash_fwd_params *params_handle = get_flash_fwd_params_handle();
dynload::fa3_clear_fwd_params_handle(params_handle);
set_params_fprop(
params_handle,
batch_size,
seqlen_q,
seqlen_k,
seqlen_q_rounded,
seqlen_k_rounded,
num_heads,
num_heads_k,
head_size,
head_size_rounded,
q,
k,
v,
out,
!is_varlen_q ? nullptr : cu_seqlens_q.data(),
!is_varlen_k ? nullptr : cu_seqlens_k.data(),
seqused_q_.is_initialized() ? const_cast<void *>(seqused_q_.get().data())
: nullptr,
seqused_k_.is_initialized() ? const_cast<void *>(seqused_k_.get().data())
: nullptr,
softmax_lse->data(),
/*p_dropout=*/0.f,
softmax_scale,
window_size_left,
window_size_right,
dprops,
softcap,
sm_margin);
dynload::fa3_fwd_params_set_total_q(params_handle, total_q);
dynload::fa3_fwd_params_set_total_k(params_handle, total_k);
dynload::fa3_fwd_params_set_b_k(params_handle, batch_size_k);
dynload::fa3_fwd_params_set_dv(params_handle, head_size_v);
dynload::fa3_fwd_params_set_dv_rounded(params_handle, head_size_v_rounded);
if (leftpad_k_
.is_initialized()) { // This needs to be set before get_pagedkv_tma
dynload::fa3_fwd_params_set_leftpad_k(params_handle,
leftpad_k_.get().data<int>());
}
if (paged_KV) {
dynload::fa3_fwd_params_set_page_table(params_handle,
page_table.data<int>());
dynload::fa3_fwd_params_set_page_table_batch_stride(
params_handle, page_table.strides()[0]);
}
dynload::fa3_fwd_params_set_page_size(params_handle, page_size);
dynload::fa3_fwd_params_set_num_pages(params_handle, num_pages);
if (k_new_.is_initialized()) { // This needs to be set before get_pagedkv_tma
DenseTensor k_new, v_new;
PADDLE_ENFORCE_EQ(
v_new_.is_initialized(),
true,
common::errors::InvalidArgument(
"If k_new is supplied, v_new must also be passed in"));
PADDLE_ENFORCE_EQ(
seqused_k_.is_initialized(),
true,
common::errors::InvalidArgument(
"If k_new is supplied, seqlens_k must also be passed in"));
PADDLE_ENFORCE_LE(
seqlen_q,
seqlen_k,
common::errors::InvalidArgument(
"If k_new is supplied, it must have seqlen <= the seqlen "
"of the KV cache"));
DenseTensor cu_seqlens_k_new;
bool const is_varlen_k_new = cu_seqlens_k_new_.is_initialized();
if (is_varlen_k_new) {
cu_seqlens_k_new = cu_seqlens_k_new_.get();
CHECK_DEVICE(cu_seqlens_k_new);
CHECK_CONTIGUOUS(cu_seqlens_k_new);
PADDLE_ENFORCE_EQ(cu_seqlens_k_new.dtype(),
DataType::INT32,
common::errors::InvalidArgument(
"cu_seqlens_k_new must have dtype paddle.int32"));
}
k_new = k_new_.get();
v_new = v_new_.get();
PADDLE_ENFORCE_EQ(k_new.dtype(),
q_type,
common::errors::InvalidArgument(
"k_new must have the same dtype as query"));
PADDLE_ENFORCE_EQ(v_new.dtype(),
q_type,
common::errors::InvalidArgument(
"v_new must have the same dtype as query"));
CHECK_DEVICE(k_new);
CHECK_DEVICE(v_new);
PADDLE_ENFORCE_EQ(k_new.strides()[k_new.strides().size() - 1],
1,
common::errors::InvalidArgument(
"k_new tensor must have contiguous last dimension"));
PADDLE_ENFORCE_EQ(v_new.strides()[v_new.strides().size() - 1],
1,
common::errors::InvalidArgument(
"v_new tensor must have contiguous last dimension"));
// We don't need max_seqlen_k_new, so seqlen_k_new can be whatever when
// is_varlen_k_new
int seqlen_k_new = !is_varlen_k_new ? k_new.dims()[1] : 0;
int total_k_new =
!is_varlen_k_new ? batch_size * k_new.dims()[1] : k_new.dims()[0];
if (!is_varlen_k_new) {
CHECK_SHAPE(k_new, batch_size, seqlen_k_new, num_heads_k, head_size);
CHECK_SHAPE(v_new, batch_size, seqlen_k_new, num_heads_k, head_size_v);
} else {
CHECK_SHAPE(k_new, total_k_new, num_heads_k, head_size);
CHECK_SHAPE(v_new, total_k_new, num_heads_k, head_size_v);
CHECK_SHAPE(cu_seqlens_k_new, batch_size + 1);
}
// umiswing: dump this to shared library
dynload::fa3_fwd_params_set_seqlen_knew(params_handle, seqlen_k_new);
dynload::fa3_fwd_params_set_total_knew(params_handle, total_k_new);
dynload::fa3_fwd_params_set_knew_ptr(params_handle,
const_cast<void *>(k_new.data()));
dynload::fa3_fwd_params_set_vnew_ptr(params_handle,
const_cast<void *>(v_new.data()));
// All stride are in elements, not bytes.
dynload::fa3_fwd_params_set_knew_row_stride(
params_handle, k_new.strides()[k_new.strides().size() - 3]);
dynload::fa3_fwd_params_set_vnew_row_stride(
params_handle, v_new.strides()[v_new.strides().size() - 3]);
dynload::fa3_fwd_params_set_knew_head_stride(
params_handle, k_new.strides()[k_new.strides().size() - 2]);
dynload::fa3_fwd_params_set_vnew_head_stride(
params_handle, v_new.strides()[v_new.strides().size() - 2]);
if (!is_varlen_k_new) {
dynload::fa3_fwd_params_set_knew_batch_stride(params_handle,
k_new.strides()[0]);
dynload::fa3_fwd_params_set_vnew_batch_stride(params_handle,
v_new.strides()[0]);
}
if (is_varlen_k_new) {
dynload::fa3_fwd_params_set_cu_seqlens_knew(params_handle,
cu_seqlens_k_new.data<int>());
}
}
// 992 = 32 * 31 is the max supported batch in prepare_varlen_num_blocks
// kernel
bool const use_dynamic_split =
is_varlen && dynload::fa3_fwd_params_get_b(params_handle) <= 992;
// Temporarily set num_splits_dynamic_ptr to 1 since get_num_splits checks it
dynload::fa3_fwd_params_set_num_splits_dynamic_ptr(
params_handle, !use_dynamic_split ? nullptr : reinterpret_cast<int *>(1));
dynload::fa3_fwd_params_set_pagedkv_tma(
params_handle, dynload::fa3_get_pagedkv_tma(params_handle));
if (num_splits <= 0) {
num_splits = dynload::fa3_get_num_splits(params_handle);
}
dynload::fa3_fwd_params_set_num_splits(params_handle, num_splits);
// Always enable PackGQA for Split, and get_pack_gqa requires
// params.num_splits to decide
const bool pack_gqa = manual_set_pack_gqa
? pack_gqa_
: dynload::fa3_get_pack_gqa(params_handle);
dynload::fa3_fwd_params_set_pack_gqa(params_handle, pack_gqa);
// This needs to be set after get_num_splits
DenseTensor tile_count_semaphore; // Contains the semaphore and optionally
// num_splits_dynamic
// We don't use the persistent scheduler if Split and not Varlen
const bool params_is_causal =
dynload::fa3_fwd_params_get_is_causal(params_handle);
const bool params_is_local =
dynload::fa3_fwd_params_get_is_local(params_handle);
const int params_num_splits =
dynload::fa3_fwd_params_get_num_splits(params_handle);
const int params_b = dynload::fa3_fwd_params_get_b(params_handle);
const int params_arch = dynload::fa3_fwd_params_get_arch(params_handle);
bool const scheduler_needs_semaphore =
params_arch >= 90 ? (((params_is_causal || params_is_local) &&
(params_num_splits == 1)) ||
is_varlen)
: ((params_is_causal && !is_varlen) ||
(is_varlen && params_num_splits > 1));
if (scheduler_needs_semaphore || use_dynamic_split) {
int metadata_size = static_cast<int>(scheduler_needs_semaphore) +
static_cast<int>(use_dynamic_split) * params_b;
dynload::fa3_fwd_params_set_skip_scheduler_metadata_computation(
params_handle, scheduler_metadata_.is_initialized());
if (scheduler_metadata_.is_initialized()) {
DenseTensor scheduler_metadata = scheduler_metadata_.get();
CHECK_DEVICE(scheduler_metadata);
CHECK_SHAPE(scheduler_metadata, metadata_size);
CHECK_CONTIGUOUS(scheduler_metadata);
PADDLE_ENFORCE_EQ(scheduler_metadata.dtype(),
DataType::INT32,
common::errors::InvalidArgument(
"scheduler_metadata must have dtype int32"));
tile_count_semaphore = scheduler_metadata;
} else {
tile_count_semaphore = Empty<int32_t>(dev_ctx, {metadata_size});
}
if (scheduler_needs_semaphore && !use_dynamic_split) {
funcs::SetConstant<Context, int32_t> set_zero;
set_zero(dev_ctx,
&tile_count_semaphore,
int32_t{0}); // If varlen we'll manually do the zero-ing
}
dynload::fa3_fwd_params_set_tile_count_semaphore(
params_handle,
scheduler_needs_semaphore
? const_cast<int *>(tile_count_semaphore.data<int>())
: nullptr);
dynload::fa3_fwd_params_set_num_splits_dynamic_ptr(
params_handle,
use_dynamic_split
? const_cast<int *>(tile_count_semaphore.data<int>()) + 1
: nullptr);
}
if (q_v_.is_initialized()) {
PADDLE_ENFORCE_LT(head_size,
64,
common::errors::InvalidArgument(
"q_v is only supported for head_size <= 64"));
PADDLE_ENFORCE_EQ(
(q_type == DataType::FLOAT16 || q_type == DataType::FLOAT16),
true,
common::errors::InvalidArgument(
"q_v is only supported for fp16 and bf16 data type"));
PADDLE_ENFORCE_EQ(params_arch,
90,
common::errors::InvalidArgument(
"q_v is only supported for Hopper GPUs"));
DenseTensor q_v = q_v_.get();
PADDLE_ENFORCE_EQ(q_v.dtype(),
q_type,
common::errors::InvalidArgument(
"q_v must have the same dtype as query"));
CHECK_DEVICE(q_v);
PADDLE_ENFORCE_EQ(q_v.strides()[q_v.strides().size() - 1],
1,
common::errors::InvalidArgument(
"q_v tensor must have contiguous last dimension"));
if (!is_varlen_q) {
CHECK_SHAPE(q_v, batch_size, seqlen_q, num_heads, head_size_v);
} else {
CHECK_SHAPE(q_v, total_q, num_heads, head_size_v);
}
dynload::fa3_fwd_params_set_qv_ptr(params_handle,
const_cast<void *>(q_v.data()));
// All stride are in elements, not bytes.
dynload::fa3_fwd_params_set_qv_row_stride(
params_handle, q_v.strides()[q_v.strides().size() - 3]);
dynload::fa3_fwd_params_set_qv_head_stride(
params_handle, q_v.strides()[q_v.strides().size() - 2]);
if (!is_varlen_q) {
dynload::fa3_fwd_params_set_qv_batch_stride(params_handle,
q_v.strides()[0]);
}
}
if (rotary_cos_.is_initialized()) {
PADDLE_ENFORCE_EQ(
k_new_.is_initialized(),
true,
common::errors::InvalidArgument(
"If rotary cos/sin are provided, new key / value to be "
"appended to KV cache must also be provided"));
DenseTensor rotary_cos = rotary_cos_.get();
CHECK_DEVICE(rotary_cos);
CHECK_CONTIGUOUS(rotary_cos);
int params_rotary_dim = rotary_cos.dims()[1] * 2;
dynload::fa3_fwd_params_set_rotary_dim(params_handle, params_rotary_dim);
PADDLE_ENFORCE_LE(
params_rotary_dim,
head_size,
common::errors::InvalidArgument("rotary_dim must be <= headdim"));
PADDLE_ENFORCE_EQ(
params_rotary_dim % 16,
0,
common::errors::InvalidArgument(
"Only rotary dimensions divisible by 16 are currently supported"));
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t seqlen_ro = rotary_cos.dims()[0];
if (paged_KV) {
PADDLE_ENFORCE_GE(
seqlen_ro,
seqlen_k,
common::errors::InvalidArgument(
"cos/sin seqlen must be at least the seqlen of KV cache"));
}
CHECK_SHAPE(rotary_cos, seqlen_ro, params_rotary_dim / 2);
PADDLE_ENFORCE_EQ(rotary_cos.dtype(),
q_type,
common::errors::InvalidArgument(
"rotary_cos must have the same dtype as query"));
PADDLE_ENFORCE_EQ(
rotary_sin_.is_initialized(),
true,
common::errors::InvalidArgument(
"If rotary cos is provided, rotary sin must also be provided"));
auto rotary_sin = rotary_sin_.get();
CHECK_DEVICE(rotary_sin);
CHECK_CONTIGUOUS(rotary_sin);
CHECK_SHAPE(rotary_sin, seqlen_ro, params_rotary_dim / 2);
PADDLE_ENFORCE_EQ(rotary_sin.dtype(),
q_type,
common::errors::InvalidArgument(
"rotary_cos must have the same dtype as query"));
dynload::fa3_fwd_params_set_rotary_cos_ptr(
params_handle, const_cast<void *>(rotary_cos.data()));
dynload::fa3_fwd_params_set_rotary_sin_ptr(
params_handle, const_cast<void *>(rotary_sin.data()));
dynload::fa3_fwd_params_set_is_rotary_interleaved(params_handle,
is_rotary_interleaved);
} else {
dynload::fa3_fwd_params_set_rotary_dim(params_handle, 0);
}
if (kv_batch_idx_.is_initialized()) {
DenseTensor kv_batch_idx = kv_batch_idx_.get();
CHECK_DEVICE(kv_batch_idx);
CHECK_CONTIGUOUS(kv_batch_idx);
PADDLE_ENFORCE_EQ(
kv_batch_idx.dtype(),
DataType::INT32,
common::errors::InvalidArgument("kv_batch_idx must have dtype int32"));
dynload::fa3_fwd_params_set_kv_batch_idx(
params_handle, reinterpret_cast<int *>(kv_batch_idx.data()));
}
if (dynload::fa3_fwd_params_get_num_splits(params_handle) > 1) {
PADDLE_ENFORCE_LE(
dynload::fa3_fwd_params_get_num_splits(params_handle),
256,
common::errors::InvalidArgument("num_splits > 256 not supported"));
if (!is_varlen_q) {
out_accum->Resize(
make_ddim({dynload::fa3_fwd_params_get_num_splits(params_handle),
batch_size,
num_heads,
seqlen_q,
head_size_v}));
softmax_lse_accum->Resize(
make_ddim({dynload::fa3_fwd_params_get_num_splits(params_handle),
batch_size,
num_heads,
seqlen_q}));
dev_ctx.template Alloc<float>(out_accum);
dev_ctx.template Alloc<float>(softmax_lse_accum);
dynload::fa3_fwd_params_set_oaccum_batch_stride(params_handle,
out_accum->strides()[1]);
dynload::fa3_fwd_params_set_lseaccum_batch_stride(
params_handle, softmax_lse_accum->strides()[1]);
} else {
out_accum->Resize(
make_ddim({dynload::fa3_fwd_params_get_num_splits(params_handle),
num_heads,
total_q,
head_size_v}));
softmax_lse_accum->Resize(
make_ddim({dynload::fa3_fwd_params_get_num_splits(params_handle),
num_heads,
total_q}));
dev_ctx.template Alloc<float>(out_accum);
dev_ctx.template Alloc<float>(softmax_lse_accum);
}
dynload::fa3_fwd_params_set_is_fp32(params_handle, false);
dynload::fa3_fwd_params_set_oaccum_ptr(
params_handle, const_cast<void *>(out_accum->data()));
dynload::fa3_fwd_params_set_softmax_lseaccum_ptr(
params_handle, const_cast<void *>(softmax_lse_accum->data()));
dynload::fa3_fwd_params_set_oaccum_split_stride(params_handle,
out_accum->strides()[0]);
dynload::fa3_fwd_params_set_oaccum_row_stride(
params_handle, out_accum->strides()[out_accum->strides().size() - 2]);
dynload::fa3_fwd_params_set_oaccum_head_stride(
params_handle, out_accum->strides()[out_accum->strides().size() - 3]);
dynload::fa3_fwd_params_set_lseaccum_split_stride(
params_handle, softmax_lse_accum->strides()[0]);
dynload::fa3_fwd_params_set_lseaccum_head_stride(
params_handle,
softmax_lse_accum->strides()[softmax_lse_accum->strides().size() - 2]);
}
if (q_type == DataType::FLOAT8_E4M3FN) {
if (q_descale_.is_initialized()) {
DenseTensor q_descale = q_descale_.get();
CHECK_DEVICE(q_descale);
CHECK_SHAPE(q_descale, batch_size, num_heads_k);
dynload::fa3_fwd_params_set_q_descale_ptr(
params_handle, const_cast<float *>(q_descale.data<float>()));
dynload::fa3_fwd_params_set_q_descale_batch_stride(
params_handle, q_descale.strides()[0]);
dynload::fa3_fwd_params_set_q_descale_head_stride(params_handle,
q_descale.strides()[1]);
} else {
dynload::fa3_fwd_params_set_q_descale_ptr(params_handle, nullptr);
}
if (k_descale_.is_initialized()) {
DenseTensor k_descale = k_descale_.get();
CHECK_DEVICE(k_descale);
CHECK_SHAPE(k_descale, batch_size, num_heads_k);
dynload::fa3_fwd_params_set_k_descale_ptr(
params_handle, const_cast<float *>(k_descale.data<float>()));
dynload::fa3_fwd_params_set_k_descale_batch_stride(
params_handle, k_descale.strides()[0]);
dynload::fa3_fwd_params_set_k_descale_head_stride(params_handle,
k_descale.strides()[1]);
} else {
dynload::fa3_fwd_params_set_k_descale_ptr(params_handle, nullptr);
}
if (v_descale_.is_initialized()) {
DenseTensor v_descale = v_descale_.get();
CHECK_DEVICE(v_descale);
CHECK_SHAPE(v_descale, batch_size, num_heads_k);
dynload::fa3_fwd_params_set_v_descale_ptr(
params_handle, const_cast<float *>(v_descale.data<float>()));
dynload::fa3_fwd_params_set_v_descale_batch_stride(
params_handle, v_descale.strides()[0]);
dynload::fa3_fwd_params_set_v_descale_head_stride(params_handle,
v_descale.strides()[1]);
} else {
dynload::fa3_fwd_params_set_v_descale_ptr(params_handle, nullptr);
}
}
#ifdef FLASHATTENTION_DISABLE_LOCAL
PADDLE_ENFORCE_EQ(
!dynload::fa3_fwd_params_get_is_local(params_handle),
true,
common::errors::InvalidArgument(
"This flash attention build does not support local attention."));
#endif
#ifdef FLASHATTENTION_DISABLE_SOFTCAP
PADDLE_ENFORCE_EQ(
dynload::fa3_fwd_params_get_softcap(params_handle),
0.0,
common::errors::InvalidArgument(
"This flash attention build does not support tanh softcapping."));
#endif
#ifdef FLASHATTENTION_DISABLE_SPLIT
PADDLE_ENFORCE_EQ(dynload::fa3_fwd_params_get_num_splits(params_handle),
1,
common::errors::InvalidArgument(
"This flash attention build does not support splits."));
#endif
#ifdef FLASHATTENTION_DISABLE_PACKGQA
PADDLE_ENFORCE_EQ(
(!dynload::fa3_fwd_params_get_pack_gqa(params_handle) ||
dynload::fa3_fwd_params_get_arch(params_handle) < 90 ||
(dynload::fa3_fwd_params_get_page_table(params_handle) &&
!dynload::fa3_fwd_params_get_pagedkv_tma(params_handle)) ||
dynload::fa3_fwd_params_get_num_splits(params_handle) > 1),
true,
common::errors::InvalidArgument(
"This flash attention build does not support pack_gqa."));
#endif
#ifdef FLASHATTENTION_DISABLE_PAGEDKV
PADDLE_ENFORCE_EQ(
(!(dynload::fa3_fwd_params_get_page_table(params_handle) &&
!dynload::fa3_fwd_params_get_pagedkv_tma(params_handle))),
true,
common::errors::InvalidArgument(
"This flash attention build does not support paged KV."));
#endif
#ifdef FLASHATTENTION_DISABLE_APPENDKV
PADDLE_ENFORCE_EQ(
!k_new_.is_initialized(),
true,
common::errors::InvalidArgument(
"This flash attention build does not support appending KV."));
#endif
if (total_q > 0 &&
(total_k + dynload::fa3_fwd_params_get_total_knew(params_handle)) > 0 &&
num_heads_k > 0) {
dynload::fa3_run_mha_fwd(params_handle, dev_ctx.stream());
if (dynload::fa3_fwd_params_get_num_splits(params_handle) > 1) {
if (out_type == DataType::BFLOAT16) {
// Since we want output in BF16. Otherwise fwd_combine will output to
// FP16
dynload::fa3_fwd_params_set_is_bf16(params_handle, true);
}
// Unless there's seqused_q, for the purpose of attn_combine, we can just
// treat it as batch=1 and seqlen = total_q, and don't need to dispatch to
// Varlen there. However, with dynamic split, each row needs to know which
// batch it belongs to to read the number of splits, so we just use the
// varlen version of combine kernel. if (is_varlen_q &&
// !seqused_q_.has_value()) { if (is_varlen_q) {
// params.b = 1;
// params.seqlen_q = total_q;
// }
// }
dynload::fa3_run_mha_fwd_combine(
params_handle, dev_ctx.stream(), true /*enable_pdl*/);
}
} else if (total_q > 0 && num_heads_k > 0) {
PADDLE_ENFORCE_EQ(
(out->dtype() == DataType::BFLOAT16 ||
out->dtype() == DataType::FLOAT16 ||
out->dtype() == DataType::FLOAT8_E4M3FN),
true,
common::errors::InvalidArgument("flash attention 3 supports bfloat16, "
"float16 and float8_e4m3fn only."));
// If seqlen_k == 0, then we have an empty tensor. We need to set the output
// to 0.
if (out->dtype() == DataType::BFLOAT16) {
funcs::SetConstant<Context, phi::bfloat16> set_zero;
set_zero(dev_ctx,
out,
phi::bfloat16{0}); // If varlen we'll manually do the zero-ing
} else if (out->dtype() == DataType::FLOAT16) {
funcs::SetConstant<Context, phi::float16> set_zero;
set_zero(dev_ctx,
out,
phi::float16{0}); // If varlen we'll manually do the zero-ing
} else if (out->dtype() == DataType::FLOAT8_E4M3FN) {
funcs::SetConstant<Context, phi::float8_e4m3fn> set_zero;
set_zero(
dev_ctx,
out,
phi::float8_e4m3fn{0}); // If varlen we'll manually do the zero-ing
}
funcs::SetConstant<Context, float> set_infinity;
set_infinity(dev_ctx, softmax_lse, std::numeric_limits<float>::infinity());
}
#else
RaiseNotSupportedError();
#endif
}
template <typename T, typename Context>
void FlashAttnV3Kernel(const Context &dev_ctx,
const DenseTensor &q,
const DenseTensor &k,
const DenseTensor &v,
const optional<DenseTensor> &q_v_,
const optional<DenseTensor> &q_descale_,
const optional<DenseTensor> &k_descale_,
const optional<DenseTensor> &v_descale_,
const float softmax_scale,
bool is_causal,
int window_size_left,
int window_size_right,
const float softcap,
int num_splits,
const bool manual_set_pack_gqa,
const bool pack_gqa_,
const int sm_margin,
DenseTensor *out,
DenseTensor *softmax_lse) {
#ifdef PADDLE_WITH_FLASHATTN_V3
// umiswing: the following options have not been fully tested
PADDLE_ENFORCE_EQ(q_v_.is_initialized(),
false,
common::errors::InvalidArgument("q_v_ is not supported"));
PADDLE_ENFORCE_EQ(
q_descale_.is_initialized(),
false,
common::errors::InvalidArgument("q_descale_ is not supported"));
PADDLE_ENFORCE_EQ(
k_descale_.is_initialized(),
false,
common::errors::InvalidArgument("k_descale_ is not supported"));
PADDLE_ENFORCE_EQ(
v_descale_.is_initialized(),
false,
common::errors::InvalidArgument("v_descale_ is not supported"));
PADDLE_ENFORCE_EQ(
window_size_left,
-1,
common::errors::InvalidArgument("window_size is not supported, please "
"set window_size_left/right to -1"));
PADDLE_ENFORCE_EQ(
window_size_right,
-1,
common::errors::InvalidArgument("window_size is not supported, please "
"set window_size_left/right to -1"));
PADDLE_ENFORCE_EQ(softcap,
0,
common::errors::InvalidArgument(
"softcap is not supported, please set softcap to 0"));
PADDLE_ENFORCE_EQ(
num_splits,
1,
common::errors::InvalidArgument(
"num_splits is not supported, please set num_splits to 1"));
PADDLE_ENFORCE_EQ(manual_set_pack_gqa,
false,
common::errors::InvalidArgument(
"manual_set_pack_gqa is not supported, please set "
"manual_set_pack_gqa to false"));
PADDLE_ENFORCE_EQ(
pack_gqa_,
false,
common::errors::InvalidArgument(
"pack_gqa_ is not supported, please set pack_gqa_ to false"));
PADDLE_ENFORCE_EQ(
sm_margin,
0,
common::errors::InvalidArgument(
"sm_margin is not supported, please set sm_margin to 0"));
DenseTensor out_accum;
DenseTensor softmax_lse_accum;
FlashAttnV3BaseKernel<T, Context>(dev_ctx,
q,
k,
v,
paddle::none, // k_new_
paddle::none, // v_new_
q_v_,
paddle::none, // out_
paddle::none, // cu_seqlens_q_
paddle::none, // cu_seqlens_k_
paddle::none, // cu_seqlens_k_new_
paddle::none, // seqused_q_
paddle::none, // seqused_k_
paddle::none, // page_table_
paddle::none, // kv_batch_idx_
paddle::none, // leftpad_k_
paddle::none, // rotary_cos_
paddle::none, // rotary_sin_
q_descale_,
k_descale_,
v_descale_,
paddle::none, // scheduler_metadata
0, // max_seqlen_q_
0, // max_seqlen_k_
softmax_scale,
is_causal,
window_size_left,
window_size_right,
softcap,
true, // is_rotary_interleaved
num_splits,
manual_set_pack_gqa,
pack_gqa_,
sm_margin,
out,
softmax_lse,
&out_accum,
&softmax_lse_accum);
#else
RaiseNotSupportedError();
#endif
}
template <typename T, typename Context>
void FlashAttnV3VarlenKernel(const Context &dev_ctx,
const DenseTensor &q,
const DenseTensor &k,
const DenseTensor &v,
const DenseTensor &cu_seqlens_q,
const DenseTensor &cu_seqlens_k,
const optional<DenseTensor> &seqused_q,
const optional<DenseTensor> &seqused_k,
const optional<DenseTensor> &qv,
const optional<DenseTensor> &q_descale,
const optional<DenseTensor> &k_descale,
const optional<DenseTensor> &v_descale,
const Scalar &max_seqlen_q,
const Scalar &max_seqlen_k,
const float softmax_scale,
const bool causal,
const int window_size_left,
const int window_size_right,
const float softcap,
const int num_splits,
const bool manual_set_pack_gqa,
const bool pack_gqa,
const int sm_margin,
DenseTensor *out,
DenseTensor *softmax_lse) {
#ifdef PADDLE_WITH_FLASHATTN_V3
// umiswing: the following options have not been fully tested
PADDLE_ENFORCE_EQ(qv.is_initialized(),
false,
common::errors::InvalidArgument("q_v_ is not supported"));
PADDLE_ENFORCE_EQ(
q_descale.is_initialized(),
false,
common::errors::InvalidArgument("q_descale is not supported"));
PADDLE_ENFORCE_EQ(
k_descale.is_initialized(),
false,
common::errors::InvalidArgument("k_descale is not supported"));
PADDLE_ENFORCE_EQ(
v_descale.is_initialized(),
false,
common::errors::InvalidArgument("v_descale is not supported"));
PADDLE_ENFORCE_EQ(softcap,
0,
common::errors::InvalidArgument(
"softcap is not supported, please set softcap to 0"));
PADDLE_ENFORCE_EQ(
num_splits,
1,
common::errors::InvalidArgument(
"num_splits is not supported, please set num_splits to 1"));
PADDLE_ENFORCE_EQ(manual_set_pack_gqa,
false,
common::errors::InvalidArgument(
"manual_set_pack_gqa is not supported, please set "
"manual_set_pack_gqa to false"));
PADDLE_ENFORCE_EQ(
pack_gqa,
false,
common::errors::InvalidArgument(
"pack_gqa is not supported, please set pack_gqa to false"));
PADDLE_ENFORCE_EQ(
sm_margin,
0,
common::errors::InvalidArgument(
"sm_margin is not supported, please set sm_margin to 0"));
DenseTensor out_accum;
DenseTensor softmax_lse_accum;
const int64_t max_seqlen_q_ = max_seqlen_q.to<int64_t>();
const int64_t max_seqlen_k_ = max_seqlen_k.to<int64_t>();
FlashAttnV3BaseKernel<T, Context>(dev_ctx,
q,
k,
v,
paddle::none, // k_new_
paddle::none, // v_new_
qv,
paddle::none, // out_
cu_seqlens_q, // cu_seqlens_q_
cu_seqlens_k, // cu_seqlens_k_
paddle::none, // cu_seqlens_k_new_
seqused_q, // seqused_q_
seqused_k, // seqused_k_
paddle::none, // page_table_
paddle::none, // kv_batch_idx_
paddle::none, // leftpad_k_
paddle::none, // rotary_cos_
paddle::none, // rotary_sin_
q_descale,
k_descale,
v_descale,
paddle::none, // scheduler_metadata
max_seqlen_q_, // max_seqlen_q_
max_seqlen_k_, // max_seqlen_k_
softmax_scale,
causal,
window_size_left,
window_size_right,
softcap,
true, // is_rotary_interleaved
num_splits,
manual_set_pack_gqa,
pack_gqa,
sm_margin,
out,
softmax_lse,
&out_accum,
&softmax_lse_accum);
#else
RaiseNotSupportedError();
#endif
}
template <typename T, typename Context>
void FlashMaskV2BaseKernel(
const Context &dev_ctx,
const DenseTensor &q,
const DenseTensor &k,
const DenseTensor &v,
const optional<DenseTensor>
&k_new_, // (b, s_k_new, h_k, d) or (total_k_new, h_k, d) if there is
// cu_seqlens_k_new
const optional<DenseTensor>
&v_new_, // (b, s_k_new, h_k, dv) or (total_k_new, h_k, dv) if there is
// cu_seqlens_k_new
const optional<DenseTensor> &q_v_, // (b, s_q, h, dv) or (total_q_new, h,
// dv) if there is cu_seqlens_q
const optional<DenseTensor>
&out_, // (b, s_q, h, dv) or (total_q, h, dv) if there is cu_seqlens_q
const optional<DenseTensor> &cu_seqlens_q_, // b+1
const optional<DenseTensor> &cu_seqlens_k_, // b+1
const optional<DenseTensor> &cu_seqlens_k_new_, // b+1
const optional<DenseTensor>
&seqused_q_, // b. If given, only this many elements of each batch
// element's queries and outputs are used.
const optional<DenseTensor>
&seqused_k_, // b. If given, only this many elements of each batch
// element's keys are used.
const optional<DenseTensor> &page_table_, // (b_k, max_num_pages_per_seq)
const optional<DenseTensor>
&kv_batch_idx_, // b. indices to index into the KV cache
const optional<DenseTensor> &leftpad_k_, // b
const optional<DenseTensor> &rotary_cos_, // seqlen_ro x (rotary_dim / 2)
const optional<DenseTensor> &rotary_sin_, // seqlen_ro x (rotary_dim / 2)
const optional<DenseTensor> &q_descale_, // (b, h_k), not (b, h)
const optional<DenseTensor> &k_descale_, // (b, h_k)
const optional<DenseTensor> &v_descale_, // (b, h_k)
const optional<DenseTensor> &scheduler_metadata_, // (b + 1)
const optional<DenseTensor> &startend_row_indices_, // （b,h,s_1,[1,2,4])
const optional<DenseTensor> &block_mask_, // （(b,h,s// 128,s // 128)
const optional<DenseTensor>
&unique_id_, // used in distributed overlap NVSHMEM init with
// unique_id (128B u8 CPU tensor)
const int
max_seqlen_q_, // if max_seqlen_q_ is set to 0, it indicates that it is
// uninitialized and should not be referenced
// TODO(tridao): check if we need max_seqlen_k
const int
max_seqlen_k_, // if max_seqlen_q_ is set to 0, it indicates that it is
// uninitialized and should not be referenced
const float softmax_scale,
bool is_causal,
int window_size_left,
int window_size_right,
const float softcap,
const bool is_rotary_interleaved, // if true, rotary combines indices 0 &
// 1, else indices 0 & rotary_dim / 2
int num_splits,
const bool manual_set_pack_gqa,
const bool
pack_gqa_, // the pack_gqa_ will be used only if manual_set_pack_gqa is
// set to True; otherwise, the internal heuristic
// get_pack_gqa() from fa3 will decide whether to pack gqa
const int sm_margin,
const int rank,
const int nranks,
DenseTensor *out,
DenseTensor *softmax_lse,
DenseTensor *out_accum,
DenseTensor *softmax_lse_accum) {
#ifdef PADDLE_WITH_FLASHATTN_V3
// TODO(umiswing): support ampere
int device_id = dev_ctx.GetPlace().GetDeviceId();
auto dprops = paddle::platform::GetDeviceProperties(device_id);
const bool is_sm90 = dprops.major == 9 && dprops.minor == 0;
PADDLE_ENFORCE_EQ(is_sm90,
true,
common::errors::Unavailable(
"FlashAttention-3 only supports Hopper GPUs."));
auto q_type = q.dtype();
PADDLE_ENFORCE_EQ(
(q_type == DataType::FLOAT16 || q_type == DataType::BFLOAT16 ||
q_type == DataType::FLOAT8_E4M3FN),
true,
common::errors::InvalidArgument(
"FlashAttention-3 only supports fp16, bf16, and fp8_e4m3 data type"));
PADDLE_ENFORCE_EQ(k.dtype(),
q_type,
common::errors::InvalidArgument(
"query and key must have the same dtype"));
PADDLE_ENFORCE_EQ(v.dtype(),
q_type,
common::errors::InvalidArgument(
"query and value must have the same dtype"));
CHECK_DEVICE(q);
CHECK_DEVICE(k);
CHECK_DEVICE(v);
PADDLE_ENFORCE_EQ(q.strides()[q.strides().size() - 1],
1,
common::errors::InvalidArgument(
"Input tensor must have contiguous last dimension"));
PADDLE_ENFORCE_EQ(k.strides()[k.strides().size() - 1],
1,
common::errors::InvalidArgument(
"Input tensor must have contiguous last dimension"));
PADDLE_ENFORCE_EQ(v.strides()[v.strides().size() - 1],
1,
common::errors::InvalidArgument(
"Input tensor must have contiguous last dimension"));
DenseTensor page_table;
// const bool paged_KV = page_table_.has_value();
// umiswing: this is stupid but idk how to use optional
const bool paged_KV = page_table_.is_initialized();
if (paged_KV) {
page_table = page_table_.get();
CHECK_DEVICE(page_table);
PADDLE_ENFORCE_EQ(page_table.dtype(),
DataType::INT32,
common::errors::InvalidArgument(
"page_table must have dtype paddle.int32"));
PADDLE_ENFORCE_EQ(page_table.strides()[page_table.strides().size() - 1],
1,
common::errors::InvalidArgument(
"page_table must have contiguous last dimension"));
}
// TODO(umiswing): support cusum
DenseTensor cu_seqlens_q;
// bool const is_varlen_q = cu_seqlens_q_.has_value();
// TODO(umiswing): this is stupid, must fix it (after understand
// optional)
const bool is_varlen_q = cu_seqlens_q_.is_initialized();
if (is_varlen_q) {
cu_seqlens_q = cu_seqlens_q_.get();
CHECK_DEVICE(cu_seqlens_q);
CHECK_CONTIGUOUS(cu_seqlens_q);
PADDLE_ENFORCE_EQ(cu_seqlens_q.dtype(),
DataType::INT32,
common::errors::InvalidArgument(
"cu_seqlens_q must have dtype paddle.int32"));
PADDLE_ENFORCE_NE(
max_seqlen_q_,
0,
common::errors::InvalidArgument(
"max_seqlen_q must be provided if cu_seqlens_q is provided"));
}
DenseTensor cu_seqlens_k;
const bool is_varlen_k = cu_seqlens_k_.is_initialized();
if (is_varlen_k) {
cu_seqlens_k = cu_seqlens_k_.get();
CHECK_DEVICE(cu_seqlens_k);
CHECK_CONTIGUOUS(cu_seqlens_k);
PADDLE_ENFORCE_EQ(cu_seqlens_k.dtype(),
DataType::INT32,
common::errors::InvalidArgument(
"cu_seqlens_k must have dtype paddle.int32"));
PADDLE_ENFORCE_NE(
max_seqlen_k_,
0,
common::errors::InvalidArgument(
"max_seqlen_k must be provided if cu_seqlens_k is provided"));
PADDLE_ENFORCE_EQ(
!paged_KV,
true,
common::errors::InvalidArgument(
"If cu_seqlens_k is passed in, then page table is not supported"));
PADDLE_ENFORCE_EQ(
!kv_batch_idx_,
true,
common::errors::InvalidArgument(
"If cu_seqlens_k is passed in, then page table is not supported"));
}
auto const sizes = q.dims();
const int batch_size = !is_varlen_q ? sizes[0] : cu_seqlens_q.dims()[0] - 1;
int seqlen_q = !is_varlen_q ? sizes[1] : max_seqlen_q_;
int total_q = !is_varlen_q ? batch_size * sizes[1] : sizes[0];
int64_t num_heads = q.dims()[q.dims().size() - 2];
int64_t const head_size = q.dims()[q.dims().size() - 1];
int const head_size_v = v.dims()[v.dims().size() - 1];
int const max_num_pages_per_seq = !paged_KV ? 0 : page_table.dims()[1];
int const num_pages = !paged_KV ? 0 : k.dims()[0];
int const page_size = !paged_KV ? 1 : k.dims()[1];
int const seqlen_k =
!is_varlen_k
? (!paged_KV ? k.dims()[1] : max_num_pages_per_seq * page_size)
: max_seqlen_k_;
int const total_k = !is_varlen_k ? batch_size * k.dims()[1] : k.dims()[0];
int const num_heads_k = k.dims()[k.dims().size() - 2];
int const batch_size_k =
!paged_KV ? (!is_varlen_k ? k.dims()[0] : cu_seqlens_k.dims()[0] - 1)
: page_table.dims()[0];
if (!kv_batch_idx_.is_initialized()) {
PADDLE_ENFORCE_EQ(batch_size,
batch_size_k,
common::errors::InvalidArgument(
"batch_size must be equal to batch_size_k"));
}
int const max_headdim = flashmaskv2_get_max_headdim();
PADDLE_ENFORCE_LE(
head_size,
max_headdim,
common::errors::InvalidArgument(
"FlashAttention forward only supports head dimension at most %d",
max_headdim));
PADDLE_ENFORCE_EQ(
num_heads % num_heads_k,
0,
common::errors::InvalidArgument(
"Number of heads in key/value must divide number of heads in query"));
if (head_size_v != head_size) {
PADDLE_ENFORCE_EQ(
((head_size > 128 && head_size <= 192 && head_size_v > 96 &&
head_size_v <= 128) ||
(head_size <= 64 && head_size_v <= 512)),
true,
common::errors::InvalidArgument(
"If V headdim is different from Q/K dim, we only support "
"Q/K headdim in (128, 192] and V headdim in (96, 128], "
"or (Q/K <= 64 and V <= 512)."));
PADDLE_ENFORCE_EQ(dprops.major,
9,
common::errors::InvalidArgument(
"Only Hopper supports different V headdim"));
if (head_size_v > 256) {
PADDLE_ENFORCE_EQ(
(q_type == DataType::FLOAT16 || q_type == DataType::BFLOAT16),
true,
common::errors::InvalidArgument(
"HeaddimV > 256 requires fp16 and bf16 data type"));
}
}
bool const is_flashmask = startend_row_indices_.is_initialized();
bool const is_blockmask = block_mask_.is_initialized();
// This needs to go before kBlockM & kBlockN since we rely on the correct
// window_size and is_causal to set kBlockM
// TODO(tridao): check this
if (window_size_left >= seqlen_k - 1) {
window_size_left = -1;
}
if (window_size_right >= seqlen_q - 1) {
window_size_right = -1;
}
// causal=true is the same as causal=false in this case
if (seqlen_q == 1 && window_size_left == -1 && window_size_right == -1) {
// Special case of hdim 128 where we want causal to have kBlockN=128, better
// for pagedKV and TMA
if (((head_size <= 64 || head_size > 128) || !paged_KV) && !is_flashmask) {
is_causal = false;
}
}
if (is_causal) {
window_size_right = 0;
}
// There's a case where is_causal=false, window_size=(-1, 0). Then
// set_params_fprop will set params.is_causal=true. If we don't have is_causal
// here matching params.is_causal, we might get the wrong kBlockM.
is_causal = window_size_left < 0 && window_size_right == 0;
if (!is_varlen_q) {
CHECK_SHAPE(q, batch_size, seqlen_q, num_heads, head_size);
} else {
CHECK_SHAPE(q, total_q, num_heads, head_size);
CHECK_SHAPE(cu_seqlens_q, batch_size + 1);
}
if (!paged_KV) {
if (!is_varlen_k) {
CHECK_SHAPE(k, batch_size_k, seqlen_k, num_heads_k, head_size);
CHECK_SHAPE(v, batch_size_k, seqlen_k, num_heads_k, head_size_v);
} else {
CHECK_SHAPE(k, total_k, num_heads_k, head_size);
CHECK_SHAPE(v, total_k, num_heads_k, head_size_v);
CHECK_SHAPE(cu_seqlens_k, batch_size + 1);
}
} else {
CHECK_SHAPE(k, num_pages, page_size, num_heads_k, head_size);
CHECK_SHAPE(v, num_pages, page_size, num_heads_k, head_size_v);
CHECK_SHAPE(page_table, batch_size_k, max_num_pages_per_seq);
}
if (seqused_q_.is_initialized()) {
auto seqused_q = seqused_q_.get();
PADDLE_ENFORCE_EQ(
seqused_q.dtype(),
DataType::INT32,
common::errors::InvalidArgument("seqused_q must have dtype int32"));
CHECK_DEVICE(seqused_q);
CHECK_CONTIGUOUS(seqused_q);
CHECK_SHAPE(seqused_q, batch_size);
}
if (seqused_k_.is_initialized()) {
auto seqused_k = seqused_k_.get();
PADDLE_ENFORCE_EQ(
seqused_k.dtype(),
DataType::INT32,
common::errors::InvalidArgument("seqused_k must have dtype int32"));
CHECK_DEVICE(seqused_k);
CHECK_CONTIGUOUS(seqused_k);
CHECK_SHAPE(seqused_k, batch_size);
}
if (leftpad_k_.is_initialized()) {
auto leftpad_k = leftpad_k_.get();
PADDLE_ENFORCE_EQ(
leftpad_k.dtype(),
DataType::INT32,
common::errors::InvalidArgument("leftpad_k must have dtype int32"));
CHECK_DEVICE(leftpad_k);
CHECK_CONTIGUOUS(leftpad_k);
CHECK_SHAPE(leftpad_k, batch_size);
}
// This is what we will template on
bool const is_varlen =
is_varlen_q || is_varlen_k || seqused_q_.is_initialized() ||
seqused_k_.is_initialized() || leftpad_k_.is_initialized();
#ifdef FLASHATTENTION_DISABLE_VARLEN
PADDLE_ENFORCE_EQ(!is_varlen,
true,
common::errors::Unavailable(
"This flash attention build does not support varlen."));
#endif
int const alignment = q_type == DataType::FLOAT8_E4M3FN ? 16 : 8;
PADDLE_ENFORCE_EQ(head_size % alignment,
0,
common::errors::InvalidArgument(
"head_size should be a multiple of %d", alignment));
PADDLE_ENFORCE_EQ(head_size_v % alignment,
0,
common::errors::InvalidArgument(
"head_size_v should be a multiple of %d", alignment));
auto out_type =
q_type == DataType::FLOAT8_E4M3FN ? DataType::BFLOAT16 : q_type;
if (out_.is_initialized()) {
*out = out_.get();
PADDLE_ENFORCE_EQ(
out->dtype(),
out_type,
common::errors::InvalidArgument(
"For FP16/BF16 input, output must have the same dtype as "
"inputs. For FP8 input, output must have dtype BF16"));
CHECK_DEVICE((*out));
PADDLE_ENFORCE_EQ(out->strides()[out->strides().size() - 1],
1,
common::errors::InvalidArgument(
"Output tensor must have contiguous last dimension"));
if (!is_varlen_q) {
CHECK_SHAPE((*out), batch_size, seqlen_q, num_heads, head_size_v);
} else {
CHECK_SHAPE((*out), total_q, num_heads, head_size_v);
}
} else {
if (!is_varlen_q) {
out->Resize({batch_size, seqlen_q, num_heads, head_size_v});
} else {
out->Resize({total_q, num_heads, head_size_v});
}
if (q_type == DataType::FLOAT8_E4M3FN) {
dev_ctx.template Alloc<phi::bfloat16>(out);
} else {
// umiswing: assuming T is Input Type
dev_ctx.template Alloc<T>(out);
}
}
auto round_multiple = [](int x, int m) { return (x + m - 1) / m * m; };
int const head_size_rounded = flashmaskv2_round_up_headdim(head_size);
int const head_size_v_rounded = flashmaskv2_round_up_headdim(head_size_v);
int const seqlen_q_rounded = round_multiple(seqlen_q, 128);
int const seqlen_k_rounded = round_multiple(seqlen_k, 128);
if (!is_varlen_q) {
softmax_lse->Resize({batch_size, num_heads, seqlen_q});
} else {
softmax_lse->Resize({num_heads, total_q});
}
dev_ctx.template Alloc<float>(softmax_lse);
FlashMask_fwd_params *params_handle = get_flashmask_fwd_params_handle();
dynload::flashmaskv2_clear_fwd_params_handle(params_handle);
set_flashmaskv2_params_fprop(
params_handle,
batch_size,
seqlen_q,
seqlen_k,
seqlen_q_rounded,
seqlen_k_rounded,
num_heads,
num_heads_k,
head_size,
head_size_rounded,
q,
k,
v,
out,
!is_varlen_q ? nullptr : cu_seqlens_q.data(),
!is_varlen_k ? nullptr : cu_seqlens_k.data(),
seqused_q_.is_initialized() ? const_cast<void *>(seqused_q_.get().data())
: nullptr,
seqused_k_.is_initialized() ? const_cast<void *>(seqused_k_.get().data())
: nullptr,
softmax_lse->data(),
/*p_dropout=*/0.f,
softmax_scale,
window_size_left,
window_size_right,
dprops,
softcap,
sm_margin);
dynload::flashmaskv2_fwd_params_set_total_q(params_handle, total_q);
dynload::flashmaskv2_fwd_params_set_total_k(params_handle, total_k);
dynload::flashmaskv2_fwd_params_set_b_k(params_handle, batch_size_k);
dynload::flashmaskv2_fwd_params_set_dv(params_handle, head_size_v);
dynload::flashmaskv2_fwd_params_set_dv_rounded(params_handle,
head_size_v_rounded);
if (leftpad_k_
.is_initialized()) { // This needs to be set before get_pagedkv_tma
dynload::flashmaskv2_fwd_params_set_leftpad_k(params_handle,
leftpad_k_.get().data<int>());
}
if (paged_KV) {
dynload::flashmaskv2_fwd_params_set_page_table(params_handle,
page_table.data<int>());
dynload::flashmaskv2_fwd_params_set_page_table_batch_stride(
params_handle, page_table.strides()[0]);
}
dynload::flashmaskv2_fwd_params_set_page_size(params_handle, page_size);
dynload::flashmaskv2_fwd_params_set_num_pages(params_handle, num_pages);
if (k_new_.is_initialized()) { // This needs to be set before get_pagedkv_tma
DenseTensor k_new, v_new;
PADDLE_ENFORCE_EQ(
v_new_.is_initialized(),
true,
common::errors::InvalidArgument(
"If k_new is supplied, v_new must also be passed in"));
PADDLE_ENFORCE_EQ(
seqused_k_.is_initialized(),
true,
common::errors::InvalidArgument(
"If k_new is supplied, seqlens_k must also be passed in"));
PADDLE_ENFORCE_LE(
seqlen_q,
seqlen_k,
common::errors::InvalidArgument(
"If k_new is supplied, it must have seqlen <= the seqlen "
"of the KV cache"));
DenseTensor cu_seqlens_k_new;
bool const is_varlen_k_new = cu_seqlens_k_new_.is_initialized();
if (is_varlen_k_new) {
cu_seqlens_k_new = cu_seqlens_k_new_.get();
CHECK_DEVICE(cu_seqlens_k_new);
CHECK_CONTIGUOUS(cu_seqlens_k_new);
PADDLE_ENFORCE_EQ(cu_seqlens_k_new.dtype(),
DataType::INT32,
common::errors::InvalidArgument(
"cu_seqlens_k_new must have dtype paddle.int32"));
}
k_new = k_new_.get();
v_new = v_new_.get();
PADDLE_ENFORCE_EQ(k_new.dtype(),
q_type,
common::errors::InvalidArgument(
"k_new must have the same dtype as query"));
PADDLE_ENFORCE_EQ(v_new.dtype(),
q_type,
common::errors::InvalidArgument(
"v_new must have the same dtype as query"));
CHECK_DEVICE(k_new);
CHECK_DEVICE(v_new);
PADDLE_ENFORCE_EQ(k_new.strides()[k_new.strides().size() - 1],
1,
common::errors::InvalidArgument(
"k_new tensor must have contiguous last dimension"));
PADDLE_ENFORCE_EQ(v_new.strides()[v_new.strides().size() - 1],
1,
common::errors::InvalidArgument(
"v_new tensor must have contiguous last dimension"));
// We don't need max_seqlen_k_new, so seqlen_k_new can be whatever when
// is_varlen_k_new
int seqlen_k_new = !is_varlen_k_new ? k_new.dims()[1] : 0;
int total_k_new =
!is_varlen_k_new ? batch_size * k_new.dims()[1] : k_new.dims()[0];
if (!is_varlen_k_new) {
CHECK_SHAPE(k_new, batch_size, seqlen_k_new, num_heads_k, head_size);
CHECK_SHAPE(v_new, batch_size, seqlen_k_new, num_heads_k, head_size_v);
} else {
CHECK_SHAPE(k_new, total_k_new, num_heads_k, head_size);
CHECK_SHAPE(v_new, total_k_new, num_heads_k, head_size_v);
CHECK_SHAPE(cu_seqlens_k_new, batch_size + 1);
}
// umiswing: dump this to shared library
dynload::flashmaskv2_fwd_params_set_seqlen_knew(params_handle,
seqlen_k_new);
dynload::flashmaskv2_fwd_params_set_total_knew(params_handle, total_k_new);
dynload::flashmaskv2_fwd_params_set_knew_ptr(params_handle, (k_new.data()));
dynload::flashmaskv2_fwd_params_set_vnew_ptr(params_handle, (v_new.data()));
// All stride are in elements, not bytes.
dynload::flashmaskv2_fwd_params_set_knew_row_stride(
params_handle, k_new.strides()[k_new.strides().size() - 3]);
dynload::flashmaskv2_fwd_params_set_vnew_row_stride(
params_handle, v_new.strides()[v_new.strides().size() - 3]);
dynload::flashmaskv2_fwd_params_set_knew_head_stride(
params_handle, k_new.strides()[k_new.strides().size() - 2]);
dynload::flashmaskv2_fwd_params_set_vnew_head_stride(
params_handle, v_new.strides()[v_new.strides().size() - 2]);
if (!is_varlen_k_new) {
dynload::flashmaskv2_fwd_params_set_knew_batch_stride(params_handle,
k_new.strides()[0]);
dynload::flashmaskv2_fwd_params_set_vnew_batch_stride(params_handle,
v_new.strides()[0]);
}
if (is_varlen_k_new) {
dynload::flashmaskv2_fwd_params_set_cu_seqlens_knew(
params_handle, cu_seqlens_k_new.data<int>());
}
}
// 992 = 32 * 31 is the max supported batch in prepare_varlen_num_blocks
// kernel
bool const use_dynamic_split =
is_varlen && dynload::flashmaskv2_fwd_params_get_b(params_handle) <= 992;
// Temporarily set num_splits_dynamic_ptr to 1 since get_num_splits checks it
dynload::flashmaskv2_fwd_params_set_num_splits_dynamic_ptr(
params_handle, !use_dynamic_split ? nullptr : reinterpret_cast<int *>(1));
dynload::flashmaskv2_fwd_params_set_pagedkv_tma(
params_handle, dynload::flashmaskv2_get_pagedkv_tma(params_handle));
if (num_splits <= 0) {
num_splits = dynload::flashmaskv2_get_num_splits(params_handle);
}
dynload::flashmaskv2_fwd_params_set_num_splits(params_handle, num_splits);
// Always enable PackGQA for Split, and get_pack_gqa requires
// params.num_splits to decide
const bool pack_gqa = manual_set_pack_gqa
? pack_gqa_
: dynload::flashmaskv2_get_pack_gqa(params_handle);
dynload::flashmaskv2_fwd_params_set_pack_gqa(params_handle, pack_gqa);
// This needs to be set after get_num_splits
DenseTensor tile_count_semaphore; // Contains the semaphore and optionally
// num_splits_dynamic
// We don't use the persistent scheduler if Split and not Varlen
const bool params_is_causal =
dynload::flashmaskv2_fwd_params_get_is_causal(params_handle);
const bool params_is_local =
dynload::flashmaskv2_fwd_params_get_is_local(params_handle);
const int params_num_splits =
dynload::flashmaskv2_fwd_params_get_num_splits(params_handle);
const int params_b = dynload::flashmaskv2_fwd_params_get_b(params_handle);
const int params_arch =
dynload::flashmaskv2_fwd_params_get_arch(params_handle);
bool const scheduler_needs_semaphore =
params_arch >= 90 ? true
: ((params_is_causal && !is_varlen) ||
(is_varlen && params_num_splits > 1));
if (scheduler_needs_semaphore || use_dynamic_split) {
int metadata_size = static_cast<int>(scheduler_needs_semaphore) +
static_cast<int>(use_dynamic_split) * params_b;
dynload::flashmaskv2_fwd_params_set_skip_scheduler_metadata_computation(
params_handle, scheduler_metadata_.is_initialized());
if (scheduler_metadata_.is_initialized()) {
DenseTensor scheduler_metadata = scheduler_metadata_.get();
CHECK_DEVICE(scheduler_metadata);
CHECK_SHAPE(scheduler_metadata, metadata_size);
CHECK_CONTIGUOUS(scheduler_metadata);
PADDLE_ENFORCE_EQ(scheduler_metadata.dtype(),
DataType::INT32,
common::errors::InvalidArgument(
"scheduler_metadata must have dtype int32"));
tile_count_semaphore = scheduler_metadata;
} else {
tile_count_semaphore = Empty<int32_t>(dev_ctx, {metadata_size});
}
if (scheduler_needs_semaphore && !use_dynamic_split) {
funcs::SetConstant<Context, int32_t> set_zero;
set_zero(dev_ctx,
&tile_count_semaphore,
int32_t{0}); // If varlen we'll manually do the zero-ing
}
dynload::flashmaskv2_fwd_params_set_tile_count_semaphore(
params_handle,
scheduler_needs_semaphore ? (tile_count_semaphore.data<int>())
: nullptr);
dynload::flashmaskv2_fwd_params_set_num_splits_dynamic_ptr(
params_handle,
use_dynamic_split ? (tile_count_semaphore.data<int>()) + 1 : nullptr);
}
if (q_v_.is_initialized()) {
PADDLE_ENFORCE_LT(head_size,
64,
common::errors::InvalidArgument(
"q_v is only supported for head_size <= 64"));
PADDLE_ENFORCE_EQ(
(q_type == DataType::FLOAT16 || q_type == DataType::FLOAT16),
true,
common::errors::InvalidArgument(
"q_v is only supported for fp16 and bf16 data type"));
PADDLE_ENFORCE_EQ(params_arch,
90,
common::errors::InvalidArgument(
"q_v is only supported for Hopper GPUs"));
DenseTensor q_v = q_v_.get();
PADDLE_ENFORCE_EQ(q_v.dtype(),
q_type,
common::errors::InvalidArgument(
"q_v must have the same dtype as query"));
CHECK_DEVICE(q_v);
PADDLE_ENFORCE_EQ(q_v.strides()[q_v.strides().size() - 1],
1,
common::errors::InvalidArgument(
"q_v tensor must have contiguous last dimension"));
if (!is_varlen_q) {
CHECK_SHAPE(q_v, batch_size, seqlen_q, num_heads, head_size_v);
} else {
CHECK_SHAPE(q_v, total_q, num_heads, head_size_v);
}
dynload::flashmaskv2_fwd_params_set_qv_ptr(params_handle, (q_v.data()));
// All stride are in elements, not bytes.
dynload::flashmaskv2_fwd_params_set_qv_row_stride(
params_handle, q_v.strides()[q_v.strides().size() - 3]);
dynload::flashmaskv2_fwd_params_set_qv_head_stride(
params_handle, q_v.strides()[q_v.strides().size() - 2]);
if (!is_varlen_q) {
dynload::flashmaskv2_fwd_params_set_qv_batch_stride(params_handle,
q_v.strides()[0]);
}
}
if (rotary_cos_.is_initialized()) {
PADDLE_ENFORCE_EQ(
k_new_.is_initialized(),
true,
common::errors::InvalidArgument(
"If rotary cos/sin are provided, new key / value to be "
"appended to KV cache must also be provided"));
DenseTensor rotary_cos = rotary_cos_.get();
CHECK_DEVICE(rotary_cos);
CHECK_CONTIGUOUS(rotary_cos);
int params_rotary_dim = rotary_cos.dims()[1] * 2;
dynload::flashmaskv2_fwd_params_set_rotary_dim(params_handle,
params_rotary_dim);
PADDLE_ENFORCE_LE(
params_rotary_dim,
head_size,
common::errors::InvalidArgument("rotary_dim must be <= headdim"));
PADDLE_ENFORCE_EQ(
params_rotary_dim % 16,
0,
common::errors::InvalidArgument(
"Only rotary dimensions divisible by 16 are currently supported"));
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t seqlen_ro = rotary_cos.dims()[0];
if (paged_KV) {
PADDLE_ENFORCE_GE(
seqlen_ro,
seqlen_k,
common::errors::InvalidArgument(
"cos/sin seqlen must be at least the seqlen of KV cache"));
}
CHECK_SHAPE(rotary_cos, seqlen_ro, params_rotary_dim / 2);
PADDLE_ENFORCE_EQ(rotary_cos.dtype(),
q_type,
common::errors::InvalidArgument(
"rotary_cos must have the same dtype as query"));
PADDLE_ENFORCE_EQ(
rotary_sin_.is_initialized(),
true,
common::errors::InvalidArgument(
"If rotary cos is provided, rotary sin must also be provided"));
auto rotary_sin = rotary_sin_.get();
CHECK_DEVICE(rotary_sin);
CHECK_CONTIGUOUS(rotary_sin);
CHECK_SHAPE(rotary_sin, seqlen_ro, params_rotary_dim / 2);
PADDLE_ENFORCE_EQ(rotary_sin.dtype(),
q_type,
common::errors::InvalidArgument(
"rotary_cos must have the same dtype as query"));
dynload::flashmaskv2_fwd_params_set_rotary_cos_ptr(params_handle,
(rotary_cos.data()));
dynload::flashmaskv2_fwd_params_set_rotary_sin_ptr(params_handle,
(rotary_sin.data()));
dynload::flashmaskv2_fwd_params_set_is_rotary_interleaved(
params_handle, is_rotary_interleaved);
} else {
dynload::flashmaskv2_fwd_params_set_rotary_dim(params_handle, 0);
}
if (kv_batch_idx_.is_initialized()) {
DenseTensor kv_batch_idx = kv_batch_idx_.get();
CHECK_DEVICE(kv_batch_idx);
CHECK_CONTIGUOUS(kv_batch_idx);
PADDLE_ENFORCE_EQ(
kv_batch_idx.dtype(),
DataType::INT32,
common::errors::InvalidArgument("kv_batch_idx must have dtype int32"));
dynload::flashmaskv2_fwd_params_set_kv_batch_idx(
params_handle, reinterpret_cast<int *>(kv_batch_idx.data()));
}
if (dynload::flashmaskv2_fwd_params_get_num_splits(params_handle) > 1) {
PADDLE_ENFORCE_LE(
dynload::flashmaskv2_fwd_params_get_num_splits(params_handle),
256,
common::errors::InvalidArgument("num_splits > 256 not supported"));
if (!is_varlen_q) {
out_accum->Resize(make_ddim(
{dynload::flashmaskv2_fwd_params_get_num_splits(params_handle),
batch_size,
num_heads,
seqlen_q,
head_size_v}));
softmax_lse_accum->Resize(make_ddim(
{dynload::flashmaskv2_fwd_params_get_num_splits(params_handle),
batch_size,
num_heads,
seqlen_q}));
dev_ctx.template Alloc<float>(out_accum);
dev_ctx.template Alloc<float>(softmax_lse_accum);
dynload::flashmaskv2_fwd_params_set_oaccum_batch_stride(
params_handle, out_accum->strides()[1]);
dynload::flashmaskv2_fwd_params_set_lseaccum_batch_stride(
params_handle, softmax_lse_accum->strides()[1]);
} else {
out_accum->Resize(make_ddim(
{dynload::flashmaskv2_fwd_params_get_num_splits(params_handle),
num_heads,
total_q,
head_size_v}));
softmax_lse_accum->Resize(make_ddim(
{dynload::flashmaskv2_fwd_params_get_num_splits(params_handle),
num_heads,
total_q}));
dev_ctx.template Alloc<float>(out_accum);
dev_ctx.template Alloc<float>(softmax_lse_accum);
}
dynload::flashmaskv2_fwd_params_set_is_fp32(params_handle, false);
dynload::flashmaskv2_fwd_params_set_oaccum_ptr(params_handle,
(out_accum->data()));
dynload::flashmaskv2_fwd_params_set_softmax_lseaccum_ptr(
params_handle, (softmax_lse_accum->data()));
dynload::flashmaskv2_fwd_params_set_oaccum_split_stride(
params_handle, out_accum->strides()[0]);
dynload::flashmaskv2_fwd_params_set_oaccum_row_stride(
params_handle, out_accum->strides()[out_accum->strides().size() - 2]);
dynload::flashmaskv2_fwd_params_set_oaccum_head_stride(
params_handle, out_accum->strides()[out_accum->strides().size() - 3]);
dynload::flashmaskv2_fwd_params_set_lseaccum_split_stride(
params_handle, softmax_lse_accum->strides()[0]);
dynload::flashmaskv2_fwd_params_set_lseaccum_head_stride(
params_handle,
softmax_lse_accum->strides()[softmax_lse_accum->strides().size() - 2]);
}
if (q_type == DataType::FLOAT8_E4M3FN) {
if (q_descale_.is_initialized()) {
DenseTensor q_descale = q_descale_.get();
CHECK_DEVICE(q_descale);
CHECK_SHAPE(q_descale, batch_size, num_heads_k);
dynload::flashmaskv2_fwd_params_set_q_descale_ptr(
params_handle, (q_descale.data<float>()));
dynload::flashmaskv2_fwd_params_set_q_descale_batch_stride(
params_handle, q_descale.strides()[0]);
dynload::flashmaskv2_fwd_params_set_q_descale_head_stride(
params_handle, q_descale.strides()[1]);
} else {
dynload::flashmaskv2_fwd_params_set_q_descale_ptr(params_handle, nullptr);
}
if (k_descale_.is_initialized()) {
DenseTensor k_descale = k_descale_.get();
CHECK_DEVICE(k_descale);
CHECK_SHAPE(k_descale, batch_size, num_heads_k);
dynload::flashmaskv2_fwd_params_set_k_descale_ptr(
params_handle, (k_descale.data<float>()));
dynload::flashmaskv2_fwd_params_set_k_descale_batch_stride(
params_handle, k_descale.strides()[0]);
dynload::flashmaskv2_fwd_params_set_k_descale_head_stride(
params_handle, k_descale.strides()[1]);
} else {
dynload::flashmaskv2_fwd_params_set_k_descale_ptr(params_handle, nullptr);
}
if (v_descale_.is_initialized()) {
DenseTensor v_descale = v_descale_.get();
CHECK_DEVICE(v_descale);
CHECK_SHAPE(v_descale, batch_size, num_heads_k);
dynload::flashmaskv2_fwd_params_set_v_descale_ptr(
params_handle, (v_descale.data<float>()));
dynload::flashmaskv2_fwd_params_set_v_descale_batch_stride(
params_handle, v_descale.strides()[0]);
dynload::flashmaskv2_fwd_params_set_v_descale_head_stride(
params_handle, v_descale.strides()[1]);
} else {
dynload::flashmaskv2_fwd_params_set_v_descale_ptr(params_handle, nullptr);
}
}
#ifdef FLASHATTENTION_DISABLE_LOCAL
PADDLE_ENFORCE_EQ(
!dynload::flashmaskv2_fwd_params_get_is_local(params_handle),
true,
common::errors::InvalidArgument(
"This flash attention build does not support local attention."));
#endif
#ifdef FLASHATTENTION_DISABLE_SOFTCAP
PADDLE_ENFORCE_EQ(
dynload::flashmaskv2_fwd_params_get_softcap(params_handle),
0.0,
common::errors::InvalidArgument(
"This flash attention build does not support tanh softcapping."));
#endif
#ifdef FLASHATTENTION_DISABLE_SPLIT
PADDLE_ENFORCE_EQ(
dynload::flashmaskv2_fwd_params_get_num_splits(params_handle),
1,
common::errors::InvalidArgument(
"This flash attention build does not support splits."));
#endif
#ifdef FLASHATTENTION_DISABLE_PACKGQA
PADDLE_ENFORCE_EQ(
(!dynload::flashmaskv2_fwd_params_get_pack_gqa(params_handle) ||
dynload::flashmaskv2_fwd_params_get_arch(params_handle) < 90 ||
(dynload::flashmaskv2_fwd_params_get_page_table(params_handle) &&
!dynload::flashmaskv2_fwd_params_get_pagedkv_tma(params_handle)) ||
dynload::flashmaskv2_fwd_params_get_num_splits(params_handle) > 1),
true,
common::errors::InvalidArgument(
"This flash attention build does not support pack_gqa."));
#endif
#ifdef FLASHATTENTION_DISABLE_PAGEDKV
PADDLE_ENFORCE_EQ(
(!(dynload::flashmaskv2_fwd_params_get_page_table(params_handle) &&
!dynload::flashmaskv2_fwd_params_get_pagedkv_tma(params_handle))),
true,
common::errors::InvalidArgument(
"This flash attention build does not support paged KV."));
#endif
#ifdef FLASHATTENTION_DISABLE_APPENDKV
PADDLE_ENFORCE_EQ(
!k_new_.is_initialized(),
true,
common::errors::InvalidArgument(
"This flash attention build does not support appending KV."));
#endif
// flashmask
DenseTensor startend_row_indices;
if (is_flashmask) startend_row_indices = startend_row_indices_.get();
DenseTensor block_mask;
if (is_blockmask) block_mask = block_mask_.get();
DenseTensor flashmask_maxmin, lt_start_row_indices, lt_end_row_indices,
ut_start_row_indices, ut_end_row_indices;
if (is_flashmask) {
PADDLE_ENFORCE_EQ(
startend_row_indices.dims().size(),
4,
common::errors::InvalidArgument(
"flashmask_attention receive startend_row_indices with dim "
"[batch_size, num_heads,seq_len, mask_bounds]"));
PADDLE_ENFORCE_EQ(startend_row_indices.dims()[3] == 1 ||
startend_row_indices.dims()[3] == 2 ||
startend_row_indices.dims()[3] == 4,
true,
common::errors::InvalidArgument(
"flashmask_attention startend_row_indices "
"mask_bounds must in [1,2,4]"));
auto flashmask_maxmin_shape = startend_row_indices.dims();
// TODO(umiswing): refine this block constraint (kBlockN % 32), since some
// of kBlockN is not divisible by 32 flashmask_maxmin_shape[2] =
// (flashmask_maxmin_shape[2] + 31) / 32 * 8;
int device_id = dev_ctx.GetPlace().GetDeviceId();
auto dprops = paddle::platform::GetDeviceProperties(device_id);
const bool is_sm90 = dprops.major == 9 && dprops.minor == 0;
if (is_sm90) {
// seqlen_k to nblock_seqlen, here we use kBlockN = 64
// as a conservative estimation (reduce allocation size)
flashmask_maxmin_shape[2] =
((flashmask_maxmin_shape[2] + 63) / 64 + 3) / 4 * 4;
// make sure this is the same with FlashMaskV3 fwd main loop
static constexpr int flashmask_buffer_length = 16 * 1024;
// estimate the upper bound of the possible chunk size
static constexpr int chunk_padded_length =
((flashmask_buffer_length + 63) / 64 + 31) & 0xffffffe0;
static constexpr int chunk_valid_length =
((flashmask_buffer_length + 63) / 64 + 3) & 0xfffffffc;
const int num_chunk =
(flashmask_maxmin_shape[2] + chunk_valid_length - 1) /
chunk_valid_length;
flashmask_maxmin_shape[2] = num_chunk * chunk_padded_length;
} else {
// seqlen_k to nblock_seqlen
flashmask_maxmin_shape[2] =
((flashmask_maxmin_shape[2] + 31) / 32 + 3) / 4 * 4;
}
flashmask_maxmin_shape[3] = 8;
flashmask_maxmin.set_type(DataType::INT32);
flashmask_maxmin.Resize(flashmask_maxmin_shape);
dev_ctx.template Alloc<int32_t>(&flashmask_maxmin);
lt_start_row_indices =
phi::Slice<int32_t>(dev_ctx, startend_row_indices, {3}, {0}, {1});
if (startend_row_indices.dims()[3] == 2) {
if (!is_causal) {
ut_end_row_indices =
phi::Slice<int32_t>(dev_ctx, startend_row_indices, {3}, {1}, {2});
} else {
lt_end_row_indices =
phi::Slice<int32_t>(dev_ctx, startend_row_indices, {3}, {1}, {2});
}
} else if (startend_row_indices.dims()[3] == 4) {
ut_end_row_indices =
phi::Slice<int32_t>(dev_ctx, startend_row_indices, {3}, {3}, {4});
lt_end_row_indices =
phi::Slice<int32_t>(dev_ctx, startend_row_indices, {3}, {1}, {2});
ut_start_row_indices =
phi::Slice<int32_t>(dev_ctx, startend_row_indices, {3}, {2}, {3});
}
}
if (is_blockmask) {
PADDLE_ENFORCE_EQ(
is_flashmask,
true,
common::errors::InvalidArgument(
"blockmask should be used with flashmask at the same time "));
PADDLE_ENFORCE_EQ(block_mask.dims().size(),
4,
common::errors::InvalidArgument(
"blockmask receive blockmask_indices with dim "
"[batch_size, num_heads, blocklen_q, blocklen_k]"));
PADDLE_ENFORCE_EQ(block_mask.dims()[2],
(seqlen_q + 127) / 128,
common::errors::InvalidArgument(
"blockmask is now only support blockdim_q = 128 "));
PADDLE_ENFORCE_EQ(block_mask.dims()[3],
(seqlen_k + 127) / 128,
common::errors::InvalidArgument(
"blockmask is now only support blockdim_k = 128 "));
PADDLE_ENFORCE_EQ(
block_mask.dims()[1],
startend_row_indices.dims()[1],
common::errors::InvalidArgument("blockmask is now only support same "
"dim num_heads with flashmask "));
}
if (is_blockmask) {
// xhy: blockmask is now only support blockdim_q k = 128
dynload::flashmaskv2_fwd_params_set_m_block_dim(params_handle, 128);
dynload::flashmaskv2_fwd_params_set_n_block_dim(params_handle, 128);
dynload::flashmaskv2_fwd_params_set_block_mask_ptr(
params_handle, (block_mask.data<int32_t>()));
}
if (is_flashmask) {
if (lt_start_row_indices.initialized())
dynload::flashmaskv2_fwd_params_set_lt_start_ptr(
params_handle, (lt_start_row_indices.data<int32_t>()));
else
dynload::flashmaskv2_fwd_params_set_lt_start_ptr(params_handle, nullptr);
if (lt_end_row_indices.initialized())
dynload::flashmaskv2_fwd_params_set_lt_end_ptr(
params_handle, (lt_end_row_indices.data<int32_t>()));
else
dynload::flashmaskv2_fwd_params_set_lt_end_ptr(params_handle, nullptr);
if (ut_start_row_indices.initialized())
dynload::flashmaskv2_fwd_params_set_ut_start_ptr(
params_handle, (ut_start_row_indices.data<int32_t>()));
else
dynload::flashmaskv2_fwd_params_set_ut_start_ptr(params_handle, nullptr);
if (ut_end_row_indices.initialized())
dynload::flashmaskv2_fwd_params_set_ut_end_ptr(
params_handle, (ut_end_row_indices.data<int32_t>()));
else
dynload::flashmaskv2_fwd_params_set_ut_end_ptr(params_handle, nullptr);
if (flashmask_maxmin.initialized())
dynload::flashmaskv2_fwd_params_set_flashmask_maxmin_ptr(
params_handle, (flashmask_maxmin.data<int32_t>()));
else
dynload::flashmaskv2_fwd_params_set_flashmask_maxmin_ptr(params_handle,
nullptr);
dynload::flashmaskv2_fwd_params_set_h_flashmask(
params_handle, startend_row_indices.dims()[1]);
dynload::flashmaskv2_fwd_params_set_h_h_flashmask_ratio(
params_handle, num_heads / startend_row_indices.dims()[1]);
// distributed settings
#ifdef PADDLE_WITH_NVSHMEM
PADDLE_ENFORCE_LE(
nranks,
64,
common::errors::InvalidArgument(
"nranks for FlashMask overlap should <= 64, got: %d", nranks));
dynload::flashmaskv2_fwd_params_set_rank(params_handle, rank);
dynload::flashmaskv2_fwd_params_set_nranks(params_handle, nranks);
if (unique_id_.is_initialized()) {
dynload::flashmaskv2_fwd_params_set_unique_id_ptr(
params_handle, unique_id_.get().data<uint8_t>());
VLOG(6) << "FlashMask overlap debug: unique_id_ptr set.";
} else {
dynload::flashmaskv2_fwd_params_set_unique_id_ptr(params_handle, nullptr);
}
VLOG(6) << "FlashMask overlap debug (rank and nranks): " << rank << ", "
<< nranks;
#else
VLOG(6) << "FlashMask overlap is not being used since PADDLE_WITH_NVSHMEM "
"is not defined.";
#endif // PADDLE_WITH_NVSHMEM
} else {
dynload::flashmaskv2_fwd_params_set_lt_start_ptr(params_handle, nullptr);
dynload::flashmaskv2_fwd_params_set_lt_end_ptr(params_handle, nullptr);
dynload::flashmaskv2_fwd_params_set_ut_start_ptr(params_handle, nullptr);
dynload::flashmaskv2_fwd_params_set_ut_end_ptr(params_handle, nullptr);
dynload::flashmaskv2_fwd_params_set_flashmask_maxmin_ptr(params_handle,
nullptr);
dynload::flashmaskv2_fwd_params_set_h_flashmask(params_handle, 0);
dynload::flashmaskv2_fwd_params_set_h_h_flashmask_ratio(params_handle, 0);
}
if (total_q > 0 &&
(total_k +
dynload::flashmaskv2_fwd_params_get_total_knew(params_handle)) > 0 &&
num_heads_k > 0) {
dynload::flashmaskv2_run_mha_fwd(params_handle, dev_ctx.stream());
if (dynload::flashmaskv2_fwd_params_get_num_splits(params_handle) > 1) {
if (out_type == DataType::BFLOAT16) {
// Since we want output in BF16. Otherwise fwd_combine will output to
// FP16
dynload::flashmaskv2_fwd_params_set_is_bf16(params_handle, true);
}
// Unless there's seqused_q, for the purpose of attn_combine, we can just
// treat it as batch=1 and seqlen = total_q, and don't need to dispatch to
// Varlen there. However, with dynamic split, each row needs to know which
// batch it belongs to to read the number of splits, so we just use the
// varlen version of combine kernel. if (is_varlen_q &&
// !seqused_q_.has_value()) { if (is_varlen_q) {
// params.b = 1;
// params.seqlen_q = total_q;
// }
// }
dynload::flashmaskv2_run_mha_fwd_combine(
params_handle, dev_ctx.stream(), true /*enable_pdl*/);
}
} else if (total_q > 0 && num_heads_k > 0) {
PADDLE_ENFORCE_EQ(
(out->dtype() == DataType::BFLOAT16 ||
out->dtype() == DataType::FLOAT16 ||
out->dtype() == DataType::FLOAT8_E4M3FN),
true,
common::errors::InvalidArgument("flash attention 3 supports bfloat16, "
"float16 and float8_e4m3fn only."));
// If seqlen_k == 0, then we have an empty tensor. We need to set the output
// to 0.
if (out->dtype() == DataType::BFLOAT16) {
funcs::SetConstant<Context, phi::bfloat16> set_zero;
set_zero(dev_ctx,
out,
phi::bfloat16{0}); // If varlen we'll manually do the zero-ing
} else if (out->dtype() == DataType::FLOAT16) {
funcs::SetConstant<Context, phi::float16> set_zero;
set_zero(dev_ctx,
out,
phi::float16{0}); // If varlen we'll manually do the zero-ing
} else if (out->dtype() == DataType::FLOAT8_E4M3FN) {
funcs::SetConstant<Context, phi::float8_e4m3fn> set_zero;
set_zero(
dev_ctx,
out,
phi::float8_e4m3fn{0}); // If varlen we'll manually do the zero-ing
}
funcs::SetConstant<Context, float> set_infinity;
set_infinity(dev_ctx, softmax_lse, std::numeric_limits<float>::infinity());
}
#else
RaiseNotSupportedError();
#endif
}
template <typename T, typename Context>
void FlashMaskV2Kernel(const Context &dev_ctx,
const DenseTensor &q,
const DenseTensor &k,
const DenseTensor &v,
const DenseTensor &startend_row_indices,
const optional<DenseTensor> &block_mask,
const optional<DenseTensor> &unique_id,
const float softmax_scale,
bool is_causal,
const int rank,
const int nranks,
DenseTensor *out,
DenseTensor *softmax_lse) {
#ifdef PADDLE_WITH_FLASHATTN_V3
// Handle 0-size tensors: return zeros without calling CUDA kernel
// to avoid invalid memory access
if (q.numel() == 0 || k.numel() == 0 || v.numel() == 0) {
if (out) {
funcs::SetConstant<Context, T> set_zero;
set_zero(dev_ctx, out, T{0});
}
if (softmax_lse) {
funcs::SetConstant<Context, float> set_infinity;
set_infinity(
dev_ctx, softmax_lse, std::numeric_limits<float>::infinity());
}
return;
}
DenseTensor out_accum;
DenseTensor softmax_lse_accum;
FlashMaskV2BaseKernel<T, Context>(dev_ctx,
q,
k,
v,
paddle::none, // k_new_
paddle::none, // v_new_
paddle::none, // q_v_
paddle::none, // out_
paddle::none, // cu_seqlens_q_
paddle::none, // cu_seqlens_k_
paddle::none, // cu_seqlens_k_new_
paddle::none, // seqused_q_
paddle::none, // seqused_k_
paddle::none, // page_table_
paddle::none, // kv_batch_idx_
paddle::none, // leftpad_k_
paddle::none, // rotary_cos_
paddle::none, // rotary_sin_
paddle::none, // q_descale_
paddle::none, // k_descale_
paddle::none, // v_descale_
paddle::none, // scheduler_metadata_
startend_row_indices,
block_mask,
unique_id,
0, // max_seqlen_q_
0, // max_seqlen_k_
softmax_scale,
is_causal,
-1, // window_size_left
-1, // window_size_right
float{0}, // softcap
true, // is_rotary_interleaved
1, // num_splits
false, // manual_set_pack_gqa
false, // pack_gqa_
0, // sm_margin
rank, // dist CP settings
nranks, // dist CP settings
out,
softmax_lse,
&out_accum,
&softmax_lse_accum);
#else
RaiseNotSupportedError();
#endif
}
template <typename T, typename Context>
void FlashMaskV2GetUniqueIdInplace(const Context &dev_ctx,
const DenseTensor &x,
DenseTensor *out) {
#if defined(PADDLE_WITH_CUDA) && defined(PADDLE_WITH_FLASHATTN_V3)
bool valid_unique_id =
dynload::flashmaskv2_get_nvshmem_unique_id(out->data<uint8_t>());
if (!valid_unique_id) {
// If FlashMask is not compiled with `WITH_DISTRIBUTED_OVERLAP` then this is
// a zero tensor
funcs::SetConstant<Context, uint8_t> set_zero;
set_zero(dev_ctx, out, uint8_t{0});
}
#else
funcs::SetConstant<Context, uint8_t> set_zero;
set_zero(dev_ctx, out, uint8_t{0});
#endif
}
} // namespace phi
PD_REGISTER_KERNEL(flashmask_get_unique_id,
CPU,
ALL_LAYOUT,
phi::FlashMaskV2GetUniqueIdInplace,
uint8_t) {
kernel->InputAt(0).SetBackend(phi::Backend::CPU);
}
PD_REGISTER_KERNEL(flash_attn_v3,
GPU,
ALL_LAYOUT,
phi::FlashAttnV3Kernel,
phi::float16,
phi::bfloat16) {}
PD_REGISTER_KERNEL(flash_attn_v3_varlen,
GPU,
ALL_LAYOUT,
phi::FlashAttnV3VarlenKernel,
phi::float16,
phi::bfloat16) {}
PD_REGISTER_KERNEL(flashmask_attention_v2,
GPU,
ALL_LAYOUT,
phi::FlashMaskV2Kernel,
phi::float16,
phi::bfloat16) {
kernel->InputAt(4).SetBackend(phi::Backend::ALL_BACKEND); // block_mask
kernel->InputAt(5).SetBackend(phi::Backend::CPU); // nvshmem unique_id
}
Reference in New Issue View Git Blame Copy Permalink