mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-19 02:01:01 +00:00
seems working
This commit is contained in:
@@ -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"
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
|
||||
@@ -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);
|
||||
|
||||
|
||||
@@ -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_,
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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
|
||||
{
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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>>,
|
||||
|
||||
@@ -376,6 +376,7 @@ struct FlatmmPipelineProblem
|
||||
}
|
||||
else
|
||||
{
|
||||
// return PackedSize * VectorLoadSize / sizeof(ADataType);
|
||||
return VectorLoadSize / sizeof(ADataType);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user