(&self, a: &Tensor, b: &Tensor)
| 984 | } |
| 985 | |
| 986 | fn tensor_matmul(&self, a: &Tensor, b: &Tensor) -> Result<Tensor> { |
| 987 | let a_dims = a.dims(); |
| 988 | let b_dims = b.dims(); |
| 989 | let a_rank = a_dims.len(); |
| 990 | let b_rank = b_dims.len(); |
| 991 | |
| 992 | let m = a_dims[a_rank - 2]; |
| 993 | let k = a_dims[a_rank - 1]; |
| 994 | let n = b_dims[b_rank - 1]; |
| 995 | |
| 996 | let a_batch: usize = a_dims[..a_rank - 2].iter().product(); |
| 997 | let b_batch: usize = b_dims[..b_rank - 2].iter().product(); |
| 998 | let batch = a_batch.max(b_batch); |
| 999 | |
| 1000 | let mk = m * k; |
| 1001 | let kn = k * n; |
| 1002 | let mn = m * n; |
| 1003 | |
| 1004 | if batch == 1 { |
| 1005 | // get_or_upload handles dtype conversion + contiguity |
| 1006 | let buf_a = self.get_or_upload(a)?; |
| 1007 | let buf_b = self.get_or_upload(b)?; |
| 1008 | // Check for F16 weight buffer (halves GEMV bandwidth) |
| 1009 | let f16_buf = if m == 1 { |
| 1010 | let vk = Self::view_key(b); |
| 1011 | vk.and_then(|key| { |
| 1012 | let cache = self.weight_cache.lock().unwrap(); |
| 1013 | cache.f16_views.get(&key).cloned() |
| 1014 | }) |
| 1015 | } else { |
| 1016 | None |
| 1017 | }; |
| 1018 | let (out, buf_out) = self.gpu_matmul(&buf_a, &buf_b, f16_buf.as_ref(), m, k, n); |
| 1019 | let mut out_shape: Vec<usize> = a_dims[..a_rank - 2].to_vec(); |
| 1020 | out_shape.push(m); |
| 1021 | out_shape.push(n); |
| 1022 | let tensor = Tensor::from_vec(out, out_shape.as_slice(), &Device::Cpu)?; |
| 1023 | self.cache_activation(tensor.id(), buf_out); |
| 1024 | return Ok(tensor); |
| 1025 | } |
| 1026 | |
| 1027 | // Batch path: record ALL batch dispatches in one command buffer (1 fence wait). |
| 1028 | let a = a.to_dtype(DType::F32)?.contiguous()?; |
| 1029 | let b = b.to_dtype(DType::F32)?.contiguous()?; |
| 1030 | let a_data: Vec<f32> = a.flatten_all()?.to_vec1()?; |
| 1031 | let b_data: Vec<f32> = b.flatten_all()?.to_vec1()?; |
| 1032 | |
| 1033 | // Upload all batch slices first |
| 1034 | let mut a_bufs = Vec::with_capacity(batch); |
| 1035 | let mut b_bufs = Vec::with_capacity(batch); |
| 1036 | let mut out_bufs = Vec::with_capacity(batch); |
| 1037 | for i in 0..batch { |
| 1038 | let a_off = if a_batch == 1 { 0 } else { i * mk }; |
| 1039 | let b_off = if b_batch == 1 { 0 } else { i * kn }; |
| 1040 | a_bufs.push(self.upload_uncached(&a_data[a_off..a_off + mk])); |
| 1041 | b_bufs.push(self.upload_uncached(&b_data[b_off..b_off + kn])); |
| 1042 | out_bufs.push(self.alloc_output(mn)); |
| 1043 | } |
no test coverage detected