[HiCache] Rework the buffer-mode storage prefetch pipeline and retry bookkeeping

[xiezhq-hermann]

Motivation

Buffer-mode HiCache (host memory used as transport staging, no L2 tier) decided a staged L3 hit's fate inside the prefill adder, after admission had already charged host-hit tokens that a shrunk device prefix could no longer splice. Under load this produced phantom host-hit accounting, holds that wedged admission, and unbounded storage-prefetch retries. PD-prefill and PP schedulers had also stopped draining HiCache events after a scheduler refactor.

Modifications

Prefill adder / scheduler

• PrefillAdder.add_one_req selects a prefill shape without allocating (_select_prefill_admission), materializes host hits under the prefix lock via init_load_back, then commits (_commit_prefill_admission). init_load_back may return None to request another admission attempt; FULL over-delivery shrinks the planned extend range, under-delivery re-selects with no host hit.
• The SWA budget gate is factored into _swa_admission_gate and keeps the ring-slot exact-fit comparison for the unified-KV SWA ring.
• Storage prefetch retries move into StoragePrefetchRetries (mem_cache/storage_prefetch.py): paced miss polls, immediate re-issues for known hits whose anchor moved, and a per-request cap via --hicache-storage-prefetch-retry-max-attempts (default 8; --hicache-storage-prefetch-retry-poll-interval default becomes 8). Requests past the cap are admitted with whatever the device holds.
• HiCache event draining is factored into _process_hicache_events and called from the PD-prefill and PP scheduler loops.

Buffer-mode pipeline

• Staged holds are anchored by a FULL-only pin (match_full_device_prefix, inc_full_pin, dec_full_pin on the Python and Rust tree cores), so the device prefix under a parked fetch cannot be evicted while the SWA segment lock is left alone.
• Aux components size prefetch staging in prepare_prefetch (PreparePrefetchResult.staging_tokens) and allocate it once the hit is known (alloc_prefetch_staging).
• Admission-time device-capacity deferrals are bounded (max_staged_admission_defers = 32); past the cap the hold is dropped and the request recomputes.
• Parked hits keep their turn ahead of newer hits, stale prefetch acks are ignored once a request re-issued its query, misses are accounted from the published hit count, and EAGLE bigram keys keep their boundary token when a prefetch span is trimmed or staged.
• swa_transient_size exposes SWA slots owned by staged holds to the scheduler invariant checker.

Adaptations to upstream

• _swa_admission_gate incorporates the _swa_req_ring exact-fit comparison added upstream in [AMD][DSV4] Reland unified-KV pool sizing and SWA ring accounting, fully gated #38192.
• The server args land in arg_groups/fields/memory.py (upstream moved the field declarations there).
• Unfulfilled-prefetch reason labels keep the existing upstream names (shrunk).
• The Rust match_full_device_prefix walks the full radix namespace (extra key + cache salt) like match_prefix does upstream.
• The Rust tree core gets the same staging_tokens contract as the Python components: SWA/Mamba PREFETCH builds size their placeholder keys from the planned staging and carry no host buffer, which the Python pipeline attaches once the hit is known (alloc_prefetch_staging). Without this the Rust adapter rejected the new keyword and the Rust-backed shared suite failed.
• Known gap (Rust tree core): buffer-mode consumption now repairs SWA tombstones under the loaded window through two new Python tree-core helpers, swa_tombstone_ranges and attach_swa_window, which have no Rust port yet. The Rust adapter exposes them as explicit NotImplementedError stubs (same pattern as its root_node), and the three shared-suite tests that exercise that path (test_buffer_only_load_back_uses_full_behind_swa_tombstone, test_buffer_only_load_back_reuses_partial_masked_full, test_buffer_only_load_back_trims_head_published_by_sibling) are skipped under the Rust backend. The last one passed under Rust before this PR, so this is a deliberate coverage regression for buffer-only + SWA + L3 on the Rust core until the helpers are ported; the Python core is unaffected.

Original commits

• 0e78bbe270
• bdc1fb6985
• 57b0a7f139

Accuracy Tests

N/A (scheduler / cache bookkeeping change; no numerics affected).

Benchmarking and Profiling

Validated on an internal buffer-mode + L3 setup: warm-cache prefix hit rate no longer collapses under high concurrency and no admission wedges were observed. No absolute numbers are included here.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging.
• Please feel free to join our Slack channel at https://slack.sglang.io to discuss your PR.

Local validation (details in the first comment): Rust cargo test --locked (865 passed incl. the new pin test), cargo check --features python-extension,inspection, cargo fmt --check, cargo clippy -D warnings; Python unit suites test_unified_radix_cache_unittest.py, test_rust_unified_radix_cache_unittest.py, test_storage_prefetch_lifecycle.py, test_prefill_adder.py, test_scheduler_hicache_events.py, test_buffer_mode_sidecar.py, test_hicache_staged_write_back_dispatch.py, test_rust_tree_core.py, test_tree_core_registry.py; ruff, ruff format, isort, codespell on changed files.

---

CI States

Latest PR Test (Base): 🚫 Run #34939088034
Latest PR Test (Extra): ❌ Run #34939087207
Latest PR Test (AMD ROCm 10): ⏳ Run #34939087808

[xiezhq-hermann]

/tag-and-rerun-ci

[xiezhq-hermann]

Local validation (GB300 devserver, prod venv):

• Rust crate rust/sglang-radix-tree: cargo test --locked 865 passed (includes the new FULL-pin test and the reworked SWA/Mamba prefetch-build tests); cargo check --features python-extension,inspection; cargo fmt --check; cargo clippy -D warnings on the default, python-extension,inspection, and --no-default-features --features tch/doc-only shapes.
• test_unified_radix_cache_unittest.py (Python core): 1251 passed, 1435 skipped, 70 subtests.
• test_rust_unified_radix_cache_unittest.py (Rust core): 2686 ran, OK, 1456 skipped (pristine main baseline: 2610 ran, OK, 1388 skipped). The 3 buffer-mode SWA window-repair tests are skipped under Rust as described above.
• test_storage_prefetch_lifecycle.py, test_prefill_adder.py, test_scheduler_hicache_events.py, test_buffer_mode_sidecar.py, test_hicache_staged_write_back_dispatch.py, test_scheduler_chunked_req_gate.py, test_invariant_checker.py, test_stat_loggers_di.py, test_rust_tree_core.py, test_tree_core_registry.py: all green (98 tests).
• ruff check --select F401,F821,UP037, ruff format --check, isort --check-only, codespell on the changed files: clean.
• Not run locally: GPU end-to-end HiCache L3 serving; covered by an internal setup, see the PR description.

[xiezhq-hermann]

Rebased onto current main (5c2de3f) to resolve conflicts. Adaptations from the rebase:

• storage_attachment.py: took upstream's buffer-pipeline anchor release from [HiCache] Release buffer prefetch anchor locks during storage cleanup #38483 (equivalent to what this PR carried).
• schedule_policy.py: re-expressed the exact-chunk-fill path from [AMD][gfx95] Fill the chunked-prefill compute budget exactly #32888 inside the refactored _select_prefill_admission / _commit_prefill_admission (raw-length fit check, verbatim chunk remainder, compute_charge at commit).
• test_storage_cleanup_releases_buffer_prefetch_anchor (new in [HiCache] Release buffer prefetch anchor locks during storage cleanup #38483) adapted to the FULL-pin anchor API (try_lock_anchor(req_id, remaining_full_tokens), release via tree_core.dec_full_pin).
• The NPU C128SidecarComponent.build_hicache_transfers override ([NPU] Support DSV4 host memory cache management #37382) now accepts the staging_tokens keyword.

Re-validated locally on the rebased tree: cargo test --locked 866 passed + check/fmt/clippy clean; Python shared suite 1351 passed; Rust-backed shared suite 2688 ran, OK; lint clean.

[xiezhq-hermann]

/tag-and-rerun-ci

[xiezhq-hermann]

Rebased onto current main (2fca6d6) over #39318 (attempt-scoped CacheRequestHandle). Conflicts were the request-id vs handle keying on the lines this PR reworked; resolution keeps this PR's logic and keys the pipeline, cache prefetch state, anchor locks and staged holds by the handle. StoragePrefetchRetries stays keyed by rid (it paces requests in the waiting queue, so an attempt change must not orphan a pending retry). Tests wrap ids in CacheRequestHandle and give fake requests a cache_request_handle.

Second commit fixes two CI failures from the first run:

• test_swa_prefetch_commit_end_to_end (Rust core) still used the old host_indices= SWA PREFETCH build contract; it now asserts the staging_tokens= contract.
• The 1-GPU unit shard OOMed on test_unified_radix_cache_unittest.py: init_hicache registers cache.shutdown with atexit, pinning every HiCache fixture until process exit (~30 GiB on a 32 GiB card on main already; this PR's extra L3 tests tipped it over). Fixtures now unregister the hook via addCleanup; peak device memory for the suite drops from ~33 GiB to ~7 GiB.

Re-validated locally on the rebased tree: shared cache suite under the Python core and the Rust core both OK, test_rust_tree_core_integration.py 105 passed, and the storage-prefetch/scheduler/dispatch/sidecar unit files pass.

[hzh0425]

Awesome

[xiezhq-hermann]

/tag-and-rerun-ci