执行摘要
- 一句话：Sparse MLA mask 向量化 128-bit 加载，B200 mask 开销降约 35%
- 推荐动作：值得精读，尤其是
dense_mask_mod 这段 cute 写法：assume + flat_divide + autovec_copy + recast_tensor 的组合是 CUTLASS 编程模型里做对齐向量加载的范式，可复用到其他 mask 或 packed 数据内核；“由 shape 函数统一保证对齐 + 测试固化不变量”的做法也值得借鉴。另一个关注点是 _build_topk_mask 的返回语义变化，合入后留意其他消费方。
功能与动机
PR body 的基准数据显示，sparse MLA 的 32-bit mask 加载在 FA4 内核中相对 unmasked baseline 带来 +53.7%~86.4% 的额外开销，是明显的性能瓶颈。作者将 mask_mod 改为一次加载 4 个 packed word（128-bit 向量），同时保留行内 padding 保证对齐，把开销压到 +17.8%~37.2%；由于只是换加载方式、不改变 mask 语义，作者声明无需模型评估，kernel 与 sparse-MLA 正确性测试即可覆盖。
实现拆解
- mask_mod 向量化（
vllm/model_executor/layers/attention/sparse_mla_mask.py）：重写 dense_mask_mod。旧实现每次 kv 位置做一次 32-bit word 加载加位移；新实现先用 cute.assume 断言 batch/query 方向 stride 可被 4 整除，再用 flat_divide 把 mask 行切成 4-word 块，autovec_copy 一次性把 128-bit 拷入寄存器，最后 recast_tensor 转成 4 个 Uint32 返回。__vec_size__ 由 32 改为 128，让上层按向量宽度展开。offset_dense_mask_mod 保持 32-bit 不变。
- 数据契约：行宽对齐（
vllm/model_executor/layers/attention/sparse_mla_attention.py）：_topk_mask_shape 在算完 num_words 后新增 triton.cdiv(num_words, 4) * 4，把每行 word 数补到 4 的倍数；_build_topk_mask 的返回值从 out[:B, :Q, :num_words] 改为 out[:B, :Q]，不在 view 层把 padding 切片掉，保证每行起始地址相对基址保持 128-bit 对齐；docstring 同步改为 “preserving padded row storage”。
- 对齐不变量的测试（
tests/v1/attention/test_sparse_mla_mask.py）：新增 test_topk_mask_rows_are_aligned_for_vectorized_loads（纯 CPU 校验 shape：129 个 key 对齐后仍为 8 words，含 reserve 场景也是 8）与 test_build_topk_mask_preserves_aligned_row_storage（CUDA 上构造 129 key 的 mask，断言结果为 8 个 word、bit 0 与 bit 128 位置正确、padding 保留）。
- 性能验证（PR 内基准）：B200、BF16、8 heads、QK dim 192、V dim 128、
num_splits=1、全 1 mask，三组 Q x KV 规模下 mask 相对开销由 32-bit 的 +56.2%/+86.4%/+53.7% 降到 128-bit 的 +21.1%/+37.2%/+17.8%；新旧 callback 输出逐位一致。
关键文件：
vllm/model_executor/layers/attention/sparse_mla_mask.py（模块 注意力内核；类别 source；类型 core-logic；符号 dense_mask_mod）: 核心 kernel 改动所在：dense_mask_mod 从 32-bit 逐字加载改为 128-bit 对齐向量加载，__vec_size__ 由 32 改为 128，是全部性能收益的来源。vllm/model_executor/layers/attention/sparse_mla_attention.py（模块 注意力层；类别 source；类型 data-contract；符号 _topk_mask_shape, _build_topk_mask）: 数据契约变更所在：_topk_mask_shape 将 num_words 对齐到 4 的倍数，_build_topk_mask 返回 view 时保留 padding，是向量化加载成立的前提。tests/v1/attention/test_sparse_mla_mask.py（模块 单元测试；类别 test；类型 test-coverage；符号 test_topk_mask_rows_are_aligned_for_vectorized_loads, test_build_topk_mask_preserves_aligned_row_storage）: 新增两个测试把“行宽 4 对齐”固化为不变量，覆盖非整 chunk 长度（129 key）场景，防止未来误伤对齐契约。
关键符号：dense_mask_mod, offset_dense_mask_mod, _topk_mask_shape, _build_topk_mask, test_topk_mask_rows_are_aligned_for_vectorized_loads, test_build_topk_mask_preserves_aligned_row_storage
关键源码片段
vllm/model_executor/layers/attention/sparse_mla_mask.py
核心 kernel 改动所在：dense_mask_mod 从 32-bit 逐字加载改为 128-bit 对齐向量加载，__vec_size__ 由 32 改为 128，是全部性能收益的来源。
@cute.jit
def dense_mask_mod(
batch: cute.TensorSSA,
head: cute.TensorSSA,
q_idx: cute.TensorSSA,
kv_idx: cute.TensorSSA,
seqlen_info,
aux_tensors: list,
) -> cute.TensorSSA:
dense_mask = aux_tensors[0]
batch_idx = utils.ssa_to_scalar(batch)
q_idx = utils.ssa_to_scalar(q_idx)
kv_idx = utils.ssa_to_scalar(kv_idx)
# 取 batch/query 两个方向的 stride，并用 assume(divby=4) 向编译器宣告
# 它们都是 4 的倍数：配合下方按 4 个 word 分组，保证 128-bit 对齐加载合法
batch_stride, query_stride, _ = dense_mask.stride
aligned_mask = cute.make_tensor(
dense_mask.iterator,
cute.make_layout(
dense_mask.shape,
stride=(
cute.assume(batch_stride, divby=4),
cute.assume(query_stride, divby=4),
1,
),
),
)
# 取当前 (batch, query) 对应的整行 mask，再按 (4,) 做 flat_divide，
# 把行切成若干 128-bit 块；kv_idx >> 5 是 32-bit word 索引，
# 再 >> 2 得到该 kv 所属的 4-word chunk 索引
mask_row = aligned_mask[batch_idx, q_idx, None]
mask_chunks = cute.flat_divide(mask_row, (4,))
mask_chunk = mask_chunks[None, (kv_idx >> 5) >> 2]
# autovec_copy 一次把 128-bit 拷入寄存器，recast_tensor 将其解释为
# 4 个 Uint32 后 load 返回，由上层按向量宽度消费，减少逐字加载与索引指令
loaded = cute.make_rmem_tensor_like(mask_chunk)
cute.autovec_copy(mask_chunk, loaded)
return cute.recast_tensor(loaded, cutlass.Uint32).load()
dense_mask_mod.__vec_size__ = 128
vllm/model_executor/layers/attention/sparse_mla_attention.py
数据契约变更所在：_topk_mask_shape 将 num_words 对齐到 4 的倍数，_build_topk_mask 返回 view 时保留 padding，是向量化加载成立的前提。
def _topk_mask_shape(
batch_size: int,
max_query_len: int,
max_key_len: int,
reserve_key_starts_word: bool = False,
) -> tuple[int, int, int]:
"""Shape of a bit-packed top-k mask, shared by every site that builds one."""
tile_m = 128 if max_query_len <= 128 else 256
padded_q_len = triton.cdiv(max_query_len, tile_m) * tile_m
num_words = triton.cdiv(max_key_len, 32) + int(reserve_key_starts_word)
# 关键对齐不变量：把每行 word 数补到 4 的倍数，使行长为 128-bit 的
# 整数倍，sparse_mla_mask.dense_mask_mod 才能安全做向量化加载
num_words = triton.cdiv(num_words, 4) * 4
return batch_size, padded_q_len, num_words
def _build_topk_mask(
topk_indices_per_req: list[torch.Tensor],
q_lens: list[int],
max_q_len: int,
max_seq_len: int,
out: torch.Tensor,
) -> torch.Tensor:
"""Build a bit-packed top-k mask while preserving padded row storage."""
batch_size = len(q_lens)
num_words = (max_seq_len + 31) // 32
total_rows = batch_size * max_q_len
if total_rows == 0:
# 空请求时同样保留完整二维 view，不裁剪 padding
return out[:batch_size, :max_q_len]
# 中间调用 _scatter_topk_single_req_kernel / _scatter_topk_kernel
# 用 atomic_or 把 top-k 索引对应的 bit 写入各 mask 行，这里省略
# 返回时保留 padding：不按 num_words 裁剪最后一维，
# 这样返回 view 的每行起始地址仍然 128-bit 对齐
return out[:batch_size, :max_q_len]
tests/v1/attention/test_sparse_mla_mask.py
新增两个测试把“行宽 4 对齐”固化为不变量，覆盖非整 chunk 长度（129 key）场景，防止未来误伤对齐契约。
def test_topk_mask_rows_are_aligned_for_vectorized_loads() -> None:
# 129 个 key 需要 5 个 32-bit word，对齐后行宽应为 8 个 word
assert _topk_mask_shape(2, 129, 129) == (2, 256, 8)
# reserve_key_starts_word 场景同样需要对齐到 8
assert _topk_mask_shape(2, 129, 129, reserve_key_starts_word=True) == (
2,
256,
8,
)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
def test_build_topk_mask_preserves_aligned_row_storage() -> None:
shape = _topk_mask_shape(1, 1, 129)
out = torch.zeros(shape, dtype=torch.int32, device="cuda")
topk = torch.tensor([[0, 128]], dtype=torch.int32, device="cuda")
mask = _build_topk_mask([topk], [1], 1, 129, out)
# 行宽保持 8：bit 0 与 bit 128 分别落在 word 0 与 word 4，
# 其余是保留的 padding，值为 0
assert mask.shape == (1, 1, 8)
assert mask[0, 0].tolist() == [1, 0, 0, 0, 1, 0, 0, 0]
评论区精华
review 中有两条值得记录的问答：
- simon-veitner-redhat 问 store 是否也能向量化：“is the store vectorised and if not can we vectorize the store aswell?”，作者回复这里没有真实内存写、会被 lower 成寄存器赋值，改用
return cute.recast_tensor(loaded, cutlass.Uint32).load() 性能相同但更可读，于是顺手用了这个写法。
-
simon-veitner-redhat 追问 Blackwell 上是否试过 256-bit 加载：“did we try to leverage 256 bit load for blackwell?”，作者回复在 hdim 256 kernel 与通用 kernel 上都没有性能提升，因此停留在 128-bit。
-
store 是否也能向量化 (performance): 不存在真实 store，改用 recast_tensor 提升可读性，性能不变。
- Blackwell 是否尝试 256-bit 加载 (performance): 256-bit 尝试后无收益，保持 128-bit 向量宽度。
风险与影响
- 风险：
- 返回形状变化：
_build_topk_mask 返回的 view 不再裁剪到 num_words，依赖 shape[-1] 与 key 长度直接换算的调用方需要核对；本 PR 未改普通 FA4 路径，影响面限定在 SM100 masked-MHA 路径。
- 对齐假设：
cute.assume(..., divby=4) 依赖所有 mask 分配都经过 _topk_mask_shape；若有调用方绕过该函数直接分配非对齐 buffer，会触发未定义行为。当前所有分配点均返回对齐 shape，测试也覆盖了 129 这类非整 chunk 长度。
- 平台限定：
_is_masked_mha_available 限定 SM100，其他 GPU 不受影响也不会受益；offset_dense_mask_mod 仍是 32-bit，与 dense_mask_mod 行为不一致但不影响正确性。
- 内存：padding 使每行最多多 3 个 word，workspace 适配检查基于
_topk_mask_shape 自动覆盖，总内存占用有极小增量。
- 影响：对用户：SM100（Blackwell）上 sparse MLA（如 DeepSeek top-k 稀疏注意力）的 mask 相关时间显著下降，大 Q x KV 场景尤其明显（相对开销最多降低约 50%），模型输出不变，无需重新评测。对团队：mask 行布局的数据契约更新，后续 kernel 若要再向量化需要保持同样的对齐约定；新增测试把“行宽 4 对齐”固化为不变量。影响范围窄，仅在 CUDA SM100 的 masked attention 内核路径。
- 风险标记：数据契约变更：mask 行保留 padding, 对齐假设依赖 divby=4, 仅 SM100 平台生效, offset_dense_mask_mod 未向量化
关联脉络
- PR #52775 [Kernel] SM120: stop routing misaligned-M blockwise FP8 GEMMs to the small-M swapAB config: 同为 SM1xx 平台内核级性能优化，反映 NVIDIA Blackwell 内核调优方向上的并行工作；虽非同一功能线，但在基准数据与内核路由策略上可互相参考。
参与讨论