github-actions[bot]
373 lines
17 KiB
Plaintext
373 lines
17 KiB
Plaintext
/**
|
|
* @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
|