CUTLASS 3.7 (#2045)
* CUTLASS 3.7 * clean up changelog --------- Co-authored-by: yuzhai <yuzhai@nvidia.com> Co-authored-by: Haicheng Wu <haichengw@nvidia.com>
This commit is contained in:
Yujia Zhai
co-authored by
yuzhai
Haicheng Wu
yuzhai
Haicheng Wu
parent
902dff3663
commit
b78588d163
@@ -0,0 +1,163 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 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 Distributed gemm device layer helpers.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cute/layout.hpp"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cutlass/cutlass.h"
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass::distributed::device::detail {
|
||||
|
||||
|
||||
cutlass::Status check_cuda_status(cudaError_t status) {
|
||||
if (status != cudaSuccess) {
|
||||
auto result = cudaGetLastError();
|
||||
CUTLASS_TRACE_HOST(" error message: " << cudaGetErrorString(result));
|
||||
return cutlass::Status::kErrorInternal;
|
||||
}
|
||||
return cutlass::Status::kSuccess;
|
||||
}
|
||||
|
||||
// DistGemmBufferHelper computes required buffer size and offsets for GEMM operands.
|
||||
template <
|
||||
typename Tiler_,
|
||||
typename ElementA_,
|
||||
typename ElementB_,
|
||||
typename ElementC_,
|
||||
typename ElementD_>
|
||||
struct DistGemmBufferHelper {
|
||||
|
||||
using Tiler = Tiler_;
|
||||
|
||||
using ElementA = ElementA_;
|
||||
using ElementB = ElementB_;
|
||||
using ElementC = ElementC_;
|
||||
using ElementD = ElementD_;
|
||||
|
||||
static constexpr int NumBuffersA = Tiler::NumBuffersA;
|
||||
static constexpr int NumBuffersB = Tiler::NumBuffersB;
|
||||
static constexpr int NumBuffersC = Tiler::NumBuffersC;
|
||||
static constexpr int NumBuffersD = Tiler::NumBuffersD;
|
||||
|
||||
template <typename ProblemShape>
|
||||
static auto
|
||||
get_buffer_size_a(ProblemShape problem_shape) {
|
||||
auto a_buffer_layout = cute::make_layout(
|
||||
cute::make_shape(NumBuffersA, Tiler::get_local_a_shape(problem_shape), sizeof(ElementA))
|
||||
);
|
||||
return size(a_buffer_layout);
|
||||
}
|
||||
|
||||
template <typename ProblemShape>
|
||||
static auto
|
||||
get_buffer_size_b(ProblemShape problem_shape) {
|
||||
auto b_buffer_layout = cute::make_layout(
|
||||
cute::make_shape(NumBuffersB, Tiler::get_local_b_shape(problem_shape), sizeof(ElementB))
|
||||
);
|
||||
return size(b_buffer_layout);
|
||||
}
|
||||
|
||||
template <typename ProblemShape>
|
||||
static auto
|
||||
get_buffer_size_c(ProblemShape problem_shape) {
|
||||
auto c_buffer_layout = cute::make_layout(
|
||||
cute::make_shape(NumBuffersC, Tiler::get_local_c_shape(problem_shape), sizeof(ElementC))
|
||||
);
|
||||
return size(c_buffer_layout);
|
||||
}
|
||||
|
||||
template <typename ProblemShape>
|
||||
static auto
|
||||
get_buffer_size_d(ProblemShape problem_shape) {
|
||||
auto d_buffer_layout = cute::make_layout(
|
||||
cute::make_shape(NumBuffersD, Tiler::get_local_d_shape(problem_shape), sizeof(ElementD))
|
||||
);
|
||||
return size(d_buffer_layout);
|
||||
}
|
||||
|
||||
template <typename ProblemShape>
|
||||
static auto
|
||||
get_buffer_size(ProblemShape problem_shape) {
|
||||
size_t buffer_size = 0;
|
||||
|
||||
if constexpr (NumBuffersA > 0) {
|
||||
buffer_size += get_buffer_size_a(problem_shape);
|
||||
}
|
||||
if constexpr (NumBuffersB > 0) {
|
||||
buffer_size += get_buffer_size_b(problem_shape);
|
||||
}
|
||||
if constexpr (NumBuffersC > 0) {
|
||||
buffer_size += get_buffer_size_c(problem_shape);
|
||||
}
|
||||
if constexpr (NumBuffersD > 0) {
|
||||
buffer_size += get_buffer_size_d(problem_shape);
|
||||
}
|
||||
|
||||
return buffer_size;
|
||||
}
|
||||
|
||||
// Buffer space: | buffer_A | buffer_B | buffer_C | buffer_D |
|
||||
// And buffer_{A,B,C,D}: | iter 1 | iter 2 | ... | iter TP - 1 |
|
||||
template <typename ProblemShape>
|
||||
static size_t
|
||||
get_buffer_offset_A(ProblemShape problem_shape) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
template <typename ProblemShape>
|
||||
static size_t
|
||||
get_buffer_offset_B(ProblemShape problem_shape) {
|
||||
return get_buffer_size_a(problem_shape);
|
||||
}
|
||||
|
||||
template <typename ProblemShape>
|
||||
static size_t
|
||||
get_buffer_offset_C(ProblemShape problem_shape) {
|
||||
return get_buffer_size_a(problem_shape) + get_buffer_size_b(problem_shape);
|
||||
}
|
||||
|
||||
template <typename ProblemShape>
|
||||
static size_t
|
||||
get_buffer_offset_D(ProblemShape problem_shape) {
|
||||
return get_buffer_size_a(problem_shape) + get_buffer_size_b(problem_shape) + get_buffer_size_c(problem_shape);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace cutlass::distributed::device::detail
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -0,0 +1,717 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 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 Distributed GEMM Device Adapter
|
||||
|
||||
Sets up local GEMM stages, the cuda graph, manages buffer and barrier spaces,
|
||||
and maps arguments to per-stage arguments.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/device_kernel.h"
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
|
||||
#include "cutlass/experimental/distributed/device/full_barrier.hpp"
|
||||
#include "cutlass/experimental/distributed/device/detail.hpp"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass::distributed::device {
|
||||
|
||||
template <class GemmKernel_>
|
||||
class DistributedGemmUniversalAdapter {
|
||||
public:
|
||||
using DeviceGemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel_>;
|
||||
using GemmKernel = GemmKernel_;
|
||||
using TileShape = typename GemmKernel::TileShape;
|
||||
using ElementA = typename GemmKernel::ElementA;
|
||||
using ElementB = typename GemmKernel::ElementB;
|
||||
using ElementC = typename GemmKernel::ElementC;
|
||||
using ElementD = typename GemmKernel::ElementD;
|
||||
using ElementAccumulator = typename GemmKernel::ElementAccumulator;
|
||||
using DispatchPolicy = typename GemmKernel::DispatchPolicy;
|
||||
using CollectiveMainloop = typename GemmKernel::CollectiveMainloop;
|
||||
using CollectiveEpilogue = typename GemmKernel::CollectiveEpilogue;
|
||||
|
||||
// "Inherit" type decls and static values from device GEMM
|
||||
using LayoutA = typename DeviceGemm::LayoutA;
|
||||
using LayoutB = typename DeviceGemm::LayoutB;
|
||||
using LayoutC = typename DeviceGemm::LayoutC;
|
||||
using LayoutD = typename DeviceGemm::LayoutD;
|
||||
|
||||
using StrideA = typename GemmKernel::StrideA;
|
||||
using StrideB = typename GemmKernel::StrideB;
|
||||
using StrideC = typename GemmKernel::StrideC;
|
||||
using StrideD = typename GemmKernel::StrideD;
|
||||
|
||||
static bool const kEnableCudaHostAdapter = DeviceGemm::kEnableCudaHostAdapter;
|
||||
|
||||
static ComplexTransform const kTransformA = DeviceGemm::kTransformA;
|
||||
static ComplexTransform const kTransformB = DeviceGemm::kTransformB;
|
||||
|
||||
using MathOperator = typename DeviceGemm::MathOperator;
|
||||
using OperatorClass = typename DeviceGemm::OperatorClass;
|
||||
using ArchTag = typename DeviceGemm::ArchTag;
|
||||
|
||||
using ThreadblockSwizzle = typename DeviceGemm::ThreadblockSwizzle;
|
||||
using ThreadblockShape = typename DeviceGemm::ThreadblockShape;
|
||||
using ClusterShape = typename DeviceGemm::ClusterShape;
|
||||
using InstructionShape = typename DeviceGemm::InstructionShape;
|
||||
|
||||
static int const kThreadCount = DeviceGemm::kThreadCount;
|
||||
static constexpr int WarpsInMma = DeviceGemm::WarpsInMma;
|
||||
static constexpr int WarpsInMmaM = DeviceGemm::WarpsInMmaM;
|
||||
static constexpr int WarpsInMmaN = DeviceGemm::WarpsInMmaN;
|
||||
|
||||
using WarpCount = typename DeviceGemm::WarpCount;
|
||||
using WarpShape = typename DeviceGemm::WarpShape;
|
||||
|
||||
static int constexpr kStages = DeviceGemm::kStages;
|
||||
|
||||
static int constexpr kAlignmentA = DeviceGemm::kAlignmentA;
|
||||
static int constexpr kAlignmentB = DeviceGemm::kAlignmentB;
|
||||
static int constexpr kAlignmentC = DeviceGemm::kAlignmentC;
|
||||
static int constexpr kAlignmentD = DeviceGemm::kAlignmentD;
|
||||
|
||||
using EpilogueOutputOp = typename DeviceGemm::EpilogueOutputOp;
|
||||
|
||||
static int constexpr kSplitKAlignment = DeviceGemm::kSplitKAlignment;
|
||||
|
||||
// Distributed GEMM types and defs
|
||||
using DistSchedule = typename GemmKernel::DistSchedule;
|
||||
static constexpr bool HasMemcpy = DistSchedule::HasMemcpy;
|
||||
using TP = typename DistSchedule::TP;
|
||||
static constexpr int TP_ = TP{};
|
||||
using ElementFlag = typename GemmKernel::ElementFlag;
|
||||
using ElementBarrier = uint32_t;
|
||||
|
||||
using BufferHelper = detail::DistGemmBufferHelper<
|
||||
DistSchedule,
|
||||
ElementA,
|
||||
ElementB,
|
||||
ElementC,
|
||||
ElementD>;
|
||||
|
||||
/// Argument structure
|
||||
using Arguments = typename GemmKernel::BaseArguments;
|
||||
using DistributedArguments = typename GemmKernel::DistributedArguments;
|
||||
using PackedArguments = typename GemmKernel::PackedArguments;
|
||||
|
||||
/// Argument structure: Kernel API
|
||||
using Params = typename GemmKernel::PackedParams;
|
||||
|
||||
struct DistributedGemmState {
|
||||
int device_idx;
|
||||
|
||||
Params params_array[TP_];
|
||||
|
||||
cudaGraph_t graph;
|
||||
cudaGraphExec_t graph_executable;
|
||||
|
||||
bool graph_created = false;
|
||||
bool graph_instantiated = false;
|
||||
|
||||
void * memcpy_source_ptr_array[TP_];
|
||||
void const * memcpy_remote_ptr_array[TP_];
|
||||
size_t memcpy_bytes[TP_];
|
||||
|
||||
cutlass::Array<ElementBarrier*, TP_> device_barrier_ptrs;
|
||||
|
||||
bool is_initialized = false;
|
||||
};
|
||||
|
||||
private:
|
||||
|
||||
DistributedGemmState state_;
|
||||
|
||||
public:
|
||||
|
||||
bool is_initialized() {
|
||||
return state_.is_initialized && state_.graph_created && state_.graph_instantiated;
|
||||
}
|
||||
|
||||
/// Determines whether the GEMM can execute the given problem.
|
||||
static Status
|
||||
can_implement(Arguments const& args) {
|
||||
if (args.epilogue.thread.beta != 0.0 && DistSchedule::RemoteC) {
|
||||
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Selected TP uses Remote C to communicate " <<
|
||||
"partial results, which do not support non-zero values for beta yet " <<
|
||||
"(epilogue must be sourceless.)\n");
|
||||
return Status::kInvalid;
|
||||
}
|
||||
|
||||
if (not DistSchedule::can_implement_global(args.problem_shape)) {
|
||||
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem shape not divisible by TP.\n");
|
||||
return Status::kInvalid;
|
||||
}
|
||||
|
||||
Arguments args_copy = args;
|
||||
args_copy.problem_shape = DistSchedule::get_local_gemm_shape(args.problem_shape);
|
||||
for (int iteration = 0; iteration < TP_; ++iteration) {
|
||||
if (not GemmKernel::can_implement(args_copy)) {
|
||||
return Status::kInvalid;
|
||||
}
|
||||
}
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
/// Gets buffer space size
|
||||
static size_t
|
||||
get_buffer_space_size(Arguments const& args) {
|
||||
size_t buffer_bytes = 0;
|
||||
|
||||
buffer_bytes = BufferHelper::get_buffer_size(args.problem_shape);
|
||||
buffer_bytes = round_nearest(buffer_bytes, MinWorkspaceAlignment);
|
||||
|
||||
return buffer_bytes;
|
||||
}
|
||||
|
||||
static auto
|
||||
get_tensor_A_for_iter(Arguments const* args_array, void** buffer_space, int device_idx, int iteration) {
|
||||
auto args = args_array[device_idx];
|
||||
auto tensor_A = make_tensor(args.mainloop.ptr_A, make_layout(
|
||||
DistSchedule::get_local_a_shape(args.problem_shape),
|
||||
args.mainloop.dA));
|
||||
|
||||
uint8_t* tensor_buffer = reinterpret_cast<uint8_t*>(buffer_space[device_idx]) +
|
||||
BufferHelper::get_buffer_offset_A(args.problem_shape);
|
||||
|
||||
return DistSchedule::get_tensor_A(tensor_A, tensor_buffer, device_idx, iteration);
|
||||
}
|
||||
|
||||
static auto
|
||||
get_tensor_B_for_iter(Arguments const* args_array, void** buffer_space, int device_idx, int iteration) {
|
||||
auto args = args_array[device_idx];
|
||||
auto tensor_B = make_tensor(args.mainloop.ptr_B, make_layout(
|
||||
DistSchedule::get_local_b_shape(args.problem_shape),
|
||||
args.mainloop.dB));
|
||||
|
||||
uint8_t* tensor_buffer = reinterpret_cast<uint8_t*>(buffer_space[device_idx]) +
|
||||
BufferHelper::get_buffer_offset_B(args.problem_shape);
|
||||
|
||||
return DistSchedule::get_tensor_B(tensor_B, tensor_buffer, device_idx, iteration);
|
||||
}
|
||||
|
||||
static auto
|
||||
get_tensor_C_for_iter(Arguments const* args_array, void** buffer_space, int device_idx, int iteration) {
|
||||
auto args = args_array[device_idx];
|
||||
auto tensor_C = make_tensor(args.epilogue.ptr_C, make_layout(
|
||||
DistSchedule::get_local_c_shape(args.problem_shape),
|
||||
args.epilogue.dC));
|
||||
|
||||
auto peer_idx_iter = DistSchedule::get_remote_peer_id(device_idx, iteration);
|
||||
void* buffer_ptr = DistSchedule::RemoteC ? buffer_space[peer_idx_iter] : buffer_space[device_idx];
|
||||
|
||||
uint8_t* tensor_buffer = reinterpret_cast<uint8_t*>(buffer_ptr) +
|
||||
BufferHelper::get_buffer_offset_C(args.problem_shape);
|
||||
|
||||
return DistSchedule::get_tensor_C(tensor_C, tensor_buffer, device_idx, iteration);
|
||||
}
|
||||
|
||||
static auto
|
||||
get_tensor_D_for_iter(Arguments const* args_array, void** buffer_space, int device_idx, int iteration) {
|
||||
auto args = args_array[device_idx];
|
||||
auto tensor_D = make_tensor(args.epilogue.ptr_D, make_layout(
|
||||
DistSchedule::get_local_d_shape(args.problem_shape),
|
||||
args.epilogue.dD));
|
||||
|
||||
// support remoteD
|
||||
uint8_t* tensor_buffer = reinterpret_cast<uint8_t*>(buffer_space[device_idx]) +
|
||||
BufferHelper::get_buffer_offset_D(args.problem_shape);
|
||||
|
||||
return DistSchedule::get_tensor_D(tensor_D, tensor_buffer, device_idx, iteration);
|
||||
}
|
||||
|
||||
static size_t
|
||||
get_workspace_size(Arguments const& args) {
|
||||
size_t workspace_bytes = 0;
|
||||
|
||||
workspace_bytes = get_buffer_space_size(args);
|
||||
|
||||
for (int iteration = 0; iteration < TP_; ++iteration) {
|
||||
// NOTE: assumes underlying kernels align up to alignment requirements on their own,
|
||||
// and that the alignment requirements of the individual kernels match.
|
||||
workspace_bytes += GemmKernel::get_workspace_size(args);
|
||||
}
|
||||
|
||||
return workspace_bytes;
|
||||
}
|
||||
|
||||
static size_t
|
||||
get_barrier_bytes() {
|
||||
return round_nearest(sizeof(ElementBarrier), 32);
|
||||
}
|
||||
|
||||
static size_t
|
||||
get_flag_bytes() {
|
||||
return round_nearest(sizeof(ElementFlag) * TP_, 32);
|
||||
}
|
||||
|
||||
static void *
|
||||
exclusive_workspace_ptr_to_flag_ptr(void * exclusive_workspace_ptr, int iteration) {
|
||||
return static_cast<void*>(
|
||||
static_cast<uint8_t*>(exclusive_workspace_ptr) +
|
||||
get_barrier_bytes() +
|
||||
(sizeof(ElementFlag) * iteration));
|
||||
}
|
||||
|
||||
static size_t
|
||||
get_exclusive_workspace_size() {
|
||||
return get_barrier_bytes() + get_flag_bytes();
|
||||
}
|
||||
|
||||
/// Initializes GEMM state from arguments.
|
||||
Status
|
||||
initialize(
|
||||
Arguments const* args,
|
||||
void** workspace_ptrs,
|
||||
void** exclusive_workspace_ptrs,
|
||||
int device_idx,
|
||||
cudaStream_t stream = nullptr,
|
||||
bool launch_with_pdl = false) {
|
||||
|
||||
CUTLASS_TRACE_HOST("DistributedGemm::initialize() - stream: " << (stream ? "non-null" : "null"));
|
||||
|
||||
state_.device_idx = device_idx;
|
||||
|
||||
for (int device = 0; device < TP_; ++device) {
|
||||
state_.device_barrier_ptrs[device] = reinterpret_cast<ElementBarrier*>(exclusive_workspace_ptrs[device]);
|
||||
}
|
||||
|
||||
// Zero out exclusive workspace
|
||||
zero_workspace(exclusive_workspace_ptrs[device_idx], get_exclusive_workspace_size(), stream, nullptr);
|
||||
|
||||
for (int iteration = 0; iteration < TP_; ++iteration) {
|
||||
|
||||
size_t workspace_iteration_offset = GemmKernel::get_workspace_size(args[device_idx]);
|
||||
uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace_ptrs[device_idx]) +
|
||||
get_buffer_space_size(args[device_idx]) +
|
||||
(iteration * workspace_iteration_offset);
|
||||
|
||||
void * workspace_iter = reinterpret_cast<void*>(workspace_ptr);
|
||||
void** buffer_space = workspace_ptrs;
|
||||
|
||||
// Set up GEMM arguments for the current stage/iteration
|
||||
auto tensor_a_iter = get_tensor_A_for_iter(args, buffer_space, device_idx, iteration);
|
||||
auto tensor_b_iter = get_tensor_B_for_iter(args, buffer_space, device_idx, iteration);
|
||||
auto tensor_c_iter = get_tensor_C_for_iter(args, buffer_space, device_idx, iteration);
|
||||
auto tensor_d_iter = get_tensor_D_for_iter(args, buffer_space, device_idx, iteration);
|
||||
|
||||
Arguments base_args = args[device_idx];
|
||||
base_args.problem_shape = DistSchedule::get_local_gemm_shape(args[device_idx].problem_shape);
|
||||
base_args.mainloop = {
|
||||
reinterpret_cast<const ElementA*>(tensor_a_iter.data()),
|
||||
tensor_a_iter.stride(),
|
||||
reinterpret_cast<const ElementB*>(tensor_b_iter.data()),
|
||||
tensor_b_iter.stride()
|
||||
};
|
||||
base_args.epilogue = {
|
||||
base_args.epilogue.thread,
|
||||
reinterpret_cast<const ElementC*>(tensor_c_iter.data()),
|
||||
tensor_c_iter.stride(),
|
||||
reinterpret_cast<const ElementD*>(tensor_d_iter.data()),
|
||||
tensor_d_iter.stride()
|
||||
};
|
||||
|
||||
if constexpr (DistSchedule::RemoteC) {
|
||||
if (iteration > 0) {
|
||||
base_args.epilogue.thread.beta = 1.0;
|
||||
}
|
||||
else if (iteration == 0){
|
||||
base_args.epilogue.thread.beta = 0.0;
|
||||
}
|
||||
}
|
||||
|
||||
auto [left_peer_idx, right_peer_idx] = DistSchedule::get_peers_for_device(device_idx);
|
||||
auto flag_peer_idx = DistSchedule::KernelWritesArrivalFlag ? right_peer_idx : device_idx;
|
||||
|
||||
void * self_flag_ptr = exclusive_workspace_ptr_to_flag_ptr(exclusive_workspace_ptrs[device_idx], iteration);
|
||||
void * peer_flag_ptr = exclusive_workspace_ptr_to_flag_ptr(exclusive_workspace_ptrs[flag_peer_idx], iteration);
|
||||
|
||||
DistributedArguments distributed_args = {
|
||||
device_idx,
|
||||
iteration,
|
||||
self_flag_ptr,
|
||||
peer_flag_ptr
|
||||
};
|
||||
PackedArguments args_iter = {base_args, distributed_args};
|
||||
|
||||
// Initialize the workspace
|
||||
Status status = GemmKernel::initialize_workspace(args_iter, workspace_iter, stream);
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
// Initialize the Params structure
|
||||
state_.params_array[iteration] = GemmKernel::to_underlying_arguments(args_iter, workspace_iter);
|
||||
|
||||
// Set up peer buffer ptrs
|
||||
if (iteration > 0 && HasMemcpy) {
|
||||
auto peer_idx_iter = DistSchedule::get_remote_peer_id(device_idx, iteration);
|
||||
|
||||
void * local_ptr_itr = nullptr;
|
||||
void const * remote_ptr_itr = nullptr;
|
||||
size_t local_size = 0;
|
||||
size_t remote_size = 0;
|
||||
|
||||
static_assert(not DistSchedule::HasMemcpy || (
|
||||
DistSchedule::MemcpyA || DistSchedule::MemcpyB),
|
||||
"Expected to either memcpy A or B when scheduler requires memcpy.");
|
||||
if constexpr (DistSchedule::MemcpyA) {
|
||||
local_size = cute::cosize(tensor_a_iter.layout()) * sizeof(ElementA);
|
||||
local_ptr_itr = reinterpret_cast<void*>(tensor_a_iter.data());
|
||||
|
||||
// Copy peer's slice in the first iteration (direct access memcpy instead of logical ring)
|
||||
auto remote_tensor_iter = get_tensor_A_for_iter(args, buffer_space, peer_idx_iter, 0);
|
||||
remote_ptr_itr = reinterpret_cast<void const*>(remote_tensor_iter.data());
|
||||
remote_size = cute::cosize(remote_tensor_iter.layout()) * sizeof(ElementA);
|
||||
}
|
||||
else if constexpr (DistSchedule::MemcpyB) {
|
||||
local_size = cute::cosize(tensor_b_iter.layout()) * sizeof(ElementB);
|
||||
local_ptr_itr = reinterpret_cast<void*>(tensor_b_iter.data());
|
||||
|
||||
// Copy peer's slice in the first iteration (direct access memcpy instead of logical ring)
|
||||
auto remote_tensor_iter = get_tensor_B_for_iter(args, buffer_space, peer_idx_iter, 0);
|
||||
remote_ptr_itr = reinterpret_cast<void const*>(remote_tensor_iter.data());
|
||||
remote_size = cute::cosize(remote_tensor_iter.layout()) * sizeof(ElementB);
|
||||
}
|
||||
|
||||
assert(local_size == remote_size && local_size > 0);
|
||||
|
||||
state_.memcpy_source_ptr_array[iteration] = local_ptr_itr;
|
||||
state_.memcpy_remote_ptr_array[iteration] = remote_ptr_itr;
|
||||
state_.memcpy_bytes[iteration] = local_size;
|
||||
}
|
||||
}
|
||||
|
||||
//
|
||||
// Account for dynamic smem capacity if needed
|
||||
//
|
||||
int smem_size = GemmKernel::SharedStorageSize;
|
||||
|
||||
if (smem_size >= (48 << 10)) {
|
||||
CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size);
|
||||
cudaError_t result = cudaFuncSetAttribute(
|
||||
device_kernel<GemmKernel>,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize,
|
||||
smem_size);
|
||||
if (cudaSuccess != result) {
|
||||
result = cudaGetLastError(); // to clear the error bit
|
||||
CUTLASS_TRACE_HOST(" cudaFuncSetAttribute() returned error: " << cudaGetErrorString(result));
|
||||
return Status::kErrorInternal;
|
||||
}
|
||||
}
|
||||
|
||||
state_.is_initialized = true;
|
||||
|
||||
// Instantiate graph
|
||||
Status status = construct_graph(launch_with_pdl);
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
Status
|
||||
construct_graph(bool launch_with_pdl) {
|
||||
#if ((__CUDACC_VER_MAJOR__ >= 12) && (__CUDACC_VER_MINOR__ >= 4))
|
||||
Status status = Status::kSuccess;
|
||||
|
||||
// Destroy existing graph, if created
|
||||
if (state_.graph_created) {
|
||||
status = detail::check_cuda_status(cudaGraphDestroy(state_.graph));
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
}
|
||||
|
||||
state_.graph_created = true;
|
||||
|
||||
cudaGraphNode_t full_barrier_node;
|
||||
|
||||
// Create dummy stream
|
||||
cudaStream_t stream;
|
||||
status = detail::check_cuda_status(cudaStreamCreate(&stream));
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
// Create graph
|
||||
status = detail::check_cuda_status(cudaGraphCreate(&state_.graph, 0));
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
// 1. Full barrier node
|
||||
status = detail::check_cuda_status(cudaStreamBeginCaptureToGraph(
|
||||
stream,
|
||||
state_.graph,
|
||||
nullptr, nullptr, 0,
|
||||
cudaStreamCaptureModeRelaxed));
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
cutlass::Array<ElementFlag*, TP_> self_flag_ptrs;
|
||||
for (int iteration = 0; iteration < TP_; ++iteration) {
|
||||
self_flag_ptrs[iteration] = state_.params_array[iteration].distributed.self_flag_ptr_;
|
||||
}
|
||||
|
||||
launch_full_barrier<TP_, ElementBarrier, TP_, ElementFlag>(
|
||||
state_.device_barrier_ptrs, self_flag_ptrs, state_.device_idx, stream, launch_with_pdl);
|
||||
|
||||
status = detail::check_cuda_status(cudaStreamEndCapture(stream, &state_.graph));
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
size_t num_nodes;
|
||||
status = detail::check_cuda_status(cudaGraphGetNodes(state_.graph, nullptr, &num_nodes));
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
if (num_nodes != 1) {
|
||||
CUTLASS_TRACE_HOST(" construct_graph() failure: expected a single node in the graph, got " << num_nodes << ".");
|
||||
return Status::kErrorInternal;
|
||||
}
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
status = detail::check_cuda_status(cudaGraphGetNodes(state_.graph, &full_barrier_node, &num_nodes));
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
// 2. Optional mem copy branch
|
||||
if constexpr (HasMemcpy) {
|
||||
|
||||
status = detail::check_cuda_status(cudaStreamBeginCaptureToGraph(
|
||||
stream,
|
||||
state_.graph,
|
||||
&full_barrier_node,
|
||||
/* dependencyData = */ nullptr,
|
||||
1,
|
||||
cudaStreamCaptureModeRelaxed));
|
||||
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
// No copies for first iter; we assume the data is already there.
|
||||
for (int iteration = 1; iteration < TP_; ++iteration) {
|
||||
|
||||
status = detail::check_cuda_status(cudaMemcpyAsync(
|
||||
state_.memcpy_source_ptr_array[iteration],
|
||||
state_.memcpy_remote_ptr_array[iteration],
|
||||
state_.memcpy_bytes[iteration],
|
||||
cudaMemcpyDeviceToDevice, stream));
|
||||
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
// Set flag to non zero
|
||||
status = detail::check_cuda_status(cudaMemsetAsync(
|
||||
reinterpret_cast<void *>(state_.params_array[iteration].distributed.peer_flag_ptr_),
|
||||
0b11111111,
|
||||
sizeof(ElementFlag),
|
||||
stream));
|
||||
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
}
|
||||
|
||||
status = detail::check_cuda_status(cudaStreamEndCapture(stream, &state_.graph));
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
}
|
||||
|
||||
// 3. Run local GEMMs
|
||||
// 3.1. Create edge between full barrier and the correct gemm stage/iteration
|
||||
cudaGraphEdgeData barrier_to_gemm_edge = {};
|
||||
barrier_to_gemm_edge.from_port = HasMemcpy ? cudaGraphKernelNodePortLaunchCompletion: cudaGraphKernelNodePortProgrammatic;
|
||||
barrier_to_gemm_edge.type = cudaGraphDependencyTypeProgrammatic;
|
||||
|
||||
status = detail::check_cuda_status(cudaStreamBeginCaptureToGraph(
|
||||
stream,
|
||||
state_.graph,
|
||||
&full_barrier_node,
|
||||
/* dependencyData = */ &barrier_to_gemm_edge,
|
||||
1,
|
||||
cudaStreamCaptureModeRelaxed));
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
for (int iteration = 0; iteration < TP_; ++iteration) {
|
||||
status = DeviceGemm::run(
|
||||
state_.params_array[iteration],
|
||||
stream,
|
||||
/* cuda_adapter = */ nullptr,
|
||||
/* launch_with_pdl = */ launch_with_pdl);
|
||||
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
}
|
||||
|
||||
status = detail::check_cuda_status(cudaStreamEndCapture(stream, &state_.graph));
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
// 4. Cleanup.
|
||||
//// Destroy dummy stream
|
||||
status = detail::check_cuda_status(cudaStreamDestroy(stream));
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
// 5. Instantiate graph
|
||||
status = detail::check_cuda_status(cudaGraphInstantiate(
|
||||
&state_.graph_executable,
|
||||
state_.graph,
|
||||
/* flags = */ 0));
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
state_.graph_instantiated = true;
|
||||
|
||||
return Status::kSuccess;
|
||||
#else
|
||||
CUTLASS_TRACE_HOST(" construct_graph() failure: target was compiled with an incompatible " <<
|
||||
"version of the CUDA toolkit. Please compile Distributed GEMM with CUDA toolkit 12.4 or later.");
|
||||
return Status::kErrorInternal;
|
||||
#endif
|
||||
}
|
||||
|
||||
Status
|
||||
update(Arguments const& args, void* workspace = nullptr) {
|
||||
CUTLASS_TRACE_HOST(" DistributedGemm does not support updating arguments yet.");
|
||||
return Status::kErrorInternal;
|
||||
}
|
||||
|
||||
// NOTE: the interface for run() is different in Distributed Gemm:
|
||||
// 1. launch_with_pdl is specified in `initialize`, where the cuda graph is being constructed,
|
||||
// 2. the state of distributed gemm is an array of params for different iterations, and a
|
||||
// cuda graph.
|
||||
// 3. Custom cuda adapters aren't supported for simplicity.
|
||||
static Status
|
||||
run(DistributedGemmState& state,
|
||||
cudaStream_t stream = nullptr) {
|
||||
CUTLASS_TRACE_HOST("DistributedGemm::run()");
|
||||
|
||||
if (not state.is_initialized) {
|
||||
CUTLASS_TRACE_HOST(" Distributed gemm was not initialized. Did you forget to call initialize()?");
|
||||
return Status::kErrorInternal;
|
||||
}
|
||||
|
||||
if (not state.graph_instantiated) {
|
||||
CUTLASS_TRACE_HOST(" Distributed gemm graph was not instantiated. Did you forget to call initialize()/construct_graph()?");
|
||||
return Status::kErrorInternal;
|
||||
}
|
||||
|
||||
cudaError_t result = cudaGraphLaunch(state.graph_executable, stream);
|
||||
if (cudaSuccess != result) {
|
||||
result = cudaGetLastError(); // to clear the error bit
|
||||
CUTLASS_TRACE_HOST(" cudaGraphLaunch() returned error: " << cudaGetErrorString(result));
|
||||
return Status::kErrorInternal;
|
||||
}
|
||||
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
//
|
||||
// Non-static launch overloads that first create and set the internal params struct of this kernel handle.
|
||||
//
|
||||
|
||||
/// Overload that allows a user to re-launch the same kernel without updating internal params struct.
|
||||
Status
|
||||
run(
|
||||
cudaStream_t stream = nullptr) {
|
||||
return run(state_, stream);
|
||||
}
|
||||
|
||||
/// Overload that allows a user to re-launch the same kernel without updating internal params struct.
|
||||
Status
|
||||
operator()(cudaStream_t stream = nullptr) {
|
||||
return run(state_, stream);
|
||||
}
|
||||
|
||||
/// Launches the kernel after first constructing Params internal state from supplied arguments.
|
||||
Status
|
||||
run(
|
||||
Arguments const* args,
|
||||
void** workspace_ptrs,
|
||||
void** exclusive_workspace_ptrs,
|
||||
int device_idx,
|
||||
cudaStream_t stream = nullptr) {
|
||||
Status status = initialize(
|
||||
args,
|
||||
workspace_ptrs,
|
||||
exclusive_workspace_ptrs,
|
||||
device_idx,
|
||||
stream);
|
||||
|
||||
if (Status::kSuccess == status) {
|
||||
status = run(stream);
|
||||
}
|
||||
return status;
|
||||
}
|
||||
|
||||
/// Launches the kernel after first constructing Params internal state from supplied arguments.
|
||||
Status
|
||||
operator()(
|
||||
Arguments const* args,
|
||||
void** workspace_ptrs,
|
||||
void** exclusive_workspace_ptrs,
|
||||
int device_idx,
|
||||
cudaStream_t stream = nullptr) {
|
||||
return run(
|
||||
args,
|
||||
workspace_ptrs,
|
||||
exclusive_workspace_ptrs,
|
||||
device_idx,
|
||||
stream);
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace cutlass::distributed::device
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,74 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 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 Device layer interface for Distributed GEMM barrier kernel.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/experimental/distributed/kernel/full_barrier.hpp"
|
||||
|
||||
namespace cutlass::distributed::device {
|
||||
|
||||
template <int NP, typename IntType, int Iterations, typename FlagType>
|
||||
void launch_full_barrier(
|
||||
cutlass::Array<IntType*, NP> device_arrival_ptrs,
|
||||
cutlass::Array<FlagType*, Iterations> iteration_flag_ptrs,
|
||||
IntType device_idx,
|
||||
cudaStream_t stream,
|
||||
bool launch_with_pdl) {
|
||||
|
||||
#if ((__CUDACC_VER_MAJOR__ >= 12) && (__CUDACC_VER_MINOR__ >= 4))
|
||||
// Legacy (kernel) launch with PDL
|
||||
cudaLaunchAttribute attributes[1];
|
||||
attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
|
||||
attributes[0].val.programmaticStreamSerializationAllowed = 1;
|
||||
|
||||
cudaLaunchConfig_t launch_config;
|
||||
launch_config.gridDim = 1;
|
||||
launch_config.blockDim = 1;
|
||||
launch_config.dynamicSmemBytes = 0;
|
||||
launch_config.stream = stream;
|
||||
launch_config.attrs = attributes;
|
||||
launch_config.numAttrs = launch_with_pdl ? 1 : 0;
|
||||
|
||||
cudaLaunchKernelEx(
|
||||
&launch_config,
|
||||
cutlass::distributed::kernel::full_barrier_kernel<NP, IntType, Iterations, FlagType>,
|
||||
device_arrival_ptrs,
|
||||
iteration_flag_ptrs,
|
||||
device_idx);
|
||||
#endif
|
||||
}
|
||||
|
||||
} // namespace cutlass::distributed::device
|
||||
|
||||
@@ -0,0 +1,72 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 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 Distributed gemm kernel layer helpers.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass::distributed::kernel::detail {
|
||||
|
||||
// Ld with CV cache hint (don’t cache and fetch again)
|
||||
// Reference:
|
||||
// https://docs.nvidia.com/cuda/parallel-thread-execution/#cache-operators
|
||||
// Used for loading arrival counts from peer devices
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void ld_without_cache(uint64_t& val, void const * ptr) {
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" ld.global.cv.u64 %0, [%1];\n"
|
||||
"}\n"
|
||||
: "=l"(val)
|
||||
: "l"(ptr));
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void ld_without_cache(uint32_t& val, void const * ptr) {
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" ld.global.cv.u32 %0, [%1];\n"
|
||||
"}\n"
|
||||
: "=r"(val)
|
||||
: "l"(ptr));
|
||||
}
|
||||
|
||||
} // namespace cutlass::distributed::kernel::detail
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
|
||||
@@ -0,0 +1,235 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 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 Distributed GEMM Kernel Wrapper
|
||||
|
||||
Prepends CUTLASS 3 GEMM kernels with barriers and other necessary instructions to exectue
|
||||
a Distributed GEMM stage.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/arch/grid_dependency_control.h"
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
|
||||
#include "cutlass/experimental/distributed/kernel/detail.hpp"
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass::distributed::kernel {
|
||||
|
||||
namespace detail {
|
||||
|
||||
// Allow all CUTLASS 3.X GEMM kernels
|
||||
template <typename GemmKernel_>
|
||||
struct SupportsDistributedGemm: cutlass::gemm::detail::IsCutlass3GemmKernel<GemmKernel_> {};
|
||||
|
||||
} // namespace detail
|
||||
|
||||
/*!
|
||||
DistributedGemmKernelWrapper is a wrapper around a GEMM kernel.
|
||||
|
||||
Depending on the underlying distribution policy/schedule, it prepends the underlying local GEMM
|
||||
kernel with a few additional instructions that gate the execution of the GEMM on buffers being
|
||||
ready for stages/iterations > 0.
|
||||
*/
|
||||
|
||||
template <class GemmKernel_, class DistSchedule_, class Enable = void>
|
||||
struct DistributedGemmKernelWrapper;
|
||||
|
||||
template <class GemmKernel_, class DistSchedule_>
|
||||
struct DistributedGemmKernelWrapper<
|
||||
GemmKernel_,
|
||||
DistSchedule_,
|
||||
cute::enable_if_t<detail::SupportsDistributedGemm<GemmKernel_>::value>
|
||||
>: GemmKernel_
|
||||
{
|
||||
using DistSchedule = DistSchedule_;
|
||||
using TP = typename DistSchedule::TP;
|
||||
|
||||
static constexpr bool KernelWritesArrivalFlag = DistSchedule::KernelWritesArrivalFlag;
|
||||
|
||||
using BaseKernel = GemmKernel_;
|
||||
using BaseArguments = typename BaseKernel::Arguments;
|
||||
using BaseParams = typename BaseKernel::Params;
|
||||
|
||||
static_assert(BaseKernel::ArchTag::kMinComputeCapability == 90, "DistGEMM only supports Hopper GEMMs for now.");
|
||||
static_assert(not cute::is_same_v<typename BaseKernel::ElementC, void>, "DistributedGEMM epilogues must have a source.");
|
||||
|
||||
using ElementFlag = uint32_t;
|
||||
|
||||
// Device side arguments
|
||||
struct DistributedArguments {
|
||||
int device_idx = 0;
|
||||
int iteration = 0;
|
||||
|
||||
void* self_flag_ptr{nullptr};
|
||||
void* peer_flag_ptr{nullptr};
|
||||
};
|
||||
|
||||
struct PackedArguments {
|
||||
BaseArguments base{};
|
||||
DistributedArguments distributed{};
|
||||
};
|
||||
|
||||
struct DistributedParams {
|
||||
int device_idx = 0;
|
||||
int iteration = 0;
|
||||
|
||||
ElementFlag* self_flag_ptr_{nullptr};
|
||||
ElementFlag* peer_flag_ptr_{nullptr};
|
||||
};
|
||||
|
||||
// Kernel entry point API
|
||||
struct PackedParams {
|
||||
BaseParams base{};
|
||||
DistributedParams distributed{};
|
||||
};
|
||||
|
||||
using Params = PackedParams;
|
||||
|
||||
// Convert to underlying arguments. In this case, a simple copy for the aliased type.
|
||||
static
|
||||
PackedParams
|
||||
to_underlying_arguments(PackedArguments const& args, void* workspace) {
|
||||
CUTLASS_TRACE_HOST("distributed::to_underlying_arguments():");
|
||||
|
||||
auto kernel_params = BaseKernel::to_underlying_arguments(args.base, workspace);
|
||||
|
||||
DistributedParams dist_params = {
|
||||
args.distributed.device_idx,
|
||||
args.distributed.iteration,
|
||||
reinterpret_cast<ElementFlag*>(args.distributed.self_flag_ptr),
|
||||
reinterpret_cast<ElementFlag*>(args.distributed.peer_flag_ptr)
|
||||
};
|
||||
|
||||
return {kernel_params, dist_params};
|
||||
}
|
||||
|
||||
static bool
|
||||
can_implement(BaseArguments const& args) {
|
||||
return BaseKernel::can_implement(args);
|
||||
}
|
||||
|
||||
static bool
|
||||
can_implement(PackedArguments const& args) {
|
||||
return BaseKernel::can_implement(args.base);
|
||||
}
|
||||
|
||||
static size_t
|
||||
get_workspace_size(BaseArguments const& args) {
|
||||
return BaseKernel::get_workspace_size(args);
|
||||
}
|
||||
|
||||
static size_t
|
||||
get_workspace_size(PackedArguments const& args) {
|
||||
return BaseKernel::get_workspace_size(args.base);
|
||||
}
|
||||
|
||||
static cutlass::Status
|
||||
initialize_workspace(BaseArguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr,
|
||||
CudaHostAdapter* cuda_adapter = nullptr) {
|
||||
return BaseKernel::initialize_workspace(args, workspace, stream, cuda_adapter);
|
||||
}
|
||||
|
||||
static cutlass::Status
|
||||
initialize_workspace(PackedArguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr,
|
||||
CudaHostAdapter* cuda_adapter = nullptr) {
|
||||
return BaseKernel::initialize_workspace(args.base, workspace, stream, cuda_adapter);
|
||||
}
|
||||
|
||||
/// Computes the grid shape
|
||||
static dim3
|
||||
get_grid_shape(PackedParams const& params) {
|
||||
return BaseKernel::get_grid_shape(params.base);
|
||||
}
|
||||
|
||||
static dim3
|
||||
get_grid_shape(BaseParams const& params) {
|
||||
return BaseKernel::get_grid_shape(params);
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void
|
||||
barrier_buffer(PackedParams const& params) {
|
||||
if (params.distributed.iteration > 0) {
|
||||
|
||||
ElementFlag comm_iter = 0;
|
||||
detail::ld_without_cache(comm_iter, params.distributed.self_flag_ptr_);
|
||||
while (comm_iter == 0) {
|
||||
detail::ld_without_cache(comm_iter, params.distributed.self_flag_ptr_);
|
||||
__nanosleep(40);
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void
|
||||
maybe_signal_arrival(PackedParams const& params) {
|
||||
if constexpr (KernelWritesArrivalFlag) {
|
||||
if (blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0 &&
|
||||
threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0 &&
|
||||
params.distributed.iteration > 0) {
|
||||
*reinterpret_cast<ElementFlag*>(params.distributed.peer_flag_ptr_) = 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void
|
||||
operator()(PackedParams const& params, char* smem_buf) {
|
||||
// Launch next grid as soon as possible
|
||||
arch::launch_dependent_grids();
|
||||
|
||||
// Wait on previous kernels to flush their memory.
|
||||
arch::wait_on_dependent_grids();
|
||||
|
||||
// Optionally write arrivals for the previous stage/iteration.
|
||||
maybe_signal_arrival(params);
|
||||
|
||||
// Spin-wait on an arrival flag, make sure the respective buffers are ready.
|
||||
// If the buffered operand is memcpied into, it would wait on its local flag.
|
||||
// If it's a remote buffer that is accessed directly, it would wait on its remote flag.
|
||||
barrier_buffer(params);
|
||||
|
||||
// Perform local gemm
|
||||
BaseKernel gemm;
|
||||
gemm(params.base, smem_buf);
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
} // namespace cutlass::distributed::kernel
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 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 Distributed GEMM barrier kernel.
|
||||
|
||||
The kernel resets the per-stage arrival flags, performs a full barrier (any-to-any),
|
||||
and also atomically resets the local barrier arrival count.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/arch/grid_dependency_control.h"
|
||||
|
||||
#include "cutlass/experimental/distributed/kernel/detail.hpp"
|
||||
|
||||
namespace cutlass::distributed::kernel {
|
||||
|
||||
template <int NP, typename IntType, int Iterations, typename FlagType>
|
||||
__global__ void full_barrier_kernel(
|
||||
cutlass::Array<IntType*, NP> device_arrival_ptrs,
|
||||
cutlass::Array<FlagType*, Iterations> iteration_flag_ptrs,
|
||||
IntType device_idx) {
|
||||
|
||||
arch::launch_dependent_grids();
|
||||
arch::wait_on_dependent_grids();
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (FlagType i = 0; i < Iterations; ++i) {
|
||||
iteration_flag_ptrs[i][0] = static_cast<FlagType>(0);
|
||||
}
|
||||
|
||||
IntType val = 1;
|
||||
IntType max_val = static_cast<IntType>(NP - 1);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (IntType d = 0; d < NP; ++d) {
|
||||
if (d != device_idx) {
|
||||
atomicAdd(device_arrival_ptrs[d], val);
|
||||
}
|
||||
}
|
||||
|
||||
IntType curr_val = 0;
|
||||
detail::ld_without_cache(curr_val, device_arrival_ptrs[device_idx]);
|
||||
while (curr_val < max_val) {
|
||||
__nanosleep(40);
|
||||
detail::ld_without_cache(curr_val, device_arrival_ptrs[device_idx]);
|
||||
}
|
||||
|
||||
atomicSub(device_arrival_ptrs[device_idx], max_val);
|
||||
}
|
||||
|
||||
} // namespace cutlass::distributed::kernel
|
||||
|
||||
@@ -0,0 +1,324 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 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 1-D Distributed GEMM Schedules
|
||||
|
||||
NOTE: This API is __experimental__ and will change heavily over time. Particularly the use of
|
||||
CuTe layouts as integer functions in defining iteration-to-tile mappings is over-expressive and
|
||||
leaves plenty of room for incorrect/unexpected behavior.
|
||||
Please proceed with caution when modifying these schedules or defining new ones.
|
||||
|
||||
Device/iteration mappings are defined with CuTe layouts,
|
||||
since they are functions from integers to integers as well.
|
||||
|
||||
Each mapping is defined as a linear function of 2 variables (rank-2 layout):
|
||||
First variable (mode) is device index, second variable (mode) is iteration.
|
||||
A constant is also added to the final result as an offset value. This is a temporary workaround
|
||||
so that identity ownership mappings in the final iteration can be guaranteed for the schedules
|
||||
currently implemented.
|
||||
How are these mappings defined?
|
||||
Each schedule represents a unique parallel matrix multiplication algorithm, which describes how
|
||||
matrices/tensors are distributed among TP GPUs.
|
||||
|
||||
Depending on the algorithm, access patterns (GPU to tile or (GPU, iteration) to tile) mappings)
|
||||
are not necessarily going to be the identity function.
|
||||
|
||||
Pitfalls:
|
||||
The current representation uses CuTe layouts as arbitrary linear functions that map
|
||||
(GPU, iteration) to tile indices.
|
||||
This approach is over-expressive, and therefore makes a lot of assumptions on the part of the
|
||||
developer in how these mappings are defined. This can easily lead to incorrect implementations
|
||||
if not handled carefully.
|
||||
|
||||
|
||||
Assumption made in all schedules: TP == number of iterations (stages)
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cute/layout.hpp"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cutlass/cutlass.h"
|
||||
|
||||
#include "cutlass/experimental/distributed/schedules/dist_gemm_base_schedule.hpp"
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass::distributed::schedules {
|
||||
|
||||
// GEMM + Reduce Scatter
|
||||
// A and B are tiled along the K mode, which means each GPU gets an [M, K / TP]-shaped slice of A,
|
||||
// and an [N, K / TP] slice of B.
|
||||
// A is further tiled along the M mode, so that each stage/iteration computes a GEMM of shape
|
||||
// [M / TP, N, K / TP], and the epilogue will perform the reduction by reading its C tensor directly
|
||||
// from the left peer's previous D buffer.
|
||||
//
|
||||
// Below is an illustration of the tiling and iteration mappings for this pattern in the TP=4 case:
|
||||
//
|
||||
// Rows correspond to the M mode, columns correspond to the K mode for A and B and N mode for
|
||||
// C and D. Because sharding is done along K, each column of tiles is owned by one GPU.
|
||||
// Values in the grid correspond to the iteration/stage accessing the tile.
|
||||
// * means the same tile is accessed in all iterations/stages.
|
||||
//
|
||||
// Tensor A Tensor B
|
||||
//
|
||||
// GPU0 GPU1 GPU2 GPU3 GPU0 GPU1 GPU2 GPU3
|
||||
// |-----|-----|-----|-----| |-----|-----|-----|-----|
|
||||
// | | | | | | | | | |
|
||||
// | 3 | 0 | 1 | 2 | | | | | |
|
||||
// |_____|_____|_____|_____| | | | | |
|
||||
// | | | | | | | | | |
|
||||
// | 2 | 3 | 0 | 1 | | | | | |
|
||||
// |_____|_____|_____|_____| | * | * | * | * |
|
||||
// | | | | | | | | | |
|
||||
// | 1 | 2 | 3 | 0 | | | | | |
|
||||
// |_____|_____|_____|_____| | | | | |
|
||||
// | | | | | | | | | |
|
||||
// | 0 | 1 | 2 | 3 | | | | | |
|
||||
// |_____|_____|_____|_____| |_____|_____|_____|_____|
|
||||
//
|
||||
// M x K N x K
|
||||
//
|
||||
//
|
||||
// Tensor C Tensor D
|
||||
// (Peer's D)
|
||||
//
|
||||
//
|
||||
// |-----------------------| |-----------------------|
|
||||
// | | | |
|
||||
// GPU0 | 1,2,3 | GPU0 | * |
|
||||
// |_______________________| |_______________________|
|
||||
// | | | |
|
||||
// GPU1 | 1,2,3 | GPU1 | * |
|
||||
// |_______________________| |_______________________|
|
||||
// | | | |
|
||||
// GPU2 | 1,2,3 | GPU2 | * |
|
||||
// |_______________________| |_______________________|
|
||||
// | | | |
|
||||
// GPU3 | 1,2,3 | GPU3 | * |
|
||||
// |_______________________| |_______________________|
|
||||
//
|
||||
// M x N M x N
|
||||
//
|
||||
//
|
||||
// Tensor A's access pattern can be expressed as follows as a function of GPU index and iteration:
|
||||
// tile_idx = ((device_idx - 1) - iter + TP) % TP
|
||||
//
|
||||
// and can be expressed with the following CuTe layout:
|
||||
// (TP, TP) : (1, -1)
|
||||
// with ProcessorOffset = -1
|
||||
//
|
||||
//
|
||||
// Note: Since this schedule does not expose any communication, iteration 0 has no reduction step,
|
||||
// therefore epilogue is sourceless in iteration 0, and in the rest of the iterations the epilogue
|
||||
// source is a remote pointer to Tensor D owned by its left peer.
|
||||
//
|
||||
// Left peer is simply (device_idx - 1 + TP) % TP, which is expressed with the following CuTe layout:
|
||||
// (TP, TP) : (1, 0)
|
||||
//
|
||||
template <class TP_>
|
||||
struct ReduceScatter1D_TilingA_RotatingC: BaseSchedule<
|
||||
TP_,
|
||||
/* ProcessorTiler_ = */ cute::Shape<_1, _1, TP_, _1>,
|
||||
/* IterationTiler_ = */ cute::Shape<TP_, _1, _1, _1>,
|
||||
/* PeerDeviceMapping_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_1, _0>>, // (left neighbor) = (device_idx + ProcessorOffset + TP) % TP, with ProcessorOffset = -1
|
||||
/* IterationMappingM_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_1, _m1>>, // = (device_idx + ProcessorOffset - iter + TP) % TP, with ProcessorOffset = -1
|
||||
/* IterationMappingN_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>, // (IterationTiler::N == 1) = 0
|
||||
/* IterationMappingK_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>, // (IterationTiler::K == 1) = 0
|
||||
/* IterationMappingL_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>, // (IterationTiler::L == 1) = 0
|
||||
/* ProcessorOffset_ = */ _m1,
|
||||
/* MemcpyA_ = */ false,
|
||||
/* MemcpyB_ = */ false,
|
||||
/* KernelWritesArrivalFlag_ = */ true,
|
||||
/* NumBuffersA_ = */ 0,
|
||||
/* NumBuffersB_ = */ 0,
|
||||
/* NumBuffersC_ = */ 0,
|
||||
/* NumBuffersD_ = */ TP_{} - 1> {};
|
||||
|
||||
// This schedule is similar to ReduceScatter1D_TilingA_RotatingC, but with the second tiling
|
||||
// done along N instead of M. All other details remain unchanged.
|
||||
template <class TP_>
|
||||
struct ReduceScatter1D_TilingB_RotatingC: BaseSchedule<
|
||||
TP_,
|
||||
/* ProcessorTiler_ = */ cute::Shape<_1, _1, TP_, _1>,
|
||||
/* IterationTiler_ = */ cute::Shape<_1, TP_, _1, _1>,
|
||||
/* PeerDeviceMapping_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_1, _0>>, // (left neighbor) = (device_idx + ProcessorOffset + TP) % TP, with ProcessorOffset = -1
|
||||
/* IterationMappingM_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>, // (IterationTiler::N == 1) = 0
|
||||
/* IterationMappingN_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_1, _m1>>, // = (device_idx + ProcessorOffset - iter + TP) % TP, with ProcessorOffset = -1
|
||||
/* IterationMappingK_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>, // (IterationTiler::K == 1) = 0
|
||||
/* IterationMappingL_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>, // (IterationTiler::L == 1) = 0
|
||||
/* ProcessorOffset_ = */ _m1,
|
||||
/* MemcpyA_ = */ false,
|
||||
/* MemcpyB_ = */ false,
|
||||
/* KernelWritesArrivalFlag_ = */ true,
|
||||
/* NumBuffersA_ = */ 0,
|
||||
/* NumBuffersB_ = */ 0,
|
||||
/* NumBuffersC_ = */ 0,
|
||||
/* NumBuffersD_ = */ TP_{} - 1> {};
|
||||
|
||||
|
||||
// AllGather + GEMM
|
||||
// A and B are tiled along the N mode, which means each GPU allgathers A,
|
||||
// and operates with an [N / TP, K] slice of B.
|
||||
// For pipelining, A is further tiled along the M mode, so that each stage/iteration computes a
|
||||
// GEMM of shape [M / TP, N / TP, K], and concurrently we copy a peer's A slice into a local buffer
|
||||
// for the next stage/iteration.
|
||||
//
|
||||
// Below is an illustration of the tiling and iteration mappings for this pattern in the TP=4 case:
|
||||
//
|
||||
// Rows correspond to the M mode, columns correspond to the K mode for A and B and N mode for
|
||||
// C and D.
|
||||
//
|
||||
// Since this is a pipelined schedule without exposed communication, the first iteration starts
|
||||
// off immediately and operates on local slices of A and B. In the rest of the iterations, each
|
||||
// GPU accesses a slice of A copied from a peer GPU while it was busy with the last stage.
|
||||
//
|
||||
// Values in the following grids correspond to the peer buffer accessed by each GPU during
|
||||
// different iterations:
|
||||
//
|
||||
// Tensor A Tensor A
|
||||
// iter 0 iter 1
|
||||
//
|
||||
// |-----------------------| |-----------------------|
|
||||
// | | | |
|
||||
// GPU0 | 0 | | 1 |
|
||||
// |_______________________| |_______________________|
|
||||
// | | | |
|
||||
// GPU1 | 1 | | 2 |
|
||||
// |_______________________| |_______________________|
|
||||
// | | | |
|
||||
// GPU2 | 2 | | 3 |
|
||||
// |_______________________| |_______________________|
|
||||
// | | | |
|
||||
// GPU3 | 3 | | 0 |
|
||||
// |_______________________| |_______________________|
|
||||
//
|
||||
// M x K M x K
|
||||
//
|
||||
// Tensor A Tensor A
|
||||
// iter 2 iter 3
|
||||
//
|
||||
// |-----------------------| |-----------------------|
|
||||
// | | | |
|
||||
// GPU0 | 2 | | 3 |
|
||||
// |_______________________| |_______________________|
|
||||
// | | | |
|
||||
// GPU1 | 3 | | 0 |
|
||||
// |_______________________| |_______________________|
|
||||
// | | | |
|
||||
// GPU2 | 0 | | 1 |
|
||||
// |_______________________| |_______________________|
|
||||
// | | | |
|
||||
// GPU3 | 1 | | 2 |
|
||||
// |_______________________| |_______________________|
|
||||
//
|
||||
// M x K M x K
|
||||
//
|
||||
// Values in the following grids correspond to the tile accessed during each iteration.
|
||||
// * means the same tile is accessed in all iterations/stages.
|
||||
//
|
||||
// Tensor B Tensor C/D
|
||||
//
|
||||
//
|
||||
// |-----------------------| |-----|-----|-----|-----|
|
||||
// | | | | | | |
|
||||
// GPU0 | * | GPU0 | 0 | 1 | 2 | 3 |
|
||||
// |_______________________| |_____|_____|_____|_____|
|
||||
// | | | | | | |
|
||||
// GPU1 | * | GPU1 | 3 | 0 | 1 | 2 |
|
||||
// |_______________________| |_____|_____|_____|_____|
|
||||
// | | | | | | |
|
||||
// GPU2 | * | GPU2 | 2 | 3 | 0 | 1 |
|
||||
// |_______________________| |_____|_____|_____|_____|
|
||||
// | | | | | | |
|
||||
// GPU3 | * | GPU3 | 1 | 2 | 3 | 0 |
|
||||
// |_______________________| |_____|_____|_____|_____|
|
||||
//
|
||||
// N x K M x N
|
||||
//
|
||||
//
|
||||
// Tensor C/D's access pattern can be expressed as follows as a function of GPU index and iteration:
|
||||
// tile_idx = (device_idx + iter) % TP
|
||||
//
|
||||
// and can be expressed with the following CuTe layout:
|
||||
// (TP, TP) : (1, 1)
|
||||
//
|
||||
// This schedule does not need a ProcessorOffset constant.
|
||||
//
|
||||
// Peer devices from which A slices are copied is also expressed with the same function and CuTe
|
||||
// layout.
|
||||
//
|
||||
template <class TP_>
|
||||
struct AllGather1D_TilingCD_RotatingA: BaseSchedule<
|
||||
TP_,
|
||||
/* ProcessorTiler_ = */ cute::Shape<_1, TP_, _1, _1>,
|
||||
/* IterationTiler_ = */ cute::Shape<TP_, _1, _1, _1>,
|
||||
/* PeerDeviceMapping_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_1, _1>>, // = device_idx + iter
|
||||
/* IterationMappingM_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_1, _1>>, // = device_idx + iter
|
||||
/* IterationMappingN_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>, // (IterationTiler::N == 1) = 0
|
||||
/* IterationMappingK_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>, // (IterationTiler::K == 1) = 0
|
||||
/* IterationMappingL_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>, // (IterationTiler::L == 1) = 0
|
||||
/* ProcessorOffset_ = */ _0,
|
||||
/* MemcpyA_ = */ true,
|
||||
/* MemcpyB_ = */ false,
|
||||
/* KernelWritesArrivalFlag_ = */ false,
|
||||
/* NumBuffersA_ = */ TP_{} - 1,
|
||||
/* NumBuffersB_ = */ 0,
|
||||
/* NumBuffersC_ = */ 0,
|
||||
/* NumBuffersD_ = */ 0>{};
|
||||
|
||||
// This schedule is similar to AllGather1D_TilingCD_RotatingA, but with the order of tiling
|
||||
// swapped from N then M to M then N. This means slices of B are rotated around GPUs instead of
|
||||
// slices of A. All other details remain unchanged.
|
||||
template <class TP_>
|
||||
struct AllGather1D_TilingCD_RotatingB: BaseSchedule<
|
||||
TP_,
|
||||
/* ProcessorTiler_ = */ cute::Shape<TP_, _1, _1, _1>,
|
||||
/* IterationTiler_ = */ cute::Shape<_1, TP_, _1, _1>,
|
||||
/* PeerDeviceMapping_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_1, _1>>, // = device_idx + iter
|
||||
/* IterationMappingM_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>, // (IterationTiler::M == 1) = 0
|
||||
/* IterationMappingN_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_1, _1>>, // = device_idx + iter
|
||||
/* IterationMappingK_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>, // (IterationTiler::K == 1) = 0
|
||||
/* IterationMappingL_ = */ cute::Layout<cute::Shape<TP_, TP_>, cute::Stride<_0, _0>>, // (IterationTiler::L == 1) = 0
|
||||
/* ProcessorOffset_ = */ _0,
|
||||
/* MemcpyA_ = */ false,
|
||||
/* MemcpyB_ = */ true,
|
||||
/* KernelWritesArrivalFlag_ = */ false,
|
||||
/* NumBuffersA_ = */ 0,
|
||||
/* NumBuffersB_ = */ TP_{} - 1,
|
||||
/* NumBuffersC_ = */ 0,
|
||||
/* NumBuffersD_ = */ 0>{};
|
||||
|
||||
|
||||
} // namespace cutlass::distributed::schedules
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -0,0 +1,538 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 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 Base Schedule for Distributed GEMM
|
||||
|
||||
Templates Distributed GEMM schedules so that they can be expressed as a set of CuTe primitives and
|
||||
other static values.
|
||||
|
||||
NOTE: This API is __experimental__ and will change heavily over time. Particularly the use of
|
||||
CuTe layouts as integer functions in defining iteration-to-tile mappings is over-expressive and
|
||||
leaves plenty of room for incorrect/unexpected behavior.
|
||||
Please proceed with caution when modifying these schedules or defining new ones.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cute/layout.hpp"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cutlass/cutlass.h"
|
||||
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass::distributed::schedules {
|
||||
|
||||
/*
|
||||
* Distributed GEMM schedules define exactly how operand tensors are tiled and sliced across
|
||||
* processors (GPUs) and stages/iterations.
|
||||
*
|
||||
* BaseSchedule's role is to ease the implementation of arbitrary Distributed GEMM schedules
|
||||
* and reduce code repetition, simply by reducing the implementation to CuTe primitives and a few
|
||||
* other static values (buffer sizes, whether tensors are rotated using memcpies or not, and the
|
||||
* like.)
|
||||
*/
|
||||
template <
|
||||
class TP_, // CuTe constant defining the number of processors / GPUs / TP value
|
||||
class ProcessorTiler_, // CuTe tiler defining how fully materialized tensors are sharded across devices
|
||||
class IterationTiler_, // CuTe tiler defining how local tensors are tiled across stages/iterations
|
||||
class PeerDeviceMapping_, // CuTe layout mapping device index and stage/iteration to the device's peer index for that stage/iteration
|
||||
class IterationMappingM_, // CuTe layout mapping device index and stage/iteration to M tile index
|
||||
class IterationMappingN_, // CuTe layout mapping device index and stage/iteration to N tile index
|
||||
class IterationMappingK_, // CuTe layout mapping device index and stage/iteration to K tile index
|
||||
class IterationMappingL_, // CuTe layout mapping device index and stage/iteration to L tile index
|
||||
class ProcessorOffset_, // Constant offset for processor / GPU index in iteration mapping
|
||||
bool MemcpyA_, // Whether tensor A is memcpied
|
||||
bool MemcpyB_, // Whether tensor B is memcpied
|
||||
bool KernelWritesArrivalFlag_, // Whether the kernel writes arrival flags (when tensors are directly accessed from peer and not memcpied)
|
||||
int NumBuffersA_, // Number of buffers required for tensor A
|
||||
int NumBuffersB_, // Number of buffers required for tensor B
|
||||
int NumBuffersC_, // Number of buffers required for tensor C
|
||||
int NumBuffersD_> // Number of buffers required for tensor D
|
||||
struct BaseSchedule {
|
||||
|
||||
using TP = TP_;
|
||||
|
||||
static_assert(
|
||||
cute::is_static<TP>::value && cute::is_integral<TP>::value && cute::rank(TP{}) == 1 && cute::depth(TP{}) == 0,
|
||||
"Only integers allowed for TP at this time.");
|
||||
|
||||
static_assert(cute::rank(ProcessorTiler_{}) == 4, "Expected rank-4 processor tiler.");
|
||||
static_assert(cute::rank(IterationTiler_{}) == 4, "Expected rank-4 iteration tiler.");
|
||||
|
||||
static_assert(cute::rank(PeerDeviceMapping_{}) == 2,
|
||||
"PeerDeviceMapping must be rank-2 (device_idx, iter)");
|
||||
|
||||
static_assert(cute::rank(IterationMappingM_{}) == 2,
|
||||
"IterationMappingM must be rank-2 (device_idx, iter).");
|
||||
static_assert(cute::rank(IterationMappingN_{}) == 2,
|
||||
"IterationMappingN must be rank-2 (device_idx, iter).");
|
||||
static_assert(cute::rank(IterationMappingK_{}) == 2,
|
||||
"IterationMappingK must be rank-2 (device_idx, iter).");
|
||||
static_assert(cute::rank(IterationMappingL_{}) == 2,
|
||||
"IterationMappingL must be rank-2 (device_idx, iter).");
|
||||
|
||||
using ProcessorTiler = ProcessorTiler_;
|
||||
using IterationTiler = IterationTiler_;
|
||||
|
||||
using PeerDeviceMapping = PeerDeviceMapping_;
|
||||
using IterationMappingM = IterationMappingM_;
|
||||
using IterationMappingN = IterationMappingN_;
|
||||
using IterationMappingK = IterationMappingK_;
|
||||
using IterationMappingL = IterationMappingL_;
|
||||
|
||||
using ProcessorOffset = ProcessorOffset_;
|
||||
|
||||
static constexpr bool KernelWritesArrivalFlag = KernelWritesArrivalFlag_;
|
||||
static constexpr bool MemcpyA = MemcpyA_;
|
||||
static constexpr bool MemcpyB = MemcpyB_;
|
||||
static constexpr bool HasMemcpy = MemcpyA || MemcpyB;
|
||||
|
||||
static constexpr int NumBuffersA = NumBuffersA_;
|
||||
static constexpr int NumBuffersB = NumBuffersB_;
|
||||
static constexpr int NumBuffersC = NumBuffersC_;
|
||||
static constexpr int NumBuffersD = NumBuffersD_;
|
||||
|
||||
static_assert(
|
||||
NumBuffersA > 0 ^
|
||||
NumBuffersB > 0 ^
|
||||
NumBuffersC > 0 ^
|
||||
NumBuffersD > 0,
|
||||
"Only one of the ABCD tensors can be buffered!");
|
||||
|
||||
static constexpr bool BufferedOutput = NumBuffersC > 0 || NumBuffersD > 0;
|
||||
static constexpr bool RemoteC = NumBuffersC == 0 && NumBuffersD > 0;
|
||||
static constexpr bool RemoteD = NumBuffersD == 0 && NumBuffersC > 0;
|
||||
|
||||
static_assert(not RemoteD, "Remote D is not supported yet.");
|
||||
|
||||
// Host-side API: can_implement based on the GLOBAL problem shape
|
||||
template <typename ProblemShape>
|
||||
static bool
|
||||
can_implement_global(ProblemShape const& global_problem_shape) {
|
||||
auto [M, N, K, L] = append<4>(global_problem_shape, 1);
|
||||
|
||||
auto [ptileM, ptileN, ptileK, ptileL] = ProcessorTiler{};
|
||||
auto [itileM, itileN, itileK, itileL] = IterationTiler{};
|
||||
|
||||
auto tileM = ptileM * itileM;
|
||||
auto tileN = ptileN * itileN;
|
||||
auto tileK = ptileK * itileK;
|
||||
auto tileL = ptileL * itileL;
|
||||
|
||||
return M % tileM == 0 && N % tileN == 0 && K % tileK == 0 && L % tileL == 0;
|
||||
}
|
||||
|
||||
template <typename ProblemShape>
|
||||
CUTLASS_HOST_DEVICE
|
||||
static auto
|
||||
get_local_gemm_shape(ProblemShape const& global_problem_shape) {
|
||||
auto problem_shape_MNKL = append<4>(global_problem_shape, 1);
|
||||
|
||||
return shape_div(
|
||||
shape_div(
|
||||
problem_shape_MNKL,
|
||||
ProcessorTiler{}),
|
||||
IterationTiler{});
|
||||
}
|
||||
|
||||
// Host-side API: determine peers
|
||||
static auto
|
||||
get_peers_for_device(int device_idx) {
|
||||
auto left_peer_id = device_idx > 0 ? device_idx - 1 : TP{} - 1;
|
||||
auto right_peer_id = device_idx < TP{} - 1 ? device_idx + 1 : 0;
|
||||
|
||||
return cute::make_tuple(left_peer_id, right_peer_id);
|
||||
}
|
||||
|
||||
// Determines peer given device index and iteration
|
||||
static int
|
||||
get_remote_peer_id(int device_idx, int iteration) {
|
||||
auto device_iter_to_peer_idx = PeerDeviceMapping{};
|
||||
auto peer_idx = (
|
||||
device_iter_to_peer_idx(device_idx + ProcessorOffset{}, iteration) + TP{}
|
||||
) % TP{};
|
||||
return peer_idx;
|
||||
}
|
||||
|
||||
// Construct tilers and index mappers for sharding across processors
|
||||
template <typename Tensor>
|
||||
CUTLASS_HOST_DEVICE
|
||||
static auto
|
||||
get_processor_tiler_a(Tensor tensor) {
|
||||
if constexpr (NumBuffersA > 0) {
|
||||
return shape_div(tensor.shape(), select<0,2,3>(IterationTiler{}));
|
||||
} else {
|
||||
return shape_div(tensor.shape(), select<0,2,3>(ProcessorTiler{}));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Tensor>
|
||||
CUTLASS_HOST_DEVICE
|
||||
static auto
|
||||
get_processor_tiler_b(Tensor tensor) {
|
||||
if constexpr (NumBuffersB > 0) {
|
||||
return shape_div(tensor.shape(), select<1,2,3>(IterationTiler{}));
|
||||
} else {
|
||||
return shape_div(tensor.shape(), select<1,2,3>(ProcessorTiler{}));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Tensor>
|
||||
CUTLASS_HOST_DEVICE
|
||||
static auto
|
||||
get_processor_tiler_c(Tensor tensor) {
|
||||
if constexpr (BufferedOutput) {
|
||||
return shape_div(tensor.shape(), select<0,1,3>(IterationTiler{}));
|
||||
} else {
|
||||
return shape_div(tensor.shape(), select<0,1,3>(ProcessorTiler{}));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Tensor>
|
||||
CUTLASS_HOST_DEVICE
|
||||
static auto
|
||||
get_processor_tiler_d(Tensor tensor) {
|
||||
return get_processor_tiler_c(tensor);
|
||||
}
|
||||
|
||||
// Construct tilers and index mappers for tiling and iterating on device
|
||||
template <typename Tensor>
|
||||
CUTLASS_HOST_DEVICE
|
||||
static auto
|
||||
get_device_tiler_a(Tensor tensor) {
|
||||
static_assert(NumBuffersA == 0, "Buffered tensors don't have device tilers!");
|
||||
return shape_div(tensor.shape(), select<0,2,3>(IterationTiler{}));
|
||||
}
|
||||
|
||||
template <typename Tensor>
|
||||
CUTLASS_HOST_DEVICE
|
||||
static auto
|
||||
get_device_tiler_b(Tensor tensor) {
|
||||
static_assert(NumBuffersB == 0, "Buffered tensors don't have device tilers!");
|
||||
return shape_div(tensor.shape(), select<1,2,3>(IterationTiler{}));
|
||||
}
|
||||
|
||||
template <typename Tensor>
|
||||
CUTLASS_HOST_DEVICE
|
||||
static auto
|
||||
get_device_tiler_c(Tensor tensor) {
|
||||
static_assert(NumBuffersC == 0 && NumBuffersD == 0, "Buffered tensors don't have device tilers!");
|
||||
return shape_div(tensor.shape(), select<0,1,3>(IterationTiler{}));
|
||||
}
|
||||
|
||||
template <typename Tensor>
|
||||
CUTLASS_HOST_DEVICE
|
||||
static auto
|
||||
get_device_tiler_d(Tensor tensor) {
|
||||
static_assert(NumBuffersC == 0 && NumBuffersD == 0, "Buffered tensors don't have device tilers!");
|
||||
return shape_div(tensor.shape(), select<0,1,3>(IterationTiler{}));
|
||||
}
|
||||
|
||||
// Map device index and iteration to tile coordinate
|
||||
// Must be implemented by children for now.
|
||||
CUTLASS_HOST_DEVICE
|
||||
static auto
|
||||
get_device_tile_idx_a(int device_idx, int iteration) {
|
||||
auto mapping_m = IterationMappingM{};
|
||||
auto mapping_k = IterationMappingK{};
|
||||
auto mapping_l = IterationMappingL{};
|
||||
auto crd_m = (mapping_m(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
|
||||
auto crd_k = (mapping_k(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
|
||||
auto crd_l = (mapping_l(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
|
||||
return make_coord(crd_m, crd_k, crd_l);
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
static auto
|
||||
get_device_tile_idx_b(int device_idx, int iteration) {
|
||||
auto mapping_n = IterationMappingN{};
|
||||
auto mapping_k = IterationMappingK{};
|
||||
auto mapping_l = IterationMappingL{};
|
||||
auto crd_n = (mapping_n(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
|
||||
auto crd_k = (mapping_k(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
|
||||
auto crd_l = (mapping_l(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
|
||||
return make_coord(crd_n, crd_k, crd_l);
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
static auto
|
||||
get_device_tile_idx_c(int device_idx, int iteration) {
|
||||
auto mapping_m = IterationMappingM{};
|
||||
auto mapping_n = IterationMappingN{};
|
||||
auto mapping_l = IterationMappingL{};
|
||||
auto crd_m = (mapping_m(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
|
||||
auto crd_n = (mapping_n(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
|
||||
auto crd_l = (mapping_l(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
|
||||
return make_coord(crd_m, crd_n, crd_l);
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
static auto
|
||||
get_device_tile_idx_d(int device_idx, int iteration) {
|
||||
auto mapping_m = IterationMappingM{};
|
||||
auto mapping_n = IterationMappingN{};
|
||||
auto mapping_l = IterationMappingL{};
|
||||
auto crd_m = (mapping_m(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
|
||||
auto crd_n = (mapping_n(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
|
||||
auto crd_l = (mapping_l(device_idx + ProcessorOffset{}, iteration) + TP{}) % TP{};
|
||||
return make_coord(crd_m, crd_n, crd_l);
|
||||
}
|
||||
|
||||
// Device Partitioners: partition non-buffered processor-resident operands.
|
||||
// Processor-resident operands fall into two categories: buffered, and not buffered.
|
||||
// Those buffered aren't expected to be further partitioned, and those
|
||||
template <typename Tensor>
|
||||
static auto
|
||||
get_tensor_A(Tensor original_tensor, void * tensor_buffer_ptr, int device_idx, int iteration) {
|
||||
static_assert(rank(original_tensor) == 3);
|
||||
|
||||
using Element = typename Tensor::value_type;
|
||||
// Recreate tensor without constness. This is to ensure return types match.
|
||||
Element* ptr = const_cast<Element*>(original_tensor.data());
|
||||
auto shape = original_tensor.shape();
|
||||
auto layout = original_tensor.layout();
|
||||
auto tensor = make_tensor(ptr, layout);
|
||||
|
||||
if constexpr (NumBuffersA == 0) {
|
||||
auto tiler = get_device_tiler_a(tensor);
|
||||
auto idx = get_device_tile_idx_a(device_idx, iteration);
|
||||
return inner_partition(tensor, tiler, idx);
|
||||
} else {
|
||||
Element* ptr_buffer = reinterpret_cast<Element*>(tensor_buffer_ptr);
|
||||
if (iteration == 0) {
|
||||
return tensor;
|
||||
}
|
||||
ptr_buffer += size(shape) * (iteration - 1);
|
||||
|
||||
return make_tensor(ptr_buffer, layout);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Tensor>
|
||||
static auto
|
||||
get_tensor_B(Tensor original_tensor, void * tensor_buffer_ptr, int device_idx, int iteration) {
|
||||
static_assert(rank(original_tensor) == 3);
|
||||
|
||||
using Element = typename Tensor::value_type;
|
||||
// Recreate tensor without constness. This is to ensure return types match.
|
||||
Element * ptr = const_cast<Element *>(original_tensor.data());
|
||||
auto shape = original_tensor.shape();
|
||||
auto layout = original_tensor.layout();
|
||||
auto tensor = make_tensor(ptr, layout);
|
||||
|
||||
if constexpr (NumBuffersB == 0) {
|
||||
auto tiler = get_device_tiler_b(tensor);
|
||||
auto idx = get_device_tile_idx_b(device_idx, iteration);
|
||||
return inner_partition(tensor, tiler, idx);
|
||||
} else {
|
||||
Element * ptr_buffer = reinterpret_cast<Element *>(tensor_buffer_ptr);
|
||||
if (iteration == 0) {
|
||||
return tensor;
|
||||
}
|
||||
ptr_buffer += size(shape) * (iteration - 1);
|
||||
|
||||
return make_tensor(ptr_buffer, layout);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Tensor>
|
||||
static auto
|
||||
get_tensor_C(Tensor original_tensor, void * tensor_buffer_ptr, int device_idx, int iteration) {
|
||||
static_assert(rank(original_tensor) == 3);
|
||||
|
||||
using Element = typename Tensor::value_type;
|
||||
// Recreate tensor without constness. This is to ensure return types match.
|
||||
Element * ptr = const_cast<Element *>(original_tensor.data());
|
||||
auto shape = original_tensor.shape();
|
||||
auto layout = original_tensor.layout();
|
||||
auto tensor = make_tensor(ptr, layout);
|
||||
|
||||
if constexpr (not BufferedOutput) {
|
||||
auto tiler = get_device_tiler_c(tensor);
|
||||
auto idx = get_device_tile_idx_c(device_idx, iteration);
|
||||
return inner_partition(tensor, tiler, idx);
|
||||
} else {
|
||||
// implement Remote D
|
||||
static_assert(RemoteC, "");
|
||||
|
||||
Element * ptr_buffer = reinterpret_cast<Element *>(tensor_buffer_ptr);
|
||||
if (iteration == 0) {
|
||||
return tensor;
|
||||
}
|
||||
ptr_buffer += size(shape) * (iteration - 1);
|
||||
|
||||
return make_tensor(ptr_buffer, layout);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Tensor>
|
||||
static auto
|
||||
get_tensor_D(Tensor original_tensor, void * tensor_buffer_ptr, int device_idx, int iteration) {
|
||||
static_assert(rank(original_tensor) == 3);
|
||||
|
||||
using Element = typename Tensor::value_type;
|
||||
// Recreate tensor without constness. This is to ensure return types match.
|
||||
Element * ptr = const_cast<Element *>(original_tensor.data());
|
||||
auto shape = original_tensor.shape();
|
||||
auto layout = original_tensor.layout();
|
||||
auto tensor = make_tensor(ptr, layout);
|
||||
|
||||
if constexpr (not BufferedOutput) {
|
||||
auto tiler = get_device_tiler_d(tensor);
|
||||
auto idx = get_device_tile_idx_d(device_idx, iteration);
|
||||
return inner_partition(tensor, tiler, idx);
|
||||
} else {
|
||||
// implement Remote D
|
||||
static_assert(RemoteC, "");
|
||||
|
||||
Element * ptr_buffer = reinterpret_cast<Element *>(tensor_buffer_ptr);
|
||||
// last iteration is the local tensor, the rest are buffers
|
||||
if (iteration == TP{} - 1) {
|
||||
return tensor;
|
||||
}
|
||||
ptr_buffer += size(shape) * iteration; // note: iteration, not iteration - 1
|
||||
|
||||
return make_tensor(ptr_buffer, layout);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename ProblemShape>
|
||||
CUTLASS_HOST_DEVICE
|
||||
static auto
|
||||
get_local_a_shape(ProblemShape problem_shape) {
|
||||
auto problem_shape_MNKL = append<4>(problem_shape, 1);
|
||||
if constexpr (NumBuffersA == 0) {
|
||||
return shape_div(
|
||||
select<0,2,3>(problem_shape_MNKL),
|
||||
select<0,2,3>(ProcessorTiler{}));
|
||||
} else {
|
||||
return shape_div(
|
||||
shape_div(
|
||||
select<0,2,3>(problem_shape_MNKL),
|
||||
select<0,2,3>(ProcessorTiler{})),
|
||||
select<0,2,3>(IterationTiler{}));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename ProblemShape>
|
||||
CUTLASS_HOST_DEVICE
|
||||
static auto
|
||||
get_local_b_shape(ProblemShape problem_shape) {
|
||||
auto problem_shape_MNKL = append<4>(problem_shape, 1);
|
||||
if constexpr (NumBuffersB == 0) {
|
||||
return shape_div(
|
||||
select<1,2,3>(problem_shape_MNKL),
|
||||
select<1,2,3>(ProcessorTiler{}));
|
||||
} else {
|
||||
return shape_div(
|
||||
shape_div(
|
||||
select<1,2,3>(problem_shape_MNKL),
|
||||
select<1,2,3>(ProcessorTiler{})),
|
||||
select<1,2,3>(IterationTiler{}));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename ProblemShape>
|
||||
CUTLASS_HOST_DEVICE
|
||||
static auto
|
||||
get_local_c_shape(ProblemShape problem_shape) {
|
||||
auto problem_shape_MNKL = append<4>(problem_shape, 1);
|
||||
if constexpr (not BufferedOutput) {
|
||||
return shape_div(
|
||||
select<0,1,3>(problem_shape_MNKL),
|
||||
select<0,1,3>(ProcessorTiler{}));
|
||||
} else {
|
||||
return shape_div(
|
||||
shape_div(
|
||||
select<0,1,3>(problem_shape_MNKL),
|
||||
select<0,1,3>(ProcessorTiler{})),
|
||||
select<0,1,3>(IterationTiler{}));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename ProblemShape>
|
||||
CUTLASS_HOST_DEVICE
|
||||
static auto
|
||||
get_local_d_shape(ProblemShape problem_shape) {
|
||||
auto problem_shape_MNKL = append<4>(problem_shape, 1);
|
||||
if constexpr (not BufferedOutput) {
|
||||
return shape_div(
|
||||
select<0,1,3>(problem_shape_MNKL),
|
||||
select<0,1,3>(ProcessorTiler{}));
|
||||
} else {
|
||||
return shape_div(
|
||||
shape_div(
|
||||
select<0,1,3>(problem_shape_MNKL),
|
||||
select<0,1,3>(ProcessorTiler{})),
|
||||
select<0,1,3>(IterationTiler{}));
|
||||
}
|
||||
}
|
||||
|
||||
// Host-side APIs: get_device_slice_{A,B,C,D}
|
||||
// Slice off a view of the GLOBAL tensor that corresponds to the shard that
|
||||
// is going to be owned by a specific device. This helps with the initial
|
||||
// distribution of the GLOBAL operands among devices.
|
||||
template <typename Tensor>
|
||||
static auto
|
||||
get_device_slice_A(Tensor tensor, int device_idx) {
|
||||
auto tiler = get_processor_tiler_a(tensor);
|
||||
return inner_partition(tensor, tiler, device_idx);
|
||||
}
|
||||
|
||||
template <typename Tensor>
|
||||
static auto
|
||||
get_device_slice_B(Tensor tensor, int device_idx) {
|
||||
auto tiler = get_processor_tiler_b(tensor);
|
||||
return inner_partition(tensor, tiler, device_idx);
|
||||
}
|
||||
|
||||
template <typename Tensor>
|
||||
static auto
|
||||
get_device_slice_C(Tensor tensor, int device_idx) {
|
||||
auto tiler = get_processor_tiler_c(tensor);
|
||||
return inner_partition(tensor, tiler, device_idx);
|
||||
}
|
||||
|
||||
template <typename Tensor>
|
||||
static auto
|
||||
get_device_slice_D(Tensor tensor, int device_idx) {
|
||||
auto tiler = get_processor_tiler_d(tensor);
|
||||
return inner_partition(tensor, tiler, device_idx);
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
|
||||
} // namespace cutlass::gemm::distributed
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
Reference in New Issue
Block a user