| 1 | // |
| 2 | // Based off of: |
| 3 | // cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp |
| 4 | // Specifically: |
| 5 | // https://github.com/NVIDIA/cutlass/tree/06b21349bcf6ddf6a1686a47a137ad1446579db9/include/cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp |
| 6 | // Referred to as upstream from in the comments |
| 7 | // |
| 8 | // The main optimization machete implements compared to upstream is to prepack |
| 9 | // the weight matrix to more closely match the shape of the wgmma instructions |
| 10 | // allowing for wider (ideally 128bit) shared memory loads. For subbyte types |
| 11 | // this is done by packing values from multiple wgmma loads (for a single |
| 12 | // thread) into a single 128bit load. This is very similar to layout used in |
| 13 | // Marlin, although specific to the wgmma instructions. |
| 14 | // |
| 15 | // Since the wgmma instructions only support sourcing from registers for the A |
| 16 | // operand, and we want to upconvert/decompress the weight values/elements |
| 17 | // before feeding them into the tensor cores in registers, we need the weight |
| 18 | // matrix to be A. To achieve this we compute the transpose of Y = XW^t as |
| 19 | // Y^t = W^tX^t. This is mostly done outside of this file in |
| 20 | // csrc/quantization/machete/machete_mm_kernel.cuh, but this why A is the |
| 21 | // quantized/narrow type and has the prepacked layout despite the API being: |
| 22 | // B_prepacked = machete_prepack_B(B) |
| 23 | // Y = machete_mm(A, B_prepacked) |
| 24 | // |
| 25 | #pragma once |
| 26 | |
| 27 | // clang-format off |
| 28 | #include "cutlass/cutlass.h" |
| 29 | #include "cutlass/numeric_conversion.h" |
| 30 | #include "cute/arch/cluster_sm90.hpp" |
| 31 | #include "cute/arch/copy_sm90.hpp" |
| 32 | #include "cutlass/gemm/gemm.h" |
| 33 | #include "cutlass/detail/dependent_false.hpp" |
| 34 | #include "cutlass/gemm/dispatch_policy.hpp" |
| 35 | #include "cutlass/detail/layout.hpp" |
| 36 | |
| 37 | #include "cute/algorithm/functional.hpp" |
| 38 | #include "cute/atom/mma_atom.hpp" |
| 39 | #include "cute/atom/copy_traits_sm90_tma.hpp" |
| 40 | #include "cute/algorithm/gemm.hpp" |
| 41 | #include "cute/numeric/arithmetic_tuple.hpp" |
| 42 | #include "cutlass/pipeline/pipeline.hpp" |
| 43 | #include "cutlass/transform/collective/sm90_wgmma_transpose.hpp" |
| 44 | #include "cutlass/trace.h" |
| 45 | |
| 46 | #include "cutlass/detail/collective.hpp" |
| 47 | // clang-format on |
| 48 | |
| 49 | #include "cutlass_extensions/cute_utils.cuh" |
| 50 | |
| 51 | namespace machete { |
| 52 | |
| 53 | using namespace cute; |
| 54 | using namespace cutlass; |
| 55 | using namespace cutlass::gemm; |
| 56 | using namespace cutlass::gemm::collective; |
| 57 | using namespace cutlass::gemm::collective::detail; |
| 58 | |
| 59 | template <class ElementATuple_, class GmemLayoutA, int AlignmentA, |
| 60 | class ElementB_, class GmemLayoutB, int AlignmentB, |
| 61 | class ElementAccumulator_, class TileShape_MNK, |
| 62 | class ClusterShape_MNK, class StageCountType, |
| 63 | class KernelScheduleType> |
| 64 | struct MacheteCollectiveMma { |
| 65 | using Schedule = KernelScheduleType; |
| 66 | static_assert( |
| 67 | cute::is_same_v<Schedule, KernelTmaWarpSpecialized> || |
| 68 | cute::is_same_v<Schedule, KernelTmaWarpSpecialized> || |
| 69 | cute::is_same_v<Schedule, KernelTmaWarpSpecializedPingpong> || |
| 70 | cute::is_same_v<Schedule, KernelTmaWarpSpecializedPingpong> || |
| 71 | cute::is_same_v<Schedule, KernelTmaWarpSpecializedCooperative> || |
| 72 | cute::is_same_v<Schedule, KernelTmaWarpSpecializedCooperative>, |
| 73 | "KernelSchedule must be one of the warp specialized policies"); |
| 74 | |
| 75 | public: |
| 76 | static constexpr bool ALayoutIsPrepacked = true; |
| 77 | |
| 78 | // Prepacked block shape (N is M in the transposed problem) |
| 79 | using PPBlockShape_MK = typename GmemLayoutA::PPBlockShape_NK; |
| 80 | // Prepacked blocks per dim for a single MMA tile |
| 81 | using PPBlocksPerTile_MK = decltype(make_shape( |
| 82 | size<0>(TileShape_MNK{}) / size<0>(PPBlockShape_MK{}), |
| 83 | size<2>(TileShape_MNK{}) / size<1>(PPBlockShape_MK{}))); |
| 84 | |
| 85 | using IlvdBlkLayout = typename GmemLayoutA::IlvdBlkLayout; |
| 86 | |
| 87 | static_assert(size<0>(TileShape_MNK{}) % size<0>(PPBlockShape_MK{}) == 0, |
| 88 | "M in PPBlockShape_MK must evenly divide M TileShape_MNK"); |
| 89 | static_assert(size<2>(TileShape_MNK{}) % size<1>(PPBlockShape_MK{}) == 0, |
| 90 | "K in PPBlockShape_MK must evenly divide K TileShape_MNK"); |
| 91 | |
| 92 | using ArchTag = arch::Sm90; |
| 93 | using TileShape = TileShape_MNK; |
| 94 | using ClusterShape = ClusterShape_MNK; |
| 95 | using ElementA = deduce_mixed_width_dtype_t<0, ElementATuple_>; |
| 96 | using StrideA = TagToStrideA_t<layout::RowMajor>; |
| 97 | using ElementB = ElementB_; |
| 98 | using StrideB = TagToStrideB_t<GmemLayoutB>; |
| 99 | using ElementAccumulator = ElementAccumulator_; |
| 100 | using ElementMma = ElementB; |
| 101 | using ElementATuple = |
| 102 | cute::conditional_t<!cute::is_tuple<ElementATuple_>::value, |
| 103 | cute::tuple<ElementA>, ElementATuple_>; |
| 104 | |
| 105 | static constexpr cute::GMMA::Major GmmaMajorA = |
| 106 | gmma_rs_tag_to_major_A<layout::RowMajor>(); |
| 107 | static constexpr cute::GMMA::Major GmmaMajorB = |
| 108 | gmma_rs_tag_to_major_B<GmemLayoutB>(); |
| 109 | |
| 110 | // For coop schedules we have two warp groups cooperatively issuing wgmma |
| 111 | // instructions so we use 2 atoms along the M dim (one for each warpgroup) |
| 112 | using AtomLayoutMNK = cute::conditional_t< |
| 113 | cute::is_same_v<KernelScheduleType, KernelTmaWarpSpecializedCooperative>, |
| 114 | Layout<Shape<_2, _1, _1>>, Layout<Shape<_1, _1, _1>>>; |
| 115 | |
| 116 | using TiledMma = decltype(cute::make_tiled_mma( |
| 117 | cute::GMMA::rs_op_selector<ElementMma, ElementMma, ElementAccumulator, |
| 118 | TileShape_MNK, GMMA::Major::K, GmmaMajorB>(), |
| 119 | AtomLayoutMNK{})); |
| 120 | |
| 121 | private: |
| 122 | // |
| 123 | // the setup section (until "section setup end") contains a combination of |
| 124 | // modified code from (used as a starting point): |
| 125 | // `cutlass/gemm/collective/builders/sm90_gmma_builder.inl` |
| 126 | // `cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp` |
| 127 | // (upstream) |
| 128 | // |
| 129 | // however in-order to simplify the code we combine a lot of the logic from |
| 130 | // `CollectiveMma` and `CollectiveBuilder` into this class, this also makes |
| 131 | // sense given that we have flexibility on layouts here. We also simplify the |
| 132 | // code by only supporting scales and zeros for A (in the transposed problem, |
| 133 | // B from an API perspective), also since we force A to be the narrow type |
| 134 | // (i.e. the type to be upconverted) we can remove all the `SwapAB` logic in |
| 135 | // the upstream also simplifying the code. This section includes new logic |
| 136 | // (compared ustream) for handling the prepacked-A layouts (in the transposed |
| 137 | // problem, B from an API perspective) |
| 138 | // |
| 139 | using ElementScale = deduce_mixed_width_dtype_t<1, ElementATuple_>; |
| 140 | using ElementZero = deduce_mixed_width_dtype_t<2, ElementATuple_>; |
| 141 | |
| 142 | static constexpr bool IsANarrow = cutlass::sizeof_bits<ElementA>::value < |
| 143 | cutlass::sizeof_bits<ElementB>::value; |
| 144 | static_assert(IsANarrow, |
| 145 | "A must be the narrow one since its the one that flows through " |
| 146 | "registers."); |
| 147 | |
| 148 | public: |
| 149 | static constexpr int PipelineStages = |
| 150 | compute_stage_count_or_override_single_affine_transformed_input< |
| 151 | sm90_smem_capacity_bytes, ElementA, ElementB, ElementScale, |
| 152 | ElementZero, TileShape_MNK>(StageCountType{}); |
| 153 | |
| 154 | struct DispatchPolicy { |
| 155 | constexpr static int Stages = PipelineStages; |
| 156 | using ClusterShape = ClusterShape_MNK; |
| 157 | using ArchTag = arch::Sm90; |
| 158 | using Schedule = KernelScheduleType; |
| 159 | }; |
| 160 | |
| 161 | using GmemTiledCopyA = |
| 162 | decltype(sm90_cluster_shape_to_tma_atom(shape<1>(ClusterShape_MNK{}))); |
| 163 | using GmemTiledCopyB = |
| 164 | decltype(sm90_cluster_shape_to_tma_atom(shape<0>(ClusterShape_MNK{}))); |
| 165 | |
| 166 | // ((T, V), (BlocksM, BlocksK), pipe) -> offset |
| 167 | using SmemLayoutA = decltype(GmemLayoutA::TVbNbKL_to_offset( |
| 168 | make_shape(size<0>(TileShape_MNK{}), size<2>(TileShape_MNK{}), |
| 169 | Int<DispatchPolicy::Stages>{}))); |
| 170 | |
| 171 | using SmemLayoutACopy = decltype(GmemLayoutA::TVbNbKL_to_offset_copy( |
| 172 | make_shape(size<0>(TileShape_MNK{}), size<2>(TileShape_MNK{}), |
| 173 | Int<DispatchPolicy::Stages>{}))); |
| 174 | |
| 175 | using SmemLayoutAtomARowMajor = |
| 176 | decltype(rs_smem_selector<GmmaMajorA, ElementA, |
| 177 | decltype(cute::get<0>(TileShape_MNK{})), |
| 178 | decltype(cute::get<2>(TileShape_MNK{}))>()); |
| 179 | |
| 180 | using SmemLayoutAtomScale = Layout< |
| 181 | Shape<decltype(cute::shape<0>(SmemLayoutAtomARowMajor{})), cute::Int<1>>>; |
| 182 | |
| 183 | using SmemLayoutAtomB = |
| 184 | decltype(rs_smem_selector<GmmaMajorB, ElementB, |
| 185 | decltype(cute::get<1>(TileShape_MNK{})), |
| 186 | decltype(cute::get<2>(TileShape_MNK{}))>()); |
| 187 | |
| 188 | using SmemCopyAtomA = Copy_Atom<cute::DefaultCopy, ElementA>; |
| 189 | using SmemCopyAtomB = void; |
| 190 | |
| 191 | // |
| 192 | // Validity checks |
| 193 | // |
| 194 | static_assert(is_static<TileShape_MNK>::value); |
| 195 | static_assert(is_static<ClusterShape_MNK>::value); |
| 196 | static_assert(is_aligned<ElementA, AlignmentA, ElementB, AlignmentB, |
| 197 | tma_alignment_bytes>(), |
| 198 | "Should meet TMA alignment requirement\n"); |
| 199 | #ifndef CUTLASS_SM90_COLLECTIVE_BUILDER_SUPPORTED |
| 200 | static_assert(cutlass::detail::dependent_false<ElementA>, |
| 201 | "Unsupported Toolkit for SM90 Collective Builder\n"); |
| 202 | #endif |
| 203 | |
| 204 | private: |
| 205 | enum class ConversionMode { |
| 206 | DirectConvert, |
| 207 | ConvertAndScale, |
| 208 | ConvertAndScaleWithZero |
| 209 | }; |
| 210 | |
| 211 | public: |
| 212 | // |
| 213 | // Type Aliases |
| 214 | // |
| 215 | using KernelSchedule = KernelScheduleType; |
| 216 | |
| 217 | // For cases where we can't have a void type, we can use this to allow the |
| 218 | // code to compile when the scale / zero is void. |
| 219 | using NonVoidElementScale = |
| 220 | cute::conditional_t<cute::is_void_v<ElementScale>, float, ElementScale>; |
| 221 | using NonVoidElementZero = |
| 222 | cute::conditional_t<cute::is_void_v<ElementZero>, float, ElementZero>; |
| 223 | |
| 224 | // These are always MN major |
| 225 | using StrideScale = cute::Stride<cute::Int<1>, int64_t, int64_t>; |
| 226 | // For cases where we can't have a void scale, we can use this to allow the |
| 227 | // code to compile when the scale is void. |
| 228 | using NonVoidStrideScale = |
| 229 | cute::conditional_t<cute::is_void_v<StrideScale>, |
| 230 | cute::Stride<_1, int64_t, int64_t>, StrideScale>; |
| 231 | |
| 232 | static_assert((cutlass::gemm::detail::is_k_major<StrideA>()), |
| 233 | "The transformed matrix (A) must be K-major."); |
| 234 | |
| 235 | static_assert((sizeof(ElementB) == 2) || |
| 236 | (cutlass::gemm::detail::is_k_major<StrideA>() && |
| 237 | cutlass::gemm::detail::is_k_major<StrideB>()), |
| 238 | "The unscaled element (matrix B) must be 2 bytes OR both " |
| 239 | "inputs must be K-major"); |
| 240 | |
| 241 | static_assert(cutlass::gemm::detail::is_mn_major<NonVoidStrideScale>(), |
| 242 | "Scale must be MN major [Col Major if A is scaled, Row Major " |
| 243 | "if B is scaled]."); |
| 244 | |
| 245 | static_assert(std::is_same_v<typename TiledMma::ValTypeC, ElementAccumulator>, |
| 246 | "TiledMma::ValTypeC must be the same as ElementAccumulator."); |
| 247 | |
| 248 | using GmemTiledCopyScale = cute::SM90_TMA_LOAD; |
| 249 | |
| 250 | using SmemCopyAtomScale = Copy_Atom<cute::DefaultCopy, NonVoidElementScale>; |
| 251 | |
| 252 | // TMA converts f32 input to tf32 when copying from GMEM to SMEM |
| 253 | // For all other types, cast to size equivalent uint type to avoid any |
| 254 | // rounding by TMA. |
| 255 | static constexpr bool ConvertF32toTF32A = cute::is_same_v<float, ElementA>; |
| 256 | static constexpr bool ConvertF32toTF32B = cute::is_same_v<float, ElementB>; |
| 257 | using InternalElementA = |
| 258 | cute::conditional_t<ConvertF32toTF32A, tfloat32_t, |
| 259 | uint_bit_t<sizeof_bits_v<ElementA>>>; |
| 260 | using InternalElementB = |
| 261 | cute::conditional_t<ConvertF32toTF32B, tfloat32_t, |
| 262 | uint_bit_t<sizeof_bits_v<ElementB>>>; |
| 263 | |
| 264 | using TransformA = cute::identity; |
| 265 | using TransformB = cute::identity; |
| 266 | |
| 267 | static constexpr int IsSubbyteA = cute::sizeof_bits_v<InternalElementA> < 8; |
| 268 | using TmaElementA = |
| 269 | cute::conditional_t<IsSubbyteA, uint8_t, InternalElementA>; |
| 270 | |
| 271 | using MainloopPipeline = cutlass::PipelineTmaAsync<DispatchPolicy::Stages>; |
| 272 | using PipelineState = cutlass::PipelineState<DispatchPolicy::Stages>; |
| 273 | |
| 274 | using PipelineParams = typename MainloopPipeline::Params; |
| 275 | |
| 276 | // One threads per CTA are producers (1 for operand tile) |
| 277 | static constexpr int NumProducerThreadEvents = 1; |
| 278 | |
| 279 | using ScaleTileShape = decltype(make_shape(shape<0>(TileShape{}), |
| 280 | shape<1>(SmemLayoutAtomScale{}))); |
| 281 | |
| 282 | static_assert(cute::rank(SmemLayoutAtomB{}) == 2, |
| 283 | "SmemLayoutAtom must be rank 2 (M/N, K)"); |
| 284 | static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0, |
| 285 | "SmemLayoutAtom must evenly divide tile shape."); |
| 286 | static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0, |
| 287 | "SmemLayoutAtom must evenly divide tile shape."); |
| 288 | |
| 289 | static_assert(rank(SmemLayoutAtomScale{}) == 2, |
| 290 | "SmemLayoutAtomScale must be rank 2"); |
| 291 | static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomScale{})) == 0, |
| 292 | "SmemLayoutAtomScale must equal the tile shape."); |
| 293 | static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomScale{})) == 0, |
| 294 | "SmemLayoutAtomScale must evenly divide tile k shape."); |
| 295 | |
| 296 | // Tile along modes in a way that maximizes the TMA box size |
| 297 | using SmemLayoutB = decltype(tile_to_shape( |
| 298 | SmemLayoutAtomB{}, |
| 299 | make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), |
| 300 | Int<DispatchPolicy::Stages>{}), |
| 301 | conditional_t<::cutlass::gemm::detail::is_major<0, StrideB>(), |
| 302 | Step<_2, _1, _3>, Step<_1, _2, _3>>{})); |
| 303 | |
| 304 | // It is assumed that the scales and zero-points share the same smem layout |
| 305 | using SmemLayoutScale = decltype(tile_to_shape( |
| 306 | SmemLayoutAtomScale{}, |
| 307 | make_shape(shape<0>(ScaleTileShape{}), shape<1>(ScaleTileShape{}), |
| 308 | Int<PipelineStages>{}))); |
| 309 | |
| 310 | // If A mn-layout and B mn-layout, transposing B matrix since WGMMA is k-major |
| 311 | // only (e.g. tf32, fp32, fp8, int8). |
| 312 | static constexpr bool IsLayoutAmnBmn = |
| 313 | cute::is_same_v<gemm::detail::StrideToLayoutTagA_t<StrideA>, |
| 314 | layout::ColumnMajor> && |
| 315 | cute::is_same_v<gemm::detail::StrideToLayoutTagB_t<StrideB>, |
| 316 | layout::RowMajor>; |
| 317 | |
| 318 | static_assert(DispatchPolicy::Stages >= 2, |
| 319 | "Specialization requires Stages set to value 2 or more."); |
| 320 | static_assert(not cute::is_base_of<cute::GMMA::DescriptorIterator, |
| 321 | typename TiledMma::FrgTypeA>::value && |
| 322 | cute::is_base_of<cute::GMMA::DescriptorIterator, |
| 323 | typename TiledMma::FrgTypeB>::value, |
| 324 | "MMA atom must source A from rmem and B operand from smem_desc " |
| 325 | "for this mainloop."); |
| 326 | static_assert(cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD> || |
| 327 | cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD_MULTICAST>, |
| 328 | "GmemTiledCopy - invalid SM90 TMA copy atom specified."); |
| 329 | static_assert(cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD> || |
| 330 | cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD_MULTICAST>, |
| 331 | "GmemTiledCopy - invalid SM90 TMA copy atom specified."); |
| 332 | |
| 333 | using GmmaSmemLayoutB = decltype(tile_to_shape( |
| 334 | SmemLayoutAtomB{}, |
| 335 | make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), |
| 336 | Int<DispatchPolicy::Stages>{}), |
| 337 | conditional_t<::cutlass::gemm::detail::is_major<0, StrideB>(), |
| 338 | Step<_2, _1, _3>, Step<_1, _2, _3>>{})); |
| 339 | |
| 340 | // These two restrictions are related, so we place the assertions together. |
| 341 | // To relax them, we need to handle loading more than 1 row of scales for |
| 342 | // every main loop iteration. We must also handle updating the pipeline |
| 343 | // transaction bytes on the fly. NOTE: Deleting this assertion without |
| 344 | // required changes will cause the code to hang. |
| 345 | static_assert(size<1>(SmemLayoutAtomScale{}) == 1, |
| 346 | "size<1>(SmemLayoutAtomScale) must be 1."); |
| 347 | |
| 348 | private: |
| 349 | static constexpr ConversionMode get_conversion_mode() { |
| 350 | if constexpr (cute::is_void_v<ElementScale>) { |
| 351 | return ConversionMode::DirectConvert; |
| 352 | } else if constexpr (cute::is_void_v<ElementZero>) { |
| 353 | return ConversionMode::ConvertAndScale; |
| 354 | } else { |
| 355 | return ConversionMode::ConvertAndScaleWithZero; |
| 356 | } |
| 357 | } |
| 358 | |
| 359 | static constexpr ConversionMode KernelConversionMode = get_conversion_mode(); |
| 360 | static constexpr bool ModeHasScales = |
| 361 | KernelConversionMode == ConversionMode::ConvertAndScale || |
| 362 | KernelConversionMode == ConversionMode::ConvertAndScaleWithZero; |
| 363 | |
| 364 | // Same as upstream, should be kept the same when possible |
| 365 | static constexpr auto elements_per_smem_scale() { |
| 366 | if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { |
| 367 | return 0; |
| 368 | } else if constexpr (ModeHasScales) { |
| 369 | return cute::cosize_v<SmemLayoutScale>; |
| 370 | } else { |
| 371 | static_assert(cutlass::detail::dependent_false<KernelSchedule>, |
| 372 | "Type not handled in scale smem allocation."); |
| 373 | } |
| 374 | } |
| 375 | |
| 376 | // Same as upstream, should be kept the same when possible |
| 377 | static constexpr auto elements_per_smem_zero() { |
| 378 | if constexpr (KernelConversionMode == ConversionMode::DirectConvert || |
| 379 | KernelConversionMode == ConversionMode::ConvertAndScale) { |
| 380 | return 0; |
| 381 | } else if constexpr (KernelConversionMode == |
| 382 | ConversionMode::ConvertAndScaleWithZero) { |
| 383 | return cute::cosize_v<SmemLayoutScale>; |
| 384 | } else { |
| 385 | static_assert(cutlass::detail::dependent_false<KernelSchedule>, |
| 386 | "Type not handled in scale smem allocation."); |
| 387 | } |
| 388 | } |
| 389 | |
| 390 | // Same as upstream, should be kept the same when possible, not formatte for |
| 391 | // easier comparison |
| 392 | // clang-format off |
| 393 | // These methods use some the public members of the class. For that reason, we define them after the public section. |
| 394 | static constexpr uint32_t |
| 395 | compute_tma_transaction_bytes_mk() { |
| 396 | constexpr uint32_t baseline_bytes = cutlass::bits_to_bytes(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast<uint32_t>(cute::sizeof_bits_v<InternalElementA>)); |
| 397 | |
| 398 | if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { |
| 399 | return baseline_bytes; |
| 400 | } |
| 401 | else if constexpr (ModeHasScales) { |
| 402 | constexpr uint32_t scale_tx_bytes = cutlass::bits_to_bytes(size<0>(SmemLayoutScale{}) * size<1>(SmemLayoutScale{}) * static_cast<uint32_t>(cute::sizeof_bits_v<ElementScale>)); |
| 403 | static_assert(scale_tx_bytes % 128 == 0, "Each scale stage must be 128B aligned."); // required by TMA |
| 404 | if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { |
| 405 | return baseline_bytes + scale_tx_bytes; |
| 406 | } |
| 407 | else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { |
| 408 | // Scale and zero share smem layout |
| 409 | constexpr uint32_t zero_tx_bytes = cutlass::bits_to_bytes(size<0>(SmemLayoutScale{}) * size<1>(SmemLayoutScale{}) * static_cast<uint32_t>(cute::sizeof_bits_v<ElementZero>)); |
| 410 | static_assert(zero_tx_bytes % 128 == 0, "Each zero stage must be 128B aligned."); // required by TMA |
| 411 | return baseline_bytes + scale_tx_bytes + zero_tx_bytes; |
| 412 | } |
| 413 | else { |
| 414 | static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Type not handled in tma transaction bytes computation."); |
| 415 | } |
| 416 | } |
| 417 | else { |
| 418 | static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Type not handled in tma transaction bytes computation."); |
| 419 | } |
| 420 | } |
| 421 | |
| 422 | static constexpr uint32_t |
| 423 | compute_tma_transaction_bytes_nk() { |
| 424 | return cutlass::bits_to_bytes(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast<uint32_t>(cute::sizeof_bits_v<InternalElementB>)); |
| 425 | } |
| 426 | // clang-format on |
| 427 | |
| 428 | // ((athrid, val), (BlocksM, BlockK), L) -> (storage_idx) |
| 429 | using PrepackedStrideA = decltype(stride(GmemLayoutA::TVbNbKL_to_offset_copy( |
| 430 | make_shape(int32_t(0), int32_t(0), int32_t(0))))); |
| 431 | |
| 432 | using ATensor = decltype(make_tensor( |
| 433 | get_logical_ptr(static_cast<InternalElementA const*>(nullptr)), |
| 434 | shape(GmemLayoutA::TVbNbKL_to_offset_copy( |
| 435 | make_shape(int32_t(0), int32_t(0), int32_t(0)))), |
| 436 | PrepackedStrideA{})); |
| 437 | |
| 438 | using BTensor = decltype(make_tensor( |
| 439 | get_logical_ptr(static_cast<InternalElementB const*>(nullptr)), |
| 440 | repeat_like(StrideB{}, int32_t(0)), StrideB{})); |
| 441 | using ScaleTensor = decltype(make_tensor( |
| 442 | get_logical_ptr(static_cast<NonVoidElementScale const*>(nullptr)), |
| 443 | repeat_like(NonVoidStrideScale{}, int32_t(0)), NonVoidStrideScale{})); |
| 444 | |
| 445 | using ZeroTensor = decltype(make_tensor( |
| 446 | get_logical_ptr(static_cast<NonVoidElementZero const*>(nullptr)), |
| 447 | repeat_like(NonVoidStrideScale{}, int32_t(0)), NonVoidStrideScale{})); |
| 448 | |
| 449 | static constexpr auto make_tma_copy_A(ATensor tensor_a = ATensor{}) { |
| 450 | return make_tma_copy<TmaElementA>( |
| 451 | GmemTiledCopyA{}, tensor_a, SmemLayoutACopy{}(_, _, cute::Int<0>{}), |
| 452 | shape(SmemLayoutACopy{}(_, _, cute::Int<0>{})), |
| 453 | size<1>(ClusterShape{})); // mcast along N mode for this M load, if any |
| 454 | } |
| 455 | |
| 456 | static constexpr auto make_tma_copy_scale( |
| 457 | ScaleTensor tensor_scale = ScaleTensor{}) { |
| 458 | return make_tma_copy(GmemTiledCopyScale{}, tensor_scale, |
| 459 | SmemLayoutScale{}(_, _, cute::Int<0>{}), |
| 460 | ScaleTileShape{}, |
| 461 | _1{}); // mcast along N mode for this M load, if any |
| 462 | } |
| 463 | |
| 464 | static constexpr auto make_tma_copy_zero( |
| 465 | ZeroTensor tensor_zero = ZeroTensor{}) { |
| 466 | return make_tma_copy(GmemTiledCopyScale{}, tensor_zero, |
| 467 | SmemLayoutScale{}(_, _, cute::Int<0>{}), |
| 468 | ScaleTileShape{}, |
| 469 | _1{}); // mcast along N mode for this M load, if any |
| 470 | } |
| 471 | |
| 472 | static constexpr auto make_tma_copy_B(BTensor tensor_b = BTensor{}) { |
| 473 | return make_tma_copy( |
| 474 | GmemTiledCopyB{}, tensor_b, SmemLayoutB{}(_, _, cute::Int<0>{}), |
| 475 | make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})), |
| 476 | size<0>(ClusterShape{})); // mcast along M mode for this N load, if any |
| 477 | } |
| 478 | |
| 479 | public: |
| 480 | // Same as upstream, should be kept the same when possible, not formatted for |
| 481 | // easier comparison |
| 482 | // with `RealInternalElementA` -> `ElementA` since we support `SwapAB` logic |
| 483 | // clang-format off |
| 484 | static constexpr size_t SmemAlignmentA = cutlass::detail::alignment_for_swizzle(SmemLayoutA{}); |
| 485 | |
| 486 | static constexpr size_t SmemAlignmentB = cutlass::detail::alignment_for_swizzle(SmemLayoutB{}); |
| 487 | |
| 488 | // Just pick the max alignment of A and B since it is required to be at least 128B |
| 489 | static constexpr size_t SmemAlignmentScale = cute::max(SmemAlignmentA, SmemAlignmentB); |
| 490 | |
| 491 | static_assert(SmemAlignmentA >= 128 and SmemAlignmentB >= 128, "Require at least 128B alignment"); |
| 492 | |
| 493 | struct SharedStorage |
| 494 | { |
| 495 | static constexpr int scale_elements = elements_per_smem_scale(); |
| 496 | static constexpr int zero_elements = elements_per_smem_zero(); |
| 497 | struct TensorStorage : cute::aligned_struct<cute::max(SmemAlignmentA, SmemAlignmentB)> { |
| 498 | cute::ArrayEngine<ElementA, cute::cosize_v<SmemLayoutA>> smem_A; |
| 499 | cute::ArrayEngine<typename TiledMma::ValTypeB, cute::cosize_v<SmemLayoutB>> smem_B; |
| 500 | cute::ArrayEngine<NonVoidElementScale, scale_elements> smem_scale; |
| 501 | cute::ArrayEngine<NonVoidElementZero, zero_elements> smem_zero; |
| 502 | } tensors; |
| 503 | |
| 504 | using PipelineStorage = typename MainloopPipeline::SharedStorage; |
| 505 | PipelineStorage pipeline; |
| 506 | }; |
| 507 | using TensorStorage = typename SharedStorage::TensorStorage; |
| 508 | using PipelineStorage = typename SharedStorage::PipelineStorage; |
| 509 | |
| 510 | // Host side kernel arguments |
| 511 | struct Arguments { |
| 512 | ElementA const* ptr_A = nullptr; |
| 513 | StrideA dA{}; |
| 514 | ElementB const* ptr_B = nullptr; |
| 515 | StrideB dB{}; |
| 516 | ElementScale const* ptr_S = nullptr; |
| 517 | NonVoidStrideScale dS{}; |
| 518 | int group_size = 0; |
| 519 | ElementZero const* ptr_Z = nullptr; |
| 520 | uint32_t mma_promotion_interval = 4; |
| 521 | }; |
| 522 | // clang-format on |
| 523 | |
| 524 | // |
| 525 | // section setup end |
| 526 | // |
| 527 | |
| 528 | // Similar (but not idendtical) to upstream, should be kept the same when |
| 529 | // possible |
| 530 | // compared to upstream we use `make_tma_copy_A`, `make_tma_copy_B` etc. to |
| 531 | // define the TMA types |
| 532 | // Device side kernel params |
| 533 | struct Params { |
| 534 | public: |
| 535 | // Assumption: StrideA is congruent with Problem_MK |
| 536 | using TMA_A = decltype(make_tma_copy_A()); |
| 537 | using TMA_Scale = decltype(make_tma_copy_scale()); |
| 538 | using TMA_Zero = decltype(make_tma_copy_zero()); |
| 539 | using TMA_B = decltype(make_tma_copy_B()); |
| 540 | |
| 541 | // required by outer loop: i.e. |
| 542 | // cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp |
| 543 | TMA_A tma_load_a; |
| 544 | TMA_B tma_load_b; |
| 545 | TMA_Scale tma_load_scale; |
| 546 | TMA_Zero tma_load_zero; |
| 547 | int64_t scale_k; |
| 548 | int group_size; |
| 549 | uint32_t tma_transaction_bytes = TmaTransactionBytes; |
| 550 | uint32_t tma_transaction_bytes_mk = TmaTransactionBytesMK; |
| 551 | uint32_t tma_transaction_bytes_nk = TmaTransactionBytesNK; |
| 552 | }; |
| 553 | |
| 554 | // |
| 555 | // Methods |
| 556 | // |
| 557 | |
| 558 | // Similar (but not idendtical) to upstream, should be kept the same when |
| 559 | // possible |
| 560 | // compared to upstream we use `make_tma_copy_A` and `TVbNbKL_to_offset` here |
| 561 | // to handle the prepacked layout |
| 562 | template <class ProblemShape> |
| 563 | static constexpr Params to_underlying_arguments( |
| 564 | ProblemShape const& problem_shape, Arguments const& args, |
| 565 | void* workspace) { |
| 566 | (void)workspace; |
| 567 | |
| 568 | // Optionally append 1s until problem shape is rank-4 (MNKL), in case it is |
| 569 | // only rank-3 (MNK) |
| 570 | auto problem_shape_MNKL = append<4>(problem_shape, 1); |
| 571 | auto [M, N, K, L] = problem_shape_MNKL; |
| 572 | |
| 573 | auto ptr_A = reinterpret_cast<InternalElementA const*>(args.ptr_A); |
| 574 | auto ptr_B = reinterpret_cast<InternalElementB const*>(args.ptr_B); |
| 575 | |
| 576 | auto make_logical_tensor = [&](auto ptr, auto shape, auto stride) { |
| 577 | return make_tensor(get_logical_ptr(ptr), make_layout(shape, stride)); |
| 578 | }; |
| 579 | |
| 580 | typename Params::TMA_A tma_load_a; |
| 581 | typename Params::TMA_B tma_load_b; |
| 582 | typename Params::TMA_Scale tma_load_scale; |
| 583 | typename Params::TMA_Zero tma_load_zero; |
| 584 | |
| 585 | auto layout = GmemLayoutA::TVbNbKL_to_offset_copy(make_shape(M, K, L)); |
| 586 | tma_load_a = make_tma_copy_A( |
| 587 | make_logical_tensor(ptr_A, shape(layout), stride(layout))); |
| 588 | |
| 589 | tma_load_b = make_tma_copy_B( |
| 590 | make_logical_tensor(ptr_B, make_shape(N, K, L), args.dB)); |
| 591 | |
| 592 | int32_t scale_k = |
| 593 | (ModeHasScales) ? (K + args.group_size - 1) / args.group_size : 0; |
| 594 | int32_t group_size = (ModeHasScales) ? args.group_size : 0; |
| 595 | |
| 596 | if constexpr (ModeHasScales) { |
| 597 | tma_load_scale = make_tma_copy_scale( |
| 598 | make_logical_tensor(args.ptr_S, make_shape(M, scale_k, L), args.dS)); |
| 599 | } |
| 600 | |
| 601 | if constexpr (KernelConversionMode == |
| 602 | ConversionMode::ConvertAndScaleWithZero) { |
| 603 | tma_load_zero = make_tma_copy_zero( |
| 604 | make_logical_tensor(args.ptr_Z, make_shape(M, scale_k, L), args.dS)); |
| 605 | } |
| 606 | |
| 607 | if constexpr (KernelConversionMode == ConversionMode::DirectConvert || |
| 608 | KernelConversionMode == ConversionMode::ConvertAndScale || |
| 609 | KernelConversionMode == |
| 610 | ConversionMode::ConvertAndScaleWithZero) { |
| 611 | return {tma_load_a, tma_load_b, tma_load_scale, |
| 612 | tma_load_zero, scale_k, group_size}; |
| 613 | } else { |
| 614 | static_assert(cutlass::detail::dependent_false<KernelSchedule>, |
| 615 | "Conversion mode not handled in to_underlying_arguments."); |
| 616 | } |
| 617 | } |
| 618 | |
| 619 | // Same as upstream, should be kept the same when possible, not formatted for |
| 620 | // easier comparison |
| 621 | // with `SwapAB ? N : M -> M` since we don't support SwapAB |
| 622 | // clang-format off |
| 623 | template<class ProblemShape> |
| 624 | static bool |
| 625 | can_implement( |
| 626 | ProblemShape const& problem_shape, |
| 627 | [[maybe_unused]] Arguments const& args) { |
| 628 | constexpr int tma_alignment_bits = 128; |
| 629 | auto problem_shape_MNKL = append<4>(problem_shape, 1); |
| 630 | auto [M,N,K,L] = problem_shape_MNKL; |
| 631 | |
| 632 | bool implementable = true; |
| 633 | constexpr int min_tma_aligned_elements_A = tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value; |
| 634 | implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_A>(cute::make_shape(M,K,L), StrideA{}); |
| 635 | constexpr int min_tma_aligned_elements_B = tma_alignment_bits / cutlass::sizeof_bits<ElementB>::value; |
| 636 | implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_B>(cute::make_shape(N,K,L), StrideB{}); |
| 637 | |
| 638 | if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { |
| 639 | implementable = implementable && (args.ptr_S == nullptr); |
| 640 | implementable = implementable && (args.ptr_Z == nullptr); |
| 641 | } |
| 642 | else if constexpr (ModeHasScales) { |
| 643 | const int scale_mn = M; |
| 644 | const int scale_k = (K + args.group_size - 1) / args.group_size; |
| 645 | constexpr int min_tma_aligned_elements_scale = tma_alignment_bits / cutlass::sizeof_bits<ElementScale>::value; |
| 646 | implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_scale>(cute::make_shape(scale_mn,scale_k,L), StrideScale{}); |
| 647 | implementable = implementable && (args.group_size == K || ((args.group_size % size<2>(TileShape{})) == 0)); |
| 648 | implementable = implementable && args.group_size != 0; |
| 649 | implementable = implementable && (args.ptr_S != nullptr); |
| 650 | |
| 651 | if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { |
| 652 | implementable = implementable && (args.ptr_Z == nullptr); |
| 653 | } |
| 654 | else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { |
| 655 | constexpr int min_tma_aligned_elements_zero = tma_alignment_bits / cutlass::sizeof_bits<ElementZero>::value; |
| 656 | implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_zero>(cute::make_shape(scale_mn,scale_k,L), StrideScale{}); |
| 657 | implementable = implementable && (args.ptr_Z != nullptr); |
| 658 | } |
| 659 | else { |
| 660 | static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in can_implement."); |
| 661 | } |
| 662 | } |
| 663 | else { |
| 664 | static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in can_implement."); |
| 665 | } |
| 666 | |
| 667 | if (!implementable) { |
| 668 | CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n"); |
| 669 | } |
| 670 | return implementable; |
| 671 | } |
| 672 | |
| 673 | static constexpr int K_PIPE_MAX = DispatchPolicy::Stages; |
| 674 | static constexpr uint32_t TmaTransactionBytesMK = compute_tma_transaction_bytes_mk(); |
| 675 | static constexpr uint32_t TmaTransactionBytesNK = compute_tma_transaction_bytes_nk(); |
| 676 | static constexpr uint32_t TmaTransactionBytes = TmaTransactionBytesMK + TmaTransactionBytesNK; |
| 677 | |
| 678 | /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance |
| 679 | CUTLASS_DEVICE |
| 680 | static void prefetch_tma_descriptors(Params const& mainloop_params) { |
| 681 | cute::prefetch_tma_descriptor(mainloop_params.tma_load_a.get_tma_descriptor()); |
| 682 | cute::prefetch_tma_descriptor(mainloop_params.tma_load_b.get_tma_descriptor()); |
| 683 | |
| 684 | if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { |
| 685 | // Nothing extra to do |
| 686 | } |
| 687 | else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { |
| 688 | cute::prefetch_tma_descriptor(mainloop_params.tma_load_scale.get_tma_descriptor()); |
| 689 | } |
| 690 | else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { |
| 691 | cute::prefetch_tma_descriptor(mainloop_params.tma_load_scale.get_tma_descriptor()); |
| 692 | cute::prefetch_tma_descriptor(mainloop_params.tma_load_zero.get_tma_descriptor()); |
| 693 | } |
| 694 | else { |
| 695 | static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in TMA prefetch."); |
| 696 | } |
| 697 | |
| 698 | } |
| 699 | // clang-format off |
| 700 | |
| 701 | // Modified from upstream, should be kept close to that when possible |
| 702 | // the main difference is special handling for the prepacked A layout |
| 703 | // |
| 704 | // Set up the data needed by this collective for load and mma. |
| 705 | // Returns a tuple of tensors. The collective and the kernel layer have the |
| 706 | // contract Returned tuple must contain at least two elements, with the first |
| 707 | // two elements being: gA_mkl - The tma tensor, A after a local tile so it |
| 708 | // has shape (TILE_V,TILE_B,m,k,l) gB_nkl - The tma tensor, B after a local |
| 709 | // tile so it has shape (TILE_N,TILE_K,n,k,l) The rest of the tensors can be |
| 710 | // specified as needed by this collective. |
| 711 | // NOTE: TILE_B is the prepacked block index within a tile. TILE_V is the |
| 712 | // values within a prepacked block. |
| 713 | template <class ProblemShape_MNKL> |
| 714 | CUTLASS_DEVICE auto load_init(ProblemShape_MNKL const& problem_shape_MNKL, |
| 715 | Params const& mainloop_params) const { |
| 716 | using X = Underscore; |
| 717 | auto M = get<0>(problem_shape_MNKL), N = get<1>(problem_shape_MNKL), |
| 718 | K = get<2>(problem_shape_MNKL), L = get<3>(problem_shape_MNKL); |
| 719 | |
| 720 | // (TILE_V,TILE_B,m,k,l) |
| 721 | auto make_gA_mkl = [&]() { |
| 722 | // ((athrid, val), (BlocksM, BlockK), L) -> (storage_idx) |
| 723 | auto layout = GmemLayoutA::TVbNbKL_to_offset_copy(make_shape(M, K, L)); |
| 724 | Tensor mA_mkl = mainloop_params.tma_load_a.get_tma_tensor(shape(layout)); |
| 725 | return local_tile(mA_mkl, |
| 726 | make_shape(size<0>(layout), PPBlocksPerTile_MK{}), |
| 727 | make_coord(0, make_coord(_, _))); |
| 728 | }; |
| 729 | |
| 730 | // (TILE_N,TILE_K,n,k,l) |
| 731 | auto make_gB_nkl = [&]() { |
| 732 | Tensor mB_nkl = |
| 733 | mainloop_params.tma_load_b.get_tma_tensor(make_shape(N, K, L)); |
| 734 | return local_tile(mB_nkl, TileShape{}, make_coord(_, _, _), |
| 735 | Step<X, _1, _1>{}); |
| 736 | }; |
| 737 | |
| 738 | // (TILE_M,TILE_Scale_K,m,scale_k,l) |
| 739 | auto make_gS_mkl = [&]() { |
| 740 | auto scale_k = mainloop_params.scale_k; |
| 741 | Tensor mS_mkl = mainloop_params.tma_load_scale.get_tma_tensor( |
| 742 | make_shape(M, scale_k, L)); |
| 743 | return local_tile(mS_mkl, ScaleTileShape{}, make_coord(_, _)); |
| 744 | }; |
| 745 | |
| 746 | // (TILE_M,TILE_Scale_K,m,scale_k,l) |
| 747 | auto make_gZ_mkl = [&]() { |
| 748 | auto scale_k = mainloop_params.scale_k; |
| 749 | Tensor mZ_mkl = mainloop_params.tma_load_zero.get_tma_tensor( |
| 750 | make_shape(M, scale_k, L)); |
| 751 | return local_tile(mZ_mkl, ScaleTileShape{}, make_coord(_, _)); |
| 752 | }; |
| 753 | |
| 754 | if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { |
| 755 | return cute::make_tuple(make_gA_mkl(), make_gB_nkl()); |
| 756 | } else if constexpr (KernelConversionMode == |
| 757 | ConversionMode::ConvertAndScale) { |
| 758 | return cute::make_tuple(make_gA_mkl(), make_gB_nkl(), make_gS_mkl()); |
| 759 | } else if constexpr (KernelConversionMode == |
| 760 | ConversionMode::ConvertAndScaleWithZero) { |
| 761 | return cute::make_tuple(make_gA_mkl(), make_gB_nkl(), make_gS_mkl(), |
| 762 | make_gZ_mkl()); |
| 763 | } else { |
| 764 | static_assert(cutlass::detail::dependent_false<KernelSchedule>, |
| 765 | "Conversion mode not handled in load_init."); |
| 766 | } |
| 767 | } |
| 768 | |
| 769 | // Similar to upstream, should be kept close to that when possible |
| 770 | // the main difference is in the layout comments |
| 771 | // clang-format off |
| 772 | /// Perform a collective-scoped matrix multiply-accumulate |
| 773 | /// Producer Perspective |
| 774 | /// This overload gets triggered when we have scales. |
| 775 | template < |
| 776 | class... Ts, |
| 777 | class KTileIterator, class BlockCoord |
| 778 | > |
| 779 | CUTLASS_DEVICE void |
| 780 | load( |
| 781 | Params const& mainloop_params, |
| 782 | MainloopPipeline pipeline, |
| 783 | PipelineState smem_pipe_write, |
| 784 | cute::tuple<Ts...> const& load_inputs, |
| 785 | BlockCoord const& blk_coord, |
| 786 | KTileIterator k_tile_iter, int k_tile_count, |
| 787 | int thread_idx, |
| 788 | uint32_t block_rank_in_cluster, |
| 789 | TensorStorage& shared_tensors) { |
| 790 | if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { |
| 791 | static_assert(sizeof... (Ts) == 2, "Direct convert needs two inputs"); |
| 792 | } |
| 793 | else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { |
| 794 | static_assert(sizeof... (Ts) == 3, "Scaled convert needs three inputs"); |
| 795 | } |
| 796 | else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { |
| 797 | static_assert(sizeof... (Ts) == 4, "Scaled and zero convert needs four inputs"); |
| 798 | } |
| 799 | else { |
| 800 | static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in TMA load."); |
| 801 | } |
| 802 | |
| 803 | int lane_predicate = cute::elect_one_sync(); |
| 804 | |
| 805 | if (lane_predicate) { |
| 806 | Tensor sA_ = make_tensor(make_smem_ptr(shared_tensors.smem_A.begin()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE) |
| 807 | Tensor sB_ = make_tensor(make_smem_ptr(shared_tensors.smem_B.begin()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE) |
| 808 | Tensor sA = as_position_independent_swizzle_tensor(sA_); // (BLK_M,BLK_K,PIPE) |
| 809 | Tensor sB = as_position_independent_swizzle_tensor(sB_); // (BLK_N,BLK_K,PIPE) |
| 810 | |
| 811 | // |
| 812 | // Prepare the TMA loads for A, B and Scales |
| 813 | // |
| 814 | |
| 815 | constexpr uint32_t cluster_shape_x = get<0>(ClusterShape()); |
| 816 | uint2 cluster_local_block_id = {block_rank_in_cluster % cluster_shape_x, block_rank_in_cluster / cluster_shape_x}; |
| 817 | |
| 818 | Tensor gA_mkl = get<0>(load_inputs); |
| 819 | Tensor gB_nkl = get<1>(load_inputs); |
| 820 | |
| 821 | auto block_tma_a = mainloop_params.tma_load_a.get_slice(cluster_local_block_id.y); |
| 822 | auto block_tma_b = mainloop_params.tma_load_b.get_slice(cluster_local_block_id.x); |
| 823 | |
| 824 | // Partition the inputs based on the current block coordinates. |
| 825 | auto [m_coord, n_coord, k_coord, l_coord] = blk_coord; |
| 826 | Tensor gA = gA_mkl(_,_,m_coord,_,l_coord); // (TILE_V,TILE_B,k) |
| 827 | Tensor gB = gB_nkl(_,_,n_coord,_,l_coord); // (TILE_N,TILE_K,k) |
| 828 | |
| 829 | // Applies the mapping from block_tma_a |
| 830 | Tensor tAgA = block_tma_a.partition_S(gA); // (TMA,TMA_M,TMA_K,k) |
| 831 | Tensor tAsA = block_tma_a.partition_D(sA); // (TMA,TMA_M,TMA_K,PIPE) |
| 832 | |
| 833 | Tensor tBgB = block_tma_b.partition_S(gB); // (TMA,TMA_N,TMA_K,k) |
| 834 | Tensor tBsB = block_tma_b.partition_D(sB); // (TMA,TMA_N,TMA_K,PIPE) |
| 835 | |
| 836 | uint16_t mcast_mask_a = 0; |
| 837 | uint16_t mcast_mask_b = 0; |
| 838 | uint16_t mcast_mask_s = 0; |
| 839 | |
| 840 | // Issue TmaLoads |
| 841 | // Maps the tile -> block, value |
| 842 | if constexpr (cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD_MULTICAST>) { |
| 843 | auto block_layout = Layout<typename DispatchPolicy::ClusterShape>{}; // (m,n) -> block_id |
| 844 | for (int n = 0; n < size<1>(block_layout); ++n) { |
| 845 | mcast_mask_a |= (uint16_t(1) << block_layout(cluster_local_block_id.x,n,Int<0>{})); |
| 846 | } |
| 847 | } |
| 848 | |
| 849 | if constexpr (cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD_MULTICAST>) { |
| 850 | auto block_layout = Layout<typename DispatchPolicy::ClusterShape>{}; // (m,n) -> block_id |
| 851 | for (int m = 0; m < size<0>(block_layout); ++m) { |
| 852 | mcast_mask_b |= (uint16_t(1) << block_layout(m,cluster_local_block_id.y,Int<0>{})); |
| 853 | } |
| 854 | } |
| 855 | |
| 856 | auto extra_input_partitions = partition_extra_tma_inputs(mainloop_params, load_inputs, shared_tensors, cluster_local_block_id, m_coord, l_coord); |
| 857 | |
| 858 | // Mainloop |
| 859 | CUTLASS_PRAGMA_NO_UNROLL |
| 860 | for ( ; k_tile_count > 0; --k_tile_count) { |
| 861 | // LOCK smem_pipe_write for _writing_ |
| 862 | pipeline.producer_acquire(smem_pipe_write); |
| 863 | |
| 864 | // |
| 865 | // Copy gmem to smem for *k_tile_iter |
| 866 | // |
| 867 | |
| 868 | using BarrierType = typename MainloopPipeline::ProducerBarrierType; |
| 869 | BarrierType* tma_barrier = pipeline.producer_get_barrier(smem_pipe_write); |
| 870 | |
| 871 | int write_stage = smem_pipe_write.index(); |
| 872 | copy(mainloop_params.tma_load_a.with(*tma_barrier, mcast_mask_a), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,write_stage)); |
| 873 | copy(mainloop_params.tma_load_b.with(*tma_barrier, mcast_mask_b), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,write_stage)); |
| 874 | |
| 875 | if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { |
| 876 | // Nothing extra to do. |
| 877 | } |
| 878 | else if constexpr (ModeHasScales) { |
| 879 | auto tSgS = get<0>(extra_input_partitions); |
| 880 | auto tSsS = get<1>(extra_input_partitions); |
| 881 | |
| 882 | // Temporary factor which will determine which k tile to reload from gmem. Needed so we don't modify tma transaction bytes |
| 883 | // on the fly. |
| 884 | // We must do a ceiling divide here to correctly handle with group_size == K. In that case, we don't require that K |
| 885 | // is a multiple of the threadblock tile K |
| 886 | const int ReloadFactor = (mainloop_params.group_size + size<2>(TileShape{}) - 1) / size<2>(TileShape{}); |
| 887 | const int scale_load_k = *k_tile_iter / ReloadFactor; // This will always be 0 when group_size == K. |
| 888 | copy(mainloop_params.tma_load_scale.with(*tma_barrier, mcast_mask_s), tSgS(_,_,_,scale_load_k), tSsS(_,_,_,write_stage)); |
| 889 | |
| 890 | if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { |
| 891 | // Nothing extra to do |
| 892 | } |
| 893 | else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { |
| 894 | auto tZgZ = get<2>(extra_input_partitions); |
| 895 | auto tZsZ = get<3>(extra_input_partitions); |
| 896 | copy(mainloop_params.tma_load_zero.with(*tma_barrier, mcast_mask_s), tZgZ(_,_,_,scale_load_k), tZsZ(_,_,_,write_stage)); |
| 897 | } |
| 898 | else { |
| 899 | static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled for TMA copy op."); |
| 900 | } |
| 901 | } |
| 902 | else { |
| 903 | static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled for TMA copy op."); |
| 904 | } |
| 905 | |
| 906 | ++k_tile_iter; |
| 907 | |
| 908 | // Advance smem_pipe_write |
| 909 | ++smem_pipe_write; |
| 910 | } |
| 911 | } |
| 912 | } |
| 913 | // clang-format off |
| 914 | |
| 915 | // Same as upstream, should be kept the same when possible, not formatted for |
| 916 | // easier comparison |
| 917 | // clang-format off |
| 918 | // Perform a Producer Epilogue to prevent early exit of blocks in a Cluster |
| 919 | CUTLASS_DEVICE void |
| 920 | load_tail(MainloopPipeline pipeline, PipelineState smem_pipe_write) { |
| 921 | int lane_predicate = cute::elect_one_sync(); |
| 922 | |
| 923 | // Issue the epilogue waits |
| 924 | if (lane_predicate) { |
| 925 | /* This helps avoid early exit of blocks in Cluster |
| 926 | * Waits for all stages to either be released (all |
| 927 | * Consumer UNLOCKs), or if the stage was never used |
| 928 | * then would just be acquired since the phase was |
| 929 | * still inverted from make_producer_start_state |
| 930 | */ |
| 931 | pipeline.producer_tail(smem_pipe_write); |
| 932 | } |
| 933 | } |
| 934 | // clang-format on |
| 935 | |
| 936 | // Modified from upstream, should be kept close to that when possible |
| 937 | // the main differences are handling the prepacked A layout, and separating |
| 938 | // the loading of A from upcoverting A |
| 939 | // |
| 940 | // Perform a collective-scoped matrix multiply-accumulate |
| 941 | // Consumer Perspective |
| 942 | template <class FrgTensorC> |
| 943 | CUTLASS_DEVICE void mma(MainloopPipeline pipeline, |
| 944 | PipelineState smem_pipe_read, FrgTensorC& accum, |
| 945 | int k_tile_count, int thread_idx, |
| 946 | TensorStorage& shared_tensors, |
| 947 | Params const& mainloop_params) { |
| 948 | static_assert(is_rmem<FrgTensorC>::value, |
| 949 | "C tensor must be rmem resident."); |
| 950 | static_assert(cute::rank(SmemLayoutB{}) == 3, |
| 951 | "Smem layout must be rank 3."); |
| 952 | static_assert(cute::rank(SmemLayoutAtomB{}) == 2, |
| 953 | "SmemLayoutAtomB must be rank 2."); |
| 954 | static_assert(!cute::is_void_v<SmemCopyAtomA>, |
| 955 | "SM90 GMMA mainloops must specify a non-void copy atom for " |
| 956 | "RF sourced instructions."); |
| 957 | static_assert(cute::is_void_v<SmemCopyAtomB>, |
| 958 | "SM90 GMMA mainloops cannot have a non-void copy atom for " |
| 959 | "smem sourced instructions."); |
| 960 | |
| 961 | // Obtain warp index |
| 962 | int warp_idx = canonical_warp_idx_sync(); |
| 963 | [[maybe_unused]] int warp_group_thread_idx = thread_idx % 128; |
| 964 | |
| 965 | // ((T, (FrgV,(RestM, RestK)), (BlocksM, BlocksK), pipe) -> offset |
| 966 | auto constexpr smem_A = SmemLayoutA{}; |
| 967 | |
| 968 | // convert: |
| 969 | // ((T, (MMA,(MMA_M, MMA_K)), (BlocksM, BlocksK), pipe) -> offset |
| 970 | // to: |
| 971 | // (T, MMA, ((MMA_M, BlocksM), (MMA_K, BlocksK)), pipe) -> offset |
| 972 | // which can be thought of as: |
| 973 | // (T, MMA, (MMA_M, MMA_K), pipe) -> offset |
| 974 | auto constexpr smem_A_mma_ = |
| 975 | make_layout(get<0, 0>(smem_A), get<0, 1, 0>(smem_A), |
| 976 | zip(get<0, 1, 1>(smem_A), get<1>(smem_A)), get<2>(smem_A)); |
| 977 | // flatten to: |
| 978 | // (T, MMA, MMA_M, MMA_K, pipe) -> offset |
| 979 | auto constexpr smem_A_mma = smem_A_mma_(_, _, make_coord(_, _), _); |
| 980 | |
| 981 | Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.begin()), |
| 982 | smem_A_mma); // (T, MMA, MMA_M, MMA_K, pipe) |
| 983 | Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.begin()), |
| 984 | SmemLayoutB{}); // (BLK_N,BLK_K,PIPE) |
| 985 | |
| 986 | // |
| 987 | // Define C accumulators and A/B partitioning |
| 988 | // |
| 989 | |
| 990 | TiledMma tiled_mma; |
| 991 | auto thread_mma = tiled_mma.get_thread_slice(thread_idx); |
| 992 | |
| 993 | Tensor tCsA = sA(thread_idx, _, _, _, _); // (MMA,MMA_M,MMA_K,PIPE) |
| 994 | Tensor tCsB = thread_mma.partition_B(sB); // (MMA,MMA_N,MMA_K,PIPE) |
| 995 | |
| 996 | // Allocate fragments and descriptors |
| 997 | Tensor tCrA_load = make_tensor<ElementA>( |
| 998 | tCsA(_, _, _, Int<0>{}).shape()); // (MMA,MMA_N,MMA_K) |
| 999 | Tensor tCrA_mma = make_fragment_like<ElementMma>(tCrA_load); |
| 1000 | |
| 1001 | Tensor tCrB = thread_mma.make_fragment_B(tCsB); // (MMA,MMA_N,MMA_K,PIPE) |
| 1002 | |
| 1003 | static constexpr int A_CPY_VEC = |
| 1004 | decltype(max_common_vector(tCsA, tCrA_load)){}; |
| 1005 | |
| 1006 | static constexpr int CONVERSION_WIDTH = |
| 1007 | std::min(A_CPY_VEC, int(size<0>(tCrA_mma))); |
| 1008 | |
| 1009 | auto load_A_to_registers = [&](int read_stage) { |
| 1010 | copy(create_auto_vectorizing_copy<ElementA, decltype(A_CPY_VEC)>(), |
| 1011 | tCsA(_, _, _, read_stage), tCrA_load(_, _, _)); |
| 1012 | }; |
| 1013 | |
| 1014 | // Partition of thread -> shared and thread -> RF |
| 1015 | auto partitioned_extra_info = |
| 1016 | partition_extra_mma_info(thread_mma, shared_tensors); |
| 1017 | auto copy_partitions_extra_info = retile_extra_mma_info( |
| 1018 | tiled_mma, partitioned_extra_info, warp_group_thread_idx); |
| 1019 | CUTE_STATIC_ASSERT_V(size<1>(tCrA_mma) == size<1>(accum)); // MMA_M |
| 1020 | CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum)); // N |
| 1021 | CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB)); // K |
| 1022 | CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB)); // PIPE |
| 1023 | CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sB)); // PIPE |
| 1024 | |
| 1025 | // |
| 1026 | // PIPELINED MAIN LOOP |
| 1027 | // |
| 1028 | |
| 1029 | auto convert_A = [&, a_vec = Int<CONVERSION_WIDTH>{}](int k_block, |
| 1030 | int read_stage) { |
| 1031 | load_extra_info_to_registers(partitioned_extra_info, |
| 1032 | copy_partitions_extra_info, k_block, |
| 1033 | read_stage); |
| 1034 | transform_A_kblock(tCrA_load, a_vec, tCrA_mma, partitioned_extra_info, |
| 1035 | k_block); |
| 1036 | }; |
| 1037 | |
| 1038 | // We release buffers to producer warps(dma load) with some mmas in flight |
| 1039 | PipelineState smem_pipe_release = smem_pipe_read; |
| 1040 | |
| 1041 | tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; |
| 1042 | |
| 1043 | warpgroup_fence_operand(accum); |
| 1044 | |
| 1045 | constexpr int K_BLOCK_MAX = size<2>(tCrA_load); |
| 1046 | |
| 1047 | ConsumerToken barrier_token = {BarrierStatus::WaitAgain}; |
| 1048 | // first k tile |
| 1049 | { |
| 1050 | barrier_token = pipeline.consumer_try_wait(smem_pipe_read); |
| 1051 | pipeline.consumer_wait(smem_pipe_read, barrier_token); |
| 1052 | |
| 1053 | int read_stage = smem_pipe_read.index(); |
| 1054 | ++smem_pipe_read; |
| 1055 | barrier_token = pipeline.consumer_try_wait(smem_pipe_read); |
| 1056 | |
| 1057 | // copy smem->rmem for A operand |
| 1058 | load_A_to_registers(read_stage); |
| 1059 | convert_A(0, read_stage); |
| 1060 | |
| 1061 | // Unroll the K mode manually to set scale D to 1 |
| 1062 | CUTLASS_PRAGMA_UNROLL |
| 1063 | for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block) { |
| 1064 | if (k_block < K_BLOCK_MAX - 1) { |
| 1065 | convert_A(k_block + 1, smem_pipe_read.index()); |
| 1066 | } |
| 1067 | warpgroup_arrive(); |
| 1068 | // (V,M) x (V,N) => (V,M,N) |
| 1069 | cute::gemm(tiled_mma, tCrA_mma(_, _, k_block), |
| 1070 | tCrB(_, _, k_block, read_stage), accum); |
| 1071 | tiled_mma.accumulate_ = GMMA::ScaleOut::One; |
| 1072 | warpgroup_commit_batch(); |
| 1073 | } |
| 1074 | |
| 1075 | --k_tile_count; |
| 1076 | if (k_tile_count > 0) { |
| 1077 | // Wait for K_BLOCK_MAX - 1 to be in flight to ensure that it is safe to |
| 1078 | // overwrite the A registers for the first mma. |
| 1079 | warpgroup_wait<K_BLOCK_MAX - 1>(); |
| 1080 | pipeline.consumer_wait(smem_pipe_read, barrier_token); |
| 1081 | load_A_to_registers(smem_pipe_read.index()); |
| 1082 | convert_A(0, smem_pipe_read.index()); |
| 1083 | } |
| 1084 | } |
| 1085 | |
| 1086 | if (k_tile_count == 0) { |
| 1087 | return; |
| 1088 | } |
| 1089 | |
| 1090 | warpgroup_fence_operand(accum); |
| 1091 | // Mainloop GMMAs |
| 1092 | CUTLASS_PRAGMA_NO_UNROLL |
| 1093 | for (; k_tile_count > 1; --k_tile_count) { |
| 1094 | // |
| 1095 | // Compute on k_tile |
| 1096 | // |
| 1097 | |
| 1098 | int read_stage = smem_pipe_read.index(); |
| 1099 | ++smem_pipe_read; |
| 1100 | |
| 1101 | warpgroup_fence_operand(accum); |
| 1102 | // Unroll the K mode manually to set scale D to 1 |
| 1103 | CUTLASS_PRAGMA_UNROLL |
| 1104 | for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block) { |
| 1105 | warpgroup_arrive(); |
| 1106 | // (V,M) x (V,N) => (V,M,N) |
| 1107 | cute::gemm(tiled_mma, tCrA_mma(_, _, k_block), |
| 1108 | tCrB(_, _, k_block, read_stage), accum); |
| 1109 | tiled_mma.accumulate_ = GMMA::ScaleOut::One; |
| 1110 | warpgroup_commit_batch(); |
| 1111 | |
| 1112 | warpgroup_wait<K_BLOCK_MAX - 1>(); |
| 1113 | if (k_block == K_BLOCK_MAX - 1) { |
| 1114 | // We have K_BLOCK_MAX - 1 GMMA instructions pending for this stage, |
| 1115 | // so we can release prior barrier |
| 1116 | pipeline.consumer_release( |
| 1117 | smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ |
| 1118 | // on it |
| 1119 | ++smem_pipe_release; |
| 1120 | } |
| 1121 | |
| 1122 | if (k_block == 0) { |
| 1123 | barrier_token = pipeline.consumer_try_wait(smem_pipe_read); |
| 1124 | } |
| 1125 | |
| 1126 | if (k_block == K_BLOCK_MAX - 1) { |
| 1127 | pipeline.consumer_wait(smem_pipe_read, barrier_token); |
| 1128 | load_A_to_registers(smem_pipe_read.index()); |
| 1129 | convert_A(0, smem_pipe_read.index()); |
| 1130 | } else { |
| 1131 | convert_A(k_block + 1, read_stage); |
| 1132 | } |
| 1133 | } |
| 1134 | warpgroup_fence_operand(accum); |
| 1135 | } |
| 1136 | |
| 1137 | warpgroup_fence_operand(accum); |
| 1138 | |
| 1139 | { |
| 1140 | // |
| 1141 | // Compute on k_tile |
| 1142 | // |
| 1143 | |
| 1144 | int read_stage = smem_pipe_read.index(); |
| 1145 | |
| 1146 | warpgroup_fence_operand(accum); |
| 1147 | |
| 1148 | // Unroll the K mode manually to set scale D to 1 |
| 1149 | CUTLASS_PRAGMA_UNROLL |
| 1150 | for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block) { |
| 1151 | warpgroup_arrive(); |
| 1152 | // (V,M) x (V,N) => (V,M,N) |
| 1153 | cute::gemm(tiled_mma, tCrA_mma(_, _, k_block), |
| 1154 | tCrB(_, _, k_block, read_stage), accum); |
| 1155 | tiled_mma.accumulate_ = GMMA::ScaleOut::One; |
| 1156 | warpgroup_commit_batch(); |
| 1157 | warpgroup_wait<K_BLOCK_MAX - 1>(); |
| 1158 | if (k_block == K_BLOCK_MAX - 1) { |
| 1159 | // release prior barrier |
| 1160 | pipeline.consumer_release( |
| 1161 | smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ |
| 1162 | // on it |
| 1163 | ++smem_pipe_release; |
| 1164 | } |
| 1165 | |
| 1166 | if (k_block < K_BLOCK_MAX - 1) { |
| 1167 | convert_A(k_block + 1, read_stage); |
| 1168 | } |
| 1169 | } |
| 1170 | } |
| 1171 | |
| 1172 | warpgroup_fence_operand(accum); |
| 1173 | } |
| 1174 | |
| 1175 | // Perform a Consumer Epilogue to release all buffers |
| 1176 | CUTLASS_DEVICE void mma_tail(MainloopPipeline pipeline, |
| 1177 | PipelineState smem_pipe_release, |
| 1178 | int k_tile_count) { |
| 1179 | // Prologue GMMAs |
| 1180 | int prologue_mma_count = 1; |
| 1181 | k_tile_count -= prologue_mma_count; |
| 1182 | |
| 1183 | smem_pipe_release.advance(k_tile_count); |
| 1184 | |
| 1185 | // Wait on all GMMAs to complete |
| 1186 | warpgroup_wait<0>(); |
| 1187 | |
| 1188 | for (int count = 0; count < prologue_mma_count; ++count) { |
| 1189 | pipeline.consumer_release( |
| 1190 | smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on |
| 1191 | // it |
| 1192 | ++smem_pipe_release; |
| 1193 | } |
| 1194 | } |
| 1195 | |
| 1196 | private: |
| 1197 | // Same as upstream, should be kept the same when possible, not formatted for |
| 1198 | // easier comparison |
| 1199 | // clang-format off |
| 1200 | /// Utilities for any additional inputs inside of the TMA load |
| 1201 | template <class... Ts> |
| 1202 | CUTLASS_DEVICE |
| 1203 | auto partition_extra_tma_inputs( |
| 1204 | Params const& mainloop_params, |
| 1205 | cute::tuple<Ts...> const& load_inputs, |
| 1206 | TensorStorage& shared_tensors, |
| 1207 | uint2 const& cluster_local_block_id, |
| 1208 | int const m_coord, |
| 1209 | int const l_coord) { |
| 1210 | |
| 1211 | if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { |
| 1212 | return cute::make_tuple(); |
| 1213 | } |
| 1214 | else if constexpr (ModeHasScales) { |
| 1215 | Tensor sS = make_tensor(make_smem_ptr(shared_tensors.smem_scale.begin()), SmemLayoutScale{}); // (BLK_M,BLK_K,PIPE) |
| 1216 | Tensor gS_mkl = get<2>(load_inputs); |
| 1217 | auto block_tma_s = mainloop_params.tma_load_scale.get_slice(cluster_local_block_id.y); |
| 1218 | Tensor gS = gS_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k) |
| 1219 | |
| 1220 | Tensor tSgS = block_tma_s.partition_S(gS); // (TMA,TMA_M,TMA_K,k) |
| 1221 | Tensor tSsS = block_tma_s.partition_D(sS); // (TMA,TMA_M,TMA_K,PIPE) |
| 1222 | if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { |
| 1223 | return cute::make_tuple(tSgS, tSsS); |
| 1224 | } |
| 1225 | else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { |
| 1226 | Tensor sZ = make_tensor(make_smem_ptr(shared_tensors.smem_zero.begin()), SmemLayoutScale{}); // (BLK_M,BLK_K,PIPE) |
| 1227 | Tensor gZ_mkl = get<3>(load_inputs); |
| 1228 | auto block_tma_z = mainloop_params.tma_load_zero.get_slice(cluster_local_block_id.y); |
| 1229 | Tensor gZ = gZ_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k) |
| 1230 | |
| 1231 | Tensor tZgZ = block_tma_z.partition_S(gZ); // (TMA,TMA_M,TMA_K,k) |
| 1232 | Tensor tZsZ = block_tma_z.partition_D(sZ); // (TMA,TMA_M,TMA_K,PIPE) |
| 1233 | return cute::make_tuple(tSgS, tSsS, tZgZ, tZsZ); |
| 1234 | } |
| 1235 | else { |
| 1236 | static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled for input partitioning."); |
| 1237 | } |
| 1238 | } |
| 1239 | else { |
| 1240 | static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled for input partitioning."); |
| 1241 | } |
| 1242 | } |
| 1243 | // clang-format off |
| 1244 | |
| 1245 | // Same as upstream, should be kept the same when possible, not formatted for |
| 1246 | // easier comparison |
| 1247 | // clang-format off |
| 1248 | /// Utilities for partitioning extra inputs for loading from smem in the mainloop. |
| 1249 | template <class ThreadMma> |
| 1250 | CUTLASS_DEVICE |
| 1251 | auto partition_extra_mma_info( |
| 1252 | ThreadMma const& mma_thread_slice, |
| 1253 | TensorStorage& shared_tensors) { |
| 1254 | |
| 1255 | if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { |
| 1256 | // nothing to do |
| 1257 | return cute::make_tuple(); |
| 1258 | } |
| 1259 | else if constexpr (ModeHasScales) { |
| 1260 | Tensor sS = make_tensor(make_smem_ptr(shared_tensors.smem_scale.begin()), SmemLayoutScale{});// (BLK_M,BLK_SCALE_K,PIPE) |
| 1261 | Tensor tCsS = mma_thread_slice.partition_A(sS); |
| 1262 | Tensor tCrS = make_tensor<ElementScale>(mma_thread_slice.partition_fragment_A(sS(_,_,Int<0>{})).shape()); |
| 1263 | |
| 1264 | if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { |
| 1265 | return cute::make_tuple(tCsS, tCrS); |
| 1266 | } |
| 1267 | else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { |
| 1268 | Tensor sZ = make_tensor(make_smem_ptr(shared_tensors.smem_zero.begin()), SmemLayoutScale{});// (BLK_M,BLK_SCALE_K,PIPE) |
| 1269 | Tensor tCsZ = mma_thread_slice.partition_A(sZ); |
| 1270 | Tensor tCrZ = make_tensor<ElementZero>(mma_thread_slice.partition_fragment_A(sZ(_,_,Int<0>{})).shape()); |
| 1271 | return cute::make_tuple(tCsS, tCrS, tCsZ, tCrZ); |
| 1272 | } |
| 1273 | else { |
| 1274 | static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in A -> RF path."); |
| 1275 | } |
| 1276 | } |
| 1277 | else { |
| 1278 | static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in A -> RF path."); |
| 1279 | } |
| 1280 | } |
| 1281 | // clang-format on |
| 1282 | |
| 1283 | // Same as upstream, should be kept the same when possible, not formatted for |
| 1284 | // easier comparison |
| 1285 | // clang-format off |
| 1286 | /// Returns the tiled copy and copy views for the extra inputs. |
| 1287 | template <class TiledMma, class... Ts> |
| 1288 | CUTLASS_DEVICE |
| 1289 | auto retile_extra_mma_info( |
| 1290 | TiledMma const& tiled_mma, |
| 1291 | cute::tuple<Ts...>& partitioned_extra_info, |
| 1292 | int const warp_group_thread_idx) { |
| 1293 | |
| 1294 | if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { |
| 1295 | // nothing to do |
| 1296 | return cute::make_tuple(); |
| 1297 | } |
| 1298 | else if constexpr (ModeHasScales) { |
| 1299 | auto smem_tiled_copy_S = make_tiled_copy_A(SmemCopyAtomScale{}, tiled_mma); |
| 1300 | auto smem_thr_copy_S = smem_tiled_copy_S.get_thread_slice(warp_group_thread_idx); |
| 1301 | Tensor tCrS_copy_view = smem_thr_copy_S.retile_D(cute::get<1>(partitioned_extra_info)); // (CPY,CPY_M,CPY_K) |
| 1302 | |
| 1303 | if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { |
| 1304 | return cute::make_tuple(smem_tiled_copy_S, tCrS_copy_view); |
| 1305 | } |
| 1306 | else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { |
| 1307 | Tensor tCrZ_copy_view = smem_thr_copy_S.retile_D(cute::get<3>(partitioned_extra_info)); // (CPY,CPY_M,CPY_K) |
| 1308 | return cute::make_tuple(smem_tiled_copy_S, tCrS_copy_view, tCrZ_copy_view); |
| 1309 | } |
| 1310 | else { |
| 1311 | static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in A -> RF path."); |
| 1312 | } |
| 1313 | } |
| 1314 | else { |
| 1315 | static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in A -> RF path."); |
| 1316 | } |
| 1317 | } |
| 1318 | // clang-format on |
| 1319 | |
| 1320 | // Similar to `copy_A_and_extra_info` upstream, should be kept the same when |
| 1321 | // possible |
| 1322 | // the main differences this only loads the extra info into registers and |
| 1323 | // not A (since we now preload more of A in the main pipeline) |
| 1324 | // Load scales and zeros into registers if required |
| 1325 | template <class... Ts, class... Us> |
| 1326 | CUTLASS_DEVICE void load_extra_info_to_registers( |
| 1327 | cute::tuple<Ts...> const& partitioned_mma_extra_info, |
| 1328 | cute::tuple<Us...> const& tiled_copy_and_views, int k_block, |
| 1329 | int read_stage) { |
| 1330 | if (k_block == 0) { |
| 1331 | // We are starting a new k-tile so copy the scale |
| 1332 | if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { |
| 1333 | // nothing to do |
| 1334 | } else if constexpr (ModeHasScales) { |
| 1335 | auto smem_tiled_copy_S = cute::get<0>(tiled_copy_and_views); |
| 1336 | auto tCrS_copy_view = cute::get<1>(tiled_copy_and_views); |
| 1337 | auto tCsS = cute::get<0>(partitioned_mma_extra_info); |
| 1338 | copy(smem_tiled_copy_S, tCsS(_, _, k_block, read_stage), |
| 1339 | tCrS_copy_view(_, _, k_block)); |
| 1340 | if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { |
| 1341 | // Nothing extra to do |
| 1342 | } else if constexpr (KernelConversionMode == |
| 1343 | ConversionMode::ConvertAndScaleWithZero) { |
| 1344 | auto tCsZ = cute::get<2>(partitioned_mma_extra_info); |
| 1345 | auto tCrZ_copy_view = cute::get<2>(tiled_copy_and_views); |
| 1346 | copy(smem_tiled_copy_S, tCsZ(_, _, k_block, read_stage), |
| 1347 | tCrZ_copy_view(_, _, k_block)); |
| 1348 | } else { |
| 1349 | static_assert(cutlass::detail::dependent_false<KernelSchedule>, |
| 1350 | "Conversion mode not handled in A -> RF path."); |
| 1351 | } |
| 1352 | } else { |
| 1353 | static_assert(cutlass::detail::dependent_false<KernelSchedule>, |
| 1354 | "Conversion mode not handled in A -> RF path."); |
| 1355 | } |
| 1356 | } |
| 1357 | } |
| 1358 | |
| 1359 | // Similar to upstream, should be kept the same when possible. |
| 1360 | // the main differences are that `convert_tensor` supports interleaved |
| 1361 | // layouts and bfloat16 has been optimized. `transform_internal_A` has also |
| 1362 | // been inlined for code simplicity. |
| 1363 | // Utilities to transform A. |
| 1364 | template <class TCrA_load, int VectorWidthA, class TCrA_mma, class... Ts> |
| 1365 | CUTLASS_DEVICE void transform_A_kblock( |
| 1366 | TCrA_load const& tCrA_load, cute::Int<VectorWidthA> vec_A, |
| 1367 | TCrA_mma& tCrA_mma, cute::tuple<Ts...> const& partitioned_extra_info, |
| 1368 | int const k_block) { |
| 1369 | auto in = tCrA_load(_, _, k_block); |
| 1370 | auto out = tCrA_mma(_, _, k_block); |
| 1371 | |
| 1372 | if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { |
| 1373 | convert_tensor<IlvdBlkLayout>(in, out, vec_A); |
| 1374 | } else if constexpr (ModeHasScales) { |
| 1375 | auto tCrS = cute::get<1>(partitioned_extra_info); |
| 1376 | auto converted_inputs = |
| 1377 | make_fragment_like<ElementScale>(tCrA_mma)(_, _, k_block); |
| 1378 | auto scales = tCrS(_, _, 0); |
| 1379 | |
| 1380 | // First, we upcast the inputs to the scale type |
| 1381 | convert_tensor<IlvdBlkLayout>(in, converted_inputs, vec_A); |
| 1382 | // Apply scales and broadcast across inputs, store in converted_inputs |
| 1383 | |
| 1384 | // We need to cast to nv_bfloat16 for the multiply since |
| 1385 | // `cutlass::bfloat16_t` has an overloaded operator* that upconverts to |
| 1386 | // float, which nvcc will not optimize to using vectorized fma |
| 1387 | // instructions (i.e. hfma.bf16_v2) |
| 1388 | if constexpr (std::is_same_v<ElementScale, cutlass::bfloat16_t>) { |
| 1389 | cute::transform( |
| 1390 | recast<nv_bfloat16>(converted_inputs), recast<nv_bfloat16>(scales), |
| 1391 | recast<nv_bfloat16>(converted_inputs), cute::multiplies{}); |
| 1392 | } else { |
| 1393 | cute::transform(converted_inputs, scales, converted_inputs, |
| 1394 | cute::multiplies{}); |
| 1395 | } |
| 1396 | |
| 1397 | // Apply zeros if required |
| 1398 | if constexpr (KernelConversionMode == |
| 1399 | ConversionMode::ConvertAndScaleWithZero) { |
| 1400 | auto tCrZ = cute::get<3>(partitioned_extra_info); |
| 1401 | auto converted_zeros = make_fragment_like<ElementScale>(tCrZ)(_, _, 0); |
| 1402 | |
| 1403 | convert_tensor<void>(tCrZ(_, _, 0), converted_zeros); |
| 1404 | if constexpr (std::is_same_v<ElementScale, cutlass::bfloat16_t>) { |
| 1405 | cute::transform(recast<nv_bfloat16>(converted_inputs), |
| 1406 | recast<nv_bfloat16>(converted_zeros), |
| 1407 | recast<nv_bfloat16>(converted_inputs), cute::plus{}); |
| 1408 | } else { |
| 1409 | cute::transform(converted_inputs, converted_zeros, converted_inputs, |
| 1410 | cute::plus{}); |
| 1411 | } |
| 1412 | } |
| 1413 | |
| 1414 | // Finally, we convert the scaled inputs to the mma type. |
| 1415 | convert_tensor<void>(converted_inputs, out); |
| 1416 | } else { |
| 1417 | static_assert(cutlass::detail::dependent_false<KernelSchedule>, |
| 1418 | "No A data is loaded."); |
| 1419 | } |
| 1420 | } |
| 1421 | |
| 1422 | // Modified from upstream, should be kept the same when possible |
| 1423 | // the main differences is that this version supports interleaved converts |
| 1424 | // Utilities for transforming the A operand prior to issuing tensorcore math. |
| 1425 | template <typename IlvdBlkLayout, class EngineIn, class EngineOut, |
| 1426 | class TensorLayout, |
| 1427 | int ConversionVectorWidth = cosize_v<TensorLayout>> |
| 1428 | CUTLASS_DEVICE void convert_tensor( |
| 1429 | Tensor<EngineIn, TensorLayout> const& in, |
| 1430 | Tensor<EngineOut, TensorLayout>& out, |
| 1431 | cute::Int<ConversionVectorWidth> width = {}) { |
| 1432 | // This is an element-wise conversion where we expect both tensors to have |
| 1433 | // the same layout. As a result, we can cast as a cutlass array to use the |
| 1434 | // fast numeric converters without worrying about indexing into the layout. |
| 1435 | constexpr int N = cosize_v<TensorLayout>; |
| 1436 | |
| 1437 | // The inputs must be backed by registers & be statically sized. |
| 1438 | static_assert(is_rmem<EngineIn>::value, |
| 1439 | "Input tensor for A conversion must come from registers"); |
| 1440 | static_assert(is_rmem<EngineOut>::value, |
| 1441 | "Output tensor for A conversion must come from registers"); |
| 1442 | static_assert(is_static_v<TensorLayout>, |
| 1443 | "Tensor layout for the conversion must be static"); |
| 1444 | static_assert(cosize_v<TensorLayout> == size(TensorLayout{}), |
| 1445 | "Cosize and size of the layout must be equal."); |
| 1446 | static_assert( |
| 1447 | N % ConversionVectorWidth == 0, |
| 1448 | "Conversion vector width must divide cosize of the tensor layout."); |
| 1449 | |
| 1450 | using SrcType = typename EngineIn::value_type; |
| 1451 | using DstType = typename EngineOut::value_type; |
| 1452 | |
| 1453 | using SrcArray = cutlass::Array<SrcType, ConversionVectorWidth>; |
| 1454 | using DstArray = cutlass::Array<DstType, ConversionVectorWidth>; |
| 1455 | |
| 1456 | constexpr cutlass::FloatRoundStyle RoundStyle = |
| 1457 | cutlass::FloatRoundStyle::round_to_nearest; |
| 1458 | |
| 1459 | using Converter = cutlass::InterleavedNumericArrayConverter< |
| 1460 | IlvdBlkLayout, DstType, SrcType, ConversionVectorWidth, RoundStyle>; |
| 1461 | |
| 1462 | constexpr int NumIterations = N / ConversionVectorWidth; |
| 1463 | |
| 1464 | for (int ii = 0; ii < NumIterations; ++ii) { |
| 1465 | SrcArray const* src_array_ptr = |
| 1466 | reinterpret_cast<SrcArray const*>(raw_pointer_cast(in.data())) + ii; |
| 1467 | DstArray* dst_array_ptr = |
| 1468 | reinterpret_cast<DstArray*>(raw_pointer_cast(out.data())) + ii; |
| 1469 | *dst_array_ptr = Converter::convert(*src_array_ptr); |
| 1470 | } |
| 1471 | } |
| 1472 | }; |
| 1473 | |
| 1474 | } // namespace machete |