Attention: probe each packed run with its own tensors; packed calls with dropout keep xformers

tests/utils/test_attention_native_routing.py

@@ -427,3 +427,45 @@ def test_probe_matches_singleton_key_length(monkeypatch):
         assert set(kv_lengths) == {128, 1}
     finally:
         ad._sdpa_gqa_choices.cache_clear()
+
+
+def test_packed_dropout_keeps_xformers(monkeypatch):
+    """Torch varlen takes no dropout: a packed call with dropout keeps xformers' single launch."""
+    monkeypatch.setattr(ad, "_ring_attention_active", lambda: False)
+    monkeypatch.setattr(ad, "HAS_FLASH_ATTENTION", False)
+    monkeypatch.setattr(ad, "HAS_XFORMERS", True)
+    monkeypatch.setattr(ad, "_NATIVE_BEATS_XFORMERS", True)
+    monkeypatch.setattr(ad, "_SDPA_PACKED_SEGMENTS", True)
+    assert ad.select_attention_backend(True) == ad.SDPA
+    assert ad.select_attention_backend(True, 0.1) == ad.XFORMERS
+    assert ad.select_attention_backend(True, dropout_p = 0.0) == ad.SDPA
+
+
+def test_packed_runs_are_probed_with_their_own_lengths(monkeypatch):
+    """A one-token run is probed as one token, not with the whole row's length."""
+    seen = []
+
+    def takes_gqa(
+        Q,
+        K,
+        is_causal,
+        has_mask,
+        V = None,
+        dropout_p = 0.0,
+    ):
+        seen.append((Q.shape[-2], K.shape[-2]))
+        return True
+
+    monkeypatch.setattr(ad, "_sdpa_takes_gqa", takes_gqa)
+    monkeypatch.setattr(
+        ad, "scaled_dot_product_attention", lambda Q, K, V, **kw: torch.zeros_like(Q)
+    )
+    packing.clear_packed_caches()
+    lengths = (5, 1)
+    T = sum(lengths)
+    Q = torch.randn(1, 4, T, 16)
+    K = torch.randn(1, 2, T, 16)
+    ad._sdpa_packed_segments(
+        Q, K, K, lengths, n_groups = 2, is_causal = True, sliding_window = None, sdpa_kwargs = {}
+    )
+    assert (1, 1) in seen and all(q == k for q, k in seen)

unsloth/models/granite.py

@@ -140,7 +140,9 @@ def GraniteAttention_fast_forward(
 
     use_varlen = attention_mask is None and seq_info is not None and not had_cache
 
-    backend = SDPA if attention_mask is not None else select_attention_backend(use_varlen)
+    backend = (
+        SDPA if attention_mask is not None else select_attention_backend(use_varlen, dropout_p)
+    )
 
     window = (kv_seq_len, kv_seq_len)
     softmax_scale = getattr(self, "scaling", None)

unsloth/utils/attention_dispatch.py

@@ -551,13 +551,15 @@ def _sdpa_packed_segments(
                 is_causal = is_causal,
             )
         kwargs["is_causal"] = is_causal and not windowed
+        Q_run = split(Q)
+        # Probe the run's own tensors: a one-token run has different eligibility than the row.
         gqa = n_groups != 1 and _sdpa_takes_gqa(
-            Q, K, kwargs["is_causal"], windowed, V, kwargs.get("dropout_p", 0.0)
+            Q_run, split(K), kwargs["is_causal"], windowed, split(V), kwargs.get("dropout_p", 0.0)
         )
         if gqa:
             kwargs["enable_gqa"] = True
         K_run, V_run = kv(gqa)
-        out = scaled_dot_product_attention(split(Q), split(K_run), split(V_run), **kwargs)
+        out = scaled_dot_product_attention(Q_run, split(K_run), split(V_run), **kwargs)
         outs.append(out.transpose(1, 2).flatten(0, 1))
         start, i = end, j
     return torch.cat(outs, dim = 0).unsqueeze(0)
@@ -712,7 +714,7 @@ _NATIVE_BEATS_XFORMERS = _native_beats_xformers(
 )
 
 
-def select_attention_backend(use_varlen: bool = False) -> str:
+def select_attention_backend(use_varlen: bool = False, dropout_p: float = 0.0) -> str:
     """Return attention backend based on availability / priority order."""
 
     if _ring_attention_active():
@@ -730,8 +732,11 @@ def select_attention_backend(use_varlen: bool = False) -> str:
     # Only the packed route prefers torch's kernels (varlen / per-segment SDPA). Dense calls keep
     # xformers: it alone consumes a caller's block-diagonal causal_mask, and expanding a GQA KV cache
     # for SDPA on fp32 decode costs memory. With UNSLOTH_SDPA_PACKED_SEGMENTS off, packed SDPA would
-    # build a dense (T, T) mask, so xformers stays there too.
-    if HAS_XFORMERS and not (use_varlen and _NATIVE_BEATS_XFORMERS and _SDPA_PACKED_SEGMENTS):
+    # build a dense (T, T) mask, so xformers stays there too. Torch varlen takes no dropout, so a
+    # packed call with dropout would become one SDPA launch per run: xformers keeps those as well.
+    if HAS_XFORMERS and not (
+        use_varlen and not dropout_p and _NATIVE_BEATS_XFORMERS and _SDPA_PACKED_SEGMENTS
+    ):
         return XFORMERS
     return SDPA