Skip to content

feat(flexkv): integrate chunked prefetch policies and hybrid restore lifecycle - #38451

Closed
XingLiu1 wants to merge 1 commit into
sgl-project:mainfrom
XingLiu1:feat/flexkv-prefetch-policies
Closed

XingLiu1 wants to merge 1 commit into
sgl-project:mainfrom
XingLiu1:feat/flexkv-prefetch-policies

Conversation

@XingLiu1

@XingLiu1 XingLiu1 commented Sep 8, 2026 •

Copy link
Copy Markdown

Superseded by XingLiu1/sglang#6. The change is now a review PR targeting agent/flexkv-dsv4-main, the branch behind upstream #31781. The mistaken direct push to that branch was reverted; prefetch is not yet incorporated. The review head 2f91f9f5f0 has the same source tree as the version that passed 167 CPU tests and 12 subtests. Companion: FlexKV #291.

Motivation

FlexKV's whole-prefix prefetch cannot stop issuing remote work when a queued request becomes a scheduling candidate or exhausts its prefetch budget. Integrate the optional chunked API from taco-project/FlexKV#291 so SGLang can consume a safely drained partial prefix and compute the remainder.

Modifications

  • Start prefetch on queue entry using the full hash-chain token input; issue demand at scheduling candidacy and propagate abort independently of HiCache backend configuration.
  • Consume terminal, leased CPU prefixes through CPU-only LOOKUP and held GET. Synchronize leader decisions across TP ranks and account only actually published L3 spans.
  • Compose FlexKV I/O with the existing hybrid radix implementation. Request-owned restore slots and SWA/state checkpoints retain one ownership path until normal cache insertion or rollback.
  • Check admission shape/capacity before H2D, retain pending restore slots across scheduler iterations, and let an empty batch retry after host-cache pressure clears.
  • Add CPU adapter/admission contract tests. No model kernels or underlying transfer protocol changes.

Depends on the companion FlexKV PR. Rebased on SGLang 5aab054ec8; the default remains unchanged unless FlexKV chunked prefetch is enabled. Configuration, interfaces, sequence diagrams and monitoring boundaries are documented in the companion PR's docs/design/chunked_prefetch_reference.md and review document. chunk_max_blocks is the single chunk-size option; remove the former experimental chunk_target_bytes field.

Accuracy Tests

  • Rebased revisions: 57 SGLang tests passed plus 12 subtests: 29 FlexKV adapter/admission cases and 28 existing PrefillAdder cases.
  • Companion FlexKV: 456 passed, 9 Python-radix SWA skips; C++ SWA cases executed. Tests used an isolated CPU/native container, not a GPU serving run.
  • Ruff/isort/formatting, registered-test validation, patch applicability and diff whitespace checks passed.

Earlier implementation snapshots passed 25 real H20/Mooncake tests, a Full+SWA byte round trip, and 210/210 formal requests each for GLM-5.2 and V4 Flash (TP8/DP1). Full generated-token comparisons were used. Those results predate this rebase and chunk-size simplification; repeat GPU/model checks on the paired PR revisions before marking ready.

Speed Tests and Profiling

No throughput or latency improvement is claimed. Earlier C1/C8/C32 workloads included shared prefixes, queueing and CPU/GPU cache hits. Same-cache-state repeated A/B, IPC/forward timeline alignment, longer load/fault/resource checks, and DP/layerwise coverage remain outstanding. Dedicated chunked-prefetch Prometheus instrumentation is follow-up work and is explicitly separated from existing metrics in the design.

Checklist

  • Format changed code and check imports/lint with repository-pinned tools.
  • Add and register unit tests; run existing PrefillAdder regressions.
  • Document paired configuration, API/lifecycle, diagrams and monitoring boundaries in the FlexKV dependency.
  • Repeat model accuracy and speed benchmarks on the rebased revisions.
  • Complete GPU/runtime and broader topology acceptance before marking ready.

CI States

Latest PR Test (Base): ❌ Run #34200927062
Latest PR Test (Extra): ❌ Run #34200926458
Latest PR Test (AMD ROCm 7.2): ❌ Run #34200926774

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

hicache Hierarchical Caching for SGLang unified-radix-cache

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant