[Refactor] mem_cache: Centralize unified KV capacity pricing and admission by ch-wan · Pull Request #46977 · sgl-project/sglang

ch-wan · 2026-10-28T06:41:45Z

This PR is part of a stack (oldest at bottom):

Motivation

Capacity decisions are spread across allocator policies and scheduler state-slot bookkeeping, making it difficult to keep admission consistent with allocation.

Changes

Price capacity from host-side band snapshots, invalidate cached snapshots when allocator state changes, and use a shared allocation path. Consolidate prefill budgeting and state-slot reservations without changing the existing chain-specific admission rules.

Run the full seeded and corner-state pricing matrix in extra-a while retaining fast capacity and layout regressions in base.

Validation

Pre-rebase cleanup validation: the complete pricing property file passed 3 tests on 1×H200 (PyTorch 2.14.1, CUDA 13). This validates the unchanged test bodies; the registration now selects extra-a. Kernel-library versions were not recorded for this run. No tests were rerun after the rebase.

Historical validation through 655d6c5605:

Four-chain replay: 0 divergences; admission: 4,608 decisions, 0 differences, also under sync-debug.
Pricing properties: 3 tests / 942 subtests passed; planted pricing/memo defects were detected.
TP1 H200 A/B (PyTorch 2.14.1, FlashInfer 0.7.2): gpt-oss-120b and Qwen3.8-27B both scored 97.0% before/after on GSM8K-200; throughput changed +0.9% / +0.2%.

Separate CPU review matched 320,000 pricing queries and 270,000 budget observations. GPU benchmarks were not rerun on this rebased head.

CI States

Latest PR Test (Base): 🚫 Run #40683990892
Latest PR Test (Extra): 🚫 Run #40683990466
Latest PR Test (AMD ROCm 10): 🚫 Run #40683990555

…hain capacity on a host-only snapshot and take through one path Every capacity decision of the chain allocator is now a pure function of a snapshot of its bands (counts and byte frontiers, never hole positions), in kv/alloc/pricing.py: the joint gap pricing, reclaim search and compaction gate of a chain of two token ends, and the in-place float pricing, end-flush certificate, float retreat target and byte bound of a chain with a floating window. The per-shape policies take a fresh snapshot at each decision and perform the same realizations in the same order, so every decision is unchanged. Token takes (alloc, extend, window-tail extend, decode) share one path that prices, realizes, takes on the owner and binds the window, and state slots are taken through the chain allocator.

…hat each capacity price is what the allocation then does States come from the corner states of the trace module, seeded operation streams, and fills punched with holes and retired windows; demands sit at and just past what fits as the chain stands. On a chain of two token ends the price must equal the outcome; on a chain with a floating window each certificate must be realized and a positive byte shortfall must refuse. Runs on CUDA: pending reuse waits on the current stream.

…ricing snapshot between decisions and price on resolved bands The chain keeps its snapshot until the capacity epoch moves, a move gate's answer changes or a band's compaction mode changes: neither the gates, which are callables, nor the mode is a capacity field. The owner's virtual-id pool now moves the epoch when it is rebound (a paged step that takes no new page leaves it as it is), and an idle-time check compares the kept snapshot with a fresh one. A snapshot is built positionally with band positions kept by the chain, each pricing function resolves its bands once, and the policies keep their role constants. Every decision is unchanged.

…dmit every unified chain through one prefill budget that owns its state-slot gate The scheduler's prefill adder priced and gated a request's state slots itself: it read a per-slot token cost from the allocator, counted the state sub-pool's recoverable slots and debited them on admission. That knowledge moves into the prefill budget, so the adder takes the per-slot cost from the budget and asks it whether slots are exhausted and whether a load-back's slots fit. Every unified chain now builds one ChainPrefillBudget from its admission rules (how window tokens are admitted: none, split or joint; the state-slot cost). The rules reproduce the three budgets the chains used before: the plain token budget for full + state, the split window budget for full + window + state, and the joint byte reservation for full + window, which this class absorbs. Static allocators keep their budget classes. No admission decision changes.

…entries 122,198,200,211,215,220,267,268,305 Re-plants the scheduler.py fork-hunks (replant/...managers__scheduler.py.fork-hunks.patch) onto the post-sgl-project#46977 admission structure. tp_worker marker consumed (15/23 B4 markers still on the tree; per-file counts, damage-state from entry 310's map). PER-HUNK DECISIONS (all VERIFIED against the merged tree, not taken on faith): - entry 122 (sgl-project#40458): grammar-sync overlap disable RE-PLANTED for mixed mode (is_decode() OR is_mixed()) + the fork's trimmed comments; upstream's host-draft-only arm was the pre-sgl-project#40458 form (origin/main verified: absent). - entry 220 (sgl-project#39765, OPEN port): _has_pending_mamba_cache_update + the early pop_and_process/short-circuited disable-overlap RE-PLANTED into event_loop_overlap; the require_mlp_sync/disable/extra-buffer gate verbatim. - entry 198 (sgl-project#37612): THE CATCHER — B3-5 landed the raiser; the try/except MambaCowAllocError around init_next_round_input (res=NO_TOKEN, request stays queued; the revert block + unconditional alloc_group_end were already in the merged tree = the no-leak halves VERIFIED PRESENT) RE-PLANTED with the import. COW component test re-anchored to the #46970x kv_storage seam the B3-5 replant moved _state() behind (stub was the pre-split mamba_allocator shape). - entry 268 (sgl-project#46643, OPEN port -- VERIFIED NEVER MERGED: git log 12dc623..2a6c204 --grep 46643 = ZERO; only sgl-project#45491/sgl-project#45435/sgl-project#45472/sgl-project#45670 of the briefed merged-list arrived in-range): the UNCONDITIONAL batch_is_full clear RE-PLANTED on the scheduler side, subsuming the entry-198 hybrid-SSM arm and the priority-preemption/is_empty arms (policy side lands in B4-3). - entry 267 (sgl-project#46496): retract_decode new_token_ratio FLOOR clamp RE-PLANTED. - entry 211: triton_load_watch.prewarm_device_loads(self.model_worker) line RE-PLANTED between install() and mark_serving_started (the utils-side function survived the merge; only the call site was dropped). - entry 200 (sgl-project#43527) CAPACITY PAIR RE-PLANTED: model_runner .logical_max_total_num_tokens passes allocator= again (the configurator's PageInterleavePoolAllocator -> allocator.size branch SURVIVED the merge byte-identical, so the merged tree had the fork predicate with upstream's consumer-side multiplications = DOUBLE-COUNT); the three scheduler * kv_shard_widening multipliers back to the fork form (pool-stats shape, load_inquirer, and the init_req one-page-per-shard subtraction); tp_worker drops page_interleave_shard_size x + import. Upstream's invariant-checker hunk: OBSOLETE (the constructor no longer takes max_total_num_tokens -- absorbed by refactor). - entry 215 (sgl-project#44351/sgl-project#44401/sgl-project#44420): VERIFIED ALREADY PRESENT at the merge (filter_batch-every-step + reset ride, num_held_rows, budget_state FIRST statement, PD-gate deletion all byte-compare to the fork form) -- DEDUPED, no edit. Entry 107's local early-revision of sgl-project#40940 stays superseded (entry 215 rule), and its re-plant here would double-count. - entry 220's/305's pause-window busy accounting + write_fence_stream + is_empty arm: VERIFIED PRESENT (entry 305 landed patch-identical via sgl-project#46127; the fence assignment is the B3-2 controller half) -- DEDUPED. TEST FIXTURE RE-ANCHORS (not B2's: no markers, pure new-upstream-contract drift, per the entry-198 fixture-adaptation precedent): the two hybrid-ssm admission stubs gain disagg_prefill_bootstrap_queue=None (upstream's routing-key bootstrap_queue arg in calc_priority arrived in-range). PRUNE vs BREAKAGE (managers −523sub audit): the three managers test files vanished in-range (test_schedule_policy_dfs_weight, test_scheduler_batch_is _full_latch, test_uno_token_accounting) were ALL deleted by upstream's OWN sgl-project#45672 prune (3438ed6). Two are byte-identical upstream copies = zero fork delta. The third carried a 3-line fork comment/stub delta whose SEMANTICS the re-planted unconditional clear pins: its case was the is_empty-arm mirror of test_scheduler_hybrid_ssm_admission_latch (present, green, hybrid_swa/non-hybrid arms included) -- NOT re-homed, recorded. ACTUALS (damage -> after, CVD=9, prod venv, PYTHONPATH=python, per-file): test_scheduler_hybrid_ssm_admission_retry 4F -> ** 4P ** test_scheduler_hybrid_ssm_admission_latch 3F -> ** 4P ** test_mamba_prefill_cache_boundary 2F/12subF -> ** 5P/12sub EXACT (220 pin) ** test_scheduler_retract_ratio 1F -> ** 2P EXACT (267 pin) ** test_scheduler_init_req_max_new_tokens 1F -> ** 10P/6sub ** test_scheduler_grammar_mixed_tail_sync 1F -> ** 4P ** adjacent scheduler-stub files (held_row, running_batch_finished, pause_generation, disagg_idle, dp_converted, hicache_events) 56P/64sub UNMOVED -- no regression from the clear/catcher/capacity edits. Lint: ruff format + F821/E9 clean on all five files; imports OK (no GPU execution, no push). Co-Authored-By: OpenCode <noreply@opencode.ai>

…46977 pricing reconciliation (entries 177,266,268,269,270 + 217) 2 markers consumed (schedule_policy.py, scheduler.py). B4 tree count now 15. prefill_budget.py needed NO edit: sgl-project#46643's fix(1) (available_chunk_tokens capped by real free KV minus the one-page over-ask, defer-not-force) SURVIVED the merge byte-identical in the base class -- deduped with evidence (diff 309-tip vs HEAD shows only upstream's additive state_slot plumbing). DECISION TABLE (per-hunk, sgl-project#46977 new admission pricing vs entries 265..270): - entry 177 prefill_chunk_alignment: RE-PLANTED as-is -- kwarg + attribute + _align_prefill_chunk (LCM w/ kv_shard_granule) verbatim; the new tree's add_chunked_req grew a kv_shard_granule floor (upstream) that idempotently composes with the LCM (both floor-to-multiple); scheduler-side get_kvcache read RE-PLANTED too (dropped at the merge = the held_row/adder TypeErrors). - entry 266 sgl-project#43484 classic-pool slot ledger: RE-ADAPTED. The unified half is SUPERSEDED-BY-EQUIVALENT: sgl-project#46977 moved the exact fork ledger onto the budget (ChainPrefillBudget.rem_state_slots = mamba_allocator.schedulable + mamba_evictable; debit rem -= state_slot_tokens//cost in reserve(); gate via budget_state's state_slots_exhausted()) -- factor 1, charge == reserve // cost, byte-for-byte identical semantics; the adder-side twin was DROPPED to avoid a double ledger. The CLASSIC half SURVIVES as fork delta (the base-class budget never sees state): isinstance(HybridReqToTokenPool) gate (NOT hasattr -- MagicMock discipline, 266(c) preserved), per-req charge 1 + mamba_initial_tracking_slots (PR-bug math substitution kept, 266(a)), budget_state gate on the FULL charge (< per_req), add_one_req required-slot check incl. the defer_staged mamba branch, and slot-unit DEBIT threading at all five commit sites + _commit_prefill_admission. - entry 267: (B4-1) done. entry 268 fix(2): RE-PLANTED -- the second-unfinished-chunk guard now applies to ALL policies (the merged shortest-prefill-first-only restriction was upstream's own pre-sgl-project#46643 form; git log: sgl-project#46643 never merged). sgl-project#45789's policy-reopening term stays the recorded PROVEN NO-OP (270(i)); test re-pins ..._is_unconditional. - entry 269 page-align: RE-PLANTED verbatim at its 3 sites, composed with sgl-project#46977's floor_chunked_tokens (DCP grid) as page-align-then-grid-then-fork. - entry 270 fork-split: RE-PLANTED (gate on get_exec().mamba.enable_mamba_ extra_buffer + SGLANG_SPLIT_PREFILL_AT_MAMBA_FORK on our env -- the env union landed complete at the merge, environ.py untouched), incl. the continuation/commit/new-chunk/ignore-eos sites, the larger-host-hit re-clamp (270's fork cut lands AFTER page-align, order-safe), the truncation_align_size threading from scheduler (270(iii): _align_prefill_chunk call KEEPS its default arg deliberately), and the 'no other prefill may take its slot' return. - entry 217 (PrefillAdder closes admission on retraction): the finish_storage_prefetch_admission sites VERIFIED PRESENT (retraction/discard x2) -- DEDUPED, no edit. - entry 265: the brp required-slot block's defer/NO_TOKEN shape above; id-set half was B4-2. BEHAVIOR (DEFAULT-ON) reaffirmed for the owed restart-watch: sgl-project#45472's sync-free planning (B4-2 verified in-tree) + sgl-project#46643's UNCONDITIONAL per-round admission re-runs (re-planting this cluster) are the two entry-198-class behavior changes owed TTFT/starvation A/B at the prod restart (still banned; zero GPU execution this session). TEST RE-ANCHORS (upstream rewrites of these files silently dropped the fork's fixture halves -- restored from the 309-tip versions, each pinned by entry 215/266's own text): held_row's _Adder.add_chunked_req truncation_ align_size kwarg + get_kvcache allocator stub (entry-177 pair, upstream sgl-project#46980 rewrite dropped them); dcp_logical's pools NS gains state=None (#46970x _InitializedPools field) and test_scheduler_consumers_receives _logical_capacity re-adapted to the new seams (pool-stats via the storage shape, invariant checker's capacity kwarg REMOVED upstream, load_inquirer direct) -- same anti-double-widening intent, kv_shard_widening=4 fixture obsolete by design after the B4-1 re-plant. ACTUALS (damage -> after, CVD=9, prod venv, per-file): test_prefill_adder 7F/32subF (39F storm) -> ** 71P/0F/94sub ** test_dcp_logical_capacity 1F/6subF -> ** 5P/6sub ** test_scheduler_held_row_admission (regressed by my scheduler wiring, then repaired) -> ** 1P/2sub EXACT (215 pin) ** hybrid_ssm latch+retry family ** 9P ** (entry-198's '8P' + B4-1's case) DIR: unit/managers 1048P/48F/9S/439sub -> ** 1055P/3F/9S/474sub ** (3F = the two order-dependent test_mm_process_config ENV F's [reproduce at base] + test_io_struct 1F [diagnosed B4-5, fixture half]). Lint: ruff format + F821/E9 clean. No GPU; no push. Co-Authored-By: OpenCode <noreply@opencode.ai>

…re-plant) SESSION 2B COMPLETE Entry 312 records: the 5 cluster commits (9bd9007..5cde7a5) with per- cluster batteries; THE sgl-project#46977 RECONCILIATION DECISION TABLE (per-hunk survive-as-is / re-adapt / supersede, incl. the git-log proof that the spec-mamba four arrived in-range while the admission six NEVER merged -- correcting the briefing's dedupe list, and sgl-project#46643 fix(1)'s reconcile as already-merged base-PrefillBudget form = the reason prefill_budget.py got ZERO B4 edits); the entry-200 CAPACITY-PAIR restoration (our configurator's allocator.size branch survived -> upstream's consumer-side multiplications double-counted -> model_runner allocator= kwarg + tp_worker/scheduler subtractive forms re-planted; invariant-checker hunk OBSOLETE by upstream refactor); the io_struct NEEDS-ATTENTION VERDICT (upstream's own pointer_head guard vs stale fork stub -- test-side one-field re-anchor, matching sibling stubs); the -523sub PRUNE-vs-BREAKAGE audit (all three vanished managers files died to upstream's OWN sgl-project#45672 prune, two byte- identical to origin/main, the third's fork delta subsumed by sgl-project#46643 -- recorded, not re-homed); the B3 cross-note closures (MambaCowAllocError CATCHER landed; 216/217 admission interplay VERIFIED absorbed); the 2err population correction (replayssm_ring_fold pruned by sgl-project#44690, gdn cutedsl probe lives in the kernels zone, the spec dir's real err was merge damage, now 0err); the runner-zone finds the damage map missed (_init_model_hooks BOOT-INTEGRITY loss -- real decode-runner builds AttributeError, and component_manager's dropped residency half that parallel_executor already CALLS: both re-planted, 7F+crash-class closed); the true marker census (23 tree-carried, NOT the briefing's 20 -- B3 lesson applied; ALL consumed, B4=0, B1/B2/B5 diff-audited intact). GATES: trailing_endpoint ** 17P = 17P ** re-verified post-B4 (DoD-d, B3's gate unmoved). DoD(e) diff audit: ZERO mem_cache/ files changed, ZERO in qsa/moe-quant zones (moe_hook diff is marker-removal-only). NEW PINS: unit/managers 1009P/81F/9S/412sub/2err -> 1056P/2F/9S/474sub (2F = the two documented order-dependent mm_process_config ENV F's; the dir's 2err never reproduced -- population is the spec row's); unit/spec 363P/20F/21S/238sub/2err -> 388P/0F/21S/238sub/0err EXACT over 7 consecutive full-dir runs. OWED RESTART-WATCH REAFFIRMED: sgl-project#45472's sync-free planning (verified present, upstream form) + sgl-project#46643's unconditional per-round admission re-runs (re-planted) ride the entry-201 restart; entry-198's TTFT/starvation measure STILL OWED. Header: STACK -> qwennext-stack-312+merge5-b4-2026-10-29; +merge5-b4 wave tag; tip -> 5cde7a5; entries 1-312; INTEGRATION OWED annotated B4 LANDED (boot ban STAYS). wanted_stack: UNTOUCHED (no S-item changed by B4 findings -- recorded here). ENV debts recorded NOT-fixed: pcie_vocab_gather GPU suite (27F identical at the merge commit), openai 2.6.1<2.8.0 serving_responses import (upstream's own beam.py), imageio (whole diffusion import chain). NOTHING PUSHED; tree clean; zero GPU execution; zero prod boots all session. Co-Authored-By: OpenCode <noreply@opencode.ai>

ch-wan requested review from hnyls2002 and hzh0425 as code owners October 28, 2026 06:41

ch-wan mentioned this pull request Oct 28, 2026

[Refactor] mem-cache: Describe KV address translation with ID spaces #46978

Merged

ch-wan requested review from alphabetc1, huangtingwei9988, ispobock and xiezhq-hermann as code owners October 28, 2026 06:41

ch-wan mentioned this pull request Oct 28, 2026

[Refactor] mem-cache: Route layer access through KV subpool views #46979

Merged

ch-wan requested a review from merrymercy as a code owner October 28, 2026 06:41

github-actions Bot added hicache Hierarchical Caching for SGLang memory-pool labels Oct 28, 2026

ch-wan requested review from Ying1123, hanming-lu and yizhang2077 as code owners October 28, 2026 06:41

This was referenced Oct 28, 2026

[Refactor] mem_cache: Build unified KV pools from layer declarations #46974

Merged

[Fix] mem_cache: Forward HiCache completion events to recurrent state #46975

Merged

[Refactor] mem_cache: Use one allocator for unified KV sub-pools #46976

Merged

ch-wan force-pushed the cheng/refactor/unified-kv-pricing-and-budget branch from 1a3f934 to ed46981 Compare October 28, 2026 08:01

ch-wan force-pushed the cheng/refactor/unified-kv-chain-allocator branch from 554979f to 17c4546 Compare October 28, 2026 08:01

ch-wan force-pushed the cheng/refactor/unified-kv-pricing-and-budget branch from ed46981 to d2a93e2 Compare October 28, 2026 11:48

ch-wan force-pushed the cheng/refactor/unified-kv-chain-allocator branch 2 times, most recently from 7b2f0b5 to 74fdbc0 Compare October 28, 2026 11:57

ch-wan requested a review from Jialin as a code owner October 28, 2026 11:57

ch-wan added 5 commits October 28, 2026 04:57

[Chore] memory: simplify comments and regression tests

6b17a73

Base automatically changed from cheng/refactor/unified-kv-chain-allocator to main October 28, 2026 11:57

ch-wan force-pushed the cheng/refactor/unified-kv-pricing-and-budget branch from d2a93e2 to 6b17a73 Compare October 28, 2026 11:57

ch-wan merged commit 5f6075f into main Oct 28, 2026
61 of 122 checks passed

ch-wan deleted the cheng/refactor/unified-kv-pricing-and-budget branch October 28, 2026 11:58

Noridom1 mentioned this pull request Nov 5, 2026

[Fix] flashattention: Resolve window KV pool for frozen-KV MTP draft (illegal memory access on Gemma-4) #48594

Open

3 tasks done

Navigation Menu

Repository navigation

[Refactor] mem_cache: Centralize unified KV capacity pricing and admission - #46977

[Refactor] mem_cache: Centralize unified KV capacity pricing and admission#46977
ch-wan merged 5 commits into
mainfrom
cheng/refactor/unified-kv-pricing-and-budget

ch-wan commented Oct 28, 2026 •

edited by github-actions Bot

Loading

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

ch-wan commented Oct 28, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Changes

Validation

CI States

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

ch-wan commented Oct 28, 2026 •

edited by github-actions Bot

Loading