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

[DavidBellamy]

Summary

Adds a --no-cache-thoughts server flag that, for reasoning requests, inserts the answer slice into the prefix cache at the original RoPE positions (skipping the thought slice) instead of caching the entire generated trajectory at contiguous positions. Lets a TITO multi-turn rollout reuse turn N's answer KV as a prefix match on turn N+1, without polluting the radix tree with request-specific reasoning content.

End-to-end validated on bbq-8b-mid3 in the agentic-rl image:

• Baseline server (no flag) — turn 2 cached_tokens = 22 (prompt prefix only)
• --no-cache-thoughts — turn 2 cached_tokens = 191 (prompt + answer, delta = answer length)

How it works (server side)

Stage	Where	What
Detect </think>	Req.update_reasoning_tokens (schedule_batch.py)	Sets req.answer_start_position when the reasoning parser emits its end token
Build the split	split_kv_for_no_cache_thoughts (common.py)	Returns (virtual_token_ids, virtual_kv_indices, virtual_positions, thought_kv_indices_to_free) — input prompt + post-</think> answer, with answer carrying its original (non-contiguous) positions; thought slots earmarked for free
Route at finish	release_kv_cache (common.py)	When the flag is on and the request has reasoning, builds the split and calls cache_finished_req(split=…)
Insert + free	RadixCache.cache_finished_req (radix_cache.py)	Inserts virtual token ids into the radix with original_positions attached, frees the thought slice directly
Round-trip positions	RadixCache.{insert,_split_node,_match_prefix_helper,match_prefix}, TreeNode.positions, InsertParams.original_positions, MatchResult.original_positions	Per-token position metadata travels through the tree; match_prefix returns it on cache hits
Plumb to Req	Req.init_next_round_input (schedule_batch.py)	Stores match_result.original_positions on req.cached_positions
Aggregate per batch	collect_cached_positions, prepare_for_extend	ScheduleBatch.cached_positions_per_req + ScheduleBatch.position_offsets populated
Forward batch positions	build_extend_positions, clamp_position (forward_batch_info.py)	Prefill and decode positions both honor non-contiguous cached positions; new tokens continue from max(cached) + 1
Backend stubs	ChunkCache, SWARadixCache, MambaRadixCache, RadixCacheCpp, LMCRadixCache, SessionAwareCache	Accept the new split= kwarg; non-RadixCache backends log a one-time warning and fall back to standard cache_finished_req (graceful no-op)

Integration spec for callers (TITO multi-turn rollouts)

The chat-completions text path cannot deliver multi-turn answer reuse even with this flag, because the chat template re-tokenizes prior assistant content via BPE and the first answer token ID after the assistant header doesn't match the model's emitted ID (leading-whitespace BPE merges differ). Callers must use the TITO protocol:

1. Turn 1: send via /v1/chat/completions with return_prompt_token_ids=True, return_completion_token_ids=True, return_meta_info=True. Capture from the response:

   • choices[0].prompt_token_ids — turn 1's input IDs
   • choices[0].completion_token_ids — turn 1's output IDs (full, including thoughts)
   • usage.reasoning_tokens — count of thought tokens (including </think>)

2. Strip the thoughts before appending to the running buffer:
   ```python
   answer_ids = completion_token_ids[reasoning_tokens:]
   running_ids = prompt_token_ids + answer_ids
   ```
   The buffer must NOT contain the thought slice; the cached entry doesn't.

3. Turn N+1: construct input_ids = running_ids + tokenize(env_delta) where env_delta is the new user message + new assistant priming, including the K2V3 boundary \n (model stops on <|im_end|> but the chat template emits <|im_end|>\n, so the delta starts with \n). Send via /v1/chat/completions with input_ids set in the request body (chat template is bypassed when input_ids is provided).

4. Verify cache hits: turn N+1's usage.prompt_tokens_details.cached_tokens should grow turn-over-turn by the prior answer's length.

For RL360: this is the same tito-patched flow already used by miles/rollout/generate_hub/agentic_tool_call.py; no client changes needed beyond ensuring the thought-stripping step runs for --no-cache-thoughts-enabled servers.

Tests

• 34 unit tests under test/registered/radix_cache/ (1 env-skipped for C++ extension). Covers: CLI flag, Req.answer_start_position, split helper, radix position round-trip, release_kv_cache routing, compute_position/clamp_position non-contiguous support, derive helpers, Req.cached_positions, ScheduleBatch.cached_positions_per_req, and the 6 non-RadixCache backend signature acceptance.
• TITO E2E at test/manual/test_no_cache_thoughts_e2e.py. Launches two bbq-8b-mid3 servers (with/without flag), drives a 2-turn protocol via raw input_ids, asserts the cached_tokens delta on turn 2. Passes on agentic-rl image.
• The pre-existing test/registered/distributed/test_parallelism_context_integration.py lacks register_cuda_ci(...) and breaks test/run_suite.py collection — separate concern, not introduced by this PR.

Known limitations / follow-up

• Chat-completions text path doesn't benefit (BPE retokenization breaks alignment regardless of the flag). The flag is a TITO-protocol feature; PR description above documents the required client behavior.
• 6 non-RadixCache backends only stub the split= kwarg (warn + fall back). Full per-backend implementation deferred until empirical results justify the engineering.
• The mid3-bundled chat_template.jinja is stale relative to LLM360/bbq-chat-template:main (the bundled copy has a 'think' default for think_tag that the upstream removed). E2E uses an explicit --chat-template override pointing at a vendored upstream copy under test/manual/chat_templates/bbq_upstream.jinja. If/when Mid3 ships with the upstream template, the override can be dropped.

[DavidBellamy]

Mapping this to the trainer side: agentic inference looks like: prefill phase -> decode phase -> environment -> prefill phase -> decode phase -> environment -> ...

For turn i in [0, 1, 2, ...] let P_i := prefill phase, D_i := decode phase, E_i := environment phase for turn i.

Example:
P_0 = [sys_msg, usr_prompt], more broadly P_i is the full trajectory history up to but not including D_i
D_0 = [think, ans] conditioned on P_0. Note that tool calls are parsed from ans.
E_0 = [tool_output] added to [P_0, D_0]

In inference engines like sglang and vllm, when the prefill phase is entered, the global KV pool (shared across all concurrent inference requests) is checked to see if those tokens' KV vectors are cached. The engine matches the largest prefix from P_i that it can find in the KV pool and returns that, then prefills the remaining suffix of P_i.

D_i then begins, conditioned on P_i. During decode, a local KV cache (scoped per inference request) is constructed for the growing decoded sequence. The global KV pool is not added to during decode. Under standard protocol - once decode finishes, the full decoded sequence is inserted into the global KV pool.

#29 alters only that last step. Once decode finishes, only the ans (i.e. non-think) tokens are added to the global KV pool. So the global KV pool for each request is updated to contain P_i + ans after one request's decode phase completes.

Note that the positional encodings (e.g. RoPE) for the ans token activations are not shifted in my current implementation. So if there are N think tokens in the first decode step, the first ans token's activation will have its original position of '6' embedded into it.

Also note that the decode phase within each turn is conditioning on the think tokens from that same turn during generation, as in standard inference.

So on the trainer side, let's say we have an agentic trajectory and want to attention mask it so the model learns how to handle this inference strategy. For all tokens in turn N, all think tokens from turns N-1, N-2, ..., 0 should be masked. That should do the trick!