Skip to content

[Optimization] Support logprob async copy #6362

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
EmmonsCurse merged 5 commits into from
Feb 9, 2026
Merged

[Optimization] Support logprob async copy #6362

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a8e07aa
support logprob async copy
Sunny-bot1 Feb 5, 2026
4b18047
fix prompt logprob
Sunny-bot1 Feb 5, 2026
3ac5902
Merge branch 'logprob_async' of https://github.com/Sunny-bot1/FastDep…
Sunny-bot1 Feb 5, 2026
3fdd70b
Merge branch 'develop' of https://github.com/PaddlePaddle/FastDeploy …
Sunny-bot1 Feb 5, 2026
a8aa492
fix xpu
Sunny-bot1 Feb 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 12 additions & 4 deletions fastdeploy/model_executor/layers/sample/sampler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -478,10 +478,18 @@ 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)
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(
self,
Expand Down
2 changes: 1 addition & 1 deletion fastdeploy/model_executor/pre_and_post_process.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
2 changes: 2 additions & 0 deletions fastdeploy/worker/gpu_model_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Copilot AI Feb 9, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

注释里的英文语法建议调整："to ensure async copy are completed" 应改为 "to ensure async copies are completed" 或 "... is completed"，避免语法错误影响可读性。

Suggested change
# Synchronize before using token_ids, logprobs and ranks to ensure async copy are completed.
# Synchronize before using token_ids, logprobs and ranks to ensure async copies are completed.

Copilot uses AI. Check for mistakes.
paddle.device.synchronize()
Comment on lines +3122 to +3123

Copilot AI Feb 9, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这里使用 paddle.device.synchronize() 会进行全设备同步，会显著降低 overlap-schedule/异步拷贝的收益（尤其在循环里每个 request 都同步）。建议改为更细粒度的同步方式：例如使用 CUDA event 只等待本次 D2H copy 对应的 stream，或将 prompt_logprobs 的目标缓冲区改为 pinned memory 并直接对目标做非阻塞 copy，然后复用 runner 里已有的 event 同步机制。

Suggested change
# Synchronize before using token_ids, logprobs and ranks to ensure async copy are completed.
paddle.device.synchronize()
# Synchronize current CUDA stream before using token_ids, logprobs and ranks
paddle.device.cuda.current_stream().synchronize()

Copilot uses AI. Check for mistakes.
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)
Expand Down
6 changes: 3 additions & 3 deletions fastdeploy/worker/output.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"),
Comment on lines +122 to +124

Copilot AI Feb 9, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LogprobsTensors.slice_rows 这里先对整块 Tensor 调用 .cpu() 再切片，并且又包了一层 paddle.to_tensor(...)，会导致不必要的数据拷贝/同步（在按 batch 循环 slice_rows 的场景下可能放大为多次整块拷贝），同时也会丢失上游可能使用的 pinned memory 优势。建议直接对原 Tensor 做切片并返回（保持原 place/pin 状态）；如果确实需要搬到 CPU，也应当先切片再 .cpu()，并避免对 Tensor 再次 to_tensor。

Suggested change
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"),
self.logprob_token_ids[start:end],
self.logprobs[start:end],
self.selected_token_ranks[start:end],

Copilot uses AI. Check for mistakes.
)


Expand Down
Loading