[P/D disagg] Decode-side radix cache for SWA hybrid models (unified radix tree) (#27770)

Co-authored-by: Shangming Cai <csmthu@gmail.com>

python/sglang/srt/disaggregation/decode.py

@@ -71,7 +71,11 @@ from sglang.srt.managers.schedule_batch import (
 from sglang.srt.managers.schedule_policy import match_prefix_for_req
 from sglang.srt.managers.utils import GenerationBatchResult
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
-from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams
+from sglang.srt.mem_cache.base_prefix_cache import (
+    BasePrefixCache,
+    DecLockRefParams,
+    EvictParams,
+)
 from sglang.srt.mem_cache.common import (
     kv_to_page_indices,
     page_align_floor,
@@ -389,6 +393,48 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
             and hasattr(self.token_to_kv_pool_allocator, "alloc_extend_swa_tail")
         )
 
+    def _release_matched_prefix_lock(self, req: Req) -> None:
+        params = DecLockRefParams(swa_uuid_for_lock=req.swa_uuid_for_lock)
+        if req.swa_prefix_lock_released:
+            self.tree_cache.dec_lock_ref(req.last_node, params, skip_swa=True)
+            req.swa_prefix_lock_released = False
+        else:
+            self.tree_cache.dec_lock_ref(req.last_node, params)
+
+    def _reclaim_swa_tail_capacity(
+        self, swa_tail_len: int, req_id: str
+    ) -> Optional[str]:
+        page_size = self.token_to_kv_pool_allocator.page_size
+        required = ceil_align(swa_tail_len, page_size)
+        available = self.token_to_kv_pool_allocator.swa_available_size()
+        if available < required:
+            self.tree_cache.evict(EvictParams(swa_num_tokens=required - available))
+            available = self.token_to_kv_pool_allocator.swa_available_size()
+
+        if available < required:
+            return (
+                f"SWA eviction insufficient: needed={required}, "
+                f"available={available}, req={req_id}"
+            )
+
+        return None
+
+    # SWA caches expose full-attention accounting through full_* accessors.
+    def _radix_full_evictable(self) -> int:
+        if self.scheduler.tp_worker.is_hybrid_swa:
+            return self.tree_cache.full_evictable_size()
+        return self.tree_cache.evictable_size()
+
+    def _radix_full_protected(self) -> int:
+        if self.scheduler.tp_worker.is_hybrid_swa:
+            return self.tree_cache.full_protected_size()
+        return self.tree_cache.protected_size()
+
+    def _radix_full_available(self) -> int:
+        if self.scheduler.tp_worker.is_hybrid_swa:
+            return self.token_to_kv_pool_allocator.full_available_size()
+        return self.token_to_kv_pool_allocator.available_size()
+
     def _swa_tail_len(self, seq_len: int) -> int:
         if not self._uses_swa_tail_prealloc() or seq_len <= 0:
             return max(seq_len, 0)
@@ -398,6 +444,16 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
             return seq_len
 
         page_size = self.token_to_kv_pool_allocator.page_size
+        if getattr(
+            self.scheduler.server_args,
+            "disaggregation_decode_enable_radix_cache",
+            False,
+        ):
+            # Keep enough SWA before the page-aligned radix-cache insert
+            # boundary for the cached key to contain a complete window.
+            # `seq_len - 1` is the last committed position.
+            window_start = max(0, seq_len - 1 - max(window_size, page_size))
+        else:
             window_start = max(0, seq_len - window_size)
         window_start = (window_start // page_size) * page_size
         return seq_len - window_start
@@ -577,8 +633,10 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
             cow_mamba=self.tree_cache.supports_mamba(),
             include_req=True,
         )
-        # Always lock to match aggregated scheduling behavior
+        # Keep aggregated scheduling semantics while preserving the SWA lock
-        self.tree_cache.inc_lock_ref(result.last_device_node)
+        # boundary needed for the matching dec_lock_ref.
+        lock_result = self.tree_cache.inc_lock_ref(result.last_device_node)
+        req.swa_uuid_for_lock = lock_result.swa_uuid_for_lock
         return self._build_decode_prefix_match(req, result)
 
     def _resolve_prefill_dp_rank(self, req: Req) -> Optional[int]:
@@ -752,7 +810,10 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
             self._pre_alloc(req)
             full_allocatable_tokens -= full_required
             if uses_swa_tail_prealloc:
-                swa_allocatable_tokens -= swa_required
+                swa_allocatable_tokens = self._swa_tail_allocatable_token_budget(
+                    count_retracted=False,
+                    extra_reserved_reqs=len(resumed_reqs),
+                )
 
             retraction_restore(
                 req,
@@ -1124,6 +1185,34 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
                 total_prefix_len = prefix_match.decode_prefix_len
 
                 fill_len = self._pre_alloc_fill_len(decode_req.req)
+
+                # Cap full-attention prefix reuse at the sliding-window start so
+                # the SWA window lands entirely in the fresh delta, keeping
+                # alloc_extend_swa_tail's tail->full mapping in range. Costs reuse
+                # of only the last ~window_size full-attention tokens.
+                if uses_swa_tail_prealloc and prefix_len > 0:
+                    swa_prefix_cap = fill_len - self._swa_tail_len(fill_len)
+                    if prefix_len > swa_prefix_cap:
+                        prefix_len = swa_prefix_cap
+                        prefix_indices = prefix_indices[:prefix_len]
+                        # Cap the prefill-committed prefix too: tokens past the
+                        # cap are not device-resident, so prefill must transfer
+                        # them.
+                        total_prefix_len = prefix_len
+
+                # Decode transfers the SWA tail fresh, so retain only the
+                # full-attention prefix lock needed for reuse.
+                if (
+                    uses_swa_tail_prealloc
+                    and prefix_match.l1_prefix_len > 0
+                    and hasattr(self.tree_cache, "dec_swa_lock_only")
+                ):
+                    self.tree_cache.dec_swa_lock_only(
+                        decode_req.req.last_node,
+                        decode_req.req.swa_uuid_for_lock,
+                    )
+                    decode_req.req.swa_prefix_lock_released = True
+
                 required_alloc_tokens = self._required_alloc_tokens(
                     fill_len=fill_len, prefix_len=prefix_len
                 )
@@ -1136,6 +1225,13 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
                     extra_reserved_reqs=len(preallocated_reqs),
                     hicache_reserved_tokens=reserved_restore_tokens,
                 )
+                if uses_swa_tail_prealloc:
+                    swa_allocatable_tokens = self._swa_tail_allocatable_token_budget(
+                        retractable_tokens=retractable_tokens,
+                        retractable_swa_tokens=retractable_swa_tokens,
+                        count_retracted=True,
+                        extra_reserved_reqs=len(preallocated_reqs),
+                    )
             else:
                 prefix_indices = None
                 prefix_len = 0
@@ -1159,12 +1255,12 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
                 )
                 > full_allocatable_tokens
             ):
-                if prefix_len > 0:
+                if prefix_match is not None and prefix_match.l1_prefix_len > 0:
-                    self.tree_cache.dec_lock_ref(decode_req.req.last_node)
+                    self._release_matched_prefix_lock(decode_req.req)
                 break
             if required_tokens_for_request > full_allocatable_tokens:
-                if prefix_len > 0:
+                if prefix_match is not None and prefix_match.l1_prefix_len > 0:
-                    self.tree_cache.dec_lock_ref(decode_req.req.last_node)
+                    self._release_matched_prefix_lock(decode_req.req)
                 break
 
             if uses_swa_tail_prealloc:
@@ -1181,10 +1277,30 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
                     )
                     > swa_allocatable_tokens
                 ):
-                    if prefix_len > 0:
+                    if prefix_match is not None and prefix_match.l1_prefix_len > 0:
-                        self.tree_cache.dec_lock_ref(decode_req.req.last_node)
+                        self._release_matched_prefix_lock(decode_req.req)
                     break
 
+                reclaim_error = self._reclaim_swa_tail_capacity(
+                    swa_len, decode_req.req.rid
+                )
+                if reclaim_error is not None:
+                    if prefix_match is not None and prefix_match.l1_prefix_len > 0:
+                        self._release_matched_prefix_lock(decode_req.req)
+                    logger.error(reclaim_error)
+                    prepare_abort(
+                        decode_req.req,
+                        reclaim_error,
+                        status_code=HTTPStatus.SERVICE_UNAVAILABLE,
+                    )
+                    self.scheduler.output_streamer.stream_output(
+                        [decode_req.req], decode_req.req.return_logprob
+                    )
+                    decode_req.kv_receiver.clear()
+                    decode_req.kv_receiver = None
+                    failed_reqs.append(decode_req)
+                    indices_to_remove.add(i)
+                    continue
             dst_kv_indices = self._pre_alloc(
                 decode_req.req,
                 prefix_indices,
@@ -1206,9 +1322,12 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
                 hicache_reserved_tokens=reserved_restore_tokens,
             )
             if uses_swa_tail_prealloc:
-                # SWA has no radix cache eviction, so decrement its
+                swa_allocatable_tokens = self._swa_tail_allocatable_token_budget(
-                # page-aligned requirement directly.
+                    retractable_tokens=retractable_tokens,
-                swa_allocatable_tokens -= swa_required
+                    retractable_swa_tokens=retractable_swa_tokens,
+                    count_retracted=True,
+                    extra_reserved_reqs=len(preallocated_reqs) + 1,
+                )
             decode_req.req.cache_protected_len = total_prefix_len
 
             page_size = self.token_to_kv_pool_allocator.page_size
@@ -1384,6 +1503,12 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
             indices_to_remove.add(i)
             decode_req.req.time_stats.set_decode_transfer_queue_entry_time()
 
+        if failed_reqs:
+            failed_ids = {id(r) for r in failed_reqs}
+            self.pending_reqs = [
+                r for r in self.pending_reqs if id(r) not in failed_ids
+            ]
+
         self.queue = [
             entry for i, entry in enumerate(self.queue) if i not in indices_to_remove
         ]
@@ -1491,13 +1616,13 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
         elif self._uses_swa_tail_prealloc():
             available_size = self.token_to_kv_pool_allocator.full_available_size()
             if self.scheduler.server_args.disaggregation_decode_enable_radix_cache:
-                available_size += self.tree_cache.evictable_size()
+                available_size += self._radix_full_evictable()
         else:
             available_size = self.token_to_kv_pool_allocator.available_size()
             # Include evictable decode-radix cache entries in the budget -- they
             # can be freed on demand before allocation.
             if self.scheduler.server_args.disaggregation_decode_enable_radix_cache:
-                available_size += self.tree_cache.evictable_size()
+                available_size += self._radix_full_evictable()
         allocatable_tokens = available_size - max(
             reserved_tokens, need_space_for_single_req
         )
@@ -1527,6 +1652,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
         count_retracted: bool = True,
         n_active: Optional[int] = None,
         reserved_tokens: Optional[int] = None,
+        extra_reserved_reqs: int = 0,
     ) -> int:
         need_swa_space_for_single_req = self._need_space_for_single_req(
             retractable_tokens
@@ -1543,7 +1669,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
             )
 
         if n_active is None:
-            n_active = self._active_req_count()
+            n_active = self._active_req_count(extra_reserved_reqs)
         if reserved_tokens is None:
             reserved_tokens = self._active_reserved_tokens(n_active)
 
@@ -1555,11 +1681,14 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
         # remaining headroom up to per-req window cap.
         window_size = self.scheduler.sliding_window_size or 0
         swa_total = self.token_to_kv_pool_allocator.size_swa
-        swa_used = swa_total - self.token_to_kv_pool_allocator.swa_available_size()
+        swa_available = self.token_to_kv_pool_allocator.swa_available_size()
+        swa_evictable = self.tree_cache.swa_evictable_size()
+        swa_used = swa_total - swa_available - swa_evictable
         swa_growth_potential = max(0, n_active * window_size - swa_used)
         swa_reserved_tokens = min(reserved_tokens, swa_growth_potential)
         swa_allocatable_tokens = (
-            self.token_to_kv_pool_allocator.swa_available_size()
+            swa_available
+            + swa_evictable
             - max(swa_reserved_tokens, need_swa_space_for_single_req)
         )
 
@@ -1640,19 +1769,17 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
         # Evict cached entries if the pool doesn't have enough free pages.
         if (
             self.scheduler.server_args.disaggregation_decode_enable_radix_cache
-            and self.token_to_kv_pool_allocator.available_size() < required_alloc_tokens
+            and self._radix_full_available() < required_alloc_tokens
         ):
-            num_to_evict = (
+            num_to_evict = required_alloc_tokens - self._radix_full_available()
-                required_alloc_tokens - self.token_to_kv_pool_allocator.available_size()
-            )
             result = self.tree_cache.evict(EvictParams(num_tokens=num_to_evict))
-            if self.token_to_kv_pool_allocator.available_size() < required_alloc_tokens:
+            if self._radix_full_available() < required_alloc_tokens:
                 logger.warning(
                     f"Eviction insufficient: needed {required_alloc_tokens} tokens, "
-                    f"available {self.token_to_kv_pool_allocator.available_size()} "
+                    f"available {self._radix_full_available()} "
                     f"after evicting {result.num_tokens_evicted}/{num_to_evict} tokens. "
-                    f"evictable_size={self.tree_cache.evictable_size()}, "
+                    f"evictable_size={self._radix_full_evictable()}, "
-                    f"protected_size={self.tree_cache.protected_size()}, "
+                    f"protected_size={self._radix_full_protected()}, "
                     f"fill_len={fill_len}, prefix_len={prefix_len}, "
                     f"total_prefix_len={total_prefix_len}, delta_len={delta_len}, "
                     f"page_size={self.token_to_kv_pool_allocator.page_size}, "
@@ -1660,6 +1787,8 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
                 )
 
         allocator = self.token_to_kv_pool_allocator
+        uses_swa_tail = self._uses_swa_tail_prealloc()
+        swa_tail_len = self._swa_tail_len(fill_len)
         if self.scheduler.enable_hisparse:
             # HiSparse is incompatible with decode-side L1 radix cache. Keep
             # this path on the upstream full-allocation semantics.
@@ -1672,8 +1801,8 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
                 allocator,
                 req=req,
                 fill_len=fill_len,
-                uses_swa_tail=self._uses_swa_tail_prealloc(),
+                uses_swa_tail=uses_swa_tail,
-                swa_tail_len=self._swa_tail_len(fill_len),
+                swa_tail_len=swa_tail_len,
             )
             # Allocate host indices for the RDMA transfer target.
             host_indices = coordinator.mem_pool_host.alloc_paged_token_slots(
@@ -1684,8 +1813,6 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
                 coordinator.host_token_len(fill_len),
             )
         else:
-            uses_swa_tail = self._uses_swa_tail_prealloc() and prefix_len == 0
-            swa_tail_len = self._swa_tail_len(fill_len)
             kv_loc = alloc_for_decode_prealloc(
                 allocator,
                 req=req,
@@ -1700,9 +1827,9 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
             )
         assert kv_loc is not None, (
             f"KV cache is full! Bug in memory estimation. "
-            f"available={self.token_to_kv_pool_allocator.available_size()}, "
+            f"available={self._radix_full_available()}, "
-            f"evictable={self.tree_cache.evictable_size()}, "
+            f"evictable={self._radix_full_evictable()}, "
-            f"protected={self.tree_cache.protected_size()}, "
+            f"protected={self._radix_full_protected()}, "
             f"required_alloc={required_alloc_tokens}, delta={delta_len}, "
             f"fill={fill_len}, prefix={prefix_len}, total_prefix={total_prefix_len}, "
             f"page_size={self.token_to_kv_pool_allocator.page_size}, "
@@ -1822,17 +1949,17 @@ def alloc_for_decode_prealloc(
                 device=device,
             )
         if uses_swa_tail:
-            # Tail-only SWA allocation: only valid when prefix_len == 0.
+            # Full-attention layers reuse prefix KV; SWA layers allocate only
-            # When prefix_len > 0 (radix cache hit), we fall back to
+            # the live window tail.
-            # alloc_extend which allocates SWA at full page count; the
-            # SWA budget in that case may slightly under-estimate.
             kv_loc = allocator.alloc_extend_swa_tail(
-                prefix_lens=torch.tensor([0], dtype=torch.int64, device=device),
+                prefix_lens=torch.tensor(
-                prefix_lens_cpu=torch.tensor([0], dtype=torch.int64),
+                    [prefix_len], dtype=torch.int64, device=device
+                ),
+                prefix_lens_cpu=torch.tensor([prefix_len], dtype=torch.int64),
                 seq_lens=torch.tensor([fill_len], dtype=torch.int64, device=device),
                 seq_lens_cpu=torch.tensor([fill_len], dtype=torch.int64),
                 last_loc=last_loc,
-                extend_num_tokens=fill_len,
+                extend_num_tokens=delta_len,
                 swa_tail_len=swa_tail_len,
                 **extra_kwargs,
             )

python/sglang/srt/mem_cache/kv_cache_builder.py

@@ -46,6 +46,7 @@ from sglang.srt.runtime_context import (
     get_parallel,
     get_schedule,
 )
+from sglang.srt.utils.tensor_bridge import use_mlx
 
 if TYPE_CHECKING:
 
@@ -259,17 +260,37 @@ def build_kv_cache(
             "Transformers backend to avoid multimodal prefix-cache mismatches."
         )
 
-    # Decode radix cache is unsupported with hybrid SWA/SSM models —
+    # Decode-side radix cache supports SWA only through the unified tree, whose
-    # these use specialized memory pools incompatible with the
+    # component pools preserve the full-attention prefix while transferring the
-    # prefix-match-and-lock allocation path.
+    # SWA window fresh. The legacy SWA cache and hybrid SSM pools remain
+    # incompatible with the prefix-match-and-lock allocation path.
     if (
         get_disagg().disaggregation_decode_enable_radix_cache
         and get_disagg().disaggregation_mode == "decode"
     ):
         if is_hybrid_swa:
+            if not (envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get() or use_mlx()):
                 raise ValueError(
-                "--disaggregation-decode-enable-radix-cache is incompatible "
+                    "--disaggregation-decode-enable-radix-cache with sliding "
-                "with sliding window attention (SWA) models"
+                    "window attention (SWA) models requires the unified radix "
+                    "tree (set SGLANG_ENABLE_UNIFIED_RADIX_TREE=1)."
+                )
+            if enable_hierarchical_cache:
+                raise ValueError(
+                    "--disaggregation-decode-enable-radix-cache with sliding "
+                    "window attention (SWA) models currently supports only "
+                    "device-resident cache and is incompatible with "
+                    "--enable-hierarchical-cache."
+                )
+            if getattr(model_config, "is_deepseek_v4_arch", False):
+                raise ValueError(
+                    "--disaggregation-decode-enable-radix-cache does not support "
+                    "DeepSeek-V4 (DSA) compressed KV (c4/c128/indexer) yet."
+                )
+            if getattr(model_config, "is_hybrid_swa_compress", False):
+                raise ValueError(
+                    "--disaggregation-decode-enable-radix-cache does not support "
+                    "SWA-compress models (e.g. Gemma4 / MiMo-V2) yet."
                 )
         if is_hybrid_ssm:
             raise ValueError(

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -901,7 +901,8 @@ class UnifiedRadixCache(BasePrefixCache):
         insert_params.value = values
         result = self.insert(insert_params)
 
-        # Match prefix
+        # Match prefix. SWA insertion retains one extra window before the
+        # page-aligned boundary, so the normal match remains safe to repoint.
         match_result = self.match_prefix(MatchPrefixParams(key=radix_key, req=req))
         new_indices = match_result.device_indices
         new_last_node = match_result.last_device_node

scripts/ci/amd/amd_ci_install_dependency.sh

@@ -96,11 +96,13 @@ docker exec ci_sglang pip install --cache-dir=/sgl-data/pip-cache --upgrade pip
 # Helper function to install with retries and fallback PyPI mirror
 install_with_retry() {
   local max_attempts=3
-  local cmd="$@"
+  local cmd=("$@")
 
   for attempt in $(seq 1 $max_attempts); do
-    echo "Attempt $attempt/$max_attempts: $cmd"
+    printf 'Attempt %s/%s:' "$attempt" "$max_attempts"
-    if eval "$cmd"; then
+    printf ' %q' "${cmd[@]}"
+    printf '\n'
+    if "${cmd[@]}"; then
       echo "Success!"
       return 0
     fi
@@ -109,9 +111,11 @@ install_with_retry() {
       echo "Failed, retrying in 5 seconds..."
       sleep 5
       # Try with alternative PyPI index on retry
-      if [[ "$cmd" =~ "pip install" ]] && [ $attempt -eq 2 ]; then
+      if [[ " ${cmd[*]} " == *" pip install "* ]] && [ $attempt -eq 2 ]; then
-        cmd="$cmd --index-url https://mirrors.aliyun.com/pypi/simple/ --trusted-host mirrors.aliyun.com"
+        cmd+=(--index-url https://mirrors.aliyun.com/pypi/simple/ --trusted-host mirrors.aliyun.com)
-        echo "Using fallback PyPI mirror: $cmd"
+        printf 'Using fallback PyPI mirror:'
+        printf ' %q' "${cmd[@]}"
+        printf '\n'
       fi
     fi
   done

test/registered/disaggregation/test_disaggregation_decode_radix_cache.py

@@ -48,11 +48,13 @@ def _has_mooncake():
 class DisaggregationDecodeRadixCacheTestMixin:
     extra_decode_args = ["--disaggregation-decode-enable-radix-cache"]
     transfer_backend_name = None
+    model_name = DEFAULT_MODEL_NAME_FOR_TEST
+    gsm8k_min_score = 0.80
 
     @classmethod
     def setUpClass(cls):
         super().setUpClass()
-        cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST)
+        cls.model = try_cached_model(cls.model_name)
         cls.transfer_backend = [
             "--disaggregation-transfer-backend",
             cls.transfer_backend_name,
@@ -117,8 +119,8 @@ class DisaggregationDecodeRadixCacheTestMixin:
         metrics_second = run_eval(args)
         print(f"Second run metrics: {metrics_second}")
 
-        self.assertGreater(metrics_first["score"], 0.80)
+        self.assertGreater(metrics_first["score"], self.gsm8k_min_score)
-        self.assertGreater(metrics_second["score"], 0.80)
+        self.assertGreater(metrics_second["score"], self.gsm8k_min_score)
 
         accuracy_drop = metrics_first["score"] - metrics_second["score"]
         self.assertLessEqual(

test/registered/disaggregation/test_disaggregation_decode_radix_cache_swa.py

@@ -0,0 +1,51 @@
+"""SWA coverage for decode-side radix cache on gpt-oss-20b.
+
+The decode worker reuses full-attention prefix KV while transferring the SWA
+window fresh per request. This path requires the unified radix tree and validates
+both multi-turn cache hits and two-pass GSM8K accuracy.
+"""
+
+import unittest
+
+from test_disaggregation_decode_radix_cache import (
+    DisaggregationDecodeRadixCacheTestMixin,
+    _has_nixl,
+)
+
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.server_fixtures.disaggregation_fixture import (
+    PDDisaggregationServerBase,
+)
+from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST_MXFP4_WITH_MOE, is_in_ci
+
+register_cuda_ci(est_time=600, stage="extra-b", runner_config="8-gpu-h200")
+
+SWA_SERVER_ARGS = ["--page-size", "64", "--attention-backend", "triton"]
+
+
+@unittest.skipUnless(
+    is_in_ci() or _has_nixl(),
+    "NIXL is required for decode radix cache disaggregation coverage.",
+)
+class TestDisaggregationDecodeRadixCacheSWANixl(
+    DisaggregationDecodeRadixCacheTestMixin, PDDisaggregationServerBase
+):
+    transfer_backend_name = "nixl"
+    model_name = DEFAULT_MODEL_NAME_FOR_TEST_MXFP4_WITH_MOE
+    # The 512-token eval cap truncates mxfp4 gpt-oss reasoning. On the fixed
+    # 500-question H200 sample, the score has a roughly 2-point standard error,
+    # so keep the original 0.45 absolute floor and rely on the two-pass
+    # non-regression check below to catch decode-cache corruption.
+    gsm8k_min_score = 0.45
+    # SWA + decode-side radix cache is gated to the unified radix tree.
+    extra_prefill_env = {"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}
+    extra_decode_env = {"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}
+    extra_prefill_args = SWA_SERVER_ARGS
+    extra_decode_args = [
+        "--disaggregation-decode-enable-radix-cache",
+        *SWA_SERVER_ARGS,
+    ]
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/radix_cache/test_mamba2_extra_buffer_kl.py

@@ -37,10 +37,10 @@ class TestMamba2ExtraBufferKL(KLDivergenceMixin, DefaultServerBase):
     # Decode-seeded reuse is the regression trigger (the graphed decode
     # track-save); the broken path fails at KL ~1.5, so 0.005 discriminates
     # cleanly while absorbing bf16 reuse noise. Prefill reuse (chunk-aligned
-    # intermediate h states) is inherently looser; threshold matches the manual
+    # intermediate h states) is inherently looser; 0.012 retains a wide margin
-    # TestNvidiaNemotronNanoV2BF16ExtraBuffer calibration.
+    # below the broken path while covering the observed bf16 calibration noise.
     kl_div_thres = 0.005
-    kl_div_thres_prefill = 0.01
+    kl_div_thres_prefill = 0.012
     kl_div_max_samples = 16
 
     other_args = [

test/registered/unit/disaggregation/test_decode_queue_cleanup.py

@@ -72,6 +72,9 @@ class TestDecodeQueueCleanup(CustomTestCase):
         queue._swa_aware_allocatable_token_budgets = MagicMock(
             return_value=(physical_available, physical_available)
         )
+        queue._swa_tail_allocatable_token_budget = MagicMock(
+            side_effect=lambda **_: physical_available
+        )
 
         def pre_alloc(_req):
             nonlocal physical_available
@@ -170,6 +173,72 @@ class TestDecodeQueueCleanup(CustomTestCase):
         self.assertTrue(all(r is not decode_req for r in queue.pending_reqs))
         self.assertIsNone(decode_req.kv_receiver)
 
+    def test_swa_reclaim_failure_rejects_only_request(self):
+        receiver = FakeReceiver()
+        req = SimpleNamespace(
+            rid="swa-reclaim-failed",
+            origin_input_ids=[1, 2, 3],
+            output_ids=[],
+            finished_reason=None,
+            return_logprob=False,
+            sampling_params=SimpleNamespace(max_new_tokens=1),
+        )
+        decode_req = SimpleNamespace(
+            req=req,
+            kv_receiver=receiver,
+            waiting_for_input=True,
+            is_rebootstrap=False,
+        )
+
+        queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
+        queue.pp_size = 1
+        queue.queue = [decode_req]
+        queue.pending_reqs = [decode_req]
+        queue.retracted_queue = []
+        queue.num_reserved_decode_tokens = 0
+        queue._resolve_pending_reqs = MagicMock()
+        queue._update_handshake_waiters = MagicMock()
+        queue._uses_swa_tail_prealloc = MagicMock(return_value=True)
+        queue._swa_aware_allocatable_token_budgets = MagicMock(
+            return_value=(1024, 1024)
+        )
+        queue._prealloc_required_tokens = MagicMock(return_value=(3, 3))
+        queue._prealloc_kv_lens = MagicMock(return_value=(3, 3))
+        queue._reclaim_swa_tail_capacity = MagicMock(
+            return_value=(
+                "SWA eviction insufficient: needed=64, available=0, "
+                "req=swa-reclaim-failed"
+            )
+        )
+        queue._hicache_pending_restore_tokens = MagicMock(return_value=0)
+        queue._pre_alloc = MagicMock()
+        queue.req_to_token_pool = MagicMock()
+        queue.req_to_token_pool.available_size.return_value = 1
+        queue.req_to_metadata_buffer_idx_allocator = MagicMock()
+        queue.req_to_metadata_buffer_idx_allocator.available_size.return_value = 1
+
+        scheduler = MagicMock()
+        scheduler.running_batch.reqs = []
+        scheduler.enable_priority_scheduling = False
+        scheduler.enable_hisparse = False
+        scheduler.server_args.disaggregation_decode_enable_radix_cache = False
+        scheduler.output_streamer = MagicMock()
+        queue.scheduler = scheduler
+
+        preallocated, failed = queue.pop_preallocated()
+
+        self.assertEqual(preallocated, [])
+        self.assertEqual(failed, [decode_req])
+        self.assertEqual(queue.queue, [])
+        self.assertEqual(queue.pending_reqs, [])
+        self.assertTrue(receiver.clear_called)
+        self.assertIsNone(decode_req.kv_receiver)
+        self.assertIsInstance(req.finished_reason, FINISH_ABORT)
+        queue._pre_alloc.assert_not_called()
+        scheduler.output_streamer.stream_output.assert_called_once_with(
+            [req], req.return_logprob
+        )
+
     def test_ensure_prefill_info_tolerates_cleared_receiver(self):
         # A req whose kv_receiver was already cleared must not crash on .abort().
         queue = DecodePreallocQueue.__new__(DecodePreallocQueue)

test/registered/unit/mem_cache/test_decode_radix_lock_ref.py

@@ -34,6 +34,7 @@ import torch
 from sglang.srt.disaggregation.decode import DecodePreallocQueue
 from sglang.srt.disaggregation.decode_hicache_mixin import DecodePrefixMatch
 from sglang.srt.mem_cache.base_prefix_cache import (
+    DecLockRefParams,
     InsertParams,
     MatchPrefixParams,
 )
@@ -95,6 +96,71 @@ def _make_req(fill_ids, req_pool_idx=0, cache_protected_len=0, last_node=None):
 class TestDecodeLockRefScenarios(unittest.TestCase):
     """Test lock_ref balance across decode transfer scenarios."""
 
+    def test_swa_tail_len_keeps_page_aligned_matchable_window(self):
+        queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
+        queue._uses_swa_tail_prealloc = MagicMock(return_value=True)
+        queue.scheduler = SimpleNamespace(
+            sliding_window_size=127,
+            server_args=SimpleNamespace(disaggregation_decode_enable_radix_cache=True),
+        )
+        queue.token_to_kv_pool_allocator = MagicMock(page_size=64)
+
+        tail_len = queue._swa_tail_len(895)
+
+        self.assertEqual(tail_len, 191)
+        swa_start = 895 - tail_len
+        radix_key_len = (895 // 64) * 64
+        self.assertGreaterEqual(radix_key_len - swa_start, 127)
+
+    def test_swa_admission_counts_evictable_capacity(self):
+        queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
+        queue.scheduler = MagicMock()
+        queue.scheduler.running_batch.reqs = []
+        queue.scheduler.sliding_window_size = 128
+        queue.scheduler.last_batch = None
+        queue.retracted_queue = []
+        queue._need_space_for_single_req = MagicMock(return_value=0)
+        queue._active_req_count = MagicMock(return_value=1)
+        queue.token_to_kv_pool_allocator = MagicMock()
+        queue.token_to_kv_pool_allocator.size_swa = 256
+        queue.token_to_kv_pool_allocator.swa_available_size.return_value = 0
+        queue.tree_cache = MagicMock()
+        queue.tree_cache.swa_evictable_size.return_value = 192
+
+        budget = queue._swa_tail_allocatable_token_budget(
+            count_retracted=False,
+            reserved_tokens=64,
+        )
+
+        # 192 reclaimable tokens minus 64 reserved for active-request growth.
+        self.assertEqual(budget, 128)
+
+    def test_reclaim_swa_tail_capacity_page_rounds(self):
+        queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
+        queue.token_to_kv_pool_allocator = MagicMock(page_size=64)
+        queue.token_to_kv_pool_allocator.swa_available_size.side_effect = [64, 192]
+        queue.tree_cache = MagicMock()
+
+        error = queue._reclaim_swa_tail_capacity(129, "req-1")
+
+        self.assertIsNone(error)
+        params = queue.tree_cache.evict.call_args.args[0]
+        self.assertEqual(params.num_tokens, 0)
+        self.assertEqual(params.swa_num_tokens, 128)
+
+    def test_reclaim_swa_tail_capacity_fails_before_allocation(self):
+        queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
+        queue.token_to_kv_pool_allocator = MagicMock(page_size=64)
+        queue.token_to_kv_pool_allocator.swa_available_size.side_effect = [64, 128]
+        queue.tree_cache = MagicMock()
+
+        error = queue._reclaim_swa_tail_capacity(129, "req-1")
+
+        self.assertEqual(
+            error,
+            "SWA eviction insufficient: needed=192, available=128, req=req-1",
+        )
+
     def _populate_prefix(self, cache, prefix_ids, prefix_values):
         """Insert a prefix into the tree so future requests can match it."""
         cache.insert(
@@ -303,6 +369,9 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
         req.last_node = object()
         req.finished_reason = None
         req.cache_protected_len = 0
+        req.swa_uuid_for_lock = 123
+        req.swa_prefix_lock_released = False
+        req.pd_rebootstrap_in_progress = False
         req.sampling_params.max_new_tokens = 16
 
         decode_req = MagicMock()
@@ -319,6 +388,10 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
         queue.num_reserved_decode_tokens = 0
         queue._resolve_pending_reqs = MagicMock()
         queue._update_handshake_waiters = MagicMock()
+        queue._uses_swa_tail_prealloc = MagicMock(return_value=True)
+        queue._swa_tail_len = MagicMock(return_value=8)
+        queue._swa_aware_allocatable_token_budgets = MagicMock(return_value=(8, 8))
+        queue._swa_tail_allocatable_token_budget = MagicMock(return_value=8)
         queue._match_prefix_and_lock = MagicMock(
             return_value=DecodePrefixMatch(
                 prefix_indices=torch.arange(4, dtype=torch.int64),
@@ -349,21 +422,32 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
         scheduler.running_batch = running_batch
         scheduler.server_args = server_args
         scheduler.enable_hisparse = False
+        scheduler.enable_decode_hicache = False
+        scheduler.enable_priority_scheduling = False
         scheduler.waiting_queue = []
         scheduler.last_batch = None
         scheduler.output_streamer = MagicMock()
         queue.scheduler = scheduler
 
-        # Initial budget says the request fits; post-lock budget says it does not.
+        # The 4-token match is locked, then capped to zero because the whole
-        queue._allocatable_token_budgets = MagicMock(side_effect=[8, 3])
+        # 8-token request is inside the SWA window. Admission rejection must
+        # still release the original matched-node lock.
+        queue._allocatable_token_budgets = MagicMock(return_value=3)
 
         preallocated, failed = queue.pop_preallocated()
 
         self.assertEqual(preallocated, [])
         self.assertEqual(failed, [])
         queue._pre_alloc.assert_not_called()
-        queue.tree_cache.dec_lock_ref.assert_called_once_with(req.last_node)
+        queue.tree_cache.dec_swa_lock_only.assert_called_once_with(req.last_node, 123)
-        self.assertEqual(queue._allocatable_token_budgets.call_count, 2)
+        queue.tree_cache.dec_lock_ref.assert_called_once_with(
+            req.last_node,
+            DecLockRefParams(swa_uuid_for_lock=123),
+            skip_swa=True,
+        )
+        self.assertFalse(req.swa_prefix_lock_released)
+        queue._swa_tail_len.assert_called_once_with(8)
+        queue._allocatable_token_budgets.assert_called_once()
 
     def test_repeated_incremental_no_leak(self):
         """Multiple incremental transfers shouldn't leak lock_refs."""

test/registered/unit/mem_cache/test_hisparse_allocator.py

@@ -116,6 +116,7 @@ class TestDeepSeekV4HiSparseAllocator(CustomTestCase):
             device=torch.device("cpu"),
             page_size=256,
             available_size=MagicMock(return_value=fill_len),
+            swa_available_size=MagicMock(return_value=swa_tail_len),
             alloc_extend_swa_tail=MagicMock(return_value=kv_loc),
             alloc_logical_only=MagicMock(return_value=kv_loc),
         )