Support duplicate cache in eagle draft model

python/sglang/srt/speculative/eagle_utils.py

@@ -744,6 +744,9 @@ def assign_draft_cache_locs(
     extend_lens,
     num_new_pages_per_topk,
     out_cache_loc,
+    source_cache_loc,
+    target_cache_loc,
+    last_page_lens_cumsum,
     pool_len: tl.constexpr,
     topk: tl.constexpr,
     speculative_num_steps: tl.constexpr,
@@ -776,39 +779,67 @@ def assign_draft_cache_locs(
     if page_size == 1 or topk == 1:
         return
 
-    # Part 2: Copy the indices for the last partial page
+    # Part 2: Copy indices into source_cache_loc and target_cache_loc
+    # Expected output: src:[8,9,10,8,9,10...] tgt:[16,17,18,24,25,26...]
     prefix_len = tl.load(seq_lens + pid)
     last_page_len = prefix_len % page_size
     offsets = tl.arange(0, page_size)
     mask = offsets < last_page_len
     num_new_pages_per_topk_ = tl.load(num_new_pages_per_topk + pid)
     prefix_base = token_pool + prefix_len - last_page_len
-
-    for topk_id in range(topk):
-        value = tl.load(prefix_base + offsets, mask=mask)
+    src_indices = tl.load(prefix_base + offsets, mask=mask)
+    last_page_lens_cumsum_ = tl.load(last_page_lens_cumsum + pid)
+    # Skip the first one since no copy is needed
+    for topk_id in range(1, topk):
         tl.store(
+            source_cache_loc
+            + (topk - 1) * (last_page_lens_cumsum_ - last_page_len)
+            + (topk_id - 1) * last_page_len
+            + offsets,
+            src_indices,
+            mask=mask,
+        )
+        tgt_indices = tl.load(
             prefix_base + topk_id * num_new_pages_per_topk_ * page_size + offsets,
-            value,
             mask=mask,
         )
-
-    # Part 3: Remove the padding in out_cache_loc
-    iter_offest = tl.arange(0, iter_upper)
+        tl.store(
+            target_cache_loc
+            + (topk - 1) * (last_page_lens_cumsum_ - last_page_len)
+            + (topk_id - 1) * last_page_len
+            + offsets,
+            tgt_indices,
+            mask=mask,
+        )
+    # Part 3: Copy and remove the used indices for duplication
+    # speculative_num_steps=5, page_size=4, num_new_pages_per_topk_=2, last_page_len=1
+    #  - xxxxx .. | - xxxxx .. |
+    #   topk=0        topk=1
+    #  "-" means prefix tokens
+    #  "x" means speculative draft tokens
+    #  "." means padded tokens
+    # we only want to copy the "x" part.
+    iter_offset = tl.arange(0, iter_upper)
     for topk_id in range(topk):
+        mask_upper = iter_offset < (speculative_num_steps + last_page_len)
+        mask_lower = iter_offset >= last_page_len
+        combined_mask = mask_upper & mask_lower
         indices = tl.load(
-            prefix_base
-            + topk_id * num_new_pages_per_topk_ * page_size
-            + last_page_len
-            + iter_offest,
-            mask=iter_offest < speculative_num_steps,
+            prefix_base + topk_id * num_new_pages_per_topk_ * page_size + iter_offset,
+            mask=combined_mask,
+            other=0,
         )
+        padding_len = (iter_upper - speculative_num_steps) * pid * topk
+        all_len = pid * num_new_pages_per_topk_ * page_size * topk
+        ptr_offset = all_len - padding_len
         tl.store(
             out_cache_loc
-            + pid * topk * speculative_num_steps
+            - last_page_len
+            + ptr_offset
             + topk_id * speculative_num_steps
-            + iter_offest,
+            + iter_offset,
             indices,
-            mask=iter_offest < speculative_num_steps,
+            mask=combined_mask,
         )
 
 

python/sglang/srt/speculative/eagle_worker.py

@@ -427,21 +427,13 @@ def _draft_preprocess_decode(self, batch: ScheduleBatch):
                 #  "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,
+                    last_page_lens,
                 ) = get_last_loc_large_page_size_large_top_k(
                     batch.req_to_token_pool.req_to_token,
                     batch.req_pool_indices,
@@ -450,7 +442,6 @@ def _draft_preprocess_decode(self, batch: ScheduleBatch):
                     self.topk,
                     self.page_size,
                 )
-
                 # TODO(lmzheng): remove this device sync
                 extend_num_tokens = torch.sum(self.extend_lens).item()
 
@@ -463,6 +454,22 @@ def _draft_preprocess_decode(self, batch: ScheduleBatch):
                     backup_state=True,
                 )
             )
+        if self.page_size > 1 and self.topk > 1:
+            last_page_lens_cumsum = torch.cumsum(last_page_lens, dim=0)
+            duplicate_cache_len = torch.sum(last_page_lens) * (self.topk - 1)
+            # TODO: Remove device sync here
+            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
+            )
+        else:
+            # When source_cache_loc is not needed, simply skip
+            duplicate_cache_len = 0
+            last_page_lens_cumsum = torch.empty(0, dtype=torch.int32, device=self.device)
+            source_cache_loc = torch.empty(0, dtype=torch.int32, device=self.device)
+            target_cache_loc = torch.empty(0, dtype=torch.int32, device=self.device)
 
         assign_draft_cache_locs[(num_seqs,)](
             batch.req_pool_indices,
@@ -471,6 +478,9 @@ def _draft_preprocess_decode(self, batch: ScheduleBatch):
             self.extend_lens,
             self.num_new_pages_per_topk,
             out_cache_loc,
+            source_cache_loc,
+            target_cache_loc,
+            last_page_lens_cumsum,
             batch.req_to_token_pool.req_to_token.shape[1],
             self.topk,
             self.speculative_num_steps,
@@ -480,6 +490,10 @@ def _draft_preprocess_decode(self, batch: ScheduleBatch):
         )
 
         if self.page_size > 1 and self.topk > 1:
+            if duplicate_cache_len > 0:
+                self.draft_model_runner.token_to_kv_pool.move_kv_cache(
+                    target_cache_loc, source_cache_loc
+                )
             # Remove padded slots
             out_cache_loc = out_cache_loc[
                 : num_seqs * self.topk * self.speculative_num_steps
@@ -533,6 +547,9 @@ def draft(self, batch: ScheduleBatch):
             # Run forward steps
             score_list, token_list, parents_list = self.draft_forward(forward_batch)
 
+        # Should this line be here?
+        # self.token_to_kv_pool_allocator.restore_state(self.token_to_kv_pool_state_backup)
+
         if batch.forward_mode.is_idle():
             return EagleVerifyInput.create_idle_input(
                 self.topk,
@@ -957,4 +974,11 @@ def get_last_loc_large_page_size_large_top_k(
         prefix_lens,
     )
 
-    return prefix_lens, seq_lens, last_loc, num_new_pages_per_topk, extend_lens
+    return (
+        prefix_lens,
+        seq_lens,
+        last_loc,
+        num_new_pages_per_topk,
+        extend_lens,
+        last_page_lens,
+    )