feat(hip-flash2): execute split-K instead of only selecting it (AIOSS-5319) - #11482
Open
zhihuidu-amd wants to merge 1 commit into
Open
feat(hip-flash2): execute split-K instead of only selecting it (AIOSS-5319)#11482zhihuidu-amd wants to merge 1 commit into
zhihuidu-amd wants to merge 1 commit into
Conversation
…-5319) Follow-up to ROCm#10977. The dispatcher already decided when to split the KV axis, but nothing ran it: buildPlan() forced params.splitK = 1 and execute() launched a single kernel, so the decision was advisory. ROCm#10977 went further and gated the credit behind K_FLASH2_SPLITK_EXECUTES = false, so selection would describe what actually happened rather than what it wished happened. This wires it up, and flips that flag. WHY B1H8S2048 non-causal was the only shape in the 21-shape suite where we lose to the CK/AITER backend -- 0.478x -- and it is exactly the shape split-K exists for. Measured on MI300X, this shape alone: nsplit=1 0.2178 ms 78.9 TFLOPS (what shipped) nsplit=2 0.1241 ms 138.4 TFLOPS nsplit=4 0.0863 ms 199.1 TFLOPS (what the rule selects) nsplit=8 0.1119 ms 153.5 TFLOPS 2.52x on the one shape we were losing. The dispatcher already clamps the split factor to [2, 4]; the nsplit=8 row is why. WHAT HipFlash2FwdPlanVariantSplitK.hip -- the split and merge kernels, built into hip_flash2_fwd_<arch>_splitk.co by the existing hip_flash2_variants target. Same S-transpose body as the single-pass variant, plus a per-chunk partial write and an exact online-softmax merge (m = max m_j, l = sum l_j*exp(m_j-m)). HipModuleGuard gains an optional second entry point. Split-K needs two kernels and both live in one object, so this is one hipModuleLoad and two hipModuleGetFunction calls -- the guard's move/unload semantics are unchanged. buildPlan() takes sel.splitK only when the split-K object is installed, and sizes the workspace with the same flash2WorkspaceBytes() helper the dispatch header already exposed. If the object is missing it falls back to single-pass, mirroring the variant probe, so a partial install degrades rather than fails. execute() carves po/pm/pl out of the caller's workspace exactly as the merge kernel indexes them, then launches split and merge on the same stream -- same-stream ordering is the synchronisation, as SdpaBwdPlan already relies on. NOTE ON THE WORKSPACE This is the first plan in hip_kernel_provider to report a non-zero getWorkspaceSize(). The path was already plumbed: the backend reads the workspace off the VariantPack and the integration harness already queries get_workspace_size() and allocates, so no test change was needed. execute() rejects a null workspace with an explicit message rather than faulting, the way SdpaBwdPlan does. No zeroing is needed and none is done: chunks with no work still write m = -inf and l = 0, so every slot the merge reads was written by the split pass. SdpaBwdPlan must memset its accumulator because that kernel accumulates atomically into leased memory; this one does not. VERIFIED on MI300X (CI does not build this engine -- ENABLE_HIP_FLASH2_ENGINE is OFF and set nowhere -- so all of this was run by hand): correctness, 8 configurations (S in {1024, 2048} x nsplit {2,4} x causal): constant-V invariant 0.000e+00 exact on every row Q=0 -> mean(V) <= 1.9e-05 vs CPU fp32 <= 7.2e-05 Causal is in the set deliberately: the split kernel's causal cap must be CTA-wide, and when an earlier revision made it per-wave the failure was causal-only corruption that non-causal tests missed entirely. object geometry: max_flat_workgroup_size 512 for flash2_split_d128/d64 and 256 for flash2_merge, matching what the plan launches. hipcc -fsyntax-only -std=c++17 -Wall -Wextra: clean on both changed TUs clang-tidy 20 vs the provider config: 7 findings, identical to the parent commit -- this change adds none TestFlash2Dispatch: 9/9, including SplitKCreditMatchesExecutionReality which now asserts the flag is true and fails if execute() and selection ever disagree again cmake-lint: clean
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
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.
feat(hip-flash2): execute split-K instead of only selecting it (AIOSS-5319)
Follow-up to #10977. The dispatcher already decided when to split the KV axis,
but nothing ran it: buildPlan() forced params.splitK = 1 and execute() launched
a single kernel, so the decision was advisory. #10977 went further and gated
the credit behind K_FLASH2_SPLITK_EXECUTES = false, so selection would describe
what actually happened rather than what it wished happened.
This wires it up, and flips that flag.
WHY
B1H8S2048 non-causal was the only shape in the 21-shape suite where we lose to
the CK/AITER backend -- 0.478x -- and it is exactly the shape split-K exists
for. Measured on MI300X, this shape alone:
2.52x on the one shape we were losing. The dispatcher already clamps the split
factor to [2, 4]; the nsplit=8 row is why.
WHAT
HipFlash2FwdPlanVariantSplitK.hip -- the split and merge kernels, built into
hip_flash2_fwd__splitk.co by the existing hip_flash2_variants target.
Same S-transpose body as the single-pass variant, plus a per-chunk partial
write and an exact online-softmax merge (m = max m_j, l = sum l_j*exp(m_j-m)).
HipModuleGuard gains an optional second entry point. Split-K needs two
kernels and both live in one object, so this is one hipModuleLoad and two
hipModuleGetFunction calls -- the guard's move/unload semantics are unchanged.
buildPlan() takes sel.splitK only when the split-K object is installed, and
sizes the workspace with the same flash2WorkspaceBytes() helper the dispatch
header already exposed. If the object is missing it falls back to single-pass,
mirroring the variant probe, so a partial install degrades rather than fails.
execute() carves po/pm/pl out of the caller's workspace exactly as the merge
kernel indexes them, then launches split and merge on the same stream --
same-stream ordering is the synchronisation, as SdpaBwdPlan already relies on.
NOTE ON THE WORKSPACE
This is the first plan in hip_kernel_provider to report a non-zero
getWorkspaceSize(). The path was already plumbed: the backend reads the
workspace off the VariantPack and the integration harness already queries
get_workspace_size() and allocates, so no test change was needed. execute()
rejects a null workspace with an explicit message rather than faulting, the
way SdpaBwdPlan does.
No zeroing is needed and none is done: chunks with no work still write
m = -inf and l = 0, so every slot the merge reads was written by the split
pass. SdpaBwdPlan must memset its accumulator because that kernel accumulates
atomically into leased memory; this one does not.
VERIFIED on MI300X (CI does not build this engine -- ENABLE_HIP_FLASH2_ENGINE
is OFF and set nowhere -- so all of this was run by hand):
correctness, 8 configurations (S in {1024, 2048} x nsplit {2,4} x causal):
constant-V invariant 0.000e+00 exact on every row
Q=0 -> mean(V) <= 1.9e-05
vs CPU fp32 <= 7.2e-05
Causal is in the set deliberately: the split kernel's causal cap must be
CTA-wide, and when an earlier revision made it per-wave the failure was
causal-only corruption that non-causal tests missed entirely.
object geometry: max_flat_workgroup_size 512 for flash2_split_d128/d64 and
256 for flash2_merge, matching what the plan launches.
hipcc -fsyntax-only -std=c++17 -Wall -Wextra: clean on both changed TUs
clang-tidy 20 vs the provider config: 7 findings, identical to the parent
commit -- this change adds none
TestFlash2Dispatch: 9/9, including SplitKCreditMatchesExecutionReality which
now asserts the flag is true and fails if execute() and selection ever
disagree again
cmake-lint: clean