SuperAPI AI 工具集

cuda-attention-kernel-patterns

编程与开发

查看项目

更新于 10-16 00:31

技能介绍

用途:面向 ONNX 域 Attention 算子(opset 23/24,实现在 core/providers/cuda/llm/attention.cc)CUDA 开发的复用知识,与 contrib 域的 MultiHeadAttention / GroupQueryAttention 相区分,二者共享 CUTLASS FMHA、Flash 内核与 attention_softmax.h 等基础设施,但调度、参数结构与资格判断不同。

主要能力与关键点:CUDA 调度按优先级 Flash → MEA(Memory Efficient)→ Unified Unfused 级联,Unfused 为兜底,覆盖 MHA 与 GQA,支持 softcap、attn_mask、output_qk、past_key/past_value、因果掩码等。Flash 仅 fp16/bf16、SM≥8.0、head_size%8==0 且 ≤256。MEA 要求 head_size 可被 8 整除且 ≤1024,bias 行步长需满足 total_sequence_length % 4 == 0;GQA 需 head_size == v_head_size,FP32 GQA 被排除。CUTLASS softmax 内部乘 kLog2e 会使 lowest() 溢出为 -inf 导致 NaN,故 MEA 路径将 mask_filter_value 上限裁至 -1.0e+30f。softcap 必须在加掩码/偏置之前应用,否则 -inf 掩码会被 tanh 变为有限值而泄漏。全掩码行输出零行(onnx#8068),CUDA 由 ZeroFullyMaskedRowsKernel 处理,CPU 为规范参考实现。

使用方式:测试可用 ScopedEnvironmentVariables 设置 ORT_DISABLE_FLASH_ATTENTION=1 或再设 ORT_DISABLE_MEMORY_EFFICIENT_ATTENTION=1 强制指定后端;用 SKIP_IF_MEA_NOT_COMPILED 避免假阳性。构建开启 onnxruntime_QUICK_BUILD 时 Flash 仅编译 head_dim=128,需注意调度变化。网格须限制在 65535 以内并使用 grid-stride 循环,核函数后用 CUDA_CALL(cudaGetLastError())。

限制:head_size=512 会被路由到 MEA 但在小共享内存架构上启动失败(#28388,修复 PR #28383 未合并),不可作为稳定测试探针;CUDA 的 qk_matmul_output_mode 超出 kNone/kQK 返回 NOT_IMPLEMENTED,不能假定 CUDA 产生 mode-3=0。其余细节资料未提供。