Unified Full KV Cache Layout in L3

[Shunkangz]

Motivation

HiCache's L3 tier names every stored KV object by who wrote it: each backend appends its own rank/topology suffix (mooncake {hash}_{tp_rank}_{pp_rank}_k, file {hash}_{model}_{tp_rank}_{tp_size}_{pp_size}_{pp_rank}_cp{r}_{s}.bin, nixl/eic/hf3fs/simm/umbp each different again). This has two costs:

No cross-topology reuse, by construction. A page written by one parallel configuration is invisible to any other, even when the reader needs exactly the bytes that are already stored (MLA KV is bit-identical on every TP rank; a TP4 reader's head shard is exactly half of a TP8 writer's two objects).

Modifications

New Server Args

--hicache-storage-key-scheme unified switches L3 keys from rank-suffixed names to a topology-free coordinate:

{page_hash}_{namespace_digest}_L{start_layer}-{end_layer}[_H{head_group_index}]_{k|v}

Field	Meaning
page_hash	The existing chained SHA-256 over token ids — already topology-free.
namespace_digest	Digest of a KVCacheNamespace descriptor: model id, logical KV dtype, page size, head grid, layer grid, object byte order, optional numerics id. Any mismatch lands in a disjoint keyspace, so deployments miss instead of colliding.
L{start}-{end}	The rank's absolute layer range. Any PP partition attaches, uneven stages included (e.g. 61-layer models); different partitions derive disjoint keys.
H{j}	The canonical kv-head-group index from the descriptor. Omitted for MLA-family pools, which replicate KV across ranks and have no head axis.
`{k	v}`

Use --hicache-storage-backend-extra-config to partition a page into finer-grained chunks. The example below makes 1 kv head x 6 layers the smallest independently stored unit:

--hicache-storage-backend-extra-config '{"head_group": 1, "layer_partition": 6}'

Both entries are fleet-wide agreements, not local tuning options: they fix chunk boundaries, and boundaries are part of the namespace identity. Every member of a fleet that should share cache must pass the same values. Pick head_group = total_kv_heads / lcm(the fleet's attn-TP sizes).

Layout Support

The L2 host memory pool supports four layouts:

layer_first:        (2, layer, page * page_size, head, dim)
page_head:          (2, page, head, page_size, layer, dim)
page_first:         (2, page * page_size, layer, head, dim)
page_first_direct:  (2, page, layer, page_size, head, dim)

With the unified scheme enabled we serve page_first_direct only, overriding --hicache-mem-layout. It is the layout friendliest to both head and layer partitioning: a chunk is one contiguous run, so the transport moves it with a single descriptor instead of one per (layer, token). With a head partition the page block becomes head-group-major:

page_first_direct_group: (2, page, head_group, layer, page_size, head_in_group, dim)

A layer partition alone leaves the data naturally contiguous, so the copy engine still works. A head partition fragments it, which the copy engine handles poorly, so the kernel path is selected automatically. To avoid a host-side repack we added a transfer kernel that converts KV cache between GPU and CPU directly.

Follow Up

• This PR supports the Mooncake store only, since it stores multiple keys per page natively. HiCache File support is a follow-up.
• Extend the unified layout to Mamba state.

Accuracy Tests

We do a small E2E test based on Qwen3-8B-FP8 model with 36 layers and 8 KV heads and compare the accuracy on GSM8K dataset.
The launch command

python3 -m sglang.launch_server \
    --model-path /model \
    --host 0.0.0.0 \
    --port 32140 \
    --attention-backend fa4 \
    --page-size 128 \
    --fp8-gemm-backend flashinfer_trtllm \
    --enable-hierarchical-cache \
    --hicache-write-policy write_through \
    --hicache-io-backend kernel \
    --hicache-mem-layout page_first \
    --hicache-size 64 \
    --hicache-storage-backend mooncake \
    --hicache-storage-key-scheme unified \
    --hicache-storage-backend-extra-config \
      '{"head_group":1,"layer_partition":6,"prefetch_threshold":1}' \
    --hicache-storage-prefetch-policy wait_complete

With unified kv cache layout, we can achieve the same 94.5% accuracy as origin.

Speed Tests and Profiling

Measured on one B200 over PCIe Gen5 x16 — 61 layers, head_dim 128, page_size 64, 2048 tokens (~488 MiB per direction at H=8), bf16. H is kv heads, HG head groups, and run B the resulting contiguous run per pass:

H	HG	run B	H2D @2 CTA	D2H @2 CTA	H2D @32 CTA	D2H @32 CTA
8	1	2048	13.6	34.2	49.9	52.4
8	2	1024	13.6	33.4	49.7	52.4
8	4	512	13.4	31.6	49.5	52.5
8	8	256	12.8	29.4	49.8	52.4
16	1	4096	13.9	35.2	50.6	52.5
16	8	512	13.4	31.7	50.4	52.5
16	16	256	12.8	29.8	50.5	52.5

All figures GB/s. Cutting the head axis costs little — at 32 CTAs the run size falls 8x with no measurable change, and even at 2 CTAs the loss is under 15%. What sets the bandwidth is the CTA count, which block_quota caps. We keep the default at 2, matching the existing kvcacheio kernels, so KV copies contend minimally with model forward for SMs; the cost is that they run well below link speed. Raise block_quota per call site to trade SM occupancy for bandwidth.

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.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ✅ Run #34124794944
Latest PR Test (Extra): ❌ Run #34124794759
Latest PR Test (AMD ROCm 7.2): ❌ Run #34124794947

[Shunkangz]

/tag-and-rerun-ci

[Shunkangz]

/tag-and-rerun-ci

[Shunkangz]

/tag-and-rerun-ci

[Shunkangz]

/rerun-failed-ci

[Shunkangz]

/rerun-failed-ci

[Shunkangz]

/rerun-failed-ci

[Shunkangz]

/rerun-failed-ci

[Shunkangz]

/rerun-failed-ci