From 83a9b5dd8886fa972324304ce0cc8acfebc0d022 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Tue, 1 Sep 2026 14:16:47 -0700 Subject: [PATCH] [mem_cache] Drop the `torch.unique` sync from the SWA page expansion (#37463) --- python/sglang/srt/mem_cache/allocator/swa.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index 568e77ab6..785f9b83a 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -407,11 +407,21 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.free_full(torch.cat(full_free_group)) def _expand_to_full_pages(self, indices: torch.Tensor) -> torch.Tensor: - pages = torch.unique(indices // self.page_size) + # Duplicates are kept: deduplicating would be a torch.unique whose + # data-dependent output shape synchronizes the scheduler stream, and + # every consumer ends in the paged free's own page dedup anyway. + base = (indices // self.page_size) * self.page_size page_offsets = torch.arange( self.page_size, dtype=indices.dtype, device=indices.device ) - return (pages[:, None] * self.page_size + page_offsets[None, :]).reshape(-1) + expanded = (base[:, None] + page_offsets[None, :]).reshape(-1) + if self.swa_attn_allocator.debug_mode: + # Reference unique on CPU: the expansion must cover exactly the + # touched pages, on every caller's real input. + got = torch.unique(expanded.cpu() // self.page_size) + ref = torch.unique(indices.cpu() // self.page_size) + assert torch.equal(got, ref), "expansion page set mismatch" + return expanded def resize(self, config) -> None: size_full = int(config.full_max_total_num_tokens)