NVIDIA / NVIDIA/TransformerEngine

[PyTorch] Integrate cuDNN GQA + DSA backend into DotProductAttention

Open
#3,028 0 comments 1 reaction 1 assignee View on GitHub

@cyanguwa is already working on this.

Since May 26, 2026.

attention
Dominant language
Python
Stars
3.5k
Forks
831
Avg merge
3d 11h
Merged PRs (30d)
65

Description

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

Transformer Engine currently does not expose a path that combines Grouped Query Attention (GQA) with DeepSeek-style sparse attention (DSA), where each query token attends only to a TopK subset of key/value tokens. Several training workloads need this combination — a GQA attention shape (many query heads sharing fewer K/V heads) with a sparsity pattern that drops attention to all but a small index list per query. Without a TE-native backend, teams either fall back to community Triton kernels, which can't reach production-scale performance, or implement sparse attention outside of TE — losing autograd integration, kernel fusion, and parity with TE's existing attention features.

Describe the solution you'd like

Add a cuDNN-backed sparse-attention path inside DotProductAttention for the PyTorch frontend that:

  • Recognizes a sparse-attention mode and dispatches to the new cuDNN GQA + DSA kernel
  • Accepts a per-query sparse_indices tensor of shape [B, S_q, topk] selecting which K/V positions each query attends to
  • Supports the standard GQA shape (num_attention_heads ≠ num_gqa_groups)
  • Supports BF16 attention at minimum (FP8 indexer extension as a follow-on if needed)
  • Integrates cleanly with TE's autograd and existing context-parallelism path
  • Ships with numerical-equivalence tests against a reference dense-attention baseline restricted to the same TopK indices

cc: @cyanguwa

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.