0619accd1c
- Document kernel gap analysis showing 77-81% CPU scheduling overhead - Identify GPU utilization at 12.8% with potential to reach 39.5% - Outline optimization directions: CUDA Graph, Triton fusion, C++ extension - Add documentation index entry in CLAUDE.md Generated with [Claude Code](https://claude.ai/code) via [Happy](https://happy.engineering) Co-Authored-By: Claude <noreply@anthropic.com> Co-Authored-By: Happy <yesreply@happy.engineering>
5.2 KiB
5.2 KiB
CPU 调度延迟分析
问题概述
在分析 nsys profile 时发现，chunked attention pipeline 中存在大量的 CPU 调度延迟，导致 GPU 利用率显著下降。
观察数据
测试环境
- GPU: NVIDIA A100-SXM4-80GB
- 模型: Llama-3.1-8B-Instruct
- 测试: RULER niah_single_1, 64K context
- Profile 文件:
ruler_8slots_test.nsys-rep - 时间段: 92.982s - 93.038s
Kernel 执行时间
| Kernel | 典型执行时间 |
|---|---|
| flash_fwd_kernel | ~138 μs |
| H2D memcpy (2MB) | ~87 μs |
| merge_lse_kernel | ~3.5 μs |
| merge_output_kernel | ~34 μs |
操作间隙分析
从 cuda_gpu_trace 观察到的间隙：
Start (ms) Dur (μs) Gap (μs) Type
------------------------------------------------------------
92984.680 138.3 378.3 flash_fwd_kernel ← GAP!
92985.051 86.8 232.9 H2D memcpy ← GAP!
92985.141 86.8 2.8 H2D memcpy
92985.587 135.9 360.0 flash_fwd_kernel ← GAP!
92986.026 3.4 302.4 merge_lse ← GAP!
92986.164 33.5 135.0 merge_output ← GAP!
92986.371 86.9 173.4 H2D memcpy ← GAP!
92986.461 86.8 2.7 H2D memcpy
92986.816 137.9 268.2 flash_fwd_kernel ← GAP!
Flash Kernel 间隙分解
| 间隙 | 总时间 | 有效工作时间 | 空闲时间 |
|---|---|---|---|
| Flash 1 → Flash 2 | 769 μs | ~174 μs (2x H2D) | ~595 μs (77%) |
| Flash 2 → Flash 3 | 1092 μs | ~211 μs (merge + H2D) | ~881 μs (81%) |
| Flash 3 → Flash 4 | 965 μs | ~211 μs (merge + H2D) | ~754 μs (78%) |
关键发现: 每个 flash kernel 之间约 77-81% 的时间是 CPU 调度空闲。
间隙来源分析
1. CPU 调度延迟类型
| 转换 | 典型延迟 | 原因 |
|---|---|---|
| Kernel 结束 → 下一个 Kernel 开始 | 100-400 μs | CPU 准备参数、调用 CUDA driver |
| Flash 结束 → H2D 开始 | ~233 μs | Python 代码执行 + CUDA launch |
| H2D 结束 → Flash 开始 | ~360 μs | 同步等待 + kernel launch |
| Flash 结束 → merge 开始 | ~302 μs | Python 代码执行 |
2. 延迟产生的代码位置
# full_policy.py: compute_chunked_prefill
for block_idx in range(num_blocks):
# 1. 等待 H2D 完成 (同步点)
offload_engine.wait_slot_layer(current_slot) # ← 可能引入延迟
# 2. 获取 KV 数据
k_block, v_block = offload_engine.get_kv_for_slot(current_slot)
# 3. 调用 flash attention (kernel launch)
block_out, block_lse = flash_attn_with_kvcache(...) # ← CPU 调度延迟
# 4. merge 操作
merge_output(...) # ← CPU 调度延迟
merge_lse(...) # ← CPU 调度延迟
# 5. 发起下一个 H2D (异步)
offload_engine.load_to_slot_layer(next_slot, ...) # ← CPU 调度延迟
3. 为什么 H2D 之间间隙小
注意到连续的 H2D memcpy 之间间隙只有 ~2.7 μs，这是因为：
- 它们在同一个 stream 上连续发起
- CUDA driver 可以批量处理
- 没有 Python 代码介入
GPU 利用率计算
基于观察数据：
| 指标 | 值 |
|---|---|
| Flash kernel 平均执行时间 | 138 μs |
| Flash kernel 平均间隔 | 942 μs |
| Flash kernel GPU 利用率 | 138 / (138 + 942) = 12.8% |
如果消除 CPU 调度延迟（仅保留必要的 H2D + merge）：
| 指标 | 值 |
|---|---|
| 必要间隔 (2x H2D + merge) | ~211 μs |
| 理论 GPU 利用率 | 138 / (138 + 211) = 39.5% |
潜在提升: 3x GPU 利用率
优化方向
1. CUDA Graph
将整个 block 处理流程编译为 CUDA Graph，消除重复的 kernel launch 开销。
# 伪代码
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
# 预录制 flash + merge 操作
block_out, block_lse = flash_attn_with_kvcache(...)
merge_output(...)
merge_lse(...)
# 运行时只需 replay
for block_idx in range(num_blocks):
graph.replay() # 单次 launch，无 Python 介入
2. 自定义 Triton Kernel
将 flash + merge 融合为单个 kernel，减少 kernel launch 次数。
3. C++ Extension
将 Python 循环移到 C++ 层，减少 Python 解释器开销。
4. 流水线重叠优化
确保 H2D 传输与前一个 block 的计算完全重叠：
Block 0: [H2D slot0] [Flash slot0] [merge]
Block 1: [H2D slot1] [Flash slot1] [merge]
Block 2: [H2D slot2] [Flash slot2] [merge]
验证方法
1. 使用 nsys 分析间隙
# 生成 profile
bash scripts/profile_offload.sh --num-gpu-blocks 8
# 查看 kernel trace
nsys stats --report cuda_gpu_trace --format csv <file>.nsys-rep | \
awk -F',' 'NR>1 && $1 >= START && $1 <= END'
2. 计算间隙
# 从 trace 数据计算
prev_end = start + duration
gap = next_start - prev_end
相关文件
nanovllm/kvcache/sparse/full_policy.py: Pipeline 实现nanovllm/kvcache/offload_engine.py: H2D/D2H 传输scripts/profile_offload.sh: Profiling 脚本