[AMD] Fix stale SWA ring buffer on radix prefix reuse for DeepSeek-V4 with unified_kv backend (#30339)

Co-authored-by: amd-danli103 <dan2.li@amd.com>

python/sglang/srt/managers/schedule_batch.py

@@ -1172,6 +1172,14 @@ class Req(ReqDllmMixin):
         if tree_cache is not None:
             if cow_mamba is None:
                 cow_mamba = tree_cache.supports_mamba()
+            # unified_kv SWA lives in a per-request ring that is not content-stable
+            # and never cached in the radix tree, so a reused prefix carries stale
+            # SWA. Cap the match by the trailing sliding window so it is re-prefilled
+            # into this request's ring. No-op for other layouts (returns 0).
+            reprefill_tail = tree_cache.swa_reprefill_tail_tokens()
+            if reprefill_tail:
+                capped = max(0, input_len - reprefill_tail)
+                key_limit = capped if key_limit is None else min(key_limit, capped)
             match_result = tree_cache.match_prefix(
                 MatchPrefixParams(
                     key=RadixKey(

python/sglang/srt/managers/schedule_policy.py

@@ -100,9 +100,15 @@ def match_prefix_for_req(
     if token_ids is None:
         token_ids = req.origin_input_ids + req.output_ids
 
+    # unified_kv SWA lives in a per-request ring (not content-stable, never cached
+    # in the radix tree), so a reused prefix carries stale SWA. Cap the match by the
+    # trailing sliding window so it is re-prefilled. No-op for other layouts.
+    reprefill_tail = tree_cache.swa_reprefill_tail_tokens()
+    key_limit = max(0, len(token_ids) - reprefill_tail) if reprefill_tail else None
+
     match_result = tree_cache.match_prefix(
         MatchPrefixParams(
-            key=RadixKey(token_ids=token_ids, extra_key=req.extra_key),
+            key=RadixKey(token_ids=token_ids, extra_key=req.extra_key, limit=key_limit),
             cow_mamba=cow_mamba,
             req=req if include_req else None,
         )

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -328,6 +328,9 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
     def supports_swa(self) -> bool:
         return False
 
+    def swa_reprefill_tail_tokens(self) -> int:
+        return 0
+
     def supports_mamba(self) -> bool:
         return False
 

python/sglang/srt/mem_cache/swa_radix_cache.py

@@ -370,6 +370,27 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
         ), "sliding_window_size must be set for SWARadixCache"
         return True
 
+    def swa_reprefill_tail_tokens(self) -> int:
+        """Tokens at the tail of a matched prefix that must NOT be reused.
+
+        The DeepSeek-V4 unified_kv layout keeps SWA in a per-request ring
+        (addressed by ``req_pool_idx * window + pos % window``), which is NOT
+        content-stable and is never stored in the radix tree. A reused prefix
+        therefore carries another request's stale SWA in the ring. Hold back the
+        trailing sliding window from the match so it gets re-prefilled into THIS
+        request's ring, making the decode window read freshly-written data.
+
+        No-op (0) for the index-addressed SWA pool, whose slots are
+        content-stable and safe to reuse.
+        """
+        from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import (
+            is_unified_kv_triton,
+        )
+
+        if self.sliding_window_size and is_unified_kv_triton():
+            return self.sliding_window_size
+        return 0
+
     def reset(self) -> None:
         self.root_node = TreeNode()
         self.root_node.key = []