zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

[CUDA] Matmul utils initial commit (#2441)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2025-08-01 14:22:25 -07:00
committed by GitHub
parent 86258f292f
commit be9bc96da4
32 changed files with 856 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
mlx/backend/cuda/gemms/cublas_batched_gemm_12_9.cu
View File

@@ -108,6 +108,7 @@ void Matmul::run_batched(
cu::set_mm_device_pointers,
cuda::ceil_div(pointers.size(), block_size),
block_size,
0,
pointers.data<int8_t*>(),
a.data<int8_t>(),
b.data<int8_t>(),
@@ -168,6 +169,7 @@ void Matmul::run_batched(
cu::set_addmm_device_pointers,
cuda::ceil_div(pointers.size(), block_size),
block_size,
0,
pointers.data<int8_t*>(),
a.data<int8_t>(),
b.data<int8_t>(),

2
mlx/backend/cuda/gemms/gemv.cu
View File

@@ -143,6 +143,7 @@ void gemv(
kernel,
num_blocks_x,
block_dims,
0,
mat,
vec,
out.data<DataType>(),
@@ -154,6 +155,7 @@ void gemv(
kernel,
dim3{num_blocks_x, batch_count},
block_dims,
0,
mat,
vec,
out.data<DataType>(),