diff --git a/custom_ops/gpu_ops/append_attn/get_block_shape_and_split_kv_block.cu b/custom_ops/gpu_ops/append_attn/get_block_shape_and_split_kv_block.cu index a46f427b998..2e05d1885e3 100644 --- a/custom_ops/gpu_ops/append_attn/get_block_shape_and_split_kv_block.cu +++ b/custom_ops/gpu_ops/append_attn/get_block_shape_and_split_kv_block.cu @@ -12,9 +12,11 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include "cute/tensor.hpp" #include "helper.h" #include "paddle/extension.h" #include "paddle/phi/core/memory/memcpy.h" +#include "utils.cuh" template __global__ void @@ -142,6 +144,94 @@ __global__ void split_q_block(const int *__restrict__ seq_lens_q, } } + +template +__global__ void search_chunk_size_for_mla( + const int *__restrict__ seq_lens_q, + const int *__restrict__ seq_lens_encoder, + const int *__restrict__ seq_lens_decoder, + int *__restrict__ num_blocks_x, + int *__restrict__ res_chunk_size, + const int bsz, + const int set_chunk_size, + const int block_size, + const int sm_cout) { + const uint32_t conf_id = threadIdx.x; + int gridx = 0; + if (set_chunk_size > 0 && conf_id == 0) { + for (uint32_t bid = 0; bid < bsz; bid++) { + int seq_len = seq_lens_q[bid]; + int seq_len_encoder = seq_lens_encoder[bid]; + int seq_len_decoder = seq_lens_decoder[bid] + seq_len; + if (seq_len == 0 || seq_len_encoder > 0) continue; + + int loop_times; + loop_times = cute::ceil_div(seq_len_decoder, set_chunk_size); + gridx += loop_times; + } + *num_blocks_x = gridx; + *res_chunk_size = set_chunk_size; + } else if (conf_id < config_size) { + __shared__ int gridx_shared[config_size]; + // chunk_size is a multiple of 64 + const int chunk_size = block_size << conf_id; + for (uint32_t bid = 0; bid < bsz; bid++) { + int seq_len = seq_lens_q[bid]; + int seq_len_encoder = seq_lens_encoder[bid]; + int seq_len_decoder = seq_lens_decoder[bid] + seq_len; + if (seq_len == 0 || seq_len_encoder > 0) continue; + + int loop_times; + loop_times = cute::ceil_div(seq_len_decoder, chunk_size); + gridx += loop_times; + } + gridx_shared[conf_id] = gridx; + __syncthreads(); + if (threadIdx.x == 0) { + uint32_t res_id = 0; + uint32_t max_last_wave_block = 0; + for (uint32_t i = 1; i < config_size; i++) { + uint32_t last_wave_block = gridx_shared[i] % sm_cout; + if (last_wave_block >= max_last_wave_block) { + res_id = i; + max_last_wave_block = last_wave_block; + } + } + *num_blocks_x = gridx_shared[res_id]; + *res_chunk_size = block_size << res_id; + } + } +} + +__global__ void split_block_for_mla(const int *__restrict__ seq_lens_q, + const int *__restrict__ seq_lens_encoder, + const int *__restrict__ seq_lens_decoder, + int *__restrict__ batch_ids, + int *__restrict__ tile_ids_per_batch, + const int bsz, + const int chunk_size) { + if (threadIdx.x == 0) { + int index = 0; + for (uint32_t bid = 0; bid < bsz; bid++) { + int seq_len = seq_lens_q[bid]; + int seq_len_encoder = seq_lens_encoder[bid]; + int seq_len_decoder = seq_lens_decoder[bid] + seq_len; + + if (seq_len == 0) continue; + + int loop_times; + loop_times = cute::ceil_div(seq_len_decoder, chunk_size); + if (seq_len_encoder > 0) { + loop_times = 0; + } + for (uint32_t tile_id = 0; tile_id < loop_times; tile_id++) { + batch_ids[index] = bid; + tile_ids_per_batch[index++] = tile_id; + } + } + } +} + __global__ void split_kv_block(const int *__restrict__ seq_lens_decoder, const int *__restrict__ seq_lens_encoder, int *__restrict__ batch_ids, @@ -230,6 +320,9 @@ std::vector GetBlockShapeAndSplitKVBlock( paddle::Tensor kv_tile_ids_per_batch; paddle::Tensor kv_num_blocks_x_cpu; /*cpu*/ paddle::Tensor max_len_kv_cpu; /*cpu*/ + paddle::Tensor decoder_num_blocks_x; + paddle::Tensor decoder_chunk_size_device; + paddle::Tensor decoder_chunk_size_cpu; /*cpu*/ auto max_len_kv = GetEmptyTensor({1}, paddle::DataType::INT32, seq_lens_decoder.place()); @@ -238,6 +331,87 @@ std::vector GetBlockShapeAndSplitKVBlock( seq_lens_decoder.data(), bsz); max_len_kv_cpu = max_len_kv.copy_to(paddle::CPUPlace(), false); + std::cout << "-----------------------------------------------------------" + << std::endl; + std::cout << "max_dec_len_this_time:================================ " + << max_dec_len_this_time << std::endl; + std::cout << "-----------------------------------------------------------" + << std::endl; + // decoder + if (max_dec_len_this_time > 0) { + const bool mla_use_tensorcore = GetMlaUseTensorcore(); + if (mla_use_tensorcore && group_size <= 64) { + const int set_chunk_size = get_mla_dec_chunk_size(bsz); + decoder_chunk_size_device = GetEmptyTensor( + {1}, paddle::DataType::INT32, seq_lens_encoder.place()); + decoder_num_blocks_x = GetEmptyTensor( + {1}, paddle::DataType::INT32, seq_lens_encoder.place()); + + int device; + cudaGetDevice(&device); + int sm_cout; + cudaDeviceGetAttribute(&sm_cout, cudaDevAttrMultiProcessorCount, device); + constexpr int config_size = + 12; // search space for chunk size:[64, 128, 256, ... 131072] + + search_chunk_size_for_mla + <<<1, 32, 0, stream>>>(seq_lens_this_time.data(), + seq_lens_encoder.data(), + seq_lens_decoder.data(), + decoder_num_blocks_x.data(), + decoder_chunk_size_device.data(), + bsz, + set_chunk_size, + block_size, + sm_cout); + + decoder_chunk_size_cpu = + decoder_chunk_size_device.copy_to(paddle::CPUPlace(), false); + const int chunk_size = decoder_chunk_size_cpu.data()[0]; + std::cout <<"-----------------------------------------------------------" << std::endl; + std::cout << "chunk size1:================================ " << chunk_size << std::endl; + std::cout <<"-----------------------------------------------------------" << std::endl; + split_block_for_mla<<<1, 32, 0, stream>>>( + seq_lens_this_time.data(), + seq_lens_encoder.data(), + seq_lens_decoder.data(), + decoder_batch_ids.data(), + decoder_tile_ids_per_batch.data(), + bsz, + chunk_size); + } else { + const uint32_t decoder_max_tile_size_per_bs_q = + div_up((decoder_step_token_num * group_size), decoder_block_shape_q); + decoder_batch_ids = GetEmptyTensor({bsz * decoder_max_tile_size_per_bs_q}, + paddle::DataType::INT32, + seq_lens_encoder.place()); + decoder_tile_ids_per_batch = + GetEmptyTensor({bsz * decoder_max_tile_size_per_bs_q}, + paddle::DataType::INT32, + seq_lens_encoder.place()); + decoder_num_blocks_x = GetEmptyTensor( + {1}, paddle::DataType::INT32, seq_lens_encoder.place()); + split_q_block<<<1, 32, 0, stream>>>( + seq_lens_this_time.data(), + seq_lens_encoder.data(), + decoder_batch_ids.data(), + decoder_tile_ids_per_batch.data(), + decoder_num_blocks_x.data(), + bsz, + decoder_block_shape_q, + group_size); + + decoder_chunk_size_cpu = paddle::full( + {1}, 131072, paddle::DataType::INT32, paddle::CPUPlace()); + } + } else { + decoder_chunk_size_cpu = + paddle::full({1}, 131072, paddle::DataType::INT32, paddle::CPUPlace()); + decoder_num_blocks_x = + paddle::full({1}, -1, paddle::DataType::INT32, paddle::GPUPlace()); + } + + // encoder if (max_enc_len_this_time > 0) { const uint32_t max_tile_size_per_bs_kv = @@ -321,6 +495,8 @@ std::vector GetBlockShapeAndSplitKVBlock( kv_batch_ids, kv_tile_ids_per_batch, kv_num_blocks_x_cpu, /*cpu*/ + decoder_num_blocks_x, + decoder_chunk_size_cpu, /*cpu*/ max_len_kv_cpu, /*cpu*/ }; } @@ -342,6 +518,8 @@ PD_BUILD_STATIC_OP(get_block_shape_and_split_kv_block) paddle::Optional("kv_batch_ids"), paddle::Optional("kv_tile_ids_per_batch"), paddle::Optional("kv_num_blocks_x_cpu"), + paddle::Optional("decoder_num_blocks_x"), + paddle::Optional("decoder_chunk_size_cpu"), "max_len_kv_cpu" }) .Attrs({ diff --git a/custom_ops/gpu_ops/cpp_extensions.cc b/custom_ops/gpu_ops/cpp_extensions.cc index b0e6e332b3d..6ec29ef6e4c 100644 --- a/custom_ops/gpu_ops/cpp_extensions.cc +++ b/custom_ops/gpu_ops/cpp_extensions.cc @@ -415,6 +415,7 @@ std::vector MultiHeadLatentAttention( const paddle::Tensor& decoder_tile_ids_per_batch, const paddle::Tensor& decoder_num_blocks, const paddle::Tensor& decoder_num_blocks_cpu, + const paddle::Tensor& decoder_chunk_size_cpu, const paddle::Tensor& max_enc_len_this_time, const paddle::Tensor& max_dec_len_this_time, const paddle::Tensor& max_len_kv, diff --git a/custom_ops/gpu_ops/env.h b/custom_ops/gpu_ops/env.h index c7db21ba8f8..33b2ec525dc 100644 --- a/custom_ops/gpu_ops/env.h +++ b/custom_ops/gpu_ops/env.h @@ -62,3 +62,21 @@ inline bool get_mla_use_tensorcore() { mla_use_tensorcore_env == nullptr ? 1 : std::stoul(std::string(mla_use_tensorcore_env)); return mla_use_tensorcore != 0 ? true : false; } + +inline int get_mla_dec_chunk_size(int bsz) { + static const char* mla_dec_chunk_size_env = + std::getenv("FLAGS_mla_dec_chunk_size"); + static const int mla_dec_chunk_size = + mla_dec_chunk_size_env == nullptr + ? -1 + : std::stoi(std::string(mla_dec_chunk_size_env)); + return bsz > 1 ? mla_dec_chunk_size : 64; +} + +inline bool get_mla_use_wg4() { + static const char* mla_use_wg4_env = std::getenv("FLAGS_mla_use_wg4"); + static const uint32_t mla_use_wg4 = + mla_use_wg4_env == nullptr ? 1 + : std::stoul(std::string(mla_use_wg4_env)); + return mla_use_wg4 != 0 ? true : false; +} \ No newline at end of file diff --git a/custom_ops/gpu_ops/helper.h b/custom_ops/gpu_ops/helper.h index 468aff1fc4b..0d4d2c73770 100644 --- a/custom_ops/gpu_ops/helper.h +++ b/custom_ops/gpu_ops/helper.h @@ -557,3 +557,11 @@ inline int GetSMVersion() { return sm_version; } + +inline bool GetMlaUseTensorcore() { + static const bool flags_mla_use_tensorcore = get_mla_use_tensorcore(); + static const bool enable_mla_tensorcore = GetSMVersion() >= 90 ? true : false; + const bool mla_use_tensorcore = + flags_mla_use_tensorcore && enable_mla_tensorcore; + return mla_use_tensorcore; +} diff --git a/custom_ops/gpu_ops/mla_attn/batch_mla_with_paged_kv_cache.cu b/custom_ops/gpu_ops/mla_attn/batch_mla_with_paged_kv_cache.cu index f7d4b8ae27e..942873d954d 100644 --- a/custom_ops/gpu_ops/mla_attn/batch_mla_with_paged_kv_cache.cu +++ b/custom_ops/gpu_ops/mla_attn/batch_mla_with_paged_kv_cache.cu @@ -60,14 +60,21 @@ template void BatchMLAWithPagedKVCacheKernel( const AppendAttnMetaData& meta_data, const paddle::Tensor& q, // [token_num, q_head_num, head_dim] - const paddle::Tensor& latent_cache, // [max_block_num, q_head_num, block_size, head_dim] + const paddle::Tensor& + latent_cache, // [max_block_num, q_head_num, block_size, head_dim] const paddle::optional& attn_mask, - const paddle::optional& cache_k_scale, // [num_kv_heads, head_dim] - const paddle::optional& cache_v_scale, // [num_kv_heads, head_dim] - const paddle::optional& cache_k_zp, // [num_kv_heads, head_dim] - const paddle::optional& cache_v_zp, // [num_kv_heads, head_dim] - const paddle::optional& shift_bias, // [num_kv_heads, head_dim] - const paddle::optional& smooth_weight, // [num_kv_heads, head_dim] + const paddle::optional& + cache_k_scale, // [num_kv_heads, head_dim] + const paddle::optional& + cache_v_scale, // [num_kv_heads, head_dim] + const paddle::optional& + cache_k_zp, // [num_kv_heads, head_dim] + const paddle::optional& + cache_v_zp, // [num_kv_heads, head_dim] + const paddle::optional& + shift_bias, // [num_kv_heads, head_dim] + const paddle::optional& + smooth_weight, // [num_kv_heads, head_dim] const paddle::Tensor& seq_lens_this_time, const paddle::Tensor& seq_lens_decoder, const paddle::Tensor& seq_lens_encoder, @@ -79,13 +86,14 @@ void BatchMLAWithPagedKVCacheKernel( const paddle::Tensor& num_blocks_x_device, const std::string& cache_quant_type_str, const int num_blocks_x, + const int chunk_size, const int max_seq_len, const int max_dec_len, const float softmax_scale, const float quant_max_bound, const float quant_min_bound, const float in_scale, - const int draft_token_num, + const int draft_total_token_num, const bool causal, cudaStream_t& stream, paddle::Tensor* out) { @@ -97,32 +105,45 @@ void BatchMLAWithPagedKVCacheKernel( const auto q_head_num = meta_data.q_num_heads; const auto max_block_num_per_seq = meta_data.max_blocks_per_seq; const auto max_block_num = bsz * max_block_num_per_seq; - const uint32_t chunk_size = get_max_partition_size(bsz); - +// const uint32_t chunk_size = get_max_partition_size(bsz); +#if CUDA_VERSION >= 12080 + constexpr bool USE_REG_EALLOC = true; +#else + constexpr bool USE_REG_EALLOC = false; +#endif int q_head_dim = meta_data.head_dims; int k_head_dim = meta_data.head_dims; int v_head_dim = meta_data.head_dims_v; // int num_chunks = max_dec_len / chunk_size; int num_chunks = div_up(max_dec_len, chunk_size); + std::cout << "-----------------------------------------------------------" + << std::endl; + std::cout << "num_chunks:================================ " << num_chunks + << std::endl; + std::cout << "-----------------------------------------------------------" + << std::endl; auto *allocator = paddle::GetAllocator(q.place()); phi::Allocator::AllocationPtr O_tmp, m_tmp, d_tmp; O_tmp = allocator->Allocate( phi::SizeOf(q.dtype()) * - static_cast(num_chunks * bsz * draft_token_num * q_head_num * v_head_dim)); + static_cast(num_chunks * bsz * draft_total_token_num * q_head_num * v_head_dim)); m_tmp = allocator->Allocate( sizeof(float) * - static_cast(num_chunks * bsz * draft_token_num * q_head_num)); + static_cast(num_chunks * bsz * draft_total_token_num * q_head_num)); d_tmp = allocator->Allocate( sizeof(float) * - static_cast(num_chunks * bsz * draft_token_num * q_head_num)); + static_cast(num_chunks * bsz * draft_total_token_num * q_head_num)); - Params params = {}; - params.Q = reinterpret_cast(const_cast(q.data())); - params.KV = reinterpret_cast(const_cast(latent_cache.data())); - params.O = reinterpret_cast(const_cast(out->data())); - params.O_tmp = reinterpret_cast(O_tmp->ptr()); + using ParamsType = + Params; + ParamsType params = {}; + params.Q = reinterpret_cast(const_cast(q.data())); + params.KV = reinterpret_cast( + const_cast(latent_cache.data())); + params.O = reinterpret_cast(const_cast(out->data())); + params.O_tmp = reinterpret_cast(O_tmp->ptr()); params.m = reinterpret_cast(m_tmp->ptr()); params.d = reinterpret_cast(d_tmp->ptr()); params.block_tables = const_cast(block_tables.data()); @@ -149,15 +170,17 @@ void BatchMLAWithPagedKVCacheKernel( params.qk_head_dim = q_head_dim; params.vo_head_dim = v_head_dim; params.block_size = block_size; - params.max_draft_token_num = draft_token_num; + params.draft_total_token_num = draft_total_token_num; params.sm_scale = softmax_scale; params.chunk_size = chunk_size; params.chunk_num = num_chunks; if (q_head_dim == 576) { - BatchMLAWithPagedKVCacheDispatched<576, 512, NV_TYPE>( - params, stream - ); + BatchMLAWithPagedKVCacheDispatched<576, + 512, + NV_TYPE, + ParamsType, + USE_REG_EALLOC>(params, stream); } else { PD_THROW("error!!! q_head_dim must be 576 !!!\n"); } @@ -166,14 +189,21 @@ void BatchMLAWithPagedKVCacheKernel( template void BatchMLAWithPagedKVCacheKernel( const AppendAttnMetaData& meta_data, const paddle::Tensor& q, // [token_num, q_head_num, head_dim] - const paddle::Tensor& latent_cache, // [max_block_num, q_head_num, block_size, head_dim] + const paddle::Tensor& + latent_cache, // [max_block_num, q_head_num, block_size, head_dim] const paddle::optional& attn_mask, - const paddle::optional& cache_k_scale, // [num_kv_heads, head_dim] - const paddle::optional& cache_v_scale, // [num_kv_heads, head_dim] - const paddle::optional& cache_k_zp, // [num_kv_heads, head_dim] - const paddle::optional& cache_v_zp, // [num_kv_heads, head_dim] - const paddle::optional& shift_bias, // [num_kv_heads, head_dim] - const paddle::optional& smooth_weight, // [num_kv_heads, head_dim] + const paddle::optional& + cache_k_scale, // [num_kv_heads, head_dim] + const paddle::optional& + cache_v_scale, // [num_kv_heads, head_dim] + const paddle::optional& + cache_k_zp, // [num_kv_heads, head_dim] + const paddle::optional& + cache_v_zp, // [num_kv_heads, head_dim] + const paddle::optional& + shift_bias, // [num_kv_heads, head_dim] + const paddle::optional& + smooth_weight, // [num_kv_heads, head_dim] const paddle::Tensor& seq_lens_this_time, const paddle::Tensor& seq_lens_decoder, const paddle::Tensor& seq_lens_encoder, @@ -185,29 +215,36 @@ template void BatchMLAWithPagedKVCacheKernel( const paddle::Tensor& num_blocks_x_device, const std::string& cache_quant_type_str, const int num_blocks_x, + const int chunk_size, const int max_seq_len, const int max_dec_len, const float softmax_scale, const float quant_max_bound, const float quant_min_bound, const float in_scale, - const int draft_token_num, + const int draft_total_token_num, const bool causal, cudaStream_t& stream, paddle::Tensor* out); - template void BatchMLAWithPagedKVCacheKernel( const AppendAttnMetaData& meta_data, const paddle::Tensor& q, // [token_num, q_head_num, head_dim] - const paddle::Tensor& latent_cache, // [max_block_num, q_head_num, block_size, head_dim] + const paddle::Tensor& + latent_cache, // [max_block_num, q_head_num, block_size, head_dim] const paddle::optional& attn_mask, - const paddle::optional& cache_k_scale, // [num_kv_heads, head_dim] - const paddle::optional& cache_v_scale, // [num_kv_heads, head_dim] - const paddle::optional& cache_k_zp, // [num_kv_heads, head_dim] - const paddle::optional& cache_v_zp, // [num_kv_heads, head_dim] - const paddle::optional& shift_bias, // [num_kv_heads, head_dim] - const paddle::optional& smooth_weight, // [num_kv_heads, head_dim] + const paddle::optional& + cache_k_scale, // [num_kv_heads, head_dim] + const paddle::optional& + cache_v_scale, // [num_kv_heads, head_dim] + const paddle::optional& + cache_k_zp, // [num_kv_heads, head_dim] + const paddle::optional& + cache_v_zp, // [num_kv_heads, head_dim] + const paddle::optional& + shift_bias, // [num_kv_heads, head_dim] + const paddle::optional& + smooth_weight, // [num_kv_heads, head_dim] const paddle::Tensor& seq_lens_this_time, const paddle::Tensor& seq_lens_decoder, const paddle::Tensor& seq_lens_encoder, @@ -219,13 +256,14 @@ template void BatchMLAWithPagedKVCacheKernel( const paddle::Tensor& num_blocks_x_device, const std::string& cache_quant_type_str, const int num_blocks_x, + const int chunk_size, const int max_seq_len, const int max_dec_len, const float softmax_scale, const float quant_max_bound, const float quant_min_bound, const float in_scale, - const int draft_token_num, + const int draft_total_token_num, const bool causal, cudaStream_t& stream, paddle::Tensor* out); diff --git a/custom_ops/gpu_ops/mla_attn/batch_mla_with_paged_kv_cache.h b/custom_ops/gpu_ops/mla_attn/batch_mla_with_paged_kv_cache.h index 97fffe39dc3..580bd5757db 100644 --- a/custom_ops/gpu_ops/mla_attn/batch_mla_with_paged_kv_cache.h +++ b/custom_ops/gpu_ops/mla_attn/batch_mla_with_paged_kv_cache.h @@ -37,14 +37,21 @@ template void BatchMLAWithPagedKVCacheKernel( const AppendAttnMetaData& meta_data, const paddle::Tensor& q, // [token_num, q_head_num, head_dim] - const paddle::Tensor& latent_cache, // [max_block_num, q_head_num, block_size, head_dim] + const paddle::Tensor& + latent_cache, // [max_block_num, q_head_num, block_size, head_dim] const paddle::optional& attn_mask, - const paddle::optional& cache_k_scale, // [num_kv_heads, head_dim] - const paddle::optional& cache_v_scale, // [num_kv_heads, head_dim] - const paddle::optional& cache_k_zp, // [num_kv_heads, head_dim] - const paddle::optional& cache_v_zp, // [num_kv_heads, head_dim] - const paddle::optional& shift_bias, // [num_kv_heads, head_dim] - const paddle::optional& smooth_weight, // [num_kv_heads, head_dim] + const paddle::optional& + cache_k_scale, // [num_kv_heads, head_dim] + const paddle::optional& + cache_v_scale, // [num_kv_heads, head_dim] + const paddle::optional& + cache_k_zp, // [num_kv_heads, head_dim] + const paddle::optional& + cache_v_zp, // [num_kv_heads, head_dim] + const paddle::optional& + shift_bias, // [num_kv_heads, head_dim] + const paddle::optional& + smooth_weight, // [num_kv_heads, head_dim] const paddle::Tensor& seq_lens_this_time, const paddle::Tensor& seq_lens_decoder, const paddle::Tensor& seq_lens_encoder, @@ -56,13 +63,14 @@ void BatchMLAWithPagedKVCacheKernel( const paddle::Tensor& num_blocks_x_device, const std::string& cache_quant_type_str, const int num_blocks_x, + const int chunk_size, const int max_seq_len, const int max_dec_len, const float softmax_scale, const float quant_max_bound, const float quant_min_bound, const float in_scale, - const int draft_token_num, + const int draft_total_token_num, const bool causal, cudaStream_t& stream, paddle::Tensor* out); diff --git a/custom_ops/gpu_ops/mla_attn/epilogue.cuh b/custom_ops/gpu_ops/mla_attn/epilogue.cuh index 72d1b557046..33fcf456d07 100644 --- a/custom_ops/gpu_ops/mla_attn/epilogue.cuh +++ b/custom_ops/gpu_ops/mla_attn/epilogue.cuh @@ -130,7 +130,7 @@ struct CollectiveEpilogue { const int tile_idx, const int kv_len, const int chunk_size, - const int max_draft_token_num, + const int draft_total_token_num, const int o_stride_bsz) { const int num_chunks = cute::ceil_div(kv_len, chunk_size); Tensor sO = make_tensor(make_smem_ptr(shared_storage.smem_o.data()), SmemLayoutO{}); @@ -149,7 +149,11 @@ struct CollectiveEpilogue { cutlass::arch::NamedBarrier::sync(NUM_MMA_THREADS, /*id=*/static_cast(NamedBarriers::kValueEmpty)); TiledCopyO gmem_tiled_copy_O; - auto O_ptr = num_chunks == 1 ? epilogue_params.O_ptr + start_token_idx * o_stride_bsz : epilogue_params.O_ptr_tmp + (tile_idx * bsz + bid) * max_draft_token_num * o_stride_bsz; + auto O_ptr = num_chunks == 1 + ? epilogue_params.O_ptr + start_token_idx * o_stride_bsz + : epilogue_params.O_ptr_tmp + (tile_idx * bsz + bid) * + draft_total_token_num * + o_stride_bsz; Tensor mO = make_tensor(make_gmem_ptr(O_ptr), epilogue_params.layout_O); Tensor gO = local_tile(mO, select<0, 1>(TileShape_PDV{}), make_coord(_, _0{}))(_, _, _0{}); Tensor cO = make_identity_tensor(gO.shape()); // (O, D) -> (o_idx, d_idx) diff --git a/custom_ops/gpu_ops/mla_attn/mainloop_load.cuh b/custom_ops/gpu_ops/mla_attn/mainloop_load.cuh index 9c67f601ff0..29447873e49 100644 --- a/custom_ops/gpu_ops/mla_attn/mainloop_load.cuh +++ b/custom_ops/gpu_ops/mla_attn/mainloop_load.cuh @@ -146,7 +146,7 @@ struct CollectiveMainloop { int o_stride_head_num; int chunk_size; int chunk_num; - int max_draft_token_num; + int draft_total_token_num; }; // Device side kernel params @@ -178,7 +178,7 @@ struct CollectiveMainloop { int o_stride_head_num; int chunk_size; int chunk_num; - int max_draft_token_num; + int draft_total_token_num; TMA_KV tma_load_KV; }; @@ -216,7 +216,7 @@ struct CollectiveMainloop { args.o_stride_head_num, args.chunk_size, args.chunk_num, - args.max_draft_token_num, + args.draft_total_token_num, tma_load_KV }; } diff --git a/custom_ops/gpu_ops/mla_attn/mainloop_mma.cuh b/custom_ops/gpu_ops/mla_attn/mainloop_mma.cuh index 77d05958304..d5114d1aada 100644 --- a/custom_ops/gpu_ops/mla_attn/mainloop_mma.cuh +++ b/custom_ops/gpu_ops/mla_attn/mainloop_mma.cuh @@ -188,7 +188,7 @@ CUTLASS_DEVICE void mma_f16(const Params& mainloop_params, if (token_idx < qo_len) { const int head_idx = token_group_idx % Ktraits::GROUP_SIZE; - const int bid_offset = mainloop_params.max_draft_token_num * Ktraits::GROUP_SIZE; + const int bid_offset = mainloop_params.draft_total_token_num * Ktraits::GROUP_SIZE; const int write_idx = bid * bid_offset + token_idx * Ktraits::GROUP_SIZE + head_idx; mM(write_idx) = static_cast(attention_updater.row_max(w_i)); mD(write_idx) = static_cast(attention_updater.row_sum(w_i)); @@ -452,7 +452,7 @@ CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params, if (token_idx < qo_len) { const int head_idx = token_group_idx % Ktraits::GROUP_SIZE; - const int bid_offset = mainloop_params.max_draft_token_num * Ktraits::GROUP_SIZE; + const int bid_offset = mainloop_params.draft_total_token_num * Ktraits::GROUP_SIZE; const int write_idx = bid * bid_offset + token_idx * Ktraits::GROUP_SIZE + head_idx; mM(write_idx) = static_cast(attention_updater.row_max(w_i)); mD(write_idx) = static_cast(attention_updater.row_sum(w_i)); diff --git a/custom_ops/gpu_ops/mla_attn/mla_hopper.cuh b/custom_ops/gpu_ops/mla_attn/mla_hopper.cuh index ba1f4b4470a..5ca4a92f6bd 100644 --- a/custom_ops/gpu_ops/mla_attn/mla_hopper.cuh +++ b/custom_ops/gpu_ops/mla_attn/mla_hopper.cuh @@ -63,8 +63,8 @@ struct Params { alignas(16) DTypeKV *KV; // [max_block_num, block_size, dim_head] alignas(16) DTypeO *O; // [token_num, head_num, dim_head] alignas(16) DTypeO *O_tmp; // [num_chunks, bsz, head_num, dim_head] - alignas(16) float *m; // [num_chunks, bsz * max_draft_token_num * head_num] - alignas(16) float *d; // [num_chunks, bsz * max_draft_token_num * head_num] + alignas(16) float *m; // [num_chunks, bsz * draft_total_token_num * head_num] + alignas(16) float *d; // [num_chunks, bsz * draft_total_token_num * head_num] alignas(16) IdType *block_tables; alignas(16) IdType *seq_lens_this_time; @@ -95,7 +95,7 @@ struct Params { int qk_head_dim; int vo_head_dim; int block_size; - int max_draft_token_num; + int draft_total_token_num; int chunk_size; int chunk_num; int num_blocks_x_int; @@ -368,7 +368,7 @@ MLAWithKVCacheKernel(CUTE_GRID_CONSTANT tile_id, seq_len_decoder_now, mainloop_params.chunk_size, - mainloop_params.max_draft_token_num, + mainloop_params.draft_total_token_num, mainloop_params.o_stride_bsz); } } else { @@ -430,7 +430,7 @@ MLAWithKVCacheKernel(CUTE_GRID_CONSTANT tile_id, seq_len_decoder_now, mainloop_params.chunk_size, - mainloop_params.max_draft_token_num, + mainloop_params.draft_total_token_num, mainloop_params.o_stride_bsz); } } @@ -453,7 +453,7 @@ cudaError_t BatchMLAWithPagedKVCacheKernelTraitsDispatched(Params& params, typename CollectiveMainloop::Params mainloop_params = CollectiveMainloop::to_underlying_arguments({ make_layout(make_shape(KernelTraits::BLOCK_SHAPE_Q, params.qk_head_dim), make_stride(params.qk_head_dim, _1{})), // layout q make_layout(make_shape(params.block_size, params.qk_head_dim, params.max_block_num), make_stride(params.qk_head_dim, _1{}, params.block_size * params.qk_head_dim)), - make_layout(make_shape(params.chunk_num, params.bsz * params.max_draft_token_num * params.q_num_head), make_stride(params.bsz * params.max_draft_token_num * params.q_num_head, _1{})), + make_layout(make_shape(params.chunk_num, params.bsz * params.draft_total_token_num * params.q_num_head), make_stride(params.bsz * params.draft_total_token_num * params.q_num_head, _1{})), params.Q, params.KV, params.m, @@ -478,7 +478,7 @@ cudaError_t BatchMLAWithPagedKVCacheKernelTraitsDispatched(Params& params, params.o_stride_head_num, params.chunk_size, params.chunk_num, - params.max_draft_token_num + params.draft_total_token_num }); typename CollectiveEpilogue::Params epilogue_params = CollectiveEpilogue::to_underlying_arguments_ntma({ params.O, @@ -535,7 +535,7 @@ cudaError_t BatchMLAWithPagedKVCacheKernelTraitsDispatched(Params& params, params.vo_head_dim, params.token_num, params.bsz, - params.max_draft_token_num + params.draft_total_token_num ); } return cudaSuccess; diff --git a/custom_ops/gpu_ops/multi_head_latent_attention.cu b/custom_ops/gpu_ops/multi_head_latent_attention.cu index 98a61e83859..f971eb76d76 100644 --- a/custom_ops/gpu_ops/multi_head_latent_attention.cu +++ b/custom_ops/gpu_ops/multi_head_latent_attention.cu @@ -38,6 +38,7 @@ std::vector MultiHeadLatentAttentionKernel( const paddle::Tensor& decoder_tile_ids_per_batch, const paddle::Tensor& decoder_num_blocks, const paddle::Tensor& decoder_num_blocks_cpu, + const paddle::Tensor& decoder_chunk_size_cpu, const paddle::Tensor& max_enc_len_this_time, const paddle::Tensor& max_dec_len_this_time, const paddle::Tensor& max_len_kv, @@ -58,63 +59,77 @@ std::vector MultiHeadLatentAttentionKernel( const float quant_max_bound, const float quant_min_bound, const float out_linear_in_scale, - const int speculate_max_draft_token_num, + const int speculate_draft_total_token_num, const bool causal, const bool speculate_decoder) { - typedef PDTraits traits_; - typedef typename traits_::data_t data_t; + typedef PDTraits traits_; + typedef typename traits_::data_t data_t; - int decoder_num_blocks_data = decoder_num_blocks_cpu.data()[0]; - int max_dec_len_this_time_data = max_dec_len_this_time.data()[0]; - int max_len_kv_data = max_len_kv.data()[0]; - - const bool mla_use_tensorcore = get_mla_use_tensorcore(); - auto sm_version = GetSMVersion(); - if ((speculate_decoder || mla_use_tensorcore) && sm_version < 90) { - PD_THROW("Please use speculate_decoder=0 and FLAGS_mla_use_tensorcore=0 when sm < 90."); - } - - auto main_stream = query.stream(); + int decoder_num_blocks_data = decoder_num_blocks_cpu.data()[0]; + int max_dec_len_this_time_data = max_dec_len_this_time.data()[0]; + int max_len_kv_data = max_len_kv.data()[0]; + int chunk_size = decoder_chunk_size_cpu.data()[0]; + std::cout << "-----------------------------------------------------------" + << std::endl; + std::cout << "chunk size2:================================ " << chunk_size + << std::endl; + std::cout << "-----------------------------------------------------------" + << std::endl; + const bool mla_use_tensorcore = get_mla_use_tensorcore(); + auto sm_version = GetSMVersion(); + if ((speculate_decoder || mla_use_tensorcore) && sm_version < 90) { + PD_THROW( + "Please use speculate_decoder=0 and FLAGS_mla_use_tensorcore=0 " + "when sm < 90."); + } - paddle::Tensor fmha_out = paddle::full( - {meta_data.token_nums, meta_data.q_num_heads * meta_data.head_dims_v}, - 0, - D, - query.place()); + auto main_stream = query.stream(); + paddle::Tensor fmha_out = paddle::full( + {meta_data.token_nums, meta_data.q_num_heads * meta_data.head_dims_v}, + 0, + D, + query.place()); + std::cout << "-----------------------------------------------------------" + << std::endl; + std::cout << "chunk size2 kernel up:================================ " + << std::endl; + std::cout << "-----------------------------------------------------------" + << std::endl; if (max_dec_len_this_time_data > 0) { if (mla_use_tensorcore) { - BatchMLAWithPagedKVCacheKernel(meta_data, - query, - key_cache, - attn_mask, - cache_k_dequant_scales, - cache_v_dequant_scales, - cache_k_zp, - cache_v_zp, - out_linear_shifts, - out_linear_smooths, - seq_lens_this_time, - seq_lens_decoder, - seq_lens_encoder, - cu_seqlens_q, - batch_id_per_token, - block_tables, - decoder_batch_ids, - decoder_tile_ids_per_batch, - decoder_num_blocks, - cache_quant_type_str, - decoder_num_blocks_data, - max_input_length, - max_len_kv_data, - softmax_scale, - quant_max_bound, - quant_min_bound, - out_linear_in_scale, - speculate_max_draft_token_num, - causal, - main_stream, - &fmha_out); + BatchMLAWithPagedKVCacheKernel(meta_data, + query, + key_cache, + attn_mask, + cache_k_dequant_scales, + cache_v_dequant_scales, + cache_k_zp, + cache_v_zp, + out_linear_shifts, + out_linear_smooths, + seq_lens_this_time, + seq_lens_decoder, + seq_lens_encoder, + cu_seqlens_q, + batch_id_per_token, + block_tables, + decoder_batch_ids, + decoder_tile_ids_per_batch, + decoder_num_blocks, + cache_quant_type_str, + decoder_num_blocks_data, + chunk_size, + max_input_length, + max_len_kv_data, + softmax_scale, + quant_max_bound, + quant_min_bound, + out_linear_in_scale, + speculate_draft_total_token_num, + causal, + main_stream, + &fmha_out); } else { DecodeMLAAttentionKernel( meta_data, @@ -138,6 +153,14 @@ std::vector MultiHeadLatentAttentionKernel( &fmha_out); } } + + std::cout << "-----------------------------------------------------------" + << std::endl; + std::cout << "chunk size2 kernel down:================================ " + << std::endl; + std::cout << "-----------------------------------------------------------" + << std::endl; + return {fmha_out}; } @@ -161,6 +184,7 @@ std::vector MultiHeadLatentAttention( const paddle::Tensor& decoder_tile_ids_per_batch, const paddle::Tensor& decoder_num_blocks, const paddle::Tensor& decoder_num_blocks_cpu, + const paddle::Tensor& decoder_chunk_size_cpu, const paddle::Tensor& max_enc_len_this_time, const paddle::Tensor& max_dec_len_this_time, const paddle::Tensor& max_len_kv, @@ -183,7 +207,7 @@ std::vector MultiHeadLatentAttention( const float quant_max_bound, const float quant_min_bound, const float out_linear_in_scale, - const int speculate_max_draft_token_num, + const int speculate_draft_total_token_num, const bool causal, const bool speculate_decoder) { AppendAttnMetaData meta_data; @@ -224,6 +248,7 @@ std::vector MultiHeadLatentAttention( decoder_tile_ids_per_batch, decoder_num_blocks, decoder_num_blocks_cpu, + decoder_chunk_size_cpu, max_enc_len_this_time, max_dec_len_this_time, max_len_kv, @@ -244,7 +269,7 @@ std::vector MultiHeadLatentAttention( quant_max_bound, quant_min_bound, out_linear_in_scale, - speculate_max_draft_token_num, + speculate_draft_total_token_num, causal, speculate_decoder); } @@ -270,6 +295,7 @@ std::vector MultiHeadLatentAttention( decoder_tile_ids_per_batch, decoder_num_blocks, decoder_num_blocks_cpu, + decoder_chunk_size_cpu, max_enc_len_this_time, max_dec_len_this_time, max_len_kv, @@ -290,7 +316,7 @@ std::vector MultiHeadLatentAttention( quant_max_bound, quant_min_bound, out_linear_in_scale, - speculate_max_draft_token_num, + speculate_draft_total_token_num, causal, speculate_decoder); } @@ -303,114 +329,6 @@ std::vector MultiHeadLatentAttention( } } -std::vector> MultiHeadLatentAttentionInferShape( - const std::vector& query_shape, - const std::vector& key_cache_shape, - const std::vector& value_cache_shape, - const std::vector& seq_lens_encoder_shape, - const std::vector& seq_lens_decoder_shape, - const std::vector& seq_lens_this_time_shape, - const std::vector& cu_seqlens_q_shape, - const std::vector& batch_id_per_token_shape, - const std::vector& block_tables_shape, - const std::vector& encoder_batch_ids_shape, - const std::vector& encoder_tile_ids_per_batch_shape, - const std::vector& encoder_num_blocks_shape, - const std::vector& kv_batch_ids_shape, - const std::vector& kv_tile_ids_per_batch_shape, - const std::vector& kv_num_blocks_shape, - const std::vector& decoder_batch_ids_shape, - const std::vector& decoder_tile_ids_per_batch_shape, - const std::vector& decoder_num_blocks_shape, - const std::vector& decoder_num_blocks_cpu_shape, - const std::vector& max_enc_len_this_time_shape, - const std::vector& max_dec_len_this_time_shape, - const std::vector& max_len_kv_shape, - const paddle::optional>& attn_mask_shape, - const paddle::optional>& query_bias_shape, - const paddle::optional>& query_out_scales_shape, - const paddle::optional>& cache_k_quant_scales_shape, - const paddle::optional>& cache_v_quant_scales_shape, - const paddle::optional>& cache_k_dequant_scales_shape, - const paddle::optional>& cache_v_dequant_scales_shape, - const paddle::optional>& cache_k_zp_shape, - const paddle::optional>& cache_v_zp_shape, - const paddle::optional>& out_linear_shifts_shape, - const paddle::optional>& out_linear_smooths_shape, - const std::string& compute_dtype, - const std::string& cache_quant_type_str, - const int nope_size, - const int max_input_length, - const float softmax_scale, - const float quant_max_bound, - const float quant_min_bound, - const float out_linear_in_scale, - const int speculate_max_draft_token_num, - const bool causal, - const bool speculate_decoder) { - const int token_num = query_shape[0]; - const int kv_num_heads = key_cache_shape[1]; - const int head_dim_qk = key_cache_shape[3]; - const int head_dim_v = nope_size; - const int q_hidden_size = query_shape[query_shape.size() - 1]; - const int num_heads = q_hidden_size / head_dim_qk; - return {{token_num, num_heads * head_dim_v}}; -} - -std::vector MultiHeadLatentAttentionInferDtype( - const paddle::DataType& query_dtype, - const paddle::DataType& key_cache_dtype, - const paddle::DataType& value_cache_dtype, - const paddle::DataType& seq_lens_encoder_dtype, - const paddle::DataType& seq_lens_decoder_dtype, - const paddle::DataType& seq_lens_this_time_dtype, - const paddle::DataType& cu_seqlens_q_dtype, - const paddle::DataType& batch_id_per_token_dtype, - const paddle::DataType& block_tables_dtype, - const paddle::DataType& encoder_batch_ids_dtype, - const paddle::DataType& encoder_tile_ids_per_batch_dtype, - const paddle::DataType& encoder_num_blocks_dtype, - const paddle::DataType& kv_batch_ids_dtype, - const paddle::DataType& kv_tile_ids_per_batch_dtype, - const paddle::DataType& kv_num_blocks_dtype, - const paddle::DataType& decoder_batch_ids_dtype, - const paddle::DataType& decoder_tile_ids_per_batch_dtype, - const paddle::DataType& decoder_num_blocks_dtype, - const paddle::DataType& decoder_num_blocks_cpu_dtype, - const paddle::DataType& max_enc_len_this_time_dtype, - const paddle::DataType& max_dec_len_this_time_dtype, - const paddle::DataType& max_len_kv_dtype, - const paddle::optional& attn_mask_dtype, - const paddle::optional& query_bias_dtype, - const paddle::optional& query_out_scales_dtype, - const paddle::optional& cache_k_quant_scales_dtype, - const paddle::optional& cache_v_quant_scales_dtype, - const paddle::optional& cache_k_dequant_scales_dtype, - const paddle::optional& cache_v_dequant_scales_dtype, - const paddle::optional& cache_k_zp_dtype, - const paddle::optional& cache_v_zp_dtype, - const paddle::optional& out_linear_shifts_dtype, - const paddle::optional& out_linear_smooths_dtype, - const std::string& compute_dtype, - const std::string& cache_quant_type_str, - const int nope_size, - const int max_input_length, - const float softmax_scale, - const float quant_max_bound, - const float quant_min_bound, - const float out_linear_in_scale, - const int speculate_max_draft_token_num, - const bool causal, - const bool speculate_decoder) { - if (compute_dtype == "bf16") { - return {paddle::DataType::BFLOAT16}; - } else if (compute_dtype == "fp16") { - return {paddle::DataType::FLOAT16}; - } else { - PD_THROW("Only supported attr of compute_dtype in ['fp16', 'bf16']."); - } -} - PD_BUILD_STATIC_OP(multi_head_latent_attention) .Inputs({"query", "key_cache", @@ -431,6 +349,7 @@ PD_BUILD_STATIC_OP(multi_head_latent_attention) "decoder_tile_ids_per_batch", "decoder_num_blocks", "decoder_num_blocks_cpu", + "decoder_chunk_size_cpu", "max_enc_len_this_time", "max_dec_len_this_time", "max_len_kv", @@ -454,9 +373,7 @@ PD_BUILD_STATIC_OP(multi_head_latent_attention) "quant_max_bound: float", "quant_min_bound: float", "out_linear_in_scale: float", - "speculate_max_draft_token_num: int", + "speculate_draft_total_token_num: int", "causal: bool", "speculate_decoder: bool"}) - .SetKernelFn(PD_KERNEL(MultiHeadLatentAttention)) - .SetInferShapeFn(PD_INFER_SHAPE(MultiHeadLatentAttentionInferShape)) - .SetInferDtypeFn(PD_INFER_DTYPE(MultiHeadLatentAttentionInferDtype)); + .SetKernelFn(PD_KERNEL(MultiHeadLatentAttention)); diff --git a/fastdeploy/demo/offline_demo.py b/fastdeploy/demo/offline_demo.py index c02bdb45c41..cc05b89e016 100644 --- a/fastdeploy/demo/offline_demo.py +++ b/fastdeploy/demo/offline_demo.py @@ -1,27 +1,52 @@ -""" -# 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. -""" - -from fastdeploy.engine.sampling_params import SamplingParams -from fastdeploy.entrypoints.llm import LLM - -model_name_or_path = "./models/llama-7b" - -# 超参设置 -sampling_params = SamplingParams(temperature=0.1, max_tokens=30) -llm = LLM(model=model_name_or_path, tensor_parallel_size=1) -output = llm.generate(prompts="who are you？", use_tqdm=True, sampling_params=sampling_params) - -print(output) +# """ +# # 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. +# """ +# # export CUTLASS_GEMM_STREAM_K=1 +# # model_name_or_path = "/root/paddlejob/workspace/env_run/output/45t04_wint2_perm_2card_all_safetensors" +# # model_name_or_path = "/cwb/models/DeepSeek-V3-0324" +# # model_name_or_path = "/cwb/ernie-4_5-300b-a47b-bf16-paddle" +# model_name_or_path = "/cwb/models/Qwen3-30B-A3B" +# prompts = ["解析三首李白的诗?"] +# from fastdeploy import LLM, SamplingParams +# sampling_params = SamplingParams(temperature=0.7, top_p=0, max_tokens=128) +# llm = LLM(model=model_name_or_path, tensor_parallel_size=1,enable_custom_all_reduce=False,quantization="wint4", use_cudagraph=False,) +# # outputs = llm.chat(messages=[{"role": "user", "content": "给我三首李白的诗"}]*20, sampling_params=sampling_params,) +# outputs = llm.generate(prompts, sampling_params) + +# print(outputs) + + + + +from fastdeploy import LLM, SamplingParams + +prompts = 20 * ["勒布朗詹姆斯是谁?"] + + +# 采样参数 +sampling_params = SamplingParams(temperature=0.7, top_p=0, max_tokens=100) + +llm = LLM(model="/cwb/DeepSeekV3-0324-5layers", tensor_parallel_size=8,engine_worker_queue_port=8008, +quantization="wint4", max_model_len=32768,max_num_seqs=35,use_cudagraph=False, enable_custom_all_reduce=True) + +# 批量进行推理（llm内部基于资源情况进行请求排队、动态插入处理） +outputs = llm.generate(prompts, sampling_params) +print("==== The outputs are:", outputs) +# 输出结果 +for output in outputs: + prompt = output.prompt + #print("===The prompt is : ", prompt) + generated_text = output.outputs.text + print("===The generated_text is : ", generated_text) \ No newline at end of file diff --git a/fastdeploy/model_executor/layers/attention/mla_attention_backend.py b/fastdeploy/model_executor/layers/attention/mla_attention_backend.py index 5279b68f6f6..4545303294f 100644 --- a/fastdeploy/model_executor/layers/attention/mla_attention_backend.py +++ b/fastdeploy/model_executor/layers/attention/mla_attention_backend.py @@ -24,6 +24,11 @@ import paddle from paddle.nn.functional.flash_attention import flash_attn_unpadded +try: + from paddle.nn.functional.flash_attention import flash_attention_v3_varlen +except: + flash_attention_v3_varlen = None + from fastdeploy.model_executor.layers.attention.ops import ( get_block_shape_and_split_kv_block, init_kv_signal_per_query, @@ -91,6 +96,7 @@ class MLAAttentionBackend(AttentionBackend): """ __infer_dynamic_dims_fields__ = ["attention_metadata"] + flash_attn_func: callable = None attention_metadata: MLAAttentionMetadata def __init__( @@ -147,6 +153,21 @@ def __init__( self.device_id: int = os.getenv("CUDA_VISIBLE_DEVICES", None) self.rank, self.device_id = init_rank_and_device_id(fd_config) + if self.flash_attn_func is None: + prop = paddle.device.cuda.get_device_properties() + cc = prop.major * 10 + prop.minor + is_current_sm_supported = cc >= 90 + is_paddle_supported = any(num >= 90 for num in paddle.version.cuda_archs()) + if is_current_sm_supported and is_paddle_supported: + self.flash_attn_func = flash_attention_v3_varlen + print("The current platform supports Flash Attention V3.") + self.flash_attn_kwargs = {} + else: + self.flash_attn_func = flash_attn_unpadded + self.flash_attn_kwargs = {"scale": self.head_dim**-0.5, "training": False} + print( + "The current platform does not support Flash Attention V3, so Flash Attention V2 will be used instead." + ) def init_attention_metadata(self, forward_meta: ForwardMeta): """Initialize attention metadata hence all layers in the forward pass can reuse it.""" @@ -173,6 +194,8 @@ def init_attention_metadata(self, forward_meta: ForwardMeta): metadata.kv_batch_ids, metadata.kv_tile_ids_per_batch, metadata.kv_num_blocks, + metadata.decoder_num_blocks, + metadata.decoder_chunk_size_cpu, metadata.max_len_kv, ) = get_block_shape_and_split_kv_block( forward_meta.seq_lens_encoder, @@ -269,17 +292,16 @@ def forward_extend( ) # Flash注意力计算 - fmha_out = flash_attn_unpadded( + fmha_out = self.flash_attn_func( q, k, v, forward_meta.cu_seqlens_q, forward_meta.cu_seqlens_k, - metadata.max_enc_len_this_time, - metadata.max_enc_len_this_time, - self.attn_softmax_scale, - causal=True, - training=False, + max_seqlen_q=forward_meta.max_len_tensor_cpu[0], + max_seqlen_k=forward_meta.max_len_tensor_cpu[3], + causal=self.causal, + **self.flash_attn_kwargs, )[0] return fmha_out @@ -310,7 +332,7 @@ def forward_decode( # 获取推测解码参数 speculate_decoder = self.speculative_method is not None - speculate_max_tokens = self.speculate_max_draft_token_num + speculate_max_tokens = self.speculate_max_draft_token_num + 1 # 写入缓存 decode_mla_write_cache( @@ -348,6 +370,7 @@ def forward_decode( forward_meta.decoder_tile_ids_per_batch, forward_meta.decoder_num_blocks_cpu, forward_meta.decoder_num_blocks_cpu, + metadata.decoder_chunk_size_cpu, metadata.max_enc_len_this_time, metadata.max_dec_len_this_time, metadata.max_len_kv, @@ -418,17 +441,16 @@ def forward_mixed( ) # FA - fmha_out = flash_attn_unpadded( + fmha_out = self.flash_attn_func( q, k, v, forward_meta.cu_seqlens_q, forward_meta.cu_seqlens_k, - metadata.max_enc_len_this_time, - metadata.max_enc_len_this_time, - self.attn_softmax_scale, - causal=True, - training=False, + max_seqlen_q=forward_meta.max_len_tensor_cpu[0], + max_seqlen_k=forward_meta.max_len_tensor_cpu[3], + causal=self.causal, + **self.flash_attn_kwargs, )[0] return fmha_out @@ -470,6 +492,7 @@ def forward_mixed( forward_meta.decoder_tile_ids_per_batch, forward_meta.decoder_num_blocks_cpu, forward_meta.decoder_num_blocks_cpu, + metadata.decoder_chunk_size_cpu, metadata.max_enc_len_this_time, metadata.max_dec_len_this_time, metadata.max_len_kv, @@ -496,5 +519,5 @@ def forward_mixed( True, # causal speculate_decoder, ) - + print("mix+++++++++++++++++fmha_out: ", fmha_out) return fmha_out diff --git a/fastdeploy/model_executor/layers/attention/ops/get_block_shape_and_split_kv_block.py b/fastdeploy/model_executor/layers/attention/ops/get_block_shape_and_split_kv_block.py index dd57b525937..6f800b61e6c 100644 --- a/fastdeploy/model_executor/layers/attention/ops/get_block_shape_and_split_kv_block.py +++ b/fastdeploy/model_executor/layers/attention/ops/get_block_shape_and_split_kv_block.py @@ -49,6 +49,8 @@ def get_block_shape_and_split_kv_block( kv_batch_ids, kv_tile_ids_per_batch, kv_num_blocks, + decoder_num_blocks_x, + decoder_chunk_size_cpu, max_len_kv_cpu, ) = get_block_shape_and_split_kv_block_cuda( seq_lens_encoder, @@ -71,6 +73,8 @@ def get_block_shape_and_split_kv_block( kv_batch_ids, kv_tile_ids_per_batch, kv_num_blocks, + decoder_num_blocks_x, + decoder_chunk_size_cpu, max_len_kv_cpu, ) else: diff --git a/fastdeploy/model_executor/models/deepseek_v3.py b/fastdeploy/model_executor/models/deepseek_v3.py index c6a838eda18..70e631858de 100644 --- a/fastdeploy/model_executor/models/deepseek_v3.py +++ b/fastdeploy/model_executor/models/deepseek_v3.py @@ -316,30 +316,23 @@ def forward( mask_encoder_batch: paddle.Tensor, ): """ """ - layernorm_out = hidden_states - fmha_out = paddle.zeros( - shape=[ - layernorm_out.shape[0], - self.num_attention_heads_tp * self.v_head_dim, - ], - dtype=layernorm_out.dtype, - ) - - if forward_meta.max_len_tensor_cpu[1]: # max_enc_len_this_time - query = self.q_a_proj(layernorm_out) - query = self.q_a_layernorm(query) - query = self.q_b_proj(query) + fmha_out = None + # NOTE: (changwenbin) Bring out the public calculation in PD MIX to avoid repeated calculation. + query = self.q_a_proj(hidden_states) + query = self.q_a_layernorm(query) + query = self.q_b_proj(query) + query = query.reshape([-1, self.num_attention_heads_tp, self.qk_head_dim]) + query_nope, query_pe = query.split([self.qk_nope_head_dim, self.qk_rope_head_dim], axis=-1) - query = query.reshape([-1, self.num_attention_heads_tp, self.qk_head_dim]) - query_nope, query_pe = query.split([self.qk_nope_head_dim, self.qk_rope_head_dim], axis=-1) + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + compressed_kv, key_pe = compressed_kv.split([self.kv_lora_rank, self.qk_rope_head_dim], axis=-1) + key_pe = key_pe.reshape([-1, 1, self.qk_rope_head_dim]) + compressed_kv = self.kv_a_layernorm(compressed_kv) - compressed_kv = self.kv_a_proj_with_mqa(layernorm_out) - compressed_kv, key_pe = compressed_kv.split([self.kv_lora_rank, self.qk_rope_head_dim], axis=-1) - key_pe = key_pe.reshape([-1, 1, self.qk_rope_head_dim]) - compressed_kv = self.kv_a_layernorm(compressed_kv) - - query_pe, key_pe = self.rotary_emb(position_ids, query_pe, key_pe) + query_pe, key_pe = self.rotary_emb(position_ids, query_pe, key_pe) + if forward_meta.max_len_tensor_cpu[1]: # max_enc_len_this_time + # NOTE: (changwenbin) We will take the public part key_value = self.kv_b_proj(compressed_kv) key_value = key_value.reshape( [ @@ -371,23 +364,10 @@ def forward( fmha_out_prefill = fmha_out_prefill.reshape([-1, self.num_attention_heads_tp * self.v_head_dim]) fmha_out_prefill = fmha_out_prefill * mask_encoder_batch.cast(fmha_out_prefill.dtype) - fmha_out = fmha_out + fmha_out_prefill - if forward_meta.max_len_tensor_cpu[2]: # max_dec_len_this_time - query = self.q_a_proj(layernorm_out) - query = self.q_a_layernorm(query) - ln_out_or_q_c = query - - compressed_kv = self.kv_a_proj_with_mqa(layernorm_out) - compressed_kv, key_pe = compressed_kv.split([self.kv_lora_rank, self.qk_rope_head_dim], axis=-1) - key_pe = key_pe.reshape([-1, 1, self.qk_rope_head_dim]) - compressed_kv = self.kv_a_layernorm(compressed_kv) - - query = self.q_b_proj(ln_out_or_q_c) - query = query.reshape([-1, self.num_attention_heads_tp, self.qk_head_dim]) - - query_nope, query_pe = query.split([self.qk_nope_head_dim, self.qk_rope_head_dim], axis=-1) - query_pe, key_pe = self.rotary_emb(position_ids, query_pe, key_pe) + fmha_out = fmha_out_prefill + if forward_meta.max_len_tensor_cpu[2]: # max_dec_len_this_time + # NOTE: (changwenbin) We will take the public part q_nope_out = self.kv_b_proj_bmm(query_nope.transpose([1, 0, 2]), proj_type="k").transpose([1, 0, 2]) q_input = paddle.concat([q_nope_out, query_pe], axis=-1) @@ -416,7 +396,10 @@ def forward( .transpose([1, 0, 2]) .reshape([-1, self.num_attention_heads_tp * self.v_head_dim]) ) - fmha_out = fmha_out + fmha_out_decode + if fmha_out is None: + fmha_out = fmha_out_decode + else: + fmha_out = fmha_out + fmha_out_decode output = self.o_proj(fmha_out) return output