[HiCache] Opt-in rank-sharded host tier for DSA under TP

[nbommi19]

Draft for code review. This branch is 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; I will rebase once the approach is agreed.

Motivation

Under TP without DP attention, every TP rank holds the same MLA latent KV and DSA index-K. HiCache therefore backs up and loads back the same bytes on every rank, and the host tier stores the same tokens once per rank. Profiling of long-context, cache-heavy serving showed write-back and load-back limited by the aggregate host transfer path, with every rank copying identical data, and the host tier filling up with duplicates.

Modifications

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.

This branch 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

Pages are restored bit-exactly from the owner rank. A multi-process gloo simulation (8 and 4 ranks) 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 the flag off, every transfer call, byte count and final device content matches upstream. A model of the queue and communicator ordering deadlocks without the forward gate and passes with it. The stacked implementation passed GSM8K under forced load-back; this standalone branch still needs an end-to-end accuracy run, which I will do before marking it ready.

Speed Tests and Profiling

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. The gain is largest where HiCache load-back and write-back dominate. Happy to share detailed profiles with reviewers.

Related: #26691 (host dedup; closed), #36800, #39565.

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): ❌ Missing run-ci label -- add it to run CI tests.
Latest PR Test (Extra): ❌ Blocked -- run-ci is required first.
Latest PR Test (AMD ROCm 10): ➖ No AMD PR run found for this commit.

[nbommi19]

Combined into #42009 (one HiCache PR), as suggested in review.