From 6adf9c5ad0114dd6244c90d8fff83fc2ac7cc460 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 17 Aug 2026 18:50:45 +0900 Subject: [PATCH 1/2] perf(speculative): add an MLX overlay switch that restores block-vs-chain exactness MLX sends M >= 2 affine quantized matmuls through qmv_wide on GPU generation 15 and later while M == 1 keeps taking qmv. The two reduce along K in a different order, so a speculative verify block stops being bitwise equal to the single-token chain at block width 2 on M5-class hardware, which closes the #1189 gate and disables MTP there outright. Generations 13 and 14 take qmv on both sides and stay equal up to get_qmv_batch_limit. Upstream exposes no way to select that path and use_qmv_wide is unchanged as of pin 9a795735, so this overlays mlx/backend/metal/quantized.cpp with a single added predicate. MLXCEL_QMV_WIDE=0 forces plain qmv everywhere; unset is upstream verbatim. The delta is one function and one call, so a pin bump refreshes the file and re-applies two hunks. Measured on M5 Max, qwen3.8-27b-4bit with the qwen3.8-27b-mtp-bf16 drafter: - The sweep turns equal exactly where the new batch limits say it should. 5120->5120, 5120->17408 and 3840->15360 hold to T=12 and break at 13 (limit 13); 3840->4096 holds past 24 (limit 25). The dense f16 control still breaks at T=2, so the switch moved only the quantized path. - The #1189 exactness probe passes with no override, three runs out of three. MTP engages on generation 17 in a shipping configuration. - qmv costs 17 to 20 percent on the verify forward and about 10 percent end to end against qmv_wide. - Per-rep ratios over four alternating reps at n=300: MTP with the switch is 1.04x classic decode, MTP without it 1.16x. So the switch is strictly better than what ships today, which is the gate declining to classic decode at 1.00x, and it costs most of what forfeiting the contract would buy. Block width 3 stays the operating point: acceptance falls from 0.65 to 0.17 across widths 3 to 12 while emitted-per-verify only moves 2.3 to 2.8. Not yet wired to anything. Setting it process-wide would also slow batched server decode, whose M lands in the same 2-to-limit window, so scoping it to the verify forward is the next step. Refs #1187, #1186 --- src/lib/mlx-cpp/CMakeLists.txt | 7 +- .../patches/mlx/backend/metal/quantized.cpp | 2161 +++++++++++++++++ 2 files changed, 2167 insertions(+), 1 deletion(-) create mode 100644 src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp diff --git a/src/lib/mlx-cpp/CMakeLists.txt b/src/lib/mlx-cpp/CMakeLists.txt index 62ffe5095..87c26619b 100644 --- a/src/lib/mlx-cpp/CMakeLists.txt +++ b/src/lib/mlx-cpp/CMakeLists.txt @@ -29,9 +29,14 @@ function(mlx_apply_source_overlays mlx_source_dir) # The steel/gemm/mma.h overlay (issue #217: upstream safe-load fix # cherry-pick, PRs #3560/#3565) was retired when the pin moved to # 2026-06-11 upstream main, which carries the fix natively (issue #222). + # quantized.cpp carries one delta, the MLXCEL_QMV_WIDE off-switch on + # `use_qmv_wide` (issue #1187). Upstream has no equivalent knob and the + # predicate is unchanged as of pin 9a795735, so it has to be overlaid here. + # Refresh the file wholesale on a bump and re-apply that one hunk. set(_metal_patch_files "mlx/backend/metal/compiled.cpp" - "mlx/backend/metal/kernels/utils.h") + "mlx/backend/metal/kernels/utils.h" + "mlx/backend/metal/quantized.cpp") foreach(_patch_file ${_metal_patch_files}) if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/patches/${_patch_file}") configure_file( diff --git a/src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp b/src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp new file mode 100644 index 000000000..4a70ee400 --- /dev/null +++ b/src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp @@ -0,0 +1,2161 @@ +// Copyright © 2023-2026 Apple Inc. +// Patched by mlxcel: `use_qmv_wide` gains an off-switch, `MLXCEL_QMV_WIDE=0`. +// Synced to upstream 9a795735; the only delta is `qmv_wide_enabled()` and the +// one call it adds to `use_qmv_wide`. Everything else is upstream verbatim, so +// a bump refreshes this file and re-applies those two hunks. +// +// Why it exists (lablup/mlxcel#1186, #1187). `use_qmv_wide` sends `M >= 2` +// affine quantized matmuls down `qmv_wide` on GPU generation 15 and later +// while `M == 1` keeps taking `qmv`. Those two kernels reduce along K in a +// different order, so a speculative verify block and the single-token decode +// chain it must match stop being bitwise equal at block width 2 on M5-class +// hardware, which is what closes the #1189 exactness gate and disables MTP +// there entirely. Generations 13 and 14 take `qmv` on both sides and stay +// equal up to `get_qmv_batch_limit`, so forcing that path is a way to buy the +// contract back. Upstream exposes no knob for this and `use_qmv_wide` is +// unchanged as of 9a795735, so the switch has to live here. +// +// Off-switch only, matching MLXCEL_METAL4_ATTENTION: unset behaves exactly as +// upstream, and the value is read once per process because this sits on the +// dispatch path of every quantized matmul. + +#include +#include + +#include "mlx/backend/common/quantized.h" +#include "mlx/backend/common/broadcasting.h" +#include "mlx/backend/common/compiled.h" +#include "mlx/backend/gpu/copy.h" +#include "mlx/backend/metal/device.h" +#include "mlx/backend/metal/kernels.h" +#include "mlx/backend/metal/reduce.h" +#include "mlx/backend/metal/unary.h" +#include "mlx/backend/metal/utils.h" +#include "mlx/fast_primitives.h" +#include "mlx/primitives.h" +#include "mlx/utils.h" + +namespace mlx::core { + +namespace { + +template +auto get_quantized_kernel_wrapped( + metal::Device& d, + const std::string& name, + const std::string& func, + const std::string& mode, + const std::string& type, + int group_size, + int bits, + Args... args) { + std::string template_def; + std::string fname = ((mode == "affine") ? "affine_" : "fp_") + func; + template_def = get_template_definition( + name, fname, type, group_size, bits, std::forward(args)...); + return get_quantized_kernel(d, name, template_def, mode); +} + +template +auto get_qmm_nax_kernel_wrapped( + metal::Device& d, + const std::string& name, + const std::string& func, + const std::string& mode, + const std::string& type, + int group_size, + int bits, + Args... args) { + std::string template_def; + std::string fname = ((mode == "affine") ? "affine_" : "fp_") + func; + template_def = get_template_definition( + name, fname, type, group_size, bits, std::forward(args)...); + return get_qmm_nax_kernel(d, name, template_def, mode); +} + +inline array +ensure_row_contiguous(const array& x, metal::Device& d, const Stream& s) { + if (!x.flags().row_contiguous) { + array x_copy = contiguous_copy_gpu(x, s); + metal::get_command_encoder(s).add_temporary(x_copy); + return x_copy; + } else { + return x; + } +} + +inline array ensure_row_contiguous_matrix( + const array& x, + metal::Device& d, + const Stream& s) { + if (x.ndim() < 2) { + if (x.strides()[0] == 1) { + return x; + } + } else { + auto stride_0 = x.strides()[x.ndim() - 2]; + auto stride_1 = x.strides()[x.ndim() - 1]; + if (stride_0 == x.shape(-1) && stride_1 == 1) { + return x; + } + } + array x_copy = contiguous_copy_gpu(x, s); + metal::get_command_encoder(s).add_temporary(x_copy); + return x_copy; +} + +inline int get_qmv_batch_limit(int D, int O, metal::Device& d) { + auto arch_size = d.get_architecture().back(); + auto arch_gen = d.get_architecture_gen(); + if (arch_gen >= 17 && arch_size != 'd') { + if (D <= 2048 && O <= 2048) { + return 33; + } else if (D <= 4096 && O <= 4096) { + return 25; + } else { + return 13; + } + } else if (arch_gen >= 15 && arch_size != 'd') { + if (D <= 2048 && O <= 2048) { + return 13; + } else if (D <= 4096 && O <= 4096) { + return 15; + } else { + return 13; + } + } else if (arch_gen >= 13) { + switch (arch_size) { + case 'd': + if (D <= 2048 && O <= 2048) { + return 32; + } else if (D <= 4096 && O <= 4096) { + return 18; + } else { + return 12; + } + default: + if (D <= 2048 && O <= 2048) { + return 14; + } else if (D <= 4096 && O <= 4096) { + return 10; + } else { + return 6; + } + } + } else { + switch (arch_size) { + case 'd': + if (D <= 2048 && O <= 2048) { + return 32; + } else if (D <= 4096 && O <= 4096) { + return 18; + } else { + return 12; + } + default: + if (D <= 2048 && O <= 2048) { + return 18; + } else if (D <= 4096 && O <= 4096) { + return 12; + } else { + return 10; + } + } + } +} + +// Must match the K step in qmv_fast_impl (kernels/quantized.h): +// pack_factor() * (bits == 2 ? 1 : 2) * SIMD_SIZE +inline int qmv_fast_k_alignment(int bits) { + return get_pack_factor(bits, 32) * (bits == 2 ? 1 : 2) * 32; +} + +inline int add_strides_and_shapes( + CommandEncoder& compute_encoder, + bool skip, + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + int offset) { + if (skip) { + return offset; + } + + // TODO: Collapse batch dimensions + + int x_batch_ndims = x.ndim() - 2; + int w_batch_ndims = w.ndim() - 2; + compute_encoder.set_bytes(x_batch_ndims, offset++); + compute_encoder.set_vector_bytes(x.shape(), offset++); + compute_encoder.set_vector_bytes(x.strides(), offset++); + compute_encoder.set_bytes(w_batch_ndims, offset++); + compute_encoder.set_vector_bytes(w.shape(), offset++); + compute_encoder.set_vector_bytes(w.strides(), offset++); + compute_encoder.set_vector_bytes(scales.strides(), offset++); + if (biases) { + compute_encoder.set_vector_bytes(biases->strides(), offset++); + } + + return offset; +} + +inline int add_gather_strides_and_shapes( + CommandEncoder& compute_encoder, + const array& lhs_indices, + const array& rhs_indices, + int offset) { + auto [shape, strides] = collapse_contiguous_dims( + lhs_indices.shape(), {lhs_indices.strides(), rhs_indices.strides()}); + int ndims = shape.size(); + + compute_encoder.set_bytes(ndims, offset++); + compute_encoder.set_vector_bytes(shape, offset++); + compute_encoder.set_vector_bytes(strides[0], offset++); + compute_encoder.set_vector_bytes(strides[1], offset++); + + return offset; +} + +auto get_quantize_kernel_dims( + MTL::ComputePipelineState* kernel, + const array& w, + const array& out, + int group_size, + int bits, + bool dequantize = false) { + // Treat uint32 as uint8 in kernel + constexpr int uint8_per_uint32 = 4; + constexpr int simd_size = 32; + int packs_per_int = (bits == 3 || bits == 5) ? 8 : bits == 6 ? 4 : 8 / bits; + int per_thread = + dequantize ? packs_per_int : std::max(group_size / simd_size, 1); + size_t nthreads = + dequantize ? out.size() / packs_per_int : w.size() / per_thread; + + NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup(); + if (thread_group_size > nthreads) { + thread_group_size = nthreads; + } + auto group_dims = MTL::Size(thread_group_size, 1, 1); + bool use_2d = nthreads > UINT_MAX; + auto grid_shape = w.shape(); + if (dequantize) { + grid_shape.back() *= uint8_per_uint32; + } else { + grid_shape.back() /= per_thread; + } + MTL::Size grid_dims = use_2d ? get_2d_grid_dims(grid_shape, w.strides()) + : MTL::Size(nthreads, 1, 1); + return std::make_tuple(grid_dims, group_dims); +} + +void quantize_impl( + const std::vector& inputs, + std::vector& outputs, + QuantizationMode mode, + int group_size, + int bits, + bool dequantize, + Stream s) { + auto& w_pre = inputs[0]; + auto& out = outputs[0]; + out.set_data(allocator::malloc(out.nbytes())); + + auto& d = metal::device(s.device); + auto& compute_encoder = metal::get_command_encoder(s); + + bool has_biases = (mode == QuantizationMode::Affine); + bool has_global_scale = !has_biases && (inputs.size() > (1 + dequantize)); + + auto w = ensure_row_contiguous(w_pre, d, s); + if (dequantize) { + auto scales = ensure_row_contiguous(inputs[1], d, s); + compute_encoder.set_input_array(w, 0); + compute_encoder.set_input_array(scales, 1); + if (has_biases) { + auto biases = ensure_row_contiguous(inputs[2], d, s); + compute_encoder.set_input_array(biases, 2); + } else if (has_global_scale) { + compute_encoder.set_input_array(inputs[2], 2); + } + compute_encoder.set_output_array(out, 3); + } else { + auto& scales = outputs[1]; + scales.set_data(allocator::malloc(scales.nbytes())); + compute_encoder.set_input_array(w, 0); + compute_encoder.set_output_array(out, 1); + compute_encoder.set_output_array(scales, 2); + if (has_biases) { + auto& biases = outputs[2]; + biases.set_data(allocator::malloc(biases.nbytes())); + compute_encoder.set_output_array(biases, 3); + } else if (has_global_scale) { + compute_encoder.set_input_array(inputs[1], 3); + } + } + + auto type_string = dequantize ? get_type_string(out.dtype()) + : get_type_string(w_pre.dtype()); + auto mode_string = quantization_mode_to_string(mode); + std::string kname; + concatenate( + kname, + mode_string + (dequantize ? "_dequantize" : "_quantize"), + "_", + type_string, + "_gs_", + group_size, + "_b_", + bits); + if (!has_biases) { + concatenate(kname, "_hgs_", has_global_scale ? "true" : "false"); + } + auto kernel = get_quantized_kernel_wrapped( + d, + kname, + dequantize ? "dequantize" : "quantize", + mode_string, + type_string, + group_size, + bits, + has_global_scale); + + auto [grid_dims, group_dims] = + get_quantize_kernel_dims(kernel, w, out, group_size, bits, dequantize); + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.dispatch_threads(grid_dims, group_dims); +} + +auto quantize_input( + const array& w, + const std::optional& global_scale, + QuantizationMode mode, + int group_size, + int bits, + metal::Device& d, + Stream s) { + auto wq_shape = w.shape(); + wq_shape.back() = w.shape(-1) * bits / 32; + auto scales_shape = w.shape(); + scales_shape.back() = w.shape(-1) / group_size; + + std::vector inputs{w}; + if (global_scale) { + inputs.push_back(*global_scale); + } + std::vector outputs{ + array(wq_shape, uint32, nullptr, {}), + array(scales_shape, uint8, nullptr, {})}; + auto& compute_encoder = metal::get_command_encoder(s); + compute_encoder.add_temporary(outputs[0]); + compute_encoder.add_temporary(outputs[1]); + quantize_impl(inputs, outputs, mode, group_size, bits, false, s); + return std::make_tuple(outputs[0], outputs[1]); +} + +void fp_quantize_dequantize( + const array& in, + const std::optional& global_scale, + array& out, + const std::string& mode, + int group_size, + int bits, + metal::Device& d, + const Stream& s) { + auto& compute_encoder = metal::get_command_encoder(s); + + auto w = ensure_row_contiguous(in, d, s); + compute_encoder.set_input_array(w, 0); + if (global_scale) { + compute_encoder.set_input_array(*global_scale, 1); + } + compute_encoder.set_output_array(out, 2); + auto type_string = get_type_string(in.dtype()); + std::string kname; + concatenate( + kname, + mode + "_quantize_dequantize_", + type_string, + "_gs_", + group_size, + "_b_", + bits, + "_hgs_", + global_scale ? "true" : "false"); + auto kernel = get_quantized_kernel_wrapped( + d, + kname, + "quantize_dequantize", + mode, + type_string, + group_size, + bits, + global_scale.has_value()); + + auto [grid_dims, group_dims] = + get_quantize_kernel_dims(kernel, w, out, group_size, bits); + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.dispatch_threads(grid_dims, group_dims); +} + +array quantize_dequantize_input( + const array& x_pre, + const std::optional& global_scale, + const std::string& mode, + int group_size, + int bits, + metal::Device& d, + Stream s) { + bool donate_x = x_pre.is_donatable(); + array x = ensure_row_contiguous(x_pre, d, s); + // If x is a copy it should be donatable + donate_x |= x.is_donatable(); + auto xhat = + donate_x ? x : array(allocator::malloc(x.nbytes()), x.shape(), x.dtype()); + if (!donate_x) { + metal::get_command_encoder(s).add_temporary(xhat); + } + fp_quantize_dequantize(x, global_scale, xhat, mode, group_size, bits, d, s); + return xhat; +} + +} // namespace + +void qmv_quad( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + array& out, + int group_size, + int bits, + int M, + int N, + int K, + metal::Device& d, + const Stream& s, + const std::string& mode) { + int B = out.size() / M / N; + + constexpr int quads_per_simd = 8; + constexpr int results_per_quadgroup = 8; + int bn = quads_per_simd * results_per_quadgroup; + int simdgroup_size = 32; + MTL::Size group_dims(simdgroup_size, 1, 1); + MTL::Size grid_dims(M, (N + bn - 1) / bn, B); + + std::string kname; + kname.reserve(64); + std::string type_string = get_type_string(x.dtype()); + + concatenate( + kname, + mode + "_qmv_quad_", + type_string, + "_gs_", + group_size, + "_b_", + bits, + "_d_", + K, + B > 1 ? "_batch_1" : "_batch_0"); + auto kernel = get_quantized_kernel_wrapped( + d, kname, "qmv_quad", mode, type_string, group_size, bits, K, B > 1); + auto& compute_encoder = metal::get_command_encoder(s); + compute_encoder.set_compute_pipeline_state(kernel); + + int c = 0; + compute_encoder.set_input_array(w, c++); + compute_encoder.set_input_array(scales, c++); + if (biases) { + compute_encoder.set_input_array(*biases, c++); + } + compute_encoder.set_input_array(x, c++); + compute_encoder.set_output_array(out, c++); + compute_encoder.set_bytes(K, c++); + compute_encoder.set_bytes(N, c++); + add_strides_and_shapes(compute_encoder, B <= 1, x, w, scales, biases, c++); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); +} + +void qmv( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const std::optional& global_scale, + array& out, + int group_size, + int bits, + int M, + int N, + int K, + metal::Device& d, + const Stream& s, + const std::string& mode) { + int B = out.size() / M / N; + + int bn = 8; + int bk = 32; + + std::string kname; + kname.reserve(64); + std::string type_string = get_type_string(x.dtype()); + bool fast = N % bn == 0 && K % qmv_fast_k_alignment(bits) == 0; + // A narrower output tile reduces register pressure for large + // floating-point quantized matrix-vector products on M5 Max GPUs. + bool use_narrow_qmv = fast && N >= 4096 && d.get_architecture_gen() == 17 && + d.get_architecture().back() == 's' && mode == "nvfp4"; + int results_per_simdgroup = use_narrow_qmv ? 2 : 4; + bn = 2 * results_per_simdgroup; + MTL::Size group_dims(bk, 2, 1); + MTL::Size grid_dims(M, (N + bn - 1) / bn, B); + + concatenate( + kname, + mode + (fast ? "_qmv_fast_" : "_qmv_"), + type_string, + "_gs_", + group_size, + "_b_", + bits, + use_narrow_qmv ? "_r_2" : "", + B > 1 ? "_batch_1" : "_batch_0", + global_scale ? "_hgs" : ""); + auto kernel = get_quantized_kernel_wrapped( + d, + kname, + (fast ? "qmv_fast" : "qmv"), + mode, + type_string, + group_size, + bits, + B > 1, + global_scale.has_value(), + results_per_simdgroup); + + auto& compute_encoder = metal::get_command_encoder(s); + compute_encoder.set_compute_pipeline_state(kernel); + + compute_encoder.set_input_array(w, 0); + compute_encoder.set_input_array(scales, 1); + if (biases) { + compute_encoder.set_input_array(*biases, 2); + } else if (global_scale) { + compute_encoder.set_input_array(*global_scale, 2); + } + int c = 3; + compute_encoder.set_input_array(x, c++); + compute_encoder.set_output_array(out, c++); + compute_encoder.set_bytes(K, c++); + compute_encoder.set_bytes(N, c++); + add_strides_and_shapes(compute_encoder, B <= 1, x, w, scales, biases, c); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); +} + +// mlxcel: MLXCEL_QMV_WIDE=0 (or false/no/off) forces every mode and every +// generation onto plain qmv. Evaluated once; a static local initializer is +// thread-safe since C++11 and this is read from the dispatch path. +inline bool qmv_wide_enabled() { + static const bool enabled = [] { + const char* raw = std::getenv("MLXCEL_QMV_WIDE"); + if (raw == nullptr) { + return true; + } + return !(std::strcmp(raw, "0") == 0 || std::strcmp(raw, "false") == 0 || + std::strcmp(raw, "no") == 0 || std::strcmp(raw, "off") == 0); + }(); + return enabled; +} + +// affine qmv_wide only beats qmv on gen-15+; fp benefits on every gen. +inline bool use_qmv_wide(const std::string& mode, metal::Device& d) { + return qmv_wide_enabled() && + (mode != "affine" || d.get_architecture_gen() >= 15); +} + +// Dispatches qmv_wide (fp modes -> fp_qmv_wide, affine -> affine_qmv_wide): +// vecs_per_tg input vectors streamed and reused per weight group. +void qmv_wide( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + array& out, + int group_size, + int bits, + int M, + int N, + int K, + metal::Device& d, + const Stream& s, + const std::string& mode) { + // vecs_per_tg is the per-threadgroup input-vector tile. Each tile re-reads + // the weights, so use the fewest tiles, then the smallest tile that fills + // them. + int n_tiles = (M + 4) / 5; // ceil(M / 5); tile size caps at 5 + int vecs_per_tg = (M + n_tiles - 1) / n_tiles; + + // k_lanes: lanes reducing K per output row (32/k_lanes rows per simdgroup). + // The affine subchunk decode has enough ALU per weight load to favor more + // rows per simdgroup (kl8); the fp modes' vectorized dot is balanced at 16. + int k_lanes = mode == "affine" ? 8 : 16; + constexpr int num_simdgroups = 2; + int B = out.size() / M / N; + bool batched = B > 1; + // Output rows per threadgroup: (32 / k_lanes) per simdgroup x num_simdgroups. + int rows_per_tg = (32 / k_lanes) * num_simdgroups; + + MTL::Size group_dims(32, num_simdgroups, 1); + MTL::Size grid_dims( + (M + vecs_per_tg - 1) / vecs_per_tg, + (N + rows_per_tg - 1) / rows_per_tg, + B); + + std::string kname; + kname.reserve(64); + std::string type_string = get_type_string(x.dtype()); + concatenate( + kname, + mode + "_qmv_wide_", + type_string, + "_gs_", + group_size, + "_b_", + bits, + "_nv_", + vecs_per_tg, + "_kl_", + k_lanes, + batched ? "_batch_1" : "_batch_0"); + auto kernel = get_quantized_kernel_wrapped( + d, + kname, + "qmv_wide", + mode, + type_string, + group_size, + bits, + vecs_per_tg, + k_lanes, + batched); + + auto& compute_encoder = metal::get_command_encoder(s); + compute_encoder.set_compute_pipeline_state(kernel); + + int c = 0; + compute_encoder.set_input_array(w, c++); + compute_encoder.set_input_array(scales, c++); + if (biases) { + compute_encoder.set_input_array(*biases, c++); + } + compute_encoder.set_input_array(x, c++); + compute_encoder.set_output_array(out, c++); + compute_encoder.set_bytes(K, c++); + compute_encoder.set_bytes(N, c++); + compute_encoder.set_bytes(M, c++); + add_strides_and_shapes(compute_encoder, !batched, x, w, scales, biases, c); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); +} + +void qvm_split_k( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + array& out, + int group_size, + int bits, + int M, + int N, + int K, + metal::Device& d, + const Stream& s, + const std::string& mode) { + auto& compute_encoder = metal::get_command_encoder(s); + + int split_k = K > 8192 ? 32 : 8; + int split_D = (K + split_k - 1) / split_k; + int B = out.size() / M / N; + B *= split_k; + + constexpr int num_simdgroups = 2; + constexpr int bk = 32; + int bn = std::min(group_size, 32) * num_simdgroups; + MTL::Size group_dims = MTL::Size(bk, num_simdgroups, 1); + MTL::Size grid_dims = MTL::Size(M, (N + bn - 1) / bn, B); + + auto x_shape = x.shape(); + auto x_strides = x.strides(); + if (x_shape.size() == 1) { + x_shape.insert(x_shape.begin(), 1); + x_strides.insert(x_strides.begin(), 0); + } + + int x_ndim = x_shape.size(); + int x_batch_ndims = x_ndim - 2; + int w_batch_ndims = w.ndim() - 2; + auto w_shape = w.shape(); + auto w_strides = w.strides(); + auto s_strides = scales.strides(); + + // Add split_k dim with reshapes + x_shape.insert(x_shape.end() - 2, split_k); + x_shape.back() /= split_k; + x_strides.insert(x_strides.end() - 2, split_D); + x_strides[x_ndim - 1] = split_D; + x_batch_ndims += 1; + + w_shape.insert(w_shape.end() - 2, split_k); + w_shape[w.ndim() - 1] /= split_k; + w_strides.insert(w_strides.end() - 2, split_D * w.shape(-1)); + w_batch_ndims += 1; + s_strides.insert(s_strides.end() - 2, split_D * scales.shape(-1)); + + int final_block_size = K - (split_k - 1) * split_D; + + auto temp_shape = out.shape(); + if (temp_shape.size() == 1) { + temp_shape.insert(temp_shape.begin(), 1); + } + temp_shape.insert(temp_shape.end() - 2, split_k); + array intermediate(temp_shape, x.dtype(), nullptr, {}); + intermediate.set_data(allocator::malloc(intermediate.nbytes())); + compute_encoder.add_temporary(intermediate); + + std::string type_string = get_type_string(x.dtype()); + std::string kname; + kname.reserve(64); + concatenate( + kname, + mode + "_qvm_split_k_", + type_string, + "_gs_", + group_size, + "_b_", + bits, + "_spk_", + split_k); + + // Encode and dispatch kernel + auto kernel = get_quantized_kernel_wrapped( + d, kname, "qvm_split_k", mode, type_string, group_size, bits, split_k); + + compute_encoder.set_compute_pipeline_state(kernel); + + int c = 0; + compute_encoder.set_input_array(w, c++); + compute_encoder.set_input_array(scales, c++); + if (biases) { + compute_encoder.set_input_array(*biases, c++); + } + compute_encoder.set_input_array(x, c++); + compute_encoder.set_output_array(intermediate, c++); + compute_encoder.set_bytes(split_D, c++); + compute_encoder.set_bytes(N, c++); + + compute_encoder.set_bytes(x_batch_ndims, c++); + compute_encoder.set_vector_bytes(x_shape, c++); + compute_encoder.set_vector_bytes(x_strides, c++); + compute_encoder.set_bytes(w_batch_ndims, c++); + compute_encoder.set_vector_bytes(w_shape, c++); + compute_encoder.set_vector_bytes(w_strides, c++); + compute_encoder.set_vector_bytes(s_strides, c++); + if (biases) { + auto b_strides = biases->strides(); + b_strides.insert(b_strides.end() - 2, split_D * biases->shape(-1)); + compute_encoder.set_vector_bytes(b_strides, c++); + } + compute_encoder.set_bytes(final_block_size, c++); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); + + int axis = intermediate.ndim() - 3; + ReductionPlan plan( + ReductionOpType::ContiguousStridedReduce, + {intermediate.shape(axis)}, + {intermediate.strides(axis)}); + strided_reduce_general_dispatch( + intermediate, out, "sum", plan, {axis}, compute_encoder, d, s); +} + +void qvm( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const std::optional& global_scale, + array& out, + int group_size, + int bits, + int M, + int N, + int K, + metal::Device& d, + const Stream& s, + const std::string& mode) { + int B = out.size() / M / N; + + constexpr int num_simdgroups = 2; + constexpr int bk = 32; + int bn = std::min(group_size, 32) * num_simdgroups; + MTL::Size group_dims(bk, num_simdgroups, 1); + MTL::Size grid_dims(M, (N + bn - 1) / bn, B); + + std::string kname; + kname.reserve(64); + std::string type_string = get_type_string(x.dtype()); + concatenate( + kname, + mode + "_qvm_", + type_string, + "_gs_", + group_size, + "_b_", + bits, + B > 1 ? "_batch_1" : "_batch_0", + global_scale ? "_hgs" : ""); + auto kernel = get_quantized_kernel_wrapped( + d, + kname, + "qvm", + mode, + type_string, + group_size, + bits, + B > 1, + global_scale.has_value()); + auto& compute_encoder = metal::get_command_encoder(s); + compute_encoder.set_compute_pipeline_state(kernel); + + compute_encoder.set_input_array(w, 0); + compute_encoder.set_input_array(scales, 1); + if (biases) { + compute_encoder.set_input_array(*biases, 2); + } else if (global_scale) { + compute_encoder.set_input_array(*global_scale, 2); + } + int c = 3; + compute_encoder.set_input_array(x, c++); + compute_encoder.set_output_array(out, c++); + compute_encoder.set_bytes(K, c++); + compute_encoder.set_bytes(N, c++); + add_strides_and_shapes(compute_encoder, B <= 1, x, w, scales, biases, c++); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); +} + +void qmm_nax( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + array& out, + bool transpose, + int group_size, + int bits, + int M, + int N, + int K, + metal::Device& d, + const Stream& s, + const std::string& mode) { + int B = out.size() / M / N; + + int wm = 2; + int wn = 2; + int bm = 64; + int bn = 64; + int bk = 64; + MTL::Size group_dims(32, wn, wm); + MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, B); + + std::string kname; + kname.reserve(64); + bool aligned = N % 64 == 0; + bool batched = B > 1; + std::string type_string = get_type_string(x.dtype()); + concatenate( + kname, + mode + (transpose ? "_qmm_t_nax_" : "_qmm_n_nax_"), + type_string, + "_gs_", + group_size, + "_b_", + bits, + "_bm", + bm, + "_bn", + bn, + "_bk", + bk, + "_wm", + wm, + "_wn", + wn, + transpose ? (aligned ? "_alN_true" : "_alN_false") : "", + batched ? "_batch_1" : "_batch_0"); + std::string template_def; + MTL::ComputePipelineState* kernel; + if (transpose) { + kernel = get_qmm_nax_kernel_wrapped( + d, + kname, + "qmm_t_nax", + mode, + type_string, + group_size, + bits, + aligned, + batched, + bm, + bk, + bn, + wm, + wn); + } else { + kernel = get_qmm_nax_kernel_wrapped( + d, + kname, + "qmm_n_nax", + mode, + type_string, + group_size, + bits, + batched, + bm, + bk, + bn, + wm, + wn); + } + auto& compute_encoder = metal::get_command_encoder(s); + compute_encoder.set_compute_pipeline_state(kernel); + + int c = 0; + compute_encoder.set_input_array(w, c++); + compute_encoder.set_input_array(scales, c++); + if (biases) { + compute_encoder.set_input_array(*biases, c++); + } + compute_encoder.set_input_array(x, c++); + compute_encoder.set_output_array(out, c++); + compute_encoder.set_bytes(K, c++); + compute_encoder.set_bytes(N, c++); + compute_encoder.set_bytes(M, c++); + add_strides_and_shapes(compute_encoder, B <= 1, x, w, scales, biases, c); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); +} + +void gather_qmm_nax( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const array& lhs_indices, + const array& rhs_indices, + array& out, + bool transpose, + int group_size, + int bits, + int M, + int N, + int K, + metal::Device& d, + const Stream& s, + const std::string& mode) { + int B = out.size() / M / N; + + int wm = 2; + int wn = 2; + int bm = 64; + int bn = 64; + // The gather qmm NAX kernels are instantiated with BK = 64 only; any + // other value here makes the kernel name lookup fail. + int bk = 64; + MTL::Size group_dims(32, wn, wm); + MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, B); + + std::string kname; + kname.reserve(64); + bool aligned = N % 64 == 0; + std::string type_string = get_type_string(x.dtype()); + concatenate( + kname, + mode + (transpose ? "_gather_qmm_t_nax_" : "_gather_qmm_n_nax_"), + type_string, + "_gs_", + group_size, + "_b_", + bits, + "_bm", + bm, + "_bn", + bn, + "_bk", + bk, + "_wm", + wm, + "_wn", + wn, + transpose ? (aligned ? "_alN_true" : "_alN_false") : ""); + MTL::ComputePipelineState* kernel; + if (transpose) { + kernel = get_qmm_nax_kernel_wrapped( + d, + kname, + "gather_qmm_t_nax_", + mode, + type_string, + group_size, + bits, + aligned, + bm, + bk, + bn, + wm, + wn); + } else { + kernel = get_qmm_nax_kernel_wrapped( + d, + kname, + "gather_qmm_n_nax_", + mode, + type_string, + group_size, + bits, + bm, + bk, + bn, + wm, + wn); + } + + auto& compute_encoder = metal::get_command_encoder(s); + compute_encoder.set_compute_pipeline_state(kernel); + + int c = 0; + compute_encoder.set_input_array(w, c++); + compute_encoder.set_input_array(scales, c++); + if (biases) { + compute_encoder.set_input_array(*biases, c++); + } + compute_encoder.set_input_array(x, c++); + compute_encoder.set_input_array(lhs_indices, c++); + compute_encoder.set_input_array(rhs_indices, c++); + compute_encoder.set_output_array(out, c++); + compute_encoder.set_bytes(K, c++); + compute_encoder.set_bytes(N, c++); + compute_encoder.set_bytes(M, c++); + c = add_strides_and_shapes(compute_encoder, false, x, w, scales, biases, c); + add_gather_strides_and_shapes(compute_encoder, lhs_indices, rhs_indices, c); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); +} + +void qmm( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + array& out, + bool transpose, + int group_size, + int bits, + int M, + int N, + int K, + metal::Device& d, + const Stream& s, + const std::string& mode) { + bool has_nax_kernel = + metal::is_nax_available() && (transpose || mode == "affine"); + if (has_nax_kernel && transpose && (K % 64 == 0) && + (env::enable_tf32() || x.dtype() != float32)) { + return qmm_nax( + /* const array& x = */ x, + /* const array& w = */ w, + /* const array& scales = */ scales, + /* const std::optional& biases = */ biases, + /* array& out = */ out, + /* bool transpose = */ transpose, + /* int group_size = */ group_size, + /* int bits = */ bits, + /* int M = */ M, + /* int N = */ N, + /* int K = */ K, + /* metal::Device& d = */ d, + /* const Stream& s = */ s, + /* const std::string& mode = */ mode); + } + + int B = out.size() / M / N; + + int wm = 2; + int wn = 2; + int bm = 32; + int bn = 32; + MTL::Size group_dims(32, wn, wm); + MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, B); + + std::string kname; + kname.reserve(64); + bool aligned = N % 32 == 0; + bool batched = B > 1; + std::string type_string = get_type_string(x.dtype()); + concatenate( + kname, + mode + (transpose ? "_qmm_t_" : "_qmm_n_"), + type_string, + "_gs_", + group_size, + "_b_", + bits, + transpose ? (aligned ? "_alN_true" : "_alN_false") : "", + batched ? "_batch_1" : "_batch_0"); + std::string template_def; + MTL::ComputePipelineState* kernel; + if (transpose) { + kernel = get_quantized_kernel_wrapped( + d, + kname, + "qmm_t", + mode, + type_string, + group_size, + bits, + aligned, + batched); + } else { + kernel = get_quantized_kernel_wrapped( + d, kname, "qmm_n", mode, type_string, group_size, bits, batched); + } + auto& compute_encoder = metal::get_command_encoder(s); + compute_encoder.set_compute_pipeline_state(kernel); + + int c = 0; + compute_encoder.set_input_array(w, c++); + compute_encoder.set_input_array(scales, c++); + if (biases) { + compute_encoder.set_input_array(*biases, c++); + } + compute_encoder.set_input_array(x, c++); + compute_encoder.set_output_array(out, c++); + compute_encoder.set_bytes(K, c++); + compute_encoder.set_bytes(N, c++); + compute_encoder.set_bytes(M, c++); + add_strides_and_shapes(compute_encoder, B <= 1, x, w, scales, biases, c); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); +} + +void qmm_splitk( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + array& out, + int group_size, + int bits, + int M, + int N, + int K, + metal::Device& d, + const Stream& s, + const std::string& mode) { + // Choose split_k to target ~512 threadgroups + int bm = 32, bn = 32; + int n_tiles = (N + bn - 1) / bn; + int m_tiles = (M + bm - 1) / bm; + int current_tgs = n_tiles * m_tiles; + int split_k = std::max(1, 512 / current_tgs); + + // Each K partition must be a whole number of BK-wide (32) K-tiles as well as + // whole quantization groups. The qmm_t_splitk kernels tile K by BK=32 and do + // not bound the K dimension, so a partition smaller than BK (e.g. nvfp4's + // group_size=16) would over-read into the next group's weights/scales. + int k_align = group_size > 32 ? group_size : 32; + split_k = std::min(split_k, K / k_align); + + // Ensure K divides evenly by split_k * k_align + while (split_k > 1 && (K % (split_k * k_align) != 0)) { + split_k--; + } + if (split_k <= 1) { + return qmm( + x, w, scales, biases, out, true, group_size, bits, M, N, K, d, s, mode); + } + + int k_partition_size = K / split_k; + int split_k_partition_stride = M * N; + + // Allocate intermediate buffer: insert split_k at the front so that + // partition_stride = M * N matches the leading stride of the buffer. + auto& compute_encoder = metal::get_command_encoder(s); + auto temp_shape = out.shape(); + if (temp_shape.size() == 1) { + temp_shape.insert(temp_shape.begin(), 1); + } + temp_shape.insert(temp_shape.begin(), split_k); + array intermediate(temp_shape, x.dtype(), nullptr, {}); + intermediate.set_data(allocator::malloc(intermediate.nbytes())); + compute_encoder.add_temporary(intermediate); + + // Grid: (N_tiles, M_tiles, split_k) + MTL::Size group_dims(32, 2, 2); + MTL::Size grid_dims(n_tiles, m_tiles, split_k); + + bool aligned = N % 32 == 0; + std::string type_string = get_type_string(x.dtype()); + std::string kname; + kname.reserve(64); + concatenate( + kname, + mode + "_qmm_t_splitk_", + type_string, + "_gs_", + group_size, + "_b_", + bits, + aligned ? "_alN_true" : "_alN_false"); + auto kernel = get_quantized_kernel_wrapped( + d, kname, "qmm_t_splitk", mode, type_string, group_size, bits, aligned); + + compute_encoder.set_compute_pipeline_state(kernel); + + int c = 0; + compute_encoder.set_input_array(w, c++); + compute_encoder.set_input_array(scales, c++); + if (biases) { + compute_encoder.set_input_array(*biases, c++); + } + compute_encoder.set_input_array(x, c++); + compute_encoder.set_output_array(intermediate, c++); + compute_encoder.set_bytes(K, c++); + compute_encoder.set_bytes(N, c++); + compute_encoder.set_bytes(M, c++); + compute_encoder.set_bytes(k_partition_size, c++); + compute_encoder.set_bytes(split_k_partition_stride, c++); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); + + // Sum across split_k dimension (axis 0) + ReductionPlan plan( + ReductionOpType::ContiguousStridedReduce, + {intermediate.shape(0)}, + {intermediate.strides(0)}); + strided_reduce_general_dispatch( + intermediate, out, "sum", plan, {0}, compute_encoder, d, s); +} + +void gather_qmm( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const array& lhs_indices, + const array& rhs_indices, + array& out, + bool transpose, + int group_size, + int bits, + int M, + int N, + int K, + metal::Device& d, + const Stream& s, + const std::string& mode) { + if (metal::is_nax_available() && transpose && (K % 64 == 0) && + (env::enable_tf32() || x.dtype() != float32)) { + return gather_qmm_nax( + /* const array& x = */ x, + /* const array& w = */ w, + /* const array& scales = */ scales, + /* const std::optional& biases = */ biases, + /* const array& lhs_indices = */ lhs_indices, + /* const array& rhs_indices = */ rhs_indices, + /* array& out = */ out, + /* bool transpose = */ transpose, + /* int group_size = */ group_size, + /* int bits = */ bits, + /* int M = */ M, + /* int N = */ N, + /* int K = */ K, + /* metal::Device& d = */ d, + /* const Stream& s = */ s, + /* const std::string& mode = */ mode); + } + + int B = out.size() / M / N; + + int wm = 2; + int wn = 2; + int bm = 32; + int bn = 32; + MTL::Size group_dims(32, wn, wm); + MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, B); + + std::string kname; + kname.reserve(64); + bool aligned = N % 32 == 0; + std::string type_string = get_type_string(x.dtype()); + concatenate( + kname, + mode + (transpose ? "_gather_qmm_t_" : "_gather_qmm_n_"), + type_string, + "_gs_", + group_size, + "_b_", + bits, + transpose ? (aligned ? "_alN_true" : "_alN_false") : ""); + MTL::ComputePipelineState* kernel; + if (transpose) { + kernel = get_quantized_kernel_wrapped( + d, kname, "gather_qmm_t", mode, type_string, group_size, bits, aligned); + } else { + kernel = get_quantized_kernel_wrapped( + d, kname, "gather_qmm_n", mode, type_string, group_size, bits); + } + + auto& compute_encoder = metal::get_command_encoder(s); + compute_encoder.set_compute_pipeline_state(kernel); + + int c = 0; + compute_encoder.set_input_array(w, c++); + compute_encoder.set_input_array(scales, c++); + if (biases) { + compute_encoder.set_input_array(*biases, c++); + } + compute_encoder.set_input_array(x, c++); + compute_encoder.set_input_array(lhs_indices, c++); + compute_encoder.set_input_array(rhs_indices, c++); + compute_encoder.set_output_array(out, c++); + compute_encoder.set_bytes(K, c++); + compute_encoder.set_bytes(N, c++); + compute_encoder.set_bytes(M, c++); + c = add_strides_and_shapes(compute_encoder, false, x, w, scales, biases, c); + add_gather_strides_and_shapes(compute_encoder, lhs_indices, rhs_indices, c); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); +} + +void gather_qmv( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const std::optional& global_scale, + const array& lhs_indices, + const array& rhs_indices, + array& out, + int group_size, + int bits, + int M, + int N, + int K, + metal::Device& d, + const Stream& s, + const std::string& mode) { + int B = out.size() / M / N; + + int bn = 8; + int bk = 32; + MTL::Size group_dims(bk, 2, 1); + MTL::Size grid_dims(M, (N + bn - 1) / bn, B); + + std::string kname; + kname.reserve(64); + std::string type_string = get_type_string(x.dtype()); + bool fast = N % bn == 0 && K % qmv_fast_k_alignment(bits) == 0; + concatenate( + kname, + mode + (fast ? "_gather_qmv_fast_" : "_gather_qmv_"), + type_string, + "_gs_", + group_size, + "_b_", + bits, + global_scale ? "_hgs" : ""); + + auto kernel = get_quantized_kernel_wrapped( + d, + kname, + (fast ? "gather_qmv_fast" : "gather_qmv"), + mode, + type_string, + group_size, + bits, + global_scale.has_value()); + + auto& compute_encoder = metal::get_command_encoder(s); + compute_encoder.set_compute_pipeline_state(kernel); + + compute_encoder.set_input_array(w, 0); + compute_encoder.set_input_array(scales, 1); + if (biases) { + compute_encoder.set_input_array(*biases, 2); + } else if (global_scale) { + compute_encoder.set_input_array(*global_scale, 2); + } + int c = 3; + compute_encoder.set_input_array(x, c++); + compute_encoder.set_input_array(lhs_indices, c++); + compute_encoder.set_input_array(rhs_indices, c++); + compute_encoder.set_output_array(out, c++); + compute_encoder.set_bytes(K, c++); + compute_encoder.set_bytes(N, c++); + c = add_strides_and_shapes(compute_encoder, false, x, w, scales, biases, c); + add_gather_strides_and_shapes(compute_encoder, lhs_indices, rhs_indices, c); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); +} + +void gather_qvm( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const std::optional& global_scale, + const array& lhs_indices, + const array& rhs_indices, + array& out, + int group_size, + int bits, + int M, + int N, + int K, + metal::Device& d, + const Stream& s, + const std::string& mode) { + int B = out.size() / M / N; + + constexpr int num_simdgroups = 2; + constexpr int bk = 32; + int bn = std::min(group_size, 32) * num_simdgroups; + MTL::Size group_dims(bk, num_simdgroups, 1); + MTL::Size grid_dims(M, (N + bn - 1) / bn, B); + + std::string kname; + kname.reserve(64); + std::string type_string = get_type_string(x.dtype()); + concatenate( + kname, + mode + "_gather_qvm_", + type_string, + "_gs_", + group_size, + "_b_", + bits, + global_scale ? "_hgs" : ""); + auto kernel = get_quantized_kernel_wrapped( + d, + kname, + "gather_qvm", + mode, + type_string, + group_size, + bits, + global_scale.has_value()); + auto& compute_encoder = metal::get_command_encoder(s); + compute_encoder.set_compute_pipeline_state(kernel); + + compute_encoder.set_input_array(w, 0); + compute_encoder.set_input_array(scales, 1); + if (biases) { + compute_encoder.set_input_array(*biases, 2); + } else if (global_scale) { + compute_encoder.set_input_array(*global_scale, 2); + } + int c = 3; + compute_encoder.set_input_array(x, c++); + compute_encoder.set_input_array(lhs_indices, c++); + compute_encoder.set_input_array(rhs_indices, c++); + compute_encoder.set_output_array(out, c++); + compute_encoder.set_bytes(K, c++); + compute_encoder.set_bytes(N, c++); + c = add_strides_and_shapes(compute_encoder, false, x, w, scales, biases, c++); + add_gather_strides_and_shapes(compute_encoder, lhs_indices, rhs_indices, c); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); +} + +void gather_qmm_rhs_nax( + const array& x_, + const array& w_, + const array& scales_, + const std::optional& biases_, + const array& indices_, + array& out, + bool transpose, + int group_size, + int bits, + int M, + int N, + int K, + metal::Device& d, + const Stream& s, + const std::string mode) { + // Start by normalizing the indices + array indices = ensure_row_contiguous(indices_, d, s); + + // Broadcast x with indices. If we are here that means lhs_indices were not + // provided so the lhs_indices are implied to be the shape of x broadcasted + // with rhs_indices. We need only broadcast x and copy it as if applying the + // lhs_indices. + auto broadcast_with_indices = [&d, &s, &indices](const array& x) { + if (x.size() / x.shape(-2) / x.shape(-1) == indices.size()) { + return ensure_row_contiguous(x, d, s); + } + + auto x_shape = indices.shape(); + x_shape.push_back(x.shape(-2)); + x_shape.push_back(x.shape(-1)); + array new_x(std::move(x_shape), x.dtype(), nullptr, {}); + broadcast(x, new_x); + return ensure_row_contiguous(new_x, d, s); + }; + + // Normalize the input arrays + array x = broadcast_with_indices(x_); + array w = ensure_row_contiguous(w_, d, s); + array scales = ensure_row_contiguous(scales_, d, s); + std::optional biases; + if (biases_) { + biases = ensure_row_contiguous(*biases_, d, s); + } + + // Use smaller bm for many experts and few tokens. + int E = w.size() / w.shape(-1) / w.shape(-2); + int bm = (M / E < 64) ? 32 : 64; + int bn = 64, bk = 64; + int wm = 2, wn = 2; + + const bool align_M = (M % bm) == 0; + const bool align_N = (N % bn) == 0; + const bool align_K = (K % bk) == 0; + + // Make the kernel name + std::string kname; + kname.reserve(64); + std::string type_string = get_type_string(x.dtype()); + concatenate( + kname, + mode + + (transpose ? "_gather_qmm_rhs_nax_nt_" : "_gather_qmm_rhs_nax_nn_"), + type_string, + "_gs_", + group_size, + "_b_", + bits, + "_bm_", + bm, + "_bn_", + bn, + "_bk_", + bk, + "_wm_", + wm, + "_wn_", + wn); + + metal::MTLFCList func_consts = { + {&align_M, MTL::DataType::DataTypeBool, 200}, + {&align_N, MTL::DataType::DataTypeBool, 201}, + {&align_K, MTL::DataType::DataTypeBool, 202}, + }; + + // And the kernel hash that includes the function constants + std::string hash_name; + hash_name.reserve(128); + concatenate( + hash_name, + kname, + "_align_M_", + align_M ? 't' : 'n', + "_align_N_", + align_N ? 't' : 'n', + "_align_K_", + align_K ? 't' : 'n'); + + // Get and set the kernel + auto& compute_encoder = metal::get_command_encoder(s); + auto kernel = get_gather_qmm_nax_kernel( + d, + kname, + hash_name, + func_consts, + x, + group_size, + bits, + mode, + bm, + bn, + bk, + wm, + wn, + transpose); + compute_encoder.set_compute_pipeline_state(kernel); + + MTL::Size group_dims(32, wn, wm); + MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, 1); + + int c = 0; + compute_encoder.set_input_array(x, c++); + compute_encoder.set_input_array(w, c++); + compute_encoder.set_input_array(scales, c++); + if (biases) { + compute_encoder.set_input_array(*biases, c++); + } + compute_encoder.set_input_array(indices, c++); + compute_encoder.set_output_array(out, c++); + compute_encoder.set_bytes(M, c++); + compute_encoder.set_bytes(N, c++); + compute_encoder.set_bytes(K, c++); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); +} + +void gather_qmm_rhs( + const array& x_, + const array& w_, + const array& scales_, + const std::optional& biases_, + const array& indices_, + array& out, + bool transpose, + int group_size, + int bits, + int M, + int N, + int K, + metal::Device& d, + const Stream& s, + const std::string mode) { + if (metal::is_nax_available() && transpose && + (env::enable_tf32() || x_.dtype() != float32)) { + return gather_qmm_rhs_nax( + /* const array& x_ = */ x_, + /* const array& w_ = */ w_, + /* const array& scales_ = */ scales_, + /* const std::optional& biases_ = */ biases_, + /* const array& indices_ = */ indices_, + /* array& out = */ out, + /* bool transpose = */ transpose, + /* int group_size = */ group_size, + /* int bits = */ bits, + /* int M = */ M, + /* int N = */ N, + /* int K = */ K, + /* metal::Device& d = */ d, + /* const Stream& s = */ s, + /* const std::string mode = */ mode); + } + + // Start by normalizing the indices + array indices = ensure_row_contiguous(indices_, d, s); + + // Broadcast x with indices. If we are here that means lhs_indices were not + // provided so the lhs_indices are implied to be the shape of x broadcasted + // with rhs_indices. We need only broadcast x and copy it as if applying the + // lhs_indices. + auto broadcast_with_indices = [&d, &s, &indices](const array& x) { + if (x.size() / x.shape(-2) / x.shape(-1) == indices.size()) { + return ensure_row_contiguous(x, d, s); + } + + auto x_shape = indices.shape(); + x_shape.push_back(x.shape(-2)); + x_shape.push_back(x.shape(-1)); + array new_x(std::move(x_shape), x.dtype(), nullptr, {}); + broadcast(x, new_x); + return ensure_row_contiguous(new_x, d, s); + }; + + // Normalize the input arrays + array x = broadcast_with_indices(x_); + array w = ensure_row_contiguous(w_, d, s); + array scales = ensure_row_contiguous(scales_, d, s); + std::optional biases; + if (biases_) { + biases = ensure_row_contiguous(*biases_, d, s); + } + + // TODO: Tune the block sizes + int bm = 16, bn = 32, bk = 32; + int wm = 1, wn = 2; + + const bool align_M = (M % bm) == 0; + const bool align_N = (N % bn) == 0; + const bool align_K = (K % bk) == 0; + + // Make the kernel name + std::string kname; + kname.reserve(64); + std::string type_string = get_type_string(x.dtype()); + concatenate( + kname, + mode + (transpose ? "_gather_qmm_rhs_nt_" : "_gather_qmm_rhs_nn_"), + type_string, + "_gs_", + group_size, + "_b_", + bits, + "_bm_", + bm, + "_bn_", + bn, + "_bk_", + bk, + "_wm_", + wm, + "_wn_", + wn); + + metal::MTLFCList func_consts = { + {&align_M, MTL::DataType::DataTypeBool, 200}, + {&align_N, MTL::DataType::DataTypeBool, 201}, + {&align_K, MTL::DataType::DataTypeBool, 202}, + }; + + // And the kernel hash that includes the function constants + std::string hash_name; + hash_name.reserve(128); + concatenate( + hash_name, + kname, + "_align_M_", + align_M ? 't' : 'n', + "_align_N_", + align_N ? 't' : 'n', + "_align_K_", + align_K ? 't' : 'n'); + + // Get and set the kernel + auto& compute_encoder = metal::get_command_encoder(s); + auto kernel = get_gather_qmm_kernel( + d, + kname, + hash_name, + func_consts, + x, + group_size, + bits, + mode, + bm, + bn, + bk, + wm, + wn, + transpose); + compute_encoder.set_compute_pipeline_state(kernel); + + MTL::Size group_dims(32, wn, wm); + MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, 1); + + int c = 0; + compute_encoder.set_input_array(x, c++); + compute_encoder.set_input_array(w, c++); + compute_encoder.set_input_array(scales, c++); + if (biases) { + compute_encoder.set_input_array(*biases, c++); + } + compute_encoder.set_input_array(indices, c++); + compute_encoder.set_output_array(out, c++); + compute_encoder.set_bytes(M, c++); + compute_encoder.set_bytes(N, c++); + compute_encoder.set_bytes(K, c++); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); +} + +void dispatch_qmv( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const std::optional& global_scale, + array& out, + int group_size, + int bits, + int M, + int N, + int K, + metal::Device& d, + const Stream& s, + const std::string& mode) { + // It is a qmv with a small inner dimension so route to qmv_quad kernel + if ((K == 128 || K == 64) && is_power_of_2(bits) && !global_scale) { + qmv_quad(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode); + return; + } + + // Small batch so route to qmv_wide, which reuses each weight group across the + // M vectors. + if (M >= 2 && use_qmv_wide(mode, d) && !global_scale) { + qmv_wide(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode); + return; + } + qmv(x, + w, + scales, + biases, + global_scale, + out, + group_size, + bits, + M, + N, + K, + d, + s, + mode); +} + +void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { + auto& s = stream(); + auto& d = metal::device(s.device); + + out.set_data(allocator::malloc(out.nbytes())); + + // Make sure the last two dims of x and w, s, b are contiguous. This should + // be relaxed for x. + array x = ensure_row_contiguous_matrix(inputs[0], d, s); + array w = ensure_row_contiguous_matrix(inputs[1], d, s); + array scales = ensure_row_contiguous_matrix(inputs[2], d, s); + std::optional biases = std::nullopt; + if (inputs.size() == 4) { + biases = ensure_row_contiguous_matrix(inputs[3], d, s); + } + + // Extract the matmul shapes + bool non_batched = w.ndim() == 2 && x.flags().row_contiguous; + int K = x.shape(-1); + int M = non_batched ? x.size() / K : x.shape(-2); + int N = out.shape(-1); + + int vector_limit = transpose_ ? get_qmv_batch_limit(K, N, d) : 4; + auto mode = quantization_mode_to_string(mode_); + // It is a matrix matrix product. + if (M >= vector_limit) { + // Use split-K qmm for small M with transposed weights (non-batched only) + int B = out.size() / M / N; + if (transpose_ && B == 1) { + qmm_splitk( + x, w, scales, biases, out, group_size_, bits_, M, N, K, d, s, mode); + return; + } + qmm(x, + w, + scales, + biases, + out, + transpose_, + group_size_, + bits_, + M, + N, + K, + d, + s, + mode); + return; + } + + // Run of the mill qmv + if (transpose_) { + dispatch_qmv( + x, + w, + scales, + biases, + std::nullopt, + out, + group_size_, + bits_, + M, + N, + K, + d, + s, + mode); + return; + } + + // Run of the mill qvm + if (K < 1024) { + qvm(x, + w, + scales, + biases, + std::nullopt, + out, + group_size_, + bits_, + M, + N, + K, + d, + s, + mode); + return; + } + + // Qvm with large dimension so route to a split K kernel for more parallelism + qvm_split_k( + x, w, scales, biases, out, group_size_, bits_, M, N, K, d, s, mode); + return; +} + +void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { + auto& s = stream(); + auto& d = metal::device(s.device); + + out.set_data(allocator::malloc(out.nbytes())); + + array x = ensure_row_contiguous_matrix(inputs[0], d, s); + array w = ensure_row_contiguous_matrix(inputs[1], d, s); + array scales = ensure_row_contiguous_matrix(inputs[2], d, s); + std::optional biases = std::nullopt; + if (inputs.size() == 6) { + biases = ensure_row_contiguous_matrix(inputs[3], d, s); + } + const array& lhs_indices = inputs[inputs.size() - 2]; + const array& rhs_indices = inputs[inputs.size() - 1]; + + int K = x.shape(-1); + int M = x.shape(-2); + int N = out.shape(-1); + int B = out.size() / M / N; + int E = w.size() / w.shape(-1) / w.shape(-2); + int vector_limit = transpose_ ? get_qmv_batch_limit(K, N, d) : 4; + auto mode = quantization_mode_to_string(mode_); + + // We are walking x in order and w is also in order so we can batch up the + // matmuls and reuse reading x and w. + // + // TODO: Tune 16 and 4 here a bit better. + if (M == 1 && B >= 16 && right_sorted_ == true && B / E >= 4) { + gather_qmm_rhs( + x, + w, + scales, + biases, + rhs_indices, + out, + transpose_, + group_size_, + bits_, + x.size() / K, + N, + K, + d, + s, + mode); + return; + } + + // It is a matrix matrix product + if (M >= vector_limit) { + gather_qmm( + x, + w, + scales, + biases, + lhs_indices, + rhs_indices, + out, + transpose_, + group_size_, + bits_, + M, + N, + K, + d, + s, + mode); + return; + } + + if (transpose_) { + gather_qmv( + x, + w, + scales, + biases, + std::nullopt, + lhs_indices, + rhs_indices, + out, + group_size_, + bits_, + M, + N, + K, + d, + s, + mode); + return; + } + + gather_qvm( + x, + w, + scales, + biases, + std::nullopt, + lhs_indices, + rhs_indices, + out, + group_size_, + bits_, + M, + N, + K, + d, + s, + mode); +} + +void QQMatmul::eval_gpu(const std::vector& inputs, array& out) { + auto& s = stream(); + auto& d = metal::device(s.device); + + const array& x_pre = inputs[0]; + const array& w_pre = inputs[1]; + auto mode = quantization_mode_to_string(mode_); + + out.set_data(allocator::malloc(out.nbytes())); + + // - 2 inputs: x, w (non-quantized w) + // - 3 inputs: x, w, scales_w (quantized w) + bool w_quantized = (inputs[1].dtype() == uint32); + int base_size = w_quantized ? 3 : 2; + // For nvfp4, global scales are optional but must be both present or both + // absent If present, they add 2 more inputs (global_scale_x, global_scale_w) + bool has_global_scales = + mode_ == QuantizationMode::Nvfp4 && inputs.size() == base_size + 2; + assert(inputs.size() == base_size || has_global_scales); + + std::optional global_scale_x; + std::optional global_scale_w; + if (has_global_scales) { + global_scale_x = inputs[inputs.size() - 2]; + global_scale_w = inputs[inputs.size() - 1]; + } + + // Quantize weights. + auto [w_q, scales_w] = !w_quantized + ? quantize_input(w_pre, global_scale_w, mode_, group_size_, bits_, d, s) + : std::make_tuple( + ensure_row_contiguous_matrix(w_pre, d, s), + ensure_row_contiguous_matrix(inputs[base_size - 1], d, s)); + + // Quantize activation. + array x = quantize_dequantize_input( + x_pre, global_scale_x, mode, group_size_, bits_, d, s); + + bool non_batched = w_q.ndim() == 2; + int K = x.shape(-1); + int M = non_batched ? x.size() / K : x.shape(-2); + int N = out.shape(-1); + dispatch_qmv( + x, + w_q, + scales_w, + std::nullopt, + global_scale_w, + out, + group_size_, + bits_, + M, + N, + K, + d, + s, + mode); +} + +void GatherQQMM::eval_gpu(const std::vector& inputs, array& out) { + auto& s = stream(); + auto& d = metal::device(s.device); + + const array& x_pre = inputs[0]; + const array& w_pre = inputs[1]; + const array& lhs_indices = ensure_row_contiguous(inputs[2], d, s); + const array& rhs_indices = ensure_row_contiguous(inputs[3], d, s); + auto mode = quantization_mode_to_string(mode_); + + out.set_data(allocator::malloc(out.nbytes())); + + // - 4 inputs: x, w (non-quantized w) + // - 5 inputs: x, w, scales_w (quantized w) + bool w_quantized = (inputs[1].dtype() == uint32); + int base_size = w_quantized ? 5 : 4; + // For nvfp4, global scales are optional but must be both present or both + // absent If present, they add 2 more inputs (global_scale_x, global_scale_w) + bool has_global_scales = + mode_ == QuantizationMode::Nvfp4 && inputs.size() == base_size + 2; + assert(inputs.size() == base_size || has_global_scales); + + std::optional global_scale_x; + std::optional global_scale_w; + if (has_global_scales) { + global_scale_x = inputs[inputs.size() - 2]; + global_scale_w = inputs[inputs.size() - 1]; + } + + // Quantize weights. + auto [w_q, scales_w] = !w_quantized + ? quantize_input(w_pre, global_scale_w, mode_, group_size_, bits_, d, s) + : std::make_tuple( + ensure_row_contiguous_matrix(w_pre, d, s), + ensure_row_contiguous_matrix(inputs[base_size - 1], d, s)); + + // Quantize activation. + array x = quantize_dequantize_input( + x_pre, global_scale_x, mode, group_size_, bits_, d, s); + + bool non_batched = w_q.ndim() == 2; + int K = x.shape(-1); + int M = non_batched ? x.size() / K : x.shape(-2); + int N = out.shape(-1); + gather_qmv( + x, + w_q, + scales_w, + std::nullopt, + global_scale_w, + lhs_indices, + rhs_indices, + out, + group_size_, + bits_, + M, + N, + K, + d, + s, + mode); +} + +void fast::Quantize::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + quantize_impl( + inputs, outputs, mode_, group_size_, bits_, dequantize_, stream()); +} + +void fast::ConvertFP8::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + auto& in = inputs[0]; + auto& out = outputs[0]; + unary_op_gpu(inputs, out, name(), stream()); +} + +} // namespace mlx::core From e161fdf33ef399cf37002f1e48e1296f349733a5 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 17 Aug 2026 19:31:44 +0900 Subject: [PATCH 2/2] feat(speculative): let the MTP gate buy back exactness by dropping qmv_wide The gate probed once and declined, which on Apple GPU generation 15 and later meant declining always: use_qmv_wide sends M >= 2 affine quantized matmuls to a kernel that reduces along K differently from the M == 1 qmv the decode chain takes, so the verify block is never byte-equal there. MTP was off on exactly the hardware it pays on. It now re-probes with qmv_wide disabled and keeps it off when that is what makes the block exact. Generations 13 and 14 already take qmv on both sides, so this puts generation 15+ on the path that was known to hold, and the new get_qmv_batch_limit values (13 for large operands on generation 17) leave the shipping block widths of 3 and 4 well inside it. Measured on M5 Max, qwen3.8-27b-4bit with the qwen3.8-27b-mtp-bf16 drafter, no environment variables set: the probe fails under qmv_wide (165506 of 496640 logit bytes), passes without it, and MTP engages with the contract intact at 37.74 tok/s against classic decode's 31 to 34. Three paths, all verified live: - nothing set: retry, engage, log what was traded - MLXCEL_QMV_WIDE=1: retry skipped, decline, and the message says why - MLXCEL_MTP_ALLOW_INEXACT=1: engage on the fast kernel, loud as before The switch is left off for the rest of the process rather than restored. Restoring it would break the block this gate just approved, and the exact arm stays correct for every other caller; it costs about 17 to 20 percent on the verify forward. That also means a server sharing the process pays it on batched decode, whose M lands in the same window, so scoping it to the verify forward is filed as follow-up rather than done here. The switch is indirected behind a thread-local under cfg(test). Toggling the real one from a unit test perturbs the numeric tests libtest runs beside it, in the direction of making them pass, which cost the suite its determinism while this was being built. Refs #1187 --- .../patches/mlx/backend/metal/quantized.cpp | 36 ++++-- src/lib/mlxcel-core/build.rs | 10 ++ src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp | 29 +++++ src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h | 11 ++ src/lib/mlxcel-core/src/lib.rs | 11 ++ src/models/speculative_exactness.rs | 118 +++++++++++++++++- src/models/speculative_exactness_tests.rs | 65 ++++++++++ 7 files changed, 267 insertions(+), 13 deletions(-) diff --git a/src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp b/src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp index 4a70ee400..e5c6208a8 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp +++ b/src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp @@ -19,6 +19,7 @@ // upstream, and the value is read once per process because this sits on the // dispatch path of every quantized matmul. +#include #include #include @@ -556,24 +557,31 @@ void qmv( compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } -// mlxcel: MLXCEL_QMV_WIDE=0 (or false/no/off) forces every mode and every -// generation onto plain qmv. Evaluated once; a static local initializer is -// thread-safe since C++11 and this is read from the dispatch path. -inline bool qmv_wide_enabled() { - static const bool enabled = [] { +// mlxcel: the qmv_wide off-switch. Seeded from MLXCEL_QMV_WIDE on first +// touch, then settable at runtime through mlxcel_set_qmv_wide so the MTP +// gate can decide by measurement instead of asking the operator. +// +// Relaxed ordering is deliberate. The flag selects between two kernels that +// compute the same function to different last-ulp results, so a stale read +// costs precision, never correctness, and this sits on the dispatch path of +// every quantized matmul. Callers must set it before the work whose kernel +// choice they mean to pin, and MLX evaluates lazily, so "before the work" +// means before the eval that forces it, not before the graph is built. +std::atomic& mlxcel_qmv_wide_flag() { + static std::atomic flag{[] { const char* raw = std::getenv("MLXCEL_QMV_WIDE"); if (raw == nullptr) { return true; } return !(std::strcmp(raw, "0") == 0 || std::strcmp(raw, "false") == 0 || std::strcmp(raw, "no") == 0 || std::strcmp(raw, "off") == 0); - }(); - return enabled; + }()}; + return flag; } // affine qmv_wide only beats qmv on gen-15+; fp benefits on every gen. inline bool use_qmv_wide(const std::string& mode, metal::Device& d) { - return qmv_wide_enabled() && + return mlxcel_qmv_wide_flag().load(std::memory_order_relaxed) && (mode != "affine" || d.get_architecture_gen() >= 15); } @@ -2158,4 +2166,16 @@ void fast::ConvertFP8::eval_gpu( unary_op_gpu(inputs, out, name(), stream()); } +// mlxcel: runtime entry point for the qmv_wide off-switch, called through +// the cxx bridge (mlxcel_core::set_qmv_wide). Declared here rather than in a +// header because upstream owns every header in this tree and an overlay that +// adds one would drift on the next pin bump. +void mlxcel_set_qmv_wide(bool enabled) { + mlxcel_qmv_wide_flag().store(enabled, std::memory_order_relaxed); +} + +bool mlxcel_qmv_wide(void) { + return mlxcel_qmv_wide_flag().load(std::memory_order_relaxed); +} + } // namespace mlx::core diff --git a/src/lib/mlxcel-core/build.rs b/src/lib/mlxcel-core/build.rs index 93587d925..af6949310 100644 --- a/src/lib/mlxcel-core/build.rs +++ b/src/lib/mlxcel-core/build.rs @@ -114,6 +114,16 @@ fn main() { // included by the generated cxx bridge, so suppress it for all profiles. .flag_if_supported("-Wno-deprecated-copy"); + // The qmv_wide off-switch (issue #1187) lives in the + // mlx/backend/metal/quantized.cpp overlay, so its symbol exists only when + // the Metal backend is actually compiled. `__APPLE__` is not that + // condition: a macOS build without the `metal` feature sets + // MLX_BUILD_METAL=OFF and would link against a symbol that was never + // emitted. Gate on the feature that decides whether the file is built. + if std::env::var("CARGO_FEATURE_METAL").is_ok() { + bridge.define("MLXCEL_BRIDGE_METAL_BACKEND", None); + } + // Add optimization flags for release builds #[cfg(not(debug_assertions))] { diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp index d00b06109..7db71059f 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp @@ -17,6 +17,19 @@ #include #include +#ifdef MLXCEL_BRIDGE_METAL_BACKEND +// Defined by the mlx/backend/metal/quantized.cpp overlay (issue #1187). +// Declared here rather than pulled from a header because upstream owns every +// header in that tree and an overlay that added one would drift on the next +// pin bump. It must sit at global scope: declaring `namespace mlx::core` +// inside `namespace mlx_cxx` creates `mlx_cxx::mlx::core` and shadows the +// real one for the rest of the file. +namespace mlx::core { +void mlxcel_set_qmv_wide(bool enabled); +bool mlxcel_qmv_wide(void); +} // namespace mlx::core +#endif + namespace mlx_cxx { using namespace mlx::core; @@ -929,6 +942,22 @@ void random_seed(uint64_t seed) { mlx::core::random::seed(seed); } +#ifdef MLXCEL_BRIDGE_METAL_BACKEND +void set_qmv_wide(bool enabled) { + ::mlx::core::mlxcel_set_qmv_wide(enabled); +} + +bool qmv_wide_enabled() { + return ::mlx::core::mlxcel_qmv_wide(); +} +#else +void set_qmv_wide(bool) {} + +bool qmv_wide_enabled() { + return true; +} +#endif + std::unique_ptr random_categorical(const MlxArray& logits, int32_t axis) { return std::make_unique(mlx::core::random::categorical(logits.inner, axis)); } diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h index a092b4924..eccd647f6 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h @@ -378,6 +378,17 @@ std::unique_ptr equal(const MlxArray& a, const MlxArray& b); // Seed the global MLX random number generator void random_seed(uint64_t seed); +// The qmv_wide off-switch carried by the mlx/backend/metal/quantized.cpp +// overlay (issue #1187). Disabling it forces every quantized matmul with +// M >= 2 onto plain qmv, the kernel M == 1 already takes, which is what makes +// a speculative verify block bitwise equal to the single-token chain on GPU +// generation 15 and later. Costs about 17 to 20 percent on that forward. +// +// No-ops and reports true on a build without the Metal backend, where the +// overlay is not compiled and the distinction does not exist. +void set_qmv_wide(bool enabled); +bool qmv_wide_enabled(); + // Random categorical sampling std::unique_ptr random_categorical(const MlxArray& logits, int32_t axis); diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index 7272b7f58..4877ef2a6 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/src/lib/mlxcel-core/src/lib.rs @@ -535,6 +535,17 @@ mod ffi { /// Seed the global MLX random number generator fn random_seed(seed: u64); + /// Force every quantized matmul with `M >= 2` onto plain `qmv` + /// instead of `qmv_wide`, which is what restores block-vs-chain + /// bitwise equality on Apple GPU generation 15 and later (#1187). + /// Process-wide and immediate. A build without the Metal backend + /// ignores it. + fn set_qmv_wide(enabled: bool); + + /// Whether `qmv_wide` is currently enabled. `true` on a build + /// without the Metal backend. + fn qmv_wide_enabled() -> bool; + /// Random categorical sampling fn random_categorical(logits: &MlxArray, axis: i32) -> UniquePtr; diff --git a/src/models/speculative_exactness.rs b/src/models/speculative_exactness.rs index 09b72c013..e8aea4ff4 100644 --- a/src/models/speculative_exactness.rs +++ b/src/models/speculative_exactness.rs @@ -177,9 +177,101 @@ fn verdict_cache() -> &'static Mutex> { /// process and never on the request path after that. The decision, not the raw verdict, is what gets cached: with /// `MLXCEL_MTP_ALLOW_INEXACT` set a diverging probe still returns `true`, /// and the log line says so. -pub fn mtp_exactness_gate(key: ProbeKey, probe: F) -> bool +/// Whether the operator pinned `MLXCEL_QMV_WIDE` themselves. +/// +/// An explicit setting wins over the gate's own retry: someone who asked for +/// a kernel selection gets it, and finds out from the decline message that +/// the contract was unreachable under it, rather than silently getting the +/// other one. +fn qmv_wide_pinned_by_operator() -> bool { + static PINNED: OnceLock = OnceLock::new(); + *PINNED.get_or_init(|| std::env::var("MLXCEL_QMV_WIDE").is_ok()) +} + +/// The `qmv_wide` switch, indirected so the gate's control flow is testable +/// without touching a process-wide kernel selection. +/// +/// Under `cfg(test)` this is a thread-local, so a test that drives the retry +/// cannot perturb the numeric tests libtest runs beside it. That matters here: +/// turning the switch off makes concurrent block-vs-chain comparisons *more* +/// exact, so a leaked toggle would flip unrelated known-failing tests to +/// passing at random and cost the suite its determinism. The live path is +/// covered by the CLI runs recorded on #1187, not by these unit tests. +#[cfg(not(test))] +fn qmv_wide_switch_get() -> bool { + mlxcel_core::qmv_wide_enabled() +} + +#[cfg(not(test))] +fn qmv_wide_switch_set(enabled: bool) { + mlxcel_core::set_qmv_wide(enabled); +} + +#[cfg(test)] +thread_local! { + static TEST_QMV_WIDE: std::cell::Cell = const { std::cell::Cell::new(true) }; +} + +#[cfg(test)] +fn qmv_wide_switch_get() -> bool { + TEST_QMV_WIDE.with(|c| c.get()) +} + +#[cfg(test)] +fn qmv_wide_switch_set(enabled: bool) { + TEST_QMV_WIDE.with(|c| c.set(enabled)); +} + +/// Re-probe with `qmv_wide` disabled, keeping it off when that is what makes +/// the block byte-identical. +/// +/// `qmv_wide` is MLX's faster kernel for `M >= 2` quantized matmuls on Apple +/// GPU generation 15 and later, and it reduces along K in a different order +/// than the `qmv` that `M == 1` takes. That difference is the whole reason a +/// verify block stops matching the single-token chain on this hardware, so +/// turning it off is the one lever that buys the contract back without +/// giving up speculative decoding (#1187). +/// +/// Returns `true` when the retry was exact, in which case the switch is left +/// off for the rest of the process. It is deliberately not restored: it is a +/// per-process kernel selection, the exact arm stays correct for every other +/// caller, and re-enabling it would break the very block this gate just +/// approved. The cost is real and measured, about 17 to 20 percent on the +/// verify forward, so the log line says what was traded for what. +fn retry_without_qmv_wide( + key: ProbeKey, + probe: &mut F, + first: &BlockChainExactness, +) -> Option +where + F: FnMut() -> BlockChainExactness, +{ + if qmv_wide_pinned_by_operator() || !qmv_wide_switch_get() { + return None; + } + + qmv_wide_switch_set(false); + let retry = probe(); + if retry.is_equal() { + tracing::info!( + block_size = key.block_size, + "MTP exactness probe failed under qmv_wide ({}) and passed without it. \ + Disabling qmv_wide for this process to keep the temperature-0 \ + byte-identity contract; the verify forward costs about 17 to 20 \ + percent more. Set MLXCEL_QMV_WIDE=1 to pin the faster kernel and \ + decline MTP instead.", + first.reason() + ); + Some(true) + } else { + qmv_wide_switch_set(true); + Some(false) + } +} + +pub fn mtp_exactness_gate(key: ProbeKey, mut probe: F) -> bool where - F: FnOnce() -> BlockChainExactness, + F: FnMut() -> BlockChainExactness, { if let Ok(cache) = verdict_cache().lock() && let Some(decision) = cache.get(&key) @@ -188,7 +280,13 @@ where } let verdict = probe(); - let decision = verdict.is_equal() || allow_inexact(); + let retried = if verdict.is_equal() { + None + } else { + retry_without_qmv_wide(key, &mut probe, &verdict) + }; + let exact = verdict.is_equal() || retried == Some(true); + let decision = exact || allow_inexact(); if verdict.is_equal() { tracing::info!( @@ -196,6 +294,8 @@ where "MTP exactness probe passed: {}", verdict.reason() ); + } else if exact { + // The retry logged what it traded; nothing to add here. } else if decision { tracing::warn!( block_size = key.block_size, @@ -205,12 +305,20 @@ where verdict.reason() ); } else { + let also_tried = match retried { + Some(false) => " Disabling qmv_wide did not make it exact either.", + None if qmv_wide_pinned_by_operator() => { + " The qmv_wide retry was skipped because MLXCEL_QMV_WIDE is pinned." + } + _ => "", + }; tracing::warn!( block_size = key.block_size, - "MTP declined: {}. Falling back to classic decode. Set \ + "MTP declined: {}.{} Falling back to classic decode. Set \ MLXCEL_MTP_ALLOW_INEXACT=1 to engage anyway and forfeit the \ temperature-0 byte-identity contract.", - verdict.reason() + verdict.reason(), + also_tried ); } diff --git a/src/models/speculative_exactness_tests.rs b/src/models/speculative_exactness_tests.rs index 8d28b2c6d..0c3d431ee 100644 --- a/src/models/speculative_exactness_tests.rs +++ b/src/models/speculative_exactness_tests.rs @@ -117,6 +117,71 @@ fn a_diverging_probe_declines_unless_the_override_is_set() { ); } +/// A probe that diverges under `qmv_wide` and agrees without it must engage +/// MTP rather than decline, because that is the whole point of the retry +/// (#1187): on Apple GPU generation 15 and later the `M >= 2` kernel split is +/// the only thing breaking the contract, and turning it off restores it. +/// +/// The control flow is what this pins, not the kernel selection. On a build +/// without the Metal backend the switch is inert, so the second call sees the +/// same hardware as the first; the stateful closure stands in for the change +/// the switch makes on hardware that has it. Both builds must reach the probe +/// exactly twice and engage. +#[test] +fn a_probe_that_only_diverges_under_qmv_wide_engages_after_the_retry() { + let k = key(9005); + let mut calls = 0; + let decision = mtp_exactness_gate(k, || { + calls += 1; + if calls == 1 { + BlockChainExactness::Diverges { + position: 0, + differing_bytes: 165_506, + total_bytes: 496_640, + } + } else { + BlockChainExactness::Equal + } + }); + assert!( + decision, + "an exact retry without qmv_wide must engage MTP, not decline" + ); + assert_eq!(calls, 2, "the gate must re-probe exactly once"); +} + +/// The retry must not fire when the operator pinned the kernel themselves. +/// +/// Only assertable when `MLXCEL_QMV_WIDE` is actually set in the ambient +/// environment, because the pin is read once per process like the other +/// switches; otherwise this asserts the ordinary diverging-probe contract, +/// which is the same thing the retry-less build does. +#[test] +fn a_pinned_qmv_wide_skips_the_retry() { + let pinned = std::env::var("MLXCEL_QMV_WIDE").is_ok(); + let k = key(9006); + let mut calls = 0; + let decision = mtp_exactness_gate(k, || { + calls += 1; + if calls == 1 { + BlockChainExactness::Diverges { + position: 0, + differing_bytes: 4, + total_bytes: 1024, + } + } else { + BlockChainExactness::Equal + } + }); + if pinned { + assert_eq!(calls, 1, "a pinned MLXCEL_QMV_WIDE must skip the re-probe"); + assert_eq!(decision, super::allow_inexact()); + } else { + assert_eq!(calls, 2, "an unpinned build must re-probe"); + assert!(decision); + } +} + #[test] fn different_block_widths_are_measured_separately() { // The whole point of probing: byte-identity holds at one block width