// Copyright (c) 2023 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 "glog/logging.h" // For VLOG() #include "paddle/phi/backends/gpu/gpu_context.h" #include "paddle/phi/common/data_type.h" #include "paddle/phi/core/enforce.h" #include "paddle/phi/core/flags.h" #include "paddle/phi/core/kernel_registry.h" #include "paddle/phi/core/tensor_utils.h" #include "paddle/phi/kernels/arange_kernel.h" #include "paddle/phi/kernels/empty_kernel.h" #include "paddle/phi/kernels/reshape_kernel.h" #ifdef PADDLE_WITH_FLASHATTN #include "paddle/phi/backends/dynload/flashattn.h" #include "paddle/phi/kernels/gpu/flash_attn_utils.h" #endif DECLARE_bool(cudnn_deterministic); namespace phi { template void FlashAttnUnpaddedKernel( const Context& ctx, const DenseTensor& q, const DenseTensor& k, const DenseTensor& v, const DenseTensor& cu_seqlens_q, const DenseTensor& cu_seqlens_k, const paddle::optional& fixed_seed_offset, int64_t max_seqlen_q, int64_t max_seqlen_k, float scale, float dropout, bool causal, bool return_softmax, bool is_test, const std::string& rng_name, DenseTensor* out, DenseTensor* softmax, DenseTensor* softmax_lse, DenseTensor* seed_offset) { #ifdef PADDLE_WITH_FLASHATTN ctx.template Alloc(out); cudaStream_t stream = ctx.stream(); // q,k,v [total_*, num_heads, head_dim] auto dims = q.dims(); PADDLE_ENFORCE_EQ( dims.size(), 3, phi::errors::InvalidArgument("flash_attn_raw receive input with dim " "[total_seq_len, num_heads, head_dim]")); const int64_t total_q = dims[0]; const int num_heads = dims[1]; const int head_size = dims[2]; const int total_k = k.dims()[0]; const int num_heads_k = k.dims()[1]; const int batch_size = cu_seqlens_q.numel() - 1; // TODO(umiswing): add deterministic in fa2. // int num_splits = 0; // 0 for an internal heuristic, which is optimal // if (FLAGS_cudnn_deterministic) { // num_splits = 1; // } // TODO(umiswing): add shape check FlashAttnFwdParamsV2 params = FlashAttnFwdParamsV2(ctx, batch_size, max_seqlen_q, max_seqlen_k, num_heads, num_heads_k, head_size, dropout, scale, causal, return_softmax, q.dtype(), is_test, rng_name, fixed_seed_offset.get_ptr(), softmax, softmax_lse, seed_offset); VLOG(4) << "FlashAttn fwd seed: " << params.seed << ", offset: " << params.offset; const bool succ = phi::dynload::flash_attn_varlen_fwd( q.data(), k.data(), v.data(), cu_seqlens_q.data(), cu_seqlens_k.data(), params.rng_state.data(), out->data(), params.return_softmax ? softmax->data() : nullptr, softmax_lse->data(), params.batch_size, params.max_seqlen_q, params.max_seqlen_k, params.seqlen_q_rounded, params.seqlen_k_rounded, params.num_heads, params.num_heads_k, params.head_size, params.head_size_rounded, params.dropout, params.scale, params.causal, params.return_softmax, params.is_bf16, stream, params.seed, params.offset); if (!succ) { PADDLE_THROW(phi::errors::External(phi::dynload::flash_attn_error())); } #else PADDLE_THROW(phi::errors::Unimplemented( "FlashAttention is unsupported, please set use_flash_attn to false.")); #endif } template void FlashAttnKernel(const Context& ctx, const DenseTensor& q, const DenseTensor& k, const DenseTensor& v, const paddle::optional& fixed_seed_offset, float dropout, bool causal, bool return_softmax, bool is_test, const std::string& rng_name, DenseTensor* out, DenseTensor* softmax, DenseTensor* softmax_lse, DenseTensor* seed_offset) { #ifdef PADDLE_WITH_FLASHATTN // q,k,v [batch_size, seq_len, num_heads, head_dim] auto dims = q.dims(); PADDLE_ENFORCE_EQ(dims.size(), 4, phi::errors::InvalidArgument( "flash_attn receive input with dim " "[batch_size, seq_len, num_heads, head_dim]")); const int batch_size = dims[0]; const int seqlen_q = dims[1]; const int num_heads = dims[2]; const int head_size = dims[3]; const int seqlen_k = k.dims()[1]; const int num_heads_k = k.dims()[2]; // TODO(umiswing): Add check shape const float scale = 1.0f / std::sqrt(head_size); FlashAttnFwdParamsV2 params = FlashAttnFwdParamsV2(ctx, batch_size, seqlen_q, seqlen_k, num_heads, num_heads_k, head_size, dropout, scale, causal, return_softmax, q.dtype(), is_test, rng_name, fixed_seed_offset.get_ptr(), softmax, softmax_lse, seed_offset); VLOG(4) << "FlashAttn fwd dims q[" << q.dims() << "], k[" << k.dims() << "], v[" << v.dims() << "]"; ctx.template Alloc(out); cudaStream_t stream = ctx.stream(); VLOG(4) << "FlashAttn fwd seed: " << params.seed << ", offset: " << params.offset; bool succ = phi::dynload::flash_attn_fwd( q.data(), k.data(), v.data(), params.rng_state.data(), out->data(), params.return_softmax ? params.softmax->data() : nullptr, params.softmax_lse->data(), params.batch_size, params.max_seqlen_q, params.max_seqlen_k, params.seqlen_q_rounded, params.seqlen_k_rounded, params.num_heads, params.num_heads_k, params.head_size, params.head_size_rounded, params.dropout, params.scale, params.causal, params.return_softmax, params.is_bf16, stream, params.seed, params.offset); PADDLE_ENFORCE_EQ( succ, true, phi::errors::External("Error in Flash-Attention-2, detail information is", phi::dynload::flash_attn_error())); #else PADDLE_THROW(phi::errors::Unimplemented( "FlashAttention is unsupported, please set use_flash_attn to false.")); #endif } } // namespace phi PD_REGISTER_KERNEL(flash_attn_unpadded, GPU, ALL_LAYOUT, phi::FlashAttnUnpaddedKernel, phi::dtype::float16, phi::dtype::bfloat16) { kernel->InputAt(5).SetBackend( phi::Backend::ALL_BACKEND); // fixed_seed_offset } PD_REGISTER_KERNEL(flash_attn, GPU, ALL_LAYOUT, phi::FlashAttnKernel, phi::dtype::float16, phi::dtype::bfloat16) { kernel->InputAt(3).SetBackend( phi::Backend::ALL_BACKEND); // fixed_seed_offset }