Skip to content

[P3-1] Tile-config table design (consteval get_tile_config) #7538

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

Merged
aosewski merged 8 commits into from
Jun 1, 2026
Merged

[P3-1] Tile-config table design (consteval get_tile_config) #7538

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
14ecf26
Added tile config with consteval lookup
chris-tsiaousis-hpc May 18, 2026
0cfd402
Added tests for consteval tile config lookup
chris-tsiaousis-hpc May 18, 2026
7bc25ce
Address copilot review comments
chris-tsiaousis-hpc May 19, 2026
c6bd381
Merge branch 'users/shumway/ck/exp-kpack' into users/chris-tsiaousis-…
chris-tsiaousis-hpc May 28, 2026
72cb60d
dqdkdv_spec: clearer consteval errors + asymmetric reject + table inv…
aosewski May 29, 2026
47c39a9
[P3-2] Remove dqdkdv_dev.hpp:105-107 d128 static_assert; dispatch via…
chris-tsiaousis-hpc Jun 1, 2026
2a93ee4
[P3-1] dqdkdv_spec: drop unsupported asymmetric head dims
aosewski Jun 1, 2026
8c1e0cd
Merge 'users/shumway/ck/exp-kpack'; unify consteval + dedup compat tests
aosewski Jun 1, 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
4 changes: 2 additions & 2 deletions ...ts/composablekernel/experimental/rocm_ck/include/rocm_ck/ops/fmha_bwd/convert_dq_spec.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ struct FmhaBwdConvertDQSpec
bool pad_seqlen_q;
bool pad_hdim_q;
int block_per_cu;
int block_size; // computed: 256 for d128
int block_size; // computed: 256 (4 warps x 64 threads)
};

// ---------------------------------------------------------------------------
Expand Down Expand Up @@ -137,7 +137,7 @@ consteval FmhaBwdConvertDQSpec makeSpec(FmhaBwdConvertDQConfig cfg)
throw "block_per_cu must be positive";

// Block size: 256 threads (4 warps x 64 threads/warp).
// Matches the CK Tile ConvertQGrad kernel configuration for d128.
// Matches the CK Tile ConvertQGrad kernel configuration.
constexpr int block_size = 256;

FmhaBwdConvertDQSpec k{sig.dtype,
Expand Down
66 changes: 32 additions & 34 deletions projects/composablekernel/experimental/rocm_ck/include/rocm_ck/ops/fmha_bwd/dqdkdv_dev.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,46 +100,44 @@ struct FmhaBwdDQDKDVTypes
K.has_bias_grad, // kHasBiasGrad
K.block_per_cu>; // kBlockPerCu

// Guard: tile geometry is only valid for d128. Other hdim values need
// different tile configs (see fmha_bwd.py get_dq_dk_dv_tiles()).
static_assert(K.hdim_q == 128 && K.hdim_v == 128,
"Tile geometry is hardcoded for d128. Other hdim values "
"require different tile configs.");

// Guard: BlockSize=256 with NumWarps=4 assumes wave64 (gfx9). On wave32
// targets (gfx10/11/12) the same NumWarps would yield BlockSize=128,
// breaking the BlockTile arithmetic and warp-level intrinsics encoded
// into the pipeline. Refuse to compile this bridge on non-gfx9 targets.
// Guard: wave64 (gfx9) only. The tile configs currently populated in
// getTileConfig() assume wavefrontSize=64. On wave32 targets (gfx10/11/12)
// the block_size arithmetic and warp-level intrinsics would be wrong.
static_assert(__AMDGCN_WAVEFRONT_SIZE == 64,
"FmhaBwdDQDKDV bridge requires wave64 (gfx9). The hardcoded "
"BlockSize=256 / NumWarps=4 tile config assumes warp_size=64.");
"FmhaBwdDQDKDV bridge requires wave64 (gfx9). "
"Add GFX11/GFX12 tile configs to enable wave32 targets.");

// --- Tile shape (hardcoded for d128 gfx9 -- Config 4 from fmha_bwd.py) ---
//
// From the gfx9 tile table for fp16/bf16 d128:
// bm0=16, bn0=128, bk0=128, bk1=16, bk2=128, bk3=16, bk4=32
// bhdq=128, bhdv=128
//
// Gemm0/Gemm2 block_warps: <1, 4, 1> warp_tile: <16, 16, 32>
// Gemm1/Gemm3 block_warps: <4, 1, 1> warp_tile: <16, 16, 16>
// Gemm4 block_warps: <1, 4, 1> warp_tile: <16, 16, min(32, 32)>
//
// NumWarps = 4, BlockSize = 256, maxSeqLenQ = 0 (unlimited)
// --- Tile shape (from consteval lookup table) ---
//
// getTileConfig() returns the architecture-specific tile geometry for
// the given (hdim_q, hdim_v, dtype). The device bridge reads plain
// integer fields and converts them to CK Tile sequence<> types.
static constexpr auto kTile = getTileConfig(K.hdim_q, K.hdim_v, K.dtype, GpuTarget::gfx942);

// BlockTile: sequence<bm0, bn0, bk0, bk1, bk2, bk3, bk4, bhdq, bhdv>
using BlockTile = ck_tile::sequence<16, 128, 128, 16, 128, 16, 32, 128, 128>;
using BlockTile = ck_tile::sequence<kTile.bm0,
kTile.bn0,
kTile.bk0,
kTile.bk1,
kTile.bk2,
kTile.bk3,
kTile.bk4,
K.hdim_q,
K.hdim_v>;

// Gemm0 & Gemm2: compute S = Q @ K^T and dP = dO @ V^T
using Gemm0BlockWarps = ck_tile::sequence<1, 4, 1>;
using Gemm0WarpTile = ck_tile::sequence<16, 16, 32>;
using Gemm0BlockWarps = ck_tile::sequence<kTile.rm0, kTile.rn0, kTile.rk0>;
using Gemm0WarpTile = ck_tile::sequence<kTile.wm0, kTile.wn0, kTile.wk0>;

// Gemm1 & Gemm3: compute dV = P^T @ dO and dK = dS^T @ Q
using Gemm1BlockWarps = ck_tile::sequence<4, 1, 1>;
using Gemm1WarpTile = ck_tile::sequence<16, 16, 16>;
using Gemm1BlockWarps = ck_tile::sequence<kTile.rm1, kTile.rn1, kTile.rk1>;
using Gemm1WarpTile = ck_tile::sequence<kTile.wm1, kTile.wn1, kTile.wk1>;

// Gemm4: compute dQ = dS @ K
using Gemm4BlockWarps = ck_tile::sequence<1, 4, 1>;
using Gemm4WarpTile = ck_tile::sequence<16, 16, 32>;
using Gemm4BlockWarps = ck_tile::sequence<kTile.rm2, kTile.rn2, kTile.rk2>;
// GEMM4 warp tile = (wm0, wn0, min(wk0, bk4)) per fmha_bwd.py
static constexpr int kWk4 = (kTile.wk0 < kTile.bk4) ? kTile.wk0 : kTile.bk4;
using Gemm4WarpTile = ck_tile::sequence<kTile.wm0, kTile.wn0, kWk4>;

// TileFmhaBwdShape: 5 GEMMs with their block_warps and warp_tiles
// G0=G2 (S/dP), G1=G3 (dV/dK), G4 (dQ)
Expand All @@ -154,7 +152,7 @@ struct FmhaBwdDQDKDVTypes
Gemm1WarpTile, // Gemm3 (same as G1)
Gemm4BlockWarps,
Gemm4WarpTile, // Gemm4
0>; // kMaxSeqLenQ (0 = unlimited)
kTile.max_seq_q>;

// --- Mask type ---
// has_mask=true -> GenericAttentionMask<true, true> (full masking)
Expand All @@ -165,11 +163,11 @@ struct FmhaBwdDQDKDVTypes

// --- Dropout type ---
// BlockDropoutBwd<IsDropout, IsWG32, IsStoreRandval>
// For d128 the warp tile wm0=16, so IsWG32=false -> wg16 variant.
// IsWG32 = (wm0 >= 32). All current gfx9 configs use wm0=16 → false.
// We never store randval in backward (IsStoreRandval=false).
using Dropout = ck_tile::BlockDropoutBwd<K.has_dropout,
false, // IsWG32 = false (wm0=16 for d128 config)
false>; // IsStoreRandval = false
(kTile.wm0 >= 32), // IsWG32
false>; // IsStoreRandval = false

// --- Pipeline problem ---
static constexpr bool kIsGroupMode = (K.mode == FmhaMode::GROUP);
Expand Down
Loading
Loading