diff --git a/tests/utils/test_attention_native_routing.py b/tests/utils/test_attention_native_routing.py index bec6f9bde1..028c1d0a6a 100644 --- a/tests/utils/test_attention_native_routing.py +++ b/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) diff --git a/unsloth/models/granite.py b/unsloth/models/granite.py index 9be2a7a85d..d26a35af19 100644 --- a/unsloth/models/granite.py +++ b/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) diff --git a/unsloth/utils/attention_dispatch.py b/unsloth/utils/attention_dispatch.py index 92ba19a00c..52c808d56a 100644 --- a/unsloth/utils/attention_dispatch.py +++ b/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