Skip to content

feat(hipdnn): add HipFlash2Engine Flash-Attention 2 V7 SDPA for gfx942/gfx950 - #9589

Merged
zhihuidu-amd merged 237 commits into
ROCm:developfrom
zhihuidu-amd:users/zhihuidu-amd/hip-flash2-sdpa-engine-v3
Aug 18, 2026
Merged

feat(hipdnn): add HipFlash2Engine Flash-Attention 2 V7 SDPA for gfx942/gfx950#9589
zhihuidu-amd merged 237 commits into
ROCm:developfrom
zhihuidu-amd:users/zhihuidu-amd/hip-flash2-sdpa-engine-v3

Conversation

@zhihuidu-amd

@zhihuidu-amd zhihuidu-amd commented Jul 19, 2026

Copy link
Copy Markdown
Contributor

Motivation

hipDNN's SDPA execution path falls back to unfused Q·Kᵀ + softmax + P·V GEMMs for prefill, which is 8–18× slower than a fused Flash-Attention kernel on MI300X/MI325X. This PR adds HipFlash2Engine as a precompiled SDPA engine plugin so hipDNN can select Flash-Attention 2 automatically for eligible shapes on gfx942 and gfx950.

JIRA ID : AIOSS-5319

Technical Details

  • Kernel: V7 Flash-Attention 2 — rocWMMA MFMA-tiled QK^T and P@V, causal tile skip (skips tiles entirely below the diagonal), fast exp2f via __builtin_amdgcn_exp2f on gfx942 and standard exp2f on gfx950 (portable guard avoids CDNA4 rounding divergence)
  • Delivery format: Precompiled .co (hsaco) binaries for gfx942 and gfx950 — no rocWMMA build-time dependency for downstream users; TheRock does DVC checkout before build
  • Engine interface: HipFlash2EngineHipFlash2FwdPlanBuilderHipFlash2FwdPlan, implementing the standard IEngine/IPlanBuilder/IPlan plugin pattern
  • Applicability: FP16 only, head_dim ∈ {64, 128}, causal/non-causal, GQA (Hq divisible by Hkv), sequence ≥ crossover threshold, no dropout/alibi/paged-attn
  • Install path: .co files installed to ${CMAKE_INSTALL_PREFIX}/${HIPDNN_RELATIVE_INSTALL_PLUGIN_ENGINE_DIR}/; overridable at runtime via HIP_FLASH2_KERNEL_DIR env var

All overlay review findings addressed: B1–B5, I1–I9.

Test Plan

  • Unit tests: TestHipFlash2Engine, TestHipFlash2FwdPlanBuilder (shape acceptance, rejection, GQA divisibility)
  • GPU integration tests: IntegrationGpuHipFlash2Forward — 11 shapes on real MI325X hardware via IntegrationGraphVerificationHarness
  • Golden reference data generated on MI325X (FP16), stored via DVC, verified against CPU FP32 reference
  • ASAN clean under hipdnn enroot image on Alola

Test Result

Correctness: MaxErr < 0.002 on all 9 golden-ref shapes vs CPU FP32 reference.

Performance (MI325X, FP16, causal, D=128, apple-to-apple same buffers):

B Sq Skv H Flash2 (ms) Unfused (ms) Speedup
1 512 512 32 0.041 0.337 8.2×
1 1024 1024 32 0.089 1.342 15.1×
1 2048 2048 32 0.301 5.365 17.8×
4 512 512 32 0.155 1.330 8.6×
4 1024 1024 32 0.351 5.310 15.1×
4 2048 2048 32 1.203 21.33 17.7×
8 1024 1024 32 0.704 10.61 15.1×

Peak throughput: 78.98 TFLOPS on MI325X (seq=4096, causal).

zhihuidu-amd and others added 30 commits June 16, 2026 12:48
…942+gfx950

Add HipFlash2Engine to hip-kernel-provider for FP16 SDPA on gfx942 and gfx950.
Implements Flash-Attention 2 V7: full rocWMMA MFMA + causal tile skip.

Performance (real hardware, FP16, seq=4096 causal D=128):
  MI300X: 71.27 TFLOPS (+8.1x vs unfused, OCI job 35381)
  MI325X: 78.98 TFLOPS (+8.1x vs unfused, Alola job 366655)

Correctness: MaxErr < 0.002 on 9/9 shapes vs CPU FP32 reference.

gfx950 (MI355X/MI350X) portability: flash2_exp2f() with #if __gfx942__ guard
replaces __builtin_amdgcn_exp2f which has different rounding on CDNA4.

Jira: AIOSS-5319
Address review comment by Daryl Hawkins: include actual file modifications
instead of a patch-as-comment file.

Changes:
- EngineNames.hpp: register HIP_FLASH2_ENGINE with HIPDNN_REGISTER_ENGINE()
- Container.cpp: add #include guards + engine factory for HipFlash2Engine
- CMakeLists.txt: add ENABLE_HIP_FLASH2_ENGINE option (ON by default)
- Remove Container_patch.diff (replaced by actual file edits above)

Build: cmake -DENABLE_HIP_FLASH2_ENGINE=ON ...
       hipcc will compile HipFlash2FwdPlan.hip with --offload-arch=gfx942;gfx950

Jira: AIOSS-5319
1. CMakeLists.txt: move ENABLE_HIP_FLASH2_ENGINE block outside of
   ENABLE_ASM_SDPA_ENGINE conditional block (fixes superbuild CI)

2. Apply clang-format (WebKit style, 4-space indent, 100 col limit)
   to all new/modified C++ files (fixes pre-commit CI)

Jira: AIOSS-5319
Apply WebKit-based clang-format style (4-space indent, 100 col limit)
to pass pre-commit CI check.

Files formatted:
  - HipFlash2Engine.hpp
  - HipFlash2FwdPlan.hip
  - HipFlash2FwdPlanBuilder.hpp
  - bench_hipdnn_sdpa.cpp
  - Container.cpp
  - EngineNames.hpp

Jira: AIOSS-5319
… OFF

Critical build fix: CMakeLists.txt referenced HipFlash2Engine.cpp in
target_sources() but the file was missing, causing all TheRock/superbuild
CI jobs to fail with 'custom container implementation failed'.

Changes:
1. ADD HipFlash2Engine.cpp — implementation of IEngine interface methods
   (id, isApplicable, getDetails, getMaxWorkspaceSize, initializeExecutionContext)
   following AsmSdpaEngine pattern exactly.

2. UPDATE HipFlash2Engine.hpp — use HIP_FLASH2_ENGINE_ID from EngineNames.hpp
   (FNV-1a hash of engine name) instead of hardcoded 0x4841544E32ULL.
   Match AsmSdpaEngine constructor/interface pattern.

3. UPDATE README.md — remove redundant sections (integration steps, files
   list, dispatch logic, PR target). Add Known Issues for gfx950 softmax
   reduction divergence under investigation with Brian Harrison.

4. FIX CMakeLists.txt — change ENABLE_HIP_FLASH2_ENGINE default ON->OFF.
   New engines should be disabled until fully validated (Brian Harrison).

Jira: AIOSS-5319
…ot initializeExecutionContext), match AsmSdpaEngine pattern
…stride checks

Finding B: src/CMakeLists.txt file(COPY) block still listed
hip_flash2_fwd_gfx950.co (line 65). Removed -- only gfx942.co exists on branch.

Finding C: isApplicable() only checked Q head_dim stride==1. The kernel
assumes contiguous innermost dimension for all four tensors (Q, K, V, O).
Extended the stride==1 guard to K, V, and O so a graph with strided K/V/O
is correctly rejected rather than silently producing wrong output.
Rebuild command: hipcc --genco -O3 --offload-arch=gfx942 -I/opt/rocm/include HipFlash2FwdPlan.hip
The --genco flag is the hipcc-documented alias for --cuda-device-only and
matches what Samuel used for verification. Correctness validated against
CPU FP32 reference and invariant probes on MI300X.

@SamuelReeder SamuelReeder left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

Validated on MI300X with -DENABLE_HIP_FLASH2_ENGINE=ON -DHIPDNN_ENABLE_SDPA=ON:

  • Build and configure clean
  • Unit: 586 pass / 1 skip (unrelated); Flash2's own 15/15
  • Flash2 integration: 6/6, executing the kernel against a CPU reference
  • Shipped .co run directly vs CPU FP32: maxErr 5.3e-5 to 3.8e-4, no NaN, across d64/d128, causal and non-causal, S=64–2048

One caveat on that integration result: it only passes with HIP_FLASH2_KERNEL_DIR set manually. The kernel path resolves to CMAKE_INSTALL_PREFIX, so a
build-tree run without installing gets 6 failures on hipModuleLoad (error 301, file not found). The ASM SDPA engine handles this with
set_tests_properties(... ENVIRONMENT ...).

Nice work on the Sᵀ reformulation!

The committed hip_flash2_fwd_gfx942.co did not match what HipFlash2FwdPlan.hip
compiles to on a current toolchain, and could not be reproduced with any ROCm
available for testing. Regenerated from the unmodified source.

Measured on MI300X (gfx942), MHA seq=4096 causal D=128:

  toolchain      d128 vgpr   TFLOPS
  ROCm 6.4       256          49.8
  ROCm 7.1.1     512          60.7   <- matches the previously committed binary
  ROCm 10.0.0    340          83.8   <- this commit

512 is the architectural maximum for a gfx942 wave (256 arch VGPRs + 256 AGPRs),
so the old build had the register allocator against the ceiling. Note the ranking
does not follow register count: ROCm 6.4 uses the fewest VGPRs and is slowest.

Kernel source is unchanged; both binaries issue the same 256 v_mfma_f32_16x16x16_f16
for d128. The difference is register allocation, not algorithm.

Launch ABI is unchanged and verified identical:
  kernarg_segment_size 112, max_flat_workgroup_size 64,
  LDS 32768 (d128) / 16384 (d64), symbols flash2_v7_hipdnn_d{64,128}

Validated on gfx942 against a CPU FP32 reference across 28 configurations
(d64 and d128, S=64..2048, causal and non-causal): all MaxErr < 0.002, no NaN.
Provider unit suite 586 pass / 1 skip, Flash2 unit 15/15, Flash2 integration 6/6.

README documents the build command, the toolchain used, and how to validate a
regenerated binary, so the artifact cannot silently drift from the source again.
@zhihuidu-amd
zhihuidu-amd force-pushed the users/zhihuidu-amd/hip-flash2-sdpa-engine-v3 branch from 5b1ae3b to 23e1495 Compare August 14, 2026 23:02
@zhihuidu-amd
zhihuidu-amd requested a review from a team as a code owner August 17, 2026 21:27
zhihuidu-amd added a commit to zhihuidu-amd/rocm-libraries that referenced this pull request Aug 17, 2026
Resolves the conflicts blocking ROCm#9589 (mergeable_state was "dirty").

Upstream landed the kernel-ingestor engine while this branch was open, and
its HIPDNN_ENABLE_KERNEL_INGESTOR blocks sit exactly where our
ENABLE_HIP_FLASH2_ENGINE blocks are -- adjacent if()/endif() pairs, so git
could not tell which endif() belonged to which. Both engines are wanted;
none of the five conflicts was a semantic disagreement.

Resolution per file:

  src/CMakeLists.txt
  src/tests/CMakeLists.txt
  src/integration_tests/CMakeLists.txt
      Keep both blocks -- our flash2 if()/endif(), then theirs. Verified
      with cmake-lint: clean.

  src/core/Container.cpp
      Took THEIRS. Upstream restructured the engine list from a flat
      initializer into a lambda that push_back()s and returns the
      definitions vector, so our side of the hunk was the obsolete form.
      Our flash2 entry sits above the conflict region and survives
      unchanged inside the new list; the ingestor's discovery loop follows
      it.

  src/tests/core/TestContainer.cpp
      Took THEIRS, then re-added our arm. Upstream replaced the inline
      engine-count expression with an expectedEngines() function. Taking
      theirs alone would have dropped flash2 from the count and made
      CopyEngineIdsReturnsExpectedEngineCount fail once the engine is
      enabled, so HIPDNN_ENGINE_HIP_FLASH2 now increments inside that
      function alongside ASM_SDPA and HIP_MLOPS.

Verified: no conflict markers remain; diffing the five files against
upstream develop shows only our flash2 additions and no deletion of
upstream work; Container.cpp compiles under hipcc -fsyntax-only -std=c++17
with HIPDNN_ENGINE_HIP_FLASH2 and HIPDNN_ENGINE_ASM_SDPA defined;
cmake-lint clean on all three CMakeLists; clang-format applied.

Note the engine still is not built by CI (ENABLE_HIP_FLASH2_ENGINE defaults
OFF and is set nowhere), so these checks were run by hand.
@github-actions github-actions Bot added rocke-lib: attention Touches rocKE library attention/FMHA kernels rocke: platform Touches rocKE platform labels Aug 17, 2026
Resolves the conflicts blocking ROCm#9589 (mergeable_state was "dirty").

Upstream landed the kernel-ingestor engine while this branch was open, and
its HIPDNN_ENABLE_KERNEL_INGESTOR blocks sit exactly where our
ENABLE_HIP_FLASH2_ENGINE blocks are -- adjacent if()/endif() pairs, so git
could not tell which endif() belonged to which. Both engines are wanted;
none of the five conflicts was a semantic disagreement.

Resolution per file:

  src/CMakeLists.txt
  src/tests/CMakeLists.txt
  src/integration_tests/CMakeLists.txt
      Keep both blocks -- our flash2 if()/endif(), then theirs. Verified
      with cmake-lint: clean.

  src/core/Container.cpp
      Took THEIRS. Upstream restructured the engine list from a flat
      initializer into a lambda that push_back()s and returns the
      definitions vector, so our side of the hunk was the obsolete form.
      Our flash2 entry sits above the conflict region and survives
      unchanged inside the new list; the ingestor's discovery loop follows
      it.

  src/tests/core/TestContainer.cpp
      Took THEIRS, then re-added our arm. Upstream replaced the inline
      engine-count expression with an expectedEngines() function. Taking
      theirs alone would have dropped flash2 from the count and made
      CopyEngineIdsReturnsExpectedEngineCount fail once the engine is
      enabled, so HIPDNN_ENGINE_HIP_FLASH2 now increments inside that
      function alongside ASM_SDPA and HIP_MLOPS.

Verified: no conflict markers remain; diffing the five files against
upstream develop shows only our flash2 additions and no deletion of
upstream work; Container.cpp compiles under hipcc -fsyntax-only -std=c++17
with HIPDNN_ENGINE_HIP_FLASH2 and HIPDNN_ENGINE_ASM_SDPA defined;
cmake-lint clean on all three CMakeLists; clang-format applied.

Note the engine still is not built by CI (ENABLE_HIP_FLASH2_ENGINE defaults
OFF and is set nowhere), so these checks were run by hand.
@zhihuidu-amd
zhihuidu-amd force-pushed the users/zhihuidu-amd/hip-flash2-sdpa-engine-v3 branch from c98921d to ca2212a Compare August 17, 2026 21:40
@zhihuidu-amd
zhihuidu-amd enabled auto-merge (squash) August 18, 2026 19:03
Comment on lines +7 to +10
- path: Small.tensor1.bin
md5: d4920f88fb5c8c942c53be0c428b3c24
size: 16777216
hash: md5

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This integration_test_bundles directory was renamed to integration-test-bundles back in #10162, and the new DVC .bin files are not excluded by a .gitignore file after that migration

D:\projects\TheRock\rocm-libraries (users/scotttodd/hipthreads-test-requirements)
λ git status
HEAD detached at b3c527c146f
Untracked files:
  (use "git add <file>..." to include in what will be committed)
        dnn-providers/integration-tests/integration_test_bundles/quick/SdpaFwd/bhsd/fp16/hd128_causal_gqa4/Small/Small.tensor0.bin
        dnn-providers/integration-tests/integration_test_bundles/quick/SdpaFwd/bhsd/fp16/hd128_causal_gqa4/Small/Small.tensor1.bin
        dnn-providers/integration-tests/integration_test_bundles/quick/SdpaFwd/bhsd/fp16/hd128_causal_gqa4/Small/Small.tensor2.bin
        dnn-providers/integration-tests/integration_test_bundles/quick/SdpaFwd/bhsd/fp16/hd128_causal_gqa4/Small/Small.tensor3.bin
        dnn-providers/integration-tests/integration_test_bundles/quick/SdpaFwd/bhsd/fp16/hd128_causal_mha/Small/Small.tensor0.bin
        dnn-providers/integration-tests/integration_test_bundles/quick/SdpaFwd/bhsd/fp16/hd128_causal_mha/Small/Small.tensor1.bin
        dnn-providers/integration-tests/integration_test_bundles/quick/SdpaFwd/bhsd/fp16/hd128_causal_mha/Small/Small.tensor2.bin
        dnn-providers/integration-tests/integration_test_bundles/quick/SdpaFwd/bhsd/fp16/hd128_causal_mha/Small/Small.tensor3.bin
        dnn-providers/integration-tests/integration_test_bundles/quick/SdpaFwd/bhsd/fp16/hd64_causal/Small/Small.tensor0.bin
        dnn-providers/integration-tests/integration_test_bundles/quick/SdpaFwd/bhsd/fp16/hd64_causal/Small/Small.tensor1.bin
        dnn-providers/integration-tests/integration_test_bundles/quick/SdpaFwd/bhsd/fp16/hd64_causal/Small/Small.tensor2.bin
        dnn-providers/integration-tests/integration_test_bundles/quick/SdpaFwd/bhsd/fp16/hd64_causal/Small/Small.tensor3.bin

See #10162 (comment)

@zhihuidu-amd

Copy link
Copy Markdown
Contributor Author

Thanks for catching this — confirmed, and fixed in #11481.

Those nine files were the only content left in integration_test_bundles/. #9589 was in flight across the #10162 rename and recreated the old path.

There's a second consequence worth flagging: the bundles were never being loaded at all. dnn-providers/integration-tests/CMakeLists.txt:133 copies only integration-test-bundles/ into the build tree and installs only that at :149, and both BundleRegistration.hpp:214 and main.cpp:134 resolve <exe>/../lib/integration-test-bundles. Nothing in the tree references the underscore path, so those three Flash2 golden-reference bundles have been inert since the rename — the untracked .bin files were the visible symptom of a dead directory.

The fix is a pure git mv: 9 renames, 0 insertions, 0 deletions. No code change needed, since the harness discovers bundles by scanning rather than from a manifest, and none of the three names already existed under the new path.

DVC is unaffected — each .tensors.dvc entry addresses its blob by md5 with path holding just the bare filename relative to the pointer, so moving the pointer doesn't change what dvc pull fetches. Verified after the move that the pointer contents are byte-identical and that git check-ignore now matches the .bin files against integration-test-bundles/.gitignore:4.

zhihuidu-amd added a commit that referenced this pull request Aug 31, 2026
…10977)

Follow-up to #9589. That PR landed the engine with a single kernel; this
adds per-shape variant selection, the parameterized source the variants
are built from, and the tests and probes for both.

WHY MORE THAN ONE VARIANT
No single tiling is fastest across attention shapes. Measured on MI300X
(gfx942, FP16) over 21 shapes against the CK/AITER backend, the best
SINGLE variant reaches 0.66x while selecting per shape reaches parity
and above. An exhaustive 18-variant x 21-shape sweep put the oracle
ceiling at 0.998x, so almost all of the gap was selection rather than
kernel quality.

Re-measured with our kernel built on ROCm 10.1 (TheRock
10.1.0a20260814), 5 alternating rounds x 3 reps x 20 iters, per-shape
medians: geomean 1.081x, median 1.104x, 16/21 wins, per-shape CV
0.14-1.20%.

One shape runs 0.478x: B1H8S2048D128 noncausal. That is the shape the
rule selects split-K for, and split-K is not yet wired through
execute(), so it runs single-pass. Excluding it the geomean is 1.126x,
but the honest headline is 1.081x -- the higher number describes code
that does not exist yet.

SELECTION KEY
The number of CTAs the shape produces, compared against the device CU
count, NOT the sequence length. Shapes that underperformed with a
sequence-length rule had nothing in common in S; they simply produced
too few CTAs to fill the GPU once a CTA is 512 threads.
ceil(S/qPerCta)*B*H captures batch, heads and sequence length together.

WHAT IS HERE
  Flash2Dispatch.hpp          the rule, as a header-only pure function
HipFlash2FwdPlanVariant.hip the parameterized kernel source (F2_WAVES /
F2_QG / F2_KG / F2_PIPE / F2_PAD / F2_PADK /
                              F2_VTMAJOR / F2_SKIPRESCALE)
CMakeLists.txt hip_flash2_variants target to regenerate the
                              .co files, off by default
  TestFlash2Dispatch.cpp      8 host-only tests, no GPU required
  tools/partial_tile_probe.cpp correctness probe for partial query tiles

All five shipped variants differ only in WAVES/QG/KG; every other knob
sits at the source default. Verified rather than recalled: rebuilding at
the defaults reproduces the shipped objects' LDS footprint exactly
(w8q3k2 18944, w8q2k4 / w8q1k4 / w8q3k4 35840).

GEOMETRY IS CHECKED AT LOAD
buildPlan() queries HIP_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK after
loading a variant and hard-errors if it disagrees with the table. The
failure is asymmetric and that is why it throws rather than warns:
too-small is a loud hipError 719, but too-large computes silently wrong
results. Thanks to S. Reeder for demonstrating the first case by copying
the legacy 64-thread object over the five variant names -- the suite
went 6/6 to 0/6.

INERT UNTIL VARIANTS SHIP
buildPlan() probes for hip_flash2_fwd_<arch>_<tag>.co and falls back to
the legacy single-kernel object with 64/64 geometry when absent, so
behaviour is unchanged from #9589 until a variant set is installed.
Variants MUST be built with the same ROCm as the shipped default: 10.0.0
is ~1.4x faster than 7.2 for identical source (S=2048 D=128 causal, 54.1
vs 38.3 TFLOPS), so a mixed set would look like an upgrade while making
most shapes slower.

PARTIAL QUERY TILES ARE SAFE -- MEASURED
isApplicable() enforces seq_len_q % 64 == 0 because a partial tile makes
__syncthreads() divergent, but that gate predates 384-query tiles:
S=2048 is a multiple of 64 and not of 384. Measured on MI300X, w8q3k2 at
S in {1024, 2048, 2560, 4096}, causal and not: constant-V invariant
exactly 0.000e+00, Q=0 -> mean(V) <= 1.6e-05, vs CPU fp32 <= 8.2e-05 --
identical to the exact-multiple rows. Negative control included, since a
test that cannot fail proves nothing: mis-declaring w8q2k4 as a
384-query tile under-covers the grid and fails every row with constV
5.000e-01.

VERIFICATION
CI does not build this engine (ENABLE_HIP_FLASH2_ENGINE defaults OFF and
is set nowhere), so the following was run by hand against this branch:
- hipcc -fsyntax-only -std=c++17 -Wall -Wextra on both changed TUs:
clean
- clang-tidy 20 vs the provider .clang-tidy: 7 findings, identical to
clean develop. This branch adds none.
  - TestFlash2Dispatch: 8/8, covering all 21 measured selections
  - cmake-lint: clean
- hip_flash2_variants target builds all five .co end-to-end on gfx942,
and each object's max_flat_workgroup_size matches the table's blockDim

FOLLOW-UPS
- wire split-K through execute() (needs a merge launch plus a workspace
pointer via IPlan); worth ~4.5 points of geomean on its own
- set_tests_properties(... ENVIRONMENT ...) so integration tests find
the .co from the build tree without exporting HIP_FLASH2_KERNEL_DIR, per
S. Reeder's note on #9589

## Motivation

<!-- Explain the purpose of this PR and the goals it aims to achieve.
-->

## Technical Details

<!-- Explain the changes along with any relevant GitHub links. -->

## Test Plan

<!-- Explain any relevant testing done to verify this PR. -->

## Test Result

<!-- Briefly summarize test outcomes. -->

## Submission Checklist

- [ ] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants