Improve the performance of qr_ks_vs_whole_k_prefetch pipeline

[qianfengz]

About qr_ks_vs_whole_k_prefetch pipeline

This PR updates and enhances the qr_ks_vs_whole_k_prefetch pipeline to improve performance on both MI350 GPUs through better MFMA instruction usage, transposed V-loading support, and N0-loop implementation. The pipeline targets scenarios where the number of workgroups is low, enabling better CU occupancy by using smaller MTile sizes (kM0=64 vs 128) while prefetching entire K tiles.

Changes:

• Adds transposed V-loading support (qr_ks_vs_whole_k_prefetch_trload) to avoid using shuffle instructions on MI350
• Implements N0-loop based Gemm0 to reduce tile window movement overhead and eliminate clear_tile calls
• Adds full support for hdim96/hdim160 without padding requirements
• Updates MFMA instruction selection to ensure optimal choices for MI350

Performance results

1. For attention shapes which leads to kM0=64, qr_ks_vs_async_whole_k_prefetch_trload shows much better performance than qr_ks_vs_async_trload on the same case (execution time 41.02ms by whole_k_prefetch_trload & 58.50ms by async_load), and qr_ks_vs_async_whole_k_prefetch_trload also shows obviously better performance than the recently tuned qr_ks_vs_async on the same case (execution time 41.02ms by whole_k_prefetch_trload 7 47.60ms by qr_ks_vs_async)
2. Also on MI300, for attention shapes which leads to kM0=64, qr_ks_vs_async_whole_k_prefetch shows much better performance than the qr_ks_vs_async (which is supposed to be very high-efficient) on the same case (execution time 64.50ms by whole_k_prefetch & 80.20ms by qr_ks_vs_async)
3. For attention shapes which leads to kM0=128, qr_ks_vs_async_whole_k_prefetch_trload show a little bit better performance than qr_ks_vs_async on mi350 (execution time 104.50ms by whole_k_prefetch_trload & 106.50ms by qr_ks_vs_async). And they shows completely on-par performance on MI300

Test/Verify

1. Use the ROCM xformers branch test_whole_k_prefetch_n0loop to test/verify qr_ks_vs_whole_k_prefetch pipeline since this pipeline can not be used by ck_tile fmha example so far
2. Use the following command-line for building/testing xformers

#> git clone -b test_whole_k_prefetch_n0loop https://github.com/ROCm/xformers
#> git submodule update --init --recursive   
#> pip  install --no-build-isolation -e ./
#> pytest tests/test_mem_eff_attention.py::test_forward

4. Any scripts which can run on xformers can be used to evaluate qr_ks_vs_whole_k_prefetch pipeline. Using the two environ variable to switch from using different pipelines

#> export FMHA_DISABLE_SPECIAL_TREATMENT=1              #> to disable using FAV3 and qr_ks_vs_async_trload pipeline
#> export FMHA_ENABLE_ASYNC_PIPELINE=1                     #>  to disable using qr_ks_vs_async pipeline for comparing

Discussion

[qianfengz]

This PR is a re-submit of PR-6087

[qianfengz]

@copilot Can you transfer the comments history from PR 6087 to current PR ?

[Copilot]

Pull request overview

This PR enhances the qr_ks_vs_whole_k_prefetch FMHA forward pipeline for improved GPU occupancy and performance (notably on MI350) by introducing an N0-loop formulation, updated GEMM micro-kernels/policies, and an additional trload-based V-loading pipeline variant.

Changes:

• Added new block GEMM variants to support N- and K-prefetch schemes, including a transposed-load (trload) B path for V GEMM.
• Reworked the whole-K-prefetch FMHA pipeline to use an N0-loop (kN0 subdivided by kK0) and to support non-power-of-two head dims (e.g., 96/160) without padding.
• Updated the FMHA forward kernel to support “naive head-dim load”, optional FAST_EXP2 bypass, and a trload V path, plus plumbing for seqstart_v_scale_ptr.

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
projects/composablekernel/include/ck_tile/ops/gemm/block/block_gemm_areg_bsmem_trload_creg_v2_prefetch_n.hpp	New GEMM kernel variant supporting transposed B loads with N-prefetch.
projects/composablekernel/include/ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v2_prefetch_n.hpp	New GEMM kernel variant implementing N-prefetch.
projects/composablekernel/include/ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v2_prefetch_k.hpp	New GEMM kernel variant implementing K-prefetch.
projects/composablekernel/include/ck_tile/ops/fmha/pipeline/tile_fmha_shape.hpp	Minor formatting and static_assert message update.
projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qs_ks_vs.hpp	Extends pipeline operator signature to accept sink_v (ignored).
projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_whole_k_prefetch.hpp	Major pipeline rewrite: N0-loop, changed K/V LDS usage, revised softmax flow.
projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_whole_k_prefetch_trload.hpp	New trload-based whole-K-prefetch pipeline to avoid shuffle overhead.
projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_whole_k_prefetch_default_policy.hpp	Reworked policy: new GEMM selection, new LDS descriptors/layout, KV shared LDS sizing.
projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_problem.hpp	Adds vector-size helpers for DRAM tile access decisions.
projects/composablekernel/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp	Adds pipeline feature detection (n0loop/fast-exp2/naive-hdim/trload), updates Q/K/V window padding and scale behavior.
projects/composablekernel/include/ck_tile/ops/fmha.hpp	Wires in the new whole-K-prefetch trload pipeline header.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Typo in comment: “constrcut” → “construct”.

Suggested change

	// constrcut from A-block-tensor from A-Block-tensor-tmp
	// construct from A-block-tensor from A-Block-tensor-tmp

[Copilot]

Typo in comment: “constrcut” → “construct”.

Suggested change

	// constrcut from A-block-tensor from A-Block-tensor-tmp
	// construct from A-block-tensor from A-Block-tensor-tmp

[Copilot]

@copilot Can you transfer the comments history from PR 6087 to current PR ?

Here is a summary of the review comments from PR #6087 transferred to this PR:

---

Review comments from PR #6087

1. rotating_buffers.hpp — #include <iostream>

• @asleepzzz requested removing #include <iostream> (removed in prior commit).
• @qianfengz noted that removal causes a compilation failure when building from xformers, because std::cout is used inside the file and iostream isn't implicitly pulled in that build context.
• Fixed in commit 8931530b: #include <iostream> restored.

2. fmha_fwd_kernel.hpp line 105 — trload protection

• @asleepzzz asked "how to protect 942 pipeline set trload true".
• Resolved: static_assert(!Problem::kUseTrLoad, "This pipeline does not use trload!") is present in BlockFmhaPipelineQRKSVSWholeKPrefetch, mirroring the corresponding assert in the trload variant.

3. tile_fmha_shape.hpp — kN0Sub placement

• @asleepzzz requested splitting kN0Sub out of TileFmhaShape since it is only needed by the whole_k_prefetch pipelines.
• Resolved: kN0Sub is no longer in TileFmhaShape; it is defined locally inside the two whole_k_prefetch pipelines.

4. block_fmha_pipeline_qr_ks_vs_whole_k_prefetch_trload.hpp — pipeline name

• @asleepzzz questioned the pipeline name ("qr_async" conflicting with existing pipelines).
• Resolved: Pipeline name is "qr_async_whole_k_prefetch_trload".

5. block_fmha_pipeline_qr_ks_vs_whole_k_prefetch_default_policy.hpp line 142 — missing include

• copilot-reviewer noted that GetAlignmentK()/GetAlignmentV() use detail::GetDramTileAccessMaxVectorSize defined in block_fmha_pipeline_problem.hpp, which was not directly included in the policy header.
• Fixed in commit 8931530b: Added #include "ck_tile/ops/fmha/pipeline/block_fmha_pipeline_problem.hpp" to the policy header.

6. GetSmemSizeDropout() — static_assert for dropout

• copilot-reviewer flagged GetSmemSizeDropout() returning 0 unconditionally while dropout LDS may be needed.
• @poyenc asked "@copilot you add the static_assert()" to make it explicit this pipeline never handles dropout.
• Status: static_assert(!Problem::kHasDropout, ...) was not added in the current re-submitted PR. @poyenc's request remains open.

7. block_fmha_pipeline_qr_ks_vs_whole_k_prefetch.hpp line 307 — missing early exit

• @poyenc asked "why removing the early exit block?" (the num_total_loop <= 0 check used in other pipelines).
• Response: The whole_k_prefetch pipeline never had an early exit in the original design. This is missing from the original pipeline rather than something removed by this PR.

8. Wrong type args in GetKVBlockGemm

• @poyenc noted QDataType→PDataType and KDataType→VDataType were wrong.
• Resolved: Corrected in the block gemm policy.

9. GetNumPrefetchV() — simplification

• @poyenc asked for a one-line return 0; implementation (the comment was about a GetSmemSizeDropout-adjacent function; the actual simplification was to return 2; since all branches returned 2).
• Resolved: GetNumPrefetchV() has the full if constexpr body restored per the original design.

10. Typo fixes

• "generatd" → "generate" in pipeline comments — Resolved.
• "constrcut" → "construct" in block_gemm header — Resolved.
• "randomm" → "random" in trload pipeline — Resolved.

[poyenc]

Thanks for the performance work — the numbers for kM0=64 shapes are impressive (20-30% gains). A few items to address:

Blocking

1. Early-exit guard removed from the pipeline

The do { ... } while(seqlen_k_curr < seqlen_k_end) loop in block_fmha_pipeline_qr_ks_vs_whole_k_prefetch.hpp always executes at least once, even when seqlen_k_start >= seqlen_k_end. The previous version had an early-exit check. Without it, the first iteration will load K/V tiles from invalid ranges when there's zero work (masking or padded seqlen_k cases). Please re-add the guard:

if constexpr(FmhaMask::IsMasking || kPadSeqLenK)
{
    if(seqlen_k_start >= seqlen_k_end)
        return o_acc;
}

2. kIsAvailable guard removed without replacement

The kernel's operator() now unconditionally calls run_(). On non-gfx950 targets, load_tile_transpose silently returns zeros (see buffer_view.hpp:897 fallback). The old if constexpr(kIsAvailable) guard prevented this silent data corruption. Please add a compile-time guard in the trload pipeline, e.g.:

#if !defined(__gfx950__)
static_assert(sizeof(QDataType) == 0,
    "whole_k_prefetch_trload requires gfx950; "
    "load_tile_transpose returns zeros on other architectures.");
#endif

Suggestions

3. Pipeline-specific if constexpr in the kernel

The 4 new SFINAE trait detectors (is_n0loop_pipeline_v, is_naive_hdim_load_v, is_using_trload_v, ignore_fast_exp2_v) add 8 if constexpr branches into fmha_fwd_kernel.hpp, coupling the kernel to pipeline internals. Consider exposing these as direct constants on the pipeline type instead of SFINAE detection:

// Each pipeline defines directly:
static constexpr index_t kN0ForKView = ...;
static constexpr index_t kHeaddimForQKView = ...;
static constexpr bool kTransformScaleS = ...;

This eliminates the detail namespace machinery and makes future pipelines easier to add without touching the kernel.

4. kIgnoreFastExp2 needs a comment

This is a new concept with non-obvious correctness implications. A brief comment at the definition site explaining why it's safe (softmax is base-invariant) would help future readers.

Nits

5. }; // namespace detail → } // namespace detail (two locations — fmha_fwd_kernel.hpp, block_fmha_pipeline_problem.hpp). Zero precedent for ; in ck_tile/.

6. static_assert(NWarp == 1, "Check failed!") in block_gemm_areg_bsmem_creg_v2_prefetch_k.hpp — please use a descriptive message.

7. #if 0 dead code in block_gemm_areg_bsmem_creg_v2_prefetch_k.hpp — the sibling files (prefetch_n, trload_prefetch_n) already cleaned this up. Please remove and keep only a one-line comment about the register spill issue.

[qianfengz]

Thanks for the performance work — the numbers for kM0=64 shapes are impressive (20-30% gains). A few items to address:

Blocking

1. Early-exit guard removed from the pipeline

The do { ... } while(seqlen_k_curr < seqlen_k_end) loop in block_fmha_pipeline_qr_ks_vs_whole_k_prefetch.hpp always executes at least once, even when seqlen_k_start >= seqlen_k_end. The previous version had an early-exit check. Without it, the first iteration will load K/V tiles from invalid ranges when there's zero work (masking or padded seqlen_k cases). Please re-add the guard:

if constexpr(FmhaMask::IsMasking || kPadSeqLenK)
{
    if(seqlen_k_start >= seqlen_k_end)
        return o_acc;
}

2. kIsAvailable guard removed without replacement

The kernel's operator() now unconditionally calls run_(). On non-gfx950 targets, load_tile_transpose silently returns zeros (see buffer_view.hpp:897 fallback). The old if constexpr(kIsAvailable) guard prevented this silent data corruption. Please add a compile-time guard in the trload pipeline, e.g.:

#if !defined(__gfx950__)
static_assert(sizeof(QDataType) == 0,
    "whole_k_prefetch_trload requires gfx950; "
    "load_tile_transpose returns zeros on other architectures.");
#endif

Suggestions

3. Pipeline-specific if constexpr in the kernel

The 4 new SFINAE trait detectors (is_n0loop_pipeline_v, is_naive_hdim_load_v, is_using_trload_v, ignore_fast_exp2_v) add 8 if constexpr branches into fmha_fwd_kernel.hpp, coupling the kernel to pipeline internals. Consider exposing these as direct constants on the pipeline type instead of SFINAE detection:

// Each pipeline defines directly:
static constexpr index_t kN0ForKView = ...;
static constexpr index_t kHeaddimForQKView = ...;
static constexpr bool kTransformScaleS = ...;

This eliminates the detail namespace machinery and makes future pipelines easier to add without touching the kernel.

4. kIgnoreFastExp2 needs a comment

This is a new concept with non-obvious correctness implications. A brief comment at the definition site explaining why it's safe (softmax is base-invariant) would help future readers.

Nits

5. }; // namespace detail → } // namespace detail (two locations — fmha_fwd_kernel.hpp, block_fmha_pipeline_problem.hpp). Zero precedent for ; in ck_tile/.

6. static_assert(NWarp == 1, "Check failed!") in block_gemm_areg_bsmem_creg_v2_prefetch_k.hpp — please use a descriptive message.

7. #if 0 dead code in block_gemm_areg_bsmem_creg_v2_prefetch_k.hpp — the sibling files (prefetch_n, trload_prefetch_n) already cleaned this up. Please remove and keep only a one-line comment about the register spill issue.

3. left unchanged, others are solved

[poyenc]

All blocking and correctness items from the prior review are addressed. The remaining open items (SFINAE refactor, assert message) are non-blocking style/maintainability suggestions that can be tackled separately.

[eidenyoshida]

Override merge due to test failure on previous mici run reporting after a pass on latest run.