diff --git a/python/sglang/kernels/ops/memory/allocator.py b/python/sglang/kernels/ops/memory/allocator.py index 5c6b4b3e1..f0f4a77b6 100644 --- a/python/sglang/kernels/ops/memory/allocator.py +++ b/python/sglang/kernels/ops/memory/allocator.py @@ -1,7 +1,107 @@ +import torch import triton import triton.language as tl +@triton.jit(do_not_specialize=["num_pages", "full_page_representatives"]) +def get_and_clear_swa_pages_kernel( + full_page_representatives, + mapping, + swa_pages, + peers_mapped, + page_mappings_valid, + num_pages, + index_stride: tl.constexpr, + page_size: tl.constexpr, + CHECK_PAGE_MAPPINGS: tl.constexpr, + BLOCK_PAGES: tl.constexpr, + BLOCK_OFFSETS: tl.constexpr, +): + rep_offsets = tl.program_id(0) * BLOCK_PAGES + tl.arange(0, BLOCK_PAGES) + rep_mask = rep_offsets < num_pages + full_reps = tl.load( + full_page_representatives + rep_offsets * index_stride, + mask=rep_mask, + other=0, + ).to(tl.int64) + + # Resolve one SWA page per FULL-page representative: + # swa_reps = mapping[full_reps] + # swa_pages = swa_reps // page_size + swa_reps = tl.load(mapping + full_reps, mask=rep_mask, other=0) + tl.store(swa_pages + rep_offsets, swa_reps // page_size, mask=rep_mask) + tl.store(peers_mapped + rep_offsets, swa_reps > 0, mask=rep_mask) + + page_offsets = tl.arange(0, BLOCK_OFFSETS) + full_page_starts = full_reps // page_size * page_size + mapping_offsets = full_page_starts[:, None] + page_offsets[None, :] + mapping_mask = rep_mask[:, None] & (page_offsets[None, :] < page_size) + if CHECK_PAGE_MAPPINGS: + page_mapping = tl.load( + mapping + mapping_offsets, + mask=mapping_mask, + other=0, + ) + # Ignore zeros; each mapped slot must share its representative's SWA page. + same_swa_page = page_mapping // page_size == swa_reps[:, None] // page_size + page_mapping_valid = (swa_reps > 0) & ( + tl.sum(((page_mapping > 0) & ~same_swa_page).to(tl.int32), axis=1) == 0 + ) + tl.store( + page_mappings_valid + rep_offsets, + page_mapping_valid, + mask=rep_mask, + ) + + # `mapping_offsets` includes `full_reps`. Finish every `mapping[full_reps]` + # load before any warp clears `mapping[mapping_offsets]`. + tl.debug_barrier() + tl.store( + mapping + mapping_offsets, + 0, + mask=mapping_mask, + ) + + +def get_and_clear_swa_pages( + full_page_representatives: torch.Tensor, + mapping: torch.Tensor, + page_size: int, + check_page_mappings: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Resolve and clear mappings; input must represent distinct FULL pages.""" + if check_page_mappings: + assert torch.all( + (full_page_representatives >= 0) + & (full_page_representatives < mapping.numel() // page_size * page_size) + ), "FULL page representative out of bounds" + num_pages = full_page_representatives.numel() + swa_pages = torch.empty(num_pages, dtype=mapping.dtype, device=mapping.device) + peers_mapped = torch.empty(num_pages, dtype=torch.bool, device=mapping.device) + page_mappings_valid = ( + torch.empty(num_pages, dtype=torch.bool, device=mapping.device) + if check_page_mappings + else None + ) + if num_pages: + block_offsets = triton.next_power_of_2(page_size) + block_pages = max(1, 256 // block_offsets) + get_and_clear_swa_pages_kernel[(triton.cdiv(num_pages, block_pages),)]( + full_page_representatives, + mapping, + swa_pages, + peers_mapped, + page_mappings_valid if page_mappings_valid is not None else peers_mapped, + num_pages, + full_page_representatives.stride(0), + page_size, + check_page_mappings, + block_pages, + block_offsets, + ) + return swa_pages, peers_mapped, page_mappings_valid + + # free_page_ptr aliases self.free_pages, which the paged allocator re-slices # after every allocation (self.free_pages = self.free_pages[num_new_pages:]). # Slicing only advances data_ptr() by num_new_pages * 8 bytes, so the pointer diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index 7e573a752..22d7e9c58 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -2,6 +2,7 @@ import logging import torch +from sglang.kernels.ops.memory.allocator import get_and_clear_swa_pages from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.allocator.paged import PagedTokenToKVPoolAllocator from sglang.srt.mem_cache.allocator.token import TokenToKVPoolAllocator @@ -500,30 +501,15 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): return self._free_swa_pages(free_index, start_pos=start_pos) - def _free_swa_pages(self, free_index: torch.Tensor, *, start_pos: int): + def _free_swa_pages(self, free_index: torch.Tensor, start_pos: int): ps = self.page_size assert start_pos % ps == 0, f"segment start {start_pos} is not page-aligned" - # First token of every page the segment touches; the caller allocated - # each one, so a dead entry means the caller wanted free_full. - reps = free_index[::ps] - swa_tokens = self.full_to_swa_index_mapping[reps] - expect(_SWA_PEER_MAPPED, swa_tokens > 0, msg="caller wants free_full") - - if ps == 1: - swa_pages = swa_tokens - mapping_indices = free_index + full_page_representatives = free_index[::ps] + # torch_npu's transfer_to_npu aliases Tensor.is_cuda to Tensor.is_npu. + if not _is_npu and free_index.is_cuda: + swa_pages = self._free_swa_pages_cuda(full_page_representatives) else: - swa_pages = swa_tokens // ps - # Both pools page in step (alloc_extend / alloc_decode drive them - # with one seq_lens), so a rep's peer page is the whole peer page. - mapping_indices = self._expand_to_full_pages(reps) - if self.swa_attn_allocator.debug_mode: - ref = self.full_to_swa_index_mapping[mapping_indices].cpu() - assert torch.equal( - torch.sort(swa_pages.cpu())[0], - torch.unique(ref[ref > 0] // ps), - ), "swa pages do not match the mapped pages" - self.clear_full_to_swa_mapping(mapping_indices) + swa_pages = self._free_swa_pages_none_cuda(full_page_representatives) if self._swa_req_ring: # Ring slots are owned by the req slot, never lent by the paged @@ -538,6 +524,52 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.swa_attn_allocator.free_page_ids(swa_pages) assert self.swa_attn_allocator.available_size() <= self.swa_attn_allocator.size + def _free_swa_pages_none_cuda( + self, full_page_representatives: torch.Tensor + ) -> torch.Tensor: + ps = self.page_size + swa_tokens = self.full_to_swa_index_mapping[full_page_representatives] + expect(_SWA_PEER_MAPPED, swa_tokens > 0, msg="caller wants free_full") + + if ps == 1: + swa_pages = swa_tokens + mapping_indices = full_page_representatives + else: + swa_pages = swa_tokens // ps + # Both pools page in step (alloc_extend / alloc_decode drive them + # with one seq_lens), so a rep's peer page is the whole peer page. + mapping_indices = self._expand_to_full_pages(full_page_representatives) + if self.swa_attn_allocator.debug_mode: + ref = self.full_to_swa_index_mapping[mapping_indices].cpu() + assert torch.equal( + torch.sort(swa_pages.cpu())[0], + torch.unique(ref[ref > 0] // ps), + ), "swa pages do not match the mapped pages" + + self.clear_full_to_swa_mapping(mapping_indices) + return swa_pages + + def _free_swa_pages_cuda( + self, full_page_representatives: torch.Tensor + ) -> torch.Tensor: + ps = self.page_size + check_page_mappings = ps > 1 and self.swa_attn_allocator.debug_mode + swa_pages, peers_mapped, page_mappings_valid = get_and_clear_swa_pages( + full_page_representatives, + self.full_to_swa_index_mapping, + ps, + check_page_mappings=check_page_mappings, + ) + expect(_SWA_PEER_MAPPED, peers_mapped, msg="caller wants free_full") + if check_page_mappings: + assert page_mappings_valid is not None + # JIT checks within pages; sorting catches duplicate SWA pages. + sorted_swa_pages = torch.sort(swa_pages).values + assert torch.all(page_mappings_valid) & torch.all( + sorted_swa_pages[1:] != sorted_swa_pages[:-1] + ), "swa pages do not match the mapped pages" + return swa_pages + def _release_swa(self, swa_indices: torch.Tensor): if self.page_size > 1: # Set-shaped frees only (see free_swa): drop the padding-slot entries diff --git a/test/registered/kernels/ops/memory/test_swa_page_free.py b/test/registered/kernels/ops/memory/test_swa_page_free.py new file mode 100644 index 000000000..adfa96196 --- /dev/null +++ b/test/registered/kernels/ops/memory/test_swa_page_free.py @@ -0,0 +1,149 @@ +import itertools +import unittest + +import torch + +from sglang.kernels.ops.memory.allocator import get_and_clear_swa_pages +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large") +register_amd_ci(est_time=10, stage="jit-kernel-unit", runner_config="amd") + + +@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA or HIP") +class TestSwaPageFree(CustomTestCase): + def test_resolve_and_clear_matches_reference(self): + generator = torch.Generator().manual_seed(0) + for page_size, num_pages, dtype, stride in itertools.product( + (1, 4, 64, 128), + (0, 1, 3, 65, 257), + (torch.int32, torch.int64), + (1, 3), + ): + with self.subTest( + page_size=page_size, + num_pages=num_pages, + dtype=dtype, + stride=stride, + ): + pool_pages = 2 * num_pages + 3 + pages = ( + torch.randperm(pool_pages - 1, generator=generator)[:num_pages] + 1 + ) + representatives = pages * page_size + ( + torch.arange(num_pages) % page_size + ) + mapping_cpu = torch.arange(pool_pages * page_size) + 7 * page_size + if num_pages > 1: + mapping_cpu[representatives[num_pages // 2]] = 0 + + expected_peers = mapping_cpu[representatives] + expected_mapping = mapping_cpu.clone() + for page in pages.tolist(): + expected_mapping[page * page_size : (page + 1) * page_size] = 0 + + indices = torch.empty(num_pages * stride, dtype=dtype, device="cuda")[ + ::stride + ] + indices.copy_(representatives) + mapping = mapping_cpu.cuda() + swa_pages, peers_mapped, page_mappings_valid = get_and_clear_swa_pages( + indices, mapping, page_size + ) + + self.assertIsNone(page_mappings_valid) + self.assertTrue( + torch.equal(swa_pages.cpu(), expected_peers // page_size) + ) + self.assertTrue(torch.equal(peers_mapped.cpu(), expected_peers > 0)) + self.assertTrue(torch.equal(mapping.cpu(), expected_mapping)) + + def test_int32_mapping_last_page_before_sentinel(self): + for page_size, check_page_mappings in itertools.product((1, 16), (False, True)): + with self.subTest( + page_size=page_size, check_page_mappings=check_page_mappings + ): + mapping_cpu = torch.arange(5 * page_size + 1, dtype=torch.int32) + mapping_cpu[-1] = -1 + representative = mapping_cpu.numel() - 2 + expected_mapping = mapping_cpu.clone() + expected_mapping[4 * page_size : 5 * page_size] = 0 + mapping = mapping_cpu.cuda() + + swa_pages, peers_mapped, page_mappings_valid = get_and_clear_swa_pages( + torch.tensor([representative], dtype=torch.int32, device="cuda"), + mapping, + page_size, + check_page_mappings=check_page_mappings, + ) + + self.assertEqual(swa_pages.dtype, torch.int32) + self.assertEqual(swa_pages.item(), 4) + self.assertTrue(peers_mapped.item()) + if check_page_mappings: + self.assertTrue(page_mappings_valid.item()) + self.assertTrue(torch.equal(mapping.cpu(), expected_mapping)) + + def test_debug_rejects_out_of_bounds(self): + mapping = torch.arange(33, device="cuda") + expected_mapping = mapping.clone() + for representative in (-1, 32, 33): + with self.subTest(representative=representative): + with self.assertRaisesRegex( + AssertionError, "FULL page representative out of bounds" + ): + get_and_clear_swa_pages( + torch.tensor([representative], device="cuda"), mapping, 4, True + ) + self.assertTrue(torch.equal(mapping, expected_mapping)) + + def test_page_mapping_validation(self): + page_size = 4 + full_page = 3 + representative = full_page * page_size + 1 + swa_page = 7 + base_mapping = torch.zeros(12 * page_size, dtype=torch.int64) + base_mapping[full_page * page_size : (full_page + 1) * page_size] = ( + torch.arange(swa_page * page_size, (swa_page + 1) * page_size) + ) + base_mapping[full_page * page_size] = 0 + + mixed_peer = full_page * page_size + 3 + for name, updates, expected_page, expected_peer, expected_valid in ( + ("valid", (), swa_page, True, True), + ("missing_representative", ((representative, 0),), 0, False, False), + ( + "multiple_peer_pages", + ((mixed_peer, base_mapping[mixed_peer] + page_size),), + swa_page, + True, + False, + ), + ): + with self.subTest(name=name): + mapping = base_mapping.clone() + for index, value in updates: + mapping[index] = value + expected_mapping = mapping.clone() + expected_mapping[ + full_page * page_size : (full_page + 1) * page_size + ] = 0 + mapping = mapping.cuda() + + swa_pages, peers_mapped, page_mappings_valid = get_and_clear_swa_pages( + torch.tensor([representative], device="cuda"), + mapping, + page_size, + check_page_mappings=True, + ) + + self.assertEqual(swa_pages.item(), expected_page) + self.assertEqual(peers_mapped.item(), expected_peer) + self.assertIsNotNone(page_mappings_valid) + self.assertEqual(page_mappings_valid.item(), expected_valid) + self.assertTrue(torch.equal(mapping.cpu(), expected_mapping)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_swa_unittest.py b/test/registered/unit/mem_cache/test_swa_unittest.py index 6b0cca4a8..0adadf515 100644 --- a/test/registered/unit/mem_cache/test_swa_unittest.py +++ b/test/registered/unit/mem_cache/test_swa_unittest.py @@ -1367,6 +1367,86 @@ class TestSWAPageRepsFree(CustomTestCase): def _sizes(self, allocator): return allocator.full_available_size(), allocator.swa_available_size() + @unittest.skipUnless(torch.cuda.is_available(), "needs a tensor with is_cuda=True") + def test_free_swa_segment_npu_uses_reference_path(self): + for page_size in (1, 4): + with self.subTest(page_size=page_size): + _, allocator, _ = _build_swa_tree( + is_eagle=False, + page_size=page_size, + kv_size=8 * page_size, + kv_size_swa=8 * page_size, + ) + available_before = allocator.swa_available_size() + full_indices = _swa_alloc(allocator, page_size) + self.assertTrue(full_indices.is_cuda) + + # transfer_to_npu makes NPU tensors report is_cuda=True as well. + with ( + patch("sglang.srt.mem_cache.allocator.swa._is_npu", True), + patch( + "sglang.srt.mem_cache.allocator.swa.get_and_clear_swa_pages", + side_effect=AssertionError("NPU free reached Triton"), + ), + ): + allocator.free_swa_segment(full_indices[:1], start_pos=0) + + self.assertEqual(allocator.swa_available_size(), available_before) + self.assertTrue( + torch.all(allocator.full_to_swa_index_mapping[full_indices] == 0) + ) + + def test_free_swa_segment_debug_rejects_invalid_page_mappings(self): + page_size = 4 + + def leading_hole(mapping, full_indices, _swa_indices): + mapping[full_indices[0]] = 0 + + def multiple_peers(mapping, full_indices, swa_indices): + mapping[full_indices[2:page_size]] = swa_indices[ + page_size + 2 : 2 * page_size + ] + + def duplicate_peer(mapping, full_indices, swa_indices): + mapping[full_indices[page_size : 2 * page_size]] = swa_indices[:page_size] + + def duplicate_representative(_mapping, full_indices, _swa_indices): + full_indices[-page_size:] = full_indices[:page_size] + + for name, mutate, num_tokens in ( + ("leading_hole", leading_hole, page_size), + ("multiple_peers", multiple_peers, page_size), + ("duplicate_peer", duplicate_peer, 2 * page_size), + # At page size 4, representatives 0 and 64 belong to separate programs. + ("duplicate_representative", duplicate_representative, 65 * page_size), + ): + with self.subTest(name=name): + num_allocated_tokens = max(2 * page_size, num_tokens) + kv_size = max(8 * page_size, num_allocated_tokens) + _, allocator, _ = _build_swa_tree( + is_eagle=False, + page_size=page_size, + kv_size=kv_size, + kv_size_swa=kv_size, + ) + full_indices = _swa_alloc(allocator, num_allocated_tokens) + mapping = allocator.full_to_swa_index_mapping + swa_indices = mapping[full_indices].clone() + mutate(mapping, full_indices, swa_indices) + allocator.swa_attn_allocator.debug_mode = True + + # Exercise debug validation without CI's fatal async assertion. + with ( + patch.dict( + "os.environ", + {"SGLANG_INVARIANT_CHECK": str(int(InvariantCheckLevel.OFF))}, + ), + self.assertRaisesRegex( + AssertionError, "swa pages do not match the mapped pages" + ), + ): + allocator.free_swa_segment(full_indices[:num_tokens], start_pos=0) + def test_segment_free_releases_the_mapped_pages_for_every_tail(self): ps = self.PS for num_tokens in (1, ps, ps + 1, 3 * ps - 1, 3 * ps):