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
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):