feat(hipdnn): add HipFlash2Engine Flash-Attention 2 V7 SDPA for gfx942/gfx950 - #9589
Conversation
…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
…ed .co (AIOSS-5319)
…execute() via HipModuleGuard
…ute() via HipModuleGuard
…ute() via HipModuleGuard
…to execute() via HipModuleGuard
… into execute() via HipModuleGuard
…to execute() via HipModuleGuard
…der.cpp to target_sources
…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
left a comment
There was a problem hiding this comment.
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.
5b1ae3b to
23e1495
Compare
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.
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.
c98921d to
ca2212a
Compare
… users/zhihuidu-amd/hip-flash2-sdpa-engine-v3
| - path: Small.tensor1.bin | ||
| md5: d4920f88fb5c8c942c53be0c428b3c24 | ||
| size: 16777216 | ||
| hash: md5 |
There was a problem hiding this comment.
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)
|
Thanks for catching this — confirmed, and fixed in #11481. Those nine files were the only content left in There's a second consequence worth flagging: the bundles were never being loaded at all. The fix is a pure DVC is unaffected — each |
…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.
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
HipFlash2Engineas 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
exp2fvia__builtin_amdgcn_exp2fon gfx942 and standardexp2fon gfx950 (portable guard avoids CDNA4 rounding divergence).co(hsaco) binaries for gfx942 and gfx950 — no rocWMMA build-time dependency for downstream users; TheRock does DVC checkout before buildHipFlash2Engine→HipFlash2FwdPlanBuilder→HipFlash2FwdPlan, implementing the standardIEngine/IPlanBuilder/IPlanplugin pattern.cofiles installed to${CMAKE_INSTALL_PREFIX}/${HIPDNN_RELATIVE_INSTALL_PLUGIN_ENGINE_DIR}/; overridable at runtime viaHIP_FLASH2_KERNEL_DIRenv varAll overlay review findings addressed: B1–B5, I1–I9.
Test Plan
TestHipFlash2Engine,TestHipFlash2FwdPlanBuilder(shape acceptance, rejection, GQA divisibility)IntegrationGpuHipFlash2Forward— 11 shapes on real MI325X hardware viaIntegrationGraphVerificationHarnessTest 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):
Peak throughput: 78.98 TFLOPS on MI325X (seq=4096, causal).