nano-vllm/docs/xattn_density_benchmark.md

XAttention Density Benchmark

GPU-only 模式下 XAttention Block Sparse Attention 的 density 测试结果。

测试配置

参数	值	说明
Model	Llama-3.1-8B-Instruct	32 layers, 32 heads, 8 KV heads
Block Size	128 tokens	BSA kernel 固定要求
Threshold	0.9 / 0.95	累积注意力阈值
Stride	4 / 8 / 16	Q/K 下采样步长
Dataset	RULER niah_single_1	Sample 0
Mode	GPU-only	无 CPU offload

Density 定义

# Density = selected_blocks / total_causal_blocks
# 在 causal attention 下，只计算下三角区域的 blocks
# Overall density = 所有层的平均值

def compute_density(mask, causal=True):
    """
    mask: [batch, heads, q_blocks, k_blocks] boolean tensor
    """
    if causal:
        causal_mask = torch.tril(torch.ones(q_blocks, k_blocks))
        total = causal_mask.sum() * batch * heads
        selected = (mask & causal_mask).sum()
    return selected / total

测试结果

threshold=0.9

Overall Density (平均)

Context	stride=4	stride=8	stride=16
4K	0.5220 (52.2%)	0.5292 (52.9%)	0.5430 (54.3%)
8K	0.5152 (51.5%)	0.5252 (52.5%)	0.5396 (54.0%)
16K	0.4682 (46.8%)	0.4775 (47.8%)	0.4888 (48.9%)
32K	0.3700 (37.0%)	0.4012 (40.1%)	0.4196 (42.0%)

Min Density (per layer)

Context	stride=4	stride=8	stride=16
4K	0.2805 (Layer 3)	0.3132 (Layer 3)	0.3376 (Layer 5)
8K	0.2886 (Layer 5)	0.2725 (Layer 5)	0.2995 (Layer 5)
16K	0.2247 (Layer 5)	0.2349 (Layer 5)	0.2451 (Layer 5)
32K	0.1799 (Layer 5)	0.1846 (Layer 5)	0.1964 (Layer 5)

threshold=0.95

Overall Density (平均)

Context	stride=4	stride=8	stride=16
4K	0.6561 (65.6%)	0.6699 (67.0%)	0.6815 (68.2%)
8K	0.6462 (64.6%)	0.6584 (65.8%)	0.6732 (67.3%)
16K	0.6004 (60.0%)	0.6114 (61.1%)	0.6193 (61.9%)
32K	0.4894 (48.9%)	0.5203 (52.0%)	0.5385 (53.9%)

Min Density (per layer)

Context	stride=4	stride=8	stride=16
4K	0.3972 (Layer 3)	0.4348 (Layer 5)	0.4517 (Layer 4)
8K	0.4004 (Layer 5)	0.3906 (Layer 5)	0.4239 (Layer 5)
16K	0.3331 (Layer 5)	0.3453 (Layer 5)	0.3589 (Layer 5)
32K	0.2656 (Layer 5)	0.2784 (Layer 5)	0.2917 (Layer 5)

threshold 对比 (stride=8)

Context	threshold=0.9	threshold=0.95	差异
4K	0.5292 (52.9%)	0.6699 (67.0%)	-14.1%
8K	0.5252 (52.5%)	0.6584 (65.8%)	-13.3%
16K	0.4775 (47.8%)	0.6114 (61.1%)	-13.4%
32K	0.4012 (40.1%)	0.5203 (52.0%)	-11.9%

关键发现

1. Context Length 影响最大

Density 随 context length 显著下降（threshold=0.9, stride=8）：

• 4K: 52.9% density
• 8K: 52.5% density
• 16K: 47.8% density
• 32K: 40.1% density

结论: 长序列有更多稀疏化机会，XAttention 的优势在长序列上更明显。

2. Threshold 影响显著

threshold=0.9 比 0.95 的 density 低约 12-14%：

• 0.9 更激进，选择更少的 blocks
• 0.95 更保守，保留更多 blocks
• 两者准确性都不受影响（RULER NIAH 全部 PASS）

3. Stride 影响较小

同一 context 下，不同 stride 的 density 差异约 2-5%：

• stride 越大 → density 略高（采样越粗糙，选择更保守）
• stride=4 最激进，stride=16 最保守

4. Min Density 集中在中间层

• 大多数情况下 min density 出现在 Layer 5
• 中间层的稀疏性最高，首尾层相对密集
• 这符合 Transformer 注意力模式的一般规律

5. 最佳稀疏化配置

32K + stride=4 + threshold=0.9 达到最低 density：

• Overall: 37.0% (节省 63% 计算)
• Min: 18.0% (Layer 5)

6. 准确性稳定

所有配置下 RULER NIAH 测试都 PASS (score=1.0)，说明：

• threshold=0.9 和 0.95 都足够保守，不损失准确性
• 不同 stride 不影响最终结果

推荐配置

场景	threshold	stride	说明
精度优先	0.95	8	保守配置，density ~52-67%
平衡	0.9	8	默认配置，density ~40-53%
性能优先	0.9	4	激进配置，density ~37-52%

测试命令

# 基本测试
CUDA_VISIBLE_DEVICES=0 PYTHONPATH=/path/to/nano-vllm:$PYTHONPATH \
    python tests/test_ruler.py \
    --data-dir tests/data/ruler_32k \
    --datasets niah_single_1 \
    --sample-indices 0 \
    --max-model-len 33792 \
    --sparse-policy XATTN_BSA \
    --sparse-threshold 0.9 \
    --sparse-stride 8 \
    --gpu-utilization 0.85

# 参数说明
# --sparse-policy XATTN_BSA    启用 XAttention Block Sparse Attention
# --sparse-threshold 0.9       累积注意力阈值 (0.9-0.99)
# --sparse-stride 8            Q/K 下采样步长 (4/8/16)

DensityObserver 使用

from nanovllm.utils.density_observer import DensityObserver

# 启用并重置
DensityObserver.enable()
DensityObserver.complete_reset()

# ... 运行 inference (compute_prefill 自动记录) ...

# 获取结果
summary = DensityObserver.get_summary()
# {
#     "mode": "gpu_only",
#     "overall_density": 0.40,  # 所有层的平均值
#     "per_layer_density": {0: 0.55, 1: 0.45, ...},
#     "num_layers": 32
# }

# 获取最低 density
min_layer, min_density = DensityObserver.get_min_density()

# 打印摘要
DensityObserver.print_summary()
# [DensityObserver] Mode: gpu_only
#   Overall density: 0.4012
#   Min density: 0.1846 (layer 5)
#   Num layers: 32

相关文件

文件	说明
nanovllm/kvcache/sparse/xattn_bsa.py	XAttention BSA Policy 实现
nanovllm/utils/density_observer.py	Density 统计 Observer
nanovllm/ops/xattn.py	xattn_estimate 核心算法
tests/test_ruler.py	RULER benchmark 测试脚本