[HiCache] Replace skip_lock_node_ids with a segment lock protocol - #36848
Conversation
39de364 to
0c94f3e
Compare
|
/tag-and-rerun-ci |
|
/tag-and-rerun-ci |
0c94f3e to
b4ea300
Compare
|
/tag-and-rerun-ci |
b4ea300 to
b9177d4
Compare
|
/tag-and-rerun-ci |
b9177d4 to
ec9e36c
Compare
|
/tag-and-rerun-ci |
| skip_lock_node_ids: dict[ComponentType, set[int]] = dataclasses.field( | ||
| default_factory=dict | ||
| ) | ||
| mamba_lock_acquired: bool = False |
There was a problem hiding this comment.
the release still depends on which node the caller passes to dec_lock_ref, not the receipt alone, so passing the wrong node releases the wrong segment and silently leaks or steals another lock. can we store the anchor node id on the receipt and assert it at release so a mispair cannot happen?
|
|
||
| def inc_lock_ref( | ||
| self, node_id: NodeId, skip_lock_components: Sequence[ComponentType] = () | ||
| self, node_id: NodeId, lock_mamba: bool = True |
There was a problem hiding this comment.
this hardcodes MAMBA into the generic core (also dec_lock_ref, dec_swa_lock_only), breaking the component polymorphism the core is built on, since the core should drive components through their interface like eviction_priority and never name a concrete type. base kept it parametric via skip_lock_components with the caller passing (ComponentType.MAMBA,), so can we keep it parametric or use a component owned predicate like decode_lock_optional?
| recomputation (sanity_check) as the per-step oracle.""" | ||
|
|
||
| cfg = CacheConfig( | ||
| components=(ComponentType.FULL, ComponentType.SWA), |
There was a problem hiding this comment.
the fuzzer is FULL+SWA only, so the hardest paths (lock_mamba=False, the lower priority cascade release, host_uuid migration) never run, and both new suites skipTest on rust. can we add MAMBA here, run the oracle on rust, and validate end to end on a real FULL+SWA+MAMBA model like Inkling?
| if cd.lock_ref == 0: | ||
| key_len = len(cur.key) | ||
| if cd.lock_ref == 0 and cd.value is not None: | ||
| key_len = len(cd.value) |
There was a problem hiding this comment.
nit: maybe we can rename key_len to value_len
| self, | ||
| node_id: NodeId, | ||
| skip_lock_components: Sequence[ComponentType] = (), | ||
| lock_mamba: bool = True, |
There was a problem hiding this comment.
Now that inc_lock_ref and dec_lock_ref need to be aware of mamba's logic, doesn't that seem a bit unreasonable?
There was a problem hiding this comment.
Suggestion:
skip_lock_components=(ComponentType.MAMBA,)
| .insert(node_handle); | ||
| } | ||
| let component = Arc::clone(&self.components[i]); | ||
| if !lock_mamba && component.component_type() == MAMBA { |
There was a problem hiding this comment.
| } | ||
| let component = Arc::clone(&self.components[i]); | ||
| component.release_component_lock(self, node_id, params, /* lock_host = */ false); | ||
| if !mamba_lock_acquired && ct == MAMBA { |
There was a problem hiding this comment.
| self, | ||
| node_id: NodeId, | ||
| params: Optional[DecLockRefParams] = None, | ||
| params: DecLockRefParams, |
There was a problem hiding this comment.
Can we store the anchor in the receipt? Right now dec_lock_ref trusts the caller's node_id, so a valid receipt + wrong anchor can release another holder's lock without the > 0 assertions catching it.
| # boundary, but the prealloc lock and the receipt on the req still | ||
| # belong to pm.last_device_node; restore the pairing so any release | ||
| # before the commit hands over the restored lock hits the right node. | ||
| dr.req.last_node = pm.last_device_node |
There was a problem hiding this comment.
The need to restore req.last_node here seems to be another symptom that the lock receipt does not fully identify the acquired lock. Could we carry the acquisition anchor in the receipt/lock handle instead? Then rematching the request wouldn't be able to silently break the node/receipt pairing.
| hicache_restored_kv_indices: Optional[torch.Tensor] = None | ||
| hicache_restored_node: Any = None | ||
| # Receipt for the inc_lock_ref held on hicache_restored_node. | ||
| hicache_restored_lock_params: Optional[DecLockRefParams] = None |
There was a problem hiding this comment.
nit: maybe renamed to hicache_restore_lock_receipt?
| self.tree_cache.dec_lock_ref(prefix_match.last_device_node) | ||
| req = decode_req.req | ||
| restored_node = decode_req.hicache_restored_node | ||
| restored_params = decode_req.hicache_restored_lock_params |
There was a problem hiding this comment.
nit: maybe renamed to restored_lock_receipt?
f50f7bc to
d90e94f
Compare
|
Rebased onto current
Validation is in the updated description, including an end-to-end run of this protocol on Inkling-Small-NVFP4 (FULL+SWA+MAMBA, tp=4) on both cores with bit-exact KL and zero lock or ledger violations. |
|
CI triage for run 34173204012 (head d90e94f):
Will re-run the two failed shards once the run finishes. |
|
Re-ran the two failed shards (attempt 2): both green.
|
|
Would it make sense to also refactor |
Lock receipts carried `skip_lock_node_ids`: a per-component set of node ids that were tombstones at acquire time, replayed at release so a short-lived lock could not consume a later load-back or request lock once that tombstone became a valid device value. The set had to be built, carried and replayed correctly at every acquire/release pair, it grew with the tree, and a lost or defaulted receipt silently released another holder's lock. Replace it with a segment lock: a lock covers the contiguous node segment from the locked node up to the boundary stamped by the SWA window walk, counting every node in it (tombstones included). The receipt alone determines the release: - `IncLockRefResult` / `DecLockRefParams` drop `skip_lock_node_ids`; the SWA uuid already on the receipt marks the segment boundary, and the receipt carries the anchor node id. `dec_lock_ref`, `dec_swa_lock_only` and `dec_host_lock_ref` assert the receipt is replayed on the node its acquire returned, so a mispaired release fails loud instead of walking another holder's segment. `Req`, `SessionSlot` and `DecodeRequest` carry the receipt object itself (`req.lock_receipt`). - The core names no component: `inc_lock_ref(node_id, skip_lock_components)` records the skipped set on the receipt (`ComponentSet` bitmask in Rust); the decode hold passes `(MAMBA,)`. Meeting `lock_ref == 0` inside the segment is a hard assertion. - Values materialized under lock (load-back, SWA rebuild, a Mamba tombstone refilled by insert) are credited to protected directly; node splits copy lock refs and migrate the boundary uuid so an outstanding receipt still releases exactly. - Every component release refreshes the evictable-leaf sets of the nodes it unlocks, and the Rust core releases auxiliaries before Full, so a leaf whose last lock is an auxiliary one is readmitted whatever the order. - The receipt is required at every release seam, including the host unlock and the component trait, in Python and Rust. A non-positive SWA window is rejected by the adapter and the Rust constructor. - `SessionSlot` persists `swa_prefix_lock_released` so closing or aborting a streaming session after a first-turn early SWA release does not release the segment twice. - The Rust TreeCore (`rust/sglang-radix-tree`), its PyO3 bindings and `rust_tree_core/adapter.py` are ported to the same protocol. The segment-lock protocol and fuzz tests read the tree through the inspection interface and run under both backends; a FULL+SWA+MAMBA fuzz exercises the Mamba opt-out receipt and the cascade. Co-authored-by: Zhiqiang Xie <zqx@meta.com> Signed-off-by: Zhiqiang Xie <xiezhq.hermann@gmail.com>
d90e94f to
72ffb39
Compare
|
Rebased onto current What the rebase changed (only the six Rust files under
Local validation on the new head (B200 box,
|
|
CI on The remaining red checks are platform workflows, each reproduced on unrelated PRs today and none on a path this PR touches:
|
I think it's worth doing it but probably not in this PR, which has been fairly large already. |
* origin/main: (27 commits) [Simulator] Give the OFFLINE/BLOCKING comparison tolerances real headroom (sgl-project#38732) [Config] msgspec.Struct for the config tier (sgl-project#38753) [AMD] ci: move the miles nightlies from rocm700 to rocm10 (sgl-project#37495) [Config] One writer for the declaration stash; no exception to the write seal (sgl-project#38752) docker(xpu): drop redundant setvars.sh from torch_memory_saver RUN (sgl-project#38665) [XPU][Fix] Pack device-pointer tables as uint64 to avoid 64-bit address overflow (sgl-project#35051) [CI] Temporarily disable GB300 tests (sgl-project#38770) [diffusion] feat: spill large tensors over shared memory like numpy arrays (sgl-project#38656) [diffusion] refactor: refactor utility ownership and document helper placement (sgl-project#38699) [NPU]Support GLM5.2 and FP8 DSA&Indexer kvcache for 950 (sgl-project#38250) [CI] Answer unrecognized slash commands instead of skipping silently (sgl-project#38736) [AMD] Parallelize aiter spec-decode KV index building over token blocks (sgl-project#37659) [DSv4] Integrate TRT-LLM DSv4 Attention for SM100/103 (sgl-project#30805) Add Opt-In for GLM-5.3 Flash breakable prefill CUDA graphs (sgl-project#38522) [CI] Install helion 1.4.0 for the KDA Helion kernel tests (sgl-project#38688) [Rust] Gate health on startup warmup completion (sgl-project#37994) [HiCache] Replace skip_lock_node_ids with a segment lock protocol (sgl-project#36848) feat: add optimized Domino rollout to DFlash V2 (sgl-project#36899) [CI] Add /run-full-ci and /run-extra-ci slash commands (sgl-project#38734) [Model] Support GLM-5.3 Flash NVFP4 loading (sgl-project#38621) ...
Conflict in schedule_batch.py: main's sgl-project#36848 replaced Req.skip_lock_node_ids (and swa_uuid_for_lock) with a single lock_receipt: DecLockRefParams. Took main's removal and kept this branch's kv_rotation_base field beside it. Ported test_page_interleave_shard.py's _GraftReq stub to the same protocol (it set the two removed attrs; auto-merge left it stale since the file is new on this branch). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Motivation
Lock receipts carried
skip_lock_node_ids: a per-component set of node idsthat were tombstones at acquire time, replayed at release so a short-lived
lock could not consume a later load-back or request lock once that tombstone
became a valid device value.
That set had to be built, carried and replayed correctly at every
acquire/release pair, it grew with the tree, and every new call site had to
remember to thread it through. The release decision depended on state captured
somewhere else, and a lost or defaulted receipt silently released another
holder's lock.
Modifications
Replace the set with a segment lock. A lock covers the contiguous node
segment from the locked node up to the boundary stamped by the SWA window walk,
counting every node in it (tombstones included). The receipt alone determines
the release:
IncLockRefResult/DecLockRefParamsdropskip_lock_node_idsand carrythe anchor node id, the SWA boundary uuid (
None= root) and theskipped component set. Every release (
dec_lock_ref,dec_swa_lock_only,dec_host_lock_ref) asserts the receipt is replayed on the node its acquirereturned, so a mispaired release fails loud instead of walking another
holder's segment.
Req,SessionSlotandDecodeRequestcarry the receiptobject itself (
req.lock_receipt) instead of scattered fields.inc_lock_ref(node_id, skip_lock_components=())replaces the oldskip_lock_componentsand the interimlock_mamba: the core names nocomponent, it drives whatever the tree registered through the same
interface, and the decode hold passes
(ComponentType.MAMBA,). Meetinglock_ref == 0inside a segment is a hard assertion instead of a silent skip.unlocks, so a leaf whose last lock is an SWA or Mamba one is readmitted
whatever the release order; the receipt is required at every release seam,
including the host unlock and the component trait.
full_component,swa_component,mamba_component,unified_tree_coreandunified_tree_core_interfacemove onto the receipt; node splits copy lockrefs and migrate the boundary uuid so an outstanding receipt still releases
exactly (the anchor stays on the deeper half). A Mamba tombstone refilled by
insert while a lock counts it is credited protected, not evictable.
swa_radix_cache,schedule_batch,schedule_policy,scheduler_components/dynamic_chunk_sizer,disaggregation/decode,disaggregation/decode_hicache_mixin,session/streaming_sessionand thescripted-runtime
lock_ref_exhausterare updated at their call sites.SessionSlotpersistsswa_prefix_lock_releasedso closing or aborting astreaming session after a first-turn early SWA release does not release the
segment twice.
rust/sglang-radix-tree): the same protocol, portedfaithfully, with a
ComponentSetbitmask for the skipped set, the anchorassert, leaf-set refresh in the auxiliary release walks, auxiliaries released
before Full, a required receipt on the host unlock, and a positive-window
assert in the SWA constructor. The PyO3 bindings and
rust_tree_core/adapter.pyfollow the new interface (the adapter alsorejects a missing or non-positive
sliding_window_sizeup front), so theshared Python/Rust cache suite runs under both backends.
Review round
The PR is one squashed commit on current
main(13469c16d3); the reviewround below is folded into it. Addresses the comments from @ispobock and @alphabetc1: anchor on the receipt
(no more trusting the caller's node; the PD hicache restore re-pair now
asserts instead of silently mispairing), no hard-coded Mamba in the generic
core (parametric skip set instead),
TestSegmentLockProtocol/TestSegmentLockFuzzported to the inspection interface so they run under theRust core too, a new
TestSegmentLockFuzzWithMamba(FULL+SWA+MAMBA: theopt-out receipt, the lower-priority cascade on early SWA release, Mamba
evictions), plus the
key_len->value_lenandhicache_restore_lock_receiptrenames.Accuracy Tests
N/A — no change to model forward paths; this is lock bookkeeping.
Speed Tests and Profiling
Neutral: the removed set was allocated, populated and replayed per
acquire/release and scaled with the tree, whereas the receipt is two scalar
fields. Tree micro-benchmarks (
test_unified_radix_cache_bench.py --bench,insert/match/evict/lock_unlock/cache_finished across FULL, FULL+SWA,
FULL+SWA+MAMBA) are within noise of the baseline; lock_unlock is parity or
better in every configuration.
Source commits (internal tree, for provenance)
0497ae074f16d896398dd0649c01637e3dChecklist
Testing
Ran locally on the rebased head
72ffb39689(B200 devbox, PyTorch 2.11,Rust 1.92 per
rust-toolchain,SGLANG_ENABLE_ASYNC_ASSERT=true):test_unified_radix_cache_unittest: 1210 passed, 0 failed,1362 skipped (non-applicable component configs). The lock consumers touched by
the receipt move (
test_streaming_session_unit,test_decode_radix_lock_ref,test_swa_lock_release_lifecycle,test_mamba_donated_alloc_ratio,test_unified_radix_cache_bench,test_swa_unittest,test_swa_eviction_boundary,test_scheduler_chunked_req_gate,test_rust_tree_core,test_rust_extension): 223 passed;test_rust_tree_core_integration(incl. [Rust TreeCore] Harden runtime and CI parity #37303's stale-handle test with thereceipt fix): 105 passed.
test_rust_unified_radix_cache_unittest.py(production + inspection extensions built from this branch): 2572 ran, OK,
1368 skipped. The segment-lock protocol and fuzz classes now run under the
Rust core; the one case that drives Python component objects directly skips.
cargo testinrust/sglang-radix-tree: 839 passed, 0 failed;cargo clippy --all-targets -- -D warningsclean in all three CI shapes(default,
python-extension,inspection,--no-default-features --features tch/doc-only);cargo fmt --checkclean;ruff check --select F401,F821,UP037,ruff format --check,isort --check-onlyandcodespellclean on thechanged files.
Mamba refill credit, the auxiliary-release leaf refresh, and the receipt
anchor mismatch all fail without their fix on both backends.
End-to-end on a real FULL+SWA+MAMBA model (this protocol as landed internally,
thinkingmachines/Inkling-Small-NVFP4, tp=4 on 4x GB300, strict pool checks on):gsm8k 0.885 / 0.860 on the Python / Rust core (threshold 0.80), bit-exact KL
0.0 on single-turn, prefill cache hit, decode cache hit and HiCache multi-turn
branching on both cores, and a shared-prefix serving stress with HiCache
write_through, ~140 retractions and outputs past the SWA window: 128/128
requests on every pass, identical cached-token counts on both cores, zero lock
or ledger violations.
Original commits
86df8c33ba4968038abdc2a563c9c96a1db3e127CI States
Latest PR Test (Base): ✅ Run #34329336067
Latest PR Test (Extra): ❌ Run #34329335686
Latest PR Test (AMD ROCm 10): ❌ Run #34329335745