ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update prepare 4d causal attention call (#2678)

Browse source
This commit is contained in:
DoubleMathew 2025-06-04 14:58:50 -05:00 • committed by GitHub
commit e160305b66
1 changed files with 24 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

32
unsloth/models/llama.py
View file

@ -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]]