确保 `cutlass` 库已经安装，然后执行以下命令进行自动调优。

- 对于89架构的 GPU，CUDA 版本至少为12.4

- 对于90架构的 GPU，CUDA 版本至少为12.0

#include "fuse_dual_gemm_act_template_3x.h"

bool fp8_fp8_gemm_scale_bias_act(GemmEpilogueAllParams params);

#pragma once

#include "fp8_common.h"

#include "cutlass/cutlass.h"

#include "cutlass/float8.h"

#include "cutlass/gemm/device/gemm_universal.h"

#include "cutlass/util/packed_stride.hpp"

#include "cutlass/gemm/dispatch_policy.hpp"

#include "cutlass/gemm/collective/collective_builder.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/gemm/kernel/tile_scheduler.hpp"

#include "cutlass_kernels/gemm/collective/collective_builder_gated.hpp"

#include "cutlass_kernels/gemm/kernel/gemm_universal_gated.hpp"

using namespace cute;

template <typename InputType, typename CTAShape, typename ClusterShape,

typename MainloopScheduleType, typename EpilogueScheduleType, typename TileSchedulerType = void,

template <class /* ElementCompute */> class Activation = cutlass::epilogue::thread::SiLu, bool SwapAB = true>

bool dispatch_dual_gemm_act_sm90(DualGemmEpilogueAllParams params) {

using ElementA = typename std::conditional_t<

std::is_same_v<InputType, phi::dtype::float8_e4m3fn>,

cutlass::float_e4m3_t,

cutlass::float_e5m2_t>;

using LayoutA = cutlass::layout::RowMajor; // Layout type for A matrix operand

static constexpr int AlignmentA

= 128 / cutlass::sizeof_bits<ElementA>::value; // Memory access granularity/alignment of A

// matrix in units of elements (up to 16 bytes)

// B matrix configuration

using ElementB = ElementA; // Element type for B matrix operand

using LayoutB = cutlass::layout::ColumnMajor; // Layout type for B matrix operand

static constexpr int AlignmentB

= 128 / cutlass::sizeof_bits<ElementB>::value; // Memory access granularity/alignment of B

// matrix in units of elements (up to 16 bytes)

using ElementC = ElementA; // Element type for C matrix operands

using LayoutC = cute::conditional_t<SwapAB, cutlass::layout::ColumnMajor, cutlass::layout::RowMajor>;

static constexpr int AlignmentC

= 128 / cutlass::sizeof_bits<ElementC>::value; // Memory access granularity/alignment of C matrices in units of

// elements (up to 16 bytes)

// Output matrix configuration

using ElementOutput = ElementA; // Element type for output matrix operands

// using LayoutOutput = cutlass::layout::RowMajor; // Layout type for output matrix operands

using LayoutOutput = cute::conditional_t<SwapAB, cutlass::layout::ColumnMajor, cutlass::layout::RowMajor>;

static constexpr int AlignmentOutput = 128 / cutlass::sizeof_bits<ElementOutput>::value;

// Multiply-accumulate blocking/pipelining details

using ElementAccumulator = float; // Element type for internal accumulation

using ElementCompute = float; // Element type for compute

using ArchTag = cutlass::arch::Sm90; // Tag indicating the minimum SM that supports the intended feature

using OperatorClass = cutlass::arch::OpClassTensorOp; // Operator class tag

using TileShape = CTAShape; // Threadblock-level tile size

using KernelSchedule = MainloopScheduleType;

using EpilogueSchedule = EpilogueScheduleType;

using TileScheduler = TileSchedulerType;

using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;

static constexpr auto RoundStyle = cutlass::FloatRoundStyle::round_to_nearest;

using FusionOperation = cutlass::epilogue::fusion::ScaledAcc<ElementOutput, ElementCompute>;

using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<ArchTag, OperatorClass,

TileShape, ClusterShape, EpilogueTileType, ElementAccumulator, ElementAccumulator, ElementC, LayoutC,

AlignmentC, ElementOutput, LayoutOutput, AlignmentOutput, EpilogueSchedule, FusionOperation>::CollectiveOp;

using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilderGated<ArchTag, OperatorClass,

ElementA, LayoutA, AlignmentA, ElementB, LayoutB, AlignmentB, ElementAccumulator, TileShape, ClusterShape,

cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(

sizeof(typename CollectiveEpilogue::SharedStorage))>,

KernelSchedule, Activation, SwapAB>::CollectiveOp;

using GemmKernel = cutlass::gemm::kernel::GemmUniversalGated<Shape<int, int, int, int>, // Indicates ProblemShape

CollectiveMainloop, CollectiveEpilogue, TileScheduler>;

using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;

using StrideA = typename Gemm::GemmKernel::StrideA;

using StrideB = typename Gemm::GemmKernel::StrideB;

using StrideC = typename Gemm::GemmKernel::StrideC;

using StrideD = typename Gemm::GemmKernel::StrideD;

int arg_m = params.M;

int arg_n = params.N;

ElementA const* ptr_A = reinterpret_cast<ElementA const*>(params.A);

ElementB const* ptr_B0 = reinterpret_cast<ElementB const*>(params.B0);

ElementB const* ptr_B1 = reinterpret_cast<ElementB const*>(params.B1);

if constexpr (SwapAB)

{

arg_m = params.N;

arg_n = params.M;

ptr_A = reinterpret_cast<ElementB const*>(params.B0);

ptr_B0 = reinterpret_cast<ElementA const*>(params.A);

}

StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(arg_m, params.K, 1));

StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(arg_n, params.K, 1));

StrideC stride_C;

StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(arg_m, arg_n, 1));

typename Gemm::Arguments arguments = {cutlass::gemm::GemmUniversalMode::kGemm, {arg_m, arg_n, params.K, 1},

{ptr_A, stride_A, ptr_B0, ptr_B1, stride_B, params.scale0, params.scale1},

{{}, // epilogue.thread

nullptr, stride_C, reinterpret_cast<ElementOutput*>(params.D), stride_D}};

arguments.epilogue.thread.alpha = params.scale_out;

Gemm gemm_op;

cutlass::Status status = gemm_op.can_implement(arguments);

if (status != cutlass::Status::kSuccess) {

std::cerr << "Gemm::can_implement() failed" << std::endl;

return false;

}

size_t workspace_size = Gemm::get_workspace_size(arguments);

phi::Allocator* allocator = paddle::GetAllocator(params.place);

auto workspace = allocator->Allocate(workspace_size);

//

// Run the GEMM

//

status = gemm_op(arguments, workspace->ptr(), params.stream);

if (status != cutlass::Status::kSuccess) {

std::cerr << "Gemm::run() failed" << std::endl;

return false;

}

return true;

}

#pragma once

#include "fp8_common.h"

#include "cutlass/cutlass.h"

#include "cutlass/epilogue/collective/default_epilogue.hpp"

#include "cutlass/epilogue/thread/linear_combination.h"

#include "cutlass/gemm/dispatch_policy.hpp"

#include "cutlass/gemm/collective/collective_builder.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/gemm/kernel/tile_scheduler.hpp"

#include "cutlass/util/packed_stride.hpp"

using namespace cute;

template <

typename InputType = phi::dtype::float8_e4m3fn,

typename OutType = phi::dtype::float16,

bool hasbias = false,

template <class> typename Activation = cutlass::epilogue::thread::Identity,

typename TileShape = Shape<_128, _128, _128>,

typename ClusterShape = Shape<_2, _1, _1>,

typename KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedPingpongFP8FastAccum,

typename EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized,

typename SM = cutlass::arch::Sm90

>

bool dispatch_fuse_gemm_act_sm90(GemmEpilogueAllParams params){

using ElementA = typename std::conditional_t<std::is_same_v<InputType, phi::dtype::float8_e4m3fn>,

cutlass::float_e4m3_t,

cutlass::float_e5m2_t>;

using ElementB = ElementA;

using ElementD = typename std::conditional_t<std::is_same_v<OutType, phi::dtype::bfloat16>,

cutlass::bfloat16_t,

cutlass::half_t>;

using ElementC = std::conditional_t<

hasbias,

ElementD,

void>;

using LayoutA = cutlass::layout::RowMajor;

using LayoutB = cutlass::layout::ColumnMajor;

using LayoutC = cutlass::layout::RowMajor;

using LayoutD = cutlass::layout::RowMajor;

using ElementAccumulator = float;

using ElementCompute = float;

using ElementScalar = float;

// 16B alignment lets us use TMA

static constexpr int AlignmentA = 16 / sizeof(ElementA);

static constexpr int AlignmentB = 16 / sizeof(ElementB);

static constexpr int AlignmentC = hasbias ? 16 / sizeof(ElementC) : 8;

static constexpr int AlignmentD = 16 / sizeof(ElementD);

static constexpr auto RoundStyle = cutlass::FloatRoundStyle::round_to_nearest;

using FusionOperation = cutlass::epilogue::fusion::LinCombEltAct<Activation, ElementD, ElementCompute, ElementC, ElementScalar, RoundStyle>;

using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<

SM, cutlass::arch::OpClassTensorOp,

TileShape,

ClusterShape,

cutlass::epilogue::collective::EpilogueTileAuto,

ElementAccumulator, ElementCompute,

ElementC, LayoutC, AlignmentC,

ElementD, LayoutD, AlignmentD,

EpilogueSchedule,

FusionOperation

>::CollectiveOp;

using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<

SM, cutlass::arch::OpClassTensorOp,

ElementA, LayoutA, AlignmentA,

ElementB, LayoutB, AlignmentB,

ElementAccumulator,

TileShape,

ClusterShape,

cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,

KernelSchedule

>::CollectiveOp;

using GemmKernel = cutlass::gemm::kernel::GemmUniversal<

Shape<int,int,int,int>,

CollectiveMainloop,

CollectiveEpilogue,

cutlass::gemm::PersistentScheduler

>;

using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;

using ProblemShapeType = typename Gemm::GemmKernel::ProblemShape;

using StrideA = typename Gemm::GemmKernel::StrideA;

using StrideB = typename Gemm::GemmKernel::StrideB;

using StrideC = typename Gemm::GemmKernel::StrideC;

using StrideD = typename Gemm::GemmKernel::StrideD;

//

// Data members

//

/// Initialization

StrideA stride_A{params.lda, cute::Int<1>{}, 0};

StrideB stride_B{params.ldb, cute::Int<1>{}, 0};

StrideC stride_C{0, cute::Int<1>{}, 0};

StrideD stride_D{params.ldd, cute::Int<1>{}, 0};

auto a_ptr = reinterpret_cast<ElementA*>(const_cast<void*>(params.A));

auto b_ptr = reinterpret_cast<ElementB*>(const_cast<void*>(params.B));

auto c_ptr = reinterpret_cast<ElementC*>(const_cast<void*>(params.bias));

auto d_ptr = reinterpret_cast<ElementD*>(params.D);

ProblemShapeType problem_size = ProblemShapeType{params.M, params.N, params.K, 1};

typename Gemm::Arguments arguments{

cutlass::gemm::GemmUniversalMode::kGemm,

problem_size,

{a_ptr, stride_A, b_ptr, stride_B},

{{params.scale}, // epilogue.thread

c_ptr, stride_C, d_ptr, stride_D}

};

if constexpr (hasbias){

arguments.epilogue.thread.beta = 1.0;

}

Gemm gemm_op;

cutlass::Status status = gemm_op.can_implement(arguments);

if (status != cutlass::Status::kSuccess) {

std::cout << "Gemm::can_implement() failed. " << cutlassGetStatusString(status) << std::endl;

return false;

}

size_t workspace_size = Gemm::get_workspace_size(arguments);

phi::Allocator* allocator = paddle::GetAllocator(params.place);

auto workspace = allocator->Allocate(workspace_size);

status = gemm_op(arguments, workspace->ptr(), params.stream);

if (status != cutlass::Status::kSuccess) {

std::cout << "Gemm::run() failed." << cutlassGetStatusString(status) << std::endl;

return false;

}

return true;

}