【Bug已解决】[Feature Request] CUDA EP: support `attention_bias` in GroupQueryAttention (last EP missing…-育师

【Bug已解决】[Feature Request] CUDA EP: support`attention_bias`in GroupQueryAttention (last EP missing it) 解决方案

一、现象长什么样

用 ONNX Runtime 跑一个带attention_bias（注意力偏置，如 ALiBi 或旋转位置编码的偏置项）的GroupQueryAttention节点，在CUDA EP上要么直接报错“不支持 attention_bias”，要么静默忽略 bias 算出错误结果；而同样的模型在 CPU EP / WebGPU EP 上是正常的。现象：

# 现象 A：CUDA EP 直接拒绝 # NotImplementedError: GroupQueryAttention with attention_bias is not # supported on CUDA EP # 现象 B：不报错但结果错（更危险） # 某些版本 CUDA EP 收下了 attention_bias 输入却没在 kernel 里加， # 输出和 CPU EP 对不上，且无任何报错 # 现象 C：只在带 bias 的 GQA 触发 # 普通 GQA（无 bias）在 CUDA EP 正常；一旦模型用了 attention_bias 就炸/错

最坑的是现象 B：能跑、不报错、但注意力算错，只有和 CPU/WebGPU EP 对拍才发现偏差，且这种偏差在长序列（ALiBi 偏置影响大）上尤其明显。

二、背景

GroupQueryAttention（GQA）在 ONNX 里支持一个可选输入attention_bias，它会被加到 attention 的QK^T / sqrt(d)分数上，实现 ALiBi 等位置偏置。ORT 的多个 EP 都实现了 GQA kernel：CPU、WebGPU、CUDA。

问题出在：CUDA EP 的 GQA kernel 是最晚补齐功能的，只实现了无 bias 的路径，当节点带attention_bias输入时：要么 kernel 没有对应的 bias-launch 分支直接拒绝（现象 A），要么（更糟）kernel 接收了 bias 这个 tensor 却没在打分公式里加上它（现象 B）。其他 EP 早就支持了，于是“最后一个 EP 缺这个功能”被单独提为 issue。

这是 EP 功能对齐审查里典型的坑：同一算子在多个 EP 上的功能覆盖不一致，CUDA EP 落后，且落后时可能静默忽略输入。

三、根因

CUDA GQA kernel 无 bias 分支：kernel 只 launch 了无 bias 的模板，遇到 bias 输入没有对应路径 → 现象 A。
bias 输入被静默忽略：kernel 接收了attention_bias这个 input 却没在score = QK^T/√d + bias里加，导致漏加 → 现象 B。
缺少跨 EP 结果对拍：CI 没把“CUDA EP 结果”和“CPU/WebGPU EP 结果”对拍，静默忽略 bias 的回归长期存在。

本质：是CUDA EP 的 GQA kernel 功能落后于其他 EP（缺 attention_bias 路径），且落后时可能静默忽略输入，缺跨 EP 对拍。

四、最小可运行复现

下面用 Python 模拟“GQA 打分：无 bias 路径忽略 bias 输入导致结果错”：

import torch def gqa_scores_buggy(q, k, bias=None): """buggy: 收了 bias 却没加。""" scores = (q @ k.transpose(-1, -2)) / (q.shape[-1] ** 0.5) # 忘了 scores = scores + bias return scores def gqa_scores_fixed(q, k, bias=None): scores = (q @ k.transpose(-1, -2)) / (q.shape[-1] ** 0.5) if bias is not None: scores = scores + bias # 正确加上 attention_bias return scores q = torch.randn(1, 4, 8, 16) k = torch.randn(1, 4, 8, 16) bias = torch.linspace(-0.1, 0.1, 8).expand(1, 4, 8, 8) b = gqa_scores_buggy(q, k, bias) f = gqa_scores_fixed(q, k, bias) print("results differ (buggy ignores bias)?", not torch.allclose(b, f, atol=1e-4)) print("max diff:", (b - f).abs().max().item()) # bias 的量级，说明漏加

buggy输出和fixed差了bias的量级，证明漏加 attention_bias。

五、解决方案（第一层：最小直接修复）

最小修复：CUDA GQA kernel 增加 bias 分支，在打分后加上attention_bias：

// 修正：CUDA GQA kernel 处理 attention_bias template <typename T> __global__ void GqaKernelWithBias(...) { // 计算 score = QK^T / sqrt(d) T score = ...; if (has_attention_bias) { score += attention_bias[batch * seq + seq_q * seq_k + seq_k_pos]; // 加偏置 } // softmax ... } // 调度：有 bias 走带 bias 的 kernel 实例化 if (attention_bias != nullptr) { LaunchGqaKernel<true>(...); // has_attention_bias=true } else { LaunchGqaKernel<false>(...); }

这一层改动最小：加 bias 分支并在打分后加偏置，结果恢复正确。但依赖“CUDA kernel 和功能对齐都维护”，下看第二层。

六、解决方案（第二层：结构性改进）

把“GQA 在 CUDA EP 必须支持 attention_bias，且与其他 EP 功能对齐”固化成单一事实来源。下面这个 dataclass 集中管理 GQA 功能能力声明，供调度与对拍使用：

from dataclasses import dataclass, field from typing import Dict, Set @dataclass class CudaGqaBiasPolicy: """单一事实来源：各 EP 的 GQA 功能能力声明（对齐契约）。""" # EP -> 支持的能力集合 _capabilities: Dict[str, Set[str]] = field(default_factory=lambda: { "CPU": {"attention_bias", "past_key", "qk_norm"}, "WebGPU": {"attention_bias", "past_key"}, "CUDA": {"past_key"}, # 初始缺 attention_bias }) def enable(self, ep: str, capability: str) -> None: self._capabilities.setdefault(ep, set()).add(capability) def supports(self, ep: str, capability: str) -> bool: return capability in self._capabilities.get(ep, set()) def assert_aligned(self, capability: str, reference_eps=("CPU", "WebGPU")) -> None: """断言某能力在所有参考 EP 上一致（CUDA 不能落后）。""" expected = all(self.supports(e, capability) for e in reference_eps) if expected and not self.supports("CUDA", capability): raise AssertionError( f"CUDA EP missing '{capability}' while {reference_eps} have it") # 用法：补齐 CUDA 的 attention_bias 后登记 policy = CudaGqaBiasPolicy() policy.enable("CUDA", "attention_bias") # 补齐功能 policy.assert_aligned("attention_bias") # 现在通过

这一层的关键收益：

能力声明集中：各 EP 的 GQA 能力集中在_capabilities，缺失一目了然；
对齐断言：assert_aligned确保 CUDA 不落后于 CPU/WebGPU，杜绝“最后一个 EP 缺功能”；
单一事实来源：所有 GQA 功能对齐约定收口在CudaGqaBiasPolicy。

七、解决方案（第三层：断言 / CI 守护）

把第二层钉成 pytest，挂进 CI，确保 CUDA GQA 支持 bias 且跨 EP 对齐：

import torch import pytest from your_package.cuda_gqa_bias import CudaGqaBiasPolicy def test_cuda_supports_bias_after_fix(): # 断言 1：修复后 CUDA EP 声明支持 attention_bias p = CudaGqaBiasPolicy() p.enable("CUDA", "attention_bias") assert p.supports("CUDA", "attention_bias") def test_alignment_assertion_catches_gap(): # 断言 2：CUDA 缺 bias 时对齐断言必报错 p = CudaGqaBiasPolicy() # CUDA 初始无 bias with pytest.raises(AssertionError): p.assert_aligned("attention_bias") def test_gqa_scores_include_bias(): # 断言 3：带 bias 的打分必须真的加上 bias（数值对拍） q = torch.randn(1, 4, 8, 16) k = torch.randn(1, 4, 8, 16) bias = torch.linspace(-0.1, 0.1, 8).expand(1, 4, 8, 8) scores = (q @ k.transpose(-1, -2)) / (q.shape[-1] ** 0.5) + bias base = (q @ k.transpose(-1, -2)) / (q.shape[-1] ** 0.5) assert not torch.allclose(scores, base, atol=1e-4) def test_cpu_webgpu_have_bias(): # 断言 4：参考 EP 都有 bias（作为对齐基准） p = CudaGqaBiasPolicy() assert p.supports("CPU", "attention_bias") assert p.supports("WebGPU", "attention_bias")

四条断言从“CUDA 支持 bias”“对齐断言抓缺口”“数值含 bias”“参考 EP 有 bias”四面把功能缺口钉死在 CI。

八、排查清单

CUDA EP 跑带 attention_bias 的 GQA 报错/结果错时：

报not supported on CUDA EP？确认 CUDA GQA kernel 是否有 bias 分支（现象 A）。
不报错但结果和 CPU/WebGPU 对不上？确认 bias 是否被静默忽略、没加进打分（现象 B）。
是否只在 CUDA EP 缺这个功能？查各 EP 的 GQA 能力是否对齐。
用第二层CudaGqaBiasPolicy：能力声明集中 +assert_aligned防落后。
加第三层 pytest，断言“CUDA 支持 bias、对齐断言抓缺口、数值含 bias、参考 EP 有 bias”。
同一算子跨 EP 必须功能对齐，新增能力时所有 EP 都要跟上，不能留“最后一个 EP 缺失”。

九、小结

CUDA EP 的 GQA 缺attention_bias支持，本质是CUDA GQA kernel 功能落后于 CPU/WebGPU EP，遇到 bias 输入要么拒绝、要么静默忽略不加进打分，导致报错或结果静默错误（长序列 ALiBi 偏差明显）；且缺跨 EP 对拍。修复分三层——第一层 CUDA kernel 加 bias 分支并在打分后加偏置；第二层用CudaGqaBiasPolicy这个 dataclass 把各 EP 的 GQA 能力声明收口成单一事实来源，assert_aligned确保 CUDA 不落后；第三层用四条 pytest 把“CUDA 支持 bias、对齐断言抓缺口、数值含 bias、参考 EP 有 bias”钉死在 CI。核心心法：同一算子在各 EP 上的功能必须对齐，能力声明应集中且用断言防‘最后一个 EP 缺失’，落后时绝不能静默忽略输入。