[HiCache] MLA/DSA host dedup with rotating layer owners, chunked D2H backup issue and idle-sync skip (opt-in)

[nbommi19]

Draft for code review. Combines #41978, #41983 and #41984 into one HiCache PR, as suggested in review. Based on 4358a16 (main as of 2026-09-13), the commit it was developed and tested on; it does not merge cleanly into current main yet, and I will rebase once the approach is agreed.

Motivation

Three costs show up in long-context, cache-heavy HiCache serving with MLA/DSA models under TP:

• Duplicate host traffic and capacity. Without DP attention every TP rank holds the same MLA latent KV and DSA index-K, so HiCache backs up and loads back the same bytes on every rank and stores the same tokens once per rank. Profiling showed write-back and load-back limited by the aggregate host transfer path, with every rank copying identical data.
• Serialized backup issue. With direct I/O and page_first_direct, each backup is issued as one large batch of per-page copies; the host issues all copy descriptors before the first copy starts, so issue time and transfer time add up.
• An idle per-step collective. check_hicache_events runs a CPU all-reduce of the ready-ack counts on every scheduler step, even when no transfer is pending on any rank.

Modifications

Three commits, each off by default (with all of them unset, the code path is unchanged):

1. Chunked D2H backup issue. SGLANG_HICACHE_D2H_ISSUE_CHUNK_PAGES=n splits each all-layer backup batch (MLA KV and DSA indexer) into n-page batches on the same stream, in page order, so each batch's copy runs while the next one is issued. Same copies, same bytes, same stream order.

2. Idle-sync skip. SGLANG_ENABLE_HICACHE_IDLE_SYNC_SKIP=1 skips the ready-count all-reduce while no ack is outstanding and no storage backend is attached. All inputs to the decision are replicated across ranks (wrapped in rank_consensus), so all ranks skip or sync together; one sync is still forced after a fixed number of skips so the write_back duplicate-reclaim digest check keeps running. SGLANG_HICACHE_EVENTS_TIMER_STEPS=N optionally logs the time spent in check_hicache_events every N calls.

3. Rank-sharded host tier. SGLANG_ENABLE_HICACHE_RANK_SHARD=1 gives each host layer of the KV and indexer host pools one owner rank (round-robin over the TP ranks):

• Backup (D2H): a rank copies only the layers it owns into a host buffer that holds only those layers. Unique host capacity grows by about tp_size at the same --hicache-size (SGLANG_HICACHE_RANK_SHARD_KEEP_CAPACITY=1 keeps the old capacity and uses less host memory instead).
• Load-back (H2D): a rank copies only its own layers into its device slots; then, layer by layer, each owner gathers the loaded pages into a staging buffer and broadcasts them over a dedicated process group on the TP ranks, and the other ranks scatter them into their slots. Each layer's load event is recorded after that layer's exchange, and the load ack after the whole exchange.
• Forward gate: the shard group's broadcasts and the forward's collectives run on different communicators, so they must never run concurrently or in a different order per rank. A sharded load-back must go through start_loading (behind the load fence; anything else raises), and before the next forward's first kernel (Scheduler.run_batch and LayerDoneCounter.set_consumer) the forward stream waits for every submitted exchange to finish. Without this gate the exchange can deadlock against the forward's collectives.
• The radix tree, host-slot allocation and device allocation stay replicated; only the bytes are sharded. L3 page accessors assert the pool is unsharded.
• Active only with TP > 1 and no DP attention, attention CP or DCP, one PP stage, one node, direct I/O with page_first_direct, no storage backend, not PD decode, and a non-FP4 MLA pool; otherwise it logs why it is inactive. Extra knobs: SGLANG_HICACHE_RANK_SHARD_STAGING_MB, SGLANG_HICACHE_RANK_SHARD_LOG_EVERY, and SGLANG_HICACHE_RANK_SHARD_VERIFY_EVERY (periodic cross-rank check that all ranks loaded the same host pages).
• Off by default. Unset, no pool is sharded, no group is created and no gate is installed; the per-layer H2D loop moves unchanged into L2TransferEngine._load_layer.
  When both 1 and 3 are enabled, the chunked issue runs over the owner-only layers.

This PR carries only the synchronous load-back path. If the async load-back issue work (#38472) lands, rank sharding needs a few small hooks there, which I can add.

Accuracy Tests

• Commit 1: copies, bytes and stream order are unchanged, so cache contents are identical.
• Commit 2: the skipped all-reduce only happens when its result is zero on every rank, so no cache state changes. A multi-process simulation (8 ranks, many seeds) gave identical decisions and radix trees across ranks and with the skip on and off.
• Commit 3: pages are restored bit-exactly from the owner rank. A multi-process gloo simulation (8 and 4 ranks, with chunked backup unset and set) matched the full stacked implementation field for field, including the transfer-op trace, with zero page mismatches; negative controls (skipping the exchange, diverging host pages) are caught. With rank sharding off, every transfer call, byte count and final device content matches upstream (and matches commit 1 alone when chunking is set). A model of the queue and communicator ordering deadlocks without the forward gate and passes with it.
• The stacked implementation passed GSM8K with forced eviction and load-back. This combined branch still needs an end-to-end accuracy run, which I will do before marking it ready.

Speed Tests and Profiling

Rank sharding cuts each rank's host copy volume to about 1/tp_size and multiplies unique host capacity by about tp_size, so far fewer requests fall out of the host tier and re-prefill; chunked issue overlaps copy issue with transfer; the idle-sync skip removes a per-step CPU collective whenever HiCache is idle. The gains are largest where HiCache load-back and write-back dominate. Happy to share detailed profiles with reviewers.

Related: #26691 (host dedup; closed), #36800, #39565, #39903 (pipelines the same all-reduce; closed), #30511, #37562.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #37351639014
Latest PR Test (Extra): ❌ Run #37351638781
Latest PR Test (AMD ROCm 10): ❌ Run #37351639104

[YAMY1234]

Thanks for putting this together. Rotating the owner per layer is a real improvement over the rank-0-only dedup: it spreads D2H/H2D traffic across ranks and still gives ~tp_size x unique host capacity. A few things I'd like to see addressed before this moves out of draft (inline comments below).

Tests. The multi-process simulations from the description aren't in the PR. Could you add unit cases for owner rotation, shard layer mapping and host capacity sizing to the existing test_mla_host_dedup_primitives.py, plus the end-to-end accuracy run (GSM8K with forced eviction/load-back) on the combined branch?

[YAMY1234]

This is nearly line-for-line the same as MLAHostDedupBroadcaster._bcast_layer (index_select -> broadcast -> index_copy_), and that primitive on main already has two users (HiCache dedup and the Mooncake linker MLA dedup). Could you add a per-layer src to _bcast_layer / broadcast_loaded_layer and build rotating ownership on top of it, instead of a separate communicator and lifecycle? That would also pick up its DP-attention support (attn_tp_group), which this PR currently excludes.

[YAMY1234]

done_event is recorded after the last layer's exchange, so whenever a sharded load-back is pending the next forward starts only after all layers are loaded and broadcast; this removes HiCache's layer-wise load/compute overlap. Could you share profiles (TTFT and load-back time: sharding off vs. rank-0 dedup vs. this PR) showing the 1/tp H2D savings outweigh the lost overlap? The Mooncake linker dedup on main issues its broadcasts on the forward stream, which keeps ordering with the forward's collectives without a gate; worth checking whether that works here too.

[YAMY1234]

If device_indices is on the GPU, di.cpu() blocks the scheduler thread on every load-back, and the result is then copied back to the device. MLAHostDedupBroadcaster.prepare_broadcast computes the same page index on device; could this do the same?

[YAMY1234]

Rank shard alone adds five env vars, plus _Stats, _GateStats, a gloo verify group and the events timer, roughly half of the new file. I'd keep only the enable flag (and maybe staging size) upstream, and leave the debug/telemetry pieces in a local branch.

[YAMY1234]

Per .claude/rules/no-getattr-defensive.md, please avoid getattr here; enable_rank_shard_forward_gate always exists on the controller, so call it directly. Similarly, the module-level mutable _d2h_issue_chunks_logged and the dict-based _events_timer would be cleaner as instance state or a small struct.