[CK_TILE] Fix typo in fmha_fwd_kernel K-dram unmerge tuple sizes - #7141
Merged
Conversation
In the qr_async_trload K-dram lambda's `else (XorLengthFold == 1)` branch,
the outer-tile dimension of the (#tiles, #aligned-blocks-per-tile, alignment)
unmerge/xor/merge triple was written as
number<kQKHeaddim / kDramTileK / kAlignmentK>{}
which divides one extra time and yields length=0 for every fp16/bf16 hdim=128
configuration (e.g. kDramTileK=128 / kAlignmentK=8 -> 128/128/8 == 0). The
product across the 3-tuple should equal kQKHeaddim, but currently evaluates
to 0 (or 16/32 for some atypical kK0 values), so unmerge -> xor -> merge no
longer round-trips to the original head dimension.
This was previously masked because the async-load path only walks via stride
and silently absorbs a length=0 outer dim. Any consumer that traverses the
descriptor (e.g. TDM on gfx1250) immediately faults on the resulting
`tuple<int, constant<0>>`.
Drop the extra `/ kAlignmentK` in all three sites (line 2692, 2705, 2715)
so the outer-tile length becomes kQKHeaddim/kDramTileK and the 3-tuple
product is restored to kQKHeaddim. Strides are unaffected, so the async
path remains bit-identical.
Validated on gfx950 (b3bdc63): test_ck_tile_fmha_fwd ctest pass/fail
counts identical pre- and post-fix (3 pass / 2 fail; the 2 pre-existing
fp16/bf16 group-mode failures are unrelated and reproduce on develop without
this patch). Total wall time 403s -> 393s, per-case latency drift within
+/-8% noise.
Bug introduced in 2cc0af6 (PR #2888 "[CK_TILE] FMHA FWD bug fix").
3 tasks
poyenc
approved these changes
May 7, 2026
assistant-librarian Bot
pushed a commit
to ROCm/composable_kernel
that referenced
this pull request
May 8, 2026
[CK_TILE] Fix typo in fmha_fwd_kernel K-dram unmerge tuple
sizes (#7141)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
## Summary
The qr_async_trload K-dram lambda's `else (XorLengthFold == 1)` branch
in `fmha_fwd_kernel.hpp` writes the outer-tile dim of its 3-tuple
unmerge/xor/merge as
```cpp
number<FmhaPipeline::kQKHeaddim / kDramTileK / FmhaPipeline::kAlignmentK>{}
```
which divides one extra time. For every fp16/bf16 hdim=128 configuration
the outer length collapses to **0**, e.g. `128 / 128 / 8 == 0`. The
3-tuple product no longer equals `kQKHeaddim`, so unmerge → xor → merge
stops round-tripping the head dimension.
This bug was masked by the async-load path: it only walks the descriptor
via stride and silently absorbs a length=0 outer dim. Any consumer that
actually traverses the descriptor (e.g. the TDM path on gfx1250)
immediately faults on the resulting `tuple<int, constant<0>>`.
The fix drops the extra `/ kAlignmentK` in all three call sites in the
same lambda so the outer dim becomes `kQKHeaddim / kDramTileK` and the
product is restored to `kQKHeaddim`. Strides are unaffected, so the
async path is bit-identical.
| Config (fp16/bf16) | hdim | kDramTileK | kAlignmentK | a (typo) | a
(fixed) | product (typo) | product (fixed) |
|---|---|---|---|---|---|---|---|
| hdim128, kKLoadOnce | 128 | 128 | 8 | 0 | 1 | **0** | **128** |
| hdim128, kK0=32 | 128 | 32 | 8 | 0 | 4 | **0** | **128** |
| hdim64, kKLoadOnce | 64 | 64 | 8 | 0 | 1 | **0** | **64** |
| hdim256, kK0=32 | 256 | 32 | 8 | 1 | 8 | **32** | **256** |
Bug introduced in 2cc0af6a815a (PR #2888 \"[CK_TILE] FMHA FWD bug
fix\"), where the original 2-tuple unmerge was generalized to a 3-tuple
and the typo slipped in.
## Test plan
- [x] Built `test_ck_tile_fmha_fwd` (umbrella, 5 gtest binaries) on
gfx950 native at develop b3bdc63a509 with `dev-gfx950` preset (clang 22,
ROCm 7.2.2). Compiles cleanly with `-Werror -Weverything`.
- [x] Ran `ctest -R test_ck_tile_fmha_fwd` on gfx950 native, baseline vs
patched: identical pass/fail (3 pass / 2 fail), identical failing case
set (114 gtest fails + 2 GPU memory access faults, all in pre-existing
fp16/bf16 group-mode `Alibi`/`Dropout` cases that reproduce on develop
without this patch). Total wall time 403s → 393s. Per-case latency drift
±8% (noise).
- [x] CI to verify on other gfx9 / gfx11 architectures.
aledudek
pushed a commit
that referenced
this pull request
May 20, 2026
## Summary
The qr_async_trload K-dram lambda's `else (XorLengthFold == 1)` branch
in `fmha_fwd_kernel.hpp` writes the outer-tile dim of its 3-tuple
unmerge/xor/merge as
```cpp
number<FmhaPipeline::kQKHeaddim / kDramTileK / FmhaPipeline::kAlignmentK>{}
```
which divides one extra time. For every fp16/bf16 hdim=128 configuration
the outer length collapses to **0**, e.g. `128 / 128 / 8 == 0`. The
3-tuple product no longer equals `kQKHeaddim`, so unmerge → xor → merge
stops round-tripping the head dimension.
This bug was masked by the async-load path: it only walks the descriptor
via stride and silently absorbs a length=0 outer dim. Any consumer that
actually traverses the descriptor (e.g. the TDM path on gfx1250)
immediately faults on the resulting `tuple<int, constant<0>>`.
The fix drops the extra `/ kAlignmentK` in all three call sites in the
same lambda so the outer dim becomes `kQKHeaddim / kDramTileK` and the
product is restored to `kQKHeaddim`. Strides are unaffected, so the
async path is bit-identical.
| Config (fp16/bf16) | hdim | kDramTileK | kAlignmentK | a (typo) | a
(fixed) | product (typo) | product (fixed) |
|---|---|---|---|---|---|---|---|
| hdim128, kKLoadOnce | 128 | 128 | 8 | 0 | 1 | **0** | **128** |
| hdim128, kK0=32 | 128 | 32 | 8 | 0 | 4 | **0** | **128** |
| hdim64, kKLoadOnce | 64 | 64 | 8 | 0 | 1 | **0** | **64** |
| hdim256, kK0=32 | 256 | 32 | 8 | 1 | 8 | **32** | **256** |
Bug introduced in 2cc0af6 (PR #2888 \"[CK_TILE] FMHA FWD bug
fix\"), where the original 2-tuple unmerge was generalized to a 3-tuple
and the typo slipped in.
## Test plan
- [x] Built `test_ck_tile_fmha_fwd` (umbrella, 5 gtest binaries) on
gfx950 native at develop b3bdc63 with `dev-gfx950` preset (clang 22,
ROCm 7.2.2). Compiles cleanly with `-Werror -Weverything`.
- [x] Ran `ctest -R test_ck_tile_fmha_fwd` on gfx950 native, baseline vs
patched: identical pass/fail (3 pass / 2 fail), identical failing case
set (114 gtest fails + 2 GPU memory access faults, all in pre-existing
fp16/bf16 group-mode `Alibi`/`Dropout` cases that reproduce on develop
without this patch). Total wall time 403s → 393s. Per-case latency drift
±8% (noise).
- [x] CI to verify on other gfx9 / gfx11 architectures.
shumway
pushed a commit
to ROCm/composable_kernel
that referenced
this pull request
May 27, 2026
[CK_TILE] Fix typo in fmha_fwd_kernel K-dram unmerge tuple sizes (#7141)
## Summary
The qr_async_trload K-dram lambda's `else (XorLengthFold == 1)` branch
in `fmha_fwd_kernel.hpp` writes the outer-tile dim of its 3-tuple
unmerge/xor/merge as
```cpp
number<FmhaPipeline::kQKHeaddim / kDramTileK / FmhaPipeline::kAlignmentK>{}
```
which divides one extra time. For every fp16/bf16 hdim=128 configuration
the outer length collapses to **0**, e.g. `128 / 128 / 8 == 0`. The
3-tuple product no longer equals `kQKHeaddim`, so unmerge → xor → merge
stops round-tripping the head dimension.
This bug was masked by the async-load path: it only walks the descriptor
via stride and silently absorbs a length=0 outer dim. Any consumer that
actually traverses the descriptor (e.g. the TDM path on gfx1250)
immediately faults on the resulting `tuple<int, constant<0>>`.
The fix drops the extra `/ kAlignmentK` in all three call sites in the
same lambda so the outer dim becomes `kQKHeaddim / kDramTileK` and the
product is restored to `kQKHeaddim`. Strides are unaffected, so the
async path is bit-identical.
| Config (fp16/bf16) | hdim | kDramTileK | kAlignmentK | a (typo) | a
(fixed) | product (typo) | product (fixed) |
|---|---|---|---|---|---|---|---|
| hdim128, kKLoadOnce | 128 | 128 | 8 | 0 | 1 | **0** | **128** |
| hdim128, kK0=32 | 128 | 32 | 8 | 0 | 4 | **0** | **128** |
| hdim64, kKLoadOnce | 64 | 64 | 8 | 0 | 1 | **0** | **64** |
| hdim256, kK0=32 | 256 | 32 | 8 | 1 | 8 | **32** | **256** |
Bug introduced in 2cc0af6a815a (PR #2888 \"[CK_TILE] FMHA FWD bug
fix\"), where the original 2-tuple unmerge was generalized to a 3-tuple
and the typo slipped in.
## Test plan
- [x] Built `test_ck_tile_fmha_fwd` (umbrella, 5 gtest binaries) on
gfx950 native at develop b3bdc63a509 with `dev-gfx950` preset (clang 22,
ROCm 7.2.2). Compiles cleanly with `-Werror -Weverything`.
- [x] Ran `ctest -R test_ck_tile_fmha_fwd` on gfx950 native, baseline vs
patched: identical pass/fail (3 pass / 2 fail), identical failing case
set (114 gtest fails + 2 GPU memory access faults, all in pre-existing
fp16/bf16 group-mode `Alibi`/`Dropout` cases that reproduce on develop
without this patch). Total wall time 403s → 393s. Per-case latency drift
±8% (noise).
- [x] CI to verify on other gfx9 / gfx11 architectures.
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
2 participants
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
The qr_async_trload K-dram lambda's
else (XorLengthFold == 1)branch infmha_fwd_kernel.hppwrites the outer-tile dim of its 3-tuple unmerge/xor/merge aswhich divides one extra time. For every fp16/bf16 hdim=128 configuration the outer length collapses to 0, e.g.
128 / 128 / 8 == 0. The 3-tuple product no longer equalskQKHeaddim, so unmerge → xor → merge stops round-tripping the head dimension.This bug was masked by the async-load path: it only walks the descriptor via stride and silently absorbs a length=0 outer dim. Any consumer that actually traverses the descriptor (e.g. the TDM path on gfx1250) immediately faults on the resulting
tuple<int, constant<0>>.The fix drops the extra
/ kAlignmentKin all three call sites in the same lambda so the outer dim becomeskQKHeaddim / kDramTileKand the product is restored tokQKHeaddim. Strides are unaffected, so the async path is bit-identical.Bug introduced in 2cc0af6 (PR #2888 "[CK_TILE] FMHA FWD bug fix"), where the original 2-tuple unmerge was generalized to a 3-tuple and the typo slipped in.
Test plan
test_ck_tile_fmha_fwd(umbrella, 5 gtest binaries) on gfx950 native at develop b3bdc63 withdev-gfx950preset (clang 22, ROCm 7.2.2). Compiles cleanly with-Werror -Weverything.ctest -R test_ck_tile_fmha_fwdon gfx950 native, baseline vs patched: identical pass/fail (3 pass / 2 fail), identical failing case set (114 gtest fails + 2 GPU memory access faults, all in pre-existing fp16/bf16 group-modeAlibi/Dropoutcases that reproduce on develop without this patch). Total wall time 403s → 393s. Per-case latency drift ±8% (noise).