Gahow Wang
2d48f25e66
- 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>
60 lines
1.8 KiB
Markdown
60 lines
1.8 KiB
Markdown
# 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
|