chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Refactor prefix cache type checking (#17028)

Browse Source
This commit is contained in:
Ke Bao
2026-01-15 11:28:13 +08:00
committed by GitHub
parent 6b065298b5
commit 7f8a58fffb
12 changed files with 71 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/srt/speculative/eagle_info_v2.py
View File

@@ -10,7 +10,6 @@ import triton.language as tl
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.managers.schedule_batch import ModelWorkerBatch, ScheduleBatch
from sglang.srt.mem_cache.chunk_cache import SWAChunkCache
from sglang.srt.mem_cache.common import (
alloc_paged_token_slots_extend,
alloc_token_slots,
@@ -80,7 +79,7 @@ def assign_draft_cache_locs_page_size_1(
@dataclass
class EagleDraftInputV2Mixin:
def prepare_for_decode(self: EagleDraftInput, batch: ScheduleBatch):
if isinstance(batch.tree_cache, SWAChunkCache):
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)

3
python/sglang/srt/speculative/eagle_worker.py
View File

@@ -19,7 +19,6 @@ from sglang.srt.managers.io_struct import UpdateWeightsFromTensorReqInput
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,
@@ -375,7 +374,7 @@ class EAGLEWorker(TpModelWorker):
)
def _draft_preprocess_decode(self, batch: ScheduleBatch):
if isinstance(batch.tree_cache, SWAChunkCache):
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)