From b73ac629cd3e2229a3feb116f82536b802ab88dc Mon Sep 17 00:00:00 2001 From: Ziming Huang Date: Thu, 18 Sep 2025 01:37:14 +0800 Subject: [PATCH] [BugFix] Fix incorrect hidden_states_tensor in pd disaggregation + eagle (#9976) --- python/sglang/srt/disaggregation/prefill.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index f128ae914..5b9255e31 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -421,9 +421,14 @@ class SchedulerDisaggregationPrefillMixin: 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() - ) + if self.spec_algorithm.is_eagle3(): + req.hidden_states_tensor = ( + batch.spec_info.hidden_states[i].cpu().clone() + ) + else: + 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