CUTLASS 3.3.0 (#1167)
* Release 3.3.0 Adds support for mixed precision GEMMs On Hopper and Ampere Adds support for < 16B aligned GEMMs on Hopper Enhancements to EVT Enhancements to Python interface Enhancements to Sub-byte type handling in CuTe Several other bug-fixes and performance improvements. * minor doc update
This commit is contained in:
Pradeep Ramani
@@ -33,7 +33,13 @@ cutlass_example_add_executable(
|
||||
ampere_gemm_universal_streamk.cu
|
||||
)
|
||||
|
||||
# Deliberately test non-square sizes to ensure that internal transpose is
|
||||
# not triggered when using SM80 EVT
|
||||
set(TEST_COMMAND_00 --m=512 --n=768 --k=1152)
|
||||
|
||||
cutlass_example_add_executable(
|
||||
47_ampere_gemm_universal_streamk_broadcast
|
||||
ampere_gemm_universal_streamk_broadcast.cu
|
||||
TEST_COMMAND_OPTIONS
|
||||
TEST_COMMAND_00
|
||||
)
|
||||
|
||||
@@ -107,16 +107,6 @@ using ClusterShape = Shape<_1,_2,_1>; // S
|
||||
using StageCountType = cutlass::gemm::collective::StageCountAuto; // Stage count maximized based on the tile size
|
||||
using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto; // Kernel to launch based on the default setting in the Collective Builder
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
ArchTag, OperatorClass,
|
||||
ElementA, LayoutA, AlignmentA,
|
||||
ElementB, LayoutB, AlignmentB,
|
||||
ElementAccumulator,
|
||||
TileShape, ClusterShape,
|
||||
cutlass::gemm::collective::StageCountAuto,
|
||||
cutlass::gemm::collective::KernelScheduleAuto
|
||||
>::CollectiveOp;
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
|
||||
TileShape, ClusterShape,
|
||||
@@ -127,6 +117,17 @@ using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBui
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto
|
||||
>::CollectiveOp;
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
ArchTag, OperatorClass,
|
||||
ElementA, LayoutA, AlignmentA,
|
||||
ElementB, LayoutB, AlignmentB,
|
||||
ElementAccumulator,
|
||||
TileShape, ClusterShape,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<
|
||||
static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int>, // Indicates ProblemShape
|
||||
CollectiveMainloop,
|
||||
@@ -397,6 +398,7 @@ int run(Options &options)
|
||||
GpuTimer timer;
|
||||
timer.start();
|
||||
for (int iter = 0; iter < options.iterations; ++iter) {
|
||||
CUTLASS_CHECK(gemm.initialize(arguments, workspace.get()));
|
||||
CUTLASS_CHECK(gemm.run());
|
||||
}
|
||||
timer.stop();
|
||||
|
||||
@@ -86,7 +86,7 @@
|
||||
CUTLASS builders make an attempt to pick the best schedule when `Auto` is provided such that the
|
||||
assembled collectives have the best performance, but this is not a guarantee. A user relying on `Auto`
|
||||
may get a free performance upgrade with newer CUTLASS releases in case we can provide more optimized
|
||||
implementations that the builder can transparently assemble for `Auto`. But a user should not rely on
|
||||
implementations that the builder can transparently assemble for `Auto`. But a user should not rely on
|
||||
`Auto` if they require a specific scheduling policy and/or stage count to be used.
|
||||
|
||||
If a user decides to let the builders pick the collective specialization via `Auto` schedules,
|
||||
@@ -289,7 +289,7 @@ struct ExampleRunner {
|
||||
// EVTs can be constructed by composing the fundamental load/store/compute visitor operations defined in include/cutlass/epilogue/fusion
|
||||
// For more complex examples of EVT construction please refer to include/cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp
|
||||
using CustomEVT = // alpha * acc + beta * C
|
||||
cutlass::epilogue::fusion::Sm90EVT<cutlass::epilogue::fusion::Sm90Compute<cutlass::multiply_add, ElementD, ElementCompute, RoundStyle>, // beta * C + (alpha * acc)
|
||||
cutlass::epilogue::fusion::Sm90EVT<cutlass::epilogue::fusion::Sm90Compute<cutlass::homogeneous_multiply_add, ElementD, ElementCompute, RoundStyle>, // beta * C + (alpha * acc)
|
||||
cutlass::epilogue::fusion::Sm90ScalarBroadcast<ElementScalar>, // beta
|
||||
cutlass::epilogue::fusion::Sm90SrcFetch, // C
|
||||
cutlass::epilogue::fusion::Sm90EVT<cutlass::epilogue::fusion::Sm90Compute<cutlass::multiplies, ElementCompute, ElementCompute, RoundStyle>, // alpha * acc
|
||||
@@ -322,7 +322,7 @@ struct ExampleRunner {
|
||||
ElementAccumulator,
|
||||
Shape<_128,_128,_64>, Shape<_2,_1,_1>,
|
||||
cute::conditional_t<cute::is_same_v<StageCountType, cutlass::gemm::collective::StageCountAuto>,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<(int)sizeof(typename CollectiveEpilogue::SharedStorage)>,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
StageCountType>,
|
||||
MainloopScheduleType
|
||||
>::CollectiveOp;
|
||||
|
||||
@@ -93,7 +93,7 @@ struct Options {
|
||||
int mode = 1; // N-mode gather/scatter by default
|
||||
|
||||
float alpha = 1.0f;
|
||||
float beta = 1.0f;
|
||||
float beta = 0.0f;
|
||||
|
||||
bool reference_check = true;
|
||||
int iterations = 20;
|
||||
@@ -179,19 +179,27 @@ struct ExampleRunner
|
||||
{
|
||||
// Useful aliases
|
||||
|
||||
using ProblemShape = Shape<int,int,int,int>;
|
||||
|
||||
using StrideA = cutlass::gemm::TagToStrideA_t<LayoutA>;
|
||||
using StrideB = cutlass::gemm::TagToStrideB_t<LayoutB>;
|
||||
using StrideC = cutlass::gemm::TagToStrideC_t<LayoutC>;
|
||||
using StrideD = cutlass::gemm::TagToStrideC_t<LayoutD>;
|
||||
|
||||
// Alias to for the epilogue type that supports gather/scatter
|
||||
using Epilogue = cutlass::epilogue::collective::EpilogueGatherScatter<
|
||||
cutlass::gemm::TagToStrideC_t<LayoutC>,
|
||||
cutlass::gemm::TagToStrideC_t<LayoutD>,
|
||||
cutlass::epilogue::thread::LinearCombination<
|
||||
ElementD, 1,
|
||||
ElementAccumulator, ElementComputeEpilogue,
|
||||
cutlass::epilogue::thread::ScaleType::Default,
|
||||
cutlass::FloatRoundStyle::round_to_nearest, ElementC
|
||||
>,
|
||||
cutlass::gemm::EpilogueDefault,
|
||||
GatherC,
|
||||
ScatterD
|
||||
using Epilogue = cutlass::epilogue::collective::detail::Sm90TmaWarpSpecializedAdapter<
|
||||
cutlass::epilogue::collective::EpilogueGatherScatter<
|
||||
StrideC, StrideD,
|
||||
cutlass::epilogue::thread::LinearCombination<
|
||||
ElementD, 1,
|
||||
ElementAccumulator, ElementComputeEpilogue,
|
||||
cutlass::epilogue::thread::ScaleType::Default,
|
||||
cutlass::FloatRoundStyle::round_to_nearest, ElementC
|
||||
>,
|
||||
cutlass::gemm::EpilogueDefault,
|
||||
GatherC,
|
||||
ScatterD
|
||||
>
|
||||
>;
|
||||
|
||||
// Alias to for the mainloop type
|
||||
@@ -202,27 +210,21 @@ struct ExampleRunner
|
||||
ElementAccumulator,
|
||||
Shape<_128,_128,_64>,
|
||||
Shape<_1,_1,_1>,
|
||||
cutlass::gemm::collective::StageCount<5>,
|
||||
cutlass::gemm::KernelMultistage
|
||||
cutlass::gemm::collective::StageCountAuto,
|
||||
cutlass::gemm::KernelCpAsyncWarpSpecialized
|
||||
>::CollectiveOp;
|
||||
|
||||
using ProblemShape = Shape<int,int,int,int>;
|
||||
|
||||
using Kernel = cutlass::gemm::kernel::GemmGather<
|
||||
ProblemShape,
|
||||
Mainloop,
|
||||
Epilogue,
|
||||
void,
|
||||
GatherA,
|
||||
GatherB
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
|
||||
|
||||
using StrideA = typename Kernel::StrideA;
|
||||
using StrideB = typename Kernel::StrideB;
|
||||
using StrideC = typename Kernel::StrideC;
|
||||
using StrideD = typename Kernel::StrideD;
|
||||
|
||||
static constexpr bool DoGatherA = not cutlass::platform::is_same<GatherA, NoGather>::value;
|
||||
static constexpr bool DoGatherB = not cutlass::platform::is_same<GatherB, NoGather>::value;
|
||||
static constexpr bool DoGatherC = not cutlass::platform::is_same<GatherC, NoGather>::value;
|
||||
@@ -250,16 +252,19 @@ struct ExampleRunner
|
||||
|
||||
using MainloopRef = Mainloop;
|
||||
|
||||
using EpilogueRef = typename cutlass::epilogue::collective::DefaultEpilogue<
|
||||
StrideC, StrideD,
|
||||
typename Epilogue::ThreadEpilogueOp,
|
||||
typename Epilogue::EpilogueSchedule
|
||||
using EpilogueRef = cutlass::epilogue::collective::detail::Sm90TmaWarpSpecializedAdapter<
|
||||
cutlass::epilogue::collective::DefaultEpilogue<
|
||||
StrideC, StrideD,
|
||||
typename Epilogue::ThreadEpilogueOp,
|
||||
typename Epilogue::EpilogueSchedule
|
||||
>
|
||||
>;
|
||||
|
||||
using KernelRef = cutlass::gemm::kernel::GemmUniversal<
|
||||
ProblemShape,
|
||||
MainloopRef,
|
||||
EpilogueRef
|
||||
ProblemShape,
|
||||
MainloopRef,
|
||||
EpilogueRef,
|
||||
void
|
||||
>;
|
||||
|
||||
using GemmRef = cutlass::gemm::device::GemmUniversalAdapter<KernelRef>;
|
||||
@@ -289,9 +294,10 @@ struct ExampleRunner
|
||||
>::CollectiveOp;
|
||||
|
||||
using KernelOpt = cutlass::gemm::kernel::GemmUniversal<
|
||||
ProblemShape,
|
||||
MainloopOpt,
|
||||
EpilogueOpt
|
||||
ProblemShape,
|
||||
MainloopOpt,
|
||||
EpilogueOpt,
|
||||
void
|
||||
>;
|
||||
|
||||
using GemmOpt = cutlass::gemm::device::GemmUniversalAdapter<KernelOpt>;
|
||||
@@ -404,6 +410,7 @@ struct ExampleRunner
|
||||
typename Epilogue::ScatterD{gather_indices.get()}
|
||||
},
|
||||
hw_info,
|
||||
{},
|
||||
typename Kernel::GatherA{gather_indices.get()},
|
||||
typename Kernel::GatherB{gather_indices.get()}
|
||||
},
|
||||
|
||||
@@ -47,9 +47,9 @@ template <
|
||||
class ProblemShape_,
|
||||
class CollectiveMainloop_,
|
||||
class CollectiveEpilogue_,
|
||||
class TileScheduler_,
|
||||
class GatherA_,
|
||||
class GatherB_,
|
||||
class TileScheduler_ = void
|
||||
class GatherB_
|
||||
>
|
||||
class GemmGather
|
||||
{
|
||||
@@ -58,8 +58,6 @@ public:
|
||||
// Type Aliases
|
||||
//
|
||||
using ProblemShape = ProblemShape_;
|
||||
using TileSchedulerTag = TileScheduler_;
|
||||
using TileScheduler = TileScheduler_;
|
||||
static_assert(rank(ProblemShape{}) == 3 or rank(ProblemShape{}) == 4,
|
||||
"ProblemShape{} should be <M,N,K> or <M,N,K,L>");
|
||||
|
||||
@@ -74,8 +72,10 @@ public:
|
||||
using StrideB = typename CollectiveMainloop::StrideB;
|
||||
using DispatchPolicy = typename CollectiveMainloop::DispatchPolicy;
|
||||
using ElementAccumulator = typename CollectiveMainloop::ElementAccumulator;
|
||||
using ClusterShape = typename DispatchPolicy::ClusterShape;
|
||||
using MainloopArguments = typename CollectiveMainloop::Arguments;
|
||||
using MainloopParams = typename CollectiveMainloop::Params;
|
||||
static_assert(ArchTag::kMinComputeCapability >= 90);
|
||||
|
||||
// Epilogue derived types
|
||||
using CollectiveEpilogue = CollectiveEpilogue_;
|
||||
@@ -85,17 +85,48 @@ public:
|
||||
using StrideD = typename CollectiveEpilogue::StrideD;
|
||||
using EpilogueArguments = typename CollectiveEpilogue::Arguments;
|
||||
using EpilogueParams = typename CollectiveEpilogue::Params;
|
||||
static_assert(std::is_same_v<ElementAccumulator, typename CollectiveEpilogue::ElementAccumulator>,
|
||||
"Mainloop and epilogue do not agree on accumulator value type.");
|
||||
|
||||
static_assert(cute::is_void_v<TileScheduler_> or cute::is_same_v<TileScheduler_, PersistentScheduler>,
|
||||
"Non-persistent warp-specialized kernel does not support specializing the tile scheduler.");
|
||||
using TileSchedulerTag = TileScheduler_;
|
||||
using TileScheduler = typename detail::TileSchedulerSelector<
|
||||
TileScheduler_, ArchTag, TileShape, ClusterShape>::Scheduler;
|
||||
using TileSchedulerArguments = typename TileScheduler::Arguments;
|
||||
|
||||
using GatherA = GatherA_;
|
||||
using GatherB = GatherB_;
|
||||
|
||||
static constexpr int SharedStorageSize = static_cast<int>(cute::max(
|
||||
sizeof(typename CollectiveMainloop::SharedStorage),
|
||||
sizeof(typename CollectiveEpilogue::SharedStorage)));
|
||||
// Kernel level shared memory storage
|
||||
struct SharedStorage {
|
||||
union TensorStorage {
|
||||
using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage;
|
||||
using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage;
|
||||
|
||||
static constexpr uint32_t MaxThreadsPerBlock = cute::size(TiledMma{});
|
||||
MainloopTensorStorage mainloop;
|
||||
EpilogueTensorStorage epilogue;
|
||||
} tensors;
|
||||
|
||||
struct PipelineStorage : cute::aligned_struct<16> {
|
||||
using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage;
|
||||
using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage;
|
||||
|
||||
alignas(16) MainloopPipelineStorage mainloop;
|
||||
alignas(16) EpiLoadPipelineStorage epi_load;
|
||||
} pipelines;
|
||||
};
|
||||
|
||||
static constexpr int SharedStorageSize = sizeof(SharedStorage);
|
||||
|
||||
using GmemTiledCopyA = typename CollectiveMainloop::GmemTiledCopyA;
|
||||
using GmemTiledCopyB = typename CollectiveMainloop::GmemTiledCopyB;
|
||||
static_assert(cute::size(GmemTiledCopyA{}) == cute::size(GmemTiledCopyB{}), "Number of threads in A/B tiled copies must be the same.");
|
||||
|
||||
static constexpr uint32_t NumLoadWarpGroups = cute::size(GmemTiledCopyA{}) / NumThreadsPerWarpGroup;
|
||||
static constexpr uint32_t NumMmaWarpGroups = cute::size(TiledMma{}) / NumThreadsPerWarpGroup;
|
||||
static constexpr uint32_t NumWarpGroups = NumLoadWarpGroups + NumMmaWarpGroups;
|
||||
static_assert(NumWarpGroups == 2 || NumWarpGroups == 3, "Number of warp groups must be 2 or 3 for good performance.");
|
||||
|
||||
static constexpr uint32_t MaxThreadsPerBlock = NumWarpGroups * NumThreadsPerWarpGroup;
|
||||
static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
|
||||
|
||||
// Device side arguments
|
||||
@@ -105,6 +136,7 @@ public:
|
||||
MainloopArguments mainloop{};
|
||||
EpilogueArguments epilogue{};
|
||||
KernelHardwareInfo hw_info{};
|
||||
TileSchedulerArguments scheduler{};
|
||||
GatherA gather_A{};
|
||||
GatherB gather_B{};
|
||||
};
|
||||
@@ -123,14 +155,20 @@ public:
|
||||
// Methods
|
||||
//
|
||||
|
||||
// Convert to underlying arguments.
|
||||
// Convert to underlying arguments. In this case, a simple copy for the aliased type.
|
||||
static
|
||||
Params
|
||||
to_underlying_arguments(Arguments const& args, void* workspace) {
|
||||
(void) workspace;
|
||||
auto problem_shape = args.problem_shape;
|
||||
if constexpr (detail::IF_SWAP_AB<CollectiveMainloop>::value) {
|
||||
// swap M/N
|
||||
get<0>(problem_shape) = get<1>(args.problem_shape);
|
||||
get<1>(problem_shape) = get<0>(args.problem_shape);
|
||||
}
|
||||
return {
|
||||
args.mode,
|
||||
args.problem_shape,
|
||||
problem_shape,
|
||||
CollectiveMainloop::to_underlying_arguments(args.problem_shape, args.mainloop, workspace),
|
||||
CollectiveEpilogue::to_underlying_arguments(args.problem_shape, args.epilogue, workspace),
|
||||
args.gather_A,
|
||||
@@ -138,17 +176,18 @@ public:
|
||||
};
|
||||
}
|
||||
|
||||
static
|
||||
Status
|
||||
initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
static
|
||||
CUTLASS_HOST_DEVICE static
|
||||
bool
|
||||
can_implement(Arguments const& args) {
|
||||
return args.mode == GemmUniversalMode::kGemm or
|
||||
(args.mode == GemmUniversalMode::kBatched && rank(ProblemShape{}) == 4);
|
||||
bool implementable = (args.mode == GemmUniversalMode::kGemm) or
|
||||
(args.mode == GemmUniversalMode::kBatched && rank(ProblemShape{}) == 4);
|
||||
if (!implementable) {
|
||||
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Arguments or Problem Shape don't meet the requirements.\n");
|
||||
return implementable;
|
||||
}
|
||||
implementable &= CollectiveMainloop::can_implement(args.problem_shape, args.mainloop);
|
||||
implementable &= CollectiveEpilogue::can_implement(args.problem_shape, args.epilogue);
|
||||
return implementable;
|
||||
}
|
||||
|
||||
static
|
||||
@@ -157,23 +196,23 @@ public:
|
||||
return 0;
|
||||
}
|
||||
|
||||
static constexpr
|
||||
dim3
|
||||
get_grid_shape(Params const& params) {
|
||||
int batch_count = 1;
|
||||
if constexpr (rank(ProblemShape{}) == 4) {
|
||||
batch_count = cute::size<3>(params.problem_shape);
|
||||
}
|
||||
|
||||
return dim3(
|
||||
cute::size(cute::ceil_div(cute::shape<0>(params.problem_shape), cute::shape<0>(TileShape{}))),
|
||||
cute::size(cute::ceil_div(cute::shape<1>(params.problem_shape), cute::shape<1>(TileShape{}))),
|
||||
batch_count
|
||||
);
|
||||
static
|
||||
cutlass::Status
|
||||
initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
static constexpr
|
||||
dim3
|
||||
// Computes the kernel launch grid shape based on runtime parameters
|
||||
static dim3
|
||||
get_grid_shape(Params const& params) {
|
||||
auto cluster_shape = Shape<_1,_1,_1>{};
|
||||
auto tile_shape = TileShape{};
|
||||
auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{});
|
||||
return TileScheduler::get_tiled_cta_shape_mnl(
|
||||
problem_shape_MNKL, tile_shape, cluster_shape);
|
||||
}
|
||||
|
||||
static dim3
|
||||
get_block_shape() {
|
||||
return dim3(MaxThreadsPerBlock, 1, 1);
|
||||
}
|
||||
@@ -184,8 +223,75 @@ public:
|
||||
using namespace cute;
|
||||
using X = Underscore;
|
||||
|
||||
// Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a.
|
||||
#if ! defined(__CUDA_ARCH_FEAT_SM90_ALL)
|
||||
if constexpr(size<0>(typename TiledMma::AtomShape_MNK{}) == 64) {
|
||||
printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n");
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
|
||||
enum class WarpGroupRole {
|
||||
Producer = 0,
|
||||
Consumer = 1,
|
||||
};
|
||||
|
||||
// Kernel level shared memory storage
|
||||
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
|
||||
|
||||
int thread_idx = int(threadIdx.x);
|
||||
int warp_group_thread_idx = thread_idx % NumThreadsPerWarpGroup;
|
||||
int warp_group_idx = canonical_warp_group_idx();
|
||||
CUTLASS_ASSERT(warp_group_idx < NumWarpGroups);
|
||||
WarpGroupRole warp_group_role = warp_group_idx < NumLoadWarpGroups ? WarpGroupRole::Producer : WarpGroupRole::Consumer;
|
||||
|
||||
// Mainloop Load pipeline
|
||||
using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline;
|
||||
typename MainloopPipeline::Params mainloop_pipeline_params;
|
||||
if (warp_group_role == WarpGroupRole::Producer) {
|
||||
mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer;
|
||||
}
|
||||
if (warp_group_role == WarpGroupRole::Consumer) {
|
||||
mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer;
|
||||
}
|
||||
mainloop_pipeline_params.producer_arv_count = NumLoadWarpGroups * NumThreadsPerWarpGroup;
|
||||
mainloop_pipeline_params.consumer_arv_count = NumMmaWarpGroups * NumThreadsPerWarpGroup;
|
||||
MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params);
|
||||
|
||||
// Epilogue Load pipeline
|
||||
using EpiLoadPipeline = typename CollectiveEpilogue::LoadPipeline;
|
||||
typename EpiLoadPipeline::Params epi_load_pipeline_params;
|
||||
if (warp_group_role == WarpGroupRole::Producer) {
|
||||
epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Producer;
|
||||
}
|
||||
if (warp_group_role == WarpGroupRole::Consumer) {
|
||||
epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Consumer;
|
||||
}
|
||||
epi_load_pipeline_params.producer_arv_count = NumLoadWarpGroups * NumThreadsPerWarpGroup;
|
||||
epi_load_pipeline_params.consumer_arv_count = NumMmaWarpGroups * NumThreadsPerWarpGroup;
|
||||
EpiLoadPipeline epi_load_pipeline(shared_storage.pipelines.epi_load, epi_load_pipeline_params);
|
||||
|
||||
// Epilogue Store pipeline
|
||||
using EpiStorePipeline = typename CollectiveEpilogue::StorePipeline;
|
||||
typename EpiStorePipeline::Params epi_store_pipeline_params;
|
||||
epi_store_pipeline_params.always_wait = true;
|
||||
EpiStorePipeline epi_store_pipeline(epi_store_pipeline_params);
|
||||
|
||||
// Initialize starting pipeline states for the collectives
|
||||
typename CollectiveMainloop::PipelineState mainloop_pipe_consumer_state;
|
||||
typename CollectiveEpilogue::LoadPipelineState epi_load_pipe_consumer_state;
|
||||
|
||||
// For the DMA Load (producer) we start with an opposite phase
|
||||
// i.e., we skip all waits since we know that the buffer is indeed empty
|
||||
PipelineState mainloop_pipe_producer_state = cutlass::make_producer_start_state<MainloopPipeline>();
|
||||
PipelineState epi_load_pipe_producer_state = cutlass::make_producer_start_state<EpiLoadPipeline>();
|
||||
PipelineState epi_store_pipe_producer_state = cutlass::make_producer_start_state<EpiStorePipeline>();
|
||||
|
||||
// Preconditions
|
||||
CUTE_STATIC_ASSERT(is_static<TileShape>::value);
|
||||
static_assert(rank(StrideA{}) == 3, "StrideA must be rank-3: [M, K, L]. If batch mode is not needed, set L stride to Int<0>.");
|
||||
static_assert(rank(StrideB{}) == 3, "StrideB must be rank-3: [N, K, L]. If batch mode is not needed, set L stride to Int<0>.");
|
||||
static_assert(rank(StrideC{}) == 3, "StrideC must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>.");
|
||||
static_assert(rank(StrideD{}) == 3, "StrideD must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>.");
|
||||
|
||||
// Separate out problem shape for convenience
|
||||
// Optionally append 1s until problem shape is rank-4 in case its is only rank-3 (MNK)
|
||||
@@ -195,69 +301,113 @@ public:
|
||||
auto K = get<2>(problem_shape_MNKL);
|
||||
auto L = get<3>(problem_shape_MNKL);
|
||||
|
||||
// Preconditions
|
||||
static_assert(rank(StrideA{}) == 3, "StrideA must be rank-3: [M, K, L]. If batch mode is not needed, set L stride to Int<0>.");
|
||||
static_assert(rank(StrideB{}) == 3, "StrideB must be rank-3: [N, K, L]. If batch mode is not needed, set L stride to Int<0>.");
|
||||
static_assert(rank(StrideC{}) == 3, "StrideC must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>.");
|
||||
static_assert(rank(StrideD{}) == 3, "StrideD must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>.");
|
||||
|
||||
// Get the appropriate blocks for this thread block -- potential for thread block locality
|
||||
int thread_idx = int(threadIdx.x);
|
||||
auto blk_shape = TileShape{}; // (BLK_M,BLK_N,BLK_K)
|
||||
auto [m_coord, n_coord, l_coord] = blockIdx;
|
||||
auto blk_coord_mnkl = make_coord(m_coord, n_coord, _, l_coord); // (m,n,k,l)
|
||||
|
||||
// Represent the full tensors
|
||||
Tensor mA_mkl = make_gather_tensor(make_gmem_ptr(params.mainloop.ptr_A), make_shape(M,K,L), params.mainloop.dA, params.gather_A); //(m,k,l)
|
||||
Tensor mB_nkl = make_gather_tensor(make_gmem_ptr(params.mainloop.ptr_B), make_shape(N,K,L), params.mainloop.dB, params.gather_B); //(n,k,l)
|
||||
|
||||
// Get batch slice
|
||||
Tensor mA_mk = mA_mkl(_,_,l_coord); // (m,k)
|
||||
Tensor mB_nk = mB_nkl(_,_,l_coord); // (n,k)
|
||||
|
||||
// Slice to get the tiles this thread block is responsible for
|
||||
Tensor gA = local_tile(mA_mk, blk_shape, take<0,3>(blk_coord_mnkl), Step<_1, X,_1>{}); // (BLK_M,BLK_K,k)
|
||||
Tensor gB = local_tile(mB_nk, blk_shape, take<0,3>(blk_coord_mnkl), Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)
|
||||
|
||||
// Compute tile residues for predication
|
||||
auto m_max_coord = M - size<0>(gA) * get<0>(blk_coord_mnkl); // M - BLK_M * m_coord
|
||||
auto n_max_coord = N - size<0>(gB) * get<1>(blk_coord_mnkl); // N - BLK_N * n_coord
|
||||
auto k_residue = K - size<1>(gA) * size<2>(gA); // K - BLK_K * k_coord_max
|
||||
auto residue_mnk = make_tuple(m_max_coord, n_max_coord, k_residue);
|
||||
|
||||
// Allocate the tiled_mma and the accumulators for the (M,N) blk_shape
|
||||
// Get the appropriate blocks for this thread block -- potential for thread block locality
|
||||
auto blk_shape = TileShape{}; // (BLK_M,BLK_N,BLK_K)
|
||||
TiledMma tiled_mma;
|
||||
Tensor accumulators = partition_fragment_C(tiled_mma, take<0,2>(blk_shape)); // (MMA,MMA_M,MMA_N)
|
||||
clear(accumulators);
|
||||
|
||||
// Make tiled views, defer the slice
|
||||
Tensor gA_mkl = local_tile(mA_mkl, blk_shape, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k,l)
|
||||
Tensor gB_nkl = local_tile(mB_nkl, blk_shape, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l)
|
||||
|
||||
// Compute m_coord, n_coord, and l_coord with their post-tiled shapes
|
||||
auto m_coord = idx2crd(int(blockIdx.x), shape<2>(gA_mkl));
|
||||
auto n_coord = idx2crd(int(blockIdx.y), shape<2>(gB_nkl));
|
||||
auto l_coord = idx2crd(int(blockIdx.z), shape<4>(gB_nkl));
|
||||
auto blk_coord = make_coord(m_coord, n_coord, _, l_coord);
|
||||
|
||||
// Slice with m_coord and n_coord
|
||||
Tensor gA = gA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
|
||||
Tensor gB = gB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k)
|
||||
|
||||
// Get pipeline iterators and increments from tensor shapes
|
||||
auto k_tile_iter = cute::make_coord_iterator(shape<2>(gA));
|
||||
int k_tile_count = size<2>(gA);
|
||||
auto k_tile_count = size<2>(gA);
|
||||
auto c_tile_count = CollectiveEpilogue::get_load_pipe_increment(blk_shape);
|
||||
auto d_tile_count = CollectiveEpilogue::get_store_pipe_increment(blk_shape);
|
||||
|
||||
// Perform the collective scoped MMA
|
||||
CollectiveMainloop collective_mma;
|
||||
collective_mma(
|
||||
accumulators,
|
||||
gA,
|
||||
gB,
|
||||
accumulators,
|
||||
k_tile_iter, k_tile_count,
|
||||
residue_mnk,
|
||||
thread_idx,
|
||||
smem_buf
|
||||
);
|
||||
// Wait for all threads in the thread block
|
||||
__syncthreads();
|
||||
|
||||
// Epilogue and write to gD
|
||||
CollectiveEpilogue epilogue{params.epilogue};
|
||||
epilogue(
|
||||
problem_shape_MNKL,
|
||||
blk_shape,
|
||||
blk_coord_mnkl,
|
||||
accumulators,
|
||||
tiled_mma,
|
||||
residue_mnk,
|
||||
thread_idx,
|
||||
smem_buf
|
||||
);
|
||||
// In a warp specialized kernel, collectives expose data movement and compute operations separately
|
||||
CollectiveMainloop collective_mainloop;
|
||||
CollectiveEpilogue collective_epilogue{params.epilogue, shared_storage.tensors.epilogue};
|
||||
|
||||
if (warp_group_role == WarpGroupRole::Producer) {
|
||||
// Compute tile residues for predication
|
||||
auto m_max_coord = M - size<0>(gA) * get<0>(blk_coord); // M - BLK_M * m_coord
|
||||
auto n_max_coord = N - size<0>(gB) * get<1>(blk_coord); // N - BLK_N * n_coord
|
||||
auto k_residue = K - size<1>(gA) * size<2>(gA); // K - BLK_K * k_coord_max
|
||||
auto residue_mnk = make_tuple(m_max_coord, n_max_coord, k_residue);
|
||||
|
||||
collective_mainloop.load(
|
||||
mainloop_pipeline,
|
||||
mainloop_pipe_producer_state,
|
||||
gA,
|
||||
gB,
|
||||
k_tile_iter, k_tile_count,
|
||||
residue_mnk,
|
||||
thread_idx,
|
||||
shared_storage.tensors.mainloop
|
||||
);
|
||||
// Update starting mainloop pipeline state for the pipeline drain
|
||||
mainloop_pipe_producer_state.advance(k_tile_count);
|
||||
// Make sure mainloop consumer has been waited upon before issuing epilogue load
|
||||
collective_mainloop.load_tail(mainloop_pipeline, mainloop_pipe_producer_state);
|
||||
|
||||
if (collective_epilogue.is_producer_load_needed()) {
|
||||
epi_load_pipe_producer_state =
|
||||
collective_epilogue.load(
|
||||
epi_load_pipeline,
|
||||
epi_load_pipe_producer_state,
|
||||
problem_shape_MNKL,
|
||||
blk_shape,
|
||||
blk_coord,
|
||||
tiled_mma,
|
||||
thread_idx,
|
||||
shared_storage.tensors.epilogue
|
||||
);
|
||||
collective_epilogue.load_tail(epi_load_pipeline, epi_load_pipe_producer_state);
|
||||
}
|
||||
}
|
||||
else if (warp_group_role == WarpGroupRole::Consumer) {
|
||||
Tensor accumulators = partition_fragment_C(tiled_mma, take<0,2>(blk_shape)); // (MMA,MMA_M,MMA_N)
|
||||
|
||||
collective_mainloop.mma(
|
||||
mainloop_pipeline,
|
||||
mainloop_pipe_consumer_state,
|
||||
accumulators,
|
||||
k_tile_count,
|
||||
warp_group_thread_idx,
|
||||
shared_storage.tensors.mainloop,
|
||||
params.mainloop
|
||||
);
|
||||
|
||||
// Make sure the math instructions are done and free buffers before entering the epilogue
|
||||
collective_mainloop.mma_tail(
|
||||
mainloop_pipeline,
|
||||
mainloop_pipe_consumer_state,
|
||||
k_tile_count
|
||||
);
|
||||
|
||||
// Epilogue and write to gD
|
||||
collective_epilogue.store(
|
||||
epi_load_pipeline,
|
||||
epi_load_pipe_consumer_state,
|
||||
epi_store_pipeline,
|
||||
epi_store_pipe_producer_state,
|
||||
problem_shape_MNKL,
|
||||
blk_shape,
|
||||
blk_coord,
|
||||
accumulators,
|
||||
tiled_mma,
|
||||
warp_group_thread_idx,
|
||||
shared_storage.tensors.epilogue
|
||||
);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -0,0 +1,558 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2023 - 2023 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/*! \file
|
||||
\brief Hopper GEMM example with different data types using CUTLASS 3.0 APIs for NVIDIA Hopper architecture
|
||||
|
||||
This example shows how to perform GEMM where the input tensors A and B have different element types. CUTLASS currently supports upcasting
|
||||
from a narrower (fewer bits) to a wider (more bits) type and utilizing the tensor core instruction for the wider type. For instance, when doing
|
||||
INT8 x FP16, CUTLASS will convert INT8 -> FP16 and do math using FP16 tensor cores. Similarly, for INT4 x INT8, it will upcast to INT8 and issue math
|
||||
using INT8 tensor cores.
|
||||
|
||||
The narrower type always passes through the register file. Therefore, in cases where the narrower type is operand B, the collective will implicitly swap
|
||||
A and B in the main loop. Consequently, it is essential to consider this when constructing the epilogue, as illustrated in this example.
|
||||
|
||||
Limitations:
|
||||
1) Only supported combinations are 16-bit x {8-bit, 4-bit, 2-bit} and {8-bit} x {4-bit, 2-bit}.
|
||||
2) The narrow type must always be in K-major format.
|
||||
3) When dealing with 8-bit x {4-bit, 2-bit}, both inputs must be in K-major format.
|
||||
4) 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 the narrow type passes through the register file, and TMA epilogues do not currently support swap + transpose operations.
|
||||
We plan to address this limitation in the future.
|
||||
|
||||
Examples:
|
||||
|
||||
$ ./examples/55_hopper_mixed_dtype_gemm/55_hopper_mixed_dtype_gemm --m=2048 --n=2048 --k=2048 --l=2
|
||||
*/
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cutlass/tensor_ref.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/util/command_line.h"
|
||||
#include "cutlass/util/distribution.h"
|
||||
#include "cutlass/util/host_tensor.h"
|
||||
#include "cutlass/util/packed_stride.hpp"
|
||||
#include "cutlass/util/tensor_view_io.h"
|
||||
#include "cutlass/util/reference/host/tensor_fill.h"
|
||||
#include "cutlass/util/reference/host/tensor_copy.h"
|
||||
#include "cutlass/util/reference/host/tensor_compare.h"
|
||||
#include "cutlass/util/reference/host/tensor_norm.h"
|
||||
#include "cutlass/util/reference/host/gett.hpp"
|
||||
|
||||
#include "helper.h"
|
||||
#include "unfused_weight_dequantize.h"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
/// GEMM kernel configurations
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
using MmaType = cutlass::half_t;
|
||||
using QuantType = int8_t;
|
||||
|
||||
// A matrix configuration
|
||||
using ElementA = MmaType; // Element type for A matrix operand
|
||||
using LayoutA = cutlass::layout::RowMajor; // Layout type for A matrix operand
|
||||
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 = QuantType; // Element type for B matrix operand
|
||||
using LayoutB = cutlass::layout::ColumnMajor; // Layout type for B matrix operand
|
||||
constexpr int AlignmentB = 128 / cutlass::sizeof_bits<ElementB>::value; // Memory access granularity/alignment of B matrix in units of elements (up to 16 bytes)
|
||||
|
||||
// C/D matrix configuration
|
||||
using ElementC = cutlass::half_t; // Element type for C and D matrix operands
|
||||
using LayoutC = cutlass::layout::RowMajor; // Layout type for C and D matrix operands
|
||||
constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value; // Memory access granularity/alignment of C matrix in units of elements (up to 16 bytes)
|
||||
|
||||
// D matrix configuration
|
||||
using ElementD = ElementC;
|
||||
using LayoutD = LayoutC;
|
||||
|
||||
// Core kernel configurations
|
||||
using ElementAccumulator = float; // Element type for internal accumulation
|
||||
using ElementCompute = float; // Element type for epilogue computation
|
||||
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 = Shape<_128,_256,_64>; // Threadblock-level tile size
|
||||
using ClusterShape = Shape<_2,_1,_1>; // Shape of the threadblocks in a cluster
|
||||
using StageCountType = cutlass::gemm::collective::StageCountAuto; // Stage count maximized based on the tile size
|
||||
using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto; // Kernel to launch based on the default setting in the Collective Builder
|
||||
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
ArchTag, OperatorClass,
|
||||
ElementA, LayoutA, AlignmentA,
|
||||
ElementB, LayoutB, AlignmentB,
|
||||
ElementAccumulator,
|
||||
TileShape, ClusterShape,
|
||||
cutlass::gemm::collective::StageCountAuto,
|
||||
cutlass::gemm::collective::KernelScheduleAuto
|
||||
>::CollectiveOp;
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
|
||||
TileShape, ClusterShape,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
ElementAccumulator, ElementAccumulator,
|
||||
// Lie here about layout of C and D since we do swap and transpose trick
|
||||
ElementC, typename cutlass::layout::LayoutTranspose<LayoutC>::type, AlignmentC,
|
||||
ElementC, typename cutlass::layout::LayoutTranspose<LayoutC>::type, AlignmentC,
|
||||
cutlass::epilogue::NoSmemWarpSpecialized // This is the only epi supporting the required swap + transpose.
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>, // Indicates ProblemShape
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
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;
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Initialization
|
||||
StrideA stride_A;
|
||||
StrideB stride_B;
|
||||
StrideC stride_C;
|
||||
StrideD stride_D;
|
||||
uint64_t seed;
|
||||
|
||||
// Initialization functions don't handle sub-byte types so we use uint8 to initialize and a separate
|
||||
// kernel to pack the data if it is necessary.
|
||||
using InitializationType = cute::conditional_t<cute::sizeof_bits_v<QuantType> < 8, uint8_t, QuantType>;
|
||||
|
||||
cutlass::HostTensor<typename Gemm::ElementA, LayoutA> tensor_A;
|
||||
cutlass::HostTensor<InitializationType, LayoutB> tensor_B_init;
|
||||
cutlass::HostTensor<typename Gemm::ElementB, LayoutB> tensor_B;
|
||||
cutlass::HostTensor<typename Gemm::ElementA, LayoutB> tensor_B_dq;
|
||||
cutlass::HostTensor<typename Gemm::ElementC, LayoutC> tensor_C;
|
||||
cutlass::HostTensor<typename Gemm::EpilogueOutputOp::ElementOutput, LayoutD> tensor_D;
|
||||
cutlass::HostTensor<typename Gemm::EpilogueOutputOp::ElementOutput, LayoutD> tensor_ref_D;
|
||||
|
||||
#endif // defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
/// Testbed utility types
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// Command line options parsing
|
||||
struct Options {
|
||||
|
||||
bool help = false;
|
||||
|
||||
float alpha = 1.0f;
|
||||
float beta = 0.0f;
|
||||
int iterations = 1000;
|
||||
int m = 5120, n = 4096, k = 4096;
|
||||
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("alpha", alpha, 1.f);
|
||||
cmd.get_cmd_line_argument("beta", beta, 0.f);
|
||||
cmd.get_cmd_line_argument("iterations", iterations);
|
||||
}
|
||||
|
||||
/// 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"
|
||||
<< "Options:\n\n"
|
||||
<< " --help If specified, displays this usage statement\n\n"
|
||||
<< " --m=<int> Sets the M extent of the GEMM\n"
|
||||
<< " --n=<int> Sets the N extent of the GEMM\n"
|
||||
<< " --k=<int> Sets the K extent of the GEMM\n"
|
||||
<< " --l=<int> The number of independent gemm problems with mnk shape\n"
|
||||
<< " --alpha=<f32> Epilogue scalar alpha\n"
|
||||
<< " --beta=<f32> Epilogue scalar beta\n\n"
|
||||
<< " --iterations=<int> Number of profiling iterations to perform.\n\n";
|
||||
|
||||
out
|
||||
<< "\n\nExamples:\n\n"
|
||||
<< "$ " << "55_hopper_warp_specialized_gemm" << " --m=1024 --n=512 --k=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)
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
/// GEMM setup and evaluation
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Helper to initialize a block of device data
|
||||
template <class Element, class Layout>
|
||||
bool initialize_tensor(
|
||||
cutlass::TensorView<Element, Layout> view,
|
||||
uint64_t seed=2023) {
|
||||
|
||||
double scope_max, scope_min;
|
||||
int bits_input = cutlass::sizeof_bits<Element>::value;
|
||||
int bits_output = cutlass::sizeof_bits<Element>::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::host::TensorFillRandomUniform(
|
||||
view, seed, scope_max, scope_min, 0);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
template <class QuantElement, typename Element, typename Layout>
|
||||
bool initialize_quant_tensor(
|
||||
cutlass::TensorView<Element, Layout> view,
|
||||
uint64_t seed=2023) {
|
||||
|
||||
Element scope_max, scope_min;
|
||||
constexpr int bits_input = cute::sizeof_bits_v<QuantElement>;
|
||||
static_assert(bits_input <= 8, "Quantization type can be at most 8 bits");
|
||||
|
||||
if constexpr (bits_input == 8) {
|
||||
// Directly init 1-byte types
|
||||
static_assert(cute::is_same_v<QuantElement, Element>, "Init type should equal quant type for 1 byte types");
|
||||
scope_max = std::numeric_limits<QuantElement>::max();
|
||||
scope_min = std::numeric_limits<QuantElement>::min();
|
||||
} else {
|
||||
static_assert(cute::is_same_v<uint8_t, Element>, "Init type should be uint8_t for sub-byte types");
|
||||
scope_max = (1 << bits_input);
|
||||
scope_min = 0;
|
||||
}
|
||||
|
||||
cutlass::reference::host::TensorFillRandomUniform(
|
||||
view, seed, scope_max, scope_min, 0);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
template <class Element, class Layout>
|
||||
bool initialize_with_one(
|
||||
cutlass::TensorView<Element, Layout> view) {
|
||||
cutlass::reference::host::TensorFill(view, Element(1.0f));
|
||||
return true;
|
||||
}
|
||||
|
||||
template <class ElementDst, class ElementSrc, class Layout, class L>
|
||||
void prepare_packed_data(cutlass::HostTensor<ElementDst, Layout> view_dst_data,
|
||||
cutlass::HostTensor<ElementSrc, Layout> view_src_data,
|
||||
const L& cute_layout) {
|
||||
if constexpr (cute::is_same_v<ElementSrc, ElementDst>) {
|
||||
view_dst_data.copy_in_device_to_device(view_src_data.device_data());
|
||||
}
|
||||
else {
|
||||
pack_data(view_dst_data.device_data(), view_src_data.device_data(), cute_layout);
|
||||
}
|
||||
}
|
||||
|
||||
/// Initialize operands to be used in the GEMM and reference GEMM
|
||||
void initialize(const Options &options) {
|
||||
|
||||
auto shape_b = cute::make_shape(options.n, options.k, options.l);
|
||||
stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(options.m, options.k, options.l));
|
||||
stride_B = cutlass::make_cute_packed_stride(StrideB{}, shape_b);
|
||||
stride_C = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(options.n, options.m, options.l));
|
||||
stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(options.n, options.m, options.l));
|
||||
|
||||
auto a_coord = cutlass::make_Coord(options.m * options.l, options.k);
|
||||
auto b_coord = cutlass::make_Coord(options.k, options.n * options.l);
|
||||
auto c_coord = cutlass::make_Coord(options.m * options.l, options.n);
|
||||
|
||||
tensor_A.resize(a_coord);
|
||||
tensor_B_init.resize(b_coord);
|
||||
tensor_B.resize(b_coord);
|
||||
tensor_B_dq.resize(b_coord);
|
||||
tensor_C.resize(c_coord);
|
||||
tensor_D.resize(c_coord);
|
||||
tensor_ref_D.resize(c_coord);
|
||||
|
||||
// We need scales since the "dequantize" kernels expects them. We just set them to 1 so the values get converted
|
||||
// to the mma type.
|
||||
cutlass::HostTensor<MmaType, cutlass::layout::RowMajor> tensor_scale;
|
||||
tensor_scale.resize({1 * options.l, options.n});
|
||||
|
||||
initialize_tensor(tensor_A.host_view(), seed + 2022);
|
||||
initialize_quant_tensor<QuantType>(tensor_B_init.host_view(), seed + 2021);
|
||||
initialize_tensor(tensor_C.host_view(), seed + 2020);
|
||||
initialize_with_one(tensor_scale.host_view());
|
||||
|
||||
tensor_A.sync_device();
|
||||
tensor_B_init.sync_device();
|
||||
tensor_C.sync_device();
|
||||
tensor_scale.sync_device();
|
||||
|
||||
auto layout_B = make_layout(shape_b, stride_B);
|
||||
prepare_packed_data(tensor_B, tensor_B_init, layout_B);
|
||||
|
||||
auto shape_scale = cute::make_shape(options.n, 1, options.l);
|
||||
auto layout_scale = make_layout(shape_scale);
|
||||
dequantize_weight(tensor_B_dq.device_data(), tensor_B.device_data(), layout_B, tensor_scale.device_data(), layout_scale);
|
||||
|
||||
tensor_B.sync_host();
|
||||
tensor_B_dq.sync_host();
|
||||
}
|
||||
|
||||
/// Populates a Gemm::Arguments structure from the given commandline options
|
||||
typename Gemm::Arguments args_from_options(const Options &options)
|
||||
{
|
||||
typename Gemm::Arguments arguments{
|
||||
cutlass::gemm::GemmUniversalMode::kGemm,
|
||||
{options.m, options.n, options.k, options.l},
|
||||
{tensor_A.device_data(), stride_A, tensor_B.device_data(), stride_B},
|
||||
{{options.alpha, options.beta}, tensor_C.device_data(), stride_C, tensor_D.device_data(), stride_D}
|
||||
};
|
||||
|
||||
return arguments;
|
||||
}
|
||||
|
||||
bool verify(const Options &options) {
|
||||
//
|
||||
// Compute reference output
|
||||
//
|
||||
|
||||
// Create instantiation for device reference gemm kernel
|
||||
auto A = cute::make_tensor(tensor_A.host_data(),
|
||||
cute::make_layout(cute::make_shape(options.m, options.k, options.l), stride_A));
|
||||
auto B = cute::make_tensor(tensor_B_dq.host_data(),
|
||||
cute::make_layout(cute::make_shape(options.n, options.k, options.l), stride_B));
|
||||
auto C = cute::make_tensor(tensor_C.host_data(),
|
||||
cute::make_layout(cute::make_shape(options.m, options.n, options.l), stride_C));
|
||||
auto D = cute::make_tensor(tensor_ref_D.host_data(),
|
||||
cute::make_layout(cute::make_shape(options.m, options.n, options.l), stride_D));
|
||||
|
||||
using unused_t = decltype(D);
|
||||
|
||||
cutlass::reference::host::GettMainloopParams<ElementAccumulator, decltype(A), decltype(B)> mainloop_params{A, B};
|
||||
|
||||
cutlass::reference::host::GettEpilogueParams<
|
||||
typename Gemm::EpilogueOutputOp::ElementScalar,
|
||||
typename Gemm::EpilogueOutputOp::ElementScalar,
|
||||
ElementAccumulator,
|
||||
ElementCompute,
|
||||
decltype(C),
|
||||
decltype(D),
|
||||
unused_t, // bias
|
||||
unused_t, // aux
|
||||
unused_t, // valpha
|
||||
unused_t // vbeta
|
||||
> epilogue_params;
|
||||
|
||||
epilogue_params.C = C;
|
||||
epilogue_params.D = D;
|
||||
epilogue_params.alpha = options.alpha;
|
||||
epilogue_params.beta = options.beta;
|
||||
|
||||
// get reference result
|
||||
cutlass::reference::host::Gemm3x(mainloop_params, epilogue_params);
|
||||
|
||||
// compare_reference
|
||||
tensor_D.sync_host();
|
||||
bool passed = cutlass::reference::host::TensorEquals(tensor_ref_D.host_view(), tensor_D.host_view());
|
||||
return passed;
|
||||
}
|
||||
|
||||
/// Execute a given example GEMM computation
|
||||
template <typename Gemm>
|
||||
int run(Options &options)
|
||||
{
|
||||
initialize(options);
|
||||
|
||||
// Instantiate CUTLASS kernel depending on templates
|
||||
Gemm gemm;
|
||||
|
||||
// Create a structure of gemm kernel arguments suitable for invoking an instance of Gemm
|
||||
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);
|
||||
|
||||
// Allocate workspace memory
|
||||
cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
|
||||
|
||||
// Check if the problem size is supported or not
|
||||
CUTLASS_CHECK(gemm.can_implement(arguments));
|
||||
|
||||
// Initialize CUTLASS kernel with arguments and workspace pointer
|
||||
CUTLASS_CHECK(gemm.initialize(arguments, workspace.get()));
|
||||
|
||||
// Correctness / Warmup iteration
|
||||
CUTLASS_CHECK(gemm.run());
|
||||
|
||||
// Check if output from CUTLASS kernel and reference kernel are equal or not
|
||||
Result result;
|
||||
result.passed = verify(options);
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
#endif // defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
int main(int argc, char const **args) {
|
||||
|
||||
// CUTLASS must be compiled with CUDA 12.0 Toolkit to run this example
|
||||
// and must have compute capability at least 90.
|
||||
if (__CUDACC_VER_MAJOR__ < 12) {
|
||||
std::cerr << "This example requires CUDA 12 or newer.\n";
|
||||
// Returning zero so this test passes on older Toolkits. Its actions are no-op.
|
||||
return 0;
|
||||
}
|
||||
|
||||
cudaDeviceProp props;
|
||||
int current_device_id;
|
||||
CUDA_CHECK(cudaGetDevice(¤t_device_id));
|
||||
CUDA_CHECK(cudaGetDeviceProperties(&props, current_device_id));
|
||||
cudaError_t error = cudaGetDeviceProperties(&props, 0);
|
||||
if (props.major < 9) {
|
||||
std::cerr
|
||||
<< "This example requires a GPU of NVIDIA's Hopper Architecture or "
|
||||
<< "later (compute capability 90 or greater).\n";
|
||||
return 0;
|
||||
}
|
||||
|
||||
//
|
||||
// Parse options
|
||||
//
|
||||
|
||||
Options options;
|
||||
|
||||
options.parse(argc, args);
|
||||
|
||||
if (options.help) {
|
||||
options.print_usage(std::cout) << std::endl;
|
||||
return 0;
|
||||
}
|
||||
|
||||
//
|
||||
// Evaluate CUTLASS kernels
|
||||
//
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)
|
||||
run<Gemm>(options);
|
||||
#endif
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
35
examples/55_hopper_mixed_dtype_gemm/CMakeLists.txt
Normal file
35
examples/55_hopper_mixed_dtype_gemm/CMakeLists.txt
Normal file
@@ -0,0 +1,35 @@
|
||||
|
||||
# Copyright (c) 2023 - 2023 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.
|
||||
|
||||
|
||||
|
||||
cutlass_example_add_executable(
|
||||
55_hopper_mixed_dtype_gemm
|
||||
55_hopper_mixed_dtype_gemm.cu
|
||||
)
|
||||
36
examples/55_hopper_mixed_dtype_gemm/README.md
Normal file
36
examples/55_hopper_mixed_dtype_gemm/README.md
Normal file
@@ -0,0 +1,36 @@
|
||||
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`.
|
||||
|
||||
This first version only supports mixed type GEMMs using TMA.
|
||||
|
||||
## Performance
|
||||
|
||||
While the example offers a harness for straightforward benchmarking, this initial implementation isn't optimized for performance in the majority of scenarios. We expect this implementation to be performant for `{fp16, bf16} x int8` for problems that are compute bound.
|
||||
|
||||
We are currently optimizing the following cases:
|
||||
1. Memory bound cases for all types
|
||||
1. Compute bound cases for `{16-bit, 8-bit} x {4-bit, 2-bit}`
|
||||
|
||||
As a result, we do not suggest using this example as a benchmarking reference until all of our optimizations are complete (this will be clearly stated in this README in a future release).
|
||||
|
||||
## Limitations
|
||||
|
||||
* The type that needs to be converted must go through the register file. This means that the collective will swap and transpose whenever the type with fewer bits is the B operand. The user must be aware of when these swaps happen to control the layout of the epilogue as shown in the example. Note that TMA epilogues currently do not support swap + transpose, so non-tma epilogues must be used in this case. We plan to relax this limitation in a future release.
|
||||
|
||||
* The layout of the narrow type must be K-major. This means the following:
|
||||
* Narrow type is the A operand: Must be Row-Major
|
||||
* Narrow type is the B operand: Must be Column-Major
|
||||
|
||||
* For 8-bit x 4-bit or 2-bit, both inputs must be K-major.
|
||||
|
||||
* TMA requires an alignment of 128 bits. As a result, for a type with `B` bits, `B x TILE_K` must be a multiple of 128 bits.
|
||||
|
||||
## Upcoming features
|
||||
|
||||
* Support for applying scales after conversion, but before issuing tensor core math (input scale fusion) is planned for v3.4.
|
||||
|
||||
* Many optimizations for SOL performance.
|
||||
178
examples/55_hopper_mixed_dtype_gemm/unfused_weight_dequantize.h
Normal file
178
examples/55_hopper_mixed_dtype_gemm/unfused_weight_dequantize.h
Normal file
@@ -0,0 +1,178 @@
|
||||
#pragma once
|
||||
|
||||
#include "cute/tensor.hpp"
|
||||
|
||||
#include <cuda.h>
|
||||
#include "helper.h"
|
||||
|
||||
template <class QuantizedElement,
|
||||
class DequantizedElement,
|
||||
class OperandLayout,
|
||||
class ElementScale,
|
||||
class ScaleBroadCastLayout,
|
||||
class ThrLayout>
|
||||
__global__ void dequantize_weight_kernel(DequantizedElement* dq_buffer,
|
||||
const QuantizedElement* q_buffer,
|
||||
const OperandLayout operand_layout,
|
||||
const ElementScale* scale_buffer,
|
||||
const ScaleBroadCastLayout 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<QuantizedElement> >= 8) {
|
||||
return make_gmem_ptr(q_buffer);
|
||||
} else {
|
||||
return subbyte_iterator<const QuantizedElement>(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);
|
||||
|
||||
// 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 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 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_op_dq = make_fragment_like(tOpDq_gOpDq(_, _, _, 0));
|
||||
|
||||
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);
|
||||
transform(rmem_op_q, rmem_op_dq, [] (const QuantizedElement& elt) { return DequantizedElement(elt); } );
|
||||
transform(rmem_op_dq, rmem_scale, rmem_op_dq, multiplies{});
|
||||
copy(rmem_op_dq, tOpDq_gOpDq(_, _, _, ii));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <class QuantizedElement,
|
||||
class DequantizedElement,
|
||||
class OperandLayout,
|
||||
class ElementScale,
|
||||
class ScaleLayout>
|
||||
void dequantize_weight(DequantizedElement* dq_buffer,
|
||||
const QuantizedElement* q_buffer,
|
||||
const OperandLayout operand_layout,
|
||||
const ElementScale* scale_buffer,
|
||||
const ScaleLayout scale_layout) {
|
||||
|
||||
using namespace cute;
|
||||
|
||||
constexpr int tpb = 128;
|
||||
auto thr_layout = make_layout(make_shape(Int<tpb>{}));
|
||||
|
||||
const auto num_rows = get<0>(shape(operand_layout));
|
||||
const auto num_cols = get<1>(shape(operand_layout)); // [MN, K, L]
|
||||
const auto batches = get<2>(shape(operand_layout)); // [MN, K, L]
|
||||
const auto num_cols_scale = get<1>(shape(scale_layout)); // [MN, G, 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);
|
||||
}
|
||||
|
||||
if (num_cols % num_cols_scale != 0) {
|
||||
std::cerr << "Invalid shape for weight / scales. Weight cols must be a multiple of scale cols."
|
||||
<< " 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(num_cols / num_cols_scale, num_cols_scale), 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 = num_cols;
|
||||
const auto blocks_y = batches;
|
||||
|
||||
dim3 blocks(blocks_x, blocks_y, 1);
|
||||
dequantize_weight_kernel<<<blocks, tpb>>>(dq_buffer, q_buffer, operand_layout, scale_buffer, scale_layout_bcast, thr_layout);
|
||||
CUDA_CHECK(cudaDeviceSynchronize());
|
||||
}
|
||||
|
||||
|
||||
template <int ELTS_PER_THREAD,
|
||||
class SubbyteType>
|
||||
__global__ void pack_data_kernel(SubbyteType* packed_data_ptr,
|
||||
const uint8_t* unpacked_data_ptr,
|
||||
const size_t max_elts) {
|
||||
using namespace cute;
|
||||
|
||||
const size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint8_t data[ELTS_PER_THREAD];
|
||||
if (tid < max_elts) {
|
||||
const uint8_t* read_ptr = unpacked_data_ptr + tid * ELTS_PER_THREAD;
|
||||
for (int ii = 0; ii < ELTS_PER_THREAD; ++ii) {
|
||||
data[ii] = read_ptr[ii];
|
||||
}
|
||||
|
||||
|
||||
using WriteType = cute::array_subbyte<SubbyteType, ELTS_PER_THREAD>;
|
||||
WriteType* write_ptr = reinterpret_cast<WriteType*>(packed_data_ptr);
|
||||
|
||||
WriteType packed_data;
|
||||
for (int ii = 0; ii < ELTS_PER_THREAD; ++ii) {
|
||||
SubbyteType elt(data[ii]);
|
||||
packed_data[ii] = elt;
|
||||
}
|
||||
write_ptr[tid] = packed_data;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
template <class SubbyteType,
|
||||
class OperandLayout>
|
||||
void pack_data(SubbyteType* packed_data, const uint8_t* unpacked_data, const OperandLayout operand_layout) {
|
||||
static_assert(cute::sizeof_bits_v<SubbyteType> < 8, "First operand must be a sub-byte type");
|
||||
constexpr int packed_elements = 8 / cute::sizeof_bits_v<SubbyteType>;
|
||||
|
||||
if (cute::stride<0>(operand_layout) == 1 && (cute::shape<0>(operand_layout) % packed_elements)) {
|
||||
std::cerr << "Invalid shape / stride for dimension 0. Contiguous dimension must be a multiple of "
|
||||
<< packed_elements << std::endl;
|
||||
exit(-1);
|
||||
}
|
||||
|
||||
if (cute::stride<1>(operand_layout) == 1 && (cute::shape<1>(operand_layout) % packed_elements)) {
|
||||
std::cerr << "Invalid shape / stride for dimension 1. Contiguous dimension must be a multiple of "
|
||||
<< packed_elements << std::endl;
|
||||
exit(-1);
|
||||
}
|
||||
|
||||
const int64_t total_threads = cute::size(operand_layout) / packed_elements;
|
||||
|
||||
const int threads_per_block = 256;
|
||||
const int64_t num_blocks = (total_threads + threads_per_block - 1) / threads_per_block;
|
||||
pack_data_kernel<packed_elements><<<int(num_blocks), threads_per_block>>>(packed_data, unpacked_data, total_threads);
|
||||
CUDA_CHECK(cudaDeviceSynchronize());
|
||||
}
|
||||
@@ -135,6 +135,7 @@ foreach(EXAMPLE
|
||||
52_hopper_gather_scatter_fusion
|
||||
53_hopper_gemm_permute
|
||||
54_hopper_fp8_warp_specialized_gemm
|
||||
55_hopper_mixed_dtype_gemm
|
||||
)
|
||||
|
||||
add_subdirectory(${EXAMPLE})
|
||||
|
||||
@@ -269,6 +269,7 @@
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"tiles = [td for td in tiles if td.threadblock_shape[0] >= 128]\n",
|
||||
"idx = random.randint(0, len(tiles)-1)\n",
|
||||
"td = tiles[idx]\n",
|
||||
"print('Tile description {} is: {}'.format(idx, td))\n",
|
||||
|
||||
@@ -29,7 +29,7 @@
|
||||
"import cutlass\n",
|
||||
"from cutlass.epilogue import relu\n",
|
||||
"from cutlass import Tensor as FakeTensor\n",
|
||||
"from cutlass.profiler import CUDAEventProfiler\n",
|
||||
"from cutlass.utils.profiler import CUDAEventProfiler\n",
|
||||
"\n",
|
||||
"# This controls whether ther C++ GEMM declaration will be printed at each step. Set to `false` to\n",
|
||||
"# omit this information.\n",
|
||||
@@ -160,10 +160,6 @@
|
||||
" return example_epilogue(accum, alpha, C, beta, aux, bias)\n",
|
||||
"\n",
|
||||
"torch_reference = TorchReference()\n",
|
||||
"if hasattr(torch, \"compile\"):\n",
|
||||
" # If the torch.compile feature is available\n",
|
||||
" torch_reference = torch.compile(torch_reference)\n",
|
||||
"\n",
|
||||
"tensor_D_ref, tensor_F_ref = torch_reference(tensor_A, tensor_B, alpha, tensor_C, beta, aux, bias)\n",
|
||||
"\n",
|
||||
"assert torch.equal(tensor_D, tensor_D_ref)\n",
|
||||
|
||||
Reference in New Issue
Block a user