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..e5c6208a8 --- /dev/null +++ b/src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp @@ -0,0 +1,2181 @@ +// 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 + +#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: 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 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 mlxcel_qmv_wide_flag().load(std::memory_order_relaxed) && + (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()); +} + +// 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