tensor-parallel-kv-cache-bench

Analytical benchmark connecting tensor parallelism and KV cache management in LLM serving.

With TP=4, each GPU holds 1/4 of the attention heads — and the KV cache is sharded. But prefix sharing, eviction, and disaggregation now require coordination across GPUs.

How does tensor parallelism change the layout, size, and transfer cost of the KV cache, and when does it help vs hurt?

---

Why This Exists

The portfolio has comm-cost-modeling for parallelism and five KV cache projects, but never connected them. This benchmark closes that gap:

• comm-cost-modeling: alpha-beta model for NVLink and InfiniBand
• kv-cache-tiering-bench: PCIe bandwidth 3.1 GB/s measured
• disaggregated-prefill-decode-sim: KV transfer cost between nodes
• radix-attention-sim: prefix sharing depends on KV co-location
• tensor-parallel-kv-cache-bench: KV sharding under TP

---

Key Results

On NVLink, all-gather is always cheaper than recompute by 1000x

Model	All-gather TP=4	Recompute	Ratio
llama3_70b	2.10 ms	8,602 ms	0.0002
llama3_8b	0.84 ms	1,147 ms	0.0007
qwen2_7b	0.37 ms	983 ms	0.0004
qwen2_0.5b	0.08 ms	164 ms	0.0005

There is no crossover point. All-gather always dominates on NVLink.

PCIe is a hard blocker for TP KV disaggregation

On PCIe gen4, shared bus contention makes total TP cost exceed TP=1 at all degrees. TP is only viable for KV disaggregation on NVLink or InfiniBand.

GQA models have a natural TP ceiling

Model	n_kv_heads	Max safe TP	Reason
llama3_70b	8	8	clean shard at all tested TP
llama3_8b	8	8	clean shard at all tested TP
qwen2_7b	4	4	TP=8 is degenerate
qwen2_0.5b	2	2	TP=4 is degenerate

Beyond n_kv_heads, some GPUs hold no KV heads and prefix sharing degrades severely.

Operational decision matrix (seq=4096)

Configuration	Recommendation
Any model, NVLink, TP=2	yes
llama3_70b, NVLink, TP=4	yes-with-note (moderate prefix degradation)
qwen2_7b, NVLink, TP=4	yes-with-note
qwen2_7b, NVLink, TP=8	no (degenerate GQA)
qwen2_0.5b, NVLink, TP=4	no (degenerate GQA)
Any model, PCIe, TP>1	no (disagg cost exceeds TP=1)

---

Models and Interconnects

Model	n_kv_heads	Layers	head_dim
qwen2_0.5b	2	24	64
qwen2_7b	4	28	128
llama3_8b	8	32	128
llama3_70b	8	80	128

Interconnect	Bandwidth	Status
NVLink	600 GB/s	viable for all TP
InfiniBand 400G	50 GB/s	viable at TP=2, borderline at TP=4
PCIe gen4	3.1 GB/s	not viable for TP disaggregation

---

Quick Start

cd ~/dev/tensor-parallel-kv-cache-bench

python3 -m venv venv
source venv/bin/activate
pip install -r requirements.txt

python -u run.py

Runtime: approximately 10 seconds. No GPU required.

---

Output Files

results/
  kv_layout.csv
  transfer_costs.csv
  prefix_sharing.csv
  gqa_analysis.csv
  tradeoff.csv
  pcie_viability.csv
  operational_decision_matrix.csv

plots/
  01_kv_per_gpu_vs_tp.png
  02_allgather_cost_vs_tp.png
  03_disagg_speedup_vs_tp.png
  04_gqa_sharding_regime.png
  05_prefix_hit_rate_degradation.png
  06_allgather_amortized.png
  07_ag_vs_recompute.png
  08_net_benefit_score.png
  09_pcie_viability.png
  10_decision_matrix.png

---

Project Structure

tensor-parallel-kv-cache-bench/
+-- src/
|   +-- config.py
|   +-- kv_layout.py
|   +-- transfer_model.py
|   +-- prefix_model.py
|   +-- tradeoff_model.py
|   +-- bench.py
|   +-- analysis.py
+-- results/
+-- plots/
+-- run.py
+-- README.md
+-- summary.txt
+-- design.md
+-- LICENSE
+-- requirements.txt

---

Methodology

This project is an analytical model calibrated from prior benchmarks.

KV layout is computed exactly from model architecture parameters. Transfer costs use an alpha-beta model with measured bandwidths. Prefix sharing degradation is modeled from eviction inconsistency probability. All-gather overlap fractions are estimated per interconnect type.

For full design details, see design.md.

---

Requirements

• Python 3.10+
• NumPy >= 1.26.0
• Pandas >= 2.0.0
• Matplotlib >= 3.8.0

No GPU required.

---

Practical Recommendation

If you need a default TP + KV cache configuration:

• use TP=2 on NVLink or InfiniBand as a safe baseline
• do not exceed n_kv_heads as your TP ceiling for GQA models
• avoid TP-based KV disaggregation on PCIe gen4
• prefer all-gather over recompute for prefix sharing on NVLink
• monitor prefix hit rate degradation as TP increases beyond TP=4

---

Documentation

• design.md -- detailed design rationale and model equations
• summary.txt -- concise high-level summary of findings
• LICENSE -- MIT License

---

License

MIT License -- Copyright (c) 2026 Joao Felipe De Souza

---

Author

Joao Felipe De Souza 2026