diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index 93b2a286b..46eef0302 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -1427,16 +1427,17 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): def _flush(self, *, urgent: bool) -> int: """One batched compaction pass; returns the number of survivor moves. - Pipeline (one D2H total, at step 3): + Pipeline (one free-list D2H plus one mapping D2H per committed move batch): 1. `_drain_pending_reuse` — return read-settled prior srcs. 2. sort the free list (or skip via env knob; either way ascending after). - 3. `.tolist()` snapshot → `all_cpu` *(the one sync)*. + 3. `.tolist()` snapshot → `all_cpu`. 4-5. `_absorb_boundary_holes` — retreat past boundary-contiguous holes; `holes_cpu` = interior holes. After this `_free_phys_pages==holes_cpu`. 6. (urgent) `_settle_inflight_forward` — wait once so the walk is race-free. 7. survivor walk — TWO-POINTER: move topmost live slot into the next hole, STOPPING when the pointers cross (band packed); batch into one - `move_kv_cache` + one v2p/p2v scatter at `_commit_move_batch`. + `move_kv_cache` + one v2p/p2v scatter at `_commit_move_batch`, which + gathers and validates all survivor virtual ids in one batch. 8-9. exit: urgent → FULL-PACK reclaim (retreat past ALL holes, empty list); non-urgent → slice consumed dsts, merge freed srcs back. @@ -1447,8 +1448,8 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): `_commit_move_batch` routes such srcs to `_pending_reuse`; urgent's settle makes them immediately reusable. - `_topmost_survivor` excludes all p2v=-1 pages, so a `v_moved < 0` in the - loop is a corrupt-state bug and raises. + `_topmost_survivor` excludes all p2v=-1 pages, so a negative virtual id in + the batched mapping lookup is a corrupt-state bug and raises. """ if not self.lazy_compaction: return 0 @@ -1463,7 +1464,7 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): if not _SORT_FREE_LIST_AFTER_MERGE and self._free_phys_pages.numel() > 1: self._free_phys_pages, _ = torch.sort(self._free_phys_pages) - all_cpu = self._free_phys_pages.tolist() # the ONE D2H sync per flush + all_cpu = self._free_phys_pages.tolist() # one batched D2H sync # `holes_cpu` = interior holes; `_free_phys_pages == holes_cpu` after. new_wm, holes_cpu = self._absorb_boundary_holes(all_cpu) @@ -1487,7 +1488,6 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): srcs: List[int] = [] dsts: List[int] = [] - v_moveds: List[int] = [] # Flush-scoped accumulator for event-FIRED srcs. `_commit_move_batch` # appends here instead of catting onto `_free_phys_pages`; the merge is @@ -1532,12 +1532,11 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): # Commit accumulated moves, then wait the forward so the # rest of the walk is race-free. self._commit_move_batch( - srcs, dsts, v_moveds, latest_event, released_fired + srcs, dsts, latest_event, released_fired ) n_moves += len(srcs) srcs.clear() dsts.clear() - v_moveds.clear() inflight = self._inflight_forward if inflight is not None: torch.cuda.current_stream().wait_event(inflight[0]) @@ -1568,22 +1567,8 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): dst_cursor -= 1 n_dst_consumed += 1 - v_moved = int(self.physical_to_virtual[src].item()) - if v_moved < 0: - # `_topmost_survivor` excludes all p2v=-1 pages — corrupt state. - raise AssertionError( - f"MultiEndedAllocator({self.sub_pool_name!r})." - f"_flush: topmost survivor p={src} has p2v=-1; " - "this should be impossible (`_topmost_survivor` " - "excludes `holes_cpu` and `_pending_reuse_pages_cpu`)." - f" State: {self.allocator_state_str()}, " - f"#holes={len(holes_cpu)}, " - f"#pending_reuse={len(self._pending_reuse_pages_cpu)}" - ) - srcs.append(src) dsts.append(dst) - v_moveds.append(v_moved) # Advance cursor strictly past the picked src. if self.grow_direction == "up": @@ -1594,7 +1579,7 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): if move_cap is not None and len(srcs) >= move_cap: break - self._commit_move_batch(srcs, dsts, v_moveds, latest_event, released_fired) + self._commit_move_batch(srcs, dsts, latest_event, released_fired) n_moves += len(srcs) if single_pass_absorb: @@ -1635,12 +1620,12 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): self, srcs: List[int], dsts: List[int], - v_moveds: List[int], latest_event: Optional[torch.cuda.Event], released_fired: List[torch.Tensor], ) -> None: """Issue ONE `move_kv_cache` + ONE bulk v2p/p2v remap for the accumulated - `(src, dst, v_moved)` triples. Fired srcs accumulate in `released_fired` + `(src, dst)` pairs. Survivor virtual ids are gathered from p2v in one + batch. Fired srcs accumulate in `released_fired` (merged by `_flush` AFTER its dst-slice, keeping the free list == holes_cpu); event-pending srcs route to `_pending_reuse` (read-race gating). """ @@ -1649,7 +1634,11 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): with record_function("MultiEndedAlloc._commit_move_batch"): src_pages_t = torch.tensor(srcs, dtype=torch.int64, device=self.device) dst_pages_t = torch.tensor(dsts, dtype=torch.int64, device=self.device) - v_moveds_t = torch.tensor(v_moveds, dtype=torch.int64, device=self.device) + v_moveds_t = self.physical_to_virtual[src_pages_t] + torch._assert_async( + (v_moveds_t >= 0).all(), + "invalid p2v mapping in MultiEndedAllocator._flush", + ) # Expand to token granularity (the move kernel is token-granular). if self.page_size == 1: src_t, dst_t = src_pages_t, dst_pages_t diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py index 115de9e21..9b8adef09 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -79,6 +79,24 @@ class _FakeKVCache: self.buf[dst_loc] = self.buf[src_loc].clone() +class _RejectScalarIndexTensor: + """Tensor proxy that rejects one-row-at-a-time mapping lookups.""" + + def __init__(self, tensor: torch.Tensor): + self.tensor = tensor + + def __getattr__(self, name): + return getattr(self.tensor, name) + + def __getitem__(self, index): + if isinstance(index, int): + raise AssertionError("physical_to_virtual was read one row at a time") + return self.tensor[index] + + def __setitem__(self, index, value): + self.tensor[index] = value + + class TestUnifiedKVPoolViews(unittest.TestCase): def test_min_slot_index_and_disjoint_bytes(self): full = _make_mha_spec("full", "up", layer_num=4) @@ -2049,6 +2067,27 @@ class TestLazyCompaction(unittest.TestCase): continue # freed self.assertEqual(int(fa.physical_to_virtual[p].item()), v) + def test_lazy_flush_gathers_survivor_mappings_as_one_batch(self): + """Compaction must not synchronize once per relocated survivor.""" + _pool, fa, kv = self._make_full(lazy=True) + values = fa.alloc(12) + self._stamp_kv(kv, fa, values) + fa.free(values[1:5].clone()) + + physical_to_virtual = fa.physical_to_virtual + fa.physical_to_virtual = _RejectScalarIndexTensor(physical_to_virtual) + try: + self.assertEqual(fa._flush(urgent=True), 4) + finally: + fa.physical_to_virtual = physical_to_virtual + + for virtual in values.tolist(): + physical = int(fa.virtual_to_physical[virtual].item()) + if physical == -1: + continue + self.assertEqual(int(fa.physical_to_virtual[physical].item()), virtual) + self.assertEqual(int(kv.buf[physical].item()), virtual) + def _replay_sequence(self, ops, lazy: bool): """Run a given alloc/free op trace under eager OR lazy mode and return the final (live virtual set, alloc-time KV stamps)."""