Skip to content

Commit 9d3551c

Browse files
RuohengMaiosmers
RuohengMa
and
iosmers
authored
[XPU] add support for rope3d (#7518)
* [XPU] add support for rope3d * support decoder --------- Co-authored-by: yinwei <yinwei_hust@163.com>
1 parent 609f649 commit 9d3551c

1 file changed

Lines changed: 90 additions & 17 deletions

File tree

‎custom_ops/xpu_ops/src/ops/block_attn_spliced.cc‎

Lines changed: 90 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -239,11 +239,7 @@ void split_kvcache_encoder(api::Context* xpu_ctx,
239239
bool use_neox_rotary_style) {
240240
int ret;
241241
int64_t real_kv_num_heads = (kv_num_heads == -1) ? q_num_heads : kv_num_heads;
242-
// TODO: spliced split kvcache should support rope3d
243-
if (FLAGS_encoder_splice && !rope_3d) {
244-
if (rope_3d) {
245-
PD_THROW("split_kvcache_encoder does not support rope_3d == true!");
246-
}
242+
if (FLAGS_encoder_splice) {
247243
paddle::Place place = qkv.place();
248244
xftblock::DataType KV_BUF_TYPE = std::is_same<bfloat16, TQKV>::value
249245
? xftblock::DataType::DT_BFLOAT16
@@ -274,6 +270,47 @@ void split_kvcache_encoder(api::Context* xpu_ctx,
274270
head_dim);
275271
PD_CHECK(ret == api::SUCCESS, "split_qkv_block failed.");
276272

273+
paddle::Tensor pos_emb_offset_xpu;
274+
paddle::Tensor pos_emb_offset_cpu;
275+
api::VectorParam<int32_t> pos_emb_offset;
276+
pos_emb_offset.len = start_tokens.len;
277+
if (rope_3d) {
278+
pos_emb_offset_xpu =
279+
paddle::empty({batch_size}, paddle::DataType::INT32, place);
280+
pos_emb_offset_cpu = paddle::empty(
281+
{batch_size}, paddle::DataType::INT32, paddle::CPUPlace());
282+
auto pos_emb_offset_xpu_ptr =
283+
const_cast<int32_t*>(pos_emb_offset_xpu.data<int32_t>());
284+
auto pos_emb_offset_cpu_ptr =
285+
const_cast<int32_t*>(pos_emb_offset_cpu.data<int32_t>());
286+
// bs_offset = real_batch.cpu[bs] * 2 * rope_max_seqlen +
287+
// start_tokens.cpu[bs];
288+
ret = api::scale(xpu_ctx,
289+
real_batch.xpu,
290+
pos_emb_offset_xpu_ptr,
291+
batch_size,
292+
true,
293+
2 * static_cast<int32_t>(rope_max_seqlen),
294+
0);
295+
PD_CHECK(ret == api::SUCCESS, "api::scale failed.");
296+
ret = api::broadcast_add<int32_t>(xpu_ctx,
297+
pos_emb_offset_xpu_ptr,
298+
start_tokens.xpu,
299+
pos_emb_offset_xpu_ptr,
300+
{pos_emb_offset.len},
301+
{pos_emb_offset.len});
302+
PD_CHECK(ret == api::SUCCESS, "api::broadcast_add failed.");
303+
for (int i = 0; i < batch_size; i++) {
304+
pos_emb_offset_cpu_ptr[i] =
305+
real_batch.cpu[i] * 2 * static_cast<int32_t>(rope_max_seqlen) +
306+
start_tokens.cpu[i];
307+
}
308+
pos_emb_offset.cpu = pos_emb_offset_cpu_ptr;
309+
pos_emb_offset.xpu = pos_emb_offset_xpu_ptr;
310+
} else {
311+
pos_emb_offset.cpu = start_tokens.cpu;
312+
pos_emb_offset.xpu = start_tokens.xpu;
313+
}
277314
if (!use_neox_rotary_style) {
278315
ret = infer_ops::vsl_rotary_embedding_gptj<TQKV, TR, TID>(
279316
xpu_ctx,
@@ -288,8 +325,8 @@ void split_kvcache_encoder(api::Context* xpu_ctx,
288325
q_num_heads,
289326
head_dim,
290327
"BLHD",
291-
start_tokens,
292-
"NORMAL",
328+
pos_emb_offset,
329+
pos_emb_type,
293330
real_kv_num_heads,
294331
false);
295332
PD_CHECK(ret == api::SUCCESS, "vsl_rotary_embedding_gptj failed.");
@@ -308,7 +345,7 @@ void split_kvcache_encoder(api::Context* xpu_ctx,
308345
head_dim,
309346
rope_head_dim,
310347
"BLHD",
311-
start_tokens,
348+
pos_emb_offset,
312349
"NORMAL",
313350
real_kv_num_heads,
314351
false);
@@ -546,12 +583,7 @@ void split_kvcache_decoder(api::Context* xpu_ctx,
546583
bool use_neox_rotary_style) {
547584
int64_t real_kv_num_heads = (kv_num_heads == -1) ? q_num_heads : kv_num_heads;
548585
int ret;
549-
// TODO: spliced split kvcache should support rope3d
550-
if (FLAGS_decoder_splice && !rope_3d) {
551-
// not yet supported
552-
if (rope_3d) {
553-
PD_THROW("split_kvcache_decoder does not support rope_3d == true!");
554-
}
586+
if (FLAGS_decoder_splice) {
555587
if (std::is_same<TKV_CACHE, int8_t>::value &&
556588
(k_cache_scale_inv == nullptr || v_cache_scale_inv == nullptr)) {
557589
PD_THROW(
@@ -591,6 +623,47 @@ void split_kvcache_decoder(api::Context* xpu_ctx,
591623
head_dim);
592624
PD_CHECK(ret == api::SUCCESS, "split_qkv_block failed.");
593625

626+
paddle::Tensor pos_emb_offset_xpu;
627+
paddle::Tensor pos_emb_offset_cpu;
628+
api::VectorParam<int32_t> pos_emb_offset;
629+
pos_emb_offset.len = start_tokens.len;
630+
if (rope_3d) {
631+
pos_emb_offset_xpu =
632+
paddle::empty({batch_size}, paddle::DataType::INT32, place);
633+
pos_emb_offset_cpu = paddle::empty(
634+
{batch_size}, paddle::DataType::INT32, paddle::CPUPlace());
635+
auto pos_emb_offset_xpu_ptr =
636+
const_cast<int32_t*>(pos_emb_offset_xpu.data<int32_t>());
637+
auto pos_emb_offset_cpu_ptr =
638+
const_cast<int32_t*>(pos_emb_offset_cpu.data<int32_t>());
639+
// bs_offset = real_batch.cpu[bs] * 2 * rope_max_seqlen +
640+
// start_tokens.cpu[bs];
641+
ret = api::scale(xpu_ctx,
642+
real_batch.xpu,
643+
pos_emb_offset_xpu_ptr,
644+
batch_size,
645+
true,
646+
2 * static_cast<int32_t>(rope_max_seqlen),
647+
0);
648+
PD_CHECK(ret == api::SUCCESS, "api::scale failed.");
649+
ret = api::broadcast_add<int32_t>(xpu_ctx,
650+
pos_emb_offset_xpu_ptr,
651+
start_tokens.xpu,
652+
pos_emb_offset_xpu_ptr,
653+
{pos_emb_offset.len},
654+
{pos_emb_offset.len});
655+
PD_CHECK(ret == api::SUCCESS, "api::broadcast_add failed.");
656+
for (int i = 0; i < batch_size; i++) {
657+
pos_emb_offset_cpu_ptr[i] =
658+
real_batch.cpu[i] * 2 * static_cast<int32_t>(rope_max_seqlen) +
659+
start_tokens.cpu[i];
660+
}
661+
pos_emb_offset.cpu = pos_emb_offset_cpu_ptr;
662+
pos_emb_offset.xpu = pos_emb_offset_xpu_ptr;
663+
} else {
664+
pos_emb_offset.cpu = start_tokens.cpu;
665+
pos_emb_offset.xpu = start_tokens.xpu;
666+
}
594667
if (!use_neox_rotary_style) {
595668
ret = infer_ops::vsl_rotary_embedding_gptj<TQKV, TR, TID>(
596669
xpu_ctx,
@@ -605,8 +678,8 @@ void split_kvcache_decoder(api::Context* xpu_ctx,
605678
q_num_heads,
606679
head_dim,
607680
"BLHD",
608-
start_tokens,
609-
"NORMAL",
681+
pos_emb_offset,
682+
pos_emb_type,
610683
real_kv_num_heads,
611684
false);
612685
PD_CHECK(ret == api::SUCCESS, "vsl_rotary_embedding_gptj failed.");
@@ -625,7 +698,7 @@ void split_kvcache_decoder(api::Context* xpu_ctx,
625698
head_dim,
626699
rope_head_dim,
627700
"BLHD",
628-
start_tokens,
701+
pos_emb_offset,
629702
"NORMAL",
630703
real_kv_num_heads,
631704
false);

0 commit comments

Comments
 (0)