feat(ck-tile): add FMHA FWD TDM Pipeline - #7755
Merged
Merged
Conversation
…>> in-place ops
Five tuple ops in CK Tile's container layer used in-place mutation into a
`tuple<Xs...> r;` of the same parametric type as the input. That pattern
fails to compile when the input is a mixed (runtime int, compile-time
constant<N>) tuple, because the runtime int can't be assigned back into
the constant<N> slot. Rewrite all five to build the result via
`generate_tuple` (or, for the scan case, recursive `container_push_front`)
so each output slot's type is freshly deduced from its expression.
This is purely a CK core fix and is independent of any consumer; it has
its own value because heterogeneous tuples are increasingly common in
descriptor / dram-view code (anywhere static dims are mixed with dynamic
ones). The discovery context was the gfx1250 FMHA TDM Q+K work — TDM's
`tile_window::get_cached_global_strides` (tile_window.hpp:1788) calls
into these ops on dram view length tuples that mix `int` (runtime stride
dim) and `constant<N>` (compile-time hdim dim), and hits a hard compile
failure on the in-place store. Fixing those five ops unblocks any TDM
consumer with such a tensor view.
Files changed
-------------
include/ck_tile/core/container/container_helper.hpp
- container_reverse_inclusive_scan: rewritten from in-place
`tuple<Xs...> y; y(i) = r;` to a recursive helper that builds a fresh
tuple via `container_push_front`. New helper
`container_reverse_inclusive_scan_impl` mirrors the existing
`container_reverse_exclusive_scan_impl` pattern just above; the
pre-existing TODO at the top of the function (`// TODO: update to like
container_reverse_exclusive_scan to deal with tuple of Number<>`)
called for exactly this rewrite, so it is removed too. Inclusive-scan
semantics preserved exactly:
y[N-1] = f(init, x[N-1])
y[N-2] = f(y[N-1], x[N-2])
...
y[0] = f(y[1], x[0])
and the reduce arg order stays `f(r_old, x[i])` to match existing
callers.
include/ck_tile/core/container/tuple.hpp (4 ops)
- operator+(tuple<Xs...>, Y) (MultiIndex + scalar-broadcast Y):
`tuple<Xs...> r; r[i] = x[i]+y[i]` → `generate_tuple(...)`. Mirrors
the already-correct `operator+(tuple<Xs>, tuple<Ys>)` overload defined
immediately below.
- operator-(tuple<Xs...>, Y): same pattern, mirrors the
`operator-(tuple<Xs>, tuple<Ys>)` overload.
- operator*(tuple<Xs...>, Y): same pattern, mirrors the
`operator*(tuple<Xs>, tuple<Ys>)` overload.
- operator*(Y, tuple<Xs...>) (scalar * MultiIndex): same pattern,
`generate_tuple([&](auto i){ return a*x[i]; }, ...)`.
Each rewritten op carries an inline comment explaining the why
(in-place store vs. mixed tuple) and pointing at the existing
correct counterpart it now mirrors, so future readers won't
re-introduce the in-place pattern out of habit.
No behavior change for homogeneous tuples
-----------------------------------------
For tuples whose element types are all the same, the old in-place path
and the new generate_tuple/recursive path produce a result of the same
type and same values. All existing call sites that work today continue
to work bit-identically; the change only enables additional inputs that
previously failed to compile.
No test added in this PR — the fix is exercised by the gfx1250 FMHA TDM
Q+K consumer (separate branch / PR), which would not compile without
this change.
…ride
Root cause
----------
`tile_window_with_static_distribution::get_cached_global_strides()` used
`glb_tensor_descriptor.get_lengths()` to derive a packed reverse-inclusive-
scan, then divided by `Traits::PackedSize`. That implicitly assumed the
tensor view was packed. For any caller whose tensor view stride differs
from the shape-implied packed default, the cached "strides" silently
fabricated a wrong byte offset.
The bug was first observed via FMHA multi-head Q under GQA: the host
sets `kargs.stride_q = h_q * hdim_q` (256 for h_q=2 / hdim_q=128) so
that adjacent rows of the per-head Q view skip across interleaved
heads, but `get_cached_global_strides()` returned `(hdim_q, 1) =
(128, 1)`, producing the wrong dram address per row. Padded GEMM
(`stride_a > K`) hits the same class of bug.
Fix
---
Read the actual stride for each top-level dimension by querying the
descriptor with a unit vector:
cached_global_strides_ = generate_array(
[&](auto i) {
auto unit_vec = make_zero_multi_index<NDimBottomTensor>();
unit_vec(i) = 1;
return max(desc.calculate_offset(unit_vec) / PackedSize, 1);
},
number<NDimBottomTensor>{});
`calculate_offset` traverses the full transform chain (embed/pad/etc.),
so the result reflects the layout the caller baked into the view.
Single path -- no packed fast-path retained. Cost is N
`calculate_offset` calls on first access; subsequent lookups still hit
the `tensor_cache_` field.
TDM consumer audit
------------------
Surveyed every caller that constructs a tile_window feeding a TDM /
prefetch path that hits `get_cached_global_strides()`. All six set the
actual stride field via `make_naive_tensor_view(p, lengths, strides)`
explicitly; none rely on the packed-default fallback for default
invocations:
| # | Consumer | View factory | Stride source / default |
|---|----------------------------------------|--------------------------------------------------------------------|--------------------------------------|
| 1 | `tdm_basic` test (`tdm_kernel.hpp`) | `make_naive_tensor_view(p, (M,N), (arg.stride_input, 1))` | TDMTestParams.normalize -> packed |
| 2 | `gemm_tdm_data_cache_prefetch` example | `make_naive_tensor_descriptor((M,K), (stride_a, 1))` (universal) | `validate_gemm_stride` default packed|
| 3 | `gemm_weight_preshuffle_tdm` example | universal kernel pattern + preshuffle invoker | same as GEMM |
| 4 | `gemm_pipeline_tdm_wmma` test | `get_default_stride(M, K, StrideA, is_row_major)` | StrideA/B/C kargs default 0 -> packed|
| 5 | `mx_gemm_pipeline_tdm_wmma` test | same util pattern | same as GEMM |
| 6 | FMHA pipeline (`fmha_fwd_kernel.hpp`) | `make_naive_tensor_view(q_ptr, (seqlen_q, hdim_q), (stride_q, 1))` | `kargs.stride_q` (= h_q*hdim_q under GQA) |
Consumer 6 is the actual triggering case: under GQA it sets a
non-packed stride and exercised the bug end-to-end. Default
invocations of consumers 1-5 all happen to pass packed strides, so
pre-fix and post-fix return identical values for them; the fix
changes behaviour only when a caller passes a non-packed stride
(which was already the documented intent of the API).
Pre-existence verify (mx_gemm fp4)
----------------------------------
Built `test_ck_tile_mx_gemm_pipeline_tdm_wmma` on `internel/gfx1250`
HEAD `62d40f9730c` (no FMHA changes from this branch). Result: 48
PASSED / 12 FAILED, identical FAIL identity set to this branch's
baseline. The 12 failures (all
`pk_float4_e2m1_t * pk_float4_e2m1_t` typed instances on
SingleTile/MidLargeM/LargeSize sub-tests) are pre-existing on
internel/gfx1250 and untouched by this fix.
Baseline diff (5 TDM consumer binaries, FAIL identity sets)
-----------------------------------------------------------
| binary | pre-fix | post-fix | identity diff |
|---------------------------------------------------|------------|------------|---------------|
| `test_tdm_basic` | 32/8/0 | 32/8/0 | empty |
| `tile_example_gemm_tdm_data_cache_prefetch` 128^3 | PASS | PASS | n/a |
| `tile_example_gemm_weight_preshuffle_tdm_*` 128^3 | PASS | PASS | n/a |
| `test_ck_tile_gemm_pipeline_tdm_wmma` (full) | 96/16/0 | 96/16/0 | SKIP set empty|
| `test_ck_tile_mx_gemm_pipeline_tdm_wmma` (full) | 48/0/12 | 48/0/12 | FAIL set empty|
PASS / SKIP / FAIL columns. No new failures, no PASS-to-FAIL flips,
no SKIP-set drift. Sim time changes are within noise (gemm_tdm
513->582ms and weight_preshuffle 124->126ms reflect the added
per-cache-init cost of N `calculate_offset` calls on first access).
Out of scope
------------
- FMHA GQA correctness re-validation belongs to QA after this lands.
This fix only addresses the cached-strides half of the GQA root
cause; the second half (GQA-aware Q distribution redesign) is
tracked as separate follow-up work.
- Dispatcher prefer-qr_tdm and re-enabling K/V padding are
independent follow-ups, unaffected by this commit.
Introduce a new forward FMHA pipeline targeting gfx1250's box-major TDM (tensor data movement) DMA, alongside its policy class: - include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm.hpp - include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp Wire the new pipeline into: - block_fmha_pipeline_enum.hpp (new enum entry) - ops/fmha.hpp (header aggregation) - example/ck_tile/01_fmha/codegen/cpp_symbol_map.py + codegen/ops/fmha_fwd.py (codegen registration; gated to a narrow trait combo) This commit only scaffolds the type and registers it with codegen; the loader still mirrors the async_load reference path and the box-major TDM intrinsic is not yet engaged. Subsequent commits make Q/K and then V go through load_tile_tdm.
Bring the scaffolded qr_ks_vs_tdm pipeline to a working state on gfx1250 by aligning its dispatch and distribution shape with the existing qr_async_trload pipeline, ahead of the box-major TDM intrinsic swap. block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp: - AsyncCopy = true - GetAlignmentK/V switched to b128 on gfx1250 - K/V dram distribution overrides mirroring qr_async_trload_policy - GetSmemSize gains the prefill / decode branch block_fmha_pipeline_qr_ks_vs_tdm.hpp: - operator() signature aligned with qr_async_trload-style dispatch - run() split into single + double buffer variants under an outer operator() wrapper - kAlignmentOacc introduced fmha_fwd_kernel.hpp: - qr_tdm dispatched into the qr_async_trload-style branch example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py: - qr_tdm codegen gated to hdim=128 (the only shape with a validated distribution override at this point) Functional tests pass on the gfx1250 ffm_lite simulator via async_load (perf -21% to -32% vs the qr baseline; expected, since the actual TDM intrinsic is not yet in use). Moving the loader to load_tile_tdm is deferred to a follow-up commit.
Move Q and K loads from async_load_tile to load_tile_tdm on gfx1250 FMHA forward; V continues to use async_load and is moved in a follow-up commit. Single-head shapes (h_q == h_k) get the full TDM acceleration; multi-head GQA falls back to the prior async_load behavior since TDM hits two latent issues there (carried as known issues). Pipeline (block_fmha_pipeline_qr_ks_vs_tdm.hpp / _policy.hpp): - K dram dist switched to trivial tile-major (mirrors GEMM v1 ColMajor B layout). Each thread's per-call footprint is now one contiguous (kN0/warpNum, kK0) tile, which the TDM box-major write lands at the matching row-major K LDS strip. - K LDS read view changed from Xor=true to plain row-major to match the writer. The Xor template parameter on MakeKLdsBlockDescriptor is removed entirely; TDM box-major writes cannot produce an XOR'd LDS layout, so the XOR branch was dead code (~80 lines). - Q dram dist switched to trivial tile-major mirroring K. Same rationale: TDM box writes need a contiguous per-thread footprint. - Q LDS read view changed Xor=true to plain row-major to match the writer; the Xor template parameter on MakeQLdsBlockDescriptor is removed entirely (same dead-code reasoning as K). - Q TDM padding disabled in GetLdsPaddingConfigQ to mirror K and V. Otherwise TDM writes Q with padding into a plain row-major LDS desc and leaves sentinel bytes in the gaps. - V keeps async_load + Xor=true read view (intentionally out of scope here). - Q load is now load_tile_tdm with TDMConfig built once at kernel start. Kernel (fmha_fwd_kernel.hpp): - Unmerge typo fix in the K dram view path. The leading dim of the unmerge was kQKHeaddim/kDramTileK/kAlignmentK, which under typical configs (e.g. fp16 hdim=128, kDramTileK=kK0=32, kAlignmentK=8) folds to 0 by integer division, producing a tuple<int, constant<0>> with product 0 instead of the real hdim_q (128). The async-load path never read those lengths so the bug stayed dormant; the TDM path reads them via get_cached_global_strides and hits a hard compile failure. Fix: drop the spurious / kAlignmentK from all three sites (unmerge, pass_through, merge) so the leading dim is kQKHeaddim/kDramTileK and the 3-axis layout stays consistent. Note: this change is logically the same as upstream develop's PR #6964 (already merged after this branch was cut). On rebase to develop, GitHub should auto-resolve the identical patch; if not, a manual conflict resolve collapses to the same final source. - TDM-aware Q/K dram view dispatch added inside make_q_dram and make_k_dram lambdas. When the pipeline is qr_tdm, the Q/K dram views skip the unmerge -> xor -> merge_v3 chain and return the affine pad-only view directly (mirroring the kQLoadOnce=false branch). The XOR'd transform chain stays intact for qr_async_trload, wrapped in `else { ... }` so it is if-constexpr discarded for qr_tdm. Why: TDM box-major DMA cannot honor a software XOR layout on the dram side, so the XOR chain was dead code for TDM. With the framework get_cached_global_strides fix (the parent commit on this branch), calculate_offset(unit_vec) walks the full transform chain and the XOR node folds the unit vector into a polluted offset (e.g. 136 = stride_q + kAlignmentQ instead of 128), which then drives the box copy to read wrong rows. Bypassing the chain returns the true byte stride. V dram is intentionally unchanged: V keeps the async_load path and never calls get_cached_global_strides, so its XOR'd dram view chain is unaffected. Codegen (fmha_fwd.py): - qr_vr emit kept disabled for d=128 fp16/bf16 single-head cases. Otherwise the runtime dispatcher selects qr_vr (no TDM acceleration) over qr_tdm. Tracked as a workaround; superseded by a later dispatcher prefer-qr_tdm change. Test (gfx1250 ffm_lite simulator, all dispatch kname qr_tdm_vr_npad): - fp16 dense: b=1 h=1 s=1023 d=128 mask=0 valid:y - fp16 GQA causal: b=1 h_q=2 h_k=1 s=1023 s_k=257 d=128 mask=2 valid:y - bf16 long: b=1 h=1 s=2047 d=128 mask=0 valid:y Multi-stride GQA (verifies K stride 256/512/1024 all resolve under the new TDM dram dispatch): h=2 h_k=2 s=1023 (no GQA, stride_q=256) -> valid:y h=4 h_k=1 s=1023 (GQA 4:1, stride_q=512) -> valid:y h=8 h_k=1 s=1023 (GQA 8:1, stride_q=1024) -> valid:y AM cycle-accurate (rocdtif r5.04, b=1 h=1 s=128 d=128 fp16 -timer=cpu): - valid:y, matching prior baseline on the same shape. Non-TDM regression (fp8 / fp8bf16 routes that bypass qr_tdm on gfx1250): kname dispatched is qr_vr_psskddv on both this commit and the prior baseline; valid:y on both with byte-identical kname suffix. GEMM identity scope: the unmerge fix and the qr_tdm dram dispatch both touch only fmha_fwd_kernel.hpp. No GEMM example or non-fmha header transitively includes that file, so GEMM binaries cannot be affected. Known issues carried forward: 1. Trivial tile-major Q dist produces wrong thread->element mapping under multi-head GQA. The qr_tdm dram dispatch sidesteps this by consuming the affine pad-only view (whose stride is the actual stride_q the host sets, e.g. h_q * hdim_q for GQA) rather than the transformed view. Remaining work: - Dispatcher prefer-qr_tdm to retire the qr_vr-disable workaround. - Re-enable K/V padding. - Move V to load_tile_tdm.
Move V load from async_load_tile + ds_load_tr to load_tile_tdm on
gfx1250 FMHA forward; Q and K were moved to TDM in the previous
commit. With this change all three operand loads go through the
unified box-major DMA path with plain row-major LDS layouts.
Pipeline (block_fmha_pipeline_qr_ks_vs_tdm.hpp / _policy.hpp):
- V dram dist: MakeVDramTileDistribution switched to trivial
tile-major mirroring the K dist (MakeKDramTileDistribution). Each
thread's per-call footprint is a contiguous (kKPerBlock/warpNum,
kNPerBlock) tile that TDM box-major DMA lands at the matching
row-major V LDS strip. The prior 5D scatter dist is incompatible
with TDM box writes — the dist projection scatters thread bytes
to LDS positions that ds_load_tr_b128 reads as garbage (a probe
during development showed 1.6% byte agreement against the
verified-working async path on the same V LDS region).
- V LDS read view: kept plain row-major (MakeVLdsBlockDescriptor
default, Xor=false), matching the writer. ds_load_tr_b128 +
MakeVRegTileDistribution + QuadInputEncoding suffix
(TransposedDstrEncode) produce per-lane VOFFSETs that read from
the plain LDS bytes to the WMMA B operand pattern; the XOR'd read
view used by the prior async path is not required for correctness
(it was a bank-conflict-avoidance perf mode that no longer
applies once V uses plain TDM writes).
- V loader: async_load_tile(v_lds_write_window, v_dram_window)
replaced by load_tile_tdm(tdm_config_v, v_lds_write_window,
v_dram_window). Mirrors Q and K dispatch.
- V sync: post-V-load barrier block_sync_lds_direct_load<0>()
(which only drains async loads) replaced by
s_wait_tensorcnt_barrier<0>() to drain the V TDM write on the
tensorcnt counter before load_tile_transpose reads.
- tdm_config_v added alongside tdm_config_q + tdm_config_k,
initialized from GetLdsPaddingConfigV. V padding remains disabled
(no change vs the prior state); see future work below.
Kernel (fmha_fwd_kernel.hpp):
- V-TDM dram view dispatch: make_v_dram lambda gains a TDM-aware
early return: when the pipeline is qr_tdm, the V dram view skips
the unmerge -> xor -> merge_v3 chain and returns the affine
pad-only view directly. Mirrors the existing qr_tdm dispatch in
make_q_dram and make_k_dram. The XOR'd transform chain stays
intact for qr_async_trload, wrapped in `else { ... }` so it is
if-constexpr discarded for qr_tdm. Same rationale as the K dram
dispatch: TDM box DMA cannot honor a software XOR layout on the
dram side, and the framework get_cached_global_strides fix walks
the full transform chain (which would otherwise fold the unit
vector into a polluted V stride).
Codegen: no change. The qr_vr-disable workaround in fmha_fwd.py
(forcing qr_tdm dispatch for d=128 fp16/bf16) continues to apply.
Test (gfx1250 ffm_lite simulator, all dispatch kname qr_tdm_vr_npad):
- fp16 dense: b=1 h=1 s=1023 d=128 mask=0 valid:y
- fp16 GQA causal: b=1 h_q=2 h_k=1 s=1023 s_k=257
d=128 mask=2 valid:y
- bf16 long: b=1 h=1 s=2047 d=128 mask=0 valid:y
Multi-stride GQA (verifies V stride 256/512/1024 all resolve under
the new TDM V dram dispatch, mask=2 causal):
h=4 h_k=1 s=1023 s_k=257 (GQA 4:1, stride_v=128) -> valid:y
h=8 h_k=1 s=1023 s_k=257 (GQA 8:1, stride_v=128) -> valid:y
h=16 h_k=1 s=1023 s_k=257 (GQA 16:1, stride_v=128) -> valid:y
Head dim sweep (d=32/64 dispatch to the d=128 qr_tdm instance with
hdim padding `pddv`; d=128 native instance):
d=32 fp16 dense s=1023 -> valid:y
d=32 bf16 dense s=2047 -> valid:y
d=64 fp16 dense s=1023 -> valid:y
d=64 bf16 dense s=2047 -> valid:y
d=128 (baseline above) -> valid:y
GEMM identity scope: V dram dispatch in make_v_dram and the
pipeline V-side changes touch only fmha_fwd_kernel.hpp /
block_fmha_pipeline_qr_ks_vs_tdm{,_policy}.hpp. No GEMM example or
non-fmha header transitively includes these files, so GEMM
binaries cannot be affected.
Resolved (vs the Q+K-only TDM state):
1. V async_load_tile + Xor=true read view replaced by V
load_tile_tdm + plain row-major read view. The earlier finding
(V LDS byte-dump probe) that ds_load_tr_b128 reads garbage from
TDM box-major writes when the V dram dist is the 5D form is
resolved by switching V dram dist to trivial tile-major, which
yields a plain row-major LDS that the standard reader machinery
consumes correctly.
2. The V dram view XOR-permuted chain is bypassed for qr_tdm via
the same affine-pad early return pattern as K. Closes the
symmetry gap between Q/K-TDM and V-TDM.
Remaining work:
- Dispatcher prefer-qr_tdm to retire the qr_vr-disable workaround.
- Re-enable K / V LDS padding. The bank-conflict-avoidance perf
mode requires a padding-aware MakeKLdsBlockDescriptor /
MakeVLdsBlockDescriptor mirroring gemm pipeline's
MakeBLdsBlockDescriptorForTrLoad (multi-transform desc); a
writer-only enable misaligns the plain reader, which is why both
GetLdsPaddingConfigK and GetLdsPaddingConfigV stay disabled here.
- d >= 192 qr_tdm instances. Currently only d=128 qr_tdm instances
are precompiled; d=192 / d=256 fall back to "not supported yet"
on the dispatcher. Codegen extension required.
- Mask=1 sliding-window sweep (only mask=0 dense and mask=2 causal
are exercised here).
Retire the prior codegen workaround that disabled qr_vr emit to force qr_tdm dispatch. Instead, reorder pipelines in Gfx125xFactory so qr_tdm is emitted before qr (qr_vr). The generated C++ dispatch is an if/else-if chain where list order = priority, so qr_tdm is now naturally preferred when both pipelines match the runtime traits. qr (qr_vr) is re-enabled as fallback for trait combos not covered by qr_tdm (bias, dropout, skip, d!=128). Verified 7/7 valid:y via ffm_lite: - d=128 fp16/bf16 single-head + GQA → qr_tdm_vr_npad (target) - d=32/64 → native d32/d64 tile qr_vr (correct, no qr_tdm at d!=128) - d=192/256 → d256 qr_vr fallback (functional, qr_tdm not yet emitted)
Replace in-development references in the TDM-touched files with neutral technical descriptions. The labels were opaque to other contributors and added no semantic information.
…pported traits
The qr_tdm pipeline does not implement sink attention (sink_v is silently
ignored), bias, or dropout. Previously codegen iterated sink in {t,f},
so sink workloads at d=128 on gfx1250 would be dispatched to qr_tdm and
produce silently-wrong results.
Restrict the codegen to sink='f' so sink workloads fall through to qr,
and add static_asserts on BiasEnum / kHasDropout / kHasSink at the top
of the pipeline so any future codegen drift fails to compile rather
than silently mis-dispatching.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds a gfx1250-targeted CK Tile FMHA forward qr_tdm pipeline that routes eligible Q/K/V loads through TDM, updates kernel dram-view handling for that pipeline, and adjusts shared framework helpers needed by TDM tile windows.
Changes:
- Adds
qr_tdmFMHA pipeline, policy, enum mapping, and public aggregation. - Updates FMHA forward kernel dram view construction and dispatcher/codegen preference for eligible gfx1250 fp16/bf16 cases.
- Fixes tuple/container helper behavior for mixed compile-time/runtime tuples and improves TDM global stride caching for non-packed views.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm.hpp |
Implements the new qr_tdm FMHA forward pipeline. |
projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp |
Defines TDM-specific FMHA tile distributions, LDS descriptors, smem sizing, and padding config. |
projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_enum.hpp |
Adds the QRKSVS_TDM pipeline enum and string mapping. |
projects/composablekernel/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp |
Routes qr_tdm through the trload-style branch and bypasses XOR dram transforms for TDM. |
projects/composablekernel/include/ck_tile/ops/fmha.hpp |
Exposes the new TDM pipeline headers. |
projects/composablekernel/include/ck_tile/core/tensor/tile_window.hpp |
Computes cached global strides from descriptor offsets instead of packed lengths. |
projects/composablekernel/include/ck_tile/core/container/tuple.hpp |
Reworks tuple arithmetic to preserve mixed element types. |
projects/composablekernel/include/ck_tile/core/container/container_helper.hpp |
Reworks tuple reverse inclusive scan to support heterogeneous tuples. |
projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py |
Emits and prioritizes qr_tdm instances for supported gfx1250 shapes. |
projects/composablekernel/example/ck_tile/01_fmha/codegen/cpp_symbol_map.py |
Maps qr_tdm to its C++ pipeline type and enum. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Mirror baseline qr_ks_vs sink logic into the qr_tdm TDM pipeline: - Sink-aware M/L initialization (pre-seed when sink_v is finite) - GetSinkTileRangeAlongX for computing sink + normal tile ranges - Sink→normal K/V dram window jump at the boundary - IsOutOfSinkBound mask for per-pixel masking in sink region - Sink-aware k_origin computation in the main loop - Propagate sink_v through operator() → run() (both overloads) Codegen: re-enable sink iteration for qr_tdm instances (reverses the sink="f" restriction from commit 7591d80). Remove receipt 3 and BUILD_TESTING --filter constraints that blocked sink instance generation. Test: - A (fp16 dense s=1023 mask=0): kname qr_tdm ✓, valid:y, 2586 ms - B (fp16 GQA causal s=1023×257 mask=2): kname qr_tdm ✓, valid:y, 371 ms - C (bf16 dense s=2047 mask=0): kname qr_tdm ✓, valid:y, 9300 ms - Sink (fp16 s=1023×257 mask=2 init_sink=2): kname qr_tdm sink ✓, valid:y, 206 ms
…support Add attention bias support to the qr_tdm TDM pipeline, mirroring the baseline qr_ks_vs implementation: - ELEMENTWISE_BIAS: load bias tile from DRAM, pre-scale s_acc by scale_s, then add bias with log2e scaling (FAST_EXP2) or raw (standard path) - ALiBi: sweep s_acc tiles with position_encoding.update() per element - Create bias_dram_window from bias_dram_block_window_tmp with the gemm_0 C-tile distribution (both single-buffer and prefill paths) - Bias window movement per iteration + sink→normal jump Codegen: iterate BIAS_MAP for qr_tdm instances (no/bias/alibi). Add "bias" to receipt 3 allowed bias types. Dropout remains deferred (requires kernel dispatch interface expansion). Test: - A (fp16 dense s=1023 mask=0): kname qr_tdm ✓, valid:y, 2582 ms - B (fp16 GQA causal s=1023×257 mask=2): kname qr_tdm ✓, valid:y, 376 ms - C (bf16 dense s=2047 mask=0): kname qr_tdm ✓, valid:y, 9206 ms - Bias (fp16 s=1023 -bias=e): kname qr_tdm bias ✓, valid:y, 3080 ms - Sink regression: kname qr_tdm sink ✓, valid:y, 205 ms
Wire up the prefill path (kM0=128) on the gfx1250 TDM FMHA forward pipeline and fix the gemm_0-C -> gemm_1-A P relayout that was only correct for decode (MIterPerWarp==1). Addresses all reviewer comments: size prefill K LDS buffers with LoadOnce=true, match the K write window lengths to the LoadOnce=true view, and correct the stale V load comments.
gino-lu
marked this pull request as ready for review
June 30, 2026 04:06
✅ 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. |
gino-lu
changed the title
Replace unicode chars in comments with ASCII to pass static checks, and add a d=128 long-seqlen case to test_fmha_fwd that exercises the qr_tdm prefill tile on gfx1250.
gino-lu
changed the title
gino-lu
changed the title
joyeamd
reviewed
Jul 15, 2026
gino-lu
enabled auto-merge (squash)
July 20, 2026 16:27
joyeamd
approved these changes
Jul 21, 2026
assistant-librarian Bot
pushed a commit
to ROCm/composable_kernel
that referenced
this pull request
Jul 21, 2026
feat(ck-tile): add FMHA FWD TDM Pipeline
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
## Motivation
Bring up a TDM (Tensor Data Movement, box-major async DMA) variant of
the CK Tile FMHA forward pipeline for gfx1250, so that Q / K / V global
loads go through the new hardware DMA path instead of the generic
`async_load_tile` / `buffer_load` path used by `qr_vr`. The goal is to
land a functionally complete pipeline that the dispatcher
prefers on gfx1250 wherever it applies, as the basis for follow-up perf
work (LDS-padding for bank-conflict avoidance, wider hdim coverage,
sched tuning).
## Technical Details
New pipeline + policy (gfx1250-targeted):
-
`include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm.hpp`
-
`include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp`
- New entry in `block_fmha_pipeline_enum.hpp`; aggregation in
`ops/fmha.hpp`
- All three operands (Q, K, V) load via `load_tile_tdm` into plain
row-major LDS, drained by `s_wait_tensorcnt_barrier`. Q is loaded once
outside the K loop; prefill uses double-buffered K, decode uses a
single-buffer K with V async via `load_tile_tdm` as well.
Kernel changes (`include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp`):
- `make_q_dram` / `make_k_dram` / `make_v_dram` gain a `qr_tdm`-aware
early return: for the TDM pipeline the dram view skips the `unmerge ->
xor -> merge_v3` chain and returns the affine pad-only view directly.
TDM box-major DMA cannot honor a software XOR layout on the dram side;
the XOR transform chain stays intact for `qr_async_trload` via an `if
constexpr` branch.
- Unmerge typo fix in the K dram path (leading dim was `kQKHeaddim /
kDramTileK / kAlignmentK` which folded to 0 in common configs; same fix
as upstream PR #6964 that landed after this branch was cut).
Framework fixes (used by FMHA but not exclusive to it):
- `include/ck_tile/core/tensor/tile_window.hpp`:
`get_cached_global_strides()` no longer assumes a packed view. It now
queries the descriptor with a unit vector (`calculate_offset`) so
non-packed views (e.g. GQA `stride_q = h_q * hdim_q`, padded GEMM
`stride_a > K`) get the actual byte stride.
- `include/ck_tile/core/container/{tuple,container_helper}.hpp`:
`operator+/-/*` and `container_reverse_inclusive_scan` for tuples now
use `generate_tuple` so mixed `tuple<int, constant<N>>` operands work
in-place.
Dispatcher:
- `Gfx125xFactory` reorders pipelines so `qr_tdm` is emitted before `qr`
(`qr_vr`). The generated dispatch is an `if/else-if` chain, so list
order = priority; this retires an earlier codegen workaround that
disabled `qr_vr` emit to force `qr_tdm` selection. `qr_vr` is re-enabled
as the fallback for trait combos not covered by `qr_tdm` (bias, dropout,
sink, hdim != 128).
Codegen (`example/ck_tile/01_fmha/codegen/`):
- `qr_tdm` emit gated to `hdim == hdim_v == 128`, `dropout == "f"`. The
pipeline carries matching `static_assert`s so any future codegen drift
fails to compile rather than silently mis-dispatching.
## Test Plan
Verified on a gfx1250 simulation environment, all dispatching to
`qr_tdm_vr_npad...` kernels:
- fp16 dense: `b=1 h=1 s=1023 d=128 mask=0`
- fp16 GQA causal: `b=1 h_q=2 h_k=1 s=1023 s_k=257 d=128 mask=2`
- bf16 long: `b=1 h=1 s=2047 d=128 mask=0`
- Multi-stride GQA: `h={4,8,16} h_k=1 s=1023 s_k=257 d=128 mask=2`
(verifies stride 256/512/1024 all resolve under the new TDM dram
dispatch)
- Head dim sweep: `d={32,64}` fp16/bf16 (fall back to padded d=128
qr_tdm instance); `d=128` native instance; `d={192,256}` fall through to
`qr_vr` (qr_tdm not yet emitted)
- Non-TDM regression: fp8 / fp8bf16 routes that bypass `qr_tdm` on
gfx1250 dispatch to `qr_vr_psskddv` byte-identically pre- and
post-this-PR.
- Prefill tile dispatch (s >= 2048)
- fp16: `b=1 h=1 s=2048 d=128 mask=0`
- bf16 causal: `b=1 h=1 s=2048 d=128 mask=2`
For the framework `get_cached_global_strides` fix, the existing TDM
consumer binaries were re-baselined (`test_tdm_basic`,
`gemm_tdm_data_cache_prefetch` example, `gemm_weight_preshuffle_tdm`
example, `test_ck_tile_gemm_pipeline_tdm_wmma`,
`test_ck_tile_mx_gemm_pipeline_tdm_wmma`); no PASS↔FAIL flips, no
SKIP-set drift.
## Test Result
- 7/7 `valid:y` on gfx1250 simulation for the full d=128 / d=32 / d=64 /
d=192 / d=256 dispatch matrix above.
- Cycle-accurate run (`b=1 h=1 s=128 d=128 fp16`): `valid:y`, matching
the pre-TDM baseline on the same shape.
- Non-TDM kname suffix identical pre/post for fp8 and fp8bf16 paths.
- TDM framework consumer audit: identical FAIL identity sets on all 5
binaries vs pre-fix baseline.
-
## Known Limitations / Follow-ups
- `qr_tdm` instances are only emitted at `d=128`; `d>=192` falls back to
`qr_vr`. Codegen extension is straightforward but deferred.
- K / V LDS padding is disabled (`GetLdsPaddingConfigK/V` returns
`(false,0,0)`). A bank-conflict-avoidance perf mode requires a
padding-aware LDS descriptor mirroring
`MakeBLdsBlockDescriptorForTrLoad`; a writer-only enable misaligns the
plain reader, hence both are off until the reader side is taught
padding.
- `dropout` are unsupported in `qr_tdm` and are routed to `qr_vr` by
codegen (and asserted at the pipeline static_assert level).
- Mask=1 (sliding window) not exercised in this PR's verification
matrix.
## Submission Checklist
- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
shumway
pushed a commit
to ROCm/composable_kernel
that referenced
this pull request
Aug 18, 2026
feat(ck-tile): add FMHA FWD TDM Pipeline
## Motivation
Bring up a TDM (Tensor Data Movement, box-major async DMA) variant of
the CK Tile FMHA forward pipeline for gfx1250, so that Q / K / V global
loads go through the new hardware DMA path instead of the generic
`async_load_tile` / `buffer_load` path used by `qr_vr`. The goal is to
land a functionally complete pipeline that the dispatcher
prefers on gfx1250 wherever it applies, as the basis for follow-up perf
work (LDS-padding for bank-conflict avoidance, wider hdim coverage,
sched tuning).
## Technical Details
New pipeline + policy (gfx1250-targeted):
-
`include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm.hpp`
-
`include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp`
- New entry in `block_fmha_pipeline_enum.hpp`; aggregation in
`ops/fmha.hpp`
- All three operands (Q, K, V) load via `load_tile_tdm` into plain
row-major LDS, drained by `s_wait_tensorcnt_barrier`. Q is loaded once
outside the K loop; prefill uses double-buffered K, decode uses a
single-buffer K with V async via `load_tile_tdm` as well.
Kernel changes (`include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp`):
- `make_q_dram` / `make_k_dram` / `make_v_dram` gain a `qr_tdm`-aware
early return: for the TDM pipeline the dram view skips the `unmerge ->
xor -> merge_v3` chain and returns the affine pad-only view directly.
TDM box-major DMA cannot honor a software XOR layout on the dram side;
the XOR transform chain stays intact for `qr_async_trload` via an `if
constexpr` branch.
- Unmerge typo fix in the K dram path (leading dim was `kQKHeaddim /
kDramTileK / kAlignmentK` which folded to 0 in common configs; same fix
as upstream PR #6964 that landed after this branch was cut).
Framework fixes (used by FMHA but not exclusive to it):
- `include/ck_tile/core/tensor/tile_window.hpp`:
`get_cached_global_strides()` no longer assumes a packed view. It now
queries the descriptor with a unit vector (`calculate_offset`) so
non-packed views (e.g. GQA `stride_q = h_q * hdim_q`, padded GEMM
`stride_a > K`) get the actual byte stride.
- `include/ck_tile/core/container/{tuple,container_helper}.hpp`:
`operator+/-/*` and `container_reverse_inclusive_scan` for tuples now
use `generate_tuple` so mixed `tuple<int, constant<N>>` operands work
in-place.
Dispatcher:
- `Gfx125xFactory` reorders pipelines so `qr_tdm` is emitted before `qr`
(`qr_vr`). The generated dispatch is an `if/else-if` chain, so list
order = priority; this retires an earlier codegen workaround that
disabled `qr_vr` emit to force `qr_tdm` selection. `qr_vr` is re-enabled
as the fallback for trait combos not covered by `qr_tdm` (bias, dropout,
sink, hdim != 128).
Codegen (`example/ck_tile/01_fmha/codegen/`):
- `qr_tdm` emit gated to `hdim == hdim_v == 128`, `dropout == "f"`. The
pipeline carries matching `static_assert`s so any future codegen drift
fails to compile rather than silently mis-dispatching.
## Test Plan
Verified on a gfx1250 simulation environment, all dispatching to
`qr_tdm_vr_npad...` kernels:
- fp16 dense: `b=1 h=1 s=1023 d=128 mask=0`
- fp16 GQA causal: `b=1 h_q=2 h_k=1 s=1023 s_k=257 d=128 mask=2`
- bf16 long: `b=1 h=1 s=2047 d=128 mask=0`
- Multi-stride GQA: `h={4,8,16} h_k=1 s=1023 s_k=257 d=128 mask=2`
(verifies stride 256/512/1024 all resolve under the new TDM dram
dispatch)
- Head dim sweep: `d={32,64}` fp16/bf16 (fall back to padded d=128
qr_tdm instance); `d=128` native instance; `d={192,256}` fall through to
`qr_vr` (qr_tdm not yet emitted)
- Non-TDM regression: fp8 / fp8bf16 routes that bypass `qr_tdm` on
gfx1250 dispatch to `qr_vr_psskddv` byte-identically pre- and
post-this-PR.
- Prefill tile dispatch (s >= 2048)
- fp16: `b=1 h=1 s=2048 d=128 mask=0`
- bf16 causal: `b=1 h=1 s=2048 d=128 mask=2`
For the framework `get_cached_global_strides` fix, the existing TDM
consumer binaries were re-baselined (`test_tdm_basic`,
`gemm_tdm_data_cache_prefetch` example, `gemm_weight_preshuffle_tdm`
example, `test_ck_tile_gemm_pipeline_tdm_wmma`,
`test_ck_tile_mx_gemm_pipeline_tdm_wmma`); no PASS↔FAIL flips, no
SKIP-set drift.
## Test Result
- 7/7 `valid:y` on gfx1250 simulation for the full d=128 / d=32 / d=64 /
d=192 / d=256 dispatch matrix above.
- Cycle-accurate run (`b=1 h=1 s=128 d=128 fp16`): `valid:y`, matching
the pre-TDM baseline on the same shape.
- Non-TDM kname suffix identical pre/post for fp8 and fp8bf16 paths.
- TDM framework consumer audit: identical FAIL identity sets on all 5
binaries vs pre-fix baseline.
-
## Known Limitations / Follow-ups
- `qr_tdm` instances are only emitted at `d=128`; `d>=192` falls back to
`qr_vr`. Codegen extension is straightforward but deferred.
- K / V LDS padding is disabled (`GetLdsPaddingConfigK/V` returns
`(false,0,0)`). A bank-conflict-avoidance perf mode requires a
padding-aware LDS descriptor mirroring
`MakeBLdsBlockDescriptorForTrLoad`; a writer-only enable misaligns the
plain reader, hence both are off until the reader side is taught
padding.
- `dropout` are unsupported in `qr_tdm` and are routed to `qr_vr` by
codegen (and asserted at the pipeline static_assert level).
- Mask=1 (sliding window) not exercised in this PR's verification
matrix.
## Submission Checklist
- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
---------
Co-authored-by: Gino Lu <gino.lu@amd.com>
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.
Motivation
Bring up a TDM (Tensor Data Movement, box-major async DMA) variant of the CK Tile FMHA forward pipeline for gfx1250, so that Q / K / V global loads go through the new hardware DMA path instead of the generic
async_load_tile/buffer_loadpath used byqr_vr. The goal is to land a functionally complete pipeline that the dispatcherprefers on gfx1250 wherever it applies, as the basis for follow-up perf work (LDS-padding for bank-conflict avoidance, wider hdim coverage, sched tuning).
Technical Details
New pipeline + policy (gfx1250-targeted):
include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm.hppinclude/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hppblock_fmha_pipeline_enum.hpp; aggregation inops/fmha.hppload_tile_tdminto plain row-major LDS, drained bys_wait_tensorcnt_barrier. Q is loaded once outside the K loop; prefill uses double-buffered K, decode uses a single-buffer K with V async viaload_tile_tdmas well.Kernel changes (
include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp):make_q_dram/make_k_dram/make_v_dramgain aqr_tdm-aware early return: for the TDM pipeline the dram view skips theunmerge -> xor -> merge_v3chain and returns the affine pad-only view directly. TDM box-major DMA cannot honor a software XOR layout on the dram side; the XOR transform chain stays intact forqr_async_trloadvia anif constexprbranch.kQKHeaddim / kDramTileK / kAlignmentKwhich folded to 0 in common configs; same fix as upstream PR [CK_TILE] Fix typo in fmha_fwd_kernel K-dram unmerge tuple sizes #6964 that landed after this branch was cut).Framework fixes (used by FMHA but not exclusive to it):
include/ck_tile/core/tensor/tile_window.hpp:get_cached_global_strides()no longer assumes a packed view. It now queries the descriptor with a unit vector (calculate_offset) so non-packed views (e.g. GQAstride_q = h_q * hdim_q, padded GEMMstride_a > K) get the actual byte stride.include/ck_tile/core/container/{tuple,container_helper}.hpp:operator+/-/*andcontainer_reverse_inclusive_scanfor tuples now usegenerate_tupleso mixedtuple<int, constant<N>>operands work in-place.Dispatcher:
Gfx125xFactoryreorders pipelines soqr_tdmis emitted beforeqr(qr_vr). The generated dispatch is anif/else-ifchain, so list order = priority; this retires an earlier codegen workaround that disabledqr_vremit to forceqr_tdmselection.qr_vris re-enabled as the fallback for trait combos not covered byqr_tdm(bias, dropout, sink, hdim != 128).Codegen (
example/ck_tile/01_fmha/codegen/):qr_tdmemit gated tohdim == hdim_v == 128,dropout == "f". The pipeline carries matchingstatic_asserts so any future codegen drift fails to compile rather than silently mis-dispatching.Test Plan
Verified on a gfx1250 simulation environment, all dispatching to
qr_tdm_vr_npad...kernels:b=1 h=1 s=1023 d=128 mask=0b=1 h_q=2 h_k=1 s=1023 s_k=257 d=128 mask=2b=1 h=1 s=2047 d=128 mask=0h={4,8,16} h_k=1 s=1023 s_k=257 d=128 mask=2(verifies stride 256/512/1024 all resolve under the new TDM dram dispatch)d={32,64}fp16/bf16 (fall back to padded d=128 qr_tdm instance);d=128native instance;d={192,256}fall through toqr_vr(qr_tdm not yet emitted)qr_tdmon gfx1250 dispatch toqr_vr_psskddvbyte-identically pre- and post-this-PR.b=1 h=1 s=2048 d=128 mask=0b=1 h=1 s=2048 d=128 mask=2For the framework
get_cached_global_stridesfix, the existing TDM consumer binaries were re-baselined (test_tdm_basic,gemm_tdm_data_cache_prefetchexample,gemm_weight_preshuffle_tdmexample,test_ck_tile_gemm_pipeline_tdm_wmma,test_ck_tile_mx_gemm_pipeline_tdm_wmma); no PASS↔FAIL flips, no SKIP-set drift.Test Result
valid:yon gfx1250 simulation for the full d=128 / d=32 / d=64 / d=192 / d=256 dispatch matrix above.b=1 h=1 s=128 d=128 fp16):valid:y, matching the pre-TDM baseline on the same shape.Known Limitations / Follow-ups
qr_tdminstances are only emitted atd=128;d>=192falls back toqr_vr. Codegen extension is straightforward but deferred.GetLdsPaddingConfigK/Vreturns(false,0,0)). A bank-conflict-avoidance perf mode requires a padding-aware LDS descriptor mirroringMakeBLdsBlockDescriptorForTrLoad; a writer-only enable misaligns the plain reader, hence both are off until the reader side is taught padding.dropoutare unsupported inqr_tdmand are routed toqr_vrby codegen (and asserted at the pipeline static_assert level).Submission Checklist