[Perf] Fuse SWA page lookup and mapping clear (#38948)

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

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
+        full_page_representatives = free_index[::ps]
-        # each one, so a dead entry means the caller wanted free_full.
+        # torch_npu's transfer_to_npu aliases Tensor.is_cuda to Tensor.is_npu.
-        reps = free_index[::ps]
+        if not _is_npu and free_index.is_cuda:
-        swa_tokens = self.full_to_swa_index_mapping[reps]
+            swa_pages = self._free_swa_pages_cuda(full_page_representatives)
-        expect(_SWA_PEER_MAPPED, swa_tokens > 0, msg="caller wants free_full")
-
-        if ps == 1:
-            swa_pages = swa_tokens
-            mapping_indices = free_index
         else:
-            swa_pages = swa_tokens // ps
+            swa_pages = self._free_swa_pages_none_cuda(full_page_representatives)
-            # 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)
 
         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

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

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