副标题： 给我们的 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 远比这复杂：
| 场景 | 需求 | 对应模型结构 |
|---|---|---|
| Decoder self-attention | Causal mask（位置 i 只能看 j ≤ i） | GPT、LLaMA、Qwen（生成阶段） |
| Encoder self-attention | 全可见，Q=K=V 等长 | BERT、Encoder-only |
| Encoder-Decoder cross-attention | Q 来自 decoder，K/V 来自 encoder，不等长 | T5、机器翻译、语音识别 |
这一篇的目标：基于上篇性能最好的 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其中 Maskij = 0（j ≤ i）或 -∞（j > i）。
在 Tiled Kernel 中的实现
我们上篇的 kernel 在每个 K/V tile 里遍历全部 TILE_KV 个位置做 online softmax。加上 causal mask 后，每个 warp（处理一行 query）只需要处理该行可见的 K/V 位置。
改动只有计算循环中的两行：
python
# 在 tile 循环中，每个 warp 检查自己 query row 的可见范围
if is_causal:
if query_row < kv_start:
compute_end = 0 # 整个 tile 不可见，跳过
else:
local_end = query_row - kv_start + 1
compute_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 统一计算：
cuda
int64_t head_base = ((int64_t)batch * H + head) * N * D;现在需要拆成两个：
cuda
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 的签名变成了：
cuda
__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
);这个签名统一了三种场景：
| 场景 | N_q vs N_kv | is_causal |
|---|---|---|
| Self-attention（全可见） | N_q = N_kv | 0 |
| Causal self-attention | N_q = N_kv | 1 |
| Cross-attention | N_q ≠ N_kv | 0 |
整个 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：
| 测试场景 | 配置 |
|---|---|
| Self-attention, non-causal | max_diff = 6.3e-07 ✅ |
| Self-attention, causal | max_diff = 8.3e-07 ✅ |
| Cross-attention (N_q=64, N_kv=128) | max_diff = 5.4e-07 ✅ |
| Cross-attention (N_q=128, N_kv=64) | max_diff = 8.3e-07 ✅ |
手工验证 causal mask 的有效性：给 V 的最后一个位置赋一个极大值（999），position 0 的输出不受影响（因为它看不到），position 3 的输出被极大值支配（因为它看得到自己）。
性能：Fused 模式 vs Padded 模式
把 VecFused 在 non-causal self-attention 下和 VecPadded 做对比，确认没有引入额外开销：
| D | N | VecPad(ms) | VecFused(ms) | 变化 |
|---|---|---|---|---|
| 32 | 4096 | 214.5 | 211.2 | -1.5% |
| 64 | 4096 | 315.6 | 312.4 | -1.0% |
| 128 | 4096 | 482.6 | 474.5 | -1.7% |
差异在 ±2% 以内，属于正常波动。 多传两个参数（N_kv、is_causal）和拆开 head_base 的计算开销完全可以忽略。
Causal Mask 带来的意外收益
由于 causal mask 跳过了一半的 K/V 计算位置，kernel 在大 N 时自动提速约 45%：
| D | N | Non-causal(ms) | Causal(ms) | 加速 |
|---|---|---|---|---|
| 32 | 4096 | 211.2 | 118.9 | -44% |
| 64 | 4096 | 312.4 | 179.7 | -43% |
| 128 | 4096 | 474.5 | 284.5 | -41% |
这不是优化带来的------这是 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%。
| D | N | FP32(ms) | FP16(ms) | 加速比 |
|---|---|---|---|---|
| 32 | 4096 | 210.9 | 232.6 | 0.91x |
| 64 | 4096 | 311.0 | 333.1 | 0.93x |
| 128 | 4096 | 472.9 | 485.0 | 0.98x |
原因很直接，而且和我们上篇的核心结论一脉相承：
- 这个 kernel 是 compute-bound，不是 memory-bound。 省一半显存带宽没用------计算才是瓶颈
- half→float 转换不是免费的。 每个从显存读进来的 half 都要转成 float 才能做点积，写完再转回去。这些转换指令在 Turing 上大约占 5-10% 的指令周期
- 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 比，还差三件事：
- 分块稀疏（block-sparse） ------Flash Attention 通过对 score 矩阵做分块稀疏化，避免了完整 N×N 的显存开销。这在大 N 时是关键优化
- 内存布局优化 ------我们一直用
[B, H, N, D]的 naive 布局。实际推理引擎会用 Interleaved 布局或 GQA/MQA 的分组策略 - 更低精度支持 ------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 做它擅长的（模式匹配、汇编），人做他擅长的（理解结构、决策取舍）。