(
&self,
q: &Tensor,
k: &Tensor,
v: &Tensor,
scale: f32,
causal: bool,
)
| 1257 | } |
| 1258 | |
| 1259 | fn attention( |
| 1260 | &self, |
| 1261 | q: &Tensor, |
| 1262 | k: &Tensor, |
| 1263 | v: &Tensor, |
| 1264 | scale: f32, |
| 1265 | causal: bool, |
| 1266 | ) -> Result<Tensor> { |
| 1267 | let orig_dtype = q.dtype(); |
| 1268 | let q = q.to_dtype(DType::F32)?; |
| 1269 | let k = k.to_dtype(DType::F32)?; |
| 1270 | let v = v.to_dtype(DType::F32)?; |
| 1271 | |
| 1272 | // Q @ K^T |
| 1273 | let attn = self.tensor_matmul(&q, &k.t()?)?; |
| 1274 | |
| 1275 | // GPU scaled softmax (fuses scale + optional causal mask + softmax) |
| 1276 | let seq_len = q.dim(q.dims().len() - 2)?; |
| 1277 | let kv_len = k.dim(k.dims().len() - 2)?; |
| 1278 | let total_rows = attn.elem_count() / kv_len; |
| 1279 | let causal_seq_len = if causal { seq_len } else { 0 }; |
| 1280 | let attn = self.dispatch_softmax(&attn, total_rows, kv_len, scale, causal_seq_len)?; |
| 1281 | |
| 1282 | // Attn @ V |
| 1283 | let out = self.tensor_matmul(&attn, &v)?; |
| 1284 | out.to_dtype(orig_dtype) |
| 1285 | } |
| 1286 | |
| 1287 | fn silu_mul(&self, gate: &Tensor, up: &Tensor) -> Result<Tensor> { |
| 1288 | let n = gate.elem_count(); |
nothing calls this directly
no test coverage detected