Skip to content

GitLab

Unverified Commit b73ac629 authored by Ziming Huang's avatar Ziming Huang Committed by GitHub
Browse files

[BugFix] Fix incorrect hidden_states_tensor in pd disaggregation + eagle (#9976)

parent 77098aea
Hide whitespace changes
Inline Side-by-side
Showing with 8 additions and 3 deletions
python/sglang/srt/disaggregation/prefill.py
View file @ b73ac629
...@@ -421,9 +421,14 @@ class SchedulerDisaggregationPrefillMixin: ...@@ -421,9 +421,14 @@ class SchedulerDisaggregationPrefillMixin:
last_hidden_index = ( last_hidden_index = (
hidden_state_offset + extend_input_len_per_req[i] - 1 hidden_state_offset + extend_input_len_per_req[i] - 1
) )
req.hidden_states_tensor = ( if self.spec_algorithm.is_eagle3():
logits_output.hidden_states[last_hidden_index].cpu().clone() 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] hidden_state_offset += extend_input_len_per_req[i]
else: else:
req.hidden_states_tensor = None req.hidden_states_tensor = None
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment