smoothquant_pipeline_one_pass.hpp Source File

smoothquant_pipeline_one_pass.hpp Source File#

Composable Kernel: smoothquant_pipeline_one_pass.hpp Source File
// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
 
#pragma once
 
#include "ck_tile/core.hpp"
#include "ck_tile/ops/rmsnorm2d/pipeline/rmsnorm2d_fwd_pipeline_default_policy.hpp"
#include <string>
#include <type_traits>
 
namespace ck_tile {
 
template <typename Problem_, typename Policy_ = SmoothquantPipelineDefaultPolicy>
struct SmoothquantPipelineOnePass
{
    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 SmoothScaleDataType = ck_tile::remove_cvref_t<typename Problem::SmoothScaleDataType>;
    using ComputeDataType     = ck_tile::remove_cvref_t<typename Problem::ComputeDataType>;
    using QYDataType          = ck_tile::remove_cvref_t<typename Problem::QYDataType>;
    using YScaleDataType      = ck_tile::remove_cvref_t<typename Problem::YScaleDataType>;
 
    static constexpr bool kNeedCrossWarpSync = Problem::kNeedCrossWarpSync;
    static constexpr bool kPadM              = false; // TODO - BlockSmoothquantProblem::kPadM
    static constexpr bool kPadN              = Problem::kPadN;
    static constexpr bool UseMax3            = true; // TODO - Move to trait
 
    static constexpr const char* name = []() {
        if constexpr(kNeedCrossWarpSync)
            return "bpr_op"; // block per row
        else
            return "wpr_op"; // warp per row
    }();
 
    CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize()
    {
        return Policy::template GetSmemSize<Problem>();
    }
 
    template <typename XWindow,
              typename SmoothScaleWindow,
              typename QYWindow,
              typename YScaleWindow>
    CK_TILE_DEVICE auto operator()(const XWindow& x_window_,
                                   const SmoothScaleWindow& smscale_window_,
                                   YScaleWindow& yscale_window,
                                   QYWindow& qy_window,
                                   ck_tile::index_t,
                                   void* smem) const
    {
        auto x_window =
            make_tile_window(x_window_, Policy::template MakeXBlockTileDistribution<Problem>());
        auto smscale_window = make_tile_window(
            smscale_window_, Policy::template MakeSmoothScaleBlockTileDistribution<Problem>());
 
        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 reduce_max_func = ReduceOp::Max{};
 
        auto block_reduce2d      = Policy::template GetBlockReduce2d<Problem>();
        auto block_reduce2d_sync = Policy::template GetBlockReduce2dSync<Problem>();
        auto block_reduce2d_cross_warp_sync =
            Policy::template GetBlockReduce2dCrossWarpSync<Problem>();
 
        const auto x       = load_tile(x_window);
        const auto smscale = load_tile(smscale_window);
        auto y             = tile_elementwise_in(
            [&](const auto& a, const auto& b) {
                return type_convert<ComputeDataType>(a) * type_convert<ComputeDataType>(b);
            },
            x,
            smscale);
 
        // compute absmax, cross-lane->cross-warp
        auto absmax = [&]() {
            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)
            {
                return block_reduce2d(y,
                                      reduce_absmax_func.GetIdentityValue<ComputeDataType>(),
                                      reduce_absmax3_func,
                                      sequence<1, 2>{});
            }
            else
            {
                return block_reduce2d(
                    y, reduce_absmax_func.GetIdentityValue<ComputeDataType>(), reduce_absmax_func);
            }
        }();
        block_reduce2d_sync(absmax, reduce_max_func);
        block_reduce2d_cross_warp_sync(absmax, smem, reduce_max_func);
 
        // ex: yscale = absmax / 127 if int8
        auto yscale = tile_elementwise_in(
            [&](const auto& v_) {
                return v_ / type_convert<ComputeDataType>(numeric<QYDataType>::max());
            },
            absmax);
        store_tile(yscale_window, cast_tile<YScaleDataType>(yscale));
 
        // quantize y to qy
        auto qy = make_static_distributed_tensor<QYDataType>(y.get_tile_distribution());
        sweep_tile(qy, [&](auto idx) {
            constexpr auto i_idx = make_tuple(idx[number<0>{}]);
            auto qy_             = y[idx] / yscale[i_idx];
            qy(idx)              = type_convert<QYDataType>(saturates<QYDataType>{}(qy_));
        });
        store_tile(qy_window, qy);
    }
};
} // namespace ck_tile