From 0b51f651269cf2d9e7fc3d18d7e56b549726beec Mon Sep 17 00:00:00 2001 From: Goni Zahavy Date: Wed, 15 Jul 2026 20:52:41 +0300 Subject: [PATCH 1/9] [Vulkan] Native BF16 fused affine-4 QMM for dense decode/prefill Dense 4-bit QuantizedMatmul was forced through float32 staging and the generic byte matvec. Add multi-column uint32-nibble matvec and tiled BF16 QMM kernels, and open the fused BF16 dispatch path for bits==4 so Qwen3.6-27B MLPs stay on-device without regressing the 8-bit MoE path. --- .../kernels/mul_mm_affine_bf16_tiled4.comp | 207 ++++++++++++++++++ .../vulkan/kernels/mul_mv_affine4.comp | 194 ++++++++++++++++ .../vulkan/kernels/vulkan-shaders-gen.cpp | 17 ++ mlx/backend/vulkan/quantized.cpp | 171 ++++++++------- 4 files changed, 514 insertions(+), 75 deletions(-) create mode 100644 mlx/backend/vulkan/kernels/mul_mm_affine_bf16_tiled4.comp create mode 100644 mlx/backend/vulkan/kernels/mul_mv_affine4.comp diff --git a/mlx/backend/vulkan/kernels/mul_mm_affine_bf16_tiled4.comp b/mlx/backend/vulkan/kernels/mul_mm_affine_bf16_tiled4.comp new file mode 100644 index 0000000000..aa7507838a --- /dev/null +++ b/mlx/backend/vulkan/kernels/mul_mm_affine_bf16_tiled4.comp @@ -0,0 +1,207 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : enable +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require + +#include "types.glsl" + +layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in; + +layout(binding = 0) readonly buffer W { + uint data_w[]; +}; +layout(binding = 1) readonly buffer SCALES { + uint16_t data_scales[]; +}; +layout(binding = 2) readonly buffer BIASES { + uint16_t data_biases[]; +}; +layout(binding = 3) readonly buffer X { + uint16_t data_x[]; +}; +layout(binding = 4) writeonly buffer OUT { + uint16_t data_out[]; +}; + +layout(push_constant) uniform parameter { + uint rows; + uint cols; + uint K; + uint packed_row_bytes; + uint x_row_stride; + uint out_row_stride; + uint scale_row_stride; + uint bias_row_stride; + uint bits; + uint group_size; + uint num_groups; +} p; + +#ifndef MLX_BM +#define MLX_BM 32 +#endif +#ifndef MLX_BN +#define MLX_BN 16 +#endif +#ifndef MLX_BK +#define MLX_BK 32 +#endif +#ifndef MLX_TM +#define MLX_TM 4 +#endif +#ifndef MLX_TN +#define MLX_TN 2 +#endif + +const uint BM = MLX_BM; +const uint BN = MLX_BN; +const uint BK = MLX_BK; +const uint TM = MLX_TM; +const uint TN = MLX_TN; + +shared float xs[BM][BK]; +shared float ws[BN][BK]; +shared float tile_scales[BN]; +shared float tile_biases[BN]; + +void main() { + const uint lx = gl_LocalInvocationID.x; + const uint ly = gl_LocalInvocationID.y; + const uint row0 = gl_WorkGroupID.y * BM + ly * TM; + const uint col0 = gl_WorkGroupID.x * BN + lx * TN; + const uint linear = ly * gl_WorkGroupSize.x + lx; + + float acc[TM][TN]; + for (uint mr = 0u; mr < TM; ++mr) { + for (uint nc = 0u; nc < TN; ++nc) { + acc[mr][nc] = 0.0f; + } + } + + const uint thread_count = gl_WorkGroupSize.x * gl_WorkGroupSize.y; + for (uint kb = 0u; kb < p.K; kb += BK) { + if (linear < BN) { + const uint src_col = gl_WorkGroupID.x * BN + linear; + const uint group = kb / p.group_size; + if (src_col < p.cols) { + const uint scale_base = src_col * p.scale_row_stride; + const uint bias_base = src_col * p.bias_row_stride; + tile_scales[linear] = + bf16_to_fp32(uint(data_scales[scale_base + group])); + tile_biases[linear] = + bf16_to_fp32(uint(data_biases[bias_base + group])); + } else { + tile_scales[linear] = 0.0f; + tile_biases[linear] = 0.0f; + } + } + barrier(); + + for (uint i = linear; i < BM * BK; i += thread_count) { + const uint r = i / BK; + const uint kk = i - r * BK; + const uint src_row = gl_WorkGroupID.y * BM + r; + const uint src_k = kb + kk; + xs[r][kk] = (src_row < p.rows && src_k < p.K) + ? bf16_to_fp32(uint(data_x[src_row * p.x_row_stride + src_k])) + : 0.0f; + } + // 4-bit: 8 nibbles per uint32. + for (uint i = linear; i < BN * (BK / 8u); i += thread_count) { + const uint c = i / (BK / 8u); + const uint packed_k = i - c * (BK / 8u); + const uint src_col = gl_WorkGroupID.x * BN + c; + const uint src_k = kb + packed_k * 8u; + const uint kk = packed_k * 8u; + if (src_col < p.cols && src_k + 7u < p.K) { + const uint packed = data_w + [src_col * (p.packed_row_bytes / 4u) + src_k / 8u]; + const float scale = tile_scales[c]; + const float bias = tile_biases[c]; + [[unroll]] for (uint n = 0u; n < 8u; ++n) { + ws[c][kk + n] = + fma(scale, float((packed >> (4u * n)) & 0xFu), bias); + } + } else { + [[unroll]] for (uint n = 0u; n < 8u; ++n) { + ws[c][kk + n] = 0.0f; + } + } + } + barrier(); + + for (uint kk = 0u; kk < BK; ++kk) { +#if MLX_TN == 2 + const float w0 = ws[lx * TN + 0u][kk]; + const float w1 = ws[lx * TN + 1u][kk]; + for (uint mr = 0u; mr < TM; ++mr) { + const float x = xs[ly * TM + mr][kk]; + acc[mr][0] = fma(x, w0, acc[mr][0]); + acc[mr][1] = fma(x, w1, acc[mr][1]); + } +#elif MLX_TN == 4 + const float w0 = ws[lx * TN + 0u][kk]; + const float w1 = ws[lx * TN + 1u][kk]; + const float w2 = ws[lx * TN + 2u][kk]; + const float w3 = ws[lx * TN + 3u][kk]; + for (uint mr = 0u; mr < TM; ++mr) { + const float x = xs[ly * TM + mr][kk]; + acc[mr][0] = fma(x, w0, acc[mr][0]); + acc[mr][1] = fma(x, w1, acc[mr][1]); + acc[mr][2] = fma(x, w2, acc[mr][2]); + acc[mr][3] = fma(x, w3, acc[mr][3]); + } +#else + for (uint mr = 0u; mr < TM; ++mr) { + const float x = xs[ly * TM + mr][kk]; + for (uint nc = 0u; nc < TN; ++nc) { + acc[mr][nc] = fma(x, ws[lx * TN + nc][kk], acc[mr][nc]); + } + } +#endif + } + barrier(); + } + + for (uint mr = 0u; mr < TM; ++mr) { + const uint row = row0 + mr; +#if MLX_TN == 2 + if (row < p.rows && col0 < p.cols) { + data_out[row * p.out_row_stride + col0] = + uint16_t(fp32_to_bf16(acc[mr][0])); + } + if (row < p.rows && col0 + 1u < p.cols) { + data_out[row * p.out_row_stride + col0 + 1u] = + uint16_t(fp32_to_bf16(acc[mr][1])); + } +#elif MLX_TN == 4 + if (row < p.rows && col0 < p.cols) { + data_out[row * p.out_row_stride + col0] = + uint16_t(fp32_to_bf16(acc[mr][0])); + } + if (row < p.rows && col0 + 1u < p.cols) { + data_out[row * p.out_row_stride + col0 + 1u] = + uint16_t(fp32_to_bf16(acc[mr][1])); + } + if (row < p.rows && col0 + 2u < p.cols) { + data_out[row * p.out_row_stride + col0 + 2u] = + uint16_t(fp32_to_bf16(acc[mr][2])); + } + if (row < p.rows && col0 + 3u < p.cols) { + data_out[row * p.out_row_stride + col0 + 3u] = + uint16_t(fp32_to_bf16(acc[mr][3])); + } +#else + if (row < p.rows) { + for (uint nc = 0u; nc < TN; ++nc) { + if (col0 + nc < p.cols) { + data_out[row * p.out_row_stride + col0 + nc] = + uint16_t(fp32_to_bf16(acc[mr][nc])); + } + } + } +#endif + } +} diff --git a/mlx/backend/vulkan/kernels/mul_mv_affine4.comp b/mlx/backend/vulkan/kernels/mul_mv_affine4.comp new file mode 100644 index 0000000000..b39ca00ab6 --- /dev/null +++ b/mlx/backend/vulkan/kernels/mul_mv_affine4.comp @@ -0,0 +1,194 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : enable +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_arithmetic : require + +#ifdef FLOAT16 +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#endif + +#include "types.glsl" + +#if !defined(TO_FLOAT_TYPE) +#define TO_FLOAT_TYPE float +#endif +#if !defined(S_TYPE) +#define S_TYPE float +#endif +#if !defined(D_TYPE) +#define D_TYPE float +#endif +#if !defined(SCALE_TO_FLOAT_TYPE) +#define SCALE_TO_FLOAT_TYPE(x) float(x) +#endif +#if !defined(FROM_FLOAT_TYPE) +#define FROM_FLOAT_TYPE(x) (x) +#endif + +// Reuse activations across output columns (llama.cpp mul_mat_vec style). +// Wave64 workgroup: one subgroup covers the reduce on AMD. +#define NUM_COLS 8 +#define BLOCK_SIZE 64 + +layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer W { + uint data_w[]; +}; +layout(binding = 1) readonly buffer SCALES { + S_TYPE data_scales[]; +}; +layout(binding = 2) readonly buffer BIASES { + S_TYPE data_biases[]; +}; +layout(binding = 3) readonly buffer X { + B_TYPE data_x[]; +}; +layout(binding = 4) writeonly buffer OUT { + D_TYPE data_out[]; +}; + +layout(push_constant) uniform parameter { + uint rows; + uint cols; + uint K; + uint packed_row_bytes; + uint x_row_stride; + uint out_row_stride; + uint scale_row_stride; + uint bias_row_stride; + uint bits; + uint group_size; + uint num_groups; +} p; + +shared float partial[NUM_COLS][BLOCK_SIZE]; + +void main() { + const uint col0 = gl_WorkGroupID.x * NUM_COLS; + const uint row = gl_WorkGroupID.y; + const uint tid = gl_LocalInvocationID.x; + + if (row >= p.rows || col0 >= p.cols) { + return; + } + + const uint num_cols = min(NUM_COLS, p.cols - col0); + const uint x_row_base = row * p.x_row_stride; + + uint w_row_words[NUM_COLS]; + uint scale_row_base[NUM_COLS]; + uint bias_row_base[NUM_COLS]; + [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { + if (c < num_cols) { + const uint col = col0 + c; + w_row_words[c] = (col * p.packed_row_bytes) >> 2; + scale_row_base[c] = col * p.scale_row_stride; + bias_row_base[c] = col * p.bias_row_stride; + } + } + + float acc[NUM_COLS]; + [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { + acc[c] = 0.0f; + } + + // 4-bit affine weights pack 8 nibbles per uint32 (low nibble first). + if ((p.K & 7u) == 0u && (p.group_size & 7u) == 0u) { + const uint packs = p.K >> 3; + for (uint pidx = tid; pidx < packs; pidx += BLOCK_SIZE) { + const uint k = pidx << 3; + const uint group = k / p.group_size; + const float x0 = TO_FLOAT_TYPE(data_x[x_row_base + k]); + const float x1 = TO_FLOAT_TYPE(data_x[x_row_base + k + 1u]); + const float x2 = TO_FLOAT_TYPE(data_x[x_row_base + k + 2u]); + const float x3 = TO_FLOAT_TYPE(data_x[x_row_base + k + 3u]); + const float x4 = TO_FLOAT_TYPE(data_x[x_row_base + k + 4u]); + const float x5 = TO_FLOAT_TYPE(data_x[x_row_base + k + 5u]); + const float x6 = TO_FLOAT_TYPE(data_x[x_row_base + k + 6u]); + const float x7 = TO_FLOAT_TYPE(data_x[x_row_base + k + 7u]); + + [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { + if (c < num_cols) { + const float s = + SCALE_TO_FLOAT_TYPE(data_scales[scale_row_base[c] + group]); + const float b = + SCALE_TO_FLOAT_TYPE(data_biases[bias_row_base[c] + group]); + const uint packed = data_w[w_row_words[c] + pidx]; + acc[c] = fma(x0, float(packed & 0xFu) * s + b, acc[c]); + acc[c] = fma( + x1, float((packed >> 4) & 0xFu) * s + b, acc[c]); + acc[c] = fma( + x2, float((packed >> 8) & 0xFu) * s + b, acc[c]); + acc[c] = fma( + x3, float((packed >> 12) & 0xFu) * s + b, acc[c]); + acc[c] = fma( + x4, float((packed >> 16) & 0xFu) * s + b, acc[c]); + acc[c] = fma( + x5, float((packed >> 20) & 0xFu) * s + b, acc[c]); + acc[c] = fma( + x6, float((packed >> 24) & 0xFu) * s + b, acc[c]); + acc[c] = fma( + x7, float((packed >> 28) & 0xFu) * s + b, acc[c]); + } + } + } + } else { + for (uint k = tid; k < p.K; k += BLOCK_SIZE) { + const uint group = k / p.group_size; + const float xval = TO_FLOAT_TYPE(data_x[x_row_base + k]); + [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { + if (c < num_cols) { + const uint word = data_w[w_row_words[c] + (k >> 3)]; + const uint q = (word >> ((k & 7u) * 4u)) & 0xFu; + const float w = + SCALE_TO_FLOAT_TYPE(data_scales[scale_row_base[c] + group]) * + float(q) + + SCALE_TO_FLOAT_TYPE(data_biases[bias_row_base[c] + group]); + acc[c] = fma(xval, w, acc[c]); + } + } + } + } + + if (gl_SubgroupSize == BLOCK_SIZE && gl_NumSubgroups == 1u) { + [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { + acc[c] = subgroupAdd(acc[c]); + } + if (tid == 0u) { + [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { + if (c < num_cols) { + data_out[row * p.out_row_stride + col0 + c] = + D_TYPE(FROM_FLOAT_TYPE(acc[c])); + } + } + } + return; + } + + [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { + partial[c][tid] = acc[c]; + } + barrier(); + for (uint stride = BLOCK_SIZE >> 1u; stride > 0u; stride >>= 1u) { + if (tid < stride) { + [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { + partial[c][tid] += partial[c][tid + stride]; + } + } + barrier(); + } + + if (tid == 0u) { + [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { + if (c < num_cols) { + data_out[row * p.out_row_stride + col0 + c] = + D_TYPE(FROM_FLOAT_TYPE(partial[c][0])); + } + } + } +} diff --git a/mlx/backend/vulkan/kernels/vulkan-shaders-gen.cpp b/mlx/backend/vulkan/kernels/vulkan-shaders-gen.cpp index 916af52f6c..1b094f7971 100644 --- a/mlx/backend/vulkan/kernels/vulkan-shaders-gen.cpp +++ b/mlx/backend/vulkan/kernels/vulkan-shaders-gen.cpp @@ -2270,6 +2270,23 @@ void process_shaders() { {"TO_FLOAT_TYPE", "bf16_to_fp32"}, {"SCALE_TO_FLOAT_TYPE(x)", "bf16_to_fp32(uint(x))"}, {"FROM_FLOAT_TYPE", "fp32_to_bf16"}}); + string_to_spv( + "fused_affine_matvec4_bf16_bf16", + "mul_mv_affine4.comp", + {{"B_TYPE", "uint16_t"}, + {"S_TYPE", "uint16_t"}, + {"D_TYPE", "uint16_t"}, + {"TO_FLOAT_TYPE", "bf16_to_fp32"}, + {"SCALE_TO_FLOAT_TYPE(x)", "bf16_to_fp32(uint(x))"}, + {"FROM_FLOAT_TYPE", "fp32_to_bf16"}}); + string_to_spv( + "fused_affine_qmm_bf16_bf16_tiled4", + "mul_mm_affine_bf16_tiled4.comp", + {}); + string_to_spv( + "fused_affine_qmm_bf16_bf16_tiled4_n32", + "mul_mm_affine_bf16_tiled4.comp", + {{"MLX_BN", "32"}, {"MLX_TN", "4"}}); string_to_spv( "fused_affine_matvec_f32_f32", "mul_mv_affine.comp", diff --git a/mlx/backend/vulkan/quantized.cpp b/mlx/backend/vulkan/quantized.cpp index e218da4267..6ff598b638 100644 --- a/mlx/backend/vulkan/quantized.cpp +++ b/mlx/backend/vulkan/quantized.cpp @@ -1829,8 +1829,9 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { return true; }(); + const bool fused_bf16_bits = bits_ == 8 || bits_ == 4; if (mode_ == QuantizationMode::Affine && enable_fused_decode_qmm && - transpose_ && bits_ == 8 && x_mat.dtype() == bfloat16 && + transpose_ && fused_bf16_bits && x_mat.dtype() == bfloat16 && out.dtype() == bfloat16 && inputs[2].dtype() == bfloat16 && inputs[3].dtype() == bfloat16 && x_mat.ndim() == 2 && w.ndim() == 2) { array scales_bf16 = ensure_row_contiguous_zero_offset(inputs[2], s); @@ -1886,82 +1887,102 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { return; } - array out_work( - (vector_lhs || flatten_lhs_batches) - ? Shape{static_cast(rows), out.shape(-1)} - : out.shape(), - bfloat16, - nullptr, - {}); - out_work.set_data(allocator::malloc(out_work.nbytes())); - if (out_work.size() != 0) { - vulkan::FusedAffineMatmulPushConstants push_constants{}; - push_constants.rows = rows; - push_constants.cols = cols; - push_constants.K = k; - push_constants.packed_row_bytes = - static_cast(w.strides(-2) * sizeof(uint32_t)); - push_constants.x_row_stride = - static_cast(x_mat.strides(-2)); - push_constants.out_row_stride = - static_cast(out_work.strides(-2)); - push_constants.scale_row_stride = - static_cast(scales_bf16.strides(-2)); - push_constants.bias_row_stride = - static_cast(biases_bf16.strides(-2)); - push_constants.bits = static_cast(bits_); - push_constants.group_size = static_cast(group_size_); - push_constants.num_groups = num_groups; - - const bool use_decode_matvec = rows == 1 || decode_lhs; - const bool use_tiled_prefill = rows > 1 && !decode_lhs && - group_size_ >= 32 && (group_size_ % 32) == 0 && - fused_affine_bf16_tiled_prefill_enabled(); - const bool use_large_n_tile = use_tiled_prefill && cols >= 1024; - const auto shader_id = use_decode_matvec - ? vulkan::StaticShaderId::fused_affine_matvec8_bf16_bf16 - : use_large_n_tile - ? vulkan::StaticShaderId::fused_affine_qmm_bf16_bf16_tiled_n32 - : use_tiled_prefill - ? vulkan::StaticShaderId::fused_affine_qmm_bf16_bf16_tiled - : vulkan::StaticShaderId::fused_affine_qmm_bf16_bf16; - const std::array grid = use_decode_matvec - ? std::array{cols, rows, 1u} - : shader_id == - vulkan::StaticShaderId:: - fused_affine_qmm_bf16_bf16_tiled_n32 - ? std::array< - uint32_t, - 3>{(cols + 31u) / 32u, (rows + 31u) / 32u, 1u} - : shader_id == - vulkan::StaticShaderId::fused_affine_qmm_bf16_bf16_tiled - ? std::array< - uint32_t, - 3>{(cols + 15u) / 16u, (rows + 31u) / 32u, 1u} - : std::array{ - (cols + 15u) / 16u, (rows + 15u) / 16u, 1u}; + const bool use_decode_matvec = rows == 1 || decode_lhs; + const bool use_tiled_prefill = rows > 1 && !decode_lhs && + group_size_ >= 32 && (group_size_ % 32) == 0 && + fused_affine_bf16_tiled_prefill_enabled(); + // 4-bit has no scalar BF16 acc kernel; require tiled prefill. + if (bits_ == 4 && !use_decode_matvec && !use_tiled_prefill) { + // Fall through to the staged float32 fused path below. + } else { + array out_work( + (vector_lhs || flatten_lhs_batches) + ? Shape{static_cast(rows), out.shape(-1)} + : out.shape(), + bfloat16, + nullptr, + {}); + out_work.set_data(allocator::malloc(out_work.nbytes())); + if (out_work.size() != 0) { + vulkan::FusedAffineMatmulPushConstants push_constants{}; + push_constants.rows = rows; + push_constants.cols = cols; + push_constants.K = k; + push_constants.packed_row_bytes = + static_cast(w.strides(-2) * sizeof(uint32_t)); + push_constants.x_row_stride = + static_cast(x_mat.strides(-2)); + push_constants.out_row_stride = + static_cast(out_work.strides(-2)); + push_constants.scale_row_stride = + static_cast(scales_bf16.strides(-2)); + push_constants.bias_row_stride = + static_cast(biases_bf16.strides(-2)); + push_constants.bits = static_cast(bits_); + push_constants.group_size = static_cast(group_size_); + push_constants.num_groups = num_groups; + + const bool use_large_n_tile = use_tiled_prefill && cols >= 1024; + vulkan::StaticShaderId shader_id; + std::array grid; + if (use_decode_matvec) { + if (bits_ == 4) { + shader_id = + vulkan::StaticShaderId::fused_affine_matvec4_bf16_bf16; + grid = {(cols + 7u) / 8u, rows, 1u}; + } else { + shader_id = + vulkan::StaticShaderId::fused_affine_matvec8_bf16_bf16; + grid = {cols, rows, 1u}; + } + } else if (bits_ == 4) { + shader_id = use_large_n_tile + ? vulkan::StaticShaderId:: + fused_affine_qmm_bf16_bf16_tiled4_n32 + : vulkan::StaticShaderId::fused_affine_qmm_bf16_bf16_tiled4; + grid = use_large_n_tile + ? std::array{ + (cols + 31u) / 32u, (rows + 31u) / 32u, 1u} + : std::array{ + (cols + 15u) / 16u, (rows + 31u) / 32u, 1u}; + } else { + shader_id = use_large_n_tile + ? vulkan::StaticShaderId::fused_affine_qmm_bf16_bf16_tiled_n32 + : use_tiled_prefill + ? vulkan::StaticShaderId::fused_affine_qmm_bf16_bf16_tiled + : vulkan::StaticShaderId::fused_affine_qmm_bf16_bf16; + grid = use_large_n_tile + ? std::array{ + (cols + 31u) / 32u, (rows + 31u) / 32u, 1u} + : use_tiled_prefill + ? std::array{ + (cols + 15u) / 16u, (rows + 31u) / 32u, 1u} + : std::array{ + (cols + 15u) / 16u, (rows + 15u) / 16u, 1u}; + } + + auto command_buffer = vulkan::begin_command_recording(s.index); + vulkan::dispatch_fused_affine_matmul_op( + w, + scales_bf16, + biases_bf16, + x_mat, + out_work, + shader_id, + command_buffer, + s, + push_constants, + grid); + vulkan::end_command_recording(s.index); + } - auto command_buffer = vulkan::begin_command_recording(s.index); - vulkan::dispatch_fused_affine_matmul_op( - w, - scales_bf16, - biases_bf16, - x_mat, - out_work, - shader_id, - command_buffer, - s, - push_constants, - grid); - vulkan::end_command_recording(s.index); + finalize_bf16_output(out_work); + trace_qmm( + "fused_bf16", + (vector_lhs || flatten_lhs_batches) ? "reshaped_output=1" + : "reshaped_output=0"); + return; } - - finalize_bf16_output(out_work); - trace_qmm( - "fused_bf16", - (vector_lhs || flatten_lhs_batches) ? "reshaped_output=1" - : "reshaped_output=0"); - return; } } } From f286dd90917b54c3fc1513e5a30fb5e9dd3a566f Mon Sep 17 00:00:00 2001 From: Goni Zahavy Date: Wed, 15 Jul 2026 21:27:58 +0300 Subject: [PATCH 2/9] [Vulkan] Multi-column dense affine8 decode matvec Restore NUM_COLS=8 + wave64 packing for fused_affine_matvec8 so 8-bit QKVO projections reuse activations across columns. Keep the 4-bit path on the same (cols+7)/8 grid. --- .../vulkan/kernels/mul_mv_affine8.comp | 115 +++++++++++++----- mlx/backend/vulkan/quantized.cpp | 7 +- 2 files changed, 89 insertions(+), 33 deletions(-) diff --git a/mlx/backend/vulkan/kernels/mul_mv_affine8.comp b/mlx/backend/vulkan/kernels/mul_mv_affine8.comp index af09ef4378..567a1fec27 100644 --- a/mlx/backend/vulkan/kernels/mul_mv_affine8.comp +++ b/mlx/backend/vulkan/kernels/mul_mv_affine8.comp @@ -1,5 +1,6 @@ #version 450 +#extension GL_EXT_control_flow_attributes : enable #extension GL_EXT_shader_16bit_storage : require #extension GL_EXT_shader_8bit_storage : require #extension GL_EXT_shader_explicit_arithmetic_types_int8 : require @@ -28,7 +29,9 @@ #define FROM_FLOAT_TYPE(x) (x) #endif -#define BLOCK_SIZE 256 +// Reuse activations across output columns. Wave64 workgroup for AMD. +#define NUM_COLS 8 +#define BLOCK_SIZE 64 layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; @@ -62,67 +65,117 @@ layout(push_constant) uniform parameter { uint num_groups; } p; -// Size to the workgroup so gl_SubgroupID is always in-range for any legal -// subgroup size (NumSubgroups <= BLOCK_SIZE). -shared float partial[BLOCK_SIZE]; +shared float partial[NUM_COLS][BLOCK_SIZE]; void main() { - const uint col = gl_WorkGroupID.x; + const uint col0 = gl_WorkGroupID.x * NUM_COLS; const uint row = gl_WorkGroupID.y; const uint tid = gl_LocalInvocationID.x; - if (row >= p.rows || col >= p.cols) { + if (row >= p.rows || col0 >= p.cols) { return; } - const uint w_row_words = (col * p.packed_row_bytes) >> 2; - const uint scale_row_base = col * p.scale_row_stride; - const uint bias_row_base = col * p.bias_row_stride; + const uint num_cols = min(NUM_COLS, p.cols - col0); const uint x_row_base = row * p.x_row_stride; - float acc = 0.0f; + uint w_row_words[NUM_COLS]; + uint scale_row_base[NUM_COLS]; + uint bias_row_base[NUM_COLS]; + [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { + if (c < num_cols) { + const uint col = col0 + c; + w_row_words[c] = (col * p.packed_row_bytes) >> 2; + scale_row_base[c] = col * p.scale_row_stride; + bias_row_base[c] = col * p.bias_row_stride; + } + } + + float acc[NUM_COLS]; + [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { + acc[c] = 0.0f; + } + + // 8-bit affine weights pack 4 values per uint32. if ((p.K & 3u) == 0u && (p.group_size & 3u) == 0u) { const uint packs = p.K >> 2; for (uint pidx = tid; pidx < packs; pidx += BLOCK_SIZE) { const uint k = pidx << 2; const uint group = k / p.group_size; - const float s = SCALE_TO_FLOAT_TYPE(data_scales[scale_row_base + group]); - const float b = SCALE_TO_FLOAT_TYPE(data_biases[bias_row_base + group]); - const uint packed = data_w[w_row_words + pidx]; - acc = fma(TO_FLOAT_TYPE(data_x[x_row_base + k]), float(packed & 0xffu) * s + b, acc); - acc = fma(TO_FLOAT_TYPE(data_x[x_row_base + k + 1u]), float((packed >> 8) & 0xffu) * s + b, acc); - acc = fma(TO_FLOAT_TYPE(data_x[x_row_base + k + 2u]), float((packed >> 16) & 0xffu) * s + b, acc); - acc = fma(TO_FLOAT_TYPE(data_x[x_row_base + k + 3u]), float((packed >> 24) & 0xffu) * s + b, acc); + const float x0 = TO_FLOAT_TYPE(data_x[x_row_base + k]); + const float x1 = TO_FLOAT_TYPE(data_x[x_row_base + k + 1u]); + const float x2 = TO_FLOAT_TYPE(data_x[x_row_base + k + 2u]); + const float x3 = TO_FLOAT_TYPE(data_x[x_row_base + k + 3u]); + + [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { + if (c < num_cols) { + const float s = + SCALE_TO_FLOAT_TYPE(data_scales[scale_row_base[c] + group]); + const float b = + SCALE_TO_FLOAT_TYPE(data_biases[bias_row_base[c] + group]); + const uint packed = data_w[w_row_words[c] + pidx]; + acc[c] = fma(x0, float(packed & 0xffu) * s + b, acc[c]); + acc[c] = fma( + x1, float((packed >> 8) & 0xffu) * s + b, acc[c]); + acc[c] = fma( + x2, float((packed >> 16) & 0xffu) * s + b, acc[c]); + acc[c] = fma( + x3, float((packed >> 24) & 0xffu) * s + b, acc[c]); + } + } } } else { for (uint k = tid; k < p.K; k += BLOCK_SIZE) { const uint group = k / p.group_size; - const uint word = data_w[w_row_words + (k >> 2)]; - const uint q = (word >> ((k & 3u) * 8u)) & 0xffu; - const float w = SCALE_TO_FLOAT_TYPE(data_scales[scale_row_base + group]) * - float(q) + SCALE_TO_FLOAT_TYPE(data_biases[bias_row_base + group]); - acc = fma(TO_FLOAT_TYPE(data_x[x_row_base + k]), w, acc); + const float xval = TO_FLOAT_TYPE(data_x[x_row_base + k]); + [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { + if (c < num_cols) { + const uint word = data_w[w_row_words[c] + (k >> 2)]; + const uint q = (word >> ((k & 3u) * 8u)) & 0xffu; + const float w = + SCALE_TO_FLOAT_TYPE(data_scales[scale_row_base[c] + group]) * + float(q) + + SCALE_TO_FLOAT_TYPE(data_biases[bias_row_base[c] + group]); + acc[c] = fma(xval, w, acc[c]); + } + } } } if (gl_SubgroupSize == BLOCK_SIZE && gl_NumSubgroups == 1u) { - acc = subgroupAdd(acc); + [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { + acc[c] = subgroupAdd(acc[c]); + } if (tid == 0u) { - data_out[row * p.out_row_stride + col] = D_TYPE(FROM_FLOAT_TYPE(acc)); + [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { + if (c < num_cols) { + data_out[row * p.out_row_stride + col0 + c] = + D_TYPE(FROM_FLOAT_TYPE(acc[c])); + } + } } return; } - acc = subgroupAdd(acc); - if (gl_SubgroupInvocationID == 0u) { - partial[gl_SubgroupID] = acc; + [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { + partial[c][tid] = acc[c]; } barrier(); + for (uint stride = BLOCK_SIZE >> 1u; stride > 0u; stride >>= 1u) { + if (tid < stride) { + [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { + partial[c][tid] += partial[c][tid + stride]; + } + } + barrier(); + } + if (tid == 0u) { - float total = 0.0f; - for (uint i = 0u; i < gl_NumSubgroups; ++i) { - total += partial[i]; + [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { + if (c < num_cols) { + data_out[row * p.out_row_stride + col0 + c] = + D_TYPE(FROM_FLOAT_TYPE(partial[c][0])); + } } - data_out[row * p.out_row_stride + col] = D_TYPE(FROM_FLOAT_TYPE(total)); } } diff --git a/mlx/backend/vulkan/quantized.cpp b/mlx/backend/vulkan/quantized.cpp index 6ff598b638..71e8c69eb5 100644 --- a/mlx/backend/vulkan/quantized.cpp +++ b/mlx/backend/vulkan/quantized.cpp @@ -1929,12 +1929,12 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { if (bits_ == 4) { shader_id = vulkan::StaticShaderId::fused_affine_matvec4_bf16_bf16; - grid = {(cols + 7u) / 8u, rows, 1u}; } else { shader_id = vulkan::StaticShaderId::fused_affine_matvec8_bf16_bf16; - grid = {cols, rows, 1u}; } + // Both affine4/affine8 decode matvecs process NUM_COLS=8. + grid = {(cols + 7u) / 8u, rows, 1u}; } else if (bits_ == 4) { shader_id = use_large_n_tile ? vulkan::StaticShaderId:: @@ -2060,10 +2060,13 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { push_constants.group_size = static_cast(group_size_); push_constants.num_groups = num_groups; + // fused_affine_matvec8 uses NUM_COLS=8; generic matvec is 1-col/WG. const std::array grid = prefill_like_rows ? std::array< uint32_t, 3>{(cols + 15u) / 16u, (rows + 31u) / 32u, 1u} + : bits_ == 8 + ? std::array{(cols + 7u) / 8u, rows, 1u} : std::array{cols, rows, 1u}; auto command_buffer = vulkan::begin_command_recording(s.index); From 80814ae71efe1c74a68b7156180e2fb249578ae3 Mon Sep 17 00:00:00 2001 From: Goni Zahavy Date: Wed, 15 Jul 2026 21:48:10 +0300 Subject: [PATCH 3/9] [Vulkan] Keep dense affine8 decode single-column MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Multi-column affine8 regressed Qwen3-0.6B-8bit generation (~121→70 tok/s). Leave NUM_COLS=8 on the new affine4 path only. --- .../vulkan/kernels/mul_mv_affine8.comp | 115 +++++------------- mlx/backend/vulkan/quantized.cpp | 7 +- 2 files changed, 33 insertions(+), 89 deletions(-) diff --git a/mlx/backend/vulkan/kernels/mul_mv_affine8.comp b/mlx/backend/vulkan/kernels/mul_mv_affine8.comp index 567a1fec27..af09ef4378 100644 --- a/mlx/backend/vulkan/kernels/mul_mv_affine8.comp +++ b/mlx/backend/vulkan/kernels/mul_mv_affine8.comp @@ -1,6 +1,5 @@ #version 450 -#extension GL_EXT_control_flow_attributes : enable #extension GL_EXT_shader_16bit_storage : require #extension GL_EXT_shader_8bit_storage : require #extension GL_EXT_shader_explicit_arithmetic_types_int8 : require @@ -29,9 +28,7 @@ #define FROM_FLOAT_TYPE(x) (x) #endif -// Reuse activations across output columns. Wave64 workgroup for AMD. -#define NUM_COLS 8 -#define BLOCK_SIZE 64 +#define BLOCK_SIZE 256 layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; @@ -65,117 +62,67 @@ layout(push_constant) uniform parameter { uint num_groups; } p; -shared float partial[NUM_COLS][BLOCK_SIZE]; +// Size to the workgroup so gl_SubgroupID is always in-range for any legal +// subgroup size (NumSubgroups <= BLOCK_SIZE). +shared float partial[BLOCK_SIZE]; void main() { - const uint col0 = gl_WorkGroupID.x * NUM_COLS; + const uint col = gl_WorkGroupID.x; const uint row = gl_WorkGroupID.y; const uint tid = gl_LocalInvocationID.x; - if (row >= p.rows || col0 >= p.cols) { + if (row >= p.rows || col >= p.cols) { return; } - const uint num_cols = min(NUM_COLS, p.cols - col0); + const uint w_row_words = (col * p.packed_row_bytes) >> 2; + const uint scale_row_base = col * p.scale_row_stride; + const uint bias_row_base = col * p.bias_row_stride; const uint x_row_base = row * p.x_row_stride; - uint w_row_words[NUM_COLS]; - uint scale_row_base[NUM_COLS]; - uint bias_row_base[NUM_COLS]; - [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { - if (c < num_cols) { - const uint col = col0 + c; - w_row_words[c] = (col * p.packed_row_bytes) >> 2; - scale_row_base[c] = col * p.scale_row_stride; - bias_row_base[c] = col * p.bias_row_stride; - } - } - - float acc[NUM_COLS]; - [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { - acc[c] = 0.0f; - } - - // 8-bit affine weights pack 4 values per uint32. + float acc = 0.0f; if ((p.K & 3u) == 0u && (p.group_size & 3u) == 0u) { const uint packs = p.K >> 2; for (uint pidx = tid; pidx < packs; pidx += BLOCK_SIZE) { const uint k = pidx << 2; const uint group = k / p.group_size; - const float x0 = TO_FLOAT_TYPE(data_x[x_row_base + k]); - const float x1 = TO_FLOAT_TYPE(data_x[x_row_base + k + 1u]); - const float x2 = TO_FLOAT_TYPE(data_x[x_row_base + k + 2u]); - const float x3 = TO_FLOAT_TYPE(data_x[x_row_base + k + 3u]); - - [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { - if (c < num_cols) { - const float s = - SCALE_TO_FLOAT_TYPE(data_scales[scale_row_base[c] + group]); - const float b = - SCALE_TO_FLOAT_TYPE(data_biases[bias_row_base[c] + group]); - const uint packed = data_w[w_row_words[c] + pidx]; - acc[c] = fma(x0, float(packed & 0xffu) * s + b, acc[c]); - acc[c] = fma( - x1, float((packed >> 8) & 0xffu) * s + b, acc[c]); - acc[c] = fma( - x2, float((packed >> 16) & 0xffu) * s + b, acc[c]); - acc[c] = fma( - x3, float((packed >> 24) & 0xffu) * s + b, acc[c]); - } - } + const float s = SCALE_TO_FLOAT_TYPE(data_scales[scale_row_base + group]); + const float b = SCALE_TO_FLOAT_TYPE(data_biases[bias_row_base + group]); + const uint packed = data_w[w_row_words + pidx]; + acc = fma(TO_FLOAT_TYPE(data_x[x_row_base + k]), float(packed & 0xffu) * s + b, acc); + acc = fma(TO_FLOAT_TYPE(data_x[x_row_base + k + 1u]), float((packed >> 8) & 0xffu) * s + b, acc); + acc = fma(TO_FLOAT_TYPE(data_x[x_row_base + k + 2u]), float((packed >> 16) & 0xffu) * s + b, acc); + acc = fma(TO_FLOAT_TYPE(data_x[x_row_base + k + 3u]), float((packed >> 24) & 0xffu) * s + b, acc); } } else { for (uint k = tid; k < p.K; k += BLOCK_SIZE) { const uint group = k / p.group_size; - const float xval = TO_FLOAT_TYPE(data_x[x_row_base + k]); - [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { - if (c < num_cols) { - const uint word = data_w[w_row_words[c] + (k >> 2)]; - const uint q = (word >> ((k & 3u) * 8u)) & 0xffu; - const float w = - SCALE_TO_FLOAT_TYPE(data_scales[scale_row_base[c] + group]) * - float(q) + - SCALE_TO_FLOAT_TYPE(data_biases[bias_row_base[c] + group]); - acc[c] = fma(xval, w, acc[c]); - } - } + const uint word = data_w[w_row_words + (k >> 2)]; + const uint q = (word >> ((k & 3u) * 8u)) & 0xffu; + const float w = SCALE_TO_FLOAT_TYPE(data_scales[scale_row_base + group]) * + float(q) + SCALE_TO_FLOAT_TYPE(data_biases[bias_row_base + group]); + acc = fma(TO_FLOAT_TYPE(data_x[x_row_base + k]), w, acc); } } if (gl_SubgroupSize == BLOCK_SIZE && gl_NumSubgroups == 1u) { - [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { - acc[c] = subgroupAdd(acc[c]); - } + acc = subgroupAdd(acc); if (tid == 0u) { - [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { - if (c < num_cols) { - data_out[row * p.out_row_stride + col0 + c] = - D_TYPE(FROM_FLOAT_TYPE(acc[c])); - } - } + data_out[row * p.out_row_stride + col] = D_TYPE(FROM_FLOAT_TYPE(acc)); } return; } - [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { - partial[c][tid] = acc[c]; + acc = subgroupAdd(acc); + if (gl_SubgroupInvocationID == 0u) { + partial[gl_SubgroupID] = acc; } barrier(); - for (uint stride = BLOCK_SIZE >> 1u; stride > 0u; stride >>= 1u) { - if (tid < stride) { - [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { - partial[c][tid] += partial[c][tid + stride]; - } - } - barrier(); - } - if (tid == 0u) { - [[unroll]] for (uint c = 0; c < NUM_COLS; ++c) { - if (c < num_cols) { - data_out[row * p.out_row_stride + col0 + c] = - D_TYPE(FROM_FLOAT_TYPE(partial[c][0])); - } + float total = 0.0f; + for (uint i = 0u; i < gl_NumSubgroups; ++i) { + total += partial[i]; } + data_out[row * p.out_row_stride + col] = D_TYPE(FROM_FLOAT_TYPE(total)); } } diff --git a/mlx/backend/vulkan/quantized.cpp b/mlx/backend/vulkan/quantized.cpp index 71e8c69eb5..6ff598b638 100644 --- a/mlx/backend/vulkan/quantized.cpp +++ b/mlx/backend/vulkan/quantized.cpp @@ -1929,12 +1929,12 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { if (bits_ == 4) { shader_id = vulkan::StaticShaderId::fused_affine_matvec4_bf16_bf16; + grid = {(cols + 7u) / 8u, rows, 1u}; } else { shader_id = vulkan::StaticShaderId::fused_affine_matvec8_bf16_bf16; + grid = {cols, rows, 1u}; } - // Both affine4/affine8 decode matvecs process NUM_COLS=8. - grid = {(cols + 7u) / 8u, rows, 1u}; } else if (bits_ == 4) { shader_id = use_large_n_tile ? vulkan::StaticShaderId:: @@ -2060,13 +2060,10 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { push_constants.group_size = static_cast(group_size_); push_constants.num_groups = num_groups; - // fused_affine_matvec8 uses NUM_COLS=8; generic matvec is 1-col/WG. const std::array grid = prefill_like_rows ? std::array< uint32_t, 3>{(cols + 15u) / 16u, (rows + 31u) / 32u, 1u} - : bits_ == 8 - ? std::array{(cols + 7u) / 8u, rows, 1u} : std::array{cols, rows, 1u}; auto command_buffer = vulkan::begin_command_recording(s.index); From a3e17ef42c44cfa29824005b4d34857faa17b836 Mon Sep 17 00:00:00 2001 From: Goni Zahavy Date: Wed, 15 Jul 2026 21:53:58 +0300 Subject: [PATCH 4/9] [Vulkan] Dense coopmat affine-4 QMM for large prefill Replace dequant+BF16 GEMM for bits==4 rows>=256 with a fused cooperative-matrix kernel that keeps weights packed. Cuts MLP gate 4096x17408x5120 from ~236 ms to ~68 ms on Strix Halo. --- .../kernels/mul_mm_affine_bf16_coop4.comp | 352 ++++++++++++++++++ .../vulkan/kernels/vulkan-shaders-gen.cpp | 6 + mlx/backend/vulkan/quantized.cpp | 89 ++++- 3 files changed, 445 insertions(+), 2 deletions(-) create mode 100644 mlx/backend/vulkan/kernels/mul_mm_affine_bf16_coop4.comp diff --git a/mlx/backend/vulkan/kernels/mul_mm_affine_bf16_coop4.comp b/mlx/backend/vulkan/kernels/mul_mm_affine_bf16_coop4.comp new file mode 100644 index 0000000000..732e1e70f8 --- /dev/null +++ b/mlx/backend/vulkan/kernels/mul_mm_affine_bf16_coop4.comp @@ -0,0 +1,352 @@ +#version 450 + +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_arithmetic : enable +#extension GL_KHR_shader_subgroup_basic : require + +#include "types.glsl" + +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer W { + uint data_w[]; +}; +layout(binding = 1) readonly buffer SCALES { + uint16_t data_scales[]; +}; +layout(binding = 2) readonly buffer BIASES { + uint16_t data_biases[]; +}; +layout(binding = 3) readonly buffer X { + uint data_x[]; +}; +layout(binding = 4) writeonly buffer OUT { + uint16_t data_out[]; +}; + +layout(push_constant) uniform parameter { + uint rows; + uint cols; + uint K; + uint packed_row_bytes; + uint x_row_stride; + uint out_row_stride; + uint scale_row_stride; + uint bias_row_stride; + uint bits; + uint group_size; + uint num_groups; +} p; + +const uint BM = 64u; +const uint BN = 128u; +const uint BK = 16u; +const uint SHMEM_STRIDE = BK / 2u + 4u; + +shared f16vec2 xs[BM * SHMEM_STRIDE]; +shared f16vec2 ws[BN * SHMEM_STRIDE]; +shared float16_t output_stage[BM * BN]; +shared uint x_scale_exponents[BM]; +shared float x_scale_factors[BM]; +shared uint w_scale_exponents[BN]; +shared float w_scales[BN]; +shared float w_biases[BN]; +shared uint tile_w_exponents[8]; + +float scale_down(const float value, const uint exponent) { + return value * uintBitsToFloat((127u - exponent) << 23u); +} + +float scale_up(const float value, const uint exponent) { + return value * uintBitsToFloat((127u + exponent) << 23u); +} + +uint fp32_biased_exponent(const float value) { + const uint exponent = (floatBitsToUint(abs(value)) >> 23u) & 0xffu; + return exponent < 255u ? exponent : 0u; +} + +uint fp16_range_scale_exponent(const uint exponent) { + return exponent >= 142u ? exponent - 141u : 0u; +} + +uint ceil_log2(const uint value) { + return value <= 1u ? 0u : uint(findMSB(value - 1u)) + 1u; +} + +uint dot_accumulator_scale_exponent( + const uint x_exponent, + const uint w_exponent, + const uint terms) { + const uint exponent_sum = x_exponent + w_exponent + ceil_log2(terms); + return exponent_sum > 268u ? exponent_sum - 268u : 0u; +} + +void main() { + const uint row_start = gl_WorkGroupID.y * BM; + if (row_start >= p.rows) { + return; + } + + const uint lane = gl_LocalInvocationID.x; + const uint subgroup = lane / 64u; + const uint subgroup_row = (subgroup / 4u) * 32u; + const uint subgroup_col = (subgroup & 3u) * 32u; + const uint col_start = gl_WorkGroupID.x * BN; + const uint valid_rows = min(BM, p.rows - row_start); + const uint packed_row_words = p.packed_row_bytes >> 2; + + // Per-row BF16 exponents for activations in this tile. + const uint row_pairs = p.K >> 1; + for (uint row = gl_SubgroupID; row < valid_rows; row += gl_NumSubgroups) { + uint max_exponent = 0u; + for (uint pair = gl_SubgroupInvocationID; pair < row_pairs; + pair += gl_SubgroupSize) { + const uint packed = + data_x[((row_start + row) * p.x_row_stride + pair * 2u) >> 1]; + const uint exponent0 = ((packed & 0x7fffu) >> 7u) & 0xffu; + const uint exponent1 = (((packed >> 16u) & 0x7fffu) >> 7u) & 0xffu; + const uint exponent = max(exponent0, exponent1); + if (exponent < 255u) { + max_exponent = max(max_exponent, exponent); + } + } + max_exponent = subgroupMax(max_exponent); + if (subgroupElect()) { + x_scale_exponents[row] = max_exponent; + } + } + + uint weight_exponent = 0u; + if (lane < BN) { + const uint src_col = col_start + lane; + if (src_col < p.cols) { + const uint scale_base = src_col * p.scale_row_stride; + const uint bias_base = src_col * p.bias_row_stride; + for (uint group = 0u; group < p.num_groups; ++group) { + const float scale = + bf16_to_fp32(uint(data_scales[scale_base + group])); + const float bias = + bf16_to_fp32(uint(data_biases[bias_base + group])); + weight_exponent = max(weight_exponent, fp32_biased_exponent(bias)); + weight_exponent = max( + weight_exponent, + fp32_biased_exponent(fma(scale, 15.0f, bias))); + } + const uint scale_exponent = + fp16_range_scale_exponent(weight_exponent); + w_scales[lane] = scale_down( + bf16_to_fp32(uint(data_scales[scale_base])), scale_exponent); + w_biases[lane] = scale_down( + bf16_to_fp32(uint(data_biases[bias_base])), scale_exponent); + } else { + w_scales[lane] = 0.0f; + w_biases[lane] = 0.0f; + } + w_scale_exponents[lane] = weight_exponent; + } + const uint subgroup_w_exponent = subgroupMax(weight_exponent); + if (subgroupElect()) { + tile_w_exponents[gl_SubgroupID] = subgroup_w_exponent; + } + barrier(); + + uint max_w_exponent = 0u; + if (lane == 0u) { + for (uint i = 0u; i < gl_NumSubgroups; ++i) { + max_w_exponent = max(max_w_exponent, tile_w_exponents[i]); + } + tile_w_exponents[0] = max_w_exponent; + } + barrier(); + max_w_exponent = tile_w_exponents[0]; + + if (lane < BM) { + const uint x_exponent = lane < valid_rows ? x_scale_exponents[lane] : 0u; + const uint max_w_scale_exponent = + fp16_range_scale_exponent(max_w_exponent); + const uint dot_scale_exponent = dot_accumulator_scale_exponent( + x_exponent, max_w_exponent, p.K); + const uint required_x_scale_exponent = + dot_scale_exponent > max_w_scale_exponent + ? dot_scale_exponent - max_w_scale_exponent + : 0u; + const uint scale_exponent = max( + fp16_range_scale_exponent(x_exponent), + required_x_scale_exponent); + x_scale_exponents[lane] = scale_exponent; + x_scale_factors[lane] = + uintBitsToFloat((127u - scale_exponent) << 23u); + } + if (lane < BN) { + w_scale_exponents[lane] = + fp16_range_scale_exponent(w_scale_exponents[lane]); + } + barrier(); + + coopmat matrix_a0; + coopmat matrix_a1; + coopmat matrix_b0; + coopmat matrix_b1; + coopmat + accum00 = coopmat(float16_t(0.0f)); + coopmat + accum10 = coopmat(float16_t(0.0f)); + coopmat + accum01 = coopmat(float16_t(0.0f)); + coopmat + accum11 = coopmat(float16_t(0.0f)); + + for (uint kb = 0u; kb < p.K; kb += BK) { + for (uint pair_idx = lane; pair_idx < BM * (BK / 2u); + pair_idx += 512u) { + const uint row = pair_idx / (BK / 2u); + const uint pair = pair_idx - row * (BK / 2u); + const uint k0 = kb + pair * 2u; + const uint src_row = row_start + row; + float16_t x0 = float16_t(0.0f); + float16_t x1 = float16_t(0.0f); + if (row < valid_rows && k0 < p.K) { + const uint packed_x = + data_x[(src_row * p.x_row_stride + k0) >> 1]; + const vec2 values = vec2( + bf16_to_fp32(packed_x & 0xffffu), + bf16_to_fp32(packed_x >> 16u)) * + x_scale_factors[row]; + x0 = float16_t(values.x); + x1 = float16_t(values.y); + } + xs[row * SHMEM_STRIDE + pair] = f16vec2(x0, x1); + } + + // 4-bit: 8 nibbles / uint32; two packs cover BK=16. + for (uint weight_idx = lane; weight_idx < BN * 2u; + weight_idx += 512u) { + const uint col = weight_idx >> 1; + const uint pack_id = weight_idx & 1u; + const uint src_col = col_start + col; + const uint k0 = kb + pack_id * 8u; + f16vec2 w01 = f16vec2(0.0f); + f16vec2 w23 = f16vec2(0.0f); + f16vec2 w45 = f16vec2(0.0f); + f16vec2 w67 = f16vec2(0.0f); + if (src_col < p.cols && k0 < p.K) { + const uint packed_w = + data_w[src_col * packed_row_words + (k0 >> 3)]; + const float scale = w_scales[col]; + const float bias = w_biases[col]; + const float q0 = float(packed_w & 0xFu); + const float q1 = float((packed_w >> 4) & 0xFu); + const float q2 = float((packed_w >> 8) & 0xFu); + const float q3 = float((packed_w >> 12) & 0xFu); + const float q4 = float((packed_w >> 16) & 0xFu); + const float q5 = float((packed_w >> 20) & 0xFu); + const float q6 = float((packed_w >> 24) & 0xFu); + const float q7 = float((packed_w >> 28) & 0xFu); + w01 = f16vec2(fma(scale, q0, bias), fma(scale, q1, bias)); + w23 = f16vec2(fma(scale, q2, bias), fma(scale, q3, bias)); + w45 = f16vec2(fma(scale, q4, bias), fma(scale, q5, bias)); + w67 = f16vec2(fma(scale, q6, bias), fma(scale, q7, bias)); + } + const uint base = col * SHMEM_STRIDE + pack_id * 4u; + ws[base + 0u] = w01; + ws[base + 1u] = w23; + ws[base + 2u] = w45; + ws[base + 3u] = w67; + } + barrier(); + + coopMatLoad( + matrix_a0, + xs, + subgroup_row * SHMEM_STRIDE, + SHMEM_STRIDE, + gl_CooperativeMatrixLayoutRowMajor); + coopMatLoad( + matrix_a1, + xs, + (subgroup_row + 16u) * SHMEM_STRIDE, + SHMEM_STRIDE, + gl_CooperativeMatrixLayoutRowMajor); + coopMatLoad( + matrix_b0, + ws, + subgroup_col * SHMEM_STRIDE, + SHMEM_STRIDE, + gl_CooperativeMatrixLayoutColumnMajor); + coopMatLoad( + matrix_b1, + ws, + (subgroup_col + 16u) * SHMEM_STRIDE, + SHMEM_STRIDE, + gl_CooperativeMatrixLayoutColumnMajor); + accum00 = coopMatMulAdd(matrix_a0, matrix_b0, accum00); + accum01 = coopMatMulAdd(matrix_a0, matrix_b1, accum01); + accum10 = coopMatMulAdd(matrix_a1, matrix_b0, accum10); + accum11 = coopMatMulAdd(matrix_a1, matrix_b1, accum11); + + const uint next_kb = kb + BK; + if (lane < BN && next_kb < p.K && (next_kb % p.group_size) == 0u) { + const uint src_col = col_start + lane; + if (src_col < p.cols) { + const uint group = next_kb / p.group_size; + const uint scale_base = src_col * p.scale_row_stride; + const uint bias_base = src_col * p.bias_row_stride; + w_scales[lane] = scale_down( + bf16_to_fp32(uint(data_scales[scale_base + group])), + w_scale_exponents[lane]); + w_biases[lane] = scale_down( + bf16_to_fp32(uint(data_biases[bias_base + group])), + w_scale_exponents[lane]); + } + } + barrier(); + } + + coopMatStore( + accum00, + output_stage, + subgroup_row * BN + subgroup_col, + BN, + gl_CooperativeMatrixLayoutRowMajor); + coopMatStore( + accum01, + output_stage, + subgroup_row * BN + subgroup_col + 16u, + BN, + gl_CooperativeMatrixLayoutRowMajor); + coopMatStore( + accum10, + output_stage, + (subgroup_row + 16u) * BN + subgroup_col, + BN, + gl_CooperativeMatrixLayoutRowMajor); + coopMatStore( + accum11, + output_stage, + (subgroup_row + 16u) * BN + subgroup_col + 16u, + BN, + gl_CooperativeMatrixLayoutRowMajor); + barrier(); + + for (uint index = lane; index < BM * BN; index += 512u) { + const uint row = index / BN; + const uint col = index - row * BN; + if (row < valid_rows && col_start + col < p.cols) { + float value = scale_up( + float(output_stage[index]), x_scale_exponents[row]); + value = scale_up(value, w_scale_exponents[col]); + data_out[(row_start + row) * p.out_row_stride + col_start + col] = + uint16_t(fp32_to_bf16(value)); + } + } +} diff --git a/mlx/backend/vulkan/kernels/vulkan-shaders-gen.cpp b/mlx/backend/vulkan/kernels/vulkan-shaders-gen.cpp index 1b094f7971..e8f5915baa 100644 --- a/mlx/backend/vulkan/kernels/vulkan-shaders-gen.cpp +++ b/mlx/backend/vulkan/kernels/vulkan-shaders-gen.cpp @@ -2348,6 +2348,12 @@ void process_shaders() { {}, true, true); + string_to_spv( + "fused_affine_qmm_bf16_bf16_coop4", + "mul_mm_affine_bf16_coop4.comp", + {}, + true, + true); #endif string_to_spv( diff --git a/mlx/backend/vulkan/quantized.cpp b/mlx/backend/vulkan/quantized.cpp index 6ff598b638..fb490e8fe2 100644 --- a/mlx/backend/vulkan/quantized.cpp +++ b/mlx/backend/vulkan/quantized.cpp @@ -373,6 +373,17 @@ bool dequantized_bf16_prefill_enabled() { return enabled; } +bool fused_affine_coop4_prefill_enabled() { + static const bool enabled = []() { + if (const char* env = std::getenv("MLX_VULKAN_AFFINE_COOP4_PREFILL"); + env != nullptr) { + return std::string_view(env) != "0"; + } + return true; + }(); + return enabled; +} + bool is_row_contiguous_zero_offset(const array& arr) { if (arr.ndim() == 0) { return arr.offset() == 0; @@ -1849,8 +1860,82 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { static_cast(w.shape(-1) * 32 / bits_) == k && num_groups == static_cast((k + group_size_ - 1) / group_size_)) { - const bool use_dequantized_prefill = rows >= 256 && !decode_lhs && - dequantized_bf16_prefill_enabled(); + const bool use_large_prefill = rows >= 256 && !decode_lhs; +#if defined(MLX_VULKAN_COOPMAT_GLSLC_SUPPORT) + const auto& vk_ctx = vulkan::VulkanContext::get(); + const auto device_limits = + vk_ctx.physical_device().getProperties().limits; + const bool supports_coop_wg = + device_limits.maxComputeWorkGroupInvocations >= 512u && + device_limits.maxComputeWorkGroupSize[0] >= 512u && + device_limits.maxComputeSharedMemorySize >= 27648u; + const bool supports_64_lane = vk_ctx.subgroup_size() == 64u || + (vk_ctx.subgroup_size_control_supported() && + vk_ctx.subgroup_min_size() <= 64u && + vk_ctx.subgroup_max_size() >= 64u); + const bool use_coop4_prefill = use_large_prefill && bits_ == 4 && + group_size_ == 64 && (k % 16u) == 0u && (k % 2u) == 0u && + vk_ctx.coopmat_f16acc_supported() && supports_64_lane && + supports_coop_wg && fused_affine_coop4_prefill_enabled(); + if (use_coop4_prefill) { + array out_work( + (vector_lhs || flatten_lhs_batches) + ? Shape{static_cast(rows), out.shape(-1)} + : out.shape(), + bfloat16, + nullptr, + {}); + out_work.set_data(allocator::malloc(out_work.nbytes())); + if (out_work.size() != 0) { + vulkan::FusedAffineMatmulPushConstants push_constants{}; + push_constants.rows = rows; + push_constants.cols = cols; + push_constants.K = k; + push_constants.packed_row_bytes = + static_cast(w.strides(-2) * sizeof(uint32_t)); + push_constants.x_row_stride = + static_cast(x_mat.strides(-2)); + push_constants.out_row_stride = + static_cast(out_work.strides(-2)); + push_constants.scale_row_stride = + static_cast(scales_bf16.strides(-2)); + push_constants.bias_row_stride = + static_cast(biases_bf16.strides(-2)); + push_constants.bits = static_cast(bits_); + push_constants.group_size = static_cast(group_size_); + push_constants.num_groups = num_groups; + + const std::array grid = { + (cols + 127u) / 128u, (rows + 63u) / 64u, 1u}; + if (!dispatch_grid_within_limits(grid[0], grid[1], grid[2])) { + throw std::runtime_error( + "[QuantizedMatmul::eval_gpu] Cooperative affine-4 dispatch grid exceeds Vulkan limits."); + } + + auto command_buffer = vulkan::begin_command_recording(s.index); + vulkan::dispatch_fused_affine_matmul_op( + w, + scales_bf16, + biases_bf16, + x_mat, + out_work, + vulkan::StaticShaderId::fused_affine_qmm_bf16_bf16_coop4_cm1, + command_buffer, + s, + push_constants, + grid); + vulkan::end_command_recording(s.index); + } + finalize_bf16_output(out_work); + trace_qmm( + "fused_coop4_prefill", + (vector_lhs || flatten_lhs_batches) ? "reshaped_output=1" + : "reshaped_output=0"); + return; + } +#endif + const bool use_dequantized_prefill = + use_large_prefill && dequantized_bf16_prefill_enabled(); if (use_dequantized_prefill) { array w_deq( expanded_quantized_shape(w, bits_), bfloat16, nullptr, {}); From 0a2164a741323078490078b2cdd4b5a5d0d7bce0 Mon Sep 17 00:00:00 2001 From: Goni Zahavy Date: Wed, 15 Jul 2026 22:14:27 +0300 Subject: [PATCH 5/9] [Vulkan] Prefer F16 coopmat1 for large BF16 GEMMs Native BF16 coopmat shaders are incorrect on current AMD drivers. Promote large BF16xBF16 matmuls to F16 coopmat1 (direct BF16 out) with a Strix Halo-tuned wave64 warptile, keeping BF16 weight transpose cache. --- mlx/backend/vulkan/matmul.cpp | 221 ++++++++++++++++++++++++++++------ 1 file changed, 187 insertions(+), 34 deletions(-) diff --git a/mlx/backend/vulkan/matmul.cpp b/mlx/backend/vulkan/matmul.cpp index 5a662c9ac9..0b5d8369bd 100644 --- a/mlx/backend/vulkan/matmul.cpp +++ b/mlx/backend/vulkan/matmul.cpp @@ -343,12 +343,64 @@ constexpr std::array kSafeMatmulSpec = constexpr std::array kLane64MatmulSpec = {64, 32, 32, 16, 32, 32, 2, 2, 2, 1, 64}; +// Cooperative-matrix warptiles (TM/TN/TK must match 16x16x16 subgroup mats). +// On Strix Halo / RDNA3.5 the llama.cpp "large" 128x128/256-thread tile is +// ~10x slower than this medium tile for dense F16 GEMMs; keep one known-good +// wave64 tile across families. +constexpr std::array kSafeCoopmatMatmulSpec = + {64, 32, 32, 16, 32, 32, 2, 16, 16, 16, 32}; +constexpr std::array kLane64CoopmatMatmulSpec = + {128, 64, 64, 16, 64, 32, 2, 16, 16, 16, 64}; + bool supports_64_lane_matmul(const vulkan::VulkanContext& ctx) { return ctx.subgroup_size() >= 64u || (ctx.subgroup_size_control_supported() && ctx.subgroup_min_size() <= 64u && ctx.subgroup_max_size() >= 64u); } +bool matmul_coopmat_env_enabled() { + if (const char* env = std::getenv("MLX_VULKAN_MATMUL_COOPMAT"); + env != nullptr && env[0] != '\0') { + return !(env[0] == '0' && env[1] == '\0'); + } + return true; +} + +bool prefer_matmul_coopmat1( + Dtype dtype, + uint32_t m, + uint32_t n, + uint32_t k) { + const auto& ctx = vulkan::VulkanContext::get(); + if (!matmul_coopmat_env_enabled() || !ctx.cooperative_matrix_supported()) { + return false; + } + // Native BF16 coopmat shaders are incorrect on current AMD drivers (cosine + // ~0.72). Use F16 coopmat instead (with BF16→F16 promote when needed). + if (dtype != float16 || !ctx.shader_float16_supported()) { + return false; + } + // Coopmat pays off on medium/large GEMMs; tiny shapes stay on scalar. + return m >= 64u && n >= 64u && k >= 64u; +} + +bool prefer_bf16_to_f16_coopmat_promote(uint32_t m, uint32_t n, uint32_t k) { + const auto& ctx = vulkan::VulkanContext::get(); + if (!ctx.shader_bfloat16_supported()) { + return false; + } + return prefer_matmul_coopmat1(float16, m, n, k); +} + +std::array coopmat_matmul_spec_for( + MatmulFamily /*family*/, + const vulkan::VulkanContext& ctx) { + if (supports_64_lane_matmul(ctx)) { + return kLane64CoopmatMatmulSpec; + } + return kSafeCoopmatMatmulSpec; +} + std::vector mul_mm_shader_candidates( Dtype dtype, bool prefer_fp32_accum) { @@ -389,37 +441,97 @@ std::vector mul_mm_shader_candidates( std::vector mul_mm_direct_shader_candidates( Dtype input_dtype, Dtype output_dtype, - bool prefer_fp32_accum) { + bool prefer_fp32_accum, + bool aligned, + bool prefer_coopmat1) { + const auto& ctx = vulkan::VulkanContext::get(); + const bool use_f16acc_cm1 = + prefer_coopmat1 && ctx.coopmat_f16acc_supported() && !prefer_fp32_accum; + + auto append_unique = [](std::vector& out, + vulkan::StaticShaderId id) { + for (auto existing : out) { + if (existing == id) { + return; + } + } + out.push_back(id); + }; + + auto with_coopmat_and_scalar = + [&](vulkan::StaticShaderId scalar, + vulkan::StaticShaderId scalar_f16acc, + vulkan::StaticShaderId aligned_scalar, + vulkan::StaticShaderId aligned_scalar_f16acc, + vulkan::StaticShaderId cm1, + vulkan::StaticShaderId cm1_f16acc, + vulkan::StaticShaderId aligned_cm1, + vulkan::StaticShaderId aligned_cm1_f16acc) { + std::vector out; + if (prefer_coopmat1) { + // Prefer fp32-accum coopmat for BF16 (matches device BF16 coopmat + // props). Prefer f16acc coopmat when requested and supported. + if (aligned) { + if (use_f16acc_cm1) { + append_unique(out, aligned_cm1_f16acc); + } + append_unique(out, aligned_cm1); + } + if (use_f16acc_cm1) { + append_unique(out, cm1_f16acc); + } + append_unique(out, cm1); + } + + if (prefer_fp32_accum) { + if (aligned) { + append_unique(out, aligned_scalar); + append_unique(out, aligned_scalar_f16acc); + } + append_unique(out, scalar); + append_unique(out, scalar_f16acc); + } else { + if (aligned) { + append_unique(out, aligned_scalar_f16acc); + append_unique(out, aligned_scalar); + } + append_unique(out, scalar_f16acc); + append_unique(out, scalar); + } + return out; + }; + switch (input_dtype) { case float16: switch (output_dtype) { case float16: - return prefer_fp32_accum - ? std::vector{ - vulkan::StaticShaderId::matmul_direct_f16, - vulkan::StaticShaderId::matmul_direct_f16_f16acc, - } - : std::vector{ - vulkan::StaticShaderId::matmul_direct_f16_f16acc, - vulkan::StaticShaderId::matmul_direct_f16, - }; + return with_coopmat_and_scalar( + vulkan::StaticShaderId::matmul_direct_f16, + vulkan::StaticShaderId::matmul_direct_f16_f16acc, + vulkan::StaticShaderId::matmul_direct_f16_aligned, + vulkan::StaticShaderId::matmul_direct_f16_aligned_f16acc, + vulkan::StaticShaderId::matmul_direct_f16_cm1, + vulkan::StaticShaderId::matmul_direct_f16_f16acc_cm1, + vulkan::StaticShaderId::matmul_direct_f16_aligned_cm1, + vulkan::StaticShaderId::matmul_direct_f16_aligned_f16acc_cm1); case bfloat16: - return prefer_fp32_accum - ? std::vector{ - vulkan::StaticShaderId::matmul_direct_f16_bf16, - vulkan::StaticShaderId::matmul_direct_f16_bf16_f16acc, - } - : std::vector{ - vulkan::StaticShaderId::matmul_direct_f16_bf16_f16acc, - vulkan::StaticShaderId::matmul_direct_f16_bf16, - }; + return with_coopmat_and_scalar( + vulkan::StaticShaderId::matmul_direct_f16_bf16, + vulkan::StaticShaderId::matmul_direct_f16_bf16_f16acc, + vulkan::StaticShaderId::matmul_direct_f16_bf16_aligned, + vulkan::StaticShaderId::matmul_direct_f16_bf16_aligned_f16acc, + vulkan::StaticShaderId::matmul_direct_f16_bf16_cm1, + vulkan::StaticShaderId::matmul_direct_f16_bf16_f16acc_cm1, + vulkan::StaticShaderId::matmul_direct_f16_bf16_aligned_cm1, + vulkan::StaticShaderId:: + matmul_direct_f16_bf16_aligned_f16acc_cm1); case float32: return mul_mm_shader_candidates(input_dtype, prefer_fp32_accum); default: return {}; } case bfloat16: - if (!vulkan::VulkanContext::get().shader_bfloat16_supported()) { + if (!ctx.shader_bfloat16_supported()) { return {}; } if (output_dtype != bfloat16) { @@ -427,15 +539,15 @@ std::vector mul_mm_direct_shader_candidates( ? mul_mm_shader_candidates(input_dtype, prefer_fp32_accum) : std::vector{}; } - return prefer_fp32_accum - ? std::vector{ - vulkan::StaticShaderId::matmul_direct_bf16, - vulkan::StaticShaderId::matmul_direct_bf16_f16acc, - } - : std::vector{ - vulkan::StaticShaderId::matmul_direct_bf16_f16acc, - vulkan::StaticShaderId::matmul_direct_bf16, - }; + return with_coopmat_and_scalar( + vulkan::StaticShaderId::matmul_direct_bf16, + vulkan::StaticShaderId::matmul_direct_bf16_f16acc, + vulkan::StaticShaderId::matmul_direct_bf16_aligned, + vulkan::StaticShaderId::matmul_direct_bf16_aligned_f16acc, + vulkan::StaticShaderId::matmul_direct_bf16_cm1, + vulkan::StaticShaderId::matmul_direct_bf16_f16acc_cm1, + vulkan::StaticShaderId::matmul_direct_bf16_aligned_cm1, + vulkan::StaticShaderId::matmul_direct_bf16_aligned_f16acc_cm1); case float32: return output_dtype == float32 ? mul_mm_shader_candidates(input_dtype, prefer_fp32_accum) @@ -612,7 +724,24 @@ select_matmul_dispatch_tuning(Dtype dtype, uint32_t m, uint32_t n, uint32_t k) { } } - if (ctx.architecture() == vulkan::GpuArchitecture::AmdRdna && + if (prefer_matmul_coopmat1(dtype, m, n, k)) { + // Keep split-K off for coopmat1 direct shaders. Prefer f16acc when the + // device advertises 16x16x16 f16-accum coopmat support. + tuning.split_k_threshold = std::max(tuning.split_k_threshold, k + 1u); + tuning.prefer_fp32_accum = !ctx.coopmat_f16acc_supported(); + const auto coop_spec = coopmat_matmul_spec_for(family, ctx); + tuning.specialization_constants.assign(coop_spec.begin(), coop_spec.end()); + if (tuning.specialization_constants.size() > 10 && + ctx.subgroup_size_control_supported()) { + uint32_t preferred = std::clamp( + supports_64_lane_matmul(ctx) ? 64u : profile.preferred_subgroup_size, + std::max(ctx.subgroup_min_size(), 1u), + std::max(ctx.subgroup_max_size(), 1u)); + preferred = std::min(preferred, tuning.specialization_constants[0]); + tuning.specialization_constants[10] = preferred; + } + } else if ( + ctx.architecture() == vulkan::GpuArchitecture::AmdRdna && dtype == bfloat16) { tuning.prefer_fp32_accum = true; } @@ -1224,16 +1353,24 @@ bool try_eval_mul_mm_vulkan( } bool b_uses_cast_scratch = false; + bool promote_bf16_to_f16_coopmat = false; if (a.dtype() == bfloat16 && !vulkan::VulkanContext::get().shader_bfloat16_supported()) { a = cast_to_float16_scratch(a, s, kMulMmACastScratchLane); b = cast_to_float16_scratch(b, s, kMulMmBCastScratchLane); b_uses_cast_scratch = true; + } else if ( + a.dtype() == bfloat16 && out.dtype() == bfloat16 && + prefer_bf16_to_f16_coopmat_promote( + static_cast(out.shape(-2)), + static_cast(out.shape(-1)), + static_cast(a.shape(-1)))) { + // Native BF16 coopmat is unreliable on current AMD drivers. Transpose the + // original BF16 weights (cacheable), then promote A/B^T to F16 and use + // matmul_direct_f16_bf16_*_cm1 (writes BF16 out directly). + promote_bf16_to_f16_coopmat = true; } - // Keep BF16 inputs in BF16 and dispatch matmul_bf16* directly. - // This matches ggml's BF16xBF16 path and avoids costly staging casts. - if (!is_row_contiguous_zero_offset(a)) { if (!ensure_vulkan_buffer(a, s)) { return false; @@ -1258,6 +1395,14 @@ bool try_eval_mul_mm_vulkan( return false; } + if (promote_bf16_to_f16_coopmat) { + a = cast_to_float16_scratch(a, s, kMulMmACastScratchLane); + b_t = cast_to_float16_scratch(b_t, s, kMulMmBCastScratchLane); + if (matmul_debug_enabled()) { + std::cerr << "[vulkan::mul_mm] promote bf16->f16 for coopmat1\n"; + } + } + if (!ensure_vulkan_buffer(a, s) || !ensure_vulkan_buffer(b_t, s)) { if (matmul_debug_enabled()) { std::cerr << "[vulkan::mul_mm] missing buffer" @@ -1496,8 +1641,16 @@ bool try_eval_mul_mm_vulkan( const bool can_direct_write = split_k == 1u && is_row_contiguous_zero_offset(out); if (can_direct_write) { + const bool use_coopmat1 = + prefer_matmul_coopmat1(a.dtype(), m, n, k) && + (out.dtype() == a.dtype() || + (a.dtype() == float16 && out.dtype() == bfloat16)); auto direct_candidates = mul_mm_direct_shader_candidates( - a.dtype(), out.dtype(), tuning.prefer_fp32_accum); + a.dtype(), + out.dtype(), + tuning.prefer_fp32_accum, + tuning.aligned, + use_coopmat1); array out_direct = out; out_direct.set_data(allocator::malloc(out_direct.nbytes())); if (ensure_vulkan_buffer(out_direct, s) && From cbb55b22776054ab7c5aed42ac2c7c67bf6400e6 Mon Sep 17 00:00:00 2001 From: Goni Zahavy Date: Thu, 16 Jul 2026 05:24:35 +0300 Subject: [PATCH 6/9] [Vulkan] Prefetch coop4 exponents; tune large RDNA matvec Dense affine-4 coopmat prefill rescanned activation/weight exponents once per column tile. Prefetch them once per QMM (like MoE gather metadata). Also use NUM_ROWS=2 for large AMD RDNA BF16 matvecs (N>=8192) to raise GDN decode bandwidth without regressing smaller shapes. --- mlx/backend/vulkan/kernels.cpp | 153 ++++++++++++++++-- mlx/backend/vulkan/kernels.h | 24 +++ .../vulkan/kernels/affine_bf16_exponents.comp | 106 ++++++++++++ .../kernels/mul_mm_affine_bf16_coop4.comp | 41 ++--- .../vulkan/kernels/vulkan-shaders-gen.cpp | 4 + mlx/backend/vulkan/quantized.cpp | 8 +- 6 files changed, 291 insertions(+), 45 deletions(-) create mode 100644 mlx/backend/vulkan/kernels/affine_bf16_exponents.comp diff --git a/mlx/backend/vulkan/kernels.cpp b/mlx/backend/vulkan/kernels.cpp index 9a51de86ae..d95326b740 100644 --- a/mlx/backend/vulkan/kernels.cpp +++ b/mlx/backend/vulkan/kernels.cpp @@ -20,16 +20,20 @@ namespace mlx::core::vulkan { constexpr uint32_t kMaxMulMatVecCols = 8; -uint32_t matvec_rows_per_workgroup() { - static const uint32_t value = []() { - if (const char* env = std::getenv("MLX_VULKAN_MATVEC_ROWS_PER_WG"); - env != nullptr) { - return std::clamp( - static_cast(std::strtoul(env, nullptr, 10)), 1u, 8u); - } - return 1u; - }(); - return value; +uint32_t matvec_rows_per_workgroup(uint32_t nrows) { + if (const char* env = std::getenv("MLX_VULKAN_MATVEC_ROWS_PER_WG"); + env != nullptr) { + return std::clamp( + static_cast(std::strtoul(env, nullptr, 10)), 1u, 8u); + } + // On RDNA, NUM_ROWS=2 reuses the activation vector across two output rows and + // lifts large GDN-style matvecs (N>=8K) from ~140GB/s toward ~180GB/s. Smaller + // N regresses, so keep the historical default of 1 elsewhere. + if (VulkanContext::get().architecture() == GpuArchitecture::AmdRdna && + nrows >= 8192u) { + return 2u; + } + return 1u; } uint32_t max_compute_work_group_invocations() { @@ -240,7 +244,8 @@ PipelineCreationOptions pipeline_creation_options( const std::vector& specialization_constants) { PipelineCreationOptions options; - if (shader_name == "gather_affine_qmm_rhs_bf16_bf16_cm1") { + if (shader_name == "gather_affine_qmm_rhs_bf16_bf16_cm1" || + shader_name == "fused_affine_qmm_bf16_bf16_coop4_cm1") { options.require_full_subgroups = true; options.required_subgroup_size = 64; return options; @@ -435,6 +440,8 @@ enum class KernelSpecId { Nvfp4Dequant, Nvfp4Quant, FusedAffineMatmul, + FusedAffineCoop4Matmul, + AffineBf16Exponents, GatherAffineMatmul, GatherAffineTileMetadata, GatherAffineCoopMatmul, @@ -469,7 +476,7 @@ KernelSpec make_kernel_spec( grid_kind}; } -const std::array kKernelSpecs = { +const std::array kKernelSpecs = { make_kernel_spec( {0, 1, 2}, sizeof(BinaryPushConstants), @@ -602,6 +609,14 @@ const std::array kKernelSpecs = { {0, 1, 2, 3, 4}, sizeof(FusedAffineMatmulPushConstants), DispatchGridKind::Linear1D), + make_kernel_spec( + {0, 1, 2, 3, 4, 5}, + sizeof(FusedAffineMatmulPushConstants), + DispatchGridKind::Linear1D), + make_kernel_spec( + {0, 1, 2, 3}, + sizeof(AffineBf16ExponentsPushConstants), + DispatchGridKind::Linear1D), make_kernel_spec( {0, 1, 2, 3, 4, 5, 6}, sizeof(GatherAffineMatmulPushConstants), @@ -3508,7 +3523,7 @@ void dispatch_mul_mat_vec_op( }}; constexpr uint32_t kMaxWorkgroupsX = 65535u; - const uint32_t rows_per_workgroup = matvec_rows_per_workgroup(); + const uint32_t rows_per_workgroup = matvec_rows_per_workgroup(nrows); const uint32_t row_groups = (nrows + rows_per_workgroup - 1u) / rows_per_workgroup; const uint32_t groups_z = @@ -4135,6 +4150,118 @@ void dispatch_fused_affine_matmul_op( matmul_specialization_constants({})); } +void dispatch_fused_affine_coop4_matmul_op( + const array& w, + const array& scales, + const array& biases, + const array& x, + array& exponents, + array& out, + StaticShaderId shader_id, + vk::CommandBuffer cmd_buffer, + const Stream& s, + const FusedAffineMatmulPushConstants& push_constants, + const std::array& grid) { + if (w.ndim() != 2 || scales.ndim() != 2 || biases.ndim() != 2 || + x.ndim() != 2 || out.ndim() != 2 || exponents.ndim() != 1) { + throw std::runtime_error( + "[vulkan::kernels] fused_affine_coop4 dispatch requires 2D tensors and 1D exponents."); + } + + const uint32_t rows = push_constants.rows; + const uint32_t cols = push_constants.cols; + if (checked_u32(exponents.shape(0), "fused_affine_coop4 exponents") != + rows + cols) { + throw std::runtime_error( + "[vulkan::kernels] fused_affine_coop4 exponents length must be rows+cols."); + } + + AffineBf16ExponentsPushConstants exp_pc{}; + exp_pc.rows = rows; + exp_pc.cols = cols; + exp_pc.K = push_constants.K; + exp_pc.x_row_stride = push_constants.x_row_stride; + exp_pc.scale_row_stride = push_constants.scale_row_stride; + exp_pc.bias_row_stride = push_constants.bias_row_stride; + exp_pc.num_groups = push_constants.num_groups; + + const std::array exp_arrays = {{ + {&x, "X"}, + {&scales, "SCALES"}, + {&biases, "BIASES"}, + {&exponents, "EXPONENTS"}, + }}; + + exp_pc.mode = 0u; + dispatch_with_spec( + StaticShaderId::affine_bf16_exponents, + KernelSpecId::AffineBf16Exponents, + exp_arrays, + exp_pc, + rows, + cmd_buffer, + s, + std::array{rows, 1u, 1u}); + + VkMemoryBarrier barrier{}; + barrier.sType = VK_STRUCTURE_TYPE_MEMORY_BARRIER; + barrier.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT; + barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT; + vkCmdPipelineBarrier( + cmd_buffer, + VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, + VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, + 0, + 1, + &barrier, + 0, + nullptr, + 0, + nullptr); + + exp_pc.mode = 1u; + dispatch_with_spec( + StaticShaderId::affine_bf16_exponents, + KernelSpecId::AffineBf16Exponents, + exp_arrays, + exp_pc, + cols, + cmd_buffer, + s, + std::array{cols, 1u, 1u}); + + barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT; + vkCmdPipelineBarrier( + cmd_buffer, + VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, + VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, + 0, + 1, + &barrier, + 0, + nullptr, + 0, + nullptr); + + const std::array matmul_arrays = {{ + {&w, "W"}, + {&scales, "SCALES"}, + {&biases, "BIASES"}, + {&x, "X"}, + {&exponents, "EXPONENTS"}, + {&out, "OUT"}, + }}; + dispatch_with_spec( + shader_id, + KernelSpecId::FusedAffineCoop4Matmul, + matmul_arrays, + push_constants, + checked_mul_u32(rows, cols, "fused_affine_coop4 output elements"), + cmd_buffer, + s, + grid); +} + void dispatch_gather_affine_matmul_op( const array& w, const array& scales, diff --git a/mlx/backend/vulkan/kernels.h b/mlx/backend/vulkan/kernels.h index 61ff9447c9..e6a9387e71 100644 --- a/mlx/backend/vulkan/kernels.h +++ b/mlx/backend/vulkan/kernels.h @@ -633,6 +633,17 @@ struct FusedAffineMatmulPushConstants { uint32_t num_groups; }; +struct AffineBf16ExponentsPushConstants { + uint32_t rows; + uint32_t cols; + uint32_t K; + uint32_t x_row_stride; + uint32_t scale_row_stride; + uint32_t bias_row_stride; + uint32_t num_groups; + uint32_t mode; +}; + struct GatherAffineMatmulPushConstants { uint32_t rows; uint32_t cols; @@ -1160,6 +1171,19 @@ void dispatch_fused_affine_matmul_op( const FusedAffineMatmulPushConstants& push_constants, const std::array& grid); +void dispatch_fused_affine_coop4_matmul_op( + const array& w, + const array& scales, + const array& biases, + const array& x, + array& exponents, + array& out, + StaticShaderId shader_id, + vk::CommandBuffer cmd_buffer, + const Stream& s, + const FusedAffineMatmulPushConstants& push_constants, + const std::array& grid); + void dispatch_gather_affine_matmul_op( const array& w, const array& scales, diff --git a/mlx/backend/vulkan/kernels/affine_bf16_exponents.comp b/mlx/backend/vulkan/kernels/affine_bf16_exponents.comp new file mode 100644 index 0000000000..c4c6897989 --- /dev/null +++ b/mlx/backend/vulkan/kernels/affine_bf16_exponents.comp @@ -0,0 +1,106 @@ +#version 450 + +#extension GL_EXT_shader_16bit_storage : require +#extension GL_KHR_shader_subgroup_arithmetic : enable +#extension GL_KHR_shader_subgroup_basic : enable + +#include "types.glsl" + +// Prefetch BF16 activation-row / weight-column exponents for dense coop4 QMM. +// Mode 0: one workgroup per activation row -> exponents[row] +// Mode 1: one workgroup per weight column -> exponents[rows + col] + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer X { + uint data_x[]; +}; +layout(binding = 1) readonly buffer SCALES { + uint16_t data_scales[]; +}; +layout(binding = 2) readonly buffer BIASES { + uint16_t data_biases[]; +}; +layout(binding = 3) writeonly buffer EXPONENTS { + uint data_exponents[]; +}; + +layout(push_constant) uniform parameter { + uint rows; + uint cols; + uint K; + uint x_row_stride; + uint scale_row_stride; + uint bias_row_stride; + uint num_groups; + uint mode; +} p; + +shared uint subgroup_max_exponents[8]; + +uint fp32_biased_exponent(const float value) { + const uint exponent = (floatBitsToUint(abs(value)) >> 23u) & 0xffu; + return exponent < 255u ? exponent : 0u; +} + +void main() { + const uint lane = gl_LocalInvocationID.x; + const uint idx = gl_WorkGroupID.x; + + if (p.mode == 0u) { + if (idx >= p.rows) { + return; + } + const uint row_pairs = p.K >> 1; + uint max_exponent = 0u; + for (uint pair = lane; pair < row_pairs; pair += 256u) { + const uint packed = + data_x[(idx * p.x_row_stride + pair * 2u) >> 1]; + const uint exponent0 = ((packed & 0x7fffu) >> 7u) & 0xffu; + const uint exponent1 = (((packed >> 16u) & 0x7fffu) >> 7u) & 0xffu; + const uint exponent = max(exponent0, exponent1); + if (exponent < 255u) { + max_exponent = max(max_exponent, exponent); + } + } + max_exponent = subgroupMax(max_exponent); + if (subgroupElect()) { + subgroup_max_exponents[gl_SubgroupID] = max_exponent; + } + barrier(); + if (lane == 0u) { + uint reduced = 0u; + for (uint i = 0u; i < gl_NumSubgroups; ++i) { + reduced = max(reduced, subgroup_max_exponents[i]); + } + data_exponents[idx] = reduced; + } + return; + } + + if (idx >= p.cols) { + return; + } + uint weight_exponent = 0u; + const uint scale_base = idx * p.scale_row_stride; + const uint bias_base = idx * p.bias_row_stride; + for (uint group = lane; group < p.num_groups; group += 256u) { + const float scale = bf16_to_fp32(uint(data_scales[scale_base + group])); + const float bias = bf16_to_fp32(uint(data_biases[bias_base + group])); + weight_exponent = max(weight_exponent, fp32_biased_exponent(bias)); + weight_exponent = + max(weight_exponent, fp32_biased_exponent(fma(scale, 15.0f, bias))); + } + weight_exponent = subgroupMax(weight_exponent); + if (subgroupElect()) { + subgroup_max_exponents[gl_SubgroupID] = weight_exponent; + } + barrier(); + if (lane == 0u) { + uint reduced = 0u; + for (uint i = 0u; i < gl_NumSubgroups; ++i) { + reduced = max(reduced, subgroup_max_exponents[i]); + } + data_exponents[p.rows + idx] = reduced; + } +} diff --git a/mlx/backend/vulkan/kernels/mul_mm_affine_bf16_coop4.comp b/mlx/backend/vulkan/kernels/mul_mm_affine_bf16_coop4.comp index 732e1e70f8..90fd210922 100644 --- a/mlx/backend/vulkan/kernels/mul_mm_affine_bf16_coop4.comp +++ b/mlx/backend/vulkan/kernels/mul_mm_affine_bf16_coop4.comp @@ -25,7 +25,12 @@ layout(binding = 2) readonly buffer BIASES { layout(binding = 3) readonly buffer X { uint data_x[]; }; -layout(binding = 4) writeonly buffer OUT { +layout(binding = 4) readonly buffer EXPONENTS { + // [0, rows) activation row exponents; [rows, rows+cols) weight column + // exponents. Precomputed once per QMM (see affine_bf16_exponents.comp). + uint data_exponents[]; +}; +layout(binding = 5) writeonly buffer OUT { uint16_t data_out[]; }; @@ -101,43 +106,17 @@ void main() { const uint valid_rows = min(BM, p.rows - row_start); const uint packed_row_words = p.packed_row_bytes >> 2; - // Per-row BF16 exponents for activations in this tile. - const uint row_pairs = p.K >> 1; - for (uint row = gl_SubgroupID; row < valid_rows; row += gl_NumSubgroups) { - uint max_exponent = 0u; - for (uint pair = gl_SubgroupInvocationID; pair < row_pairs; - pair += gl_SubgroupSize) { - const uint packed = - data_x[((row_start + row) * p.x_row_stride + pair * 2u) >> 1]; - const uint exponent0 = ((packed & 0x7fffu) >> 7u) & 0xffu; - const uint exponent1 = (((packed >> 16u) & 0x7fffu) >> 7u) & 0xffu; - const uint exponent = max(exponent0, exponent1); - if (exponent < 255u) { - max_exponent = max(max_exponent, exponent); - } - } - max_exponent = subgroupMax(max_exponent); - if (subgroupElect()) { - x_scale_exponents[row] = max_exponent; - } + // Load precomputed activation / weight exponents for this tile. + if (lane < valid_rows) { + x_scale_exponents[lane] = data_exponents[row_start + lane]; } - uint weight_exponent = 0u; if (lane < BN) { const uint src_col = col_start + lane; if (src_col < p.cols) { + weight_exponent = data_exponents[p.rows + src_col]; const uint scale_base = src_col * p.scale_row_stride; const uint bias_base = src_col * p.bias_row_stride; - for (uint group = 0u; group < p.num_groups; ++group) { - const float scale = - bf16_to_fp32(uint(data_scales[scale_base + group])); - const float bias = - bf16_to_fp32(uint(data_biases[bias_base + group])); - weight_exponent = max(weight_exponent, fp32_biased_exponent(bias)); - weight_exponent = max( - weight_exponent, - fp32_biased_exponent(fma(scale, 15.0f, bias))); - } const uint scale_exponent = fp16_range_scale_exponent(weight_exponent); w_scales[lane] = scale_down( diff --git a/mlx/backend/vulkan/kernels/vulkan-shaders-gen.cpp b/mlx/backend/vulkan/kernels/vulkan-shaders-gen.cpp index e8f5915baa..0c7b0c9652 100644 --- a/mlx/backend/vulkan/kernels/vulkan-shaders-gen.cpp +++ b/mlx/backend/vulkan/kernels/vulkan-shaders-gen.cpp @@ -2348,6 +2348,10 @@ void process_shaders() { {}, true, true); + string_to_spv( + "affine_bf16_exponents", + "affine_bf16_exponents.comp", + {}); string_to_spv( "fused_affine_qmm_bf16_bf16_coop4", "mul_mm_affine_bf16_coop4.comp", diff --git a/mlx/backend/vulkan/quantized.cpp b/mlx/backend/vulkan/quantized.cpp index fb490e8fe2..e3a1b6ed43 100644 --- a/mlx/backend/vulkan/quantized.cpp +++ b/mlx/backend/vulkan/quantized.cpp @@ -1912,12 +1912,17 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { "[QuantizedMatmul::eval_gpu] Cooperative affine-4 dispatch grid exceeds Vulkan limits."); } + array exponents( + {static_cast(rows + cols)}, uint32, nullptr, {}); + exponents.set_data(allocator::malloc(exponents.nbytes())); + auto command_buffer = vulkan::begin_command_recording(s.index); - vulkan::dispatch_fused_affine_matmul_op( + vulkan::dispatch_fused_affine_coop4_matmul_op( w, scales_bf16, biases_bf16, x_mat, + exponents, out_work, vulkan::StaticShaderId::fused_affine_qmm_bf16_bf16_coop4_cm1, command_buffer, @@ -1925,6 +1930,7 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { push_constants, grid); vulkan::end_command_recording(s.index); + vulkan::retain_array_for_stream(s, exponents); } finalize_bf16_output(out_work); trace_qmm( From e6b880b1e7f2e5cc1657f45d7367ee4bc3bd8136 Mon Sep 17 00:00:00 2001 From: Goni Zahavy Date: Thu, 16 Jul 2026 05:40:44 +0300 Subject: [PATCH 7/9] [Vulkan] Harden coopmat aligned/K and coop4 subgroup safety Gate aligned matmul on full BK (k%32), require real 16x16 f16 coopmat props before cm1 tuning, size exponent scratch for 16 subgroups, and force subgroup size 64 for coop4 / exponent pipelines. --- mlx/backend/vulkan/kernels.cpp | 4 +++- .../vulkan/kernels/affine_bf16_exponents.comp | 4 +++- mlx/backend/vulkan/matmul.cpp | 12 +++++++++++- 3 files changed, 17 insertions(+), 3 deletions(-) diff --git a/mlx/backend/vulkan/kernels.cpp b/mlx/backend/vulkan/kernels.cpp index d95326b740..9dc5760945 100644 --- a/mlx/backend/vulkan/kernels.cpp +++ b/mlx/backend/vulkan/kernels.cpp @@ -245,7 +245,9 @@ PipelineCreationOptions pipeline_creation_options( PipelineCreationOptions options; if (shader_name == "gather_affine_qmm_rhs_bf16_bf16_cm1" || - shader_name == "fused_affine_qmm_bf16_bf16_coop4_cm1") { + shader_name == "fused_affine_qmm_bf16_bf16_coop4_cm1" || + shader_name == "fused_affine_qmm_bf16_bf16_coop4" || + shader_name == "affine_bf16_exponents") { options.require_full_subgroups = true; options.required_subgroup_size = 64; return options; diff --git a/mlx/backend/vulkan/kernels/affine_bf16_exponents.comp b/mlx/backend/vulkan/kernels/affine_bf16_exponents.comp index c4c6897989..b88d2cf41e 100644 --- a/mlx/backend/vulkan/kernels/affine_bf16_exponents.comp +++ b/mlx/backend/vulkan/kernels/affine_bf16_exponents.comp @@ -36,7 +36,9 @@ layout(push_constant) uniform parameter { uint mode; } p; -shared uint subgroup_max_exponents[8]; +// Workgroup is 256 threads. With a 16-lane default subgroup (and no required +// size), gl_NumSubgroups can be 16 — size for that worst case. +shared uint subgroup_max_exponents[16]; uint fp32_biased_exponent(const float value) { const uint exponent = (floatBitsToUint(abs(value)) >> 23u) & 0xffu; diff --git a/mlx/backend/vulkan/matmul.cpp b/mlx/backend/vulkan/matmul.cpp index 0b5d8369bd..69adb4c736 100644 --- a/mlx/backend/vulkan/matmul.cpp +++ b/mlx/backend/vulkan/matmul.cpp @@ -375,6 +375,13 @@ bool prefer_matmul_coopmat1( if (!matmul_coopmat_env_enabled() || !ctx.cooperative_matrix_supported()) { return false; } + // mul_mm.comp cm1 paths require a real 16x16x16 f16 A/B coopmat mode. + // Extension presence alone is not enough: without it we would still replace + // scalar TM/TN/TK with 16 and break the scalar fallback candidates. + if (!ctx.coopmat_f16acc_supported() && + !ctx.coopmat_flash_attention_f32acc_supported()) { + return false; + } // Native BF16 coopmat shaders are incorrect on current AMD drivers (cosine // ~0.72). Use F16 coopmat instead (with BF16→F16 promote when needed). if (dtype != float16 || !ctx.shader_float16_supported()) { @@ -646,7 +653,10 @@ MatmulFamily classify_matmul_family(uint32_t m, uint32_t n, uint32_t k) { } bool matmul_inputs_aligned(uint32_t m, uint32_t n, uint32_t k) { - return (m % 4u) == 0 && (n % 8u) == 0 && (k % 8u) == 0; + // Aligned mul_mm load paths omit end_k checks and assume a full BK tile. + // F16/F32 shaders hardcode BK=32; BF16 uses BK=16. Require k%32 so aligned + // candidates never read past the logical K dimension (e.g. K=72). + return (m % 4u) == 0 && (n % 8u) == 0 && (k % 32u) == 0; } uint32_t round_up_div(uint32_t value, uint32_t divisor) { From 6ae98f9c4cd8c93afc9c794e4035609a6fe9bd47 Mon Sep 17 00:00:00 2001 From: Goni Zahavy Date: Thu, 16 Jul 2026 06:41:45 +0300 Subject: [PATCH 8/9] [Vulkan] Cache F16-promoted B^T for BF16 coopmat1 GEMMs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Large BF16×BF16 coopmat1 paths already cache BF16 B^T but re-cast to F16 every call. Store the F16 promote on the same transpose-cache entry and only scratch-cast dynamic A. --- mlx/backend/vulkan/matmul.cpp | 114 ++++++++++++++++++++++++++++++++-- 1 file changed, 110 insertions(+), 4 deletions(-) diff --git a/mlx/backend/vulkan/matmul.cpp b/mlx/backend/vulkan/matmul.cpp index 69adb4c736..fc94cacafa 100644 --- a/mlx/backend/vulkan/matmul.cpp +++ b/mlx/backend/vulkan/matmul.cpp @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -49,6 +50,8 @@ struct MulMmTransposeCacheEntry { int64_t source_offset{0}; Dtype source_dtype; array transposed; + // Optional F16 copy of transposed for BF16→F16 coopmat1 promote. + std::optional f16_promoted; }; std::mutex& mul_mm_transpose_cache_mutex() { @@ -1064,9 +1067,14 @@ void cache_mul_mm_transpose(const array& source, const array& transposed) { std::lock_guard lock(mul_mm_transpose_cache_mutex()); auto& cache = mul_mm_transpose_cache(); + std::optional keep_f16; for (auto it = cache.begin(); it != cache.end();) { - if (it->source.expired() || - same_mul_mm_transpose_source(*it, source_data, source)) { + if (it->source.expired()) { + it = cache.erase(it); + continue; + } + if (same_mul_mm_transpose_source(*it, source_data, source)) { + keep_f16 = std::move(it->f16_promoted); it = cache.erase(it); } else { ++it; @@ -1080,12 +1088,109 @@ void cache_mul_mm_transpose(const array& source, const array& transposed) { source.strides(), source.offset(), source.dtype(), - transposed}); + transposed, + std::move(keep_f16)}); + while (cache.size() > kMulMmTransposeCacheLimit) { + cache.pop_front(); + } +} + +std::optional find_cached_mul_mm_f16_promote(const array& source) { + auto source_data = source.data_shared_ptr(); + if (source_data == nullptr) { + return std::nullopt; + } + + std::lock_guard lock(mul_mm_transpose_cache_mutex()); + auto& cache = mul_mm_transpose_cache(); + for (auto it = cache.begin(); it != cache.end();) { + if (it->source.expired()) { + it = cache.erase(it); + continue; + } + if (same_mul_mm_transpose_source(*it, source_data, source) && + it->f16_promoted.has_value()) { + auto cached = *it->f16_promoted; + auto entry = std::move(*it); + cache.erase(it); + cache.push_back(std::move(entry)); + return cached; + } + ++it; + } + return std::nullopt; +} + +void cache_mul_mm_f16_promote( + const array& source, + const array& b_t_bf16, + const array& f16_bt) { + if (!cacheable_mul_mm_transpose_source(source) || source.dtype() != bfloat16) { + return; + } + auto source_data = source.data_shared_ptr(); + if (source_data == nullptr || b_t_bf16.data_shared_ptr() == nullptr || + f16_bt.data_shared_ptr() == nullptr || f16_bt.dtype() != float16) { + return; + } + + std::lock_guard lock(mul_mm_transpose_cache_mutex()); + auto& cache = mul_mm_transpose_cache(); + for (auto it = cache.begin(); it != cache.end();) { + if (it->source.expired()) { + it = cache.erase(it); + continue; + } + if (same_mul_mm_transpose_source(*it, source_data, source)) { + it->f16_promoted = f16_bt; + auto entry = std::move(*it); + cache.erase(it); + cache.push_back(std::move(entry)); + return; + } + ++it; + } + cache.push_back( + MulMmTransposeCacheEntry{ + source_data, + source_data.get(), + source.shape(), + source.strides(), + source.offset(), + source.dtype(), + b_t_bf16, + f16_bt}); while (cache.size() > kMulMmTransposeCacheLimit) { cache.pop_front(); } } +array cast_to_float16_contiguous(const array& arr, Stream s) { + array out(arr.shape(), float16, nullptr, {}); + out.set_data(allocator::malloc(out.nbytes())); + copy_gpu(arr, out, CopyType::General, s); + return out; +} + +// Cacheable BF16→F16 promote of a materialized B^T. Keyed by original B so the +// F16 copy lives with the BF16 transpose cache entry (A stays per-call scratch). +array materialize_mul_mm_f16_promoted_transpose( + const array& source_b, + const array& b_t_bf16, + Stream s) { + if (detail::in_tracing() || detail::retain_graph() || + !cacheable_mul_mm_transpose_source(source_b) || + source_b.dtype() != bfloat16) { + return cast_to_float16_contiguous(b_t_bf16, s); + } + if (auto cached = find_cached_mul_mm_f16_promote(source_b)) { + return *cached; + } + array f16_bt = cast_to_float16_contiguous(b_t_bf16, s); + cache_mul_mm_f16_promote(source_b, b_t_bf16, f16_bt); + return f16_bt; +} + array materialize_mul_mm_transpose(array b, Stream s, bool allow_cache = true) { array b_t = swapaxes_in_eval(b, -1, -2); if (is_row_contiguous_zero_offset(b_t)) { @@ -1407,7 +1512,8 @@ bool try_eval_mul_mm_vulkan( if (promote_bf16_to_f16_coopmat) { a = cast_to_float16_scratch(a, s, kMulMmACastScratchLane); - b_t = cast_to_float16_scratch(b_t, s, kMulMmBCastScratchLane); + // B^T F16 promote is keyed by original BF16 B (same cache as transpose). + b_t = materialize_mul_mm_f16_promoted_transpose(b, b_t, s); if (matmul_debug_enabled()) { std::cerr << "[vulkan::mul_mm] promote bf16->f16 for coopmat1\n"; } From 05644d634e94acc2029821c462fae1535ed821d4 Mon Sep 17 00:00:00 2001 From: Goni Zahavy Date: Thu, 16 Jul 2026 07:04:50 +0300 Subject: [PATCH 9/9] [Vulkan] Fall back SDPA VJP to composed graph Custom ScaledDotProductAttentionVJP SIGFPEs on AMD while evaluating dQ. Match Metal and differentiate through the fallback graph instead. --- mlx/backend/vulkan/fast.cpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/mlx/backend/vulkan/fast.cpp b/mlx/backend/vulkan/fast.cpp index 19a01f3563..2f6efcfd8c 100644 --- a/mlx/backend/vulkan/fast.cpp +++ b/mlx/backend/vulkan/fast.cpp @@ -2691,7 +2691,10 @@ bool ScaledDotProductAttention::supports_bool_mask() { } bool ScaledDotProductAttentionVJP::use_fallback(const array& q, Stream s) { - return s.device == Device::cpu; + // Match Metal: differentiate through the composed fallback graph. + // The custom Vulkan SDPA VJP path currently SIGFPEs on AMD (Radeon 8060S) + // while evaluating dQ for float16 attention (see mlx-vulkan#66). + return true; } void ScaledDotProductAttention::eval_gpu(