Hacktoberfest 2026: los issues que los mantenedores marcaron para octubre, abiertos y aptos para principiantes. Explorar issues de Hacktoberfest

Fused attention should distinguish bias input from dBias request

Abierto
  • ¿Algo no está bien?
#3,082 0 comentarios 0 reacciones 1 asignado Ver en GitHub

Los mantenedores suelen responder en 2 días

@cyanguwa ya está trabajando en esto.

Desde el 4/6/2026.

Evaluación

Este issue todavía no se ha evaluado.

Descripción

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.

Lenguaje dominante
Python
Estrellas
3.6k
Forks
851
Merge medio
5 d 1 h
PR fusionados (30 d)
52

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de NVIDIA/TransformerEngine

Todos los issues de NVIDIA/TransformerEngine

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.