chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[PD] Transfer hidden states for mtp when disaggregation (#7242)

Browse Source
This commit is contained in:
Atream
2025-06-20 02:22:47 +08:00
committed by GitHub
parent d20a073bc3
commit 4f838c09cd
6 changed files with 43 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
python/sglang/srt/disaggregation/prefill.py
View File

@@ -393,6 +393,8 @@ class SchedulerDisaggregationPrefillMixin:
logits_output.input_token_logprobs = tuple(
logits_output.input_token_logprobs.tolist()
)
hidden_state_offset = 0
for i, (req, next_token_id) in enumerate(
zip(batch.reqs, next_token_ids, strict=True)
):
@@ -402,6 +404,16 @@ class SchedulerDisaggregationPrefillMixin:
req.output_ids.append(next_token_id)
self.tree_cache.cache_unfinished_req(req) # update the tree and lock
self.disagg_prefill_inflight_queue.append(req)
if logits_output.hidden_states is not None:
last_hidden_index = (
hidden_state_offset + extend_input_len_per_req[i] - 1
)
req.hidden_states_tensor = (
logits_output.hidden_states[last_hidden_index].cpu().clone()
)
hidden_state_offset += extend_input_len_per_req[i]
else:
req.hidden_states_tensor = None
if req.return_logprob:
assert extend_logprob_start_len_per_req is not None
assert extend_input_len_per_req is not None