diff --git a/python/sglang/kernels/ops/kvcache/kv_read_table.py b/python/sglang/kernels/ops/kvcache/kv_read_table.py index 04e4b90ae..3995c8cb2 100644 --- a/python/sglang/kernels/ops/kvcache/kv_read_table.py +++ b/python/sglang/kernels/ops/kvcache/kv_read_table.py @@ -32,6 +32,11 @@ cuda-graph buffer be refreshed in place. Readers bound themselves by A `-1` in `req_to_token` and a freed (`-1`) v2p row both clamp to entry 0, the reserved padding slot, so a kernel dereferences padding, not a wild address. + +The grid is sized from `bs` alone and each program strides over the columns it +owns, bounded by the device-side `seq_lens`. A cuda-graph capture bakes the +grid, so a grid spanning `max_pages` would replay `max_context_len`/BLOCK column +blocks every step no matter how short the sequences actually are. """ from __future__ import annotations @@ -40,7 +45,11 @@ import torch import triton import triton.language as tl -_BLOCK_COLS = 256 +_BLOCK_COLS = 512 +_NUM_WARPS = 8 +# Enough blocks to fill the device without oversubscribing the column loop; +# measured on H100 over bs 1..256 x seq 1k..128k, flat within ~10% either side. +_TARGET_BLOCKS = 1024 @triton.jit @@ -53,28 +62,29 @@ def build_kv_read_table_kernel( req_stride, # runtime: req_to_token row stride (elements) out_stride, # runtime: out row stride (elements) mult, # runtime: kernel_page_multiplier of the target sub-pool + col_stride, # runtime: columns one program advances per loop trip PAGE_SIZE: tl.constexpr, BLOCK: tl.constexpr, ): bid = tl.program_id(0) - blk = tl.program_id(1) req = tl.load(req_pool_indices_ptr + bid).to(tl.int64) seqlen = tl.load(seq_lens_ptr + bid) n_pages = (seqlen + PAGE_SIZE - 1) // PAGE_SIZE + row_in = req_to_token_ptr + req * req_stride + row_out = out_ptr + bid.to(tl.int64) * out_stride - cols = blk * BLOCK + tl.arange(0, BLOCK) - mask = cols < n_pages - tok = tl.load( - req_to_token_ptr + req * req_stride + cols.to(tl.int64) * PAGE_SIZE, - mask=mask, - other=0, - ).to(tl.int64) - # Triton's `//` truncates toward zero, so `-1 // ps` is 0 for ps > 1 but - # -1 at ps == 1, which would read one element BEFORE `v2p`. - page = tl.where(tok < 0, 0, tok // PAGE_SIZE) - phys = tl.load(v2p_ptr + page, mask=mask, other=0) - entry = tl.maximum(phys * mult, 0).to(tl.int32) - tl.store(out_ptr + bid.to(tl.int64) * out_stride + cols, entry, mask=mask) + for start in range(tl.program_id(1) * BLOCK, n_pages, col_stride): + cols = start + tl.arange(0, BLOCK) + mask = cols < n_pages + tok = tl.load(row_in + cols.to(tl.int64) * PAGE_SIZE, mask=mask, other=0).to( + tl.int64 + ) + # Triton's `//` truncates toward zero, so `-1 // ps` is 0 for ps > 1 but + # -1 at ps == 1, which would read one element BEFORE `v2p`. + page = tl.where(tok < 0, 0, tok // PAGE_SIZE) + phys = tl.load(v2p_ptr + page, mask=mask, other=0) + entry = tl.maximum(phys * mult, 0).to(tl.int32) + tl.store(row_out + cols, entry, mask=mask) def build_kv_read_table( @@ -124,8 +134,10 @@ def build_kv_read_table( dst.copy_(torch.where(live, entry, dst)) return out - grid = (bs, triton.cdiv(max_pages, _BLOCK_COLS)) - build_kv_read_table_kernel[grid]( + col_programs = min( + triton.cdiv(_TARGET_BLOCKS, bs), triton.cdiv(max_pages, _BLOCK_COLS) + ) + build_kv_read_table_kernel[(bs, col_programs)]( req_to_token, req_pool_indices, seq_lens, @@ -134,7 +146,9 @@ def build_kv_read_table( req_to_token.stride(0), out.stride(0), multiplier, + col_programs * _BLOCK_COLS, PAGE_SIZE=page_size, BLOCK=_BLOCK_COLS, + num_warps=_NUM_WARPS, ) return out diff --git a/python/sglang/kernels/ops/memory/__init__.py b/python/sglang/kernels/ops/memory/__init__.py index 1e9ae9d1e..c40c8ba53 100644 --- a/python/sglang/kernels/ops/memory/__init__.py +++ b/python/sglang/kernels/ops/memory/__init__.py @@ -16,6 +16,8 @@ _TRITON_KERNELS = [ ("common", "get_last_loc_triton"), ("common", "get_last_loc_triton_safe"), ("virtual_slot", "alloc_bind_inplace"), + ("virtual_slot", "free_unbind_inplace"), + ("virtual_slot", "bind_inplace"), ] for _mod, _fn in _TRITON_KERNELS: register_kernel( diff --git a/python/sglang/kernels/ops/memory/virtual_slot.py b/python/sglang/kernels/ops/memory/virtual_slot.py index b54c2da97..c62fc1ac9 100644 --- a/python/sglang/kernels/ops/memory/virtual_slot.py +++ b/python/sglang/kernels/ops/memory/virtual_slot.py @@ -94,3 +94,98 @@ def alloc_bind_inplace( BLOCK=ALLOC_BIND_BLOCK, ) return phys_pages + + +@triton.jit +def free_unbind_inplace_kernel( + v_pages_ptr, # in: [N] int64 — virtual page ids being freed + v2p_ptr, # in/out: int64 — virtual_to_physical table + p2v_ptr, # in/out: int64 — physical_to_virtual table + out_phys_ptr, # out: [N] int64 — physical page ids released + N, # runtime: number of pages to free + BLOCK: tl.constexpr, +): + """Fused inverse of `alloc_bind_inplace_kernel`: v2p read + both tombstones. + + Each lane owns one virtual page, so the read-then-tombstone of `v2p[v]` has + no cross-lane dependency. That holds only because the caller's ids are + unique (`_free_lazy` dedups at ps>1 and takes uniqueness from its contract + at ps==1); duplicates would race on `p2v[p]`. + """ + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < N + + v = tl.load(v_pages_ptr + offs, mask=mask, other=0).to(tl.int64) + p = tl.load(v2p_ptr + v, mask=mask, other=0).to(tl.int64) + + tl.store(out_phys_ptr + offs, p, mask=mask) + tl.store(v2p_ptr + v, -1, mask=mask) + tl.store(p2v_ptr + p, -1, mask=mask) + + +@triton.jit +def bind_inplace_kernel( + v_pages_ptr, # in: [N] int64 — virtual page ids + p_pages_ptr, # in: [N] int64 — physical page ids to bind them to + v2p_ptr, # in/out: int64 + p2v_ptr, # in/out: int64 + N, # runtime: number of pages + BLOCK: tl.constexpr, +): + """`alloc_bind_inplace_kernel` for a caller-supplied physical range. + + The fast path generates an ascending range in-kernel; the hole-draining + slow path already holds the physical ids, so it passes them instead. + """ + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < N + + v = tl.load(v_pages_ptr + offs, mask=mask, other=0).to(tl.int64) + p = tl.load(p_pages_ptr + offs, mask=mask, other=0).to(tl.int64) + + tl.store(v2p_ptr + v, p, mask=mask) + tl.store(p2v_ptr + p, v, mask=mask) + + +def free_unbind_inplace( + v_pages: torch.Tensor, + v2p: torch.Tensor, + p2v: torch.Tensor, +) -> torch.Tensor: + """Tombstone `v_pages` in both tables and return the physical pages freed.""" + N = int(v_pages.numel()) + if N == 0: + return torch.empty(0, dtype=torch.int64, device=v_pages.device) + v = v_pages.to(torch.int64) + if not v_pages.is_cuda: + # Pure-torch CPU reference for the CUDA-only kernel. + phys_pages = v2p[v].clone() + v2p.index_fill_(0, v, -1) + p2v.index_fill_(0, phys_pages, -1) + return phys_pages + phys_pages = torch.empty(N, dtype=torch.int64, device=v_pages.device) + grid = (triton.cdiv(N, ALLOC_BIND_BLOCK),) + free_unbind_inplace_kernel[grid](v, v2p, p2v, phys_pages, N, BLOCK=ALLOC_BIND_BLOCK) + return phys_pages + + +def bind_inplace( + v_pages: torch.Tensor, + p_pages: torch.Tensor, + v2p: torch.Tensor, + p2v: torch.Tensor, +) -> None: + """Bind `v_pages` to `p_pages` in both tables.""" + N = int(v_pages.numel()) + if N == 0: + return + v = v_pages.to(torch.int64) + p = p_pages.to(torch.int64) + if not v_pages.is_cuda: + v2p[v] = p + p2v[p] = v + return + grid = (triton.cdiv(N, ALLOC_BIND_BLOCK),) + bind_inplace_kernel[grid](v, p, v2p, p2v, N, BLOCK=ALLOC_BIND_BLOCK) diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index 06e026af7..3d8ce76c6 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -40,7 +40,11 @@ from typing import ( import torch from torch.profiler import record_function -from sglang.kernels.ops.memory.virtual_slot import alloc_bind_inplace +from sglang.kernels.ops.memory.virtual_slot import ( + alloc_bind_inplace, + bind_inplace, + free_unbind_inplace, +) from sglang.srt.environ import envs from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.allocator.paged import ( @@ -909,8 +913,12 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): def bind(self, virtual_ids: torch.Tensor, physical_ids: torch.Tensor) -> None: """Bind page-granular virtual ids to physical ids.""" with record_function("MultiEndedAlloc.bind"): - self.virtual_to_physical[virtual_ids] = physical_ids - self.physical_to_virtual[physical_ids] = virtual_ids + bind_inplace( + virtual_ids, + physical_ids, + self.virtual_to_physical, + self.physical_to_virtual, + ) def bind_pages( self, virtual_pages: torch.Tensor, physical_pages: torch.Tensor @@ -1432,26 +1440,22 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): """ self._stats_n_free_lazy += 1 with record_function("MultiEndedAlloc._free_lazy"): - with record_function("MultiEndedAlloc._free_lazy.v2p_lookup"): - free_v_pages_raw = free_index.detach().to(torch.int64) - if pages is not None: - # `free_segment` already derived these by stride slicing. - free_v_pages = pages - elif self.page_size == 1: - free_v_pages = free_v_pages_raw - else: - free_v_pages = torch.unique(free_v_pages_raw // self.page_size) - freed_p_pages = self.virtual_to_physical[free_v_pages] - # Disjoint-element scatters — no barrier (a freed v has no live reader; - # per-element scatter writes are atomic). - # `index_fill_`, NOT `t[idx] = -1`: the scalar form makes torch - # materialise -1 as a CPU tensor and copy it H2D, and a pageable - # H2D copy is host-BLOCKING -- the scheduler parks behind the - # in-flight forward until the stream drains (~16 ms per free on an - # 8192-token prefill). `index_fill_` takes the scalar through the - # ATen Scalar overload: one device kernel, no host sync. - self.virtual_to_physical.index_fill_(0, free_v_pages, -1) - self.physical_to_virtual.index_fill_(0, freed_p_pages, -1) + free_v_pages_raw = free_index.detach().to(torch.int64) + if pages is not None: + # `free_segment` already derived these by stride slicing. + free_v_pages = pages + elif self.page_size == 1: + free_v_pages = free_v_pages_raw + else: + free_v_pages = torch.unique(free_v_pages_raw // self.page_size) + # One kernel for the v2p read and both tombstones. Disjoint-element + # scatters need no barrier (a freed v has no live reader), and the + # tombstone value never crosses the host -- the scalar `t[idx] = -1` + # form would materialise -1 on the CPU and block the scheduler on a + # pageable H2D copy (~16 ms per free on an 8192-token prefill). + freed_p_pages = free_unbind_inplace( + free_v_pages, self.virtual_to_physical, self.physical_to_virtual + ) if self.is_id_owner: self.free_virtual_ids = torch.cat([self.free_virtual_ids, free_v_pages]) self._free_phys_pages = torch.cat([self._free_phys_pages, freed_p_pages]) diff --git a/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py b/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py index 9e6245706..0bee23ad4 100644 --- a/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py +++ b/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py @@ -81,6 +81,12 @@ _TOMBSTONE_METHODS = [ ] +# Ways to write `-1` into a table without the value crossing the bus: +# `index_fill_` takes the scalar as an argument torch keeps off the host, and +# the fused launchers store it from inside the kernel. +_NO_SYNC_TOMBSTONE_FORMS = ("index_fill_", "free_unbind_inplace") + + def _allocators_in_module(): """Every allocator class DEFINED in multi_ended_allocator (not imported).""" return sorted( @@ -219,11 +225,19 @@ class TestTombstonesDoNotCrossTheBus(unittest.TestCase): ), ) - def test_free_paths_actually_use_index_fill(self): - """Positive form, so deleting the scatter entirely cannot pass.""" + def test_free_paths_actually_write_a_tombstone(self): + """Positive form, so deleting the scatter entirely cannot pass. The + mechanism is not the point -- keeping the tombstone value off the host + is -- so this lists the sanctioned ways to do that and a new one is + added here deliberately.""" for cls, name in _TOMBSTONE_METHODS: with self.subTest(method=f"{cls.__name__}.{name}"): - self.assertIn("index_fill_", inspect.getsource(getattr(cls, name))) + src = inspect.getsource(getattr(cls, name)) + self.assertTrue( + any(form in src for form in _NO_SYNC_TOMBSTONE_FORMS), + f"{cls.__name__}.{name} writes no tombstone through any of " + f"{_NO_SYNC_TOMBSTONE_FORMS}", + ) def test_index_fill_matches_scalar_assign_semantics(self): """Behaviour-preserving, including the edge cases the free path hands @@ -494,5 +508,77 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase): alloc.free_swa(v, start_pos=0) # all tombstoned -> filtered to empty +@unittest.skipUnless( + torch.cuda.is_available(), "the fused tombstone is a Triton kernel" +) +class TestFusedTombstoneWritesBothTables(unittest.TestCase): + """The source scan above accepts `free_unbind_inplace` as a no-sync + mechanism; this is what makes that acceptance mean something. On CPU the + launcher takes its pure-torch reference path, so nothing else in the suite + ever runs the kernel that does the tombstoning. + """ + + def test_matches_the_reference_over_randomized_bindings(self): + from sglang.kernels.ops.memory.virtual_slot import ( + bind_inplace, + free_unbind_inplace, + ) + + g = torch.Generator(device="cuda").manual_seed(11) + for trial in range(20): + n_pages = int(torch.randint(4, 400, (1,), generator=g, device="cuda")) + n_free = int( + torch.randint(1, n_pages + 1, (1,), generator=g, device="cuda") + ) + phys = torch.randperm(n_pages, device="cuda", generator=g).to(torch.int64) + virt = torch.randperm(n_pages, device="cuda", generator=g).to(torch.int64) + v2p = torch.full((n_pages,), -1, dtype=torch.int64, device="cuda") + p2v = torch.full((n_pages,), -1, dtype=torch.int64, device="cuda") + bind_inplace(virt, phys, v2p, p2v) + self.assertTrue(torch.equal(v2p[virt], phys), f"bind trial {trial}") + self.assertTrue(torch.equal(p2v[phys], virt), f"bind trial {trial}") + + freed_v = virt[:n_free] + want_p = v2p[freed_v].clone() + got_p = free_unbind_inplace(freed_v, v2p, p2v) + + self.assertTrue(torch.equal(got_p, want_p), f"freed pages, trial {trial}") + self.assertTrue( + torch.all(v2p[freed_v] == -1), f"v2p not tombstoned, trial {trial}" + ) + self.assertTrue( + torch.all(p2v[want_p] == -1), f"p2v not tombstoned, trial {trial}" + ) + live_v = virt[n_free:] + self.assertTrue( + torch.equal(v2p[live_v], phys[n_free:]), + f"a live binding was disturbed, trial {trial}", + ) + + def test_cuda_agrees_with_the_cpu_reference(self): + from sglang.kernels.ops.memory.virtual_slot import free_unbind_inplace + + v = torch.tensor([3, 0, 5], dtype=torch.int64) + cpu_v2p = torch.tensor([1, 2, 3, 4, 5, 0], dtype=torch.int64) + cpu_p2v = torch.tensor([5, 0, 1, 2, 3, 4], dtype=torch.int64) + cu_v2p, cu_p2v = cpu_v2p.cuda(), cpu_p2v.cuda() + cpu_out = free_unbind_inplace(v, cpu_v2p, cpu_p2v) + cu_out = free_unbind_inplace(v.cuda(), cu_v2p, cu_p2v) + self.assertTrue(torch.equal(cpu_out, cu_out.cpu())) + self.assertTrue(torch.equal(cpu_v2p, cu_v2p.cpu())) + self.assertTrue(torch.equal(cpu_p2v, cu_p2v.cpu())) + + def test_empty_free_is_a_noop(self): + from sglang.kernels.ops.memory.virtual_slot import free_unbind_inplace + + v2p = torch.arange(4, dtype=torch.int64, device="cuda") + before = v2p.clone() + out = free_unbind_inplace( + torch.empty(0, dtype=torch.int64, device="cuda"), v2p, v2p.clone() + ) + self.assertEqual(int(out.numel()), 0) + self.assertTrue(torch.equal(v2p, before)) + + if __name__ == "__main__": unittest.main()