diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 00cf04097089..032ab23fcd92 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -21,6 +21,7 @@ from __future__ import annotations import logging +import math import time from collections import deque from dataclasses import dataclass @@ -69,7 +70,13 @@ 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, + MatchPrefixParams, + zero_match_result, +) from sglang.srt.mem_cache.common import ( kv_to_page_indices, page_align_floor, @@ -81,6 +88,7 @@ KVCache, ReqToTokenPool, ) +from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.observability.req_time_stats import ( set_schedule_time_batch, @@ -124,18 +132,18 @@ def __init__( max_context_len: int, device: str, enable_memory_saver: bool, - pre_alloc_size: int, + pre_alloc_size: Optional[int], ): memory_saver_adapter = TorchMemorySaverAdapter.create( enable=enable_memory_saver ) self.size = size + self.pre_alloc_size = pre_alloc_size if pre_alloc_size is not None else 0 # +1 padding row at index 0; see ReqToTokenPool for rationale. - self._alloc_size = size + pre_alloc_size + 1 + self._alloc_size = size + self.pre_alloc_size + 1 self.max_context_len = max_context_len self.device = device - self.pre_alloc_size = pre_alloc_size with memory_saver_adapter.region(tag=GPU_MEMORY_TYPE_KV_CACHE): self.req_to_token = torch.zeros( (self._alloc_size, max_context_len), @@ -201,7 +209,7 @@ def __init__( mamba_layer_ids: List[int], speculative_num_draft_tokens: int, enable_mamba_extra_buffer: bool, - pre_alloc_size: int, + pre_alloc_size: Optional[int], enable_overlap_schedule: bool, mamba_size: int = None, start_layer: int = None, @@ -219,12 +227,13 @@ def __init__( self.mamba_ping_pong_track_buffer_size = 2 if enable_overlap_schedule else 1 self.enable_mamba_extra_buffer = enable_mamba_extra_buffer self.enable_memory_saver = enable_memory_saver + _pre_alloc = pre_alloc_size if pre_alloc_size is not None else 0 # Each request needs 1 main mamba slot + ping-pong slots when extra_buffer is enabled. # Cap the pool at max concurrent requests * slots_per_req to avoid allocating failed. slots_per_req = 1 + ( self.mamba_ping_pong_track_buffer_size if enable_mamba_extra_buffer else 0 ) - max_slots_needed = (size + pre_alloc_size) * slots_per_req + max_slots_needed = (size + _pre_alloc) * slots_per_req if mamba_size is not None: effective_mamba_size = max(mamba_size, max_slots_needed) if mamba_size < max_slots_needed: @@ -233,7 +242,7 @@ def __init__( "raising effective_mamba_size to %d", mamba_size, max_slots_needed, - size + pre_alloc_size, + size + _pre_alloc, slots_per_req, effective_mamba_size, ) @@ -243,7 +252,7 @@ def __init__( self.layer_transfer_counter = None self._init_mamba_pool( mamba_size=effective_mamba_size, - mamba_spec_state_size=size + pre_alloc_size, + mamba_spec_state_size=size + _pre_alloc, cache_params=cache_params, mamba_layer_ids=mamba_layer_ids, device=device, @@ -370,6 +379,44 @@ def _uses_swa_tail_prealloc(self) -> bool: 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 getattr(req, "swa_prefix_lock_released", False): + 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) -> None: + page_size = self.token_to_kv_pool_allocator.page_size + required = ((swa_tail_len + page_size - 1) // page_size) * 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: + raise RuntimeError( + f"SWA eviction insufficient: needed={required}, " + f"available={available}, req={req_id}" + ) + + # 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) @@ -379,7 +426,10 @@ def _swa_tail_len(self, seq_len: int) -> int: return seq_len page_size = self.token_to_kv_pool_allocator.page_size - window_start = max(0, seq_len - window_size) + # Match the radix-cache eviction margin: keep enough SWA before the + # page-aligned 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)) window_start = (window_start // page_size) * page_size return seq_len - window_start @@ -394,6 +444,90 @@ def _prealloc_kv_lens(self, req: Req) -> Tuple[int, int]: return allocated_kv_len, self._swa_tail_len(allocated_kv_len) return allocated_kv_len, allocated_kv_len + def _uses_dsv4_decode_radix_cache(self) -> bool: + return ( + self.scheduler.server_args.disaggregation_decode_enable_radix_cache + and isinstance(self.token_to_kv_pool, DeepSeekV4TokenToKVPool) + ) + + def _dsv4_safe_prefix_len(self, prefix_len: int) -> int: + """Avoid splitting reused prefixes through compressed DSV4 blocks.""" + if not self._uses_dsv4_decode_radix_cache() or prefix_len <= 0: + return prefix_len + + compression_ratios = self.token_to_kv_pool.compression_ratios + max_compression_ratio = max([r for r in compression_ratios if r > 0], default=1) + page_size = self.token_to_kv_pool_allocator.page_size + alignment = math.lcm(page_size, max_compression_ratio) + return (prefix_len // alignment) * alignment + + def _dsv4_singleflight_min_prefix_len(self) -> int: + compression_ratios = self.token_to_kv_pool.compression_ratios + max_compression_ratio = max([r for r in compression_ratios if r > 0], default=1) + alignment = math.lcm( + self.token_to_kv_pool_allocator.page_size, max_compression_ratio + ) + return max(4096, alignment) + + @staticmethod + def _common_prefix_len(lhs: List[int], rhs: List[int], limit: int) -> int: + for prefix_len in range(limit): + if lhs[prefix_len] != rhs[prefix_len]: + return prefix_len + return min(len(lhs), len(rhs), limit) + + def _dsv4_inflight_prompt_reqs( + self, preallocated_reqs: List[DecodeRequest] + ) -> List[Req]: + if not self._uses_dsv4_decode_radix_cache(): + return [] + + inflight_reqs = [ + decode_req.req + for decode_req in self.transfer_queue.queue + if getattr(decode_req.req, "dsv4_decode_radix_cache_prompt_once", False) + ] + inflight_reqs.extend( + req + for req in self.scheduler.running_batch.reqs + if getattr(req, "dsv4_decode_radix_cache_prompt_once", False) + ) + inflight_reqs.extend( + decode_req.req + for decode_req in preallocated_reqs + if getattr(decode_req.req, "dsv4_decode_radix_cache_prompt_once", False) + ) + return inflight_reqs + + def _should_wait_for_dsv4_inflight_prompt( + self, + req: Req, + *, + prefix_len: int, + preallocated_reqs: List[DecodeRequest], + ) -> bool: + min_prefix_len = self._dsv4_singleflight_min_prefix_len() + for inflight_req in self._dsv4_inflight_prompt_reqs(preallocated_reqs): + if inflight_req is req: + continue + common_len = self._common_prefix_len( + req.origin_input_ids, + inflight_req.origin_input_ids, + min(len(req.origin_input_ids), len(inflight_req.origin_input_ids)), + ) + safe_common_len = self._dsv4_safe_prefix_len(common_len) + if safe_common_len - prefix_len >= min_prefix_len: + return True + return False + + def _release_decode_radix_match(self, req: Req) -> None: + self.tree_cache.dec_lock_ref( + req.last_node, + DecLockRefParams(swa_uuid_for_lock=req.swa_uuid_for_lock), + ) + req.last_node = self.tree_cache.root_node + req.swa_uuid_for_lock = None + def _prealloc_required_tokens(self, req: Req) -> Tuple[int, int]: full_len, swa_len = self._prealloc_kv_lens(req) swa_reserved = self.num_reserved_decode_tokens @@ -541,15 +675,48 @@ def _match_prefix_and_lock(self, req: Req) -> DecodePrefixMatch: Match a request against the decode-side radix cache, lock the matched node to prevent eviction, and return the matched prefix information. """ - result = match_prefix_for_req( - self.tree_cache, - req, - req.origin_input_ids, - cow_mamba=self.tree_cache.supports_mamba(), - include_req=True, - ) - # Always lock to match aggregated scheduling behavior - self.tree_cache.inc_lock_ref(result.last_device_node) + if self._uses_dsv4_decode_radix_cache(): + # DSV4 prompt donation creates full-only leaves: the SWA component is + # intentionally tombstoned because decode only needs the long full + # prefix while the SWA tail is recomputed/transferred. Use the full + # match here; the generic SWA-window-safe match would truncate these + # full-only leaves to zero and force full KV transfer again. + result = self.tree_cache.match_prefix( + MatchPrefixParams( + key=RadixKey(req.origin_input_ids, req.extra_key), + cow_mamba=self.tree_cache.supports_mamba(), + req=req, + return_full_match=True, + ) + ) + if envs.SGLANG_RADIX_FORCE_MISS.get(): + result = zero_match_result(self.tree_cache, result) + ( + req.prefix_indices, + req.last_node, + req.last_host_node, + req.best_match_node, + req.host_hit_length, + ) = ( + result.device_indices, + result.last_device_node, + result.last_host_node, + result.best_match_node, + result.host_hit_length, + ) + else: + result = match_prefix_for_req( + self.tree_cache, + req, + req.origin_input_ids, + cow_mamba=self.tree_cache.supports_mamba(), + include_req=True, + ) + # Always lock to match aggregated scheduling behavior. SWA locks only + # span the sliding window and return a boundary uuid; store it so the + # matching dec_lock_ref stops there instead of underflowing toward root. + 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]: @@ -711,7 +878,10 @@ def resume_retracted_reqs( 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), + ) # load from cpu, release the cpu copy req.load_kv_cache(self.req_to_token_pool, self.token_to_kv_pool_allocator) @@ -984,8 +1154,63 @@ def pop_preallocated( # gap is filled by HiCache loadback later. prefix_len = prefix_match.l1_prefix_len total_prefix_len = prefix_match.decode_prefix_len + locked_prefix_len = prefix_len + # Align prefix_len down to page boundary so both prefill and + # decode agree on the page-aligned split point for KV transfer. + page_size = self.token_to_kv_pool_allocator.page_size + if page_size > 1 and prefix_len % page_size != 0: + prefix_len = page_align_floor(prefix_len, page_size) + prefix_indices = prefix_indices[:prefix_len] + total_prefix_len = min(total_prefix_len, prefix_len) + + fill_len = origin_input_len + max(len(decode_req.req.output_ids) - 1, 0) + + # 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 = min(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 + + dsv4_safe_prefix_len = self._dsv4_safe_prefix_len(prefix_len) + if dsv4_safe_prefix_len < prefix_len: + prefix_len = dsv4_safe_prefix_len + prefix_indices = prefix_indices[:prefix_len] + total_prefix_len = min(total_prefix_len, prefix_len) + + if self._should_wait_for_dsv4_inflight_prompt( + decode_req.req, + prefix_len=prefix_len, + preallocated_reqs=preallocated_reqs, + ): + if locked_prefix_len > 0: + self._release_decode_radix_match(decode_req.req) + continue + + decode_req.req.cache_protected_len = prefix_len + if locked_prefix_len > 0 and prefix_len == 0: + self._release_decode_radix_match(decode_req.req) - fill_len = self._pre_alloc_fill_len(decode_req.req) required_alloc_tokens = self._required_alloc_tokens( fill_len=fill_len, prefix_len=prefix_len ) @@ -998,6 +1223,13 @@ def pop_preallocated( 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 @@ -1021,12 +1253,12 @@ def pop_preallocated( ) > full_allocatable_tokens ): - if prefix_len > 0: - self.tree_cache.dec_lock_ref(decode_req.req.last_node) + if prefix_match is not None and prefix_match.l1_prefix_len > 0: + self._release_matched_prefix_lock(decode_req.req) break if required_tokens_for_request > full_allocatable_tokens: - if prefix_len > 0: - self.tree_cache.dec_lock_ref(decode_req.req.last_node) + if prefix_match is not None and prefix_match.l1_prefix_len > 0: + self._release_matched_prefix_lock(decode_req.req) break if uses_swa_tail_prealloc: @@ -1043,8 +1275,8 @@ def pop_preallocated( ) > swa_allocatable_tokens ): - if prefix_len > 0: - self.tree_cache.dec_lock_ref(decode_req.req.last_node) + if prefix_match is not None and prefix_match.l1_prefix_len > 0: + self._release_matched_prefix_lock(decode_req.req) break dst_kv_indices = self._pre_alloc( @@ -1068,9 +1300,12 @@ def pop_preallocated( hicache_reserved_tokens=reserved_restore_tokens, ) if uses_swa_tail_prealloc: - # SWA budget uses simple decrement (no radix cache eviction in - # the SWA pool, so page-rounding drift is negligible). - swa_allocatable_tokens -= swa_required + 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) + 1, + ) decode_req.req.cache_protected_len = total_prefix_len page_size = self.token_to_kv_pool_allocator.page_size @@ -1178,6 +1413,21 @@ def _c128_state_payload(): page_indices = kv_to_page_indices(kv_indices, kv_transfer_page_size).astype( np.int32 ) + if ( + self._uses_dsv4_decode_radix_cache() + and envs.SGLANG_DEBUG_DSV4_DECODE_RADIX_TRANSFER.get() + ): + logger.info( + "DSV4 decode radix transfer stats: rid=%s " + "origin_input_len=%d decode_prefix_len=%d " + "transfer_tokens=%d transfer_pages=%d page_size=%d", + decode_req.req.rid, + origin_input_len, + total_prefix_len, + origin_input_len - total_prefix_len, + len(page_indices), + kv_transfer_page_size, + ) if ( self.transfer_queue.enable_staging and hasattr(decode_req.kv_receiver, "require_staging") @@ -1298,13 +1548,13 @@ def _allocatable_token_budgets( 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 ) @@ -1334,6 +1584,7 @@ def _swa_tail_allocatable_token_budget( 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 @@ -1350,7 +1601,7 @@ def _swa_tail_allocatable_token_budget( ) 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) @@ -1362,11 +1613,14 @@ def _swa_tail_allocatable_token_budget( # 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) ) @@ -1447,19 +1701,17 @@ def _pre_alloc( # 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 = ( - required_alloc_tokens - self.token_to_kv_pool_allocator.available_size() - ) + num_to_evict = required_alloc_tokens - self._radix_full_available() 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"protected_size={self.tree_cache.protected_size()}, " + f"evictable_size={self._radix_full_evictable()}, " + 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}, " @@ -1467,6 +1719,10 @@ def _pre_alloc( ) 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 uses_swa_tail: + self._reclaim_swa_tail_capacity(swa_tail_len, req.rid) if self.scheduler.enable_hisparse: # HiSparse is incompatible with decode-side L1 radix cache. Keep # this path on the upstream full-allocation semantics. @@ -1479,8 +1735,8 @@ def _pre_alloc( allocator, req=req, fill_len=fill_len, - uses_swa_tail=self._uses_swa_tail_prealloc(), - swa_tail_len=self._swa_tail_len(fill_len), + uses_swa_tail=uses_swa_tail, + swa_tail_len=swa_tail_len, ) # Allocate host indices for the RDMA transfer target. host_indices = coordinator.mem_pool_host.alloc_paged_token_slots( @@ -1506,9 +1762,9 @@ def _pre_alloc( ) 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"evictable={self.tree_cache.evictable_size()}, " - f"protected={self.tree_cache.protected_size()}, " + f"available={self._radix_full_available()}, " + f"evictable={self._radix_full_evictable()}, " + 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}, " @@ -1527,6 +1783,15 @@ def _pre_alloc( # inserts committed KV into the radix tree. The last output token # hasn't had KV committed yet (output_ids is 1 ahead). req.full_untruncated_fill_ids = req.origin_input_ids + req.output_ids + if self._uses_dsv4_decode_radix_cache(): + # DSV4 compressed sidecars are not yet safe to reinsert from decode + # workers after generation starts. Defer the one prompt insert + # until the prebuilt forward has finished; process_prebuilt() runs + # before forward and must not free duplicate prompt pages still + # referenced by the current batch's out_cache_loc. + req.dsv4_decode_radix_cache_prompt_len = req.kv_committed_len + req.dsv4_decode_radix_cache_prompt_once = True + req.skip_radix_cache_insert = True # Set prefix_indices so downstream consumers (init_next_round_input, # prepare_for_extend) see the correct prefix length. In the agg path # this is done inside init_next_round_input, but decode-disagg needs @@ -1609,17 +1874,16 @@ def alloc_for_decode_prealloc( else torch.tensor([-1], dtype=torch.int64, device=device) ) if uses_swa_tail: - # Tail-only SWA allocation: only valid when prefix_len == 0. - # When prefix_len > 0 (radix cache hit), we fall back to - # alloc_extend which allocates SWA at full page count; the - # SWA budget in that case may slightly under-estimate. + # Full-attention layers reuse prefix KV; SWA layers allocate + # only the live window tail. Supports prefix_len > 0 (radix cache + # hit) by passing the real prefix_len and extending only the delta. kv_loc = allocator.alloc_extend_swa_tail( - prefix_lens=torch.tensor([0], dtype=torch.int64, device=device), - prefix_lens_cpu=torch.tensor([0], dtype=torch.int64), + prefix_lens=torch.tensor([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, ) req.kv.swa_evicted_seqlen = fill_len - swa_tail_len diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index b4ff61057eb6..3a6ce2177ccb 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index 248a92939e28..d6f8100d946c 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/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] diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index 750683b8e2a3..4a7c4c947bc1 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/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 = [] diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 244eef333b70..3f3acf586158 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -53,6 +53,9 @@ class MatchPrefixParams: # Mamba specific cow_mamba: bool = False req: Optional[Req] = None + # Return full-attention KV indices even if SWA tombstones cap the normal + # match. Used by decode-side radix repointing; non-SWA caches are unchanged. + return_full_match: bool = False @dataclasses.dataclass @@ -68,6 +71,7 @@ class InsertParams: # SWA specific prev_prefix_len: int = 0 swa_evicted_seqlen: int = 0 + force_leaf_creation: bool = False # General chunked: bool = False diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 25d3f96e81f7..b244d9bbbd74 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -97,7 +97,10 @@ def free_swa_out_of_window_slots( def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs): if getattr(req, "skip_radix_cache_insert", False): - return + if getattr(req, "allow_radix_cache_insert_once", False): + req.allow_radix_cache_insert_once = False + else: + return tree_cache.cache_unfinished_req(req, **kwargs) diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index 3e6ac06394ac..182f6fbb3199 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/python/sglang/srt/mem_cache/kv_cache_builder.py @@ -37,6 +37,7 @@ class KVCacheBuildResult: from sglang.srt.mem_cache.registry import TreeCacheBuildContext, create_tree_cache from sglang.srt.model_loader.utils import get_resolved_model_impl from sglang.srt.runtime_context import get_parallel +from sglang.srt.utils.tensor_bridge import use_mlx if TYPE_CHECKING: @@ -127,6 +128,19 @@ def maybe_register_hicache_draft( tree_cache.cache_controller.set_draft_kv_pool(pool, draft_host_pool) +def is_supported_dsv4_decode_radix_mtp( + *, spec_algorithm: SpeculativeAlgorithm, server_args: ServerArgs +) -> bool: + return ( + spec_algorithm.is_eagle() + and not spec_algorithm.is_eagle3() + and not spec_algorithm.is_frozen_kv_mtp() + and server_args.speculative_eagle_topk == 1 + and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get() + and envs.SGLANG_EXPERIMENTAL_ONLINE_C128_MTP.get() + ) + + def build_kv_cache( *, server_args: ServerArgs, @@ -182,18 +196,58 @@ def build_kv_cache( "Transformers backend to avoid multimodal prefix-cache mismatches." ) - # Decode radix cache is unsupported with hybrid SWA/SSM models — - # these use specialized memory pools incompatible with the - # prefix-match-and-lock allocation path. + # Decode-side radix cache: SWA is supported only via the unified radix tree + # (its component pools handle the prefix-match-and-lock path); the default + # SWARadixCache and the Mamba/SSM pools are not supported. if ( server_args.disaggregation_decode_enable_radix_cache and server_args.disaggregation_mode == "decode" ): if is_hybrid_swa: - raise ValueError( - "--disaggregation-decode-enable-radix-cache is incompatible " - "with sliding window attention (SWA) models" - ) + if not (envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get() or use_mlx()): + raise ValueError( + "--disaggregation-decode-enable-radix-cache with sliding " + "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." + ) + # Compressed-KV SWA variants need model-specific sidecar guarantees. + # DSV4 has a conservative experimental L1-only path below. + if getattr(model_config, "is_deepseek_v4_arch", False): + if not envs.SGLANG_EXPERIMENTAL_DSV4_DECODE_RADIX_CACHE.get(): + raise ValueError( + "--disaggregation-decode-enable-radix-cache with " + "DeepSeek-V4 (DSA compressed KV) is experimental. Set " + "SGLANG_EXPERIMENTAL_DSV4_DECODE_RADIX_CACHE=1 to enable " + "the conservative L1-only path." + ) + if enable_hierarchical_cache: + raise ValueError( + "DeepSeek-V4 decode-side radix cache currently supports " + "only device-resident L1 cache. Disable hierarchical " + "cache / HiCache storage for this experimental path." + ) + if not spec_algorithm.is_none(): + if not is_supported_dsv4_decode_radix_mtp( + spec_algorithm=spec_algorithm, server_args=server_args + ): + raise ValueError( + "DeepSeek-V4 decode-side radix cache currently supports " + "only the experimental EAGLE topk=1 online c128 MTP " + "path. Set SGLANG_OPT_USE_ONLINE_COMPRESS=1 and " + "SGLANG_EXPERIMENTAL_ONLINE_C128_MTP=1." + ) + 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( "--disaggregation-decode-enable-radix-cache is incompatible " diff --git a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py index 12621dce7581..0eb93b8d7564 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py @@ -630,6 +630,9 @@ def prepare_for_caching_req( # Unfinished requests can already have an SWA-evicted prefix; preserve # that boundary so insertion creates a tombstone instead of live SWA KV. insert_params.swa_evicted_seqlen = req.kv.swa_evicted_seqlen + insert_params.force_leaf_creation = getattr( + req, "force_radix_leaf_creation", False + ) return None def free_out_of_window_slots( diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 69477a27a42f..737d1180e6db 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -715,8 +715,13 @@ def cache_unfinished_req(self, req: Req, chunked: bool = False, **kwargs) -> Non insert_params.value = values result = self.insert(insert_params) - # Match prefix - match_result = self.match_prefix(MatchPrefixParams(key=radix_key)) + # Repoint by full-attention residency, not the SWA-window-safe match. + # A reused prefix can remain full-resident after its SWA window is + # tombstoned, so the window-safe match may be shorter than the KV + # indices that must be preserved. + match_result = self.match_prefix( + MatchPrefixParams(key=radix_key, return_full_match=True) + ) new_indices = match_result.device_indices new_last_node = match_result.last_device_node new_prefix_len = result.prefix_len diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 4ec37bf78744..1a75bc9e7b01 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -7105,6 +7105,78 @@ def _check_format(has_params, has_consolidated, has_hf_weights) -> bool: except Exception: return False + def _handle_pd_disaggregation(self): + if self.disaggregation_mode == "decode": + if self.disaggregation_decode_enable_radix_cache: + if self.enable_hisparse: + raise ValueError( + "--disaggregation-decode-enable-radix-cache is incompatible " + "with --enable-hisparse" + ) + if self.disaggregation_transfer_backend not in ("nixl", "mooncake"): + raise ValueError( + "--disaggregation-decode-enable-radix-cache currently " + "requires --disaggregation-transfer-backend in " + "('nixl', 'mooncake'), but got " + f"{self.disaggregation_transfer_backend!r}" + ) + if ( + self.speculative_algorithm is not None + and not self._allow_dsv4_decode_radix_speculative() + ): + raise ValueError( + "--disaggregation-decode-enable-radix-cache is incompatible " + "with speculative decoding " + f"(--speculative-algorithm {self.speculative_algorithm})" + ) + if self.enable_dp_attention: + logger.warning( + "EXPERIMENTAL: Decode radix cache with DP attention. " + "Requires prefix-aware DP rank routing for optimal cache hits." + ) + self.disable_radix_cache = False + logger.warning("EXPERIMENTAL: Radix cache is enabled for decode server") + else: + self.disable_radix_cache = True + logger.warning("KV cache is forced as chunk cache for decode server") + if self.enable_mamba_extra_buffer(): + logger.warning( + "Mamba extra_buffer is disabled because decode disaggregation " + "currently forces chunk cache. Falling back to no_buffer." + ) + self.mamba_scheduler_strategy = "no_buffer" + + elif self.disaggregation_mode == "prefill": + assert ( + self.disaggregation_transfer_backend != "fake" + ), "Prefill server does not support 'fake' as the transfer backend" + + if getattr(self, "disable_piecewise_cuda_graph", False): + self.disable_cuda_graph = True + logger.warning( + "Cuda graph is disabled for prefill server when piecewise cuda graph is not enabled." + ) + + if self.disaggregation_mode in ("prefill", "decode"): + if ( + envs.SGLANG_DISAGG_STAGING_BUFFER.get() + and self.disaggregation_transfer_backend not in ("mooncake", "nixl") + ): + raise ValueError( + f"SGLANG_DISAGG_STAGING_BUFFER requires " + f"disaggregation_transfer_backend='mooncake' or 'nixl', " + f"got '{self.disaggregation_transfer_backend}'." + ) + + def _allow_dsv4_decode_radix_speculative(self) -> bool: + return ( + envs.SGLANG_EXPERIMENTAL_DSV4_DECODE_RADIX_CACHE.get() + and self.speculative_algorithm.upper() == "EAGLE" + and self.speculative_eagle_topk == 1 + and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get() + and envs.SGLANG_EXPERIMENTAL_ONLINE_C128_MTP.get() + ) + def _handle_encoder_disaggregation(self): if self.enable_prefix_mm_cache and not self.encoder_only: raise ValueError( diff --git a/test/manual/unit/disaggregation/test_dsv4_decode_radix_cache.py b/test/manual/unit/disaggregation/test_dsv4_decode_radix_cache.py new file mode 100644 index 000000000000..f76833c4f027 --- /dev/null +++ b/test/manual/unit/disaggregation/test_dsv4_decode_radix_cache.py @@ -0,0 +1,254 @@ +from collections import namedtuple +from types import SimpleNamespace + +import torch + +from sglang.srt.disaggregation.decode import DecodePreallocQueue +from sglang.srt.environ import envs +from sglang.srt.managers.scheduler_components.batch_result_processor import ( + SchedulerBatchResultProcessor, +) +from sglang.srt.mem_cache.base_prefix_cache import MatchResult +from sglang.srt.mem_cache.common import maybe_cache_unfinished_req +from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool +from sglang.srt.mem_cache.kv_cache_builder import is_supported_dsv4_decode_radix_mtp +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + +Range = namedtuple("Range", ["start", "end"]) + + +def _make_queue(*, page_size: int, enable_decode_radix: bool = True): + queue = DecodePreallocQueue.__new__(DecodePreallocQueue) + queue.scheduler = SimpleNamespace( + server_args=SimpleNamespace( + disaggregation_decode_enable_radix_cache=enable_decode_radix + ), + running_batch=SimpleNamespace(reqs=[]), + ) + queue.token_to_kv_pool_allocator = SimpleNamespace(page_size=page_size) + queue.token_to_kv_pool = DeepSeekV4TokenToKVPool.__new__(DeepSeekV4TokenToKVPool) + queue.token_to_kv_pool.compression_ratios = [0, 4, 128] + queue.transfer_queue = SimpleNamespace(queue=[]) + return queue + + +def test_dsv4_decode_radix_prefix_is_c128_safe(): + queue = _make_queue(page_size=64) + + assert queue._dsv4_safe_prefix_len(127) == 0 + assert queue._dsv4_safe_prefix_len(128) == 128 + assert queue._dsv4_safe_prefix_len(383) == 256 + + +def test_dsv4_decode_radix_prefix_respects_page_alignment(): + queue = _make_queue(page_size=256) + + assert queue._dsv4_safe_prefix_len(255) == 0 + assert queue._dsv4_safe_prefix_len(256) == 256 + assert queue._dsv4_safe_prefix_len(511) == 256 + + +def test_dsv4_decode_radix_prefix_unchanged_when_disabled(): + queue = _make_queue(page_size=64, enable_decode_radix=False) + + assert queue._dsv4_safe_prefix_len(383) == 383 + + +def test_dsv4_singleflight_waits_for_large_inflight_shared_prefix(): + queue = _make_queue(page_size=256) + shared_prefix = list(range(5000)) + inflight_req = SimpleNamespace( + origin_input_ids=shared_prefix + [1], + dsv4_decode_radix_cache_prompt_once=True, + ) + waiting_req = SimpleNamespace(origin_input_ids=shared_prefix + [2]) + queue.transfer_queue.queue = [SimpleNamespace(req=inflight_req)] + + assert queue._should_wait_for_dsv4_inflight_prompt( + waiting_req, + prefix_len=0, + preallocated_reqs=[], + ) + assert not queue._should_wait_for_dsv4_inflight_prompt( + waiting_req, + prefix_len=4096, + preallocated_reqs=[], + ) + + +def test_dsv4_singleflight_ignores_short_shared_prefix(): + queue = _make_queue(page_size=256) + shared_prefix = list(range(1024)) + inflight_req = SimpleNamespace( + origin_input_ids=shared_prefix + [1], + dsv4_decode_radix_cache_prompt_once=True, + ) + waiting_req = SimpleNamespace(origin_input_ids=shared_prefix + [2]) + queue.transfer_queue.queue = [SimpleNamespace(req=inflight_req)] + + assert not queue._should_wait_for_dsv4_inflight_prompt( + waiting_req, + prefix_len=0, + preallocated_reqs=[], + ) + + +def test_allow_radix_cache_insert_once_bypasses_skip_once(): + req = SimpleNamespace( + skip_radix_cache_insert=True, + allow_radix_cache_insert_once=True, + ) + tree_cache = SimpleNamespace(num_cache_unfinished_req=0) + + def cache_unfinished_req(_req, **_kwargs): + tree_cache.num_cache_unfinished_req += 1 + + tree_cache.cache_unfinished_req = cache_unfinished_req + + maybe_cache_unfinished_req(req, tree_cache) + maybe_cache_unfinished_req(req, tree_cache) + + assert tree_cache.num_cache_unfinished_req == 1 + assert req.allow_radix_cache_insert_once is False + + +def test_dsv4_decode_radix_mtp_guard_allows_eagle_topk1_online_c128(): + server_args = SimpleNamespace(speculative_eagle_topk=1) + + with envs.SGLANG_OPT_USE_ONLINE_COMPRESS.override(True): + with envs.SGLANG_EXPERIMENTAL_ONLINE_C128_MTP.override(True): + assert is_supported_dsv4_decode_radix_mtp( + spec_algorithm=SpeculativeAlgorithm.EAGLE, + server_args=server_args, + ) + + +def test_dsv4_decode_radix_mtp_guard_rejects_non_minimal_paths(): + server_args = SimpleNamespace(speculative_eagle_topk=1) + + with envs.SGLANG_OPT_USE_ONLINE_COMPRESS.override(True): + with envs.SGLANG_EXPERIMENTAL_ONLINE_C128_MTP.override(True): + assert not is_supported_dsv4_decode_radix_mtp( + spec_algorithm=SpeculativeAlgorithm.EAGLE3, + server_args=server_args, + ) + assert not is_supported_dsv4_decode_radix_mtp( + spec_algorithm=SpeculativeAlgorithm.FROZEN_KV_MTP, + server_args=server_args, + ) + + server_args.speculative_eagle_topk = 2 + with envs.SGLANG_OPT_USE_ONLINE_COMPRESS.override(True): + with envs.SGLANG_EXPERIMENTAL_ONLINE_C128_MTP.override(True): + assert not is_supported_dsv4_decode_radix_mtp( + spec_algorithm=SpeculativeAlgorithm.EAGLE, + server_args=server_args, + ) + + +def test_dsv4_prompt_insert_uses_prompt_snapshot_and_restores_fill_ids(): + req = SimpleNamespace( + dsv4_decode_radix_cache_prompt_once=True, + dsv4_decode_radix_cache_prompt_len=3, + skip_radix_cache_insert=True, + allow_radix_cache_insert_once=False, + extend_range=Range(0, 5), + full_untruncated_fill_ids=[1, 2, 3, 4, 5], + extra_key=None, + req_pool_idx=0, + cache_protected_len=0, + swa_evicted_seqlen=4, + rid="test", + ) + req.get_fill_ids = lambda: [1, 2, 3, 4, 5][: req.extend_range.end] + + tree_cache = SimpleNamespace( + inserted_fill_ids=[], + inserted_swa_evicted_seqlens=[], + inserted_force_leaf_creation=[], + page_size=2, + sliding_window_size=2, + is_eagle=True, + ) + + def cache_unfinished_req(inserted_req, **_kwargs): + tree_cache.inserted_fill_ids.append(list(inserted_req.get_fill_ids())) + tree_cache.inserted_swa_evicted_seqlens.append(inserted_req.swa_evicted_seqlen) + tree_cache.inserted_force_leaf_creation.append( + inserted_req.force_radix_leaf_creation + ) + inserted_req.cache_protected_len = 2 + + tree_cache.cache_unfinished_req = cache_unfinished_req + token_to_kv_pool_allocator = SimpleNamespace(freed_swa_indices=[]) + + def free_swa(indices): + token_to_kv_pool_allocator.freed_swa_indices.append(indices.tolist()) + + token_to_kv_pool_allocator.free_swa = free_swa + req_to_token_pool = SimpleNamespace( + req_to_token=torch.tensor([[11, 12, 13, 14, 15]], dtype=torch.int64) + ) + processor = SimpleNamespace( + tree_cache=tree_cache, + token_to_kv_pool_allocator=token_to_kv_pool_allocator, + req_to_token_pool=req_to_token_pool, + ) + + SchedulerBatchResultProcessor._maybe_insert_dsv4_decode_radix_prompt(processor, req) + SchedulerBatchResultProcessor._maybe_insert_dsv4_decode_radix_prompt(processor, req) + + assert tree_cache.inserted_fill_ids == [[1, 2, 3]] + assert tree_cache.inserted_swa_evicted_seqlens == [2] + assert tree_cache.inserted_force_leaf_creation == [True] + assert token_to_kv_pool_allocator.freed_swa_indices == [[11, 12]] + assert req.extend_range.end == 5 + assert req.cache_protected_len == 2 + assert req.swa_evicted_seqlen == 4 + assert req.force_radix_leaf_creation is False + assert req.allow_radix_cache_insert_once is False + assert req.dsv4_decode_radix_cache_prompt_once is False + + +def test_dsv4_decode_radix_match_uses_full_match_for_full_only_leaf(): + queue = _make_queue(page_size=2) + node = object() + tree_cache = SimpleNamespace(captured_params=None) + + def supports_mamba(): + return False + + def match_prefix(params): + tree_cache.captured_params = params + return MatchResult( + device_indices=torch.tensor([1, 2], dtype=torch.int64), + last_device_node=node, + last_host_node=node, + best_match_node=node, + ) + + def inc_lock_ref(_node): + return SimpleNamespace(swa_uuid_for_lock=None) + + tree_cache.supports_mamba = supports_mamba + tree_cache.match_prefix = match_prefix + tree_cache.inc_lock_ref = inc_lock_ref + queue.tree_cache = tree_cache + + req = SimpleNamespace( + origin_input_ids=[1, 2, 3], + extra_key=None, + prefix_indices=None, + last_node=None, + last_host_node=None, + best_match_node=None, + host_hit_length=0, + swa_uuid_for_lock=None, + ) + + prefix_indices, prefix_len = DecodePreallocQueue._match_prefix_and_lock(queue, req) + + assert prefix_indices.tolist() == [1, 2] + assert prefix_len == 2 + assert tree_cache.captured_params.return_full_match is True + assert req.last_node is node diff --git a/test/registered/disaggregation/test_disaggregation_decode_radix_cache.py b/test/registered/disaggregation/test_disaggregation_decode_radix_cache.py index aa2796caa5c0..6c5f4ca0567e 100644 --- a/test/registered/disaggregation/test_disaggregation_decode_radix_cache.py +++ b/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 @@ def test_gsm8k_accuracy_two_passes(self): metrics_second = run_eval(args) print(f"Second run metrics: {metrics_second}") - self.assertGreater(metrics_first["score"], 0.80) - self.assertGreater(metrics_second["score"], 0.80) + self.assertGreater(metrics_first["score"], self.gsm8k_min_score) + self.assertGreater(metrics_second["score"], self.gsm8k_min_score) accuracy_drop = metrics_first["score"] - metrics_second["score"] self.assertLessEqual( diff --git a/test/registered/disaggregation/test_disaggregation_decode_radix_cache_swa.py b/test/registered/disaggregation/test_disaggregation_decode_radix_cache_swa.py new file mode 100644 index 000000000000..34cf04dd38e1 --- /dev/null +++ b/test/registered/disaggregation/test_disaggregation_decode_radix_cache_swa.py @@ -0,0 +1,49 @@ +"""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="base-c", runner_config="8-gpu-h20") + +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, so use a lower + # absolute floor while checking that the cached second pass does not regress. + 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() diff --git a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py index 7cba0af202b7..1f6cf2ec24e7 100644 --- a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py +++ b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py @@ -35,6 +35,7 @@ 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,63 @@ 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 = MagicMock(sliding_window_size=127) + 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() + + queue._reclaim_swa_tail_capacity(129, "req-1") + + 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() + + with self.assertRaisesRegex(RuntimeError, "needed=192, available=128"): + queue._reclaim_swa_tail_capacity(129, "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( @@ -302,6 +360,9 @@ def test_pop_preallocated_rechecks_budget_after_lock(self): 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() @@ -318,6 +379,10 @@ def test_pop_preallocated_rechecks_budget_after_lock(self): 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), @@ -348,21 +413,32 @@ def test_pop_preallocated_rechecks_budget_after_lock(self): 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. - queue._allocatable_token_budgets = MagicMock(side_effect=[8, 3]) + # The 4-token match is locked, then capped to zero because the whole + # 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) - self.assertEqual(queue._allocatable_token_budgets.call_count, 2) + queue.tree_cache.dec_swa_lock_only.assert_called_once_with(req.last_node, 123) + 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.""" diff --git a/test/registered/unit/mem_cache/test_hisparse_allocator.py b/test/registered/unit/mem_cache/test_hisparse_allocator.py index 592a5d528a14..51c836d6ee91 100644 --- a/test/registered/unit/mem_cache/test_hisparse_allocator.py +++ b/test/registered/unit/mem_cache/test_hisparse_allocator.py @@ -110,6 +110,7 @@ def write(self, indices, values): device=torch.device("cpu"), page_size=64, 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), ) diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 677e13b80511..1ebb63b9636d 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -238,6 +238,50 @@ def test_pd_decode_radix_cache_allows_mooncake_tcp(self): self.assertFalse(server_args.disable_radix_cache) self.assertEqual(server_args.disaggregation_transfer_backend, "mooncake") + def test_pd_decode_radix_cache_rejects_speculative_by_default(self): + with self.assertRaises(ValueError) as context: + ServerArgs( + model_path="dummy", + disaggregation_mode="decode", + disaggregation_decode_enable_radix_cache=True, + disaggregation_transfer_backend="mooncake", + speculative_algorithm="EAGLE", + speculative_eagle_topk=1, + ) + + self.assertIn("speculative decoding", str(context.exception)) + + def test_pd_decode_radix_cache_allows_dsv4_eagle_topk1_online_c128(self): + with envs.SGLANG_EXPERIMENTAL_DSV4_DECODE_RADIX_CACHE.override(True): + with envs.SGLANG_OPT_USE_ONLINE_COMPRESS.override(True): + with envs.SGLANG_EXPERIMENTAL_ONLINE_C128_MTP.override(True): + server_args = ServerArgs( + model_path="dummy", + disaggregation_mode="decode", + disaggregation_decode_enable_radix_cache=True, + disaggregation_transfer_backend="mooncake", + speculative_algorithm="EAGLE", + speculative_eagle_topk=1, + ) + + self.assertFalse(server_args.disable_radix_cache) + + def test_pd_decode_radix_cache_rejects_dsv4_eagle_topk_gt1(self): + with envs.SGLANG_EXPERIMENTAL_DSV4_DECODE_RADIX_CACHE.override(True): + with envs.SGLANG_OPT_USE_ONLINE_COMPRESS.override(True): + with envs.SGLANG_EXPERIMENTAL_ONLINE_C128_MTP.override(True): + with self.assertRaises(ValueError) as context: + ServerArgs( + model_path="dummy", + disaggregation_mode="decode", + disaggregation_decode_enable_radix_cache=True, + disaggregation_transfer_backend="mooncake", + speculative_algorithm="EAGLE", + speculative_eagle_topk=2, + ) + + self.assertIn("speculative decoding", str(context.exception)) + class TestSkipTokenizerInit(unittest.TestCase): def test_skip_tokenizer_worker_counts(self):