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

[unified memory] Support DSPARK speculative decoding + fix two NaN root causes (page hand-out zeroing, CuTe int32 slot-stride wrap) (#33974)

Browse Source
This commit is contained in:
Cheng Wan
2026-08-10 10:35:06 -07:00
committed by GitHub
parent ec9babe36c
commit 7738062294
13 changed files with 661 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/kimi_k3/kda_decode_mtp.py
+3
View File
@@ -279,6 +279,9 @@ def kda_decode_mtp_kernel(
# staged if (UNSUP_EARLY_EXIT).
nvvm.exit()
# int64: `slot * stride` overflows int32 on envelope-strided pools.
slot = cutlass.Int64(slot)
# q/k/g each run on P1_JOB_WARPS warps split by token parity and the v-conv
# takes the rest. Each token's conv is an independent window over globals,
# so the split needs no cross-warp communication.
python/sglang/kernels/ops/kvcache/zero_pages.py
+48
View File
@@ -0,0 +1,48 @@
"""Zero whole page envelopes of the unified pool by physical page id.
The pool is viewed as int64 words (the MLA page envelope is always
8-byte-aligned: entry bytes per layer = kv_cache_dim * itemsize, a multiple
of 8), one wide element per lane; grid = (num_pages, page word blocks).
"""
from __future__ import annotations
import torch
import triton
import triton.language as tl
@triton.jit
def _zero_pages_kernel(
buf_ptr, # int64 view of the raw pool buffer
pages_ptr, # int64 [M] physical page ids to zero
page_words, # int64 words per page envelope
BLOCK: tl.constexpr,
):
m = tl.program_id(0)
blk = tl.program_id(1)
pg = tl.load(pages_ptr + m).to(tl.int64)
offs = blk * BLOCK + tl.arange(0, BLOCK)
mask = offs < page_words
tl.store(buf_ptr + pg * page_words + offs, 0, mask=mask)
_BLOCK = 2048
def zero_pages(
raw: torch.Tensor,
pages: torch.Tensor,
num_pages: int,
page_bytes: int,
) -> None:
"""Zero the listed physical PAGE envelopes of the uint8 pool `raw`."""
m = int(pages.numel())
if m == 0:
return
assert raw.dtype == torch.uint8, f"expected uint8 pool, got {raw.dtype}"
assert page_bytes % 8 == 0, f"page_bytes {page_bytes} not int64-aligned"
page_words = page_bytes // 8
words = raw[: num_pages * page_bytes].view(torch.int64)
grid = (m, triton.cdiv(page_words, _BLOCK))
_zero_pages_kernel[grid](words, pages.to(torch.int64), page_words, BLOCK=_BLOCK)
python/sglang/kernels/ops/mamba/mamba_state_scatter_triton.py
+24 -9
View File
@@ -11,6 +11,25 @@ import triton
import triton.language as tl
def _require_entry_contiguous_dst(
dst: torch.Tensor, entry_start_dim: int, fn_name: str
) -> None:
"""dst layout contract: the kernels index through the real layer/slot
strides (int64) plus a FLAT element offset within one (layer, slot)
entry — layer/slot strides may be arbitrary (envelope-strided unified
pool views), but the trailing entry dims must be contiguous.
"""
expected = 1
for i in range(dst.ndim - 1, entry_start_dim - 1, -1):
if dst.shape[i] != 1 and dst.stride(i) != expected:
raise ValueError(
f"{fn_name}: dst entry dims (dims {entry_start_dim}.."
f"{dst.ndim - 1}) must be contiguous; got "
f"shape={tuple(dst.shape)} strides={tuple(dst.stride())}"
)
expected *= dst.shape[i]
@triton.jit
def track_mamba_state_if_needed_kernel(
conv_states_ptr,
@@ -271,9 +290,7 @@ def fused_mamba_state_scatter_with_mask(
dst_indices_raw = dst_indices_raw.to(torch.int32).contiguous()
step_indices_raw = step_indices_raw.to(torch.int32).contiguous()
# Ensure tensors are contiguous
if not dst.is_contiguous():
raise ValueError("dst tensor must be contiguous")
_require_entry_contiguous_dst(dst, 2, "fused_mamba_state_scatter_with_mask")
if not src.is_contiguous():
raise ValueError("src tensor must be contiguous")
@@ -420,12 +437,10 @@ def fused_conv_window_scatter_with_mask(
src_step_size = src.shape[2]
dst_req_size = dst.shape[1]
# `dst` stays contiguous; `src` is an intentionally non-contiguous (overlapping)
# view, so we do NOT assert src contiguity here (unlike the dense scatter).
if not dst.is_contiguous():
raise ValueError(
"dst tensor in fused_conv_window_scatter_with_mask must be contiguous"
)
# `src` is an intentionally non-contiguous (overlapping) view indexed per
# dim through its real strides, so we do NOT assert src contiguity here
# (unlike the dense scatter).
_require_entry_contiguous_dst(dst, 2, "fused_conv_window_scatter_with_mask")
dst_indices_raw = dst_indices_raw.to(torch.int32).contiguous()
step_indices_raw = step_indices_raw.to(torch.int32).contiguous()
python/sglang/srt/environ.py
+2
View File
@@ -334,6 +334,8 @@ class Envs:
SGLANG_GRAPH_BATCH_CAPTURE = EnvBool(False)
SGLANG_FORCE_SHUTDOWN = EnvBool(False)
SGLANG_DEBUG_MEMORY_POOL = EnvBool(False)
# NaN-fill the unified memory pool at boot (debug repro switch).
SGLANG_DEBUG_POISON_POOL = EnvBool(False)
SGLANG_DSPARK_DEBUG_CONFIDENCE_PREFIX_SCHEDULER = EnvBool(False)
SGLANG_DSPARK_DEBUG_CONFIDENCE_METRICS = EnvBool(False)
SGLANG_DSPARK_DEBUG_DUMP = EnvTuple(tuple())
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
+6
View File
@@ -1204,6 +1204,12 @@ class HybridLinearAttnBackend(AttentionBackend):
del req_pool_indices
request_number = last_correct_step_indices.shape[0]
# `mamba_track_indices` is VIRTUAL; the scatter writes physical views.
if mamba_track_indices is not None:
mamba_track_indices = self.linear_attn_backend._translate_mamba_indices(
mamba_track_indices
)
state_indices_tensor = (
self.linear_attn_backend.forward_metadata.mamba_cache_indices[
:request_number
python/sglang/srt/layers/attention/trtllm_mla_backend.py
+13 -6
View File
@@ -620,9 +620,11 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
# Unified pool: precompute the DENSE KV write loc into the capture-stable
# buffer (both capture and each replay-prep run this out of the graph),
# so the in-graph set_mla_kv_buffer writes a dense loc without capturing a
# translate. Only decode writes KV under unified (spec is gated off).
if self._unified_mla and forward_mode.is_decode_or_idle():
# so the in-graph set_mla_kv_buffer writes a dense loc without capturing
# a translate.
if self._unified_mla 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_dense[:n]
@@ -1243,9 +1245,14 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
assert (
k is not None and k_rope is not None
), "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None."
self.token_to_kv_pool.set_mla_kv_buffer(
layer, forward_batch.out_cache_loc, k, k_rope
)
if self._decode_dense_loc is not None:
self.token_to_kv_pool.set_mla_kv_buffer(
layer, self._decode_dense_loc, k, k_rope, loc_is_dense=True
)
else:
self.token_to_kv_pool.set_mla_kv_buffer(
layer, forward_batch.out_cache_loc, k, k_rope
)
# TODO refactor to avoid code duplication
# Prepare query tensor inline
python/sglang/srt/mem_cache/kv_cache_configurator.py
+48
View File
@@ -386,6 +386,41 @@ class KVCacheConfigurator:
unified_memory_pool=bundle.unified_memory_pool,
)
# The unified allocator hands out VIRTUAL token ids from the whole
# virtual space (> max_total_num_tokens); the direct-indexed draft
# pool must be sized by that space.
draft_virtual_id_space: Optional[int] = None
if self.is_draft_worker and token_to_kv_pool_allocator is not None:
from sglang.srt.mem_cache.multi_ended_allocator import (
UnifiedMambaTokenToKVPoolAllocator,
UnifiedSWATokenToKVPoolAllocator,
)
if isinstance(token_to_kv_pool_allocator, UnifiedSWATokenToKVPoolAllocator):
raise ValueError(
"Speculative decoding with --enable-unified-memory is only "
"supported for hybrid-Mamba targets; the unified hybrid-SWA "
"pool's draft sizing (virtual-id space) is not wired yet."
)
if isinstance(
token_to_kv_pool_allocator, UnifiedMambaTokenToKVPoolAllocator
):
draft_virtual_id_space = token_to_kv_pool_allocator.size_full
assert draft_virtual_id_space >= sizes.max_total_num_tokens, (
"unified allocator virtual space smaller than the token "
f"budget: size_full={draft_virtual_id_space} < "
f"max_total_num_tokens={sizes.max_total_num_tokens}"
)
# Round UP to page alignment (paged draft backends view the
# pool as (-1, page_size, H, D); size_full is not aligned).
page = max(int(self.pool_page_size or 1), 1)
draft_virtual_id_space = (
(draft_virtual_id_space + page - 1) // page * page
)
sizes = msgspec.structs.replace(
sizes, max_total_num_tokens=draft_virtual_id_space
)
# Initialize req_to_token_pool
if req_to_token_pool is None:
req_to_token_pool = self._build_req_to_token_pool(
@@ -428,6 +463,15 @@ class KVCacheConfigurator:
req_to_token_pool=req_to_token_pool,
)
if draft_virtual_id_space is not None:
assert token_to_kv_pool.size >= draft_virtual_id_space, (
"draft token_to_kv_pool smaller than the shared unified "
f"allocator's virtual-id space: pool size="
f"{token_to_kv_pool.size} < size_full={draft_virtual_id_space}; "
"verify-window writes at high virtual ids would go out of "
"bounds."
)
token_to_kv_pool_allocator = self._build_token_to_kv_pool_allocator(
sizes=sizes,
token_to_kv_pool=token_to_kv_pool,
@@ -870,6 +914,10 @@ class KVCacheConfigurator:
# default keeps upstream's per-layer layout. The Mamba state pool is routed
# separately via `mamba_envelope_layout` on the req-to-token pool above.
enable_page_major = get_memory().enable_page_major_kv_layout
if self.is_draft_worker and get_memory().enable_unified_memory:
# Page-major is a target-pool layout choice; the draft backend
# reads the plain per-layer contiguous layout.
enable_page_major = False
mha_pool_class = (
PageMajorMHATokenToKVPool if enable_page_major else MHATokenToKVPool
)
python/sglang/srt/mem_cache/multi_ended_allocator.py
+20 -1
View File
@@ -38,7 +38,10 @@ from sglang.srt.mem_cache.allocator.paged import (
alloc_extend_kernel,
)
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.unified_memory_pool import UnifiedKVPool
from sglang.srt.mem_cache.unified_memory_pool import (
UnifiedKVPool,
UnifiedMLATokenToKVPool,
)
from sglang.srt.utils.common import get_num_new_pages, next_power_of_2
logger = logging.getLogger(__name__)
@@ -136,6 +139,8 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
# once page ids are scaled by layer_num — `translate_kv_loc_dense` emits
# that space. 1 for sub-pools whose kernels take real physical ids.
self.kernel_page_multiplier = kernel_page_multiplier
# Zero page envelopes on hand-out — see _maybe_zero_pages.
self._zero_pages_on_alloc = isinstance(kvcache, UnifiedMLATokenToKVPool)
# Overlap mode: `free` drops a wait_stream(forward_stream) barrier so its
# v2p writes + move kernel serialize after the in-flight forward.
self.forward_stream = forward_stream
@@ -617,6 +622,7 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
if self.lazy_compaction: # live_page_count tracked only in lazy mode
self.live_page_count += N
self._maybe_zero_pages(phys_pages)
return phys_pages
# SLOW PATH: holes exist — drain them first, then bind.
@@ -624,8 +630,21 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
if phys_pages is None:
return None
self.bind(v_pages, phys_pages)
self._maybe_zero_pages(phys_pages)
return phys_pages
def _maybe_zero_pages(self, phys_pages: torch.Tensor) -> None:
"""Zero the page ENVELOPES on hand-out (MLA-dense full pool only):
the MLA kernels arithmetically mask the rows beyond seq_len, so
never-written page bytes must read as finite values. Runs on the
schedule stream, ordered before the consuming forward by the
run_batch wait_stream fence.
"""
if not self._zero_pages_on_alloc or phys_pages.numel() == 0:
return
with record_function("MultiEndedAlloc._maybe_zero_pages"):
self._kvcache.zero_physical_pages(phys_pages)
# -- translate (virtual TOKEN ids -> physical TOKEN ids) --
def translate_kv_loc(
python/sglang/srt/mem_cache/unified_memory_pool.py
+22 -1
View File
@@ -33,7 +33,9 @@ import triton
from torch.profiler import record_function
from sglang.kernels.ops.kvcache.cache_move import store_cache_4d_kernel
from sglang.kernels.ops.kvcache.zero_pages import zero_pages
from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
from sglang.srt.environ import envs
from sglang.srt.mem_cache.layout.page_major import (
build_dense_mla_views,
build_page_major_mamba_views,
@@ -258,7 +260,16 @@ class UnifiedKVPool:
self._raw = torch.empty(
total_bytes + view_tail_pad_bytes, dtype=torch.uint8, device=device
)
self._raw.zero_() # unset slots must read as zeros (matches non-shared)
if envs.SGLANG_DEBUG_POISON_POOL.get():
# Debug: bf16-NaN-fill so NaN-unsafe reads of never-written bytes
# fail deterministically.
self._raw.view(torch.int16).fill_(0x7FC1)
logger.warning(
"[unified-memory-pool] POISONED: pool filled with bf16-NaN "
"patterns (SGLANG_DEBUG_POISON_POOL)"
)
else:
self._raw.zero_() # unset slots must read as zeros (matches non-shared)
self._max_slots: Dict[str, int] = {}
self._anchor_bytes: Dict[str, int] = {}
@@ -671,6 +682,16 @@ class UnifiedMLATokenToKVPool(MLATokenToKVPool):
)
env[tgt_pages] = env[src_pages]
def zero_physical_pages(self, phys_pages: torch.Tensor) -> None:
"""Zero whole page envelopes (PHYSICAL page ids) on allocator
hand-out."""
zero_pages(
self._unified_buffer._raw,
phys_pages,
self._num_pages,
self._page_bytes,
)
class UnifiedMambaPool(MambaPool):
"""Mamba state pool whose conv/temporal state are strided views into a `UnifiedKVPool`.
python/sglang/srt/server_args.py
+24 -3
View File
@@ -8010,10 +8010,31 @@ class ServerArgs:
assert self.disaggregation_mode == "null", (
"--enable-unified-memory is not yet compatible with PD " "disaggregation."
)
assert self.speculative_algorithm is None, (
"--enable-unified-memory is not yet compatible with speculative "
"decoding."
assert self.speculative_algorithm in (None, "DSPARK"), (
"--enable-unified-memory only supports --speculative-algorithm "
"DSPARK (chain draft); other speculative algorithms are not yet "
"audited for the unified pool's virtual/dense loc translation. Got "
f"--speculative-algorithm={self.speculative_algorithm!r}."
)
if self.speculative_algorithm == "DSPARK":
assert self.speculative_eagle_topk in (None, 1), (
"--enable-unified-memory + DSPARK supports a linear draft "
"chain only (--speculative-eagle-topk in {None, 1}); tree "
"verify is not audited for the unified pool. Got "
f"--speculative-eagle-topk={self.speculative_eagle_topk!r}."
)
# Both roles: verify routes to either backend depending on
# --speculative-attention-mode.
spec_allowed = {"triton", "trtllm_mla", "cutedsl_mla", "tokenspeed_mla"}
spec_backends = set(self._resolved_attention_backends())
spec_backends.discard(None)
assert spec_backends <= spec_allowed, (
"--enable-unified-memory + DSPARK requires spec-verify-audited "
f"attention backends {sorted(spec_allowed)} for both prefill "
f"and decode; got {sorted(spec_backends)}. flashinfer / fa3 do "
"not translate speculative verify indices to the unified "
"pool's dense space yet."
)
assert not (self.enable_hierarchical_cache or self.enable_lmcache), (
"--enable-unified-memory is not yet compatible with hierarchical / "
"host-tiered KV cache (--enable-hierarchical-cache / --enable-lmcache): "