From 3accb89df613cab8467fbf2654d59154b58fbbcd Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 13 Mar 2025 17:57:46 -0700
Subject: [PATCH] vLLM fixes

---
 unsloth/models/llama.py  | 20 ++++++++++++++++++++
 unsloth/models/loader.py |  4 +---
 unsloth/models/vision.py | 14 +++++++++++++-
 3 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 7000739850..0bb8c4a771 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -1663,6 +1663,10 @@ class FastLlamaModel:
             if platform.system().lower() == 'windows':
                 print("Unsloth: vLLM does not work in Windows! Will use Unsloth inference!")
                 fast_inference = False
+            major_version, minor_version = torch.cuda.get_device_capability()
+            if major_version < 7:
+                print("Unsloth: vLLM does not work on older GPUs - will switch to Unsloth inference!")
+                fast_inference = False
         pass
 
         if token is None: token = get_token()
@@ -1786,6 +1790,8 @@ class FastLlamaModel:
                 attn_implementation     = "eager",
                 **kwargs,
             )
+            model.fast_generate = model.generate
+            model.fast_generate_batches = None
         else:
             from unsloth_zoo.vllm_utils import (
                 load_vllm,
@@ -1804,6 +1810,7 @@ class FastLlamaModel:
                 enable_lora            = True,
                 max_lora_rank          = max_lora_rank,
                 disable_log_stats      = disable_log_stats,
+                use_bitsandbytes       = load_in_4bit,
             )
             for allowed_arg in allowed_args:
                 if allowed_arg not in load_vllm_kwargs and allowed_arg in kwargs:
@@ -2651,6 +2658,19 @@ class FastLlamaModel:
             torch.cuda.empty_cache()
         pass
 
+        # Patch for fast inference
+        vllm_engine = getattr(model, "vllm_engine")
+        if vllm_engine is not None:
+            model.vllm_engine = vllm_engine
+            model.fast_generate = vllm_fast_generate
+            model.fast_generate_batches = vllm_fast_generate_batches
+
+            # Also saving and loading LoRA
+            from unsloth_zoo.vllm_utils import save_lora, load_lora
+            model.save_lora = functools.partial(save_lora, model)
+            model.load_lora = functools.partial(load_lora, model)
+        pass
+
         # Add for_inference and for_training
         model.for_training  = functools.partial(FastLlamaModel.for_training,  model)
         model.for_inference = functools.partial(FastLlamaModel.for_inference, model)
diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index 1b54c8c7fc..ae9e9dfad2 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -405,7 +405,6 @@ class FastLanguageModel(FastLlamaModel):
         if is_peft:
             # From https://github.com/huggingface/peft/issues/184
             # Now add PEFT adapters
-            model.enable_input_require_grads()
             model = PeftModel.from_pretrained(
                 model,
                 old_model_name,
@@ -668,7 +667,7 @@ class FastModel(FastBaseModel):
             use_gradient_checkpointing = use_gradient_checkpointing,
             *args, **kwargs,
         )
-        
+
         if resize_model_vocab is not None:
             model.resize_token_embeddings(resize_model_vocab)
         pass
@@ -703,7 +702,6 @@ class FastModel(FastBaseModel):
         if is_peft:
             # From https://github.com/huggingface/peft/issues/184
             # Now add PEFT adapters
-            model.enable_input_require_grads()
             model = PeftModel.from_pretrained(
                 model,
                 old_model_name,
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index 2ef9d2ee99..f0d5a0930c 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -68,6 +68,9 @@ __all__ = [
 global FORCE_FLOAT32
 FORCE_FLOAT32 = ["gemma3"]
 
+global FORCE_EAGER_ATTENTION
+FORCE_EAGER_ATTENTION = ["pixtral"]
+
 
 def unsloth_base_fast_generate(
     self,
@@ -193,6 +196,15 @@ class FastBaseModel:
                 break
         pass
 
+        global FORCE_EAGER_ATTENTION
+        attn_implementation = "sdpa"
+        for disable_sdpa_name in FORCE_EAGER_ATTENTION:
+            if disable_sdpa_name.lower() == model_type_arch.lower():
+                print(f"Unsloth: {model_type_arch} does not support SDPA - switching to eager!")
+                attn_implementation = "eager"
+                break
+        pass
+
         bnb_config = None
         if full_finetuning and (load_in_4bit or load_in_8bit):
             print("Unsloth: You selected full finetuning support, but 4bit / 8bit is enabled - disabling LoRA / QLoRA.")
@@ -249,7 +261,7 @@ class FastBaseModel:
             # quantization_config   = bnb_config,
             token                   = token,
             trust_remote_code       = trust_remote_code,
-            attn_implementation     = "sdpa", #[TODO] Pixtral for eg fails
+            attn_implementation     = attn_implementation,
             **kwargs,
         )
         # Return old flag