[P/D Disagg] DeepSeek-V4 decode radix cache with MTP support

python/sglang/srt/environ.py

@@ -938,6 +938,8 @@ class Envs:
     # ====================================================================
     # DeepSeek V4
     SGLANG_OPT_DPSK_V4_RADIX = EnvBool(True)
+    SGLANG_EXPERIMENTAL_DSV4_DECODE_RADIX_CACHE = EnvBool(False)
+    SGLANG_DEBUG_DSV4_DECODE_RADIX_TRANSFER = EnvBool(False)
     SGLANG_OPT_USE_OLD_COMPRESSOR = EnvBool(False)
     SGLANG_OPT_USE_TRITON_SWA_PREPARE = EnvBool(True)
     SGLANG_OPT_USE_AITER_MHC_PRE = EnvBool(True)

python/sglang/srt/managers/scheduler_components/batch_result_processor.py

@@ -27,6 +27,7 @@
     maybe_cache_unfinished_req,
     release_kv_cache,
 )
+from sglang.srt.mem_cache.radix_cache import RadixKey
 from sglang.srt.runtime_context import get_server_args
 from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
 from sglang.srt.state_capturer.indexer_topk import get_global_indexer_capturer
@@ -146,6 +147,79 @@ def _maybe_collect_routed_experts(self, req: Req):
                 req.routed_experts_start_len,
             )
 
+    def _maybe_insert_dsv4_decode_radix_prompt(self, req: Req):
+        if not getattr(req, "dsv4_decode_radix_cache_prompt_once", False):
+            return
+
+        # process_batch_result_prebuilt() runs before the first real decode
+        # forward and the batch still references request-owned prompt pages via
+        # out_cache_loc. Insert only after a decode forward completes. For DSV4
+        # we donate the prompt pages to radix cache; protect the prompt snapshot
+        # before insertion so the generic overlap path does not free those pages
+        # again, and later request release skips the donated prefix. Keep the key
+        # bounded to the prefill-committed prompt snapshot so MTP accepted/draft
+        # deltas never enter the tree.
+        req.dsv4_decode_radix_cache_prompt_once = False
+        req.allow_radix_cache_insert_once = True
+        prompt_len = getattr(req, "dsv4_decode_radix_cache_prompt_len", None)
+        if prompt_len is None:
+            maybe_cache_unfinished_req(req, self.tree_cache)
+            return
+
+        page_size = self.tree_cache.page_size
+        # prompt-once flag was armed (decode.py), so get_fill_ids() already
+        # returns exactly the committed prompt snapshot.
+        prompt_fill_ids = req.get_fill_ids()
+        radix_key_len = len(
+            RadixKey(
+                prompt_fill_ids,
+                req.extra_key,
+                is_bigram=self.tree_cache.is_eagle,
+            ).page_aligned(page_size)
+        )
+        if radix_key_len <= 0:
+            req.allow_radix_cache_insert_once = False
+            return
+
+        old_cache_protected_len = req.cache_protected_len
+        old_swa_evicted_seqlen = req.kv.swa_evicted_seqlen
+        old_force_leaf_creation = getattr(req, "force_radix_leaf_creation", False)
+
+        # DSV4 prompt donation only needs a full-attention radix leaf. Mark the
+        # whole donated key as SWA-evicted so the SWA component stays tombstoned,
+        # but force full leaf creation so later matches can reuse the full prefix.
+        # Do not pre-protect the whole radix key here: when the prefix already
+        # exists, the generic overlap path must free this request's duplicate
+        # prompt pages and repoint it to the existing radix leaf.
+        req.kv.swa_evicted_seqlen = radix_key_len
+        req.force_radix_leaf_creation = True
+        try:
+            maybe_cache_unfinished_req(req, self.tree_cache)
+            protected_len = min(req.cache_protected_len, radix_key_len)
+            if protected_len > 0 and hasattr(
+                self.token_to_kv_pool_allocator, "free_swa"
+            ):
+                donated_full_indices = self.req_to_token_pool.req_to_token[
+                    req.req_pool_idx, :protected_len
+                ]
+                # The donated DSV4 prompt leaf is full-only. Release any
+                # request-private SWA tail mapped from those full indices; the
+                # full pages stay owned by radix cache.
+                self.token_to_kv_pool_allocator.free_swa(donated_full_indices)
+            if envs.SGLANG_DEBUG_DSV4_DECODE_RADIX_TRANSFER.get():
+                logger.info(
+                    "DSV4 decode radix prompt inserted: rid=%s "
+                    "prompt_len=%d radix_key_len=%d",
+                    req.rid,
+                    prompt_len,
+                    radix_key_len,
+                )
+        finally:
+            req.kv.swa_evicted_seqlen = old_swa_evicted_seqlen
+            req.force_radix_leaf_creation = old_force_leaf_creation
+            if req.cache_protected_len < old_cache_protected_len:
+                req.cache_protected_len = old_cache_protected_len
+
     def _maybe_collect_indexer_topk(self, req: Req):
         capturer = get_global_indexer_capturer()
         if capturer is None:
@@ -773,6 +847,8 @@ def process_batch_result_decode(
                 # And all the over-allocated tokens will be freed in `release_kv_cache`.
                 continue
 
+            self._maybe_insert_dsv4_decode_radix_prompt(req)
+
             # next_token_id is a per-req list: 1 token for non-spec, the verified
             # run for spec (already grammar-truncated in _resolve_spec_v2_tokens).
             next_token_id = next_token_ids[i]

python/sglang/srt/mem_cache/allocator/swa.py

@@ -10,8 +10,6 @@
 _is_npu = is_npu()
 
 if _is_npu:
-    import torch_npu
-
     from sglang.srt.hardware_backend.npu.allocator_npu import (
         NPUPagedTokenToKVPoolAllocator,
     )
@@ -112,15 +110,6 @@ def full_available_size(self):
     def swa_available_size(self):
         return self.swa_attn_allocator.available_size()
 
-    # Slot-conservation views for the leak invariant. On the non-shared allocator
-    # the static budget IS physical (conserve == physical); the shared composite
-    # overrides these with the static-cap view.
-    def _conserve_full_available_size(self):
-        return self.full_available_size()
-
-    def _conserve_swa_available_size(self):
-        return self.swa_available_size()
-
     @property
     def size(self):
         return min(self._size_full, self._size_swa)
@@ -160,17 +149,14 @@ def alloc(self, need_size: int):
         assert alloc_full_indices is not None
         assert alloc_swa_indices is not None
 
-        self.set_full_to_swa_mapping(alloc_full_indices, alloc_swa_indices)
+        if _is_npu:
+            self.full_to_swa_index_mapping[alloc_full_indices.to(torch.int64)] = (
+                alloc_swa_indices.to(torch.int64)
+            )
+        else:
+            self.full_to_swa_index_mapping[alloc_full_indices] = alloc_swa_indices
         return alloc_full_indices
 
-    def new_pages_available(self, num_full_pages: int, num_swa_pages: int) -> bool:
-        return (
-            num_full_pages
-            <= self.full_attn_allocator.available_size() // self.page_size
-            and num_swa_pages
-            <= self.swa_attn_allocator.available_size() // self.page_size
-        )
-
     def alloc_extend(
         self,
         prefix_lens: torch.Tensor,
@@ -185,7 +171,9 @@ def alloc_extend(
         num_new_pages = get_num_new_pages(
             seq_lens=seq_lens_cpu, page_size=self.page_size, prefix_lens=prefix_lens_cpu
         )
-        if not self.new_pages_available(num_new_pages, num_new_pages):
+        if num_new_pages > self.full_attn_allocator.available_size() // self.page_size:
+            return None
+        if num_new_pages > self.swa_attn_allocator.available_size() // self.page_size:
             return None
 
         swa_last_loc = self.translate_loc_from_full_to_swa(last_loc)
@@ -211,7 +199,12 @@ def alloc_extend(
         assert alloc_full_indices is not None
         assert alloc_swa_indices is not None
 
-        self.set_full_to_swa_mapping(alloc_full_indices, alloc_swa_indices)
+        if _is_npu:
+            self.full_to_swa_index_mapping[alloc_full_indices.to(torch.int64)] = (
+                alloc_swa_indices.to(torch.int64)
+            )
+        else:
+            self.full_to_swa_index_mapping[alloc_full_indices] = alloc_swa_indices
 
         return alloc_full_indices
 
@@ -240,7 +233,9 @@ def alloc_extend_swa_tail(
             seq_lens=seq_lens_cpu, page_size=self.page_size, prefix_lens=prefix_lens_cpu
         )
         num_swa_pages = (swa_tail_len + self.page_size - 1) // self.page_size
-        if not self.new_pages_available(num_full_pages, num_swa_pages):
+        if num_full_pages > self.full_attn_allocator.available_size() // self.page_size:
+            return None
+        if num_swa_pages > self.swa_attn_allocator.available_size() // self.page_size:
             return None
 
         alloc_full_indices = self.full_attn_allocator.alloc_extend(
@@ -250,7 +245,6 @@ def alloc_extend_swa_tail(
             seq_lens_cpu,
             last_loc,
             extend_num_tokens,
-            num_new_pages=num_full_pages,
         )
         assert alloc_full_indices is not None
 
@@ -271,13 +265,12 @@ def alloc_extend_swa_tail(
             swa_seq_lens_cpu,
             swa_last_loc,
             swa_tail_len,
-            num_new_pages=num_swa_pages,
         )
         assert alloc_swa_indices is not None
 
-        self.set_full_to_swa_mapping(
-            alloc_full_indices[-swa_tail_len:], alloc_swa_indices
-        )
+        self.full_to_swa_index_mapping[
+            alloc_full_indices[-swa_tail_len:].to(torch.int64)
+        ] = alloc_swa_indices.to(torch.int64)
         if swa_tail_len < extend_num_tokens:
             self.full_to_swa_index_mapping[
                 alloc_full_indices[:-swa_tail_len].to(torch.int64)
@@ -304,24 +297,43 @@ def alloc_decode(
             return None
 
         if _is_npu:
-            indices_2d = alloc_full_indices.to(torch.int64).unsqueeze(-1)
-            torch_npu.npu_scatter_nd_update_(
-                self.full_to_swa_index_mapping,
-                indices_2d,
-                alloc_swa_indices.to(torch.int64),
+            self.full_to_swa_index_mapping[alloc_full_indices.to(torch.int64)] = (
+                alloc_swa_indices.to(torch.int64)
             )
         else:
             self.full_to_swa_index_mapping[alloc_full_indices] = alloc_swa_indices
 
         return alloc_full_indices
 
+    def _filter_unreleased_indices(
+        self, allocator, indices: torch.Tensor
+    ) -> torch.Tensor:
+        if indices.numel() == 0 or self.page_size == 1:
+            return indices
+
+        unavailable_pages = []
+        if allocator.free_pages.numel() > 0:
+            unavailable_pages.append(allocator.free_pages)
+        if allocator.release_pages.numel() > 0:
+            unavailable_pages.append(allocator.release_pages)
+        if not unavailable_pages:
+            return indices
+
+        pages = indices // self.page_size
+        unavailable_pages = torch.cat(unavailable_pages)
+        return indices[~torch.isin(pages, unavailable_pages)]
+
     def free(self, free_index: torch.Tensor):
         if free_index.numel() == 0:
             return
 
         # NOTE: the API is not idempotent.
         if self.is_not_in_free_group:
-            self.full_attn_allocator.free(free_index)
+            full_free_index = self._filter_unreleased_indices(
+                self.full_attn_allocator, free_index
+            )
+            if full_free_index.numel() > 0:
+                self.full_attn_allocator.free(full_free_index)
             self.free_swa(free_index)
         else:
             self.free_group.append(free_index)
@@ -340,30 +352,22 @@ def set_full_to_swa_mapping(
         if full_indices.numel() == 0:
             return
         assert full_indices.numel() == swa_indices.numel()
-        full_indices = full_indices.to(torch.int64)
-        swa_indices = swa_indices.to(self.full_to_swa_index_mapping.dtype)
-        self.full_to_swa_index_mapping[full_indices] = swa_indices
-
-    def free_swa(self, free_index: torch.Tensor):
-        if free_index.numel() == 0:
-            return
-
-        if self.page_size == 1:
-            mapping_indices = free_index
+        if _is_npu:
+            self.full_to_swa_index_mapping[full_indices.to(torch.int64)] = (
+                swa_indices.to(torch.int64)
+            )
         else:
-            mapping_indices = self._expand_to_full_pages(free_index)
+            self.full_to_swa_index_mapping[full_indices] = swa_indices
 
-        swa_indices = self.full_to_swa_index_mapping[mapping_indices]
+    def free_swa(self, free_index: torch.Tensor):
+        swa_indices = self.full_to_swa_index_mapping[free_index]
         swa_indices = swa_indices[swa_indices > 0]
-        self.swa_attn_allocator.free(swa_indices)
-        self.full_to_swa_index_mapping[mapping_indices] = 0
-
-    def _expand_to_full_pages(self, indices: torch.Tensor) -> torch.Tensor:
-        pages = torch.unique(indices // self.page_size)
-        page_offsets = torch.arange(
-            self.page_size, dtype=indices.dtype, device=indices.device
+        swa_indices = self._filter_unreleased_indices(
+            self.swa_attn_allocator, swa_indices
         )
-        return (pages[:, None] * self.page_size + page_offsets[None, :]).reshape(-1)
+        if swa_indices.numel() > 0:
+            self.swa_attn_allocator.free(swa_indices)
+        self.full_to_swa_index_mapping[free_index] = 0
 
     def backup_state(self):
         return [
@@ -452,73 +456,3 @@ def __init__(
         self._kvcache = kvcache
         self.swa_attn_allocator.clear()
         self._kvcache.register_mapping(self.full_to_swa_index_mapping)
-
-    def available_size(self):
-        return self.swa_attn_allocator.available_size()
-
-    def full_available_size(self):
-        return self.swa_attn_allocator.available_size()
-
-    def swa_available_size(self):
-        return self.swa_attn_allocator.available_size()
-
-    def new_pages_available(self, num_full_pages: int, num_swa_pages: int) -> bool:
-        avail = self.swa_attn_allocator.available_size() // self.page_size
-        return num_full_pages <= avail and num_swa_pages <= avail
-
-    def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor):
-        return kv_indices
-
-    def alloc(self, need_size: int):
-        assert self.page_size == 1
-        return self.swa_attn_allocator.alloc(need_size)
-
-    def alloc_extend(self, *args, **kwargs):
-        raise NotImplementedError(
-            "PureSWATokenToKVPoolAllocator does not support page_size > 1."
-        )
-
-    def alloc_decode(self, *args, **kwargs):
-        raise NotImplementedError(
-            "PureSWATokenToKVPoolAllocator does not support page_size > 1."
-        )
-
-    def alloc_extend_swa_tail(self, *args, **kwargs):
-        raise NotImplementedError(
-            "PureSWATokenToKVPoolAllocator does not support page_size > 1."
-        )
-
-    def free(self, free_index: torch.Tensor):
-        if free_index.numel() == 0:
-            return
-        if self.is_not_in_free_group:
-            self.swa_attn_allocator.free(free_index[free_index > 0])
-        else:
-            self.free_group.append(free_index)
-        assert self.swa_attn_allocator.available_size() <= self.swa_attn_allocator.size
-
-    def free_swa(self, free_index: torch.Tensor):
-        if free_index.numel() == 0:
-            return
-        self.swa_attn_allocator.free(free_index[free_index > 0])
-
-    def free_group_begin(self):
-        self.is_not_in_free_group = False
-        self.free_group = []
-
-    def free_group_end(self):
-        self.is_not_in_free_group = True
-        if self.free_group:
-            self.free(torch.cat(self.free_group))
-        self.free_group = []
-
-    def backup_state(self):
-        return self.swa_attn_allocator.backup_state()
-
-    def restore_state(self, state):
-        self.swa_attn_allocator.restore_state(state)
-
-    def clear(self):
-        self.swa_attn_allocator.clear()
-        self.is_not_in_free_group = True
-        self.free_group = []