From a60f83a782cdfa90a76b0e31ab30304a986e3d25 Mon Sep 17 00:00:00 2001 From: yanxishi Date: Sun, 11 Oct 2026 02:30:55 +0000 Subject: [PATCH 01/10] perf(fmha/fwd): 16-row short-query tile for the gfx1250 qr_tdm forward At max_seqlen_q <= 16 the 64-row tile idles at least 3/4 of its M rows. Add a d=128 fp16/bf16 tile of 16 rows that spends its four warps on N (1x4x1 for both gemms), dispatched first for max_seqlen_q <= 16 and only paired with qr_tdm: qr's gemm_1 slices P along M and cannot form a slice narrower than one warp tile. Warps on N route S through LDS, a path no shipped qr_tdm tile exercised, and it needs three fixes: - MakeSRegTileDistribution built S's layout from the QK gemm's warp config, but S is consumed by gemm_1 as its A operand. The layouts differ, so P's columns reached gemm_1 permuted within each row: the row sum and LSE stay right while O is off by ~1%. Use the P distribution, which embeds the PV warp gemm's A encoding. - On the double K/V LDS buffer path S aliased K0, which the K prefetch two iterations ahead overwrites. Give S its own slot after V1 when it is in LDS at all (kSBytes > 0); tiles with the warps on M are unchanged. - The progressive K reload walks four N-subtiles per warp and needs the gemm_0 warps all on M, so it is only enabled for kNWarp == 1. The LDS padding selection also accepts the 16-row, warps-on-N tile, which otherwise ran with unpadded K/V rows. test_qr_tdm_lds_padding checks the padding selection and the arena (S at 71680, 76032 bytes) for that tile. gfx1250, bf16, d=128, bshd, b=1280 h=2, TFLOPS develop -> this: s_q=8 s_k=532 96.6 -> 111.5 (1.15x), causal 95.9 -> 110.3 (1.15x) s_q=8 s_k=232 73.7 -> 87.0 (1.18x), causal 73.3 -> 85.3 (1.16x) Shapes with s_q > 16 keep their tiles and are unchanged. --- .../ck_tile/01_fmha/codegen/ops/fmha_fwd.py | 20 +++++- .../block_fmha_pipeline_qr_ks_vs_tdm.hpp | 14 ++-- ...lock_fmha_pipeline_qr_ks_vs_tdm_policy.hpp | 72 +++++++++---------- .../ck_tile/fmha/test_qr_tdm_lds_padding.cpp | 40 +++++++++++ 4 files changed, 100 insertions(+), 46 deletions(-) diff --git a/projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py b/projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py index f657b5b50ec0..8de617e13d20 100644 --- a/projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py +++ b/projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py @@ -1680,7 +1680,19 @@ def check_gemm0_k( return k0_loops == 1 return k0_loops >= 2 - rules.extend([check_gemm0_k, check_qr_tdm_sched_codegen_compatibility]) + def check_short_query_tile( + problem_ctx: ProblemContext, kernel_ctx: KernelContext + ) -> bool: + # The 16-row tile spends its warps on N, which only qr_tdm handles: + # qr's gemm_1 slices P along M and cannot form a slice narrower than + # one warp tile, so pairing the tile with qr fails to compile. + if kernel_ctx.tile.F_bm0 != 16: + return True + return kernel_ctx.pipeline.tag == "qr_tdm" + + rules.extend( + [check_gemm0_k, check_qr_tdm_sched_codegen_compatibility, check_short_query_tile] + ) return rules @classmethod @@ -1716,7 +1728,11 @@ def get_hdim_tile_size_dict(cls, dtype: str) -> Optional[dict]: FmhaFwdTileSize(128, 64, 32, 64, 32, 64, 4, 1, 1, 4, 1, 1, 16, 16, 32, 16, 16, 32, -1)], ( 96, 96) : [FmhaFwdTileSize( 64, 64, 32, 96, 32, 96, 4, 1, 1, 4, 1, 1, 16, 16, 32, 16, 16, 32, -1, _gfx125_qr_tdm_crossover_constraint(( 96, 96))), FmhaFwdTileSize(128, 64, 32, 96, 32, 96, 4, 1, 1, 4, 1, 1, 16, 16, 32, 16, 16, 32, -1)], - (128, 128) : [FmhaFwdTileSize( 64, 64, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 32, 16, 16, 32, -1, CppConstraint(m64_d128_constraint)), + # Short query (s_q <= 16): the 64-row tile would idle 3/4 of its M + # rows, so this tile spends its four warps on N instead (1,4,1). + # Listed first: dispatch takes the first tile whose constraint holds. + (128, 128) : [FmhaFwdTileSize( 16, 64, 32, 128, 32, 128, 1, 4, 1, 1, 4, 1, 16, 16, 32, 16, 16, 32, -1, CppConstraint("a.max_seqlen_q <= 16")), + FmhaFwdTileSize( 64, 64, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 32, 16, 16, 32, -1, CppConstraint(m64_d128_constraint)), FmhaFwdTileSize(128, 64, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 16, 16, 32, 16, 16, 32, -1)], (160, 160) : [FmhaFwdTileSize( 64, 64, 32, 160, 32, 160, 4, 1, 1, 4, 1, 1, 16, 16, 32, 16, 16, 32, -1, _gfx125_qr_tdm_crossover_constraint((160, 160))), FmhaFwdTileSize(128, 64, 32, 160, 32, 160, 4, 1, 1, 4, 1, 1, 16, 16, 32, 16, 16, 32, -1)], diff --git a/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm.hpp b/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm.hpp index 033806ba37c5..4890fe70f155 100644 --- a/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm.hpp +++ b/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm.hpp @@ -362,8 +362,13 @@ struct BlockFmhaPipelineQRKSVSTdm // The progressive M128 traversal has the fragment shape required by the // staged hand-off below. Element types and attention features do not alter // those LDS fragment lifetimes. + // The progressive K reload walks four N-subtiles per warp, so it needs the + // gemm_0 warps all on M. Tiles that spend their warps on N (the short-query + // tile) keep the plain whole-slice K reload. + static constexpr bool kProgressiveDsLoadK = Problem::kProgressiveDsLoadK && kNWarp == 1; + static constexpr bool kStagedKPairs = - Problem::kProgressiveDsLoadK && kM0 == 128 && kQKHeaddim == 128 && kK0 == 32; + kProgressiveDsLoadK && kM0 == 128 && kQKHeaddim == 128 && kK0 == 32; static constexpr bool kUseIglpBulkScheduling = CK_TILE_FMHA_TDM_IGLP_BULK && !kStagedKPairs; // Unchecked TDM bounds need a stronger, independent range proof. Keep this @@ -1462,8 +1467,7 @@ struct BlockFmhaPipelineQRKSVSTdm // S tile in LDS auto s_lds = make_tensor_view( - reinterpret_cast(reinterpret_cast(smem_ptrk0) + - Policy::template GetSmemSizeK()), + reinterpret_cast(static_cast(smem_arena) + Layout::kSOffset), Policy::template MakeSLdsBlockDescriptor()); auto s_write_lds_window = make_tile_window( s_lds, Policy::template MakeSLdsBlockDescriptor().get_lengths(), {0, 0}); @@ -1673,7 +1677,7 @@ struct BlockFmhaPipelineQRKSVSTdm // loop over along the [K]ey head dimension move_tile_window(k_lds_read_window, {0, kK0}); using Gemm0 = remove_cvref_t; - if constexpr(Problem::kProgressiveDsLoadK) + if constexpr(kProgressiveDsLoadK) { static_assert( std::is_same_v); @@ -1730,7 +1734,7 @@ struct BlockFmhaPipelineQRKSVSTdm move_tile_window(k_lds_read_window, {0, -kK0 * (k0_loops - 1)}); } - if constexpr(!kStagedKPairs && Problem::kProgressiveDsLoadK) + if constexpr(!kStagedKPairs && kProgressiveDsLoadK) { // Retain the reload-overlap traversal when consuming the final // head-dimension slice, even though this slice has no reload. diff --git a/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp b/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp index 47aa4f8e0d2b..fed6ef7b6393 100644 --- a/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp +++ b/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp @@ -211,7 +211,8 @@ inline constexpr bool is_qr_tdm_padding_enabled_problem_v = std::is_same_v && std::is_same_v)) && Problem::BlockFmhaShape::NumWarps == 4 && - (Problem::BlockFmhaShape::kM0 == 64 || Problem::BlockFmhaShape::kM0 == 128) && + (Problem::BlockFmhaShape::kM0 == 16 || Problem::BlockFmhaShape::kM0 == 64 || + Problem::BlockFmhaShape::kM0 == 128) && Problem::BlockFmhaShape::kN0 == 64 && Problem::BlockFmhaShape::kK0 == 32 && Problem::BlockFmhaShape::kK1 == 32 && // Head dim must be stored in LDS without tile-length rounding (excludes the @@ -230,8 +231,11 @@ inline constexpr bool is_qr_tdm_padding_enabled_problem_v = (Problem::BlockFmhaShape::kN1 * sizeof(typename Problem::VDataType)) % kQrTdmLdsAccessBytes == 0 && Problem::BlockFmhaShape::IsVLayoutRowMajor && - std::is_same_v> && - std::is_same_v> && + // warps all on M, or all on N for the short-query tile + (std::is_same_v> || + std::is_same_v>) && + std::is_same_v && std::is_same_v> && std::is_same_v> && numeric_traits::PackedSize == 1 && @@ -803,37 +807,14 @@ struct BlockFmhaPipelineQRKSVSTdmDefaultPolicy template CK_TILE_HOST_DEVICE static constexpr auto MakeSRegTileDistribution() { - using BlockGemm = remove_cvref_t())>; - - constexpr auto config = BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); - using WG = remove_cvref_t())>; - constexpr index_t MWarp = config.template at<1>(); - constexpr index_t NWarp = config.template at<2>(); - - constexpr index_t kMPerBlock = Problem::BlockFmhaShape::kM0; - constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK1; - constexpr index_t kTileK = Problem::BlockFmhaShape::kN0; - - // K2 is equal to Impl::kABKPerLane * kKIterPerWarpGemm - constexpr index_t K3 = WG::kK / WG::WarpGemmAttribute::Impl::kABKLane; - constexpr index_t K2 = WG::WarpGemmAttribute::Impl::kABKLane; - constexpr index_t K1 = kKPerBlock / (K2 * K3); - constexpr index_t K0 = kTileK / kKPerBlock; - constexpr index_t M2 = WG::WarpGemmAttribute::Impl::kAMLane; - constexpr index_t M1 = MWarp; - constexpr index_t M0 = kMPerBlock / (M2 * M1); - - constexpr auto s2_block_dstr_encoding = - tile_distribution_encoding, - tuple, sequence>, - tuple, sequence<2, 1>>, - tuple, sequence<2, 2>>, - sequence<1, 2, 2, 2>, - sequence<0, 0, 1, 3>>{}; - - constexpr auto s2_block_dstr = make_static_tile_distribution(s2_block_dstr_encoding); - - return s2_block_dstr; + // S comes back out of LDS to be softmaxed and then fed into gemm_1 as its A + // operand, so it has to land in exactly the layout gemm_1 expects. Building + // the encoding from the QK block gemm's warp config -- a different gemm than + // the PV one that consumes S -- permutes P's columns within a row: the row + // sum (and hence LSE) still comes out right, but O picks up a ~1% error + // because each key is paired with the wrong value row. Reuse the P + // distribution instead, which embeds the PV warp gemm's own A encoding. + return MakePRegTileDistribution(); } template @@ -936,9 +917,9 @@ struct QrTdmLdsArenaLayout static constexpr index_t kK0Offset = 0; static constexpr index_t kK1Offset = kUseDoubleKVLdsBuffer ? integer_least_multiple(kK0Offset + kKBytes, kRegionAlignment) : 0; - static constexpr index_t kSOffset = - kUseDoubleKVLdsBuffer ? 0 - : integer_least_multiple(kK0Offset + kKBytes, kSRequiredAlignment); + // Single K/V LDS buffer: S sits between K and V. + static constexpr index_t kSingleBufferSOffset = + integer_least_multiple(kK0Offset + kKBytes, kSRequiredAlignment); static constexpr index_t kV0Offset = [] { if constexpr(kUseDoubleKVLdsBuffer) { @@ -947,16 +928,27 @@ struct QrTdmLdsArenaLayout } else { - return integer_least_multiple(kSOffset + kSBytes, kRegionAlignment); + return integer_least_multiple(kSingleBufferSOffset + kSBytes, kRegionAlignment); } }(); static constexpr index_t kV1Offset = kUseDoubleKVLdsBuffer ? integer_least_multiple(kV0Offset + kVBytes, kRegionAlignment) : kV0Offset; + // S only lives in LDS when the gemm_0 warps sit on N (kSBytes > 0). On the + // double K/V buffer path every K/V slot holds a tile that is being read or is + // still in flight from the TDM, so S gets its own slot after V1: aliased into + // K0, the K prefetch two iterations ahead would overwrite it. + static constexpr index_t kSOffset = [] { + if constexpr(kUseDoubleKVLdsBuffer) + return kSBytes > 0 ? integer_least_multiple(kV1Offset + kVBytes, kRegionAlignment) : 0; + else + return kSingleBufferSOffset; + }(); static constexpr index_t kArenaBytes = [] { if constexpr(kUseDoubleKVLdsBuffer) { - return integer_least_multiple(kV1Offset + kVBytes, kArenaAlignment); + return integer_least_multiple(max(kV1Offset + kVBytes, kSOffset + kSBytes), + kArenaAlignment); } else { @@ -978,6 +970,8 @@ struct QrTdmLdsArenaLayout static_assert(!kUseDoubleKVLdsBuffer || kQOffset + kQBytes <= kV0Offset); static_assert(kUseDoubleKVLdsBuffer || kK0Offset + kKBytes <= kSOffset); static_assert(kUseDoubleKVLdsBuffer || kSOffset + kSBytes <= kV0Offset); + static_assert(!kUseDoubleKVLdsBuffer || kSBytes == 0 || kV1Offset + kVBytes <= kSOffset); + static_assert(kSOffset + kSBytes <= kArenaBytes); static_assert(!kUseDoubleKVLdsBuffer || !KPadding::kEnabled || (kK1Offset - kK0Offset) % KPadding::kIntervalBytes == 0); static_assert(!kUseDoubleKVLdsBuffer || !VPadding::kEnabled || diff --git a/projects/composablekernel/test/ck_tile/fmha/test_qr_tdm_lds_padding.cpp b/projects/composablekernel/test/ck_tile/fmha/test_qr_tdm_lds_padding.cpp index b8542a072bbe..b2d464c86aa0 100644 --- a/projects/composablekernel/test/ck_tile/fmha/test_qr_tdm_lds_padding.cpp +++ b/projects/composablekernel/test/ck_tile/fmha/test_qr_tdm_lds_padding.cpp @@ -67,6 +67,15 @@ using TestFmhaShape = ck_tile::TileFmhaShape, true>; +// Same geometry as TestFmhaShape, with explicit gemm_0 / gemm_1 warp layouts. +template +using TestFmhaShapeWarps = ck_tile::TileFmhaShape, + Gemm0Warps, + ck_tile::sequence<16, 16, 32>, + Gemm1Warps, + ck_tile::sequence<16, 16, 32>, + true>; + using TestFmhaTraits = ck_tile::TileFmhaTraits, 128, 64, 32, 32, 64>>; +// The 16-row short-query tile: four warps on N for both gemms. +using ShortQueryShape = TestFmhaShapeWarps<16, ck_tile::sequence<1, 4, 1>>; +using ShortQueryProblem = + TestProblemWithShape, ShortQueryShape>; +// gemm_0 and gemm_1 must agree on where the warps sit. +using MixedWarpsProblem = TestProblemWithShape< + SelectionBaseProblem, + TestFmhaShapeWarps<128, ck_tile::sequence<1, 4, 1>, ck_tile::sequence<4, 1, 1>>>; using VColumnMajorProblem = TestProblemWithShape, false>>; @@ -205,6 +222,12 @@ static_assert(std::is_same_v: ck_tile::detail::LdsPaddingConfig>); static_assert(std::is_same_v::V, ck_tile::detail::LdsPaddingConfig>); +// The short-query tile gets the d=128 K/V padding; Q stays unpadded (only the +// progressive M128 traversal pads Q). +static_assert(std::is_same_v::Q, NoPad>); +static_assert(std::is_same_v::K, QKPad>); +static_assert(std::is_same_v::V, VPad>); +static_assert(is_disabled_selection>()); static_assert(is_disabled_selection>()); static_assert(is_disabled_selection>()); static_assert(is_disabled_selection>()); @@ -446,6 +469,23 @@ constexpr bool validate_arena_layouts() static_assert(DecodeAll::kV0Offset == 4352); static_assert(DecodeAll::kArenaBytes == 22784); + // Warps on M keep S in registers (no S region), so the layout above is + // unchanged. Warps on N (the 16-row short-query tile) stage S through LDS: on + // the double K/V buffer path it gets its own slot after V1, since every K/V + // slot holds a tile that is being read or is still in flight from the TDM. + using ShortQueryProblem = + TestProblemWithShape, + TestFmhaShapeWarps<16, ck_tile::sequence<1, 4, 1>>>; + using ShortQueryAll = + typename Policy::template LdsArenaLayout; + static_assert(PrefillAll::kSBytes == 0); + static_assert(ShortQueryAll::kSBytes > 0); + static_assert(ShortQueryAll::kK1Offset == 17408); + static_assert(ShortQueryAll::kV0Offset == 34816); + static_assert(ShortQueryAll::kV1Offset == 53248); + static_assert(ShortQueryAll::kSOffset == 71680); + static_assert(ShortQueryAll::kArenaBytes == 76032); + using PrefillNone = typename Policy::template LdsArenaLayout; using PrefillQKV = typename Policy::template LdsArenaLayout; From 9fb13088f9d78bde9c5de6df8545b110238ca730 Mon Sep 17 00:00:00 2001 From: yanxishi Date: Sun, 11 Oct 2026 09:17:30 +0000 Subject: [PATCH 02/10] fix(fmha/fwd): count S in the qr_tdm policy's double-buffer GetSmemSize With the gemm_0 warps on N, S has its own slot in the double K/V buffer arena, but the policy's GetSmemSize still returned the K/V buffers only. The kernel sizes its LDS from the arena layout, so nothing was under- allocated; keep the policy's estimate consistent with it anyway. --- .../pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp b/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp index fed6ef7b6393..e9eb6dff1452 100644 --- a/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp +++ b/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp @@ -853,8 +853,10 @@ struct BlockFmhaPipelineQRKSVSTdmDefaultPolicy { if constexpr(Problem::kUseDoubleKVLdsBuffer) { - // Double K/V LDS buffers: two K buffers (ping/pong) and two V buffers (ping/pong). - return 2 * GetSmemSizeK() + 2 * GetSmemSizeV(); + // Double K/V LDS buffers: two K buffers (ping/pong) and two V buffers (ping/pong), + // plus S when the gemm_0 warps sit on N (GetSmemSizeS() is 0 otherwise). + return 2 * GetSmemSizeK() + 2 * GetSmemSizeV() + + GetSmemSizeS(); } else { From 402cd34af50e11c40ab51b2b43ee37fcdf7d3274 Mon Sep 17 00:00:00 2001 From: yanxishi Date: Sun, 11 Oct 2026 09:47:50 +0000 Subject: [PATCH 03/10] test(fmha/fwd): end-to-end check of the gfx1250 16-row short-query tile Run s_q=8 against s_k=509 with LSE and assert the b16x64 qr_tdm instance is selected. A mis-distributed S on its way through LDS permutes P within a row, which keeps the row sums (and LSE) right while O is wrong, so this case needs the O comparison rather than LSE alone. --- .../test/ck_tile/fmha/test_fmha_fwd.cpp | 52 +++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/projects/composablekernel/test/ck_tile/fmha/test_fmha_fwd.cpp b/projects/composablekernel/test/ck_tile/fmha/test_fmha_fwd.cpp index 926db57949cf..73f152e3a433 100644 --- a/projects/composablekernel/test/ck_tile/fmha/test_fmha_fwd.cpp +++ b/projects/composablekernel/test/ck_tile/fmha/test_fmha_fwd.cpp @@ -411,6 +411,58 @@ TEST(TestCkTileFmhaFwd, QrTdmLdsArenaPrefill) } } +// max_seqlen_q <= 16 takes a tile that spends its warps on N instead of M, which +// sends S through LDS on its way from gemm_0 to gemm_1. A distribution mismatch +// there permutes P's columns within a row, which leaves the row sums -- and so +// the LSE -- correct while O is wrong, hence lse=true below. +TEST(TestCkTileFmhaFwd, QrTdmShortQuery) +{ + if constexpr(ck_tile::is_any_of::value) + { + if(!ck_tile::is_gfx125_supported()) + GTEST_SKIP() << "qr_tdm is only supported on gfx1250"; + + std::string selected_kernel; + auto result = fmha_fwd_run(mode_enum::batch, + 1, + 4, + 2, + {8}, + {509}, + 128, + 128, + 0, + {-1}, + {-1}, + {}, + {}, + 0, + true, + true, + 0, + 0, + true, + false, + 0, + false, + "n", + 0.0f, + 0, + 0, + false, + "0", + qscale_str, + true, + 1, + COMMON_ARGS, + std::nullopt, + &selected_kernel); + ASSERT_EQ(result, fwd_result::success); + EXPECT_NE(selected_kernel.find("_qr_tdm_"), std::string::npos); + EXPECT_NE(selected_kernel.find("b16x64"), std::string::npos); + } +} + // gfx125x-only: qr_tdm correctness at the head dims this PR added but that the // generic General sweep does not cover -- (160,160) and (96,96). The generic // suite already exercises 32/64/128/192-128 through qr_tdm but never asserts From 95855fd1d49e4d7f3aaeee9705e35531cdb50232 Mon Sep 17 00:00:00 2001 From: yanxishi Date: Mon, 12 Oct 2026 09:08:28 +0000 Subject: [PATCH 04/10] test(fmha/fwd): rename a local alias that shadows the file-scope ShortQueryProblem -Wshadow (an error in the CK build) flags the function-local alias in the LDS arena checks, which reused the name of the file-scope short-query problem. --- .../test/ck_tile/fmha/test_qr_tdm_lds_padding.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/projects/composablekernel/test/ck_tile/fmha/test_qr_tdm_lds_padding.cpp b/projects/composablekernel/test/ck_tile/fmha/test_qr_tdm_lds_padding.cpp index b2d464c86aa0..ebd4c45669f0 100644 --- a/projects/composablekernel/test/ck_tile/fmha/test_qr_tdm_lds_padding.cpp +++ b/projects/composablekernel/test/ck_tile/fmha/test_qr_tdm_lds_padding.cpp @@ -473,11 +473,11 @@ constexpr bool validate_arena_layouts() // unchanged. Warps on N (the 16-row short-query tile) stage S through LDS: on // the double K/V buffer path it gets its own slot after V1, since every K/V // slot holds a tile that is being read or is still in flight from the TDM. - using ShortQueryProblem = + using ShortQueryDoubleBufferProblem = TestProblemWithShape, TestFmhaShapeWarps<16, ck_tile::sequence<1, 4, 1>>>; using ShortQueryAll = - typename Policy::template LdsArenaLayout; + typename Policy::template LdsArenaLayout; static_assert(PrefillAll::kSBytes == 0); static_assert(ShortQueryAll::kSBytes > 0); static_assert(ShortQueryAll::kK1Offset == 17408); From 94227a0ea4a69bc72558bbec6c4ac34b683d911c Mon Sep 17 00:00:00 2001 From: yanxishi Date: Tue, 13 Oct 2026 00:44:06 +0000 Subject: [PATCH 05/10] fix(fmha/fwd): return the LDS arena size from the qr_tdm policy's GetSmemSize The hand-summed region sizes dropped the 256-byte region alignment (and the double-buffer S slot is placed after V1), so they disagreed with the arena the pipeline actually allocates. The arena layout owns the offsets; use it. --- .../block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp | 16 +++------------- 1 file changed, 3 insertions(+), 13 deletions(-) diff --git a/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp b/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp index e9eb6dff1452..e4d1694b0678 100644 --- a/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp +++ b/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp @@ -851,19 +851,9 @@ struct BlockFmhaPipelineQRKSVSTdmDefaultPolicy template CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize() { - if constexpr(Problem::kUseDoubleKVLdsBuffer) - { - // Double K/V LDS buffers: two K buffers (ping/pong) and two V buffers (ping/pong), - // plus S when the gemm_0 warps sit on N (GetSmemSizeS() is 0 otherwise). - return 2 * GetSmemSizeK() + 2 * GetSmemSizeV() + - GetSmemSizeS(); - } - else - { - // Single K/V LDS buffers: Q, K, S, V laid out sequentially. - return max(GetSmemSizeQ(), - GetSmemSizeK() + GetSmemSizeS() + GetSmemSizeV()); - } + // The arena layout owns the region offsets, alignment and aliasing; summing the + // region sizes here would drop the 256-byte region alignment. + return LdsArenaLayout::kArenaBytes; } }; From 0d89b8b7316adb9add2fa07fab5ca23939b8de51 Mon Sep 17 00:00:00 2001 From: yanxishi Date: Tue, 13 Oct 2026 05:35:24 +0000 Subject: [PATCH 06/10] fix(fmha/fwd): keep the qr_tdm_sched compatibility rule last among the Gfx125 rules Its codegen tests address it as get_rules()[-1]. --- .../example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py b/projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py index 8de617e13d20..79c5bd01de8e 100644 --- a/projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py +++ b/projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py @@ -1690,8 +1690,9 @@ def check_short_query_tile( return True return kernel_ctx.pipeline.tag == "qr_tdm" + # The sched rule stays last (its tests address it as get_rules()[-1]). rules.extend( - [check_gemm0_k, check_qr_tdm_sched_codegen_compatibility, check_short_query_tile] + [check_gemm0_k, check_short_query_tile, check_qr_tdm_sched_codegen_compatibility] ) return rules From 71530c7acac976482cf90aa26eeaba45319b054a Mon Sep 17 00:00:00 2001 From: yanxishi Date: Tue, 13 Oct 2026 05:47:16 +0000 Subject: [PATCH 07/10] test(fmha/fwd): d128 dispatch expectations for the 16-row short-query tile max_seqlen_q <= 16 takes the 16-row tile, not the bm0=64 one. --- .../test/ck_tile/fmha/test_fmha_fwd_codegen.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/projects/composablekernel/test/ck_tile/fmha/test_fmha_fwd_codegen.py b/projects/composablekernel/test/ck_tile/fmha/test_fmha_fwd_codegen.py index 9e6e0bc997ef..e36f6df89846 100644 --- a/projects/composablekernel/test/ck_tile/fmha/test_fmha_fwd_codegen.py +++ b/projects/composablekernel/test/ck_tile/fmha/test_fmha_fwd_codegen.py @@ -943,7 +943,16 @@ def test_real_dispatch_d128_priority_and_sequence_boundary(self): ] ) for mode, group in (("batch", 0), ("group", 1)): - for length, bm0 in ((1, 64), (127, 64), (128, 64), (2047, 64), (2048, 128)): + # max_seqlen_q <= 16 takes the 16-row short-query tile. + for length, bm0 in ( + (1, 16), + (16, 16), + (17, 64), + (127, 64), + (128, 64), + (2047, 64), + (2048, 128), + ): expected = dispatch.expected_id( mode=mode, hdim=128, @@ -1773,11 +1782,11 @@ def test_d128_filtered_priority_and_fallback(self): candidate = dispatch.expected_id( pipeline="qr_tdm_sched", **features ) - for size in (1, 127, 128, 129, 2047, 2048, 32768): + for size in (1, 16, 17, 127, 128, 129, 2047, 2048, 32768): fallback = dispatch.expected_id( pipeline="qr_tdm", spad="t" if group else "f", - bm0=64 if size < 128 else 128, + bm0=16 if size <= 16 else 64 if size < 128 else 128, dpad="f", **features, ) From 04ab967ef6e943f6f704891fcdae7992f025bb9d Mon Sep 17 00:00:00 2001 From: yanxishi Date: Tue, 13 Oct 2026 09:06:52 +0000 Subject: [PATCH 08/10] test(fmha/fwd): cover the gfx1250 16-row short-query tile across lengths and masks QrTdmShortQuery ran one batch-mode s_q=8 dense case with LSE off, while its comment asked for LSE on. It now runs six representative cases with LSE on: s_q = 1, 8 and 16, seqlen_k on and off a tile multiple, no mask, both causal masks and a sliding window, batch and group mode, and checks the tile in batch mode. MHA keeps pack-GQA from folding the heads into seqlen_q. Also: ruff-format the rule list, put the progressive-K comment back on kStagedKPairs, and word the S distribution comment for the code as it is. --- .../ck_tile/01_fmha/codegen/ops/fmha_fwd.py | 6 ++- .../block_fmha_pipeline_qr_ks_vs_tdm.hpp | 6 +-- ...lock_fmha_pipeline_qr_ks_vs_tdm_policy.hpp | 13 +++--- .../test/ck_tile/fmha/test_fmha_fwd.cpp | 45 ++++++++++++++----- 4 files changed, 46 insertions(+), 24 deletions(-) diff --git a/projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py b/projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py index 79c5bd01de8e..3386e7550780 100644 --- a/projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py +++ b/projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py @@ -1692,7 +1692,11 @@ def check_short_query_tile( # The sched rule stays last (its tests address it as get_rules()[-1]). rules.extend( - [check_gemm0_k, check_short_query_tile, check_qr_tdm_sched_codegen_compatibility] + [ + check_gemm0_k, + check_short_query_tile, + check_qr_tdm_sched_codegen_compatibility, + ] ) return rules diff --git a/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm.hpp b/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm.hpp index 4890fe70f155..45c76c57c147 100644 --- a/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm.hpp +++ b/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm.hpp @@ -359,14 +359,14 @@ struct BlockFmhaPipelineQRKSVSTdm static constexpr const char* name = "qr_tdm"; - // The progressive M128 traversal has the fragment shape required by the - // staged hand-off below. Element types and attention features do not alter - // those LDS fragment lifetimes. // The progressive K reload walks four N-subtiles per warp, so it needs the // gemm_0 warps all on M. Tiles that spend their warps on N (the short-query // tile) keep the plain whole-slice K reload. static constexpr bool kProgressiveDsLoadK = Problem::kProgressiveDsLoadK && kNWarp == 1; + // The progressive M128 traversal has the fragment shape required by the + // staged hand-off below. Element types and attention features do not alter + // those LDS fragment lifetimes. static constexpr bool kStagedKPairs = kProgressiveDsLoadK && kM0 == 128 && kQKHeaddim == 128 && kK0 == 32; static constexpr bool kUseIglpBulkScheduling = CK_TILE_FMHA_TDM_IGLP_BULK && !kStagedKPairs; diff --git a/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp b/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp index e4d1694b0678..30521dd90cd6 100644 --- a/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp +++ b/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp @@ -808,12 +808,10 @@ struct BlockFmhaPipelineQRKSVSTdmDefaultPolicy CK_TILE_HOST_DEVICE static constexpr auto MakeSRegTileDistribution() { // S comes back out of LDS to be softmaxed and then fed into gemm_1 as its A - // operand, so it has to land in exactly the layout gemm_1 expects. Building - // the encoding from the QK block gemm's warp config -- a different gemm than - // the PV one that consumes S -- permutes P's columns within a row: the row - // sum (and hence LSE) still comes out right, but O picks up a ~1% error - // because each key is paired with the wrong value row. Reuse the P - // distribution instead, which embeds the PV warp gemm's own A encoding. + // operand, so it takes P's distribution, which embeds the PV warp gemm's A + // encoding. An encoding built from the QK gemm's warp config would permute + // P's columns within a row once the warps sit on N: the row sums (and LSE) + // stay right while each key meets the wrong value row. return MakePRegTileDistribution(); } @@ -851,8 +849,7 @@ struct BlockFmhaPipelineQRKSVSTdmDefaultPolicy template CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize() { - // The arena layout owns the region offsets, alignment and aliasing; summing the - // region sizes here would drop the 256-byte region alignment. + // The arena layout owns the region offsets, alignment and aliasing. return LdsArenaLayout::kArenaBytes; } }; diff --git a/projects/composablekernel/test/ck_tile/fmha/test_fmha_fwd.cpp b/projects/composablekernel/test/ck_tile/fmha/test_fmha_fwd.cpp index 73f152e3a433..020386622434 100644 --- a/projects/composablekernel/test/ck_tile/fmha/test_fmha_fwd.cpp +++ b/projects/composablekernel/test/ck_tile/fmha/test_fmha_fwd.cpp @@ -413,22 +413,41 @@ TEST(TestCkTileFmhaFwd, QrTdmLdsArenaPrefill) // max_seqlen_q <= 16 takes a tile that spends its warps on N instead of M, which // sends S through LDS on its way from gemm_0 to gemm_1. A distribution mismatch -// there permutes P's columns within a row, which leaves the row sums -- and so -// the LSE -- correct while O is wrong, hence lse=true below. -TEST(TestCkTileFmhaFwd, QrTdmShortQuery) +// there permutes P's columns within a row: the row sums -- and so the LSE -- stay +// right while O is wrong, so O is what these cases have to check. seqlen_k covers +// a tile multiple and a padded tail; the masks cover the causal and windowed +// paths. +class QrTdmShortQuery : public TestWithParam> +{ +}; + +INSTANTIATE_TEST_SUITE_P(TestCkTileFmhaFwd, + QrTdmShortQuery, + // mode, seqlen_q, seqlen_k, mask + Values(std::tuple{mode_enum::batch, 1, 509, "0"}, + std::tuple{mode_enum::batch, 8, 64, "0"}, + std::tuple{mode_enum::batch, 16, 509, "1"}, + std::tuple{mode_enum::batch, 8, 509, "2"}, + std::tuple{mode_enum::batch, 16, 64, "t:50,64"}, + std::tuple{mode_enum::group, 8, 509, "2"})); + +TEST_P(QrTdmShortQuery, DataTypeConfig) { if constexpr(ck_tile::is_any_of::value) { if(!ck_tile::is_gfx125_supported()) GTEST_SKIP() << "qr_tdm is only supported on gfx1250"; + auto [mode, seqlen_q, seqlen_k, mask_str] = GetParam(); + std::string selected_kernel; - auto result = fmha_fwd_run(mode_enum::batch, - 1, - 4, + // MHA: pack-GQA would fold the heads into seqlen_q and move the tile choice. + auto result = fmha_fwd_run(mode, 2, - {8}, - {509}, + 4, + -1, + {seqlen_q}, + {seqlen_k}, 128, 128, 0, @@ -442,7 +461,7 @@ TEST(TestCkTileFmhaFwd, QrTdmShortQuery) 0, 0, true, - false, + true, 0, false, "n", @@ -450,7 +469,7 @@ TEST(TestCkTileFmhaFwd, QrTdmShortQuery) 0, 0, false, - "0", + mask_str, qscale_str, true, 1, @@ -458,8 +477,10 @@ TEST(TestCkTileFmhaFwd, QrTdmShortQuery) std::nullopt, &selected_kernel); ASSERT_EQ(result, fwd_result::success); - EXPECT_NE(selected_kernel.find("_qr_tdm_"), std::string::npos); - EXPECT_NE(selected_kernel.find("b16x64"), std::string::npos); + EXPECT_NE(selected_kernel.find("_qr_tdm_"), std::string::npos) << selected_kernel; + // Group mode dispatches on the packed total q length (see QrTdmNewHeadDim). + if(mode == mode_enum::batch) + EXPECT_NE(selected_kernel.find("_b16x64"), std::string::npos) << selected_kernel; } } From b6c8dad925e2d8df36821ba116f06a08b9116a91 Mon Sep 17 00:00:00 2001 From: yanxishi Date: Tue, 13 Oct 2026 10:24:18 +0000 Subject: [PATCH 09/10] refactor(fmha/fwd): name the pipeline's progressive-K switch apart from the problem's The pipeline's kProgressiveDsLoadK (the problem's flag, off for tiles with the warps on N) shadowed Problem::kProgressiveDsLoadK, which the static_assert above and the policy still read; call it kUseProgressiveDsLoadK. No functional change. --- .../fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm.hpp | 8 ++++---- .../composablekernel/test/ck_tile/fmha/test_fmha_fwd.cpp | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm.hpp b/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm.hpp index 45c76c57c147..6c94d3ccbd83 100644 --- a/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm.hpp +++ b/projects/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm.hpp @@ -362,13 +362,13 @@ struct BlockFmhaPipelineQRKSVSTdm // The progressive K reload walks four N-subtiles per warp, so it needs the // gemm_0 warps all on M. Tiles that spend their warps on N (the short-query // tile) keep the plain whole-slice K reload. - static constexpr bool kProgressiveDsLoadK = Problem::kProgressiveDsLoadK && kNWarp == 1; + static constexpr bool kUseProgressiveDsLoadK = Problem::kProgressiveDsLoadK && kNWarp == 1; // The progressive M128 traversal has the fragment shape required by the // staged hand-off below. Element types and attention features do not alter // those LDS fragment lifetimes. static constexpr bool kStagedKPairs = - kProgressiveDsLoadK && kM0 == 128 && kQKHeaddim == 128 && kK0 == 32; + kUseProgressiveDsLoadK && kM0 == 128 && kQKHeaddim == 128 && kK0 == 32; static constexpr bool kUseIglpBulkScheduling = CK_TILE_FMHA_TDM_IGLP_BULK && !kStagedKPairs; // Unchecked TDM bounds need a stronger, independent range proof. Keep this @@ -1677,7 +1677,7 @@ struct BlockFmhaPipelineQRKSVSTdm // loop over along the [K]ey head dimension move_tile_window(k_lds_read_window, {0, kK0}); using Gemm0 = remove_cvref_t; - if constexpr(kProgressiveDsLoadK) + if constexpr(kUseProgressiveDsLoadK) { static_assert( std::is_same_v); @@ -1734,7 +1734,7 @@ struct BlockFmhaPipelineQRKSVSTdm move_tile_window(k_lds_read_window, {0, -kK0 * (k0_loops - 1)}); } - if constexpr(!kStagedKPairs && kProgressiveDsLoadK) + if constexpr(!kStagedKPairs && kUseProgressiveDsLoadK) { // Retain the reload-overlap traversal when consuming the final // head-dimension slice, even though this slice has no reload. diff --git a/projects/composablekernel/test/ck_tile/fmha/test_fmha_fwd.cpp b/projects/composablekernel/test/ck_tile/fmha/test_fmha_fwd.cpp index 020386622434..5f112e34468e 100644 --- a/projects/composablekernel/test/ck_tile/fmha/test_fmha_fwd.cpp +++ b/projects/composablekernel/test/ck_tile/fmha/test_fmha_fwd.cpp @@ -413,7 +413,7 @@ TEST(TestCkTileFmhaFwd, QrTdmLdsArenaPrefill) // max_seqlen_q <= 16 takes a tile that spends its warps on N instead of M, which // sends S through LDS on its way from gemm_0 to gemm_1. A distribution mismatch -// there permutes P's columns within a row: the row sums -- and so the LSE -- stay +// there permutes P's columns within a row: the row sums, and so the LSE, stay // right while O is wrong, so O is what these cases have to check. seqlen_k covers // a tile multiple and a padded tail; the masks cover the causal and windowed // paths. From c98808a2f8ea563b85cb9e76cb8c155662366678 Mon Sep 17 00:00:00 2001 From: yanxishi Date: Tue, 13 Oct 2026 10:47:08 +0000 Subject: [PATCH 10/10] fix(fmha/test): build test_ck_tile_fmha_reference with the test_ck_tile_fmha target When a PR touches example/ck_tile/01_fmha, smart_build_ci.sh builds the test_ck_tile_fmha umbrella and runs every ctest matching test_ck_tile_fmha. That includes test_ck_tile_fmha_reference, which the umbrella did not build, so CI reported it as "Not Run" on gfx1201 and the other non-gfx125 targets. --- projects/composablekernel/test/ck_tile/fmha/CMakeLists.txt | 2 ++ 1 file changed, 2 insertions(+) diff --git a/projects/composablekernel/test/ck_tile/fmha/CMakeLists.txt b/projects/composablekernel/test/ck_tile/fmha/CMakeLists.txt index c3a44b1cabb4..12ee8d1b908c 100644 --- a/projects/composablekernel/test/ck_tile/fmha/CMakeLists.txt +++ b/projects/composablekernel/test/ck_tile/fmha/CMakeLists.txt @@ -248,6 +248,8 @@ get_test_property(test_ck_tile_fmha_reference LABELS FMHA_REFERENCE_COMMON_LABEL set_tests_properties(test_ck_tile_fmha_reference PROPERTIES LABELS "${FMHA_REFERENCE_COMMON_LABELS};${TEST_NAME};CK_TILE_FMHA_TESTS") list(APPEND FMHA_QR_TDM_SCHED_TEST_TARGETS test_ck_tile_fmha_reference) +# CI runs every test_ck_tile_fmha* ctest after building the ${TEST_NAME} target. +add_dependencies(${TEST_NAME} test_ck_tile_fmha_reference) # Python integration test for the fmha_fwd codegen path. #