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
python/sglang/srt/arg_groups/kv_cache_hook.py
+6
View File
@@ -249,6 +249,12 @@ def handle_unified_memory_pool(server_args: Any) -> None:
"not translate speculative verify indices to the unified "
"pool's kernel-facing space yet."
)
assert not cfg.enable_two_batch_overlap, (
"--enable-unified-memory does not support --enable-two-batch-overlap: "
"TBO's replay split hands each child a view without the pre-translate "
"write loc, so a captured decode replay raises. "
"TODO(ch-wan): carry out_cache_loc_virtual into the child view."
)
assert not (cfg.enable_hierarchical_cache or cfg.enable_lmcache), (
"--enable-unified-memory is not yet compatible with hierarchical / "
"host-tiered KV cache (--enable-hierarchical-cache / --enable-lmcache): "
python/sglang/srt/batch_overlap/two_batch_overlap.py
+5
View File
@@ -693,6 +693,11 @@ class TboForwardBatchPreparer:
)
output_dict[key] = old_value[start_token_index:end_token_index]
if batch.out_cache_loc_virtual is not None:
output_dict["out_cache_loc_virtual"] = batch.out_cache_loc_virtual[
start_token_index:end_token_index
]
attention_tp_size = get_parallel().attn_tp_size
_tbo_padded_len = (
(end_token_index - start_token_index - 1) // attention_tp_size + 1
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
+5 -8
View File
@@ -70,15 +70,11 @@ class MambaAttnBackendBase(AttentionBackend):
# backends on runners without a real model_config.
self._model_runner = model_runner
self._mamba_chunk_size: Optional[int] = None
# Fused replay-prep state-indices fast path (fused_replay_state_indices):
# requires the static hybrid pool whose v2p translate is the identity —
# the unified pool overrides translate_mamba_indices with an allocator
# lookup that is not a flat table gather.
pool = self.req_to_token_pool
self._fused_state_indices_ok = (
str(self.device).startswith("cuda")
and isinstance(self.req_to_token_pool, HybridReqToTokenPool)
and type(self.req_to_token_pool).translate_mamba_indices
is HybridReqToTokenPool.translate_mamba_indices
torch.device(self.device).type == "cuda"
and isinstance(pool, HybridReqToTokenPool)
and pool.mamba_translate_is_fusable
)
self.forward_metadata: ForwardMetadata = None
self.state_indices_list = []
@@ -643,6 +639,7 @@ class MambaAttnBackendBase(AttentionBackend):
out_state_indices=self.state_indices_list[bs - 1],
valid_bs=bs - int(num_padding),
total_bs=bs,
v2p=self.req_to_token_pool.mamba_v2p_table,
)
else:
# Make sure forward metadata is correctly handled for padding reqs
python/sglang/srt/layers/attention/triton_backend.py
+8 -14
View File
@@ -748,22 +748,16 @@ class TritonAttnBackend(AttentionBackend):
def _fill_cuda_graph_write_locs(
self, forward_batch: ForwardBatch, bs: int
) -> Optional[torch.Tensor]:
"""Copy the cuda-graph WRITE loc into the capture-stable buffer and
return the ``[:n]`` view; no-op for non-unified pools.
Runs BEFORE graph.replay() so it reads the live post-compaction v2p.
The capture batch is runner-built with zeros, which is safe because
slot 0 is the reserved sink in every id space.
"""
"""Runs BEFORE graph.replay(), so it reads the live post-compaction
v2p; no-op for non-unified pools."""
# The buffer exists only for a translating pool; return before naming it.
if not self.kv_index_translator.is_translating:
return None
out_cache_loc = forward_batch.out_cache_loc
n = out_cache_loc.shape[0]
# Zero the padded tail first: a smaller replay batch leaves [n:] holding
# stale ids that the captured store would write; send them to slot 0 (sink).
self.cuda_graph_out_cache_loc_full_physical[n:].zero_()
self.cuda_graph_out_cache_loc_full_physical[:n].copy_(out_cache_loc)
return self.cuda_graph_out_cache_loc_full_physical[:n]
return self.kv_index_translator.fill_capture_write_loc(
out=self.cuda_graph_out_cache_loc_full_physical,
forward_batch=forward_batch,
width=self.cuda_graph_out_cache_loc_full_physical.numel(),
)
def init_forward_metadata(self, forward_batch: ForwardBatch):
"""Init auxiliary variables for triton attention backend."""
python/sglang/srt/layers/attention/trtllm_mla_backend.py
+7 -13
View File
@@ -766,19 +766,13 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
if self.kv_index_translator.is_translating and (
forward_mode.is_decode_or_idle() or forward_mode.is_target_verify()
):
out_cache_loc = forward_batch.out_cache_loc
n = out_cache_loc.shape[0]
dst = self.cuda_graph_out_cache_loc_kernel[:n]
dst.copy_(out_cache_loc)
# Replay-prep receives the RAW (unpadded) out_cache_loc
# (build_replay_fb_view), but the captured write kernel consumes the
# full captured tier of this buffer. Zero the tail so pad rows write
# to the sink (row 0) instead of stale kernel-facing locs left by
# earlier larger replays — a stale tail scatters pad-row garbage into
# live KV pages. Mirrors the runner's PaddingPolicy.ZERO on its own
# out_cache_loc slot.
self.cuda_graph_out_cache_loc_kernel[n:].zero_()
self._decode_kernel_loc = dst
# The captured kernel consumes the whole buffer, so the tail a
# shorter replay leaves must go to slot 0 rather than live pages.
self._decode_kernel_loc = self.kv_index_translator.fill_capture_write_loc(
out=self.cuda_graph_out_cache_loc_kernel,
forward_batch=forward_batch,
width=self.cuda_graph_out_cache_loc_kernel.numel(),
)
else:
self._decode_kernel_loc = None
python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py
+4 -1
View File
@@ -331,10 +331,13 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
loc: torch.Tensor,
*,
out: Optional[torch.Tensor] = None,
out_width: Optional[int] = None,
) -> torch.Tensor:
"""Widened virtual WRITE loc -> kernel-facing id. DCP is rejected for this
composite at argument validation, so it coincides with the read translate."""
return self.full_attn_allocator.translate_write_loc_for_kernel(loc, out=out)
return self.full_attn_allocator.translate_write_loc_for_kernel(
loc, out=out, out_width=out_width
)
@property
def swa_kernel_page_multiplier(self) -> int:
python/sglang/srt/mem_cache/allocator/unified_mamba.py
+4 -1
View File
@@ -260,9 +260,12 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
loc: torch.Tensor,
*,
out: Optional[torch.Tensor] = None,
out_width: Optional[int] = None,
) -> torch.Tensor:
"""Widened virtual WRITE loc -> DENSE id; see the sub-allocator's copy."""
return self.full_attn_allocator.translate_write_loc_for_kernel(loc, out=out)
return self.full_attn_allocator.translate_write_loc_for_kernel(
loc, out=out, out_width=out_width
)
def translate_kv_indices_for_transfer(
self, kv_indices: torch.Tensor
python/sglang/srt/mem_cache/allocator/unified_sub_pool.py
+36 -29
View File
@@ -43,6 +43,7 @@ from sglang.kernels.ops.memory.virtual_slot import (
alloc_bind_inplace,
bind_inplace,
free_unbind_inplace,
write_loc_to_kernel_ids,
)
from sglang.srt.environ import envs
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
@@ -990,39 +991,47 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
clamp to kernel-facing id 0, the page-0 sink. int64 out; a consumer whose
kernel ABI wants int32 narrows where it fills that buffer.
"""
ps = self.pool_page_size
stride = ps * self.kernel_page_multiplier
with record_function("MultiEndedAlloc.translate_kv_loc_for_kernel"):
pages = virt_tokens if ps == 1 else virt_tokens // ps
offsets = None if ps == 1 else virt_tokens % ps
if out is None:
phys = self.virtual_to_physical[pages]
ids = phys * stride if offsets is None else phys * stride + offsets
return ids.clamp_(min=0)
return self._translate_loc_fused(virt_tokens, dcp_size=1, out=out)
def _translate_loc_fused(
self,
loc: torch.Tensor,
*,
dcp_size: int,
dcp_rank: int = 0,
out: Optional[torch.Tensor] = None,
out_width: Optional[int] = None,
) -> torch.Tensor:
"""One launch for the read and write conversions alike; see
`write_loc_to_kernel_ids`."""
if out is not None:
assert out.dtype == torch.int64, (
f"translate_kv_loc_for_kernel: out= dtype must be int64 (matches v2p), "
f"got {out.dtype}"
)
assert out.shape == virt_tokens.shape, (
if out_width is None:
assert out.shape == loc.shape, (
f"translate_kv_loc_for_kernel: out= shape {tuple(out.shape)} must "
f"match virt_tokens shape {tuple(virt_tokens.shape)}"
f"match virt_tokens shape {tuple(loc.shape)}"
)
return write_loc_to_kernel_ids(
loc=loc,
v2p=self.virtual_to_physical,
page_size=self.pool_page_size,
stride=self.pool_page_size * self.kernel_page_multiplier,
dcp_size=dcp_size,
dcp_rank=dcp_rank,
out=out,
out_width=out_width,
)
if pages.dtype != torch.int64:
pages = pages.to(torch.int64)
if pages is virt_tokens:
out.copy_(torch.take(self.virtual_to_physical, pages))
else:
torch.take(self.virtual_to_physical, pages, out=out)
out.mul_(stride)
if offsets is not None:
out.add_(offsets)
return out.clamp_(min=0)
def translate_write_loc_for_kernel(
self,
widened_loc: torch.Tensor,
*,
out: Optional[torch.Tensor] = None,
out_width: Optional[int] = None,
) -> torch.Tensor:
"""Widened virtual WRITE loc (`out_cache_loc`) -> kernel-facing id.
@@ -1032,16 +1041,14 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
"""
parallel = get_parallel()
dcp_size = parallel.attn_dcp_size if self.shards_under_dcp else 1
if dcp_size == 1:
return self.translate_kv_loc_for_kernel(widened_loc, out=out)
with record_function("MultiEndedAlloc.translate_write_loc_for_kernel"):
owned = (widened_loc % dcp_size) == parallel.attn_dcp_rank
dense = self.translate_kv_loc_for_kernel(widened_loc // dcp_size)
dense = torch.where(owned, dense, torch.zeros_like(dense))
if out is not None:
out.copy_(dense)
return out
return dense
return self._translate_loc_fused(
widened_loc,
dcp_size=dcp_size,
dcp_rank=parallel.attn_dcp_rank,
out=out,
out_width=out_width,
)
# -- alloc --
python/sglang/srt/mem_cache/kv_index_translator.py
+38 -3
View File
@@ -432,20 +432,55 @@ class KVIndexTranslator:
def rebind_write_loc(self, forward_batch) -> None:
"""Phase 1 of the WRITE contract: translate the batch's write loc to
FULL-side kernel-facing ids exactly once, at ForwardBatch
construction. No-op on non-unified pools.
FULL-side kernel-facing ids, once, at ForwardBatch construction.
REBIND, never mutate: the translate returns a FRESH tensor, so the
ScheduleBatch's aliased tensor stays VIRTUAL for the radix / accept /
in-flight machinery that reads it.
in-flight machinery that reads it. The pre-translate tensor stays on
the batch for `fill_capture_write_loc`.
"""
self._index_table_memo = None
if not self.is_translating or forward_batch.out_cache_loc is None:
return
forward_batch.out_cache_loc_virtual = forward_batch.out_cache_loc
forward_batch.out_cache_loc = self._translate_write_full(
forward_batch.out_cache_loc
)
def fill_capture_write_loc(
self,
*,
out: torch.Tensor,
forward_batch,
width: Optional[int] = None,
) -> Optional[torch.Tensor]:
"""Translate this batch's WRITE loc straight into ``out``, a backend's
capture-stable buffer, and return the live ``[:n]`` view. One launch
fills the live prefix and clears the tail a shorter replay leaves;
None when this pool needs no translation.
Must run at metadata-init time: `out` is reused every step, so filling
it sooner would race a still-pending previous step under overlap
scheduling.
"""
if not self.is_translating:
return None
virtual = forward_batch.out_cache_loc_virtual
if virtual is None:
loc = forward_batch.out_cache_loc
if loc is None:
return None
# The runner builds the capture batch outside `init_new`, so no
# rebind marked its virtual source; bake it holding sink ids.
width = int(loc.numel()) if width is None else int(width)
out[:width].zero_()
return out[: int(loc.numel())]
n = int(virtual.numel())
width = n if width is None else int(width)
buf = out[:width]
self._translate_write_full(virtual, out=buf, out_width=width)
return buf[:n]
def sliding_window_write_loc_for(
self, out_cache_loc: Optional[torch.Tensor]
) -> Optional[torch.Tensor]:
python/sglang/srt/mem_cache/memory_pool.py
+22
View File
@@ -1385,6 +1385,28 @@ class HybridReqToTokenPool(ReqToTokenPool):
def get_mamba_indices(self, req_indices: torch.Tensor) -> torch.Tensor:
return self.req_index_to_mamba_index_mapping[req_indices]
@property
def mamba_v2p_table(self) -> Optional[torch.Tensor]:
"""The mamba virtual->physical slot table, or None when the ids this
pool hands out are already physical."""
return None
@property
def mamba_translate_is_fusable(self) -> bool:
"""Whether `fused_replay_state_indices` can reproduce this pool's
`translate_mamba_indices` in its own launch.
The kernel expresses exactly two shapes: the identity, and one gather
through `mamba_v2p_table`. A subclass that replaces the translate with
anything else is excluded here rather than silently mis-served.
"""
if self.mamba_v2p_table is not None:
return True
return (
type(self).translate_mamba_indices
is HybridReqToTokenPool.translate_mamba_indices
)
def translate_mamba_indices(self, mamba_indices: torch.Tensor) -> torch.Tensor:
"""Virtual->physical mamba-slot translate. Identity for a static pool
(slots are physical); UnifiedHybridReqToTokenPool overrides it for the
python/sglang/srt/mem_cache/unified_memory_pool.py
+9
View File
@@ -1107,6 +1107,15 @@ class UnifiedHybridReqToTokenPool(HybridReqToTokenPool):
)
)
@property
def mamba_v2p_table(self) -> Optional[torch.Tensor]:
"""This pool's ids ARE virtual; page_size is 1, so the translate is the
plain gather this table serves, which keeps `mamba_translate_is_fusable`
true despite the override."""
if self.mamba_allocator is None:
return None
return self.mamba_allocator.virtual_to_physical
def translate_mamba_indices(self, virtual_ids: torch.Tensor) -> torch.Tensor:
"""Virtual mamba ids -> physical slot ids."""
return self.mamba_allocator.translate(virtual_ids).to(torch.int32)
python/sglang/srt/model_executor/forward_batch_info.py
+5
View File
@@ -414,6 +414,9 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
# The original sequence length without being chunked. Qwen-1M related.
orig_seq_lens: Optional[torch.Tensor] = None
# The write loc before `rebind_write_loc` replaced it with kernel-facing
# ids; a backend re-derives from it into its capture-stable buffer.
out_cache_loc_virtual: Optional[torch.Tensor] = None
# DSV4-NPU only: per-pool slot bundle from DSV4NPUTokenToKVPoolAllocator,
# consumed by the Ascend backend for PA_ND block tables. None elsewhere.
out_cache_loc_dsv4: Optional[DSV4OutCacheLoc] = None
@@ -1836,6 +1839,8 @@ def build_inner_fb_view(
seq_lens_cpu=forward_batch.seq_lens_cpu,
encoder_lens=encoder_lens,
out_cache_loc=getattr(forward_batch, "out_cache_loc", None),
# A caller may hand in another view that does not carry this field.
out_cache_loc_virtual=getattr(forward_batch, "out_cache_loc_virtual", None),
out_cache_loc_dsv4=getattr(forward_batch, "out_cache_loc_dsv4", None),
spec_info=forward_batch.spec_info,
)
python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
+1
View File
@@ -192,6 +192,7 @@ def build_replay_fb_view(
num_padding=bs - raw_bs,
encoder_lens=buffers.encoder_lens[:bs] if is_encoder_decoder else None,
out_cache_loc=getattr(forward_batch, "out_cache_loc", None),
out_cache_loc_virtual=forward_batch.out_cache_loc_virtual,
out_cache_loc_dsv4=getattr(forward_batch, "out_cache_loc_dsv4", None),
# The mamba-track registry slot (VIRTUAL ids) is the v2p translate SOURCE
# for the backend, which copies the result into its own static buffer and
test/registered/dcp/test_trtllm_mla_family_dcp_metadata.py
+83
View File
@@ -432,6 +432,89 @@ class TestFusedFp8WriteGate(CustomTestCase):
)
@unittest.skipUnless(torch.cuda.is_available(), "the fused translate is Triton")
class TestFusedWriteLocTranslateCuda(CustomTestCase):
"""The fused write-loc translate must agree with its own CPU branch.
The two implementations must not drift: Triton truncates division toward
zero where torch floors it, so a negative loc and a tombstoned v2p row are
where a divergence would appear -- and the CPU branch is all the CPU suites
ever exercise.
"""
def test_cuda_matches_the_cpu_branch(self):
from sglang.kernels.ops.memory.virtual_slot import write_loc_to_kernel_ids
for page_size in (1, 64):
span = page_size * 4
loc = torch.tensor(
[-1, 0, 1, page_size, span, span + 1, 2 * span + 3, 5 * page_size],
dtype=torch.int64,
)
# Size the table past the highest page any loc can name, then
# scramble it and tombstone one row (-1), so neither a dropped
# clamp nor a skipped gather can coincide with the right answer.
num_pages = int(loc.max()) // page_size + 2
v2p = torch.tensor(
[(5 * i + 2) % num_pages for i in range(num_pages)] + [-1],
dtype=torch.int64,
)
v2p[1] = -1
for dcp_size, dcp_rank in ((1, 0), (2, 1), (4, 2)):
kw = dict(
page_size=page_size,
stride=page_size * 3,
dcp_size=dcp_size,
dcp_rank=dcp_rank,
)
cpu = write_loc_to_kernel_ids(loc=loc, v2p=v2p, **kw)
gpu = write_loc_to_kernel_ids(loc=loc.cuda(), v2p=v2p.cuda(), **kw)
self.assertEqual(
cpu.tolist(),
gpu.cpu().tolist(),
f"ps={page_size} dcp_size={dcp_size} rank={dcp_rank}",
)
def test_wide_out_clears_the_stale_tail(self):
"""`out_width` past the batch must zero the tail in the same launch.
This is what lets a backend hand in its whole capture-stable buffer:
a shorter replay leaves stale kernel-facing ids past the batch, and the
captured write kernel consumes the full buffer, so an uncleared tail
scatters pad rows into live KV pages.
"""
from sglang.kernels.ops.memory.virtual_slot import write_loc_to_kernel_ids
v2p = torch.tensor([2, 5, 1, 3, 4], dtype=torch.int64, device="cuda")
loc = torch.tensor([0, 64, 128], dtype=torch.int64, device="cuda")
width = 8
# Poison the whole buffer so an unwritten or uncleared cell is visible.
buf = torch.full((width,), -999, dtype=torch.int64, device="cuda")
write_loc_to_kernel_ids(
loc=loc, v2p=v2p, page_size=64, stride=64 * 2, out=buf, out_width=width
)
self.assertEqual(buf[:3].tolist(), [2 * 128, 5 * 128, 1 * 128])
self.assertEqual(buf[3:].tolist(), [0] * (width - 3))
# And it must agree with the narrow call on the live prefix.
narrow = write_loc_to_kernel_ids(loc=loc, v2p=v2p, page_size=64, stride=64 * 2)
self.assertEqual(narrow.tolist(), buf[:3].tolist())
def test_out_is_written_in_place(self):
# The captured decode path hands in a capture-stable buffer; rebinding
# instead of filling it would leave the graph on a stale pointer.
from sglang.kernels.ops.memory.virtual_slot import write_loc_to_kernel_ids
v2p = torch.tensor([2, 5, -1, 3], dtype=torch.int64, device="cuda")
loc = torch.tensor([0, 64, 128, 192], dtype=torch.int64, device="cuda")
dst = torch.full_like(loc, -7)
ret = write_loc_to_kernel_ids(
loc=loc, v2p=v2p, page_size=64, stride=64 * 2, out=dst
)
self.assertIs(ret, dst)
self.assertEqual(dst.tolist(), [2 * 128, 5 * 128, 0, 3 * 128])
class TestDcpDecodeLayout(CustomTestCase):
"""Rank-local length math the decode page table above is built from."""
test/registered/kernels/ops/mamba/test_fused_replay_state_indices.py
+78 -1
View File
@@ -49,11 +49,15 @@ def _reference_chain(
out_buf: torch.Tensor,
valid_bs: int,
total_bs: int,
v2p: torch.Tensor = None,
) -> None:
"""Replicates the _replay_metadata reference ops, in order, in place."""
req_pool_indices[valid_bs:total_bs] = 0
mamba_indices = mapping[req_pool_indices[:total_bs]]
# static pool: _translate_mamba_indices is the identity
if v2p is not None:
# Unified pool: virtual->physical slot translate, and it runs BEFORE
# the padding sentinel -- captured kernels read physical ids.
mamba_indices = v2p[mamba_indices].to(torch.int32)
mamba_indices[valid_bs:] = -1
out_buf[: len(mamba_indices)].copy_(mamba_indices)
@@ -139,6 +143,79 @@ class TestFusedReplayStateIndices(CustomTestCase):
f"{name} guard tail clobbered ({case}): {buf[total_bs:].tolist()}",
)
def _run_v2p_case(self, total_bs: int, num_padding: int, seed: int) -> None:
"""Same equivalence, with the unified pool's virtual slot ids.
The mapping yields VIRTUAL slots there and the kernel folds the v2p
gather in, so the two must still agree element for element.
"""
device = torch.device("cuda")
gen = torch.Generator(device="cpu").manual_seed(seed + 1000)
valid_bs = total_bs - num_padding
req_pool = torch.randint(
0, _REQ_POOL_SIZE, (total_bs + _GUARD,), generator=gen, dtype=torch.int64
)
req_pool[total_bs:] = _GUARD_SENTINEL
mapping = torch.randint(
0, _MAMBA_POOL_SIZE, (_REQ_POOL_SIZE,), generator=gen, dtype=torch.int32
)
# Scrambled table with a tombstone, so a skipped gather cannot pass.
v2p = torch.randperm(_MAMBA_POOL_SIZE + 1, generator=gen).to(torch.int64)
v2p[7] = -1
out = torch.full((total_bs + _GUARD,), _OUT_POISON, dtype=torch.int32)
req_ref, req_fused = req_pool.clone().to(device), req_pool.clone().to(device)
out_ref, out_fused = out.clone().to(device), out.clone().to(device)
mapping_d, v2p_d = mapping.to(device), v2p.to(device)
_reference_chain(
req_pool_indices=req_ref,
mapping=mapping_d,
out_buf=out_ref,
valid_bs=valid_bs,
total_bs=total_bs,
v2p=v2p_d,
)
returned = fused_replay_state_indices(
req_pool_indices=req_fused,
mamba_index_mapping=mapping_d,
out_state_indices=out_fused,
valid_bs=valid_bs,
total_bs=total_bs,
v2p=v2p_d,
)
case = f"v2p {total_bs=} {num_padding=} {seed=}"
self.assertTrue(
torch.equal(out_ref[:total_bs], out_fused[:total_bs]),
f"state indices mismatch ({case}):\n"
f" ref {out_ref[:total_bs].tolist()}\n"
f" fused {out_fused[:total_bs].tolist()}",
)
self.assertTrue(torch.equal(returned, out_fused[:total_bs]), case)
self.assertTrue(
torch.equal(req_pool_ref_tail := req_ref[total_bs:], req_fused[total_bs:]),
f"guard tail diverged ({case}): {req_pool_ref_tail.tolist()}",
)
self.assertTrue(
(out_fused[total_bs:] == _OUT_POISON).all(),
f"out guard tail clobbered ({case})",
)
def test_v2p_matrix(self):
for total_bs in (1, 7, 32, 33):
paddings = sorted(
{0, 1, total_bs // 2, total_bs - 1} & set(range(total_bs))
)
for num_padding in paddings:
for seed in (0, 1):
with self.subTest(
total_bs=total_bs, num_padding=num_padding, seed=seed
):
self._run_v2p_case(
total_bs=total_bs, num_padding=num_padding, seed=seed
)
def test_matrix(self):
# Non-power-of-two sizes (7, 33) exercise the BS_UPPER in_range mask;
# num_padding sweeps none / one / half / all-but-one padded rows.
test/registered/unit/mem_cache/test_multi_ended_allocator.py
+77
View File
@@ -3100,5 +3100,82 @@ class TestDcpWidening(unittest.TestCase):
self.assertLess(base_cost * full_entry, mamba_bytes * dcp_size)
class TestFusedWriteLocTranslate(unittest.TestCase):
"""`write_loc_to_kernel_ids` must equal the arithmetic it stands for.
Nothing downstream can tell the two apart except by value, so the reference
here is the definition rather than a recorded expectation. Triton truncates
division toward zero where torch floors it, so the negative-loc and
tombstoned-page cases are the ones that matter.
"""
def _reference(self, loc, v2p, page_size, stride, dcp_size, dcp_rank):
out = []
for raw in loc.tolist():
if raw < 0 or (dcp_size > 1 and raw % dcp_size != dcp_rank):
out.append(0)
continue
collapsed = raw // dcp_size
page = collapsed // page_size
offset = collapsed % page_size if page_size > 1 else 0
out.append(max(int(v2p[page]) * stride + offset, 0))
return out
def _check(self, *, page_size, multiplier, dcp_size, dcp_rank, device):
from sglang.kernels.ops.memory.virtual_slot import write_loc_to_kernel_ids
span = page_size * dcp_size
locs = [0, 1, span - 1, span, 2 * span + 3, 5 * span + dcp_rank, -1]
locs += [3 * span + dcp_rank] # lands on the tombstoned page
loc = torch.tensor(locs, dtype=torch.int64, device=device)
# Table sized past the highest page any loc can name, scrambled, with
# one tombstone (-1) so a missing clamp shows up.
num_pages = max(locs) // (page_size * dcp_size) + 2
v2p = torch.tensor(
[(5 * i + 2) % num_pages for i in range(num_pages)] + [-1],
dtype=torch.int64,
device=device,
)
v2p[min(3, num_pages - 1)] = -1
stride = page_size * multiplier
got = write_loc_to_kernel_ids(
loc=loc,
v2p=v2p,
page_size=page_size,
stride=stride,
dcp_size=dcp_size,
dcp_rank=dcp_rank,
)
want = self._reference(
loc.cpu(), v2p.cpu(), page_size, stride, dcp_size, dcp_rank
)
self.assertEqual(got.tolist(), want, f"ps={page_size} dcp={dcp_size}")
# `out=` must write in place and agree (the cuda-graph-stable path).
dst = torch.full_like(loc, -7)
ret = write_loc_to_kernel_ids(
loc=loc,
v2p=v2p,
page_size=page_size,
stride=stride,
dcp_size=dcp_size,
dcp_rank=dcp_rank,
out=dst,
)
self.assertIs(ret, dst)
self.assertEqual(dst.tolist(), want)
def test_matches_reference_on_cpu(self):
for page_size in (1, 64):
for dcp_size, dcp_rank in ((1, 0), (2, 1), (4, 2)):
self._check(
page_size=page_size,
multiplier=7,
dcp_size=dcp_size,
dcp_rank=dcp_rank,
device="cpu",
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/server_args/test_unified_prefill_cuda_graph_gate.py
+1
View File
@@ -55,6 +55,7 @@ def _run_handler(*, prefill_backend, explicit):
"speculative_eagle_topk": None,
"enable_hierarchical_cache": False,
"enable_lmcache": False,
"enable_two_batch_overlap": False,
"dcp_size": 1,
"cuda_graph_config": cg,
"cuda_graph_backend_prefill": prefill_backend if explicit else None,
test/registered/unit/server_args/test_unified_tbo_gate.py
+69
View File
@@ -0,0 +1,69 @@
# Copyright 2023-2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""`--enable-unified-memory` refuses `--enable-two-batch-overlap`.
BUG REGRESSION. The combination launches and captures fine, then dies in the
forward path on the first captured decode replay. Nothing else rejects the
pair, so without this gate a running server crashes mid-serving.
python -m pytest test/registered/unit/server_args/test_unified_tbo_gate.py -v
"""
import unittest
from types import SimpleNamespace
from sglang.srt.arg_groups.kv_cache_hook import handle_unified_memory_pool
from sglang.srt.model_executor.cuda_graph_config import Backend
from sglang.srt.server_args import ServerArgs
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
def _run_handler(*, unified, tbo):
"""Run just `handle_unified_memory_pool` over a minimal stand-in."""
sa = ServerArgs.__new__(ServerArgs)
for name, value in {
"enable_unified_memory": unified,
"enable_two_batch_overlap": tbo,
"disaggregation_mode": "null",
"speculative_algorithm": None,
"speculative_eagle_topk": None,
"enable_hierarchical_cache": False,
"enable_lmcache": False,
"dcp_size": 1,
"cuda_graph_config": SimpleNamespace(
prefill=SimpleNamespace(backend=Backend.DISABLED),
decode=SimpleNamespace(backend=Backend.FULL),
),
"cuda_graph_backend_prefill": Backend.DISABLED,
}.items():
object.__setattr__(sa, name, value)
handle_unified_memory_pool(sa)
class TestUnifiedTboGate(unittest.TestCase):
def test_tbo_with_unified_memory_is_refused(self):
with self.assertRaises(AssertionError) as ctx:
_run_handler(unified=True, tbo=True)
self.assertIn("two-batch-overlap", str(ctx.exception))
def test_gate_fires_only_on_the_pair(self):
"""An inverted condition here would reject every unified launch."""
_run_handler(unified=True, tbo=False)
_run_handler(unified=False, tbo=True)
if __name__ == "__main__":
unittest.main()