📝 docs: add XAttention performance analysis documentation
Add comprehensive performance analysis for XAttention: - NVTX marker locations and usage - Block size impact on offload mode (4096 vs 1024) - Detailed timing breakdown for estimate vs compute phases - softmax_fuse_block_sum_kernel analysis - Optimization recommendations Key findings: - block_size=4096 is 2x faster than 1024 for 64K context - find_blocks_chunked is bottleneck (40%) at block_size=4096 - estimate_gemm becomes bottleneck (24%) at block_size=1024 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>
This commit is contained in:
Zijie Tian
@@ -30,6 +30,7 @@ Nano-vLLM is a lightweight vLLM implementation (~1,200 lines) for fast offline L
|
||||
| [`docs/bench_offload_results.md`](docs/bench_offload_results.md) | 📊 BENCH: CPU offload 性能测试结果，Full vs XAttention 对比 (32K/128K) |
|
||||
| [`docs/cpu_offload_optimization_strategies.md`](docs/cpu_offload_optimization_strategies.md) | 🚀 OPT: CPU offload 优化策略：chunk size、CUDA Graph、前沿研究(InfiniGen/ShadowKV) |
|
||||
| [`docs/gpu_only_xattn_guide.md`](docs/gpu_only_xattn_guide.md) | 🚀 GPU-Only XAttention: 内存预分配、性能分析 (32K +15%, 64K +41%)、CUDA Graph 限制 |
|
||||
| [`docs/xattn_performance_analysis.md`](docs/xattn_performance_analysis.md) | 📊 XAttention 性能分析: NVTX 标记、block size 影响、estimate vs compute 耗时对比 |
|
||||
|
||||
## Rules Index
|
||||
|
||||
|
||||
170
docs/xattn_performance_analysis.md
Normal file
170
docs/xattn_performance_analysis.md
Normal file
@@ -0,0 +1,170 @@
|
||||
# XAttention Performance Analysis
|
||||
|
||||
本文档记录 XAttention 在不同配置下的性能分析结果，包括 NVTX 标记位置、block size 影响和性能瓶颈。
|
||||
|
||||
## NVTX 标记
|
||||
|
||||
XAttention 代码中添加了 NVTX 标记用于 nsys profiling，便于分析 estimate 和 compute 阶段的性能。
|
||||
|
||||
### 标记位置
|
||||
|
||||
| 模式 | 标记名称 | 文件位置 | 说明 |
|
||||
|------|---------|---------|------|
|
||||
| GPU-only | `xattn_estimate` | `xattn_bsa.py:compute_prefill` | xattn_estimate 调用 |
|
||||
| GPU-only | `xattn_bsa_compute` | `xattn_bsa.py:compute_prefill` | BSA kernel 调用 |
|
||||
| Offload | `xattn_estimate_gemm` | `xattn_bsa.py:select_blocks` | flat_group_gemm 循环 |
|
||||
| Offload | `xattn_estimate_softmax` | `xattn_bsa.py:select_blocks` | softmax_fuse_block_sum |
|
||||
| Offload | `xattn_estimate_find_blocks` | `xattn_bsa.py:select_blocks` | find_blocks_chunked |
|
||||
| Offload | `xattn_compute_historical` | `xattn_bsa.py:compute_chunked_prefill` | 历史 chunks attention |
|
||||
| Offload | `xattn_compute_current` | `xattn_bsa.py:compute_chunked_prefill` | 当前 chunk attention |
|
||||
| Offload | `xattn_compute_merge` | `xattn_bsa.py:compute_chunked_prefill` | merge 操作 |
|
||||
|
||||
### 查看 NVTX 统计
|
||||
|
||||
```bash
|
||||
# 生成 profile
|
||||
bash scripts/profile_offload.sh --policy xattn --ctx-len 64k --block-size 4096 --gpu 0
|
||||
|
||||
# 查看 NVTX 统计
|
||||
nsys stats --report nvtx_pushpop_sum results/nsys/<filename>.nsys-rep
|
||||
```
|
||||
|
||||
## Block Size 对 Offload 模式的影响
|
||||
|
||||
### 测试配置
|
||||
|
||||
- Model: Llama-3.1-8B-Instruct
|
||||
- Context: 64K tokens
|
||||
- Mode: xattn + offload
|
||||
- GPU: A100 40GB
|
||||
|
||||
### 性能对比
|
||||
|
||||
| 指标 | block_size=4096 | block_size=1024 | 变化 |
|
||||
|------|----------------|-----------------|------|
|
||||
| **总时间** | 27.7s | 55.5s | **2x 慢** |
|
||||
| **Chunks 数量** | 16 | 64 | 4x |
|
||||
| **CPU blocks** | 18 | 71 | ~4x |
|
||||
|
||||
### 各阶段耗时分布
|
||||
|
||||
#### block_size=4096
|
||||
|
||||
| 阶段 | 占比 | 总时间 | 平均时间 | 调用次数 |
|
||||
|-----|------|--------|---------|---------|
|
||||
| **xattn_estimate_find_blocks** | **39.7%** | 18.0s | 37.6ms | 480 |
|
||||
| xattn_compute_historical | 4.4% | 2.0s | 4.2ms | 480 |
|
||||
| xattn_estimate_gemm | 3.4% | 1.5s | 3.2ms | 480 |
|
||||
| xattn_compute_current | 0.2% | 113ms | 0.22ms | 512 |
|
||||
| xattn_compute_merge | 0.2% | 96ms | 0.19ms | 512 |
|
||||
| xattn_estimate_softmax | 0.2% | 88ms | 0.18ms | 480 |
|
||||
|
||||
#### block_size=1024
|
||||
|
||||
| 阶段 | 占比 | 总时间 | 平均时间 | 调用次数 |
|
||||
|-----|------|--------|---------|---------|
|
||||
| **xattn_estimate_gemm** | **23.6%** | 22.6s | 11.4ms | 1984 |
|
||||
| **xattn_compute_historical** | **16.9%** | 16.2s | 8.0ms | 2016 |
|
||||
| xattn_estimate_find_blocks | 1.4% | 1.3s | 0.66ms | 1984 |
|
||||
| xattn_compute_current | 0.5% | 433ms | 0.21ms | 2048 |
|
||||
| xattn_compute_merge | 0.4% | 373ms | 0.18ms | 2048 |
|
||||
| xattn_estimate_softmax | 0.2% | 222ms | 0.11ms | 1984 |
|
||||
|
||||
### 关键发现
|
||||
|
||||
1. **Block size 对性能影响显著**
|
||||
- block_size=1024 比 4096 慢约 2x
|
||||
- 更小的 block size 导致更多的 chunks，增加调用次数
|
||||
|
||||
2. **性能瓶颈随 block size 变化**
|
||||
- **block_size=4096**: 瓶颈是 `find_blocks_chunked` (39.7%)
|
||||
- **block_size=1024**: 瓶颈转移到 `estimate_gemm` (23.6%) 和 `compute_historical` (16.9%)
|
||||
|
||||
3. **Amortization 效应**
|
||||
- 大 block size 虽然单次 `find_blocks` 更慢 (37.6ms vs 0.66ms)
|
||||
- 但调用次数少 (480 vs 1984)，总时间反而更少
|
||||
|
||||
4. **find_blocks_chunked 的特殊性**
|
||||
- 该函数主要在 CPU 上执行 block 选择逻辑
|
||||
- 处理更大的数据量时开销显著增加
|
||||
- block_size=4096 时占用 40% 时间，是主要优化目标
|
||||
|
||||
## softmax_fuse_block_sum_kernel 性能分析
|
||||
|
||||
`softmax_fuse_block_sum_kernel_non_causal` 是 XAttention 估计阶段的核心 Triton kernel。
|
||||
|
||||
### Kernel 结构
|
||||
|
||||
```python
|
||||
# 每个 thread block 处理的数据形状
|
||||
工作负载: [block_size, segment_size] # 单个 Q block 对所有 K 的注意力
|
||||
|
||||
# Pass 1: 计算全局 softmax 参数 (m_i, l_i)
|
||||
for iter in range(num_iters): # num_iters = k_len / segment_size
|
||||
X = load [block_size, segment_size]
|
||||
compute max, sum for softmax normalization
|
||||
|
||||
# Pass 2: Normalize + Block Sum
|
||||
for iter in range(num_iters):
|
||||
X = load [block_size, segment_size]
|
||||
X = softmax(X)
|
||||
X = reshape(X, [block_size, segment_size/block_size, block_size])
|
||||
X = sum(X, axis=2) # → [block_size, segment_size/block_size]
|
||||
X = sum(X, axis=0) # → [segment_size/block_size]
|
||||
store output
|
||||
```
|
||||
|
||||
### 性能随 block_size 变化的因素
|
||||
|
||||
| 因素 | 小 block_size (64) | 大 block_size (256) |
|
||||
|------|-------------------|---------------------|
|
||||
| Grid 并行度 | 高 (更多 blocks) | 低 (更少 blocks) |
|
||||
| 寄存器使用 | 低 | 高 (可能 spill) |
|
||||
| L2 Cache 复用 | 差 | 好 |
|
||||
| 输出大小 | 大 | 小 |
|
||||
|
||||
### 典型性能曲线
|
||||
|
||||
```
|
||||
Performance
|
||||
│
|
||||
│ ┌─────┐
|
||||
│ / \
|
||||
│ / \
|
||||
│ / \
|
||||
│ / \
|
||||
└────/───────────────\────────→ block_size
|
||||
64 128 256 512
|
||||
|
||||
最优点通常在 128-256 之间
|
||||
```
|
||||
|
||||
## 优化建议
|
||||
|
||||
1. **优先使用 block_size=4096**
|
||||
- 减少 chunk 数量，降低调度开销
|
||||
- 更好的 amortization 效果
|
||||
|
||||
2. **优化 find_blocks_chunked**
|
||||
- 当前是 block_size=4096 的主要瓶颈
|
||||
- 考虑 GPU 加速或批量处理
|
||||
|
||||
3. **Pipeline 优化**
|
||||
- 利用多 slot 的 ring buffer 实现计算和传输 overlap
|
||||
- 当前已实现，但 find_blocks 是 CPU 操作，无法 overlap
|
||||
|
||||
## 测试命令
|
||||
|
||||
```bash
|
||||
# GPU-only 模式 (需要 40GB+ VRAM)
|
||||
bash scripts/profile_offload.sh --policy xattn --ctx-len 64k --no-offload --gpu 0
|
||||
|
||||
# Offload 模式，block_size=4096
|
||||
bash scripts/profile_offload.sh --policy xattn --ctx-len 64k --block-size 4096 --gpu 0
|
||||
|
||||
# Offload 模式，block_size=1024
|
||||
bash scripts/profile_offload.sh --policy xattn --ctx-len 64k --block-size 1024 --gpu 0
|
||||
|
||||
# 128K context
|
||||
bash scripts/profile_offload.sh --policy xattn --ctx-len 128k --block-size 4096 --gpu 0
|
||||
```
|
||||
Reference in New Issue
Block a user