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

Update llama.py

Browse source
This commit is contained in:
Daniel Han 2025-03-21 16:42:32 -07:00
commit 1e103a3e23
1 changed files with 2 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

2
unsloth/models/llama.py
View file

@ -929,6 +929,8 @@ def LlamaModel_fast_forward_inference(
temp_mlp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda:0")
temp_gate, temp_up = temp_mlp[0], temp_mlp[1]
print(type(past_key_values), len(past_key_values))
seq_len = past_key_values[0][0].shape[-2]
if bsz != 1:
attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(