Skip to content

feat(connector): support tensor-parallel KV cache#91

Merged
GentleCold merged 4 commits into
masterfrom
feat/multi-tp
Jul 12, 2026
Merged

feat(connector): support tensor-parallel KV cache#91
GentleCold merged 4 commits into
masterfrom
feat/multi-tp

Conversation

@GentleCold

@GentleCold GentleCold commented Jul 11, 2026

Copy link
Copy Markdown
Owner

Summary

  • isolate every tensor-parallel rank in its own KV storage lane and publish chunks only after the full rank commit quorum
  • wait for pending stores during drain, preserve async suffix allocation state, and apply staging-pool backpressure until a real lease is released
  • derive benchmark KV capacity from the selected model, block size, and TP geometry instead of fixed model constants
  • default benchmark entry points to 128-token blocks while preserving 16-token legacy manifest decoding
  • provide the tensor-parallel rank API in CPU-only vLLM test stubs so PR CI can collect connector tests

Model geometry

For glm-4-9b-chat-1m with TP=2 and block size 128:

  • KV heads: 4 (multi_query_group_num)
  • head dimension: 128 (kv_channels)
  • layers: 40
  • aggregate slot size: 10,485,760 bytes
  • 200 x 10,000-token workload: 15,800 blocks
  • no-evict L1 sizing: 231.99 GiB, skip_l2=true

The model-aware sizing keeps total KV capacity comparable across block sizes and avoids the earlier fixed-geometry rejection.

Benchmark

python benchmarks/run_bench.py \
  --backend baseline,lmcache,daser-prefix \
  --load-generator vllm-bench \
  --model <glm-4-9b-chat-1m-path> \
  --store-dir <data-scratch> \
  --gpu-id 0,2 \
  --gpu-util 0.85 \
  --max-num-seqs 32 \
  --tensor-parallel-size 2 \
  --trust-remote-code \
  --block-size 128 \
  --bench-num-prompts 200 \
  --bench-input-len 10000 \
  --bench-output-len 1 \
  --bench-request-rate inf \
  --bench-max-concurrency 16 \
  --bench-random-range-ratio 0 \
  --bench-seed 42

All rows processed 200 requests / 1,999,600 input tokens with 0 errors.

Backend Cold prompt tok/s Warm prompt tok/s Cold mean TTFT Warm mean TTFT Warm P99 TTFT Warm backend hit
baseline 18,858.55 n/a 8,176.94 ms n/a n/a n/a
LMCache 17,533.48 344,205.91 8,793.90 ms 449.50 ms 567.66 ms 100%
DaseR prefix 17,861.61 353,695.45 8,646.56 ms 436.37 ms 536.41 ms 100%

At block size 128, DaseR warm mean TTFT is 2.92% lower than LMCache and warm prompt throughput is 2.76% higher. Relative to baseline, LMCache and DaseR reduce mean TTFT by 18.19x and 18.74x respectively.

Block-size comparison

Backend Block size Warm prompt tok/s Warm mean TTFT Warm P99 TTFT
LMCache 16 106,805.30 1,456.36 ms 1,892.92 ms
LMCache 128 344,205.91 449.50 ms 567.66 ms
DaseR prefix 16 355,202.06 435.29 ms 558.66 ms
DaseR prefix 128 353,695.45 436.37 ms 536.41 ms

Increasing the chunk/block size from 16 to 128 improves LMCache warm throughput by 3.22x and reduces mean TTFT by 69.14%. DaseR remains within 0.5%, showing that the earlier LMCache gap was dominated by its 16-token MP chunk granularity rather than cache misses or connector-version selection.

Correctness

  • baseline, LMCache cold, and DaseR cold outputs match 200/200
  • LMCache and DaseR cold outputs match 200/200
  • LMCache and DaseR warm outputs match 200/200
  • LMCache cold/warm exact match: 196/200 (98%)
  • DaseR cold/warm exact match: 196/200 (98%)
  • both backends have the same four repeat-run mismatch indices: 96, 135, 159, and 165
  • no transfer, commit, staging, CUDA, OOM, traceback, or EngineCore errors were present in the run logs

The matching cross-backend outputs and identical repeat-run drift show no DaseR-specific correctness regression.

Validation

  • CPU CI suite: 345 passed
  • benchmark/config focused suite: 116 passed
  • Ruff check and format: passed
  • targeted mypy for changed benchmark modules: passed
  • CPU-only missing-vLLM stub import simulation: passed
  • commit hooks: passed

- isolate each TP rank in a contiguous storage lane and validate runtime geometry
- publish chunks only after all distinct TP rank stores complete
- bind background CUDA work and registered staging buffers to the owning worker
- propagate TP and remote-code options through the benchmark and server startup paths
- cover TP layout, commit quorum, IPC mapping, and benchmark command contracts
- wait for pending store commits during backend drain
- retain async prefix request tokens until suffix allocation completes
- wait for a real staging lease when store buffers are exhausted
- cover TP commit drain, suffix population, and staging backpressure regressions
- size benchmark caches from model, block, and tensor-parallel geometry
- propagate TP and remote-code options to benchmark load generators
- reject unsupported FP8 and incompatible TP KV-head layouts
- cover GLM sizing and model-aware benchmark command contracts
- share the 128-token default across benchmark entry points
- preserve 16-token decoding for legacy benchmark manifests
- provide the TP-rank API in CPU-only vLLM test stubs
@GentleCold
GentleCold merged commit eb5ea84 into master Jul 12, 2026
6 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant