Skip to content

[Others]add logporbs chunked process #8098

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
Jiang-Jia-Jun merged 1 commit into from
Jul 14, 2026
Merged

[Others]add logporbs chunked process #8098

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
48 changes: 28 additions & 20 deletions fastdeploy/worker/gpu_model_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3642,35 +3642,43 @@ def _get_prompt_logprobs_list(
# In the == case, there are no more prompt logprobs to produce
# but we want to defer returning them to the next step where we
# have new generated tokens to return.
num_logits = num_tokens
num_logits_total = num_tokens
else:
# This is the last chunk of prompt tokens to return.
num_logits = num_remaining_tokens
num_logits_total = num_remaining_tokens
completed_prefill_reqs.append(request)
prompt_logprobs_list[batch_id] = logprobs_tensors
if num_logits <= 0:
if num_logits_total <= 0:
# This can happen for the final chunk if we prefilled exactly
# (num_prompt_tokens - 1) tokens for this request in the prior
# step. There are no more prompt logprobs to produce.
continue
offset = self.share_inputs["cu_seqlens_q"][batch_id]
prompt_hidden_states = hidden_states[offset : offset + num_logits]
logits = self.model.compute_logits(prompt_hidden_states, self.forward_meta)
prompt_token_ids = request.prompt_token_ids[start_tok : start_tok + num_logits]
prompt_token_ids_tensor = paddle.to_tensor(prompt_token_ids, dtype="int64")
if logprobs_mode == "raw_logprobs":
raw_logprobs = self.sampler.compute_logprobs(logits)
elif logprobs_mode == "raw_logits":
raw_logprobs = logits
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)
logprobs_tensors.selected_token_ranks[chunk_slice].copy_(ranks, False)
max_prompt_logprobs_chunk_size = 8 * 1024

@PaddlePaddle-bot PaddlePaddle-bot Jul 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

🟡 建议 这里仅给 GPU runner 的 prompt logprobs 加了分块，但同一逻辑在 XPU/Metax runner 仍保持整段计算 logits。

fastdeploy/worker/xpu_model_runner.py:276-293 和 fastdeploy/worker/metax_model_runner.py:2817-2832 仍会一次性对 num_logits 个 prompt token 计算 [num_logits, vocab]，长 prompt + prompt_logprobs 场景下会保留本 PR 要解决的显存/内存峰值问题。FastDeploy checklist A6 要求通用 GPU runner 变更检查多硬件同步。

建议把同样的 max_prompt_logprobs_chunk_size 分块处理同步到这些同名实现；如果本次只支持 CUDA，请在代码/PR 描述中明确范围，并补充对应硬件不受支持或待同步的说明。

for chunk_offset in range(0, num_logits_total, max_prompt_logprobs_chunk_size):
chunk_num_logits = min(max_prompt_logprobs_chunk_size, num_logits_total - chunk_offset)
prompt_hidden_states = hidden_states[offset + chunk_offset : offset + chunk_offset + chunk_num_logits]
logits = self.model.compute_logits(prompt_hidden_states, self.forward_meta)
prompt_token_ids = request.prompt_token_ids[
start_tok + chunk_offset : start_tok + chunk_offset + chunk_num_logits
]
prompt_token_ids_tensor = paddle.to_tensor(prompt_token_ids, dtype="int64")
if logprobs_mode == "raw_logprobs":
raw_logprobs = self.sampler.compute_logprobs(logits)
elif logprobs_mode == "raw_logits":
raw_logprobs = logits
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 + chunk_offset,
start_idx + chunk_offset + chunk_num_logits,
)
logprobs_tensors.logprob_token_ids[chunk_slice].copy_(token_ids, False)
logprobs_tensors.logprobs[chunk_slice].copy_(logprobs, False)
logprobs_tensors.selected_token_ranks[chunk_slice].copy_(ranks, False)

for req in completed_prefill_reqs:
del self.prompt_logprobs_reqs[req.request_id]
Expand Down
Loading