269 lines
12 KiB
Plaintext
269 lines
12 KiB
Plaintext
// JIT build of pytorch/pytorch#178230 all_to_all_permute (Ulysses-style fused
|
|
// permute-free all-to-all over NCCL symmetric memory).
|
|
//
|
|
// Kernel is verbatim from the PR; the host entry point is adapted to the
|
|
// NCCLDevCommManager API shipped in torch 2.11 (we obtain the host ncclComm_t
|
|
// via the manager and create our own ncclDevComm with enough LSA barriers,
|
|
// cached per group — the PR's keyed-devcomm manager API is not in 2.11).
|
|
|
|
#include <c10/cuda/CUDAGuard.h>
|
|
#include <c10/cuda/CUDAException.h>
|
|
#include <c10/macros/Macros.h>
|
|
#include <ATen/native/cuda/MemoryAccess.cuh>
|
|
#include <ATen/cuda/CUDAContext.h>
|
|
#include <torch/csrc/distributed/c10d/NCCLUtils.hpp>
|
|
#include <torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp>
|
|
#include <torch/csrc/distributed/c10d/symm_mem/nccl_devcomm_manager.hpp>
|
|
#include <torch/csrc/distributed/c10d/symm_mem/NCCLSymmetricMemory.hpp>
|
|
#include <torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp>
|
|
#include <torch/library.h>
|
|
|
|
#include <map>
|
|
#include <mutex>
|
|
#include <string>
|
|
|
|
namespace {
|
|
|
|
using namespace c10d::symmetric_memory;
|
|
|
|
#ifndef NCCL_HAS_SYMMEM_DEVICE_SUPPORT
|
|
#error "Build requires NCCL >= 2.28 with symmetric-memory device support"
|
|
#endif
|
|
|
|
void copy_rows(const at::Tensor& input, at::Tensor& out) {
|
|
TORCH_CHECK(input.is_cuda() && out.is_cuda(), "a2a_permute: copy_rows requires CUDA tensors");
|
|
TORCH_CHECK(input.device() == out.device(), "a2a_permute: copy_rows device mismatch");
|
|
TORCH_CHECK(input.scalar_type() == out.scalar_type(), "a2a_permute: copy_rows dtype mismatch");
|
|
TORCH_CHECK(input.dim() == 3, "a2a_permute: copy_rows input must be 3-D");
|
|
TORCH_CHECK(out.sizes() == input.sizes(), "a2a_permute: copy_rows shape mismatch");
|
|
TORCH_CHECK(out.is_contiguous(), "a2a_permute: copy_rows output must be contiguous");
|
|
TORCH_CHECK(
|
|
input.stride(2) == 1 && input.stride(1) == input.size(2),
|
|
"a2a_permute: copy_rows requires contiguous rows");
|
|
|
|
const int64_t row_elements = input.size(1) * input.size(2);
|
|
TORCH_CHECK(input.stride(0) >= row_elements, "a2a_permute: copy_rows input rows overlap");
|
|
const size_t rows = static_cast<size_t>(input.size(0));
|
|
const size_t row_bytes = static_cast<size_t>(row_elements) * input.element_size();
|
|
const size_t input_pitch =
|
|
static_cast<size_t>(input.stride(0) * input.element_size());
|
|
|
|
c10::cuda::CUDAGuard guard(input.device());
|
|
auto stream = at::cuda::getCurrentCUDAStream();
|
|
C10_CUDA_CHECK(cudaMemcpy2DAsync(
|
|
out.data_ptr(),
|
|
row_bytes,
|
|
input.data_ptr(),
|
|
input_pitch,
|
|
row_bytes,
|
|
rows,
|
|
cudaMemcpyDeviceToDevice,
|
|
stream));
|
|
}
|
|
|
|
// ---- kernel (verbatim from PR) ----------------------------------------------
|
|
__device__ inline void copy_bytes_vec16_aligned(
|
|
const char* src, char* dst, size_t nbytes, size_t tid, size_t stride) {
|
|
const size_t n_vec = nbytes / 16;
|
|
constexpr int kUnroll = 4;
|
|
size_t vec_idx = tid;
|
|
for (; vec_idx + static_cast<size_t>(kUnroll - 1) * stride < n_vec;
|
|
vec_idx += static_cast<size_t>(kUnroll) * stride) {
|
|
at::native::memory::Vec<16> chunk[kUnroll];
|
|
#pragma unroll 4
|
|
for (int u = 0; u < kUnroll; ++u) {
|
|
const size_t i = vec_idx + static_cast<size_t>(u) * stride;
|
|
chunk[u] = at::native::memory::ld_vec<16>(src + i * 16);
|
|
}
|
|
#pragma unroll 4
|
|
for (int u = 0; u < kUnroll; ++u) {
|
|
const size_t i = vec_idx + static_cast<size_t>(u) * stride;
|
|
at::native::memory::st_vec<16>(dst + i * 16, chunk[u]);
|
|
}
|
|
}
|
|
for (; vec_idx < n_vec; vec_idx += stride) {
|
|
auto v = at::native::memory::ld_vec<16>(src + vec_idx * 16);
|
|
at::native::memory::st_vec<16>(dst + vec_idx * 16, v);
|
|
}
|
|
}
|
|
|
|
constexpr int A2A_MAX_SLOTS = 64;
|
|
constexpr int A2A_MAX_CTAS_PER_SLOT = 128;
|
|
constexpr int A2A_THREADS_PER_CTA = 128;
|
|
constexpr int A2A_MAX_CTA_COUNT = A2A_MAX_SLOTS * A2A_MAX_CTAS_PER_SLOT;
|
|
|
|
__global__ void all_to_all_lsa_kernel(
|
|
ncclWindow_t window,
|
|
size_t base_src_byte_offset,
|
|
unsigned char* out,
|
|
int num_rows,
|
|
size_t src_row_stride_bytes,
|
|
size_t copy_row_bytes,
|
|
size_t peer_stride_bytes,
|
|
size_t dst_row_stride_bytes,
|
|
ncclDevComm devComm) {
|
|
const int peer_idx = blockIdx.x;
|
|
const int local_block = blockIdx.y;
|
|
const ncclCoopCta coop{};
|
|
|
|
ncclLsaBarrierSession<ncclCoopCta> bar{
|
|
coop, devComm, ncclTeamLsa(devComm), devComm.lsaBarrier,
|
|
blockIdx.x * gridDim.y + blockIdx.y};
|
|
bar.sync(coop, cuda::memory_order_acquire);
|
|
|
|
unsigned char* dst_peer_base =
|
|
out + static_cast<size_t>(peer_idx) * peer_stride_bytes;
|
|
CUDA_KERNEL_ASSERT((base_src_byte_offset & 15) == 0);
|
|
CUDA_KERNEL_ASSERT((reinterpret_cast<uintptr_t>(dst_peer_base) & 15) == 0);
|
|
|
|
for (int i = local_block; i < num_rows; i += gridDim.y) {
|
|
const size_t row_byte_offset =
|
|
base_src_byte_offset + static_cast<size_t>(i) * src_row_stride_bytes;
|
|
const void* src_row = ncclGetLsaPointer(window, row_byte_offset, peer_idx);
|
|
unsigned char* dst_row =
|
|
dst_peer_base + static_cast<size_t>(i) * dst_row_stride_bytes;
|
|
copy_bytes_vec16_aligned(
|
|
reinterpret_cast<const char*>(src_row),
|
|
reinterpret_cast<char*>(dst_row),
|
|
copy_row_bytes, coop.thread_rank(), coop.size());
|
|
}
|
|
bar.sync(coop, cuda::memory_order_release);
|
|
}
|
|
|
|
// ---- own devcomm cache (replaces PR's keyed manager API) --------------------
|
|
ncclDevComm get_or_create_devcomm(ncclComm_t comm, const std::string& group_name) {
|
|
static std::mutex mu;
|
|
static std::map<std::string, ncclDevComm> cache;
|
|
std::lock_guard<std::mutex> lk(mu);
|
|
auto it = cache.find(group_name);
|
|
if (it != cache.end()) return it->second;
|
|
#ifdef NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER
|
|
ncclDevCommRequirements reqs = NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER;
|
|
#else
|
|
ncclDevCommRequirements reqs;
|
|
memset(&reqs, 0, sizeof(ncclDevCommRequirements));
|
|
#endif
|
|
reqs.lsaBarrierCount = A2A_MAX_CTA_COUNT;
|
|
ncclDevComm dc;
|
|
ncclResult_t res = ncclDevCommCreate(comm, &reqs, &dc);
|
|
TORCH_CHECK(res == ncclSuccess, "ncclDevCommCreate (a2a_permute) failed: ", ncclGetErrorString(res));
|
|
cache.emplace(group_name, dc);
|
|
return dc;
|
|
}
|
|
|
|
// ---- host entry -------------------------------------------------------------
|
|
void all_to_all_permute(
|
|
const at::Tensor& input,
|
|
at::Tensor& out,
|
|
int64_t scatter_dim,
|
|
int64_t gather_dim,
|
|
std::string group_name) {
|
|
TORCH_CHECK(input.stride(-1) == 1, "a2a_permute: innermost dim must be contiguous");
|
|
const bool col_scatter = (scatter_dim == 1 && gather_dim == 0);
|
|
const bool row_scatter = (scatter_dim == 0 && gather_dim == 1);
|
|
TORCH_CHECK(col_scatter || row_scatter,
|
|
"a2a_permute: unsupported (scatter_dim, gather_dim)=(", scatter_dim, ",", gather_dim, ")");
|
|
|
|
auto symm_mem = c10d::symmetric_memory::rendezvous(input, group_name);
|
|
TORCH_CHECK(symm_mem != nullptr, "a2a_permute: input must be NCCL symmetric memory");
|
|
auto* nccl_hdl = dynamic_cast<NCCLSymmetricMemory*>(symm_mem.get());
|
|
TORCH_CHECK(nccl_hdl != nullptr, "a2a_permute: requires NCCL symmetric memory backend");
|
|
|
|
c10::cuda::CUDAGuard guard(input.device());
|
|
auto stream = at::cuda::getCurrentCUDAStream();
|
|
auto device = input.device();
|
|
|
|
auto& manager = NCCLDevCommManager::get(device);
|
|
ncclComm_t comm = manager.get_comm(group_name);
|
|
ncclDevComm devcomm = get_or_create_devcomm(comm, group_name);
|
|
|
|
const int my_rank = nccl_hdl->get_rank();
|
|
const int p = nccl_hdl->get_world_size();
|
|
TORCH_CHECK(p <= A2A_MAX_SLOTS, "a2a_permute: group size ", p, " exceeds ", A2A_MAX_SLOTS);
|
|
TORCH_CHECK(out.is_contiguous(), "a2a_permute: out must be contiguous");
|
|
TORCH_CHECK(out.scalar_type() == input.scalar_type(), "a2a_permute: dtype mismatch");
|
|
|
|
auto window = nccl_hdl->get_window();
|
|
TORCH_CHECK(window != nullptr, "a2a_permute: NCCL window is null");
|
|
|
|
const size_t window_base_offset = nccl_hdl->get_offset();
|
|
const int64_t esize = input.element_size();
|
|
const size_t tensor_leading_offset =
|
|
window_base_offset + static_cast<size_t>(input.storage_offset()) * static_cast<size_t>(esize);
|
|
TORCH_CHECK(tensor_leading_offset % 16 == 0, "a2a_permute: tensor byte offset must be 16B aligned");
|
|
TORCH_CHECK(reinterpret_cast<uintptr_t>(input.data_ptr()) % 16 == 0, "a2a_permute: input ptr must be 16B aligned");
|
|
TORCH_CHECK(reinterpret_cast<uintptr_t>(out.data_ptr()) % 16 == 0, "a2a_permute: out ptr must be 16B aligned");
|
|
|
|
const int unroll = 4 * 16 / static_cast<int>(esize);
|
|
const int elems_per_cta = A2A_THREADS_PER_CTA * unroll;
|
|
int ctas_per_slot = 1;
|
|
|
|
if (col_scatter) {
|
|
const int rows = static_cast<int>(input.size(0));
|
|
int64_t total_cols = 0; int local_cols = 0;
|
|
if (input.dim() == 2) {
|
|
total_cols = input.size(1);
|
|
TORCH_CHECK(total_cols % p == 0, "a2a_permute: cols must divide group size");
|
|
local_cols = static_cast<int>(total_cols / p);
|
|
} else {
|
|
TORCH_CHECK(input.dim() == 3, "a2a_permute: input must be 2-D or 3-D");
|
|
TORCH_CHECK(input.size(1) == p, "a2a_permute: 3-D input dim1 must equal group size");
|
|
const int64_t lc = input.size(2);
|
|
total_cols = static_cast<int64_t>(p) * lc;
|
|
TORCH_CHECK(input.stride(1) == lc && input.stride(0) == total_cols,
|
|
"a2a_permute: 3-D input must be row-major contiguous");
|
|
local_cols = static_cast<int>(lc);
|
|
}
|
|
const bool ok3 = out.dim() == 3 && out.size(0) == p && out.size(1) == rows && out.size(2) == local_cols;
|
|
const bool ok2 = out.dim() == 2 && out.size(0) == (int64_t)p * rows && out.size(1) == local_cols;
|
|
TORCH_CHECK(ok3 || ok2, "a2a_permute: bad out shape for (1,0)");
|
|
const size_t row_bytes = (size_t)local_cols * (size_t)esize;
|
|
TORCH_CHECK(row_bytes % 16 == 0, "a2a_permute: local_cols*esize must be 16B-divisible");
|
|
ctas_per_slot = std::max(1, std::min((rows * local_cols + elems_per_cta - 1) / elems_per_cta, A2A_MAX_CTAS_PER_SLOT));
|
|
const size_t esz = (size_t)esize;
|
|
const size_t base_src = tensor_leading_offset + (size_t)my_rank * (size_t)local_cols * esz;
|
|
all_to_all_lsa_kernel<<<dim3(p, ctas_per_slot), A2A_THREADS_PER_CTA, 0, stream>>>(
|
|
window, base_src, reinterpret_cast<unsigned char*>(out.data_ptr()), rows,
|
|
(size_t)total_cols * esz, row_bytes, (size_t)rows * local_cols * esz, (size_t)local_cols * esz, devcomm);
|
|
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
|
} else {
|
|
int local_rows = 0; int cols = 0;
|
|
if (input.dim() == 2) {
|
|
const int64_t total_rows = input.size(0);
|
|
cols = static_cast<int>(input.size(1));
|
|
TORCH_CHECK(total_rows % p == 0, "a2a_permute: rows must divide group size");
|
|
local_rows = static_cast<int>(total_rows / p);
|
|
} else {
|
|
TORCH_CHECK(input.dim() == 3, "a2a_permute: input must be 2-D or 3-D");
|
|
TORCH_CHECK(input.size(0) == p, "a2a_permute: 3-D input dim0 must equal group size");
|
|
local_rows = static_cast<int>(input.size(1));
|
|
const int64_t c = input.size(2); cols = static_cast<int>(c);
|
|
const int64_t s01 = (int64_t)local_rows * c;
|
|
TORCH_CHECK(input.stride(1) == c && input.stride(0) == s01, "a2a_permute: 3-D input must be row-major contiguous");
|
|
}
|
|
const bool ok3 = out.dim() == 3 && out.size(0) == local_rows && out.size(1) == p && out.size(2) == cols;
|
|
const bool ok2 = out.dim() == 2 && out.size(0) == local_rows && out.size(1) == (int64_t)p * cols;
|
|
TORCH_CHECK(ok3 || ok2, "a2a_permute: bad out shape for (0,1)");
|
|
const size_t row_bytes = (size_t)cols * (size_t)esize;
|
|
TORCH_CHECK(row_bytes % 16 == 0, "a2a_permute: cols*esize must be 16B-divisible");
|
|
ctas_per_slot = std::max(1, std::min((local_rows * cols + elems_per_cta - 1) / elems_per_cta, A2A_MAX_CTAS_PER_SLOT));
|
|
const size_t esz = (size_t)esize; const size_t cu = (size_t)cols;
|
|
const size_t base_src = tensor_leading_offset + (size_t)(my_rank * local_rows) * cu * esz;
|
|
all_to_all_lsa_kernel<<<dim3(p, ctas_per_slot), A2A_THREADS_PER_CTA, 0, stream>>>(
|
|
window, base_src, reinterpret_cast<unsigned char*>(out.data_ptr()), local_rows,
|
|
cu * esz, row_bytes, cu * esz, (size_t)p * cu * esz, devcomm);
|
|
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
|
|
TORCH_LIBRARY_FRAGMENT(a2ap, m) {
|
|
m.def("copy_rows(Tensor input, Tensor(a!) out) -> ()");
|
|
m.def("all_to_all_permute(Tensor input, Tensor(a!) out, int scatter_dim, int gather_dim, str group_name) -> ()");
|
|
}
|
|
TORCH_LIBRARY_IMPL(a2ap, CUDA, m) {
|
|
m.impl("copy_rows", TORCH_FN(copy_rows));
|
|
m.impl("all_to_all_permute", TORCH_FN(all_to_all_permute));
|
|
}
|