diff --git a/python/sglang/srt/beam_search/coordinator.py b/python/sglang/srt/beam_search/coordinator.py index 8b224c970..7d3ade09b 100644 --- a/python/sglang/srt/beam_search/coordinator.py +++ b/python/sglang/srt/beam_search/coordinator.py @@ -308,7 +308,7 @@ class BeamCoordinator(msgspec.Struct, kw_only=True): if final: # Prefill-terminated (max_new_tokens == 1): no spawn, but the relay # slot still needs a token so an overshoot step has a valid input. - self._stash_next_tokens([req.req_pool_idx], next_tokens[:1]) + self._stash_next_tokens([req.kv.req_pool_idx], next_tokens[:1]) return self._spawn_member_rows(group, req) req.output_ids.append(0) # length placeholder; DAG owns history @@ -339,13 +339,13 @@ class BeamCoordinator(msgspec.Struct, kw_only=True): alias_members_prompt_kv( self.req_to_token_pool.req_to_token, member_rows, - leader.req_pool_idx, + leader.kv.req_pool_idx, group.prompt_len, ) group.member_rows = member_rows group.member_rows_cpu = torch.tensor(rows, dtype=torch.int64) leader_row = torch.tensor( - [leader.req_pool_idx], dtype=torch.int64, device=device + [leader.kv.req_pool_idx], dtype=torch.int64, device=device ) group.all_rows = torch.cat([leader_row, member_rows]) diff --git a/python/sglang/srt/disaggregation/common/staging_handler.py b/python/sglang/srt/disaggregation/common/staging_handler.py index 900b3eb42..e6b77bbca 100644 --- a/python/sglang/srt/disaggregation/common/staging_handler.py +++ b/python/sglang/srt/disaggregation/common/staging_handler.py @@ -411,7 +411,7 @@ class DecodeStagingHandler: staging_view = self.staging_allocator.buffer.buffer[staging_offset:] - req_pool_idx = decode_req.req.req_pool_idx + req_pool_idx = decode_req.req.kv.req_pool_idx # page_start is suffix-relative (pages after the decode-side cached # prefix); req_to_token rows are absolute. prefix_tokens = decode_req.req.kv.cache_protected_len diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 033726dda..d75239116 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -188,7 +188,7 @@ class DecodeReqToTokenPool: def alloc(self, reqs: List[Req]) -> Optional[List[int]]: # Indices of reqs that already have a req_pool_idx and will reuse # their existing slot (e.g. chunked prefill continuing across chunks). - reusing = [i for i, r in enumerate(reqs) if r.req_pool_idx is not None] + reusing = [i for i, r in enumerate(reqs) if r.kv.req_pool_idx is not None] assert ( len(reusing) <= 1 ), "only one chunked request may reuse req_pool_idx in a batch" @@ -204,16 +204,16 @@ class DecodeReqToTokenPool: self.free_slots = self.free_slots[need_size:] offset = 0 for r in reqs: - if r.req_pool_idx is None: - r.req_pool_idx = select_index[offset] - self.req_generation[r.req_pool_idx] += 1 + if r.kv.req_pool_idx is None: + r.kv.req_pool_idx = select_index[offset] + self.req_generation[r.kv.req_pool_idx] += 1 offset += 1 - return [r.req_pool_idx for r in reqs] + return [r.kv.req_pool_idx for r in reqs] def free(self, req: Req): - assert req.req_pool_idx is not None, "request must have req_pool_idx" - self.free_slots.append(req.req_pool_idx) - req.req_pool_idx = None + assert req.kv.req_pool_idx is not None, "request must have req_pool_idx" + self.free_slots.append(req.kv.req_pool_idx) + req.kv.req_pool_idx = None def clear(self): self.free_slots = list(range(1, self._alloc_size)) @@ -1376,7 +1376,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): else: # Only send delta indices (beyond prefix) to prefill. kv_indices = self.req_to_token_pool.req_to_token[ - decode_req.req.req_pool_idx + decode_req.req.kv.req_pool_idx ][total_prefix_len:origin_input_len] kv_indices = ( self.token_to_kv_pool_allocator.translate_kv_indices_for_transfer( @@ -1390,7 +1390,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): return [ self.req_to_token_pool.translate_mamba_indices( self.req_to_token_pool.req_index_to_mamba_index_mapping[ - decode_req.req.req_pool_idx + decode_req.req.kv.req_pool_idx ] ) .cpu() @@ -1402,7 +1402,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): window_start = max(total_prefix_len, seq_len - window_size) window_start = page_align_floor(window_start, page_size) window_kv_indices_full = self.req_to_token_pool.req_to_token[ - decode_req.req.req_pool_idx, window_start:seq_len + decode_req.req.kv.req_pool_idx, window_start:seq_len ] window_kv_indices_swa = ( self.token_to_kv_pool_allocator.translate_loc_from_full_to_swa( @@ -1413,7 +1413,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): def _full_kv_pages_payload(): kv_indices_full = self.req_to_token_pool.req_to_token[ - decode_req.req.req_pool_idx, :seq_len + decode_req.req.kv.req_pool_idx, :seq_len ] # Indexer lives on device pool; always use device page_size device_page_size = self.token_to_kv_pool.page_size @@ -1426,7 +1426,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): window_size = self.token_to_kv_pool.unified_swa_window window_start = max(0, seq_len - window_size) positions = np.arange(window_start, seq_len, dtype=np.int64) - state_slot = int(decode_req.req.req_pool_idx) + state_slot = int(decode_req.req.kv.req_pool_idx) ring_rows = state_slot * ring_stride + (positions % ring_stride) return ring_rows.astype(np.int32) @@ -1434,7 +1434,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): online = is_dsv4_c128_online_enabled() ring_size = 1 if online else self.token_to_kv_pool.get_ring_size(128) return get_dsv4_c128_state_indices( - int(decode_req.req.req_pool_idx), + int(decode_req.req.kv.req_pool_idx), seq_len, online=online, ring_size=ring_size, @@ -1446,7 +1446,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): self.token_to_kv_pool, "clear_c128_req_state", None ) if clear_c128_state is not None: - clear_c128_state(int(decode_req.req.req_pool_idx)) + clear_c128_state(int(decode_req.req.kv.req_pool_idx)) payloads = { StateType.MAMBA: _mamba_payload, StateType.SWA: _swa_payload, @@ -1465,7 +1465,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): payloads.update( dsv4_state_payloads( self.req_to_token_pool, - decode_req.req.req_pool_idx, + decode_req.req.kv.req_pool_idx, seq_len, self.token_to_kv_pool_allocator.page_size, prefix_len=total_prefix_len, @@ -1494,7 +1494,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): # the C4 sparse physical-slot mapping; carry their logical page IDs # alongside the independently allocated C4 host page IDs. full_kv_indices = self.req_to_token_pool.req_to_token[ - decode_req.req.req_pool_idx, + decode_req.req.kv.req_pool_idx, prefix_len:origin_input_len, ] device_page_indices = kv_to_page_indices( @@ -1787,7 +1787,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): if prefix_len > 0: self.req_to_token_pool.write( - (req.req_pool_idx, slice(0, prefix_len)), prefix_indices + (req.kv.req_pool_idx, slice(0, prefix_len)), prefix_indices ) # TODO(retraction): when retraction is implemented with radix cache @@ -1842,7 +1842,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): host_indices = coordinator.mem_pool_host.alloc_paged_token_slots( coordinator.req_to_host_pool, coordinator.req_to_host_pool_allocated_len, - req.req_pool_idx, + req.kv.req_pool_idx, 0, coordinator.host_token_len(fill_len), ) @@ -1872,7 +1872,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): self.req_to_token_pool.write( ( - req.req_pool_idx, + req.kv.req_pool_idx, slice(total_prefix_len, total_prefix_len + len(kv_loc)), ), kv_loc, diff --git a/python/sglang/srt/disaggregation/decode_hicache_mixin.py b/python/sglang/srt/disaggregation/decode_hicache_mixin.py index 588056fa3..fb39cf3e4 100644 --- a/python/sglang/srt/disaggregation/decode_hicache_mixin.py +++ b/python/sglang/srt/disaggregation/decode_hicache_mixin.py @@ -306,7 +306,7 @@ class DecodeHiCacheTransferMixin: self.tree_cache.req_to_token_pool.write( ( - decode_req.req.req_pool_idx, + decode_req.req.kv.req_pool_idx, slice(prefix_match.l1_prefix_len, prefix_match.decode_prefix_len), ), decode_req.hicache_restored_kv_indices, diff --git a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py index c441a88d5..69cc858ef 100644 --- a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py +++ b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py @@ -123,10 +123,10 @@ class DecodeKVCacheOffloadManager: if self.cache_controller is None or self.decode_host_mem_pool is None: return False - if req.req_pool_idx == -1 or len(req.output_ids) == 0: + if req.kv.req_pool_idx == -1 or len(req.output_ids) == 0: return False - token_indices = self.req_to_token_pool.req_to_token[req.req_pool_idx] + token_indices = self.req_to_token_pool.req_to_token[req.kv.req_pool_idx] if token_indices.dim() == 0 or token_indices.numel() == 0: return False @@ -251,7 +251,7 @@ class DecodeKVCacheOffloadManager: # so a previously-released request must be skipped here to avoid # non-idempotent side effects (e.g. tree_cache.protected_size_ # double-decrement, host pool double-free). - if req.req_pool_idx is None or req.req_pool_idx == -1: + if req.kv.req_pool_idx is None or req.kv.req_pool_idx == -1: return kv_committed_len = req.effective_kv_committed_len() @@ -264,13 +264,13 @@ class DecodeKVCacheOffloadManager: state = self.offloaded_state.get(req) if state is not None and state.prefill_len > 0: prefill_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, : state.prefill_len + req.kv.req_pool_idx, : state.prefill_len ] self.token_to_kv_pool_allocator.free(prefill_indices) start = start_offset end = kv_committed_len # Free the incremental part of the request (DSA-aware) - kv_indices = self.req_to_token_pool.req_to_token[req.req_pool_idx, start:end] + kv_indices = self.req_to_token_pool.req_to_token[req.kv.req_pool_idx, start:end] self.token_to_kv_pool_allocator.free(kv_indices) # Free over-allocated KV cache slots (e.g. from speculative decoding v2). @@ -280,7 +280,7 @@ class DecodeKVCacheOffloadManager: start_p = ceil_align(start_p, self.page_size) if start_p < end_p: overalloc_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, start_p:end_p + req.kv.req_pool_idx, start_p:end_p ] self.token_to_kv_pool_allocator.free(overalloc_indices) @@ -329,7 +329,7 @@ class DecodeKVCacheOffloadManager: """Free any remaining tail KV that was not offloaded due to non-aligned length.""" # ReqToTokenPool.free sets req_pool_idx to None on release, so # guard against both sentinels here. - if req.req_pool_idx is None or req.req_pool_idx == -1: + if req.kv.req_pool_idx is None or req.kv.req_pool_idx == -1: return state = self.offloaded_state.get(req) if state is None: diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index c813c30cd..836c93332 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py @@ -43,10 +43,10 @@ class ScheduleBatchDisaggregationDecodeMixin: # Fill the tensor in one pass offset = 0 for i, req in enumerate(reqs): - req_pool_indices.append(req.req_pool_idx) + req_pool_indices.append(req.kv.req_pool_idx) pre_len = len(req.prefix_indices) - chunk = self.req_to_token_pool.req_to_token[req.req_pool_idx][ + chunk = self.req_to_token_pool.req_to_token[req.kv.req_pool_idx][ pre_len : pre_len + req.extend_range.length ] assert ( diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index b2310f133..4bc257981 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -1036,7 +1036,7 @@ class SchedulerDisaggregationPrefillMixin: else: logger.warning(error_message) req.time_stats.trace_ctx.abort(abort_info={"reason": error_message}) - if req.is_holding_kv or req.mamba_pool_idx is not None: + if req.kv.is_held or req.mamba_pool_idx is not None: release_kv_cache(req, self.tree_cache) maybe_release_metadata_buffer(req, self.req_to_metadata_buffer_idx_allocator) req.pending_bootstrap = False @@ -1215,7 +1215,7 @@ class SchedulerDisaggregationPrefillMixin: return [ self.req_to_token_pool.translate_mamba_indices( self.req_to_token_pool.req_index_to_mamba_index_mapping[ - req.req_pool_idx + req.kv.req_pool_idx ] ) .cpu() @@ -1227,7 +1227,7 @@ class SchedulerDisaggregationPrefillMixin: window_start = max(req.disagg_decode_prefix_len, seq_len - window_size) window_start = (window_start // page_size) * page_size window_kv_indices_full = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, window_start:seq_len + req.kv.req_pool_idx, window_start:seq_len ] window_kv_indices_swa = ( self.token_to_kv_pool_allocator.translate_loc_from_full_to_swa( @@ -1238,7 +1238,7 @@ class SchedulerDisaggregationPrefillMixin: def _full_kv_pages_payload(): kv_indices_full = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :seq_len + req.kv.req_pool_idx, :seq_len ] return kv_to_page_indices(kv_indices_full, page_size) @@ -1251,7 +1251,7 @@ class SchedulerDisaggregationPrefillMixin: window_size = _pool.unified_swa_window window_start = max(0, seq_len - window_size) positions = np.arange(window_start, seq_len, dtype=np.int64) - state_slot = int(req.req_pool_idx) + state_slot = int(req.kv.req_pool_idx) ring_rows = state_slot * ring_stride + (positions % ring_stride) return ring_rows.astype(np.int32) @@ -1265,7 +1265,7 @@ class SchedulerDisaggregationPrefillMixin: ) ) return get_dsv4_c128_state_indices( - int(req.req_pool_idx), + int(req.kv.req_pool_idx), c128_seq_len, online=online, ring_size=ring_size, @@ -1295,7 +1295,7 @@ class SchedulerDisaggregationPrefillMixin: payloads.update( dsv4_state_payloads( self.req_to_token_pool, - req.req_pool_idx, + req.kv.req_pool_idx, seq_len, page_size, prefix_len=req.disagg_decode_prefix_len, @@ -1321,7 +1321,7 @@ class SchedulerDisaggregationPrefillMixin: for seg_start, seg_end in segments: is_final_segment = seg_end == end_idx kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, seg_start:seg_end + req.kv.req_pool_idx, seg_start:seg_end ] # Unified memory: req_to_token holds VIRTUAL ids; the transfer needs # physical ones. Per segment, since each is its own gather. diff --git a/python/sglang/srt/hardware_backend/mlx/tp_worker.py b/python/sglang/srt/hardware_backend/mlx/tp_worker.py index edd5256f5..813b4a335 100644 --- a/python/sglang/srt/hardware_backend/mlx/tp_worker.py +++ b/python/sglang/srt/hardware_backend/mlx/tp_worker.py @@ -502,7 +502,7 @@ class MlxTpModelWorker(TpModelWorker): full_token_ids=full_token_ids, prefix_slot_ids=prefix_slot_ids, new_slot_ids=req_new_slots, - req_pool_idx=req.req_pool_idx, + req_pool_idx=req.kv.req_pool_idx, req=req, needs_logits=self._chunk_needs_logits(req), logit_edit_row=edit_rows[req.rid] if edit_rows else None, diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/c128_sidecar_component.py b/python/sglang/srt/hardware_backend/npu/dsv4/c128_sidecar_component.py index b4fda7839..a7b338d98 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/c128_sidecar_component.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/c128_sidecar_component.py @@ -259,7 +259,7 @@ class C128SidecarComponent(TreeComponent): cache_len = logical_len // group_tokens * group_tokens num_pages = cache_len // group_tokens insert_params.c128_value = self.cache.req_to_token_pool.req_to_c128_sidecar[ - int(req.req_pool_idx), :num_pages + int(req.kv.req_pool_idx), :num_pages ].clone() return cache_len + 1 if self.tree_core.is_eagle and cache_len > 0 else cache_len diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py index 0737508de..6b9d8913d 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py @@ -466,7 +466,7 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): if req is None or req_to_token_pool is None: return kv_len = max(req.kv.kv_committed_len, req.kv.kv_allocated_len) - req_pool_idx = req.req_pool_idx + req_pool_idx = req.kv.req_pool_idx if kv_len <= 0 or req_pool_idx is None: return diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_common_hooks.py b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_common_hooks.py index 5cc5761aa..3cda73847 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_common_hooks.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_common_hooks.py @@ -104,7 +104,7 @@ def dsv4_prealloc_kwargs(allocator, req, fill_len, req_to_token_pool, *, device) return {} return dict( req_pool_indices=torch.tensor( - [req.req_pool_idx], dtype=torch.int64, device=device + [req.kv.req_pool_idx], dtype=torch.int64, device=device ), req_to_token_pool=req_to_token_pool, ) @@ -130,7 +130,7 @@ def write_dsv4_prealloc_tables( prealloc path (no ScheduleBatch); no-op without bundle / DSV4 tables.""" if bundle is None or not hasattr(req_to_token_pool, "write_c128"): return - rp = torch.tensor([req.req_pool_idx]) + rp = torch.tensor([req.kv.req_pool_idx]) pl = torch.tensor([prefix_len]) sl = torch.tensor([fill_len]) diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_req_to_token_pool.py b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_req_to_token_pool.py index 35369cc17..00b40f081 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_req_to_token_pool.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_req_to_token_pool.py @@ -83,15 +83,15 @@ class DSV4ReqToTokenTablesMixin: Prefix matching can happen before a request slot is allocated, so the page ids are temporarily carried by ``Req`` and installed by ``alloc``. """ - if req.req_pool_idx is None: + if req.kv.req_pool_idx is None: req.c128_prefix_page_ids = page_ids return self._dsv4_allocator.replace_req_c128_prefix( - int(req.req_pool_idx), page_ids, self + int(req.kv.req_pool_idx), page_ids, self ) def alloc(self, reqs): - fresh = [req.req_pool_idx is None for req in reqs] + fresh = [req.kv.req_pool_idx is None for req in reqs] indices = super().alloc(reqs) if indices is None: return None diff --git a/python/sglang/srt/managers/hisparse_coordinator.py b/python/sglang/srt/managers/hisparse_coordinator.py index 430bd0692..061d7545c 100644 --- a/python/sglang/srt/managers/hisparse_coordinator.py +++ b/python/sglang/srt/managers/hisparse_coordinator.py @@ -349,7 +349,7 @@ class HiSparseCoordinator: req.hisparse_staging = True full_kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, : req.extend_range.end + req.kv.req_pool_idx, : req.extend_range.end ].to(dtype=torch.int64, copy=True) device_indices = ( self.mem_pool_device.translate_loc_from_full_to_hisparse_device( @@ -361,7 +361,7 @@ class HiSparseCoordinator: host_indices = self.mem_pool_host.alloc_paged_token_slots( self.req_to_host_pool, self.req_to_host_pool_allocated_len, - req.req_pool_idx, + req.kv.req_pool_idx, 0, prefill_len, ) @@ -411,11 +411,11 @@ class HiSparseCoordinator: # Long sequence: reset device_buffer_tokens to -1 so the kernel # sees all slots as empty -> every top-k lookup is a miss -> host load. self.req_device_buffer_tokens[ - :, req.req_pool_idx, : self.device_buffer_size + :, req.kv.req_pool_idx, : self.device_buffer_size ] = -1 req.hisparse_staging = False - self._skip_first_backup[req.req_pool_idx] = True + self._skip_first_backup[req.kv.req_pool_idx] = True logger.debug("HiSparse: admitting request %s directly", req.rid) def host_token_len(self, kv_allocated_len: int) -> int: @@ -426,8 +426,8 @@ class HiSparseCoordinator: def _preload_to_device_buffer(self, req: Req) -> None: """Preload all tokens from host pool into the device buffer.""" n = self.host_token_len(req.kv.kv_allocated_len) - host_indices = self.req_to_host_pool[req.req_pool_idx, :n] - device_locs = self.req_to_device_buffer[req.req_pool_idx, :n] + host_indices = self.req_to_host_pool[req.kv.req_pool_idx, :n] + device_locs = self.req_to_device_buffer[req.kv.req_pool_idx, :n] for layer_id in range(self.mem_pool_device.layer_num): self.mem_pool_host.load_to_device_per_layer( @@ -456,7 +456,7 @@ class HiSparseCoordinator: compressed_logical_indices = ( self.mem_pool_device.translate_loc_from_full_to_compressed( - self.req_to_token_pool.req_to_token[req.req_pool_idx, :allocated_len] + self.req_to_token_pool.req_to_token[req.kv.req_pool_idx, :allocated_len] ) ) compressed_len = len(compressed_logical_indices) @@ -475,13 +475,13 @@ class HiSparseCoordinator: raise RuntimeError("HiSparse alloc_device_buffer returned None") buffer_indices = buffer_indices.to(torch.int32) - self.req_to_device_buffer[req.req_pool_idx, :alloc_size] = buffer_indices - self.req_device_buffer_size[req.req_pool_idx] = alloc_size + self.req_to_device_buffer[req.kv.req_pool_idx, :alloc_size] = buffer_indices + self.req_device_buffer_size[req.kv.req_pool_idx] = alloc_size self.req_device_buffer_tokens[ - :, req.req_pool_idx, : self.device_buffer_size + :, req.kv.req_pool_idx, : self.device_buffer_size ] = self._device_buffer_arange_i32 - self.req_device_buffer_token_locs[:, req.req_pool_idx, :alloc_size] = ( + self.req_device_buffer_token_locs[:, req.kv.req_pool_idx, :alloc_size] = ( buffer_indices[:alloc_size] ) @@ -584,7 +584,7 @@ class HiSparseCoordinator: _, _, req = self.ack_staging_queue.pop(0) # prepare device buffer and update req self.alloc_device_buffer(req) - self._skip_first_backup[req.req_pool_idx] = True + self._skip_first_backup[req.kv.req_pool_idx] = True req.hisparse_staging = False finish_count -= 1 ready_reqs.append(req) @@ -863,21 +863,21 @@ class HiSparseCoordinator: prefill_len = req.extend_range.end allocated_locs = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :prefill_len + req.kv.req_pool_idx, :prefill_len ] self.token_to_kv_pool_allocator.free_hisparse(allocated_locs) # Free host memory that was allocated during admit_request_into_staging host_indices = self.mem_pool_host.allocated_host_indices( self.req_to_host_pool, - req.req_pool_idx, - self.req_to_host_pool_allocated_len[req.req_pool_idx], + req.kv.req_pool_idx, + self.req_to_host_pool_allocated_len[req.kv.req_pool_idx], ) if host_indices.numel() > 0: self.mem_pool_host.free(host_indices) - self.req_to_host_pool[req.req_pool_idx, :] = -1 - self.req_to_host_pool_allocated_len[req.req_pool_idx] = 0 - self._skip_first_backup[req.req_pool_idx] = False + self.req_to_host_pool[req.kv.req_pool_idx, :] = -1 + self.req_to_host_pool_allocated_len[req.kv.req_pool_idx] = 0 + self._skip_first_backup[req.kv.req_pool_idx] = False req.hisparse_staging = False def retract_req(self, req: Req) -> None: @@ -901,15 +901,15 @@ class HiSparseCoordinator: allocated_len = req.kv.kv_allocated_len # release memory -- only free actually-allocated buffer indices - current_cap = int(self.req_device_buffer_size[req.req_pool_idx]) + current_cap = int(self.req_device_buffer_size[req.kv.req_pool_idx]) if current_cap > 0: - side_buf_hi = self.req_to_device_buffer[req.req_pool_idx, :current_cap] + side_buf_hi = self.req_to_device_buffer[req.kv.req_pool_idx, :current_cap] all_hi = torch.unique(side_buf_hi[side_buf_hi > 0]) if all_hi.numel() > 0: self.token_to_kv_pool_allocator.free_hisparse_indices(all_hi) allocated_locs = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :allocated_len + req.kv.req_pool_idx, :allocated_len ] compressed_locs = self.mem_pool_device.translate_loc_from_full_to_compressed( allocated_locs @@ -918,21 +918,21 @@ class HiSparseCoordinator: host_indices = self.mem_pool_host.allocated_host_indices( self.req_to_host_pool, - req.req_pool_idx, - self.req_to_host_pool_allocated_len[req.req_pool_idx], + req.kv.req_pool_idx, + self.req_to_host_pool_allocated_len[req.kv.req_pool_idx], ) if host_indices.numel() > 0: self.mem_pool_host.free(host_indices) # clear req info - self.req_device_buffer_tokens[:, req.req_pool_idx, :] = -1 - self.req_device_buffer_token_locs[:, req.req_pool_idx, :] = -1 - self.req_to_device_buffer[req.req_pool_idx, :] = 0 - self.req_device_buffer_size[req.req_pool_idx] = 0 - self.req_to_host_pool[req.req_pool_idx, :] = -1 - self.req_to_host_pool_allocated_len[req.req_pool_idx] = 0 - self.lru_slots[:, req.req_pool_idx, :].copy_(self._lru_init) - self._skip_first_backup[req.req_pool_idx] = False + self.req_device_buffer_tokens[:, req.kv.req_pool_idx, :] = -1 + self.req_device_buffer_token_locs[:, req.kv.req_pool_idx, :] = -1 + self.req_to_device_buffer[req.kv.req_pool_idx, :] = 0 + self.req_device_buffer_size[req.kv.req_pool_idx] = 0 + self.req_to_host_pool[req.kv.req_pool_idx, :] = -1 + self.req_to_host_pool_allocated_len[req.kv.req_pool_idx] = 0 + self.lru_slots[:, req.kv.req_pool_idx, :].copy_(self._lru_init) + self._skip_first_backup[req.kv.req_pool_idx] = False def _run_swap_in_kernel( self, diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 7f4733887..c3e5a562b 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -816,7 +816,8 @@ class ReqLogprob: @dataclasses.dataclass(slots=True, kw_only=True) class ReqKvInfo: # Device KV a request holds outside the prefix cache. Always present on the Req; - # whether any KV is held is `req.req_pool_idx is not None` (Req.is_holding_kv). + # whether any KV is held is `is_held` (a row is registered). + req_pool_idx: Optional[int] = None # req_to_token row, the register for the slots # 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) @@ -835,6 +836,10 @@ class ReqKvInfo: lo = ceil_align(lo, page_size) return lo + @property + def is_held(self) -> bool: + return self.req_pool_idx is not None + @property def is_released(self) -> bool: return self.kv_allocated_len == 0 and self.swa_evicted_seqlen == 0 @@ -970,7 +975,6 @@ class Req(ReqDllmMixin): self.routing_key = routing_key # Memory pool info - self.req_pool_idx: Optional[int] = None self.mamba_pool_idx: Optional[torch.Tensor] = None # shape (1) self.mamba_ping_pong_track_buffer: Optional[torch.Tensor] = None # shape (2) self.mamba_next_track_idx: Optional[int] = None # 0 or 1 @@ -1278,10 +1282,6 @@ class Req(ReqDllmMixin): or self.mamba_host_hit_length > 0 ) - @property - def is_holding_kv(self) -> bool: - return self.req_pool_idx is not None - def effective_kv_committed_len(self) -> int: # Report only the prompt prefix so thinking + answer fall into the # overallocated range and are reclaimed by release_kv_cache. #22373. @@ -1748,7 +1748,7 @@ class Req(ReqDllmMixin): self.mamba_cow_src_index = None self.mamba_needs_clear = False self.already_computed = 0 - assert not self.is_holding_kv, "expect it is already released" + assert not self.kv.is_held, "expect it is already released" self.kv.kv_committed_len = 0 self.extend_batch_idx = 0 self.decode_batch_idx = 0 @@ -1773,7 +1773,7 @@ class Req(ReqDllmMixin): def offload_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator): token_indices = req_to_token_pool.req_to_token[ - self.req_pool_idx, : self.seqlen - 1 + self.kv.req_pool_idx, : self.seqlen - 1 ] # Copies over both the kv cache and mamba state if available mamba_pool = self._mamba_pool_needing_backup( @@ -1793,7 +1793,7 @@ class Req(ReqDllmMixin): def load_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator): assert self.retraction_backup is not None token_indices = req_to_token_pool.req_to_token[ - self.req_pool_idx, : self.seqlen - 1 + self.kv.req_pool_idx, : self.seqlen - 1 ] # Loads both the kv cache and mamba state if exists mamba_cpu = self.retraction_backup.mamba_cpu @@ -3492,7 +3492,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # seqlen progress is monotonic per KV handle. if ( req.decode_batch_idx >= 1 - and req.is_holding_kv + and req.kv.is_held and req.seqlen - 1 - sliding_window_size >= req.kv.swa_evicted_seqlen + eviction_interval ): diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 8dc2f21b9..e9c0ec8df 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -3124,7 +3124,7 @@ class Scheduler( if self.chunked_req is not req: # Already past chunked prefill; the running-batch abort path handles # it. Drop the marker once the request is actually gone. - if req.finished() or req.req_pool_idx is None: + if req.finished() or req.kv.req_pool_idx is None: self._pending_chunked_abort_req = None return @@ -3161,7 +3161,7 @@ class Scheduler( spec_algorithm=self.spec_algorithm, ) - req_pool_indices = [r.req_pool_idx for r in reqs] + req_pool_indices = [r.kv.req_pool_idx for r in reqs] batch.req_pool_indices = torch.tensor( req_pool_indices, dtype=torch.int64, device=device ) @@ -4868,7 +4868,7 @@ class Scheduler( _make_abort_req(req), req ) if ( - req.req_pool_idx is not None + req.kv.req_pool_idx is not None or getattr(req, "mamba_pool_idx", None) is not None ): release_kv_cache(req, self.tree_cache, is_insert=False) diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index 832c222fe..d1fba515d 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -136,7 +136,7 @@ class SchedulerBatchResultProcessor: start_len = req.routed_experts_start_len seqlen = len(req.origin_input_ids) + len(req.output_ids_through_stop) req.routed_experts = capturer.get_topk( - req_pool_idx=req.req_pool_idx, + req_pool_idx=req.kv.req_pool_idx, seqlen=seqlen, req_to_token_pool=self.req_to_token_pool, start_len=start_len, @@ -166,7 +166,7 @@ class SchedulerBatchResultProcessor: return seqlen = len(req.origin_input_ids) + len(req.output_ids_through_stop) req.indexer_topk = capturer.get_topk( - req_pool_idx=req.req_pool_idx, + req_pool_idx=req.kv.req_pool_idx, seqlen=seqlen, req_to_token_pool=self.req_to_token_pool, ) diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index a3557302c..39ea63bec 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -252,7 +252,7 @@ class SchedulerInvariantChecker: swa_uncached = 0 for batch in batches: for req in batch.reqs: - if not req.is_holding_kv: + if not req.kv.is_held: continue allocated_len = req.kv.kv_allocated_len @@ -324,23 +324,23 @@ class SchedulerInvariantChecker: batch = self.get_last_batch() if batch is not None: for req in batch.reqs: - if not req.is_holding_kv: + if not req.kv.is_held: continue _add_owner( req, f"req {req.rid}", - req.req_pool_idx, + req.kv.req_pool_idx, req.kv.kv_committed_len, req.kv.kv_allocated_len, ) sess = getattr(self.tree_cache, "slots", None) if sess: for sid, slot in sess.items(): - if getattr(slot, "is_holding_kv", False): + if slot.kv.is_held: _add_owner( slot, f"slot {sid[:8]}", - slot.req_pool_idx, + slot.kv.req_pool_idx, slot.kv.kv_committed_len, slot.kv.kv_allocated_len, ) diff --git a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py index dd9cdc7d7..f653c5f68 100644 --- a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py +++ b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py @@ -172,8 +172,8 @@ class SchedulerPoolStatsObserver: if batch is None or batch.is_empty(): continue for req in batch.reqs: - if req.req_pool_idx is not None: - idxs.add(req.req_pool_idx) + if req.kv.req_pool_idx is not None: + idxs.add(req.kv.req_pool_idx) return idxs def session_held_tokens(self) -> int: diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index c58b1f6d3..e49224e49 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -723,9 +723,9 @@ class SchedulerPPMixin: latencies.append(latency_ms) # Release KV and Mamba cache - if req.is_holding_kv: + if req.kv.is_held: kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, : req.extend_range.end + req.kv.req_pool_idx, : req.extend_range.end ] self.token_to_kv_pool_allocator.free(kv_indices) if req.mamba_pool_idx is not None: diff --git a/python/sglang/srt/mem_cache/allocation.py b/python/sglang/srt/mem_cache/allocation.py index 5e3044757..fe8a2e858 100644 --- a/python/sglang/srt/mem_cache/allocation.py +++ b/python/sglang/srt/mem_cache/allocation.py @@ -297,7 +297,7 @@ def alloc_for_extend( reuse_kv = None if batch.is_dllm(): reuse_kv = [ - r.req_pool_idx is not None and bool(r.dllm_incomplete_ids) + r.kv.req_pool_idx is not None and bool(r.dllm_incomplete_ids) for r in batch.reqs ] diff --git a/python/sglang/srt/mem_cache/chunk_cache.py b/python/sglang/srt/mem_cache/chunk_cache.py index 854d35ff4..cdc79ce49 100644 --- a/python/sglang/srt/mem_cache/chunk_cache.py +++ b/python/sglang/srt/mem_cache/chunk_cache.py @@ -82,13 +82,13 @@ class ChunkCache(BasePrefixCache): # For decode server: if req.output_ids is empty, we want to free all req.origin_input_ids # The protected prefix is not this req's to free. kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, req.kv.cache_protected_len : kv_len_to_handle + req.kv.req_pool_idx, req.kv.cache_protected_len : kv_len_to_handle ] self.token_to_kv_pool_allocator.free(kv_indices) def cache_unfinished_req(self, req: Req, chunked=False): kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, : req.extend_range.end + req.kv.req_pool_idx, : req.extend_range.end ] # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True) @@ -158,7 +158,7 @@ class PureSWAChunkCache(SWAChunkCache): ): kv_committed_len = kv_len_to_handle kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :kv_committed_len + req.kv.req_pool_idx, :kv_committed_len ] # The cache_protected_len prefix is not this req's to free. protected_len = req.kv.cache_protected_len diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 93d814f38..5627a27e4 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -62,7 +62,7 @@ def free_swa_out_of_window_slots( is_chunk_cache: bool = False, retain_floor: int | None = None, ) -> None: - if not req.is_holding_kv: + if not req.kv.is_held: return # For swa radix cache, we need to evict the tokens that are not in the tree cache and also not in the sliding window @@ -99,7 +99,7 @@ def free_swa_out_of_window_slots( if new_swa_evicted_seqlen > req.kv.swa_evicted_seqlen: free_slots = req_to_token_pool.req_to_token[ - req.req_pool_idx, req.kv.swa_evicted_seqlen : new_swa_evicted_seqlen + req.kv.req_pool_idx, req.kv.swa_evicted_seqlen : new_swa_evicted_seqlen ] token_to_kv_pool_allocator.free_swa(free_slots) req.kv.swa_evicted_seqlen = new_swa_evicted_seqlen @@ -199,9 +199,9 @@ def retraction_discard(req: Req, tree_cache: BasePrefixCache, backend: str) -> N def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = True): - assert (not req.is_holding_kv) == req.kv.is_released + assert (not req.kv.is_held) == req.kv.is_released # MambaRadixCache may alloc mamba state before alloc KV cache - if not req.is_holding_kv: + if not req.kv.is_held: assert ( tree_cache.supports_mamba() ), "Only MambaRadixCache allow freeing before alloc" @@ -222,8 +222,8 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr # StreamingSession.cache_finished_req handles speculative tail trim # internally, then sets req_pool_idx = None. - assert (not req.is_holding_kv) == req.kv.is_released - if not req.is_holding_kv: + assert (not req.kv.is_held) == req.kv.is_released + if not req.kv.is_held: return start_p, end_p = effective_kv_committed_len, req.kv.kv_allocated_len @@ -261,9 +261,9 @@ def _release_overallocated_kv_indices( start_p = ceil_align(start_p, page_size) if start_p < end_p: - indices_to_free = tree_cache.req_to_token_pool.req_to_token[req.req_pool_idx][ - start_p:end_p - ] + indices_to_free = tree_cache.req_to_token_pool.req_to_token[ + req.kv.req_pool_idx + ][start_p:end_p] # start_p is aligned to the allocator's physical page size above, so it # never shares a page with cache_finished_req's tail free in this group. allocator.free_segment(indices_to_free, start_pos=start_p) diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 125a9847b..14b432087 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -548,7 +548,7 @@ class MambaRadixCache(BasePrefixCache): """Cache request when it finishes.""" if self.disable: kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :kv_len_to_handle + req.kv.req_pool_idx, :kv_len_to_handle ] self.token_to_kv_pool_allocator.free_segment(kv_indices, start_pos=0) self.req_to_token_pool.free_mamba_cache(req) @@ -556,7 +556,7 @@ class MambaRadixCache(BasePrefixCache): token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle] kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :kv_len_to_handle + req.kv.req_pool_idx, :kv_len_to_handle ] if is_insert: @@ -676,7 +676,7 @@ class MambaRadixCache(BasePrefixCache): def _skip_cache_unfinished_req(req: Req) -> None: kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, : req.extend_range.end + req.kv.req_pool_idx, : req.extend_range.end ] # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later @@ -693,7 +693,7 @@ class MambaRadixCache(BasePrefixCache): return _skip_cache_unfinished_req(req) kv_indices_orig = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, : len(token_ids) + req.kv.req_pool_idx, : len(token_ids) ] # kv_indices is the kv indices to be cached kv_indices = kv_indices_orig[:cache_len] @@ -786,7 +786,7 @@ class MambaRadixCache(BasePrefixCache): ), f"{new_prefix_len=}, {len(new_indices)=}" self.req_to_token_pool.write( - (req.req_pool_idx, slice(req.kv.cache_protected_len, len(new_indices))), + (req.kv.req_pool_idx, slice(req.kv.cache_protected_len, len(new_indices))), new_indices[req.kv.cache_protected_len :], ) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 5191304aa..03df6688b 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -292,7 +292,7 @@ class ReqToTokenPool: def alloc(self, reqs: list[Req]) -> Optional[List[int]]: # Indices of reqs that already have a req_pool_idx and will reuse # their existing slot (e.g. chunked prefill continuing across chunks). - reusing = [i for i, r in enumerate(reqs) if r.req_pool_idx is not None] + reusing = [i for i, r in enumerate(reqs) if r.kv.req_pool_idx is not None] # NOTE: this check is relaxed temporarily # https://github.com/sgl-project/sglang/pull/20476 # if not any(r.is_dllm() for r in reqs): @@ -309,10 +309,10 @@ class ReqToTokenPool: return None offset = 0 for r in reqs: - if r.req_pool_idx is None: - r.req_pool_idx = select_index[offset] + if r.kv.req_pool_idx is None: + r.kv.req_pool_idx = select_index[offset] offset += 1 - return [r.req_pool_idx for r in reqs] + return [r.kv.req_pool_idx for r in reqs] def alloc_rows(self, need_size: int) -> Optional[List[int]]: """Take need_size rows and bump their generation, with no Req bound to @@ -338,9 +338,9 @@ class ReqToTokenPool: self.free_slots.extend(indices) def free(self, req: Req): - assert req.req_pool_idx is not None, "request must have req_pool_idx" - self.free_rows([req.req_pool_idx]) - req.req_pool_idx = None + assert req.kv.req_pool_idx is not None, "request must have req_pool_idx" + self.free_rows([req.kv.req_pool_idx]) + req.kv.req_pool_idx = None def clear(self): self.free_slots = list(range(1, self._alloc_size)) @@ -1493,7 +1493,7 @@ class HybridReqToTokenPool(ReqToTokenPool): set_mamba_track_indices_from_reqs reads correct slot indices. """ req.mamba_ping_pong_track_buffer[idx] = value - self.req_index_to_mamba_ping_pong_track_buffer_mapping[req.req_pool_idx] = ( + self.req_index_to_mamba_ping_pong_track_buffer_mapping[req.kv.req_pool_idx] = ( req.mamba_ping_pong_track_buffer ) @@ -1530,7 +1530,9 @@ class HybridReqToTokenPool(ReqToTokenPool): if self.enable_mamba_extra_buffer: mamba_ping_pong_track_buffer_to_free = ( - self.req_index_to_mamba_ping_pong_track_buffer_mapping[req.req_pool_idx] + self.req_index_to_mamba_ping_pong_track_buffer_mapping[ + req.kv.req_pool_idx + ] ) if mamba_ping_pong_track_buffer_to_keep is not None: assert mamba_ping_pong_track_buffer_to_keep in [ diff --git a/python/sglang/srt/mem_cache/pure_swa_radix_cache.py b/python/sglang/srt/mem_cache/pure_swa_radix_cache.py index 256300554..071d155c1 100644 --- a/python/sglang/srt/mem_cache/pure_swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/pure_swa_radix_cache.py @@ -78,14 +78,14 @@ class PureSWARadixCache(RadixCache): kv_committed_len = kv_len_to_handle if self.disable: kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :kv_committed_len + req.kv.req_pool_idx, :kv_committed_len ] self.token_to_kv_pool_allocator.free(kv_indices) return token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len] kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :kv_committed_len + req.kv.req_pool_idx, :kv_committed_len ] radix_key = RadixKey( diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 6552854e0..8cebee66a 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -467,7 +467,7 @@ class RadixCache(BasePrefixCache): if self.disable: # The protected prefix is not this req's to free. kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, req.kv.cache_protected_len : kv_len_to_handle + req.kv.req_pool_idx, req.kv.cache_protected_len : kv_len_to_handle ] self.token_to_kv_pool_allocator.free_segment( kv_indices, start_pos=req.kv.cache_protected_len @@ -476,7 +476,7 @@ class RadixCache(BasePrefixCache): token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle] kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, : len(token_ids) + req.kv.req_pool_idx, : len(token_ids) ] radix_key = RadixKey( @@ -520,7 +520,7 @@ class RadixCache(BasePrefixCache): token_ids = req.get_fill_ids() kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, : len(token_ids) + req.kv.req_pool_idx, : len(token_ids) ] radix_key = RadixKey( @@ -558,7 +558,7 @@ class RadixCache(BasePrefixCache): ), f"{len(new_indices)=}, {len(radix_key)=}" self.req_to_token_pool.write( - (req.req_pool_idx, slice(req.kv.cache_protected_len, len(new_indices))), + (req.kv.req_pool_idx, slice(req.kv.cache_protected_len, len(new_indices))), new_indices[req.kv.cache_protected_len :], ) diff --git a/python/sglang/srt/mem_cache/radix_cache_cpp.py b/python/sglang/srt/mem_cache/radix_cache_cpp.py index ca5bcedcf..547214104 100644 --- a/python/sglang/srt/mem_cache/radix_cache_cpp.py +++ b/python/sglang/srt/mem_cache/radix_cache_cpp.py @@ -185,10 +185,10 @@ class RadixCacheCpp(BasePrefixCache): ): """Cache request when it finishes.""" self._reject_cache_salt(req.cache_salt) - assert req.req_pool_idx is not None + assert req.kv.req_pool_idx is not None token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle] kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :kv_len_to_handle + req.kv.req_pool_idx, :kv_len_to_handle ].to(dtype=torch.int64, copy=True) # NOTE: our C++ implementation don't need `token_ids` and `kv_indices` to be page-aligned @@ -223,11 +223,11 @@ class RadixCacheCpp(BasePrefixCache): def cache_unfinished_req(self, req: Req, chunked=False): """Cache request when it is unfinished.""" self._reject_cache_salt(req.cache_salt) - assert req.req_pool_idx is not None + assert req.kv.req_pool_idx is not None token_ids = req.get_fill_ids() prefill_len = len(token_ids) # prefill only (maybe chunked) kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :prefill_len + req.kv.req_pool_idx, :prefill_len ].to(dtype=torch.int64, copy=True) # NOTE: our C++ implementation don't need `token_ids` and `kv_indices` to be page-aligned @@ -254,7 +254,7 @@ class RadixCacheCpp(BasePrefixCache): ) reused_indices = new_indices[old_prefix_len:new_prefix_len] self.req_to_token_pool.req_to_token[ - req.req_pool_idx, old_prefix_len:new_prefix_len + req.kv.req_pool_idx, old_prefix_len:new_prefix_len ] = reused_indices if req.last_node != new_last_node: diff --git a/python/sglang/srt/mem_cache/sparsity/core/sparse_coordinator.py b/python/sglang/srt/mem_cache/sparsity/core/sparse_coordinator.py index 5d2e4849b..cb5d91755 100644 --- a/python/sglang/srt/mem_cache/sparsity/core/sparse_coordinator.py +++ b/python/sglang/srt/mem_cache/sparsity/core/sparse_coordinator.py @@ -142,18 +142,18 @@ class SparseCoordinator: Registers the request in the state tracker to enable sparse attention processing. """ - if req.req_pool_idx is not None: - self.states.register(req.req_pool_idx, len(req.origin_input_ids)) + if req.kv.req_pool_idx is not None: + self.states.register(req.kv.req_pool_idx, len(req.origin_input_ids)) def on_request_end(self, req: "Req") -> None: """ Handle request end event. Called when a request is completed or aborted. Cleans up request-specific state and releases resources. """ - if req.req_pool_idx is None: + if req.kv.req_pool_idx is None: return - self.states.clear(req.req_pool_idx) + self.states.clear(req.kv.req_pool_idx) # TODO: Implement request end handling # - Release host indices if any were allocated for offloading diff --git a/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py index e1c582eb3..7eb53ddf7 100644 --- a/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py +++ b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py @@ -409,7 +409,7 @@ class FlexKVRadixCache(RadixCache): if not token_ids: return kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :kv_committed_len + req.kv.req_pool_idx, :kv_committed_len ] # Anchor on the new last_device_node so FlexKV's lock matches diff --git a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py index 54f70e618..e07754690 100644 --- a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py +++ b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py @@ -463,7 +463,7 @@ class LMCRadixCache(RadixCache): token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len] kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :kv_committed_len + req.kv.req_pool_idx, :kv_committed_len ] # Use super() to avoid a redundant LOOKUP — we only need new_last_node from radix. diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 361c2517a..86374c907 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -465,14 +465,14 @@ class SWARadixCache(BasePrefixCache): """Cache request when it finishes.""" if self.disable: kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :kv_len_to_handle + req.kv.req_pool_idx, :kv_len_to_handle ] self.token_to_kv_pool_allocator.free(kv_indices) return token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle] kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :kv_len_to_handle + req.kv.req_pool_idx, :kv_len_to_handle ] radix_key = RadixKey( @@ -516,7 +516,7 @@ class SWARadixCache(BasePrefixCache): """Cache request when it is unfinished.""" if self.disable: kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, : req.extend_range.end + req.kv.req_pool_idx, : req.extend_range.end ] # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later @@ -525,7 +525,7 @@ class SWARadixCache(BasePrefixCache): token_ids = req.get_fill_ids() kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, : len(token_ids) + req.kv.req_pool_idx, : len(token_ids) ] radix_key = RadixKey( @@ -561,7 +561,7 @@ class SWARadixCache(BasePrefixCache): assert old_prefix_len <= len(new_indices), f"{old_prefix_len=}, {new_indices=}" assert new_prefix_len <= len(new_indices), f"{new_prefix_len=}, {new_indices=}" self.req_to_token_pool.write( - (req.req_pool_idx, slice(old_prefix_len, len(new_indices))), + (req.kv.req_pool_idx, slice(old_prefix_len, len(new_indices))), new_indices[old_prefix_len:], ) diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index f4371bb41..ff80ca33e 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -825,7 +825,7 @@ class UnifiedRadixCache(BasePrefixCache): if self.disable: kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :kv_len_to_handle + req.kv.req_pool_idx, :kv_len_to_handle ] self.token_to_kv_pool_allocator.free_segment(kv_indices, start_pos=0) for comp in self._components_tuple: @@ -834,7 +834,7 @@ class UnifiedRadixCache(BasePrefixCache): token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle] kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :kv_len_to_handle + req.kv.req_pool_idx, :kv_len_to_handle ] result = None @@ -918,13 +918,13 @@ class UnifiedRadixCache(BasePrefixCache): if self.disable: kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, : len(token_ids) + req.kv.req_pool_idx, : len(token_ids) ] req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True) return kv_indices_orig = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, : len(token_ids) + req.kv.req_pool_idx, : len(token_ids) ] # components prepare insert data + return effective cache_len @@ -991,7 +991,7 @@ class UnifiedRadixCache(BasePrefixCache): new_indices ), f"{new_prefix_len=}, {len(new_indices)=}" self.req_to_token_pool.write( - (req.req_pool_idx, slice(req.kv.cache_protected_len, len(new_indices))), + (req.kv.req_pool_idx, slice(req.kv.cache_protected_len, len(new_indices))), new_indices[req.kv.cache_protected_len :], ) @@ -1160,7 +1160,7 @@ class UnifiedRadixCache(BasePrefixCache): ) -> tuple[torch.Tensor, list[PoolTransfer]]: num_tokens = req.seqlen - 1 full_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :num_tokens + req.kv.req_pool_idx, :num_tokens ].to(torch.int64) full_indices = self._pad_retraction_indices(full_indices, self.page_size) @@ -1171,7 +1171,7 @@ class UnifiedRadixCache(BasePrefixCache): window_start = max(0, num_tokens - self.sliding_window_size) window_start = window_start // self.page_size * self.page_size window_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, window_start:num_tokens + req.kv.req_pool_idx, window_start:num_tokens ].to(torch.int64) swa_indices = kv_cache.translate_loc_from_full_to_swa(window_indices) assert bool( diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index 631256934..f46a62d36 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/python/sglang/srt/session/streaming_session.py @@ -45,7 +45,6 @@ class SessionSlot: virtual_node: _VirtualNode = field(default_factory=_VirtualNode) # KV pool state - req_pool_idx: Optional[int] = None kv: ReqKvInfo = field(default_factory=ReqKvInfo) # First req's radix tree node (for dec_lock_ref on session close) @@ -63,15 +62,8 @@ class SessionSlot: mamba_last_track_seqlen: Any = None mamba_branching_seqlen: Any = None - @property - def is_holding_kv(self) -> bool: - """Whether this slot currently holds KV pool resources.""" - return self.req_pool_idx is not None - def save_from_req(self, req: Req, is_first: bool): """Save KV state from a finishing request into this slot.""" - self.req_pool_idx = req.req_pool_idx - if is_first: self.last_node = req.last_node self.swa_uuid_for_lock = req.swa_uuid_for_lock @@ -93,7 +85,6 @@ class SessionSlot: # Ownership moved to the slot; clear the req's references so a later # alloc/retract path cannot mistake slot-owned mamba state for its own. - req.req_pool_idx = None req.kv = ReqKvInfo() req.mamba_pool_idx = None req.mamba_ping_pong_track_buffer = None @@ -104,7 +95,6 @@ class SessionSlot: def restore_to_req(self, req: Req): """Restore KV state from this slot into an incoming request.""" - req.req_pool_idx = self.req_pool_idx req.kv = copy.copy(self.kv) req.swa_uuid_for_lock = self.swa_uuid_for_lock req.skip_lock_node_ids = self.skip_lock_node_ids @@ -189,7 +179,7 @@ class StreamingSession(BasePrefixCache): return session_id in self.slots def any_holding_kv(self) -> bool: - return any(s.is_holding_kv for s in self.slots.values()) + return any(s.kv.is_held for s in self.slots.values()) # -- Try-handle entries for composition (see class docstring) -- @@ -217,7 +207,7 @@ class StreamingSession(BasePrefixCache): if not _is_streaming(req): return None slot = self.slots.get(req.session.session_id) - if slot is None or not slot.is_holding_kv: + if slot is None or not slot.kv.is_held: return None if req.to_finish is not None: req.session.abort_req() @@ -290,7 +280,7 @@ class StreamingSession(BasePrefixCache): self._free_tail(slot, req, prefix_len) device_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :prefix_len + req.kv.req_pool_idx, :prefix_len ].to(dtype=torch.int64) return MatchResult( @@ -330,7 +320,6 @@ class StreamingSession(BasePrefixCache): # return the (possibly extra_buffer ping-pong) slots to # the mamba pool; otherwise the abort orphans them. slot = SessionSlot( - req_pool_idx=req.req_pool_idx, kv=copy.copy(req.kv), last_node=req.last_node, swa_uuid_for_lock=req.swa_uuid_for_lock, @@ -347,7 +336,6 @@ class StreamingSession(BasePrefixCache): slot.kv.kv_allocated_len, req.kv.kv_allocated_len ) self.release_session(session_id) - req.req_pool_idx = None req.kv = ReqKvInfo() req.session.abort_req() return True @@ -386,7 +374,7 @@ class StreamingSession(BasePrefixCache): return False if chunked: kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, : req.extend_range.end + req.kv.req_pool_idx, : req.extend_range.end ] req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True) return True @@ -441,9 +429,7 @@ class StreamingSession(BasePrefixCache): protected_len = slot.kv.cache_protected_len lock_node = slot.last_node tokens_freed = ( - max(0, slot.kv.kv_allocated_len - protected_len) - if slot.is_holding_kv - else 0 + max(0, slot.kv.kv_allocated_len - protected_len) if slot.kv.is_held else 0 ) logger.info( "Session KV released: %s (%d tokens freed)", session_id, tokens_freed @@ -458,12 +444,12 @@ class StreamingSession(BasePrefixCache): ), ) - if slot.is_holding_kv: + if slot.kv.is_held: start = protected_len end = slot.kv.kv_allocated_len if start < end: kv_indices = self.req_to_token_pool.req_to_token[ - slot.req_pool_idx, start:end + slot.kv.req_pool_idx, start:end ] self.token_to_kv_pool_allocator.free(kv_indices) self.req_to_token_pool.free(slot) @@ -483,9 +469,10 @@ class StreamingSession(BasePrefixCache): total = 0 for slot in self.slots.values(): in_batch = ( - active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs + active_pool_idxs is not None + and slot.kv.req_pool_idx in active_pool_idxs ) - if slot.is_holding_kv and not in_batch: + if slot.kv.is_held and not in_batch: allocated = ceil_align(slot.kv.kv_allocated_len, self.page_size) total += allocated - slot.kv.cache_protected_len return total @@ -499,9 +486,10 @@ class StreamingSession(BasePrefixCache): total = 0 for slot in self.slots.values(): in_batch = ( - active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs + active_pool_idxs is not None + and slot.kv.req_pool_idx in active_pool_idxs ) - if slot.is_holding_kv and not in_batch: + if slot.kv.is_held and not in_batch: allocated = ceil_align(slot.kv.kv_allocated_len, self.page_size) total += allocated - max( slot.kv.cache_protected_len, slot.kv.swa_evicted_seqlen @@ -513,9 +501,9 @@ class StreamingSession(BasePrefixCache): def _owned(s): in_batch = ( - active_pool_idxs is not None and s.req_pool_idx in active_pool_idxs + active_pool_idxs is not None and s.kv.req_pool_idx in active_pool_idxs ) - return s.is_holding_kv and not in_batch + return s.kv.is_held and not in_batch return sum(_owned(s) for s in self.slots.values()) @@ -529,7 +517,8 @@ class StreamingSession(BasePrefixCache): total = 0 for slot in self.slots.values(): in_batch = ( - active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs + active_pool_idxs is not None + and slot.kv.req_pool_idx in active_pool_idxs ) if in_batch: continue @@ -559,7 +548,9 @@ class StreamingSession(BasePrefixCache): decoding pushes allocated above committed, or when retract retry's logit-reserve pulls prefix_len below committed. """ - self._free_kv_aligned(slot.req_pool_idx, prefix_len, slot.kv.kv_allocated_len) + self._free_kv_aligned( + slot.kv.req_pool_idx, prefix_len, slot.kv.kv_allocated_len + ) slot.kv.kv_allocated_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) @@ -574,7 +565,7 @@ class StreamingSession(BasePrefixCache): be released to avoid token/KV mismatch. """ 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.kv.req_pool_idx, target, req.kv.kv_allocated_len) req.kv.kv_allocated_len = min(req.kv.kv_allocated_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) diff --git a/python/sglang/test/scripted_runtime/req_handle.py b/python/sglang/test/scripted_runtime/req_handle.py index d3459d8e9..a758fd603 100644 --- a/python/sglang/test/scripted_runtime/req_handle.py +++ b/python/sglang/test/scripted_runtime/req_handle.py @@ -42,7 +42,7 @@ class ScriptedReqHandle: @property def kv_pages(self) -> int: req = self.req - if req is None or not req.is_holding_kv: + if req is None or not req.kv.is_held: return 0 page_size = self.context.scheduler.page_size return (req.kv.kv_allocated_len + page_size - 1) // page_size diff --git a/test/manual/chunked_prefill/test_scripted_abort.py b/test/manual/chunked_prefill/test_scripted_abort.py index fe399a222..2214133c8 100644 --- a/test/manual/chunked_prefill/test_scripted_abort.py +++ b/test/manual/chunked_prefill/test_scripted_abort.py @@ -23,7 +23,7 @@ def _drain_until_released(t: ScriptedContext, *handles: ScriptedReqHandle): if all( h.kv_pages == 0 and h.lock_refs == 0 - and (h.req is None or h.req.req_pool_idx is None) + and (h.req is None or h.req.kv.req_pool_idx is None) for h in handles ): return @@ -57,7 +57,7 @@ class TestAbortBasic(ScriptedTestCase): r.kv_pages == 0 ), f"abort must release KV; r.kv_pages={r.kv_pages} after abort" assert ( - r.req is None or r.req.req_pool_idx is None + r.req is None or r.req.kv.req_pool_idx is None ), f"abort must release row; r.req={r.req} after abort" assert ( r.lock_refs == 0 @@ -78,7 +78,7 @@ class TestAbortBasic(ScriptedTestCase): yield from _drain_until_released(t, r) assert r.kv_pages == 0 - assert r.req is None or r.req.req_pool_idx is None + assert r.req is None or r.req.kv.req_pool_idx is None def test_abort_at_chunk_mid(self): self.server.execute_script(self._script_abort_at_chunk_mid) @@ -124,7 +124,7 @@ class TestAbortBasic(ScriptedTestCase): t.abort(r) yield from _drain_until_released(t, r) assert r.kv_pages == 0 - assert r.req is None or r.req.req_pool_idx is None + assert r.req is None or r.req.kv.req_pool_idx is None assert r.lock_refs == 0 def test_abort_at_admission_step(self): @@ -137,7 +137,7 @@ class TestAbortBasic(ScriptedTestCase): t.abort(r) yield from _drain_until_released(t, r) assert r.kv_pages == 0 - assert r.req is None or r.req.req_pool_idx is None + assert r.req is None or r.req.kv.req_pool_idx is None def test_abort_then_start_same_step_new_rid(self): self.server.execute_script(self._script_abort_then_start_same_step_new_rid) @@ -191,7 +191,7 @@ class TestAbortBasic(ScriptedTestCase): yield from _drain_until_released(t, *reqs) for r in reqs: assert r.kv_pages == 0 - assert r.req is None or r.req.req_pool_idx is None + assert r.req is None or r.req.kv.req_pool_idx is None def test_abort_unknown_rid_noop(self): self.server.execute_script(self._script_abort_unknown_rid_noop) @@ -264,7 +264,7 @@ class TestAbortBasic(ScriptedTestCase): t.abort(r) yield from _drain_until_released(t, r) assert r.kv_pages == 0 - assert r.req is None or r.req.req_pool_idx is None + assert r.req is None or r.req.kv.req_pool_idx is None assert r.lock_refs == 0 def test_double_abort_idempotent(self): @@ -351,7 +351,7 @@ class TestAbortBasic(ScriptedTestCase): f"aborted req revived and ran another chunk; " f"chunks_done went {chunks_after_abort} -> {r.chunks_done}" ) - assert r.req is None or r.req.req_pool_idx is None + assert r.req is None or r.req.kv.req_pool_idx is None def test_abort_mid_chunk_no_extra_radix_node(self): self.server.execute_script(self._script_abort_mid_chunk_no_extra_radix_node) @@ -369,7 +369,7 @@ class TestAbortBasic(ScriptedTestCase): yield from _drain_until_released(t, r) assert r.kv_pages == 0 - assert r.req is None or r.req.req_pool_idx is None + assert r.req is None or r.req.kv.req_pool_idx is None chunks_after_release = r.chunks_done for _ in range(4): yield @@ -406,7 +406,7 @@ class TestAbortBasic(ScriptedTestCase): yield from run_until_finished(r2) assert r2.finished, "resubmit under same rid must complete independently" assert r1.kv_pages == 0, "aborted r1 must release KV before resubmit" - assert r1.req is None or r1.req.req_pool_idx is None + assert r1.req is None or r1.req.kv.req_pool_idx is None assert r1.lock_refs == 0 def test_abort_during_gap_inflight_middle_chunks_positive(self): @@ -429,7 +429,7 @@ class TestAbortBasic(ScriptedTestCase): yield from _drain_until_released(t, r) assert r.kv_pages == 0 - assert r.req is None or r.req.req_pool_idx is None + assert r.req is None or r.req.kv.req_pool_idx is None assert not r.is_chunking, "aborted gap req must not re-enter chunking" yield @@ -511,7 +511,7 @@ class TestAbortBasic(ScriptedTestCase): assert r1.kv_pages == 0, ( f"force_retract + abort same yield must release KV; got " f"{r1.kv_pages}" ) - assert r1.req is None or r1.req.req_pool_idx is None, ( + assert r1.req is None or r1.req.kv.req_pool_idx is None, ( f"force_retract + abort same yield must release row; got " f"{r1.req}" ) assert r1.lock_refs == 0, ( @@ -540,7 +540,7 @@ class TestAbortBasic(ScriptedTestCase): yield from run_until(r2, lambda h: h.is_chunking) assert r1.kv_pages == 0 - assert r1.req is None or r1.req.req_pool_idx is None + assert r1.req is None or r1.req.kv.req_pool_idx is None assert r1.lock_refs == 0 yield from run_until_finished(r2) assert r2.finished, "baton handoff must let r2 complete" @@ -570,7 +570,7 @@ class TestAbortPP(ScriptedTestCase): yield from _drain_until_released(t, r) assert r.kv_pages == 0 - assert r.req is None or r.req.req_pool_idx is None + assert r.req is None or r.req.kv.req_pool_idx is None assert r.lock_refs == 0 assert r.finished diff --git a/test/manual/chunked_prefill/test_scripted_kv_pressure.py b/test/manual/chunked_prefill/test_scripted_kv_pressure.py index 751f1098a..81be5f5f8 100644 --- a/test/manual/chunked_prefill/test_scripted_kv_pressure.py +++ b/test/manual/chunked_prefill/test_scripted_kv_pressure.py @@ -409,13 +409,13 @@ class TestKVPressureSmallPool(ScriptedTestCase): if ( r_chunk.kv_pages == 0 and r_chunk.lock_refs == 0 - and (r_chunk.req is None or r_chunk.req.req_pool_idx is None) + and (r_chunk.req is None or r_chunk.req.kv.req_pool_idx is None) ): break yield assert r_chunk.kv_pages == 0, f"kv_pages={r_chunk.kv_pages}" assert r_chunk.lock_refs == 0, f"lock_refs={r_chunk.lock_refs}" - assert r_chunk.req is None or r_chunk.req.req_pool_idx is None + assert r_chunk.req is None or r_chunk.req.kv.req_pool_idx is None t.abort(ballast) for _ in range(200): diff --git a/test/manual/chunked_prefill/test_scripted_lifecycle.py b/test/manual/chunked_prefill/test_scripted_lifecycle.py index bffc2784c..8cc136bca 100644 --- a/test/manual/chunked_prefill/test_scripted_lifecycle.py +++ b/test/manual/chunked_prefill/test_scripted_lifecycle.py @@ -18,7 +18,7 @@ def _drain_until_released(t, *handles): if all( h.kv_pages == 0 and h.lock_refs == 0 - and (h.req is None or h.req.req_pool_idx is None) + and (h.req is None or h.req.kv.req_pool_idx is None) for h in handles ): return @@ -223,7 +223,7 @@ class TestLifecycleBasic(ScriptedTestCase): r = t.start_req(prompt_len=16, max_new_tokens=2, ignore_eos=True) yield from run_until_finished(r) assert r.finished - assert r.req.req_pool_idx is None + assert r.req.kv.req_pool_idx is None assert r.kv_pages == 0 assert r.lock_refs == 0 @@ -235,7 +235,7 @@ class TestLifecycleBasic(ScriptedTestCase): r1 = t.start_req(prompt_len=16, max_new_tokens=2, ignore_eos=True) yield from run_until_finished(r1) yield from _drain_until_released(t, r1) - assert r1.req.req_pool_idx is None and r1.kv_pages == 0 and r1.lock_refs == 0 + assert r1.req.kv.req_pool_idx is None and r1.kv_pages == 0 and r1.lock_refs == 0 r1_output_len = len(r1.req.output_ids) r2 = t.start_req(prompt_len=16, max_new_tokens=2, ignore_eos=True) @@ -243,7 +243,7 @@ class TestLifecycleBasic(ScriptedTestCase): yield from _drain_until_released(t, r2) assert r1.finished and r2.finished assert r1_output_len == 2 and len(r2.req.output_ids) == 2 - assert r2.req.req_pool_idx is None and r2.kv_pages == 0 and r2.lock_refs == 0 + assert r2.req.kv.req_pool_idx is None and r2.kv_pages == 0 and r2.lock_refs == 0 def test_five_seq_clean(self): self.server.execute_script(self._script_five_seq_clean) @@ -256,7 +256,7 @@ class TestLifecycleBasic(ScriptedTestCase): yield from run_until_finished(r) assert r.finished assert len(r.req.output_ids) == 2 - assert r.req.req_pool_idx is None + assert r.req.kv.req_pool_idx is None assert r.kv_pages == 0 assert r.lock_refs == 0 reqs.append(r) @@ -300,7 +300,9 @@ class TestLifecycleBasic(ScriptedTestCase): yield from run_until_finished(r) assert r.finished assert len(r.req.output_ids) == 2 - assert r.req.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0 + assert ( + r.req.kv.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0 + ) if prompt == VERY_LONG_PROMPT_LEN: assert r.chunks_done == 8 else: @@ -319,7 +321,7 @@ class TestLifecycleBasic(ScriptedTestCase): assert r.finished assert len(r.req.output_ids) == 1 yield from _drain_until_released(t, r) - assert r.req is None or r.req.req_pool_idx is None + assert r.req is None or r.req.kv.req_pool_idx is None assert r.kv_pages == 0 and r.lock_refs == 0 if L > DEFAULT_CHUNK_SIZE: assert ( @@ -341,7 +343,7 @@ class TestLifecycleBasic(ScriptedTestCase): assert r.finished assert len(r.req.output_ids) == 1 yield from _drain_until_released(t, r) - assert r.req is None or r.req.req_pool_idx is None + assert r.req is None or r.req.kv.req_pool_idx is None assert r.kv_pages == 0 and r.lock_refs == 0 if L > DEFAULT_CHUNK_SIZE: assert ( @@ -360,7 +362,9 @@ class TestLifecycleBasic(ScriptedTestCase): yield from run_until_finished(r) assert r.finished assert len(r.req.output_ids) == 2 - assert r.req.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0 + assert ( + r.req.kv.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0 + ) for _ in range(20): yield @@ -377,7 +381,9 @@ class TestLifecycleBasic(ScriptedTestCase): yield from run_until_finished(r) assert r.finished assert len(r.req.output_ids) == 2 - assert r.req.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0 + assert ( + r.req.kv.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0 + ) if L == VERY_LONG_PROMPT_LEN: assert r.chunks_done == 8 else: @@ -394,7 +400,9 @@ class TestLifecycleBasic(ScriptedTestCase): yield from run_until_finished(r) assert r.finished assert len(r.req.output_ids) == 2 - assert r.req.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0 + assert ( + r.req.kv.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0 + ) for _ in range(5): yield t.flush_cache() @@ -432,7 +440,7 @@ class TestLifecycleBasic(ScriptedTestCase): ) assert r.finished or _error_message(r) is not None assert r.kv_pages == 0 - assert r.req is None or r.req.req_pool_idx is None + assert r.req is None or r.req.kv.req_pool_idx is None assert r.lock_refs == 0 diff --git a/test/manual/chunked_prefill/test_scripted_lora.py b/test/manual/chunked_prefill/test_scripted_lora.py index 612683937..13e6087fe 100644 --- a/test/manual/chunked_prefill/test_scripted_lora.py +++ b/test/manual/chunked_prefill/test_scripted_lora.py @@ -222,7 +222,9 @@ class TestLoRAAdapterEviction(ScriptedTestCase): t.abort(r_a) for _ in range(12): - if r_a.kv_pages == 0 and (r_a.req is None or r_a.req.req_pool_idx is None): + if r_a.kv_pages == 0 and ( + r_a.req is None or r_a.req.kv.req_pool_idx is None + ): break yield diff --git a/test/manual/chunked_prefill/test_scripted_multi_req.py b/test/manual/chunked_prefill/test_scripted_multi_req.py index 27db67252..d9ba9f207 100644 --- a/test/manual/chunked_prefill/test_scripted_multi_req.py +++ b/test/manual/chunked_prefill/test_scripted_multi_req.py @@ -173,7 +173,7 @@ class TestMultiReqBasic(ScriptedTestCase): for _ in range(5): yield assert r.kv_pages == 0 - assert r.req is None or r.req.req_pool_idx is None + assert r.req is None or r.req.kv.req_pool_idx is None def test_rid_reuse_after_finish(self): self.server.execute_script(self._script_rid_reuse_after_finish) diff --git a/test/manual/chunked_prefill/test_scripted_pp.py b/test/manual/chunked_prefill/test_scripted_pp.py index bb5cc9e23..736bb06ce 100644 --- a/test/manual/chunked_prefill/test_scripted_pp.py +++ b/test/manual/chunked_prefill/test_scripted_pp.py @@ -31,7 +31,7 @@ def _expected_chunks(prompt_len: int, chunk_size: int) -> int: def _drain_until_released(t, *handles): for _ in range(16): if all( - h.kv_pages == 0 and (h.req is None or h.req.req_pool_idx is None) + h.kv_pages == 0 and (h.req is None or h.req.kv.req_pool_idx is None) for h in handles ): return diff --git a/test/manual/chunked_prefill/test_scripted_priority.py b/test/manual/chunked_prefill/test_scripted_priority.py index c164342b6..090f4aec4 100644 --- a/test/manual/chunked_prefill/test_scripted_priority.py +++ b/test/manual/chunked_prefill/test_scripted_priority.py @@ -168,13 +168,13 @@ class TestPriorityBasic(ScriptedTestCase): if ( r.kv_pages == 0 and r.lock_refs == 0 - and (r.req is None or r.req.req_pool_idx is None) + and (r.req is None or r.req.kv.req_pool_idx is None) ): break yield assert r.kv_pages == 0 assert r.lock_refs == 0 - assert r.req is None or r.req.req_pool_idx is None + assert r.req is None or r.req.kv.req_pool_idx is None t.continue_generation() yield assert r.kv_pages == 0 and r.lock_refs == 0 diff --git a/test/manual/chunked_prefill/test_scripted_regression.py b/test/manual/chunked_prefill/test_scripted_regression.py index 2c9e11780..29f0c1731 100644 --- a/test/manual/chunked_prefill/test_scripted_regression.py +++ b/test/manual/chunked_prefill/test_scripted_regression.py @@ -19,7 +19,7 @@ def _drain_until_released(t, *handles): if all( h.kv_pages == 0 and h.lock_refs == 0 - and (h.req is None or h.req.req_pool_idx is None) + and (h.req is None or h.req.kv.req_pool_idx is None) for h in handles ): return @@ -41,7 +41,7 @@ class TestRegressionBasic(ScriptedTestCase): yield from _drain_until_released(t, r) assert r.kv_pages == 0 - assert r.req.req_pool_idx is None + assert r.req.kv.req_pool_idx is None assert r.lock_refs == 0 assert not r.is_chunking assert r.req.inflight_middle_chunks == 0 @@ -58,7 +58,7 @@ class TestRegressionBasic(ScriptedTestCase): yield assert r.kv_pages == 0 - assert r.req.req_pool_idx is None + assert r.req.kv.req_pool_idx is None assert r.lock_refs == 0 assert not r.is_chunking @@ -240,7 +240,7 @@ class TestRegressionBasic(ScriptedTestCase): 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) - assert r.req.req_pool_idx is not None, "row must be held mid-chunk" + assert r.req.kv.req_pool_idx is not None, "row must be held mid-chunk" assert r.kv_pages > 0, "committed KV must be held mid-chunk" assert r.lock_refs >= 1, "radix lock_ref must be held mid-chunk" @@ -248,8 +248,8 @@ class TestRegressionBasic(ScriptedTestCase): yield from _drain_until_released(t, r) assert ( - r.req.req_pool_idx is None - ), f"96d4749094: abort must release row; got row_idx={r.req.req_pool_idx!r}" + r.req.kv.req_pool_idx is None + ), f"96d4749094: abort must release row; got row_idx={r.req.kv.req_pool_idx!r}" assert ( r.kv_pages == 0 ), f"96d4749094: abort must release KV; got kv_pages={r.kv_pages}" @@ -269,14 +269,14 @@ class TestRegressionBasic(ScriptedTestCase): def _script_pause_retract_releases_waiting_chunked_resume(t: ScriptedContext): 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) - assert r.req.req_pool_idx is not None and r.kv_pages > 0 and r.lock_refs >= 1 + assert r.req.kv.req_pool_idx is not None and r.kv_pages > 0 and r.lock_refs >= 1 t.pause_generation(mode="retract") yield - assert r.req.req_pool_idx is None, ( + assert r.req.kv.req_pool_idx is None, ( f"f38e69f87d: pause(retract) must release waiting " - f"chunked-resume row; got row_idx={r.req.req_pool_idx!r}" + f"chunked-resume row; got row_idx={r.req.kv.req_pool_idx!r}" ) assert r.kv_pages == 0 assert r.lock_refs == 0 diff --git a/test/manual/scheduler/test_scripted_pp_abort.py b/test/manual/scheduler/test_scripted_pp_abort.py index 06ab4505a..02a0de2c4 100644 --- a/test/manual/scheduler/test_scripted_pp_abort.py +++ b/test/manual/scheduler/test_scripted_pp_abort.py @@ -14,7 +14,7 @@ def _drain_until_released(t: ScriptedContext, *handles: ScriptedReqHandle): if all( h.kv_pages == 0 and h.lock_refs == 0 - and (h.req is None or h.req.req_pool_idx is None) + and (h.req is None or h.req.kv.req_pool_idx is None) for h in handles ): return diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index 3fbaea191..1419a5827 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -102,7 +102,7 @@ class TestDisaggregationWire(unittest.TestCase): def test_prebuilt_skips_unused_prompt_tensor(self): req = SimpleNamespace( - req_pool_idx=0, + kv=SimpleNamespace(req_pool_idx=0), prefix_indices=[0, 1], extend_range=SimpleNamespace(length=3), origin_input_ids=[0, 1, 2, 3, 4], diff --git a/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py b/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py index f2595158b..5700987a6 100644 --- a/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py +++ b/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py @@ -33,9 +33,10 @@ def _make_mock_req( """Create a mock Req with the KV cache state needed for testing.""" req = MagicMock() req.rid = rid - req.req_pool_idx = req_pool_idx req.kv = ReqKvInfo( - kv_committed_len=kv_committed_len, kv_allocated_len=kv_allocated_len + req_pool_idx=req_pool_idx, + kv_committed_len=kv_committed_len, + kv_allocated_len=kv_allocated_len, ) req.prefix_indices = list(range(prefix_indices_len)) req.effective_kv_committed_len = lambda: req.kv.kv_committed_len diff --git a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py index 11c816d1f..b9a1f9b93 100644 --- a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py +++ b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py @@ -719,7 +719,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase): full_token_ids=[11, 12, 13], prefix_slot_ids=[2, 3], new_slot_ids=[4], - req_pool_idx=req.req_pool_idx, + req_pool_idx=req.kv.req_pool_idx, ) runner.eval_pending(pending) runner.prefill_finalize(pending) @@ -775,7 +775,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase): full_token_ids=token_ids, prefix_slot_ids=[], new_slot_ids=list(range(1, 71)), - req_pool_idx=req.req_pool_idx, + req_pool_idx=req.kv.req_pool_idx, req=req, ) runner.eval_pending(pending) @@ -837,7 +837,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase): full_token_ids=token_ids, prefix_slot_ids=list(range(1, 65)), new_slot_ids=list(range(65, 258)), - req_pool_idx=req.req_pool_idx, + req_pool_idx=req.kv.req_pool_idx, req=req, ) runner.eval_pending(pending) @@ -917,7 +917,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase): req = FakeRequest() req_indices = pool.alloc([req]) - auxiliary_state_idx = pool.get_auxiliary_state_indices(req.req_pool_idx) + auxiliary_state_idx = pool.get_auxiliary_state_indices(req.kv.req_pool_idx) pool.free(req) # Which free slot a fresh alloc gets is not semantically meaningful @@ -925,7 +925,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase): self.assertEqual(len(req_indices), 1) self.assertIn(req_indices[0], range(1, pool.size + 1)) self.assertIsNotNone(auxiliary_state_idx) - self.assertIsNone(req.req_pool_idx) + self.assertIsNone(req.kv.req_pool_idx) self.assertIsNotNone(req.mamba_pool_idx) self.assertIs(pool.mamba_allocator, pool.mamba_pool) self.assertEqual(pool.auxiliary_state_pool.available_size(), 3) @@ -1518,7 +1518,7 @@ if _HAS_MLX: class FakeRequest: def __init__(self): - self.req_pool_idx = None + self.kv = SimpleNamespace(req_pool_idx=None) self.mamba_pool_idx = None self.inflight_middle_chunks = 0 diff --git a/test/registered/unit/hardware_backend/mlx/test_max_running_requests.py b/test/registered/unit/hardware_backend/mlx/test_max_running_requests.py index a75eccdea..48f05dfa3 100644 --- a/test/registered/unit/hardware_backend/mlx/test_max_running_requests.py +++ b/test/registered/unit/hardware_backend/mlx/test_max_running_requests.py @@ -119,8 +119,8 @@ def _hybrid_stub_for_initialize( def _fake_req(): return SimpleNamespace( - req_pool_idx=None, inflight_middle_chunks=0, + kv=SimpleNamespace(req_pool_idx=None), mamba_pool_idx=None, mamba_ping_pong_track_buffer=None, ) diff --git a/test/registered/unit/hardware_backend/mlx/test_tp_worker_routing.py b/test/registered/unit/hardware_backend/mlx/test_tp_worker_routing.py index c361618f1..b4274a7c4 100644 --- a/test/registered/unit/hardware_backend/mlx/test_tp_worker_routing.py +++ b/test/registered/unit/hardware_backend/mlx/test_tp_worker_routing.py @@ -169,7 +169,7 @@ class _FakeReq: self.rid = rid self.prefix_indices = torch.empty(0, dtype=torch.long) self.fill_ids = [0] - self.req_pool_idx = req_pool_idx + self.kv = SimpleNamespace(req_pool_idx=req_pool_idx) # Mirrors Req's chunk-finality contract read by # MlxTpModelWorker._chunk_needs_logits: extend_range=None means # "not truncated" (final chunk / plain prefill). diff --git a/test/registered/unit/layers/test_minicpm_sparse_cache.py b/test/registered/unit/layers/test_minicpm_sparse_cache.py index 9d33d8f25..2690ad7a1 100644 --- a/test/registered/unit/layers/test_minicpm_sparse_cache.py +++ b/test/registered/unit/layers/test_minicpm_sparse_cache.py @@ -7,6 +7,7 @@ import torch from sglang.srt.layers.attention.minicpm.cache import ( attach_compressed_cache, ) +from sglang.srt.managers.schedule_batch import ReqKvInfo from sglang.srt.managers.scheduler_components.invariant_checker import ( SchedulerInvariantChecker, ) @@ -66,8 +67,8 @@ def make_pool_and_req(capacity: int = 64): enable_memory_saver=False, ) req = SimpleNamespace( - req_pool_idx=None, inflight_middle_chunks=0, + kv=SimpleNamespace(req_pool_idx=None), ) req_pool_idx = pool.alloc([req])[0] return pool, req, req_pool_idx, allocator @@ -224,8 +225,7 @@ def test_streaming_session_release_frees_compressed_slots(): ) ) session.slots["session-a"] = SessionSlot( - req_pool_idx=req_pool_idx, - kv=SimpleNamespace(kv_allocated_len=16, cache_protected_len=0), + kv=ReqKvInfo(req_pool_idx=req_pool_idx, kv_allocated_len=16), ) session.release_session("session-a") @@ -293,7 +293,7 @@ def test_partial_failure_rolls_back_and_free_releases_every_slot(): assert len(cache.free_slots) == 0 pool.free(req) - assert req.req_pool_idx is None + assert req.kv.req_pool_idx is None assert allocator.available_size() == 12 assert len(cache.free_slots) == 8 diff --git a/test/registered/unit/managers/test_hisparse_unit.py b/test/registered/unit/managers/test_hisparse_unit.py index be5586464..c2b5bdf99 100644 --- a/test/registered/unit/managers/test_hisparse_unit.py +++ b/test/registered/unit/managers/test_hisparse_unit.py @@ -49,8 +49,7 @@ def _make_req(rid="test-req-0", origin_input_ids=None, output_ids=None): output_ids=output_ids, fill_ids=origin_input_ids + output_ids, seqlen=len(origin_input_ids) + len(output_ids), - req_pool_idx=None, - kv=SimpleNamespace(kv_allocated_len=0, kv_committed_len=0), + kv=SimpleNamespace(req_pool_idx=None, kv_allocated_len=0, kv_committed_len=0), finished_reason=None, hisparse_staging=False, staging=False, @@ -190,11 +189,11 @@ class TestHiSparseUnit(unittest.TestCase): """Allocate a req_pool_idx for the request.""" indices = self.req_to_token_pool.alloc([req]) self.assertIsNotNone(indices, "Failed to allocate req pool slot") - return req.req_pool_idx + return req.kv.req_pool_idx def _free_req_slot(self, req): """Free the req_pool_idx.""" - if req.req_pool_idx is not None: + if req.kv.req_pool_idx is not None: self.req_to_token_pool.free(req) def _alloc_kv(self, req, fill_len, *, logical_only=False): @@ -216,7 +215,9 @@ class TestHiSparseUnit(unittest.TestCase): extend_num_tokens=fill_len, ) 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.kv.req_pool_idx, slice(0, len(kv_loc))), kv_loc + ) req.kv.kv_allocated_len = fill_len req.kv.kv_committed_len = fill_len req.full_untruncated_fill_ids = array("q", range(fill_len)) @@ -254,8 +255,8 @@ class TestHiSparseUnit(unittest.TestCase): host_indices = host_pool.alloc(fill_len) self.assertIsNotNone(host_indices, "Host alloc failed") host_indices = host_indices.to(device="cuda") - self.coordinator.req_to_host_pool[req.req_pool_idx, :fill_len] = host_indices - self.coordinator.req_to_host_pool_allocated_len[req.req_pool_idx] = fill_len + self.coordinator.req_to_host_pool[req.kv.req_pool_idx, :fill_len] = host_indices + self.coordinator.req_to_host_pool_allocated_len[req.kv.req_pool_idx] = fill_len for lid in range(LAYER_NUM): for i in range(fill_len): host_pool.kv_buffer[lid][host_indices[i]] = self._kv_pattern(lid, i) @@ -291,7 +292,7 @@ class TestHiSparseUnit(unittest.TestCase): def _make_batch_tensors(self, reqs, fill_lens): """Build (req_pool_indices [int64], seq_lens [int32]) on cuda.""" rpi = torch.tensor( - [r.req_pool_idx for r in reqs], dtype=torch.int64, device="cuda" + [r.kv.req_pool_idx for r in reqs], dtype=torch.int64, device="cuda" ) sls = torch.tensor(fill_lens, dtype=torch.int32, device="cuda") return rpi, sls @@ -566,7 +567,7 @@ class TestHiSparseUnit(unittest.TestCase): kv_loc = self._alloc_kv(req, fill_len) self.coordinator.alloc_device_buffer(req) - self.coordinator._skip_first_backup[req.req_pool_idx] = True + self.coordinator._skip_first_backup[req.kv.req_pool_idx] = True out_loc = self.allocator.alloc(1) self.assertIsNotNone(out_loc) @@ -576,7 +577,7 @@ class TestHiSparseUnit(unittest.TestCase): self.assertTrue(torch.all(stale_loc > 0), "Temporary mapping should exist") 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.kv.req_pool_idx, fill_len), out_loc) req.kv.kv_allocated_len = seq_len req.kv.kv_committed_len = seq_len @@ -584,10 +585,10 @@ class TestHiSparseUnit(unittest.TestCase): seq_lens=torch.tensor([seq_len], dtype=torch.int64, device=device), out_cache_loc=out_loc, req_pool_indices=torch.tensor( - [req.req_pool_idx], dtype=torch.int64, device=device + [req.kv.req_pool_idx], dtype=torch.int64, device=device ), seq_lens_cpu=torch.tensor([seq_len], dtype=torch.int64), - req_pool_indices_cpu=torch.tensor([req.req_pool_idx], dtype=torch.int64), + req_pool_indices_cpu=torch.tensor([req.kv.req_pool_idx], dtype=torch.int64), ) remapped_loc = self.allocator.full_to_hisparse_device_index_mapping[out_loc] @@ -626,7 +627,7 @@ class TestHiSparseUnit(unittest.TestCase): ready = self.coordinator.collect_ready_reqs() self.assertEqual(len(ready), 1) self.assertFalse(req.hisparse_staging) - self.assertTrue(self.coordinator._skip_first_backup[req.req_pool_idx]) + self.assertTrue(self.coordinator._skip_first_backup[req.kv.req_pool_idx]) tokens = self._build_topk_tokens(fill_len) batch = tokens.unsqueeze(0) @@ -662,11 +663,11 @@ class TestHiSparseUnit(unittest.TestCase): ready = self.coordinator.collect_ready_reqs() self.assertEqual(ready, [req]) - host_row = self.coordinator.req_to_host_pool[req.req_pool_idx, :rounded_len] + host_row = self.coordinator.req_to_host_pool[req.kv.req_pool_idx, :rounded_len] self.assertTrue(torch.all(host_row >= 0)) self.assertEqual(torch.unique(host_row).numel(), rounded_len) self.assertEqual( - int(self.coordinator.req_to_host_pool_allocated_len[req.req_pool_idx]), + int(self.coordinator.req_to_host_pool_allocated_len[req.kv.req_pool_idx]), rounded_len, ) @@ -674,7 +675,7 @@ class TestHiSparseUnit(unittest.TestCase): next_host_index = self.coordinator.mem_pool_host.alloc_paged_token_slots( self.coordinator.req_to_host_pool, self.coordinator.req_to_host_pool_allocated_len, - req.req_pool_idx, + req.kv.req_pool_idx, fill_len, 1, ) @@ -691,8 +692,8 @@ class TestHiSparseUnit(unittest.TestCase): expected_total = rounded_len + expected_new_pages * self.page_size allocated_host_indices = self.coordinator.mem_pool_host.allocated_host_indices( self.coordinator.req_to_host_pool, - req.req_pool_idx, - int(self.coordinator.req_to_host_pool_allocated_len[req.req_pool_idx]), + req.kv.req_pool_idx, + int(self.coordinator.req_to_host_pool_allocated_len[req.kv.req_pool_idx]), ) self.assertEqual(allocated_host_indices.numel(), expected_total) @@ -714,9 +715,9 @@ class TestHiSparseUnit(unittest.TestCase): self.coordinator.admit_request_direct(req) self.assertFalse(req.staging) - self.assertTrue(self.coordinator._skip_first_backup[req.req_pool_idx]) + self.assertTrue(self.coordinator._skip_first_backup[req.kv.req_pool_idx]) buf_tokens = self.coordinator.req_device_buffer_tokens[ - :, req.req_pool_idx, :DEVICE_BUFFER_SIZE + :, req.kv.req_pool_idx, :DEVICE_BUFFER_SIZE ] self.assertTrue(torch.all(buf_tokens == -1)) @@ -765,7 +766,7 @@ class TestHiSparseUnit(unittest.TestCase): self.assertTrue( torch.equal( host_indices, - self.coordinator.req_to_host_pool[req.req_pool_idx, :fill_len], + self.coordinator.req_to_host_pool[req.kv.req_pool_idx, :fill_len], ) ) self.assertEqual(req.kv.kv_allocated_len, fill_len) @@ -774,18 +775,18 @@ class TestHiSparseUnit(unittest.TestCase): rounded_len = (fill_len + self.page_size - 1) // self.page_size * self.page_size self.assertEqual( - int(self.coordinator.req_to_host_pool_allocated_len[req.req_pool_idx]), + int(self.coordinator.req_to_host_pool_allocated_len[req.kv.req_pool_idx]), rounded_len, ) allocated_host_indices = self.coordinator.mem_pool_host.allocated_host_indices( self.coordinator.req_to_host_pool, - req.req_pool_idx, - int(self.coordinator.req_to_host_pool_allocated_len[req.req_pool_idx]), + req.kv.req_pool_idx, + int(self.coordinator.req_to_host_pool_allocated_len[req.kv.req_pool_idx]), ) self.assertEqual(allocated_host_indices.numel(), rounded_len) kv_loc = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, : req.kv.kv_allocated_len + req.kv.req_pool_idx, : req.kv.kv_allocated_len ].clone() self._cleanup_req(req, kv_loc, logical_only=True) self._assert_sizes_restored(initial, "pd_decode_prealloc_hisparse") diff --git a/test/registered/unit/managers/test_kv_page_invariants.py b/test/registered/unit/managers/test_kv_page_invariants.py index b9629b6f4..83a670274 100644 --- a/test/registered/unit/managers/test_kv_page_invariants.py +++ b/test/registered/unit/managers/test_kv_page_invariants.py @@ -5,6 +5,7 @@ from types import SimpleNamespace import torch +from sglang.srt.managers.schedule_batch import ReqKvInfo from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -43,23 +44,15 @@ def _make_checker(page_size=_PAGE_SIZE, row_width=4096, num_reqs=8, free_pages=N return _FakeChecker(), rtt, tc, alloc -class _FakeReq: - def __init__(self, rid, rpi, committed, allocated): +class _FakeOwner: + # A req or a session slot; the checker reads only `kv` (and `rid` for reqs). + def __init__(self, rpi, committed, allocated, rid=None): self.rid = rid - self.req_pool_idx = rpi - self.kv = SimpleNamespace( - kv_committed_len=committed, kv_allocated_len=allocated, swa_evicted_seqlen=0 + self.kv = ReqKvInfo( + req_pool_idx=rpi, + kv_committed_len=committed, + kv_allocated_len=allocated, ) - self.is_holding_kv = True - - -class _FakeSlot: - def __init__(self, rpi, committed, allocated): - self.req_pool_idx = rpi - self.kv = SimpleNamespace( - kv_committed_len=committed, kv_allocated_len=allocated, swa_evicted_seqlen=0 - ) - self.is_holding_kv = True class TestKVPageInvariants(CustomTestCase): @@ -70,21 +63,23 @@ class TestKVPageInvariants(CustomTestCase): rtt[0, :256] = torch.arange(_PAGE_SIZE) # req 0 owns page 0 rtt[1, :256] = torch.arange(_PAGE_SIZE, 2 * _PAGE_SIZE) # req 1 owns page 1 chk.get_last_batch = lambda: SimpleNamespace( - reqs=[_FakeReq("a", 0, 256, 256), _FakeReq("b", 1, 200, 256)] + reqs=[_FakeOwner(0, 256, 256, rid="a"), _FakeOwner(1, 200, 256, rid="b")] ) chk._check_kv_page_invariants() self.assertEqual(chk.count_memory_leak_warnings, 0) def test_committed_gt_allocated_raises(self): chk, rtt, tc, alloc = _make_checker() - chk.get_last_batch = lambda: SimpleNamespace(reqs=[_FakeReq("a", 0, 145, 144)]) + chk.get_last_batch = lambda: SimpleNamespace( + reqs=[_FakeOwner(0, 145, 144, rid="a")] + ) with self.assertRaises(AssertionError): chk._check_kv_page_invariants() def test_slot_committed_gt_allocated_raises(self): chk, rtt, tc, alloc = _make_checker() chk.get_last_batch = lambda: None - tc.slots = {"s1": _FakeSlot(0, 145, 144)} + tc.slots = {"s1": _FakeOwner(0, 145, 144)} with self.assertRaises(AssertionError): chk._check_kv_page_invariants() @@ -94,14 +89,18 @@ class TestKVPageInvariants(CustomTestCase): rtt[0, :3] = torch.tensor( [5 * _PAGE_SIZE, 5 * _PAGE_SIZE + 1, 5 * _PAGE_SIZE + 2] ) - chk.get_last_batch = lambda: SimpleNamespace(reqs=[_FakeReq("a", 0, 3, 3)]) + chk.get_last_batch = lambda: SimpleNamespace( + reqs=[_FakeOwner(0, 3, 3, rid="a")] + ) with self.assertRaises(ValueError): chk._check_kv_page_invariants() def test_free_pool_duplicate_raises(self): chk, rtt, tc, alloc = _make_checker(free_pages=torch.tensor([3, 3, 4])) rtt[0, :1] = torch.tensor([10 * _PAGE_SIZE]) # owner page 10, not in free - chk.get_last_batch = lambda: SimpleNamespace(reqs=[_FakeReq("a", 0, 1, 1)]) + chk.get_last_batch = lambda: SimpleNamespace( + reqs=[_FakeOwner(0, 1, 1, rid="a")] + ) with self.assertRaises(ValueError): chk._check_kv_page_invariants() diff --git a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py index 2a208d037..41e332759 100644 --- a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py +++ b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py @@ -93,11 +93,10 @@ class TestDecodePreallocQueuePriority(unittest.TestCase): priority=priority, origin_input_ids=[1, 2, 3], output_ids=[], - req_pool_idx=int(priority) % 8, finished_reason=FINISH_ABORT("failed") if failed else None, return_logprob=False, sampling_params=SimpleNamespace(max_new_tokens=8), - kv=SimpleNamespace(cache_protected_len=0), + kv=SimpleNamespace(req_pool_idx=int(priority) % 8, cache_protected_len=0), time_stats=MagicMock(), ) return SimpleNamespace( diff --git a/test/registered/unit/managers/test_schedule_batch_req_pool_indices.py b/test/registered/unit/managers/test_schedule_batch_req_pool_indices.py index 1e129fbfc..1c7c5a5e1 100644 --- a/test/registered/unit/managers/test_schedule_batch_req_pool_indices.py +++ b/test/registered/unit/managers/test_schedule_batch_req_pool_indices.py @@ -18,7 +18,7 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu") def _make_req(req_pool_idx, origin_input_ids, output_ids): return types.SimpleNamespace( - req_pool_idx=req_pool_idx, + kv=types.SimpleNamespace(req_pool_idx=req_pool_idx), origin_input_ids=origin_input_ids, output_ids=output_ids, return_logprob=False, diff --git a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py index dc7c45e5c..efb1de522 100644 --- a/test/registered/unit/managers/test_scheduler_chunked_req_gate.py +++ b/test/registered/unit/managers/test_scheduler_chunked_req_gate.py @@ -34,11 +34,10 @@ def _make_req( req.output_ids = array("q") req.full_untruncated_fill_ids = array("q", fill_ids) req.prefix_indices = prefix_indices - req.req_pool_idx = req_pool_idx req.extend_range = Range(fill_len - extend_input_len, fill_len) req.inflight_middle_chunks = 0 req.host_hit_length = 0 - req.kv = ReqKvInfo() + req.kv = ReqKvInfo(req_pool_idx=req_pool_idx) req.skip_radix_cache_insert = False req.last_node = None req.swa_uuid_for_lock = None 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 f574baa5a..0f9957b7c 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 @@ -75,13 +75,13 @@ class MockReq: "q", fill_ids[:-1] if len(fill_ids) > 1 else fill_ids ) self.output_ids = array("q", [fill_ids[-1]] if len(fill_ids) > 1 else []) - self.req_pool_idx = req_pool_idx self.last_node = last_node self.extra_key = None self.cache_salt = None self.prefix_indices = torch.empty(0, dtype=torch.int64) self.priority = 0 self.kv = SimpleNamespace( + req_pool_idx=req_pool_idx, kv_committed_len=len(fill_ids), kv_allocated_len=len(fill_ids), cache_protected_len=cache_protected_len, diff --git a/test/registered/unit/mem_cache/test_decode_retraction_backup.py b/test/registered/unit/mem_cache/test_decode_retraction_backup.py index 4330a2bd9..24e53ffee 100644 --- a/test/registered/unit/mem_cache/test_decode_retraction_backup.py +++ b/test/registered/unit/mem_cache/test_decode_retraction_backup.py @@ -129,12 +129,14 @@ class TestDecodeRetractionBackup(unittest.TestCase): ) def _admit_req(self, env, num_tokens: int): - req = SimpleNamespace(rid="request", req_pool_idx=None, seqlen=num_tokens + 1) + req = SimpleNamespace( + rid="request", kv=SimpleNamespace(req_pool_idx=None), seqlen=num_tokens + 1 + ) self.assertIsNotNone(env.req_to_token_pool.alloc([req])) source_indices = env.allocator.alloc(num_tokens) self.assertIsNotNone(source_indices) env.req_to_token_pool.write( - (req.req_pool_idx, slice(0, num_tokens)), source_indices + (req.kv.req_pool_idx, slice(0, num_tokens)), source_indices ) return req, source_indices @@ -200,7 +202,7 @@ class TestDecodeRetractionBackup(unittest.TestCase): self.assertIsNotNone(destination_indices) self.assertFalse(torch.equal(source_indices, destination_indices)) req_to_token_pool.write( - (req.req_pool_idx, slice(0, self.num_tokens)), destination_indices + (req.kv.req_pool_idx, slice(0, self.num_tokens)), destination_indices ) cache.retraction_restore(req, backup) diff --git a/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py b/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py index b9aae98ee..fb9879a13 100644 --- a/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py +++ b/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py @@ -61,10 +61,10 @@ def _make_req(rid, prefix, block_size, *, req_pool_idx=None, reuse=False): return SimpleNamespace( rid=rid, prefix_indices=torch.tensor(prefix, dtype=torch.int32), - req_pool_idx=req_pool_idx, 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, kv=SimpleNamespace( + req_pool_idx=req_pool_idx, 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 @@ -75,8 +75,8 @@ def _make_req(rid, prefix, block_size, *, req_pool_idx=None, reuse=False): def _remove_allocated_req_slots(pool, *reqs): for req in reqs: - if req.req_pool_idx in pool.free_slots: - pool.free_slots.remove(req.req_pool_idx) + if req.kv.req_pool_idx in pool.free_slots: + pool.free_slots.remove(req.kv.req_pool_idx) def _make_batch(pool, allocator, reqs, extend_lens): @@ -105,8 +105,8 @@ def _make_batch(pool, allocator, reqs, extend_lens): def _seed_retained_block(pool, req, values): prefix_len = len(req.prefix_indices) - pool.req_to_token[req.req_pool_idx, :prefix_len] = req.prefix_indices - pool.req_to_token[req.req_pool_idx, prefix_len : prefix_len + len(values)] = ( + pool.req_to_token[req.kv.req_pool_idx, :prefix_len] = req.prefix_indices + pool.req_to_token[req.kv.req_pool_idx, prefix_len : prefix_len + len(values)] = ( torch.tensor(values, dtype=torch.int32) ) diff --git a/test/registered/unit/mem_cache/test_hisparse_allocator.py b/test/registered/unit/mem_cache/test_hisparse_allocator.py index 7b6337981..7ea10c1f4 100644 --- a/test/registered/unit/mem_cache/test_hisparse_allocator.py +++ b/test/registered/unit/mem_cache/test_hisparse_allocator.py @@ -106,7 +106,7 @@ class TestDeepSeekV4HiSparseAllocator(CustomTestCase): def alloc(self, reqs): for item in reqs: - item.req_pool_idx = 0 + item.kv.req_pool_idx = 0 return torch.tensor([0], dtype=torch.int64) def write(self, indices, values): @@ -160,7 +160,7 @@ class TestDeepSeekV4HiSparseAllocator(CustomTestCase): regular_host_alloc.assert_called_once_with( coordinator.req_to_host_pool, coordinator.req_to_host_pool_allocated_len, - req.req_pool_idx, + req.kv.req_pool_idx, 0, len(host_indices), ) diff --git a/test/registered/unit/mem_cache/test_paged_free_segment.py b/test/registered/unit/mem_cache/test_paged_free_segment.py index b2630cb8e..6203677fc 100644 --- a/test/registered/unit/mem_cache/test_paged_free_segment.py +++ b/test/registered/unit/mem_cache/test_paged_free_segment.py @@ -131,7 +131,7 @@ class TestFreeSegment(unittest.TestCase): token_to_kv_pool_allocator=alloc, req_to_token_pool=SimpleNamespace(req_to_token=row.unsqueeze(0)), ) - req = SimpleNamespace(req_pool_idx=0) + req = SimpleNamespace(kv=SimpleNamespace(req_pool_idx=0)) before = len(alloc.free_pages) alloc.free_group_begin() diff --git a/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py b/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py index e2a358efd..438159acb 100644 --- a/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py +++ b/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py @@ -21,11 +21,12 @@ class _FakeAllocator: class _FakeReq: - req_pool_idx = 0 - def __init__(self): self.kv = SimpleNamespace( - swa_evicted_seqlen=6, swa_evict_floor=3, cache_protected_len=0 + req_pool_idx=0, + swa_evicted_seqlen=6, + swa_evict_floor=3, + cache_protected_len=0, ) def pop_committed_kv_cache(self): diff --git a/test/registered/unit/mem_cache/test_radix_cache_unit.py b/test/registered/unit/mem_cache/test_radix_cache_unit.py index 0998f3827..cb792e31b 100644 --- a/test/registered/unit/mem_cache/test_radix_cache_unit.py +++ b/test/registered/unit/mem_cache/test_radix_cache_unit.py @@ -518,8 +518,7 @@ class TestRadixCache(unittest.TestCase): ) cache.req_to_token_pool = ReqToTokenPool(request_indices.clone()) req = unittest.mock.Mock( - req_pool_idx=0, - kv=SimpleNamespace(cache_protected_len=0), + kv=SimpleNamespace(req_pool_idx=0, cache_protected_len=0), extra_key=None, cache_salt=None, priority=0, diff --git a/test/registered/unit/mem_cache/test_retraction_mamba_backup.py b/test/registered/unit/mem_cache/test_retraction_mamba_backup.py index 7b678f35a..fc4ef7dce 100644 --- a/test/registered/unit/mem_cache/test_retraction_mamba_backup.py +++ b/test/registered/unit/mem_cache/test_retraction_mamba_backup.py @@ -2,7 +2,7 @@ import unittest import torch -from sglang.srt.managers.schedule_batch import Req +from sglang.srt.managers.schedule_batch import Req, ReqKvInfo from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool from sglang.test.ci.ci_register import register_cpu_ci @@ -39,7 +39,7 @@ class _Allocator: def _req_and_pool(): req = object.__new__(Req) - req.req_pool_idx = 0 + req.kv = ReqKvInfo(req_pool_idx=0) req.origin_input_ids = [1, 2] req.output_ids = [3] req.mamba_pool_idx = torch.tensor(1) diff --git a/test/registered/unit/mem_cache/test_streaming_session_unit.py b/test/registered/unit/mem_cache/test_streaming_session_unit.py index 2bed0f237..34657afc6 100644 --- a/test/registered/unit/mem_cache/test_streaming_session_unit.py +++ b/test/registered/unit/mem_cache/test_streaming_session_unit.py @@ -2,7 +2,7 @@ from types import SimpleNamespace import torch -from sglang.srt.managers.schedule_batch import FINISH_ABORT +from sglang.srt.managers.schedule_batch import FINISH_ABORT, ReqKvInfo from sglang.srt.mem_cache.base_prefix_cache import MatchResult from sglang.srt.session.streaming_session import SessionSlot, StreamingSession from sglang.test.ci.ci_register import register_cpu_ci @@ -24,8 +24,8 @@ class _FakeReqToTokenPool: self.free_slots = [] def free(self, req): - self.free_slots.append(req.req_pool_idx) - req.req_pool_idx = None + self.free_slots.append(req.kv.req_pool_idx) + req.kv.req_pool_idx = None class _FakeInnerCache: @@ -67,8 +67,8 @@ class _FakeReq: abort_req=lambda: None, _inflight=False, ) - self.req_pool_idx = req_pool_idx - self.kv = SimpleNamespace( + self.kv = ReqKvInfo( + req_pool_idx=req_pool_idx, kv_committed_len=committed, kv_allocated_len=allocated, swa_evicted_seqlen=0, @@ -112,8 +112,8 @@ def test_preabort_detaches_session_and_preserves_slot(): ) tree_cache = StreamingSession(inner) tree_cache.slots["session-a"] = SessionSlot( - req_pool_idx=0, - kv=SimpleNamespace( + kv=ReqKvInfo( + req_pool_idx=0, kv_committed_len=48, kv_allocated_len=48, swa_evicted_seqlen=0, @@ -135,7 +135,7 @@ def test_preabort_detaches_session_and_preserves_slot(): assert req.session is None # Slot untouched. slot = tree_cache.slots["session-a"] - assert slot.req_pool_idx == 0 + assert slot.kv.req_pool_idx == 0 assert slot.kv.kv_committed_len == 48 assert slot.kv.kv_allocated_len == 48 assert len(result.device_indices) == 0 @@ -160,7 +160,7 @@ def test_first_mid_abort_nukes_ephemeral_slot(): # Slot must NOT be created. assert "session-a" not in tree_cache.slots # Transient pool slot freed. - assert req.req_pool_idx is None + assert req.kv.req_pool_idx is None assert req_to_token_pool.free_slots == [0] assert len(allocator.freed) == 1 assert allocator.freed[0].tolist() == list(range(20)) @@ -179,8 +179,8 @@ def test_nth_mid_abort_nukes_session_slot(): # Session already has a slot from a previous turn. tree_cache.slots["session-a"] = SessionSlot( - req_pool_idx=0, - kv=SimpleNamespace( + kv=ReqKvInfo( + req_pool_idx=0, kv_committed_len=50, kv_allocated_len=50, swa_evicted_seqlen=0, @@ -202,7 +202,7 @@ def test_nth_mid_abort_nukes_session_slot(): assert allocator.freed[0].tolist() == list(range(65)) # Pool slot returned. assert req_to_token_pool.free_slots == [0] - assert req.req_pool_idx is None + assert req.kv.req_pool_idx is None def test_release_session_threads_mamba_skip_ids(): @@ -220,8 +220,8 @@ def test_release_session_threads_mamba_skip_ids(): lock_node = SimpleNamespace(id=42) tree_cache.slots["session-a"] = SessionSlot( - req_pool_idx=0, - kv=SimpleNamespace( + kv=ReqKvInfo( + req_pool_idx=0, kv_committed_len=50, kv_allocated_len=50, swa_evicted_seqlen=0, diff --git a/test/registered/unit/mem_cache/test_swa_eviction_boundary.py b/test/registered/unit/mem_cache/test_swa_eviction_boundary.py index 1ca7f063b..94a04d52a 100644 --- a/test/registered/unit/mem_cache/test_swa_eviction_boundary.py +++ b/test/registered/unit/mem_cache/test_swa_eviction_boundary.py @@ -103,11 +103,11 @@ def _build_swa_tree(page_size, sliding_window_size, kv_size=1024, kv_size_swa=51 def _make_req(req_pool_idx, token_ids, cache_protected_len, tree): """Mock Req with fields needed by _evict_swa and cache_finished_req.""" req = SimpleNamespace( - req_pool_idx=req_pool_idx, - is_holding_kv=True, origin_input_ids=token_ids, output_ids=[], - kv=ReqKvInfo(cache_protected_len=cache_protected_len), + kv=ReqKvInfo( + req_pool_idx=req_pool_idx, cache_protected_len=cache_protected_len + ), extra_key=None, cache_salt=None, last_node=tree.root_node, diff --git a/test/registered/unit/mem_cache/test_swa_unittest.py b/test/registered/unit/mem_cache/test_swa_unittest.py index 9595926e3..b01e5ca86 100644 --- a/test/registered/unit/mem_cache/test_swa_unittest.py +++ b/test/registered/unit/mem_cache/test_swa_unittest.py @@ -669,13 +669,13 @@ class TestSWA(unittest.TestCase): # Case 1: is_insert=True should pass bigram key and use cache_protected_len. req = _DummyReq() - req.req_pool_idx = 0 + req.kv.req_pool_idx = 0 req.origin_input_ids = array("q", [1, 2, 3, 4, 5, 6]) req.output_ids = array("q") req._kv_committed_len = len(req.origin_input_ids) kv_indices = allocator.alloc(req._kv_committed_len) req_to_token_pool.write( - (req.req_pool_idx, slice(0, req._kv_committed_len)), kv_indices + (req.kv.req_pool_idx, slice(0, req._kv_committed_len)), kv_indices ) req.extra_key = None req.cache_salt = None @@ -707,13 +707,13 @@ class TestSWA(unittest.TestCase): # Case 2: is_insert=False should free [cache_protected_len:page_aligned_len] # even when len(prefix_indices) is intentionally larger. req2 = _DummyReq() - req2.req_pool_idx = 1 + req2.kv.req_pool_idx = 1 req2.origin_input_ids = array("q", [11, 12, 13, 14, 15, 16]) req2.output_ids = array("q") req2._kv_committed_len = len(req2.origin_input_ids) kv_indices2 = allocator.alloc(req2._kv_committed_len) req_to_token_pool.write( - (req2.req_pool_idx, slice(0, req2._kv_committed_len)), kv_indices2 + (req2.kv.req_pool_idx, slice(0, req2._kv_committed_len)), kv_indices2 ) req2.extra_key = None req2.cache_salt = None @@ -911,7 +911,7 @@ class TestCacheUnfinishedReqEvictedPrefix(CustomTestCase): token_ids = array("q", range(1, num_tokens + 1)) req = _DummyReq() - req.req_pool_idx = 0 + req.kv.req_pool_idx = 0 req.origin_input_ids = token_ids req.output_ids = array("q") req.get_fill_ids = lambda: token_ids diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py index 527a8c3a6..90083afa2 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py @@ -646,7 +646,7 @@ def bench_cache_finished( req.kv.kv_committed_len = len(seq) if hasattr(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.kv.req_pool_idx, : len(kv_indices)] = kv_indices req_items.append(req) if not req_items: diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 834dd2e81..f991e0c9d 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -1235,7 +1235,7 @@ class UnifiedRadixCacheSuite: req.output_ids = array("q", output_ids) kv_len = len(input_ids) + len(output_ids) 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.kv.req_pool_idx, slice(0, kv_len)), kv_indices) req.kv.kv_committed_len = kv_len req.last_node = cache.root_node_handle() req.kv.cache_protected_len = 0 @@ -1275,7 +1275,7 @@ class UnifiedRadixCacheSuite: ) kv_len = req.extend_range.end 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.kv.req_pool_idx, slice(0, kv_len)), kv_indices) req.kv.kv_committed_len = kv_len req.kv.kv_allocated_len = kv_len req.last_node = cache.root_node_handle() @@ -1297,7 +1297,7 @@ class UnifiedRadixCacheSuite: start_p = ((start_p + ps - 1) // ps) * ps if start_p < end_p: allocator.free( - req_to_token_pool.req_to_token[req.req_pool_idx][start_p:end_p] + req_to_token_pool.req_to_token[req.kv.req_pool_idx][start_p:end_p] ) prompt_aligned = (len(prompt_ids) // ps) * ps @@ -1320,7 +1320,7 @@ class UnifiedRadixCacheSuite: req.output_ids = array("q") kv_len = len(tokens) 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.kv.req_pool_idx, slice(0, kv_len)), kv_indices) req.kv.kv_committed_len = kv_len req.last_node = cache.root_node_handle() req.kv.cache_protected_len = 0 @@ -1355,7 +1355,7 @@ class UnifiedRadixCacheSuite: ) kv_len = len(tokens) 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.kv.req_pool_idx, slice(0, kv_len)), kv_indices) req.kv.kv_committed_len = kv_len req.last_node = cache.root_node_handle() req.kv.cache_protected_len = 0 @@ -1392,7 +1392,9 @@ class UnifiedRadixCacheSuite: req.full_untruncated_fill_ids = array("q", tokens) req.set_extend_range(0, len(req.full_untruncated_fill_ids)) 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.kv.req_pool_idx, slice(0, len(tokens))), kv_indices + ) req.kv.kv_committed_len = len(tokens) req.last_node = cache.root_node_handle() req.kv.cache_protected_len = 0 @@ -1494,7 +1496,7 @@ class UnifiedRadixCacheSuite: req.output_ids = array("q") kv_len = len(input_ids) 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.kv.req_pool_idx, slice(0, kv_len)), kv_indices) req.kv.kv_committed_len = kv_len req.last_node = cache.root_node_handle() req.kv.cache_protected_len = 0 @@ -1619,7 +1621,7 @@ class UnifiedRadixCacheSuite: req.set_extend_range(0, len(req.full_untruncated_fill_ids)) kv_len = len(tokens) 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.kv.req_pool_idx, slice(0, kv_len)), fresh_value) req.kv.kv_committed_len = kv_len req.last_node = cache.root_node_handle() req.kv.cache_protected_len = 0 @@ -2206,7 +2208,7 @@ class UnifiedRadixCacheSuite: req.full_untruncated_fill_ids = array("q", tokens) req.set_extend_range(0, len(req.full_untruncated_fill_ids)) 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.kv.req_pool_idx, slice(0, pre_len)), kv_indices) req.kv.kv_committed_len = pre_len req.last_node = cache.root_node_handle() req.kv.cache_protected_len = 0 @@ -2295,7 +2297,7 @@ class UnifiedRadixCacheSuite: req.full_untruncated_fill_ids = array("q", tokens) req.set_extend_range(0, len(req.full_untruncated_fill_ids)) 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.kv.req_pool_idx, slice(0, pre_len)), kv_indices) req.kv.kv_committed_len = pre_len req.last_node = cache.root_node_handle() req.kv.cache_protected_len = 0 @@ -6790,7 +6792,9 @@ class TestUnifiedRadixCacheInt8MambaCheckpoint(CustomTestCase): req = self._make_req(req_to_token_pool, tokens) kv_indices = allocator.alloc(len(tokens)) self.assertIsNotNone(kv_indices) - req_to_token_pool.write((req.req_pool_idx, slice(0, len(tokens))), kv_indices) + req_to_token_pool.write( + (req.kv.req_pool_idx, slice(0, len(tokens))), kv_indices + ) req.last_node = cache.root_node_handle() cache.cache_finished_req( @@ -8104,7 +8108,7 @@ class TestSWAWindowUnderBigramKey(CustomTestCase): req.full_untruncated_fill_ids = array("q", tokens) req.set_extend_range(0, len(req.full_untruncated_fill_ids)) 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.kv.req_pool_idx, slice(0, seq_len)), kv_indices) req.kv.kv_committed_len = seq_len req.last_node = cache.root_node_handle() req.kv.cache_protected_len = 0 diff --git a/test/registered/xpu/test_lmcache_radix_cache.py b/test/registered/xpu/test_lmcache_radix_cache.py index bf0100cd6..51b414d9c 100644 --- a/test/registered/xpu/test_lmcache_radix_cache.py +++ b/test/registered/xpu/test_lmcache_radix_cache.py @@ -66,13 +66,16 @@ def _make_req(rid, req_pool_idx, token_ids, tree): SimpleNamespace pattern in test_swa_eviction_boundary.py).""" req = SimpleNamespace( rid=rid, - req_pool_idx=req_pool_idx, origin_input_ids=token_ids, output_ids=[], extra_key=None, cache_salt=None, last_node=tree.root_node, - kv=SimpleNamespace(cache_protected_len=0, kv_committed_len=len(token_ids)), + kv=SimpleNamespace( + req_pool_idx=req_pool_idx, + cache_protected_len=0, + kv_committed_len=len(token_ids), + ), priority=0, kv_committed_freed=False, ) @@ -162,7 +165,11 @@ class TestLMCRadixCacheXPU(unittest.TestCase): # commit it as a finished request (inserts into radix + stores to # LMCache on tree.store_stream). req_pool_idx = req_to_token_pool.alloc( - [SimpleNamespace(req_pool_idx=None, inflight_middle_chunks=0)] + [ + SimpleNamespace( + inflight_middle_chunks=0, kv=SimpleNamespace(req_pool_idx=None) + ) + ] )[0] kv_slots = allocator.alloc(self.INPUT_LEN) self.assertIsNotNone(kv_slots)