diff --git a/bench.py b/bench.py index d1dac2e..3358157 100644 --- a/bench.py +++ b/bench.py @@ -58,6 +58,8 @@ def main(): help="Enable sparse policy routing (FullAttentionPolicy by default)") parser.add_argument("--gpu-util", type=float, default=0.9, help="GPU memory utilization (default: 0.9)") + parser.add_argument("--block-size", type=int, default=1024, + help="KV cache block size (default: 1024)") parser.add_argument("--enforce-eager", action="store_true", help="Disable CUDA graphs (default: False)") args = parser.parse_args() @@ -83,6 +85,7 @@ def main(): max_num_batched_tokens=max_len, sparse_policy=sparse_policy, gpu_memory_utilization=args.gpu_util, + kvcache_block_size=args.block_size, ) # Warmup diff --git a/docs/bench_offload_results.md b/docs/bench_offload_results.md index 9a3ff7b..e35e206 100644 --- a/docs/bench_offload_results.md +++ b/docs/bench_offload_results.md @@ -9,8 +9,6 @@ | GPU | NVIDIA A100-SXM4-80GB | | 模型 | Llama-3.1-8B-Instruct | | GPU slots | 4 | -| Block size | 1024 tokens | -| Chunk size | 2048 tokens | ## Sparse Policy 配置 @@ -21,64 +19,114 @@ ## 测试结果 -### 32K 上下文 +### Block Size 4096 (推荐) -| 策略 | 输入长度 | 耗时 | 吞吐量 | 相对性能 | -|------|----------|------|--------|----------| -| Full Attention | 32767 tok | 20.64s | **1587.74 tok/s** | baseline | -| XAttention BSA | 32767 tok | 27.95s | **1172.33 tok/s** | 0.74x | +#### GPU-only 模式 -### 128K 上下文 +| 上下文 | Full Attention | XAttention | 相对性能 | +|--------|----------------|------------|----------| +| 32K | 4863 tok/s | 5587 tok/s | **+14.9%** ✅ | +| 64K | 3373 tok/s | 4766 tok/s | **+41.3%** ✅ | -| 策略 | 输入长度 | 耗时 | 吞吐量 | 相对性能 | -|------|----------|------|--------|----------| -| Full Attention | 131071 tok | 237.18s | **552.63 tok/s** | baseline | -| XAttention BSA | 131071 tok | 281.17s | **466.17 tok/s** | 0.84x | +#### CPU Offload 模式 -### KV Cache 配置 +| 上下文 | Full Attention | XAttention | 相对性能 | +|--------|----------------|------------|----------| +| 32K | 4648 tok/s | 4002 tok/s | **-13.9%** ❌ | +| 64K | 3329 tok/s | 2642 tok/s | **-20.6%** ❌ | +| 128K | 2122 tok/s | 867 tok/s | **-59.1%** ❌ | -| 上下文 | GPU Memory | CPU Memory | Total | -|--------|------------|------------|-------| -| 32K | 512 MB (4 blocks) | 4096 MB (32 blocks) | 4608 MB | -| 128K | 512 MB (4 blocks) | 16384 MB (128 blocks) | 16896 MB | +### Block Size 256 (小 block 测试) -## 分析 +#### CPU Offload 模式 (64K) -### XAttention 性能特点 +| 策略 | 耗时 | 吞吐量 | 相对性能 | +|------|------|--------|----------| +| Full Attention | 401.04s | 163.41 tok/s | baseline | +| XAttention BSA | 390.35s | 167.89 tok/s | **+2.7%** ✅ | -1. **32K 上下文**: XAttention 比 Full 慢 26% -2. **128K 上下文**: XAttention 比 Full 慢 16% +### Block Size 1024 (历史测试) -随着上下文增长,XAttention 的相对性能有所提升(74% → 84%),但仍未超过 Full Attention。 +#### CPU Offload 模式 -### 原因分析 +| 上下文 | Full Attention | XAttention | 相对性能 | +|--------|----------------|------------|----------| +| 32K | 1587.74 tok/s | 1172.33 tok/s | -26% | +| 128K | 552.63 tok/s | 466.17 tok/s | -16% | -1. **tau=0.95 阈值较高**: 需要覆盖 95% 累积注意力,实际跳过的 block 较少 -2. **估计开销**: `xattn_estimate_chunked` 需要对每个 chunk 计算稀疏 mask -3. **BSA kernel overhead**: Block sparse kernel 有额外的 mask 处理和索引开销 -4. **Offload 瓶颈**: CPU→GPU 传输是主要瓶颈,稀疏注意力节省的是计算而非传输 +## 关键发现 -### 适用场景 +### 1. GPU-only vs CPU Offload 模式差异 -XAttention BSA 更适合以下场景: -- 更长的上下文(256K+),稀疏收益更明显 -- 计算密集型任务(非 offload 模式),传输不是瓶颈 -- 较低的 tau 阈值(如 0.8),增加稀疏性 +| 模式 | XAttention 效果 | 原因 | +|------|-----------------|------| +| **GPU-only** | ✅ 显著加速 (+15% ~ +41%) | 计算是瓶颈,稀疏注意力减少 FLOPs | +| **CPU Offload** | ❌ 性能下降 (-14% ~ -59%) | 传输是瓶颈,稀疏估计增加额外开销 | + +### 2. Block Size 对性能的影响 + +| Block Size | 64K Full (Offload) | 特点 | +|------------|-------------------|------| +| 4096 | 3329 tok/s | ⭐ 最佳性能 | +| 1024 | ~1500 tok/s | 中等 | +| 256 | 163 tok/s | 极慢(20x 下降) | + +**原因**: 更小的 block = 更多的 blocks = 更多 H2D 传输开销 + +### 3. XAttention 在小 Block Size 下反转 + +当 block size = 256 时,XAttention 反而略有优势 (+2.7%): +- 256 个 blocks (vs 16 个 @ 4096) +- 稀疏跳过的 blocks 比例更明显 +- 但绝对性能极差,不推荐使用 + +### 4. 性能下降随上下文增长加剧 + +``` +Offload 模式 XAttention 相对性能: +32K: -14% (传输占 ~60%) +64K: -21% (传输占 ~70%) +128K: -59% (传输占 ~80%) +``` + +原因: +- 传输占比随上下文增长 +- XAttention 估计开销 O(num_chunks) 线性增长 +- 节省的计算量被传输瓶颈掩盖 + +## 结论 + +### 推荐配置 + +| 场景 | 推荐策略 | Block Size | +|------|----------|------------| +| GPU-only (VRAM 充足) | XAttention | 4096 | +| CPU Offload | Full Attention | 4096 | + +### XAttention 适用条件 + +✅ **适合**: +- GPU-only 模式(计算密集) +- 长上下文(64K+)收益更大 + +❌ **不适合**: +- CPU Offload 模式(传输密集) +- 短上下文(<32K)收益不明显 ## 运行命令 ```bash -# Full Attention (32K) -CUDA_VISIBLE_DEVICES=0 python bench_offload.py --max-len 32768 +# GPU-only 模式 +CUDA_VISIBLE_DEVICES=0 python bench.py --max-len 65536 --block-size 4096 --gpu-util 0.7 +CUDA_VISIBLE_DEVICES=0 python bench.py --max-len 65536 --block-size 4096 --gpu-util 0.7 --policy xattn -# XAttention BSA (32K) -CUDA_VISIBLE_DEVICES=0 python bench_offload.py --max-len 32768 --enable-xattn +# CPU Offload 模式 (推荐 block-size 4096) +CUDA_VISIBLE_DEVICES=0 python bench_offload.py --max-len 65536 --block-size 4096 +CUDA_VISIBLE_DEVICES=0 python bench_offload.py --max-len 65536 --block-size 4096 --enable-xattn -# Full Attention (128K) -CUDA_VISIBLE_DEVICES=0 python bench_offload.py --max-len 131072 - -# XAttention BSA (128K) -CUDA_VISIBLE_DEVICES=0 python bench_offload.py --max-len 131072 --enable-xattn +# CPU Offload 模式 (小 block size 测试) +CUDA_VISIBLE_DEVICES=0 python bench_offload.py --max-len 65536 --block-size 256 +CUDA_VISIBLE_DEVICES=0 python bench_offload.py --max-len 65536 --block-size 256 --enable-xattn # 调整 XAttention 参数 CUDA_VISIBLE_DEVICES=0 python bench_offload.py --enable-xattn --xattn-threshold 0.8 --xattn-stride 16 @@ -86,4 +134,5 @@ CUDA_VISIBLE_DEVICES=0 python bench_offload.py --enable-xattn --xattn-threshold ## 更新记录 +- 2026-01-27: 添加 GPU-only vs Offload 对比,block size 影响分析 - 2026-01-27: 初始测试,Llama-3.1-8B-Instruct, A100 80GB