huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 18:48:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

use softmax_last_dim (metal and cuda kernel) in llama attention layer (#2572)

Browse Source
This commit is contained in:
Zack Angelo
2024-10-23 11:07:09 -07:00
committed by GitHub
parent 7c09215ef4
commit a2e9d41b20
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
candle-transformers/src/models/llama.rs
View File

@ -341,7 +341,8 @@ impl CausalSelfAttention {
let mask = cache.mask(seq_len)?.broadcast_as(att.shape())?; let mask = cache.mask(seq_len)?.broadcast_as(att.shape())?;
masked_fill(&att, &mask, f32::NEG_INFINITY)? masked_fill(&att, &mask, f32::NEG_INFINITY)?
}; };
let att = candle_nn::ops::softmax(&att, D::Minus1)?;
let att = candle_nn::ops::softmax_last_dim(&att)?;
// Convert to contiguous as matmul doesn't support strided vs for now. // Convert to contiguous as matmul doesn't support strided vs for now.
att.matmul(&v.contiguous()?)?.to_dtype(in_dtype)? att.matmul(&v.contiguous()?)?.to_dtype(in_dtype)?
}; };