diff --git a/tests/test_fa2_fast_generate_bypass.py b/tests/test_fa2_fast_generate_bypass.py new file mode 100644 index 0000000000..376364f39c --- /dev/null +++ b/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") diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 729c191e83..dc8032e7ba 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -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