Commit graph

1 commit

Author SHA1 Message Date
Piotr Wasiewicz
62d3438b99
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>
2026-07-26 23:07:33 -07:00