Evict swa kv cache during decoding (#17220)

python/sglang/srt/speculative/eagle_info_v2.py

@@ -79,9 +79,7 @@ def assign_draft_cache_locs_page_size_1(
 @dataclass
 class EagleDraftInputV2Mixin:
     def prepare_for_decode(self: EagleDraftInput, batch: ScheduleBatch):
-        if batch.tree_cache.supports_swa() and batch.tree_cache.is_chunk_cache():
-            for req in batch.reqs:
-                batch.tree_cache.evict_swa(req, req.seqlen - 1)
+        batch.maybe_evict_swa()
 
         from sglang.srt.speculative.spec_utils import assign_req_to_token_pool_func
 
@@ -101,6 +99,7 @@ class EagleDraftInputV2Mixin:
             nxt_kv_lens_cpu.append(r.kv_allocated_len + x)
             num_needed_tokens += x
             r.kv_allocated_len += x
+            r.decode_batch_idx += 1
 
         cur_kv_lens_cpu = torch.tensor(cur_kv_lens_cpu, dtype=torch.int32, device="cpu")
         nxt_kv_lens_cpu = torch.tensor(nxt_kv_lens_cpu, dtype=torch.int32, device="cpu")

python/sglang/srt/speculative/eagle_worker.py

@@ -374,9 +374,9 @@ class EAGLEWorker(TpModelWorker):
         )
 
     def _draft_preprocess_decode(self, batch: ScheduleBatch):
-        if batch.tree_cache.supports_swa() and batch.tree_cache.is_chunk_cache():
-            for req in batch.reqs:
-                batch.tree_cache.evict_swa(req, req.seqlen - 1)
+        batch.maybe_evict_swa()
+        for req in batch.reqs:
+            req.decode_batch_idx += 1
 
         # Parse args
         num_seqs = batch.batch_size()