perf(rocm): defer SDPA accumulator rescale, stage Q once per block

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

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

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))"

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.

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"]