minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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
test/registered/unit/layers/test_kda_decode_mtp_slot_stride.py
+165
View File
@@ -0,0 +1,165 @@
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=60, stage="base-b", runner_config="1-gpu-small")
import importlib.util
import unittest
import torch
TILE_K = 128
def _sm100():
return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10
def _strided_replica(shape, slot_stride, dtype, device):
"""A tensor whose slot dim (dim 0) has an ARTIFICIALLY large stride, with
zeroed gap bytes — the unified pool's envelope-strided state layout, scaled
so `slot * stride` exceeds int32 at small slot ids."""
inner = 1
for s in shape[1:]:
inner *= s
reach = (shape[0] - 1) * slot_stride + inner
base = torch.zeros(reach, dtype=dtype, device=device)
strides = [slot_stride]
acc = inner
for s in shape[1:]:
acc //= s
strides.append(acc)
return base.as_strided(tuple(shape), tuple(strides))
@unittest.skipUnless(_sm100(), "SM100-only CuTe kernel")
@unittest.skipUnless(
importlib.util.find_spec("cutlass") is not None, "nvidia-cutlass-dsl required"
)
class TestKdaDecodeMtpSlotStride(unittest.TestCase):
"""Root-cause guard: `slot * stride` must be computed in int64.
The DSPARK KDA verify kernel compiles with STATIC CuTe layouts, so a
state-pool slot stride that individually fits int32 folds into 32-bit
arithmetic and `slot * stride` wraps mod 2^32 once the product exceeds
int32 — reads land inside other slots (silent corruption) or off the
allocation (illegal access). The unified pool's envelope-strided KDA
views reach that regime at slot ids ~153 (conv) / ~306 (ssm). This test
reproduces the regime with an artificially large ssm slot stride at a
small slot id and asserts bitwise parity against a contiguous pool."""
def test_wrap_regime_matches_contiguous(self):
from sglang.kernels.ops.kimi_k3.kda_decode_mtp import (
fused_kda_decode_mtp_dspark,
)
device = "cuda"
torch.manual_seed(3)
H, num_spec = 2, 7
T, N = 1 + num_spec, 1
dim = H * TILE_K
# slot * stride crosses 2^31 elements at slot 8. The stride must NOT
# be a power of two: pow2 constants lower to shifts, which dodge the
# 32-bit imul this test pins (the real pool strides — e.g. K3's
# 14,042,880 ssm / 28,085,760 conv — are not pow2). Multiple of 4
# (wrapper's cp.async alignment contract). Both state families get
# the huge stride: in the production repro the conv direct-index path
# (cs_q[slot, ch, w]) wrapped at lower slot ids than the ssm tiled
# copy, so pinning only one path can silently pass.
slot_id, slots = 8, 9
ssm_slot_stride = (1 << 28) + 12_344 # fp32 base ~8.6 GB
conv_slot_stride = (1 << 28) + 23_448 # bf16 base ~4.3 GB x3
free = torch.cuda.mem_get_info()[0]
if free < 26 << 30:
self.skipTest(f"needs ~26GB free GPU memory, have {free >> 30}GB")
def acts(shape, dtype=torch.bfloat16):
return (torch.randn(shape, device=device, dtype=torch.float32) * 0.1).to(
dtype
)
x_q, x_k, x_v, g = (acts((1, T, H, TILE_K)) for _ in range(4))
beta = acts((1, T, H))
w = torch.randn(3 * dim, 4, device=device, dtype=torch.float32) * 0.1
w_q, w_k, w_v = w.split([dim, dim, dim], dim=0)
A_log = torch.randn(H, device=device, dtype=torch.float32) * 0.1
dt_bias = torch.randn(dim, device=device, dtype=torch.float32) * 0.1
state_c = torch.randn(
slots, H, TILE_K, TILE_K, device=device, dtype=torch.float32
)
# conv pool in the backend's post-split/transpose shape [slots, dim, 3]
# with the production stride pattern (slot_stride, 1, dim): the
# underlying envelope is [slots, 3, dim] and the backend transposes.
conv_c = [
(torch.randn(slots, 3, dim, device=device, dtype=torch.float32) * 0.1)
.to(torch.bfloat16)
.transpose(-1, -2)
for _ in range(3)
]
inter_ssm = torch.zeros(
2, T, H, TILE_K, TILE_K, device=device, dtype=torch.float32
)
inter_conv = [
torch.zeros(2, T, dim, 3, device=device, dtype=torch.bfloat16)
for _ in range(3)
]
common = dict(
x_q=x_q,
x_k=x_k,
x_v=x_v,
w_q=w_q,
w_k=w_k,
w_v=w_v,
g=g,
beta=beta,
A_log=A_log,
dt_bias=dt_bias,
intermediate_state_indices=torch.zeros(N, dtype=torch.int32, device=device),
ssm_state_indices=torch.full(
(N,), slot_id, dtype=torch.int32, device=device
),
cu_seqlens=torch.tensor([0, T], dtype=torch.int32, device=device),
lower_bound=-5.0,
)
def run(state, conv, issm, iconv):
out = fused_kda_decode_mtp_dspark(
recurrent_state=state,
cs_q=conv[0],
cs_k=conv[1],
cs_v=conv[2],
intermediate_ssm=issm,
intermediate_conv_q=iconv[0],
intermediate_conv_k=iconv[1],
intermediate_conv_v=iconv[2],
**common,
)
torch.cuda.synchronize()
return out
ref = run(state_c, conv_c, inter_ssm.clone(), [c.clone() for c in inter_conv])
state_s = _strided_replica(
(slots, H, TILE_K, TILE_K), ssm_slot_stride, torch.float32, device
)
state_s.copy_(state_c)
conv_s = []
for c in conv_c:
v = _strided_replica(
(slots, 3, dim), conv_slot_stride, torch.bfloat16, device
).transpose(-1, -2)
v.copy_(c)
conv_s.append(v)
issm_s = inter_ssm.clone()
iconv_s = [c.clone() for c in inter_conv]
got = run(state_s, conv_s, issm_s, iconv_s)
# Pre-fix: 32-bit `slot * stride` wraps (8 * 2^28 = 2^31) and the read
# lands at offset 0 of the pool — silently returning slot 0's state —
# or off the allocation. Post-fix: bit-exact.
torch.testing.assert_close(got, ref, rtol=0, atol=0)
if __name__ == "__main__": # pragma: no cover
unittest.main()
test/registered/unit/layers/test_mamba_state_scatter_triton.py
+153 -1
View File
@@ -1,7 +1,13 @@
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.ci.ci_register import (
register_amd_ci,
register_cpu_ci,
register_cuda_ci,
)
register_cuda_ci(est_time=7, stage="base-b", runner_config="1-gpu-small")
register_amd_ci(est_time=7, suite="stage-b-test-1-gpu-small-amd-mi35x")
# The dst layout-contract tests run on CPU (no kernel launch).
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
import unittest
@@ -9,14 +15,23 @@ import torch
try:
from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
_require_entry_contiguous_dst,
fused_conv_window_scatter_with_mask,
fused_mamba_state_scatter_with_mask,
)
_FUSED_IMPORT_ERROR = None
except Exception as e: # pragma: no cover
_require_entry_contiguous_dst = None
fused_conv_window_scatter_with_mask = None
fused_mamba_state_scatter_with_mask = None
_FUSED_IMPORT_ERROR = e
from sglang.srt.mem_cache.layout.page_major import (
build_page_major_mamba_views,
mamba_entry_bytes,
)
def _ref_scatter(dst, src, dst_indices, src_indices, step_indices):
"""Reference implementation using PyTorch advanced indexing."""
@@ -213,5 +228,142 @@ class TestMambaStateScatterCorrectness(unittest.TestCase):
torch.testing.assert_close(conv_fused, conv_ref)
def _make_envelope_views(device="cpu"):
"""Envelope-strided conv/temporal views, exactly as UnifiedMambaPool /
the page-major MambaPool serve them ((num_layers, max_slots, *inner) with
slot stride = the multi-layer entry envelope). Mirrors
test_flashkda_strided_state_access.py's setup."""
layers, slots = 2, 16
temporal_shape = (2, 4, 4) # (H, V, K)
conv_shapes = ((8, 3),) # (dim, K-1) as fused_conv_window_scatter expects
conv_dtype = torch.bfloat16
temporal_dtype = torch.float32
entry = mamba_entry_bytes(
layer_num=layers,
conv_state_shapes=conv_shapes,
conv_dtype=conv_dtype,
temporal_state_shape=temporal_shape,
temporal_dtype=temporal_dtype,
)
raw = torch.zeros(slots * entry, dtype=torch.uint8, device=device)
conv_views, temporal = build_page_major_mamba_views(
raw,
layer_num=layers,
conv_state_shapes=conv_shapes,
conv_dtype=conv_dtype,
temporal_state_shape=temporal_shape,
temporal_dtype=temporal_dtype,
max_slots=slots,
)
return conv_views, temporal
class TestScatterDstLayoutContract(unittest.TestCase):
"""The scatter wrappers' dst contract (CPU, no kernel launch).
Derived property: the Triton kernels index dst through its REAL
``stride(0)``/``stride(1)`` plus a FLAT in-entry element offset, so the
layout contract is "arbitrary layer/slot strides, contiguous trailing
entry dims" — NOT ``dst.is_contiguous()``. The blanket contiguity assert
the wrappers used to carry rejected the unified pool's envelope-strided
views (DSPARK verify commit under --enable-unified-memory); the relaxed
check must keep accepting them while still rejecting a dst whose entry
dims the kernels would mis-address."""
def setUp(self):
if _require_entry_contiguous_dst is None:
self.skipTest(f"import failed: {_FUSED_IMPORT_ERROR}")
def test_envelope_strided_views_accepted(self):
conv_views, temporal = _make_envelope_views()
# Precondition: the views really are envelope-strided (else the
# property below is vacuous).
self.assertFalse(temporal.is_contiguous())
self.assertFalse(conv_views[0].is_contiguous())
# dst = temporal (5-D) for the dense scatter, conv (4-D) for the
# conv-window scatter; entry dims start at 2 for both.
_require_entry_contiguous_dst(temporal, 2, "test")
_require_entry_contiguous_dst(conv_views[0], 2, "test")
def test_entry_noncontiguous_dst_rejected(self):
# A dst whose ENTRY dims are strided (inner transpose) would be
# mis-addressed by the flat in-entry offset; the check must not have
# degraded to always-pass.
dst = torch.zeros(2, 4, 8, 3).transpose(-1, -2) # entry dims strided
with self.assertRaises(ValueError):
_require_entry_contiguous_dst(dst, 2, "test")
class TestMambaStateScatterEnvelopeDst(unittest.TestCase):
"""End-to-end: both scatter wrappers accept the unified pool's
envelope-strided dst views and address slots through the real strides
(bug regression: the wrappers used to raise 'dst tensor must be
contiguous' on these views)."""
@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for this test.")
def test_fused_scatter_envelope_strided_dst(self):
if fused_mamba_state_scatter_with_mask is None:
self.skipTest(f"import failed: {_FUSED_IMPORT_ERROR}")
torch.manual_seed(7)
device = torch.device("cuda")
conv_views, temporal = _make_envelope_views(device=device)
layers, slots = temporal.shape[0], temporal.shape[1]
temporal_shape = tuple(temporal.shape[2:]) # (H, V, K)
dim, km1 = conv_views[0].shape[2], conv_views[0].shape[3]
B, D = 5, 3
temporal[:] = torch.randn_like(temporal)
conv_views[0][:] = torch.randn_like(conv_views[0])
temporal_before = temporal.clone()
conv_before = conv_views[0].clone()
# Dense SSM scatter: contiguous per-step src (the intermediate cache).
src_ssm = torch.randn(
(layers, B, D) + temporal_shape, device=device, dtype=temporal.dtype
)
# Conv-window scatter: overlapping as_strided src over a shared
# [dim, D+K-2] buffer per (layer, slot) — window t = shared[:, t:t+K-1].
shared = torch.randn(
(layers, B, dim, D + km1 - 1), device=device, dtype=conv_views[0].dtype
)
src_conv = shared.as_strided(
(layers, B, D, dim, km1),
(
shared.stride(0),
shared.stride(1),
1, # step: window slides by one position
shared.stride(2),
1, # within-window
),
)
dst_indices = torch.randperm(slots, device=device, dtype=torch.int64)[:B].to(
torch.int32
)
step_indices = torch.randint(0, D, (B,), device=device, dtype=torch.int64)
step_indices[0] = -1 # one rejected row must be skipped
fused_mamba_state_scatter_with_mask(
temporal, src_ssm, dst_indices, step_indices
)
fused_conv_window_scatter_with_mask(
conv_views[0], src_conv, dst_indices, step_indices
)
# Reference via advanced indexing (layout-agnostic).
valid = step_indices >= 0
d = dst_indices[valid].long()
s = torch.arange(B, device=device)[valid]
t = step_indices[valid]
expect_temporal = temporal_before.clone()
expect_temporal[:, d] = src_ssm[:, s, t]
expect_conv = conv_before.clone()
expect_conv[:, d] = src_conv[:, s, t]
torch.testing.assert_close(temporal, expect_temporal)
torch.testing.assert_close(conv_views[0], expect_conv)
if __name__ == "__main__": # pragma: no cover
unittest.main()
test/registered/unit/mem_cache/test_unified_handout_zeroing.py
+133
View File
@@ -0,0 +1,133 @@
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=4, stage="base-b", runner_config="1-gpu-small")
import unittest
import torch
from sglang.srt.mem_cache.multi_ended_allocator import MultiEndedAllocator
from sglang.srt.mem_cache.unified_memory_pool import (
MambaSubPoolSpec,
MLASubPoolSpec,
UnifiedKVPool,
UnifiedMLATokenToKVPool,
)
BF16_NAN = 0x7FC1 # LE bf16 NaN bit pattern, as SGLANG_DEBUG_POISON_POOL fills
def _build(device, page_size=1, kernel_page_multiplier=None):
"""A tiny MLA+mamba unified pool + full-side allocator.
Mirrors init_unified_mamba_pools' construction just enough for the
allocator hand-out path (the piece under test)."""
layer_num = 2
full_spec = MLASubPoolSpec(
name="full",
layer_num=layer_num,
grow_direction="up",
kv_lora_rank=16,
qk_rope_head_dim=8,
store_dtype=torch.bfloat16,
)
mamba_spec = MambaSubPoolSpec(
name="mamba",
layer_num=1,
grow_direction="down",
conv_state_shapes=((8, 3),),
conv_dtype=torch.bfloat16,
temporal_state_shape=(2, 4, 4),
temporal_dtype=torch.float32,
)
total_bytes = 4096 * full_spec.entry_bytes()
buf = UnifiedKVPool(
total_bytes=total_bytes,
sub_pool_specs=[full_spec, mamba_spec],
device=device,
enable_memory_saver=False,
page_size=page_size,
view_tail_pad_bytes=page_size * full_spec.entry_bytes(),
)
kvcache = UnifiedMLATokenToKVPool(
unified_buffer=buf,
sub_pool_name="full",
kv_cache_dtype=torch.bfloat16,
page_size=page_size,
)
allocator = MultiEndedAllocator(
kvcache=kvcache,
unified_buffer=buf,
sub_pool_name="full",
device=device,
is_id_owner=True,
page_size=page_size,
kernel_page_multiplier=(
layer_num if kernel_page_multiplier is None else kernel_page_multiplier
),
)
return buf, kvcache, allocator
@unittest.skipUnless(torch.cuda.is_available(), "CUDA required (fused alloc kernel)")
class TestUnifiedHandoutZeroing(unittest.TestCase):
"""Root-cause guard: pages must leave the allocator ZEROED.
The trtllm MLA kernel arithmetically masks (NaN-unsafe) the unwritten
tail rows of a request's last partial page, so recycled / fresh page
bytes must never carry NaN bit patterns. Static pools get this from
torch.zeros; the unified pool must re-establish it at every hand-out."""
def _poison(self, buf):
buf._raw.view(torch.int16).fill_(BF16_NAN)
def _env(self, buf, kvcache):
return buf._raw[: kvcache._num_pages * kvcache._page_bytes].view(
kvcache._num_pages, kvcache._page_bytes
)
def _phys_pages(self, allocator, virt_tokens):
return (allocator.translate_kv_loc(virt_tokens) // allocator.page_size).unique()
def test_fresh_and_recycled_pages_zeroed(self):
buf, kvcache, allocator = _build("cuda")
env = self._env(buf, kvcache)
# Fresh hand-out over a poisoned pool (the deterministic form of
# "freed GPU heap happened to contain NaN patterns").
self._poison(buf)
out = allocator.alloc(16)
self.assertIsNotNone(out)
pages = self._phys_pages(allocator, out)
self.assertTrue((env[pages] == 0).all().item())
# Untouched pages must still be poisoned, else the assert above is
# vacuous (a whole-pool memset would also pass it).
wm_page = int(pages.max().item()) + 2
self.assertFalse((env[wm_page] == 0).all().item())
# Recycle: free, re-poison the raw bytes (data only; v2p bookkeeping
# is separate storage), re-alloc — recycled pages must be zeroed too.
allocator.free(out)
self._poison(buf)
out2 = allocator.alloc(16)
self.assertIsNotNone(out2)
pages2 = self._phys_pages(allocator, out2)
self.assertTrue((env[pages2] == 0).all().item())
def test_zeroing_enabled_for_single_layer_multiplier(self):
# A shard owning exactly ONE full-attention MLA layer has
# kernel_page_multiplier == 1 but its pool is still
# UnifiedMLATokenToKVPool with the same NaN-unsafe partial-page
# reads — zeroing must key on the pool type, not on multiplier > 1.
buf, kvcache, allocator = _build("cuda", kernel_page_multiplier=1)
self.assertTrue(allocator._zero_pages_on_alloc)
self._poison(buf)
out = allocator.alloc(8)
self.assertIsNotNone(out)
env = self._env(buf, kvcache)
pages = self._phys_pages(allocator, out)
self.assertTrue((env[pages] == 0).all().item())
if __name__ == "__main__": # pragma: no cover
unittest.main()