[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())
|
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()
|
self.token_to_kv_pool_allocator.free_group_begin()
|
||||||
for idx, req in enumerate(self.reqs):
|
for idx, req in enumerate(self.reqs):
|
||||||
if self.forward_mode.is_decode():
|
if self.forward_mode.is_decode():
|
||||||
# We set evict_swa condition here with two reasons:
|
# 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.
|
# 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.
|
# 2. Evict only once >= eviction_interval tokens have slid
|
||||||
if swa_maintenance_step and req.decode_batch_idx >= 1:
|
# 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)
|
self._evict_swa(req, req.seqlen - 1)
|
||||||
|
|
||||||
# DSV4-NPU only (no-op elsewhere): the small paged compress-state
|
# DSV4-NPU only (no-op elsewhere): the small paged compress-state
|
||||||
|
|||||||
Reference in New Issue
Block a user