ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-13 18:51:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

store

Browse Source
This commit is contained in:
fsx950223
2025-04-10 02:33:51 +00:00
parent b000f23356
commit 33a5f34fe1
6 changed files with 245 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
example/ck_tile/36_quant/CMakeLists.txt
View File

@@ -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)

9
example/ck_tile/36_quant/example_dynamic_per_tensor_quant.cpp
View File

@@ -136,15 +136,16 @@ int main()
constexpr dim3 blocks = Kernel::BlockSize();
auto kargs = Kernel::MakeKargs(a);
ck_tile::launch_kernel(
{}, ck_tile::make_kernel<blocks.x, kBlockPerCu>(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<blocks.x, kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
std::cout<<"time: "<<time<<std::endl;
bool pass = true;
scale_buf.FromDevice(scale_host.data());
ck_tile::reference_per_tensor_quantization2d<XDataType, ScaleDataType, QXDataType>(
x_host, scale_host, qx_host_ref);
// std::cout<<scale_host(0)<<std::endl;
qx_buf.FromDevice(qx_host_dev.data());
auto [rtol, atol] = get_elimit<QXDataType>();

6
example/ck_tile/36_quant/example_static_per_tensor_quant.cpp
View File

@@ -138,9 +138,11 @@ int main()
constexpr dim3 blocks = Kernel::BlockSize();
auto kargs = Kernel::MakeKargs(a);
ck_tile::launch_kernel(
{}, ck_tile::make_kernel<blocks.x, kBlockPerCu>(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<blocks.x, kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
std::cout<<"time: "<<time<<std::endl;
bool pass = true;
ck_tile::reference_per_tensor_quantization2d<XDataType, ScaleDataType, QXDataType>(

20
include/ck_tile/host/reference/reference_rowwise_quantization2d.hpp
View File

@@ -50,4 +50,24 @@ CK_TILE_HOST void reference_per_tensor_quantization2d(const HostTensor<XDataType
x_m_n.mDesc.get_lengths()[0])(std::thread::hardware_concurrency());
}
template <typename XDataType, typename ScaleDataType, typename QXDataType>
CK_TILE_HOST void reference_per_token_quantization2d(const HostTensor<XDataType>& x_m_n,
const HostTensor<ScaleDataType>& scale_m,
HostTensor<QXDataType>& 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<QXDataType>(saturates<QXDataType>{}(v_qx));
}
};
make_ParallelTensorFunctor(f,
x_m_n.mDesc.get_lengths()[0])(std::thread::hardware_concurrency());
}
} // namespace ck_tile

126
include/ck_tile/ops/quant/kernel/quant_kernel.hpp
View File

@@ -116,4 +116,130 @@ struct PerTensorQuant
Pipeline{}(x_window, scale, kargs.n, qx_window, smem);
}
};
template <typename Pipeline_>
struct PerTokenQuant
{
using Pipeline = remove_cvref_t<Pipeline_>;
using Problem = typename Pipeline::Problem;
using XDataType = remove_cvref_t<typename Problem::XDataType>;
using ScaleDataType = remove_cvref_t<typename Problem::ScaleDataType>;
using ComputeDataType = remove_cvref_t<typename Problem::ComputeDataType>;
using QXDataType = remove_cvref_t<typename Problem::QXDataType>;
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 <typename T> struct t2s;
template <> struct t2s<float> { static constexpr const char * name = "fp32"; };
template <> struct t2s<ck_tile::fp16_t> { static constexpr const char * name = "fp16"; };
template <> struct t2s<ck_tile::bf16_t> { static constexpr const char * name = "bf16"; };
template <> struct t2s<ck_tile::fp8_t> { static constexpr const char * name = "fp8"; };
template <> struct t2s<ck_tile::bf8_t> { 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<XDataType>::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<address_space_enum::global>(
static_cast<const XDataType*>(kargs.p_x),
make_tuple(kargs.m, kargs.n),
make_tuple(kargs.x_stride, 1),
number<Vector_N>{},
number<1>{});
const auto tmp2_ = pad_tensor_view(
tmp_, make_tuple(number<Block_M>{}, number<Block_N>{}), sequence<kPadM, kPadN>{});
return make_tile_window(
tmp2_, make_tuple(number<Block_M>{}, number<Block_N>{}), {iM, 0});
}();
auto scale_window = [&]() {
auto tmp_ = make_naive_tensor_view<address_space_enum::global>(
static_cast<ScaleDataType*>(kargs.p_scale),
make_tuple(kargs.m),
make_tuple(1),
number<1>{},
number<1>{});
auto tmp2_ = pad_tensor_view(
tmp_, make_tuple(number<Block_M>{}), sequence<kPadM>{});
return make_tile_window(
tmp2_, make_tuple(number<Block_M>{}), {iM});
}();
auto qx_window = [&]() {
auto tmp_ = make_naive_tensor_view<address_space_enum::global>(
static_cast<QXDataType*>(kargs.p_qx),
make_tuple(kargs.m, kargs.n),
make_tuple(kargs.x_stride, 1),
number<Vector_N>{},
number<1>{});
auto tmp2_ = pad_tensor_view(
tmp_, make_tuple(number<Block_M>{}, number<Block_N>{}), sequence<kPadM, kPadN>{});
return make_tile_window(
tmp2_, make_tuple(number<Block_M>{}, number<Block_N>{}), {iM, 0});
}();
__shared__ char smem[GetSmemSize()];
Pipeline{}(x_window, scale_window, kargs.n, qx_window, smem);
}
};
} // namespace ck_tile

91
include/ck_tile/ops/quant/pipeline/quant_pipeline.hpp
View File

@@ -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 <typename Problem_, typename Policy_ = PerTensorQuantPipelineDefaultPolicy>
struct DynamicPerTokenQuantPipeline
{
using Problem = ck_tile::remove_cvref_t<Problem_>;
using Policy = ck_tile::remove_cvref_t<Policy_>;
using XDataType = ck_tile::remove_cvref_t<typename Problem::XDataType>;
using ScaleDataType = ck_tile::remove_cvref_t<typename Problem::ScaleDataType>;
using ComputeDataType = ck_tile::remove_cvref_t<typename Problem::ComputeDataType>;
using QXDataType = ck_tile::remove_cvref_t<typename Problem::QXDataType>;
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 <typename XWindow,
typename ScaleWindow,
typename QXWindow>
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<Problem>());
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<Problem>();
using XTensorType = decltype(cast_tile<ComputeDataType>(load_tile(x_window)));
auto absmax = block_reduce2d.template MakeYBlockTile<XTensorType>();
set_tile(absmax, reduce_absmax_func.GetIdentityValue<ComputeDataType>());
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<ComputeDataType, float> &&
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<ComputeDataType>(numeric<QXDataType>::max());
},
absmax);
store_tile(scale_window, cast_tile<ScaleDataType>(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<ComputeDataType>(load_tile(x_window));
auto qx = make_static_distributed_tensor<QXDataType>(x.get_tile_distribution());
sweep_tile(x, [&](auto idx) {
constexpr auto i_idx = make_tuple(idx[number<0>{}]);
qx(idx) = type_convert<QXDataType>(saturates<QXDataType>{}(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