Repository navigation
perf(rocm): defer SDPA accumulator rescale, stage Q once per block - #2306
Merged
Merged
Conversation
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
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
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
inureyes
force-pushed
the
perf/issue-2251-flash-sdpa-prefill
branch
from
October 14, 2026 04:30
818144f to
661dc19
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Defer the WMMA flash-attention O-accumulator rescale so it runs only when a row's block max outgrows its softmax frame by more than a 10-point margin, and stage Q once per block instead of on every K block (issue #2251). The 16K-key chunked-prefill call on gfx1151 measured 540.1 ms/call (0.95 TFLOP/s) before the change and 429.2 ms/call (1.20 TFLOP/s) after, 1.26x, with 1.28x at 2K keys and 1.30x at 8K; the non-WMMA fallback arm is another 3.5 to 5.7x slower.
What changed
src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/flash_attention_wmma.hip: the O tiles and l of a row stay in the softmax frame the row was last rescaled to, so the LDS store/scale/reload of all accumulator tiles (3 barriers per tile, 24 barriers per K block at D=128) runs only when a row's block max outgrows its frame by more thankFaWmmaRescaleMargin(10.0), never in the steady state of a chunked prefill; P = exp(S - f) stays bounded by exp(10) = 22,026, inside the bf16 score cap. Q is staged once per block when the head dim fits one DW-wide slice (D <= 128).tests/rocm_sdpa_flash_wmma_prefill.rs: causal correctness against an explicit f32 reference at the 2K and 16K key shapes for f16 and bf16, a score ramp that drives rows past the rescale margin, 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).src/lib/mlxcel-core/examples/sdpa_prefill_probe.rs: probe timing the maskless causal call at the three key lengths against a dense f16 GEMM bound;docs/benchmark_results/sdpa-flash-prefill-gfx1151-2026-10-13.mdrecords the numbers.Measurements
rocm-smi --showpidsempty and load under 0.1 before and after each arm). Before: 24.4 / 161.7 / 540.1 ms (1.41 / 1.49 / 0.95 TFLOP/s); after: 19.0 / 124.5 / 429.2 ms (1.81 / 1.93 / 1.20 TFLOP/s) at 2K / 8K / 16K keys.MLX_SDPA_NO_WMMA=1fallback arm measured 103.2 / 709.0 / 1516.8 ms (0.33 to 0.34 TFLOP/s). A gate run underscripts/rocm_gpu_guard.sh --holdmeasured the fixed 16K call at 392.6 ms/call (1.31 TFLOP/s).Test plan
make verify-rocm-heldon the gfx1151 host: 12,316 tests passed, 0 failed, 404 ignored,[verify-rocm] OKcargo test --profile test-fast --features rocm --test rocm_sdpa_flash_wmma_prefill -- --test-threads=1: all three tests passcargo clippy --workspace --all-targets --features rocm -- -D warningsclean,cargo fmtcleanCloses #2251