From b5766336d4fa1f235a05e864b5c5be6fd13cd5e5 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Mon, 7 Sep 2026 01:10:44 -0700 Subject: [PATCH] [Perf] Unified memory: close the DCP decode gap on Blackwell (#37926) --- .../ops/mamba/mamba_state_indices_triton.py | 18 ++- python/sglang/kernels/ops/memory/__init__.py | 1 + .../sglang/kernels/ops/memory/virtual_slot.py | 133 ++++++++++++++++++ python/sglang/srt/arg_groups/kv_cache_hook.py | 6 + .../srt/batch_overlap/two_batch_overlap.py | 5 + .../attention/hybrid_linear_attn_backend.py | 13 +- .../srt/layers/attention/triton_backend.py | 22 ++- .../layers/attention/trtllm_mla_backend.py | 20 +-- .../mem_cache/allocator/unified_hybrid_swa.py | 5 +- .../srt/mem_cache/allocator/unified_mamba.py | 5 +- .../mem_cache/allocator/unified_sub_pool.py | 69 +++++---- .../srt/mem_cache/kv_index_translator.py | 41 +++++- python/sglang/srt/mem_cache/memory_pool.py | 22 +++ .../srt/mem_cache/unified_memory_pool.py | 9 ++ .../srt/model_executor/forward_batch_info.py | 5 + .../runner/decode_cuda_graph_runner.py | 1 + .../test_trtllm_mla_family_dcp_metadata.py | 83 +++++++++++ .../mamba/test_fused_replay_state_indices.py | 79 ++++++++++- .../mem_cache/test_multi_ended_allocator.py | 77 ++++++++++ .../test_unified_prefill_cuda_graph_gate.py | 1 + .../unit/server_args/test_unified_tbo_gate.py | 69 +++++++++ 21 files changed, 609 insertions(+), 75 deletions(-) create mode 100644 test/registered/unit/server_args/test_unified_tbo_gate.py diff --git a/python/sglang/kernels/ops/mamba/mamba_state_indices_triton.py b/python/sglang/kernels/ops/mamba/mamba_state_indices_triton.py index 553ec66f6..a5fe15d43 100644 --- a/python/sglang/kernels/ops/mamba/mamba_state_indices_triton.py +++ b/python/sglang/kernels/ops/mamba/mamba_state_indices_triton.py @@ -14,6 +14,8 @@ MTP inter-phase seam: This module fuses that chain into a single launch. """ +from typing import Optional + import torch import triton import triton.language as tl @@ -24,15 +26,20 @@ def _fused_replay_state_indices_kernel( req_pool_indices_ptr, # (total_bs,) int64 — static replay buffer mamba_map_ptr, # (req_pool_size,) int32 — req_index_to_mamba_index_mapping out_ptr, # (total_bs,) int32 — state_indices_list[bs - 1] + v2p_ptr, # (num_slots + 1,) int64 — mamba virtual->physical, or unused valid_bs, total_bs, BS_UPPER: tl.constexpr, + HAS_V2P: tl.constexpr, ): offs = tl.arange(0, BS_UPPER) in_range = offs < total_bs valid = offs < valid_bs req = tl.load(req_pool_indices_ptr + offs, mask=valid, other=0) idx = tl.load(mamba_map_ptr + req, mask=valid, other=0) + if HAS_V2P: + # Must gather before the padding sentinel, as the reference chain does. + idx = tl.load(v2p_ptr + idx, mask=valid, other=0) out_val = tl.where(valid, idx.to(tl.int32), -1) tl.store(out_ptr + offs, out_val, mask=in_range) # Preserve the reference chain's side effect: padded rows of the static @@ -49,6 +56,7 @@ def fused_replay_state_indices( out_state_indices: torch.Tensor, valid_bs: int, total_bs: int, + v2p: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Fill the captured replay state-indices buffer in one launch. @@ -58,9 +66,11 @@ def fused_replay_state_indices( the ``-1`` sentinel (mamba kernels skip ``state_idx < 0``) and their ``req_pool_indices`` entries are zeroed. - Callers must supply an identity virtual->physical mapping (the static - hybrid pool); the unified pool's allocator translate is not a flat table - gather and has to take the reference chain. + ``v2p`` is the mamba virtual->physical slot table, for a pool whose slot + ids are virtual (the unified memory pool). Pass None when the mapping + already yields physical slots (the static hybrid pool). The unified + allocator runs the mamba sub-pool at page_size 1, so its translate is a + plain table gather and folds into this launch. Returns the filled ``out_state_indices[:total_bs]`` view. """ @@ -68,8 +78,10 @@ def fused_replay_state_indices( req_pool_indices, mamba_index_mapping, out_state_indices, + v2p, valid_bs, total_bs, BS_UPPER=triton.next_power_of_2(total_bs), + HAS_V2P=v2p is not None, ) return out_state_indices[:total_bs] diff --git a/python/sglang/kernels/ops/memory/__init__.py b/python/sglang/kernels/ops/memory/__init__.py index c40c8ba53..7a4aaa0d7 100644 --- a/python/sglang/kernels/ops/memory/__init__.py +++ b/python/sglang/kernels/ops/memory/__init__.py @@ -18,6 +18,7 @@ _TRITON_KERNELS = [ ("virtual_slot", "alloc_bind_inplace"), ("virtual_slot", "free_unbind_inplace"), ("virtual_slot", "bind_inplace"), + ("virtual_slot", "write_loc_to_kernel_ids"), ] for _mod, _fn in _TRITON_KERNELS: register_kernel( diff --git a/python/sglang/kernels/ops/memory/virtual_slot.py b/python/sglang/kernels/ops/memory/virtual_slot.py index c62fc1ac9..e19c1dc69 100644 --- a/python/sglang/kernels/ops/memory/virtual_slot.py +++ b/python/sglang/kernels/ops/memory/virtual_slot.py @@ -15,6 +15,8 @@ from __future__ import annotations +from typing import Optional + import torch import triton import triton.language as tl @@ -189,3 +191,134 @@ def bind_inplace( return grid = (triton.cdiv(N, ALLOC_BIND_BLOCK),) bind_inplace_kernel[grid](v, p, v2p, p2v, N, BLOCK=ALLOC_BIND_BLOCK) + + +WRITE_LOC_BLOCK = 512 + + +@triton.jit +def write_loc_to_kernel_id_kernel( + loc_ptr, # in: [N] int64 — WIDENED virtual token ids + v2p_ptr, # in: [num_pages + 1] int64 — virtual->physical page table + out_ptr, # out: [N] int64 — kernel-facing ids + N, # runtime: live element count + W, # runtime: lanes to write; [N, W) get 0 + stride, # runtime: pool_page_size * kernel_page_multiplier + PAGE_SIZE: tl.constexpr, + DCP_SIZE: tl.constexpr, + DCP_RANK: tl.constexpr, + BLOCK: tl.constexpr, +): + """``kernel_id(t) = v2p[t // ps] * ps * mult + t % ps``, clamped at 0. + + Under DCP the incoming id is WIDENED: ``loc % dcp_size`` names its owner + and ``loc // dcp_size`` is the row. Ids this rank does not own resolve to + kernel id 0, the padding sink every write kernel skips. + + Triton truncates ``//`` toward zero where torch floors it, so a negative + loc is tested explicitly rather than left to the division; it resolves to + 0, as the torch path does. + + Writing ``W > N`` lanes fills ``[N, W)`` with 0, the padding sink, so a + caller may hand in a capture-stable buffer wider than this batch and have + the stale tail cleared in the same launch. + """ + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + in_range = offs < W + mask = offs < N + loc = tl.load(loc_ptr + offs, mask=mask, other=0).to(tl.int64) + + keep = mask & (loc >= 0) + if DCP_SIZE > 1: + keep = keep & ((loc % DCP_SIZE) == DCP_RANK) + loc = loc // DCP_SIZE + + page = loc // PAGE_SIZE if PAGE_SIZE > 1 else loc + offset = loc % PAGE_SIZE if PAGE_SIZE > 1 else 0 + # `keep` already excludes negatives, so the gather index is in range. + phys = tl.load(v2p_ptr + tl.where(keep, page, 0), mask=mask, other=0).to(tl.int64) + ids = tl.maximum(phys * stride + offset, 0) + tl.store(out_ptr + offs, tl.where(keep, ids, 0), mask=in_range) + + +def write_loc_to_kernel_ids( + *, + loc: torch.Tensor, + v2p: torch.Tensor, + page_size: int, + stride: int, + dcp_size: int = 1, + dcp_rank: int = 0, + out: Optional[torch.Tensor] = None, + out_width: Optional[int] = None, +) -> torch.Tensor: + """One launch for the whole write-loc conversion; see the kernel. + + ``out`` is written in place when given (a captured graph records the + gather against a fixed ``data_ptr``), else a fresh int64 tensor is + returned. Cuda-graph safe: no ``.item()``, no host sync, no allocation on + the ``out=`` path. + + ``out_width`` writes that many lanes rather than ``loc.numel()``, zeroing + the ones past the batch; pass the captured tier's width to clear a stale + tail here. + """ + N = int(loc.numel()) + # Flat-indexed as `ptr + offs`, so a strided view is mis-addressed. + assert loc.is_contiguous(), ( + f"write_loc_to_kernel_ids: loc must be contiguous, got shape " + f"{tuple(loc.shape)} stride {tuple(loc.stride())}" + ) + if out is None: + out = torch.empty_like(loc, dtype=torch.int64) + width = N if out_width is None else int(out_width) + assert out.dtype == torch.int64, ( + f"write_loc_to_kernel_ids: out dtype must be int64 (matches v2p), " + f"got {out.dtype}" + ) + if out_width is None: + # `out` mirrors `loc` whatever its shape; a 2-D page table is legal. + assert out.shape == loc.shape and out.is_contiguous(), ( + f"write_loc_to_kernel_ids: out shape {tuple(out.shape)} must match " + f"loc shape {tuple(loc.shape)}" + ) + else: + assert out.dim() == 1 and out.is_contiguous() and out.numel() >= width, ( + f"write_loc_to_kernel_ids: out_width needs a packed 1-D out of at " + f"least {width}, got {tuple(out.shape)}" + ) + assert width >= N, ( + f"write_loc_to_kernel_ids: out_width {width} is under the batch's " + f"{N} locs, which would drop live rows" + ) + if width == 0: + return out + if not loc.is_cuda: + # Pure-torch reference; the allocator's unit tests run on CPU. + big = loc.to(torch.int64) + keep = big >= 0 + if dcp_size > 1: + keep = keep & (big % dcp_size == dcp_rank) + big = torch.div(big, dcp_size, rounding_mode="floor") + page = torch.where(keep, torch.div(big, page_size, rounding_mode="floor"), 0) + offset = big % page_size if page_size > 1 else 0 + ids = (v2p[page] * stride + offset).clamp_(min=0) + out[:N].copy_(torch.where(keep, ids, torch.zeros_like(ids))) + if width > N: + out[N:width].zero_() + return out + grid = (triton.cdiv(width, WRITE_LOC_BLOCK),) + write_loc_to_kernel_id_kernel[grid]( + loc, + v2p, + out, + N, + width, + stride, + PAGE_SIZE=page_size, + DCP_SIZE=dcp_size, + DCP_RANK=dcp_rank, + BLOCK=WRITE_LOC_BLOCK, + ) + return out diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index 157a4fde5..c1d07e4b2 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -249,6 +249,12 @@ def handle_unified_memory_pool(server_args: Any) -> None: "not translate speculative verify indices to the unified " "pool's kernel-facing space yet." ) + assert not cfg.enable_two_batch_overlap, ( + "--enable-unified-memory does not support --enable-two-batch-overlap: " + "TBO's replay split hands each child a view without the pre-translate " + "write loc, so a captured decode replay raises. " + "TODO(ch-wan): carry out_cache_loc_virtual into the child view." + ) assert not (cfg.enable_hierarchical_cache or cfg.enable_lmcache), ( "--enable-unified-memory is not yet compatible with hierarchical / " "host-tiered KV cache (--enable-hierarchical-cache / --enable-lmcache): " diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index b6ff7b99f..cc45bce98 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/python/sglang/srt/batch_overlap/two_batch_overlap.py @@ -693,6 +693,11 @@ class TboForwardBatchPreparer: ) output_dict[key] = old_value[start_token_index:end_token_index] + if batch.out_cache_loc_virtual is not None: + output_dict["out_cache_loc_virtual"] = batch.out_cache_loc_virtual[ + start_token_index:end_token_index + ] + attention_tp_size = get_parallel().attn_tp_size _tbo_padded_len = ( (end_token_index - start_token_index - 1) // attention_tp_size + 1 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 fe494bafe..673e1cb3f 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -70,15 +70,11 @@ class MambaAttnBackendBase(AttentionBackend): # backends on runners without a real model_config. self._model_runner = model_runner self._mamba_chunk_size: Optional[int] = None - # Fused replay-prep state-indices fast path (fused_replay_state_indices): - # requires the static hybrid pool whose v2p translate is the identity — - # the unified pool overrides translate_mamba_indices with an allocator - # lookup that is not a flat table gather. + pool = self.req_to_token_pool self._fused_state_indices_ok = ( - str(self.device).startswith("cuda") - and isinstance(self.req_to_token_pool, HybridReqToTokenPool) - and type(self.req_to_token_pool).translate_mamba_indices - is HybridReqToTokenPool.translate_mamba_indices + torch.device(self.device).type == "cuda" + and isinstance(pool, HybridReqToTokenPool) + and pool.mamba_translate_is_fusable ) self.forward_metadata: ForwardMetadata = None self.state_indices_list = [] @@ -643,6 +639,7 @@ class MambaAttnBackendBase(AttentionBackend): out_state_indices=self.state_indices_list[bs - 1], valid_bs=bs - int(num_padding), total_bs=bs, + v2p=self.req_to_token_pool.mamba_v2p_table, ) else: # Make sure forward metadata is correctly handled for padding reqs diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index 699ae34a8..8046af61b 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -748,22 +748,16 @@ class TritonAttnBackend(AttentionBackend): def _fill_cuda_graph_write_locs( self, forward_batch: ForwardBatch, bs: int ) -> Optional[torch.Tensor]: - """Copy the cuda-graph WRITE loc into the capture-stable buffer and - return the ``[:n]`` view; no-op for non-unified pools. - - Runs BEFORE graph.replay() so it reads the live post-compaction v2p. - The capture batch is runner-built with zeros, which is safe because - slot 0 is the reserved sink in every id space. - """ + """Runs BEFORE graph.replay(), so it reads the live post-compaction + v2p; no-op for non-unified pools.""" + # The buffer exists only for a translating pool; return before naming it. if not self.kv_index_translator.is_translating: return None - out_cache_loc = forward_batch.out_cache_loc - n = out_cache_loc.shape[0] - # Zero the padded tail first: a smaller replay batch leaves [n:] holding - # stale ids that the captured store would write; send them to slot 0 (sink). - self.cuda_graph_out_cache_loc_full_physical[n:].zero_() - self.cuda_graph_out_cache_loc_full_physical[:n].copy_(out_cache_loc) - return self.cuda_graph_out_cache_loc_full_physical[:n] + return self.kv_index_translator.fill_capture_write_loc( + out=self.cuda_graph_out_cache_loc_full_physical, + forward_batch=forward_batch, + width=self.cuda_graph_out_cache_loc_full_physical.numel(), + ) def init_forward_metadata(self, forward_batch: ForwardBatch): """Init auxiliary variables for triton attention backend.""" diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 8bf2eff96..da171f782 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -766,19 +766,13 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): if self.kv_index_translator.is_translating and ( forward_mode.is_decode_or_idle() or forward_mode.is_target_verify() ): - out_cache_loc = forward_batch.out_cache_loc - n = out_cache_loc.shape[0] - dst = self.cuda_graph_out_cache_loc_kernel[:n] - dst.copy_(out_cache_loc) - # Replay-prep receives the RAW (unpadded) out_cache_loc - # (build_replay_fb_view), but the captured write kernel consumes the - # full captured tier of this buffer. Zero the tail so pad rows write - # to the sink (row 0) instead of stale kernel-facing locs left by - # earlier larger replays — a stale tail scatters pad-row garbage into - # live KV pages. Mirrors the runner's PaddingPolicy.ZERO on its own - # out_cache_loc slot. - self.cuda_graph_out_cache_loc_kernel[n:].zero_() - self._decode_kernel_loc = dst + # The captured kernel consumes the whole buffer, so the tail a + # shorter replay leaves must go to slot 0 rather than live pages. + self._decode_kernel_loc = self.kv_index_translator.fill_capture_write_loc( + out=self.cuda_graph_out_cache_loc_kernel, + forward_batch=forward_batch, + width=self.cuda_graph_out_cache_loc_kernel.numel(), + ) else: self._decode_kernel_loc = None diff --git a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py index fb964cfca..759c22682 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py +++ b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py @@ -331,10 +331,13 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): loc: torch.Tensor, *, out: Optional[torch.Tensor] = None, + out_width: Optional[int] = None, ) -> torch.Tensor: """Widened virtual WRITE loc -> kernel-facing id. DCP is rejected for this composite at argument validation, so it coincides with the read translate.""" - return self.full_attn_allocator.translate_write_loc_for_kernel(loc, out=out) + return self.full_attn_allocator.translate_write_loc_for_kernel( + loc, out=out, out_width=out_width + ) @property def swa_kernel_page_multiplier(self) -> int: diff --git a/python/sglang/srt/mem_cache/allocator/unified_mamba.py b/python/sglang/srt/mem_cache/allocator/unified_mamba.py index 5d78eabc5..bc2ac7c51 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_mamba.py +++ b/python/sglang/srt/mem_cache/allocator/unified_mamba.py @@ -260,9 +260,12 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): loc: torch.Tensor, *, out: Optional[torch.Tensor] = None, + out_width: Optional[int] = None, ) -> torch.Tensor: """Widened virtual WRITE loc -> DENSE id; see the sub-allocator's copy.""" - return self.full_attn_allocator.translate_write_loc_for_kernel(loc, out=out) + return self.full_attn_allocator.translate_write_loc_for_kernel( + loc, out=out, out_width=out_width + ) def translate_kv_indices_for_transfer( self, kv_indices: torch.Tensor diff --git a/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py b/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py index 9ee3d1685..a5286fd58 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py +++ b/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py @@ -43,6 +43,7 @@ from sglang.kernels.ops.memory.virtual_slot import ( alloc_bind_inplace, bind_inplace, free_unbind_inplace, + write_loc_to_kernel_ids, ) from sglang.srt.environ import envs from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator @@ -990,39 +991,47 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): clamp to kernel-facing id 0, the page-0 sink. int64 out; a consumer whose kernel ABI wants int32 narrows where it fills that buffer. """ - ps = self.pool_page_size - stride = ps * self.kernel_page_multiplier with record_function("MultiEndedAlloc.translate_kv_loc_for_kernel"): - pages = virt_tokens if ps == 1 else virt_tokens // ps - offsets = None if ps == 1 else virt_tokens % ps - if out is None: - phys = self.virtual_to_physical[pages] - ids = phys * stride if offsets is None else phys * stride + offsets - return ids.clamp_(min=0) + return self._translate_loc_fused(virt_tokens, dcp_size=1, out=out) + + def _translate_loc_fused( + self, + loc: torch.Tensor, + *, + dcp_size: int, + dcp_rank: int = 0, + out: Optional[torch.Tensor] = None, + out_width: Optional[int] = None, + ) -> torch.Tensor: + """One launch for the read and write conversions alike; see + `write_loc_to_kernel_ids`.""" + if out is not None: assert out.dtype == torch.int64, ( f"translate_kv_loc_for_kernel: out= dtype must be int64 (matches v2p), " f"got {out.dtype}" ) - assert out.shape == virt_tokens.shape, ( - f"translate_kv_loc_for_kernel: out= shape {tuple(out.shape)} must " - f"match virt_tokens shape {tuple(virt_tokens.shape)}" - ) - if pages.dtype != torch.int64: - pages = pages.to(torch.int64) - if pages is virt_tokens: - out.copy_(torch.take(self.virtual_to_physical, pages)) - else: - torch.take(self.virtual_to_physical, pages, out=out) - out.mul_(stride) - if offsets is not None: - out.add_(offsets) - return out.clamp_(min=0) + if out_width is None: + assert out.shape == loc.shape, ( + f"translate_kv_loc_for_kernel: out= shape {tuple(out.shape)} must " + f"match virt_tokens shape {tuple(loc.shape)}" + ) + return write_loc_to_kernel_ids( + loc=loc, + v2p=self.virtual_to_physical, + page_size=self.pool_page_size, + stride=self.pool_page_size * self.kernel_page_multiplier, + dcp_size=dcp_size, + dcp_rank=dcp_rank, + out=out, + out_width=out_width, + ) def translate_write_loc_for_kernel( self, widened_loc: torch.Tensor, *, out: Optional[torch.Tensor] = None, + out_width: Optional[int] = None, ) -> torch.Tensor: """Widened virtual WRITE loc (`out_cache_loc`) -> kernel-facing id. @@ -1032,16 +1041,14 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): """ parallel = get_parallel() dcp_size = parallel.attn_dcp_size if self.shards_under_dcp else 1 - if dcp_size == 1: - return self.translate_kv_loc_for_kernel(widened_loc, out=out) with record_function("MultiEndedAlloc.translate_write_loc_for_kernel"): - owned = (widened_loc % dcp_size) == parallel.attn_dcp_rank - dense = self.translate_kv_loc_for_kernel(widened_loc // dcp_size) - dense = torch.where(owned, dense, torch.zeros_like(dense)) - if out is not None: - out.copy_(dense) - return out - return dense + return self._translate_loc_fused( + widened_loc, + dcp_size=dcp_size, + dcp_rank=parallel.attn_dcp_rank, + out=out, + out_width=out_width, + ) # -- alloc -- diff --git a/python/sglang/srt/mem_cache/kv_index_translator.py b/python/sglang/srt/mem_cache/kv_index_translator.py index fe9ed87ae..2e1b97068 100644 --- a/python/sglang/srt/mem_cache/kv_index_translator.py +++ b/python/sglang/srt/mem_cache/kv_index_translator.py @@ -432,20 +432,55 @@ class KVIndexTranslator: def rebind_write_loc(self, forward_batch) -> None: """Phase 1 of the WRITE contract: translate the batch's write loc to - FULL-side kernel-facing ids exactly once, at ForwardBatch - construction. No-op on non-unified pools. + FULL-side kernel-facing ids, once, at ForwardBatch construction. REBIND, never mutate: the translate returns a FRESH tensor, so the ScheduleBatch's aliased tensor stays VIRTUAL for the radix / accept / - in-flight machinery that reads it. + in-flight machinery that reads it. The pre-translate tensor stays on + the batch for `fill_capture_write_loc`. """ self._index_table_memo = None if not self.is_translating or forward_batch.out_cache_loc is None: return + forward_batch.out_cache_loc_virtual = forward_batch.out_cache_loc forward_batch.out_cache_loc = self._translate_write_full( forward_batch.out_cache_loc ) + def fill_capture_write_loc( + self, + *, + out: torch.Tensor, + forward_batch, + width: Optional[int] = None, + ) -> Optional[torch.Tensor]: + """Translate this batch's WRITE loc straight into ``out``, a backend's + capture-stable buffer, and return the live ``[:n]`` view. One launch + fills the live prefix and clears the tail a shorter replay leaves; + None when this pool needs no translation. + + Must run at metadata-init time: `out` is reused every step, so filling + it sooner would race a still-pending previous step under overlap + scheduling. + """ + if not self.is_translating: + return None + virtual = forward_batch.out_cache_loc_virtual + if virtual is None: + loc = forward_batch.out_cache_loc + if loc is None: + return None + # The runner builds the capture batch outside `init_new`, so no + # rebind marked its virtual source; bake it holding sink ids. + width = int(loc.numel()) if width is None else int(width) + out[:width].zero_() + return out[: int(loc.numel())] + n = int(virtual.numel()) + width = n if width is None else int(width) + buf = out[:width] + self._translate_write_full(virtual, out=buf, out_width=width) + return buf[:n] + def sliding_window_write_loc_for( self, out_cache_loc: Optional[torch.Tensor] ) -> Optional[torch.Tensor]: diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 55c46a742..5607bcd07 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1385,6 +1385,28 @@ class HybridReqToTokenPool(ReqToTokenPool): def get_mamba_indices(self, req_indices: torch.Tensor) -> torch.Tensor: return self.req_index_to_mamba_index_mapping[req_indices] + @property + def mamba_v2p_table(self) -> Optional[torch.Tensor]: + """The mamba virtual->physical slot table, or None when the ids this + pool hands out are already physical.""" + return None + + @property + def mamba_translate_is_fusable(self) -> bool: + """Whether `fused_replay_state_indices` can reproduce this pool's + `translate_mamba_indices` in its own launch. + + The kernel expresses exactly two shapes: the identity, and one gather + through `mamba_v2p_table`. A subclass that replaces the translate with + anything else is excluded here rather than silently mis-served. + """ + if self.mamba_v2p_table is not None: + return True + return ( + type(self).translate_mamba_indices + is HybridReqToTokenPool.translate_mamba_indices + ) + def translate_mamba_indices(self, mamba_indices: torch.Tensor) -> torch.Tensor: """Virtual->physical mamba-slot translate. Identity for a static pool (slots are physical); UnifiedHybridReqToTokenPool overrides it for the diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index 0b7870046..5699a0f99 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -1107,6 +1107,15 @@ class UnifiedHybridReqToTokenPool(HybridReqToTokenPool): ) ) + @property + def mamba_v2p_table(self) -> Optional[torch.Tensor]: + """This pool's ids ARE virtual; page_size is 1, so the translate is the + plain gather this table serves, which keeps `mamba_translate_is_fusable` + true despite the override.""" + if self.mamba_allocator is None: + return None + return self.mamba_allocator.virtual_to_physical + def translate_mamba_indices(self, virtual_ids: torch.Tensor) -> torch.Tensor: """Virtual mamba ids -> physical slot ids.""" return self.mamba_allocator.translate(virtual_ids).to(torch.int32) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 2191348fa..45fc4effa 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -414,6 +414,9 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): # The original sequence length without being chunked. Qwen-1M related. orig_seq_lens: Optional[torch.Tensor] = None + # The write loc before `rebind_write_loc` replaced it with kernel-facing + # ids; a backend re-derives from it into its capture-stable buffer. + out_cache_loc_virtual: Optional[torch.Tensor] = None # DSV4-NPU only: per-pool slot bundle from DSV4NPUTokenToKVPoolAllocator, # consumed by the Ascend backend for PA_ND block tables. None elsewhere. out_cache_loc_dsv4: Optional[DSV4OutCacheLoc] = None @@ -1836,6 +1839,8 @@ def build_inner_fb_view( seq_lens_cpu=forward_batch.seq_lens_cpu, encoder_lens=encoder_lens, out_cache_loc=getattr(forward_batch, "out_cache_loc", None), + # A caller may hand in another view that does not carry this field. + out_cache_loc_virtual=getattr(forward_batch, "out_cache_loc_virtual", None), out_cache_loc_dsv4=getattr(forward_batch, "out_cache_loc_dsv4", None), spec_info=forward_batch.spec_info, ) diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index 2287ccaef..cef7ac162 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -192,6 +192,7 @@ def build_replay_fb_view( num_padding=bs - raw_bs, encoder_lens=buffers.encoder_lens[:bs] if is_encoder_decoder else None, out_cache_loc=getattr(forward_batch, "out_cache_loc", None), + out_cache_loc_virtual=forward_batch.out_cache_loc_virtual, out_cache_loc_dsv4=getattr(forward_batch, "out_cache_loc_dsv4", None), # The mamba-track registry slot (VIRTUAL ids) is the v2p translate SOURCE # for the backend, which copies the result into its own static buffer and diff --git a/test/registered/dcp/test_trtllm_mla_family_dcp_metadata.py b/test/registered/dcp/test_trtllm_mla_family_dcp_metadata.py index 252c4dad8..94de9c3db 100644 --- a/test/registered/dcp/test_trtllm_mla_family_dcp_metadata.py +++ b/test/registered/dcp/test_trtllm_mla_family_dcp_metadata.py @@ -432,6 +432,89 @@ class TestFusedFp8WriteGate(CustomTestCase): ) +@unittest.skipUnless(torch.cuda.is_available(), "the fused translate is Triton") +class TestFusedWriteLocTranslateCuda(CustomTestCase): + """The fused write-loc translate must agree with its own CPU branch. + + The two implementations must not drift: Triton truncates division toward + zero where torch floors it, so a negative loc and a tombstoned v2p row are + where a divergence would appear -- and the CPU branch is all the CPU suites + ever exercise. + """ + + def test_cuda_matches_the_cpu_branch(self): + from sglang.kernels.ops.memory.virtual_slot import write_loc_to_kernel_ids + + for page_size in (1, 64): + span = page_size * 4 + loc = torch.tensor( + [-1, 0, 1, page_size, span, span + 1, 2 * span + 3, 5 * page_size], + dtype=torch.int64, + ) + # Size the table past the highest page any loc can name, then + # scramble it and tombstone one row (-1), so neither a dropped + # clamp nor a skipped gather can coincide with the right answer. + num_pages = int(loc.max()) // page_size + 2 + v2p = torch.tensor( + [(5 * i + 2) % num_pages for i in range(num_pages)] + [-1], + dtype=torch.int64, + ) + v2p[1] = -1 + for dcp_size, dcp_rank in ((1, 0), (2, 1), (4, 2)): + kw = dict( + page_size=page_size, + stride=page_size * 3, + dcp_size=dcp_size, + dcp_rank=dcp_rank, + ) + cpu = write_loc_to_kernel_ids(loc=loc, v2p=v2p, **kw) + gpu = write_loc_to_kernel_ids(loc=loc.cuda(), v2p=v2p.cuda(), **kw) + self.assertEqual( + cpu.tolist(), + gpu.cpu().tolist(), + f"ps={page_size} dcp_size={dcp_size} rank={dcp_rank}", + ) + + def test_wide_out_clears_the_stale_tail(self): + """`out_width` past the batch must zero the tail in the same launch. + + This is what lets a backend hand in its whole capture-stable buffer: + a shorter replay leaves stale kernel-facing ids past the batch, and the + captured write kernel consumes the full buffer, so an uncleared tail + scatters pad rows into live KV pages. + """ + from sglang.kernels.ops.memory.virtual_slot import write_loc_to_kernel_ids + + v2p = torch.tensor([2, 5, 1, 3, 4], dtype=torch.int64, device="cuda") + loc = torch.tensor([0, 64, 128], dtype=torch.int64, device="cuda") + width = 8 + # Poison the whole buffer so an unwritten or uncleared cell is visible. + buf = torch.full((width,), -999, dtype=torch.int64, device="cuda") + write_loc_to_kernel_ids( + loc=loc, v2p=v2p, page_size=64, stride=64 * 2, out=buf, out_width=width + ) + self.assertEqual(buf[:3].tolist(), [2 * 128, 5 * 128, 1 * 128]) + self.assertEqual(buf[3:].tolist(), [0] * (width - 3)) + + # And it must agree with the narrow call on the live prefix. + narrow = write_loc_to_kernel_ids(loc=loc, v2p=v2p, page_size=64, stride=64 * 2) + self.assertEqual(narrow.tolist(), buf[:3].tolist()) + + def test_out_is_written_in_place(self): + # The captured decode path hands in a capture-stable buffer; rebinding + # instead of filling it would leave the graph on a stale pointer. + from sglang.kernels.ops.memory.virtual_slot import write_loc_to_kernel_ids + + v2p = torch.tensor([2, 5, -1, 3], dtype=torch.int64, device="cuda") + loc = torch.tensor([0, 64, 128, 192], dtype=torch.int64, device="cuda") + dst = torch.full_like(loc, -7) + ret = write_loc_to_kernel_ids( + loc=loc, v2p=v2p, page_size=64, stride=64 * 2, out=dst + ) + self.assertIs(ret, dst) + self.assertEqual(dst.tolist(), [2 * 128, 5 * 128, 0, 3 * 128]) + + class TestDcpDecodeLayout(CustomTestCase): """Rank-local length math the decode page table above is built from.""" diff --git a/test/registered/kernels/ops/mamba/test_fused_replay_state_indices.py b/test/registered/kernels/ops/mamba/test_fused_replay_state_indices.py index 93d059283..9410aeea6 100644 --- a/test/registered/kernels/ops/mamba/test_fused_replay_state_indices.py +++ b/test/registered/kernels/ops/mamba/test_fused_replay_state_indices.py @@ -49,11 +49,15 @@ def _reference_chain( out_buf: torch.Tensor, valid_bs: int, total_bs: int, + v2p: torch.Tensor = None, ) -> None: """Replicates the _replay_metadata reference ops, in order, in place.""" req_pool_indices[valid_bs:total_bs] = 0 mamba_indices = mapping[req_pool_indices[:total_bs]] - # static pool: _translate_mamba_indices is the identity + if v2p is not None: + # Unified pool: virtual->physical slot translate, and it runs BEFORE + # the padding sentinel -- captured kernels read physical ids. + mamba_indices = v2p[mamba_indices].to(torch.int32) mamba_indices[valid_bs:] = -1 out_buf[: len(mamba_indices)].copy_(mamba_indices) @@ -139,6 +143,79 @@ class TestFusedReplayStateIndices(CustomTestCase): f"{name} guard tail clobbered ({case}): {buf[total_bs:].tolist()}", ) + def _run_v2p_case(self, total_bs: int, num_padding: int, seed: int) -> None: + """Same equivalence, with the unified pool's virtual slot ids. + + The mapping yields VIRTUAL slots there and the kernel folds the v2p + gather in, so the two must still agree element for element. + """ + device = torch.device("cuda") + gen = torch.Generator(device="cpu").manual_seed(seed + 1000) + valid_bs = total_bs - num_padding + + req_pool = torch.randint( + 0, _REQ_POOL_SIZE, (total_bs + _GUARD,), generator=gen, dtype=torch.int64 + ) + req_pool[total_bs:] = _GUARD_SENTINEL + mapping = torch.randint( + 0, _MAMBA_POOL_SIZE, (_REQ_POOL_SIZE,), generator=gen, dtype=torch.int32 + ) + # Scrambled table with a tombstone, so a skipped gather cannot pass. + v2p = torch.randperm(_MAMBA_POOL_SIZE + 1, generator=gen).to(torch.int64) + v2p[7] = -1 + out = torch.full((total_bs + _GUARD,), _OUT_POISON, dtype=torch.int32) + + req_ref, req_fused = req_pool.clone().to(device), req_pool.clone().to(device) + out_ref, out_fused = out.clone().to(device), out.clone().to(device) + mapping_d, v2p_d = mapping.to(device), v2p.to(device) + + _reference_chain( + req_pool_indices=req_ref, + mapping=mapping_d, + out_buf=out_ref, + valid_bs=valid_bs, + total_bs=total_bs, + v2p=v2p_d, + ) + returned = fused_replay_state_indices( + req_pool_indices=req_fused, + mamba_index_mapping=mapping_d, + out_state_indices=out_fused, + valid_bs=valid_bs, + total_bs=total_bs, + v2p=v2p_d, + ) + case = f"v2p {total_bs=} {num_padding=} {seed=}" + self.assertTrue( + torch.equal(out_ref[:total_bs], out_fused[:total_bs]), + f"state indices mismatch ({case}):\n" + f" ref {out_ref[:total_bs].tolist()}\n" + f" fused {out_fused[:total_bs].tolist()}", + ) + self.assertTrue(torch.equal(returned, out_fused[:total_bs]), case) + self.assertTrue( + torch.equal(req_pool_ref_tail := req_ref[total_bs:], req_fused[total_bs:]), + f"guard tail diverged ({case}): {req_pool_ref_tail.tolist()}", + ) + self.assertTrue( + (out_fused[total_bs:] == _OUT_POISON).all(), + f"out guard tail clobbered ({case})", + ) + + def test_v2p_matrix(self): + for total_bs in (1, 7, 32, 33): + paddings = sorted( + {0, 1, total_bs // 2, total_bs - 1} & set(range(total_bs)) + ) + for num_padding in paddings: + for seed in (0, 1): + with self.subTest( + total_bs=total_bs, num_padding=num_padding, seed=seed + ): + self._run_v2p_case( + total_bs=total_bs, num_padding=num_padding, seed=seed + ) + def test_matrix(self): # Non-power-of-two sizes (7, 33) exercise the BS_UPPER in_range mask; # num_padding sweeps none / one / half / all-but-one padded rows. diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py index d7abd02a2..b7cac8288 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -3100,5 +3100,82 @@ class TestDcpWidening(unittest.TestCase): self.assertLess(base_cost * full_entry, mamba_bytes * dcp_size) +class TestFusedWriteLocTranslate(unittest.TestCase): + """`write_loc_to_kernel_ids` must equal the arithmetic it stands for. + + Nothing downstream can tell the two apart except by value, so the reference + here is the definition rather than a recorded expectation. Triton truncates + division toward zero where torch floors it, so the negative-loc and + tombstoned-page cases are the ones that matter. + """ + + def _reference(self, loc, v2p, page_size, stride, dcp_size, dcp_rank): + out = [] + for raw in loc.tolist(): + if raw < 0 or (dcp_size > 1 and raw % dcp_size != dcp_rank): + out.append(0) + continue + collapsed = raw // dcp_size + page = collapsed // page_size + offset = collapsed % page_size if page_size > 1 else 0 + out.append(max(int(v2p[page]) * stride + offset, 0)) + return out + + def _check(self, *, page_size, multiplier, dcp_size, dcp_rank, device): + from sglang.kernels.ops.memory.virtual_slot import write_loc_to_kernel_ids + + span = page_size * dcp_size + locs = [0, 1, span - 1, span, 2 * span + 3, 5 * span + dcp_rank, -1] + locs += [3 * span + dcp_rank] # lands on the tombstoned page + loc = torch.tensor(locs, dtype=torch.int64, device=device) + # Table sized past the highest page any loc can name, scrambled, with + # one tombstone (-1) so a missing clamp shows up. + num_pages = max(locs) // (page_size * dcp_size) + 2 + v2p = torch.tensor( + [(5 * i + 2) % num_pages for i in range(num_pages)] + [-1], + dtype=torch.int64, + device=device, + ) + v2p[min(3, num_pages - 1)] = -1 + stride = page_size * multiplier + + got = write_loc_to_kernel_ids( + loc=loc, + v2p=v2p, + page_size=page_size, + stride=stride, + dcp_size=dcp_size, + dcp_rank=dcp_rank, + ) + want = self._reference( + loc.cpu(), v2p.cpu(), page_size, stride, dcp_size, dcp_rank + ) + self.assertEqual(got.tolist(), want, f"ps={page_size} dcp={dcp_size}") + # `out=` must write in place and agree (the cuda-graph-stable path). + dst = torch.full_like(loc, -7) + ret = write_loc_to_kernel_ids( + loc=loc, + v2p=v2p, + page_size=page_size, + stride=stride, + dcp_size=dcp_size, + dcp_rank=dcp_rank, + out=dst, + ) + self.assertIs(ret, dst) + self.assertEqual(dst.tolist(), want) + + def test_matches_reference_on_cpu(self): + for page_size in (1, 64): + for dcp_size, dcp_rank in ((1, 0), (2, 1), (4, 2)): + self._check( + page_size=page_size, + multiplier=7, + dcp_size=dcp_size, + dcp_rank=dcp_rank, + device="cpu", + ) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/server_args/test_unified_prefill_cuda_graph_gate.py b/test/registered/unit/server_args/test_unified_prefill_cuda_graph_gate.py index 0cde03f33..c951f5747 100644 --- a/test/registered/unit/server_args/test_unified_prefill_cuda_graph_gate.py +++ b/test/registered/unit/server_args/test_unified_prefill_cuda_graph_gate.py @@ -55,6 +55,7 @@ def _run_handler(*, prefill_backend, explicit): "speculative_eagle_topk": None, "enable_hierarchical_cache": False, "enable_lmcache": False, + "enable_two_batch_overlap": False, "dcp_size": 1, "cuda_graph_config": cg, "cuda_graph_backend_prefill": prefill_backend if explicit else None, diff --git a/test/registered/unit/server_args/test_unified_tbo_gate.py b/test/registered/unit/server_args/test_unified_tbo_gate.py new file mode 100644 index 000000000..d3c290ad4 --- /dev/null +++ b/test/registered/unit/server_args/test_unified_tbo_gate.py @@ -0,0 +1,69 @@ +# Copyright 2023-2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""`--enable-unified-memory` refuses `--enable-two-batch-overlap`. + +BUG REGRESSION. The combination launches and captures fine, then dies in the +forward path on the first captured decode replay. Nothing else rejects the +pair, so without this gate a running server crashes mid-serving. + + python -m pytest test/registered/unit/server_args/test_unified_tbo_gate.py -v +""" + +import unittest +from types import SimpleNamespace + +from sglang.srt.arg_groups.kv_cache_hook import handle_unified_memory_pool +from sglang.srt.model_executor.cuda_graph_config import Backend +from sglang.srt.server_args import ServerArgs +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _run_handler(*, unified, tbo): + """Run just `handle_unified_memory_pool` over a minimal stand-in.""" + sa = ServerArgs.__new__(ServerArgs) + for name, value in { + "enable_unified_memory": unified, + "enable_two_batch_overlap": tbo, + "disaggregation_mode": "null", + "speculative_algorithm": None, + "speculative_eagle_topk": None, + "enable_hierarchical_cache": False, + "enable_lmcache": False, + "dcp_size": 1, + "cuda_graph_config": SimpleNamespace( + prefill=SimpleNamespace(backend=Backend.DISABLED), + decode=SimpleNamespace(backend=Backend.FULL), + ), + "cuda_graph_backend_prefill": Backend.DISABLED, + }.items(): + object.__setattr__(sa, name, value) + handle_unified_memory_pool(sa) + + +class TestUnifiedTboGate(unittest.TestCase): + def test_tbo_with_unified_memory_is_refused(self): + with self.assertRaises(AssertionError) as ctx: + _run_handler(unified=True, tbo=True) + self.assertIn("two-batch-overlap", str(ctx.exception)) + + def test_gate_fires_only_on_the_pair(self): + """An inverted condition here would reject every unified launch.""" + _run_handler(unified=True, tbo=False) + _run_handler(unified=False, tbo=True) + + +if __name__ == "__main__": + unittest.main()