| 1 | #pragma once |
| 2 | |
| 3 | #include <torch/headeronly/util/Exception.h> |
| 4 | |
| 5 | // clang-format off |
| 6 | // The cutlass include order matters (annoyingly) |
| 7 | #include "cutlass/cutlass.h" |
| 8 | |
| 9 | #include "cute/tensor.hpp" |
| 10 | #include "cutlass/tensor_ref.h" |
| 11 | #include "cutlass/epilogue/collective/default_epilogue.hpp" |
| 12 | #include "cutlass/epilogue/thread/linear_combination.h" |
| 13 | #include "cutlass/gemm/dispatch_policy.hpp" |
| 14 | #include "cutlass/gemm/collective/collective_builder.hpp" |
| 15 | #include "cutlass/epilogue/collective/collective_builder.hpp" |
| 16 | #include "cutlass/gemm/device/gemm_universal_adapter.h" |
| 17 | #include "cutlass/gemm/kernel/gemm_universal.hpp" |
| 18 | // clang-format on |
| 19 | |
| 20 | #include "cutlass_extensions/cute_utils.cuh" |
| 21 | #include "libtorch_stable/cutlass_extensions/vllm_numeric_conversion.cuh" |
| 22 | #include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" |
| 23 | #include "libtorch_stable/cutlass_extensions/torch_utils.hpp" |
| 24 | #include "machete_collective_builder.cuh" |
| 25 | #include "machete_prepacked_layout.cuh" |
| 26 | #include "machete_interleaving_utils.cuh" |
| 27 | |
| 28 | namespace machete { |
| 29 | |
| 30 | using namespace cute; |
| 31 | |
| 32 | // NOTE This kernel computes D = alpha * A * B + beta * C by computing |
| 33 | // D^t = alpha * B^t * A^t + beta * C^t, this is because the wgmma |
| 34 | // instructions only support sourcing from registers for the left-hand |
| 35 | // operand, we want to upconvert/decompress the quantized operand in |
| 36 | // register. Since the primary use case we want to support is Y = XW^t where |
| 37 | // W is quantized, in this situation or right-hand operand is quantized so |
| 38 | // we compute the transpose to move it to the left-hand side. |
| 39 | template <typename ElementA_, typename ElementB_, typename ElementD_, |
| 40 | typename AccumulatorT, typename GroupScaleT, typename GroupZeroT, |
| 41 | typename ChannelScaleT, typename TokenScaleT, class KernelSchedule, |
| 42 | typename ScheduleConfig> |
| 43 | struct MacheteKernelTemplate { |
| 44 | static constexpr bool with_C = false; // not ever used |
| 45 | static constexpr bool with_group_scales = !std::is_same_v<GroupScaleT, void>; |
| 46 | static constexpr bool with_group_zeropoints = |
| 47 | !std::is_same_v<GroupZeroT, void>; |
| 48 | static constexpr bool with_channel_scales = |
| 49 | !std::is_same_v<ChannelScaleT, void>; |
| 50 | static constexpr bool with_token_scales = !std::is_same_v<TokenScaleT, void>; |
| 51 | |
| 52 | using MmaType = ElementA_; |
| 53 | using ElementA = ElementA_; |
| 54 | using ElementB = ElementB_; |
| 55 | using ElementD = ElementD_; |
| 56 | using ElementC = cute::conditional_t<with_C, ElementD, void>; |
| 57 | using ElementAccumulator = AccumulatorT; |
| 58 | using ElementCompute = AccumulatorT; // For Epilogue |
| 59 | // Use dummy values when we don't have scales or zeropoints |
| 60 | using ElementZGroup = |
| 61 | cute::conditional_t<with_group_zeropoints, GroupZeroT, MmaType>; |
| 62 | using ElementSGroup = |
| 63 | cute::conditional_t<with_group_scales, GroupScaleT, MmaType>; |
| 64 | using ElementConvertGroup = |
| 65 | cute::conditional_t<with_group_scales, GroupScaleT, MmaType>; |
| 66 | using ElementSChannel = |
| 67 | cute::conditional_t<with_channel_scales, ChannelScaleT, AccumulatorT>; |
| 68 | using ElementSToken = |
| 69 | cute::conditional_t<with_token_scales, TokenScaleT, AccumulatorT>; |
| 70 | |
| 71 | using BTypeTuple = cute::conditional_t< |
| 72 | with_group_scales, |
| 73 | cute::conditional_t<with_group_zeropoints, |
| 74 | cute::tuple<ElementB, ElementSGroup, ElementZGroup>, |
| 75 | cute::tuple<ElementB, ElementSGroup>>, |
| 76 | ElementB>; |
| 77 | |
| 78 | using LayoutA = cutlass::layout::RowMajor; |
| 79 | using LayoutC = cutlass::layout::RowMajor; |
| 80 | using LayoutD = LayoutC; |
| 81 | using LayoutScale = cutlass::layout::RowMajor; |
| 82 | // not actually used since B has the prepacked layout, but required by cutlass |
| 83 | using _LayoutB = cutlass::layout::ColumnMajor; |
| 84 | |
| 85 | // Interface strides expected by create_arguments (will get transposed) |
| 86 | using StrideA = cutlass::detail::TagToStrideA_t<LayoutA>; |
| 87 | using StrideC = cutlass::detail::TagToStrideA_t<LayoutC>; |
| 88 | using StrideD = cutlass::detail::TagToStrideA_t<LayoutD>; |
| 89 | using StrideSGroup = cutlass::detail::TagToStrideA_t<LayoutScale>; |
| 90 | using StrideZGroup = StrideSGroup; |
| 91 | |
| 92 | using LayoutA_Transpose = |
| 93 | typename cutlass::layout::LayoutTranspose<LayoutA>::type; |
| 94 | using LayoutC_Transpose = |
| 95 | typename cutlass::layout::LayoutTranspose<LayoutC>::type; |
| 96 | using LayoutD_Transpose = |
| 97 | typename cutlass::layout::LayoutTranspose<LayoutD>::type; |
| 98 | |
| 99 | using ArchTag = cutlass::arch::Sm90; |
| 100 | using OperatorClass = cutlass::arch::OpClassTensorOp; |
| 101 | |
| 102 | using PrepackedLayoutB = |
| 103 | PrepackedLayoutBTemplate<ElementA_, ElementB_, ElementConvertGroup, |
| 104 | AccumulatorT, LayoutA_Transpose, KernelSchedule>; |
| 105 | |
| 106 | static int constexpr TileShapeK = |
| 107 | 128 * 8 / cutlass::sizeof_bits<MmaType>::value; |
| 108 | static int constexpr AlignmentA = 128 / cutlass::sizeof_bits_v<ElementA>; |
| 109 | static int constexpr AlignmentB = 128 / cutlass::sizeof_bits_v<ElementB>; |
| 110 | static int constexpr AlignmentC = |
| 111 | (with_C) ? 128 / cutlass::sizeof_bits_v<ElementC> : 0; |
| 112 | static int constexpr AlignmentD = 128 / cutlass::sizeof_bits_v<ElementD>; |
| 113 | |
| 114 | using TileShape = decltype(append(typename ScheduleConfig::TileShapeNM{}, |
| 115 | cute::Int<TileShapeK>{})); |
| 116 | using ClusterShape = typename ScheduleConfig::ClusterShape; |
| 117 | using EpilogueSchedule = typename ScheduleConfig::EpilogueSchedule; |
| 118 | using EpilogueTileType = typename ScheduleConfig::EpilogueTileType; |
| 119 | using TileScheduler = typename ScheduleConfig::TileScheduler; |
| 120 | |
| 121 | static_assert( |
| 122 | (!with_channel_scales && !with_token_scales) || |
| 123 | ((with_channel_scales && with_token_scales) && |
| 124 | std::is_same_v<ElementSChannel, ElementSToken>), |
| 125 | "Currently token and channel scales (if present) must be the same type"); |
| 126 | |
| 127 | // Currently only supports float scales |
| 128 | using ChTokScalesEpilogue = |
| 129 | typename vllm::c3x::ScaledEpilogue<ElementAccumulator, ElementD, |
| 130 | TileShape>; |
| 131 | static_assert((with_channel_scales || with_token_scales) || |
| 132 | (std::is_same_v<ElementSChannel, float> && |
| 133 | std::is_same_v<ElementSToken, float>), |
| 134 | "Currently token and channel scales (if present) must be float " |
| 135 | "(and if one is present the other must be too)"); |
| 136 | |
| 137 | using StoreEpilogueCompute = typename cutlass::epilogue::fusion::Sm90EVT< |
| 138 | cutlass::epilogue::fusion::Sm90AccFetch>; |
| 139 | |
| 140 | using EVTCompute = |
| 141 | std::conditional_t<with_channel_scales || with_token_scales, |
| 142 | typename ChTokScalesEpilogue::EVTCompute, |
| 143 | StoreEpilogueCompute>; |
| 144 | |
| 145 | // EVTCompute |
| 146 | using CollectiveEpilogue = |
| 147 | typename cutlass::epilogue::collective::CollectiveBuilder< |
| 148 | ArchTag, OperatorClass, TileShape, ClusterShape, EpilogueTileType, |
| 149 | ElementAccumulator, ElementSChannel, ElementC, LayoutC_Transpose, |
| 150 | AlignmentC, ElementD, LayoutD_Transpose, AlignmentD, EpilogueSchedule, |
| 151 | EVTCompute>::CollectiveOp; |
| 152 | |
| 153 | using CollectiveMainloop = |
| 154 | typename cutlass::gemm::collective::VLLMCollectiveBuilder< |
| 155 | cutlass::gemm::collective::MacheteKernelTag, ArchTag, OperatorClass, |
| 156 | BTypeTuple, PrepackedLayoutB, AlignmentB, ElementA, LayoutA_Transpose, |
| 157 | AlignmentA, ElementAccumulator, TileShape, ClusterShape, |
| 158 | cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>( |
| 159 | sizeof(typename CollectiveEpilogue::SharedStorage))>, |
| 160 | KernelSchedule>::CollectiveOp; |
| 161 | |
| 162 | using GemmKernel = cutlass::gemm::kernel::GemmUniversal< |
| 163 | Shape<int, int, int, int>, // Indicates ProblemShape |
| 164 | CollectiveMainloop, CollectiveEpilogue, TileScheduler>; |
| 165 | using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>; |
| 166 | |
| 167 | // stride_B is unused (since B is prepacked), but still required by cutlass |
| 168 | using _StrideB = cutlass::detail::TagToStrideB_t<_LayoutB>; |
| 169 | |
| 170 | using Arguments = typename Gemm::Arguments; |
| 171 | using MainloopArguments = typename GemmKernel::MainloopArguments; |
| 172 | using EpilogueArguments = typename GemmKernel::EpilogueArguments; |
| 173 | |
| 174 | static Arguments create_arguments( |
| 175 | cudaStream_t stream, |
| 176 | torch::stable::Tensor const& A, // MxK matrix |
| 177 | torch::stable::Tensor const& B, // KxN prepacked matrix |
| 178 | torch::stable::Tensor& D, // MxN matrix |
| 179 | std::optional<torch::stable::Tensor> const& |
| 180 | maybe_g_scales, // scale_KxN matrix |
| 181 | std::optional<torch::stable::Tensor> const& |
| 182 | maybe_g_zeros, // scale_KxN matrix |
| 183 | std::optional<int64_t> maybe_group_size, |
| 184 | std::optional<torch::stable::Tensor> const& |
| 185 | maybe_ch_scales, // len N vector |
| 186 | std::optional<torch::stable::Tensor> const& |
| 187 | maybe_tok_scales) // len M vector |
| 188 | { |
| 189 | static_assert(!with_group_zeropoints || with_group_scales); |
| 190 | |
| 191 | int M = A.size(0), N = B.size(1), K = A.size(1); |
| 192 | STD_TORCH_CHECK(D.size(0) == M && D.size(1) == N); |
| 193 | |
| 194 | auto layout_A = make_cute_layout<StrideA>(A, "A"); |
| 195 | auto layout_D = make_cute_layout<StrideD>(D, "D"); |
| 196 | auto layout_S_group = |
| 197 | maybe_make_cute_layout<StrideSGroup>(maybe_g_scales, "group_scales"); |
| 198 | auto layout_Z_group = |
| 199 | maybe_make_cute_layout<StrideZGroup>(maybe_g_zeros, "group_zeros"); |
| 200 | int64_t numel_S_channel = maybe_ch_scales ? maybe_ch_scales->numel() : 0; |
| 201 | int64_t numel_S_token = maybe_tok_scales ? maybe_tok_scales->numel() : 0; |
| 202 | |
| 203 | auto unwrap = [](auto const& t) { |
| 204 | return t ? t->const_data_ptr() : nullptr; |
| 205 | }; |
| 206 | auto A_ptr = static_cast<ElementA const*>(A.const_data_ptr()); |
| 207 | auto B_ptr = static_cast<ElementB const*>(B.const_data_ptr()); |
| 208 | auto D_ptr = static_cast<ElementD*>(D.mutable_data_ptr()); |
| 209 | auto S_group_ptr = |
| 210 | static_cast<ElementSGroup const*>(unwrap(maybe_g_scales)); |
| 211 | auto Z_group_ptr = static_cast<ElementZGroup const*>(unwrap(maybe_g_zeros)); |
| 212 | auto S_channel_ptr = |
| 213 | static_cast<ElementSChannel const*>(unwrap(maybe_ch_scales)); |
| 214 | auto S_token_ptr = |
| 215 | static_cast<ElementSToken const*>(unwrap(maybe_tok_scales)); |
| 216 | |
| 217 | int const group_size = |
| 218 | maybe_group_size == -1 ? K : maybe_group_size.value_or(K); |
| 219 | int const scale_k = (K + group_size - 1) / group_size; |
| 220 | |
| 221 | STD_TORCH_CHECK(size<0>(layout_A) == M && size<1>(layout_A) == K); |
| 222 | STD_TORCH_CHECK(size<0>(layout_D) == M && size<1>(layout_D) == N); |
| 223 | |
| 224 | if constexpr (with_group_scales) { |
| 225 | STD_TORCH_CHECK(S_group_ptr && layout_S_group); |
| 226 | STD_TORCH_CHECK((size<0>(*layout_S_group) == scale_k && |
| 227 | size<1>(*layout_S_group) == N)); |
| 228 | } else { |
| 229 | STD_TORCH_CHECK(!S_group_ptr, "Scales not supported"); |
| 230 | } |
| 231 | |
| 232 | if constexpr (with_group_zeropoints) { |
| 233 | STD_TORCH_CHECK(Z_group_ptr && layout_Z_group); |
| 234 | STD_TORCH_CHECK((size<0>(*layout_Z_group) == scale_k && |
| 235 | size<1>(*layout_Z_group) == N)); |
| 236 | STD_TORCH_CHECK(layout_S_group && *layout_Z_group == *layout_S_group, |
| 237 | "Scales and zeros must have the same layout"); |
| 238 | } else { |
| 239 | STD_TORCH_CHECK(!Z_group_ptr, "Zeropoints not supported"); |
| 240 | } |
| 241 | |
| 242 | if constexpr (with_channel_scales || with_token_scales) { |
| 243 | STD_TORCH_CHECK( |
| 244 | (maybe_ch_scales->numel() == N || maybe_ch_scales->numel() == 1) && |
| 245 | (maybe_tok_scales->numel() == M || maybe_tok_scales->numel() == 1)); |
| 246 | } |
| 247 | |
| 248 | // Transpose A and D |
| 249 | // A doesn't need to be transposed since cutlass expects a NxK matrix |
| 250 | // for B (which is At) |
| 251 | auto stride_At = layout_A.stride(); |
| 252 | auto stride_Dt = permute_layout<1, 0, 2>(layout_D).stride(); |
| 253 | |
| 254 | MainloopArguments mainloop_arguments{}; |
| 255 | // {Accum, C, C_layout, D, D} |
| 256 | EpilogueArguments epilogue_arguments{}; |
| 257 | |
| 258 | if constexpr (with_channel_scales || with_token_scales) { |
| 259 | epilogue_arguments = |
| 260 | EpilogueArguments{ChTokScalesEpilogue::prepare_args( |
| 261 | *maybe_ch_scales, *maybe_tok_scales), |
| 262 | nullptr, |
| 263 | {}, |
| 264 | D_ptr, |
| 265 | stride_Dt}; |
| 266 | } else { |
| 267 | epilogue_arguments = EpilogueArguments{{}, nullptr, {}, D_ptr, stride_Dt}; |
| 268 | } |
| 269 | |
| 270 | if constexpr (with_group_scales && with_group_zeropoints) { |
| 271 | auto stride_S_group = permute_layout<1, 0, 2>(*layout_S_group).stride(); |
| 272 | mainloop_arguments = MainloopArguments{ |
| 273 | B_ptr, _StrideB{}, A_ptr, stride_At, |
| 274 | S_group_ptr, stride_S_group, group_size, Z_group_ptr}; |
| 275 | } else if constexpr (with_group_scales) { |
| 276 | auto stride_S_group = permute_layout<1, 0, 2>(*layout_S_group).stride(); |
| 277 | mainloop_arguments = |
| 278 | MainloopArguments{B_ptr, _StrideB{}, A_ptr, stride_At, |
| 279 | S_group_ptr, stride_S_group, group_size}; |
| 280 | } else { |
| 281 | mainloop_arguments = |
| 282 | MainloopArguments{B_ptr, _StrideB{}, A_ptr, stride_At}; |
| 283 | } |
| 284 | |
| 285 | return Arguments{cutlass::gemm::GemmUniversalMode::kGemm, |
| 286 | {N, M, K, 1}, |
| 287 | mainloop_arguments, |
| 288 | epilogue_arguments}; |
| 289 | }; |
| 290 | |
| 291 | static size_t get_workspace_size(Arguments const& args) { |
| 292 | return Gemm::get_workspace_size(args); |
| 293 | } |
| 294 | |
| 295 | static bool can_implement(Arguments const& args) { |
| 296 | return Gemm::can_implement(args) == cutlass::Status::kSuccess; |
| 297 | } |
| 298 | |
| 299 | static void run(Arguments const& args, void* workspace, cudaStream_t stream) { |
| 300 | Gemm gemm_op; |
| 301 | |
| 302 | cutlass::Status status = gemm_op.initialize(args, workspace, stream); |
| 303 | STD_TORCH_CHECK(status == cutlass::Status::kSuccess, |
| 304 | "Machete kernel failed to initialize workspace"); |
| 305 | |
| 306 | status = gemm_op.run(stream); |
| 307 | STD_TORCH_CHECK(status == cutlass::Status::kSuccess, |
| 308 | "Machete kernel failed"); |
| 309 | } |
| 310 | }; |
| 311 | |
| 312 | }; // namespace machete |