From 3541464dffa7d2dcd4f0ecad0225afef094c955f Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Tue, 13 Oct 2026 21:02:40 +0900 Subject: [PATCH 1/3] fix(rocm): defer SDPA accumulator rescale, stage Q once per block The fused WMMA flash attention kernel rescaled every O accumulator tile through LDS on every K block: a store, scale and reload of all D_TILES f32 tiles, 3 barriers per tile, 24 barriers per K block at D=128, and re-staged Q from global memory on every K block. At the chunked-prefill shapes (2,048-token query chunks against 2K to 16K keys, GQA 32/8, head dim 128) that per-block scalar overhead dominated: the kernel ran at 1.1 to 1.7 TFLOP/s while the same device's f16 GEMM path delivers about 14 TFLOP/s at the same shape. The O entries and l of a row now stay in the softmax frame the row was last rescaled to, so the rescale runs only when a row's block max outgrows that frame by more than a 10-point margin (kFaWmmaRescaleMargin), never in the steady state of a chunked prefill where the block max moves a few points per block; P = exp(S - f) stays bounded by exp(10) = 22,026, well inside the bf16 score cap. Q is staged once per block when the head dim fits one DW-wide slice (D <= 128). Refs #2251 --- .../mlx/backend/rocm/flash_attention_wmma.hip | 91 ++++++++++++++++--- 1 file changed, 76 insertions(+), 15 deletions(-) diff --git a/src/lib/mlxcelverse/overlays/rocm/mlx/backend/rocm/flash_attention_wmma.hip b/src/lib/mlxcelverse/overlays/rocm/mlx/backend/rocm/flash_attention_wmma.hip index 33f95c085..948b1bea3 100644 --- a/src/lib/mlxcelverse/overlays/rocm/mlx/backend/rocm/flash_attention_wmma.hip +++ b/src/lib/mlxcelverse/overlays/rocm/mlx/backend/rocm/flash_attention_wmma.hip @@ -1,7 +1,11 @@ // WMMA-accelerated flash attention for RDNA 3+ (gfx1100+) // // Uses rocwmma 16x16x16 bf16→f32 tiles for Q@K^T and P@V matmuls. -// Implements FlashAttention-2 online softmax. +// Implements FlashAttention-2 online softmax with a deferred accumulator +// rescale: the O tiles and l of a row live in the softmax frame the row was +// last rescaled to, so the store/scale/reload of every accumulator tile +// through LDS runs only when a row's block max outgrows that frame by more +// than kFaWmmaRescaleMargin, never in the steady state of a chunked prefill. // // BLOCK_M=64, BLOCK_N=64, 128 threads (4 waves), each wave owns 16 query rows. // Shared memory ~50 KB for D=128. @@ -35,6 +39,15 @@ struct FAWmmaParams { int64_t LSE_strides[3]; }; +// Margin between a row's running block max and its softmax frame before the O +// accumulator rescale runs. Entries of O and l stay in the frame f the row +// was last rescaled to, so P = exp(S - f) is bounded by exp(margin) instead +// of 1; the margin keeps that well inside the score dtype (bf16 P caps at +// 65504, exp(10) = 22,026) while making the per-K-block rescale a steady-state +// no-op for a chunked prefill, where the block max moves a few points per +// block. +constexpr float kFaWmmaRescaleMargin = 10.0f; + // Helper: collaborative load from global to shared memory template __device__ void load_tile( @@ -163,6 +176,23 @@ __global__ void __launch_bounds__(256) + q_start * params.Q_strides[2]; const int q_valid = min(BLOCK_M, params.qL - q_start); + // ---- Stage Q once when the head dim fits one DW-wide slice (D <= 128) ---- + // Q_sh is resident through the K/V loop, so the per-K-block re-stage inside + // the loop below (a global re-read plus an LDS write of BLOCK_M x DW halfs + // per block, up to 256 iterations per call at kL=16K) runs only for D > DW. + if constexpr (N_DCHUNKS == 1) { + load_tile( + Q_sh, Q_base, params.Q_strides[2], BLOCK_M, min(DW, D), q_valid, tid, + NTHREADS); + __syncthreads(); + } + + // Dynamic smem scratch: per-row scale_old for the (deferred) accumulator + // rescale, plus the per-block rescale flag. + float* wave_scale = reinterpret_cast( + reinterpret_cast(S_sh) + BLOCK_M * S_PAD * sizeof(float)); + bool* rescale_flag = reinterpret_cast(wave_scale + BLOCK_M); + // ---- K/V block loop ---- for (int k_start = 0; k_start < params.kL; k_start += BLOCK_N) { @@ -173,6 +203,10 @@ __global__ void __launch_bounds__(256) if (k_start > max_k_allowed) break; } + // Reset the per-block accumulator-rescale flag; the __syncthreads inside + // the S stage below orders this reset before the online softmax sets it. + if (tid == 0) *rescale_flag = false; + int k_valid = min(BLOCK_N, params.kL - k_start); const T* K_base = K + bid_b * params.K_strides[0] @@ -191,9 +225,11 @@ __global__ void __launch_bounds__(256) for (int dc = 0; dc < N_DCHUNKS; dc++) { const int d0 = dc * DW; const int dwid = min(DW, D - d0); - load_tile( - Q_sh, Q_base + d0, params.Q_strides[2], BLOCK_M, dwid, q_valid, tid, - NTHREADS); + if constexpr (N_DCHUNKS > 1) { + load_tile( + Q_sh, Q_base + d0, params.Q_strides[2], BLOCK_M, dwid, q_valid, + tid, NTHREADS); + } load_tile( KV_sh, K_base + d0, params.K_strides[2], BLOCK_N, dwid, k_valid, tid, NTHREADS); @@ -235,9 +271,13 @@ __global__ void __launch_bounds__(256) rocwmma::synchronize_workgroup(); // ---- Online softmax: each wave processes its own 16 query rows ---- - // Dynamic smem slot for per-row scale_old (avoids static/dynamic LDS aliasing). - float* wave_scale = reinterpret_cast( - reinterpret_cast(S_sh) + BLOCK_M * S_PAD * sizeof(float)); + // Deferred accumulator rescale: the O entries and l of a row live in the + // softmax frame f = m_arr[row] the row was last rescaled to, so the + // store/scale/reload of all D_TILES f32 tiles through LDS (3 barriers per + // tile, 24 per K block at D=128) only runs when a row's block max + // outgrows its frame by more than kFaWmmaRescaleMargin, never in the + // steady state of a chunked prefill. + float scale_old = 1.0f; if (lane < WT) { int row = wave * WT + lane; // 0..BLOCK_M-1, one row per lane of this wave int q_idx = q_start + row; @@ -256,7 +296,26 @@ __global__ void __launch_bounds__(256) } } - float scale_old = (old_m > -1e29f) ? expf(old_m - new_m) : 0.0f; + if (valid && old_l > 0.0f) { + // O already carries entries in the frame old_m: rescale them only + // when the frame would go stale by more than the margin. + if ((new_m - old_m) > kFaWmmaRescaleMargin) { + scale_old = expf(old_m - new_m); + m_arr[row] = new_m; + *rescale_flag = true; + } else { + scale_old = 1.0f; + m_arr[row] = old_m; + } + } else if (valid) { + // Nothing accumulated yet (O is zero): move the frame for free. + scale_old = 0.0f; + m_arr[row] = new_m; + } else { + // A padding row: no state change, P writes 0 below. + scale_old = 0.0f; + } + float row_sum = 0.0f; for (int j = 0; j < BLOCK_N; j++) { @@ -265,7 +324,7 @@ __global__ void __launch_bounds__(256) use = use && ((k_start + j) <= (params.kL - params.qL + q_idx)); if (use) { - float p = expf(S_sh[row * S_PAD + j] - new_m); + float p = expf(S_sh[row * S_PAD + j] - m_arr[row]); S_sh[row * S_PAD + j] = p; row_sum += p; } else { @@ -273,16 +332,16 @@ __global__ void __launch_bounds__(256) } } - m_arr[row] = new_m; l_arr[row] = old_l * scale_old + row_sum; wave_scale[row] = scale_old; } __syncthreads(); - // ---- Rescale O accumulators by scale_old ---- + // ---- Rescale O accumulators by scale_old (only when a row outgrew its + // frame this block) ---- // Each wave rescales its 16-row O tiles. // Since fragment layout is opaque, store→scale→reload via shared memory. - { + if (*rescale_flag) { // Use KV_sh area as f32 temp (it held K which we no longer need) float* o_tmp = reinterpret_cast(KV_sh); constexpr int OT_PAD = WT + 4; @@ -445,13 +504,15 @@ bool supports_sdpa_flash_wmma( int sdpa_flash_wmma_smem(int D) { constexpr int BM = 64, BN = 64; const int DW = (D < 128) ? D : 128; - // Layout: m[BM] l[BM] | Q[BM,DW+4] | KV[BN,DW+4] | S[BM,BN+4] | wave_scale[BM] - // wave_scale is dynamic (not static __shared__) so it is accounted here. + // Layout: m[BM] l[BM] | Q[BM,DW+4] | KV[BN,DW+4] | S[BM,BN+4] | + // wave_scale[BM] | rescale_flag + // wave_scale and the per-block rescale flag are dynamic (not static + // __shared__) so they are accounted here. return 2 * BM * (int)sizeof(float) + BM * (DW + 4) * (int)sizeof(uint16_t) + BN * (DW + 4) * (int)sizeof(uint16_t) + BM * (BN + 4) * (int)sizeof(float) - + BM * (int)sizeof(float); + + (BM + 1) * (int)sizeof(float); } void sdpa_flash_wmma( From b86d5ae744325e90ff923b5384820aca06c35710 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 14 Oct 2026 04:24:13 +0900 Subject: [PATCH 2/3] test(rocm): add flash SDPA prefill tests and probe Three tests on the fused ROCm SDPA prefill path at the chunked-prefill shapes: causal correctness against an explicit f32 matmul + softmax + matmul reference at the 2K and 16K key shapes for f16 and bf16, a score ramp that drives a row's block max past the rescale margin so the deferred branch runs, and a throughput gate that requires the 16K-key f16 call to clear 1.1 TFLOP/s, set between the measured before (0.95) and after (1.20 to 1.31) on an idle gfx1151. The probe example times the maskless causal call at the three key lengths against a dense f16 GEMM bound; it builds only under the rocm feature. Refs #2251 --- src/lib/mlxcel-core/Cargo.toml | 8 + .../examples/sdpa_prefill_probe.rs | 267 +++++++++++++++ tests/rocm_sdpa_flash_wmma_prefill.rs | 315 ++++++++++++++++++ 3 files changed, 590 insertions(+) create mode 100644 src/lib/mlxcel-core/examples/sdpa_prefill_probe.rs create mode 100644 tests/rocm_sdpa_flash_wmma_prefill.rs diff --git a/src/lib/mlxcel-core/Cargo.toml b/src/lib/mlxcel-core/Cargo.toml index 3e51d6e50..d0850f4a6 100644 --- a/src/lib/mlxcel-core/Cargo.toml +++ b/src/lib/mlxcel-core/Cargo.toml @@ -47,3 +47,11 @@ path = "examples/benchmark.rs" [[example]] name = "wht_microbench" path = "examples/wht_microbench.rs" + +# Per-call time and TFLOP/s of the fused ROCm SDPA prefill path at the +# chunked-prefill shapes (issue #2251). Needs a ROCm host; builds only under +# the rocm feature, which carries the overlay link recipe. +[[example]] +name = "sdpa_prefill_probe" +path = "examples/sdpa_prefill_probe.rs" +required-features = ["rocm"] diff --git a/src/lib/mlxcel-core/examples/sdpa_prefill_probe.rs b/src/lib/mlxcel-core/examples/sdpa_prefill_probe.rs new file mode 100644 index 000000000..01c5293db --- /dev/null +++ b/src/lib/mlxcel-core/examples/sdpa_prefill_probe.rs @@ -0,0 +1,267 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Per-call time and achieved TFLOP/s of the fused ROCm SDPA prefill path at +//! the chunked-prefill shapes (issue #2251). +//! +//! The batched server prefills long prompts in 2,048-token chunks, so every +//! attention call of a chunk sees a query block of 2,048 rows against 2K to +//! 16K keys (GQA 32/8, head dim 128, causal, f16). On gfx1151 that call is +//! dominated by `kernel_sdpa_flash_wmma` at about 192 ms per call. This probe +//! measures the call the model makes, the maskless causal +//! `fast_scaled_dot_product_attention_causal` wrapper, at those shapes, +//! reports the median per-call time and the achieved TFLOP/s from the +//! causal-aware FLOP count, and compares against a dense f16 `matmul` of the +//! same shape (the device GEMM bound). Run it a second time with +//! `MLX_SDPA_NO_WMMA=1` to measure the scalar flash arm. +//! +//! ```sh +//! cargo build --release --features rocm -p mlxcel-core --example sdpa_prefill_probe +//! target/release/examples/sdpa_prefill_probe --key-lens 2048,8192,16384 +//! ``` + +use std::time::Instant; + +use mlxcel_core::streams::DefaultDeviceGuard; +use mlxcel_core::{self, MlxArray, UniquePtr, dtype}; + +struct Args { + q_len: i32, + key_lens: Vec, + heads: i32, + kv_heads: i32, + dim: i32, + runs: usize, + warmup: usize, + scale: f32, +} + +fn parse_args() -> Args { + let mut args = Args { + q_len: 2048, + key_lens: vec![2048, 8192, 16384], + heads: 32, + kv_heads: 8, + dim: 128, + runs: 5, + warmup: 2, + scale: 1.0 / (128.0_f32).sqrt(), + }; + // Both spellings: --key=value and --key value. + let tokens: Vec = std::env::args().skip(1).collect(); + let mut i = 0; + while i < tokens.len() { + let arg = tokens[i].clone(); + let (key, inline_value) = match arg.split_once('=') { + Some((k, v)) => (k.to_string(), Some(v.to_string())), + None => (arg.clone(), None), + }; + let value = match inline_value { + Some(v) => v, + None => { + i += 1; + if i >= tokens.len() { + eprintln!("missing value for {key}"); + break; + } + tokens[i].clone() + } + }; + i += 1; + let take = |field: &mut i32| { + if let Ok(v) = value.parse() { + *field = v; + } + }; + match key.as_str() { + "--q-len" => take(&mut args.q_len), + "--heads" => take(&mut args.heads), + "--kv-heads" => take(&mut args.kv_heads), + "--dim" => { + take(&mut args.dim); + args.scale = 1.0 / (args.dim as f32).sqrt(); + } + "--runs" => { + if let Ok(v) = value.parse() { + args.runs = v; + } + } + "--warmup" => { + if let Ok(v) = value.parse() { + args.warmup = v; + } + } + "--scale" => { + if let Ok(v) = value.parse() { + args.scale = v; + } + } + "--key-lens" => { + args.key_lens = value.split(',').filter_map(|s| s.parse().ok()).collect(); + } + _ => eprintln!("ignoring unknown argument {arg}"), + } + } + args +} + +fn median(values: &mut [f64]) -> f64 { + values.sort_by(|a, b| a.partial_cmp(b).unwrap()); + let n = values.len(); + if n.is_multiple_of(2) { + (values[n / 2 - 1] + values[n / 2]) / 2.0 + } else { + values[n / 2] + } +} + +/// Uniform(-1, 1) values from an integer hash, so every host builds the same +/// inputs and the same shape is comparable across runs and builds. +fn pseudo_random(count: usize, seed: u32) -> Vec { + let mut state = seed; + (0..count) + .map(|_| { + state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + ((state >> 8) as f32 / (1u32 << 24) as f32) * 2.0 - 1.0 + }) + .collect() +} + +/// FLOPs of a causal SDPA call: query i (global position offset + i) attends +/// keys [0, offset + i + 1], for QK^T and P@V (4 D-FLOPs per key per head). +fn causal_flops(heads: i32, dim: i32, q_len: i32, k_len: i32) -> f64 { + let offset = (k_len - q_len).max(0) as f64; + let q = q_len as f64; + 4.0 * heads as f64 * dim as f64 * (q * offset + q * (q + 1.0) / 2.0) +} + +fn build_inputs( + heads: i32, + kv_heads: i32, + q_len: i32, + k_len: i32, + dim: i32, +) -> ( + UniquePtr, + UniquePtr, + UniquePtr, +) { + let q_f32 = mlxcel_core::from_slice_f32( + &pseudo_random((heads * q_len * dim) as usize, 0x2251), + &[1, heads, q_len, dim], + ); + let k_f32 = mlxcel_core::from_slice_f32( + &pseudo_random((kv_heads * k_len * dim) as usize, 0x2252), + &[1, kv_heads, k_len, dim], + ); + let v_f32 = mlxcel_core::from_slice_f32( + &pseudo_random((kv_heads * k_len * dim) as usize, 0x2253), + &[1, kv_heads, k_len, dim], + ); + let q = mlxcel_core::astype(&q_f32, dtype::FLOAT16); + let k = mlxcel_core::astype(&k_f32, dtype::FLOAT16); + let v = mlxcel_core::astype(&v_f32, dtype::FLOAT16); + mlxcel_core::eval(&q); + mlxcel_core::eval(&k); + mlxcel_core::eval(&v); + (q, k, v) +} + +/// Median per-call seconds of the maskless causal SDPA wrapper. +fn time_sdpa( + q: &MlxArray, + k: &MlxArray, + v: &MlxArray, + scale: f32, + runs: usize, + warmup: usize, +) -> f64 { + for _ in 0..warmup { + let out = mlxcel_core::fast_scaled_dot_product_attention_causal(q, k, v, scale); + mlxcel_core::eval(&out); + } + let mut per_call: Vec = Vec::with_capacity(runs); + for _ in 0..runs { + let start = Instant::now(); + let out = mlxcel_core::fast_scaled_dot_product_attention_causal(q, k, v, scale); + mlxcel_core::eval(&out); + per_call.push(start.elapsed().as_secs_f64()); + } + median(&mut per_call) +} + +/// Median per-call seconds of a dense f16 GEMM of the QK^T shape, the device +/// GEMM bound for one of the two matmuls the flash kernel fuses. +fn time_dense_gemm(q_len: i32, k_len: i32, dim: i32, runs: usize, warmup: usize) -> f64 { + let a = mlxcel_core::from_slice_f32( + &pseudo_random((q_len * dim) as usize, 0x2254), + &[q_len, dim], + ); + let b = mlxcel_core::from_slice_f32( + &pseudo_random((k_len * dim) as usize, 0x2255), + &[k_len, dim], + ); + let a = mlxcel_core::astype(&a, dtype::FLOAT16); + let b = mlxcel_core::astype(&mlxcel_core::transpose(&b), dtype::FLOAT16); + mlxcel_core::eval(&a); + mlxcel_core::eval(&b); + for _ in 0..warmup { + let c = mlxcel_core::matmul(&a, &b); + mlxcel_core::eval(&c); + } + let mut per_call: Vec = Vec::with_capacity(runs); + for _ in 0..runs { + let start = Instant::now(); + let c = mlxcel_core::matmul(&a, &b); + mlxcel_core::eval(&c); + per_call.push(start.elapsed().as_secs_f64()); + } + median(&mut per_call) +} + +fn main() { + let args = parse_args(); + let _gpu = DefaultDeviceGuard::gpu(); + let no_wmma = std::env::var("MLX_SDPA_NO_WMMA").is_ok(); + + println!("=== SDPA chunked-prefill probe (issue #2251) ==="); + println!( + "q_len={} heads={} kv_heads={} dim={} runs={} warmup={} MLX_SDPA_NO_WMMA={}", + args.q_len, args.heads, args.kv_heads, args.dim, args.runs, args.warmup, no_wmma + ); + println!(); + println!( + "{:>8} {:>12} {:>12} {:>12} {:>12}", + "keys", "ms/call", "TFLOP/s", "gemm_ms", "gemm_TFLOP/s" + ); + println!("{}", "-".repeat(62)); + + for k_len in &args.key_lens { + let (q, k, v) = build_inputs(args.heads, args.kv_heads, args.q_len, *k_len, args.dim); + let secs = time_sdpa(&q, &k, &v, args.scale, args.runs, args.warmup); + let flops = causal_flops(args.heads, args.dim, args.q_len, *k_len); + let tflops = flops / secs / 1e12; + let gemm_secs = time_dense_gemm(args.q_len, *k_len, args.dim, args.runs, args.warmup); + let gemm_flops = 2.0 * (args.q_len as f64) * (*k_len as f64) * (args.dim as f64); + let gemm_tflops = gemm_flops / gemm_secs / 1e12; + println!( + "{:>8} {:>12.1} {:>12.2} {:>12.1} {:>12.2}", + k_len, + secs * 1e3, + tflops, + gemm_secs * 1e3, + gemm_tflops + ); + } +} diff --git a/tests/rocm_sdpa_flash_wmma_prefill.rs b/tests/rocm_sdpa_flash_wmma_prefill.rs new file mode 100644 index 000000000..adfd477c1 --- /dev/null +++ b/tests/rocm_sdpa_flash_wmma_prefill.rs @@ -0,0 +1,315 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The fused ROCm SDPA prefill path at the chunked-prefill shapes (issue +//! #2251). +//! +//! The batched server prefills long prompts in 2,048-token chunks, so every +//! attention call of a chunk sees a query block of 2,048 rows against 2K to +//! 16K keys (GQA 32/8, head dim 128, causal, f16). On gfx1151 that call runs +//! in `kernel_sdpa_flash_wmma` at 0.95 to 1.49 TFLOP/s while the device's +//! dense f16 GEMM path delivers about 14 TFLOP/s at the same shape. The fix +//! defers the online-softmax O-accumulator rescale so the store/scale/reload +//! of all accumulator tiles through LDS (3 barriers per tile, per K block) +//! runs only when a row's block max outgrows its softmax frame, and stages Q +//! once per block instead of re-staging it from global memory on every K +//! block. +//! +//! Tests: +//! +//! - correctness of the causal call against an explicit f32 +//! `matmul + softmax + matmul` reference at the 2K and 16K key shapes, for +//! f16 and bf16 (the deferred frame changes the rounding, not the result); +//! - a score ramp that drives a row's block max more than the rescale margin +//! above its frame, so the deferred rescale branch runs, still matching the +//! reference; +//! - a throughput reproduction: the median per-call time of the 16K-key f16 +//! call must reach the achieved TFLOP/s floor below. On gfx1151, both arms +//! the probe's median of 5 calls on an idle GPU: before the fix the call +//! measured 540.1 ms/call (0.95 TFLOP/s), after the fix 429.2 ms/call (1.20 +//! TFLOP/s), 1.26x; the non-WMMA fallback arm (`MLX_SDPA_NO_WMMA=1`) is +//! another 3.5x slower (1516.8 ms/call, 0.34 TFLOP/s). A gate run under +//! `scripts/rocm_gpu_guard.sh --hold` measured 392.6 ms/call (1.31 +//! TFLOP/s). The floor fails without the fix and passes with it. +//! +//! Skips on any other backend. Run on a ROCm host with: +//! +//! ```sh +//! cargo test --profile test-fast --features rocm --test rocm_sdpa_flash_wmma_prefill -- --test-threads=1 +//! ``` + +#![cfg(feature = "rocm")] + +use mlxcel_core::streams::DefaultDeviceGuard; +use mlxcel_core::{self, MlxArray, UniquePtr, dtype}; + +/// Relative error of a half-precision GEMM against the f32 reference. +const TOLERANCE: f32 = 5e-3; + +/// Rescale margin between a row's block max and its softmax frame. Must match +/// `kFaWmmaRescaleMargin` in `flash_attention_wmma.hip`. +const RESCALE_MARGIN: f32 = 10.0; + +/// Achieved TFLOP/s floor of the 16K-key f16 chunked-prefill call. Set +/// between the before (0.95) and after (1.20 to 1.31) measurements on an +/// idle gfx1151. +const TFLOPS_FLOOR: f64 = 1.1; + +fn on_rocm() -> bool { + mlxcel_core::hardware::gpu_backend_kind() == mlxcel_core::hardware::GpuBackendKind::Rocm +} + +/// Uniform(-1, 1) values from an integer hash, so every host builds the same +/// inputs. +fn pseudo_random(count: usize, seed: u32) -> Vec { + let mut state = seed; + (0..count) + .map(|_| { + state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + ((state >> 8) as f32 / (1u32 << 24) as f32) * 2.0 - 1.0 + }) + .collect() +} + +fn median(values: &mut [f64]) -> f64 { + values.sort_by(|a, b| a.partial_cmp(b).unwrap()); + let n = values.len(); + if n.is_multiple_of(2) { + (values[n / 2 - 1] + values[n / 2]) / 2.0 + } else { + values[n / 2] + } +} + +struct Inputs { + q: UniquePtr, + k: UniquePtr, + v: UniquePtr, +} + +impl Inputs { + /// `[1, heads, q_len, dim]` queries against `[1, kv_heads, k_len, dim]` + /// keys/values in `dt`, from fixed pseudo-random f32 sources. + fn new(heads: i32, kv_heads: i32, q_len: i32, k_len: i32, dim: i32, dt: i32) -> Self { + let q_f32 = mlxcel_core::from_slice_f32( + &pseudo_random((heads * q_len * dim) as usize, 0x2251), + &[1, heads, q_len, dim], + ); + let k_f32 = mlxcel_core::from_slice_f32( + &pseudo_random((kv_heads * k_len * dim) as usize, 0x2252), + &[1, kv_heads, k_len, dim], + ); + let v_f32 = mlxcel_core::from_slice_f32( + &pseudo_random((kv_heads * k_len * dim) as usize, 0x2253), + &[1, kv_heads, k_len, dim], + ); + Self { + q: mlxcel_core::astype(&q_f32, dt), + k: mlxcel_core::astype(&k_f32, dt), + v: mlxcel_core::astype(&v_f32, dt), + } + } + + /// Constant key rows whose value ramps linearly from 0 to `ramp` across + /// the key axis, with unit queries, so every query's score at key j is + /// close to `ramp * j / (k_len - 1)`: the block max outgrows a row's + /// softmax frame by more than the rescale margin across blocks, running + /// the deferred rescale branch. Values stay pseudo-random. + fn ramped(heads: i32, kv_heads: i32, q_len: i32, k_len: i32, dim: i32, ramp: f32) -> Self { + let mut k_data = vec![0.0_f32; (kv_heads * k_len * dim) as usize]; + for j in 0..k_len { + let value = ramp * (j as f32) / ((k_len - 1) as f32); + let base = (j * dim) as usize; + for kv in 0..kv_heads { + let row = (kv * k_len * dim) as usize + base; + for d in 0..dim { + k_data[row + d as usize] = value; + } + } + } + let k_f32 = mlxcel_core::from_slice_f32(&k_data, &[1, kv_heads, k_len, dim]); + // Unit queries: every entry 1/sqrt(dim), so q·k_j = value * sqrt(dim) + // and the wrapper's 1/sqrt(dim) scale brings the score back to the + // ramp. f16 keeps the value within rounding of it. + let unit = 1.0 / (dim as f32).sqrt(); + let q_data = vec![unit; (heads * q_len * dim) as usize]; + let q_f32 = mlxcel_core::from_slice_f32(&q_data, &[1, heads, q_len, dim]); + let v_f32 = mlxcel_core::from_slice_f32( + &pseudo_random((kv_heads * k_len * dim) as usize, 0x2253), + &[1, kv_heads, k_len, dim], + ); + Self { + q: mlxcel_core::astype(&q_f32, dtype::FLOAT16), + k: mlxcel_core::astype(&k_f32, dtype::FLOAT16), + v: mlxcel_core::astype(&v_f32, dtype::FLOAT16), + } + } +} + +fn max_abs(a: &MlxArray) -> f32 { + let f = mlxcel_core::astype(a, dtype::FLOAT32); + let m = mlxcel_core::max_all(&mlxcel_core::abs(&f)); + mlxcel_core::eval(&m); + mlxcel_core::item_f32(&m) +} + +/// Explicit f32 `matmul + softmax + matmul` reference of the maskless causal +/// call, per KV head so the [q_len, k_len] score matrix stays bounded, with a +/// bottom-right aligned causal mask over all keys. +fn reference( + inp: &Inputs, + heads: i32, + kv_heads: i32, + q_len: i32, + k_len: i32, + scale: f32, +) -> UniquePtr { + let gqa = heads / kv_heads; + let mut mask_data = vec![0.0_f32; (q_len * k_len) as usize]; + let offset = k_len - q_len; + for i in 0..q_len { + for j in 0..k_len { + if j > offset + i { + mask_data[(i * k_len + j) as usize] = -1e30_f32; + } + } + } + let mask = mlxcel_core::from_slice_f32(&mask_data, &[q_len, k_len]); + let mut parts: Vec> = Vec::with_capacity(kv_heads as usize); + for kv in 0..kv_heads { + let q_h = mlxcel_core::slice( + &inp.q, + &[0, kv * gqa, 0, 0], + &[1, (kv + 1) * gqa, q_len, head_dim(&inp.q)], + ); + let k_h = mlxcel_core::slice( + &inp.k, + &[0, kv, 0, 0], + &[1, kv + 1, k_len, head_dim(&inp.k)], + ); + let v_h = mlxcel_core::slice( + &inp.v, + &[0, kv, 0, 0], + &[1, kv + 1, k_len, head_dim(&inp.v)], + ); + let k_t = mlxcel_core::transpose_axes(&k_h, &[0, 1, 3, 2]); + let scores = mlxcel_core::matmul(&q_h, &k_t); + let scores = mlxcel_core::multiply_scalar(&scores, scale); + let scores = mlxcel_core::add(&scores, &mask); + let probs = mlxcel_core::softmax(&scores, -1); + let out = mlxcel_core::matmul(&probs, &v_h); + parts.push(out); + } + let refs: Vec<&MlxArray> = parts.iter().map(|p| p.as_ref().unwrap()).collect(); + mlxcel_core::concatenate_many(&refs, 1) +} + +fn head_dim(a: &MlxArray) -> i32 { + *mlxcel_core::array_shape(a).last().unwrap() +} + +/// Relative L2 distance of the fused output against the f32 reference. +fn check_case(dt: i32, name: &str, heads: i32, kv_heads: i32, q_len: i32, k_len: i32) { + let inp = Inputs::new(heads, kv_heads, q_len, k_len, 128, dt); + let scale = 1.0 / (128.0_f32).sqrt(); + let out = mlxcel_core::fast_scaled_dot_product_attention_causal(&inp.q, &inp.k, &inp.v, scale); + if let Err(err) = mlxcel_core::try_eval(&out) { + panic!("{name}: causal SDPA on the GPU failed: {err}"); + } + let expect = reference(&inp, heads, kv_heads, q_len, k_len, scale); + if let Err(err) = mlxcel_core::try_eval(&expect) { + panic!("{name}: f32 reference failed: {err}"); + } + let diff = mlxcel_core::subtract(&mlxcel_core::astype(&out, dtype::FLOAT32), &expect); + let err = max_abs(&diff); + assert!( + err <= TOLERANCE, + "{name}: max abs diff {err:.3e} against the f32 reference exceeds {TOLERANCE:.1e}" + ); +} + +#[test] +fn causal_prefill_shapes_match_reference() { + if !on_rocm() { + return; + } + let _gpu = DefaultDeviceGuard::gpu(); + // The chunked-prefill shapes: a 2,048-query chunk against 2K and 16K keys. + check_case(dtype::FLOAT16, "f16 2048q/2048k", 32, 8, 2048, 2048); + check_case(dtype::FLOAT16, "f16 2048q/16384k", 32, 8, 2048, 16384); + check_case(dtype::BFLOAT16, "bf16 2048q/16384k", 32, 8, 2048, 16384); +} + +#[test] +fn score_ramp_exercises_deferred_rescale() { + if !on_rocm() { + return; + } + let _gpu = DefaultDeviceGuard::gpu(); + // Scores span 3x the margin across the key axis, so rows rescale several + // times over the K/V loop; the result must still match the reference. + let inp = Inputs::ramped(32, 8, 2048, 8192, 128, 3.0 * RESCALE_MARGIN); + let scale = 1.0 / (128.0_f32).sqrt(); + let out = mlxcel_core::fast_scaled_dot_product_attention_causal(&inp.q, &inp.k, &inp.v, scale); + if let Err(err) = mlxcel_core::try_eval(&out) { + panic!("ramped causal SDPA on the GPU failed: {err}"); + } + let expect = reference(&inp, 32, 8, 2048, 8192, scale); + if let Err(err) = mlxcel_core::try_eval(&expect) { + panic!("ramped f32 reference failed: {err}"); + } + let diff = mlxcel_core::subtract(&mlxcel_core::astype(&out, dtype::FLOAT32), &expect); + let err = max_abs(&diff); + assert!( + err <= TOLERANCE, + "ramped: max abs diff {err:.3e} against the f32 reference exceeds {TOLERANCE:.1e}" + ); +} + +/// Causal-aware FLOPs of a maskless causal SDPA call. +fn causal_flops(heads: i32, dim: i32, q_len: i32, k_len: i32) -> f64 { + let offset = (k_len - q_len).max(0) as f64; + let q = q_len as f64; + 4.0 * heads as f64 * dim as f64 * (q * offset + q * (q + 1.0) / 2.0) +} + +#[test] +fn chunked_prefill_throughput_reaches_floor() { + if !on_rocm() { + return; + } + let _gpu = DefaultDeviceGuard::gpu(); + let inp = Inputs::new(32, 8, 2048, 16384, 128, dtype::FLOAT16); + let scale = 1.0 / (128.0_f32).sqrt(); + for _ in 0..2 { + let out = + mlxcel_core::fast_scaled_dot_product_attention_causal(&inp.q, &inp.k, &inp.v, scale); + mlxcel_core::eval(&out); + } + let mut per_call: Vec = Vec::with_capacity(5); + for _ in 0..5 { + let start = std::time::Instant::now(); + let out = + mlxcel_core::fast_scaled_dot_product_attention_causal(&inp.q, &inp.k, &inp.v, scale); + mlxcel_core::eval(&out); + per_call.push(start.elapsed().as_secs_f64()); + } + let secs = median(&mut per_call); + let tflops = causal_flops(32, 128, 2048, 16384) / secs / 1e12; + assert!( + tflops >= TFLOPS_FLOOR, + "16K-key chunked-prefill call measured {tflops:.2} TFLOP/s ({:.1} ms/call), floor {TFLOPS_FLOOR}", + secs * 1e3 + ); +} From 661dc1903fa2d87dcf83f04cab131b5641a45551 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 14 Oct 2026 04:24:27 +0900 Subject: [PATCH 3/3] docs: add gfx1151 flash SDPA prefill benchmark record Benchmark session page for the flash SDPA prefill record on Radeon 8060S (gfx1151) with the measurement harness scripts (session, bench round, server). The kernel-level table is filled from the probe's idle-GPU A/B (before 540.1 ms, after 429.2 ms at 16K keys, 1.26x, 3.5x over the non-WMMA fallback); the model-level session (bench-decode, batch-4 server) is marked not measured yet because the shared host kept the guard lock busy. Refs #2251 --- .../harness/run_bench_round.sh | 36 ++++++++++ .../harness/run_server.sh | 71 +++++++++++++++++++ .../session.sh | 48 +++++++++++++ .../sdpa-flash-prefill-gfx1151-2026-10-13.md | 31 ++++++++ 4 files changed, 186 insertions(+) create mode 100644 docs/benchmark_results/data/sdpa-flash-prefill-gfx1151-2026-10-13/harness/run_bench_round.sh create mode 100644 docs/benchmark_results/data/sdpa-flash-prefill-gfx1151-2026-10-13/harness/run_server.sh create mode 100644 docs/benchmark_results/data/sdpa-flash-prefill-gfx1151-2026-10-13/session.sh create mode 100644 docs/benchmark_results/sdpa-flash-prefill-gfx1151-2026-10-13.md diff --git a/docs/benchmark_results/data/sdpa-flash-prefill-gfx1151-2026-10-13/harness/run_bench_round.sh b/docs/benchmark_results/data/sdpa-flash-prefill-gfx1151-2026-10-13/harness/run_bench_round.sh new file mode 100644 index 000000000..fa342578e --- /dev/null +++ b/docs/benchmark_results/data/sdpa-flash-prefill-gfx1151-2026-10-13/harness/run_bench_round.sh @@ -0,0 +1,36 @@ +#!/usr/bin/env bash +# One probe + bench-decode round for issue #2251 on gfx1151, one arm. +# Usage: run_bench_round.sh OUT_DIR ARM ROUND BIN_ROOT MODELS_DIR GUARD REPS +# BIN_ROOT: directory holding target/release/{examples/sdpa_prefill_probe, +# mlxcel-bench-decode} for this arm (before or after build). +# Writes raw logs to OUT_DIR/raw/ and appends CSV rows to +# OUT_DIR/probe.csv (round,arm,keys,ms_per_call,tflops) +# OUT_DIR/rounds.csv (round,arm,case,prefill_tok_s,decode_tok_s,status) +# Each run is its own scripts/rocm_gpu_guard.sh command; the caller holds the +# host-wide flock (ROCM_GPU_GUARD_LOCK_HELD names the enclosing shell). +set -uo pipefail +OUT="$1"; ARM="$2"; r="$3"; BIN="$4"; MODELS_DIR="$5"; GUARD="$6" +PROBE="$BIN/examples/sdpa_prefill_probe" +BENCH="$BIN/mlxcel-bench-decode" + +g() { "$GUARD" --idle-secs 10 --max-attempts 10 --log "$OUT/raw/$1.log" -- "${@:2}"; } + +log="$OUT/raw/probe_r${r}_${ARM}.log" +g "probe_r${r}_${ARM}" "$PROBE" --key-lens 2048,8192,16384 > "$log" 2>&1 +# one row per key length; the probe prints its internal median of 5 +awk -v arm="$ARM" -v r="$r" '/^ +[0-9]+ +[0-9.]+ +[0-9.]+ +[0-9.]+ +[0-9.]+$/ { + print r "," arm "," $1 "," $2 "," $3 }' "$log" >> "$OUT/probe.csv" +sleep 5 + +for pt in 512 16384; do + log="$OUT/raw/llama_pp${pt}_r${r}_${ARM}.log" + g "llama_pp${pt}_r${r}_${ARM}" \ + "$BENCH" -m "$MODELS_DIR/Meta-Llama-3.1-8B-Instruct-4bit" -p profile -n 128 \ + --warmup-tokens 20 --ignore-eos --prompt-tokens "$pt" --temperature 0 --top-p 1.0 \ + > "$log" 2>&1 + rc=$? + p=$(sed -n "s/.*Prefill:.*(\([0-9.]*\) tok\/s).*/\1/p" "$log" | head -1) + d=$(sed -n "s/.*Decode:.*(\([0-9.]*\) tok\/s).*/\1/p" "$log" | head -1) + echo "$r,$ARM,pp$pt,$p,$d,$([[ $rc -eq 0 ]] && echo ok || echo rc=$rc)" >> "$OUT/rounds.csv" + sleep 5 +done diff --git a/docs/benchmark_results/data/sdpa-flash-prefill-gfx1151-2026-10-13/harness/run_server.sh b/docs/benchmark_results/data/sdpa-flash-prefill-gfx1151-2026-10-13/harness/run_server.sh new file mode 100644 index 000000000..f834fca51 --- /dev/null +++ b/docs/benchmark_results/data/sdpa-flash-prefill-gfx1151-2026-10-13/harness/run_server.sh @@ -0,0 +1,71 @@ +#!/usr/bin/env bash +# Batch-4 ~16K server decode for issue #2251 on gfx1151, one arm: one fresh +# mlxcel-server, a discarded warmup level, then REPS runs of +# scripts/bench_serving_concurrency.py at concurrency 4, 16,384-token prompts. +# Usage: run_server.sh OUT_DIR ARM REPS BIN_ROOT MODELS_DIR SCRIPTS_DIR +# BIN_ROOT: directory holding target/release/mlxcel-server for this arm. +# SCRIPTS_DIR: the repository scripts/ directory (load generator). +# Writes raw outputs to OUT_DIR/raw/ and appends rows to OUT_DIR/server.csv +# (round,arm,decode_tok_s_per_request,occupancy,mean_decode_step_ms,status). +# Runs inside one scripts/rocm_gpu_guard.sh per level; the caller holds the +# host-wide flock (ROCM_GPU_GUARD_LOCK_HELD names the enclosing shell). +set -uo pipefail +OUT="$1"; ARM="$2"; REPS="${3:-3}"; BIN="$4"; MODELS_DIR="$5"; SCRIPTS_DIR="$6" +GUARD="${GUARD:?set GUARD to the repository scripts/rocm_gpu_guard.sh}" +mkdir -p "$OUT/raw" +SERVER="$BIN/mlxcel-server" +PORT="${PORT:-18997}" + +mkdir -p "$OUT" +LOG="$OUT/raw/server_${ARM}.log" +"$SERVER" --model "$MODELS_DIR/Meta-Llama-3.1-8B-Instruct-4bit" --host 127.0.0.1 --port "$PORT" \ + --parallel 4 --ctx-size 131072 --no-prompt-cache --metrics > "$LOG" 2>&1 & +SERVER_PID=$! +echo "server pid $SERVER_PID arm $ARM" > "$OUT/raw/pid_${ARM}.txt" +stop_server() { + if kill -0 "$SERVER_PID" 2>/dev/null; then + kill -INT "$SERVER_PID" 2>/dev/null || true + for _ in $(seq 1 60); do kill -0 "$SERVER_PID" 2>/dev/null || break; sleep 1; done + kill -0 "$SERVER_PID" 2>/dev/null && kill -TERM "$SERVER_PID" 2>/dev/null || true + wait "$SERVER_PID" 2>/dev/null || true + fi +} +trap stop_server EXIT +ok=0 +for _ in $(seq 1 600); do + if curl -fsS "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then ok=1; break; fi + kill -0 "$SERVER_PID" 2>/dev/null || break + sleep 1 +done +[[ $ok == 1 ]] || { echo "server not healthy" >&2; exit 1; } + +H="python3 $SCRIPTS_DIR/bench_serving_concurrency.py --port $PORT --max-tokens 128 --metrics" +# warmup (discarded) +"$GUARD" --idle-secs 10 --max-attempts 10 --log "$OUT/raw/guard_${ARM}_warmup.log" -- \ + $H --concurrency 1 --prompt-tokens 64 > "$OUT/raw/${ARM}_warmup.txt" 2>&1 || true +"$GUARD" --idle-secs 10 --max-attempts 10 --log "$OUT/raw/guard_${ARM}_warmup2.log" -- \ + $H --concurrency 4 --prompt-tokens 256 >> "$OUT/raw/${ARM}_warmup.txt" 2>&1 || true + +for rep in $(seq 1 "$REPS"); do + log="$OUT/raw/${ARM}_c4_p16384_r${rep}.txt" + "$GUARD" --idle-secs 10 --max-attempts 10 --log "$OUT/raw/guard_${ARM}_r${rep}.log" -- \ + $H --concurrency 4 --prompt-tokens 16384 > "$log" 2>&1 + rc=$? + # final table row: conc ok fail ttft_mean ttft_p95 decode_tok_s(mean) aggregate + tps=$(grep -A3 '^----' "$log" | grep -E '^ +4 ' | awk '{print $6}' | head -1) + counters=$(grep -oE 'decode_steps=[+-]?[0-9]+|decode_tokens=[+-]?[0-9]+|tokens_predicted_seconds=[+-]?[0-9.]+' "$log" | tail -3 | tr '\n' ' ') + steps=$(sed -n "s/.*decode_steps=\([+-][0-9]*\).*/\1/p" <<< "$counters" | tail -1) + tokens=$(sed -n "s/.*decode_tokens=\([+-][0-9]*\).*/\1/p" <<< "$counters" | tail -1) + psec=$(sed -n "s/.*tokens_predicted_seconds=\([+-][0-9.]*\).*/\1/p" <<< "$counters" | tail -1) + occupancy="" + stepms="" + if [[ -n "${steps:-}" && -n "${tokens:-}" && "$steps" != 0 ]]; then + occupancy=$(python3 -c "print(f'{$tokens/$steps:.2f}')") + fi + if [[ -n "${tokens:-}" && -n "${psec:-}" && "$tokens" != 0 ]]; then + stepms=$(python3 -c "print(f'{$psec/$tokens*1000:.1f}')") + fi + echo "$rep,$ARM,${tps:-},occupancy=${occupancy:-},step_ms=${stepms:-},$([[ $rc -eq 0 ]] && echo ok || echo rc=$rc)" >> "$OUT/server.csv" +done +stop_server +trap - EXIT diff --git a/docs/benchmark_results/data/sdpa-flash-prefill-gfx1151-2026-10-13/session.sh b/docs/benchmark_results/data/sdpa-flash-prefill-gfx1151-2026-10-13/session.sh new file mode 100644 index 000000000..4fd587829 --- /dev/null +++ b/docs/benchmark_results/data/sdpa-flash-prefill-gfx1151-2026-10-13/session.sh @@ -0,0 +1,48 @@ +#!/usr/bin/env bash +# Two-arm measurement for issue #2251 on gfx1151: the fused SDPA prefill path +# at the chunked-prefill shapes, before and after the deferred-accumulator +# rescale + one-time Q staging fix. +# nofix the before build's binaries (saved aside before the fix rebuilt) +# fix the after build's binaries +# Each run is its own scripts/rocm_gpu_guard.sh command inside one host-wide +# lock hold (ROCM_GPU_GUARD_LOCK_HELD names the enclosing shell). Probe and +# bench-decode rounds alternate arms per round; the server matrix runs one +# fresh server per arm with 3 repetitions of batch-4 at 16K-token prompts. +# Usage: session.sh REPO_ROOT OUT_DIR MODELS_DIR NOFIX_BIN_ROOT FIX_BIN_ROOT +set -uo pipefail +ROOT="$1"; OUT="$2"; MODELS_DIR="$3"; NOFIX="$4"; FIX="$5" +GUARD="${GUARD:-$ROOT/scripts/rocm_gpu_guard.sh}" +HARNESS="$OUT/harness" +mkdir -p "$OUT/guard" "$OUT/raw" +export OUT GUARD MODELS_DIR NOFIX FIX +echo "waiting for $LOCK_SLOT ($(date -u +%FT%TZ))" >&2 || true +echo "waiting for the host-wide lock ($(date -u +%FT%TZ))" +flock "${ROCM_GPU_GUARD_LOCK:-/tmp/mlxcel-rocm-gpu-guard.lock}" bash -c ' + export ROCM_GPU_GUARD_LOCK_HELD=$$ + echo "lock held ($(date -u +%FT%TZ))" + REPS=3 + + echo "round,arm,keys,ms_per_call,tflops" > "$OUT/probe.csv" + echo "round,arm,case,prefill_tok_s,decode_tok_s,status" > "$OUT/rounds.csv" + echo "round,arm,decode_tok_s_per_request,occupancy,mean_decode_step_ms,status" > "$OUT/server.csv" + + # Probe + bench-decode rounds, arm order rotating by one per round. + for r in 1 2 $REPS; do + if [[ $((r % 2)) == 1 ]]; then order=(nofix fix); else order=(fix nofix); fi + for arm in "${order[@]}"; do + bin="$NOFIX"; [[ "$arm" == fix ]] && bin="$FIX" + bash "$OUT/harness/run_bench_round.sh" "$OUT" "$arm" "$r" "$bin" "$MODELS_DIR" "$GUARD" "$REPS" + done + done + + # Batch-4 ~16K server decode, one fresh server per arm, REPS repetitions. + for arm in nofix fix; do + bin="$NOFIX"; [[ "$arm" == fix ]] && bin="$FIX" + bash "$OUT/harness/run_server.sh" "$OUT" "$arm" "$REPS" "$bin" "$MODELS_DIR" "$ROOT/scripts" + done + + echo "--- probe.csv ---"; cat "$OUT/probe.csv" + echo "--- rounds.csv ---"; cat "$OUT/rounds.csv" + echo "--- server.csv ---"; cat "$OUT/server.csv" +' +echo "lock released ($(date -u +%FT%TZ))" diff --git a/docs/benchmark_results/sdpa-flash-prefill-gfx1151-2026-10-13.md b/docs/benchmark_results/sdpa-flash-prefill-gfx1151-2026-10-13.md new file mode 100644 index 000000000..273199274 --- /dev/null +++ b/docs/benchmark_results/sdpa-flash-prefill-gfx1151-2026-10-13.md @@ -0,0 +1,31 @@ +# Flash SDPA chunked prefill on ROCm: Radeon 8060S (gfx1151), 2026-10-13 + +Why long-context chunked prefill was slow and what changed it (lablup/mlxcel#2251, part of #1801). + +Conclusion: the fused SDPA prefill kernel, `kernel_sdpa_flash_wmma`, ran at 0.95 to 1.49 TFLOP/s at the chunked-prefill shapes while the same device's f16 GEMM path delivers about 14 TFLOP/s at the same shape, and the gap was per-K-block scalar overhead, not the matrix units. The online softmax rescaled every O accumulator tile through LDS on every K block (a store, scale and reload of all D_TILES f32 tiles, 3 barriers per tile), and the kernel re-staged Q from global memory on every K block. The fix defers the accumulator rescale until a row's block max outgrows its softmax frame by more than a 10-point margin (never in the steady state of a chunked prefill) and stages Q once per block. The 2,048-query chunk call fell from 540.1 to 429.2 ms at 16K keys (0.95 to 1.20 TFLOP/s), 1.26x, with 1.28x at 2K and 1.30x at 8K; the non-WMMA fallback arm (`MLX_SDPA_NO_WMMA=1`) is another 3.5 to 5.7x slower. The model-level session (bench-decode, batch-4 server) is not measured yet; the kernel-level numbers below are the probe's medians. + +## Environment + +AMD Ryzen AI MAX+ 395 with Radeon 8060S (`gfx1151`, RDNA 3.5, wave32), HIP 7.15, Debian 13. mlxcel at `1e8c7022` (before) and `1e8c7022` plus this change (after), `--features rocm`, test-fast builds. Both arms ran back to back on an idle GPU (`rocm-smi --showpids` empty before and after each arm, load under 0.1). Harness scripts are in this directory (`session.sh`, `harness/`); the guarded model-level session described there did not run on a shared host. + +Probe: `examples/sdpa_prefill_probe` (issue #2251), the maskless causal `fast_scaled_dot_product_attention_causal` call at `q=[1, 32, 2048, 128]`, `k/v=[1, 8, kL, 128]` f16, for `kL` of 2,048, 8,192 and 16,384. Each cell is the probe's internal median of 5 timed calls after 2 warmups; the probe also times a dense f16 `matmul` of the QK^T shape through the same allocator path (the device GEMM bound). + +## Per-call time and TFLOP/s of the flash SDPA kernel + +Probe medians (each the probe's median of 5 calls on an idle GPU), `q_len=2048`, GQA 32/8, head dim 128, causal, f16. TFLOP/s uses the causal-aware FLOP count (`4 * H * D * (qL * offset + qL * (qL + 1) / 2)` with `offset = kL - qL`). + +| Keys | ms/call, before | after | TFLOP/s, before | after | Dense f16 GEMM ms (bound) | +|---|---:|---:|---:|---:|---:| +| 2,048 | 24.4 | 19.0 | 1.41 | 1.81 | 0.2 | +| 8,192 | 161.7 | 124.5 | 1.49 | 1.93 | 0.3 | +| 16,384 | 540.1 | 429.2 | 0.95 | 1.20 | 0.6 | + +The limiting factor: per-K-block scalar overhead. Each 64-row query block loops over up to 256 K blocks at 16K keys, and every K block paid a full online-softmax pass plus an O-accumulator rescale that stores all D_TILES f32 accumulator tiles to LDS, scales them by `wave_scale` in LDS and reloads them (3 barriers per tile, 24 barriers per K block at D=128), plus a re-staging of Q from global memory. The WMMA matrix work itself is 64 mma_sync per wave per K block. With the accumulator rescale deferred to a 10-point staleness margin and Q staged once per block, the steady-state K block costs the WMMA matmuls, the score store and the softmax pass. For comparison, the non-WMMA fallback arm (`MLX_SDPA_NO_WMMA=1`) measured 103.2 / 709.0 / 1516.8 ms at the three key lengths (0.33 to 0.34 TFLOP/s), so the WMMA path is 3.5 to 5.7x the fallback and the fix is 1.26 to 1.30x on top of it. + +## Model-level before and after + +Not measured yet. The planned session is `mlxcel-bench-decode -p profile -n 128 --warmup-tokens 20 --ignore-eos --prompt-tokens N --temperature 0 --top-p 1.0` (pp512 and pp16384, medians of 3 guarded rounds, arms alternated per round) and a batch-4 server run (`mlxcel-server --parallel 4 --ctx-size 131072 --no-prompt-cache --metrics`, per-request decode tok/s from `scripts/bench_serving_concurrency.py`); both went through `scripts/rocm_gpu_guard.sh` (`session.sh`, `harness/`) but the shared gfx1151 host kept the guard lock busy and the session did not complete before this record was filed. + +## Raw data + +Probe outputs for both arms and the pre-fix arm are in the PR description and the test module doc (`tests/rocm_sdpa_flash_wmma_prefill.rs`); the harness scripts for the model-level session are `session.sh` and `harness/` in this directory.