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

Quantized matmul fix (#677)

Browse Source 
* Fix qmv for small or unaligned matrices

* Fix qmm
This commit is contained in:
Angelos Katharopoulos
2024-02-12 18:54:21 -08:00
committed by GitHub
parent 4cc70290f7
commit 40c108766b
3 changed files with 81 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
mlx/backend/metal/quantized.cpp
View File

@@ -55,7 +55,7 @@ void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
int bo = std::min(32, O);
int bd = 32;
MTL::Size group_dims = MTL::Size(bd, bo, 1);
MTL::Size grid_dims = MTL::Size(1, O / bo, B);
MTL::Size grid_dims = MTL::Size(1, (O + bo - 1) / bo, B);
set_array_buffer(compute_encoder, w, 0);
set_array_buffer(compute_encoder, scales, 1);