Skip to content

Commit 31b12ee

Browse files
Sunny-bot1Sunny-bot1
authored
[Cherry-Pick][Optimization] Reduce logprob processing overhead by using actual topk instead of fixed K+1 (#7860) (#7861)
* opt logprob process * fix test_token_processor * fix xpu
1 parent 8c4f5a6 commit 31b12ee

6 files changed

Lines changed: 62 additions & 37 deletions

File tree

‎custom_ops/gpu_ops/get_output_msg_with_topk.cc‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -88,13 +88,17 @@ void GetOutputTopK(const paddle::Tensor& x,
8888
return;
8989
}
9090

91-
int bsz = msg_rcv.mtext[1];
91+
// Unpack bsz (low 16 bits) and actual_topk (high 16 bits) from mtext[1].
92+
// This matches the packing in save_output_msg_with_topk.cc:
93+
// mtext[1] = bsz | (max_num_logprobs << 16)
94+
int bsz = msg_rcv.mtext[1] & 0xFFFF;
95+
int actual_topk = (msg_rcv.mtext[1] >> 16) & 0xFFFF;
9296
out_data[0] = (int64_t)msg_rcv.mtext[0];
93-
out_data[1] = (int64_t)msg_rcv.mtext[1];
97+
out_data[1] = (int64_t)msg_rcv.mtext[1]; // keep packed value; Python unpacks
9498

9599
for (int i = 0; i < bsz; i++) {
96-
for (int j = 0; j < k + 1; j++) {
97-
const int64_t offset = i * (K + 1) + j;
100+
for (int j = 0; j < actual_topk; j++) {
101+
const int64_t offset = i * actual_topk + j;
98102
out_data[offset + 2] = (int64_t)msg_rcv.mtext[offset + 2];
99103
scores_data[offset] = msg_rcv.mtext_f[offset];
100104
}

‎custom_ops/gpu_ops/save_output_msg_with_topk.cc‎

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -109,20 +109,21 @@ void SaveOutMmsgTopK(const paddle::Tensor& x,
109109
: -inference_msg_id_from_env;
110110
int bsz = x.shape()[0];
111111
int max_num_logprobs = logprob_token_ids.shape()[1];
112-
msg_sed.mtext[1] = bsz;
112+
// Pack bsz (low 16 bits) and max_num_logprobs (high 16 bits) into mtext[1].
113+
// token_processor unpacks both fields to avoid reading unused topk slots.
114+
msg_sed.mtext[1] = bsz | (max_num_logprobs << 16);
113115
for (int i = 0; i < bsz; i++) {
114-
for (int j = 0; j < K + 1; j++) {
115-
const int64_t offset = i * (K + 1) + j;
116+
// Loop only over actual logprob columns (max_num_logprobs) instead of the
117+
// fixed K+1=21, and use max_num_logprobs as the stride so data is packed
118+
// densely in the message buffer.
119+
for (int j = 0; j < max_num_logprobs; j++) {
120+
const int64_t offset = i * max_num_logprobs + j;
116121
if (j == 0) {
117122
msg_sed.mtext[offset + 2] = (int)x_data[i];
118-
msg_sed.mtext_f[offset] = logprob_scores_data[i * max_num_logprobs + j];
119-
} else if (j < max_num_logprobs) {
120-
msg_sed.mtext[offset + 2] =
121-
(int)logprob_token_ids_data[i * max_num_logprobs + j];
122-
msg_sed.mtext_f[offset] = logprob_scores_data[i * max_num_logprobs + j];
123+
msg_sed.mtext_f[offset] = logprob_scores_data[offset];
123124
} else {
124-
msg_sed.mtext[offset + 2] = -1;
125-
msg_sed.mtext_f[offset] = 0.0;
125+
msg_sed.mtext[offset + 2] = (int)logprob_token_ids_data[offset];
126+
msg_sed.mtext_f[offset] = logprob_scores_data[offset];
126127
}
127128
if (preempted_idx_data[i] == 1) {
128129
msg_sed.mtext[offset + 2] = -9;

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

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -82,13 +82,17 @@ void GetOutputTopK(const paddle::Tensor& x,
8282
return;
8383
}
8484

85-
int bsz = msg_rcv.mtext[1];
85+
// Unpack bsz (low 16 bits) and actual_topk (high 16 bits) from mtext[1].
86+
// This matches the packing in save_output_msg_with_topk.cc:
87+
// mtext[1] = bsz | (max_num_logprobs << 16)
88+
int bsz = msg_rcv.mtext[1] & 0xFFFF;
89+
int actual_topk = (msg_rcv.mtext[1] >> 16) & 0xFFFF;
8690
out_data[0] = (int64_t)msg_rcv.mtext[0];
87-
out_data[1] = (int64_t)msg_rcv.mtext[1];
91+
out_data[1] = (int64_t)msg_rcv.mtext[1]; // keep packed value; Python unpacks
8892

8993
for (int i = 0; i < bsz; i++) {
90-
for (int j = 0; j < k + 1; j++) {
91-
const int64_t offset = i * (K + 1) + j;
94+
for (int j = 0; j < actual_topk; j++) {
95+
const int64_t offset = i * actual_topk + j;
9296
out_data[offset + 2] = (int64_t)msg_rcv.mtext[offset + 2];
9397
scores_data[offset] = msg_rcv.mtext_f[offset];
9498
}

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

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -109,20 +109,21 @@ void SaveOutMmsgTopK(const paddle::Tensor& x,
109109
: -inference_msg_id_from_env;
110110
int bsz = x.shape()[0];
111111
int max_num_logprobs = logprob_token_ids.shape()[1];
112-
msg_sed.mtext[1] = bsz;
112+
// Pack bsz (low 16 bits) and max_num_logprobs (high 16 bits) into mtext[1].
113+
// token_processor unpacks both fields to avoid reading unused topk slots.
114+
msg_sed.mtext[1] = bsz | (max_num_logprobs << 16);
113115
for (int i = 0; i < bsz; i++) {
114-
for (int j = 0; j < K + 1; j++) {
115-
const int64_t offset = i * (K + 1) + j;
116+
// Loop only over actual logprob columns (max_num_logprobs) instead of the
117+
// fixed K+1=21, and use max_num_logprobs as the stride so data is packed
118+
// densely in the message buffer.
119+
for (int j = 0; j < max_num_logprobs; j++) {
120+
const int64_t offset = i * max_num_logprobs + j;
116121
if (j == 0) {
117122
msg_sed.mtext[offset + 2] = (int)x_data[i];
118-
msg_sed.mtext_f[offset] = logprob_scores_data[i * max_num_logprobs + j];
119-
} else if (j < max_num_logprobs) {
120-
msg_sed.mtext[offset + 2] =
121-
(int)logprob_token_ids_data[i * max_num_logprobs + j];
122-
msg_sed.mtext_f[offset] = logprob_scores_data[i * max_num_logprobs + j];
123+
msg_sed.mtext_f[offset] = logprob_scores_data[offset];
123124
} else {
124-
msg_sed.mtext[offset + 2] = -1;
125-
msg_sed.mtext_f[offset] = 0.0;
125+
msg_sed.mtext[offset + 2] = (int)logprob_token_ids_data[offset];
126+
msg_sed.mtext_f[offset] = logprob_scores_data[offset];
126127
}
127128
if (preempted_idx_data[i] == 1) {
128129
msg_sed.mtext[offset + 2] = -9;

‎fastdeploy/output/token_processor.py‎

Lines changed: 20 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -846,10 +846,22 @@ def _process_batch_output(self):
846846
batch = self.output_tokens[1]
847847
accept_num = tokens[2 : batch + 2]
848848
elif self.use_logprobs:
849-
batch = self.output_tokens[1, 0]
850-
tokens = tokens[2 : batch * (K + 1) + 2].reshape([batch, K + 1])[:, : (K + 1)]
851-
scores = self.output_scores[: batch * (K + 1)].numpy().reshape([batch, K + 1])[:, : (K + 1)]
849+
# mtext[1] packs bsz (low 16 bits) and actual_topk (high 16 bits).
850+
# actual_topk = max_num_logprobs written by save_output_topk, which
851+
# equals the actual number of logprob columns in this step's message
852+
# (top_logprobs+1 across the batch). Using actual_topk as stride
853+
# avoids processing the K+1=21 fixed-size slots when fewer are needed.
854+
packed = int(self.output_tokens[1, 0])
855+
batch = packed & 0xFFFF
856+
actual_topk = (packed >> 16) & 0xFFFF
857+
tokens = tokens[2 : batch * actual_topk + 2].reshape([batch, actual_topk])
858+
scores = self.output_scores[: batch * actual_topk].numpy().reshape([batch, actual_topk])
852859
ranks = self.output_ranks[:batch].numpy()
860+
# Pre-convert the full [batch, actual_topk] arrays to Python lists once,
861+
# avoiding per-row .tolist() calls inside the loop below.
862+
tokens_lists = tokens.tolist()
863+
scores_lists = scores.tolist()
864+
ranks_list = ranks.tolist()
853865
else:
854866
batch = self.output_tokens[1, 0]
855867
tokens = tokens[2 : batch + 2]
@@ -1022,10 +1034,11 @@ def _process_batch_output(self):
10221034
topk_logprobs = scores[i, batch_token_index, :].tolist()
10231035
sampled_rank = ranks[i, batch_token_index].item()
10241036
else:
1025-
result.outputs.logprob = float(scores[i, 0])
1026-
topk_token_ids = tokens[i, :].tolist()
1027-
topk_logprobs = scores[i, :].tolist()
1028-
sampled_rank = ranks[i].item()
1037+
# Use pre-converted lists (batch .tolist() done before the loop).
1038+
result.outputs.logprob = scores_lists[i][0]
1039+
topk_token_ids = tokens_lists[i]
1040+
topk_logprobs = scores_lists[i]
1041+
sampled_rank = ranks_list[i]
10291042

10301043
if result.outputs.top_logprobs is None:
10311044
result.outputs.top_logprobs = LogprobsLists(

‎tests/output/test_token_processor.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -719,7 +719,8 @@ def test_process_batch_output_logprob_records_topk_and_caching():
719719
task.trace_carrier = None
720720
rm.tasks_list[0] = task
721721
rm.req_dict[task.request_id] = task
722-
processor.output_tokens[1, 0] = 1
722+
# mtext[1] packs bsz (low 16 bits) | actual_topk (high 16 bits)
723+
processor.output_tokens[1, 0] = 1 | ((K + 1) << 16)
723724
token_block = np.arange(K + 1, dtype=np.int64) + 3
724725
processor.output_tokens[2 : 2 + K + 1] = paddle.to_tensor(token_block.reshape([-1, 1]))
725726
processor.output_scores[: K + 1] = paddle.ones([K + 1, 1], dtype="float32")
@@ -842,7 +843,8 @@ def test_process_batch_output_prefill_chunk_and_adapter_skip():
842843
task.get = lambda key, default=None: getattr(task, key, default)
843844
rm.tasks_list[0] = task
844845
rm.req_dict[task.request_id] = task
845-
processor.output_tokens[1, 0] = 1
846+
# mtext[1] packs bsz (low 16 bits) | actual_topk (high 16 bits)
847+
processor.output_tokens[1, 0] = 1 | ((K + 1) << 16)
846848
processor.output_tokens[2 : 2 + K + 1] = paddle.to_tensor(np.ones([K + 1, 1], dtype=np.int64))
847849
processor.output_scores[: K + 1] = paddle.ones([K + 1, 1], dtype="float32")
848850
processor.output_ranks[0] = paddle.to_tensor(0, dtype="int64")

0 commit comments

Comments
 (0)