Skip to content

mlx-quantization-metal-kernels: fix batched inputs in host dispatch #1038

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
fenilsonani wants to merge 1 commit into from
Closed

mlx-quantization-metal-kernels: fix batched inputs in host dispatch #1038

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
81 changes: 58 additions & 23 deletions mlx-quantization-metal-kernels/quantization_mlx/fp_quantized.mm
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,41 +86,54 @@ void set_tensor(id<MTLComputeCommandEncoder> enc,
atIndex:index];
}

// Set batch metadata: ndim, shape, strides for the first `batch_dims` dims.
// If tensor has no batch dims, passes ndim=1, shape={1}, strides={0}.
// Batch metadata: batch_ndim = dim - matrix_dims (min 1); the full
// shape/strides are sent so kernels can read shape[batch_ndim] (qmv's M).
void set_batch_dims(id<MTLComputeCommandEncoder> enc,
const torch::Tensor& t,
int matrix_dims, // 2 for matmul, 1 for matvec
int matrix_dims,
int& idx) {
int total = t.dim();
int batch_ndim = std::max(total - matrix_dims, 0);
if (batch_ndim == 0) batch_ndim = 1;

std::vector<int> shape(batch_ndim, 1);
std::vector<int64_t> strides(batch_ndim, 0);
if (total > matrix_dims) {
for (int i = 0; i < total - matrix_dims; ++i) {
shape[i] = static_cast<int>(t.size(i));
strides[i] = t.stride(i);
}
std::vector<int> shape;
std::vector<int64_t> strides;
if (batch_ndim == 0) {
batch_ndim = 1;
shape.push_back(1);
strides.push_back(0);
}
for (int i = 0; i < total; ++i) {
shape.push_back(static_cast<int>(t.size(i)));
strides.push_back(t.stride(i));
}
[enc setBytes:&batch_ndim length:sizeof(int) atIndex:idx++];
[enc setBytes:shape.data() length:batch_ndim * sizeof(int) atIndex:idx++];
[enc setBytes:strides.data() length:batch_ndim * sizeof(int64_t) atIndex:idx++];
[enc setBytes:&batch_ndim length:sizeof(int) atIndex:idx++];
[enc setBytes:shape.data() length:shape.size() * sizeof(int) atIndex:idx++];
[enc setBytes:strides.data() length:strides.size() * sizeof(int64_t) atIndex:idx++];
}

// Set only strides for the first `batch_ndim` dimensions.
// Batch strides only; zeros when the tensor is shared across the batch.
void set_strides(id<MTLComputeCommandEncoder> enc,
const torch::Tensor& t,
int matrix_dims,
int batch_ndim,
int& idx) {
std::vector<int64_t> strides(batch_ndim, 0);
for (int i = 0; i < std::min(batch_ndim, (int)t.dim()); ++i) {
int batch_dims = std::max((int)t.dim() - matrix_dims, 0);
for (int i = 0; i < std::min(batch_dims, batch_ndim); ++i) {
strides[i] = t.stride(i);
}
[enc setBytes:strides.data() length:batch_ndim * sizeof(int64_t) atIndex:idx++];
}

// Kernels assume row-contiguous matrix dims; only batch dims may be strided.
torch::Tensor ensure_matrix_contiguous(const torch::Tensor& t) {
if (t.dim() < 2) {
return t.is_contiguous() ? t : t.contiguous();
}
bool ok = t.stride(-1) == 1 && t.stride(-2) == t.size(-1);
return ok ? t : t.contiguous();
}

} // namespace

// ===========================================================================
Expand Down Expand Up @@ -164,7 +177,7 @@ static void mxfp4_qmm_n_dispatch(
set_batch_dims(encoder, w, 2, idx);

// scales strides (13)
set_strides(encoder, scales, std::max((int)w.dim() - 2, 1), idx);
set_strides(encoder, scales, 2, std::max((int)w.dim() - 2, 1), idx);

// Grid
int grid_x = (N + 31) / 32;
Expand All @@ -185,10 +198,12 @@ static void mxfp4_qmv_dispatch(
const torch::Tensor& w,
const torch::Tensor& scales,
torch::Tensor& y) {
// For qmv, w is [N, K_packed], x is [..., K], y is [..., N]
bool batched = x.dim() > 1;
// For qmv, w is [N, K_packed], x is [..., M, K], y is [..., M, N].
// Rows are covered by grid.x (tid.x offsets), batch elements by grid.z.
bool batched = x.dim() > 2;
int K = static_cast<int>(x.size(-1));
int N = static_cast<int>(y.size(-1));
int M = x.dim() >= 2 ? static_cast<int>(x.size(-2)) : 1;

std::stringstream ss;
ss << "mxfp4_qmv_" << type_str(x.scalar_type())
Expand All @@ -212,18 +227,18 @@ static void mxfp4_qmv_dispatch(
[encoder setBytes:&N length:sizeof(int) atIndex:idx++]; // 5: out_vec_size

// Batch metadata (6,7,8 = x; 9,10,11 = w; 12 = s_strides)
set_batch_dims(encoder, x, 1, idx);
set_batch_dims(encoder, x, 2, idx);
set_batch_dims(encoder, w, 2, idx);
int w_batch_ndim = std::max((int)w.dim() - 2, 1);
set_strides(encoder, scales, w_batch_ndim, idx);
set_strides(encoder, scales, 2, w_batch_ndim, idx);

// Grid: each threadgroup handles results_per_simdgroup=4 output rows,
// with num_simdgroups=2.
int rows_per_tg = 8; // num_simdgroups(2) * results_per_simdgroup(4)
int grid_y = (N + rows_per_tg - 1) / rows_per_tg;
int grid_z = batched ? static_cast<int>(x.numel() / K) : 1;
int grid_z = batched ? static_cast<int>(x.numel() / (M * K)) : 1;

[encoder dispatchThreadgroups:MTLSizeMake(1, grid_y, grid_z)
[encoder dispatchThreadgroups:MTLSizeMake(M, grid_y, grid_z)
threadsPerThreadgroup:MTLSizeMake(32 * 2, 1, 1)]; // 2 simdgroups
}
}
Expand All @@ -234,6 +249,16 @@ static void mxfp4_qmv_dispatch(

at::Tensor mxfp4_qmm_n(
at::Tensor x, at::Tensor w, at::Tensor scales, int64_t output_features) {
x = ensure_matrix_contiguous(x);
w = ensure_matrix_contiguous(w);
scales = ensure_matrix_contiguous(scales);
// Dense batches run as one 2D GEMM; the grid.z path is ~Bx slower at M=1.
if (x.dim() > 2 && x.is_contiguous()) {
auto y = mxfp4_qmm_n(x.reshape({-1, x.size(-1)}), w, scales, output_features);
auto out_sizes = x.sizes().vec();
out_sizes.back() = y.size(-1);
return y.view(out_sizes);
}
// x: [..., M, K], y: [..., M, N]
auto out_sizes = x.sizes().vec();
out_sizes.back() = output_features;
Expand All @@ -244,6 +269,16 @@ static void mxfp4_qmv_dispatch(

at::Tensor mxfp4_qmv(
at::Tensor x, at::Tensor w, at::Tensor scales, int64_t output_features) {
x = ensure_matrix_contiguous(x);
w = ensure_matrix_contiguous(w);
scales = ensure_matrix_contiguous(scales);
// Dense batches run as one 2D GEMM; the grid.z path is ~Bx slower at M=1.
if (x.dim() > 2 && x.is_contiguous()) {
auto y = mxfp4_qmv(x.reshape({-1, x.size(-1)}), w, scales, output_features);
auto out_sizes = x.sizes().vec();
out_sizes.back() = y.size(-1);
return y.view(out_sizes);
}
// x: [..., K], y: [..., N]
auto out_sizes = x.sizes().vec();
out_sizes.back() = output_features;
Expand Down
96 changes: 73 additions & 23 deletions mlx-quantization-metal-kernels/quantization_mlx/quantized.mm
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,39 +86,54 @@ void set_tensor(id<MTLComputeCommandEncoder> enc,
atIndex:index];
}

// Batch dims: ndim, shape, strides of the first (tensor.dim() - matrix_dims) dims.
// Batch metadata: batch_ndim = dim - matrix_dims (min 1); the full
// shape/strides are sent so kernels can read shape[batch_ndim] (qmv's M).
void set_batch_dims(id<MTLComputeCommandEncoder> enc,
const torch::Tensor& t,
int matrix_dims,
int& idx) {
int total = t.dim();
int batch_ndim = std::max(total - matrix_dims, 0);
if (batch_ndim == 0) batch_ndim = 1;

std::vector<int> shape(batch_ndim, 1);
std::vector<int64_t> strides(batch_ndim, 0);
if (total > matrix_dims) {
for (int i = 0; i < total - matrix_dims; ++i) {
shape[i] = static_cast<int>(t.size(i));
strides[i] = t.stride(i);
}
std::vector<int> shape;
std::vector<int64_t> strides;
if (batch_ndim == 0) {
batch_ndim = 1;
shape.push_back(1);
strides.push_back(0);
}
[enc setBytes:&batch_ndim length:sizeof(int) atIndex:idx++];
[enc setBytes:shape.data() length:batch_ndim * sizeof(int) atIndex:idx++];
[enc setBytes:strides.data() length:batch_ndim * sizeof(int64_t) atIndex:idx++];
for (int i = 0; i < total; ++i) {
shape.push_back(static_cast<int>(t.size(i)));
strides.push_back(t.stride(i));
}
[enc setBytes:&batch_ndim length:sizeof(int) atIndex:idx++];
[enc setBytes:shape.data() length:shape.size() * sizeof(int) atIndex:idx++];
[enc setBytes:strides.data() length:strides.size() * sizeof(int64_t) atIndex:idx++];
}

// Batch strides only; zeros when the tensor is shared across the batch.
void set_strides(id<MTLComputeCommandEncoder> enc,
const torch::Tensor& t,
int matrix_dims,
int batch_ndim,
int& idx) {
std::vector<int64_t> strides(batch_ndim, 0);
for (int i = 0; i < std::min(batch_ndim, (int)t.dim()); ++i) {
int batch_dims = std::max((int)t.dim() - matrix_dims, 0);
for (int i = 0; i < std::min(batch_dims, batch_ndim); ++i) {
strides[i] = t.stride(i);
}
[enc setBytes:strides.data() length:batch_ndim * sizeof(int64_t) atIndex:idx++];
}

// Kernels assume row-contiguous matrix dims; only batch dims may be strided.
torch::Tensor ensure_matrix_contiguous(const torch::Tensor& t) {
if (t.dim() < 2) {
return t.is_contiguous() ? t : t.contiguous();
}
bool ok = t.stride(-1) == 1 && t.stride(-2) == t.size(-1);
return ok ? t : t.contiguous();
}

} // namespace

// ===========================================================================
Expand All @@ -139,9 +154,11 @@ static void affine_qmv_dispatch(
torch::Tensor& y,
int group_size,
int bits) {
bool batched = x.dim() > 1;
// Rows are covered by grid.x (tid.x offsets), batch elements by grid.z.
bool batched = x.dim() > 2;
int K = static_cast<int>(x.size(-1));
int N = static_cast<int>(y.size(-1));
int M = x.dim() >= 2 ? static_cast<int>(x.size(-2)) : 1;

std::stringstream ss;
ss << "affine_qmv_" << type_str(x.scalar_type())
Expand All @@ -167,21 +184,21 @@ static void affine_qmv_dispatch(
[encoder setBytes:&N length:sizeof(int) atIndex:idx++]; // 6: out_vec_size

// x batch dims (7,8,9)
set_batch_dims(encoder, x, 1, idx);
set_batch_dims(encoder, x, 2, idx);
// w batch dims (10,11,12)
set_batch_dims(encoder, w, 2, idx);
// scales strides (13)
int w_batch_ndim = std::max((int)w.dim() - 2, 1);
set_strides(encoder, scales, w_batch_ndim, idx);
set_strides(encoder, scales, 2, w_batch_ndim, idx);
// biases strides (14)
set_strides(encoder, biases, w_batch_ndim, idx);
set_strides(encoder, biases, 2, w_batch_ndim, idx);

// Grid: each TG handles 8 output rows (num_simdgroups=2 * results_per=4)
int rows_per_tg = 8;
int grid_y = (N + rows_per_tg - 1) / rows_per_tg;
int grid_z = batched ? static_cast<int>(x.numel() / K) : 1;
int grid_z = batched ? static_cast<int>(x.numel() / (M * K)) : 1;

[encoder dispatchThreadgroups:MTLSizeMake(1, grid_y, grid_z)
[encoder dispatchThreadgroups:MTLSizeMake(M, grid_y, grid_z)
threadsPerThreadgroup:MTLSizeMake(32 * 2, 1, 1)];
}
}
Expand Down Expand Up @@ -241,9 +258,9 @@ static void affine_qmm_t_dispatch(
set_batch_dims(encoder, w, 2, idx);
// scales strides (14)
int w_batch_ndim = std::max((int)w.dim() - 2, 1);
set_strides(encoder, scales, w_batch_ndim, idx);
set_strides(encoder, scales, 2, w_batch_ndim, idx);
// biases strides (15)
set_strides(encoder, biases, w_batch_ndim, idx);
set_strides(encoder, biases, 2, w_batch_ndim, idx);

// BM=BK=BN=32 (kernel defaults)
int grid_x = (N + 31) / 32;
Expand Down Expand Up @@ -303,9 +320,9 @@ static void affine_qmm_n_dispatch(
set_batch_dims(encoder, w, 2, idx);
// scales strides (14)
int w_batch_ndim = std::max((int)w.dim() - 2, 1);
set_strides(encoder, scales, w_batch_ndim, idx);
set_strides(encoder, scales, 2, w_batch_ndim, idx);
// biases strides (15)
set_strides(encoder, biases, w_batch_ndim, idx);
set_strides(encoder, biases, 2, w_batch_ndim, idx);

int grid_x = (N + 31) / 32;
int grid_y = (M + 31) / 32;
Expand All @@ -323,6 +340,17 @@ static void affine_qmm_n_dispatch(
at::Tensor affine_qmv(
at::Tensor x, at::Tensor w, at::Tensor scales, at::Tensor biases,
int64_t group_size, int64_t bits, int64_t output_features) {
x = ensure_matrix_contiguous(x);
w = ensure_matrix_contiguous(w);
scales = ensure_matrix_contiguous(scales);
biases = ensure_matrix_contiguous(biases);
// Dense batches run as one 2D GEMM; the grid.z path is ~Bx slower at M=1.
if (x.dim() > 2 && x.is_contiguous()) {
auto y = affine_qmv(x.reshape({-1, x.size(-1)}), w, scales, biases, group_size, bits, output_features);
auto out_sizes = x.sizes().vec();
out_sizes.back() = y.size(-1);
return y.view(out_sizes);
}
auto out_sizes = x.sizes().vec();
out_sizes.back() = output_features;
auto y = torch::zeros(out_sizes, x.options());
Expand All @@ -334,6 +362,17 @@ static void affine_qmm_n_dispatch(
at::Tensor affine_qmm_t(
at::Tensor x, at::Tensor w, at::Tensor scales, at::Tensor biases,
int64_t group_size, int64_t bits) {
x = ensure_matrix_contiguous(x);
w = ensure_matrix_contiguous(w);
scales = ensure_matrix_contiguous(scales);
biases = ensure_matrix_contiguous(biases);
// Dense batches run as one 2D GEMM; the grid.z path is ~Bx slower at M=1.
if (x.dim() > 2 && x.is_contiguous()) {
auto y = affine_qmm_t(x.reshape({-1, x.size(-1)}), w, scales, biases, group_size, bits);
auto out_sizes = x.sizes().vec();
out_sizes.back() = y.size(-1);
return y.view(out_sizes);
}
// For transposed weight: w = [N, K_packed], N = w.size(0)
int64_t N = w.size(0);
auto out_sizes = x.sizes().vec();
Expand All @@ -347,6 +386,17 @@ static void affine_qmm_n_dispatch(
at::Tensor affine_qmm_n(
at::Tensor x, at::Tensor w, at::Tensor scales, at::Tensor biases,
int64_t group_size, int64_t bits, int64_t output_features) {
x = ensure_matrix_contiguous(x);
w = ensure_matrix_contiguous(w);
scales = ensure_matrix_contiguous(scales);
biases = ensure_matrix_contiguous(biases);
// Dense batches run as one 2D GEMM; the grid.z path is ~Bx slower at M=1.
if (x.dim() > 2 && x.is_contiguous()) {
auto y = affine_qmm_n(x.reshape({-1, x.size(-1)}), w, scales, biases, group_size, bits, output_features);
auto out_sizes = x.sizes().vec();
out_sizes.back() = y.size(-1);
return y.view(out_sizes);
}
// For non-transposed weight: w = [K_packed, N_packed]
// output_features = N (the logical unpacked output dimension)
auto out_sizes = x.sizes().vec();
Expand Down
Loading
Loading