phase 11: GPU-resident KV cache

- GpuKVCache: pre-allocated GPU buffers, D2D copy append at offset
- Per-head strided layout [num_kv_heads, max_seq_len, head_dim]
- Fixed critical bug: seq_len must advance AFTER all layers write
  (not inside the loop per-layer)
- GpuBuffer::copy_from_device_at for offset-based D2D copy
- Tensor::from_storage constructor for wrapping raw GPU buffers
- Exported Storage and Dims from xserv-tensor

Correctness: GPU KV cache vs CPU KV cache = 50/50 bit-identical
Performance: ~neutral (KV cache was never the main bottleneck —
reshape/merge/transpose CPU round-trips dominate for Qwen3-8B)

TTFT: 122ms, TBT: 142ms, 7.0 tok/s (marginal change from 7.3)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-05-22 11:50:12 +08:00
parent be5c64ea8a
commit 2d48f25e66
9 changed files with 274 additions and 14 deletions

View File

@@ -0,0 +1,59 @@
# Phase 11: Paged Attention + KV Cache Manager — Design Document
## Goal
将 KV cache 从 CPU Vec 迁移到 GPU使用 block-based paging 管理显存。消除每步 decode 的 CPU round-trip当前 KV cache 最大性能瓶颈之一)。
## 当前问题
每步 decode 的 KV cache 路径:
```
GPU tensor (K_new) → CPU (per-head Vec append) → reconstruct → CPU tensor → GPU tensor
```
这涉及 2 次 GPU↔CPU 拷贝 × 36 层 × 2(K,V) = 144 次 transfer/token。
## 目标设计
KV cache 直接存在 GPU 上decode 时只做 GPU→GPU append
```
GPU tensor (K_new) → GPU KV cache (in-place append, no CPU)
```
## 实现方案
### GPU KV Cache简化版非 paged
先实现连续分配的 GPU KV cache预分配 max_seq_len消除 CPU round-trip。Paged allocation 留待后续优化。
```rust
pub struct GpuKVCache {
// 预分配: [num_layers, 2, num_kv_heads, max_seq_len, head_dim] on GPU
k_caches: Vec<Tensor>, // per layer: [1, num_kv_heads, max_seq_len, head_dim]
v_caches: Vec<Tensor>,
seq_len: usize, // 当前已填充的长度
max_seq_len: usize,
}
```
### Append 操作
用 cudaMemcpy D2D 将新 K/V 写入 cache 的正确偏移位置:
```
k_cache[layer][0, :, seq_len:seq_len+new, :] = k_new[0, :, :, :]
```
### 读取操作
不需要拷贝——直接用 view/slice 返回 [0, :, 0:seq_len, :] 的 GPU tensor。
## 需要的新功能
1. Tensor slice 支持view into sub-range of a dimension
2. GPU D2D copy at offset写入 cache 指定位置)
3. 去掉 Qwen3/GPT-2 forward 中的 CPU round-trip KV cache 路径
## Test Plan
- [ ] GPU KV cache 输出与 CPU KV cache bit-identical
- [ ] Benchmark: TBT 应显著降低(消除 144 次 CPU round-trip
- [ ] 50-prompt correctness re-validation