huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-20 20:09:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #1606 from FL33TW00D/feature/larger-batches

Browse Source 
fix: larger batches
This commit is contained in:
Christopher Fleetwood
2024-01-29 15:31:10 +00:00
committed by GitHub
parent fd7c856564 b6afb46601
commit 6d83d42efb
2 changed files with 6 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
candle-metal-kernels/src/lib.rs
View File

@ -1364,13 +1364,12 @@ pub fn call_gemm(
// TODO byte_stride_d // TODO byte_stride_d
let byte_stride_d = 0; let byte_stride_d = 0;
let mut buffer: Vec<u64> = Vec::with_capacity(b * 4); let buffer: Vec<u64> = vec![
for i in 0..b { byte_stride_a as _,
buffer.push((i * byte_stride_a) as u64); byte_stride_b as _,
buffer.push((i * byte_stride_b) as u64); byte_stride_c as _,
buffer.push((i * byte_stride_c) as u64); byte_stride_d as _,
buffer.push((i * byte_stride_d) as u64); ];
}
encoder.set_bytes( encoder.set_bytes(
10, 10,
(buffer.len() * core::mem::size_of::<u64>()) as NSUInteger, (buffer.len() * core::mem::size_of::<u64>()) as NSUInteger,

BIN
candle-metal-kernels/src/libMetalFlashAttention.metallib
View File

Binary file not shown.