Skip to content

feat: Compatible K when not a multiple of 8 #11537

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
illsilin merged 6 commits into from
Sep 21, 2026
Merged

feat: Compatible K when not a multiple of 8 #11537

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
3c6c293
Add A,B,C vector sizes as parameters to GemmConfig
Sep 4, 2026
fee9935
Add tests and address review comments
Sep 8, 2026
9d94c74
Broaden arch scope beyond gfx1250 (address review)
Sep 9, 2026
3e30eba
Merge branch 'develop' into users/fjimenez/ck/meta_gfx1250_K_dimension
JP-Fernando Sep 15, 2026
63b4de9
Fix test compilation in CI
Sep 15, 2026
5970868
Merge branch 'develop' into users/fjimenez/ck/meta_gfx1250_K_dimension
JP-Fernando Sep 21, 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
5 changes: 4 additions & 1 deletion projects/composablekernel/example/ck_tile/03_gemm/gemm_splitk_two_stage_invoker.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,10 @@ struct SplitKTwoStageInvoker
ck_tile::element_wise::PassThrough,
ck_tile::element_wise::PassThrough,
AComputeDataType,
BComputeDataType>;
BComputeDataType,
GemmConfig::FixedVectorSize,
GemmConfig::VectorSizeA,
GemmConfig::VectorSizeB>;
Comment thread
chris-tsiaousis-hpc marked this conversation as resolved.
Resolved
using WorkspaceType = ck_tile::remove_cvref_t<typename GemmConfig::WorkspaceType>;

using GemmPipeline = typename PipelineTypeTraits<
Expand Down
32 changes: 32 additions & 0 deletions projects/composablekernel/example/ck_tile/03_gemm/gemm_utils.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,36 @@ struct GemmConfigBase
static constexpr ck_tile::DataCachePrefetchKind DataCachePrefetchB =
ck_tile::DataCachePrefetchKind::None;
static constexpr bool Async = false;

static constexpr bool FixedVectorSize = false;
// If FixedVectorSize==true: use these vector sizes for A/B loads and C store
static constexpr ck_tile::index_t VectorSizeA = 1;
static constexpr ck_tile::index_t VectorSizeB = 1;
static constexpr ck_tile::index_t VectorSizeC = 1;

// Enable for gfx1250 RCR
Comment thread
chris-tsiaousis-hpc marked this conversation as resolved.
Resolved
static constexpr bool EnableKPadFallback = false;
};

// A,B vector sizes must be multiples of K
Comment thread
chris-tsiaousis-hpc marked this conversation as resolved.
Resolved
template <typename GemmConfig,
ck_tile::index_t VectorSizeA_,
ck_tile::index_t VectorSizeB_,
ck_tile::index_t VectorSizeC_>
struct GemmConfigFixedVectorSize : public GemmConfig
{
// Split-K partitions are aligned to K_Warp_Tile: preserve split-K remainder
static_assert(GemmConfig::K_Warp_Tile % VectorSizeA_ == 0 &&
GemmConfig::K_Warp_Tile % VectorSizeB_ == 0,
"A/B vector width must divide K_Warp_Tile");

// Enable K padding
static constexpr bool kPadK = true;

static constexpr bool FixedVectorSize = true;
static constexpr ck_tile::index_t VectorSizeA = VectorSizeA_;
static constexpr ck_tile::index_t VectorSizeB = VectorSizeB_;
static constexpr ck_tile::index_t VectorSizeC = VectorSizeC_;
};

template <typename PrecType>
Expand Down Expand Up @@ -193,6 +223,8 @@ struct GemmConfigComputeV3_WMMA : public GemmConfigBase
static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V3;

static constexpr int kBlockPerCu = 2;

static constexpr bool EnableKPadFallback = true;
Comment thread
chris-tsiaousis-hpc marked this conversation as resolved.
Resolved
};

template <typename PrecType>
Expand Down
136 changes: 111 additions & 25 deletions projects/composablekernel/example/ck_tile/03_gemm/run_gemm_example.inc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -347,31 +347,117 @@ int run_gemm_example_with_layouts(ck_tile::ArgParser& arg_parser,
c_m_n_dev_buf.SetZero();
c_m_n_dev_result.SetZero();

float ave_time = invoke_gemm<GemmConfig,
Invoker,
ADataTypeBuf,
BDataTypeBuf,
ck_tile::tuple<>,
AccDataType,
CDataType,
ALayout,
BLayout,
ck_tile::tuple<>,
CLayout>(a_m_k_dev_buf,
b_k_n_dev_buf,
c_m_n_dev_buf,
M,
N,
K,
stride_A,
stride_B,
stride_C,
kbatch,
n_warmup,
n_repeat,
persistent,
flush_cache,
rotating_count);
#ifdef CK_USE_GFX1250
Comment thread
chris-tsiaousis-hpc marked this conversation as resolved.
Resolved
// Enable fallback: gfx1250, RCR, PackedSize=1
constexpr bool k_pad_fallback_eligible =
GemmConfig::EnableKPadFallback && !GemmConfig::FixedVectorSize &&
std::is_same_v<ADataTypeBuf, BDataTypeBuf> &&
std::is_same_v<ALayout, ck_tile::tensor_layout::gemm::RowMajor> &&
std::is_same_v<BLayout, ck_tile::tensor_layout::gemm::ColumnMajor> &&
std::is_same_v<CLayout, ck_tile::tensor_layout::gemm::RowMajor> &&
ck_tile::numeric_traits<ADataTypeBuf>::PackedSize == 1 &&
ck_tile::numeric_traits<BDataTypeBuf>::PackedSize == 1;
#else
constexpr bool k_pad_fallback_eligible = false;
#endif

float ave_time;
if constexpr(k_pad_fallback_eligible)
{
using KPadConfig =
GemmConfigFixedVectorSize<GemmConfig,
1,
1,
static_cast<ck_tile::index_t>(16 / sizeof(CDataType))>;

if(K % (16 / static_cast<int>(sizeof(ADataTypeBuf))) != 0)
Comment thread
chris-tsiaousis-hpc marked this conversation as resolved.
Resolved
{
// Unaligned K
ave_time = invoke_gemm<KPadConfig,
Invoker,
ADataTypeBuf,
BDataTypeBuf,
ck_tile::tuple<>,
AccDataType,
CDataType,
ALayout,
BLayout,
ck_tile::tuple<>,
CLayout>(a_m_k_dev_buf,
b_k_n_dev_buf,
c_m_n_dev_buf,
M,
N,
K,
stride_A,
stride_B,
stride_C,
kbatch,
n_warmup,
n_repeat,
persistent,
flush_cache,
rotating_count);
}
else
{
// Aligned K
ave_time = invoke_gemm<GemmConfig,
Invoker,
ADataTypeBuf,
BDataTypeBuf,
ck_tile::tuple<>,
AccDataType,
CDataType,
ALayout,
BLayout,
ck_tile::tuple<>,
CLayout>(a_m_k_dev_buf,
b_k_n_dev_buf,
c_m_n_dev_buf,
M,
N,
K,
stride_A,
stride_B,
stride_C,
kbatch,
n_warmup,
n_repeat,
persistent,
flush_cache,
rotating_count);
}
Comment thread
chris-tsiaousis-hpc marked this conversation as resolved.
Resolved
}
else
{
// Do not use KPad Fallback
ave_time = invoke_gemm<GemmConfig,
Invoker,
ADataTypeBuf,
BDataTypeBuf,
ck_tile::tuple<>,
AccDataType,
CDataType,
ALayout,
BLayout,
ck_tile::tuple<>,
CLayout>(a_m_k_dev_buf,
b_k_n_dev_buf,
c_m_n_dev_buf,
M,
N,
K,
stride_A,
stride_B,
stride_C,
kbatch,
n_warmup,
n_repeat,
persistent,
flush_cache,
rotating_count);
}

c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data());

Expand Down
9 changes: 6 additions & 3 deletions projects/composablekernel/example/ck_tile/03_gemm/universal_gemm_invoker.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,10 @@ struct UniversalInvoker
ck_tile::element_wise::PassThrough,
ck_tile::element_wise::PassThrough,
AComputeDataType,
BComputeDataType>;
BComputeDataType,
GemmConfig::FixedVectorSize,
GemmConfig::VectorSizeA,
GemmConfig::VectorSizeB>;

using GemmPipeline = typename PipelineTypeTraits<
GemmConfig::Pipeline>::template GemmPipeline<UniversalGemmProblem>;
Expand All @@ -127,8 +130,8 @@ struct UniversalInvoker
GemmConfig::K_Warp_Tile,
UniversalGemmProblem::TransposeC,
GemmConfig::NumWaveGroups,
false, /*FixedVectorSize_*/
1, /*VectorSizeC_*/
GemmConfig::FixedVectorSize,
GemmConfig::VectorSizeC,
1, /*BlockedXDLN_PerWarp_*/
GemmConfig::DoubleSmemBuffer, /*DoubleSmemBuffer*/
AComputeDataType,
Expand Down
Loading