diff --git a/mlx-quantization-metal-kernels/quantization_mlx/fp_quantized.mm b/mlx-quantization-metal-kernels/quantization_mlx/fp_quantized.mm index f9c6dfc6..2d564cc8 100644 --- a/mlx-quantization-metal-kernels/quantization_mlx/fp_quantized.mm +++ b/mlx-quantization-metal-kernels/quantization_mlx/fp_quantized.mm @@ -86,41 +86,54 @@ void set_tensor(id 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 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 shape(batch_ndim, 1); - std::vector strides(batch_ndim, 0); - if (total > matrix_dims) { - for (int i = 0; i < total - matrix_dims; ++i) { - shape[i] = static_cast(t.size(i)); - strides[i] = t.stride(i); - } + std::vector shape; + std::vector 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(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 enc, const torch::Tensor& t, + int matrix_dims, int batch_ndim, int& idx) { std::vector 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 // =========================================================================== @@ -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; @@ -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(x.size(-1)); int N = static_cast(y.size(-1)); + int M = x.dim() >= 2 ? static_cast(x.size(-2)) : 1; std::stringstream ss; ss << "mxfp4_qmv_" << type_str(x.scalar_type()) @@ -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(x.numel() / K) : 1; + int grid_z = batched ? static_cast(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 } } @@ -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; @@ -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; diff --git a/mlx-quantization-metal-kernels/quantization_mlx/quantized.mm b/mlx-quantization-metal-kernels/quantization_mlx/quantized.mm index bddce90c..53d9b15c 100644 --- a/mlx-quantization-metal-kernels/quantization_mlx/quantized.mm +++ b/mlx-quantization-metal-kernels/quantization_mlx/quantized.mm @@ -86,39 +86,54 @@ void set_tensor(id 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 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 shape(batch_ndim, 1); - std::vector strides(batch_ndim, 0); - if (total > matrix_dims) { - for (int i = 0; i < total - matrix_dims; ++i) { - shape[i] = static_cast(t.size(i)); - strides[i] = t.stride(i); - } + std::vector shape; + std::vector 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(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 enc, const torch::Tensor& t, + int matrix_dims, int batch_ndim, int& idx) { std::vector 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 // =========================================================================== @@ -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(x.size(-1)); int N = static_cast(y.size(-1)); + int M = x.dim() >= 2 ? static_cast(x.size(-2)) : 1; std::stringstream ss; ss << "affine_qmv_" << type_str(x.scalar_type()) @@ -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(x.numel() / K) : 1; + int grid_z = batched ? static_cast(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)]; } } @@ -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; @@ -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; @@ -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()); @@ -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(); @@ -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(); diff --git a/mlx-quantization-metal-kernels/quantization_mlx/quantized_nax.mm b/mlx-quantization-metal-kernels/quantization_mlx/quantized_nax.mm index 8f4365ab..cb0a83d6 100644 --- a/mlx-quantization-metal-kernels/quantization_mlx/quantized_nax.mm +++ b/mlx-quantization-metal-kernels/quantization_mlx/quantized_nax.mm @@ -83,38 +83,54 @@ void set_tensor(id enc, atIndex:index]; } +// 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 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 shape(batch_ndim, 1); - std::vector strides(batch_ndim, 0); - if (total > matrix_dims) { - for (int i = 0; i < total - matrix_dims; ++i) { - shape[i] = static_cast(t.size(i)); - strides[i] = t.stride(i); - } + std::vector shape; + std::vector 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(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 enc, const torch::Tensor& t, + int matrix_dims, int batch_ndim, int& idx) { std::vector 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 // =========================================================================== @@ -170,9 +186,9 @@ static void affine_qmm_t_nax_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=BN=BK=64 for NAX int grid_x = (N + 63) / 64; @@ -233,9 +249,9 @@ static void affine_qmm_n_nax_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 + 63) / 64; int grid_y = (M + 63) / 64; @@ -308,6 +324,17 @@ static void affine_gather_qmm_rhs_nax_dispatch( at::Tensor affine_qmm_t_nax( 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_nax(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); + } // Transposed weight: w = [N, K_packed], N = w.size(0) int64_t N = w.size(0); auto out_sizes = x.sizes().vec(); @@ -321,6 +348,17 @@ static void affine_gather_qmm_rhs_nax_dispatch( at::Tensor affine_qmm_n_nax( 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_nax(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()); @@ -334,6 +372,10 @@ static void affine_gather_qmm_rhs_nax_dispatch( at::Tensor indices, int64_t group_size, int64_t bits, int64_t output_features, bool transpose) { + x = ensure_matrix_contiguous(x); + w = ensure_matrix_contiguous(w); + scales = ensure_matrix_contiguous(scales); + biases = ensure_matrix_contiguous(biases); // x: [M, K], y: [M, N] int64_t M = x.size(0); auto y = torch::zeros({M, output_features}, x.options()); diff --git a/mlx-quantization-metal-kernels/tests/test_quantization.py b/mlx-quantization-metal-kernels/tests/test_quantization.py index c8e7cc71..655b7b2d 100644 --- a/mlx-quantization-metal-kernels/tests/test_quantization.py +++ b/mlx-quantization-metal-kernels/tests/test_quantization.py @@ -233,6 +233,74 @@ def test_affine_qmv(K=256, N=128, group_size=128, bits=4, dtype=torch.float16): return ok +def test_affine_qmm_t_batched(shape=(4, 1, 256), K=256, N=128, group_size=128, bits=4, + dtype=torch.float16): + """Every batch element must match its own reference, not just element 0.""" + print(f"\n--- affine_qmm_t batched (x={shape}, N={N}, gs={group_size}, bits={bits}, {dtype}) ---") + + device = "mps" + x = torch.randn(*shape, dtype=dtype, device=device) + + w_float = torch.randn(N, K, dtype=torch.float32) + w_packed_cpu, scales_cpu, biases_cpu = affine_quantize(w_float, group_size, bits) + + w_packed = w_packed_cpu.to(device) + scales = scales_cpu.to(dtype).to(device) + biases = biases_cpu.to(dtype).to(device) + + y = quantization_mlx.affine_qmm_t(x, w_packed, scales, biases, group_size, bits) + + w_deq = affine_dequantize(w_packed_cpu, scales_cpu.float(), biases_cpu.float(), group_size, bits) + w_deq = w_deq.to(dtype).to(device) + y_ref = (x.reshape(-1, K) @ w_deq.T).reshape(*shape[:-1], N) + + # worst row, not mean: the bug corrupts whole batch elements + diff = (y.float() - y_ref.float()).abs().amax(dim=-1) + denom = y_ref.float().abs().amax(dim=-1) + 1e-6 + worst = (diff / denom).max().item() + + print(f" output shape: {tuple(y.shape)}") + print(f" worst row rel err: {worst:.6f}") + + ok = worst < 0.1 + print(" PASS" if ok else " FAIL") + return ok + + +def test_affine_qmv_batched(shape=(2, 1, 256), N=128, group_size=128, bits=4, + dtype=torch.float16): + """Multi-row / batched qmv: every row and batch element must be computed.""" + K = shape[-1] + print(f"\n--- affine_qmv batched (x={shape}, N={N}, gs={group_size}, bits={bits}, {dtype}) ---") + + device = "mps" + x = torch.randn(*shape, dtype=dtype, device=device) + + w_float = torch.randn(N, K, dtype=torch.float32) + w_packed_cpu, scales_cpu, biases_cpu = affine_quantize(w_float, group_size, bits) + + w_packed = w_packed_cpu.to(device) + scales = scales_cpu.to(dtype).to(device) + biases = biases_cpu.to(dtype).to(device) + + y = quantization_mlx.affine_qmv(x, w_packed, scales, biases, N, group_size, bits) + + w_deq = affine_dequantize(w_packed_cpu, scales_cpu.float(), biases_cpu.float(), group_size, bits) + w_deq = w_deq.to(dtype).to(device) + y_ref = (x.reshape(-1, K) @ w_deq.T).reshape(*shape[:-1], N) + + diff = (y.float() - y_ref.float()).abs().amax(dim=-1) + denom = y_ref.float().abs().amax(dim=-1) + 1e-6 + worst = (diff / denom).max().item() + + print(f" output shape: {tuple(y.shape)}") + print(f" worst row rel err: {worst:.6f}") + + ok = worst < 0.1 + print(" PASS" if ok else " FAIL") + return ok + + # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- @@ -259,6 +327,16 @@ def test_affine_qmv(K=256, N=128, group_size=128, bits=4, dtype=torch.float16): results.append(test_affine_qmv(K=256, N=128, bits=4)) results.append(test_affine_qmv(K=512, N=256, bits=4)) + # batched tests + for bits in (4, 8): + for dt in (torch.float16, torch.bfloat16): + results.append(test_affine_qmm_t_batched(shape=(4, 1, 256), K=256, N=128, bits=bits, dtype=dt)) + results.append(test_affine_qmm_t_batched(shape=(2, 4, 256), K=256, N=128, bits=bits, dtype=dt)) + results.append(test_affine_qmm_t_batched(shape=(2, 2, 2, 256), K=256, N=128, bits=bits, dtype=dt)) + results.append(test_affine_qmv_batched(shape=(4, 256), N=128, bits=bits, dtype=dt)) + results.append(test_affine_qmv_batched(shape=(2, 1, 256), N=128, bits=bits, dtype=dt)) + results.append(test_affine_qmv_batched(shape=(2, 4, 256), N=128, bits=bits, dtype=dt)) + print(f"\n{'='*50}") passed = sum(results) total = len(results)