apps/FastDeploy
1
0
Fork 0
mirror of https://github.com/PaddlePaddle/FastDeploy.git synced 2025-12-24 13:28:13 +08:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

[Feature][MTP] Support cacheKV transfer in per_chunk mode (#2890)

Browse Source 
* support chunk_prefill both normal and speculative_decoding(mtp)

* optimize pd-disaggregation config

* fix bug
This commit is contained in:
freeliuzc
2025-07-17 17:58:08 +08:00
committed by GitHub
parent 67180c1ff9
commit d49f8fb30a
10 changed files with 110 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
custom_ops/gpu_ops/get_output_ep.cc
View File

@@ -36,9 +36,9 @@ void GetOutputKVSignal(const paddle::Tensor& x,
int* out_data = const_cast<int*>(x.data<int>());
int ret = -1;
if (!wait_flag) {
ret = msgrcv(msgid, &msg_rcv, (MAX_BSZ * 2 + 2) * 4, 0, IPC_NOWAIT);
ret = msgrcv(msgid, &msg_rcv, (MAX_BSZ * 3 + 2) * 4, 0, IPC_NOWAIT);
} else {
ret = msgrcv(msgid, &msg_rcv, (MAX_BSZ * 2 + 2) * 4, 0, 0);
ret = msgrcv(msgid, &msg_rcv, (MAX_BSZ * 3 + 2) * 4, 0, 0);
}
if (ret == -1) {
out_data[0] = -1;
@@ -47,7 +47,7 @@ void GetOutputKVSignal(const paddle::Tensor& x,
}
int encoder_count = msg_rcv.mtext[0];
for (int i = 0; i < encoder_count * 2 + 2; i++) {
for (int i = 0; i < encoder_count * 3 + 2; i++) {
out_data[i] = msg_rcv.mtext[i];
}
return;

2
custom_ops/gpu_ops/msg_utils.h
View File

@@ -35,5 +35,5 @@ struct msgdata {
struct msgdatakv {
long mtype;
int mtext[MAX_BSZ * 2 + 2]; // encoder_count, layer_id, bid- pair
int mtext[MAX_BSZ * 3 + 2]; // encoder_count, layer_id, bid- pair
};

7
custom_ops/gpu_ops/remote_cache_kv_ipc.h
View File

@@ -64,9 +64,10 @@ struct RemoteCacheKvIpc {
int encoder_count = 0;
for (int i = 0; i < real_bsz; i++) {
if (seq_lens_encoder[i] > 0) {
msg_sed.mtext[3 * encoder_count + 2] = i;
msg_sed.mtext[3 * encoder_count + 3] = seq_lens_decoder[i];
msg_sed.mtext[3 * encoder_count + 4] = seq_lens_encoder[i];
encoder_count++;
msg_sed.mtext[2 * i + 2] = i;
msg_sed.mtext[2 * i + 3] = seq_lens_decoder[i];
}
}
msg_sed.mtext[0] = encoder_count;
@@ -82,7 +83,7 @@ struct RemoteCacheKvIpc {
void CUDART_CB send_signal() {
msg_sed.mtext[1] = layer_id_;
if ((msgsnd(msgid, &msg_sed, (MAX_BSZ * 2 + 2) * 4, 0)) == -1) {
if ((msgsnd(msgid, &msg_sed, (MAX_BSZ * 3 + 2) * 4, 0)) == -1) {
printf("kv signal full msg buffer\n");
}
layer_id_ = (layer_id_ + 1);