[Feature] Spec-Overlap supporting DP-ATTN; PD-Disaggregation; npugraph mode (#12443)
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user