indigo/LTX-2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Automated PR - 2026-07-07

Browse Source
This commit is contained in:
github-actions[bot]
2026-07-07 16:57:50 +00:00
parent 780984275f
commit 63fd9a4f86
157 changed files with 15976 additions and 5043 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ltx-kernels/csrc/all2all/all2all.cpp
+424
View File
@@ -0,0 +1,424 @@
/**
* @file all2all.cpp
* @brief Implementation of All2All communication primitives for multi-GPU tensor parallelism.
*
* This file implements the All2All class which provides efficient inter-GPU communication
* using CUDA IPC (Inter-Process Communication). The implementation supports:
* - Head redistribution for tensor-parallel attention (send_recv_heads, gather_heads)
* - Sequence gathering for cross-rank aggregation (allgather)
*
* All operations use a barrier-based synchronization protocol where each GPU writes
* directly to remote GPU memory via IPC, then signals completion through atomic
* operations on barrier counters.
*/
#include <ATen/cuda/CUDAContext.h>
#include <ATen/cuda/CUDADataType.h>
#include <c10/cuda/CUDAGuard.h>
#include <chrono>
#include <cuda_runtime.h>
#include <memory>
#include <pybind11/functional.h>
#include <torch/python.h>
#include "all2all.hpp"
#include "cuda/api.cuh"
#include "cuda/configs.cuh"
namespace ltx_kernels {
namespace all2all {
/**
* Constructs the All2All communication manager.
*
* Memory Allocation Strategy:
* The constructor allocates a single contiguous GPU memory block that contains:
* 1. Data buffer (tensor_bytes): Space for tensor data exchange
* 2. Barrier signals (MAX_NUM_PEERS * sizeof(int)): Per-rank completion counters
* 3. Buffer pointers (MAX_NUM_PEERS * sizeof(void*)): GPU-accessible pointer array
* 4. Barrier pointer array (MAX_NUM_PEERS * sizeof(int*)): GPU-accessible signal pointers
*
* This layout minimizes memory allocations and allows the entire region to be
* shared via a single IPC handle.
*/
All2All::All2All(int rank, int world_size, int num_tokens, int hidden_dim, int num_sms, at::ScalarType tensor_dtype,
double timeout_seconds)
: rank(rank), world_size(world_size), num_sms(num_sms), max_tokens(num_tokens), num_elems(0), tensor_bytes(0),
tensor_dtype(tensor_dtype) {
num_elems = int64_t(num_tokens) * int64_t(hidden_dim);
tensor_bytes = num_elems * elementSize(tensor_dtype);
// Derive the barrier timeout from the device's peak SM clock so the wall-clock guard is
// correct on any GPU (the kernel counts SM cycles via clock64). Use cudaDeviceGetAttribute,
// not cudaDeviceProp::clockRate, which was removed in CUDA 13. The attribute is in kHz.
int device = 0;
CUDA_CHECK(cudaGetDevice(&device));
int sm_clock_khz = 0;
CUDA_CHECK(cudaDeviceGetAttribute(&sm_clock_khz, cudaDevAttrClockRate, device));
sm_clock_hz_ = static_cast<double>(sm_clock_khz) * 1e3;
set_timeout_seconds(timeout_seconds);
// Calculate sizes for each region of the shared memory block
int64_t ptrs_bytes = MAX_NUM_PEERS * sizeof(void *);
int64_t barrier_signal_bytes = MAX_NUM_PEERS * sizeof(int);
int64_t barrier_signal_ptrs_bytes = MAX_NUM_PEERS * sizeof(int *);
// Allocate GPU memory for token count arrays (used by kernels)
CUDA_CHECK(cudaMalloc(reinterpret_cast<void **>(&rank_tokens_gpu), sizeof(int) * MAX_NUM_PEERS));
CUDA_CHECK(cudaMalloc(reinterpret_cast<void **>(&prefix_rank_tokens_gpu), sizeof(int) * MAX_NUM_PEERS));
// Allocate the main shared memory block and create IPC handle
// Layout: [data_buffer | barrier_signals | buffer_ptrs | barrier_signal_ptrs]
CUDA_CHECK(
cudaMalloc(&buffer_ptrs[rank], tensor_bytes + barrier_signal_bytes + ptrs_bytes + barrier_signal_ptrs_bytes));
CUDA_CHECK(cudaIpcGetMemHandle(&ipc_handlers[rank], buffer_ptrs[rank]));
// Set up pointers to each region within the allocated block
buffer_ptrs_gpu =
reinterpret_cast<void **>(static_cast<uint8_t *>(buffer_ptrs[rank]) + tensor_bytes + barrier_signal_bytes);
barrier_signal_ptrs[rank] = reinterpret_cast<int *>(static_cast<uint8_t *>(buffer_ptrs[rank]) + tensor_bytes);
barrier_signal_ptrs_gpu = reinterpret_cast<int **>(static_cast<uint8_t *>(buffer_ptrs[rank]) + tensor_bytes +
barrier_signal_bytes + ptrs_bytes);
// Initialize barrier signals to zero
CUDA_CHECK(cudaMemset(barrier_signal_ptrs[rank], 0, barrier_signal_bytes));
}
All2All::~All2All() noexcept(false) {
if (!destroyed) {
printf("WARNING: destroy() was not called, which can leak resources.\n");
fflush(stdout);
destroy();
}
}
/**
* Releases all allocated resources.
*
* This must be called explicitly before destruction to ensure proper cleanup of:
* - IPC memory mappings to remote GPUs
* - Local GPU memory allocations
*
* The method synchronizes the device to ensure all pending operations complete
* before releasing resources.
*/
void All2All::destroy() {
if (destroyed) {
return;
}
CUDA_CHECK(cudaDeviceSynchronize());
// Close IPC mappings to remote GPU memory (skip our own rank)
// Only close handles that were actually opened via sync()
for (int i = 0; i < world_size; i++) {
if (i != rank && buffer_ptrs[i] != nullptr) {
CUDA_CHECK(cudaIpcCloseMemHandle(buffer_ptrs[i]));
}
}
// Free local GPU memory allocations
CUDA_CHECK(cudaFree(buffer_ptrs[rank]));
CUDA_CHECK(cudaFree(rank_tokens_gpu));
CUDA_CHECK(cudaFree(prefix_rank_tokens_gpu));
destroyed = true;
}
/**
* Opens IPC memory mappings to all peer GPUs.
*
* This method processes IPC handles gathered from all ranks and opens memory
* mappings to enable direct GPU-to-GPU memory access. After calling this method,
* each GPU can read/write directly to any other GPU's buffer via buffer_ptrs.
*
* The barrier_signal_ptrs are also set up to point to the correct offset within
* each peer's shared memory block.
*/
void All2All::sync(const std::vector<std::optional<pybind11::bytearray>> &all_gathered_handles) {
for (int i = 0; i < world_size; i++) {
auto handle_str = std::string(all_gathered_handles[i].value());
EP_HOST_ASSERT(handle_str.size() == CUDA_IPC_HANDLE_SIZE);
if (i != rank) {
// Open IPC mapping to remote GPU's memory
std::memcpy(ipc_handlers[i].reserved, handle_str.c_str(), CUDA_IPC_HANDLE_SIZE);
CUDA_CHECK(cudaIpcOpenMemHandle(&buffer_ptrs[i], ipc_handlers[i], cudaIpcMemLazyEnablePeerAccess));
// Calculate offset to barrier signals in remote buffer
barrier_signal_ptrs[i] = reinterpret_cast<int *>(static_cast<uint8_t *>(buffer_ptrs[i]) + tensor_bytes);
} else {
// Verify our own handle matches what we sent
EP_HOST_ASSERT(std::memcmp(ipc_handlers[i].reserved, handle_str.c_str(), CUDA_IPC_HANDLE_SIZE) == 0);
}
}
// Copy pointer arrays to GPU for kernel access
CUDA_CHECK(cudaMemcpy(buffer_ptrs_gpu, buffer_ptrs, sizeof(void *) * world_size, cudaMemcpyHostToDevice));
CUDA_CHECK(
cudaMemcpy(barrier_signal_ptrs_gpu, barrier_signal_ptrs, sizeof(int *) * world_size, cudaMemcpyHostToDevice));
CUDA_CHECK(cudaDeviceSynchronize());
}
pybind11::bytearray All2All::get_local_ipc_handle() const {
return {ipc_handlers[rank].reserved, CUDA_IPC_HANDLE_SIZE};
}
/**
* Configures token distribution across ranks for the current batch.
*
* This method computes prefix sums needed by the kernels to calculate source
* and destination offsets. It must be called before any communication operation
* when the token distribution changes between batches.
*
* Example: For rank_num_tokens = {128, 96, 128, 64}
* - rank_tokens = {128, 96, 128, 64}
* - prefix_rank_tokens = {0, 128, 224, 352}
* - total_tokens = 416
*/
void All2All::set_rank_tokens(const std::vector<int> &rank_num_tokens) {
EP_HOST_ASSERT(static_cast<int>(rank_num_tokens.size()) == world_size);
// Initialize prefix sums to zero
for (int i = 0; i < world_size; i++) {
prefix_rank_tokens[i] = 0;
}
// Compute prefix sums (exclusive scan)
for (int i = 0; i < world_size; i++) {
rank_tokens[i] = rank_num_tokens[i];
if (i > 0) {
prefix_rank_tokens[i] = prefix_rank_tokens[i - 1] + rank_tokens[i - 1];
}
}
// Total tokens is the sum of all rank tokens
total_tokens = prefix_rank_tokens[world_size - 1] + rank_tokens[world_size - 1];
// Copy to GPU for kernel access
CUDA_CHECK(cudaMemcpy(rank_tokens_gpu, rank_tokens, sizeof(int) * MAX_NUM_PEERS, cudaMemcpyHostToDevice));
CUDA_CHECK(
cudaMemcpy(prefix_rank_tokens_gpu, prefix_rank_tokens, sizeof(int) * MAX_NUM_PEERS, cudaMemcpyHostToDevice));
CUDA_CHECK(cudaDeviceSynchronize());
}
/**
* Creates a tensor from the local IPC buffer.
*
* This helper method returns either a zero-copy view of the IPC buffer or
* a newly allocated tensor with the data copied. The zero-copy mode is more
* efficient but the tensor lifetime is tied to the All2All instance.
*
* @note The buffer pointer is cast to the template type T for proper interpretation.
*/
at::Tensor All2All::get_local_buffer_tensor(at::Tensor &x, int batch_size, int out_tokens, int out_heads, int head_size,
bool should_copy, cudaStream_t stream) {
auto ptr = buffer_ptrs[rank];
if (should_copy) {
// Allocate new tensor and copy data from IPC buffer
auto out_tensor = torch::empty({batch_size, out_tokens, out_heads, head_size}, x.options());
CUDA_CHECK(cudaMemcpyAsync(out_tensor.data_ptr(), ptr,
int64_t(batch_size) * int64_t(out_tokens) * int64_t(out_heads) * int64_t(head_size) *
int64_t(elementSize(x.scalar_type())),
cudaMemcpyDeviceToDevice, stream));
return out_tensor;
} else {
// Return a view directly into the IPC buffer (zero-copy)
auto out_tensor = torch::from_blob(ptr, {batch_size, out_tokens, out_heads, head_size}, x.options());
return out_tensor;
}
}
/**
* All2All communication to redistribute attention heads across GPUs.
*
* This operation is used in tensor-parallel transformers to exchange attention heads:
* - Before: Each GPU has all tokens but only a subset of heads
* - After: Each GPU has all tokens with heads redistributed
*
* Tensor Layout Transformation:
* Input: [batch, local_tokens, all_heads, head_size] per GPU
* Output: [batch, all_tokens, heads_per_rank, head_size] per GPU
*
* The operation partitions heads evenly: heads_per_rank = all_heads / world_size
* GPU i receives heads [i*heads_per_rank : (i+1)*heads_per_rank] from all GPUs.
*/
at::Tensor All2All::send_recv_heads(at::Tensor &x, bool copy_output) {
// Validate input tensor properties
EP_HOST_ASSERT(x.dim() == 4 and x.is_contiguous());
EP_HOST_ASSERT(x.dtype() == tensor_dtype);
EP_HOST_ASSERT(x.device().is_cuda());
EP_HOST_ASSERT(x.device().index() == rank);
int batch_size = x.size(0);
int num_tokens = x.size(1);
int num_heads = x.size(2);
int head_size = x.size(3);
// Output dimensions after redistribution
int out_tokens = total_tokens; // All tokens from all ranks
int out_heads = num_heads / world_size; // Each rank gets 1/world_size of heads
EP_HOST_ASSERT(int64_t(batch_size) * int64_t(out_tokens) * int64_t(out_heads) * int64_t(head_size) *
int64_t(elementSize(x.scalar_type())) <=
tensor_bytes);
at::cuda::CUDAGuard device_guard{x.device()};
auto stream = at::cuda::getCurrentCUDAStream().stream();
// Launch the All2All kernel
all2all_cuda::all2all_head_launch(buffer_ptrs_gpu, barrier_signal_ptrs_gpu, x.data_ptr(), prefix_rank_tokens_gpu,
rank, world_size, batch_size, total_tokens, num_tokens, num_heads, head_size,
stream, num_sms, tensor_dtype, timeout_cycles_);
return get_local_buffer_tensor(x, batch_size, out_tokens, out_heads, head_size, copy_output, stream);
}
/**
* Inverse All2All to gather heads back to original distribution.
*
* This is the inverse operation of send_recv_heads(). It redistributes data
* so each GPU gets back its original tokens with all attention heads.
*
* Tensor Layout Transformation:
* Input: [batch, all_tokens, heads_per_rank, head_size] per GPU
* Output: [batch, local_tokens, all_heads, head_size] per GPU
*
* Each GPU sends its portion of tokens to the originating rank, reconstructing
* the original head distribution.
*/
at::Tensor All2All::gather_heads(at::Tensor &x, bool copy_output) {
// Validate input tensor properties
EP_HOST_ASSERT(x.dim() == 4 and x.is_contiguous());
EP_HOST_ASSERT(x.dtype() == tensor_dtype);
EP_HOST_ASSERT(x.device().is_cuda());
EP_HOST_ASSERT(x.device().index() == rank);
at::cuda::CUDAGuard device_guard{x.device()};
auto stream = at::cuda::getCurrentCUDAStream().stream();
int batch_size = x.size(0);
int num_heads = x.size(2) * world_size; // Reconstruct total head count
int head_size = x.size(3);
// Output dimensions: this rank's tokens with all heads
int out_tokens = rank_tokens[rank];
int out_heads = num_heads;
EP_HOST_ASSERT(int64_t(batch_size) * int64_t(out_tokens) * int64_t(out_heads) * int64_t(head_size) *
int64_t(elementSize(x.scalar_type())) <=
tensor_bytes);
// Launch the gather kernel
all2all_cuda::all2all_head_gather_launch(buffer_ptrs_gpu, barrier_signal_ptrs_gpu, x.data_ptr(), rank_tokens_gpu,
prefix_rank_tokens_gpu, rank, world_size, batch_size, total_tokens,
num_heads, head_size, stream, num_sms, tensor_dtype, timeout_cycles_);
return get_local_buffer_tensor(x, batch_size, out_tokens, out_heads, head_size, copy_output, stream);
}
/**
* AllGather operation to collect sequence tokens from all ranks.
*
* Each GPU contributes its local sequence tokens, which are gathered into
* a complete sequence replicated on all GPUs. This is typically used after
* tensor-parallel operations to reconstruct the full sequence.
*
* Tensor Layout Transformation:
* Input: [batch, local_seqlen, heads, head_size] per GPU
* Output: [batch, total_seqlen, heads, head_size] per GPU (identical on all GPUs)
*
* Each GPU's tokens are placed at offset prefix_rank_tokens[rank] in the output.
*/
at::Tensor All2All::allgather(at::Tensor &x, bool copy_output) {
// Validate input tensor properties
EP_HOST_ASSERT(x.dim() == 4 and x.is_contiguous());
EP_HOST_ASSERT(x.dtype() == tensor_dtype);
EP_HOST_ASSERT(x.device().is_cuda());
EP_HOST_ASSERT(x.device().index() == rank);
at::cuda::CUDAGuard device_guard{x.device()};
auto stream = at::cuda::getCurrentCUDAStream().stream();
int batch_size = x.size(0);
int seqlen = x.size(1);
int num_heads = x.size(2);
int head_size = x.size(3);
// Output contains all tokens from all ranks
int out_tokens = total_tokens;
int out_heads = num_heads;
int hidden_dim = num_heads * head_size;
EP_HOST_ASSERT(int64_t(batch_size) * int64_t(out_tokens) * int64_t(out_heads) * int64_t(head_size) *
int64_t(elementSize(x.scalar_type())) <=
tensor_bytes);
// Launch the allgather kernel
all2all_cuda::allgather_launch(buffer_ptrs_gpu, barrier_signal_ptrs_gpu, x.data_ptr(), prefix_rank_tokens_gpu, rank,
world_size, batch_size, seqlen, hidden_dim, total_tokens, stream, num_sms,
tensor_dtype, timeout_cycles_);
return get_local_buffer_tensor(x, batch_size, out_tokens, out_heads, head_size, copy_output, stream);
}
} // namespace all2all
} // namespace ltx_kernels
/**
* Python bindings for the All2All communication library.
*
* Usage from Python:
* import all2all_cpp
*
* # Create instance (one per GPU)
* comm = all2all_cpp.All2All(rank, world_size, max_tokens, hidden_dim, num_sms, dtype)
*
* # Exchange IPC handles and synchronize
* handle = comm.get_local_ipc_handle()
* # ... gather handles via NCCL ...
* comm.sync(all_handles)
*
* # Set token distribution
* comm.set_rank_tokens([128, 128, 128, 128])
*
* # Perform operations
* output = comm.send_recv_heads(input_tensor, copy_output=False)
*
* # Cleanup
* comm.destroy()
*/
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.doc() = "High-performance All2All communication library for multi-GPU tensor parallelism.\n\n"
"This library provides IPC-based All2All operations optimized for transformer models.\n"
"Supported operations:\n"
" - send_recv_heads: Redistribute attention heads across GPUs\n"
" - gather_heads: Inverse of send_recv_heads\n"
" - allgather: Gather sequence tokens from all ranks\n";
pybind11::class_<ltx_kernels::all2all::All2All>(
m, "All2All",
"Manages All2All communication state for multi-GPU operations.\n\n"
"Args:\n"
" rank: This GPU's rank (0 to world_size-1)\n"
" world_size: Total number of GPUs\n"
" num_tokens: Maximum tokens per rank\n"
" hidden_dim: Hidden dimension (heads * head_size)\n"
" num_sms: Number of SMs for kernel launches\n"
" tensor_dtype: Tensor data type (torch.bfloat16 or torch.float8_e4m3fn)\n"
" timeout_seconds: Optional initial barrier timeout in seconds (defaults to the kernel default)")
.def(pybind11::init<int, int, int, int, int, at::ScalarType>())
.def(pybind11::init<int, int, int, int, int, at::ScalarType, double>())
.def("get_local_ipc_handle", &ltx_kernels::all2all::All2All::get_local_ipc_handle,
"Returns the IPC handle for this rank's buffer.")
.def("sync", &ltx_kernels::all2all::All2All::sync, "Opens IPC mappings to all peer GPUs using gathered handles.")
.def("destroy", &ltx_kernels::all2all::All2All::destroy,
"Releases all GPU resources. Must be called before destruction.")
.def("send_recv_heads", &ltx_kernels::all2all::All2All::send_recv_heads,
"All2All operation to redistribute attention heads.")
.def("gather_heads", &ltx_kernels::all2all::All2All::gather_heads,
"Inverse All2All to gather heads back to original distribution.")
.def("allgather", &ltx_kernels::all2all::All2All::allgather, "Gathers sequence tokens from all ranks.")
.def("set_rank_tokens", &ltx_kernels::all2all::All2All::set_rank_tokens,
"Sets token counts per rank for the current batch.")
.def("set_timeout_seconds", &ltx_kernels::all2all::All2All::set_timeout_seconds,
"Sets the barrier timeout in seconds (converted to cycles via the device peak SM clock).");
}
packages/ltx-kernels/csrc/all2all/all2all.hpp
+265
View File
@@ -0,0 +1,265 @@
/**
* @file all2all.hpp
* @brief High-performance All2All communication primitives for multi-GPU tensor parallelism.
*
* This library provides efficient All2All communication operations optimized for transformer
* models using tensor parallelism. It uses CUDA IPC (Inter-Process Communication) for
* zero-copy data transfer between GPUs in the same node.
*
* ## Architecture Overview
*
* The All2All class manages shared memory buffers accessible by all GPUs via IPC handles.
* Each GPU allocates a contiguous memory region containing:
* - Data buffer: Stores tensor data for exchange
* - Barrier signals: Synchronization counters for coordination
* - GPU pointer arrays: Device-accessible pointers to all peer buffers
*
* Memory Layout (per GPU):
* ```
* |<---- tensor_bytes ---->|<-- barrier signals -->|<-- buffer_ptrs_gpu -->|<-- barrier_signal_ptrs_gpu -->|
* | Data Buffer | MAX_PEERS * int | MAX_PEERS * void* | MAX_PEERS * int* |
* ```
*
* ## Supported Operations
*
* 1. **send_recv_heads**: Redistributes attention heads across GPUs (All2All)
* - Input: [batch, tokens, heads, head_size] on each GPU
* - Output: [batch, total_tokens, heads/world_size, head_size] on each GPU
*
* 2. **gather_heads**: Inverse of send_recv_heads
* - Gathers distributed heads back to original distribution
*
* 3. **allgather**: Gathers sequence data from all ranks
* - Each GPU contributes its local tokens to form the complete sequence
*
* ## Thread Safety
*
* - The class is NOT thread-safe. Each thread/process should have its own instance.
* - Multiple CUDA streams may use the same instance sequentially.
* - The `destroy()` method MUST be called before destruction to properly release IPC handles.
*
* ## Usage Example
*
* ```cpp
* // Initialize on each GPU
* auto comm = All2All(rank, world_size, max_tokens, hidden_dim, num_sms, dtype);
*
* // Exchange IPC handles (via NCCL or other collective)
* auto my_handle = comm.get_local_ipc_handle();
* // ... gather all handles ...
* comm.sync(all_handles);
*
* // Set token distribution for current batch
* comm.set_rank_tokens({128, 128, 128, 128}); // tokens per rank
*
* // Perform All2All on attention heads
* auto result = comm.send_recv_heads(input_tensor, copy_output=false);
*
* // Clean up
* comm.destroy();
* ```
*/
#pragma once
#include "cuda/configs.cuh"
#include "event.hpp"
#include <cmath>
#include <limits>
#include <pybind11/pybind11.h>
#include <pybind11/pytypes.h>
#include <stdexcept>
#include <torch/types.h>
#include <tuple>
#include <vector>
namespace ltx_kernels {
namespace all2all {
/**
* @class All2All
* @brief Manages All2All communication state and operations for multi-GPU tensor parallelism.
*
* This class encapsulates the IPC-based communication infrastructure needed for
* efficient All2All operations. It maintains shared memory buffers, barrier signals,
* and provides methods for head-parallel tensor redistribution.
*/
struct All2All {
private:
int rank; ///< This GPU's rank (0 to world_size-1)
int world_size; ///< Total number of GPUs in the communication group
int num_sms; ///< Number of SMs to use for kernel launches
int max_tokens; ///< Maximum number of tokens the buffer was allocated for
int64_t num_elems; ///< Number of elements in the data buffer (tokens * hidden_dim)
int64_t tensor_bytes; ///< Size of the data buffer in bytes
/// Host array of pointers to each rank's data buffer (GPU memory)
void *buffer_ptrs[MAX_NUM_PEERS] = {nullptr};
/// Device-accessible array of buffer pointers (copied to GPU)
void **buffer_ptrs_gpu = nullptr;
/// Host array of pointers to each rank's barrier signal buffer
int *barrier_signal_ptrs[MAX_NUM_PEERS] = {nullptr};
/// Device-accessible array of barrier signal pointers
int **barrier_signal_ptrs_gpu = nullptr;
/// IPC handles for sharing memory between processes
cudaIpcMemHandle_t ipc_handlers[MAX_NUM_PEERS];
at::ScalarType tensor_dtype; ///< Data type of tensors (BFloat16 or Float8_e4m3fn)
bool destroyed = false; ///< Flag to track if resources have been released
int total_tokens; ///< Sum of tokens across all ranks for current batch
int rank_tokens[MAX_NUM_PEERS]; ///< Number of tokens on each rank
int prefix_rank_tokens[MAX_NUM_PEERS]; ///< Cumulative sum of tokens (for offset calculation)
int *rank_tokens_gpu = nullptr; ///< Device copy of rank_tokens
int *prefix_rank_tokens_gpu = nullptr; ///< Device copy of prefix_rank_tokens
/// Device peak SM clock in Hz (from cudaDeviceGetAttribute(cudaDevAttrClockRate)), queried
/// once at construction. Used to convert a wall-clock timeout in seconds to barrier cycles.
double sm_clock_hz_ = 0.0;
/// All2All barrier timeout in GPU clock cycles. The constructor sets it from
/// DEFAULT_BARRIER_TIMEOUT_SECONDS and the queried SM clock; raise it (set_timeout_seconds)
/// to tolerate large cross-rank kernel-launch skew during the first torch.compile forward,
/// where one rank's recompile can delay its launch past the steady-state timeout.
uint64_t timeout_cycles_ = 0;
public:
/**
* @brief Constructs an All2All communication manager.
*
* Allocates GPU memory for the local data buffer, barrier signals, and pointer arrays.
* The IPC handle for the local buffer is created and can be retrieved via get_local_ipc_handle().
*
* @param rank This GPU's rank in the communication group (0-indexed)
* @param world_size Total number of GPUs/ranks
* @param num_tokens Maximum number of tokens this rank will handle
* @param hidden_dim Hidden dimension size (heads * head_size)
* @param num_sms Number of CUDA SMs to use for kernel execution
* @param tensor_dtype Data type for tensors (BFloat16 or Float8_e4m3fn)
* @param timeout_seconds Initial barrier timeout in seconds (see set_timeout_seconds); may be
* raised/reset at runtime for the first torch.compile forward
*/
All2All(int rank, int world_size, int num_tokens, int hidden_dim, int num_sms, at::ScalarType tensor_dtype,
double timeout_seconds = DEFAULT_BARRIER_TIMEOUT_SECONDS);
/**
* @brief Destructor - warns if destroy() was not called.
*
* @warning Always call destroy() explicitly before the destructor to properly
* release IPC handles. Failing to do so may leak resources.
*/
~All2All() noexcept(false);
/**
* @brief Synchronizes IPC handles from all ranks and opens remote memory mappings.
*
* This method must be called after all ranks have created their All2All instances
* and exchanged IPC handles via an external collective (e.g., NCCL allgather).
*
* @param all_gathered_handles Vector of IPC handles from all ranks (indexed by rank)
*/
void sync(const std::vector<std::optional<pybind11::bytearray>> &all_gathered_handles);
/**
* @brief Returns the IPC handle for this rank's shared buffer.
*
* The returned handle should be gathered across all ranks and passed to sync().
*
* @return pybind11::bytearray containing the CUDA IPC handle (CUDA_IPC_HANDLE_SIZE bytes)
*/
pybind11::bytearray get_local_ipc_handle() const;
/**
* @brief Creates a tensor view or copy of the local output buffer.
*
* @param x Reference tensor for options (dtype, device)
* @param batch_size Batch dimension size
* @param out_tokens Output token dimension size
* @param out_heads Output heads dimension size
* @param head_size Head dimension size
* @param should_copy If true, copies data to a new tensor; if false, returns a view
* @param stream CUDA stream for async copy
* @return Tensor with shape [batch_size, out_tokens, out_heads, head_size]
*/
at::Tensor get_local_buffer_tensor(at::Tensor &x, int batch_size, int out_tokens, int out_heads, int head_size,
bool should_copy, cudaStream_t stream);
/**
* @brief Releases all GPU resources and closes IPC handles.
*
* This method MUST be called before the object is destroyed. It synchronizes
* the device, closes remote IPC mappings, and frees local GPU memory.
*/
void destroy();
/**
* @brief Performs All2All communication to redistribute attention heads.
*
* Redistributes tensor from [batch, local_tokens, all_heads, head_size] to
* [batch, all_tokens, local_heads, head_size]. Each rank sends its portion
* of heads to the corresponding target rank.
*
* @param x Input tensor with shape [batch, num_tokens, num_heads, head_size]
* @param copy_output If true, returns a copy; if false, returns a view of the IPC buffer
* @return Tensor with shape [batch, total_tokens, num_heads/world_size, head_size]
*/
at::Tensor send_recv_heads(at::Tensor &x, bool copy_output);
/**
* @brief Performs inverse All2All to gather heads back to original distribution.
*
* Inverse of send_recv_heads(). Redistributes from [batch, all_tokens, local_heads, head_size]
* back to [batch, local_tokens, all_heads, head_size].
*
* @param x Input tensor with shape [batch, total_tokens, heads_per_rank, head_size]
* @param copy_output If true, returns a copy; if false, returns a view of the IPC buffer
* @return Tensor with shape [batch, rank_tokens[rank], num_heads, head_size]
*/
at::Tensor gather_heads(at::Tensor &x, bool copy_output);
/**
* @brief Gathers sequence tokens from all ranks.
*
* Each rank contributes its local sequence tokens, which are gathered into
* a complete sequence on all ranks.
*
* @param x Input tensor with shape [batch, seqlen, num_heads, head_size]
* @param copy_output If true, returns a copy; if false, returns a view of the IPC buffer
* @return Tensor with shape [batch, total_tokens, num_heads, head_size]
*/
at::Tensor allgather(at::Tensor &x, bool copy_output);
/**
* @brief Sets the token count for each rank in the current batch.
*
* Must be called before send_recv_heads(), gather_heads(), or allgather()
* to configure the token distribution. This allows variable-length sequences
* across ranks.
*
* @param rank_num_tokens Vector of token counts, one per rank (must have world_size elements)
*/
void set_rank_tokens(const std::vector<int> &rank_num_tokens);
/**
* @brief Sets the all2all barrier timeout in seconds.
*
* Converted to GPU clock cycles using the device's peak SM clock (queried at construction).
* Relaxes deadlock detection during the first torch.compile forward, where asymmetric
* per-rank recompilation can delay a rank's kernel launch beyond the steady-state timeout.
* Reset to the default for steady-state replay.
*/
void set_timeout_seconds(double seconds) {
if (!std::isfinite(seconds) || seconds < 0.0) {
throw std::invalid_argument("All2All timeout (seconds) must be finite and non-negative");
}
// Saturate rather than overflow the float->uint64 cast (out-of-range conversion is UB).
const double cycles = seconds * sm_clock_hz_;
const double max_cycles = static_cast<double>(std::numeric_limits<uint64_t>::max());
timeout_cycles_ = cycles >= max_cycles ? std::numeric_limits<uint64_t>::max() : static_cast<uint64_t>(cycles);
}
};
} // namespace all2all
} // namespace ltx_kernels
packages/ltx-kernels/csrc/all2all/cuda/all2all_heads.cu
+372
View File
@@ -0,0 +1,372 @@
/**
* @file all2all_heads.cu
* @brief CUDA kernels for All2All attention head redistribution.
*
* This file implements the GPU kernels for redistributing attention heads across
* multiple GPUs using IPC-based direct memory access. The kernels are designed
* for tensor-parallel transformer models where attention heads need to be
* exchanged between GPUs.
*
* ## Algorithm Overview
*
* The kernels use a direct-write approach where each GPU writes its data directly
* to the target GPU's memory buffer via IPC. This avoids intermediate copies and
* achieves near-peak memory bandwidth utilization.
*
* ## SM Work Distribution (Round-Robin)
*
* SMs are distributed round-robin among target ranks to handle non-divisible SM counts:
* - SM i writes to rank (i % world_size)
* - With 132 SMs and 8 GPUs: ranks 0-3 get 17 SMs, ranks 4-7 get 16 SMs
* - Each SM group processes all tokens for its assigned target rank
* - Within each group, SMs cooperate to cover all tokens in strided fashion
*
* ## Synchronization Protocol
*
* After data transfer, a barrier synchronization ensures all ranks have completed:
* 1. Each SM atomically increments the target rank's barrier counter for this rank
* 2. SM 0 waits until it has received signals from all ranks
* 3. Barrier counters are reset for the next operation
*/
#include "cuda/configs.cuh"
#include "cuda/exceptions.cuh"
#include "cuda/utils.cuh"
#include <ATen/cuda/CUDADataType.h>
namespace ltx_kernels {
namespace all2all {
namespace all2all_cuda {
/**
* @brief All2All kernel for redistributing attention heads across GPUs.
*
* This kernel performs the "send" phase of All2All: each GPU writes its assigned
* subset of attention heads to all other GPUs. The data layout transformation is:
*
* Source: [batch, num_tokens, num_heads, head_size]
* Dest: [batch, total_tokens, heads_per_rank, head_size]
*
* Each GPU writes heads [target_rank * heads_per_rank : (target_rank+1) * heads_per_rank]
* to target_rank's buffer at token offset prefix_rank_tokens[rank].
*
* ## Memory Layout
*
* Input tensor x (row-major, contiguous):
* - Batch dimension: outermost
* - Token dimension: batch_stride = num_tokens * num_heads * head_size
* - Head dimension: token_stride = num_heads * head_size
* - Head element: head_stride = head_size
*
* Output buffer (per target rank):
* - Similar layout but with heads_per_rank instead of num_heads
* - Tokens from this rank placed at offset prefix_rank_tokens[rank]
*
* ## Thread Block Organization
*
* Each thread block handles multiple tokens cooperatively:
* - Threads are organized in a 2D logical grid (rows=tokens, cols=elements)
* - Each thread copies 16 bytes (int4) per iteration
* - num_threads_per_token = (heads_per_rank * head_size) / elements_per_thread
* - num_tokens_per_copy = num_threads / num_threads_per_token
*
* @tparam ELEM_T Element type (at::BFloat16 or at::Float8_e4m3fn)
* @param buffer_ptrs Device array of pointers to each rank's data buffer
* @param barrier_signal_ptrs Device array of pointers to each rank's barrier signals
* @param x Source tensor data pointer
* @param rank This GPU's rank
* @param world_size Total number of GPUs
* @param batch_size Number of batches
* @param num_tokens Number of tokens on this rank
* @param num_heads Total number of attention heads
* @param head_size Size of each attention head
* @param total_tokens Sum of tokens across all ranks
* @param prefix_rank_tokens Cumulative token counts for offset calculation
*/
template <typename ELEM_T>
__global__ void send_recv_all2all(void **buffer_ptrs, int **barrier_signal_ptrs, void *x, int rank, int world_size,
int batch_size, int num_tokens, int num_heads, int head_size, int total_tokens,
int *prefix_rank_tokens, uint64_t timeout_cycles) {
// Grid dimensions
int num_sms = gridDim.x;
int sm_id = blockIdx.x;
int num_threads = blockDim.x;
// === SM Work Distribution (Round-Robin) ===
// Use modular assignment to handle num_sms not divisible by world_size.
// This ensures all SMs are utilized: some ranks get ceil(num_sms/world_size)
// SMs, others get floor(num_sms/world_size) SMs.
int64_t target_rank = get_target_rank(sm_id, world_size);
int64_t rank_local_sm_id = get_rank_local_sm_id(sm_id, world_size);
int64_t num_sms_for_this_rank = get_num_sms_for_rank(target_rank, num_sms, world_size);
// === Head Assignment ===
// Heads are partitioned evenly: rank i gets heads [i*hpr : (i+1)*hpr]
int64_t heads_per_rank = num_heads / world_size;
int64_t head_id = target_rank * heads_per_rank; // Starting head for target rank
// === Thread Mapping ===
// Each thread copies an int4 (16 bytes) per memory operation
// Threads form a 2D grid: (tokens_per_copy, threads_per_token)
int64_t num_elems_per_thread = sizeof(int4) / sizeof(ELEM_T);
int64_t num_threads_per_token = heads_per_rank * head_size / num_elems_per_thread;
int64_t num_tokens_per_copy = num_threads / num_threads_per_token;
// 2D thread coordinates within the logical grid
int64_t copy_thr_col_idx = threadIdx.x % num_threads_per_token; // Element offset
int64_t copy_thr_row_idx = threadIdx.x / num_threads_per_token; // Token offset
// Get target rank's buffer pointer
auto ptr = reinterpret_cast<void *>(static_cast<int8_t *>(buffer_ptrs[target_rank]));
// Use 64-bit arithmetic to avoid overflow for large tensors
int64_t num_tokens_64b = int64_t(num_tokens);
int64_t num_heads_64b = int64_t(num_heads);
int64_t head_size_64b = int64_t(head_size);
// === Main Copy Loop ===
// Iterate over batches and tokens, with SMs in the same group
// working on different token ranges in strided fashion
for (int64_t batch_ind = 0; batch_ind < batch_size; batch_ind++) {
// Strided token iteration: each SM in the group handles different token ranges
for (int64_t token_idx = rank_local_sm_id * num_tokens_per_copy; token_idx < num_tokens;
token_idx += num_tokens_per_copy * num_sms_for_this_rank) {
int64_t copy_token_idx = token_idx + copy_thr_row_idx;
// Destination token index accounts for this rank's offset in the global sequence
int64_t dst_token_idx = prefix_rank_tokens[rank] + copy_token_idx;
if (copy_token_idx >= num_tokens)
break;
// === Pointer Arithmetic ===
// Source: Read from this rank's input tensor at [batch, token, head_id:head_id+hpr, :]
// Note: We read a contiguous chunk of heads starting at head_id
int4 *shuffled_x_ptr =
reinterpret_cast<int4 *>(reinterpret_cast<uint8_t *>(x) +
batch_ind * num_tokens_64b * num_heads_64b * head_size_64b * sizeof(ELEM_T) +
copy_token_idx * num_heads_64b * head_size_64b * sizeof(ELEM_T) +
head_id * head_size_64b * sizeof(ELEM_T)) +
copy_thr_col_idx;
// Destination: Write to target rank's buffer at [batch, dst_token, :, :]
// The buffer has layout [batch, total_tokens, heads_per_rank, head_size]
int4 *shuffled_buffer_ptr =
reinterpret_cast<int4 *>(reinterpret_cast<uint8_t *>(ptr) +
batch_ind * total_tokens * heads_per_rank * head_size_64b * sizeof(ELEM_T) +
dst_token_idx * heads_per_rank * head_size_64b * sizeof(ELEM_T)) +
copy_thr_col_idx;
// Non-allocating store to avoid polluting L1 cache
st_na_global(shuffled_buffer_ptr, __ldg(shuffled_x_ptr));
}
}
// === Barrier Synchronization ===
// Signal completion to target rank and wait for all ranks to finish
barrier_wait_and_reset_roundrobin(barrier_signal_ptrs, target_rank, rank, world_size, num_sms, sm_id, threadIdx.x,
timeout_cycles);
}
/**
* @brief All2All kernel for gathering attention heads back to original distribution.
*
* This kernel performs the inverse of send_recv_all2all: it gathers heads from
* all ranks back to reconstruct the original tensor layout. Each GPU reads from
* its local buffer and writes its portion of heads to all target ranks.
*
* Data layout transformation:
* Source: [batch, total_tokens, heads_per_rank, head_size] (per GPU)
* Dest: [batch, rank_tokens[target], num_heads, head_size] (per target GPU)
*
* ## Memory Layout
*
* Input tensor x (this rank's portion after send_recv_all2all):
* - Contains all tokens but only heads_per_rank heads
* - Layout: [batch, total_tokens, heads_per_rank, head_size]
*
* Output buffer (per target rank):
* - Contains only that rank's tokens but all heads
* - Layout: [batch, rank_tokens[target], num_heads, head_size]
* - This rank writes heads [rank * heads_per_rank : (rank+1) * heads_per_rank]
*
* @tparam ELEM_T Element type (at::BFloat16 or at::Float8_e4m3fn)
* @param buffer_ptrs Device array of pointers to each rank's data buffer
* @param barrier_signal_ptrs Device array of pointers to barrier signals
* @param x Source tensor data (this rank's buffer after send_recv)
* @param rank This GPU's rank
* @param world_size Total number of GPUs
* @param batch_size Number of batches
* @param num_heads Total number of heads (reconstructed)
* @param head_size Size of each attention head
* @param rank_tokens Number of tokens for each rank
* @param total_tokens Sum of tokens across all ranks
* @param prefix_rank_tokens Cumulative token counts for offset calculation
*/
template <typename ELEM_T>
__global__ void gather_heads(void **buffer_ptrs, int **barrier_signal_ptrs, void *x, int rank, int world_size,
int batch_size, int num_heads, int head_size, const int *__restrict__ rank_tokens,
int total_tokens, int *prefix_rank_tokens, uint64_t timeout_cycles) {
// Grid dimensions
int num_sms = gridDim.x;
int sm_id = blockIdx.x;
int num_threads = blockDim.x;
// === SM Work Distribution (Round-Robin) ===
// Same partitioning as send_recv_all2all
int64_t target_rank = get_target_rank(sm_id, world_size);
int64_t rank_local_sm_id = get_rank_local_sm_id(sm_id, world_size);
int64_t num_sms_for_this_rank = get_num_sms_for_rank(target_rank, num_sms, world_size);
int64_t heads_per_rank = num_heads / world_size;
// === Thread Mapping ===
int64_t num_elems_per_thread = sizeof(int4) / sizeof(ELEM_T);
int64_t num_threads_per_token = heads_per_rank * head_size / num_elems_per_thread;
int64_t num_tokens_per_copy = num_threads / num_threads_per_token;
int64_t copy_thr_col_idx = threadIdx.x % num_threads_per_token;
int64_t copy_thr_row_idx = threadIdx.x / num_threads_per_token;
// Number of tokens owned by target rank
const int64_t tgt_tokens = int64_t(rank_tokens[target_rank]);
// This rank writes its heads at offset [rank * heads_per_rank] in the output
int64_t head_idx = rank * heads_per_rank;
int64_t num_heads_64b = int64_t(num_heads);
int64_t head_size_64b = int64_t(head_size);
int64_t total_tokens_64b = int64_t(total_tokens);
// Get target rank's buffer pointer
auto ptr = reinterpret_cast<void *>(static_cast<int8_t *>(buffer_ptrs[target_rank]));
// === Main Copy Loop ===
// Process target rank's tokens: read from global position, write to local position
for (int64_t batch_idx = 0; batch_idx < batch_size; batch_idx++) {
for (int64_t token_idx = rank_local_sm_id * num_tokens_per_copy; token_idx < tgt_tokens;
token_idx += num_tokens_per_copy * num_sms_for_this_rank) {
int64_t copy_token = token_idx + copy_thr_row_idx;
if (copy_token >= tgt_tokens)
break;
// Source: Read from global token position (target rank's tokens in our buffer)
int64_t src_token_idx = prefix_rank_tokens[target_rank] + copy_token;
// Destination: Write to local token position in target's buffer
int64_t dst_token_idx = copy_token;
// Source pointer: our input tensor at [batch, src_token, :, :]
int4 *shuffled_x_ptr =
reinterpret_cast<int4 *>(reinterpret_cast<uint8_t *>(x) +
batch_idx * total_tokens_64b * heads_per_rank * head_size_64b * sizeof(ELEM_T) +
src_token_idx * heads_per_rank * head_size_64b * sizeof(ELEM_T)) +
copy_thr_col_idx;
// Destination pointer: target's buffer at [batch, dst_token, head_idx:head_idx+hpr, :]
int4 *shuffled_buffer_ptr =
reinterpret_cast<int4 *>(reinterpret_cast<uint8_t *>(ptr) +
batch_idx * tgt_tokens * num_heads_64b * head_size_64b * sizeof(ELEM_T) +
dst_token_idx * num_heads_64b * head_size_64b * sizeof(ELEM_T) +
head_idx * head_size_64b * sizeof(ELEM_T)) +
copy_thr_col_idx;
st_na_global(shuffled_buffer_ptr, __ldg(shuffled_x_ptr));
}
}
// === Barrier Synchronization ===
barrier_wait_and_reset_roundrobin(barrier_signal_ptrs, target_rank, rank, world_size, num_sms, sm_id, threadIdx.x,
timeout_cycles);
}
/**
* @brief Host function to launch the gather_heads kernel.
*
* Selects the appropriate template instantiation based on tensor data type
* and launches the kernel with the specified number of SMs.
*
* @param buffer_ptrs Device array of buffer pointers
* @param barrier_signal_ptrs Device array of barrier signal pointers
* @param x Input tensor data pointer
* @param rank_tokens Token count per rank (device memory)
* @param prefix_rank_tokens Cumulative token counts (device memory)
* @param rank This GPU's rank
* @param world_size Total number of GPUs
* @param batch_size Number of batches
* @param total_tokens Sum of tokens across all ranks
* @param num_heads Total number of attention heads
* @param head_size Size of each attention head
* @param stream CUDA stream for async execution
* @param num_sms Number of SMs to launch
* @param tensor_dtype Data type (BFloat16 or Float8_e4m3fn)
*/
void all2all_head_gather_launch(void **buffer_ptrs, int **barrier_signal_ptrs, void *x, const int *rank_tokens,
int *prefix_rank_tokens, int rank, int world_size, int batch_size, int total_tokens,
int num_heads, int head_size, cudaStream_t stream, int num_sms,
at::ScalarType tensor_dtype, uint64_t timeout_cycles) {
do {
if (tensor_dtype == at::ScalarType::BFloat16) {
gather_heads<at::BFloat16><<<num_sms, DEFAULT_KERNEL_THREADS, 0, stream>>>(
buffer_ptrs, barrier_signal_ptrs, x, rank, world_size, batch_size, num_heads, head_size, rank_tokens,
total_tokens, prefix_rank_tokens, timeout_cycles);
} else if (tensor_dtype == at::ScalarType::Float8_e4m3fn) {
gather_heads<at::Float8_e4m3fn><<<num_sms, DEFAULT_KERNEL_THREADS, 0, stream>>>(
buffer_ptrs, barrier_signal_ptrs, x, rank, world_size, batch_size, num_heads, head_size, rank_tokens,
total_tokens, prefix_rank_tokens, timeout_cycles);
}
// Check for kernel launch errors
cudaError_t e = cudaGetLastError();
if (e != cudaSuccess) {
EPException cuda_exception("CUDA", __FILE__, __LINE__, cudaGetErrorString(e));
fprintf(stderr, "%s\n", cuda_exception.what());
throw cuda_exception;
}
} while (0);
}
/**
* @brief Host function to launch the send_recv_all2all kernel.
*
* Selects the appropriate template instantiation based on tensor data type
* and launches the kernel with the specified number of SMs.
*
* @param buffer_ptrs Device array of buffer pointers
* @param barrier_signal_ptrs Device array of barrier signal pointers
* @param x Input tensor data pointer
* @param prefix_rank_tokens Cumulative token counts (device memory)
* @param rank This GPU's rank
* @param world_size Total number of GPUs
* @param batch_size Number of batches
* @param total_tokens Sum of tokens across all ranks
* @param num_tokens Number of tokens on this rank
* @param num_heads Total number of attention heads
* @param head_size Size of each attention head
* @param stream CUDA stream for async execution
* @param num_sms Number of SMs to launch
* @param tensor_dtype Data type (BFloat16 or Float8_e4m3fn)
*/
void all2all_head_launch(void **buffer_ptrs, int **barrier_signal_ptrs, void *x, int *prefix_rank_tokens, int rank,
int world_size, int batch_size, int total_tokens, int num_tokens, int num_heads, int head_size,
cudaStream_t stream, int num_sms, at::ScalarType tensor_dtype, uint64_t timeout_cycles) {
do {
if (tensor_dtype == at::ScalarType::BFloat16) {
send_recv_all2all<at::BFloat16><<<num_sms, DEFAULT_KERNEL_THREADS, 0, stream>>>(
buffer_ptrs, barrier_signal_ptrs, x, rank, world_size, batch_size, num_tokens, num_heads, head_size,
total_tokens, prefix_rank_tokens, timeout_cycles);
} else if (tensor_dtype == at::ScalarType::Float8_e4m3fn) {
send_recv_all2all<at::Float8_e4m3fn><<<num_sms, DEFAULT_KERNEL_THREADS, 0, stream>>>(
buffer_ptrs, barrier_signal_ptrs, x, rank, world_size, batch_size, num_tokens, num_heads, head_size,
total_tokens, prefix_rank_tokens, timeout_cycles);
}
// Check for kernel launch errors
cudaError_t e = cudaGetLastError();
if (e != cudaSuccess) {
EPException cuda_exception("CUDA", __FILE__, __LINE__, cudaGetErrorString(e));
fprintf(stderr, "%s\n", cuda_exception.what());
throw cuda_exception;
}
} while (0);
}
} // namespace all2all_cuda
} // namespace all2all
} // namespace ltx_kernels
packages/ltx-kernels/csrc/all2all/cuda/allgather.cu
+198
View File
@@ -0,0 +1,198 @@
/**
* @file allgather.cu
* @brief CUDA kernel for AllGather operation using IPC-based direct memory access.
*
* This file implements the GPU kernel for gathering sequence tokens from all GPUs
* into a complete sequence on each GPU. Unlike the head redistribution kernels,
* this kernel preserves the head dimension and only gathers across the token
* (sequence) dimension.
*
* ## Algorithm Overview
*
* Each GPU broadcasts its local tokens to all other GPUs' buffers:
* - GPU i writes its tokens to position [prefix_rank_tokens[i]] in each buffer
* - After completion, all buffers contain the full sequence [0:total_tokens]
*
* ## Use Case
*
* This is typically used after tensor-parallel computation to reconstruct the
* full sequence for operations that require global context (e.g., output projection).
*/
#include "cuda/configs.cuh"
#include "cuda/exceptions.cuh"
#include "cuda/utils.cuh"
#include <ATen/cuda/CUDADataType.h>
namespace ltx_kernels {
namespace all2all {
namespace all2all_cuda {
/**
* @brief AllGather kernel to collect sequence tokens from all ranks.
*
* Each GPU writes its local sequence tokens to all other GPUs' buffers at the
* appropriate offset. After synchronization, all GPUs have the complete sequence.
*
* Data layout transformation:
* Input per GPU: [batch, seqlen, hidden_dim]
* Output per GPU: [batch, total_tokens, hidden_dim] (identical on all GPUs)
*
* ## Memory Layout
*
* Input tensor x (contiguous):
* - Shape: [batch, seqlen, hidden_dim]
* - hidden_dim = num_heads * head_size (flattened)
*
* Output buffer (per target rank, after gather):
* - Shape: [batch, total_tokens, hidden_dim]
* - This rank's tokens placed at offset rank_tokens_prefix[rank]
*
* ## Thread Mapping
*
* Similar to all2all_heads, threads cooperate to copy tokens:
* - Each thread copies 16 bytes (int4)
* - Threads per token = hidden_dim * sizeof(ELEM_T) / sizeof(int4)
* - Multiple tokens processed per thread block
*
* @tparam ELEM_T Element type (__nv_bfloat16 or at::Float8_e4m3fn)
* @param x Source tensor data pointer (this rank's tokens)
* @param buffer_ptrs Device array of pointers to each rank's data buffer
* @param barrier_signal_ptrs Device array of pointers to barrier signals
* @param batch_size Number of batches
* @param seqlen Number of tokens on this rank
* @param hidden_dim Hidden dimension size (num_heads * head_size)
* @param world_size Total number of GPUs
* @param rank This GPU's rank
* @param total_tokens Sum of tokens across all ranks
* @param rank_tokens_prefix Cumulative token counts (device memory)
*/
template <typename ELEM_T>
__global__ void allgather(void *x, void **buffer_ptrs, int **barrier_signal_ptrs, int batch_size, int seqlen,
int hidden_dim, int world_size, int rank, int total_tokens, int *rank_tokens_prefix,
uint64_t timeout_cycles) {
// Grid dimensions
int num_sms = gridDim.x;
int sm_id = blockIdx.x;
int num_threads = blockDim.x;
// === SM Work Distribution (Round-Robin) ===
// Use modular assignment to handle num_sms not divisible by world_size.
// This ensures all SMs are utilized: some ranks get ceil(num_sms/world_size)
// SMs, others get floor(num_sms/world_size) SMs.
int tgt_rank = get_target_rank(sm_id, world_size);
int rank_local_sm_id = get_rank_local_sm_id(sm_id, world_size);
int num_sms_for_this_rank = get_num_sms_for_rank(tgt_rank, num_sms, world_size);
// Get target rank's buffer pointer
auto ptr = reinterpret_cast<void *>(static_cast<int8_t *>(buffer_ptrs[tgt_rank]));
// === Thread Mapping ===
// Each thread copies one int4 (16 bytes)
int64_t num_elems_per_thread = sizeof(int4) / sizeof(ELEM_T);
int64_t num_threads_per_token = hidden_dim / num_elems_per_thread;
int64_t num_tokens_per_copy = num_threads / num_threads_per_token;
// 2D thread coordinates
int64_t copy_thr_col_idx = threadIdx.x % num_threads_per_token; // Element offset
int64_t copy_thr_row_idx = threadIdx.x / num_threads_per_token; // Token offset
// Use 64-bit arithmetic to avoid overflow
int64_t hidden_dim_64b = int64_t(hidden_dim);
int64_t total_tokens_64b = int64_t(total_tokens);
int64_t seqlen_64b = int64_t(seqlen);
// === Main Copy Loop ===
// Broadcast this rank's tokens to all target ranks' buffers
for (int64_t batch_idx = 0; batch_idx < batch_size; batch_idx++) {
// Strided token iteration within SM group for this target rank
for (int64_t token_idx = rank_local_sm_id * num_tokens_per_copy; token_idx < seqlen;
token_idx += num_tokens_per_copy * num_sms_for_this_rank) {
int64_t copy_token = token_idx + copy_thr_row_idx;
if (copy_token >= seqlen)
break;
// Source: local token index in input tensor
int64_t src_token_idx = copy_token;
// Destination: global token index in output buffer
// This rank's tokens start at prefix_rank_tokens[rank]
int64_t dst_token_idx = copy_token + rank_tokens_prefix[rank];
// Source pointer: input tensor at [batch, src_token, :]
int4 *shuffled_x_ptr = reinterpret_cast<int4 *>(reinterpret_cast<uint8_t *>(x) +
batch_idx * seqlen_64b * hidden_dim_64b * sizeof(ELEM_T) +
src_token_idx * hidden_dim_64b * sizeof(ELEM_T)) +
copy_thr_col_idx;
// Destination pointer: target buffer at [batch, dst_token, :]
int4 *shuffled_buffer_ptr =
reinterpret_cast<int4 *>(reinterpret_cast<uint8_t *>(ptr) +
batch_idx * total_tokens_64b * hidden_dim_64b * sizeof(ELEM_T) +
dst_token_idx * hidden_dim_64b * sizeof(ELEM_T)) +
copy_thr_col_idx;
// Non-allocating store for better cache behavior
st_na_global(shuffled_buffer_ptr, __ldg(shuffled_x_ptr));
}
}
// === Barrier Synchronization ===
// Signal completion to target rank and wait for all ranks
// Use round-robin variant since SM counts per rank may differ
barrier_wait_and_reset_roundrobin(barrier_signal_ptrs, tgt_rank, rank, world_size, num_sms, sm_id, threadIdx.x,
timeout_cycles);
}
/**
* @brief Host function to launch the allgather kernel.
*
* Launches the AllGather kernel with the specified configuration.
* Uses ALLGATHER_KERNEL_THREADS (1024) threads per block for higher
* occupancy than the All2All kernels.
*
* @param buffer_ptrs Device array of buffer pointers
* @param barrier_signal_ptrs Device array of barrier signal pointers
* @param x Input tensor data pointer
* @param prefix_rank_tokens Cumulative token counts (device memory)
* @param rank This GPU's rank
* @param world_size Total number of GPUs
* @param batch_size Number of batches
* @param seqlen Number of tokens on this rank
* @param hidden_dim Hidden dimension size
* @param total_tokens Sum of tokens across all ranks
* @param stream CUDA stream for async execution
* @param num_sms Number of SMs to launch
* @param tensor_dtype Data type (BFloat16 or Float8_e4m3fn)
*/
void allgather_launch(void **buffer_ptrs, int **barrier_signal_ptrs, void *x, int *prefix_rank_tokens, int rank,
int world_size, int batch_size, int seqlen, int hidden_dim, int total_tokens, cudaStream_t stream,
int num_sms, at::ScalarType tensor_dtype, uint64_t timeout_cycles) {
do {
if (tensor_dtype == at::ScalarType::BFloat16) {
allgather<at::BFloat16><<<num_sms, ALLGATHER_KERNEL_THREADS, 0, stream>>>(
x, buffer_ptrs, barrier_signal_ptrs, batch_size, seqlen, hidden_dim, world_size, rank, total_tokens,
prefix_rank_tokens, timeout_cycles);
} else if (tensor_dtype == at::ScalarType::Float8_e4m3fn) {
allgather<at::Float8_e4m3fn><<<num_sms, ALLGATHER_KERNEL_THREADS, 0, stream>>>(
x, buffer_ptrs, barrier_signal_ptrs, batch_size, seqlen, hidden_dim, world_size, rank, total_tokens,
prefix_rank_tokens, timeout_cycles);
} else {
EPException dtype_exception("allgather_launch", __FILE__, __LINE__, "Unsupported dtype");
fprintf(stderr, "%s\n", dtype_exception.what());
throw dtype_exception;
}
// Check for kernel launch errors
cudaError_t e = cudaGetLastError();
if (e != cudaSuccess) {
EPException cuda_exception("CUDA", __FILE__, __LINE__, cudaGetErrorString(e));
fprintf(stderr, "%s\n", cuda_exception.what());
throw cuda_exception;
}
} while (0);
}
} // namespace all2all_cuda
} // namespace all2all
} // namespace ltx_kernels
packages/ltx-kernels/csrc/all2all/cuda/api.cuh
+99
View File
@@ -0,0 +1,99 @@
/**
* @file api.cuh
* @brief CUDA kernel launch function declarations for All2All operations.
*
* This header provides the host-callable interface for launching the All2All
* CUDA kernels. These functions handle template instantiation and kernel
* configuration based on the tensor data type.
*/
#pragma once
#include <ATen/cuda/CUDADataType.h>
#include <vector>
namespace ltx_kernels {
namespace all2all {
namespace all2all_cuda {
/**
* @brief Launches the All2All head redistribution kernel.
*
* Redistributes attention heads across GPUs:
* Input: [batch, num_tokens, num_heads, head_size] per GPU
* Output: [batch, total_tokens, num_heads/world_size, head_size] per GPU
*
* @param buffer_ptrs Device array of pointers to each rank's data buffer
* @param barrier_signal_ptrs Device array of pointers to barrier signals
* @param x Source tensor data pointer
* @param prefix_rank_tokens Cumulative token counts per rank (device memory)
* @param rank This GPU's rank (0 to world_size-1)
* @param world_size Total number of GPUs
* @param batch_size Batch dimension size
* @param total_tokens Sum of tokens across all ranks
* @param num_tokens Number of tokens on this rank
* @param num_heads Total number of attention heads
* @param head_size Size of each attention head
* @param stream CUDA stream for async execution
* @param num_sms Number of SMs to use for the kernel
* @param tensor_dtype Data type (BFloat16 or Float8_e4m3fn)
*/
void all2all_head_launch(void **buffer_ptrs, int **barrier_signal_ptrs, void *x, int *prefix_rank_tokens, int rank,
int world_size, int batch_size, int total_tokens, int num_tokens, int num_heads, int head_size,
cudaStream_t stream, int num_sms, at::ScalarType tensor_dtype, uint64_t timeout_cycles);
/**
* @brief Launches the gather heads kernel (inverse of all2all_head_launch).
*
* Redistributes tokens back to original head distribution:
* Input: [batch, total_tokens, heads_per_rank, head_size] per GPU
* Output: [batch, rank_tokens[rank], num_heads, head_size] per GPU
*
* @param buffer_ptrs Device array of pointers to each rank's data buffer
* @param barrier_signal_ptrs Device array of pointers to barrier signals
* @param x Source tensor data pointer
* @param rank_tokens Token count for each rank (device memory)
* @param prefix_rank_tokens Cumulative token counts (device memory)
* @param rank This GPU's rank
* @param world_size Total number of GPUs
* @param batch_size Batch dimension size
* @param total_tokens Sum of tokens across all ranks
* @param num_heads Total number of attention heads (reconstructed)
* @param head_size Size of each attention head
* @param stream CUDA stream for async execution
* @param num_sms Number of SMs to use for the kernel
* @param tensor_dtype Data type (BFloat16 or Float8_e4m3fn)
*/
void all2all_head_gather_launch(void **buffer_ptrs, int **barrier_signal_ptrs, void *x, const int *rank_tokens,
int *prefix_rank_tokens, int rank, int world_size, int batch_size, int total_tokens,
int num_heads, int head_size, cudaStream_t stream, int num_sms,
at::ScalarType tensor_dtype, uint64_t timeout_cycles);
/**
* @brief Launches the AllGather kernel for sequence tokens.
*
* Gathers sequence tokens from all ranks:
* Input: [batch, seqlen, hidden_dim] per GPU
* Output: [batch, total_tokens, hidden_dim] per GPU (identical on all)
*
* @param buffer_ptrs Device array of pointers to each rank's data buffer
* @param barrier_signal_ptrs Device array of pointers to barrier signals
* @param x Source tensor data pointer
* @param prefix_rank_tokens Cumulative token counts (device memory)
* @param rank This GPU's rank
* @param world_size Total number of GPUs
* @param batch_size Batch dimension size
* @param seqlen Number of tokens on this rank
* @param hidden_dim Hidden dimension size (num_heads * head_size)
* @param total_tokens Sum of tokens across all ranks
* @param stream CUDA stream for async execution
* @param num_sms Number of SMs to use for the kernel
* @param tensor_dtype Data type (BFloat16 or Float8_e4m3fn)
*/
void allgather_launch(void **buffer_ptrs, int **barrier_signal_ptrs, void *x, int *prefix_rank_tokens, int rank,
int world_size, int batch_size, int seqlen, int hidden_dim, int total_tokens, cudaStream_t stream,
int num_sms, at::ScalarType tensor_dtype, uint64_t timeout_cycles);
} // namespace all2all_cuda
} // namespace all2all
} // namespace ltx_kernels