From a8e07aa8c7c6e7b8a0980ac8108e9cd6c93dee88 Mon Sep 17 00:00:00 2001 From: Sunny-bot1 <592045536@qq.com> Date: Thu, 5 Feb 2026 16:59:49 +0800 Subject: [PATCH 1/3] support logprob async copy --- fastdeploy/model_executor/layers/sample/sampler.py | 11 +++++++---- fastdeploy/model_executor/pre_and_post_process.py | 2 +- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/fastdeploy/model_executor/layers/sample/sampler.py b/fastdeploy/model_executor/layers/sample/sampler.py index 19b1fcdfb12..281f1efed00 100644 --- a/fastdeploy/model_executor/layers/sample/sampler.py +++ b/fastdeploy/model_executor/layers/sample/sampler.py @@ -478,10 +478,13 @@ def gather_logprobs( else: indices = token_ids top_logprobs = token_logprobs - indices = indices.cpu() - top_logprobs = top_logprobs.cpu() - token_ranks = token_ranks.cpu() - return LogprobsTensors(indices, top_logprobs, token_ranks) + indices_cpu = paddle.empty_like(indices, device="cpu").pin_memory() + top_logprobs_cpu = paddle.empty_like(top_logprobs, device="cpu").pin_memory() + token_ranks_cpu = paddle.empty_like(token_ranks, device="cpu").pin_memory() + indices_cpu.copy_(indices, False) + top_logprobs_cpu.copy_(top_logprobs, False) + token_ranks_cpu.copy_(token_ranks, False) + return LogprobsTensors(indices_cpu, top_logprobs_cpu, token_ranks_cpu) def forward_cuda( self, diff --git a/fastdeploy/model_executor/pre_and_post_process.py b/fastdeploy/model_executor/pre_and_post_process.py index af269987ce5..9a65a12c9e2 100644 --- a/fastdeploy/model_executor/pre_and_post_process.py +++ b/fastdeploy/model_executor/pre_and_post_process.py @@ -464,7 +464,7 @@ def save_output_normal( sampler_output, model_output.index_to_batch_id, model_output.enable_pd_reorder ) save_output_topk( - sampler_output.sampled_token_ids, + share_inputs["sampled_token_ids"], sampler_output.logprobs_tensors.logprob_token_ids, sampler_output.logprobs_tensors.logprobs, sampler_output.logprobs_tensors.selected_token_ranks, From 4b1804724abd7a3a8e8510f17e0fe5535d076751 Mon Sep 17 00:00:00 2001 From: Sunny-bot1 <592045536@qq.com> Date: Thu, 5 Feb 2026 20:47:02 +0800 Subject: [PATCH 2/3] fix prompt logprob --- fastdeploy/model_executor/layers/sample/sampler.py | 11 +++++++---- fastdeploy/model_executor/pre_and_post_process.py | 2 +- fastdeploy/worker/gpu_model_runner.py | 2 ++ fastdeploy/worker/output.py | 6 +++--- 4 files changed, 13 insertions(+), 8 deletions(-) diff --git a/fastdeploy/model_executor/layers/sample/sampler.py b/fastdeploy/model_executor/layers/sample/sampler.py index 19b1fcdfb12..281f1efed00 100644 --- a/fastdeploy/model_executor/layers/sample/sampler.py +++ b/fastdeploy/model_executor/layers/sample/sampler.py @@ -478,10 +478,13 @@ def gather_logprobs( else: indices = token_ids top_logprobs = token_logprobs - indices = indices.cpu() - top_logprobs = top_logprobs.cpu() - token_ranks = token_ranks.cpu() - return LogprobsTensors(indices, top_logprobs, token_ranks) + indices_cpu = paddle.empty_like(indices, device="cpu").pin_memory() + top_logprobs_cpu = paddle.empty_like(top_logprobs, device="cpu").pin_memory() + token_ranks_cpu = paddle.empty_like(token_ranks, device="cpu").pin_memory() + indices_cpu.copy_(indices, False) + top_logprobs_cpu.copy_(top_logprobs, False) + token_ranks_cpu.copy_(token_ranks, False) + return LogprobsTensors(indices_cpu, top_logprobs_cpu, token_ranks_cpu) def forward_cuda( self, diff --git a/fastdeploy/model_executor/pre_and_post_process.py b/fastdeploy/model_executor/pre_and_post_process.py index af269987ce5..9a65a12c9e2 100644 --- a/fastdeploy/model_executor/pre_and_post_process.py +++ b/fastdeploy/model_executor/pre_and_post_process.py @@ -464,7 +464,7 @@ def save_output_normal( sampler_output, model_output.index_to_batch_id, model_output.enable_pd_reorder ) save_output_topk( - sampler_output.sampled_token_ids, + share_inputs["sampled_token_ids"], sampler_output.logprobs_tensors.logprob_token_ids, sampler_output.logprobs_tensors.logprobs, sampler_output.logprobs_tensors.selected_token_ranks, diff --git a/fastdeploy/worker/gpu_model_runner.py b/fastdeploy/worker/gpu_model_runner.py index 94813313a3f..1a72f0c5019 100644 --- a/fastdeploy/worker/gpu_model_runner.py +++ b/fastdeploy/worker/gpu_model_runner.py @@ -3119,6 +3119,8 @@ def _get_prompt_logprobs_list( token_ids, logprobs, ranks = self.sampler.gather_logprobs( raw_logprobs, num_prompt_logprobs, prompt_token_ids_tensor ) + # Synchronize before using token_ids, logprobs and ranks to ensure async copy are completed. + paddle.device.synchronize() chunk_slice = slice(start_idx, start_idx + num_logits) logprobs_tensors.logprob_token_ids[chunk_slice].copy_(token_ids, False) logprobs_tensors.logprobs[chunk_slice].copy_(logprobs, False) diff --git a/fastdeploy/worker/output.py b/fastdeploy/worker/output.py index 5d2608439a0..9d5501636d3 100644 --- a/fastdeploy/worker/output.py +++ b/fastdeploy/worker/output.py @@ -119,9 +119,9 @@ def slice_rows(self, start: int, end: int): """ with paddle.no_grad(): return LogprobsTensors( - paddle.to_tensor(self.logprob_token_ids[start:end], place=self.logprob_token_ids.place), - paddle.to_tensor(self.logprobs[start:end], place=self.logprob_token_ids.place), - paddle.to_tensor(self.selected_token_ranks[start:end], place=self.logprob_token_ids.place), + paddle.to_tensor(self.logprob_token_ids.cpu()[start:end], place="cpu"), + paddle.to_tensor(self.logprobs.cpu()[start:end], place="cpu"), + paddle.to_tensor(self.selected_token_ranks.cpu()[start:end], place="cpu"), ) From a8aa4921938e0998549a8cb3bc6bb8f8e2174904 Mon Sep 17 00:00:00 2001 From: Sunny-bot1 <592045536@qq.com> Date: Sat, 7 Feb 2026 12:07:59 +0800 Subject: [PATCH 3/3] fix xpu --- .../model_executor/layers/sample/sampler.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/fastdeploy/model_executor/layers/sample/sampler.py b/fastdeploy/model_executor/layers/sample/sampler.py index 281f1efed00..b1685838a74 100644 --- a/fastdeploy/model_executor/layers/sample/sampler.py +++ b/fastdeploy/model_executor/layers/sample/sampler.py @@ -478,12 +478,17 @@ def gather_logprobs( else: indices = token_ids top_logprobs = token_logprobs - indices_cpu = paddle.empty_like(indices, device="cpu").pin_memory() - top_logprobs_cpu = paddle.empty_like(top_logprobs, device="cpu").pin_memory() - token_ranks_cpu = paddle.empty_like(token_ranks, device="cpu").pin_memory() - indices_cpu.copy_(indices, False) - top_logprobs_cpu.copy_(top_logprobs, False) - token_ranks_cpu.copy_(token_ranks, False) + if current_platform.is_cuda(): + indices_cpu = paddle.empty_like(indices, device="cpu").pin_memory() + top_logprobs_cpu = paddle.empty_like(top_logprobs, device="cpu").pin_memory() + token_ranks_cpu = paddle.empty_like(token_ranks, device="cpu").pin_memory() + indices_cpu.copy_(indices, False) + top_logprobs_cpu.copy_(top_logprobs, False) + token_ranks_cpu.copy_(token_ranks, False) + else: + indices_cpu = indices.cpu() + top_logprobs_cpu = top_logprobs.cpu() + token_ranks_cpu = token_ranks.cpu() return LogprobsTensors(indices_cpu, top_logprobs_cpu, token_ranks_cpu) def forward_cuda(