Layerwise Split KV Cache Sharding

Commit: 494f293b5629 · PR: #2260 · Upstream: xLLM
Adaptation: Iluvatar BI-V100 (PCIe topology)
LOC: +455 −10 across 19 files

Problem

For models with heterogeneous layer structures (e.g., DeepSeek-V3 with dense attention layers interleaved with MoE layers), the KV cache is uniformly sharded across all tensor-parallel ranks. Each rank stores KV for all layers, even though different layers may have vastly different head counts.

On Iluvatar BI-V100 (32 GB HBM per card), this wastes memory on ranks that serve layers with fewer KV heads and prevents optimal utilisation of each card's HBM.

Solution

Introduce layerwise split KV cache sharding: a new KV cache layout strategy where each layer's KV cache can be sharded independently across a configurable subset of TP ranks.

Key Components

#	Component	Files	Description
1	`LayerwiseSplitLayout`	`core/framework/kv_cache/layerwise_split_layout.h`	Per-layer KV shard mappings. Dense layers spread across all TP ranks; MoE layers concentrate on fewer ranks.
2	Layerwise allocation	`core/framework/kv_cache/kv_cache_layerwise.{h,cpp}`	`allocate_kv_caches_layerwise()` — allocates per-layer shard sizes from layout. Handles the ILU/MLU transposed cache layout `[n_blocks, n_heads, block_size, head_dim]`.
3	Memory estimation	`core/framework/kv_cache/kv_cache_estimation_layerwise.{h,cpp}`	Reports peak/average per-rank memory; computes savings vs uniform.
4	ILU topology mapping	`core/framework/parallel_state/mapping_ilu.{h,cpp}`	PCIe-aware assignment: MoE layers placed on ranks sharing a PCIe switch to maximise intra-group bandwidth.
5	Engine integration	`core/distributed_runtime/layerwise_split_{engine_ext,master}.{h,cpp}`	Master computes layout at startup; engines propagate to workers.
6	Worker init	`core/runtime/worker_layerwise_init.{h,cpp}`	Workers receive and apply per-layer KV shard assignments.
7	Config flag	`core/config/parallel_config_layerwise.{h,cpp}`	`--enable_layerwise_split` gflag (default: false).

Iluvatar BI-V100 Hardware Context (verified via ixsmi + debug_warpsize.py)

4× BI-V100, Bus-Id 4B:00.0 – 4E:00.0, NUMA node 1
Warp size: 64 (NOT 32 — verified via CUDA kernel warpSize builtin)
32768 MiB HBM per card, 1500 MHz SM clock, 1200 MHz mem clock
Flat PIX topology — all pairs connected via single PCIe bridge (equal BW)
IX-ML 3.2.3, Driver 3.2.1, CUDA 10.2 (CoreX)
CoreX SDK at /usr/local/corex/
NCCL for collective communication (same process group as CUDA)
KV cache tensor layout (ILU): [n_blocks, n_heads, block_size, head_dim] (axis 1 = heads)
All verified constants centralized in core/config/ilu_hw_constants.h

Usage

# Enable layerwise split KV cache
./xllm_server --model deepseek-v3 --enable_layerwise_split=true

# Disable (default — uniform sharding, no regression)
./xllm_server --model deepseek-v3 --enable_layerwise_split=false

Test Plan

ID	Level	Description	Criteria
TC-01	L1	Allocation correctness	Per-rank allocation matches layout; unassigned layers get zero KV; total equals sum
TC-02	L1	Memory estimation accuracy	Layerwise peak ≤ uniform; estimation error ≤ 5%
TC-03	L1	ILU PCIe topology mapping	All layers assigned; MoE layers on same-switch ranks; no oversubscription
TC-04	L1	Engine layout propagation	All 8 workers receive consistent layout; full layer coverage
TC-05	L1	Worker KV shard application	KV populated for assigned layers; zero for unassigned; ASAN clean
TC-06	L2	Fallback when disabled	Uniform allocation identical to pre-feature behaviour
TC-07	L2	Speculative engine	No crash; KV correctly partitioned per model

File Summary

core/
├── CMakeLists.txt
├── config/
│   ├── ilu_hw_constants.h
│   ├── parallel_config_layerwise.cpp
│   └── parallel_config_layerwise.h
├── distributed_runtime/
│   ├── layerwise_split_engine_ext.cpp
│   ├── layerwise_split_engine_ext.h
│   ├── layerwise_split_master.cpp
│   └── layerwise_split_master.h
├── framework/
│   ├── kv_cache/
│   │   ├── kv_cache_estimation_layerwise.cpp
│   │   ├── kv_cache_estimation_layerwise.h
│   │   ├── kv_cache_layerwise.cpp
│   │   ├── kv_cache_layerwise.h
│   │   └── layerwise_split_layout.h
│   └── parallel_state/
│       ├── mapping_ilu.cpp
│       └── mapping_ilu.h
└── runtime/
    ├── worker_layerwise_init.cpp
    └── worker_layerwise_init.h
docs/
└── LAYERWISE_SPLIT_KV_CACHE.md
tests/core/
└── test_layerwise_split_kv_cache.cpp

4.8 KiB Raw Permalink Blame History Unescape Escape