35 lines
967 B
Python
35 lines
967 B
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, List
|
|
|
|
if TYPE_CHECKING:
|
|
import torch
|
|
|
|
from sglang.srt.managers.scheduler import Scheduler
|
|
|
|
|
|
class ScriptedKvPoolExhauster:
|
|
|
|
def __init__(self, scheduler: Scheduler) -> None:
|
|
self.scheduler = scheduler
|
|
self._held: List[torch.Tensor] = []
|
|
|
|
def exhaust(self, *, leave_pages: int) -> None:
|
|
allocator = self.scheduler.token_to_kv_pool_allocator
|
|
|
|
leave_tokens = leave_pages * self.scheduler.page_size
|
|
need = allocator.available_size() - leave_tokens
|
|
if need <= 0:
|
|
return
|
|
|
|
held = allocator.alloc(need)
|
|
assert (
|
|
held is not None
|
|
), f"exhaust_kv: allocator could not grab {need} tokens to create pressure"
|
|
self._held.append(held)
|
|
|
|
def release(self) -> None:
|
|
for held in self._held:
|
|
self.scheduler.token_to_kv_pool_allocator.free(held)
|
|
self._held.clear()
|