[Refactor] mem-cache: Centralize recurrent state ownership and access by ch-wan · Pull Request #46980 · sgl-project/sglang

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

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

Motivation

Request rows, recurrent-state allocation and per-layer storage access currently overlap across request tables and KV pools. Centralizing those owners makes request admission rollback and layer routing consistent for static and unified storage.

Changes

Build recurrent state with the pool builder, bind state parts to the layer router, and translate state IDs through their allocator-owned space.
Acquire request rows and auxiliary state together, preserving existing ownership when allocation fails.
Migrate recurrent, short-convolution, N-gram, host-transfer and platform consumers to the new owners.
Route GDN prefill-graph and EAGLE plugin bindings through the state owner and layer router, with fixtures that exercise noncontiguous layers and storage order.

Clarify state ownership and routing comments, and use CustomTestCase in the affected tests.

Validation

The cleanup changes comments and test base classes; no new state/router execution result is claimed. Tests were not rerun after the rebase.

Historical validation through 2788de7080:

Four allocator traces: 0 divergences; admission: 4,608 unchanged decisions, including sync-debug.
TP1 H200, unified HiCache / FA3, GSM8K-200: gpt-oss-120b 96.5% → 97.0%, Qwen3.8-27B 97.0% → 98.0%, zero errors. Per-run kernel versions were not fully recorded.

At published f647f7bede, CPU regression probes reproduced the old GDN/EAGLE failures and passed after repair, including layer mapping, FP8 scales and tensor identity. GPU serving and graph execution were not rerun on this head.

CI States

Latest PR Test (Base): 🚫 Run #40684177627
Latest PR Test (Extra): 🚫 Run #40684177620
Latest PR Test (AMD ROCm 10): ❌ Run #40684177628

…ecurrent state through the KV pool's layer router The layer-wise load-back counter was registered twice: on the KV pool, and on the request table when the host stack also restores Mamba state. State reads waited on the table's copy, KV reads on the pool's. The counter now has one owner, the KV pool. The request table binds its state layers to the pool's layer router as named parts (the recurrent state, and Qwen4-Exp's short-conv and N-gram side states), and attention backends, Jet-Nemotron and Qwen4-Exp read state through the router. KV reads wait whenever a counter is registered, as before; state reads wait only once the host stack says it restores state layer by layer, which is when the table's counter used to be registered. The N-gram context still waits for the first state layer, and the copy-on-write state copy for the last.

…a table's recurrent state in one StateStorage The state pool, its slot allocator, the layer map, the PLE side states and the int8 checkpoint pool move into StateStorage, which also takes the slot work of a request's lifecycle: the state block of the table's alloc, the ping-pong tracking buffer, free_mamba_cache's slot frees, the copy and the clear. The table keeps its id columns and the id-only ping-pong bookkeeping, and still answers the old attribute names from its storage while their callers move.

… recurrent state through the KV pool's router and the token allocator Readers stop going through the request table for the state storage. The forward side (linear-attention backends, graph runners, speculative workers) takes the state pool from the KV pool's router, which binds it with the state routes; the scheduler, the radix tree and its Mamba component, HiCache, the external linker, the observers and the runner's copy-on-write take the state and its slot allocator from the token allocator, which the KV cache configurator binds to the request table's state. A draft worker shares the target's allocator and keeps its own router binding.

…a request's row and state slots through one acquire The state table's alloc stops handing out state slots; it binds rows only, as every request table does. acquire binds the rows and then, when the model has recurrent state, the state and tracking slots, from the state the token allocator carries. The scheduler's admission (alloc_req_slots) and PD decode's pre-allocation call it; the capacity check reads the same state. If the state slots run out after the rows were bound, acquire gives back the rows it bound and the slots it handed out before the failure surfaces, where the table's alloc used to stop at an assertion with both taken. The radix tree's state component gives a finished request's slots back through the state as well (StateStorage.release).

…e): Build the unified state storage in the pool builder The unified request table built its recurrent-state storage (the state pool, the disabled PLE side states, the layer map) inside its own constructor, and the builder then took it from the table. The table now holds the request rows and their slot-id columns only: the unified pool builder constructs the StateStorage, wires the state band's slot allocator into it, and hands it to the chain allocator, which also clears it with its bands, and to the KV pool. A fused draft's state is no longer a shallow clone of the target's table: the state binder builds a StateStorage over the draft block that shares the target's slots, and the KV cache configurator binds the routes of whichever state the pool init returned (the table's own under static placement, the shared allocator's for a draft that shares a unified target's table). The request table's state delegates go with their last callers: speculative commits and the SConv and NPU backends read the speculative buffers from the state pool, Qwen4-Exp writes its N-gram context through the router's part view, and PD registers the state buffers from the allocator's state pool.

… Ask the KV pool's router what PD and HiCache dispatch on PD and HiCache chose their paths by the KV pool's class: SWAKVPool or BaseSWAKVPool for a pool with a sliding-window sub-pool, HybridLinearKVPool for a pool that keeps its full-attention KV in one inner view beside the state, and, in the decode offload manager, the unified window pool's own class. The unified pools subclass those classes only so that these checks pass. The checks now ask the pool's layer router: whether it has a sliding-window sub-pool, the full-attention view, and the one token KV view a pool wraps (KVLayerRouter.wrapped_view). DeepSeek-V4, which describes its window state per entry without a window sub-pool, is named where the old check covered it through BaseSWAKVPool. HiCache sizes a wrapping pool by its sub-pool views. Every pool answers as its class did, so the paths taken are the same.

… combined kind checks outside the pools as plan facts A few checks outside the KV pools paired two kind predicates to name a combination: - EAGLE under unified memory needs a unified target: "hybrid SWA or a Mamba hybrid". - Post-capture KV sizing under unified memory exists for one chain only: "hybrid SWA and no Mamba". - A multi-layer MTP draft that owns the layer at its depth was flagged "hybrid SWA MTP draft", and that flag gated the draft's own Mamba state. The first two now read the lifetimes the model declares (kv.decl.declared_lifetime_types) and a fact of the plan (kv.plan.resized_after_capture: only the full plus window chain resizes after capture). The draft flag becomes the fact it describes, draft_owns_its_depth, without the window kind; the only model that numbers its MTP layers by depth is a hybrid SWA model, so it takes the same value.

PR sgl-project#46980 (306a699) removed HybridReqToTokenPool.layer_transfer_counter and routes recurrent/PLE load-back waits through the KV pool. Read get_token_to_kv_pool().layer_transfer_counter for the existing no-transfer eligibility guard. Keep the fused path when no counter is installed and the native fallback when transfer is active; no kernel arithmetic, launch, flag or synchronization change. CPU regression cases use a real main HybridReqToTokenPool without the obsolete attribute.

sgl-project/sglang#46980 removed the request pool's mamba2_layer_cache and short_conv_layer_* accessors; SGLang reads a layer's recurrent state through token_to_kv_pool.layer_router.read(layer_id, part). Read them the same way.

… pools by mamba_map - the runner sized the shared scratch with K1 at one bank row; KDA layers read up to MAX_BLOCKS = 8 of them - at eight rows K1 needs 789 KB against the 738 KB buffer, so it overran into other allocations and faulted the GPU - k1.scratch_bytes now sizes for MAX_BLOCKS, and kda_pre refuses a scratch shorter than its build's layout - SGLang dropped HybridReqToTokenPool.mamba2_layer_index (sgl-project/sglang#46980); K1 now reads mamba_map - both stayed hidden while K1 declined every KDA layer of the served checkpoints

…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>

…traverse — entries 102,103,119-supersede,120,194,195,261-264 3 markers consumed (eagle_utils RE-PLANTED; spec_utils re-plant; eagle_info VERIFIED-ABSORBED). B4 tree count now 12. PER-FILE: - eagle_utils.py (entry 103): RE-PLANTED the tensors-after-early-return graft (num_needed_tokens == 0 returns before any torch.tensor -- the 17F were the fail-fast + alloc-path rebuilding CPU tensors every round; topk>1 max_alloc_len back to the list-builtin max) MERGED with upstream's in-range mirrorless block: it rides AFTER the tensors now (it reads cur_kv_lens_cpu) with a getattr-guarded seq_lens_cpu read -- identical semantics for the real ScheduleBatch (attr always present), stale-stub-proof. - spec_utils.py (entries 102,120,195; 261 deduped): * entry 261 sgl-project#45491: DEDUPED with evidence -- git log 12dc623..2a6c204 proves 573b571 IN RANGE; its verify-window semantics SURVIVE upstream's later refactors (sgl-project#45504 kernel move + sgl-project#46990 typed views collapsed the triton assign_extend_cache_locs pair) as fill_front_out_cache_loc_func(num_correct_drafts+1, batch.out_cache_loc, ...) carrying sgl-project#45491's OWN comment text; the fork-hunks window-form is thus superseded-in-place (origin/main has neither old name -- the replacement landed within the same range). * entry 195 sgl-project#44317: RE-PLANTED the traverse_tree host walk (tolist + numpy bitmask view; int words) + the needs_vocab_mask() skip (the backend-side methods ALREADY landed upstream in constrained/ -- consumer half was the dropped one) + the leftover-grammar_mask clear on the grammars-present-but-unconstrained path (the merged tree had the clear on the SUCCESS path only -- the stale-mask class this entry exists for). traverse_tree 3F -> green. * entry 102 (PLE on the ReplaySSM branches) + entry 120 (circular-ring fold before a plain extend): RE-PLANTED, RETARGETED to sgl-project#45504/sgl-project#46980 reality: kernel import path attention.fla.* -> attention.linear.* (upstream's kernels regroup), and the fold reads mamba_pool.get_speculative_mamba2_params_all_layers() (upstream sgl-project#46980 deleted the req-pool delegate; same route the verify-commit uses now). test_replayssm_extend_fold collection 1err -> 2P (the documented 'pair': its sibling test_replayssm_ring_fold.py was PRUNED by upstream's own sgl-project#44690 (file gone; a stale .pyc remains) and the 'gdn cutedsl probe' err belongs to kernels/ops/attention (6 CUDA-import collection errs there -- B1/B5 zone, not this dir): the spec-dir population is now 0err. test_replayssm_verify_ple_commit 1F -> green (stub re-anchored: router part_view seam + PLE pools off-pin restored -- the entry-102 pin the upstream file rewrite dropped, present in the 309-tip test verbatim). - eagle_info.py (entry 194): VERIFIED-ABSORBED -- the call sites pass num_token_blocks=None + kv_lens_sum, and runner_ids.pack_read_stream auto-sizes with the fork's OWN helper (spec_kv_index_token_blocks, now hosted in kernels/ops/attention/kv_indices) -- the 'AUTO-SIZED upgrade' B3-4 recorded; identical tiling math. - Entry-119 supersession chain CLOSED as briefed: sgl-project#45435's upstream form (tails padded -1 + mamba_lazy_undo_prealloc_at_boundary + untracked_decode_batch_idx gate) is the tree's reality (B4-2 audit); entry 119's production half stays dead. - test_ngram_mamba_verify_update.py (B2-marked file, B2 marker INTACT): restored only the two PLE off-pins + the entry-102 comment that the 309-tip fork version carried and upstream's rewrite dropped -- my re-plant re-creates the condition they exist for. ACTUALS (damage -> after, CVD=9, prod venv, per-file): test_eagle_decode_allocation 17F -> ** 18P ** (16-matrix + capacity + 1) test_spec_utils_traverse_tree 3F -> green (file pin incl.) test_replayssm_extend_fold 1err -> ** 2P ** test_replayssm_verify_ple_commit 1F -> green DIR unit/spec 363P/20F/21S/238sub/2err -> ** 388P/0F/21S/238sub/0err ** (two consecutive full-dir re-runs EXACT; DoD(b): zero non-env F). Lint: ruff format + F821/E9/F401 clean on the six files. No GPU; no push. Co-Authored-By: OpenCode <noreply@opencode.ai>

… backend (#94) ## Motivation sgl-project/sglang#46980 removes `HybridReqToTokenPool.mamba2_layer_cache`. glm53_flash calls it in `_bind_kda` on every routed step, and `SGLANG_AMD_MONO_CHECK` calls it in `_regions`. On SGLang with that change, for example `lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261028` (SGLang `5d8e90a0f9`), a GLM-5.3-Flash server with the monokernel exits at CUDA-graph capture: ``` AttributeError: 'HybridReqToTokenPool' object has no attribute 'mamba2_layer_cache' ``` `test_contract.py` does not catch it: the name is still in `memory_pool.py`, on `MambaPool`. ## Modifications - `spec.py`: `_bind_kda` takes the linear-attention backend and reads its `state_layer_cache(layer_id)`, as SGLang's KDA backend does since #46980. - `check.py`: `_regions` reads the same cache through `get_attn_backend().linear_attn_backend`. - `SGLANG_READS`: `state_layer_cache` in `hybrid_linear_attn_backend.py`; `mamba2_layer_cache` and `get_req_to_token_pool` go. This needs SGLang with sgl-project/sglang#46980. ## Accuracy Tests MI355X, `lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261028` (SGLang `5d8e90a0f9`), `amd/GLM-5.3-Flash-Quark-MXFP4`, TP4, `test/e2e/test_glm53_flash_server.py`: | Test | `main` | This PR | | --- | --- | --- | | `TestMonokernelDecode` (2 tests) | error: the server exits at graph capture | passed | | `TestMonokernelCheck` | not run | passed | | `TestMonokernelHiddenStates` | not run | passed | | `TestMonokernelFault` | not run | fails | `main` is eb0f20d; glm53_flash has not changed since. `TestMonokernelFault` also fails on `main` on SGLang `5044779f`, which is before #46980, so this PR does not cause it. ## Speed Tests and Profiling None: the launches and their arguments do not change. ## Checklist - [x] `pre-commit run --all-files` passes. - [x] Tests under `test/` cover the change. (`test_glm53_flash_server.py` fails before this PR and passes after, on SGLang with #46980.) - [x] The SGLang commit or PR this was tested against is named above. - [x] Accuracy and speed results are attached when kernel or forward code changes. - [ ] A new model adds one line to `.github/CODEOWNERS` that names its owner. (No new model.)

- test_mamba2_replay_sizing.py::test_puzzle_uses_stock_path_in_hooks_and_pool_dispatch built NemotronHConfig() with defaults, but mamba2_cache_params reads n_groups, which only checkpoint configs set. Set it on both configs. - test_mamba2_spec_replay.py::test_mixer_and_backend_commit: main's sgl-project#46980 made update_mamba_state_after_mtp_verify read linear_attn_backend.state_pool. Add it to the test's fake backend. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

ch-wan mentioned this pull request Oct 28, 2026

[Fix] hicache: Build hybrid recurrent host pools consistently #46981

Merged

ch-wan requested review from Qiaolin-Yu, Ying1123, hnyls2002 and merrymercy as code owners October 28, 2026 06:42

github-actions Bot added hicache Hierarchical Caching for SGLang npu apple-silicon unified-radix-cache labels Oct 28, 2026

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

github-actions Bot added memory-pool mooncake labels Oct 28, 2026

ch-wan requested review from ByronHsu, Duyi-Wang, HaiShaw, ShangmingCai, alphabetc1, hanming-lu, hebiao064, huangtingwei9988, hzh0425, iforgetmyname, ispobock, ping1jing2, stmatengss, whybeyoung, xiezhq-hermann, yeahdongcn, yizhang2077 and yuan-luo as code owners October 28, 2026 06:42

ch-wan added 11 commits October 28, 2026 04:59

Route GDN prefill graph state through the KV layer view

98b4f93

Update GDN graph fixtures for state views and identifier spaces

bd22d32

fix: bind EAGLE buffers through KV layer routes

d0db427

[Chore] memory: simplify comments and regression tests

3256b64

ch-wan force-pushed the cheng/refactor/unified-kv-state-access branch from de1027c to 3256b64 Compare October 28, 2026 11:59

ch-wan merged commit 306a699 into main Oct 28, 2026
25 of 30 checks passed

ch-wan deleted the cheng/refactor/unified-kv-state-access branch October 28, 2026 11:59

akhilg-nv mentioned this pull request Oct 28, 2026

[Fix] mem_cache: Build an empty KVLayerRouter from empty layer maps #47127

Merged

4 of 5 tasks

amd-cgilli mentioned this pull request Oct 29, 2026

[Feat] [MI355X] qwen38_flash_next: Add the Qwen3.8-Flash-Next 180B monokernel sgl-project/sgl-amd-mono#29

Merged

5 tasks done

ChangLiu0709 mentioned this pull request Oct 29, 2026

[ROCm][AMD] Qwen3.8-Flash-Next gfx950 / MI355X Performance Optimization #44304

Open

ChangLiu0709 mentioned this pull request Oct 29, 2026

[AMD][MI35X] Qwen3.8-Flash-Next FP4 SGLang agentic MTP: 20261031 image, monokernel at c1-c8 / [AMD][MI35X] Qwen3.8-Flash-Next：20261031 镜像，c1-c8 启用 monokernel SemiAnalysisAI/InferenceX#4542

Merged

6 of 11 tasks

sha7doww mentioned this pull request Oct 30, 2026

[Fix] mem_cache: Make PLE logical clearing safe while KV memory is offloaded #40480

Open

3 of 5 tasks

eppaneamd mentioned this pull request Oct 30, 2026

[Feat] [MI355X] qwen38_flash_next: Serve TP2 decode of the MXFP4 model sgl-project/sgl-amd-mono#93

Draft

3 of 5 tasks

This was referenced Oct 30, 2026

[Fix] glm53_flash: Read the KDA layer cache from the linear-attention backend sgl-project/sgl-amd-mono#94

Merged

[Feat] [MI355X] glm53_flash: Serve MTP verify steps on the monokernel sgl-project/sgl-amd-mono#71

Merged

nvpohanh mentioned this pull request Nov 2, 2026

[Feat] [NVIDIA] mamba2: Add compact speculative-state replay #41968

Open

Navigation Menu

Repository navigation

[Refactor] mem-cache: Centralize recurrent state ownership and access - #46980

[Refactor] mem-cache: Centralize recurrent state ownership and access#46980
ch-wan merged 14 commits into
mainfrom
cheng/refactor/unified-kv-state-access

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