[KV Cache] Hold a finished turn's prefix until its next turn is admitted
With a prefix unit finer than the block, a hybrid follow-up resumes from a single GDN state at the previous prompt's tail, and that key sits in the top attention block. Losing any one block of that prefix is a total miss. The waiting-request pins only protect a hit that is still whole when the follow-up is first looked up. Tracing the remaining misses under a long-context multi-turn agent load shows they lost their top attention block before that, in the gap between the previous turn finishing and the next turn arriving. In that gap nothing referenced the prefix. LRU then consumed it ahead of dead blocks freed later (other conversations' handed-over CoW sources and superseded states). A finished turn (stop, length, repetition, or an abort after decoding) now keeps one extra ref on each of its hashed blocks, including its GDN holds, as a grace pin. The ref is taken on both free paths, the immediate free and the deferred pop_blocks_for_free behind an in-flight step, so the prefix never passes through the free queue. A grace is released when: - the next turn is admitted on it. That is the grace holding the deepest hit block of every group, with a hit that reaches its prompt to within two prefix units, so a new turn sharing only a system prompt releases nothing. The table takes the blocks over, with no double ref. - anyone needs blocks. Graces yield before waiter pins, oldest first and one per failed allocation, so they never cause a preemption. - the bound (16) is exceeded. The oldest grace goes. - reset_prefix_cache runs. A head waiter whose admission fails pins its own hit first. It then yields graces and other pins one at a time, and only when everything yieldable covers its shortfall. Otherwise yielding would drop other conversations' prefixes and still not admit it. The shortfall is measured after the pin: the pin takes the proof block past the hit (and any hit block the lookup did not return) out of the free queue, so the shortfall of the failed attempt can be short by those blocks and the guard would yield another conversation's grace for nothing. Its own pinned blocks outside the hit it is admitted on, the proof block past it, go back first: once the hit is taken over they protect nothing, and they cost nobody else anything. At an exact fit they are the blocks it needs, and nothing else ever yields a head waiter's own pin, so without this a request that needs the whole pool waited forever. While any pin or grace is held the scheduler reports no common prefix to cascade attention. Those refs can make a block's count equal the number of running requests when only some of them share it, and cascade attention would then read it for all of them. Graces and waiter pins now follow a producer's GDN hash into its CoW copy, as holds already do. Otherwise they would keep a second ref on a running state block that its owner must own exclusively. Graces apply only to fine-grained hybrid hits (coordinator.enable_partial_hash_hits) and are off with a KV connector; a block-sized unit (the 4096 lane) is unchanged. When graces are active the scheduler logs one startup line saying so, with the bound, the hash unit and the block size. Tests: the arrival-gap reproduction fails on the previous commit (hit 0 instead of 168, both free paths). It passes here. Further tests cover the handover refcounts, yield order and bound, finish statuses, a running request at an exactly full pool, yieldable accounting, moved-hash following, the pin's extras contract, and six runs on the real scheduler: among them a head waiter whose pin raises its shortfall (it gives its proof block back and the one grace that covers the rest yields), a request that needs the whole pool (admitted on its own proof block; it waited forever before), and the common-prefix count under a grace (none reported; the count returns once the pins are gone). Co-Authored-By: Hinata <hinata@idihore.id> Signed-off-by: MOVZX <movzx@yahoo.com>
This commit is contained in:
@@ -237,6 +237,37 @@ class KVCacheManager:
|
||||
# whenever a running request would otherwise fail to allocate
|
||||
# (``yield_pinned_hits``), so a pin never causes a preemption.
|
||||
self._pinned_hits: dict[str, list[KVCacheBlock]] = {}
|
||||
# A finished request's hashed blocks (its attention prefix and
|
||||
# the GDN states it holds) stay referenced as a "grace" pin until the
|
||||
# conversation's next turn is admitted on them. With waiter pins alone,
|
||||
# every remaining total miss had lost its top attention block BEFORE
|
||||
# the next turn was first looked up -- in the gap between the
|
||||
# previous turn's free and the next turn's arrival, or while its
|
||||
# admission kept failing -- when nothing referenced the prefix and
|
||||
# the running set's allocations popped it; once the hybrid hit is 0
|
||||
# a waiter pin has nothing to pin. Grace pins yield first (oldest
|
||||
# first) whenever anyone needs blocks, so they never cost a running
|
||||
# request or an admission anything a free queue would have given.
|
||||
# A grace is one extra ref per block, taken on whichever free path
|
||||
# the request takes (immediate or deferred behind an in-flight step).
|
||||
self._grace_pins: dict[str, list[KVCacheBlock]] = {}
|
||||
self._grace_block_ids: dict[str, frozenset[int]] = {}
|
||||
self._grace_prompt_tokens: dict[str, int] = {}
|
||||
# Only for fine-grained hybrid hits (align Mamba hashed finer than its
|
||||
# block): there the next turn's key sits in the top attention block
|
||||
# and one lost block is a total miss; at a block-sized unit the top
|
||||
# block is a spare and nothing changes. The scheduler also turns
|
||||
# graces off when a KV connector is configured (as for waiter pins).
|
||||
self.enable_grace_pins = (
|
||||
self.enable_caching and self.coordinator.enable_partial_hash_hits
|
||||
)
|
||||
# Shortfall (blocks) of the last allocate_slots that returned None.
|
||||
self.last_alloc_shortfall = 0
|
||||
# A producer that moves its live GDN tail's hash to a CoW copy keeps
|
||||
# writing the source; pins and graces follow the hash to the copy.
|
||||
for manager in self.coordinator.single_type_managers:
|
||||
if isinstance(manager, MambaManager):
|
||||
manager.moved_hash_listeners.append(self._pins_follow_moved_hash)
|
||||
|
||||
@property
|
||||
def usage(self) -> float:
|
||||
@@ -551,6 +582,9 @@ class KVCacheManager:
|
||||
)
|
||||
required_blocks = num_blocks_to_allocate + watermark_blocks
|
||||
if required_blocks > self.block_pool.get_num_free_blocks():
|
||||
self.last_alloc_shortfall = (
|
||||
required_blocks - self.block_pool.get_num_free_blocks()
|
||||
)
|
||||
return None
|
||||
|
||||
num_tokens_main_model = total_computed_tokens + num_new_tokens
|
||||
@@ -591,7 +625,9 @@ class KVCacheManager:
|
||||
required_blocks = num_blocks_to_allocate + watermark_blocks
|
||||
if required_blocks > available_blocks:
|
||||
# Cannot allocate new blocks
|
||||
self.last_alloc_shortfall = required_blocks - available_blocks
|
||||
return None
|
||||
self.last_alloc_shortfall = 0
|
||||
|
||||
if (
|
||||
new_computed_block_list is not self.empty_kv_cache_blocks.blocks
|
||||
@@ -618,6 +654,10 @@ class KVCacheManager:
|
||||
# remote. Update state for locally cached blocks.
|
||||
if delay_cache_blocks:
|
||||
self.release_pinned_hit(request.request_id)
|
||||
if self._grace_pins and num_new_computed_tokens > 0:
|
||||
self._release_grace_for_hit(
|
||||
new_computed_block_list, num_local_computed_tokens
|
||||
)
|
||||
return self.create_kv_cache_blocks(new_blocks)
|
||||
|
||||
# NOTE(woosuk): We want to commit (cache) up to num_local_computed_tokens
|
||||
@@ -631,8 +671,13 @@ class KVCacheManager:
|
||||
)
|
||||
self.coordinator.cache_blocks(request, num_tokens_to_cache)
|
||||
|
||||
# Admitted: the table now references the hit, the pin is redundant.
|
||||
# Admitted: the table now references the hit, the pins are redundant.
|
||||
self.release_pinned_hit(request.request_id)
|
||||
if self._grace_pins and num_new_computed_tokens > 0:
|
||||
# ... and so does the grace pin of the turn it resumes.
|
||||
self._release_grace_for_hit(
|
||||
new_computed_block_list, num_local_computed_tokens
|
||||
)
|
||||
return self.create_kv_cache_blocks(new_blocks)
|
||||
|
||||
def free(self, request: Request) -> None:
|
||||
@@ -645,7 +690,220 @@ class KVCacheManager:
|
||||
|
||||
"""
|
||||
self.release_pinned_hit(request.request_id)
|
||||
self.coordinator.free(request.request_id)
|
||||
if not self._grace_eligible(request):
|
||||
self.coordinator.free(request.request_id)
|
||||
return
|
||||
# Same blocks and order as coordinator.free (the deferred path frees
|
||||
# this list as-is too); the grace takes its ref before the table's
|
||||
# goes, so the hashed blocks never pass through the free queue.
|
||||
blocks = self.coordinator.pop_blocks_for_free(request.request_id)
|
||||
self._take_grace_pin(request, blocks)
|
||||
self.block_pool.free_blocks(blocks)
|
||||
|
||||
# Finished requests whose hashed blocks stay pinned for the next turn
|
||||
# (oldest released first beyond this). One per conversation is the
|
||||
# need; twice the deployed max_num_seqs.
|
||||
MAX_GRACE_PINS = 16
|
||||
|
||||
@property
|
||||
def num_grace_pins(self) -> int:
|
||||
return len(self._grace_pins)
|
||||
|
||||
@property
|
||||
def num_grace_pinned_blocks(self) -> int:
|
||||
return sum(len(blocks) for blocks in self._grace_pins.values())
|
||||
|
||||
def grace_pinned_blocks(self, request_id: str) -> list[KVCacheBlock]:
|
||||
return list(self._grace_pins.get(request_id, ()))
|
||||
|
||||
def _grace_eligible(self, request: Request) -> bool:
|
||||
"""A turn that finished with output (stop, length, repetition, or
|
||||
an abort after decoding) may have a next turn that resumes it. A
|
||||
preemption (still RUNNING), an aborted waiter (no output), an error
|
||||
or an ignored request leaves nothing worth holding."""
|
||||
return (
|
||||
self.enable_grace_pins
|
||||
and self.enable_caching
|
||||
and self.prefix_cache_lookup_enabled(request)
|
||||
and request.is_finished()
|
||||
and request.status
|
||||
not in (RequestStatus.FINISHED_ERROR, RequestStatus.FINISHED_IGNORED)
|
||||
and request.num_output_tokens > 0
|
||||
)
|
||||
|
||||
def _take_grace_pin(self, request: Request, blocks: list[KVCacheBlock]) -> int:
|
||||
"""Take one extra ref on every hashed block of a finished request's
|
||||
just-popped blocks (``blocks`` is in eviction order: each group's
|
||||
table tail first, the held GDN states last). The caller then drops
|
||||
the table refs, now or after the in-flight fence; the hashed blocks
|
||||
stay out of the free queue until the grace is released. Unhashed
|
||||
blocks (running state, MTP scratch, an unkeyed last block) go back
|
||||
with the table refs as before. Returns the number of blocks held."""
|
||||
keep: list[KVCacheBlock] = []
|
||||
seen: set[int] = set()
|
||||
for block in blocks:
|
||||
if block.is_null or block.block_hash is None or id(block) in seen:
|
||||
continue
|
||||
# Still referenced by the request's table or hold: the grace ref
|
||||
# never moves a block in or out of a free queue (any pool).
|
||||
assert block.ref_cnt > 0
|
||||
seen.add(id(block))
|
||||
keep.append(block)
|
||||
if not keep:
|
||||
return 0
|
||||
request_id = request.request_id
|
||||
self.release_grace_pin(request_id) # a repeated free: replace
|
||||
for block in keep:
|
||||
block.ref_cnt += 1
|
||||
self._grace_pins[request_id] = keep
|
||||
self._grace_block_ids[request_id] = frozenset(seen)
|
||||
self._grace_prompt_tokens[request_id] = request.num_prompt_tokens
|
||||
while len(self._grace_pins) > self.MAX_GRACE_PINS:
|
||||
self.release_grace_pin(next(iter(self._grace_pins)))
|
||||
return len(keep)
|
||||
|
||||
def release_grace_pin(self, request_id: str) -> int:
|
||||
"""Return a finished request's grace-pinned blocks to the pool (in
|
||||
the eviction order a free would have used). Returns the count."""
|
||||
blocks = self._grace_pins.pop(request_id, None)
|
||||
self._grace_block_ids.pop(request_id, None)
|
||||
self._grace_prompt_tokens.pop(request_id, None)
|
||||
if not blocks:
|
||||
return 0
|
||||
self.block_pool.free_blocks(blocks)
|
||||
return len(blocks)
|
||||
|
||||
def _release_grace_for_hit(
|
||||
self,
|
||||
computed_blocks: tuple[Sequence[KVCacheBlock], ...],
|
||||
num_hit_tokens: int,
|
||||
) -> None:
|
||||
"""An admitted hit now references its blocks through the table, so
|
||||
the grace of the turn it resumes has done its job. That turn is a
|
||||
grace that holds the deepest hit block of every group whose deepest
|
||||
hit block any grace holds (the GDN state it resumes from, the
|
||||
attention block holding the key), and whose prompt the hit reaches
|
||||
to within two prefix units (the EAGLE-dropped tail key). A new turn
|
||||
that only shares a system prompt with finished conversations
|
||||
releases none of theirs."""
|
||||
candidates: set[str] | None = None
|
||||
for group in computed_blocks:
|
||||
deepest = next((b for b in reversed(group) if not b.is_null), None)
|
||||
if deepest is None:
|
||||
continue
|
||||
owners = {
|
||||
request_id
|
||||
for request_id, block_ids in self._grace_block_ids.items()
|
||||
if id(deepest) in block_ids
|
||||
}
|
||||
if owners:
|
||||
candidates = owners if candidates is None else candidates & owners
|
||||
if not candidates:
|
||||
return
|
||||
reach = num_hit_tokens + 2 * self.block_pool.hash_block_size
|
||||
for request_id in candidates:
|
||||
if reach >= self._grace_prompt_tokens.get(request_id, 0):
|
||||
self.release_grace_pin(request_id)
|
||||
|
||||
def _pins_follow_moved_hash(
|
||||
self, source_block: KVCacheBlock, cow_block: KVCacheBlock
|
||||
) -> None:
|
||||
"""A Mamba producer moved ``source_block``'s hashes to ``cow_block``
|
||||
and keeps writing ``source_block``. A grace or waiter pin on the
|
||||
source follows to the copy (as a hold does), so it protects the
|
||||
hashed state and never leaves an extra ref on a block its owner must
|
||||
hold exclusively (speculative relocation asserts on that)."""
|
||||
for request_id, blocks in self._grace_pins.items():
|
||||
if id(source_block) in self._grace_block_ids[request_id]:
|
||||
self._swap_pinned_block(blocks, source_block, cow_block)
|
||||
self._grace_block_ids[request_id] = frozenset(
|
||||
id(block) for block in blocks
|
||||
)
|
||||
for blocks in self._pinned_hits.values():
|
||||
self._swap_pinned_block(blocks, source_block, cow_block)
|
||||
|
||||
def _swap_pinned_block(
|
||||
self,
|
||||
blocks: list[KVCacheBlock],
|
||||
source_block: KVCacheBlock,
|
||||
cow_block: KVCacheBlock,
|
||||
) -> None:
|
||||
for i, block in enumerate(blocks):
|
||||
if block is source_block:
|
||||
if any(other is cow_block for other in blocks):
|
||||
del blocks[i]
|
||||
else:
|
||||
assert cow_block.ref_cnt > 0
|
||||
cow_block.ref_cnt += 1
|
||||
blocks[i] = cow_block
|
||||
self.block_pool.free_blocks([source_block])
|
||||
return
|
||||
|
||||
def num_yieldable_blocks(
|
||||
self,
|
||||
keep: str | None = None,
|
||||
keep_hit: tuple[Sequence[KVCacheBlock], ...] | None = None,
|
||||
) -> int:
|
||||
"""Blocks that would return to the free queue if every grace and
|
||||
every waiter pin except ``keep``'s were released now. With
|
||||
``keep_hit`` (the hit ``keep`` is being admitted on), ``keep``'s
|
||||
pinned blocks outside that hit (the proof block past it) count too:
|
||||
``release_pinned_extras`` can give them back. Blocks of that hit
|
||||
never count: freed, they are taken straight back by the admission
|
||||
(a hit that grew past the pin snapshot would otherwise be counted
|
||||
through the grace that still holds it)."""
|
||||
pins: dict[int, list] = {}
|
||||
holders = [*self._grace_pins.values()] + [
|
||||
blocks
|
||||
for request_id, blocks in self._pinned_hits.items()
|
||||
if request_id != keep
|
||||
]
|
||||
hit_ids: set[int] = set()
|
||||
if keep is not None and keep_hit is not None:
|
||||
holders.append(self._pinned_extras(keep, keep_hit))
|
||||
hit_ids = {id(block) for group in keep_hit for block in group}
|
||||
for blocks in holders:
|
||||
for block in blocks:
|
||||
entry = pins.setdefault(id(block), [block, 0])
|
||||
entry[1] += 1
|
||||
return sum(
|
||||
1
|
||||
for block_id, (block, n) in pins.items()
|
||||
if block.ref_cnt == n and block_id not in hit_ids
|
||||
)
|
||||
|
||||
def _pinned_extras(
|
||||
self, request_id: str, hit: tuple[Sequence[KVCacheBlock], ...]
|
||||
) -> list[KVCacheBlock]:
|
||||
hit_ids = {id(block) for group in hit for block in group}
|
||||
return [
|
||||
block
|
||||
for block in self._pinned_hits.get(request_id, ())
|
||||
if id(block) not in hit_ids
|
||||
]
|
||||
|
||||
def release_pinned_extras(
|
||||
self, request_id: str, hit: tuple[Sequence[KVCacheBlock], ...]
|
||||
) -> int:
|
||||
"""Drop from the request's pin the blocks outside ``hit``: the proof
|
||||
block past it, and any hit block a later lookup added. The request
|
||||
is being admitted on ``hit`` now, so they protect nothing it still
|
||||
needs; they cost it blocks. Returns the number dropped."""
|
||||
extras = self._pinned_extras(request_id, hit)
|
||||
if not extras:
|
||||
return 0
|
||||
extra_ids = {id(block) for block in extras}
|
||||
rest = [
|
||||
block
|
||||
for block in self._pinned_hits[request_id]
|
||||
if id(block) not in extra_ids
|
||||
]
|
||||
if rest:
|
||||
self._pinned_hits[request_id] = rest
|
||||
else:
|
||||
del self._pinned_hits[request_id]
|
||||
self.block_pool.free_blocks(extras)
|
||||
return len(extras)
|
||||
|
||||
# Waiting requests whose prefix-cache hit is pinned at once (arrival
|
||||
# order). Pins never block anyone (they yield before a running request
|
||||
@@ -671,11 +929,10 @@ class KVCacheManager:
|
||||
|
||||
Known gap: a pin is a snapshot of the hit at pin time. It is not
|
||||
refreshed if the hit grows while the request waits (a sibling
|
||||
computing more of the same prefix) or if a producer moves the hash of
|
||||
a pinned live block into its CoW copy (a resume at a running
|
||||
producer's tail, unreachable under the EAGLE drop); the admission
|
||||
lookup still finds whatever is cached then, so the only cost is a
|
||||
smaller hit, never a wrong one."""
|
||||
computing more of the same prefix); the admission lookup still finds
|
||||
whatever is cached then, so the only cost is a smaller hit, never a
|
||||
wrong one. A producer moving the hash of a pinned live GDN block into
|
||||
its CoW copy is followed (``_pins_follow_moved_hash``)."""
|
||||
request_id = request.request_id
|
||||
if (
|
||||
request_id in self._pinned_hits
|
||||
@@ -734,13 +991,21 @@ class KVCacheManager:
|
||||
self, keep: str | None = None, limit: int | None = None
|
||||
) -> int:
|
||||
"""Drop pins (except ``keep``'s) so a running request can allocate,
|
||||
or a waiter ahead in line can be admitted: the most recently pinned
|
||||
waiters first, at most ``limit`` of them per call (all by default;
|
||||
the scheduler passes 1 and retries, so a decode step that needs a
|
||||
block costs one waiter its pin, not every waiter). Returns the number
|
||||
of blocks released; 0 means there was nothing to yield."""
|
||||
or a waiter ahead in line can be admitted: grace pins of finished
|
||||
turns first (oldest first), then the most recently pinned waiters, at
|
||||
most ``limit`` pins per call (all by default; the scheduler passes 1
|
||||
and retries, so a decode step that needs a block costs one pin, not
|
||||
every pin). Returns the number of blocks released; 0 means there was
|
||||
nothing to yield."""
|
||||
released = 0
|
||||
yielded = 0
|
||||
# Grace pins first (oldest finished conversation first), then the
|
||||
# waiters' pins (most recently pinned first).
|
||||
for request_id in list(self._grace_pins):
|
||||
if limit is not None and yielded >= limit:
|
||||
return released
|
||||
released += self.release_grace_pin(request_id)
|
||||
yielded += 1
|
||||
for request_id in reversed(list(self._pinned_hits)):
|
||||
if request_id == keep:
|
||||
continue
|
||||
@@ -787,7 +1052,12 @@ class KVCacheManager:
|
||||
# A pin is never written by an in-flight step (the waiter has no
|
||||
# table), so it is released now, not with the deferred blocks.
|
||||
self.release_pinned_hit(request.request_id)
|
||||
return self.coordinator.pop_blocks_for_free(request.request_id)
|
||||
blocks = self.coordinator.pop_blocks_for_free(request.request_id)
|
||||
# A finished turn whose last step is still in flight: the grace takes
|
||||
# its ref now; the table refs go when the fence drains.
|
||||
if self._grace_eligible(request):
|
||||
self._take_grace_pin(request, blocks)
|
||||
return blocks
|
||||
|
||||
def evict_blocks(self, block_ids: set[int]) -> None:
|
||||
"""Evict blocks from the prefix cache by their block IDs.
|
||||
@@ -850,6 +1120,12 @@ class KVCacheManager:
|
||||
group.
|
||||
|
||||
"""
|
||||
if self._pinned_hits or self._grace_pins:
|
||||
# A pinned or graced block carries a ref that is not a running
|
||||
# request's, so the ref count can equal the running count for a
|
||||
# block only some of them share. Report no common prefix (cascade
|
||||
# attention is an optimisation; a wrong prefix would be read).
|
||||
return [0] * len(self.kv_cache_config.kv_cache_groups)
|
||||
return self.coordinator.get_num_common_prefix_blocks(running_request_id)
|
||||
|
||||
def take_events(self) -> list[KVCacheEvent]:
|
||||
|
||||
@@ -366,6 +366,9 @@ class Scheduler(SchedulerInterface):
|
||||
# Bind after construction so connectors can access the cache manager.
|
||||
if self.connector is not None:
|
||||
self.connector.bind_kv_cache_manager(self.kv_cache_manager)
|
||||
# Finished turns' grace pins are off with a connector, as
|
||||
# waiter pins are (_pin_waiting_hits).
|
||||
self.kv_cache_manager.enable_grace_pins = False
|
||||
# Ring blocks an async KV load is about to fill. Their tag clear
|
||||
# is skipped for the same reason as zeroing (it could land after the
|
||||
# out-of-band write and wipe the producer's committed tags). Tracked
|
||||
@@ -583,6 +586,15 @@ class Scheduler(SchedulerInterface):
|
||||
self.hash_block_size,
|
||||
self.block_size,
|
||||
)
|
||||
if manager.enable_grace_pins:
|
||||
logger.info(
|
||||
"Prefix cache: a finished turn's cached prefix stays referenced "
|
||||
"until its next turn is admitted (up to %d turns; hash unit %d, "
|
||||
"block %d)",
|
||||
manager.MAX_GRACE_PINS,
|
||||
self.hash_block_size,
|
||||
self.block_size,
|
||||
)
|
||||
|
||||
def _pin_waiting_hits(self) -> None:
|
||||
"""Pin the prefix-cache hits of the requests at the front of the
|
||||
@@ -1326,12 +1338,89 @@ class Scheduler(SchedulerInterface):
|
||||
has_scheduled_reqs=bool(self.running),
|
||||
)
|
||||
|
||||
if new_blocks is None and self.kv_cache_manager.yield_pinned_hits(
|
||||
keep=request_id
|
||||
# Finished turns' grace pins (oldest first), then the
|
||||
# pins of waiters behind this one, yield to it one at a time
|
||||
# (arrival order). Its own hit is pinned first (a turn that
|
||||
# arrived this step is not pinned yet and its hit may sit in
|
||||
# a grace that yields), so only blocks it does not resume
|
||||
# from go back. And only when releasing them all would cover
|
||||
# the shortfall: otherwise yielding costs other conversations
|
||||
# their prefix and still does not admit this one.
|
||||
kv_cache_manager = self.kv_cache_manager
|
||||
yield_for_head = new_blocks is None and (
|
||||
kv_cache_manager.last_alloc_shortfall > 0
|
||||
and (
|
||||
kv_cache_manager.num_grace_pins > 0
|
||||
or kv_cache_manager.num_pinned_waiters > 0
|
||||
)
|
||||
)
|
||||
if yield_for_head:
|
||||
if (
|
||||
self.connector is None
|
||||
and request.num_computed_tokens == 0
|
||||
and kv_cache_manager.pin_waiting_hit(request) > 0
|
||||
):
|
||||
# The pin just took blocks out of the free queue that
|
||||
# the failed attempt counted as free: the proof block
|
||||
# past the hit, and any hit block the lookup above did
|
||||
# not return. Recompute the shortfall on the pinned
|
||||
# state before measuring it against the others' pins.
|
||||
new_blocks = self.kv_cache_manager.allocate_slots(
|
||||
request,
|
||||
num_new_tokens,
|
||||
num_new_computed_tokens=num_new_local_computed_tokens,
|
||||
new_computed_blocks=new_computed_blocks,
|
||||
num_lookahead_tokens=effective_lookahead_tokens,
|
||||
num_external_computed_tokens=num_external_computed_tokens,
|
||||
delay_cache_blocks=load_kv_async,
|
||||
num_encoder_tokens=num_encoder_tokens,
|
||||
full_sequence_must_fit=self.scheduler_reserve_full_isl,
|
||||
reserved_blocks=reserved_blocks,
|
||||
has_scheduled_reqs=bool(self.running),
|
||||
)
|
||||
yield_for_head = new_blocks is None and (
|
||||
kv_cache_manager.last_alloc_shortfall
|
||||
<= kv_cache_manager.num_yieldable_blocks(
|
||||
keep=request_id, keep_hit=new_computed_blocks.blocks
|
||||
)
|
||||
)
|
||||
if (
|
||||
yield_for_head
|
||||
and kv_cache_manager.release_pinned_extras(
|
||||
request_id, new_computed_blocks.blocks
|
||||
)
|
||||
> 0
|
||||
):
|
||||
# Its own pinned blocks outside the hit it is admitted
|
||||
# on (the proof block past it) go first: they cost
|
||||
# nobody else anything, and at an exact fit they are
|
||||
# the blocks it needs. Retry on them before the others
|
||||
# yield.
|
||||
new_blocks = self.kv_cache_manager.allocate_slots(
|
||||
request,
|
||||
num_new_tokens,
|
||||
num_new_computed_tokens=num_new_local_computed_tokens,
|
||||
new_computed_blocks=new_computed_blocks,
|
||||
num_lookahead_tokens=effective_lookahead_tokens,
|
||||
num_external_computed_tokens=num_external_computed_tokens,
|
||||
delay_cache_blocks=load_kv_async,
|
||||
num_encoder_tokens=num_encoder_tokens,
|
||||
full_sequence_must_fit=self.scheduler_reserve_full_isl,
|
||||
reserved_blocks=reserved_blocks,
|
||||
has_scheduled_reqs=bool(self.running),
|
||||
)
|
||||
# Coverage is re-read before every yield: a retry can fail
|
||||
# by more than the attempt before it (the full-sequence gate
|
||||
# passes, the lookahead gate does not).
|
||||
while (
|
||||
yield_for_head
|
||||
and new_blocks is None
|
||||
and kv_cache_manager.last_alloc_shortfall
|
||||
<= kv_cache_manager.num_yieldable_blocks(
|
||||
keep=request_id, keep_hit=new_computed_blocks.blocks
|
||||
)
|
||||
and kv_cache_manager.yield_pinned_hits(keep=request_id, limit=1)
|
||||
):
|
||||
# Pins of waiters behind this one yield to it
|
||||
# (arrival order); its own pin stays, those blocks are
|
||||
# its hit. One retry with them back in the free queue.
|
||||
new_blocks = self.kv_cache_manager.allocate_slots(
|
||||
request,
|
||||
num_new_tokens,
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -3,7 +3,7 @@
|
||||
import itertools
|
||||
from abc import ABC, abstractmethod
|
||||
from collections import defaultdict
|
||||
from collections.abc import Sequence
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import TYPE_CHECKING, ClassVar
|
||||
|
||||
from vllm import envs
|
||||
@@ -1705,6 +1705,12 @@ class MambaManager(SingleTypeKVCacheManager):
|
||||
# therefore hold their hashed endpoint here too: the moved tail for
|
||||
# a running request, and the hit source a new request copies from.
|
||||
self._retained_checkpoints: dict[str, tuple[int, KVCacheBlock]] = {}
|
||||
# Callbacks (source_block, cow_block) run after a running
|
||||
# request moved its live tail's hash to the CoW copy, for refs held
|
||||
# outside this manager (the KVCacheManager's grace and waiter pins).
|
||||
self.moved_hash_listeners: list[
|
||||
Callable[[KVCacheBlock, KVCacheBlock], None]
|
||||
] = []
|
||||
if self.mamba_cache_mode == "align" and self.enable_caching:
|
||||
# The A/B evidence line: boot logs must show which arm ran.
|
||||
logger.info_once(
|
||||
@@ -1801,6 +1807,8 @@ class MambaManager(SingleTypeKVCacheManager):
|
||||
self.block_pool.touch([cow_block])
|
||||
self._retained_checkpoints[request_id] = (position, cow_block)
|
||||
self.block_pool.free_blocks([source_block])
|
||||
for listener in self.moved_hash_listeners:
|
||||
listener(source_block, cow_block)
|
||||
|
||||
def pop_retained_checkpoints(self, request_id: str) -> list[KVCacheBlock]:
|
||||
if self.mamba_cache_mode != "align":
|
||||
|
||||
Reference in New Issue
Block a user