[Feature] Spec-Overlap supporting DP-ATTN; PD-Disaggregation; npugraph mode (#12443)

This commit is contained in:
Even Zhou
2025-11-15 21:51:07 +08:00
committed by GitHub
parent 0d41ddfbd0
commit 2aec8b6e1b
24 changed files with 657 additions and 366 deletions
+23 -7
View File
@@ -114,17 +114,33 @@ class FutureMap:
else:
_resolve_future_token_ids(model_worker_batch.input_ids, self.token_ids_buf)
def is_empty_slice(self, s: slice) -> bool:
start, stop, step = s.indices(self.future_buffer_len)
if step > 0:
return start >= stop
else:
return start <= stop
def store_to_map(
self, future_indices: FutureIndices, batch_result: GenerationBatchResult
):
intv = future_indices.interval
if self.spec_algo.is_eagle():
draft_input: EagleDraftInput = batch_result.next_draft_input
self._lazy_init_buf(draft_input)
self.topk_p_buf[intv] = draft_input.topk_p
self.topk_index_buf[intv] = draft_input.topk_index
self.hidden_states_buf[intv] = draft_input.hidden_states
self.verified_id_buf[intv] = draft_input.verified_id
self.new_seq_lens_buf[intv] = draft_input.new_seq_lens
self.store_to_map_for_new_batch(future_indices, draft_input)
else:
intv = future_indices.interval
self.token_ids_buf[intv] = batch_result.next_token_ids
def store_to_map_for_new_batch(
self, future_indices: FutureIndices, draft_input: EagleDraftInput
):
intv = future_indices.interval
# idle indices do not need store info
if self.is_empty_slice(intv):
return
self._lazy_init_buf(draft_input)
self.topk_p_buf[intv] = draft_input.topk_p
self.topk_index_buf[intv] = draft_input.topk_index
self.hidden_states_buf[intv] = draft_input.hidden_states
self.verified_id_buf[intv] = draft_input.verified_id
self.new_seq_lens_buf[intv] = draft_input.new_seq_lens