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. 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/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( 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 + ); +}