f3e4611e3b
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>
171 lines
5.9 KiB
Markdown
171 lines
5.9 KiB
Markdown
# 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
|
||
```
|