diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 09254f314b0c..fc0c7c995c49 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -554,6 +554,19 @@ def merge(self, other: MultimodalInputs): # other args would be kept intact +def collect_cached_positions(reqs): + """Aggregate per-request cached non-contiguous RoPE positions across a batch. + + Returns a per-request list of Optional[torch.Tensor] when at least one request + has cached positions; None otherwise (signal to ForwardBatch.init_new that the + legacy contiguous-positions path applies). + """ + positions = [getattr(r, "cached_positions", None) for r in reqs] + if all(p is None for p in positions): + return None + return positions + + class Req(ReqDllmMixin): """The input and output status of a request.""" @@ -640,6 +653,18 @@ def __init__( # State indicating whether the reasoning phase has finished (only meaningful when require_reasoning is True) self._is_reasoning_over = False self.reasoning_tokens = 0 + # Absolute position (in the origin_input_ids + output_ids index space) of the + # first token after . Set by update_reasoning_tokens when the + # boundary is detected; consumed at request finish under --no-cache-thoughts to + # split the request's KV between freed-only thoughts and radix-inserted answer. + self.answer_start_position: Optional[int] = None + # Per-token original RoPE positions for the prefix matched in the radix cache. + # Set by init_next_round_input when match_prefix returns non-None positions + # (i.e. the cached entry was inserted with non-contiguous positions, e.g. via + # the --no-cache-thoughts split path). Consumed at batch construction so the + # ForwardBatch's positions tensor lines up with the rotation baked into the + # cached K vectors. None means "use legacy contiguous positions". + self.cached_positions: Optional[torch.Tensor] = None # Sampling info if isinstance(sampling_params.custom_params, dict): @@ -996,6 +1021,7 @@ def init_next_round_input( match_result.host_hit_length, match_result.mamba_branching_seqlen, ) + self.cached_positions = match_result.original_positions if match_result.cache_protected_len is not None: self.cache_protected_len = match_result.cache_protected_len else: @@ -1291,6 +1317,10 @@ def update_reasoning_tokens(self, token_id, think_end_id): end_pos = token_id.index(think_end_id) self.reasoning_tokens += end_pos + 1 self._is_reasoning_over = True + # The answer begins immediately after . Position is in the absolute + # token-index space (origin_input_ids + output_ids), which equals the RoPE + # position when the request was decoded with contiguous positions. + self.answer_start_position = len(self.origin_input_ids) + self.reasoning_tokens except ValueError: self.reasoning_tokens += len(token_id) @@ -1375,6 +1405,14 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # For extend and mixed chunekd prefill prefix_lens: List[int] = None + # Per-request original RoPE positions for the matched prefix when the cache hit + # carried non-contiguous positions (e.g. from --no-cache-thoughts). None means no + # request in the batch has such positions; the contiguous-positions path applies. + cached_positions_per_req: Optional[List[Optional[torch.Tensor]]] = None + # Per-request RoPE offset (one int per req, device tensor) added to (seq_len - 1) + # at decode time so decode positions continue from where a non-contiguous prefill + # cache hit left off. None when no req in the batch needs an offset. + position_offsets: Optional[torch.Tensor] = None extend_lens: List[int] = None extend_num_tokens: Optional[int] = None decoding_reqs: List[Req] = None @@ -1606,6 +1644,17 @@ def prepare_for_extend(self): # Set batch fields needed by alloc_for_extend self.prefix_lens = prefix_lens + self.cached_positions_per_req = collect_cached_positions(reqs) + if self.cached_positions_per_req is not None: + from sglang.srt.mem_cache.common import derive_position_offsets + + offsets_list = derive_position_offsets( + prefix_lens, self.cached_positions_per_req + ) + if offsets_list is not None: + self.position_offsets = torch.tensor( + offsets_list, dtype=torch.int64, pin_memory=_pin + ).to(self.device, non_blocking=True) self.extend_lens = extend_lens self.seq_lens = seq_lens_tensor self.seq_lens_cpu = seq_lens_cpu diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 122ddb387486..e3b07b510328 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1778,7 +1778,6 @@ def handle_generate_request( time_stats=recv_req.time_stats, ) req.tokenizer = self.tokenizer - if self.disaggregation_mode != DisaggregationMode.NULL: # Invalid request for disaggregated mode if ( diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index be219339cef6..8713fa841e57 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -61,6 +61,13 @@ class InsertParams: chunked: bool = False priority: int = 0 + # Per-token original RoPE positions for the inserted tokens. When None (default), + # the cache uses standard contiguous positions and behavior is unchanged. When + # provided, length must equal len(key.token_ids); positions are stored alongside + # the kv_indices so subsequent match_prefix calls can return them, letting the + # scheduler use non-contiguous positions in the ForwardBatch. + original_positions: Optional[torch.Tensor] = None + @dataclasses.dataclass class InsertResult: @@ -145,6 +152,7 @@ class MatchResult(NamedTuple): host_hit_length: int = 0 mamba_branching_seqlen: Optional[int] = None cache_protected_len: Optional[int] = None + original_positions: Optional[torch.Tensor] = None class BasePrefixCache(ABC, PrefixCacheTrait): diff --git a/python/sglang/srt/mem_cache/chunk_cache.py b/python/sglang/srt/mem_cache/chunk_cache.py index 5d58bcde530c..2afd5bfa4696 100644 --- a/python/sglang/srt/mem_cache/chunk_cache.py +++ b/python/sglang/srt/mem_cache/chunk_cache.py @@ -65,7 +65,14 @@ def insert(self, params: InsertParams) -> InsertResult: # ChunkCache does not support prefix caching, so insert is a no-op return InsertResult(prefix_len=0) - def cache_finished_req(self, req: Req, is_insert: bool = True): + def cache_finished_req(self, req: Req, is_insert: bool = True, split=None): + # ChunkCache does not implement prefix caching, so it cannot honor a + # split-insertion request from --no-cache-thoughts. Fall back to the + # default behavior; the caller's split is dropped. + if split is not None: + from sglang.srt.mem_cache.common import warn_split_unsupported_once + + warn_split_unsupported_once("ChunkCache") kv_committed_len = req.pop_committed_kv_cache() # For decode server: if req.output_ids is empty, we want to free all req.origin_input_ids kv_indices = self.req_to_token_pool.req_to_token[ diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 5a759ed11bcd..e8d73bd88093 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -1,7 +1,8 @@ from __future__ import annotations +import dataclasses import logging -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, List, Optional import torch import triton @@ -17,6 +18,213 @@ if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req, ScheduleBatch + +_split_unsupported_warned = set() + + +def warn_split_unsupported_once(backend_name: str) -> None: + """Emit a one-time warning when a prefix-cache backend that does not implement + --no-cache-thoughts split insertion receives a split kwarg. The split is dropped + and the backend falls back to its default cache_finished_req behavior, so the + feature gracefully no-ops on that backend. + """ + if backend_name in _split_unsupported_warned: + return + _split_unsupported_warned.add(backend_name) + logger.warning( + "%s does not implement --no-cache-thoughts split insertion; " + "thoughts will be cached normally on this backend (no-op).", + backend_name, + ) + + +def derive_position_offsets( + extend_prefix_lens: List[int], + cached_positions_per_req: List[Optional["torch.Tensor"]], +) -> Optional[List[int]]: + """Given per-request cached RoPE positions, return a per-request offset to add + to (seq_len - 1) at decode time so decode positions continue from where the + non-contiguous prefill cache hit left off. + + The offset captures the gap in RoPE space caused by tokens that exist in the + cached entry's position layout but not in the contiguous-token-count layout: + offset[i] = max(cached_positions[i]) - (extend_prefix_lens[i] - 1) + A request without cached positions contributes 0 (no offset). + + Returns None if every entry is None (i.e. no request in the batch needs an + offset; the legacy clamp(seq_lens - 1) path applies). + """ + if all(p is None for p in cached_positions_per_req): + return None + offsets: List[int] = [] + for prefix_len, positions in zip(extend_prefix_lens, cached_positions_per_req): + if positions is None or len(positions) == 0: + offsets.append(0) + else: + offsets.append(int(positions.max().item()) - (int(prefix_len) - 1)) + return offsets + + +def derive_extend_position_start( + extend_prefix_lens: List[int], + cached_positions_per_req: List[Optional["torch.Tensor"]], +) -> Optional[List[int]]: + """Given per-request cached RoPE positions, return the starting RoPE position for + each request's extend (prefill) tokens. + + Args: + extend_prefix_lens: per-request count of cached prefix tokens. + cached_positions_per_req: per-request tensor of cached original positions, + or None if no positions are cached for that request. + + Returns: + None if every entry is None (i.e. no request has non-contiguous cached + positions, so the legacy contiguous-positions path applies). Otherwise a + list of per-request integer starts: max(cached_positions) + 1 when cached + positions exist, else extend_prefix_lens[i] (legacy). + """ + if all(p is None for p in cached_positions_per_req): + return None + starts: List[int] = [] + for prefix_len, positions in zip(extend_prefix_lens, cached_positions_per_req): + if positions is None or len(positions) == 0: + starts.append(int(prefix_len)) + else: + starts.append(int(positions.max().item()) + 1) + return starts + + +@dataclasses.dataclass +class NoCacheThoughtsSplit: + """Plan for caching a finished reasoning request without its thought span. + + Dropping the generated ... tokens shifts every answer token's index + in the cached sequence DOWN by len(thoughts), but its KV still physically sits in + the slot decode gave it. Paged KV requires slot % page_size == index % page_size, + so the answer's KV must be RELOCATED into the slots the thoughts are vacating (which + are page-congruent to the answer's new indices) before it can be cached and safely + extended later. cache_finished_req consumes this plan as: + 1. move_kv_cache(move_dst, move_src) -> slide the answer's KV left + 2. insert virtual_token_ids/positions, values from virtual_kv_indices[:aligned] + 3. free virtual_kv_indices[aligned:] -> one page-aligned cut for the dead tail + """ + + # input + post- answer (thoughts removed); the cached key sequence. + virtual_token_ids: List[int] + # The request's FULL original contiguous slot span S[0:total_len]. After the move, + # S[0:kept_len] holds input+answer and S[kept_len:] is the dead tail (stale thought + # slots + the answer's page-unaligned remainder), freed in one page-aligned cut. + virtual_kv_indices: torch.Tensor + # Original RoPE positions of the kept tokens (gapped where thoughts were); the K + # vectors already encode these, so positions are preserved while only slots move. + virtual_positions: torch.Tensor + # Relocation: move the answer's current slots (move_src = S[answer_start:total_len]) + # into the page-congruent destination (move_dst = S[input_len:kept_len]). Empty when + # there is no thought span or no answer (nothing to relocate). + move_src: torch.Tensor + move_dst: torch.Tensor + + +def split_kv_for_no_cache_thoughts( + origin_input_ids: List[int], + output_ids: List[int], + req_to_token_slot: torch.Tensor, + answer_start_position: int, + committed_len: int, +) -> NoCacheThoughtsSplit: + """Compute the split-insertion tensors for a finished reasoning request. + + When --no-cache-thoughts is enabled and a request emits ``, the tokens + between input_end and the `` boundary are thoughts that must NOT be + registered in the cross-request radix cache; the post-`` answer must + be inserted with the original RoPE positions preserved. + + Args: + origin_input_ids: input token ids (positions 0..len(input)-1). + output_ids: generated token ids (positions len(input)..len(input)+len(output)-1). + req_to_token_slot: 1D tensor of kv_indices, one per token in the request's + sequence (length must be >= committed_len). + answer_start_position: absolute position (in the input+output index space) + of the first answer token, i.e. the token immediately after ``. + committed_len: number of tokens whose KV is actually committed + (``req.kv_committed_len``). The final generated token can appear in + output_ids while its KV slot is uncommitted (overlap scheduling), in which + case its req_to_token entry is the zero/unwritten sentinel (reserved page 0). + Walking past committed_len would move/free that sentinel and double-free + page 0, so we cap here exactly as the normal cache_finished_req path does. + + Returns: + NoCacheThoughtsSplit (see that dataclass for field semantics): the cached key + sequence (input + answer), the request's committed slot span, the kept tokens' + original RoPE positions, and the move_src/move_dst slot tensors that relocate + the answer's KV into page-congruent slots. + + When answer_start_position >= committed_len, the answer hasn't been committed + (e.g. the request was cut off mid-thought) and the result caches only the input + prompt slice (no relocation). + """ + input_len = len(origin_input_ids) + # Cap at the committed KV length (see committed_len arg): never touch the + # uncommitted trailing token's unwritten (page-0) slot. + total_len = min(input_len + len(output_ids), committed_len) + + # Clamp answer_start so we behave sanely if reasoning never finished. + answer_start = min(answer_start_position, total_len) + + think_len = answer_start - input_len # decoded thought tokens to drop + answer_count = max(total_len - answer_start, 0) + answer_output_offset = answer_start - input_len # index into output_ids + kept_len = input_len + answer_count # cached sequence length (= total_len - think_len) + + # The full input prompt (including any \n priming tail from + # add_generation_prompt) stays in the cached entry — TITO rollouts feed + # turn N+1's input as raw token IDs that include turn N's prompt verbatim, + # so keeping the priming preserves cache alignment in that flow. The answer slice + # is bounded by answer_count so an uncommitted trailing token is never cached. + virtual_token_ids = list(origin_input_ids) + list( + output_ids[answer_output_offset : answer_output_offset + answer_count] + if answer_count > 0 + else [] + ) + + # Original RoPE positions of the kept tokens: input is contiguous, then the answer + # keeps its ORIGINAL positions (a gap where the thoughts were). The K vectors already + # encode these positions, so we must not renumber them; only the physical slots move. + kept_positions = list(range(input_len)) + list(range(answer_start, total_len)) + + device = req_to_token_slot.device + slots = req_to_token_slot.to(torch.int64) + + # Hand cache_finished_req the request's FULL original slot span. It takes cached + # values from slots[:page_aligned(kept_len)] (after the move) and frees + # slots[page_aligned(kept_len):] in a single page-aligned cut — that one cut covers + # both the stale thought slots and the answer's unaligned tail with no boundary + # double-free. + virtual_kv_indices = slots[:total_len].clone() + + # Relocate the answer's KV LEFT by think_len slots so each answer token lands on the + # slot that originally held its NEW index (slot % page_size == index % page_size). + # The destination slots[input_len:kept_len] are exactly the thought slots plus the + # answer's leading slots — all owned privately by this finished request, so the move + # cannot disturb the shared input prefix already in the radix. Skip when there is no + # thought span (already aligned) or no answer (nothing to relocate). + if think_len > 0 and answer_count > 0: + move_src = slots[answer_start:total_len].clone() + move_dst = slots[input_len:kept_len].clone() + else: + move_src = torch.empty((0,), dtype=torch.int64, device=device) + move_dst = torch.empty((0,), dtype=torch.int64, device=device) + + virtual_positions = torch.tensor(kept_positions, dtype=torch.int64, device=device) + + return NoCacheThoughtsSplit( + virtual_token_ids=virtual_token_ids, + virtual_kv_indices=virtual_kv_indices, + virtual_positions=virtual_positions, + move_src=move_src, + move_dst=move_dst, + ) + # Needs 2 + 1 slots for mamba request with prefix cache. 2 for ping pong cache, 1 for running mamba state. MAMBA_STATE_PER_REQ_PREFIX_CACHE = 3 MAMBA_STATE_PER_REQ_NO_CACHE = 1 @@ -476,7 +684,27 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr req.mamba_pool_idx = None return - tree_cache.cache_finished_req(req, is_insert=is_insert) + server_args = get_global_server_args() + if ( + is_insert + and getattr(server_args, "no_cache_thoughts", False) + and getattr(req, "require_reasoning", False) + and getattr(req, "answer_start_position", None) is not None + ): + # Skip the thought tokens from the shared prefix cache; insert only the + # input + post- answer slice, preserving original RoPE positions + # for the answer (the input prompt keeps its contiguous positions). + req_to_token_slot = tree_cache.req_to_token_pool.req_to_token[req.req_pool_idx] + split = split_kv_for_no_cache_thoughts( + origin_input_ids=req.origin_input_ids, + output_ids=req.output_ids, + req_to_token_slot=req_to_token_slot, + answer_start_position=req.answer_start_position, + committed_len=req.kv_committed_len, + ) + tree_cache.cache_finished_req(req, is_insert=is_insert, split=split) + else: + tree_cache.cache_finished_req(req, is_insert=is_insert) # FIXME: SessionAwareCache.cache_finished_req sets req_pool_idx = None to # transfer KV ownership to the SessionSlot, so we skip the remaining diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index d07702cf1efd..8fae60d0a982 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -514,8 +514,18 @@ def insert(self, params: InsertParams) -> InsertResult: ) return InsertResult(prefix_len=prefix_len, mamba_exist=mamba_exist) - def cache_finished_req(self, req: Req, is_insert: bool = True) -> None: - """Cache request when it finishes.""" + def cache_finished_req( + self, req: Req, is_insert: bool = True, split=None + ) -> None: + """Cache request when it finishes. + + MambaRadixCache does not yet implement split insertion; when --no-cache-thoughts + passes a split, fall back to the default behavior (thoughts cached normally). + """ + if split is not None: + from sglang.srt.mem_cache.common import warn_split_unsupported_once + + warn_split_unsupported_once("MambaRadixCache") kv_committed_len = req.pop_committed_kv_cache() if self.disable: kv_indices = self.req_to_token_pool.req_to_token[ diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index e4c158cda9f6..1fbe205ff168 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1547,6 +1547,21 @@ def get_value_buffer(self, layer_id: int): def get_kv_buffer(self, layer_id: int): return self.get_key_buffer(layer_id), self.get_value_buffer(layer_id) + def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): + # Relocate the compressed MLA latent for a set of slots. Used by + # --no-cache-thoughts to slide a finished request's answer KV into + # page-congruent slots before caching it. Native indexed copy over every + # layer buffer; the advanced-indexed RHS (kv_cache[src]) is materialized + # before the scatter write, so this is correct even when src and tgt + # overlap. Works for any store_dtype, including the NSA FP8 byte layout, + # since it copies whole slot rows verbatim. + if tgt_loc.numel() == 0: + return + tgt = tgt_loc.view(-1).long() + src = src_loc.view(-1).long() + for kv_cache in self.kv_buffer: + kv_cache[tgt] = kv_cache[src] + def set_kv_buffer( self, layer: RadixAttention, diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 7d1616037243..f2ca18a10000 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -127,6 +127,9 @@ def __init__(self, id: Optional[int] = None, priority: int = 0): self.parent: TreeNode = None self.key: RadixKey = None self.value: Optional[torch.Tensor] = None + # Per-token original RoPE positions, parallel to `value`. None when the node was + # inserted without positions (standard behavior). + self.positions: Optional[torch.Tensor] = None self.lock_ref = 0 self.last_access_time = time.monotonic() self.creation_time = time.monotonic() @@ -432,15 +435,30 @@ def empty_match_result(): if len(key) == 0: return empty_match_result() - value, last_node = self._match_prefix_helper(self.root_node, key) + value, positions, last_node = self._match_prefix_helper(self.root_node, key) if value: value = torch.cat(value) else: value = torch.empty((0,), dtype=torch.int64, device=self.device) + # If any matched node carried original positions, concatenate and return them. + # Otherwise return None for backwards compatibility. + if positions and any(p is not None for p in positions): + # Replace any None entries (mixed-mode tree) with contiguous fallback positions. + # This should be rare; tree builders typically use positions consistently. + concat_positions = torch.cat( + [ + p if p is not None else torch.empty((0,), dtype=torch.int64, device=self.device) + for p in positions + ] + ) + original_positions = concat_positions + else: + original_positions = None return MatchResult( device_indices=value, last_device_node=last_node, last_host_node=last_node, + original_positions=original_positions, ) def insert(self, params: InsertParams) -> InsertResult: @@ -451,22 +469,95 @@ def insert(self, params: InsertParams) -> InsertResult: value = params.value priority = params.priority chunked = params.chunked + positions = params.original_positions if value is None: value = torch.tensor(key.token_ids, dtype=torch.int64) + if positions is not None and len(positions) != len(key.token_ids): + raise ValueError( + f"original_positions length {len(positions)} does not match key " + f"token_ids length {len(key.token_ids)}" + ) + key, value = self.maybe_bigram_convert(key, value) - prefix_len = self._insert_helper(self.root_node, key, value, priority, chunked) + prefix_len = self._insert_helper( + self.root_node, key, value, priority, chunked, positions + ) return InsertResult(prefix_len=prefix_len) - def cache_finished_req(self, req: Req, is_insert: bool = True): - """Cache request when it finishes.""" + def cache_finished_req(self, req: Req, is_insert: bool = True, split=None): + """Cache request when it finishes. + + Args: + req: the finished request whose KV is being committed. + is_insert: when False, free the request's kv_indices without inserting them + into the radix tree (e.g. abort / retract paths). + split: when provided (a NoCacheThoughtsSplit from + ``sglang.srt.mem_cache.common.split_kv_for_no_cache_thoughts``), cache the + thought-stripped sequence: relocate the answer's KV into page-congruent + slots (``split.move_dst`` <- ``split.move_src``), insert the answer with + its original RoPE positions preserved, then free the page-aligned dead + tail (stale thoughts + unaligned answer remainder). + """ # In deterministic mode, disable finished request insertion to radix cache if self.disable_finished_insert: is_insert = False kv_committed_len = req.pop_committed_kv_cache() + + if split is not None: + # Skip the per-req KV-pool lookup; use the pre-computed plan. + # split.virtual_kv_indices is the request's FULL original slot span, so a + # single free here returns everything when we are not inserting. + if self.disable or not is_insert: + self.token_to_kv_pool_allocator.free(split.virtual_kv_indices) + self.dec_lock_ref(req.last_node) + return + # Relocate the answer's KV left into the slots the thoughts are vacating, so + # each cached answer token sits in a slot page-congruent to its NEW index in + # the thought-stripped sequence. Paged KV requires slot % page == index % + # page; without this the answer is unreachable for safe extension and trips + # the allocator's page-alignment assert. The move reads its source fully + # before writing (overlap-safe); the vacated thought slots are reclaimed below + # as part of the page-aligned dead tail. + if split.move_src.numel() > 0: + self.token_to_kv_pool_allocator.get_kvcache().move_kv_cache( + split.move_dst, split.move_src + ) + keys = ( + convert_to_bigram_key(split.virtual_token_ids) + if self.is_eagle + else split.virtual_token_ids + ) + keys = page_align_keys(keys, self.page_size) + values = split.virtual_kv_indices[: len(keys)].to( + dtype=torch.int64, copy=True + ) + positions = split.virtual_positions[: len(keys)].to( + dtype=torch.int64, copy=True + ) + radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle) + priority = getattr(req, "priority", 0) or 0 + result = self.insert( + InsertParams( + key=radix_key, + value=values, + priority=priority, + original_positions=positions, + ) + ) + new_prefix_len = result.prefix_len + self.token_to_kv_pool_allocator.free( + split.virtual_kv_indices[req.cache_protected_len : new_prefix_len] + ) + # One page-aligned cut frees the dead tail: the stale thought slots plus the + # answer's page-unaligned remainder. Single free() -> no boundary double-free. + self.token_to_kv_pool_allocator.free(split.virtual_kv_indices[len(keys) :]) + self.dec_lock_ref(req.last_node) + return + if self.disable: kv_indices = self.req_to_token_pool.req_to_token[ req.req_pool_idx, :kv_committed_len @@ -671,6 +762,7 @@ def _match_prefix_helper(self, node: TreeNode, key: RadixKey): child_key = self.get_child_key_fn(key) value = [] + positions = [] # Parallel list of per-node positions tensors (or Nones). while len(key) > 0 and child_key in node.children.keys(): child = node.children[child_key] child.last_access_time = access_time @@ -678,17 +770,19 @@ def _match_prefix_helper(self, node: TreeNode, key: RadixKey): if prefix_len < len(child.key): new_node = self._split_node(child.key, child, prefix_len) value.append(new_node.value) + positions.append(new_node.positions) node = new_node break else: value.append(child.value) + positions.append(child.positions) node = child key = key[prefix_len:] if len(key): child_key = self.get_child_key_fn(key) - return value, node + return value, positions, node def _split_node(self, key: RadixKey, child: TreeNode, split_len: int): # new_node -> child @@ -700,6 +794,10 @@ def _split_node(self, key: RadixKey, child: TreeNode, split_len: int): new_node.lock_ref = child.lock_ref new_node.key = child.key[:split_len] new_node.value = child.value[:split_len].clone() + # Split positions in lockstep with value. + if child.positions is not None: + new_node.positions = child.positions[:split_len].clone() + child.positions = child.positions[split_len:].clone() child.parent = new_node child.key = child.key[split_len:] child.value = child.value[split_len:].clone() @@ -727,6 +825,7 @@ def _insert_helper( value, priority: int = 0, chunked: bool = False, + positions: Optional[torch.Tensor] = None, ): # Convert None priority to 0 if priority is None: @@ -748,6 +847,8 @@ def _insert_helper( total_prefix_length += prefix_len key = key[prefix_len:] value = value[prefix_len:] + if positions is not None: + positions = positions[prefix_len:] if prefix_len < len(node.key): new_node = self._split_node(node.key, node, prefix_len) @@ -765,6 +866,8 @@ def _insert_helper( new_node.parent = node new_node.key = key new_node.value = value.clone() + if positions is not None: + new_node.positions = positions.clone() self._inc_hit_count(new_node, chunked) node.children[child_key] = new_node self.evictable_size_ += len(key) diff --git a/python/sglang/srt/mem_cache/radix_cache_cpp.py b/python/sglang/srt/mem_cache/radix_cache_cpp.py index 66f9fad96ad7..62eabdf19c21 100644 --- a/python/sglang/srt/mem_cache/radix_cache_cpp.py +++ b/python/sglang/srt/mem_cache/radix_cache_cpp.py @@ -168,8 +168,16 @@ def protected_size(self): def total_size(self): return self.tree.total_size() - def cache_finished_req(self, req: Req, is_insert: bool = True): - """Cache request when it finishes.""" + def cache_finished_req(self, req: Req, is_insert: bool = True, split=None): + """Cache request when it finishes. + + RadixCacheCpp does not yet implement split insertion; when --no-cache-thoughts + passes a split, fall back to the default behavior (thoughts cached normally). + """ + if split is not None: + from sglang.srt.mem_cache.common import warn_split_unsupported_once + + warn_split_unsupported_once("RadixCacheCpp") assert req.req_pool_idx is not None kv_committed_len = req.pop_committed_kv_cache() token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len] diff --git a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py index 9a82aa31f4ef..31f4613dca7c 100644 --- a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py +++ b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py @@ -211,10 +211,21 @@ def match_prefix(self, params: MatchPrefixParams) -> MatchResult: # type: ignor return base_res - def cache_finished_req(self, req: Req, is_insert: bool = True) -> None: # type: ignore[override] - """On request completion, insert device KV into radix and store to LMCache.""" + def cache_finished_req( + self, req: Req, is_insert: bool = True, split=None + ) -> None: # type: ignore[override] + """On request completion, insert device KV into radix and store to LMCache. + + When --no-cache-thoughts passes a split, the inner radix accepts it (and stores + only the answer slice). The LMCache offload that follows reads req.fill_ids + directly and may offload more tokens than the radix retained; treat this as a + soft-no-op for the offload portion until proper split-aware offload lands. + """ + if split is not None: + from sglang.srt.mem_cache.common import warn_split_unsupported_once - super().cache_finished_req(req, is_insert=is_insert) + warn_split_unsupported_once("LMCRadixCache") + super().cache_finished_req(req, is_insert=is_insert, split=split) if not is_insert: return diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 34df1617dd61..c6cb5f19aa65 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -441,8 +441,18 @@ def insert(self, params: InsertParams) -> InsertResult: ) return InsertResult(prefix_len=prefix_len) - def cache_finished_req(self, req: Req, is_insert: bool = True) -> None: - """Cache request when it finishes.""" + def cache_finished_req( + self, req: Req, is_insert: bool = True, split=None + ) -> None: + """Cache request when it finishes. + + SWARadixCache does not yet implement split insertion; when --no-cache-thoughts + passes a split, fall back to the default behavior (thoughts cached normally). + """ + if split is not None: + from sglang.srt.mem_cache.common import warn_split_unsupported_once + + warn_split_unsupported_once("SWARadixCache") kv_committed_len = req.pop_committed_kv_cache() if self.disable: kv_indices = self.req_to_token_pool.req_to_token[ diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index eaecdc54bcf4..ce4716efb2b9 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -544,7 +544,10 @@ def init_new( # Init position information if ret.forward_mode.is_decode() or ret.forward_mode.is_target_verify(): if ret.positions is None: - ret.positions = clamp_position(batch.seq_lens) + ret.positions = clamp_position( + batch.seq_lens, + position_offsets=getattr(batch, "position_offsets", None), + ) else: assert isinstance(batch.extend_seq_lens, list) assert isinstance(batch.extend_prefix_lens, list) @@ -555,11 +558,18 @@ def init_new( batch.extend_prefix_lens, dtype=torch.int32 ).to(device, non_blocking=True) ret.extend_num_tokens = batch.extend_num_tokens - positions, ret.extend_start_loc = compute_position( - model_runner.server_args.attention_backend, - ret.extend_prefix_lens, - ret.extend_seq_lens, - ret.extend_num_tokens, + # Honor per-request cached non-contiguous positions from cache hits when + # available; otherwise behaves identically to compute_position. + positions, ret.extend_start_loc = build_extend_positions( + attn_backend=model_runner.server_args.attention_backend, + extend_prefix_lens=ret.extend_prefix_lens, + extend_seq_lens=ret.extend_seq_lens, + extend_num_tokens=ret.extend_num_tokens, + extend_prefix_lens_cpu=batch.extend_prefix_lens, + cached_positions_per_req=getattr( + batch, "cached_positions_per_req", None + ), + device=device, ) if ret.positions is None: ret.positions = positions @@ -1100,13 +1110,63 @@ def __repr__(self) -> str: return f"PPProxyTensors(tensors={self.tensors})" +def build_extend_positions( + attn_backend: str, + extend_prefix_lens: torch.Tensor, + extend_seq_lens: torch.Tensor, + extend_num_tokens: int, + extend_prefix_lens_cpu: List[int], + cached_positions_per_req: Optional[List[Optional[torch.Tensor]]], + device, +): + """Build extend-token positions, honoring per-request cached non-contiguous positions. + + Returns (positions, extend_start_loc) matching compute_position's signature. When + cached_positions_per_req is None or contains only None entries, positions are + contiguous (legacy behavior). Otherwise the per-request start position is + max(cached_positions[i]) + 1 for cached requests, and extend_prefix_lens_cpu[i] + for non-cached requests. + """ + from sglang.srt.mem_cache.common import derive_extend_position_start + + extend_position_start_tensor = None + if cached_positions_per_req is not None: + starts = derive_extend_position_start( + extend_prefix_lens_cpu, cached_positions_per_req + ) + if starts is not None: + extend_position_start_tensor = torch.tensor( + starts, dtype=torch.int64 + ).to(device, non_blocking=True) + + return compute_position( + attn_backend, + extend_prefix_lens, + extend_seq_lens, + extend_num_tokens, + extend_position_start=extend_position_start_tensor, + ) + + def compute_position( attn_backend: str, extend_prefix_lens: torch.Tensor, extend_seq_lens: torch.Tensor, extend_seq_lens_sum: int, + extend_position_start: Optional[torch.Tensor] = None, ): - if support_triton(attn_backend): + """Compute positions for the extend (prefill) tokens. + + extend_position_start (optional): per-request override for the starting RoPE + position of the extend tokens. When None, positions are contiguous and start + at extend_prefix_lens[i]. When provided, positions for request i are + [extend_position_start[i], extend_position_start[i] + 1, ...]; used by cache + hits whose cached entries carry non-contiguous original positions. + """ + if support_triton(attn_backend) and extend_position_start is None: + # The fused triton kernel uses extend_prefix_lens as the position start. + # The override path is not yet implemented there; fall through to the + # torch path when an override is requested. positions, extend_start_loc = compute_position_triton( extend_prefix_lens, extend_seq_lens, @@ -1114,7 +1174,7 @@ def compute_position( ) else: positions, extend_start_loc = compute_position_torch( - extend_prefix_lens, extend_seq_lens + extend_prefix_lens, extend_seq_lens, extend_position_start ) return positions, extend_start_loc @@ -1176,14 +1236,32 @@ def compute_position_kernel( def compute_position_torch( - extend_prefix_lens: torch.Tensor, extend_seq_lens: torch.Tensor + extend_prefix_lens: torch.Tensor, + extend_seq_lens: torch.Tensor, + extend_position_start: Optional[torch.Tensor] = None, ): + """Compute per-token positions for the extend (prefill) tokens of a batch. + + Args: + extend_prefix_lens: per-request count of cached prefix tokens (KV slot count). + extend_seq_lens: per-request count of new tokens being prefilled. + extend_position_start: optional per-request override for the first RoPE + position of the extend tokens. When None, positions start at + extend_prefix_lens[i] (contiguous: token i sits at RoPE position i). + When provided, positions for request i are + [extend_position_start[i], extend_position_start[i] + 1, ...]. This + supports cache hits whose stored kv has non-contiguous positions + (e.g. with gaps where thoughts were skipped). + """ + starts = ( + extend_position_start if extend_position_start is not None else extend_prefix_lens + ) positions = torch.cat( [ torch.arange( - prefix_len, prefix_len + extend_len, device=extend_prefix_lens.device + start, start + extend_len, device=extend_prefix_lens.device ) - for prefix_len, extend_len in zip(extend_prefix_lens, extend_seq_lens) + for start, extend_len in zip(starts, extend_seq_lens) ], axis=0, ) @@ -1192,13 +1270,34 @@ def compute_position_torch( return positions.to(torch.int64), extend_start_loc -def _clamp_position_native(seq_lens): - return torch.clamp((seq_lens - 1), min=0).to(torch.int64) +def _clamp_position_native(seq_lens, position_offsets: Optional[torch.Tensor] = None): + """Per-token decode position = clamp(seq_lens - 1, 0). + + Args: + seq_lens: per-request sequence length (token count). + position_offsets: optional per-request RoPE offset that shifts each + position. Used after a cache hit whose cached entry carried + non-contiguous RoPE positions: the offset is + max(cached_positions) - (prefix_token_count - 1), capturing the + gap in RoPE space caused by skipped thought tokens. When None, + behavior matches the legacy contiguous-positions path. + """ + base = torch.clamp((seq_lens - 1), min=0).to(torch.int64) + if position_offsets is not None: + base = base + position_offsets.to(torch.int64) + return base if is_cuda() or is_hip(): from sglang.jit_kernel.clamp_position import clamp_position_cuda - clamp_position = clamp_position_cuda + def clamp_position( + seq_lens: torch.Tensor, + position_offsets: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + base = clamp_position_cuda(seq_lens) + if position_offsets is not None: + base = base + position_offsets.to(torch.int64) + return base else: clamp_position = _clamp_position_native diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index a6baa4817ace..236d0990a6ba 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -695,6 +695,9 @@ def _init_pools(self: ModelRunner): enable_alt_stream=not self.server_args.enable_pdmux, enable_kv_cache_copy=( self.server_args.speculative_algorithm is not None + # --no-cache-thoughts relocates the answer KV via + # move_kv_cache when it drops a finished think span + or self.server_args.no_cache_thoughts ), ) else: @@ -714,6 +717,9 @@ def _init_pools(self: ModelRunner): enable_alt_stream=not self.server_args.enable_pdmux, enable_kv_cache_copy=( self.server_args.speculative_algorithm is not None + # --no-cache-thoughts relocates the answer KV via + # move_kv_cache when it drops a finished think span + or self.server_args.no_cache_thoughts ), ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 1bd07eec2992..2d63a93ac2d7 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -436,6 +436,7 @@ class ServerArgs: file_storage_path: str = "sglang_storage" enable_cache_report: bool = False reasoning_parser: Optional[str] = None + no_cache_thoughts: bool = False tool_call_parser: Optional[str] = None tool_server: Optional[str] = None sampling_defaults: str = "model" @@ -4508,6 +4509,17 @@ def add_cli_args(parser: argparse.ArgumentParser): default=ServerArgs.reasoning_parser, help=f"Specify the parser for reasoning models, supported parsers are: {list(ReasoningParser.DetectorMap.keys())}.", ) + parser.add_argument( + "--no-cache-thoughts", + action="store_true", + default=ServerArgs.no_cache_thoughts, + help=( + "Skip inserting reasoning (thought) tokens into the shared prefix cache. " + "Answer tokens after are inserted with their original RoPE positions " + "preserved so that thought-infused K/V representations remain reusable across " + "turns. Requires --reasoning-parser." + ), + ) tool_call_parser_choices = list(FunctionCallParser.ToolCallParserEnum.keys()) parser.add_argument( "--tool-call-parser", diff --git a/test/manual/test_bbq_smoke.py b/test/manual/test_bbq_smoke.py new file mode 100644 index 000000000000..fbd62602a5e7 --- /dev/null +++ b/test/manual/test_bbq_smoke.py @@ -0,0 +1,102 @@ +"""Smoke-load BBQ-8B-Mid1 in SGLang with --reasoning-parser k2_v3 to confirm the +model architecture and reasoning parser are wired before designing the +--no-cache-thoughts E2E test around BBQ. + +Run manually on a GPU host: + python -m sglang.launch_server ... (see below) + +Or just invoke this module: + python test/manual/test_bbq_smoke.py +""" + +import json +import os +import subprocess +import sys +import time + +import requests + +BBQ_PATH = "/mnt/weka/shrd/k2m/suqi.sun/bbq_image/bbq-8b-mid3-final" +PORT = 30000 +BASE_URL = f"http://127.0.0.1:{PORT}" + + +def main() -> int: + assert os.path.isdir(BBQ_PATH), f"missing model dir: {BBQ_PATH}" + + cmd = [ + sys.executable, + "-m", + "sglang.launch_server", + "--model-path", + BBQ_PATH, + "--reasoning-parser", + "k2_v3", + "--enable-cache-report", + "--host", + "127.0.0.1", + "--port", + str(PORT), + "--mem-fraction-static", + "0.85", + "--trust-remote-code", + ] + print("Launching:", " ".join(cmd), flush=True) + proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) + try: + # Wait for /health (up to 5 minutes). + deadline = time.time() + 300 + while time.time() < deadline: + try: + r = requests.get(f"{BASE_URL}/health", timeout=2) + if r.status_code == 200: + break + except Exception: + pass + if proc.poll() is not None: + # Server died. + out = proc.stdout.read().decode("utf-8", "replace") if proc.stdout else "" + print("SERVER EXITED EARLY:\n" + out[-4000:], flush=True) + return 1 + time.sleep(2) + else: + print("HEALTH POLL TIMED OUT", flush=True) + return 1 + print("Server is up.", flush=True) + + # Use the OpenAI-compatible chat completions endpoint so the chat template + # applies (which for BBQ-Mid3 primes the assistant turn with \n). + payload = { + "model": BBQ_PATH, + "messages": [ + { + "role": "user", + "content": "What is 12 * 7? Reason step by step.", + } + ], + "temperature": 0.0, + "max_tokens": 512, + } + r = requests.post(f"{BASE_URL}/v1/chat/completions", json=payload, timeout=120) + r.raise_for_status() + body = r.json() + choice = body["choices"][0]["message"] + content = choice.get("content", "") + reasoning = choice.get("reasoning_content", "") + print("=== reasoning_content ===\n", reasoning[:2000], flush=True) + print("=== content ===\n", content[:2000], flush=True) + print("=== usage ===\n", json.dumps(body.get("usage"), indent=2), flush=True) + # Hard check: the model emitted , so the K2V3 parser separated reasoning. + assert reasoning, "No reasoning_content — model didn't emit ..." + return 0 + finally: + proc.terminate() + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + proc.kill() + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/test/manual/test_no_cache_thoughts_e2e.py b/test/manual/test_no_cache_thoughts_e2e.py new file mode 100644 index 000000000000..d0318139f2ab --- /dev/null +++ b/test/manual/test_no_cache_thoughts_e2e.py @@ -0,0 +1,306 @@ +"""End-to-end test for --no-cache-thoughts on BBQ-8B-Mid3, TITO-style. + +What this verifies +------------------ +Two SGLang servers on the same BBQ checkpoint, one with --no-cache-thoughts +and one without. Both serve the same turn-1 reasoning request. Turn 2 builds +its input via the TITO protocol — raw token IDs, never re-tokenizing prior +content through the chat template — and excludes turn 1's thought slice +(output_token_ids[:reasoning_tokens]) from the running buffer. + +Expected behavior on turn 2: + * Baseline server (no flag): cached_tokens covers the original user prompt + only. The cached path from turn 1's finish includes the priming + + thoughts, but turn 2's buffer has the priming followed by the answer + (no thoughts), so the match dies at the first thought slot. + * --no-cache-thoughts server: the cached path was inserted with non- + contiguous positions and excludes both thoughts AND the priming tail. + Turn 2's buffer aligns: prompt-without-priming + answer-only matches + the cached entry up through the entire answer. cached_tokens covers + roughly len(prompt - priming) + len(answer). + +Hard assertion: cached_tokens(--no-cache-thoughts) > cached_tokens(baseline). + +Why TITO-style and not OpenAI chat-completions text passthrough +--------------------------------------------------------------- +Round-tripping the model's answer through text and back through the chat +template's tokenizer drifts (BPE merges differ across string-concat +boundaries), so the answer's first token ID in turn 2's input does not +equal the first token ID stored in the cached path. The radix match dies +right after the assistant header — observable as cached_tokens stuck at +the prompt-prefix length regardless of whether the flag is set. The TITO +protocol prevents this by passing input_ids directly and never letting the +chat template re-tokenize prior assistant content. See reference_tito.md. + +How to run (inside the agentic-rl image; needs 2 GPUs) +------------------------------------------------------ + srun --partition=main --time=00:30:00 -N 1 --gres=gpu:2 \\ + --container-image=/mnt/weka/shrd/k2pta/agentic_rl_images/agentic-rl-2eff86d1.sqsh \\ + --container-mounts=/mnt/weka:/mnt/weka,$PWD:/sglang \\ + bash -c "pip install --no-deps --break-system-packages -e /sglang/python && \\ + cd /sglang/test/manual && \\ + python3 -m unittest test_no_cache_thoughts_e2e -v" +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import time +import unittest + +import requests + +BBQ_PATH = "/mnt/weka/shrd/k2m/suqi.sun/bbq_image/bbq-8b-mid3-final" +CHAT_TEMPLATE = os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "chat_templates", + "bbq_upstream.jinja", +) +PORT_NO_CACHE = 30000 +PORT_BASELINE = 30001 +BASE_URL_NO_CACHE = f"http://127.0.0.1:{PORT_NO_CACHE}" +BASE_URL_BASELINE = f"http://127.0.0.1:{PORT_BASELINE}" +HEALTH_TIMEOUT = 600 # bbq-mid3 takes 1-3 min to load +REQUEST_TIMEOUT = 300 + + +def _launch(port: int, extra_args: list[str], gpu: str) -> subprocess.Popen: + env = os.environ.copy() + env["CUDA_VISIBLE_DEVICES"] = gpu + cmd = [ + sys.executable, + "-m", + "sglang.launch_server", + "--model-path", + BBQ_PATH, + "--reasoning-parser", + "k2_v3", + "--chat-template", + CHAT_TEMPLATE, + "--enable-cache-report", + "--trust-remote-code", + "--host", + "127.0.0.1", + "--port", + str(port), + "--mem-fraction-static", + "0.80", + ] + list(extra_args) + log_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".e2e_logs") + os.makedirs(log_dir, exist_ok=True) + log_path = os.path.join(log_dir, f"sglang_port{port}.log") + log_fd = open(log_path, "w") + proc = subprocess.Popen(cmd, env=env, stdout=log_fd, stderr=subprocess.STDOUT) + proc._log_path = log_path # type: ignore[attr-defined] + return proc + + +def _wait_healthy(base_url: str, proc: subprocess.Popen) -> None: + deadline = time.time() + HEALTH_TIMEOUT + while time.time() < deadline: + try: + r = requests.get(f"{base_url}/health", timeout=2) + if r.status_code == 200: + return + except Exception: + pass + if proc.poll() is not None: + log_path = getattr(proc, "_log_path", None) + tail = "" + if log_path: + try: + with open(log_path) as f: + tail = f.read()[-4000:] + except Exception: + pass + raise RuntimeError(f"server at {base_url} exited early:\n{tail}") + time.sleep(3) + raise TimeoutError(f"server at {base_url} never became healthy") + + +def _kill(proc: subprocess.Popen) -> None: + try: + proc.terminate() + try: + proc.wait(timeout=15) + except subprocess.TimeoutExpired: + proc.kill() + except Exception: + pass + + +def _chat_text(base_url: str, messages: list[dict], max_tokens: int) -> dict: + """Turn 1 — chat completions with text, asking the server to return raw token IDs.""" + payload = { + "model": BBQ_PATH, + "messages": messages, + "temperature": 0.0, + "max_tokens": max_tokens, + "return_prompt_token_ids": True, + "return_completion_token_ids": True, + "return_meta_info": True, + } + r = requests.post( + f"{base_url}/v1/chat/completions", json=payload, timeout=REQUEST_TIMEOUT + ) + r.raise_for_status() + return r.json() + + +def _chat_token_ids(base_url: str, input_ids: list[int], max_tokens: int) -> dict: + """Turn 2+ — chat completions with input_ids, bypassing the chat template.""" + payload = { + "model": BBQ_PATH, + # messages is required by the OpenAI shape but is ignored for tokenization + # when input_ids is set; SGLang still uses it to derive stop tokens. + "messages": [{"role": "user", "content": "ignored when input_ids is set"}], + "input_ids": input_ids, + "temperature": 0.0, + "max_tokens": max_tokens, + "return_prompt_token_ids": True, + "return_completion_token_ids": True, + "return_meta_info": True, + } + r = requests.post( + f"{base_url}/v1/chat/completions", json=payload, timeout=REQUEST_TIMEOUT + ) + r.raise_for_status() + return r.json() + + +def _tokenize(text: str) -> list[int]: + """Tokenize a string fragment via HF AutoTokenizer with add_special_tokens=False.""" + from transformers import AutoTokenizer + + tok = AutoTokenizer.from_pretrained(BBQ_PATH, trust_remote_code=True) + return tok.encode(text, add_special_tokens=False) + + +def _build_turn2_input_ids(turn1_resp: dict, new_user_text: str) -> list[int]: + """Construct turn 2's input_ids by: + 1. Taking turn 1's prompt_token_ids verbatim (no re-tokenization). + 2. Appending turn 1's answer slice — output_token_ids[reasoning_tokens:] — + so the thought tokens are stripped before they enter turn 2's buffer. + 3. Appending the env-delta tokens for the new user turn + assistant + generation prompt. The K2V3 boundary patch (append \\n after + <|im_end|>) is included because the model stops on <|im_end|> + without emitting the trailing \\n that the chat template would have. + """ + prompt_ids = turn1_resp["choices"][0]["prompt_token_ids"] + completion_ids = turn1_resp["choices"][0]["completion_token_ids"] + reasoning_tokens = turn1_resp["usage"]["reasoning_tokens"] + + # Strip the thought slice — keep only what follows . + answer_ids = completion_ids[reasoning_tokens:] + + # K2V3 / Qwen3 boundary patch: model stops on <|im_end|>; template emits + # <|im_end|>\n. The missing \n is part of the env-delta tokenization below + # (it's the leading \n of the env-delta string). + env_delta = ( + f"\n<|im_start|>user\n{new_user_text}<|im_end|>\n" + f"<|im_start|>assistant\n\n" + ) + env_delta_ids = _tokenize(env_delta) + + return list(prompt_ids) + list(answer_ids) + list(env_delta_ids) + + +class TestNoCacheThoughtsE2ETito(unittest.TestCase): + + @classmethod + def setUpClass(cls): + assert os.path.isdir(BBQ_PATH), f"missing model dir: {BBQ_PATH}" + cls.proc_no_cache = _launch( + PORT_NO_CACHE, ["--no-cache-thoughts"], gpu="0" + ) + cls.proc_baseline = _launch(PORT_BASELINE, [], gpu="1") + try: + _wait_healthy(BASE_URL_NO_CACHE, cls.proc_no_cache) + _wait_healthy(BASE_URL_BASELINE, cls.proc_baseline) + except Exception: + _kill(cls.proc_no_cache) + _kill(cls.proc_baseline) + raise + + @classmethod + def tearDownClass(cls): + _kill(cls.proc_no_cache) + _kill(cls.proc_baseline) + + def _dump_logs(self, label: str, proc: subprocess.Popen) -> None: + log_path = getattr(proc, "_log_path", None) + if log_path is None: + return + try: + with open(log_path) as f: + out = f.read() + print( + f"=== {label} server log (last 6KB) from {log_path} ===\n{out[-6000:]}", + flush=True, + ) + except Exception as e: + print(f"(failed to read {label} log: {e})", flush=True) + + def test_tito_cached_tokens_delta_on_turn2(self): + user_turn1 = {"role": "user", "content": "What is 12 * 7? Reason carefully."} + new_user_text = "Now multiply that result by 3." + + # Turn 1: chat completions over text, with Tito flags so we get raw IDs back. + try: + resp_nc = _chat_text(BASE_URL_NO_CACHE, [user_turn1], max_tokens=512) + resp_bl = _chat_text(BASE_URL_BASELINE, [user_turn1], max_tokens=512) + except Exception: + self._dump_logs("no_cache (turn 1)", self.proc_no_cache) + self._dump_logs("baseline (turn 1)", self.proc_baseline) + raise + + # Turn 2: build input_ids via the Tito protocol, send with input_ids in body. + input_ids_nc = _build_turn2_input_ids(resp_nc, new_user_text) + input_ids_bl = _build_turn2_input_ids(resp_bl, new_user_text) + + try: + resp_nc_t2 = _chat_token_ids( + BASE_URL_NO_CACHE, input_ids_nc, max_tokens=256 + ) + resp_bl_t2 = _chat_token_ids( + BASE_URL_BASELINE, input_ids_bl, max_tokens=256 + ) + except Exception: + self._dump_logs("no_cache (turn 2)", self.proc_no_cache) + self._dump_logs("baseline (turn 2)", self.proc_baseline) + raise + + cached_nc = resp_nc_t2["usage"]["prompt_tokens_details"]["cached_tokens"] + cached_bl = resp_bl_t2["usage"]["prompt_tokens_details"]["cached_tokens"] + prompt_nc = resp_nc_t2["usage"]["prompt_tokens"] + prompt_bl = resp_bl_t2["usage"]["prompt_tokens"] + + print(f"baseline: prompt_tokens={prompt_bl} cached_tokens={cached_bl}") + print(f"no-cache-thoughts: prompt_tokens={prompt_nc} cached_tokens={cached_nc}") + + # Both servers should see the same prompt_tokens (we constructed both inputs + # identically; the answer text was identical at temperature=0). + self.assertEqual( + prompt_nc, + prompt_bl, + "turn 2 input lengths diverged — answer differed across servers?", + ) + + # Core assertion: --no-cache-thoughts caches more of turn 2's input than + # baseline. Baseline's turn 1 cache included the priming + thoughts, + # which turn 2's Tito buffer doesn't contain at the same slot, so its + # cache match dies after the user prompt. --no-cache-thoughts's split + # path put the answer at non-contiguous positions and stripped the + # priming — turn 2 aligns through the answer. + self.assertGreater( + cached_nc, + cached_bl, + f"--no-cache-thoughts cached_tokens ({cached_nc}) should exceed " + f"baseline ({cached_bl}) by approximately len(turn-1 answer)", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/radix_cache/test_build_extend_positions.py b/test/registered/radix_cache/test_build_extend_positions.py new file mode 100644 index 000000000000..f20e23d14261 --- /dev/null +++ b/test/registered/radix_cache/test_build_extend_positions.py @@ -0,0 +1,57 @@ +"""Tests for the call-site helper that builds extend-token positions while honoring +per-request cached non-contiguous positions. + +This is the bridge between ScheduleBatch state (per-request cached_positions on +cache hits) and ForwardBatch's positions tensor. +""" + +import unittest + +import torch + +from sglang.srt.model_executor.forward_batch_info import build_extend_positions + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=5, suite="stage-b-test-1-gpu-small") + + +class TestBuildExtendPositions(unittest.TestCase): + def test_legacy_path_when_no_cached_positions(self): + """When cached_positions_per_req is None (or all None entries), positions are + contiguous starting from extend_prefix_lens[i] — i.e. unchanged from today.""" + positions, _ = build_extend_positions( + attn_backend="torch_native", # forces torch path (not triton-supported) + extend_prefix_lens=torch.tensor([2, 4], dtype=torch.int64), + extend_seq_lens=torch.tensor([3, 1], dtype=torch.int64), + extend_num_tokens=4, + extend_prefix_lens_cpu=[2, 4], + cached_positions_per_req=None, + device="cpu", + ) + # Req 0: arange(2, 5) = [2,3,4]. Req 1: arange(4, 5) = [4]. + self.assertEqual(positions.tolist(), [2, 3, 4, 4]) + + def test_cached_positions_override_extends_from_max_plus_one(self): + """When a request has cached non-contiguous positions ending at p, its extend + tokens start at p+1, not at extend_prefix_lens. Other requests fall back to + extend_prefix_lens (legacy).""" + cached_positions_per_req = [ + torch.tensor([0, 1, 6, 7, 8], dtype=torch.int64), # max 8 -> extend starts at 9 + None, # legacy fallback -> extend starts at 4 + ] + positions, _ = build_extend_positions( + attn_backend="torch_native", + extend_prefix_lens=torch.tensor([5, 4], dtype=torch.int64), + extend_seq_lens=torch.tensor([3, 1], dtype=torch.int64), + extend_num_tokens=4, + extend_prefix_lens_cpu=[5, 4], + cached_positions_per_req=cached_positions_per_req, + device="cpu", + ) + # Req 0: arange(9, 12) = [9,10,11]. Req 1: arange(4, 5) = [4]. + self.assertEqual(positions.tolist(), [9, 10, 11, 4]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/radix_cache/test_cache_finished_req_split.py b/test/registered/radix_cache/test_cache_finished_req_split.py new file mode 100644 index 000000000000..f34921eac3cd --- /dev/null +++ b/test/registered/radix_cache/test_cache_finished_req_split.py @@ -0,0 +1,79 @@ +"""Tests that RadixCache.cache_finished_req accepts a pre-computed split, bypassing +the per-req KV-pool lookup and inserting [prompt + post- answer] with original +positions preserved. + +The split is constructed by split_kv_for_no_cache_thoughts in the caller; this test +pins down the cache-side contract that consumes it. +""" + +import unittest +from unittest.mock import MagicMock + +import torch + +from sglang.srt.mem_cache.base_prefix_cache import MatchPrefixParams +from sglang.srt.mem_cache.common import split_kv_for_no_cache_thoughts +from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=5, suite="stage-b-test-1-gpu-small") + + +class TestCacheFinishedReqSplit(unittest.TestCase): + def test_split_inserts_virtual_slice_with_positions(self): + # Mock allocator records free() calls so we can assert the thought slice was freed. + mock_allocator = MagicMock() + tree = RadixCache.create_simulated(mock_allocator=mock_allocator) + + # Synthetic finished reasoning request: + # positions: 0 1 2 3 4 5 6 7 8 + # tokens: A B T1 T2 X Y Z + # prompt=[A,B] thoughts=[,T1,T2,] answer=[X,Y,Z]; answer_start_position=6 + split = split_kv_for_no_cache_thoughts( + origin_input_ids=[101, 102], + output_ids=[201, 202, 203, 204, 301, 302, 303], + req_to_token_slot=torch.tensor( + [1000, 1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008], + dtype=torch.int64, + ), + answer_start_position=6, + ) + + # Build a minimal Req-like stub: cache_finished_req(req, ..., split=split) should + # use split.virtual_* directly and ignore req.req_to_token_pool/req_pool_idx. + req_stub = MagicMock() + req_stub.extra_key = None + req_stub.priority = 0 + req_stub.pop_committed_kv_cache.return_value = 0 # bookkeeping no-op + req_stub.last_node = tree.root_node + req_stub.cache_protected_len = 0 + + tree.cache_finished_req(req_stub, is_insert=True, split=split) + + # The radix tree should now contain a path matching the virtual token ids + # (skipping the thought slice). + match = tree.match_prefix( + MatchPrefixParams( + key=RadixKey(token_ids=[101, 102, 301, 302, 303], extra_key=None) + ) + ) + self.assertEqual(match.device_indices.tolist(), [1000, 1001, 1006, 1007, 1008]) + self.assertIsNotNone(match.original_positions) + self.assertEqual(match.original_positions.tolist(), [0, 1, 6, 7, 8]) + + # The thought-slice kv_indices ([1002, 1003, 1004, 1005]) should have been freed. + freed_calls = mock_allocator.free.call_args_list + freed_indices = [] + for call in freed_calls: + arg = call.args[0] if call.args else call.kwargs.get("indices") + if isinstance(arg, torch.Tensor): + freed_indices.extend(arg.tolist()) + self.assertEqual( + sorted(set(freed_indices) & {1002, 1003, 1004, 1005}), + [1002, 1003, 1004, 1005], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/radix_cache/test_chunk_cache_split_kwarg.py b/test/registered/radix_cache/test_chunk_cache_split_kwarg.py new file mode 100644 index 000000000000..a0719c0bfbbb --- /dev/null +++ b/test/registered/radix_cache/test_chunk_cache_split_kwarg.py @@ -0,0 +1,44 @@ +"""ChunkCache.cache_finished_req must accept the split kwarg used by the +--no-cache-thoughts code path. ChunkCache doesn't do prefix caching, so the +behavior is to ignore split entirely and fall back to its default cleanup. +""" + +import unittest +from unittest.mock import MagicMock + +import torch + +from sglang.srt.mem_cache.chunk_cache import ChunkCache +from sglang.srt.mem_cache.common import NoCacheThoughtsSplit + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=5, suite="stage-b-test-1-gpu-small") + + +class TestChunkCacheSplitKwarg(unittest.TestCase): + def test_cache_finished_req_accepts_split(self): + cache = ChunkCache.__new__(ChunkCache) + cache.req_to_token_pool = MagicMock() + cache.req_to_token_pool.req_to_token = torch.tensor( + [[10, 11, 12]], dtype=torch.int64 + ) + cache.token_to_kv_pool_allocator = MagicMock() + + req = MagicMock() + req.pop_committed_kv_cache.return_value = 3 + req.req_pool_idx = 0 + + split = NoCacheThoughtsSplit( + virtual_token_ids=[1, 2], + virtual_kv_indices=torch.tensor([10, 12], dtype=torch.int64), + virtual_positions=torch.tensor([0, 5], dtype=torch.int64), + thought_kv_indices_to_free=torch.tensor([11], dtype=torch.int64), + ) + + # Must not raise TypeError on the new kwarg. + cache.cache_finished_req(req, is_insert=True, split=split) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/radix_cache/test_clamp_position_with_offsets.py b/test/registered/radix_cache/test_clamp_position_with_offsets.py new file mode 100644 index 000000000000..bc7b1feeccb8 --- /dev/null +++ b/test/registered/radix_cache/test_clamp_position_with_offsets.py @@ -0,0 +1,39 @@ +"""Tests for clamp_position honoring per-request position offsets so decode tokens +after a non-contiguous prefill cache hit get the right RoPE positions. + +At decode step N the next token's RoPE position should be: + (seq_lens[i] - 1) + position_offsets[i] +where position_offsets[i] accounts for the gap in RoPE-space caused by thought +tokens that were skipped from the cached entry. +""" + +import unittest + +import torch + +from sglang.srt.model_executor.forward_batch_info import _clamp_position_native + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=5, suite="stage-b-test-1-gpu-small") + + +class TestClampPositionWithOffsets(unittest.TestCase): + def test_offset_shifts_position(self): + # Single req, seq_len=6 (so legacy position is 5), with a 4-position gap + # in RoPE space from skipped thoughts -> next decode position should be 9. + seq_lens = torch.tensor([6], dtype=torch.int64) + offsets = torch.tensor([4], dtype=torch.int64) + + positions = _clamp_position_native(seq_lens, position_offsets=offsets) + self.assertEqual(positions.tolist(), [9]) + + def test_legacy_behavior_when_offsets_none(self): + # Without offsets, behavior must match today's clamp(seq_lens - 1, min=0). + seq_lens = torch.tensor([5, 0, 3], dtype=torch.int64) + positions = _clamp_position_native(seq_lens) + self.assertEqual(positions.tolist(), [4, 0, 2]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/radix_cache/test_collect_cached_positions.py b/test/registered/radix_cache/test_collect_cached_positions.py new file mode 100644 index 000000000000..24fb26b175b0 --- /dev/null +++ b/test/registered/radix_cache/test_collect_cached_positions.py @@ -0,0 +1,32 @@ +"""Tests for the helper that aggregates Req.cached_positions into a per-request list +suitable for ForwardBatch.init_new to consume via build_extend_positions. +""" + +import unittest +from unittest.mock import MagicMock + +import torch + +from sglang.srt.managers.schedule_batch import collect_cached_positions + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=5, suite="stage-b-test-1-gpu-small") + + +class TestCollectCachedPositions(unittest.TestCase): + def test_returns_none_when_no_req_has_cached_positions(self): + reqs = [MagicMock(cached_positions=None), MagicMock(cached_positions=None)] + self.assertIsNone(collect_cached_positions(reqs)) + + def test_returns_list_when_any_req_has_cached_positions(self): + r1 = MagicMock(cached_positions=torch.tensor([0, 1, 6, 7, 8], dtype=torch.int64)) + r2 = MagicMock(cached_positions=None) + out = collect_cached_positions([r1, r2]) + self.assertIsNotNone(out) + self.assertEqual(out[0].tolist(), [0, 1, 6, 7, 8]) + self.assertIsNone(out[1]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/radix_cache/test_compute_position_noncontig.py b/test/registered/radix_cache/test_compute_position_noncontig.py new file mode 100644 index 000000000000..790660994cf0 --- /dev/null +++ b/test/registered/radix_cache/test_compute_position_noncontig.py @@ -0,0 +1,68 @@ +"""Tests for non-contiguous extend-positions in forward_batch_info.compute_position_torch. + +When a request hits a cached entry with non-contiguous original positions (e.g. +[0, 1, 6, 7, 8] — gap where thoughts used to live), the new extend tokens must +continue from max(cached_positions) + 1, not from len(cached). This test pins +down the extended API that supports that case. +""" + +import unittest + +import torch + +from sglang.srt.model_executor.forward_batch_info import ( + compute_position, + compute_position_torch, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=5, suite="stage-b-test-1-gpu-small") + + +class TestComputePositionNonContiguous(unittest.TestCase): + def test_extend_position_start_overrides_prefix_len(self): + """When extend_position_start is provided, positions for each request's + extend tokens start at extend_position_start[i] rather than extend_prefix_lens[i].""" + # Single request: cached 5 tokens at positions [0, 1, 6, 7, 8], now extending by 3. + # Standard behavior would put extend positions at [5, 6, 7]; with the override, + # they should be at [9, 10, 11]. + extend_prefix_lens = torch.tensor([5], dtype=torch.int64) + extend_seq_lens = torch.tensor([3], dtype=torch.int64) + extend_position_start = torch.tensor([9], dtype=torch.int64) + + positions, _ = compute_position_torch( + extend_prefix_lens, extend_seq_lens, extend_position_start + ) + self.assertEqual(positions.tolist(), [9, 10, 11]) + + def test_none_override_preserves_legacy_behavior(self): + """When extend_position_start is None, positions start at extend_prefix_lens (unchanged).""" + extend_prefix_lens = torch.tensor([2, 4], dtype=torch.int64) + extend_seq_lens = torch.tensor([3, 1], dtype=torch.int64) + + positions, _ = compute_position_torch(extend_prefix_lens, extend_seq_lens) + # Request 0: starts at 2 -> [2, 3, 4]. Request 1: starts at 4 -> [4]. + self.assertEqual(positions.tolist(), [2, 3, 4, 4]) + + def test_compute_position_wrapper_forwards_override(self): + """compute_position(...) (the wrapper) must forward extend_position_start to the + underlying torch / triton implementation.""" + extend_prefix_lens = torch.tensor([5], dtype=torch.int64) + extend_seq_lens = torch.tensor([3], dtype=torch.int64) + extend_position_start = torch.tensor([9], dtype=torch.int64) + + # Use the non-triton backend name to force the torch path; the wrapper still + # routes to compute_position_triton when support_triton(attn_backend) is True + # on CUDA hosts. The torch path is the unambiguous behavioral test. + positions, _ = compute_position( + attn_backend="torch_native", # not triton-supported -> takes torch path + extend_prefix_lens=extend_prefix_lens, + extend_seq_lens=extend_seq_lens, + extend_seq_lens_sum=3, + extend_position_start=extend_position_start, + ) + self.assertEqual(positions.tolist(), [9, 10, 11]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/radix_cache/test_derive_extend_position_start.py b/test/registered/radix_cache/test_derive_extend_position_start.py new file mode 100644 index 000000000000..6286615706b9 --- /dev/null +++ b/test/registered/radix_cache/test_derive_extend_position_start.py @@ -0,0 +1,45 @@ +"""Tests for the helper that derives per-request extend_position_start from cached +positions returned by cache hits. + +This helper is what the scheduler / ForwardBatch call site uses to bridge between +the radix tree (which returns non-contiguous original_positions on cache hits) and +compute_position's extend_position_start parameter. +""" + +import unittest + +import torch + +from sglang.srt.mem_cache.common import derive_extend_position_start + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=5, suite="stage-b-test-1-gpu-small") + + +class TestDeriveExtendPositionStart(unittest.TestCase): + def test_returns_none_when_all_requests_lack_cached_positions(self): + """When no request has cached positions (e.g. flag off, or no cache hit), the + helper returns None — signaling that compute_position should use the legacy + contiguous behavior.""" + out = derive_extend_position_start( + extend_prefix_lens=[3, 5], + cached_positions_per_req=[None, None], + ) + self.assertIsNone(out) + + def test_uses_max_plus_one_for_cached_request(self): + """A request with cached non-contiguous positions returns max(positions) + 1; + a request without cached positions falls back to extend_prefix_lens (legacy).""" + out = derive_extend_position_start( + extend_prefix_lens=[5, 3], + cached_positions_per_req=[ + torch.tensor([0, 1, 6, 7, 8], dtype=torch.int64), # max 8 -> start 9 + None, # legacy fallback -> start 3 + ], + ) + self.assertEqual(out, [9, 3]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/radix_cache/test_derive_position_offsets.py b/test/registered/radix_cache/test_derive_position_offsets.py new file mode 100644 index 000000000000..1e8680911338 --- /dev/null +++ b/test/registered/radix_cache/test_derive_position_offsets.py @@ -0,0 +1,41 @@ +"""Tests for the helper that computes per-request RoPE position offsets so decode +positions continue from where the non-contiguous prefill cache hit left off. +""" + +import unittest + +import torch + +from sglang.srt.mem_cache.common import derive_position_offsets + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=5, suite="stage-b-test-1-gpu-small") + + +class TestDerivePositionOffsets(unittest.TestCase): + def test_returns_none_when_no_cached_positions(self): + out = derive_position_offsets( + extend_prefix_lens=[3, 5], + cached_positions_per_req=[None, None], + ) + self.assertIsNone(out) + + def test_offset_equals_max_minus_prefix_minus_one_plus_one(self): + """Per-req offset = max(cached_positions) - (prefix_len - 1). + + Example: prefix_len=5 (cached 5 tokens), cached_positions=[0,1,6,7,8] + -> last cached position is 8, legacy max for 5 tokens is 4, offset is 4. + """ + out = derive_position_offsets( + extend_prefix_lens=[5, 3], + cached_positions_per_req=[ + torch.tensor([0, 1, 6, 7, 8], dtype=torch.int64), # max 8 -> offset 4 + None, # no cache positions -> offset 0 + ], + ) + self.assertEqual(out, [4, 0]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/radix_cache/test_no_cache_thoughts_cli.py b/test/registered/radix_cache/test_no_cache_thoughts_cli.py new file mode 100644 index 000000000000..b20910ca1537 --- /dev/null +++ b/test/registered/radix_cache/test_no_cache_thoughts_cli.py @@ -0,0 +1,26 @@ +"""Tests for the --no-cache-thoughts CLI flag on ServerArgs.""" + +import argparse +import unittest + +from sglang.srt.server_args import ServerArgs + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=5, suite="stage-b-test-1-gpu-small") + + +class TestNoCacheThoughtsCliFlag(unittest.TestCase): + def test_default_is_false(self): + s = ServerArgs(model_path="dummy") + self.assertFalse(s.no_cache_thoughts) + + def test_argparse_sets_to_true(self): + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + ns = parser.parse_args(["--model-path", "dummy", "--no-cache-thoughts"]) + self.assertTrue(ns.no_cache_thoughts) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/radix_cache/test_no_cache_thoughts_split.py b/test/registered/radix_cache/test_no_cache_thoughts_split.py new file mode 100644 index 000000000000..d5c29d9e5a97 --- /dev/null +++ b/test/registered/radix_cache/test_no_cache_thoughts_split.py @@ -0,0 +1,132 @@ +"""Tests for the --no-cache-thoughts split-insertion helper. + +When a reasoning request finishes with --no-cache-thoughts enabled, the request's +KV must be split: the input + post- answer is inserted into the radix tree +with original positions preserved; the thought-slice KV pages are freed directly. + +This test pins down the helper function's contract: given a finished Req's metadata +and its full per-request KV slot, produce the virtual token list / kv_indices / +positions that should be inserted, plus the kv_indices that should be freed. + +Tests run without a GPU or model. +""" + +import unittest + +import torch + +from sglang.srt.mem_cache.common import split_kv_for_no_cache_thoughts +from sglang.test.test_utils import CustomTestCase + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=5, suite="stage-b-test-1-gpu-small") + + +class TestNoCacheThoughtsSplit(CustomTestCase): + """Validate the split-insertion helper for --no-cache-thoughts.""" + + def test_split_basic_case(self): + """ + Setup mirrors a typical reasoning request: + positions: 0 1 2 3 4 5 6 7 8 + tokens: A B T1 T2 X Y Z + └─prompt─┘ └────thoughts────┘ └─answer─┘ + answer_start_position = 6 (position right after ) + kv_indices in the per-request slot: [100..108], one per token. + """ + origin_input_ids = [101, 102] # A, B + output_ids = [201, 202, 203, 204, 301, 302, 303] # think, T1, T2, end, X, Y, Z + req_to_token_slot = torch.tensor( + [1000, 1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008], dtype=torch.int64 + ) + answer_start_position = 6 + + result = split_kv_for_no_cache_thoughts( + origin_input_ids=origin_input_ids, + output_ids=output_ids, + req_to_token_slot=req_to_token_slot, + answer_start_position=answer_start_position, + ) + + # Virtual token list: [A, B, X, Y, Z] + self.assertEqual(result.virtual_token_ids, [101, 102, 301, 302, 303]) + # Virtual kv_indices: pointers for slots [0, 1, 6, 7, 8] + self.assertEqual( + result.virtual_kv_indices.tolist(), [1000, 1001, 1006, 1007, 1008] + ) + # Virtual positions: prompt positions + answer original positions + self.assertEqual( + result.virtual_positions.tolist(), [0, 1, 6, 7, 8] + ) + # Thought kv_indices to free: slots [2, 3, 4, 5] + self.assertEqual( + result.thought_kv_indices_to_free.tolist(), [1002, 1003, 1004, 1005] + ) + + def test_split_no_answer_yet(self): + """If answer_start_position == total_len (i.e. emitted last, no answer + tokens generated yet), virtual sequence is just the prompt and there's no answer + slice.""" + origin_input_ids = [101, 102] + output_ids = [201, 202, 203, 204] # think, T1, T2, end_think — no answer yet + req_to_token_slot = torch.tensor( + [1000, 1001, 1002, 1003, 1004, 1005], dtype=torch.int64 + ) + # at position 5; answer starts at 6, but seq ends at 5. + answer_start_position = 6 + + result = split_kv_for_no_cache_thoughts( + origin_input_ids=origin_input_ids, + output_ids=output_ids, + req_to_token_slot=req_to_token_slot, + answer_start_position=answer_start_position, + ) + + # Virtual list contains only the input. + self.assertEqual(result.virtual_token_ids, [101, 102]) + self.assertEqual(result.virtual_kv_indices.tolist(), [1000, 1001]) + self.assertEqual(result.virtual_positions.tolist(), [0, 1]) + # All output tokens are thoughts to free. + self.assertEqual( + result.thought_kv_indices_to_free.tolist(), [1002, 1003, 1004, 1005] + ) + + def test_split_long_answer(self): + """Multi-token answer with a longer thought slice.""" + origin_input_ids = [10, 11, 12] # 3-token prompt at positions 0-2 + # 5-token thoughts at positions 3-7, then 4-token answer at positions 8-11 + output_ids = [20, 21, 22, 23, 24, 30, 31, 32, 33] + req_to_token_slot = torch.tensor( + [500, 501, 502, 503, 504, 505, 506, 507, 508, 509, 510, 511], + dtype=torch.int64, + ) + answer_start_position = 8 + + result = split_kv_for_no_cache_thoughts( + origin_input_ids=origin_input_ids, + output_ids=output_ids, + req_to_token_slot=req_to_token_slot, + answer_start_position=answer_start_position, + ) + + # Virtual list: [10, 11, 12, 30, 31, 32, 33] + self.assertEqual(result.virtual_token_ids, [10, 11, 12, 30, 31, 32, 33]) + # Virtual kv_indices: slots [0, 1, 2, 8, 9, 10, 11] + self.assertEqual( + result.virtual_kv_indices.tolist(), + [500, 501, 502, 508, 509, 510, 511], + ) + # Virtual positions: prompt [0, 1, 2] + answer [8, 9, 10, 11] + self.assertEqual( + result.virtual_positions.tolist(), [0, 1, 2, 8, 9, 10, 11] + ) + # Thoughts: slots [3, 4, 5, 6, 7] + self.assertEqual( + result.thought_kv_indices_to_free.tolist(), + [503, 504, 505, 506, 507], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/radix_cache/test_other_backends_split_kwarg.py b/test/registered/radix_cache/test_other_backends_split_kwarg.py new file mode 100644 index 000000000000..4546b729ffab --- /dev/null +++ b/test/registered/radix_cache/test_other_backends_split_kwarg.py @@ -0,0 +1,65 @@ +"""Signature-level test: every non-RadixCache prefix-cache backend's cache_finished_req +must accept the split kwarg (either explicitly or via **kwargs) so the --no-cache-thoughts +routing in release_kv_cache doesn't raise TypeError on these backends. +""" + +import inspect +import unittest + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=5, suite="stage-b-test-1-gpu-small") + + +def _accepts_split(cls) -> bool: + sig = inspect.signature(cls.cache_finished_req) + has_split = "split" in sig.parameters + has_kwargs = any( + p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values() + ) + return has_split or has_kwargs + + +class TestOtherBackendsAcceptSplit(unittest.TestCase): + def test_swa_radix_cache(self): + from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache + + self.assertTrue(_accepts_split(SWARadixCache), "SWARadixCache rejects split kwarg") + + def test_mamba_radix_cache(self): + from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache + + self.assertTrue( + _accepts_split(MambaRadixCache), "MambaRadixCache rejects split kwarg" + ) + + def test_session_aware_cache(self): + from sglang.srt.mem_cache.session_aware_cache import SessionAwareCache + + self.assertTrue( + _accepts_split(SessionAwareCache), "SessionAwareCache rejects split kwarg" + ) + + def test_radix_cache_cpp(self): + try: + from sglang.srt.mem_cache.radix_cache_cpp import RadixCacheCpp + except Exception as e: + self.skipTest(f"RadixCacheCpp not importable in this env: {e}") + self.assertTrue( + _accepts_split(RadixCacheCpp), "RadixCacheCpp rejects split kwarg" + ) + + def test_lmc_radix_cache(self): + try: + from sglang.srt.mem_cache.storage.lmcache.lmc_radix_cache import ( + LMCRadixCache, + ) + except Exception as e: + self.skipTest(f"LMCRadixCache not importable in this env: {e}") + self.assertTrue( + _accepts_split(LMCRadixCache), "LMCRadixCache rejects split kwarg" + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/radix_cache/test_radix_position_preservation.py b/test/registered/radix_cache/test_radix_position_preservation.py new file mode 100644 index 000000000000..a318f1051ff0 --- /dev/null +++ b/test/registered/radix_cache/test_radix_position_preservation.py @@ -0,0 +1,137 @@ +"""Tests for radix-tree round-tripping of per-token original RoPE positions. + +The radix tree must accept original_positions on insert and return them on +match_prefix, so callers can preserve non-contiguous positions (e.g. when a +generated thought slice was excluded from the cached entry) across cache hits. + +These tests run without a GPU or model — they use RadixCache.create_simulated(). +""" + +import unittest + +import torch + +from sglang.srt.mem_cache.base_prefix_cache import InsertParams, MatchPrefixParams +from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=5, suite="stage-b-test-1-gpu-small") + + +class TestRadixPositionPreservation(CustomTestCase): + """Radix tree must carry per-token original_positions through insert and match.""" + + def setUp(self): + self.tree = RadixCache.create_simulated() + + def test_insert_accepts_original_positions(self): + """InsertParams must accept an original_positions tensor matching key length.""" + token_ids = [10, 11, 12, 13, 14] + positions = torch.tensor([0, 1, 6, 7, 8], dtype=torch.int64) + kv_indices = torch.tensor([100, 101, 102, 103, 104], dtype=torch.int64) + + result = self.tree.insert( + InsertParams( + key=RadixKey(token_ids=token_ids, extra_key=None), + value=kv_indices, + original_positions=positions, + ) + ) + # Insert is expected to succeed; prefix_len reflects pre-existing tree overlap (here, 0). + self.assertEqual(result.prefix_len, 0) + + def test_match_returns_non_contiguous_positions(self): + """After inserting with non-contiguous positions, match_prefix must return them.""" + token_ids = [10, 11, 12, 13, 14] + positions = torch.tensor([0, 1, 6, 7, 8], dtype=torch.int64) + kv_indices = torch.tensor([100, 101, 102, 103, 104], dtype=torch.int64) + + self.tree.insert( + InsertParams( + key=RadixKey(token_ids=token_ids, extra_key=None), + value=kv_indices, + original_positions=positions, + ) + ) + + match = self.tree.match_prefix( + MatchPrefixParams(key=RadixKey(token_ids=token_ids, extra_key=None)) + ) + + self.assertIsNotNone(match.original_positions) + self.assertEqual(match.original_positions.tolist(), [0, 1, 6, 7, 8]) + # device_indices must still match the inserted kv_indices. + self.assertEqual(match.device_indices.tolist(), [100, 101, 102, 103, 104]) + + def test_match_returns_none_positions_for_legacy_insert(self): + """Backwards-compat: insert without original_positions returns None on match.""" + token_ids = [20, 21, 22] + kv_indices = torch.tensor([200, 201, 202], dtype=torch.int64) + + self.tree.insert( + InsertParams( + key=RadixKey(token_ids=token_ids, extra_key=None), + value=kv_indices, + ) + ) + + match = self.tree.match_prefix( + MatchPrefixParams(key=RadixKey(token_ids=token_ids, extra_key=None)) + ) + + self.assertIsNone(match.original_positions) + self.assertEqual(match.device_indices.tolist(), [200, 201, 202]) + + def test_partial_match_returns_position_prefix(self): + """If only a prefix of the cached entry matches, returned positions cover that prefix.""" + token_ids = [10, 11, 12, 13, 14] + positions = torch.tensor([0, 1, 6, 7, 8], dtype=torch.int64) + kv_indices = torch.tensor([100, 101, 102, 103, 104], dtype=torch.int64) + + self.tree.insert( + InsertParams( + key=RadixKey(token_ids=token_ids, extra_key=None), + value=kv_indices, + original_positions=positions, + ) + ) + + # Query with only the first 3 tokens; expect positions [0, 1, 6] + match = self.tree.match_prefix( + MatchPrefixParams(key=RadixKey(token_ids=[10, 11, 12], extra_key=None)) + ) + + self.assertIsNotNone(match.original_positions) + self.assertEqual(match.original_positions.tolist(), [0, 1, 6]) + self.assertEqual(match.device_indices.tolist(), [100, 101, 102]) + + def test_extend_existing_path_with_positions(self): + """Inserting a longer sequence with positions extends an existing prefix path.""" + # First, insert the contiguous prompt [A, B] at positions [0, 1]. + self.tree.insert( + InsertParams( + key=RadixKey(token_ids=[1, 2], extra_key=None), + value=torch.tensor([100, 101], dtype=torch.int64), + original_positions=torch.tensor([0, 1], dtype=torch.int64), + ) + ) + + # Then, insert [A, B, X, Y, Z] at positions [0, 1, 6, 7, 8] (gap where thoughts were). + self.tree.insert( + InsertParams( + key=RadixKey(token_ids=[1, 2, 3, 4, 5], extra_key=None), + value=torch.tensor([100, 101, 102, 103, 104], dtype=torch.int64), + original_positions=torch.tensor([0, 1, 6, 7, 8], dtype=torch.int64), + ) + ) + + match = self.tree.match_prefix( + MatchPrefixParams(key=RadixKey(token_ids=[1, 2, 3, 4, 5], extra_key=None)) + ) + self.assertIsNotNone(match.original_positions) + self.assertEqual(match.original_positions.tolist(), [0, 1, 6, 7, 8]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/radix_cache/test_release_kv_cache_routing.py b/test/registered/radix_cache/test_release_kv_cache_routing.py new file mode 100644 index 000000000000..01949bcc413a --- /dev/null +++ b/test/registered/radix_cache/test_release_kv_cache_routing.py @@ -0,0 +1,95 @@ +"""Tests that release_kv_cache routes through the split helper when --no-cache-thoughts +is enabled and the request has a recorded answer_start_position. + +The test mocks the tree_cache and req objects narrowly enough to observe the routing +decision without needing a real KV pool. The next-cycle test will exercise the +non-flag path to ensure no regression. +""" + +import unittest +from unittest.mock import MagicMock, patch + +import torch + +from sglang.srt.mem_cache.common import ( + NoCacheThoughtsSplit, + release_kv_cache, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=5, suite="stage-b-test-1-gpu-small") + + +class TestReleaseKvCacheRouting(unittest.TestCase): + def _make_tree_cache_mock(self): + tree = MagicMock() + tree.supports_mamba.return_value = False + # Per-req KV slot used by the split helper. + tree.req_to_token_pool.req_to_token = torch.tensor( + [[1000, 1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008]], + dtype=torch.int64, + ) + # The split helper builds its slot indexing from this tensor. + return tree + + def _make_req_mock(self): + req = MagicMock() + req.req_pool_idx = 0 + req.require_reasoning = True + req.answer_start_position = 6 + req.origin_input_ids = [101, 102] + req.output_ids = [201, 202, 203, 204, 301, 302, 303] + req.pop_overallocated_kv_cache.return_value = (0, 0) + req.mamba_pool_idx = None + return req + + def test_routes_through_split_when_flag_on(self): + tree = self._make_tree_cache_mock() + req = self._make_req_mock() + + fake_server_args = MagicMock() + fake_server_args.no_cache_thoughts = True + fake_server_args.page_size = 1 + fake_server_args.speculative_algorithm = None + + with patch( + "sglang.srt.mem_cache.common.get_global_server_args", + return_value=fake_server_args, + ): + release_kv_cache(req, tree, is_insert=True) + + # cache_finished_req must have been called with a split kwarg. + call = tree.cache_finished_req.call_args + self.assertIsNotNone(call, "cache_finished_req was not called") + split = call.kwargs.get("split") + self.assertIsNotNone(split, "split kwarg was not passed") + self.assertIsInstance(split, NoCacheThoughtsSplit) + # Sanity-check the split contents match the synthetic Req. + self.assertEqual(split.virtual_token_ids, [101, 102, 301, 302, 303]) + self.assertEqual(split.virtual_positions.tolist(), [0, 1, 6, 7, 8]) + + def test_no_split_when_flag_off(self): + tree = self._make_tree_cache_mock() + req = self._make_req_mock() # has require_reasoning + answer_start_position set + + fake_server_args = MagicMock() + fake_server_args.no_cache_thoughts = False # flag off + fake_server_args.page_size = 1 + fake_server_args.speculative_algorithm = None + + with patch( + "sglang.srt.mem_cache.common.get_global_server_args", + return_value=fake_server_args, + ): + release_kv_cache(req, tree, is_insert=True) + + call = tree.cache_finished_req.call_args + self.assertIsNotNone(call) + self.assertIsNone( + call.kwargs.get("split"), + "split kwarg should not be passed when --no-cache-thoughts is off", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/radix_cache/test_req_answer_start_position.py b/test/registered/radix_cache/test_req_answer_start_position.py new file mode 100644 index 000000000000..6d82be3410bb --- /dev/null +++ b/test/registered/radix_cache/test_req_answer_start_position.py @@ -0,0 +1,35 @@ +"""Tests for Req.answer_start_position tracking via update_reasoning_tokens.""" + +import unittest + +from sglang.srt.managers.schedule_batch import Req +from sglang.srt.sampling.sampling_params import SamplingParams + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=5, suite="stage-b-test-1-gpu-small") + + +class TestReqAnswerStartPosition(unittest.TestCase): + def test_set_when_think_end_detected(self): + """When update_reasoning_tokens sees the token, answer_start_position + is set to len(input) + reasoning_tokens, i.e. the position right after .""" + # Prompt is 2 tokens [10, 11] at positions 0, 1. + # Thoughts are 4 tokens [20, 21, 22, 99] at positions 2, 3, 4, 5 — 99 is . + # Answer should start at position 6. + req = Req(rid="r1", origin_input_text="hi", origin_input_ids=[10, 11], + sampling_params=SamplingParams(), require_reasoning=True) + think_end_id = 99 + # Feed thought tokens one at a time; not yet the id. + req.update_reasoning_tokens(20, think_end_id) + req.update_reasoning_tokens(21, think_end_id) + req.update_reasoning_tokens(22, think_end_id) + self.assertIsNone(req.answer_start_position) + # Feed the token. + req.update_reasoning_tokens(99, think_end_id) + self.assertTrue(req._is_reasoning_over) + self.assertEqual(req.answer_start_position, 6) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/radix_cache/test_req_cached_positions.py b/test/registered/radix_cache/test_req_cached_positions.py new file mode 100644 index 000000000000..51c54056c864 --- /dev/null +++ b/test/registered/radix_cache/test_req_cached_positions.py @@ -0,0 +1,52 @@ +"""Tests that Req captures the cached non-contiguous positions returned by match_prefix. + +When the prefix cache has an entry with original_positions set (e.g. because a prior +turn was inserted via the split path), a future request that hits that entry must +record those positions on the Req so the scheduler can build the right +extend_position_start for compute_position. +""" + +import unittest + +import torch + +from sglang.srt.managers.schedule_batch import Req +from sglang.srt.mem_cache.base_prefix_cache import InsertParams +from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey +from sglang.srt.sampling.sampling_params import SamplingParams + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=5, suite="stage-b-test-1-gpu-small") + + +class TestReqCachedPositions(unittest.TestCase): + def test_match_with_non_contiguous_positions_stored_on_req(self): + # Seed the radix tree with [A, B, X, Y, Z] at positions [0, 1, 6, 7, 8] — + # simulating a prior turn that was inserted via the split path. + tree = RadixCache.create_simulated() + tree.insert( + InsertParams( + key=RadixKey(token_ids=[10, 11, 30, 31, 32], extra_key=None), + value=torch.tensor([100, 101, 102, 103, 104], dtype=torch.int64), + original_positions=torch.tensor([0, 1, 6, 7, 8], dtype=torch.int64), + ) + ) + + # New request whose input matches the cached entry. + req = Req( + rid="r1", + origin_input_text="...", + origin_input_ids=[10, 11, 30, 31, 32, 99], # +1 trailing token so we don't truncate + sampling_params=SamplingParams(), + ) + req.init_next_round_input(tree_cache=tree) + + self.assertIsNotNone(req.cached_positions) + # The first 5 tokens should hit the cached prefix; positions reflect the original + # non-contiguous layout. (Match may stop one token short to enable logprob compute.) + self.assertEqual(req.cached_positions.tolist()[:5], [0, 1, 6, 7, 8]) + + +if __name__ == "__main__": + unittest.main()