Skip flex_attention on pre-Ampere GPUs (T4, V100) by danielhanchen · Pull Request #4302 · unslothai/unsloth

danielhanchen · 2026-03-15T11:42:44Z

Summary

Skip flex_attention on GPUs with compute capability < 8.0 (pre-Ampere: T4/sm75, V100/sm70)
On these GPUs the Triton flex_attention kernel cannot run, so PyTorch falls back to sdpa_dense_backward which has a dtype mismatch under fp16 autocast (Half @ Float matmul at line 904 of flex_attention.py)
Falls back to sdpa instead, which works correctly on all GPU architectures

Root cause

prefer_flex_attn_if_supported() only checks torch >= 2.5.0 and the model's _supports_flex_attn flag, but not GPU compute capability. The flex_attention Triton kernels require sm80+.

Reproduction

Ministral-3B on T4 with transformers >= 5.0:

from unsloth import FastModel
model, tokenizer = FastModel.from_pretrained(
    model_name="unsloth/Ministral-3-3B-Instruct-2512-unsloth-bnb-4bit",
    max_seq_length=2048, load_in_4bit=True,
)

Crashes during backward with:

RuntimeError: expected scalar type Float but found Half

at torch/_higher_order_ops/flex_attention.py:904 in sdpa_dense_backward.

Test plan

Verified prefer_flex_attn_if_supported returns None on T4 (sm75)
Model loads with sdpa attention instead of flex_attention
Ministral-3B SFT training completes on T4 with transformers 5.2.0

Fixes #4295

flex_attention Triton kernels require sm80+ (Ampere). On older GPUs the dense Python fallback runs instead, but sdpa_dense_backward has a dtype mismatch under fp16 autocast -- the matmul at line 904 of torch/_higher_order_ops/flex_attention.py does softmax_scores.to(query.dtype) @ grad_out where query.dtype is Half and grad_out is Float, producing "expected scalar type Float but found Half". This affected Ministral-3B/8B on T4 GPUs (issue #4295) and potentially any model using flex_attention on pre-Ampere hardware. Fix: check torch.cuda.get_device_capability() >= (8, 0) before enabling flex_attention. Falls back to sdpa on older GPUs.

for more information, see https://pre-commit.ci

gemini-code-assist · 2026-03-15T11:42:57Z

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request resolves a critical issue where flex_attention caused runtime errors on older NVIDIA GPUs (pre-Ampere architectures like T4 and V100) due to its Triton kernels requiring compute capability 8.0 or higher. By introducing a GPU compute capability check, the system now intelligently defaults to sdpa attention on these unsupported architectures, preventing crashes and ensuring stable operation, especially when using fp16 autocast.

Highlights

GPU Compatibility Check: Implemented a check within prefer_flex_attn_if_supported to identify and skip flex_attention on GPUs with compute capability less than 8.0 (pre-Ampere architectures like T4/sm75 and V100/sm70).
Error Prevention: Addressed a RuntimeError (expected scalar type Float but found Half) that occurred during the backward pass with fp16 autocast on unsupported GPUs, caused by a dtype mismatch in sdpa_dense_backward when flex_attention fell back to its dense Python implementation.
Stable Fallback: Ensured that on these older, unsupported GPU architectures, the system now correctly falls back to using sdpa attention, which functions without issues.

Changelog

unsloth/models/_utils.py
- Added logic to prefer_flex_attn_if_supported to check the CUDA device's compute capability.
- Introduced a condition to return None (disabling flex_attention) if the GPU's major compute capability is less than 8.

Activity

No human activity (comments, reviews, etc.) has been recorded for this pull request yet.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request correctly adds a check to skip flex_attention on older GPUs. My feedback includes a small refactoring to improve the code's conciseness and readability by removing a redundant import and simplifying the conditional logic.

gemini-code-assist · 2026-03-15T11:44:46Z

+        import torch
+
+        if torch.cuda.is_available():
+            major, _ = torch.cuda.get_device_capability()
+            if major < 8:
+                return None
+        else:


This block can be made more concise and readable. The import torch is redundant since it's already imported at the top of the file, and the conditional logic can be simplified by checking for the negative case first and returning early. This flattens the code structure.

Suggested change

import torch

if torch.cuda.is_available():

major, _ = torch.cuda.get_device_capability()

if major < 8:

return None

else:

# Check for CUDA availability and compute capability.

# Return early if not supported to avoid nested ifs.

if not torch.cuda.is_available():

return None

major, _ = torch.cuda.get_device_capability()

if major < 8:

return None

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 45feffde6f

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

chatgpt-codex-connector · 2026-03-15T11:45:19Z

+        import torch
+
+        if torch.cuda.is_available():
+            major, _ = torch.cuda.get_device_capability()


Evaluate compute capability across all visible GPUs

This check only reads torch.cuda.get_device_capability() for the current CUDA device, but prefer_flex_attn_if_supported sets one global attention implementation for the whole model. In mixed-GPU runs (for example, device_map="auto" with both Ampere and pre-Ampere cards), if the current device is sm80+ and another shard lands on sm70/sm75, flex_attention will still be enabled and the same backward dtype mismatch this patch aims to avoid can still occur on the older shard.

Useful? React with 👍 / 👎.

Datta0 · 2026-03-18T05:51:15Z

We should find a way to combine #4210
Current impl prefers flex over flash which is potentially slower right?

Macmill-340 · 2026-04-14T03:26:33Z

Any updates on this?

danielhanchen requested a review from mmathew23 as a code owner March 15, 2026 11:42

[pre-commit.ci] auto fixes from pre-commit.com hooks

45feffd

for more information, see https://pre-commit.ci

danielhanchen mentioned this pull request Mar 15, 2026

Add compile fallback for fused CE loss (fix GB0149) unslothai/unsloth-zoo#547

Merged

5 tasks

gemini-code-assist Bot reviewed Mar 15, 2026

View reviewed changes

chatgpt-codex-connector Bot reviewed Mar 15, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Skip flex_attention on pre-Ampere GPUs (T4, V100)#4302

Skip flex_attention on pre-Ampere GPUs (T4, V100)#4302
danielhanchen wants to merge 2 commits into
mainfrom
fix/flex-attention-pre-ampere

danielhanchen commented Mar 15, 2026

Uh oh!

gemini-code-assist Bot commented Mar 15, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Mar 15, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

chatgpt-codex-connector Bot Mar 15, 2026

Uh oh!

Datta0 commented Mar 18, 2026

Uh oh!

Macmill-340 commented Apr 14, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

danielhanchen commented Mar 15, 2026

Summary

Root cause

Reproduction

Test plan

Uh oh!

gemini-code-assist Bot commented Mar 15, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Mar 15, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Mar 15, 2026

Choose a reason for hiding this comment

Uh oh!

Datta0 commented Mar 18, 2026

Uh oh!

Macmill-340 commented Apr 14, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants