Beta spec-overlap for EAGLE (#11398)
Co-authored-by: Lianmin Zheng <15100009+merrymercy@users.noreply.github.com> Co-authored-by: Hanming Lu <69857889+hanming-lu@users.noreply.github.com>
This commit is contained in:
Liangsheng Yin
co-authored by
Lianmin Zheng
Hanming Lu
Lianmin Zheng
Hanming Lu
parent
47c606d3dc
commit
20a6c0a63d
@@ -1,7 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from typing import TYPE_CHECKING, List, Optional, Tuple, Union
|
||||
|
||||
@@ -200,6 +199,28 @@ class SchedulerOutputProcessorMixin:
|
||||
|
||||
self.stream_output(batch.reqs, batch.return_logprob, skip_stream_req)
|
||||
|
||||
def hacky_process_eagle_overlap_result(
|
||||
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()
|
||||
next_token_ids = result.next_token_ids.tolist()
|
||||
|
||||
predict_tokens = []
|
||||
num_draft_tokens = 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]
|
||||
]
|
||||
)
|
||||
# FIXME(lsyin): move this update elsewhere
|
||||
req.spec_verify_ct += 1
|
||||
|
||||
return last_batch_allocate_lens_cpu, accept_lens_cpu, predict_tokens
|
||||
|
||||
def process_batch_result_decode(
|
||||
self: Scheduler,
|
||||
batch: ScheduleBatch,
|
||||
@@ -220,6 +241,17 @@ class SchedulerOutputProcessorMixin:
|
||||
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)
|
||||
|
||||
# FIXME(lsyin): we suppose we have already got the num_accepted_tokens in result
|
||||
if not self.spec_algorithm.is_none():
|
||||
self.update_spec_metrics(batch.batch_size(), result.num_accepted_tokens)
|
||||
|
||||
self.token_to_kv_pool_allocator.free_group_begin()
|
||||
|
||||
@@ -227,29 +259,74 @@ class SchedulerOutputProcessorMixin:
|
||||
# NOTE: the length of reqs and next_token_ids don't match if it is spec decoding.
|
||||
# We should ignore using next_token_ids for spec decoding cases.
|
||||
for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)):
|
||||
req: Req
|
||||
if req.is_retracted:
|
||||
continue
|
||||
|
||||
if self.enable_overlap and req.finished():
|
||||
# Free the one extra delayed token
|
||||
if self.page_size == 1:
|
||||
self.token_to_kv_pool_allocator.free(batch.out_cache_loc[i : i + 1])
|
||||
else:
|
||||
# Only free when the extra token is in a new page
|
||||
if (
|
||||
len(req.origin_input_ids) + len(req.output_ids) - 1
|
||||
) % self.page_size == 0:
|
||||
if batch.spec_algorithm.is_eagle():
|
||||
from sglang.srt.speculative.eagle_worker_v2 import (
|
||||
free_spec_dec_tokens_page_size_1,
|
||||
)
|
||||
|
||||
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,
|
||||
)
|
||||
else:
|
||||
# Free the one extra delayed token
|
||||
self.token_to_kv_pool_allocator.free(
|
||||
batch.out_cache_loc[i : i + 1]
|
||||
)
|
||||
else:
|
||||
if batch.spec_algorithm.is_eagle():
|
||||
# TODO(lsyin): 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]
|
||||
)
|
||||
continue
|
||||
|
||||
if batch.spec_algorithm.is_none():
|
||||
# speculative worker will solve the output_ids in speculative decoding
|
||||
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!!!
|
||||
req.output_ids.extend(next_token_id)
|
||||
|
||||
req.check_finished()
|
||||
if req.finished():
|
||||
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,
|
||||
)
|
||||
|
||||
if self.server_args.disaggregation_decode_enable_offload_kvcache:
|
||||
# Asynchronously offload KV cache; cache_finished_req will be called after Device->Host transfer completes
|
||||
if not self.decode_offload_manager.offload_kv_cache(req):
|
||||
|
||||
Reference in New Issue
Block a user