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

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

Browse Source
This commit is contained in:
Jialin Ouyang
2026-09-18 15:51:07 -07:00
committed by GitHub
parent 81a199f56a
commit f3851486cb
4 changed files with 382 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/memory/allocator.py
+100
View File
@@ -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
+53 -21
View File
@@ -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
test/registered/kernels/ops/memory/test_swa_page_free.py
+149
View File
@@ -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
+80
View File
@@ -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):