feat: --no-cache-thoughts for reasoning-model multi-turn cache alignment

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 </think>. Set by update_reasoning_tokens when the </think>
+        # 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 </think>. 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

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 (

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):

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[

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 <think>...</think> 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-</think> 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 `</think>`, the tokens
+    between input_end and the `</think>` boundary are thoughts that must NOT be
+    registered in the cross-request radix cache; the post-`</think>` 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 `</think>`.
+        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 <think>\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-</think> 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