Skip to content

Default fully-async recipes to retract weight-update pauses - #2782

Open
fzyzcjy wants to merge 1 commit into
mainfrom
tom/fix-inplace-stale-radix-cache
Open

fzyzcjy wants to merge 1 commit into
mainfrom
tom/fix-inplace-stale-radix-cache

Conversation

@fzyzcjy

@fzyzcjy fzyzcjy commented Aug 27, 2026

Copy link
Copy Markdown
Collaborator

TODO: show cache weight version, fix the unbounded issue

Summary

--pause-generation-mode in_place leaves the SGLang radix (prefix) cache poisoned with pre-update KV after every weight update, so later rollouts silently generate on a mixture of new weights and old-policy KV. Every fully-async recipe on main currently defaults to or hardcodes in_place. This PR flips them to retract, which recomputes retracted requests' KV under the new weights and flushes the prefix cache on every update.

Mechanism

  • pause_generation(mode="in_place") freezes in-flight requests so their KV survives the update — that is the mode's purpose. Because those frozen requests hold KV, the engine-side cache flush cannot run: scheduler.flush_cache refuses while requests hold pages, and the engine's update paths assert on flush success. sglang's own io_struct.py documents it: "In 'inplace' mode, flush_cache will fail if there are any requests".
  • The miles update paths therefore skip the flush entirely under in_place (update_weight_from_distributed/mixin.py:315-316, same guard in delta.py and update_weight_from_tensor.py).
  • Skipping the flush is forced, but its fallout is unhandled: nothing ever invalidates the radix tree. Prefixes cached under pre-update weights keep serving every later request that shares them, across arbitrarily many updates. A radix hit reuses the cached KV without recompute, and LRU keeps hot prefixes (shared system prompt, the per-prompt prefix under --n-samples-per-prompt, earlier turns of a multi-turn conversation) alive indefinitely — so the poisoning is unbounded in time and covers all subsequent traffic, not just the requests frozen across one update.

Symptom

  • Rollouts become a hybrid behavior policy: fresh weights on the un-cached suffix, stale KV on the cached prefix. The recorded rollout_log_probs describe that hybrid, while the trainer recomputes log probs under the current weights, so train/train_rollout_logprob_abs_diff grows with every update instead of sitting at cross-implementation numeric noise, and the effective behavior policy lags the trainer roughly in proportion to the cached-prefix token share.
  • The existing rollout/weight_version/mixed_version_ratio metric stays near zero throughout: per-request weight-version events only track requests alive across an update, and radix reuse by new requests is invisible to them. The corruption is therefore silent on the dashboard.
  • Severity scales with radix hit rate × per-update policy movement — worst for exactly the standard RL shape: shared system prompt, many samples per prompt, --update-weights-interval 1.

Why retract is the right default now

Changes

file change
examples/infra_features/fully_async/run_qwen3_30b_a3b_fully_async.py default in_placeretract
examples/infra_features/fully_async/run_qwen3_5_4b_fully_async_eval.py hardcoded → retract
examples/infra_features/random_async/run_random_async_3node.py default → retract
examples/swe-agent-harbor-docker/run-glm47-flash-agentic-async.py default → retract
examples/experimental/openenv/glm52_tbench2/run_glm5_2_744b_a40b_daytona.py hardcoded → retract
scripts/run_inkling.py hardcoded → retract
docs/user-guide/fully-async.md, docs/examples/infra-features/random-async.md, examples/infra_features/random_async/README.md document the default and the reason

Deliberate exception: examples/multi_lora/run_multi_lora.py keeps in_place — its adapter upsert relies on it (an unload would deadlock behind paused in-flight requests), and its --use-tis corrects the training-side bias since the recorded log probs faithfully describe the (lagged) behavior policy.

Follow-up

  • Engine-side long-term fix, tracked separately: evict the evictable radix nodes when the weight version bumps. Frozen requests' nodes are lock-protected, so that would make in_place safe again (including for multi-turn workloads, which lose cross-turn prefix reuse under any flush-based scheme).

Tests

  • pre-commit run --all-files clean. The change is recipe defaults plus docs; behavior is covered by the existing fully-async paths that already run retract.

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

Tip: disable this comment in your organization's Code Review settings.

pause-generation-mode=in_place freezes in-flight requests across a weight
update so their KV survives, which also means the engine-side prefix-cache
flush cannot run (flush_cache refuses while requests hold KV, and the
engine asserts on it). The update path therefore skips the flush in this
mode -- but nothing ever invalidates the radix tree, so prefixes cached
under pre-update weights keep serving every later request that shares
them. Rollouts silently mix old-policy KV into new-policy generations,
the recorded rollout log probs describe that hybrid, and the
train-vs-rollout log-prob gap grows with every update. The per-request
weight_version bookkeeping (mixed_version_ratio) only tracks requests
alive across an update, so it stays near zero while the prefix reuse
poisons all subsequent traffic.

retract does not have this problem: every retracted request recomputes
its KV under the new weights and the update path flushes the prefix
cache, which actually succeeds under load since #1750 and
sgl-project/sglang#31962. Flip every recipe that defaulted or hardcoded
in_place over to retract, and document why in the fully-async guide.
examples/multi_lora keeps in_place deliberately: adapter upsert relies on
it (an unload would deadlock behind paused requests) and its --use-tis
corrects the training-side bias.
@fzyzcjy
fzyzcjy force-pushed the tom/fix-inplace-stale-radix-cache branch from bef6eef to 80f79db Compare August 27, 2026 07:58
@fzyzcjy fzyzcjy changed the title Disable the radix cache when in-place weight updates cannot flush it Default fully-async recipes to retract weight-update pauses Aug 27, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant