[None][feat] Share Kimi K3 FP4 modeling and add NVFP4 support on SM103

cpp/tensorrt_llm/common/lamportUtils.cuh

@@ -26,8 +26,6 @@
 #include <cuda_runtime.h>
 #include <type_traits>
 
-#include <cooperative_groups.h>
-
 #include "tensorrt_llm/common/cudaTypeUtils.cuh"
 
 TRTLLM_NAMESPACE_BEGIN
@@ -196,17 +194,16 @@ struct LamportBufferLayout
                 (lamportIndex * numStages + stageIndex) * static_cast<size_t>(bytesPerBuffer / numStages)));
     }
 };
+
 // Current Index
 // Dirty Index
 // bytes_per_buffer
 // Dirty num_stages
 // Dirty bytes_to_clear = {stage0, stage1, stage2, stage3}  # We fix this to 4 stages
 // offset_access_ptr
 
-namespace cg = cooperative_groups;
-
 // PackedType is the one used in kernel for Lamport buffer (LDG.128 or LDG.64)
-template <typename PackedType = float4, bool UseCGA = false>
+template <typename PackedType = float4>
 struct __attribute__((aligned(32))) LamportFlags
 {
 public:
@@ -258,63 +255,28 @@ public:
 
     __device__ void ctaArrive()
     {
-        if constexpr (UseCGA)
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700))
+        uint32_t const barrierThreads
+            = ((static_cast<uint32_t>(blockDim.x) + kWARP_SIZE - 1U) / kWARP_SIZE) * kWARP_SIZE;
+        if (threadIdx.x < kWARP_SIZE)
         {
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-            cg::cluster_group cluster = cg::this_cluster();
-            __cluster_barrier_arrive();
-            if (cluster.block_rank() == 0 && threadIdx.x < kWARP_SIZE)
-            {
-                __cluster_barrier_wait();
-                arriveCounter(threadIdx.x);
-            }
-#else
-            __syncthreads();
+            asm volatile("barrier.cta.sync 1, %0;" ::"r"(barrierThreads) : "memory");
             arriveCounter(threadIdx.x);
-#endif
         }
         else
         {
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700))
-            uint32_t const barrierThreads
-                = ((static_cast<uint32_t>(blockDim.x) + kWARP_SIZE - 1U) / kWARP_SIZE) * kWARP_SIZE;
-            if (threadIdx.x < kWARP_SIZE)
-            {
-                asm volatile("barrier.cta.sync 1, %0;" ::"r"(barrierThreads) : "memory");
-                arriveCounter(threadIdx.x);
-            }
-            else
-            {
-                asm volatile("barrier.cta.arrive 1, %0;" ::"r"(barrierThreads) : "memory");
-            }
+            asm volatile("barrier.cta.arrive 1, %0;" ::"r"(barrierThreads) : "memory");
+        }
 #else
-            __syncthreads();
-            arriveCounter(threadIdx.x);
+        __syncthreads();
+        arriveCounter(threadIdx.x);
 #endif
-        }
     }
 
     __device__ void waitAndUpdate(uint4 bytesToClearPerStage)
     {
-        bool isLastCtaT0{false};
-        int targetCount{0};
-        if constexpr (UseCGA)
-        {
-#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;
-            targetCount = grid.num_clusters();
-#else
-            isLastCtaT0 = threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0;
-            targetCount = gridDim.x * gridDim.y * gridDim.z;
-#endif
-        }
-        else
-        {
-            isLastCtaT0 = threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0;
-            targetCount = gridDim.x * gridDim.y * gridDim.z;
-        }
+        bool const isLastCtaT0 = threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0;
+        int const targetCount = gridDim.x * gridDim.y * gridDim.z;
         if (isLastCtaT0)
         {
             uint4* flagPtr = reinterpret_cast<uint4*>(mBufferFlagsPtr);

cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu

@@ -336,6 +336,7 @@ struct MnnvlAllReduceKernelParams
     uint32_t* bufferFlags;
     bool waitForResults;
     QuantizationSFLayout layout;
+    bool earlyTrigger;
 };
 
 template <typename PackedType, typename T>
@@ -590,15 +591,20 @@ __global__ void __launch_bounds__(1024) oneshotAllreduceFusionKernel(MnnvlAllRed
     int threadOffset = token * params.tokenDim + packedIdx * kELTS_PER_THREAD;
 
     cudaGridDependencySynchronize();
+    if (params.earlyTrigger)
+    {
+        cudaTriggerProgrammaticLaunchCompletion();
+    }
 #else
     int packedIdx = blockIdx.y * blockDim.x + threadIdx.x;
     int token = blockIdx.x;
     // Offset w.r.t. the input shard
     int threadOffset = token * params.tokenDim + packedIdx * kELTS_PER_THREAD;
 #endif
 
-    // We only use 1 stage for the oneshot allreduce
-    LamportFlags<PackedType, true> flag(params.bufferFlags, 1);
+    // We only use 1 stage for the oneshot allreduce. Arrivals are counted per CTA: the RMSNorm epilogue's cluster
+    // reduction below must be each thread's only phase of the cluster barrier.
+    LamportFlags<PackedType> flag(params.bufferFlags, 1);
     T* stagePtrMcast = reinterpret_cast<T*>(flag.getCurLamportBuf(params.mcastPtr, 0));
     T* stagePtrLocal = reinterpret_cast<T*>(flag.getCurLamportBuf(params.inputPtrs[params.rank], 0));
     bool const inBounds = packedIdx * kELTS_PER_THREAD < params.tokenDim;
@@ -710,11 +716,12 @@ __global__ void __launch_bounds__(1024) oneshotAllreduceFusionKernel(MnnvlAllRed
             fullSum = 0.F;
             // Need to reduce over the entire cluster
             int const blockRank = cluster.block_rank();
+            // Every CTA of the cluster must have started before another CTA writes its shared memory.
+            cluster.sync();
             if (threadIdx.x < numBlocks)
             {
                 cluster.map_shared_rank(&sharedVal[0], threadIdx.x)[blockRank] = blockSum;
             }
-            // cluster.sync();
             cluster.barrier_wait(cluster.barrier_arrive());
             for (int i = 0; i < numBlocks; ++i)
             {
@@ -841,7 +848,7 @@ void oneshotAllreduceFusionOp(AllReduceFusionParams const& params)
         MnnvlAllReduceKernelParams<T> kernelParams{output, residualOut, input, residualIn, gamma, ucPtrs,
             reinterpret_cast<T*>(params.bufferPtrLocal), mcPtr, params.quantOut, params.scaleOut, params.scaleFactor,
             numTokens, tokenDim, params.nRanks, params.rank, static_cast<float>(params.epsilon), params.bufferFlags,
-            false, params.layout};
+            false, params.layout, params.earlyTrigger};
 
         switch (params.nRanks)
         {
@@ -890,7 +897,7 @@ __global__ __launch_bounds__(128) void twoshotAllreduceKernel(MnnvlAllReduceKern
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
     cudaGridDependencySynchronize();
 #endif
-    LamportFlags<PackedType, false> flag(params.bufferFlags, MNNVLTwoShotStage::NUM_STAGES);
+    LamportFlags<PackedType> flag(params.bufferFlags, MNNVLTwoShotStage::NUM_STAGES);
 
     T* scatterBufLocal
         = reinterpret_cast<T*>(flag.getCurLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::SCATTER));
@@ -1016,7 +1023,9 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(MnnvlAllReduceKernelParam
     T* smemResidual = reinterpret_cast<T*>(&smem[smemBufferSize]);
     T* smemGamma = reinterpret_cast<T*>(&smem[2 * smemBufferSize]);
 
-    LamportFlags<float4, UseCGA> flag(params.bufferFlags, MNNVLTwoShotStage::NUM_STAGES);
+    // Arrivals are counted per CTA: with UseCGA the RMSNorm cluster reduction below must be each thread's only phase
+    // of the cluster barrier.
+    LamportFlags<float4> flag(params.bufferFlags, MNNVLTwoShotStage::NUM_STAGES);
     T* input = reinterpret_cast<T*>(
         flag.getCurLamportBuf(reinterpret_cast<void*>(params.bufferInputPtr), MNNVLTwoShotStage::BROADCAST));
 
@@ -1137,11 +1146,12 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(MnnvlAllReduceKernelParam
             fullSum = 0.F;
             // Need to reduce over the entire cluster
             int const blockRank = cluster.block_rank();
+            // Every CTA of the cluster must have started before another CTA writes its shared memory.
+            cluster.sync();
             if (threadIdx.x < numBlocks)
             {
                 cluster.map_shared_rank(&sharedVal[0], threadIdx.x)[blockRank] = blockSum;
             }
-            // cluster.sync();
             cluster.barrier_wait(cluster.barrier_arrive());
             for (int i = 0; i < numBlocks; ++i)
             {

cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h

@@ -68,6 +68,7 @@ struct AllReduceFusionParams
     void* scaleOut = nullptr;           //!< NVFP4 scale-factor output (used by NVFP4 fusion patterns)
     float const* scaleFactor = nullptr; //!< Quantization scale factor
     QuantizationSFLayout layout = QuantizationSFLayout::SWIZZLED; //!< NVFP4 scale-factor layout
+    bool earlyTrigger = false;
     cudaStream_t stream;                                          //!< CUDA stream for asynchronous kernel execution
 
     //! @}

cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu

@@ -1046,12 +1046,18 @@ __global__ void __launch_bounds__(256, 1) attn_res_fwd_s1_single_cta_kernel(bf16
     bf16_t const* __restrict__ layer_res, bf16_t const* __restrict__ layer_res_add, bf16_t const* __restrict__ res_w,
     bf16_t const* __restrict__ rms_w, bf16_t const* __restrict__ output_rms_w, bf16_t* __restrict__ updated_layer_res,
     bf16_t* __restrict__ output, float* __restrict__ rsigma_out, float* __restrict__ probs_out,
-    float* __restrict__ logits_out, float rms_eps, float output_rms_eps, int num_tokens)
+    float* __restrict__ logits_out, float rms_eps, float output_rms_eps, int num_tokens, bool early_trigger)
 {
 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
     if constexpr (ENABLE_PDL)
     {
         cudaGridDependencySynchronize();
+        // The dependent grid still waits for this one to complete before reading its
+        // output; triggering here lets it launch and stream its weights while this kernel runs.
+        if (early_trigger)
+        {
+            cudaTriggerProgrammaticLaunchCompletion();
+        }
     }
 
     constexpr int H = 7168;
@@ -1321,7 +1327,10 @@ __global__ void __launch_bounds__(256, 1) attn_res_fwd_s1_single_cta_kernel(bf16
 
     if constexpr (ENABLE_PDL)
     {
-        cudaTriggerProgrammaticLaunchCompletion();
+        if (!early_trigger)
+        {
+            cudaTriggerProgrammaticLaunchCompletion();
+        }
     }
 #else
     if (cute::thread0())
@@ -1331,11 +1340,17 @@ __global__ void __launch_bounds__(256, 1) attn_res_fwd_s1_single_cta_kernel(bf16
 #endif
 }
 
+// Decode handoff options for the s1 kernels (see AttnResFwdParams).
+struct S1Handoff
+{
+    bool early_trigger = false;
+};
+
 template <int N, bool FUSE_OUTPUT_RMS_NORM = false, bool FUSE_LAYER_ADD = false>
 static void launch_s1_single_cta(bf16_t const* block_residual, bf16_t const* layer_residual,
     bf16_t const* layer_residual_add, bf16_t const* res_weight, bf16_t const* rms_weight,
     bf16_t const* output_rms_weight, bf16_t* updated_layer_residual, bf16_t* output, float* rsigma, float* probs,
-    float* logits, float rms_eps, float output_rms_eps, int num_tokens, cudaStream_t stream)
+    float* logits, float rms_eps, float output_rms_eps, int num_tokens, cudaStream_t stream, S1Handoff handoff = {})
 {
     if (tensorrt_llm::common::getEnvEnablePDL())
     {
@@ -1351,14 +1366,14 @@ static void launch_s1_single_cta(bf16_t const* block_residual, bf16_t const* lay
         config.numAttrs = 1;
         cudaLaunchKernelEx(&config, kernel, block_residual, layer_residual, layer_residual_add, res_weight, rms_weight,
             output_rms_weight, updated_layer_residual, output, rsigma, probs, logits, rms_eps, output_rms_eps,
-            num_tokens);
+            num_tokens, handoff.early_trigger);
     }
     else
     {
         attn_res_fwd_s1_single_cta_kernel<N, FUSE_OUTPUT_RMS_NORM, FUSE_LAYER_ADD, false>
             <<<num_tokens, 256, 0, stream>>>(block_residual, layer_residual, layer_residual_add, res_weight, rms_weight,
                 output_rms_weight, updated_layer_residual, output, rsigma, probs, logits, rms_eps, output_rms_eps,
-                num_tokens);
+                num_tokens, false);
     }
 }
 
@@ -1380,13 +1395,22 @@ __global__ void __launch_bounds__(256, 1) attn_res_fwd_s1_splitk_kernel(bf16_t c
     bf16_t const* __restrict__ layer_res, bf16_t const* __restrict__ layer_res_add, bf16_t const* __restrict__ res_w,
     bf16_t const* __restrict__ rms_w, bf16_t const* __restrict__ output_rms_w, bf16_t* __restrict__ updated_layer_res,
     bf16_t* __restrict__ output, float* __restrict__ rsigma_out, float* __restrict__ probs_out,
-    float* __restrict__ logits_out, float rms_eps, float output_rms_eps, int num_tokens)
+    float* __restrict__ logits_out, float rms_eps, float output_rms_eps, int num_tokens, bool early_trigger)
 {
 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
-    if constexpr (ENABLE_PDL)
+    auto const wait_for_previous_grid = [&]
     {
-        cudaGridDependencySynchronize();
-    }
+        if constexpr (ENABLE_PDL)
+        {
+            cudaGridDependencySynchronize();
+            // See attn_res_fwd_s1_single_cta_kernel.
+            if (early_trigger)
+            {
+                cudaTriggerProgrammaticLaunchCompletion();
+            }
+        }
+    };
+    wait_for_previous_grid();
 
     namespace cg = cooperative_groups;
     constexpr int H = 7168;
@@ -1439,6 +1463,7 @@ __global__ void __launch_bounds__(256, 1) attn_res_fwd_s1_splitk_kernel(bf16_t c
 
     float sq[N] = {};
     float dot[N] = {};
+    // Every candidate in one pass, so all loads are in flight together.
 #pragma unroll
     for (int ki = tid; ki < K_PER_CTA; ki += THREADS)
     {
@@ -1638,7 +1663,10 @@ __global__ void __launch_bounds__(256, 1) attn_res_fwd_s1_splitk_kernel(bf16_t c
 
     if constexpr (ENABLE_PDL)
     {
-        cudaTriggerProgrammaticLaunchCompletion();
+        if (!early_trigger)
+        {
+            cudaTriggerProgrammaticLaunchCompletion();
+        }
     }
 #else
     if (cute::thread0())
@@ -1652,13 +1680,14 @@ template <int N, int GROUPS = 8, bool FUSE_OUTPUT_RMS_NORM = false, bool FUSE_LA
 static void launch_s1_splitk(bf16_t const* block_residual, bf16_t const* layer_residual,
     bf16_t const* layer_residual_add, bf16_t const* res_weight, bf16_t const* rms_weight,
     bf16_t const* output_rms_weight, bf16_t* updated_layer_residual, bf16_t* output, float* rsigma, float* probs,
-    float* logits, float rms_eps, float output_rms_eps, int num_tokens, cudaStream_t stream)
+    float* logits, float rms_eps, float output_rms_eps, int num_tokens, cudaStream_t stream, S1Handoff handoff = {})
 {
     constexpr int K_PER_CTA = 7168 / GROUPS;
     constexpr int WARPS = 8;
     constexpr size_t smem_size
         = (size_t) N * K_PER_CTA * sizeof(float) + (size_t) WARPS * N * sizeof(float2) + (size_t) N * sizeof(float);
     bool const enable_pdl = tensorrt_llm::common::getEnvEnablePDL();
+    bool early_trigger = enable_pdl && handoff.early_trigger;
     auto kernel = enable_pdl ? &attn_res_fwd_s1_splitk_kernel<N, GROUPS, FUSE_OUTPUT_RMS_NORM, FUSE_LAYER_ADD, true>
                              : &attn_res_fwd_s1_splitk_kernel<N, GROUPS, FUSE_OUTPUT_RMS_NORM, FUSE_LAYER_ADD, false>;
     static std::once_flag attrs_set[2][64];
@@ -1678,7 +1707,7 @@ static void launch_s1_splitk(bf16_t const* block_residual, bf16_t const* layer_r
     void* args[] = {const_cast<bf16_t**>(&block_residual), const_cast<bf16_t**>(&layer_residual),
         const_cast<bf16_t**>(&layer_residual_add), const_cast<bf16_t**>(&res_weight), const_cast<bf16_t**>(&rms_weight),
         const_cast<bf16_t**>(&output_rms_weight), &updated_layer_residual, &output, &rsigma, &probs, &logits, &rms_eps,
-        &output_rms_eps, &num_tokens};
+        &output_rms_eps, &num_tokens, &early_trigger};
     cudaLaunchConfig_t config{};
     // One cluster per token. clusterDim stays GROUPS, so clusters are contiguous
     // spans of the grid and each token's DSM exchange stays within its own.
@@ -1887,18 +1916,19 @@ static void launchAttnResDecodeRmsNorm(AttnResFwdParams const& params, cudaStrea
 
     auto const* layer_residual_add = FUSE_LAYER_ADD ? params.layerResidualAdd : nullptr;
     auto* updated_layer_residual = FUSE_LAYER_ADD ? params.updatedLayerResidual : nullptr;
+    S1Handoff const handoff{params.earlyTrigger};
 
     if constexpr (N <= 4)
     {
         launch_s1_single_cta<N, true, FUSE_LAYER_ADD>(params.blockResidual, params.layerResidual, layer_residual_add,
             params.resWeight, params.rmsWeight, params.outputRmsWeight, updated_layer_residual, params.output, nullptr,
-            nullptr, nullptr, params.rmsEps, params.outputRmsEps, params.seqLen, stream);
+            nullptr, nullptr, params.rmsEps, params.outputRmsEps, params.seqLen, stream, handoff);
     }
     else
     {
         launch_s1_splitk<N, 8, true, FUSE_LAYER_ADD>(params.blockResidual, params.layerResidual, layer_residual_add,
             params.resWeight, params.rmsWeight, params.outputRmsWeight, updated_layer_residual, params.output, nullptr,
-            nullptr, nullptr, params.rmsEps, params.outputRmsEps, params.seqLen, stream);
+            nullptr, nullptr, params.rmsEps, params.outputRmsEps, params.seqLen, stream, handoff);
     }
 }
 

cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.h

@@ -53,6 +53,7 @@ struct AttnResFwdParams
     int hiddenSize;                        // H
     float rmsEps;
     float outputRmsEps;
+    bool earlyTrigger = false;
 };
 
 //! Launches the fused attention-residual forward on the supplied stream.