Bypass fast_generate for flash_attention_2 models (StaticCache + FA2 produces gibberish) (#7429)

* Bypass fast_generate for flash_attention_2 models (frozen KV / gibberish)

unsloth_base_fast_generate forces cache_implementation="static", which
pre-allocates the full prompt+max_new_tokens KV buffer. With SDPA the
not-yet-filled slots are masked out; flash_attention_2 does not receive such
a mask, so decoding attends over uninitialized cache memory and produces
incoherent output (observed: coherent prompt echo followed by gibberish
rollouts on Phi-4-mini-instruct during TRL GRPO training; the KV length
appears frozen at the pre-allocated size). Note that on transformers >=
4.56 UNSLOTH_DISABLE_STATIC_GENERATION=1 still selects the static cache, so
the env-var escape hatch does not help either.

Fall back to the wrapped model's original generate when the config reports
_attn_implementation == "flash_attention_2" - plain HF generate is correct
with FA2 (validated: prefill q=13/kv=13, cache grows 14, 15, ..., coherent
output; equivalent to UNSLOTH_DISABLE_FAST_GENERATION=1 but scoped to FA2).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* Fix FA2 vision generation fallback

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

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

* Detect FA2 in VLM llm configs

* Fix default FlashAttention config detection

* Honor language attention overrides

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

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

* Handle nested FA2 configs and cache cleanup

* Pin a dynamic cache on the FlashAttention fallback for PR #7429

* Cover the explicit cache kwarg and caller caches in the FA2 fallback for PR #7429

* Tighten the FlashAttention fallback comments for PR #7429

---------

Co-authored-by: Piotr Wąsiewicz <piotrwasiewicz72@mail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <danielhanchen@gmail.com>

tests/test_fa2_fast_generate_bypass.py

@@ -0,0 +1,343 @@
+"""Regression coverage for the FlashAttention generation fallback."""
+
+import ast
+import inspect
+import os
+from contextlib import nullcontext
+from pathlib import Path
+from types import SimpleNamespace
+
+
+VISION_PATH = Path(__file__).parents[1] / "unsloth" / "models" / "vision.py"
+
+
+def _load_function(name, namespace):
+    tree = ast.parse(VISION_PATH.read_text(encoding = "utf-8"))
+    function = next(
+        node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == name
+    )
+    exec(compile(ast.Module(body = [function], type_ignores = []), str(VISION_PATH), "exec"), namespace)
+    return namespace[name]
+
+
+uses_flash_attention = _load_function(
+    "_uses_flash_attention_for_generation",
+    {
+        "_config_get": lambda config, field, default = None: (
+            config.get(field, default)
+            if isinstance(config, dict)
+            else getattr(config, field, default)
+        ),
+        "_is_flash_attention_requested": lambda value: (
+            isinstance(value, str) and value.startswith("flash_attention")
+        ),
+    },
+)
+clear_generation_caches = _load_function("_clear_generation_caches", {})
+
+
+def test_top_level_flash_attention_is_detected():
+    config = SimpleNamespace(_attn_implementation = "flash_attention_2")
+    assert uses_flash_attention(config)
+
+
+def test_per_backbone_text_flash_attention_is_detected():
+    private_config = SimpleNamespace(
+        _attn_implementation = {
+            "vision_config": "sdpa",
+            "text_config": "flash_attention_2",
+        }
+    )
+    public_config = SimpleNamespace(
+        attn_implementation = {
+            "vision_config": "sdpa",
+            "text_config": "flash_attention_2",
+        }
+    )
+    assert uses_flash_attention(private_config)
+    assert uses_flash_attention(public_config)
+
+
+def test_per_backbone_llm_flash_attention_is_detected():
+    config = SimpleNamespace(
+        _attn_implementation = {
+            "vision_config": "sdpa",
+            "llm_config": "flash_attention_2",
+        }
+    )
+    assert uses_flash_attention(config)
+
+
+def test_default_backbone_flash_attention_is_detected():
+    config = SimpleNamespace(
+        _attn_implementation = {
+            "": "flash_attention_2",
+            "vision_config": "sdpa",
+        }
+    )
+    assert uses_flash_attention(config)
+
+
+def test_explicit_language_backend_overrides_default_backend():
+    config = SimpleNamespace(
+        _attn_implementation = {
+            "": "flash_attention_2",
+            "text_config": "sdpa",
+        }
+    )
+    assert not uses_flash_attention(config)
+
+
+def test_nested_language_backend_overrides_normalized_default_backend():
+    config = SimpleNamespace(
+        _attn_implementation = "flash_attention_2",
+        text_config = SimpleNamespace(_attn_implementation = "sdpa"),
+    )
+    assert not uses_flash_attention(config)
+
+    nested_text = SimpleNamespace(_attn_implementation = "sdpa")
+    thinker_config = SimpleNamespace(
+        _attn_implementation = "flash_attention_2",
+        sub_configs = {"text_config": object},
+        text_config = nested_text,
+        get_text_config = lambda: nested_text,
+    )
+    assert not uses_flash_attention(SimpleNamespace(thinker_config = thinker_config))
+
+
+def test_nested_text_and_decoder_configs_are_detected():
+    nested_text = SimpleNamespace(attn_implementation = "flash_attention_2")
+    assert uses_flash_attention(
+        SimpleNamespace(_attn_implementation = "sdpa", text_config = nested_text)
+    )
+    assert uses_flash_attention(
+        SimpleNamespace(decoder_config = {"_attn_implementation": "flash_attention_2"})
+    )
+
+
+def test_nested_llm_config_is_detected():
+    config = SimpleNamespace(llm_config = SimpleNamespace(_attn_implementation = "flash_attention_2"))
+    assert uses_flash_attention(config)
+
+
+def test_get_text_config_is_detected():
+    nested_text = SimpleNamespace(_attn_implementation = "flash_attention_2")
+    config = SimpleNamespace(get_text_config = lambda: nested_text)
+    assert uses_flash_attention(config)
+
+
+def test_declared_custom_generation_subconfig_is_detected():
+    nested_text = SimpleNamespace(_attn_implementation = "flash_attention_2")
+    custom_generation = SimpleNamespace(
+        sub_configs = {"text_config": object},
+        text_config = nested_text,
+    )
+    config = SimpleNamespace(
+        sub_configs = {"custom_generation_config": object},
+        custom_generation_config = custom_generation,
+    )
+    assert uses_flash_attention(config)
+    assert uses_flash_attention(
+        SimpleNamespace(
+            _attn_implementation = {
+                "thinker_config": "flash_attention_2",
+                "vision_config": "sdpa",
+            }
+        )
+    )
+
+
+def test_vision_only_flash_attention_does_not_bypass_text_generation():
+    config = SimpleNamespace(
+        _attn_implementation = {
+            "vision_config": "flash_attention_2",
+            "text_config": "sdpa",
+        }
+    )
+    assert not uses_flash_attention(config)
+
+
+def test_non_flash_attention_does_not_bypass_fast_generation():
+    assert not uses_flash_attention(SimpleNamespace(_attn_implementation = "sdpa"))
+    assert not uses_flash_attention(SimpleNamespace())
+
+
+def test_wrapper_dispatch_preserves_normalization_and_selects_expected_path():
+    events = []
+
+    class FakeTensor:
+        shape = (1, 3)
+
+        def __init__(self):
+            self.converted_to = None
+
+        def to(self, dtype):
+            self.converted_to = dtype
+            return self
+
+    class FailIfUsed:
+        def __getattr__(self, name):
+            raise AssertionError(f"fast-generation path unexpectedly used torch._dynamo.{name}")
+
+    fake_torch = SimpleNamespace(
+        Tensor = FakeTensor,
+        bfloat16 = "bfloat16",
+        float16 = "float16",
+        _dynamo = FailIfUsed(),
+        inference_mode = nullcontext,
+        autocast = lambda **kwargs: nullcontext(),
+    )
+
+    class FakeFastBaseModel:
+        @staticmethod
+        def for_inference(model):
+            events.append("for_inference")
+
+    architecture = "Qwen3VLForConditionalGeneration"
+    namespace = {
+        "torch": fake_torch,
+        "os": os,
+        "inspect": inspect,
+        "FastBaseModel": FakeFastBaseModel,
+        "dtype_from_config": lambda config: "bfloat16",
+        "_get_dtype": lambda dtype: dtype,
+        "_unsloth_generate_accepts_kwarg": lambda model, name: False,
+        "NUM_LOGITS_TO_KEEP": {architecture: None},
+        "DEVICE_TYPE_TORCH": "cuda",
+        "_uses_flash_attention_for_generation": uses_flash_attention,
+        "_clear_generation_caches": clear_generation_caches,
+    }
+    fast_generate = _load_function("unsloth_base_fast_generate", namespace)
+
+    captured = {}
+    cache_module = SimpleNamespace(_flex_attention_cache = object())
+
+    class Model:
+        config = SimpleNamespace(
+            architectures = [architecture],
+            eos_token_id = 2,
+            text_config = SimpleNamespace(_attn_implementation = "flash_attention_2"),
+        )
+
+        def forward(self, input_ids = None):
+            return input_ids
+
+        def named_modules(self):
+            return [("cache", cache_module)]
+
+        def _old_generate(self, *args, **kwargs):
+            assert not hasattr(cache_module, "_flex_attention_cache")
+            captured.update(kwargs)
+            cache_module._flex_attention_cache = object()
+            return "fallback-result"
+
+    input_ids = FakeTensor()
+    pixel_values = FakeTensor()
+    result = fast_generate(
+        Model(),
+        input_ids = input_ids,
+        pixel_values = pixel_values,
+        mm_token_type_ids = FakeTensor(),
+    )
+
+    assert result == "fallback-result"
+    assert events == ["for_inference"]
+    assert "mm_token_type_ids" not in captured
+    assert captured["pixel_values"] is pixel_values
+    assert pixel_values.converted_to == "bfloat16"
+    assert not hasattr(cache_module, "_flex_attention_cache")
+
+    class FastPathReached(Exception):
+        pass
+
+    class ExpectFastPath:
+        @staticmethod
+        def mark_static(*args, **kwargs):
+            raise FastPathReached
+
+    fake_torch._dynamo = ExpectFastPath()
+    Model.config._attn_implementation = "flash_attention_2"
+    Model.config.text_config._attn_implementation = "sdpa"
+    captured.clear()
+    try:
+        fast_generate(Model(), input_ids = FakeTensor())
+    except FastPathReached:
+        pass
+    else:
+        raise AssertionError("non-FlashAttention generation did not enter the fast path")
+    assert captured == {}
+
+
+def test_flash_attention_fallback_pins_a_dynamic_cache():
+    # Delegating is not enough on its own: a static cache still reaches FlashAttention via
+    # an explicit kwarg, the caller's generation_config, or the model default.
+    namespace = {
+        "torch": SimpleNamespace(
+            Tensor = type("FakeTensor", (), {"shape": (1, 3)}),
+            bfloat16 = "bfloat16",
+            float16 = "float16",
+            inference_mode = nullcontext,
+            autocast = lambda **kwargs: nullcontext(),
+        ),
+        "os": os,
+        "inspect": inspect,
+        "FastBaseModel": SimpleNamespace(for_inference = lambda model: None),
+        "dtype_from_config": lambda config: "bfloat16",
+        "_get_dtype": lambda dtype: dtype,
+        "_unsloth_generate_accepts_kwarg": lambda model, name: False,
+        "NUM_LOGITS_TO_KEEP": {"Qwen3VLForConditionalGeneration": None},
+        "DEVICE_TYPE_TORCH": "cuda",
+        "_uses_flash_attention_for_generation": uses_flash_attention,
+        "_clear_generation_caches": clear_generation_caches,
+    }
+    fast_generate = _load_function("unsloth_base_fast_generate", namespace)
+
+    captured = {}
+
+    class Model:
+        config = SimpleNamespace(
+            architectures = ["Qwen3VLForConditionalGeneration"],
+            eos_token_id = 2,
+            _attn_implementation = "flash_attention_2",
+        )
+
+        def forward(self, input_ids = None):
+            return input_ids
+
+        def named_modules(self):
+            return []
+
+        def _old_generate(self, *args, **kwargs):
+            captured.clear()
+            captured.update(kwargs)
+            return "fallback-result"
+
+    input_ids = namespace["torch"].Tensor()
+
+    fast_generate(Model(), input_ids = input_ids)
+    assert captured["cache_implementation"] == "dynamic"
+
+    # The kwarg wins over a supplied generation_config, since update() applies it last.
+    generation_config = SimpleNamespace(cache_implementation = "static")
+    fast_generate(Model(), input_ids = input_ids, generation_config = generation_config)
+    assert captured["cache_implementation"] == "dynamic"
+
+    fast_generate(Model(), input_ids = input_ids, cache_implementation = "static")
+    assert captured["cache_implementation"] == "dynamic"
+
+    # generate() rejects a caller cache combined with any cache_implementation.
+    cache = object()
+    fast_generate(Model(), input_ids = input_ids, past_key_values = cache)
+    assert "cache_implementation" not in captured
+    assert captured["past_key_values"] is cache
+
+
+if __name__ == "__main__":
+    tests = [
+        value
+        for name, value in sorted(globals().items())
+        if name.startswith("test_") and callable(value)
+    ]
+    for test in tests:
+        test()
+    print(f"OK: {len(tests)} FA2 fallback regression tests passed")