Repository navigation
Conversation
github-actions
Bot
added
project: hip-kernel-provider
rocke-lib: attention
✅ 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. |
arai713
changed the title
arai713
marked this pull request as ready for review
October 21, 2026 16:14
arai713
force-pushed
the
users/arai/attn_free_stride
branch
from
October 26, 2026 17:46
cd39c1d to
f9b8478
Compare
Add fmha_sdpa_fwd, the net-new MFMA-body layout shell (hipDNN requirement 2), reusing mfma_attention_fwd_inner_body for the attention math and adding only the addressing layer. This scaffold covers the contiguous BHSD dense case: batch is threaded through the grid z-axis plus the inner body's K/V offset hooks, with Q/O batch folded into the tile base. BSHD, general strides, and packed QKV are declared on the spec but gated off in is_valid_spec as Phase 3. CPU-only build/smoke tests cover the validator gates, ABI signature, grid, and that a valid spec lowers to a KernelDef.
Make fmha_sdpa_fwd read Q/K/V/O from any memory layout by rebasing each tensor's base pointer to the CTA's batch slice (batch_idx * stride_*_batch) and letting the shared body address head/token through its runtime strides. Batch, head and token offsets are all runtime, so BSHD, BHSD, general B/H/S strides and packed QKV build the same kernel, with nothing baked and the contiguous Q/O fold removed. Add _sdpa_fwd_strides.contiguous_strides as the launcher's single source of truth for the (batch, head, token) element strides of the bhsd/bshd orders, plus host tests: the layout contract and the updated validator/build smoke tests.
Add an opt-in free_stride mode that rebases Q/K/V/O base pointers per batch and declares per-tensor batch strides, so one kernel reads any layout (BSHD/BHSD/general/packed) instead of the BSHD-contiguous seqlen fold.The free-stride path is Python-only for now. Retire the fmha_sdpa_fwd scaffold, add a CPU ABI test and a GPU verify harness.
Extract the verify harness core into a reusable run_verify() and add a CDNA-gated pytest (skips off gfx942/gfx950) that checks fmha_mfma(free_stride= True) against the numpy reference for BSHD and BHSD, confirming layout invariance of the per-batch pointer rebase.
arai713
force-pushed
the
users/arai/attn_free_stride
branch
from
October 27, 2026 17:53
f9b8478 to
0bdfdac
Compare
Contributor
Author
|
Moved to PR 13953 |
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
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.
JIRA ID: AICK-2618
Motivation
hipDNN needs attention forward to read Q/K/V from arbitrary memory layouts, not one fixed contiguous order. The production forward kernel (fmha_mfma) only served the BSHD-contiguous case: it folded the batch axis in as batch_idx * seqlen_q rows, which is expressible only when the batch stride is a whole number of token strides. BHSD, general B/H/S strides and packed QKV could not be served without a host-side repack. This generalizes fmha_mfma so one kernel reads any layout.
Technical Details
free_stride: bool = Falsefield to FmhaMfmaSpec (defaulted, so existing descriptors / AOT packs hydrate to the current behavior).global_ptr_add(ptr, batch_idx * stride_*_batch * elem_bytes)). The shared body then addresses head and token through the runtime strides it already takes. Batch, head and token offsets are all runtime, so BSHD, BHSD, general strides and packed QKV all build the same kernel with nothing baked._fs._sdpa_fwd_strides.contiguous_stridessingle-sources the (batch, head, token) element strides the launcher passes for the bhsd/bshd orders. The earlier standalone free-stride scaffold (fmha_sdpa_fwd) is retired, now subsumed by fmha_mfma.Test Plan
test_fmha_mfma_free_stride.py): default ABI unchanged (15 params, no batch strides), free_stride appends exactly the four batch strides (19), name tagging, builds for gfx942/gfx950 and GQA.test_sdpa_free_stride_layouts.py): every layout (BHSD/BSHD/padded/packed) recovers the same values and attention output via the kernel's scalar address formula;contiguous_stridesmatches numpy.test_fmha_mfma_free_stride_numeric.py, CDNA-gated): fmha_mfma(free_stride=True) matches the numpy reference for BSHD and BHSD.tools/check_byte_identity.py).Test Result
Submission Checklist