From 6ad2b8f3ca356b6251794aecb4b3068c0b1340d3 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Tue, 7 Jul 2026 00:07:11 +0900 Subject: [PATCH 1/2] perf(cuda): extend sdpa_vector decode kernels to head_dim 256/288 (#675) Extend the CUDA sdpa_vector decode kernels and the supports_sdpa_vector gate to head_dim 256 and 288 through the mlxcel MLX overlay, so head_dim > 128 models (gemma family 256, qwen3.5/3.6 256, baichuan-m1 256, paligemma2 288) take the fused vector path at decode instead of MLX's unfused materializing SDPA fallback. This is a single full-file overlay of scaled_dot_product_attention.cu; the .cpp is untouched because cudnn rejects head_dim > 128, so widening this gate alone routes 256/288 decode to the vector kernels. dispatch_headdim gains case 256/288 plus a throwing default (a gate/dispatch desync previously launched no kernel and left the output buffer as garbage), the D=256/288 instantiations get __launch_bounds__ so their larger per-thread register footprint still fits the 1024/256/1024 thread launches, and the path sits behind the MLXCEL_SDPA_VECTOR_LARGE_D env kill switch (default on) for A/B and rollback without a rebuild. Measured on GB10 (same bench harness, env toggle): the fused path is a real win on full-attention head_dim-256 families, qwen3.5-2b-4bit at 8192 context decodes 129.06 vs 114.89 tok/s (1.123x) and qwen3.5-9b-4bit 38.55 vs 36.92 tok/s (1.044x), growing with attention share. gemma-4-12b-it-4bit stays flat (0.993/1.001/0.994/1.007x at 512/2048/8192/32768) because only 8 of its 48 layers are full-attention and its decode is dominated by model-specific fixed overhead rather than attention; that overhead is tracked as follow-up issue #680, which carries the original 2x goal. Output is byte-identical to the fallback on a chat-template greedy prompt, and the near-tie divergence seen on degenerate pad prompts is FP-ordering noise present in both paths. Numbers are recorded in benchmarks/cuda_gb10_sdpav_675_2026-07-06.csv, and the overlay is a full-file copy synced to upstream e9463bb that must be re-synced on MLX pin bumps. --- benchmarks/cuda_gb10_sdpav_675_2026-07-06.csv | 13 + .../mlx-pin-upgrade-2026-07-03.md | 1 + .../cuda/scaled_dot_product_attention.cu | 836 ++++++++++++++++++ 3 files changed, 850 insertions(+) create mode 100644 benchmarks/cuda_gb10_sdpav_675_2026-07-06.csv create mode 100644 src/lib/mlx-cpp/patches/mlx/backend/cuda/scaled_dot_product_attention.cu diff --git a/benchmarks/cuda_gb10_sdpav_675_2026-07-06.csv b/benchmarks/cuda_gb10_sdpav_675_2026-07-06.csv new file mode 100644 index 000000000..5cbc1769b --- /dev/null +++ b/benchmarks/cuda_gb10_sdpav_675_2026-07-06.csv @@ -0,0 +1,13 @@ +model,prompt_tokens,generated_tokens,prefill_chunk,prefill_ms,prefill_tok_s,decode_ms,decode_tok_s,mlx_peak_gb,sys_peak_gb,date,hardware,mlx_version,build_type,notes +gemma-4-12b-it-4bit,512,64,,,,4377.57,14.62,,,2026-07-06,NVIDIA_GB10_122GB,0.3.3,release,sdpa_vector path (#675 overlay); 1-pass kernel; decode_ms derived from tok/s +gemma-4-12b-it-4bit,512,64,,,,4347.83,14.72,,,2026-07-06,NVIDIA_GB10_122GB,0.3.3,release,materializing fallback (MLXCEL_SDPA_VECTOR_LARGE_D=0); decode_ms derived from tok/s +gemma-4-12b-it-4bit,2048,64,,,,4466.15,14.33,,,2026-07-06,NVIDIA_GB10_122GB,0.3.3,release,sdpa_vector path (#675 overlay); 2-pass kernel; decode_ms derived from tok/s +gemma-4-12b-it-4bit,2048,64,,,,4469.27,14.32,,,2026-07-06,NVIDIA_GB10_122GB,0.3.3,release,materializing fallback (MLXCEL_SDPA_VECTOR_LARGE_D=0); decode_ms derived from tok/s +gemma-4-12b-it-4bit,8192,64,,9768.20,838.64,4481.07,14.28,14.51,,2026-07-06,NVIDIA_GB10_122GB,0.3.3,release,sdpa_vector path (#675 overlay); 2-pass kernel +gemma-4-12b-it-4bit,8192,64,,10003.80,818.89,4454.60,14.37,14.51,,2026-07-06,NVIDIA_GB10_122GB,0.3.3,release,materializing fallback (MLXCEL_SDPA_VECTOR_LARGE_D=0) +gemma-4-12b-it-4bit,32768,64,,42398.79,772.85,4643.21,13.78,16.65,,2026-07-07,NVIDIA_GB10_122GB,0.3.3,release,sdpa_vector path (#675 overlay); only 8/48 layers are full-attention +gemma-4-12b-it-4bit,32768,64,,43317.00,756.47,4678.69,13.68,16.65,,2026-07-07,NVIDIA_GB10_122GB,0.3.3,release,materializing fallback (MLXCEL_SDPA_VECTOR_LARGE_D=0) +qwen3.5-2b-4bit,8192,64,,2517.63,3253.85,495.89,129.06,5.48,,2026-07-07,NVIDIA_GB10_122GB,0.3.3,release,sdpa_vector path (#675 overlay); full attention every layer; 1.123x vs fallback +qwen3.5-2b-4bit,8192,64,,2553.53,3208.11,557.06,114.89,5.48,,2026-07-07,NVIDIA_GB10_122GB,0.3.3,release,materializing fallback (MLXCEL_SDPA_VECTOR_LARGE_D=0) +qwen3.5-9b-4bit,8192,64,,10923.69,749.93,1660.21,38.55,12.30,,2026-07-07,NVIDIA_GB10_122GB,0.3.3,release,sdpa_vector path (#675 overlay); full attention every layer; 1.044x vs fallback +qwen3.5-9b-4bit,8192,64,,11060.52,740.65,1733.28,36.92,12.30,,2026-07-07,NVIDIA_GB10_122GB,0.3.3,release,materializing fallback (MLXCEL_SDPA_VECTOR_LARGE_D=0) diff --git a/docs/benchmark_results/mlx-pin-upgrade-2026-07-03.md b/docs/benchmark_results/mlx-pin-upgrade-2026-07-03.md index 5af6bde4b..3cdc15c48 100644 --- a/docs/benchmark_results/mlx-pin-upgrade-2026-07-03.md +++ b/docs/benchmark_results/mlx-pin-upgrade-2026-07-03.md @@ -40,6 +40,7 @@ in `src/lib/mlx-cpp/CMakeLists.txt`. Each was reconciled three-way | `patches/mlx/backend/cuda/reduce/init_reduce.cu` | kept | Base unchanged. Output buffer typed as T so it matches out.dtype() for bf16. | | `patches/mlx/backend/cuda/reduce/reduce_ops.cuh` | kept | Base unchanged. `ReduceResult` accumulate-in-fp32 specializations. | | `patches/mlx/backend/cuda/reduce/row_reduce.cu` | kept | Base unchanged. Same V/U split. | +| `patches/mlx/backend/cuda/scaled_dot_product_attention.cu` | added (#675) | New overlay, not part of the #625 rebase. Extends the `sdpa_vector` decode kernels and the `supports_sdpa_vector` gate to head_dim 256/288 so head_dim > 128 models (gemma family, qwen3.5/3.6, baichuan-m1, paligemma2) take the fused vector path at decode instead of the unfused materializing SDPA fallback. Gated behind the `MLXCEL_SDPA_VECTOR_LARGE_D` env kill switch (default on); `__launch_bounds__` caps the larger per-thread register footprint of the D=256/288 instantiations. Base is pristine e9463bb, so future pin bumps must rebase this file. | | `patches/mlx/backend/metal/compiled.cpp` | kept | Base unchanged. static_cast insertion for mixed-dtype compiled ops. | | `patches/mlx/backend/metal/kernels/utils.h` | kept | Base unchanged. `` include and `metal::vec` qualification. | | `patches-cuda/dtype.cpp` | kept | Base unchanged. bf16+fp32 -> bf16 promotion-table patch (CUDA-only). | diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/scaled_dot_product_attention.cu b/src/lib/mlx-cpp/patches/mlx/backend/cuda/scaled_dot_product_attention.cu new file mode 100644 index 000000000..a82f50755 --- /dev/null +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/scaled_dot_product_attention.cu @@ -0,0 +1,836 @@ +// Copyright © 2025 Apple Inc. +// +// Patched by mlxcel: extend the sdpa_vector decode kernels and the +// supports_sdpa_vector gate to head_dim 256 and 288, so head_dim > 128 models +// (gemma family 256, qwen3.5/3.6 256, baichuan-m1 256, paligemma2 288) take the +// fused vector path at decode instead of the unfused materializing SDPA +// fallback (see lablup/mlxcel#675). Three changes: (A) dispatch_headdim gains +// case 256/288 plus a default that throws on a gate/dispatch desync (turns a +// silent no-launch with garbage output into a loud error); (B) +// supports_sdpa_vector accepts head_dim 256/288 behind the +// MLXCEL_SDPA_VECTOR_LARGE_D env kill switch (default on) so A/B and rollback +// need no rebuild; (C) __launch_bounds__ caps the larger per-thread register +// footprint of the D=256/288 instantiations to the 1024/256/1024 thread +// launches (a no-op for D <= 128 which already fit). The .cpp is unchanged: +// cudnn rejects head_dim > 128, so widening this gate alone routes 256/288 +// decode to the vector kernels. Synced to upstream e9463bb. + +// Required for using M_LOG2E in MSVC. +#define _USE_MATH_DEFINES + +#include "mlx/backend/cuda/device.h" +#include "mlx/backend/cuda/device/config.h" +#include "mlx/backend/cuda/device/utils.cuh" +#include "mlx/backend/cuda/kernel_utils.cuh" +#include "mlx/backend/gpu/copy.h" +#include "mlx/dtype_utils.h" +#include "mlx/utils.h" // mlxcel: env::get_var for the MLXCEL_SDPA_VECTOR_LARGE_D kill switch + +#include +#include + +#include // mlxcel: std::runtime_error in the dispatch_headdim default +#include // mlxcel: std::to_string in the dispatch_headdim default + +namespace mlx::core { + +namespace cu { + +namespace cg = cooperative_groups; + +#define PRAGMA_LOOP_UNROLL #pragma unroll + +struct AttnParams { + int B; + int H; + int D; + + int qL; + int kL; + + int gqa_factor; + float scale; + + int64_t Q_strides[3]; + int64_t K_strides[3]; + int64_t V_strides[3]; + int64_t O_strides[3]; +}; + +template +__global__ __launch_bounds__(1024) void kernel_sdpav_1pass( + const T* Q, + const T* K, + const T* V, + T* O, + const T* sinks, + __grid_constant__ const AttnParams params) { + constexpr int BN = 32; + constexpr int BD = 32; + + constexpr int v_per_thread = D / BD; + + const int inner_k_stride = BN * int(params.K_strides[2]); + const int inner_v_stride = BN * int(params.V_strides[2]); + + typedef float U; + + U q[v_per_thread]; + U k[v_per_thread]; + U o[v_per_thread]; + + __shared__ U outputs[BN][BD + 1]; + __shared__ U max_scores[BN]; + __shared__ U sum_exp_scores[BN]; + + const U scale_log2 = params.scale * M_LOG2E; + + auto block = cg::this_thread_block(); + auto warp = cg::tiled_partition<32>(block); + + const int lane_idx = warp.thread_rank(); + const int warp_idx = warp.meta_group_rank(); + + // Adjust to thread block and thread + const int batch_idx = blockIdx.z; + const int head_idx = blockIdx.x; + const int kv_head_idx = head_idx / params.gqa_factor; + + const int q_seq_idx = blockIdx.y; + const int kv_seq_idx = warp_idx; + + Q += batch_idx * params.Q_strides[0] + // Batch + head_idx * params.Q_strides[1] + // Head + q_seq_idx * params.Q_strides[2]; // Sequence + + K += batch_idx * params.K_strides[0] + // Batch + kv_head_idx * params.K_strides[1] + // Head + kv_seq_idx * params.K_strides[2]; // Sequence + + V += batch_idx * params.V_strides[0] + // Batch + kv_head_idx * params.V_strides[1] + // Head + kv_seq_idx * params.V_strides[2]; // Sequence + + O += batch_idx * params.O_strides[0] + // Batch + head_idx * params.O_strides[1] + // Head + q_seq_idx * params.O_strides[2]; // Sequence + + // Read the query and 0 the output accumulator + PRAGMA_LOOP_UNROLL + for (int i = 0; i < v_per_thread; i++) { + q[i] = scale_log2 * static_cast(Q[v_per_thread * lane_idx + i]); + } + + PRAGMA_LOOP_UNROLL + for (int i = 0; i < v_per_thread; i++) { + o[i] = 0.f; + } + + U max_score = Limits::finite_min(); + U sum_exp_score = 0.f; + if (sinks && warp_idx == 0) { + max_score = M_LOG2E * static_cast(sinks[head_idx]); + sum_exp_score = 1.f; + } + + // For each key + for (int i = kv_seq_idx; i < params.kL; i += BN) { + bool use_key = true; + if constexpr (do_causal) { + use_key = i <= (params.kL - params.qL + q_seq_idx); + } + + if (use_key) { + // Read the key + PRAGMA_LOOP_UNROLL + for (int j = 0; j < v_per_thread; j++) { + k[j] = K[v_per_thread * lane_idx + j]; + } + + // Compute the i-th score + U score = 0.f; + PRAGMA_LOOP_UNROLL + for (int j = 0; j < v_per_thread; j++) { + score += q[j] * k[j]; + } + + // Warp sum + score = cg::reduce(warp, score, cg::plus()); + + // Update the accumulators + U new_max = max(max_score, score); + U factor = exp2f(max_score - new_max); + U exp_score = exp2f(score - new_max); + + max_score = new_max; + sum_exp_score = sum_exp_score * factor + exp_score; + + // Update the output accumulator + PRAGMA_LOOP_UNROLL + for (int j = 0; j < v_per_thread; j++) { + o[j] = o[j] * factor + + exp_score * static_cast(V[v_per_thread * lane_idx + j]); + } + } + + // Move the pointers to the next kv + K += inner_k_stride; + V += inner_v_stride; + } + + if (lane_idx == 0) { + max_scores[warp_idx] = max_score; + sum_exp_scores[warp_idx] = sum_exp_score; + } + block.sync(); + + max_score = max_scores[lane_idx]; + U new_max = cg::reduce(warp, max_score, cg::greater()); + U factor = exp2f(max_score - new_max); + sum_exp_score = + cg::reduce(warp, sum_exp_scores[lane_idx] * factor, cg::plus()); + sum_exp_score = sum_exp_score == 0 ? 0 : __frcp_rn(sum_exp_score); + + // Now we need to aggregate all the outputs + PRAGMA_LOOP_UNROLL + for (int i = 0; i < v_per_thread; i++) { + outputs[lane_idx][warp_idx] = o[i]; + block.sync(); + U ot = outputs[warp_idx][lane_idx] * factor; + o[i] = cg::reduce(warp, ot, cg::plus()) * sum_exp_score; + block.sync(); + } + + // And write the output + if (lane_idx == 0) { + PRAGMA_LOOP_UNROLL + for (int i = 0; i < v_per_thread; i++) { + O[v_per_thread * warp_idx + i] = static_cast(o[i]); + } + } +} + +template +__global__ __launch_bounds__(256) void kernel_sdpav_2pass_1( + const T* Q, + const T* K, + const T* V, + const T* sinks, + float* partials, + float* sums, + float* maxs, + __grid_constant__ const AttnParams params) { + constexpr int BN = 8; + constexpr int BD = 32; + constexpr int blocks = 32; + + constexpr int v_per_thread = D / BD; + + const int inner_k_stride = blocks * BN * int(params.K_strides[2]); + const int inner_v_stride = blocks * BN * int(params.V_strides[2]); + + typedef float U; + + U q[v_per_thread]; + U k[v_per_thread]; + U o[v_per_thread]; + + __shared__ U outputs[BN][BD + 1]; + __shared__ U max_scores[BN]; + __shared__ U sum_exp_scores[BN]; + + const U scale_log2 = params.scale * 1.44269504089f; + + auto block = cg::this_thread_block(); + auto warp = cg::tiled_partition<32>(block); + + const int lane_idx = warp.thread_rank(); + const int warp_idx = warp.meta_group_rank(); + + // Adjust to thread block and thread + const int batch_idx = blockIdx.z / blocks; + const int block_idx = blockIdx.z % blocks; + const int head_idx = blockIdx.x; + const int kv_head_idx = head_idx / params.gqa_factor; + + const int q_seq_idx = blockIdx.y; + const int kv_seq_idx = block_idx * BN + warp_idx; + + Q += batch_idx * params.Q_strides[0] + // Batch + head_idx * params.Q_strides[1] + // Head + q_seq_idx * params.Q_strides[2]; // Sequence + + K += batch_idx * params.K_strides[0] + // Batch + kv_head_idx * params.K_strides[1] + // Head + kv_seq_idx * params.K_strides[2]; // Sequence + + V += batch_idx * params.V_strides[0] + // Batch + kv_head_idx * params.V_strides[1] + // Head + kv_seq_idx * params.V_strides[2]; // Sequence + + const int p_stride_s = blocks; + const int p_stride_h = params.qL * p_stride_s; + const int p_stride_b = params.H * p_stride_h; + const int p_offset = batch_idx * p_stride_b + // Batch + head_idx * p_stride_h + // Head + q_seq_idx * p_stride_s + // Sequence + block_idx; // Block + + partials += p_offset * D; + sums += p_offset; + maxs += p_offset; + + // Read the query and 0 the output accumulator + PRAGMA_LOOP_UNROLL + for (int i = 0; i < v_per_thread; i++) { + q[i] = scale_log2 * static_cast(Q[v_per_thread * lane_idx + i]); + } + + PRAGMA_LOOP_UNROLL + for (int i = 0; i < v_per_thread; i++) { + o[i] = 0.f; + } + + U max_score = Limits::finite_min(); + U sum_exp_score = 0.f; + if (sinks && warp_idx == 0 && block_idx == 0) { + max_score = M_LOG2E * static_cast(sinks[head_idx]); + sum_exp_score = 1.f; + } + + // For each key + for (int i = kv_seq_idx; i < params.kL; i += blocks * BN) { + bool use_key = true; + if constexpr (do_causal) { + use_key = i <= (params.kL - params.qL + q_seq_idx); + } + + if (use_key) { + // Read the key + PRAGMA_LOOP_UNROLL + for (int j = 0; j < v_per_thread; j++) { + k[j] = K[v_per_thread * lane_idx + j]; + } + + // Compute the i-th score + U score = 0.f; + PRAGMA_LOOP_UNROLL + for (int j = 0; j < v_per_thread; j++) { + score += q[j] * k[j]; + } + + // Warp sum + score = cg::reduce(warp, score, cg::plus()); + + // Update the accumulators + U new_max = max(max_score, score); + U factor = exp2f(max_score - new_max); + U exp_score = exp2f(score - new_max); + + max_score = new_max; + sum_exp_score = sum_exp_score * factor + exp_score; + + // Update the output accumulator + PRAGMA_LOOP_UNROLL + for (int j = 0; j < v_per_thread; j++) { + o[j] = o[j] * factor + + exp_score * static_cast(V[v_per_thread * lane_idx + j]); + } + } + + // Move the pointers to the next kv + K += inner_k_stride; + V += inner_v_stride; + } + + if (lane_idx == 0) { + max_scores[warp_idx] = max_score; + sum_exp_scores[warp_idx] = sum_exp_score; + } + + block.sync(); + + max_score = (lane_idx < BN) ? max_scores[lane_idx] : -1e9; + U new_max = cg::reduce(warp, max_score, cg::greater()); + U factor = exp2f(max_score - new_max); + sum_exp_score = (lane_idx < BN) ? sum_exp_scores[lane_idx] : 0.f; + sum_exp_score = cg::reduce(warp, sum_exp_score * factor, cg::plus()); + + // Write the sum and new max + if (warp_idx == 0) { + sums[0] = sum_exp_score; + maxs[0] = new_max; + } + + // Now we need to aggregate all the outputs + auto ff = exp2f(max_scores[warp_idx] - new_max); + PRAGMA_LOOP_UNROLL + for (int i = 0; i < v_per_thread; i++) { + outputs[warp_idx][lane_idx] = o[i] * ff; + block.sync(); + + if (warp_idx == 0) { + U ot = outputs[0][lane_idx]; + PRAGMA_LOOP_UNROLL + for (int j = 1; j < BN; j++) { + ot += outputs[j][lane_idx]; + warp.sync(); + } + o[i] = ot; + } + block.sync(); + } + + if (warp_idx == 0) { + PRAGMA_LOOP_UNROLL + for (int i = 0; i < v_per_thread; i++) { + partials[v_per_thread * lane_idx + i] = o[i]; + } + } +} + +template +__global__ __launch_bounds__(1024) void kernel_sdpav_2pass_2( + const float* partials, + const float* sums, + const float* maxs, + T* O, + __grid_constant__ const AttnParams params) { + constexpr int BN = 32; + constexpr int BD = 32; + constexpr int blocks = 32; + + constexpr int v_per_thread = D / BD; + + typedef float U; + + U o[v_per_thread]; + __shared__ U outputs[BN][BD + 1]; + + auto block = cg::this_thread_block(); + auto warp = cg::tiled_partition<32>(block); + + const int lane_idx = warp.thread_rank(); + const int warp_idx = warp.meta_group_rank(); + + // Adjust to thread block and thread + const int batch_idx = blockIdx.z; + const int head_idx = blockIdx.x; + const int q_seq_idx = blockIdx.y; + + const int p_stride_s = blocks; + const int p_stride_h = params.qL * p_stride_s; + const int p_stride_b = params.H * p_stride_h; + const int p_offset = batch_idx * p_stride_b + // Batch + head_idx * p_stride_h + // Head + q_seq_idx * p_stride_s; // Sequence + + partials += p_offset * D + warp_idx * D; + sums += p_offset; + maxs += p_offset; + + O += batch_idx * params.O_strides[0] + // Batch + head_idx * params.O_strides[1] + // Head + q_seq_idx * params.O_strides[2]; // Sequence + + U max_score = maxs[lane_idx]; + U new_max = cg::reduce(warp, max_score, cg::greater()); + U factor = exp2f(max_score - new_max); + U sum_exp_score = cg::reduce(warp, sums[lane_idx] * factor, cg::plus()); + sum_exp_score = sum_exp_score == 0 ? 0 : __frcp_rn(sum_exp_score); + + PRAGMA_LOOP_UNROLL + for (int i = 0; i < v_per_thread; i++) { + o[i] = partials[v_per_thread * lane_idx + i]; + } + + // Now we need to aggregate all the outputs + PRAGMA_LOOP_UNROLL + for (int i = 0; i < v_per_thread; i++) { + outputs[lane_idx][warp_idx] = o[i]; + block.sync(); + U ot = outputs[warp_idx][lane_idx] * factor; + o[i] = cg::reduce(warp, ot, cg::plus()) * sum_exp_score; + block.sync(); + } + + // And write the output + if (lane_idx == 0) { + PRAGMA_LOOP_UNROLL + for (int i = 0; i < v_per_thread; i++) { + O[v_per_thread * warp_idx + i] = static_cast(o[i]); + } + } +} + +} // namespace cu + +namespace { + +template +void dispatch_headdim(int n, F&& f) { + switch (n) { + case 64: + f(std::integral_constant{}); + break; + case 96: + f(std::integral_constant{}); + break; + case 128: + f(std::integral_constant{}); + break; + // mlxcel: head_dim 256 (gemma family, qwen3.5/3.6, baichuan-m1) and 288 + // (paligemma2). Both are divisible by BD=32, so v_per_thread and the fixed + // shared-memory tiles are the only D-dependent quantities and the existing + // kernels handle them unchanged. + case 256: + f(std::integral_constant{}); + break; + case 288: + f(std::integral_constant{}); + break; + // mlxcel: reachable only if supports_sdpa_vector and this dispatch ever + // desync. Upstream had no default, so an unhandled head_dim launched no + // kernel and left the output buffer as uninitialized garbage. Fail loudly. + default: + throw std::runtime_error( + "sdpa_vector: unsupported head_dim " + std::to_string(n)); + } +} + +void sdpa_vector_1pass_fallback( + const Stream& s, + cu::CommandEncoder& encoder, + const array& q, + const array& k, + const array& v, + const float scale, + array& o, + bool do_causal, + const std::optional& sinks) { + encoder.set_input_array(q); + encoder.set_input_array(k); + encoder.set_input_array(v); + if (sinks) { + encoder.set_input_array(*sinks); + } + encoder.set_output_array(o); + + cu::AttnParams params{ + /* int B = */ q.shape(0), + /* int H = */ q.shape(1), + /* int D = */ q.shape(3), + + /* int qL = */ q.shape(2), + /* int kL = */ k.shape(2), + + /* int gqa_factor = */ q.shape(1) / k.shape(1), + /* float scale = */ scale, + + /* int64_t Q_strides[3] = */ {q.strides(0), q.strides(1), q.strides(2)}, + /* int64_t K_strides[3] = */ {k.strides(0), k.strides(1), k.strides(2)}, + /* int64_t V_strides[3] = */ {v.strides(0), v.strides(1), v.strides(2)}, + /* int64_t O_strides[3] = */ {o.strides(0), o.strides(1), o.strides(2)}}; + + dim3 grid_dim(params.H, params.qL, params.B); + dim3 block_dim(1024, 1, 1); + + dispatch_float_types(o.dtype(), "kernel_sdpav_1pass", [&](auto type_tag) { + dispatch_bool(do_causal, [&](auto do_causal) { + dispatch_headdim(params.D, [&](auto headdim) { + using DataType = cuda_type_t; + + auto kernel = + cu::kernel_sdpav_1pass; + encoder.add_kernel_node( + kernel, + grid_dim, + block_dim, + gpu_ptr(q), + gpu_ptr(k), + gpu_ptr(v), + gpu_ptr(o), + sinks ? gpu_ptr(*sinks) : nullptr, + params); + }); + }); + }); +} + +void sdpa_vector_2pass_fallback( + const Stream& s, + cu::CommandEncoder& encoder, + const array& q, + const array& k, + const array& v, + const float scale, + array& o, + bool do_causal, + const std::optional& sinks) { + cu::AttnParams params{ + /* int B = */ q.shape(0), + /* int H = */ q.shape(1), + /* int D = */ q.shape(3), + + /* int qL = */ q.shape(2), + /* int kL = */ k.shape(2), + + /* int gqa_factor = */ q.shape(1) / k.shape(1), + /* float scale = */ scale, + + /* int64_t Q_strides[3] = */ {q.strides(0), q.strides(1), q.strides(2)}, + /* int64_t K_strides[3] = */ {k.strides(0), k.strides(1), k.strides(2)}, + /* int64_t V_strides[3] = */ {v.strides(0), v.strides(1), v.strides(2)}, + /* int64_t O_strides[3] = */ {o.strides(0), o.strides(1), o.strides(2)}}; + + // Allocate the intermediates + int blocks = 32; + + Shape intermediate_shape; + intermediate_shape.reserve(o.ndim() + 1); + intermediate_shape.insert( + intermediate_shape.end(), o.shape().begin(), o.shape().end() - 1); + intermediate_shape.push_back(blocks); + intermediate_shape.push_back(o.shape().back()); + + array intermediate(intermediate_shape, float32, nullptr, {}); + intermediate_shape.pop_back(); + array sums(intermediate_shape, float32, nullptr, {}); + array maxs(std::move(intermediate_shape), float32, nullptr, {}); + + intermediate.set_data(cu::malloc_async(intermediate.nbytes(), encoder)); + sums.set_data(cu::malloc_async(sums.nbytes(), encoder)); + maxs.set_data(cu::malloc_async(maxs.nbytes(), encoder)); + + encoder.add_temporary(intermediate); + encoder.add_temporary(sums); + encoder.add_temporary(maxs); + + dispatch_float_types(o.dtype(), "kernel_sdpav_2pass", [&](auto type_tag) { + dispatch_bool(do_causal, [&](auto do_causal) { + dispatch_headdim(params.D, [&](auto headdim) { + using DataType = cuda_type_t; + + { + auto kernel = cu:: + kernel_sdpav_2pass_1; + + encoder.set_input_array(q); + encoder.set_input_array(k); + encoder.set_input_array(v); + if (sinks) { + encoder.set_input_array(*sinks); + } + + encoder.set_output_array(intermediate); + encoder.set_output_array(sums); + encoder.set_output_array(maxs); + + dim3 grid_dim(params.H, params.qL, params.B * 32); + dim3 block_dim(8 * 32, 1, 1); + + encoder.add_kernel_node( + kernel, + grid_dim, + block_dim, + gpu_ptr(q), + gpu_ptr(k), + gpu_ptr(v), + sinks ? gpu_ptr(*sinks) : nullptr, + gpu_ptr(intermediate), + gpu_ptr(sums), + gpu_ptr(maxs), + params); + } + + { + auto kernel = cu:: + kernel_sdpav_2pass_2; + + encoder.set_input_array(intermediate); + encoder.set_input_array(sums); + encoder.set_input_array(maxs); + encoder.set_output_array(o); + + dim3 grid_dim(params.H, params.qL, params.B); + dim3 block_dim(1024, 1, 1); + + encoder.add_kernel_node( + kernel, + grid_dim, + block_dim, + gpu_ptr(intermediate), + gpu_ptr(sums), + gpu_ptr(maxs), + gpu_ptr(o), + params); + } + }); + }); + }); +} + +void sdpa_vector_fallback( + const Stream& s, + cu::CommandEncoder& encoder, + const array& q, + const array& k, + const array& v, + const float scale, + array& o, + bool do_causal, + const std::optional& sinks) { + int kL = k.shape(2); + + if (kL > 1024) { + return sdpa_vector_2pass_fallback( + s, encoder, q, k, v, scale, o, do_causal, sinks); + } else { + return sdpa_vector_1pass_fallback( + s, encoder, q, k, v, scale, o, do_causal, sinks); + } +} + +} // namespace + +bool supports_sdpa_vector( + const array& q, + const array& k, + const array& v, + bool has_arr_mask, + bool output_logsumexp) { + if (output_logsumexp) { + return false; + } + + const int value_head_dim = v.shape(-1); + const int query_head_dim = q.shape(-1); + const int query_sequence_length = q.shape(2); + const int key_sequence_length = k.shape(2); + + // mlxcel: head_dim 256/288 (the head_dim > 128 models) are gated behind a + // static env kill switch, default on. It lets token-identity and perf A/B + // against the old materializing fallback, and a rollback, run without a + // rebuild. cudnn rejects head_dim > 128, so this widened gate alone routes + // their decode to the vector kernels; the .cpp needs no change. + static bool large_d_enabled = env::get_var("MLXCEL_SDPA_VECTOR_LARGE_D", 1); + const bool sdpa_supported_head_dim = query_head_dim == value_head_dim && + (query_head_dim == 64 || query_head_dim == 96 || query_head_dim == 128 || + (large_d_enabled && + (query_head_dim == 256 || query_head_dim == 288))); + + const bool supported_vector_config = + sdpa_supported_head_dim && query_sequence_length < 4; + + return supported_vector_config && !has_arr_mask; +} + +void sdpa_vector( + const array& q_pre, + const array& k_pre, + const array& v_pre, + float scale, + array& o, + bool do_causal, + const std::optional& sinks_pre, + Stream s) { + auto& encoder = cu::get_command_encoder(s); + std::vector copies; + + // Define some copy functions to ensure the layout of the inputs is as + // expected. + copies.reserve(4); + auto copy_unless = [&copies, &s]( + auto predicate, const array& arr) -> const array& { + if (!predicate(arr)) { + array arr_copy = contiguous_copy_gpu(arr, s); + copies.push_back(std::move(arr_copy)); + return copies.back(); + } else { + return arr; + } + }; + + // Checks that the headdim dimension has stride 1. + auto is_matrix_contiguous = [](const array& arr) { + return arr.strides(-1) == 1; + }; + + std::optional sinks = std::nullopt; + if (sinks_pre) { + sinks = copy_unless(is_matrix_contiguous, sinks_pre.value()); + } + + // We are in vector mode ie single query + if (q_pre.shape(2) < 4) { + auto q_copy_unless = [](const array& arr) { + if (arr.flags().row_contiguous) { + return true; + } + auto& strides = arr.strides(); + auto& shape = arr.shape(); + if (shape[0] == 1 || shape[1] == 1) { + // If either the batch or head dimension is a singleton, the other can + // be transposed with the sequence dimension + auto bidx = shape[0] == 1 ? 1 : 0; + return (strides[3] == 1) && (strides[2] == shape[3] * shape[bidx]) && + (strides[bidx] == shape[3]); + } + return false; + }; + + auto kv_copy_unless = [](const array& arr) { + // keys and values should be copied if: + // - the last dimension is not contiguous + // - the batch and head dim are not contiguous + auto& strides = arr.strides(); + auto& shape = arr.shape(); + if (strides.back() != 1) { + return false; + } + if (shape[0] == 1 || shape[1] == 1) { + return true; + } + return (strides[0] == strides[1] * shape[1]); + }; + + const auto& q = copy_unless(q_copy_unless, q_pre); + const auto& k = copy_unless(kv_copy_unless, k_pre); + const auto& v = copy_unless(kv_copy_unless, v_pre); + + // Donate the query if possible + if (q.is_donatable() && q.flags().row_contiguous && q.size() == o.size()) { + o.copy_shared_buffer(q); + } else { + int64_t str_oD = 1; + int64_t str_oH = o.shape(3); + int64_t str_oL = o.shape(1) * str_oH; + int64_t str_oB = o.shape(2) * str_oL; + + array::Flags flags{ + /* bool contiguous = */ 1, + /* bool row_contiguous = */ o.shape(2) == 1, + /* bool col_contiguous = */ o.size() == o.shape(3), + }; + + o.set_data( + cu::malloc_async(o.nbytes(), encoder), + o.size(), + {str_oB, str_oH, str_oL, str_oD}, + flags); + } + + for (const auto& cp : copies) { + encoder.add_temporary(cp); + } + + sdpa_vector_fallback(s, encoder, q, k, v, scale, o, do_causal, sinks); + } + + // Full attention mode should never reach here + else { + throw std::runtime_error("Doesn't support matrix yet."); + } +} + +} // namespace mlx::core From 3b126a81ee4335bcdd4ae2b1efb5b06baa5d6648 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Tue, 7 Jul 2026 00:34:29 +0900 Subject: [PATCH 2/2] docs: document MLXCEL_SDPA_VECTOR_LARGE_D env var (#675) Add the missing row to the environment-variables reference page for the CUDA sdpa_vector large-head_dim kill switch introduced in #681, alongside the other attention-path diagnostic toggles. --- docs/environment-variables.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 29c3b0ca0..ef0cec72d 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -207,6 +207,7 @@ recommended as normal deployment settings. | `MLXCEL_DISABLE_SOFTCAP_GQA_DECODE_GROUPED` | `1` disables, `0` enables | unset | Legacy rollback/override for grouped softcap-GQA decode. | | `MLXCEL_DISABLE_SINGLE_QUERY_MASKLESS` | truthy disables | maskless path on | Disables the single-query maskless attention path. | | `MLXCEL_EXPERIMENTAL_BOOL_CAUSAL_MASK` | truthy enables | off | Enables an experimental boolean causal-mask path. | +| `MLXCEL_SDPA_VECTOR_LARGE_D` | `0`/`false`/`off`/`no` disable; any other value or unset enables | on | **CUDA only.** Gates whether the CUDA `supports_sdpa_vector` check accepts head_dim 256/288 (gemma family, qwen3.5/3.6, baichuan-m1, paligemma2), routing their decode to the fused `sdpa_vector` kernels instead of the materializing SDPA fallback (issue #675). Disabling restores the prior fallback with no rebuild; used for the A/B in `benchmarks/cuda_gb10_sdpav_675_2026-07-06.csv`. | | `MLXCEL_PIPELINE_GRANULARITY` | `off`, `layer`, `block:N` | `off` | Inserts layer-boundary async-eval hints for pipeline experiments. | | `MLXCEL_FUSED_MOE` | `0`/`false`/`off`/`no` disable; any other value or unset enables | on | Fused single-token decode-MoE kernel (#268), on by default since #282 (Metal) and #319 (CUDA, via `mx.fast.cuda_kernel`); validated on M1 Ultra, M5, and GB10. Set to `0` to force the proven `gather_qmm`/`SwitchGLU` path. Active for qwen3_moe, qwen3_next, dots.llm1, gemma4, qwen2_moe, mixtral, phimoe, lfm2, qwen3_vl_moe, and olmoe decode. | | `MLXCEL_FUSED_MOE_SGY` | `1`-`32` | `8` | Simdgroups (Metal) / warps-per-block (CUDA) per threadgroup for the fused decode-MoE kernel; tune per hardware. |