---
title: "Block-Sparse Attention Mask (分块稀疏注意力掩码)"
created: 2025-06-02
updated: 2025-06-02
type: concept
tags: [attention, training-optimization, multi-turn-reasoning]
sources: ["[[goru-one-pass-to-reason-2025]]"]
---

# Block-Sparse Attention Mask

> [[goru-one-pass-to-reason-2025|One-Pass to Reason]] 中设计的自定义注意力掩码，通过为不同 token 类型指定不同的可见性规则，在单次前向传播中同时满足"生成时可见、上下文时隐藏"的冲突需求。

## 可见性规则

在多轮推理对话中，token 被分为四种角色：

| Token 类型 | 角色 | 可见范围 |
|-----------|------|---------|
| hi | 人类消息 | A(H_{<i}) — 只看历史 |
| ti | 推理 token | A(H_{<i}, hi) — 看历史+当前人类消息 |
| ri_in | 上下文副本 | A(H_{<i}, hi) — 不看推理 token |
| ri_out | 生成副本 | A(H_{<i}, hi, ti) — 看全部 |

## 结构特点

1. **分块稀疏**：整个掩码按对话轮次和 token 类型划分为 block，每个 block 有不同的稀疏模式
2. **因果性保持**：所有 token 仍遵循因果约束（不偷看未来）
3. **选择性可见**：核心创新——ri_in 和 ri_out 虽然内容相同，但对 ti 的可见性不同

## 与标准掩码的区别

- **标准因果掩码**：所有 token 看到所有之前的 token
- **Block-Sparse Mask**：特定 token 类型跳过特定块（ri_in 跳过 ti）

## 高效生成

因为掩码需要在 GPU 上为每个 batch 动态生成，论文提出了向量化生成算法：

```python
# 伪代码
M[i][j] = can_see(role[i], role[j]) AND causal(i, j)
# role 映射: 0=pad, 1=human, 2=thinking, 3=response_in, 4=response_out
```

用卡诺图（Karnaugh map）化简布尔逻辑，最小化逻辑运算次数。

## 实现后端

使用 [[flex-attention|PyTorch FlexAttention]]（Dong et al., 2024），因为 FlashAttention-2 不支持传入自定义掩码。

## 与 packing 的兼容

可与 [[sequence-packing]] 叠加：打包掩码（防跨样本污染）与自定义注意力掩码通过逻辑 AND 结合。

## 相关

- [[one-pass-fine-tuning]]
- [[token-duplication]]
- [[visibility-constraint]]
- [[flex-attention]]
- [[goru-one-pass-to-reason-2025|One-Pass to Reason 论文]]
