perf(training): share prefix caches wherever they are cleared at publication

[BobbyZhouZijian]

Closes #207.

Motivation

Reef disabled SGLang's radix cache for every deployment. Agent requests repeatedly send system prompts, tool schemas, and conversation history, so this prevented useful prefix reuse.

A cached KV entry must not survive a weight publication: an entry carries no runtime-load-ID identity, so a prefix encoded under old weights would be matched by a request generating under the new ones. This change couples prefix sharing to publications that retract in-flight requests and clear the cache before the weights change.

Changes

After #432 moved inference ownership into reef.inference.sglang and publication into reef.runtime, this branch was re-ported onto those owners (the earlier rollout manager, preflight and bridge edits no longer exist):

• reef/inference/sglang/config.py: SGLangConfig owns the sharing decision. When disable_radix_cache is unset, sharing is enabled only if pause_mode == "retract" and, under LoRA, SGLang keys radix entries by adapter. An explicit disable_radix_cache: false (global or per group) is rejected otherwise. A retract pause mode rejects prefill/decode groups, because the pinned scheduler cannot retract disaggregated requests.
• reef/inference/sglang/lora_schema.py: adapter_scoped_prefix_cache_supported() exercises SGLang's request and radix-key APIs: identical tokens must share a key under one adapter and differ under another. No signature inspection is used. Any failure answers "no".
• reef/train/slime_backend/inference.py + slime_arguments.py: colocated engines, or regular disjoint engines with --disjoint-prefix-sharing, select the retract pause mode. The translation no longer forces the radix flag. It keeps an explicit opt-out (inference.options.disable-radix-cache: true) and otherwise leaves the default to SGLangConfig.
• reef/inference/sglang/worker.py: a retract pause flushes each engine's cache after pausing. A failed flush raises, and the publication then retires the engines.
• reef/inference/sglang/plugin.py: while paused in retract mode, the destructive-memory idle check ignores the grammar queue as well as waiting requests. Grammar polling stops during a pause, so requests still compiling would otherwise block every flush. Queued requests and their compilation futures are preserved, and all GPU idle checks remain. The inspect-based signature probing is replaced by the pinned ignore_waiting API.
• reef/runtime/publication.py: a failed pause barrier aborts, which retires the engines. TrainingPublication.publish and prepare_recovery now move a CHECKPOINT marker to UPDATING_WEIGHTS in that case. The retry therefore rebuilds the engines and forces a full publication instead of replaying CHECKPOINT against retired handles. Previously the marker stayed CHECKPOINT, and the health check reported the failure as not recoverable.

Compatibility and operational impact

• Default change: colocated engines, including colocated LoRA on the pinned SGLang, now share prefixes. The trade-off is cache occupancy plus a flush at each publication. Colocated engines already dropped their KV at publication.
• New option: training.options.disjoint-prefix-sharing (off by default). Regular disjoint engines keep in-place pauses and no sharing unless they opt in. With the option, each publication retracts in-flight requests and re-prefills them under the new weights.
• Configuration validation: SGLangConfig now accepts disable_radix_cache: false only where it is safe; it previously rejected it always. It also rejects unknown pause_mode values, and PD groups under retract.
• External engines: these are still verified to match the configured disable_radix_cache.
• Marker semantics: a pause-barrier failure now leaves UPDATING_WEIGHTS, which is replayable in place, instead of CHECKPOINT.
• Wire and persisted formats: no wire or persisted-format change.
• Dependencies: no new dependencies.

Verification

CPU, at head 873c49c0 (Python 3.12, pinned Slime and SGLang installed):

pytest tests/slime_backend/test_sglang_engine.py tests/slime_backend/test_rollout_recovery.py \
       tests/slime_backend/test_driver_runtime_env.py tests/slime_backend/test_independent_sglang.py \
       tests/reef_service/test_slime_bridge.py tests/reef_service/test_training_publication.py \
       tests/reef_service/test_training_restart.py tests/reef_service/test_inference_recovery.py tests/inference/
322 passed, 4 skipped
pre-commit run --files <changed files>   # all hooks passed
python -m mypy                           # no issues in 318 source files
pytest tests/reef_service tests/slime_backend tests/inference tests/plugin_contracts
4286 passed, 131 skipped, 19 failed

The same 19 tests also fail on unmodified origin/main in this environment. The main cause is that git-lfs is not installed, plus several CLI and connector startup tests. None of them touch this change.

Regression coverage includes:

• the default and opt-in configuration matrix;
• YAML overrides for disjoint, colocated and LoRA engines;
• CLI and YAML rejection of PD topologies;
• behavior of the adapter-scoped check with broken, missing or unexpected SGLang APIs;
• the retract pause flushing, and the in-place pause not flushing;
• preservation of pending grammar requests together with the GPU idle guards;
• the marker after a publication or startup barrier failure, and the retry that rebuilds the engines.

GPU: 1× NVIDIA RTX PRO 6000 Blackwell per scenario, driver 580.173.02, torch 2.11.0+cu130, Ray 2.58.0, SGLang f7ea06e1e519c569c1d6cb7e1c7b1128cbdbc39f, Qwen2.5-0.5B-Instruct.

Each scenario runs Reef's ReefSGLangEngine as a Ray actor with the Reef SGLang plugin enabled, builds its configuration through the Slime translation and SGLangConfig, and publishes in Reef's order: _SGLangInferenceEngines.pause() (pause plus flush), optional memory release and resume, update_weights_from_disk with a new runtime load ID, then continue_generation. The prompt shares a prefix of about 2.2k tokens.

Scenario	Result
Disjoint default	in_place, radix disabled, 0 cached tokens on a repeated prefix; 4 in-flight 3000-token requests continue across an in-place publication ✅
Disjoint + --disjoint-prefix-sharing	retract, radix enabled; second request reuses 2187/2189 prompt tokens; 0 cached tokens right after a publication, 2187 again afterwards ✅
↳ in-flight publication	8 × 3000-token requests plus 16 JSON-schema requests in flight; the pause-and-flush barrier succeeds; every request completes (3000/3000 tokens, 16/16 valid JSON); per-token runtime load IDs are engine:2 for the 872 tokens before the barrier and only engine:3 for the 2128 after ✅
Colocated	retract, radix and memory saver enabled; pause, flush, full memory release and resume, then weights; 6 × 3000-token requests plus 8 JSON-schema requests complete; tokens split 955 old / 2045 new; cache is empty after a publication (2187 → 0) ✅
Colocated LoRA (two random rank-8 adapters)	the adapter-scoped check passes on the pinned SGLang, so radix is enabled; the same adapter reuses 2187 tokens; a different adapter on the same prompt gets 0 cached tokens; adapter entries are cleared by a publication ✅
Grammar still compiling at pause (about 800-field schemas, about 7 s xgrammar compile)	Reef engine: a paused flush returns HTTP 200 and the 4 queued requests complete after resume. Negative control, plain SGLang without the Reef plugin: the same paused flush returns HTTP 400 ✅

The pinned SGLang's paused flush_cache already ignores the waiting queue (is_fully_idle(ignore_waiting=self._engine_paused)). The plugin's contribution is the grammar queue, which the negative control confirms. The GPU host needed its GCE gIB NCCL net shim removed from LD_LIBRARY_PATH and its NCCL_* overrides unset; this is a single-node host setting unrelated to Reef.

Not covered on GPU:

• a full Megatron training loop through Reef's coordinator;
• TP > 1 or multi-node engines;
• throughput measurements.

The marker changes are covered by the CPU tests above.

AI assistance

• Claude Code (Opus 5) implemented the initial change.
• Codex audited it against the pinned dependencies, implemented the first publication-barrier fixes, and added and ran the original regression tests.
• Claude Code (Opus 5) ported the change onto the post-feat: let Reef own inference and training coordination #432 runtime while resolving the merge with main, updated the tests, and ran the CPU and GPU validation above.

Checklist

• The change is focused and contains no unrelated cleanup.
• Tests cover behavior changes, or this pull request does not change behavior.
• Public interface changes include contract tests, or no public interface changes are present.
• Affected user and developer documentation is updated, or no documentation update is required.
• Root README changes update both languages and README.i18n.yaml, or do not affect README content.
• An accepted RFC is linked, or this change does not require an RFC.
• Relevant pre-commit and test checks pass.
• Non-trivial AI assistance is disclosed above, or none was used.

Review and merge

The Merge Oncall is assigned automatically. See the
maintenance model
for review and merge requirements.

🤖 Generated with Claude Code

[github-actions]

Merge Oncall: @BobbyZhouZijian

The merge oncall will route reviews and keep this pull request moving toward a clear decision.