diff --git a/csrc/nv_internal/tensorrt_llm/common/envUtils.h b/csrc/nv_internal/tensorrt_llm/common/envUtils.h index cdbdd8c4146..253b9613478 100644 --- a/csrc/nv_internal/tensorrt_llm/common/envUtils.h +++ b/csrc/nv_internal/tensorrt_llm/common/envUtils.h @@ -19,6 +19,10 @@ #include #include #include +#include + +#include +#include "flashinfer/exception.h" namespace tensorrt_llm::common { // Useful when you want to inject some debug code controllable with env var. @@ -101,4 +105,23 @@ int getEnvMoeA2ADispatchBlockSize(); // Block size (threads per block) for MoE A2A Combine kernels (default 256 if unset or invalid) int getEnvMoeA2ACombineBlockSize(); +template +inline void launchWithPdlWhenEnabled(char const* name, bool enable_pdl, KernelFn kernelFn, + dim3 grid, dim3 block, size_t dynamicShmSize, cudaStream_t stream, Args&&... args) +{ + cudaLaunchConfig_t kernelConfig; + kernelConfig.gridDim = grid; + kernelConfig.blockDim = block; + kernelConfig.dynamicSmemBytes = dynamicShmSize; + kernelConfig.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; + kernelConfig.attrs = attrs; + kernelConfig.numAttrs = 1; + cudaError_t e = cudaLaunchKernelEx(&kernelConfig, kernelFn, std::forward(args)...); + FLASHINFER_CHECK(e == cudaSuccess, "cudaLaunchKernelEx (", name, ") failed: ", + cudaGetErrorString(e)); +} + } // namespace tensorrt_llm::common diff --git a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu b/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu index c0f0e918966..8cad41b5413 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu +++ b/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu @@ -14,6 +14,7 @@ * limitations under the License. */ #include +#include #include #include @@ -27,6 +28,8 @@ namespace tensorrt_llm::kernels::moe_alltoall { +using tensorrt_llm::common::launchWithPdlWhenEnabled; + #define ENABLE_DEBUG_PRINT 0 #define DISABLE_SYNC_FOR_PROFILING 0 @@ -114,6 +117,11 @@ __host__ __device__ inline T ceilDiv(T m, T n) { __VA_ARGS__; \ break; \ } \ + case nvinfer1::DataType::kFP8: { \ + using TYPE = __nv_fp8_e4m3; \ + __VA_ARGS__; \ + break; \ + } \ default: { \ FLASHINFER_CHECK(false, "Unsupported dtype for moe_a2a_combine"); \ } \ @@ -282,7 +290,11 @@ __device__ void vectorized_dispatch(uint8_t const* src_ptr, int bytes_per_token, } __global__ void moeA2APrepareDispatchKernel(int* send_counters, int* local_token_counter, - int ep_size, uint32_t* flag_val_ptr) { + int ep_size, uint32_t* flag_val_ptr, bool enable_pdl) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + if (enable_pdl) cudaGridDependencySynchronize(); + cudaTriggerProgrammaticLaunchCompletion(); +#endif int idx = blockIdx.x * blockDim.x + threadIdx.x; // Zero send_counters if (idx < ep_size) { @@ -310,7 +322,7 @@ __global__ void moeA2ADispatchKernel( const DispatchKernelPointers ptrs, // Struct containing all kernel pointers int num_payloads, // Number of payloads int max_tokens_per_rank, // Maximum tokens per rank - int local_num_tokens, int rank_id, int ep_size, int num_experts_per_rank) { + int local_num_tokens, int rank_id, int ep_size, int num_experts_per_rank, bool enable_pdl) { int thread_idx = ThreadingPolicy::offset(); int local_token_idx = ThreadingPolicy::token_idx(); @@ -319,10 +331,17 @@ __global__ void moeA2ADispatchKernel( // we need to keep the threads where local_token_idx == 0 alive to participate in the // synchronization. Other threads should return. if (local_token_idx > 0) return; +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + if (enable_pdl) cudaGridDependencySynchronize(); +#endif } else { // Threads that do not have a token to process should return. if (local_token_idx >= local_num_tokens) return; +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + if (enable_pdl) cudaGridDependencySynchronize(); +#endif + // Prepare per-policy shared-memory tiles for this token extern __shared__ int smem[]; int* smem_topk_target_ranks; @@ -393,6 +412,10 @@ __global__ void moeA2ADispatchKernel( ThreadingPolicy::sync(); } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + bool is_first_warp = threadIdx.x / warpSize == 0; if (is_first_warp) { int lane_id = threadIdx.x % warpSize; @@ -468,8 +491,9 @@ __global__ void moeA2ADispatchKernel( } void moe_a2a_prepare_dispatch_launch(MoeA2ADispatchParams const& params) { - moeA2APrepareDispatchKernel<<<1, params.ep_size, 0, params.stream>>>( - params.send_counters, params.local_token_counter, params.ep_size, params.flag_val); + launchWithPdlWhenEnabled("moeA2APrepareDispatchKernel", params.enable_pdl, + moeA2APrepareDispatchKernel, 1, params.ep_size, 0, params.stream, params.send_counters, + params.local_token_counter, params.ep_size, params.flag_val, params.enable_pdl); } // ============================================================================ @@ -526,12 +550,13 @@ void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params) { grid_size = 1; } int shared_bytes = 2 * params.top_k * (int)sizeof(int); - SWITCH_TOP_K(params.top_k, TOP_K, - moeA2ADispatchKernel - <<>>( - params.token_selected_experts, kernel_ptrs, params.num_payloads, - params.max_tokens_per_rank, params.local_num_tokens, params.ep_rank, - params.ep_size, params.num_experts_per_rank)) + SWITCH_TOP_K(params.top_k, TOP_K, { + auto kernel_fn = moeA2ADispatchKernel; + launchWithPdlWhenEnabled("moeA2ADispatchKernel", params.enable_pdl, kernel_fn, grid_size, + kBlockSize, shared_bytes, params.stream, params.token_selected_experts, kernel_ptrs, + params.num_payloads, params.max_tokens_per_rank, params.local_num_tokens, params.ep_rank, + params.ep_size, params.num_experts_per_rank, params.enable_pdl); + }) } else { int grid_size = ceilDiv(params.local_num_tokens, kWarpsPerBlock); // If local_num_tokens is 0, we still need to launch a minimal kernel to participate in the @@ -540,12 +565,13 @@ void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params) { grid_size = 1; } int shared_bytes = 2 * kWarpsPerBlock * params.top_k * (int)sizeof(int); - SWITCH_TOP_K(params.top_k, TOP_K, - moeA2ADispatchKernel - <<>>( - params.token_selected_experts, kernel_ptrs, params.num_payloads, - params.max_tokens_per_rank, params.local_num_tokens, params.ep_rank, - params.ep_size, params.num_experts_per_rank)) + SWITCH_TOP_K(params.top_k, TOP_K, { + auto kernel_fn = moeA2ADispatchKernel; + launchWithPdlWhenEnabled("moeA2ADispatchKernel", params.enable_pdl, kernel_fn, grid_size, + kBlockSize, shared_bytes, params.stream, params.token_selected_experts, kernel_ptrs, + params.num_payloads, params.max_tokens_per_rank, params.local_num_tokens, params.ep_rank, + params.ep_size, params.num_experts_per_rank, params.enable_pdl); + }) } } @@ -553,38 +579,44 @@ void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params) { // Combine kernels // ============================================================================ -template -__device__ __forceinline__ void accumulate_vec(T* dst, T const* src) { -#pragma unroll - for (int j = 0; j < ELEMS_PER_VEC; ++j) { - dst[j] += src[j]; - } -} - -// Accumulate across all valid ranks into registers, then store once per segment -template -__device__ void vectorized_combine_impl(T* dst_typed_base, int size_per_token, int rank_id, +// Accumulate across all valid ranks into float32 registers, then store as T. +// InT: input element type in recv buffer (defaults to T for same-type accumulation). +// T: output element type written to dst. +// +// Unified path: load VEC_SIZE bytes, reinterpret as InT[elems_per_vec], accumulate as float32, +// store as T. Works for same-type (InT==T: half/bf16/float) and cross-type +// (e.g. InT=fp8_e4m3, T=bf16). sizeof(InT) must divide VEC_SIZE. +template +__device__ void vectorized_combine_impl(T* dst_typed_base, int size_per_token, + int stride_per_token, int rank_id, int max_tokens_per_rank, CombineKernelPointers const& ptrs) { - constexpr int elems_per_vec = VEC_SIZE / sizeof(T); using flashinfer::vec_t; - uint8_t* dst_bytes = reinterpret_cast(dst_typed_base); + // elems_per_vec: number of InT elements per VEC_SIZE-byte load (constexpr). + constexpr int elems_per_vec = VEC_SIZE / static_cast(sizeof(InT)); int const stride = ThreadingPolicy::stride() * VEC_SIZE; int const local_token_idx = ThreadingPolicy::token_idx(); + // offset is a byte offset into the recv buffer, stepping by VEC_SIZE bytes. for (int offset = ThreadingPolicy::offset() * VEC_SIZE; offset < size_per_token; offset += stride) { - vec_t acc[TOP_K]; - -// Unrolled K accumulation using compact top-k lists + // Per-k vec_t accumulators, zero-initialised via fill(). + // Using vec_t enables cast_store() for the output, emitting a vectorized int4 write. + vec_t acc[TOP_K]; + + // Pass 1: issue all TOP_K loads back-to-back without any type conversion. + // Raw InT bytes are loaded directly into acc[k]'s register storage, reinterpreted as + // vec_t (VEC_SIZE bytes, fitting in the low end of acc[k]'s + // sizeof(float)*elems_per_vec allocation). Separating load from cast lets the compiler + // schedule all VEC_SIZE-byte global loads consecutively, hiding memory latency across k. #pragma unroll for (int k = 0; k < TOP_K; ++k) { int target_rank = ptrs.topk_target_ranks[local_token_idx * TOP_K + k]; int dst_idx = ptrs.topk_send_indices[local_token_idx * TOP_K + k]; if (dst_idx < 0) { - acc[k].fill(0); + acc[k].fill(0.0f); continue; } @@ -592,227 +624,350 @@ __device__ void vectorized_combine_impl(T* dst_typed_base, int size_per_token, i size_t base_source_rank = static_cast(rank_id) * static_cast(max_tokens_per_rank) + static_cast(dst_idx); - size_t base_token = base_source_rank * static_cast(size_per_token); + // stride_per_token: byte distance between tokens in the recv buffer. + // Equals size_per_token for normal cases; may differ for FP8 in-place + // (BF16-stride workspace but FP8-sized payload). + size_t base_token = base_source_rank * static_cast(stride_per_token); - // Load directly into the per-k accumulator; reduce across k below - acc[k].load(recv_buffer + base_token + offset); + reinterpret_cast&>(acc[k]).load( + reinterpret_cast(recv_buffer + base_token + offset)); } - // Reduce acc[TOP_K] into acc[0] + // Pass 2: in-place cast InT → float, iterating j in descending order. + // float[j] occupies bytes [j*4, j*4+3]; InT[j] occupies [j*sizeof(InT), ...). + // For sizeof(InT) < sizeof(float), high-j float writes land above all remaining + // InT bytes, so descending order is always write-after-read safe. +#pragma unroll + for (int k = 0; k < TOP_K; ++k) { + if (ptrs.topk_send_indices[local_token_idx * TOP_K + k] < 0) + continue; // acc[k] already holds 0.0f from fill() above +#pragma unroll + for (int j = elems_per_vec - 1; j >= 0; --j) + acc[k][j] = static_cast(reinterpret_cast(&acc[k])[j]); + } + // Reduce acc[TOP_K] into acc[0] via unrolled tree-reduction. + // acc[k][j] uses vec_t::operator[] which returns float& — no indirection overhead. if constexpr (TOP_K == 22) { - T* a0 = reinterpret_cast(&acc[0]); - T* a1 = reinterpret_cast(&acc[1]); - T* a2 = reinterpret_cast(&acc[2]); - T* a3 = reinterpret_cast(&acc[3]); - T* a4 = reinterpret_cast(&acc[4]); - T* a5 = reinterpret_cast(&acc[5]); - T* a6 = reinterpret_cast(&acc[6]); - T* a7 = reinterpret_cast(&acc[7]); - T* a8 = reinterpret_cast(&acc[8]); - T* a9 = reinterpret_cast(&acc[9]); - T* a10 = reinterpret_cast(&acc[10]); - T* a11 = reinterpret_cast(&acc[11]); - T* a12 = reinterpret_cast(&acc[12]); - T* a13 = reinterpret_cast(&acc[13]); - T* a14 = reinterpret_cast(&acc[14]); - T* a15 = reinterpret_cast(&acc[15]); - T* a16 = reinterpret_cast(&acc[16]); - T* a17 = reinterpret_cast(&acc[17]); - T* a18 = reinterpret_cast(&acc[18]); - T* a19 = reinterpret_cast(&acc[19]); - T* a20 = reinterpret_cast(&acc[20]); - T* a21 = reinterpret_cast(&acc[21]); - accumulate_vec(a0, a1); - accumulate_vec(a2, a3); - accumulate_vec(a4, a5); - accumulate_vec(a6, a7); - accumulate_vec(a8, a9); - accumulate_vec(a10, a11); - accumulate_vec(a12, a13); - accumulate_vec(a14, a15); - accumulate_vec(a16, a17); - accumulate_vec(a18, a19); - accumulate_vec(a20, a21); - - accumulate_vec(a0, a2); - accumulate_vec(a4, a6); - accumulate_vec(a8, a10); - accumulate_vec(a12, a14); - accumulate_vec(a16, a18); - - accumulate_vec(a0, a4); - accumulate_vec(a8, a12); - accumulate_vec(a16, a20); - - accumulate_vec(a0, a8); - accumulate_vec(a0, a16); +#pragma unroll + for (int j = 0; j < elems_per_vec; ++j) { + acc[0][j] += acc[1][j]; + acc[2][j] += acc[3][j]; + acc[4][j] += acc[5][j]; + acc[6][j] += acc[7][j]; + acc[8][j] += acc[9][j]; + acc[10][j] += acc[11][j]; + acc[12][j] += acc[13][j]; + acc[14][j] += acc[15][j]; + acc[16][j] += acc[17][j]; + acc[18][j] += acc[19][j]; + acc[20][j] += acc[21][j]; + } +#pragma unroll + for (int j = 0; j < elems_per_vec; ++j) { + acc[0][j] += acc[2][j]; + acc[4][j] += acc[6][j]; + acc[8][j] += acc[10][j]; + acc[12][j] += acc[14][j]; + acc[16][j] += acc[18][j]; + } +#pragma unroll + for (int j = 0; j < elems_per_vec; ++j) { + acc[0][j] += acc[4][j]; + acc[8][j] += acc[12][j]; + acc[16][j] += acc[20][j]; + } +#pragma unroll + for (int j = 0; j < elems_per_vec; ++j) { + acc[0][j] += acc[8][j]; + acc[0][j] += acc[16][j]; + } } else if constexpr (TOP_K == 16) { - T* a0 = reinterpret_cast(&acc[0]); - T* a1 = reinterpret_cast(&acc[1]); - T* a2 = reinterpret_cast(&acc[2]); - T* a3 = reinterpret_cast(&acc[3]); - T* a4 = reinterpret_cast(&acc[4]); - T* a5 = reinterpret_cast(&acc[5]); - T* a6 = reinterpret_cast(&acc[6]); - T* a7 = reinterpret_cast(&acc[7]); - T* a8 = reinterpret_cast(&acc[8]); - T* a9 = reinterpret_cast(&acc[9]); - T* a10 = reinterpret_cast(&acc[10]); - T* a11 = reinterpret_cast(&acc[11]); - T* a12 = reinterpret_cast(&acc[12]); - T* a13 = reinterpret_cast(&acc[13]); - T* a14 = reinterpret_cast(&acc[14]); - T* a15 = reinterpret_cast(&acc[15]); - accumulate_vec(a0, a1); - accumulate_vec(a2, a3); - accumulate_vec(a4, a5); - accumulate_vec(a6, a7); - accumulate_vec(a8, a9); - accumulate_vec(a10, a11); - accumulate_vec(a12, a13); - accumulate_vec(a14, a15); - - accumulate_vec(a0, a2); - accumulate_vec(a4, a6); - accumulate_vec(a8, a10); - accumulate_vec(a12, a14); - - accumulate_vec(a0, a4); - accumulate_vec(a8, a12); - - accumulate_vec(a0, a8); +#pragma unroll + for (int j = 0; j < elems_per_vec; ++j) { + acc[0][j] += acc[1][j]; + acc[2][j] += acc[3][j]; + acc[4][j] += acc[5][j]; + acc[6][j] += acc[7][j]; + acc[8][j] += acc[9][j]; + acc[10][j] += acc[11][j]; + acc[12][j] += acc[13][j]; + acc[14][j] += acc[15][j]; + } +#pragma unroll + for (int j = 0; j < elems_per_vec; ++j) { + acc[0][j] += acc[2][j]; + acc[4][j] += acc[6][j]; + acc[8][j] += acc[10][j]; + acc[12][j] += acc[14][j]; + } +#pragma unroll + for (int j = 0; j < elems_per_vec; ++j) { + acc[0][j] += acc[4][j]; + acc[8][j] += acc[12][j]; + } +#pragma unroll + for (int j = 0; j < elems_per_vec; ++j) { + acc[0][j] += acc[8][j]; + } } else if constexpr (TOP_K == 10) { - T* a0 = reinterpret_cast(&acc[0]); - T* a1 = reinterpret_cast(&acc[1]); - T* a2 = reinterpret_cast(&acc[2]); - T* a3 = reinterpret_cast(&acc[3]); - T* a4 = reinterpret_cast(&acc[4]); - T* a5 = reinterpret_cast(&acc[5]); - T* a6 = reinterpret_cast(&acc[6]); - T* a7 = reinterpret_cast(&acc[7]); - T* a8 = reinterpret_cast(&acc[8]); - T* a9 = reinterpret_cast(&acc[9]); - accumulate_vec(a0, a1); - accumulate_vec(a2, a3); - accumulate_vec(a4, a5); - accumulate_vec(a6, a7); - accumulate_vec(a8, a9); - - accumulate_vec(a0, a2); - accumulate_vec(a4, a6); - - accumulate_vec(a0, a4); - accumulate_vec(a0, a8); +#pragma unroll + for (int j = 0; j < elems_per_vec; ++j) { + acc[0][j] += acc[1][j]; + acc[2][j] += acc[3][j]; + acc[4][j] += acc[5][j]; + acc[6][j] += acc[7][j]; + acc[8][j] += acc[9][j]; + } +#pragma unroll + for (int j = 0; j < elems_per_vec; ++j) { + acc[0][j] += acc[2][j]; + acc[4][j] += acc[6][j]; + } +#pragma unroll + for (int j = 0; j < elems_per_vec; ++j) { + acc[0][j] += acc[4][j]; + acc[0][j] += acc[8][j]; + } } else if constexpr (TOP_K == 8) { - T* a0 = reinterpret_cast(&acc[0]); - T* a1 = reinterpret_cast(&acc[1]); - T* a2 = reinterpret_cast(&acc[2]); - T* a3 = reinterpret_cast(&acc[3]); - T* a4 = reinterpret_cast(&acc[4]); - T* a5 = reinterpret_cast(&acc[5]); - T* a6 = reinterpret_cast(&acc[6]); - T* a7 = reinterpret_cast(&acc[7]); - accumulate_vec(a0, a1); - accumulate_vec(a2, a3); - accumulate_vec(a4, a5); - accumulate_vec(a6, a7); - accumulate_vec(a0, a2); - accumulate_vec(a4, a6); - accumulate_vec(a0, a4); +#pragma unroll + for (int j = 0; j < elems_per_vec; ++j) { + acc[0][j] += acc[1][j]; + acc[2][j] += acc[3][j]; + acc[4][j] += acc[5][j]; + acc[6][j] += acc[7][j]; + } +#pragma unroll + for (int j = 0; j < elems_per_vec; ++j) { + acc[0][j] += acc[2][j]; + acc[4][j] += acc[6][j]; + } +#pragma unroll + for (int j = 0; j < elems_per_vec; ++j) { + acc[0][j] += acc[4][j]; + } } else if constexpr (TOP_K == 6) { - T* a0 = reinterpret_cast(&acc[0]); - T* a1 = reinterpret_cast(&acc[1]); - T* a2 = reinterpret_cast(&acc[2]); - T* a3 = reinterpret_cast(&acc[3]); - T* a4 = reinterpret_cast(&acc[4]); - T* a5 = reinterpret_cast(&acc[5]); - accumulate_vec(a0, a1); - accumulate_vec(a2, a3); - accumulate_vec(a4, a5); - accumulate_vec(a0, a2); - accumulate_vec(a0, a4); +#pragma unroll + for (int j = 0; j < elems_per_vec; ++j) { + acc[0][j] += acc[1][j]; + acc[2][j] += acc[3][j]; + acc[4][j] += acc[5][j]; + } +#pragma unroll + for (int j = 0; j < elems_per_vec; ++j) { + acc[0][j] += acc[2][j]; + acc[0][j] += acc[4][j]; + } } else if constexpr (TOP_K == 4) { - T* a0 = reinterpret_cast(&acc[0]); - T* a1 = reinterpret_cast(&acc[1]); - T* a2 = reinterpret_cast(&acc[2]); - T* a3 = reinterpret_cast(&acc[3]); - accumulate_vec(a0, a1); - accumulate_vec(a2, a3); - accumulate_vec(a0, a2); +#pragma unroll + for (int j = 0; j < elems_per_vec; ++j) { + acc[0][j] += acc[1][j]; + acc[2][j] += acc[3][j]; + } +#pragma unroll + for (int j = 0; j < elems_per_vec; ++j) { + acc[0][j] += acc[2][j]; + } } else if constexpr (TOP_K == 2) { - T* a0 = reinterpret_cast(&acc[0]); - T* a1 = reinterpret_cast(&acc[1]); - accumulate_vec(a0, a1); +#pragma unroll + for (int j = 0; j < elems_per_vec; ++j) { + acc[0][j] += acc[1][j]; + } } else if constexpr (TOP_K == 1) { // nothing to do } else { - // Fallback for any future unspecialized TOP_K instantiations. - T* a0 = reinterpret_cast(&acc[0]); + // Generic fallback: accumulate all into acc[0] #pragma unroll for (int k = 1; k < TOP_K; ++k) { - T* ak = reinterpret_cast(&acc[k]); - accumulate_vec(a0, ak); +#pragma unroll + for (int j = 0; j < elems_per_vec; ++j) { + acc[0][j] += acc[k][j]; + } } } - acc[0].store(dst_bytes + offset); + // cast_store: converts float→T element-by-element then writes via vectorized int4 store. + acc[0].cast_store(dst_typed_base + offset / static_cast(sizeof(InT))); } } -// Wrapper that selects vector width based on size_per_token alignment -template -__device__ void vectorized_combine(T* dst_typed_base, int size_per_token, int rank_id, - int max_tokens_per_rank, CombineKernelPointers const& ptrs) { +// Wrapper that selects vector width based on size_per_token alignment. +// stride_per_token: byte distance between tokens in the recv buffer (may differ from +// size_per_token when FP8 in-place uses BF16-stride workspace with FP8-sized payload). +// InT: input element type in recv buffer (defaults to T for same-type accumulation) +template +__device__ void vectorized_combine(T* dst_typed_base, int size_per_token, int stride_per_token, + int rank_id, int max_tokens_per_rank, + CombineKernelPointers const& ptrs) { + // Each branch is guarded by if constexpr (sizeof(InT) <= VEC_SIZE) so that the compiler + // never instantiates vectorized_combine_impl with elems_per_vec=0. + // Branches where VEC_SIZE < sizeof(InT) are unreachable at runtime because size_per_token + // is always a multiple of sizeof(InT), so a larger alignment branch is taken first. if (size_per_token % 16 == 0) { - vectorized_combine_impl<16, TOP_K, ThreadingPolicy, T>(dst_typed_base, size_per_token, rank_id, - max_tokens_per_rank, ptrs); + if constexpr (static_cast(sizeof(InT)) <= 16) + vectorized_combine_impl<16, TOP_K, ThreadingPolicy, T, InT>( + dst_typed_base, size_per_token, stride_per_token, rank_id, max_tokens_per_rank, ptrs); } else if (size_per_token % 8 == 0) { - vectorized_combine_impl<8, TOP_K, ThreadingPolicy, T>(dst_typed_base, size_per_token, rank_id, - max_tokens_per_rank, ptrs); + if constexpr (static_cast(sizeof(InT)) <= 8) + vectorized_combine_impl<8, TOP_K, ThreadingPolicy, T, InT>( + dst_typed_base, size_per_token, stride_per_token, rank_id, max_tokens_per_rank, ptrs); } else if (size_per_token % 4 == 0) { - vectorized_combine_impl<4, TOP_K, ThreadingPolicy, T>(dst_typed_base, size_per_token, rank_id, - max_tokens_per_rank, ptrs); + if constexpr (static_cast(sizeof(InT)) <= 4) + vectorized_combine_impl<4, TOP_K, ThreadingPolicy, T, InT>( + dst_typed_base, size_per_token, stride_per_token, rank_id, max_tokens_per_rank, ptrs); } else if (size_per_token % 2 == 0) { - vectorized_combine_impl<2, TOP_K, ThreadingPolicy, T>(dst_typed_base, size_per_token, rank_id, - max_tokens_per_rank, ptrs); + if constexpr (static_cast(sizeof(InT)) <= 2) + vectorized_combine_impl<2, TOP_K, ThreadingPolicy, T, InT>( + dst_typed_base, size_per_token, stride_per_token, rank_id, max_tokens_per_rank, ptrs); } else { - vectorized_combine_impl<1, TOP_K, ThreadingPolicy, T>(dst_typed_base, size_per_token, rank_id, - max_tokens_per_rank, ptrs); + if constexpr (static_cast(sizeof(InT)) <= 1) + vectorized_combine_impl<1, TOP_K, ThreadingPolicy, T, InT>( + dst_typed_base, size_per_token, stride_per_token, rank_id, max_tokens_per_rank, ptrs); } } +// ---- vec_convert: per-vector type conversion, specialized by PTX where available ---- +// Generic: SrcT → float → DstT (all architectures, all type combinations). +template +__device__ __forceinline__ void vec_convert( + flashinfer::vec_t& out, flashinfer::vec_t const& in) { +#pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) + out[j] = DstT(static_cast(in[j])); +} + +// BF16 → FP8 e4m3: use CUDA intrinsic (SM100+, Blackwell). +// Inline PTX "cvt.rn.satfinite.e4m3x2.bf16x2 %h, %r" is rejected by SM100a ptxas +// ("Unexpected instruction types for cvt") because SM100a requires a 32-bit output +// register for this instruction. __nv_fp8x2_e4m3(bfloat162) emits the correct form. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +template = 0> +__device__ __forceinline__ void vec_convert( + flashinfer::vec_t<__nv_fp8_e4m3, VEC_SIZE>& out, + flashinfer::vec_t<__nv_bfloat16, VEC_SIZE> const& in) { + __nv_fp8x2_e4m3* out_fp8x2 = reinterpret_cast<__nv_fp8x2_e4m3*>(&out); + __nv_bfloat162 const* in_bf16x2 = reinterpret_cast<__nv_bfloat162 const*>(&in); +#pragma unroll + for (int p = 0; p < static_cast(VEC_SIZE) / 2; ++p) + out_fp8x2[p] = __nv_fp8x2_e4m3(in_bf16x2[p]); +} +#endif + +// FP16 → FP8 e4m3: paired PTX cvt.rn.satfinite.e4m3x2.f16x2 (SM89+, Hopper). +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 890) +template = 0> +__device__ __forceinline__ void vec_convert( + flashinfer::vec_t<__nv_fp8_e4m3, VEC_SIZE>& out, + flashinfer::vec_t const& in) { + uint32_t const* src_u32 = reinterpret_cast(&in); + uint16_t* dst_u16 = reinterpret_cast(&out); +#pragma unroll + for (int p = 0; p < VEC_SIZE / 2; ++p) { + uint16_t d; + asm volatile("cvt.rn.satfinite.e4m3x2.f16x2 %0, %1;" : "=h"(d) : "r"(src_u32[p])); + dst_u16[p] = d; + } +} +#endif + +// ---- vectorized_quant_impl: load → sync → convert → store ---- +// VEC_SIZE is in elements (not bytes), so both SrcT and DstT vectors hold VEC_SIZE values. +template +__device__ void vectorized_quant_impl(DstT* dst, SrcT const* src, int num_elements) { + using flashinfer::vec_t; + + int const stride = ThreadingPolicy::stride() * VEC_SIZE; + + for (int e = ThreadingPolicy::offset() * VEC_SIZE; e < num_elements; e += stride) { + vec_t in_vec; + in_vec.load(src + e); + + // Sync to ensure all threads have loaded their input vectors before any thread starts writing output. + // This avoids write-after-read hazards in the FP8 in-place case where the output of this kernel is + // read by the next iteration as input. Without this sync, some threads might start writing their + // output (DstT) before other threads have loaded their input (SrcT), causing the load to read partially + // updated data. + ThreadingPolicy::sync(); + + vec_t out_vec; + vec_convert(out_vec, in_vec); + out_vec.store(dst + e); + } +} + +template +__device__ void vectorized_quant(DstT* dst, SrcT const* src, int num_elements) { + if (num_elements % 16 == 0) + vectorized_quant_impl<16, ThreadingPolicy, SrcT, DstT>(dst, src, num_elements); + else if (num_elements % 8 == 0) + vectorized_quant_impl<8, ThreadingPolicy, SrcT, DstT>(dst, src, num_elements); + else if (num_elements % 4 == 0) + vectorized_quant_impl<4, ThreadingPolicy, SrcT, DstT>(dst, src, num_elements); + else if (num_elements % 2 == 0) + vectorized_quant_impl<2, ThreadingPolicy, SrcT, DstT>(dst, src, num_elements); + else + vectorized_quant_impl<1, ThreadingPolicy, SrcT, DstT>(dst, src, num_elements); +} + +// LOW_PRECISION=false: vectorized byte-copy (SrcT = payload dtype). +// LOW_PRECISION=true: vectorized SrcT→FP8 quantization via vectorized_quant. +// stride_per_token: byte distance between tokens in recv_buffer_bytes (host-computed, avoids +// per-thread recomputation): +// - FP8 external payload: elements_per_token × 1 (compact FP8 layout) +// - FP8 in-place / byte-copy: elements_per_token × sizeof(SrcT) (payload-dtype stride) // Copy payload to recv buffer using vectorized copy; supports warp/block token mapping -template -__global__ void moeA2APrepareCombineKernel(uint8_t* recv_buffer_bytes, uint8_t const* payload_bytes, - int bytes_per_token, int ep_size, +template +__global__ void moeA2APrepareCombineKernel(uint8_t* recv_buffer_bytes, void const* payload, + int elements_per_token, int ep_size, int max_tokens_per_rank, uint32_t* flag_val_ptr, - int const* recv_counters) { + int const* recv_counters, int stride_per_token, + bool enable_pdl) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + if (enable_pdl) cudaGridDependencySynchronize(); + cudaTriggerProgrammaticLaunchCompletion(); +#endif + if (blockIdx.x == 0 && threadIdx.x == 0) { // Increment flag_val for this combine round *flag_val_ptr = *flag_val_ptr + 1; } - if (payload_bytes == nullptr) return; + // Copy path: null payload means data is already in workspace — nothing to do. + if (!LOW_PRECISION && payload == nullptr) return; - int slot_idx = ThreadingPolicy::token_idx(); + int global_token_idx = ThreadingPolicy::token_idx(); - int total_slots = ep_size * max_tokens_per_rank; - if (slot_idx >= total_slots) return; + int global_token_num = ep_size * max_tokens_per_rank; + if (global_token_idx >= global_token_num) return; - // Map global token to (source_rank, token_idx) - int source_rank = slot_idx / max_tokens_per_rank; - int token_idx = slot_idx % max_tokens_per_rank; + // Map global_token_idx to (rank_idx, local_token_idx) + int rank_idx = global_token_idx / max_tokens_per_rank; + int local_token_idx = global_token_idx % max_tokens_per_rank; - // Skip invalid tokens beyond per-source recv count - if (token_idx >= recv_counters[source_rank]) return; + // Skip invalid tokens beyond per-rank recv count + if (local_token_idx >= recv_counters[rank_idx]) return; - // Calculate source and destination pointers for this token - size_t slot_offset = static_cast(slot_idx) * bytes_per_token; - uint8_t* dst_ptr = recv_buffer_bytes + slot_offset; - uint8_t const* src_ptr = payload_bytes + slot_offset; + size_t const token_offset = static_cast(global_token_idx) * stride_per_token; - // Copy one token's data using vectorized copy with policy - vectorized_copy(dst_ptr, src_ptr, bytes_per_token); + if constexpr (LOW_PRECISION) { + // Source pointer: external payload or in-place from workspace. + SrcT const* src_ptr = + (payload != nullptr) + ? static_cast(payload) + + static_cast(global_token_idx) * elements_per_token + : reinterpret_cast(recv_buffer_bytes + token_offset); + + // Destination: stride_per_token encodes the correct layout for both paths + // (compact FP8 for external, payload-dtype stride for in-place). + __nv_fp8_e4m3* dst_ptr = + reinterpret_cast<__nv_fp8_e4m3*>(recv_buffer_bytes + token_offset); + + vectorized_quant(dst_ptr, src_ptr, elements_per_token); + } else { + // Generic byte copy (payload guaranteed non-null by early return above). + vectorized_copy(recv_buffer_bytes + token_offset, + static_cast(payload) + token_offset, + stride_per_token); + } } // ============================================================================ @@ -822,10 +977,10 @@ __global__ void moeA2APrepareCombineKernel(uint8_t* recv_buffer_bytes, uint8_t c template __global__ void moeA2ACombineKernel( const CombineKernelPointers ptrs, // Combine-specific struct, src_data_ptrs[0] is output - int max_tokens_per_rank, int elements_per_token, int local_num_tokens, int rank_id, - int ep_size) { + int max_tokens_per_rank, int elements_per_token, int stride_per_token, int local_num_tokens, + int rank_id, int ep_size, bool enable_pdl) { int local_token_idx = ThreadingPolicy::token_idx(); - int const size_per_token = elements_per_token * sizeof(T); + int const size_per_token = elements_per_token * static_cast(sizeof(T)); if (local_num_tokens == 0) { // Special case: If local_num_tokens == 0, @@ -837,6 +992,10 @@ __global__ void moeA2ACombineKernel( if (local_token_idx >= local_num_tokens) return; } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + if (enable_pdl) cudaGridDependencySynchronize(); +#endif + #if !DISABLE_SYNC_FOR_PROFILING // In-kernel readiness synchronization at start of combine: // - One warp signals readiness to all peers with current flag_val. @@ -896,52 +1055,66 @@ __global__ void moeA2ACombineKernel( if (local_num_tokens == 0) return; - // Get output location for this token (using src_data_ptrs[0] as output) - T* token_output = static_cast(ptrs.src_data_ptrs[0]) + local_token_idx * elements_per_token; + // Dispatch to FP8→BF16 or same-type combine path + if constexpr (std::is_same_v) { + // FP8 recv buffer → BF16 output + // src_data_ptrs[0] points to a BF16 output buffer (set by moeA2ACombineOp) + auto* token_output = reinterpret_cast<__nv_bfloat16*>(ptrs.src_data_ptrs[0]) + + local_token_idx * elements_per_token; + vectorized_combine( + token_output, size_per_token, stride_per_token, rank_id, max_tokens_per_rank, ptrs); + } else { + // Get output location for this token (using src_data_ptrs[0] as output) + T* token_output = + static_cast(ptrs.src_data_ptrs[0]) + local_token_idx * elements_per_token; + // Accumulate across ranks in registers, then store once per segment + vectorized_combine(token_output, size_per_token, stride_per_token, + rank_id, max_tokens_per_rank, ptrs); + } - // Accumulate across ranks in registers, then store once per segment - vectorized_combine(token_output, size_per_token, rank_id, - max_tokens_per_rank, ptrs); +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif } void moe_a2a_prepare_combine_launch(MoeA2ACombineParams const& params) { constexpr int kBlockSize = 256; constexpr int kWarpsPerBlock = kBlockSize / 32; // 8 warps per block - // Calculate bytes per token based on dtype - int element_size; - switch (params.dtype) { - case nvinfer1::DataType::kHALF: - element_size = sizeof(half); - break; - case nvinfer1::DataType::kBF16: - element_size = sizeof(__nv_bfloat16); - break; - case nvinfer1::DataType::kFLOAT: - element_size = sizeof(float); - break; - default: - FLASHINFER_CHECK(false, "Unsupported dtype for combine prepare"); - return; - } - - int bytes_per_token = params.elements_per_token * element_size; - int total_slots = - params.prepare_payload == nullptr ? 1 : params.ep_size * params.max_tokens_per_rank; - int grid_size_warp = ceilDiv(total_slots, kWarpsPerBlock); - int grid_size_block = total_slots; // one block per token - - if (params.one_block_per_token) { - moeA2APrepareCombineKernel<<>>( - static_cast(const_cast(params.recv_buffers[params.ep_rank])), - static_cast(params.prepare_payload), bytes_per_token, params.ep_size, - params.max_tokens_per_rank, params.flag_val, params.recv_counters); - } else { - moeA2APrepareCombineKernel<<>>( - static_cast(const_cast(params.recv_buffers[params.ep_rank])), - static_cast(params.prepare_payload), bytes_per_token, params.ep_size, - params.max_tokens_per_rank, params.flag_val, params.recv_counters); - } + // FP8 in-place (payload_in_workspace=true, prepare_payload==nullptr): each CTA writes + // FP8 at the BF16-stride position, so CTAs never race — all tokens must be processed. + // Copy path with null payload is a no-op; 1 block suffices for the flag increment only. + int global_token_num = (params.use_low_precision || params.prepare_payload != nullptr) + ? params.ep_size * params.max_tokens_per_rank + : 1; + int grid_size_warp = ceilDiv(global_token_num, kWarpsPerBlock); + int grid_size_block = global_token_num; // one block per token + int grid = params.one_block_per_token ? grid_size_block : grid_size_warp; + + uint8_t* recv_buffer_bytes = + static_cast(const_cast(params.recv_buffers[params.ep_rank])); + void const* payload = params.prepare_payload; + + // stride_per_token is computed once on the host and passed to the kernel to avoid + // per-thread recomputation: + // FP8 external: EPT × 1 (compact FP8, dst packed tightly) + // FP8 in-place / byte-copy: EPT × sizeof(SrcT) (payload-dtype stride) + SWITCH_BOOL(params.use_low_precision, LOW_PRECISION, { + SWITCH_DTYPE(params.dtype, SrcT, { + bool const low_precision_staged = LOW_PRECISION && (params.prepare_payload != nullptr); + int const stride_per_token = low_precision_staged + ? params.elements_per_token + : params.elements_per_token * static_cast(sizeof(SrcT)); + auto kernel_fn = + params.one_block_per_token + ? moeA2APrepareCombineKernel + : moeA2APrepareCombineKernel; + launchWithPdlWhenEnabled("moeA2APrepareCombineKernel", params.enable_pdl, kernel_fn, grid, + kBlockSize, 0, params.stream, recv_buffer_bytes, payload, params.elements_per_token, + params.ep_size, params.max_tokens_per_rank, params.flag_val, params.recv_counters, + stride_per_token, params.enable_pdl); + }); + }); } // ============================================================================ @@ -990,19 +1163,33 @@ void moe_a2a_combine_launch(MoeA2ACombineParams const& params) { kernel_ptrs.topk_target_ranks = params.topk_target_ranks; kernel_ptrs.topk_send_indices = params.topk_send_indices; + int grid = params.one_block_per_token ? grid_size_block : grid_size_warp; + + // stride_per_token: byte distance between tokens in the recv buffer. + // FP8 external payload: EPT × 1 (compact FP8 layout) + // FP8 in-place / non-FP8: EPT × sizeof(PayloadT) (payload-dtype stride) + bool const low_precision_staged = params.use_low_precision && (params.prepare_payload != nullptr); + int stride_per_token; + SWITCH_DTYPE(params.dtype, PayloadT, { + stride_per_token = low_precision_staged + ? params.elements_per_token + : params.elements_per_token * static_cast(sizeof(PayloadT)); + }); + + // When use_low_precision is set the recv buffers contain FP8 data regardless of params.dtype, + // so dispatch the FP8 accumulation kernel in that case. + auto const effective_dtype = + params.use_low_precision ? nvinfer1::DataType::kFP8 : params.dtype; + // Launch appropriate kernel with compact macros - SWITCH_DTYPE(params.dtype, TKernelType, { + SWITCH_DTYPE(effective_dtype, TKernelType, { SWITCH_POLICY(params.one_block_per_token, Policy, { SWITCH_TOP_K(params.top_k, TOP_K, { - auto launch = [&](int grid_blocks, int block_threads) { - moeA2ACombineKernel - <<>>( - kernel_ptrs, params.max_tokens_per_rank, params.elements_per_token, - params.local_num_tokens, params.ep_rank, params.ep_size); - }; - int grid = params.one_block_per_token ? grid_size_block : grid_size_warp; - int cta = kBlockSize; - launch(grid, cta); + auto kernel_fn = moeA2ACombineKernel; + launchWithPdlWhenEnabled("moeA2ACombineKernel", params.enable_pdl, kernel_fn, grid, + kBlockSize, 0, params.stream, kernel_ptrs, params.max_tokens_per_rank, + params.elements_per_token, stride_per_token, params.local_num_tokens, params.ep_rank, + params.ep_size, params.enable_pdl); }); }); }); @@ -1012,7 +1199,11 @@ void moe_a2a_combine_launch(MoeA2ACombineParams const& params) { __global__ void moeA2ASanitizeExpertIdsKernel(int32_t* expert_ids_ptr, int32_t const* recv_counters_ptr, int ep_size, int max_tokens_per_rank, int top_k, - int32_t invalid_id) { + int32_t invalid_id, bool enable_pdl) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + if (enable_pdl) cudaGridDependencySynchronize(); + cudaTriggerProgrammaticLaunchCompletion(); +#endif int tid = blockIdx.x * blockDim.x + threadIdx.x; int total_tokens = ep_size * max_tokens_per_rank; if (tid >= total_tokens) return; @@ -1030,12 +1221,13 @@ __global__ void moeA2ASanitizeExpertIdsKernel(int32_t* expert_ids_ptr, void moe_a2a_sanitize_expert_ids_launch(int32_t* expert_ids, int32_t const* recv_counters, int32_t invalid_id, int ep_size, int max_tokens_per_rank, - int top_k, cudaStream_t stream) { + int top_k, cudaStream_t stream, bool enable_pdl) { constexpr int kBlockSize = 256; int total_tokens = ep_size * max_tokens_per_rank; int grid = ceilDiv(total_tokens, kBlockSize); - moeA2ASanitizeExpertIdsKernel<<>>( - expert_ids, recv_counters, ep_size, max_tokens_per_rank, top_k, invalid_id); + launchWithPdlWhenEnabled("moeA2ASanitizeExpertIdsKernel", enable_pdl, + moeA2ASanitizeExpertIdsKernel, grid, kBlockSize, 0, stream, expert_ids, recv_counters, + ep_size, max_tokens_per_rank, top_k, invalid_id, enable_pdl); } } // namespace tensorrt_llm::kernels::moe_alltoall diff --git a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h b/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h index fb1c476e509..8be4c25b5f9 100644 --- a/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h +++ b/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h @@ -77,6 +77,7 @@ struct CombineKernelPointers { // Dispatch phase parameters struct MoeA2ADispatchParams { bool one_block_per_token; // True: one block per token, False: one warp per token + bool enable_pdl; // True: launch with programmatic dependent launch // Threading policy // EP configuration @@ -126,6 +127,7 @@ void moe_a2a_prepare_dispatch_launch(MoeA2ADispatchParams const& params); // Combine phase parameters struct MoeA2ACombineParams { bool one_block_per_token; // True: one block per token, False: one warp per token + bool enable_pdl; // True: launch with programmatic dependent launch // EP configuration int ep_size; // Number of EP ranks @@ -137,6 +139,10 @@ struct MoeA2ACombineParams { // runtime_max_tokens_per_rank int top_k; // Number of experts per token + // If true, recv buffers contain FP8 data (either pre-staged or quantized in-place by prepare). + // The combine kernel reads FP8 and writes BF16 output. + bool use_low_precision; + // Prepare-only field: original payload tensor pointer used to stage into workspace void const* prepare_payload; @@ -175,6 +181,6 @@ void moe_a2a_prepare_combine_launch(MoeA2ACombineParams const& params); // invalid_id: value to fill for invalid tokens' expert ids void moe_a2a_sanitize_expert_ids_launch(int32_t* expert_ids, int32_t const* recv_counters, int32_t invalid_id, int ep_size, int max_tokens_per_rank, - int top_k, cudaStream_t stream); + int top_k, cudaStream_t stream, bool enable_pdl); } // namespace tensorrt_llm::kernels::moe_alltoall diff --git a/csrc/trtllm_moe_alltoall.cu b/csrc/trtllm_moe_alltoall.cu index a772529d85b..b89f363f39b 100644 --- a/csrc/trtllm_moe_alltoall.cu +++ b/csrc/trtllm_moe_alltoall.cu @@ -119,7 +119,7 @@ Tensor moeA2AInitializeOp(TensorView workspace, int64_t epRank, int64_t epSize, Tuple, Array, int64_t> moeA2ADispatchOp( TensorView tokenSelectedExperts, Array inputPayloads, TensorView workspace, TensorView metainfo, int64_t runtimeMaxTokensPerRank, int64_t epRank, int64_t epSize, - int64_t topK, int64_t numExperts) { + int64_t topK, int64_t numExperts, bool enablePdl) { using tl_throughput::PayloadDescriptor; CHECK_INPUT(tokenSelectedExperts); @@ -197,6 +197,7 @@ Tuple, Array, int64_t> moeA2ADispatchOp( tl_throughput::MoeA2ADispatchParams params{}; params.one_block_per_token = tensorrt_llm::common::getEnvMoeA2AOneBlockPerToken(); + params.enable_pdl = enablePdl; params.ep_size = static_cast(epSize); params.ep_rank = static_cast(epRank); params.num_experts_per_rank = static_cast(numExperts / epSize); @@ -265,6 +266,9 @@ nvinfer1::DataType toNvDataType(DLDataType dtype) { if (code == float32_code) { return nvinfer1::DataType::kFLOAT; } + if (code == float8_e4m3fn_code) { + return nvinfer1::DataType::kFP8; + } TVM_FFI_LOG_AND_THROW(TypeError) << "Unsupported dtype for MoE combine"; return nvinfer1::DataType::kFLOAT; } @@ -272,7 +276,7 @@ nvinfer1::DataType toNvDataType(DLDataType dtype) { Tensor moeA2ACombineOp(TensorView payload, int64_t localNumTokens, TensorView workspace, TensorView metainfo, int64_t runtimeMaxTokensPerRank, int64_t epRank, int64_t epSize, int64_t topK, int64_t combinePayloadOffset, - bool payloadInWorkspace) { + bool payloadInWorkspace, bool useLowPrecision, bool enablePdl) { using tl_throughput::MoeA2ACombineParams; CHECK_INPUT(payload); TVM_FFI_ICHECK_EQ(payload.ndim(), 3) @@ -313,16 +317,19 @@ Tensor moeA2ACombineOp(TensorView payload, int64_t localNumTokens, TensorView wo << " != " << (void*)expectedPtr; } - Tensor output = - alloc_tensor({localNumTokens, elementsPerToken}, payload.dtype(), payload.device()); + // For FP8 combine, recv buffers hold FP8 but output is BF16 (upcast during accumulation). + DLDataType outputDtype = useLowPrecision ? dl_bfloat16 : payload.dtype(); + Tensor output = alloc_tensor({localNumTokens, elementsPerToken}, outputDtype, payload.device()); MoeA2ACombineParams params{}; params.one_block_per_token = tensorrt_llm::common::getEnvMoeA2AOneBlockPerToken(); + params.enable_pdl = enablePdl; params.ep_size = static_cast(epSize); params.ep_rank = static_cast(epRank); params.local_num_tokens = static_cast(localNumTokens); params.max_tokens_per_rank = static_cast(runtimeMaxTokensPerRank); params.top_k = static_cast(topK); + params.use_low_precision = useLowPrecision; params.prepare_payload = payloadInWorkspace ? nullptr : payload.data_ptr(); params.output_data = output.data_ptr(); params.elements_per_token = static_cast(elementsPerToken); @@ -354,7 +361,7 @@ Tensor moeA2ACombineOp(TensorView payload, int64_t localNumTokens, TensorView wo } void moeA2ASanitizeExpertIdsOp(TensorView expertIds, TensorView workspace, TensorView metainfo, - int64_t epRank, int64_t invalidExpertId) { + int64_t epRank, int64_t invalidExpertId, bool enablePdl) { CHECK_INPUT(expertIds); CHECK_INPUT_TYPE(expertIds, dl_int32); TVM_FFI_ICHECK_EQ(expertIds.ndim(), 3); @@ -380,7 +387,8 @@ void moeA2ASanitizeExpertIdsOp(TensorView expertIds, TensorView workspace, Tenso tl_throughput::moe_a2a_sanitize_expert_ids_launch( static_cast(expertIds.data_ptr()), recvCounters, static_cast(invalidExpertId), static_cast(epSize), - static_cast(runtimeMaxTokensPerRank), static_cast(topK), get_current_stream()); + static_cast(runtimeMaxTokensPerRank), static_cast(topK), get_current_stream(), + enablePdl); auto err = cudaGetLastError(); TVM_FFI_ICHECK(err == cudaSuccess) diff --git a/flashinfer/comm/trtllm_moe_alltoall.py b/flashinfer/comm/trtllm_moe_alltoall.py index da85ea918cd..ee1149fa9e4 100644 --- a/flashinfer/comm/trtllm_moe_alltoall.py +++ b/flashinfer/comm/trtllm_moe_alltoall.py @@ -17,7 +17,7 @@ from .mnnvl import MnnvlMemory, MnnvlConfig from .mapping import Mapping from ..jit.comm import gen_moe_alltoall_module -from ..utils import register_custom_op +from ..utils import register_custom_op, device_support_pdl @dataclass @@ -60,6 +60,7 @@ def moe_a2a_dispatch( ep_size: int, top_k: int, num_experts: int, + enable_pdl: bool, ): """ Dispatch tokens and payloads to expert ranks. @@ -74,6 +75,7 @@ def moe_a2a_dispatch( ep_size: Total expert parallel size top_k: Number of experts per token num_experts: Total number of experts + enable_pdl: Whether to use programmatic dependent launch Returns: recv_offsets: List of offsets for each payload in the workspace @@ -90,6 +92,7 @@ def moe_a2a_dispatch( ep_size, top_k, num_experts, + enable_pdl, ) @register_custom_op( @@ -106,7 +109,9 @@ def moe_a2a_combine( ep_size: int, top_k: int, combine_payload_offset: int, - payload_in_workspace: bool = False, + payload_in_workspace: bool, + use_low_precision: bool, + enable_pdl: bool, ) -> torch.Tensor: """ Combine expert outputs back to originating tokens. @@ -122,6 +127,7 @@ def moe_a2a_combine( top_k: Number of experts per token combine_payload_offset: Offset from dispatch payload_in_workspace: If True, payload is workspace-backed + enable_pdl: Whether to use programmatic dependent launch Returns: output: [local_num_tokens, elements_per_token] tensor @@ -137,6 +143,8 @@ def moe_a2a_combine( top_k, combine_payload_offset, payload_in_workspace, + use_low_precision, + enable_pdl, ) @register_custom_op( @@ -149,9 +157,10 @@ def moe_a2a_sanitize_expert_ids( metainfo: torch.Tensor, ep_rank: int, invalid_expert_id: int, + enable_pdl: bool, ): return module.moe_a2a_sanitize_expert_ids( - expert_ids, workspace, metainfo, ep_rank, invalid_expert_id + expert_ids, workspace, metainfo, ep_rank, invalid_expert_id, enable_pdl ) @register_custom_op( @@ -264,6 +273,7 @@ def moe_a2a_dispatch( ep_size: int, top_k: int, num_experts: int, + enable_pdl: Optional[bool] = None, ): """ Dispatch tokens and payloads to expert ranks. @@ -278,11 +288,14 @@ def moe_a2a_dispatch( ep_size: Total expert parallel size top_k: Number of experts per token num_experts: Total number of experts + enable_pdl: Whether to use programmatic dependent launch (default: auto-detect from GPU) Returns: output_payloads: List of payloads for this rank, backed by data in the workspace combine_payload_offset: The offset to place the combine payload in the workspace """ + if enable_pdl is None: + enable_pdl = device_support_pdl(token_selected_experts.device) recv_offsets, recv_sizes, combine_payload_offset = ( get_moe_alltoall_module().moe_a2a_dispatch( token_selected_experts, @@ -294,6 +307,7 @@ def moe_a2a_dispatch( ep_size, top_k, num_experts, + enable_pdl, ) ) @@ -327,7 +341,11 @@ def moe_a2a_combine( top_k: int, combine_payload_offset: int, payload_in_workspace: bool = False, + use_low_precision: bool = False, + enable_pdl: Optional[bool] = None, ) -> torch.Tensor: + if enable_pdl is None: + enable_pdl = device_support_pdl(payload.device) return get_moe_alltoall_module().moe_a2a_combine( payload, local_num_tokens, @@ -339,6 +357,8 @@ def moe_a2a_combine( top_k, combine_payload_offset, payload_in_workspace, + use_low_precision, + enable_pdl, ) @@ -349,9 +369,12 @@ def moe_a2a_sanitize_expert_ids( metainfo: torch.Tensor, ep_rank: int, invalid_expert_id: int, + enable_pdl: Optional[bool] = None, ): + if enable_pdl is None: + enable_pdl = device_support_pdl(expert_ids.device) return get_moe_alltoall_module().moe_a2a_sanitize_expert_ids( - expert_ids, workspace, metainfo, ep_rank, invalid_expert_id + expert_ids, workspace, metainfo, ep_rank, invalid_expert_id, enable_pdl ) @@ -658,6 +681,7 @@ def combine( payload: torch.Tensor, runtime_max_tokens_per_rank: int, payload_in_workspace: bool = False, + use_low_precision: bool = False, ) -> torch.Tensor: """ Perform MoE all-to-all combine operation. @@ -666,6 +690,7 @@ def combine( payload: [ep_size, max_tokens, elements_per_token] tensor runtime_max_tokens_per_rank: Max tokens per rank in this batch payload_in_workspace: If True, payload is workspace-backed (skip staging) + use_low_precision: If True, quantize payload to FP8 before combine Returns: output: [local_num_tokens, elements_per_token] tensor @@ -688,6 +713,7 @@ def combine( self.top_k, self._state.combine_payload_offset, payload_in_workspace, + use_low_precision, ) # Reset state for next round diff --git a/tests/comm/test_trtllm_moe_alltoall.py b/tests/comm/test_trtllm_moe_alltoall.py index 4f008d38e0a..43ddb287e34 100644 --- a/tests/comm/test_trtllm_moe_alltoall.py +++ b/tests/comm/test_trtllm_moe_alltoall.py @@ -184,6 +184,7 @@ def dispatch_from_single_rank( num_experts, num_tokens, hidden_state_index=None, + enable_pdl=True, ): payloads = input_tensors total_payload_size_per_element = [x[0].numel() * x.itemsize for x in payloads] @@ -244,6 +245,7 @@ def dispatch_from_single_rank( ep_size=world_size, top_k=rank_token_selected_experts.shape[-1], num_experts=num_experts, + enable_pdl=enable_pdl, ) output_tensors.append(output) combine_payload_offsets.append(offset) @@ -263,6 +265,7 @@ def sanitize_expert_ids_from_single_rank( metainfo, world_size, invalid_expert_id, + enable_pdl=True, ): for rank in range(world_size): trtllm_moe_alltoall.moe_a2a_sanitize_expert_ids( @@ -271,6 +274,7 @@ def sanitize_expert_ids_from_single_rank( metainfo[rank], rank, invalid_expert_id, + enable_pdl=enable_pdl, ) return output_tensors @@ -284,6 +288,8 @@ def combine_from_single_rank( world_size, combine_payload_offsets, payload_in_workspace, + use_low_precision=False, + enable_pdl=True, ): combine_results = [] @@ -304,6 +310,8 @@ def combine_from_single_rank( top_k=top_k, combine_payload_offset=combine_payload_offsets[rank], payload_in_workspace=payload_in_workspace, + use_low_precision=use_low_precision, + enable_pdl=enable_pdl, ) ) @@ -465,11 +473,12 @@ def fake_moe( return processed_states.view(target_shape) +@pytest.mark.parametrize("enable_pdl", [True, False]) @pytest.mark.parametrize( "world_size,num_tokens,vector_dim,top_k,dtype,payload_in_workspace", COMBINE_PARAMS ) def test_moe_combine_multi_rank_single_gpu( - world_size, num_tokens, vector_dim, top_k, dtype, payload_in_workspace + world_size, num_tokens, vector_dim, top_k, dtype, payload_in_workspace, enable_pdl ): torch.cuda.set_device(0) check_sufficient_sm_count(num_tokens, world_size) @@ -512,6 +521,7 @@ def test_moe_combine_multi_rank_single_gpu( num_experts, num_tokens, hidden_state_index, + enable_pdl=enable_pdl, ) ) @@ -523,6 +533,7 @@ def test_moe_combine_multi_rank_single_gpu( metainfo, world_size, -1, + enable_pdl=enable_pdl, ) inplace_combine_tensors = [] @@ -570,6 +581,7 @@ def test_moe_combine_multi_rank_single_gpu( world_size, combine_payload_offsets, payload_in_workspace=payload_in_workspace, + enable_pdl=enable_pdl, ) reference_result = fake_moe( @@ -585,6 +597,109 @@ def test_moe_combine_multi_rank_single_gpu( ) +# FP8 combine: use_low_precision=True quantizes the BF16 payload to FP8 in the recv buffer, +# then accumulates FP8 values and writes BF16 output. Tolerance is looser to account for +# the BF16→FP8→BF16 round-trip (FP8 e4m3 has ~3-bit mantissa). +FP8_COMBINE_PARAMS = [ + (4, 16, 4096, 2, torch.bfloat16, True), # payload in workspace (in-place FP8) + (8, 16, 7168, 8, torch.bfloat16, True), # DeepSeek-V3 scale + (8, 16, 4096, 8, torch.bfloat16, False), # payload not in workspace (external copy) +] + + +@pytest.mark.skipif( + torch.cuda.get_device_capability(0) < (8, 9), + reason="FP8 (e4m3) requires SM>=89 (Ada Lovelace or newer)", +) +@pytest.mark.parametrize( + "world_size,num_tokens,vector_dim,top_k,dtype,payload_in_workspace", FP8_COMBINE_PARAMS +) +def test_moe_combine_fp8(world_size, num_tokens, vector_dim, top_k, dtype, payload_in_workspace): + """Test FP8 combine path (use_low_precision=True). + + Dispatches BF16 hidden states, then combine quantizes to FP8 in the recv buffer + before accumulating and writing BF16 output. Validates both workspace-backed + (in-place quantization) and external-payload (copy+quantize) paths. + """ + torch.cuda.set_device(0) + check_sufficient_sm_count(num_tokens, world_size) + + num_experts = world_size * top_k + token_selected_experts_index = 0 + hidden_state_index = 1 + + token_selected_experts = torch.empty( + num_tokens * world_size, top_k, dtype=torch.int32, device=torch.device("cuda") + ) + for i in range(num_tokens * world_size): + token_selected_experts[i] = torch.randperm( + num_experts, dtype=torch.int32, device=torch.device("cuda") + )[:top_k] + token_selected_experts = token_selected_experts.contiguous() + + # Scale inputs to FP8 e4m3 representable range (max ~448); randn values are O(1). + reference_tensor = make_payload(num_tokens * world_size, vector_dim, dtype) + input_tensors = [ + token_selected_experts, + reference_tensor, + make_payload(num_tokens * world_size, 1, torch.uint8), + ] + + output_tensors, all_workspaces, metainfo, combine_payload_offsets = ( + dispatch_from_single_rank( + input_tensors, token_selected_experts, world_size, num_experts, num_tokens, + hidden_state_index, + ) + ) + output_tensors = sanitize_expert_ids_from_single_rank( + output_tensors, token_selected_experts_index, all_workspaces, metainfo, world_size, -1, + ) + + inplace_combine_tensors = [] + for rank in range(world_size): + if payload_in_workspace: + inplace_combine_tensors.append( + trtllm_moe_alltoall.moe_a2a_wrap_payload_tensor_in_workspace( + all_workspaces[rank, :], + [world_size, num_tokens], + combine_payload_offsets[rank], + combine_payload_offsets[rank] + world_size * num_tokens * vector_dim * dtype.itemsize, + dtype, + ) + ) + else: + inplace_combine_tensors.append( + torch.empty(world_size, num_tokens, vector_dim, dtype=dtype, device="cuda") + ) + inplace_combine_tensors[rank].copy_( + fake_moe( + output_tensors[rank][hidden_state_index], + output_tensors[rank][token_selected_experts_index], + num_experts, is_ep=True, ep_rank=rank, + num_experts_per_rank=num_experts // world_size, + ) + ) + + combine_results = combine_from_single_rank( + inplace_combine_tensors, num_tokens, top_k, all_workspaces, metainfo, + world_size, combine_payload_offsets, payload_in_workspace=payload_in_workspace, + use_low_precision=True, + ) + + # Reference computed in float32 (same as fake_moe internals), cast to bf16 for comparison. + reference_result = fake_moe( + input_tensors[hidden_state_index], token_selected_experts, num_experts + ) + + for rank in range(world_size): + # FP8 e4m3 has ~3 mantissa bits; tolerate quantization error in each accumulated value. + torch.testing.assert_close( + combine_results[rank], + reference_result[rank * num_tokens : (rank + 1) * num_tokens], + atol=0.5, rtol=0.2, + ) + + @pytest.mark.skipif( not mnnvl_available(), reason="Mnnvl memory is not supported on this platform or container lacks SYS_PTRACE capability",