Fix FlashAttention fp32 crash with DoRA (use_dora=True) (#6526)

* Fix FlashAttention fp32 crash with DoRA (use_dora=True)

DoRA upcasts lora_magnitude_vector to fp32 for the optimizer, which promotes
the q/k/v_proj output to fp32. FlashAttention only accepts fp16/bf16, so the
fp32 q/k/v raised 'FlashAttention only support fp16 and bf16 data type'.
Downcast q/k/v to the compute dtype before the flash kernels.

Fixes #1013

* Apply kwarg-spacing format hook to DoRA dtype test (pre-commit)

* DoRA+FA2: downcast any fp32 among Q/K/V and clamp to a flash-supported dtype

* Tighten code comments (no logic change)

---------

Co-authored-by: danielhanchen <michaelhan2050@gmail.com>

tests/utils/test_attention_dispatch_dora_dtype.py

@@ -0,0 +1,73 @@
+"""Regression test for unslothai/unsloth#1013: run_attention must downcast DoRA's
+fp32 q/k/v for FlashAttention while leaving already-16bit tensors untouched."""
+
+import torch
+import unsloth  # noqa: F401
+
+from unsloth.utils import attention_dispatch as ad
+
+
+def _run(monkeypatch, qkv_dtype, backend):
+    captured = {}
+
+    def _check(Q):
+        captured["dtype"] = Q.dtype
+        # Mirror the real kernel constraint so an unfixed dispatch fails loudly.
+        if Q.dtype not in (torch.float16, torch.bfloat16):
+            raise RuntimeError("FlashAttention only support fp16 and bf16 data type")
+
+    def fake_flash_dense(Q, K, V, **kwargs):
+        _check(Q)
+        return torch.zeros_like(Q)
+
+    def fake_flash_varlen(Q, K, V, *args, **kwargs):
+        _check(Q)
+        return torch.zeros_like(Q)
+
+    monkeypatch.setattr(ad, "flash_attn_func", fake_flash_dense, raising = False)
+    monkeypatch.setattr(ad, "flash_attn_varlen_func", fake_flash_varlen, raising = False)
+
+    bsz, n_heads, q_len, head_dim = 1, 2, 4, 8
+    Q = torch.randn(bsz, n_heads, q_len, head_dim, dtype = qkv_dtype)
+    K = torch.randn(bsz, n_heads, q_len, head_dim, dtype = qkv_dtype)
+    V = torch.randn(bsz, n_heads, q_len, head_dim, dtype = qkv_dtype)
+
+    seq_info = None
+    if backend == ad.FLASH_VARLEN:
+        cu = torch.tensor([0, q_len], dtype = torch.int32)
+        seq_info = (None, cu, q_len)
+
+    config = ad.AttentionConfig(
+        backend = backend,
+        n_kv_heads = n_heads,
+        n_groups = 1,
+        flash_dense_kwargs = {"causal": True},
+        flash_varlen_kwargs = {"dropout_p": 0.0, "causal": True},
+    )
+    context = ad.AttentionContext(
+        bsz = bsz,
+        q_len = q_len,
+        kv_seq_len = q_len,
+        n_heads = n_heads,
+        head_dim = head_dim,
+        requires_grad = False,
+        seq_info = seq_info,
+        attention_mask = None,
+        causal_mask = None,
+    )
+    ad.run_attention(config = config, context = context, Q = Q, K = K, V = V)
+    return captured["dtype"]
+
+
+def test_dense_flash_downcasts_fp32_qkv(monkeypatch):
+    # fp32 DoRA output must be downcast to a flash-compatible dtype.
+    assert _run(monkeypatch, torch.float32, ad.FLASH_DENSE) in (torch.bfloat16, torch.float16)
+
+
+def test_varlen_flash_downcasts_fp32_qkv(monkeypatch):
+    assert _run(monkeypatch, torch.float32, ad.FLASH_VARLEN) in (torch.bfloat16, torch.float16)
+
+
+def test_bf16_qkv_left_untouched(monkeypatch):
+    # Standard LoRA path (already bf16) must not be altered.
+    assert _run(monkeypatch, torch.bfloat16, ad.FLASH_DENSE) == torch.bfloat16

unsloth/utils/attention_dispatch.py

@@ -137,6 +137,25 @@ def run_attention(
     requires_grad = context.requires_grad
     sliding_window = context.sliding_window
 
+    # DoRA promotes q/k/v_proj outputs to fp32, which FlashAttention rejects, so
+    # downcast any fp32 Q/K/V to a flash-supported dtype (#1013).
+    if backend in (FLASH_DENSE, FLASH_VARLEN) and torch.float32 in (
+        Q.dtype,
+        K.dtype,
+        V.dtype,
+    ):
+        # Prefer the autocast dtype, else a non-fp32 input's dtype, then clamp.
+        if torch.is_autocast_enabled():
+            try:
+                flash_dtype = torch.get_autocast_dtype("cuda")
+            except (AttributeError, TypeError):
+                flash_dtype = torch.get_autocast_gpu_dtype()
+        else:
+            flash_dtype = next((d for d in (Q.dtype, K.dtype, V.dtype) if d != torch.float32), None)
+        if flash_dtype not in (torch.float16, torch.bfloat16):
+            flash_dtype = torch.bfloat16 if SUPPORTS_BFLOAT16 else torch.float16
+        Q, K, V = Q.to(flash_dtype), K.to(flash_dtype), V.to(flash_dtype)
+
     if backend == FLASH_VARLEN:
         Q_f = Q.transpose(1, 2).reshape(bsz * q_len, n_heads, head_dim)
         K_f = K.transpose(1, 2).reshape(bsz * q_len, config.n_kv_heads, head_dim)