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

[Perf] Unified memory: close the DCP decode gap on Blackwell (#37926)

Browse Source
This commit is contained in:
Cheng Wan
2026-09-07 01:10:44 -07:00
committed by GitHub
parent a8edafff7c
commit b5766336d4
21 changed files with 609 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/mamba/mamba_state_indices_triton.py
+15 -3
View File
@@ -14,6 +14,8 @@ MTP inter-phase seam:
This module fuses that chain into a single launch.
"""
from typing import Optional
import torch
import triton
import triton.language as tl
@@ -24,15 +26,20 @@ def _fused_replay_state_indices_kernel(
req_pool_indices_ptr, # (total_bs,) int64 — static replay buffer
mamba_map_ptr, # (req_pool_size,) int32 — req_index_to_mamba_index_mapping
out_ptr, # (total_bs,) int32 — state_indices_list[bs - 1]
v2p_ptr, # (num_slots + 1,) int64 — mamba virtual->physical, or unused
valid_bs,
total_bs,
BS_UPPER: tl.constexpr,
HAS_V2P: tl.constexpr,
):
offs = tl.arange(0, BS_UPPER)
in_range = offs < total_bs
valid = offs < valid_bs
req = tl.load(req_pool_indices_ptr + offs, mask=valid, other=0)
idx = tl.load(mamba_map_ptr + req, mask=valid, other=0)
if HAS_V2P:
# Must gather before the padding sentinel, as the reference chain does.
idx = tl.load(v2p_ptr + idx, mask=valid, other=0)
out_val = tl.where(valid, idx.to(tl.int32), -1)
tl.store(out_ptr + offs, out_val, mask=in_range)
# Preserve the reference chain's side effect: padded rows of the static
@@ -49,6 +56,7 @@ def fused_replay_state_indices(
out_state_indices: torch.Tensor,
valid_bs: int,
total_bs: int,
v2p: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Fill the captured replay state-indices buffer in one launch.
@@ -58,9 +66,11 @@ def fused_replay_state_indices(
the ``-1`` sentinel (mamba kernels skip ``state_idx < 0``) and their
``req_pool_indices`` entries are zeroed.
Callers must supply an identity virtual->physical mapping (the static
hybrid pool); the unified pool's allocator translate is not a flat table
gather and has to take the reference chain.
``v2p`` is the mamba virtual->physical slot table, for a pool whose slot
ids are virtual (the unified memory pool). Pass None when the mapping
already yields physical slots (the static hybrid pool). The unified
allocator runs the mamba sub-pool at page_size 1, so its translate is a
plain table gather and folds into this launch.
Returns the filled ``out_state_indices[:total_bs]`` view.
"""
@@ -68,8 +78,10 @@ def fused_replay_state_indices(
req_pool_indices,
mamba_index_mapping,
out_state_indices,
v2p,
valid_bs,
total_bs,
BS_UPPER=triton.next_power_of_2(total_bs),
HAS_V2P=v2p is not None,
)
return out_state_indices[:total_bs]
python/sglang/kernels/ops/memory/__init__.py
+1
View File
@@ -18,6 +18,7 @@ _TRITON_KERNELS = [
("virtual_slot", "alloc_bind_inplace"),
("virtual_slot", "free_unbind_inplace"),
("virtual_slot", "bind_inplace"),
("virtual_slot", "write_loc_to_kernel_ids"),
]
for _mod, _fn in _TRITON_KERNELS:
register_kernel(
python/sglang/kernels/ops/memory/virtual_slot.py
+133
View File
@@ -15,6 +15,8 @@
from __future__ import annotations
from typing import Optional
import torch
import triton
import triton.language as tl
@@ -189,3 +191,134 @@ def bind_inplace(
return
grid = (triton.cdiv(N, ALLOC_BIND_BLOCK),)
bind_inplace_kernel[grid](v, p, v2p, p2v, N, BLOCK=ALLOC_BIND_BLOCK)
WRITE_LOC_BLOCK = 512
@triton.jit
def write_loc_to_kernel_id_kernel(
loc_ptr, # in: [N] int64 — WIDENED virtual token ids
v2p_ptr, # in: [num_pages + 1] int64 — virtual->physical page table
out_ptr, # out: [N] int64 — kernel-facing ids
N, # runtime: live element count
W, # runtime: lanes to write; [N, W) get 0
stride, # runtime: pool_page_size * kernel_page_multiplier
PAGE_SIZE: tl.constexpr,
DCP_SIZE: tl.constexpr,
DCP_RANK: tl.constexpr,
BLOCK: tl.constexpr,
):
"""``kernel_id(t) = v2p[t // ps] * ps * mult + t % ps``, clamped at 0.
Under DCP the incoming id is WIDENED: ``loc % dcp_size`` names its owner
and ``loc // dcp_size`` is the row. Ids this rank does not own resolve to
kernel id 0, the padding sink every write kernel skips.
Triton truncates ``//`` toward zero where torch floors it, so a negative
loc is tested explicitly rather than left to the division; it resolves to
0, as the torch path does.
Writing ``W > N`` lanes fills ``[N, W)`` with 0, the padding sink, so a
caller may hand in a capture-stable buffer wider than this batch and have
the stale tail cleared in the same launch.
"""
pid = tl.program_id(0)
offs = pid * BLOCK + tl.arange(0, BLOCK)
in_range = offs < W
mask = offs < N
loc = tl.load(loc_ptr + offs, mask=mask, other=0).to(tl.int64)
keep = mask & (loc >= 0)
if DCP_SIZE > 1:
keep = keep & ((loc % DCP_SIZE) == DCP_RANK)
loc = loc // DCP_SIZE
page = loc // PAGE_SIZE if PAGE_SIZE > 1 else loc
offset = loc % PAGE_SIZE if PAGE_SIZE > 1 else 0
# `keep` already excludes negatives, so the gather index is in range.
phys = tl.load(v2p_ptr + tl.where(keep, page, 0), mask=mask, other=0).to(tl.int64)
ids = tl.maximum(phys * stride + offset, 0)
tl.store(out_ptr + offs, tl.where(keep, ids, 0), mask=in_range)
def write_loc_to_kernel_ids(
*,
loc: torch.Tensor,
v2p: torch.Tensor,
page_size: int,
stride: int,
dcp_size: int = 1,
dcp_rank: int = 0,
out: Optional[torch.Tensor] = None,
out_width: Optional[int] = None,
) -> torch.Tensor:
"""One launch for the whole write-loc conversion; see the kernel.
``out`` is written in place when given (a captured graph records the
gather against a fixed ``data_ptr``), else a fresh int64 tensor is
returned. Cuda-graph safe: no ``.item()``, no host sync, no allocation on
the ``out=`` path.
``out_width`` writes that many lanes rather than ``loc.numel()``, zeroing
the ones past the batch; pass the captured tier's width to clear a stale
tail here.
"""
N = int(loc.numel())
# Flat-indexed as `ptr + offs`, so a strided view is mis-addressed.
assert loc.is_contiguous(), (
f"write_loc_to_kernel_ids: loc must be contiguous, got shape "
f"{tuple(loc.shape)} stride {tuple(loc.stride())}"
)
if out is None:
out = torch.empty_like(loc, dtype=torch.int64)
width = N if out_width is None else int(out_width)
assert out.dtype == torch.int64, (
f"write_loc_to_kernel_ids: out dtype must be int64 (matches v2p), "
f"got {out.dtype}"
)
if out_width is None:
# `out` mirrors `loc` whatever its shape; a 2-D page table is legal.
assert out.shape == loc.shape and out.is_contiguous(), (
f"write_loc_to_kernel_ids: out shape {tuple(out.shape)} must match "
f"loc shape {tuple(loc.shape)}"
)
else:
assert out.dim() == 1 and out.is_contiguous() and out.numel() >= width, (
f"write_loc_to_kernel_ids: out_width needs a packed 1-D out of at "
f"least {width}, got {tuple(out.shape)}"
)
assert width >= N, (
f"write_loc_to_kernel_ids: out_width {width} is under the batch's "
f"{N} locs, which would drop live rows"
)
if width == 0:
return out
if not loc.is_cuda:
# Pure-torch reference; the allocator's unit tests run on CPU.
big = loc.to(torch.int64)
keep = big >= 0
if dcp_size > 1:
keep = keep & (big % dcp_size == dcp_rank)
big = torch.div(big, dcp_size, rounding_mode="floor")
page = torch.where(keep, torch.div(big, page_size, rounding_mode="floor"), 0)
offset = big % page_size if page_size > 1 else 0
ids = (v2p[page] * stride + offset).clamp_(min=0)
out[:N].copy_(torch.where(keep, ids, torch.zeros_like(ids)))
if width > N:
out[N:width].zero_()
return out
grid = (triton.cdiv(width, WRITE_LOC_BLOCK),)
write_loc_to_kernel_id_kernel[grid](
loc,
v2p,
out,
N,
width,
stride,
PAGE_SIZE=page_size,
DCP_SIZE=dcp_size,
DCP_RANK=dcp_rank,
BLOCK=WRITE_LOC_BLOCK,
)
return out