[KongchangAI]
· 2 min read· 1,438 words

BasisKV: How Low-Rank KV Cache Compression Optimizes Both Tensor Parallel Communication and Key Page Routing

BasisKV: How Low-Rank KV Cache Compression Optimizes Both Tensor Parallel Communication and Key Page Routing

BasisKV binds low-rank Value compression to the output projection, cutting both KV cache and tensor parallel communication, and uses Value latents to assist Key routing.

BasisKV is a new paper addressing KV cache bloat in distributed long-context LLM inference. It notes that prior low-rank KV compression mostly optimized compression ratios in single-GPU settings, ignoring tensor parallel communication overhead. BasisKV binds low-rank Value representations to the output projection so the same factors both compress Value cache width and reduce TP communication, and jointly optimizes factors across GPU partitions against the complete attention output. It also reuses the compressed Value cache to predict Key page routing, reading back original Keys only for selected pages to compute exact QK scores. The paper reports up to 1.42x end-to-end speedup under 8-GPU tensor parallelism, with code open-sourced.

Long-context LLM inference is under mounting memory and bandwidth pressure from KV cache bloat. Low-rank K/V projection is seen as an effective way to compress the KV cache, but prior work has mostly optimized compression ratios in single-GPU or non-distributed settings, overlooking the communication overhead that arises when attention computation is split across multiple GPUs under tensor parallelism (TP) deployment. A new paper, BasisKV, targets exactly this gap: it binds low-rank Value representations to the output projection, so that the same set of low-rank factors both compresses the Value cache width and reduces tensor parallel communication volume; at the same time, it reuses the compressed Value cache to predict part of the information needed for Key routing, reading back original Keys only for the selected Key pages to compute exact QK scores.

BasisKV paper source screenshot

The Overlooked Communication Cost in Distributed Inference

The paper makes two key observations. First, when low-rank KV cache compression is deployed in distributed inference environments, tensor parallel communication itself is a non-negligible cost. In TP mode, attention heads are partitioned across different GPUs, and each card holds only the low-rank factors for its portion of the attention partition. If each partition independently fits its own low-rank representation, the gains from compression are partially offset by cross-GPU communication.

Second, Value compression can be tied to the output projection. Since the attention output ultimately passes through the output projection matrix, jointly designing the low-rank Value representation with the output projection allows the same set of low-rank factors to serve two goals at once: narrowing the Value cache width and reducing the amount of data that must be exchanged in tensor parallelism. This coupling is the core design that distinguishes BasisKV from previous low-rank compression schemes.

Tensor Parallelism (TP) is a model parallelism strategy that splits the parameters and computation of a single Transformer layer across multiple GPUs. Taking the attention layer as an example, different attention heads (or dimensions within the same head) of multi-head attention are assigned to different devices, each card computes the output for only part of the heads, and the results are aggregated through collective communication (such as All-Reduce). The advantage of TP is that parallelism happens within a single layer, making it suitable for single-machine multi-GPU setups; the cost is that every forward and backward pass of each layer requires cross-GPU communication, and the communication volume is proportional to the size of the activation tensors along the split dimension. In long-context inference, the KV cache and attention output tensors grow with sequence length, so TP communication overhead rises significantly and becomes a non-negligible bottleneck. BasisKV directly reduces the size of the activation tensors that must be exchanged across GPUs by compressing the Value cache width, thereby cutting TP communication volume while compressing the cache.

Joint Optimization Rather Than Per-Partition Fitting

BasisKV's optimization strategy emphasizes global consistency. Existing low-rank methods typically let each GPU partition separately fit its own attention output, but per-partition optimality does not equal optimality of the overall attention output. BasisKV instead jointly optimizes the low-rank factors across GPU attention partitions against the complete attention output, making the compressed representation closer to the original attention result at the global level.

This joint optimization is especially important for long-context scenarios: the longer the context, the larger the KV cache, and the more pronounced the error accumulation of low-rank approximation becomes; per-partition fitting easily amplifies deviations during cross-GPU aggregation.

Low-Rank Approximation is built on the idea of approximating an original large matrix by the product of two smaller matrices (such as U and V), thereby representing the main information of the original matrix with far fewer parameters. In KV cache compression, the Key or Value cache matrix is typically decomposed into low-rank factors, and only the low-rank factors are stored at inference time rather than the full cache. The quality of low-rank approximation depends on whether the factors capture the dominant energy directions of the original matrix; when the matrix is split by rows or columns across devices, the factors obtained from independent SVD or fitting on each partition are only locally optimal, and after cross-partition aggregation they may deviate from the global optimum. BasisKV's emphasis on jointly optimizing low-rank factors across partitions against the complete attention output essentially upgrades low-rank decomposition from "independent per block" to "global coordination," reducing error amplification during distributed aggregation.

Using Value Latents to Assist Key Routing

A classic challenge in KV cache compression is this: once Keys are compressed, how do you accurately determine which Key pages are worth accessing? BasisKV's approach reuses the compressed Value cache as a routing signal. The low-rank Value representation itself carries some Key-related information, so it can be used to predict part of what Key routing needs; for Key information that cannot be recovered from Value, a Key-specific residual is stored.

At inference time, the system reads back original Keys only for the Key pages selected by routing and computes exact QK scores. This avoids the cost of reading back all Keys while preserving the accuracy of exact scoring. The Value cache plays a dual role here: it is both the compressed attention input and the predictor for Key page selection.

Key Routing refers to first predicting, at low cost, which Key pages are most relevant to the current Query in sparse attention or cache retrieval scenarios, and then computing exact QK dot products only for those pages. The difficulty of routing lies in the fact that the prediction signal must be cheap enough and reasonably accurate: if full Keys are used for routing, the point of sparsification is lost; if the prediction is too coarse, important Keys are missed and quality degrades. BasisKV uses the compressed low-rank Value representation as the routing signal because Value and Key share interaction with the same set of Queries in attention computation, and the Value latent indirectly encodes some Key-related semantic information. For Key-specific information that Value cannot cover, a Key-specific residual is stored, and at inference time original Keys are read back only for the pages selected by routing for exact scoring. This two-stage design of "coarse routing + exact scoring" preserves QK score accuracy while avoiding the cost of reading back all Keys.

Quality-Efficiency Trade-off Performance

The paper evaluates BasisKV across multiple LLM families and system settings, reporting that its quality-efficiency trade-off outperforms existing low-rank baselines and achieves up to 1.42x end-to-end speedup under 8-GPU tensor parallel inference. This figure comes from the paper abstract; specific test configurations and baseline details should be checked against the original paper and open-source code.

From a design standpoint, BasisKV's speedup comes from more than one source: low-rank Value compression reduces cache width and TP communication volume, while Key routing reduces read-back of original Keys; only the combination of the two produces end-to-end gains. This also explains why low-rank schemes that optimize compression ratio alone without considering communication cost struggle to achieve comparable speedups in distributed deployments.

Implications for Long-Context Inference Deployment

BasisKV's value lies in turning KV cache compression from a single-point optimization into a system-level co-design. As inference clusters commonly adopt tensor parallelism to serve long-context models, compression algorithms should be evaluated not only on cache footprint or perplexity loss but also with communication overhead taken into account. Binding Value compression to the output projection and using Value latents to assist Key routing offer two reusable directions for future research.

The code is open-sourced on GitHub (alexz949/BasisKV), and interested readers can reproduce its behavior under different tensor parallel scales.

One-line Summary

BasisKV binds low-rank Value compression to the output projection, cutting both KV cache and tensor parallel communication, and uses Value latents to assist Key routing.

Paragraph Summary

BasisKV is a new paper addressing KV cache bloat in distributed long-context LLM inference. It points out that prior low-rank KV compression mostly optimized compression ratios in single-GPU settings, ignoring the communication overhead that arises when attention computation is split across multiple GPUs under tensor parallel deployment. BasisKV has two core designs: first, it binds low-rank Value representations to the output projection, so the same set of low-rank factors both compresses Value cache width and reduces TP communication volume; second, it jointly optimizes low-rank factors across GPU partitions against the complete attention output, avoiding the global bias of per-partition fitting. For Key routing, it reuses the compressed Value cache to predict Key page selection and reads back original Keys only for the selected pages to compute exact QK scores. The paper reports up to 1.42x end-to-end speedup under 8-GPU tensor parallelism, with a quality-efficiency trade-off superior to existing low-rank baselines, and the code is open-sourced.

#BasisKV#KV cache compression#low-rank approximation#tensor parallelism#long-context LLM inference#Key routing
Share:

Related articles