perf(ck_tile/fmha): [gfx1250 opt 5] optimize qr_tdm_pingpong, extend it to every head dim and simplify the dispatch

projects/composablekernel/example/ck_tile/01_fmha/README.md

@@ -43,32 +43,50 @@ There are 2 template parameters for this kernel template.
 ## codegen
 To speed up compile time, we instantiate the kernels into separate file. In this way we can benefit from parallel building from CMake/Make system. This is achieved by `generate.py` script. Besides, you can look into this script to learn how to instantiate a kernel instance step by step, which is described in `FMHA_FWD_KERNEL_BODY` variable.
 
+### Ping-pong TDM forward pipeline
+
+`qr_tdm_pingpong` (gfx125, FP16/BF16 Q/K/V) is the default forward pipeline for
+head dims 32, 64, 96, 128, 160, 192/128, 256 and 512, at any query and key length.
+It stages a whole K tile per trip in two LDS buffers: trip i reads one buffer while
+the TDM fills the other with tile i + 1. On the 128-row tiles V has two buffers as well,
+issued a trip ahead. Two or more resident blocks per CU overlap one block's softmax
+with another's gemms on the same SIMD; `qr_tdm` and `qr_tdm_sched` instead prefetch
+and schedule within a block.
+
+It has 16-, 64- and 128-row query tiles (d256/d512 have no 128-row tile); the tile rules
+are `_gfx125_qr_tdm_pingpong_tiles` in `codegen/ops/fmha_fwd.py`. Q/K rows are staged at
+the head dim's own length (d32 is not padded to 64).
+
+The LDS read schedule is chosen per head dim (`QrTdmPingPongLdsReadSchedule`,
+`GFX125_QR_TDM_PINGPONG_LDS_READ_SCHEDULE` in `codegen/ops/fmha_fwd.py`): prefetching
+the next K/V slice while the current slice's WMMAs run, and reading V's first slice
+before the softmax, each from the tile size on where its VGPRs leave the occupancy
+unchanged. Instance names carry `_sp<rows>` / `_vbs<rows>` for the parts that take effect
+on their tile.
+
+It supports masks, LSE, logits soft cap, elementwise bias and ALiBi, dropout (not on
+the 16-row tile), sink tokens, column-major V, batch and group mode. The one query it
+leaves to `qr_tdm_sched` is d128 with at least 128 query rows and fewer than two
+128-row query blocks per CU (`batch * nhead_q * max_seqlen_q < 256 * num_cus`), where a
+single resident block cannot overlap softmax and gemms. Among its head dims `qr_tdm` is
+generated for d128 only, as the fallback of the calls `qr_tdm_sched` refuses; skip goes
+to the generic `qr` pipeline.
+
 ### Scheduled TDM forward pipeline
 
 `qr_tdm_sched` shares the gfx125 TDM architecture requirements with `qr_tdm`.
-Its generated configurations use FP16 or BF16 Q/K/V with D64/V64, D128/V128,
-or D192/V128, and row-major V. They use no bias, dropout, logits
-soft cap, quantization, or sink tokens; LSE output is optional. Other configurations retain
-the existing pipeline candidates.
-Generated `qr_tdm_sched` requires a null runtime `sink_ptr`; a finite virtual-sink
-logit uses the existing `qr_tdm` fallback. This runtime pointer is independent
-of the compile-time sink-token flag.
-
-The generated API tries matching scheduled traits first. D64/V64 uses a
-128-by-64 Q/K sequence tile, matching the existing `qr_tdm` M128 geometry and
-compiler minimum-waves-per-EU hint of three. The launch API retains the
-historical `MinBlockPerCu` name; this hint does not establish achieved workgroup
-residency. Aligned batch K lengths use unpadded
-sequence traits; tails and group mode use padded traits. D64 uses compiler-scheduled
-softmax maximum chains and places each next-stage LDS read before the current
-WMMA, preserving read counts and stage-tail waits. Generated D64/D128/D192
-scheduled sources enable expert scheduling mode. D128 uses a
-128-by-128 Q/K sequence tile. D192 tries a 128-by-128 tile when average K length
-is at least 512, the Q grid fits within the CU count, logical length pointers
-are absent, and the mask is dense or square causal. Otherwise it uses a
-128-by-64 tile. Group mode compares packed K length against 512 times the batch
-count; its square-causal check compares packed Q/K totals. When
-`max_seqlen_q` is below 128, the existing fallback candidates are used.
+On gfx125 it is generated for FP16/BF16 D128/V128 with row-major V and serves the
+small-grid d128 queries `qr_tdm_pingpong` leaves to it (see above). It takes masks,
+LSE, logits soft cap, elementwise bias and ALiBi, dropout and sink tokens. Its
+variants without sink tokens require a null runtime `sink_ptr`; a sink logit without
+sink tokens uses the `qr_tdm` fallback. This runtime pointer is independent of the
+compile-time sink-token flag. The pipeline also implements D64/V64 and D192/V128
+geometries (described below), which are not generated on gfx125 since
+`qr_tdm_pingpong` serves those head dims.
+
+D128 uses a 128-by-128 Q/K sequence tile. Aligned batch K lengths use unpadded
+sequence traits; tails and group mode use padded traits. Generated scheduled sources
+enable expert scheduling mode.
 
 WG indexing follows the existing shared kernel rules. For batch BHSD square
 causal attention with one V tile and no dropout, sched D64/D128 visit the longest
@@ -123,8 +141,7 @@ counts, and rescale modes.
 For a new placement, supply a policy's third template argument (the tile
 schedule) and implement its local callbacks. The executor executes those
 callbacks directly. TDM loads and hardware synchronization remain explicit
-pipeline phase boundaries. The pipeline and policy reject attention bias;
-other pipelines provide its existing dispatch fallback. Sink operations remain
+pipeline phase boundaries. Sink operations remain
 for the shared kernel entry. A runtime sink logit is independent of the
 compile-time sink-token flag.
 
@@ -220,6 +237,13 @@ Currently we support both `batch mode` and `group mode` (or `varlen`, in FA's te
 ### MQA/GQA
 By setting `-h`(nhead for q) and `-h_k`(nhead for k/v) with different number, you can achieve MQA/GQA. Please pay attention that `h % h_K == 0` when you set different numbers.
 
+`-pack_gqa=1` (the default) folds the `h / h_k` query heads that share a K/V head into the
+query length for non-causal batch-mode runs without bias or sink, so each K/V tile is
+loaded once per group. On gfx125 fp16/bf16 with `d_v >= 64` it is skipped (with a note)
+for 2..16 query rows that it would fold into more than 16 rows on a grid under two
+(batch, head) pairs per CU: the query would leave the 16-row decode tile for a tile with
+`h / h_k` times fewer blocks.
+
 ### input/output permute, and `b*s*3*h*d`
 If you look at the kernel argument inside `fmha_fwd_kernel.hpp`, we support providing arbitrary stride for seqlen(stride_q/k/v), nhead, batch of q/k/v matrix, hence it is very flexible to support `b*h*s*d` or `b*s*h*d` input/output permute. The `-iperm=0/1`, `-operm=0/1` is a convenient way to achieve this through the executable. We didn't provide a command-line arg to test `b*s*3*h*d` layout which is by default used by torch/FA, but it's trivial to achieve this if one set the proper `stride_q/k/v` value as `3*h*d`.
 

projects/composablekernel/example/ck_tile/01_fmha/cmake/fmha_fwd_source_flags.cmake

@@ -3,11 +3,10 @@
 
 function(ck_tile_fmha_fwd_get_source_compile_options fwd_blob_name out_var)
   set(compile_options)
-  # Keep the compiler scheduling mode used by the validated qr_tdm_sched and
-  # qr_tdm_pingpong instances. The qr_tdm_pingpong instances that store the LSE
-  # keep the default: built in this mode, which leaves the ALU dependency checks
-  # to the compiler, they produce intermittently wrong O and LSE values.
-  if(fwd_blob_name MATCHES "^fmha_fwd_d(64|128|192)_(fp16|bf16)_.*_qr_tdm_(sched|pingpong_.*_nlse)(_.*)?_gfx125\\.cpp$")
+  # Keep the compiler scheduling mode used by the validated qr_tdm_sched instances,
+  # and use it for every qr_tdm_pingpong instance.
+  if(fwd_blob_name MATCHES "^fmha_fwd_d(64|128|192)_(fp16|bf16)_.*_qr_tdm_sched(_.*)?_gfx125\\.cpp$" OR
+     fwd_blob_name MATCHES "^fmha_fwd_d[0-9]+_(fp16|bf16)_.*_qr_tdm_pingpong(_.*)?_gfx125\\.cpp$")
     list(APPEND compile_options -mllvm -amdgpu-expert-scheduling-mode)
   endif()
   set(${out_var} "${compile_options}" PARENT_SCOPE)