seems working

This commit is contained in:
yadaish
2025-12-31 12:33:48 +00:00
parent 1530e9a9dc
commit ae31f7c0d7
11 changed files with 95 additions and 47 deletions

View File

@@ -602,7 +602,6 @@ include_directories(BEFORE
SET(BUILD_DEV ON CACHE BOOL "BUILD_DEV")
if(BUILD_DEV)
add_compile_options(-Werror)
add_compile_options(-Weverything)
endif()
message(STATUS "CMAKE_CXX_FLAGS: ${CMAKE_CXX_FLAGS}")
@@ -670,7 +669,6 @@ option(MIOPEN_REQ_LIBS_ONLY "Build only the MIOpen required libraries" OFF)
option(DISABLE_OFFLOAD_COMPRESS "Disable offload compress compiler flag when building instances" OFF)
option(BUILD_MHA_LIB "Build the static library for flash attention" OFF)
add_subdirectory(library)
if(NOT GPU_ARCHS AND USER_GPU_TARGETS AND NOT MIOPEN_REQ_LIBS_ONLY)
rocm_package_setup_component(tests
@@ -684,14 +682,11 @@ if(NOT GPU_ARCHS AND USER_GPU_TARGETS AND NOT MIOPEN_REQ_LIBS_ONLY)
)
add_subdirectory(example)
add_subdirectory(tutorial)
rocm_package_setup_component(tutorials
LIBRARY_NAME composablekernel
PACKAGE_NAME tutorials
)
add_subdirectory(tile_engine)
if(BUILD_TESTING)
add_subdirectory(test)
endif()
endif()
@@ -700,15 +695,12 @@ if (NOT MIOPEN_REQ_LIBS_ONLY)
LIBRARY_NAME composablekernel
PACKAGE_NAME ckprofiler
)
add_subdirectory(profiler)
endif()
if (CK_EXPERIMENTAL_BUILDER)
add_subdirectory(experimental/builder)
endif()
if(CK_USE_CODEGEN AND (SUPPORTED_GPU_TARGETS MATCHES "gfx9" OR GPU_ARCHS))
add_subdirectory(codegen)
endif()
#Create an interface target for the include only files and call it "composablekernels"

View File

@@ -27,6 +27,9 @@ if(has_supported_gpu)
add_executable(tile_example_grouped_flatmm grouped_flatmm.cpp)
target_compile_options(tile_example_grouped_flatmm PRIVATE ${EXAMPLE_FLATMM_COMPILE_OPTIONS})
add_executable(tile_example_a4w4_moe_flatmm mixed_prec/a4w4_moe_flatmm.cpp)
target_compile_options(tile_example_a4w4_moe_flatmm PRIVATE ${EXAMPLE_FLATMM_COMPILE_OPTIONS})
if (GPU_TARGETS MATCHES "gfx95")
add_executable(tile_example_mixed_prec_flatmm mixed_prec/mixed_prec_flatmm.cpp)
target_compile_options(tile_example_mixed_prec_flatmm PRIVATE ${EXAMPLE_FLATMM_COMPILE_OPTIONS})
@@ -37,8 +40,8 @@ if(has_supported_gpu)
add_executable(tile_example_a8w4_moe_flatmm mixed_prec/a8w4_moe_flatmm.cpp)
target_compile_options(tile_example_a8w4_moe_flatmm PRIVATE ${EXAMPLE_FLATMM_COMPILE_OPTIONS})
add_executable(tile_example_a4w4_moe_flatmm mixed_prec/a4w4_moe_flatmm.cpp)
target_compile_options(tile_example_a4w4_moe_flatmm PRIVATE ${EXAMPLE_FLATMM_COMPILE_OPTIONS})
# add_executable(tile_example_a4w4_moe_flatmm mixed_prec/a4w4_moe_flatmm.cpp)
# target_compile_options(tile_example_a4w4_moe_flatmm PRIVATE ${EXAMPLE_FLATMM_COMPILE_OPTIONS})
include(mxgemm/mx_flatmm_instance.cmake)
mx_flatmm_instance_generate(EXAMPLE_MX_FLATMM_FILES)

View File

@@ -99,6 +99,9 @@ float a4w4_moe_flatmm(const MoeFlatmmHostArgs& args, const ck_tile::stream_confi
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);
std::cerr << "k_grain: " << k_grain << ", K_split: " << K_split
<< ", num_loop: " << num_loop << ", has_hot_loop: " << has_hot_loop
<< ", tail_num: " << static_cast<int>(tail_num) << std::endl;
float ave_time{0};
const auto Run = [&](const auto has_hot_loop_,
@@ -123,7 +126,7 @@ float a4w4_moe_flatmm(const MoeFlatmmHostArgs& args, const ck_tile::stream_confi
using GemmEpilogue = ck_tile::CShuffleEpilogue<
ck_tile::CShuffleEpilogueProblem<ADataType,
BDataType,
CDataType,
DsDatatype,
AccDataType,
CDataType,
DsLayout,
@@ -275,8 +278,8 @@ void shuffle_mxfp4_weight(const IterSrc src, IterDst dst, int experts_cnt, int N
}
}
template <typename FlatmmConfig, ck_tile::MoeFlatmmKind moe_kind, bool KLast,Ltypename T>
auto shuffle_mxfp4_scale(const ck_tile::HostTensor<T>& scale, int experts_cnt)
template <typename FlatmmConfig, ck_tile::MoeFlatmmKind moe_kind, bool KLast, typename T>
auto shuffle_mxfp4_scale(const ck_tile::HostTensor<T>& scale, int experts_cnt, bool is_activation = false)
{
assert(scale.get_lengths().size() == 2);
int n_ = scale.get_lengths()[1];
@@ -308,6 +311,19 @@ auto shuffle_mxfp4_scale(const ck_tile::HostTensor<T>& scale, int experts_cnt)
std::copy(scale.begin(), scale.end(), shfl_scale.begin());
return ck_tile::reference_permute(shfl_scale, {0, 5, 1, 3, 6, 2, 4});
}
else if (is_activation)
{
ck_tile::HostTensor<T> shfl_scale({
n_ / FlatmmConfig::N_Warp_Tile / N_Pack,
N_Pack,
FlatmmConfig::N_Warp_Tile,
k_per_expert / K_Pack / K_Lane,
K_Pack,
K_Lane,
});
std::copy(scale.begin(), scale.end(), shfl_scale.begin());
return ck_tile::reference_permute(shfl_scale, {0, 3, 5, 2, 4, 1});
}
else
{
ck_tile::HostTensor<T> shfl_scale({
@@ -324,7 +340,7 @@ auto shuffle_mxfp4_scale(const ck_tile::HostTensor<T>& scale, int experts_cnt)
}
}
#include "run_a8w4_moe_flatmm_example.inc"
#include "run_a4w4_moe_flatmm_example.inc"
template <typename FlatmmConfig>
int run_a4w4_moe_flatmm_example(int argc, char* argv[])
@@ -402,7 +418,7 @@ int main(int argc, char* argv[])
int warp_tile = arg_parser.get_int("warp_tile");
if(warp_tile == 0)
{
return !run_mx_moe_flatmm_example<A4W4_FlatmmConfig16>(argc, argv);
return !run_a4w4_moe_flatmm_example<A4W4_FlatmmConfig16>(argc, argv);
}
else
{
@@ -414,4 +430,4 @@ int main(int argc, char* argv[])
std::cerr << "Runtime error: " << e.what() << '\n';
return EXIT_FAILURE;
}
}
}

View File

@@ -13,7 +13,7 @@
// GEMM config with 16x16 warp tile for FP4×FP4 MoE
struct A4W4_FlatmmConfig16
{
static constexpr ck_tile::index_t M_Tile = 128;
static constexpr ck_tile::index_t M_Tile = 32;
static constexpr ck_tile::index_t N_Tile = 256;
static constexpr ck_tile::index_t K_Tile = 256;

View File

@@ -91,10 +91,11 @@ int run_a4w4_moe_flatmm_with_layouts(int argc,
const ck_tile::index_t MPerBlock = FlatmmConfig::M_Tile;
ck_tile::index_t sorted_tile_num = (num_tokens + MPerBlock - 1) / MPerBlock * MPerBlock * topk;
ck_tile::index_t valid_tile_num = sorted_tile_num;
ck_tile::index_t valid_tile_num = 1;
ck_tile::index_t sorted_size = sorted_tile_num * MPerBlock;
const ck_tile::index_t M = sorted_tile_num * MPerBlock;
const ck_tile::index_t M = valid_tile_num * MPerBlock;
std::cerr << "M = " << M << " " << "num_tokens = " << num_tokens << std::endl;
const ck_tile::index_t outputN = kind == ck_tile::MoeFlatmmKind::kFFN_gemm1_gate_up ? N / 2 : N;
static_assert(std::is_same_v<BLayout, ck_tile::tensor_layout::gemm::ColumnMajor>);
@@ -115,6 +116,7 @@ int run_a4w4_moe_flatmm_with_layouts(int argc,
auto c_m_n_tensor = ck_tile::HostTensor<CDataType>(ck_tile::host_tensor_descriptor(
IsInputGemm ? num_tokens * topk : num_tokens, outputN, stride_C, is_row_major(CLayout{})));
// TODO(yadai), fix me for gemm1
ck_tile::HostTensor<ScaleType> scale_a(ck_tile::HostTensorDescriptor(
{ M / ScaleGranularityM, K / ScaleGranularityK},
{K / ScaleGranularityK, 1}));
@@ -123,29 +125,38 @@ int run_a4w4_moe_flatmm_with_layouts(int argc,
if(init_method == 0)
{
ck_tile::FillUniformDistribution<ADataType>{0.0f, 1.0f}(a_m_k_tensor);
ck_tile::FillUniformDistribution<BDataType>{-.5f, .5f}(b_k_n_tensor);
ck_tile::FillUniformDistribution<ScaleType>{0.f, 1.f}(scale_a);
ck_tile::FillUniformDistribution<ScaleType>{0.f, 1.f}(scale_b);
ck_tile::FillUniformDistribution<ADataType>{0.f, 1.f}(a_m_k_tensor);
ck_tile::FillUniformDistribution<BDataType>{0.f, 1.f}(b_k_n_tensor);
ck_tile::FillUniformDistribution<ScaleType>{0.2f, 1.f}(scale_a);
ck_tile::FillUniformDistribution<ScaleType>{0.2f, 1.f}(scale_b);
}
else
{
ck_tile::FillUniformDistribution<ADataType>{1.0f, 1.0f}(a_m_k_tensor);
ck_tile::FillUniformDistribution<BDataType>{1.0f, 1.0f}(b_k_n_tensor);
ck_tile::FillUniformDistribution<ScaleType>{1.0f, 1.0f}(scale_a);
ck_tile::FillUniformDistribution<ScaleType>{0.0f, 1.0f}(scale_a);
ck_tile::FillUniformDistribution<ScaleType>{1.0f, 1.0f}(scale_b);
}
/*
for (auto i = 0; i < M / ScaleGranularityM; ++i) {
for (auto j = 0; j < K / ScaleGranularityK; ++j) {
scale_a.data()[i * K / ScaleGranularityK + j] = ck_tile::type_convert<ScaleType>(1.0f * (i + 1));
}
}
*/
ck_tile::HostTensor<BDataType> b_shuffle_host(
ck_tile::host_tensor_descriptor(K, experts * N, stride_B, is_row_major(b_layout)));
shuffle_mxfp4_weight<FlatmmConfig, kind>(
b_k_n_tensor.begin(), b_shuffle_host.begin(), experts, N, K);
ck_tile::HostTensor<ScaleType> scale_a_shuffle =
shuffle_mxfp4_scale<FlatmmConfig, kind, false>(scale_a, 1);
shuffle_mxfp4_scale<FlatmmConfig, kind, false>(scale_a, 1, true);
ck_tile::HostTensor<ScaleType> scale_b_shuffle =
shuffle_mxfp4_scale<FlatmmConfig, kind, true>(scale_b, experts);
ck_tile::DeviceMem scale_a_dev_buf(scale_a.get_element_space_size_in_bytes());
ck_tile::DeviceMem scale_a_shuffle_dev_buf(scale_a_shuffle.get_element_space_size_in_bytes());
ck_tile::DeviceMem scale_b_shuffle_dev_buf(scale_b_shuffle.get_element_space_size_in_bytes());
@@ -222,6 +233,8 @@ int run_a4w4_moe_flatmm_with_layouts(int argc,
max_token_id_dev.ToDevice(max_token_id.data());
expert_weight_dev.ToDevice(expert_weight.data());
expert_bias_dev.ToDevice(expert_bias.data());
scale_a_dev_buf.ToDevice(scale_a.data());
scale_a_shuffle_dev_buf.ToDevice(scale_a_shuffle.data());
scale_b_shuffle_dev_buf.ToDevice(scale_b_shuffle.data());
@@ -243,6 +256,7 @@ int run_a4w4_moe_flatmm_with_layouts(int argc,
static_cast<float*>(scale_b_shuffle_dev_buf.GetDeviceBuffer()), N / ScaleGranularityN};
auto exp_bias_dev_ptr = ck_tile::FlatmmScalePointer<1>{
static_cast<float*>(expert_bias_dev.GetDeviceBuffer()), experts * N};
using MoeFlatmmArgs = ck_tile::MoeFlatmmHostArgs<
ck_tile::FlatmmScalePointer<ScaleGranularityM, ScaleGranularityK>,
@@ -268,8 +282,7 @@ int run_a4w4_moe_flatmm_with_layouts(int argc,
scale_a_shuffle_dev_ptr,
scale_b_shuffle_dev_ptr,
exp_bias_dev_ptr};
invoke_mx_moe_flatmm<FlatmmConfig,
invoke_a4w4_moe_flatmm<FlatmmConfig,
ADataType,
BDataType,
ck_tile::tuple<>,
@@ -280,7 +293,6 @@ int run_a4w4_moe_flatmm_with_layouts(int argc,
ck_tile::tuple<>,
CLayout,
kind>(warmup, repeat, gemm_desc);
c_m_n_dev_buf.FromDevice(c_m_n_tensor.data());
bool pass{true};
@@ -292,15 +304,23 @@ int run_a4w4_moe_flatmm_with_layouts(int argc,
stride_C,
is_row_major(CLayout{})));
c_m_n_host_ref.SetZero();
std::cerr << "tag0" << std::endl;
// Convert scale_a from e8m0 to float
ck_tile::HostTensor<AccDataType> scale_a_float(ck_tile::HostTensorDescriptor(
{(IsInputGemm ? num_tokens : M) / ScaleGranularityM, K / ScaleGranularityK},
{K / ScaleGranularityK, 1}));
/*
ck_tile::HostTensor<AccDataType> scale_a_float(ck_tile::HostTensorDescriptor(
{(M) / ScaleGranularityM, K / ScaleGranularityK},
{K / ScaleGranularityK, 1}));
*/
std::cerr << "tag0-0" << std::endl;
std::copy(scale_a.begin(), scale_a.end(), scale_a_float.begin());
std::cerr << "tag0-1" << std::endl;
ck_tile::DeviceMem scale_a_float_dev_buf(scale_a_float.get_element_space_size_in_bytes());
scale_a_float_dev_buf.ToDevice(scale_a_float.data());
std::cerr << "tag1" << std::endl;
// Convert scale_b from e8m0 to float
ck_tile::HostTensor<AccDataType> scale_b_float(ck_tile::HostTensorDescriptor(
{K * experts / ScaleGranularityK, N / ScaleGranularityN}, {N / ScaleGranularityN, 1}));
@@ -311,7 +331,7 @@ int run_a4w4_moe_flatmm_with_layouts(int argc,
std::unique_ptr<ck_tile::DeviceMem> c_m_n_ref_buf =
std::make_unique<ck_tile::DeviceMem>(c_m_n_tensor.get_element_space_size_in_bytes());
c_m_n_ref_buf->SetZero();
std::cerr << "tag2" << std::endl;
ck_tile::reference_moe_gemm_gpu<ADataType,
BDataType,
AccDataType,
@@ -343,12 +363,12 @@ int run_a4w4_moe_flatmm_with_layouts(int argc,
static_cast<float*>(scale_a_float_dev_buf.GetDeviceBuffer()),
static_cast<float*>(scale_b_float_dev_buf.GetDeviceBuffer()),
static_cast<float*>(expert_bias_dev.GetDeviceBuffer()));
std::cerr << "tag3" << std::endl;
c_m_n_ref_buf->FromDevice(c_m_n_host_ref.data());
std::cerr << "tag4" << std::endl;
const float rtol = std::is_same_v<ADataType, ck_tile::pk_fp4_t> && IsInputGemm ? 1e-2 : 1e-2;
const float atol = std::is_same_v<ADataType, ck_tile::pk_fp4_t> && IsInputGemm ? 1e-2 : 1e-2;
std::cerr << "tag5" << std::endl;
pass = ck_tile::check_err(
c_m_n_tensor, c_m_n_host_ref, "Error: Incorrect results!", rtol, atol);

View File

@@ -410,7 +410,6 @@ struct tile_scatter_gather
constexpr index_t d =
tile_dstr.get_ys_to_d_descriptor().calculate_offset(idx_ys) /
Traits::PackedSize;
dst_tensor.get_thread_buffer().template at<d>() =
vec_value.template get_as<DataType>()[j / Traits::PackedSize];
});
@@ -440,6 +439,10 @@ struct tile_scatter_gather
}
});
});
/*
for (auto i = 0; i < dst_tensor.get_thread_buffer().size(); ++i)
printf("gather info %s:%d %d %d %d %f\n", __FILE__, __LINE__, blockIdx.x, threadIdx.x, i, type_convert<float>(dst_tensor.get_thread_buffer()[i]));
*/
}
template <typename LdsTileWindow_,

View File

@@ -19,7 +19,7 @@
namespace ck_tile {
/** @brief Maximum number of error values to display when checking errors */
constexpr int ERROR_DETAIL_LIMIT = 128;
constexpr int ERROR_DETAIL_LIMIT = 1000000;
/** @brief 8-bit floating point type */
using F8 = ck_tile::fp8_t;

View File

@@ -83,9 +83,9 @@ __global__ void moe_gemm_kernel(const ck_tile::index_t* p_sorted_token_ids_,
AccDataType acc_temp = 0.0;
AccDataType acc_up_temp = 0.0;
float scale_A = 0;
float scale_B = 0;
float scale_B_up = 0;
float scale_A = 1.0;
float scale_B = 1.0;
float scale_B_up = 1.0;
index_t scale_A_stride = (M + scale_granularity_m - 1) / scale_granularity_m;
index_t scale_B_stride = (N + scale_granularity_n - 1) / scale_granularity_n;
@@ -102,8 +102,8 @@ __global__ void moe_gemm_kernel(const ck_tile::index_t* p_sorted_token_ids_,
acc_temp = 0.0;
acc_up_temp = 0.0;
// update scale factors
scale_A = scale_A_ptr[(gather_token_id / scale_granularity_m) +
(k / scale_granularity_k) * scale_A_stride];
scale_A = scale_A_ptr[(gather_token_id / scale_granularity_m) * (K / scale_granularity_k) +
(k / scale_granularity_k)];
scale_B =
scale_B_ptr[expert_id * scale_B_expert_stride + col / scale_granularity_n +
(k / scale_granularity_k) * scale_B_stride];
@@ -144,7 +144,7 @@ __global__ void moe_gemm_kernel(const ck_tile::index_t* p_sorted_token_ids_,
}
else if constexpr(std::is_same_v<ADataType, pk_fp4_t>)
{
const fp32x2_t fp32_val = pk_fp4_to_fp32x2(A[a_index / packed_size_a]);
const fp32x2_t fp32_val = pk_fp4_to_fp32x2(A[a_index / packed_size_a], 1.0f);
if(k % 2 == 1)
v_a = fp32_val.hi;
else
@@ -217,6 +217,10 @@ __global__ void moe_gemm_kernel(const ck_tile::index_t* p_sorted_token_ids_,
{
C[c_index] = ck_tile::type_convert<CDataType>(
ActivationOp{}(acc + bias, MoeGemmKind == 1 ? acc_up + bias_up : 1));
if (c_index == 0) {
printf("%s:%d acc = %f, acc_up = %f, bias = %f, bias_up = %d\n", __FILE__, __LINE__,
ck_tile::type_convert<float>(acc), ck_tile::type_convert<float>(acc_up), ck_tile::type_convert<float>(bias), ck_tile::type_convert<float>(bias_up));
}
}
else
{

View File

@@ -922,6 +922,8 @@ struct MoeFlatmmKernel
? gather_token_id * kargs.stride_A
: gather_token_id;
// printf("lane, %d a_offset%d \n", threadIdx.x, a_offsets[m0]);
printf("%s:%d row_idx = %ld, gather_token_id=%ld, a_offset=%ld\n", __FILE__, __LINE__,
row_idx, gather_token_id, a_offsets[m0]);
});
const SplitKBatchOffset splitk_batch_offset(kargs);

View File

@@ -190,10 +190,10 @@ struct F16xMXF4FlatmmPipelineAgBgCrPolicy : UniversalFlatmmPipelineAgBgCrPolicy
return make_static_tile_distribution(
tile_distribution_encoding<
sequence<WaveRepeat>, // ?
tuple<sequence<NWavePerBlk, N_Pack>, // second
sequence<WaveRepeat>, // ? // 1
tuple<sequence<NWavePerBlk, N_Pack>, // second // 4, 2
// direction
sequence<KWavePerBlk, KThdPerWave, KBPerLoad>>, // first direction
sequence<KWavePerBlk, KThdPerWave, KBPerLoad>>, // first direction // 1, 32, 32
// wave in blk, // thd in wave
// <M, K> // <M, K>
tuple<sequence<0, 1, 2>, sequence<2>>, // which direction
@@ -855,6 +855,10 @@ struct F8xMXF4FlatmmPipelineAgBgCrPolicy : UniversalFlatmmPipelineAgBgCrPolicy
constexpr index_t num_access_v = 1;
constexpr int K1 = K_Thread / num_access_v; // 16
// printf("debug tile distribution = %d %d\n", num_access_v, K1);
if (blockIdx.x == 0 && threadIdx.x == 0)
printf("%s:%d num_access_v = %d %d %d\n", __FILE__, __LINE__, num_access_v, K_Lane, K1, M_warps, N_warps);
return make_static_tile_distribution(
std::conditional_t<
num_access_v == 1,
@@ -891,16 +895,19 @@ struct F8xMXF4FlatmmPipelineAgBgCrPolicy : UniversalFlatmmPipelineAgBgCrPolicy
constexpr index_t NWavePerBlk = TileShape::BlockWarps::at(number<1>{}); // N_Warp
constexpr index_t WaveRepeat = WaveNum / TileShape::flatNPerWarp;
constexpr index_t WaveRepeat = WaveNum / TileShape::flatNPerWarp; // 1
constexpr index_t kKPerThread = 32;
constexpr index_t num_access_v = static_cast<index_t>(wg_attr_num_access<Problem>);
constexpr index_t K2 = kKPerThread / num_access_v;
if (blockIdx.x == 0 && threadIdx.x == 0)
printf("debug B tile distribution = %d %d %d\n", num_access_v, K0, K2);
return make_static_tile_distribution(
std::conditional_t< //
num_access_v == 1,
tile_distribution_encoding< //
sequence<WaveRepeat>,
sequence<WaveRepeat>, // 1
tuple<sequence<NWavePerBlk, NXdlPack>, // 4 2
sequence<K0, K1, K2>>, // 1 64 32
tuple<sequence<0, 1, 2>, sequence<2>>,

View File

@@ -376,6 +376,7 @@ struct FlatmmPipelineProblem
}
else
{
// return PackedSize * VectorLoadSize / sizeof(ADataType);
return VectorLoadSize / sizeof(ADataType);
}
}