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