ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source 
* 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>
This commit is contained in:
Piotr Wasiewicz 2026-07-27 08:07:33 +02:00 • committed by GitHub
commit 62d3438b99
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 477 additions and 32 deletions
Download patch file Download diff file Expand all files Collapse all files

343
tests/test_fa2_fast_generate_bypass.py Normal file
View file

@ -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")

166
unsloth/models/vision.py
View file

@ -36,6 +36,8 @@ from ._utils import (
resolve_attention_implementation,
_get_text_only_config,
_is_family_text_decoder,
_config_get,
_is_flash_attention_requested,
_apply_text_only_key_mapping,
_select_moe_detection_targets,
set_task_config_attr,
@ -226,8 +228,7 @@ def _attach_bnb_multidevice_hooks(
param.__dict__[key] = val
logger.info(
f"Unsloth: Attached accelerate AlignDevicesHook ({desc}) "
f"for bnb multi-GPU inference."
f"Unsloth: Attached accelerate AlignDevicesHook ({desc}) for bnb multi-GPU inference."
)
except Exception as exc:
warnings.warn(
@ -345,6 +346,117 @@ except:
torch_compiler_set_stance = None
def _uses_flash_attention_for_generation(config):
language_config_names = (
"text_config",
"llm_config",
"decoder_config",
"language_config",
"thinker_config",
"talker_config",
"decoder",
"generator",
)
non_language_config_names = (
"vision_config",
"audio_config",
"vision_encoder_config",
"audio_encoder_config",
"encoder_config",
"text_encoder",
)
def _mapping_uses_flash_attention(attn_implementation):
if not isinstance(attn_implementation, dict):
return _is_flash_attention_requested(attn_implementation)
language_implementations = [
implementation
for config_name, implementation in attn_implementation.items()
if config_name not in ("", *non_language_config_names) and implementation is not None
]
if language_implementations:
return any(map(_is_flash_attention_requested, language_implementations))
return _is_flash_attention_requested(attn_implementation.get(""))
def _get_text_config(current_config):
get_text_config = _config_get(current_config, "get_text_config", None)
if not callable(get_text_config):
return None
try:
return get_text_config()
except Exception:
return None
language_configs = []
pending_configs = [config]
visited_config_ids = set()
while pending_configs:
current_config = pending_configs.pop()
if id(current_config) in visited_config_ids:
continue
visited_config_ids.add(id(current_config))
text_config = _get_text_config(current_config)
if (
text_config is not None
and text_config is not current_config
and all(text_config is not item for item in language_configs)
):
language_configs.append(text_config)
nested_config_names = list(language_config_names)
declared_sub_configs = _config_get(current_config, "sub_configs", None)
if isinstance(declared_sub_configs, dict):
nested_config_names.extend(
config_name
for config_name in declared_sub_configs
if config_name not in nested_config_names
)
for config_name in nested_config_names:
nested_config = _config_get(current_config, config_name, None)
if nested_config is None or nested_config is current_config:
continue
pending_configs.append(nested_config)
nested_text_config = _get_text_config(nested_config)
if (
config_name in language_config_names
and (nested_text_config is None or nested_text_config is nested_config)
and all(nested_config is not item for item in language_configs)
):
language_configs.append(nested_config)
language_implementations = [
_config_get(language_config, config_field, None)
for language_config in language_configs
for config_field in ("_attn_implementation", "attn_implementation")
]
language_implementations = [
implementation for implementation in language_implementations if implementation is not None
]
if language_implementations:
return any(map(_mapping_uses_flash_attention, language_implementations))
return any(
_mapping_uses_flash_attention(_config_get(config, config_field, None))
for config_field in ("_attn_implementation", "attn_implementation")
)
def _clear_generation_caches(model):
for name, module in model.named_modules():
if hasattr(module, "_flex_attention_cache"):
try:
del module._flex_attention_cache
except:
pass
# Solves AttributeError: 'SlidingWindowLayer' object has no attribute 'max_batch_size'
if hasattr(module, "_cache") and "cache_utils" in str(module._cache.__class__):
try:
del module._cache
except:
pass
def unsloth_base_fast_generate(self, *args, **kwargs):
if len(args) != 0:
input_ids = args[0]
@ -444,6 +556,21 @@ def unsloth_base_fast_generate(self, *args, **kwargs):
# Prepare LoRA
# state_dict = convert_lora_modules(self, dtype = dtype)
# FlashAttention breaks on the forced static cache below (unfilled slots stay
# unmasked while decoding), so delegate after normalization but before it.
_clear_generation_caches(self)
if _uses_flash_attention_for_generation(self.config):
# Pin the literal "dynamic": None is merged back to the model default, and a
# static cache still arrives via kwargs / the caller's generation_config (TRL).
# The kwarg wins (update runs last); skip it when the caller passed a cache.
if kwargs.get("past_key_values") is None:
kwargs["cache_implementation"] = "dynamic"
try:
with torch.inference_mode(), autocaster:
return self._old_generate(*args, **kwargs)
finally:
_clear_generation_caches(self)
# Set compile dynamic shapes
torch._dynamo.mark_static(input_ids, 0)
torch._dynamo.mark_dynamic(input_ids, 1)
@ -491,36 +618,11 @@ def unsloth_base_fast_generate(self, *args, **kwargs):
if cache_implementation is not None:
kwargs["compile_config"] = _compile_config
# Delete cached Flex Attention masks to reset inference
for name, module in self.named_modules():
if hasattr(module, "_flex_attention_cache"):
try:
del module._flex_attention_cache
except:
pass
# Solves AttributeError: 'SlidingWindowLayer' object has no attribute 'max_batch_size'
if hasattr(module, "_cache") and "cache_utils" in str(module._cache.__class__):
try:
del module._cache
except:
pass
with torch.inference_mode(), autocaster:
output = self._old_generate(*args, **kwargs)
# Delete cached Flex Attention masks to reset inference
for name, module in self.named_modules():
if hasattr(module, "_flex_attention_cache"):
try:
del module._flex_attention_cache
except:
pass
# Solves AttributeError: 'SlidingWindowLayer' object has no attribute 'max_batch_size'
if hasattr(module, "_cache") and "cache_utils" in str(module._cache.__class__):
try:
del module._cache
except:
pass
try:
with torch.inference_mode(), autocaster:
output = self._old_generate(*args, **kwargs)
finally:
_clear_generation_caches(self)
# FastBaseModel.for_training(self)
return output