From 2c5e2340e5ae7032795f682f8e105eabc0ecf17a Mon Sep 17 00:00:00 2001 From: Shiyu Li Date: Wed, 20 May 2026 11:22:35 -0700 Subject: [PATCH 01/15] Port latcomm AllReduce optimizations to trtllm_mnnvl_allreduce Brings non-quant latency-communication optimizations from the latcomm reference impl into FlashInfer's trtllm_mnnvl_allreduce kernels: - add typed kernel params and raw-bit Lamport dirty polling - keep OOB threads in CTA arrival and dirty-buffer clearing - skip local-rank polling and sanitize polled Lamport payloads - fuse twoshot allreduce with residual add and RMSNorm - stage fused RMSNorm inputs with SM90+ TMA/mbarrier helpers --- .../comm/trtllm_mnnvl_allreduce.cuh | 1026 ++++++++++------- 1 file changed, 626 insertions(+), 400 deletions(-) diff --git a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh index 75c157147a1..d3fc1fad80d 100644 --- a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh +++ b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh @@ -54,6 +54,25 @@ struct AllReduceFusionParams { cudaStream_t stream = nullptr; }; +template +struct AllReduceKernelParams { + T* outputPtr; + T* prenormedPtr; + T const* shardPtr; + T const* residualInPtr; + T const* gammaPtr; + T** inputPtrs; + T* mcastPtr; + T* localBufferPtr; + int nRanks; + int rank; + int numTokens; + int tokenDim; + float epsilon; + float weightBias; + uint32_t* bufferFlags; +}; + namespace utils { constexpr uint16_t kNEGZERO_FP16 = 0x8000U; @@ -323,6 +342,16 @@ union PackedVec { } }; +template +inline __device__ void sanitizeLamportPayload(PackedVec& value) { +#pragma unroll + for (int i = 0; i < sizeof(PackedType) / sizeof(T); i++) { + if (isNegZero(value.elements[i])) { + value.elements[i] = fromFloat(0.f); + } + } +} + template inline __device__ PackedType loadPacked(T* ptr) { return *reinterpret_cast(ptr); @@ -334,29 +363,44 @@ inline __device__ const PackedType loadPacked(T const* ptr) { } template -inline __device__ PackedType loadPackedVolatile(void const* ptr) { +union VolatilePackedLoad { + PackedType packed; + uint32_t words[sizeof(PackedType) / sizeof(uint32_t)]; +}; + +template +inline __device__ VolatilePackedLoad loadPackedVolatile(void const* ptr) { static_assert(sizeof(PackedType) == 0, "Not implemented"); - return PackedType{}; + return {}; } template <> -inline __device__ float4 loadPackedVolatile(void const* ptr) { - float4 returnValue; - asm volatile("ld.volatile.global.v4.f32 {%0, %1, %2, %3}, [%4];\n" - : "=f"(returnValue.x), "=f"(returnValue.y), "=f"(returnValue.z), "=f"(returnValue.w) - : "l"(ptr)); +inline __device__ VolatilePackedLoad loadPackedVolatile(void const* ptr) { + VolatilePackedLoad returnValue; + asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(returnValue.words[0]), "=r"(returnValue.words[1]), + "=r"(returnValue.words[2]), "=r"(returnValue.words[3]) + : "l"(ptr) + : "memory"); return returnValue; } template <> -inline __device__ float2 loadPackedVolatile(void const* ptr) { - float2 returnValue; - asm volatile("ld.volatile.global.v2.f32 {%0, %1}, [%2];\n" - : "=f"(returnValue.x), "=f"(returnValue.y) - : "l"(ptr)); +inline __device__ VolatilePackedLoad loadPackedVolatile(void const* ptr) { + VolatilePackedLoad returnValue; + asm volatile("ld.volatile.global.v2.u32 {%0, %1}, [%2];\n" + : "=r"(returnValue.words[0]), "=r"(returnValue.words[1]) + : "l"(ptr) + : "memory"); return returnValue; } +template +inline __device__ bool isLamportDirty(VolatilePackedLoad const& value) { + // The dirty sentinel is a raw word; typed fp compares can flush nearby bit patterns. + return value.words[0] == kNEGZERO_FP32; +} + template inline __device__ void copyF4(T_IN* dst, T_IN const* src) { float4* dst4 = reinterpret_cast(dst); @@ -364,6 +408,49 @@ inline __device__ void copyF4(T_IN* dst, T_IN const* src) { __pipeline_memcpy_async(dst4, src4, sizeof(float4)); } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) +inline __device__ void mbarrierInit(void* barrier, uint32_t count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" + : + : "r"(static_cast(__cvta_generic_to_shared(barrier))), "r"(count) + : "memory"); +} + +inline __device__ void mbarrierArriveExpectTx(void* barrier, uint32_t bytes) { + asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;\n" + : + : "r"(static_cast(__cvta_generic_to_shared(barrier))), "r"(bytes) + : "memory"); +} + +inline __device__ void mbarrierWait(void* barrier, int phase) { + uint32_t barrierPtr = static_cast(__cvta_generic_to_shared(barrier)); + asm volatile("{\n\t" + ".reg .pred P1;\n\t" + "WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni WAIT;\n\t" + "DONE:\n\t" + "}" + : + : "r"(barrierPtr), "r"(phase) + : "memory"); +} + +inline __device__ void cpAsyncBulkGlobalToShared(void* dst, void const* src, void* barrier, + uint32_t bytes) { + uint32_t smemPtr = static_cast(__cvta_generic_to_shared(dst)); + uint64_t gmemPtr = static_cast(__cvta_generic_to_global(const_cast(src))); + uint32_t barrierPtr = static_cast(__cvta_generic_to_shared(barrier)); + asm volatile( + "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];\n" + : + : "r"(smemPtr), "l"(gmemPtr), "r"(bytes), "r"(barrierPtr) + : "memory"); +} +#endif + uint32_t constexpr kWARP_SIZE = 32U; uint32_t constexpr kLOG2_WARP_SIZE = 5U; uint32_t constexpr kLANE_ID_MASK = 0x1f; @@ -490,7 +577,11 @@ std::tuple adjustGridConfig(int numTokens, int dim, int eltsPerTh if (clusterSize < 8) { clusterSize = clusterSize << 1; } else { - break; + if (loadsPerThread < 8) { + loadsPerThread += 1; + } else { + break; + } } } else { if (loadsPerThread < 8) { @@ -501,29 +592,31 @@ std::tuple adjustGridConfig(int numTokens, int dim, int eltsPerTh } blockSize = ceil_div(threadsNeeded, clusterSize * loadsPerThread); } + while (smVersionMajor >= 9 && blockSize > 1024 && loadsPerThread < 8) { + loadsPerThread += 1; + blockSize = ceil_div(threadsNeeded, clusterSize * loadsPerThread); + } return {blockSize, clusterSize, loadsPerThread}; } }; // namespace utils using utils::blockReduceSum; using utils::fromFloat; +using utils::isLamportDirty; using utils::isNegZero; using utils::LamportFlags; using utils::loadPacked; using utils::loadPackedVolatile; using utils::PackedVec; +using utils::sanitizeLamportPayload; using utils::toFloat; template __global__ void __launch_bounds__(1024) - oneshotAllreduceFusionKernel(T* outputPtr, T* prenormedPtr, T const* shardPtr, - T const* residualInPtr, T const* gammaPtr, T** inputPtrs, - T* mcastPtr, int const numTokens, int const tokenDim, - float epsilon, float weightBias, int const rank, - uint32_t* bufferFlags) { + oneshotAllreduceFusionKernel(AllReduceKernelParams params) { constexpr int kELTS_PER_THREAD = sizeof(PackedType) / sizeof(T); - constexpr int kLAMPORT_ELTS_PER_PACKED = sizeof(PackedType) / sizeof(float); constexpr uint32_t kELT_SIZE = sizeof(T); + int const tokenDim = params.tokenDim; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) namespace cg = cooperative_groups; cg::cluster_group cluster = cg::this_cluster(); @@ -538,47 +631,49 @@ __global__ void __launch_bounds__(1024) // Offset w.r.t. the input shard int threadOffset = token * tokenDim + packedIdx * kELTS_PER_THREAD; #endif + bool const inBounds = packedIdx * kELTS_PER_THREAD < tokenDim; // We only use 1 stage for the oneshot allreduce - LamportFlags flag(bufferFlags, 1); - T* stagePtrMcast = reinterpret_cast(flag.getCurLamportBuf(mcastPtr, 0)); - T* stagePtrLocal = reinterpret_cast(flag.getCurLamportBuf(inputPtrs[rank], 0)); - - if (packedIdx * kELTS_PER_THREAD >= tokenDim) { - flag.ctaArrive(); - flag.clearDirtyLamportBuf(inputPtrs[rank], -1); - return; - } + LamportFlags flag(params.bufferFlags, 1); + T* stagePtrMcast = reinterpret_cast(flag.getCurLamportBuf(params.mcastPtr, 0)); + T* stagePtrLocal = reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[params.rank], 0)); // ==================== Broadcast tokens to each rank ============================= PackedVec val; - val.packed = loadPacked(&shardPtr[threadOffset]); + if (inBounds) { + val.packed = loadPacked(¶ms.shardPtr[threadOffset]); #pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) { - if (isNegZero(val.elements[i])) val.elements[i] = fromFloat(0.f); - } + for (int i = 0; i < kELTS_PER_THREAD; i++) { + if (isNegZero(val.elements[i])) val.elements[i] = fromFloat(0.f); + } - reinterpret_cast( - &stagePtrMcast[token * tokenDim * WorldSize + rank * tokenDim])[packedIdx] = val.packed; + reinterpret_cast( + &stagePtrMcast[token * tokenDim * WorldSize + params.rank * tokenDim])[packedIdx] = + val.packed; + } + // OOB threads still arrive and clear so the CTA/grid Lamport handoff cannot deadlock. flag.ctaArrive(); // ======================= Lamport Sync and clear the output buffer from previous iteration // ============================= - flag.clearDirtyLamportBuf(inputPtrs[rank], -1); + flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], -1); + + if (!inBounds) { + return; + } - PackedVec valuesLamport[WorldSize]; + PackedVec valuesLamport[WorldSize - 1]; while (1) { bool valid = true; #pragma unroll - for (int r = 0; r < WorldSize; r++) { - valuesLamport[r].packed = loadPackedVolatile( - &stagePtrLocal[token * tokenDim * WorldSize + r * tokenDim + + for (int r = 0; r < WorldSize - 1; r++) { + int const rankToLoad = (r + params.rank + 1) % WorldSize; + auto loaded = loadPackedVolatile( + &stagePtrLocal[token * tokenDim * WorldSize + rankToLoad * tokenDim + packedIdx * kELTS_PER_THREAD]); - -#pragma unroll - for (int i = 0; i < kLAMPORT_ELTS_PER_PACKED; i++) { - valid &= !isNegZero(valuesLamport[r].elements[i]); - } + valuesLamport[r].packed = loaded.packed; + // The local rank was just loaded from the shard, so only poll remote ranks. + valid &= !isLamportDirty(loaded); } if (valid) { break; @@ -592,13 +687,9 @@ __global__ void __launch_bounds__(1024) #pragma unroll for (int i = 0; i < kELTS_PER_THREAD; i++) { - accum[i] = toFloat(values[0].elements[i]); - } - + accum[i] = toFloat(val.elements[i]); #pragma unroll - for (int r = 1; r < WorldSize; r++) { -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) { + for (int r = 0; r < WorldSize - 1; r++) { accum[i] += toFloat(values[r].elements[i]); } } @@ -613,12 +704,13 @@ __global__ void __launch_bounds__(1024) if constexpr (RMSNormFusion) { // =============================== Residual =============================== PackedVec residualIn; - residualIn.packed = *reinterpret_cast(&residualInPtr[threadOffset]); + residualIn.packed = *reinterpret_cast(¶ms.residualInPtr[threadOffset]); packedAccum += residualIn; - *reinterpret_cast(&prenormedPtr[threadOffset]) = packedAccum.packed; + *reinterpret_cast(¶ms.prenormedPtr[threadOffset]) = packedAccum.packed; // =============================== Rmsnorm ================================ PackedVec gamma; - gamma.packed = *reinterpret_cast(&gammaPtr[packedIdx * kELTS_PER_THREAD]); + gamma.packed = + *reinterpret_cast(¶ms.gammaPtr[packedIdx * kELTS_PER_THREAD]); float threadSum = 0.F; #pragma unroll @@ -631,7 +723,6 @@ __global__ void __launch_bounds__(1024) __shared__ float sharedVal[8]; // Temporary variable to share the sum within block float fullSum = blockSum; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - namespace cg = cooperative_groups; cg::cluster_group cluster = cg::this_cluster(); int const numBlocks = cluster.num_blocks(); if (numBlocks > 1) { @@ -647,16 +738,16 @@ __global__ void __launch_bounds__(1024) } } #endif - float rcpRms = rsqrtf(fullSum / tokenDim + epsilon); + float rcpRms = rsqrtf(fullSum / tokenDim + params.epsilon); #pragma unroll for (int i = 0; i < kELTS_PER_THREAD; i++) { packedAccum.elements[i] = fromFloat(toFloat(packedAccum.elements[i]) * rcpRms * - (weightBias + toFloat(gamma.elements[i]))); + (params.weightBias + toFloat(gamma.elements[i]))); } } - reinterpret_cast(&outputPtr[threadOffset])[0] = packedAccum.packed; + reinterpret_cast(¶ms.outputPtr[threadOffset])[0] = packedAccum.packed; flag.waitAndUpdate( - {static_cast(numTokens * tokenDim * WorldSize * kELT_SIZE), 0, 0, 0}); + {static_cast(params.numTokens * tokenDim * WorldSize * kELT_SIZE), 0, 0, 0}); } using utils::adjustGridConfig; @@ -702,11 +793,9 @@ cudaError_t oneshotAllreduceFusionDispatch(AllReduceFusionParams const& params) .numAttrs = kSMVersionMajor >= 9 ? 2 : 1, }; -#define LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, RMSNORM) \ - FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ - &config, &oneshotAllreduceFusionKernel, output, residualOut, input, \ - residualIn, gamma, ucPtrs, mcPtr, numTokens, tokenDim, static_cast(params.epsilon), \ - params.weightBias, params.rank, params.bufferFlags)); +#define LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, RMSNORM) \ + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ + &config, &oneshotAllreduceFusionKernel, kernelParams)); #define DISPATCH_ALLREDUCE_KERNEL(WORLD_SIZE) \ if (params.rmsNormFusion) { \ LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, true); \ @@ -714,16 +803,25 @@ cudaError_t oneshotAllreduceFusionDispatch(AllReduceFusionParams const& params) LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, false); \ } - T** ucPtrs = reinterpret_cast(params.bufferPtrsDev); - T* mcPtr = reinterpret_cast(params.multicastPtr); - T* output = reinterpret_cast(params.output); - T* residualOut = reinterpret_cast(params.residualOut); - T const* input = reinterpret_cast(params.input); - T const* residualIn = reinterpret_cast(params.residualIn); - T const* gamma = reinterpret_cast(params.gamma); + AllReduceKernelParams kernelParams{ + .outputPtr = reinterpret_cast(params.output), + .prenormedPtr = reinterpret_cast(params.residualOut), + .shardPtr = reinterpret_cast(params.input), + .residualInPtr = reinterpret_cast(params.residualIn), + .gammaPtr = reinterpret_cast(params.gamma), + .inputPtrs = reinterpret_cast(params.bufferPtrsDev), + .mcastPtr = reinterpret_cast(params.multicastPtr), + .localBufferPtr = reinterpret_cast(params.bufferPtrLocal), + .nRanks = params.nRanks, + .rank = params.rank, + .numTokens = params.numTokens, + .tokenDim = params.tokenDim, + .epsilon = static_cast(params.epsilon), + .weightBias = params.weightBias, + .bufferFlags = params.bufferFlags, + }; switch (params.nRanks) { - // FIXME: Do we need other world sizes? case 2: DISPATCH_ALLREDUCE_KERNEL(2); break; @@ -759,150 +857,115 @@ enum MNNVLTwoShotStage : uint8_t { template __global__ __launch_bounds__(128) void twoshotAllreduceKernel( - T* outputPtr, T const* shardPtr, T** inputPtrs, T* mcastPtr, uint32_t const numTokens, - uint32_t const tokenDim, uint32_t const rank, uint32_t* bufferFlags, - bool const wait_for_results) { + AllReduceKernelParams params) { constexpr int kELTS_PER_THREAD = sizeof(PackedType) / sizeof(T); - constexpr int kLAMPORT_ELTS_PER_PACKED = sizeof(PackedType) / sizeof(float); constexpr uint32_t kELT_SIZE = sizeof(T); - int packedIdx = blockIdx.y * blockDim.x + threadIdx.x; - int token = blockIdx.x; - // Offset w.r.t. the input shard - int threadOffset = token * tokenDim + packedIdx * kELTS_PER_THREAD; + int const packedIdx = blockIdx.y * blockDim.x + threadIdx.x; + int const token = blockIdx.x; + int const threadOffset = token * params.tokenDim + packedIdx * kELTS_PER_THREAD; + bool const inBounds = packedIdx * kELTS_PER_THREAD < params.tokenDim; + int const destRank = token % WorldSize; + int const destTokenOffset = token / WorldSize; - int destRank = token % WorldSize; - int destTokenOffset = token / WorldSize; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaGridDependencySynchronize(); #endif - LamportFlags flag(bufferFlags, MNNVLTwoShotStage::NUM_STAGES); + LamportFlags flag(params.bufferFlags, MNNVLTwoShotStage::NUM_STAGES); T* scatterBufLocal = - reinterpret_cast(flag.getCurLamportBuf(inputPtrs[rank], MNNVLTwoShotStage::SCATTER)); + reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[params.rank], + MNNVLTwoShotStage::SCATTER)); T* scatterBufDest = - reinterpret_cast(flag.getCurLamportBuf(inputPtrs[destRank], MNNVLTwoShotStage::SCATTER)); + reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[destRank], + MNNVLTwoShotStage::SCATTER)); T* broadcastBufW = - reinterpret_cast(flag.getCurLamportBuf(mcastPtr, MNNVLTwoShotStage::BROADCAST)); + reinterpret_cast(flag.getCurLamportBuf(params.mcastPtr, MNNVLTwoShotStage::BROADCAST)); T* broadcastBufR = - reinterpret_cast(flag.getCurLamportBuf(inputPtrs[rank], MNNVLTwoShotStage::BROADCAST)); + reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[params.rank], + MNNVLTwoShotStage::BROADCAST)); -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif - // Make sure the clear function is called before OOB thread exits - if (packedIdx * kELTS_PER_THREAD >= tokenDim) { - flag.clearDirtyLamportBuf(inputPtrs[rank], -1); - return; - } - - // =============================== Scatter =============================== - - // Load vectorized data PackedVec val; - val.packed = loadPacked(&shardPtr[threadOffset]); + if (inBounds) { + val.packed = loadPacked(¶ms.shardPtr[threadOffset]); #pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) { - if (isNegZero(val.elements[i])) { - val.elements[i] = fromFloat(0.F); + for (int i = 0; i < kELTS_PER_THREAD; i++) { + if (isNegZero(val.elements[i])) { + val.elements[i] = fromFloat(0.F); + } } - } - - // Store vectorized data - reinterpret_cast( - &scatterBufDest[destTokenOffset * tokenDim * WorldSize + rank * tokenDim])[packedIdx] = - val.packed; - flag.clearDirtyLamportBuf(inputPtrs[rank], MNNVLTwoShotStage::SCATTER); - - // =============================== Reduction and Broadcast =============================== + reinterpret_cast( + &scatterBufDest[destTokenOffset * params.tokenDim * WorldSize + + params.rank * params.tokenDim])[packedIdx] = val.packed; + } - if ((token % WorldSize) == rank) { - int localToken = token / WorldSize; - float accum[kELTS_PER_THREAD] = {0.F}; +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::SCATTER); - // Use float as we only check each float value for validity - PackedVec valuesLamport[WorldSize]; + if (inBounds && destRank == params.rank) { + int const localToken = token / WorldSize; + PackedVec valuesLamport[WorldSize - 1]; while (1) { bool valid = true; #pragma unroll - for (int r = 0; r < WorldSize; r++) { - valuesLamport[r].packed = loadPackedVolatile( - &scatterBufLocal[localToken * tokenDim * WorldSize + r * tokenDim + - packedIdx * kELTS_PER_THREAD]); - - // Check validity across all elements -#pragma unroll - for (int i = 0; i < kLAMPORT_ELTS_PER_PACKED; i++) { - valid &= !isNegZero(valuesLamport[r].elements[i]); - } + for (int r = 0; r < WorldSize - 1; r++) { + int const rankToLoad = (r + params.rank + 1) % WorldSize; + auto loaded = loadPackedVolatile( + &scatterBufLocal[localToken * params.tokenDim * WorldSize + + rankToLoad * params.tokenDim + packedIdx * kELTS_PER_THREAD]); + valuesLamport[r].packed = loaded.packed; + valid &= !isLamportDirty(loaded); } if (valid) { break; } } - // Now we view it as the value for reduction auto values = reinterpret_cast*>(valuesLamport); -#pragma unroll - for (int r = 0; r < WorldSize; r++) { -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) { - accum[i] += toFloat(values[r].elements[i]); - } - } - - // Store vectorized result PackedVec packedAccum; #pragma unroll for (int i = 0; i < kELTS_PER_THREAD; i++) { - packedAccum.elements[i] = fromFloat(accum[i]); + float accum = toFloat(val.elements[i]); +#pragma unroll + for (int r = 0; r < WorldSize - 1; r++) { + accum += toFloat(values[r].elements[i]); + } + packedAccum.elements[i] = fromFloat(accum); } - reinterpret_cast(&broadcastBufW[token * tokenDim])[packedIdx] = packedAccum.packed; + // Reduced values can round to the dirty sentinel; sanitize before broadcast polling. + sanitizeLamportPayload(packedAccum); + reinterpret_cast(&broadcastBufW[token * params.tokenDim])[packedIdx] = + packedAccum.packed; } - flag.clearDirtyLamportBuf(inputPtrs[rank], MNNVLTwoShotStage::BROADCAST); - - // Optionally wait for results if the next layer isn't doing the Lamport check - if (wait_for_results) { - // Update the atomic counter to indicate the block has read the offsets - flag.ctaArrive(); + flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::BROADCAST); - PackedVec valLamport; - valLamport.packed = loadPackedVolatile(&broadcastBufR[threadOffset]); - while (isNegZero(valLamport.elements[0])) { - valLamport.packed = loadPackedVolatile(&broadcastBufR[threadOffset]); - } - if (outputPtr) { - reinterpret_cast(&outputPtr[threadOffset])[0] = valLamport.packed; + // OOB threads still arrive and clear so waitAndUpdate sees every CTA. + flag.ctaArrive(); + if (inBounds) { + auto loaded = loadPackedVolatile(&broadcastBufR[threadOffset]); + while (isLamportDirty(loaded)) { + loaded = loadPackedVolatile(&broadcastBufR[threadOffset]); } - - // Update the buffer flags - flag.waitAndUpdate( - {static_cast(round_up(numTokens, WorldSize) * tokenDim * - kELT_SIZE), // Clear Size for scatter stage - static_cast(numTokens * tokenDim * kELT_SIZE), // Clear Size for broadcast stage - 0, 0}); - // If not wait for results, we will rely on the following kernel to update the buffer + reinterpret_cast(¶ms.outputPtr[threadOffset])[0] = loaded.packed; } + + flag.waitAndUpdate({static_cast(round_up(params.numTokens, WorldSize) * + params.tokenDim * kELT_SIZE), + static_cast(params.numTokens * params.tokenDim * kELT_SIZE), 0, 0}); } using utils::copyF4; -// This kernel works performant when loads_per_thread is 1. -// For this mode, we are able to support up to 1024 (threads) x 8 (elements) = 8192 hidden -// dimension. There are two options for further scaling up: -// 1. Use CGA if supported. It expands the hidden dimension to 8k x 8 = 64k. -// 2. Set loads_per_thread >1. Which can be used if CGA is not supported. Note that this will -// be limited by the shared memory size and register count. -template -__global__ __launch_bounds__(1024) void rmsNormLamport(T_IN* outputPreNorm, T_OUT* outputNorm, - T_IN* bufferInput, T_IN const* gamma, - float epsilon, float weightBias, - T_IN const* residual, uint32_t numTokens, - uint32_t dim, uint32_t worldSize, - uint32_t* bufferFlags) { - static_assert(std::is_same_v, "T_IN and T_OUT must be the same type"); - static int const kELTS_PER_LOAD = sizeof(float4) / sizeof(T_IN); + +template +__global__ __launch_bounds__(1024) void twoshotAllreduceRMSNormFusionKernel( + AllReduceKernelParams params) { + constexpr int kELTS_PER_LOAD = sizeof(PackedType) / sizeof(T); + constexpr uint32_t kELT_SIZE = sizeof(T); uint32_t const token = blockIdx.x; uint32_t const blockSize = blockDim.x; @@ -910,157 +973,264 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(T_IN* outputPreNorm, T_OU uint32_t numThreads = blockSize; uint32_t clusterSize = 1; - uint32_t blockOffset = 0; + uint32_t clusterBlockRank = 0; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - namespace cg = cooperative_groups; - cg::cluster_group cluster = cg::this_cluster(); - numThreads = cluster.num_threads(); - clusterSize = cluster.num_blocks(); - blockOffset = cluster.block_rank(); + if constexpr (UseCGA) { + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + numThreads = cluster.num_threads(); + clusterSize = cluster.num_blocks(); + clusterBlockRank = cluster.block_rank(); + } #endif - uint32_t const dimPadded = round_up(dim, kELTS_PER_LOAD * numThreads); + + uint32_t const dimPadded = round_up(params.tokenDim, kELTS_PER_LOAD * numThreads); uint32_t const elemsPerThread = dimPadded / numThreads; uint32_t const loadStride = blockSize; + uint32_t const blockChunkSize = + ceil_div(params.tokenDim, clusterSize * kELTS_PER_LOAD) * kELTS_PER_LOAD; + uint32_t const baseTokenOffset = clusterBlockRank * blockChunkSize; + uint32_t const destRank = token % WorldSize; + uint32_t const destTokenOffset = token / WorldSize; extern __shared__ uint8_t smem[]; - float rInput[LoadsPerThread * kELTS_PER_LOAD]; - uint32_t offsets[LoadsPerThread * kELTS_PER_LOAD]; + uint32_t const smemBufferSize = blockSize * elemsPerThread * sizeof(T); + T* smemInput = reinterpret_cast(&smem[0]); + T* smemResidual = reinterpret_cast(&smem[smemBufferSize]); + T* smemGamma = reinterpret_cast(&smem[2 * smemBufferSize]); - uint32_t const smemBufferSize = blockSize * elemsPerThread * sizeof(T_IN); - T_IN* smemInput = (T_IN*)&smem[0]; - T_IN* smemResidual = (T_IN*)&smem[smemBufferSize]; - T_IN* smemGamma = (T_IN*)&smem[2 * smemBufferSize]; + alignas(16) __shared__ uint64_t residualStageBarrier; + alignas(16) __shared__ uint64_t gammaStageBarrier; +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + if (threadIdx.x == 0) { + utils::mbarrierInit(&residualStageBarrier, 1); + utils::mbarrierInit(&gammaStageBarrier, 1); + } + __syncthreads(); + cudaGridDependencySynchronize(); +#endif + + LamportFlags flag(params.bufferFlags, MNNVLTwoShotStage::NUM_STAGES); + T* scatterBufLocal = + reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[params.rank], + MNNVLTwoShotStage::SCATTER)); + T* scatterBufDest = + reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[destRank], + MNNVLTwoShotStage::SCATTER)); + T* broadcastBufW = + reinterpret_cast(flag.getCurLamportBuf(params.mcastPtr, MNNVLTwoShotStage::BROADCAST)); + T* broadcastBufR = + reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[params.rank], + MNNVLTwoShotStage::BROADCAST)); - LamportFlags flag(bufferFlags, MNNVLTwoShotStage::NUM_STAGES); - T_IN* input = reinterpret_cast( - flag.getCurLamportBuf(reinterpret_cast(bufferInput), MNNVLTwoShotStage::BROADCAST)); + PackedVec shardVals[LoadsPerThread]; +#pragma unroll + for (uint32_t i = 0; i < LoadsPerThread; i++) { + uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; + uint32_t const tokenOffset = baseTokenOffset + chunkOffset; + if (tokenOffset < params.tokenDim) { + shardVals[i].packed = + loadPacked(¶ms.shardPtr[token * params.tokenDim + tokenOffset]); +#pragma unroll + for (int j = 0; j < kELTS_PER_LOAD; j++) { + if (isNegZero(shardVals[i].elements[j])) { + shardVals[i].elements[j] = fromFloat(0.F); + } + } + reinterpret_cast( + &scatterBufDest[destTokenOffset * params.tokenDim * WorldSize + + params.rank * params.tokenDim])[tokenOffset / kELTS_PER_LOAD] = + shardVals[i].packed; + } + } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); #endif - // The offset that current thread should load from. Note that the hidden dimension is split by CGA - // size and each block loads a contiguous chunk; The size of chunk that each block processes - uint32_t const blockChunkSize = ceil_div(dim, clusterSize * kELTS_PER_LOAD) * kELTS_PER_LOAD; - uint32_t const blockLoadOffset = token * dim + blockOffset * blockChunkSize; + flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::SCATTER); + if (destRank == params.rank) { + uint32_t const localToken = token / WorldSize; #pragma unroll - for (uint32_t i = 0; i < LoadsPerThread; i++) { - // Each block load a contiguous chunk of tokens - uint32_t const threadLoadOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - offsets[i] = blockLoadOffset + threadLoadOffset; + for (uint32_t i = 0; i < LoadsPerThread; i++) { + uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; + uint32_t const tokenOffset = baseTokenOffset + chunkOffset; + if (tokenOffset < params.tokenDim) { + PackedVec valuesLamport[WorldSize - 1]; + while (1) { + bool valid = true; +#pragma unroll + for (int r = 0; r < WorldSize - 1; r++) { + int const rankToLoad = (r + params.rank + 1) % WorldSize; + auto loaded = loadPackedVolatile( + &scatterBufLocal[localToken * params.tokenDim * WorldSize + + rankToLoad * params.tokenDim + tokenOffset]); + valuesLamport[r].packed = loaded.packed; + valid &= !isLamportDirty(loaded); + } + if (valid) { + break; + } + } + + PackedVec* values = + reinterpret_cast*>(valuesLamport); + PackedVec packedAccum; +#pragma unroll + for (int j = 0; j < kELTS_PER_LOAD; j++) { + float accum = toFloat(shardVals[i].elements[j]); +#pragma unroll + for (int r = 0; r < WorldSize - 1; r++) { + accum += toFloat(values[r].elements[j]); + } + packedAccum.elements[j] = fromFloat(accum); + } + // Reduced values can round to the dirty sentinel; sanitize before broadcast polling. + sanitizeLamportPayload(packedAccum); + reinterpret_cast(&broadcastBufW[token * params.tokenDim + tokenOffset])[0] = + packedAccum.packed; + } + } } + flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::BROADCAST); + flag.ctaArrive(); + + uint32_t const stageElems = + baseTokenOffset < params.tokenDim + ? (blockChunkSize < params.tokenDim - baseTokenOffset ? blockChunkSize + : params.tokenDim - baseTokenOffset) + : 0; + uint32_t const stageBytes = stageElems * sizeof(T); +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + if (stageBytes > 0 && threadIdx.x == 0) { + T const* residualSrc = ¶ms.residualInPtr[token * params.tokenDim + baseTokenOffset]; + T const* gammaSrc = ¶ms.gammaPtr[baseTokenOffset]; + // The mbarrier expect_tx must be issued before the TMA copy that completes it. + utils::mbarrierArriveExpectTx(&residualStageBarrier, stageBytes); + utils::cpAsyncBulkGlobalToShared(smemResidual, residualSrc, &residualStageBarrier, stageBytes); + utils::mbarrierArriveExpectTx(&gammaStageBarrier, stageBytes); + utils::cpAsyncBulkGlobalToShared(smemGamma, gammaSrc, &gammaStageBarrier, stageBytes); + } +#else #pragma unroll for (uint32_t i = 0; i < LoadsPerThread; i++) { - uint32_t const threadLoadOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - if (blockOffset * blockChunkSize + threadLoadOffset < dim) { - copyF4(&smemResidual[threadLoadOffset], &residual[blockLoadOffset + threadLoadOffset]); + uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; + uint32_t const tokenOffset = baseTokenOffset + chunkOffset; + if (tokenOffset < params.tokenDim) { + copyF4(&smemResidual[chunkOffset], + ¶ms.residualInPtr[token * params.tokenDim + tokenOffset]); } } __pipeline_commit(); #pragma unroll for (uint32_t i = 0; i < LoadsPerThread; i++) { - uint32_t const threadLoadOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - if (blockOffset * blockChunkSize + threadLoadOffset < dim) { - copyF4(&smemGamma[threadLoadOffset], &gamma[blockOffset * blockChunkSize + threadLoadOffset]); + uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; + uint32_t const tokenOffset = baseTokenOffset + chunkOffset; + if (tokenOffset < params.tokenDim) { + copyF4(&smemGamma[chunkOffset], ¶ms.gammaPtr[tokenOffset]); } } __pipeline_commit(); +#endif - flag.ctaArrive(); - bool valid = false; - // ACQBLK if not lamport - while (!valid) { - valid = true; #pragma unroll - for (uint32_t i = 0; i < LoadsPerThread; i++) { - uint32_t threadLoadOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - - if (blockOffset * blockChunkSize + threadLoadOffset < dim) { - float4* dst4 = reinterpret_cast(&smemInput[threadLoadOffset]); - float4 const* src4 = reinterpret_cast(&input[offsets[i]]); - - float4 value = loadPackedVolatile(src4); - // Assume that the 16B were written atomically, so we only need to check one value - valid &= !isNegZero(value.x); - *dst4 = value; + for (uint32_t i = 0; i < LoadsPerThread; i++) { + uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; + uint32_t const tokenOffset = baseTokenOffset + chunkOffset; + if (tokenOffset < params.tokenDim) { + auto loaded = + loadPackedVolatile(&broadcastBufR[token * params.tokenDim + tokenOffset]); + while (isLamportDirty(loaded)) { + loaded = + loadPackedVolatile(&broadcastBufR[token * params.tokenDim + tokenOffset]); } + reinterpret_cast(&smemInput[chunkOffset])[0] = loaded.packed; } } + float rInput[LoadsPerThread * kELTS_PER_LOAD]; + float threadSum = 0.f; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + if (stageElems > 0) { + utils::mbarrierWait(&residualStageBarrier, 0); + } +#else __pipeline_wait_prior(1); - __syncthreads(); +#endif - float threadSum = 0.f; #pragma unroll - for (int i = 0; i < LoadsPerThread; i++) { - int threadLoadOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - if (blockOffset * blockChunkSize + threadLoadOffset < dim) { - PackedVec inp{.packed = loadPacked(&smemInput[threadLoadOffset])}; - PackedVec res{.packed = loadPacked(&smemResidual[threadLoadOffset])}; + for (uint32_t i = 0; i < LoadsPerThread; i++) { + uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; + uint32_t const tokenOffset = baseTokenOffset + chunkOffset; + if (tokenOffset < params.tokenDim) { + PackedVec inp{.packed = loadPacked(&smemInput[chunkOffset])}; + PackedVec res{.packed = loadPacked(&smemResidual[chunkOffset])}; + PackedVec inpPlusRes = inp + res; + reinterpret_cast(¶ms.prenormedPtr[token * params.tokenDim + tokenOffset])[0] = + inpPlusRes.packed; - PackedVec inp_plus_res = inp + res; #pragma unroll for (int j = 0; j < kELTS_PER_LOAD; j++) { - rInput[i * kELTS_PER_LOAD + j] = toFloat(inp_plus_res.elements[j]); - threadSum += toFloat(inp_plus_res.elements[j] * inp_plus_res.elements[j]); + float const value = toFloat(inpPlusRes.elements[j]); + rInput[i * kELTS_PER_LOAD + j] = value; + threadSum += value * value; } - - *reinterpret_cast(&outputPreNorm[blockLoadOffset + threadLoadOffset]) = - inp_plus_res.packed; } } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + if (stageElems > 0) { + utils::mbarrierWait(&gammaStageBarrier, 0); + } +#else __pipeline_wait_prior(0); +#endif float blockSum = blockReduceSum(threadSum); - float fullSum = blockSum; __shared__ float sharedVal[8]; - // Use CGA Reduction if supported #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - int const numBlocks = cluster.num_blocks(); - if (numBlocks > 1) { - fullSum = 0.F; - // Need to reduce over the entire cluster - int const blockRank = cluster.block_rank(); - if (threadIdx.x < numBlocks) { - cluster.map_shared_rank(&sharedVal[0], threadIdx.x)[blockRank] = blockSum; - } - cluster.barrier_wait(cluster.barrier_arrive()); - for (int i = 0; i < numBlocks; ++i) { - fullSum += sharedVal[i]; + if constexpr (UseCGA) { + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + int const numBlocks = cluster.num_blocks(); + if (numBlocks > 1) { + fullSum = 0.F; + int const blockRank = cluster.block_rank(); + if (threadIdx.x < numBlocks) { + cluster.map_shared_rank(&sharedVal[0], threadIdx.x)[blockRank] = blockSum; + } + cluster.barrier_wait(cluster.barrier_arrive()); + for (int i = 0; i < numBlocks; ++i) { + fullSum += sharedVal[i]; + } } } #endif - float rcpRms = rsqrtf(fullSum / dim + epsilon); - + float const rcpRms = rsqrtf(fullSum / params.tokenDim + params.epsilon); #pragma unroll - for (int i = 0; i < LoadsPerThread; i++) { - PackedVec r_out; - uint32_t threadLoadOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - if (blockOffset * blockChunkSize + threadLoadOffset < dim) { - PackedVec gamma = {.packed = loadPacked(&smemGamma[threadLoadOffset])}; - + for (uint32_t i = 0; i < LoadsPerThread; i++) { + uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; + uint32_t const tokenOffset = baseTokenOffset + chunkOffset; + if (tokenOffset < params.tokenDim) { + PackedVec gamma{.packed = loadPacked(&smemGamma[chunkOffset])}; + PackedVec out; #pragma unroll - for (uint32_t j = 0; j < kELTS_PER_LOAD; j++) { - r_out.elements[j] = fromFloat((weightBias + toFloat(gamma.elements[j])) * - rInput[i * kELTS_PER_LOAD + j] * rcpRms); + for (int j = 0; j < kELTS_PER_LOAD; j++) { + out.elements[j] = fromFloat((params.weightBias + toFloat(gamma.elements[j])) * + rInput[i * kELTS_PER_LOAD + j] * rcpRms); } - - *reinterpret_cast(&outputNorm[blockLoadOffset + threadLoadOffset]) = r_out.packed; + reinterpret_cast(¶ms.outputPtr[token * params.tokenDim + tokenOffset])[0] = + out.packed; } } - constexpr int kELTS_SIZE = sizeof(T_IN); - // Assume the previous kernel does not modify the buffer_flags. -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaGridDependencySynchronize(); -#endif - // Update the buffer pointers - flag.waitAndUpdate({static_cast(round_up(numTokens, worldSize) * dim * kELTS_SIZE), - static_cast(numTokens * dim * kELTS_SIZE), 0, 0}); + flag.waitAndUpdate({static_cast(round_up(params.numTokens, WorldSize) * + params.tokenDim * kELT_SIZE), + static_cast(params.numTokens * params.tokenDim * kELT_SIZE), 0, 0}); } template @@ -1071,152 +1241,208 @@ cudaError_t twoshotAllreduceFusionDispatch(AllReduceFusionParams const& params) FLASHINFER_CHECK(tokenDim % numEltsPerThread == 0, "[MNNVL AllReduceTwoShot] token_dim must be divisible by %d", numEltsPerThread); - int const arNumThreads = ceil_div(tokenDim, numEltsPerThread); - int const arNumBlocksPerToken = ceil_div(arNumThreads, 128); + AllReduceKernelParams kernelParams{ + .outputPtr = reinterpret_cast(params.output), + .prenormedPtr = reinterpret_cast(params.residualOut), + .shardPtr = reinterpret_cast(params.input), + .residualInPtr = reinterpret_cast(params.residualIn), + .gammaPtr = reinterpret_cast(params.gamma), + .inputPtrs = reinterpret_cast(params.bufferPtrsDev), + .mcastPtr = reinterpret_cast(params.multicastPtr), + .localBufferPtr = reinterpret_cast(params.bufferPtrLocal), + .nRanks = params.nRanks, + .rank = params.rank, + .numTokens = params.numTokens, + .tokenDim = params.tokenDim, + .epsilon = static_cast(params.epsilon), + .weightBias = params.weightBias, + .bufferFlags = params.bufferFlags, + }; - dim3 arGrid(numTokens, arNumBlocksPerToken); + if (!params.rmsNormFusion) { + int const arNumThreads = ceil_div(tokenDim, numEltsPerThread); + int const arNumBlocksPerToken = ceil_div(arNumThreads, 128); - cudaLaunchAttribute arAttrs[1]; - arAttrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - arAttrs[0].val.programmaticStreamSerializationAllowed = params.launchWithPdl ? 1 : 0; + dim3 arGrid(numTokens, arNumBlocksPerToken); - cudaLaunchConfig_t arConfig{ - .gridDim = arGrid, - .blockDim = 128, - .dynamicSmemBytes = 0, - .stream = params.stream, - .attrs = arAttrs, - .numAttrs = 1, - }; + cudaLaunchAttribute arAttrs[1]; + arAttrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + arAttrs[0].val.programmaticStreamSerializationAllowed = params.launchWithPdl ? 1 : 0; - FLASHINFER_LOG_DEBUG("[MNNVL AllReduceTwoShot] Dispatch: grid size: (%d, %d, 1), block_size: 128", - numTokens, arNumBlocksPerToken); - -#define LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE) \ - FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ - &arConfig, &twoshotAllreduceKernel, output, input, ucPtrs, mcastPtr, \ - numTokens, tokenDim, params.rank, params.bufferFlags, (!params.rmsNormFusion))); - T** ucPtrs = reinterpret_cast(params.bufferPtrsDev); - T* mcastPtr = reinterpret_cast(params.multicastPtr); - T* output = reinterpret_cast(params.output); - T const* input = reinterpret_cast(params.input); - switch (params.nRanks) { - case 2: - LAUNCH_ALLREDUCE_KERNEL(2); - break; - case 4: - LAUNCH_ALLREDUCE_KERNEL(4); - break; - case 8: - LAUNCH_ALLREDUCE_KERNEL(8); - break; - case 16: - LAUNCH_ALLREDUCE_KERNEL(16); - break; - case 32: - LAUNCH_ALLREDUCE_KERNEL(32); - break; - case 64: - LAUNCH_ALLREDUCE_KERNEL(64); - break; - default: - FLASHINFER_ERROR("[MNNVL AllReduceTwoShot] Unsupported world_size" + - std::to_string(params.nRanks) + ". Supported sizes: {2, 4, 8, 16, 32, 64}"); - return cudaErrorInvalidValue; - } + cudaLaunchConfig_t arConfig{ + .gridDim = arGrid, + .blockDim = 128, + .dynamicSmemBytes = 0, + .stream = params.stream, + .attrs = arAttrs, + .numAttrs = 1, + }; + + FLASHINFER_LOG_DEBUG( + "[MNNVL AllReduceTwoShot] Dispatch: grid size: (%d, %d, 1), block_size: 128", + numTokens, arNumBlocksPerToken); + +#define LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE) \ + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ + &arConfig, &twoshotAllreduceKernel, kernelParams)); + switch (params.nRanks) { + case 2: + LAUNCH_ALLREDUCE_KERNEL(2); + break; + case 4: + LAUNCH_ALLREDUCE_KERNEL(4); + break; + case 8: + LAUNCH_ALLREDUCE_KERNEL(8); + break; + case 16: + LAUNCH_ALLREDUCE_KERNEL(16); + break; + case 32: + LAUNCH_ALLREDUCE_KERNEL(32); + break; + case 64: + LAUNCH_ALLREDUCE_KERNEL(64); + break; + default: + FLASHINFER_ERROR("[MNNVL AllReduceTwoShot] Unsupported world_size" + + std::to_string(params.nRanks) + + ". Supported sizes: {2, 4, 8, 16, 32, 64}"); + return cudaErrorInvalidValue; + } #undef LAUNCH_ALLREDUCE_KERNEL + return cudaSuccess; + } - // Launch the rmsnorm lamport kernel if fusion is enabled - if (params.rmsNormFusion) { - static const int kSMVersionMajor = GetCudaComputeCapability().first; - auto gridConfig = adjustGridConfig(numTokens, tokenDim, numEltsPerThread, kSMVersionMajor); - int rnBlockSize = std::get<0>(gridConfig); - int rnClusterSize = std::get<1>(gridConfig); - int rnLoadsPerThread = std::get<2>(gridConfig); - - int rnNumThreads = rnClusterSize * rnBlockSize; - dim3 rnGrid(numTokens, rnClusterSize, 1); - cudaLaunchConfig_t rnConfig; - cudaLaunchAttribute rnAttrs[2]; - rnConfig.stream = params.stream; - rnConfig.gridDim = rnGrid; - rnConfig.blockDim = rnBlockSize; - rnConfig.attrs = rnAttrs; - rnAttrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - rnAttrs[0].val.programmaticStreamSerializationAllowed = params.launchWithPdl ? 1 : 0; + static const int kSMVersionMajor = GetCudaComputeCapability().first; + auto gridConfig = adjustGridConfig(numTokens, tokenDim, numEltsPerThread, kSMVersionMajor); + int rnBlockSize = std::get<0>(gridConfig); + int rnClusterSize = std::get<1>(gridConfig); + int rnLoadsPerThread = std::get<2>(gridConfig); + + bool rnUseCGA = kSMVersionMajor >= 9 && rnClusterSize > 1 && rnLoadsPerThread == 1; + if (kSMVersionMajor >= 9 && rnClusterSize > 1 && !rnUseCGA) { + gridConfig = adjustGridConfig(numTokens, tokenDim, numEltsPerThread, 0); + rnBlockSize = std::get<0>(gridConfig); + rnClusterSize = std::get<1>(gridConfig); + rnLoadsPerThread = std::get<2>(gridConfig); + rnUseCGA = false; + } + + if (rnBlockSize > 1024 || rnLoadsPerThread > 8) { + FLASHINFER_ERROR("[MNNVL AllReduceTwoShotRMSNorm] Unsupported hidden dimension " + + std::to_string(tokenDim)); + return cudaErrorInvalidValue; + } + + int const rnNumThreads = rnClusterSize * rnBlockSize; + int const dimPadded = round_up(tokenDim, numEltsPerThread * rnNumThreads); + int const iters = dimPadded / rnNumThreads; + size_t const smemSize = 3 * rnBlockSize * iters * sizeof(T); + + dim3 rnGrid(numTokens, rnClusterSize, 1); + cudaLaunchConfig_t rnConfig; + cudaLaunchAttribute rnAttrs[2]; + rnConfig.stream = params.stream; + rnConfig.gridDim = rnGrid; + rnConfig.blockDim = rnBlockSize; + rnConfig.dynamicSmemBytes = smemSize; + rnConfig.attrs = rnAttrs; + rnAttrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + rnAttrs[0].val.programmaticStreamSerializationAllowed = params.launchWithPdl ? 1 : 0; + rnConfig.numAttrs = 1; + if (rnUseCGA) { rnAttrs[1].id = cudaLaunchAttributeClusterDimension; rnAttrs[1].val.clusterDim.x = 1; rnAttrs[1].val.clusterDim.y = rnClusterSize; rnAttrs[1].val.clusterDim.z = 1; - rnConfig.numAttrs = kSMVersionMajor >= 9 ? 2 : 1; + rnConfig.numAttrs = 2; + } - bool const rnUseCGA = kSMVersionMajor >= 9 && rnClusterSize > 1; - int const dimPadded = round_up(tokenDim, numEltsPerThread * rnNumThreads); - int const iters = dimPadded / rnNumThreads; + FLASHINFER_LOG_DEBUG( + "[MNNVL AllReduceTwoShotRMSNorm] Dispatch: grid size: (%d, %d, 1), block_size: %d, " + "cluster_size: %d, loads_per_thread: %d, threads_needed: %d", + numTokens, rnClusterSize, rnBlockSize, rnClusterSize, rnLoadsPerThread, + ceil_div(tokenDim, numEltsPerThread)); + +#define LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, USE_CGA, LOADS_PER_THREAD) \ + FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( \ + &twoshotAllreduceRMSNormFusionKernel, \ + cudaFuncAttributeMaxDynamicSharedMemorySize, smemSize)); \ + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ + &rnConfig, &twoshotAllreduceRMSNormFusionKernel, \ + kernelParams)); + +#define DISPATCH_LOADS(WORLD_SIZE, USE_CGA) \ + switch (rnLoadsPerThread) { \ + case 1: \ + LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, USE_CGA, 1); \ + break; \ + case 2: \ + LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 2); \ + break; \ + case 3: \ + LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 3); \ + break; \ + case 4: \ + LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 4); \ + break; \ + case 5: \ + LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 5); \ + break; \ + case 6: \ + LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 6); \ + break; \ + case 7: \ + LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 7); \ + break; \ + case 8: \ + LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 8); \ + break; \ + default: \ + FLASHINFER_ERROR("[MNNVL AllReduceTwoShotRMSNorm] Unsupported loads_per_thread " + \ + std::to_string(rnLoadsPerThread) + \ + ". Supported sizes: {1, 2, 3, 4, 5, 6, 7, 8}"); \ + return cudaErrorInvalidValue; \ + } - size_t const smemSize = 3 * rnBlockSize * iters * sizeof(T); +#define DISPATCH_WORLD_SIZE(USE_CGA) \ + switch (params.nRanks) { \ + case 2: \ + DISPATCH_LOADS(2, USE_CGA); \ + break; \ + case 4: \ + DISPATCH_LOADS(4, USE_CGA); \ + break; \ + case 8: \ + DISPATCH_LOADS(8, USE_CGA); \ + break; \ + case 16: \ + DISPATCH_LOADS(16, USE_CGA); \ + break; \ + case 32: \ + DISPATCH_LOADS(32, USE_CGA); \ + break; \ + case 64: \ + DISPATCH_LOADS(64, USE_CGA); \ + break; \ + default: \ + FLASHINFER_ERROR("[MNNVL AllReduceTwoShotRMSNorm] Unsupported world_size" + \ + std::to_string(params.nRanks) + \ + ". Supported sizes: {2, 4, 8, 16, 32, 64}"); \ + return cudaErrorInvalidValue; \ + } - FLASHINFER_LOG_DEBUG( - "[MNNVL AllReduceTwoShotRMSNorm] Dispatch: grid size: (%d, %d, 1), block_size: %d, " - "cluster_size: %d, " - "loads_per_thread: %d, " - "threads_needed: %d", - numTokens, rnClusterSize, rnBlockSize, rnClusterSize, rnLoadsPerThread, - ceil_div(tokenDim, numEltsPerThread)); - -#define RUN_RMSNORM_KERNEL(LOADS_PER_THREAD) \ - FLASHINFER_CUDA_CALL(cudaFuncSetAttribute(&rmsNormLamport, \ - cudaFuncAttributeMaxDynamicSharedMemorySize, \ - smemSize)); \ - rnConfig.dynamicSmemBytes = smemSize; \ - FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ - &rnConfig, &rmsNormLamport, residualOut, output, bufferInput, gamma, \ - static_cast(params.epsilon), params.weightBias, residualIn, numTokens, tokenDim, \ - params.nRanks, params.bufferFlags)); - - T* residualOut = reinterpret_cast(params.residualOut); - T* output = reinterpret_cast(params.output); - T* bufferInput = reinterpret_cast(params.bufferPtrLocal); - T const* gamma = reinterpret_cast(params.gamma); - T const* residualIn = reinterpret_cast(params.residualIn); - if (rnUseCGA) { - RUN_RMSNORM_KERNEL(1); - } else { - switch (rnLoadsPerThread) { - case 1: - RUN_RMSNORM_KERNEL(1); - break; - case 2: - RUN_RMSNORM_KERNEL(2); - break; - case 3: - RUN_RMSNORM_KERNEL(3); - break; - case 4: - RUN_RMSNORM_KERNEL(4); - break; - case 5: - RUN_RMSNORM_KERNEL(5); - break; - case 6: - RUN_RMSNORM_KERNEL(6); - break; - case 7: - RUN_RMSNORM_KERNEL(7); - break; - case 8: - RUN_RMSNORM_KERNEL(8); - break; - default: - FLASHINFER_ERROR("[MNNVL AllReduceTwoShotRMSNorm] Unsupported loads_per_thread" + - std::to_string(rnLoadsPerThread) + - ". Supported sizes: {1, 2, 3, 4, 5, 6, 7, 8}"); - return cudaErrorInvalidValue; - } // switch (rnLoadsPerThread) - } // if (rnUseCGA) -#undef RUN_RMSNORM_KERNEL - - } // if (params.rmsNormFusion) + if (rnUseCGA) { + DISPATCH_WORLD_SIZE(true); + } else { + DISPATCH_WORLD_SIZE(false); + } + +#undef DISPATCH_WORLD_SIZE +#undef DISPATCH_LOADS +#undef LAUNCH_TWOSHOT_FUSION_TYPED return cudaSuccess; } } // namespace trtllm_mnnvl_allreduce From 153da029accc4670395eb5377dea874288783575 Mon Sep 17 00:00:00 2001 From: Shiyu Li Date: Wed, 20 May 2026 11:45:42 -0700 Subject: [PATCH 02/15] Use rank-stable reduction order in MNNVL allreduce Keep the local-rank Lamport poll skip, but store the local shard and remote payloads in rank-indexed slots so oneshot and twoshot reductions accumulate ranks in the same 0..world_size-1 order everywhere. --- .../comm/trtllm_mnnvl_allreduce.cuh | 57 ++++++++++--------- 1 file changed, 31 insertions(+), 26 deletions(-) diff --git a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh index d3fc1fad80d..ba8de5af11b 100644 --- a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh +++ b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh @@ -662,17 +662,20 @@ __global__ void __launch_bounds__(1024) return; } - PackedVec valuesLamport[WorldSize - 1]; + PackedVec values[WorldSize]; + values[params.rank].packed = val.packed; while (1) { bool valid = true; #pragma unroll - for (int r = 0; r < WorldSize - 1; r++) { - int const rankToLoad = (r + params.rank + 1) % WorldSize; + for (int r = 0; r < WorldSize; r++) { + if (r == params.rank) { + continue; + } auto loaded = loadPackedVolatile( - &stagePtrLocal[token * tokenDim * WorldSize + rankToLoad * tokenDim + + &stagePtrLocal[token * tokenDim * WorldSize + r * tokenDim + packedIdx * kELTS_PER_THREAD]); - valuesLamport[r].packed = loaded.packed; - // The local rank was just loaded from the shard, so only poll remote ranks. + values[r].packed = loaded.packed; + // Skip the local poll, but keep values indexed by rank for deterministic sums. valid &= !isLamportDirty(loaded); } if (valid) { @@ -680,16 +683,15 @@ __global__ void __launch_bounds__(1024) } } - auto values = reinterpret_cast*>(valuesLamport); // ======================= Reduction ============================= float accum[kELTS_PER_THREAD]; PackedVec packedAccum; #pragma unroll for (int i = 0; i < kELTS_PER_THREAD; i++) { - accum[i] = toFloat(val.elements[i]); + accum[i] = 0.f; #pragma unroll - for (int r = 0; r < WorldSize - 1; r++) { + for (int r = 0; r < WorldSize; r++) { accum[i] += toFloat(values[r].elements[i]); } } @@ -907,16 +909,19 @@ __global__ __launch_bounds__(128) void twoshotAllreduceKernel( if (inBounds && destRank == params.rank) { int const localToken = token / WorldSize; - PackedVec valuesLamport[WorldSize - 1]; + PackedVec values[WorldSize]; + values[params.rank].packed = val.packed; while (1) { bool valid = true; #pragma unroll - for (int r = 0; r < WorldSize - 1; r++) { - int const rankToLoad = (r + params.rank + 1) % WorldSize; + for (int r = 0; r < WorldSize; r++) { + if (r == params.rank) { + continue; + } auto loaded = loadPackedVolatile( &scatterBufLocal[localToken * params.tokenDim * WorldSize + - rankToLoad * params.tokenDim + packedIdx * kELTS_PER_THREAD]); - valuesLamport[r].packed = loaded.packed; + r * params.tokenDim + packedIdx * kELTS_PER_THREAD]); + values[r].packed = loaded.packed; valid &= !isLamportDirty(loaded); } if (valid) { @@ -924,13 +929,12 @@ __global__ __launch_bounds__(128) void twoshotAllreduceKernel( } } - auto values = reinterpret_cast*>(valuesLamport); PackedVec packedAccum; #pragma unroll for (int i = 0; i < kELTS_PER_THREAD; i++) { - float accum = toFloat(val.elements[i]); + float accum = 0.f; #pragma unroll - for (int r = 0; r < WorldSize - 1; r++) { + for (int r = 0; r < WorldSize; r++) { accum += toFloat(values[r].elements[i]); } packedAccum.elements[i] = fromFloat(accum); @@ -1056,16 +1060,19 @@ __global__ __launch_bounds__(1024) void twoshotAllreduceRMSNormFusionKernel( uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; uint32_t const tokenOffset = baseTokenOffset + chunkOffset; if (tokenOffset < params.tokenDim) { - PackedVec valuesLamport[WorldSize - 1]; + PackedVec values[WorldSize]; + values[params.rank].packed = shardVals[i].packed; while (1) { bool valid = true; #pragma unroll - for (int r = 0; r < WorldSize - 1; r++) { - int const rankToLoad = (r + params.rank + 1) % WorldSize; + for (int r = 0; r < WorldSize; r++) { + if (r == params.rank) { + continue; + } auto loaded = loadPackedVolatile( &scatterBufLocal[localToken * params.tokenDim * WorldSize + - rankToLoad * params.tokenDim + tokenOffset]); - valuesLamport[r].packed = loaded.packed; + r * params.tokenDim + tokenOffset]); + values[r].packed = loaded.packed; valid &= !isLamportDirty(loaded); } if (valid) { @@ -1073,14 +1080,12 @@ __global__ __launch_bounds__(1024) void twoshotAllreduceRMSNormFusionKernel( } } - PackedVec* values = - reinterpret_cast*>(valuesLamport); PackedVec packedAccum; #pragma unroll for (int j = 0; j < kELTS_PER_LOAD; j++) { - float accum = toFloat(shardVals[i].elements[j]); + float accum = 0.f; #pragma unroll - for (int r = 0; r < WorldSize - 1; r++) { + for (int r = 0; r < WorldSize; r++) { accum += toFloat(values[r].elements[j]); } packedAccum.elements[j] = fromFloat(accum); From a42d4119dc0947a8f111c4e2805fa49bae6b25e7 Mon Sep 17 00:00:00 2001 From: Shiyu Li Date: Wed, 20 May 2026 13:54:19 -0700 Subject: [PATCH 03/15] Staging optimizations. --- .../comm/trtllm_mnnvl_allreduce.cuh | 138 ++++++++---------- 1 file changed, 64 insertions(+), 74 deletions(-) diff --git a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh index ba8de5af11b..8c105723bc1 100644 --- a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh +++ b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh @@ -247,27 +247,14 @@ __device__ struct __attribute__((aligned(32))) LamportFlags { } __device__ void ctaArrive() { - int tid{0}; -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - - cg::cluster_group cluster = cg::this_cluster(); - // We update the atomic counter per cluster - tid = cluster.thread_rank(); - cluster.sync(); -#else - tid = threadIdx.x; - __syncthreads(); -#endif - if (tid == 0) { -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) - asm volatile("red.async.release.global.gpu.add.u32 [%0], %1;" ::"l"(mFlagAccessPtr), "r"(1) - : "memory"); -#elif (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700)) - asm volatile("red.release.global.gpu.add.u32 [%0], %1;" ::"l"(mFlagAccessPtr), "r"(1) - : "memory"); -#else - atomicAdd(mFlagAccessPtr, 1); -#endif + uint32_t const barrierThreads = round_up(static_cast(blockDim.x), 32u); + // Named CTA barrier avoids a full __syncthreads() while still ordering payload stores + // before the per-CTA Lamport arrival counter update. + if (threadIdx.x < 32) { + asm volatile("barrier.cta.sync 1, %0;" ::"r"(barrierThreads) : "memory"); + arriveCounter(threadIdx.x); + } else { + asm volatile("barrier.cta.arrive 1, %0;" ::"r"(barrierThreads) : "memory"); } } @@ -278,10 +265,10 @@ __device__ struct __attribute__((aligned(32))) LamportFlags { cg::grid_group grid = cg::this_grid(); // Use the first thread instead of the last thread as the last thread may exit early isLastCtaT0 = grid.thread_rank() == 0; - targetCount = grid.num_clusters(); + targetCount = gridDim.x * gridDim.y * gridDim.z; #else isLastCtaT0 = threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0; - targetCount = gridDim.x * gridDim.y; + targetCount = gridDim.x * gridDim.y * gridDim.z; #endif if (isLastCtaT0) { uint4* flagPtr = reinterpret_cast(mBufferFlagsPtr); @@ -306,6 +293,20 @@ __device__ struct __attribute__((aligned(32))) LamportFlags { alignas(16) std::array mBytesToClear; LamportBufferLayout mCurBufferLayout, mDirtyBufferLayout; + __device__ void arriveCounter(int tid) { + if (tid == 0) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) + asm volatile("red.async.release.global.gpu.add.u32 [%0], %1;" ::"l"(mFlagAccessPtr), "r"(1) + : "memory"); +#elif (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700)) + asm volatile("red.release.global.gpu.add.u32 [%0], %1;" ::"l"(mFlagAccessPtr), "r"(1) + : "memory"); +#else + atomicAdd(mFlagAccessPtr, 1); +#endif + } + } + inline __device__ void clearPackedBuf(void* bufferBasePtr, uint32_t globalTid, uint32_t numThreads, uint32_t bytesToClear, uint8_t dirtyIndex, uint8_t stageIdx) { @@ -631,51 +632,45 @@ __global__ void __launch_bounds__(1024) // Offset w.r.t. the input shard int threadOffset = token * tokenDim + packedIdx * kELTS_PER_THREAD; #endif - bool const inBounds = packedIdx * kELTS_PER_THREAD < tokenDim; - // We only use 1 stage for the oneshot allreduce LamportFlags flag(params.bufferFlags, 1); T* stagePtrMcast = reinterpret_cast(flag.getCurLamportBuf(params.mcastPtr, 0)); T* stagePtrLocal = reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[params.rank], 0)); + if (packedIdx * kELTS_PER_THREAD >= tokenDim) { + flag.ctaArrive(); + flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], -1); + return; + } + // ==================== Broadcast tokens to each rank ============================= PackedVec val; - if (inBounds) { - val.packed = loadPacked(¶ms.shardPtr[threadOffset]); + val.packed = loadPacked(¶ms.shardPtr[threadOffset]); #pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) { - if (isNegZero(val.elements[i])) val.elements[i] = fromFloat(0.f); - } - - reinterpret_cast( - &stagePtrMcast[token * tokenDim * WorldSize + params.rank * tokenDim])[packedIdx] = - val.packed; + for (int i = 0; i < kELTS_PER_THREAD; i++) { + if (isNegZero(val.elements[i])) val.elements[i] = fromFloat(0.f); } - // OOB threads still arrive and clear so the CTA/grid Lamport handoff cannot deadlock. + reinterpret_cast( + &stagePtrMcast[token * tokenDim * WorldSize + params.rank * tokenDim])[packedIdx] = + val.packed; + flag.ctaArrive(); // ======================= Lamport Sync and clear the output buffer from previous iteration // ============================= flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], -1); - if (!inBounds) { - return; - } - - PackedVec values[WorldSize]; - values[params.rank].packed = val.packed; + PackedVec remoteValues[WorldSize - 1]; while (1) { bool valid = true; #pragma unroll - for (int r = 0; r < WorldSize; r++) { - if (r == params.rank) { - continue; - } + for (int r = 0; r < WorldSize - 1; r++) { + int const rankToLoad = (r + params.rank + 1) % WorldSize; auto loaded = loadPackedVolatile( - &stagePtrLocal[token * tokenDim * WorldSize + r * tokenDim + + &stagePtrLocal[token * tokenDim * WorldSize + rankToLoad * tokenDim + packedIdx * kELTS_PER_THREAD]); - values[r].packed = loaded.packed; - // Skip the local poll, but keep values indexed by rank for deterministic sums. + remoteValues[r].packed = loaded.packed; + // Keep the local value in registers and poll remote ranks in latcomm's cyclic order. valid &= !isLamportDirty(loaded); } if (valid) { @@ -689,10 +684,10 @@ __global__ void __launch_bounds__(1024) #pragma unroll for (int i = 0; i < kELTS_PER_THREAD; i++) { - accum[i] = 0.f; + accum[i] = toFloat(val.elements[i]); #pragma unroll - for (int r = 0; r < WorldSize; r++) { - accum[i] += toFloat(values[r].elements[i]); + for (int r = 0; r < WorldSize - 1; r++) { + accum[i] += toFloat(remoteValues[r].elements[i]); } } @@ -909,19 +904,16 @@ __global__ __launch_bounds__(128) void twoshotAllreduceKernel( if (inBounds && destRank == params.rank) { int const localToken = token / WorldSize; - PackedVec values[WorldSize]; - values[params.rank].packed = val.packed; + PackedVec remoteValues[WorldSize - 1]; while (1) { bool valid = true; #pragma unroll - for (int r = 0; r < WorldSize; r++) { - if (r == params.rank) { - continue; - } + for (int r = 0; r < WorldSize - 1; r++) { + int const rankToLoad = (r + params.rank + 1) % WorldSize; auto loaded = loadPackedVolatile( &scatterBufLocal[localToken * params.tokenDim * WorldSize + - r * params.tokenDim + packedIdx * kELTS_PER_THREAD]); - values[r].packed = loaded.packed; + rankToLoad * params.tokenDim + packedIdx * kELTS_PER_THREAD]); + remoteValues[r].packed = loaded.packed; valid &= !isLamportDirty(loaded); } if (valid) { @@ -932,10 +924,10 @@ __global__ __launch_bounds__(128) void twoshotAllreduceKernel( PackedVec packedAccum; #pragma unroll for (int i = 0; i < kELTS_PER_THREAD; i++) { - float accum = 0.f; + float accum = toFloat(val.elements[i]); #pragma unroll - for (int r = 0; r < WorldSize; r++) { - accum += toFloat(values[r].elements[i]); + for (int r = 0; r < WorldSize - 1; r++) { + accum += toFloat(remoteValues[r].elements[i]); } packedAccum.elements[i] = fromFloat(accum); } @@ -1060,19 +1052,16 @@ __global__ __launch_bounds__(1024) void twoshotAllreduceRMSNormFusionKernel( uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; uint32_t const tokenOffset = baseTokenOffset + chunkOffset; if (tokenOffset < params.tokenDim) { - PackedVec values[WorldSize]; - values[params.rank].packed = shardVals[i].packed; + PackedVec remoteValues[WorldSize - 1]; while (1) { bool valid = true; #pragma unroll - for (int r = 0; r < WorldSize; r++) { - if (r == params.rank) { - continue; - } + for (int r = 0; r < WorldSize - 1; r++) { + int const rankToLoad = (r + params.rank + 1) % WorldSize; auto loaded = loadPackedVolatile( &scatterBufLocal[localToken * params.tokenDim * WorldSize + - r * params.tokenDim + tokenOffset]); - values[r].packed = loaded.packed; + rankToLoad * params.tokenDim + tokenOffset]); + remoteValues[r].packed = loaded.packed; valid &= !isLamportDirty(loaded); } if (valid) { @@ -1083,10 +1072,10 @@ __global__ __launch_bounds__(1024) void twoshotAllreduceRMSNormFusionKernel( PackedVec packedAccum; #pragma unroll for (int j = 0; j < kELTS_PER_LOAD; j++) { - float accum = 0.f; + float accum = toFloat(shardVals[i].elements[j]); #pragma unroll - for (int r = 0; r < WorldSize; r++) { - accum += toFloat(values[r].elements[j]); + for (int r = 0; r < WorldSize - 1; r++) { + accum += toFloat(remoteValues[r].elements[j]); } packedAccum.elements[j] = fromFloat(accum); } @@ -1099,7 +1088,6 @@ __global__ __launch_bounds__(1024) void twoshotAllreduceRMSNormFusionKernel( } flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::BROADCAST); - flag.ctaArrive(); uint32_t const stageElems = baseTokenOffset < params.tokenDim @@ -1139,6 +1127,8 @@ __global__ __launch_bounds__(1024) void twoshotAllreduceRMSNormFusionKernel( __pipeline_commit(); #endif + flag.ctaArrive(); + #pragma unroll for (uint32_t i = 0; i < LoadsPerThread; i++) { uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; From c8cc74015c7a192643c075f5f87b59a491421d1f Mon Sep 17 00:00:00 2001 From: Shiyu Li Date: Wed, 20 May 2026 20:11:45 -0700 Subject: [PATCH 04/15] Staging optimizations: cga barrier, ts fusion. --- .../comm/trtllm_mnnvl_allreduce.cuh | 600 +++++++++--------- 1 file changed, 286 insertions(+), 314 deletions(-) diff --git a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh index 8c105723bc1..2b226f5d597 100644 --- a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh +++ b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh @@ -202,7 +202,7 @@ struct LamportBufferLayout { namespace cg = cooperative_groups; // PackedType is the one used in kernel for Lamport buffer (LDG.128 or LDG.64) -template +template __device__ struct __attribute__((aligned(32))) LamportFlags { public: __device__ explicit LamportFlags(uint32_t* bufferFlags, uint32_t numStages = 1) @@ -247,6 +247,17 @@ __device__ struct __attribute__((aligned(32))) LamportFlags { } __device__ void ctaArrive() { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + if constexpr (UseCGA) { + cg::cluster_group cluster = cg::this_cluster(); + __cluster_barrier_arrive(); + if (cluster.block_rank() == 0 && threadIdx.x < 32) { + __cluster_barrier_wait(); + arriveCounter(threadIdx.x); + } + return; + } +#endif uint32_t const barrierThreads = round_up(static_cast(blockDim.x), 32u); // Named CTA barrier avoids a full __syncthreads() while still ordering payload stores // before the per-CTA Lamport arrival counter update. @@ -265,7 +276,12 @@ __device__ struct __attribute__((aligned(32))) LamportFlags { cg::grid_group grid = cg::this_grid(); // Use the first thread instead of the last thread as the last thread may exit early isLastCtaT0 = grid.thread_rank() == 0; - targetCount = gridDim.x * gridDim.y * gridDim.z; + if constexpr (UseCGA) { + cg::cluster_group cluster = cg::this_cluster(); + targetCount = gridDim.x * gridDim.y * gridDim.z / cluster.num_blocks(); + } else { + targetCount = gridDim.x * gridDim.y * gridDim.z; + } #else isLastCtaT0 = threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0; targetCount = gridDim.x * gridDim.y * gridDim.z; @@ -379,8 +395,8 @@ template <> inline __device__ VolatilePackedLoad loadPackedVolatile(void const* ptr) { VolatilePackedLoad returnValue; asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];\n" - : "=r"(returnValue.words[0]), "=r"(returnValue.words[1]), - "=r"(returnValue.words[2]), "=r"(returnValue.words[3]) + : "=r"(returnValue.words[0]), "=r"(returnValue.words[1]), "=r"(returnValue.words[2]), + "=r"(returnValue.words[3]) : "l"(ptr) : "memory"); return returnValue; @@ -426,17 +442,18 @@ inline __device__ void mbarrierArriveExpectTx(void* barrier, uint32_t bytes) { inline __device__ void mbarrierWait(void* barrier, int phase) { uint32_t barrierPtr = static_cast(__cvta_generic_to_shared(barrier)); - asm volatile("{\n\t" - ".reg .pred P1;\n\t" - "WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni WAIT;\n\t" - "DONE:\n\t" - "}" - : - : "r"(barrierPtr), "r"(phase) - : "memory"); + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni WAIT;\n\t" + "DONE:\n\t" + "}" + : + : "r"(barrierPtr), "r"(phase) + : "memory"); } inline __device__ void cpAsyncBulkGlobalToShared(void* dst, void const* src, void* barrier, @@ -561,6 +578,7 @@ std::tuple adjustGridConfig(int numTokens, int dim, int eltsPerTh while (threadsNeeded % clusterSize != 0 && clusterSize > 1) { clusterSize /= 2; } + int const maxDivisibleClusterSize = clusterSize; blockSize = ceil_div(threadsNeeded, clusterSize); while (blockSize < 128 && clusterSize >= 2) { blockSize *= 2; @@ -571,6 +589,17 @@ std::tuple adjustGridConfig(int numTokens, int dim, int eltsPerTh blockSize *= 2; clusterSize /= 2; } + // If the token grid still underfills the GPU, prefer more CTAs per token even when that + // makes each CTA smaller than the usual 128-thread target. + while (clusterSize < maxDivisibleClusterSize) { + int const candidateClusterSize = clusterSize * 2; + int const candidateBlockSize = ceil_div(threadsNeeded, candidateClusterSize); + if (candidateBlockSize < 64 || numTokens * candidateClusterSize > smCount) { + break; + } + clusterSize = candidateClusterSize; + blockSize = candidateBlockSize; + } } // Trying to scale up use multiple loads or CGA while (blockSize > 1024) { @@ -633,7 +662,12 @@ __global__ void __launch_bounds__(1024) int threadOffset = token * tokenDim + packedIdx * kELTS_PER_THREAD; #endif // We only use 1 stage for the oneshot allreduce - LamportFlags flag(params.bufferFlags, 1); +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + constexpr bool kUseLamportCGA = true; +#else + constexpr bool kUseLamportCGA = false; +#endif + LamportFlags flag(params.bufferFlags, 1); T* stagePtrMcast = reinterpret_cast(flag.getCurLamportBuf(params.mcastPtr, 0)); T* stagePtrLocal = reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[params.rank], 0)); @@ -790,8 +824,8 @@ cudaError_t oneshotAllreduceFusionDispatch(AllReduceFusionParams const& params) .numAttrs = kSMVersionMajor >= 9 ? 2 : 1, }; -#define LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, RMSNORM) \ - FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ +#define LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, RMSNORM) \ + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ &config, &oneshotAllreduceFusionKernel, kernelParams)); #define DISPATCH_ALLREDUCE_KERNEL(WORLD_SIZE) \ if (params.rmsNormFusion) { \ @@ -852,126 +886,27 @@ enum MNNVLTwoShotStage : uint8_t { NUM_STAGES = 2, }; -template -__global__ __launch_bounds__(128) void twoshotAllreduceKernel( - AllReduceKernelParams params) { - constexpr int kELTS_PER_THREAD = sizeof(PackedType) / sizeof(T); - constexpr uint32_t kELT_SIZE = sizeof(T); - - int const packedIdx = blockIdx.y * blockDim.x + threadIdx.x; - int const token = blockIdx.x; - int const threadOffset = token * params.tokenDim + packedIdx * kELTS_PER_THREAD; - bool const inBounds = packedIdx * kELTS_PER_THREAD < params.tokenDim; - int const destRank = token % WorldSize; - int const destTokenOffset = token / WorldSize; - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaGridDependencySynchronize(); -#endif - LamportFlags flag(params.bufferFlags, MNNVLTwoShotStage::NUM_STAGES); - - T* scatterBufLocal = - reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[params.rank], - MNNVLTwoShotStage::SCATTER)); - T* scatterBufDest = - reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[destRank], - MNNVLTwoShotStage::SCATTER)); - T* broadcastBufW = - reinterpret_cast(flag.getCurLamportBuf(params.mcastPtr, MNNVLTwoShotStage::BROADCAST)); - T* broadcastBufR = - reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[params.rank], - MNNVLTwoShotStage::BROADCAST)); - - PackedVec val; - if (inBounds) { - val.packed = loadPacked(¶ms.shardPtr[threadOffset]); -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) { - if (isNegZero(val.elements[i])) { - val.elements[i] = fromFloat(0.F); - } - } - - reinterpret_cast( - &scatterBufDest[destTokenOffset * params.tokenDim * WorldSize + - params.rank * params.tokenDim])[packedIdx] = val.packed; - } - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif - flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::SCATTER); - - if (inBounds && destRank == params.rank) { - int const localToken = token / WorldSize; - PackedVec remoteValues[WorldSize - 1]; - while (1) { - bool valid = true; -#pragma unroll - for (int r = 0; r < WorldSize - 1; r++) { - int const rankToLoad = (r + params.rank + 1) % WorldSize; - auto loaded = loadPackedVolatile( - &scatterBufLocal[localToken * params.tokenDim * WorldSize + - rankToLoad * params.tokenDim + packedIdx * kELTS_PER_THREAD]); - remoteValues[r].packed = loaded.packed; - valid &= !isLamportDirty(loaded); - } - if (valid) { - break; - } - } - - PackedVec packedAccum; -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) { - float accum = toFloat(val.elements[i]); -#pragma unroll - for (int r = 0; r < WorldSize - 1; r++) { - accum += toFloat(remoteValues[r].elements[i]); - } - packedAccum.elements[i] = fromFloat(accum); - } - // Reduced values can round to the dirty sentinel; sanitize before broadcast polling. - sanitizeLamportPayload(packedAccum); - reinterpret_cast(&broadcastBufW[token * params.tokenDim])[packedIdx] = - packedAccum.packed; - } - - flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::BROADCAST); - - // OOB threads still arrive and clear so waitAndUpdate sees every CTA. - flag.ctaArrive(); - if (inBounds) { - auto loaded = loadPackedVolatile(&broadcastBufR[threadOffset]); - while (isLamportDirty(loaded)) { - loaded = loadPackedVolatile(&broadcastBufR[threadOffset]); - } - reinterpret_cast(¶ms.outputPtr[threadOffset])[0] = loaded.packed; - } - - flag.waitAndUpdate({static_cast(round_up(params.numTokens, WorldSize) * - params.tokenDim * kELT_SIZE), - static_cast(params.numTokens * params.tokenDim * kELT_SIZE), 0, 0}); -} - using utils::copyF4; -template -__global__ __launch_bounds__(1024) void twoshotAllreduceRMSNormFusionKernel( +template +__global__ __launch_bounds__((RMSNormFusion ? 1024 : 128)) void twoshotAllreduceKernel( AllReduceKernelParams params) { constexpr int kELTS_PER_LOAD = sizeof(PackedType) / sizeof(T); constexpr uint32_t kELT_SIZE = sizeof(T); + constexpr bool kExtraAR = RMSNormFusion && ExtraAR; + constexpr uint32_t kARLoadsPerThread = kExtraAR ? 1 : LoadsPerThread; uint32_t const token = blockIdx.x; uint32_t const blockSize = blockDim.x; uint32_t const threadOffset = threadIdx.x; + bool const runRMSNormTail = !kExtraAR || blockIdx.y == 0; uint32_t numThreads = blockSize; uint32_t clusterSize = 1; uint32_t clusterBlockRank = 0; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if constexpr (UseCGA) { + if constexpr (RMSNormFusion && UseCGA) { namespace cg = cooperative_groups; cg::cluster_group cluster = cg::this_cluster(); numThreads = cluster.num_threads(); @@ -985,7 +920,12 @@ __global__ __launch_bounds__(1024) void twoshotAllreduceRMSNormFusionKernel( uint32_t const loadStride = blockSize; uint32_t const blockChunkSize = ceil_div(params.tokenDim, clusterSize * kELTS_PER_LOAD) * kELTS_PER_LOAD; - uint32_t const baseTokenOffset = clusterBlockRank * blockChunkSize; + uint32_t const baseTokenOffset = + RMSNormFusion ? clusterBlockRank * blockChunkSize : blockIdx.y * blockSize * kELTS_PER_LOAD; + uint32_t const arBaseTokenOffset = + kExtraAR ? blockIdx.y * blockSize * kELTS_PER_LOAD : baseTokenOffset; + uint32_t const rmsBaseTokenOffset = kExtraAR ? 0 : baseTokenOffset; + uint32_t const rmsBlockChunkSize = kExtraAR ? params.tokenDim : blockChunkSize; uint32_t const destRank = token % WorldSize; uint32_t const destTokenOffset = token / WorldSize; @@ -998,32 +938,32 @@ __global__ __launch_bounds__(1024) void twoshotAllreduceRMSNormFusionKernel( alignas(16) __shared__ uint64_t residualStageBarrier; alignas(16) __shared__ uint64_t gammaStageBarrier; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if (threadIdx.x == 0) { - utils::mbarrierInit(&residualStageBarrier, 1); - utils::mbarrierInit(&gammaStageBarrier, 1); + if constexpr (RMSNormFusion) { + if (runRMSNormTail && threadIdx.x == 0) { + utils::mbarrierInit(&residualStageBarrier, 1); + utils::mbarrierInit(&gammaStageBarrier, 1); + } + __syncthreads(); } - __syncthreads(); cudaGridDependencySynchronize(); #endif - LamportFlags flag(params.bufferFlags, MNNVLTwoShotStage::NUM_STAGES); - T* scatterBufLocal = - reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[params.rank], - MNNVLTwoShotStage::SCATTER)); - T* scatterBufDest = - reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[destRank], - MNNVLTwoShotStage::SCATTER)); + LamportFlags flag(params.bufferFlags, + MNNVLTwoShotStage::NUM_STAGES); + T* scatterBufLocal = reinterpret_cast( + flag.getCurLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::SCATTER)); + T* scatterBufDest = reinterpret_cast( + flag.getCurLamportBuf(params.inputPtrs[destRank], MNNVLTwoShotStage::SCATTER)); T* broadcastBufW = reinterpret_cast(flag.getCurLamportBuf(params.mcastPtr, MNNVLTwoShotStage::BROADCAST)); - T* broadcastBufR = - reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[params.rank], - MNNVLTwoShotStage::BROADCAST)); + T* broadcastBufR = reinterpret_cast( + flag.getCurLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::BROADCAST)); - PackedVec shardVals[LoadsPerThread]; + PackedVec shardVals[kARLoadsPerThread]; #pragma unroll - for (uint32_t i = 0; i < LoadsPerThread; i++) { + for (uint32_t i = 0; i < kARLoadsPerThread; i++) { uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - uint32_t const tokenOffset = baseTokenOffset + chunkOffset; + uint32_t const tokenOffset = arBaseTokenOffset + chunkOffset; if (tokenOffset < params.tokenDim) { shardVals[i].packed = loadPacked(¶ms.shardPtr[token * params.tokenDim + tokenOffset]); @@ -1043,14 +983,17 @@ __global__ __launch_bounds__(1024) void twoshotAllreduceRMSNormFusionKernel( #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); #endif + if constexpr (!kExtraAR) { + flag.ctaArrive(); + } flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::SCATTER); if (destRank == params.rank) { uint32_t const localToken = token / WorldSize; #pragma unroll - for (uint32_t i = 0; i < LoadsPerThread; i++) { + for (uint32_t i = 0; i < kARLoadsPerThread; i++) { uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - uint32_t const tokenOffset = baseTokenOffset + chunkOffset; + uint32_t const tokenOffset = arBaseTokenOffset + chunkOffset; if (tokenOffset < params.tokenDim) { PackedVec remoteValues[WorldSize - 1]; while (1) { @@ -1089,50 +1032,62 @@ __global__ __launch_bounds__(1024) void twoshotAllreduceRMSNormFusionKernel( flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::BROADCAST); - uint32_t const stageElems = - baseTokenOffset < params.tokenDim - ? (blockChunkSize < params.tokenDim - baseTokenOffset ? blockChunkSize - : params.tokenDim - baseTokenOffset) - : 0; - uint32_t const stageBytes = stageElems * sizeof(T); + if (!runRMSNormTail) { + flag.ctaArrive(); + return; + } + + uint32_t stageElems = 0; + uint32_t stageBytes = 0; + if constexpr (RMSNormFusion) { + stageElems = rmsBaseTokenOffset < params.tokenDim + ? (rmsBlockChunkSize < params.tokenDim - rmsBaseTokenOffset + ? rmsBlockChunkSize + : params.tokenDim - rmsBaseTokenOffset) + : 0; + stageBytes = stageElems * sizeof(T); + } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if (stageBytes > 0 && threadIdx.x == 0) { - T const* residualSrc = ¶ms.residualInPtr[token * params.tokenDim + baseTokenOffset]; - T const* gammaSrc = ¶ms.gammaPtr[baseTokenOffset]; - // The mbarrier expect_tx must be issued before the TMA copy that completes it. - utils::mbarrierArriveExpectTx(&residualStageBarrier, stageBytes); - utils::cpAsyncBulkGlobalToShared(smemResidual, residualSrc, &residualStageBarrier, stageBytes); - utils::mbarrierArriveExpectTx(&gammaStageBarrier, stageBytes); - utils::cpAsyncBulkGlobalToShared(smemGamma, gammaSrc, &gammaStageBarrier, stageBytes); + if constexpr (RMSNormFusion) { + if (stageBytes > 0 && threadIdx.x == 0) { + T const* residualSrc = ¶ms.residualInPtr[token * params.tokenDim + rmsBaseTokenOffset]; + T const* gammaSrc = ¶ms.gammaPtr[rmsBaseTokenOffset]; + // The mbarrier expect_tx must be issued before the TMA copy that completes it. + utils::mbarrierArriveExpectTx(&residualStageBarrier, stageBytes); + utils::cpAsyncBulkGlobalToShared(smemResidual, residualSrc, &residualStageBarrier, + stageBytes); + utils::mbarrierArriveExpectTx(&gammaStageBarrier, stageBytes); + utils::cpAsyncBulkGlobalToShared(smemGamma, gammaSrc, &gammaStageBarrier, stageBytes); + } } #else + if constexpr (RMSNormFusion) { #pragma unroll - for (uint32_t i = 0; i < LoadsPerThread; i++) { - uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - uint32_t const tokenOffset = baseTokenOffset + chunkOffset; - if (tokenOffset < params.tokenDim) { - copyF4(&smemResidual[chunkOffset], - ¶ms.residualInPtr[token * params.tokenDim + tokenOffset]); + for (uint32_t i = 0; i < LoadsPerThread; i++) { + uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; + uint32_t const tokenOffset = rmsBaseTokenOffset + chunkOffset; + if (tokenOffset < params.tokenDim) { + copyF4(&smemResidual[chunkOffset], + ¶ms.residualInPtr[token * params.tokenDim + tokenOffset]); + } } - } - __pipeline_commit(); + __pipeline_commit(); #pragma unroll - for (uint32_t i = 0; i < LoadsPerThread; i++) { - uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - uint32_t const tokenOffset = baseTokenOffset + chunkOffset; - if (tokenOffset < params.tokenDim) { - copyF4(&smemGamma[chunkOffset], ¶ms.gammaPtr[tokenOffset]); + for (uint32_t i = 0; i < LoadsPerThread; i++) { + uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; + uint32_t const tokenOffset = rmsBaseTokenOffset + chunkOffset; + if (tokenOffset < params.tokenDim) { + copyF4(&smemGamma[chunkOffset], ¶ms.gammaPtr[tokenOffset]); + } } + __pipeline_commit(); } - __pipeline_commit(); #endif - flag.ctaArrive(); - #pragma unroll for (uint32_t i = 0; i < LoadsPerThread; i++) { uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - uint32_t const tokenOffset = baseTokenOffset + chunkOffset; + uint32_t const tokenOffset = rmsBaseTokenOffset + chunkOffset; if (tokenOffset < params.tokenDim) { auto loaded = loadPackedVolatile(&broadcastBufR[token * params.tokenDim + tokenOffset]); @@ -1140,92 +1095,102 @@ __global__ __launch_bounds__(1024) void twoshotAllreduceRMSNormFusionKernel( loaded = loadPackedVolatile(&broadcastBufR[token * params.tokenDim + tokenOffset]); } - reinterpret_cast(&smemInput[chunkOffset])[0] = loaded.packed; + if constexpr (RMSNormFusion) { + reinterpret_cast(&smemInput[chunkOffset])[0] = loaded.packed; + } else { + reinterpret_cast(¶ms.outputPtr[token * params.tokenDim + tokenOffset])[0] = + loaded.packed; + } } } - float rInput[LoadsPerThread * kELTS_PER_LOAD]; - float threadSum = 0.f; + if constexpr (RMSNormFusion) { + float rInput[LoadsPerThread * kELTS_PER_LOAD]; + float threadSum = 0.f; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if (stageElems > 0) { - utils::mbarrierWait(&residualStageBarrier, 0); - } + if (stageElems > 0) { + utils::mbarrierWait(&residualStageBarrier, 0); + } #else - __pipeline_wait_prior(1); + __pipeline_wait_prior(1); #endif #pragma unroll - for (uint32_t i = 0; i < LoadsPerThread; i++) { - uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - uint32_t const tokenOffset = baseTokenOffset + chunkOffset; - if (tokenOffset < params.tokenDim) { - PackedVec inp{.packed = loadPacked(&smemInput[chunkOffset])}; - PackedVec res{.packed = loadPacked(&smemResidual[chunkOffset])}; - PackedVec inpPlusRes = inp + res; - reinterpret_cast(¶ms.prenormedPtr[token * params.tokenDim + tokenOffset])[0] = - inpPlusRes.packed; + for (uint32_t i = 0; i < LoadsPerThread; i++) { + uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; + uint32_t const tokenOffset = rmsBaseTokenOffset + chunkOffset; + if (tokenOffset < params.tokenDim) { + PackedVec inp{.packed = loadPacked(&smemInput[chunkOffset])}; + PackedVec res{.packed = loadPacked(&smemResidual[chunkOffset])}; + PackedVec inpPlusRes = inp + res; + reinterpret_cast( + ¶ms.prenormedPtr[token * params.tokenDim + tokenOffset])[0] = inpPlusRes.packed; #pragma unroll - for (int j = 0; j < kELTS_PER_LOAD; j++) { - float const value = toFloat(inpPlusRes.elements[j]); - rInput[i * kELTS_PER_LOAD + j] = value; - threadSum += value * value; + for (int j = 0; j < kELTS_PER_LOAD; j++) { + float const value = toFloat(inpPlusRes.elements[j]); + rInput[i * kELTS_PER_LOAD + j] = value; + threadSum += value * value; + } } } - } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if (stageElems > 0) { - utils::mbarrierWait(&gammaStageBarrier, 0); - } + if (stageElems > 0) { + utils::mbarrierWait(&gammaStageBarrier, 0); + } #else - __pipeline_wait_prior(0); + __pipeline_wait_prior(0); #endif - float blockSum = blockReduceSum(threadSum); - float fullSum = blockSum; - __shared__ float sharedVal[8]; + float blockSum = blockReduceSum(threadSum); + float fullSum = blockSum; + __shared__ float sharedVal[8]; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if constexpr (UseCGA) { - namespace cg = cooperative_groups; - cg::cluster_group cluster = cg::this_cluster(); - int const numBlocks = cluster.num_blocks(); - if (numBlocks > 1) { - fullSum = 0.F; - int const blockRank = cluster.block_rank(); - if (threadIdx.x < numBlocks) { - cluster.map_shared_rank(&sharedVal[0], threadIdx.x)[blockRank] = blockSum; - } - cluster.barrier_wait(cluster.barrier_arrive()); - for (int i = 0; i < numBlocks; ++i) { - fullSum += sharedVal[i]; + if constexpr (UseCGA) { + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + int const numBlocks = cluster.num_blocks(); + if (numBlocks > 1) { + fullSum = 0.F; + int const blockRank = cluster.block_rank(); + if (threadIdx.x < numBlocks) { + cluster.map_shared_rank(&sharedVal[0], threadIdx.x)[blockRank] = blockSum; + } + cluster.barrier_wait(cluster.barrier_arrive()); + for (int i = 0; i < numBlocks; ++i) { + fullSum += sharedVal[i]; + } } } - } #endif - float const rcpRms = rsqrtf(fullSum / params.tokenDim + params.epsilon); + float const rcpRms = rsqrtf(fullSum / params.tokenDim + params.epsilon); #pragma unroll - for (uint32_t i = 0; i < LoadsPerThread; i++) { - uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - uint32_t const tokenOffset = baseTokenOffset + chunkOffset; - if (tokenOffset < params.tokenDim) { - PackedVec gamma{.packed = loadPacked(&smemGamma[chunkOffset])}; - PackedVec out; + for (uint32_t i = 0; i < LoadsPerThread; i++) { + uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; + uint32_t const tokenOffset = rmsBaseTokenOffset + chunkOffset; + if (tokenOffset < params.tokenDim) { + PackedVec gamma{.packed = loadPacked(&smemGamma[chunkOffset])}; + PackedVec out; #pragma unroll - for (int j = 0; j < kELTS_PER_LOAD; j++) { - out.elements[j] = fromFloat((params.weightBias + toFloat(gamma.elements[j])) * - rInput[i * kELTS_PER_LOAD + j] * rcpRms); + for (int j = 0; j < kELTS_PER_LOAD; j++) { + out.elements[j] = fromFloat((params.weightBias + toFloat(gamma.elements[j])) * + rInput[i * kELTS_PER_LOAD + j] * rcpRms); + } + reinterpret_cast(¶ms.outputPtr[token * params.tokenDim + tokenOffset])[0] = + out.packed; } - reinterpret_cast(¶ms.outputPtr[token * params.tokenDim + tokenOffset])[0] = - out.packed; } } - flag.waitAndUpdate({static_cast(round_up(params.numTokens, WorldSize) * - params.tokenDim * kELT_SIZE), - static_cast(params.numTokens * params.tokenDim * kELT_SIZE), 0, 0}); + if constexpr (kExtraAR) { + flag.ctaArrive(); + } + flag.waitAndUpdate( + {static_cast(round_up(params.numTokens, WorldSize) * params.tokenDim * kELT_SIZE), + static_cast(params.numTokens * params.tokenDim * kELT_SIZE), 0, 0}); } template @@ -1235,6 +1200,8 @@ cudaError_t twoshotAllreduceFusionDispatch(AllReduceFusionParams const& params) int const numEltsPerThread = sizeof(float4) / sizeof(T); FLASHINFER_CHECK(tokenDim % numEltsPerThread == 0, "[MNNVL AllReduceTwoShot] token_dim must be divisible by %d", numEltsPerThread); + int const arNumThreads = ceil_div(tokenDim, numEltsPerThread); + int const arNumBlocksPerToken = ceil_div(arNumThreads, 128); AllReduceKernelParams kernelParams{ .outputPtr = reinterpret_cast(params.output), @@ -1255,9 +1222,6 @@ cudaError_t twoshotAllreduceFusionDispatch(AllReduceFusionParams const& params) }; if (!params.rmsNormFusion) { - int const arNumThreads = ceil_div(tokenDim, numEltsPerThread); - int const arNumBlocksPerToken = ceil_div(arNumThreads, 128); - dim3 arGrid(numTokens, arNumBlocksPerToken); cudaLaunchAttribute arAttrs[1]; @@ -1274,12 +1238,12 @@ cudaError_t twoshotAllreduceFusionDispatch(AllReduceFusionParams const& params) }; FLASHINFER_LOG_DEBUG( - "[MNNVL AllReduceTwoShot] Dispatch: grid size: (%d, %d, 1), block_size: 128", - numTokens, arNumBlocksPerToken); + "[MNNVL AllReduceTwoShot] Dispatch: grid size: (%d, %d, 1), block_size: 128", numTokens, + arNumBlocksPerToken); -#define LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE) \ - FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ - &arConfig, &twoshotAllreduceKernel, kernelParams)); +#define LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE) \ + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ + &arConfig, &twoshotAllreduceKernel, kernelParams)); switch (params.nRanks) { case 2: LAUNCH_ALLREDUCE_KERNEL(2); @@ -1330,17 +1294,23 @@ cudaError_t twoshotAllreduceFusionDispatch(AllReduceFusionParams const& params) return cudaErrorInvalidValue; } + bool const useExtraAR = params.nRanks == 8 && !rnUseCGA && rnLoadsPerThread == 1 && + arNumBlocksPerToken > rnClusterSize && arNumBlocksPerToken <= 8; + int const launchBlockSize = useExtraAR ? 128 : rnBlockSize; + int const launchGridY = useExtraAR ? arNumBlocksPerToken : rnClusterSize; + int const launchLoadsPerThread = useExtraAR ? arNumBlocksPerToken : rnLoadsPerThread; int const rnNumThreads = rnClusterSize * rnBlockSize; - int const dimPadded = round_up(tokenDim, numEltsPerThread * rnNumThreads); - int const iters = dimPadded / rnNumThreads; - size_t const smemSize = 3 * rnBlockSize * iters * sizeof(T); + int const launchNumThreads = useExtraAR ? launchBlockSize : rnNumThreads; + int const dimPadded = round_up(tokenDim, numEltsPerThread * launchNumThreads); + int const iters = dimPadded / launchNumThreads; + size_t const smemSize = 3 * launchBlockSize * iters * sizeof(T); - dim3 rnGrid(numTokens, rnClusterSize, 1); + dim3 rnGrid(numTokens, launchGridY, 1); cudaLaunchConfig_t rnConfig; cudaLaunchAttribute rnAttrs[2]; rnConfig.stream = params.stream; rnConfig.gridDim = rnGrid; - rnConfig.blockDim = rnBlockSize; + rnConfig.blockDim = launchBlockSize; rnConfig.dynamicSmemBytes = smemSize; rnConfig.attrs = rnAttrs; rnAttrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; @@ -1356,83 +1326,85 @@ cudaError_t twoshotAllreduceFusionDispatch(AllReduceFusionParams const& params) FLASHINFER_LOG_DEBUG( "[MNNVL AllReduceTwoShotRMSNorm] Dispatch: grid size: (%d, %d, 1), block_size: %d, " - "cluster_size: %d, loads_per_thread: %d, threads_needed: %d", - numTokens, rnClusterSize, rnBlockSize, rnClusterSize, rnLoadsPerThread, - ceil_div(tokenDim, numEltsPerThread)); - -#define LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, USE_CGA, LOADS_PER_THREAD) \ - FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( \ - &twoshotAllreduceRMSNormFusionKernel, \ - cudaFuncAttributeMaxDynamicSharedMemorySize, smemSize)); \ - FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ - &rnConfig, &twoshotAllreduceRMSNormFusionKernel, \ + "cluster_size: %d, loads_per_thread: %d, threads_needed: %d, extra_ar: %d", + numTokens, launchGridY, launchBlockSize, rnClusterSize, launchLoadsPerThread, + ceil_div(tokenDim, numEltsPerThread), static_cast(useExtraAR)); + +#define LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, USE_CGA, LOADS_PER_THREAD, EXTRA_AR) \ + FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( \ + &twoshotAllreduceKernel, \ + cudaFuncAttributeMaxDynamicSharedMemorySize, smemSize)); \ + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ + &rnConfig, \ + &twoshotAllreduceKernel, \ kernelParams)); -#define DISPATCH_LOADS(WORLD_SIZE, USE_CGA) \ - switch (rnLoadsPerThread) { \ - case 1: \ - LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, USE_CGA, 1); \ - break; \ - case 2: \ - LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 2); \ - break; \ - case 3: \ - LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 3); \ - break; \ - case 4: \ - LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 4); \ - break; \ - case 5: \ - LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 5); \ - break; \ - case 6: \ - LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 6); \ - break; \ - case 7: \ - LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 7); \ - break; \ - case 8: \ - LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 8); \ - break; \ - default: \ +#define DISPATCH_LOADS(WORLD_SIZE, USE_CGA, EXTRA_AR) \ + switch (launchLoadsPerThread) { \ + case 1: \ + LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, USE_CGA, 1, EXTRA_AR); \ + break; \ + case 2: \ + LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 2, EXTRA_AR); \ + break; \ + case 3: \ + LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 3, EXTRA_AR); \ + break; \ + case 4: \ + LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 4, EXTRA_AR); \ + break; \ + case 5: \ + LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 5, EXTRA_AR); \ + break; \ + case 6: \ + LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 6, EXTRA_AR); \ + break; \ + case 7: \ + LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 7, EXTRA_AR); \ + break; \ + case 8: \ + LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 8, EXTRA_AR); \ + break; \ + default: \ FLASHINFER_ERROR("[MNNVL AllReduceTwoShotRMSNorm] Unsupported loads_per_thread " + \ - std::to_string(rnLoadsPerThread) + \ - ". Supported sizes: {1, 2, 3, 4, 5, 6, 7, 8}"); \ - return cudaErrorInvalidValue; \ + std::to_string(launchLoadsPerThread) + \ + ". Supported sizes: {1, 2, 3, 4, 5, 6, 7, 8}"); \ + return cudaErrorInvalidValue; \ } -#define DISPATCH_WORLD_SIZE(USE_CGA) \ - switch (params.nRanks) { \ - case 2: \ - DISPATCH_LOADS(2, USE_CGA); \ - break; \ - case 4: \ - DISPATCH_LOADS(4, USE_CGA); \ - break; \ - case 8: \ - DISPATCH_LOADS(8, USE_CGA); \ - break; \ - case 16: \ - DISPATCH_LOADS(16, USE_CGA); \ - break; \ - case 32: \ - DISPATCH_LOADS(32, USE_CGA); \ - break; \ - case 64: \ - DISPATCH_LOADS(64, USE_CGA); \ - break; \ - default: \ - FLASHINFER_ERROR("[MNNVL AllReduceTwoShotRMSNorm] Unsupported world_size" + \ - std::to_string(params.nRanks) + \ - ". Supported sizes: {2, 4, 8, 16, 32, 64}"); \ - return cudaErrorInvalidValue; \ +#define DISPATCH_WORLD_SIZE(USE_CGA, EXTRA_AR) \ + switch (params.nRanks) { \ + case 2: \ + DISPATCH_LOADS(2, USE_CGA, EXTRA_AR); \ + break; \ + case 4: \ + DISPATCH_LOADS(4, USE_CGA, EXTRA_AR); \ + break; \ + case 8: \ + DISPATCH_LOADS(8, USE_CGA, EXTRA_AR); \ + break; \ + case 16: \ + DISPATCH_LOADS(16, USE_CGA, EXTRA_AR); \ + break; \ + case 32: \ + DISPATCH_LOADS(32, USE_CGA, EXTRA_AR); \ + break; \ + case 64: \ + DISPATCH_LOADS(64, USE_CGA, EXTRA_AR); \ + break; \ + default: \ + FLASHINFER_ERROR("[MNNVL AllReduceTwoShotRMSNorm] Unsupported world_size" + \ + std::to_string(params.nRanks) + \ + ". Supported sizes: {2, 4, 8, 16, 32, 64}"); \ + return cudaErrorInvalidValue; \ } - if (rnUseCGA) { - DISPATCH_WORLD_SIZE(true); + if (useExtraAR) { + DISPATCH_WORLD_SIZE(false, true); + } else if (rnUseCGA) { + DISPATCH_WORLD_SIZE(true, false); } else { - DISPATCH_WORLD_SIZE(false); + DISPATCH_WORLD_SIZE(false, false); } #undef DISPATCH_WORLD_SIZE From 0789210fd63d0a9be81f0704b557da17cfeb5951 Mon Sep 17 00:00:00 2001 From: Shiyu Li Date: Wed, 20 May 2026 20:33:27 -0700 Subject: [PATCH 05/15] Revert to separate kernels for 2-shot --- .../comm/trtllm_mnnvl_allreduce.cuh | 664 +++++++++--------- 1 file changed, 313 insertions(+), 351 deletions(-) diff --git a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh index 2b226f5d597..e638b138a8b 100644 --- a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh +++ b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh @@ -888,68 +888,26 @@ enum MNNVLTwoShotStage : uint8_t { using utils::copyF4; -template -__global__ __launch_bounds__((RMSNormFusion ? 1024 : 128)) void twoshotAllreduceKernel( +template +__global__ __launch_bounds__(128) void twoshotAllreduceKernel( AllReduceKernelParams params) { - constexpr int kELTS_PER_LOAD = sizeof(PackedType) / sizeof(T); + constexpr int kELTS_PER_THREAD = sizeof(PackedType) / sizeof(T); constexpr uint32_t kELT_SIZE = sizeof(T); - constexpr bool kExtraAR = RMSNormFusion && ExtraAR; - constexpr uint32_t kARLoadsPerThread = kExtraAR ? 1 : LoadsPerThread; - uint32_t const token = blockIdx.x; - uint32_t const blockSize = blockDim.x; - uint32_t const threadOffset = threadIdx.x; - bool const runRMSNormTail = !kExtraAR || blockIdx.y == 0; + int const packedIdx = blockIdx.y * blockDim.x + threadIdx.x; + int const token = blockIdx.x; + int const tokenOffset = packedIdx * kELTS_PER_THREAD; + int const threadOffset = token * params.tokenDim + tokenOffset; + bool const inBounds = tokenOffset < params.tokenDim; + int const destRank = token % WorldSize; + int const destTokenOffset = token / WorldSize; - uint32_t numThreads = blockSize; - uint32_t clusterSize = 1; - uint32_t clusterBlockRank = 0; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if constexpr (RMSNormFusion && UseCGA) { - namespace cg = cooperative_groups; - cg::cluster_group cluster = cg::this_cluster(); - numThreads = cluster.num_threads(); - clusterSize = cluster.num_blocks(); - clusterBlockRank = cluster.block_rank(); - } -#endif - - uint32_t const dimPadded = round_up(params.tokenDim, kELTS_PER_LOAD * numThreads); - uint32_t const elemsPerThread = dimPadded / numThreads; - uint32_t const loadStride = blockSize; - uint32_t const blockChunkSize = - ceil_div(params.tokenDim, clusterSize * kELTS_PER_LOAD) * kELTS_PER_LOAD; - uint32_t const baseTokenOffset = - RMSNormFusion ? clusterBlockRank * blockChunkSize : blockIdx.y * blockSize * kELTS_PER_LOAD; - uint32_t const arBaseTokenOffset = - kExtraAR ? blockIdx.y * blockSize * kELTS_PER_LOAD : baseTokenOffset; - uint32_t const rmsBaseTokenOffset = kExtraAR ? 0 : baseTokenOffset; - uint32_t const rmsBlockChunkSize = kExtraAR ? params.tokenDim : blockChunkSize; - uint32_t const destRank = token % WorldSize; - uint32_t const destTokenOffset = token / WorldSize; - - extern __shared__ uint8_t smem[]; - uint32_t const smemBufferSize = blockSize * elemsPerThread * sizeof(T); - T* smemInput = reinterpret_cast(&smem[0]); - T* smemResidual = reinterpret_cast(&smem[smemBufferSize]); - T* smemGamma = reinterpret_cast(&smem[2 * smemBufferSize]); - - alignas(16) __shared__ uint64_t residualStageBarrier; - alignas(16) __shared__ uint64_t gammaStageBarrier; -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if constexpr (RMSNormFusion) { - if (runRMSNormTail && threadIdx.x == 0) { - utils::mbarrierInit(&residualStageBarrier, 1); - utils::mbarrierInit(&gammaStageBarrier, 1); - } - __syncthreads(); - } cudaGridDependencySynchronize(); #endif - LamportFlags flag(params.bufferFlags, - MNNVLTwoShotStage::NUM_STAGES); + LamportFlags flag(params.bufferFlags, MNNVLTwoShotStage::NUM_STAGES); T* scatterBufLocal = reinterpret_cast( flag.getCurLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::SCATTER)); T* scatterBufDest = reinterpret_cast( @@ -959,237 +917,273 @@ __global__ __launch_bounds__((RMSNormFusion ? 1024 : 128)) void twoshotAllreduce T* broadcastBufR = reinterpret_cast( flag.getCurLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::BROADCAST)); - PackedVec shardVals[kARLoadsPerThread]; -#pragma unroll - for (uint32_t i = 0; i < kARLoadsPerThread; i++) { - uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - uint32_t const tokenOffset = arBaseTokenOffset + chunkOffset; - if (tokenOffset < params.tokenDim) { - shardVals[i].packed = - loadPacked(¶ms.shardPtr[token * params.tokenDim + tokenOffset]); + PackedVec val; + if (inBounds) { + val.packed = loadPacked(¶ms.shardPtr[threadOffset]); #pragma unroll - for (int j = 0; j < kELTS_PER_LOAD; j++) { - if (isNegZero(shardVals[i].elements[j])) { - shardVals[i].elements[j] = fromFloat(0.F); - } + for (int i = 0; i < kELTS_PER_THREAD; i++) { + if (isNegZero(val.elements[i])) { + val.elements[i] = fromFloat(0.F); } - reinterpret_cast( - &scatterBufDest[destTokenOffset * params.tokenDim * WorldSize + - params.rank * params.tokenDim])[tokenOffset / kELTS_PER_LOAD] = - shardVals[i].packed; } + + reinterpret_cast( + &scatterBufDest[destTokenOffset * params.tokenDim * WorldSize + + params.rank * params.tokenDim])[packedIdx] = val.packed; } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); #endif - if constexpr (!kExtraAR) { - flag.ctaArrive(); - } flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::SCATTER); - if (destRank == params.rank) { - uint32_t const localToken = token / WorldSize; -#pragma unroll - for (uint32_t i = 0; i < kARLoadsPerThread; i++) { - uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - uint32_t const tokenOffset = arBaseTokenOffset + chunkOffset; - if (tokenOffset < params.tokenDim) { - PackedVec remoteValues[WorldSize - 1]; - while (1) { - bool valid = true; + if (inBounds && destRank == params.rank) { + int const localToken = token / WorldSize; + PackedVec remoteValues[WorldSize - 1]; + while (1) { + bool valid = true; #pragma unroll - for (int r = 0; r < WorldSize - 1; r++) { - int const rankToLoad = (r + params.rank + 1) % WorldSize; - auto loaded = loadPackedVolatile( - &scatterBufLocal[localToken * params.tokenDim * WorldSize + - rankToLoad * params.tokenDim + tokenOffset]); - remoteValues[r].packed = loaded.packed; - valid &= !isLamportDirty(loaded); - } - if (valid) { - break; - } - } + for (int r = 0; r < WorldSize - 1; r++) { + int const rankToLoad = (r + params.rank + 1) % WorldSize; + auto loaded = loadPackedVolatile( + &scatterBufLocal[localToken * params.tokenDim * WorldSize + + rankToLoad * params.tokenDim + tokenOffset]); + remoteValues[r].packed = loaded.packed; + valid &= !isLamportDirty(loaded); + } + if (valid) { + break; + } + } - PackedVec packedAccum; + PackedVec packedAccum; #pragma unroll - for (int j = 0; j < kELTS_PER_LOAD; j++) { - float accum = toFloat(shardVals[i].elements[j]); + for (int i = 0; i < kELTS_PER_THREAD; i++) { + float accum = toFloat(val.elements[i]); #pragma unroll - for (int r = 0; r < WorldSize - 1; r++) { - accum += toFloat(remoteValues[r].elements[j]); - } - packedAccum.elements[j] = fromFloat(accum); - } - // Reduced values can round to the dirty sentinel; sanitize before broadcast polling. - sanitizeLamportPayload(packedAccum); - reinterpret_cast(&broadcastBufW[token * params.tokenDim + tokenOffset])[0] = - packedAccum.packed; + for (int r = 0; r < WorldSize - 1; r++) { + accum += toFloat(remoteValues[r].elements[i]); } + packedAccum.elements[i] = fromFloat(accum); } + // Reduced values can round to the dirty sentinel; sanitize before broadcast polling. + sanitizeLamportPayload(packedAccum); + reinterpret_cast(&broadcastBufW[token * params.tokenDim])[packedIdx] = + packedAccum.packed; } flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::BROADCAST); - if (!runRMSNormTail) { + if constexpr (WaitForResults) { + // OOB threads still arrive so waitAndUpdate sees every CTA. flag.ctaArrive(); - return; + + if (inBounds) { + auto loaded = loadPackedVolatile(&broadcastBufR[threadOffset]); + while (isLamportDirty(loaded)) { + loaded = loadPackedVolatile(&broadcastBufR[threadOffset]); + } + reinterpret_cast(¶ms.outputPtr[threadOffset])[0] = loaded.packed; + } + + flag.waitAndUpdate( + {static_cast(round_up(params.numTokens, WorldSize) * params.tokenDim * + kELT_SIZE), + static_cast(params.numTokens * params.tokenDim * kELT_SIZE), 0, 0}); } +} - uint32_t stageElems = 0; - uint32_t stageBytes = 0; - if constexpr (RMSNormFusion) { - stageElems = rmsBaseTokenOffset < params.tokenDim - ? (rmsBlockChunkSize < params.tokenDim - rmsBaseTokenOffset - ? rmsBlockChunkSize - : params.tokenDim - rmsBaseTokenOffset) - : 0; - stageBytes = stageElems * sizeof(T); +template +__global__ __launch_bounds__(1024) void rmsNormLamport(AllReduceKernelParams params) { + constexpr int kELTS_PER_LOAD = sizeof(PackedType) / sizeof(T); + constexpr uint32_t kELT_SIZE = sizeof(T); + + uint32_t const token = blockIdx.x; + uint32_t const blockSize = blockDim.x; + uint32_t const threadOffset = threadIdx.x; + + uint32_t numThreads = blockSize; + uint32_t clusterSize = 1; + uint32_t clusterBlockRank = 0; +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + if constexpr (UseCGA) { + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + numThreads = cluster.num_threads(); + clusterSize = cluster.num_blocks(); + clusterBlockRank = cluster.block_rank(); } +#endif + + uint32_t const dimPadded = round_up(params.tokenDim, kELTS_PER_LOAD * numThreads); + uint32_t const elemsPerThread = dimPadded / numThreads; + uint32_t const loadStride = blockSize; + uint32_t const blockChunkSize = + ceil_div(params.tokenDim, clusterSize * kELTS_PER_LOAD) * kELTS_PER_LOAD; + uint32_t const baseTokenOffset = clusterBlockRank * blockChunkSize; + + extern __shared__ uint8_t smem[]; + uint32_t const smemBufferSize = blockSize * elemsPerThread * sizeof(T); + T* smemInput = reinterpret_cast(&smem[0]); + T* smemResidual = reinterpret_cast(&smem[smemBufferSize]); + T* smemGamma = reinterpret_cast(&smem[2 * smemBufferSize]); + + alignas(16) __shared__ uint64_t residualStageBarrier; + alignas(16) __shared__ uint64_t gammaStageBarrier; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if constexpr (RMSNormFusion) { - if (stageBytes > 0 && threadIdx.x == 0) { - T const* residualSrc = ¶ms.residualInPtr[token * params.tokenDim + rmsBaseTokenOffset]; - T const* gammaSrc = ¶ms.gammaPtr[rmsBaseTokenOffset]; - // The mbarrier expect_tx must be issued before the TMA copy that completes it. - utils::mbarrierArriveExpectTx(&residualStageBarrier, stageBytes); - utils::cpAsyncBulkGlobalToShared(smemResidual, residualSrc, &residualStageBarrier, - stageBytes); - utils::mbarrierArriveExpectTx(&gammaStageBarrier, stageBytes); - utils::cpAsyncBulkGlobalToShared(smemGamma, gammaSrc, &gammaStageBarrier, stageBytes); - } + if (threadIdx.x == 0) { + utils::mbarrierInit(&residualStageBarrier, 1); + utils::mbarrierInit(&gammaStageBarrier, 1); + } + __syncthreads(); + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + LamportFlags flag(params.bufferFlags, MNNVLTwoShotStage::NUM_STAGES); + T* input = reinterpret_cast( + flag.getCurLamportBuf(params.localBufferPtr, MNNVLTwoShotStage::BROADCAST)); + + uint32_t const stageElems = + baseTokenOffset < params.tokenDim + ? (blockChunkSize < params.tokenDim - baseTokenOffset + ? blockChunkSize + : params.tokenDim - baseTokenOffset) + : 0; + uint32_t const stageBytes = stageElems * sizeof(T); +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + if (stageBytes > 0 && threadIdx.x == 0) { + T const* residualSrc = ¶ms.residualInPtr[token * params.tokenDim + baseTokenOffset]; + T const* gammaSrc = ¶ms.gammaPtr[baseTokenOffset]; + // The mbarrier expect_tx must be issued before the TMA copy that completes it. + utils::mbarrierArriveExpectTx(&residualStageBarrier, stageBytes); + utils::cpAsyncBulkGlobalToShared(smemResidual, residualSrc, &residualStageBarrier, stageBytes); + utils::mbarrierArriveExpectTx(&gammaStageBarrier, stageBytes); + utils::cpAsyncBulkGlobalToShared(smemGamma, gammaSrc, &gammaStageBarrier, stageBytes); } #else - if constexpr (RMSNormFusion) { #pragma unroll - for (uint32_t i = 0; i < LoadsPerThread; i++) { - uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - uint32_t const tokenOffset = rmsBaseTokenOffset + chunkOffset; - if (tokenOffset < params.tokenDim) { - copyF4(&smemResidual[chunkOffset], - ¶ms.residualInPtr[token * params.tokenDim + tokenOffset]); - } + for (uint32_t i = 0; i < LoadsPerThread; i++) { + uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; + uint32_t const tokenOffset = baseTokenOffset + chunkOffset; + if (tokenOffset < params.tokenDim) { + copyF4(&smemResidual[chunkOffset], + ¶ms.residualInPtr[token * params.tokenDim + tokenOffset]); } - __pipeline_commit(); + } + __pipeline_commit(); #pragma unroll - for (uint32_t i = 0; i < LoadsPerThread; i++) { - uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - uint32_t const tokenOffset = rmsBaseTokenOffset + chunkOffset; - if (tokenOffset < params.tokenDim) { - copyF4(&smemGamma[chunkOffset], ¶ms.gammaPtr[tokenOffset]); - } + for (uint32_t i = 0; i < LoadsPerThread; i++) { + uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; + uint32_t const tokenOffset = baseTokenOffset + chunkOffset; + if (tokenOffset < params.tokenDim) { + copyF4(&smemGamma[chunkOffset], ¶ms.gammaPtr[tokenOffset]); } - __pipeline_commit(); } + __pipeline_commit(); #endif + flag.ctaArrive(); + #pragma unroll for (uint32_t i = 0; i < LoadsPerThread; i++) { uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - uint32_t const tokenOffset = rmsBaseTokenOffset + chunkOffset; + uint32_t const tokenOffset = baseTokenOffset + chunkOffset; if (tokenOffset < params.tokenDim) { auto loaded = - loadPackedVolatile(&broadcastBufR[token * params.tokenDim + tokenOffset]); + loadPackedVolatile(&input[token * params.tokenDim + tokenOffset]); while (isLamportDirty(loaded)) { - loaded = - loadPackedVolatile(&broadcastBufR[token * params.tokenDim + tokenOffset]); - } - if constexpr (RMSNormFusion) { - reinterpret_cast(&smemInput[chunkOffset])[0] = loaded.packed; - } else { - reinterpret_cast(¶ms.outputPtr[token * params.tokenDim + tokenOffset])[0] = - loaded.packed; + loaded = loadPackedVolatile(&input[token * params.tokenDim + tokenOffset]); } + reinterpret_cast(&smemInput[chunkOffset])[0] = loaded.packed; } } - if constexpr (RMSNormFusion) { - float rInput[LoadsPerThread * kELTS_PER_LOAD]; - float threadSum = 0.f; + float rInput[LoadsPerThread * kELTS_PER_LOAD]; + float threadSum = 0.f; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if (stageElems > 0) { - utils::mbarrierWait(&residualStageBarrier, 0); - } + if (stageElems > 0) { + utils::mbarrierWait(&residualStageBarrier, 0); + } #else - __pipeline_wait_prior(1); + __pipeline_wait_prior(1); #endif #pragma unroll - for (uint32_t i = 0; i < LoadsPerThread; i++) { - uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - uint32_t const tokenOffset = rmsBaseTokenOffset + chunkOffset; - if (tokenOffset < params.tokenDim) { - PackedVec inp{.packed = loadPacked(&smemInput[chunkOffset])}; - PackedVec res{.packed = loadPacked(&smemResidual[chunkOffset])}; - PackedVec inpPlusRes = inp + res; - reinterpret_cast( - ¶ms.prenormedPtr[token * params.tokenDim + tokenOffset])[0] = inpPlusRes.packed; + for (uint32_t i = 0; i < LoadsPerThread; i++) { + uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; + uint32_t const tokenOffset = baseTokenOffset + chunkOffset; + if (tokenOffset < params.tokenDim) { + PackedVec inp{.packed = loadPacked(&smemInput[chunkOffset])}; + PackedVec res{.packed = loadPacked(&smemResidual[chunkOffset])}; + PackedVec inpPlusRes = inp + res; + reinterpret_cast(¶ms.prenormedPtr[token * params.tokenDim + tokenOffset])[0] = + inpPlusRes.packed; #pragma unroll - for (int j = 0; j < kELTS_PER_LOAD; j++) { - float const value = toFloat(inpPlusRes.elements[j]); - rInput[i * kELTS_PER_LOAD + j] = value; - threadSum += value * value; - } + for (int j = 0; j < kELTS_PER_LOAD; j++) { + float const value = toFloat(inpPlusRes.elements[j]); + rInput[i * kELTS_PER_LOAD + j] = value; + threadSum += value * value; } } + } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if (stageElems > 0) { - utils::mbarrierWait(&gammaStageBarrier, 0); - } + if (stageElems > 0) { + utils::mbarrierWait(&gammaStageBarrier, 0); + } #else - __pipeline_wait_prior(0); + __pipeline_wait_prior(0); #endif - float blockSum = blockReduceSum(threadSum); - float fullSum = blockSum; - __shared__ float sharedVal[8]; + float blockSum = blockReduceSum(threadSum); + float fullSum = blockSum; + __shared__ float sharedVal[8]; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if constexpr (UseCGA) { - namespace cg = cooperative_groups; - cg::cluster_group cluster = cg::this_cluster(); - int const numBlocks = cluster.num_blocks(); - if (numBlocks > 1) { - fullSum = 0.F; - int const blockRank = cluster.block_rank(); - if (threadIdx.x < numBlocks) { - cluster.map_shared_rank(&sharedVal[0], threadIdx.x)[blockRank] = blockSum; - } - cluster.barrier_wait(cluster.barrier_arrive()); - for (int i = 0; i < numBlocks; ++i) { - fullSum += sharedVal[i]; - } + if constexpr (UseCGA) { + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + int const numBlocks = cluster.num_blocks(); + if (numBlocks > 1) { + fullSum = 0.F; + int const blockRank = cluster.block_rank(); + if (threadIdx.x < numBlocks) { + cluster.map_shared_rank(&sharedVal[0], threadIdx.x)[blockRank] = blockSum; + } + cluster.barrier_wait(cluster.barrier_arrive()); + for (int i = 0; i < numBlocks; ++i) { + fullSum += sharedVal[i]; } } + } #endif - float const rcpRms = rsqrtf(fullSum / params.tokenDim + params.epsilon); + float const rcpRms = rsqrtf(fullSum / params.tokenDim + params.epsilon); #pragma unroll - for (uint32_t i = 0; i < LoadsPerThread; i++) { - uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; - uint32_t const tokenOffset = rmsBaseTokenOffset + chunkOffset; - if (tokenOffset < params.tokenDim) { - PackedVec gamma{.packed = loadPacked(&smemGamma[chunkOffset])}; - PackedVec out; + for (uint32_t i = 0; i < LoadsPerThread; i++) { + uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; + uint32_t const tokenOffset = baseTokenOffset + chunkOffset; + if (tokenOffset < params.tokenDim) { + PackedVec gamma{.packed = loadPacked(&smemGamma[chunkOffset])}; + PackedVec out; #pragma unroll - for (int j = 0; j < kELTS_PER_LOAD; j++) { - out.elements[j] = fromFloat((params.weightBias + toFloat(gamma.elements[j])) * - rInput[i * kELTS_PER_LOAD + j] * rcpRms); - } - reinterpret_cast(¶ms.outputPtr[token * params.tokenDim + tokenOffset])[0] = - out.packed; + for (int j = 0; j < kELTS_PER_LOAD; j++) { + out.elements[j] = fromFloat((params.weightBias + toFloat(gamma.elements[j])) * + rInput[i * kELTS_PER_LOAD + j] * rcpRms); } + reinterpret_cast(¶ms.outputPtr[token * params.tokenDim + tokenOffset])[0] = + out.packed; } } - if constexpr (kExtraAR) { - flag.ctaArrive(); - } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaGridDependencySynchronize(); +#endif + flag.waitAndUpdate( - {static_cast(round_up(params.numTokens, WorldSize) * params.tokenDim * kELT_SIZE), + {static_cast(round_up(params.numTokens, params.nRanks) * params.tokenDim * + kELT_SIZE), static_cast(params.numTokens * params.tokenDim * kELT_SIZE), 0, 0}); } @@ -1221,57 +1215,63 @@ cudaError_t twoshotAllreduceFusionDispatch(AllReduceFusionParams const& params) .bufferFlags = params.bufferFlags, }; - if (!params.rmsNormFusion) { - dim3 arGrid(numTokens, arNumBlocksPerToken); - - cudaLaunchAttribute arAttrs[1]; - arAttrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - arAttrs[0].val.programmaticStreamSerializationAllowed = params.launchWithPdl ? 1 : 0; - - cudaLaunchConfig_t arConfig{ - .gridDim = arGrid, - .blockDim = 128, - .dynamicSmemBytes = 0, - .stream = params.stream, - .attrs = arAttrs, - .numAttrs = 1, - }; - - FLASHINFER_LOG_DEBUG( - "[MNNVL AllReduceTwoShot] Dispatch: grid size: (%d, %d, 1), block_size: 128", numTokens, - arNumBlocksPerToken); - -#define LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE) \ - FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ - &arConfig, &twoshotAllreduceKernel, kernelParams)); - switch (params.nRanks) { - case 2: - LAUNCH_ALLREDUCE_KERNEL(2); - break; - case 4: - LAUNCH_ALLREDUCE_KERNEL(4); - break; - case 8: - LAUNCH_ALLREDUCE_KERNEL(8); - break; - case 16: - LAUNCH_ALLREDUCE_KERNEL(16); - break; - case 32: - LAUNCH_ALLREDUCE_KERNEL(32); - break; - case 64: - LAUNCH_ALLREDUCE_KERNEL(64); - break; - default: - FLASHINFER_ERROR("[MNNVL AllReduceTwoShot] Unsupported world_size" + - std::to_string(params.nRanks) + - ". Supported sizes: {2, 4, 8, 16, 32, 64}"); - return cudaErrorInvalidValue; - } -#undef LAUNCH_ALLREDUCE_KERNEL + dim3 arGrid(numTokens, arNumBlocksPerToken); + + cudaLaunchAttribute arAttrs[1]; + arAttrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + arAttrs[0].val.programmaticStreamSerializationAllowed = params.launchWithPdl ? 1 : 0; + + cudaLaunchConfig_t arConfig{ + .gridDim = arGrid, + .blockDim = 128, + .dynamicSmemBytes = 0, + .stream = params.stream, + .attrs = arAttrs, + .numAttrs = 1, + }; + + FLASHINFER_LOG_DEBUG( + "[MNNVL AllReduceTwoShot] Dispatch: grid size: (%d, %d, 1), block_size: 128", numTokens, + arNumBlocksPerToken); + +#define LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, WAIT_FOR_RESULTS) \ + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ + &arConfig, &twoshotAllreduceKernel, kernelParams)); +#define DISPATCH_ALLREDUCE_KERNEL(WAIT_FOR_RESULTS) \ + switch (params.nRanks) { \ + case 2: \ + LAUNCH_ALLREDUCE_KERNEL(2, WAIT_FOR_RESULTS); \ + break; \ + case 4: \ + LAUNCH_ALLREDUCE_KERNEL(4, WAIT_FOR_RESULTS); \ + break; \ + case 8: \ + LAUNCH_ALLREDUCE_KERNEL(8, WAIT_FOR_RESULTS); \ + break; \ + case 16: \ + LAUNCH_ALLREDUCE_KERNEL(16, WAIT_FOR_RESULTS); \ + break; \ + case 32: \ + LAUNCH_ALLREDUCE_KERNEL(32, WAIT_FOR_RESULTS); \ + break; \ + case 64: \ + LAUNCH_ALLREDUCE_KERNEL(64, WAIT_FOR_RESULTS); \ + break; \ + default: \ + FLASHINFER_ERROR("[MNNVL AllReduceTwoShot] Unsupported world_size" + \ + std::to_string(params.nRanks) + \ + ". Supported sizes: {2, 4, 8, 16, 32, 64}"); \ + return cudaErrorInvalidValue; \ + } + + if (params.rmsNormFusion) { + DISPATCH_ALLREDUCE_KERNEL(false); + } else { + DISPATCH_ALLREDUCE_KERNEL(true); return cudaSuccess; } +#undef DISPATCH_ALLREDUCE_KERNEL +#undef LAUNCH_ALLREDUCE_KERNEL static const int kSMVersionMajor = GetCudaComputeCapability().first; auto gridConfig = adjustGridConfig(numTokens, tokenDim, numEltsPerThread, kSMVersionMajor); @@ -1294,23 +1294,17 @@ cudaError_t twoshotAllreduceFusionDispatch(AllReduceFusionParams const& params) return cudaErrorInvalidValue; } - bool const useExtraAR = params.nRanks == 8 && !rnUseCGA && rnLoadsPerThread == 1 && - arNumBlocksPerToken > rnClusterSize && arNumBlocksPerToken <= 8; - int const launchBlockSize = useExtraAR ? 128 : rnBlockSize; - int const launchGridY = useExtraAR ? arNumBlocksPerToken : rnClusterSize; - int const launchLoadsPerThread = useExtraAR ? arNumBlocksPerToken : rnLoadsPerThread; int const rnNumThreads = rnClusterSize * rnBlockSize; - int const launchNumThreads = useExtraAR ? launchBlockSize : rnNumThreads; - int const dimPadded = round_up(tokenDim, numEltsPerThread * launchNumThreads); - int const iters = dimPadded / launchNumThreads; - size_t const smemSize = 3 * launchBlockSize * iters * sizeof(T); + int const dimPadded = round_up(tokenDim, numEltsPerThread * rnNumThreads); + int const iters = dimPadded / rnNumThreads; + size_t const smemSize = 3 * rnBlockSize * iters * sizeof(T); - dim3 rnGrid(numTokens, launchGridY, 1); + dim3 rnGrid(numTokens, rnClusterSize, 1); cudaLaunchConfig_t rnConfig; cudaLaunchAttribute rnAttrs[2]; rnConfig.stream = params.stream; rnConfig.gridDim = rnGrid; - rnConfig.blockDim = launchBlockSize; + rnConfig.blockDim = rnBlockSize; rnConfig.dynamicSmemBytes = smemSize; rnConfig.attrs = rnAttrs; rnAttrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; @@ -1326,90 +1320,58 @@ cudaError_t twoshotAllreduceFusionDispatch(AllReduceFusionParams const& params) FLASHINFER_LOG_DEBUG( "[MNNVL AllReduceTwoShotRMSNorm] Dispatch: grid size: (%d, %d, 1), block_size: %d, " - "cluster_size: %d, loads_per_thread: %d, threads_needed: %d, extra_ar: %d", - numTokens, launchGridY, launchBlockSize, rnClusterSize, launchLoadsPerThread, - ceil_div(tokenDim, numEltsPerThread), static_cast(useExtraAR)); - -#define LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, USE_CGA, LOADS_PER_THREAD, EXTRA_AR) \ - FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( \ - &twoshotAllreduceKernel, \ - cudaFuncAttributeMaxDynamicSharedMemorySize, smemSize)); \ - FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ - &rnConfig, \ - &twoshotAllreduceKernel, \ - kernelParams)); - -#define DISPATCH_LOADS(WORLD_SIZE, USE_CGA, EXTRA_AR) \ - switch (launchLoadsPerThread) { \ - case 1: \ - LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, USE_CGA, 1, EXTRA_AR); \ - break; \ - case 2: \ - LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 2, EXTRA_AR); \ - break; \ - case 3: \ - LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 3, EXTRA_AR); \ - break; \ - case 4: \ - LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 4, EXTRA_AR); \ - break; \ - case 5: \ - LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 5, EXTRA_AR); \ - break; \ - case 6: \ - LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 6, EXTRA_AR); \ - break; \ - case 7: \ - LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 7, EXTRA_AR); \ - break; \ - case 8: \ - LAUNCH_TWOSHOT_FUSION_TYPED(WORLD_SIZE, false, 8, EXTRA_AR); \ - break; \ - default: \ + "cluster_size: %d, loads_per_thread: %d, threads_needed: %d", + numTokens, rnClusterSize, rnBlockSize, rnClusterSize, rnLoadsPerThread, + ceil_div(tokenDim, numEltsPerThread)); + +#define LAUNCH_RMSNORM_KERNEL(USE_CGA, LOADS_PER_THREAD) \ + FLASHINFER_CUDA_CALL(cudaFuncSetAttribute(&rmsNormLamport, \ + cudaFuncAttributeMaxDynamicSharedMemorySize, \ + smemSize)); \ + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ + &rnConfig, &rmsNormLamport, kernelParams)); + +#define DISPATCH_LOADS(USE_CGA) \ + switch (rnLoadsPerThread) { \ + case 1: \ + LAUNCH_RMSNORM_KERNEL(USE_CGA, 1); \ + break; \ + case 2: \ + LAUNCH_RMSNORM_KERNEL(false, 2); \ + break; \ + case 3: \ + LAUNCH_RMSNORM_KERNEL(false, 3); \ + break; \ + case 4: \ + LAUNCH_RMSNORM_KERNEL(false, 4); \ + break; \ + case 5: \ + LAUNCH_RMSNORM_KERNEL(false, 5); \ + break; \ + case 6: \ + LAUNCH_RMSNORM_KERNEL(false, 6); \ + break; \ + case 7: \ + LAUNCH_RMSNORM_KERNEL(false, 7); \ + break; \ + case 8: \ + LAUNCH_RMSNORM_KERNEL(false, 8); \ + break; \ + default: \ FLASHINFER_ERROR("[MNNVL AllReduceTwoShotRMSNorm] Unsupported loads_per_thread " + \ - std::to_string(launchLoadsPerThread) + \ - ". Supported sizes: {1, 2, 3, 4, 5, 6, 7, 8}"); \ - return cudaErrorInvalidValue; \ + std::to_string(rnLoadsPerThread) + \ + ". Supported sizes: {1, 2, 3, 4, 5, 6, 7, 8}"); \ + return cudaErrorInvalidValue; \ } -#define DISPATCH_WORLD_SIZE(USE_CGA, EXTRA_AR) \ - switch (params.nRanks) { \ - case 2: \ - DISPATCH_LOADS(2, USE_CGA, EXTRA_AR); \ - break; \ - case 4: \ - DISPATCH_LOADS(4, USE_CGA, EXTRA_AR); \ - break; \ - case 8: \ - DISPATCH_LOADS(8, USE_CGA, EXTRA_AR); \ - break; \ - case 16: \ - DISPATCH_LOADS(16, USE_CGA, EXTRA_AR); \ - break; \ - case 32: \ - DISPATCH_LOADS(32, USE_CGA, EXTRA_AR); \ - break; \ - case 64: \ - DISPATCH_LOADS(64, USE_CGA, EXTRA_AR); \ - break; \ - default: \ - FLASHINFER_ERROR("[MNNVL AllReduceTwoShotRMSNorm] Unsupported world_size" + \ - std::to_string(params.nRanks) + \ - ". Supported sizes: {2, 4, 8, 16, 32, 64}"); \ - return cudaErrorInvalidValue; \ - } - - if (useExtraAR) { - DISPATCH_WORLD_SIZE(false, true); - } else if (rnUseCGA) { - DISPATCH_WORLD_SIZE(true, false); + if (rnUseCGA) { + DISPATCH_LOADS(true); } else { - DISPATCH_WORLD_SIZE(false, false); + DISPATCH_LOADS(false); } -#undef DISPATCH_WORLD_SIZE #undef DISPATCH_LOADS -#undef LAUNCH_TWOSHOT_FUSION_TYPED +#undef LAUNCH_RMSNORM_KERNEL return cudaSuccess; } } // namespace trtllm_mnnvl_allreduce From 2ea7f2e0d86d35b3e4acd343301f477dc58d131f Mon Sep 17 00:00:00 2001 From: Shiyu Li Date: Wed, 20 May 2026 21:07:08 -0700 Subject: [PATCH 06/15] Clean up; Use JIT CUDA Arch guard. --- flashinfer/jit/comm.py | 6 + .../comm/trtllm_mnnvl_allreduce.cuh | 118 +----------------- 2 files changed, 7 insertions(+), 117 deletions(-) diff --git a/flashinfer/jit/comm.py b/flashinfer/jit/comm.py index 834f77f921b..52250f5454d 100644 --- a/flashinfer/jit/comm.py +++ b/flashinfer/jit/comm.py @@ -35,11 +35,17 @@ def gen_comm_alltoall_module() -> JitSpec: def gen_trtllm_mnnvl_comm_module() -> JitSpec: + # This MNNVL path is supported on Hopper/Blackwell datacenter targets only. + # Thor (SM11x) and RTX/consumer (SM12x) are intentionally excluded here. + nvcc_flags = current_compilation_context.get_nvcc_flags_list( + supported_major_versions=[9, 10] + ) return gen_jit_spec( "trtllm_mnnvl_comm", [ jit_env.FLASHINFER_CSRC_DIR / "trtllm_mnnvl_allreduce.cu", ], + extra_cuda_cflags=nvcc_flags, ) diff --git a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh index e638b138a8b..71846b4dffe 100644 --- a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh +++ b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh @@ -247,7 +247,6 @@ __device__ struct __attribute__((aligned(32))) LamportFlags { } __device__ void ctaArrive() { -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) if constexpr (UseCGA) { cg::cluster_group cluster = cg::this_cluster(); __cluster_barrier_arrive(); @@ -257,7 +256,6 @@ __device__ struct __attribute__((aligned(32))) LamportFlags { } return; } -#endif uint32_t const barrierThreads = round_up(static_cast(blockDim.x), 32u); // Named CTA barrier avoids a full __syncthreads() while still ordering payload stores // before the per-CTA Lamport arrival counter update. @@ -272,7 +270,6 @@ __device__ struct __attribute__((aligned(32))) LamportFlags { __device__ void waitAndUpdate(uint4 bytesToClearPerStage) { bool isLastCtaT0{false}; int targetCount{0}; -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cg::grid_group grid = cg::this_grid(); // Use the first thread instead of the last thread as the last thread may exit early isLastCtaT0 = grid.thread_rank() == 0; @@ -282,10 +279,6 @@ __device__ struct __attribute__((aligned(32))) LamportFlags { } else { targetCount = gridDim.x * gridDim.y * gridDim.z; } -#else - isLastCtaT0 = threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0; - targetCount = gridDim.x * gridDim.y * gridDim.z; -#endif if (isLastCtaT0) { uint4* flagPtr = reinterpret_cast(mBufferFlagsPtr); while (*reinterpret_cast(mFlagAccessPtr) < targetCount) { @@ -425,50 +418,6 @@ inline __device__ void copyF4(T_IN* dst, T_IN const* src) { __pipeline_memcpy_async(dst4, src4, sizeof(float4)); } -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) -inline __device__ void mbarrierInit(void* barrier, uint32_t count) { - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" - : - : "r"(static_cast(__cvta_generic_to_shared(barrier))), "r"(count) - : "memory"); -} - -inline __device__ void mbarrierArriveExpectTx(void* barrier, uint32_t bytes) { - asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;\n" - : - : "r"(static_cast(__cvta_generic_to_shared(barrier))), "r"(bytes) - : "memory"); -} - -inline __device__ void mbarrierWait(void* barrier, int phase) { - uint32_t barrierPtr = static_cast(__cvta_generic_to_shared(barrier)); - asm volatile( - "{\n\t" - ".reg .pred P1;\n\t" - "WAIT:\n\t" - "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1;\n\t" - "@P1 bra.uni DONE;\n\t" - "bra.uni WAIT;\n\t" - "DONE:\n\t" - "}" - : - : "r"(barrierPtr), "r"(phase) - : "memory"); -} - -inline __device__ void cpAsyncBulkGlobalToShared(void* dst, void const* src, void* barrier, - uint32_t bytes) { - uint32_t smemPtr = static_cast(__cvta_generic_to_shared(dst)); - uint64_t gmemPtr = static_cast(__cvta_generic_to_global(const_cast(src))); - uint32_t barrierPtr = static_cast(__cvta_generic_to_shared(barrier)); - asm volatile( - "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];\n" - : - : "r"(smemPtr), "l"(gmemPtr), "r"(bytes), "r"(barrierPtr) - : "memory"); -} -#endif - uint32_t constexpr kWARP_SIZE = 32U; uint32_t constexpr kLOG2_WARP_SIZE = 5U; uint32_t constexpr kLANE_ID_MASK = 0x1f; @@ -647,7 +596,6 @@ __global__ void __launch_bounds__(1024) constexpr int kELTS_PER_THREAD = sizeof(PackedType) / sizeof(T); constexpr uint32_t kELT_SIZE = sizeof(T); int const tokenDim = params.tokenDim; -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) namespace cg = cooperative_groups; cg::cluster_group cluster = cg::this_cluster(); int packedIdx = cluster.thread_rank(); @@ -655,18 +603,8 @@ __global__ void __launch_bounds__(1024) int threadOffset = token * tokenDim + packedIdx * kELTS_PER_THREAD; cudaGridDependencySynchronize(); -#else - int packedIdx = blockIdx.y * blockDim.x + threadIdx.x; - int token = blockIdx.x; - // Offset w.r.t. the input shard - int threadOffset = token * tokenDim + packedIdx * kELTS_PER_THREAD; -#endif // We only use 1 stage for the oneshot allreduce -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) constexpr bool kUseLamportCGA = true; -#else - constexpr bool kUseLamportCGA = false; -#endif LamportFlags flag(params.bufferFlags, 1); T* stagePtrMcast = reinterpret_cast(flag.getCurLamportBuf(params.mcastPtr, 0)); T* stagePtrLocal = reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[params.rank], 0)); @@ -729,9 +667,7 @@ __global__ void __launch_bounds__(1024) for (int i = 0; i < kELTS_PER_THREAD; i++) { packedAccum.elements[i] = fromFloat(accum[i]); } -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); -#endif if constexpr (RMSNormFusion) { // =============================== Residual =============================== PackedVec residualIn; @@ -753,8 +689,6 @@ __global__ void __launch_bounds__(1024) __shared__ float sharedVal[8]; // Temporary variable to share the sum within block float fullSum = blockSum; -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cg::cluster_group cluster = cg::this_cluster(); int const numBlocks = cluster.num_blocks(); if (numBlocks > 1) { fullSum = 0.F; @@ -768,7 +702,6 @@ __global__ void __launch_bounds__(1024) fullSum += sharedVal[i]; } } -#endif float rcpRms = rsqrtf(fullSum / tokenDim + params.epsilon); #pragma unroll for (int i = 0; i < kELTS_PER_THREAD; i++) { @@ -903,9 +836,7 @@ __global__ __launch_bounds__(128) void twoshotAllreduceKernel( int const destRank = token % WorldSize; int const destTokenOffset = token / WorldSize; -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaGridDependencySynchronize(); -#endif LamportFlags flag(params.bufferFlags, MNNVLTwoShotStage::NUM_STAGES); T* scatterBufLocal = reinterpret_cast( @@ -932,9 +863,7 @@ __global__ __launch_bounds__(128) void twoshotAllreduceKernel( params.rank * params.tokenDim])[packedIdx] = val.packed; } -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); -#endif flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::SCATTER); if (inBounds && destRank == params.rank) { @@ -1005,7 +934,6 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(AllReduceKernelParams uint32_t numThreads = blockSize; uint32_t clusterSize = 1; uint32_t clusterBlockRank = 0; -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) if constexpr (UseCGA) { namespace cg = cooperative_groups; cg::cluster_group cluster = cg::this_cluster(); @@ -1013,7 +941,6 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(AllReduceKernelParams clusterSize = cluster.num_blocks(); clusterBlockRank = cluster.block_rank(); } -#endif uint32_t const dimPadded = round_up(params.tokenDim, kELTS_PER_LOAD * numThreads); uint32_t const elemsPerThread = dimPadded / numThreads; @@ -1028,39 +955,12 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(AllReduceKernelParams T* smemResidual = reinterpret_cast(&smem[smemBufferSize]); T* smemGamma = reinterpret_cast(&smem[2 * smemBufferSize]); - alignas(16) __shared__ uint64_t residualStageBarrier; - alignas(16) __shared__ uint64_t gammaStageBarrier; -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if (threadIdx.x == 0) { - utils::mbarrierInit(&residualStageBarrier, 1); - utils::mbarrierInit(&gammaStageBarrier, 1); - } - __syncthreads(); cudaTriggerProgrammaticLaunchCompletion(); -#endif LamportFlags flag(params.bufferFlags, MNNVLTwoShotStage::NUM_STAGES); T* input = reinterpret_cast( flag.getCurLamportBuf(params.localBufferPtr, MNNVLTwoShotStage::BROADCAST)); - uint32_t const stageElems = - baseTokenOffset < params.tokenDim - ? (blockChunkSize < params.tokenDim - baseTokenOffset - ? blockChunkSize - : params.tokenDim - baseTokenOffset) - : 0; - uint32_t const stageBytes = stageElems * sizeof(T); -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if (stageBytes > 0 && threadIdx.x == 0) { - T const* residualSrc = ¶ms.residualInPtr[token * params.tokenDim + baseTokenOffset]; - T const* gammaSrc = ¶ms.gammaPtr[baseTokenOffset]; - // The mbarrier expect_tx must be issued before the TMA copy that completes it. - utils::mbarrierArriveExpectTx(&residualStageBarrier, stageBytes); - utils::cpAsyncBulkGlobalToShared(smemResidual, residualSrc, &residualStageBarrier, stageBytes); - utils::mbarrierArriveExpectTx(&gammaStageBarrier, stageBytes); - utils::cpAsyncBulkGlobalToShared(smemGamma, gammaSrc, &gammaStageBarrier, stageBytes); - } -#else #pragma unroll for (uint32_t i = 0; i < LoadsPerThread; i++) { uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; @@ -1080,7 +980,6 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(AllReduceKernelParams } } __pipeline_commit(); -#endif flag.ctaArrive(); @@ -1101,13 +1000,8 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(AllReduceKernelParams float rInput[LoadsPerThread * kELTS_PER_LOAD]; float threadSum = 0.f; -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if (stageElems > 0) { - utils::mbarrierWait(&residualStageBarrier, 0); - } -#else + // Residual and gamma staging use normal cp.async groups; residual is consumed first. __pipeline_wait_prior(1); -#endif #pragma unroll for (uint32_t i = 0; i < LoadsPerThread; i++) { @@ -1129,18 +1023,11 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(AllReduceKernelParams } } -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - if (stageElems > 0) { - utils::mbarrierWait(&gammaStageBarrier, 0); - } -#else __pipeline_wait_prior(0); -#endif float blockSum = blockReduceSum(threadSum); float fullSum = blockSum; __shared__ float sharedVal[8]; -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) if constexpr (UseCGA) { namespace cg = cooperative_groups; cg::cluster_group cluster = cg::this_cluster(); @@ -1157,7 +1044,6 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(AllReduceKernelParams } } } -#endif float const rcpRms = rsqrtf(fullSum / params.tokenDim + params.epsilon); #pragma unroll @@ -1177,9 +1063,7 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(AllReduceKernelParams } } -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaGridDependencySynchronize(); -#endif flag.waitAndUpdate( {static_cast(round_up(params.numTokens, params.nRanks) * params.tokenDim * From 0f2e0fc045a9391f2d2b9cdce6a1d7bf1c4efbf2 Mon Sep 17 00:00:00 2001 From: Shiyu Li Date: Wed, 20 May 2026 21:18:09 -0700 Subject: [PATCH 07/15] Simplify the grid dispatch logic and add comments. --- .../comm/trtllm_mnnvl_allreduce.cuh | 68 ++++++++----------- 1 file changed, 29 insertions(+), 39 deletions(-) diff --git a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh index 71846b4dffe..a8b5c7075af 100644 --- a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh +++ b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh @@ -511,35 +511,41 @@ inline __device__ T blockReduceSum(T val) { return blockReduceSumFull(val); } } -// A helper function to tune the grid configuration for fused oneshot and rmsnorm kernels -// Return (block_size, cluster_size, loads_per_thread) +// Tune the grid configuration for fused oneshot and RMSNorm kernels. +// Return (block_size, cluster_size, loads_per_thread). std::tuple adjustGridConfig(int numTokens, int dim, int eltsPerThread, - int smVersionMajor) { - // Start with preferred block_size and cluster_size - int clusterSize = smVersionMajor >= 9 ? 8 : 1; + bool useCluster) { + // Step 1: start from the widest cluster we are willing to launch. MNNVL JIT only + // targets SM90/SM100, but RMSNorm may request a no-CGA fallback for multi-load rows. + int clusterSize = useCluster ? 8 : 1; int blockSize = 128; - // ========================== Adjust the grid configuration ========================== int threadsNeeded = ceil_div(dim, eltsPerThread); int loadsPerThread = 1; blockSize = ceil_div(threadsNeeded, clusterSize); - if (smVersionMajor >= 9) { + if (useCluster) { + // Step 2: shrink the cluster until the hidden dimension partitions cleanly across CTAs. + // This keeps each CTA responsible for an integral chunk of packed elements. while (threadsNeeded % clusterSize != 0 && clusterSize > 1) { clusterSize /= 2; } int const maxDivisibleClusterSize = clusterSize; blockSize = ceil_div(threadsNeeded, clusterSize); + // Step 3: if divisibility leaves each CTA too small, trade cluster width for at least + // a 128-thread CTA. This improves occupancy and avoids tiny CTAs when the row is narrow. while (blockSize < 128 && clusterSize >= 2) { blockSize *= 2; clusterSize /= 2; } int smCount = GetCudaMultiProcessorCount(); + // Step 4: if the token grid already has enough CTAs to cover the GPU, reduce cluster + // width and make CTAs larger. This avoids over-partitioning one token across too many CTAs. while (numTokens * clusterSize > smCount && clusterSize > 1 && blockSize <= 512) { blockSize *= 2; clusterSize /= 2; } - // If the token grid still underfills the GPU, prefer more CTAs per token even when that - // makes each CTA smaller than the usual 128-thread target. + // Step 5: if the token grid still underfills the GPU, restore cluster width up to the + // divisibility limit. We accept 64-thread CTAs here to expose more CTAs per token. while (clusterSize < maxDivisibleClusterSize) { int const candidateClusterSize = clusterSize * 2; int const candidateBlockSize = ceil_div(threadsNeeded, candidateClusterSize); @@ -550,31 +556,18 @@ std::tuple adjustGridConfig(int numTokens, int dim, int eltsPerTh blockSize = candidateBlockSize; } } - // Trying to scale up use multiple loads or CGA + // Step 6: for very wide rows, first increase cluster width on SM90+ and then increase + // per-thread loads. The goal is to keep block_size within CUDA's 1024-thread limit. while (blockSize > 1024) { - if (smVersionMajor >= 9) { - if (clusterSize < 8) { - clusterSize = clusterSize << 1; - } else { - if (loadsPerThread < 8) { - loadsPerThread += 1; - } else { - break; - } - } + if (useCluster && clusterSize < 8) { + clusterSize = clusterSize << 1; + } else if (loadsPerThread < 8) { + loadsPerThread += 1; } else { - if (loadsPerThread < 8) { - loadsPerThread += 1; - } else { - break; - } + break; } blockSize = ceil_div(threadsNeeded, clusterSize * loadsPerThread); } - while (smVersionMajor >= 9 && blockSize > 1024 && loadsPerThread < 8) { - loadsPerThread += 1; - blockSize = ceil_div(threadsNeeded, clusterSize * loadsPerThread); - } return {blockSize, clusterSize, loadsPerThread}; } }; // namespace utils @@ -722,10 +715,8 @@ cudaError_t oneshotAllreduceFusionDispatch(AllReduceFusionParams const& params) int const tokenDim = params.tokenDim; int const eltsPerThread = sizeof(float4) / sizeof(T); - static const int kSMVersionMajor = GetCudaComputeCapability().first; - auto [blockSize, clusterSize, loadsPerThread] = - adjustGridConfig(numTokens, tokenDim, eltsPerThread, kSMVersionMajor); + adjustGridConfig(numTokens, tokenDim, eltsPerThread, true); dim3 grid(numTokens, clusterSize, 1); FLASHINFER_LOG_DEBUG( @@ -738,7 +729,7 @@ cudaError_t oneshotAllreduceFusionDispatch(AllReduceFusionParams const& params) FLASHINFER_CHECK(blockSize <= 1024 && loadsPerThread == 1, "Hidden Dimension %d exceeds the maximum supported hidden dimension (%d)", - tokenDim, 1024 * (kSMVersionMajor >= 9 ? 8 : 1) * eltsPerThread); + tokenDim, 1024 * 8 * eltsPerThread); cudaLaunchAttribute attrs[2]; attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; @@ -754,7 +745,7 @@ cudaError_t oneshotAllreduceFusionDispatch(AllReduceFusionParams const& params) .dynamicSmemBytes = 0, .stream = params.stream, .attrs = attrs, - .numAttrs = kSMVersionMajor >= 9 ? 2 : 1, + .numAttrs = 2, }; #define LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, RMSNORM) \ @@ -1157,15 +1148,14 @@ cudaError_t twoshotAllreduceFusionDispatch(AllReduceFusionParams const& params) #undef DISPATCH_ALLREDUCE_KERNEL #undef LAUNCH_ALLREDUCE_KERNEL - static const int kSMVersionMajor = GetCudaComputeCapability().first; - auto gridConfig = adjustGridConfig(numTokens, tokenDim, numEltsPerThread, kSMVersionMajor); + auto gridConfig = adjustGridConfig(numTokens, tokenDim, numEltsPerThread, true); int rnBlockSize = std::get<0>(gridConfig); int rnClusterSize = std::get<1>(gridConfig); int rnLoadsPerThread = std::get<2>(gridConfig); - bool rnUseCGA = kSMVersionMajor >= 9 && rnClusterSize > 1 && rnLoadsPerThread == 1; - if (kSMVersionMajor >= 9 && rnClusterSize > 1 && !rnUseCGA) { - gridConfig = adjustGridConfig(numTokens, tokenDim, numEltsPerThread, 0); + bool rnUseCGA = rnClusterSize > 1 && rnLoadsPerThread == 1; + if (rnClusterSize > 1 && !rnUseCGA) { + gridConfig = adjustGridConfig(numTokens, tokenDim, numEltsPerThread, false); rnBlockSize = std::get<0>(gridConfig); rnClusterSize = std::get<1>(gridConfig); rnLoadsPerThread = std::get<2>(gridConfig); From bc659085d4f08c1da03c45f4e3b785e5079f13f6 Mon Sep 17 00:00:00 2001 From: Shiyu Li Date: Wed, 20 May 2026 21:23:16 -0700 Subject: [PATCH 08/15] Fix pre-commit error. --- .../comm/trtllm_mnnvl_allreduce.cuh | 154 +++++++++--------- 1 file changed, 74 insertions(+), 80 deletions(-) diff --git a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh index a8b5c7075af..7004c27e17b 100644 --- a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh +++ b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh @@ -812,10 +812,8 @@ enum MNNVLTwoShotStage : uint8_t { using utils::copyF4; -template -__global__ __launch_bounds__(128) void twoshotAllreduceKernel( - AllReduceKernelParams params) { +template +__global__ __launch_bounds__(128) void twoshotAllreduceKernel(AllReduceKernelParams params) { constexpr int kELTS_PER_THREAD = sizeof(PackedType) / sizeof(T); constexpr uint32_t kELT_SIZE = sizeof(T); @@ -849,9 +847,9 @@ __global__ __launch_bounds__(128) void twoshotAllreduceKernel( } } - reinterpret_cast( - &scatterBufDest[destTokenOffset * params.tokenDim * WorldSize + - params.rank * params.tokenDim])[packedIdx] = val.packed; + reinterpret_cast(&scatterBufDest[destTokenOffset * params.tokenDim * WorldSize + + params.rank * params.tokenDim])[packedIdx] = + val.packed; } cudaTriggerProgrammaticLaunchCompletion(); @@ -907,8 +905,7 @@ __global__ __launch_bounds__(128) void twoshotAllreduceKernel( } flag.waitAndUpdate( - {static_cast(round_up(params.numTokens, WorldSize) * params.tokenDim * - kELT_SIZE), + {static_cast(round_up(params.numTokens, WorldSize) * params.tokenDim * kELT_SIZE), static_cast(params.numTokens * params.tokenDim * kELT_SIZE), 0, 0}); } } @@ -979,8 +976,7 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(AllReduceKernelParams uint32_t const chunkOffset = (i * loadStride + threadOffset) * kELTS_PER_LOAD; uint32_t const tokenOffset = baseTokenOffset + chunkOffset; if (tokenOffset < params.tokenDim) { - auto loaded = - loadPackedVolatile(&input[token * params.tokenDim + tokenOffset]); + auto loaded = loadPackedVolatile(&input[token * params.tokenDim + tokenOffset]); while (isLamportDirty(loaded)) { loaded = loadPackedVolatile(&input[token * params.tokenDim + tokenOffset]); } @@ -1002,8 +998,8 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(AllReduceKernelParams PackedVec inp{.packed = loadPacked(&smemInput[chunkOffset])}; PackedVec res{.packed = loadPacked(&smemResidual[chunkOffset])}; PackedVec inpPlusRes = inp + res; - reinterpret_cast(¶ms.prenormedPtr[token * params.tokenDim + tokenOffset])[0] = - inpPlusRes.packed; + reinterpret_cast( + ¶ms.prenormedPtr[token * params.tokenDim + tokenOffset])[0] = inpPlusRes.packed; #pragma unroll for (int j = 0; j < kELTS_PER_LOAD; j++) { @@ -1056,10 +1052,9 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(AllReduceKernelParams cudaGridDependencySynchronize(); - flag.waitAndUpdate( - {static_cast(round_up(params.numTokens, params.nRanks) * params.tokenDim * - kELT_SIZE), - static_cast(params.numTokens * params.tokenDim * kELT_SIZE), 0, 0}); + flag.waitAndUpdate({static_cast(round_up(params.numTokens, params.nRanks) * + params.tokenDim * kELT_SIZE), + static_cast(params.numTokens * params.tokenDim * kELT_SIZE), 0, 0}); } template @@ -1105,38 +1100,37 @@ cudaError_t twoshotAllreduceFusionDispatch(AllReduceFusionParams const& params) .numAttrs = 1, }; - FLASHINFER_LOG_DEBUG( - "[MNNVL AllReduceTwoShot] Dispatch: grid size: (%d, %d, 1), block_size: 128", numTokens, - arNumBlocksPerToken); + FLASHINFER_LOG_DEBUG("[MNNVL AllReduceTwoShot] Dispatch: grid size: (%d, %d, 1), block_size: 128", + numTokens, arNumBlocksPerToken); #define LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, WAIT_FOR_RESULTS) \ FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ &arConfig, &twoshotAllreduceKernel, kernelParams)); -#define DISPATCH_ALLREDUCE_KERNEL(WAIT_FOR_RESULTS) \ - switch (params.nRanks) { \ - case 2: \ - LAUNCH_ALLREDUCE_KERNEL(2, WAIT_FOR_RESULTS); \ - break; \ - case 4: \ - LAUNCH_ALLREDUCE_KERNEL(4, WAIT_FOR_RESULTS); \ - break; \ - case 8: \ - LAUNCH_ALLREDUCE_KERNEL(8, WAIT_FOR_RESULTS); \ - break; \ - case 16: \ - LAUNCH_ALLREDUCE_KERNEL(16, WAIT_FOR_RESULTS); \ - break; \ - case 32: \ - LAUNCH_ALLREDUCE_KERNEL(32, WAIT_FOR_RESULTS); \ - break; \ - case 64: \ - LAUNCH_ALLREDUCE_KERNEL(64, WAIT_FOR_RESULTS); \ - break; \ - default: \ - FLASHINFER_ERROR("[MNNVL AllReduceTwoShot] Unsupported world_size" + \ - std::to_string(params.nRanks) + \ - ". Supported sizes: {2, 4, 8, 16, 32, 64}"); \ - return cudaErrorInvalidValue; \ +#define DISPATCH_ALLREDUCE_KERNEL(WAIT_FOR_RESULTS) \ + switch (params.nRanks) { \ + case 2: \ + LAUNCH_ALLREDUCE_KERNEL(2, WAIT_FOR_RESULTS); \ + break; \ + case 4: \ + LAUNCH_ALLREDUCE_KERNEL(4, WAIT_FOR_RESULTS); \ + break; \ + case 8: \ + LAUNCH_ALLREDUCE_KERNEL(8, WAIT_FOR_RESULTS); \ + break; \ + case 16: \ + LAUNCH_ALLREDUCE_KERNEL(16, WAIT_FOR_RESULTS); \ + break; \ + case 32: \ + LAUNCH_ALLREDUCE_KERNEL(32, WAIT_FOR_RESULTS); \ + break; \ + case 64: \ + LAUNCH_ALLREDUCE_KERNEL(64, WAIT_FOR_RESULTS); \ + break; \ + default: \ + FLASHINFER_ERROR("[MNNVL AllReduceTwoShot] Unsupported world_size" + \ + std::to_string(params.nRanks) + \ + ". Supported sizes: {2, 4, 8, 16, 32, 64}"); \ + return cudaErrorInvalidValue; \ } if (params.rmsNormFusion) { @@ -1200,42 +1194,42 @@ cudaError_t twoshotAllreduceFusionDispatch(AllReduceFusionParams const& params) #define LAUNCH_RMSNORM_KERNEL(USE_CGA, LOADS_PER_THREAD) \ FLASHINFER_CUDA_CALL(cudaFuncSetAttribute(&rmsNormLamport, \ - cudaFuncAttributeMaxDynamicSharedMemorySize, \ - smemSize)); \ - FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ - &rnConfig, &rmsNormLamport, kernelParams)); - -#define DISPATCH_LOADS(USE_CGA) \ - switch (rnLoadsPerThread) { \ - case 1: \ - LAUNCH_RMSNORM_KERNEL(USE_CGA, 1); \ - break; \ - case 2: \ - LAUNCH_RMSNORM_KERNEL(false, 2); \ - break; \ - case 3: \ - LAUNCH_RMSNORM_KERNEL(false, 3); \ - break; \ - case 4: \ - LAUNCH_RMSNORM_KERNEL(false, 4); \ - break; \ - case 5: \ - LAUNCH_RMSNORM_KERNEL(false, 5); \ - break; \ - case 6: \ - LAUNCH_RMSNORM_KERNEL(false, 6); \ - break; \ - case 7: \ - LAUNCH_RMSNORM_KERNEL(false, 7); \ - break; \ - case 8: \ - LAUNCH_RMSNORM_KERNEL(false, 8); \ - break; \ - default: \ + cudaFuncAttributeMaxDynamicSharedMemorySize, \ + smemSize)); \ + FLASHINFER_CUDA_CALL( \ + cudaLaunchKernelEx(&rnConfig, &rmsNormLamport, kernelParams)); + +#define DISPATCH_LOADS(USE_CGA) \ + switch (rnLoadsPerThread) { \ + case 1: \ + LAUNCH_RMSNORM_KERNEL(USE_CGA, 1); \ + break; \ + case 2: \ + LAUNCH_RMSNORM_KERNEL(false, 2); \ + break; \ + case 3: \ + LAUNCH_RMSNORM_KERNEL(false, 3); \ + break; \ + case 4: \ + LAUNCH_RMSNORM_KERNEL(false, 4); \ + break; \ + case 5: \ + LAUNCH_RMSNORM_KERNEL(false, 5); \ + break; \ + case 6: \ + LAUNCH_RMSNORM_KERNEL(false, 6); \ + break; \ + case 7: \ + LAUNCH_RMSNORM_KERNEL(false, 7); \ + break; \ + case 8: \ + LAUNCH_RMSNORM_KERNEL(false, 8); \ + break; \ + default: \ FLASHINFER_ERROR("[MNNVL AllReduceTwoShotRMSNorm] Unsupported loads_per_thread " + \ - std::to_string(rnLoadsPerThread) + \ - ". Supported sizes: {1, 2, 3, 4, 5, 6, 7, 8}"); \ - return cudaErrorInvalidValue; \ + std::to_string(rnLoadsPerThread) + \ + ". Supported sizes: {1, 2, 3, 4, 5, 6, 7, 8}"); \ + return cudaErrorInvalidValue; \ } if (rnUseCGA) { From 99b1ec7cdb66f44a5274c27b03a0fde1361873dc Mon Sep 17 00:00:00 2001 From: Shiyu Li Date: Thu, 21 May 2026 10:57:53 -0700 Subject: [PATCH 09/15] Add FP8/NVFP4 quant fusion --- csrc/trtllm_mnnvl_allreduce.cu | 84 +++- flashinfer/comm/allreduce.py | 102 ++++- flashinfer/comm/trtllm_mnnvl_ar.py | 194 +++++++- .../comm/trtllm_mnnvl_allreduce.cuh | 279 +++++++++-- tests/comm/test_trtllm_mnnvl_allreduce.py | 433 ++++++++++++++++++ 5 files changed, 1044 insertions(+), 48 deletions(-) diff --git a/csrc/trtllm_mnnvl_allreduce.cu b/csrc/trtllm_mnnvl_allreduce.cu index 03badc47b72..de8cbc963bc 100644 --- a/csrc/trtllm_mnnvl_allreduce.cu +++ b/csrc/trtllm_mnnvl_allreduce.cu @@ -3,6 +3,7 @@ using namespace flashinfer::trtllm_mnnvl_allreduce; +using flashinfer::QuantizationSFLayout; using tvm::ffi::Optional; #define DISPATCH_FLOATING_TYPES_FOR_MNNVL_ALLREDUCE(dtype, c_type, ...) \ @@ -30,9 +31,12 @@ void trtllm_mnnvl_allreduce_fusion(TensorView input, int64_t multicast_buffer_pt int64_t buffer_ptrs_dev, int64_t buffer_ptr_local, TensorView buffer_flags_mnnvl, int64_t nranks, int64_t rank, bool rmsnorm_fusion, bool launch_with_pdl, bool use_oneshot, - TensorView output, Optional residual_out, + Optional output, Optional residual_out, Optional residual_in, Optional gamma, - Optional epsilon, Optional weight_bias) { + Optional epsilon, Optional weight_bias, + Optional quant_type, Optional quant_out, + Optional sf_out, Optional output_scale, + Optional layout_code) { ffi::CUDADeviceGuard device_guard(input.device().device_id); auto stream = get_stream(input.device()); @@ -40,13 +44,23 @@ void trtllm_mnnvl_allreduce_fusion(TensorView input, int64_t multicast_buffer_pt // Extract parameters from tensors int64_t num_tokens = input.size(0); int64_t token_dim = input.size(1); + auto quant_type_enum = + quant_type.has_value() ? static_cast(quant_type.value()) : QuantType::kNone; + auto sf_layout = layout_code.has_value() + ? static_cast(layout_code.value()) + : QuantizationSFLayout::SWIZZLED_128x4; // Validate input parameters TVM_FFI_ICHECK_EQ(token_dim % (sizeof(float4) / sizeof(c_type)), 0) << "token_dim must be divisible by " << sizeof(float4) / sizeof(c_type); - TVM_FFI_ICHECK(output.size(0) == input.size(0) && output.size(1) == input.size(1)) - << "output shape mismatch: expected (" << input.size(0) << ", " << input.size(1) - << ") but got (" << output.size(0) << ", " << output.size(1) << ")"; + TVM_FFI_ICHECK(output.has_value() || quant_type_enum != QuantType::kNone) + << "output must be provided unless quantization fusion is enabled"; + if (output.has_value()) { + TVM_FFI_ICHECK(output.value().size(0) == input.size(0) && + output.value().size(1) == input.size(1)) + << "output shape mismatch: expected (" << input.size(0) << ", " << input.size(1) + << ") but got (" << output.value().size(0) << ", " << output.value().size(1) << ")"; + } TVM_FFI_ICHECK(nranks >= 2 && nranks <= 64) << "nranks must be between 2 and 64, got " << nranks; TVM_FFI_ICHECK(rank >= 0 && rank < nranks) @@ -56,6 +70,17 @@ void trtllm_mnnvl_allreduce_fusion(TensorView input, int64_t multicast_buffer_pt !rmsnorm_fusion) << "residual_in, residual_out, gamma, and epsilon must be provided if rmsnorm_fusion is " "true"; + TVM_FFI_ICHECK(quant_type_enum == QuantType::kNone || rmsnorm_fusion) + << "MNNVL quantization fusion requires rmsnorm_fusion=true"; + TVM_FFI_ICHECK(quant_type_enum == QuantType::kNone || + (output_scale.has_value() && + encode_dlpack_dtype(output_scale.value().dtype()) == float32_code)) + << "output_scale must be provided for MNNVL quantization fusion and must be float32"; + TVM_FFI_ICHECK(quant_type_enum == QuantType::kNone || + sf_layout != QuantizationSFLayout::SWIZZLED_8x4) + << "MNNVL quantization fusion supports SWIZZLED_128x4 or LINEAR scale layouts"; + TVM_FFI_ICHECK(quant_type_enum == QuantType::kNone || quant_out.has_value()) + << "quant_out must be provided when quantization fusion is enabled"; if (rmsnorm_fusion) { TVM_FFI_ICHECK(residual_in.value().size(0) == num_tokens && @@ -72,6 +97,44 @@ void trtllm_mnnvl_allreduce_fusion(TensorView input, int64_t multicast_buffer_pt << "gamma must have the same shape as token dimension (" << token_dim << ") but got (" << gamma.value().size(0) << ")"; } + switch (quant_type_enum) { + case QuantType::kNone: + break; + case QuantType::kFP8: + TVM_FFI_ICHECK(quant_out.value().size(0) == num_tokens && + quant_out.value().size(1) == token_dim) + << "quant_out shape mismatch for FP8: expected (" << num_tokens << ", " << token_dim + << ") but got (" << quant_out.value().size(0) << ", " << quant_out.value().size(1) + << ")"; + TVM_FFI_ICHECK(encode_dlpack_dtype(quant_out.value().dtype()) == float8_e4m3fn_code) + << "quant_out for FP8 must have dtype float8_e4m3fn"; + break; + case QuantType::kFP4: + TVM_FFI_ICHECK(sizeof(c_type) == 2) + << "NVFP4 MNNVL quantization fusion is only supported for FP16/BF16 inputs"; + TVM_FFI_ICHECK(token_dim % 16 == 0) + << "NVFP4 MNNVL quantization fusion requires token_dim divisible by 16"; + TVM_FFI_ICHECK(quant_out.value().size(0) == num_tokens && + quant_out.value().size(1) == token_dim / 2) + << "quant_out shape mismatch for FP4: expected (" << num_tokens << ", " << token_dim / 2 + << ") but got (" << quant_out.value().size(0) << ", " << quant_out.value().size(1) + << ")"; + TVM_FFI_ICHECK(encode_dlpack_dtype(quant_out.value().dtype()) == uint8_code || + encode_dlpack_dtype(quant_out.value().dtype()) == + encode_dlpack_dtype(dl_float4_e2m1fn_x2)) + << "quant_out for FP4 must have dtype uint8 or float4_e2m1fn_x2"; + TVM_FFI_ICHECK(sf_out.has_value()) + << "sf_out must be provided for NVFP4 MNNVL quantization fusion"; + TVM_FFI_ICHECK(encode_dlpack_dtype(sf_out.value().dtype()) == float8_e4m3fn_code) + << "sf_out for FP4 must have dtype float8_e4m3fn"; + TVM_FFI_ICHECK(sf_out.value().numel() >= num_tokens * token_dim / 16) + << "sf_out is too small for FP4: expected at least " << num_tokens * token_dim / 16 + << " elements but got " << sf_out.value().numel(); + break; + default: + TVM_FFI_LOG_AND_THROW(NotImplementedError) + << "Unsupported MNNVL quantization type " << static_cast(quant_type_enum); + } // Create the parameters struct AllReduceFusionParams params; @@ -87,6 +150,8 @@ void trtllm_mnnvl_allreduce_fusion(TensorView input, int64_t multicast_buffer_pt params.bufferFlags = reinterpret_cast(buffer_flags_mnnvl.data_ptr()); params.rmsNormFusion = rmsnorm_fusion; params.launchWithPdl = launch_with_pdl; + params.sfLayout = sf_layout; + params.quantType = quant_type_enum; // input data params.input = const_cast(input.data_ptr()); @@ -95,11 +160,18 @@ void trtllm_mnnvl_allreduce_fusion(TensorView input, int64_t multicast_buffer_pt params.gamma = gamma.has_value() ? const_cast(gamma.value().data_ptr()) : nullptr; params.epsilon = epsilon.has_value() ? epsilon.value() : 1e-5; params.weightBias = weight_bias.has_value() ? static_cast(weight_bias.value()) : 0.0f; + params.outputScale = output_scale.has_value() + ? reinterpret_cast(output_scale.value().data_ptr()) + : nullptr; // output data - params.output = const_cast(output.data_ptr()); + params.output = output.has_value() ? const_cast(output.value().data_ptr()) : nullptr; params.residualOut = residual_out.has_value() ? const_cast(residual_out.value().data_ptr()) : nullptr; + params.quantOut = + quant_out.has_value() ? reinterpret_cast(quant_out.value().data_ptr()) : nullptr; + params.scalingFactorOut = + sf_out.has_value() ? reinterpret_cast(sf_out.value().data_ptr()) : nullptr; params.stream = stream; cudaError_t status; diff --git a/flashinfer/comm/allreduce.py b/flashinfer/comm/allreduce.py index c0bb663bd58..2665436956f 100644 --- a/flashinfer/comm/allreduce.py +++ b/flashinfer/comm/allreduce.py @@ -74,10 +74,13 @@ # Note: AllReduceFusionPattern and QuantizationSFLayout are pseudo-types (classes with int constants) # Import them for runtime use but type hint as int for mypy compatibility from .trtllm_ar import AllReduceFusionPattern +from .trtllm_ar import QuantizationSFLayout from .trtllm_mnnvl_ar import MNNVLAllReduceFusionWorkspace from .trtllm_mnnvl_ar import MNNVLAllreduceFusionStrategy +from .trtllm_mnnvl_ar import MNNVLQuantType from .trtllm_mnnvl_ar import trtllm_mnnvl_allreduce from .trtllm_mnnvl_ar import trtllm_mnnvl_fused_allreduce_add_rmsnorm +from .trtllm_mnnvl_ar import trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant # ============================================================================ # WORKSPACE IMPLEMENTATIONS @@ -520,7 +523,8 @@ def allreduce_fusion( - kMoEFinalizeARResidualRMSNorm = 7 (trtllm only) - kARResidualRMSNormPerTokenGroupFP8PackedQuant = 8 (trtllm only) - kARResidualRMSNormOutPerTokenGroupFP8PackedQuant = 9 (trtllm only) - Note: MNNVL only supports patterns 0 and 1 + Note: MNNVL supports patterns 0-5; MoE and packed group quant + patterns remain TRTLLM-only. Note: MOE patterns (6-7) only support trtllm backend launch_with_pdl: Use Programmatic Dependent Launch trigger_completion_at_end: [trtllm only] Controls when PDL completion is signaled. @@ -533,8 +537,9 @@ def allreduce_fusion( output: AllReduce output [token_num, hidden_dim] residual_out: Prenorm output (after residual add, before norm) [token_num, hidden_dim] norm_out: Normalized output [token_num, hidden_dim] - quant_out: Quantized output [token_num, hidden_dim] [trtllm only] - scale_out: Quantization scale factors [trtllm only] + quant_out: Quantized output [token_num, hidden_dim] for FP8 or + [token_num, hidden_dim / 2] for NVFP4 + scale_out: NVFP4 quantization scale factors # ===== INPUT parameters ===== residual_in: Residual tensor to ADD [token_num, hidden_dim] @@ -546,8 +551,8 @@ def allreduce_fusion( Supported by both TRTLLM and MNNVL backends for kARResidualRMSNorm plus all TRTLLM RMSNorm variants (quant + MoE Reduction/Finalize). Ignored for kAllReduce (no normalization). - scale_factor: Input scale factor for quantization [trtllm only] - layout_code: Scale factor layout (QuantizationSFLayout) [trtllm only] + scale_factor: Input scale factor for quantization + layout_code: Scale factor layout (QuantizationSFLayout) # ===== Control parameters ===== use_oneshot: Use oneshot strategy vs twoshot @@ -683,9 +688,9 @@ def allreduce_fusion( residual_in=residual_in, rms_gamma=rms_gamma, rms_eps=rms_eps, - scale_factor=scale_factor - if isinstance(scale_factor, (int, float)) - else 1.0, + scale_factor=( + scale_factor if isinstance(scale_factor, (int, float)) else 1.0 + ), moe_reduction_device_num_experts=moe_reduction_device_num_experts, moe_reduction_scale_input=moe_reduction_scale_input, moe_reduction_active_experts_token_input=moe_reduction_active_experts_token_input, @@ -865,17 +870,51 @@ def _flatten_checked(t, name): return output elif isinstance(workspace, MNNVLAllReduceFusionWorkspace): - if ( - pattern != AllReduceFusionPattern.kARResidualRMSNorm - and pattern != AllReduceFusionPattern.kAllReduce + strategy = ( + MNNVLAllreduceFusionStrategy.AUTO + if use_oneshot is None + else ( + MNNVLAllreduceFusionStrategy.ONESHOT + if use_oneshot + else MNNVLAllreduceFusionStrategy.TWOSHOT + ) + ) + mnnvl_quant_patterns = { + AllReduceFusionPattern.kARResidualRMSNormFP8Quant: ( + MNNVLQuantType.FP8, + False, + ), + AllReduceFusionPattern.kARResidualRMSNormFP4Quant: ( + MNNVLQuantType.NVFP4, + False, + ), + AllReduceFusionPattern.kARResidualRMSNormOutFP8Quant: ( + MNNVLQuantType.FP8, + True, + ), + AllReduceFusionPattern.kARResidualRMSNormOutFP4Quant: ( + MNNVLQuantType.NVFP4, + True, + ), + } + if pattern not in ( + AllReduceFusionPattern.kAllReduce, + AllReduceFusionPattern.kARResidualRMSNorm, + *mnnvl_quant_patterns.keys(), ): raise ValueError( f"MNNVL AllReduce+RMS fusion does not support pattern {pattern}. Please try the TRTLLM backend instead." ) - if layout_code is not None: + mnnvl_layout_code = ( + QuantizationSFLayout.SWIZZLED_128x4 if layout_code is None else layout_code + ) + if ( + pattern in mnnvl_quant_patterns + and mnnvl_layout_code == QuantizationSFLayout.SWIZZLED_8x4 + ): raise ValueError( - "MNNVL AllReduce does not support quantization fusion and thus no layout_code" + "MNNVL quantization fusion supports SWIZZLED_128x4 or LINEAR scale layouts, not SWIZZLED_8x4" ) # MNNVL backend implementation @@ -888,6 +927,7 @@ def _flatten_checked(t, name): workspace=workspace, launch_with_pdl=launch_with_pdl, output=output, + strategy=strategy, ) return output @@ -915,10 +955,46 @@ def _flatten_checked(t, name): output=norm_out, residual_out=residual_out, launch_with_pdl=launch_with_pdl, + strategy=strategy, weight_bias=weight_bias, ) return norm_result + elif pattern in mnnvl_quant_patterns: + if residual_in is None: + raise ValueError( + "MNNVL quantized AllReduce+RMS fusion requires residual_in" + ) + if rms_gamma is None: + raise ValueError( + "MNNVL quantized AllReduce+RMS fusion requires rms_gamma" + ) + + quant_type, has_norm_out = mnnvl_quant_patterns[pattern] + if has_norm_out and norm_out is None: + norm_out = torch.empty_like(input) + if residual_out is None: + residual_out = torch.empty_like(input) + + quant_result, _, _, _ = trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant( + input=input, + residual_in=residual_in, + gamma=rms_gamma, + workspace=workspace, + epsilon=rms_eps, + output=norm_out if has_norm_out else None, + residual_out=residual_out, + quant_out=quant_out, + scale_out=scale_out, + output_scale=scale_factor, + layout_code=mnnvl_layout_code, + quant_type=quant_type, + launch_with_pdl=launch_with_pdl, + strategy=strategy, + weight_bias=weight_bias, + ) + return quant_result + else: raise ValueError(f"Unsupported pattern for MNNVL backend: {pattern}") diff --git a/flashinfer/comm/trtllm_mnnvl_ar.py b/flashinfer/comm/trtllm_mnnvl_ar.py index 32464ca98f1..3109fb210f8 100644 --- a/flashinfer/comm/trtllm_mnnvl_ar.py +++ b/flashinfer/comm/trtllm_mnnvl_ar.py @@ -7,7 +7,7 @@ import math import logging from types import SimpleNamespace -from typing import Optional, Tuple +from typing import Optional, Tuple, Union from enum import Enum import torch @@ -18,7 +18,9 @@ from ..jit import gen_trtllm_mnnvl_comm_module from ..utils import register_custom_op +from ..fp4_quantization import _compute_swizzled_layout_sf_size from .mnnvl import CommBackend, MPIBackend +from .trtllm_ar import QuantizationSFLayout from .workspace_base import AllReduceFusionWorkspace from .torch_symmetric_memory import _alloc_symm_buffer_bytes @@ -50,6 +52,12 @@ def select_strategy( MNNVL_ONE_SHOT_THRESHOLD = 64 * 1024 * 8 * 2 +class MNNVLQuantType: + NONE = 0 + FP8 = 1 + NVFP4 = 2 + + class MNNVLAllReduceFusionWorkspace(AllReduceFusionWorkspace): NUM_LAMPORT_BUFFERS = 3 @@ -280,6 +288,11 @@ def get_trtllm_mnnvl_comm_module(): "residual_in", "gamma", "epsilon", + "quant_type", + "quant_out", + "sf_out", + "output_scale", + "layout_code", ], ) def trtllm_mnnvl_allreduce_fusion( @@ -293,12 +306,17 @@ def trtllm_mnnvl_allreduce_fusion( rmsnorm_fusion: bool, launch_with_pdl: bool, use_oneshot: bool, - output: torch.Tensor, + output: Optional[torch.Tensor], residual_out: Optional[torch.Tensor], residual_in: Optional[torch.Tensor], gamma: Optional[torch.Tensor], epsilon: Optional[float], weight_bias: Optional[float] = None, + quant_type: int = MNNVLQuantType.NONE, + quant_out: Optional[torch.Tensor] = None, + sf_out: Optional[torch.Tensor] = None, + output_scale: Optional[torch.Tensor] = None, + layout_code: int = QuantizationSFLayout.SWIZZLED_128x4, ) -> None: """ Perform a multi-node NVLink all-reduce operation with fusion. @@ -335,6 +353,11 @@ def trtllm_mnnvl_allreduce_fusion( gamma, epsilon, weight_bias, + quant_type, + quant_out, + sf_out, + output_scale, + layout_code, ) return SimpleNamespace( @@ -519,6 +542,173 @@ def trtllm_mnnvl_fused_allreduce_add_rmsnorm( return output, residual_out +def trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant( + input: torch.Tensor, + residual_in: torch.Tensor, + gamma: torch.Tensor, + workspace: MNNVLAllReduceFusionWorkspace, + epsilon: Optional[float] = None, + output: Optional[torch.Tensor] = None, + residual_out: Optional[torch.Tensor] = None, + quant_out: Optional[torch.Tensor] = None, + scale_out: Optional[torch.Tensor] = None, + output_scale: Union[torch.Tensor, float, None] = None, + layout_code: int = QuantizationSFLayout.SWIZZLED_128x4, + quant_type: int = MNNVLQuantType.NONE, + launch_with_pdl: bool = False, + strategy: MNNVLAllreduceFusionStrategy = MNNVLAllreduceFusionStrategy.AUTO, + weight_bias: float = 0.0, +) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, Optional[torch.Tensor]]: + """Performs MNNVL AllReduce + Residual + RMSNorm + FP8/NVFP4 quantization.""" + + if epsilon is None: + epsilon = torch.finfo(input.dtype).eps + if output_scale is None: + output_scale = 1.0 + + if len(input.shape) != 2: + raise ValueError( + f"The input tensor must be 2D, got {len(input.shape)}D. The shape is {input.shape}." + ) + if len(residual_in.shape) != 2: + raise ValueError( + f"The residual input tensor must be 2D, got {len(residual_in.shape)}D. The shape is {residual_in.shape}." + ) + if residual_in.shape != input.shape: + raise ValueError( + f"residual_in shape must match input shape, got {residual_in.shape} and {input.shape}." + ) + if gamma.numel() != input.shape[1]: + raise ValueError( + f"The gamma tensor must have the same number of elements as the hidden dimension, got {gamma.numel()} elements but expected {input.shape[1]} elements." + ) + + if layout_code == QuantizationSFLayout.SWIZZLED_8x4: + raise ValueError( + "MNNVL NVFP4 quantization supports SWIZZLED_128x4 or LINEAR scale layouts, not SWIZZLED_8x4." + ) + + if residual_out is None: + residual_out = torch.empty_like(input) + elif residual_out.shape != input.shape: + raise ValueError( + f"residual_out shape must match input shape, got {residual_out.shape} and {input.shape}." + ) + + if output is not None and output.shape != input.shape: + raise ValueError( + f"output shape must match input shape, got {output.shape} and {input.shape}." + ) + + if isinstance(output_scale, torch.Tensor): + if output_scale.numel() < 1: + raise ValueError("output_scale must contain at least one element") + output_scale_tensor = output_scale.to(device=input.device, dtype=torch.float32) + else: + output_scale_tensor = torch.tensor( + [float(output_scale)], dtype=torch.float32, device=input.device + ) + + token_num, hidden_dim = input.shape + if quant_type == MNNVLQuantType.FP8: + if quant_out is None: + quant_out = torch.empty_like(input, dtype=torch.float8_e4m3fn) + elif quant_out.shape != input.shape: + raise ValueError( + f"quant_out shape must be {tuple(input.shape)} for FP8, got {tuple(quant_out.shape)}." + ) + if quant_out.dtype != torch.float8_e4m3fn: + raise ValueError( + f"quant_out dtype for FP8 must be float8_e4m3fn, got {quant_out.dtype}." + ) + scale_out = None + elif quant_type == MNNVLQuantType.NVFP4: + if input.dtype == torch.float32: + raise ValueError("MNNVL NVFP4 quantization requires fp16 or bf16 input.") + if hidden_dim % 16 != 0: + raise ValueError( + f"MNNVL NVFP4 quantization requires hidden_dim divisible by 16, got {hidden_dim}." + ) + expected_quant_shape = (token_num, hidden_dim // 2) + fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None) + if quant_out is None: + quant_out = torch.empty( + expected_quant_shape, dtype=torch.uint8, device=input.device + ) + elif tuple(quant_out.shape) != expected_quant_shape: + raise ValueError( + f"quant_out shape must be {expected_quant_shape} for NVFP4, got {tuple(quant_out.shape)}." + ) + if quant_out.dtype != torch.uint8 and not ( + fp4_dtype is not None and quant_out.dtype == fp4_dtype + ): + raise ValueError( + f"quant_out dtype for NVFP4 must be uint8 or float4_e2m1fn_x2, got {quant_out.dtype}." + ) + + if scale_out is None: + if layout_code == QuantizationSFLayout.LINEAR: + scale_out = torch.empty( + (token_num, hidden_dim // 16), + dtype=torch.float8_e4m3fn, + device=input.device, + ) + else: + scale_out = torch.empty( + _compute_swizzled_layout_sf_size(token_num, hidden_dim // 16), + dtype=torch.float8_e4m3fn, + device=input.device, + ) + elif scale_out.numel() < token_num * hidden_dim // 16: + raise ValueError( + f"scale_out is too small for NVFP4: got {scale_out.numel()} elements, need at least {token_num * hidden_dim // 16}." + ) + if scale_out.dtype != torch.float8_e4m3fn: + raise ValueError( + f"scale_out dtype for NVFP4 must be float8_e4m3fn, got {scale_out.dtype}." + ) + else: + raise ValueError(f"Unsupported MNNVL quant_type: {quant_type}") + + module = get_trtllm_mnnvl_comm_module() + + if strategy == MNNVLAllreduceFusionStrategy.AUTO: + strategy = MNNVLAllreduceFusionStrategy.select_strategy( + workspace.tp_size, token_num, hidden_dim, input.dtype + ) + if not workspace.is_buffer_size_sufficient( + workspace.tp_size, token_num, hidden_dim, input.dtype, strategy + ): + raise ValueError( + f"The buffer size in the given workspace is insufficient for the given problem size. Buffer: {workspace.buffer_size_bytes} bytes, Required: {workspace.get_required_buffer_size_bytes(workspace.tp_size, token_num, hidden_dim, input.dtype, strategy)} bytes." + ) + + module.trtllm_mnnvl_allreduce_fusion( + input, + workspace.mc_ptr, + workspace.uc_ptrs_dev, + workspace.uc_ptr_local, + workspace.buffer_flags, + workspace.tp_size, + workspace.rank, + True, + launch_with_pdl, + strategy == MNNVLAllreduceFusionStrategy.ONESHOT, + output, + residual_out, + residual_in, + gamma, + epsilon, + weight_bias, + quant_type, + quant_out, + scale_out, + output_scale_tensor, + layout_code, + ) + return quant_out, scale_out, residual_out, output + + # Legacy API that has been deprecated; Left for backward compatibility @deprecated( "get_allreduce_mnnvl_workspace is deprecated, use MNNVLAllReduceFusionWorkspace class to manage the workspace instead" diff --git a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh index 7004c27e17b..40d7f7baec0 100644 --- a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh +++ b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh @@ -14,21 +14,37 @@ * limitations under the License. */ +#pragma once + #include #include #include +#include #include #include +#include #include +#include #include #include "../exception.h" +#include "../fp4_layout.cuh" #include "../logging.h" #include "../utils.cuh" +#include "trtllm_allreduce_fusion.cuh" + namespace flashinfer { namespace trtllm_mnnvl_allreduce { +using flashinfer::QuantizationSFLayout; + +enum class QuantType : int { + kNone = 0, + kFP8 = 1, + kFP4 = 2, +}; + struct AllReduceFusionParams { int nRanks; int rank; @@ -48,9 +64,14 @@ struct AllReduceFusionParams { // 0 for standard RMSNorm (out = gamma * x * rsqrt(...)), // 1 for Gemma / Qwen3.5 (out = (1 + gamma) * x * rsqrt(...)). float weightBias = 0.f; + float* outputScale = nullptr; + QuantizationSFLayout sfLayout = QuantizationSFLayout::SWIZZLED_128x4; + QuantType quantType = QuantType::kNone; void* residualOut; void* output; + void* quantOut = nullptr; + void* scalingFactorOut = nullptr; cudaStream_t stream = nullptr; }; @@ -70,6 +91,10 @@ struct AllReduceKernelParams { int tokenDim; float epsilon; float weightBias; + float* outputScalePtr; + void* quantOutPtr; + void* scalingFactorOutPtr; + QuantizationSFLayout sfLayout; uint32_t* bufferFlags; }; @@ -583,9 +608,64 @@ using utils::PackedVec; using utils::sanitizeLamportPayload; using utils::toFloat; -template +namespace quant { + +template +inline __device__ void quant_fp8(PackedVec packedAccum, void* quantOutPtr, + float invOutputScale, uint32_t threadOffset) { + static_assert(ELTS_PER_THREAD == 8 || ELTS_PER_THREAD == 4, "ELTS_PER_THREAD must be 8 or 4"); + using QuantizedPackedType = std::conditional_t; + + auto quantOut = reinterpret_cast<__nv_fp8_e4m3*>(quantOutPtr); + PackedVec quantizedAccum; +#pragma unroll + for (int i = 0; i < ELTS_PER_THREAD; i++) { + quantizedAccum.elements[i] = + __nv_fp8_e4m3(toFloat(packedAccum.elements[i]) * invOutputScale); + } + reinterpret_cast(&quantOut[threadOffset])[0] = quantizedAccum.packed; +} + +template +inline __device__ void quant_nvfp4(PackedVec packedAccum, void* quantOutPtr, + void* sfOutPtr, float* outputScale, uint32_t tokenIdx, + uint32_t tokenDim, uint32_t packedIdx, + QuantizationSFLayout sfLayout) { +#if CUDA_VERSION >= 12080 + static_assert( + ELTS_PER_THREAD == 8 && (std::is_same_v || std::is_same_v), + "NVFP4 quantization fusion is only supported for FP16/BF16!"); + + auto packedAccumVec = *reinterpret_cast*>(&packedAccum); + auto sfOut = trtllm_allreduce_fusion::utils::cvt_quant_to_fp4_get_sf_out_offset< + uint32_t, trtllm_allreduce_fusion::details::CVT_FP4_SF_VEC_SIZE / ELTS_PER_THREAD>( + cuda::std::nullopt, tokenIdx, packedIdx, cuda::std::nullopt, tokenDim, + reinterpret_cast(sfOutPtr), sfLayout); + + uint32_t quantOutOffset = tokenIdx * tokenDim / ELTS_PER_THREAD + packedIdx; + reinterpret_cast(quantOutPtr)[quantOutOffset] = + trtllm_allreduce_fusion::utils::cvt_warp_fp16_to_fp4(packedAccumVec, + *outputScale, sfOut); +#else + (void)packedAccum; + (void)quantOutPtr; + (void)sfOutPtr; + (void)outputScale; + (void)tokenIdx; + (void)tokenDim; + (void)packedIdx; + (void)sfLayout; +#endif +} + +} // namespace quant + +template __global__ void __launch_bounds__(1024) oneshotAllreduceFusionKernel(AllReduceKernelParams params) { + static_assert(QType == QuantType::kNone || RMSNormFusion, + "Quant-only pattern without RMSNorm is not supported!"); constexpr int kELTS_PER_THREAD = sizeof(PackedType) / sizeof(T); constexpr uint32_t kELT_SIZE = sizeof(T); int const tokenDim = params.tokenDim; @@ -702,7 +782,21 @@ __global__ void __launch_bounds__(1024) (params.weightBias + toFloat(gamma.elements[i]))); } } - reinterpret_cast(¶ms.outputPtr[threadOffset])[0] = packedAccum.packed; + if (params.outputPtr != nullptr) { + reinterpret_cast(¶ms.outputPtr[threadOffset])[0] = packedAccum.packed; + } + if constexpr (QType == QuantType::kFP8) { + float invOutputScale = 1.0f / (*params.outputScalePtr); + quant::quant_fp8(packedAccum, params.quantOutPtr, + invOutputScale, threadOffset); + } +#if CUDA_VERSION >= 12080 + else if constexpr (QType == QuantType::kFP4) { + quant::quant_nvfp4( + packedAccum, params.quantOutPtr, params.scalingFactorOutPtr, params.outputScalePtr, token, + tokenDim, packedIdx, params.sfLayout); + } +#endif flag.waitAndUpdate( {static_cast(params.numTokens * tokenDim * WorldSize * kELT_SIZE), 0, 0, 0}); } @@ -731,6 +825,36 @@ cudaError_t oneshotAllreduceFusionDispatch(AllReduceFusionParams const& params) "Hidden Dimension %d exceeds the maximum supported hidden dimension (%d)", tokenDim, 1024 * 8 * eltsPerThread); + if (!params.rmsNormFusion && params.quantType != QuantType::kNone) { + FLASHINFER_ERROR("[MNNVL AllReduceOneShot] Quantization requires RMSNorm fusion"); + return cudaErrorInvalidValue; + } + if (params.quantType != QuantType::kNone) { + if (params.quantOut == nullptr || params.outputScale == nullptr) { + FLASHINFER_ERROR( + "[MNNVL AllReduceOneShot] quantOut and outputScale must be non-null when " + "quantization is enabled"); + return cudaErrorInvalidValue; + } + } + if (params.quantType == QuantType::kFP4) { +#if CUDA_VERSION < 12080 + FLASHINFER_ERROR("[MNNVL AllReduceOneShot] FP4 quantization requires CUDA 12.8+"); + return cudaErrorInvalidValue; +#else + if (params.scalingFactorOut == nullptr) { + FLASHINFER_ERROR( + "[MNNVL AllReduceOneShot] scalingFactorOut is required for FP4 quantization"); + return cudaErrorInvalidValue; + } + if (tokenDim % trtllm_allreduce_fusion::details::CVT_FP4_SF_VEC_SIZE != 0) { + FLASHINFER_ERROR("[MNNVL AllReduceOneShot] FP4 quantization requires tokenDim divisible by " + + std::to_string(trtllm_allreduce_fusion::details::CVT_FP4_SF_VEC_SIZE)); + return cudaErrorInvalidValue; + } +#endif + } + cudaLaunchAttribute attrs[2]; attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; attrs[0].val.programmaticStreamSerializationAllowed = params.launchWithPdl ? 1 : 0; @@ -748,14 +872,35 @@ cudaError_t oneshotAllreduceFusionDispatch(AllReduceFusionParams const& params) .numAttrs = 2, }; -#define LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, RMSNORM) \ - FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ - &config, &oneshotAllreduceFusionKernel, kernelParams)); -#define DISPATCH_ALLREDUCE_KERNEL(WORLD_SIZE) \ - if (params.rmsNormFusion) { \ - LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, true); \ - } else { \ - LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, false); \ +#define LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, RMSNORM, QTYPE) \ + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ + &config, &oneshotAllreduceFusionKernel, kernelParams)); +#define DISPATCH_ALLREDUCE_KERNEL(WORLD_SIZE) \ + if (params.rmsNormFusion) { \ + switch (params.quantType) { \ + case QuantType::kFP8: \ + LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, true, QuantType::kFP8); \ + break; \ + case QuantType::kFP4: \ + if constexpr (std::is_same_v || std::is_same_v) { \ + LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, true, QuantType::kFP4); \ + } else { \ + FLASHINFER_ERROR( \ + "[MNNVL AllReduceOneShot] FP4 quantization is only " \ + "supported for FP16/BF16"); \ + return cudaErrorInvalidValue; \ + } \ + break; \ + case QuantType::kNone: \ + LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, true, QuantType::kNone); \ + break; \ + default: \ + FLASHINFER_ERROR("[MNNVL AllReduceOneShot] Unsupported quant type " + \ + std::to_string(static_cast(params.quantType))); \ + return cudaErrorInvalidValue; \ + } \ + } else { \ + LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, false, QuantType::kNone); \ } AllReduceKernelParams kernelParams{ @@ -773,6 +918,10 @@ cudaError_t oneshotAllreduceFusionDispatch(AllReduceFusionParams const& params) .tokenDim = params.tokenDim, .epsilon = static_cast(params.epsilon), .weightBias = params.weightBias, + .outputScalePtr = params.outputScale, + .quantOutPtr = params.quantOut, + .scalingFactorOutPtr = params.scalingFactorOut, + .sfLayout = params.sfLayout, .bufferFlags = params.bufferFlags, }; @@ -800,6 +949,7 @@ cudaError_t oneshotAllreduceFusionDispatch(AllReduceFusionParams const& params) ". Supported sizes: {2, 4, 8, 16, 32, 64}"); return cudaErrorInvalidValue; } +#undef DISPATCH_ALLREDUCE_KERNEL #undef LAUNCH_ALLREDUCE_KERNEL return cudaSuccess; } @@ -910,7 +1060,8 @@ __global__ __launch_bounds__(128) void twoshotAllreduceKernel(AllReduceKernelPar } } -template +template __global__ __launch_bounds__(1024) void rmsNormLamport(AllReduceKernelParams params) { constexpr int kELTS_PER_LOAD = sizeof(PackedType) / sizeof(T); constexpr uint32_t kELT_SIZE = sizeof(T); @@ -1045,8 +1196,22 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(AllReduceKernelParams out.elements[j] = fromFloat((params.weightBias + toFloat(gamma.elements[j])) * rInput[i * kELTS_PER_LOAD + j] * rcpRms); } - reinterpret_cast(¶ms.outputPtr[token * params.tokenDim + tokenOffset])[0] = - out.packed; + if (params.outputPtr != nullptr) { + reinterpret_cast(¶ms.outputPtr[token * params.tokenDim + tokenOffset])[0] = + out.packed; + } + if constexpr (QType == QuantType::kFP8) { + float invOutputScale = 1.0f / (*params.outputScalePtr); + quant::quant_fp8(out, params.quantOutPtr, invOutputScale, + token * params.tokenDim + tokenOffset); + } +#if CUDA_VERSION >= 12080 + else if constexpr (QType == QuantType::kFP4) { + quant::quant_nvfp4( + out, params.quantOutPtr, params.scalingFactorOutPtr, params.outputScalePtr, token, + params.tokenDim, tokenOffset / kELTS_PER_LOAD, params.sfLayout); + } +#endif } } @@ -1064,6 +1229,37 @@ cudaError_t twoshotAllreduceFusionDispatch(AllReduceFusionParams const& params) int const numEltsPerThread = sizeof(float4) / sizeof(T); FLASHINFER_CHECK(tokenDim % numEltsPerThread == 0, "[MNNVL AllReduceTwoShot] token_dim must be divisible by %d", numEltsPerThread); + + if (!params.rmsNormFusion && params.quantType != QuantType::kNone) { + FLASHINFER_ERROR("[MNNVL AllReduceTwoShot] Quantization requires RMSNorm fusion"); + return cudaErrorInvalidValue; + } + if (params.quantType != QuantType::kNone) { + if (params.quantOut == nullptr || params.outputScale == nullptr) { + FLASHINFER_ERROR( + "[MNNVL AllReduceTwoShot] quantOut and outputScale must be non-null when " + "quantization is enabled"); + return cudaErrorInvalidValue; + } + } + if (params.quantType == QuantType::kFP4) { +#if CUDA_VERSION < 12080 + FLASHINFER_ERROR("[MNNVL AllReduceTwoShot] FP4 quantization requires CUDA 12.8+"); + return cudaErrorInvalidValue; +#else + if (params.scalingFactorOut == nullptr) { + FLASHINFER_ERROR( + "[MNNVL AllReduceTwoShot] scalingFactorOut is required for FP4 quantization"); + return cudaErrorInvalidValue; + } + if (tokenDim % trtllm_allreduce_fusion::details::CVT_FP4_SF_VEC_SIZE != 0) { + FLASHINFER_ERROR("[MNNVL AllReduceTwoShot] FP4 quantization requires tokenDim divisible by " + + std::to_string(trtllm_allreduce_fusion::details::CVT_FP4_SF_VEC_SIZE)); + return cudaErrorInvalidValue; + } +#endif + } + int const arNumThreads = ceil_div(tokenDim, numEltsPerThread); int const arNumBlocksPerToken = ceil_div(arNumThreads, 128); @@ -1082,6 +1278,10 @@ cudaError_t twoshotAllreduceFusionDispatch(AllReduceFusionParams const& params) .tokenDim = params.tokenDim, .epsilon = static_cast(params.epsilon), .weightBias = params.weightBias, + .outputScalePtr = params.outputScale, + .quantOutPtr = params.quantOut, + .scalingFactorOutPtr = params.scalingFactorOut, + .sfLayout = params.sfLayout, .bufferFlags = params.bufferFlags, }; @@ -1192,38 +1392,62 @@ cudaError_t twoshotAllreduceFusionDispatch(AllReduceFusionParams const& params) numTokens, rnClusterSize, rnBlockSize, rnClusterSize, rnLoadsPerThread, ceil_div(tokenDim, numEltsPerThread)); -#define LAUNCH_RMSNORM_KERNEL(USE_CGA, LOADS_PER_THREAD) \ - FLASHINFER_CUDA_CALL(cudaFuncSetAttribute(&rmsNormLamport, \ - cudaFuncAttributeMaxDynamicSharedMemorySize, \ - smemSize)); \ - FLASHINFER_CUDA_CALL( \ - cudaLaunchKernelEx(&rnConfig, &rmsNormLamport, kernelParams)); +#define LAUNCH_RMSNORM_KERNEL(USE_CGA, LOADS_PER_THREAD, QTYPE) \ + FLASHINFER_CUDA_CALL(cudaFuncSetAttribute(&rmsNormLamport, \ + cudaFuncAttributeMaxDynamicSharedMemorySize, \ + smemSize)); \ + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ + &rnConfig, &rmsNormLamport, kernelParams)); + +#define DISPATCH_QUANT(USE_CGA, LOADS_PER_THREAD) \ + switch (params.quantType) { \ + case QuantType::kFP8: \ + LAUNCH_RMSNORM_KERNEL(USE_CGA, LOADS_PER_THREAD, QuantType::kFP8); \ + break; \ + case QuantType::kFP4: \ + if constexpr (std::is_same_v || std::is_same_v) { \ + LAUNCH_RMSNORM_KERNEL(USE_CGA, LOADS_PER_THREAD, QuantType::kFP4); \ + } else { \ + FLASHINFER_ERROR( \ + "[MNNVL AllReduceTwoShotRMSNorm] FP4 quantization is only " \ + "supported for FP16/BF16"); \ + return cudaErrorInvalidValue; \ + } \ + break; \ + case QuantType::kNone: \ + LAUNCH_RMSNORM_KERNEL(USE_CGA, LOADS_PER_THREAD, QuantType::kNone); \ + break; \ + default: \ + FLASHINFER_ERROR("[MNNVL AllReduceTwoShotRMSNorm] Unsupported quant type " + \ + std::to_string(static_cast(params.quantType))); \ + return cudaErrorInvalidValue; \ + } #define DISPATCH_LOADS(USE_CGA) \ switch (rnLoadsPerThread) { \ case 1: \ - LAUNCH_RMSNORM_KERNEL(USE_CGA, 1); \ + DISPATCH_QUANT(USE_CGA, 1); \ break; \ case 2: \ - LAUNCH_RMSNORM_KERNEL(false, 2); \ + DISPATCH_QUANT(false, 2); \ break; \ case 3: \ - LAUNCH_RMSNORM_KERNEL(false, 3); \ + DISPATCH_QUANT(false, 3); \ break; \ case 4: \ - LAUNCH_RMSNORM_KERNEL(false, 4); \ + DISPATCH_QUANT(false, 4); \ break; \ case 5: \ - LAUNCH_RMSNORM_KERNEL(false, 5); \ + DISPATCH_QUANT(false, 5); \ break; \ case 6: \ - LAUNCH_RMSNORM_KERNEL(false, 6); \ + DISPATCH_QUANT(false, 6); \ break; \ case 7: \ - LAUNCH_RMSNORM_KERNEL(false, 7); \ + DISPATCH_QUANT(false, 7); \ break; \ case 8: \ - LAUNCH_RMSNORM_KERNEL(false, 8); \ + DISPATCH_QUANT(false, 8); \ break; \ default: \ FLASHINFER_ERROR("[MNNVL AllReduceTwoShotRMSNorm] Unsupported loads_per_thread " + \ @@ -1239,6 +1463,7 @@ cudaError_t twoshotAllreduceFusionDispatch(AllReduceFusionParams const& params) } #undef DISPATCH_LOADS +#undef DISPATCH_QUANT #undef LAUNCH_RMSNORM_KERNEL return cudaSuccess; } diff --git a/tests/comm/test_trtllm_mnnvl_allreduce.py b/tests/comm/test_trtllm_mnnvl_allreduce.py index e0f78775026..c6c1cbaa4f8 100644 --- a/tests/comm/test_trtllm_mnnvl_allreduce.py +++ b/tests/comm/test_trtllm_mnnvl_allreduce.py @@ -6,9 +6,11 @@ import torch import torch.distributed as dist +import flashinfer.comm as comm import flashinfer.comm.trtllm_mnnvl_ar as trtllm_mnnvl_ar from flashinfer.comm.mapping import Mapping from flashinfer.comm.mnnvl import TorchDistBackend +from flashinfer.fp4_quantization import _compute_swizzled_layout_sf_size # Use flashinfer.norm.rmsnorm as reference implementation. from flashinfer.norm import rmsnorm @@ -17,10 +19,27 @@ from tests.test_helpers.comm import ( init_torch_distributed_from_mpi, ) +from tests.test_helpers.utils_fp4 import ( + cast_from_fp4, + recover_swizzled_scales, + ref_fp4_quant, +) # Note: torch.distributed cleanup is handled by tests/comm/conftest.py +def fp8_quant(input: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + finfo = torch.finfo(torch.float8_e4m3fn) + qinput = (input.float() * scale.reciprocal()).clamp(min=finfo.min, max=finfo.max) + return qinput.to(torch.float8_e4m3fn) + + +def dequant( + input: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype +) -> torch.Tensor: + return (input.to(torch.float32) * scale).to(dtype) + + @torch.inference_mode() def row_linear_residual_norm_fusion_forward( x: torch.Tensor, @@ -496,6 +515,420 @@ def run_mnnvl_ar_full( dist.barrier() +def _prepare_quant_test_data( + seq_len: int, hidden_size: int, dtype: torch.dtype, eps: float +): + rank = dist.get_rank() + world_size = dist.get_world_size() + + if rank == 0: + x_full = torch.randn((world_size, seq_len, hidden_size), dtype=dtype) + residual = torch.randn((seq_len, hidden_size), dtype=dtype) + norm_weight = torch.randn((hidden_size,), dtype=dtype) + else: + x_full = None + residual = None + norm_weight = None + + data_list = [x_full, residual, norm_weight] + dist.broadcast_object_list(data_list, src=0) + x_full, residual, norm_weight = data_list + + x_full = x_full.cuda() + residual = residual.cuda() + norm_weight = norm_weight.cuda() + x_local = x_full[rank, :, :] + + allreduce_result = torch.sum(x_full, dim=0) + residual_out = allreduce_result + residual + norm_out = rmsnorm(residual_out, norm_weight, eps, enable_pdl=False) + return (x_local, residual, norm_weight), (norm_out, residual_out) + + +def _assert_quant_close( + quant_out: torch.Tensor, + scale_out: Optional[torch.Tensor], + ref_norm_out: torch.Tensor, + quant_type: int, + global_scale: torch.Tensor, + layout_code: int, +): + if quant_type == trtllm_mnnvl_ar.MNNVLQuantType.FP8: + ref_dequant = dequant( + fp8_quant(ref_norm_out, global_scale), global_scale, torch.float32 + ) + out_dequant = dequant(quant_out, global_scale, torch.float32) + mismatch_tol = 0.002 + else: + assert scale_out is not None + ref_fp4, ref_sf = ref_fp4_quant(ref_norm_out, global_scale, block_size=16) + if layout_code == comm.QuantizationSFLayout.SWIZZLED_128x4: + padded_rows = ((ref_norm_out.shape[0] + 127) // 128) * 128 + padded_cols = ((ref_norm_out.shape[1] // 16 + 3) // 4) * 4 + scale_out = recover_swizzled_scales( + scale_out.reshape(padded_rows, padded_cols), + ref_norm_out.shape[0], + ref_norm_out.shape[1], + 16, + ) + ref_dequant = ( + ref_fp4.to(torch.float32) + * ref_sf.repeat_interleave(16, dim=-1).to(torch.float32) + * global_scale.to(torch.float32) + ) + out_fp4 = cast_from_fp4(quant_out.view(torch.uint8)).to(ref_norm_out.device) + out_dequant = ( + out_fp4.to(torch.float32) + * scale_out.to(torch.float32).repeat_interleave(16, dim=-1) + * global_scale.to(torch.float32) + ) + mismatch_tol = 0.01 + + rtol, atol = 0.05, 0.15 + diff = torch.abs(out_dequant - ref_dequant) + max_diff = torch.maximum( + torch.abs(ref_dequant) * rtol, torch.tensor(atol, device=ref_dequant.device) + ) + mismatch_ratio = (diff > max_diff).sum().item() / out_dequant.numel() + assert mismatch_ratio <= mismatch_tol, ( + f"Mismatch ratio {mismatch_ratio:.4%} exceeds {mismatch_tol:.4%} threshold" + ) + + +# Quant fusion reuses the same MNNVL allreduce and RMSNorm paths that the +# exhaustive non-quant tests cover below. Keep this matrix pruned to the axes +# that are quant-specific or easy to regress: FP8/NVFP4, oneshot/twoshot, +# fp16/bf16, norm_out optionality, FP4 scale layout, and representative +# sequence/hidden sizes for small, mixed, wide, and large-token offsets. +MNNVL_QUANT_TEST_CASES = [ + pytest.param( + [1], + comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.ONESHOT, + torch.float16, + 2880, + [comm.QuantizationSFLayout.SWIZZLED_128x4], + id="fp8-oneshot-fp16-h2880-s1", + ), + pytest.param( + [4], + comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.TWOSHOT, + torch.bfloat16, + 5120, + [comm.QuantizationSFLayout.SWIZZLED_128x4], + id="fp8-twoshot-bf16-h5120-s4", + ), + pytest.param( + [15], + comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.TWOSHOT, + torch.float16, + 7168, + [comm.QuantizationSFLayout.SWIZZLED_128x4], + id="fp8-twoshot-fp16-h7168-s15", + ), + pytest.param( + [27, 11, 24, 256], + comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.ONESHOT, + torch.bfloat16, + 8192, + [comm.QuantizationSFLayout.SWIZZLED_128x4], + id="fp8-oneshot-bf16-h8192-mixed", + ), + pytest.param( + [127], + comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.ONESHOT, + torch.float16, + 16384, + [comm.QuantizationSFLayout.SWIZZLED_128x4], + id="fp8-oneshot-fp16-h16384-s127", + ), + pytest.param( + [998, 2048], + comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.TWOSHOT, + torch.bfloat16, + 8192, + [comm.QuantizationSFLayout.SWIZZLED_128x4], + id="fp8-twoshot-bf16-h8192-large", + ), + pytest.param( + [1], + comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.ONESHOT, + torch.float16, + 2880, + [comm.QuantizationSFLayout.LINEAR, comm.QuantizationSFLayout.SWIZZLED_128x4], + id="fp4-oneshot-fp16-h2880-s1", + ), + pytest.param( + [4], + comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.TWOSHOT, + torch.bfloat16, + 5120, + [comm.QuantizationSFLayout.LINEAR, comm.QuantizationSFLayout.SWIZZLED_128x4], + id="fp4-twoshot-bf16-h5120-s4", + ), + pytest.param( + [15], + comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.TWOSHOT, + torch.float16, + 7168, + [comm.QuantizationSFLayout.LINEAR], + id="fp4-twoshot-fp16-h7168-s15", + ), + pytest.param( + [27, 11, 24, 256], + comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.ONESHOT, + torch.bfloat16, + 8192, + [comm.QuantizationSFLayout.SWIZZLED_128x4], + id="fp4-oneshot-bf16-h8192-mixed", + ), + pytest.param( + [127], + comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.ONESHOT, + torch.float16, + 16384, + [comm.QuantizationSFLayout.LINEAR], + id="fp4-oneshot-fp16-h16384-s127", + ), + pytest.param( + [998, 2048], + comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.TWOSHOT, + torch.bfloat16, + 8192, + [comm.QuantizationSFLayout.SWIZZLED_128x4], + id="fp4-twoshot-bf16-h8192-large", + ), + pytest.param( + [4], + comm.AllReduceFusionPattern.kARResidualRMSNormOutFP8Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.ONESHOT, + torch.bfloat16, + 2880, + [comm.QuantizationSFLayout.SWIZZLED_128x4], + id="out-fp8-oneshot-bf16-h2880-s4", + ), + pytest.param( + [127], + comm.AllReduceFusionPattern.kARResidualRMSNormOutFP8Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.TWOSHOT, + torch.float16, + 7168, + [comm.QuantizationSFLayout.SWIZZLED_128x4], + id="out-fp8-twoshot-fp16-h7168-s127", + ), + pytest.param( + [4], + comm.AllReduceFusionPattern.kARResidualRMSNormOutFP4Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.TWOSHOT, + torch.bfloat16, + 2880, + [comm.QuantizationSFLayout.LINEAR], + id="out-fp4-twoshot-bf16-h2880-s4", + ), + pytest.param( + [127], + comm.AllReduceFusionPattern.kARResidualRMSNormOutFP4Quant, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.ONESHOT, + torch.float16, + 7168, + [comm.QuantizationSFLayout.SWIZZLED_128x4], + id="out-fp4-oneshot-fp16-h7168-s127", + ), +] + + +@pytest.mark.parametrize( + "seq_lens,pattern,strategy,dtype,hidden_size,layout_codes", + MNNVL_QUANT_TEST_CASES, +) +def test_mnnvl_allreduce_quant_unified( + seq_lens: list[int], + pattern: int, + strategy: trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy, + dtype: torch.dtype, + hidden_size: int, + layout_codes: list[int], +): + if torch.cuda.device_count() == 0: + pytest.skip("MNNVL quant test requires CUDA") + + init_torch_distributed_from_mpi() + rank = dist.get_rank() + world_size = dist.get_world_size() + if world_size < 2: + pytest.skip(f"This test requires at least 2 ranks, got {world_size}") + + mapping = Mapping( + world_size=world_size, + rank=rank, + gpus_per_node=torch.cuda.device_count(), + tp_size=world_size, + ) + torch.cuda.set_device(mapping.local_rank) + comm_backend = TorchDistBackend() + eps = torch.finfo(dtype).eps + + workspace = trtllm_mnnvl_ar.MNNVLAllReduceFusionWorkspace( + mapping, + max_num_tokens=max(seq_lens), + hidden_dim=hidden_size, + dtype=dtype, + comm_backend=comm_backend, + buffer_size_in_bytes=trtllm_mnnvl_ar.MNNVLAllReduceFusionWorkspace.get_required_buffer_size_bytes( + world_size, max(seq_lens), hidden_size, dtype, strategy + ), + ) + try: + is_fp4 = pattern in ( + comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant, + comm.AllReduceFusionPattern.kARResidualRMSNormOutFP4Quant, + ) + if is_fp4 and torch.cuda.get_device_capability()[0] < 10: + pytest.skip("MNNVL NVFP4 quantization requires SM100+") + + for seq_len in seq_lens: + (x_local, residual, norm_weight), (ref_norm, ref_residual) = ( + _prepare_quant_test_data(seq_len, hidden_size, dtype, eps) + ) + global_scale = torch.ones(1, dtype=torch.float32, device=x_local.device) + residual_out = torch.empty_like(x_local) + has_norm_out = pattern in ( + comm.AllReduceFusionPattern.kARResidualRMSNormOutFP8Quant, + comm.AllReduceFusionPattern.kARResidualRMSNormOutFP4Quant, + ) + norm_out = torch.empty_like(x_local) if has_norm_out else None + + for layout_code in layout_codes: + if is_fp4: + quant_type = trtllm_mnnvl_ar.MNNVLQuantType.NVFP4 + quant_out = torch.empty( + (seq_len, hidden_size // 2), + dtype=torch.uint8, + device=x_local.device, + ) + if layout_code == comm.QuantizationSFLayout.LINEAR: + scale_out = torch.empty( + (seq_len, hidden_size // 16), + dtype=torch.float8_e4m3fn, + device=x_local.device, + ) + else: + scale_out = torch.empty( + _compute_swizzled_layout_sf_size( + seq_len, hidden_size // 16 + ), + dtype=torch.float8_e4m3fn, + device=x_local.device, + ) + else: + quant_type = trtllm_mnnvl_ar.MNNVLQuantType.FP8 + quant_out = torch.empty_like(x_local, dtype=torch.float8_e4m3fn) + scale_out = None + + result = comm.allreduce_fusion( + input=x_local, + workspace=workspace, + pattern=pattern, + residual_in=residual, + residual_out=residual_out, + norm_out=norm_out, + quant_out=quant_out, + scale_out=scale_out, + rms_gamma=norm_weight, + rms_eps=eps, + scale_factor=global_scale, + layout_code=layout_code, + use_oneshot=strategy + == trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.ONESHOT, + ) + assert result is quant_out + torch.testing.assert_close( + residual_out, ref_residual, rtol=0.05, atol=0.15 + ) + if has_norm_out: + assert norm_out is not None + torch.testing.assert_close(norm_out, ref_norm, rtol=0.05, atol=0.15) + _assert_quant_close( + quant_out, + scale_out, + ref_norm, + quant_type, + global_scale, + layout_code, + ) + dist.barrier() + finally: + workspace.destroy() + dist.barrier() + + +@pytest.mark.parametrize( + "strategy", + [ + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.ONESHOT, + trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.TWOSHOT, + ], +) +def test_mnnvl_nvfp4_rejects_fp32( + strategy: trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy, +): + if torch.cuda.device_count() == 0: + pytest.skip("MNNVL quant test requires CUDA") + + init_torch_distributed_from_mpi() + rank = dist.get_rank() + world_size = dist.get_world_size() + if world_size < 2: + pytest.skip(f"This test requires at least 2 ranks, got {world_size}") + + mapping = Mapping( + world_size=world_size, + rank=rank, + gpus_per_node=torch.cuda.device_count(), + tp_size=world_size, + ) + torch.cuda.set_device(mapping.local_rank) + workspace = trtllm_mnnvl_ar.MNNVLAllReduceFusionWorkspace( + mapping, + max_num_tokens=4, + hidden_dim=2048, + dtype=torch.float32, + comm_backend=TorchDistBackend(), + buffer_size_in_bytes=trtllm_mnnvl_ar.MNNVLAllReduceFusionWorkspace.get_required_buffer_size_bytes( + world_size, 4, 2048, torch.float32, strategy + ), + ) + try: + x = torch.randn((4, 2048), dtype=torch.float32, device="cuda") + residual = torch.randn_like(x) + gamma = torch.randn((2048,), dtype=torch.float32, device="cuda") + with pytest.raises(ValueError, match="NVFP4"): + comm.allreduce_fusion( + input=x, + workspace=workspace, + pattern=comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant, + residual_in=residual, + rms_gamma=gamma, + scale_factor=torch.ones(1, dtype=torch.float32, device="cuda"), + use_oneshot=strategy + == trtllm_mnnvl_ar.MNNVLAllreduceFusionStrategy.ONESHOT, + ) + finally: + workspace.destroy() + dist.barrier() + + """Test with default workspace size""" # Multi-gpu test: mpirun -np 4 pytest tests/comm/test_trtllm_mnnvl_allreduce.py -vv -s From 4a64aaef16444d2255f2acafe66441c7aa4f9c81 Mon Sep 17 00:00:00 2001 From: Shiyu Li Date: Thu, 21 May 2026 11:34:17 -0700 Subject: [PATCH 10/15] Add fallback reduction order; Update docstrings. --- docs/api/comm.rst | 1 + flashinfer/comm/allreduce.py | 43 ++++---- flashinfer/comm/trtllm_mnnvl_ar.py | 101 +++++++++++++++--- flashinfer/jit/comm.py | 15 +++ .../comm/trtllm_mnnvl_allreduce.cuh | 52 +++++++-- 5 files changed, 170 insertions(+), 42 deletions(-) diff --git a/docs/api/comm.rst b/docs/api/comm.rst index 1ac4b1669cb..54edefad483 100644 --- a/docs/api/comm.rst +++ b/docs/api/comm.rst @@ -140,6 +140,7 @@ TensorRT-LLM MNNVL AllReduce trtllm_mnnvl_all_reduce trtllm_mnnvl_allreduce trtllm_mnnvl_fused_allreduce_add_rmsnorm + trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant trtllm_mnnvl_fused_allreduce_rmsnorm mpi_barrier diff --git a/flashinfer/comm/allreduce.py b/flashinfer/comm/allreduce.py index 2665436956f..aec144126a8 100644 --- a/flashinfer/comm/allreduce.py +++ b/flashinfer/comm/allreduce.py @@ -273,7 +273,8 @@ def _workspace_creation_heuristic( # Single-node scenarios # From benchmarking data, we can see that MNNVL is either on par (smaller problem sizes) or significantly faster than TRTLLM (larger problem sizes such as hidden_dim=8192, token_num=64 for TP=4), for single-node scenarios. - # However, trtllm has a larger support surface (more fusion patterns, more quantization support, etc.) + # TRTLLM still has the larger specialized-fusion surface, such as MoE + # patterns and packed group FP8 quantization. if "mnnvl" in suitable_backends: return ["mnnvl"] else: @@ -496,14 +497,15 @@ def allreduce_fusion( weight_bias: float = 0.0, ) -> torch.Tensor: """ - AllReduce + RMSNorm fusion operation. + AllReduce + RMSNorm fusion operation, with optional FP8/NVFP4 + quantization for supported backends. Backend is automatically determined from workspace type. If you need another backend, create the workspace for the desired backend. Supports multiple fusion patterns: - AllReduce only - AllReduce + Residual + RMSNorm - - AllReduce + Residual + RMSNorm + Quantization (FP8/FP4) + - AllReduce + Residual + RMSNorm + Quantization (FP8/NVFP4) **Note on Workspace Reusability:** You can reuse the same workspace with different (token_num, hidden_dim) combinations @@ -512,7 +514,7 @@ def allreduce_fusion( Args: input: Input tensor [token_num, hidden_dim] workspace: Workspace object (type determines backend, see create_allreduce_fusion_workspace) - pattern: Fusion pattern (AllReduceFusionPattern constant, 0-7) + pattern: Fusion pattern (AllReduceFusionPattern constant, 0-9) - kAllReduce = 0 - kARResidualRMSNorm = 1 - kARResidualRMSNormFP8Quant = 2 @@ -523,9 +525,8 @@ def allreduce_fusion( - kMoEFinalizeARResidualRMSNorm = 7 (trtllm only) - kARResidualRMSNormPerTokenGroupFP8PackedQuant = 8 (trtllm only) - kARResidualRMSNormOutPerTokenGroupFP8PackedQuant = 9 (trtllm only) - Note: MNNVL supports patterns 0-5; MoE and packed group quant - patterns remain TRTLLM-only. - Note: MOE patterns (6-7) only support trtllm backend + MNNVL supports the standard FP8/NVFP4 quant patterns (2-5). + MoE and packed group quant patterns remain TRTLLM-only. launch_with_pdl: Use Programmatic Dependent Launch trigger_completion_at_end: [trtllm only] Controls when PDL completion is signaled. True (default): signal completion after the kernel finishes (safe, no overlap). @@ -537,9 +538,10 @@ def allreduce_fusion( output: AllReduce output [token_num, hidden_dim] residual_out: Prenorm output (after residual add, before norm) [token_num, hidden_dim] norm_out: Normalized output [token_num, hidden_dim] - quant_out: Quantized output [token_num, hidden_dim] for FP8 or - [token_num, hidden_dim / 2] for NVFP4 - scale_out: NVFP4 quantization scale factors + quant_out: Quantized output. FP8 uses shape [token_num, hidden_dim] + and NVFP4 uses shape [token_num, hidden_dim / 2]. + scale_out: NVFP4 quantization scale factors. Not used by per-tensor + FP8 quantization. # ===== INPUT parameters ===== residual_in: Residual tensor to ADD [token_num, hidden_dim] @@ -548,11 +550,14 @@ def allreduce_fusion( weight_bias: Bias added to rms_gamma before scaling. 0.0 (default) -> standard RMSNorm (out = gamma * x * rsqrt(...)). 1.0 -> Gemma / Qwen3.5 RMSNorm (out = (1 + gamma) * x * rsqrt(...)). - Supported by both TRTLLM and MNNVL backends for kARResidualRMSNorm - plus all TRTLLM RMSNorm variants (quant + MoE Reduction/Finalize). + Supported by both TRTLLM and MNNVL backends for standard + RMSNorm and quant patterns (1-5), and by TRTLLM for MoE + RMSNorm variants. Ignored for kAllReduce (no normalization). - scale_factor: Input scale factor for quantization - layout_code: Scale factor layout (QuantizationSFLayout) + scale_factor: Output scale used by FP8/NVFP4 quantization. + layout_code: NVFP4 scale factor layout (QuantizationSFLayout). + MNNVL supports SWIZZLED_128x4 and LINEAR; SWIZZLED_8x4 + remains TRTLLM-only. # ===== Control parameters ===== use_oneshot: Use oneshot strategy vs twoshot @@ -574,7 +579,9 @@ def allreduce_fusion( shared_expert_output: Optional shared expert output to add [token_num, hidden_dim] Returns: - Output tensor (typically norm_out for fusion cases, output otherwise) + Output tensor for the selected pattern. Quant patterns return + quant_out, RMSNorm patterns return norm_out, and kAllReduce returns + output. Examples: >>> # Basic AllReduce + Residual + RMSNorm @@ -604,17 +611,15 @@ def allreduce_fusion( ... ) >>> # output == normed (final result) - >>> # With FP8 quantization + >>> # With FP8 quantization using either TRTLLM or MNNVL standard quant + >>> # patterns, depending on the workspace backend. >>> quant = torch.empty_like(hidden_states, dtype=torch.float8_e4m3fn) - >>> scales = torch.empty(token_num * hidden_dim // 16, dtype=torch.float16) >>> >>> output = allreduce_fusion( ... input=hidden_states, ... workspace=workspace, ... pattern=AllReduceFusionPattern.kARResidualRMSNormFP8Quant, - ... norm_out=normed, ... quant_out=quant, - ... scale_out=scales, ... residual_in=residual, ... rms_gamma=norm_weight, ... scale_factor=scale_tensor diff --git a/flashinfer/comm/trtllm_mnnvl_ar.py b/flashinfer/comm/trtllm_mnnvl_ar.py index 3109fb210f8..1d0bd6541fb 100644 --- a/flashinfer/comm/trtllm_mnnvl_ar.py +++ b/flashinfer/comm/trtllm_mnnvl_ar.py @@ -372,18 +372,47 @@ def trtllm_mnnvl_allreduce( output: Optional[torch.Tensor] = None, strategy: MNNVLAllreduceFusionStrategy = MNNVLAllreduceFusionStrategy.AUTO, ) -> torch.Tensor: - """Perform a multi-node NVLink all-reduce operation across multiple GPUs. - - This function performs an all-reduce (sum) operation using NVIDIA's multi-node NVLink (MNNVL) - technology to efficiently combine tensors across multiple GPUs and nodes. - - There are 2 variants: One-shot and Two-shot: - - One-shot: Each rank stores local shard to all other ranks. Each ranks will receive all shards at the end of the communication round and perfom local reduction. Suitable for small data size and is optimized for low latency. - - Two-shot: There will be 3 steps: - 1. Scatter each GPU's input shard to other ranks. Each rank will received all shards of a slice of tokens. - 2. Each rank perform reduction on the local tokens. - 3. Each rank broadcast the result to all ranks. - Suitable for large data size and is optimized for balancing throughput and latency. + """Perform an MNNVL all-reduce sum across tensor-parallel ranks. + + ``input`` must be a 2-D local shard with shape ``[num_tokens, hidden_dim]``. + The result has the same shape and is written to ``output`` when provided, or + to a newly allocated tensor otherwise. ``workspace`` must be an + :class:`MNNVLAllReduceFusionWorkspace` created for the same tensor-parallel + group and with enough buffer capacity for the selected strategy. + + MNNVL supports two execution strategies: + + * ``ONESHOT`` stores each rank's local shard to the peer-visible workspace + and each rank performs the reduction locally. This is the low-latency path + used for smaller payloads. + * ``TWOSHOT`` scatters token slices, reduces each slice on its destination + rank, then broadcasts the reduced result. This is the throughput-oriented + path for larger payloads. + + With ``AUTO``, FlashInfer selects the strategy from the payload size. The + selected strategy can change the floating-point accumulation order, so + bitwise results are not guaranteed to match between ``ONESHOT`` and + ``TWOSHOT`` for non-associative floating-point sums. + + This allreduce-only helper does not perform quantization. Use + :func:`trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant`, or the unified + :func:`flashinfer.comm.allreduce_fusion` API with FP8/NVFP4 + ``AllReduceFusionPattern`` values, when the post-RMSNorm quantized output is + needed. + + Reduction-order policy: + + * ``TWOSHOT`` always reduces ranks in stable rank order + ``0..world_size-1``. This keeps the result independent of the caller's + rank. + * ``ONESHOT`` uses a faster cyclic remote-rank order by default while + keeping the local rank's value in registers. This can produce + rank-dependent last-bit differences. + * Set ``FLASHINFER_MNNVL_STABLE_REDUCTION_ORDER=1`` before the MNNVL JIT + module is compiled to make ``ONESHOT`` use the same stable rank order. + Use a fresh ``FLASHINFER_WORKSPACE_BASE`` or clear the JIT cache when + changing this environment variable, because the choice is compiled into + the generated CUDA module. Args: input: Local Input Shard [num_tokens, hidden_dim] @@ -559,7 +588,44 @@ def trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant( strategy: MNNVLAllreduceFusionStrategy = MNNVLAllreduceFusionStrategy.AUTO, weight_bias: float = 0.0, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, Optional[torch.Tensor]]: - """Performs MNNVL AllReduce + Residual + RMSNorm + FP8/NVFP4 quantization.""" + """Perform MNNVL AllReduce + Residual + RMSNorm + FP8/NVFP4 quantization. + + Quantization is applied after RMSNorm. ``output`` is optional; pass it only + when the normalized non-quantized tensor is also needed. The quantized result + is always returned as ``quant_out``. + + Args: + input: Input tensor with shape ``[num_tokens, hidden_dim]``. + residual_in: Residual tensor with the same shape as ``input``. + gamma: RMSNorm gamma tensor with shape ``[hidden_dim]``. + workspace: MNNVL workspace for the tensor-parallel group. + epsilon: RMSNorm epsilon. Defaults to ``torch.finfo(input.dtype).eps``. + output: Optional normalized output tensor with shape + ``[num_tokens, hidden_dim]``. + residual_out: Optional residual output tensor with shape + ``[num_tokens, hidden_dim]``. + quant_out: Optional quantized output. For FP8, shape must be + ``[num_tokens, hidden_dim]`` and dtype ``torch.float8_e4m3fn``. For + NVFP4, shape must be ``[num_tokens, hidden_dim // 2]`` and dtype + ``torch.uint8`` or ``torch.float4_e2m1fn_x2``. + scale_out: Optional NVFP4 scale output. For ``LINEAR`` layout, shape is + ``[num_tokens, hidden_dim // 16]``. For ``SWIZZLED_128x4``, provide a + 1-D tensor large enough for the padded swizzled scale layout. FP8 + ignores this argument. + output_scale: Scalar float or float32 tensor used as the quantization + output scale. Defaults to ``1.0``. + layout_code: NVFP4 scale layout. MNNVL supports ``SWIZZLED_128x4`` and + ``LINEAR``; ``SWIZZLED_8x4`` is not supported. + quant_type: ``MNNVLQuantType.FP8`` or ``MNNVLQuantType.NVFP4``. + launch_with_pdl: Whether to launch with PDL. + strategy: MNNVL execution strategy. ``AUTO`` uses internal heuristics. + weight_bias: Bias added to gamma before scaling. ``0.0`` for standard + RMSNorm; ``1.0`` for Gemma / Qwen3.5 RMSNorm. + + Returns: + A tuple ``(quant_out, scale_out, residual_out, output)``. ``scale_out`` + is ``None`` for FP8, and ``output`` is ``None`` unless requested. + """ if epsilon is None: epsilon = torch.finfo(input.dtype).eps @@ -776,7 +842,14 @@ def trtllm_mnnvl_all_reduce( launch_with_pdl: bool, out: Optional[torch.Tensor] = None, ) -> None: - """Perform a multi-node NVLink all-reduce operation across multiple GPUs. + """Deprecated pointer-based MNNVL all-reduce API. + + Deprecated: + Use :func:`trtllm_mnnvl_allreduce` with + :class:`MNNVLAllReduceFusionWorkspace` instead. This legacy function is + kept for compatibility and may be removed in a future release. + + Perform a multi-node NVLink all-reduce operation across multiple GPUs. This function performs an all-reduce (sum) operation using NVIDIA's multi-node NVLink (MNNVL) technology to efficiently combine tensors across multiple GPUs and nodes. diff --git a/flashinfer/jit/comm.py b/flashinfer/jit/comm.py index 52250f5454d..e9fd81fac26 100644 --- a/flashinfer/jit/comm.py +++ b/flashinfer/jit/comm.py @@ -40,6 +40,21 @@ def gen_trtllm_mnnvl_comm_module() -> JitSpec: nvcc_flags = current_compilation_context.get_nvcc_flags_list( supported_major_versions=[9, 10] ) + # Reduction-order policy for the JIT-compiled MNNVL allreduce kernels: + # - twoshot always uses the original rank-stable reduction order + # (0..world_size-1). Its measured cost is noise-level, and keeping this + # path deterministic avoids rank-dependent floating-point roundoff. + # - oneshot defaults to latcomm's faster cyclic remote-rank order while + # keeping the local value in registers. + # - set FLASHINFER_MNNVL_STABLE_REDUCTION_ORDER to a nonzero value before + # JIT to make oneshot use the original rank-stable order too. Use a fresh + # FLASHINFER_WORKSPACE_BASE when switching this knob so the cached module + # is rebuilt with the desired macro. + if os.environ.get("FLASHINFER_MNNVL_STABLE_REDUCTION_ORDER", "0") not in ( + "0", + "", + ): + nvcc_flags += ["-DFLASHINFER_MNNVL_ONESHOT_STABLE_REDUCTION_ORDER"] return gen_jit_spec( "trtllm_mnnvl_comm", [ diff --git a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh index 40d7f7baec0..5d773fb97f6 100644 --- a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh +++ b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh @@ -705,6 +705,28 @@ __global__ void __launch_bounds__(1024) // ============================= flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], -1); + // Default path keeps the local value in registers and polls remote ranks in + // latcomm's cyclic order. The fallback preserves the original rank-stable + // reduction order by doing one extra local-rank Lamport load; otherwise we + // would need a dynamic `values[params.rank]` subscript to insert the register + // value into the 0..WorldSize-1 reduction sequence. +#ifdef FLASHINFER_MNNVL_ONESHOT_STABLE_REDUCTION_ORDER + PackedVec values[WorldSize]; + while (1) { + bool valid = true; +#pragma unroll + for (int r = 0; r < WorldSize; r++) { + auto loaded = loadPackedVolatile( + &stagePtrLocal[token * tokenDim * WorldSize + r * tokenDim + + packedIdx * kELTS_PER_THREAD]); + values[r].packed = loaded.packed; + valid &= !isLamportDirty(loaded); + } + if (valid) { + break; + } + } +#else PackedVec remoteValues[WorldSize - 1]; while (1) { bool valid = true; @@ -722,6 +744,7 @@ __global__ void __launch_bounds__(1024) break; } } +#endif // ======================= Reduction ============================= float accum[kELTS_PER_THREAD]; @@ -729,11 +752,19 @@ __global__ void __launch_bounds__(1024) #pragma unroll for (int i = 0; i < kELTS_PER_THREAD; i++) { +#ifdef FLASHINFER_MNNVL_ONESHOT_STABLE_REDUCTION_ORDER + accum[i] = 0.f; +#pragma unroll + for (int r = 0; r < WorldSize; r++) { + accum[i] += toFloat(values[r].elements[i]); + } +#else accum[i] = toFloat(val.elements[i]); #pragma unroll for (int r = 0; r < WorldSize - 1; r++) { accum[i] += toFloat(remoteValues[r].elements[i]); } +#endif } #pragma unroll @@ -1007,16 +1038,19 @@ __global__ __launch_bounds__(128) void twoshotAllreduceKernel(AllReduceKernelPar if (inBounds && destRank == params.rank) { int const localToken = token / WorldSize; - PackedVec remoteValues[WorldSize - 1]; + // Twoshot intentionally always uses the original rank-stable order. The + // sweep showed no meaningful speedup from the cyclic remote-only path here, + // so keep deterministic 0..WorldSize-1 accumulation and pay the one extra + // local-rank Lamport load instead of dynamically placing `val`. + PackedVec values[WorldSize]; while (1) { bool valid = true; #pragma unroll - for (int r = 0; r < WorldSize - 1; r++) { - int const rankToLoad = (r + params.rank + 1) % WorldSize; + for (int r = 0; r < WorldSize; r++) { auto loaded = loadPackedVolatile( - &scatterBufLocal[localToken * params.tokenDim * WorldSize + - rankToLoad * params.tokenDim + tokenOffset]); - remoteValues[r].packed = loaded.packed; + &scatterBufLocal[localToken * params.tokenDim * WorldSize + r * params.tokenDim + + tokenOffset]); + values[r].packed = loaded.packed; valid &= !isLamportDirty(loaded); } if (valid) { @@ -1027,10 +1061,10 @@ __global__ __launch_bounds__(128) void twoshotAllreduceKernel(AllReduceKernelPar PackedVec packedAccum; #pragma unroll for (int i = 0; i < kELTS_PER_THREAD; i++) { - float accum = toFloat(val.elements[i]); + float accum = 0.f; #pragma unroll - for (int r = 0; r < WorldSize - 1; r++) { - accum += toFloat(remoteValues[r].elements[i]); + for (int r = 0; r < WorldSize; r++) { + accum += toFloat(values[r].elements[i]); } packedAccum.elements[i] = fromFloat(accum); } From 7b4402e06ef89fb3f44d2e97cfff66bee151fb3a Mon Sep 17 00:00:00 2001 From: Shiyu Li Date: Thu, 21 May 2026 11:59:47 -0700 Subject: [PATCH 11/15] Fix scale factor padding check. --- flashinfer/comm/trtllm_mnnvl_ar.py | 9 +++++-- tests/comm/test_trtllm_mnnvl_allreduce.py | 29 +++++++++++++++++++++++ 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/flashinfer/comm/trtllm_mnnvl_ar.py b/flashinfer/comm/trtllm_mnnvl_ar.py index 1d0bd6541fb..832064d344f 100644 --- a/flashinfer/comm/trtllm_mnnvl_ar.py +++ b/flashinfer/comm/trtllm_mnnvl_ar.py @@ -712,6 +712,11 @@ def trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant( f"quant_out dtype for NVFP4 must be uint8 or float4_e2m1fn_x2, got {quant_out.dtype}." ) + expected_scale_out_numel = ( + token_num * hidden_dim // 16 + if layout_code == QuantizationSFLayout.LINEAR + else _compute_swizzled_layout_sf_size(token_num, hidden_dim // 16) + ) if scale_out is None: if layout_code == QuantizationSFLayout.LINEAR: scale_out = torch.empty( @@ -725,9 +730,9 @@ def trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant( dtype=torch.float8_e4m3fn, device=input.device, ) - elif scale_out.numel() < token_num * hidden_dim // 16: + elif scale_out.numel() < expected_scale_out_numel: raise ValueError( - f"scale_out is too small for NVFP4: got {scale_out.numel()} elements, need at least {token_num * hidden_dim // 16}." + f"scale_out is too small for NVFP4: got {scale_out.numel()} elements, need at least {expected_scale_out_numel}." ) if scale_out.dtype != torch.float8_e4m3fn: raise ValueError( diff --git a/tests/comm/test_trtllm_mnnvl_allreduce.py b/tests/comm/test_trtllm_mnnvl_allreduce.py index c6c1cbaa4f8..59b664447dc 100644 --- a/tests/comm/test_trtllm_mnnvl_allreduce.py +++ b/tests/comm/test_trtllm_mnnvl_allreduce.py @@ -40,6 +40,35 @@ def dequant( return (input.to(torch.float32) * scale).to(dtype) +def test_mnnvl_nvfp4_default_swizzled_scale_out_requires_padded_size(): + token_num = 1 + hidden_dim = 16 + linear_scale_size = token_num * hidden_dim // 16 + swizzled_scale_size = _compute_swizzled_layout_sf_size(token_num, hidden_dim // 16) + assert linear_scale_size < swizzled_scale_size + + input_tensor = torch.randn((token_num, hidden_dim), dtype=torch.float16) + residual = torch.randn_like(input_tensor) + gamma = torch.randn((hidden_dim,), dtype=torch.float16) + scale_out = torch.empty(linear_scale_size, dtype=torch.float8_e4m3fn) + + with pytest.raises( + ValueError, + match=( + f"scale_out is too small for NVFP4: got {linear_scale_size} " + f"elements, need at least {swizzled_scale_size}" + ), + ): + trtllm_mnnvl_ar.trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant( + input=input_tensor, + residual_in=residual, + gamma=gamma, + workspace=object(), # type: ignore[arg-type] + scale_out=scale_out, + quant_type=trtllm_mnnvl_ar.MNNVLQuantType.NVFP4, + ) + + @torch.inference_mode() def row_linear_residual_norm_fusion_forward( x: torch.Tensor, From c9eed294b360473ee64dfd5e277a5d877cae59ab Mon Sep 17 00:00:00 2001 From: Shiyu Li Date: Thu, 21 May 2026 12:14:09 -0700 Subject: [PATCH 12/15] Revert lamport polling behavior --- include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh index 5d773fb97f6..a2b0273a88c 100644 --- a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh +++ b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh @@ -433,7 +433,14 @@ inline __device__ VolatilePackedLoad loadPackedVolatile(void con template inline __device__ bool isLamportDirty(VolatilePackedLoad const& value) { // The dirty sentinel is a raw word; typed fp compares can flush nearby bit patterns. - return value.words[0] == kNEGZERO_FP32; + // ptx memeory model only gaurantee atomicity for 64bits granularity so we check every element + // here to make sure the packed payload is consumed. + bool dirty = false; +#pragma unroll + for (int i = 0; i < sizeof(PackedType) / sizeof(uint32_t); i++) { + dirty |= value.words[i] == kNEGZERO_FP32; + } + return dirty; } template From 1246cc5ce803d11b4e454d9cedf0f9386060c11f Mon Sep 17 00:00:00 2001 From: Shiyu Li Date: Thu, 21 May 2026 14:28:07 -0700 Subject: [PATCH 13/15] Address review comment --- csrc/trtllm_mnnvl_allreduce.cu | 14 ++++--- flashinfer/comm/trtllm_mnnvl_ar.py | 14 ++++++- .../comm/trtllm_mnnvl_allreduce.cuh | 22 +++++------ tests/comm/test_trtllm_mnnvl_allreduce.py | 37 ++++++++++++------- 4 files changed, 53 insertions(+), 34 deletions(-) diff --git a/csrc/trtllm_mnnvl_allreduce.cu b/csrc/trtllm_mnnvl_allreduce.cu index de8cbc963bc..0bdecba8883 100644 --- a/csrc/trtllm_mnnvl_allreduce.cu +++ b/csrc/trtllm_mnnvl_allreduce.cu @@ -46,9 +46,14 @@ void trtllm_mnnvl_allreduce_fusion(TensorView input, int64_t multicast_buffer_pt int64_t token_dim = input.size(1); auto quant_type_enum = quant_type.has_value() ? static_cast(quant_type.value()) : QuantType::kNone; - auto sf_layout = layout_code.has_value() - ? static_cast(layout_code.value()) - : QuantizationSFLayout::SWIZZLED_128x4; + auto sf_layout = QuantizationSFLayout::SWIZZLED_128x4; + if (layout_code.has_value()) { + auto const sf_layout_code = layout_code.value(); + TVM_FFI_ICHECK(sf_layout_code == static_cast(QuantizationSFLayout::LINEAR) || + sf_layout_code == static_cast(QuantizationSFLayout::SWIZZLED_128x4)) + << "MNNVL quantization fusion supports SWIZZLED_128x4 or LINEAR scale layouts"; + sf_layout = static_cast(sf_layout_code); + } // Validate input parameters TVM_FFI_ICHECK_EQ(token_dim % (sizeof(float4) / sizeof(c_type)), 0) @@ -76,9 +81,6 @@ void trtllm_mnnvl_allreduce_fusion(TensorView input, int64_t multicast_buffer_pt (output_scale.has_value() && encode_dlpack_dtype(output_scale.value().dtype()) == float32_code)) << "output_scale must be provided for MNNVL quantization fusion and must be float32"; - TVM_FFI_ICHECK(quant_type_enum == QuantType::kNone || - sf_layout != QuantizationSFLayout::SWIZZLED_8x4) - << "MNNVL quantization fusion supports SWIZZLED_128x4 or LINEAR scale layouts"; TVM_FFI_ICHECK(quant_type_enum == QuantType::kNone || quant_out.has_value()) << "quant_out must be provided when quantization fusion is enabled"; diff --git a/flashinfer/comm/trtllm_mnnvl_ar.py b/flashinfer/comm/trtllm_mnnvl_ar.py index 832064d344f..2c808e0b5df 100644 --- a/flashinfer/comm/trtllm_mnnvl_ar.py +++ b/flashinfer/comm/trtllm_mnnvl_ar.py @@ -741,8 +741,6 @@ def trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant( else: raise ValueError(f"Unsupported MNNVL quant_type: {quant_type}") - module = get_trtllm_mnnvl_comm_module() - if strategy == MNNVLAllreduceFusionStrategy.AUTO: strategy = MNNVLAllreduceFusionStrategy.select_strategy( workspace.tp_size, token_num, hidden_dim, input.dtype @@ -754,6 +752,18 @@ def trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant( f"The buffer size in the given workspace is insufficient for the given problem size. Buffer: {workspace.buffer_size_bytes} bytes, Required: {workspace.get_required_buffer_size_bytes(workspace.tp_size, token_num, hidden_dim, input.dtype, strategy)} bytes." ) + for tensor, name in ( + (residual_in, "residual_in"), + (gamma, "gamma"), + (output, "output"), + (residual_out, "residual_out"), + (quant_out, "quant_out"), + (scale_out, "scale_out"), + ): + if tensor is not None and tensor.device != input.device: + raise ValueError(f"{name} must be on {input.device}, got {tensor.device}.") + + module = get_trtllm_mnnvl_comm_module() module.trtllm_mnnvl_allreduce_fusion( input, workspace.mc_ptr, diff --git a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh index a2b0273a88c..3cd65caf3cc 100644 --- a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh +++ b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh @@ -698,10 +698,7 @@ __global__ void __launch_bounds__(1024) // ==================== Broadcast tokens to each rank ============================= PackedVec val; val.packed = loadPacked(¶ms.shardPtr[threadOffset]); -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) { - if (isNegZero(val.elements[i])) val.elements[i] = fromFloat(0.f); - } + sanitizeLamportPayload(val); reinterpret_cast( &stagePtrMcast[token * tokenDim * WorldSize + params.rank * tokenDim])[packedIdx] = @@ -784,7 +781,9 @@ __global__ void __launch_bounds__(1024) PackedVec residualIn; residualIn.packed = *reinterpret_cast(¶ms.residualInPtr[threadOffset]); packedAccum += residualIn; - *reinterpret_cast(¶ms.prenormedPtr[threadOffset]) = packedAccum.packed; + if (params.prenormedPtr != nullptr) { + *reinterpret_cast(¶ms.prenormedPtr[threadOffset]) = packedAccum.packed; + } // =============================== Rmsnorm ================================ PackedVec gamma; gamma.packed = @@ -1028,12 +1027,7 @@ __global__ __launch_bounds__(128) void twoshotAllreduceKernel(AllReduceKernelPar PackedVec val; if (inBounds) { val.packed = loadPacked(¶ms.shardPtr[threadOffset]); -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) { - if (isNegZero(val.elements[i])) { - val.elements[i] = fromFloat(0.F); - } - } + sanitizeLamportPayload(val); reinterpret_cast(&scatterBufDest[destTokenOffset * params.tokenDim * WorldSize + params.rank * params.tokenDim])[packedIdx] = @@ -1190,8 +1184,10 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(AllReduceKernelParams PackedVec inp{.packed = loadPacked(&smemInput[chunkOffset])}; PackedVec res{.packed = loadPacked(&smemResidual[chunkOffset])}; PackedVec inpPlusRes = inp + res; - reinterpret_cast( - ¶ms.prenormedPtr[token * params.tokenDim + tokenOffset])[0] = inpPlusRes.packed; + if (params.prenormedPtr != nullptr) { + reinterpret_cast( + ¶ms.prenormedPtr[token * params.tokenDim + tokenOffset])[0] = inpPlusRes.packed; + } #pragma unroll for (int j = 0; j < kELTS_PER_LOAD; j++) { diff --git a/tests/comm/test_trtllm_mnnvl_allreduce.py b/tests/comm/test_trtllm_mnnvl_allreduce.py index 59b664447dc..1e10e36df5f 100644 --- a/tests/comm/test_trtllm_mnnvl_allreduce.py +++ b/tests/comm/test_trtllm_mnnvl_allreduce.py @@ -11,6 +11,7 @@ from flashinfer.comm.mapping import Mapping from flashinfer.comm.mnnvl import TorchDistBackend from flashinfer.fp4_quantization import _compute_swizzled_layout_sf_size +from flashinfer.utils import get_compute_capability # Use flashinfer.norm.rmsnorm as reference implementation. from flashinfer.norm import rmsnorm @@ -28,16 +29,24 @@ # Note: torch.distributed cleanup is handled by tests/comm/conftest.py -def fp8_quant(input: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: +def fp8_quant(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: finfo = torch.finfo(torch.float8_e4m3fn) - qinput = (input.float() * scale.reciprocal()).clamp(min=finfo.min, max=finfo.max) + qinput = (x.float() * scale.reciprocal()).clamp(min=finfo.min, max=finfo.max) return qinput.to(torch.float8_e4m3fn) -def dequant( - input: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype -) -> torch.Tensor: - return (input.to(torch.float32) * scale).to(dtype) +def dequant(x: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + return (x.to(torch.float32) * scale).to(dtype) + + +def _mnnvl_nvfp4_supported_on_all_ranks(device: torch.device) -> bool: + local_supported = get_compute_capability(device)[0] >= 10 + if not dist.is_initialized(): + return local_supported + + support_flags = [False] * dist.get_world_size() + dist.all_gather_object(support_flags, local_supported) + return all(support_flags) def test_mnnvl_nvfp4_default_swizzled_scale_out_requires_padded_size(): @@ -805,6 +814,15 @@ def test_mnnvl_allreduce_quant_unified( tp_size=world_size, ) torch.cuda.set_device(mapping.local_rank) + is_fp4 = pattern in ( + comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant, + comm.AllReduceFusionPattern.kARResidualRMSNormOutFP4Quant, + ) + if is_fp4 and not _mnnvl_nvfp4_supported_on_all_ranks( + torch.device("cuda", torch.cuda.current_device()) + ): + pytest.skip("MNNVL NVFP4 quantization requires SM100+ on all ranks") + comm_backend = TorchDistBackend() eps = torch.finfo(dtype).eps @@ -819,13 +837,6 @@ def test_mnnvl_allreduce_quant_unified( ), ) try: - is_fp4 = pattern in ( - comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant, - comm.AllReduceFusionPattern.kARResidualRMSNormOutFP4Quant, - ) - if is_fp4 and torch.cuda.get_device_capability()[0] < 10: - pytest.skip("MNNVL NVFP4 quantization requires SM100+") - for seq_len in seq_lens: (x_local, residual, norm_weight), (ref_norm, ref_residual) = ( _prepare_quant_test_data(seq_len, hidden_size, dtype, eps) From a30023f25ed2245d8b0473318cfc62484866d2b1 Mon Sep 17 00:00:00 2001 From: Shiyu Li Date: Fri, 22 May 2026 15:30:46 -0700 Subject: [PATCH 14/15] Use fast path + determinisitic reduction order. --- flashinfer/comm/trtllm_mnnvl_ar.py | 25 +-- flashinfer/jit/comm.py | 20 +- .../comm/trtllm_mnnvl_allreduce.cuh | 211 +++++++++++++----- 3 files changed, 167 insertions(+), 89 deletions(-) diff --git a/flashinfer/comm/trtllm_mnnvl_ar.py b/flashinfer/comm/trtllm_mnnvl_ar.py index 2c808e0b5df..5b2d2783517 100644 --- a/flashinfer/comm/trtllm_mnnvl_ar.py +++ b/flashinfer/comm/trtllm_mnnvl_ar.py @@ -389,10 +389,8 @@ def trtllm_mnnvl_allreduce( rank, then broadcasts the reduced result. This is the throughput-oriented path for larger payloads. - With ``AUTO``, FlashInfer selects the strategy from the payload size. The - selected strategy can change the floating-point accumulation order, so - bitwise results are not guaranteed to match between ``ONESHOT`` and - ``TWOSHOT`` for non-associative floating-point sums. + With ``AUTO``, FlashInfer selects the strategy from the payload size. Both + ``ONESHOT`` and ``TWOSHOT`` are fully deterministic across ranks. This allreduce-only helper does not perform quantization. Use :func:`trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant`, or the unified @@ -400,19 +398,16 @@ def trtllm_mnnvl_allreduce( ``AllReduceFusionPattern`` values, when the post-RMSNorm quantized output is needed. - Reduction-order policy: + Determinism: - * ``TWOSHOT`` always reduces ranks in stable rank order - ``0..world_size-1``. This keeps the result independent of the caller's + * ``ONESHOT`` and ``TWOSHOT`` use the exact same reduction order on each rank. - * ``ONESHOT`` uses a faster cyclic remote-rank order by default while - keeping the local rank's value in registers. This can produce - rank-dependent last-bit differences. - * Set ``FLASHINFER_MNNVL_STABLE_REDUCTION_ORDER=1`` before the MNNVL JIT - module is compiled to make ``ONESHOT`` use the same stable rank order. - Use a fresh ``FLASHINFER_WORKSPACE_BASE`` or clear the JIT cache when - changing this environment variable, because the choice is compiled into - the generated CUDA module. + * ``ONESHOT`` keeps the local rank's value in registers; only remote ranks + are volatile-loaded from the Lamport buffer. + * ``ONESHOT`` uses a rank-specialized fast path for ``world_size <= 8``. + Larger world sizes use a compact deterministic fallback because the + runtime benefit is thin and specializing every rank significantly + increases JIT compile time. Args: input: Local Input Shard [num_tokens, hidden_dim] diff --git a/flashinfer/jit/comm.py b/flashinfer/jit/comm.py index e9fd81fac26..e43ec428e65 100644 --- a/flashinfer/jit/comm.py +++ b/flashinfer/jit/comm.py @@ -40,21 +40,11 @@ def gen_trtllm_mnnvl_comm_module() -> JitSpec: nvcc_flags = current_compilation_context.get_nvcc_flags_list( supported_major_versions=[9, 10] ) - # Reduction-order policy for the JIT-compiled MNNVL allreduce kernels: - # - twoshot always uses the original rank-stable reduction order - # (0..world_size-1). Its measured cost is noise-level, and keeping this - # path deterministic avoids rank-dependent floating-point roundoff. - # - oneshot defaults to latcomm's faster cyclic remote-rank order while - # keeping the local value in registers. - # - set FLASHINFER_MNNVL_STABLE_REDUCTION_ORDER to a nonzero value before - # JIT to make oneshot use the original rank-stable order too. Use a fresh - # FLASHINFER_WORKSPACE_BASE when switching this knob so the cached module - # is rebuilt with the desired macro. - if os.environ.get("FLASHINFER_MNNVL_STABLE_REDUCTION_ORDER", "0") not in ( - "0", - "", - ): - nvcc_flags += ["-DFLASHINFER_MNNVL_ONESHOT_STABLE_REDUCTION_ORDER"] + # MNNVL allreduce is fully deterministic across ranks. Oneshot specializes + # the common TP<=8 cases to keep the local value in registers and + # volatile-load only peers from the Lamport buffer. Larger world sizes use a + # compact deterministic fallback: the runtime benefit is thin there, while + # rank-specializing every case significantly increases JIT compile time. return gen_jit_spec( "trtllm_mnnvl_comm", [ diff --git a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh index 3cd65caf3cc..e3e45e642d1 100644 --- a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh +++ b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh @@ -27,6 +27,7 @@ #include #include #include +#include #include "../exception.h" #include "../fp4_layout.cuh" @@ -443,6 +444,102 @@ inline __device__ bool isLamportDirty(VolatilePackedLoad const& valu return dirty; } +template +inline __device__ bool pollOneshotRemoteRank(PackedVec* remoteValues, + T* stagePtrLocal, int token, int tokenDim, + int packedIdx) { + if constexpr (Rank == LocalRank) { + return true; + } else { + auto loaded = loadPackedVolatile( + &stagePtrLocal[token * tokenDim * WorldSize + Rank * tokenDim + + packedIdx * kELTS_PER_THREAD]); + remoteValues[Rank].packed = loaded.packed; + return !isLamportDirty(loaded); + } +} + +template +inline __device__ bool pollOneshotRemoteRanks(PackedVec* remoteValues, + T* stagePtrLocal, int token, int tokenDim, + int packedIdx, std::integer_sequence) { + bool valid = true; + ((valid &= pollOneshotRemoteRank( + remoteValues, stagePtrLocal, token, tokenDim, packedIdx)), + ...); + return valid; +} + +template +inline __device__ void accumulatePacked(float (&accum)[kELTS_PER_THREAD], + PackedVec const& value) { +#pragma unroll + for (int i = 0; i < kELTS_PER_THREAD; i++) { + accum[i] += toFloat(value.elements[i]); + } +} + +template +inline __device__ void accumulateOneshotRank(float (&accum)[kELTS_PER_THREAD], + PackedVec const* remoteValues, + PackedVec const& localValue) { + if constexpr (Rank == LocalRank) { + accumulatePacked(accum, localValue); + } else { + accumulatePacked(accum, remoteValues[Rank]); + } +} + +template +inline __device__ void accumulateOneshotRanks(float (&accum)[kELTS_PER_THREAD], + PackedVec const* remoteValues, + PackedVec const& localValue, + std::integer_sequence) { + (accumulateOneshotRank(accum, remoteValues, + localValue), + ...); +} + +template +inline __device__ void waitOneshotRemoteRanks(PackedVec* remoteValues, + T* stagePtrLocal, int token, int tokenDim, + int packedIdx) { + static_assert(LocalRank < WorldSize); + while (1) { + bool const valid = + pollOneshotRemoteRanks( + remoteValues, stagePtrLocal, token, tokenDim, packedIdx, + std::make_integer_sequence{}); + if (valid) { + break; + } + } +} + +template +inline __device__ PackedVec reduceOneshotDeterministic( + PackedVec const* remoteValues, PackedVec const& localValue) { + static_assert(LocalRank < WorldSize); + float accum[kELTS_PER_THREAD]; +#pragma unroll + for (int i = 0; i < kELTS_PER_THREAD; i++) { + accum[i] = 0.f; + } + accumulateOneshotRanks( + accum, remoteValues, localValue, std::make_integer_sequence{}); + + PackedVec packedAccum; +#pragma unroll + for (int i = 0; i < kELTS_PER_THREAD; i++) { + packedAccum.elements[i] = fromFloat(accum[i]); + } + return packedAccum; +} + template inline __device__ void copyF4(T_IN* dst, T_IN const* src) { float4* dst4 = reinterpret_cast(dst); @@ -709,72 +806,71 @@ __global__ void __launch_bounds__(1024) // ============================= flag.clearDirtyLamportBuf(params.inputPtrs[params.rank], -1); - // Default path keeps the local value in registers and polls remote ranks in - // latcomm's cyclic order. The fallback preserves the original rank-stable - // reduction order by doing one extra local-rank Lamport load; otherwise we - // would need a dynamic `values[params.rank]` subscript to insert the register - // value into the 0..WorldSize-1 reduction sequence. -#ifdef FLASHINFER_MNNVL_ONESHOT_STABLE_REDUCTION_ORDER - PackedVec values[WorldSize]; - while (1) { - bool valid = true; -#pragma unroll - for (int r = 0; r < WorldSize; r++) { - auto loaded = loadPackedVolatile( - &stagePtrLocal[token * tokenDim * WorldSize + r * tokenDim + - packedIdx * kELTS_PER_THREAD]); - values[r].packed = loaded.packed; - valid &= !isLamportDirty(loaded); - } - if (valid) { - break; + // ======================= Reduction ============================= + // Fully deterministic: every rank uses the exact same reduction order. + // For WorldSize <= 8, specialize the local slot so the fast path reuses `val` + // from registers without a dynamic `remoteValues[params.rank]` store. Larger + // world sizes use the compact fallback because the benefit is thin but + // specializing every rank significantly increases JIT compile time. + PackedVec packedAccum; + if constexpr (WorldSize <= 8) { + packedAccum = val; +#define RUN_ONESHOT_LOCAL_RANK(LOCAL_RANK) \ + case LOCAL_RANK: \ + if constexpr (WorldSize > LOCAL_RANK) { \ + PackedVec remoteValues[WorldSize]; \ + utils::waitOneshotRemoteRanks( \ + remoteValues, stagePtrLocal, token, tokenDim, packedIdx); \ + packedAccum = utils::reduceOneshotDeterministic(remoteValues, val); \ + } \ + break + + switch (params.rank) { + RUN_ONESHOT_LOCAL_RANK(0); + RUN_ONESHOT_LOCAL_RANK(1); + RUN_ONESHOT_LOCAL_RANK(2); + RUN_ONESHOT_LOCAL_RANK(3); + RUN_ONESHOT_LOCAL_RANK(4); + RUN_ONESHOT_LOCAL_RANK(5); + RUN_ONESHOT_LOCAL_RANK(6); + RUN_ONESHOT_LOCAL_RANK(7); } - } -#else - PackedVec remoteValues[WorldSize - 1]; - while (1) { - bool valid = true; +#undef RUN_ONESHOT_LOCAL_RANK + } else { + // Compact deterministic fallback for larger instantiations. + PackedVec values[WorldSize]; + while (1) { + bool valid = true; #pragma unroll - for (int r = 0; r < WorldSize - 1; r++) { - int const rankToLoad = (r + params.rank + 1) % WorldSize; - auto loaded = loadPackedVolatile( - &stagePtrLocal[token * tokenDim * WorldSize + rankToLoad * tokenDim + - packedIdx * kELTS_PER_THREAD]); - remoteValues[r].packed = loaded.packed; - // Keep the local value in registers and poll remote ranks in latcomm's cyclic order. - valid &= !isLamportDirty(loaded); - } - if (valid) { - break; + for (int r = 0; r < WorldSize; r++) { + auto loaded = loadPackedVolatile( + &stagePtrLocal[token * tokenDim * WorldSize + r * tokenDim + + packedIdx * kELTS_PER_THREAD]); + values[r].packed = loaded.packed; + valid &= !isLamportDirty(loaded); + } + if (valid) { + break; + } } - } -#endif - - // ======================= Reduction ============================= - float accum[kELTS_PER_THREAD]; - PackedVec packedAccum; + float accum[kELTS_PER_THREAD]; #pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) { -#ifdef FLASHINFER_MNNVL_ONESHOT_STABLE_REDUCTION_ORDER - accum[i] = 0.f; + for (int i = 0; i < kELTS_PER_THREAD; i++) { + accum[i] = 0.f; #pragma unroll - for (int r = 0; r < WorldSize; r++) { - accum[i] += toFloat(values[r].elements[i]); + for (int r = 0; r < WorldSize; r++) { + accum[i] += toFloat(values[r].elements[i]); + } } -#else - accum[i] = toFloat(val.elements[i]); + #pragma unroll - for (int r = 0; r < WorldSize - 1; r++) { - accum[i] += toFloat(remoteValues[r].elements[i]); + for (int i = 0; i < kELTS_PER_THREAD; i++) { + packedAccum.elements[i] = fromFloat(accum[i]); } -#endif } -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) { - packedAccum.elements[i] = fromFloat(accum[i]); - } cudaTriggerProgrammaticLaunchCompletion(); if constexpr (RMSNormFusion) { // =============================== Residual =============================== @@ -1039,10 +1135,7 @@ __global__ __launch_bounds__(128) void twoshotAllreduceKernel(AllReduceKernelPar if (inBounds && destRank == params.rank) { int const localToken = token / WorldSize; - // Twoshot intentionally always uses the original rank-stable order. The - // sweep showed no meaningful speedup from the cyclic remote-only path here, - // so keep deterministic 0..WorldSize-1 accumulation and pay the one extra - // local-rank Lamport load instead of dynamically placing `val`. + // Fully deterministic: every rank uses the exact same reduction order. PackedVec values[WorldSize]; while (1) { bool valid = true; From 6a384bfbb0cca2eb41fe0b22dae70f975e37201a Mon Sep 17 00:00:00 2001 From: Shiyu Li Date: Sat, 23 May 2026 01:39:52 -0700 Subject: [PATCH 15/15] Chunked reduction to avoid register spilling. --- .../comm/trtllm_mnnvl_allreduce.cuh | 118 +++++++++--------- 1 file changed, 60 insertions(+), 58 deletions(-) diff --git a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh index e3e45e642d1..429555a05bc 100644 --- a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh +++ b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh @@ -540,6 +540,57 @@ inline __device__ PackedVec reduceOneshotDeterministic( return packedAccum; } +template +inline __device__ PackedVec reduceLamportRanksChunked(T* buffer, int token, + int tokenDim, + int tokenOffset) { + static_assert(kRankChunk > 0, "kRankChunk must be positive"); + static_assert(WorldSize % kRankChunk == 0, "WorldSize must be divisible by kRankChunk"); + // Chunk ranks for large world sizes to avoid register spills from keeping every rank payload + // live at once. + float accum[kELTS_PER_THREAD]; +#pragma unroll + for (int i = 0; i < kELTS_PER_THREAD; i++) { + accum[i] = 0.f; + } + +#pragma unroll 1 + for (int rankBase = 0; rankBase < WorldSize; rankBase += kRankChunk) { + float chunkAccum[kELTS_PER_THREAD]; + while (1) { + bool valid = true; +#pragma unroll + for (int i = 0; i < kELTS_PER_THREAD; i++) { + chunkAccum[i] = 0.f; + } +#pragma unroll + for (int rr = 0; rr < kRankChunk; rr++) { + int const r = rankBase + rr; + auto loaded = loadPackedVolatile( + &buffer[token * tokenDim * WorldSize + r * tokenDim + tokenOffset]); + PackedVec value; + value.packed = loaded.packed; + valid &= !isLamportDirty(loaded); + accumulatePacked(chunkAccum, value); + } + if (valid) { + break; + } + } +#pragma unroll + for (int i = 0; i < kELTS_PER_THREAD; i++) { + accum[i] += chunkAccum[i]; + } + } + + PackedVec packedAccum; +#pragma unroll + for (int i = 0; i < kELTS_PER_THREAD; i++) { + packedAccum.elements[i] = fromFloat(accum[i]); + } + return packedAccum; +} + template inline __device__ void copyF4(T_IN* dst, T_IN const* src) { float4* dst4 = reinterpret_cast(dst); @@ -838,37 +889,9 @@ __global__ void __launch_bounds__(1024) } #undef RUN_ONESHOT_LOCAL_RANK } else { - // Compact deterministic fallback for larger instantiations. - PackedVec values[WorldSize]; - while (1) { - bool valid = true; -#pragma unroll - for (int r = 0; r < WorldSize; r++) { - auto loaded = loadPackedVolatile( - &stagePtrLocal[token * tokenDim * WorldSize + r * tokenDim + - packedIdx * kELTS_PER_THREAD]); - values[r].packed = loaded.packed; - valid &= !isLamportDirty(loaded); - } - if (valid) { - break; - } - } - - float accum[kELTS_PER_THREAD]; -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) { - accum[i] = 0.f; -#pragma unroll - for (int r = 0; r < WorldSize; r++) { - accum[i] += toFloat(values[r].elements[i]); - } - } - -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) { - packedAccum.elements[i] = fromFloat(accum[i]); - } + // Chunk large-world reductions to avoid register spills from a live values[WorldSize] array. + packedAccum = utils::reduceLamportRanksChunked( + stagePtrLocal, token, tokenDim, packedIdx * kELTS_PER_THREAD); } cudaTriggerProgrammaticLaunchCompletion(); @@ -1135,33 +1158,12 @@ __global__ __launch_bounds__(128) void twoshotAllreduceKernel(AllReduceKernelPar if (inBounds && destRank == params.rank) { int const localToken = token / WorldSize; - // Fully deterministic: every rank uses the exact same reduction order. - PackedVec values[WorldSize]; - while (1) { - bool valid = true; -#pragma unroll - for (int r = 0; r < WorldSize; r++) { - auto loaded = loadPackedVolatile( - &scatterBufLocal[localToken * params.tokenDim * WorldSize + r * params.tokenDim + - tokenOffset]); - values[r].packed = loaded.packed; - valid &= !isLamportDirty(loaded); - } - if (valid) { - break; - } - } - - PackedVec packedAccum; -#pragma unroll - for (int i = 0; i < kELTS_PER_THREAD; i++) { - float accum = 0.f; -#pragma unroll - for (int r = 0; r < WorldSize; r++) { - accum += toFloat(values[r].elements[i]); - } - packedAccum.elements[i] = fromFloat(accum); - } + // Fully deterministic: every rank uses the exact same reduction order. Chunking avoids + // register spills for large world sizes. + constexpr int kRankChunk = WorldSize < 16 ? WorldSize : 16; + PackedVec packedAccum = + utils::reduceLamportRanksChunked( + scatterBufLocal, localToken, params.tokenDim, tokenOffset); // Reduced values can round to the dirty sentinel; sanitize before broadcast polling. sanitizeLamportPayload(packedAccum); reinterpret_cast(&broadcastBufW[token * params.tokenDim])[packedIdx] =