Abstraction for spec worker and code cleanup (#11643)
This commit is contained in:
Liangsheng Yin
@@ -42,23 +42,21 @@ class SchedulerOutputProcessorMixin:
|
||||
skip_stream_req = None
|
||||
|
||||
if self.is_generation:
|
||||
if result.copy_done is not None:
|
||||
result.copy_done.synchronize()
|
||||
|
||||
(
|
||||
logits_output,
|
||||
next_token_ids,
|
||||
extend_input_len_per_req,
|
||||
extend_logprob_start_len_per_req,
|
||||
copy_done,
|
||||
) = (
|
||||
result.logits_output,
|
||||
result.next_token_ids,
|
||||
result.extend_input_len_per_req,
|
||||
result.extend_logprob_start_len_per_req,
|
||||
result.copy_done,
|
||||
)
|
||||
|
||||
if copy_done is not None:
|
||||
copy_done.synchronize()
|
||||
|
||||
# Move next_token_ids and logprobs to cpu
|
||||
next_token_ids = next_token_ids.tolist()
|
||||
if batch.return_logprob:
|
||||
@@ -199,57 +197,52 @@ class SchedulerOutputProcessorMixin:
|
||||
|
||||
self.stream_output(batch.reqs, batch.return_logprob, skip_stream_req)
|
||||
|
||||
def hacky_process_eagle_overlap_result(
|
||||
def _resolve_spec_overlap_token_ids(
|
||||
self: Scheduler, result: GenerationBatchResult, batch: ScheduleBatch
|
||||
):
|
||||
# TODO(lsyin): try use a copy stream to share SMs with forward
|
||||
# FIXME(lsyin): better organize this token free logic in eagle-overlap
|
||||
last_batch_allocate_lens_cpu = result.last_batch_allocate_lens.tolist()
|
||||
accept_lens_cpu = result.accept_lens.tolist()
|
||||
) -> List[List[int]]:
|
||||
"""Resolve the padding next token ids for speculative decoding with overlap."""
|
||||
assert result.next_token_ids.is_cpu
|
||||
assert result.accept_lens.is_cpu
|
||||
assert result.allocate_lens.is_cpu
|
||||
|
||||
next_token_ids = result.next_token_ids.tolist()
|
||||
accept_lens = result.accept_lens.tolist()
|
||||
result.num_accepted_tokens = sum(accept_lens)
|
||||
|
||||
predict_tokens = []
|
||||
num_draft_tokens = self.draft_worker.speculative_num_draft_tokens
|
||||
stride = self.draft_worker.speculative_num_draft_tokens
|
||||
for i, req in enumerate(batch.reqs):
|
||||
predict_tokens.append(
|
||||
next_token_ids[
|
||||
i * num_draft_tokens : i * num_draft_tokens + accept_lens_cpu[i]
|
||||
]
|
||||
next_token_ids[i * stride : i * stride + accept_lens[i]]
|
||||
)
|
||||
# FIXME(lsyin): move this update elsewhere
|
||||
req.spec_verify_ct += 1
|
||||
|
||||
return last_batch_allocate_lens_cpu, accept_lens_cpu, predict_tokens
|
||||
return predict_tokens
|
||||
|
||||
def process_batch_result_decode(
|
||||
self: Scheduler,
|
||||
batch: ScheduleBatch,
|
||||
result: GenerationBatchResult,
|
||||
):
|
||||
logits_output, next_token_ids, can_run_cuda_graph, copy_done = (
|
||||
if result.copy_done is not None:
|
||||
result.copy_done.synchronize()
|
||||
|
||||
logits_output, next_token_ids, can_run_cuda_graph = (
|
||||
result.logits_output,
|
||||
result.next_token_ids,
|
||||
result.can_run_cuda_graph,
|
||||
result.copy_done,
|
||||
)
|
||||
self.num_generated_tokens += len(batch.reqs)
|
||||
|
||||
if copy_done is not None:
|
||||
copy_done.synchronize()
|
||||
|
||||
if batch.spec_algorithm.is_none():
|
||||
next_token_ids = next_token_ids.tolist()
|
||||
if batch.return_logprob:
|
||||
next_token_logprobs = logits_output.next_token_logprobs.tolist()
|
||||
elif batch.is_v2_eagle:
|
||||
(
|
||||
last_batch_allocate_lens_cpu,
|
||||
accept_lens_cpu,
|
||||
next_token_ids,
|
||||
) = self.hacky_process_eagle_overlap_result(result, batch)
|
||||
result.num_accepted_tokens = sum(accept_lens_cpu)
|
||||
next_token_ids = self._resolve_spec_overlap_token_ids(result, batch)
|
||||
allocate_lens_list = result.allocate_lens.tolist()
|
||||
accept_lens_list = result.accept_lens.tolist()
|
||||
|
||||
# FIXME(lsyin): we suppose we have already got the num_accepted_tokens in result
|
||||
self.num_generated_tokens += len(batch.reqs)
|
||||
if not self.spec_algorithm.is_none():
|
||||
self.update_spec_metrics(batch.batch_size(), result.num_accepted_tokens)
|
||||
|
||||
@@ -264,43 +257,38 @@ class SchedulerOutputProcessorMixin:
|
||||
continue
|
||||
|
||||
if self.enable_overlap and req.finished():
|
||||
indices_to_free = None
|
||||
if self.page_size == 1:
|
||||
if batch.spec_algorithm.is_eagle():
|
||||
from sglang.srt.speculative.eagle_worker_v2 import (
|
||||
free_spec_dec_tokens_page_size_1,
|
||||
)
|
||||
from sglang.srt.speculative.eagle_info import EagleDraftInput
|
||||
|
||||
free_spec_dec_tokens_page_size_1(
|
||||
self.req_to_token_pool,
|
||||
self.token_to_kv_pool_allocator,
|
||||
req,
|
||||
last_batch_allocate_lens_cpu[i],
|
||||
None,
|
||||
)
|
||||
end_p = allocate_lens_list[i]
|
||||
start_p = end_p - EagleDraftInput.ALLOC_LEN_PER_DECODE
|
||||
indices_to_free = self.req_to_token_pool.req_to_token[
|
||||
req.req_pool_idx
|
||||
][start_p:end_p]
|
||||
else:
|
||||
# Free the one extra delayed token
|
||||
self.token_to_kv_pool_allocator.free(
|
||||
batch.out_cache_loc[i : i + 1]
|
||||
)
|
||||
indices_to_free = batch.out_cache_loc[i : i + 1]
|
||||
else:
|
||||
if batch.spec_algorithm.is_eagle():
|
||||
# TODO(lsyin): support eagle with page_size > 1
|
||||
# TODO(spec-v2): support eagle with page_size > 1
|
||||
raise NotImplementedError()
|
||||
else:
|
||||
if (
|
||||
len(req.origin_input_ids) + len(req.output_ids) - 1
|
||||
) % self.page_size == 0:
|
||||
# Only free when the extra token is in a new page
|
||||
self.token_to_kv_pool_allocator.free(
|
||||
batch.out_cache_loc[i : i + 1]
|
||||
)
|
||||
indices_to_free = batch.out_cache_loc[i : i + 1]
|
||||
|
||||
if indices_to_free is not None:
|
||||
self.token_to_kv_pool_allocator.free(indices_to_free)
|
||||
continue
|
||||
|
||||
if batch.spec_algorithm.is_none():
|
||||
req.output_ids.append(next_token_id)
|
||||
elif batch.is_v2_eagle:
|
||||
# FIXME(lsyin): non-overlap spec worker will solve the output_ids in speculative decoding
|
||||
# !!!unify the logic here!!!
|
||||
# Only v2 eagle's output_ids are updated here.
|
||||
req.output_ids.extend(next_token_id)
|
||||
|
||||
req.check_finished()
|
||||
@@ -308,24 +296,13 @@ class SchedulerOutputProcessorMixin:
|
||||
if batch.is_v2_eagle and self.cur_batch.forward_mode.is_extend():
|
||||
# FIXME(lsyin): fix the messy logic here
|
||||
# 1) when not overlap (v2 impl), we free the extra tokens in the req
|
||||
# 2) when overlap and current batch is extend, we free the extra tokens in the req of the previous batch
|
||||
from sglang.srt.speculative.eagle_worker_v2 import (
|
||||
free_spec_dec_tokens_page_size_1,
|
||||
)
|
||||
|
||||
new_seq_len = len(req.origin_input_ids) + len(req.output_ids) - 1
|
||||
# FIXME(lsyin): remove this assert
|
||||
assert new_seq_len == int(
|
||||
batch.seq_lens_cpu[i] + accept_lens_cpu[i]
|
||||
), f"{new_seq_len=} vs {batch.seq_lens_cpu[i] + accept_lens_cpu[i]=}"
|
||||
|
||||
free_spec_dec_tokens_page_size_1(
|
||||
self.req_to_token_pool,
|
||||
self.token_to_kv_pool_allocator,
|
||||
req,
|
||||
last_batch_allocate_lens_cpu[i],
|
||||
new_seq_len,
|
||||
)
|
||||
# 2) overlap eagle and the current batch is prefill. This seq will not run extra iteration.
|
||||
start_p = batch.seq_lens_cpu[i] + accept_lens_list[i]
|
||||
end_p = allocate_lens_list[i]
|
||||
indices_to_free = self.req_to_token_pool.req_to_token[
|
||||
req.req_pool_idx
|
||||
][start_p:end_p]
|
||||
self.token_to_kv_pool_allocator.free(indices_to_free)
|
||||
|
||||
if self.server_args.disaggregation_decode_enable_offload_kvcache:
|
||||
# Asynchronously offload KV cache; cache_finished_req will be called after Device->Host transfer completes
|
||||
|
||||
Reference in New Issue
Block a user