zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-25 23:38:10 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

start

Browse Source
This commit is contained in:
Alex Barron
2024-12-02 16:19:29 -08:00
parent 9d40e521d7
commit 890fdd1ef0
4 changed files with 59 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
mlx/backend/metal/kernels/quantized.h
View File

@@ -2006,7 +2006,8 @@ template <typename T, const int group_size, const int bits>
uint2 index [[thread_position_in_grid]],
uint2 grid_dim [[threads_per_grid]]) {
constexpr T eps = T(1e-7);
constexpr int simd_size = 32;
constexpr bool use_quads = group_size <= 16;
constexpr int simd_size = use_quads ? 4 : 32;
constexpr T n_bins = (1 << bits) - 1;
constexpr int packs_per_int = bits == 3 ? 8 : bits == 6 ? 4 : 8 / bits;
constexpr int values_per_reduce = group_size / simd_size;
@@ -2038,8 +2039,13 @@ template <typename T, const int group_size, const int bits>
w_max = max(w_max, val);
}
w_min = simd_min(w_min);
w_max = simd_max(w_max);
if (use_quads) {
w_min = quad_min(w_min);
w_max = quad_max(w_max);
} else {
w_min = simd_min(w_min);
w_max = simd_max(w_max);
}
T scale = max((w_max - w_min) / n_bins, eps);
bool side = abs(w_min) > abs(w_max);

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

@@ -308,13 +308,13 @@ void qmm_op(
group_dims = MTL::Size(simdgroup_size, 1, 1);
grid_dims = MTL::Size((O + bo - 1) / bo, B, N);
quad = true;
} else if (B < 6 && O % 8 == 0 && D % 512 == 0 && D >= 512) {
} else if (B > 0 && O % 8 == 0 && D % 512 == 0 && D >= 512) {
name += "qmv_fast";
int bo = 8;
int bd = 32;
group_dims = MTL::Size(bd, 2, 1);
grid_dims = MTL::Size(O / bo, B, N);
} else if (B < 6) {
} else if (B > 0) {
name += "qmv";
int bo = 8;
int bd = 32;
@@ -445,7 +445,8 @@ void fast::AffineQuantize::eval_gpu(
// Treat uint32 as uint8 in kernel
constexpr int uint8_per_uint32 = 4;
constexpr int simd_size = 32;
// Use quads for small group sizes
int simd_size = group_size_ <= 16 ? 4 : 32;
int packs_per_int = bits_ == 3 ? 8 : bits_ == 6 ? 4 : 8 / bits_;
int per_thread = dequantize_ ? packs_per_int : group_size_ / simd_size;
size_t nthreads =

3
mlx/fast.cpp
View File

@@ -729,7 +729,8 @@ std::tuple<array, array, array>
affine_quantize(const array& w, int group_size, int bits, StreamOrDevice s_) {
auto s = to_stream(s_);
if (group_size != 32 && group_size != 64 && group_size != 128) {
if (group_size != 16 && group_size != 32 && group_size != 64 &&
group_size != 128) {
std::ostringstream msg;
msg << "[quantize] The requested group size " << group_size
<< " is not supported. The supported group sizes are 64 and 128.";