minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 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
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):