[XPU] support absorb mla, separate the model networking from GPU (#10353)
* 独立的，不和 GPU 混用的组网，现在还缺算子迁移过来 * update csrc * fix * fix * 移除多余代码 * support infer param * fix save output * add warm script * 组网剥离 * fix style * update setup and remove useless files * fused block attn * hack 实现 int_wo MOE * 多线程 mmlu 脚本 * mmlu 8K 输出 * add inference time for benchmark * fix warm max_dec_len * Revert "mmlu 8K 输出" This reverts commit0c89ea0fa7. * Revert "多线程 mmlu 脚本" This reverts commitfe6bfc5a3b. * add high speed moe * change rope cos sin to cpu * fix save output
This commit is contained in:
yulangz
2
.gitignore
vendored
2
.gitignore
vendored
@@ -140,4 +140,4 @@ autogen/
|
||||
#fp8
|
||||
ops/csrc/fp8/deep_gemm/include/cutlass
|
||||
ops/csrc/fp8/deep_gemm/include/cute
|
||||
|
||||
.ccls-cache
|
||||
|
||||
219
csrc/xpu/src/absorb_mla_block_mha_decoder_xpu.cc
Normal file
219
csrc/xpu/src/absorb_mla_block_mha_decoder_xpu.cc
Normal file
@@ -0,0 +1,219 @@
|
||||
// Copyright (c) 2024 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/backends/xpu/xpu_context.h>
|
||||
|
||||
#include "paddle/extension.h"
|
||||
#include "paddle/phi/core/enforce.h"
|
||||
#include "xpu/plugin.h"
|
||||
#include <core/ctx_manager.h>
|
||||
#include <core/xft_check.h>
|
||||
#include <core/xft_event.h>
|
||||
#include <core/xft_params.h>
|
||||
#include <xft/xdnn_plugin.h>
|
||||
#include <xft/operation/page_attn.h>
|
||||
#include <xft/operation/fmha.h>
|
||||
#include <flash_api.h> // link xfa
|
||||
#include "ops.h"
|
||||
|
||||
namespace xftkernel = baidu::xpu::xftkernel;
|
||||
|
||||
template <typename T>
|
||||
struct kl3_pa_TL_trait {
|
||||
using TL = T;
|
||||
};
|
||||
template <>
|
||||
struct kl3_pa_TL_trait<bfloat16> {
|
||||
using TL = float;
|
||||
};
|
||||
std::vector<paddle::Tensor> MlaDeAttn(
|
||||
const paddle::Tensor& q,
|
||||
const paddle::Tensor& kv_cache,
|
||||
const paddle::Tensor& decoder_context_len,
|
||||
const paddle::Tensor& decoder_batch_map,
|
||||
const paddle::Tensor& decoder_context_len_cpu,
|
||||
const paddle::Tensor& decoder_batch_map_cpu,
|
||||
const paddle::Tensor& dec_batch_tensor,
|
||||
const paddle::Tensor& padding_offsets,
|
||||
const paddle::Tensor& cum_offsets,
|
||||
const paddle::Tensor& block_tables,
|
||||
const float softmax_scale,
|
||||
const int block_size,
|
||||
const int num_head,
|
||||
const int kv_lora_rank,
|
||||
const int rope_head_dim,
|
||||
const int dim_qk,
|
||||
const int dim_v) {
|
||||
baidu::xpu::api::plugin::print_times("[TIME BEGIN] MlaDeAttn" );
|
||||
phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
|
||||
auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
|
||||
auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
|
||||
xpu::ctx_guard RAII_GUARD(xpu_ctx->x_context());
|
||||
|
||||
using QType = typename XPUTypeTrait<bfloat16>::Type;
|
||||
using CacheType = typename XPUTypeTrait<bfloat16>::Type;
|
||||
typedef paddle::bfloat16 qdata_t, cache_t;
|
||||
const auto& input_dims = q.dims();
|
||||
const int bsz = cum_offsets.dims()[0];
|
||||
const int token_num = input_dims[0];
|
||||
const int block_batch = block_tables.dims()[0]; // TODO参数含义 block_batch_ PageParam page_param_
|
||||
const int max_block_per_seq = block_tables.dims()[1];
|
||||
const int max_seq_len = block_size * max_block_per_seq;
|
||||
int dec_batch = dec_batch_tensor.data<int32_t>()[0];
|
||||
// 初始化输入：q k v
|
||||
auto q_xft = baidu::xpu::xft::xftTensor<QType, 3>(
|
||||
reinterpret_cast<QType*>(const_cast<paddle::bfloat16*>(q.data<qdata_t>())),
|
||||
std::array<int64_t, 3>{q.shape()[0],
|
||||
q.shape()[1],
|
||||
q.shape()[2]});
|
||||
// 初始化输入：k cache
|
||||
auto kv_cache_xft = baidu::xpu::xft::xftTensor<CacheType, 4>(
|
||||
reinterpret_cast<CacheType*>(const_cast<paddle::bfloat16*>(kv_cache.data<cache_t>())),
|
||||
std::array<int64_t, 4>{kv_cache.shape()[0],
|
||||
kv_cache.shape()[1],
|
||||
kv_cache.shape()[2],
|
||||
kv_cache.shape()[3]});
|
||||
// 初始化输入：block table
|
||||
auto block_tables_xft = baidu::xpu::xft::xftTensor<int, 2>(
|
||||
reinterpret_cast<int*>(const_cast<int*>(block_tables.data<int>())),
|
||||
std::array<int64_t, 2>{block_tables.shape()[0],
|
||||
block_tables.shape()[1]});
|
||||
// 初始化输出tensor
|
||||
auto fmha_out = paddle::full({q.shape()[0], num_head * kv_lora_rank}, -2, q.type(), q.place());
|
||||
auto fmha_out_xft = baidu::xpu::xft::xftTensor<QType, 2>(
|
||||
reinterpret_cast<QType*>(const_cast<paddle::bfloat16*>(fmha_out.data<qdata_t>())),
|
||||
std::array<int64_t, 2>{fmha_out.shape()[0],
|
||||
fmha_out.shape()[1]});
|
||||
|
||||
// decoder
|
||||
if(dec_batch > 0){
|
||||
// context_len
|
||||
baidu::xpu::api::VectorParam<int32_t> context_len_vp{const_cast<int32_t*>(decoder_context_len_cpu.data<int32_t>()), dec_batch, const_cast<int32_t*>(decoder_context_len.data<int32_t>())};
|
||||
// real batch
|
||||
baidu::xpu::api::VectorParam<int32_t> valid_batch_vp{const_cast<int32_t*>(decoder_batch_map_cpu.data<int32_t>()), dec_batch, const_cast<int32_t*>(decoder_batch_map.data<int32_t>())};
|
||||
|
||||
// multi_latent_attention
|
||||
using TQ = bfloat16;
|
||||
using TKVCACHE = bfloat16;
|
||||
using TO = TQ;
|
||||
using TGEMM = float;
|
||||
using TEW = float;
|
||||
using TID = int;
|
||||
constexpr int quant_mode = 0;
|
||||
// xpu_ctx->x_context().set_debug_level(0xa1);
|
||||
int ret = baidu::xpu::xfa::multi_latent_attention<
|
||||
TQ,
|
||||
TKVCACHE,
|
||||
TO,
|
||||
TGEMM,
|
||||
TEW,
|
||||
TID,
|
||||
quant_mode>(
|
||||
xpu_ctx->x_context(),
|
||||
fmha_out_xft.data(),
|
||||
q_xft.data(),
|
||||
kv_cache_xft.data(),
|
||||
block_tables_xft.data(),
|
||||
context_len_vp,
|
||||
valid_batch_vp,
|
||||
block_batch,
|
||||
max_seq_len,
|
||||
num_head,
|
||||
kv_lora_rank,
|
||||
rope_head_dim,
|
||||
nullptr, // attn_mask
|
||||
softmax_scale, // 0.13523377478122711f, // scale
|
||||
block_size,
|
||||
max_block_per_seq,
|
||||
-1,
|
||||
nullptr,
|
||||
nullptr,
|
||||
nullptr);
|
||||
}
|
||||
baidu::xpu::api::plugin::print_times("[TIME END] MlaDeAttn");
|
||||
return {fmha_out};
|
||||
}
|
||||
|
||||
std::vector<std::vector<int64_t>> MlaDeAttnInferShape(
|
||||
const std::vector<int64_t>& q_shape,
|
||||
const std::vector<int64_t>& kv_cache_shape,
|
||||
const std::vector<int64_t>& decoder_context_len_shape,
|
||||
const std::vector<int64_t>& decoder_batch_map_shape,
|
||||
const std::vector<int64_t>& decoder_context_len_cpu_shape,
|
||||
const std::vector<int64_t>& decoder_batch_map_cpu_shape,
|
||||
const std::vector<int64_t>& dec_batch_tensor_shape,
|
||||
const std::vector<int64_t>& padding_offsets_shape,
|
||||
const std::vector<int64_t>& cum_offsets_shape,
|
||||
const std::vector<int64_t>& block_tables_shape,
|
||||
const float softmax_scale,
|
||||
const int block_size,
|
||||
const int num_head,
|
||||
const int kv_lora_rank,
|
||||
const int rope_head_dim,
|
||||
const int dim_qk,
|
||||
const int dim_v) {
|
||||
return {{q_shape[0], num_head * kv_lora_rank}};
|
||||
}
|
||||
|
||||
std::vector<paddle::DataType> MlaDeAttnInferDtype(
|
||||
const paddle::DataType& q_dtype,
|
||||
const paddle::DataType& kv_cache_dtype,
|
||||
const paddle::DataType& decoder_context_len_dtype,
|
||||
const paddle::DataType& decoder_batch_map_dtype,
|
||||
const paddle::DataType& decoder_context_len_cpu_dtype,
|
||||
const paddle::DataType& decoder_batch_map_cpu_dtype,
|
||||
const paddle::DataType& dec_batch_tensor_dtype,
|
||||
const paddle::DataType& padding_offsets_dtype,
|
||||
const paddle::DataType& cum_offsets_dtype,
|
||||
const paddle::DataType& block_tables_dtype,
|
||||
const float softmax_scale,
|
||||
const int block_size,
|
||||
const int num_head,
|
||||
const int kv_lora_rank,
|
||||
const int rope_head_dim,
|
||||
const int dim_qk,
|
||||
const int dim_v) {
|
||||
if (q_dtype == paddle::DataType::FLOAT16) {
|
||||
return {paddle::DataType::FLOAT16};
|
||||
} else if(q_dtype == paddle::DataType::BFLOAT16){
|
||||
return {paddle::DataType::BFLOAT16};
|
||||
}
|
||||
else {
|
||||
PD_THROW("Only supported attr of compute_dtype in ['fp16','bfp16'].");
|
||||
}
|
||||
}
|
||||
|
||||
PD_BUILD_OP(absorb_mla_block_mha_decoder_xpu)
|
||||
.Inputs({"q",
|
||||
"kv_cache",
|
||||
"decoder_context_len",
|
||||
"decoder_batch_map",
|
||||
"decoder_context_len_cpu",
|
||||
"decoder_batch_map_cpu",
|
||||
"dec_batch_tensor",
|
||||
"padding_offsets",
|
||||
"cum_offsets",
|
||||
"block_tables"})
|
||||
.Outputs({"fmha_out"})
|
||||
.Attrs({"softmax_scale: float",
|
||||
"block_size: int",
|
||||
"num_head: int",
|
||||
"kv_lora_rank: int",
|
||||
"rope_head_dim: int",
|
||||
"dim_qk: int",
|
||||
"dim_v: int"})
|
||||
.SetKernelFn(PD_KERNEL(MlaDeAttn))
|
||||
.SetInferShapeFn(PD_INFER_SHAPE(MlaDeAttnInferShape))
|
||||
.SetInferDtypeFn(PD_INFER_DTYPE(MlaDeAttnInferDtype));
|
||||
|
||||
251
csrc/xpu/src/absorb_mla_block_mha_encoder_xpu.cc
Normal file
251
csrc/xpu/src/absorb_mla_block_mha_encoder_xpu.cc
Normal file
@@ -0,0 +1,251 @@
|
||||
// Copyright (c) 2024 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/backends/xpu/xpu_context.h>
|
||||
|
||||
#include "paddle/extension.h"
|
||||
#include "paddle/phi/core/enforce.h"
|
||||
#include "xpu/plugin.h"
|
||||
#include <core/ctx_manager.h>
|
||||
#include <core/xft_check.h>
|
||||
#include <core/xft_event.h>
|
||||
#include <core/xft_params.h>
|
||||
#include <xft/xdnn_plugin.h>
|
||||
#include <xft/operation/page_attn.h>
|
||||
#include <xft/operation/fmha.h>
|
||||
#include <flash_api.h> // link xfa
|
||||
#include "ops.h"
|
||||
|
||||
namespace xftkernel = baidu::xpu::xftkernel;
|
||||
|
||||
template <typename T>
|
||||
struct kl3_pa_TL_trait {
|
||||
using TL = T;
|
||||
};
|
||||
template <>
|
||||
struct kl3_pa_TL_trait<bfloat16> {
|
||||
using TL = float;
|
||||
};
|
||||
std::vector<paddle::Tensor> MlaEnAttn(
|
||||
const paddle::Tensor& q,
|
||||
const paddle::Tensor& k,
|
||||
const paddle::Tensor& v,
|
||||
const paddle::Tensor& encoder_seq_lod,
|
||||
const paddle::Tensor& encoder_batch_map,
|
||||
const paddle::Tensor& encoder_seq_lod_cpu,
|
||||
const paddle::Tensor& encoder_batch_map_cpu,
|
||||
const paddle::Tensor& enc_batch_tensor,
|
||||
const paddle::Tensor& padding_offsets,
|
||||
const paddle::Tensor& cum_offsets,
|
||||
const paddle::Tensor& block_tables,
|
||||
const paddle::Tensor& max_enc_len_this_time,
|
||||
const paddle::Tensor& max_dec_len_this_time,
|
||||
const float softmax_scale,
|
||||
const int block_size,
|
||||
const int num_head,
|
||||
const int dim_qk,
|
||||
const int dim_v) {
|
||||
|
||||
baidu::xpu::api::plugin::print_times("[TIME BEGIN] MlaEnAttn");
|
||||
phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
|
||||
auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
|
||||
auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
|
||||
|
||||
xpu::ctx_guard RAII_GUARD(xpu_ctx->x_context());
|
||||
using QType = typename XPUTypeTrait<bfloat16>::Type;
|
||||
using CacheType = typename XPUTypeTrait<bfloat16>::Type;
|
||||
typedef paddle::bfloat16 qdata_t, cache_t;
|
||||
const auto& input_dims = q.dims();
|
||||
const int bsz = cum_offsets.dims()[0];
|
||||
const int token_num = input_dims[0];
|
||||
const int block_batch = block_tables.dims()[0]; // TODO参数含义 block_batch_ PageParam page_param_
|
||||
const int max_block_per_seq = block_tables.dims()[1];
|
||||
const int max_seq_len = block_size * max_block_per_seq;
|
||||
int enc_batch = enc_batch_tensor.data<int32_t>()[0];
|
||||
// 初始化输入：q k v
|
||||
auto q_xft = baidu::xpu::xft::xftTensor<QType, 3>(
|
||||
reinterpret_cast<QType*>(const_cast<paddle::bfloat16*>(q.data<qdata_t>())),
|
||||
std::array<int64_t, 3>{q.shape()[0],
|
||||
q.shape()[1],
|
||||
q.shape()[2]});
|
||||
auto k_xft = baidu::xpu::xft::xftTensor<QType, 3>(
|
||||
reinterpret_cast<QType*>(const_cast<paddle::bfloat16*>(k.data<qdata_t>())),
|
||||
std::array<int64_t, 3>{k.shape()[0],
|
||||
k.shape()[1],
|
||||
k.shape()[2]});
|
||||
auto v_xft = baidu::xpu::xft::xftTensor<QType, 3>(
|
||||
reinterpret_cast<QType*>(const_cast<paddle::bfloat16*>(v.data<qdata_t>())),
|
||||
std::array<int64_t, 3>{v.shape()[0],
|
||||
v.shape()[1],
|
||||
v.shape()[2]});
|
||||
|
||||
// 初始化输出tensor
|
||||
auto fmha_out = paddle::full({q.shape()[0], num_head * dim_v}, -2, q.type(), q.place());
|
||||
auto fmha_out_xft = baidu::xpu::xft::xftTensor<QType, 2>(
|
||||
reinterpret_cast<QType*>(const_cast<paddle::bfloat16*>(fmha_out.data<qdata_t>())),
|
||||
std::array<int64_t, 2>{fmha_out.shape()[0],
|
||||
fmha_out.shape()[1]});
|
||||
|
||||
// encoder
|
||||
if(max_enc_len_this_time.data<int>()[0] > 0){
|
||||
// q_lod
|
||||
baidu::xpu::api::VectorParam<int32_t> context_len_vp{const_cast<int32_t*>(encoder_seq_lod_cpu.data<int32_t>()), enc_batch + 1, const_cast<int32_t*>(encoder_seq_lod.data<int32_t>())};
|
||||
// real batch（encoder阶段 kv cache写需要）
|
||||
baidu::xpu::api::VectorParam<int32_t> valid_batch_vp{const_cast<int32_t*>(encoder_batch_map_cpu.data<int32_t>()), enc_batch, const_cast<int32_t*>(encoder_batch_map.data<int32_t>())};
|
||||
// prefix (not support)
|
||||
baidu::xpu::api::VectorParam<int32_t> prefix_lens_vp = baidu::xpu::api::VectorParam<int32_t>();
|
||||
// kv_lod (非prefix cache情况与q_lod一致)
|
||||
baidu::xpu::api::VectorParam<int32_t> encoder_kv_lods_vp{const_cast<int32_t*>(encoder_seq_lod_cpu.data<int32_t>()), enc_batch + 1, const_cast<int32_t*>(encoder_seq_lod.data<int32_t>())};
|
||||
// page_param_
|
||||
baidu::xpu::xft::PageParam page_param_(block_batch, // block_batch_
|
||||
max_enc_len_this_time.data<int>()[0], // max_enc_len_this_time.data<int>()[0],max_seq_len // max_context_len_
|
||||
block_size, // block_size_
|
||||
max_block_per_seq, // max_num_blocks_per_seq_
|
||||
false); // v_trans
|
||||
// attn_param
|
||||
baidu::xpu::xft::PageAttnParam attn_param_(dim_qk, // head_dim(q,k)
|
||||
num_head, // q_head_num
|
||||
num_head, // kv_head_num
|
||||
false, // vp_lod_flag
|
||||
context_len_vp, // context_len_vp
|
||||
valid_batch_vp, // valid_batch_vp
|
||||
prefix_lens_vp, // prefix_lens_vp
|
||||
encoder_kv_lods_vp); // encoder_kv_lods_vp
|
||||
// fmha op
|
||||
using FMHA_Type = typename baidu::xpu::xft::FMHA_QBF16_KVBF16;
|
||||
using TQ = typename FMHA_Type::Qtype;
|
||||
using TK = typename FMHA_Type::Ktype;
|
||||
using TV = typename FMHA_Type::Vtype;
|
||||
using TBIAS = typename FMHA_Type::Btype;
|
||||
using TO = typename FMHA_Type::Otype;
|
||||
using TMASK = typename FMHA_Type::Mtype;
|
||||
int ret = baidu::xpu::xfa::flash_attention_context_vllm<
|
||||
TQ,
|
||||
TK,
|
||||
TV,
|
||||
TBIAS,
|
||||
TO,
|
||||
float, // T_GEMM0
|
||||
float, // T_GEMM1
|
||||
float, // T_EW
|
||||
int,
|
||||
0>(
|
||||
xpu_ctx->x_context(),
|
||||
q_xft.data(), // q
|
||||
std::is_same<TK, int8_t>::value ? nullptr : reinterpret_cast<TK*>(k_xft.data()), // k
|
||||
std::is_same<TK, int8_t>::value ? nullptr : reinterpret_cast<TV*>(v_xft.data()), // v
|
||||
nullptr, // p_bias == nullptr ? nullptr : p_bias->data(), // mask_bias
|
||||
fmha_out_xft.data(), //o
|
||||
enc_batch, // batch_size
|
||||
page_param_.max_context_len_, // max_seq_q
|
||||
page_param_.max_context_len_, // max_seq_kv, （prefix cache是会不一样）
|
||||
attn_param_.q_head_num_,
|
||||
attn_param_.head_dim_,
|
||||
attn_param_.kv_head_num_,
|
||||
nullptr,
|
||||
attn_param_.context_len_vp_,
|
||||
attn_param_.encoder_kv_lods_vp_,
|
||||
{}, // mask_shape_
|
||||
{}, // alibi_slopes_shape
|
||||
true, // is_causal_mask_
|
||||
false, // is_qkv_fusion_
|
||||
0x0010, // param.qkv_layout_ = AttnQKVLayout_t::ATTN_BLHD
|
||||
0x10, // vsl_flag_ = VslConverter::FMHA_LVSL
|
||||
nullptr, // q_maxptr
|
||||
nullptr, // k_maxptr
|
||||
nullptr, // v_maxptr
|
||||
nullptr,
|
||||
dim_v, // vo_head_dim
|
||||
nullptr, // p_cache_k->data(),
|
||||
nullptr, // p_cache_v->data(),
|
||||
nullptr, // std::is_same<TK, int8_t>::value ? p_kcache_perhead_scale->data() : p_cache_k->max_data(),
|
||||
nullptr, // std::is_same<TK, int8_t>::value ? p_vcache_perhead_scale->data() : p_cache_v->max_data(),
|
||||
block_size, // block_size
|
||||
max_block_per_seq, // max_blocks_per_seq (prefix cache)
|
||||
page_param_.max_context_len_, // prefill_len
|
||||
nullptr,
|
||||
softmax_scale * sqrt(dim_qk)); // block_tables (prefix cache)
|
||||
|
||||
// std::cout << "fmha kernel done " <<std::endl;
|
||||
}
|
||||
baidu::xpu::api::plugin::print_times("[TIME END] MlaEnAttn:");
|
||||
return {fmha_out};
|
||||
}
|
||||
|
||||
std::vector<std::vector<int64_t>> MlaEnAttnInferShape(
|
||||
const std::vector<int64_t>& q_shape,
|
||||
const std::vector<int64_t>& k_shape,
|
||||
const std::vector<int64_t>& v_shape,
|
||||
const std::vector<int64_t>& encoder_seq_lod_shape,
|
||||
const std::vector<int64_t>& encoder_batch_map_shape,
|
||||
const std::vector<int64_t>& encoder_seq_lod_cpu_shape,
|
||||
const std::vector<int64_t>& encoder_batch_map_cpu_shape,
|
||||
const std::vector<int64_t>& enc_batch_tensor_shape,
|
||||
const std::vector<int64_t>& padding_offsets_shape,
|
||||
const std::vector<int64_t>& cum_offsets_shape,
|
||||
const std::vector<int64_t>& block_tables_shape,
|
||||
const std::vector<int64_t>& max_enc_len_this_time_shape,
|
||||
const std::vector<int64_t>& max_dec_len_this_time_shape) {
|
||||
|
||||
return {v_shape};
|
||||
}
|
||||
|
||||
std::vector<paddle::DataType> MlaEnAttnInferDtype(
|
||||
const paddle::DataType& q_dtype,
|
||||
const paddle::DataType& k_dtype,
|
||||
const paddle::DataType& v_dtype,
|
||||
const paddle::DataType& encoder_seq_lod_dtype,
|
||||
const paddle::DataType& encoder_batch_map_dtype,
|
||||
const paddle::DataType& encoder_seq_lod_cpu_dtype,
|
||||
const paddle::DataType& encoder_batch_map_cpu_dtype,
|
||||
const paddle::DataType& enc_batch_tensor_dtype,
|
||||
const paddle::DataType& padding_offsets_dtype,
|
||||
const paddle::DataType& cum_offsets_dtype,
|
||||
const paddle::DataType& block_tables_dtype,
|
||||
const paddle::DataType& max_enc_len_this_time_dtype,
|
||||
const paddle::DataType& max_dec_len_this_time_dtype) {
|
||||
if (q_dtype == paddle::DataType::FLOAT16) {
|
||||
return {paddle::DataType::FLOAT16};
|
||||
} else if(q_dtype == paddle::DataType::BFLOAT16){
|
||||
return {paddle::DataType::BFLOAT16};
|
||||
}
|
||||
else {
|
||||
PD_THROW("Only supported attr of compute_dtype in ['fp16','bfp16'].");
|
||||
}
|
||||
}
|
||||
|
||||
PD_BUILD_OP(absorb_mla_block_mha_encoder_xpu)
|
||||
.Inputs({"q",
|
||||
"k",
|
||||
"v",
|
||||
"encoder_seq_lod",
|
||||
"encoder_batch_map",
|
||||
"encoder_seq_lod_cpu",
|
||||
"encoder_batch_map_cpu",
|
||||
"enc_batch_tensor",
|
||||
"padding_offsets",
|
||||
"cum_offsets",
|
||||
"block_tables",
|
||||
"max_enc_len_this_time",
|
||||
"max_dec_len_this_time"})
|
||||
.Outputs({"fmha_out"})
|
||||
.Attrs({"softmax_scale: float",
|
||||
"block_size: int",
|
||||
"num_head: int",
|
||||
"dim_qk: int",
|
||||
"dim_v: int"})
|
||||
.SetKernelFn(PD_KERNEL(MlaEnAttn))
|
||||
.SetInferShapeFn(PD_INFER_SHAPE(MlaEnAttnInferShape))
|
||||
.SetInferDtypeFn(PD_INFER_DTYPE(MlaEnAttnInferDtype));
|
||||
|
||||
@@ -19,63 +19,35 @@
|
||||
namespace xftkernel = baidu::xpu::xftkernel;
|
||||
std::vector<paddle::Tensor> AdjustBatch(const paddle::Tensor& tmp_out, // [token_num, dim_embed]
|
||||
const paddle::Tensor& cum_offsets, // [bsz, 1]
|
||||
const paddle::Tensor& seq_lens_decoder,
|
||||
const paddle::Tensor& seq_lens_encoder,
|
||||
const paddle::Tensor& encoder_seq_lod,
|
||||
const paddle::Tensor& encoder_batch_idx,
|
||||
const paddle::Tensor& decoder_batch_idx,
|
||||
const paddle::Tensor& encoder_seq_lod_cpu,
|
||||
const paddle::Tensor& encoder_batch_idx_cpu,
|
||||
const paddle::Tensor& decoder_batch_idx_cpu,
|
||||
const paddle::Tensor& enc_batch_tensor,
|
||||
const paddle::Tensor& dec_batch_tensor,
|
||||
const paddle::optional<paddle::Tensor>& output_padding_offset,
|
||||
int max_input_length) {
|
||||
|
||||
phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
|
||||
auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
|
||||
auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
|
||||
xpu::ctx_guard RAII_GUARD(xpu_ctx->x_context());
|
||||
using XPUType = typename XPUTypeTrait<bfloat16>::Type; // only support bfloat16
|
||||
typedef paddle::bfloat16 data_t;
|
||||
const int token_num = tmp_out.dims()[0];
|
||||
const int dim = tmp_out.dims()[1];
|
||||
const int bsz = cum_offsets.shape()[0];
|
||||
const int bsz = cum_offsets.shape()[0];
|
||||
int enc_batch = enc_batch_tensor.data<int32_t>()[0];
|
||||
int dec_batch = dec_batch_tensor.data<int32_t>()[0];
|
||||
|
||||
baidu::xpu::api::VectorParam<int32_t> encoder_seqs_lods_vp{const_cast<int32_t*>(encoder_seq_lod_cpu.data<int32_t>()), enc_batch + 1, const_cast<int32_t*>(encoder_seq_lod.data<int32_t>())};
|
||||
baidu::xpu::api::VectorParam<int32_t> encoder_batch_map_vp{const_cast<int32_t*>(encoder_batch_idx_cpu.data<int32_t>()), enc_batch, const_cast<int32_t*>(encoder_batch_idx.data<int32_t>())};
|
||||
baidu::xpu::api::VectorParam<int32_t> decoder_batch_map_vp{const_cast<int32_t*>(decoder_batch_idx_cpu.data<int32_t>()), dec_batch, const_cast<int32_t*>(decoder_batch_idx.data<int32_t>())};
|
||||
|
||||
std::vector<int> seq_lens_encoder_cpu(bsz, 0);
|
||||
std::vector<int> seq_lens_decoder_cpu(bsz, 0);
|
||||
std::vector<int> encoder_batch_idx; // 去除空隙的batch map
|
||||
std::vector<int> decoder_batch_idx; // 去除空隙的batch map
|
||||
std::vector<int> encoder_seq_lod;
|
||||
int r = xpu_memcpy(seq_lens_encoder_cpu.data(),
|
||||
seq_lens_encoder.data<int>(),
|
||||
sizeof(int32_t) * bsz,
|
||||
XPUMemcpyKind::XPU_DEVICE_TO_HOST);
|
||||
r = xpu_memcpy(seq_lens_decoder_cpu.data(),
|
||||
seq_lens_decoder.data<int>(),
|
||||
sizeof(int32_t) * bsz,
|
||||
XPUMemcpyKind::XPU_DEVICE_TO_HOST);
|
||||
int enc_batch = 0, dec_batch = 0;
|
||||
int batch_offset = 0;
|
||||
encoder_seq_lod.push_back(0);
|
||||
for(int i = 0; i < bsz; ++i){
|
||||
if(seq_lens_encoder_cpu[i] > 0){
|
||||
enc_batch++;
|
||||
encoder_batch_idx.push_back(i - batch_offset);
|
||||
encoder_seq_lod.push_back(seq_lens_encoder_cpu[i]);
|
||||
encoder_seq_lod[enc_batch] += encoder_seq_lod[enc_batch - 1];
|
||||
}
|
||||
else if(seq_lens_decoder_cpu[i] > 0){
|
||||
dec_batch++;
|
||||
decoder_batch_idx.push_back(i - batch_offset);
|
||||
}
|
||||
else{
|
||||
batch_offset++;
|
||||
}
|
||||
}
|
||||
baidu::xpu::api::VectorParam<int32_t> encoder_seqs_lods_vp =
|
||||
baidu::xpu::api::VectorParam<int32_t>{encoder_seq_lod.data(), enc_batch + 1, nullptr}
|
||||
.to_xpu(RAII_GUARD);
|
||||
baidu::xpu::api::VectorParam<int32_t> encoder_batch_map_vp =
|
||||
baidu::xpu::api::VectorParam<int32_t>{encoder_batch_idx.data(), enc_batch, nullptr}
|
||||
.to_xpu(RAII_GUARD);
|
||||
baidu::xpu::api::VectorParam<int32_t> decoder_batch_map_vp =
|
||||
baidu::xpu::api::VectorParam<int32_t>{decoder_batch_idx.data(), dec_batch, nullptr}
|
||||
.to_xpu(RAII_GUARD);
|
||||
auto out = paddle::full({token_num, dim}, -2, tmp_out.type(), tmp_out.place());
|
||||
|
||||
r = xftkernel::xft_eb_adjust_batch<XPUType, XPUType>(
|
||||
int r = xftkernel::xft_eb_adjust_batch<XPUType, XPUType>(
|
||||
xpu_ctx->x_context(),
|
||||
reinterpret_cast<const XPUType*>(tmp_out.data<data_t>()),
|
||||
reinterpret_cast<XPUType*>(out.data<data_t>()),
|
||||
@@ -88,8 +60,14 @@ std::vector<paddle::Tensor> AdjustBatch(const paddle::Tensor& tmp_out, // [token
|
||||
|
||||
std::vector<std::vector<int64_t>> AdjustBatchInferShape(const std::vector<int64_t>& tmp_out_shape,
|
||||
const std::vector<int64_t>& cum_offsets_shape,
|
||||
const std::vector<int64_t>& seq_lens_decoder_shape,
|
||||
const std::vector<int64_t>& seq_lens_encoder_shape,
|
||||
const std::vector<int64_t>& encoder_seq_lod_shape,
|
||||
const std::vector<int64_t>& encoder_batch_idx_shape,
|
||||
const std::vector<int64_t>& decoder_batch_idx_shape,
|
||||
const std::vector<int64_t>& encoder_seq_lod_cpu_shape,
|
||||
const std::vector<int64_t>& encoder_batch_idx_cpu_shape,
|
||||
const std::vector<int64_t>& decoder_batch_idx_cpu_shape,
|
||||
const std::vector<int64_t>& enc_batch_tensor_shape,
|
||||
const std::vector<int64_t>& dec_batch_tensor_shape,
|
||||
const paddle::optional<std::vector<int64_t>>& output_padding_offset_shape) {
|
||||
if (output_padding_offset_shape) {
|
||||
PD_THROW("speculative decoding is not supported in XPU.");
|
||||
@@ -101,14 +79,20 @@ std::vector<std::vector<int64_t>> AdjustBatchInferShape(const std::vector<int64_
|
||||
|
||||
std::vector<paddle::DataType> AdjustBatchInferDtype(const paddle::DataType& tmp_out_dtype,
|
||||
const paddle::DataType& cum_offsets_dtype,
|
||||
const paddle::DataType& seq_lens_decoder_dtype,
|
||||
const paddle::DataType& seq_lens_encoder_dtype,
|
||||
const paddle::DataType& encoder_seq_lod_dtype,
|
||||
const paddle::DataType& encoder_batch_idx_dtype,
|
||||
const paddle::DataType& decoder_batch_idx_dtype,
|
||||
const paddle::DataType& encoder_seq_lod_cpu_dtype,
|
||||
const paddle::DataType& encoder_batch_idx_cpu_dtype,
|
||||
const paddle::DataType& decoder_batch_idx_cpu_dtype,
|
||||
const paddle::DataType& enc_batch_tensor_dtype,
|
||||
const paddle::DataType& dec_batch_tensor_dtype,
|
||||
const paddle::optional<paddle::DataType>& output_padding_offset_dtype) {
|
||||
return {tmp_out_dtype};
|
||||
}
|
||||
|
||||
PD_BUILD_OP(adjust_batch)
|
||||
.Inputs({"tmp_out", "cum_offsets", "seq_lens_decoder", "seq_lens_encoder", paddle::Optional("output_padding_offset")})
|
||||
.Inputs({"tmp_out", "cum_offsets", "encoder_seq_lod", "encoder_batch_idx", "decoder_batch_idx", "encoder_seq_lod_cpu", "encoder_batch_idx_cpu", "decoder_batch_idx_cpu", "enc_batch_tensor", "dec_batch_tensor", paddle::Optional("output_padding_offset")})
|
||||
.Outputs({"out"})
|
||||
.Attrs({"max_input_length: int"})
|
||||
.SetKernelFn(PD_KERNEL(AdjustBatch))
|
||||
|
||||
76
csrc/xpu/src/batch_matmul_xpu.cc
Normal file
76
csrc/xpu/src/batch_matmul_xpu.cc
Normal file
@@ -0,0 +1,76 @@
|
||||
// Copyright (c) 2024 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/backends/xpu/xpu_context.h>
|
||||
#include "paddle/extension.h"
|
||||
#include "xpu/plugin.h"
|
||||
#include <xft/xdnn_plugin.h>
|
||||
#include "cublasLt.h"
|
||||
#include "ops.h"
|
||||
|
||||
namespace xftkernel = baidu::xpu::xftkernel;
|
||||
std::vector<paddle::Tensor> Bmm(const paddle::Tensor& input,
|
||||
const paddle::Tensor& weight) {
|
||||
phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
|
||||
auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
|
||||
auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
|
||||
xpu::ctx_guard RAII_GUARD(xpu_ctx->x_context());
|
||||
using XPUType = typename XPUTypeTrait<bfloat16>::Type; // only support bfloat16
|
||||
typedef paddle::bfloat16 data_t;
|
||||
const int batch = input.dims()[0];
|
||||
const int m = input.dims()[1];
|
||||
const int n = weight.dims()[2];
|
||||
const int k = weight.dims()[1];
|
||||
|
||||
auto out = paddle::full({batch, m, n}, -2, input.type(), input.place());
|
||||
int ret = baidu::xpu::xblas::fc_batched<XPUType, XPUType, XPUType, float, float, 0>(
|
||||
xpu_ctx->x_context(),
|
||||
batch,
|
||||
false,
|
||||
false,
|
||||
m,
|
||||
n,
|
||||
k,
|
||||
1.0f,
|
||||
reinterpret_cast<const XPUType*>(input.data<data_t>()),
|
||||
m * k,
|
||||
reinterpret_cast<const XPUType*>(weight.data<data_t>()),
|
||||
k * n,
|
||||
0.0f,
|
||||
reinterpret_cast<XPUType*>(out.data<data_t>()),
|
||||
m * n,
|
||||
nullptr,
|
||||
nullptr);
|
||||
return {out};
|
||||
}
|
||||
|
||||
std::vector<std::vector<int64_t>> BmmInferShape(const std::vector<int64_t>& input_shape,
|
||||
const std::vector<int64_t>& weight_shape) {
|
||||
int64_t batch = input_shape[0];
|
||||
int64_t m = input_shape[1];
|
||||
int64_t n = weight_shape[2];
|
||||
return {{batch, m, n}};
|
||||
}
|
||||
|
||||
std::vector<paddle::DataType> BmmInferDtype(const paddle::DataType& input_dtype,
|
||||
const paddle::DataType& weight_dtype) {
|
||||
return {input_dtype};
|
||||
}
|
||||
|
||||
PD_BUILD_OP(batch_matmul_xpu)
|
||||
.Inputs({"imput", "weight"})
|
||||
.Outputs({"out"})
|
||||
.SetKernelFn(PD_KERNEL(Bmm))
|
||||
.SetInferShapeFn(PD_INFER_SHAPE(BmmInferShape))
|
||||
.SetInferDtypeFn(PD_INFER_DTYPE(BmmInferDtype));
|
||||
448
csrc/xpu/src/block_attn.cc
Normal file
448
csrc/xpu/src/block_attn.cc
Normal file
@@ -0,0 +1,448 @@
|
||||
// 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/backends/xpu/xpu_context.h>
|
||||
#include "paddle/extension.h"
|
||||
#include "paddle/phi/core/enforce.h"
|
||||
#include "xpu/plugin.h"
|
||||
#include "ops.h"
|
||||
|
||||
#include <core/ctx_manager.h>
|
||||
#include <core/xft_check.h>
|
||||
#include <core/xft_event.h>
|
||||
#include <core/xft_params.h>
|
||||
#include <xft/xdnn_plugin.h>
|
||||
#include <xft/operation/page_attn.h>
|
||||
#include <xft/operation/fmha.h>
|
||||
#include <flash_api.h> // link xfa
|
||||
|
||||
namespace xftkernel = baidu::xpu::xftkernel;
|
||||
namespace api = baidu::xpu::api;
|
||||
|
||||
std::vector<paddle::Tensor> MlaBlockAttn(
|
||||
const paddle::Tensor& attn_input, // [m, 7168]
|
||||
const paddle::Tensor& q_a_proj_weights, // self.q_a_proj_weights[i] [7168, 1536] 【enc】【dec】
|
||||
const paddle::Tensor& q_b_proj_weights, // self.q_b_proj_weights[i] [1536, 24576] 【enc】【dec】
|
||||
const paddle::Tensor& q_a_layernorm_weights, // self.q_a_layernorm_weights[i] [1536] 【enc】【dec】
|
||||
const paddle::Tensor& kv_a_proj_with_mqa_weights, // self.kv_a_proj_with_mqa_weights[i] [7168, 576] 【enc】
|
||||
const paddle::Tensor& kv_a_layernorm_weights, // self.kv_a_layernorm_weights[i] [512] 【enc】【dec】
|
||||
const paddle::Tensor& kv_b_proj_weights, // self.kv_b_proj_weights[i] [512, 32768] 【enc】
|
||||
const paddle::Tensor& linear_weights, // self.linear_weights[i] [16384, 7168] 【enc】【dec】
|
||||
const paddle::Tensor& k_b_proj_weights, // self.k_b_proj_weights[i] [128, 128, 512] 【dec】
|
||||
const paddle::Tensor& v_b_proj_weights, // self.v_b_proj_weights[i] [128, 512, 128] 【dec】
|
||||
const paddle::Tensor& cos_sin_cache, // self.config.rotary_emb.cos_sin_cache [163840, 64]
|
||||
const paddle::Tensor& position_ids, // self.position_ids [m]
|
||||
const paddle::Tensor& kv_cache,
|
||||
const paddle::Tensor& block_tables,
|
||||
const paddle::Tensor& encoder_seq_lod,
|
||||
const paddle::Tensor& encoder_batch_map,
|
||||
const paddle::Tensor& start_token_raw,
|
||||
const paddle::Tensor& decoder_context_len_cache,
|
||||
const paddle::Tensor& decoder_context_len,
|
||||
const paddle::Tensor& decoder_batch_map,
|
||||
const paddle::Tensor& kv_seq_lod_raw,
|
||||
const paddle::Tensor& encoder_seq_lod_cpu,
|
||||
const paddle::Tensor& encoder_batch_map_cpu,
|
||||
const paddle::Tensor& start_token_raw_cpu,
|
||||
const paddle::Tensor& decoder_context_len_cache_cpu,
|
||||
const paddle::Tensor& decoder_context_len_cpu,
|
||||
const paddle::Tensor& decoder_batch_map_cpu,
|
||||
const paddle::Tensor& kv_seq_lod_raw_cpu,
|
||||
const paddle::Tensor& enc_batch_tensor,
|
||||
const paddle::Tensor& dec_batch_tensor,
|
||||
const paddle::Tensor& total_enc_len_tensor,
|
||||
const paddle::Tensor& padding_offsets,
|
||||
const paddle::Tensor& cum_offsets,
|
||||
const paddle::Tensor& max_enc_len_this_time,
|
||||
const paddle::Tensor& max_dec_len_this_time,
|
||||
const std::string& cache_quant_type_str,
|
||||
const float softmax_scale,
|
||||
const int block_size,
|
||||
const int num_heads, // self.num_heads
|
||||
const int q_lora_rank, // 传 self.config.mla_config.q_lora_rank = 1536
|
||||
const int qk_head_dim, // self.config.mla_config.qk_head_dim 192
|
||||
const int qk_nope_head_dim, // self.config.mla_config.qk_nope_head_dim 128
|
||||
const int qk_rope_head_dim, // self.config.mla_config.qk_rope_head_dim 64
|
||||
const int kv_lora_rank, // self.config.mla_config.kv_lora_rank 512
|
||||
const int v_head_dim // self.config.mla_config.v_head_dim 128
|
||||
) {
|
||||
|
||||
phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
|
||||
auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
|
||||
auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
|
||||
xpu::ctx_guard RAII_GUARD(xpu_ctx->x_context());
|
||||
using CType = typename XPUTypeTrait<bfloat16>::Type;
|
||||
using CacheType = typename XPUTypeTrait<bfloat16>::Type;
|
||||
typedef paddle::bfloat16 data_t, cache_t;
|
||||
|
||||
const int bsz = encoder_batch_map.dims()[0]; // batch map tensor长度都是max_batch_size, 有效长度根据
|
||||
// const int block_batch = block_tables.dims()[0]; // TODO参数含义 block_batch_ PageParam page_param_
|
||||
const int max_block_per_seq = block_tables.dims()[1];
|
||||
const int max_seq_len = block_size * max_block_per_seq;
|
||||
const int token_num = attn_input.dims()[0];
|
||||
const int hidden_dim = attn_input.dims()[1];
|
||||
int enc_batch = enc_batch_tensor.data<int32_t>()[0];
|
||||
int dec_batch = dec_batch_tensor.data<int32_t>()[0];
|
||||
int total_enc_len = total_enc_len_tensor.data<int32_t>()[0];
|
||||
|
||||
// 输出tensor
|
||||
auto block_attn_out = paddle::full({token_num, hidden_dim}, -1, attn_input.type(), attn_input.place());
|
||||
|
||||
// encoder
|
||||
if(enc_batch > 0){
|
||||
// split and copy
|
||||
auto ln_out_encoder = paddle::experimental::slice(attn_input, {0}, {0}, {total_enc_len}, {0}, {});
|
||||
auto position_ids_enc = paddle::experimental::slice(position_ids, {0}, {0}, {total_enc_len}, {0}, {});
|
||||
// encoder no absorb fc
|
||||
if (q_lora_rank <= 0) {
|
||||
assert("Not supported in block attn");
|
||||
}
|
||||
auto query_0 = paddle::matmul(ln_out_encoder, q_a_proj_weights);
|
||||
auto query_1 = std::get<0>(paddle::experimental::rms_norm(query_0, NULL, NULL, q_a_layernorm_weights, NULL, 1e-06, 1, -1, 0, 0, 0));
|
||||
auto query_2 = paddle::matmul(query_1, q_b_proj_weights);
|
||||
|
||||
// auto query_nope = paddle::full({total_enc_len, num_heads, qk_nope_head_dim}, -1, query_2.type(), query_2.place());
|
||||
auto query_pe = paddle::full({total_enc_len, num_heads, qk_rope_head_dim}, -1, query_2.type(), query_2.place());
|
||||
int r = api::split<CType>(
|
||||
xpu_ctx->x_context(),
|
||||
reinterpret_cast<const CType*>(query_2.data<data_t>()),
|
||||
{nullptr, reinterpret_cast<CType*>(query_pe.data<data_t>())},
|
||||
{total_enc_len, num_heads, qk_head_dim},
|
||||
{qk_nope_head_dim, qk_rope_head_dim},
|
||||
2); // axis = 2
|
||||
PD_CHECK(r == 0, "api::split failed.");
|
||||
|
||||
auto compressed_kv = paddle::matmul(ln_out_encoder, kv_a_proj_with_mqa_weights);
|
||||
|
||||
auto compressed_kv_0 = paddle::full({total_enc_len, kv_lora_rank}, -1, compressed_kv.type(), compressed_kv.place());
|
||||
auto key_pe = paddle::full({total_enc_len, qk_rope_head_dim}, -1, compressed_kv.type(), compressed_kv.place());
|
||||
r = api::split<CType>(
|
||||
xpu_ctx->x_context(),
|
||||
reinterpret_cast<const CType*>(compressed_kv.data<data_t>()),
|
||||
{reinterpret_cast<CType*>(compressed_kv_0.data<data_t>()), reinterpret_cast<CType*>(key_pe.data<data_t>())},
|
||||
{total_enc_len, kv_lora_rank + qk_rope_head_dim},
|
||||
{kv_lora_rank, qk_rope_head_dim},
|
||||
1); // axis = 1
|
||||
PD_CHECK(r == 0, "api::split failed.");
|
||||
|
||||
auto compressed_kv_1 = std::get<0>(paddle::experimental::rms_norm(compressed_kv_0, NULL, NULL, kv_a_layernorm_weights, NULL, 1e-06, 1, -1, 0, 0, 0));
|
||||
FusedRotaryPositionEncoding(query_pe, key_pe, position_ids_enc, cos_sin_cache, qk_rope_head_dim, false);
|
||||
PrefillMLAWriteCacheKernel(
|
||||
compressed_kv_1,
|
||||
key_pe,
|
||||
kv_cache,
|
||||
encoder_seq_lod,
|
||||
encoder_batch_map,
|
||||
start_token_raw,
|
||||
encoder_seq_lod_cpu,
|
||||
encoder_batch_map_cpu,
|
||||
start_token_raw_cpu,
|
||||
enc_batch_tensor,
|
||||
dec_batch_tensor,
|
||||
padding_offsets,
|
||||
cum_offsets,
|
||||
block_tables,
|
||||
cache_quant_type_str,
|
||||
max_seq_len,
|
||||
num_heads);
|
||||
|
||||
auto key_value = paddle::matmul(compressed_kv_1, kv_b_proj_weights);
|
||||
|
||||
auto key_nope = paddle::full({total_enc_len, num_heads, qk_nope_head_dim}, -1, key_value.type(), key_value.place());
|
||||
auto value = paddle::full({total_enc_len, num_heads, v_head_dim}, -1, key_value.type(), key_value.place());
|
||||
r = api::split<CType>(
|
||||
xpu_ctx->x_context(),
|
||||
reinterpret_cast<const CType*>(key_value.data<data_t>()),
|
||||
{reinterpret_cast<CType*>(key_nope.data<data_t>()), reinterpret_cast<CType*>(value.data<data_t>())},
|
||||
{total_enc_len, num_heads, qk_nope_head_dim + v_head_dim},
|
||||
{qk_nope_head_dim, v_head_dim},
|
||||
2); // axis = 2
|
||||
PD_CHECK(r == 0, "api::split failed.");
|
||||
|
||||
r = api::strided_slice_view_update<CType>(
|
||||
xpu_ctx->x_context(),
|
||||
reinterpret_cast<CType*>(query_pe.data<data_t>()),
|
||||
reinterpret_cast<CType*>(query_2.data<data_t>()),
|
||||
{total_enc_len, num_heads, qk_rope_head_dim},
|
||||
{total_enc_len, num_heads, qk_head_dim},
|
||||
{0, 0, qk_nope_head_dim}, // start
|
||||
{total_enc_len, num_heads, qk_head_dim}, // end
|
||||
{1, 1, 1}); // stride
|
||||
PD_CHECK(r == 0, "api::strided_slice_view_update failed.");
|
||||
|
||||
auto key = paddle::full({total_enc_len, num_heads, qk_head_dim}, -1, query_2.type(), query_2.place());
|
||||
r = api::strided_slice_view_update<CType>(
|
||||
xpu_ctx->x_context(),
|
||||
reinterpret_cast<CType*>(key_nope.data<data_t>()),
|
||||
reinterpret_cast<CType*>(key.data<data_t>()),
|
||||
{total_enc_len, num_heads, qk_nope_head_dim},
|
||||
{total_enc_len, num_heads, qk_head_dim},
|
||||
{0, 0, 0}, // start
|
||||
{total_enc_len, num_heads, qk_nope_head_dim}, // end
|
||||
{1, 1, 1}); // stride
|
||||
PD_CHECK(r == 0, "api::strided_slice_view_update failed.");
|
||||
|
||||
auto key_pe_broadcast = paddle::full({total_enc_len, num_heads, qk_rope_head_dim}, -1, query_2.type(), query_2.place());
|
||||
r = api::broadcast<CType>(
|
||||
xpu_ctx->x_context(),
|
||||
reinterpret_cast<CType*>(key_pe.data<data_t>()),
|
||||
reinterpret_cast<CType*>(key_pe_broadcast.data<data_t>()),
|
||||
{total_enc_len, 1, qk_rope_head_dim},
|
||||
{total_enc_len, num_heads, qk_rope_head_dim});
|
||||
PD_CHECK(r == 0, "api::broadcast failed.");
|
||||
|
||||
r = api::strided_slice_view_update<CType>(
|
||||
xpu_ctx->x_context(),
|
||||
reinterpret_cast<CType*>(key_pe_broadcast.data<data_t>()),
|
||||
reinterpret_cast<CType*>(key.data<data_t>()),
|
||||
{total_enc_len, num_heads, qk_rope_head_dim},
|
||||
{total_enc_len, num_heads, qk_head_dim},
|
||||
{0, 0, qk_nope_head_dim}, // start
|
||||
{total_enc_len, num_heads, qk_head_dim}, // end
|
||||
{1, 1, 1}); // stride
|
||||
PD_CHECK(r == 0, "api::strided_slice_view_update failed.");
|
||||
|
||||
// attn
|
||||
auto fmha_out_prefill = MlaEnAttn(
|
||||
query_2,
|
||||
key,
|
||||
value,
|
||||
encoder_seq_lod,
|
||||
encoder_batch_map,
|
||||
encoder_seq_lod_cpu,
|
||||
encoder_batch_map_cpu,
|
||||
enc_batch_tensor,
|
||||
padding_offsets,
|
||||
cum_offsets,
|
||||
block_tables,
|
||||
max_enc_len_this_time,
|
||||
max_dec_len_this_time,
|
||||
softmax_scale,
|
||||
block_size,
|
||||
num_heads,
|
||||
qk_head_dim,
|
||||
v_head_dim)[0];
|
||||
|
||||
auto out_linear_out_prefill = paddle::matmul(fmha_out_prefill, linear_weights);
|
||||
|
||||
r = api::copy<CType>(
|
||||
xpu_ctx->x_context(),
|
||||
reinterpret_cast<const CType*>(out_linear_out_prefill.data<data_t>()),
|
||||
reinterpret_cast<CType*>(block_attn_out.data<data_t>()),
|
||||
{total_enc_len * hidden_dim});
|
||||
PD_CHECK(r == 0, "api::copy failed.");
|
||||
|
||||
}
|
||||
|
||||
// decoder
|
||||
if(dec_batch > 0){
|
||||
// split
|
||||
auto ln_out_decoder = paddle::experimental::slice(attn_input, {0}, {total_enc_len}, {token_num}, {0}, {});
|
||||
auto position_ids_dec = paddle::experimental::slice(position_ids, {0}, {total_enc_len}, {token_num}, {0}, {});
|
||||
// decoder absorb fc
|
||||
auto compressed_kv = paddle::matmul(ln_out_decoder, kv_a_proj_with_mqa_weights);
|
||||
|
||||
auto compressed_kv_0 = paddle::full({dec_batch, kv_lora_rank}, -1, compressed_kv.type(), compressed_kv.place());
|
||||
auto key_pe = paddle::full({dec_batch, qk_rope_head_dim}, -1, compressed_kv.type(), compressed_kv.place());
|
||||
int r = api::split<CType>(
|
||||
xpu_ctx->x_context(),
|
||||
reinterpret_cast<const CType*>(compressed_kv.data<data_t>()),
|
||||
{reinterpret_cast<CType*>(compressed_kv_0.data<data_t>()), reinterpret_cast<CType*>(key_pe.data<data_t>())},
|
||||
{dec_batch, kv_lora_rank + qk_rope_head_dim},
|
||||
{kv_lora_rank, qk_rope_head_dim},
|
||||
1); // axis = 1
|
||||
PD_CHECK(r == 0, "api::split failed.");
|
||||
|
||||
auto compressed_kv_1 = std::get<0>(paddle::experimental::rms_norm(compressed_kv_0, NULL, NULL, kv_a_layernorm_weights, NULL, 1e-06, 1, -1, 0, 0, 0)); // TODO inplace
|
||||
|
||||
if (q_lora_rank <= 0) {
|
||||
assert("Not supported in block attn");
|
||||
}
|
||||
auto query_0 = paddle::matmul(ln_out_decoder, q_a_proj_weights);
|
||||
auto query_1 = std::get<0>(paddle::experimental::rms_norm(query_0, NULL, NULL, q_a_layernorm_weights, NULL, 1e-06, 1, -1, 0, 0, 0));
|
||||
auto query_2 = paddle::matmul(query_1, q_b_proj_weights);
|
||||
|
||||
auto query_nope = paddle::full({dec_batch, num_heads, qk_nope_head_dim}, -1, query_2.type(), query_2.place());
|
||||
auto query_pe = paddle::full({dec_batch, num_heads, qk_rope_head_dim}, -1, query_2.type(), query_2.place());
|
||||
r = api::split<CType>(
|
||||
xpu_ctx->x_context(),
|
||||
reinterpret_cast<const CType*>(query_2.data<data_t>()),
|
||||
{reinterpret_cast<CType*>(query_nope.data<data_t>()), reinterpret_cast<CType*>(query_pe.data<data_t>())},
|
||||
{dec_batch, num_heads, qk_head_dim},
|
||||
{qk_nope_head_dim, qk_rope_head_dim},
|
||||
2); // axis = 2
|
||||
PD_CHECK(r == 0, "api::split failed.");
|
||||
|
||||
auto query_nope_trans = paddle::full({num_heads, dec_batch, qk_nope_head_dim}, -1, query_2.type(), query_2.place());
|
||||
r = api::transpose<CType>(
|
||||
xpu_ctx->x_context(),
|
||||
reinterpret_cast<const CType*>(query_nope.data<data_t>()),
|
||||
reinterpret_cast<CType*>(query_nope_trans.data<data_t>()),
|
||||
{dec_batch, num_heads, qk_nope_head_dim},
|
||||
{1, 0, 2});
|
||||
PD_CHECK(r == 0, "api::transpose failed.");
|
||||
|
||||
auto q_nope_bmm = Bmm(query_nope_trans, k_b_proj_weights)[0];
|
||||
|
||||
auto q_nope_out = paddle::full({dec_batch, num_heads, kv_lora_rank}, -1, query_2.type(), query_2.place());
|
||||
r = api::transpose<CType>(
|
||||
xpu_ctx->x_context(),
|
||||
reinterpret_cast<const CType*>(q_nope_bmm.data<data_t>()),
|
||||
reinterpret_cast<CType*>(q_nope_out.data<data_t>()),
|
||||
{num_heads, dec_batch, kv_lora_rank},
|
||||
{1, 0, 2});
|
||||
PD_CHECK(r == 0, "api::transpose failed.");
|
||||
|
||||
FusedRotaryPositionEncoding(query_pe, key_pe, position_ids_dec, cos_sin_cache, qk_rope_head_dim, false);
|
||||
DecodeMLAWriteCacheKernel(
|
||||
compressed_kv_1,
|
||||
key_pe,
|
||||
kv_cache,
|
||||
decoder_context_len_cache,
|
||||
decoder_batch_map,
|
||||
kv_seq_lod_raw,
|
||||
decoder_context_len_cache_cpu,
|
||||
decoder_batch_map_cpu,
|
||||
kv_seq_lod_raw_cpu,
|
||||
enc_batch_tensor,
|
||||
dec_batch_tensor,
|
||||
padding_offsets,
|
||||
cum_offsets,
|
||||
block_tables,
|
||||
cache_quant_type_str,
|
||||
max_seq_len,
|
||||
num_heads,
|
||||
false); // self.config.speculate_config.speculate_method is not None
|
||||
|
||||
auto q_input = paddle::full({dec_batch, num_heads * (kv_lora_rank + qk_rope_head_dim)}, -1, query_2.type(), query_2.place()); // TODO inplace
|
||||
r = api::concat<CType>(
|
||||
xpu_ctx->x_context(),
|
||||
{reinterpret_cast<const CType*>(q_nope_out.data<data_t>()), reinterpret_cast<const CType*>(query_pe.data<data_t>())},
|
||||
reinterpret_cast<CType*>(q_input.data<data_t>()),
|
||||
{{dec_batch, num_heads, kv_lora_rank}, {dec_batch, num_heads, qk_rope_head_dim}},
|
||||
2); // axis = 2
|
||||
|
||||
// attn
|
||||
auto fmha_out_decode = MlaDeAttn( // [dec_batch, 128*512]
|
||||
q_input,
|
||||
kv_cache,
|
||||
decoder_context_len,
|
||||
decoder_batch_map,
|
||||
decoder_context_len_cpu,
|
||||
decoder_batch_map_cpu,
|
||||
dec_batch_tensor,
|
||||
padding_offsets,
|
||||
cum_offsets,
|
||||
block_tables,
|
||||
softmax_scale,
|
||||
block_size,
|
||||
num_heads,
|
||||
kv_lora_rank,
|
||||
qk_rope_head_dim,
|
||||
qk_head_dim,
|
||||
v_head_dim)[0];
|
||||
|
||||
auto fmha_out_decode_trans = paddle::full({num_heads, dec_batch, kv_lora_rank}, -1, query_2.type(), query_2.place()); // shape要跟transpose后的对齐，下面的BMM会读取shape做计算
|
||||
r = api::transpose<CType>(
|
||||
xpu_ctx->x_context(),
|
||||
reinterpret_cast<const CType*>(fmha_out_decode.data<data_t>()),
|
||||
reinterpret_cast<CType*>(fmha_out_decode_trans.data<data_t>()),
|
||||
{dec_batch, num_heads, kv_lora_rank},
|
||||
{1, 0, 2});
|
||||
PD_CHECK(r == 0, "api::transpose failed.");
|
||||
|
||||
auto fmha_out_decode_bmm = Bmm(fmha_out_decode_trans, v_b_proj_weights)[0]; // fmha_out_decode_bmm = [dec_batch, 128*128]
|
||||
|
||||
auto fmha_out_decode_out = paddle::full({dec_batch, num_heads * v_head_dim}, -1, query_2.type(), query_2.place());
|
||||
r = api::transpose<CType>(
|
||||
xpu_ctx->x_context(),
|
||||
reinterpret_cast<const CType*>(fmha_out_decode_bmm.data<data_t>()),
|
||||
reinterpret_cast<CType*>(fmha_out_decode_out.data<data_t>()),
|
||||
{num_heads, dec_batch, v_head_dim},
|
||||
{1, 0, 2});
|
||||
PD_CHECK(r == 0, "api::transpose failed.");
|
||||
|
||||
auto out_linear_out_decode = paddle::matmul(fmha_out_decode_out, linear_weights);
|
||||
|
||||
r = api::copy<CType>(
|
||||
xpu_ctx->x_context(),
|
||||
reinterpret_cast<const CType*>(out_linear_out_decode.data<data_t>()),
|
||||
reinterpret_cast<CType*>(block_attn_out.data<data_t>()) + total_enc_len * hidden_dim,
|
||||
{dec_batch * hidden_dim});
|
||||
PD_CHECK(r == 0, "api::copy failed.");
|
||||
}
|
||||
return {block_attn_out};
|
||||
|
||||
}
|
||||
|
||||
|
||||
std::vector<std::vector<int64_t>> MlaBlockAttnInferShape(const std::vector<int64_t>& attn_input_shape) {
|
||||
return {attn_input_shape};
|
||||
}
|
||||
|
||||
std::vector<paddle::DataType> MlaBlockAttnInferDtype(const paddle::DataType& attn_input_dtype) {
|
||||
return {attn_input_dtype};
|
||||
}
|
||||
|
||||
PD_BUILD_OP(absorb_mla_block_attention_xpu)
|
||||
.Inputs({"attn_input",
|
||||
"q_a_proj_weights",
|
||||
"q_b_proj_weights",
|
||||
"q_a_layernorm_weights",
|
||||
"kv_a_proj_with_mqa_weights",
|
||||
"kv_a_layernorm_weights",
|
||||
"kv_b_proj_weights",
|
||||
"linear_weights",
|
||||
"k_b_proj_weights",
|
||||
"v_b_proj_weights",
|
||||
"cos_sin_cache",
|
||||
"position_ids",
|
||||
"kv_cache",
|
||||
"block_tables",
|
||||
"encoder_seq_lod",
|
||||
"encoder_batch_map",
|
||||
"start_token_raw",
|
||||
"decoder_context_len_cache",
|
||||
"decoder_context_len",
|
||||
"decoder_batch_map",
|
||||
"kv_seq_lod_raw",
|
||||
"encoder_seq_lod_cpu",
|
||||
"encoder_batch_map_cpu",
|
||||
"start_token_raw_cpu",
|
||||
"decoder_context_len_cache_cpu",
|
||||
"decoder_context_len_cpu",
|
||||
"decoder_batch_map_cpu",
|
||||
"kv_seq_lod_raw_cpu",
|
||||
"enc_batch_tensor",
|
||||
"dec_batch_tensor",
|
||||
"total_enc_len_tensor",
|
||||
"padding_offsets",
|
||||
"cum_offsets",
|
||||
"max_enc_len_this_time",
|
||||
"max_dec_len_this_time",})
|
||||
.Outputs({"block_attn_out"})
|
||||
// .SetInplaceMap({{"kv_cache", "kv_cache_out"}})
|
||||
.Attrs({"cache_quant_type_str: std::string",
|
||||
"softmax_scale: float",
|
||||
"block_size: int",
|
||||
"num_heads: int",
|
||||
"q_lora_rank: int",
|
||||
"qk_head_dim: int",
|
||||
"qk_nope_head_dim: int",
|
||||
"qk_rope_head_dim: int",
|
||||
"kv_lora_rank: int",
|
||||
"v_head_dim: int"})
|
||||
.SetKernelFn(PD_KERNEL(MlaBlockAttn))
|
||||
.SetInferShapeFn(PD_INFER_SHAPE(MlaBlockAttnInferShape))
|
||||
.SetInferDtypeFn(PD_INFER_DTYPE(MlaBlockAttnInferDtype));
|
||||
102
csrc/xpu/src/dynamic_set_split_value.cc
Normal file
102
csrc/xpu/src/dynamic_set_split_value.cc
Normal file
@@ -0,0 +1,102 @@
|
||||
// Copyright (c) 2024 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 <cassert>
|
||||
#include <paddle/phi/backends/xpu/xpu_context.h>
|
||||
#include "paddle/extension.h"
|
||||
#include "xpu/plugin.h"
|
||||
#include <xft/xdnn_plugin.h>
|
||||
namespace xftkernel = baidu::xpu::xftkernel;
|
||||
std::vector<paddle::Tensor> DynamicSetValue(const paddle::Tensor& input,
|
||||
const paddle::Tensor& index,
|
||||
paddle::Tensor* output,
|
||||
bool is_before_index) { // true: [ : index, :] ; false: [index : , :]
|
||||
|
||||
phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
|
||||
auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
|
||||
auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
|
||||
using XPUType = typename XPUTypeTrait<bfloat16>::Type; // only support bfloat16
|
||||
typedef paddle::bfloat16 data_t;
|
||||
|
||||
int index_num = index.data<int64_t>()[0];
|
||||
int total_len = output->dims()[0];
|
||||
int dim = output->dims()[1];
|
||||
int copy_num = input.numel();
|
||||
assert(dim == input.dims()[1] && "input.dim != output.dim! ");
|
||||
assert((index_num >= 0 && index_num <= total_len) && "index out of range! ");
|
||||
if (is_before_index) {
|
||||
assert((index_num - 0) * dim == copy_num && "copy_num wrong! ");
|
||||
}
|
||||
else{
|
||||
assert((total_len - index_num) * dim == copy_num && "copy_num wrong! ");
|
||||
}
|
||||
int r = baidu::xpu::api::copy<XPUType>(
|
||||
xpu_ctx->x_context(),
|
||||
reinterpret_cast<const XPUType*>(input.data<data_t>()),
|
||||
reinterpret_cast<XPUType*>(output->data<data_t>() + (is_before_index ? 0 : index_num * dim)),
|
||||
copy_num);
|
||||
|
||||
return {};
|
||||
}
|
||||
|
||||
PD_BUILD_OP(dynamic_set_value)
|
||||
.Inputs({"input", "index", "output"})
|
||||
.Outputs({"output_"})
|
||||
.SetInplaceMap({{"output_", "output"}})
|
||||
.Attrs({"is_before_index: bool"})
|
||||
.SetKernelFn(PD_KERNEL(DynamicSetValue));
|
||||
|
||||
|
||||
|
||||
// std::vector<paddle::Tensor> DynamicSplitValue(const paddle::Tensor& input,
|
||||
// const paddle::Tensor& index,
|
||||
// bool is_before_index) { // true: [ : index, :] ; false: [index : , :]
|
||||
|
||||
// phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
|
||||
// auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
|
||||
// auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
|
||||
// using XPUType = typename XPUTypeTrait<bfloat16>::Type; // only support bfloat16
|
||||
// typedef paddle::bfloat16 data_t;
|
||||
|
||||
|
||||
|
||||
|
||||
// int index_num = index.data<int64_t>()[0];
|
||||
// int total_len = input.dims()[0];
|
||||
// int dim = input.dims()[1];
|
||||
// int copy_len = 0;
|
||||
|
||||
// if (is_before_index) {
|
||||
// copy_len = index_num
|
||||
// }
|
||||
// else{
|
||||
// copy_len = total_len - index_num
|
||||
// }
|
||||
|
||||
// auto output = paddle::full({copy_len, dim}, 0, seq_lens_encoder.type(), seq_lens_encoder.place());
|
||||
// assert(dim == input.dims()[1] && "input.dim != output.dim! ");
|
||||
// assert((index_num >= 0 && index_num <= total_len) && "index out of range! ");
|
||||
// if (is_before_index) {
|
||||
// assert((index_num - 0) * dim == copy_len && "copy_len wrong! ");
|
||||
// }
|
||||
// else{
|
||||
// assert((total_len - index_num) * dim == copy_len && "copy_len wrong! ");
|
||||
// }
|
||||
// int r = baidu::xpu::api::copy<XPUType>(
|
||||
// xpu_ctx->x_context(),
|
||||
// reinterpret_cast<const XPUType*>(input.data<data_t>()),
|
||||
// reinterpret_cast<XPUType*>(output->data<data_t>() + (is_before_index ? 0 : index_num * dim)),
|
||||
// copy_len);
|
||||
|
||||
// return {};
|
||||
// }
|
||||
72
csrc/xpu/src/fused_rotary_position_encoding.cc
Normal file
72
csrc/xpu/src/fused_rotary_position_encoding.cc
Normal file
@@ -0,0 +1,72 @@
|
||||
// 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/backends/xpu/xpu_context.h>
|
||||
#include "paddle/extension.h"
|
||||
#include "xpu/plugin.h"
|
||||
#include "ops.h"
|
||||
#include <xft/xdnn_plugin.h>
|
||||
|
||||
|
||||
|
||||
void FusedRotaryPositionEncoding(
|
||||
paddle::Tensor& query, // [num_tokens, num_heads, head_size] or
|
||||
// [num_tokens, num_heads * head_size]
|
||||
paddle::Tensor& key,
|
||||
// [num_tokens, num_kv_heads, head_size] or [num_tokens, num_kv_heads *
|
||||
// head_size]
|
||||
const paddle::Tensor& position_ids, // [num_tokens]
|
||||
const paddle::Tensor& cos_sin_cache, // [max_position, rot_dim]
|
||||
int head_size,
|
||||
bool is_neox) {
|
||||
|
||||
baidu::xpu::api::plugin::print_times("[TIME BEGIN] FusedRotaryPositionEncoding" );
|
||||
|
||||
phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
|
||||
auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
|
||||
auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
|
||||
xpu::ctx_guard RAII_GUARD(xpu_ctx->x_context());
|
||||
|
||||
typedef paddle::bfloat16 data_t;
|
||||
using XPUType = typename XPUTypeTrait<data_t>::Type;
|
||||
|
||||
int64_t num_tokens = query.dims()[0];
|
||||
int num_heads = query.numel() / num_tokens / head_size;
|
||||
int num_kv_heads = key.numel() / num_tokens / head_size;
|
||||
int rot_dim = cos_sin_cache.dims()[1];
|
||||
int64_t query_stride = num_heads * head_size;
|
||||
int64_t key_stride = num_kv_heads * head_size;
|
||||
|
||||
baidu::xpu::api::plugin::rotary_embedding_neox<XPUType, XPUType>(xpu_ctx->x_context(),
|
||||
position_ids.data<int>(),
|
||||
reinterpret_cast<XPUType*>(query.data<data_t>()),
|
||||
reinterpret_cast<XPUType*>(key.data<data_t>()),
|
||||
reinterpret_cast<const XPUType*>(cos_sin_cache.data<data_t>()),
|
||||
rot_dim,
|
||||
num_heads,
|
||||
num_kv_heads,
|
||||
head_size,
|
||||
num_tokens);
|
||||
baidu::xpu::api::plugin::print_times("[TIME END] FusedRotaryPositionEncoding" );
|
||||
|
||||
}
|
||||
|
||||
|
||||
PD_BUILD_OP(fused_rotary_position_encoding)
|
||||
.Inputs({"query", "key", "position_ids", "cos_sin_cache"})
|
||||
.Outputs({"query_out", "key_out"})
|
||||
.Attrs({"head_size: int", "is_neox: bool"})
|
||||
.SetInplaceMap({{"query", "query_out"}, {"key", "key_out"}})
|
||||
.SetKernelFn(PD_KERNEL(FusedRotaryPositionEncoding));
|
||||
|
||||
@@ -19,62 +19,33 @@
|
||||
namespace xftkernel = baidu::xpu::xftkernel;
|
||||
std::vector<paddle::Tensor> GatherNextToken(const paddle::Tensor& tmp_out, // [token_num, dim_embed]
|
||||
const paddle::Tensor& cum_offsets, // [bsz, 1]
|
||||
const paddle::Tensor& seq_lens_decoder,
|
||||
const paddle::Tensor& seq_lens_encoder,
|
||||
const paddle::Tensor& encoder_seq_lod,
|
||||
const paddle::Tensor& encoder_batch_map,
|
||||
const paddle::Tensor& decoder_batch_map,
|
||||
const paddle::Tensor& encoder_seq_lod_cpu,
|
||||
const paddle::Tensor& encoder_batch_map_cpu,
|
||||
const paddle::Tensor& decoder_batch_map_cpu,
|
||||
const paddle::Tensor& enc_batch_tensor,
|
||||
const paddle::Tensor& dec_batch_tensor,
|
||||
const paddle::optional<paddle::Tensor>& output_padding_offset,
|
||||
int max_input_length) {
|
||||
phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
|
||||
auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
|
||||
auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
|
||||
xpu::ctx_guard RAII_GUARD(xpu_ctx->x_context());
|
||||
using XPUType = typename XPUTypeTrait<bfloat16>::Type; // only support bfloat16
|
||||
typedef paddle::bfloat16 data_t;
|
||||
const int dim = tmp_out.dims()[1];
|
||||
const int bsz = cum_offsets.shape()[0];
|
||||
int enc_batch = enc_batch_tensor.data<int32_t>()[0];
|
||||
int dec_batch = dec_batch_tensor.data<int32_t>()[0];
|
||||
|
||||
baidu::xpu::api::VectorParam<int32_t> encoder_seqs_lods_vp{const_cast<int32_t*>(encoder_seq_lod_cpu.data<int32_t>()), enc_batch + 1, const_cast<int32_t*>(encoder_seq_lod.data<int32_t>())};
|
||||
baidu::xpu::api::VectorParam<int32_t> encoder_batch_map_vp{const_cast<int32_t*>(encoder_batch_map_cpu.data<int32_t>()), enc_batch, const_cast<int32_t*>(encoder_batch_map.data<int32_t>())};
|
||||
baidu::xpu::api::VectorParam<int32_t> decoder_batch_map_vp{const_cast<int32_t*>(decoder_batch_map_cpu.data<int32_t>()), dec_batch, const_cast<int32_t*>(decoder_batch_map.data<int32_t>())};
|
||||
|
||||
std::vector<int> seq_lens_encoder_cpu(bsz, 0);
|
||||
std::vector<int> seq_lens_decoder_cpu(bsz, 0);
|
||||
std::vector<int> encoder_batch_map; // 去除空隙的batch map
|
||||
std::vector<int> decoder_batch_map; // 去除空隙的batch map
|
||||
std::vector<int> encoder_seq_lod;
|
||||
int r = xpu_memcpy(seq_lens_encoder_cpu.data(),
|
||||
seq_lens_encoder.data<int>(),
|
||||
sizeof(int32_t) * bsz,
|
||||
XPUMemcpyKind::XPU_DEVICE_TO_HOST);
|
||||
r = xpu_memcpy(seq_lens_decoder_cpu.data(),
|
||||
seq_lens_decoder.data<int>(),
|
||||
sizeof(int32_t) * bsz,
|
||||
XPUMemcpyKind::XPU_DEVICE_TO_HOST);
|
||||
int enc_batch = 0, dec_batch = 0;
|
||||
int batch_offset = 0;
|
||||
encoder_seq_lod.push_back(0);
|
||||
for(int i = 0; i < bsz; ++i){
|
||||
if(seq_lens_encoder_cpu[i] > 0){
|
||||
enc_batch++;
|
||||
encoder_batch_map.push_back(i);
|
||||
encoder_seq_lod.push_back(seq_lens_encoder_cpu[i]);
|
||||
encoder_seq_lod[enc_batch] += encoder_seq_lod[enc_batch - 1];
|
||||
}
|
||||
else if(seq_lens_decoder_cpu[i] > 0){
|
||||
dec_batch++;
|
||||
decoder_batch_map.push_back(i);
|
||||
}
|
||||
else{
|
||||
batch_offset++;
|
||||
}
|
||||
}
|
||||
baidu::xpu::api::VectorParam<int32_t> encoder_seqs_lods_vp =
|
||||
baidu::xpu::api::VectorParam<int32_t>{encoder_seq_lod.data(), enc_batch + 1, nullptr}
|
||||
.to_xpu(RAII_GUARD);
|
||||
baidu::xpu::api::VectorParam<int32_t> encoder_batch_map_vp =
|
||||
baidu::xpu::api::VectorParam<int32_t>{encoder_batch_map.data(), enc_batch, nullptr}
|
||||
.to_xpu(RAII_GUARD);
|
||||
baidu::xpu::api::VectorParam<int32_t> decoder_batch_map_vp =
|
||||
baidu::xpu::api::VectorParam<int32_t>{decoder_batch_map.data(), dec_batch, nullptr}
|
||||
.to_xpu(RAII_GUARD);
|
||||
auto out = paddle::full({bsz, dim}, -2, tmp_out.type(), tmp_out.place());
|
||||
|
||||
r = xftkernel::xft_eb_gather_next_token<XPUType, XPUType>(
|
||||
int r = xftkernel::xft_eb_gather_next_token<XPUType, XPUType>(
|
||||
xpu_ctx->x_context(),
|
||||
reinterpret_cast<const XPUType*>(tmp_out.data<data_t>()),
|
||||
reinterpret_cast<XPUType*>(out.data<data_t>()),
|
||||
@@ -87,8 +58,14 @@ std::vector<paddle::Tensor> GatherNextToken(const paddle::Tensor& tmp_out, // [t
|
||||
|
||||
std::vector<std::vector<int64_t>> GatherNextTokenInferShape(const std::vector<int64_t>& tmp_out_shape,
|
||||
const std::vector<int64_t>& cum_offsets_shape,
|
||||
const std::vector<int64_t>& seq_lens_decoder_shape,
|
||||
const std::vector<int64_t>& seq_lens_encoder_shape,
|
||||
const std::vector<int64_t>& encoder_seq_lod_shape,
|
||||
const std::vector<int64_t>& encoder_batch_map_shape,
|
||||
const std::vector<int64_t>& decoder_batch_map_shape,
|
||||
const std::vector<int64_t>& encoder_seq_lod_cpu_shape,
|
||||
const std::vector<int64_t>& encoder_batch_map_cpu_shape,
|
||||
const std::vector<int64_t>& decoder_batch_map_cpu_shape,
|
||||
const std::vector<int64_t>& enc_batch_tensor_shape,
|
||||
const std::vector<int64_t>& dec_batch_tensor_shape,
|
||||
const paddle::optional<std::vector<int64_t>>& output_padding_offset_shape) {
|
||||
if (output_padding_offset_shape) {
|
||||
PD_THROW("speculative decoding is not supported in XPU.");
|
||||
@@ -100,14 +77,20 @@ std::vector<std::vector<int64_t>> GatherNextTokenInferShape(const std::vector<in
|
||||
|
||||
std::vector<paddle::DataType> GatherNextTokenInferDtype(const paddle::DataType& tmp_out_dtype,
|
||||
const paddle::DataType& cum_offsets_dtype,
|
||||
const paddle::DataType& seq_lens_decoder_dtype,
|
||||
const paddle::DataType& seq_lens_encoder_dtype,
|
||||
const paddle::DataType& encoder_seq_lod_dtype,
|
||||
const paddle::DataType& encoder_batch_map_dtype,
|
||||
const paddle::DataType& decoder_batch_map_dtype,
|
||||
const paddle::DataType& encoder_seq_lod_cpu_dtype,
|
||||
const paddle::DataType& encoder_batch_map_cpu_dtype,
|
||||
const paddle::DataType& decoder_batch_map_cpu_dtype,
|
||||
const paddle::DataType& enc_batch_tensor_dtype,
|
||||
const paddle::DataType& dec_batch_tensor_dtype,
|
||||
const paddle::optional<paddle::DataType>& output_padding_offset_dtype) {
|
||||
return {tmp_out_dtype};
|
||||
}
|
||||
|
||||
PD_BUILD_OP(gather_next_token)
|
||||
.Inputs({"tmp_out", "cum_offsets", "seq_lens_decoder", "seq_lens_encoder", paddle::Optional("output_padding_offset")})
|
||||
.Inputs({"tmp_out", "cum_offsets", "encoder_seq_lod", "encoder_batch_map", "decoder_batch_map","encoder_seq_lod_cpu", "encoder_batch_map_cpu", "decoder_batch_map_cpu", "enc_batch_tensor", "dec_batch_tensor", paddle::Optional("output_padding_offset")})
|
||||
.Outputs({"out"})
|
||||
.Attrs({"max_input_length: int"})
|
||||
.SetKernelFn(PD_KERNEL(GatherNextToken))
|
||||
|
||||
196
csrc/xpu/src/get_infer_param.cc
Normal file
196
csrc/xpu/src/get_infer_param.cc
Normal file
@@ -0,0 +1,196 @@
|
||||
// Copyright (c) 2024 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/backends/xpu/xpu_context.h>
|
||||
#include "paddle/extension.h"
|
||||
#include "xpu/plugin.h"
|
||||
#include <xft/xdnn_plugin.h>
|
||||
namespace xftkernel = baidu::xpu::xftkernel;
|
||||
std::vector<paddle::Tensor> GetInferParam(const paddle::Tensor& seq_lens_encoder,
|
||||
const paddle::Tensor& seq_lens_decoder) {
|
||||
baidu::xpu::api::plugin::print_times("[TIME BEGIN] GetInferParam" );
|
||||
|
||||
phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
|
||||
auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
|
||||
auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
|
||||
const int bsz = seq_lens_encoder.dims()[0];
|
||||
// 判断逻辑
|
||||
std::vector<int32_t> seq_lens_encoder_vec(bsz, 0); // input
|
||||
std::vector<int32_t> seq_lens_decoder_vec(bsz, 0); // input
|
||||
|
||||
std::vector<int32_t> encoder_batch_map_vec(bsz, 0); //
|
||||
std::vector<int32_t> decoder_batch_map_vec(bsz, 0); //
|
||||
std::vector<int32_t> encoder_batch_idx_vec(bsz, 0); // 去除空隙的batch map
|
||||
std::vector<int32_t> decoder_batch_idx_vec(bsz, 0); // 去除空隙的batch map
|
||||
std::vector<int32_t> encoder_seq_lod_vec(bsz + 1, 0);
|
||||
std::vector<int32_t> decoder_context_len_vec(bsz, 0);
|
||||
std::vector<int32_t> decoder_context_len_cache_vec(bsz, 0);
|
||||
xpu_wait(xpu_ctx->x_context()->xpu_stream); // 是否需要！！！！TODO
|
||||
int r = xpu_memcpy(seq_lens_encoder_vec.data(),
|
||||
seq_lens_encoder.data<int32_t>(),
|
||||
sizeof(int32_t) * bsz,
|
||||
XPUMemcpyKind::XPU_DEVICE_TO_HOST);
|
||||
r = xpu_memcpy(seq_lens_decoder_vec.data(),
|
||||
seq_lens_decoder.data<int32_t>(),
|
||||
sizeof(int32_t) * bsz,
|
||||
XPUMemcpyKind::XPU_DEVICE_TO_HOST);
|
||||
|
||||
int enc_batch = 0, dec_batch = 0;
|
||||
int total_enc_len = 0;
|
||||
int batch_offset = 0;
|
||||
for(int i = 0; i < bsz; ++i){
|
||||
if(seq_lens_encoder_vec[i] > 0){
|
||||
enc_batch++;
|
||||
total_enc_len += seq_lens_encoder_vec[i];
|
||||
encoder_batch_map_vec[enc_batch - 1] = i;
|
||||
encoder_batch_idx_vec[enc_batch - 1] = i - batch_offset;
|
||||
encoder_seq_lod_vec[enc_batch] = seq_lens_encoder_vec[i] + encoder_seq_lod_vec[enc_batch - 1];
|
||||
}
|
||||
else if(seq_lens_decoder_vec[i] > 0){
|
||||
dec_batch++;
|
||||
decoder_batch_map_vec[dec_batch - 1] = i;
|
||||
decoder_batch_idx_vec[dec_batch - 1] = i - batch_offset;
|
||||
decoder_context_len_vec[dec_batch - 1] = seq_lens_decoder_vec[i] + 1;
|
||||
decoder_context_len_cache_vec[dec_batch - 1] = seq_lens_decoder_vec[i];
|
||||
}
|
||||
else{
|
||||
batch_offset++;
|
||||
}
|
||||
}
|
||||
|
||||
auto encoder_batch_map_xpu = paddle::full({encoder_batch_map_vec.size()}, 0, seq_lens_encoder.type(), seq_lens_encoder.place());
|
||||
auto decoder_batch_map_xpu = paddle::full({decoder_batch_map_vec.size()}, 0, seq_lens_encoder.type(), seq_lens_encoder.place());
|
||||
auto encoder_batch_idx_xpu = paddle::full({encoder_batch_idx_vec.size()}, 0, seq_lens_encoder.type(), seq_lens_encoder.place());
|
||||
auto decoder_batch_idx_xpu = paddle::full({decoder_batch_idx_vec.size()}, 0, seq_lens_encoder.type(), seq_lens_encoder.place());
|
||||
auto encoder_seq_lod_xpu = paddle::full({encoder_seq_lod_vec.size()}, 0, seq_lens_encoder.type(), seq_lens_encoder.place());
|
||||
auto decoder_context_len_xpu = paddle::full({decoder_context_len_vec.size()}, 0, seq_lens_encoder.type(), seq_lens_encoder.place());
|
||||
auto decoder_context_len_cache_xpu = paddle::full({decoder_context_len_cache_vec.size()}, 0, seq_lens_encoder.type(), seq_lens_encoder.place());
|
||||
|
||||
auto encoder_batch_map_cpu = paddle::full({encoder_batch_map_vec.size()}, 0, seq_lens_encoder.type(), paddle::CPUPlace());
|
||||
auto decoder_batch_map_cpu = paddle::full({decoder_batch_map_vec.size()}, 0, seq_lens_encoder.type(), paddle::CPUPlace());
|
||||
auto encoder_batch_idx_cpu = paddle::full({encoder_batch_idx_vec.size()}, 0, seq_lens_encoder.type(), paddle::CPUPlace());
|
||||
auto decoder_batch_idx_cpu = paddle::full({decoder_batch_idx_vec.size()}, 0, seq_lens_encoder.type(), paddle::CPUPlace());
|
||||
auto encoder_seq_lod_cpu = paddle::full({encoder_seq_lod_vec.size()}, 0, seq_lens_encoder.type(), paddle::CPUPlace());
|
||||
auto decoder_context_len_cpu = paddle::full({decoder_context_len_vec.size()}, 0, seq_lens_encoder.type(), paddle::CPUPlace());
|
||||
auto decoder_context_len_cache_cpu = paddle::full({decoder_context_len_cache_vec.size()}, 0, seq_lens_encoder.type(), paddle::CPUPlace());
|
||||
|
||||
int ret = 0;
|
||||
ret = xpu_memcpy(reinterpret_cast<int32_t*>(const_cast<int32_t*>(encoder_batch_map_xpu.data<int32_t>())), encoder_batch_map_vec.data(), sizeof(int32_t) * encoder_batch_map_vec.size(), XPUMemcpyKind::XPU_HOST_TO_DEVICE);
|
||||
ret = xpu_memcpy(reinterpret_cast<int32_t*>(const_cast<int32_t*>(decoder_batch_map_xpu.data<int32_t>())), decoder_batch_map_vec.data(), sizeof(int32_t) * decoder_batch_map_vec.size(), XPUMemcpyKind::XPU_HOST_TO_DEVICE);
|
||||
ret = xpu_memcpy(reinterpret_cast<int32_t*>(const_cast<int32_t*>(encoder_batch_idx_xpu.data<int32_t>())), encoder_batch_idx_vec.data(), sizeof(int32_t) * encoder_batch_idx_vec.size(), XPUMemcpyKind::XPU_HOST_TO_DEVICE);
|
||||
ret = xpu_memcpy(reinterpret_cast<int32_t*>(const_cast<int32_t*>(decoder_batch_idx_xpu.data<int32_t>())), decoder_batch_idx_vec.data(), sizeof(int32_t) * decoder_batch_idx_vec.size(), XPUMemcpyKind::XPU_HOST_TO_DEVICE);
|
||||
ret = xpu_memcpy(reinterpret_cast<int32_t*>(const_cast<int32_t*>(encoder_seq_lod_xpu.data<int32_t>())), encoder_seq_lod_vec.data(), sizeof(int32_t) * encoder_seq_lod_vec.size(), XPUMemcpyKind::XPU_HOST_TO_DEVICE);
|
||||
ret = xpu_memcpy(reinterpret_cast<int32_t*>(const_cast<int32_t*>(decoder_context_len_xpu.data<int32_t>())), decoder_context_len_vec.data(), sizeof(int32_t) * decoder_context_len_vec.size(), XPUMemcpyKind::XPU_HOST_TO_DEVICE);
|
||||
ret = xpu_memcpy(reinterpret_cast<int32_t*>(const_cast<int32_t*>(decoder_context_len_cache_xpu.data<int32_t>())), decoder_context_len_cache_vec.data(), sizeof(int32_t) * decoder_context_len_cache_vec.size(), XPUMemcpyKind::XPU_HOST_TO_DEVICE);
|
||||
|
||||
std::memcpy(encoder_batch_map_cpu.data<int32_t>(), encoder_batch_map_vec.data(), sizeof(int32_t) * encoder_batch_map_vec.size());
|
||||
std::memcpy(decoder_batch_map_cpu.data<int32_t>(), decoder_batch_map_vec.data(), sizeof(int32_t) * decoder_batch_map_vec.size());
|
||||
std::memcpy(encoder_batch_idx_cpu.data<int32_t>(), encoder_batch_idx_vec.data(), sizeof(int32_t) * encoder_batch_idx_vec.size());
|
||||
std::memcpy(decoder_batch_idx_cpu.data<int32_t>(), decoder_batch_idx_vec.data(), sizeof(int32_t) * decoder_batch_idx_vec.size());
|
||||
std::memcpy(encoder_seq_lod_cpu.data<int32_t>(), encoder_seq_lod_vec.data(), sizeof(int32_t) * encoder_seq_lod_vec.size());
|
||||
std::memcpy(decoder_context_len_cpu.data<int32_t>(), decoder_context_len_vec.data(), sizeof(int32_t) * decoder_context_len_vec.size());
|
||||
std::memcpy(decoder_context_len_cache_cpu.data<int32_t>(), decoder_context_len_cache_vec.data(), sizeof(int32_t) * decoder_context_len_cache_vec.size());
|
||||
|
||||
// auto batch_tensor = paddle::full({3}, 0, seq_lens_encoder.type(), paddle::CPUPlace());
|
||||
// auto batch_tensor_data = const_cast<int32_t*>(batch_tensor.data<int32_t>());
|
||||
// batch_tensor_data[0] = enc_batch;
|
||||
// batch_tensor_data[1] = dec_batch;
|
||||
// batch_tensor_data[2] = encoder_seq_lod_vec[enc_batch];
|
||||
auto enc_batch_tensor = paddle::full({1}, enc_batch, seq_lens_encoder.type(), paddle::CPUPlace());
|
||||
auto dec_batch_tensor = paddle::full({1}, dec_batch, seq_lens_encoder.type(), paddle::CPUPlace());
|
||||
auto total_enc_len_tensor = paddle::full({1}, total_enc_len, seq_lens_encoder.type(), paddle::CPUPlace());
|
||||
|
||||
|
||||
baidu::xpu::api::plugin::print_times("[TIME END] GetInferParam" );
|
||||
return {encoder_batch_map_xpu,
|
||||
decoder_batch_map_xpu,
|
||||
encoder_batch_idx_xpu,
|
||||
decoder_batch_idx_xpu,
|
||||
encoder_seq_lod_xpu,
|
||||
decoder_context_len_xpu,
|
||||
decoder_context_len_cache_xpu,
|
||||
encoder_batch_map_cpu,
|
||||
decoder_batch_map_cpu,
|
||||
encoder_batch_idx_cpu,
|
||||
decoder_batch_idx_cpu,
|
||||
encoder_seq_lod_cpu,
|
||||
decoder_context_len_cpu,
|
||||
decoder_context_len_cache_cpu,
|
||||
enc_batch_tensor,
|
||||
dec_batch_tensor,
|
||||
total_enc_len_tensor};
|
||||
}
|
||||
|
||||
std::vector<std::vector<int64_t>> GetInferParamInferShape(const std::vector<int64_t>& seq_lens_encoder_shape,
|
||||
const std::vector<int64_t>& seq_lens_decoder_shape) {
|
||||
return {seq_lens_encoder_shape,
|
||||
seq_lens_encoder_shape,
|
||||
seq_lens_encoder_shape,
|
||||
seq_lens_encoder_shape,
|
||||
{seq_lens_encoder_shape[0] + 1},
|
||||
seq_lens_encoder_shape,
|
||||
seq_lens_encoder_shape,
|
||||
seq_lens_encoder_shape,
|
||||
seq_lens_encoder_shape,
|
||||
seq_lens_encoder_shape,
|
||||
seq_lens_encoder_shape,
|
||||
seq_lens_encoder_shape,
|
||||
seq_lens_encoder_shape,
|
||||
seq_lens_encoder_shape,
|
||||
{1},{1},{1}};
|
||||
}
|
||||
|
||||
std::vector<paddle::DataType> GetInferParamInferDtype(const paddle::DataType& seq_lens_encoder_dtype,
|
||||
const paddle::DataType& seq_lens_decoder_dtype) {
|
||||
return {seq_lens_encoder_dtype,
|
||||
seq_lens_encoder_dtype,
|
||||
seq_lens_encoder_dtype,
|
||||
seq_lens_encoder_dtype,
|
||||
seq_lens_encoder_dtype,
|
||||
seq_lens_encoder_dtype,
|
||||
seq_lens_encoder_dtype,
|
||||
seq_lens_encoder_dtype,
|
||||
seq_lens_encoder_dtype,
|
||||
seq_lens_encoder_dtype,
|
||||
seq_lens_encoder_dtype,
|
||||
seq_lens_encoder_dtype,
|
||||
seq_lens_encoder_dtype,
|
||||
seq_lens_encoder_dtype,
|
||||
seq_lens_encoder_dtype,
|
||||
seq_lens_encoder_dtype,
|
||||
seq_lens_encoder_dtype};
|
||||
}
|
||||
|
||||
PD_BUILD_OP(get_infer_param)
|
||||
.Inputs({"seq_lens_encoder", "seq_lens_decoder"})
|
||||
.Outputs({"encoder_batch_map_xpu",
|
||||
"decoder_batch_map_xpu",
|
||||
"encoder_batch_idx_xpu",
|
||||
"decoder_batch_idx_xpu",
|
||||
"encoder_seq_lod_xpu",
|
||||
"decoder_context_len_xpu",
|
||||
"decoder_context_len_cache_xpu",
|
||||
"encoder_batch_map_cpu",
|
||||
"decoder_batch_map_cpu",
|
||||
"encoder_batch_idx_cpu",
|
||||
"decoder_batch_idx_cpu",
|
||||
"encoder_seq_lod_cpu",
|
||||
"decoder_context_len_cpu",
|
||||
"decoder_context_len_cache_cpu",
|
||||
"enc_batch_tensor",
|
||||
"dec_batch_tensor",
|
||||
"total_enc_len_tensor"})
|
||||
.SetKernelFn(PD_KERNEL(GetInferParam))
|
||||
.SetInferShapeFn(PD_INFER_SHAPE(GetInferParamInferShape))
|
||||
.SetInferDtypeFn(PD_INFER_DTYPE(GetInferParamInferDtype));
|
||||
70
csrc/xpu/src/get_output.cc
Normal file
70
csrc/xpu/src/get_output.cc
Normal file
@@ -0,0 +1,70 @@
|
||||
// Copyright (c) 2024 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 <stdio.h>
|
||||
#include <string.h>
|
||||
#include <sys/ipc.h>
|
||||
#include <sys/msg.h>
|
||||
#include <sys/types.h>
|
||||
#include "paddle/extension.h"
|
||||
|
||||
#define MAX_BSZ 512
|
||||
|
||||
struct msgdata {
|
||||
long mtype;
|
||||
int mtext[MAX_BSZ + 2]; // stop_flag, bsz, tokens
|
||||
};
|
||||
|
||||
void GetOutput(const paddle::Tensor& x,
|
||||
const paddle::Tensor& msg_queue_id,
|
||||
int64_t rank_id,
|
||||
bool wait_flag) {
|
||||
if (rank_id > 0) return;
|
||||
|
||||
static struct msgdata msg_rcv;
|
||||
int msg_queue_id_val = msg_queue_id.data<int>()[0];
|
||||
static key_t key = ftok("./", msg_queue_id_val);
|
||||
|
||||
static int msgid = msgget(key, IPC_CREAT | 0666);
|
||||
|
||||
int64_t *out_data = const_cast<int64_t*>(x.data<int64_t>());
|
||||
int ret = -1;
|
||||
if (!wait_flag) {
|
||||
ret = msgrcv(msgid, &msg_rcv, (MAX_BSZ + 2) * 4, 0, IPC_NOWAIT);
|
||||
} else {
|
||||
ret = msgrcv(msgid, &msg_rcv, (MAX_BSZ + 2) * 4, 0, 0);
|
||||
}
|
||||
if(ret == -1)
|
||||
{
|
||||
// read none
|
||||
out_data[0] = -2;
|
||||
out_data[1] = 0;
|
||||
return;
|
||||
}
|
||||
|
||||
int bsz = msg_rcv.mtext[1];
|
||||
|
||||
for (int64_t i = 0; i < bsz + 2; i++) {
|
||||
out_data[i] = (int64_t)msg_rcv.mtext[i];
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
PD_BUILD_OP(get_output)
|
||||
.Inputs({"x", "msg_queue_id"})
|
||||
.Attrs({"rank_id: int64_t",
|
||||
"wait_flag: bool"})
|
||||
.Outputs({"x_out"})
|
||||
.SetInplaceMap({{"x", "x_out"}})
|
||||
.SetKernelFn(PD_KERNEL(GetOutput));
|
||||
@@ -28,6 +28,8 @@ std::vector<paddle::Tensor> GetPositionIdsKernelV2(
|
||||
const paddle::Tensor& seq_lens_decoder,
|
||||
const paddle::Tensor& seq_lens_this_time) {
|
||||
|
||||
baidu::xpu::api::plugin::print_times("[TIME BEGIN] GetPositionIdsKernelV2" );
|
||||
|
||||
phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
|
||||
auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
|
||||
auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
|
||||
@@ -54,7 +56,9 @@ std::vector<paddle::Tensor> GetPositionIdsKernelV2(
|
||||
const_cast<int32_t*>(position_ids_out.data<int32_t>()),
|
||||
bs
|
||||
);
|
||||
PD_CHECK(ret == 0, "api::plugin::get_position_ids failed");
|
||||
PD_CHECK(ret == 0, "api::plugin::get_position_ids failed");
|
||||
|
||||
baidu::xpu::api::plugin::print_times("[TIME END] GetPositionIdsKernelV2" );
|
||||
return {position_ids_out};
|
||||
}
|
||||
|
||||
|
||||
@@ -27,6 +27,7 @@ void TokenPenaltyMultiScores(const paddle::Tensor& pre_ids,
|
||||
const paddle::Tensor& cur_len,
|
||||
const paddle::Tensor& min_len,
|
||||
const paddle::Tensor& eos_token_id) {
|
||||
baidu::xpu::api::plugin::print_times("[TIME BEGIN] TokenPenaltyMultiScores");
|
||||
phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
|
||||
auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
|
||||
auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
|
||||
@@ -90,8 +91,11 @@ void TokenPenaltyMultiScores(const paddle::Tensor& pre_ids,
|
||||
"Only float16 and float32 are supported. ");
|
||||
break;
|
||||
}
|
||||
|
||||
baidu::xpu::api::plugin::print_times("[TIME END] TokenPenaltyMultiScores");
|
||||
}
|
||||
|
||||
|
||||
PD_BUILD_OP(get_token_penalty_multi_scores_v2)
|
||||
.Inputs({"pre_ids",
|
||||
"logits",
|
||||
|
||||
@@ -323,7 +323,8 @@ std::vector<paddle::Tensor> MlaAttn(
|
||||
block_size, // block_size
|
||||
max_block_per_seq, // max_blocks_per_seq (prefix cache)
|
||||
page_param_.max_context_len_, // prefill_len
|
||||
nullptr); // block_tables (prefix cache)
|
||||
nullptr,
|
||||
softmax_scale * sqrt(dim_qk)); // block_tables (prefix cache)
|
||||
|
||||
// std::cout << "fmha kernel done " <<std::endl;
|
||||
}
|
||||
@@ -425,7 +426,7 @@ std::vector<paddle::Tensor> MlaAttn(
|
||||
const_cast<TKV*>(key_cache_xft.data()),
|
||||
const_cast<TKV*>(value_cache_xft.data()),
|
||||
attn_param_.kv_head_num_,
|
||||
attn_param_.scale_,
|
||||
softmax_scale,
|
||||
block_tables_xft.data(),
|
||||
attn_param_.context_len_vp_,
|
||||
attn_param_.valid_batch_vp_,
|
||||
|
||||
477
csrc/xpu/src/mla_cache_kernel_xpu.cc
Normal file
477
csrc/xpu/src/mla_cache_kernel_xpu.cc
Normal file
@@ -0,0 +1,477 @@
|
||||
// Copyright (c) 2024 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 <core/ctx_manager.h>
|
||||
#include <core/xft_check.h>
|
||||
#include <core/xft_event.h>
|
||||
#include <core/xft_params.h>
|
||||
#include <paddle/phi/backends/xpu/xpu_context.h>
|
||||
#include <xft/operation/xft_fc_helper.h>
|
||||
#include <xft/xdnn_plugin.h>
|
||||
|
||||
#include "paddle/extension.h"
|
||||
#include "paddle/phi/core/enforce.h"
|
||||
#include "xpu/plugin.h"
|
||||
#include "ops.h"
|
||||
|
||||
namespace xftkernel = baidu::xpu::xftkernel;
|
||||
|
||||
template <paddle::DataType D>
|
||||
class PaddleTypeToXPUType;
|
||||
|
||||
template <>
|
||||
class PaddleTypeToXPUType<paddle::DataType::FLOAT32> {
|
||||
public:
|
||||
typedef float DataType;
|
||||
typedef float data_t;
|
||||
};
|
||||
|
||||
template <>
|
||||
class PaddleTypeToXPUType<paddle::DataType::FLOAT16> {
|
||||
public:
|
||||
typedef float16 DataType;
|
||||
typedef paddle::float16 data_t;
|
||||
};
|
||||
|
||||
template <>
|
||||
class PaddleTypeToXPUType<paddle::DataType::BFLOAT16> {
|
||||
public:
|
||||
typedef bfloat16 DataType;
|
||||
typedef paddle::bfloat16 data_t;
|
||||
};
|
||||
|
||||
template <paddle::DataType T>
|
||||
std::vector<paddle::Tensor> PrefillMLAWriteCache(
|
||||
const paddle::Tensor& kv_nope,
|
||||
const paddle::Tensor& kv_pe,
|
||||
const paddle::Tensor& encoder_seq_lod,
|
||||
const paddle::Tensor& encoder_batch_map,
|
||||
const paddle::Tensor& start_token_raw, // encoder cache写时 start token为0 （不算prefix cache）
|
||||
const paddle::Tensor& encoder_seq_lod_cpu,
|
||||
const paddle::Tensor& encoder_batch_map_cpu,
|
||||
const paddle::Tensor& start_token_raw_cpu, // encoder cache写时 start token为0 （不算prefix cache）
|
||||
const paddle::Tensor& enc_batch_tensor,
|
||||
const paddle::Tensor& dec_batch_tensor,
|
||||
const paddle::Tensor& padding_offsets,
|
||||
const paddle::Tensor& cum_offsets,
|
||||
const paddle::Tensor& block_tables,
|
||||
const int max_seq_len,
|
||||
const int kv_num_heads,
|
||||
paddle::Tensor* kv_cache) {
|
||||
baidu::xpu::api::plugin::print_times("[TIME BEGIN] PrefillMLAWriteCache");
|
||||
|
||||
phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
|
||||
auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
|
||||
auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
|
||||
typedef PaddleTypeToXPUType<T> traits_;
|
||||
typedef typename traits_::DataType DataType_;
|
||||
typedef typename traits_::data_t data_t;
|
||||
|
||||
const auto& kv_nope_dims = kv_nope.dims();
|
||||
const auto& kv_pe_dims = kv_pe.dims();
|
||||
const auto& kv_cache_dims = kv_cache->dims();
|
||||
auto num_tokens = kv_nope_dims[0];
|
||||
auto block_size = kv_cache_dims[2];
|
||||
auto block_batch = block_tables.dims()[0];
|
||||
auto max_blocks_per_seq = block_tables.dims()[1];
|
||||
auto kv_lora_rank = kv_nope_dims[kv_nope_dims.size() - 1];
|
||||
auto rope_head_dim = kv_pe_dims[kv_pe_dims.size() - 1];
|
||||
auto bsz = cum_offsets.dims()[0];
|
||||
int enc_batch = enc_batch_tensor.data<int32_t>()[0];
|
||||
int dec_batch = dec_batch_tensor.data<int32_t>()[0];
|
||||
// 初始化输入
|
||||
auto kv_nope_xft = baidu::xpu::xft::xftTensor<DataType_, 2>(
|
||||
reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_nope.data<data_t>())),
|
||||
std::array<int64_t, 2>{num_tokens,
|
||||
kv_lora_rank});
|
||||
auto kv_pe_xft = baidu::xpu::xft::xftTensor<DataType_, 2>(
|
||||
reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_pe.data<data_t>())),
|
||||
std::array<int64_t, 2>{num_tokens,
|
||||
rope_head_dim});
|
||||
// 初始化输入 cache
|
||||
auto cache_xft = baidu::xpu::xft::xftTensor<DataType_, 4>(
|
||||
reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_cache->data<data_t>())),
|
||||
std::array<int64_t, 4>{kv_cache->shape()[0],
|
||||
kv_cache->shape()[1],
|
||||
kv_cache->shape()[2],
|
||||
kv_cache->shape()[3]});
|
||||
// 初始化输入：block table
|
||||
auto block_tables_xft = baidu::xpu::xft::xftTensor<int, 2>(
|
||||
reinterpret_cast<int*>(const_cast<int*>(block_tables.data<int>())),
|
||||
std::array<int64_t, 2>{block_tables.shape()[0],
|
||||
block_tables.shape()[1]});
|
||||
// 初始化concat
|
||||
std::vector<const DataType_*> concat_input;
|
||||
concat_input.push_back(kv_nope_xft.data());
|
||||
concat_input.push_back(kv_pe_xft.data());
|
||||
auto concat_out = paddle::full({num_tokens, kv_lora_rank + rope_head_dim}, -2, kv_nope.type(), kv_nope.place());
|
||||
auto concat_out_xft = baidu::xpu::xft::xftTensor<DataType_, 2>(
|
||||
reinterpret_cast<DataType_*>(const_cast<data_t*>(concat_out.data<data_t>())),
|
||||
std::array<int64_t, 2>{concat_out.shape()[0],
|
||||
concat_out.shape()[1]});
|
||||
|
||||
// 拼接kv_nope和kv_pe
|
||||
baidu::xpu::api::concat<DataType_>(xpu_ctx->x_context(),
|
||||
concat_input,
|
||||
concat_out_xft.data(),
|
||||
{{num_tokens, kv_lora_rank},{num_tokens, rope_head_dim}},
|
||||
1);
|
||||
|
||||
|
||||
baidu::xpu::api::VectorParam<int32_t> context_len_vp{const_cast<int32_t*>(encoder_seq_lod_cpu.data<int32_t>()), enc_batch + 1, const_cast<int32_t*>(encoder_seq_lod.data<int32_t>())};
|
||||
baidu::xpu::api::VectorParam<int32_t> valid_batch_vp{const_cast<int32_t*>(encoder_batch_map_cpu.data<int32_t>()), enc_batch, const_cast<int32_t*>(encoder_batch_map.data<int32_t>())};
|
||||
baidu::xpu::api::VectorParam<int32_t> start_tokens_vp{const_cast<int32_t*>(start_token_raw_cpu.data<int32_t>()), enc_batch, const_cast<int32_t*>(start_token_raw.data<int32_t>())};
|
||||
int ret_cache = xftkernel::xft_reshape_cached_kv<float16, float16, int>(
|
||||
xpu_ctx->x_context(),
|
||||
reinterpret_cast<float16*>(concat_out_xft.data()),
|
||||
reinterpret_cast<float16*>(cache_xft.data()),
|
||||
block_tables_xft.data(),
|
||||
context_len_vp,
|
||||
start_tokens_vp,
|
||||
valid_batch_vp,
|
||||
enc_batch,
|
||||
1, // num_head == 1 吸收版
|
||||
kv_lora_rank + rope_head_dim,
|
||||
block_batch,
|
||||
block_size,
|
||||
max_blocks_per_seq,
|
||||
"BLHD", // qkv_layout,
|
||||
"HLD", // page_layout,
|
||||
nullptr, // scale,
|
||||
nullptr, // batch_max_ptrs,
|
||||
nullptr); // max_ptrs
|
||||
|
||||
baidu::xpu::api::plugin::print_times("[TIME END] PrefillMLAWriteCache");
|
||||
|
||||
return {};
|
||||
}
|
||||
|
||||
std::vector<paddle::Tensor> PrefillMLAWriteCacheKernel(
|
||||
const paddle::Tensor& kv_nope,
|
||||
const paddle::Tensor& kv_pe,
|
||||
const paddle::Tensor& kv_cache,
|
||||
const paddle::Tensor& encoder_seq_lod,
|
||||
const paddle::Tensor& encoder_batch_map,
|
||||
const paddle::Tensor& start_token_raw, // encoder cache写时 start token为0 （不算prefix cache）
|
||||
const paddle::Tensor& encoder_seq_lod_cpu,
|
||||
const paddle::Tensor& encoder_batch_map_cpu,
|
||||
const paddle::Tensor& start_token_raw_cpu,
|
||||
const paddle::Tensor& enc_batch_tensor,
|
||||
const paddle::Tensor& dec_batch_tensor,
|
||||
const paddle::Tensor& padding_offsets,
|
||||
const paddle::Tensor& cum_offsets,
|
||||
const paddle::Tensor& block_tables,
|
||||
const std::string& cache_quant_type_str,
|
||||
const int max_seq_len,
|
||||
const int kv_num_heads) {
|
||||
|
||||
switch (kv_pe.dtype()) {
|
||||
case paddle::DataType::BFLOAT16: {
|
||||
return PrefillMLAWriteCache<paddle::DataType::BFLOAT16>(
|
||||
kv_nope,
|
||||
kv_pe,
|
||||
encoder_seq_lod,
|
||||
encoder_batch_map,
|
||||
start_token_raw,
|
||||
encoder_seq_lod_cpu,
|
||||
encoder_batch_map_cpu,
|
||||
start_token_raw_cpu,
|
||||
enc_batch_tensor,
|
||||
dec_batch_tensor,
|
||||
padding_offsets,
|
||||
cum_offsets,
|
||||
block_tables,
|
||||
max_seq_len,
|
||||
kv_num_heads,
|
||||
const_cast<paddle::Tensor*>(&kv_cache));
|
||||
}
|
||||
case paddle::DataType::FLOAT16: {
|
||||
return PrefillMLAWriteCache<paddle::DataType::FLOAT16>(
|
||||
kv_nope,
|
||||
kv_pe,
|
||||
encoder_seq_lod,
|
||||
encoder_batch_map,
|
||||
start_token_raw,
|
||||
encoder_seq_lod_cpu,
|
||||
encoder_batch_map_cpu,
|
||||
start_token_raw_cpu,
|
||||
enc_batch_tensor,
|
||||
dec_batch_tensor,
|
||||
padding_offsets,
|
||||
cum_offsets,
|
||||
block_tables,
|
||||
max_seq_len,
|
||||
kv_num_heads,
|
||||
const_cast<paddle::Tensor*>(&kv_cache));
|
||||
}
|
||||
}
|
||||
return {};
|
||||
}
|
||||
|
||||
template <paddle::DataType T>
|
||||
std::vector<paddle::Tensor> DecodeMLAWriteCache(
|
||||
const paddle::Tensor& kv_nope,
|
||||
const paddle::Tensor& kv_pe,
|
||||
const paddle::Tensor& decoder_context_len_cache,
|
||||
const paddle::Tensor& decoder_batch_map,
|
||||
const paddle::Tensor& kv_seq_lod_raw, // decoder cache写时，lod为0，1，2，3，……
|
||||
const paddle::Tensor& decoder_context_len_cache_cpu,
|
||||
const paddle::Tensor& decoder_batch_map_cpu,
|
||||
const paddle::Tensor& kv_seq_lod_raw_cpu, // decoder cache写时，lod为0，1，2，3，……
|
||||
const paddle::Tensor& enc_batch_tensor,
|
||||
const paddle::Tensor& dec_batch_tensor,
|
||||
const paddle::Tensor& padding_offsets,
|
||||
const paddle::Tensor& cum_offsets,
|
||||
const paddle::Tensor& block_tables,
|
||||
const int max_seq_len,
|
||||
const int kv_num_heads,
|
||||
const bool speculate_decoder,
|
||||
paddle::Tensor* kv_cache) {
|
||||
|
||||
baidu::xpu::api::plugin::print_times("[TIME BEGIN] DecodeMLAWriteCache");
|
||||
|
||||
phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
|
||||
auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
|
||||
auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
|
||||
xpu::ctx_guard RAII_GUARD(xpu_ctx->x_context());
|
||||
typedef PaddleTypeToXPUType<T> traits_;
|
||||
typedef typename traits_::DataType DataType_;
|
||||
typedef typename traits_::data_t data_t;
|
||||
|
||||
const auto& kv_nope_dims = kv_nope.dims();
|
||||
const auto& kv_pe_dims = kv_pe.dims();
|
||||
const auto& kv_cache_dims = kv_cache->dims();
|
||||
auto num_tokens = kv_nope_dims[0];
|
||||
auto block_size = kv_cache_dims[2];
|
||||
auto block_batch = block_tables.dims()[0];
|
||||
auto max_blocks_per_seq = block_tables.dims()[1];
|
||||
auto kv_lora_rank = kv_nope_dims[kv_nope_dims.size() - 1];
|
||||
auto rope_head_dim = kv_pe_dims[kv_pe_dims.size() - 1];
|
||||
auto bsz = cum_offsets.dims()[0];
|
||||
int enc_batch = enc_batch_tensor.data<int32_t>()[0];
|
||||
int dec_batch = dec_batch_tensor.data<int32_t>()[0];
|
||||
|
||||
// 初始化输入
|
||||
auto kv_nope_xft = baidu::xpu::xft::xftTensor<DataType_, 2>(
|
||||
reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_nope.data<data_t>())),
|
||||
std::array<int64_t, 2>{num_tokens,
|
||||
kv_lora_rank});
|
||||
auto kv_pe_xft = baidu::xpu::xft::xftTensor<DataType_, 2>(
|
||||
reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_pe.data<data_t>())),
|
||||
std::array<int64_t, 2>{num_tokens,
|
||||
rope_head_dim});
|
||||
// 初始化输入 cache
|
||||
auto cache_xft = baidu::xpu::xft::xftTensor<DataType_, 4>(
|
||||
reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_cache->data<data_t>())),
|
||||
std::array<int64_t, 4>{kv_cache->shape()[0],
|
||||
kv_cache->shape()[1],
|
||||
kv_cache->shape()[2],
|
||||
kv_cache->shape()[3]});
|
||||
// 初始化输入：block table
|
||||
auto block_tables_xft = baidu::xpu::xft::xftTensor<int, 2>(
|
||||
reinterpret_cast<int*>(const_cast<int*>(block_tables.data<int>())),
|
||||
std::array<int64_t, 2>{block_tables.shape()[0],
|
||||
block_tables.shape()[1]});
|
||||
// 初始化concat
|
||||
// std::vector<DataType_*> concat_input = {kv_nope_xft.data(), kv_pe_xft.data()};
|
||||
std::vector<const DataType_*> concat_input;
|
||||
concat_input.push_back(kv_nope_xft.data());
|
||||
concat_input.push_back(kv_pe_xft.data());
|
||||
auto concat_out = paddle::full({num_tokens, kv_lora_rank + rope_head_dim}, -2, kv_nope.type(), kv_nope.place());
|
||||
auto concat_out_xft = baidu::xpu::xft::xftTensor<DataType_, 2>(
|
||||
reinterpret_cast<DataType_*>(const_cast<data_t*>(concat_out.data<data_t>())),
|
||||
std::array<int64_t, 2>{concat_out.shape()[0],
|
||||
concat_out.shape()[1]});
|
||||
|
||||
// 拼接kv_nope和kv_pe
|
||||
baidu::xpu::api::concat<DataType_>(xpu_ctx->x_context(),
|
||||
concat_input,
|
||||
concat_out_xft.data(),
|
||||
{{num_tokens, kv_lora_rank},{num_tokens, rope_head_dim}},
|
||||
1);
|
||||
|
||||
|
||||
|
||||
|
||||
// context_len
|
||||
baidu::xpu::api::VectorParam<int32_t> context_len_vp_cache{const_cast<int32_t*>(decoder_context_len_cache_cpu.data<int32_t>()), dec_batch, const_cast<int32_t*>(decoder_context_len_cache.data<int32_t>())};
|
||||
// real batch
|
||||
baidu::xpu::api::VectorParam<int32_t> valid_batch_vp{const_cast<int32_t*>(decoder_batch_map_cpu.data<int32_t>()), dec_batch, const_cast<int32_t*>(decoder_batch_map.data<int32_t>())};
|
||||
// baidu::xpu::api::VectorParam<int32_t> kv_seq_lod_vp{const_cast<int32_t*>(kv_seq_lod_raw_cpu.data<int32_t>()), dec_batch + 1, const_cast<int32_t*>(kv_seq_lod_raw.data<int32_t>())};
|
||||
|
||||
// baidu::xpu::api::VectorParam<int32_t> kv_seq_lod_vp = baidu::xpu::api::VectorParam<int32_t>{const_cast<int32_t*>(kv_seq_lod_raw_cpu.data<int32_t>()), dec_batch + 1, const_cast<int32_t*>(kv_seq_lod_raw.data<int32_t>())};
|
||||
|
||||
std::vector<int> kv_seq_lod(dec_batch + 1);
|
||||
std::iota(kv_seq_lod.begin(), kv_seq_lod.end(), 0);
|
||||
baidu::xpu::api::VectorParam<int32_t> kv_seq_lod_vp =
|
||||
baidu::xpu::api::VectorParam<int32_t>{kv_seq_lod.data(), dec_batch + 1, nullptr}.to_xpu(RAII_GUARD);
|
||||
|
||||
|
||||
// std::cout << "Tensor dtype: " << kv_seq_lod_raw_cpu.dtype() << std::endl;
|
||||
// std::cout << "Tensor place: " << kv_seq_lod_raw_cpu.place() << std::endl;
|
||||
|
||||
// // 打印形状信息
|
||||
// std::cout << "Tensor shape: [";
|
||||
// for (int i = 0; i < kv_seq_lod_raw_cpu.dims().size(); ++i) {
|
||||
// if (i > 0) std::cout << ", ";
|
||||
// std::cout << kv_seq_lod_raw_cpu.dims()[i];
|
||||
// }
|
||||
// std::cout << "]" << std::endl;
|
||||
|
||||
// std::cout << "Tensor dtype: " << kv_seq_lod_raw.dtype() << std::endl;
|
||||
// std::cout << "Tensor place: " << kv_seq_lod_raw.place() << std::endl;
|
||||
|
||||
// // 打印形状信息
|
||||
// std::cout << "Tensor shape: [";
|
||||
// for (int i = 0; i < kv_seq_lod_raw.dims().size(); ++i) {
|
||||
// if (i > 0) std::cout << ", ";
|
||||
// std::cout << kv_seq_lod_raw.dims()[i];
|
||||
// }
|
||||
// std::cout << "]" << std::endl;
|
||||
|
||||
|
||||
int ret_cache = xftkernel::xft_reshape_cached_kv<float16, float16, int>(
|
||||
xpu_ctx->x_context(),
|
||||
reinterpret_cast<float16*>(concat_out_xft.data()),
|
||||
reinterpret_cast<float16*>(cache_xft.data()),
|
||||
block_tables_xft.data(),
|
||||
kv_seq_lod_vp,
|
||||
context_len_vp_cache,
|
||||
valid_batch_vp,
|
||||
dec_batch,
|
||||
1, // num_head == 1
|
||||
kv_lora_rank + rope_head_dim,
|
||||
block_batch,
|
||||
block_size,
|
||||
max_blocks_per_seq,
|
||||
"BLHD", // qkv_layout,
|
||||
"HLD", // page_layout,
|
||||
nullptr, // scale,
|
||||
nullptr, // batch_max_ptrs,
|
||||
nullptr); // max_ptrs
|
||||
|
||||
|
||||
baidu::xpu::api::plugin::print_times("[TIME END] DecodeMLAWriteCache");
|
||||
|
||||
return {};
|
||||
}
|
||||
|
||||
std::vector<paddle::Tensor> DecodeMLAWriteCacheKernel(
|
||||
const paddle::Tensor& kv_nope,
|
||||
const paddle::Tensor& kv_pe,
|
||||
const paddle::Tensor& kv_cache,
|
||||
const paddle::Tensor& decoder_context_len_cache,
|
||||
const paddle::Tensor& decoder_batch_map,
|
||||
const paddle::Tensor& kv_seq_lod_raw,
|
||||
const paddle::Tensor& decoder_context_len_cache_cpu,
|
||||
const paddle::Tensor& decoder_batch_map_cpu,
|
||||
const paddle::Tensor& kv_seq_lod_raw_cpu,
|
||||
const paddle::Tensor& enc_batch_tensor,
|
||||
const paddle::Tensor& dec_batch_tensor,
|
||||
const paddle::Tensor& padding_offsets,
|
||||
const paddle::Tensor& cum_offsets,
|
||||
const paddle::Tensor& block_tables,
|
||||
const std::string& cache_quant_type_str,
|
||||
const int max_seq_len,
|
||||
const int kv_num_heads,
|
||||
const bool speculate_decoder) {
|
||||
switch (kv_pe.dtype()) {
|
||||
case paddle::DataType::BFLOAT16: {
|
||||
return DecodeMLAWriteCache<paddle::DataType::BFLOAT16>(
|
||||
kv_nope,
|
||||
kv_pe,
|
||||
decoder_context_len_cache,
|
||||
decoder_batch_map,
|
||||
kv_seq_lod_raw,
|
||||
decoder_context_len_cache_cpu,
|
||||
decoder_batch_map_cpu,
|
||||
kv_seq_lod_raw_cpu,
|
||||
enc_batch_tensor,
|
||||
dec_batch_tensor,
|
||||
padding_offsets,
|
||||
cum_offsets,
|
||||
block_tables,
|
||||
max_seq_len,
|
||||
kv_num_heads,
|
||||
speculate_decoder,
|
||||
const_cast<paddle::Tensor*>(&kv_cache));
|
||||
}
|
||||
case paddle::DataType::FLOAT16: {
|
||||
return DecodeMLAWriteCache<paddle::DataType::FLOAT16>(
|
||||
kv_nope,
|
||||
kv_pe,
|
||||
decoder_context_len_cache,
|
||||
decoder_batch_map,
|
||||
kv_seq_lod_raw,
|
||||
decoder_context_len_cache_cpu,
|
||||
decoder_batch_map_cpu,
|
||||
kv_seq_lod_raw_cpu,
|
||||
enc_batch_tensor,
|
||||
dec_batch_tensor,
|
||||
padding_offsets,
|
||||
cum_offsets,
|
||||
block_tables,
|
||||
max_seq_len,
|
||||
kv_num_heads,
|
||||
speculate_decoder,
|
||||
const_cast<paddle::Tensor*>(&kv_cache));
|
||||
}
|
||||
}
|
||||
return {};
|
||||
}
|
||||
|
||||
|
||||
PD_BUILD_OP(prefill_mla_write_cache_xpu)
|
||||
.Inputs({"kv_nope",
|
||||
"kv_pe",
|
||||
"kv_cache",
|
||||
"encoder_seq_lod",
|
||||
"encoder_batch_map",
|
||||
"start_token_raw",
|
||||
"encoder_seq_lod_cpu",
|
||||
"encoder_batch_map_cpu",
|
||||
"start_token_raw_cpu",
|
||||
"enc_batch_tensor",
|
||||
"dec_batch_tensor",
|
||||
"padding_offsets",
|
||||
"cum_offsets",
|
||||
"block_tables"})
|
||||
.Outputs({"kv_cache_out"})
|
||||
.SetInplaceMap({{"kv_cache", "kv_cache_out"}})
|
||||
.Attrs({"cache_quant_type_str: std::string",
|
||||
"max_seq_len: int",
|
||||
"kv_num_heads: int"})
|
||||
.SetKernelFn(PD_KERNEL(PrefillMLAWriteCacheKernel));
|
||||
|
||||
PD_BUILD_OP(decode_mla_write_cache_xpu)
|
||||
.Inputs({"kv_nope",
|
||||
"kv_pe",
|
||||
"kv_cache",
|
||||
"decoder_context_len_cache",
|
||||
"decoder_batch_map",
|
||||
"kv_seq_lod_raw",
|
||||
"decoder_context_len_cache_cpu",
|
||||
"decoder_batch_map_cpu",
|
||||
"kv_seq_lod_raw_cpu",
|
||||
"enc_batch_tensor",
|
||||
"dec_batch_tensor",
|
||||
"padding_offsets",
|
||||
"cum_offsets",
|
||||
"block_tables"})
|
||||
.Outputs({"kv_cache_out"})
|
||||
.SetInplaceMap({{"kv_cache", "kv_cache_out"}})
|
||||
.Attrs({"cache_quant_type_str: std::string",
|
||||
"max_seq_len: int",
|
||||
"kv_num_heads: int",
|
||||
"speculate_decoder: bool"})
|
||||
.SetKernelFn(PD_KERNEL(DecodeMLAWriteCacheKernel));
|
||||
173
csrc/xpu/src/moe_fused.cc
Normal file
173
csrc/xpu/src/moe_fused.cc
Normal file
@@ -0,0 +1,173 @@
|
||||
// Copyright (c) 2024 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 <core/ctx_manager.h>
|
||||
#include <core/xft_tensor.h>
|
||||
#include <core/xft_check.h>
|
||||
#include <core/xft_event.h>
|
||||
#include <core/xft_params.h>
|
||||
#include <paddle/phi/backends/xpu/xpu_context.h>
|
||||
#include <xft/operation/ffn_moe_vllm.h>
|
||||
#include <xft/xdnn_plugin.h>
|
||||
|
||||
#include "paddle/extension.h"
|
||||
#include "paddle/phi/core/enforce.h"
|
||||
#include "xpu/plugin.h"
|
||||
|
||||
namespace xftkernel = baidu::xpu::xftkernel;
|
||||
namespace xft = baidu::xpu::xft;
|
||||
|
||||
inline uint32_t get_xpu_moe_speed_layers() {
|
||||
static const char* xpu_moe_speed_layers_env = std::getenv("FLAGS_xpu_moe_speed_layers");
|
||||
static const uint32_t xpu_moe_speed_layers =
|
||||
xpu_moe_speed_layers_env == nullptr ? 0 : std::stoul(std::string(xpu_moe_speed_layers_env));
|
||||
return xpu_moe_speed_layers;
|
||||
}
|
||||
|
||||
std::vector<paddle::Tensor> MoeFusedKernel(
|
||||
const paddle::Tensor& input, //[980. 7168]
|
||||
const paddle::Tensor& gate_weight,//[7168. 256] ->[256,7168]
|
||||
const paddle::Tensor& ffn_inter_weights,//[256, 7168,4096] reshape
|
||||
const paddle::Tensor& ffn_output_weights,//[256, 2048, 7168] reshape
|
||||
const paddle::Tensor& ffn1_weights_scale,// [256, 4096]
|
||||
const paddle::Tensor& ffn2_weights_scale,// [256, 7168]
|
||||
const paddle::Tensor& score_bias, // 256
|
||||
int moe_top_k,
|
||||
int expert_group_num,
|
||||
int moe_topk_group,
|
||||
int layer) {
|
||||
|
||||
baidu::xpu::api::plugin::print_times("[TIME BEGIN] MoeFusedKernel");
|
||||
|
||||
phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
|
||||
auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
|
||||
auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
|
||||
|
||||
using XPUType = typename XPUTypeTrait<bfloat16>::Type;
|
||||
typedef paddle::bfloat16 data_t;
|
||||
|
||||
const int64_t m = input.shape()[0];
|
||||
const int64_t hidden_size = input.shape()[1];
|
||||
const int64_t expert_num = gate_weight.shape()[0];
|
||||
const int64_t ffn_inter_hidden_size = ffn_inter_weights.shape()[2];
|
||||
const int64_t ffn_output_hidden_size = ffn_output_weights.shape()[2];
|
||||
|
||||
xft::Tensor input_tensor(
|
||||
reinterpret_cast<XPUType*>(const_cast<data_t*>(input.data<data_t>())),
|
||||
xft::DataType::DT_BFLOAT16,
|
||||
input.shape());
|
||||
xft::Tensor gate_weight_tensor(
|
||||
const_cast<float*>(gate_weight.data<float>()),
|
||||
xft::DataType::DT_FLOAT,
|
||||
gate_weight.shape());
|
||||
xft::Tensor ffn_inter_weights_tensor(
|
||||
const_cast<int8_t*>(ffn_inter_weights.data<int8_t>()),
|
||||
nullptr,
|
||||
const_cast<float*>(ffn1_weights_scale.data<float>()),
|
||||
xft::DataType::DT_INT8,
|
||||
{expert_num, ffn_inter_hidden_size, hidden_size});
|
||||
xft::Tensor ffn_output_weights_tensor(
|
||||
const_cast<int8_t*>(ffn_output_weights.data<int8_t>()),
|
||||
nullptr,
|
||||
const_cast<float*>(ffn2_weights_scale.data<float>()),
|
||||
xft::DataType::DT_INT8,
|
||||
{expert_num, ffn_output_hidden_size, ffn_inter_hidden_size / 2});
|
||||
xft::Tensor score_bias_tensor(const_cast<float*>(score_bias.data<float>()),
|
||||
xft::DataType::DT_FLOAT,
|
||||
score_bias.shape());
|
||||
|
||||
auto output = paddle::full({m, hidden_size}, -1, input.type(), input.place());
|
||||
xft::Tensor output_tensor(reinterpret_cast<XPUType*>(output.data<data_t>()),
|
||||
xft::DataType::DT_BFLOAT16,
|
||||
output.shape());
|
||||
|
||||
xft::MoeFFNWeight moe_weight{&gate_weight_tensor,
|
||||
&ffn_inter_weights_tensor,
|
||||
nullptr,
|
||||
&ffn_output_weights_tensor,
|
||||
nullptr,
|
||||
nullptr,
|
||||
nullptr,
|
||||
&score_bias_tensor};
|
||||
xft::MoeFFNParam moe_param{
|
||||
expert_num,
|
||||
moe_top_k,
|
||||
false,
|
||||
true,
|
||||
expert_group_num,
|
||||
moe_topk_group,
|
||||
true,
|
||||
false,
|
||||
1,
|
||||
0,
|
||||
"sigmoid"
|
||||
};
|
||||
|
||||
if (layer >= get_xpu_moe_speed_layers()) {
|
||||
xft::xft_moe_ffn_block_sorted<XPUType, int8_t, XPUType, float>(
|
||||
xpu_ctx->x_context(), &input_tensor, &output_tensor, moe_weight, moe_param);
|
||||
} else {
|
||||
xft::xft_moe_ffn_block_sorted<XPUType, int8_t, XPUType, int8_wo_t>(
|
||||
xpu_ctx->x_context(), &input_tensor, &output_tensor, moe_weight, moe_param);
|
||||
}
|
||||
|
||||
baidu::xpu::api::plugin::print_times("[TIME END] MoeFusedKernel");
|
||||
|
||||
return {
|
||||
output,
|
||||
};
|
||||
}
|
||||
|
||||
|
||||
std::vector<std::vector<int64_t>> MoeFusedInferShape(
|
||||
const std::vector<int64_t>& input_shape,
|
||||
const std::vector<int64_t>& gate_weight_shape,
|
||||
const std::vector<int64_t>& ffn_inter_weights_shape,
|
||||
const std::vector<int64_t>& ffn_output_weights_shape,
|
||||
const std::vector<int64_t>& ffn1_weights_scale_shape,
|
||||
const std::vector<int64_t>& ffn2_weights_scale_shape,
|
||||
const std::vector<int64_t>& score_bias_shape) {
|
||||
return {input_shape};
|
||||
}
|
||||
|
||||
std::vector<paddle::DataType> MoeFusedInferDtype(
|
||||
const paddle::DataType& input_type,
|
||||
const paddle::DataType& gate_weight_type,
|
||||
const paddle::DataType& ffn_inter_weights_type,
|
||||
const paddle::DataType& ffn_output_weights_type,
|
||||
const paddle::DataType& ffn1_weights_scale_type,
|
||||
const paddle::DataType& ffn2_weights_scale_type,
|
||||
const paddle::DataType& score_bias_type,
|
||||
const int moe_top_k,
|
||||
const int expert_group_num,
|
||||
const int moe_topk_group,
|
||||
const int layer) {
|
||||
return {input_type};
|
||||
}
|
||||
|
||||
PD_BUILD_OP(moe_fused_xpu)
|
||||
.Inputs({
|
||||
"input",
|
||||
"gate_weight",
|
||||
"ffn_inter_weights",
|
||||
"ffn_outer_weights",
|
||||
"ffn1_weights_scale",
|
||||
"ffn2_weights_scale",
|
||||
"score_bias",
|
||||
})
|
||||
.Outputs({"output"})
|
||||
.Attrs({"moe_top_k: int", "expert_group_num: int", "moe_topk_group: int", "layer: int"})
|
||||
.SetKernelFn(PD_KERNEL(MoeFusedKernel))
|
||||
.SetInferShapeFn(PD_INFER_SHAPE(MoeFusedInferShape))
|
||||
.SetInferDtypeFn(PD_INFER_DTYPE(MoeFusedInferDtype));
|
||||
113
csrc/xpu/src/ops.h
Normal file
113
csrc/xpu/src/ops.h
Normal file
@@ -0,0 +1,113 @@
|
||||
// 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/backends/xpu/xpu_context.h>
|
||||
#include "paddle/extension.h"
|
||||
#include "xpu/plugin.h"
|
||||
#include <xft/xdnn_plugin.h>
|
||||
|
||||
void FusedRotaryPositionEncoding(
|
||||
paddle::Tensor& query, // [num_tokens, num_heads, head_size] or
|
||||
// [num_tokens, num_heads * head_size]
|
||||
paddle::Tensor& key,
|
||||
// [num_tokens, num_kv_heads, head_size] or [num_tokens, num_kv_heads *
|
||||
// head_size]
|
||||
const paddle::Tensor& position_ids, // [num_tokens]
|
||||
const paddle::Tensor& cos_sin_cache, // [max_position, rot_dim]
|
||||
int head_size,
|
||||
bool is_neox);
|
||||
|
||||
|
||||
std::vector<paddle::Tensor> MlaDeAttn(
|
||||
const paddle::Tensor& q,
|
||||
const paddle::Tensor& kv_cache,
|
||||
const paddle::Tensor& decoder_context_len,
|
||||
const paddle::Tensor& decoder_batch_map,
|
||||
const paddle::Tensor& decoder_context_len_cpu,
|
||||
const paddle::Tensor& decoder_batch_map_cpu,
|
||||
const paddle::Tensor& dec_batch_tensor,
|
||||
const paddle::Tensor& padding_offsets,
|
||||
const paddle::Tensor& cum_offsets,
|
||||
const paddle::Tensor& block_tables,
|
||||
const float softmax_scale,
|
||||
const int block_size,
|
||||
const int num_head,
|
||||
const int kv_lora_rank,
|
||||
const int rope_head_dim,
|
||||
const int dim_qk,
|
||||
const int dim_v);
|
||||
|
||||
|
||||
std::vector<paddle::Tensor> MlaEnAttn(
|
||||
const paddle::Tensor& q,
|
||||
const paddle::Tensor& k,
|
||||
const paddle::Tensor& v,
|
||||
const paddle::Tensor& encoder_seq_lod,
|
||||
const paddle::Tensor& encoder_batch_map,
|
||||
const paddle::Tensor& encoder_seq_lod_cpu,
|
||||
const paddle::Tensor& encoder_batch_map_cpu,
|
||||
const paddle::Tensor& enc_batch_tensor,
|
||||
const paddle::Tensor& padding_offsets,
|
||||
const paddle::Tensor& cum_offsets,
|
||||
const paddle::Tensor& block_tables,
|
||||
const paddle::Tensor& max_enc_len_this_time,
|
||||
const paddle::Tensor& max_dec_len_this_time,
|
||||
const float softmax_scale,
|
||||
const int block_size,
|
||||
const int num_head,
|
||||
const int dim_qk,
|
||||
const int dim_v);
|
||||
|
||||
|
||||
std::vector<paddle::Tensor> DecodeMLAWriteCacheKernel(
|
||||
const paddle::Tensor& kv_nope,
|
||||
const paddle::Tensor& kv_pe,
|
||||
const paddle::Tensor& kv_cache,
|
||||
const paddle::Tensor& decoder_context_len_cache,
|
||||
const paddle::Tensor& decoder_batch_map,
|
||||
const paddle::Tensor& kv_seq_lod_raw,
|
||||
const paddle::Tensor& decoder_context_len_cache_cpu,
|
||||
const paddle::Tensor& decoder_batch_map_cpu,
|
||||
const paddle::Tensor& kv_seq_lod_raw_cpu,
|
||||
const paddle::Tensor& enc_batch_tensor,
|
||||
const paddle::Tensor& dec_batch_tensor,
|
||||
const paddle::Tensor& padding_offsets,
|
||||
const paddle::Tensor& cum_offsets,
|
||||
const paddle::Tensor& block_tables,
|
||||
const std::string& cache_quant_type_str,
|
||||
const int max_seq_len,
|
||||
const int kv_num_heads,
|
||||
const bool speculate_decoder);
|
||||
|
||||
std::vector<paddle::Tensor> PrefillMLAWriteCacheKernel(
|
||||
const paddle::Tensor& kv_nope,
|
||||
const paddle::Tensor& kv_pe,
|
||||
const paddle::Tensor& kv_cache,
|
||||
const paddle::Tensor& encoder_seq_lod,
|
||||
const paddle::Tensor& encoder_batch_map,
|
||||
const paddle::Tensor& start_token_raw, // encoder cache写时 start token为0 （不算prefix cache）
|
||||
const paddle::Tensor& encoder_seq_lod_cpu,
|
||||
const paddle::Tensor& encoder_batch_map_cpu,
|
||||
const paddle::Tensor& start_token_raw_cpu,
|
||||
const paddle::Tensor& enc_batch_tensor,
|
||||
const paddle::Tensor& dec_batch_tensor,
|
||||
const paddle::Tensor& padding_offsets,
|
||||
const paddle::Tensor& cum_offsets,
|
||||
const paddle::Tensor& block_tables,
|
||||
const std::string& cache_quant_type_str,
|
||||
const int max_seq_len,
|
||||
const int kv_num_heads);
|
||||
|
||||
std::vector<paddle::Tensor> Bmm(const paddle::Tensor& input,
|
||||
const paddle::Tensor& weight);
|
||||
@@ -17,11 +17,20 @@
|
||||
|
||||
#pragma once
|
||||
#include "xpu/xdnn.h"
|
||||
|
||||
#include <chrono>
|
||||
namespace baidu {
|
||||
namespace xpu {
|
||||
namespace api {
|
||||
namespace plugin {
|
||||
static std::string print_times(std::string str){
|
||||
#ifdef XPU_PRINT_OP_TIME
|
||||
auto now = std::chrono::system_clock::now();
|
||||
auto duration = now.time_since_epoch();
|
||||
auto milliseconds = std::chrono::duration_cast<std::chrono::milliseconds>(duration).count();
|
||||
std::cout<<str<<":" <<milliseconds<<std::endl;
|
||||
#endif
|
||||
return "";
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
DLL_EXPORT int set_stop_value_multi_ends(Context* ctx,
|
||||
@@ -84,20 +93,6 @@ DLL_EXPORT int get_position_ids(Context *ctx,
|
||||
int *position_ids,
|
||||
const int bs);
|
||||
|
||||
DLL_EXPORT int update_inputs(Context* ctx,
|
||||
bool* not_need_stop,
|
||||
int* seq_lens_this_time,
|
||||
int* seq_lens_encoder,
|
||||
int* seq_lens_decoder,
|
||||
int64_t* input_ids,
|
||||
const int64_t* stop_nums,
|
||||
const bool* stop_flags,
|
||||
const bool* is_block_step,
|
||||
const int64_t* next_tokens,
|
||||
const int bsz,
|
||||
const int max_bsz,
|
||||
const int input_ids_stride);
|
||||
|
||||
template <typename T>
|
||||
DLL_EXPORT int rebuild_padding(Context *ctx,
|
||||
T *output_data, // [bs, dim_embed]
|
||||
@@ -172,6 +167,19 @@ DLL_EXPORT int free_and_dispatch_block(Context *ctx,
|
||||
const int block_size,
|
||||
const int block_num_per_seq,
|
||||
const int max_decoder_block_num);
|
||||
|
||||
template <typename T, typename TR>
|
||||
DLL_EXPORT int rotary_embedding_neox(
|
||||
Context* ctx,
|
||||
const int* positions, // [num_tokens]
|
||||
T* query, // [num_tokens, num_heads, head_size]
|
||||
T* key, // [num_tokens, num_kv_heads, head_size]
|
||||
const TR* cos_sin_cache, // [max_position, 2, rot_dim // 2]
|
||||
const int rot_dim,
|
||||
const int num_heads,
|
||||
const int num_kv_heads,
|
||||
const int head_size,
|
||||
const int32_t num_tokens);
|
||||
} // namespace plugin
|
||||
} // namespace api
|
||||
} // namespace xpu
|
||||
|
||||
@@ -0,0 +1,96 @@
|
||||
#include "xpu/kernel/cluster.h"
|
||||
#include "xpu/kernel/cluster_partition.h"
|
||||
#include "xpu/kernel/cluster_primitive.h"
|
||||
|
||||
namespace xpu3 {
|
||||
namespace plugin {
|
||||
|
||||
static inline __device__ void batch_load(
|
||||
_global_ptr_ const int32_t* kv_seq_lod,
|
||||
_global_ptr_ const int32_t* start_token,
|
||||
_global_ptr_ const int32_t* real_batch,
|
||||
int* kv_seq_lod_lm,
|
||||
int* start_token_lm,
|
||||
int* real_batch_lm,
|
||||
int batch_begin,
|
||||
int read_batch_len) {
|
||||
GM2LM_ASYNC(kv_seq_lod + batch_begin, kv_seq_lod_lm, (read_batch_len + 1) * sizeof(int));
|
||||
GM2LM_ASYNC(start_token + batch_begin, start_token_lm, read_batch_len * sizeof(int));
|
||||
GM2LM_ASYNC(real_batch + batch_begin, real_batch_lm, read_batch_len * sizeof(int));
|
||||
mfence();
|
||||
}
|
||||
|
||||
template <typename TI, typename TO, typename TID>
|
||||
__global__ void absorb_mla_cache_write(
|
||||
const TI* src,
|
||||
TO* dst,
|
||||
const TID* block_table,
|
||||
int32_t* kv_seq_lod,
|
||||
int32_t* start_token,
|
||||
int32_t* real_batch,
|
||||
int64_t batch_size,
|
||||
int64_t kv_head_num,
|
||||
int64_t head_dim,
|
||||
int64_t max_seq_num,
|
||||
int64_t block_size,
|
||||
int64_t max_num_blocks_per_seq) {
|
||||
constexpr int loop_batch_size = 32;
|
||||
int cid = core_id();
|
||||
int ncores = core_num();
|
||||
int tid = cid * cluster_num() + cluster_id();
|
||||
int nthreads = ncores * cluster_num();
|
||||
// prepare lm buffer
|
||||
__local__ TI local_scr[2048];
|
||||
__local__ int lod_lm[loop_batch_size];
|
||||
__local__ int kv_seq_lod_lm[loop_batch_size];
|
||||
__local__ int start_token_lm[loop_batch_size];
|
||||
__local__ int real_batch_lm[loop_batch_size];
|
||||
int table_cnt = max_seq_num * max_num_blocks_per_seq;
|
||||
batch_load(
|
||||
kv_seq_lod,
|
||||
start_token,
|
||||
real_batch,
|
||||
kv_seq_lod_lm,
|
||||
start_token_lm,
|
||||
real_batch_lm,
|
||||
0,
|
||||
batch_size);
|
||||
int token_num = kv_seq_lod[batch_size];
|
||||
for(int i = tid; i < token_num; i += nthreads){
|
||||
int batch = 0;
|
||||
for(int b = 0; b < batch_size; b++){
|
||||
if(i < lod_lm[b + 1]){
|
||||
batch = b;
|
||||
break;
|
||||
}
|
||||
GM2LM(src + i * kv_head_num * head_dim, local_scr, kv_head_num * head_dim * sizeof(TI));
|
||||
TID block_id = block_table[real_batch_lm[batch] * max_num_blocks_per_seq + (start_token_lm[batch] + i - lod_lm[b]) / block_size];
|
||||
LM2GM(local_scr, dst + block_id * kv_head_num * block_size * head_dim + (start_token_lm[batch] + i - lod_lm[b]) % block_size * head_dim, kv_head_num * head_dim * sizeof(TO));
|
||||
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
#define _XPU_DEF__ABSORB_MLA_CACHE_WRITE(TI, TO, TID) \
|
||||
template __global__ void absorb_mla_cache_write<TI, TO, TID>( \
|
||||
const TI*, \
|
||||
TO*, \
|
||||
const TID*, \
|
||||
int32_t*, \
|
||||
int32_t*, \
|
||||
int32_t*, \
|
||||
int64_t, \
|
||||
int64_t, \
|
||||
int64_t, \
|
||||
int64_t, \
|
||||
int64_t, \
|
||||
int64_t);
|
||||
|
||||
_XPU_DEF__ABSORB_MLA_CACHE_WRITE(float16, float16, int32_t);
|
||||
_XPU_DEF__ABSORB_MLA_CACHE_WRITE(bfloat16, bfloat16, int32_t);
|
||||
|
||||
} // namespace plugin
|
||||
} // namespace xpu3
|
||||
@@ -0,0 +1,238 @@
|
||||
// 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 "xpu/kernel/cluster.h"
|
||||
#include "xpu/kernel/cluster_partition.h"
|
||||
#include "xpu/kernel/cluster_primitive.h"
|
||||
#include "xpu/kernel/cluster_debug.h"
|
||||
|
||||
namespace xpu3 {
|
||||
namespace plugin {
|
||||
|
||||
static __device__ void apply_rotary_pos_emb_nonexo(
|
||||
float* qk,
|
||||
float* cos_sin,
|
||||
int rot_dim,
|
||||
int head_size,
|
||||
int element_num) {
|
||||
const int emb_dim = rot_dim / 2;
|
||||
for(int i = 0; i < emb_dim; i++){
|
||||
float cos_ind = cos_sin[i];
|
||||
float sin_ind = cos_sin[i + emb_dim];
|
||||
for (int j = 0; j < element_num; j++) {
|
||||
// if(i == 0){
|
||||
// printf("input %f %f %f %f ",qk[j * head_size],qk[j * head_size + 1],cos_ind , sin_ind);
|
||||
// }
|
||||
float x_ind = qk[j * head_size + i * 2];
|
||||
float y_ind = qk[j * head_size + i * 2 + 1];
|
||||
qk[j * head_size + i * 2] = x_ind * cos_ind - y_ind * sin_ind;
|
||||
qk[j * head_size + i * 2 + 1] = y_ind * cos_ind + x_ind * sin_ind;
|
||||
// if(i == 0){
|
||||
// printf("output %f %f \n ",qk[j * head_size],qk[j * head_size + 1]);
|
||||
// }
|
||||
}
|
||||
}
|
||||
|
||||
mfence();
|
||||
}
|
||||
|
||||
|
||||
static __device__ void apply_rotary_pos_emb_nexo(
|
||||
float* qk,
|
||||
float* cos_sin,
|
||||
int rot_dim,
|
||||
int head_size,
|
||||
int element_num) {
|
||||
float32x16_t vqx_0;
|
||||
float32x16_t vqx_1;
|
||||
float32x16_t vqy_0;
|
||||
float32x16_t vqy_1;
|
||||
float32x16_t vcos_0;
|
||||
float32x16_t vcos_1;
|
||||
float32x16_t vsin_0;
|
||||
float32x16_t vsin_1;
|
||||
float32x16_t v_x_cos;
|
||||
float32x16_t v_y_cos;
|
||||
const float var = -1.0f;
|
||||
const int emb_dim = rot_dim / 2;
|
||||
// TODO: fixme. It assumes that rot_dim is the multiplier of 32.
|
||||
for (int i = 0; i < emb_dim; i += 32) {
|
||||
vload2_lm(cos_sin + i, vcos_0, vcos_1);
|
||||
vload2_lm(cos_sin + i + emb_dim, vsin_0, vsin_1);
|
||||
for (int j = 0; j < element_num; j++) {
|
||||
int offset = j * head_size;
|
||||
vload2_lm(qk + offset + i, vqx_0, vqx_1);
|
||||
vload2_lm(qk + offset + i + emb_dim, vqy_0, vqy_1);
|
||||
// qx
|
||||
v_x_cos = vvmul_float32x16(vqx_0, vcos_0);
|
||||
vqx_0 = vvmul_float32x16(vqx_0, vsin_0);
|
||||
// qy * cons
|
||||
v_y_cos = vvmul_float32x16(vqy_0, vcos_0);
|
||||
// -qy
|
||||
vqy_0 = svmul_float32x16(var, vqy_0);
|
||||
// -qy * sin
|
||||
vqy_0 = vvmul_float32x16(vqy_0, vsin_0);
|
||||
// add
|
||||
vqy_0 = vvadd_float32x16(v_x_cos, vqy_0); // out_left
|
||||
vqx_0 = vvadd_float32x16(vqx_0, v_y_cos); // out_right
|
||||
// qy
|
||||
v_x_cos = vvmul_float32x16(vqx_1, vcos_1);
|
||||
vqx_1 = vvmul_float32x16(vqx_1, vsin_1);
|
||||
// qy * cons
|
||||
v_y_cos = vvmul_float32x16(vqy_1, vcos_1);
|
||||
// -qy
|
||||
vqy_1 = svmul_float32x16(var, vqy_1);
|
||||
// -qy * sin
|
||||
vqy_1 = vvmul_float32x16(vqy_1, vsin_1);
|
||||
// add
|
||||
vqy_1 = vvadd_float32x16(v_x_cos, vqy_1); // out_left
|
||||
vqx_1 = vvadd_float32x16(vqx_1, v_y_cos); // out_right
|
||||
vstore2_lm(qk + offset + i, vqy_0, vqy_1);
|
||||
vstore2_lm(qk + offset + i + emb_dim, vqx_0, vqx_1);
|
||||
mfence();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename TX, typename TY>
|
||||
static __device__ void fast_cast(TX* x, TY* y, int64_t n) {
|
||||
primitive_cast<TX, TY>(x, y, n);
|
||||
}
|
||||
|
||||
template <>
|
||||
__device__ void fast_cast<bfloat16, float>(bfloat16* x, float* y, int64_t n) {
|
||||
int start = (n - 1) / 32 * 32;
|
||||
x = x + start;
|
||||
y = y + start;
|
||||
float32x16_t xl;
|
||||
float32x16_t xh;
|
||||
for (int i = start; i >= 0; i -= 32) {
|
||||
vload2_lm(x, xl, xh);
|
||||
vstore2_lm(y, xl, xh);
|
||||
x -= 32;
|
||||
y -= 32;
|
||||
}
|
||||
mfence_lm();
|
||||
}
|
||||
template <>
|
||||
__device__ void fast_cast<float, bfloat16>(float* x, bfloat16* y, int64_t n) {
|
||||
float32x16_t xl;
|
||||
float32x16_t xh;
|
||||
for (int i = 0; i < n; i += 32) {
|
||||
vload2_lm(x, xl, xh);
|
||||
vstore2_lm(y, xl, xh);
|
||||
x += 32;
|
||||
y += 32;
|
||||
}
|
||||
mfence_lm();
|
||||
}
|
||||
|
||||
template <typename T, typename TR>
|
||||
__global__ void rotary_embedding_neox(
|
||||
const int* positions, // [num_tokens]
|
||||
T* query, // [num_tokens, num_heads, head_size]
|
||||
T* key, // [num_tokens, num_kv_heads, head_size]
|
||||
const TR* cos_sin_cache, // [max_position, 2, rot_dim // 2]
|
||||
const int rot_dim,
|
||||
const int num_heads,
|
||||
const int num_kv_heads,
|
||||
const int head_size,
|
||||
const int32_t num_tokens) {
|
||||
int cid = core_id();
|
||||
int ncores = core_num();
|
||||
if (cid >= ncores) {
|
||||
return;
|
||||
}
|
||||
int tid = cluster_id() * ncores + cid;
|
||||
int nthreads = ncores * cluster_num();
|
||||
constexpr int sm_buffer_size = 64 * 512;
|
||||
int head_num = 2048 / sizeof(T) / head_size; // max DMA is 2048 Bytes
|
||||
int lm_buffer_size = head_num * head_size;
|
||||
__simd__ float data_lm[lm_buffer_size];
|
||||
__simd__ float cos_sin_lm[rot_dim];
|
||||
__simd__ __shared__ int pos_sm[sm_buffer_size];
|
||||
|
||||
int32_t cluster_start_idx = 0;
|
||||
int32_t cluster_end_idx = 0;
|
||||
int32_t start_idx = 0;
|
||||
int32_t end_idx = 0;
|
||||
partition(cluster_id(), cluster_num(), num_tokens, 1, &cluster_start_idx, &cluster_end_idx);
|
||||
if(cid == 0 && cluster_end_idx - cluster_start_idx > 0){
|
||||
GM2SM(positions + cluster_start_idx, pos_sm, (cluster_end_idx - cluster_start_idx) * sizeof(int));
|
||||
}
|
||||
partition(cid, ncores, cluster_end_idx - cluster_start_idx, 1, &start_idx, &end_idx);
|
||||
start_idx = start_idx + cluster_start_idx;
|
||||
end_idx = end_idx + cluster_start_idx;
|
||||
mfence();
|
||||
sync_all();
|
||||
// printf("id %d %d %d %d\n ",cluster_id(), cluster_start_idx, cluster_end_idx, pos_sm[0]);
|
||||
for (int64_t token_idx = start_idx; token_idx < end_idx; token_idx += 1) {
|
||||
int64_t pos = pos_sm[token_idx - cluster_start_idx];
|
||||
// printf("poss %lld %lld \n ",token_idx, pos);
|
||||
|
||||
_global_ptr_ const TR* cache_ptr = cos_sin_cache + pos * rot_dim;
|
||||
int64_t offset_q = token_idx * num_heads * head_size;
|
||||
int64_t offset_k = token_idx * num_kv_heads * head_size;
|
||||
// get cos & sin
|
||||
if (num_tokens > token_idx) {
|
||||
GM2LM(cache_ptr, cos_sin_lm, rot_dim * sizeof(TR));
|
||||
fast_cast<TR, float>((TR*)cos_sin_lm, cos_sin_lm, rot_dim);
|
||||
}
|
||||
// generate query
|
||||
for (int64_t i = 0; i < num_heads; i += head_num) {
|
||||
// int loop_size = min<int64_t>(n - nstart, lm_buffer_size);
|
||||
int64_t element_num = min<int64_t>(num_heads - i, head_num);
|
||||
GM2LM(query + offset_q, data_lm, element_num * head_size * sizeof(T));
|
||||
fast_cast<T, float>((T*)data_lm, data_lm, element_num * head_size);
|
||||
// calculation
|
||||
apply_rotary_pos_emb_nonexo(data_lm, cos_sin_lm, rot_dim, head_size, element_num);
|
||||
// write back to GM
|
||||
fast_cast<float, T>(data_lm, (T*)data_lm, element_num * head_size);
|
||||
LM2GM(data_lm, query + offset_q, element_num * head_size * sizeof(T));
|
||||
offset_q += (head_size * element_num);
|
||||
}
|
||||
// generate key
|
||||
for (int64_t j = 0; j < num_kv_heads; j += head_num) {
|
||||
int64_t element_num = min<int64_t>(num_kv_heads - j, head_num);
|
||||
GM2LM(key + offset_k, data_lm, element_num * head_size * sizeof(T));
|
||||
fast_cast<T, float>((T*)data_lm, data_lm, element_num * head_size);
|
||||
// calculation
|
||||
apply_rotary_pos_emb_nonexo(data_lm, cos_sin_lm, rot_dim, head_size, element_num);
|
||||
// write back to GM
|
||||
fast_cast<float, T>(data_lm, (T*)data_lm, element_num * head_size);
|
||||
LM2GM(data_lm, key + offset_k, element_num * head_size * sizeof(T));
|
||||
offset_k += (head_size * element_num);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#define _XPU_DEF_(T, TR) \
|
||||
template __global__ void rotary_embedding_neox<T, TR>( \
|
||||
const int* positions, \
|
||||
T* query, \
|
||||
T* key, \
|
||||
const TR* cos_sin_cache, \
|
||||
const int rot_dim, \
|
||||
const int num_heads, \
|
||||
const int num_kv_heads, \
|
||||
const int head_size, \
|
||||
const int32_t num_tokens);
|
||||
|
||||
_XPU_DEF_(float, float);
|
||||
_XPU_DEF_(float16, float16);
|
||||
_XPU_DEF_(bfloat16, float);
|
||||
_XPU_DEF_(bfloat16, bfloat16);
|
||||
|
||||
}
|
||||
}
|
||||
@@ -6,10 +6,10 @@ namespace xpu3 {
|
||||
namespace plugin {
|
||||
|
||||
__global__ void get_position_ids(const int *seq_lens_encoder,
|
||||
const int *seq_lens_decoder,
|
||||
const int *seq_lens_this_time,
|
||||
int *position_ids,
|
||||
const int bs) {
|
||||
const int *seq_lens_decoder,
|
||||
const int *seq_lens_this_time,
|
||||
int *position_ids,
|
||||
const int bs) {
|
||||
int cid = core_id();
|
||||
int ncores = core_num();
|
||||
int clusterid = cluster_id();
|
||||
@@ -26,31 +26,39 @@ __global__ void get_position_ids(const int *seq_lens_encoder,
|
||||
GM2LM_ASYNC(seq_lens_encoder, seq_lens_encoder_lm, bs * sizeof(int));
|
||||
GM2LM_ASYNC(seq_lens_decoder, seq_lens_decoder_lm, bs * sizeof(int));
|
||||
GM2LM(seq_lens_this_time, seq_lens_this_time_lm, bs * sizeof(int));
|
||||
|
||||
int offset = 0;
|
||||
for(int ind = tid; ind < bs; ind += ncores * nclusters){
|
||||
int encode_sum = 0;
|
||||
int is_encode = 0;
|
||||
for (int i = 0; i < bs; i++) {
|
||||
if (seq_lens_encoder_lm[i] > 0) {
|
||||
encode_sum = encode_sum + seq_lens_encoder_lm[i];
|
||||
}
|
||||
}
|
||||
for (int ind = tid; ind < bs; ind += ncores * nclusters) {
|
||||
int decode_offset = encode_sum;
|
||||
int encode_offset = 0;
|
||||
for (int i = 0; i < ind; i++) {
|
||||
offset += seq_lens_encoder_lm[i];
|
||||
if (seq_lens_decoder_lm[i] > 0) {
|
||||
offset += seq_lens_this_time_lm[i];
|
||||
if (seq_lens_encoder_lm[i] > 0) {
|
||||
encode_offset += seq_lens_encoder_lm[i];
|
||||
} else {
|
||||
decode_offset += seq_lens_this_time_lm[i];
|
||||
}
|
||||
}
|
||||
int encoder_len = seq_lens_encoder_lm[ind];
|
||||
int decoder_len = seq_lens_decoder_lm[ind];
|
||||
int seq_len_this_time = decoder_len > 0 ? seq_lens_this_time_lm[ind] : 0;
|
||||
int offset = decoder_len > 0 ? decode_offset : encode_offset;
|
||||
|
||||
int begin_offset = offset;
|
||||
for (int i = 0; i < encoder_len; i++) {
|
||||
position_ids_lm[offset + i] = i;
|
||||
}
|
||||
offset += encoder_len;
|
||||
if (decoder_len > 0) {
|
||||
for (int i = 0; i < seq_len_this_time; i++) {
|
||||
position_ids_lm[offset + i] = decoder_len + i; // 使用 decoder 长度本身
|
||||
for (int i = 0; i < encoder_len + seq_len_this_time; i += buf_len) {
|
||||
int read_len = min(encoder_len + seq_len_this_time - i, buf_len);
|
||||
for (int j = 0; j < buf_len; j++) {
|
||||
int ind_id = i + j;
|
||||
if (ind_id < encoder_len) {
|
||||
position_ids_lm[j] = ind_id;
|
||||
} else {
|
||||
position_ids_lm[j] = decoder_len + ind_id;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (encoder_len + seq_len_this_time > 0) {
|
||||
LM2GM(position_ids_lm + begin_offset, position_ids + begin_offset, (encoder_len + seq_len_this_time) * sizeof(int));
|
||||
LM2GM(position_ids_lm, position_ids + offset + i, read_len * sizeof(int));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -27,7 +27,7 @@ __global__ void update_repeat_times(const int64_t *pre_ids,
|
||||
int nclusters = cluster_num();
|
||||
int tid = clusterid * ncores + cid;
|
||||
|
||||
const int max_sm_len = 256 * 1024 / sizeof(int);
|
||||
const int max_sm_len = 128 * 1024 / sizeof(int);
|
||||
__shared__ int repeated_times_sm[max_sm_len];
|
||||
int64_t pre_id_lm;
|
||||
int n_length = (length + max_sm_len - 1) / max_sm_len;
|
||||
|
||||
168
csrc/xpu/src/plugin/src/wrapper/absorb_mla_cache_write.cpp
Normal file
168
csrc/xpu/src/plugin/src/wrapper/absorb_mla_cache_write.cpp
Normal file
@@ -0,0 +1,168 @@
|
||||
// Copyright (c) 2024 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 <algorithm>
|
||||
#include <numeric>
|
||||
#include "xpu/plugin.h"
|
||||
#include "xpu/refactor/impl_public/wrapper_check.h"
|
||||
|
||||
|
||||
namespace xpu3 {
|
||||
namespace plugin {
|
||||
|
||||
template <typename TI, typename TO, typename TID>
|
||||
__attribute__((global)) void absorb_mla_cache_write(const TI* src,
|
||||
TO* dst,
|
||||
const TID* block_table,
|
||||
int32_t* kv_seq_lod,
|
||||
int32_t* start_token,
|
||||
int32_t* real_batch,
|
||||
int64_t batch_size,
|
||||
int64_t kv_head_num,
|
||||
int64_t head_dim,
|
||||
int64_t max_seq_num,
|
||||
int64_t block_size,
|
||||
int64_t max_num_blocks_per_seq);
|
||||
|
||||
} // namespace plugin
|
||||
} // namespace xpu3
|
||||
|
||||
namespace baidu {
|
||||
namespace xpu {
|
||||
namespace api {
|
||||
namespace plugin {
|
||||
|
||||
template <typename TI, typename TO, typename TID>
|
||||
static int xpu3_wrapper(
|
||||
Context* ctx,
|
||||
const TI* x,
|
||||
TO* y,
|
||||
const TID* block_table,
|
||||
const VectorParam<int32_t>& kv_seq_lod_xpu,
|
||||
const VectorParam<int32_t>& start_tokens_xpu,
|
||||
const VectorParam<int32_t>& real_batch_xpu,
|
||||
int64_t batch_size,
|
||||
int64_t kv_head_num,
|
||||
int64_t head_dim,
|
||||
int64_t max_seq_num,
|
||||
int64_t block_size,
|
||||
int64_t max_num_blocks_per_seq) {
|
||||
|
||||
auto func = &xpu3::plugin::absorb_mla_cache_write<TI, TO, TID>;
|
||||
func<<<ctx->ncluster(), 64, ctx->xpu_stream>>>(
|
||||
x,
|
||||
y,
|
||||
block_table,
|
||||
kv_seq_lod_xpu.xpu,
|
||||
start_tokens_xpu.xpu,
|
||||
real_batch_xpu.xpu,
|
||||
batch_size,
|
||||
kv_head_num,
|
||||
head_dim,
|
||||
max_seq_num,
|
||||
block_size,
|
||||
max_num_blocks_per_seq);
|
||||
|
||||
return api::SUCCESS;
|
||||
|
||||
}
|
||||
|
||||
template <typename TI, typename TO, typename TID>
|
||||
int absorb_mla_cache_write_xpu(Context *ctx,
|
||||
const TI* x,
|
||||
TO* y,
|
||||
const TID* block_table,
|
||||
const VectorParam<int32_t>& kv_seq_lod,
|
||||
const VectorParam<int32_t>& start_tokens,
|
||||
const VectorParam<int32_t>& real_batch,
|
||||
int64_t batch_size,
|
||||
int64_t kv_head_num,
|
||||
int64_t head_dim,
|
||||
int64_t max_seq_num,
|
||||
int64_t block_size,
|
||||
int64_t max_num_blocks_per_seq) {
|
||||
WRAPPER_CHECK_CTX(ctx);
|
||||
WRAPPER_DUMP_FUNCTION_T3(ctx, "absorb_mla_cache_write", TI, TO, TID);
|
||||
WRAPPER_DUMP_PARAM6(ctx, x, y, block_table, kv_seq_lod, start_tokens, real_batch);
|
||||
WRAPPER_DUMP_PARAM6(ctx, batch_size, kv_head_num, head_dim, max_seq_num, block_size, max_num_blocks_per_seq);
|
||||
WRAPPER_DUMP(ctx);
|
||||
// check vector param size
|
||||
WRAPPER_ASSERT_EQ(ctx, batch_size + 1, kv_seq_lod.len);
|
||||
WRAPPER_ASSERT_EQ(ctx, batch_size, start_tokens.len);
|
||||
WRAPPER_ASSERT_EQ(ctx, batch_size, real_batch.len);
|
||||
WRAPPER_ASSERT_EQ(ctx, kv_head_num, 1);
|
||||
WRAPPER_ASSERT_LE(ctx, batch_size, 32);
|
||||
// check input ptr shape
|
||||
int64_t kv_seqlen_sum = kv_seq_lod.cpu[batch_size];
|
||||
int64_t hidden_dim = kv_head_num * head_dim;
|
||||
// check kv_cache and block_table ptr shape
|
||||
int64_t table_count = max_seq_num * max_num_blocks_per_seq; // block number in a block table
|
||||
int64_t kv_cache_size = table_count * kv_head_num * head_dim * block_size;
|
||||
// x 'shape is supposed to be [kv_seqlen_sum, hidden_dim]
|
||||
WRAPPER_CHECK_PTR(ctx, TI, kv_seqlen_sum * hidden_dim, x);
|
||||
// y is kv cache ptr, which is supposed to be [table_count, kv_head_num, block_size, head_dim]
|
||||
WRAPPER_CHECK_PTR(ctx, TO, kv_cache_size, y);
|
||||
|
||||
// int64_t single_block_occupy = kv_head_num * head_dim * block_size; // single_block_occupy represent how many bytes
|
||||
// occupied in a block
|
||||
// which is different from block size
|
||||
// block_size means the seqlen a block can hold
|
||||
// kv_cache 'shape is supposed to be [block_num, hidden_dim, block_size]
|
||||
// FIXME: in real cases, kv cache may not be full
|
||||
// WRAPPER_CHECK_PTR(ctx, TO, table_count * single_block_occupy, y);
|
||||
// block_table 'shape is supposed to be [max_seq_num, max_num_blocks_per_seq]
|
||||
WRAPPER_CHECK_PTR(ctx, TID, table_count, block_table);
|
||||
|
||||
if (ctx->dev().type() == api::kXPU3) {
|
||||
return xpu3_wrapper<TI, TO, TID>(ctx,
|
||||
x,
|
||||
y,
|
||||
block_table,
|
||||
kv_seq_lod,
|
||||
start_tokens,
|
||||
real_batch,
|
||||
batch_size,
|
||||
kv_head_num,
|
||||
head_dim,
|
||||
max_seq_num,
|
||||
block_size,
|
||||
max_num_blocks_per_seq);
|
||||
}
|
||||
WRAPPER_UNIMPLEMENTED(ctx);
|
||||
}
|
||||
|
||||
#define INSTANTIATION_ABSORB_MLA_CACHE_WRITE(TI, TO, TID) \
|
||||
template int absorb_mla_cache_write_xpu<TI, TO, TID>(api::Context*, \
|
||||
const TI*, \
|
||||
TO*, \
|
||||
const TID*, \
|
||||
const api::VectorParam<int32_t>&, \
|
||||
const api::VectorParam<int32_t>&, \
|
||||
const api::VectorParam<int32_t>&, \
|
||||
int64_t, \
|
||||
int64_t, \
|
||||
int64_t, \
|
||||
int64_t, \
|
||||
int64_t, \
|
||||
int64_t);
|
||||
|
||||
INSTANTIATION_ABSORB_MLA_CACHE_WRITE(float16, float16, int32_t);
|
||||
INSTANTIATION_ABSORB_MLA_CACHE_WRITE(bfloat16, bfloat16, int32_t);
|
||||
|
||||
|
||||
|
||||
} // namespace plugin
|
||||
} // namespace api
|
||||
} // namespace xpu
|
||||
} // namespace baidu
|
||||
@@ -0,0 +1,166 @@
|
||||
// Copyright (c) 2024 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 <algorithm>
|
||||
#include <numeric>
|
||||
|
||||
#include "xpu/plugin.h"
|
||||
#include "xpu/refactor/impl_public/wrapper_check.h"
|
||||
|
||||
namespace xpu3 {
|
||||
namespace plugin {
|
||||
template <typename T, typename TR>
|
||||
__attribute__((global)) void rotary_embedding_neox(
|
||||
const int* positions, // [num_tokens]
|
||||
T* query, // [num_tokens, num_heads, head_size]
|
||||
T* key, // [num_tokens, num_kv_heads, head_size]
|
||||
const TR* cos_sin_cache, // [max_position, 2, rot_dim // 2]
|
||||
const int rot_dim,
|
||||
const int num_heads,
|
||||
const int num_kv_heads,
|
||||
const int head_size,
|
||||
const int32_t num_tokens);
|
||||
} // namespace plugin
|
||||
} // namespace xpu3
|
||||
|
||||
namespace baidu {
|
||||
namespace xpu {
|
||||
namespace api {
|
||||
namespace plugin {
|
||||
|
||||
// template <typename T>
|
||||
// static int cpu_wrapper() {
|
||||
// }
|
||||
|
||||
template <typename T, typename TR>
|
||||
static int xpu3_wrapper(Context* ctx,
|
||||
const int* positions, // [num_tokens]
|
||||
T* query, // [num_tokens, num_heads, head_size]
|
||||
T* key, // [num_tokens, num_kv_heads, head_size]
|
||||
const TR* cos_sin_cache, // [max_position, 2, rot_dim // 2]
|
||||
const int rot_dim,
|
||||
const int num_heads,
|
||||
const int num_kv_heads,
|
||||
const int head_size,
|
||||
const int32_t num_tokens) {
|
||||
using XPU_TID = typename XPUIndexType<T>::type;
|
||||
using XPU_TRID = typename XPUIndexType<TR>::type;
|
||||
using XPU_INT64 = typename XPUIndexType<int64_t>::type;
|
||||
auto rotary_embedding_neox = xpu3::plugin::rotary_embedding_neox<XPU_TID, XPU_TRID>;
|
||||
rotary_embedding_neox<<<ctx->ncluster(), 64, ctx->xpu_stream>>>(
|
||||
positions,
|
||||
reinterpret_cast<XPU_TID *>(query),
|
||||
reinterpret_cast<XPU_TID *>(key),
|
||||
reinterpret_cast<const XPU_TRID *>(cos_sin_cache),
|
||||
rot_dim,
|
||||
num_heads,
|
||||
num_kv_heads,
|
||||
head_size,
|
||||
num_tokens);
|
||||
return api::SUCCESS;
|
||||
}
|
||||
|
||||
// template <typename T>
|
||||
// int set_stop_value_multi_ends(Context* ctx,
|
||||
// bool* stop_flags,
|
||||
// T* topk_ids,
|
||||
// T* next_tokens,
|
||||
// const T* end_ids,
|
||||
// const int* seq_lens,
|
||||
// const int bs,
|
||||
// const int end_length,
|
||||
// const bool beam_search)
|
||||
|
||||
|
||||
template <typename T, typename TR>
|
||||
int rotary_embedding_neox(
|
||||
Context* ctx,
|
||||
const int* positions, // [num_tokens]
|
||||
T* query, // [num_tokens, num_heads, head_size]
|
||||
T* key, // [num_tokens, num_kv_heads, head_size]
|
||||
const TR* cos_sin_cache, // [max_position, 2, rot_dim // 2]
|
||||
const int rot_dim,
|
||||
const int num_heads,
|
||||
const int num_kv_heads,
|
||||
const int head_size,
|
||||
const int32_t num_tokens)
|
||||
{
|
||||
WRAPPER_CHECK_CTX(ctx);
|
||||
WRAPPER_DUMP_FUNCTION_T2(ctx, "rotary_embedding_neox", T, TR);
|
||||
WRAPPER_DUMP_PARAM5(
|
||||
ctx, positions, query, key, cos_sin_cache, rot_dim);
|
||||
WRAPPER_DUMP_PARAM4(ctx, num_heads, num_kv_heads, head_size,num_tokens);
|
||||
WRAPPER_DUMP(ctx);
|
||||
|
||||
WRAPPER_CHECK_PTR(ctx, int, num_tokens, positions);
|
||||
WRAPPER_CHECK_PTR(ctx, T, num_tokens * num_heads * head_size, query);
|
||||
WRAPPER_CHECK_PTR(ctx, T, num_tokens * num_kv_heads * head_size, key);
|
||||
WRAPPER_CHECK_PTR(ctx, TR, num_tokens * rot_dim, cos_sin_cache);
|
||||
|
||||
if (ctx->dev().type() == api::kXPU2 || ctx->dev().type() == api::kXPU3) {
|
||||
return xpu3_wrapper<T, TR>(ctx,
|
||||
positions,
|
||||
query,
|
||||
key,
|
||||
cos_sin_cache,
|
||||
rot_dim,
|
||||
num_heads,
|
||||
num_kv_heads,
|
||||
head_size,
|
||||
num_tokens);
|
||||
}
|
||||
WRAPPER_UNIMPLEMENTED(ctx);
|
||||
}
|
||||
|
||||
|
||||
template int rotary_embedding_neox<float, float>(
|
||||
Context* ctx,
|
||||
const int* positions, // [num_tokens]
|
||||
float* query, // [num_tokens, num_heads, head_size]
|
||||
float* key, // [num_tokens, num_kv_heads, head_size]
|
||||
const float* cos_sin_cache, // [max_position, 2, rot_dim // 2]
|
||||
const int rot_dim,
|
||||
const int num_heads,
|
||||
const int num_kv_heads,
|
||||
const int head_size,
|
||||
const int32_t num_tokens);
|
||||
|
||||
template int rotary_embedding_neox<float16, float16>(
|
||||
Context* ctx,
|
||||
const int* positions, // [num_tokens]
|
||||
float16* query, // [num_tokens, num_heads, head_size]
|
||||
float16* key, // [num_tokens, num_kv_heads, head_size]
|
||||
const float16* cos_sin_cache, // [max_position, 2, rot_dim // 2]
|
||||
const int rot_dim,
|
||||
const int num_heads,
|
||||
const int num_kv_heads,
|
||||
const int head_size,
|
||||
const int32_t num_tokens);
|
||||
|
||||
template int rotary_embedding_neox<bfloat16, bfloat16>(
|
||||
Context* ctx,
|
||||
const int* positions, // [num_tokens]
|
||||
bfloat16* query, // [num_tokens, num_heads, head_size]
|
||||
bfloat16* key, // [num_tokens, num_kv_heads, head_size]
|
||||
const bfloat16* cos_sin_cache, // [max_position, 2, rot_dim // 2]
|
||||
const int rot_dim,
|
||||
const int num_heads,
|
||||
const int num_kv_heads,
|
||||
const int head_size,
|
||||
const int32_t num_tokens);
|
||||
|
||||
} // namespace plugin
|
||||
} // namespace api
|
||||
} // namespace xpu
|
||||
} // namespace baidu
|
||||
@@ -163,7 +163,7 @@ int set_value_by_flags_and_idx(Context* ctx,
|
||||
WRAPPER_CHECK_SHAPE(ctx, &seq_lens_encoder_len, {bs});
|
||||
WRAPPER_CHECK_SHAPE(ctx, &seq_lens_decoder_len, {bs});
|
||||
WRAPPER_CHECK_SHAPE(ctx, &step_idx_len, {bs});
|
||||
WRAPPER_CHECK_PTR(ctx, int64_t, stop_flags_len, stop_flags);
|
||||
WRAPPER_CHECK_PTR(ctx, bool, stop_flags_len, stop_flags);
|
||||
WRAPPER_CHECK_PTR(ctx, int64_t, pre_ids_all_len, pre_ids_all);
|
||||
WRAPPER_CHECK_PTR(ctx, int64_t, input_ids_len, input_ids);
|
||||
WRAPPER_CHECK_PTR(ctx, int, seq_lens_encoder_len, seq_lens_encoder);
|
||||
|
||||
@@ -1,192 +0,0 @@
|
||||
// Copyright (c) 2024 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 <algorithm>
|
||||
#include <numeric>
|
||||
#include "xpu/plugin.h"
|
||||
#include "xpu/refactor/impl_public/wrapper_check.h"
|
||||
|
||||
namespace xpu2 {
|
||||
namespace plugin {
|
||||
|
||||
__attribute__((global)) void update_inputs(bool *not_need_stop,
|
||||
int *seq_lens_this_time,
|
||||
int *seq_lens_encoder,
|
||||
int *seq_lens_decoder,
|
||||
int64_t *input_ids,
|
||||
const int64_t *stop_nums,
|
||||
const bool *stop_flags,
|
||||
const bool *is_block_step,
|
||||
const int64_t *next_tokens,
|
||||
const int bsz,
|
||||
const int max_bsz,
|
||||
const int input_ids_stride);
|
||||
|
||||
} // namespace plugin
|
||||
} // namespace xpu2
|
||||
|
||||
namespace xpu3 {
|
||||
namespace plugin {
|
||||
|
||||
__attribute__((global)) void update_inputs(bool *not_need_stop,
|
||||
int *seq_lens_this_time,
|
||||
int *seq_lens_encoder,
|
||||
int *seq_lens_decoder,
|
||||
int64_t *input_ids,
|
||||
const int64_t *stop_nums,
|
||||
const bool *stop_flags,
|
||||
const bool *is_block_step,
|
||||
const int64_t *next_tokens,
|
||||
const int bsz,
|
||||
const int max_bsz,
|
||||
const int input_ids_stride);
|
||||
|
||||
} // namespace plugin
|
||||
} // namespace xpu3
|
||||
|
||||
namespace baidu {
|
||||
namespace xpu {
|
||||
namespace api {
|
||||
namespace plugin {
|
||||
|
||||
static int cpu_wrapper(Context *ctx,
|
||||
bool *not_need_stop,
|
||||
int *seq_lens_this_time,
|
||||
int *seq_lens_encoder,
|
||||
int *seq_lens_decoder,
|
||||
int64_t *input_ids,
|
||||
const int64_t *stop_nums,
|
||||
const bool *stop_flags,
|
||||
const bool *is_block_step,
|
||||
const int64_t *next_tokens,
|
||||
const int bsz,
|
||||
const int max_bsz,
|
||||
const int input_ids_stride) {
|
||||
std::vector<int64_t> stop_flag_now_int(max_bsz, 1);
|
||||
for (int i = 0; i < bsz; i++) {
|
||||
bool stop_flags_now = stop_flags[i];
|
||||
stop_flag_now_int[i] = is_block_step[i] ? 0 : stop_flags_now;
|
||||
const int seq_len_encoder = seq_lens_encoder[i];
|
||||
const int seq_len_decoder = seq_lens_decoder[i];
|
||||
|
||||
seq_lens_decoder[i] =
|
||||
stop_flags[i]
|
||||
? 0
|
||||
: (seq_len_decoder == 0 ? seq_len_encoder : seq_len_decoder + 1);
|
||||
|
||||
seq_lens_this_time[i] = stop_flags[i] ? 0 : 1;
|
||||
seq_lens_encoder[i] = 0;
|
||||
int64_t *input_ids_now = input_ids + i * input_ids_stride;
|
||||
input_ids_now[0] = next_tokens[i];
|
||||
}
|
||||
int64_t stop_sum = 0;
|
||||
for (size_t i = 0; i < stop_flag_now_int.size(); i++) {
|
||||
stop_sum += stop_flag_now_int[i];
|
||||
}
|
||||
not_need_stop[0] = stop_sum < stop_nums[0];
|
||||
return api::SUCCESS;
|
||||
}
|
||||
|
||||
static int xpu2or3_wrapper(Context *ctx,
|
||||
bool *not_need_stop,
|
||||
int *seq_lens_this_time,
|
||||
int *seq_lens_encoder,
|
||||
int *seq_lens_decoder,
|
||||
int64_t *input_ids,
|
||||
const int64_t *stop_nums,
|
||||
const bool *stop_flags,
|
||||
const bool *is_block_step,
|
||||
const int64_t *next_tokens,
|
||||
const int bsz,
|
||||
const int max_bsz,
|
||||
const int input_ids_stride) {
|
||||
using XPU_INT64 = typename XPUIndexType<int64_t>::type;
|
||||
auto update_inputs = xpu2::plugin::update_inputs;
|
||||
update_inputs<<<ctx->ncluster(), 64, ctx->xpu_stream>>>(
|
||||
not_need_stop,
|
||||
seq_lens_this_time,
|
||||
seq_lens_encoder,
|
||||
seq_lens_decoder,
|
||||
reinterpret_cast<XPU_INT64 *>(input_ids),
|
||||
reinterpret_cast<const XPU_INT64 *>(stop_nums),
|
||||
stop_flags,
|
||||
is_block_step,
|
||||
reinterpret_cast<const XPU_INT64 *>(next_tokens),
|
||||
bsz,
|
||||
max_bsz,
|
||||
input_ids_stride);
|
||||
return api::SUCCESS;
|
||||
}
|
||||
|
||||
int update_inputs(Context *ctx,
|
||||
bool *not_need_stop,
|
||||
int *seq_lens_this_time,
|
||||
int *seq_lens_encoder,
|
||||
int *seq_lens_decoder,
|
||||
int64_t *input_ids,
|
||||
const int64_t *stop_nums,
|
||||
const bool *stop_flags,
|
||||
const bool *is_block_step,
|
||||
const int64_t *next_tokens,
|
||||
const int bsz,
|
||||
const int max_bsz,
|
||||
const int input_ids_stride) {
|
||||
WRAPPER_CHECK_CTX(ctx);
|
||||
WRAPPER_DUMP_FUNCTION_T1(ctx, "update_inputs", int);
|
||||
WRAPPER_DUMP_PARAM5(ctx,
|
||||
not_need_stop,
|
||||
seq_lens_this_time,
|
||||
seq_lens_encoder,
|
||||
seq_lens_decoder,
|
||||
input_ids);
|
||||
WRAPPER_DUMP_PARAM4(ctx, stop_nums, stop_flags, is_block_step, next_tokens);
|
||||
WRAPPER_DUMP_PARAM3(ctx, bsz, max_bsz, input_ids_stride);
|
||||
WRAPPER_DUMP(ctx);
|
||||
if (ctx->dev().type() == api::kCPU) {
|
||||
return cpu_wrapper(ctx,
|
||||
not_need_stop,
|
||||
seq_lens_this_time,
|
||||
seq_lens_encoder,
|
||||
seq_lens_decoder,
|
||||
input_ids,
|
||||
stop_nums,
|
||||
stop_flags,
|
||||
is_block_step,
|
||||
next_tokens,
|
||||
bsz,
|
||||
max_bsz,
|
||||
input_ids_stride);
|
||||
}
|
||||
if (ctx->dev().type() == api::kXPU2 || ctx->dev().type() == api::kXPU3) {
|
||||
return xpu2or3_wrapper(ctx,
|
||||
not_need_stop,
|
||||
seq_lens_this_time,
|
||||
seq_lens_encoder,
|
||||
seq_lens_decoder,
|
||||
input_ids,
|
||||
stop_nums,
|
||||
stop_flags,
|
||||
is_block_step,
|
||||
next_tokens,
|
||||
bsz,
|
||||
max_bsz,
|
||||
input_ids_stride);
|
||||
}
|
||||
WRAPPER_UNIMPLEMENTED(ctx);
|
||||
}
|
||||
|
||||
} // namespace plugin
|
||||
} // namespace api
|
||||
} // namespace xpu
|
||||
} // namespace baidu
|
||||
@@ -191,17 +191,35 @@ int update_inputs_v2(Context* ctx,
|
||||
stop_flags,
|
||||
seq_lens_this_time,
|
||||
seq_lens_encoder);
|
||||
|
||||
WRAPPER_DUMP_PARAM5(ctx,
|
||||
seq_lens_decoder,
|
||||
next_tokens,
|
||||
kwargs_next_tokens,
|
||||
input_ids,
|
||||
end_ids);
|
||||
|
||||
WRAPPER_DUMP_PARAM3(ctx, stop_nums, is_block_step, max_dec_len);
|
||||
WRAPPER_DUMP_PARAM4(ctx, now_bsz, max_bsz, input_ids_stride, end_length);
|
||||
|
||||
|
||||
WRAPPER_DUMP(ctx);
|
||||
WRAPPER_ASSERT_LE(ctx, max_bsz, 1024);
|
||||
WRAPPER_ASSERT_LE(ctx, now_bsz, max_bsz);
|
||||
WRAPPER_CHECK_PTR(ctx, bool, 1, not_need_stop);
|
||||
WRAPPER_CHECK_PTR(ctx, int64_t, max_bsz, step_idx);
|
||||
WRAPPER_CHECK_PTR(ctx, bool, max_bsz, stop_flags);
|
||||
WRAPPER_CHECK_PTR(ctx, int, max_bsz, seq_lens_this_time);
|
||||
WRAPPER_CHECK_PTR(ctx, int, max_bsz, seq_lens_encoder);
|
||||
WRAPPER_CHECK_PTR(ctx, int, max_bsz, seq_lens_decoder);
|
||||
WRAPPER_CHECK_PTR(ctx, int64_t, max_bsz, next_tokens);
|
||||
WRAPPER_CHECK_PTR(ctx, int64_t, max_bsz, kwargs_next_tokens);
|
||||
WRAPPER_CHECK_PTR(ctx, int64_t, max_bsz, input_ids);
|
||||
WRAPPER_CHECK_PTR(ctx, int64_t, end_length, end_ids);
|
||||
WRAPPER_CHECK_PTR(ctx, int64_t, 1, stop_nums);
|
||||
WRAPPER_CHECK_PTR(ctx, bool, max_bsz, is_block_step);
|
||||
WRAPPER_CHECK_PTR(ctx, int64_t, max_bsz, max_dec_len);
|
||||
|
||||
|
||||
// std::cout << "wht --- ctx debug level is " << ctx->debug_level() << std::endl;
|
||||
// TODO(mayang02): check ptrs
|
||||
if (ctx->dev().type() == api::kCPU) {
|
||||
return cpu_wrapper(ctx,
|
||||
|
||||
166
csrc/xpu/src/save_with_output.cc
Normal file
166
csrc/xpu/src/save_with_output.cc
Normal file
@@ -0,0 +1,166 @@
|
||||
// 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 <iostream>
|
||||
#include <fstream>
|
||||
#include <memory>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
|
||||
#include "stdlib.h"
|
||||
#include <stdio.h>
|
||||
#include <dlfcn.h> // dladdr
|
||||
#include <sys/time.h>
|
||||
#include <sys/stat.h>
|
||||
#include "paddle/extension.h"
|
||||
|
||||
constexpr char kSEP = '/';
|
||||
|
||||
std::string DirName(const std::string &filepath) {
|
||||
auto pos = filepath.rfind(kSEP);
|
||||
if (pos == std::string::npos) {
|
||||
return "";
|
||||
}
|
||||
return filepath.substr(0, pos);
|
||||
}
|
||||
|
||||
bool FileExists(const std::string &filepath) {
|
||||
struct stat buffer;
|
||||
return (stat(filepath.c_str(), &buffer) == 0);
|
||||
}
|
||||
|
||||
void MkDir(const char *path) {
|
||||
std::string path_error(path);
|
||||
path_error += " mkdir failed!";
|
||||
if (mkdir(path, 0755)) {
|
||||
if (errno != EEXIST) {
|
||||
throw std::runtime_error(path_error);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void MkDirRecursively(const char *fullpath) {
|
||||
if (*fullpath == '\0') return; // empty string
|
||||
if (FileExists(fullpath)) return;
|
||||
MkDirRecursively(DirName(fullpath).c_str());
|
||||
MkDir(fullpath);
|
||||
}
|
||||
|
||||
|
||||
template<typename data_t>
|
||||
void saveToFile(std::ostream & os, const void* x_data, std::vector<int64_t> shape, int64_t x_numel, const char type_id) {
|
||||
// 1.type
|
||||
os.write(reinterpret_cast<const char *>(&type_id),sizeof(type_id));
|
||||
// 2.data
|
||||
uint64_t size = x_numel * sizeof(data_t);
|
||||
os.write(static_cast<const char*>(x_data),static_cast<std::streamsize>(size));
|
||||
|
||||
}
|
||||
|
||||
template<typename data_t>
|
||||
void save_with_output_kernel(const paddle::Tensor& x,
|
||||
const paddle::Tensor& batch_idx,
|
||||
const paddle::Tensor& step_idx,
|
||||
std::string file_path,
|
||||
int64_t rank_id,
|
||||
char type_id) {
|
||||
std::vector<int64_t> x_shape = x.shape();
|
||||
|
||||
if(rank_id >= 0) {
|
||||
file_path += "_rank_" + std::to_string(rank_id);
|
||||
}
|
||||
|
||||
int batch_idx_data = -1, step_idx_data = -1;
|
||||
|
||||
if(batch_idx.is_gpu()) {
|
||||
paddle::Tensor batch_idx_cpu = batch_idx.copy_to<int32_t>(paddle::CPUPlace());
|
||||
batch_idx_data = batch_idx_cpu.data<int32_t>()[0];
|
||||
} else {
|
||||
batch_idx_data = batch_idx.data<int32_t>()[0];
|
||||
}
|
||||
if(step_idx.is_gpu()) {
|
||||
paddle::Tensor step_idx_cpu = step_idx.copy_to<int64_t>(paddle::CPUPlace());
|
||||
step_idx_data = step_idx_cpu.data<int64_t>()[0];
|
||||
} else {
|
||||
step_idx_data = step_idx.data<int64_t>()[0];
|
||||
}
|
||||
auto x_data = x.data<data_t>();
|
||||
|
||||
if(batch_idx_data >= 0) {
|
||||
file_path += "_batch_" + std::to_string(batch_idx_data);
|
||||
}
|
||||
if(step_idx_data >= 0) {
|
||||
file_path += "_step_" + std::to_string(step_idx_data);
|
||||
}
|
||||
MkDirRecursively(DirName(file_path).c_str());
|
||||
std::ofstream fout(file_path, std::ios::binary);
|
||||
fout.write("0",1);
|
||||
saveToFile<data_t>(fout, x_data, x_shape, x.numel(),type_id);
|
||||
fout.seekp(std::ios::beg);
|
||||
fout.write("1",1);
|
||||
fout.close();
|
||||
|
||||
}
|
||||
|
||||
void print_shape(const paddle::Tensor& tmp, char *tmp_str){
|
||||
std::vector<int64_t> shape = tmp.shape();
|
||||
printf("%s's shape: \n", tmp_str);
|
||||
for(int i=0; i < shape.size(); i++) {
|
||||
printf("%d ", (int)shape[i]);
|
||||
}
|
||||
printf("\n");
|
||||
}
|
||||
|
||||
std::vector<paddle::Tensor> SaveWithOutputForward(const paddle::Tensor& x,
|
||||
const paddle::Tensor& batch_idx,
|
||||
const paddle::Tensor& step_idx,
|
||||
std::string file_path,
|
||||
int64_t rank_id) {
|
||||
auto out = x.copy_to(paddle::CPUPlace(), false);
|
||||
switch(x.type()) {
|
||||
case paddle::DataType::FLOAT32:
|
||||
save_with_output_kernel<float>(out, batch_idx, step_idx, file_path, rank_id, '0');
|
||||
break;
|
||||
case paddle::DataType::INT64:
|
||||
save_with_output_kernel<int64_t>(out, batch_idx, step_idx, file_path, rank_id,'1');
|
||||
break;
|
||||
case paddle::DataType::INT32:
|
||||
save_with_output_kernel<int32_t>(out, batch_idx, step_idx, file_path, rank_id, '2');
|
||||
break;
|
||||
default:
|
||||
PD_THROW("function SaveWithOutputForward is not implemented for data type");
|
||||
}
|
||||
return {out};
|
||||
}
|
||||
|
||||
std::vector<std::vector<int64_t>> SaveWithOutputInferShape(const std::vector<int64_t>& x_shape,
|
||||
const std::vector<int64_t>& batch_idx_shape,
|
||||
const std::vector<int64_t>& step_idx_shape) {
|
||||
return {x_shape};
|
||||
}
|
||||
|
||||
std::vector<paddle::DataType> SaveWithOutputInferDtype(const paddle::DataType& x_dtype,
|
||||
const paddle::DataType& batch_idx_dtype,
|
||||
const paddle::DataType& step_idx_dtype) {
|
||||
return {x_dtype};
|
||||
}
|
||||
|
||||
PD_BUILD_OP(save_with_output)
|
||||
.Inputs({"x", "batch_idx", "step_idx"})
|
||||
.Attrs({"file_path: std::string",
|
||||
"rank_id: int64_t"})
|
||||
.Outputs({"out"})
|
||||
.SetKernelFn(PD_KERNEL(SaveWithOutputForward))
|
||||
.SetInferShapeFn(PD_INFER_SHAPE(SaveWithOutputInferShape))
|
||||
.SetInferDtypeFn(PD_INFER_DTYPE(SaveWithOutputInferDtype));
|
||||
@@ -1,11 +1,11 @@
|
||||
// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
||||
//
|
||||
// Copyright (c) 2024 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.
|
||||
@@ -18,41 +18,54 @@
|
||||
#include <sys/msg.h>
|
||||
#include <sys/types.h>
|
||||
#include "paddle/extension.h"
|
||||
#include "xpu/plugin.h"
|
||||
|
||||
#define MAX_BSZ 256
|
||||
#define MAX_BSZ 512
|
||||
|
||||
struct msgdata {
|
||||
long mtype; // NOLINT
|
||||
int mtext[MAX_BSZ + 2]; // stop_flag, bsz, tokens
|
||||
long mtype;
|
||||
int mtext[MAX_BSZ + 2]; // stop_flag, bsz, tokens
|
||||
};
|
||||
|
||||
void SaveOutMmsg(const paddle::Tensor& x,
|
||||
const paddle::Tensor& not_need_stop,
|
||||
const paddle::Tensor& not_need_stop, // cpu
|
||||
const paddle::Tensor& msg_queue_id, // cpu
|
||||
int64_t rank_id) {
|
||||
if (rank_id > 0) return;
|
||||
auto x_cpu = x.copy_to(paddle::CPUPlace(), false);
|
||||
int64_t* x_data = x_cpu.data<int64_t>();
|
||||
static struct msgdata msg_sed;
|
||||
static key_t key = ftok("./", 1);
|
||||
static int msgid = msgget(key, IPC_CREAT | 0666);
|
||||
|
||||
msg_sed.mtype = 1;
|
||||
bool not_need_stop_data = not_need_stop.data<bool>()[0];
|
||||
msg_sed.mtext[0] = not_need_stop_data ? 1 : -1;
|
||||
int bsz = x.shape()[0];
|
||||
msg_sed.mtext[1] = bsz;
|
||||
for (int i = 2; i < bsz + 2; i++) {
|
||||
msg_sed.mtext[i] = static_cast<int>(x_data[i - 2]);
|
||||
}
|
||||
if ((msgsnd(msgid, &msg_sed, (MAX_BSZ + 2) * 4, IPC_NOWAIT)) == -1) {
|
||||
printf("full msg buffer\n");
|
||||
}
|
||||
return;
|
||||
baidu::xpu::api::plugin::print_times("[TIME BEGIN] SaveOutMmsg");
|
||||
if (rank_id > 0) return;
|
||||
auto x_cpu = x.copy_to(paddle::CPUPlace(), false);
|
||||
int64_t *x_data = x_cpu.data<int64_t>();
|
||||
auto not_need_stop_data = not_need_stop.data<bool>()[0];
|
||||
|
||||
std::cout << "not_need_stop_data: " << not_need_stop_data << std::endl;
|
||||
|
||||
static struct msgdata msg_sed;
|
||||
int msg_queue_id_val = msg_queue_id.data<int>()[0];
|
||||
|
||||
std::cout << "msg_queue_id_val: " << msg_queue_id_val << std::endl;
|
||||
|
||||
static key_t key = ftok("./", msg_queue_id_val);
|
||||
static int msgid = msgget(key, IPC_CREAT | 0666);
|
||||
|
||||
msg_sed.mtype = 1;
|
||||
msg_sed.mtext[0] = not_need_stop_data ? 1 : -1;
|
||||
int bsz = x.shape()[0];
|
||||
msg_sed.mtext[1] = bsz;
|
||||
for (int i = 2; i < bsz + 2; i++) {
|
||||
msg_sed.mtext[i] = (int)x_data[i - 2];
|
||||
}
|
||||
if ((msgsnd(msgid, &msg_sed, (MAX_BSZ + 2) * 4, 0)) == -1) {
|
||||
// printf("full msg buffer\n");
|
||||
}
|
||||
|
||||
baidu::xpu::api::plugin::print_times("[TIME END] SaveOutMmsg");
|
||||
return;
|
||||
}
|
||||
|
||||
PD_BUILD_OP(save_output)
|
||||
.Inputs({"x", "not_need_stop"})
|
||||
.Inputs({"x", "not_need_stop", "msg_queue_id"})
|
||||
.Attrs({"rank_id: int64_t"})
|
||||
.Outputs({"x_out"})
|
||||
.SetInplaceMap({{"x", "x_out"}})
|
||||
.SetKernelFn(PD_KERNEL(SaveOutMmsg));
|
||||
.SetKernelFn(PD_KERNEL(SaveOutMmsg));
|
||||
@@ -56,7 +56,7 @@ void SetPreidsTokenPenaltyMultiScores(const paddle::Tensor& pre_ids,
|
||||
seq_lens_decoder.data<int>(),
|
||||
step_idx.data<int64_t>(),
|
||||
bs,
|
||||
length,
|
||||
length_id,
|
||||
length_input_ids);
|
||||
PD_CHECK(r == 0, "xpu::plugin::set_value_by_flags_and_idx failed.");
|
||||
|
||||
|
||||
@@ -29,9 +29,15 @@ PADDLE_PATH = os.path.dirname(paddle.__file__)
|
||||
PADDLE_INCLUDE_PATH = os.path.join(PADDLE_PATH, "include")
|
||||
PADDLE_LIB_PATH = os.path.join(PADDLE_PATH, "libs")
|
||||
|
||||
XPU_LIB = os.getenv("XPU_LIB")
|
||||
BKCL_PATH = os.getenv("BKCL_PATH")
|
||||
if BKCL_PATH is None:
|
||||
BKCL_INC_PATH = os.path.join(PADDLE_INCLUDE_PATH, "xpu")
|
||||
BKCL_LIB_PATH = os.path.join(PADDLE_LIB_PATH, "libbkcl.so")
|
||||
else:
|
||||
BKCL_INC_PATH = os.path.join(BKCL_PATH, "include")
|
||||
BKCL_LIB_PATH = os.path.join(BKCL_PATH, "so", "libbkcl.so")
|
||||
|
||||
XFT_PATH = f"{XPU_LIB}/xft_output"
|
||||
XFT_PATH = os.getenv("XFT_PATH")
|
||||
if XFT_PATH is None:
|
||||
XFT_INC_PATH = os.path.join(PADDLE_INCLUDE_PATH, "xft")
|
||||
XFT_LIB_PATH = os.path.join(PADDLE_LIB_PATH, "libxft.so")
|
||||
@@ -39,18 +45,29 @@ else:
|
||||
XFT_INC_PATH = os.path.join(XFT_PATH, "include")
|
||||
XFT_LIB_PATH = os.path.join(XFT_PATH, "so", "libxft.so")
|
||||
|
||||
XRE_PATH = os.getenv("XRE_PATH")
|
||||
if XRE_PATH is None:
|
||||
XRE_INC_PATH = os.path.join(PADDLE_INCLUDE_PATH, "xre")
|
||||
XRE_LIB_PATH = os.path.join(PADDLE_LIB_PATH, "libxpucuda.so")
|
||||
else:
|
||||
XRE_INC_PATH = os.path.join(XRE_PATH, "include")
|
||||
XRE_LIB_PATH = os.path.join(XRE_PATH, "so", "libxpucuda.so")
|
||||
|
||||
XRE_PATH = f"{XPU_LIB}/xre"
|
||||
XRE_INC_PATH = os.path.join(XRE_PATH, "include")
|
||||
XRE_LIB_PATH = os.path.join(XRE_PATH, "so", "libcudart.so")
|
||||
XFA_PATH = os.getenv("XFA_PATH")
|
||||
if XFA_PATH is None:
|
||||
XFA_INC_PATH = os.path.join(PADDLE_INCLUDE_PATH, "xhpc", "xfa")
|
||||
XFA_LIB_PATH = os.path.join(PADDLE_LIB_PATH, "libxpu_flash_attention.so")
|
||||
else:
|
||||
XFA_INC_PATH = os.path.join(XFA_PATH, "include")
|
||||
XFA_LIB_PATH = os.path.join(XFA_PATH, "so", "libxpu_flash_attention.so")
|
||||
|
||||
XFA_PATH = f"{XPU_LIB}/xhpc/xfa"
|
||||
XFA_INC_PATH = os.path.join(XFA_PATH, "include")
|
||||
XFA_LIB_PATH = os.path.join(XFA_PATH, "so", "libxpu_flash_attention.so")
|
||||
|
||||
XBLAS_PATH = f"{XPU_LIB}/xhpc/xblas"
|
||||
XBLAS_INC_PATH = os.path.join(XBLAS_PATH, "include")
|
||||
XBLAS_LIB_PATH = os.path.join(XBLAS_PATH, "so", "libxpu_blas.so")
|
||||
XBLAS_PATH = os.getenv("XBLAS_PATH")
|
||||
if XBLAS_PATH is None:
|
||||
XBLAS_INC_PATH = os.path.join(PADDLE_INCLUDE_PATH, "xhpc", "xblas")
|
||||
XBLAS_LIB_PATH = os.path.join(PADDLE_LIB_PATH, "libxpu_blas.so")
|
||||
else:
|
||||
XBLAS_INC_PATH = os.path.join(XBLAS_PATH, "include")
|
||||
XBLAS_LIB_PATH = os.path.join(XBLAS_PATH, "so", "libxpu_blas.so")
|
||||
|
||||
setup(
|
||||
name="paddlenlp_ops",
|
||||
@@ -63,24 +80,46 @@ setup(
|
||||
"./set_value_by_flags_and_idx_v2.cc",
|
||||
"./get_token_penalty_multi_scores_v2.cc",
|
||||
"./get_padding_offset_v2.cc",
|
||||
# "./update_inputs.cc",
|
||||
"./rebuild_padding_v2.cc",
|
||||
"../../gpu/save_with_output.cc",
|
||||
"../../gpu/save_with_output_msg.cc",
|
||||
"../../gpu/get_output.cc",
|
||||
"./save_with_output.cc",
|
||||
"./save_with_output_msg.cc",
|
||||
"./get_output.cc",
|
||||
"./moe_dispatch.cc",
|
||||
"./moe_ffn.cc",
|
||||
"./moe_reduce.cc",
|
||||
"./moe_fused.cc",
|
||||
"./mla_block_multihead_attention_xpu.cc",
|
||||
"./weight_only_linear.cc",
|
||||
"./absorb_mla_block_mha_encoder_xpu.cc",
|
||||
"./absorb_mla_block_mha_decoder_xpu.cc",
|
||||
"./get_position_ids.cc",
|
||||
"./get_position_ids_v2.cc",
|
||||
"./adjust_batch.cc",
|
||||
"./gather_next_token.cc",
|
||||
"step.cc",
|
||||
"./step.cc",
|
||||
"./fused_rotary_position_encoding.cc",
|
||||
"./mla_cache_kernel_xpu.cc",
|
||||
"./batch_matmul_xpu.cc",
|
||||
"./get_infer_param.cc",
|
||||
"./block_attn.cc",
|
||||
# "./dynamic_set_split_value.cc"
|
||||
],
|
||||
include_dirs=[
|
||||
".",
|
||||
"./plugin/include",
|
||||
XRE_INC_PATH,
|
||||
XFT_INC_PATH,
|
||||
BKCL_LIB_PATH,
|
||||
XFA_INC_PATH,
|
||||
XBLAS_INC_PATH,
|
||||
],
|
||||
extra_objects=[
|
||||
"./plugin/build/libxpuplugin.a",
|
||||
XRE_LIB_PATH,
|
||||
XFT_LIB_PATH,
|
||||
BKCL_LIB_PATH,
|
||||
XFA_LIB_PATH,
|
||||
XBLAS_LIB_PATH,
|
||||
],
|
||||
include_dirs=[".", "./plugin/include", XRE_INC_PATH, XFT_INC_PATH, XFA_INC_PATH, XBLAS_INC_PATH],
|
||||
extra_objects=["./plugin/build/libxpuplugin.a", XRE_LIB_PATH, XFT_LIB_PATH, XFA_LIB_PATH, XBLAS_LIB_PATH],
|
||||
extra_compile_args={"cxx": ["-D_GLIBCXX_USE_CXX11_ABI=1", "-DPADDLE_WITH_XPU"]},
|
||||
)
|
||||
],
|
||||
|
||||
@@ -41,7 +41,10 @@ void StepPaddle(const paddle::Tensor& stop_flags,
|
||||
const paddle::Tensor& first_token_ids,
|
||||
const int block_size,
|
||||
const int encoder_decoder_block_num) {
|
||||
phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
|
||||
|
||||
baidu::xpu::api::plugin::print_times("[TIME BEGIN] StepPaddle");
|
||||
|
||||
phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
|
||||
auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
|
||||
auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
|
||||
|
||||
@@ -81,7 +84,7 @@ void StepPaddle(const paddle::Tensor& stop_flags,
|
||||
auto recover_lens_cpu = recover_lens.copy_to(paddle::CPUPlace(), false);
|
||||
int recover_lens_cpu_data = recover_lens_cpu.data<int>()[0];
|
||||
if (recover_lens_cpu_data > 0) {
|
||||
r = baidu::xpu::api::plugin::recover_block(
|
||||
r = baidu::xpu::api::plugin::recover_block(
|
||||
xpu_ctx->x_context(),
|
||||
const_cast<int*>(recover_block_list.data<int>()),
|
||||
const_cast<int*>(recover_lens.data<int>()),
|
||||
@@ -105,7 +108,8 @@ void StepPaddle(const paddle::Tensor& stop_flags,
|
||||
length,
|
||||
pre_id_length);
|
||||
PD_CHECK(r == 0, "recover_block failed.");
|
||||
}
|
||||
}
|
||||
baidu::xpu::api::plugin::print_times("[TIME END] StepPaddle");
|
||||
}
|
||||
|
||||
PD_BUILD_OP(step_paddle)
|
||||
|
||||
@@ -26,6 +26,9 @@ void UpdateInputes(const paddle::Tensor& stop_flags,
|
||||
const paddle::Tensor& stop_nums,
|
||||
const paddle::Tensor& next_tokens,
|
||||
const paddle::Tensor& is_block_step) {
|
||||
baidu::xpu::api::plugin::print_times("[TIME BEGIN] UpdateInputes");
|
||||
|
||||
|
||||
phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
|
||||
auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
|
||||
auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
|
||||
@@ -53,7 +56,11 @@ void UpdateInputes(const paddle::Tensor& stop_flags,
|
||||
max_bsz,
|
||||
input_ids_stride);
|
||||
PD_CHECK(r == 0, "baidu::xpu::api::plugin::update_inputs failed.");
|
||||
}
|
||||
|
||||
|
||||
baidu::xpu::api::plugin::print_times("[TIME END] UpdateInputes");
|
||||
}
|
||||
|
||||
|
||||
PD_BUILD_OP(update_inputs)
|
||||
.Inputs({"stop_flags",
|
||||
|
||||
@@ -1,142 +0,0 @@
|
||||
// Copyright (c) 2024 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/backends/xpu/xpu_context.h>
|
||||
#include "paddle/extension.h"
|
||||
#include "paddle/phi/core/enforce.h"
|
||||
#include "xpu/plugin.h"
|
||||
#include <xft/xdnn_plugin.h>
|
||||
#include "xblas_legacy_api.h"
|
||||
|
||||
namespace xftkernel = baidu::xpu::xftkernel;
|
||||
namespace api = baidu::xpu::api;
|
||||
// namespace xblas = baidu::xpu::xblas;
|
||||
|
||||
std::vector<paddle::Tensor> WeightOnlyLinearKernel(
|
||||
const paddle::Tensor& x,
|
||||
const paddle::Tensor& weight,
|
||||
const paddle::optional<paddle::Tensor>& weight_scale,
|
||||
const paddle::optional<paddle::Tensor>& bias,
|
||||
const std::string& weight_dtype="int8",
|
||||
const int32_t arch=0,
|
||||
const int32_t group_size=64
|
||||
) {
|
||||
|
||||
phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());
|
||||
auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(place);
|
||||
auto xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
|
||||
|
||||
using XPUType = typename XPUTypeTrait<bfloat16>::Type;
|
||||
typedef paddle::bfloat16 data_t;
|
||||
|
||||
int64_t n = weight.dims()[0];
|
||||
int64_t k = weight.dims()[1];
|
||||
int64_t m = x.numel() / k;
|
||||
if(weight_dtype == "int4") {
|
||||
n = n * 2;
|
||||
}
|
||||
|
||||
paddle::Tensor bias_fp32;
|
||||
if (bias.is_initialized() && bias.get().dtype() == phi::DataType::FLOAT16) {
|
||||
bias_fp32 = paddle::full(bias->shape(), -1, paddle::DataType::FLOAT32, bias->place());
|
||||
int r = baidu::xpu::api::cast<XPUType, float>(
|
||||
xpu_ctx->x_context(),
|
||||
reinterpret_cast<const XPUType*>(
|
||||
bias.get().data<phi::dtype::float16>()),
|
||||
bias_fp32.data<float>(),
|
||||
n);
|
||||
}
|
||||
|
||||
paddle::Tensor weight_scale_fp32;
|
||||
if (weight_scale.is_initialized() && weight_scale.get().dtype() == phi::DataType::FLOAT16) {
|
||||
weight_scale_fp32 = paddle::full(weight_scale->shape(), -1, paddle::DataType::FLOAT32, weight_scale->place());
|
||||
int r = baidu::xpu::api::cast<XPUType, float>(
|
||||
xpu_ctx->x_context(),
|
||||
reinterpret_cast<const XPUType*>(
|
||||
weight_scale.get().data<phi::dtype::float16>()),
|
||||
weight_scale_fp32.data<float>(),
|
||||
n);
|
||||
}
|
||||
|
||||
auto output = paddle::full({m, n}, -1, x.type(), x.place());
|
||||
|
||||
auto input_x = reinterpret_cast<const XPUType*>(x.data<data_t>());
|
||||
auto input_y = reinterpret_cast<XPUType*>(output.data<data_t>());
|
||||
|
||||
baidu::xpu::xblas::FcFusionTensor<const XPUType>
|
||||
tensor_x{input_x, nullptr, m, k, k, false};
|
||||
baidu::xpu::xblas::FcFusionTensor<const XPUType> tensor_y_const{
|
||||
input_y, nullptr, m, n, n, false};
|
||||
baidu::xpu::xblas::FcFusionTensor<XPUType> tensor_y{
|
||||
input_y, nullptr, m, n, n, false};
|
||||
|
||||
baidu::xpu::xblas::FcFusionEpilogue<float, float> epilogue{
|
||||
api::Activation_t::LINEAR,
|
||||
bias ? bias_fp32.data<float>() : nullptr,
|
||||
nullptr,
|
||||
weight_scale ? weight_scale_fp32.data<float>() : nullptr, // weight_scale_fp32.data<float>(),
|
||||
0,
|
||||
1,
|
||||
nullptr};
|
||||
if(weight_dtype == "int4") {
|
||||
baidu::xpu::xblas::FcFusionDesc<int4_wo_int15, float, XPUType> desc{1.0f, 0.0f};
|
||||
baidu::xpu::xblas::FcFusionTensor<const int4_t> tensor_w{
|
||||
reinterpret_cast<const int4_t*>(weight.data<int8_t>()), nullptr, n, k, k, true};
|
||||
int r1 = baidu::xpu::xblas::fc_fusion<XPUType,
|
||||
int4_t,
|
||||
XPUType,
|
||||
XPUType,
|
||||
int4_wo_int15, // int8_wo_t
|
||||
float,
|
||||
XPUType,
|
||||
float,
|
||||
float>(xpu_ctx->x_context(),
|
||||
tensor_x,
|
||||
tensor_w,
|
||||
tensor_y_const,
|
||||
tensor_y,
|
||||
desc,
|
||||
epilogue);
|
||||
PD_CHECK(r1 == 0, "xblas::fc_fusion failed");
|
||||
} else {
|
||||
baidu::xpu::xblas::FcFusionDesc<float, float, float> desc{1.0f, 0.0f};
|
||||
baidu::xpu::xblas::FcFusionTensor<const XPUType> tensor_w{
|
||||
// reinterpret_cast<const XPUType*>(weight.data<phi::dtype::bfloat16>()), nullptr, k, n, n, true};
|
||||
reinterpret_cast<const XPUType*>(weight.data<phi::dtype::bfloat16>()), nullptr, n, k, k, true};
|
||||
int r1 = baidu::xpu::xblas::fc_fusion<XPUType,
|
||||
bfloat16,
|
||||
XPUType,
|
||||
XPUType,
|
||||
float,
|
||||
float,
|
||||
float,
|
||||
float,
|
||||
float,
|
||||
XPUType>(xpu_ctx->x_context(),
|
||||
tensor_x,
|
||||
tensor_w,
|
||||
tensor_y_const,
|
||||
tensor_y,
|
||||
desc,
|
||||
epilogue);
|
||||
PD_CHECK(r1 == 0, "xblas::fc_fusion failed");
|
||||
}
|
||||
return {output};
|
||||
}
|
||||
|
||||
PD_BUILD_OP(weight_only_linear_kernel)
|
||||
.Inputs({"x", "weight", paddle::Optional("weight_scale"), paddle::Optional("bias")})
|
||||
.Outputs({"output"})
|
||||
.Attrs({"weight_dtype: std::string", "arch: int", "group_size: int"})
|
||||
.SetKernelFn(PD_KERNEL(WeightOnlyLinearKernel));
|
||||
@@ -20,24 +20,20 @@ import paddle
|
||||
np.random.seed(2023)
|
||||
|
||||
|
||||
class GetUpdateInputsTest(unittest.TestCase):
|
||||
def test_update_inputs(self):
|
||||
def test_update_inputs():
|
||||
seq_lens_encoder = paddle.to_tensor(
|
||||
[[0], [0], [0], [0]], dtype="int32", place=paddle.XPUPlace(0), stop_gradient=True
|
||||
)
|
||||
seq_lens_decoder = paddle.to_tensor(
|
||||
[[27], [29], [31], [31]], dtype="int32", place=paddle.XPUPlace(0), stop_gradient=True
|
||||
)
|
||||
batch_size = paddle.to_tensor(
|
||||
[0, 8191, 16382, 24573], dtype="int32", place=paddle.XPUPlace(0), stop_gradient=True
|
||||
)
|
||||
|
||||
seq_lens_encoder = paddle.to_tensor(
|
||||
[[0], [0], [0], [0]], dtype="int32", place=paddle.XPUPlace(0), stop_gradient=True
|
||||
)
|
||||
seq_lens_decoder = paddle.to_tensor(
|
||||
[[27], [29], [31], [31]], dtype="int32", place=paddle.XPUPlace(0), stop_gradient=True
|
||||
)
|
||||
batch_size = paddle.to_tensor(
|
||||
[0, 8191, 16382, 24573], dtype="int32", place=paddle.XPUPlace(0), stop_gradient=True
|
||||
)
|
||||
a, b = paddle.incubate.nn.functional.blha_get_max_len(seq_lens_encoder, seq_lens_decoder, batch_size)
|
||||
|
||||
a, b = paddle.incubate.nn.functional.blha_get_max_len(seq_lens_encoder, seq_lens_decoder, batch_size)
|
||||
print(a)
|
||||
print(b)
|
||||
|
||||
print(a)
|
||||
print(b)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
test_update_inputs()
|
||||
110
csrc/xpu/test/python/test_get_position_ids.py
Normal file
110
csrc/xpu/test/python/test_get_position_ids.py
Normal file
@@ -0,0 +1,110 @@
|
||||
# Copyright (c) 2024 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.
|
||||
|
||||
import numpy as np
|
||||
import paddle
|
||||
import unittest
|
||||
from paddlenlp_ops import update_inputs
|
||||
|
||||
np.random.seed(2023)
|
||||
class GetUpdateInputsTest(unittest.TestCase):
|
||||
def test_update_inputs(self):
|
||||
bs = 48
|
||||
max_bs = 64
|
||||
max_input_length = 6144
|
||||
|
||||
stop_flags = np.random.randint(0, 2, max_bs).astype("bool")
|
||||
not_need_stop = np.array([1], "bool")
|
||||
seq_lens_this_time = np.zeros([bs], "int32")
|
||||
seq_lens_encoder = np.zeros([max_bs], "int32")
|
||||
seq_lens_decoder = np.zeros([max_bs], "int32")
|
||||
for i in range(bs):
|
||||
if i % 2 == 0:
|
||||
seq_lens_encoder[i] = i
|
||||
seq_lens_this_time[i] = i
|
||||
else:
|
||||
seq_lens_decoder[i] = i
|
||||
seq_lens_this_time[i] = 1
|
||||
input_ids_np = np.random.randint(1, 10, [max_bs, max_input_length], "int64")
|
||||
stop_nums = np.array([max_bs], "int64")
|
||||
next_tokens = np.random.randint(1, 10, [max_bs], "int64")
|
||||
is_block_step = np.random.randint(0, 2, [max_bs]).astype("bool")
|
||||
|
||||
stop_flags = paddle.to_tensor(stop_flags)
|
||||
not_need_stop = paddle.to_tensor(not_need_stop, place=paddle.CPUPlace())
|
||||
seq_lens_this_time = paddle.to_tensor(seq_lens_this_time)
|
||||
seq_lens_encoder = paddle.to_tensor(seq_lens_encoder)
|
||||
seq_lens_decoder = paddle.to_tensor(seq_lens_decoder)
|
||||
input_ids = paddle.to_tensor(input_ids_np)
|
||||
stop_nums = paddle.to_tensor(stop_nums)
|
||||
next_tokens = paddle.to_tensor(next_tokens)
|
||||
is_block_step = paddle.to_tensor(is_block_step)
|
||||
|
||||
print("stop_flags:\n", stop_flags)
|
||||
print("not_need_stop:\n", not_need_stop)
|
||||
print("seq_lens_this_time:\n", seq_lens_this_time)
|
||||
print("seq_lens_encoder:\n", seq_lens_encoder)
|
||||
print("seq_lens_decoder:\n", seq_lens_decoder)
|
||||
print("input_ids:\n", input_ids)
|
||||
print("stop_nums:\n", stop_nums)
|
||||
print("next_tokens:\n", next_tokens)
|
||||
print("is_block_step:\n", is_block_step)
|
||||
|
||||
update_inputs(
|
||||
stop_flags,
|
||||
not_need_stop,
|
||||
seq_lens_this_time,
|
||||
seq_lens_encoder,
|
||||
seq_lens_decoder,
|
||||
input_ids,
|
||||
stop_nums,
|
||||
next_tokens,
|
||||
is_block_step
|
||||
)
|
||||
|
||||
print("-" * 50)
|
||||
print("stop_flags:\n", stop_flags)
|
||||
print("not_need_stop:\n", not_need_stop)
|
||||
print("seq_lens_this_time:\n", seq_lens_this_time)
|
||||
print("seq_lens_encoder:\n", seq_lens_encoder)
|
||||
print("seq_lens_decoder:\n", seq_lens_decoder)
|
||||
print("input_ids:\n", input_ids)
|
||||
print("stop_nums:\n", stop_nums)
|
||||
print("next_tokens:\n", next_tokens)
|
||||
|
||||
ref_not_need_stop_out = np.array([True])
|
||||
ref_seq_lens_this_time_out = np.array([0, 0, 1, 0, 0, 1, 0, 1, 1, 1, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1,
|
||||
1, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 0, 1, 1, 1], "int32")
|
||||
ref_seq_lens_encoder_out = np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
|
||||
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
|
||||
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], "int32")
|
||||
ref_seq_lens_decoder_out = np.array([0, 0, 2, 0, 0, 6, 0, 8, 8, 10, 0, 12, 12, 0, 0, 0, 0, 0, 0, 0, 20, 22, 0, 24,
|
||||
24, 0, 26, 28, 0, 0, 0, 32, 32, 0, 34, 0, 0, 38, 0, 40, 0, 0, 42, 0, 0, 46, 46, 48,
|
||||
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], "int32")
|
||||
input_ids_np[:, 0] = np.array([6, 5, 9, 8, 6, 2, 8, 1, 3, 1, 3, 6, 9, 8, 1, 9, 1, 8, 8, 6, 7, 6, 5, 3,
|
||||
5, 9, 3, 6, 3, 9, 8, 8, 8, 8, 4, 8, 7, 4, 2, 3, 5, 8, 4, 2, 5, 6, 8, 9,
|
||||
6, 7, 4, 2, 4, 6, 2, 3, 4, 9, 7, 2, 1, 8, 7, 8], "int64")
|
||||
|
||||
assert not_need_stop.numpy() == ref_not_need_stop_out, 'Check not_need_stop failed.'
|
||||
assert np.all(seq_lens_this_time.numpy()
|
||||
== ref_seq_lens_this_time_out), 'Check seq_lens_this_time failed.'
|
||||
assert np.all(seq_lens_encoder.numpy()
|
||||
== ref_seq_lens_encoder_out), 'Check seq_lens_encoder failed.'
|
||||
assert np.all(seq_lens_decoder.numpy()
|
||||
== ref_seq_lens_decoder_out), 'Check seq_lens_decoder failed.'
|
||||
assert np.all(input_ids.numpy()
|
||||
== input_ids_np), 'Check input_ids failed.'
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -89,6 +89,12 @@ def infer(
|
||||
is_first = False
|
||||
|
||||
total_token_latency = time.time() - start
|
||||
inference_first_token_latency = float(
|
||||
json.loads(chunks[0])["inference_time_cost"]
|
||||
)
|
||||
inference_total_token_latency = float(
|
||||
json.loads(chunks[-1])["inference_time_cost"]
|
||||
)
|
||||
|
||||
if backend == "vllm":
|
||||
try:
|
||||
@@ -108,7 +114,7 @@ def infer(
|
||||
token_num = len(chunks)
|
||||
|
||||
stats.append(
|
||||
[first_token_latency, total_token_latency, input_seqlen, output_seqlen, token_num]
|
||||
[first_token_latency, total_token_latency, input_seqlen, output_seqlen, token_num, inference_first_token_latency, inference_total_token_latency]
|
||||
)
|
||||
print(f"Request queue size: {req_que.qsize()}, Real return tokens: {token_num}, Request Chunks: {len(chunks)}, Label out_seq_len: {output_seqlen}")
|
||||
|
||||
@@ -276,6 +282,16 @@ def save_results_to_file(stats: np.ndarray, output_file: str):
|
||||
"max": stats[:, 4].max(),
|
||||
"avg": stats[:, 4].mean(),
|
||||
},
|
||||
"inference_first_token_latency": {
|
||||
"min": stats[:, 5].min(),
|
||||
"max": stats[:, 5].max(),
|
||||
"avg": stats[:, 5].mean(),
|
||||
},
|
||||
"inference_total_token_latency": {
|
||||
"min": stats[:, 6].min(),
|
||||
"max": stats[:, 6].max(),
|
||||
"avg": stats[:, 6].mean(),
|
||||
},
|
||||
"qps": len(stats) / (stats[:, 1].sum() / len(stats)),
|
||||
"real_output_tokens_per_second": len(stats) / (stats[:, 1].sum() / len(stats)) * stats[:, 4].mean(),
|
||||
}
|
||||
@@ -347,7 +363,7 @@ def main(
|
||||
if len(stats[i + 1]) > 0:
|
||||
all_stat.append(stats[i + 1])
|
||||
|
||||
all_stat = np.concatenate(all_stat).reshape(-1, 5)
|
||||
all_stat = np.concatenate(all_stat).reshape(-1, 7)
|
||||
|
||||
print(f"Processed {len(all_stat)} requests in {elapsed_time:.2f} seconds.")
|
||||
print(
|
||||
@@ -365,6 +381,12 @@ def main(
|
||||
print(
|
||||
f"Real Output Sequence Length (min, max, avg): {all_stat[:, 4].min():.2f}, {all_stat[:, 4].max():.2f}, {all_stat[:, 4].mean():.2f}"
|
||||
)
|
||||
print(
|
||||
f"Inference First Token Latency (min, max, avg): {all_stat[:, 5].min():.2f}, {all_stat[:, 5].max():.2f}, {all_stat[:, 5].mean():.2f}"
|
||||
)
|
||||
print(
|
||||
f"Inference Total Token Latency (min, max, avg): {all_stat[:, 6].min():.2f}, {all_stat[:, 6].max():.2f}, {all_stat[:, 6].mean():.2f}"
|
||||
)
|
||||
print(f"QPS: {len(all_stat) / elapsed_time:.2f}")
|
||||
print(f"Real Output Tokens/s: {len(all_stat) / elapsed_time * all_stat[:, 4].mean():.2f}")
|
||||
|
||||
|
||||
@@ -1676,6 +1676,7 @@ class StaticGraphBlockInferencePredictor(BlockInferencePredictorMixin):
|
||||
xpu_config.l3_autotune_size = 0
|
||||
config.set_xpu_config(xpu_config)
|
||||
config.switch_ir_optim(True)
|
||||
config.delete_pass("fc_xpu_fuse_pass")
|
||||
# config.enable_memory_optim()
|
||||
else:
|
||||
device_id = int(os.environ.get("FLAGS_selected_gpus", 0))
|
||||
|
||||
@@ -105,7 +105,7 @@ LOG_REDIRECT=""
|
||||
if [ "$OUTPUT_LOG_TO_CONSOLE" == "1" ]; then
|
||||
LOG_REDIRECT="> log/console.log 2>&1"
|
||||
fi
|
||||
eval tritonserver --exit-timeout-secs 100000 --cuda-memory-pool-byte-size 0:0 --cuda-memory-pool-byte-size 1:0 \
|
||||
eval nohup tritonserver --exit-timeout-secs 100000 --cuda-memory-pool-byte-size 0:0 --cuda-memory-pool-byte-size 1:0 \
|
||||
--cuda-memory-pool-byte-size 2:0 --cuda-memory-pool-byte-size 3:0 --cuda-memory-pool-byte-size 4:0 \
|
||||
--cuda-memory-pool-byte-size 5:0 --cuda-memory-pool-byte-size 6:0 --cuda-memory-pool-byte-size 7:0 \
|
||||
--pinned-memory-pool-byte-size 0 --model-repository llm_model/ \
|
||||
|
||||
151
llm/server/server/scripts/warmup.py
Normal file
151
llm/server/server/scripts/warmup.py
Normal file
@@ -0,0 +1,151 @@
|
||||
"""
|
||||
warmup client
|
||||
"""
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import queue
|
||||
import time
|
||||
from functools import partial
|
||||
|
||||
import numpy as np
|
||||
import tritonclient.grpc as grpcclient
|
||||
from tritonclient.utils import *
|
||||
|
||||
parse = argparse.ArgumentParser(description='A program to generate config')
|
||||
parse.add_argument('--batch_size',
|
||||
'-bs',
|
||||
type=int,
|
||||
help='batch size',
|
||||
default=2)
|
||||
parse.add_argument('--input_length',
|
||||
type=int,
|
||||
help='input length',
|
||||
default=6144)
|
||||
args = parse.parse_args()
|
||||
batch_size = args.batch_size
|
||||
input_length = args.input_length
|
||||
|
||||
GRPC_PORT = int(os.getenv('GRPC_PORT', default=8811))
|
||||
|
||||
|
||||
class UserData:
|
||||
"""
|
||||
Init UserData.
|
||||
|
||||
Args:
|
||||
None.
|
||||
|
||||
Returns:
|
||||
None.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._completed_requests = queue.Queue()
|
||||
|
||||
|
||||
def callback(user_data, result, error):
|
||||
"""
|
||||
Asynchronous callback function to handle request results.
|
||||
|
||||
Args:
|
||||
user_data (Any): User-defined data to be passed to the callback.
|
||||
result (Any): The result of the request.
|
||||
error (Any): An error that occurred during the request, if any.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
if error:
|
||||
user_data._completed_requests.put(error)
|
||||
else:
|
||||
user_data._completed_requests.put(result)
|
||||
|
||||
|
||||
model_name = "model"
|
||||
inputs = [grpcclient.InferInput("IN", [1], np_to_triton_dtype(np.object_))]
|
||||
outputs = [grpcclient.InferRequestedOutput("OUT")]
|
||||
|
||||
user_data = UserData()
|
||||
|
||||
|
||||
def has_list(obj):
|
||||
"""
|
||||
has_list func
|
||||
|
||||
Args:
|
||||
obj: object.
|
||||
|
||||
Returns:
|
||||
True or False.
|
||||
|
||||
"""
|
||||
if isinstance(obj, list):
|
||||
return True
|
||||
if isinstance(obj, dict):
|
||||
for value in obj.values():
|
||||
if has_list(value):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
with grpcclient.InferenceServerClient(url="localhost:{}".format(GRPC_PORT),
|
||||
verbose=False) as triton_client:
|
||||
# Establish stream
|
||||
triton_client.start_stream(callback=partial(callback, user_data))
|
||||
|
||||
assert input_length > 2, "input tokens num should be larger than 2"
|
||||
input_ids = [855] * input_length
|
||||
input_ids[0] = 128803
|
||||
input_ids[-1] = 201
|
||||
input_ids[-2] = 128798
|
||||
input_ids[-3] = 128804
|
||||
|
||||
data1 = {
|
||||
"text": "<cls>" + "北" + "<sep>",
|
||||
"input_ids": input_ids,
|
||||
"req_id": 0,
|
||||
"seq_len": 1,
|
||||
"min_dec_len": 1,
|
||||
"max_dec_len": 256,
|
||||
"penalty_score": 1.0,
|
||||
"temperature": 0.8,
|
||||
"topp": 0,
|
||||
"frequency_score": 0.1,
|
||||
# "eos_token_ids": [
|
||||
# 2,
|
||||
# ],
|
||||
"presence_score": 0.0
|
||||
}
|
||||
data_list = [data1.copy() for _ in range(batch_size)]
|
||||
for idx, d in enumerate(data_list):
|
||||
d['req_id'] = idx
|
||||
|
||||
# Send specified many requests in parallel
|
||||
for i in range(1):
|
||||
req_dict = json.dumps(data_list)
|
||||
in_data = np.array([req_dict], dtype=np.object_)
|
||||
inputs[0].set_data_from_numpy(in_data)
|
||||
|
||||
triton_client.async_stream_infer(model_name=model_name,
|
||||
inputs=inputs,
|
||||
request_id="{}".format(i),
|
||||
outputs=outputs)
|
||||
# Retrieve results...
|
||||
while True:
|
||||
data_item = user_data._completed_requests.get(timeout=1200)
|
||||
if type(data_item) == InferenceServerException:
|
||||
print('Exception:', 'status', data_item.status(), 'msg',
|
||||
data_item.message())
|
||||
|
||||
else:
|
||||
request_id = data_item.get_response().id
|
||||
results = data_item.as_numpy("OUT")
|
||||
result = results[0]
|
||||
json_res = json.loads(result)
|
||||
if has_list(json_res):
|
||||
data = json_res[0]
|
||||
print(data)
|
||||
if data["is_end"] == 1:
|
||||
break
|
||||
os._exit(0)
|
||||
12
llm/server/server/scripts/warmup.sh
Normal file
12
llm/server/server/scripts/warmup.sh
Normal file
@@ -0,0 +1,12 @@
|
||||
WARMUP_SEQ_LEN=${1-"8192"}
|
||||
WARMUP_BATCH_SIZE=${2-"1"}
|
||||
|
||||
python warmup.py --batch_size=1 --input_length=${WARMUP_SEQ_LEN}
|
||||
sleep 1
|
||||
for ((i=1; i<=WARMUP_BATCH_SIZE; i++)); do
|
||||
curl 127.0.0.1:9965/v1/chat/completions \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{"text": "hello, llm"}' &
|
||||
done
|
||||
sleep 5
|
||||
python warmup.py --batch_size=1 --input_length=${WARMUP_SEQ_LEN}
|
||||
@@ -279,7 +279,7 @@ class ModelRunner:
|
||||
self.share_inputs["input_ids"] = paddle.full(
|
||||
shape=[self.args.max_batch_size, self.args.max_seq_len], fill_value=self.pad_token_id, dtype="int64"
|
||||
)
|
||||
self.share_inputs["msg_queue_id"] = paddle.full(shape=[1], fill_value=1, dtype="int32")
|
||||
self.share_inputs["msg_queue_id"] = paddle.full(shape=[1], fill_value=1, dtype="int32").cpu()
|
||||
self.share_inputs["top_p"] = paddle.full(
|
||||
shape=[self.args.max_batch_size, 1], fill_value=self.top_p, dtype="float32"
|
||||
)
|
||||
@@ -743,10 +743,11 @@ class InferenceEngine(object):
|
||||
config.set_xpu_device_id(device_id)
|
||||
xpu_config = paddle.inference.XpuConfig()
|
||||
xpu_config.device_id = device_id
|
||||
xpu_config.l3_size = 0
|
||||
xpu_config.l3_size = 0
|
||||
xpu_config.l3_autotune_size = 0
|
||||
config.set_xpu_config(xpu_config)
|
||||
config.switch_ir_optim(True)
|
||||
config.delete_pass("fc_xpu_fuse_pass")
|
||||
else:
|
||||
config.enable_use_gpu(100, device_id)
|
||||
|
||||
|
||||
@@ -28,6 +28,7 @@ from paddlenlp.experimental.transformers.fused_transformer_layers import (
|
||||
FusedBlockMultiTransformerFP8DynamicQuant,
|
||||
FusedBlockMultiTransformerWeightOnly,
|
||||
FusedMultiTransformerConfig,
|
||||
FusedMultiTransformerXPU,
|
||||
MLAConfig,
|
||||
MoeConfig,
|
||||
SpeculateConfig,
|
||||
@@ -143,9 +144,8 @@ class DeepseekScalingRotaryEmbedding(nn.Layer):
|
||||
|
||||
|
||||
class DeepseekScalingRotaryEmbeddingXPU(nn.Layer):
|
||||
"""RotaryEmbedding extended with YaRN method.
|
||||
|
||||
Credits to Peng et al. github.com/jquesnelle/yarn
|
||||
"""
|
||||
RotaryEmbedding XPU Implemention. In XPU, cos and sin must be computed in cpu.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -162,9 +162,9 @@ class DeepseekScalingRotaryEmbeddingXPU(nn.Layer):
|
||||
mscale: float = 1,
|
||||
mscale_all_dim: float = 0,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
ori_device = paddle.device.get_device()
|
||||
paddle.device.set_device("cpu")
|
||||
super().__init__()
|
||||
self._dtype = paddle.get_default_dtype()
|
||||
|
||||
self.rotary_dim = rotary_dim
|
||||
@@ -183,13 +183,12 @@ class DeepseekScalingRotaryEmbeddingXPU(nn.Layer):
|
||||
* attn_factor
|
||||
)
|
||||
|
||||
cos_cache, sin_cache = self._compute_cos_sin_cache()
|
||||
|
||||
self.cos_cache: paddle.Tensor
|
||||
self.register_buffer("cos_cache", cos_cache, persistable=True)
|
||||
self.sin_cache: paddle.Tensor
|
||||
self.register_buffer("sin_cache", sin_cache, persistable=True)
|
||||
cache = self._compute_cos_sin_cache()
|
||||
paddle.device.set_device(ori_device)
|
||||
cache = cache.to("xpu")
|
||||
|
||||
self.cos_sin_cache: paddle.Tensor
|
||||
self.register_buffer("cos_sin_cache", cache, persistable=True)
|
||||
|
||||
def _compute_inv_freq(self, scaling_factor: float) -> paddle.Tensor:
|
||||
pos_freqs = self.base ** (paddle.arange(0, self.rotary_dim, 2, dtype=paddle.float32) / self.rotary_dim)
|
||||
@@ -208,13 +207,11 @@ class DeepseekScalingRotaryEmbeddingXPU(nn.Layer):
|
||||
def _compute_cos_sin_cache(self) -> paddle.Tensor:
|
||||
inv_freq = self._compute_inv_freq(self.scaling_factor)
|
||||
t = paddle.arange(self.max_position_embeddings * self.scaling_factor, dtype=paddle.float32)
|
||||
|
||||
freqs = paddle.outer(t, inv_freq)
|
||||
emb = paddle.concat((freqs, freqs), axis=-1)
|
||||
cos = emb.cos() * self.mscale
|
||||
sin = emb.sin() * self.mscale
|
||||
|
||||
return cos.cast(self._dtype), sin.cast(self._dtype)
|
||||
freqs = paddle.einsum("i,j->ij", t, inv_freq)
|
||||
cos = freqs.cos() * self.mscale
|
||||
sin = freqs.sin() * self.mscale
|
||||
cache = paddle.concat((cos, sin), axis=-1)
|
||||
return cache.cast(self._dtype)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
@@ -222,23 +219,13 @@ class DeepseekScalingRotaryEmbeddingXPU(nn.Layer):
|
||||
query: paddle.Tensor,
|
||||
key: paddle.Tensor,
|
||||
) -> Tuple[paddle.Tensor, paddle.Tensor]:
|
||||
cos = self.cos_cache[position_ids].unsqueeze(1)
|
||||
sin = self.sin_cache[position_ids].unsqueeze(1)
|
||||
import os
|
||||
|
||||
def rotate_half(x):
|
||||
"""Rotates half the hidden axiss of the input."""
|
||||
x1 = x[..., : x.shape[-1] // 2]
|
||||
x2 = x[..., x.shape[-1] // 2 :]
|
||||
return paddle.concat([-x2, x1], axis=-1) # shape is the same as x
|
||||
from paddlenlp_ops import fused_rotary_position_encoding
|
||||
|
||||
s, h, d = query.shape
|
||||
query = query.reshape([s, h, d // 2, 2]).transpose([0, 1, 3, 2]).reshape([s, h, d])
|
||||
|
||||
s, h, d = key.shape
|
||||
key = key.reshape([s, h, d // 2, 2]).transpose([0, 1, 3, 2]).reshape([s, h, d])
|
||||
|
||||
query = (query * cos) + (rotate_half(query) * sin)
|
||||
key = (key * cos) + (rotate_half(key) * sin)
|
||||
# In-place operations that update the query and key tensors.
|
||||
os.environ["stride_in_no_check_dy2st_diff"] = "1"
|
||||
fused_rotary_position_encoding(query, key, position_ids, self.cos_sin_cache, self.rotary_dim, False)
|
||||
|
||||
return query, key
|
||||
|
||||
@@ -334,22 +321,13 @@ class DeepseekV2BlockInferenceModel(DeepseekV2PretrainedModel):
|
||||
for k, v in config.rope_scaling.items()
|
||||
if k in ("extrapolation_factor", "attn_factor", "beta_fast", "beta_slow", "mscale", "mscale_all_dim")
|
||||
}
|
||||
if paddle.is_compiled_with_xpu():
|
||||
self.rotary_emb = DeepseekScalingRotaryEmbeddingXPU(
|
||||
config.qk_rope_head_dim,
|
||||
original_max_position,
|
||||
config.rope_theta,
|
||||
scaling_factor,
|
||||
**extra_kwargs,
|
||||
)
|
||||
else:
|
||||
self.rotary_emb = DeepseekScalingRotaryEmbedding(
|
||||
config.qk_rope_head_dim,
|
||||
original_max_position,
|
||||
config.rope_theta,
|
||||
scaling_factor,
|
||||
**extra_kwargs,
|
||||
)
|
||||
self.rotary_emb = DeepseekScalingRotaryEmbedding(
|
||||
config.qk_rope_head_dim,
|
||||
original_max_position,
|
||||
config.rope_theta,
|
||||
scaling_factor,
|
||||
**extra_kwargs,
|
||||
)
|
||||
|
||||
# get ring_id
|
||||
ring_id = -1
|
||||
@@ -1129,20 +1107,6 @@ class DeepseekV2BlockInferenceModel(DeepseekV2PretrainedModel):
|
||||
ffn2_weights.append(ffn2_quanted_weight.view(paddle.uint8))
|
||||
ffn1_scales.append(ffn1_weight_scale)
|
||||
ffn2_scales.append(ffn2_weight_scale)
|
||||
elif self.moe_quant_type in ["weight_only_int8"]:
|
||||
assert paddle.is_compiled_with_xpu()
|
||||
ffn1_quanted_weight, ffn1_weight_scale = weight_quantize(
|
||||
ffn1_weight.cast("float16"), algo=self.moe_quant_type, group_size=-1, arch=70
|
||||
)
|
||||
ffn2_quanted_weight, ffn2_weight_scale = weight_quantize(
|
||||
ffn2_weight.cast("float16"), algo=self.moe_quant_type, group_size=-1, arch=70
|
||||
)
|
||||
ffn1_weight_scale = ffn1_weight_scale.cast("float16")
|
||||
ffn2_weight_scale = ffn2_weight_scale.cast("float16")
|
||||
ffn1_weights.append(ffn1_quanted_weight.reshape([self.transformer_block.config.embed_dim, -1]))
|
||||
ffn2_weights.append(ffn2_quanted_weight.reshape([-1, self.transformer_block.config.embed_dim]))
|
||||
ffn1_scales.append(ffn1_weight_scale)
|
||||
ffn2_scales.append(ffn2_weight_scale)
|
||||
else:
|
||||
ffn1_weights.append(ffn1_weight)
|
||||
ffn2_weights.append(ffn2_weight)
|
||||
@@ -1160,7 +1124,7 @@ class DeepseekV2BlockInferenceModel(DeepseekV2PretrainedModel):
|
||||
).cast("float32")
|
||||
self.transformer_block.e_score_correction_biases[idx].set_value(e_score_correction_bias)
|
||||
|
||||
if self.use_weight_only or self.moe_quant_type in ["weight_only_int8"]:
|
||||
if self.use_weight_only:
|
||||
self.transformer_block.ffn1_weights[idx].set_value(fused_moe_ffn1_weight)
|
||||
self.transformer_block.ffn2_weights[idx].set_value(fused_moe_ffn2_weight)
|
||||
elif "fp8" in self.quant_type:
|
||||
@@ -1177,7 +1141,7 @@ class DeepseekV2BlockInferenceModel(DeepseekV2PretrainedModel):
|
||||
self.transformer_block.ffn2_weights[idx].set_value(fused_moe_ffn2_weight)
|
||||
self.transformer_block.gate_weights[idx].set_value(gate_weight)
|
||||
|
||||
if self.use_weight_only or self.moe_quant_type in ["weight_only_int8"]:
|
||||
if self.use_weight_only:
|
||||
self.transformer_block.ffn1_weights_scale[idx].set_value(fused_moe_ffn1_weight_scale)
|
||||
self.transformer_block.ffn2_weights_scale[idx].set_value(fused_moe_ffn2_weight_scale)
|
||||
elif "fp8" in self.quant_type:
|
||||
@@ -1346,6 +1310,614 @@ class DeepseekV2BlockInferenceModel(DeepseekV2PretrainedModel):
|
||||
)
|
||||
|
||||
|
||||
@register_base_model
|
||||
class DeepseekV2BlockInferenceModelXPU(DeepseekV2BlockInferenceModel):
|
||||
def __init__(self, config: DeepseekV2Config, base_model_prefix: str):
|
||||
# super(DeepseekV2PretrainedModel, self).__init__(config)
|
||||
DeepseekV2PretrainedModel.__init__(self, config)
|
||||
self.base_model_prefix = base_model_prefix
|
||||
|
||||
self.config = config
|
||||
|
||||
self.max_seq_len = config.max_seq_len
|
||||
self.block_size = config.block_size
|
||||
|
||||
self.vocab_size = config.vocab_size
|
||||
self.hidden_size = config.hidden_size
|
||||
self.intermediate_size = config.intermediate_size
|
||||
self.num_attention_heads = config.num_attention_heads
|
||||
self.num_key_value_heads = config.num_key_value_heads
|
||||
self.num_layers = config.num_hidden_layers
|
||||
self.rms_norm_eps = config.rms_norm_eps
|
||||
self.quant_type = config.quant_type
|
||||
self.weight_block_size = config.weight_block_size
|
||||
self.moe_quant_type = config.moe_quant_type
|
||||
self.rope_theta = config.rope_theta
|
||||
self.return_full_hidden_states = config.get("return_full_hidden_states", False)
|
||||
|
||||
self.use_weight_only = False
|
||||
self.weightonly_group_size = -1
|
||||
if self.quant_type == "weight_only_int8":
|
||||
self.use_weight_only = True
|
||||
self.quant_algo = "weight_only_int8"
|
||||
self.weightonly_group_size = config.weightonly_group_size
|
||||
|
||||
if self.use_weight_only:
|
||||
assert (
|
||||
self.quant_type == "weight_only_int8"
|
||||
), f"Expected quant_type equal to 'weight_only_int8', but received {self.quant_type}"
|
||||
|
||||
assert self.config.mla_use_matrix_absorption
|
||||
|
||||
self.first_k_dense_replace = config.first_k_dense_replace
|
||||
self.n_routed_experts = config.n_routed_experts
|
||||
|
||||
if config.tensor_parallel_degree > config.n_routed_experts:
|
||||
raise ValueError(
|
||||
f"Tensor parallel size {config.tensor_parallel_degree} is greater than "
|
||||
f"the number of experts {config.n_routed_experts}."
|
||||
)
|
||||
|
||||
if config.tensor_parallel_degree > 1 and config.vocab_size % config.tensor_parallel_degree == 0:
|
||||
self.embed_tokens = fleet.meta_parallel.VocabParallelEmbedding(
|
||||
self.vocab_size,
|
||||
self.hidden_size,
|
||||
weight_attr=paddle.ParamAttr(initializer=nn.initializer.XavierNormal()),
|
||||
)
|
||||
else:
|
||||
self.embed_tokens = nn.Embedding(
|
||||
self.vocab_size,
|
||||
self.hidden_size,
|
||||
)
|
||||
|
||||
self.norm = DeepseekV2RMSNorm(config)
|
||||
|
||||
scaling_factor = config.rope_scaling.get("factor", 1)
|
||||
original_max_position = config.rope_scaling.get("original_max_position_embeddings", 4096)
|
||||
extra_kwargs = {
|
||||
k: v
|
||||
for k, v in config.rope_scaling.items()
|
||||
if k in ("extrapolation_factor", "attn_factor", "beta_fast", "beta_slow", "mscale", "mscale_all_dim")
|
||||
}
|
||||
|
||||
self.rotary_emb = DeepseekScalingRotaryEmbeddingXPU(
|
||||
config.qk_rope_head_dim,
|
||||
original_max_position,
|
||||
config.rope_theta,
|
||||
scaling_factor,
|
||||
**extra_kwargs,
|
||||
)
|
||||
# get ring_id
|
||||
ring_id = -1
|
||||
try:
|
||||
hcg = fleet.get_hybrid_communicate_group()
|
||||
model_parallel_group = hcg.get_model_parallel_group()
|
||||
ring_id = model_parallel_group.id
|
||||
except:
|
||||
pass
|
||||
|
||||
ln_scale_attrs = [
|
||||
paddle.ParamAttr(name=f"fuse{self.base_model_prefix}.{idx}.ln_scale") for idx in range(self.num_layers)
|
||||
]
|
||||
|
||||
q_a_proj_weight_attrs = None
|
||||
q_a_layernorm_weight_attrs = None
|
||||
q_b_proj_weight_attrs = None
|
||||
q_proj_weight_attrs = None
|
||||
|
||||
if self.config.q_lora_rank is not None:
|
||||
q_a_proj_weight_attrs = [
|
||||
paddle.ParamAttr(
|
||||
name=f"fuse{self.base_model_prefix}.{idx}.q_a_proj_weight",
|
||||
initializer=paddle.nn.initializer.Constant(value=0),
|
||||
)
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
q_a_layernorm_weight_attrs = [
|
||||
paddle.ParamAttr(
|
||||
name=f"fuse{self.base_model_prefix}.{idx}.q_a_layernorm_weight",
|
||||
initializer=paddle.nn.initializer.Constant(value=1.0),
|
||||
)
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
q_b_proj_weight_attrs = [
|
||||
paddle.ParamAttr(
|
||||
name=f"fuse{self.base_model_prefix}.{idx}.q_b_proj_weight",
|
||||
initializer=paddle.nn.initializer.Constant(value=0),
|
||||
)
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
else:
|
||||
q_proj_weight_attrs = [
|
||||
paddle.ParamAttr(
|
||||
name=f"fuse{self.base_model_prefix}.{idx}.q_proj_weight",
|
||||
initializer=paddle.nn.initializer.Constant(value=0),
|
||||
)
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
|
||||
kv_a_proj_with_mqa_weight_attrs = [
|
||||
paddle.ParamAttr(
|
||||
name=f"fuse{self.base_model_prefix}.{idx}.kv_a_proj_with_mqa_weight",
|
||||
initializer=paddle.nn.initializer.Constant(value=0),
|
||||
)
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
kv_a_layernorm_weight_attrs = [
|
||||
paddle.ParamAttr(
|
||||
name=f"fuse{self.base_model_prefix}.{idx}.kv_a_layernorm_weight",
|
||||
initializer=paddle.nn.initializer.Constant(value=1.0),
|
||||
)
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
kv_b_proj_weight_attrs = [
|
||||
paddle.ParamAttr(
|
||||
name=f"fuse{self.base_model_prefix}.{idx}.kv_b_proj_weight",
|
||||
initializer=paddle.nn.initializer.Constant(value=0),
|
||||
)
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
k_b_proj_weight_attrs = None
|
||||
v_b_proj_weight_attrs = None
|
||||
if config.mla_use_matrix_absorption:
|
||||
k_b_proj_weight_attrs = [
|
||||
paddle.ParamAttr(
|
||||
name=f"fuse{self.base_model_prefix}.{idx}.k_b_proj_weight",
|
||||
initializer=paddle.nn.initializer.Constant(value=0),
|
||||
)
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
v_b_proj_weight_attrs = [
|
||||
paddle.ParamAttr(
|
||||
name=f"fuse{self.base_model_prefix}.{idx}.v_b_proj_weight",
|
||||
initializer=paddle.nn.initializer.Constant(value=0),
|
||||
)
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
|
||||
out_proj_weight_attrs = [
|
||||
paddle.ParamAttr(
|
||||
name=f"fuse{self.base_model_prefix}.{idx}.out_proj_weight",
|
||||
initializer=paddle.nn.initializer.Constant(value=0),
|
||||
)
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
ffn_ln_scale_attrs = [
|
||||
paddle.ParamAttr(name=f"fuse{self.base_model_prefix}.{idx}.ffn_ln_scale") for idx in range(self.num_layers)
|
||||
]
|
||||
ffn1_weight_attrs = [
|
||||
paddle.ParamAttr(
|
||||
name=f"fuse{self.base_model_prefix}.{idx}.ffn1_weight",
|
||||
initializer=paddle.nn.initializer.Constant(value=0),
|
||||
)
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
ffn2_weight_attrs = [
|
||||
paddle.ParamAttr(
|
||||
name=f"fuse{self.base_model_prefix}.{idx}.ffn2_weight",
|
||||
initializer=paddle.nn.initializer.Constant(value=0),
|
||||
)
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
gate_weight_attrs = [
|
||||
paddle.ParamAttr(
|
||||
name=f"fuse{self.base_model_prefix}.{idx}.gate_weight",
|
||||
initializer=paddle.nn.initializer.Constant(value=0),
|
||||
)
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
|
||||
e_score_correction_bias_attrs = None
|
||||
if self.base_model_prefix.startswith("deepseek_v3"):
|
||||
e_score_correction_bias_attrs = [
|
||||
paddle.ParamAttr(
|
||||
name=f"fuse{self.base_model_prefix}.{idx}.e_score_correction_bias",
|
||||
initializer=paddle.nn.initializer.Constant(value=0),
|
||||
)
|
||||
if idx >= self.config.first_k_dense_replace
|
||||
else None
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
|
||||
shared_expert_ffn1_weight_attrs = [
|
||||
paddle.ParamAttr(
|
||||
name=f"fuse{self.base_model_prefix}.{idx}.shared_expert_ffn1_weight",
|
||||
initializer=paddle.nn.initializer.Constant(value=0),
|
||||
)
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
shared_expert_ffn2_weight_attrs = [
|
||||
paddle.ParamAttr(
|
||||
name=f"fuse{self.base_model_prefix}.{idx}.shared_expert_ffn2_weight",
|
||||
initializer=paddle.nn.initializer.Constant(value=0),
|
||||
)
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
|
||||
q_proj_weight_scale_attrs = None
|
||||
q_a_proj_weight_scale_attrs = None
|
||||
q_b_proj_weight_scale_attrs = None
|
||||
kv_a_proj_with_mqa_weight_scale_attrs = None
|
||||
kv_b_proj_weight_scale_attrs = None
|
||||
k_b_proj_weight_scale_attrs = None
|
||||
v_b_proj_weight_scale_attrs = None
|
||||
|
||||
out_proj_weight_scale_attrs = None
|
||||
ffn1_weight_scale_attrs = None
|
||||
ffn2_weight_scale_attrs = None
|
||||
shared_expert_ffn1_weight_scale_attrs = None
|
||||
shared_expert_ffn2_weight_scale_attrs = None
|
||||
|
||||
if self.use_weight_only:
|
||||
if self.config.q_lora_rank is not None:
|
||||
q_a_proj_weight_scale_attrs = [
|
||||
paddle.ParamAttr(
|
||||
name=f"fuse{self.base_model_prefix}.{idx}.q_a_proj_weight_scale",
|
||||
)
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
q_b_proj_weight_scale_attrs = [
|
||||
paddle.ParamAttr(
|
||||
name=f"fuse{self.base_model_prefix}.{idx}.q_b_proj_weight_scale",
|
||||
)
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
else:
|
||||
q_proj_weight_scale_attrs = [
|
||||
paddle.ParamAttr(
|
||||
name=f"fuse{self.base_model_prefix}.{idx}.q_proj_weight_scale",
|
||||
)
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
|
||||
if self.config.mla_use_matrix_absorption:
|
||||
kv_a_proj_with_mqa_weight_scale_attrs = [
|
||||
paddle.ParamAttr(
|
||||
name=f"fuse{self.base_model_prefix}.{idx}.kv_a_proj_with_mqa_weight_scale",
|
||||
)
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
kv_b_proj_weight_scale_attrs = [
|
||||
paddle.ParamAttr(
|
||||
name=f"fuse{self.base_model_prefix}.{idx}.kv_b_proj_weight_scale",
|
||||
)
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
|
||||
out_proj_weight_scale_attrs = [
|
||||
paddle.ParamAttr(name=f"fuse{self.base_model_prefix}.{idx}.out_proj_weight_scale")
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
ffn1_weight_scale_attrs = [
|
||||
paddle.ParamAttr(name=f"fuse{self.base_model_prefix}.{idx}.ffn1_weight_scale")
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
ffn2_weight_scale_attrs = [
|
||||
paddle.ParamAttr(name=f"fuse{self.base_model_prefix}.{idx}.ffn2_weight_scale")
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
shared_expert_ffn1_weight_scale_attrs = [
|
||||
paddle.ParamAttr(name=f"fuse{self.base_model_prefix}.{idx}.shared_expert_ffn1_weight_scale")
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
shared_expert_ffn2_weight_scale_attrs = [
|
||||
paddle.ParamAttr(name=f"fuse{self.base_model_prefix}.{idx}.shared_expert_ffn2_weight_scale")
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
|
||||
if self.config.mla_use_matrix_absorption:
|
||||
k_b_proj_weight_scale_attrs = [
|
||||
paddle.ParamAttr(name=f"fuse{self.base_model_prefix}.{idx}.k_b_proj_weight_scale")
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
v_b_proj_weight_scale_attrs = [
|
||||
paddle.ParamAttr(name=f"fuse{self.base_model_prefix}.{idx}.v_b_proj_weight_scale")
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
|
||||
mla_config = MLAConfig(
|
||||
use_matrix_absorption=self.config.mla_use_matrix_absorption,
|
||||
q_lora_rank=self.config.q_lora_rank,
|
||||
kv_lora_rank=self.config.kv_lora_rank,
|
||||
qk_nope_head_dim=self.config.qk_nope_head_dim,
|
||||
qk_rope_head_dim=self.config.qk_rope_head_dim,
|
||||
v_head_dim=self.config.v_head_dim,
|
||||
mscale=yarn_get_mscale(scaling_factor, float(config.rope_scaling.get("mscale_all_dim", 1.0))),
|
||||
q_proj_weight_attrs=q_proj_weight_attrs,
|
||||
q_proj_weight_scale_attrs=q_proj_weight_scale_attrs,
|
||||
q_a_proj_weight_attrs=q_a_proj_weight_attrs,
|
||||
q_a_proj_weight_scale_attrs=q_a_proj_weight_scale_attrs,
|
||||
q_a_layernorm_weight_attrs=q_a_layernorm_weight_attrs,
|
||||
q_b_proj_weight_attrs=q_b_proj_weight_attrs,
|
||||
q_b_proj_weight_scale_attrs=q_b_proj_weight_scale_attrs,
|
||||
kv_a_proj_with_mqa_weight_attrs=kv_a_proj_with_mqa_weight_attrs,
|
||||
kv_a_proj_with_mqa_weight_scale_attrs=kv_a_proj_with_mqa_weight_scale_attrs,
|
||||
kv_a_layernorm_weight_attrs=kv_a_layernorm_weight_attrs,
|
||||
kv_b_proj_weight_attrs=kv_b_proj_weight_attrs,
|
||||
kv_b_proj_weight_scale_attrs=kv_b_proj_weight_scale_attrs,
|
||||
k_b_proj_weight_attrs=k_b_proj_weight_attrs,
|
||||
k_b_proj_weight_scale_attrs=k_b_proj_weight_scale_attrs,
|
||||
v_b_proj_weight_attrs=v_b_proj_weight_attrs,
|
||||
v_b_proj_weight_scale_attrs=v_b_proj_weight_scale_attrs,
|
||||
)
|
||||
|
||||
moe_config = MoeConfig(
|
||||
num_experts=self.n_routed_experts,
|
||||
top_k=self.config.num_experts_per_tok,
|
||||
topk_group=self.config.topk_group,
|
||||
norm_topk_prob=self.config.norm_topk_prob,
|
||||
routed_scaling_factor=self.config.routed_scaling_factor,
|
||||
num_expert_group=self.config.n_group,
|
||||
topk_method=self.config.topk_method,
|
||||
moe_intermediate_size=self.config.moe_intermediate_size,
|
||||
first_k_dense_replace=self.first_k_dense_replace,
|
||||
shared_expert_with_gate=False,
|
||||
shared_expert_intermediate_size=self.config.moe_intermediate_size * self.config.n_shared_experts,
|
||||
shared_expert_ffn1_weight_attrs=shared_expert_ffn1_weight_attrs,
|
||||
shared_expert_ffn1_weight_scale_attrs=shared_expert_ffn1_weight_scale_attrs,
|
||||
shared_expert_ffn2_weight_attrs=shared_expert_ffn2_weight_attrs,
|
||||
shared_expert_ffn2_weight_scale_attrs=shared_expert_ffn2_weight_scale_attrs,
|
||||
)
|
||||
|
||||
speculate_config = SpeculateConfig(
|
||||
speculate_method=config.get("speculate_method", None),
|
||||
speculate_max_draft_token_num=config.get("speculate_max_draft_token_num", 5),
|
||||
return_full_hidden_states=config.get("return_full_hidden_states", False),
|
||||
)
|
||||
|
||||
transformer_config = FusedMultiTransformerConfig(
|
||||
embed_dim=self.hidden_size,
|
||||
num_heads=self.num_attention_heads,
|
||||
kv_num_heads=self.num_key_value_heads,
|
||||
intermediate_size=self.intermediate_size,
|
||||
quant_type=self.quant_type,
|
||||
weight_block_size=self.weight_block_size,
|
||||
moe_quant_type=self.moe_quant_type,
|
||||
weightonly_group_size=self.weightonly_group_size,
|
||||
activation="swiglu",
|
||||
num_layers=config.num_hidden_layers,
|
||||
tp_degree=config.tensor_parallel_degree,
|
||||
ring_id=ring_id,
|
||||
ln_scale_attrs=ln_scale_attrs,
|
||||
linear_weight_attrs=out_proj_weight_attrs,
|
||||
linear_weight_scale_attrs=out_proj_weight_scale_attrs,
|
||||
ffn_ln_scale_attrs=ffn_ln_scale_attrs,
|
||||
gate_weight_attrs=gate_weight_attrs,
|
||||
ffn1_weight_attrs=ffn1_weight_attrs,
|
||||
ffn1_weight_scale_attrs=ffn1_weight_scale_attrs,
|
||||
ffn2_weight_attrs=ffn2_weight_attrs,
|
||||
ffn2_weight_scale_attrs=ffn2_weight_scale_attrs,
|
||||
e_score_correction_bias_attrs=e_score_correction_bias_attrs,
|
||||
epsilon=self.rms_norm_eps,
|
||||
rope_theta=self.rope_theta,
|
||||
rotary_emb=self.rotary_emb,
|
||||
norm_type="rmsnorm",
|
||||
rank_id=config.tensor_parallel_rank,
|
||||
moe_config=moe_config,
|
||||
mla_config=mla_config,
|
||||
append_attn=config.append_attn,
|
||||
speculate_config=speculate_config,
|
||||
)
|
||||
|
||||
self.set_transformer_block(transformer_config)
|
||||
|
||||
@paddle.no_grad()
|
||||
def set_state_dict(self, state_dict):
|
||||
self.transformer_block.init_weight()
|
||||
|
||||
dtype = paddle.get_default_dtype()
|
||||
embed_tokens_weight = paddle.to_tensor(state_dict[f"{self.base_model_prefix}.embed_tokens.weight"]).cast(
|
||||
self.embed_tokens.weight.dtype
|
||||
)
|
||||
norm_weight = paddle.to_tensor(state_dict[f"{self.base_model_prefix}.norm.weight"]).cast(
|
||||
self.norm.weight.dtype
|
||||
)
|
||||
self.embed_tokens.weight.set_value(embed_tokens_weight)
|
||||
self.norm.weight.set_value(norm_weight)
|
||||
|
||||
for idx in range(self.num_layers):
|
||||
logger.info(f"set state for layer {idx}")
|
||||
|
||||
ln_scale = paddle.to_tensor(
|
||||
state_dict[f"{self.base_model_prefix}.layers.{idx}.input_layernorm.weight"]
|
||||
).cast(self.transformer_block.ln_scales[idx].dtype)
|
||||
self.transformer_block.ln_scales[idx].set_value(ln_scale)
|
||||
|
||||
if self.config.q_lora_rank is not None:
|
||||
q_a_proj_weight = paddle.to_tensor(
|
||||
state_dict[f"{self.base_model_prefix}.layers.{idx}.self_attn.q_a_proj.weight"]
|
||||
).cast(dtype)
|
||||
q_a_layernorm_weight = paddle.to_tensor(
|
||||
state_dict[f"{self.base_model_prefix}.layers.{idx}.self_attn.q_a_layernorm.weight"]
|
||||
).cast(self.transformer_block.q_a_layernorm_weights[idx].dtype)
|
||||
q_b_proj_weight = paddle.to_tensor(
|
||||
state_dict[f"{self.base_model_prefix}.layers.{idx}.self_attn.q_b_proj.weight"]
|
||||
).cast(dtype)
|
||||
|
||||
self.transformer_block.q_a_proj_weights[idx].set_value(q_a_proj_weight)
|
||||
self.transformer_block.q_a_layernorm_weights[idx].set_value(q_a_layernorm_weight)
|
||||
self.transformer_block.q_b_proj_weights[idx].set_value(q_b_proj_weight)
|
||||
else:
|
||||
q_proj_weight = paddle.to_tensor(
|
||||
state_dict[f"{self.base_model_prefix}.layers.{idx}.self_attn.q_proj.weight"]
|
||||
).cast(dtype)
|
||||
|
||||
self.transformer_block.q_proj_weights[idx].set_value(q_proj_weight)
|
||||
|
||||
kv_a_proj_with_mqa_weight = paddle.to_tensor(
|
||||
state_dict[f"{self.base_model_prefix}.layers.{idx}.self_attn.kv_a_proj_with_mqa.weight"]
|
||||
).cast(dtype)
|
||||
kv_a_layernorm_weight = paddle.to_tensor(
|
||||
state_dict[f"{self.base_model_prefix}.layers.{idx}.self_attn.kv_a_layernorm.weight"]
|
||||
).cast(self.transformer_block.kv_a_layernorm_weights[idx].dtype)
|
||||
kv_b_proj_weight = paddle.to_tensor(
|
||||
state_dict[f"{self.base_model_prefix}.layers.{idx}.self_attn.kv_b_proj.weight"]
|
||||
).cast(dtype)
|
||||
linear_weight = paddle.to_tensor(
|
||||
state_dict[f"{self.base_model_prefix}.layers.{idx}.self_attn.o_proj.weight"]
|
||||
).cast(dtype)
|
||||
|
||||
if self.config.mla_use_matrix_absorption:
|
||||
# if self.config.q_lora_rank is None:
|
||||
# q_proj_weight_inner = q_proj_weight.reshape(
|
||||
# shape=[
|
||||
# -1,
|
||||
# self.num_attention_heads // self.config.tensor_parallel_degree,
|
||||
# self.config.qk_nope_head_dim + self.config.qk_rope_head_dim,
|
||||
# ]
|
||||
# )
|
||||
# else:
|
||||
# q_proj_weight_inner = q_b_proj_weight.reshape(
|
||||
# shape=[
|
||||
# -1,
|
||||
# self.num_attention_heads // self.config.tensor_parallel_degree,
|
||||
# self.config.qk_nope_head_dim + self.config.qk_rope_head_dim,
|
||||
# ]
|
||||
# )
|
||||
|
||||
kv_b_proj_weight_inner = kv_b_proj_weight.reshape(
|
||||
shape=[
|
||||
self.config.kv_lora_rank,
|
||||
self.num_attention_heads // self.config.tensor_parallel_degree,
|
||||
-1,
|
||||
]
|
||||
)
|
||||
# linear_weight_inner = linear_weight.T.reshape(
|
||||
# shape=[
|
||||
# -1,
|
||||
# self.num_attention_heads // self.config.tensor_parallel_degree,
|
||||
# self.config.v_head_dim,
|
||||
# ]
|
||||
# )
|
||||
|
||||
# W_Q = q_proj_weight_inner[..., : self.config.qk_nope_head_dim]
|
||||
# W_QR = q_proj_weight_inner[..., self.config.qk_nope_head_dim :].flatten(start_axis=1)
|
||||
W_UK, W_UV = kv_b_proj_weight_inner.split(
|
||||
[self.config.qk_nope_head_dim, self.config.v_head_dim], axis=-1
|
||||
)
|
||||
# W_O = linear_weight_inner
|
||||
|
||||
if self.use_weight_only:
|
||||
assert False
|
||||
pass
|
||||
else:
|
||||
self.transformer_block.k_b_proj_weights[idx].set_value(W_UK.transpose([1, 2, 0]))
|
||||
self.transformer_block.v_b_proj_weights[idx].set_value(W_UV.transpose([1, 0, 2]))
|
||||
|
||||
self.transformer_block.kv_a_proj_with_mqa_weights[idx].set_value(kv_a_proj_with_mqa_weight)
|
||||
self.transformer_block.kv_a_layernorm_weights[idx].set_value(kv_a_layernorm_weight)
|
||||
self.transformer_block.kv_b_proj_weights[idx].set_value(kv_b_proj_weight)
|
||||
|
||||
self.transformer_block.linear_weights[idx].set_value(linear_weight)
|
||||
|
||||
ffn_ln_scale = paddle.to_tensor(
|
||||
state_dict[f"{self.base_model_prefix}.layers.{idx}.post_attention_layernorm.weight"],
|
||||
).cast(
|
||||
self.transformer_block.ffn_ln_scales[idx].dtype,
|
||||
)
|
||||
self.transformer_block.ffn_ln_scales[idx].set_value(ffn_ln_scale)
|
||||
if idx < self.first_k_dense_replace:
|
||||
concated_ffn1_weight = np.concatenate(
|
||||
[
|
||||
state_dict[f"{self.base_model_prefix}.layers.{idx}.mlp.gate_proj.weight"],
|
||||
state_dict[f"{self.base_model_prefix}.layers.{idx}.mlp.up_proj.weight"],
|
||||
],
|
||||
axis=-1,
|
||||
)
|
||||
ffn1_weight_tensor = paddle.to_tensor(concated_ffn1_weight).cast(paddle.get_default_dtype())
|
||||
|
||||
self.transformer_block.ffn1_weights[idx].set_value(ffn1_weight_tensor)
|
||||
|
||||
ffn2_weight_tensor = paddle.to_tensor(
|
||||
state_dict[f"{self.base_model_prefix}.layers.{idx}.mlp.down_proj.weight"]
|
||||
).cast(paddle.get_default_dtype())
|
||||
|
||||
self.transformer_block.ffn2_weights[idx].set_value(ffn2_weight_tensor)
|
||||
else:
|
||||
ffn1_weights = []
|
||||
ffn2_weights = []
|
||||
ffn1_scales = []
|
||||
ffn2_scales = []
|
||||
|
||||
for expert_idx in range(self.n_routed_experts):
|
||||
concated_gate_up_weight = np.concatenate(
|
||||
[
|
||||
state_dict[
|
||||
f"{self.base_model_prefix}.layers.{idx}.mlp.experts.{expert_idx}.gate_proj.weight"
|
||||
],
|
||||
state_dict[
|
||||
f"{self.base_model_prefix}.layers.{idx}.mlp.experts.{expert_idx}.up_proj.weight"
|
||||
],
|
||||
],
|
||||
axis=-1,
|
||||
)
|
||||
ffn1_weight = paddle.to_tensor(concated_gate_up_weight).cast(dtype)
|
||||
ffn2_weight = paddle.to_tensor(
|
||||
state_dict[f"{self.base_model_prefix}.layers.{idx}.mlp.experts.{expert_idx}.down_proj.weight"]
|
||||
).cast(dtype)
|
||||
|
||||
# quant moe
|
||||
ffn1_quanted_weight, ffn1_weight_scale = weight_quantize(
|
||||
ffn1_weight, algo=self.moe_quant_type, group_size=-1, arch=70
|
||||
)
|
||||
ffn2_quanted_weight, ffn2_weight_scale = weight_quantize(
|
||||
ffn2_weight, algo=self.moe_quant_type, group_size=-1, arch=70
|
||||
)
|
||||
ffn1_weight_scale = ffn1_weight_scale
|
||||
ffn2_weight_scale = ffn2_weight_scale
|
||||
ffn1_weights.append(
|
||||
ffn1_quanted_weight.transpose((1, 0)).reshape([self.transformer_block.config.embed_dim, -1])
|
||||
)
|
||||
ffn2_weights.append(
|
||||
ffn2_quanted_weight.transpose((1, 0)).reshape([-1, self.transformer_block.config.embed_dim])
|
||||
)
|
||||
ffn1_scales.append(ffn1_weight_scale)
|
||||
ffn2_scales.append(ffn2_weight_scale)
|
||||
|
||||
fused_moe_ffn1_weight = paddle.to_tensor(ffn1_weights)
|
||||
fused_moe_ffn2_weight = paddle.to_tensor(ffn2_weights)
|
||||
|
||||
# 这里的 paddle.to_tensor，默认转为 bf16，而不是 ffn1_scales.dtype
|
||||
fused_moe_ffn1_weight_scale = paddle.to_tensor(ffn1_scales).cast("float32")
|
||||
fused_moe_ffn2_weight_scale = paddle.to_tensor(ffn2_scales).cast("float32")
|
||||
gate_weight = paddle.to_tensor(
|
||||
state_dict[f"{self.base_model_prefix}.layers.{idx}.mlp.gate.weight"]
|
||||
).cast("float32")
|
||||
|
||||
if self.base_model_prefix.startswith("deepseek_v3"):
|
||||
e_score_correction_bias = paddle.to_tensor(
|
||||
state_dict[f"{self.base_model_prefix}.layers.{idx}.mlp.gate.e_score_correction_bias"]
|
||||
).cast("float32")
|
||||
self.transformer_block.e_score_correction_biases[idx].set_value(e_score_correction_bias)
|
||||
|
||||
self.transformer_block.ffn1_weights[idx].set_value(fused_moe_ffn1_weight)
|
||||
self.transformer_block.ffn2_weights[idx].set_value(fused_moe_ffn2_weight)
|
||||
|
||||
self.transformer_block.gate_weights[idx].set_value(gate_weight)
|
||||
|
||||
self.transformer_block.ffn1_weights_scale[idx].set_value(fused_moe_ffn1_weight_scale)
|
||||
self.transformer_block.ffn2_weights_scale[idx].set_value(fused_moe_ffn2_weight_scale)
|
||||
|
||||
concated_gate_up_weight = np.concatenate(
|
||||
[
|
||||
state_dict[f"{self.base_model_prefix}.layers.{idx}.mlp.shared_experts.gate_proj.weight"],
|
||||
state_dict[f"{self.base_model_prefix}.layers.{idx}.mlp.shared_experts.up_proj.weight"],
|
||||
],
|
||||
axis=-1,
|
||||
)
|
||||
shared_expert_ffn1_weight = paddle.to_tensor(concated_gate_up_weight).cast(dtype)
|
||||
shared_expert_ffn2_weight = paddle.to_tensor(
|
||||
state_dict[f"{self.base_model_prefix}.layers.{idx}.mlp.shared_experts.down_proj.weight"]
|
||||
).cast(dtype)
|
||||
|
||||
self.transformer_block.shared_expert_ffn1_weights[idx].set_value(shared_expert_ffn1_weight)
|
||||
self.transformer_block.shared_expert_ffn2_weights[idx].set_value(shared_expert_ffn2_weight)
|
||||
|
||||
def set_transformer_block(self, transformer_config):
|
||||
assert paddle.is_compiled_with_xpu()
|
||||
self.transformer_block = FusedMultiTransformerXPU(transformer_config)
|
||||
|
||||
|
||||
@register_base_model
|
||||
class MTPDeepseekV2BlockInferenceModel(DeepseekV2BlockInferenceModel):
|
||||
def __init__(self, config: DeepseekV2Config, base_model_prefix: str):
|
||||
@@ -1432,7 +2004,10 @@ class DeepseekV2ForCausalLMBlockInferenceModel(GenerationBlockInferenceModel, De
|
||||
self.max_seq_len = config.max_seq_len
|
||||
self.return_full_hidden_states = config.get("return_full_hidden_states", False)
|
||||
|
||||
self.deepseek_v2 = DeepseekV2BlockInferenceModel(config, base_model_prefix)
|
||||
if paddle.is_compiled_with_xpu():
|
||||
self.deepseek_v2 = DeepseekV2BlockInferenceModelXPU(config, base_model_prefix)
|
||||
else:
|
||||
self.deepseek_v2 = DeepseekV2BlockInferenceModel(config, base_model_prefix)
|
||||
if config.tie_word_embeddings:
|
||||
self.lm_head = DeepseekV2LMHead(
|
||||
config, embedding_weights=self.deepseek_v2.embed_tokens.weight, transpose_y=True
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user