diff --git a/python/sglang/kernels/ops/kimi_k3/kda_decode_mtp.py b/python/sglang/kernels/ops/kimi_k3/kda_decode_mtp.py index 554dd753c..8bd3d4326 100644 --- a/python/sglang/kernels/ops/kimi_k3/kda_decode_mtp.py +++ b/python/sglang/kernels/ops/kimi_k3/kda_decode_mtp.py @@ -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. diff --git a/python/sglang/kernels/ops/kvcache/zero_pages.py b/python/sglang/kernels/ops/kvcache/zero_pages.py new file mode 100644 index 000000000..afa3c0071 --- /dev/null +++ b/python/sglang/kernels/ops/kvcache/zero_pages.py @@ -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) diff --git a/python/sglang/kernels/ops/mamba/mamba_state_scatter_triton.py b/python/sglang/kernels/ops/mamba/mamba_state_scatter_triton.py index 5c050cf6b..945579893 100644 --- a/python/sglang/kernels/ops/mamba/mamba_state_scatter_triton.py +++ b/python/sglang/kernels/ops/mamba/mamba_state_scatter_triton.py @@ -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() diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 27fdbe107..b95e9d6ef 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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()) diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index 54a06c73f..08f2eb296 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -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 diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index df9754762..0da94696f 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -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 diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index c2642914a..b8f72b146 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -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 ) diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index 075c2e930..6f70d23a4 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -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( diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index 1575616ac..6c9aa877b 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -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`. diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 27793f396..477a648b9 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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): " diff --git a/test/registered/unit/layers/test_kda_decode_mtp_slot_stride.py b/test/registered/unit/layers/test_kda_decode_mtp_slot_stride.py new file mode 100644 index 000000000..45c3b2e78 --- /dev/null +++ b/test/registered/unit/layers/test_kda_decode_mtp_slot_stride.py @@ -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() diff --git a/test/registered/unit/layers/test_mamba_state_scatter_triton.py b/test/registered/unit/layers/test_mamba_state_scatter_triton.py index aaf5ce55a..74be3c848 100644 --- a/test/registered/unit/layers/test_mamba_state_scatter_triton.py +++ b/test/registered/unit/layers/test_mamba_state_scatter_triton.py @@ -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() diff --git a/test/registered/unit/mem_cache/test_unified_handout_zeroing.py b/test/registered/unit/mem_cache/test_unified_handout_zeroing.py new file mode 100644 index 000000000..6349f29e8 --- /dev/null +++ b/test/registered/unit/mem_cache/test_unified_handout_zeroing.py @@ -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()