Add support for topk metadata transferring for PD (#10616)

Signed-off-by: Shangming Cai <csmthu@gmail.com>
2025-09-28 00:09:38 +08:00
parent 51f7c6bd3c
commit e23e280e16
6 changed files with 60 additions and 20 deletions
@@ -421,6 +421,8 @@ class SchedulerDisaggregationPrefillMixin:
                    last_hidden_index = (
                        hidden_state_offset + extend_input_len_per_req[i] - 1
                    )
+                    req.output_topk_p = batch.spec_info.topk_p[i]
+                    req.output_topk_index = batch.spec_info.topk_index[i]
                    if self.spec_algorithm.is_eagle3():
                        req.hidden_states_tensor = (
                            batch.spec_info.hidden_states[i].cpu().clone()