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:
59
docs/11-paged-attention.md
Normal file
59
docs/11-paged-attention.md
Normal 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
|
||||
Reference in New Issue
Block a user