zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-19 15:41:13 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Change order in weight packing

Browse Source
This commit is contained in:
Angelos Katharopoulos 2024-12-14 22:51:41 -08:00
parent bf6dc54110
commit fd161aa31f
2 changed files with 7 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -2171,11 +2171,10 @@ inline vec<U, 4> partial_qdot_vec(const thread U* x, vec<uint32_t, 4> w) {
else if (bits == 4) {
for (int i = 0; i < 4; i++) {
auto ws = as_type<vec<uint16_t, 2>>(w[i]);
for (int j = 0; j < 2; j++) {
accum[i] +=
(x[4 * j + 0] * (ws[j] & 0x000f) + x[4 * j + 1] * (ws[j] & 0x00f0) +
x[4 * j + 2] * (ws[j] & 0x0f00) + x[4 * j + 3] * (ws[j] & 0xf000));
auto ws = as_type<vec<uint8_t, 4>>(w[i]);
for (int j = 0; j < 4; j++) {
accum[j] +=
x[2 * i + 0] * (ws[j] & 0x0f) + x[2 * i + 1] * (ws[j] & 0xf0);
}
}
}
@ -2184,7 +2183,7 @@ inline vec<U, 4> partial_qdot_vec(const thread U* x, vec<uint32_t, 4> w) {
for (int i = 0; i < 4; i++) {
auto ws = as_type<vec<uint8_t, 4>>(w[i]);
for (int j = 0; j < 4; j++) {
accum[i] += x[j] * ws[j];
accum[j] += x[i] * ws[j];
}
}
}

2
mlx/ops.cpp
View File

@ -3795,9 +3795,11 @@ std::tuple<array, array, std::optional<array>> quantize(
scales = moveaxis(scales, -2, -1, s);
scales = flatten(scales, -2, -1, s);
wq = view(wq, uint8, s);
wq = unflatten(wq, -2, {-1, 4}, s);
wq = moveaxis(wq, -2, -1, s);
wq = flatten(wq, -2, -1, s);
wq = view(wq, uint32, s);
return std::make_tuple(wq, scales, std::nullopt);
}