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