[mem_cache] Move kv_committed_len into ReqKvInfo (#37078)
This commit is contained in:
@@ -442,7 +442,7 @@ class BeamCoordinator(msgspec.Struct, kw_only=True):
|
|||||||
prefix_len=group.prompt_len,
|
prefix_len=group.prompt_len,
|
||||||
# All rows are synchronized; the leader's committed length
|
# All rows are synchronized; the leader's committed length
|
||||||
# covers the KV computed through this step.
|
# covers the KV computed through this step.
|
||||||
seq_len=group.leader.kv_committed_len,
|
seq_len=group.leader.kv.kv_committed_len,
|
||||||
)
|
)
|
||||||
group.pending_orphans.append(StagedOrphans(tick, old_map, new_map))
|
group.pending_orphans.append(StagedOrphans(tick, old_map, new_map))
|
||||||
# Length placeholder only; the DAG owns history and member rows
|
# Length placeholder only; the DAG owns history and member rows
|
||||||
|
|||||||
@@ -93,7 +93,7 @@ def free_member_rows(group, req_to_token_pool, token_to_kv_pool_allocator) -> No
|
|||||||
# release frees this decode region a second time.
|
# release frees this decode region a second time.
|
||||||
slots = req_to_token_pool.req_to_token[group.all_rows, start:end]
|
slots = req_to_token_pool.req_to_token[group.all_rows, start:end]
|
||||||
token_to_kv_pool_allocator.free(slots.flatten().unique())
|
token_to_kv_pool_allocator.free(slots.flatten().unique())
|
||||||
leader.kv_committed_len = start
|
leader.kv.kv_committed_len = start
|
||||||
leader.kv.kv_allocated_len = start
|
leader.kv.kv_allocated_len = start
|
||||||
req_to_token_pool.free_rows(group.member_rows_cpu.tolist())
|
req_to_token_pool.free_rows(group.member_rows_cpu.tolist())
|
||||||
group.member_rows = None
|
group.member_rows = None
|
||||||
|
|||||||
@@ -193,7 +193,7 @@ class DecodeReqToTokenPool:
|
|||||||
len(reusing) <= 1
|
len(reusing) <= 1
|
||||||
), "only one chunked request may reuse req_pool_idx in a batch"
|
), "only one chunked request may reuse req_pool_idx in a batch"
|
||||||
assert all(
|
assert all(
|
||||||
reqs[i].inflight_middle_chunks > 0 or reqs[i].kv_committed_len > 0
|
reqs[i].inflight_middle_chunks > 0 or reqs[i].kv.kv_committed_len > 0
|
||||||
for i in reusing
|
for i in reusing
|
||||||
), "reusing request must be chunked or have committed KV"
|
), "reusing request must be chunked or have committed KV"
|
||||||
|
|
||||||
@@ -1783,7 +1783,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
|
|||||||
), "req_pool_indices is full! There is a bug in memory estimation."
|
), "req_pool_indices is full! There is a bug in memory estimation."
|
||||||
|
|
||||||
fill_len = self._pre_alloc_fill_len(req)
|
fill_len = self._pre_alloc_fill_len(req)
|
||||||
req.kv_committed_len = fill_len
|
req.kv.kv_committed_len = fill_len
|
||||||
|
|
||||||
if prefix_len > 0:
|
if prefix_len > 0:
|
||||||
self.req_to_token_pool.write(
|
self.req_to_token_pool.write(
|
||||||
@@ -1889,7 +1889,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
|
|||||||
req.prefix_indices = (
|
req.prefix_indices = (
|
||||||
prefix_indices if prefix_len > 0 else torch.empty((0,), dtype=torch.int64)
|
prefix_indices if prefix_len > 0 else torch.empty((0,), dtype=torch.int64)
|
||||||
)
|
)
|
||||||
req.set_extend_range(total_prefix_len, req.kv_committed_len)
|
req.set_extend_range(total_prefix_len, req.kv.kv_committed_len)
|
||||||
|
|
||||||
# Return the transfer destination indices:
|
# Return the transfer destination indices:
|
||||||
if self.scheduler.enable_hisparse:
|
if self.scheduler.enable_hisparse:
|
||||||
@@ -2634,8 +2634,10 @@ class SchedulerDisaggregationDecodeMixin:
|
|||||||
# Truncate fill_len to kv_committed_len so cache_unfinished_req
|
# Truncate fill_len to kv_committed_len so cache_unfinished_req
|
||||||
# only sees committed KV (full array includes one uncommitted
|
# only sees committed KV (full array includes one uncommitted
|
||||||
# token because init_next_round_input rebuilt it as full).
|
# token because init_next_round_input rebuilt it as full).
|
||||||
if req.kv_committed_len is not None:
|
if req.kv.kv_committed_len is not None:
|
||||||
req.set_extend_range(len(req.prefix_indices), req.kv_committed_len)
|
req.set_extend_range(
|
||||||
|
len(req.prefix_indices), req.kv.kv_committed_len
|
||||||
|
)
|
||||||
else:
|
else:
|
||||||
waiting_queue.append(req)
|
waiting_queue.append(req)
|
||||||
|
|
||||||
|
|||||||
@@ -465,7 +465,7 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
|
|||||||
|
|
||||||
if req is None or req_to_token_pool is None:
|
if req is None or req_to_token_pool is None:
|
||||||
return
|
return
|
||||||
kv_len = max(req.kv_committed_len, req.kv.kv_allocated_len)
|
kv_len = max(req.kv.kv_committed_len, req.kv.kv_allocated_len)
|
||||||
req_pool_idx = req.req_pool_idx
|
req_pool_idx = req.req_pool_idx
|
||||||
if kv_len <= 0 or req_pool_idx is None:
|
if kv_len <= 0 or req_pool_idx is None:
|
||||||
return
|
return
|
||||||
|
|||||||
@@ -820,6 +820,7 @@ class ReqKvInfo:
|
|||||||
|
|
||||||
# The request's own KV is [cache_protected_len, kv_allocated_len).
|
# The request's own KV is [cache_protected_len, kv_allocated_len).
|
||||||
cache_protected_len: int = 0 # tree cache owns [0, here) (matched or inserted)
|
cache_protected_len: int = 0 # tree cache owns [0, here) (matched or inserted)
|
||||||
|
kv_committed_len: int = 0 # KV content committed up to here, <= kv_allocated_len
|
||||||
kv_allocated_len: int = 0
|
kv_allocated_len: int = 0
|
||||||
|
|
||||||
# SWA slots in [swa_dead_lo(page_size), swa_evicted_seqlen) are already freed.
|
# SWA slots in [swa_dead_lo(page_size), swa_evicted_seqlen) are already freed.
|
||||||
@@ -922,7 +923,6 @@ class Req(ReqDllmMixin):
|
|||||||
self.multi_item_delimiter_indices = multi_item_delimiter_indices
|
self.multi_item_delimiter_indices = multi_item_delimiter_indices
|
||||||
|
|
||||||
# For req-level memory management
|
# For req-level memory management
|
||||||
self.kv_committed_len = 0
|
|
||||||
self.kv = ReqKvInfo()
|
self.kv = ReqKvInfo()
|
||||||
self.retraction_backup: Optional[RetractionBackup] = None
|
self.retraction_backup: Optional[RetractionBackup] = None
|
||||||
|
|
||||||
@@ -1286,8 +1286,8 @@ class Req(ReqDllmMixin):
|
|||||||
# Report only the prompt prefix so thinking + answer fall into the
|
# Report only the prompt prefix so thinking + answer fall into the
|
||||||
# overallocated range and are reclaimed by release_kv_cache. #22373.
|
# overallocated range and are reclaimed by release_kv_cache. #22373.
|
||||||
if get_serving().strip_thinking_cache and self.reasoning_tokens > 0:
|
if get_serving().strip_thinking_cache and self.reasoning_tokens > 0:
|
||||||
return min(self.kv_committed_len, len(self.origin_input_ids))
|
return min(self.kv.kv_committed_len, len(self.origin_input_ids))
|
||||||
return self.kv_committed_len
|
return self.kv.kv_committed_len
|
||||||
|
|
||||||
def update_spec_correct_drafts_histogram(self, num_correct_drafts: int):
|
def update_spec_correct_drafts_histogram(self, num_correct_drafts: int):
|
||||||
"""Record one step accepted draft count (excludes bonus token) into the histogram."""
|
"""Record one step accepted draft count (excludes bonus token) into the histogram."""
|
||||||
@@ -1749,7 +1749,7 @@ class Req(ReqDllmMixin):
|
|||||||
self.mamba_needs_clear = False
|
self.mamba_needs_clear = False
|
||||||
self.already_computed = 0
|
self.already_computed = 0
|
||||||
assert not self.is_holding_kv, "expect it is already released"
|
assert not self.is_holding_kv, "expect it is already released"
|
||||||
self.kv_committed_len = 0
|
self.kv.kv_committed_len = 0
|
||||||
self.extend_batch_idx = 0
|
self.extend_batch_idx = 0
|
||||||
self.decode_batch_idx = 0
|
self.decode_batch_idx = 0
|
||||||
|
|
||||||
@@ -1937,7 +1937,7 @@ def mamba_lazy_spec_in_window(
|
|||||||
kv_committed_len lags device seq_lens by up to one verify under overlap;
|
kv_committed_len lags device seq_lens by up to one verify under overlap;
|
||||||
the 2x window absorbs it.
|
the 2x window absorbs it.
|
||||||
"""
|
"""
|
||||||
seq_len = req.kv_committed_len
|
seq_len = req.kv.kv_committed_len
|
||||||
window = 2 * max_draft_tokens
|
window = 2 * max_draft_tokens
|
||||||
return seq_len // mamba_track_interval != (seq_len + window) // mamba_track_interval
|
return seq_len // mamba_track_interval != (seq_len + window) // mamba_track_interval
|
||||||
|
|
||||||
@@ -2513,7 +2513,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
|||||||
# TODO(th4): co-locate this req.kv bookkeeping with the real KV
|
# TODO(th4): co-locate this req.kv bookkeeping with the real KV
|
||||||
# allocation in alloc_for_extend above; they are currently a few
|
# allocation in alloc_for_extend above; they are currently a few
|
||||||
# steps apart and should become one owned-kv allocation step.
|
# steps apart and should become one owned-kv allocation step.
|
||||||
req.kv_committed_len = seq_len
|
req.kv.kv_committed_len = seq_len
|
||||||
|
|
||||||
# If input_embeds are available, store them
|
# If input_embeds are available, store them
|
||||||
if req.input_embeds is not None:
|
if req.input_embeds is not None:
|
||||||
@@ -2902,7 +2902,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
|||||||
)
|
)
|
||||||
|
|
||||||
if self.spec_algorithm.is_none():
|
if self.spec_algorithm.is_none():
|
||||||
new_pages = sum(1 for r in requests if r.kv_committed_len % page_size == 0)
|
new_pages = sum(
|
||||||
|
1 for r in requests if r.kv.kv_committed_len % page_size == 0
|
||||||
|
)
|
||||||
return new_pages * page_size + num_beam_member_rows(requests)
|
return new_pages * page_size + num_beam_member_rows(requests)
|
||||||
|
|
||||||
return self._new_tokens_required_next_decode_spec_v2(requests, page_size)
|
return self._new_tokens_required_next_decode_spec_v2(requests, page_size)
|
||||||
@@ -2912,7 +2914,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
|||||||
reserve = get_alloc_reserve_per_decode()
|
reserve = get_alloc_reserve_per_decode()
|
||||||
total = 0
|
total = 0
|
||||||
for r in requests:
|
for r in requests:
|
||||||
x = max(0, r.kv_committed_len + reserve - r.kv.kv_allocated_len)
|
x = max(0, r.kv.kv_committed_len + reserve - r.kv.kv_allocated_len)
|
||||||
cur = r.kv.kv_allocated_len
|
cur = r.kv.kv_allocated_len
|
||||||
nxt = cur + x
|
nxt = cur + x
|
||||||
total += ceil_align(nxt, page_size) - ceil_align(cur, page_size)
|
total += ceil_align(nxt, page_size) - ceil_align(cur, page_size)
|
||||||
@@ -3222,7 +3224,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
|||||||
# Update req-level memory management fields
|
# Update req-level memory management fields
|
||||||
for req in self.reqs:
|
for req in self.reqs:
|
||||||
req.decode_batch_idx += 1
|
req.decode_batch_idx += 1
|
||||||
req.kv_committed_len += 1
|
req.kv.kv_committed_len += 1
|
||||||
|
|
||||||
# New-tensor avoids racing model_worker_batch refs queued for
|
# New-tensor avoids racing model_worker_batch refs queued for
|
||||||
# overlap forward.
|
# overlap forward.
|
||||||
|
|||||||
@@ -748,7 +748,7 @@ class SchedulerBatchResultProcessor:
|
|||||||
|
|
||||||
# Commit the full accepted run (drafts + bonus).
|
# Commit the full accepted run (drafts + bonus).
|
||||||
num_accept_tokens = len(accept_tokens)
|
num_accept_tokens = len(accept_tokens)
|
||||||
req.kv_committed_len += num_accept_tokens
|
req.kv.kv_committed_len += num_accept_tokens
|
||||||
req.spec_verify_ct += 1
|
req.spec_verify_ct += 1
|
||||||
|
|
||||||
num_correct_drafts = result.num_correct_drafts_per_req_cpu[i]
|
num_correct_drafts = result.num_correct_drafts_per_req_cpu[i]
|
||||||
@@ -1126,7 +1126,7 @@ class SchedulerBatchResultProcessor:
|
|||||||
|
|
||||||
if completed_mamba_boundary and not lazy:
|
if completed_mamba_boundary and not lazy:
|
||||||
req.mamba_last_track_idx = batch.mamba_track_buffer_indices[i]
|
req.mamba_last_track_idx = batch.mamba_track_buffer_indices[i]
|
||||||
req.mamba_last_track_seqlen = req.kv_committed_len - lookahead
|
req.mamba_last_track_seqlen = req.kv.kv_committed_len - lookahead
|
||||||
elif (
|
elif (
|
||||||
req.finished()
|
req.finished()
|
||||||
and lazy
|
and lazy
|
||||||
@@ -1222,7 +1222,7 @@ class SchedulerBatchResultProcessor:
|
|||||||
lazy = mamba_extra_buffer_lazy_enabled()
|
lazy = mamba_extra_buffer_lazy_enabled()
|
||||||
if known_boundary:
|
if known_boundary:
|
||||||
self._mamba_assert_committed_len_lookahead(req)
|
self._mamba_assert_committed_len_lookahead(req)
|
||||||
track_seqlen = req.kv_committed_len
|
track_seqlen = req.kv.kv_committed_len
|
||||||
assert track_seqlen % mamba_track_grid(self.tree_cache.page_size) == 0
|
assert track_seqlen % mamba_track_grid(self.tree_cache.page_size) == 0
|
||||||
at_boundary = True
|
at_boundary = True
|
||||||
else:
|
else:
|
||||||
@@ -1317,8 +1317,8 @@ class SchedulerBatchResultProcessor:
|
|||||||
f"(req {req.rid}); output_ids is empty"
|
f"(req {req.rid}); output_ids is empty"
|
||||||
)
|
)
|
||||||
token_seq_len = len(req.origin_input_ids) + len(req.output_ids) - 1
|
token_seq_len = len(req.origin_input_ids) + len(req.output_ids) - 1
|
||||||
assert (req.kv_committed_len - token_seq_len) in (0, 1), (
|
assert (req.kv.kv_committed_len - token_seq_len) in (0, 1), (
|
||||||
f"mamba track boundary: kv_committed_len={req.kv_committed_len} "
|
f"mamba track boundary: kv_committed_len={req.kv.kv_committed_len} "
|
||||||
f"leads seq_len={token_seq_len} by more than one (req {req.rid}); "
|
f"leads seq_len={token_seq_len} by more than one (req {req.rid}); "
|
||||||
"overlap lookahead wider than assumed"
|
"overlap lookahead wider than assumed"
|
||||||
)
|
)
|
||||||
@@ -1340,7 +1340,7 @@ class SchedulerBatchResultProcessor:
|
|||||||
|
|
||||||
if batch.spec_algorithm.is_none():
|
if batch.spec_algorithm.is_none():
|
||||||
lookahead = req.decode_batch_idx - batch.mamba_decode_batch_idx_cpu[i]
|
lookahead = req.decode_batch_idx - batch.mamba_decode_batch_idx_cpu[i]
|
||||||
committed_len = req.kv_committed_len - lookahead
|
committed_len = req.kv.kv_committed_len - lookahead
|
||||||
if committed_len % interval == 0:
|
if committed_len % interval == 0:
|
||||||
return True, committed_len
|
return True, committed_len
|
||||||
elif result.num_correct_drafts_per_req_cpu is not None:
|
elif result.num_correct_drafts_per_req_cpu is not None:
|
||||||
|
|||||||
@@ -330,7 +330,7 @@ class SchedulerInvariantChecker:
|
|||||||
req,
|
req,
|
||||||
f"req {req.rid}",
|
f"req {req.rid}",
|
||||||
req.req_pool_idx,
|
req.req_pool_idx,
|
||||||
req.kv_committed_len,
|
req.kv.kv_committed_len,
|
||||||
req.kv.kv_allocated_len,
|
req.kv.kv_allocated_len,
|
||||||
)
|
)
|
||||||
sess = getattr(self.tree_cache, "slots", None)
|
sess = getattr(self.tree_cache, "slots", None)
|
||||||
@@ -341,7 +341,7 @@ class SchedulerInvariantChecker:
|
|||||||
slot,
|
slot,
|
||||||
f"slot {sid[:8]}",
|
f"slot {sid[:8]}",
|
||||||
slot.req_pool_idx,
|
slot.req_pool_idx,
|
||||||
slot.kv_committed_len,
|
slot.kv.kv_committed_len,
|
||||||
slot.kv.kv_allocated_len,
|
slot.kv.kv_allocated_len,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -69,7 +69,7 @@ def page_aligned_decode_alloc_lens(
|
|||||||
cur = r.kv.kv_allocated_len
|
cur = r.kv.kv_allocated_len
|
||||||
nxt = max(
|
nxt = max(
|
||||||
cur,
|
cur,
|
||||||
(r.kv_committed_len + reserve + page_size - 1) // page_size * page_size,
|
(r.kv.kv_committed_len + reserve + page_size - 1) // page_size * page_size,
|
||||||
)
|
)
|
||||||
cur_kv_lens[i] = cur
|
cur_kv_lens[i] = cur
|
||||||
nxt_kv_lens[i] = nxt
|
nxt_kv_lens[i] = nxt
|
||||||
|
|||||||
@@ -255,7 +255,7 @@ def _release_overallocated_kv_indices(
|
|||||||
if spec_algo is None and not get_serving().strip_thinking_cache:
|
if spec_algo is None and not get_serving().strip_thinking_cache:
|
||||||
assert (
|
assert (
|
||||||
start_p == end_p
|
start_p == end_p
|
||||||
), f"Unexpected overallocated KV cache, {req.kv_committed_len=}, {req.kv.kv_allocated_len=}"
|
), f"Unexpected overallocated KV cache, {req.kv.kv_committed_len=}, {req.kv.kv_allocated_len=}"
|
||||||
|
|
||||||
if page_size > 1:
|
if page_size > 1:
|
||||||
start_p = ceil_align(start_p, page_size)
|
start_p = ceil_align(start_p, page_size)
|
||||||
|
|||||||
@@ -300,7 +300,7 @@ class ReqToTokenPool:
|
|||||||
# sum(1 for i in reusing if reqs[i].inflight_middle_chunks > 0) <= 1
|
# sum(1 for i in reusing if reqs[i].inflight_middle_chunks > 0) <= 1
|
||||||
# ), "only one chunked request may reuse req_pool_idx in a batch"
|
# ), "only one chunked request may reuse req_pool_idx in a batch"
|
||||||
assert all(
|
assert all(
|
||||||
reqs[i].inflight_middle_chunks > 0 or reqs[i].kv_committed_len > 0
|
reqs[i].inflight_middle_chunks > 0 or reqs[i].kv.kv_committed_len > 0
|
||||||
for i in reusing
|
for i in reusing
|
||||||
), "reusing request must be chunked or have committed KV"
|
), "reusing request must be chunked or have committed KV"
|
||||||
|
|
||||||
|
|||||||
@@ -399,7 +399,7 @@ class FlexKVRadixCache(RadixCache):
|
|||||||
topk = get_spec().speculative_eagle_topk
|
topk = get_spec().speculative_eagle_topk
|
||||||
enable_kv_committed_len = topk is None or topk == 1
|
enable_kv_committed_len = topk is None or topk == 1
|
||||||
if enable_kv_committed_len:
|
if enable_kv_committed_len:
|
||||||
kv_committed_len = req.kv_committed_len
|
kv_committed_len = req.kv.kv_committed_len
|
||||||
else:
|
else:
|
||||||
kv_committed_len = len(req.origin_input_ids) + max(
|
kv_committed_len = len(req.origin_input_ids) + max(
|
||||||
len(req.output_ids) - 1, 0
|
len(req.output_ids) - 1, 0
|
||||||
|
|||||||
@@ -455,7 +455,7 @@ class LMCRadixCache(RadixCache):
|
|||||||
topk = get_spec().speculative_eagle_topk
|
topk = get_spec().speculative_eagle_topk
|
||||||
enable_kv_committed_len = topk is None or topk == 1
|
enable_kv_committed_len = topk is None or topk == 1
|
||||||
if enable_kv_committed_len:
|
if enable_kv_committed_len:
|
||||||
kv_committed_len = req.kv_committed_len
|
kv_committed_len = req.kv.kv_committed_len
|
||||||
else:
|
else:
|
||||||
kv_committed_len = len(req.origin_input_ids) + max(
|
kv_committed_len = len(req.origin_input_ids) + max(
|
||||||
len(req.output_ids) - 1, 0
|
len(req.output_ids) - 1, 0
|
||||||
|
|||||||
@@ -46,7 +46,6 @@ class SessionSlot:
|
|||||||
|
|
||||||
# KV pool state
|
# KV pool state
|
||||||
req_pool_idx: Optional[int] = None
|
req_pool_idx: Optional[int] = None
|
||||||
kv_committed_len: int = 0
|
|
||||||
kv: ReqKvInfo = field(default_factory=ReqKvInfo)
|
kv: ReqKvInfo = field(default_factory=ReqKvInfo)
|
||||||
|
|
||||||
# First req's radix tree node (for dec_lock_ref on session close)
|
# First req's radix tree node (for dec_lock_ref on session close)
|
||||||
@@ -72,7 +71,6 @@ class SessionSlot:
|
|||||||
def save_from_req(self, req: Req, is_first: bool):
|
def save_from_req(self, req: Req, is_first: bool):
|
||||||
"""Save KV state from a finishing request into this slot."""
|
"""Save KV state from a finishing request into this slot."""
|
||||||
self.req_pool_idx = req.req_pool_idx
|
self.req_pool_idx = req.req_pool_idx
|
||||||
self.kv_committed_len = req.kv_committed_len
|
|
||||||
|
|
||||||
if is_first:
|
if is_first:
|
||||||
self.last_node = req.last_node
|
self.last_node = req.last_node
|
||||||
@@ -107,7 +105,6 @@ class SessionSlot:
|
|||||||
def restore_to_req(self, req: Req):
|
def restore_to_req(self, req: Req):
|
||||||
"""Restore KV state from this slot into an incoming request."""
|
"""Restore KV state from this slot into an incoming request."""
|
||||||
req.req_pool_idx = self.req_pool_idx
|
req.req_pool_idx = self.req_pool_idx
|
||||||
req.kv_committed_len = self.kv_committed_len
|
|
||||||
req.kv = copy.copy(self.kv)
|
req.kv = copy.copy(self.kv)
|
||||||
req.swa_uuid_for_lock = self.swa_uuid_for_lock
|
req.swa_uuid_for_lock = self.swa_uuid_for_lock
|
||||||
req.skip_lock_node_ids = self.skip_lock_node_ids
|
req.skip_lock_node_ids = self.skip_lock_node_ids
|
||||||
@@ -252,7 +249,7 @@ class StreamingSession(BasePrefixCache):
|
|||||||
# fall back to radix cache (full prefill). Once context >= page_size,
|
# fall back to radix cache (full prefill). Once context >= page_size,
|
||||||
# streaming session kicks in with page-aligned KV reuse.
|
# streaming session kicks in with page-aligned KV reuse.
|
||||||
if is_npu() and self.page_size > 1:
|
if is_npu() and self.page_size > 1:
|
||||||
expected_prefix_len = min(slot.kv_committed_len, len(params.key))
|
expected_prefix_len = min(slot.kv.kv_committed_len, len(params.key))
|
||||||
aligned_prefix_len = (
|
aligned_prefix_len = (
|
||||||
expected_prefix_len // self.page_size
|
expected_prefix_len // self.page_size
|
||||||
) * self.page_size
|
) * self.page_size
|
||||||
@@ -270,7 +267,7 @@ class StreamingSession(BasePrefixCache):
|
|||||||
# token_ids = get_fill_ids()[:input_len-1] (1-token logit reserve
|
# token_ids = get_fill_ids()[:input_len-1] (1-token logit reserve
|
||||||
# already applied). min handles retract retry where committed_len
|
# already applied). min handles retract retry where committed_len
|
||||||
# can exceed len(token_ids) by 1.
|
# can exceed len(token_ids) by 1.
|
||||||
prefix_len = min(req.kv_committed_len, len(params.key))
|
prefix_len = min(req.kv.kv_committed_len, len(params.key))
|
||||||
|
|
||||||
# Streaming sessions are append-only (session_controller rollback
|
# Streaming sessions are append-only (session_controller rollback
|
||||||
# ensures req_nodes always points to the last successful req).
|
# ensures req_nodes always points to the last successful req).
|
||||||
@@ -282,8 +279,8 @@ class StreamingSession(BasePrefixCache):
|
|||||||
# Floor-align prefix_len to page boundary (NPU workaround).
|
# Floor-align prefix_len to page boundary (NPU workaround).
|
||||||
if is_npu() and self.page_size > 1:
|
if is_npu() and self.page_size > 1:
|
||||||
prefix_len = (prefix_len // self.page_size) * self.page_size
|
prefix_len = (prefix_len // self.page_size) * self.page_size
|
||||||
req.kv_committed_len = min(req.kv_committed_len, prefix_len)
|
req.kv.kv_committed_len = min(req.kv.kv_committed_len, prefix_len)
|
||||||
slot.kv_committed_len = min(slot.kv_committed_len, prefix_len)
|
slot.kv.kv_committed_len = min(slot.kv.kv_committed_len, prefix_len)
|
||||||
|
|
||||||
# Free orphaned tail: alloc_for_extend will overwrite
|
# Free orphaned tail: alloc_for_extend will overwrite
|
||||||
# req_to_token[prefix_len:] with new indices. The range
|
# req_to_token[prefix_len:] with new indices. The range
|
||||||
@@ -370,7 +367,7 @@ class StreamingSession(BasePrefixCache):
|
|||||||
# req clock (under overlap + honest committed the clock lags the in-flight
|
# req clock (under overlap + honest committed the clock lags the in-flight
|
||||||
# verify by ~1, which would short-change inheritance). Clamp to allocated
|
# verify by ~1, which would short-change inheritance). Clamp to allocated
|
||||||
# to keep committed <= allocated for prepare_for_decode.
|
# to keep committed <= allocated for prepare_for_decode.
|
||||||
slot.kv_committed_len = min(target, slot.kv.kv_allocated_len)
|
slot.kv.kv_committed_len = min(target, slot.kv.kv_allocated_len)
|
||||||
|
|
||||||
# Update req_nodes to this successfully finished request.
|
# Update req_nodes to this successfully finished request.
|
||||||
req.session.finish_req(req)
|
req.session.finish_req(req)
|
||||||
@@ -564,10 +561,10 @@ class StreamingSession(BasePrefixCache):
|
|||||||
"""
|
"""
|
||||||
self._free_kv_aligned(slot.req_pool_idx, prefix_len, slot.kv.kv_allocated_len)
|
self._free_kv_aligned(slot.req_pool_idx, prefix_len, slot.kv.kv_allocated_len)
|
||||||
slot.kv.kv_allocated_len = prefix_len
|
slot.kv.kv_allocated_len = prefix_len
|
||||||
slot.kv_committed_len = min(slot.kv_committed_len, prefix_len)
|
slot.kv.kv_committed_len = min(slot.kv.kv_committed_len, prefix_len)
|
||||||
slot.kv.swa_evicted_seqlen = min(slot.kv.swa_evicted_seqlen, prefix_len)
|
slot.kv.swa_evicted_seqlen = min(slot.kv.swa_evicted_seqlen, prefix_len)
|
||||||
req.kv.kv_allocated_len = prefix_len
|
req.kv.kv_allocated_len = prefix_len
|
||||||
req.kv_committed_len = min(req.kv_committed_len, prefix_len)
|
req.kv.kv_committed_len = min(req.kv.kv_committed_len, prefix_len)
|
||||||
req.kv.swa_evicted_seqlen = min(req.kv.swa_evicted_seqlen, prefix_len)
|
req.kv.swa_evicted_seqlen = min(req.kv.swa_evicted_seqlen, prefix_len)
|
||||||
|
|
||||||
def _trim_overshoot(self, req: Req, finished_len: int) -> None:
|
def _trim_overshoot(self, req: Req, finished_len: int) -> None:
|
||||||
@@ -579,7 +576,7 @@ class StreamingSession(BasePrefixCache):
|
|||||||
target = len(req.origin_input_ids) + finished_len
|
target = len(req.origin_input_ids) + finished_len
|
||||||
self._free_kv_aligned(req.req_pool_idx, target, req.kv.kv_allocated_len)
|
self._free_kv_aligned(req.req_pool_idx, target, req.kv.kv_allocated_len)
|
||||||
req.kv.kv_allocated_len = min(req.kv.kv_allocated_len, target)
|
req.kv.kv_allocated_len = min(req.kv.kv_allocated_len, target)
|
||||||
req.kv_committed_len = min(req.kv_committed_len, target)
|
req.kv.kv_committed_len = min(req.kv.kv_committed_len, target)
|
||||||
req.kv.swa_evicted_seqlen = min(req.kv.swa_evicted_seqlen, target)
|
req.kv.swa_evicted_seqlen = min(req.kv.swa_evicted_seqlen, target)
|
||||||
req.output_ids = req.output_ids[:finished_len]
|
req.output_ids = req.output_ids[:finished_len]
|
||||||
|
|
||||||
|
|||||||
@@ -159,7 +159,7 @@ class DFlashDraftInputV2(SpecInput):
|
|||||||
nxt_kv_lens_sum = 0
|
nxt_kv_lens_sum = 0
|
||||||
committed_seq_lens_sum = 0
|
committed_seq_lens_sum = 0
|
||||||
for i, (req, cur, nxt) in enumerate(zip(batch.reqs, cur_kv_lens, nxt_kv_lens)):
|
for i, (req, cur, nxt) in enumerate(zip(batch.reqs, cur_kv_lens, nxt_kv_lens)):
|
||||||
committed_len = int(req.kv_committed_len)
|
committed_len = int(req.kv.kv_committed_len)
|
||||||
committed_seq_lens_sum += committed_len
|
committed_seq_lens_sum += committed_len
|
||||||
top_k = int(req.sampling_params.top_k)
|
top_k = int(req.sampling_params.top_k)
|
||||||
|
|
||||||
|
|||||||
@@ -117,7 +117,7 @@ def remaining_prompt_tokens(ctx: ScriptedContext, rid: str) -> int:
|
|||||||
req = find_req_by_rid(ctx, rid)
|
req = find_req_by_rid(ctx, rid)
|
||||||
if req is None:
|
if req is None:
|
||||||
return 0
|
return 0
|
||||||
return max(0, len(req.origin_input_ids) - req.kv_committed_len)
|
return max(0, len(req.origin_input_ids) - req.kv.kv_committed_len)
|
||||||
|
|
||||||
|
|
||||||
def chunks_done(ctx: ScriptedContext, rid: str) -> int:
|
def chunks_done(ctx: ScriptedContext, rid: str) -> int:
|
||||||
|
|||||||
@@ -96,10 +96,10 @@ class TestSWABasic(ScriptedTestCase):
|
|||||||
r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
|
r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
|
||||||
for _ in range(400):
|
for _ in range(400):
|
||||||
if r.is_chunking:
|
if r.is_chunking:
|
||||||
assert len(r.req.prefix_indices) <= r.req.kv_committed_len, (
|
assert len(r.req.prefix_indices) <= r.req.kv.kv_committed_len, (
|
||||||
f"prefix_indices must be bounded by kv_committed_len, "
|
f"prefix_indices must be bounded by kv_committed_len, "
|
||||||
f"got prefix_indices_len={len(r.req.prefix_indices)}, "
|
f"got prefix_indices_len={len(r.req.prefix_indices)}, "
|
||||||
f"kv_committed_len={r.req.kv_committed_len}"
|
f"kv_committed_len={r.req.kv.kv_committed_len}"
|
||||||
)
|
)
|
||||||
if r.finished:
|
if r.finished:
|
||||||
break
|
break
|
||||||
|
|||||||
@@ -509,7 +509,7 @@ class TestRegressionGptOss(ScriptedTestCase):
|
|||||||
r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
|
r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
|
||||||
yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1)
|
yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1)
|
||||||
|
|
||||||
committed = r.req.kv_committed_len
|
committed = r.req.kv.kv_committed_len
|
||||||
assert committed > 0
|
assert committed > 0
|
||||||
|
|
||||||
assert len(r.req.prefix_indices) <= committed, (
|
assert len(r.req.prefix_indices) <= committed, (
|
||||||
|
|||||||
@@ -225,10 +225,10 @@ class TestSpecialCaseBasic(ScriptedTestCase):
|
|||||||
)
|
)
|
||||||
for _ in range(DEFAULT_MAX_STEPS):
|
for _ in range(DEFAULT_MAX_STEPS):
|
||||||
if r.is_chunking:
|
if r.is_chunking:
|
||||||
assert len(r.req.prefix_indices) <= r.req.kv_committed_len, (
|
assert len(r.req.prefix_indices) <= r.req.kv.kv_committed_len, (
|
||||||
f"streaming-session chunked stash must stay bounded by "
|
f"streaming-session chunked stash must stay bounded by "
|
||||||
f"kv_committed_len; prefix_indices_len={len(r.req.prefix_indices)}, "
|
f"kv_committed_len; prefix_indices_len={len(r.req.prefix_indices)}, "
|
||||||
f"kv_committed_len={r.req.kv_committed_len}"
|
f"kv_committed_len={r.req.kv.kv_committed_len}"
|
||||||
)
|
)
|
||||||
if r.finished:
|
if r.finished:
|
||||||
break
|
break
|
||||||
|
|||||||
@@ -119,8 +119,9 @@ class _FakeAllocator:
|
|||||||
class TestFreeMemberRows(CustomTestCase):
|
class TestFreeMemberRows(CustomTestCase):
|
||||||
def _make_group(self, req_to_token, allocated_len):
|
def _make_group(self, req_to_token, allocated_len):
|
||||||
leader = SimpleNamespace(
|
leader = SimpleNamespace(
|
||||||
kv=SimpleNamespace(kv_allocated_len=allocated_len),
|
kv=SimpleNamespace(
|
||||||
kv_committed_len=allocated_len,
|
kv_allocated_len=allocated_len, kv_committed_len=allocated_len
|
||||||
|
),
|
||||||
)
|
)
|
||||||
return SimpleNamespace(
|
return SimpleNamespace(
|
||||||
leader=leader,
|
leader=leader,
|
||||||
@@ -146,7 +147,7 @@ class TestFreeMemberRows(CustomTestCase):
|
|||||||
# Leader rewound to the prompt: its own release must not free the
|
# Leader rewound to the prompt: its own release must not free the
|
||||||
# decode region a second time.
|
# decode region a second time.
|
||||||
self.assertEqual(leader.kv.kv_allocated_len, 5)
|
self.assertEqual(leader.kv.kv_allocated_len, 5)
|
||||||
self.assertEqual(leader.kv_committed_len, 5)
|
self.assertEqual(leader.kv.kv_committed_len, 5)
|
||||||
self.assertEqual(sorted(pool.freed), [1, 2])
|
self.assertEqual(sorted(pool.freed), [1, 2])
|
||||||
self.assertIsNone(group.member_rows)
|
self.assertIsNone(group.member_rows)
|
||||||
self.assertIsNone(group.member_rows_cpu)
|
self.assertIsNone(group.member_rows_cpu)
|
||||||
@@ -207,7 +208,7 @@ class TestRetireReclaimsStagedOrphans(CustomTestCase):
|
|||||||
allocator = _FakeAllocator()
|
allocator = _FakeAllocator()
|
||||||
group = SimpleNamespace(
|
group = SimpleNamespace(
|
||||||
leader=SimpleNamespace(
|
leader=SimpleNamespace(
|
||||||
kv=SimpleNamespace(kv_allocated_len=8), kv_committed_len=8
|
kv=SimpleNamespace(kv_allocated_len=8, kv_committed_len=8),
|
||||||
),
|
),
|
||||||
prompt_len=5,
|
prompt_len=5,
|
||||||
member_rows=torch.tensor([1, 2], dtype=torch.int64),
|
member_rows=torch.tensor([1, 2], dtype=torch.int64),
|
||||||
|
|||||||
@@ -34,10 +34,11 @@ def _make_mock_req(
|
|||||||
req = MagicMock()
|
req = MagicMock()
|
||||||
req.rid = rid
|
req.rid = rid
|
||||||
req.req_pool_idx = req_pool_idx
|
req.req_pool_idx = req_pool_idx
|
||||||
req.kv_committed_len = kv_committed_len
|
req.kv = ReqKvInfo(
|
||||||
req.kv = ReqKvInfo(kv_allocated_len=kv_allocated_len)
|
kv_committed_len=kv_committed_len, kv_allocated_len=kv_allocated_len
|
||||||
|
)
|
||||||
req.prefix_indices = list(range(prefix_indices_len))
|
req.prefix_indices = list(range(prefix_indices_len))
|
||||||
req.effective_kv_committed_len = lambda: req.kv_committed_len
|
req.effective_kv_committed_len = lambda: req.kv.kv_committed_len
|
||||||
return req
|
return req
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -1521,7 +1521,6 @@ if _HAS_MLX:
|
|||||||
self.req_pool_idx = None
|
self.req_pool_idx = None
|
||||||
self.mamba_pool_idx = None
|
self.mamba_pool_idx = None
|
||||||
self.inflight_middle_chunks = 0
|
self.inflight_middle_chunks = 0
|
||||||
self.kv_committed_len = 0
|
|
||||||
|
|
||||||
class FakeTpWorker:
|
class FakeTpWorker:
|
||||||
def __init__(self, next_token_ids):
|
def __init__(self, next_token_ids):
|
||||||
|
|||||||
@@ -121,7 +121,6 @@ def _fake_req():
|
|||||||
return SimpleNamespace(
|
return SimpleNamespace(
|
||||||
req_pool_idx=None,
|
req_pool_idx=None,
|
||||||
inflight_middle_chunks=0,
|
inflight_middle_chunks=0,
|
||||||
kv_committed_len=0,
|
|
||||||
mamba_pool_idx=None,
|
mamba_pool_idx=None,
|
||||||
mamba_ping_pong_track_buffer=None,
|
mamba_ping_pong_track_buffer=None,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -68,7 +68,6 @@ def make_pool_and_req(capacity: int = 64):
|
|||||||
req = SimpleNamespace(
|
req = SimpleNamespace(
|
||||||
req_pool_idx=None,
|
req_pool_idx=None,
|
||||||
inflight_middle_chunks=0,
|
inflight_middle_chunks=0,
|
||||||
kv_committed_len=0,
|
|
||||||
)
|
)
|
||||||
req_pool_idx = pool.alloc([req])[0]
|
req_pool_idx = pool.alloc([req])[0]
|
||||||
return pool, req, req_pool_idx, allocator
|
return pool, req, req_pool_idx, allocator
|
||||||
|
|||||||
@@ -33,7 +33,7 @@ def _make_batch() -> tuple[Req, ScheduleBatch]:
|
|||||||
vocab_size=128,
|
vocab_size=128,
|
||||||
)
|
)
|
||||||
req.output_ids.append(3)
|
req.output_ids.append(3)
|
||||||
req.kv_committed_len = 2
|
req.kv.kv_committed_len = 2
|
||||||
|
|
||||||
batch = ScheduleBatch(reqs=[req])
|
batch = ScheduleBatch(reqs=[req])
|
||||||
batch.tree_cache = SimpleNamespace(page_size=TRACK_INTERVAL)
|
batch.tree_cache = SimpleNamespace(page_size=TRACK_INTERVAL)
|
||||||
|
|||||||
@@ -92,7 +92,7 @@ def _make_req(terminate_after: int) -> Req:
|
|||||||
sampling_params=sp,
|
sampling_params=sp,
|
||||||
)
|
)
|
||||||
req.grammar = _FakeGrammar(terminate_after=terminate_after)
|
req.grammar = _FakeGrammar(terminate_after=terminate_after)
|
||||||
req.kv_committed_len = 0
|
req.kv.kv_committed_len = 0
|
||||||
return req
|
return req
|
||||||
|
|
||||||
|
|
||||||
@@ -121,7 +121,7 @@ class TestSpecV2GrammarTruncation(CustomTestCase):
|
|||||||
|
|
||||||
self.assertEqual(predict_tokens, [[101, 102]])
|
self.assertEqual(predict_tokens, [[101, 102]])
|
||||||
# No pre-claim: commit the full retained run (no -1 refund).
|
# No pre-claim: commit the full retained run (no -1 refund).
|
||||||
self.assertEqual(req.kv_committed_len, 2)
|
self.assertEqual(req.kv.kv_committed_len, 2)
|
||||||
|
|
||||||
def test_resolve_keeps_all_when_grammar_not_terminated(self):
|
def test_resolve_keeps_all_when_grammar_not_terminated(self):
|
||||||
req = _make_req(terminate_after=99)
|
req = _make_req(terminate_after=99)
|
||||||
@@ -131,7 +131,7 @@ class TestSpecV2GrammarTruncation(CustomTestCase):
|
|||||||
predict_tokens = proc._resolve_spec_v2_tokens(result, _FakeBatch([req]))
|
predict_tokens = proc._resolve_spec_v2_tokens(result, _FakeBatch([req]))
|
||||||
|
|
||||||
self.assertEqual(predict_tokens, [[201, 202, 203]])
|
self.assertEqual(predict_tokens, [[201, 202, 203]])
|
||||||
self.assertEqual(req.kv_committed_len, 3)
|
self.assertEqual(req.kv.kv_committed_len, 3)
|
||||||
|
|
||||||
|
|
||||||
class TestReasoningTokenAccounting(CustomTestCase):
|
class TestReasoningTokenAccounting(CustomTestCase):
|
||||||
|
|||||||
@@ -50,8 +50,7 @@ def _make_req(rid="test-req-0", origin_input_ids=None, output_ids=None):
|
|||||||
fill_ids=origin_input_ids + output_ids,
|
fill_ids=origin_input_ids + output_ids,
|
||||||
seqlen=len(origin_input_ids) + len(output_ids),
|
seqlen=len(origin_input_ids) + len(output_ids),
|
||||||
req_pool_idx=None,
|
req_pool_idx=None,
|
||||||
kv=SimpleNamespace(kv_allocated_len=0),
|
kv=SimpleNamespace(kv_allocated_len=0, kv_committed_len=0),
|
||||||
kv_committed_len=0,
|
|
||||||
finished_reason=None,
|
finished_reason=None,
|
||||||
hisparse_staging=False,
|
hisparse_staging=False,
|
||||||
staging=False,
|
staging=False,
|
||||||
@@ -219,7 +218,7 @@ class TestHiSparseUnit(unittest.TestCase):
|
|||||||
self.assertIsNotNone(kv_loc, "KV alloc failed")
|
self.assertIsNotNone(kv_loc, "KV alloc failed")
|
||||||
self.req_to_token_pool.write((req.req_pool_idx, slice(0, len(kv_loc))), kv_loc)
|
self.req_to_token_pool.write((req.req_pool_idx, slice(0, len(kv_loc))), kv_loc)
|
||||||
req.kv.kv_allocated_len = fill_len
|
req.kv.kv_allocated_len = fill_len
|
||||||
req.kv_committed_len = fill_len
|
req.kv.kv_committed_len = fill_len
|
||||||
req.full_untruncated_fill_ids = array("q", range(fill_len))
|
req.full_untruncated_fill_ids = array("q", range(fill_len))
|
||||||
req.extend_range = Range(0, fill_len)
|
req.extend_range = Range(0, fill_len)
|
||||||
return kv_loc
|
return kv_loc
|
||||||
@@ -579,7 +578,7 @@ class TestHiSparseUnit(unittest.TestCase):
|
|||||||
seq_len = fill_len + 1
|
seq_len = fill_len + 1
|
||||||
self.req_to_token_pool.write((req.req_pool_idx, fill_len), out_loc)
|
self.req_to_token_pool.write((req.req_pool_idx, fill_len), out_loc)
|
||||||
req.kv.kv_allocated_len = seq_len
|
req.kv.kv_allocated_len = seq_len
|
||||||
req.kv_committed_len = seq_len
|
req.kv.kv_committed_len = seq_len
|
||||||
|
|
||||||
self.coordinator.map_last_loc_to_buffer(
|
self.coordinator.map_last_loc_to_buffer(
|
||||||
seq_lens=torch.tensor([seq_len], dtype=torch.int64, device=device),
|
seq_lens=torch.tensor([seq_len], dtype=torch.int64, device=device),
|
||||||
@@ -770,7 +769,7 @@ class TestHiSparseUnit(unittest.TestCase):
|
|||||||
)
|
)
|
||||||
)
|
)
|
||||||
self.assertEqual(req.kv.kv_allocated_len, fill_len)
|
self.assertEqual(req.kv.kv_allocated_len, fill_len)
|
||||||
self.assertEqual(req.kv_committed_len, fill_len)
|
self.assertEqual(req.kv.kv_committed_len, fill_len)
|
||||||
self.assertEqual(req.extend_range.length, fill_len)
|
self.assertEqual(req.extend_range.length, fill_len)
|
||||||
|
|
||||||
rounded_len = (fill_len + self.page_size - 1) // self.page_size * self.page_size
|
rounded_len = (fill_len + self.page_size - 1) // self.page_size * self.page_size
|
||||||
|
|||||||
@@ -47,16 +47,18 @@ class _FakeReq:
|
|||||||
def __init__(self, rid, rpi, committed, allocated):
|
def __init__(self, rid, rpi, committed, allocated):
|
||||||
self.rid = rid
|
self.rid = rid
|
||||||
self.req_pool_idx = rpi
|
self.req_pool_idx = rpi
|
||||||
self.kv_committed_len = committed
|
self.kv = SimpleNamespace(
|
||||||
self.kv = SimpleNamespace(kv_allocated_len=allocated, swa_evicted_seqlen=0)
|
kv_committed_len=committed, kv_allocated_len=allocated, swa_evicted_seqlen=0
|
||||||
|
)
|
||||||
self.is_holding_kv = True
|
self.is_holding_kv = True
|
||||||
|
|
||||||
|
|
||||||
class _FakeSlot:
|
class _FakeSlot:
|
||||||
def __init__(self, rpi, committed, allocated):
|
def __init__(self, rpi, committed, allocated):
|
||||||
self.req_pool_idx = rpi
|
self.req_pool_idx = rpi
|
||||||
self.kv_committed_len = committed
|
self.kv = SimpleNamespace(
|
||||||
self.kv = SimpleNamespace(kv_allocated_len=allocated, swa_evicted_seqlen=0)
|
kv_committed_len=committed, kv_allocated_len=allocated, swa_evicted_seqlen=0
|
||||||
|
)
|
||||||
self.is_holding_kv = True
|
self.is_holding_kv = True
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -18,8 +18,7 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
|||||||
def _make_req():
|
def _make_req():
|
||||||
return types.SimpleNamespace(
|
return types.SimpleNamespace(
|
||||||
decode_batch_idx=0,
|
decode_batch_idx=0,
|
||||||
kv_committed_len=3,
|
kv=types.SimpleNamespace(kv_committed_len=3, kv_allocated_len=3),
|
||||||
kv_allocated_len=3,
|
|
||||||
beam_group=None,
|
beam_group=None,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -81,9 +81,10 @@ class MockReq:
|
|||||||
self.cache_salt = None
|
self.cache_salt = None
|
||||||
self.prefix_indices = torch.empty(0, dtype=torch.int64)
|
self.prefix_indices = torch.empty(0, dtype=torch.int64)
|
||||||
self.priority = 0
|
self.priority = 0
|
||||||
self.kv_committed_len = len(fill_ids)
|
|
||||||
self.kv = SimpleNamespace(
|
self.kv = SimpleNamespace(
|
||||||
kv_allocated_len=len(fill_ids), cache_protected_len=cache_protected_len
|
kv_committed_len=len(fill_ids),
|
||||||
|
kv_allocated_len=len(fill_ids),
|
||||||
|
cache_protected_len=cache_protected_len,
|
||||||
)
|
)
|
||||||
|
|
||||||
def get_fill_ids(self):
|
def get_fill_ids(self):
|
||||||
@@ -211,7 +212,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
|
|||||||
cache.cache_unfinished_req(req)
|
cache.cache_unfinished_req(req)
|
||||||
|
|
||||||
# Step 3: cache_finished_req with is_insert=True (dec lock)
|
# Step 3: cache_finished_req with is_insert=True (dec lock)
|
||||||
cache.cache_finished_req(req, kv_len_to_handle=req.kv_committed_len)
|
cache.cache_finished_req(req, kv_len_to_handle=req.kv.kv_committed_len)
|
||||||
|
|
||||||
# Verify: all non-root nodes should have lock_ref == 0
|
# Verify: all non-root nodes should have lock_ref == 0
|
||||||
# (root always has lock_ref == 1)
|
# (root always has lock_ref == 1)
|
||||||
@@ -260,7 +261,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
|
|||||||
cache.cache_unfinished_req(req)
|
cache.cache_unfinished_req(req)
|
||||||
|
|
||||||
# Step 3: cache_finished_req (dec leaf)
|
# Step 3: cache_finished_req (dec leaf)
|
||||||
cache.cache_finished_req(req, kv_len_to_handle=req.kv_committed_len)
|
cache.cache_finished_req(req, kv_len_to_handle=req.kv.kv_committed_len)
|
||||||
|
|
||||||
# Root lock unchanged, all nodes unlocked
|
# Root lock unchanged, all nodes unlocked
|
||||||
self.assertEqual(cache.root_node.lock_ref, root_lock_before)
|
self.assertEqual(cache.root_node.lock_ref, root_lock_before)
|
||||||
@@ -304,7 +305,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
|
|||||||
# Transfer fails -> cache_finished_req with is_insert=False
|
# Transfer fails -> cache_finished_req with is_insert=False
|
||||||
# This frees delta tokens and dec_lock_ref on last_node
|
# This frees delta tokens and dec_lock_ref on last_node
|
||||||
cache.cache_finished_req(
|
cache.cache_finished_req(
|
||||||
req, is_insert=False, kv_len_to_handle=req.kv_committed_len
|
req, is_insert=False, kv_len_to_handle=req.kv.kv_committed_len
|
||||||
)
|
)
|
||||||
|
|
||||||
# The prefix node should be unlocked (back to evictable)
|
# The prefix node should be unlocked (back to evictable)
|
||||||
@@ -351,7 +352,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
|
|||||||
# Transfer fails -> cache_finished_req with is_insert=False
|
# Transfer fails -> cache_finished_req with is_insert=False
|
||||||
# dec_lock_ref(root) is a no-op
|
# dec_lock_ref(root) is a no-op
|
||||||
cache.cache_finished_req(
|
cache.cache_finished_req(
|
||||||
req, is_insert=False, kv_len_to_handle=req.kv_committed_len
|
req, is_insert=False, kv_len_to_handle=req.kv.kv_committed_len
|
||||||
)
|
)
|
||||||
|
|
||||||
# Root lock unchanged, nothing protected or evictable
|
# Root lock unchanged, nothing protected or evictable
|
||||||
@@ -482,7 +483,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
|
|||||||
)
|
)
|
||||||
|
|
||||||
cache.cache_unfinished_req(req)
|
cache.cache_unfinished_req(req)
|
||||||
cache.cache_finished_req(req, kv_len_to_handle=req.kv_committed_len)
|
cache.cache_finished_req(req, kv_len_to_handle=req.kv.kv_committed_len)
|
||||||
|
|
||||||
# After all iterations, root lock should be 1, no protected nodes
|
# After all iterations, root lock should be 1, no protected nodes
|
||||||
self.assertEqual(cache.root_node.lock_ref, 1)
|
self.assertEqual(cache.root_node.lock_ref, 1)
|
||||||
|
|||||||
@@ -64,9 +64,11 @@ def _make_req(rid, prefix, block_size, *, req_pool_idx=None, reuse=False):
|
|||||||
req_pool_idx=req_pool_idx,
|
req_pool_idx=req_pool_idx,
|
||||||
dllm_incomplete_ids=array("q", range(block_size)) if reuse else array("q"),
|
dllm_incomplete_ids=array("q", range(block_size)) if reuse else array("q"),
|
||||||
inflight_middle_chunks=1 if req_pool_idx is not None else 0,
|
inflight_middle_chunks=1 if req_pool_idx is not None else 0,
|
||||||
kv_committed_len=len(prefix) if req_pool_idx is not None else 0,
|
|
||||||
kv=SimpleNamespace(
|
kv=SimpleNamespace(
|
||||||
kv_allocated_len=len(prefix) + block_size if req_pool_idx is not None else 0
|
kv_committed_len=len(prefix) if req_pool_idx is not None else 0,
|
||||||
|
kv_allocated_len=(
|
||||||
|
len(prefix) + block_size if req_pool_idx is not None else 0
|
||||||
|
),
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -153,7 +153,7 @@ class TestDeepSeekV4HiSparseAllocator(CustomTestCase):
|
|||||||
self.assertEqual(kwargs["swa_tail_len"], swa_tail_len)
|
self.assertEqual(kwargs["swa_tail_len"], swa_tail_len)
|
||||||
self.assertEqual(req.kv.swa_evicted_seqlen, fill_len - swa_tail_len)
|
self.assertEqual(req.kv.swa_evicted_seqlen, fill_len - swa_tail_len)
|
||||||
self.assertEqual(req.kv.kv_allocated_len, fill_len)
|
self.assertEqual(req.kv.kv_allocated_len, fill_len)
|
||||||
self.assertEqual(req.kv_committed_len, fill_len)
|
self.assertEqual(req.kv.kv_committed_len, fill_len)
|
||||||
self.assertEqual(req.extend_range.length, fill_len)
|
self.assertEqual(req.extend_range.length, fill_len)
|
||||||
self.assertEqual(len(req_to_token_pool.writes), 1)
|
self.assertEqual(len(req_to_token_pool.writes), 1)
|
||||||
coordinator.host_token_len.assert_called_once_with(fill_len)
|
coordinator.host_token_len.assert_called_once_with(fill_len)
|
||||||
|
|||||||
@@ -166,7 +166,6 @@ def register(cache, token_ids, session_id, generation=None):
|
|||||||
).last_device_node,
|
).last_device_node,
|
||||||
origin_input_ids=array("q", token_ids),
|
origin_input_ids=array("q", token_ids),
|
||||||
output_ids=array("q"),
|
output_ids=array("q"),
|
||||||
kv_committed_len=len(token_ids),
|
|
||||||
extra_key=None,
|
extra_key=None,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -68,8 +68,8 @@ class _FakeReq:
|
|||||||
_inflight=False,
|
_inflight=False,
|
||||||
)
|
)
|
||||||
self.req_pool_idx = req_pool_idx
|
self.req_pool_idx = req_pool_idx
|
||||||
self.kv_committed_len = committed
|
|
||||||
self.kv = SimpleNamespace(
|
self.kv = SimpleNamespace(
|
||||||
|
kv_committed_len=committed,
|
||||||
kv_allocated_len=allocated,
|
kv_allocated_len=allocated,
|
||||||
swa_evicted_seqlen=0,
|
swa_evicted_seqlen=0,
|
||||||
cache_protected_len=0,
|
cache_protected_len=0,
|
||||||
@@ -113,9 +113,11 @@ def test_preabort_detaches_session_and_preserves_slot():
|
|||||||
tree_cache = StreamingSession(inner)
|
tree_cache = StreamingSession(inner)
|
||||||
tree_cache.slots["session-a"] = SessionSlot(
|
tree_cache.slots["session-a"] = SessionSlot(
|
||||||
req_pool_idx=0,
|
req_pool_idx=0,
|
||||||
kv_committed_len=48,
|
|
||||||
kv=SimpleNamespace(
|
kv=SimpleNamespace(
|
||||||
kv_allocated_len=48, swa_evicted_seqlen=0, cache_protected_len=16
|
kv_committed_len=48,
|
||||||
|
kv_allocated_len=48,
|
||||||
|
swa_evicted_seqlen=0,
|
||||||
|
cache_protected_len=16,
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -134,7 +136,7 @@ def test_preabort_detaches_session_and_preserves_slot():
|
|||||||
# Slot untouched.
|
# Slot untouched.
|
||||||
slot = tree_cache.slots["session-a"]
|
slot = tree_cache.slots["session-a"]
|
||||||
assert slot.req_pool_idx == 0
|
assert slot.req_pool_idx == 0
|
||||||
assert slot.kv_committed_len == 48
|
assert slot.kv.kv_committed_len == 48
|
||||||
assert slot.kv.kv_allocated_len == 48
|
assert slot.kv.kv_allocated_len == 48
|
||||||
assert len(result.device_indices) == 0
|
assert len(result.device_indices) == 0
|
||||||
|
|
||||||
@@ -178,9 +180,11 @@ def test_nth_mid_abort_nukes_session_slot():
|
|||||||
# Session already has a slot from a previous turn.
|
# Session already has a slot from a previous turn.
|
||||||
tree_cache.slots["session-a"] = SessionSlot(
|
tree_cache.slots["session-a"] = SessionSlot(
|
||||||
req_pool_idx=0,
|
req_pool_idx=0,
|
||||||
kv_committed_len=50,
|
|
||||||
kv=SimpleNamespace(
|
kv=SimpleNamespace(
|
||||||
kv_allocated_len=50, swa_evicted_seqlen=0, cache_protected_len=0
|
kv_committed_len=50,
|
||||||
|
kv_allocated_len=50,
|
||||||
|
swa_evicted_seqlen=0,
|
||||||
|
cache_protected_len=0,
|
||||||
),
|
),
|
||||||
last_node=None,
|
last_node=None,
|
||||||
)
|
)
|
||||||
@@ -217,9 +221,11 @@ def test_release_session_threads_mamba_skip_ids():
|
|||||||
lock_node = SimpleNamespace(id=42)
|
lock_node = SimpleNamespace(id=42)
|
||||||
tree_cache.slots["session-a"] = SessionSlot(
|
tree_cache.slots["session-a"] = SessionSlot(
|
||||||
req_pool_idx=0,
|
req_pool_idx=0,
|
||||||
kv_committed_len=50,
|
|
||||||
kv=SimpleNamespace(
|
kv=SimpleNamespace(
|
||||||
kv_allocated_len=50, swa_evicted_seqlen=0, cache_protected_len=0
|
kv_committed_len=50,
|
||||||
|
kv_allocated_len=50,
|
||||||
|
swa_evicted_seqlen=0,
|
||||||
|
cache_protected_len=0,
|
||||||
),
|
),
|
||||||
last_node=lock_node,
|
last_node=lock_node,
|
||||||
skip_lock_node_ids={ComponentType.MAMBA: {42}},
|
skip_lock_node_ids={ComponentType.MAMBA: {42}},
|
||||||
@@ -265,7 +271,7 @@ def test_trim_overshoot_postcondition():
|
|||||||
tree_cache._trim_overshoot(req, finished_len=12)
|
tree_cache._trim_overshoot(req, finished_len=12)
|
||||||
|
|
||||||
target = 38
|
target = 38
|
||||||
assert req.kv_committed_len == target
|
assert req.kv.kv_committed_len == target
|
||||||
assert req.kv.kv_allocated_len == target
|
assert req.kv.kv_allocated_len == target
|
||||||
assert req.kv.swa_evicted_seqlen == target
|
assert req.kv.swa_evicted_seqlen == target
|
||||||
assert len(req.output_ids) == 12
|
assert len(req.output_ids) == 12
|
||||||
|
|||||||
@@ -643,7 +643,7 @@ def bench_cache_finished(
|
|||||||
)
|
)
|
||||||
req.last_node = node
|
req.last_node = node
|
||||||
req.kv.cache_protected_len = matched_len
|
req.kv.cache_protected_len = matched_len
|
||||||
req.kv_committed_len = len(seq)
|
req.kv.kv_committed_len = len(seq)
|
||||||
if hasattr(lr, "swa_uuid_for_lock"):
|
if hasattr(lr, "swa_uuid_for_lock"):
|
||||||
req.swa_uuid_for_lock = lr.swa_uuid_for_lock
|
req.swa_uuid_for_lock = lr.swa_uuid_for_lock
|
||||||
env.rtp.req_to_token[req.req_pool_idx, : len(kv_indices)] = kv_indices
|
env.rtp.req_to_token[req.req_pool_idx, : len(kv_indices)] = kv_indices
|
||||||
@@ -657,7 +657,7 @@ def bench_cache_finished(
|
|||||||
"cache_finished",
|
"cache_finished",
|
||||||
lambda: req_items,
|
lambda: req_items,
|
||||||
lambda req: env.tree.cache_finished_req(
|
lambda req: env.tree.cache_finished_req(
|
||||||
req, is_insert=True, kv_len_to_handle=req.kv_committed_len
|
req, is_insert=True, kv_len_to_handle=req.kv.kv_committed_len
|
||||||
),
|
),
|
||||||
len(req_items) - warmup,
|
len(req_items) - warmup,
|
||||||
env.avg_tokens,
|
env.avg_tokens,
|
||||||
|
|||||||
@@ -1226,7 +1226,7 @@ class UnifiedRadixCacheSuite:
|
|||||||
kv_len = len(input_ids) + len(output_ids)
|
kv_len = len(input_ids) + len(output_ids)
|
||||||
kv_indices = self._alloc(allocator, kv_len)
|
kv_indices = self._alloc(allocator, kv_len)
|
||||||
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
|
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
|
||||||
req.kv_committed_len = kv_len
|
req.kv.kv_committed_len = kv_len
|
||||||
req.last_node = cache.root_node_handle()
|
req.last_node = cache.root_node_handle()
|
||||||
req.kv.cache_protected_len = 0
|
req.kv.cache_protected_len = 0
|
||||||
req.swa_uuid_for_lock = None
|
req.swa_uuid_for_lock = None
|
||||||
@@ -1266,7 +1266,7 @@ class UnifiedRadixCacheSuite:
|
|||||||
kv_len = req.extend_range.end
|
kv_len = req.extend_range.end
|
||||||
kv_indices = self._alloc(allocator, kv_len)
|
kv_indices = self._alloc(allocator, kv_len)
|
||||||
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
|
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
|
||||||
req.kv_committed_len = kv_len
|
req.kv.kv_committed_len = kv_len
|
||||||
req.kv.kv_allocated_len = kv_len
|
req.kv.kv_allocated_len = kv_len
|
||||||
req.last_node = cache.root_node_handle()
|
req.last_node = cache.root_node_handle()
|
||||||
req.kv.cache_protected_len = 0
|
req.kv.cache_protected_len = 0
|
||||||
@@ -1311,7 +1311,7 @@ class UnifiedRadixCacheSuite:
|
|||||||
kv_len = len(tokens)
|
kv_len = len(tokens)
|
||||||
kv_indices = self._alloc(allocator, kv_len)
|
kv_indices = self._alloc(allocator, kv_len)
|
||||||
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
|
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
|
||||||
req.kv_committed_len = kv_len
|
req.kv.kv_committed_len = kv_len
|
||||||
req.last_node = cache.root_node_handle()
|
req.last_node = cache.root_node_handle()
|
||||||
req.kv.cache_protected_len = 0
|
req.kv.cache_protected_len = 0
|
||||||
req.swa_uuid_for_lock = None
|
req.swa_uuid_for_lock = None
|
||||||
@@ -1346,7 +1346,7 @@ class UnifiedRadixCacheSuite:
|
|||||||
kv_len = len(tokens)
|
kv_len = len(tokens)
|
||||||
kv_indices = self._alloc(allocator, kv_len)
|
kv_indices = self._alloc(allocator, kv_len)
|
||||||
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
|
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
|
||||||
req.kv_committed_len = kv_len
|
req.kv.kv_committed_len = kv_len
|
||||||
req.last_node = cache.root_node_handle()
|
req.last_node = cache.root_node_handle()
|
||||||
req.kv.cache_protected_len = 0
|
req.kv.cache_protected_len = 0
|
||||||
req.swa_uuid_for_lock = None
|
req.swa_uuid_for_lock = None
|
||||||
@@ -1383,7 +1383,7 @@ class UnifiedRadixCacheSuite:
|
|||||||
req.set_extend_range(0, len(req.full_untruncated_fill_ids))
|
req.set_extend_range(0, len(req.full_untruncated_fill_ids))
|
||||||
kv_indices = self._alloc(allocator, len(tokens))
|
kv_indices = self._alloc(allocator, len(tokens))
|
||||||
req_to_token_pool.write((req.req_pool_idx, slice(0, len(tokens))), kv_indices)
|
req_to_token_pool.write((req.req_pool_idx, slice(0, len(tokens))), kv_indices)
|
||||||
req.kv_committed_len = len(tokens)
|
req.kv.kv_committed_len = len(tokens)
|
||||||
req.last_node = cache.root_node_handle()
|
req.last_node = cache.root_node_handle()
|
||||||
req.kv.cache_protected_len = 0
|
req.kv.cache_protected_len = 0
|
||||||
req.swa_uuid_for_lock = None
|
req.swa_uuid_for_lock = None
|
||||||
@@ -1485,7 +1485,7 @@ class UnifiedRadixCacheSuite:
|
|||||||
kv_len = len(input_ids)
|
kv_len = len(input_ids)
|
||||||
kv_indices = self._alloc(allocator, kv_len)
|
kv_indices = self._alloc(allocator, kv_len)
|
||||||
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
|
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
|
||||||
req.kv_committed_len = kv_len
|
req.kv.kv_committed_len = kv_len
|
||||||
req.last_node = cache.root_node_handle()
|
req.last_node = cache.root_node_handle()
|
||||||
req.kv.cache_protected_len = 0
|
req.kv.cache_protected_len = 0
|
||||||
req.swa_uuid_for_lock = None
|
req.swa_uuid_for_lock = None
|
||||||
@@ -1610,7 +1610,7 @@ class UnifiedRadixCacheSuite:
|
|||||||
kv_len = len(tokens)
|
kv_len = len(tokens)
|
||||||
fresh_value = self._alloc(allocator, kv_len)
|
fresh_value = self._alloc(allocator, kv_len)
|
||||||
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), fresh_value)
|
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), fresh_value)
|
||||||
req.kv_committed_len = kv_len
|
req.kv.kv_committed_len = kv_len
|
||||||
req.last_node = cache.root_node_handle()
|
req.last_node = cache.root_node_handle()
|
||||||
req.kv.cache_protected_len = 0
|
req.kv.cache_protected_len = 0
|
||||||
req.swa_uuid_for_lock = None
|
req.swa_uuid_for_lock = None
|
||||||
@@ -2197,7 +2197,7 @@ class UnifiedRadixCacheSuite:
|
|||||||
req.set_extend_range(0, len(req.full_untruncated_fill_ids))
|
req.set_extend_range(0, len(req.full_untruncated_fill_ids))
|
||||||
kv_indices = self._alloc(allocator, pre_len)
|
kv_indices = self._alloc(allocator, pre_len)
|
||||||
req_to_token_pool.write((req.req_pool_idx, slice(0, pre_len)), kv_indices)
|
req_to_token_pool.write((req.req_pool_idx, slice(0, pre_len)), kv_indices)
|
||||||
req.kv_committed_len = pre_len
|
req.kv.kv_committed_len = pre_len
|
||||||
req.last_node = cache.root_node_handle()
|
req.last_node = cache.root_node_handle()
|
||||||
req.kv.cache_protected_len = 0
|
req.kv.cache_protected_len = 0
|
||||||
req.swa_uuid_for_lock = None
|
req.swa_uuid_for_lock = None
|
||||||
@@ -2286,7 +2286,7 @@ class UnifiedRadixCacheSuite:
|
|||||||
req.set_extend_range(0, len(req.full_untruncated_fill_ids))
|
req.set_extend_range(0, len(req.full_untruncated_fill_ids))
|
||||||
kv_indices = self._alloc(allocator, pre_len)
|
kv_indices = self._alloc(allocator, pre_len)
|
||||||
req_to_token_pool.write((req.req_pool_idx, slice(0, pre_len)), kv_indices)
|
req_to_token_pool.write((req.req_pool_idx, slice(0, pre_len)), kv_indices)
|
||||||
req.kv_committed_len = pre_len
|
req.kv.kv_committed_len = pre_len
|
||||||
req.last_node = cache.root_node_handle()
|
req.last_node = cache.root_node_handle()
|
||||||
req.kv.cache_protected_len = 0
|
req.kv.cache_protected_len = 0
|
||||||
req.swa_uuid_for_lock = None
|
req.swa_uuid_for_lock = None
|
||||||
@@ -6768,7 +6768,7 @@ class TestUnifiedRadixCacheInt8MambaCheckpoint(CustomTestCase):
|
|||||||
)
|
)
|
||||||
req_to_token_pool.alloc([req])
|
req_to_token_pool.alloc([req])
|
||||||
req.output_ids = array("q")
|
req.output_ids = array("q")
|
||||||
req.kv_committed_len = len(tokens)
|
req.kv.kv_committed_len = len(tokens)
|
||||||
req.kv.kv_allocated_len = len(tokens)
|
req.kv.kv_allocated_len = len(tokens)
|
||||||
req.kv.cache_protected_len = 0
|
req.kv.cache_protected_len = 0
|
||||||
req.swa_uuid_for_lock = None
|
req.swa_uuid_for_lock = None
|
||||||
@@ -8095,7 +8095,7 @@ class TestSWAWindowUnderBigramKey(CustomTestCase):
|
|||||||
req.set_extend_range(0, len(req.full_untruncated_fill_ids))
|
req.set_extend_range(0, len(req.full_untruncated_fill_ids))
|
||||||
kv_indices = self._alloc_paged(allocator, seq_len)
|
kv_indices = self._alloc_paged(allocator, seq_len)
|
||||||
req_to_token_pool.write((req.req_pool_idx, slice(0, seq_len)), kv_indices)
|
req_to_token_pool.write((req.req_pool_idx, slice(0, seq_len)), kv_indices)
|
||||||
req.kv_committed_len = seq_len
|
req.kv.kv_committed_len = seq_len
|
||||||
req.last_node = cache.root_node_handle()
|
req.last_node = cache.root_node_handle()
|
||||||
req.kv.cache_protected_len = 0
|
req.kv.cache_protected_len = 0
|
||||||
req.swa_uuid_for_lock = None
|
req.swa_uuid_for_lock = None
|
||||||
|
|||||||
@@ -104,9 +104,7 @@ _OWNER_SITES = {
|
|||||||
"free_member_rows",
|
"free_member_rows",
|
||||||
"kv_allocated_len",
|
"kv_allocated_len",
|
||||||
): 1,
|
): 1,
|
||||||
# streaming session slot save/restore and tail trimming
|
# streaming session tail trimming
|
||||||
(_SS, "SessionSlot.save_from_req", "kv_committed_len"): 1,
|
|
||||||
(_SS, "SessionSlot.restore_to_req", "kv_committed_len"): 1,
|
|
||||||
(_SS, "StreamingSession._free_tail", "kv_committed_len"): 2,
|
(_SS, "StreamingSession._free_tail", "kv_committed_len"): 2,
|
||||||
(_SS, "StreamingSession._free_tail", "kv_allocated_len"): 2,
|
(_SS, "StreamingSession._free_tail", "kv_allocated_len"): 2,
|
||||||
(_SS, "StreamingSession._trim_overshoot", "kv_committed_len"): 1,
|
(_SS, "StreamingSession._trim_overshoot", "kv_committed_len"): 1,
|
||||||
|
|||||||
@@ -72,10 +72,9 @@ def _make_req(rid, req_pool_idx, token_ids, tree):
|
|||||||
extra_key=None,
|
extra_key=None,
|
||||||
cache_salt=None,
|
cache_salt=None,
|
||||||
last_node=tree.root_node,
|
last_node=tree.root_node,
|
||||||
kv=SimpleNamespace(cache_protected_len=0),
|
kv=SimpleNamespace(cache_protected_len=0, kv_committed_len=len(token_ids)),
|
||||||
priority=0,
|
priority=0,
|
||||||
kv_committed_freed=False,
|
kv_committed_freed=False,
|
||||||
kv_committed_len=len(token_ids),
|
|
||||||
)
|
)
|
||||||
req.pop_committed_kv_cache = lambda: len(token_ids)
|
req.pop_committed_kv_cache = lambda: len(token_ids)
|
||||||
return req
|
return req
|
||||||
@@ -163,11 +162,7 @@ class TestLMCRadixCacheXPU(unittest.TestCase):
|
|||||||
# commit it as a finished request (inserts into radix + stores to
|
# commit it as a finished request (inserts into radix + stores to
|
||||||
# LMCache on tree.store_stream).
|
# LMCache on tree.store_stream).
|
||||||
req_pool_idx = req_to_token_pool.alloc(
|
req_pool_idx = req_to_token_pool.alloc(
|
||||||
[
|
[SimpleNamespace(req_pool_idx=None, inflight_middle_chunks=0)]
|
||||||
SimpleNamespace(
|
|
||||||
req_pool_idx=None, inflight_middle_chunks=0, kv_committed_len=0
|
|
||||||
)
|
|
||||||
]
|
|
||||||
)[0]
|
)[0]
|
||||||
kv_slots = allocator.alloc(self.INPUT_LEN)
|
kv_slots = allocator.alloc(self.INPUT_LEN)
|
||||||
self.assertIsNotNone(kv_slots)
|
self.assertIsNotNone(kv_slots)
|
||||||
|
|||||||
Reference in New Issue
Block a user