Repository navigation
Conversation
…ect linker (sgl-project#41343) Squash of sgl-project#41343 at 9417ac9 (5 commits by alphabetc1, base 3e13f1b), applied unchanged. The K3 / UMBP / DCP commits above run on sgl-project#41343's Mamba linker layer; drop this commit once sgl-project#41343 lands on main.
…on the tree page size, reject DCP on hybrid SWA Carried over unchanged from sgl-project#39511 (5ed7075, 456b7ce, 4557fecaff), rebased onto current main: - kv_cache_configurator: the paged allocator keeps the DCP-widened loc space (max_total * loc_span, pool_page_size * loc_span with loc_span = attn_dcp_size // loc_space_scale). Narrowing it strands 7/8 of the KV tensor; multiplying by attn_dcp_size unconditionally double-scales a replicated draft worker to dcp_size**2. - scheduler: PrefillAdder chunks against tree_cache.page_size, which is the widened page under DCP, so it never hands the tree a range that straddles a DCP page. - kv_cache_hook: reject --dcp-size > 1 on hybrid sliding-window models; their allocators do not widen the loc space and nothing downstream collapses it.
…ker, collapse DCP-widened KV rows Builds on the Mamba component support of the Mooncake direct linker (sgl-project#41343) and adds what that PR leaves out: hybrid linear models without SWA, the UMBP (mori) backend, and DCP. Carried over from sgl-project#39511. - linker_pool_assembler: _build_mamba_device_pool_group for _MambaStrategy (HybridLinearKVPool): the MLA KV entry plus a slot-addressed MAMBA entry packing every conv then every temporal state of a slot into one object. Not rank_replicated: the MLA KV is TP-replicated but the KDA state is not. Shares the Mamba rejections (unified memory, int8 checkpoints, side states) and the dense-row check with the FULL + SWA + MAMBA builder, and uses stage-local layer ids. - DevicePoolEntry takes dcp_size: token-addressed indices are DCP-widened locs, so a widened page collapses to page_size // dcp_size rows at loc // dcp_size. Without it the linker walks widened locs as rows and runs off the buffer. The DSA group rejects DCP instead of collapsing a replicated draft pool with the sharded target. index_mapper is dropped, nothing set it. - UMBPDirectLinker: chunk-level RPC/local calls run in parallel (UMBP_RPC_PARALLELISM), CHUNK_PAGES is tunable; load completion tracks the request ids. - UMBPStore: the MAMBA object-key branch dispatches on the host pool type instead of getattr defaults; keys are unchanged.
…cking (sgl-project#43567) (cherry picked from commit 0923bcd)
…ba UMBP group
A FULL + MAMBA pool group is not rank-replicated because the KDA state is
TP-sharded, so every object key carried a tp{rank} term and each attention
TP rank stored its own copy of the byte-identical MLA KV. At TP8 that is
eight copies of every KV page; with the per-rank Mamba states next to them
a 1160 GiB tier fills in about 25 minutes of the agentic trace, after which
local LRU eviction drops prefixes the device no longer holds.
Mark the group's KV pool as rank-replicated when DCP does not shard it, and
key only that pool without the tp term. Every rank still writes and the
store keeps one object. The Mamba and draft pools keep a keyspace per rank.
…linker group
With --speculative-algorithm DSPARK the linker refused to start on Kimi K3
("only supports packed draft KV caches"): the draft's MHA KV covers the
same token locs as the target KV, and a group that does not store it would
restore prefixes whose draft KV was never written.
For hybrid linear targets the draft plan is now PACKED under the linker,
and _build_mamba_device_pool_group adds the draft KV as its own DRAFT pool,
indexed by the target KV's transfer (indices_from_pool=KV). Draft layer d
rides on the d-th full-attention layer: the draft forward runs after the
target forward, and one buffer per mapped layer is what UMBP's layer-wise
loader requires. DCP is rejected, since the draft KV is TP- but not
DCP-sharded. Every other pool group still rejects non-packed drafts; the
HiCache path is unchanged.
Each rank probes the Rust toolchain on its own (`cargo --version --verbose`, `rustc -vV`, 10 s timeout). When rust-toolchain.toml pins a toolchain the image does not ship, rustup installs it on first use: ranks that probe while the install runs time out and fall back to Python, and a later rank gets Rust. Kimi K3 at TP8 came up with 7 Python and 1 Rust TreeCore in every default-config run on MI355X, and a Rust-only limitation then stopped startup on that one rank. Reduce the Rust decision over the cache's rank groups with MIN, so one rank's fallback applies to all of them. The reduction runs only when Rust was requested, which every rank derives from the same configuration.
sgl-project#43479) (cherry picked from commit 58cb1a2)
…amba pool A queued offload holds inc_lock_ref on its node until the store acknowledges it, and the MAMBA LRU skips locked nodes. Nothing bounds the queue, so once the store drains slower than prefill donates checkpoints, every mamba checkpoint in the tree ends up locked, evict_for_alloc frees nothing, and the next _alloc_mamba_slot asserts "Can not alloc mamba cache" on every rank. Measured on Kimi K3 at DCP8, conc=8, with a 120-slot pool: 117 slots locked by 117 pending offloads, the LRU walk starting from None, and the crash after 22 requests. offload_nodes now skips an offload that carries a mamba state while the queue holds three quarters of the pool. The checkpoint stays on the device and evictable; it is only not persisted, as when offload() declines. The check reads the queue length and the node's transfers, which every rank shares, so all ranks skip the same nodes.
UMBP's local eviction ranks objects by its own commit and read recency. With the direct linker the prefixes the device tree keeps hitting are never read back from UMBP, so they look coldest there and are evicted first; the next request that needs one after the device has dropped it recomputes the whole prefix. Measured on Kimi K3 TP8 with a 580 GiB DRAM tier: 46% of the evicted KV objects were probed again later, most within five minutes. UnifiedCacheLinkerWrapper.match() now offers each device hit to an optional linker hook, on_device_hit(); a linker without the method (Mooncake, ...) is not called. The UMBP direct linker implements it: at most once per UMBP_TOUCH_INTERVAL_S per device node it hashes the hit prefix the way _tail_hashes does, expands it to this rank's object keys with each pool's own LOOKUP keys -- every page for KV and draft KV, only the boundary page for Mamba, whose earlier checkpoints should keep aging -- and sends them to mori's batch_touch on a thread of its own. A touch moves the objects to the recent end of the store's eviction order and nothing else: no read, no lease, and exists stays side-effect free. Off by default (interval 0), and off when the mori client has no batch_touch. With UMBP_TOUCH_INTERVAL_S=30 on the 580 GiB tier the prefix hit rate went from 82.4% to 96.4% (TP8 without the KV-once change) and from 87.5% to 96.7% (with it), on par with HiCache; the hook costs under 1 ms of scheduler time per hint.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why
The agreed order is sgl-project#41343 and sgl-project#40896 first, then sgl-project#39511. Until they land, Kimi K3 / UMBP / DCP testing and fixes go on this branch. It stacks our work on sgl-project#41343 as-is, because sgl-project#41343's Mamba COW at COMMIT is what makes a Mamba state restored from UMBP correct.
6fc8d9da32. The fork'smainwas fast-forwarded to it so this PR lists only the stack. Upstream main has moved on by 111 commits; not rebased yet.9417ac9a9e, still current on 2026-10-12.Commits (bottom to top)
fe4ffebca94f5966a4473e058b7a3a76cb3232b84d5b023b28dc97ecc4a985d5657f92fix/tree-core-rank-consistent(30492de4c5)42e707b7eeWhat the fixes address
4d5b023b28). K3's FULL + Mamba group is not rank-replicated (the KDA state is TP-sharded), so every key carriedtp{rank}and TP8 stored each MLA KV page 8 times. fp8 KV is about 12.9 KB per token per rank; a Mamba state is 52.7 MB per rank. The 1160 GiB UMBP DRAM tier filled after about 25 minutes of the agentic trace, and its local LRU eviction then dropped prefixes that were still hot on the device, so prefix hits declined. Now only the KV pool's key drops thetpterm (pool_suffix_overrides={'kv': 'cp0_pp0'}in the TP8 startup log). Mamba and draft keys stay per rank, and DCP keeps per-rank KV.dc97ecc4a9). With--speculative-algorithm DSPARKthe linker refused to start on K3. The draft KV now rides along as its own DRAFT pool, indexed by the target KV's transfer. DSPARK with DCP is still rejected: the draft KV is TP-sharded but not DCP-sharded.85d5657f92). Each rank probed the Rust toolchain on its own with a 10 s timeout, and K3 TP8 came up with 7 Python and 1 Rust TreeCore in every default-config run. The Rust decision is now MIN-reduced over the cache's rank groups.76cb3232b8). When the ping-pong slot allocation failed, a lazy Mamba checkpoint could be overwritten in place, and a later hit restored the wrong state. extra_buffer and DSPARK verify both take this path.42e707b7ee). The 0.5 s wall-clock gate on refreshing waiting prefixes (from [Scheduler] Refresh waiting requests' cached prefixes under FCFS-like policies with LRU eviction sgl-project/sglang#41593, in our base) fires on different scheduling passes on different TP/CP ranks, so LRU order and evictions can diverge across ranks. Our configs take this path: defaultfcfspolicy, LRU eviction, and the unified cache reportssupports_fast_match_prefix() == False.Validation
Kimi K3, MI355X x8, InferenceX agentic trace, 3600 s per run, mori
fix6image (--ulimit memlock=-1,HSA_USERPTR_FOR_PAGED_MEM=0on the standalone server).TP8, conc=8, DSPARK block 3, UMBP:
HiCache
write_throughreaches 96.7% on the same trace.85d5657f92, the base image runs 8 Python TreeCores and therust192image 8 Rust ones, and the TP8 probes pass on both.dc97ecc4a9plus local diagnostic counters.42e707b7eehas unit tests only; no GPU nodes are available right now.Not included
5d95ba3c8a) still removesExternalLinkerLoadPhase.COMMIT, where [Feat] unified_cache: Support hybrid SWA + Mamba models in the Mooncake direct linker sgl-project/sglang#41343 copies the loaded Mamba state into the request's slot (_cow_node_state_into_req). Stacking the two means reworking [Feat] unified_cache: Support hybrid SWA + Mamba models in the Mooncake direct linker sgl-project/sglang#41343's Mamba load onto request-private slots and proving accuracy again. Without that copy, a probe that restored 17,536 tokens from UMBP deviated 8.7x the noise floor.pr39511-umbp-lookup-timeout(lookup deadline plus offload back-pressure) is experimental. It held up on TP8, but DCP8 never recovered from a ~27 min store freeze. With the mori fix, lookups take 1–5 ms (max 32 ms).3de4a7db24and541412ce53fromk3-hybrid-mamba-linker-rebased-20261008are not needed on the [Feat] unified_cache: Support hybrid SWA + Mamba models in the Mooncake direct linker sgl-project/sglang#41343 line. The UMBP lookup checks the Mamba state on the trailing page of every candidate length, and DCP8 conc=8 with a 120-slot Mamba pool ran 900 s with noCan not alloc mamba cache.*-diag*branches) are for experiments only.RecoverEvictedSWAsgl-project/sglang#43721 and [MemCache] Add shared component size and path APIs to the tree core interface sgl-project/sglang#43722 come with the next rebase.Known open issues
PageBackend) works on last write/read time, independent of the radix tree, andBatchContains(exists) does not refresh it. Shared prefixes that stay hot on the device therefore age out of UMBP. Storing KV once postpones the DRAM fill to about an hour but does not remove the mismatch; run 3's small decline may be this. Being investigated separately on the mori side.umbp_direct_linker.pyraises when an object is evicted between lookup and load, which stops every rank. [Unified Cache] Abort requests whose external linker load fails instead of crashing the scheduler sgl-project/sglang#40896 turns load failures into a per-request abort.Plan
fe4ffebca9and rebase onto main; the two cherry-picks fall away.3e058b7a3a) onto request-private load slots and re-run the probes.