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

Size the unified read-table grid from bs, and fuse the allocator's tombstone scatters (#37511)

Browse Source 
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-09-02 16:54:26 -07:00
committed by GitHub
co-authored by Claude Opus 5
parent c05f8ae830
commit 18d5ffb42a
5 changed files with 244 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/kvcache/kv_read_table.py
+31 -17
View File
@@ -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
python/sglang/kernels/ops/memory/__init__.py
+2
View File
@@ -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(
python/sglang/kernels/ops/memory/virtual_slot.py
+95
View File
@@ -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)
python/sglang/srt/mem_cache/multi_ended_allocator.py
+27 -23
View File
@@ -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])
test/registered/unit/mem_cache/test_unified_free_no_host_sync.py
+89 -3
View File
@@ -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()