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
......@@ -421,6 +421,11 @@ class SchedulerDisaggregationPrefillMixin:
last_hidden_index = (
hidden_state_offset + extend_input_len_per_req[i] - 1
)
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()
)
......
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