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

5
python/sglang/srt/disaggregation/decode.py
View File

@@ -541,6 +541,7 @@ class DecodeTransferQueue:
self.metadata_buffers = metadata_buffers
self.scheduler = scheduler
self.tree_cache = tree_cache
self.spec_algorithm = scheduler.spec_algorithm
def add(self, decode_req: DecodeRequest) -> None:
self.queue.append(decode_req)
@@ -582,6 +583,7 @@ class DecodeTransferQueue:
idx = decode_req.metadata_buffer_index
(
output_id,
output_hidden_states,
output_token_logprobs_val,
output_token_logprobs_idx,
output_top_logprobs_val,
@@ -589,7 +591,8 @@ class DecodeTransferQueue:
) = self.metadata_buffers.get_buf(idx)
decode_req.req.output_ids.append(output_id[0].item())
if not self.spec_algorithm.is_none():
decode_req.req.hidden_states_tensor = output_hidden_states
if decode_req.req.return_logprob:
decode_req.req.output_token_logprobs_val.append(
output_token_logprobs_val[0].item()