✨ feat: add XAttention BSA support to bench_offload.py

- Add --model parameter (default: Llama-3.1-8B-Instruct) - Add --enable-xattn flag for XAttention BSA sparse prefill - Add --xattn-threshold and --xattn-stride parameters - Change default num-gpu-blocks from 6 to 4 - Add benchmark results doc with Full vs XAttn comparison (32K/128K) 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>
2026-01-27 04:20:16 +08:00
parent 924a0d2bfa
commit 73c9dc46ff
3 changed files with 115 additions and 4 deletions
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -27,6 +27,7 @@ Nano-vLLM is a lightweight vLLM implementation (~1,200 lines) for fast offline L
 | [`docs/chunked_attention_solutions.md`](docs/chunked_attention_solutions.md) | 🔧 SOLUTIONS: Chunked attention 准确性问题的代码分析和解决方案 |
 | [`docs/nsys_wrong_event_order_bug.md`](docs/nsys_wrong_event_order_bug.md) | 🐛 NSYS BUG: Ring buffer pipeline 触发 nsys 时间戳乱序问题的调试记录 |
 | [`docs/cpu_scheduling_latency_analysis.md`](docs/cpu_scheduling_latency_analysis.md) | ⚡ PERF: CPU 调度延迟分析，kernel 间隙来源，GPU 利用率优化方向 |
+| [`docs/bench_offload_results.md`](docs/bench_offload_results.md) | 📊 BENCH: CPU offload 性能测试结果，Full vs XAttention 对比 (32K/128K) |

 ## Rules Index

--- a/bench_offload.py
+++ b/bench_offload.py
@@ -46,24 +46,41 @@ def main():
    from nanovllm.config import SparsePolicyType

    parser = argparse.ArgumentParser(description="Benchmark CPU offload performance")
-    parser.add_argument("--enable-quest", action="store_true", help="Enable Quest sparse attention for decode")
+    parser.add_argument("--model", type=str, default="~/models/Llama-3.1-8B-Instruct",
+                        help="Model path (default: ~/models/Llama-3.1-8B-Instruct)")
+    # Sparse policy selection (mutually exclusive)
+    sparse_group = parser.add_mutually_exclusive_group()
+    sparse_group.add_argument("--enable-quest", action="store_true",
+                              help="Enable Quest sparse attention (decode only, prefill uses full)")
+    sparse_group.add_argument("--enable-xattn", action="store_true",
+                              help="Enable XAttention BSA (prefill only, decode uses full)")
+    # Quest parameters
    parser.add_argument("--topk", type=int, default=16, help="Top-K blocks for Quest (default: 16)")
    parser.add_argument("--threshold", type=int, default=4, help="Apply sparse only when blocks > threshold (default: 4)")
+    # XAttention parameters
+    parser.add_argument("--xattn-threshold", type=float, default=0.95,
+                        help="XAttention cumulative attention threshold (default: 0.95)")
+    parser.add_argument("--xattn-stride", type=int, default=8,
+                        help="XAttention Q/K downsampling stride (default: 8)")
+    # General parameters
    parser.add_argument("--input-len", type=int, default=None, help="Input length in tokens")
    parser.add_argument("--output-len", type=int, default=64, help="Output length for decode benchmark (default: 64)")
-    parser.add_argument("--num-gpu-blocks", type=int, default=6, help="Number of GPU blocks (default: 6)")
+    parser.add_argument("--num-gpu-blocks", type=int, default=4, help="Number of GPU blocks (default: 4)")
    parser.add_argument("--max-len", type=int, default=32*1024, help="Max model length (default: 32K)")
    parser.add_argument("--bench-decode", action="store_true", help="Run decode benchmark (default: prefill only)")
    parser.add_argument("--bench-all", action="store_true", help="Run both prefill and decode benchmarks")
    args = parser.parse_args()

-    path = os.path.expanduser("~/models/Qwen3-4B-Instruct-2507/")
+    path = os.path.expanduser(args.model)
    max_len = args.max_len

    # Setup policy configuration
    if args.enable_quest:
        sparse_policy = SparsePolicyType.QUEST
-        print(f"\n[Quest Sparse Attention] topk={args.topk}, threshold={args.threshold}")
+        print(f"\n[Quest Sparse Attention] decode: Quest (topk={args.topk}, threshold={args.threshold}), prefill: Full")
+    elif args.enable_xattn:
+        sparse_policy = SparsePolicyType.XATTN_BSA
+        print(f"\n[XAttention BSA] prefill: XAttn (tau={args.xattn_threshold}, stride={args.xattn_stride}), decode: Full")
    else:
        sparse_policy = SparsePolicyType.FULL
        print("\n[Full Attention] baseline (no sparse)")
@@ -78,8 +95,12 @@ def main():
        enable_cpu_offload=True,
        num_gpu_blocks=args.num_gpu_blocks,
        sparse_policy=sparse_policy,
+        # Quest parameters
        sparse_topk_blocks=args.topk,
        sparse_threshold_blocks=args.threshold,
+        # XAttention parameters
+        sparse_threshold=args.xattn_threshold,
+        sparse_stride=args.xattn_stride,
    )

    # Warmup
--- a/docs/bench_offload_results.md
+++ b/docs/bench_offload_results.md
@@ -0,0 +1,89 @@
+# CPU Offload Benchmark Results
+
+本文档记录 `bench_offload.py` 在不同配置下的性能测试结果。
+
+## 测试环境
+
+| 参数 | 值 |
+|------|-----|
+| GPU | NVIDIA A100-SXM4-80GB |
+| 模型 | Llama-3.1-8B-Instruct |
+| GPU slots | 4 |
+| Block size | 1024 tokens |
+| Chunk size | 2048 tokens |
+
+## Sparse Policy 配置
+
+| 策略 | Prefill | Decode | 说明 |
+|------|---------|--------|------|
+| FULL | Full Attention | Full Attention | 基线，加载所有 blocks |
+| XATTN_BSA | XAttention (tau=0.95, stride=8) | Full Attention (fallback) | 稀疏 prefill |
+
+## 测试结果
+
+### 32K 上下文
+
+| 策略 | 输入长度 | 耗时 | 吞吐量 | 相对性能 |
+|------|----------|------|--------|----------|
+| Full Attention | 32767 tok | 20.64s | **1587.74 tok/s** | baseline |
+| XAttention BSA | 32767 tok | 27.95s | **1172.33 tok/s** | 0.74x |
+
+### 128K 上下文
+
+| 策略 | 输入长度 | 耗时 | 吞吐量 | 相对性能 |
+|------|----------|------|--------|----------|
+| Full Attention | 131071 tok | 237.18s | **552.63 tok/s** | baseline |
+| XAttention BSA | 131071 tok | 281.17s | **466.17 tok/s** | 0.84x |
+
+### KV Cache 配置
+
+| 上下文 | 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 |
+
+## 分析
+
+### XAttention 性能特点
+
+1. **32K 上下文**: XAttention 比 Full 慢 26%
+2. **128K 上下文**: XAttention 比 Full 慢 16%
+
+随着上下文增长，XAttention 的相对性能有所提升（74% → 84%），但仍未超过 Full Attention。
+
+### 原因分析
+
+1. **tau=0.95 阈值较高**: 需要覆盖 95% 累积注意力，实际跳过的 block 较少
+2. **估计开销**: `xattn_estimate_chunked` 需要对每个 chunk 计算稀疏 mask
+3. **BSA kernel overhead**: Block sparse kernel 有额外的 mask 处理和索引开销
+4. **Offload 瓶颈**: CPU→GPU 传输是主要瓶颈，稀疏注意力节省的是计算而非传输
+
+### 适用场景
+
+XAttention BSA 更适合以下场景：
+- 更长的上下文（256K+），稀疏收益更明显
+- 计算密集型任务（非 offload 模式），传输不是瓶颈
+- 较低的 tau 阈值（如 0.8），增加稀疏性
+
+## 运行命令
+
+```bash
+# Full Attention (32K)
+CUDA_VISIBLE_DEVICES=0 python bench_offload.py --max-len 32768
+
+# XAttention BSA (32K)
+CUDA_VISIBLE_DEVICES=0 python bench_offload.py --max-len 32768 --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
+
+# 调整 XAttention 参数
+CUDA_VISIBLE_DEVICES=0 python bench_offload.py --enable-xattn --xattn-threshold 0.8 --xattn-stride 16
+```
+
+## 更新记录
+
+- 2026-01-27: 初始测试，Llama-3.1-8B-Instruct, A100 80GB