diff --git a/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_bf16_gemm.cu b/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_bf16_gemm.cu index 9346734a..c1d3caeb 100644 --- a/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_bf16_gemm.cu +++ b/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_bf16_gemm.cu @@ -41,12 +41,15 @@ Note that in this example, we explicitly swap A and B in order to use TMA epilogues. We do this since TMA epilogues are more performant on problem sizes of interest. As an additional optimization, we can reorder the narrow data type tensor such that elements read into register file by the same thread are contiguous in global and shared memory. - This promotes vectorization of shared memory loads and removes additional instructions on the critical path. For example, when MMA is performed in FP8 data type, each thread reads + This promotes vectorization of shared memory loads and removes additional instructions on the critical path. For example, when MMA is performed in 16-bit data type, each thread reads 4 groups of 2 elements that are logically contiguous in the same row (refer to https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#wgmma-64n16-a for thread-value layout). If the narrow type is INT4 and tensor is major in K dim, only 8 bits can be read at a time, leading to extra load instructions and suboptimal utilization of shared memory throughput. If we reorder the data offline to place all 16 elements read by a thread contiguously in memory, a single 64-bit load is sufficient. This reordering is often feasible when the quantized - tensor is static (e.g. weight tensor of a NN layer at inference time). This example demonstrates how such a reordering can be performed and communicated to the kernel when the macro - OPTIMIZE_WEIGHT_LAYOUT is set to 1. + tensor is static (e.g. weight tensor of a NN layer at inference time). This example demonstrates how such a reordering can be performed and communicated to the kernel when the options.shuffle is set to true. + + Furthermore, the conversion from {INT4, UINT4} to {FP16, BF16} can benefit from pre-shuffling the weights in the order [0,2,4,6,1,3,5,7]. This allows multiple nibbles to be efficiently extracted and up-converted + in parallel. The reordering is enabled by defining the layout type `ValueShuffle`. Refer to the partial specializations of `NumericArrayShuffleConverter` in "include/cutlass/detail/collective/mixed_input_utils.hpp" + for more details. It is expected that the scale's K dimension be scale_k = ceil_div(problem_k, group_size). @@ -58,12 +61,11 @@ equal to the gemm problem K. Limitations: - 1) Only supports INT4 x { FP16, BF16 }. The scales must be the same as mma Type. Scale with zero-point mode is not supported. - 2) The INT4 weights have additional encoding requirements. - 3) The scales must be MN major. That means if A is scaled, it must be column major, but if B is scaled it must be row major. - 4) The scales must have the same layout and groupsize. - 5) The groupsize must be greater or equal to the tile shape k. - 6) Currently, TMA epilogues cannot be used when the narrow type is the B operand. This limitation arises because the implementation always swaps the + 1) The INT4 weights have additional encoding requirements. + 2) The scales must be MN major. That means if A is scaled, it must be column major, but if B is scaled it must be row major. + 3) The scales must have the same layout and groupsize. + 4) The groupsize must be greater or equal to the tile shape k. + 5) Currently, TMA epilogues cannot be used when the narrow type is the B operand. This limitation arises because the implementation always swaps the operands to ensure that the narrow type passes through the register file, and TMA epilogues do not currently support implicit swap + transpose operations. We plan to address this limitation in the future. However, we address this in the example by explicitly swapping and transposing the operands. @@ -103,14 +105,12 @@ #include "cutlass/util/reference/device/tensor_compare.h" #include "helper.h" -#include "unfused_weight_dequantize.hpp" +#include "mixed_dtype_utils.hpp" #include "packed_scale.hpp" #include "reorder_utils.hpp" using namespace cute; -#define OPTIMIZE_WEIGHT_LAYOUT 1 - #if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -137,16 +137,18 @@ using LayoutB_Transpose = typename cutlass::layout::LayoutTranspose::ty using StrideA = cutlass::detail::TagToStrideA_t; using StrideB = cutlass::detail::TagToStrideB_t; -#if OPTIMIZE_WEIGHT_LAYOUT // Define the CuTe layout for reoredered quantized tensor B // LayoutAtomQuant places values that will be read by the same thread in contiguous locations in global memory. // It specifies the reordering within a single warp's fragment -using LayoutAtomQuant = decltype(compute_memory_reordering_atom()); +//using ValueShuffle = Layout<_1>; // no value reordering +using ValueShuffle = Layout, Stride<_4,_1>>; // order [0,2,4,6,1,3,5,7] +int constexpr NumShuffleAtoms = 1; +using MmaAtomShape = Layout>>; +using LayoutAtomQuant = decltype(compute_memory_reordering_atom()); using LayoutB_Reordered = decltype(tile_to_shape(LayoutAtomQuant{}, Layout, StrideB>{})); -#endif using ElementScale = MmaType; -using ElementZero = ElementScale; // only for verify +using ElementZero = ElementScale; using LayoutScale = cutlass::layout::RowMajor; // C/D matrix configuration @@ -166,7 +168,7 @@ using ArchTag = cutlass::arch::Sm90; // T using OperatorClass = cutlass::arch::OpClassTensorOp; // Operator class tag using TileShape = Shape<_128,_128,cute::Int>; // Threadblock-level tile size using ClusterShape = Shape<_1,_1,_1>; // Shape of the threadblocks in a cluster -using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedCooperativeMixedInput; // Kernel to launch based on the default setting in the Collective Builder +using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedCooperative; // Kernel to launch based on the default setting in the Collective Builder using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative; using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto; @@ -183,15 +185,54 @@ using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBui EpilogueSchedule // This is the only epi supporting the required swap + transpose. >::CollectiveOp; +// ============================================================ MIXED INPUT NO SCALES ============================================================================ +// The collective will infer that the narrow type should be upcasted to the wide type. +// We swap A and B operands to the builder here +using CollectiveMainloopConvertOnly = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ElementB, LayoutB_Transpose, AlignmentB, + ElementA, LayoutA_Transpose, AlignmentA, + ElementAccumulator, + TileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage)) + >, + KernelSchedule + >::CollectiveOp; + +using GemmKernelConvertOnly = cutlass::gemm::kernel::GemmUniversal< + Shape, // Indicates ProblemShape + CollectiveMainloopConvertOnly, + CollectiveEpilogue +>; + +using GemmConvertOnly = cutlass::gemm::device::GemmUniversalAdapter; + +using CollectiveMainloopConvertOnlyShuffled = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ElementB, LayoutB_Reordered, AlignmentB, + ElementA, LayoutA_Transpose, AlignmentA, + ElementAccumulator, + TileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage)) + >, + KernelSchedule + >::CollectiveOp; + +using GemmKernelConvertOnlyShuffled = cutlass::gemm::kernel::GemmUniversal< + Shape, // Indicates ProblemShape + CollectiveMainloopConvertOnlyShuffled, + CollectiveEpilogue +>; + +using GemmConvertOnlyShuffled = cutlass::gemm::device::GemmUniversalAdapter; + // =========================================================== MIXED INPUT WITH SCALES =========================================================================== // The Scale information must get paired with the operand that will be scaled. In this example, B is scaled so we make a tuple of B's information and the scale information. using CollectiveMainloopScaleOnly = typename cutlass::gemm::collective::CollectiveBuilder< ArchTag, OperatorClass, -#if OPTIMIZE_WEIGHT_LAYOUT - cute::tuple, LayoutB_Reordered, AlignmentB, -#else cute::tuple, LayoutB_Transpose, AlignmentB, -#endif ElementA, LayoutA_Transpose, AlignmentA, ElementAccumulator, TileShape, ClusterShape, @@ -209,6 +250,69 @@ using GemmKernelScaleOnly = cutlass::gemm::kernel::GemmUniversal< using GemmScaleOnly = cutlass::gemm::device::GemmUniversalAdapter; +using CollectiveMainloopScaleOnlyShuffled = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + cute::tuple, LayoutB_Reordered, AlignmentB, + ElementA, LayoutA_Transpose, AlignmentA, + ElementAccumulator, + TileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage)) + >, + KernelSchedule + >::CollectiveOp; + +using GemmKernelScaleOnlyShuffled = cutlass::gemm::kernel::GemmUniversal< + Shape, // Indicates ProblemShape + CollectiveMainloopScaleOnlyShuffled, + CollectiveEpilogue +>; + +using GemmScaleOnlyShuffled = cutlass::gemm::device::GemmUniversalAdapter; + +// =========================================================== MIXED INPUT WITH SCALES AND ZEROS ================================================================== +// We specify scale + zero elements to indicate that we require both. Scales and biases have the same format. +using CollectiveMainloopScaleWithZeroPoint = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + cute::tuple, LayoutB_Transpose, AlignmentB, + ElementA, LayoutA_Transpose, AlignmentA, + ElementAccumulator, + TileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage)) + >, + KernelSchedule + >::CollectiveOp; + +using GemmKernelScaleWithZeroPoint = cutlass::gemm::kernel::GemmUniversal< + Shape, // Indicates ProblemShape + CollectiveMainloopScaleWithZeroPoint, + CollectiveEpilogue +>; + +using GemmScaleWithZeroPoint = cutlass::gemm::device::GemmUniversalAdapter; + +using CollectiveMainloopScaleWithZeroPointShuffled = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + cute::tuple, LayoutB_Reordered, AlignmentB, + ElementA, LayoutA_Transpose, AlignmentA, + ElementAccumulator, + TileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage)) + >, + KernelSchedule + >::CollectiveOp; + +using GemmKernelScaleWithZeroPointShuffled = cutlass::gemm::kernel::GemmUniversal< + Shape, // Indicates ProblemShape + CollectiveMainloopScaleWithZeroPointShuffled, + CollectiveEpilogue +>; + +using GemmScaleWithZeroPointShuffled = cutlass::gemm::device::GemmUniversalAdapter; +// ================================================================================================================================================================= + using StrideC = typename GemmKernelScaleOnly::StrideC; using StrideD = typename GemmKernelScaleOnly::StrideD; @@ -228,9 +332,7 @@ StrideD stride_D; StrideD_ref stride_D_ref; uint64_t seed; -#if OPTIMIZE_WEIGHT_LAYOUT LayoutB_Reordered layout_B_reordered; -#endif using StrideS = typename CollectiveMainloopScaleOnly::StrideScale; using StrideS_ref = cutlass::detail::TagToStrideB_t; @@ -253,41 +355,22 @@ cutlass::DeviceAllocationMixedDtypeOptions::parse(argc, args); } /// Prints the usage statement. std::ostream & print_usage(std::ostream &out) const { - out << "55_hopper_warp_specialized_gemm\n\n" - << " Hopper FP32 GEMM using a Warp Specialized kernel.\n\n" + out << "55_hopper_int4_bf16_gemm\n\n" + << " Hopper Mixed Data Type GEMM using a Warp Specialized kernel.\n\n" << "Options:\n\n" << " --help If specified, displays this usage statement\n\n" << " --m= Sets the M extent of the GEMM\n" @@ -295,36 +378,19 @@ struct Options { << " --k= Sets the K extent of the GEMM\n" << " --l= The number of independent gemm problems with mnk shape\n" << " --g= The size of each group for the scales. To broadcast a vector of scales or zeros, set the group size to K.\n" + << " --mode= The mode to run the gemm. 0 does (A @ B), 1 means A @ (scale * B), 2 means A @ (scale * B + zero-point).\n" << " --alpha= Epilogue scalar alpha\n" << " --beta= Epilogue scalar beta\n\n" - << " --iterations= Number of profiling iterations to perform.\n\n"; + << " --iterations= Number of profiling iterations to perform.\n\n" + << " --warmup= Number of warmup iterations to perform.\n\n" + << " --shuffle= Enable the offline layout swizzling.\n\n"; out << "\n\nExamples:\n\n" - << "$ " << "55_hopper_warp_specialized_gemm" << " --m=1024 --n=512 --k=1024 -g 0 --l=10 --alpha=2 --mode=2 --beta=0.707 \n\n"; + << "$ " << "55_hopper_int4_bf16_gemm" << " --m=1024 --n=512 --k=1024 -g=1024 --l=10 --alpha=2 --mode=2 --beta=0.707 \n\n"; return out; } - - /// Compute performance in GFLOP/s - double gflops(double runtime_s) const - { - // Two flops per multiply-add - uint64_t flop = uint64_t(2) * m * n * k * l; - double gflop = double(flop) / double(1.0e9); - return gflop / runtime_s; - } -}; - -/// Result structure -struct Result -{ - double avg_runtime_ms = 0.0; - double gflops = 0.0; - cutlass::Status status = cutlass::Status::kSuccess; - cudaError_t error = cudaSuccess; - bool passed = false; - }; #if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) @@ -333,78 +399,6 @@ struct Result /// GEMM setup and evaluation ///////////////////////////////////////////////////////////////////////////////////////////////// -/// Helper to initialize a block of device data -template -bool initialize_tensor( - cutlass::DeviceAllocation& block, - uint64_t seed=2023) { - - double scope_max, scope_min; - int bits_input = cutlass::sizeof_bits::value; - int bits_output = cutlass::sizeof_bits::value; - - if (bits_input == 1) { - scope_max = 2; - scope_min = 0; - } - else if (bits_input <= 8) { - scope_max = 2; - scope_min = -2; - } - else if (bits_output == 16) { - scope_max = 5; - scope_min = -5; - } - else { - scope_max = 8; - scope_min = -8; - } - cutlass::reference::device::BlockFillRandomUniform( - block.get(), block.size(), seed, Element(scope_max), Element(scope_min)); - - return true; -} - -template -bool initialize_quant_tensor( - cutlass::DeviceAllocation& block, - uint64_t seed=2023) { - - float scope_min = float(cutlass::platform::numeric_limits::lowest()); - float scope_max = float(cutlass::platform::numeric_limits::max()); - - cutlass::reference::device::BlockFillRandomUniform( - block.get(), block.size(), seed, Element(scope_max), Element(scope_min)); - - return true; -} - -template -bool initialize_scale( - cutlass::DeviceAllocation& block, - Options const& options) { - - float elt_max_f = float(cutlass::platform::numeric_limits::max()); - float const max_dequant_val = 4.f; - float const min_dequant_val = 0.5f; - - float scope_max(max_dequant_val / elt_max_f); - float scope_min(min_dequant_val / elt_max_f); - - cutlass::reference::device::BlockFillRandomUniform( - block.get(), block.size(), seed, Element(scope_max), Element(scope_min)); - return true; -} - -template -bool initialize_zero( - cutlass::DeviceAllocation& block, - Options const& options) { - std::vector stage(block.size(), Element(0.0f)); - block.copy_from_host(stage.data()); - return true; -} - /// Initialize operands to be used in the GEMM and reference GEMM void initialize(Options const& options) { @@ -448,33 +442,61 @@ void initialize(Options const& options) { dequantize_weight(block_B_dq.get(), block_B.get(), layout_B, block_scale.get(), block_zero.get(), layout_scale_zero, options.g); -#if OPTIMIZE_WEIGHT_LAYOUT - // Repeat the reorder layout atom to tile the whole tensor shape - layout_B_reordered = tile_to_shape(LayoutAtomQuant{}, shape_B); - reorder_tensor(block_B.get(), layout_B, layout_B_reordered); + if (options.shuffle) { + // Repeat the reorder layout atom to tile the whole tensor shape + layout_B_reordered = tile_to_shape(LayoutAtomQuant{}, shape_B); + reorder_tensor(block_B.get(), layout_B, layout_B_reordered); - print("Quantized tensor layout: "); - print(layout_B_reordered); - print("\n"); -#endif + print("Quantized tensor layout: "); + print(layout_B_reordered); + print("\n"); + } } /// Populates a Gemm::Arguments structure from the given commandline options -template -Args args_from_options(Options const& options) +/// Swap the A and B tensors, as well as problem shapes here. +template +typename Gemm::Arguments args_from_options(Options const& options) { -// Swap the A and B tensors, as well as problem shapes here. - - return Args { - cutlass::gemm::GemmUniversalMode::kGemm, - {options.n, options.m, options.k, options.l}, -#if OPTIMIZE_WEIGHT_LAYOUT - {block_B.get(), layout_B_reordered, block_A.get(), stride_A, block_scale.get(), stride_S, options.g}, -#else - {block_B.get(), stride_B, block_A.get(), stride_A, block_scale.get(), stride_S, options.g}, -#endif - {{options.alpha, options.beta}, block_C.get(), stride_C, block_D.get(), stride_D} - }; + using Args = typename Gemm::Arguments; + auto&& dB = [&]() { + if constexpr (cute::is_same_v || + cute::is_same_v || + cute::is_same_v) { + // offline swizzling is enabled. + return layout_B_reordered; + } + else { + return stride_B; + } + }(); + if (options.mode == MixedDtypeGemmMode::ConvertOnly) { + return Args { + cutlass::gemm::GemmUniversalMode::kGemm, + {options.n, options.m, options.k, options.l}, + {block_B.get(), dB, block_A.get(), stride_A}, + {{options.alpha, options.beta}, block_C.get(), stride_C, block_D.get(), stride_D} + }; + } + else if (options.mode == MixedDtypeGemmMode::ScaleOnly) { + return Args { + cutlass::gemm::GemmUniversalMode::kGemm, + {options.n, options.m, options.k, options.l}, + {block_B.get(), dB, block_A.get(), stride_A, block_scale.get(), stride_S, options.g}, + {{options.alpha, options.beta}, block_C.get(), stride_C, block_D.get(), stride_D} + }; + } + else if (options.mode == MixedDtypeGemmMode::ScaleWithZeroPoint) { + return Args { + cutlass::gemm::GemmUniversalMode::kGemm, + {options.n, options.m, options.k, options.l}, + {block_B.get(), dB, block_A.get(), stride_A, block_scale.get(), stride_S, options.g, block_zero.get()}, + {{options.alpha, options.beta}, block_C.get(), stride_C, block_D.get(), stride_D} + }; + } else { + std::cerr << "Invalid mode " << options.mode << ". Must be 0, 1 or 2." << std::endl; + exit(-1); + } } bool verify(Options const& options) { @@ -543,7 +565,7 @@ int run(Options &options) Gemm gemm; // Create a structure of gemm kernel arguments suitable for invoking an instance of Gemm - auto arguments = args_from_options(options); + auto arguments = args_from_options(options); // Using the arguments, query for extra workspace required for matrix multiplication computation size_t workspace_size = Gemm::get_workspace_size(arguments); @@ -561,35 +583,14 @@ int run(Options &options) CUTLASS_CHECK(gemm.run()); // Check if output from CUTLASS kernel and reference kernel are equal or not - Result result; + MixedDtypeResult result; result.passed = verify(options); - + mixed_dtype_profiling(gemm, options, result); std::cout << " Disposition: " << (result.passed ? "Passed" : "Failed") << std::endl; - if (!result.passed) { exit(-1); } - // Run profiling loop - if (options.iterations > 0) - { - GpuTimer timer; - timer.start(); - for (int iter = 0; iter < options.iterations; ++iter) { - CUTLASS_CHECK(gemm.run()); - } - timer.stop(); - - // Compute average runtime and GFLOPs. - float elapsed_ms = timer.elapsed_millis(); - result.avg_runtime_ms = double(elapsed_ms) / double(options.iterations); - result.gflops = options.gflops(result.avg_runtime_ms / 1000.0); - - std::cout << " Problem Size: " << options.m << 'x' << options.n << 'x' << options.k << 'x' << options.l << std::endl; - std::cout << " Avg runtime: " << result.avg_runtime_ms << " ms" << std::endl; - std::cout << " GFLOPS: " << result.gflops << std::endl; - } - return 0; } @@ -618,12 +619,6 @@ int main(int argc, char const **args) { << "later (compute capability 90 or greater).\n"; return 0; } - // {$nv-internal-release begin} - else if (props.major != 9 || props.minor != 0) { - std::cerr << "This example requires a GPU of NVIDIA's Hopper Architecture (compute capability 90).\n"; - return 0; - } - // {$nv-internal-release end} // // Parse options @@ -643,12 +638,44 @@ int main(int argc, char const **args) { // #if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) - if (options.g == options.k) { - std::cout << "Running in per-column scale mode." << std::endl; - } else { - std::cout << "Running in group scale mode." << std::endl; + if (options.mode == MixedDtypeGemmMode::ConvertOnly) { + std::cout << "Running in no scale mode." << std::endl; + if (options.shuffle) { + std::cout << "Offline shuffle enabled." << std::endl; + run(options); + } else { + std::cout << "Offline shuffle disabled." << std::endl; + run(options); + } + } + else if (options.mode == MixedDtypeGemmMode::ScaleOnly) { + if (options.g == options.k) { + std::cout << "Running in per-column scale mode." << std::endl; + } else { + std::cout << "Running in group scale mode." << std::endl; + } + if (options.shuffle) { + std::cout << "Offline shuffle enabled." << std::endl; + run(options); + } else { + std::cout << "Offline shuffle disabled." << std::endl; + run(options); + } + } + else if (options.mode == MixedDtypeGemmMode::ScaleWithZeroPoint) { + if (options.g == options.k) { + std::cout << "Running in per-column scale and zero mode." << std::endl; + } else { + std::cout << "Running in group scale and zero mode." << std::endl; + } + if (options.shuffle) { + std::cout << "Offline shuffle enabled." << std::endl; + run(options); + } else { + std::cout << "Offline shuffle disabled." << std::endl; + run(options); + } } - run(options); #endif return 0; diff --git a/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_fp8_gemm.cu b/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_fp8_gemm.cu index eee10e01..5fc96d4e 100644 --- a/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_fp8_gemm.cu +++ b/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_fp8_gemm.cu @@ -36,7 +36,7 @@ between INT4 and FP8. To trigger this method, use cutlass::Array as the scale type in the collective's arguments. However, this algorithm requires changes to the encoding of INT4 weights and scale factors. These changes must happen before launching the GEMM. See the helper functions - `unify_quant_encoding`, `initialize_packed_scale`, and header `fp8_packed_scale.hpp` for details. + `unify_quant_encoding`, `initialize_packed_scale` in the header `fp8_packed_scale.hpp` for details. In a nutshell, the positive values of INT4 weights need to be encoded in the same way as negative values except for the sign bit. For each scale factor, 8 negative results (-8 x scale, -7 x scale, ... -1 x scale) are packed together, forming a cutlass::Array value. @@ -52,8 +52,7 @@ 4 groups of 4 elements that are logically contiguous in the same row (refer to https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#wgmma-64n32-a for thread-value layout). If the narrow type is INT4 and tensor is major in K dim, only 16 bits can be read at a time, leading to extra load instructions and suboptimal utilization of shared memory throughput. If we reorder the data offline to place all 16 elements read by a thread contiguously in memory, a single 64-bit load is sufficient. This reordering is often feasible when the quantized - tensor is static (e.g. weight tensor of a NN layer at inference time). This example demonstrates how such a reordering can be performed and communicated to the kernel when the macro - OPTIMIZE_WEIGHT_LAYOUT is set to 1. + tensor is static (e.g. weight tensor of a NN layer at inference time). This example demonstrates how such a reordering can be performed and communicated to the kernel when the options.shuffle is set to true. It is expected that the scale's K dimension be scale_k = ceil_div(problem_k, group_size). @@ -110,14 +109,12 @@ #include "cutlass/util/reference/device/tensor_compare.h" #include "helper.h" -#include "unfused_weight_dequantize.hpp" +#include "mixed_dtype_utils.hpp" #include "packed_scale.hpp" #include "reorder_utils.hpp" using namespace cute; -#define OPTIMIZE_WEIGHT_LAYOUT 1 - #if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -144,13 +141,11 @@ using LayoutB_Transpose = typename cutlass::layout::LayoutTranspose::ty using StrideA = cutlass::detail::TagToStrideA_t; using StrideB = cutlass::detail::TagToStrideB_t; -#if OPTIMIZE_WEIGHT_LAYOUT // Define the CuTe layout for reoredered quantized tensor B // LayoutAtomQuant places values that will be read by the same thread in contiguous locations in global memory. // It specifies the reordering within a single warp's fragment using LayoutAtomQuant = decltype(compute_memory_reordering_atom()); using LayoutB_Reordered = decltype(tile_to_shape(LayoutAtomQuant{}, Layout, StrideB>{})); -#endif using ElementScale = MmaType; using ElementZero = ElementScale; // only for verify @@ -173,7 +168,7 @@ using ArchTag = cutlass::arch::Sm90; // T using OperatorClass = cutlass::arch::OpClassTensorOp; // Operator class tag using TileShape = Shape<_128,_128,cute::Int>; // Threadblock-level tile size using ClusterShape = Shape<_1,_1,_1>; // Shape of the threadblocks in a cluster -using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedCooperativeMixedInput; // Kernel to launch based on the default setting in the Collective Builder +using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedCooperative; // Kernel to launch based on the default setting in the Collective Builder using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative; using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto; @@ -194,11 +189,7 @@ using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBui // The Scale information must get paired with the operand that will be scaled. In this example, B is scaled so we make a tuple of B's information and the scale information. using CollectiveMainloopScaleOnly = typename cutlass::gemm::collective::CollectiveBuilder< ArchTag, OperatorClass, -#if OPTIMIZE_WEIGHT_LAYOUT - cute::tuple>, LayoutB_Reordered, AlignmentB, -#else cute::tuple>, LayoutB_Transpose, AlignmentB, -#endif ElementA, LayoutA_Transpose, AlignmentA, ElementAccumulator, TileShape, ClusterShape, @@ -214,7 +205,26 @@ using GemmKernelScaleOnly = cutlass::gemm::kernel::GemmUniversal< CollectiveEpilogue >; +using CollectiveMainloopShuffled = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + cute::tuple>, LayoutB_Reordered, AlignmentB, + ElementA, LayoutA_Transpose, AlignmentA, + ElementAccumulator, + TileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage)) + >, + KernelSchedule + >::CollectiveOp; + +using GemmKernelShuffled = cutlass::gemm::kernel::GemmUniversal< + Shape, // Indicates ProblemShape + CollectiveMainloopShuffled, + CollectiveEpilogue +>; + using GemmScaleOnly = cutlass::gemm::device::GemmUniversalAdapter; +using GemmShuffled = cutlass::gemm::device::GemmUniversalAdapter; using StrideC = typename GemmKernelScaleOnly::StrideC; using StrideD = typename GemmKernelScaleOnly::StrideD; @@ -235,9 +245,7 @@ StrideD stride_D; StrideD_ref stride_D_ref; uint64_t seed; -#if OPTIMIZE_WEIGHT_LAYOUT LayoutB_Reordered layout_B_reordered; -#endif using StrideS = typename CollectiveMainloopScaleOnly::StrideScale; using StrideS_ref = cutlass::detail::TagToStrideB_t; @@ -262,41 +270,24 @@ cutlass::DeviceAllocationMixedDtypeOptions::parse(argc, args); - cmd.get_cmd_line_argument("m", m); - cmd.get_cmd_line_argument("n", n); - cmd.get_cmd_line_argument("k", k); - cmd.get_cmd_line_argument("l", l); - cmd.get_cmd_line_argument("g", g); - cmd.get_cmd_line_argument("alpha", alpha, 1.f); - cmd.get_cmd_line_argument("beta", beta, 0.f); - cmd.get_cmd_line_argument("iterations", iterations); + mode = 1; // override the mode value to always be scale only mode } /// Prints the usage statement. std::ostream & print_usage(std::ostream &out) const { - out << "55_hopper_warp_specialized_gemm\n\n" - << " Hopper FP32 GEMM using a Warp Specialized kernel.\n\n" + out << "55_hopper_int4_fp8_gemm\n\n" + << " Hopper Mixed Data Type GEMM using a Warp Specialized kernel.\n\n" << "Options:\n\n" << " --help If specified, displays this usage statement\n\n" << " --m= Sets the M extent of the GEMM\n" @@ -306,34 +297,16 @@ struct Options { << " --g= The size of each group for the scales. To broadcast a vector of scales or zeros, set the group size to K.\n" << " --alpha= Epilogue scalar alpha\n" << " --beta= Epilogue scalar beta\n\n" - << " --iterations= Number of profiling iterations to perform.\n\n"; + << " --iterations= Number of profiling iterations to perform.\n\n" + << " --warmup= Number of warmup iterations to perform.\n\n" + << " --shuffle= Enable the offline layout swizzling.\n\n"; out << "\n\nExamples:\n\n" - << "$ " << "55_hopper_warp_specialized_gemm" << " --m=1024 --n=512 --k=1024 -g 0 --l=10 --alpha=2 --mode=2 --beta=0.707 \n\n"; + << "$ " << "55_hopper_int4_fp8_gemm" << " --m=1024 --n=512 --k=1024 -g=1024 --l=10 --alpha=2 --beta=0.707 \n\n"; return out; } - - /// Compute performance in GFLOP/s - double gflops(double runtime_s) const - { - // Two flops per multiply-add - uint64_t flop = uint64_t(2) * m * n * k * l; - double gflop = double(flop) / double(1.0e9); - return gflop / runtime_s; - } -}; - -/// Result structure -struct Result -{ - double avg_runtime_ms = 0.0; - double gflops = 0.0; - cutlass::Status status = cutlass::Status::kSuccess; - cudaError_t error = cudaSuccess; - bool passed = false; - }; #if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) @@ -342,150 +315,6 @@ struct Result /// GEMM setup and evaluation ///////////////////////////////////////////////////////////////////////////////////////////////// -/// Helper to initialize a block of device data -template -bool initialize_tensor( - cutlass::DeviceAllocation& block, - uint64_t seed=2023) { - - double scope_max, scope_min; - int bits_input = cutlass::sizeof_bits::value; - int bits_output = cutlass::sizeof_bits::value; - - if (bits_input == 1) { - scope_max = 2; - scope_min = 0; - } - else if (bits_input <= 8) { - scope_max = 2; - scope_min = -2; - } - else if (bits_output == 16) { - scope_max = 5; - scope_min = -5; - } - else { - scope_max = 8; - scope_min = -8; - } - cutlass::reference::device::BlockFillRandomUniform( - block.get(), block.size(), seed, Element(scope_max), Element(scope_min)); - - return true; -} - -template -bool initialize_quant_tensor( - cutlass::DeviceAllocation& block, - uint64_t seed=2023) { - - float scope_min = float(cutlass::platform::numeric_limits::lowest()); - float scope_max = float(cutlass::platform::numeric_limits::max()); - - cutlass::reference::device::BlockFillRandomUniform( - block.get(), block.size(), seed, Element(scope_max), Element(scope_min)); - - return true; -} - -// In the mainloop, PRMT selects 1 byte from only 8 bytes so the sign bit is handled in an extra PRMT. -// Here the encodings of positive values and negative values are unified (except for the sign bit). -// For instance, 1 becomes 0b0111, which is the same encoding as -1 (0b1111). -bool unify_quant_encoding( - cutlass::DeviceAllocation const& block_in, - cutlass::DeviceAllocation& block_out) { - - using StorageType = cutlass::int4b_t::Storage; - - if (block_in.size() != block_out.size()) { - std::cerr << "block_in and block_out must have same size.\n"; - return false; - } - constexpr int pack = sizeof_bits_v / 4; - std::vector data(block_in.size() / pack); - cutlass::device_memory::copy_to_host(data.data(), (StorageType*)block_in.get(), block_in.size() / pack); - - for (auto&& d : data) { - StorageType out = 0; - StorageType mask = 0x0f; - for (int i = 0; i < pack; ++i) { - cutlass::int4b_t curr; - curr.storage = (d >> (i * 4)) & 0x0f; - switch (curr) { - case 1: curr.storage = StorageType(0b0111); break; // 2's complement - case 2: curr.storage = StorageType(0b0110); break; // 2's complement - case 3: curr.storage = StorageType(0b0101); break; // 2's complement - case 4: curr.storage = StorageType(0b0100); break; // 2's complement - case 5: curr.storage = StorageType(0b0011); break; // 2's complement - case 6: curr.storage = StorageType(0b0010); break; // 2's complement - case 7: curr.storage = StorageType(0b0001); break; // 2's complement - default: break; - } - out |= (curr.storage << (4 * i)) & mask; - mask <<= 4; - } - d = out; - } - - cutlass::device_memory::copy_to_device((StorageType*)block_out.get(), data.data(), block_out.size() / pack); - return true; -} - -template -bool initialize_scale( - cutlass::DeviceAllocation& block, - Options const& options) { - - float elt_max_f = float(cutlass::platform::numeric_limits::max()); - float const max_dequant_val = 4.f; - float const min_dequant_val = 0.5f; - - float scope_max(max_dequant_val / elt_max_f); - float scope_min(min_dequant_val / elt_max_f); - - cutlass::reference::device::BlockFillRandomUniform( - block.get(), block.size(), seed, Element(scope_max), Element(scope_min)); - return true; -} - -bool initialize_packed_scale( - cutlass::DeviceAllocation const& block_in, - cutlass::DeviceAllocation > & block_out) { - - std::vector data_in(block_in.size()); - std::vector > data_out(block_in.size()); - try { - block_in.copy_to_host(data_in.data()); - } catch (cutlass::cuda_exception const& e) - { - std::cerr << "CUDA Error: " << cudaGetErrorString(e.cudaError()) << std::endl; - return false; - } - for (size_t i = 0; i < block_in.size(); ++i) - { - cutlass::packed_scale_t tmp(data_in[i]); - data_out[i] = reinterpret_cast const&>(tmp); - // std::cout << data_in[i] << ":" << std::hex << static_cast(data_in[i].storage) << ",\t" << -data_in[i] << ":" << std::hex << static_cast((-data_in[i]).storage) << std::endl; - } - try { - block_out.copy_from_host(data_out.data()); - } catch (cutlass::cuda_exception const& e) - { - std::cerr << "CUDA Error: " << cudaGetErrorString(e.cudaError()) << std::endl; - return false; - } - return true; -} - -template -bool initialize_zero( - cutlass::DeviceAllocation& block, - Options const& options) { - std::vector stage(block.size(), Element(0.0f)); - block.copy_from_host(stage.data()); - return true; -} - /// Initialize operands to be used in the GEMM and reference GEMM void initialize(Options const& options) { @@ -533,31 +362,35 @@ void initialize(Options const& options) { dequantize_weight(block_B_dq.get(), block_B.get(), layout_B, block_scale.get(), block_zero.get(), layout_scale_zero, options.g); - #if OPTIMIZE_WEIGHT_LAYOUT - // Repeat the reorder layout atom to tile the whole tensor shape - layout_B_reordered = tile_to_shape(LayoutAtomQuant{}, shape_B); - reorder_tensor(block_B_modified.get(), layout_B, layout_B_reordered); + if (options.shuffle) { + // Repeat the reorder layout atom to tile the whole tensor shape + layout_B_reordered = tile_to_shape(LayoutAtomQuant{}, shape_B); + reorder_tensor(block_B_modified.get(), layout_B, layout_B_reordered); - print("Quantized tensor layout: "); - print(layout_B_reordered); - print("\n"); -#endif + print("Quantized tensor layout: "); + print(layout_B_reordered); + print("\n"); + } } /// Populates a Gemm::Arguments structure from the given commandline options -template -Args args_from_options(Options const& options) +/// Swap the A and B tensors, as well as problem shapes here. +template +typename Gemm::Arguments args_from_options(Options const& options) { -// Swap the A and B tensors, as well as problem shapes here. - + using Args = typename Gemm::Arguments; + auto&& dB = [&]() { + if constexpr (cute::is_same_v) { // offline swizzling is enabled. + return layout_B_reordered; + } + else { + return stride_B; + } + }(); return Args { cutlass::gemm::GemmUniversalMode::kGemm, {options.n, options.m, options.k, options.l}, -#if OPTIMIZE_WEIGHT_LAYOUT - {block_B_modified.get(), layout_B_reordered, block_A.get(), stride_A, block_scale_packed.get(), stride_S, options.g}, -#else - {block_B_modified.get(), stride_B, block_A.get(), stride_A, block_scale_packed.get(), stride_S, options.g}, -#endif + {block_B_modified.get(), dB, block_A.get(), stride_A, block_scale_packed.get(), stride_S, options.g}, {{options.alpha, options.beta}, block_C.get(), stride_C, block_D.get(), stride_D} }; } @@ -637,7 +470,7 @@ int run(Options &options) Gemm gemm; // Create a structure of gemm kernel arguments suitable for invoking an instance of Gemm - auto arguments = args_from_options(options); + auto arguments = args_from_options(options); // Using the arguments, query for extra workspace required for matrix multiplication computation size_t workspace_size = Gemm::get_workspace_size(arguments); @@ -655,35 +488,14 @@ int run(Options &options) CUTLASS_CHECK(gemm.run()); // Check if output from CUTLASS kernel and reference kernel are equal or not - Result result; + MixedDtypeResult result; result.passed = verify(options); - + mixed_dtype_profiling(gemm, options, result); std::cout << " Disposition: " << (result.passed ? "Passed" : "Failed") << std::endl; - if (!result.passed) { exit(-1); } - // Run profiling loop - if (options.iterations > 0) - { - GpuTimer timer; - timer.start(); - for (int iter = 0; iter < options.iterations; ++iter) { - CUTLASS_CHECK(gemm.run()); - } - timer.stop(); - - // Compute average runtime and GFLOPs. - float elapsed_ms = timer.elapsed_millis(); - result.avg_runtime_ms = double(elapsed_ms) / double(options.iterations); - result.gflops = options.gflops(result.avg_runtime_ms / 1000.0); - - std::cout << " Problem Size: " << options.m << 'x' << options.n << 'x' << options.k << 'x' << options.l << std::endl; - std::cout << " Avg runtime: " << result.avg_runtime_ms << " ms" << std::endl; - std::cout << " GFLOPS: " << result.gflops << std::endl; - } - return 0; } @@ -735,10 +547,16 @@ int main(int argc, char const **args) { } else { std::cout << "Running in group scale mode." << std::endl; } - run(options); + if (options.shuffle) { + std::cout << "Offline shuffle enabled." << std::endl; + run(options); + } else { + std::cout << "Offline shuffle disabled." << std::endl; + run(options); + } #endif return 0; } -///////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////////////////////////////////////////////////// \ No newline at end of file diff --git a/examples/55_hopper_mixed_dtype_gemm/55_hopper_mixed_dtype_gemm.cu b/examples/55_hopper_mixed_dtype_gemm/55_hopper_mixed_dtype_gemm.cu index 6a3a8f80..b482d0d1 100644 --- a/examples/55_hopper_mixed_dtype_gemm/55_hopper_mixed_dtype_gemm.cu +++ b/examples/55_hopper_mixed_dtype_gemm/55_hopper_mixed_dtype_gemm.cu @@ -102,19 +102,12 @@ #include "cutlass/util/reference/device/tensor_compare.h" #include "helper.h" -#include "unfused_weight_dequantize.hpp" +#include "mixed_dtype_utils.hpp" using namespace cute; #if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) -// This is just an example, so we use a regular enum so we can compare directly to the command-line int. -enum GemmMode { - ConvertOnly, - ScaleOnly, - ScaleWithZeroPoint -}; - ///////////////////////////////////////////////////////////////////////////////////////////////// /// GEMM kernel configurations ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -157,7 +150,7 @@ using ArchTag = cutlass::arch::Sm90; // T using OperatorClass = cutlass::arch::OpClassTensorOp; // Operator class tag using TileShape = Shape<_128,_128,cute::Int>; // Threadblock-level tile size using ClusterShape = Shape<_1,_1,_1>; // Shape of the threadblocks in a cluster -using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedCooperativeMixedInput; // Kernel to launch based on the default setting in the Collective Builder +using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedCooperative; // Kernel to launch based on the default setting in the Collective Builder using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative; using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto; @@ -284,178 +277,14 @@ cutlass::DeviceAllocation Sets the M extent of the GEMM\n" - << " --n= Sets the N extent of the GEMM\n" - << " --k= Sets the K extent of the GEMM\n" - << " --l= The number of independent gemm problems with mnk shape\n" - << " --g= The size of each group for the scales and zeros. To broadcast a vector of scales or zeros, set the group size to K.\n" - << " --mode= The mode to run the gemm. 0 does (A @ B), 1 means A @ (scale * B), 2 means A @ (scale * B + zero-point).\n" - << " --alpha= Epilogue scalar alpha\n" - << " --beta= Epilogue scalar beta\n\n" - << " --iterations= Number of profiling iterations to perform.\n\n"; - - out - << "\n\nExamples:\n\n" - << "$ " << "55_hopper_warp_specialized_gemm" << " --m=1024 --n=512 --k=1024 -g 0 --l=10 --alpha=2 --mode=2 --beta=0.707 \n\n"; - - return out; - } - - /// Compute performance in GFLOP/s - double gflops(double runtime_s) const - { - // Two flops per multiply-add - uint64_t flop = uint64_t(2) * m * n * k * l; - double gflop = double(flop) / double(1.0e9); - return gflop / runtime_s; - } -}; - -/// Result structure -struct Result -{ - double avg_runtime_ms = 0.0; - double gflops = 0.0; - cutlass::Status status = cutlass::Status::kSuccess; - cudaError_t error = cudaSuccess; - bool passed = false; - -}; - #if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) ///////////////////////////////////////////////////////////////////////////////////////////////// /// GEMM setup and evaluation ///////////////////////////////////////////////////////////////////////////////////////////////// -/// Helper to initialize a block of device data -template -bool initialize_tensor( - cutlass::DeviceAllocation& block, - uint64_t seed=2023) { - - double scope_max, scope_min; - int bits_input = cutlass::sizeof_bits::value; - int bits_output = cutlass::sizeof_bits::value; - - if (bits_input == 1) { - scope_max = 2; - scope_min = 0; - } - else if (bits_input <= 8) { - scope_max = 2; - scope_min = -2; - } - else if (bits_output == 16) { - scope_max = 5; - scope_min = -5; - } - else { - scope_max = 8; - scope_min = -8; - } - cutlass::reference::device::BlockFillRandomUniform( - block.get(), block.size(), seed, Element(scope_max), Element(scope_min)); - - return true; -} - -template -bool initialize_quant_tensor( - cutlass::DeviceAllocation& block, - uint64_t seed=2023) { - - float scope_min = float(cutlass::platform::numeric_limits::lowest()); - float scope_max = float(cutlass::platform::numeric_limits::max()); - - cutlass::reference::device::BlockFillRandomUniform( - block.get(), block.size(), seed, Element(scope_max), Element(scope_min)); - - return true; -} - -template -bool initialize_scale( - cutlass::DeviceAllocation& block, - Options const& options) { - - if (options.mode == GemmMode::ConvertOnly) { - // No scales, so just initialize with 1 so we can use the same kernel to dequantize the data. - std::vector stage(block.size(), Element(1.0f)); - block.copy_from_host(stage.data()); - } - else { - float elt_max_f = float(cutlass::platform::numeric_limits::max()); - const float max_dequant_val = 4.f; - const float min_dequant_val = 0.5f; - - float scope_max(max_dequant_val / elt_max_f); - float scope_min(min_dequant_val / elt_max_f); - - cutlass::reference::device::BlockFillRandomUniform( - block.get(), block.size(), seed, Element(scope_max), Element(scope_min)); - } - return true; -} - -template -bool initialize_zero( - cutlass::DeviceAllocation& block, - Options const& options) { - - if (options.mode == GemmMode::ScaleWithZeroPoint) { - cutlass::reference::device::BlockFillRandomUniform( - block.get(), block.size(), seed, Element(2.0f), Element(-2.0f)); - } else { - // No bias, so just initialize with 1 so we can use the same kernel to dequantize the data. - std::vector stage(block.size(), Element(0.0f)); - block.copy_from_host(stage.data()); - } - return true; -} - /// Initialize operands to be used in the GEMM and reference GEMM -void initialize(Options const& options) { +void initialize(MixedDtypeOptions const& options) { auto shape_b = cute::make_shape(options.n, options.k, options.l); int const scale_k = (options.k + options.g - 1) / options.g; @@ -500,10 +329,10 @@ void initialize(Options const& options) { /// Populates a Gemm::Arguments structure from the given commandline options template -Args args_from_options(Options const& options) +Args args_from_options(MixedDtypeOptions const& options) { // Swap the A and B tensors, as well as problem shapes here. - if (options.mode == GemmMode::ConvertOnly) { + if (options.mode == MixedDtypeGemmMode::ConvertOnly) { return Args { cutlass::gemm::GemmUniversalMode::kGemm, {options.n, options.m, options.k, options.l}, @@ -511,7 +340,7 @@ Args args_from_options(Options const& options) {{options.alpha, options.beta}, block_C.get(), stride_C, block_D.get(), stride_D} }; } - else if (options.mode == GemmMode::ScaleOnly) { + else if (options.mode == MixedDtypeGemmMode::ScaleOnly) { return Args { cutlass::gemm::GemmUniversalMode::kGemm, {options.n, options.m, options.k, options.l}, @@ -519,7 +348,7 @@ Args args_from_options(Options const& options) {{options.alpha, options.beta}, block_C.get(), stride_C, block_D.get(), stride_D} }; } - else if (options.mode == GemmMode::ScaleWithZeroPoint) { + else if (options.mode == MixedDtypeGemmMode::ScaleWithZeroPoint) { return Args { cutlass::gemm::GemmUniversalMode::kGemm, {options.n, options.m, options.k, options.l}, @@ -532,7 +361,7 @@ Args args_from_options(Options const& options) } } -bool verify(const Options &options) { +bool verify(MixedDtypeOptions const& options) { // // Compute reference output // @@ -598,7 +427,7 @@ bool verify(const Options &options) { /// Execute a given example GEMM computation template -int run(Options &options) +int run(MixedDtypeOptions &options) { initialize(options); @@ -624,35 +453,14 @@ int run(Options &options) CUTLASS_CHECK(gemm.run()); // Check if output from CUTLASS kernel and reference kernel are equal or not - Result result; + MixedDtypeResult result; result.passed = verify(options); - + mixed_dtype_profiling(gemm, options, result); std::cout << " Disposition: " << (result.passed ? "Passed" : "Failed") << std::endl; - if (!result.passed) { exit(-1); } - // Run profiling loop - if (options.iterations > 0) - { - GpuTimer timer; - timer.start(); - for (int iter = 0; iter < options.iterations; ++iter) { - CUTLASS_CHECK(gemm.run()); - } - timer.stop(); - - // Compute average runtime and GFLOPs. - float elapsed_ms = timer.elapsed_millis(); - result.avg_runtime_ms = double(elapsed_ms) / double(options.iterations); - result.gflops = options.gflops(result.avg_runtime_ms / 1000.0); - - std::cout << " Problem Size: " << options.m << 'x' << options.n << 'x' << options.k << 'x' << options.l << std::endl; - std::cout << " Avg runtime: " << result.avg_runtime_ms << " ms" << std::endl; - std::cout << " GFLOPS: " << result.gflops << std::endl; - } - return 0; } @@ -685,7 +493,7 @@ int main(int argc, char const **args) { // Parse options // - Options options; + MixedDtypeOptions options; options.parse(argc, args); @@ -699,11 +507,11 @@ int main(int argc, char const **args) { // #if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) - if (options.mode == GemmMode::ConvertOnly) { + if (options.mode == MixedDtypeGemmMode::ConvertOnly) { std::cout << "Running in no scale mode." << std::endl; run(options); } - else if (options.mode == GemmMode::ScaleOnly) { + else if (options.mode == MixedDtypeGemmMode::ScaleOnly) { if (options.g == options.k) { std::cout << "Running in per-column scale mode." << std::endl; } else { @@ -711,7 +519,7 @@ int main(int argc, char const **args) { } run(options); } - else if (options.mode == GemmMode::ScaleWithZeroPoint) { + else if (options.mode == MixedDtypeGemmMode::ScaleWithZeroPoint) { if (options.g == options.k) { std::cout << "Running in per-column scale and zero mode." << std::endl; } else { @@ -724,4 +532,4 @@ int main(int argc, char const **args) { return 0; } -///////////////////////////////////////////////////////////////////////////////////////////////// \ No newline at end of file +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/examples/55_hopper_mixed_dtype_gemm/README.md b/examples/55_hopper_mixed_dtype_gemm/README.md index 07265f0d..8d22c75f 100644 --- a/examples/55_hopper_mixed_dtype_gemm/README.md +++ b/examples/55_hopper_mixed_dtype_gemm/README.md @@ -3,7 +3,7 @@ This example shows how to do mixed types GEMMs in CUTLASS. ## High level overview This example shows how to perform GEMMs on Hopper when A and B have different types. This implementation always passes the type with fewer bits through the register file and upcasts to the type with the higher bit count. -When relying on `KernelScheduleAuto`, the main loop supporting different A and B types will be selected whenever the bit count of A is not equal to the bit count of B. Users can manually select the mixed type main loop and explicitly choose the scheduling policy by specifying one of the following schedules to the `CollectiveBuilder`: `KernelTmaWarpSpecializedMixedInput`, `KernelTmaWarpSpecializedPingpongMixedInput` or `KernelTmaWarpSpecializedCooperativeMixedInput`. +When relying on `KernelScheduleAuto`, the main loop supporting different A and B types will be selected whenever the bit count of A is not equal to the bit count of B. Users can manually select the mixed type main loop and explicitly choose the scheduling policy by specifying one of the following schedules to the `CollectiveBuilder`: `KernelTmaWarpSpecialized`, `KernelTmaWarpSpecializedPingpong` or `KernelTmaWarpSpecializedCooperative`. This first version only supports mixed type GEMMs using TMA. @@ -36,4 +36,4 @@ We are currently optimizing the following cases: * Optimizations for memory bound cases. -* Optimizations for scale and zero-point loading when the group size is not equal to the threadblock-k size. +* Optimizations for scale and zero-point loading when the group size is not equal to the threadblock-k size. \ No newline at end of file diff --git a/examples/55_hopper_mixed_dtype_gemm/mixed_dtype_utils.hpp b/examples/55_hopper_mixed_dtype_gemm/mixed_dtype_utils.hpp new file mode 100644 index 00000000..fdb31316 --- /dev/null +++ b/examples/55_hopper_mixed_dtype_gemm/mixed_dtype_utils.hpp @@ -0,0 +1,391 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/default_epilogue.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/util/command_line.h" +#include "cutlass/util/reference/device/tensor_fill.h" +#include "cutlass/util/reference/device/tensor_compare.h" + +#include "cute/tensor.hpp" + +#include +#include +#include "helper.h" + +enum MixedDtypeGemmMode { + ConvertOnly, + ScaleOnly, + ScaleWithZeroPoint +}; + +/// Command line options parsing +struct MixedDtypeOptions { + + bool help = false; + + float alpha = 1.0f; + float beta = 0.0f; + int iterations = 1000; + int warmup = 1000; + int mode = 1; + int m = 5120, n = 4096, k = 4096; + int g = 128; + int l = 1; + + // Parses the command line + void parse(int argc, char const **args) { + cutlass::CommandLine cmd(argc, args); + + if (cmd.check_cmd_line_flag("help")) { + help = true; + return; + } + + cmd.get_cmd_line_argument("m", m); + cmd.get_cmd_line_argument("n", n); + cmd.get_cmd_line_argument("k", k); + cmd.get_cmd_line_argument("l", l); + cmd.get_cmd_line_argument("g", g); + cmd.get_cmd_line_argument("mode", mode); + cmd.get_cmd_line_argument("alpha", alpha, 1.f); + cmd.get_cmd_line_argument("beta", beta, 0.f); + cmd.get_cmd_line_argument("iterations", iterations); + cmd.get_cmd_line_argument("warmup", warmup); + } + + /// Prints the usage statement. + std::ostream & print_usage(std::ostream &out) const { + + out << "55_hopper_mixed_dtype_gemm\n\n" + << " Hopper Mixed Data Type GEMM using a Warp Specialized kernel.\n\n" + << "Options:\n\n" + << " --help If specified, displays this usage statement\n\n" + << " --m= Sets the M extent of the GEMM\n" + << " --n= Sets the N extent of the GEMM\n" + << " --k= Sets the K extent of the GEMM\n" + << " --l= The number of independent gemm problems with mnk shape\n" + << " --g= The size of each group for the scales and zeros. To broadcast a vector of scales or zeros, set the group size to K.\n" + << " --mode= The mode to run the gemm. 0 does (A @ B), 1 means A @ (scale * B), 2 means A @ (scale * B + zero-point).\n" + << " --alpha= Epilogue scalar alpha\n" + << " --beta= Epilogue scalar beta\n\n" + << " --iterations= Number of profiling iterations to perform.\n\n" + << " --warmup= Number of warmup iterations to perform.\n\n"; + + out + << "\n\nExamples:\n\n" + << "$ " << "55_hopper_mixed_dtype_gemm" << " --m=1024 --n=512 --k=1024 -g=1024 --l=10 --alpha=2 --mode=2 --beta=0.707 \n\n"; + + return out; + } + + /// Compute performance in GFLOP/s + double gflops(double runtime_s) const + { + // Two flops per multiply-add + uint64_t flop = uint64_t(2) * m * n * k * l; + double gflop = double(flop) / double(1.0e9); + return gflop / runtime_s; + } +}; + +/// Result structure +struct MixedDtypeResult +{ + double avg_runtime_ms = 0.0; + double gflops = 0.0; + cutlass::Status status = cutlass::Status::kSuccess; + cudaError_t error = cudaSuccess; + bool passed = false; + +}; + +/// Profiling Loop +template +void mixed_dtype_profiling( + Gemm& gemm, + MixedDtypeOptions const& options, + MixedDtypeResult& result) { + + if (options.iterations <= 0) return; + + cudaEvent_t start, stop; + cudaEventCreate(&start); + cudaEventCreate(&stop); + + std::vector runtimes; + runtimes.reserve(options.iterations); + + for (int iter = 0; iter < options.warmup + options.iterations; ++iter) { + cudaEventRecord(start); + CUTLASS_CHECK(gemm.run()); + cudaEventRecord(stop); + cudaEventSynchronize(stop); + + if (iter >= options.warmup) { + float milliseconds = 0; + cudaEventElapsedTime(&milliseconds, start, stop); + runtimes.push_back(milliseconds); + } + } + + cudaEventDestroy(start); + cudaEventDestroy(stop); + + // Compute average setup and runtime and GFLOPs. + result.avg_runtime_ms = std::accumulate(runtimes.begin(), runtimes.end(), 0.0f) / runtimes.size(); + result.gflops = options.gflops(result.avg_runtime_ms / 1000.0); + + std::cout << " Problem Size: " << options.m << 'x' << options.n << 'x' << options.k << 'x' << options.l << std::endl; + std::cout << " Avg runtime: " << result.avg_runtime_ms << " ms" << std::endl; + std::cout << " GFLOPS: " << result.gflops << std::endl; + +} + +/// Helpers to initialize a block of device data +template +bool initialize_tensor( + cutlass::DeviceAllocation& block, + uint64_t seed = 2023) { + + double scope_max, scope_min; + int bits_input = cutlass::sizeof_bits::value; + int bits_output = cutlass::sizeof_bits::value; + + if (bits_input == 1) { + scope_max = 2; + scope_min = 0; + } + else if (bits_input <= 8) { + scope_max = 2; + scope_min = -2; + } + else if (bits_output == 16) { + scope_max = 5; + scope_min = -5; + } + else { + scope_max = 8; + scope_min = -8; + } + cutlass::reference::device::BlockFillRandomUniform( + block.get(), block.size(), seed, Element(scope_max), Element(scope_min)); + + return true; +} + +template +bool initialize_quant_tensor( + cutlass::DeviceAllocation& block, + uint64_t seed = 2023) { + + float scope_min = float(cutlass::platform::numeric_limits::lowest()); + float scope_max = float(cutlass::platform::numeric_limits::max()); + + cutlass::reference::device::BlockFillRandomUniform( + block.get(), block.size(), seed, Element(scope_max), Element(scope_min)); + + return true; +} + +template +bool initialize_scale( + cutlass::DeviceAllocation& block, + MixedDtypeOptions const& options, + uint64_t seed = 2023) { + + if (options.mode == MixedDtypeGemmMode::ConvertOnly) { + // No scales, so just initialize with 1 so we can use the same kernel to dequantize the data. + std::vector stage(block.size(), Element(1.0f)); + block.copy_from_host(stage.data()); + } + else { + float elt_max_f = float(cutlass::platform::numeric_limits::max()); + const float max_dequant_val = 4.f; + const float min_dequant_val = 0.5f; + + float scope_max(max_dequant_val / elt_max_f); + float scope_min(min_dequant_val / elt_max_f); + + cutlass::reference::device::BlockFillRandomUniform( + block.get(), block.size(), seed, Element(scope_max), Element(scope_min)); + } + return true; +} + +template +bool initialize_zero( + cutlass::DeviceAllocation& block, + MixedDtypeOptions const& options, + uint64_t seed = 2023) { + + if (options.mode == MixedDtypeGemmMode::ScaleWithZeroPoint) { + cutlass::reference::device::BlockFillRandomUniform( + block.get(), block.size(), seed, Element(2.0f), Element(-2.0f)); + } else { + // No bias, so just initialize with 1 so we can use the same kernel to dequantize the data. + std::vector stage(block.size(), Element(0.0f)); + block.copy_from_host(stage.data()); + } + return true; +} + +/// Dequantize the weights for verification + +template +__global__ void dequantize_weight_kernel(DequantizedElement* dq_buffer, + QuantizedElement const* q_buffer, + OperandLayout const operand_layout, + ElementScale const* scale_buffer, + ElementZero const* zero_buffer, + ScaleBroadCastLayout const broadcasted_scale_layout, + ThrLayout thr_layout) { + using namespace cute; + + // Represent the full tensors to gmem elements. + // These are expected to have shape [MN, K, L] + cute::Tensor gmem_op_dq = cute::make_tensor(cute::make_gmem_ptr(dq_buffer), operand_layout); + auto init_quantized_iterator = [&]() { + if constexpr (cute::sizeof_bits_v >= 8) { + return cute::make_gmem_ptr(q_buffer); + } else { + return cute::subbyte_iterator(q_buffer); + } + }; + cute::Tensor gmem_op_q = cute::make_tensor(init_quantized_iterator(), operand_layout); + // While the scales are expected to have shape [MN, G, L] but with a stride to allow broadcasting + // It is expected that K % G == 0 + cute::Tensor gmem_scale_broadcasted = cute::make_tensor(make_gmem_ptr(scale_buffer), broadcasted_scale_layout); + cute::Tensor gmem_zero_broadcasted = cute::make_tensor(make_gmem_ptr(zero_buffer), broadcasted_scale_layout); + + // Assign 1 thread per element in the thread block + auto blk_shape = make_shape(size<0>(thr_layout), _1{}, _1{}); // + auto blk_coord = make_coord(_, blockIdx.x, blockIdx.y); // (MN, K, L) + + // Tile across the block + auto gOp_dq = cute::local_tile(gmem_op_dq, blk_shape, blk_coord); + auto gScale = cute::local_tile(gmem_scale_broadcasted, blk_shape, blk_coord); + auto gZero = cute::local_tile(gmem_zero_broadcasted, blk_shape, blk_coord); + auto gOp_q = cute::local_tile(gmem_op_q, blk_shape, blk_coord); + + auto tOpDq_gOpDq = cute::local_partition(gOp_dq, thr_layout, threadIdx.x); + auto tScale_gScale = cute::local_partition(gScale, thr_layout, threadIdx.x); + auto tZero_gZero = cute::local_partition(gZero, thr_layout, threadIdx.x); + auto tOpQ_gOpQ = cute::local_partition(gOp_q, thr_layout, threadIdx.x); + + // Make a fragment of registers to hold gmem loads + cute::Tensor rmem_op_q = cute::make_fragment_like(tOpQ_gOpQ(_, _, _, 0)); + cute::Tensor rmem_scale = cute::make_fragment_like(tScale_gScale(_, _, _, 0)); + cute::Tensor rmem_zero = cute::make_fragment_like(tZero_gZero(_, _, _, 0)); + cute::Tensor rmem_op_dq = cute::make_fragment_like(tOpDq_gOpDq(_, _, _, 0)); + cute::Tensor rmem_op_scaled = cute::make_fragment_like(rmem_op_dq); + cute::Tensor rmem_zero_buf = cute::make_fragment_like(rmem_zero); + + cute::Tensor pred_id = cute::make_identity_tensor(shape(operand_layout)); + auto pred_blk_tile = cute::local_tile(pred_id, blk_shape, blk_coord); + auto pred_thr_partition = cute::local_partition(pred_blk_tile, thr_layout, threadIdx.x); + + const auto num_iters = cute::size<3>(tOpDq_gOpDq); + + for (int ii = 0; ii < num_iters; ++ii) { + const auto thread_offset = cute::get<0>(pred_thr_partition(0, 0, 0, ii)); + if (thread_offset < cute::size<0>(operand_layout)) { + cute::copy(tOpQ_gOpQ(_, _, _, ii), rmem_op_q); + cute::copy(tScale_gScale(_, _, _, ii), rmem_scale); + cute::copy(tZero_gZero(_, _, _, ii), rmem_zero); + cute::transform(rmem_op_q, rmem_op_scaled, [] (const QuantizedElement& elt) { return ElementScale(elt); } ); + cute::transform(rmem_zero, rmem_zero_buf, [] (const ElementZero& elt) { return ElementScale(elt); } ); + cute::transform(rmem_op_scaled, rmem_scale, rmem_op_scaled, multiplies{}); + cute::transform(rmem_op_scaled, rmem_zero_buf, rmem_op_scaled, plus{}); + cute::transform(rmem_op_scaled, rmem_op_dq, [] (const ElementScale& elt) { return DequantizedElement(elt); } ); + cute::copy(rmem_op_dq, tOpDq_gOpDq(_, _, _, ii)); + } + } +} + +template +void dequantize_weight(DequantizedElement* dq_buffer, + QuantizedElement const* q_buffer, + OperandLayout const operand_layout, + ElementScale const* scale_buffer, + ElementZero const* zero_buffer, + ScaleLayout const scale_layout, + int const group_size) { + + using namespace cute; + + constexpr int tpb = 128; + auto thr_layout = make_layout(make_shape(Int{})); + + const auto num_rows = get<0>(shape(operand_layout)); + const auto gemm_k = get<1>(shape(operand_layout)); // [MN, K, L] + const auto batches = get<2>(shape(operand_layout)); // [MN, K, L] + const auto scale_k = get<1>(shape(scale_layout)); // [MN, Scale_K, L] + + if (num_rows != size<0>(scale_layout)) { + std::cerr << "Invalid first dimension for scales. Must match first dim for weights." + << " But got shapes " << shape(operand_layout) << " " << shape(scale_layout) + << std::endl; + exit(-1); + } + + const auto scale_stride0 = get<0>(stride(scale_layout)); + const auto scale_stride1 = get<1>(stride(scale_layout)); + const auto scale_stride2 = get<2>(stride(scale_layout)); + + auto scale_shape_bcast = make_shape(num_rows, make_shape(group_size, scale_k), batches); + auto scale_stride_bcast = make_stride(scale_stride0, make_stride(0, scale_stride1), scale_stride2); + auto scale_layout_bcast = make_layout(scale_shape_bcast, scale_stride_bcast); + + const auto blocks_x = gemm_k; + const auto blocks_y = batches; + + dim3 blocks(blocks_x, blocks_y, 1); + dequantize_weight_kernel<<>>(dq_buffer, q_buffer, operand_layout, scale_buffer, zero_buffer, scale_layout_bcast, thr_layout); + CUDA_CHECK(cudaDeviceSynchronize()); +} diff --git a/examples/55_hopper_mixed_dtype_gemm/packed_scale.hpp b/examples/55_hopper_mixed_dtype_gemm/packed_scale.hpp index 7d732dcd..7f701126 100644 --- a/examples/55_hopper_mixed_dtype_gemm/packed_scale.hpp +++ b/examples/55_hopper_mixed_dtype_gemm/packed_scale.hpp @@ -34,6 +34,10 @@ #include #include "cutlass/float8.h" +#include "cutlass/util/reference/device/tensor_fill.h" + +#include "cute/tensor.hpp" +#include "cute/util/type_traits.hpp" namespace cutlass { @@ -129,3 +133,78 @@ private: } }; } + +/// Helpers to initialize scale lookup table + +// In the mainloop, PRMT selects 1 byte from only 8 bytes so the sign bit is handled in an extra PRMT. +// Here the encodings of positive values and negative values are unified (except for the sign bit). +// For instance, 1 becomes 0b0111, which is the same encoding as -1 (0b1111). +bool unify_quant_encoding( + cutlass::DeviceAllocation const& block_in, + cutlass::DeviceAllocation& block_out) { + + using StorageType = cutlass::int4b_t::Storage; + + if (block_in.size() != block_out.size()) { + std::cerr << "block_in and block_out must have same size.\n"; + return false; + } + constexpr int pack = cute::sizeof_bits_v / 4; + std::vector data(block_in.size() / pack); + cutlass::device_memory::copy_to_host(data.data(), (StorageType*)block_in.get(), block_in.size() / pack); + + for (auto&& d : data) { + StorageType out = 0; + StorageType mask = 0x0f; + for (int i = 0; i < pack; ++i) { + cutlass::int4b_t curr; + curr.storage = (d >> (i * 4)) & 0x0f; + switch (curr) { + case 1: curr.storage = StorageType(0b0111); break; // 2's complement + case 2: curr.storage = StorageType(0b0110); break; // 2's complement + case 3: curr.storage = StorageType(0b0101); break; // 2's complement + case 4: curr.storage = StorageType(0b0100); break; // 2's complement + case 5: curr.storage = StorageType(0b0011); break; // 2's complement + case 6: curr.storage = StorageType(0b0010); break; // 2's complement + case 7: curr.storage = StorageType(0b0001); break; // 2's complement + default: break; + } + out |= (curr.storage << (4 * i)) & mask; + mask <<= 4; + } + d = out; + } + + cutlass::device_memory::copy_to_device((StorageType*)block_out.get(), data.data(), block_out.size() / pack); + return true; +} + +template +bool initialize_packed_scale( + cutlass::DeviceAllocation const& block_in, + cutlass::DeviceAllocation > & block_out) { + + std::vector data_in(block_in.size()); + std::vector > data_out(block_in.size()); + try { + block_in.copy_to_host(data_in.data()); + } catch (cutlass::cuda_exception const& e) + { + std::cerr << "CUDA Error: " << cudaGetErrorString(e.cudaError()) << std::endl; + return false; + } + for (size_t i = 0; i < block_in.size(); ++i) + { + cutlass::packed_scale_t tmp(data_in[i]); + data_out[i] = reinterpret_cast const&>(tmp); + // std::cout << data_in[i] << ":" << std::hex << static_cast(data_in[i].storage) << ",\t" << -data_in[i] << ":" << std::hex << static_cast((-data_in[i]).storage) << std::endl; + } + try { + block_out.copy_from_host(data_out.data()); + } catch (cutlass::cuda_exception const& e) + { + std::cerr << "CUDA Error: " << cudaGetErrorString(e.cudaError()) << std::endl; + return false; + } + return true; +} \ No newline at end of file diff --git a/examples/55_hopper_mixed_dtype_gemm/reorder_utils.hpp b/examples/55_hopper_mixed_dtype_gemm/reorder_utils.hpp index 2be42551..97df3b87 100644 --- a/examples/55_hopper_mixed_dtype_gemm/reorder_utils.hpp +++ b/examples/55_hopper_mixed_dtype_gemm/reorder_utils.hpp @@ -31,6 +31,7 @@ #include "cute/layout.hpp" #include "cute/tensor.hpp" +#include "cute/arch/mma_sm90.hpp" #include "cutlass/util/device_memory.h" @@ -38,62 +39,99 @@ // owned by each thread in contiguous memory locations. This improves smem load vectorization, // particularly for mixed dtype GEMMs where a narrow type is loaded in the thread/value order // of the wider type and may result in inefficient sub-bank (8-bit or 16-bit) accesses. -template -auto compute_memory_reordering_atom() +// In addition, we can reorder the values across several MMA instructions to get even wider +// vectorization (AtomLayout parameter) and permute the values within each instruction to get +// more optimal conversion instruction sequences (ValLayout parameter). +template, + class ValLayout = cute::Layout> +constexpr auto compute_memory_reordering_atom(AtomLayout atom_layout = {}, ValLayout val_layout = {}) { using namespace cute; + static_assert(is_static_v, "ValLayout must be static"); + static_assert(is_static_v, "AtomLayout must be static"); + // 1. Choose an MMA atom to access TV layout and MN shape // Note: parameters like GMMA Major, TileShape, ElementC don't affect TV layout of A, use arbitrary - using MmaAtom = decltype(SM90::GMMA::rs_op_selector>()); + using MmaAtom = decltype(SM90::GMMA::rs_op_selector>()); using MmaTraits = MMA_Traits; - auto shape_MK = select<0,2>(typename MmaTraits::Shape_MNK{}); + auto mk_shape_mma = select<0,2>(typename MmaTraits::Shape_MNK{}); auto tv_layout_mma = typename MmaTraits::ALayout{}; + static_assert(size<1>(tv_layout_mma) % size(val_layout) == 0, "Value layout must evenly divide the MMA value layout"); - // 2. Create a single warp's TV layout from that of the whole MMA + // 2. Create a single warp's TV layout from that of the whole MMA and invert to get (m,k -> thr,val) // Note: this assumes A is partitioned between warps along M mode - auto tile_TV_warp = make_shape(Int<32>{}, size<1>(tv_layout_mma)); - auto tv_layout_mma_warp = make_layout_like(composition(tv_layout_mma, tile_TV_warp)); + auto tv_tiler_warp = make_shape(Int<32>{}, size<1>(tv_layout_mma)); + auto mk_shape_warp = shape_div(mk_shape_mma, size(typename MmaTraits::ThrID{}) / Int<32>{}); + auto tv_layout_mma_warp = make_layout_like(composition(tv_layout_mma, tv_tiler_warp)); + auto mk_layout_mma_warp = right_inverse(tv_layout_mma_warp).with_shape(mk_shape_warp); - // 3. Invert warp's TV layout to get MK layout (m,k -> thr,val) - auto shape_MK_warp = shape_div(shape_MK, size(typename MmaTraits::ThrID{}) / Int<32>{}); - auto mk_layout_mma_warp = right_inverse(tv_layout_mma_warp).with_shape(shape_MK_warp); + // 3. Repeat the warp layout NumAtoms times along K mode to get wider vectorization + auto mk_layout_mma_trgt = blocked_product(mk_layout_mma_warp, atom_layout); // 4. Compose with a contiguous layout of values in each thread (required for smem vectorization) - auto tv_to_offset = make_ordered_layout(shape(tv_layout_mma_warp), Step<_1,_0>{}); - auto layout_atom = composition(tv_to_offset, mk_layout_mma_warp); + auto val_to_offset = logical_product(val_layout, size<1>(tv_layout_mma) / size(val_layout) * size(atom_layout)); + auto thr_to_offset = make_layout(size<0>(tv_layout_mma_warp)); + auto tv_to_offset = select<1,0>(logical_product(val_to_offset, thr_to_offset)); + auto layout_atom = composition(tv_to_offset, mk_layout_mma_trgt); return layout_atom; } -template +template __global__ void reorder_tensor_kernel( - cute::Tensor src, - cute::Tensor dst) + cute::Tensor S, + cute::Tensor D, + TiledCopy tiled_copy) { - auto i = blockIdx.x; - auto k = blockIdx.y; - for (int j = threadIdx.x; j < cute::size<1>(src); j += blockDim.x) { - dst(i,j,k) = src(i,j,k); - } + using namespace cute; + + using T = typename EngineDst::value_type; + + Tensor gS = local_tile(S, TileShape{}, make_coord(blockIdx.x, _, blockIdx.z)); + Tensor gD = local_tile(D, TileShape{}, make_coord(blockIdx.x, _, blockIdx.z)); + + auto thread_copy = tiled_copy.get_slice(threadIdx.x); + Tensor tS = thread_copy.partition_S(gS); + Tensor tD = thread_copy.partition_D(gD); + + copy(tiled_copy, tS, tD); } template void reorder_tensor( - cute::Tensor t_src, - cute::Tensor t_dst) + cute::Tensor S, + cute::Tensor D) { + using namespace cute; + using T = typename EngineDst::value_type; - static_assert(cute::is_same_v, T>, "Type mismatch"); - using V = cute::uint_bit_t)>; + static_assert(is_same_v, T>, "Type mismatch"); - cute::Tensor v_src = cute::recast(t_src); - cute::Tensor v_dst = cute::recast(t_dst); + // Construct a value layout that assigns at least 8 bits of contiguous elements in destination tensor to a thread + // This avoids a race condition when writing out subbyte types (e.g. int4b_t). + auto has_major_mode = [](auto s) { + return any_of(s, [](auto a){ return is_constant<1, decltype(a)>{}; }); + }; + static_assert(has_major_mode(stride<0>(LayoutDst{})) ^ has_major_mode(stride<1>(LayoutDst{})), + "Could not find stride-1 mode in destination layout"); + constexpr int N = shape_div(Int<8>{}, sizeof_bits{}); + auto val_layout = conditional_return(LayoutDst{}))>( + make_layout(make_shape(Int{}, Int<1>{}), GenColMajor{}), + make_layout(make_shape(Int<1>{}, Int{}), GenRowMajor{})); - int threads = 256; - dim3 blocks{unsigned(cute::size<0>(v_src)), unsigned(cute::size<2>(v_src)), 1u}; + // Make a tiled copy with a simple row-major thread order and above layout + int constexpr NumThreads = 128; + auto const thr_layout = make_layout(make_shape(Int<1>{}, Int{})); + auto tiled_copy = make_tiled_copy(Copy_Atom{}, thr_layout, val_layout); - reorder_tensor_kernel<<>>(v_src, v_dst); + // Assign a group of 16 rows to a threadblock; this matches the shuffle atom size for Hopper + using TileShape = Shape<_16>; + auto tiled_D = group_modes<3,rank_v>(tiled_divide(D, TileShape{})); + dim3 blocks{unsigned(size<1>(tiled_D)), 1u, unsigned(size<3>(tiled_D))}; + + reorder_tensor_kernel<<>>(S, D, tiled_copy); CUDA_CHECK(cudaDeviceSynchronize()); } @@ -105,8 +143,9 @@ void reorder_tensor( T * dst, LayoutDst const& layout_dst) { - reorder_tensor(make_tensor(src, layout_src), - make_tensor(dst, layout_dst)); + using namespace cute; + reorder_tensor(make_tensor(make_gmem_ptr(src), layout_src), + make_tensor(make_gmem_ptr(dst), layout_dst)); } // In-place version @@ -116,7 +155,8 @@ void reorder_tensor( LayoutSrc const& layout_src, LayoutDst const& layout_dst) { - cutlass::DeviceAllocation temp(cute::size(layout_src)); + using namespace cute; + cutlass::DeviceAllocation temp(size(layout_src)); reorder_tensor(data, layout_src, temp.get(), layout_dst); - cutlass::device_memory::copy_device_to_device(data, temp.get(), static_cast(cute::size(layout_src))); + cutlass::device_memory::copy_device_to_device(data, temp.get(), static_cast(size(layout_src))); } \ No newline at end of file diff --git a/examples/55_hopper_mixed_dtype_gemm/unfused_weight_dequantize.hpp b/examples/55_hopper_mixed_dtype_gemm/unfused_weight_dequantize.hpp deleted file mode 100644 index e7b5f5a7..00000000 --- a/examples/55_hopper_mixed_dtype_gemm/unfused_weight_dequantize.hpp +++ /dev/null @@ -1,161 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ - -#pragma once - -#include "cute/tensor.hpp" - -#include -#include "helper.h" - -template -__global__ void dequantize_weight_kernel(DequantizedElement* dq_buffer, - QuantizedElement const* q_buffer, - OperandLayout const operand_layout, - ElementScale const* scale_buffer, - ElementZero const* zero_buffer, - ScaleBroadCastLayout const broadcasted_scale_layout, - ThrLayout thr_layout) { - using namespace cute; - - // Represent the full tensors to gmem elements. - // These are expected to have shape [MN, K, L] - Tensor gmem_op_dq = make_tensor(make_gmem_ptr(dq_buffer), operand_layout); - auto init_quantized_iterator = [&]() { - if constexpr (cute::sizeof_bits_v >= 8) { - return make_gmem_ptr(q_buffer); - } else { - return subbyte_iterator(q_buffer); - } - }; - Tensor gmem_op_q = make_tensor(init_quantized_iterator(), operand_layout); - // While the scales are expected to have shape [MN, G, L] but with a stride to allow broadcasting - // It is expected that K % G == 0 - Tensor gmem_scale_broadcasted = make_tensor(make_gmem_ptr(scale_buffer), broadcasted_scale_layout); - Tensor gmem_zero_broadcasted = make_tensor(make_gmem_ptr(zero_buffer), broadcasted_scale_layout); - - // Assign 1 thread per element in the thread block - auto blk_shape = make_shape(size<0>(thr_layout), _1{}, _1{}); // - auto blk_coord = make_coord(_, blockIdx.x, blockIdx.y); // (MN, K, L) - - // Tile across the block - auto gOp_dq = local_tile(gmem_op_dq, blk_shape, blk_coord); - auto gScale = local_tile(gmem_scale_broadcasted, blk_shape, blk_coord); - auto gZero = local_tile(gmem_zero_broadcasted, blk_shape, blk_coord); - auto gOp_q = local_tile(gmem_op_q, blk_shape, blk_coord); - - auto tOpDq_gOpDq = local_partition(gOp_dq, thr_layout, threadIdx.x); - auto tScale_gScale = local_partition(gScale, thr_layout, threadIdx.x); - auto tZero_gZero = local_partition(gZero, thr_layout, threadIdx.x); - auto tOpQ_gOpQ = local_partition(gOp_q, thr_layout, threadIdx.x); - - // Make a fragment of registers to hold gmem loads - Tensor rmem_op_q = make_fragment_like(tOpQ_gOpQ(_, _, _, 0)); - Tensor rmem_scale = make_fragment_like(tScale_gScale(_, _, _, 0)); - Tensor rmem_zero = make_fragment_like(tZero_gZero(_, _, _, 0)); - Tensor rmem_op_dq = make_fragment_like(tOpDq_gOpDq(_, _, _, 0)); - Tensor rmem_op_scaled = make_fragment_like(rmem_op_dq); - Tensor rmem_zero_buf = make_fragment_like(rmem_zero); - - Tensor pred_id = make_identity_tensor(shape(operand_layout)); - auto pred_blk_tile = local_tile(pred_id, blk_shape, blk_coord); - auto pred_thr_partition = local_partition(pred_blk_tile, thr_layout, threadIdx.x); - - const auto num_iters = size<3>(tOpDq_gOpDq); - - for (int ii = 0; ii < num_iters; ++ii) { - const auto thread_offset = get<0>(pred_thr_partition(0, 0, 0, ii)); - if (thread_offset < size<0>(operand_layout)) { - copy(tOpQ_gOpQ(_, _, _, ii), rmem_op_q); - copy(tScale_gScale(_, _, _, ii), rmem_scale); - copy(tZero_gZero(_, _, _, ii), rmem_zero); - transform(rmem_op_q, rmem_op_scaled, [] (const QuantizedElement& elt) { return ElementScale(elt); } ); - transform(rmem_zero, rmem_zero_buf, [] (const ElementZero& elt) { return ElementScale(elt); } ); - transform(rmem_op_scaled, rmem_scale, rmem_op_scaled, multiplies{}); - transform(rmem_op_scaled, rmem_zero_buf, rmem_op_scaled, plus{}); - transform(rmem_op_scaled, rmem_op_dq, [] (const ElementScale& elt) { return DequantizedElement(elt); } ); - copy(rmem_op_dq, tOpDq_gOpDq(_, _, _, ii)); - } - } -} - -template -void dequantize_weight(DequantizedElement* dq_buffer, - QuantizedElement const* q_buffer, - OperandLayout const operand_layout, - ElementScale const* scale_buffer, - ElementZero const* zero_buffer, - ScaleLayout const scale_layout, - int const group_size) { - - using namespace cute; - - constexpr int tpb = 128; - auto thr_layout = make_layout(make_shape(Int{})); - - const auto num_rows = get<0>(shape(operand_layout)); - const auto gemm_k = get<1>(shape(operand_layout)); // [MN, K, L] - const auto batches = get<2>(shape(operand_layout)); // [MN, K, L] - const auto scale_k = get<1>(shape(scale_layout)); // [MN, Scale_K, L] - - if (num_rows != size<0>(scale_layout)) { - std::cerr << "Invalid first dimension for scales. Must match first dim for weights." - << " But got shapes " << shape(operand_layout) << " " << shape(scale_layout) - << std::endl; - exit(-1); - } - - const auto scale_stride0 = get<0>(stride(scale_layout)); - const auto scale_stride1 = get<1>(stride(scale_layout)); - const auto scale_stride2 = get<2>(stride(scale_layout)); - - auto scale_shape_bcast = make_shape(num_rows, make_shape(group_size, scale_k), batches); - auto scale_stride_bcast = make_stride(scale_stride0, make_stride(0, scale_stride1), scale_stride2); - auto scale_layout_bcast = make_layout(scale_shape_bcast, scale_stride_bcast); - - const auto blocks_x = gemm_k; - const auto blocks_y = batches; - - dim3 blocks(blocks_x, blocks_y, 1); - dequantize_weight_kernel<<>>(dq_buffer, q_buffer, operand_layout, scale_buffer, zero_buffer, scale_layout_bcast, thr_layout); - CUDA_CHECK(cudaDeviceSynchronize()); -} diff --git a/include/cutlass/detail/collective/mixed_input_utils.hpp b/include/cutlass/detail/collective/mixed_input_utils.hpp new file mode 100644 index 00000000..c78cf4ba --- /dev/null +++ b/include/cutlass/detail/collective/mixed_input_utils.hpp @@ -0,0 +1,1017 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/numeric_conversion.h" + +#include "cute/util/type_traits.hpp" +#include "cute/arch/copy_sm90.hpp" +#include "cute/numeric/arithmetic_tuple.hpp" + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { + +// The universal converter +template < + class SrcType, + class DstType, + class LayoutIn, + class LayoutOut +> +struct LayoutAwareConvertImpl { + template + CUTLASS_DEVICE + static void convert( + cute::Tensor const& src, + cute::Tensor & dst) { + + static_assert(cute::is_same_v && + cute::is_same_v); + static_assert(cute::cosize_v == cute::cosize_v); + constexpr int N = decltype(cute::max_common_vector(LayoutIn{}, LayoutOut{})){}; + using SrcArray = cutlass::Array; + using DstArray = cutlass::Array; + using Converter = cutlass::NumericArrayConverter; + auto&& src_vm = cute::recast(src); + auto&& dst_vm = cute::recast(dst); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i BF16 with [02461357] value order +template <> +struct LayoutAwareConvertImpl< + cutlass::int4b_t, + cutlass::bfloat16_t, + cute::Layout, cute::Stride<_4,_1>>, + cute::Layout<_8> +> { + template + CUTLASS_DEVICE + static void convert( + cute::Tensor, cute::Stride<_4,_1>> + > const& src, + cute::Tensor + >& dst) { + + static_assert(cute::is_same_v && + cute::is_same_v); + using SrcArray = cutlass::Array; + using DstArray = cutlass::Array; + using RegArray = cutlass::AlignedArray; + + auto&& src_reg = cute::recast(src)(0); + auto&& r = cute::recast(dst)(0); + CUTLASS_PRAGMA_UNROLL + for (size_t ii = 0; ii < RegArray::kElements; ++ii) { + r[ii] = src_reg >> (4 * (ii)); + static constexpr uint32_t xor_mask = 0x43084308; + static constexpr uint32_t lo_mask = 0x000F000F; + static constexpr uint32_t immLut = (0xf0 & 0xcc) ^ 0xaa; + asm volatile( + "{\n" + " lop3.b32 %0, %0, %1, %2, %3;\n" + "}\n" + : "+r"(r[ii]) + : "n"(lo_mask), "n"(xor_mask), "n"(immLut)); + static constexpr uint32_t lo_bias = xor_mask; // 0x43084308, {136, 136} + { + __nv_bfloat162& bf16x2_val = reinterpret_cast<__nv_bfloat162&>(r[ii]); + bf16x2_val = __hsub2(bf16x2_val, + reinterpret_cast(lo_bias)); + } + } + } +}; + +// Specialization for UINT4 -> BF16 with [02461357] value order +template <> +struct LayoutAwareConvertImpl< + cutlass::uint4b_t, + cutlass::bfloat16_t, + cute::Layout, cute::Stride<_4,_1>>, + cute::Layout<_8> +> { + template + CUTLASS_DEVICE + static void convert( + cute::Tensor, cute::Stride<_4,_1>> + > const& src, + cute::Tensor + >& dst) { + + static_assert(cute::is_same_v && + cute::is_same_v); + using SrcArray = cutlass::Array; + using DstArray = cutlass::Array; + using RegArray = cutlass::AlignedArray; + + auto&& src_reg = cute::recast(src)(0); + auto&& r = cute::recast(dst)(0); + CUTLASS_PRAGMA_UNROLL + for (size_t ii = 0; ii < RegArray::kElements; ++ii) { + r[ii] = src_reg >> (4 * (ii)); + static constexpr uint32_t or_mask = 0x43004300; + static constexpr uint32_t lo_mask = 0x000F000F; + static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; + asm volatile( + "{\n" + " lop3.b32 %0, %0, %1, %2, %3;\n" + "}\n" + : "+r"(r[ii]) + : "n"(lo_mask), "n"(or_mask), "n"(immLut)); + static constexpr uint32_t lo_bias = or_mask; // 0x43004300, {128, 128} + { + __nv_bfloat162& bf16x2_val = reinterpret_cast<__nv_bfloat162&>(r[ii]); + bf16x2_val = __hsub2(bf16x2_val, + reinterpret_cast(lo_bias)); + } + } + } +}; + +// Specialization for INT4 -> FP16 with [02461357] value order +template <> +struct LayoutAwareConvertImpl< + cutlass::int4b_t, + cutlass::half_t, + cute::Layout, cute::Stride<_4,_1>>, + cute::Layout<_8> +> { + template + CUTLASS_DEVICE + static void convert( + cute::Tensor, cute::Stride<_4,_1>> + > const& src, + cute::Tensor + >& dst) { + + static_assert(cute::is_same_v && + cute::is_same_v); + using SrcArray = cutlass::Array; + using DstArray = cutlass::Array; + using RegArray = cutlass::AlignedArray; + + auto&& src_reg = cute::recast(src)(0); + auto&& r = cute::recast(dst)(0); + CUTLASS_PRAGMA_UNROLL + for (int ii = 0; ii < RegArray::kElements; ii += 2) { + auto src_ = src_reg >> (4 * (ii)); + r[ii + 0] = src_; + r[ii + 1] = src_; + static constexpr uint32_t lo_xor_mask = 0x64086408; + static constexpr uint32_t hi_xor_mask = 0x64806480; + static constexpr uint32_t lo_mask = 0x000F000F; + static constexpr uint32_t hi_mask = 0x00F000F0; + static constexpr uint32_t immLut = (0xf0 & 0xcc) ^ 0xaa; + asm volatile( + "{\n" + " lop3.b32 %0, %0, %1, %2, %3;\n" + "}\n" + : "+r"(r[ii + 0]) + : "n"(lo_mask), "n"(lo_xor_mask), "n"(immLut)); + asm volatile( + "{\n" + " lop3.b32 %0, %0, %1, %2, %3;\n" + "}\n" + : "+r"(r[ii + 1]) + : "n"(hi_mask), "n"(hi_xor_mask), "n"(immLut)); + static constexpr uint32_t lo_bias = 0x64086408; // {1032, 1032} + static constexpr uint32_t hi_bias = 0xD480D480; // {-72, -72} + static constexpr uint32_t hi_scale = 0x2C002C00; // {1/16, 1/16} + { + half2& fp16x2_val = reinterpret_cast<__half2&>(r[ii + 0]); + fp16x2_val = __hsub2(fp16x2_val, + reinterpret_cast(lo_bias)); + } + { + half2& fp16x2_val = reinterpret_cast<__half2&>(r[ii + 1]); + fp16x2_val = __hfma2(fp16x2_val, + reinterpret_cast(hi_scale), + reinterpret_cast(hi_bias)); + } + } + } +}; + +// Specialization for UINT4 -> FPF16 with [02461357] value order +template <> +struct LayoutAwareConvertImpl< + cutlass::uint4b_t, + cutlass::half_t, + cute::Layout, cute::Stride<_4,_1>>, + cute::Layout<_8> +> { + template + CUTLASS_DEVICE + static void convert( + cute::Tensor, cute::Stride<_4,_1>> + > const& src, + cute::Tensor + >& dst) { + + static_assert(cute::is_same_v && + cute::is_same_v); + using SrcArray = cutlass::Array; + using DstArray = cutlass::Array; + using RegArray = cutlass::AlignedArray; + + auto&& src_reg = cute::recast(src)(0); + auto&& r = cute::recast(dst)(0); + CUTLASS_PRAGMA_UNROLL + for (int ii = 0; ii < RegArray::kElements; ii += 2) { + auto src_ = src_reg >> (4 * (ii)); + r[ii + 0] = src_; + r[ii + 1] = src_; + static constexpr uint32_t or_mask = 0x64006400; + static constexpr uint32_t lo_mask = 0x000F000F; + static constexpr uint32_t hi_mask = 0x00F000F0; + static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; + asm volatile( + "{\n" + " lop3.b32 %0, %0, %1, %2, %3;\n" + "}\n" + : "+r"(r[ii]) + : "n"(lo_mask), "n"(or_mask), "n"(immLut)); + asm volatile( + "{\n" + " lop3.b32 %0, %0, %1, %2, %3;\n" + "}\n" + : "+r"(r[ii + 1]) + : "n"(hi_mask), "n"(or_mask), "n"(immLut)); + static constexpr uint32_t lo_bias = or_mask; // 0x64006400, {1024, 1024} + static constexpr uint32_t hi_bias = 0xD400D400; // {-64, -64} + static constexpr uint32_t hi_scale = 0x2C002C00; // {1/16, 1/16} + { + half2& fp16x2_val = reinterpret_cast<__half2&>(r[ii + 0]); + fp16x2_val = __hsub2(fp16x2_val, + reinterpret_cast(lo_bias)); + } + { + half2& fp16x2_val = reinterpret_cast<__half2&>(r[ii + 1]); + fp16x2_val = __hfma2(fp16x2_val, + reinterpret_cast(hi_scale), + reinterpret_cast(hi_bias)); + } + } + } +}; + +// Specialization for E5M2 -> FP16 with [3120] value order +template <> +struct LayoutAwareConvertImpl< + cutlass::float_e5m2_t, + cutlass::half_t, + cute::Layout, cute::Stride<_2,_1>>, + cute::Layout<_4> +> { + template + CUTLASS_DEVICE + static void convert( + cute::Tensor, cute::Stride<_2,_1>> + > const& src, + cute::Tensor + >& dst) { + + static_assert(cute::is_same_v && + cute::is_same_v); + using SrcArray = cutlass::Array; + using DstArray = cutlass::Array; + using RegArray = cutlass::AlignedArray; + + auto&& src_reg = cute::recast(src)(0); + auto&& r = cute::recast(dst)(0); + CUTLASS_PRAGMA_UNROLL + for (int ii = 0; ii < RegArray::kElements; ++ii) { + // in registers: a3, a1, a2, a0 + r[RegArray::kElements - ii - 1] = src_reg << (8 * (ii)); + + static constexpr uint32_t and_mask = 0xFF00FF00; + asm volatile( + "{\n" + " and.b32 %0, %0, %1;\n" + "}\n" + : "+r"(r[ii]) + : "n"(and_mask)); + } + } +}; + +// Specialization for INT8 -> BF16 with [3120] value order +template <> +struct LayoutAwareConvertImpl< + cutlass::int8_t, + cutlass::half_t, + cute::Layout, cute::Stride<_2,_1>>, + cute::Layout<_4> +> { + template + CUTLASS_DEVICE + static void convert( + cute::Tensor, cute::Stride<_2,_1>> + > const& src, + cute::Tensor + >& dst) { + + static_assert(cute::is_same_v && + cute::is_same_v); + using SrcArray = cutlass::Array; + using DstArray = cutlass::Array; + using RegArray = cutlass::AlignedArray; + + auto&& src_reg = cute::recast(src)(0); + auto&& r = cute::recast(dst)(0); + CUTLASS_PRAGMA_UNROLL + for (int ii = 0; ii < RegArray::kElements; ++ii) { + uint32_t tmp0, tmp1; + r[ii] = src_reg >> (8 * (ii)); + static constexpr uint32_t or_mask = 0x43004300; + static constexpr uint32_t and_mask_0 = 0x007F007F; + static constexpr uint32_t and_mask_1 = 0x00800080; + static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; + asm volatile( + "{\n" + " lop3.b32 %0, %1, %2, %3, %4;\n" + "}\n" + : "=r"(tmp0) + : "r"(r[ii]), "n"(and_mask_0), "n"(or_mask), "n"(immLut)); + asm volatile( + "{\n" + " lop3.b32 %0, %1, %2, %3, %4;\n" + "}\n" + : "=r"(tmp1) + : "r"(r[ii]), "n"(and_mask_1), "n"(or_mask), "n"(immLut)); + { + __nv_bfloat162& bf16x2_val = reinterpret_cast<__nv_bfloat162&>(r[ii]); + bf16x2_val = __hsub2(reinterpret_cast<__nv_bfloat162 const&>(tmp0), + reinterpret_cast<__nv_bfloat162 const&>(tmp1)); + } + } + } +}; + +// Specialization for INT8 -> FP16 with [3120] value order +template <> +struct LayoutAwareConvertImpl< + cutlass::int8_t, + cutlass::bfloat16_t, + cute::Layout, cute::Stride<_2,_1>>, + cute::Layout<_4> +> { + template + CUTLASS_DEVICE + static void convert( + cute::Tensor, cute::Stride<_2,_1>> + > const& src, + cute::Tensor + >& dst) { + + static_assert(cute::is_same_v && + cute::is_same_v); + using SrcArray = cutlass::Array; + using DstArray = cutlass::Array; + using RegArray = cutlass::AlignedArray; + + auto&& src_reg = cute::recast(src)(0); + auto&& r = cute::recast(dst)(0); + CUTLASS_PRAGMA_UNROLL + for (int ii = 0; ii < RegArray::kElements; ++ii) { + r[ii] = src_reg >> (8 * (ii)); + static constexpr uint32_t xor_mask = 0x64806480; + static constexpr uint32_t and_mask = 0x00FF00FF; + static constexpr uint32_t immLut = (0xf0 & 0xcc) ^ 0xaa; + asm volatile( + "{\n" + " lop3.b32 %0, %0, %1, %2, %3;\n" + "}\n" + : "+r"(r[ii]) + : "n"(and_mask), "n"(xor_mask), "n"(immLut)); + { + static constexpr uint32_t bias = 0x64806480; + __half2& fp16x2_val = reinterpret_cast<__half2&>(r[ii]); + fp16x2_val = __hsub2(fp16x2_val, + reinterpret_cast<__half2 const&>(bias)); + } + } + } +}; + +template < + class EngineIn, + class EngineOut, + class LayoutIn, + class LayoutOut +> +CUTLASS_DEVICE +void LayoutAwareConvert( // Accept mutable temporaries + cute::Tensor const& src, + cute::Tensor && dst) { + + LayoutAwareConvert(src, dst); +} +template < + class EngineIn, + class EngineOut, + class LayoutIn, + class LayoutOut +> +CUTLASS_DEVICE +void LayoutAwareConvert( + cute::Tensor const& src, + cute::Tensor & dst) { + + using SrcType = typename EngineIn::value_type; + using DstType = typename EngineOut::value_type; + Tensor src_vm = coalesce(src); + Tensor dst_vm = coalesce(dst); + Layout src_layout = src_vm.layout(); + Layout dst_layout = dst_vm.layout(); + LayoutAwareConvertImpl::convert(src_vm, dst_vm); +} + +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::gemm::collective::detail { + +template +static constexpr +CUTLASS_HOST_DEVICE +auto get_logical_ptr(PointerType const* ptr) { + if constexpr (cute::sizeof_bits_v < 8) { + return subbyte_iterator(ptr); + } + else { + return ptr; + } +} +template +static constexpr +CUTLASS_HOST_DEVICE +auto get_smem_layout(LayoutAtom layout_atom, TileShape const& tile_shape, Stride const& stride) { + if constexpr (not cute::is_layout::value) { + return tile_to_shape( + layout_atom, + append(tile_shape, Int{}), + cute::conditional_t< ::cutlass::gemm::detail::is_major<0,Stride>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}); + } + else { + auto gmem_tile = composition(stride, tile_shape); + return make_layout_like(append(gmem_tile, make_layout(Int{}, 0))); + } +} +template +static constexpr +CUTLASS_HOST_DEVICE +auto get_gmem_layout(Shape const& shape, Stride const& stride) { + if constexpr (not cute::is_layout::value) { + return make_layout(shape, stride); + } + else { + return stride; + } +} + +template +struct MixedInputUtils { +private: + using KernelSchedule = typename Collective::KernelSchedule; + using ConversionMode = typename Collective::ConversionMode; + using SmemLayoutA = typename Collective::SmemLayoutA; + using SmemLayoutB = typename Collective::SmemLayoutB; + using SmemLayoutScale = typename Collective::SmemLayoutScale; + using SwappedElementA = typename Collective::SwappedElementA; + using SwappedElementB = typename Collective::SwappedElementB; + using RealSwappedElementA = typename Collective::RealSwappedElementA; + using RealSwappedElementB = typename Collective::RealSwappedElementB; + using ElementScale = typename Collective::ElementScale; + using ElementZero = typename Collective::ElementZero; + using SmemCopyAtomScale = typename Collective::SmemCopyAtomScale; + static constexpr auto KernelConversionMode = Collective::KernelConversionMode; + static constexpr auto ModeHasScales = Collective::ModeHasScales; + static constexpr auto UseScaleLookupTable = Collective::UseScaleLookupTable; + +public: + static constexpr auto + elements_per_smem_scale() { + if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { + return 0; + } + else if constexpr (ModeHasScales) { + return cute::cosize_v; + } + else { + static_assert(cutlass::detail::dependent_false, "Type not handled in scale smem allocation."); + } + } + + static constexpr auto + elements_per_smem_zero() { + if constexpr (KernelConversionMode == ConversionMode::DirectConvert || + KernelConversionMode == ConversionMode::ConvertAndScale ) { + return 0; + } + else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { + return cute::cosize_v; + } + else { + static_assert(cutlass::detail::dependent_false, "Type not handled in scale smem allocation."); + } + } + + // These methods use some the public members of the class. For that reason, we define them after the public section. + static constexpr uint32_t + compute_tma_transaction_bytes_mk() { + return cutlass::bits_to_bytes(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast(cute::sizeof_bits_v)); + } + + static constexpr uint32_t + compute_tma_transaction_bytes_nk() { + return cutlass::bits_to_bytes(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast(cute::sizeof_bits_v)); + } + + static constexpr uint32_t + compute_tma_transaction_bytes_extra() { + if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { + return 0; + } + else if constexpr (ModeHasScales) { + constexpr uint32_t scale_tx_bytes = cutlass::bits_to_bytes(size<0>(SmemLayoutScale{}) * size<1>(SmemLayoutScale{}) * static_cast(cute::sizeof_bits_v)); + static_assert(scale_tx_bytes % 128 == 0, "Each scale stage must be 128B aligned."); // required by TMA + if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { + return scale_tx_bytes; + } + else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { + // Scale and zero share smem layout + constexpr uint32_t zero_tx_bytes = cutlass::bits_to_bytes(size<0>(SmemLayoutScale{}) * size<1>(SmemLayoutScale{}) * static_cast(cute::sizeof_bits_v)); + static_assert(zero_tx_bytes % 128 == 0, "Each zero stage must be 128B aligned."); // required by TMA + return scale_tx_bytes + zero_tx_bytes; + } + else { + static_assert(cutlass::detail::dependent_false, "Type not handled in tma transaction bytes computation."); + } + } + else { + static_assert(cutlass::detail::dependent_false, "Type not handled in tma transaction bytes computation."); + } + } + + /// Utilities to copy A and extra inputs from smem to RF + template + CUTLASS_DEVICE + static void copy_tensors_MK( + SmemTiledCopyA const& smem_tiled_copy_A, + TensorASmemView const& tCsA, + TensorACopyView& tCrA_copy_view, + cute::tuple const& partitioned_mma_extra_info, + cute::tuple const& tiled_copy_and_views, + int k_block, + int read_stage) { + + copy(smem_tiled_copy_A, tCsA(_,_,k_block,read_stage), tCrA_copy_view(_,_,k_block)); + + if (k_block == 0) { + // We are starting a new k-tile so copy the scale + if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { + // nothing to do + } + else if constexpr (ModeHasScales) { + auto smem_tiled_copy_S = cute::get<0>(tiled_copy_and_views); + auto tCrS_copy_view = cute::get<1>(tiled_copy_and_views); + auto tCsS = cute::get<0>(partitioned_mma_extra_info); + copy(smem_tiled_copy_S, tCsS(_,_,k_block,read_stage), tCrS_copy_view(_,_,k_block)); + if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { + // Nothing extra to do + } else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { + auto tCsZ = cute::get<2>(partitioned_mma_extra_info); + auto tCrZ_copy_view = cute::get<2>(tiled_copy_and_views); + copy(smem_tiled_copy_S, tCsZ(_,_,k_block,read_stage), tCrZ_copy_view(_,_,k_block)); + } else { + static_assert(cutlass::detail::dependent_false, "Conversion mode not handled in A -> RF path."); + } + } + else { + static_assert(cutlass::detail::dependent_false, "Conversion mode not handled in A -> RF path."); + } + } + } + + // The core converter uses a lookup table to converts i4 -> 8 bit value. + template + CUTLASS_DEVICE + static void lookup_table_convert( // Accept mutable temporaries + Tensor const& src, + Tensor && dst, + Tensor const& scales_neg, + Tensor const& scales_pos) { + + lookup_table_convert(src, dst, scales_neg, scales_pos); + } + template + CUTLASS_DEVICE + static void lookup_table_convert( + Tensor const& src, + Tensor & dst, + Tensor const& scales_neg, + Tensor const& scales_pos) { + + constexpr int N = cute::cosize(LayoutIn{}); + static_assert(N == 4 || N == 8); + static_assert(cosize(LayoutScale{}) <= N / 4, + "at least 4 consecutive weights must share the same scale."); + using SrcArray = cutlass::Array; + using DstArray = cutlass::Array; + using RegArray = cutlass::AlignedArray; + + // View the input as reg + auto&& src_reg = cute::recast(src)(0); + auto&& r = cute::recast(dst)(0); + + // Determines if to get from the signed or unsigned candidates + static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; + uint32_t sign; // ((reg & 0x88888888) | 0x64206420) >> 1 + asm volatile( + "{\n" + " lop3.b32 %0, %1, %2, %3, %4;\n" \ + "}\n" + : "=r"(sign) + : "r"(src_reg), "n"(0x88888888), "n"(0x64206420), "n"(immLut) + ); + sign = sign >> 1; + + // Ignore sign bit when indexing into LUT + uint32_t lut_idx = src_reg & 0x77777777; + Tensor scales_neg_ = cute::filter(scales_neg); + Tensor scales_pos_ = cute::filter(scales_pos); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 4; ++i, lut_idx >>=16, sign >>=16) { + auto&& scale_neg_ = reinterpret_cast const&>(scales_neg_(i)); + auto&& scale_pos_ = reinterpret_cast const&>(scales_pos_(i)); + asm volatile( + "{\n" + " .reg .b32 pos, neg ;\n" \ + " prmt .b32 neg, %3, %4, %1 ;\n" \ + " prmt .b32 pos, %5, %6, %1 ;\n" \ + " prmt .b32 %0, pos, neg, %2 ;\n" \ + "}\n" + : "=r"(r[i]) + : "r"(lut_idx), "r"(sign), "r"(scale_neg_[0]), "r"(scale_neg_[1]), "r"(scale_pos_[0]), "r"(scale_pos_[1]) + ); + } + } + + /// Utilities to dequantize A. + template + CUTLASS_DEVICE + static void static_check_scale(Layout const& tensor) { + static_assert(shape<0>(Layout{}) >= 4 && stride<0>(Layout{}) == 0, "At least 4 adjacent weights in a thread must share the same scale."); + } + template + CUTLASS_DEVICE + static void static_check_scale(Tensor const& tensor) { + static_check_scale(flatten(Layout{})); + } + template + CUTLASS_DEVICE + static void dequantize_A_kblock( + Tensor const& tCrA_load, + Tensor& tCrA_mma, + cute::tuple& partitioned_extra_info, + int const k_block) { + + + static_assert(is_rmem::value, "Input tensor for A conversion must come from registers"); + static_assert(is_rmem::value, "Output tensor for A conversion must come from registers"); + static_assert(cosize_v == cosize_v); + static_assert(size_v == cosize_v); + static_assert(size_v == cosize_v); + using SrcType = typename EngineIn::value_type; + using DstType = typename EngineOut::value_type; + + Tensor src = tCrA_load(_, _, k_block); + Tensor dst = tCrA_mma(_, _, k_block); + + CUTE_STATIC_ASSERT_V(size(src(_, 0)) == cosize(src(_, 0).layout()), + "The first mode of tensor src must be contiguous in memory"); + // try to make the size of the first mode equal to 32bit + int constexpr NumValPerSrcReg = cute::min(decltype(size(src(_, 0)))::value, + ceil_div(32, sizeof_bits_v)); + Tensor src_vm = cute::group_modes<1,-1>(cute::zipped_divide(src, Int{})); + Tensor dst_vm = cute::group_modes<1,-1>(cute::zipped_divide(dst, Int{})); + + if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<1>(dst_vm); ++i) { + LayoutAwareConvert(src_vm(_, i), dst_vm(_, i)); + } + } + else if constexpr (UseScaleLookupTable) { + constexpr int num_elements = decltype(size(src))::value; + static_assert(is_same_v, "Lookup table only supports int4 being the quant type now."); + static_assert(sizeof_bits_v == 64, "Lookup table only supports 8 8bit scale values now."); + static_assert(num_elements % 4 == 0 && num_elements >= 4, "Lookup table requires a vector size of 4x when converting."); + + Tensor tCrS_neg = cute::get<1>(partitioned_extra_info); + auto&& tCrS_pos = cute::get<2>(partitioned_extra_info); // modification to its value is needed + Tensor scales_neg = tCrS_neg(_, _, k_block); + Tensor scales_pos = tCrS_pos(_, _, k_block); + CUTE_STATIC_ASSERT_V(cute::size(src) == cute::size(scales_neg)); + + static_check_scale(scales_neg); + static_check_scale(scales_pos); + Tensor scales_neg_vm = cute::group_modes<1,-1>(cute::zipped_divide(scales_neg, Int{})); + Tensor scales_pos_vm = cute::group_modes<1,-1>(cute::zipped_divide(scales_pos, Int{})); + + if (k_block == 0) { + Tensor scales_neg_vm_ = filter(scales_neg_vm); + Tensor scales_pos_vm_ = filter(scales_pos_vm); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(scales_neg_vm_.layout()); ++i) + { + auto&& scale_neg_ = reinterpret_cast const&>(scales_neg_vm_(i)); + auto&& scale_pos_ = reinterpret_cast &>(scales_pos_vm_(i)); + asm volatile( + "{\n" + " and .b32 %0, %2, %4 ;\n" \ + " and .b32 %1, %3, %5 ;\n" \ + "}\n" + : "=r"(scale_pos_[0]), "=r"(scale_pos_[1]) + : "r"(scale_neg_[0]), "r"(scale_neg_[1]), "n"(0x7F7F7F00), "n"(0x7F7F7F7F) + ); + } + } + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<1>(dst_vm); ++i) { + lookup_table_convert(src_vm(_, i), dst_vm(_, i), scales_neg_vm(_, i), scales_pos_vm(_, i)); + } + } + else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { + Tensor scales = cute::get<1>(partitioned_extra_info)(_, _, k_block); + CUTE_STATIC_ASSERT_V(size(src) == size(scales)); + Tensor scales_vm = cute::group_modes<1,-1>(cute::zipped_divide(scales, Int{})); + + if constexpr (is_same_v) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<1>(dst_vm); ++i) { + LayoutAwareConvert(src_vm(_, i), dst_vm(_, i)); + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<0>(dst_vm); ++j) { + dst_vm(j, i) *= scales_vm(j, i); + } + } + } + else { + auto stage = make_tensor_like(src_vm); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<1>(dst_vm); ++i) { + LayoutAwareConvert(src_vm(_, i), stage); + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<0>(dst_vm); ++j) { + stage(j) *= scales_vm(j, i); + } + LayoutAwareConvert(stage, dst_vm(_, i)); + } + } + } + else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { + static_assert(is_same_v, "ElementScale and ElementZero must be the same."); + Tensor scales = cute::get<1>(partitioned_extra_info)(_, _, k_block); + Tensor zeros = cute::get<3>(partitioned_extra_info)(_, _, k_block); + CUTE_STATIC_ASSERT_V(size(src) == size(scales)); + CUTE_STATIC_ASSERT_V(size(src) == size(zeros)); + Tensor scales_vm = cute::group_modes<1,-1>(cute::zipped_divide(scales, Int{})); + Tensor zeros_vm = cute::group_modes<1,-1>(cute::zipped_divide(zeros, Int{})); + + if constexpr (is_same_v) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<1>(dst_vm); ++i) { + LayoutAwareConvert(src_vm(_, i), dst_vm(_, i)); + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<0>(dst_vm); ++j) { + dst_vm(j, i) = dst_vm(j, i) * scales_vm(j, i) + zeros_vm(j, i); + } + } + } + else { + auto stage = make_tensor_like(src_vm); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<1>(dst_vm); ++i) { + LayoutAwareConvert(src_vm(_, i), stage); + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size<0>(dst_vm); ++j) { + stage(j) = stage(j) * scales_vm(j, i) + zeros_vm(j, i); + } + LayoutAwareConvert(stage, dst_vm(_, i)); + } + } + } + else { + static_assert(cutlass::detail::dependent_false, "No A data is loaded."); + } + } + + /// Utilities for any additional inputs inside of the TMA load + template < + class Params, + class TensorStorage, + class... Ts + > + CUTLASS_DEVICE + static auto partition_extra_tma_inputs( + Params const& mainloop_params, + cute::tuple const& load_inputs, + TensorStorage& shared_tensors, + uint2 const& cluster_local_block_id, + int const m_coord, + int const l_coord) { + + if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { + return cute::make_tuple(); + } + else if constexpr (ModeHasScales) { + Tensor sS = make_tensor(make_smem_ptr(shared_tensors.smem_scale.begin()), SmemLayoutScale{}); // (BLK_M,BLK_K,PIPE) + Tensor gS_mkl = get<2>(load_inputs); + auto block_tma_s = mainloop_params.tma_load_scale.get_slice(cluster_local_block_id.y); + Tensor gS = gS_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k) + + Tensor tSgS = block_tma_s.partition_S(gS); // (TMA,TMA_M,TMA_K,k) + Tensor tSsS = block_tma_s.partition_D(sS); // (TMA,TMA_M,TMA_K,PIPE) + if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { + return cute::make_tuple(tSgS, tSsS); + } + else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { + Tensor sZ = make_tensor(make_smem_ptr(shared_tensors.smem_zero.begin()), SmemLayoutScale{}); // (BLK_M,BLK_K,PIPE) + Tensor gZ_mkl = get<3>(load_inputs); + auto block_tma_z = mainloop_params.tma_load_zero.get_slice(cluster_local_block_id.y); + Tensor gZ = gZ_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k) + + Tensor tZgZ = block_tma_z.partition_S(gZ); // (TMA,TMA_M,TMA_K,k) + Tensor tZsZ = block_tma_z.partition_D(sZ); // (TMA,TMA_M,TMA_K,PIPE) + return cute::make_tuple(tSgS, tSsS, tZgZ, tZsZ); + } + else { + static_assert(cutlass::detail::dependent_false, "Conversion mode not handled for input partitioning."); + } + } + else { + static_assert(cutlass::detail::dependent_false, "Conversion mode not handled for input partitioning."); + } + } + + /// Utilities for partitioning extra inputs for loading from smem in the mainloop. + template < + class ThreadMma, + class TensorStorage + > + CUTLASS_DEVICE + static auto partition_extra_mma_info( + ThreadMma const& mma_thread_slice, + TensorStorage& shared_tensors) { + + if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { + // nothing to do + return cute::make_tuple(); + } + else if constexpr (UseScaleLookupTable) { + Tensor sS = make_tensor(make_smem_ptr(shared_tensors.smem_scale.begin()), SmemLayoutScale{});// (BLK_M,BLK_SCALE_K,PIPE) + Tensor tCsS = mma_thread_slice.partition_A(sS); + Tensor tCrS_neg = make_tensor(mma_thread_slice.partition_fragment_A(sS(_,_,Int<0>{})).layout()); + Tensor tCrS_pos = make_tensor(mma_thread_slice.partition_fragment_A(sS(_,_,Int<0>{})).layout()); + + if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { + return cute::make_tuple(tCsS, tCrS_neg, tCrS_pos); + } + } + else if constexpr (ModeHasScales) { + Tensor sS = make_tensor(make_smem_ptr(shared_tensors.smem_scale.begin()), SmemLayoutScale{});// (BLK_M,BLK_SCALE_K,PIPE) + Tensor tCsS = mma_thread_slice.partition_A(sS); + Tensor tCrS = make_tensor(mma_thread_slice.partition_fragment_A(sS(_,_,Int<0>{})).layout()); + + if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { + return cute::make_tuple(tCsS, tCrS); + } + else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { + Tensor sZ = make_tensor(make_smem_ptr(shared_tensors.smem_zero.begin()), SmemLayoutScale{});// (BLK_M,BLK_SCALE_K,PIPE) + Tensor tCsZ = mma_thread_slice.partition_A(sZ); + Tensor tCrZ = make_tensor(mma_thread_slice.partition_fragment_A(sZ(_,_,Int<0>{})).layout()); + return cute::make_tuple(tCsS, tCrS, tCsZ, tCrZ); + } + else { + static_assert(cutlass::detail::dependent_false, "Conversion mode not handled in A -> RF path."); + } + } + else { + static_assert(cutlass::detail::dependent_false, "Conversion mode not handled in A -> RF path."); + } + } + + /// Returns the tiled copy and copy views for the extra inputs. + template + CUTLASS_DEVICE + static auto retile_extra_mma_info( + TiledMma const& tiled_mma, + cute::tuple& partitioned_extra_info, + int const warp_group_thread_idx) { + + if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { + // nothing to do + return cute::make_tuple(); + } + else if constexpr (ModeHasScales) { + auto smem_tiled_copy_S = make_tiled_copy_A(SmemCopyAtomScale{}, tiled_mma); + auto smem_thr_copy_S = smem_tiled_copy_S.get_thread_slice(warp_group_thread_idx); + Tensor tCrS_copy_view = smem_thr_copy_S.retile_D(cute::get<1>(partitioned_extra_info)); // (CPY,CPY_M,CPY_K) + + if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { + return cute::make_tuple(smem_tiled_copy_S, tCrS_copy_view); + } + else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { + Tensor tCrZ_copy_view = smem_thr_copy_S.retile_D(cute::get<3>(partitioned_extra_info)); // (CPY,CPY_M,CPY_K) + return cute::make_tuple(smem_tiled_copy_S, tCrS_copy_view, tCrZ_copy_view); + } + else { + static_assert(cutlass::detail::dependent_false, "Conversion mode not handled in A -> RF path."); + } + } + else { + static_assert(cutlass::detail::dependent_false, "Conversion mode not handled in A -> RF path."); + } + } +}; + +} // cutlass::gemm::collective::detail diff --git a/include/cutlass/gemm/collective/builders/sm90_gmma_builder.inl b/include/cutlass/gemm/collective/builders/sm90_gmma_builder.inl index 8657aad2..64e27a8d 100644 --- a/include/cutlass/gemm/collective/builders/sm90_gmma_builder.inl +++ b/include/cutlass/gemm/collective/builders/sm90_gmma_builder.inl @@ -50,36 +50,40 @@ namespace cutlass::gemm::collective { namespace detail { // Returns the maximum number of smem tiles that can be used with a given smem capacity, or overrides with manual count. -template +template constexpr int compute_stage_count_or_override(StageCount stage_count) { return stages; } // Returns the maximum number of smem tiles that can be used with a given smem capacity, or overrides with manual count. -template +template constexpr int compute_stage_count_or_override(cute::Int stage_count) { return stages; } // Returns the maximum number of smem tiles that can be used with a given smem capacity, or overrides with manual count. -template +template constexpr int -compute_stage_count_or_override(StageCountAutoCarveout stage_count) { +compute_stage_count_or_override(StageCountAutoCarveout stage_count) { constexpr auto mainloop_pipeline_bytes = sizeof(typename cutlass::PipelineTmaAsync<1>::SharedStorage); constexpr auto a_bits = cute::sizeof_bits_v; constexpr auto b_bits = cute::sizeof_bits_v; - constexpr int stage_bytes = + constexpr int stage_bytes_ = cutlass::bits_to_bytes(a_bits * size<0>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) + - cutlass::bits_to_bytes(b_bits * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) + - static_cast(mainloop_pipeline_bytes); + cutlass::bits_to_bytes(b_bits * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{})); - return (CapacityBytes - carveout_bytes) / stage_bytes; + constexpr int stage_bytes = cutlass::round_up(stage_bytes_, alignment) + + static_cast(mainloop_pipeline_bytes); + constexpr int carveout_bytes = cutlass::round_up(carveout_bytes_, alignment); + constexpr int capacity_bytes = capacity_bytes_ / alignment * alignment; + + return (capacity_bytes - carveout_bytes) / stage_bytes; } // Returns the maximum number of smem tiles that can be used with a given smem capacity (with an optional scale matrix), or overrides with manual count. -template +template constexpr int compute_stage_count_or_override_single_affine_transformed_input(StageCount stage_count) { return stages; @@ -96,9 +100,9 @@ constexpr int get_bits_for_possibly_void_element() { } // Returns the maximum number of smem tiles that can be used with a given smem capacity (with an optional scale matrix), or overrides with manual count. -template +template constexpr int -compute_stage_count_or_override_single_affine_transformed_input(StageCountAutoCarveout stage_count) { +compute_stage_count_or_override_single_affine_transformed_input(StageCountAutoCarveout stage_count) { // 32 bytes to account for barriers etc. constexpr auto mainloop_pipeline_bytes = sizeof(typename cutlass::PipelineTmaAsync<1>::SharedStorage); @@ -114,12 +118,17 @@ compute_stage_count_or_override_single_affine_transformed_input(StageCountAutoCa static_assert(zero_bytes % 128 == 0, "Zero bytes must be a multiple of 128"); // When scales are void, s_bits will be 0 so no smem will be allocated for scales. - constexpr int stage_bytes = + constexpr int stage_bytes_ = cutlass::bits_to_bytes(a_bits * size<0>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) + cutlass::bits_to_bytes(b_bits * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) + - static_cast(scale_bytes + zero_bytes + mainloop_pipeline_bytes); + scale_bytes + zero_bytes; - return (CapacityBytes - carveout_bytes) / stage_bytes; + constexpr int stage_bytes = cutlass::round_up(stage_bytes_, alignment) + + static_cast(mainloop_pipeline_bytes); + constexpr int carveout_bytes = cutlass::round_up(carveout_bytes_, alignment); + constexpr int capacity_bytes = capacity_bytes_ / alignment * alignment; + + return (capacity_bytes - carveout_bytes) / stage_bytes; } template @@ -228,10 +237,9 @@ struct CollectiveBuilder< using SmemLayoutAtomB = decltype(detail::ss_smem_selector< GmmaMajorB, ElementBMma, decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>()); - static constexpr size_t TensorMapStorage = IsArrayOfPointersGemm ? sizeof(cute::TmaDescriptor) * 2 /* for A and B */ : 0; - static constexpr int KernelSmemCarveout = static_cast(TensorMapStorage); + static constexpr int Sm90ReducedSmemCapacityBytes = detail::sm90_smem_capacity_bytes; - static constexpr int PipelineStages = detail::compute_stage_count_or_override(StageCountType{}); using DispatchPolicy = cute::conditional_t, @@ -264,107 +272,12 @@ struct CollectiveBuilder< ///////////////////////////////////////////////////////////////////////////////////////////////// -// GMMA_TMA_WS_RS +// GMMA_TMA_WS_RS template < - class ElementA, - class GmemLayoutATag, - int AlignmentA, - class ElementB, - class GmemLayoutBTag, - int AlignmentB, - class ElementAccumulator, - class TileShape_MNK, - class ClusterShape_MNK, - class StageCountType, - class KernelScheduleType -> -struct CollectiveBuilder< - arch::Sm90, - arch::OpClassTensorOp, - ElementA, - GmemLayoutATag, - AlignmentA, - ElementB, - GmemLayoutBTag, - AlignmentB, - ElementAccumulator, - TileShape_MNK, - ClusterShape_MNK, - StageCountType, - KernelScheduleType, - cute::enable_if_t< - (cute::is_same_v || - cute::is_same_v || - cute::is_same_v) && - detail::is_use_rmem_A()> -> { - static_assert(is_static::value); - static_assert(is_static::value); - static_assert(detail::is_aligned(), - "Should meet TMA alignment requirement\n"); -#ifndef CUTLASS_SM90_COLLECTIVE_BUILDER_SUPPORTED - static_assert(cutlass::detail::dependent_false, "Unsupported Toolkit for SM90 Collective Builder\n"); -#endif - static constexpr cute::GMMA::Major GmmaMajorA = detail::gmma_rs_tag_to_major_A(); - static constexpr cute::GMMA::Major GmmaMajorB = detail::gmma_rs_tag_to_major_B(); - static constexpr bool SwapAB = detail::is_swapAB(); - static constexpr bool IsWarpSpecializedTransposeB = detail::is_warpspecialized_transpose_B< - ElementA, GmemLayoutATag, ElementB, GmemLayoutBTag, KernelScheduleType>(); - - // For fp32 types, map to tf32 MMA value type - using ElementAMma = cute::conditional_t, tfloat32_t, ElementA>; - using ElementBMma = cute::conditional_t, tfloat32_t, ElementB>; - - using AtomLayoutMNK = cute::conditional_t, - Layout>, Layout>>; - - using TiledMma = decltype(cute::make_tiled_mma(cute::GMMA::rs_op_selector< - ElementAMma, ElementBMma, ElementAccumulator, TileShape_MNK, GMMA::Major::K, GMMA::Major::K>(), AtomLayoutMNK{})); - - using GmemTiledCopyA = decltype(detail::sm90_cluster_shape_to_tma_atom(shape<1>(ClusterShape_MNK{}))); - using GmemTiledCopyB = decltype(detail::sm90_cluster_shape_to_tma_atom(shape<0>(ClusterShape_MNK{}))); - - using SmemLayoutAtomA = decltype(detail::rs_smem_selector(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{})), IsWarpSpecializedTransposeB>()); - using SmemLayoutAtomB = decltype(detail::rs_smem_selector(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{})), IsWarpSpecializedTransposeB>()); - - static constexpr int PipelineStages = detail::compute_stage_count_or_override(StageCountType{}); - - using DispatchPolicy = MainloopSm90TmaGmmaRmemAWarpSpecialized< - PipelineStages, ClusterShape_MNK, KernelScheduleType>; - - using SmemCopyAtomA = cute::conditional_t>; - using SmemCopyAtomB = cute::conditional_t, void>; - - using CollectiveOp = CollectiveMma< - DispatchPolicy, - TileShape_MNK, - ElementA, - TagToStrideA_t, - ElementB, - TagToStrideB_t, - TiledMma, - GmemTiledCopyA, - SmemLayoutAtomA, - SmemCopyAtomA, - cute::identity, - GmemTiledCopyB, - SmemLayoutAtomB, - SmemCopyAtomB, - cute::identity - >; -}; - -///////////////////////////////////////////////////////////////////////////////////////////////// - -// GMMA_TMA_WS_RS Mixed Scaled GEMM -template < - class ElementPairA_, + class ElementA_, class GmemLayoutATag_, int AlignmentA, - class ElementPairB_, + class ElementB_, class GmemLayoutBTag_, int AlignmentB, class ElementAccumulator, @@ -376,10 +289,10 @@ template < struct CollectiveBuilder< arch::Sm90, arch::OpClassTensorOp, - ElementPairA_, + ElementA_, GmemLayoutATag_, AlignmentA, - ElementPairB_, + ElementB_, GmemLayoutBTag_, AlignmentB, ElementAccumulator, @@ -388,30 +301,45 @@ struct CollectiveBuilder< StageCountType, KernelScheduleType, cute::enable_if_t< - (cute::is_same_v || - cute::is_same_v || - cute::is_same_v)> + (cute::is_same_v || + cute::is_same_v || + cute::is_same_v || + cute::is_same_v || + cute::is_same_v) && + (detail::is_use_rmem_A() || + // ConvertAndScale and ConvertAndScaleWithZero + cute::is_tuple::value || cute::is_tuple::value || + // DirectConvert + sizeof_bits::value != sizeof_bits::value)> > { private: - using ScaleA = detail::deduce_mixed_width_dtype_t<1, ElementPairA_>; - using ScaleB = detail::deduce_mixed_width_dtype_t<1, ElementPairB_>; - using ZeroA = detail::deduce_mixed_width_dtype_t<2, ElementPairA_>; - using ZeroB = detail::deduce_mixed_width_dtype_t<2, ElementPairB_>; - static constexpr bool NeitherIsTuple = !cute::is_tuple::value && !cute::is_tuple::value; + using ScaleA = detail::deduce_mixed_width_dtype_t<1, ElementA_>; + using ScaleB = detail::deduce_mixed_width_dtype_t<1, ElementB_>; + using ZeroA = detail::deduce_mixed_width_dtype_t<2, ElementA_>; + using ZeroB = detail::deduce_mixed_width_dtype_t<2, ElementB_>; + static constexpr bool NeitherIsTuple = !cute::is_tuple::value && !cute::is_tuple::value; + // Determine if mixed input types. + static constexpr bool IsMixedInput = + cute::sizeof_bits_v> != cute::sizeof_bits_v>; + static constexpr bool IsArrayOfPointersGemm = cute::is_any_of_v; + static_assert(IsMixedInput || !IsArrayOfPointersGemm, "Only mixed input grouped RS GEMM is supported."); public: - using ElementA = detail::deduce_mixed_width_dtype_t<0, ElementPairA_>; - using ElementB = detail::deduce_mixed_width_dtype_t<0, ElementPairB_>; - static_assert(cute::is_tuple::value ^ cute::is_tuple::value || - (NeitherIsTuple && (sizeof_bits::value != sizeof_bits::value)), + using ElementA = detail::deduce_mixed_width_dtype_t<0, ElementA_>; + using ElementB = detail::deduce_mixed_width_dtype_t<0, ElementB_>; + + static_assert(!IsMixedInput || (cute::is_tuple::value ^ cute::is_tuple::value || + (NeitherIsTuple && (sizeof_bits::value != sizeof_bits::value))), "Either A OR B must be a tuple or the widths of A and B must be different."); static constexpr bool IsANarrow = sizeof_bits::value < sizeof_bits::value; template static auto get_stride(T const& t) { - if constexpr (not cute::is_layout::value) { + if constexpr (not cute::is_layout>::value) { return t; } else { @@ -422,8 +350,8 @@ public: using GmemLayoutATag = decltype(get_stride(GmemLayoutATag_{})); using GmemLayoutBTag = decltype(get_stride(GmemLayoutBTag_{})); - using ElementPairA = cute::conditional_t, ElementPairA_>; - using ElementPairB = cute::conditional_t, ElementPairB_>; + using ElementPairA = cute::conditional_t, ElementA_>; + using ElementPairB = cute::conditional_t, ElementB_>; static constexpr bool IsATransformed = cute::is_tuple::value; using ElementScale = cute::conditional_t; @@ -438,44 +366,71 @@ public: #endif static constexpr cute::GMMA::Major GmmaMajorA = detail::gmma_rs_tag_to_major_A(); static constexpr cute::GMMA::Major GmmaMajorB = detail::gmma_rs_tag_to_major_B(); + // If A is scaled, then we don't need to swap. Otherwise, we must ensure B goes to rmem and we must swap the operands. + static constexpr bool SwapAB = IsMixedInput ? !IsATransformed : detail::is_swapAB(); static constexpr bool IsWarpSpecializedTransposeB = detail::is_warpspecialized_transpose_B< ElementA, GmemLayoutATag, ElementB, GmemLayoutBTag, KernelScheduleType>(); - static_assert(!IsWarpSpecializedTransposeB, "Mixed input GEMM does not support WS transpose B."); - - // If A is scaled, then we don't need to swap. Otherwise, we must ensure B goes to RF and we must swap the operands. - static constexpr bool SwapAB = !IsATransformed; + static_assert(!IsMixedInput || !IsWarpSpecializedTransposeB, "Mixed input GEMM does not support WS transpose B."); // When we relax the above assertion, we must handle setting the tile mma GmmaMajorB correctly. static constexpr cute::GMMA::Major TiledMmaGmmaMajorB = SwapAB ? GmmaMajorA : GmmaMajorB; - using ElementMma = cute::conditional_t; - using AtomLayoutMNK = cute::conditional_t, + // For fp32 types, map to tf32 MMA value type. + using ElementAMma = cute::conditional_t, tfloat32_t, ElementA>; + using ElementBMma = cute::conditional_t, tfloat32_t, ElementB>; + + // Handle mixed dtypes and MMA. + using RealElementA = cute::conditional_t; + using RealElementB = cute::conditional_t; + using RealElementAMma = cute::conditional_t; + // Always the same for element B. + using RealElementBMma = RealElementB; + + static_assert(!IsMixedInput || TiledMmaGmmaMajorB == GMMA::Major::K || sizeof_bits::value == 16, + "Mixed input GEMM does not support MN major layout except for 16bit"); + + using AtomLayoutMNK = cute::conditional_t< + cute::is_any_of_v, Layout>, Layout>>; using TiledMma = decltype(cute::make_tiled_mma(cute::GMMA::rs_op_selector< - ElementMma, ElementMma, ElementAccumulator, TileShape_MNK, GMMA::Major::K, TiledMmaGmmaMajorB>(), AtomLayoutMNK{})); + RealElementAMma, RealElementBMma, ElementAccumulator, TileShape_MNK, GMMA::Major::K, GMMA::Major::K>(), AtomLayoutMNK{})); using GmemTiledCopyA = decltype(detail::sm90_cluster_shape_to_tma_atom(shape<1>(ClusterShape_MNK{}))); using GmemTiledCopyB = decltype(detail::sm90_cluster_shape_to_tma_atom(shape<0>(ClusterShape_MNK{}))); - using SmemLayoutAtomA = decltype(detail::rs_smem_selector(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{})), IsWarpSpecializedTransposeB>()); - using SmemLayoutAtomB = decltype(detail::rs_smem_selector(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{})), IsWarpSpecializedTransposeB>()); - using RealElementA = cute::conditional_t; - using RealElementB = cute::conditional_t; - static constexpr int PipelineStages = detail::compute_stage_count_or_override_single_affine_transformed_input(StageCountType{}); + static constexpr size_t SmemAlignmentA = cutlass::detail::alignment_for_swizzle(SmemLayoutAtomA{}); + static constexpr size_t SmemAlignmentB = cutlass::detail::alignment_for_swizzle(SmemLayoutAtomB{}); + static constexpr int SmemAlignment = static_cast(cute::max(SmemAlignmentA, SmemAlignmentB)); + + // Handle mixed dtype array GEMM's size of tensor map storage. + static constexpr size_t TensorMapStorage = sizeof(cute::TmaDescriptor) * size_t(IsMixedInput) * 4; + static constexpr int KernelSmemCarveout = static_cast(TensorMapStorage); + static constexpr int Sm90ReducedSmemCapacityBytes = detail::sm90_smem_capacity_bytes - KernelSmemCarveout; + + static constexpr int PipelineStages = IsMixedInput ? + detail::compute_stage_count_or_override_single_affine_transformed_input(StageCountType{}) : + detail::compute_stage_count_or_override(StageCountType{}); + + using DispatchPolicy = cute::conditional_t, + MainloopSm90TmaGmmaRmemAWarpSpecialized>; using SmemCopyAtomA = cute::conditional_t>; using SmemCopyAtomB = cute::conditional_t, void>; - - using DispatchPolicy = MainloopSm90TmaGmmaRmemAWarpSpecializedMixedInput; - + // We pack the scale data with the operand that will be optionally scaled and converted before MMA. - using StrideA = cute::conditional_t::value, GmemLayoutATag_, TagToStrideA_t>; - using StrideB = cute::conditional_t::value, GmemLayoutBTag_, TagToStrideB_t>; + using StrideA = cute::conditional_t>::value, GmemLayoutATag_, TagToStrideA_t>; + using StrideB = cute::conditional_t>::value, GmemLayoutBTag_, TagToStrideB_t>; using CollectiveOp = CollectiveMma< DispatchPolicy, @@ -495,6 +450,7 @@ public: cute::identity >; + static_assert(SmemAlignment == static_cast(cute::max(CollectiveOp::SmemAlignmentA, CollectiveOp::SmemAlignmentB))); }; ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -1008,15 +964,10 @@ static constexpr bool IsMixedWidthInput = IsDifferentWidth || (IsDifferentWidth #if ((__CUDACC_VER_MAJOR__ > 12) || ((__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ >= 1))) // Persistent schedules perform best for CUDA Toolkits with version >= 12.1 // KernelTmaWarpSpecializedCooperative requires TileShape_M to be at least 128 - using KernelTmaWarpSpecializedScheduleSameInput = cute::conditional_t(TileShape_MNK{}) == Int<64>{}, + using KernelTmaWarpSpecializedSchedule = cute::conditional_t(TileShape_MNK{}) == Int<64>{}, KernelTmaWarpSpecializedPingpong, KernelTmaWarpSpecializedCooperative>; - - using KernelTmaWarpSpecializedScheduleMixedInput = cute::conditional_t(TileShape_MNK{}) == Int<64>{}, - KernelTmaWarpSpecializedPingpongMixedInput, KernelTmaWarpSpecializedCooperativeMixedInput>; - - using KernelTmaWarpSpecializedSchedule = cute::conditional_t; #else - using KernelTmaWarpSpecializedSchedule = cute::conditional_t; + using KernelTmaWarpSpecializedSchedule = KernelTmaWarpSpecialized; #endif // Non-persistent schedule is a safer choice for CpAsync kernels due to register pressure diff --git a/include/cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp b/include/cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp index a3efc67e..6292f6e4 100644 --- a/include/cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp +++ b/include/cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp @@ -42,6 +42,7 @@ #include "cutlass/pipeline/pipeline.hpp" #include "cutlass/trace.h" #include "cutlass/detail/collective.hpp" +#include "cutlass/detail/collective/mixed_input_utils.hpp" #include "cute/arch/cluster_sm90.hpp" #include "cute/arch/copy_sm90.hpp" @@ -63,7 +64,7 @@ using namespace cute; template < int Stages, class ClusterShape, - class KernelSchedule, + class KernelSchedule_, class TileShape_, class ElementAOptionalTuple, class StrideA_, @@ -79,7 +80,7 @@ template < class SmemCopyAtomB_, class TransformB_> struct CollectiveMma< - MainloopSm90TmaGmmaRmemAWarpSpecializedMixedInput, + MainloopSm90TmaGmmaRmemAWarpSpecializedMixedInput, TileShape_, ElementAOptionalTuple, StrideA_, @@ -95,23 +96,31 @@ struct CollectiveMma< SmemCopyAtomB_, TransformB_> { -private: - template - static constexpr auto - get_logical_ptr(PointerType const* ptr) { - if constexpr (cute::sizeof_bits_v < 8) { - return subbyte_iterator(ptr); - } - else { - return ptr; - } - } - +public: enum class ConversionMode { DirectConvert, ConvertAndScale, ConvertAndScaleWithZero }; + + // + // Type Aliases + // + using DispatchPolicy = MainloopSm90TmaGmmaRmemAWarpSpecializedMixedInput; + using TileShape = TileShape_; + using KernelSchedule = KernelSchedule_; + +private: + template friend struct detail::MixedInputUtils; + using CollectiveType = CollectiveMma; + using Utils = detail::MixedInputUtils; using ScaleA = detail::deduce_mixed_width_dtype_t<1, ElementAOptionalTuple>; using ScaleB = detail::deduce_mixed_width_dtype_t<1, ElementBOptionalTuple>; @@ -119,12 +128,6 @@ private: using ZeroB = detail::deduce_mixed_width_dtype_t<2, ElementBOptionalTuple>; public: - // - // Type Aliases - // - using DispatchPolicy = MainloopSm90TmaGmmaRmemAWarpSpecializedMixedInput; - using TileShape = TileShape_; - static_assert(cute::is_tuple::value ^ cute::is_tuple::value, "Either A OR B must be a tuple. It must take the from {ElementOperand, [ElementScale]," "[ElementZero]}. Inputs in [] are optional."); @@ -146,14 +149,14 @@ public: using NonVoidStrideScale = cute::conditional_t< cute::is_void_v, cute::Stride<_1, int64_t, int64_t>, StrideScale>; - static_assert((IsATransformed && cutlass::gemm::detail::is_k_major()) || - (!IsATransformed && cutlass::gemm::detail::is_k_major()), + static_assert(( IsATransformed && (cutlass::gemm::detail::is_k_major() || is_layout::value)) || + (!IsATransformed && (cutlass::gemm::detail::is_k_major() || is_layout::value)), "The transformed type must be K-major."); static_assert(( IsATransformed && (sizeof(ElementB) == 2)) || (!IsATransformed && (sizeof(ElementA) == 2)) || - (cutlass::gemm::detail::is_k_major() && - cutlass::gemm::detail::is_k_major()), + ((cutlass::gemm::detail::is_k_major() || is_layout::value) && + (cutlass::gemm::detail::is_k_major() || is_layout::value)), "The unscaled element must be 2 bytes OR both inputs must be K-major"); static_assert(cutlass::gemm::detail::is_mn_major(), @@ -178,10 +181,10 @@ public: // We must ensure the type to be scaled goes to RF static constexpr bool SwapAB = !IsATransformed; - using InternalSmemLayoutAtomA = cute::conditional_t; - using InternalSmemLayoutAtomB = cute::conditional_t; - using InternalSmemCopyAtomA = cute::conditional_t; - using InternalSmemCopyAtomB = cute::conditional_t; + using SwappedSmemLayoutAtomA = cute::conditional_t; + using SwappedSmemLayoutAtomB = cute::conditional_t; + using SwappedSmemCopyAtomA = cute::conditional_t; + using SwappedSmemCopyAtomB = cute::conditional_t; // TMA converts f32 input to tf32 when copying from GMEM to SMEM // For all other types, cast to size equivalent uint type to avoid any rounding by TMA. @@ -189,20 +192,20 @@ public: static constexpr bool ConvertF32toTF32B = cute::is_same_v; using ConvertedElementA = cute::conditional_t>>; using ConvertedElementB = cute::conditional_t>>; - using RealInternalElementA = cute::conditional_t; - using RealInternalElementB = cute::conditional_t; - using InternalElementA = cute::conditional_t; - using InternalElementB = cute::conditional_t; - using InternalStrideA = cute::conditional_t; - using InternalStrideB = cute::conditional_t; + using RealSwappedElementA = cute::conditional_t; + using RealSwappedElementB = cute::conditional_t; + using SwappedElementA = cute::conditional_t; + using SwappedElementB = cute::conditional_t; + using SwappedStrideA = cute::conditional_t; + using SwappedStrideB = cute::conditional_t; using TransformA = TransformA_; using TransformB = TransformB_; - using InternalTransformA = cute::conditional_t; - using InternalTransformB = cute::conditional_t; + using SwappedTransformA = cute::conditional_t; + using SwappedTransformB = cute::conditional_t; - static constexpr int IsSubbyteA = cute::sizeof_bits_v < 8; - using TmaElementA = cute::conditional_t; + static constexpr int IsSubbyteA = cute::sizeof_bits_v < 8; + using TmaElementA = cute::conditional_t; using TmaElementScale = uint_bit_t >; // in case we have array. translating to uint to satisfy tma descriptor's specialization using ArchTag = typename DispatchPolicy::ArchTag; @@ -213,16 +216,16 @@ public: using PipelineParams = typename MainloopPipeline::Params; - using SmemLayoutAtomScale = Layout(InternalSmemLayoutAtomA{})), cute::Int<1>>>; + using SmemLayoutAtomScale = Layout(SwappedSmemLayoutAtomA{})), cute::Int<1>>>; using ScaleTileShape = decltype(make_shape(shape<0>(TileShape{}), shape<1>(SmemLayoutAtomScale{}))); - static_assert(cute::rank(InternalSmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); - static_assert((size<0>(TileShape{}) % size<0>(InternalSmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); - static_assert((size<2>(TileShape{}) % size<1>(InternalSmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert(cute::rank(SwappedSmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert((size<0>(TileShape{}) % size<0>(SwappedSmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert((size<2>(TileShape{}) % size<1>(SwappedSmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); - static_assert(cute::rank(InternalSmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); - static_assert((size<1>(TileShape{}) % size<0>(InternalSmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); - static_assert((size<2>(TileShape{}) % size<1>(InternalSmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert(cute::rank(SwappedSmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert((size<1>(TileShape{}) % size<0>(SwappedSmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert((size<2>(TileShape{}) % size<1>(SwappedSmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); static_assert(rank(SmemLayoutAtomScale{}) == 2, "SmemLayoutAtomScale must be rank 2"); static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomScale{})) == 0, "SmemLayoutAtomScale must equal the tile shape."); @@ -230,24 +233,8 @@ public: // Tile along modes in a way that maximizes the TMA box size. - template - static constexpr - CUTLASS_HOST_DEVICE - auto get_smem_layout(LayoutAtom layout_atom, TileShape const& tile_shape, Stride const& stride) { - if constexpr (not cute::is_layout::value) { - return tile_to_shape( - layout_atom, - append(tile_shape, Int{}), - cute::conditional_t< ::cutlass::gemm::detail::is_major<0,Stride>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}); - } - else { - auto gmem_tile = composition(stride, tile_shape); - return make_layout_like(append(gmem_tile, make_layout(Int{}, 0))); - } - } - - using SmemLayoutA = decltype(get_smem_layout(InternalSmemLayoutAtomA{}, select<0,2>(TileShape{}), InternalStrideA{})); - using SmemLayoutB = decltype(get_smem_layout(InternalSmemLayoutAtomB{}, select<1,2>(TileShape{}), InternalStrideB{})); + using SmemLayoutA = decltype(detail::get_smem_layout(SwappedSmemLayoutAtomA{}, select<0,2>(TileShape{}), SwappedStrideA{})); + using SmemLayoutB = decltype(detail::get_smem_layout(SwappedSmemLayoutAtomB{}, select<1,2>(TileShape{}), SwappedStrideB{})); // It is assumed that the scales and zero-points share the same smem layout using SmemLayoutScale = decltype(tile_to_shape( @@ -283,77 +270,12 @@ private: } } +public: static constexpr ConversionMode KernelConversionMode = get_conversion_mode(); static constexpr bool ModeHasScales = KernelConversionMode == ConversionMode::ConvertAndScale || KernelConversionMode == ConversionMode::ConvertAndScaleWithZero; static constexpr bool UseScaleLookupTable = KernelConversionMode == ConversionMode::ConvertAndScale && cutlass::detail::is_Array_v; - - static constexpr auto - elements_per_smem_scale() { - if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { - return 0; - } - else if constexpr (ModeHasScales) { - return cute::cosize_v; - } - else { - static_assert(cutlass::detail::dependent_false, "Type not handled in scale smem allocation."); - } - } - - static constexpr auto - elements_per_smem_zero() { - if constexpr (KernelConversionMode == ConversionMode::DirectConvert || - KernelConversionMode == ConversionMode::ConvertAndScale ) { - return 0; - } - else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { - return cute::cosize_v; - } - else { - static_assert(cutlass::detail::dependent_false, "Type not handled in scale smem allocation."); - } - } - - // These methods use some the public members of the class. For that reason, we define them after the public section. - static constexpr uint32_t - compute_tma_transaction_bytes_mk() { - return cutlass::bits_to_bytes(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast(cute::sizeof_bits_v)); - } - - static constexpr uint32_t - compute_tma_transaction_bytes_nk() { - return cutlass::bits_to_bytes(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast(cute::sizeof_bits_v)); - } - - static constexpr uint32_t - compute_tma_transaction_bytes_extra() { - if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { - return 0; - } - else if constexpr (ModeHasScales) { - constexpr uint32_t scale_tx_bytes = cutlass::bits_to_bytes(size<0>(SmemLayoutScale{}) * size<1>(SmemLayoutScale{}) * static_cast(cute::sizeof_bits_v)); - static_assert(scale_tx_bytes % 128 == 0, "Each scale stage must be 128B aligned."); // required by TMA - if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { - return scale_tx_bytes; - } - else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { - // Scale and zero share smem layout - constexpr uint32_t zero_tx_bytes = cutlass::bits_to_bytes(size<0>(SmemLayoutScale{}) * size<1>(SmemLayoutScale{}) * static_cast(cute::sizeof_bits_v)); - static_assert(zero_tx_bytes % 128 == 0, "Each zero stage must be 128B aligned."); // required by TMA - return scale_tx_bytes + zero_tx_bytes; - } - else { - static_assert(cutlass::detail::dependent_false, "Type not handled in tma transaction bytes computation."); - } - } - else { - static_assert(cutlass::detail::dependent_false, "Type not handled in tma transaction bytes computation."); - } - } - -public: static constexpr size_t SmemAlignmentA = cutlass::detail::alignment_for_swizzle(SmemLayoutA{}); static constexpr size_t SmemAlignmentB = cutlass::detail::alignment_for_swizzle(SmemLayoutB{}); @@ -365,11 +287,11 @@ public: struct SharedStorage { - static constexpr int scale_elements = elements_per_smem_scale(); - static constexpr int zero_elements = elements_per_smem_zero(); - struct TensorStorage : cute::aligned_struct { - cute::ArrayEngine> smem_A; - cute::ArrayEngine> smem_B; + static constexpr int scale_elements = Utils::elements_per_smem_scale(); + static constexpr int zero_elements = Utils::elements_per_smem_zero(); + struct TensorStorage { + CUTE_ALIGNAS(SmemAlignmentA) cute::ArrayEngine> smem_A; + CUTE_ALIGNAS(SmemAlignmentB) cute::ArrayEngine> smem_B; cute::ArrayEngine smem_scale; cute::ArrayEngine smem_zero; } tensors; @@ -393,53 +315,31 @@ public: uint32_t mma_promotion_interval = 4; }; - template - static constexpr - CUTLASS_HOST_DEVICE - auto get_gmem_layout(Shape const& shape, Stride const& stride) { - if constexpr (not cute::is_layout::value) { - return make_layout(shape, stride); - } - else { - return stride; - } - } - // Device side kernel params struct Params { - private: - using Outer = CollectiveMma; - public: // Assumption: StrideA is congruent with Problem_MK - using LayoutA = decltype(get_gmem_layout(repeat_like(InternalStrideA{}, int32_t(0)), InternalStrideA{})); - using LayoutB = decltype(get_gmem_layout(repeat_like(InternalStrideB{}, int32_t(0)), InternalStrideB{})); + using LayoutA = decltype(detail::get_gmem_layout(repeat_like(SwappedStrideA{}, int32_t(0)), SwappedStrideA{})); + using LayoutB = decltype(detail::get_gmem_layout(repeat_like(SwappedStrideB{}, int32_t(0)), SwappedStrideB{})); using TMA_A = decltype(make_tma_copy_A_sm90( GmemTiledCopyA{}, - make_tensor(Outer::get_logical_ptr(static_cast(nullptr)), LayoutA{}), + make_tensor(detail::get_logical_ptr(static_cast(nullptr)), LayoutA{}), SmemLayoutA{}(_,_,cute::Int<0>{}), TileShape{}, ClusterShape{})); // mcast along N mode for this M load, if any using TMA_Scale = decltype(make_tma_copy( GmemTiledCopyScale{}, - make_tensor(Outer::get_logical_ptr(static_cast(nullptr)), repeat_like(NonVoidStrideScale{}, int32_t(0)), NonVoidStrideScale{}), + make_tensor(detail::get_logical_ptr(static_cast(nullptr)), repeat_like(NonVoidStrideScale{}, int32_t(0)), NonVoidStrideScale{}), SmemLayoutScale{}(_,_,cute::Int<0>{}), ScaleTileShape{}, _1{})); // mcast along N mode for this M load, if any. Scale is ALWAYS loaded with A for RF kernel using TMA_Zero = decltype(make_tma_copy( GmemTiledCopyScale{}, - make_tensor(Outer::get_logical_ptr(static_cast(nullptr)), repeat_like(NonVoidStrideScale{}, int32_t(0)), NonVoidStrideScale{}), + make_tensor(detail::get_logical_ptr(static_cast(nullptr)), repeat_like(NonVoidStrideScale{}, int32_t(0)), NonVoidStrideScale{}), SmemLayoutScale{}(_,_,cute::Int<0>{}), ScaleTileShape{}, _1{})); // mcast along N mode for this M load, if any. Scale is ALWAYS loaded with A for RF kernel @@ -447,7 +347,7 @@ public: // Assumption: StrideB is congruent with Problem_NK using TMA_B = decltype(make_tma_copy_B_sm90( GmemTiledCopyB{}, - make_tensor(Outer::get_logical_ptr(static_cast(nullptr)), LayoutB{}), + make_tensor(detail::get_logical_ptr(static_cast(nullptr)), LayoutB{}), SmemLayoutB{}(_,_,cute::Int<0>{}), TileShape{}, ClusterShape{})); // mcast along M mode for this N load, if any @@ -459,8 +359,8 @@ public: int group_size; uint32_t tma_transaction_bytes = TmaTransactionBytes; int reload_factor = (group_size + size<2>(TileShape{}) - 1) / size<2>(TileShape{}); - InternalStrideA dA; - InternalStrideB dB; + SwappedStrideA dA; + SwappedStrideB dB; }; // @@ -481,26 +381,26 @@ public: N = get<0>(problem_shape_MNKL); } - InternalElementA const* ptr_A; - InternalStrideA dA; - InternalElementB const* ptr_B; - InternalStrideB dB; + SwappedElementA const* ptr_A; + SwappedStrideA dA; + SwappedElementB const* ptr_B; + SwappedStrideB dB; if constexpr (not SwapAB) { - ptr_A = reinterpret_cast(args.ptr_A); - ptr_B = reinterpret_cast(args.ptr_B); + ptr_A = reinterpret_cast(args.ptr_A); + ptr_B = reinterpret_cast(args.ptr_B); dA = args.dA; dB = args.dB; } else { - ptr_A = reinterpret_cast(args.ptr_B); - ptr_B = reinterpret_cast(args.ptr_A); + ptr_A = reinterpret_cast(args.ptr_B); + ptr_B = reinterpret_cast(args.ptr_A); dA = args.dB; dB = args.dA; } - Tensor tensor_a = make_tensor(get_logical_ptr(ptr_A), get_gmem_layout(make_shape(M,K,L), dA)); - Tensor tensor_b = make_tensor(get_logical_ptr(ptr_B), get_gmem_layout(make_shape(N,K,L), dB)); + Tensor tensor_a = make_tensor(detail::get_logical_ptr(ptr_A), detail::get_gmem_layout(make_shape(M,K,L), dA)); + Tensor tensor_b = make_tensor(detail::get_logical_ptr(ptr_B), detail::get_gmem_layout(make_shape(N,K,L), dB)); typename Params::TMA_A tma_load_a = make_tma_copy_A_sm90( GmemTiledCopyA{}, tensor_a, @@ -526,7 +426,7 @@ public: auto scale_k = (K + args.group_size - 1) / args.group_size; ElementScale const* ptr_S = args.ptr_S; StrideScale dS = args.dS; - Tensor tensor_scale = make_tensor(get_logical_ptr(ptr_S), make_layout(make_shape(M,scale_k,L), dS)); + Tensor tensor_scale = make_tensor(detail::get_logical_ptr(ptr_S), make_layout(make_shape(M,scale_k,L), dS)); tma_load_scale = make_tma_copy( GmemTiledCopyScale{}, tensor_scale, @@ -538,7 +438,7 @@ public: return { tma_load_a, tma_load_b, tma_load_scale, tma_load_zero, scale_k, args.group_size, tma_transaction_bytes + TmaTransactionBytesExtra, (args.group_size + size<2>(TileShape{}) - 1) / size<2>(TileShape{}), dA, dB }; } else if constexpr(KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { - Tensor tensor_zero = make_tensor(get_logical_ptr(args.ptr_Z), make_layout(make_shape(M,scale_k,L), dS)); + Tensor tensor_zero = make_tensor(detail::get_logical_ptr(args.ptr_Z), make_layout(make_shape(M,scale_k,L), dS)); tma_load_zero = make_tma_copy( GmemTiledCopyScale{}, tensor_zero, @@ -565,10 +465,10 @@ public: auto [M,N,K,L] = problem_shape_MNKL; constexpr int min_tma_aligned_elements_A = tma_alignment_bits / cutlass::sizeof_bits::value; - bool check_aligned_A = cutlass::detail::check_alignment(get_gmem_layout(cute::make_shape(M,K,L), args.dA)); + bool check_aligned_A = cutlass::detail::check_alignment(detail::get_gmem_layout(cute::make_shape(M,K,L), args.dA)); constexpr int min_tma_aligned_elements_B = tma_alignment_bits / cutlass::sizeof_bits::value; - bool check_aligned_B = cutlass::detail::check_alignment(get_gmem_layout(cute::make_shape(N,K,L), args.dB)); + bool check_aligned_B = cutlass::detail::check_alignment(detail::get_gmem_layout(cute::make_shape(N,K,L), args.dB)); bool check_aligned_S = true; bool check_aligned_Z = true; @@ -623,9 +523,9 @@ public: } static constexpr int K_PIPE_MAX = DispatchPolicy::Stages; - static constexpr uint32_t TmaTransactionBytesMK = compute_tma_transaction_bytes_mk(); - static constexpr uint32_t TmaTransactionBytesNK = compute_tma_transaction_bytes_nk(); - static constexpr uint32_t TmaTransactionBytesExtra = compute_tma_transaction_bytes_extra(); + static constexpr uint32_t TmaTransactionBytesMK = Utils::compute_tma_transaction_bytes_mk(); + static constexpr uint32_t TmaTransactionBytesNK = Utils::compute_tma_transaction_bytes_nk(); + static constexpr uint32_t TmaTransactionBytesExtra = Utils::compute_tma_transaction_bytes_extra(); static constexpr uint32_t TmaTransactionBytes = TmaTransactionBytesMK + TmaTransactionBytesNK + TmaTransactionBytesExtra; /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance @@ -665,8 +565,8 @@ public: // TMA requires special handling of strides to deal with coord codomain mapping // Represent the full tensors -- get these from TMA - Tensor mA_mkl = mainloop_params.tma_load_a.get_tma_tensor(shape(get_gmem_layout(make_shape(M,K,L), mainloop_params.dA))); // (m,k,l) - Tensor mB_nkl = mainloop_params.tma_load_b.get_tma_tensor(shape(get_gmem_layout(make_shape(N,K,L), mainloop_params.dB))); // (n,k,l) + Tensor mA_mkl = mainloop_params.tma_load_a.get_tma_tensor(shape(detail::get_gmem_layout(make_shape(M,K,L), mainloop_params.dA))); // (m,k,l) + Tensor mB_nkl = mainloop_params.tma_load_b.get_tma_tensor(shape(detail::get_gmem_layout(make_shape(N,K,L), mainloop_params.dB))); // (n,k,l) // Make tiled views, defer the slice Tensor gA_mkl = local_tile(mA_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k,l) @@ -777,7 +677,7 @@ public: } } - auto extra_input_partitions = partition_extra_tma_inputs(mainloop_params, load_inputs, shared_tensors, cluster_local_block_id, m_coord, l_coord); + auto extra_input_partitions = Utils::partition_extra_tma_inputs(mainloop_params, load_inputs, shared_tensors, cluster_local_block_id, m_coord, l_coord); // Mainloop CUTLASS_PRAGMA_NO_UNROLL @@ -866,11 +766,11 @@ public: static_assert(is_rmem::value, "C tensor must be rmem resident."); static_assert(cute::rank(SmemLayoutA{}) == 3, "Smem layout must be rank 3."); static_assert(cute::rank(SmemLayoutB{}) == 3, "Smem layout must be rank 3."); - static_assert(cute::rank(InternalSmemLayoutAtomA{}) == 2, "InternalSmemLayoutAtomA must be rank 2."); - static_assert(cute::rank(InternalSmemLayoutAtomB{}) == 2, "InternalSmemLayoutAtomB must be rank 2."); - static_assert(!cute::is_void_v, + static_assert(cute::rank(SwappedSmemLayoutAtomA{}) == 2, "SwappedSmemLayoutAtomA must be rank 2."); + static_assert(cute::rank(SwappedSmemLayoutAtomB{}) == 2, "SwappedSmemLayoutAtomB must be rank 2."); + static_assert(!cute::is_void_v, "SM90 GMMA mainloops must specify a non-void copy atom for RF sourced instructions."); - static_assert(cute::is_void_v, + static_assert(cute::is_void_v, "SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions."); // Obtain warp index @@ -905,7 +805,17 @@ public: // Allocate fragments and descriptors Tensor tCrA_mma = mma_thread_slice.partition_fragment_A(sA(_,_,Int<0>{})); // (MMA,MMA_M,MMA_K,PIPE) - Tensor tCrA_load = make_fragment_like(tCrA_mma); + + Tensor tCrA_load = [&]{ + if constexpr (not is_layout::value) { + // Make register tensor with MMA layout + return make_fragment_like(tCrA_mma); + } + else { + // Make register tensor matching smem layout, converter will take care of de-swizzling + return make_tensor_like(tCsA(_,_,_,Int<0>{})); + } + }(); Tensor tCsB = mma_warpgroup_slice.partition_B(sB); // (MMA,MMA_N,MMA_K,PIPE) Tensor tCrB = mma_warpgroup_slice.make_fragment_B(tCsB); // (MMA,MMA_N,MMA_K,PIPE) @@ -913,14 +823,14 @@ public: // // Copy Atom A retiling // - auto smem_tiled_copy_A = make_tiled_copy_A(InternalSmemCopyAtomA{}, tiled_mma); + auto smem_tiled_copy_A = make_tiled_copy_A(SwappedSmemCopyAtomA{}, tiled_mma); auto smem_thr_copy_A = smem_tiled_copy_A.get_thread_slice(warp_group_thread_idx); Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA_load); // (CPY,CPY_M,CPY_K) // Partition of thread -> shared and thread -> RF - auto partitioned_extra_info = partition_extra_mma_info(mma_thread_slice, shared_tensors); - auto copy_partitions_extra_info = retile_extra_mma_info(tiled_mma, partitioned_extra_info, warp_group_thread_idx); + auto partitioned_extra_info = Utils::partition_extra_mma_info(mma_thread_slice, shared_tensors); + auto copy_partitions_extra_info = Utils::retile_extra_mma_info(tiled_mma, partitioned_extra_info, warp_group_thread_idx); CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view)); // CPY_M CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCrA_copy_view)); // CPY_K @@ -943,7 +853,9 @@ public: warpgroup_fence_operand(accum); constexpr int K_BLOCK_MAX = size<2>(tCrA_load); - + constexpr int K_WAIT_MAX = cute::min(K_BLOCK_MAX - 1, 7); + static_assert(K_BLOCK_MAX >= 4, "Consider increasing TileShapeK"); + ConsumerToken barrier_token = {BarrierStatus::WaitAgain}; // first k tile { @@ -956,13 +868,13 @@ public: barrier_token = pipeline.consumer_try_wait(smem_pipe_read); // copy smem->rmem for A operand - copy_A_and_extra_info(smem_tiled_copy_A, tCsA, tCrA_copy_view, + Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view, partitioned_extra_info, copy_partitions_extra_info, 0, read_stage); if (K_BLOCK_MAX > 1) { // prefetch next block - copy_A_and_extra_info(smem_tiled_copy_A, tCsA, tCrA_copy_view, + Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view, partitioned_extra_info, copy_partitions_extra_info, 1, read_stage); } - transform_A_kblock(tCrA_load, tCrA_mma, partitioned_extra_info, 0); + Utils::dequantize_A_kblock(tCrA_load, tCrA_mma, partitioned_extra_info, 0); // Unroll the K mode manually to set scale D to 1 CUTLASS_PRAGMA_UNROLL @@ -974,11 +886,11 @@ public: warpgroup_commit_batch(); if (k_block < K_BLOCK_MAX - 2) { // prefetch next block - copy_A_and_extra_info(smem_tiled_copy_A, tCsA, tCrA_copy_view, + Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view, partitioned_extra_info, copy_partitions_extra_info, k_block + 2, read_stage); } if (k_block < K_BLOCK_MAX - 1) { - transform_A_kblock(tCrA_load, tCrA_mma, partitioned_extra_info, k_block + 1); + Utils::dequantize_A_kblock(tCrA_load, tCrA_mma, partitioned_extra_info, k_block + 1); } } @@ -986,14 +898,14 @@ public: if (k_tile_count > 0) { // Wait for K_BLOCK_MAX - 1 to be in flight to ensure that it is safe to overwrite the A registers for the first mma. pipeline.consumer_wait(smem_pipe_read, barrier_token); - copy_A_and_extra_info(smem_tiled_copy_A, tCsA, tCrA_copy_view, + Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view, partitioned_extra_info, copy_partitions_extra_info, 0, smem_pipe_read.index()); if (K_BLOCK_MAX > 1) { // prefetch next block - copy_A_and_extra_info(smem_tiled_copy_A, tCsA, tCrA_copy_view, + Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view, partitioned_extra_info, copy_partitions_extra_info, 1, smem_pipe_read.index()); } - warpgroup_wait(); - transform_A_kblock(tCrA_load, tCrA_mma, partitioned_extra_info, 0); + warpgroup_wait(); + Utils::dequantize_A_kblock(tCrA_load, tCrA_mma, partitioned_extra_info, 0); } } @@ -1024,7 +936,7 @@ public: tiled_mma.accumulate_ = GMMA::ScaleOut::One; warpgroup_commit_batch(); - warpgroup_wait(); // We have K_BLOCK_MAX - 1 GMMA instructions pending for this stage, so we can release prior barrier + warpgroup_wait(); // We have K_BLOCK_MAX - 1 GMMA instructions pending for this stage, so we can release prior barrier if (k_block == K_BLOCK_MAX - 1) { pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it ++smem_pipe_release; @@ -1036,20 +948,20 @@ public: if (k_block == K_BLOCK_MAX - 1) { pipeline.consumer_wait(smem_pipe_read, barrier_token); - copy_A_and_extra_info(smem_tiled_copy_A, tCsA, tCrA_copy_view, + Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view, partitioned_extra_info, copy_partitions_extra_info, 0, smem_pipe_read.index()); if (K_BLOCK_MAX > 1) { // prefetch next block - copy_A_and_extra_info(smem_tiled_copy_A, tCsA, tCrA_copy_view, + Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view, partitioned_extra_info, copy_partitions_extra_info, 1, smem_pipe_read.index()); } - transform_A_kblock(tCrA_load, tCrA_mma, partitioned_extra_info, 0); + Utils::dequantize_A_kblock(tCrA_load, tCrA_mma, partitioned_extra_info, 0); } else { if (k_block < K_BLOCK_MAX - 2) { // prefetch next block - copy_A_and_extra_info(smem_tiled_copy_A, tCsA, tCrA_copy_view, + Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view, partitioned_extra_info, copy_partitions_extra_info, k_block + 2, read_stage); } - transform_A_kblock(tCrA_load, tCrA_mma, partitioned_extra_info, k_block + 1); + Utils::dequantize_A_kblock(tCrA_load, tCrA_mma, partitioned_extra_info, k_block + 1); } } warpgroup_fence_operand(accum); @@ -1077,27 +989,25 @@ public: tiled_mma.accumulate_ = GMMA::ScaleOut::One; warpgroup_commit_batch(); - warpgroup_wait(); + warpgroup_wait(); if (k_block == K_BLOCK_MAX - 1) { // release prior barrier pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it ++smem_pipe_release; } if (k_block < K_BLOCK_MAX - 2) { // prefetch next block - copy_A_and_extra_info(smem_tiled_copy_A, tCsA, tCrA_copy_view, + Utils::copy_tensors_MK(smem_tiled_copy_A, tCsA, tCrA_copy_view, partitioned_extra_info, copy_partitions_extra_info, k_block + 2, read_stage); } if (k_block < K_BLOCK_MAX - 1) { - copy_A_and_extra_info(smem_tiled_copy_A, tCsA, tCrA_copy_view, - partitioned_extra_info, copy_partitions_extra_info, k_block + 1, read_stage); - transform_A_kblock(tCrA_load, tCrA_mma, partitioned_extra_info, k_block + 1); + Utils::dequantize_A_kblock(tCrA_load, tCrA_mma, partitioned_extra_info, k_block + 1); } } } warpgroup_fence_operand(accum); } - + /// Perform a Consumer Epilogue to release all buffers CUTLASS_DEVICE void mma_tail(MainloopPipeline pipeline, PipelineState smem_pipe_release, int k_tile_count) { @@ -1115,442 +1025,6 @@ public: ++smem_pipe_release; } } - -private: - /// Utilities for any additional inputs inside of the TMA load - template - CUTLASS_DEVICE - auto partition_extra_tma_inputs( - Params const& mainloop_params, - cute::tuple const& load_inputs, - TensorStorage& shared_tensors, - uint2 const& cluster_local_block_id, - int const m_coord, - int const l_coord) { - - if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { - return cute::make_tuple(); - } - else if constexpr (ModeHasScales) { - Tensor sS = make_tensor(make_smem_ptr(shared_tensors.smem_scale.begin()), SmemLayoutScale{}); // (BLK_M,BLK_K,PIPE) - Tensor gS_mkl = get<2>(load_inputs); - auto block_tma_s = mainloop_params.tma_load_scale.get_slice(cluster_local_block_id.y); - Tensor gS = gS_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k) - - Tensor tSgS = block_tma_s.partition_S(gS); // (TMA,TMA_M,TMA_K,k) - Tensor tSsS = block_tma_s.partition_D(sS); // (TMA,TMA_M,TMA_K,PIPE) - if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { - return cute::make_tuple(tSgS, tSsS); - } - else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { - Tensor sZ = make_tensor(make_smem_ptr(shared_tensors.smem_zero.begin()), SmemLayoutScale{}); // (BLK_M,BLK_K,PIPE) - Tensor gZ_mkl = get<3>(load_inputs); - auto block_tma_z = mainloop_params.tma_load_zero.get_slice(cluster_local_block_id.y); - Tensor gZ = gZ_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k) - - Tensor tZgZ = block_tma_z.partition_S(gZ); // (TMA,TMA_M,TMA_K,k) - Tensor tZsZ = block_tma_z.partition_D(sZ); // (TMA,TMA_M,TMA_K,PIPE) - return cute::make_tuple(tSgS, tSsS, tZgZ, tZsZ); - } - else { - static_assert(cutlass::detail::dependent_false, "Conversion mode not handled for input partitioning."); - } - } - else { - static_assert(cutlass::detail::dependent_false, "Conversion mode not handled for input partitioning."); - } - } - - /// Utilities for partitioning extra inputs for loading from smem in the mainloop. - template - CUTLASS_DEVICE - auto partition_extra_mma_info( - ThreadMma const& mma_thread_slice, - TensorStorage& shared_tensors) { - - if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { - // nothing to do - return cute::make_tuple(); - } - else if constexpr (UseScaleLookupTable) { - Tensor sS = make_tensor(make_smem_ptr(shared_tensors.smem_scale.begin()), SmemLayoutScale{});// (BLK_M,BLK_SCALE_K,PIPE) - Tensor tCsS = mma_thread_slice.partition_A(sS); - Tensor tCrS_neg = make_tensor(mma_thread_slice.partition_fragment_A(sS(_,_,Int<0>{})).layout()); - Tensor tCrS_pos = make_tensor(mma_thread_slice.partition_fragment_A(sS(_,_,Int<0>{})).layout()); - - if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { - return cute::make_tuple(tCsS, tCrS_neg, tCrS_pos); - } - } - else if constexpr (ModeHasScales) { - Tensor sS = make_tensor(make_smem_ptr(shared_tensors.smem_scale.begin()), SmemLayoutScale{});// (BLK_M,BLK_SCALE_K,PIPE) - Tensor tCsS = mma_thread_slice.partition_A(sS); - Tensor tCrS = make_tensor(mma_thread_slice.partition_fragment_A(sS(_,_,Int<0>{})).layout()); - - if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { - return cute::make_tuple(tCsS, tCrS); - } - else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { - Tensor sZ = make_tensor(make_smem_ptr(shared_tensors.smem_zero.begin()), SmemLayoutScale{});// (BLK_M,BLK_SCALE_K,PIPE) - Tensor tCsZ = mma_thread_slice.partition_A(sZ); - Tensor tCrZ = make_tensor(mma_thread_slice.partition_fragment_A(sZ(_,_,Int<0>{})).layout()); - return cute::make_tuple(tCsS, tCrS, tCsZ, tCrZ); - } - else { - static_assert(cutlass::detail::dependent_false, "Conversion mode not handled in A -> RF path."); - } - } - else { - static_assert(cutlass::detail::dependent_false, "Conversion mode not handled in A -> RF path."); - } - } - - /// Returns the tiled copy and copy views for the extra inputs. - template - CUTLASS_DEVICE - auto retile_extra_mma_info( - TiledMma const& tiled_mma, - cute::tuple& partitioned_extra_info, - int const warp_group_thread_idx) { - - if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { - // nothing to do - return cute::make_tuple(); - } - else if constexpr (ModeHasScales) { - auto smem_tiled_copy_S = make_tiled_copy_A(SmemCopyAtomScale{}, tiled_mma); - auto smem_thr_copy_S = smem_tiled_copy_S.get_thread_slice(warp_group_thread_idx); - Tensor tCrS_copy_view = smem_thr_copy_S.retile_D(cute::get<1>(partitioned_extra_info)); // (CPY,CPY_M,CPY_K) - - if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { - return cute::make_tuple(smem_tiled_copy_S, tCrS_copy_view); - } - else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { - Tensor tCrZ_copy_view = smem_thr_copy_S.retile_D(cute::get<3>(partitioned_extra_info)); // (CPY,CPY_M,CPY_K) - return cute::make_tuple(smem_tiled_copy_S, tCrS_copy_view, tCrZ_copy_view); - } - else { - static_assert(cutlass::detail::dependent_false, "Conversion mode not handled in A -> RF path."); - } - } - else { - static_assert(cutlass::detail::dependent_false, "Conversion mode not handled in A -> RF path."); - } - } - - /// Utilities to copy A and extra inputs from smem to RF - template - CUTLASS_DEVICE - void copy_A_and_extra_info( - SmemTiledCopyA const& smem_tiled_copy_A, - TensorASmemView const& tCsA, - TensorACopyView& tCrA_copy_view, - cute::tuple const& partitioned_mma_extra_info, - cute::tuple const& tiled_copy_and_views, - int k_block, - int read_stage) { - - copy(smem_tiled_copy_A, tCsA(_,_,k_block,read_stage), tCrA_copy_view(_,_,k_block)); - - if (k_block == 0) { - // We are starting a new k-tile so copy the scale - if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { - // nothing to do - } - else if constexpr (ModeHasScales) { - auto smem_tiled_copy_S = cute::get<0>(tiled_copy_and_views); - auto tCrS_copy_view = cute::get<1>(tiled_copy_and_views); - auto tCsS = cute::get<0>(partitioned_mma_extra_info); - copy(smem_tiled_copy_S, tCsS(_,_,k_block,read_stage), tCrS_copy_view(_,_,k_block)); - if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { - // Nothing extra to do - } else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { - auto tCsZ = cute::get<2>(partitioned_mma_extra_info); - auto tCrZ_copy_view = cute::get<2>(tiled_copy_and_views); - copy(smem_tiled_copy_S, tCsZ(_,_,k_block,read_stage), tCrZ_copy_view(_,_,k_block)); - } else { - static_assert(cutlass::detail::dependent_false, "Conversion mode not handled in A -> RF path."); - } - } - else { - static_assert(cutlass::detail::dependent_false, "Conversion mode not handled in A -> RF path."); - } - } - } - - // Helper functions to select packing for conversion - template - struct select_packing { // Naive packing policy - static constexpr auto value() { - return Int, sizeof_bits_v))>{}; - } - }; - - CUTLASS_DEVICE - static uint32_t to_reg(Array const& source) { - return static_cast( - reinterpret_cast(source)); - } - CUTLASS_DEVICE - static uint32_t to_reg(Array const& source) { - return reinterpret_cast(source); - } - // The core converter uses a lookup table to converts i4 -> 8 bit value. - template - CUTLASS_DEVICE - static Array lookup_table_convert( - cute::Int _, - Array const& source, - TensorPos const& scale_neg, - TensorNeg const& scale_pos, - int scale_idx) { - - static_assert(N == 4 || N == 8); - uint32_t res[N / 4]; - - // View the input as reg - uint32_t reg = to_reg(source); - - // Determines if to get from the signed or unsigned candidates - static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; - uint32_t sign; // ((reg & 0x88888888) | 0x64206420) >> 1 - asm volatile( - "{\n" - " lop3.b32 %0, %1, %2, %3, %4;\n" \ - "}\n" - : "=r"(sign) - : "r"(reg), "n"(0x88888888), "n"(0x64206420), "n"(immLut) - ); - sign = sign >> 1; - - // Ignore sign bit when indexing into LUT - uint32_t lut_idx = reg & 0x77777777; - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < N / 4; ++i, lut_idx >>=16, sign >>=16) { - Array const& _scale_neg = reinterpret_cast const&>(scale_neg[scale_idx + i * 4]); - Array const& _scale_pos = reinterpret_cast const&>(scale_pos[scale_idx + i * 4]); - asm volatile( - "{\n" - " .reg .b32 pos, neg ;\n" \ - " prmt .b32 neg, %3, %4, %1 ;\n" \ - " prmt .b32 pos, %5, %6, %1 ;\n" \ - " prmt .b32 %0, pos, neg, %2 ;\n" \ - "}\n" - : "=r"(res[i]) - : "r"(lut_idx), "r"(sign), "r"(_scale_neg[0]), "r"(_scale_neg[1]), "r"(_scale_pos[0]), "r"(_scale_pos[1]) - ); - } - return reinterpret_cast&>(res); - } - - template - CUTLASS_DEVICE - static void static_check_scale(Layout const& tensor) { - static_assert(shape<0>(Layout{}) >= 4 && stride<0>(Layout{}) == 0, "At least 4 adjacent weights in a thread must share the same scale."); - } - template - CUTLASS_DEVICE - static void static_check_scale(Tensor const& tensor) { - static_check_scale(flatten(Layout{})); - } - - /// Utilities to transform A. - template - CUTLASS_DEVICE - void transform_A_kblock( - Tensor const& tCrA_load, - Tensor& tCrA_mma, - cute::tuple const& partitioned_extra_info, - int const k_block) { - - static_assert(is_rmem::value, "Input tensor for A conversion must come from registers"); - static_assert(is_rmem::value, "Output tensor for A conversion must come from registers"); - static_assert(cosize_v == cosize_v); - static_assert(size_v == cosize_v); - static_assert(size_v == cosize_v); - using SrcType = typename EngineIn::value_type; - using DstType = typename EngineOut::value_type; - - auto const& src = tCrA_load(_, _, k_block); - auto const& dst = tCrA_mma(_, _, k_block); - auto pSrc = raw_pointer_cast(src.data()); - auto pDst = const_cast(raw_pointer_cast(dst.data())); - constexpr int num_elements = decltype(size(src))::value; - - if constexpr (KernelConversionMode == ConversionMode::DirectConvert) { - constexpr int pack = decltype(select_packing::value())::value; - using Converter = cutlass::NumericArrayConverter; - using SrcArray = cutlass::Array; - using DstArray = cutlass::Array; - constexpr int iters = num_elements / pack; - - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < iters; ++i) { - SrcArray const* pSrcArr = reinterpret_cast(pSrc) + i; - DstArray* pDstArr = reinterpret_cast(pDst) + i; - *pDstArr = Converter::convert(*pSrcArr); - } - } - else if constexpr (UseScaleLookupTable) { - static_assert(is_same_v, "Lookup table only supports int4 being the quant type now."); - static_assert(sizeof_bits_v == 64, "Lookup table only supports 8 8bit scale values now."); - static_assert(num_elements % 4 == 0 && num_elements >= 4, "Lookup table requires a vector size of 4x when converting."); - constexpr int pack = num_elements % 8 == 0? 8 : 4; - constexpr int iters = num_elements / pack; - using SrcArray = cutlass::Array; - using DstArray = cutlass::Array; - - auto const& tCrS_neg = cute::get<1>(partitioned_extra_info); - auto const& tCrS_pos = cute::get<2>(partitioned_extra_info); - auto const& scale_neg = tCrS_neg(_, _, k_block); - auto const& scale_pos = tCrS_pos(_, _, k_block); - CUTE_STATIC_ASSERT_V(size(src) == size(scale_neg)); - - static_check_scale(scale_neg); - static_check_scale(scale_pos); - if (k_block == 0) { - auto pNeg = raw_pointer_cast(tCrS_neg.data()); - auto pPos = const_cast(raw_pointer_cast(tCrS_pos.data())); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < cosize(tCrS_neg.layout()); ++i) - { - // pPos[i] = pNeg[i] & 0x7F7F7F7F7F7F7F00; - cutlass::Array const& _scale_neg = reinterpret_cast const&>(pNeg[i]); - cutlass::Array & _scale_pos = reinterpret_cast &>(pPos[i]); - asm volatile( - "{\n" - " and .b32 %0, %2, %4 ;\n" \ - " and .b32 %1, %3, %5 ;\n" \ - "}\n" - : "=r"(_scale_pos[0]), "=r"(_scale_pos[1]) - : "r"(_scale_neg[0]), "r"(_scale_neg[1]), "n"(0x7F7F7F00), "n"(0x7F7F7F7F) - ); - } - } - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < iters; i ++) { - SrcArray const* pSrcArr = reinterpret_cast(raw_pointer_cast(src.data())) + i; - DstArray* pDstArr = reinterpret_cast(raw_pointer_cast(dst.data())) + i; - - *pDstArr = lookup_table_convert(Int{}, *pSrcArr, scale_neg, scale_pos, i * pack); - } - } - else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { - auto const& scales = cute::get<1>(partitioned_extra_info)(_, _, k_block); - CUTE_STATIC_ASSERT_V(size(src) == size(scales)); - - if constexpr (is_same_v) { - constexpr int pack = decltype(select_packing::value())::value; - using Converter = cutlass::NumericArrayConverter; - using SrcArray = cutlass::Array; - using DstArray = cutlass::Array; - constexpr int iters = num_elements / pack; - - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < iters; ++i) { - SrcArray const* pSrcArr = reinterpret_cast(pSrc) + i; - DstArray* pDstArr = reinterpret_cast(pDst) + i; - *pDstArr = Converter::convert(*pSrcArr); - CUTLASS_PRAGMA_UNROLL - for (int j = 0; j < pack; ++j) { - (*pDstArr)[j] = (*pDstArr)[j] * scales[i*pack + j]; - } - } - } - else { - constexpr int pack1 = decltype(select_packing::value())::value; - constexpr int pack2 = decltype(select_packing::value())::value; - constexpr int pack = cute::gcd(pack1, pack2); - using Converter1 = cutlass::NumericArrayConverter; - using Converter2 = cutlass::NumericArrayConverter; - using SrcArray = cutlass::Array; - using DstArray = cutlass::Array; - using StageArray = cutlass::Array; - constexpr int iters = num_elements / pack; - - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < iters; ++i) { - SrcArray const* pSrcArr = reinterpret_cast(pSrc) + i; - DstArray* pDstArr = reinterpret_cast(pDst) + i; - StageArray stageArr; - stageArr = Converter1::convert(*pSrcArr); - CUTLASS_PRAGMA_UNROLL - for (int j = 0; j < pack; ++j) { - stageArr[j] = stageArr[j] * scales[i*pack + j]; - } - *pDstArr = Converter2::convert(stageArr); - } - } - } - else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { - static_assert(is_same_v, "ElementScale and ElementZero must be the same."); - auto const& scales = cute::get<1>(partitioned_extra_info)(_, _, k_block); - auto const& zeros = cute::get<3>(partitioned_extra_info)(_, _, k_block); - CUTE_STATIC_ASSERT_V(size(src) == size(scales)); - CUTE_STATIC_ASSERT_V(size(src) == size(zeros)); - - if constexpr (is_same_v) { - constexpr int pack = decltype(select_packing::value())::value; - using Converter = cutlass::NumericArrayConverter; - using SrcArray = cutlass::Array; - using DstArray = cutlass::Array; - constexpr int iters = num_elements / pack; - - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < iters; ++i) { - SrcArray const* pSrcArr = reinterpret_cast(pSrc) + i; - DstArray* pDstArr = reinterpret_cast(pDst) + i; - *pDstArr = Converter::convert(*pSrcArr); - CUTLASS_PRAGMA_UNROLL - for (int j = 0; j < pack; ++j) { - (*pDstArr)[j] = (*pDstArr)[j] * scales[i*pack + j] + zeros[i*pack + j]; - } - } - } - else { - constexpr int pack1 = decltype(select_packing::value())::value; - constexpr int pack2 = decltype(select_packing::value())::value; - constexpr int pack = cute::gcd(pack1, pack2); - using Converter1 = cutlass::NumericArrayConverter; - using Converter2 = cutlass::NumericArrayConverter; - using SrcArray = cutlass::Array; - using DstArray = cutlass::Array; - using StageArray = cutlass::Array; - constexpr int iters = num_elements / pack; - - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < iters; ++i) { - SrcArray const* pSrcArr = reinterpret_cast(pSrc) + i; - DstArray* pDstArr = reinterpret_cast(pDst) + i; - StageArray stageArr; - stageArr = Converter1::convert(*pSrcArr); - CUTLASS_PRAGMA_UNROLL - for (int j = 0; j < pack; ++j) { - stageArr[j] = stageArr[j] * scales[i*pack + j] + zeros[i*pack + j]; - } - *pDstArr = Converter2::convert(stageArr); - } - } - return; - } - else { - static_assert(cutlass::detail::dependent_false, "No A data is loaded."); - } - } }; ///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/gemm/dispatch_policy.hpp b/include/cutlass/gemm/dispatch_policy.hpp index 904e6af3..fa275bda 100644 --- a/include/cutlass/gemm/dispatch_policy.hpp +++ b/include/cutlass/gemm/dispatch_policy.hpp @@ -126,11 +126,6 @@ struct KernelTmaWarpSpecializedCooperativeFP8FastAccum: KernelTmaWarpSpecialized struct KernelPtrArrayTmaWarpSpecializedCooperativeFP8FastAccum : KernelPtrArrayTmaWarpSpecializedCooperative { }; struct KernelPtrArrayTmaWarpSpecializedPingpongFP8FastAccum : KernelPtrArrayTmaWarpSpecializedPingpong { }; -// Policies to opt into mixed type GEMMs -struct KernelTmaWarpSpecializedMixedInput : KernelTmaWarpSpecialized { }; -struct KernelTmaWarpSpecializedPingpongMixedInput : KernelTmaWarpSpecializedPingpong { }; -struct KernelTmaWarpSpecializedCooperativeMixedInput: KernelTmaWarpSpecializedCooperative { }; - ////////////////////////////////////////////////////////////////////////////// // Policies for dispatch of epilogue @@ -264,11 +259,8 @@ struct MainloopSm90TmaGmmaRmemAWarpSpecializedMixedInput { using Schedule = KernelSchedule; static_assert( cute::is_same_v || - cute::is_same_v || cute::is_same_v || - cute::is_same_v || - cute::is_same_v || - cute::is_same_v, + cute::is_same_v, "KernelSchedule must be one of the warp specialized policies"); };