ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-27 08:25:46 +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

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