From 2a980cbf10db36c6052042cc3f92b4a80c4047b6 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Thu, 3 Sep 2026 13:28:33 -0700 Subject: [PATCH] [mem_cache] Require page-aligned starts in `free_segment` and drop the boundary trim (#37729) --- .../decode_kvcache_offload_manager.py | 23 +------ python/sglang/srt/mem_cache/allocator/base.py | 31 +++++---- .../sglang/srt/mem_cache/allocator/paged.py | 27 ++++---- .../srt/mem_cache/multi_ended_allocator.py | 48 +++++--------- .../test_specv2_kvcache_offloading.py | 64 +++++-------------- .../unit/mem_cache/test_paged_free_segment.py | 60 +++++++++-------- .../test_unified_free_no_host_sync.py | 28 ++++---- 7 files changed, 113 insertions(+), 168 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py index eed819e08..b58adc483 100644 --- a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py +++ b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py @@ -27,7 +27,6 @@ from sglang.srt.runtime_context import ( get_schedule, get_serving, ) -from sglang.srt.utils.common import ceil_align if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req @@ -245,25 +244,9 @@ class DecodeKVCacheOffloadManager: if req.kv.req_pool_idx is None or req.kv.req_pool_idx == -1: return - kv_committed_len = req.effective_kv_committed_len() - - # Prefill-aligned slots are freed only here, at request finish; freeing - # them mid-decode races with concurrent admission over live slots. - prefill_len = self._prefill_offloaded_len(req) - ranges = [] - if prefill_len > 0: - ranges.append((0, prefill_len)) - # The incremental part of the request (DSA-aware) - ranges.append((prefill_len, kv_committed_len)) - - # Over-allocated KV cache slots (e.g. from speculative decoding v2). - # Without spec v2, start_p == end_p so this contributes nothing. - start_p, end_p = kv_committed_len, req.kv.kv_allocated_len - if self.page_size > 1: - start_p = ceil_align(start_p, self.page_size) - if start_p < end_p: - ranges.append((start_p, end_p)) - self.tree_cache.free_kv_row(req.kv, ranges) + # Released only at request finish; a mid-decode free races with + # concurrent admission over live slots. + self.tree_cache.free_kv_row(req.kv, [(0, req.kv.kv_allocated_len)]) self.req_to_token_pool.free(req) req.kv.mark_kv_released() diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index fc1cbe257..c9826e068 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -174,25 +174,32 @@ class BaseTokenToKVPoolAllocator(abc.ABC): self.free(free_index) def free_segment(self, free_index: torch.Tensor, *, start_pos: int): - """Free ``kv_row[start_pos : start_pos + n]`` of one request (or a - page-aligned copy); subclasses may use ``start_pos`` to skip the - data-dependent dedup. Default: plain free().""" + """Free ``kv_row[start_pos : start_pos + n]`` of one request. + + In page units the segment is ``[start_pos // ps, ceil(end / ps))``: + ``start_pos`` sits on a page boundary, the end may fall mid-page, and + the whole last page is released. Default: plain free().""" + assert start_pos % self.page_size == 0, ( + f"segment start {start_pos} is not page-aligned" + ) self.free(free_index) def free_segments(self, segments): - """Free disjoint ascending ``(free_index, start_pos)`` segments of one - request's kv row; a boundary page shared by consecutive segments is - emitted once (the later segment's head is trimmed).""" + """Free several ``(free_index, start_pos)`` segments of one request's + kv row. + + Each segment covers the pages ``[start_pos // ps, ceil(end / ps))``. + Starts sit on page boundaries, ends may fall mid-page, and the page + ranges of consecutive segments do not overlap -- so in page units the + segments are aligned and disjoint, and every page is released once.""" ps = self.page_size prev_end = None for free_index, start_pos in segments: n = free_index.numel() if n == 0: continue - seg_end = start_pos + n - if prev_end is not None and start_pos // ps == (prev_end - 1) // ps: - boundary = (start_pos // ps + 1) * ps - free_index = free_index[boundary - start_pos :] - start_pos = boundary - prev_end = seg_end + assert prev_end is None or start_pos // ps > (prev_end - 1) // ps, ( + f"segment at {start_pos} shares a page with the one ending at {prev_end}" + ) + prev_end = start_pos + n self.free_segment(free_index, start_pos=start_pos) diff --git a/python/sglang/srt/mem_cache/allocator/paged.py b/python/sglang/srt/mem_cache/allocator/paged.py index b50272d26..fcb2b3ee2 100755 --- a/python/sglang/srt/mem_cache/allocator/paged.py +++ b/python/sglang/srt/mem_cache/allocator/paged.py @@ -282,36 +282,33 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self._debug_check_no_duplicate_pages() def free_segment(self, free_index: torch.Tensor, *, start_pos: int): - """Fixed-shape counterpart of free(): a page's tokens sit consecutively - in the kv row, so page representatives are stride slices -- no - torch.unique, whose data-dependent output shape forces a device sync. - Contract: see base; a page must be freed by only one call per group.""" + """Fixed-shape counterpart of free(). + + The segment starts on a page boundary and a page's tokens sit + consecutively in the kv row, so ``free_index[::page_size]`` is one + token from each page the segment covers -- including a partial last + page. No torch.unique, whose data-dependent output shape forces a + device sync. Contract: see base.""" if free_index.numel() == 0: return ps = self.page_size - offset = start_pos % ps - if offset == 0: - pieces = (free_index[::ps],) - else: - pieces = (free_index[:1], free_index[ps - offset :: ps]) + assert start_pos % ps == 0, f"segment start {start_pos} is not page-aligned" + reps = free_index[::ps] if self.debug_mode: # reference unique on CPU: the NPU subclass deliberately avoids device unique - page_ids = torch.cat([p // ps for p in pieces]) assert torch.equal( - torch.sort(page_ids.cpu())[0], + torch.sort(reps.cpu() // ps)[0], torch.unique(free_index.cpu() // ps), ) if self.free_group is None: - self._release_page_ids(*(p // ps for p in pieces)) + self._release_page_ids(reps // ps) if self.debug_mode: self._debug_check_no_duplicate_pages() else: - self.free_page_reps_group.extend( - self._copy_for_free_group(piece) for piece in pieces - ) + self.free_page_reps_group.append(self._copy_for_free_group(reps)) def _debug_check_no_duplicate_pages(self): pages = self.get_all_free_pages() diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index a5d371395..c6964dcab 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -1393,27 +1393,15 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): self.free_virtual_ids = torch.cat([self.free_virtual_ids, free_v_pages]) self._compact_pending(freed_p_pages) - def _page_reps_pieces( - self, free_index: torch.Tensor, start_pos: int - ) -> Tuple[torch.Tensor, ...]: - """Page-representative TOKEN slices of one kv-row segment. - - Mirrors `PagedTokenToKVPoolAllocator.free_segment`: a page's tokens sit - consecutively in the kv row, so with `start_pos` known on the host the - representatives are stride slices -- no `torch.unique`, whose - data-dependent output shape forces a device sync. - - Exact for any segment shape: a partial head page is the `[:1]` term, a - partial tail page the final stride step. - """ + def _page_reps(self, free_index: torch.Tensor, start_pos: int) -> torch.Tensor: + """One token of every page touched by a page-aligned kv-row segment: + the fixed-shape stand-in for `unique(free_index // page_size)`.""" ps = self.page_size - offset = start_pos % ps - if offset == 0: - return (free_index[::ps],) - return (free_index[:1], free_index[ps - offset :: ps]) + assert start_pos % ps == 0, f"segment start {start_pos} is not page-aligned" + return free_index[::ps] def free_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None: - """Fixed-shape counterpart of `free()`; see `_page_reps_pieces`. + """Fixed-shape counterpart of `free()`; see `_page_reps`. Contract: see base; a page must be freed by only one call per group. """ @@ -1423,12 +1411,11 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): # token == page: nothing to dedup, the plain path is already exact. self.free(free_index) return - pieces = self._page_reps_pieces(free_index.detach().to(torch.int64), start_pos) + reps = self._page_reps(free_index.detach().to(torch.int64), start_pos) if self.free_page_reps_group is None: - reps = pieces[0] if len(pieces) == 1 else torch.cat(pieces) self.free(reps, _pages=reps // self.page_size) else: - self.free_page_reps_group.extend(pieces) + self.free_page_reps_group.append(reps) def _free_lazy( self, free_index: torch.Tensor, pages: Optional[torch.Tensor] = None @@ -3054,7 +3041,7 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): def free_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None: """Fixed-shape counterpart of `free()`; see - `MultiEndedAllocator._page_reps_pieces`. The mamba sub-pool is + `MultiEndedAllocator._page_reps`. The mamba sub-pool is slot-granular and untouched by a token free, so only the full side needs the representatives. """ @@ -3063,13 +3050,13 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): if self.page_size == 1: self.free(free_index) return - pieces = self.full_attn_allocator._page_reps_pieces( + reps = self.full_attn_allocator._page_reps( free_index.detach().to(torch.int64), start_pos ) if self.free_page_reps_group is None: - self._release_page_reps(pieces) + self._release_page_reps((reps,)) else: - self.free_page_reps_group.extend(pieces) + self.free_page_reps_group.append(reps) def _release_page_reps(self, pieces: Sequence[torch.Tensor]) -> None: reps = pieces[0] if len(pieces) == 1 else torch.cat(tuple(pieces)) @@ -3637,8 +3624,7 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): v = free_index.detach().to(torch.int64) ps = self.page_size if start_pos is not None and ps > 1: - pieces = self.swa_attn_allocator._page_reps_pieces(v, start_pos) - reps = pieces[0] if len(pieces) == 1 else torch.cat(pieces) + reps = self.swa_attn_allocator._page_reps(v, start_pos) # Keep only pages still bound on swa (freeing a tombstoned one # would corrupt the hole list). `> 0` strict: -1 = tombstoned, # page 0 = padding sink (never freeable). @@ -3702,7 +3688,7 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): def free_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None: """Fixed-shape counterpart of `free()`; see - `MultiEndedAllocator._page_reps_pieces`. Both sides share one + `MultiEndedAllocator._page_reps`. Both sides share one derivation -- neither repeats the position-less dedup. """ if free_index is None or free_index.numel() == 0: @@ -3710,13 +3696,13 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): if self.page_size == 1: self.free(free_index) return - pieces = self.full_attn_allocator._page_reps_pieces( + reps = self.full_attn_allocator._page_reps( free_index.detach().to(torch.int64), start_pos ) if self.free_page_reps_group is None: - self._release_page_reps(pieces) + self._release_page_reps((reps,)) else: - self.free_page_reps_group.extend(pieces) + self.free_page_reps_group.append(reps) def _release_page_reps(self, pieces: Sequence[torch.Tensor]) -> None: reps = pieces[0] if len(pieces) == 1 else torch.cat(tuple(pieces)) diff --git a/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py b/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py index f915b68ec..70628017a 100644 --- a/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py +++ b/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py @@ -127,14 +127,12 @@ class TestReleaseFinishedReq(unittest.TestCase): manager._release_finished_req(req) - # Prefill [0:8] and committed [8:20]; no overalloc free. - self.assertEqual(len(freed), 2) - self.assertTrue(torch.equal(freed[0], torch.arange(0, 8, dtype=torch.int64))) - self.assertTrue(torch.equal(freed[1], torch.arange(8, 20, dtype=torch.int64))) + self.assertEqual(len(freed), 1) + self.assertTrue(torch.equal(freed[0], torch.arange(0, 20, dtype=torch.int64))) manager.req_to_token_pool.free.assert_called_once_with(req) def test_with_overallocation(self): - """With spec v2, overallocated slots [committed:allocated] must be freed.""" + """With spec v2, the over-allocated slots go back with the row.""" manager, freed = _make_manager(pool_size=32) req = _make_mock_req( req_pool_idx=0, @@ -145,15 +143,12 @@ class TestReleaseFinishedReq(unittest.TestCase): manager._release_finished_req(req) - # Prefill [0:8], committed [8:20], overallocated [20:28]. - self.assertEqual(len(freed), 3) - self.assertTrue(torch.equal(freed[0], torch.arange(0, 8, dtype=torch.int64))) - self.assertTrue(torch.equal(freed[1], torch.arange(8, 20, dtype=torch.int64))) - self.assertTrue(torch.equal(freed[2], torch.arange(20, 28, dtype=torch.int64))) + self.assertEqual(len(freed), 1) + self.assertTrue(torch.equal(freed[0], torch.arange(0, 28, dtype=torch.int64))) manager.req_to_token_pool.free.assert_called_once_with(req) - def test_overallocation_with_page_alignment(self): - """With page_size > 1, start of overallocated range is ceil-aligned.""" + def test_unaligned_committed_len_frees_the_whole_row(self): + """A mid-page committed length needs no alignment arithmetic here.""" page_size = 4 manager, freed = _make_manager(pool_size=32, page_size=page_size) req = _make_mock_req( @@ -165,30 +160,8 @@ class TestReleaseFinishedReq(unittest.TestCase): manager._release_finished_req(req) - # Prefill [0:4], committed [4:10], - # overallocated: start_p = ceil_align(10, 4) = 12, end_p = 28 => [12:28] - self.assertEqual(len(freed), 3) - self.assertTrue(torch.equal(freed[0], torch.arange(0, 4, dtype=torch.int64))) - self.assertTrue(torch.equal(freed[1], torch.arange(4, 10, dtype=torch.int64))) - self.assertTrue(torch.equal(freed[2], torch.arange(12, 28, dtype=torch.int64))) - - def test_overallocation_page_aligned_noop(self): - """When ceil_align(committed, page_size) >= allocated, no overalloc free.""" - page_size = 4 - manager, freed = _make_manager(pool_size=32, page_size=page_size) - req = _make_mock_req( - req_pool_idx=0, - kv_committed_len=10, # ceil_align(10, 4) = 12 - kv_allocated_len=12, # same as aligned start - origin_len=4, - ) - - manager._release_finished_req(req) - - # Prefill [0:4] and committed [4:10]; no overalloc since start_p == end_p - self.assertEqual(len(freed), 2) - self.assertTrue(torch.equal(freed[0], torch.arange(0, 4, dtype=torch.int64))) - self.assertTrue(torch.equal(freed[1], torch.arange(4, 10, dtype=torch.int64))) + self.assertEqual(len(freed), 1) + self.assertTrue(torch.equal(freed[0], torch.arange(0, 28, dtype=torch.int64))) def test_prefix_indices_decremented(self): """protected_size_ is decremented by len(req.prefix_indices).""" @@ -223,10 +196,8 @@ class TestReleaseFinishedReq(unittest.TestCase): manager._release_finished_req(req) - # Two frees in order: prefill [0:8] then committed [8:20]. - self.assertEqual(len(freed), 2) - self.assertTrue(torch.equal(freed[0], torch.arange(0, 8, dtype=torch.int64))) - self.assertTrue(torch.equal(freed[1], torch.arange(8, 20, dtype=torch.int64))) + self.assertEqual(len(freed), 1) + self.assertTrue(torch.equal(freed[0], torch.arange(0, 20, dtype=torch.int64))) # State entry is removed at the end of _release_finished_req. self.assertNotIn(req, manager.offloaded_state) @@ -267,12 +238,8 @@ class TestReleaseFinishedReq(unittest.TestCase): manager.finalize_release_on_finish(req) - # _release_finished_req frees prefill [0:12] then committed [12:13]. - self.assertEqual(len(freed), 2) - expected_prefill = torch.arange(0, 12, dtype=torch.int64) - expected_committed = torch.arange(12, 13, dtype=torch.int64) - self.assertTrue(torch.equal(freed[0], expected_prefill)) - self.assertTrue(torch.equal(freed[1], expected_committed)) + self.assertEqual(len(freed), 1) + self.assertTrue(torch.equal(freed[0], torch.arange(0, 13, dtype=torch.int64))) # No state entry is left behind. self.assertNotIn(req, manager.offloaded_state) @@ -452,9 +419,8 @@ class TestReleaseFinishedReq(unittest.TestCase): manager._check_offload_progress(1) - self.assertEqual(len(freed), 2) - self.assertTrue(torch.equal(freed[0], torch.arange(0, 4, dtype=torch.int64))) - self.assertTrue(torch.equal(freed[1], torch.arange(4, 20, dtype=torch.int64))) + self.assertEqual(len(freed), 1) + self.assertTrue(torch.equal(freed[0], torch.arange(0, 20, dtype=torch.int64))) manager.req_to_token_pool.free.assert_called_once_with(req) self.assertNotIn(req, manager.offloaded_state) self.assertNotIn(req, manager.offload_inflight) 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 4f0b125fc..b1cc88cef 100644 --- a/test/registered/unit/mem_cache/test_paged_free_segment.py +++ b/test/registered/unit/mem_cache/test_paged_free_segment.py @@ -1,5 +1,5 @@ -"""free_segment / free_segments vs the torch.unique reference: stride page -extraction over all segment alignments, plus boundary-page dedup and free-group +"""free_segment / free_segments vs the torch.unique reference: page-aligned +starts over every tail alignment, the page-disjoint contract, and free-group deferral. See PagedTokenToKVPoolAllocator.free_segment for why unique is avoided. python -m pytest test/registered/unit/mem_cache/test_paged_free_segment.py -v @@ -43,11 +43,9 @@ def _make_kv_row(alloc, num_tokens): class TestFreeSegment(unittest.TestCase): - def test_matches_unique_over_alignments(self): - # Sweep (start, end) so segments cover: aligned/unaligned head and - # tail, single partial page, full row. + def test_matches_unique_over_tail_alignments(self): for num_tokens in (1, PAGE_SIZE, PAGE_SIZE + 1, 3 * PAGE_SIZE - 1): - for start in range(num_tokens): + for start in range(0, num_tokens, PAGE_SIZE): for end in range(start + 1, num_tokens + 1): alloc = _make_allocator() row = _make_kv_row(alloc, num_tokens) @@ -67,6 +65,13 @@ class TestFreeSegment(unittest.TestCase): alloc.free_segment(row[:0], start_pos=0) self.assertEqual(len(alloc.free_pages), before) + def test_unaligned_start_is_rejected(self): + alloc = _make_allocator() + row = _make_kv_row(alloc, 2 * PAGE_SIZE) + for start in (1, PAGE_SIZE - 1, PAGE_SIZE + 1): + with self.assertRaises(AssertionError): + alloc.free_segment(row[start:], start_pos=start) + def test_need_sort_defers_released_pages(self): alloc = _make_allocator(need_sort=True) row = _make_kv_row(alloc, 2 * PAGE_SIZE) @@ -169,30 +174,22 @@ class TestFreeSegments(unittest.TestCase): reference = torch.unique(torch.cat([row[a:b] for a, b in spans]) // PAGE_SIZE) return freed, reference - def test_adjacent_segments_share_boundary_page(self): - # [0, 6) and [6, 11) with page_size 4: page 1 spans both segments and - # must be freed exactly once. - freed, reference = self._freed_by_segments(11, [(0, 6), (6, 11)]) + def test_partial_tail_then_next_page(self): + # [0, 5) releases page 1 whole; [8, 11) starts on page 2. + freed, reference = self._freed_by_segments(11, [(0, 5), (8, 11)]) self.assertTrue(torch.equal(torch.sort(freed)[0], reference)) - def test_disjoint_segments_share_boundary_page(self): - # [0, 5) and [7, 11): gap [5, 7) stays within page 1, which both - # segments touch. - freed, reference = self._freed_by_segments(11, [(0, 5), (7, 11)]) - self.assertTrue(torch.equal(torch.sort(freed)[0], reference)) - - def test_second_segment_inside_shared_page_is_skipped(self): - # [0, 5) and [5, 7): the second segment lies entirely in page 1, - # already emitted by the first. - freed, reference = self._freed_by_segments(7, [(0, 5), (5, 7)]) - self.assertTrue(torch.equal(torch.sort(freed)[0], reference)) - - def test_page_aligned_segments_no_trim(self): + def test_page_aligned_segments(self): freed, reference = self._freed_by_segments( 3 * PAGE_SIZE, [(0, PAGE_SIZE), (PAGE_SIZE, 3 * PAGE_SIZE)] ) self.assertTrue(torch.equal(torch.sort(freed)[0], reference)) + def test_segments_sharing_a_page_are_rejected(self): + for spans in ([(0, 5), (5, 8)], [(0, 5), (7, 11)], [(0, 6), (4, 11)]): + with self.assertRaises(AssertionError): + self._freed_by_segments(11, spans) + class _RecordingBaseAllocator(BaseTokenToKVPoolAllocator): """Base-fallback allocator: free_segment inherits the default (ignore @@ -220,15 +217,24 @@ class _RecordingBaseAllocator(BaseTokenToKVPoolAllocator): class TestBaseFallbackFreeSegments(unittest.TestCase): - def test_trim_dedups_boundary_page_before_fallback_free(self): - # fallback allocators (UnifiedMamba/SWA) dedup per free() call at best; - # the shared boundary page must reach free() in exactly one call + def test_fallback_forwards_page_disjoint_segments(self): + # base fallback: each segment reaches free() as-is, no cross-segment dedup alloc = _RecordingBaseAllocator() row = torch.arange(11) # position i lives on page i // PAGE_SIZE - alloc.free_segments([(row[0:6], 0), (row[6:11], 6)]) + alloc.free_segments([(row[0:6], 0), (row[8:11], 8)]) per_call_pages = [set((t // PAGE_SIZE).tolist()) for t in alloc.freed] self.assertEqual(per_call_pages, [{0, 1}, {2}]) + def test_fallback_rejects_shared_page_and_unaligned_start(self): + alloc = _RecordingBaseAllocator() + row = torch.arange(11) + with self.assertRaises(AssertionError): + alloc.free_segments([(row[0:6], 0), (row[6:11], 6)]) + with self.assertRaises(AssertionError): + alloc.free_segment(row[1:], start_pos=1) + self.assertEqual(len(alloc.freed), 1) + self.assertTrue(torch.equal(alloc.freed[0], row[0:6])) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py b/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py index b39989a47..364fdde98 100644 --- a/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py +++ b/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py @@ -264,11 +264,10 @@ class TestTombstonesDoNotCrossTheBus(unittest.TestCase): class TestFreeSegment(unittest.TestCase): """Mirrors `test_paged_free_segment.TestFreeSegment`.""" - def test_matches_unique_over_alignments(self): - """Sweep (start, end) so segments cover aligned/unaligned head and - tail, a single partial page, and the full row.""" + def test_matches_unique_over_tail_alignments(self): + """Page-aligned starts against every tail alignment.""" for num_tokens in (1, PAGE_SIZE, PAGE_SIZE + 1, 3 * PAGE_SIZE - 1): - for start in range(0, num_tokens, max(1, num_tokens // 4)): + for start in range(0, num_tokens, PAGE_SIZE): for end in (start + 1, num_tokens): if end <= start: continue @@ -286,7 +285,7 @@ class TestFreeSegment(unittest.TestCase): def test_never_calls_unique(self): """The decisive check -- make `torch.unique` explode. A textual guard can be fooled; this cannot.""" - for start in (0, 1, PAGE_SIZE - 1, PAGE_SIZE, PAGE_SIZE + 3): + for start in (0, PAGE_SIZE, 2 * PAGE_SIZE): alloc = _paged_allocator(lazy=True) row = alloc.alloc(3 * PAGE_SIZE) with self.subTest(start_pos=start): @@ -295,6 +294,12 @@ class TestFreeSegment(unittest.TestCase): ): alloc.free_segment(row[start : start + PAGE_SIZE], start_pos=start) + def test_unaligned_start_is_rejected(self): + alloc = _paged_allocator(lazy=True) + row = alloc.alloc(3 * PAGE_SIZE) + with self.assertRaises(AssertionError): + alloc.free_segment(row[1 : PAGE_SIZE + 1], start_pos=1) + def test_empty_segment_is_noop(self): alloc = _paged_allocator(lazy=True) before = alloc._free_phys_pages.numel() @@ -342,9 +347,7 @@ class TestFreeGroupKeepsPositions(unittest.TestCase): row = alloc.alloc(3 * PAGE_SIZE) alloc.free_group_begin() alloc.free_segment(row[:PAGE_SIZE], start_pos=0) - alloc.free_segment( - row[PAGE_SIZE + 3 : 2 * PAGE_SIZE + 3], start_pos=PAGE_SIZE + 3 - ) + alloc.free_segment(row[PAGE_SIZE : 2 * PAGE_SIZE + 3], start_pos=PAGE_SIZE) with mock.patch.object( torch, "unique", side_effect=AssertionError("sync path taken") ): @@ -468,14 +471,11 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase): alloc.free_swa(v[: 4 * self.PS], start_pos=0) alloc.free_swa(v[4 * self.PS :], start_pos=4 * self.PS) - def test_unaligned_start_pos_still_no_sync(self): - """`_page_reps_pieces` covers a misaligned start with a second piece; - the sync-free property must not depend on alignment.""" + def test_unaligned_start_pos_is_rejected(self): + """A mid-page start must fail loudly, not release the head page whole.""" alloc = self._swa_composite(lazy=True) v = alloc.alloc(8 * self.PS) - with mock.patch.object( - torch, "unique", side_effect=AssertionError("unique = host sync") - ): + with self.assertRaises(AssertionError): alloc.free_swa(v[1 : 5 * self.PS], start_pos=1) def test_start_pos_path_matches_the_fallback_end_state(self):