diff --git a/docs/cuda_graph_memory_guide.md b/docs/cuda_graph_memory_guide.md new file mode 100644 index 0000000..936b31e --- /dev/null +++ b/docs/cuda_graph_memory_guide.md @@ -0,0 +1,152 @@ +# CUDA Graph 内存机制指南 + +本文档基于对 Qwen3-4B 模型的实际测试,详细分析 CUDA Graph 在 LLM 推理中的内存行为。 + +## 概述 + +CUDA Graph 通过捕获 GPU kernel 执行序列并重放来减少 CPU 开销,从而提升推理性能。本指南重点分析其内存特性。 + +## 性能提升 + +| 模式 | Decode 吞吐量 | 说明 | +|------|--------------|------| +| Eager | ~25 tok/s | 每次推理重新调度 kernel | +| CUDA Graph | ~70 tok/s | 重放预录制的 kernel 序列 | +| **加速比** | **2.80x** | | + +## 内存阶段分析 + +基于 Qwen3-4B (bf16) 在 RTX 3090 上的测试结果: + +### 各阶段内存变化 + +| 阶段 | 内存 (MB) | 增量 | 说明 | +|------|-----------|------|------| +| 模型加载 | 7672 | +7672 | 模型权重 | +| StaticCache 分配 | 7816 | +144 | **主要开销** | +| Warmup (3次) | 7825 | +8 | 激活值缓存 | +| Graph 捕获 | 7833 | +8 | 存储 kernel 序列 | +| Graph Replay | 7833 | **0** | 零额外分配 | + +### 关键发现 + +1. **Graph 捕获开销很小**:仅约 8 MB,用于存储 kernel 调用序列 + +2. **StaticCache 是主要开销**: + ``` + size = num_layers × 2 × batch_size × num_kv_heads × max_cache_len × head_dim × dtype_size + ``` + - Qwen3-4B (1024 tokens): 36 × 2 × 1 × 8 × 1024 × 128 × 2 = **144 MB** + +3. **Graph Replay 零分配**:所有张量地址在 capture 时已固定,replay 只重放 kernel + +## Cache 长度与内存关系 + +| Cache 长度 | 总开销 | 每 1K tokens | +|------------|--------|--------------| +| 256 | 53 MB | 206 MB | +| 512 | 89 MB | 174 MB | +| 1024 | 161 MB | 157 MB | +| 2048 | 305 MB | 149 MB | +| 4096 | 593 MB | 145 MB | + +内存开销与 cache 长度近似线性关系,每 1K tokens 约需 145-160 MB。 + +## CUDA Graph 工作原理 + +### 核心要求:固定内存地址 + +CUDA Graph 要求所有张量在 capture 时地址固定,之后只能通过 `copy_()` 更新值: + +```python +# 分配固定地址的张量 +static_input_ids = torch.zeros(batch_size, 1, dtype=torch.long, device=device) +static_cache_position = torch.tensor([0], dtype=torch.long, device=device) + +# Capture 时使用这些张量 +with torch.cuda.graph(graph): + outputs = model(input_ids=static_input_ids, ...) + +# Replay 时通过 copy_() 更新值(地址不变) +static_input_ids.copy_(new_token) # 更新输入 +static_cache_position.fill_(position) # 更新位置 +graph.replay() # 重放 +``` + +### StaticCache vs DynamicCache + +| 特性 | DynamicCache | StaticCache | +|------|--------------|-------------| +| 内存分配 | 按需增长 | 预分配固定大小 | +| 地址稳定性 | 不稳定 | 稳定 | +| CUDA Graph 兼容 | ❌ | ✅ | +| 内存效率 | 高(按需) | 低(预分配) | + +### 典型工作流程 + +``` +1. Prefill (Eager) + └── 使用 DynamicCache 处理变长输入 + +2. 创建 StaticCache + └── 预分配 max_cache_len 大小的缓存 + +3. 复制 Prefill KV 到 StaticCache + └── 将 DynamicCache 内容拷贝到固定地址 + +4. Warmup (3次) + └── 确保所有 lazy initialization 完成 + +5. Capture Graph + └── 录制 decode 的 kernel 序列 + +6. Decode Loop + └── 更新输入 → graph.replay() → 读取输出 +``` + +## 多 Batch Size Graph 的内存问题 + +如果为多个 batch size 分别捕获 graph(如 nanovllm 的设计),内存会快速增长: + +| Batch Size | StaticCache (1024 tokens) | 累计 | +|------------|---------------------------|------| +| 1 | 144 MB | 144 MB | +| 2 | 288 MB | 432 MB | +| 4 | 576 MB | 1,008 MB | +| 8 | 1,152 MB | 2,160 MB | +| 16 | 2,304 MB | 4,464 MB | +| ... | ... | ... | + +这是因为每个 batch size 需要独立的 StaticCache。实际系统(如 nanovllm)使用 PagedAttention 共享 KV cache 来避免此问题。 + +## 测试脚本 + +提供了测试脚本用于验证以上结论: + +```bash +# 基本内存分析 +CUDA_VISIBLE_DEVICES=0 python tests/test_cudagraph_memory.py + +# 指定 cache 长度 +CUDA_VISIBLE_DEVICES=0 python tests/test_cudagraph_memory.py --max-cache-len 2048 + +# 测试 cache 长度缩放 +CUDA_VISIBLE_DEVICES=0 python tests/test_cudagraph_memory.py --test-scaling +``` + +性能对比演示: + +```bash +# Eager vs CUDA Graph 性能对比 +CUDA_VISIBLE_DEVICES=0 python tests/data/test_cudagraph_demo.py --mode both +``` + +## 总结 + +| 项目 | 结论 | +|------|------| +| 性能提升 | ~2.8x decode 吞吐量 | +| Graph 捕获开销 | ~8 MB(很小) | +| 主要内存开销 | StaticCache(与 cache_len 成正比) | +| Replay 内存 | 零额外分配 | +| 核心要求 | 固定张量地址 |