VSA 实现详解：三维分块、均值路由与块稀疏 Attention

发表于 2026-10-23 分类于深度学习阅读次数：本文字数： 20k 阅读时长 ≈ 1:14

VSA（Video Sparse Attention）先把视频 token 按 (T,H,W) 时空邻域分块，对每块的 Q、K 分别求均值，再用两个均值向量的点积决定保留哪些区块连接。选块完成后，正式 attention 仍用原始 Q/K/V 计算。

下文默认读者熟悉 FlashAttention，按“全量 Q/K → 分块均值 → 小 S → 每块 TopK → 稀疏 attention”的顺序讲。例子固定为 (T,H,W)=(4,4,4) 的 64 个 token 切成四个 (4,2,2) 区块，同一批 token 会依次经过池化、选块和 GPU 读取。

先看完整计算顺序

先为所有 token 生成完整的 Q、K、V；把 Q、K 按块分别平均成小 Q、小 K，视频按三维邻域分块；两者相乘得到小 S；每个视频 query 块按小 S 选 TopK；再用原始 Q/K/V 计算选中连接的 attention。



%%{init: {"flowchart": {"nodeSpacing": 24, "rankSpacing": 26, "padding": 12}, "themeVariables": {"fontSize": "14px"}}}%%
flowchart TB
    full["① 所有 token 的 Q、K、V<br/>每个都是 N×D"]
    full -->|Q、K 分别按块求均值| pool["② 小 Q、小 K<br/>每个都是 J×D"]
    pool -->|小 Q × 小 Kᵀ / √D| score["③ 小 S：J×J<br/>一行对应一个 query 块"]
    score -->|逐行、逐视频区域取 TopK| select["④ 每个视频 query 块<br/>得到自己的 KV 块索引"]
    select --> fine["⑤ 只计算保留的块连接<br/>原始 QKᵀ → softmax → PV"]
    full -. 原始 Q/K/V 保留给正式计算 .-> fine

上图省略 batch 和 head 维；N 是 token 数，J 是块数。“全量 Q/K”指完整的向量序列，不必先算出 N×N 的完整 QKᵀ。下文把小 S 记作路由矩阵 E，以区别于正式 attention 的 token 级 logits。

一个“视频组”就是同一 query 块。块内的 token 共用一组选中的 KV 块，各 token 的权重仍分别计算；选块由当前 Q/K 决定，可随 head、层和去噪步变化。参与选择的是 query block，不是 value 张量 V。Context 的全连接与 R2V 的分组规则见第 5.3 节。

1. 先理解 (T,H,W) 分块：一个块包含哪些 token

T、H、W 是模型中视频 token 网格的时间、高度和宽度，已经过编码和 patch 化，不能直接当作原视频的帧数与像素尺寸。为避免与高度 H 混淆，下文用 h 表示 attention head 数，D 表示每个 head 的维度。

设 token 网格为 (4,4,4)，分块大小为 (4,2,2)：

1
2
3

整个网格：4 个时间位置 × 4 行 × 4 列 = 64 个 token
每个区块：4 个时间位置 × 2 行 × 2 列 = 16 个 token
区块数量：(4/4) × (4/2) × (4/2) = 4

一个块跨越四个时间位置，在每个时间位置覆盖同一个 2×2 空间窗口。四个块分别命名为 A、B、C、D：左上、右上、左下、右下。

t = 0

0123456789101112131415

t = 1

16171819202122232425262728293031

t = 2

32333435363738394041424344454647

t = 3

48495051525354555657585960616263

A 块B 块C 块D 块

数字是原始 token 编号；同一种颜色跨四个时间切片构成一个 (4,2,2) 块。蓝色 A 共 16 个 token。文中图表都用教学数据，编号直接显示在格子中，颜色表示区块归属。

例如 A 的范围为 t=0…3、y=0…1、x=0…1。按 x 最快、y 次之、t 最慢展平，原始编号为：

index(t,y,x) = t × H × W + y × W + x

A 在 t=0： 0,  1,  4,  5
A 在 t=1：16, 17, 20, 21
A 在 t=2：32, 33, 36, 37
A 在 t=3：48, 49, 52, 53

所以，“三维相邻的一块”在原始序列中通常分散在多个位置。直接把序列每 16 个连续 token 分成一组，会得到一个完整的时间切片，几何含义完全不同。

一般地，块大小为 (bT,bH,bW) 时，完整块容量 M=bT×bH×bW，块数为：

J=\left\lceil\frac{T}{b_T}\right\rceil\left\lceil\frac{H}{b_H}\right\rceil\left\lceil\frac{W}{b_W}\right\rceil

边缘块可能少于 M 个 token，要单独记录真实长度 nᵢ；这会同时影响池化分母和正式 attention 的有效 key 范围。

这里用 (4,2,2) 展示三维池化，便于逐个标注 token。它是教学几何；FastVideo H3 后端支持 64、128、256-token 块，具体配置见第 5 节。

2. 路由分数：分别平均 Q、K，再计算相关性

2.1 平均掉三个时空维度，保留 D 维特征

固定一个 batch 和一个 head，A 块的 Q 看成 [4,2,2,D]。池化只平均 T、H、W 三个维度：

1
2
3

# Q_A、K_A 都是 [4,2,2,D]
q_mean_A = Q_A.mean(dim=(0, 1, 2))  # [D]
k_mean_A = K_A.mean(dim=(0, 1, 2))  # [D]

另一种做法是先把块内三维坐标展平为 16 个 token，再对 token 维求均值，结果相同：

\begin{aligned}\bar{\mathbf q}_A&=\frac1{16}\sum_{a\in A}\mathbf q_a\\\bar{\mathbf k}_A&=\frac1{16}\sum_{a\in A}\mathbf k_a\end{aligned}

每块得到一个 D 维 Q 代表向量和一个 D 维 K 代表向量。head 之间不混合，D 维通道也不相互求平均。原始 Q/K 会继续保留。

输入	沿哪些维度平均	输出
A 块的 Q：`[4,2,2,D]`	T、H、W，共 16 个 token	一个 `[D]` 向量
B 块的 K：`[4,2,2,D]`	T、H、W，共 16 个 token	一个 `[D]` 向量
两个均值向量	点积，再除以 √D	一个路由分数

Q、K 各自池化完成后，再计算代表向量的相关性。



%%{init: {"flowchart": {"nodeSpacing": 24, "rankSpacing": 28, "padding": 12}, "themeVariables": {"fontSize": "14px"}}}%%
flowchart TB
    q["Q 块 A · 4×2×2×D"] -->|沿 T、H、W 求均值| qm["Q 代表向量 · D"]
    k["K 块 B · 4×2×2×D"] -->|沿 T、H、W 求均值| km["K 代表向量 · D"]
    qm --> dot["点积，再除以 √D"]
    km --> dot
    dot --> score["路由分数 E［A,B］"]

用 D=2 手算一次。假设 A 块在每个时间片内的四个 Q 向量都相同，B 块的 K 也按下表取值：

时间位置	A 内四个 Q 都为	B 内四个 K 都为
t=0	`[2,0]`	`[1,0]`
t=1	`[0,2]`	`[1,0]`
t=2	`[2,2]`	`[3,2]`
t=3	`[0,0]`	`[3,2]`

每行有四个 token，所以整个块的均值为：

\begin{aligned}\bar{\mathbf q}_A&=[1,1],\quad\bar{\mathbf k}_B=[2,1]\\E_{AB}&=\frac{\bar{\mathbf q}_A^{\mathsf T}\bar{\mathbf k}_B}{\sqrt D}=\frac3{\sqrt2}\end{aligned}

先分别压缩 Q、K，再用代表向量相乘，块内全部 token 对的点积就不必再算。

2.2 用 J×J 的矩阵决定区块连接

对所有区块重复上述过程：

1
2
3

按块组织的 Q/K：[B,h,J,M,D]
       块内平均 → [B,h,J,D]
       均值 Q @ 均值 K.T → E：[B,h,J,J]

E 就是上面的小 S。本例 J=4，路由矩阵只有 4×4；每个元素 E[i,j] 表示 query 块 i 与 key 块 j 的相关性。若 A 对四个 key blocks 的分数为：

1 2	A B C D Query A 0.2 3.1 −0.5 2.4

保留两个块，就选择 B、D。这是一组独立的 TopK 教学分数，与前面的 D=2 手算数据分开。

每个 head、每个 query block 独立选块。A 选 B，并不要求 B 也选 A；同一个块内的 query 共用一份选择，但后续每个 query 的 token 权重分别计算。

只做 TopK 时可以直接对 logits 排序，softmax 不改变同一候选集合内的大小关系。若某视频区域有 Jᵣ 块，舍弃比例为 sᵣ，该 H3 实现的保留预算为：

R_r=\max\!\left(1,\min\!\left(\left\lceil(1-s_r)J_r\right\rceil,J_r\right)\right)

例如 256 块、舍弃 75%，保留 64 块；块数较少时，取整会使实际比例发生变化。预算与选块源码

2.3 均值为什么有效，又可能漏掉什么

由点积的线性性质：

\bar{\mathbf q}_i^{\mathsf T}\bar{\mathbf k}_j=\frac1{n_i n_j}\sum_{a\in i}\sum_{b\in j}\mathbf q_a^{\mathsf T}\mathbf k_b

它刻画的是两个区块间的平均匹配。相邻时空位置的特征如果较接近，一份共享路由就可能代表整块的主要需求。

但平均会稀释局部强匹配。例如一个 query 块中的两个一维向量为 [10,-10]，均值为 0；对 keys [10,10]，路由分数为 0，第一个 query 的两个 token 点积却都是 100。块内不同 token 的需求可能被均值掩盖，这种偏差要靠训练适配和质量验证来弥补。

3. pack：让同一块的 token 在内存中连续

三维分块规定“哪些 token 属于同一块”，pack 则把这些成员实际排到一起，方便 kernel 读取。回到前面同一组编号：

原始顺序

0123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263

按区块 pack 后

0145161720213233363748495253236718192223343538395051545589121324252829404144455657606110111415262730314243464758596263

两幅图均按行读成一维序列，每行展示 8 个位置。pack 后每两行组成一个 16-token 区块；悬停可查看序列位置和原 token 编号。

例如 A 会变成 packed 序列的前 16 个位置：

1 2	packed 位置： 0 1 2 3 4 5 6 7 … 12 13 14 15 原 token 号： 0 1 4 5 16 17 20 21 … 48 49 52 53

随后是 B 的 16 个，再是 C、D。Q/K/V 使用同一张位置映射；若位置编码已进入 Q/K，重排移动的是已有向量，不会为 packed 位置重新分配原始时空坐标。

池化和 pack 在数学上都可以直接按几何索引完成。工程里常先 pack 成 [J,M,D]，再沿 M 求均值，让路由和主计算共用一份布局。所以按几何索引理解选块规则就够了；实现里先重排再选块，只是为了复用同一份布局。

输出必须按逆映射放回去。例如 packed 位置 2 存放原 token 4 的结果，最终就写回输出序列的第 4 个位置。padding 输出被丢弃，所有有效 query 都有输出。

一张带编号的 mask：只看 A → B

先暂时忽略其他连接，只看“A 块的 query 访问 B 块的 key”。从三维分块图中放大 t=0 的前两行：

空间中的第一行：0(A)  1(A)  2(B)  3(B)
空间中的第二行：4(A)  5(A)  6(B)  7(B)

A 的 queries：0、1、4、5
B 的 keys：   2、3、6、7

mask 的每个绿色格子表示一条允许计算的 query–key 连接。例如 query 4 与 key 2 的交叉格是绿色。

重排前

q∖k	0	1	2	3	4	5	6	7
0
1
2
3
4			★
5
6
7

重排后

q∖k	0	1	4	5	2	3	6	7
0
1
4					★
5
2
3
6
7

行是 query，列是 key；★ 始终标记 q₄ → k₂。

两图的轴标签都保留原 token 编号。重排后把 A 的编号排到一起、B 的编号排到一起，黄色星号始终表示 q₄ → k₂。这只是所选 8 个 token 的局部视图，灰色表示本图未展示的连接，不能据此推断完整 TopK mask。

以 q₄ → k₂ 为例。重排前它在行 4、列 2，属于 A→B，保留。重排后编号换了位置，按轴上的原编号找到行 4、列 2，交叉格仍是绿色。把 A 的所有行、B 的所有列这样重排一遍，原来分散的绿色格子就拼成了一个方块。

上图是一个时间片内的 4×4 条连接。完整 A、B 各有 16 个 token，A→B 一共保留 16×16=256 条连接，包括跨时间位置的配对。pack 让完整矩阵中的这 256 格连续，既没有增加连接，也没有删掉 query。

4. TopK 之后：把选块结果交给稀疏 kernel

回到 A 选中 B、D 的例子。区块编号 A、B、C、D 分别是 0、1、2、3，A 的 block mask 为：

1 2	A B C D Query A 0 1 0 1

一个布尔元素代表整对区块。上层将它转换为 KV 区块索引和有效数量：

1 2	q2k_idx[A] = [1,3,−1,−1] q2k_num[A] = 2

该版本的 map_to_index 按候选块总数分配行宽，按块编号递增填入有效项，剩余位置为 −1。所以“选中两块”和“索引表分配四列”能同时成立；kernel 只读前两项。转换源码

绿色柱对应保留的 B、D。pooled Q/K 完成选路后，主计算读取这两块的原始 K/V。

对完整区块，A→B 的计算为：

1	Q_A: [16,D] @ K_B.T: [D,16] → logits: [16,16]

A 的每个 query 都可以访问 B 的全部 16 个 keys。同理，再访问 D 的 16 个 keys，每行总共看到 32 个 keys，输出 shape 仍为 [16,D]。

块内共享的是可见 key 集合，各 token 的 logits 和权重由原始 Q/K 单独计算。路由分数 E[A,B] 不会直接成为整个 16×16 子矩阵的权重。

5. H3 的 sequence 布局：从多模态序列到块序列

FastH3 是基于 MiniMax-H3 的音视频生成模型，它的 attention 同时处理文本、条件图像、音频和视频。读 VSA kernel 之前先分清三点：模态在原始 sequence 中的顺序、token 被重排到哪个块、张量的维度排列。

5.1 原始 packed sequence：多种模态共同做 attention

常见输入是 [文本 | 条件关键帧 | 音频 | 生成视频]。带参考视频的任务会多出视频区域，具体顺序由调用方传入的 segments 决定。H3 模型用 text_indices、video_indices、audio_indices 把各模态 embedding 写入同一条序列，并携带各 token 的位置和模态信息。

这层 attention 是联合双向的：视频 query 能看音频、文本，非视频 query 也能看视频，没有语言模型解码时的 causal 三角限制。模型组装入口

5.2 VSA tiled sequence：dense prefix 在前，视频块在后

metadata builder 接收有序的 segment 描述：整数表示一个 dense segment 的 token 数，三元组表示一段视频的 (T,H,W) 网格。随后执行两类分块：

文本、音频、参考图片按一维连续 token 切块；不同 segment 不拼进同一个块，末块单独 padding。
参考视频和目标视频分别按三维窗口切块，再把块内成员排成连续序列；每段视频保留独立的块范围。

所有 dense chunks 排到前面，形成 prefix；视频区域依次跟在后面，目标视频是最后一个区域。prefix 在这里只指计算布局中的前缀，不表示“只包含文本”或“只能被单向访问”。

下面用一个含参考视频的 segment 顺序作示例，块容量取 128。图中上半部分是原始顺序，下半部分展开成 kernel 读取的块顺序：



%%{init: {"flowchart": {"nodeSpacing": 24, "rankSpacing": 28, "padding": 12}, "themeVariables": {"fontSize": "14px"}}}%%
flowchart TB
    orig["原始 sequence<br/>文本 → 参考视频 → 图片<br/>→ 音频 → 目标视频"]
    orig -->|按 segment 分块并重排| prefix["dense prefix · 块 0–4<br/>文本 2 块 · 图片 1 块<br/>音频 2 块"]
    prefix --> ref["参考视频 · 块 5–8<br/>4 个三维块"]
    ref --> target["目标视频 · 块 9–16<br/>8 个三维块"]
    target -.-> partner["伙伴块 · 块 17<br/>有效长度 0<br/>仅用于双块调度"]

按这组长度算出的 metadata：

字段	示例值	含义
`num_prefix_tiles`	5	文本 2 块、图片 1 块、音频 2 块
`video_tile_spans`	`[(5,9),(9,17)]`	参考视频与目标视频的半开块区间
`variable_block_sizes`	`[128,2,60,128,72,128,…,128]`	17 个逻辑块各自的有效 token 数
原始有效长度	1,926	所有 segment 的真实 token 数之和
逻辑 tiled 长度	2,176	17×128，包含块内 padding
native kernel 长度	2,304	增加一个伙伴块，18×128

两种 padding 要分开看。块内 padding 补足每块的容量，池化时除以真实长度，主计算时屏蔽无效 keys；伙伴块用于原生 kernel 的双 Q 块调度，有效长度为 0、mask 行列全为 false，不进入逻辑路由与输出。模型为并行或容量预留的尾部空行在分块之前就已排除，不应再算成视频 token。

输出通过逆映射回到原始 sequence；例如原来位于音频 segment 的 token，计算完成后仍回到它的音频位置。布局与逆映射

5.3 从 attention 矩阵看普通生成与 R2V

先看不含参考视频的情况，sequence 为 [Context | OutputVideo]。Context 指文本、参考图片、音频等 dense prefix；横轴是 key，纵轴是 query。

Context query 看全局，OutputVideo query 保留全部 Context keys，只对 Video keys 做 sparse attention，因此只有矩阵右下角需要 TopK：

Key → 每个格子代表一对 block

Query ↓		Context		Output Video
Query ↓		C0	C1	O0	O1	O2	O3	O4	O5
Context	C0	●	●	●	●	●	●	●	●
Context	C1	●	●	●	●	●	●	●	●
Output Video	O0	●	●	●	·	●	·	·	·
	O1 ▸	●	●	·	●	·	·	●	·
	O2	●	●	●	·	·	·	·	●
	O3	●	●	·	●	·	●	·	·
	O4	●	●	·	·	●	·	·	●
	O5	●	●	●	·	·	·	●	·

● 全量保留● TopK 保留· 跳过

普通生成：Context 看全局，OutputVideo 看全部 Context，并在 Video keys 中选 TopK。此处每行保留 2/6 个视频块；选择位置为教学示例。

标有 ▸ 的 O1 行保留了 C0、C1，以及小 S 中分数最高的两个视频块 O1、O4。每个实心格内部仍执行完整的 token 级 QK、softmax、PV，空格对应的块计算被跳过。mask 可以不对称，也没有 causal 三角限制。

R2V：每段参考视频都是一个独立的候选组

R2V（Reference-to-Video）把参考视频作为条件输入。假设有两段参考视频，tiled sequence 为 [Context | RefVideo 1 | RefVideo 2 | OutputVideo]。所有视频 query 块共用一个小 S，但选块沿 key 轴分组，在每组内部单独取 TopK：

对一个视频 query 块 i：
保留的 KV 块 = 全部 Context
            ∪ TopK(S_small[i, RefVideo 1], k₁)
            ∪ TopK(S_small[i, RefVideo 2], k₂)
            ∪ TopK(S_small[i, OutputVideo], kₒ)

Key → 每个格子代表一对 block

Query ↓		Context		Ref 1			Ref 2			Output Video
Query ↓		C0	C1	R0	R1	R2	S0	S1	S2	O0	O1	O2	O3
Context	C0	●	●	●	●	●	●	●	●	●	●	●	●
Context	C1	●	●	●	●	●	●	●	●	●	●	●	●
Ref 1	R0	●	●	●	·	·	·	●	·	●	·	●	·
	R1	●	●	·	●	·	·	·	●	·	●	·	●
	R2	●	●	·	·	●	●	·	·	●	·	●	·
Ref 2	S0	●	●	●	·	·	·	●	·	·	●	·	●
	S1	●	●	·	●	·	·	·	●	●	·	●	·
	S2	●	●	·	·	●	●	·	·	·	●	·	●
Output Video	O0	●	●	●	·	·	·	●	·	●	·	●	·
	O1 ▸	●	●	·	●	·	·	·	●	●	·	·	●
	O2	●	●	·	·	●	●	·	·	●	·	●	·
	O3	●	●	●	·	·	·	●	·	·	●	·	●

● 全量保留● TopK 保留· 跳过

R2V 分组路由：每个视频 query 在 Ref 1、Ref 2、OutputVideo 中分别取 TopK。图中预算分别为 1/3、1/3、2/4；不同区域可以配置不同保留比例。

O1 行保留全部 Context，再从 Ref 1 选 R1、从 Ref 2 选 S2、从输出视频选 O0 和 O3，共 6 个 KV 块。参考组的配额独立保留，即使输出视频的分数整体更高，也不会挤掉参考组的名额。选完后，正式 attention 在这些保留 keys 的并集上做一次 softmax 归一化，不是各组分别归一化再相加。

参考图片和参考视频要分开：参考图片属于 dense prefix，参考视频属于视频区域。H3 默认的 exempt 实现里，参考视频 query 和输出视频 query 都按上面的分组规则路由，只有非视频 Context 行保持全局 dense。所以“只有 OutputVideo 行稀疏”说的是上一张没有参考视频 query 的图。分组 mask 源码

套用第 5.2 节的长度示例，若参考视频的 4 块、目标视频的 8 块都保留 25%，每个视频 query 访问 5 + 1 + 2 = 8 块，Context query 访问全部 17 块。q2k_num 因此随 query 类型变化，video_tile_spans 也必须保留区域边界。另一种 compete 策略会让 prefix keys 参与竞争，下面都按默认的 exempt 规则讲。

5.4 BSHD 与 BHSD：维度交换不等于重新分块

上层常用 [B,L,h,D]（BSHD），原生 Blackwell 路径接收连续的 [B,h,L,D]（BHSD）。transpose(1,2).contiguous() 改变维度排列，使一个 head 的整段序列连续；L 轴里的块顺序在更早的阶段就已确定。具体执行路径可以融合 scatter 与布局转换，但语义相同。

三维块几何与 GPU kernel 容量也要对应：

配置	`(bT,bH,bW)`	块容量 M	作用
开头的小例子	`(4,2,2)`	16	展示几何、池化和 mask
H3 tile64	`(4,4,4)`	64	64-token 稀疏路径
H3 tile128	`(4,4,8)`	128	下节展开的原生 kernel
H3 tile256	`(4,8,8)`	256	FA4 CuTe / Triton 路径

下面只看 ptx_blackwell 的 BF16、D=128、tile128 前向，对应数据中心 Blackwell。64、128、256 是三套不同的执行配置，不能把其中一条的 warp 数和缓冲参数推广到所有 VSA 实现。后端选择

6. FastH3 的稀疏 kernel：拿到索引之后做什么

6.1 kernel 外完成路由，kernel 内完成细粒度 attention

“kernel 外”指主 block-sparse attention CUDA kernel 之外。这些工作仍可能由 GPU 上的 PyTorch 算子、Triton kernel 或编译后的融合算子完成，不等于在 CPU 上算。

阶段	执行位置	产物
Q/K/V 投影、QK norm、位置编码	attention 主 kernel 外	带位置语义的 Q/K/V
segment 分块、索引和有效长度	metadata 构建阶段，可缓存几何	分块映射、区域范围、真实长度
scatter / pack、补齐、布局整理	主 kernel 外	连续 BHSD Q/K/V
Q/K 块内均值、相关性、区域 TopK	主 kernel 外	布尔 block mask
`map_to_index`	外部 Triton kernel	`q2k_idx`、`q2k_num`
TMA 描述符、输出分配、launch	C++ host wrapper	GPU 启动参数
按索引读 KV、QK、softmax、PV、写回	主 CUDA kernel 内	tiled 输出 O，可选 LSE
coarse 分支、gate 合并、逆重排	主 kernel 外	原始顺序的输出

主 kernel 不接收视频的 T/H/W，也不在里面求均值或做 TopK。它只知道“当前 Q 块应读取哪些 KV 块、每块多少有效 token”。这样，文本前缀、参考视频和目标视频的策略都能由上层通过同一套索引表达。

6.2 输入参数：数据、索引、长度分开传

Python 包装层的调用是：

out, lse = block_sparse_attn_ptx_blackwell(
    q, k, v,
    q2k_idx, q2k_num,
    variable_block_sizes,
    need_lse=False,
)

记补齐后的块数为 Jₑ，L=M×Jₑ；tile128 时 M=D=128：

参数	Shape / 类型	kernel 用它做什么
`q, k, v`	各 `[B,h,L,128]`，连续 BF16	原始 token 向量，按块读取
`q2k_idx`	`[B,h,Jₑ,Kmax]`，int32	每个 Q 块的 KV 块号列表
`q2k_num`	`[B,h,Jₑ]`，int32	当前行有多少有效块号
`variable_block_sizes`	`[Jₑ]`，int32	每个 KV 块有效的 key 数
`sm_scale`	浮点数，包装层设为 1/√D	QK 的缩放系数
`need_lse`	布尔值	是否输出反向计算所需的归一化统计量
返回 `out`	与 Q 相同	每个 query 的 D 维输出
返回 `lse`	可选 `[B,h,L]`，FP32	此实现采用 log₂ 约定的 LSE

map_to_index 分配的 Kmax 等于候选块数，不等于 TopK 保留数。上节例子中，补齐后的索引 shape 是 [B,h,18,18]；prefix 行有 17 个有效项，视频行有 8 个，伙伴行有 0 个。每行剩余位置填 −1，kernel 不把它们当成地址。

C++ wrapper 把张量指针、shape、索引指针装进 BlockSparseVsaArgs，再为 Q/K/V/O 创建 TMA tensor map，描述基址、维度、stride、块形状和 shared-memory swizzle，让搬运单元按坐标取块。描述符不复制 Q/K/V。包装层还把缩放系数转成 scale_log2 = sm_scale × log₂(e)，供 exp2 实现 softmax。

接口保留了可选 v_t 参数，但这条调用传入 None，当前实现直接读取自然布局的 V，不要求预先存一份转置 V。Python 包装、C++ 参数与 launch

6.3 双 Q 工作项：各自取索引，共享流水线轮数

这个 kernel 的工作项是 (batch, head, Q块对)：一个 CTA 同时处理相邻两个 Q 块，各有自己的索引、有效数量、softmax 状态与输出。源码直接定义 M_TILES_PER_CTA=2，并取 mtile0=2*p、mtile1=2*p+1；这是实际调度粒度。同一 CTA 还可以继续领取后续工作项，并非处理一对 Q 后必然退出。

下面仅用单 Q 块伪代码表示计算语义；真实 kernel 由多组 warp 配合，按双块流水线推进：

Q_i = load_query_block(i)
m, l, u = init_online_softmax()
for p in range(q2k_num[i]):
    j = q2k_idx[i, p]
    K_j, V_j = load_kv_block(j)
    S = Q_i @ K_j.T / sqrt(D)
    S[:, variable_block_sizes[j]:] = -inf
    m, l, u = online_update(m, l, u, S, V_j)
O_i = u / l

和 dense 循环的差别在读取地址这一步：j 来自稀疏索引，没选中的 KV 块通常不会进入主循环。所有 Q 块共用全局 K/V，不为每个 Q 复制一套数据。

双块流水线还有一个边界：两行索引数量不同，CTA 的公共循环次数取两者较大值，且至少为 1。短行的额外轮次使用安全的重复地址，随后全部 mask 掉，以保持生产者和消费者的 barrier 次数一致。

例如上节 Q 块 4 是音频 prefix，保留 17 块；Q 块 5 属于参考视频，只保留 8 块。这一对按 17 轮调度，短行多出的 9 轮不贡献输出。两行共享的是调度节奏，不是 mask。所以数学上少算的配对，不能直接换算成同倍率的实测加速。工作项与计数

6.4 Warp 分工：区分指令提交、指数权重和输出累加

以下假设读者熟悉 FA/FA2 与 online softmax，直接看 tile128、D=128 路径的线程与存储分工。每个 CTA 有 16 个 warp、512 个线程；图中的数字是 warp 编号。

0Q₀ ·

\widetilde{P}

/ m / ℓ

1Q₀ ·

\widetilde{P}

/ m / ℓ

2Q₀ ·

\widetilde{P}

/ m / ℓ

3Q₀ ·

\widetilde{P}

/ m / ℓ

4Q₁ ·

\widetilde{P}

/ m / ℓ

5Q₁ ·

\widetilde{P}

/ m / ℓ

6Q₁ ·

\widetilde{P}

/ m / ℓ

7Q₁ ·

\widetilde{P}

/ m / ℓ

8U 缩放 / 归一化

9U 缩放 / 归一化

10U 缩放 / 归一化

11U 缩放 / 归一化

12提交 QK /

\widetilde{P}\,V

13TMA 写回 O

14索引 / TMA 加载

15CLC 工作调度

0–3 服务 Q₀，4–7 服务 Q₁；8–11 依次处理两个 Q 块的输出状态。

\widetilde{P}

是未归一化的指数权重，U 是输出分子。

Warp	实际工作	主要读写位置
0–3、4–7	mask、更新参考值 m、生成 $\widetilde{P}$ 、α 和行和 ℓ；可选输出 LSE	TMEM 的 S → 寄存器； $\widetilde{P}$ 写回 TMEM，α/ℓ 写入 shared memory
8–11	按 α 修正旧 U；结束时计算 O=U/ℓ，转换为 BF16	TMEM 的 U ↔ 寄存器；最终 O 写入 shared memory
12	管理 QK/ $\widetilde{P}\,V$ 指令序列、等待依赖并提交异步 MMA	通过描述符和 TMEM 地址指定操作数，由 Tensor Core 执行
13	提交 O 的 TMA store，确认源缓冲可复用	shared memory → 全局内存
14	读取两行 KV 索引，计算地址，提交 Q/K/V 的 TMA load	全局内存 → shared memory
15	发出 CLC 请求，发布后续工作项编号	shared memory 中的 CLC 响应与同步记录

Shared memory 放 Q、K/V 环形缓冲和待写出的 O；Tensor Memory（TMEM）放 S/ $\widetilde{P}$ 与输出累加器 U。TMEM 是 Blackwell 的专用片上存储，不是显存或普通 shared memory。S 被读出后，其 TMEM 区域会复用来存当前 tile 的 BF16 $\widetilde{P}$ ，没有完整 N×N 矩阵的显存落地。

下图表示一轮稳态更新的数据依赖，不表示所有阶段必须串行执行；两个 Q 块及相邻轮次可以交错。



%%{init: {"flowchart": {"nodeSpacing": 24, "rankSpacing": 28, "padding": 12}, "themeVariables": {"fontSize": "14px"}}}%%
flowchart TB
    load["Warp 14 · TMA<br/>Q / K / V → shared memory"] --> qk["Warp 12 提交 QK<br/>Tensor Core 写 S → TMEM"]
    qk --> online["Warp 0–7<br/>mask、指数、m / ℓ"]
    online -->|α| corr["Warp 8–11<br/>U ← αU"]
    online -->|指数权重 分段写入 TMEM| pv["Warp 12 提交 权重 × V<br/>Tensor Core 累加到 U"]
    corr -->|旧 U 已修正| pv
    pv -->|最后一轮完成| norm["Warp 8–11<br/>O = U / ℓ → shared memory"]
    norm --> store["Warp 13 · TMA<br/>O → 全局内存"]

6.5 一轮计算中，各组 warp 究竟做了什么

Warp 0–7：每线程一行，生成指数权重与归一化统计

四个 warp 共 128 个线程，正好覆盖一个 Q 块的 128 行。源码中 m_tile = warp_id < 4 ? 0 : 1，sm_tid = (warp_id & 3) * 32 + lane。每个线程负责一个 query 行，在每轮读取该行的 128 个 logits，在寄存器中完成 mask、行最大值与行和计算；不是四个 warp 合作归约同一行。

它先等待 full_bar_spo，通过 tcgen05.ld 把 S 从 TMEM 读到 s_regs[128]。这里 S 是尚未乘 sm_scale 的 QK 点积。用源码的缩放约定，令 γ=sm_scale × log₂(e)，一轮状态更新可写为：

\begin{aligned} \alpha&=2^{(m-m')\gamma},\qquad \widetilde P=2^{(S-m')\gamma}\\ \ell'&=\alpha\ell+\operatorname{rowsum}(\widetilde P)\\ U'&=\alpha U+\widetilde P V,\qquad O=U/\ell\\ \mathrm{LSE}_2&=m\gamma+\log_2\ell \end{aligned}

0–7 负责前两行中的指数与行统计，不负责 U 的缩放或 $\widetilde{P}\,V$ 矩阵乘。 $\widetilde{P}$ 没有除以 ℓ。线程在 FP32 中累加行和，把 $\widetilde{P}$ 转成 BF16 写入 TMEM，供后续 MMA 消费。当前配置还混用硬件 exp2 与软件近似指数路径，以调整指数计算的执行压力。指数与权重写出

源码的 m′ 也不一定每轮都更新为最新最大值：先求候选最大值，再检查旧参考值到新值的变化；若 acc_scale >= -8，沿用旧 m、令 α=1。随后 $\widetilde{P}$ 和 ℓ 都按这个参考值计算，省去部分输出缩放。空 key 和短行补轮通过有效长度 mask 排除；所有保留的 prefix/video keys 共用同一行统计。

LSE 是循环结束后的可选输出：tile128 由这组 warp 写出 m_run * scale_log2 + log2(l_run)。need_lse=False 只省掉 LSE 写出， $\widetilde{P}$ 、m、ℓ 的计算仍然需要。指数近似、BF16 舍入与参考值更新策略属于数值实现细节，不能据上述公式宣称与理想实数计算逐位一致。

Warp 8–11：修正 U，再完成最终归一化

这四个 warp 共 128 个线程，依次服务 Q₀、Q₁ 的 128 行。等待 full_bar_alpha 后，从 shared memory 读取 α；若一个 warp 内全部 α=1，就跳过该 warp 对应行的缩放，否则用 tcgen05.ld/st 读写 TMEM 中的 FP32 U，执行 U ← αU。

每轮由 Tensor Core 累加的是 $\widetilde{P}\,V$ ，8–11 负责的是累加器的尺度修正。最后一轮，它们等待 full_bar_o_acc 与 full_bar_l，取最终 ℓ，用倒数乘法完成 O=U/ℓ，转换为 BF16 并写入 shared-memory 输出缓冲，再通知 epilogue。全部 keys 被屏蔽的行使用零缩放因子，输出为零。correction 与最终归一化

Warp 12：提交指令，由 Tensor Core 计算

elect_one_sync() 从 MMA warp 中选出一个线程，由它发射 tcgen05.mma。QK 读取 shared memory 中的 Q/K，结果写入 TMEM 的 S； $\widetilde{P}\,V$ 读取 TMEM 的 $\widetilde{P}$ 和 shared memory 的 V，累加到 TMEM 的 U。一个 tile 由多条 MMA 指令完成，warp 还要推进描述符、等待 barrier，并提交完成通知。

这里的单线程是指令发射粒度，矩阵乘加由 Tensor Core 硬件执行；warp 0–7 则要逐项执行指数、比较、求和与数据转换。因此不能用“softmax 8 个 warp、MMA 1 个 warp”比较两部分的计算量。MMA 指令序列、tcgen05 发射语义

SPLIT_P 与 barrier：权重就绪和累加器可写是两件事

当前配置启用 SPLIT_P：每行 128 个 BF16 权重打包为 64 个 32-bit 单元。0–7 先写前 48 个单元，即 96 列 $\widetilde{P}$ ，通知 empty_bar_spo；再写剩余 16 个单元，即 32 列，通知 full_bar_p_last。MMA 可以先消费前一段，推进到最后 32 列时才等待后一段就绪。

稳态下的 empty_bar_spo 不是只由 softmax 一方触发：它汇合该 Q 块的 128 个指数/统计线程与 128 个 correction 线程。这同时保证“ $\widetilde{P}$ 前段已写好”和“旧 U 的缩放已完成”，随后才能安全开始本轮 $\widetilde{P}\,V$ 。首轮通过初始化提供 correction 侧的就绪信号，因为还没有旧 U 需要缩放。

同步点	约束的数据依赖
`full_bar_spo`	QK 结果已可供 0–7 读取
`full_bar_alpha`	本轮缩放信息已发布，8–11 可以修正 U
`empty_bar_spo`	$\widetilde{P}$ 前段与 U 均已准备好，允许 MMA 继续
`full_bar_p_last`	$\widetilde{P}$ 尾段已准备好，允许完成剩余 PV
`full_bar_o_acc` + `full_bar_l`	最终累加器和分母都已就绪，允许归一化

这些 barrier 随轮次切换 phase；共享位置可以反复使用，而不会把上一轮的完成信号当成下一轮。分段写入、PV 尾段等待

Warp 14、13、15：加载、写回与跨工作项衔接

Load warp 14 为 Q₀、Q₁ 分别读索引；每次由 32 个 lanes 预读一个索引窗口，再用 shuffle 取得当前块号。它等待空槽后提交 TMA load，硬件完成搬运后满足 full barrier，MMA 才能读取。tile128 的 KV ring 是 3 个 32 KiB 槽，分别承接 K 或 V 搬运，并非三份完整 K/V。Q 使用独立缓冲；短行越界轮次夹取最后一个有效索引，空行取安全块号 0，再由 0–7 全部 mask 掉。

Epilogue warp 13 等待 O 在 shared memory 中就绪后提交 TMA store，并等待源数据已被读取，再释放对应输出槽位。它不重新计算归一化。

Scheduler warp 15 通过 CLC（Cluster Launch Control）尝试取消同一次 launch 中尚未执行的 CTA，并取得其工作编号。各角色消费同一份 CLC 响应，更新局部 workitem_id，继续处理下一对 Q；原 CTA 的 blockIdx 不变，也没有再次请求 CPU launch。缓冲区的 full/empty barrier 约束前后工作项的重叠，而非要求所有 warp 同步退出再重启。加载与工作项推进、写回与 CLC 调度

6.6 相比 dense kernel，收益从哪里来

与现代 dense FlashAttention 相比，新增收益主要来自缩短 KV 遍历集合。TMA、Tensor Core、在线 softmax、warp 分工本身也可以用于 dense kernel，不能把它们全算成 VSA 独有的收益。

层面	dense attention	此稀疏路径
KV 遍历	对双向 attention 遍历全部 key blocks	按每个 Q 块的索引列表读取
QK / PV	每个候选块都做矩阵乘	跳过未选块的主计算
KV 搬运	为全部候选块搬运数据	通常随选中块减少，双 Q 块补轮有额外开销
中间矩阵	FlashAttention 已避免把完整 N×N 矩阵写回显存	延续相同的融合思路
负载	每行遍历长度较一致	prefix dense、视频 sparse；需要独立计数与调度
额外成本	无 VSA 路由	pool、TopK、pack、索引转换及 padding

CLC（Cluster Launch Control）让完成当前工作的 CTA 获取尚未执行的工作项，帮助处理长度不齐的工作负载；它不会消除同一 Q 块对内部的补轮。吞吐还受数据复用、带宽、占用率和边缘块比例影响，不能只看一个稀疏率。

6.7 coarse 分支也在主 kernel 外

FastH3 还可以通过 pooled V 形成粗粒度补偿：

\begin{aligned} C&=\operatorname{softmax}(E)\,\bar V\\ O&=O_{\mathrm{sparse}}+g\odot\operatorname{broadcast}(C) \end{aligned}

它对所有逻辑区块做 coarse attention，再把每个 Q 块的结果广播给块内 token，由学习到的 gate 控制加入多少。FastH3 的 VSA 蒸馏模型可携带训练好的 gate；基础 H3 没有该参数时，加载器将它初始化为零。此分支不改变主 kernel 的索引接口，也不保证恢复 dense 的 token 级细节。coarse 合并

7. 收益与代价：看配对数量，也看实际质量

对完整等大块，N=MJ，每个 query block 保留 R 个 key blocks。忽略 batch 和 head，参与计算的配对数：

计算	配对数量	开头的小例子：M=16、J=4、R=2
dense QK	N²	4,096
pooled 路由	J²	16
主稀疏 QK	J×R×M²	2,048

这个例子里，主 QK/PV 保留约一半工作量；路由点积数量是 dense 的 1/M²。这个比值只描述数学规模，实际耗时还包括 pool、TopK、pack、索引转换和访存。

块越大，路由矩阵越小，却会为局部重要 token 连带计算更多连接；块越小，选择更精细，但路由和调度成本增加。边缘块、dense prefix 也会改变实际保留比例，不能把块数比例直接当成端到端加速比。

质量同样要单独验证。若保留集合覆盖了某行 dense 概率的比例 c，令 O_kept、O_drop 分别为两组归一化后的输出，则在 0<c<1 时：

\begin{aligned}O_{\mathrm{dense}}&=cO_{\mathrm{kept}}+(1-c)O_{\mathrm{drop}}\\O_{\mathrm{sparse}}&=O_{\mathrm{kept}}\end{aligned}

误差既取决于舍弃的概率质量，也取决于两组 V 的差异。评估时要把这几项分开：路由覆盖率、attention 输出误差、完整视频质量，以及拆分耗时与端到端延迟。

8. 一个可运行的参考实现

把上面的顺序压成伪代码，Q/K/V 已按块排成 [J,M,D]：

# 用 (bT,bH,bW) 几何建立分块映射，再 pack
Qp, Kp, Vp, sizes, inverse = pack_3d(Q, K, V, grid, block)

# 每块分别平均 Q、K；D 不变
q_mean = mean_valid_tokens(Qp, sizes)  # [J,D]
k_mean = mean_valid_tokens(Kp, sizes)  # [J,D]
E = q_mean @ k_mean.T / sqrt(D)       # [J,J]
selected = topk(E, keep_blocks, dim=-1)

# 对选中的原始 K/V 做细粒度 attention
Op = block_sparse_attention(Qp, Kp, Vp, selected, sizes)
O = restore_original_order(Op, inverse)

下载 Python 教学实现，只依赖标准库：

1	python3 vsa-reference.py

程序先复现前面的 (4,4,4) 网格与 (4,2,2) 分块，打印 A 的成员编号；随后用 (5,3,4) 网格检验不完整边缘块。两组输入都验证“pack → 均值选块 → 在线 softmax → 逆重排”与原始顺序下的显式 masked attention 一致，并验证全保留时恢复为 dense。代码用于核对数学和索引，不模拟 GPU 性能。

想继续读源码，可以从这三个文件入手：

三维分块与逆映射：确认每个 token 属于哪个块、被放到哪里。
H3 路由与前向调用：跟踪 pooled Q/K、区域 TopK 和可选 coarse 分支。
Blackwell 稀疏计算核：观察每个 query block 如何按索引读取 KV，并维护自己的在线 softmax 状态。