副标题给我们的 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-attentionCausal mask位置 i 只能看 j ≤ iGPT、LLaMA、Qwen生成阶段Encoder self-attention全可见QKV 等长BERT、Encoder-onlyEncoder-Decoder cross-attentionQ 来自 decoderK/V 来自 encoder不等长T5、机器翻译、语音识别这一篇的目标基于上篇性能最好的 VecPadded kernel一次性支持上面所有场景。二、Causal Mask为什么需要怎么实现为什么 decoder 需要 causal mask自回归生成时位置 i 的输出只能依赖位置 1…i 的输入——它看不到未来的 token。如果允许看到未来的 token模型就相当于提前知道了答案训练时 loss 会 collapse生成时无法自回归。实现上就是在 softmax 之前把未来位置设为 -infAttention(Q, K, V) softmax(QK^T / √D Mask) · V其中 Mask[i][j] 0j ≤ 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_rowkv_start:compute_end0# 整个 tile 不可见跳过else:local_endquery_row-kv_start1compute_endmin(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 来自 decoderK/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 );这个签名统一了三种场景场景N_q vs N_kvis_causalSelf-attention全可见N_q N_kv0Causal self-attentionN_q N_kv1Cross-attentionN_q ≠ N_kv0整个 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_basecausal 就是在计算循环里加个判断。通过 Agent review 确认没有遗漏。对Agent 写代码的补充思考上篇我们提到 Agent 能一次写对 CUDA kernel是因为 NVIDIA 的生态太好了——编程指南、官方 Sample、StackOverflow 百万问答都在训练数据里。这篇的经验进一步说明了一个更微妙的点Agent 擅长的是模式匹配——在熟悉的语境下从零组装新代码。但它不太擅长理解已有代码然后在正确的位置插入修改。后者需要的不只是模式识别还需要理解局部结构和全局约束的关系。这恰好对应算子开发中的两类工作绿字段落新架构、新算子Agent 很强因为它见过相似的 pattern棕字段落扩展现有算子、修 bugAgent 会犯错因为它不理解代码的上下文所以一个人机协作的合理分工是让 Agent 写新 kernel让人来改已有代码——或者反过来人改代码、Agent 做 review。五、正确性与性能正确性测试结果对比 PyTorch 的F.scaled_dot_product_attention(q, k, v, is_causal...)4 个场景全 PASS测试场景配置Self-attention, non-causalmax_diff 6.3e-07 ✅Self-attention, causalmax_diff 8.3e-07 ✅Cross-attention (N_q64, N_kv128)max_diff 5.4e-07 ✅Cross-attention (N_q128, N_kv64)max_diff 8.3e-07 ✅手工验证 causal mask 的有效性给 V 的最后一个位置赋一个极大值999position 0 的输出不受影响因为它看不到position 3 的输出被极大值支配因为它看得到自己。性能Fused 模式 vs Padded 模式把 VecFused 在 non-causal self-attention 下和 VecPadded 做对比确认没有引入额外开销DNVecPad(ms)VecFused(ms)变化324096214.5211.2-1.5%644096315.6312.4-1.0%1284096482.6474.5-1.7%差异在 ±2% 以内属于正常波动。多传两个参数N_kv、is_causal和拆开 head_base 的计算开销完全可以忽略。Causal Mask 带来的意外收益由于 causal mask 跳过了一半的 K/V 计算位置kernel 在大 N 时自动提速约45%DNNon-causal(ms)Causal(ms)加速324096211.2118.9-44%644096312.4179.7-43%1284096474.5284.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%。DNFP32(ms)FP16(ms)加速比324096210.9232.60.91x644096311.0333.10.93x1284096472.9485.00.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 做它擅长的模式匹配、汇编人做他擅长的理解结构、决策取舍。