diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker.py b/python/sglang/srt/speculative/multi_layer_eagle_worker.py index 6aa79e7a8..305eff355 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker.py @@ -26,11 +26,6 @@ from sglang.srt.layers.utils.logprob import add_output_logprobs_for_spec_v1 from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.managers.tp_worker import TpModelWorker -from sglang.srt.mem_cache.chunk_cache import SWAChunkCache -from sglang.srt.mem_cache.common import ( - alloc_paged_token_slots_extend, - alloc_token_slots, -) from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, ForwardBatch, @@ -47,28 +42,19 @@ from sglang.srt.speculative.eagle_utils import ( build_tree_kernel_efficient, organize_draft_results, ) -from sglang.srt.speculative.eagle_worker import get_last_loc_large_page_size_top_k_1 from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import ( MultiLayerEagleDraftExtendCudaGraphRunner, ) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_utils import ( - assign_draft_cache_locs, detect_nan, draft_tp_context, fast_topk, generate_token_bitmask, - get_last_loc_large_page_size_large_top_k, load_token_map, select_top_k_tokens, ) -from sglang.srt.utils import ( - empty_context, - get_available_gpu_memory, - is_cuda, - is_npu, - next_power_of_2, -) +from sglang.srt.utils import empty_context, get_available_gpu_memory, is_cuda, is_npu _is_npu = is_npu() @@ -346,156 +332,16 @@ class MultiLayerEagleWorker(TpModelWorker): ) def _draft_preprocess_decode(self, batch: ScheduleBatch): - if isinstance(batch.tree_cache, SWAChunkCache): - for req in batch.reqs: - batch.tree_cache.evict_swa(req, req.seqlen - 1) + from sglang.srt.speculative.eagle_worker import EAGLEWorker - # Parse args - num_seqs = batch.batch_size() - spec_info = batch.spec_info - - # Accumulate penalty - if batch.sampling_info.penalizer_orchestrator.is_required: - # This is a relaxed version of penalties for speculative decoding. - batch.sampling_info.penalizer_orchestrator.cumulate_output_tokens( - spec_info.verified_id.to(torch.int64) - ) - - # Allocate cache locations - # Layout of the out_cache_loc - # [ topk 0 ] [ topk 1 ] - # [iter=0, iter=1, iter=2] [iter=0, iter=1, iter=2] - if self.page_size == 1: - out_cache_loc, token_to_kv_pool_state_backup = alloc_token_slots( - batch.tree_cache, - num_seqs * self.speculative_num_steps * self.topk, - backup_state=True, - ) - duplicate_cache_len = 0 - source_cache_loc, target_cache_loc, last_page_lens_cumsum = None, None, None - else: - if self.topk == 1: - prefix_lens, seq_lens, last_loc = get_last_loc_large_page_size_top_k_1( - batch.req_to_token_pool.req_to_token, - batch.req_pool_indices, - batch.seq_lens, - self.speculative_num_steps, - ) - prefix_lens_cpu = batch.seq_lens_cpu - seq_lens_cpu = batch.seq_lens_cpu + self.speculative_num_steps - extend_num_tokens = num_seqs * self.speculative_num_steps - duplicate_cache_len = 0 - source_cache_loc, target_cache_loc, last_page_lens_cumsum = ( - None, - None, - None, - ) - else: - # In this case, the last partial page needs to be duplicated. - # KV cache layout in batch.req_to_token_pool.req_to_token: - # - # | -------- | -- xxxx .. | -- xxxx .. | -- xxxx .. | - # prefix top-k = 0 tok-k = 1 top-k = 2 - # - # "-" means prefix tokens - # "x" means speculative draft tokens - # "." means padded tokens - - # TODO(lmzheng): The current implementation is still a fake support - # for page size > 1. In the `assign_draft_cache_locs` below, - # we directly move the indices instead of the real kv cache. - # This only works when the kernel backend runs with page size = 1. - # If the kernel backend runs with page size > 1, we need to - # duplicate the real KV cache. The overhead of duplicating KV - # cache seems okay because the draft KV cache only has one layer. - # see a related copy operation in MHATokenToKVPool::move_kv_cache. - - ( - prefix_lens, - seq_lens, - last_loc, - self.num_new_pages_per_topk, - self.extend_lens, - _, - ) = get_last_loc_large_page_size_large_top_k( - batch.req_to_token_pool.req_to_token, - batch.req_pool_indices, - batch.seq_lens, - self.speculative_num_steps, - self.topk, - self.page_size, - ) - prefix_lens_cpu = batch.seq_lens_cpu - last_page_lens = prefix_lens_cpu % self.page_size - num_new_pages_per_topk = ( - last_page_lens + self.speculative_num_steps + self.page_size - 1 - ) // self.page_size - seq_lens_cpu = ( - prefix_lens_cpu // self.page_size * self.page_size - + num_new_pages_per_topk * (self.page_size * self.topk) - ) - extend_num_tokens = torch.sum((seq_lens_cpu - prefix_lens_cpu)).item() - - out_cache_loc, token_to_kv_pool_state_backup = ( - alloc_paged_token_slots_extend( - batch.tree_cache, - prefix_lens, - prefix_lens_cpu, - seq_lens, - seq_lens_cpu, - last_loc, - extend_num_tokens, - backup_state=True, - ) - ) - last_page_lens_cumsum = torch.cumsum(last_page_lens, dim=0) - duplicate_cache_len = torch.sum(last_page_lens).item() * (self.topk - 1) - target_cache_loc = torch.zeros( - duplicate_cache_len, dtype=torch.int32, device=self.device - ) - source_cache_loc = torch.zeros( - duplicate_cache_len, dtype=torch.int32, device=self.device - ) - - assign_draft_cache_locs[(num_seqs,)]( - batch.req_pool_indices, - batch.req_to_token_pool.req_to_token, - batch.seq_lens, - self.extend_lens, - self.num_new_pages_per_topk, - out_cache_loc, - source_cache_loc, - target_cache_loc, - last_page_lens_cumsum, - duplicate_cache_len, - batch.req_to_token_pool.req_to_token.shape[1], - self.topk, - self.speculative_num_steps, - self.page_size, - next_power_of_2(num_seqs), - next_power_of_2(self.speculative_num_steps), - ) - - if self.page_size > 1 and self.topk > 1: - # Remove padded slots - out_cache_loc = out_cache_loc[ - : num_seqs * self.topk * self.speculative_num_steps - ] - - batch.out_cache_loc = out_cache_loc - batch.seq_lens_sum = torch.sum(batch.seq_lens).item() - batch.return_hidden_states = False - spec_info.positions = batch.seq_lens.repeat_interleave(self.topk, dim=0) - self.token_to_kv_pool_allocator.restore_state(token_to_kv_pool_state_backup) + # FIXME: migrate multi-layer eagle worker to eagle worker + return EAGLEWorker._draft_preprocess_decode(self, batch) def _draft_preprocess_idle(self, batch: ScheduleBatch): - batch.spec_info = EagleDraftInput.create_idle_input( - device=self.device, - hidden_size=self.model_config.hidden_size, - dtype=self.model_config.dtype, - topk=self.topk * self.speculative_num_steps, - capture_hidden_mode=CaptureHiddenMode.LAST, - ) + from sglang.srt.speculative.eagle_worker import EAGLEWorker + + # FIXME: migrate multi-layer eagle worker to eagle worker + return EAGLEWorker._draft_preprocess_idle(self, batch) def draft(self, batch: ScheduleBatch): # Parse args diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 425f78ac5..44bb2f0de 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -27,11 +27,7 @@ from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, Forw from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.base_spec_worker import BaseDraftWorker, BaseSpecWorker from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput -from sglang.srt.speculative.eagle_info_v2 import ( - assign_extend_cache_locs, - fill_accepted_out_cache_loc, - fill_new_verified_id, -) +from sglang.srt.speculative.eagle_info_v2 import fill_new_verified_id from sglang.srt.speculative.eagle_utils import TreeMaskMode, build_tree_kernel_efficient from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import ( MultiLayerEagleMultiStepDraftExtendCudaGraphRunner, @@ -46,7 +42,7 @@ from sglang.srt.speculative.spec_utils import ( draft_tp_context, select_top_k_tokens, ) -from sglang.srt.utils.common import empty_context, fast_topk, next_power_of_2 +from sglang.srt.utils.common import empty_context, fast_topk if TYPE_CHECKING: from sglang.srt.model_executor.model_runner import ModelRunnerOutput @@ -129,7 +125,6 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker): ) # Alias for better readability - # self.draft_runner = self.draft_worker.model_runner self.draft_runner_list = self.draft_worker.model_runner_list self.init_lm_head() @@ -709,47 +704,3 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker): next_draft_input=next_draft_input, accept_lens=accept_length, ) - - def move_accepted_tokens_to_target_kvcache( - self, - batch: ModelWorkerBatch, - accept_index: torch.Tensor, - accept_length: torch.Tensor, - ): - """ - Move accepted tokens to the target KV cache. - - Args: - batch: The batch to run. - accept_index: The index of the accepted tokens. - accept_length: The length of the accepted tokens. - """ - bs = len(batch.seq_lens) - size = bs * self.speculative_num_draft_tokens - - tgt_cache_loc = torch.zeros( - size, - dtype=torch.int64, - device=self.device, - ) - accepted_out_cache_loc = torch.zeros( - size, dtype=torch.int64, device=self.device - ) - assign_extend_cache_locs[(bs,)]( - batch.req_pool_indices, - self.req_to_token_pool.req_to_token, - batch.seq_lens, - batch.seq_lens + accept_length, - tgt_cache_loc, - self.req_to_token_pool.req_to_token.shape[1], - next_power_of_2(bs), - ) - fill_accepted_out_cache_loc[(size,)]( - accept_index, - batch.out_cache_loc, - accepted_out_cache_loc, - next_power_of_2(size), - ) - self.token_to_kv_pool_allocator.get_kvcache().move_kv_cache( - tgt_cache_loc, accepted_out_cache_loc - )