Hacktoberfest 2026：维护者为十月标记出来的 issue，仍然开放、适合新手。 浏览 Hacktoberfest issue

Fused attention should distinguish bias input from dBias request

未关闭
  • 哪里不对？
#3,082 0 条评论 0 个 reaction 已指派 1 人 在 GitHub 查看

维护者通常 2 天内回复

@cyanguwa 已经在做这个了。

开始于 2026年6月4日。

评估

这个 Issue 还没有评估数据。

描述

attention

Is your feature request related to a problem? Please describe.

TE fused attention currently treats bias_type != NO_BIAS as implying that backward should request dBias from cuDNN FE. However, cuDNN FE models bias input and dBias output independently (set_bias vs set_dbias), and frameworks can have bias tensors that affect attention but do not require gradients. This matters for training cases where an additive attention bias is used as a fixed mask or frozen score modifier. We still need backward for Q/K/V, but do not need dBias. Requesting dBias unnecessarily can disable otherwise-supported cuDNN kernels or trigger plan-build failures for kernels that support bias input but not dBias.

Describe the solution you'd like
TE fused attention should distinguish “bias input is present” from “bias gradient is requested”.
Concretely, plumb a dbias_requested / bias_requires_grad flag through backend selection and backward execution:
• Backend selection should allow POST_SCALE_BIAS when dbias_requested == false for cuDNN kernels that support bias input
but not dBias.
• Backward graph construction should call sdpa_backward_options.set_bias(bias) whenever bias input is present.
• Backward graph construction should call sdpa_backward_options.set_dbias(dBias) only when the framework actually requests
a bias gradient.
• PyTorch can derive this from core_attention_bias.requires_grad.
• JAX can derive this from whether bias is included in value_and_grad(..., argnums=...).

So, in the cuDNN FE graph construction:

  if (bias_type != NVTE_NO_BIAS) {
      sdpa_backward_options.set_bias(bias);
      if (dbias_requested) {
          sdpa_backward_options.set_dbias(dBias);
      }
  }

For PyTorch usage, a frozen bias should still use fused attention without requesting dBias:

bias = make_additive_attention_bias(...)
bias.requires_grad_(False)
out = transformer_engine.pytorch.DotProductAttention(...)(q, k, v, core_attention_bias=bias)

Backward should compute dQ/dK/dV, but not request dBias from cuDNN FE.

Describe alternatives you've considered
One workaround is to disable fused attention whenever bias is present for kernels that do not support dBias. This is safe
but overly conservative, because cuDNN FE may support bias input even when it does not support dBias.

Another workaround is for users to encode fixed masks through attn_mask_type instead of additive bias. That only works for
built-in mask patterns and does not cover arbitrary dense score modifiers, frozen relative-position bias, or other
application-specific additive biases.

A third option is to keep using bias_type alone and infer dBias from bias shape. This is fragile because whether dBias is
needed is an autograd property, not a tensor-shape property.

Additional context

This came up while enabling D=256 backward fused attention on Blackwell/SM10x. The cuDNN FE path can distinguish
set_bias(...) from set_dbias(...), but TE’s common fused attention path currently does not expose that distinction in
backend selection. As a result, TE may reject or fail to use kernels that would be valid for bias-input-only training.

主要语言
Python
星标
3.6k
派生
851
平均合并
4 天 15 小时
30 天内合并 PR
51

环境准备

从这里开始

  1. 先读完整个 Issue，再读项目的贡献指南。
  2. 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
  3. Fork 仓库，在一个分支上完成修改。
  4. 提交 Pull Request，并在描述里引用这个 Issue 编号。

NVIDIA/TransformerEngine 的其他 Issue

查看 NVIDIA/TransformerEngine 的全部 Issue

相似的 Issue

更多 Python Issue

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。