diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 53c956ce1a..cba8a68565 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -121,16 +121,32 @@ def _fast_prepare_inputs_for_generation(self, input_ids, attention_mask=None, **
                 base_model = getattr(base_model, base_model.base_model_prefix)
                 
             if hasattr(base_model, "_prepare_4d_causal_attention_mask_with_cache_position"):
+                def needs_device_kw(fn) -> bool:
+                    try:
+                        sig = inspect.signature(inspect.unwrap(fn))
+                        return "device" in sig.parameters
+                    except:
+                        # transformers <= 4.51.3 includes device arg but > 4.51.3 does not
+                        return transformers_version < Version("4.52.0")
+
+                kwargs = {
+                    "sequence_length": 1,
+                    "target_length": cache_length,
+                    "dtype": self.dtype,
+                    "cache_position": torch.arange(cache_length, cache_length+1, device=input_ids.device),
+                    "batch_size": bs,
+                    "config": self.config,
+                    "past_key_values": past_key_values,
+                }
+                try:
+                    if needs_device_kw(base_model._prepare_4d_causal_attention_mask_with_cache_position):
+                        kwargs["device"] = input_ids.device
+                except:
+                    print(f"Unsloth: Could not inspect signature of {base_model._prepare_4d_causal_attention_mask_with_cache_position}")
+
                 attention_mask = base_model._prepare_4d_causal_attention_mask_with_cache_position(
                     attention_mask,
-                    sequence_length=1,
-                    target_length=cache_length,
-                    dtype=self.dtype,
-                    device=input_ids.device,
-                    cache_position=torch.arange(cache_length, cache_length+1, device=input_ids.device),
-                    batch_size=bs,
-                    config=self.config,
-                    past_key_values=past_key_values,
+                    **kwargs,
                 )
             else:
                 attention_mask = attention_mask[:,[-1]]