Skip to content

add blackwell support filter for 9.7<=cudnn<9.18.1 #2775

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sudhakarsingh27 wants to merge 8 commits into
base: main
Choose a base branch
from
Open

add blackwell support filter for 9.7<=cudnn<9.18.1 #2775

Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
c9daeba
add blackwell support filter for 9.7<=cudnn<9.18.1
sudhakarsingh27 Mar 17, 2026
9c30493
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Mar 17, 2026
37800b9
simplify conditionals
sudhakarsingh27 Mar 17, 2026
25da675
merge conflict
sudhakarsingh27 Mar 17, 2026
342c68a
fix conditionals again
sudhakarsingh27 Mar 17, 2026
73880c5
fix conditionals again
sudhakarsingh27 Mar 17, 2026
481d712
Merge branch 'main' of github.com:NVIDIA/TransformerEngine into fix_f…
sudhakarsingh27 Mar 19, 2026
9e4f358
update the error log
sudhakarsingh27 Mar 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions transformer_engine/pytorch/attention/dot_product_attention/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1090,6 +1090,19 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
logger.debug("Disabling FusedAttention for determinism reasons with post_scale_bias")
use_fused_attention = False
fused_attention_backend = None
if (
fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"]
and is_training
and cudnn_version >= (9, 7, 0)
and cudnn_version < (9, 18, 1)
and device_compute_capability >= (10, 0)
):
logger.debug(
"Disabling FusedAtttention because determinism is not supported on Blackwell for "
"FP16/BF16 with 9.7 <= cuDNN < 9.18.1"
)
use_fused_attention = False
fused_attention_backend = None
Comment on lines +1093 to +1105
Copy link
Contributor

@greptile-apps greptile-apps bot Mar 17, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Missing is_training guard — may incorrectly disable FusedAttention for inference

Every other determinism filter in this same if use_fused_attention and deterministic: block guards against is_training (see lines 1070–1080 for FP8 and 1081–1092 for F16_arbitrary_seqlen), conveying that those non-determinism issues are backward-pass–specific. The new Blackwell / cuDNN-version filter does not include and is_training, so it will also disable FusedAttention during deterministic inference on Blackwell GPUs with cuDNN 9.7–9.18.1.

If the cuDNN bug only manifests during training (backward pass), the filter is overly broad and will unnecessarily fall back to a slower backend during inference. If it truly affects the forward pass as well, a comment explaining that would help reviewers and future maintainers understand the deviation from the existing pattern.

Consider either:

        if (
            fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"]
            and is_training
            and (cudnn_version >= (9, 7) and cudnn_version < (9, 18, 1))
            and device_compute_capability >= (10, 0)
        ):

or, if inference is also affected, add a comment explaining why is_training is deliberately omitted.

Copy link
Collaborator Author

@sudhakarsingh27 sudhakarsingh27 Mar 19, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

updated, check again


# use_flash_attention may have been set above
use_flash_attention_2 = use_flash_attention and use_flash_attention_2
Expand Down
Loading