minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[mem_cache] Require page-aligned starts in free_segment and drop the boundary trim (#37729)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-09-03 13:28:33 -07:00
committed by GitHub
parent 3ffacf949b
commit 2a980cbf10
7 changed files with 113 additions and 168 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py
+3 -20
View File
@@ -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()
python/sglang/srt/mem_cache/allocator/base.py
+19 -12
View File
@@ -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)
python/sglang/srt/mem_cache/allocator/paged.py
+12 -15
View File
@@ -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()
python/sglang/srt/mem_cache/multi_ended_allocator.py
+17 -31
View File
@@ -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))
test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py
+15 -49
View File
@@ -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)
test/registered/unit/mem_cache/test_paged_free_segment.py
+33 -27
View File
@@ -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()
test/registered/unit/mem_cache/test_unified_free_no_host_sync.py
+14 -14
View File
@@ -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):