diff --git a/benchmarks/bench_dcp_alltoall.py b/benchmarks/bench_dcp_alltoall.py new file mode 100644 index 0000000000..0a00df06c4 --- /dev/null +++ b/benchmarks/bench_dcp_alltoall.py @@ -0,0 +1,336 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +DCP All-to-All Microbenchmark: Native LL128 FIFO vs NCCL Baseline + +Measures single kernel-level latency for the DCP A2A communication op: + - Native: decode_cp_a2a_alltoall (fused LL128 FIFO kernel via MNNVL) + - NCCL baseline: 2x torch.distributed.all_to_all_single (partial_o + softmax_stats) + +This is NOT an end-to-end pipeline benchmark. It measures raw communication +kernel time only. + +Launch: + mpirun --allow-run-as-root --oversubscribe -np 4 \ + python benchmarks/bench_dcp_alltoall.py + + mpirun --allow-run-as-root --oversubscribe -np 2 \ + python benchmarks/bench_dcp_alltoall.py --batch_sizes 1 16 64 + +Options: + --batch_sizes : Batch sizes to benchmark (default: 1 16 64 128) + --head_dim : Head dimension D (default: 128) + --stats_dim : Stats dimension S (default: 2) + --warmup : Warmup iterations (default: 50) + --iters : Timed iterations (default: 200) + --skip_nccl : Skip NCCL baseline (only run native) + --skip_native : Skip native (only run NCCL baseline) + +Requires: + - SM90+ GPU (Hopper/Blackwell) + - MNNVL support (multi-GPU fabric memory) + - mpi4py +""" + +import argparse +import os +import socket + +import numpy as np +import pynvml +import torch +import torch.distributed as dist +from mpi4py import MPI + +from flashinfer.comm import ( + decode_cp_a2a_alltoall, + decode_cp_a2a_init_workspace, + decode_cp_a2a_workspace_size, +) +from flashinfer.comm.mapping import Mapping +from flashinfer.comm.mnnvl import MnnvlMemory, MpiComm + + +def _to_torch(t): + """Convert a tvm_ffi.core.Tensor (or any DLPack object) to torch.Tensor.""" + if isinstance(t, torch.Tensor): + return t + return torch.from_dlpack(t) + + +def setup_mpi(): + """Initialize MPI and set CUDA device. Returns (rank, world_size, comm, local_rank).""" + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + world_size = comm.Get_size() + + # Compute local rank from hostname + hostname = socket.gethostname() + all_hostnames = comm.allgather(hostname) + local_rank = sum(1 for i in range(rank) if all_hostnames[i] == hostname) + torch.cuda.set_device(local_rank) + + return rank, world_size, comm, local_rank + + +def setup_nccl(rank, world_size): + """Initialize torch.distributed NCCL backend.""" + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "29500") + dist.init_process_group(backend="nccl", rank=rank, world_size=world_size) + + +def allocate_mnnvl_workspace(rank, cp_size, mpi_comm): + """Allocate MNNVL workspace for native DCP A2A.""" + pynvml.nvmlInit() + MnnvlMemory.initialize() + MnnvlMemory.comm = MpiComm() + + mapping = Mapping( + world_size=cp_size, + rank=rank, + cp_size=cp_size, + tp_size=1, + pp_size=1, + ) + + ws_bytes = decode_cp_a2a_workspace_size(cp_size) + mnnvl_mem = MnnvlMemory(mapping, ws_bytes) + workspace = mnnvl_mem.as_torch_strided_tensor(torch.int64) + workspace._mnnvl_mem = mnnvl_mem # prevent GC + return workspace + + +def bench_native( + workspace, + rank, + cp_size, + batch_size, + head_dim, + stats_dim, + dtype, + warmup, + iters, + mpi_comm, +): + """Benchmark native decode_cp_a2a_alltoall. Returns list of per-iteration times in ms.""" + # Init workspace once — FIFO supports reuse across iterations + decode_cp_a2a_init_workspace(workspace, rank, cp_size) + torch.cuda.synchronize() + mpi_comm.Barrier() + + partial_o = torch.randn(batch_size, cp_size, head_dim, dtype=dtype, device="cuda") + softmax_stats = torch.randn( + batch_size, cp_size, stats_dim, dtype=torch.float32, device="cuda" + ) + + # Warmup + for _ in range(warmup): + recv_o, recv_s = decode_cp_a2a_alltoall( + partial_o, softmax_stats, workspace, rank, cp_size + ) + torch.cuda.synchronize() + mpi_comm.Barrier() + + # Timed iterations + times = [] + for _ in range(iters): + torch.cuda.synchronize() + mpi_comm.Barrier() + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + + start.record() + recv_o, recv_s = decode_cp_a2a_alltoall( + partial_o, softmax_stats, workspace, rank, cp_size + ) + end.record() + torch.cuda.synchronize() + + times.append(start.elapsed_time(end)) + mpi_comm.Barrier() + + return times + + +def bench_nccl( + rank, cp_size, batch_size, head_dim, stats_dim, dtype, warmup, iters, mpi_comm +): + """Benchmark NCCL 2x all_to_all_single. Returns list of per-iteration times in ms.""" + group = dist.group.WORLD + + # partial_o: [B, cp_size, D] — each rank sends chunk [B, 1, D] to each peer + partial_o = torch.randn(batch_size, cp_size, head_dim, dtype=dtype, device="cuda") + softmax_stats = torch.randn( + batch_size, cp_size, stats_dim, dtype=torch.float32, device="cuda" + ) + + # Flatten for all_to_all_single: [B * cp_size * D] split into cp_size equal chunks + send_o = partial_o.contiguous().view(-1) + recv_o = torch.empty_like(send_o) + send_s = softmax_stats.contiguous().view(-1) + recv_s = torch.empty_like(send_s) + + # Warmup + for _ in range(warmup): + dist.all_to_all_single(recv_o, send_o, group=group) + dist.all_to_all_single(recv_s, send_s, group=group) + torch.cuda.synchronize() + mpi_comm.Barrier() + + # Timed iterations + times = [] + for _ in range(iters): + torch.cuda.synchronize() + mpi_comm.Barrier() + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + + start.record() + dist.all_to_all_single(recv_o, send_o, group=group) + dist.all_to_all_single(recv_s, send_s, group=group) + end.record() + torch.cuda.synchronize() + + times.append(start.elapsed_time(end)) + mpi_comm.Barrier() + + return times + + +def compute_stats(all_times, iters): + """Compute p50/p95/mean from gathered per-rank times. + + Takes max across ranks per iteration (communication is synchronous), + then reports percentiles. + """ + per_iter_max = [ + max(all_times[r][i] for r in range(len(all_times))) for i in range(iters) + ] + return { + "p50": float(np.percentile(per_iter_max, 50)), + "p95": float(np.percentile(per_iter_max, 95)), + "mean": float(np.mean(per_iter_max)), + } + + +def main(): + parser = argparse.ArgumentParser(description="DCP A2A Microbenchmark") + parser.add_argument("--batch_sizes", type=int, nargs="+", default=[1, 16, 64, 128]) + parser.add_argument("--head_dim", type=int, default=128) + parser.add_argument("--stats_dim", type=int, default=2) + parser.add_argument("--warmup", type=int, default=50) + parser.add_argument("--iters", type=int, default=200) + parser.add_argument("--skip_nccl", action="store_true") + parser.add_argument("--skip_native", action="store_true") + args = parser.parse_args() + + rank, world_size, mpi_comm, local_rank = setup_mpi() + cp_size = world_size + dtype = torch.bfloat16 + + if rank == 0: + print(f"=== DCP A2A Benchmark (cp_size={cp_size}, {cp_size} GPUs) ===") + print() + + # Initialize NCCL if needed + if not args.skip_nccl: + setup_nccl(rank, world_size) + + # Allocate MNNVL workspace if needed + workspace = None + if not args.skip_native: + workspace = allocate_mnnvl_workspace(rank, cp_size, mpi_comm) + + for batch_size in args.batch_sizes: + nccl_stats = None + native_stats = None + + # NCCL baseline + if not args.skip_nccl: + times = bench_nccl( + rank, + cp_size, + batch_size, + args.head_dim, + args.stats_dim, + dtype, + args.warmup, + args.iters, + mpi_comm, + ) + all_times = mpi_comm.allgather(times) + nccl_stats = compute_stats(all_times, args.iters) + + # Native DCP A2A + if not args.skip_native and workspace is not None: + times = bench_native( + workspace, + rank, + cp_size, + batch_size, + args.head_dim, + args.stats_dim, + dtype, + args.warmup, + args.iters, + mpi_comm, + ) + all_times = mpi_comm.allgather(times) + native_stats = compute_stats(all_times, args.iters) + + # Print results (rank 0 only) + if rank == 0: + print( + f" batch={batch_size}, head_dim={args.head_dim}, " + f"stats_dim={args.stats_dim}, dtype=bf16" + ) + if nccl_stats: + print( + f" NCCL (2x all_to_all_single): " + f"p50={nccl_stats['p50']:.3f}ms " + f"p95={nccl_stats['p95']:.3f}ms " + f"mean={nccl_stats['mean']:.3f}ms" + ) + if native_stats: + print( + f" Native (decode_cp_a2a_alltoall): " + f"p50={native_stats['p50']:.3f}ms " + f"p95={native_stats['p95']:.3f}ms " + f"mean={native_stats['mean']:.3f}ms" + ) + if nccl_stats and native_stats and native_stats["p50"] > 0: + speedup = nccl_stats["p50"] / native_stats["p50"] + print(f" Speedup: {speedup:.1f}x") + print() + + # Cleanup + if not args.skip_nccl and dist.is_initialized(): + dist.destroy_process_group() + + # Prevent segfault at exit: MnnvlMemory uses a bump allocator that + # doesn't support individual frees. If Python GC destroys the workspace + # tensor during interpreter shutdown, it triggers a segfault in + # TensorImpl::~TensorImpl. Calling os._exit() skips GC entirely. + mpi_comm.Barrier() + MPI.Finalize() + os._exit(0) + + +if __name__ == "__main__": + main() diff --git a/csrc/nv_internal/tensorrt_llm/kernels/cudaAsyncOps.cuh b/csrc/nv_internal/tensorrt_llm/kernels/cudaAsyncOps.cuh new file mode 100644 index 0000000000..0bf87186c5 --- /dev/null +++ b/csrc/nv_internal/tensorrt_llm/kernels/cudaAsyncOps.cuh @@ -0,0 +1,200 @@ +/* + * Copyright (c) 2019-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include + +#include "tensorrt_llm/kernels/moeCommKernelsCommon.h" + +namespace tensorrt_llm { +namespace kernels { + +// ============================================================================ +// Address Conversion Utilities +// ============================================================================ + +static __device__ __forceinline__ uint32_t __as_ptr_smem(void const* __ptr) { + // Consider adding debug asserts here. + return static_cast(__cvta_generic_to_shared(__ptr)); +} + +static __device__ __forceinline__ uint64_t __as_ptr_gmem(void const* __ptr) { + // Consider adding debug asserts here. + return static_cast(__cvta_generic_to_global(__ptr)); +} + +// ============================================================================ +// Memory Fence Operations +// ============================================================================ + +__device__ __forceinline__ void fence_release_sys() { + asm volatile("fence.release.sys;" : : : "memory"); +} + +// ============================================================================ +// Memory Barrier Operations (mbarrier) +// ============================================================================ + +__device__ __forceinline__ void mbarrier_init(uint64_t* addr, uint32_t const& count) { +#if defined(__CUDACC__) && __CUDA_ARCH__ >= 800 + asm("mbarrier.init.shared.b64 [%0], %1;" : : "r"(__as_ptr_smem(addr)), "r"(count) : "memory"); +#endif +} + +__device__ __forceinline__ void mbarrier_expect_tx(uint64_t* addr, const uint32_t txCount) { +#if defined(__CUDACC__) && __CUDA_ARCH__ >= 900 + asm("mbarrier.expect_tx.relaxed.cta.shared::cta.b64 [%0], %1;" + : + : "r"(__as_ptr_smem(addr)), "r"(txCount) + : "memory"); +#endif +} + +__device__ __forceinline__ uint64_t mbarrier_arrive(uint64_t* addr) { +#if defined(__CUDACC__) && __CUDA_ARCH__ >= 800 + uint64_t state; + asm("mbarrier.arrive.shared.b64 %0, [%1];" : "=l"(state) : "r"(__as_ptr_smem(addr)) : "memory"); + return state; +#else + return 0; +#endif +} + +__device__ __forceinline__ uint64_t mbarrier_arrive_expect_tx(uint64_t* addr, + const uint32_t txCount) { +#if defined(__CUDACC__) && __CUDA_ARCH__ >= 900 + uint64_t state; + asm("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 %0, [%1], %2;" + : "=l"(state) + : "r"(__as_ptr_smem(addr)), "r"(txCount) + : "memory"); + return state; +#else + return 0; +#endif +} + +__device__ __forceinline__ bool mbarrier_try_wait_parity(uint64_t* addr, + uint32_t const& phaseParity) { +#if defined(__CUDACC__) && __CUDA_ARCH__ >= 900 + uint32_t waitComplete; + asm("{\n\t .reg .pred P_OUT; \n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P_OUT, [%1], %2;\n\t" + "selp.b32 %0, 1, 0, P_OUT; \n" + "}" + : "=r"(waitComplete) + : "r"(__as_ptr_smem(addr)), "r"(phaseParity) + : "memory"); + return static_cast(waitComplete); +#else + return false; +#endif +} + +// ============================================================================ +// Async Copy Operations (cp.async for SM80+) +// ============================================================================ + +template +__device__ __forceinline__ void ldgsts(int* dstShm, int const* srcMem, bool predGuard) { +#if defined(__CUDACC__) && __CUDA_ARCH__ >= 800 + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.ca.shared.global [%1], [%2], %3;\n" + "}\n" ::"r"((int)predGuard), + "r"(__as_ptr_smem(dstShm)), "l"(__as_ptr_gmem(srcMem)), "n"(COPY_SIZE)); +#endif +} + +__device__ __forceinline__ void cp_async_commit_group() { +#if defined(__CUDACC__) && __CUDA_ARCH__ >= 800 + asm volatile("cp.async.commit_group;" : : :); +#endif +} + +template +__device__ __forceinline__ void cp_async_wait_group() { +#if defined(__CUDACC__) && __CUDA_ARCH__ >= 800 + asm volatile("cp.async.wait_group %0;" : : "n"(N) : "memory"); +#endif +} + +// ============================================================================ +// Bulk Async Copy Operations (cp.async.bulk for SM90+) +// ============================================================================ + +__device__ __forceinline__ void cp_async_bulk_g2s(void* dstMem, void const* srcMem, int copySize, + uint64_t* smemBar) { +#if defined(__CUDACC__) && __CUDA_ARCH__ >= 900 + asm("cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" + : + : "r"(__as_ptr_smem(dstMem)), "l"(__as_ptr_gmem(srcMem)), "r"(copySize), + "r"(__as_ptr_smem(smemBar)) + : "memory"); +#endif +} + +__device__ __forceinline__ void cp_async_bulk_s2g(void* dstMem, void const* srcMem, int copySize) { +#if defined(__CUDACC__) && __CUDA_ARCH__ >= 900 + asm("cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;" + : + : "l"(__as_ptr_gmem(dstMem)), "r"(__as_ptr_smem(srcMem)), "r"(copySize) + : "memory"); +#endif +} + +__device__ __forceinline__ void cp_async_bulk_commit_group() { +#if defined(__CUDACC__) && __CUDA_ARCH__ >= 900 + asm volatile("cp.async.bulk.commit_group;" : : :); +#endif +} + +template +__device__ __forceinline__ void cp_async_bulk_wait_group() { +#if defined(__CUDACC__) && __CUDA_ARCH__ >= 900 + asm volatile("cp.async.bulk.wait_group %0;" : : "n"(N) : "memory"); +#endif +} + +template +__device__ __forceinline__ void cp_async_bulk_wait_group_read() { +#if defined(__CUDACC__) && __CUDA_ARCH__ >= 900 + asm volatile("cp.async.bulk.wait_group.read %0;" : : "n"(N) : "memory"); +#endif +} + +// ============================================================================ +// Shared Memory Barrier Helpers +// ============================================================================ + +__device__ __forceinline__ void initSmemBar(uint64_t* smemBar, int laneId) { + if (laneId == 0) { + mbarrier_init(smemBar, WARP_SIZE); + } + __syncwarp(); +} + +__device__ __forceinline__ void smemBarWait(uint64_t* smemBar, uint32_t* phaseParity) { + while (!mbarrier_try_wait_parity(smemBar, *phaseParity)) { + } + *phaseParity = 1 - *phaseParity; +} + +} // namespace kernels +} // namespace tensorrt_llm diff --git a/csrc/nv_internal/tensorrt_llm/kernels/helixAllToAll.cu b/csrc/nv_internal/tensorrt_llm/kernels/helixAllToAll.cu new file mode 100644 index 0000000000..1751d30756 --- /dev/null +++ b/csrc/nv_internal/tensorrt_llm/kernels/helixAllToAll.cu @@ -0,0 +1,638 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/cudaAsyncOps.cuh" +#include "tensorrt_llm/kernels/helixAllToAll.h" +#include "tensorrt_llm/kernels/ll128Proto.cuh" +#include "tensorrt_llm/kernels/moeCommKernelsCommon.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels { + +namespace { + +// ============================================================================ +// Structure declarations and definitions +// ============================================================================ + +// ALIGN_256 is defined in moeCommKernelsCommon.h + +struct ALIGN_256 HelixFifoInfo { + volatile int64_t head; + volatile int64_t tail; +}; + +// ============================================================================ +// Helix-specific FIFO constants +// Note: Helix uses 128KB FIFO entries vs 256KB in FusedMoe +// ============================================================================ + +constexpr int HELIX_FIFO_DEPTH = 4; +constexpr int HELIX_FIFO_ENTRY_BYTES = 128 * 1024; +constexpr int HELIX_FIFO_TOTAL_BYTES = HELIX_FIFO_ENTRY_BYTES * HELIX_FIFO_DEPTH; +constexpr int HELIX_FIFO_ENTRY_128B_COUNT = HELIX_FIFO_ENTRY_BYTES / BYTES_PER_128B_BLOCK; +constexpr int HELIX_FIFO_TOTAL_U64 = HELIX_FIFO_TOTAL_BYTES / sizeof(uint64_t); + +// ============================================================================ +// Implementation-only structures +// ============================================================================ + +struct HelixPairInfo { + int senderRank; + int receiverRank; + int channel; + int runChannelCount; +}; + +// WARP_SIZE, WARP_MASK, and other constants are defined in moeCommKernelsCommon.h + +// ============================================================================ +// Helper Functions +// ============================================================================ + +__host__ __device__ inline int getFieldSize(HelixFieldInfo const& fieldInfo) { + return fieldInfo.elementCount * fieldInfo.elementSize; +} + +__host__ __device__ inline uint8_t* getPtr(HelixFieldInfo const& fieldInfo, int blockIdx) { + return fieldInfo.dataPtr + blockIdx * fieldInfo.stride; +} + +__device__ __forceinline__ void waitG2sAllFields(uint64_t* smemBar, uint32_t* phaseParity) { + cp_async_wait_group<0>(); + smemBarWait(smemBar, phaseParity); +} + +// Align size to 128 bytes +__host__ __device__ __forceinline__ int align128(int size) { + return align_up(size, BYTES_PER_128B_BLOCK); +} + +// ============================================================================ +// G2S (Global to Shared) Operations +// ============================================================================ + +__device__ __forceinline__ void g2sField(HelixFieldInfo const& fieldInfo, int dataIndex, + uint8_t* shmemBase, int shmemOffset, uint64_t* smemBar, + int laneId) { + int copySize = getFieldSize(fieldInfo); + if (copySize > 0 && laneId == 0) { + uint8_t* srcPtr = getPtr(fieldInfo, dataIndex); + uint8_t* dstPtr = shmemBase + shmemOffset; + cp_async_bulk_g2s(dstPtr, srcPtr, copySize, smemBar); + } +} + +template +__device__ __forceinline__ int g2sAllFields(HelixFieldInfo const* fieldInfo, int dataIndex, + uint8_t* shmemBase, uint64_t* smemBar, int laneId) { + int totalSize = 0; + + // Load field 0 (variable size half) + g2sField(fieldInfo[0], dataIndex, shmemBase, 0, smemBar, laneId); + int field0Size = getFieldSize(fieldInfo[0]); + totalSize += field0Size; + + // Load field 1 (single float2) + if constexpr (ALLOW_VARIABLE_FIELD1) { + g2sField(fieldInfo[1], dataIndex, shmemBase, totalSize, smemBar, laneId); + totalSize += getFieldSize(fieldInfo[1]); + } else { + ldgsts<8>(reinterpret_cast(shmemBase + totalSize), + reinterpret_cast(getPtr(fieldInfo[1], dataIndex)), laneId == 0); + cp_async_commit_group(); + } + + return totalSize; +} + +// ============================================================================ +// S2G (Shared to Global) Operations +// ============================================================================ + +__device__ __forceinline__ void s2gField(HelixFieldInfo const& fieldInfo, int dataIndex, + uint8_t* shmemBase, int shmemOffset, int laneId) { + int copySize = getFieldSize(fieldInfo); + if (copySize > 0 && laneId == 0) { + uint8_t* srcPtr = shmemBase + shmemOffset; + uint8_t* dstPtr = getPtr(fieldInfo, dataIndex); + cp_async_bulk_s2g(dstPtr, srcPtr, copySize); + } +} + +template +__device__ __forceinline__ void s2gAllFields(HelixFieldInfo const* fieldInfo, int dataIndex, + uint8_t* shmemBase, int laneId) { + int offset = 0; + + // Store field 0 (variable size half) + s2gField(fieldInfo[0], dataIndex, shmemBase, offset, laneId); + int field0Size = getFieldSize(fieldInfo[0]); + offset += field0Size; + + // Store field 1 (single float2) + if constexpr (ALLOW_VARIABLE_FIELD1) { + s2gField(fieldInfo[1], dataIndex, shmemBase, offset, laneId); + offset += getFieldSize(fieldInfo[1]); + } else { + if (laneId == 0) { + auto* srcPtr = reinterpret_cast(reinterpret_cast(shmemBase) + offset); + auto* dstPtr = reinterpret_cast(getPtr(fieldInfo[1], dataIndex)); + dstPtr[0] = srcPtr[0]; + } + } + cp_async_bulk_commit_group(); +} + +// ============================================================================ +// Workspace FIFO Operations +// ============================================================================ + +__device__ __forceinline__ uint64_t* getFifoBasePtr(HelixAllToAllParams const& params, + HelixPairInfo const& pairInfo) { + // FIFO is physically located at receiver rank + int mappedMemoryRank = pairInfo.receiverRank; + int rankInsideMappedMemory = pairInfo.senderRank; + + auto* mappedMemory = params.workspace + mappedMemoryRank * params.workspaceStrideInU64; + // Navigate to the right FIFO: [peer_rank][channel] + size_t fifoOffset = rankInsideMappedMemory * params.maxChannelCount * HELIX_FIFO_TOTAL_U64; + fifoOffset += pairInfo.channel * HELIX_FIFO_TOTAL_U64; + + return mappedMemory + fifoOffset; +} + +__device__ __forceinline__ HelixFifoInfo* getSenderHelixFifoInfo(HelixAllToAllParams const& params, + HelixPairInfo const& pairInfo) { + // SenderSideHelixFifoInfo is physically located at sender rank + int mappedMemoryRank = pairInfo.senderRank; + int rankInsideMappedMemory = pairInfo.receiverRank; + + auto* mappedMemory = + reinterpret_cast(params.workspace + mappedMemoryRank * params.workspaceStrideInU64); + size_t fieldOffset = + static_cast(HELIX_FIFO_TOTAL_BYTES) * params.cpSize * params.maxChannelCount; + mappedMemory += fieldOffset; + mappedMemory += rankInsideMappedMemory * params.maxChannelCount * sizeof(HelixFifoInfo); + mappedMemory += pairInfo.channel * sizeof(HelixFifoInfo); + + return reinterpret_cast(mappedMemory); +} + +__device__ __forceinline__ HelixFifoInfo* getReceiverHelixFifoInfo( + HelixAllToAllParams const& params, HelixPairInfo const& pairInfo) { + // ReceiverSideHelixFifoInfo is physically located at receiver rank + int mappedMemoryRank = pairInfo.receiverRank; + int rankInsideMappedMemory = pairInfo.senderRank; + + auto* mappedMemory = + reinterpret_cast(params.workspace + mappedMemoryRank * params.workspaceStrideInU64); + size_t fieldOffset = + static_cast(HELIX_FIFO_TOTAL_BYTES) * params.cpSize * params.maxChannelCount; + fieldOffset += sizeof(HelixFifoInfo) * params.cpSize * params.maxChannelCount; + mappedMemory += fieldOffset; + mappedMemory += rankInsideMappedMemory * params.maxChannelCount * sizeof(HelixFifoInfo); + mappedMemory += pairInfo.channel * sizeof(HelixFifoInfo); + + return reinterpret_cast(mappedMemory); +} + +__device__ __forceinline__ void startWorkspaceS2G(uint64_t* fifoEntry, uint8_t* shmemBase, + int send128ByteCount, int fifo128ByteOffset, + int laneId) { + int copyByteCount = send128ByteCount * BYTES_PER_128B_BLOCK; + if (laneId == 0) { + cp_async_bulk_s2g(fifoEntry + fifo128ByteOffset * BYTES_PER_128B_BLOCK / sizeof(uint64_t), + shmemBase, copyByteCount); + } + cp_async_bulk_commit_group(); +} + +__device__ __forceinline__ void startWorkspaceS2GReg(uint64_t* fifoEntry, uint8_t* sharedMemoryBase, + int send128ByteCount, int fifo128ByteOffset, + int laneId) { + int copyInt4Count = send128ByteCount * BYTES_PER_128B_BLOCK / sizeof(int4); + int4* sharedMemoryInt4 = reinterpret_cast(sharedMemoryBase); + uint64_t* fifoPtr = fifoEntry + fifo128ByteOffset * UINT64_PER_128B_BLOCK; + int4* fifoPtrInt4 = reinterpret_cast(fifoPtr); +#pragma unroll 4 + for (int i = laneId; i < copyInt4Count; i += WARP_SIZE) { + fifoPtrInt4[i] = sharedMemoryInt4[i]; + } +} + +__device__ __forceinline__ uint64_t startWorkspaceG2S(uint8_t* shmemBase, uint64_t* fifoEntry, + int allLoad128ByteCount, + int fifo128ByteOffset, int loaded128ByteCount, + uint64_t* smemBar, int laneId) { + int copyByteCount = (allLoad128ByteCount - loaded128ByteCount) * BYTES_PER_128B_BLOCK; + if (laneId == 0) { + cp_async_bulk_g2s(shmemBase + loaded128ByteCount * BYTES_PER_128B_BLOCK, + fifoEntry + (fifo128ByteOffset + loaded128ByteCount) * UINT64_PER_128B_BLOCK, + copyByteCount, smemBar); + } + return mbarrier_arrive_expect_tx(smemBar, laneId == 0 ? copyByteCount : 0); +} + +// LL128Proto is now defined in ll128Proto.cuh + +// ============================================================================ +// Size helpers +// ============================================================================ + +// Compute total size needed for both fields +__host__ __device__ __forceinline__ int computeTotalUnpackedSize(HelixFieldInfo const* fields) { + int size = 0; + // Field 0: note it must be aligned to 16 bytes + size += align_up(getFieldSize(fields[0]), 16); + // Field 1: single float2 + size += align_up(getFieldSize(fields[1]), 16); + return align128(size); +} + +__host__ __device__ __forceinline__ int computeTotalPackedSize(HelixFieldInfo const* fields) { + // because field 0 must be aligned to 16 bytes, this is the same as unpacked + return computeTotalUnpackedSize(fields); +} + +__host__ __device__ __forceinline__ int computeProtoTransferSize(HelixFieldInfo const* fields) { + return LL128Proto::computeProtoTransfer128ByteAlignedSize(computeTotalPackedSize(fields)); +} + +// ============================================================================ +// Main All-to-All Kernel +// ============================================================================ + +template +__global__ void helixAllToAllKernel(HelixAllToAllParams params) { + extern __shared__ uint8_t allWarpShmem[]; + __shared__ uint64_t allWarpSmemBar[MAX_GROUP_COUNT_PER_BLOCK]; + + bool isSender = (blockIdx.z == 0); + // Each warp is a group handling a different peer rank + int group = __shfl_sync(WARP_MASK, threadIdx.y, 0); + int laneId = threadIdx.x % WARP_SIZE; + int runChannelCount = gridDim.y; + + // Compute peer rank: blockIdx.x determines which set of peers, group + // determines which peer in that set + int peerRank = blockIdx.x * blockDim.y + group; + + if (peerRank >= params.cpSize) { + return; + } + + // Setup pair info for this communication + HelixPairInfo pairInfo; + pairInfo.channel = blockIdx.y; + pairInfo.runChannelCount = runChannelCount; + pairInfo.senderRank = isSender ? params.cpRank : peerRank; + pairInfo.receiverRank = isSender ? peerRank : params.cpRank; + + // Initialize barrier for this group + initSmemBar(&allWarpSmemBar[group], laneId); + uint32_t phaseParity = 0; + + // Get shared memory for this group + int singlePackedSize = computeTotalPackedSize(params.sendFields); + int singlePacked128ByteCount = singlePackedSize / BYTES_PER_128B_BLOCK; + int singleUnpackedSize = computeTotalUnpackedSize(params.sendFields); + int singleProtoTransferSize = computeProtoTransferSize(params.sendFields); + int singleProtoTransfer128ByteCount = singleProtoTransferSize / BYTES_PER_128B_BLOCK; + int singleShmSize = std::max(singleUnpackedSize, singleProtoTransferSize); + uint8_t* shmem = allWarpShmem + group * singleShmSize; + + // Get FIFO pointers + uint64_t* fifoBase = getFifoBasePtr(params, pairInfo); + HelixFifoInfo* senderFifo = getSenderHelixFifoInfo(params, pairInfo); + HelixFifoInfo* receiverFifo = getReceiverHelixFifoInfo(params, pairInfo); + + int fifoEntry128ByteIndexBase = HELIX_FIFO_ENTRY_128B_COUNT; + int fifoEntryIndex = -1; + + // regardless of sender or receiver, we wait for the previous kernel here + // receiver blocks do not need to wait at all, but they should not start + // to stress the memory system regardless +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaGridDependencySynchronize(); +#endif + + if (isSender) { + // sender blocks should trigger next kernel immediately, s.t. they + // do not block the next kernel from starting +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + // Sender logic: send data from cpRank's slice to peerRank + int64_t head = senderFifo->head; + int64_t tail = senderFifo->tail; + + // Each channel processes entries with stride + // Start at channel index, increment by total channel count + for (int entryIdx = pairInfo.channel; entryIdx < params.entryCount; + entryIdx += runChannelCount) { + // dataIndex points to the data for peerRank in this entry + int dataIndex = entryIdx * params.cpSize + peerRank; + + // Load data from global to shared, then arrive on barrier + int loadedSize = g2sAllFields(params.sendFields, dataIndex, shmem, + &allWarpSmemBar[group], laneId); + uint64_t arriveState = + mbarrier_arrive_expect_tx(&allWarpSmemBar[group], laneId == 0 ? loadedSize : 0); + + // update FIFO entry index and head if needed + if (fifoEntry128ByteIndexBase + singleProtoTransfer128ByteCount > + HELIX_FIFO_ENTRY_128B_COUNT) { + if (fifoEntryIndex >= 0) { + head++; + __syncwarp(); + senderFifo->head = head; + } + fifoEntryIndex = head % HELIX_FIFO_DEPTH; + fifoEntry128ByteIndexBase = 0; + while (tail + HELIX_FIFO_DEPTH <= head) { + tail = senderFifo->tail; + } + __syncwarp(); + } + + // wait for data to be loaded into shared memory + waitG2sAllFields(&allWarpSmemBar[group], &phaseParity); + // note: we don't need to pack anything, fields are already packed in + // shared memory + + LL128Proto::protoPack(shmem, head, singlePacked128ByteCount, fifoEntry128ByteIndexBase, + laneId); + + uint64_t* fifoEntry = fifoBase + fifoEntryIndex * (HELIX_FIFO_ENTRY_BYTES / sizeof(uint64_t)); + + // Copy from shared to workspace FIFO + startWorkspaceS2GReg(fifoEntry, shmem, singleProtoTransfer128ByteCount, + fifoEntry128ByteIndexBase, laneId); + + fifoEntry128ByteIndexBase += singleProtoTransfer128ByteCount; + + // ensure that we can over-write shmem in next iteration + // (it must be fully read by all threads when doing S2G above) + __syncwarp(); + } + if (fifoEntry128ByteIndexBase > 0) { + head++; + senderFifo->head = head; + } + } else { + // Receiver logic: receive data from peerRank to cpRank's slice + int64_t tail = receiverFifo->tail; + bool needRelease = false; + + // Each channel processes entries with stride + // Start at channel index, increment by total channel count + for (int entryIdx = pairInfo.channel; entryIdx < params.entryCount; + entryIdx += runChannelCount) { + // receiver blocks should trigger next kernel at last iteration + // note: some blocks might not even go into this for-loop, but they + // would exit which is equivalent to the pre-exit trigger +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + if (entryIdx + runChannelCount >= params.entryCount) { + cudaTriggerProgrammaticLaunchCompletion(); + } +#endif + // dataIndex points to where we receive data from peerRank in this entry + int dataIndex = entryIdx * params.cpSize + peerRank; + int loaded128ByteCount = 0; + + if (fifoEntry128ByteIndexBase + singleProtoTransfer128ByteCount > + HELIX_FIFO_ENTRY_128B_COUNT) { + if (fifoEntryIndex >= 0) { + tail++; + needRelease = true; + } + fifoEntryIndex = tail % HELIX_FIFO_DEPTH; + fifoEntry128ByteIndexBase = 0; + // receiver doesn't need to wait on FIFO entry being readable: it's + // always readable + __syncwarp(); + } + + uint64_t* fifoEntry = fifoBase + fifoEntryIndex * (HELIX_FIFO_ENTRY_BYTES / sizeof(uint64_t)); + while (loaded128ByteCount < singleProtoTransfer128ByteCount) { + startWorkspaceG2S(shmem, fifoEntry, singleProtoTransfer128ByteCount, + fifoEntry128ByteIndexBase, loaded128ByteCount, &allWarpSmemBar[group], + laneId); + if (needRelease) { + receiverFifo->tail = tail; + senderFifo->tail = tail; + needRelease = false; + } + smemBarWait(&allWarpSmemBar[group], &phaseParity); + loaded128ByteCount += LL128Proto::template checkDataReceivedInShm( + shmem, tail, singleProtoTransfer128ByteCount, fifoEntry128ByteIndexBase, + loaded128ByteCount, laneId); + } + + LL128Proto::protoUnpack(shmem, tail, singlePacked128ByteCount, fifoEntry128ByteIndexBase, + loaded128ByteCount, laneId); + + // note: fields are already unpacked in shared memory + s2gAllFields(params.recvFields, dataIndex, shmem, laneId); + // wait for data to be read from shared memory + cp_async_bulk_wait_group_read<0>(); + + // note: LL128Proto doesn't need rearm + // rearmFifoBuffer(); + fifoEntry128ByteIndexBase += singleProtoTransfer128ByteCount; + } + if (fifoEntry128ByteIndexBase > 0) { + tail++; + receiverFifo->tail = tail; + senderFifo->tail = tail; + } + } +} + +// ============================================================================ +// Compute actual channel count +// ============================================================================ + +struct hash_cache_key { + size_t operator()(std::tuple const& x) const { + return std::get<0>(x) ^ std::get<1>(x) ^ std::get<2>(x); + } +}; + +template +std::tuple computeChannelAndGroupCount(int cpSize, HelixFieldInfo const* fields) { + static std::unordered_map, std::tuple, hash_cache_key> + cache; + int deviceId = 0; + TLLM_CUDA_CHECK(cudaGetDevice(&deviceId)); + int singleShmSize = std::max(computeTotalUnpackedSize(fields), computeProtoTransferSize(fields)); + auto key = std::make_tuple(deviceId, cpSize, singleShmSize); + auto it = cache.find(key); + if (it != cache.end()) { + return it->second; + } + + int maxGroupCountPerCta = std::min(cpSize, MAX_GROUP_COUNT_PER_BLOCK); + int groupCountPerCta = maxGroupCountPerCta; // Start with max + int totalDynamicShmemSize = singleShmSize * groupCountPerCta; + int maxDynamicShmSize = 0; + TLLM_CUDA_CHECK(cudaDeviceGetAttribute(&maxDynamicShmSize, + cudaDevAttrMaxSharedMemoryPerBlockOptin, deviceId)); + + while (totalDynamicShmemSize > maxDynamicShmSize) { + groupCountPerCta--; + totalDynamicShmemSize = singleShmSize * groupCountPerCta; + } + + TLLM_CHECK_WITH_INFO(totalDynamicShmemSize <= maxDynamicShmSize, + "Single packed size %d exceeds limit %d", singleShmSize, maxDynamicShmSize); + + // Set shared memory attribute if needed + if (totalDynamicShmemSize > 48 * 1024) { + TLLM_CUDA_CHECK(cudaFuncSetAttribute(helixAllToAllKernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + totalDynamicShmemSize)); + } + + int blockCountPerChannel = ceil_div(cpSize, groupCountPerCta); + blockCountPerChannel *= 2; // for send and recv + + int smCount = 0; + TLLM_CUDA_CHECK(cudaDeviceGetAttribute(&smCount, cudaDevAttrMultiProcessorCount, deviceId)); + // TODO: we might only want to use half the SMs to overlap with other kernels. + // note that overlap with FMHA is almost impossible because it must use + // all SMs and probably uses >50% shmem per SM. + // overlap with the subsequent BMM / out proj GEMMs might be possible, + // so we need experiments to see whether it makes sense. + int channelCount = std::max(smCount / blockCountPerChannel, 1); + auto value = std::make_tuple(channelCount, groupCountPerCta, totalDynamicShmemSize); + cache[key] = value; + return value; +} + +// ============================================================================ +// Host Launch Function +// ============================================================================ + +template +void launchHelixAllToAllImpl(HelixAllToAllParams const& params, bool enablePdl, + cudaStream_t stream) { + int maxChannelCount = computeHelixMaxChannelCount(params.cpSize); + TLLM_CHECK_WITH_INFO(params.maxChannelCount == maxChannelCount, + "maxChannelCount %d does not match computed maxChannelCount %d", + params.maxChannelCount, maxChannelCount); + auto [channelCount, groupCountPerCta, totalDynamicShmemSize] = + computeChannelAndGroupCount(params.cpSize, params.sendFields); + if (params.channelCount > 0) { + channelCount = params.channelCount; + TLLM_CHECK_WITH_INFO(channelCount <= maxChannelCount, + "channelCount %d exceeds maxChannelCount %d", channelCount, + maxChannelCount); + } + + // Compute grid dimensions + // grid.x = blocks per channel (how many blocks needed to cover all peer + // ranks) grid.y = number of channels (parallel channels) grid.z = 2 (sender + // and receiver) + int ctaPerChannel = ceil_div(params.cpSize, groupCountPerCta); + + auto* kernel_instance = &helixAllToAllKernel; + cudaLaunchConfig_t config; + config.gridDim = dim3(ctaPerChannel, channelCount, 2); + config.blockDim = dim3(WARP_SIZE, groupCountPerCta); + config.dynamicSmemBytes = totalDynamicShmemSize; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enablePdl; + config.numAttrs = 1; + config.attrs = attrs; + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel_instance, params)); +} + +} // anonymous namespace + +// ============================================================================ +// Public API Functions +// ============================================================================ + +int computeHelixMaxChannelCount(int cpSize, int smCount) { + if (smCount == 0) { + int deviceId = 0; + TLLM_CUDA_CHECK(cudaGetDevice(&deviceId)); + TLLM_CUDA_CHECK(cudaDeviceGetAttribute(&smCount, cudaDevAttrMultiProcessorCount, deviceId)); + } + + int blockCountPerChannel = ceil_div(cpSize, MAX_GROUP_COUNT_PER_BLOCK); + blockCountPerChannel *= 2; // for send and recv + + int preferredChannel = smCount / blockCountPerChannel; + return std::max(preferredChannel, 1); // at least one channel +} + +size_t computeHelixWorkspaceSizePerRank(int cpSize) { + int maxChannelCount = computeHelixMaxChannelCount(cpSize); + + // FIFO buffers: cpSize * channelCount pairs + size_t fifoSize = static_cast(HELIX_FIFO_TOTAL_BYTES) * cpSize * maxChannelCount; + + // Sender and receiver FIFO info structures + size_t senderInfoSize = sizeof(HelixFifoInfo) * cpSize * maxChannelCount; + size_t receiverInfoSize = sizeof(HelixFifoInfo) * cpSize * maxChannelCount; + + return fifoSize + senderInfoSize + receiverInfoSize; +} + +void launchHelixAllToAll(HelixAllToAllParams const& params, bool allowVariableField1, + bool enablePdl, cudaStream_t stream) { + if (allowVariableField1) { + launchHelixAllToAllImpl(params, enablePdl, stream); + } else { + launchHelixAllToAllImpl(params, enablePdl, stream); + } +} + +// ============================================================================ +// Workspace Initialization +// ============================================================================ + +void initializeHelixWorkspace(uint64_t* local_workspace_ptr, int cpSize, cudaStream_t stream) { + int maxChannelCount = computeHelixMaxChannelCount(cpSize); + // Calculate sizes with channel dimension + size_t fifoSize = static_cast(HELIX_FIFO_TOTAL_BYTES) * cpSize * maxChannelCount; + size_t senderInfoSize = sizeof(HelixFifoInfo) * cpSize * maxChannelCount; + size_t receiverInfoSize = sizeof(HelixFifoInfo) * cpSize * maxChannelCount; + + // Initialize FIFO buffers to 0xFFFFFFFF (-1 for signed integer types) + TLLM_CUDA_CHECK(cudaMemsetAsync(local_workspace_ptr, 0xFF, fifoSize, stream)); + + // Initialize sender and receiver info to zero (single call for both) + uint8_t* infoPtr = reinterpret_cast(local_workspace_ptr) + fifoSize; + TLLM_CUDA_CHECK(cudaMemsetAsync(infoPtr, 0, senderInfoSize + receiverInfoSize, stream)); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/csrc/nv_internal/tensorrt_llm/kernels/helixAllToAll.h b/csrc/nv_internal/tensorrt_llm/kernels/helixAllToAll.h new file mode 100644 index 0000000000..641b9276e1 --- /dev/null +++ b/csrc/nv_internal/tensorrt_llm/kernels/helixAllToAll.h @@ -0,0 +1,92 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include + +#include +#include + +#include "tensorrt_llm/common/config.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels { + +struct HelixFieldInfo { + uint8_t* dataPtr; + int elementCount; // Number of elements (e.g., kv_lora_rank for field 0, 1 for + // field 1) + int elementSize; // Size of each element in bytes (2 for half, 8 for float2) + int stride; // Stride between rows in bytes +}; + +struct HelixAllToAllParams { + HelixFieldInfo sendFields[2]; + HelixFieldInfo recvFields[2]; + int entryCount; // Number of entries per peer rank to process + uint64_t* workspace; + size_t workspaceStrideInU64; + int cpRank; + int cpSize; + int channelCount; // use 0 to auto-compute + int maxChannelCount; +}; + +// ============================================================================ +// Workspace Management Functions +// ============================================================================ + +/** + * Compute number of channels for communication based on cpSize. + * + * @param cpSize Number of context parallel ranks + * @param smCount Number of SMs available (0 = auto-detect) + * @return Number of channels to use + */ +int computeHelixMaxChannelCount(int cpSize, int smCount = 0); + +/** + * Compute the workspace size required per rank for the all-to-all operation. + * + * @param cpSize Number of context parallel ranks + * @return Size in bytes + */ +size_t computeHelixWorkspaceSizePerRank(int cpSize); + +/** + * Initialize workspace memory for a given rank. + * Should be called once during setup. + * + * @param workspace Pointer to workspace memory (per-rank view) + * @param cpSize Number of context parallel ranks + * @param stream CUDA stream for asynchronous operations + */ +void initializeHelixWorkspace(uint64_t* workspace, int cpSize, cudaStream_t stream); + +/** + * Launch the helix all-to-all kernel. + * + * @param params Kernel parameters including field info and workspace + * @param allowVariableField1 Whether to allow variable field 1 + * @param stream CUDA stream for kernel launch + */ +void launchHelixAllToAll(HelixAllToAllParams const& params, bool allowVariableField1, + bool enablePdl, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/csrc/nv_internal/tensorrt_llm/kernels/ll128Proto.cuh b/csrc/nv_internal/tensorrt_llm/kernels/ll128Proto.cuh new file mode 100644 index 0000000000..2ccd929e55 --- /dev/null +++ b/csrc/nv_internal/tensorrt_llm/kernels/ll128Proto.cuh @@ -0,0 +1,153 @@ +/* + * Copyright (c) 2019-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include + +#include "tensorrt_llm/kernels/moeCommKernelsCommon.h" + +namespace tensorrt_llm { +namespace kernels { + +class LL128Proto { + public: + static constexpr uint32_t INITIALIZED_VALUE = 0xFFFFFFFFU; + + template + static __device__ __forceinline__ int checkDataReceivedInShm(uint8_t* sharedMemoryBase, + uint64_t step, int countIn128Bytes, + int fifoEntry128ByteIndexBase, + int loaded128ByteCount, int laneId) { + // return value should be how many package already been received. + // 0 means no data received, -1 means has received finish package(should be the very first 128 + // Byte). + uint64_t* aligned128BytesShm = reinterpret_cast(sharedMemoryBase); + int totalValidCount = 0; + for (int idxBase = loaded128ByteCount; idxBase < countIn128Bytes; idxBase += WARP_SIZE) { + int idx = idxBase + laneId; + bool valid = false; + bool finish = false; + if (idx < countIn128Bytes) { + int indexInFifoEntry = fifoEntry128ByteIndexBase + idx; + uint64_t value = aligned128BytesShm[idx * UINT64_PER_128B_BLOCK + + indexInFifoEntry % UINT64_PER_128B_BLOCK]; + if (USE_FINISH) { + finish = (value == (step & (1ULL << 63ULL))); + valid = (value == step) || finish; + } else { + valid = (value == step); + } + } + __syncwarp(); + unsigned validMask = __ballot_sync(WARP_MASK, valid); + // here we check valid in order, if previous valid is not true, we ignore the current valid. + int validCount = (validMask == WARP_MASK) ? WARP_SIZE : (__ffs(~validMask) - 1); + if (USE_FINISH) { + unsigned finishedMask = __ballot_sync(WARP_MASK, finish); + // finish should be the very first 128 Byte. + if (finishedMask & 0x1) { + return -1; + } + } + totalValidCount += validCount; + + if (validCount != WARP_SIZE) { + break; + } + } + return totalValidCount; + } + + static __device__ __forceinline__ void protoPack(uint8_t* sharedMemoryBase, uint64_t step, + int countIn128Bytes, + int fifoEntry128ByteIndexBase, int laneId) { + uint64_t* aligned128BytesShm = reinterpret_cast(sharedMemoryBase); + int halfLaneId = laneId % 16; + int halfIndex = laneId / 16; + int tailOffsetIn128Bytes = countIn128Bytes + halfIndex; + // for LL128 15 * 128 Bytes will be packed to 16 * 128 Bytes, each 16 threads is used for one 15 + // * 128 bytes. + for (int idxIn128BytesBase = halfIndex * 15; idxIn128BytesBase < countIn128Bytes; + idxIn128BytesBase += 30) { + int tailFlagIndexFromFifoEntry = fifoEntry128ByteIndexBase + tailOffsetIn128Bytes; + int tailFlagInnerIndex = tailFlagIndexFromFifoEntry % UINT64_PER_128B_BLOCK; + int idxIn128Bytes = idxIn128BytesBase + halfLaneId; + int idxFromFifoEntry = fifoEntry128ByteIndexBase + idxIn128Bytes; + uint64_t tailValue = step; + uint64_t tailInnerIndex = (halfLaneId >= tailFlagInnerIndex) ? halfLaneId + 1 : halfLaneId; + if (halfLaneId == 15) { + tailInnerIndex = tailFlagInnerIndex; + } + int targetTailIndex = tailOffsetIn128Bytes * UINT64_PER_128B_BLOCK + tailInnerIndex; + if (idxIn128Bytes < countIn128Bytes && halfLaneId < 15) { + int flagIndex = + idxIn128Bytes * UINT64_PER_128B_BLOCK + idxFromFifoEntry % UINT64_PER_128B_BLOCK; + tailValue = aligned128BytesShm[flagIndex]; + aligned128BytesShm[flagIndex] = step; + } + aligned128BytesShm[targetTailIndex] = tailValue; + tailOffsetIn128Bytes += 2; + } + __syncwarp(); + } + + static __device__ __forceinline__ void protoUnpack(uint8_t* sharedMemoryBase, uint64_t step, + int countIn128Bytes, + int fifoEntry128ByteIndexBase, + int loaded128ByteCount, int laneId) { + uint64_t* aligned128BytesShm = reinterpret_cast(sharedMemoryBase); + int halfLaneId = laneId % 16; + int halfIndex = laneId / 16; + int tailOffsetIn128Bytes = countIn128Bytes + halfIndex; + for (int idxIn128BytesBase = halfIndex * 15; idxIn128BytesBase < countIn128Bytes; + idxIn128BytesBase += 30) { + int tailFlagIndexFromFifoEntry = fifoEntry128ByteIndexBase + tailOffsetIn128Bytes; + int tailFlagInnerIndex = tailFlagIndexFromFifoEntry % UINT64_PER_128B_BLOCK; + int idxIn128Bytes = idxIn128BytesBase + halfLaneId; + int idxFromFifoEntry = fifoEntry128ByteIndexBase + idxIn128Bytes; + uint64_t tailValue = 0; + int tailInnerIndex = (halfLaneId >= tailFlagInnerIndex) ? halfLaneId + 1 : halfLaneId; + int targetTailIndex = tailOffsetIn128Bytes * UINT64_PER_128B_BLOCK + tailInnerIndex; + if (halfLaneId < 15) { + tailValue = aligned128BytesShm[targetTailIndex]; + } + if (idxIn128Bytes < countIn128Bytes && halfLaneId < 15) { + int flagIndex = + idxIn128Bytes * UINT64_PER_128B_BLOCK + idxFromFifoEntry % UINT64_PER_128B_BLOCK; + aligned128BytesShm[flagIndex] = tailValue; + } + tailOffsetIn128Bytes += 2; + } + __syncwarp(); + } + + static __device__ __forceinline__ void rearm(uint32_t* u32FifoPtr, uint64_t step, + int countIn128Bytes, int fifoEntry128ByteIndexBase, + int laneId) { + // LL128 don't need rearm + } + + static __device__ __host__ __forceinline__ int computeProtoTransfer128ByteAlignedSize( + int compact128ByteSizeBeforeProto) { + // each 15 * 128 byte need one tail 128 byte + int tail128ByteSize = (compact128ByteSizeBeforeProto + 15 * 128 - 1) / (15 * 128) * 128; + return compact128ByteSizeBeforeProto + tail128ByteSize; + } +}; + +} // namespace kernels +} // namespace tensorrt_llm diff --git a/csrc/nv_internal/tensorrt_llm/kernels/moeCommKernelsCommon.h b/csrc/nv_internal/tensorrt_llm/kernels/moeCommKernelsCommon.h new file mode 100644 index 0000000000..30024dac08 --- /dev/null +++ b/csrc/nv_internal/tensorrt_llm/kernels/moeCommKernelsCommon.h @@ -0,0 +1,102 @@ +/* + * Copyright (c) 2019-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include + +#include "tensorrt_llm/common/config.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels { + +// ============================================================================ +// Alignment Macro +// ============================================================================ + +#ifdef __CUDACC__ +#define ALIGN_256 __align__(256) +#else +#define ALIGN_256 alignas(256) +#endif + +// ============================================================================ +// Warp Constants +// ============================================================================ + +constexpr int WARP_SIZE = 32; +constexpr uint32_t WARP_MASK = 0xffffffff; + +// ============================================================================ +// Memory Block Constants +// ============================================================================ + +// Size of a 128-byte aligned block (used for bulk async copies) +constexpr int BYTES_PER_128B_BLOCK = 128; + +// Size of a 16-byte aligned block (used for field alignment) +constexpr int BYTES_PER_16B_BLOCK = 16; + +// Number of int elements per 128-byte block +constexpr int INTS_PER_128B_BLOCK = BYTES_PER_128B_BLOCK / sizeof(int); + +// Number of uint64_t elements per 128-byte block +constexpr int UINT64_PER_128B_BLOCK = BYTES_PER_128B_BLOCK / sizeof(uint64_t); + +// ============================================================================ +// Block Organization Constants +// ============================================================================ + +// Maximum number of groups (warps) per CTA for MoE communication kernels +constexpr int MAX_GROUP_COUNT_PER_BLOCK = 8; + +// ============================================================================ +// Utility Functions +// ============================================================================ + +/** + * Ceiling division: compute ceil(a / b) for integers + */ +template +inline constexpr T ceil_div(T a, T b) { + return (a + b - 1) / b; +} + +/** + * Align value up to nearest multiple of alignment + */ +template +inline constexpr T align_up(T value, T alignment) { + return ceil_div(value, alignment) * alignment; +} + +// ============================================================================ +// MoE Parallel Info Structures +// ============================================================================ + +struct MoeEpWorldInfo { + int epSize; + int epRank; +}; + +struct MoeExpertParallelInfo { + int expertCount = -1; + int topK = 1; +}; + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/csrc/trtllm_dcp_alltoall.cu b/csrc/trtllm_dcp_alltoall.cu new file mode 100644 index 0000000000..76851e09d2 --- /dev/null +++ b/csrc/trtllm_dcp_alltoall.cu @@ -0,0 +1,165 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +#include +#include +#include + +#include "tensorrt_llm/kernels/helixAllToAll.h" +#include "tvm_ffi_utils.h" + +using tvm::ffi::Tensor; +using tvm::ffi::TensorView; + +namespace { + +static int getEnvChannelCount() { + static int cached = -1; + if (cached < 0) { + const char* env = std::getenv("DCP_A2A_CHANNEL_COUNT"); + cached = (env && std::string(env) != "0") ? std::atoi(env) : 0; + } + return cached; +} + +int64_t getDcpWorkspaceSizePerRank(int64_t cp_size) { + return static_cast( + tensorrt_llm::kernels::computeHelixWorkspaceSizePerRank(static_cast(cp_size))); +} + +void initializeDcpWorkspaceOp(TensorView workspace, int64_t cp_rank, int64_t cp_size) { + CHECK_INPUT_TYPE(workspace, dl_int64); + TVM_FFI_ICHECK_EQ(workspace.ndim(), 2) << "workspace must be 2D"; + TVM_FFI_ICHECK_EQ(workspace.size(0), cp_size) << "workspace first dim must equal cp_size"; + TVM_FFI_ICHECK(cp_rank >= 0 && cp_rank < cp_size) << "cp_rank must be in [0, cp_size)"; + + auto stream = get_current_stream(); + auto* global_ptr = reinterpret_cast(workspace.data_ptr()); + auto* local_ptr = global_ptr + cp_rank * workspace.stride(0); + + tensorrt_llm::kernels::initializeHelixWorkspace(local_ptr, static_cast(cp_size), stream); +} + +tvm::ffi::Tuple alltoallDcpNativeOp(TensorView partial_o, TensorView softmax_stats, + TensorView workspace, int64_t cp_rank, + int64_t cp_size, bool enable_pdl) { + CHECK_INPUT(partial_o); + CHECK_INPUT(softmax_stats); + CHECK_CUDA(workspace); + + auto po_dtype_code = encode_dlpack_dtype(partial_o.dtype()); + TVM_FFI_ICHECK(po_dtype_code == float16_code || po_dtype_code == bfloat16_code) + << "partial_o must be half or bfloat16"; + CHECK_INPUT_TYPE(softmax_stats, dl_float32); + CHECK_INPUT_TYPE(workspace, dl_int64); + + TVM_FFI_ICHECK(partial_o.ndim() >= 2) << "partial_o must have at least 2 dimensions"; + TVM_FFI_ICHECK(softmax_stats.ndim() >= 2) << "softmax_stats must have at least 2 dimensions"; + TVM_FFI_ICHECK_EQ(partial_o.ndim(), softmax_stats.ndim()) + << "partial_o and softmax_stats must have same number of dimensions"; + + int64_t kv_lora_rank = partial_o.size(-1); + TVM_FFI_ICHECK_EQ(partial_o.size(-2), cp_size) + << "partial_o second-to-last dim must equal cp_size"; + TVM_FFI_ICHECK_EQ(softmax_stats.size(-2), cp_size) + << "softmax_stats second-to-last dim must equal cp_size"; + TVM_FFI_ICHECK(softmax_stats.size(-1) % 2 == 0 && softmax_stats.size(-1) >= 2) + << "softmax_stats last dim must be even and >= 2"; + bool allowVariableField1 = softmax_stats.size(-1) > 2; + + for (int i = 0; i < partial_o.ndim() - 2; i++) { + TVM_FFI_ICHECK_EQ(partial_o.size(i), softmax_stats.size(i)) << "batch dimensions must match"; + } + + int64_t po_elem_size = get_element_size(partial_o); + TVM_FFI_ICHECK(kv_lora_rank * po_elem_size % 16 == 0) + << "partial_o last dim must be 16-byte aligned"; + + TVM_FFI_ICHECK_EQ(workspace.ndim(), 2) << "workspace must be 2D"; + TVM_FFI_ICHECK_EQ(workspace.size(0), cp_size) << "workspace first dim must equal cp_size"; + + int64_t entry_count = 1; + for (int i = 0; i < partial_o.ndim() - 2; i++) { + entry_count *= partial_o.size(i); + } + + // Build output shapes matching inputs + std::vector po_shape(partial_o.ndim()); + for (int i = 0; i < partial_o.ndim(); i++) { + po_shape[i] = partial_o.size(i); + } + std::vector ss_shape(softmax_stats.ndim()); + for (int i = 0; i < softmax_stats.ndim(); i++) { + ss_shape[i] = softmax_stats.size(i); + } + + Tensor partial_o_out = + alloc_tensor(tvm::ffi::Shape(po_shape), partial_o.dtype(), partial_o.device()); + Tensor softmax_stats_out = + alloc_tensor(tvm::ffi::Shape(ss_shape), softmax_stats.dtype(), softmax_stats.device()); + + int64_t ss_last = softmax_stats.size(-1); + int64_t ss_elem_size = get_element_size(softmax_stats); + + // stride(1) of a contiguous [entry_count, cp_size, D] view = D + // HelixFieldInfo.stride is in bytes: D * elem_size + tensorrt_llm::kernels::HelixAllToAllParams params; + + params.sendFields[0].dataPtr = static_cast(partial_o.data_ptr()); + params.sendFields[0].elementCount = static_cast(kv_lora_rank); + params.sendFields[0].elementSize = static_cast(po_elem_size); + params.sendFields[0].stride = static_cast(kv_lora_rank * po_elem_size); + + params.recvFields[0].dataPtr = static_cast(partial_o_out.data_ptr()); + params.recvFields[0].elementCount = static_cast(kv_lora_rank); + params.recvFields[0].elementSize = static_cast(po_elem_size); + params.recvFields[0].stride = static_cast(kv_lora_rank * po_elem_size); + + params.sendFields[1].dataPtr = static_cast(softmax_stats.data_ptr()); + params.sendFields[1].elementCount = static_cast(ss_last); + params.sendFields[1].elementSize = static_cast(ss_elem_size); + params.sendFields[1].stride = static_cast(ss_last * ss_elem_size); + + params.recvFields[1].dataPtr = static_cast(softmax_stats_out.data_ptr()); + params.recvFields[1].elementCount = static_cast(ss_last); + params.recvFields[1].elementSize = static_cast(ss_elem_size); + params.recvFields[1].stride = static_cast(ss_last * ss_elem_size); + + params.entryCount = static_cast(entry_count); + params.workspace = reinterpret_cast(workspace.data_ptr()); + params.workspaceStrideInU64 = workspace.stride(0); + + params.cpRank = static_cast(cp_rank); + params.cpSize = static_cast(cp_size); + params.channelCount = getEnvChannelCount(); + params.maxChannelCount = + tensorrt_llm::kernels::computeHelixMaxChannelCount(static_cast(cp_size)); + + auto stream = get_current_stream(); + tensorrt_llm::kernels::launchHelixAllToAll(params, allowVariableField1, enable_pdl, stream); + + return tvm::ffi::Tuple(partial_o_out, softmax_stats_out); +} + +} // namespace + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_dcp_workspace_size_per_rank, getDcpWorkspaceSizePerRank); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(initialize_dcp_workspace, initializeDcpWorkspaceOp); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(alltoall_dcp_native, alltoallDcpNativeOp); diff --git a/flashinfer/aot.py b/flashinfer/aot.py index dfb05150a8..d26d540796 100644 --- a/flashinfer/aot.py +++ b/flashinfer/aot.py @@ -543,6 +543,7 @@ def gen_all_modules( if add_comm: from .jit.comm import ( gen_comm_alltoall_module, + gen_dcp_alltoall_module, gen_moe_alltoall_module, gen_trtllm_comm_module, gen_trtllm_mnnvl_comm_module, @@ -554,6 +555,11 @@ def gen_all_modules( jit_specs.append(gen_trtllm_comm_module()) jit_specs.append(gen_trtllm_mnnvl_comm_module()) jit_specs.append(gen_moe_alltoall_module()) + # dcp_alltoall: kernel itself supports SM90+, but ptxas 12.6.0 has + # a known state-space inference bug on cp.async.bulk that aborts + # compilation. has_sm100 implies CUDA >= 12.8, which avoids the bug. + # SM90/SM12x users still get this via JIT. + jit_specs.append(gen_dcp_alltoall_module()) jit_specs.append(gen_vllm_comm_module()) if add_misc: diff --git a/flashinfer/comm/__init__.py b/flashinfer/comm/__init__.py index 5f186002dc..2be8e124b5 100644 --- a/flashinfer/comm/__init__.py +++ b/flashinfer/comm/__init__.py @@ -65,4 +65,12 @@ moe_a2a_wrap_payload_tensor_in_workspace as moe_a2a_wrap_payload_tensor_in_workspace, ) +# DCP A2A (Decode Context Parallel Attention Reduction) +from .dcp_alltoall import decode_cp_a2a_alltoall as decode_cp_a2a_alltoall +from .dcp_alltoall import ( + decode_cp_a2a_allocate_workspace as decode_cp_a2a_allocate_workspace, +) +from .dcp_alltoall import decode_cp_a2a_init_workspace as decode_cp_a2a_init_workspace +from .dcp_alltoall import decode_cp_a2a_workspace_size as decode_cp_a2a_workspace_size + # from .mnnvl import MnnvlMemory, MnnvlMoe, MoEAlltoallInfo diff --git a/flashinfer/comm/dcp_alltoall.py b/flashinfer/comm/dcp_alltoall.py new file mode 100644 index 0000000000..8a1eea1e8e --- /dev/null +++ b/flashinfer/comm/dcp_alltoall.py @@ -0,0 +1,254 @@ +""" +DCP All-to-All Operations for DCP Attention Reduction + +Provides the DCP LL128 FIFO-based all-to-all kernel for context-parallel +attention reduction. Uses SM90+ features (TMA, mbarrier). + +Usage protocol:: + + # 1. Query workspace size + ws_bytes = decode_cp_a2a_workspace_size(cp_size) + + # 2. Allocate workspace (MNNVL or plain device memory) + workspace = decode_cp_a2a_allocate_workspace(cp_size, cp_rank, mapping=mapping) + + # 3. Initialize workspace (synchronous — includes stream sync) + decode_cp_a2a_init_workspace(workspace, cp_rank, cp_size) + + # 4. Cross-rank barrier (REQUIRED before first alltoall) + dist.barrier(group) + + # 5. Run all-to-all + recv_o, recv_stats = decode_cp_a2a_alltoall( + partial_o, softmax_stats, workspace, cp_rank, cp_size + ) + +.. important:: + All ranks MUST complete ``decode_cp_a2a_init_workspace`` and execute a + cross-rank barrier before ANY rank calls ``decode_cp_a2a_alltoall``. + Failure to do so causes a deadlock on MNNVL workspaces. + +Tensor specifications: + +- ``partial_o``: ``[..., cp_size, D]`` — half or bfloat16, + ``D * element_size`` must be 16-byte aligned. +- ``softmax_stats``: ``[..., cp_size, S]`` — float32, ``S >= 2`` and even. + Batch dims must match ``partial_o``. +- ``workspace``: ``[cp_size, ws_elems_per_rank]`` — int64, from + :func:`decode_cp_a2a_allocate_workspace`. +""" + +import functools +import logging +from types import SimpleNamespace +from typing import Optional + +import torch + +from ..api_logging import flashinfer_api +from ..jit.comm import gen_dcp_alltoall_module +from ..utils import device_support_pdl, register_custom_op +from .mapping import Mapping +from .mnnvl import MnnvlConfig, MnnvlMemory + +logger = logging.getLogger(__name__) + + +# ─── Cached JIT Loader ─────────────────────────────────────────────────── + + +@functools.cache +def get_dcp_alltoall_module(): + """Build (once) and return the DCP A2A JIT module with custom op wrappers.""" + module = gen_dcp_alltoall_module().build_and_load() + + @register_custom_op( + "flashinfer::decode_cp_a2a_init_workspace", + mutates_args=("workspace",), + ) + def decode_cp_a2a_init_workspace( + workspace: torch.Tensor, + cp_rank: int, + cp_size: int, + ): + module.initialize_dcp_workspace(workspace, cp_rank, cp_size) + + @register_custom_op( + "flashinfer::decode_cp_a2a_alltoall", + mutates_args=("workspace",), + ) + def decode_cp_a2a_alltoall( + partial_o: torch.Tensor, + softmax_stats: torch.Tensor, + workspace: torch.Tensor, + cp_rank: int, + cp_size: int, + enable_pdl: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + return module.alltoall_dcp_native( + partial_o, softmax_stats, workspace, cp_rank, cp_size, enable_pdl + ) + + return SimpleNamespace( + get_workspace_size_per_rank=module.get_dcp_workspace_size_per_rank, + initialize_workspace=decode_cp_a2a_init_workspace, + alltoall=decode_cp_a2a_alltoall, + ) + + +# ─── Public API ─────────────────────────────────────────────────────────── + + +@flashinfer_api +def decode_cp_a2a_workspace_size(cp_size: int) -> int: + """Return the workspace size **in bytes** per rank for the given CP group size. + + Args: + cp_size: Context-parallel group size (number of ranks). + + Returns: + Workspace size in bytes per rank. + + Example:: + + >>> decode_cp_a2a_workspace_size(4) + 16778240 + """ + return get_dcp_alltoall_module().get_workspace_size_per_rank(cp_size) + + +@flashinfer_api +def decode_cp_a2a_allocate_workspace( + cp_size: int, + cp_rank: int, + *, + mapping: Optional[Mapping] = None, + mnnvl_config: Optional[MnnvlConfig] = None, +) -> torch.Tensor: + """Allocate a workspace tensor of shape ``[cp_size, ws_elems_per_rank]``. + + After allocation, call :func:`decode_cp_a2a_init_workspace` followed by a + cross-rank barrier before the first :func:`decode_cp_a2a_alltoall` call. + + Two allocation modes: + + - **MNNVL** (``mapping`` provided): Cross-rank visible GPU memory via + FlashInfer's ``MnnvlMemory``. Required for multi-node or when ranks + cannot see each other's device memory directly. + - **Plain device memory** (``mapping=None``): Standard ``torch.zeros`` + allocation. Sufficient for single-node with NVLink P2P. + + Args: + cp_size: Context-parallel group size. + cp_rank: This rank's position in the CP group. + mapping: Mapping object for MNNVL allocation. If provided, MNNVL is + used. The mapping must have ``cp_size`` set correctly. The + communicator is split using ``mapping.pp_rank``, ``mapping.cp_rank``, + and ``mapping.tp_rank``. + mnnvl_config: Configuration for the MNNVL communication backend. + Required when using MNNVL with ``torch.distributed`` (pass + ``MnnvlConfig(comm_backend=TorchDistBackend(group))``). + + Returns: + ``torch.int64`` tensor of shape ``[cp_size, ws_elems_per_rank]``. + """ + ws_bytes = decode_cp_a2a_workspace_size(cp_size) + + if mapping is not None: + MnnvlMemory.initialize() + if mnnvl_config: + MnnvlMemory.set_comm_from_config(mapping, mnnvl_config) + + mnnvl_mem = MnnvlMemory(mapping, ws_bytes) + workspace = mnnvl_mem.as_torch_strided_tensor(torch.int64) + workspace._mnnvl_mem = mnnvl_mem # prevent GC of MNNVL handle + logger.info( + "Rank %d: DCP MNNVL workspace allocated — shape=%s, stride=%s", + cp_rank, + list(workspace.shape), + list(workspace.stride()), + ) + return workspace + + ws_elems_per_rank = (ws_bytes + 7) // 8 + return torch.zeros(cp_size, ws_elems_per_rank, dtype=torch.int64, device="cuda") + + +@flashinfer_api +def decode_cp_a2a_init_workspace( + workspace: torch.Tensor, + cp_rank: int, + cp_size: int, +) -> None: + """Initialize the workspace FIFO buffers. Call once before the first alltoall. + + Resets the FIFO buffers in the **local** workspace row + (``workspace[cp_rank]``). This function is **synchronous**: when it + returns, the GPU memset is guaranteed to have completed. + + .. important:: + With MNNVL workspaces, **all ranks** must complete + ``decode_cp_a2a_init_workspace`` and execute a cross-rank barrier + (e.g. ``dist.barrier(group)``) before **any** rank calls + :func:`decode_cp_a2a_alltoall`. Without the barrier, a rank may + start writing to a peer's FIFO before that peer has finished + initializing → deadlock. + + Args: + workspace: ``[cp_size, ws_elems_per_rank]`` int64 tensor from + :func:`decode_cp_a2a_allocate_workspace`. + cp_rank: This rank's position in the CP group. + cp_size: Context-parallel group size. + """ + get_dcp_alltoall_module().initialize_workspace(workspace, cp_rank, cp_size) + # CRITICAL: The C++ op uses cudaMemsetAsync. Without this sync, a + # subsequent cross-GPU alltoall can race with the unfinished memset + # on MNNVL memory, causing a deadlock. + torch.cuda.current_stream().synchronize() + + +@flashinfer_api +def decode_cp_a2a_alltoall( + partial_o: torch.Tensor, + softmax_stats: torch.Tensor, + workspace: torch.Tensor, + cp_rank: int, + cp_size: int, + enable_pdl: Optional[bool] = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Perform the DCP all-to-all exchange. + + Each rank sends its ``partial_o[..., peer, :]`` slice to the + corresponding peer and receives all peers' contributions into the + output tensors. + + Args: + partial_o: ``[..., cp_size, D]`` — half or bfloat16. + ``D * element_size`` must be 16-byte aligned. + softmax_stats: ``[..., cp_size, S]`` — float32, ``S >= 2`` and even. + Batch dimensions must match ``partial_o``. + workspace: ``[cp_size, ws_elems_per_rank]`` int64 tensor from + :func:`decode_cp_a2a_allocate_workspace`, already initialized. + cp_rank: This rank's position in the CP group. + cp_size: Context-parallel group size. + enable_pdl: Enable Programmatic Dependent Launch (SM90+). + Defaults to ``True`` on SM90+ GPUs, ``False`` otherwise. + + Returns: + ``(partial_o_out, softmax_stats_out)`` with the same shapes and + dtypes as the inputs. Each output contains the gathered data from + all peers for this rank. + """ + if enable_pdl is None: + enable_pdl = device_support_pdl(partial_o.device) + return get_dcp_alltoall_module().alltoall( + partial_o, softmax_stats, workspace, cp_rank, cp_size, enable_pdl + ) + + +__all__ = [ + "decode_cp_a2a_workspace_size", + "decode_cp_a2a_allocate_workspace", + "decode_cp_a2a_init_workspace", + "decode_cp_a2a_alltoall", +] diff --git a/flashinfer/jit/__init__.py b/flashinfer/jit/__init__.py index 7f36a31474..8378e0ab74 100644 --- a/flashinfer/jit/__init__.py +++ b/flashinfer/jit/__init__.py @@ -82,6 +82,7 @@ from .comm import gen_trtllm_comm_module as gen_trtllm_comm_module from .comm import gen_vllm_comm_module as gen_vllm_comm_module from .comm import gen_moe_alltoall_module as gen_moe_alltoall_module +from .comm import gen_dcp_alltoall_module as gen_dcp_alltoall_module from .dsv3_optimizations import ( gen_dsv3_router_gemm_module as gen_dsv3_router_gemm_module, ) diff --git a/flashinfer/jit/comm.py b/flashinfer/jit/comm.py index 2f270582f3..5dbe7eb8f1 100644 --- a/flashinfer/jit/comm.py +++ b/flashinfer/jit/comm.py @@ -227,3 +227,35 @@ def gen_moe_alltoall_module() -> JitSpec: str(jit_env.FLASHINFER_CSRC_DIR / "nv_internal" / "include"), ], ) + + +def gen_dcp_alltoall_module() -> JitSpec: + nvcc_flags = current_compilation_context.get_nvcc_flags_list( + supported_major_versions=[9, 10, 11, 12] + ) + return gen_jit_spec( + "dcp_alltoall", + [ + jit_env.FLASHINFER_CSRC_DIR / "trtllm_dcp_alltoall.cu", + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal" + / "tensorrt_llm" + / "kernels" + / "helixAllToAll.cu", + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal" + / "cpp" + / "common" + / "envUtils.cpp", + jit_env.FLASHINFER_CSRC_DIR + / "nv_internal" + / "cpp" + / "common" + / "tllmException.cpp", + ], + extra_include_paths=[ + str(jit_env.FLASHINFER_CSRC_DIR / "nv_internal"), + str(jit_env.FLASHINFER_CSRC_DIR / "nv_internal" / "include"), + ], + extra_cuda_cflags=nvcc_flags, + ) diff --git a/tests/comm/test_dcp_alltoall.py b/tests/comm/test_dcp_alltoall.py new file mode 100644 index 0000000000..847e6f5d29 --- /dev/null +++ b/tests/comm/test_dcp_alltoall.py @@ -0,0 +1,404 @@ +# Copyright (c) 2024 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for flashinfer.comm.dcp_alltoall — DCP LL128 FIFO All-to-All. + +Single-GPU multi-rank pattern: simulates cp_size ranks on one GPU using +separate CUDA streams for the alltoall phase. All ranks share a single +workspace tensor of shape [cp_size, ws_elems_per_rank]. + +Run: python -m pytest tests/comm/test_dcp_alltoall.py -v -s +""" + +import pytest +import torch + +from flashinfer.comm import ( + decode_cp_a2a_alltoall, + decode_cp_a2a_allocate_workspace, + decode_cp_a2a_init_workspace, + decode_cp_a2a_workspace_size, +) + + +# ─── SM90/SM100 gate ───────────────────────────────────────────────────── + + +def _dcp_alltoall_supported() -> bool: + try: + if not torch.cuda.is_available(): + return False + major, _ = torch.cuda.get_device_capability(0) + # DCP A2A requires TMA + mbarrier + PDL (SM90 baseline); the kernel + # builds for SM90 / SM10x / SM11x / SM12x. + return major in (9, 10, 11, 12) + except Exception: + return False + + +pytestmark = pytest.mark.skipif( + not _dcp_alltoall_supported(), + reason="Requires SM90+ GPU (Hopper or Blackwell family)", +) + + +@pytest.fixture(autouse=True, scope="session") +def setup_test_environment(): + """Set torch seed for deterministic tests.""" + torch.manual_seed(0xA2A) + yield + + +# ─── Helper ────────────────────────────────────────────────────────────── + + +def _to_torch(t): + """Convert a tvm_ffi.core.Tensor (or any DLPack object) to torch.Tensor.""" + if isinstance(t, torch.Tensor): + return t + return torch.from_dlpack(t) + + +def _run_single_gpu_alltoall(cp_size, batch_size, head_dim, stats_dim, dtype): + """Simulate cp_size ranks on one GPU and return (inputs, outputs, workspace). + + 1. Allocate shared workspace (plain device memory). + 2. Generate random partial_o [B, cp_size, D] and softmax_stats [B, cp_size, S] + per rank. + 3. Init workspace for all ranks (default stream, sequential). + 4. torch.cuda.synchronize() — cross-rank barrier. + 5. Alltoall for each rank on separate CUDA streams. + 6. Sync all streams and return results. + """ + torch.cuda.set_device(0) + + workspace = decode_cp_a2a_allocate_workspace(cp_size, cp_rank=0) + + all_partial_o = [] + all_softmax_stats = [] + for _ in range(cp_size): + po = torch.randn(batch_size, cp_size, head_dim, dtype=dtype, device="cuda") + ss = torch.randn( + batch_size, cp_size, stats_dim, dtype=torch.float32, device="cuda" + ) + all_partial_o.append(po) + all_softmax_stats.append(ss) + + for r in range(cp_size): + decode_cp_a2a_init_workspace(workspace, r, cp_size) + + torch.cuda.synchronize() + + streams = [torch.cuda.Stream() for _ in range(cp_size)] + recv_o = [None] * cp_size + recv_s = [None] * cp_size + + for r in range(cp_size): + with torch.cuda.stream(streams[r]): + o, s = decode_cp_a2a_alltoall( + all_partial_o[r], + all_softmax_stats[r], + workspace, + r, + cp_size, + ) + recv_o[r] = _to_torch(o) + recv_s[r] = _to_torch(s) + + for stream in streams: + stream.synchronize() + torch.cuda.synchronize() + + return all_partial_o, all_softmax_stats, recv_o, recv_s, workspace + + +def _verify_transpose(cp_size, all_partial_o, all_softmax_stats, recv_o, recv_s): + """Assert the transpose property for all (rank, peer) pairs. + + recv_o[r][..., peer, :] == all_partial_o[peer][..., r, :] + recv_s[r][..., peer, :] == all_softmax_stats[peer][..., r, :] + """ + for r in range(cp_size): + for peer in range(cp_size): + torch.testing.assert_close( + recv_o[r][..., peer, :], + all_partial_o[peer][..., r, :], + atol=0, + rtol=0, + ) + torch.testing.assert_close( + recv_s[r][..., peer, :], + all_softmax_stats[peer][..., r, :], + atol=0, + rtol=0, + ) + + +# ─── Workspace Lifecycle Tests ─────────────────────────────────────────── + + +class TestWorkspaceLifecycle: + """Verify workspace sizing, allocation, and initialization.""" + + def test_workspace_size_positive(self): + for cp_size in [2, 4, 8]: + ws = decode_cp_a2a_workspace_size(cp_size) + assert isinstance(ws, int) + assert ws > 0, f"cp_size={cp_size}: workspace_size should be positive" + + def test_workspace_size_monotonic(self): + ws2 = decode_cp_a2a_workspace_size(2) + ws4 = decode_cp_a2a_workspace_size(4) + ws8 = decode_cp_a2a_workspace_size(8) + assert ws4 > ws2, "ws(4) should be > ws(2)" + assert ws8 > ws4, "ws(8) should be > ws(4)" + + def test_allocate_returns_correct_shape_and_dtype(self): + for cp_size in [2, 4]: + ws_bytes = decode_cp_a2a_workspace_size(cp_size) + workspace = decode_cp_a2a_allocate_workspace(cp_size, cp_rank=0) + assert workspace.dtype == torch.int64 + assert workspace.shape[0] == cp_size + assert workspace.shape[1] == (ws_bytes + 7) // 8 + + def test_init_workspace_does_not_hang(self): + for cp_size in [2, 4]: + workspace = decode_cp_a2a_allocate_workspace(cp_size, cp_rank=0) + for r in range(cp_size): + decode_cp_a2a_init_workspace(workspace, r, cp_size) + torch.cuda.synchronize() + + +# ─── Correctness Tests (Single-GPU Multi-Rank) ────────────────────────── + +CORRECTNESS_PARAMS = [ + # cp_size=2 (minimum) + pytest.param(2, 1, 128, 2, torch.bfloat16, id="cp2-B1-D128-S2-bf16"), + pytest.param(2, 16, 128, 2, torch.bfloat16, id="cp2-B16-D128-S2-bf16"), + pytest.param(2, 128, 128, 2, torch.bfloat16, id="cp2-B128-D128-S2-bf16"), + pytest.param(2, 16, 256, 2, torch.bfloat16, id="cp2-B16-D256-S2-bf16"), + pytest.param(2, 16, 128, 4, torch.bfloat16, id="cp2-B16-D128-S4-bf16"), + pytest.param(2, 16, 128, 2, torch.float16, id="cp2-B16-D128-S2-fp16"), + # cp_size=4 (common single-node) + pytest.param(4, 1, 128, 2, torch.bfloat16, id="cp4-B1-D128-S2-bf16"), + pytest.param(4, 16, 128, 2, torch.bfloat16, id="cp4-B16-D128-S2-bf16"), + pytest.param(4, 128, 128, 2, torch.bfloat16, id="cp4-B128-D128-S2-bf16"), + pytest.param(4, 16, 256, 4, torch.bfloat16, id="cp4-B16-D256-S4-bf16"), + pytest.param(4, 16, 128, 2, torch.float16, id="cp4-B16-D128-S2-fp16"), + pytest.param(4, 16, 256, 2, torch.float16, id="cp4-B16-D256-S2-fp16"), + # cp_size=8 (full single-node, e.g. 8xH200) + pytest.param(8, 1, 128, 2, torch.bfloat16, id="cp8-B1-D128-S2-bf16"), + pytest.param(8, 16, 128, 2, torch.bfloat16, id="cp8-B16-D128-S2-bf16"), + pytest.param(8, 64, 128, 2, torch.bfloat16, id="cp8-B64-D128-S2-bf16"), + pytest.param(8, 16, 256, 4, torch.bfloat16, id="cp8-B16-D256-S4-bf16"), + pytest.param(8, 16, 128, 2, torch.float16, id="cp8-B16-D128-S2-fp16"), +] + + +@pytest.mark.parametrize( + "cp_size,batch_size,head_dim,stats_dim,dtype", + CORRECTNESS_PARAMS, +) +def test_alltoall_correctness(cp_size, batch_size, head_dim, stats_dim, dtype): + """Verify the transpose property: recv[r][.., peer, :] == input[peer][.., r, :].""" + all_po, all_ss, recv_o, recv_s, _ = _run_single_gpu_alltoall( + cp_size, batch_size, head_dim, stats_dim, dtype + ) + _verify_transpose(cp_size, all_po, all_ss, recv_o, recv_s) + + +# ─── FIFO Reuse (Repeated Alltoall without re-init) ───────────────────── + +FIFO_REUSE_PARAMS = [ + pytest.param(2, 16, 128, 2, torch.bfloat16, 3, id="cp2-3rounds"), + pytest.param(4, 16, 128, 2, torch.bfloat16, 3, id="cp4-3rounds"), +] + + +@pytest.mark.parametrize( + "cp_size,batch_size,head_dim,stats_dim,dtype,num_rounds", + FIFO_REUSE_PARAMS, +) +def test_repeated_alltoall(cp_size, batch_size, head_dim, stats_dim, dtype, num_rounds): + """Multiple alltoall calls on the same workspace without re-init (FIFO reuse).""" + torch.cuda.set_device(0) + + workspace = decode_cp_a2a_allocate_workspace(cp_size, cp_rank=0) + + for r in range(cp_size): + decode_cp_a2a_init_workspace(workspace, r, cp_size) + torch.cuda.synchronize() + + for _round_idx in range(num_rounds): + all_po = [ + torch.randn(batch_size, cp_size, head_dim, dtype=dtype, device="cuda") + for _ in range(cp_size) + ] + all_ss = [ + torch.randn( + batch_size, cp_size, stats_dim, dtype=torch.float32, device="cuda" + ) + for _ in range(cp_size) + ] + + streams = [torch.cuda.Stream() for _ in range(cp_size)] + recv_o = [None] * cp_size + recv_s = [None] * cp_size + + for r in range(cp_size): + with torch.cuda.stream(streams[r]): + o, s = decode_cp_a2a_alltoall( + all_po[r], all_ss[r], workspace, r, cp_size + ) + recv_o[r] = _to_torch(o) + recv_s[r] = _to_torch(s) + + for stream in streams: + stream.synchronize() + torch.cuda.synchronize() + + _verify_transpose(cp_size, all_po, all_ss, recv_o, recv_s) + + +# ─── Edge Cases ────────────────────────────────────────────────────────── + + +class TestEdgeCases: + """Edge cases: cp_size=1, batch_size=0.""" + + def test_cp_size_1_is_identity(self): + """cp_size=1: output should equal input (no peers to exchange with).""" + cp_size, batch_size, head_dim, stats_dim = 1, 16, 128, 2 + dtype = torch.bfloat16 + torch.cuda.set_device(0) + + workspace = decode_cp_a2a_allocate_workspace(cp_size, cp_rank=0) + po = torch.randn(batch_size, cp_size, head_dim, dtype=dtype, device="cuda") + ss = torch.randn( + batch_size, cp_size, stats_dim, dtype=torch.float32, device="cuda" + ) + + decode_cp_a2a_init_workspace(workspace, 0, cp_size) + torch.cuda.synchronize() + + o, s = decode_cp_a2a_alltoall(po, ss, workspace, 0, cp_size) + o = _to_torch(o) + s = _to_torch(s) + + torch.testing.assert_close(o, po, atol=0, rtol=0) + torch.testing.assert_close(s, ss, atol=0, rtol=0) + + def test_batch_size_0(self): + """batch_size=0: should not crash (zero entries → no work).""" + cp_size, batch_size, head_dim, stats_dim = 2, 0, 128, 2 + dtype = torch.bfloat16 + torch.cuda.set_device(0) + + workspace = decode_cp_a2a_allocate_workspace(cp_size, cp_rank=0) + po = torch.randn(batch_size, cp_size, head_dim, dtype=dtype, device="cuda") + ss = torch.randn( + batch_size, cp_size, stats_dim, dtype=torch.float32, device="cuda" + ) + + for r in range(cp_size): + decode_cp_a2a_init_workspace(workspace, r, cp_size) + torch.cuda.synchronize() + + # Should not crash; output shape should match input shape + o, s = decode_cp_a2a_alltoall(po, ss, workspace, 0, cp_size) + o = _to_torch(o) + s = _to_torch(s) + assert o.shape == po.shape + assert s.shape == ss.shape + + +# ─── Input Validation (Error Handling) ─────────────────────────────────── + + +class TestInputValidation: + """Verify that invalid inputs are rejected with errors, not silent corruption.""" + + def test_wrong_dtype_float64(self): + """partial_o with float64 should be rejected.""" + cp_size = 2 + workspace = decode_cp_a2a_allocate_workspace(cp_size, cp_rank=0) + for r in range(cp_size): + decode_cp_a2a_init_workspace(workspace, r, cp_size) + torch.cuda.synchronize() + + po = torch.randn(16, cp_size, 128, dtype=torch.float64, device="cuda") + ss = torch.randn(16, cp_size, 2, dtype=torch.float32, device="cuda") + + with pytest.raises(RuntimeError): + decode_cp_a2a_alltoall(po, ss, workspace, 0, cp_size) + + def test_wrong_dtype_float32(self): + """partial_o with float32 should be rejected (must be half/bfloat16).""" + cp_size = 2 + workspace = decode_cp_a2a_allocate_workspace(cp_size, cp_rank=0) + for r in range(cp_size): + decode_cp_a2a_init_workspace(workspace, r, cp_size) + torch.cuda.synchronize() + + po = torch.randn(16, cp_size, 128, dtype=torch.float32, device="cuda") + ss = torch.randn(16, cp_size, 2, dtype=torch.float32, device="cuda") + + with pytest.raises(RuntimeError): + decode_cp_a2a_alltoall(po, ss, workspace, 0, cp_size) + + def test_stats_dim_1_odd_alignment(self): + """stats_dim=1 violates 'even and >= 2' constraint — should error.""" + cp_size = 2 + workspace = decode_cp_a2a_allocate_workspace(cp_size, cp_rank=0) + for r in range(cp_size): + decode_cp_a2a_init_workspace(workspace, r, cp_size) + torch.cuda.synchronize() + + po = torch.randn(16, cp_size, 128, dtype=torch.bfloat16, device="cuda") + ss = torch.randn(16, cp_size, 1, dtype=torch.float32, device="cuda") + + with pytest.raises(RuntimeError): + decode_cp_a2a_alltoall(po, ss, workspace, 0, cp_size) + + def test_mismatched_batch_dims(self): + """partial_o and softmax_stats with different batch sizes should error.""" + cp_size = 2 + workspace = decode_cp_a2a_allocate_workspace(cp_size, cp_rank=0) + for r in range(cp_size): + decode_cp_a2a_init_workspace(workspace, r, cp_size) + torch.cuda.synchronize() + + po = torch.randn(16, cp_size, 128, dtype=torch.bfloat16, device="cuda") + ss = torch.randn(32, cp_size, 2, dtype=torch.float32, device="cuda") + + with pytest.raises(RuntimeError): + decode_cp_a2a_alltoall(po, ss, workspace, 0, cp_size) + + def test_wrong_stats_dtype(self): + """softmax_stats with half instead of float32 should error.""" + cp_size = 2 + workspace = decode_cp_a2a_allocate_workspace(cp_size, cp_rank=0) + for r in range(cp_size): + decode_cp_a2a_init_workspace(workspace, r, cp_size) + torch.cuda.synchronize() + + po = torch.randn(16, cp_size, 128, dtype=torch.bfloat16, device="cuda") + ss = torch.randn(16, cp_size, 2, dtype=torch.float16, device="cuda") + + with pytest.raises(RuntimeError): + decode_cp_a2a_alltoall(po, ss, workspace, 0, cp_size) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/tests/comm/test_mnnvl_dcp_alltoall.py b/tests/comm/test_mnnvl_dcp_alltoall.py new file mode 100644 index 0000000000..bd020b9566 --- /dev/null +++ b/tests/comm/test_mnnvl_dcp_alltoall.py @@ -0,0 +1,350 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Multi-GPU tests for DCP LL128 FIFO All-to-All with MNNVL workspace. + +Complements test_dcp_alltoall.py (single-GPU simulation) by running +real multi-process, multi-GPU alltoall via MPI. This catches bugs +that single-GPU tests cannot: + - MNNVL workspace allocation and communicator grouping + - Cross-GPU memory visibility (LL128 FIFO writes to peer memory) + - Workspace shape [cp_size, ws_elems] with real multi-rank segments + +Run: + mpirun -np 2 pytest tests/comm/test_mnnvl_dcp_alltoall.py -v -s + mpirun -np 4 pytest tests/comm/test_mnnvl_dcp_alltoall.py -v -s +""" + +import socket + +import pynvml +import pytest +import torch + +from flashinfer.comm import ( + decode_cp_a2a_alltoall, + decode_cp_a2a_allocate_workspace, + decode_cp_a2a_init_workspace, + decode_cp_a2a_workspace_size, +) +from flashinfer.comm.mapping import Mapping +from flashinfer.comm.mnnvl import MnnvlMemory, MpiComm + +from .conftest import mnnvl_available + +pynvml.nvmlInit() + + +# ─── SM90+ gate ────────────────────────────────────────────────────────── + + +def _dcp_alltoall_supported() -> bool: + if not torch.cuda.is_available(): + return False + from flashinfer.utils import get_compute_capability + + major, _ = get_compute_capability(torch.device("cuda")) + return major in (9, 10, 11, 12) + + +def _mpi4py_available() -> bool: + try: + import mpi4py # noqa: F401 + + return True + except ImportError: + return False + + +pytestmark = [ + pytest.mark.skipif( + not _dcp_alltoall_supported(), + reason="Requires SM90+ GPU (Hopper or Blackwell family)", + ), + pytest.mark.skipif( + not mnnvl_available(), + reason="MNNVL not supported on this platform or container lacks SYS_PTRACE", + ), + pytest.mark.skipif( + not _mpi4py_available(), + reason="mpi4py not installed (run with: mpirun -np N pytest ...)", + ), +] + + +# ─── Helper ────────────────────────────────────────────────────────────── + + +def _to_torch(t): + """Convert a tvm_ffi.core.Tensor (or any DLPack object) to torch.Tensor.""" + if isinstance(t, torch.Tensor): + return t + return torch.from_dlpack(t) + + +def _setup_rank(): + """Initialize MPI rank and CUDA device. Returns (rank, world_size, comm).""" + comm = MpiComm() + rank = comm.Get_rank() + world_size = comm.Get_size() + + # Get local rank from hostname + hostname = socket.gethostname() + all_hostnames = comm.allgather(hostname) + local_ranks_before_me = sum(1 for i in range(rank) if all_hostnames[i] == hostname) + local_rank = local_ranks_before_me + torch.cuda.set_device(local_rank) + + return rank, world_size, comm + + +# ─── Module-level MNNVL workspace (allocated once, reused across tests) ── + +# Guard module-level allocation: pytestmark skipif conditions are not +# enforced during module import (collection phase). If we allocate +# unconditionally, CI environments without SYS_PTRACE will fail at +# collection time instead of gracefully skipping. +if _dcp_alltoall_supported() and mnnvl_available() and _mpi4py_available(): + _rank, _cp_size, _comm = _setup_rank() + + def _allocate_mnnvl_workspace_once(): + """Allocate MNNVL workspace once at module level. + + MnnvlMemory uses a global bump allocator that doesn't support + individual frees. Allocating per-test causes segfaults when + workspace tensors from previous tests get GC'd. So we allocate + once and reuse. + """ + MnnvlMemory.initialize() + MnnvlMemory.comm = _comm + + mapping = Mapping( + world_size=_cp_size, + rank=_rank, + cp_size=_cp_size, + tp_size=1, + pp_size=1, + ) + + ws_bytes = decode_cp_a2a_workspace_size(_cp_size) + mnnvl_mem = MnnvlMemory(mapping, ws_bytes) + workspace = mnnvl_mem.as_torch_strided_tensor(torch.int64) + workspace._mnnvl_mem = mnnvl_mem # prevent GC + return workspace + + _mnnvl_workspace = _allocate_mnnvl_workspace_once() +else: + _rank, _cp_size, _comm = 0, 1, None + _mnnvl_workspace = None + + +# ─── Tests ─────────────────────────────────────────────────────────────── + + +class TestMnnvlDcpWorkspace: + """Test MNNVL workspace allocation for DCP A2A.""" + + @pytest.fixture(autouse=True) + def setup(self): + torch.manual_seed(0xA2A + _rank) + yield + + def test_workspace_shape(self): + """MNNVL workspace must have shape [cp_size, ws_elems_per_rank].""" + try: + assert _mnnvl_workspace.shape[0] == _cp_size, ( + f"Expected workspace.shape[0] == {_cp_size}, got {_mnnvl_workspace.shape[0]}" + ) + + ws_bytes = decode_cp_a2a_workspace_size(_cp_size) + expected_elems = (ws_bytes + 7) // 8 # int64 elements + assert _mnnvl_workspace.shape[1] == expected_elems + assert _mnnvl_workspace.dtype == torch.int64 + finally: + _comm.Barrier() + + def test_workspace_cross_rank_visible(self): + """Each rank can write to its own segment and peers can read it.""" + try: + # Each rank writes a unique pattern to its own workspace segment + pattern = torch.full_like(_mnnvl_workspace[_rank], fill_value=_rank + 1) + _mnnvl_workspace[_rank].copy_(pattern) + torch.cuda.synchronize() + _comm.Barrier() + + # Each rank reads all segments and verifies the pattern + for peer in range(_cp_size): + expected = peer + 1 + actual = _mnnvl_workspace[peer][0].item() + assert actual == expected, ( + f"Rank {_rank}: workspace[{peer}][0] = {actual}, expected {expected}" + ) + finally: + _comm.Barrier() + + +class TestMnnvlDcpAlltoall: + """Multi-GPU correctness tests for DCP A2A alltoall.""" + + @pytest.fixture(autouse=True) + def setup(self): + torch.manual_seed(0xA2A) + yield + + def _run_alltoall(self, batch_size, head_dim, stats_dim, dtype): + """Run DCP A2A alltoall on real multi-GPU and verify correctness. + + The transpose property must hold: + recv_o[rank][.., peer, :] == send_o[peer][.., rank, :] + recv_s[rank][.., peer, :] == send_s[peer][.., rank, :] + """ + try: + workspace = _mnnvl_workspace + + decode_cp_a2a_init_workspace(workspace, _rank, _cp_size) + torch.cuda.synchronize() + _comm.Barrier() + + # Generate input with deterministic seed per rank + torch.manual_seed(0xA2A + _rank) + partial_o = torch.randn( + batch_size, _cp_size, head_dim, dtype=dtype, device="cuda" + ) + softmax_stats = torch.randn( + batch_size, _cp_size, stats_dim, dtype=torch.float32, device="cuda" + ) + + # Run alltoall + recv_o, recv_s = decode_cp_a2a_alltoall( + partial_o, softmax_stats, workspace, _rank, _cp_size + ) + recv_o = _to_torch(recv_o) + recv_s = _to_torch(recv_s) + torch.cuda.synchronize() + _comm.Barrier() + + # Gather all inputs to all ranks for verification + all_partial_o = _comm.allgather(partial_o.cpu()) + all_softmax_stats = _comm.allgather(softmax_stats.cpu()) + + # Verify transpose property + for peer in range(_cp_size): + expected_o = all_partial_o[peer][..., _rank, :].cuda() + expected_s = all_softmax_stats[peer][..., _rank, :].cuda() + + torch.testing.assert_close( + recv_o[..., peer, :], + expected_o, + atol=0, + rtol=0, + ) + torch.testing.assert_close( + recv_s[..., peer, :], + expected_s, + atol=0, + rtol=0, + ) + finally: + _comm.Barrier() + + @pytest.mark.parametrize( + "batch_size,head_dim,stats_dim,dtype", + [ + pytest.param(1, 128, 2, torch.bfloat16, id="B1-D128-S2-bf16"), + pytest.param(16, 128, 2, torch.bfloat16, id="B16-D128-S2-bf16"), + pytest.param(128, 128, 2, torch.bfloat16, id="B128-D128-S2-bf16"), + pytest.param(16, 256, 4, torch.bfloat16, id="B16-D256-S4-bf16"), + pytest.param(16, 128, 2, torch.float16, id="B16-D128-S2-fp16"), + ], + ) + def test_alltoall_correctness(self, batch_size, head_dim, stats_dim, dtype): + """Verify transpose property across real GPUs.""" + self._run_alltoall(batch_size, head_dim, stats_dim, dtype) + + def test_repeated_alltoall(self): + """Multiple alltoall calls on the same workspace (FIFO reuse).""" + try: + workspace = _mnnvl_workspace + + decode_cp_a2a_init_workspace(workspace, _rank, _cp_size) + torch.cuda.synchronize() + _comm.Barrier() + + for round_idx in range(3): + torch.manual_seed(0xA2A + _rank * 100 + round_idx) + partial_o = torch.randn( + 16, _cp_size, 128, dtype=torch.bfloat16, device="cuda" + ) + softmax_stats = torch.randn( + 16, _cp_size, 2, dtype=torch.float32, device="cuda" + ) + + recv_o, recv_s = decode_cp_a2a_alltoall( + partial_o, softmax_stats, workspace, _rank, _cp_size + ) + recv_o = _to_torch(recv_o) + recv_s = _to_torch(recv_s) + torch.cuda.synchronize() + _comm.Barrier() + + all_partial_o = _comm.allgather(partial_o.cpu()) + all_softmax_stats = _comm.allgather(softmax_stats.cpu()) + + for peer in range(_cp_size): + torch.testing.assert_close( + recv_o[..., peer, :], + all_partial_o[peer][..., _rank, :].cuda(), + atol=0, + rtol=0, + ) + torch.testing.assert_close( + recv_s[..., peer, :], + all_softmax_stats[peer][..., _rank, :].cuda(), + atol=0, + rtol=0, + ) + finally: + _comm.Barrier() + + +class TestMnnvlDcpDeviceMemoryFallback: + """Test that non-MNNVL (device memory) path also works multi-GPU. + + Uses decode_cp_a2a_allocate_workspace without MNNVL mapping. This only + works when all ranks are on the same GPU (single-GPU simulation) + or with IPC. Included here to verify the workspace API contract. + """ + + @pytest.fixture(autouse=True) + def setup(self): + torch.manual_seed(0xA2A) + yield + + def test_device_workspace_shape(self): + """Device workspace has correct shape [cp_size, ws_elems].""" + try: + workspace = decode_cp_a2a_allocate_workspace(_cp_size, cp_rank=_rank) + assert workspace.shape[0] == _cp_size + + ws_bytes = decode_cp_a2a_workspace_size(_cp_size) + expected_elems = (ws_bytes + 7) // 8 + assert workspace.shape[1] == expected_elems + assert workspace.dtype == torch.int64 + finally: + _comm.Barrier() + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"])