diff --git a/example/ck_tile/36_quant/CMakeLists.txt b/example/ck_tile/36_quant/CMakeLists.txt index 6529f8a447..ce6e15136b 100644 --- a/example/ck_tile/36_quant/CMakeLists.txt +++ b/example/ck_tile/36_quant/CMakeLists.txt @@ -23,4 +23,5 @@ endfunction(add_per_tensor_quant_example TARGET_NAME MAIN_SRC) add_per_tensor_quant_example(tile_example_static_per_tensor_quant example_static_per_tensor_quant.cpp) add_per_tensor_quant_example(tile_example_dynamic_per_tensor_quant example_dynamic_per_tensor_quant.cpp) +add_per_tensor_quant_example(tile_example_dynamic_per_token_quant example_dynamic_per_token_quant.cpp) diff --git a/example/ck_tile/36_quant/example_dynamic_per_tensor_quant.cpp b/example/ck_tile/36_quant/example_dynamic_per_tensor_quant.cpp index 15d00995c5..2d78f9331a 100644 --- a/example/ck_tile/36_quant/example_dynamic_per_tensor_quant.cpp +++ b/example/ck_tile/36_quant/example_dynamic_per_tensor_quant.cpp @@ -136,15 +136,16 @@ int main() constexpr dim3 blocks = Kernel::BlockSize(); auto kargs = Kernel::MakeKargs(a); - ck_tile::launch_kernel( - {}, ck_tile::make_kernel(Kernel{}, grids, blocks, 0, kargs)); - + auto s = ck_tile::stream_config{nullptr, true, 1, 5, 10}; + auto time = ck_tile::launch_kernel( + s, ck_tile::make_kernel(Kernel{}, grids, blocks, 0, kargs)); + std::cout<<"time: "<( x_host, scale_host, qx_host_ref); - // std::cout<(); diff --git a/example/ck_tile/36_quant/example_static_per_tensor_quant.cpp b/example/ck_tile/36_quant/example_static_per_tensor_quant.cpp index 8fb06e3229..4e3490310d 100644 --- a/example/ck_tile/36_quant/example_static_per_tensor_quant.cpp +++ b/example/ck_tile/36_quant/example_static_per_tensor_quant.cpp @@ -138,9 +138,11 @@ int main() constexpr dim3 blocks = Kernel::BlockSize(); auto kargs = Kernel::MakeKargs(a); - ck_tile::launch_kernel( - {}, ck_tile::make_kernel(Kernel{}, grids, blocks, 0, kargs)); + auto s = ck_tile::stream_config{nullptr, true, 1, 5, 10}; + auto time = ck_tile::launch_kernel( + s, ck_tile::make_kernel(Kernel{}, grids, blocks, 0, kargs)); + std::cout<<"time: "<( diff --git a/include/ck_tile/host/reference/reference_rowwise_quantization2d.hpp b/include/ck_tile/host/reference/reference_rowwise_quantization2d.hpp index 9b7eeb06b5..54949f3e1f 100644 --- a/include/ck_tile/host/reference/reference_rowwise_quantization2d.hpp +++ b/include/ck_tile/host/reference/reference_rowwise_quantization2d.hpp @@ -50,4 +50,24 @@ CK_TILE_HOST void reference_per_tensor_quantization2d(const HostTensor +CK_TILE_HOST void reference_per_token_quantization2d(const HostTensor& x_m_n, + const HostTensor& scale_m, + HostTensor& qx_m_n) +{ + auto f = [&](auto m) { + const int N = x_m_n.mDesc.get_lengths()[1]; + + for(int n = 0; n < N; ++n) + { + auto v_x = x_m_n(m, n); + auto v_qx = v_x / scale_m(m); + qx_m_n(m, n) = type_convert(saturates{}(v_qx)); + } + }; + + make_ParallelTensorFunctor(f, + x_m_n.mDesc.get_lengths()[0])(std::thread::hardware_concurrency()); +} + } // namespace ck_tile diff --git a/include/ck_tile/ops/quant/kernel/quant_kernel.hpp b/include/ck_tile/ops/quant/kernel/quant_kernel.hpp index 70a4a7eeef..d4d66b844c 100644 --- a/include/ck_tile/ops/quant/kernel/quant_kernel.hpp +++ b/include/ck_tile/ops/quant/kernel/quant_kernel.hpp @@ -116,4 +116,130 @@ struct PerTensorQuant Pipeline{}(x_window, scale, kargs.n, qx_window, smem); } }; + +template +struct PerTokenQuant +{ + using Pipeline = remove_cvref_t; + using Problem = typename Pipeline::Problem; + + using XDataType = remove_cvref_t; + using ScaleDataType = remove_cvref_t; + using ComputeDataType = remove_cvref_t; + using QXDataType = remove_cvref_t; + + static constexpr index_t Block_M = Problem::BlockShape::Block_M; + static constexpr index_t Block_N = Problem::BlockShape::Block_N; + static constexpr bool kPadM = false; // always no need to pad along M + static constexpr bool kPadN = Problem::kPadN; + + static constexpr index_t ThreadPerWarp_N = Problem::BlockShape::ThreadPerWarp_N; + static constexpr index_t Vector_N = Problem::BlockShape::Vector_N; + static constexpr index_t Repeat_N = Problem::BlockShape::Repeat_N; + + struct Kargs + { + const void* p_x; + void* p_scale; + void* p_qx; + + index_t m; + index_t n; + index_t x_stride; // input row_stride + }; + + CK_TILE_HOST static constexpr Kargs MakeKargs(const Kargs& kargs) + { + return kargs; + } + + CK_TILE_HOST static constexpr auto GridSize(const Kargs& kargs) + { + return dim3(integer_divide_ceil(kargs.m, Block_M)); + } + + CK_TILE_HOST static constexpr auto BlockSize() { return Problem::BlockShape::BlockSize; } + + // clang-format off + template struct t2s; + template <> struct t2s { static constexpr const char * name = "fp32"; }; + template <> struct t2s { static constexpr const char * name = "fp16"; }; + template <> struct t2s { static constexpr const char * name = "bf16"; }; + template <> struct t2s { static constexpr const char * name = "fp8"; }; + template <> struct t2s { static constexpr const char * name = "bf8"; }; + // clang-format on + + // in byte + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() { return Pipeline::GetSmemSize(); } + + CK_TILE_HOST static std::string GetName() + { + // clang-format off + using S_ = typename Problem::BlockShape; + auto surfix = [&] () { + std::string n; + if (kPadN) n += "_pn"; + return n; }(); + + #define _SS_ std::string + #define _TS_ std::to_string + return _SS_("quant_fwd_") + _SS_(t2s::name) + "_" + + _TS_(S_::Block_M) + "x" + _TS_(S_::Block_N) + "_" + _TS_(S_::WarpPerBlock_M) + "x" + _TS_(S_::WarpPerBlock_N) + "_" + + _TS_(S_::Warp_M) + "x" + _TS_(S_::Warp_N) + "_" + _TS_(S_::Vector_M) + "x" + _TS_(S_::Vector_N) + "_" + + _SS_(Pipeline::name) + surfix; + #undef _SS_ + #undef _TS_ + // clang-format on + } + + CK_TILE_DEVICE void operator()(Kargs kargs) const + { + const auto iM = get_block_id() * Block_M; + + const auto x_window = [&]() { + const auto tmp_ = make_naive_tensor_view( + static_cast(kargs.p_x), + make_tuple(kargs.m, kargs.n), + make_tuple(kargs.x_stride, 1), + number{}, + number<1>{}); + + const auto tmp2_ = pad_tensor_view( + tmp_, make_tuple(number{}, number{}), sequence{}); + return make_tile_window( + tmp2_, make_tuple(number{}, number{}), {iM, 0}); + }(); + + auto scale_window = [&]() { + auto tmp_ = make_naive_tensor_view( + static_cast(kargs.p_scale), + make_tuple(kargs.m), + make_tuple(1), + number<1>{}, + number<1>{}); + + auto tmp2_ = pad_tensor_view( + tmp_, make_tuple(number{}), sequence{}); + return make_tile_window( + tmp2_, make_tuple(number{}), {iM}); + }(); + + auto qx_window = [&]() { + auto tmp_ = make_naive_tensor_view( + static_cast(kargs.p_qx), + make_tuple(kargs.m, kargs.n), + make_tuple(kargs.x_stride, 1), + number{}, + number<1>{}); + + auto tmp2_ = pad_tensor_view( + tmp_, make_tuple(number{}, number{}), sequence{}); + return make_tile_window( + tmp2_, make_tuple(number{}, number{}), {iM, 0}); + }(); + __shared__ char smem[GetSmemSize()]; + Pipeline{}(x_window, scale_window, kargs.n, qx_window, smem); + } +}; + } // namespace ck_tile diff --git a/include/ck_tile/ops/quant/pipeline/quant_pipeline.hpp b/include/ck_tile/ops/quant/pipeline/quant_pipeline.hpp index 756282623a..c31a9c40a9 100644 --- a/include/ck_tile/ops/quant/pipeline/quant_pipeline.hpp +++ b/include/ck_tile/ops/quant/pipeline/quant_pipeline.hpp @@ -22,7 +22,7 @@ struct StaticPerTensorQuantPipeline static constexpr const char* name = []() { - return "static_quant_op"; + return "static_per_tensor_quant_op"; }(); CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() @@ -72,7 +72,7 @@ struct DynamicPerTensorQuantPipeline static constexpr bool UseMax3 = true; static constexpr const char* name = []() { - return "dynamic_quant_op"; + return "dynamic_per_tensorquant_op"; }(); CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() @@ -144,4 +144,91 @@ struct DynamicPerTensorQuantPipeline } }; +template +struct DynamicPerTokenQuantPipeline +{ + using Problem = ck_tile::remove_cvref_t; + using Policy = ck_tile::remove_cvref_t; + + using XDataType = ck_tile::remove_cvref_t; + using ScaleDataType = ck_tile::remove_cvref_t; + using ComputeDataType = ck_tile::remove_cvref_t; + using QXDataType = ck_tile::remove_cvref_t; + + static constexpr bool UseMax3 = true; + + static constexpr const char* name = []() { + return "dynamic_per_token_quant_op"; + }(); + + CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize() + { + return 1; + } + + template + CK_TILE_DEVICE auto operator()(const XWindow& x_window_, + ScaleWindow& scale_window, + ck_tile::index_t row_size, + QXWindow& qx_window, + void*) const + { + auto x_window = + make_tile_window(x_window_, Policy::template MakeXBlockTileDistribution()); + auto origin = x_window.get_window_origin(); + static constexpr index_t Block_N = Problem::BlockShape::Block_N; + index_t num_n_tile_iteration = + __builtin_amdgcn_readfirstlane(integer_divide_ceil(row_size, Block_N)); + + auto reduce_absmax_func = ReduceOp::AbsMax{}; + auto reduce_absmax3_func = [](auto acc_, auto v_0_, auto v_1_) { + float rtn; + asm volatile("v_max3_f32 %0, %1, abs(%2), abs(%3)" + : "=v"(rtn) + : "v"(acc_), "v"(v_0_), "v"(v_1_)); + return rtn; + }; + + auto block_reduce2d = Policy::template GetBlockReduce2d(); + + using XTensorType = decltype(cast_tile(load_tile(x_window))); + auto absmax = block_reduce2d.template MakeYBlockTile(); + set_tile(absmax, reduce_absmax_func.GetIdentityValue()); + + for(int iN = __builtin_amdgcn_readfirstlane(0); iN < num_n_tile_iteration; ++iN){ + const auto x = load_tile(x_window); + constexpr auto x_size_per_row = + x.get_tile_distribution().get_ys_to_d_descriptor().get_lengths().at(number<1>{}); + if constexpr(UseMax3 && std::is_same_v && + x_size_per_row % 2 == 0) + block_reduce2d(x, absmax, reduce_absmax3_func, sequence<1, 2>{}); + else + block_reduce2d(x, absmax, reduce_absmax_func); + move_tile_window(x_window, {0, Block_N}); + } + + auto scale_tile = tile_elementwise_in( + [&](const auto& v_) { + return v_ / type_convert(numeric::max()); + }, + absmax); + + store_tile(scale_window, cast_tile(scale_tile)); + + x_window.set_window_origin(origin); + for(int iN = __builtin_amdgcn_readfirstlane(0); iN < num_n_tile_iteration; ++iN){ + const auto x = cast_tile(load_tile(x_window)); + auto qx = make_static_distributed_tensor(x.get_tile_distribution()); + sweep_tile(x, [&](auto idx) { + constexpr auto i_idx = make_tuple(idx[number<0>{}]); + qx(idx) = type_convert(saturates{}(x[idx] / scale_tile[i_idx])); + }); + store_tile(qx_window, qx); + move_tile_window(x_window, {0, Block_N}); + move_tile_window(qx_window, {0, Block_N}); + } + } +}; } // namespace ck_tile