TB-Scale LLM Inference: Zero-Redundancy Weight Deployment & 80% KV Cache Compression
Baidu Baige presents extreme VRAM optimization for TB-scale model inference: zero-redundancy weight deployment via DP Attn-sharding (releasing 16–30 GiB/card) and full-route CPP (50% throughput gain), plus KV Cache compression via LU-KV (80% reduction, 0.52% accuracy loss), Draft SWA, and PD handshake optimization (20% VRAM release).
Memory Pressure in TB-Scale Flagship Models
Using Kimi-K3 (2.8T parameters) as a representative flagship model, the speaker breaks down VRAM composition: ~74% for model weights, ~18% for KV Cache, leaving only ~8% for activations and runtime overhead. Parameter counts continue growing (GLM-5.2 744B → DeepSeek-V4-Pro 1.6T → Kimi-K3 2.8T, potentially 5–10T), while activated parameters per token grow slowly (40B → 100B), creating a widening gap between storage demand and effective compute. Long-context Agent workloads further increase dynamic KV Cache pressure; at 60K context, single-request KV Cache exceeds 3 GB for GLM-5.2 and Kimi-K3, while DeepSeek-V4-Pro compresses to ~300 MB.
VRAM capacity becomes the primary bottleneck limiting batch size (BS) and thus compute utilization. For Kimi-K3 under long-context, VRAM constraints cap BS at ~40 (16.8k tokens/s), whereas unconstrained BS could reach 256+ (38.6k tokens/s) — actual throughput only 44% of ideal. The goal is not merely fitting the model, but freeing VRAM for larger BS and higher throughput.
Weight Optimization: Zero-Redundancy Distributed Deployment
Root Cause: Multi-DP Redundancy in MoE Large-EP Deployments
Flagship models adopt MoE sparse architectures. Large Expert Parallelism (EP) combined with multi-Data Parallelism (DP) is a standard throughput optimization: each DP replica holds a full copy of Attention weights, enabling independent request processing, larger effective BS for Group GEMM, and balanced per-GPU BS reducing activation memory. However, Attention weights are replicated across DP replicas, creating significant VRAM redundancy that limits BS scaling.
Decode Path: DP Attn-Sharding
Retain multi-DP throughput benefits but shard Attention weights evenly across GPUs (each GPU holds 1/N). During computation, AllGather temporarily reassembles full weights. This eliminates Attention weight redundancy while preserving multi-DP advantages. On DeepSeek-V4-Pro, each GPU releases ~16.57 GiB; on Kimi-K3, ~29.8 GiB.
Key Implementation Optimizations
Weight rearrangement & communication merging: Discrete sharded weights reorganized into contiguous storage; multiple AllGather ops merged into one per layer, reducing communication count from N to 1 and cutting single-layer communication latency ~20%.
Double-buffer pipelining: While one buffer supplies weights for current-layer compute, the other receives next-layer weights, overlapping communication with compute. Extra VRAM ~4% of released memory, reusable across layers.
Compute-communication overlap enhancements: (1) Dedicated Copy Engine offloads data movement from SMs; (2) Cross-layer prefetch scheduling launches AllGather for layer N+1 during layer N compute; (3) FP8 precision for communicated weights halves data volume. Combined, these hide communication overhead with no measurable end-to-end performance loss.
Prefill Path: Single-DP Convergence + Full-Route CPP (Chunked Pipeline Parallelism)
Merge multiple DP replicas into one, eliminating DP-level weight redundancy at the architecture level. Then deploy model weights evenly across pipeline stages via full-route CPP (each PP stage occupies one GPU exclusively). Advantages: removes TP and inter-DP communication; near-uniform weight sharding; flexible linear scaling with node count to adapt to model size and GPU memory.
Experimental results: TTFT decreases near-linearly with node count. Main challenge is pipeline bubble reduction. Optimizations across three phases yield ~50% throughput improvement and ~60% TTFT reduction vs. community baseline.
Pre-processing: Centralized Async Consensus
Original ring-based synchronous consensus latency grew linearly with stage count (~1/3 of total bubble). Replaced with centralized async consensus: all stages report readiness in parallel to PP0 via background threads, reducing consensus time to ~one communication round, nearly eliminating pre-processing bubble.
Execution: FLOPs-Based Layer Partitioning & Dynamic Chunking
Heterogeneous layer structures (e.g., alternating KDA and MLA layers) make layer-count partitioning uneven in compute. Solution: partition by measured FLOPs per layer, so stages may hold different layer counts but balanced compute. Inside each stage, replace static quadratic-model-based Dynamic Chunking with runtime dynamic estimation considering request count, prefix cache hits, and model structure, enabling uniform micro-chunk compute load. Combined, execution bubble reduced >70%.
Post-processing: Decouple Token Output & KV Cache Transfer
Original SGLang PP forced serialized KV Cache sends across stages (waiting for prior stage send completion), creating ~1/3 of end-to-end bubble. New design decouples token output dependency from KV transfer: only PP0 and last stage retain minimal dependency for token output; other stages send KV Cache immediately after macro-chunk compute. Tail overhead equals one forward pass, diminishing proportionally with more stages.
KV Cache Optimization: Algorithmic, Post-Training, Runtime
Algorithmic: LU-KV (Long-Horizon Utility KV)
Evolved from ESS (arXiv:2512.10576) which managed cache at system level. LU-KV shifts to algorithmic layer: evaluates long-horizon utility (LU) of each token and attention head, globally planning VRAM budget across layers/heads under a given memory budget. Three insights: (1) Instant/local scores are noisy; (2) Head-level importance varies per token; (3) Prefill token value should be assessed from decode-generation feedback. LU metric generalizes across datasets.
Formulated as NP-hard non-convex discrete optimization; solved via convex relaxation and greedy approximation. System uses offline profiling (collect data, compute budget configs for various VRAM budgets) and online lookup (direct config retrieval, zero runtime overhead). On LongBench, accuracy loss only 0.52%, KV Cache compressed ~80%. Flexible across VRAM configs/compression rates, dataset-agnostic, compatible with Top-K etc. Accepted to ICML 2026; open-sourced at https://github.com/baidu-baige/LU-KV.
Post-Training: Draft Model Sliding Window Attention (SWA)
Speculative decoding uses a small Draft model generating candidates verified by Target model. Both maintain KV Cache. Draft model Attention redesigned to fixed-window SWA (window=128), making Draft KV Cache constant w.r.t. context length. In experiments, single-machine VRAM saved ~100 GB. Post-training adaptation preserves acceptance rate, no speculative gain loss.
Runtime: SGLang PD Handshake Optimization
In Prefill-Decode separation, Decode (D) pre-reserves VRAM to overlap KV transfer from Prefill (P) with decode compute. This reserved VRAM sits idle before transfer, limiting BS. Optimization: shift handshake/transfer to later timing, avoiding long pre-reservation. Decode VRAM usable for compute most of its lifecycle. Result: ~20% overall VRAM release, D-side BS reaches ideal configuration level.
Summary
Weight side: Decode uses DP Attn-sharding (16–30 GiB/GPU released); Prefill uses single-DP + full-route CPP (throughput +50%, TTFT -60%). KV Cache side: LU-KV compresses ~80% with 0.52% accuracy loss; Draft SWA saves ~100 GB; PD handshake optimization releases ~20% VRAM. Collectively, these techniques convert saved VRAM into larger batch sizes and higher throughput for TB-scale inference.
Signed-in readers can open the original source through BestHub's protected redirect.
This article has been distilled and summarized from source material, then republished for learning and reference. If you believe it infringes your rights, please contactand we will review it promptly.
Baidu Intelligent Cloud Tech Hub
We share the cloud tech topics you care about. Feel free to leave a message and tell us what you'd like to learn.
How this landed with the community
Was this worth your time?
0 Comments
Thoughtful readers leave field notes, pushback, and hard-won operational detail here.
