Files

Gahow Wang 2d48f25e66 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>

2026-05-22 11:50:12 +08:00

1.8 KiB

Raw Permalink Blame History

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 留待后续优化。

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。

需要的新功能

Tensor slice 支持（view into sub-range of a dimension）
GPU D2D copy at offset（写入 cache 指定位置）
去掉 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

1.8 KiB Raw Permalink Blame History Unescape Escape