Tiny cleanup duplicate code for multi-layer eagle worker. (#16004)

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

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
-        )