[unified memory] Support DSPARK speculative decoding + fix two NaN root causes (page hand-out zeroing, CuTe int32 slot-stride wrap) (#33974)
This commit is contained in:
Cheng Wan
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
@@ -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()
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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`.
|
||||
|
||||
@@ -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): "
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user