[Scheduler] Gate SWA eviction on accumulated tokens (#33404)
This commit is contained in:
@@ -3252,14 +3252,23 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
||||
)
|
||||
|
||||
eviction_interval = max(1, envs.SGLANG_SWA_EVICTION_INTERVAL.get())
|
||||
swa_maintenance_step = (self.forward_iter or 0) % eviction_interval == 0
|
||||
self.token_to_kv_pool_allocator.free_group_begin()
|
||||
for idx, req in enumerate(self.reqs):
|
||||
if self.forward_mode.is_decode():
|
||||
# We set evict_swa condition here with two reasons:
|
||||
# 1. In overlap scheduler, we cannot evict swa when req.decode_batch_idx == 0 since the prev extend batch is still running.
|
||||
# 2. Evict swa every eviction_interval iterations to reduce the overhead.
|
||||
if swa_maintenance_step and req.decode_batch_idx >= 1:
|
||||
# 2. Evict only once >= eviction_interval tokens have slid
|
||||
# out of the window, amortizing eviction work while keeping
|
||||
# each request's overshoot within the interval the pool
|
||||
# budget reserves. Gating on accumulated tokens (rather
|
||||
# than an iteration-counter phase) cannot starve because
|
||||
# seqlen progress is monotonic per KV handle.
|
||||
if (
|
||||
req.decode_batch_idx >= 1
|
||||
and req.kv is not None
|
||||
and req.seqlen - 1 - sliding_window_size
|
||||
>= req.kv.swa_evicted_seqlen + eviction_interval
|
||||
):
|
||||
self._evict_swa(req, req.seqlen - 1)
|
||||
|
||||
# DSV4-NPU only (no-op elsewhere): the small paged compress-state
|
||||
|
||||
Reference in New Issue
Block a user