From 8886f3aa3197922361580ca7f2870c415ee28afe Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 18 Mar 2025 23:42:26 -0700
Subject: [PATCH] Update vision.py

---
 unsloth/models/vision.py | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index 66e133fa27..8dcd36790e 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -76,7 +76,13 @@ NUM_LOGITS_TO_KEEP = dict()
 global PROMPT_LOOPKUP
 PROMPT_LOOPKUP = dict()
 
-from transformers import GenerationConfig
+from transformers import GenerationConfig, CompileConfig, HybridCache
+_compile_config = CompileConfig(
+    fullgraph = False,
+    dynamic = None,
+    mode = "reduce-overhead",
+)
+_compile_config.disable = True # Must set manually
 
 def unsloth_base_fast_generate(
     self,
@@ -158,6 +164,16 @@ def unsloth_base_fast_generate(
         torch._dynamo.mark_dynamic(kwargs["attention_mask"], 1)
     pass
 
+    # Fix generation_config
+    cache_implementation = getattr(self.config, "cache_implementation", "static")
+    if "generation_config" in kwargs:
+        kwargs["generation_config"].cache_implementation = cache_implementation
+        kwargs["generation_config"].compile_config = _compile_config
+    else:
+        kwargs["cache_implementation"] = cache_implementation
+        kwargs["compile_config"] = _compile_config
+    pass
+
     # Mixed precision autocast
     if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1":
         autocaster = torch.autocast(device_type = "cuda", dtype = dtype)