diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 228edd0b9..b717256cf 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -62,11 +62,13 @@ from sglang.srt.mem_cache.common import ( page_align_floor, release_kv_cache, ) +from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.mem_cache.memory_pool import ( HybridReqToTokenPool, KVCache, ReqToTokenPool, ) +from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.observability.req_time_stats import ( set_schedule_time_batch, set_time_batch, @@ -175,7 +177,6 @@ class DecodeReqToTokenPool: class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool): - def __init__( self, size: int, @@ -316,13 +317,55 @@ class DecodePreallocQueue: if self.enable_staging: self.transfer_queue._init_staging_handler(self.kv_manager) - if self.scheduler.tp_worker.is_hybrid_swa: - # FIXME: current SWA allocation allocate full kv cache size in prefill + if ( + self.scheduler.tp_worker.is_hybrid_swa + and not self._uses_swa_tail_prealloc() + ): + # Fallback for SWA allocators that still allocate the SWA pool at + # full prompt length. self.max_total_num_tokens = min( self.max_total_num_tokens, self.scheduler.tp_worker.model_runner.swa_max_total_num_tokens, ) + def _uses_swa_tail_prealloc(self) -> bool: + return ( + isinstance(self.token_to_kv_pool, (SWAKVPool, DeepSeekV4TokenToKVPool)) + and self.token_to_kv_pool_allocator.page_size > 1 + and hasattr(self.token_to_kv_pool_allocator, "alloc_extend_swa_tail") + ) + + def _swa_tail_len(self, seq_len: int) -> int: + if not self._uses_swa_tail_prealloc() or seq_len <= 0: + return max(seq_len, 0) + + window_size = self.scheduler.sliding_window_size + if window_size is None or window_size <= 0: + return seq_len + + page_size = self.token_to_kv_pool_allocator.page_size + window_start = max(0, seq_len - window_size) + window_start = (window_start // page_size) * page_size + return seq_len - window_start + + def _swa_retractable_len(self, req: Req) -> int: + if not self._uses_swa_tail_prealloc(): + return len(req.origin_input_ids) + len(req.output_ids) + return self._swa_tail_len(len(req.origin_input_ids)) + len(req.output_ids) + + def _prealloc_kv_lens(self, req: Req) -> Tuple[int, int]: + allocated_kv_len = len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0) + if self._uses_swa_tail_prealloc(): + return allocated_kv_len, self._swa_tail_len(allocated_kv_len) + return allocated_kv_len, allocated_kv_len + + def _prealloc_required_tokens(self, req: Req) -> Tuple[int, int]: + full_len, swa_len = self._prealloc_kv_lens(req) + return ( + full_len + self.num_reserved_decode_tokens, + swa_len + self.num_reserved_decode_tokens, + ) + def _init_kv_manager(self) -> CommonKVManager: kv_args_class = get_kv_class(self.transfer_backend, KVClassType.KVARGS) kv_args = kv_args_class() @@ -486,6 +529,18 @@ class DecodePreallocQueue: prepare_abort(req, message, status_code=HTTPStatus.BAD_REQUEST) self.scheduler.stream_output([req], req.return_logprob) return True + if self._uses_swa_tail_prealloc(): + _, swa_required = self._prealloc_required_tokens(req) + swa_capacity = self.token_to_kv_pool_allocator.size_swa + if swa_required > swa_capacity: + message = ( + f"Request {req.rid} requires too many SWA KV tokens for " + f"decode preallocation: {swa_required} > {swa_capacity}" + ) + logger.error(message) + prepare_abort(req, message, status_code=HTTPStatus.BAD_REQUEST) + self.scheduler.stream_output([req], req.return_logprob) + return True return False def extend(self, reqs: List[Req], is_retracted: bool = False) -> None: @@ -501,7 +556,15 @@ class DecodePreallocQueue: # allocate memory resumed_reqs = [] indices_to_remove = set() - allocatable_tokens = self._allocatable_tokens(count_retracted=False) + uses_swa_tail_prealloc = self._uses_swa_tail_prealloc() + if uses_swa_tail_prealloc: + full_allocatable_tokens, swa_allocatable_tokens = ( + self._swa_aware_allocatable_token_budgets(count_retracted=False) + ) + else: + full_allocatable_tokens = self._allocatable_token_budgets( + count_retracted=False + ) for i, req in enumerate(self.retracted_queue): if rids_to_check is not None and req.rid not in rids_to_check: @@ -510,19 +573,19 @@ class DecodePreallocQueue: if self.req_to_token_pool.available_size() <= 0: break - required_tokens_for_request = ( - len(req.origin_input_ids) - + len(req.output_ids) - + self.num_reserved_decode_tokens - ) - if required_tokens_for_request > allocatable_tokens: + full_required, swa_required = self._prealloc_required_tokens(req) + if full_required > full_allocatable_tokens: + break + if uses_swa_tail_prealloc and swa_required > swa_allocatable_tokens: break resumed_reqs.append(req) indices_to_remove.add(i) req.is_retracted = False self._pre_alloc(req) - allocatable_tokens -= required_tokens_for_request + full_allocatable_tokens -= full_required + if uses_swa_tail_prealloc: + swa_allocatable_tokens -= swa_required # load from cpu, release the cpu copy req.load_kv_cache(self.req_to_token_pool, self.token_to_kv_pool_allocator) @@ -677,9 +740,24 @@ class DecodePreallocQueue: len(r.origin_input_ids) + len(r.output_ids) for r in self.scheduler.running_batch.reqs ) - allocatable_tokens = self._allocatable_tokens( - retractable_tokens=retractable_tokens, count_retracted=True - ) + uses_swa_tail_prealloc = self._uses_swa_tail_prealloc() + swa_allocatable_tokens = 0 + if uses_swa_tail_prealloc: + retractable_swa_tokens = sum( + self._swa_retractable_len(r) for r in self.scheduler.running_batch.reqs + ) + full_allocatable_tokens, swa_allocatable_tokens = ( + self._swa_aware_allocatable_token_budgets( + retractable_tokens=retractable_tokens, + retractable_swa_tokens=retractable_swa_tokens, + count_retracted=True, + ) + ) + else: + retractable_swa_tokens = 0 + full_allocatable_tokens = self._allocatable_token_budgets( + retractable_tokens=retractable_tokens, count_retracted=True + ) # First, remove all failed requests from the queue for i, decode_req in enumerate(self.queue): if rids_to_check is not None and decode_req.req.rid not in rids_to_check: @@ -746,7 +824,7 @@ class DecodePreallocQueue: # Matching may lock previously-evictable radix pages, so refresh # the admission budget against the post-lock pool state before we # decide whether this request still fits. - allocatable_tokens = self._allocatable_tokens( + full_allocatable_tokens = self._allocatable_token_budgets( retractable_tokens=retractable_tokens, count_retracted=True, extra_reserved_reqs=len(preallocated_reqs), @@ -771,25 +849,47 @@ class DecodePreallocQueue: ) - retractable_tokens, ) - > allocatable_tokens + > full_allocatable_tokens ): if prefix_len > 0: self.tree_cache.dec_lock_ref(decode_req.req.last_node) break - if required_tokens_for_request > allocatable_tokens: + if required_tokens_for_request > full_allocatable_tokens: if prefix_len > 0: self.tree_cache.dec_lock_ref(decode_req.req.last_node) break + if uses_swa_tail_prealloc: + _, swa_required = self._prealloc_required_tokens(decode_req.req) + _, swa_len = self._prealloc_kv_lens(decode_req.req) + max_new_tokens = min( + decode_req.req.sampling_params.max_new_tokens, + CLIP_MAX_NEW_TOKEN, + ) + if ( + max( + swa_required, + swa_len + max_new_tokens - retractable_swa_tokens, + ) + > swa_allocatable_tokens + ): + if prefix_len > 0: + self.tree_cache.dec_lock_ref(decode_req.req.last_node) + break + dst_kv_indices = self._pre_alloc(decode_req.req, prefix_indices, prefix_len) hisparse_req_budget -= 1 # Recompute from actual pool state for the next queue entry. # This accounts for page rounding and newly locked evictable cache. - allocatable_tokens = self._allocatable_tokens( + full_allocatable_tokens = self._allocatable_token_budgets( retractable_tokens=retractable_tokens, count_retracted=True, extra_reserved_reqs=len(preallocated_reqs) + 1, ) + if uses_swa_tail_prealloc: + # SWA budget uses simple decrement (no radix cache eviction in + # the SWA pool, so page-rounding drift is negligible). + swa_allocatable_tokens -= swa_required decode_req.req.cache_protected_len = prefix_len if self.scheduler.enable_hisparse: @@ -896,11 +996,8 @@ class DecodePreallocQueue: len(decode_req.req.fill_ids) for decode_req in self.transfer_queue.queue ) - def _allocatable_tokens( - self, - retractable_tokens: Optional[int] = None, - count_retracted: bool = True, - extra_reserved_reqs: int = 0, + def _need_space_for_single_req( + self, retractable_tokens: Optional[int] = None ) -> int: need_space_for_single_req = ( max( @@ -915,12 +1012,69 @@ class DecodePreallocQueue: and len(self.scheduler.running_batch.reqs) > 0 else 0 ) + return need_space_for_single_req + + def _active_req_count(self, extra_reserved_reqs: int = 0) -> int: + return ( + len(self.scheduler.running_batch.reqs) + + len(self.transfer_queue.queue) + + len(self.scheduler.waiting_queue) + + extra_reserved_reqs + ) + + def _active_reserved_tokens( + self, n_active: Optional[int] = None, extra_reserved_reqs: int = 0 + ) -> int: + if n_active is None: + n_active = self._active_req_count(extra_reserved_reqs) + return self.num_reserved_decode_tokens * n_active + + def _swa_aware_allocatable_token_budgets( + self, + retractable_tokens: Optional[int] = None, + retractable_swa_tokens: Optional[int] = None, + count_retracted: bool = True, + ) -> Tuple[int, int]: + n_active = self._active_req_count() + reserved_tokens = self._active_reserved_tokens(n_active) + + full_allocatable_tokens = self._allocatable_token_budgets( + retractable_tokens=retractable_tokens, + count_retracted=count_retracted, + reserved_tokens=reserved_tokens, + ) + + return full_allocatable_tokens, self._swa_tail_allocatable_token_budget( + retractable_tokens=retractable_tokens, + retractable_swa_tokens=retractable_swa_tokens, + count_retracted=count_retracted, + n_active=n_active, + reserved_tokens=reserved_tokens, + ) + + def _allocatable_token_budgets( + self, + retractable_tokens: Optional[int] = None, + count_retracted: bool = True, + extra_reserved_reqs: int = 0, + reserved_tokens: Optional[int] = None, + ) -> int: + need_space_for_single_req = self._need_space_for_single_req(retractable_tokens) + if reserved_tokens is None: + reserved_tokens = self._active_reserved_tokens( + extra_reserved_reqs=extra_reserved_reqs + ) + if self.scheduler.enable_hisparse: # HiSparse pre-alloc only allocates logical indices (alloc_logical_only), # so the logical pool is the binding constraint for admission control. available_size = ( self.token_to_kv_pool_allocator.logical_attn_allocator.available_size() ) + elif self._uses_swa_tail_prealloc(): + available_size = self.token_to_kv_pool_allocator.full_available_size() + if self.scheduler.server_args.disaggregation_decode_enable_radix_cache: + available_size += self.tree_cache.evictable_size() else: available_size = self.token_to_kv_pool_allocator.available_size() # Include evictable decode-radix cache entries in the budget -- they @@ -928,16 +1082,7 @@ class DecodePreallocQueue: if self.scheduler.server_args.disaggregation_decode_enable_radix_cache: available_size += self.tree_cache.evictable_size() allocatable_tokens = available_size - max( - # preserve some space for future decode - self.num_reserved_decode_tokens - * ( - len(self.scheduler.running_batch.reqs) - + len(self.transfer_queue.queue) - + len(self.scheduler.waiting_queue) - + extra_reserved_reqs - ), - # make sure each request can finish if reach max_tokens with all other requests retracted - need_space_for_single_req, + reserved_tokens, need_space_for_single_req ) # Note: if the last prebuilt extend just finishes, and we enter `pop_preallocated` immediately in the next iteration @@ -951,16 +1096,75 @@ class DecodePreallocQueue: ) if count_retracted: - allocatable_tokens -= sum( - [ - len(req.origin_input_ids) - + len(req.output_ids) - + self.num_reserved_decode_tokens - for req in self.retracted_queue - ] - ) + for req in self.retracted_queue: + full_required, _ = self._prealloc_required_tokens(req) + allocatable_tokens -= full_required + return allocatable_tokens + def _swa_tail_allocatable_token_budget( + self, + retractable_tokens: Optional[int] = None, + retractable_swa_tokens: Optional[int] = None, + count_retracted: bool = True, + n_active: Optional[int] = None, + reserved_tokens: Optional[int] = None, + ) -> int: + need_swa_space_for_single_req = self._need_space_for_single_req( + retractable_tokens + ) + if ( + retractable_swa_tokens is not None + and len(self.scheduler.running_batch.reqs) > 0 + ): + need_swa_space_for_single_req = max( + self._swa_tail_len(len(x.origin_input_ids)) + + min(x.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKEN) + - retractable_swa_tokens + for x in self.scheduler.running_batch.reqs + ) + + if n_active is None: + n_active = self._active_req_count() + if reserved_tokens is None: + reserved_tokens = self._active_reserved_tokens(n_active) + + # SWA growth is bounded by the sliding window: once a req's SWA + # footprint reaches `sliding_window_size`, further decode tokens + # evict old ones and net growth is zero. The linear reservation + # `num_reserved_decode_tokens * n_active` (correct for the full + # pool) over-reserves SWA in steady state. Cap by the actual + # remaining headroom up to per-req window cap. + window_size = self.scheduler.sliding_window_size or 0 + swa_total = self.token_to_kv_pool_allocator.size_swa + swa_used = swa_total - self.token_to_kv_pool_allocator.swa_available_size() + swa_growth_potential = max(0, n_active * window_size - swa_used) + swa_reserved_tokens = min(reserved_tokens, swa_growth_potential) + swa_allocatable_tokens = ( + self.token_to_kv_pool_allocator.swa_available_size() + - max(swa_reserved_tokens, need_swa_space_for_single_req) + ) + + # Note: if the last prebuilt extend just finishes, and we enter `pop_preallocated` immediately in the next iteration + # the extend batch is not in any queue, so we need to explicitly add the tokens slots here + if ( + self.scheduler.last_batch + and self.scheduler.last_batch.forward_mode.is_prebuilt() + ): + prebuilt_reserved_tokens = self.num_reserved_decode_tokens * len( + self.scheduler.last_batch.reqs + ) + prebuilt_n = len(self.scheduler.last_batch.reqs) + prebuilt_swa_growth = max(0, prebuilt_n * window_size - swa_used) + swa_allocatable_tokens -= min(prebuilt_reserved_tokens, prebuilt_swa_growth) + + if count_retracted: + for req in self.retracted_queue: + _, swa_required = self._prealloc_required_tokens(req) + swa_allocatable_tokens -= swa_required + + return swa_allocatable_tokens + def _required_alloc_tokens(self, *, fill_len: int, prefix_len: int) -> int: page_size = self.token_to_kv_pool_allocator.page_size if page_size == 1: @@ -1062,16 +1266,31 @@ class DecodePreallocQueue: if prefix_len > 0 else torch.tensor([-1], dtype=torch.int64, device=device) ) - kv_loc = self.token_to_kv_pool_allocator.alloc_extend( - prefix_lens=torch.tensor( - [prefix_len], dtype=torch.int64, device=device - ), - prefix_lens_cpu=torch.tensor([prefix_len], dtype=torch.int64), - seq_lens=torch.tensor([fill_len], dtype=torch.int64, device=device), - seq_lens_cpu=torch.tensor([fill_len], dtype=torch.int64), - last_loc=last_loc, - extend_num_tokens=delta_len, - ) + if self._uses_swa_tail_prealloc() and prefix_len == 0: + # Tail-only SWA allocation: only valid when prefix_len == 0. + # When prefix_len > 0 (radix cache hit), we fall back to + # alloc_extend which allocates SWA at full page count; the + # SWA budget in that case may slightly under-estimate. + kv_loc = self.token_to_kv_pool_allocator.alloc_extend_swa_tail( + prefix_lens=torch.tensor([0], dtype=torch.int64, device=device), + prefix_lens_cpu=torch.tensor([0], dtype=torch.int64), + seq_lens=torch.tensor([fill_len], dtype=torch.int64, device=device), + seq_lens_cpu=torch.tensor([fill_len], dtype=torch.int64), + last_loc=last_loc, + extend_num_tokens=fill_len, + swa_tail_len=self._swa_tail_len(fill_len), + ) + else: + kv_loc = self.token_to_kv_pool_allocator.alloc_extend( + prefix_lens=torch.tensor( + [prefix_len], dtype=torch.int64, device=device + ), + prefix_lens_cpu=torch.tensor([prefix_len], dtype=torch.int64), + seq_lens=torch.tensor([fill_len], dtype=torch.int64, device=device), + seq_lens_cpu=torch.tensor([fill_len], dtype=torch.int64), + last_loc=last_loc, + extend_num_tokens=delta_len, + ) assert kv_loc is not None, ( f"KV cache is full! Bug in memory estimation. " @@ -1333,7 +1552,6 @@ class DecodeTransferQueue: class SchedulerDisaggregationDecodeMixin: - @torch.no_grad() def event_loop_normal_disagg_decode(self: Scheduler): """A normal scheduler loop for decode worker in disaggregation mode.""" diff --git a/python/sglang/srt/mem_cache/swa_memory_pool.py b/python/sglang/srt/mem_cache/swa_memory_pool.py index 393738ec3..62b1b11e0 100644 --- a/python/sglang/srt/mem_cache/swa_memory_pool.py +++ b/python/sglang/srt/mem_cache/swa_memory_pool.py @@ -214,20 +214,55 @@ class SWAKVPool(BaseSWAKVPool): src_loc_swa = self.translate_loc_from_full_to_swa(src_loc) self.swa_kv_pool.move_kv_cache(tgt_loc_swa, src_loc_swa) + def _filter_swa_cpu_copy(self, swa_kv_cpu, row_mask: torch.Tensor): + if swa_kv_cpu is None: + return None + if row_mask is None or bool(torch.all(row_mask).item()): + return swa_kv_cpu + + chunk_size = getattr( + self.swa_kv_pool, "cpu_offloading_chunk_size", len(row_mask) + ) + filtered = [] + for layer_chunks in swa_kv_cpu: + if len(layer_chunks) == 0: + filtered.append([]) + continue + + k_cpu = torch.cat([chunk[0] for chunk in layer_chunks], dim=0) + v_cpu = torch.cat([chunk[1] for chunk in layer_chunks], dim=0) + k_cpu = k_cpu[row_mask] + v_cpu = v_cpu[row_mask] + + filtered_layer = [] + for i in range(0, len(k_cpu), chunk_size): + filtered_layer.append( + [k_cpu[i : i + chunk_size], v_cpu[i : i + chunk_size]] + ) + filtered.append(filtered_layer) + return filtered + def get_cpu_copy(self, indices, mamba_indices=None): # For SWA, we need to copy KV cache from both full and SWA pools # The indices are for the full pool, and we use mapping to get SWA indices full_kv_cpu = self.full_kv_pool.get_cpu_copy(indices) - # Get SWA indices through the mapping - # Note: SWA allocation always creates 1:1 mapping, so no need to filter + swa_mask = None if self.full_to_swa_index_mapping is not None: swa_indices = self.full_to_swa_index_mapping[indices] - swa_kv_cpu = self.swa_kv_pool.get_cpu_copy(swa_indices) + # Slot 0 is reserved as a dummy slot. Tail-only SWA allocations leave + # the out-of-window full KV indices unmapped, so only copy mapped SWA + # tokens and keep their positions for load_cpu_copy(). + swa_mask = swa_indices > 0 + if torch.any(swa_mask): + swa_kv_cpu = self.swa_kv_pool.get_cpu_copy(swa_indices[swa_mask]) + swa_mask = swa_mask.cpu() + else: + swa_kv_cpu = None else: swa_kv_cpu = None - return {"full": full_kv_cpu, "swa": swa_kv_cpu} + return {"full": full_kv_cpu, "swa": swa_kv_cpu, "swa_mask": swa_mask} def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): # Load KV cache back from CPU to both full and SWA pools @@ -241,6 +276,20 @@ class SWAKVPool(BaseSWAKVPool): # Load SWA KV cache if it exists if swa_kv_cpu is not None and self.full_to_swa_index_mapping is not None: swa_indices = self.full_to_swa_index_mapping[indices] + new_swa_mask = swa_indices > 0 + old_swa_mask = kv_cache_cpu.get("swa_mask") + if old_swa_mask is not None: + old_swa_mask = old_swa_mask.to(indices.device) + row_mask = new_swa_mask[old_swa_mask].cpu() + swa_indices = swa_indices[old_swa_mask][row_mask.to(indices.device)] + else: + row_mask = new_swa_mask.cpu() + swa_indices = swa_indices[new_swa_mask] + + if swa_indices.numel() == 0: + return + + swa_kv_cpu = self._filter_swa_cpu_copy(swa_kv_cpu, row_mask) self.swa_kv_pool.load_cpu_copy(swa_kv_cpu, swa_indices) @@ -437,6 +486,71 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): return alloc_full_indices + def alloc_extend_swa_tail( + self, + prefix_lens: torch.Tensor, + prefix_lens_cpu: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, # last_loc for full layers + extend_num_tokens: int, + swa_tail_len: int, + ): + """Allocate full KV for the whole extend and SWA KV only for the tail. + + This is used by disaggregated decode preallocation: decode receives full + prompt KV for full-attention layers, but only the sliding-window state is + transferred for SWA layers. + """ + assert self.page_size > 1 + assert len(seq_lens_cpu) == 1, "SWA tail allocation currently supports bs=1" + assert len(prefix_lens_cpu) == 1 + assert 0 <= swa_tail_len <= extend_num_tokens + + num_full_pages = get_num_new_pages( + seq_lens=seq_lens_cpu, page_size=self.page_size, prefix_lens=prefix_lens_cpu + ) + num_swa_pages = (swa_tail_len + self.page_size - 1) // self.page_size + if num_full_pages > self.full_attn_allocator.available_size() // self.page_size: + return None + if num_swa_pages > self.swa_attn_allocator.available_size() // self.page_size: + return None + + alloc_full_indices = self.full_attn_allocator.alloc_extend( + prefix_lens, + prefix_lens_cpu, + seq_lens, + seq_lens_cpu, + last_loc, + extend_num_tokens, + ) + assert alloc_full_indices is not None + + if swa_tail_len == 0: + return alloc_full_indices + + device = self.device + swa_prefix_lens = torch.zeros((1,), dtype=torch.int64, device=device) + swa_prefix_lens_cpu = torch.zeros((1,), dtype=torch.int64) + swa_seq_lens = torch.tensor([swa_tail_len], dtype=torch.int64, device=device) + swa_seq_lens_cpu = torch.tensor([swa_tail_len], dtype=torch.int64) + swa_last_loc = torch.tensor([-1], dtype=torch.int64, device=device) + + alloc_swa_indices = self.swa_attn_allocator.alloc_extend( + swa_prefix_lens, + swa_prefix_lens_cpu, + swa_seq_lens, + swa_seq_lens_cpu, + swa_last_loc, + swa_tail_len, + ) + assert alloc_swa_indices is not None + + self.full_to_swa_index_mapping[alloc_full_indices[-swa_tail_len:]] = ( + alloc_swa_indices + ) + return alloc_full_indices + def alloc_decode( self, seq_lens: torch.Tensor, diff --git a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py index 4dca23453..eba5e34ef 100644 --- a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py +++ b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py @@ -319,6 +319,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase): queue.req_to_token_pool.available_size.return_value = 1 queue.req_to_metadata_buffer_idx_allocator = MagicMock() queue.req_to_metadata_buffer_idx_allocator.available_size.return_value = 1 + queue.token_to_kv_pool = MagicMock() queue.token_to_kv_pool_allocator = MagicMock() queue.token_to_kv_pool_allocator.page_size = 4 @@ -336,7 +337,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase): queue.scheduler = scheduler # Initial budget says the request fits; post-lock budget says it does not. - queue._allocatable_tokens = MagicMock(side_effect=[8, 3]) + queue._allocatable_token_budgets = MagicMock(side_effect=[8, 3]) preallocated, failed = queue.pop_preallocated() @@ -344,7 +345,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase): self.assertEqual(failed, []) queue._pre_alloc.assert_not_called() queue.tree_cache.dec_lock_ref.assert_called_once_with(req.last_node) - self.assertEqual(queue._allocatable_tokens.call_count, 2) + self.assertEqual(queue._allocatable_token_budgets.call_count, 2) def test_repeated_incremental_no_leak(self): """Multiple incremental transfers shouldn't leak lock_refs."""