ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-14 02:57:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

fixed vector load siz for fp4

Browse Source
This commit is contained in:
Sami Remes
2026-01-16 12:04:34 -05:00
parent 16ca5cb532
commit f09e10936d
7 changed files with 66 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

60
example/ck_tile/42_mx_gemm/mx_gemm.cpp
View File

@@ -60,40 +60,40 @@ float invoke_mx_gemm(ck_tile::DeviceMem& a_dev_buf,
scale_m,
scale_n);
using GemmShape = ck_tile::TileGemmShape<
ck_tile::sequence<GemmConfig::M_Tile, GemmConfig::N_Tile, GemmConfig::K_Tile>,
ck_tile::sequence<GemmConfig::M_Warp, GemmConfig::N_Warp, GemmConfig::K_Warp>,
ck_tile::sequence<GemmConfig::M_Warp_Tile,
GemmConfig::N_Warp_Tile,
GemmConfig::K_Warp_Tile>>;
// using GemmShape = ck_tile::TileGemmShape<
// ck_tile::sequence<GemmConfig::M_Tile, GemmConfig::N_Tile, GemmConfig::K_Tile>,
// ck_tile::sequence<GemmConfig::M_Warp, GemmConfig::N_Warp, GemmConfig::K_Warp>,
// ck_tile::sequence<GemmConfig::M_Warp_Tile,
// GemmConfig::N_Warp_Tile,
// GemmConfig::K_Warp_Tile>>;
using TilePartitioner =
ck_tile::GemmSpatiallyLocalTilePartitioner<GemmShape,
GemmConfig::TileParitionerGroupNum,
GemmConfig::TileParitionerM01>;
// using TilePartitioner =
// ck_tile::GemmSpatiallyLocalTilePartitioner<GemmShape,
// GemmConfig::TileParitionerGroupNum,
// GemmConfig::TileParitionerM01>;
using MXGemmTraits = ck_tile::TileGemmUniversalTraits<GemmConfig::kPadM,
GemmConfig::kPadN,
GemmConfig::kPadK,
GemmConfig::DoubleSmemBuffer,
ALayout,
BLayout,
CLayout,
GemmConfig::TransposeC,
GemmConfig::UseStructuredSparsity,
UsePersistentKernel,
GemmConfig::NumWaveGroups,
false>;
// using MXGemmTraits = ck_tile::TileGemmUniversalTraits<GemmConfig::kPadM,
// GemmConfig::kPadN,
// GemmConfig::kPadK,
// GemmConfig::DoubleSmemBuffer,
// ALayout,
// BLayout,
// CLayout,
// GemmConfig::TransposeC,
// GemmConfig::UseStructuredSparsity,
// UsePersistentKernel,
// GemmConfig::NumWaveGroups,
// false>;
using MXPipelineProblem = MXGemmPipelineProblem<ADataType,
BDataType,
AccDataType,
GemmShape,
MXGemmTraits,
GemmConfig::Scheduler>;
// using MXPipelineProblem = MXGemmPipelineProblem<ADataType,
// BDataType,
// AccDataType,
// GemmShape,
// MXGemmTraits,
// GemmConfig::Scheduler>;
// Use the new MX comp_async pipeline with MX scaling support
using MXGemmPipeline = ck_tile::MXGemmPipelineAgBgCrCompAsync<MXPipelineProblem>;
// // Use the new MX comp_async pipeline with MX scaling support
// using MXGemmPipeline = ck_tile::MXGemmPipelineAgBgCrCompAsync<MXPipelineProblem>;
// Simplified invocation - comp_async handles hot loop and tail internally
auto invoke_splitk_path = [&](auto split_k_) {

4
example/ck_tile/42_mx_gemm/mx_gemm.hpp
View File

@@ -43,7 +43,7 @@ struct MXGemmHostArgs : ck_tile::UniversalGemmHostArgs<1, 1, 0>
struct MxGemmConfig
{
static constexpr ck_tile::index_t M_Tile = 128;
static constexpr ck_tile::index_t N_Tile = 256;
static constexpr ck_tile::index_t N_Tile = 128;
static constexpr ck_tile::index_t K_Tile = 512;
static constexpr ck_tile::index_t M_Warp = 1;
@@ -74,7 +74,7 @@ struct MxGemmConfig
struct MXfp4_GemmConfig16 : MxGemmConfig
{
static constexpr ck_tile::index_t M_Tile = 128;
static constexpr ck_tile::index_t N_Tile = 256;
static constexpr ck_tile::index_t N_Tile = 128;
static constexpr ck_tile::index_t K_Tile = 256;
};

18
example/ck_tile/42_mx_gemm/run_mx_gemm.inc
View File

@@ -154,17 +154,17 @@ int run_mx_gemm_example(int argc, char* argv[])
MXfp4_GemmConfig16,
true>(argc, argv, Row{}, Col{}, Row{});
}
else if(mx_prec == "fp8" || mx_prec == "fp8xfp8")
{
return run_mx_gemm_with_layouts<ck_tile::fp8_t,
ck_tile::fp8_t,
float,
MXfp8_GemmConfig16,
false>(argc, argv, Row{}, Col{}, Row{});
}
// else if(mx_prec == "fp8" || mx_prec == "fp8xfp8")
// {
// return run_mx_gemm_with_layouts<ck_tile::fp8_t,
// ck_tile::fp8_t,
// float,
// MXfp8_GemmConfig16,
// false>(argc, argv, Row{}, Col{}, Row{});
// }
else
{
throw std::runtime_error("Only fp4 and fp8 is supported currently!");
throw std::runtime_error("Only fp4 is supported currently!");
}
}
else

4
include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp
View File

@@ -673,8 +673,8 @@ struct UniversalGemmKernel
using AiLayout = remove_cvref_t<std::tuple_element_t<i.value, AsLayout>>;
using AiDataType = remove_cvref_t<std::tuple_element_t<i.value, AsDataType>>;
static_assert(GemmPipeline::GetVectorSizeA() == GemmPipeline::GetVectorSizeB(), "Vector size of A and B must be the same!");
static_assert(GemmPipeline::GetVectorSizeA() == 16, "Vector size of A must be 16!");
static_assert(GemmPipeline::GetVectorSizeB() == 16, "Vector size of B must be 16!");
static_assert(GemmPipeline::GetVectorSizeA() == 32, "Vector size of A must be 16!");
static_assert(GemmPipeline::GetVectorSizeB() == 32, "Vector size of B must be 16!");
if constexpr(std::is_same_v<AiLayout, tensor_layout::gemm::RowMajor>)
{
return make_naive_tensor_view<address_space_enum::global>(

6
include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp
View File

@@ -843,7 +843,7 @@ struct UniversalGemmBasePolicy
}
template <typename Problem>
CK_TILE_DEVICE static constexpr index_t GetSmemSizeA()
CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSizeA()
{
using ADataType = remove_cvref_t<typename Problem::ADataType>;
constexpr auto a_lds_block_desc = Derived::template MakeALdsBlockDescriptor<Problem>();
@@ -853,7 +853,7 @@ struct UniversalGemmBasePolicy
}
template <typename Problem>
CK_TILE_DEVICE static constexpr index_t GetSmemSizeB()
CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSizeB()
{
using BDataType =
std::conditional_t<std::is_same_v<typename Problem::BDataType, pk_fp4_raw_t>,
@@ -866,7 +866,7 @@ struct UniversalGemmBasePolicy
}
template <typename Problem>
CK_TILE_DEVICE static constexpr index_t GetSmemSize()
CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize()
{
constexpr index_t smem_size_a = GetSmemSizeA<Problem>();
constexpr index_t smem_size_b = GetSmemSizeB<Problem>();

4
include/ck_tile/ops/gemm_mx/pipeline/gemm_pipeline_ag_bg_cr_comp_async.hpp
View File

@@ -316,10 +316,10 @@ struct MXGemmPipelineAgBgCrCompAsync : public BaseMXGemmPipelineAgBgCrCompAsync<
number<BsLayout::size()>{});
/// Check tile window traits for vector size
using ATileDstr = remove_cvref_t<decltype(Policy::template MakeADramTileDistribution<Problem>())>;
// using ATileDstr = remove_cvref_t<decltype(Policy::template MakeADramTileDistribution<Problem>())>;
// static_assert(ATileDstr::LargestVec >= 16, "wrong! not implemented vector size");
// static_assert(ATileDstr::X1 >= 16, "wrong! not implemented vector size");
using BTileDstr = remove_cvref_t<decltype(Policy::template MakeBDramTileDistribution<Problem>())>;
// using BTileDstr = remove_cvref_t<decltype(Policy::template MakeBDramTileDistribution<Problem>())>;
// static_assert(BTileDstr::LargestVec >= 16, "wrong! not implemented vector size");
// static_assert(BTileDstr::X1 >= 16, "wrong! not implemented vector size");
using ATileType = remove_cvref_t<decltype(a_tile_windows[number<0>{}])>;

20
include/ck_tile/ops/gemm_mx/pipeline/gemm_pipeline_ag_bg_cr_comp_async_default_policy.hpp
View File

@@ -39,7 +39,15 @@ struct MXGemmPipelineAgBgCrCompAsyncDefaultPolicy
// constexpr index_t vector_size_for_16_bytes = 16 / sizeof(ADataType);
// return vector_size_for_16_bytes;
return 16;
static_assert(std::is_same_v<ADataType, pk_fp4_t>, "ADataType must be pk_fp4_t or pk_fp4_raw_t");
if constexpr(std::is_same_v<ADataType, pk_fp4_t> || std::is_same_v<ADataType, pk_fp4_raw_t>)
{
return 32;
}
else
{
return 16;
}
}
template <typename Problem, bool IsWave32Host = false>
@@ -55,7 +63,15 @@ struct MXGemmPipelineAgBgCrCompAsyncDefaultPolicy
// constexpr index_t vector_size_for_16_bytes = 16 / sizeof(BDataType);
// return vector_size_for_16_bytes;
return 16;
static_assert(std::is_same_v<BDataType, pk_fp4_t>, "BDataType must be pk_fp4_t or pk_fp4_raw_t");
if constexpr(std::is_same_v<BDataType, pk_fp4_t> || std::is_same_v<BDataType, pk_fp4_raw_t>)
{
return 32;
}
else
{
return 16;
}
}
// Override DRAM tile distributions to use the constrained vector sizes