Agent 帮我写 CUDA（续）：当 Self-Attention 遇上 Causal Mask 和 Cross-Attention

副标题：给我们的 CUDA Attention 算子加上因果掩码和交叉注意力——从"能跑"到"能用"

一、引子：上一篇的 kernel，能跑但不能用

上一篇我们走完了 Softmax Attention CUDA 算子的完整开发流程：从 Naive 到 Tiled，从 float4 向量化到 bank conflict padding，最终在 GTX 1660 Ti 上实测了四版 kernel 的性能。

但那个 kernel 有个致命的问题——它只能在最简单的场景下工作：

• Q/K/V 必须等长（N_q = N_kv）
• 没有 causal mask（每个位置可以看到所有位置，包括未来的）
• 只能做 self-attention，不能做 cross-attention

换句话说，它能在 MNIST 上跑，但不能在 transformer 里用。

现实中的 Attention 远比这复杂：

这一篇的目标：基于上篇性能最好的 VecPadded kernel，一次性支持上面所有场景。

二、Causal Mask：为什么需要，怎么实现

为什么 decoder 需要 causal mask？

自回归生成时，位置 i 的输出只能依赖位置 1…i 的输入——它"看不到"未来的 token。如果允许看到未来的 token，模型就相当于提前知道了答案，训练时 loss 会 collapse，生成时无法自回归。

实现上就是在 softmax 之前把未来位置设为 -inf：

Attention(Q, K, V) = softmax(QK^T / √D + Mask) · V

其中 Mask[i][j] = 0（j ≤ i）或 -∞（j > i）。

在 Tiled Kernel 中的实现

我们上篇的 kernel 在每个 K/V tile 里遍历全部 TILE_KV 个位置做 online softmax。加上 causal mask 后，每个 warp（处理一行 query）只需要处理该行可见的 K/V 位置。

改动只有计算循环中的两行：

# 在 tile 循环中，每个 warp 检查自己 query row 的可见范围ifis_causal:ifquery_row<kv_start:compute_end=0# 整个 tile 不可见，跳过else:local_end=query_row-kv_start+1compute_end=min(num_rows_in_tile,local_end)

K/V 的加载不受影响——所有线程照常从显存读取数据到 shared memory，因为不同 warp 可看到的 tile 可能不同，但 shared memory 加载需要全员参与。只有计算阶段跳过被 mask 的位置。

这带来一个有趣的副效应：causal mask 不仅正确，还免费加速了 kernel——平均一半的 K/V 位置被跳过，计算量直接减半。

三、Cross-Attention：当 Q 和 K/V 不再等长

Encoder-Decoder 架构的内存布局

Cross-attention 中，Q 来自 decoder，K/V 来自 encoder 的输出：

Q: [B, H, N_q, D] — decoder 侧，长度随生成步数变化 K: [B, H, N_kv, D] — encoder 侧，长度固定（输入序列长） V: [B, H, N_kv, D] — 同上 O: [B, H, N_q, D] — 输出，长度 = N_q

关键变化：Q 和 K/V 在batch × head 内的全局偏移量不同。之前的代码用一个head_base统一计算：

int64_t head_base = ((int64_t)batch * H + head) * N * D;

现在需要拆成两个：

int64_t q_head_base = ((int64_t)batch * H + head) * N_q * D; int64_t kv_head_base = ((int64_t)batch * H + head) * N_kv * D;

Grid 维度也从B × H × ceil(N / TILE_ROWS)变为B × H × ceil(N_q / TILE_ROWS)——因为要处理的 query 行数由 N_q 决定。

Self-Attention vs Cross-Attention 的代码统一

把这两个改动合到一起，最终 kernel 的签名变成了：

__global__ void vec_fused_attention_kernel( const float* Q, const float* K, const float* V, float* O, int B, int H, int N_q, int N_kv, int D, int is_causal );

这个签名统一了三种场景：

整个 kernel 从 vec_padded 到 vec_fused 的改动只有约 30 行：拆分 head_base、改 grid、改 Q/K/V 地址计算、加 causal check。模板代码和优化逻辑一个字没动。

四、Agent 在这步的角色：代码扩展 vs 从零写 kernel

上一篇我们展示了 Agent 从零写 CUDA kernel 的能力。这一篇的场景不同——不是"写新代码"，而是"修改已有代码"。

有意思的发现：Agent 在"代码扩展"任务上，比"从零写"更容易出错。

原因在于：

• 从零写：Agent 会调用训练数据中最熟悉的模式——online softmax、Tiling、float4 加载——这些在公开文档和代码里被反复讨论过
• 改已有代码：Agent 需要理解现有代码的结构（extern shared memory 的布局、__syncthreads()的位置约束、头文件宏的传递），然后在正确的位置插入改动。它不会"看到"你的代码——它只能根据你的描述做修改

这次的实际体验：我（人类）直接改了代码，Agent 做 review。

因为代码是我们自己一行行写（和上篇 Agent 一起）的，我很清楚哪里该改。跨注意力其实就是两个序列长度 + 两个 head_base，causal 就是在计算循环里加个判断。通过 Agent review 确认没有遗漏。

对"Agent 写代码"的补充思考

上篇我们提到 Agent 能一次写对 CUDA kernel，是因为 NVIDIA 的生态太好了——编程指南、官方 Sample、StackOverflow 百万问答都在训练数据里。

这篇的经验进一步说明了一个更微妙的点：

Agent 擅长的是"模式匹配"——在熟悉的语境下从零组装新代码。但它不太擅长"理解已有代码然后在正确的位置插入修改"。后者需要的不只是模式识别，还需要理解局部结构和全局约束的关系。

这恰好对应算子开发中的两类工作：

• 绿字段落（新架构、新算子）：Agent 很强，因为它见过相似的 pattern
• 棕字段落（扩展现有算子、修 bug）：Agent 会犯错，因为它不理解代码的"上下文"

所以一个人机协作的合理分工是：让 Agent 写新 kernel，让人来改已有代码——或者反过来，人改代码、Agent 做 review。

五、正确性与性能

正确性测试结果

对比 PyTorch 的F.scaled_dot_product_attention(q, k, v, is_causal=...)，4 个场景全 PASS：

手工验证 causal mask 的有效性：给 V 的最后一个位置赋一个极大值（999），position 0 的输出不受影响（因为它看不到），position 3 的输出被极大值支配（因为它看得到自己）。

性能：Fused 模式 vs Padded 模式

把 VecFused 在 non-causal self-attention 下和 VecPadded 做对比，确认没有引入额外开销：

差异在 ±2% 以内，属于正常波动。多传两个参数（N_kv、is_causal）和拆开 head_base 的计算开销完全可以忽略。

Causal Mask 带来的意外收益

由于 causal mask 跳过了一半的 K/V 计算位置，kernel 在大 N 时自动提速约45%：

这不是优化带来的——这是 causal mask 本身自带的"免费"收益：K/V 加载仍然需要做（因为 masked 掉的 warp 也需要参与 shared memory 加载的同步），但计算阶段跳过了一半的位置。

理论上如果只考虑计算，causal 应该省 50%。实测 45% 左右，差距来自 K/V 加载的开销。

顺带一提：FP16 版本

我们顺手把 fused kernel 改了个 FP16 I/O 版本——Q/K/V/O 从显存以__half精度读写，内部计算保留 float32。

结果出人意料：FP16 比 FP32 慢了 2-12%。

原因很直接，而且和我们上篇的核心结论一脉相承：

1. 这个 kernel 是 compute-bound，不是 memory-bound。省一半显存带宽没用——计算才是瓶颈
2. half→float 转换不是免费的。每个从显存读进来的 half 都要转成 float 才能做点积，写完再转回去。这些转换指令在 Turing 上大约占 5-10% 的指令周期
3. Turing 的 Tensor Core 不能用。FP16 如果走 Tensor Core 可以大幅加速计算路径。但在 1660 Ti 上手动写 Tensor Core 调用需要 WMMA API，复杂度远超普通 kernel

这个实验的价值是反面的：它用数据证明了"换低精度就能快"是一个普遍的误解。精度提升速度的前提是你的 kernel 是 memory-bound 的。一旦是 compute-bound，省带宽没意义，转换开销反而拖后腿。

所以，我们的 fused kernel 在 1660 Ti 上，FP32 就是最终性能上限了。要突破这个上限，需要的不是换精度，而是换计算方式——比如用 Flash Attention 的分块策略减少总计算量。

六、总结

我们做了什么

• 把之前只能做等长全可见 self-attention 的 kernel，扩展成了支持causal mask + cross-attention的统一版本
• 总代码改动量约30 行，核心思想极简：head_base 拆成两个，计算循环加个判断
• 验证了 4 种场景的正确性，确认了 non-causal 模式下零性能损失

一个可用的 Attention kernel 还差什么？

到现在为止，我们的 kernel 支持了标准 Attention 的所有核心功能。但和 cuDNN 的flash_attention比，还差三件事：

1. 分块稀疏（block-sparse）——Flash Attention 通过对 score 矩阵做分块稀疏化，避免了完整 N×N 的显存开销。这在大 N 时是关键优化
2. 内存布局优化——我们一直用[B, H, N, D]的 naive 布局。实际推理引擎会用 Interleaved 布局或 GQA/MQA 的分组策略
3. 更低精度支持——float16 / bfloat16 / FP8 的支持。在 1660 Ti 上没有 Tensor Core，但 float16 可以省一半带宽

这些已经超出了"手写 CUDA kernel"的范畴，进入了"和推理引擎深度集成"的领域——这正是flash_attn和vLLM等库在做的事。

但站在 Agent 辅助开发的角度看，这个演进路径是清晰的：从能跑的 naive → 能用的 tiled → 功能完整的 fused → 高性能的 flash——每一步 Agent 都可以参与，但每一步的参与形式不同。从"替人类写"到"帮人类改"再到"和人类一起调"，Agent 的角色越来越像一个平等的协作者，而不是抢活的工具。

上一篇我们说：Agent 把算子开发的门槛从"手写每一行"降到了"review 和决策"。这一篇我们看到，即使只做 review 和决策，这个门槛也比"从零写"高——因为它需要理解代码的结构和约束，而不只是调用熟悉的 pattern。

这恰恰是最健康的协作状态：Agent 做它擅长的（模式匹配、汇编），人做他擅长的（理解结构、决策取舍）。