[rocm-libraries] ROCm/rocm-libraries#7755 (commit e5143a9)

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.