Skip to content

perf(ck_tile/fmha): [gfx1250 opt 3] 16-row short-query tile for the gfx1250 qr_tdm forward #13469

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Yanxing-Shi wants to merge 10 commits into from
Closed

perf(ck_tile/fmha): [gfx1250 opt 3] 16-row short-query tile for the gfx1250 qr_tdm forward #13469

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
a60f83a
perf(fmha/fwd): 16-row short-query tile for the gfx1250 qr_tdm forward
Yanxing-Shi Oct 11, 2026
9fb1308
fix(fmha/fwd): count S in the qr_tdm policy's double-buffer GetSmemSize
Yanxing-Shi Oct 11, 2026
402cd34
test(fmha/fwd): end-to-end check of the gfx1250 16-row short-query tile
Yanxing-Shi Oct 11, 2026
95855fd
test(fmha/fwd): rename a local alias that shadows the file-scope Shor…
Yanxing-Shi Oct 12, 2026
94227a0
fix(fmha/fwd): return the LDS arena size from the qr_tdm policy's Get…
Yanxing-Shi Oct 13, 2026
0d89b8b
fix(fmha/fwd): keep the qr_tdm_sched compatibility rule last among th…
Yanxing-Shi Oct 13, 2026
71530c7
test(fmha/fwd): d128 dispatch expectations for the 16-row short-query…
Yanxing-Shi Oct 13, 2026
04ab967
test(fmha/fwd): cover the gfx1250 16-row short-query tile across leng…
Yanxing-Shi Oct 13, 2026
b6c8dad
refactor(fmha/fwd): name the pipeline's progressive-K switch apart fr…
Yanxing-Shi Oct 13, 2026
c98808a
fix(fmha/test): build test_ck_tile_fmha_reference with the test_ck_ti…
Yanxing-Shi Oct 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 23 additions & 2 deletions projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1680,7 +1680,24 @@ 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.

@poyenc poyenc Oct 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This check keys on the tile's F_bm0 dimension rather than the geometry invariant it's actually protecting. The comment above explains the real constraint: qr's gemm_1 slices P along M and can't form a slice narrower than one warp tile, which only matters when gemm_1's warps are split along N instead of M. That's exactly what F_rn1 encodes (the tile dataclass defines it as "number of warps for gemm1 along head dim v"). Today's single F_bm0==16 tile happens to be the only N-split tile, so the check works, but it only works by coincidence: a future N-split tile at a different F_bm0 would skip this guard and still fail to compile with qr, while an unrelated M-split tile that happens to use F_bm0==16 would be needlessly barred from qr.

Check the actual invariant instead:

if kernel_ctx.tile.F_rn1 <= 1:
    return True
return kernel_ctx.pipeline.tag == "qr_tdm"

if kernel_ctx.tile.F_bm0 != 16:
return True
return kernel_ctx.pipeline.tag == "qr_tdm"

# The sched rule stays last (its tests address it as get_rules()[-1]).

@poyenc poyenc Oct 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This comment documents that check_qr_tdm_sched_codegen_compatibility must stay last in the list because two tests reach it positionally through get_rules()[-1]. Since get_rules()'s output is only ever consumed as a conjunction over all rules, list order has no effect on which kernels get generated — the ordering constraint exists purely to satisfy those two tests' indexing, and this comment is the production code accommodating that rather than the tests being fixed. The next engineer adding a compatibility rule has to remember to append before this one, with only a comment enforcing it; a reorder or an unrelated append placed after it will silently break both tests again, the same way this PR's own first commit did when it introduced check_short_query_tile.

Have those two tests select the rule by name or function identity instead of position, e.g. next(r for r in get_rules() if r.__name__ == "check_qr_tdm_sched_codegen_compatibility"), then this comment and the ordering constraint on rules.extend(...) here can both be removed.

rules.extend(
[
check_gemm0_k,
check_short_query_tile,
check_qr_tdm_sched_codegen_compatibility,
]
)
return rules

@classmethod
Expand Down Expand Up @@ -1716,7 +1733,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)],
Expand Down
14 changes: 9 additions & 5 deletions ...s/composablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -359,11 +359,16 @@ struct BlockFmhaPipelineQRKSVSTdm

static constexpr const char* name = "qr_tdm";

// 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 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 =
Problem::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
Expand Down Expand Up @@ -1462,8 +1467,7 @@ struct BlockFmhaPipelineQRKSVSTdm

// S tile in LDS
auto s_lds = make_tensor_view<address_space_enum::lds>(
reinterpret_cast<SaccDataType*>(reinterpret_cast<char*>(smem_ptrk0) +
Policy::template GetSmemSizeK<Problem>()),
reinterpret_cast<SaccDataType*>(static_cast<char*>(smem_arena) + Layout::kSOffset),
Policy::template MakeSLdsBlockDescriptor<Problem>());
auto s_write_lds_window = make_tile_window(
s_lds, Policy::template MakeSLdsBlockDescriptor<Problem>().get_lengths(), {0, 0});
Expand Down Expand Up @@ -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<decltype(gemm_0)>;
if constexpr(Problem::kProgressiveDsLoadK)
if constexpr(kUseProgressiveDsLoadK)
{
static_assert(
std::is_same_v<Policy, BlockFmhaPipelineQRKSVSTdmDefaultPolicy>);
Expand Down Expand Up @@ -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 && kUseProgressiveDsLoadK)
{
// Retain the reload-overlap traversal when consuming the final
// head-dimension slice, even though this slice has no reload.
Expand Down
83 changes: 33 additions & 50 deletions ...sablekernel/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_tdm_policy.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -211,7 +211,8 @@ inline constexpr bool is_qr_tdm_padding_enabled_problem_v =
std::is_same_v<typename Problem::KDataType, half_t> &&
std::is_same_v<typename Problem::VDataType, half_t>)) &&
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
Expand All @@ -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<typename Problem::BlockFmhaShape::Gemm0BlockWarps, sequence<4, 1, 1>> &&
std::is_same_v<typename Problem::BlockFmhaShape::Gemm1BlockWarps, sequence<4, 1, 1>> &&
// warps all on M, or all on N for the short-query tile
(std::is_same_v<typename Problem::BlockFmhaShape::Gemm0BlockWarps, sequence<4, 1, 1>> ||
std::is_same_v<typename Problem::BlockFmhaShape::Gemm0BlockWarps, sequence<1, 4, 1>>) &&
std::is_same_v<typename Problem::BlockFmhaShape::Gemm1BlockWarps,
typename Problem::BlockFmhaShape::Gemm0BlockWarps> &&
std::is_same_v<typename Problem::BlockFmhaShape::Gemm0WarpTile, sequence<16, 16, 32>> &&
std::is_same_v<typename Problem::BlockFmhaShape::Gemm1WarpTile, sequence<16, 16, 32>> &&
numeric_traits<typename Problem::QDataType>::PackedSize == 1 &&
Expand Down Expand Up @@ -803,37 +807,12 @@ struct BlockFmhaPipelineQRKSVSTdmDefaultPolicy
template <typename Problem>
CK_TILE_HOST_DEVICE static constexpr auto MakeSRegTileDistribution()
{
using BlockGemm = remove_cvref_t<decltype(GetKVBlockGemm<Problem>())>;

constexpr auto config = BlockGemm::Policy::template GetWarpGemmMWarpNWarp<Problem>();
using WG = remove_cvref_t<decltype(config.template at<0>())>;
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<sequence<NWarp>,
tuple<sequence<M0, M1, M2>, sequence<K0, K1, K2, K3>>,
tuple<sequence<1, 0>, sequence<2, 1>>,
tuple<sequence<1, 0>, 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 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<Problem>();
}

template <typename Problem>
Expand Down Expand Up @@ -870,17 +849,8 @@ struct BlockFmhaPipelineQRKSVSTdmDefaultPolicy
template <typename Problem>
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).
return 2 * GetSmemSizeK<Problem, true>() + 2 * GetSmemSizeV<Problem>();
}
else
{
// Single K/V LDS buffers: Q, K, S, V laid out sequentially.
return max(GetSmemSizeQ<Problem>(),
GetSmemSizeK<Problem>() + GetSmemSizeS<Problem>() + GetSmemSizeV<Problem>());
}
// The arena layout owns the region offsets, alignment and aliasing.
return LdsArenaLayout<Problem>::kArenaBytes;
}
};

Expand Down Expand Up @@ -936,9 +906,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)
{
Expand All @@ -947,16 +917,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
{
Expand All @@ -978,6 +959,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 ||
Expand Down
2 changes: 2 additions & 0 deletions projects/composablekernel/test/ck_tile/fmha/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
#
Expand Down
73 changes: 73 additions & 0 deletions projects/composablekernel/test/ck_tile/fmha/test_fmha_fwd.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -411,6 +411,79 @@ 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
// 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<std::tuple<mode_enum, int, int, const char*>>
{
};

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<DataTypeConfig, FmhaFwdFp16, FmhaFwdBf16>::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;
// MHA: pack-GQA would fold the heads into seqlen_q and move the tile choice.
auto result = fmha_fwd_run<DataTypeConfig>(mode,
2,
4,
-1,
{seqlen_q},
{seqlen_k},
128,
128,
0,
{-1},
{-1},
{},
{},
0,
true,
true,
0,
0,
true,
true,
0,
false,
"n",
0.0f,
0,
0,
false,
mask_str,
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) << 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;
}
}

// 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
Expand Down
15 changes: 12 additions & 3 deletions projects/composablekernel/test/ck_tile/fmha/test_fmha_fwd_codegen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
)
Expand Down
Loading
Loading