diff --git a/cpp/include/tensorrt_llm/runtime/ipcUtils.h b/cpp/include/tensorrt_llm/runtime/ipcUtils.h index 0a2e9fb87267..8663be993a3f 100644 --- a/cpp/include/tensorrt_llm/runtime/ipcUtils.h +++ b/cpp/include/tensorrt_llm/runtime/ipcUtils.h @@ -69,6 +69,7 @@ class AllReduceBuffers bool const fakeBuffers = false); TensorPtr mAllReduceCommPtrs; + TensorPtr mFlagPtrs; std::vector mIpcMemoryHandles; }; diff --git a/cpp/tensorrt_llm/kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/CMakeLists.txt index 8a43d66fc567..24fb05ae0d0d 100644 --- a/cpp/tensorrt_llm/kernels/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/CMakeLists.txt @@ -82,3 +82,4 @@ add_subdirectory(userbuffers) add_subdirectory(trtllmGenKernels) add_subdirectory(fusedLayernormKernels) add_subdirectory(groupRmsNormKernels) +add_subdirectory(llama4MinLatencyKernels) diff --git a/cpp/tensorrt_llm/kernels/customAllReduceKernels.cu b/cpp/tensorrt_llm/kernels/customAllReduceKernels.cu index f08319360830..39911eac613e 100644 --- a/cpp/tensorrt_llm/kernels/customAllReduceKernels.cu +++ b/cpp/tensorrt_llm/kernels/customAllReduceKernels.cu @@ -22,6 +22,7 @@ #include "tensorrt_llm/common/dataType.h" #include "tensorrt_llm/common/envUtils.h" #include +#include #include #include @@ -104,6 +105,16 @@ struct PackedOn16Bytes<__nv_bfloat16> #endif +__inline__ __device__ bool thread0() +{ + return !threadIdx.x && !threadIdx.y && !threadIdx.z; +} + +__inline__ __device__ bool block0() +{ + return !blockIdx.x && !blockIdx.y && !blockIdx.z; +} + // add two 128b data template inline __device__ int4 add128b(T& a, T& b) @@ -145,7 +156,7 @@ __inline__ __device__ void multi_gpu_barrier(uint32_t** signals, uint32_t const } __inline__ __device__ void block_barrier(uint32_t** signals, uint32_t const flag, size_t const local_rank, - size_t const world_size, int const tidx, int const bidx, int const grid_size) + size_t const world_size, int const tidx, int const bidx) { // After this function, the block of id == bidx of each GPU has reached the barrier if (tidx < world_size) @@ -155,11 +166,11 @@ __inline__ __device__ void block_barrier(uint32_t** signals, uint32_t const flag // Dimension 0 is the "listening" dimension, dimension 3 is "emitting" dimension // Block broadcast its flag (local_rank on emitting dimension) to all receivers - uint32_t flag_block_offset = world_size + bidx * world_size; + uint32_t flag_block_offset = (bidx + 1) * world_size; if (flag % 2 == 1) { - flag_block_offset += (grid_size + 1) * world_size; + flag_block_offset += (MAX_ALL_REDUCE_BLOCKS + 1) * world_size; } st_flag_release(flag, signals[tidx] + flag_block_offset + local_rank); @@ -175,6 +186,24 @@ __inline__ __device__ void block_barrier(uint32_t** signals, uint32_t const flag __syncthreads(); } +__inline__ __device__ void update_barrier_flag(uint32_t* barrier_flag_ptr, uint32_t* barrier_flag_counter_ptr) +{ + if (thread0()) + { + atomicAdd(barrier_flag_counter_ptr, 1); + + if (block0()) + { + auto blockNum = gridDim.x * gridDim.y * gridDim.z; + while (*reinterpret_cast(barrier_flag_counter_ptr) != blockNum) + { + } + *barrier_flag_ptr = ((*barrier_flag_ptr) + 1) % MAX_ALL_REDUCE_MODULES; + *barrier_flag_counter_ptr = 0; + } + } +} + namespace reduce_fusion { @@ -637,8 +666,9 @@ struct Reducer static __device__ __forceinline__ int4 allreduce(AllReduceParams& params, int global_offset) { using PackedStruct = typename PackedOn16Bytes::Type; - int ping = params.barrier_flag % 3; - int pong = (params.barrier_flag + 2) % 3; + auto const barrier_flag = *params.barrier_flag_ptr; + int ping = barrier_flag % 3; + int pong = (barrier_flag + 2) % 3; T const* local_input_buffer = reinterpret_cast(params.local_input_buffer_ptr); T* local_shared_buffer = reinterpret_cast( params.fusion_params.lamport_peer_comm_buffer_ptrs[params.local_rank + ping * MAX_RANKS_PER_NODE]); @@ -704,8 +734,9 @@ struct Reducer static __device__ __forceinline__ int4 allreduce(AllReduceParams& params, int global_offset) { using PackedStruct = typename PackedOn16Bytes::Type; - int ping = params.barrier_flag % 3; - int pong = (params.barrier_flag + 2) % 3; + auto const barrier_flag = *params.barrier_flag_ptr; + int ping = barrier_flag % 3; + int pong = (barrier_flag + 2) % 3; T const* local_input_buffer = reinterpret_cast(params.local_input_buffer_ptr); T* local_shared_buffer = reinterpret_cast( params.fusion_params.lamport_peer_comm_buffer_ptrs[params.local_rank + ping * MAX_RANKS_PER_NODE]); @@ -827,8 +858,8 @@ static __global__ void lamport_style_one_shot_all_reduce_norm_kernel(AllReducePa float denom = rsqrtf(acc / params.fusion_params.hidden_size + params.fusion_params.eps); sum_vec.packed = rms_norm(denom, sum_vec, weight_vec); *reinterpret_cast(local_final_output_buffer) = sum_vec.packed; - cudaTriggerProgrammaticLaunchCompletion(); + update_barrier_flag(params.barrier_flag_ptr, params.barrier_flag_counter_ptr); #endif } @@ -916,7 +947,6 @@ static __global__ void __launch_bounds__(1024, 1) one_shot_all_reduce_norm_kerne T const* bias_buffer = reinterpret_cast(params.fusion_params.bias_buffer); T const* residual_buffer = reinterpret_cast(params.fusion_params.residual_buffer); T const* weight_buffer = reinterpret_cast(params.fusion_params.weight_buffer); - T* local_shared_buffer = reinterpret_cast(params.peer_comm_buffer_ptrs[params.local_rank]); T* local_final_output_buffer = reinterpret_cast(params.local_output_buffer_ptr); T* intermediate_buffer = reinterpret_cast(params.fusion_params.intermediate_buffer); @@ -925,21 +955,25 @@ static __global__ void __launch_bounds__(1024, 1) one_shot_all_reduce_norm_kerne local_input_buffer += block_offset; residual_buffer += block_offset; - local_shared_buffer += block_offset; local_final_output_buffer += block_offset; intermediate_buffer += block_offset; +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + cudaGridDependencySynchronize(); +#endif + + auto const barrier_flag = *params.barrier_flag_ptr; + auto const buffer_offset = (barrier_flag % 2 == 0) ? 0 : params.ranks_per_node; + T* local_shared_buffer = reinterpret_cast(params.peer_comm_buffer_ptrs[params.local_rank + buffer_offset]); + T* buffers[RanksPerNode]; #pragma unroll for (int ii = 0; ii < RanksPerNode; ++ii) { int rank = (params.local_rank + ii) % RanksPerNode; - buffers[ii] = reinterpret_cast(params.peer_comm_buffer_ptrs[rank]); + buffers[ii] = reinterpret_cast(params.peer_comm_buffer_ptrs[rank + buffer_offset]); } - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) - cudaGridDependencySynchronize(); -#endif + local_shared_buffer += block_offset; for (int offset = thread_offset; offset < norm_this_block * params.fusion_params.hidden_size; offset += blockDim.x * kPackedSize) @@ -947,8 +981,7 @@ static __global__ void __launch_bounds__(1024, 1) one_shot_all_reduce_norm_kerne *reinterpret_cast(&local_shared_buffer[offset]) = *reinterpret_cast(&local_input_buffer[offset]); } - block_barrier( - params.peer_barrier_ptrs_in, params.barrier_flag, params.local_rank, RanksPerNode, tid, bid, gridDim.x); + block_barrier(params.peer_barrier_ptrs_in, barrier_flag, params.local_rank, RanksPerNode, tid, bid); for (int norm_idx = 0; norm_idx < norm_this_block; ++norm_idx) { int norm_offset = norm_idx * params.fusion_params.hidden_size; @@ -1004,6 +1037,7 @@ static __global__ void __launch_bounds__(1024, 1) one_shot_all_reduce_norm_kerne #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) cudaTriggerProgrammaticLaunchCompletion(); #endif + update_barrier_flag(params.barrier_flag_ptr, params.barrier_flag_counter_ptr); } template @@ -1023,7 +1057,6 @@ static __global__ void __launch_bounds__(1024, 1) one_shot_prenorm_all_reduce_no T const* weight_buffer = reinterpret_cast(params.fusion_params.weight_buffer); T const* weight_buffer_pre_residual_norm = reinterpret_cast(params.fusion_params.weight_buffer_pre_residual_norm); - T* local_shared_buffer = reinterpret_cast(params.peer_comm_buffer_ptrs[params.local_rank]); T* local_final_output_buffer = reinterpret_cast(params.local_output_buffer_ptr); T* intermediate_buffer = reinterpret_cast(params.fusion_params.intermediate_buffer); @@ -1032,21 +1065,25 @@ static __global__ void __launch_bounds__(1024, 1) one_shot_prenorm_all_reduce_no local_input_buffer += block_offset; residual_buffer += block_offset; - local_shared_buffer += block_offset; local_final_output_buffer += block_offset; intermediate_buffer += block_offset; +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + cudaGridDependencySynchronize(); +#endif + + auto const barrier_flag = *params.barrier_flag_ptr; + auto const buffer_offset = (barrier_flag % 2 == 0) ? 0 : params.ranks_per_node; + T* local_shared_buffer = reinterpret_cast(params.peer_comm_buffer_ptrs[params.local_rank + buffer_offset]); + T* buffers[RanksPerNode]; #pragma unroll for (int ii = 0; ii < RanksPerNode; ++ii) { int rank = (params.local_rank + ii) % RanksPerNode; - buffers[ii] = reinterpret_cast(params.peer_comm_buffer_ptrs[rank]); + buffers[ii] = reinterpret_cast(params.peer_comm_buffer_ptrs[rank + buffer_offset]); } - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) - cudaGridDependencySynchronize(); -#endif + local_shared_buffer += block_offset; for (int offset = thread_offset; offset < norm_this_block * params.fusion_params.hidden_size; offset += blockDim.x * kPackedSize) @@ -1054,8 +1091,7 @@ static __global__ void __launch_bounds__(1024, 1) one_shot_prenorm_all_reduce_no *reinterpret_cast(&local_shared_buffer[offset]) = *reinterpret_cast(&local_input_buffer[offset]); } - block_barrier( - params.peer_barrier_ptrs_in, params.barrier_flag, params.local_rank, RanksPerNode, tid, bid, gridDim.x); + block_barrier(params.peer_barrier_ptrs_in, barrier_flag, params.local_rank, RanksPerNode, tid, bid); for (int norm_idx = 0; norm_idx < norm_this_block; ++norm_idx) { int norm_offset = norm_idx * params.fusion_params.hidden_size; @@ -1117,6 +1153,7 @@ static __global__ void __launch_bounds__(1024, 1) one_shot_prenorm_all_reduce_no #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) cudaTriggerProgrammaticLaunchCompletion(); #endif + update_barrier_flag(params.barrier_flag_ptr, params.barrier_flag_counter_ptr); } template @@ -1303,7 +1340,7 @@ void lamport_initialize_kernel_launcher(void* buffer, size_t size, cudaStream_t } }; // namespace reduce_fusion -template +template static __global__ void oneShotAllReduceKernel(AllReduceParams params) { // Suppose that two GPUs participate in the AR exchange, and we start four blocks. @@ -1318,28 +1355,30 @@ static __global__ void oneShotAllReduceKernel(AllReduceParams params) // 2. B0 on GPU 0 and B0 on GPU 1 wait for each other (block_barrier) // 3. B0 on GPU 0 pull and sum the chunk from GPU 1, writes the result to local_output // - // With COPY_INPUT == false, skip step 1. and use gpu_barrier instead of block barrier during step 2. - // We only to know if the other GPU as arrived at the AR kernel, that would mean that data is ready - // // With PUSH_MODE, we consider that the shared buffer is of size: - // params.peer_comm_buffer_ptrs: [world_size, world_size, message_size] + // params.peer_comm_buffer_ptrs: [world_size * 2, world_size, message_size] + // Even plugins use ping buffers, odd plugins use pong. + // That way, we don't need to wait for other GPUs to be done + // before copying input tensor to workspace. + // For each plugin, the buffer is of size: [world_size, world_size, message_size] // // Here the step-by-step behavior of one block: // 1. B0 push the chunk is it responsible for into all other GPUs: - // params.peer_comm_buffer_ptrs[:, local_gpu, B0 slice] + // peer_comm_buffer_ptrs[:, local_gpu, B0 slice] // 2. block sync so the block is shared by other GPUs - // 3. Reduce along second dimension params.peer_comm_buffer_ptrs[local_gpu, :, B0 slice] + // 3. Reduce along second dimension peer_comm_buffer_ptrs[local_gpu, :, B0 slice] int const bidx = blockIdx.x; int const tidx = threadIdx.x; - int const grid_size = gridDim.x; + auto const barrier_flag = *params.barrier_flag_ptr; + auto const buffer_offset = (barrier_flag % 2 == 0) ? 0 : params.ranks_per_node; // The number of elements packed into one for comms static constexpr int PACKED_ELTS = 16 / sizeof(T); using PackedStruct = typename PackedOn16Bytes::Type; T const* local_input_buffer = reinterpret_cast(params.local_input_buffer_ptr); - T* local_shared_buffer = reinterpret_cast(params.peer_comm_buffer_ptrs[params.local_rank]); + T* local_shared_buffer = reinterpret_cast(params.peer_comm_buffer_ptrs[params.local_rank + buffer_offset]); T* local_output_buffer = reinterpret_cast(params.local_output_buffer_ptr); // Start and end offsets of the thread @@ -1352,45 +1391,35 @@ static __global__ void oneShotAllReduceKernel(AllReduceParams params) { // buffers[0] is always the local buffers. Helps load balancing reads. int rank = (params.local_rank + ii) % RANKS_PER_NODE; - buffers[ii] = reinterpret_cast(params.peer_comm_buffer_ptrs[rank]); + buffers[ii] = reinterpret_cast(params.peer_comm_buffer_ptrs[rank + buffer_offset]); } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) cudaGridDependencySynchronize(); #endif - if constexpr (PUSH_MODE || COPY_INPUT) + // Copy from local buffer to shareable buffer + for (size_t iter_offset = chunk_start; iter_offset < chunk_end; iter_offset += blockDim.x * PACKED_ELTS) { - // Copy from local buffer to shareable buffer - for (size_t iter_offset = chunk_start; iter_offset < chunk_end; iter_offset += blockDim.x * PACKED_ELTS) + if constexpr (PUSH_MODE) { - if constexpr (PUSH_MODE) - { #pragma unroll - for (int ii = 0; ii < RANKS_PER_NODE; ++ii) - { - *reinterpret_cast(&buffers[ii][params.local_rank * params.elts_total + iter_offset]) - = *reinterpret_cast(&local_input_buffer[iter_offset]); - } - } - else + for (int ii = 0; ii < RANKS_PER_NODE; ++ii) { - *reinterpret_cast(&local_shared_buffer[iter_offset]) + *reinterpret_cast(&buffers[ii][params.local_rank * params.elts_total + iter_offset]) = *reinterpret_cast(&local_input_buffer[iter_offset]); } } - - // wait for equivalent blocks of other GPUs to have copied data to their shareable buffer - block_barrier( - params.peer_barrier_ptrs_in, params.barrier_flag, params.local_rank, RANKS_PER_NODE, tidx, bidx, grid_size); - } - else - { - // In the non-copy case, we assume that once the kernel has been started, data is ready to be consumed - multi_gpu_barrier( - params.peer_barrier_ptrs_in, params.barrier_flag, params.local_rank, RANKS_PER_NODE, tidx, bidx); + else + { + *reinterpret_cast(&local_shared_buffer[iter_offset]) + = *reinterpret_cast(&local_input_buffer[iter_offset]); + } } + // wait for equivalent blocks of other GPUs to have copied data to their shareable buffer + block_barrier(params.peer_barrier_ptrs_in, barrier_flag, params.local_rank, RANKS_PER_NODE, tidx, bidx); + // Each block accumulates the values from the different GPUs on the same node. for (size_t iter_offset = chunk_start; iter_offset < chunk_end; iter_offset += blockDim.x * PACKED_ELTS) { @@ -1427,10 +1456,10 @@ static __global__ void oneShotAllReduceKernel(AllReduceParams params) #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) cudaTriggerProgrammaticLaunchCompletion(); #endif + update_barrier_flag(params.barrier_flag_ptr, params.barrier_flag_counter_ptr); } -template +template static __global__ void __launch_bounds__(512, 1) twoShotAllReduceKernel(AllReduceParams params) { // Suppose that two GPUs participate in the AR exchange, and we start two blocks. @@ -1452,40 +1481,45 @@ static __global__ void __launch_bounds__(512, 1) twoShotAllReduceKernel(AllReduc // 5. B0 writes result to local_output. It gathers each chunk from its responsible GPU. // For example, here it reads the first chunk from GPU 0 and second chunk from GPU 1. // - // With COPY_INPUT == false, skip step 1. and use gpu_barrier instead of block barrier during step 2. - // We only to know if the other GPU as arrived at the AR kernel, that would mean that data is ready - // to be read. - // // Note that compared to one-shot, one block (CTA) writes multiple input chunks and write multiple output chunks. // However, it's only responsible for the summation of a single chunk. // // With PUSH_MODE, we consider that the shared buffer is of size: // params.peer_comm_buffer_ptrs: [world_size, world_size, message_size / world_size] + // Even plugins use ping buffers, odd plugins use pong. + // That way, we don't need to wait for other GPUs to be done + // before copying input tensor to workspace. + // For each plugin, the buffer is of size: [world_size, world_size, message_size / world_size] // // Here the step-by-step behavior of one block: // 1. B0 push the chunks is it responsible for into the corresponding GPUs: - // params.peer_comm_buffer_ptrs[target_gpu, local_gpu, current B0 slice] + // peer_comm_buffer_ptrs[target_gpu, local_gpu, current B0 slice] // 2. block sync so the blocks have been shared by other GPUs - // 3. Reduce along second dimension params.peer_comm_buffer_ptrs[local_gpu, :, B0 slice] + // 3. Reduce along second dimension peer_comm_buffer_ptrs[local_gpu, :, B0 slice] // 4. block barrier (corresponding blocks have finished reduction) - // 5. pull and write on local buffer, by reading params.peer_comm_buffer_ptrs[:, 0, B0 slice] (reduction result is + // 5. pull and write on local buffer, by reading peer_comm_buffer_ptrs[:, 0, B0 slice] (reduction result is // written at index 0 of 2nd dim) int const bidx = blockIdx.x; int const tidx = threadIdx.x; - int const grid_size = gridDim.x; - // The number of elements packed into one for comms static constexpr int PACKED_ELTS = 16 / sizeof(T); using PackedType = typename PackedOn16Bytes::Type; T const* local_input_buffer = reinterpret_cast(params.local_input_buffer_ptr); - T* local_shared_buffer = reinterpret_cast(params.peer_comm_buffer_ptrs[params.local_rank]); T* local_output_buffer = reinterpret_cast(params.local_output_buffer_ptr); size_t const chunk_start = bidx * params.elts_per_block + tidx * PACKED_ELTS; size_t const chunk_end = min(chunk_start + params.elts_per_block, params.elts_per_rank); +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + cudaGridDependencySynchronize(); +#endif + + auto const barrier_flag = *params.barrier_flag_ptr; + auto const buffer_offset = (barrier_flag % 2 == 0) ? 0 : params.ranks_per_node; + T* local_shared_buffer = reinterpret_cast(params.peer_comm_buffer_ptrs[params.local_rank + buffer_offset]); + T* buffers[RANKS_PER_NODE]; int ranks[RANKS_PER_NODE]; #pragma unroll @@ -1494,48 +1528,34 @@ static __global__ void __launch_bounds__(512, 1) twoShotAllReduceKernel(AllReduc // A mapping of the ranks to scatter reads as much as possible int rank = (params.local_rank + ii) % RANKS_PER_NODE; ranks[ii] = rank; - buffers[ii] = reinterpret_cast(params.peer_comm_buffer_ptrs[rank]); + buffers[ii] = reinterpret_cast(params.peer_comm_buffer_ptrs[rank + buffer_offset]); } -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) - cudaGridDependencySynchronize(); -#endif - - if constexpr (PUSH_MODE || COPY_INPUT) + // Copy all blocks from local buffer to shareable buffer + for (size_t local_offset = chunk_start; local_offset < chunk_end; local_offset += blockDim.x * PACKED_ELTS) { - // Copy all blocks from local buffer to shareable buffer - for (size_t local_offset = chunk_start; local_offset < chunk_end; local_offset += blockDim.x * PACKED_ELTS) - { #pragma unroll - for (int ii = 0; ii < RANKS_PER_NODE; ++ii) + for (int ii = 0; ii < RANKS_PER_NODE; ++ii) + { + size_t offset_rank = ranks[ii] * params.elts_per_rank + local_offset; + if (offset_rank >= params.elts_total) { - size_t offset_rank = ranks[ii] * params.elts_per_rank + local_offset; - if (offset_rank >= params.elts_total) - { - continue; - } + continue; + } - if constexpr (PUSH_MODE) - { - *reinterpret_cast(&buffers[ii][params.local_rank * params.elts_per_rank + local_offset]) - = *reinterpret_cast(&local_input_buffer[offset_rank]); - } - else - { - *reinterpret_cast(&local_shared_buffer[offset_rank]) - = *reinterpret_cast(&local_input_buffer[offset_rank]); - } + if constexpr (PUSH_MODE) + { + *reinterpret_cast(&buffers[ii][params.local_rank * params.elts_per_rank + local_offset]) + = *reinterpret_cast(&local_input_buffer[offset_rank]); + } + else + { + *reinterpret_cast(&local_shared_buffer[offset_rank]) + = *reinterpret_cast(&local_input_buffer[offset_rank]); } } - block_barrier( - params.peer_barrier_ptrs_in, params.barrier_flag, params.local_rank, RANKS_PER_NODE, tidx, bidx, grid_size); - } - else - { - // In the non-copy case, we assume that once the kernel has been started, data is ready to be consumed - multi_gpu_barrier( - params.peer_barrier_ptrs_in, params.barrier_flag, params.local_rank, RANKS_PER_NODE, tidx, bidx); } + block_barrier(params.peer_barrier_ptrs_in, barrier_flag, params.local_rank, RANKS_PER_NODE, tidx, bidx); // Each block accumulates the values from the different GPUs on the same node. for (size_t local_offset = chunk_start; local_offset < chunk_end; local_offset += blockDim.x * PACKED_ELTS) @@ -1580,8 +1600,7 @@ static __global__ void __launch_bounds__(512, 1) twoShotAllReduceKernel(AllReduc } } - block_barrier( - params.peer_barrier_ptrs_out, params.barrier_flag, params.local_rank, RANKS_PER_NODE, tidx, bidx, grid_size); + block_barrier(params.peer_barrier_ptrs_out, barrier_flag, params.local_rank, RANKS_PER_NODE, tidx, bidx); // Gather all needed elts from other intra-node ranks for (size_t local_offset = chunk_start; local_offset < chunk_end; local_offset += blockDim.x * PACKED_ELTS) @@ -1630,10 +1649,10 @@ static __global__ void __launch_bounds__(512, 1) twoShotAllReduceKernel(AllReduc *reinterpret_cast(&local_output_buffer[offset_rank]) = sums.packed; } } - #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) cudaTriggerProgrammaticLaunchCompletion(); #endif + update_barrier_flag(params.barrier_flag_ptr, params.barrier_flag_counter_ptr); } bool configurationSupported(AllReduceStrategyType algo, size_t msg_size, size_t n_ranks, nvinfer1::DataType type) @@ -1696,8 +1715,7 @@ std::tuple kernelLaunchConfig(AllReduceStrategyType algo, AllReducePar return std::make_tuple(blocks_per_grid, threads_per_block); } -template +template void AllReduceNormKernelLaunch(AllReduceStrategyType algo, AllReduceStrategyConfig config, AllReduceFusionOp fusionOp, AllReduceParams& params, cudaStream_t stream) { @@ -1711,14 +1729,8 @@ void AllReduceNormKernelLaunch(AllReduceStrategyType algo, AllReduceStrategyConf } else { - TLLM_CHECK_WITH_INFO(!(USE_MEMCPY && PUSH_MODE), "Memcpy cannot be used with PUSH_MODE."); size_t elts_per_thread = 16 / sizeof(T); auto [blocks_per_grid, threads_per_block] = kernelLaunchConfig(algo, params, elts_per_thread); - if (USE_MEMCPY) - { - cudaMemcpyAsync(params.peer_comm_buffer_ptrs[params.local_rank], params.local_input_buffer_ptr, - params.elts_total * sizeof(T), cudaMemcpyDeviceToDevice, stream); - } auto output_ptr = params.local_output_buffer_ptr; params.local_output_buffer_ptr = params.fusion_params.intermediate_buffer; @@ -1738,11 +1750,11 @@ void AllReduceNormKernelLaunch(AllReduceStrategyType algo, AllReduceStrategyConf kernelConfig.numAttrs = 1; TLLM_CUDA_CHECK(cudaLaunchKernelEx( - &kernelConfig, twoShotAllReduceKernel, params)); + &kernelConfig, twoShotAllReduceKernel, params)); } else { - twoShotAllReduceKernel + twoShotAllReduceKernel <<>>(params); } params.local_output_buffer_ptr = output_ptr; @@ -1750,48 +1762,38 @@ void AllReduceNormKernelLaunch(AllReduceStrategyType algo, AllReduceStrategyConf } } -template +template void AllReduceNormDispatch(AllReduceStrategyType algo, AllReduceStrategyConfig config, AllReduceFusionOp fusionOp, AllReduceParams& params, cudaStream_t stream) { if (params.fusion_params.bias_buffer && params.fusion_params.weight_buffer) { - AllReduceNormKernelLaunch( - algo, config, fusionOp, params, stream); + AllReduceNormKernelLaunch(algo, config, fusionOp, params, stream); } else if (params.fusion_params.bias_buffer && !params.fusion_params.weight_buffer) { - AllReduceNormKernelLaunch( - algo, config, fusionOp, params, stream); + AllReduceNormKernelLaunch(algo, config, fusionOp, params, stream); } else if (!params.fusion_params.bias_buffer && params.fusion_params.weight_buffer) { - AllReduceNormKernelLaunch( - algo, config, fusionOp, params, stream); + AllReduceNormKernelLaunch(algo, config, fusionOp, params, stream); } else { - AllReduceNormKernelLaunch( - algo, config, fusionOp, params, stream); + AllReduceNormKernelLaunch(algo, config, fusionOp, params, stream); } } -template +template void AllReduceDispatch(AllReduceStrategyType algo, AllReduceStrategyConfig config, AllReduceFusionOp fusionOp, AllReduceParams& params, cudaStream_t stream) { TLLM_CHECK(fusionOp == AllReduceFusionOp::NONE); - TLLM_CHECK_WITH_INFO(!(USE_MEMCPY && PUSH_MODE), "Memcpy cannot be used with PUSH_MODE."); size_t elts_per_thread = 16 / sizeof(T); auto [blocks_per_grid, threads_per_block] = kernelLaunchConfig(algo, params, elts_per_thread); - if (USE_MEMCPY) - { - cudaMemcpyAsync(params.peer_comm_buffer_ptrs[params.local_rank], params.local_input_buffer_ptr, - params.elts_total * sizeof(T), cudaMemcpyDeviceToDevice, stream); - } if (algo == AllReduceStrategyType::ONESHOT) { - auto* kernel_instance = &oneShotAllReduceKernel; + auto* kernel_instance = &oneShotAllReduceKernel; cudaLaunchConfig_t config; config.gridDim = blocks_per_grid; config.blockDim = threads_per_block; @@ -1806,7 +1808,7 @@ void AllReduceDispatch(AllReduceStrategyType algo, AllReduceStrategyConfig confi } else { - auto* kernel_instance = &twoShotAllReduceKernel; + auto* kernel_instance = &twoShotAllReduceKernel; cudaLaunchConfig_t config; config.gridDim = blocks_per_grid; config.blockDim = threads_per_block; @@ -1844,11 +1846,18 @@ void AllReduceDispatchPushMode(AllReduceStrategyType algo, AllReduceStrategyConf if (static_cast>(config) & static_cast>(AllReduceStrategyConfig::USE_MEMCPY)) { - AllReduceDispatchMemcpy(algo, config, fusionOp, params, stream); + TLLM_LOG_DEBUG("USE_MEMCPY is deprecated and has no effect. "); + } + + if (fusionOp == AllReduceFusionOp::NONE) + { + TLLM_LOG_DEBUG("AllReduceDispatch enabled"); + AllReduceDispatch(algo, config, fusionOp, params, stream); } else { - AllReduceDispatchMemcpy(algo, config, fusionOp, params, stream); + TLLM_LOG_DEBUG("AllReduceNormDispatch enabled"); + AllReduceNormDispatch(algo, config, fusionOp, params, stream); } } @@ -1897,20 +1906,12 @@ AllReduceParams AllReduceParams::deserialize(int64_t* buffer, size_t tpSize, siz flag_offset = 1; } auto const flag_ptr - = &buffer[tensorrt_llm::utils::customAllReduceUtils::NUM_POINTERS_PER_RANK * tpSize + flag_offset]; - // cannot use 0 since 0 represents released state for barrier - *flag_ptr += 1; - TLLM_LOG_TRACE("AllReduceParams's flag value is %d, flag offset %d", *flag_ptr, flag_offset); - uint32_t flag_value = *flag_ptr; + = buffer[tensorrt_llm::utils::customAllReduceUtils::NUM_POINTERS_PER_RANK * tpSize + flag_offset]; AllReduceParams params; - // Even plugins use ping buffers, odd plugins use pong. - // That way, we don't need to wait for other GPUs to be done - // before copying input tensor to workspace. - auto const buffer_offset = (flag_value % 2 == 0) ? 0 : tpSize; - for (int i = 0; i < tpSize; ++i) + for (int i = 0; i < tpSize * 2; ++i) { - params.peer_comm_buffer_ptrs[i] = buffer_ptrs[buffer_offset + i]; + params.peer_comm_buffer_ptrs[i] = buffer_ptrs[i]; } for (int i = 0; i < tpSize; ++i) { @@ -1920,7 +1921,9 @@ AllReduceParams AllReduceParams::deserialize(int64_t* buffer, size_t tpSize, siz { params.peer_barrier_ptrs_out[i] = reinterpret_cast(buffer_ptrs[3 * tpSize + i]); } - params.barrier_flag = flag_value; + params.barrier_flag_ptr = reinterpret_cast(flag_ptr); + params.barrier_flag_counter_ptr = reinterpret_cast( + buffer[tensorrt_llm::utils::customAllReduceUtils::NUM_POINTERS_PER_RANK * tpSize + 2]); params.ranks_per_node = tpSize; params.local_rank = tpRank; diff --git a/cpp/tensorrt_llm/kernels/customAllReduceKernels.h b/cpp/tensorrt_llm/kernels/customAllReduceKernels.h index 0a3ec6ae4f0d..b4bd22e49e3e 100644 --- a/cpp/tensorrt_llm/kernels/customAllReduceKernels.h +++ b/cpp/tensorrt_llm/kernels/customAllReduceKernels.h @@ -19,6 +19,7 @@ #include #include #include +#include #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" @@ -28,6 +29,9 @@ namespace tensorrt_llm::kernels constexpr size_t WARP_SIZE = 32; constexpr size_t MAX_ALL_REDUCE_BLOCKS = 24; +// Use max modules to avoid overflow and ABA problem when the block num changes for barrier_flag +// Not a perfect solution, but it has large chance that it is correct +constexpr size_t MAX_ALL_REDUCE_MODULES = std::numeric_limits::max() / 6 * 6; constexpr size_t MAX_RANKS_PER_NODE = 16; constexpr size_t DEFAULT_BLOCK_SIZE = 512; @@ -131,10 +135,11 @@ struct AllReduceParams size_t rank_offset; size_t ranks_per_node; size_t local_rank; - uint32_t barrier_flag; + uint32_t* barrier_flag_ptr; + uint32_t* barrier_flag_counter_ptr; uint32_t* peer_barrier_ptrs_in[MAX_RANKS_PER_NODE]; uint32_t* peer_barrier_ptrs_out[MAX_RANKS_PER_NODE]; - void* peer_comm_buffer_ptrs[MAX_RANKS_PER_NODE]; + void* peer_comm_buffer_ptrs[MAX_RANKS_PER_NODE * 2]; void* local_output_buffer_ptr; void const* local_input_buffer_ptr; @@ -154,4 +159,9 @@ void residualRmsNorm( void lamportInitialize(void* buffer, size_t size, nvinfer1::DataType dataType, cudaStream_t stream); +namespace reduce_fusion +{ +bool is_lamport_supported(nvinfer1::DataType dataType, int token_num, int hidden_size); +} + } // namespace tensorrt_llm::kernels diff --git a/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/CMakeLists.txt new file mode 100644 index 000000000000..472798c24be0 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/CMakeLists.txt @@ -0,0 +1,25 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2025-2025 NVIDIA CORPORATION & +# AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not +# use this file except in compliance with the License. You may obtain a copy of +# the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations under +# the License. +# + +file(GLOB_RECURSE SRC_CPP *.cpp) +file(GLOB_RECURSE SRC_CU *.cu) + +add_library(llama4_min_latency_kernels OBJECT ${SRC_CPP} ${SRC_CU}) +set_property(TARGET llama4_min_latency_kernels + PROPERTY POSITION_INDEPENDENT_CODE ON) +set_property(TARGET llama4_min_latency_kernels + PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) diff --git a/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Bf16Bf16Gemm.cu b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Bf16Bf16Gemm.cu new file mode 100644 index 000000000000..1c7d8a25a722 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Bf16Bf16Gemm.cu @@ -0,0 +1,128 @@ +/* + * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Bf16Bf16Gemm.h" +#include "tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Utils.cuh" + +namespace tensorrt_llm::kernels::llama4_min_latency::llama4_bf16_bf16_gemm +{ + +struct __align__(8) aligned_bf16x4 +{ + __align__(8) __nv_bfloat16 data[VEC_SIZE]; +}; + +__global__ void llama4_bf16_bf16_gemm_kernel(int num_tokens, + __nv_bfloat16 const* __restrict__ A, // Input vector [num_tokens][5120] + __nv_bfloat16 const* __restrict__ B, // Input matrix [128][5120] + __nv_bfloat16* __restrict__ C // Output vector [num_tokens][128] +) +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + // Shared memory for block reduction + __shared__ float reduce_buffer[BLOCK_SIZE]; + + // Each thread accumulates its partial sum + float2 thread_sum; + thread_sum.x = 0.0f; + thread_sum.y = 0.0f; + + // Each thread processes 4 elements at a time, 5 times + int const token_idx = blockIdx.x / NUM_EXPERTS; + int const row = blockIdx.x % NUM_EXPERTS; // Matrix row / Output element index + int const tid = threadIdx.x; // Thread ID within the block + + // FDL prefetch all B data + aligned_bf16x4 b_vec[GEMM_K / BLOCK_SIZE / VEC_SIZE]; +#pragma unroll + for (int chunk = 0; chunk < GEMM_K / BLOCK_SIZE / VEC_SIZE; chunk++) + { + // Base index for this chunk + int base_idx = chunk * BLOCK_SIZE + tid; + + // Load 4 elements at once + b_vec[chunk] = reinterpret_cast(B)[row * GEMM_K / VEC_SIZE + base_idx]; + } + + asm volatile("griddepcontrol.wait;" ::: "memory"); + + // Process 5 chunks of 4 elements each +#pragma unroll + for (int chunk = 0; chunk < GEMM_K / BLOCK_SIZE / VEC_SIZE; chunk++) + { + // Base index for this chunk + int base_idx = chunk * BLOCK_SIZE + tid; + + // Load 4 elements at once + aligned_bf16x4 a_vec = reinterpret_cast(A)[token_idx * GEMM_K / VEC_SIZE + base_idx]; +#pragma unroll + for (int i = 0; i < VEC_SIZE; i += 2) + { + + float2 a_val = make_float2(a_vec.data[i], a_vec.data[i + 1]); + float2 b_val = make_float2(b_vec[chunk].data[i], b_vec[chunk].data[i + 1]); + + thread_sum = ffma2(a_val, b_val, thread_sum); + } + } + + // Warp-level reduction + float warp_sum = thread_sum.x + thread_sum.y; + for (int offset = warpSize / 2; offset > 0; offset >>= 1) + { + warp_sum += __shfl_down_sync(0xffffffff, warp_sum, offset); + } + + // First thread in each warp writes to shared memory + if (tid % warpSize == 0) + { + reduce_buffer[tid / warpSize] = warp_sum; + } + __syncthreads(); + + // Final thread reduces across warps and writes the result + if (tid == 0) + { + float block_sum = 0.0f; + for (int i = 0; i < BLOCK_SIZE / warpSize; i++) + { + block_sum += reduce_buffer[i]; + } + C[token_idx * NUM_EXPERTS + row] = __float2bfloat16(block_sum); + } +#endif +} + +void llama4_bf16_bf16_gemm_launcher( + int num_tokens, __nv_bfloat16 const* A, __nv_bfloat16 const* B, __nv_bfloat16* C, cudaStream_t stream) +{ + + int const grid_size = NUM_EXPERTS * num_tokens; + + void* args[] = {(void*) &num_tokens, (void*) &A, (void*) &B, (void*) &C}; + launch_kernel_fdl(dim3(grid_size), dim3(BLOCK_SIZE), stream, (void*) llama4_bf16_bf16_gemm_kernel, args, 4); +} + +void llama4_bf16_bf16_gemm_op(int num_tokens, void const* A, void const* B, void* C, cudaStream_t stream) +{ + __nv_bfloat16 const* A_bf16 = static_cast<__nv_bfloat16 const*>(A); + __nv_bfloat16 const* B_bf16 = static_cast<__nv_bfloat16 const*>(B); + __nv_bfloat16* C_bf16 = static_cast<__nv_bfloat16*>(C); + + llama4_bf16_bf16_gemm_launcher(num_tokens, A_bf16, B_bf16, C_bf16, stream); +} + +} // namespace tensorrt_llm::kernels::llama4_min_latency::llama4_bf16_bf16_gemm diff --git a/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Bf16Bf16Gemm.h b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Bf16Bf16Gemm.h new file mode 100644 index 000000000000..18104f2a2bdd --- /dev/null +++ b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Bf16Bf16Gemm.h @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2025-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once +#include +#include +#include + +namespace tensorrt_llm::kernels::llama4_min_latency::llama4_bf16_bf16_gemm +{ + +void llama4_bf16_bf16_gemm_op(int num_tokens, void const* A, void const* B, void* C, cudaStream_t stream); + +} // namespace tensorrt_llm::kernels::llama4_min_latency::llama4_bf16_bf16_gemm diff --git a/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Bf16Gemm.cu b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Bf16Gemm.cu new file mode 100644 index 000000000000..1bd1016a9650 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Bf16Gemm.cu @@ -0,0 +1,189 @@ +/* + * Copyright (c) 2025-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Bf16Gemm.h" +#include "tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Bf16GemmAttnScalingPerBlockTemplate.cuh" +#include "tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Bf16GemmPerBlockTemplate.cuh" +#include "tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Bf16GemmPerWarpTemplate.cuh" +#include "tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Utils.cuh" +#include + +namespace tensorrt_llm::kernels::llama4_min_latency::llama4_fp8_bf16_gemm +{ + +DEFINE_GET_PER_BLOCK_FUNC_PTR(/*HIDDEN_IN=*/5120, /*ALIGNED=*/true); +DEFINE_GET_PER_BLOCK_ATTN_SCALING_FUNC_PTR(/*HIDDEN_IN=*/5120, /*ALIGNED=*/true, /*POS_IDS_INT64=*/false); +DEFINE_GET_PER_BLOCK_ATTN_SCALING_FUNC_PTR(/*HIDDEN_IN=*/5120, /*ALIGNED=*/true, /*POS_IDS_INT64=*/true); + +void llama4_fp8_bf16_gemm_launcher(void const* A, void const* B, void* C, void const* scaling_factor, int num_tokens, + int hidden_in, int hidden_out, cudaStream_t stream) +{ + void* args[] = {(void*) &A, (void*) &B, (void*) &C, (void*) &scaling_factor, (void*) &num_tokens, + (void*) &hidden_in, (void*) &hidden_out}; + if (num_tokens == 1) + { + // When num_tokens == 1, the best tiling size is tile_token == 1 and tile_out == 1. + dim3 const grid_size = dim3(div_up(hidden_out, 1), div_up(num_tokens, 1), 1); + void* kernel_func = get_per_block_func_ptr_aligned_true_5120_(1, 1); + launch_kernel_fdl(dim3(grid_size), dim3(BLOCK_SIZE), stream, kernel_func, args, 7); + } + else if (num_tokens == 2) + { + // When num_tokens == 2, the best tiling size is tile_token == 2 and tile_out == 1. + dim3 const grid_size = dim3(div_up(hidden_out, 1), div_up(num_tokens, 2), 1); + void* kernel_func = get_per_block_func_ptr_aligned_true_5120_(2, 1); + launch_kernel_fdl(dim3(grid_size), dim3(BLOCK_SIZE), stream, kernel_func, args, 7); + } + else if (num_tokens == 3) + { + // When num_tokens == 3, the best tiling size is tile_token == 1 and tile_out == 4. + dim3 const grid_size = dim3(div_up(hidden_out, 4), div_up(num_tokens, 1), 1); + void* kernel_func = get_per_block_func_ptr_aligned_true_5120_(1, 4); + launch_kernel_fdl(dim3(grid_size), dim3(BLOCK_SIZE), stream, kernel_func, args, 7); + } + else if (num_tokens == 4) + { + // When num_tokens == 4, the best tiling size is tile_token == 2 and tile_out == 2. + dim3 const grid_size = dim3(div_up(hidden_out, 2), div_up(num_tokens, 2), 1); + void* kernel_func = get_per_block_func_ptr_aligned_true_5120_(2, 2); + launch_kernel_fdl(dim3(grid_size), dim3(BLOCK_SIZE), stream, kernel_func, args, 7); + } + else if (num_tokens == 5) + { + // When num_tokens == 5, the best tiling size is tile_token == 1 and tile_out == 4. + dim3 const grid_size = dim3(div_up(hidden_out, 4), div_up(num_tokens, 1), 1); + void* kernel_func = get_per_block_func_ptr_aligned_true_5120_(1, 4); + launch_kernel_fdl(dim3(grid_size), dim3(BLOCK_SIZE), stream, kernel_func, args, 7); + } + else if (num_tokens == 6) + { + // When num_tokens == 6, the best tiling size is tile_token == 3 and tile_out == 4. + dim3 const grid_size = dim3(div_up(hidden_out, 4), div_up(num_tokens, 3), 1); + void* kernel_func = get_per_block_func_ptr_aligned_true_5120_(3, 4); + launch_kernel_fdl(dim3(grid_size), dim3(BLOCK_SIZE), stream, kernel_func, args, 7); + } + else if (num_tokens == 7) + { + // When num_tokens == 7, the best tiling size is tile_token == 1 and tile_out == 4. + dim3 const grid_size = dim3(div_up(hidden_out, 4), div_up(num_tokens, 1), 1); + void* kernel_func = get_per_block_func_ptr_aligned_true_5120_(1, 4); + launch_kernel_fdl(dim3(grid_size), dim3(BLOCK_SIZE), stream, kernel_func, args, 7); + } + else if (num_tokens == 8) + { + // When num_tokens == 8, the best tiling size is tile_token == 2 and tile_out == 4. + dim3 const grid_size = dim3(div_up(hidden_out, 4), div_up(num_tokens, 2), 1); + void* kernel_func = get_per_block_func_ptr_aligned_true_5120_(2, 4); + launch_kernel_fdl(dim3(grid_size), dim3(BLOCK_SIZE), stream, kernel_func, args, 7); + } + else + { + throw std::runtime_error("llama4QKVGemm: num_tokens > 8 are not supported."); + } +} + +void* get_kernel_func(int tile_token, int tile_out, bool pos_ids_int64) +{ + if (pos_ids_int64) + { + return get_per_block_attn_scaling_func_ptr_aligned_true_pos_int64_true_5120_(tile_token, tile_out); + } + return get_per_block_attn_scaling_func_ptr_aligned_true_pos_int64_false_5120_(tile_token, tile_out); +} + +void llama4_fp8_bf16_gemm_attn_scaling_launcher(void const* A, void const* B, void* C, void const* scaling_factor, + void const* pos_ids, bool pos_ids_int64, float floor_scale, float attn_scale, int num_tokens, int hidden_in, + int hidden_out, int q_hidden_out, cudaStream_t stream) +{ + void* args[] = {(void*) &A, (void*) &B, (void*) &C, (void*) &scaling_factor, (void*) &pos_ids, (void*) &floor_scale, + (void*) &attn_scale, (void*) &num_tokens, (void*) &hidden_in, (void*) &hidden_out, (void*) &q_hidden_out}; + if (num_tokens == 1) + { + // When num_tokens == 1, the best tiling size is tile_token == 1 and tile_out == 1. + dim3 const grid_size = dim3(div_up(hidden_out, 1), div_up(num_tokens, 1), 1); + void* kernel_func = get_kernel_func(1, 1, pos_ids_int64); + launch_kernel_fdl(dim3(grid_size), dim3(BLOCK_SIZE), stream, kernel_func, args, 11); + } + else if (num_tokens == 2) + { + // When num_tokens == 2, the best tiling size is tile_token == 2 and tile_out == 2. + dim3 const grid_size = dim3(div_up(hidden_out, 2), div_up(num_tokens, 2), 1); + void* kernel_func = get_kernel_func(2, 2, pos_ids_int64); + launch_kernel_fdl(dim3(grid_size), dim3(BLOCK_SIZE), stream, kernel_func, args, 11); + } + else if (num_tokens == 3) + { + // When num_tokens == 3, the best tiling size is tile_token == 1 and tile_out == 4. + dim3 const grid_size = dim3(div_up(hidden_out, 4), div_up(num_tokens, 1), 1); + void* kernel_func = get_kernel_func(1, 4, pos_ids_int64); + launch_kernel_fdl(dim3(grid_size), dim3(BLOCK_SIZE), stream, kernel_func, args, 11); + } + else if (num_tokens == 4) + { + // When num_tokens == 4, the best tiling size is tile_token == 2 and tile_out == 2. + dim3 const grid_size = dim3(div_up(hidden_out, 2), div_up(num_tokens, 2), 1); + void* kernel_func = get_kernel_func(2, 2, pos_ids_int64); + launch_kernel_fdl(dim3(grid_size), dim3(BLOCK_SIZE), stream, kernel_func, args, 11); + } + else if (num_tokens == 5) + { + // When num_tokens == 5, the best tiling size is tile_token == 1 and tile_out == 4. + dim3 const grid_size = dim3(div_up(hidden_out, 4), div_up(num_tokens, 1), 1); + void* kernel_func = get_kernel_func(1, 4, pos_ids_int64); + launch_kernel_fdl(dim3(grid_size), dim3(BLOCK_SIZE), stream, kernel_func, args, 11); + } + else if (num_tokens == 6) + { + // When num_tokens == 6, the best tiling size is tile_token == 2 and tile_out == 4. + dim3 const grid_size = dim3(div_up(hidden_out, 4), div_up(num_tokens, 2), 1); + void* kernel_func = get_kernel_func(2, 4, pos_ids_int64); + launch_kernel_fdl(dim3(grid_size), dim3(BLOCK_SIZE), stream, kernel_func, args, 11); + } + else if (num_tokens == 7) + { + // When num_tokens == 7, the best tiling size is tile_token == 1 and tile_out == 4. + dim3 const grid_size = dim3(div_up(hidden_out, 4), div_up(num_tokens, 1), 1); + void* kernel_func = get_kernel_func(1, 4, pos_ids_int64); + launch_kernel_fdl(dim3(grid_size), dim3(BLOCK_SIZE), stream, kernel_func, args, 11); + } + else if (num_tokens == 8) + { + // When num_tokens == 8, the best tiling size is tile_token == 2 and tile_out == 4. + dim3 const grid_size = dim3(div_up(hidden_out, 4), div_up(num_tokens, 2), 1); + void* kernel_func = get_kernel_func(2, 4, pos_ids_int64); + launch_kernel_fdl(dim3(grid_size), dim3(BLOCK_SIZE), stream, kernel_func, args, 11); + } + else + { + throw std::runtime_error("llama4_fp8_bf16_gemm: num_tokens > 8 are not supported."); + } +} + +void llama4_fp8_bf16_gemm_op(void const* A, void const* B, void* C, void const* scaling_factor, void const* pos_ids, + bool pos_ids_int64, int num_tokens, int hidden_in, int hidden_out, cudaStream_t stream) +{ + if (pos_ids != nullptr) + { + llama4_fp8_bf16_gemm_attn_scaling_launcher(A, B, C, scaling_factor, pos_ids, pos_ids_int64, FLOOR_SCALE, + ATTN_SCALE, num_tokens, hidden_in, hidden_out, Q_HIDDEN_OUT, stream); + } + else + { + llama4_fp8_bf16_gemm_launcher(A, B, C, scaling_factor, num_tokens, hidden_in, hidden_out, stream); + } +} + +} // namespace tensorrt_llm::kernels::llama4_min_latency::llama4_fp8_bf16_gemm diff --git a/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Bf16Gemm.h b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Bf16Gemm.h new file mode 100644 index 000000000000..709d56d3bf54 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Bf16Gemm.h @@ -0,0 +1,30 @@ +/* + * Copyright (c) 2025-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include + +namespace tensorrt_llm::kernels::llama4_min_latency::llama4_fp8_bf16_gemm +{ + +void llama4_fp8_bf16_gemm_op(void const* A, void const* B, void* C, void const* scaling_factor, void const* pos_ids, + bool pos_ids_int64, int num_tokens, int hidden_in, int hidden_out, cudaStream_t stream); + +} // namespace tensorrt_llm::kernels::llama4_min_latency::llama4_fp8_bf16_gemm diff --git a/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Bf16GemmAttnScalingPerBlockTemplate.cuh b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Bf16GemmAttnScalingPerBlockTemplate.cuh new file mode 100644 index 000000000000..b330908d090a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Bf16GemmAttnScalingPerBlockTemplate.cuh @@ -0,0 +1,360 @@ +/* + * Copyright (c) 2025-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Utils.cuh" + +#include +#include +#include + +namespace tensorrt_llm::kernels::llama4_min_latency::llama4_fp8_bf16_gemm +{ + +// Grid size is num_tokens / TILE_TOKEN * hidden_out / TILE_OUT. +// Each block processes TILE_TOKEN tokens and TILE_OUT rows. +// within each block, it steps through hidden_in in steps of BLOCK_SIZE * VEC_SIZE. +template +__launch_bounds__(BLOCK_SIZE) __global__ void llama4_fp8_bf16_gemm_attn_scaling_per_block_kernel( + __nv_fp8_e4m3 const* __restrict__ A, // Input tensor [num_tokens][hidden_in] + __nv_fp8_e4m3 const* __restrict__ B, // Input tensor [hidden_out][hidden_in] + __nv_bfloat16* __restrict__ C, // Output tensor [num_tokens][hidden_out] + float* __restrict__ scaling_factor, POS_IDS_TYPE const* __restrict__ pos_ids, float const floor_scale, + float const attn_scale, int num_tokens, int hidden_in, int hidden_out, int q_hidden_out) +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + // Shared memory for block reduction + __shared__ float reduce_buffer[TILE_TOKEN][TILE_OUT][BLOCK_SIZE]; + + // Each thread accumulates its partial sum + float2 thread_sum[TILE_TOKEN][TILE_OUT]; +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + thread_sum[tile_token_idx][tile_out_idx].x = 0.0f; + thread_sum[tile_token_idx][tile_out_idx].y = 0.0f; + } + } + + int const token_idx = blockIdx.y; + int const row_idx = blockIdx.x; + int const tid = threadIdx.x; + + // Calculate attn scaling factor. + float attn_scaling_factors[TILE_TOKEN]; +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { + int current_token = token_idx * TILE_TOKEN + tile_token_idx; + float const floor = floorf((static_cast(pos_ids[current_token]) + 1.0f) / floor_scale); + attn_scaling_factors[tile_token_idx] = (__logf(floor + 1.0f) * attn_scale) + 1.0f; + } + + int chunk = 0; +#if ENABLE_ACQBULK && ENABLE_PREFETCH +#pragma unroll 9 + for (; chunk < ((HIDDEN_IN > 0 ? HIDDEN_IN : hidden_in) / BLOCK_SIZE / VEC_SIZE - 1); chunk++) + { +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + int current_row = row_idx * TILE_OUT + tile_out_idx; + int base_idx = chunk * BLOCK_SIZE + tid; + asm volatile("prefetch.global.L2 [%0];" ::"l"(&B[current_row * hidden_in / VEC_SIZE + base_idx]) + : "memory"); + } + } + { +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + int current_row = row_idx * TILE_OUT + tile_out_idx; + int base_idx = chunk * BLOCK_SIZE + tid; + if (ALIGNED || base_idx * VEC_SIZE < hidden_in) + { + asm volatile("prefetch.global.L2 [%0];" ::"l"(&B[current_row * hidden_in / VEC_SIZE + base_idx]) + : "memory"); + } + } + } +#endif + +#if ENABLE_ACQBULK + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + + // Processing 8 elements each + chunk = 0; +#pragma unroll 9 + for (; chunk < ((HIDDEN_IN > 0 ? HIDDEN_IN : hidden_in) / BLOCK_SIZE / VEC_SIZE - 1); chunk++) + { + int base_idx = chunk * BLOCK_SIZE + tid; + // Load values from tensor B + aligned_fp8x8 b_vec[TILE_OUT]; +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + int current_row = row_idx * TILE_OUT + tile_out_idx; + b_vec[tile_out_idx] + = reinterpret_cast(B)[current_row * hidden_in / VEC_SIZE + base_idx]; + } + + // Load values from tensor A + aligned_fp8x8 a_vec[TILE_TOKEN]; +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { + int current_token = token_idx * TILE_TOKEN + tile_token_idx; + a_vec[tile_token_idx] + = reinterpret_cast(A)[current_token * hidden_in / VEC_SIZE + base_idx]; + } + + // Compute partial sum +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { + aligned_fp8x8 a_vec_current = a_vec[tile_token_idx]; + aligned_fp8x8 b_vec_current = b_vec[tile_out_idx]; +#pragma unroll + for (int i = 0; i < VEC_SIZE / 4; i++) + { + float4 a_val = float4(a_vec_current.data[i]); + float4 b_val = float4(b_vec_current.data[i]); + + thread_sum[tile_token_idx][tile_out_idx] = ffma2(make_float2(a_val.x, a_val.y), + make_float2(b_val.x, b_val.y), thread_sum[tile_token_idx][tile_out_idx]); + thread_sum[tile_token_idx][tile_out_idx] = ffma2(make_float2(a_val.z, a_val.w), + make_float2(b_val.z, b_val.w), thread_sum[tile_token_idx][tile_out_idx]); + } + } + } + } + + // The last chunk may be partial. + { + int base_idx = chunk * BLOCK_SIZE + tid; + if (ALIGNED || base_idx * VEC_SIZE < hidden_in) + { + // Load values from tensor B + aligned_fp8x8 b_vec[TILE_OUT]; +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + int current_row = row_idx * TILE_OUT + tile_out_idx; + b_vec[tile_out_idx] + = reinterpret_cast(B)[current_row * hidden_in / VEC_SIZE + base_idx]; + } + + // Load values from tensor A + aligned_fp8x8 a_vec[TILE_TOKEN]; +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { + int current_token = token_idx * TILE_TOKEN + tile_token_idx; + a_vec[tile_token_idx] + = reinterpret_cast(A)[current_token * hidden_in / VEC_SIZE + base_idx]; + } + + // Compute partial sum +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { + aligned_fp8x8 a_vec_current = a_vec[tile_token_idx]; + aligned_fp8x8 b_vec_current = b_vec[tile_out_idx]; +#pragma unroll + for (int i = 0; i < VEC_SIZE / 4; i++) + { + float4 a_val = float4(a_vec_current.data[i]); + float4 b_val = float4(b_vec_current.data[i]); + + thread_sum[tile_token_idx][tile_out_idx] = ffma2(make_float2(a_val.x, a_val.y), + make_float2(b_val.x, b_val.y), thread_sum[tile_token_idx][tile_out_idx]); + thread_sum[tile_token_idx][tile_out_idx] = ffma2(make_float2(a_val.z, a_val.w), + make_float2(b_val.z, b_val.w), thread_sum[tile_token_idx][tile_out_idx]); + } + } + } + } + } + + // Reduce partial sums using warp-level reduction. +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { + float warp_sum = thread_sum[tile_token_idx][tile_out_idx].x + thread_sum[tile_token_idx][tile_out_idx].y; +#pragma unroll + for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) + { + warp_sum += __shfl_down_sync(0xffffffff, warp_sum, offset); + } + // First thread in each warp writes to shared memory + if (tid % WARP_SIZE == 0) + { + reduce_buffer[tile_token_idx][tile_out_idx][tid / WARP_SIZE] = warp_sum; + } + } + } + + __syncthreads(); + + if (tid == 0) + { +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + float block_sum = 0.0f; +#pragma unroll + for (int i = 0; i < BLOCK_SIZE / WARP_SIZE; i++) + { + block_sum += reduce_buffer[tile_token_idx][tile_out_idx][i]; + } + int current_row = row_idx * TILE_OUT + tile_out_idx; + int current_token = token_idx * TILE_TOKEN + tile_token_idx; + float attn_scaling_factor = current_row < q_hidden_out ? attn_scaling_factors[tile_token_idx] : 1.0f; + C[current_token * hidden_out + current_row] + = __float2bfloat16(block_sum * attn_scaling_factor * scaling_factor[0]); + } + } + } + +#if ENABLE_PREEXIT + asm volatile("griddepcontrol.launch_dependents;"); +#endif +#endif +} + +#define DISPATCH_PER_BLOCK_FC_FP8_BF16_ATTN_SCALING_TILE_TOKEN( \ + HIDDEN_IN, TILE_TOKEN, TILE_OUT, ALIGNED, POS_IDS_INT64) \ + do \ + { \ + if (TILE_TOKEN == 1) \ + { \ + if (POS_IDS_INT64) \ + { \ + return reinterpret_cast( \ + llama4_fp8_bf16_gemm_attn_scaling_per_block_kernel); \ + } \ + else \ + { \ + return reinterpret_cast( \ + llama4_fp8_bf16_gemm_attn_scaling_per_block_kernel); \ + } \ + } \ + if (TILE_TOKEN == 2) \ + { \ + if (POS_IDS_INT64) \ + { \ + return reinterpret_cast( \ + llama4_fp8_bf16_gemm_attn_scaling_per_block_kernel); \ + } \ + else \ + { \ + return reinterpret_cast( \ + llama4_fp8_bf16_gemm_attn_scaling_per_block_kernel); \ + } \ + } \ + if (TILE_TOKEN == 3) \ + { \ + if (POS_IDS_INT64) \ + { \ + return reinterpret_cast( \ + llama4_fp8_bf16_gemm_attn_scaling_per_block_kernel); \ + } \ + else \ + { \ + return reinterpret_cast( \ + llama4_fp8_bf16_gemm_attn_scaling_per_block_kernel); \ + } \ + } \ + if (TILE_TOKEN == 4) \ + { \ + if (POS_IDS_INT64) \ + { \ + return reinterpret_cast( \ + llama4_fp8_bf16_gemm_attn_scaling_per_block_kernel); \ + } \ + else \ + { \ + return reinterpret_cast( \ + llama4_fp8_bf16_gemm_attn_scaling_per_block_kernel); \ + } \ + } \ + if (TILE_TOKEN == 8) \ + { \ + if (POS_IDS_INT64) \ + { \ + return reinterpret_cast( \ + llama4_fp8_bf16_gemm_attn_scaling_per_block_kernel); \ + } \ + else \ + { \ + return reinterpret_cast( \ + llama4_fp8_bf16_gemm_attn_scaling_per_block_kernel); \ + } \ + } \ + throw std::invalid_argument("Invalid tile token"); \ + } while (0) + +#define DISPATCH_PER_BLOCK_FC_FP8_BF16_ATTN_SCALING_TILE_OUT(HIDDEN_IN, TILE_TOKEN, TILE_OUT, ALIGNED, POS_IDS_INT64) \ + do \ + { \ + if (TILE_OUT == 1) \ + { \ + DISPATCH_PER_BLOCK_FC_FP8_BF16_ATTN_SCALING_TILE_TOKEN(HIDDEN_IN, TILE_TOKEN, 1, ALIGNED, POS_IDS_INT64); \ + } \ + if (TILE_OUT == 2) \ + { \ + DISPATCH_PER_BLOCK_FC_FP8_BF16_ATTN_SCALING_TILE_TOKEN(HIDDEN_IN, TILE_TOKEN, 2, ALIGNED, POS_IDS_INT64); \ + } \ + if (TILE_OUT == 3) \ + { \ + DISPATCH_PER_BLOCK_FC_FP8_BF16_ATTN_SCALING_TILE_TOKEN(HIDDEN_IN, TILE_TOKEN, 3, ALIGNED, POS_IDS_INT64); \ + } \ + if (TILE_OUT == 4) \ + { \ + DISPATCH_PER_BLOCK_FC_FP8_BF16_ATTN_SCALING_TILE_TOKEN(HIDDEN_IN, TILE_TOKEN, 4, ALIGNED, POS_IDS_INT64); \ + } \ + if (TILE_OUT == 8) \ + { \ + DISPATCH_PER_BLOCK_FC_FP8_BF16_ATTN_SCALING_TILE_TOKEN(HIDDEN_IN, TILE_TOKEN, 8, ALIGNED, POS_IDS_INT64); \ + } \ + throw std::invalid_argument("Invalid tile token"); \ + } while (0) + +#define DEFINE_GET_PER_BLOCK_ATTN_SCALING_FUNC_PTR(HIDDEN_IN, ALIGNED, POS_IDS_INT64) \ + void* get_per_block_attn_scaling_func_ptr_aligned_##ALIGNED##_pos_int64_##POS_IDS_INT64##_##HIDDEN_IN##_( \ + int tile_token, int tile_out) \ + { \ + DISPATCH_PER_BLOCK_FC_FP8_BF16_ATTN_SCALING_TILE_OUT(HIDDEN_IN, tile_token, tile_out, ALIGNED, POS_IDS_INT64); \ + } + +} // namespace tensorrt_llm::kernels::llama4_min_latency::llama4_fp8_bf16_gemm diff --git a/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Bf16GemmPerBlockTemplate.cuh b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Bf16GemmPerBlockTemplate.cuh new file mode 100644 index 000000000000..eac5a41399af --- /dev/null +++ b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Bf16GemmPerBlockTemplate.cuh @@ -0,0 +1,300 @@ +/* + * Copyright (c) 2025-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Utils.cuh" + +#include +#include +#include + +namespace tensorrt_llm::kernels::llama4_min_latency::llama4_fp8_bf16_gemm +{ + +// Grid size is num_tokens / TILE_TOKEN * hidden_out / TILE_OUT. +// Each block processes TILE_TOKEN tokens and TILE_OUT rows. +// Within each block, it steps through hidden_in in steps of BLOCK_SIZE * VEC_SIZE. +template +__launch_bounds__(BLOCK_SIZE) __global__ void llama4_fp8_bf16_gemm_per_block_kernel( + __nv_fp8_e4m3 const* __restrict__ A, // Input tensor [num_tokens][hidden_in] + __nv_fp8_e4m3 const* __restrict__ B, // Input tensor [hidden_out][hidden_in] + __nv_bfloat16* __restrict__ C, // Output tensor [num_tokens][hidden_out] + float* __restrict__ scaling_factor, int num_tokens, int hidden_in, int hidden_out) +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + // Shared memory for block reduction + __shared__ float reduce_buffer[TILE_TOKEN][TILE_OUT][BLOCK_SIZE]; + + // Each thread accumulates its partial sum + float2 thread_sum[TILE_TOKEN][TILE_OUT]; +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + thread_sum[tile_token_idx][tile_out_idx].x = 0.0f; + thread_sum[tile_token_idx][tile_out_idx].y = 0.0f; + } + } + + int const token_idx = blockIdx.y; + int const row_idx = blockIdx.x; + int const tid = threadIdx.x; + + int chunk = 0; +#if ENABLE_ACQBULK && ENABLE_PREFETCH +#pragma unroll 9 + for (; chunk < ((HIDDEN_IN > 0 ? HIDDEN_IN : hidden_in) / BLOCK_SIZE / VEC_SIZE - 1); chunk++) + { +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + int current_row = row_idx * TILE_OUT + tile_out_idx; + int base_idx = chunk * BLOCK_SIZE + tid; + asm volatile("prefetch.global.L2 [%0];" ::"l"(&B[current_row * hidden_in / VEC_SIZE + base_idx]) + : "memory"); + } + } + { +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + int current_row = row_idx * TILE_OUT + tile_out_idx; + int base_idx = chunk * BLOCK_SIZE + tid; + if (ALIGNED || base_idx * VEC_SIZE < hidden_in) + { + asm volatile("prefetch.global.L2 [%0];" ::"l"(&B[current_row * hidden_in / VEC_SIZE + base_idx]) + : "memory"); + } + } + } +#endif + +#if ENABLE_ACQBULK + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + + // Processing 8 elements each + chunk = 0; +#pragma unroll 9 + for (; chunk < ((HIDDEN_IN > 0 ? HIDDEN_IN : hidden_in) / BLOCK_SIZE / VEC_SIZE - 1); chunk++) + { + int base_idx = chunk * BLOCK_SIZE + tid; + // Load values from tensor B + aligned_fp8x8 b_vec[TILE_OUT]; +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + int current_row = row_idx * TILE_OUT + tile_out_idx; + b_vec[tile_out_idx] + = reinterpret_cast(B)[current_row * hidden_in / VEC_SIZE + base_idx]; + } + + // Load values from tensor A + aligned_fp8x8 a_vec[TILE_TOKEN]; +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { + int current_token = token_idx * TILE_TOKEN + tile_token_idx; + a_vec[tile_token_idx] + = reinterpret_cast(A)[current_token * hidden_in / VEC_SIZE + base_idx]; + } + + // Compute partial sum +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { + aligned_fp8x8 a_vec_current = a_vec[tile_token_idx]; + aligned_fp8x8 b_vec_current = b_vec[tile_out_idx]; +#pragma unroll + for (int i = 0; i < VEC_SIZE / 4; i++) + { + float4 a_val = float4(a_vec_current.data[i]); + float4 b_val = float4(b_vec_current.data[i]); + + thread_sum[tile_token_idx][tile_out_idx] = ffma2(make_float2(a_val.x, a_val.y), + make_float2(b_val.x, b_val.y), thread_sum[tile_token_idx][tile_out_idx]); + thread_sum[tile_token_idx][tile_out_idx] = ffma2(make_float2(a_val.z, a_val.w), + make_float2(b_val.z, b_val.w), thread_sum[tile_token_idx][tile_out_idx]); + } + } + } + } + + // The last chunk may be partial. + { + int base_idx = chunk * BLOCK_SIZE + tid; + if (ALIGNED || base_idx * VEC_SIZE < hidden_in) + { + // Load values from tensor B + aligned_fp8x8 b_vec[TILE_OUT]; +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + int current_row = row_idx * TILE_OUT + tile_out_idx; + b_vec[tile_out_idx] + = reinterpret_cast(B)[current_row * hidden_in / VEC_SIZE + base_idx]; + } + + // Load values from tensor A + aligned_fp8x8 a_vec[TILE_TOKEN]; +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { + int current_token = token_idx * TILE_TOKEN + tile_token_idx; + a_vec[tile_token_idx] + = reinterpret_cast(A)[current_token * hidden_in / VEC_SIZE + base_idx]; + } + + // Compute partial sum +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { + aligned_fp8x8 a_vec_current = a_vec[tile_token_idx]; + aligned_fp8x8 b_vec_current = b_vec[tile_out_idx]; +#pragma unroll + for (int i = 0; i < VEC_SIZE / 4; i++) + { + float4 a_val = float4(a_vec_current.data[i]); + float4 b_val = float4(b_vec_current.data[i]); + + thread_sum[tile_token_idx][tile_out_idx] = ffma2(make_float2(a_val.x, a_val.y), + make_float2(b_val.x, b_val.y), thread_sum[tile_token_idx][tile_out_idx]); + thread_sum[tile_token_idx][tile_out_idx] = ffma2(make_float2(a_val.z, a_val.w), + make_float2(b_val.z, b_val.w), thread_sum[tile_token_idx][tile_out_idx]); + } + } + } + } + } + + // Reduce partial sums using warp-level reduction. +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { + float warp_sum = thread_sum[tile_token_idx][tile_out_idx].x + thread_sum[tile_token_idx][tile_out_idx].y; +#pragma unroll + for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) + { + warp_sum += __shfl_down_sync(0xffffffff, warp_sum, offset); + } + // First thread in each warp writes to shared memory + if (tid % WARP_SIZE == 0) + { + reduce_buffer[tile_token_idx][tile_out_idx][tid / WARP_SIZE] = warp_sum; + } + } + } + + __syncthreads(); + + if (tid == 0) + { +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + float block_sum = 0.0f; +#pragma unroll + for (int i = 0; i < BLOCK_SIZE / WARP_SIZE; i++) + { + block_sum += reduce_buffer[tile_token_idx][tile_out_idx][i]; + } + int current_row = row_idx * TILE_OUT + tile_out_idx; + int current_token = token_idx * TILE_TOKEN + tile_token_idx; + C[current_token * hidden_out + current_row] = __float2bfloat16(block_sum * scaling_factor[0]); + } + } + } + +#if ENABLE_PREEXIT + asm volatile("griddepcontrol.launch_dependents;"); +#endif +#endif +} + +#define DISPATCH_PER_BLOCK_FC_FP8_BF16_TILE_TOKEN(HIDDEN_IN, TILE_TOKEN, TILE_OUT, ALIGNED) \ + do \ + { \ + if (TILE_TOKEN == 1) \ + { \ + return reinterpret_cast(llama4_fp8_bf16_gemm_per_block_kernel); \ + } \ + if (TILE_TOKEN == 2) \ + { \ + return reinterpret_cast(llama4_fp8_bf16_gemm_per_block_kernel); \ + } \ + if (TILE_TOKEN == 3) \ + { \ + return reinterpret_cast(llama4_fp8_bf16_gemm_per_block_kernel); \ + } \ + if (TILE_TOKEN == 4) \ + { \ + return reinterpret_cast(llama4_fp8_bf16_gemm_per_block_kernel); \ + } \ + if (TILE_TOKEN == 8) \ + { \ + return reinterpret_cast(llama4_fp8_bf16_gemm_per_block_kernel); \ + } \ + throw std::invalid_argument("Invalid tile token"); \ + } while (0) + +#define DISPATCH_PER_BLOCK_FC_FP8_BF16_TILE_OUT(HIDDEN_IN, TILE_TOKEN, TILE_OUT, ALIGNED) \ + do \ + { \ + if (TILE_OUT == 1) \ + { \ + DISPATCH_PER_BLOCK_FC_FP8_BF16_TILE_TOKEN(HIDDEN_IN, TILE_TOKEN, 1, ALIGNED); \ + } \ + if (TILE_OUT == 2) \ + { \ + DISPATCH_PER_BLOCK_FC_FP8_BF16_TILE_TOKEN(HIDDEN_IN, TILE_TOKEN, 2, ALIGNED); \ + } \ + if (TILE_OUT == 3) \ + { \ + DISPATCH_PER_BLOCK_FC_FP8_BF16_TILE_TOKEN(HIDDEN_IN, TILE_TOKEN, 3, ALIGNED); \ + } \ + if (TILE_OUT == 4) \ + { \ + DISPATCH_PER_BLOCK_FC_FP8_BF16_TILE_TOKEN(HIDDEN_IN, TILE_TOKEN, 4, ALIGNED); \ + } \ + if (TILE_OUT == 8) \ + { \ + DISPATCH_PER_BLOCK_FC_FP8_BF16_TILE_TOKEN(HIDDEN_IN, TILE_TOKEN, 8, ALIGNED); \ + } \ + throw std::invalid_argument("Invalid tile token"); \ + } while (0) + +#define DEFINE_GET_PER_BLOCK_FUNC_PTR(HIDDEN_IN, ALIGNED) \ + void* get_per_block_func_ptr_aligned_##ALIGNED##_##HIDDEN_IN##_(int tile_token, int tile_out) \ + { \ + DISPATCH_PER_BLOCK_FC_FP8_BF16_TILE_OUT(HIDDEN_IN, tile_token, tile_out, ALIGNED); \ + } + +} // namespace tensorrt_llm::kernels::llama4_min_latency::llama4_fp8_bf16_gemm diff --git a/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Bf16GemmPerWarpTemplate.cuh b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Bf16GemmPerWarpTemplate.cuh new file mode 100644 index 000000000000..592995dc4abb --- /dev/null +++ b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Bf16GemmPerWarpTemplate.cuh @@ -0,0 +1,326 @@ +/* + * Copyright (c) 2025-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Utils.cuh" + +#include +#include +#include + +namespace tensorrt_llm::kernels::llama4_min_latency::llama4_fp8_bf16_gemm +{ + +// Grid size is num_tokens / TILE_TOKEN * hidden_out / TILE_OUT / WARP_PER_BLOCK. +// Each warp processes TILE_TOKEN tokens and TILE_OUT rows. +// Within each warp, it loops through hidden_in in steps of WARP_SIZE * VEC_SIZE. +template +__launch_bounds__(BLOCK_SIZE) __global__ void llama4_fp8_bf16_gemm_per_warp_kernel( + __nv_fp8_e4m3 const* __restrict__ A, // Input tensor [num_tokens][hidden_in] + __nv_fp8_e4m3 const* __restrict__ B, // Input tensor [hidden_out][hidden_in] + __nv_bfloat16* __restrict__ C, // Output tensor [num_tokens][hidden_out] + float* __restrict__ scaling_factor, int num_tokens, int hidden_in, int hidden_out) +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + // Shared memory for tensor A values. + __shared__ aligned_fp8x8 shared_a[TILE_TOKEN][BLOCK_SIZE]; + + // Each thread accumulates its partial sum + float2 thread_sum[TILE_TOKEN][TILE_OUT]; +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + thread_sum[tile_token_idx][tile_out_idx].x = 0.0f; + thread_sum[tile_token_idx][tile_out_idx].y = 0.0f; + } + } + + int const token_idx = blockIdx.y; // Token ID in the unit of TILE_TOKEN + int const warp_idx = threadIdx.x / WARP_SIZE; // Warp ID within the block + int const row_idx = blockIdx.x * WARP_PER_BLOCK + warp_idx; // row ID in the unit of TILE_OUT + int const tid = threadIdx.x; + int const lane_idx = tid % WARP_SIZE; + + int chunk = 0; +#if ENABLE_ACQBULK && ENABLE_PREFETCH +#pragma unroll + for (; chunk < ((HIDDEN_IN > 0 ? HIDDEN_IN : hidden_in) / WARP_SIZE / VEC_SIZE - 1); chunk++) + { +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + int current_row = row_idx * TILE_OUT + tile_out_idx; + int base_idx = chunk * WARP_SIZE + lane_idx; + asm volatile("prefetch.global.L2 [%0];" ::"l"(&B[current_row * HIDDEN_IN / VEC_SIZE + base_idx]) + : "memory"); + } + } + { +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + int current_row = row_idx * TILE_OUT + tile_out_idx; + int base_idx = chunk * WARP_SIZE + lane_idx; + if (ALIGNED || base_idx * VEC_SIZE < hidden_in) + { + asm volatile("prefetch.global.L2 [%0];" ::"l"(&B[current_row * HIDDEN_IN / VEC_SIZE + base_idx]) + : "memory"); + } + } + } +#endif + +#if ENABLE_ACQBULK + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + + // Processing 8 elements each + chunk = 0; +#pragma unroll + for (; chunk < ((HIDDEN_IN > 0 ? HIDDEN_IN : hidden_in) / BLOCK_SIZE / VEC_SIZE) - 1; chunk++) + { + int base_idx_a = chunk * BLOCK_SIZE + tid; + int base_idx_b = chunk * BLOCK_SIZE + lane_idx; + + // Load values from tensor A into shared memory +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { + int current_token = token_idx * TILE_TOKEN + tile_token_idx; + shared_a[tile_token_idx][tid] + = reinterpret_cast(A)[current_token * hidden_in / VEC_SIZE + base_idx_a]; + } + + // Load values from tensor B + aligned_fp8x8 b_vec[WARP_PER_BLOCK][TILE_OUT]; +#pragma unroll + for (int tile_warp_idx = 0; tile_warp_idx < WARP_PER_BLOCK; tile_warp_idx++) + { +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + int current_col = base_idx_b + tile_warp_idx * WARP_SIZE; + int current_row = row_idx * TILE_OUT + tile_out_idx; + b_vec[tile_warp_idx][tile_out_idx] + = reinterpret_cast(B)[current_row * hidden_in / VEC_SIZE + current_col]; + } + } + + __syncthreads(); + + // Compute partial sum +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { +#pragma unroll + for (int tile_warp_idx = 0; tile_warp_idx < WARP_PER_BLOCK; tile_warp_idx++) + { + aligned_fp8x8 a_vec_current = shared_a[tile_token_idx][tile_warp_idx * WARP_SIZE + lane_idx]; + aligned_fp8x8 b_vec_current = b_vec[tile_warp_idx][tile_out_idx]; +#pragma unroll + for (int i = 0; i < VEC_SIZE / 4; i++) + { + float4 a_val = float4(a_vec_current.data[i]); + float4 b_val = float4(b_vec_current.data[i]); + + thread_sum[tile_token_idx][tile_out_idx] = ffma2(make_float2(a_val.x, a_val.y), + make_float2(b_val.x, b_val.y), thread_sum[tile_token_idx][tile_out_idx]); + thread_sum[tile_token_idx][tile_out_idx] = ffma2(make_float2(a_val.z, a_val.w), + make_float2(b_val.z, b_val.w), thread_sum[tile_token_idx][tile_out_idx]); + } + } + } + } + + __syncthreads(); + } + { + int base_idx_a = chunk * BLOCK_SIZE + tid; + int base_idx_b = chunk * BLOCK_SIZE + lane_idx; + + if (ALIGNED || base_idx_a * VEC_SIZE < hidden_in) + { + // Load values from tensor A into shared memory +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { + int current_token = token_idx * TILE_TOKEN + tile_token_idx; + shared_a[tile_token_idx][tid] + = reinterpret_cast(A)[current_token * hidden_in / VEC_SIZE + base_idx_a]; + } + } + + // Load values from tensor B + aligned_fp8x8 b_vec[WARP_PER_BLOCK][TILE_OUT]; +#pragma unroll + for (int tile_warp_idx = 0; tile_warp_idx < WARP_PER_BLOCK; tile_warp_idx++) + { + int current_col = base_idx_b + tile_warp_idx * WARP_SIZE; + if (ALIGNED || current_col < hidden_in) + { +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + int current_row = row_idx * TILE_OUT + tile_out_idx; + b_vec[tile_warp_idx][tile_out_idx] + = reinterpret_cast(B)[current_row * hidden_in / VEC_SIZE + current_col]; + } + } + } + + __syncthreads(); + + // Compute partial sum +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { +#pragma unroll + for (int tile_warp_idx = 0; tile_warp_idx < WARP_PER_BLOCK; tile_warp_idx++) + { + int current_col = base_idx_b + tile_warp_idx * WARP_SIZE; + if (ALIGNED || current_col < hidden_in) + { + aligned_fp8x8 a_vec_current = shared_a[tile_token_idx][tile_warp_idx * WARP_SIZE + lane_idx]; + aligned_fp8x8 b_vec_current = b_vec[tile_warp_idx][tile_out_idx]; +#pragma unroll + for (int i = 0; i < VEC_SIZE / 4; i++) + { + float4 a_val = float4(a_vec_current.data[i]); + float4 b_val = float4(b_vec_current.data[i]); + + thread_sum[tile_token_idx][tile_out_idx] = ffma2(make_float2(a_val.x, a_val.y), + make_float2(b_val.x, b_val.y), thread_sum[tile_token_idx][tile_out_idx]); + thread_sum[tile_token_idx][tile_out_idx] = ffma2(make_float2(a_val.z, a_val.w), + make_float2(b_val.z, b_val.w), thread_sum[tile_token_idx][tile_out_idx]); + } + } + } + } + } + + __syncthreads(); + } + + // Reduce partial sums using warp-level reduction. +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { + float warp_sum = thread_sum[tile_token_idx][tile_out_idx].x + thread_sum[tile_token_idx][tile_out_idx].y; +#pragma unroll + for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) + { + warp_sum += __shfl_down_sync(0xffffffff, warp_sum, offset); + } + + // Broadcast the result to all threads in the warp + thread_sum[tile_token_idx][tile_out_idx].x = __shfl_sync(0xffffffff, warp_sum, 0); + } + } + + if (lane_idx == 0) + { +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + int current_token = token_idx * TILE_TOKEN + tile_token_idx; + int current_row = row_idx * TILE_OUT + tile_out_idx; + C[current_token * hidden_out + current_row] + = __float2bfloat16(thread_sum[tile_token_idx][tile_out_idx].x * scaling_factor[0]); + } + } + } + +#if ENABLE_PREEXIT + asm volatile("griddepcontrol.launch_dependents;"); +#endif +#endif +} + +#define DISPATCH_PER_WARP_FC_FP8_BF16_TILE_TOKEN(HIDDEN_IN, TILE_TOKEN, TILE_OUT, ALIGNED) \ + do \ + { \ + if (TILE_TOKEN == 1) \ + { \ + return reinterpret_cast(llama4_fp8_bf16_gemm_per_warp_kernel); \ + } \ + if (TILE_TOKEN == 2) \ + { \ + return reinterpret_cast(llama4_fp8_bf16_gemm_per_warp_kernel); \ + } \ + if (TILE_TOKEN == 3) \ + { \ + return reinterpret_cast(llama4_fp8_bf16_gemm_per_warp_kernel); \ + } \ + if (TILE_TOKEN == 4) \ + { \ + return reinterpret_cast(llama4_fp8_bf16_gemm_per_warp_kernel); \ + } \ + if (TILE_TOKEN == 8) \ + { \ + return reinterpret_cast(llama4_fp8_bf16_gemm_per_warp_kernel); \ + } \ + throw std::invalid_argument("Invalid tile token"); \ + } while (0) + +#define DISPATCH_PER_WARP_FC_FP8_BF16_TILE_OUT(HIDDEN_IN, TILE_TOKEN, TILE_OUT, ALIGNED) \ + do \ + { \ + if (TILE_OUT == 1) \ + { \ + DISPATCH_PER_WARP_FC_FP8_BF16_TILE_TOKEN(HIDDEN_IN, TILE_TOKEN, 1, ALIGNED); \ + } \ + if (TILE_OUT == 2) \ + { \ + DISPATCH_PER_WARP_FC_FP8_BF16_TILE_TOKEN(HIDDEN_IN, TILE_TOKEN, 2, ALIGNED); \ + } \ + if (TILE_OUT == 3) \ + { \ + DISPATCH_PER_WARP_FC_FP8_BF16_TILE_TOKEN(HIDDEN_IN, TILE_TOKEN, 3, ALIGNED); \ + } \ + if (TILE_OUT == 4) \ + { \ + DISPATCH_PER_WARP_FC_FP8_BF16_TILE_TOKEN(HIDDEN_IN, TILE_TOKEN, 4, ALIGNED); \ + } \ + if (TILE_OUT == 8) \ + { \ + DISPATCH_PER_WARP_FC_FP8_BF16_TILE_TOKEN(HIDDEN_IN, TILE_TOKEN, 8, ALIGNED); \ + } \ + throw std::invalid_argument("Invalid tile token"); \ + } while (0) + +#define DEFINE_GET_PER_WARP_FUNC_PTR(HIDDEN_IN, ALIGNED) \ + void* get_per_warp_func_ptr_aligned_##ALIGNED##_##HIDDEN_IN##_(int tile_token, int tile_out) \ + { \ + DISPATCH_PER_WARP_FC_FP8_BF16_TILE_OUT(HIDDEN_IN, tile_token, tile_out, ALIGNED); \ + } + +} // namespace tensorrt_llm::kernels::llama4_min_latency::llama4_fp8_bf16_gemm diff --git a/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Fp8GemmSwiGLU.cu b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Fp8GemmSwiGLU.cu new file mode 100644 index 000000000000..ee6a54fd6d14 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Fp8GemmSwiGLU.cu @@ -0,0 +1,239 @@ +/* + * Copyright (c) 2025-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Fp8GemmSwiGLUPerBlockTemplate.cuh" +#include "tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Utils.cuh" + +#include +#include + +namespace tensorrt_llm::kernels::llama4_min_latency::llama4_fp8_fp8_gemm_swiglu +{ + +DEFINE_GET_FUNC_PTR(5120, true); + +dim3 get_grid_size(int num_tokens, int hidden_in, int hidden_out, int tile_token, int tile_out) +{ + if (num_tokens % tile_token != 0) + { + throw std::runtime_error("num_tokens must be divisible by tile_token"); + } + if (hidden_out % tile_out != 0) + { + throw std::runtime_error("hidden_out must be divisible by tile_out"); + } + if (hidden_in != 5120) + { + throw std::runtime_error("hidden_in must be 5120"); + } + return dim3(div_up(hidden_out, tile_out), div_up(num_tokens, tile_token), 1); +} + +bool is_supported(int num_tokens, int hidden_in, int hidden_out, int tile_token, int tile_out) +{ + if (hidden_in != 5120) + { + return false; + } + if (num_tokens % tile_token != 0) + { + return false; + } + if (hidden_out % tile_out != 0) + { + return false; + } + return true; +} + +template +void dispatch_llama4_fp8_fp8_gemm_swiglu_hidden_in(void const* __restrict__ A, void const* __restrict__ B, + void* __restrict__ C, void const* __restrict__ in_scale, void const* __restrict__ out_scale_inv, int num_tokens, + int hidden_in, int hidden_out, cudaStream_t stream) +{ + void* func_ptr; + constexpr int step = BLOCK_SIZE * VEC_SIZE; + + if (hidden_in == 5 * step) + { + func_ptr = get_func_ptr_aligned_true_5120_(TILE_TOKEN, TILE_OUT); + } + else + { + throw std::runtime_error("Unsupported hidden_in size " + std::to_string(hidden_in)); + } + + dim3 const grid_size = get_grid_size(num_tokens, hidden_in, hidden_out, TILE_TOKEN, TILE_OUT); + + void* args[] = {(void*) &A, (void*) &B, (void*) &C, (void*) &in_scale, (void*) &out_scale_inv, (void*) &num_tokens, + (void*) &hidden_in, (void*) &hidden_out}; + launch_kernel_fdl(grid_size, dim3(BLOCK_SIZE), stream, func_ptr, args, 8); +} + +template +void dispatch_llama4_fp8_fp8_gemm_swiglu_tile_out(void const* __restrict__ A, void const* __restrict__ B, + void* __restrict__ C, void const* __restrict__ in_scale, void const* __restrict__ out_scale_inv, int num_tokens, + int hidden_in, int hidden_out, int tile_out, cudaStream_t stream) +{ + if (tile_out == 1) + { + dispatch_llama4_fp8_fp8_gemm_swiglu_hidden_in( + A, B, C, in_scale, out_scale_inv, num_tokens, hidden_in, hidden_out, stream); + } + else if (tile_out == 2) + { + dispatch_llama4_fp8_fp8_gemm_swiglu_hidden_in( + A, B, C, in_scale, out_scale_inv, num_tokens, hidden_in, hidden_out, stream); + } + else if (tile_out == 4) + { + dispatch_llama4_fp8_fp8_gemm_swiglu_hidden_in( + A, B, C, in_scale, out_scale_inv, num_tokens, hidden_in, hidden_out, stream); + } + else + { + throw std::runtime_error("Unsupported tile_out size " + std::to_string(tile_out)); + } +} + +void dispatch_llama4_fp8_fp8_gemm_swiglu_tile_token(void const* __restrict__ A, void const* __restrict__ B, + void* __restrict__ C, void const* __restrict__ in_scale, void const* __restrict__ out_scale_inv, int num_tokens, + int hidden_in, int hidden_out, int tile_token, int tile_out, cudaStream_t stream) +{ + if (tile_token == 1) + { + dispatch_llama4_fp8_fp8_gemm_swiglu_tile_out<1>( + A, B, C, in_scale, out_scale_inv, num_tokens, hidden_in, hidden_out, tile_out, stream); + } + else if (tile_token == 2) + { + dispatch_llama4_fp8_fp8_gemm_swiglu_tile_out<2>( + A, B, C, in_scale, out_scale_inv, num_tokens, hidden_in, hidden_out, tile_out, stream); + } + else if (tile_token == 3) + { + dispatch_llama4_fp8_fp8_gemm_swiglu_tile_out<3>( + A, B, C, in_scale, out_scale_inv, num_tokens, hidden_in, hidden_out, tile_out, stream); + } + else if (tile_token == 4) + { + dispatch_llama4_fp8_fp8_gemm_swiglu_tile_out<4>( + A, B, C, in_scale, out_scale_inv, num_tokens, hidden_in, hidden_out, tile_out, stream); + } + else if (tile_token == 8) + { + dispatch_llama4_fp8_fp8_gemm_swiglu_tile_out<8>( + A, B, C, in_scale, out_scale_inv, num_tokens, hidden_in, hidden_out, tile_out, stream); + } + else + { + throw std::runtime_error("Unsupported tile_token size " + std::to_string(tile_token)); + } +} + +using Tactic = std::pair; + +int get_grid_size_total(int num_tokens, int hidden_in, int hidden_out, Tactic tactic) +{ + dim3 grid_size = get_grid_size(num_tokens, hidden_in, hidden_out, tactic.first, tactic.second); + return grid_size.x * grid_size.y * grid_size.z; +} + +std::optional get_hard_coded_heuristic_tactic(int num_tokens, int hidden_in, int hidden_out) +{ + + std::map, Tactic> hard_coded_tactics = { + {{1, 5120, 2048}, {1, 2}}, + {{1, 5120, 1024}, {1, 1}}, + {{2, 5120, 2048}, {2, 2}}, + {{2, 5120, 1024}, {1, 2}}, + {{4, 5120, 2048}, {2, 2}}, + {{4, 5120, 1024}, {2, 2}}, + {{8, 5120, 2048}, {4, 1}}, + {{8, 5120, 1024}, {4, 1}}, + }; + + auto it = hard_coded_tactics.find({num_tokens, hidden_in, hidden_out}); + if (it != hard_coded_tactics.end()) + { + return it->second; + } + + return std::nullopt; +} + +Tactic get_heuristic_tactic(int num_tokens, int hidden_in, int hidden_out) +{ + // Try hard-coded heuristics first. + auto hard_coded_tactic = get_hard_coded_heuristic_tactic(num_tokens, hidden_in, hidden_out); + if (hard_coded_tactic) + { + return hard_coded_tactic.value(); + } + + // Fall back to analysis-based heuristics. + int tile_token = 1; + int tile_out = 1; + + std::vector const tile_tokens = {1, 2, 3, 4, 8}; + for (auto t = tile_tokens.rbegin(); t != tile_tokens.rend(); ++t) + { + int const cta_threshold = 1024; + if (num_tokens % *t == 0 + && get_grid_size_total(num_tokens, hidden_in, hidden_out, Tactic{*t, tile_out}) >= cta_threshold) + { + tile_token = *t; + break; + } + } + + std::vector const tile_outs = {1, 2, 4}; + for (auto t = tile_outs.rbegin(); t != tile_outs.rend(); ++t) + { + int const cta_threshold = 1024; + if (hidden_in % *t == 0 + && get_grid_size_total(num_tokens, hidden_in, hidden_out, Tactic{tile_token, *t}) >= cta_threshold) + { + tile_out = *t; + break; + } + } + + Tactic tactic = {tile_token, tile_out}; + if (!is_supported(num_tokens, hidden_in, hidden_out, tactic.first, tactic.second)) + { + throw std::runtime_error("llama4Fp8Fp8GemmSwiGLU: unsupported tactic: " + std::to_string(tactic.first) + "-" + + std::to_string(tactic.second)); + } + + return tactic; +} + +void llama4_fp8_fp8_gemm_swiglu_op(int num_tokens, int hidden_in, int hidden_out, void const* A, void const* B, void* C, + void const* in_scale, void const* out_scale_inv, cudaStream_t stream) +{ + Tactic tactic = get_heuristic_tactic(num_tokens, hidden_in, hidden_out); + if (!is_supported(num_tokens, hidden_in, hidden_out, tactic.first, tactic.second)) + { + throw std::runtime_error("llama4Fp8Fp8GemmSwiGLU: unsupported tactic: " + std::to_string(tactic.first) + "-" + + std::to_string(tactic.second)); + } + + dispatch_llama4_fp8_fp8_gemm_swiglu_tile_token( + A, B, C, in_scale, out_scale_inv, num_tokens, hidden_in, hidden_out, tactic.first, tactic.second, stream); +} + +} // namespace tensorrt_llm::kernels::llama4_min_latency::llama4_fp8_fp8_gemm_swiglu diff --git a/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Fp8GemmSwiGLU.h b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Fp8GemmSwiGLU.h new file mode 100644 index 000000000000..aa11c4485d45 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Fp8GemmSwiGLU.h @@ -0,0 +1,31 @@ +/* + * Copyright (c) 2025-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include +#include +#include + +namespace tensorrt_llm::kernels::llama4_min_latency::llama4_fp8_fp8_gemm_swiglu +{ + +void llama4_fp8_fp8_gemm_swiglu_op(int num_tokens, int hidden_in, int hidden_out, void const* A, void const* B, void* C, + void const* in_scale, void const* out_scale_inv, cudaStream_t stream); + +} // namespace tensorrt_llm::kernels::llama4_min_latency::llama4_fp8_fp8_gemm_swiglu diff --git a/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Fp8GemmSwiGLUPerBlockTemplate.cuh b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Fp8GemmSwiGLUPerBlockTemplate.cuh new file mode 100644 index 000000000000..e0a459656b70 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Fp8GemmSwiGLUPerBlockTemplate.cuh @@ -0,0 +1,340 @@ +/* + * Copyright (c) 2025-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Utils.cuh" + +#include + +namespace tensorrt_llm::kernels::llama4_min_latency::llama4_fp8_fp8_gemm_swiglu +{ + +// Grid size is num_tokens / TILE_TOKEN * hidden_out / TILE_OUT. +// Each block processes TILE_TOKEN tokens and TILE_OUT rows. +// within each block, it steps through hidden_in in steps of BLOCK_SIZE * VEC_SIZE. +template +__launch_bounds__(BLOCK_SIZE) __global__ void llama4_fp8_fp8_gemm_swiglu_per_block_kernel( + __nv_fp8_e4m3 const* __restrict__ A, // Input tensor [num_tokens][hidden_in] + __nv_fp8_e4m3 const* __restrict__ B, // Input tensor [2 * hidden_out][hidden_in] + __nv_fp8_e4m3* __restrict__ C, // Output tensor [num_tokens][hidden_out] + float const* __restrict__ in_scale, float const* __restrict__ out_scale_inv, int num_tokens, int hidden_in, + int hidden_out) +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + // Shared memory for block reduction + __shared__ float reduce_buffer_gate[TILE_TOKEN][TILE_OUT][BLOCK_SIZE]; + __shared__ float reduce_buffer_linear[TILE_TOKEN][TILE_OUT][BLOCK_SIZE]; + + // Each thread accumulates its partial sum + float2 thread_sum_gate[TILE_TOKEN][TILE_OUT]; + float2 thread_sum_linear[TILE_TOKEN][TILE_OUT]; +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + thread_sum_gate[tile_token_idx][tile_out_idx].x = 0.0f; + thread_sum_gate[tile_token_idx][tile_out_idx].y = 0.0f; + thread_sum_linear[tile_token_idx][tile_out_idx].x = 0.0f; + thread_sum_linear[tile_token_idx][tile_out_idx].y = 0.0f; + } + } + + int const token_idx = blockIdx.y; + int const row_idx = blockIdx.x; + int const tid = threadIdx.x; + + int chunk = 0; +#if ENABLE_ACQBULK && ENABLE_PREFETCH +#pragma unroll 9 + for (; chunk < ((HIDDEN_IN > 0 ? HIDDEN_IN : hidden_in) / BLOCK_SIZE / VEC_SIZE - 1); chunk++) + { +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + int current_row = row_idx * TILE_OUT + tile_out_idx; + int base_idx = chunk * BLOCK_SIZE + tid; + asm volatile("prefetch.global.L2 [%0];" ::"l"(&B[current_row * hidden_in / VEC_SIZE + base_idx]) + : "memory"); + asm volatile( + "prefetch.global.L2 [%0];" ::"l"(&B[(current_row + hidden_out) * hidden_in / VEC_SIZE + base_idx]) + : "memory"); + } + } + { +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + int current_row = row_idx * TILE_OUT + tile_out_idx; + int base_idx = chunk * BLOCK_SIZE + tid; + if (ALIGNED || base_idx * VEC_SIZE < hidden_in) + { + asm volatile("prefetch.global.L2 [%0];" ::"l"(&B[current_row * hidden_in / VEC_SIZE + base_idx]) + : "memory"); + asm volatile( + "prefetch.global.L2 [%0];" ::"l"(&B[(current_row + hidden_out) * hidden_in / VEC_SIZE + base_idx]) + : "memory"); + } + } + } +#endif + +#if ENABLE_ACQBULK + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + + // Processing 8 elements each + chunk = 0; +#pragma unroll 9 + for (; chunk < ((HIDDEN_IN > 0 ? HIDDEN_IN : hidden_in) / BLOCK_SIZE / VEC_SIZE - 1); chunk++) + { + int base_idx = chunk * BLOCK_SIZE + tid; + // Load values from tensor B + aligned_fp8x8 b_vec_gate[TILE_OUT]; + aligned_fp8x8 b_vec_linear[TILE_OUT]; +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + int current_row = row_idx * TILE_OUT + tile_out_idx; + b_vec_gate[tile_out_idx] + = reinterpret_cast(B)[current_row * hidden_in / VEC_SIZE + base_idx]; + b_vec_linear[tile_out_idx] = reinterpret_cast( + B)[(current_row + hidden_out) * hidden_in / VEC_SIZE + base_idx]; + } + + // Load values from tensor A + aligned_fp8x8 a_vec[TILE_TOKEN]; +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { + int current_token = token_idx * TILE_TOKEN + tile_token_idx; + a_vec[tile_token_idx] + = reinterpret_cast(A)[current_token * hidden_in / VEC_SIZE + base_idx]; + } + + // Compute partial sum +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { + aligned_fp8x8 a_vec_current = a_vec[tile_token_idx]; + aligned_fp8x8 b_vec_current_gate = b_vec_gate[tile_out_idx]; + aligned_fp8x8 b_vec_current_linear = b_vec_linear[tile_out_idx]; +#pragma unroll + for (int i = 0; i < VEC_SIZE / 4; i++) + { + float4 a_val = float4(a_vec_current.data[i]); + float4 b_val_gate = float4(b_vec_current_gate.data[i]); + float4 b_val_linear = float4(b_vec_current_linear.data[i]); + + thread_sum_gate[tile_token_idx][tile_out_idx] = ffma2(make_float2(a_val.x, a_val.y), + make_float2(b_val_gate.x, b_val_gate.y), thread_sum_gate[tile_token_idx][tile_out_idx]); + thread_sum_gate[tile_token_idx][tile_out_idx] = ffma2(make_float2(a_val.z, a_val.w), + make_float2(b_val_gate.z, b_val_gate.w), thread_sum_gate[tile_token_idx][tile_out_idx]); + thread_sum_linear[tile_token_idx][tile_out_idx] = ffma2(make_float2(a_val.x, a_val.y), + make_float2(b_val_linear.x, b_val_linear.y), thread_sum_linear[tile_token_idx][tile_out_idx]); + thread_sum_linear[tile_token_idx][tile_out_idx] = ffma2(make_float2(a_val.z, a_val.w), + make_float2(b_val_linear.z, b_val_linear.w), thread_sum_linear[tile_token_idx][tile_out_idx]); + } + } + } + } + + // The last chunk may be partial. + { + int base_idx = chunk * BLOCK_SIZE + tid; + if (ALIGNED || base_idx * VEC_SIZE < hidden_in) + { + // Load values from tensor B + aligned_fp8x8 b_vec_gate[TILE_OUT]; + aligned_fp8x8 b_vec_linear[TILE_OUT]; +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + int current_row = row_idx * TILE_OUT + tile_out_idx; + b_vec_gate[tile_out_idx] + = reinterpret_cast(B)[current_row * hidden_in / VEC_SIZE + base_idx]; + b_vec_linear[tile_out_idx] = reinterpret_cast( + B)[(current_row + hidden_out) * hidden_in / VEC_SIZE + base_idx]; + } + + // Load values from tensor A + aligned_fp8x8 a_vec[TILE_TOKEN]; +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { + int current_token = token_idx * TILE_TOKEN + tile_token_idx; + a_vec[tile_token_idx] + = reinterpret_cast(A)[current_token * hidden_in / VEC_SIZE + base_idx]; + } + + // Compute partial sum +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { + aligned_fp8x8 a_vec_current = a_vec[tile_token_idx]; + aligned_fp8x8 b_vec_current_gate = b_vec_gate[tile_out_idx]; + aligned_fp8x8 b_vec_current_linear = b_vec_linear[tile_out_idx]; +#pragma unroll + for (int i = 0; i < VEC_SIZE / 4; i++) + { + float4 a_val = float4(a_vec_current.data[i]); + float4 b_val_gate = float4(b_vec_current_gate.data[i]); + float4 b_val_linear = float4(b_vec_current_linear.data[i]); + + thread_sum_gate[tile_token_idx][tile_out_idx] = ffma2(make_float2(a_val.x, a_val.y), + make_float2(b_val_gate.x, b_val_gate.y), thread_sum_gate[tile_token_idx][tile_out_idx]); + thread_sum_gate[tile_token_idx][tile_out_idx] = ffma2(make_float2(a_val.z, a_val.w), + make_float2(b_val_gate.z, b_val_gate.w), thread_sum_gate[tile_token_idx][tile_out_idx]); + thread_sum_linear[tile_token_idx][tile_out_idx] + = ffma2(make_float2(a_val.x, a_val.y), make_float2(b_val_linear.x, b_val_linear.y), + thread_sum_linear[tile_token_idx][tile_out_idx]); + thread_sum_linear[tile_token_idx][tile_out_idx] + = ffma2(make_float2(a_val.z, a_val.w), make_float2(b_val_linear.z, b_val_linear.w), + thread_sum_linear[tile_token_idx][tile_out_idx]); + } + } + } + } + } + + // Reduce partial sums using warp-level reduction. +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { + float warp_sum_gate + = thread_sum_gate[tile_token_idx][tile_out_idx].x + thread_sum_gate[tile_token_idx][tile_out_idx].y; + float warp_sum_linear + = thread_sum_linear[tile_token_idx][tile_out_idx].x + thread_sum_linear[tile_token_idx][tile_out_idx].y; +#pragma unroll + for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) + { + warp_sum_gate += __shfl_down_sync(0xffffffff, warp_sum_gate, offset); + warp_sum_linear += __shfl_down_sync(0xffffffff, warp_sum_linear, offset); + } + // First thread in each warp writes to shared memory + if (tid % WARP_SIZE == 0) + { + reduce_buffer_gate[tile_token_idx][tile_out_idx][tid / WARP_SIZE] = warp_sum_gate; + reduce_buffer_linear[tile_token_idx][tile_out_idx][tid / WARP_SIZE] = warp_sum_linear; + } + } + } + + __syncthreads(); + + if (tid == 0) + { +#pragma unroll + for (int tile_token_idx = 0; tile_token_idx < TILE_TOKEN; tile_token_idx++) + { +#pragma unroll + for (int tile_out_idx = 0; tile_out_idx < TILE_OUT; tile_out_idx++) + { + float block_sum_gate = 0.0f; + float block_sum_linear = 0.0f; +#pragma unroll + for (int i = 0; i < BLOCK_SIZE / WARP_SIZE; i++) + { + block_sum_gate += reduce_buffer_gate[tile_token_idx][tile_out_idx][i]; + block_sum_linear += reduce_buffer_linear[tile_token_idx][tile_out_idx][i]; + } + int current_row = row_idx * TILE_OUT + tile_out_idx; + int current_token = token_idx * TILE_TOKEN + tile_token_idx; + float in_scale_val = in_scale[0]; + float out_scale_inv_val = out_scale_inv[0]; + C[current_token * hidden_out + current_row] = __nv_fp8_e4m3( + silu(block_sum_gate * in_scale_val) * block_sum_linear * in_scale_val * out_scale_inv_val); + } + } + } + +#if ENABLE_PREEXIT + asm volatile("griddepcontrol.launch_dependents;"); +#endif +#endif +} + +#define DISPATCH_FC_FP8_BF16_TILE_TOKEN(HIDDEN_IN, TILE_TOKEN, TILE_OUT, ALIGNED) \ + do \ + { \ + if (TILE_TOKEN == 1) \ + { \ + return reinterpret_cast( \ + llama4_fp8_fp8_gemm_swiglu_per_block_kernel); \ + } \ + if (TILE_TOKEN == 2) \ + { \ + return reinterpret_cast( \ + llama4_fp8_fp8_gemm_swiglu_per_block_kernel); \ + } \ + if (TILE_TOKEN == 3) \ + { \ + return reinterpret_cast( \ + llama4_fp8_fp8_gemm_swiglu_per_block_kernel); \ + } \ + if (TILE_TOKEN == 4) \ + { \ + return reinterpret_cast( \ + llama4_fp8_fp8_gemm_swiglu_per_block_kernel); \ + } \ + if (TILE_TOKEN == 8) \ + { \ + return reinterpret_cast( \ + llama4_fp8_fp8_gemm_swiglu_per_block_kernel); \ + } \ + throw std::invalid_argument("Invalid tile token"); \ + } while (0) + +#define DISPATCH_FC_FP8_BF16_TILE_OUT(HIDDEN_IN, TILE_TOKEN, TILE_OUT, ALIGNED) \ + do \ + { \ + if (TILE_OUT == 1) \ + { \ + DISPATCH_FC_FP8_BF16_TILE_TOKEN(HIDDEN_IN, TILE_TOKEN, 1, ALIGNED); \ + } \ + if (TILE_OUT == 2) \ + { \ + DISPATCH_FC_FP8_BF16_TILE_TOKEN(HIDDEN_IN, TILE_TOKEN, 2, ALIGNED); \ + } \ + if (TILE_OUT == 3) \ + { \ + DISPATCH_FC_FP8_BF16_TILE_TOKEN(HIDDEN_IN, TILE_TOKEN, 3, ALIGNED); \ + } \ + if (TILE_OUT == 4) \ + { \ + DISPATCH_FC_FP8_BF16_TILE_TOKEN(HIDDEN_IN, TILE_TOKEN, 4, ALIGNED); \ + } \ + throw std::invalid_argument("Invalid tile token"); \ + } while (0) + +#define DEFINE_GET_FUNC_PTR(HIDDEN_IN, ALIGNED) \ + void* get_func_ptr_aligned_##ALIGNED##_##HIDDEN_IN##_(int tile_token, int tile_out) \ + { \ + DISPATCH_FC_FP8_BF16_TILE_OUT(HIDDEN_IN, tile_token, tile_out, ALIGNED); \ + } + +} // namespace tensorrt_llm::kernels::llama4_min_latency::llama4_fp8_fp8_gemm_swiglu diff --git a/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4MinLatencyMoEOp.cu b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4MinLatencyMoEOp.cu new file mode 100644 index 000000000000..2198885d5005 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4MinLatencyMoEOp.cu @@ -0,0 +1,354 @@ +/* + * Copyright (c) 2020-2024, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/llama4MinLatencyKernels/llama4MinLatencyMoEOp.h" +#include "tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Utils.cuh" +#include +#include +#include +#include + +#define NUM_EXPERTS 128 +#define HIDDEN_SIZE 5120 +#define INTER_SIZE 1024 +#define BLOCK_SIZE 128 +#define WARP_SIZE 32 +#define TILE_ROW 4 +#define VEC_SIZE 8 + +#define ENABLE_ACQBULK 1 +#define ENABLE_PREFETCH 1 +#define ENABLE_PREEXIT 1 + +namespace tensorrt_llm::kernels::llama4_min_latency::llama4_moe +{ + +#define TOPK_VEC_SIZE 4 +static_assert(NUM_EXPERTS == TOPK_VEC_SIZE * WARP_SIZE, "NUM_EXPERTS must be equal to TOPK_VEC_SIZE * WARP_SIZE"); + +// This is the hand-optimized kernel by Po-Han. +// The computation is: +// C = silu(AxB_gated * in_scale * sigmoid(logit)) * (AxB_linear * in_scale * sigmoid(logit)) * out_scale_inv +// The out_scale_inv cannot be fused with in_scale because silu() is non-linear. +// Also, Llama-4 applies score scaling, which is sigmoid(logit), on tensor A. +__global__ void llama4_moe_fc13_swiglu_fp8_kernel(int num_tokens, + __nv_fp8_e4m3 const* __restrict__ A, // Input tensor [num_tokens][HIDDEN_SIZE] + __nv_fp8_e4m3 const* __restrict__ B, // Input tensor [num_experts][INTER_SIZE*2][HIDDEN_SIZE] + __nv_bfloat16 const* __restrict__ logits, // Input tensor logits [num_tokens][num_experts] + __nv_fp8_e4m3* __restrict__ C, // Output tensor [num_tokens][INTER_SIZE] + int* __restrict__ exp_idx, // Output tensor [num_tokens] + float const* __restrict__ in_scales, // Input scales [num_experts] + float const* __restrict__ out_scale_inv // Output scale [1] +) +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + // Shared memory for block reduction + __shared__ float reduce_buffer[2][BLOCK_SIZE / WARP_SIZE]; + + // Each thread accumulates its partial sum + float2 thread_sum_linear; + thread_sum_linear.x = 0.0f; + thread_sum_linear.y = 0.0f; + + float2 thread_sum_gate; + thread_sum_gate.x = 0.0f; + thread_sum_gate.y = 0.0f; + + // Each thread processes 8 elements at a time, 5 times + int const token_idx = blockIdx.x / INTER_SIZE; + int const row = blockIdx.x % INTER_SIZE; // Matrix row / Output element index + int const tid = threadIdx.x; // Thread ID within the block + int const lane_id = tid % WARP_SIZE; // Lane ID within the warp + + // Preload the scaling factors before ACQBULK. + __shared__ float in_scales_shared[NUM_EXPERTS]; + in_scales_shared[tid] = in_scales[tid]; + + // Logits depends on the previous kernel, so we cannot prefetch anything. +#if ENABLE_ACQBULK + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + + // Perform top1 within the current thread, which processes 4 experts. + aligned_bfloat16x4 logits_vec; + logits_vec = reinterpret_cast(logits)[token_idx * NUM_EXPERTS / TOPK_VEC_SIZE + lane_id]; + + __nv_bfloat16 best_logit = logits_vec.data[0]; + int base_exp = lane_id * TOPK_VEC_SIZE; + int best_exp = base_exp; +#pragma unroll + for (int i = 1; i < TOPK_VEC_SIZE; i++) + { + __nv_bfloat16 current_logit = logits_vec.data[i]; + if (current_logit >= best_logit) + { + best_logit = current_logit; + best_exp = base_exp + i; + } + } + + // Perform top1 across threads using Warp reduction. + // We pack logit and expert index into an int so that we can use integer max op for reduction. + int best_result + = ((int) (__bfloat16_as_short(best_logit) ^ (best_logit < __nv_bfloat16(0.f) ? 0x7fff : 0)) << 16) | best_exp; + +#pragma unroll + for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) + { + best_result = max(best_result, __shfl_down_sync(0xffffffff, best_result, offset)); + } + + // Broadcast the best result to all threads in the warp. + best_result = __shfl_sync(0xffffffff, best_result, 0); + + // Extract the expert index and score. + int expert_idx = best_result & 0xffff; + float top_logit = __bfloat162float(__short_as_bfloat16(short((best_result >> 16) & 0xffff))); + + // Select the corresponding expert weight. + int expert_weight_offset = expert_idx * 2 * INTER_SIZE * HIDDEN_SIZE / VEC_SIZE; + + // Process 5 chunks of 8 elements each +#pragma unroll + for (int chunk = 0; chunk < HIDDEN_SIZE / BLOCK_SIZE / VEC_SIZE; chunk++) + { + int base_idx = chunk * BLOCK_SIZE + tid; + // Load 8 elements at once + aligned_fp8x8 a_vec = reinterpret_cast(A)[token_idx * HIDDEN_SIZE / VEC_SIZE + base_idx]; + aligned_fp8x8 b_vec_linear + = reinterpret_cast(B)[row * HIDDEN_SIZE / VEC_SIZE + base_idx + expert_weight_offset]; + aligned_fp8x8 b_vec_gate = reinterpret_cast( + B)[(row + INTER_SIZE) * HIDDEN_SIZE / VEC_SIZE + base_idx + expert_weight_offset]; + +#pragma unroll + for (int i = 0; i < VEC_SIZE / 4; i++) + { + float4 a_val = float4(a_vec.data[i]); + float4 b_val_linear = float4(b_vec_linear.data[i]); + float4 b_val_gate = float4(b_vec_gate.data[i]); + + thread_sum_linear + = ffma2(make_float2(a_val.x, a_val.y), make_float2(b_val_linear.x, b_val_linear.y), thread_sum_linear); + thread_sum_linear + = ffma2(make_float2(a_val.z, a_val.w), make_float2(b_val_linear.z, b_val_linear.w), thread_sum_linear); + thread_sum_gate + = ffma2(make_float2(a_val.x, a_val.y), make_float2(b_val_gate.x, b_val_gate.y), thread_sum_gate); + thread_sum_gate + = ffma2(make_float2(a_val.z, a_val.w), make_float2(b_val_gate.z, b_val_gate.w), thread_sum_gate); + } + } + + // Block reduction + float warp_sum_linear = thread_sum_linear.x + thread_sum_linear.y; + float warp_sum_gate = thread_sum_gate.x + thread_sum_gate.y; +#pragma unroll + for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) + { + warp_sum_linear += __shfl_down_sync(0xffffffff, warp_sum_linear, offset); + warp_sum_gate += __shfl_down_sync(0xffffffff, warp_sum_gate, offset); + } + + if (tid % WARP_SIZE == 0) + { + reduce_buffer[0][tid / WARP_SIZE] = warp_sum_linear; + reduce_buffer[1][tid / WARP_SIZE] = warp_sum_gate; + } + __syncthreads(); + + if (tid == 0) + { + float block_sum_linear = warp_sum_linear; + float block_sum_gate = warp_sum_gate; +#pragma unroll + for (int i = 1; i < BLOCK_SIZE / WARP_SIZE; i++) + { + block_sum_linear += reduce_buffer[0][i]; + block_sum_gate += reduce_buffer[1][i]; + } + float fused_in_scale = in_scales_shared[expert_idx] * sigmoid(top_logit); + C[token_idx * INTER_SIZE + row] = __nv_fp8_e4m3( + silu(block_sum_gate * fused_in_scale) * block_sum_linear * fused_in_scale * out_scale_inv[0]); + if (row == 0) + { + exp_idx[token_idx] = expert_idx; + } + } + +#if ENABLE_PREEXIT + asm volatile("griddepcontrol.launch_dependents;"); +#endif +#endif +} + +// Launch llama4_moe_fc13_swiglu_fp8_kernel +void launch_llama4_moe_fc13_swiglu_fp8_kernel(int num_tokens, int num_experts, + void const* __restrict__ A, // Input tensor A [num_tokens][HIDDEN_SIZE] + void const* __restrict__ B, // Input tensor B [num_experts][INTER_SIZE*2][HIDDEN_SIZE] + void const* __restrict__ logits, // Input tensor logits [num_tokens][num_experts] + void* __restrict__ C, // Output tensor [num_tokens][INTER_SIZE] + int* __restrict__ exp_idx, // Output tensor [num_tokens] + float const* __restrict__ in_scales, // Input scales [num_experts] + float const* __restrict__ out_scale_inv, // Output scale [1] + cudaStream_t stream) +{ + int const grid_size = num_tokens * INTER_SIZE; + if (num_experts != NUM_EXPERTS) + { + printf("The implementation currently assumes num_experts = %d\n", NUM_EXPERTS); + exit(1); + } + + void* args[] = {(void*) &num_tokens, (void*) &A, (void*) &B, (void*) &logits, (void*) &C, (void*) &exp_idx, + (void*) &in_scales, (void*) &out_scale_inv}; + launch_kernel_fdl(dim3(grid_size), dim3(BLOCK_SIZE), stream, (void*) llama4_moe_fc13_swiglu_fp8_kernel, args, 8); +} + +// This is the hand-optimized kernel by Po-Han. +__global__ void llama4_moe_fc2_fp8_kernel(int num_tokens, + __nv_fp8_e4m3 const* __restrict__ A, // Input tensor A [num_tokens][INTER_SIZE] + __nv_fp8_e4m3 const* __restrict__ B, // Input tensor B [num_experts][HIDDEN_SIZE][INTER_SIZE] + int const* __restrict__ exp_idx, // Input tensor exp_idx [num_tokens]. + __nv_bfloat16* __restrict__ C, // Output tensor [num_tokens][HIDDEN_SIZE] + float const* __restrict__ scaling_factors // Scaling factors [num_experts] +) +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + // Shared memory for block reduction + __shared__ float reduce_buffer[TILE_ROW][BLOCK_SIZE / WARP_SIZE]; + + // Each thread processes 8 elements at a time, 5 times + int const token_idx = blockIdx.x / (HIDDEN_SIZE / TILE_ROW); + int const row = blockIdx.x % (HIDDEN_SIZE / TILE_ROW); // Matrix row / Output element index + int const tid = threadIdx.x; // Thread ID within the block + + // Preload the scaling factors before ACQBULK. + __shared__ float scaling_factors_shared[NUM_EXPERTS]; + scaling_factors_shared[tid] = scaling_factors[tid]; + +#if ENABLE_ACQBULK + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + + // Select the corresponding expert weight. + int expert_idx = exp_idx[token_idx]; + int expert_weight_offset = expert_idx * HIDDEN_SIZE * INTER_SIZE / VEC_SIZE; + int base_idx = tid; + + // Load 8 elements at once. + aligned_fp8x8 a_vec = reinterpret_cast(A)[token_idx * INTER_SIZE / VEC_SIZE + base_idx]; + aligned_fp8x8 b_vec[TILE_ROW]; +#pragma unroll + for (int tile_row_idx = 0; tile_row_idx < TILE_ROW; tile_row_idx++) + { + int row_current = tile_row_idx + row * TILE_ROW; + b_vec[tile_row_idx] = reinterpret_cast( + B)[row_current * INTER_SIZE / VEC_SIZE + base_idx + expert_weight_offset]; + } + + // Loop over TILE_ROW times to compute the gemm result. +#pragma unroll + for (int tile_row_idx = 0; tile_row_idx < TILE_ROW; tile_row_idx++) + { + // Each thread accumulates its partial sum + float2 thread_sum = {0.0f, 0.0f}; + thread_sum.x = 0.0f; + thread_sum.y = 0.0f; + +#pragma unroll + for (int i = 0; i < VEC_SIZE / 4; i++) + { + float4 a_val = float4(a_vec.data[i]); + float4 b_val = float4(b_vec[tile_row_idx].data[i]); + thread_sum = ffma2(make_float2(a_val.x, a_val.y), make_float2(b_val.x, b_val.y), thread_sum); + thread_sum = ffma2(make_float2(a_val.z, a_val.w), make_float2(b_val.z, b_val.w), thread_sum); + } + + // Warp reduction + float warp_sum = thread_sum.x + thread_sum.y; +#pragma unroll + for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) + { + warp_sum += __shfl_down_sync(0xffffffff, warp_sum, offset); + } + + if (tid % WARP_SIZE == 0) + { + reduce_buffer[tile_row_idx][tid / WARP_SIZE] = warp_sum; + } + } + + __syncthreads(); + + // Use the first TILE_ROW threads to do block reduction and writes the result. + if (tid < TILE_ROW) + { + int row_current = tid + row * TILE_ROW; + float block_sum = 0.0f; +#pragma unroll + for (int i = 0; i < BLOCK_SIZE / WARP_SIZE; i++) + { + block_sum += reduce_buffer[tid][i]; + } + float scaling_factor = scaling_factors_shared[expert_idx]; + C[token_idx * HIDDEN_SIZE + row_current] = __float2bfloat16(block_sum * scaling_factor); + } + +#if ENABLE_PREEXIT + asm volatile("griddepcontrol.launch_dependents;"); +#endif +#endif +} + +void launch_llama4_moe_fc2_fp8_kernel(int num_tokens, int num_experts, + void const* __restrict__ A, // Input tensor A [num_tokens][INTER_SIZE] + void const* __restrict__ B, // Input tensor B [num_experts][HIDDEN_SIZE][INTER_SIZE] + int const* __restrict__ exp_idx, // Input tensor exp_idx [num_tokens]. + void* __restrict__ C, // Output tensor [num_tokens][HIDDEN_SIZE] + float const* __restrict__ scaling_factors, // Scaling factors [num_experts] + cudaStream_t stream) +{ + if (num_experts != NUM_EXPERTS) + { + printf("Current implementation assumes num_experts == %d\n", NUM_EXPERTS); + exit(1); + } + int const grid_size = num_tokens * HIDDEN_SIZE / TILE_ROW; + + void* args[] + = {(void*) &num_tokens, (void*) &A, (void*) &B, (void*) &exp_idx, (void*) &C, (void*) &scaling_factors}; + launch_kernel_fdl(dim3(grid_size), dim3(BLOCK_SIZE), stream, (void*) llama4_moe_fc2_fp8_kernel, args, 6); +} + +void run_moe_llama4_tp8ep1_min_latency(int num_tokens, int num_experts, + void const* __restrict__ input_activations_void, // Input tensor FP8 [num_tokens][HIDDEN_SIZE] + void const* __restrict__ router_logits_void, // Router logits tensor BF16 [num_tokens][num_experts] + void const* __restrict__ fc1_expert_weights_void, // FC13 weight tensor FP8 [num_experts][2*INTER_SIZE][HIDDEN_SIZE] + void const* __restrict__ fc2_expert_weights_void, // FC2 weight tensor FP8 [num_experts][HIDDEN_SIZE][INTER_SIZE] + float const* __restrict__ dequant_fc1, // FC1 out scale factor FP32 [num_experts] + float const* __restrict__ quant_fc2, // FC2 input scaling factor FP32 [1] + float const* __restrict__ dequant_fc2, // FC2 out scaling factor FP32 [num_experts] + void* __restrict__ fc2_input_activations_void, // FC2 input tensor FP8 [num_tokens][INTER_SIZE] + int* __restrict__ exp_idx, // Expert indexes INT [num_tokens] + void* __restrict__ output_void, // FC2 output tensor BF16 [num_tokens][HIDDEN_SIZE] + cudaStream_t stream) +{ + launch_llama4_moe_fc13_swiglu_fp8_kernel(num_tokens, num_experts, input_activations_void, fc1_expert_weights_void, + router_logits_void, fc2_input_activations_void, exp_idx, dequant_fc1, quant_fc2, stream); + launch_llama4_moe_fc2_fp8_kernel(num_tokens, num_experts, fc2_input_activations_void, fc2_expert_weights_void, + exp_idx, output_void, dequant_fc2, stream); +} + +} // namespace tensorrt_llm::kernels::llama4_min_latency::llama4_moe diff --git a/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4MinLatencyMoEOp.h b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4MinLatencyMoEOp.h new file mode 100644 index 000000000000..7d0d52c68306 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4MinLatencyMoEOp.h @@ -0,0 +1,40 @@ +/* + * Copyright (c) 2025-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include +#include +#include + +namespace tensorrt_llm::kernels::llama4_min_latency::llama4_moe +{ + +// Launch moe_mlp_fc13_swiglu_fp8_5120 and moe_fc_fp8_bf16_1024. +void run_moe_llama4_tp8ep1_min_latency(int num_tokens, int num_experts, + void const* __restrict__ input_activations_void, // Input tensor FP8 [num_tokens][HIDDEN_SIZE] + void const* __restrict__ router_logits_void, // Router logits tensor BF16 [num_tokens][num_experts] + void const* __restrict__ fc1_expert_weights_void, // FC13 weight tensor FP8 [num_experts][2*INTER_SIZE][HIDDEN_SIZE] + void const* __restrict__ fc2_expert_weights_void, // FC2 weight tensor FP8 [num_experts][HIDDEN_SIZE][INTER_SIZE] + float const* __restrict__ dequant_fc1, // FC1 out scale factor FP32 [num_experts] + float const* __restrict__ quant_fc2, // FC2 input scaling factor FP32 [1] + float const* __restrict__ dequant_fc2, // FC2 out scaling factor FP32 [num_experts] + void* __restrict__ fc2_input_activations_void, // FC2 input tensor FP8 [num_tokens][INTER_SIZE] + int* __restrict__ exp_idx, // Expert indexes INT [num_tokens] + void* __restrict__ output_void, // FC2 output tensor BF16 [num_tokens][HIDDEN_SIZE] + cudaStream_t stream); + +} // namespace tensorrt_llm::kernels::llama4_min_latency::llama4_moe diff --git a/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Utils.cuh b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Utils.cuh new file mode 100644 index 000000000000..1dabd2dc64a8 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Utils.cuh @@ -0,0 +1,122 @@ +/* + * Copyright (c) 2025-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include "tensorrt_llm/common/envUtils.h" + +namespace tensorrt_llm::kernels::llama4_min_latency +{ + +namespace llama4_bf16_bf16_gemm +{ +constexpr int GEMM_K = 5120; +constexpr int BLOCK_SIZE = 256; +constexpr int NUM_EXPERTS = 128; +constexpr int VEC_SIZE = 4; + +} // namespace llama4_bf16_bf16_gemm + +namespace llama4_fp8_bf16_gemm +{ + +constexpr int HIDDEN_IN = 5120; +constexpr int HIDDEN_OUT = 896; +constexpr int Q_HIDDEN_OUT = 640; + +constexpr float FLOOR_SCALE = 8192.0; +constexpr float ATTN_SCALE = 0.1; + +constexpr int BLOCK_SIZE = 128; +constexpr int WARP_SIZE = 32; +constexpr int WARP_PER_BLOCK = BLOCK_SIZE / WARP_SIZE; +constexpr int VEC_SIZE = 8; + +constexpr bool ENABLE_ACQBULK = 1; +constexpr bool ENABLE_PREFETCH = 1; +constexpr bool ENABLE_PREEXIT = 0; + +} // namespace llama4_fp8_bf16_gemm + +namespace llama4_fp8_fp8_gemm_swiglu +{ + +constexpr int BLOCK_SIZE = 128; +constexpr int WARP_SIZE = 32; +constexpr int VEC_SIZE = 8; + +constexpr bool ENABLE_ACQBULK = 1; +constexpr bool ENABLE_PREFETCH = 0; +constexpr bool ENABLE_PREEXIT = 0; + +} // namespace llama4_fp8_fp8_gemm_swiglu + +inline void launch_kernel_fdl( + dim3 grid_dim, dim3 block_dim, cudaStream_t stream, void* kernel_func, void* args[], int num_args) +{ + cudaLaunchConfig_t config; + config.gridDim = grid_dim; + config.blockDim = block_dim; + config.dynamicSmemBytes = 0; + config.stream = stream; + + cudaLaunchAttribute attrs[1]; + config.attrs = attrs; + config.numAttrs = 0; + attrs[config.numAttrs].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[config.numAttrs++].val.programmaticStreamSerializationAllowed + = (tensorrt_llm::common::getEnvEnablePDL() ? 1 : 0); + + cudaLaunchKernelExC(&config, (void const*) kernel_func, args); +} + +inline int div_up(int x, int y) +{ + return (x + y - 1) / y; +} + +__device__ __forceinline__ float2 ffma2(float2 x, float2 y, float2 acc) +{ +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000)) + return __ffma2_rn(x, y, acc); +#else + return make_float2(x.x * y.x + acc.x, x.y * y.y + acc.y); +#endif +} + +__device__ __forceinline__ float silu(float x) +{ + return x / (1.0f + __expf(-x)); +} + +__device__ __forceinline__ float sigmoid(float x) +{ + return 1.0f / (1.0f + __expf(-x)); +} + +struct __align__(8) aligned_fp8x8 +{ + __align__(8) __nv_fp8x4_e4m3 data[2]; +}; + +struct __align__(8) aligned_bfloat16x4 +{ + __align__(8) __nv_bfloat16 data[4]; +}; + +} // namespace tensorrt_llm::kernels::llama4_min_latency diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/GemmInterface.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/GemmInterface.h index feb213aab26b..7a3296c2df96 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/GemmInterface.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/GemmInterface.h @@ -444,7 +444,9 @@ int32_t GemmInterface::run(GemmConfig const& config, void* workspace, GemmData c // Run the kernel. auto result = trtllm::gen::launchKernel((void*) &kernelParams, cudaStream, config.mSharedMemSize, cuFunction, - block3, grid3, cluster3, config.mOptions.mGridWaitForPrimary); + block3, grid3, cluster3, + config.mOptions.mGridWaitForPrimaryEarlyExit | config.mOptions.mGridWaitForPrimaryA + | config.mOptions.mGridWaitForPrimaryB); if (result != CUDA_SUCCESS) { return -1; diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/GemmOptions.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/GemmOptions.h index 13273d2a184a..24624ee0aa2e 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/GemmOptions.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/GemmOptions.h @@ -50,7 +50,7 @@ void printArgs(T first, Args... args) #define TLLM_LOG_ERROR(...) TLLM_CHECK_ERROR(false, __VA_ARGS__) -#define TLLM_CHECK_ERROR_FMT(...) TLLM_CHECK_ERROR(false, __VA_ARGS__) +#define TLLM_CHECK_ERROR_FMT(cond, ...) TLLM_CHECK_ERROR(cond, __VA_ARGS__) #define TLLM_CHECK_WARNING(cond, ...) \ if (!(cond)) \ @@ -100,18 +100,24 @@ struct GemmOptions tg::Dtype mDtypeC{tg::Dtype::Void}; // Whether to enable early exit. bool mEnablesEarlyExit{false}; + // Whether to enable early exit. + bool mEnablesDelayedEarlyExit{false}; // Whether to enable the global PTX knobs for guiding the compiler optimizations. bool mEnablesGlobalPtxKnobs{true}; // Tile size for the epilogue in M dimension. int mEpilogueTileM{128}; // Tile size for the epilogue in N dimension. int mEpilogueTileN{32}; - // Whether A or B load task triggers grid dependency controls. - bool mGridDepTriggerA{false}; - // Whether we should trigger a grid dependency. - bool mGridTriggerSecondary{false}; - // Whether we should wait on a grid dependency. - bool mGridWaitForPrimary{false}; + // Whether load task A triggers the next grid. + bool mGridTriggerSecondaryA{false}; + // Whether load task B triggers the next grid. + bool mGridTriggerSecondaryB{false}; + // Whether the loads that check for an early exit should wait on a grid dependency. + bool mGridWaitForPrimaryEarlyExit{true}; + // Whether the load of A should wait on a grid dependency. + bool mGridWaitForPrimaryA{true}; + // Whether the load of B should wait on a grid dependency. + bool mGridWaitForPrimaryB{true}; // Whether to hoist the mbarrier try_waits (e.g., mma.prodAcq, smemAb.consWait) in the MMA task. bool mHoistMmaTaskTryWaits{false}; // The K dimension of GEMM. @@ -139,10 +145,15 @@ struct GemmOptions int mNumSlicesForSliceK{1}; // The depth of the mainloop pipeline. int mNumStages{2}; - // The depth of the mma pipeline. + // The depth of the mma pipeline. Equals numStagesMmaWithinWorkTile * numStagesMmaAcrossWorkTile. int mNumStagesMma{1}; - // The depth of the work id pipeline. - int mNumStagesWorkId{2}; + // The depth of the mma pipeline within work tile. Only GmemC classes with "WithAccInReg" suffix + // are allowed to be greater than 1. + int mNumStagesMmaWithinWorkTile{-1}; + // The depth of the mma pipeline across work tiles in the persistent loop. + int mNumStagesMmaAcrossWorkTile{-1}; + // The depth of the work id pipeline and the work throttle pipeline. + int mNumStagesWorkId{3}; // Whether to output debug tensors. bool mOutputDebugTensors{false}; // Reorder rows/cols in the A matrix for the better memory accesses in the M-major epilogue. @@ -163,6 +174,10 @@ struct GemmOptions bool mUseUnrollLoop2xForMma{true}; // Use custom MMA schedule optimized for low-latency. bool mUseCustomMmaSchedule{false}; + // The purpose of hoisting trywaits is to opportunistically peek at the availability of the next + // k-block. It benefits when the next k-block is already available and thus sustaining the + // momentum, but it adds latency to the first k-block for smaller k-loop. + bool mUseHoistTryWaitForCustomMmaSchedule{false}; // Use DeepSeek Fp8. bool mUseDeepSeekFp8{false}; // Apply per-token scales from A @@ -256,12 +271,15 @@ inline std::string dumpOptions(GemmOptions const& options) << "trtllm::gen::Dtype(" << static_cast(options.mDtypeC) << ")" << "," << std::endl; ss << "mEnablesEarlyExit=" << options.mEnablesEarlyExit << "," << std::endl; + ss << "mEnablesDelayedEarlyExit=" << options.mEnablesDelayedEarlyExit << "," << std::endl; ss << "mEnablesGlobalPtxKnobs=" << options.mEnablesGlobalPtxKnobs << "," << std::endl; ss << "mEpilogueTileM=" << options.mEpilogueTileM << "," << std::endl; ss << "mEpilogueTileN=" << options.mEpilogueTileN << "," << std::endl; - ss << "mGridDepTriggerA=" << options.mGridDepTriggerA << "," << std::endl; - ss << "mGridTriggerSecondary=" << options.mGridTriggerSecondary << "," << std::endl; - ss << "mGridWaitForPrimary=" << options.mGridWaitForPrimary << "," << std::endl; + ss << "mGridTriggerSecondaryA=" << options.mGridTriggerSecondaryA << "," << std::endl; + ss << "mGridTriggerSecondaryB=" << options.mGridTriggerSecondaryB << "," << std::endl; + ss << "mGridWaitForPrimaryEarlyExit=" << options.mGridWaitForPrimaryEarlyExit << "," << std::endl; + ss << "mGridWaitForPrimaryA=" << options.mGridWaitForPrimaryA << "," << std::endl; + ss << "mGridWaitForPrimaryB=" << options.mGridWaitForPrimaryB << "," << std::endl; ss << "mHoistMmaTaskTryWaits=" << options.mHoistMmaTaskTryWaits << "," << std::endl; ss << "mK=" << options.mK << "," << std::endl; ss << "mKernelTraits={}" @@ -276,6 +294,8 @@ inline std::string dumpOptions(GemmOptions const& options) ss << "mNumSlicesForSliceK=" << options.mNumSlicesForSliceK << "," << std::endl; ss << "mNumStages=" << options.mNumStages << "," << std::endl; ss << "mNumStagesMma=" << options.mNumStagesMma << "," << std::endl; + ss << "mNumStagesMmaWithinWorkTile=" << options.mNumStagesMmaWithinWorkTile << "," << std::endl; + ss << "mNumStagesMmaAcrossWorkTile=" << options.mNumStagesMmaAcrossWorkTile << "," << std::endl; ss << "mNumStagesWorkId=" << options.mNumStagesWorkId << "," << std::endl; ss << "mOutputDebugTensors=" << options.mOutputDebugTensors << "," << std::endl; ss << "mUseShuffledMatrixA=" << options.mUseShuffledMatrixA << "," << std::endl; @@ -289,6 +309,7 @@ inline std::string dumpOptions(GemmOptions const& options) ss << "mTileK=" << options.mTileK << "," << std::endl; ss << "mUseUnrollLoop2xForMma=" << options.mUseUnrollLoop2xForMma << "," << std::endl; ss << "mUseCustomMmaSchedule=" << options.mUseCustomMmaSchedule << "," << std::endl; + ss << "mUseHoistTryWaitForCustomMmaSchedule=" << options.mUseHoistTryWaitForCustomMmaSchedule << "," << std::endl; ss << "mUseDeepSeekFp8=" << options.mUseDeepSeekFp8 << "," << std::endl; ss << "mUsePerTokenSfA=" << options.mUsePerTokenSfA << "," << std::endl; ss << "mUsePerTokenSfB=" << options.mUsePerTokenSfB << "," << std::endl; @@ -350,20 +371,41 @@ inline bool checkAndUpdateGemmOptions( } } - // NvFp4 constraints - if (options.mDtypeElt == tg::Dtype::E2m1) + // Constraints for NvFp4 and MxFp8. + if ((options.mDtypeElt == tg::Dtype::E2m1 || options.mDtypeElt == tg::Dtype::MxE4m3 + || options.mDtypeC == tg::Dtype::MxE4m3) + && options.mMmaM != 128) + { + // MMA M must be 128 when the input uses block scaling, or when the output is an Mx format. + int newTileM = 128 * divUp(options.mTileM, 128); + TLLM_LOG_WARNING("Unsupported MmaM (", options.mMmaM, ") for dtypeElt=", gemm::toString(options.mDtypeElt), + ", dtypeC=", gemm::toString(options.mDtypeC), ". Setting MmaM to 128 and TileM to ", newTileM); + if (updateOptions) + { + options.mMmaM = 128; + options.mTileM = newTileM; + } + else + { + return false; + } + } + if (options.mDtypeElt == tg::Dtype::E2m1 || options.mDtypeElt == tg::Dtype::MxE4m3) { - TLLM_CHECK_ERROR(isBlackwell, "FP4 is only supported on Blackwell"); + TLLM_CHECK_ERROR(isBlackwell, "Block scaling is only supported on Blackwell"); + TLLM_CHECK_ERROR(options.mSfLayoutB == tg::SfLayout::R128c4 || options.mSfLayoutB == tg::SfLayout::R8c4, "Only the 128x4 and 8x4 SF layouts are supported for B, got ", tg::sfLayoutToString(options.mSfLayoutB)); - if (options.mMmaK != 64) + + int const mmaK = (options.mDtypeElt == tg::Dtype::E2m1) ? 64 : 32; + if (options.mMmaK != mmaK) { - int newTileK = 64 * divUp(options.mTileK, 64); + int newTileK = mmaK * divUp(options.mTileK, mmaK); TLLM_LOG_WARNING("Unsupported MmaK (", options.mMmaK, ") for ", gemm::toString(options.mDtypeElt), - ". Setting MmaK to 64 and TileK to ", newTileK); + ". Setting MmaK to ", mmaK, " and TileK to ", newTileK); if (updateOptions) { - options.mMmaK = 64; + options.mMmaK = mmaK; options.mTileK = newTileK; } else @@ -371,33 +413,22 @@ inline bool checkAndUpdateGemmOptions( return false; } } - if (options.mMmaM != 128) - { - int newTileM = 128 * divUp(options.mTileM, 128); - TLLM_LOG_WARNING("Unsupported MmaM (", options.mMmaM, ") for ", gemm::toString(options.mDtypeElt), - ". Setting MmaM to 128 and TileM to ", newTileM); - if (updateOptions) - { - options.mMmaM = 128; - options.mTileM = newTileM; - } - else - { - return false; - } - } + // TileN must be a multiple of the number of rows per SF tile. - // TODO: relax this restriction. int const numSfTileRowsB = options.mSfLayoutB == tg::SfLayout::R128c4 ? 128 : 8; TLLM_CHECK_ERROR(options.mTileN % numSfTileRowsB == 0, "TileN (", options.mTileN, ") must be a multiple of ", numSfTileRowsB, " for B SF layout ", tg::sfLayoutToString(options.mSfLayoutB)); // The MMA N may only be smaller than 64 if it is equal to the tile N. TLLM_CHECK_ERROR(options.mMmaN >= 64 || options.mMmaN == options.mTileN, "MmaN (", options.mMmaN, ") must be >= 64 or equal to TileN (", options.mTileN, ") for ", gemm::toString(options.mDtypeElt)); + + int numEltsPerSf = tg::dtypeNumEltsPerSf(options.mDtypeElt); + TLLM_CHECK_ERROR(options.mTileK % (4 * numEltsPerSf) == 0, "TileK (", options.mTileK, + ") must be a multiple of ", (4 * numEltsPerSf), " for type ", gemm::toString(options.mDtypeElt)); } - if (options.mDtypeC == tg::Dtype::E2m1) + if (options.mDtypeC == tg::Dtype::E2m1 || options.mDtypeC == tg::Dtype::MxE4m3) { - TLLM_CHECK_ERROR(isBlackwell, "FP4 is only supported on Blackwell"); + TLLM_CHECK_ERROR(isBlackwell, "Block scaling is only supported on Blackwell"); TLLM_CHECK_ERROR(options.mSfLayoutC == tg::SfLayout::R128c4 || options.mSfLayoutC == tg::SfLayout::R8c4, "Only the 128x4 and 8x4 SF layouts are supported for C."); @@ -406,9 +437,11 @@ inline bool checkAndUpdateGemmOptions( numSfTileRowsC, " for C SF layout ", tg::sfLayoutToString(options.mSfLayoutC)); int const hiddenDim = options.mTransposeMmaOutput ? options.mM : options.mN; - TLLM_CHECK_ERROR(hiddenDim % 64 == 0, "Hidden dim (", hiddenDim, ") must be a multiple of 64 for FP4 outputs."); + int const hiddenGranularity = 4 * tg::dtypeNumEltsPerSf(options.mDtypeC); + TLLM_CHECK_ERROR(hiddenDim % hiddenGranularity == 0, "Hidden dim (", hiddenDim, ") must be a multiple of ", + hiddenGranularity, " for block-scaled outputs."); TLLM_CHECK_ERROR(!options.mTransposeMmaOutput || options.mUseShuffledMatrixA, - "Transposing FP4 outputs requires shuffled A."); + "Transposing block-scaled outputs requires shuffled A."); } // If dtypeC is unspecified (Dtype::Void), assign to the input dtype. @@ -589,10 +622,61 @@ inline bool checkAndUpdateGemmOptions( "CGA size must be equal to the number of slices in split-k"); } + // Maps numStagesMma to (stagesWithinWorkTile, stagesAcrossWorkTile) if not already set. + // If (-1, -1) -> (numStagesMma / min(2, numStagesMma), min(2, numStagesMma)) + // If ( m, -1) -> (m, numStagesMma / m) + // If (-1, n) -> (numStagesMma / n, n) + if (options.mNumStagesMmaWithinWorkTile == -1 && options.mNumStagesMmaAcrossWorkTile == -1) + { + if (updateOptions) + { + options.mNumStagesMmaAcrossWorkTile = std::min(2, options.mNumStagesMma); + options.mNumStagesMmaWithinWorkTile = options.mNumStagesMma / options.mNumStagesMmaAcrossWorkTile; + } + else + { + return false; + } + } + else if (options.mNumStagesMmaWithinWorkTile == -1) + { + if (updateOptions) + { + options.mNumStagesMmaWithinWorkTile = options.mNumStagesMma / options.mNumStagesMmaAcrossWorkTile; + } + else + { + return false; + } + } + else if (options.mNumStagesMmaAcrossWorkTile == -1) + { + if (updateOptions) + { + options.mNumStagesMmaAcrossWorkTile = options.mNumStagesMma / options.mNumStagesMmaWithinWorkTile; + } + else + { + return false; + } + } + // Check mma stages. + TLLM_CHECK_ERROR_FMT( + options.mNumStagesMmaWithinWorkTile * options.mNumStagesMmaAcrossWorkTile == options.mNumStagesMma + && options.mNumStagesMmaAcrossWorkTile <= 2, + "Condition numStagesMmaWithinWorkTile (%d) * numStagesMmaAcrossWorkTile " + "(%d) == numStagesMma (%d) && numStagesMmaAcrossWorkTile (%d) <= 2 must be " + "satisfied. Check arguments.", + options.mNumStagesMmaWithinWorkTile, options.mNumStagesMmaAcrossWorkTile, options.mNumStagesMma, + options.mNumStagesMmaAcrossWorkTile); + // Mma stage must be 1 for pre-Hopper. + TLLM_CHECK_ERROR( + isBlackwell || options.mNumStagesMma == 1, "Mma stage must be 1 for pre-Hopper. Found ", options.mNumStagesMma); // DeepSeek Fp8 if (!options.mUseDeepSeekFp8) { - TLLM_CHECK_ERROR(options.mNumStagesMma == 1, "Non-DeepSeekFp8 requires numStagesMma == 1"); + TLLM_CHECK_ERROR( + options.mNumStagesMmaWithinWorkTile == 1, "Non-DeepSeekFp8 requires numStagesMmaWithinWorkTile == 1"); } if (options.mUseDeepSeekFp8) { @@ -610,7 +694,8 @@ inline bool checkAndUpdateGemmOptions( auto hiddenDimPerEpilogueTile = options.mTransposeMmaOutput ? options.mEpilogueTileM : options.mEpilogueTileN; auto hiddenDimPerMma = options.mTransposeMmaOutput ? options.mMmaM : options.mMmaN; auto hiddenDimName = options.mTransposeMmaOutput ? "M" : "N"; - TLLM_CHECK_WARNING(options.mNumStagesMma > 1, "DeepSeekFp8 recommends >1 MMA accumulator stages."); + TLLM_CHECK_WARNING(options.mNumStagesMmaWithinWorkTile > 1, + "DeepSeekFp8 recommends setting \"-numStagesMmaWithinWorkTile 2\"."); // Update the number of stages of the MMA accumulator pipeline. TODO: enable by default for // deepseek. // options.mNumStagesMma = 2; @@ -708,6 +793,42 @@ inline bool checkAndUpdateGemmOptions( } } + if (options.mEnablesDelayedEarlyExit && options.mEnablesEarlyExit) + { + TLLM_LOG_WARNING( + "Only one of early exit and delayed early exit should be enabled. Disabling " + "delayed early exit"); + if (updateOptions) + { + options.mEnablesDelayedEarlyExit = false; + } + else + { + return false; + } + } + + // This check prevents the triggering of the secondary (PREEXIT) from executing before the wait + // for primary (ACQBULK). This could lead to the following confusing situation, which we want to + // avoid: + // + // Kernel 3 is written with the assumption that it can read the output of + // kernel 1 *without* ACQBULK and the output of kernel 2 *with* ACQBULK. + // However, when we allow PREEXIT and ACQBULK to be executed out of order, + // this is not guaranteed. + // + // Time: ----> + // + // Kernel 1: ----PREEXIT-----------FLUSH + // Kernel 2: -------PREEXIT----ACQBULK---FLUSH + // Kernel 3: Warp 0: ---- (!) Output of 1,2 is not yet visible ----------------------- + // Warp 1: ---- (!) We normally assume that 1 is visible is not yet visible- + // Warp 2: -------------------ACQBULK-- Kernel 1,2 output visible ---------- + TLLM_CHECK_ERROR((options.mGridWaitForPrimaryA || !options.mGridTriggerSecondaryA), + "A: If a task triggers a secondary kernel, it must also wait for primary kernel."); + TLLM_CHECK_ERROR((options.mGridWaitForPrimaryB || !options.mGridTriggerSecondaryB), + "B: If a task triggers a secondary kernel, it must also wait for primary kernel."); + if (updateOptions) { // Init kernel traits. diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/KernelMetaInfo.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/KernelMetaInfo.h index 0cd8dd1a0917..02507702b652 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/KernelMetaInfo.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/KernelMetaInfo.h @@ -25,7 +25,7 @@ namespace kernels { // clang-format off -#define TLLM_GEN_COMMIT "3d9075a-dirty" +#define TLLM_GEN_COMMIT "23d32a5" #define TLLM_GEN_EXPORT_VERSION "0.0" static constexpr size_t tllmGenGemmListLen = 12; @@ -71,12 +71,15 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mDtypeElt */ trtllm::gen::Dtype(17826819) , /* mDtypeC */ trtllm::gen::Dtype(1052672) , /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 , /* mEnablesGlobalPtxKnobs */ 1 , /* mEpilogueTileM */ 128 , /* mEpilogueTileN */ 128 -, /* mGridDepTriggerA */ 0 -, /* mGridTriggerSecondary */ 0 -, /* mGridWaitForPrimary */ 0 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 0 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 1 +, /* mGridWaitForPrimaryB */ 1 , /* mHoistMmaTaskTryWaits */ 0 , /* mK */ 2048 , /* mKernelTraits */ {} @@ -90,7 +93,9 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mNumSlicesForSliceK */ 1 , /* mNumStages */ 3 , /* mNumStagesMma */ 1 -, /* mNumStagesWorkId */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 , /* mOutputDebugTensors */ 0 , /* mUseShuffledMatrixA */ 0 , /* mSliceK */ 0 @@ -101,6 +106,7 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mTileK */ 256 , /* mUseUnrollLoop2xForMma */ 1 , /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 , /* mUseDeepSeekFp8 */ 0 , /* mUsePerTokenSfA */ 0 , /* mUsePerTokenSfB */ 0 @@ -120,12 +126,15 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mDtypeElt */ trtllm::gen::Dtype(1050630) , /* mDtypeC */ trtllm::gen::Dtype(1052672) , /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 , /* mEnablesGlobalPtxKnobs */ 1 , /* mEpilogueTileM */ 128 , /* mEpilogueTileN */ 128 -, /* mGridDepTriggerA */ 0 -, /* mGridTriggerSecondary */ 0 -, /* mGridWaitForPrimary */ 0 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 0 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 1 +, /* mGridWaitForPrimaryB */ 1 , /* mHoistMmaTaskTryWaits */ 0 , /* mK */ 2048 , /* mKernelTraits */ {} @@ -139,7 +148,9 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mNumSlicesForSliceK */ 1 , /* mNumStages */ 2 , /* mNumStagesMma */ 1 -, /* mNumStagesWorkId */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 , /* mOutputDebugTensors */ 0 , /* mUseShuffledMatrixA */ 0 , /* mSliceK */ 0 @@ -150,6 +161,7 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mTileK */ 256 , /* mUseUnrollLoop2xForMma */ 1 , /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 , /* mUseDeepSeekFp8 */ 0 , /* mUsePerTokenSfA */ 0 , /* mUsePerTokenSfB */ 0 @@ -161,7 +173,7 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mSfLayoutC */ trtllm::gen::SfLayout(3) , /* mTileScheduler */ gemm::TileScheduler(0) }, gemm::SmVersion::Sm100a }, -{GemmKernel_Bfloat16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin, GemmKernel_Bfloat16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin_len, 216064, "gemmKernel_Bfloat16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a", 224, { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +{GemmKernel_Bfloat16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin, GemmKernel_Bfloat16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin_len, 217088, "gemmKernel_Bfloat16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a", 224, { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) , /* mClusterDimX */ 1 , /* mClusterDimY */ 1 , /* mClusterDimZ */ 1 @@ -169,12 +181,15 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mDtypeElt */ trtllm::gen::Dtype(1050630) , /* mDtypeC */ trtllm::gen::Dtype(1052672) , /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 , /* mEnablesGlobalPtxKnobs */ 1 , /* mEpilogueTileM */ 128 , /* mEpilogueTileN */ 8 -, /* mGridDepTriggerA */ 1 -, /* mGridTriggerSecondary */ 1 -, /* mGridWaitForPrimary */ 1 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 , /* mHoistMmaTaskTryWaits */ 0 , /* mK */ 2048 , /* mKernelTraits */ {} @@ -188,7 +203,9 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mNumSlicesForSliceK */ 1 , /* mNumStages */ 3 , /* mNumStagesMma */ 1 -, /* mNumStagesWorkId */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 , /* mOutputDebugTensors */ 0 , /* mUseShuffledMatrixA */ 1 , /* mSliceK */ 0 @@ -199,6 +216,7 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mTileK */ 512 , /* mUseUnrollLoop2xForMma */ 1 , /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 , /* mUseDeepSeekFp8 */ 0 , /* mUsePerTokenSfA */ 0 , /* mUsePerTokenSfB */ 0 @@ -210,7 +228,7 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mSfLayoutC */ trtllm::gen::SfLayout(3) , /* mTileScheduler */ gemm::TileScheduler(0) }, gemm::SmVersion::Sm100a }, -{GemmKernel_Bfloat16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin, GemmKernel_Bfloat16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin_len, 214016, "gemmKernel_Bfloat16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a", 224, { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +{GemmKernel_Bfloat16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin, GemmKernel_Bfloat16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin_len, 215040, "gemmKernel_Bfloat16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a", 224, { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) , /* mClusterDimX */ 1 , /* mClusterDimY */ 1 , /* mClusterDimZ */ 2 @@ -218,12 +236,15 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mDtypeElt */ trtllm::gen::Dtype(1050630) , /* mDtypeC */ trtllm::gen::Dtype(1052672) , /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 , /* mEnablesGlobalPtxKnobs */ 1 , /* mEpilogueTileM */ 128 , /* mEpilogueTileN */ 8 -, /* mGridDepTriggerA */ 1 -, /* mGridTriggerSecondary */ 1 -, /* mGridWaitForPrimary */ 1 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 , /* mHoistMmaTaskTryWaits */ 0 , /* mK */ 2048 , /* mKernelTraits */ {} @@ -237,7 +258,9 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mNumSlicesForSliceK */ 1 , /* mNumStages */ 3 , /* mNumStagesMma */ 1 -, /* mNumStagesWorkId */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 , /* mOutputDebugTensors */ 0 , /* mUseShuffledMatrixA */ 1 , /* mSliceK */ 0 @@ -248,6 +271,7 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mTileK */ 512 , /* mUseUnrollLoop2xForMma */ 1 , /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 , /* mUseDeepSeekFp8 */ 0 , /* mUsePerTokenSfA */ 0 , /* mUsePerTokenSfB */ 0 @@ -267,12 +291,15 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mDtypeElt */ trtllm::gen::Dtype(1050630) , /* mDtypeC */ trtllm::gen::Dtype(1050630) , /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 , /* mEnablesGlobalPtxKnobs */ 1 , /* mEpilogueTileM */ 128 , /* mEpilogueTileN */ 128 -, /* mGridDepTriggerA */ 0 -, /* mGridTriggerSecondary */ 0 -, /* mGridWaitForPrimary */ 0 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 0 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 1 +, /* mGridWaitForPrimaryB */ 1 , /* mHoistMmaTaskTryWaits */ 0 , /* mK */ 2048 , /* mKernelTraits */ {} @@ -286,7 +313,9 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mNumSlicesForSliceK */ 1 , /* mNumStages */ 3 , /* mNumStagesMma */ 1 -, /* mNumStagesWorkId */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 , /* mOutputDebugTensors */ 0 , /* mUseShuffledMatrixA */ 0 , /* mSliceK */ 0 @@ -297,6 +326,7 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mTileK */ 256 , /* mUseUnrollLoop2xForMma */ 1 , /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 , /* mUseDeepSeekFp8 */ 0 , /* mUsePerTokenSfA */ 0 , /* mUsePerTokenSfB */ 0 @@ -308,7 +338,7 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mSfLayoutC */ trtllm::gen::SfLayout(3) , /* mTileScheduler */ gemm::TileScheduler(0) }, gemm::SmVersion::Sm100a }, -{GemmKernel_E4m3_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin, GemmKernel_E4m3_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin_len, 215040, "gemmKernel_E4m3_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a", 224, { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +{GemmKernel_E4m3_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin, GemmKernel_E4m3_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin_len, 216064, "gemmKernel_E4m3_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a", 224, { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) , /* mClusterDimX */ 1 , /* mClusterDimY */ 1 , /* mClusterDimZ */ 1 @@ -316,12 +346,15 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mDtypeElt */ trtllm::gen::Dtype(1050630) , /* mDtypeC */ trtllm::gen::Dtype(1050630) , /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 , /* mEnablesGlobalPtxKnobs */ 1 , /* mEpilogueTileM */ 128 , /* mEpilogueTileN */ 8 -, /* mGridDepTriggerA */ 1 -, /* mGridTriggerSecondary */ 1 -, /* mGridWaitForPrimary */ 1 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 , /* mHoistMmaTaskTryWaits */ 0 , /* mK */ 2048 , /* mKernelTraits */ {} @@ -335,7 +368,9 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mNumSlicesForSliceK */ 1 , /* mNumStages */ 3 , /* mNumStagesMma */ 1 -, /* mNumStagesWorkId */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 , /* mOutputDebugTensors */ 0 , /* mUseShuffledMatrixA */ 1 , /* mSliceK */ 0 @@ -346,6 +381,7 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mTileK */ 512 , /* mUseUnrollLoop2xForMma */ 1 , /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 , /* mUseDeepSeekFp8 */ 0 , /* mUsePerTokenSfA */ 0 , /* mUsePerTokenSfB */ 0 @@ -357,7 +393,7 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mSfLayoutC */ trtllm::gen::SfLayout(3) , /* mTileScheduler */ gemm::TileScheduler(0) }, gemm::SmVersion::Sm100a }, -{GemmKernel_E4m3_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin, GemmKernel_E4m3_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin_len, 214016, "gemmKernel_E4m3_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a", 224, { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +{GemmKernel_E4m3_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin, GemmKernel_E4m3_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin_len, 215040, "gemmKernel_E4m3_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a", 224, { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) , /* mClusterDimX */ 1 , /* mClusterDimY */ 1 , /* mClusterDimZ */ 2 @@ -365,12 +401,15 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mDtypeElt */ trtllm::gen::Dtype(1050630) , /* mDtypeC */ trtllm::gen::Dtype(1050630) , /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 , /* mEnablesGlobalPtxKnobs */ 1 , /* mEpilogueTileM */ 128 , /* mEpilogueTileN */ 8 -, /* mGridDepTriggerA */ 1 -, /* mGridTriggerSecondary */ 1 -, /* mGridWaitForPrimary */ 1 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 , /* mHoistMmaTaskTryWaits */ 0 , /* mK */ 2048 , /* mKernelTraits */ {} @@ -384,7 +423,9 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mNumSlicesForSliceK */ 1 , /* mNumStages */ 3 , /* mNumStagesMma */ 1 -, /* mNumStagesWorkId */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 , /* mOutputDebugTensors */ 0 , /* mUseShuffledMatrixA */ 1 , /* mSliceK */ 0 @@ -395,6 +436,7 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mTileK */ 512 , /* mUseUnrollLoop2xForMma */ 1 , /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 , /* mUseDeepSeekFp8 */ 0 , /* mUsePerTokenSfA */ 0 , /* mUsePerTokenSfB */ 0 @@ -414,12 +456,15 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mDtypeElt */ trtllm::gen::Dtype(17826819) , /* mDtypeC */ trtllm::gen::Dtype(1052680) , /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 , /* mEnablesGlobalPtxKnobs */ 1 , /* mEpilogueTileM */ 128 , /* mEpilogueTileN */ 128 -, /* mGridDepTriggerA */ 0 -, /* mGridTriggerSecondary */ 0 -, /* mGridWaitForPrimary */ 0 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 0 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 1 +, /* mGridWaitForPrimaryB */ 1 , /* mHoistMmaTaskTryWaits */ 0 , /* mK */ 2048 , /* mKernelTraits */ {} @@ -433,7 +478,9 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mNumSlicesForSliceK */ 1 , /* mNumStages */ 3 , /* mNumStagesMma */ 1 -, /* mNumStagesWorkId */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 , /* mOutputDebugTensors */ 0 , /* mUseShuffledMatrixA */ 0 , /* mSliceK */ 0 @@ -444,6 +491,7 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mTileK */ 256 , /* mUseUnrollLoop2xForMma */ 1 , /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 , /* mUseDeepSeekFp8 */ 0 , /* mUsePerTokenSfA */ 0 , /* mUsePerTokenSfB */ 0 @@ -463,12 +511,15 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mDtypeElt */ trtllm::gen::Dtype(1050630) , /* mDtypeC */ trtllm::gen::Dtype(1052680) , /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 , /* mEnablesGlobalPtxKnobs */ 1 , /* mEpilogueTileM */ 128 , /* mEpilogueTileN */ 128 -, /* mGridDepTriggerA */ 0 -, /* mGridTriggerSecondary */ 0 -, /* mGridWaitForPrimary */ 0 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 0 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 1 +, /* mGridWaitForPrimaryB */ 1 , /* mHoistMmaTaskTryWaits */ 0 , /* mK */ 2048 , /* mKernelTraits */ {} @@ -482,7 +533,9 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mNumSlicesForSliceK */ 1 , /* mNumStages */ 2 , /* mNumStagesMma */ 1 -, /* mNumStagesWorkId */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 , /* mOutputDebugTensors */ 0 , /* mUseShuffledMatrixA */ 0 , /* mSliceK */ 0 @@ -493,6 +546,7 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mTileK */ 256 , /* mUseUnrollLoop2xForMma */ 1 , /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 , /* mUseDeepSeekFp8 */ 0 , /* mUsePerTokenSfA */ 0 , /* mUsePerTokenSfB */ 0 @@ -504,7 +558,7 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mSfLayoutC */ trtllm::gen::SfLayout(3) , /* mTileScheduler */ gemm::TileScheduler(0) }, gemm::SmVersion::Sm100a }, -{GemmKernel_Fp16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin, GemmKernel_Fp16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin_len, 216064, "gemmKernel_Fp16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a", 224, { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +{GemmKernel_Fp16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin, GemmKernel_Fp16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin_len, 217088, "gemmKernel_Fp16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a", 224, { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) , /* mClusterDimX */ 1 , /* mClusterDimY */ 1 , /* mClusterDimZ */ 1 @@ -512,12 +566,15 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mDtypeElt */ trtllm::gen::Dtype(1050630) , /* mDtypeC */ trtllm::gen::Dtype(1052680) , /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 , /* mEnablesGlobalPtxKnobs */ 1 , /* mEpilogueTileM */ 128 , /* mEpilogueTileN */ 8 -, /* mGridDepTriggerA */ 1 -, /* mGridTriggerSecondary */ 1 -, /* mGridWaitForPrimary */ 1 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 , /* mHoistMmaTaskTryWaits */ 0 , /* mK */ 2048 , /* mKernelTraits */ {} @@ -531,7 +588,9 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mNumSlicesForSliceK */ 1 , /* mNumStages */ 3 , /* mNumStagesMma */ 1 -, /* mNumStagesWorkId */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 , /* mOutputDebugTensors */ 0 , /* mUseShuffledMatrixA */ 1 , /* mSliceK */ 0 @@ -542,6 +601,7 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mTileK */ 512 , /* mUseUnrollLoop2xForMma */ 1 , /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 , /* mUseDeepSeekFp8 */ 0 , /* mUsePerTokenSfA */ 0 , /* mUsePerTokenSfB */ 0 @@ -553,7 +613,7 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mSfLayoutC */ trtllm::gen::SfLayout(3) , /* mTileScheduler */ gemm::TileScheduler(0) }, gemm::SmVersion::Sm100a }, -{GemmKernel_Fp16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin, GemmKernel_Fp16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin_len, 214016, "gemmKernel_Fp16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a", 224, { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +{GemmKernel_Fp16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin, GemmKernel_Fp16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin_len, 215040, "gemmKernel_Fp16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a", 224, { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) , /* mClusterDimX */ 1 , /* mClusterDimY */ 1 , /* mClusterDimZ */ 2 @@ -561,12 +621,15 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mDtypeElt */ trtllm::gen::Dtype(1050630) , /* mDtypeC */ trtllm::gen::Dtype(1052680) , /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 , /* mEnablesGlobalPtxKnobs */ 1 , /* mEpilogueTileM */ 128 , /* mEpilogueTileN */ 8 -, /* mGridDepTriggerA */ 1 -, /* mGridTriggerSecondary */ 1 -, /* mGridWaitForPrimary */ 1 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 , /* mHoistMmaTaskTryWaits */ 0 , /* mK */ 2048 , /* mKernelTraits */ {} @@ -580,7 +643,9 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mNumSlicesForSliceK */ 1 , /* mNumStages */ 3 , /* mNumStagesMma */ 1 -, /* mNumStagesWorkId */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 , /* mOutputDebugTensors */ 0 , /* mUseShuffledMatrixA */ 1 , /* mSliceK */ 0 @@ -591,6 +656,7 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mTileK */ 512 , /* mUseUnrollLoop2xForMma */ 1 , /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 , /* mUseDeepSeekFp8 */ 0 , /* mUsePerTokenSfA */ 0 , /* mUsePerTokenSfB */ 0 @@ -610,12 +676,15 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mDtypeElt */ trtllm::gen::Dtype(17826819) , /* mDtypeC */ trtllm::gen::Dtype(1056777) , /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 , /* mEnablesGlobalPtxKnobs */ 1 , /* mEpilogueTileM */ 128 , /* mEpilogueTileN */ 128 -, /* mGridDepTriggerA */ 0 -, /* mGridTriggerSecondary */ 0 -, /* mGridWaitForPrimary */ 0 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 0 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 1 +, /* mGridWaitForPrimaryB */ 1 , /* mHoistMmaTaskTryWaits */ 0 , /* mK */ 2048 , /* mKernelTraits */ {} @@ -629,7 +698,9 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mNumSlicesForSliceK */ 1 , /* mNumStages */ 3 , /* mNumStagesMma */ 1 -, /* mNumStagesWorkId */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 , /* mOutputDebugTensors */ 0 , /* mUseShuffledMatrixA */ 0 , /* mSliceK */ 0 @@ -640,6 +711,7 @@ static const gemm::GemmConfig tllmGenGemmList[] = { , /* mTileK */ 256 , /* mUseUnrollLoop2xForMma */ 1 , /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 , /* mUseDeepSeekFp8 */ 0 , /* mUsePerTokenSfA */ 0 , /* mUsePerTokenSfB */ 0 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/KernelParams.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/KernelParams.h index 4d19c5cabe71..f00d4b189a3a 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/KernelParams.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/KernelParams.h @@ -75,12 +75,14 @@ struct KernelParams // Must be setup using gemm::buildSfTmaDescriptor with shapes and strides from // makeTmaShapeStrideSfAb. // The layout of scaling factors for A is always R128c4 - // M must be a multiple of 128. - // K must be a multiple of 64. - // The "logical" shape is: [M, K / 16]. - // The R128c4 layout is: [M / 128, K / 16 / 4, 512]. - // The shape we use for TMA is: [M / 128, K / 16 / 4, 2, 256]. - // Dtype is Dtype::E4m3. + // + // Let P be the number of elements per SF. P=16 for NvFp4, P=32 for Mx formats. + // K must be a multiple of 4P. + // The "logical" shape is: [M, K / P]. + // The R128c4 layout is: [⌈M / 128⌉, K / P / 4, 512]. + // The shape we use for TMA is: [⌈M / 128⌉, K / P / 4, 2, 256]. + // + // Dtype is Dtype::E4m3 for NvFp4, Dtype::UE8m0 for Mx formats. CUtensorMap tmaSfA; // TMA descriptor for the block scaling factors for B, for MxFp{4,8} and NvFp4 formats. @@ -88,22 +90,21 @@ struct KernelParams // makeTmaShapeStrideSfAb. // The layout of scaling factors for B is controlled by options.mSfLayoutB. // - // The "logical" shape is: [N, K / 16] + // Let P be the number of elements per SF. P=16 for NvFp4, P=32 for Mx formats. + // The "logical" shape is: [N, K / P] // // If the layout is R128c4, - // N must be a multiple of 128. - // K must be a multiple of 64. - // The R128c4 layout is: [N / 128, K / 16 / 4, 512] - // The shape we use for TMA is: [N / 128, K / 16 / 4, 2, 256] + // K must be a multiple of 4P. + // The R128c4 layout is: [⌈N / 128⌉, K / P / 4, 512] + // The shape we use for TMA is: [⌈N / 128⌉, K / P / 4, 2, 256] // // If the layout is R8c4, - // N must be a multiple of 8. - // K must be a multiple of 64. - // The R8c4 layout is: [N / 8, K / 16 / 4, 32] - // The shape we use for TMA is: [N / 8, K / 16 / 4 / repeats, repeats * 32] - // where repeats = min(tileK / 16 / 4, 8) + // K must be a multiple of 4P. + // The R8c4 layout is: [⌈N / 8⌉, K / P / 4, 32] + // The shape we use for TMA is: [⌈N / 8⌉, K / P / 4 / r, r * 32] + // where r = min(tileK / P / 4, 8) // - // Dtype is Dtype::E4m3. + // Dtype is Dtype::E4m3 for NvFp4, Dtype::UE8m0 for Mx formats. CUtensorMap tmaSfB; // The output matrix C. The data type is controlled by options.mDtypeC. @@ -317,7 +318,7 @@ struct KernelParams // The inner tile dimension. auto hiddenSizePerTile = options.mTileK; // Number of elements per scaling factor. - int32_t const numEltsPerSf = 16; + int32_t const numEltsPerSf = (options.mDtypeElt == tg::Dtype::E2m1) ? 16 : 32; switch (layout) { @@ -330,9 +331,9 @@ struct KernelParams // Additionally, we have to meet constraints of TMA that the box dimensions are less // than 256 and boxDim[0] is a multiple of 16B. // - // The "logical" tensor is: [outer, inner / numEltsPerSf] - // The aforementioned format is: [outer / 128, inner / numEltsPerSf / 4, 512] - // The shape we use for TMA is: [outer / 128, inner / numEltsPerSf / 4, 2, 256] + // The "logical" tensor is: [outer, inner / numEltsPerSf] + // The aforementioned format is: [⌈outer / 128⌉, inner / (4 * numEltsPerSf), 512] + // The shape we use for TMA is: [⌈outer / 128⌉, inner / (4 * numEltsPerSf), 2, 256] auto shape = std::vector{256, 2, static_cast(tg::ceilDiv(hiddenSize, numEltsPerSf * 4)), static_cast(tg::ceilDiv(numTokens, 128))}; @@ -357,11 +358,11 @@ struct KernelParams // // As the inner dimension (k) is required to be a multiple of the tile size, we // can reshape to use fewer read requests, if the tile dimensions allow. - // I.e., let's define repeats = min(hiddenSizePerTile / numEltsPerSf / 4, 8) + // I.e., let's define r = min(⌈hiddenSizePerTile / (numEltsPerSf * 4)⌉, 8) // - // The "logical" tensor is: [outer, inner / numEltsPerSf] - // The 8x4 SF layout is: [outer / 8, inner / numEltsPerSf / 4, 32] - // The TMA tensor shape is: [outer / 8, inner / numEltsPerSf / 4 / repeats, repeats * 32] + // The "logical" tensor is: [outer, inner / numEltsPerSf] + // The 8x4 SF layout is: [⌈outer / 128⌉, inner / (4 * numEltsPerSf), 32] + // The TMA tensor shape is: [⌈outer / 128⌉, inner / (4 * numEltsPerSf * r), r * 32] int const repeats = std::min(tg::ceilDiv(hiddenSizePerTile, numEltsPerSf * 4), 8); @@ -420,9 +421,9 @@ struct KernelParams const_cast(ptrB), /* swizzle */ !options.mSliceK); - if (options.mDtypeElt == tg::Dtype::E2m1) + if (options.mDtypeElt == tg::Dtype::E2m1 || options.mDtypeElt == tg::Dtype::MxE4m3) { - tg::Dtype const dTypeSf = tg::Dtype::E4m3; + tg::Dtype const dTypeSf = (options.mDtypeElt == tg::Dtype::E2m1) ? tg::Dtype::E4m3 : tg::Dtype::UE8m0; // Build TMA descriptor for gmem A block scaling factors. auto [shapeSfA, strideSfA, tileShapesSfA] diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/KernelTraits.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/KernelTraits.h index 2ed49f33e6b5..1c3d4581c497 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/KernelTraits.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/KernelTraits.h @@ -38,9 +38,11 @@ class MemAllocatorHelper MemAllocatorHelper() {} // Constructor to initialize chunk sizes, alignments, and reuse flags - MemAllocatorHelper(std::vector> const& sizes, std::vector const& reuse) + MemAllocatorHelper(std::vector> const& sizes, std::vector const& reuse, + std::vector const& names) : mNumBytesAndAlignmentPerSmemChunk(sizes) , mFirstChunkReuse(reuse) + , mSmemChunkNames(names) { } @@ -92,6 +94,23 @@ class MemAllocatorHelper return getOffsetBeforeChunk(static_cast(mNumBytesAndAlignmentPerSmemChunk.size())); } + // Returns the first chunk reuse flag for the ith chunk. + int getFirstChunkReuseFlag(int32_t ii) const + { + return mFirstChunkReuse[ii]; + } + + // Print the contents of this object. + void print() const + { + for (size_t ii = 0; ii < mNumBytesAndAlignmentPerSmemChunk.size(); ++ii) + { + printf("Chunk %zd %s: %d bytes, %d alignment, reuse %s, offset %d\n", ii, mSmemChunkNames[ii].c_str(), + mNumBytesAndAlignmentPerSmemChunk[ii].first, mNumBytesAndAlignmentPerSmemChunk[ii].second, + mFirstChunkReuse[ii] ? "true" : "false", getChunkOffset(ii)); + } + } + private: // Helper function to calculate padded size int32_t getSizePaddedToAlignment(int32_t size, int32_t alignment) const @@ -107,6 +126,8 @@ class MemAllocatorHelper std::vector> mNumBytesAndAlignmentPerSmemChunk; // Chunk reuse configuration. True at ith position means that ith chunk starts at smemOffset = 0. std::vector mFirstChunkReuse; + // Buffer names for inspection purposes. + std::vector mSmemChunkNames; }; //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -145,6 +166,8 @@ class KernelTraits std::vector> numBytesAndAlignmentPerSmemChunk; std::vector firstChunkReuseSmem; + // Buffer names for inspection purposes. + std::vector smemChunkNames; // LoadA { @@ -154,8 +177,8 @@ class KernelTraits auto const numBytesAlignmentLoadA = 1024; // loadA is already at first chunk. No need to reuse it. auto const reuseChunksSmemLoadA = false; - // Add info. + smemChunkNames.emplace_back("smemLoadA"); numBytesAndAlignmentPerSmemChunk.emplace_back( std::make_pair(numSmemBytesLoadA, numBytesAlignmentLoadA)); firstChunkReuseSmem.emplace_back(reuseChunksSmemLoadA); @@ -169,8 +192,8 @@ class KernelTraits auto const numBytesAlignmentLoadB = 1024; // No need to reuse the first chunk. auto const reuseChunksSmemLoadB = false; - // Add info. + smemChunkNames.emplace_back("smemLoadB"); numBytesAndAlignmentPerSmemChunk.emplace_back( std::make_pair(numSmemBytesLoadB, numBytesAlignmentLoadB)); firstChunkReuseSmem.emplace_back(reuseChunksSmemLoadB); @@ -192,6 +215,7 @@ class KernelTraits auto const reuseChunksSmemLoadB = false; // Add info. + smemChunkNames.emplace_back("smemBShuffle"); numBytesAndAlignmentPerSmemChunk.emplace_back( std::make_pair(numSmemBytesLoadB, numBytesAlignmentLoadB)); firstChunkReuseSmem.emplace_back(reuseChunksSmemLoadB); @@ -236,6 +260,7 @@ class KernelTraits = doesSplitKUseDsmem(splitK) && resIdx == 0 && !usePersistentScheduler; // Add info. + smemChunkNames.emplace_back("smemGmemC" + std::to_string(resIdx)); numBytesAndAlignmentPerSmemChunk.emplace_back( std::make_pair(numBytesSmemStoreC, numBytesAlignmentStoreC)); firstChunkReuseSmem.emplace_back(reuseFirstChunksSmemStoreC); @@ -252,6 +277,7 @@ class KernelTraits auto const numBytesAlignmentRowMax = 16; // Add info. + smemChunkNames.emplace_back("smemRowMax"); numBytesAndAlignmentPerSmemChunk.emplace_back( std::make_pair(numBytesSmemRowMax, numBytesAlignmentRowMax)); firstChunkReuseSmem.emplace_back(false); @@ -268,6 +294,7 @@ class KernelTraits auto const numBytesAlignmentTile = 16; // Add info. + smemChunkNames.emplace_back("smemSliceK"); numBytesAndAlignmentPerSmemChunk.emplace_back(std::make_pair(numBytesSmemTile, numBytesAlignmentTile)); firstChunkReuseSmem.emplace_back(false); } @@ -280,13 +307,41 @@ class KernelTraits // Number of bytes alignment for per-token scale factors auto const numBytesAlignmentPerTokenSf = 16; // Add info. + smemChunkNames.emplace_back("smemPerTokenSf"); numBytesAndAlignmentPerSmemChunk.emplace_back( std::make_pair(numBytesSmemPerTokenSf, numBytesAlignmentPerTokenSf)); firstChunkReuseSmem.emplace_back(false); } + // Per-block absolute maximum for multi-warp reduction. + { + // Number of bytes: number of epilogue warps * number of tile columns. + // TODO: avoid allocating this memory when it's not needed (it's only for MxFp8 + fusedAct) + auto const numBytesSmemBlockAmax = transposeMmaOutput ? 4 * tileN * sizeof(float) : 0; + // Number of bytes alignment. + auto const numBytesAlignmentBlockAmax = 16; + // Add info. + smemChunkNames.emplace_back("smemBlockAmax"); + numBytesAndAlignmentPerSmemChunk.emplace_back( + std::make_pair(numBytesSmemBlockAmax, numBytesAlignmentBlockAmax)); + firstChunkReuseSmem.emplace_back(false); + } + // Create SMEM helper object. - mSmemAllocatorHelper = MemAllocatorHelper(numBytesAndAlignmentPerSmemChunk, firstChunkReuseSmem); + mSmemAllocatorHelper + = MemAllocatorHelper(numBytesAndAlignmentPerSmemChunk, firstChunkReuseSmem, smemChunkNames); +#if 0 + // E.g., + // Chunk 0 smemLoadA: 32768 bytes, 1024 alignment, false, offset 0 + // Chunk 1 smemLoadB: 32768 bytes, 1024 alignment, false, offset 32768 + // Chunk 2 smemBShuffle: 0 bytes, 1024 alignment, false, offset 65536 + // Chunk 3 smemGmemC0: 65536 bytes, 1024 alignment, true, offset 0 + // Chunk 4 smemGmemC1: 65536 bytes, 1024 alignment, false, offset 65536 + // Chunk 5 smemRowMax: 512 bytes, 16 alignment, false, offset 131072 + // Chunk 6 smemSliceK: 0 bytes, 16 alignment, false, offset 131584 + // Chunk 7 smemPerTokenSf: 0 bytes, 16 alignment, false, offset 131584 + mSmemAllocatorHelper.print(); +#endif } // @@ -296,7 +351,7 @@ class KernelTraits { std::vector> numBytesAndAlignmentPerTmemChunk; std::vector firstChunkReuseTmem; - + std::vector tmemChunkNames; // Matrix D { // Number of columns for accumulators. @@ -308,6 +363,7 @@ class KernelTraits auto const reuseChunksTmemD = false; // Add info. + tmemChunkNames.emplace_back("tmemD"); numBytesAndAlignmentPerTmemChunk.emplace_back(std::make_pair(numTmemColsD, numColsAlignmentD)); firstChunkReuseTmem.emplace_back(reuseChunksTmemD); } @@ -324,6 +380,7 @@ class KernelTraits auto const reuseChunksTmemA = false; // Add info. + tmemChunkNames.emplace_back("tmemA"); numBytesAndAlignmentPerTmemChunk.emplace_back(std::make_pair(numTmemColsA, numColsAlignmentA)); firstChunkReuseTmem.emplace_back(reuseChunksTmemA); } @@ -341,6 +398,7 @@ class KernelTraits auto const reuseChunksTmemSfA = false; // Add info. + tmemChunkNames.emplace_back("tmemSfA"); numBytesAndAlignmentPerTmemChunk.emplace_back(std::make_pair(numTmemColsSfA, numColsAlignmentSfA)); firstChunkReuseTmem.emplace_back(reuseChunksTmemSfA); } @@ -356,12 +414,14 @@ class KernelTraits auto const reuseChunksTmemSfB = false; // Add info. + tmemChunkNames.emplace_back("tmemSfB"); numBytesAndAlignmentPerTmemChunk.emplace_back(std::make_pair(numTmemColsSfB, numColsAlignmentSfB)); firstChunkReuseTmem.emplace_back(reuseChunksTmemSfB); } // Create TMEM helper object. - mTmemAllocatorHelper = MemAllocatorHelper(numBytesAndAlignmentPerTmemChunk, firstChunkReuseTmem); + mTmemAllocatorHelper + = MemAllocatorHelper(numBytesAndAlignmentPerTmemChunk, firstChunkReuseTmem, tmemChunkNames); } } @@ -446,6 +506,21 @@ inline int32_t getSmemOffsetPerTokenSf(KernelTraits traits) return traits.mSmemAllocatorHelper.getChunkOffset(7); } +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetBlockAmax(KernelTraits traits) +{ + return traits.mSmemAllocatorHelper.getChunkOffset(8); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t isSmemAbRepurposedToGmemC(KernelTraits traits, int resIdx = 0) +{ + // Be conscious that the index (3 + resIdx) should match the index in getSmemOffsetGmemC(). + return traits.mSmemAllocatorHelper.getFirstChunkReuseFlag(3 + resIdx); +} + //////////////////////////////////////////////////////////////////////////////////////////////////// // // Starting address of each TMEM buffer. diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/TmaDescriptor.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/TmaDescriptor.h index 46ee0176078e..8d26c4b972bb 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/TmaDescriptor.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/TmaDescriptor.h @@ -41,8 +41,8 @@ inline CUtensorMap buildNdTmaDescriptor(tg::Dtype dtype, std::vector c { CUtensorMap desc{}; // The data type. - CUtensorMapDataType tmaDataFormat; - if (dtype == tg::Dtype::E4m3) + CUtensorMapDataType tmaDataFormat{CU_TENSOR_MAP_DATA_TYPE_FLOAT32}; + if (dtype == tg::Dtype::E4m3 || dtype == tg::Dtype::MxE4m3) { tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_UINT8; } @@ -64,7 +64,7 @@ inline CUtensorMap buildNdTmaDescriptor(tg::Dtype dtype, std::vector c } else { - std::cerr << "Unexpected dtype " << static_cast(dtype) << std::endl; + std::cerr << "buildNdTmaDescriptor: unexpected dtype " << static_cast(dtype) << std::endl; assert(false); } @@ -87,7 +87,7 @@ inline CUtensorMap buildNdTmaDescriptor(tg::Dtype dtype, std::vector c } else { - std::cerr << "Unexpected tileKSizeInBytes " << tileKSizeInBytes << std::endl; + std::cerr << "buildNdTmaDescriptor: unexpected tileKSizeInBytes " << tileKSizeInBytes << std::endl; assert(false); } } @@ -187,13 +187,13 @@ inline CUtensorMap buildSfTmaDescriptor(tg::Dtype dtype, std::vector c { CUtensorMap desc{}; CUtensorMapDataType tmaDataFormat; - if (dtype == tg::Dtype::E4m3) + if (dtype == tg::Dtype::E4m3 || dtype == tg::Dtype::UE8m0) { tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_UINT8; } else { - std::cerr << "Unexpected dtype " << static_cast(dtype) << std::endl; + std::cerr << "buildSfTmaDescriptor: unexpected dtype " << static_cast(dtype) << std::endl; assert(false); } diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/config.json b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/config.json index f7bcd45051a3..d051020f3de5 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/config.json +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/config.json @@ -25,7 +25,10 @@ "useTmaStore": true, "useCustomMmaSchedule": true, "sfLayoutB": "8x4", - "sfLayoutC": "8x4" + "sfLayoutC": "8x4", + "gridTriggerSecondaryB": true, + "gridWaitForPrimaryA": false, + "gridWaitForPrimaryB": true }, "GemmFp4Throughput": { "dtypeElt": "e2m1", @@ -77,7 +80,10 @@ "useShuffledMatrixA": true, "useDeepSeekFp8": true, "useTmaStore": true, - "useCustomMmaSchedule": true + "useCustomMmaSchedule": true, + "gridTriggerSecondaryB": true, + "gridWaitForPrimaryA": false, + "gridWaitForPrimaryB": true }, "GemmPerTensorScalingFp8Throughput": { "dtypeElt": "e4m3", @@ -128,9 +134,9 @@ "useDeepSeekFp8": false, "useTmaStore": true, "useCustomMmaSchedule": true, - "gridTriggerSecondary": true, - "gridWaitForPrimary": true, - "gridDepTriggerA": true + "gridTriggerSecondaryB": true, + "gridWaitForPrimaryA": false, + "gridWaitForPrimaryB": true } }, "configs": [ diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Bfloat16_E2m1_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x64_cluster1x1x1_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Bfloat16_E2m1_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x64_cluster1x1x1_sm100a_cubin.cpp index 64842471ced0..dd8d72a64bdb 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Bfloat16_E2m1_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x64_cluster1x1x1_sm100a_cubin.cpp +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Bfloat16_E2m1_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x64_cluster1x1x1_sm100a_cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:829f0962c056bb9b68a3d73587f3fd2f20ed2e0842bcaa26ff4d15c8fa6117ce -size 317759 +oid sha256:2d9c11ea6ceb8091c57e32bb7470095ba3bf57e8143d6953767fe36dcdd61dd2 +size 320127 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Bfloat16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Bfloat16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin.cpp new file mode 100644 index 000000000000..2ce775bc43fa --- /dev/null +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Bfloat16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin.cpp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:07bae7058bb985c753104e61b2fe7d7db65b7b087b0b3eb4c4b180fc4b3fa676 +size 343799 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Bfloat16_E2m1_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x64_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Bfloat16_E2m1_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x64_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp new file mode 100644 index 000000000000..2828ef45b2ec --- /dev/null +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Bfloat16_E2m1_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x64_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:65dd3f72fa61b68d3a201186bf31a5b83de648f60952c1e63a816e4c2e41c29d +size 323557 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Bfloat16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Bfloat16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp index ebb19151d0d1..5ddbe9a7e509 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Bfloat16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Bfloat16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:06ab65c8b780087a9c23c05465a4dc07bc9895f1123e39ef0580d414bec28c70 -size 242623 +oid sha256:795c48bdb4acc0f97561afb92eab8380a0d611109ee9ab9504ed41ada9774e50 +size 247853 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Bfloat16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Bfloat16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp index e18225550053..e8ffd10a71ab 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Bfloat16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Bfloat16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:a655d6409521c88aeaadfe4c57516d03088e83b6a9b7d2695cc13ee8e69933fe -size 230169 +oid sha256:19268e7f225df0eee153cf8ae118cf02940563de0705cf7c4b6e19d9d5260760 +size 234657 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Bfloat16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Bfloat16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin.cpp index 736d8702add2..6c3747f38350 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Bfloat16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin.cpp +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Bfloat16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:2ed42a19e269f9a255f576d2206f6df99d17c1a32f5c2d243dc217d381547719 -size 257613 +oid sha256:d789c4534f03a88d00a42ba1318427263253d47528a20ec08edbba824cb39fb4 +size 263189 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_E4m3_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_E4m3_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp index a204cd5dc730..c1bdfd690359 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_E4m3_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_E4m3_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:57b3cb8cd7fe11fc97ac2e637f4900f7d169d21ee583d04b0da5fb8c8acd45ea -size 254307 +oid sha256:5b9d2d0eb0730fc8e3f0a217bd96c8d77e61840e98a2e775b653ef99f567ecbf +size 257317 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_E4m3_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_E4m3_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp index 198c2715294c..eb0d672f71cd 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_E4m3_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_E4m3_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:3a30e23b1873e5d1d1e483214c2fe18446ededd142ff72464707daf189a796a0 -size 227793 +oid sha256:3dec737f46811775ce3634e067eb2208131987a879fe2026db923967167b3c28 +size 233071 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_E4m3_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_E4m3_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin.cpp index 30831d07ec93..b2f18d6c080a 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_E4m3_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin.cpp +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_E4m3_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:fce0feb99bddf8dcd6033455d5d49577a9585dfa21fe490eb296074f73adcb7b -size 256027 +oid sha256:25530f8f894b64c0903a620654e1a03f2da41bfe0975f3e1f185a2092773f080 +size 260813 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp16_E2m1_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x64_cluster1x1x1_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp16_E2m1_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x64_cluster1x1x1_sm100a_cubin.cpp index 119f2cf17e61..f06edd86fe5a 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp16_E2m1_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x64_cluster1x1x1_sm100a_cubin.cpp +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp16_E2m1_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x64_cluster1x1x1_sm100a_cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:29cbe332b5644245a0c13ba06b552b59dee42a01647e466ee26e1c027d543a39 -size 323325 +oid sha256:9cd9fc2392a932daf866e081ae37178e3c078ed56f542d9513d13e4e69e9b94c +size 325693 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin.cpp new file mode 100644 index 000000000000..db8e309fe638 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin.cpp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:334d0748444d2f8ffed2325d66892d99fb1f27ef61718645b9e3b0936606f9e8 +size 343001 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp16_E2m1_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x64_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp16_E2m1_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x64_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp new file mode 100644 index 000000000000..0c9a33849823 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp16_E2m1_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x64_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2760f9740e2791c03396539bcf34bfb5148b240b16c5c617a348df96819db775 +size 321969 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp index 7a7ce4c63144..b710337185ca 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:1af24c450b0b7f58889c37b6866d60375a322faace23dbb01413282539099f67 -size 247401 +oid sha256:0890ff69e276aac4d09aaa4a46bffe00d729ff359d269750d2bf3a0c6f82f8ff +size 252631 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp index f1c191068979..88f398da0b0f 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:3dd3c99ac543786bfc98b286d70a95ed4477aadf9d6cefe202fe1fe33a6dd269 -size 228581 +oid sha256:9373eede77a9ca316fa0ae26ed93d05ad43be2ff8dc8b4aa0e7210680e03031a +size 233071 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin.cpp index 6b910965286b..c74e46844e3b 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin.cpp +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp16_E4m3_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x32_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:a2d67b92d0c6afe39e2c05f6e08d5dfdb38556a0e04068b21c3735d9e9f2f2fa -size 256817 +oid sha256:d52df1a9d43c7d16c797f6a192c8748470b2200c7e984291cfb4fa0b943787df +size 261601 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp32_E2m1_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x64_cluster1x1x1_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp32_E2m1_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x64_cluster1x1x1_sm100a_cubin.cpp index 18bbdb8a5205..6ab784b22015 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp32_E2m1_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x64_cluster1x1x1_sm100a_cubin.cpp +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp32_E2m1_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x64_cluster1x1x1_sm100a_cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:a427654d9a8802677467488a89b5e3d1bc136e1b06d5dbae195bc6172b703faf -size 321895 +oid sha256:930cacc02bfc082aff9db437fc1c9c0fd19671c71a6d2ab1d049babf40eece7a +size 324263 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp32_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp32_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin.cpp new file mode 100644 index 000000000000..f2da12f30b9d --- /dev/null +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp32_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x2_splitK2_transposeMmaOutput_sm100a_cubin.cpp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8b39a92aa4d7f7ce06ce75f97ee0f81def44225c404ae7334925b65e2bf1e45b +size 343989 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp32_E2m1_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x64_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp32_E2m1_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x64_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp new file mode 100644 index 000000000000..6325fd6a73c1 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/cubins/GemmKernel_Fp32_E2m1_Fp32_tile128x8x512_epilogueTile128x8_mma128x8x64_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:71dc2688f471b76e4a6b7979f0133f295ef0f30ca5d2839664763f6b59558458 +size 323745 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/trtllm/gen/DtypeDecl.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/trtllm/gen/DtypeDecl.h index dd8eb6f3924c..a6892f12ca09 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/trtllm/gen/DtypeDecl.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_export/trtllm/gen/DtypeDecl.h @@ -67,13 +67,14 @@ enum class Dtype : uint32_t Int32 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 1u, /*bits*/ 32u, /*uid*/ 11u), Int64 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 1u, /*bits*/ 64u, /*uid*/ 12u), MxE2m1 = TLLM_ENCODE_DTYPE(/*block*/ 1u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 4u, /*uid*/ 13u), - UE8m0 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 0u, /*bits*/ 8u, /*uid*/ 14u), - UInt8 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 8u, /*uid*/ 15u), - UInt16 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 16u, /*uid*/ 16u), - UInt32 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 32u, /*uid*/ 17u), - UInt64 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 64u, /*uid*/ 18u), - UInt128 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 128u, /*uid*/ 19u), - Void = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 0u, /*uid*/ 20u), + MxE4m3 = TLLM_ENCODE_DTYPE(/*block*/ 1u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 8u, /*uid*/ 14u), + UE8m0 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 0u, /*bits*/ 8u, /*uid*/ 15u), + UInt8 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 8u, /*uid*/ 16u), + UInt16 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 16u, /*uid*/ 17u), + UInt32 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 32u, /*uid*/ 18u), + UInt64 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 64u, /*uid*/ 19u), + UInt128 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 128u, /*uid*/ 20u), + Void = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 0u, /*uid*/ 21u), // clang-format on #undef TLLM_ENCODE_DTYPE @@ -151,6 +152,7 @@ inline std::string dtypeToString(Dtype dtype) case Dtype::Int8: return "Int8"; case Dtype::Int32: return "Int32"; case Dtype::Int64: return "Int64"; + case Dtype::MxE4m3: return "MxE4m3"; case Dtype::UE8m0: return "UE8m0"; case Dtype::UInt8: return "UInt8"; case Dtype::UInt16: return "UInt16"; @@ -164,5 +166,44 @@ inline std::string dtypeToString(Dtype dtype) //////////////////////////////////////////////////////////////////////////////////////////////////// +inline Dtype dtypeEltType(Dtype dtype) +{ + switch (dtype) + { + case Dtype::MxE2m1: return Dtype::E2m1; + case Dtype::MxE4m3: return Dtype::E4m3; + default: return dtype; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int dtypeNumEltsPerSf(Dtype dtype) +{ + switch (dtype) + { + case Dtype::E2m1: return 16; + case Dtype::MxE2m1: + case Dtype::MxE4m3: return 32; + default: assert(false); return -1; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Returns the dtype of scaling factors, if applicable. +inline Dtype dtypeGetBlockSfType(Dtype dtype) +{ + switch (dtype) + { + case Dtype::E2m1: return Dtype::E4m3; + case Dtype::MxE2m1: + case Dtype::MxE4m3: return Dtype::UE8m0; + default: assert(false); return Dtype::Void; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + } // namespace gen } // namespace trtllm diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/GemmGatedActInterface.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/GemmGatedActInterface.h index 77948782db46..7bd170f7364b 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/GemmGatedActInterface.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/GemmGatedActInterface.h @@ -369,9 +369,8 @@ int32_t GemmGatedActInterface::run(GemmGatedActConfig const& config, void* works // Create kernel params. auto kernelParams = gemmGatedAct::KernelParams::setKernelParams(options, data.mInputBuffers.mPtrA, - reinterpret_cast(data.mInputBuffers.mPtrSfA), data.mInputBuffers.mPtrPerTokenSfA, - data.mInputBuffers.mPtrB, reinterpret_cast(data.mInputBuffers.mPtrSfB), - data.mInputBuffers.mPtrPerTokenSfB, data.mOutputBuffers.mPtrC, + data.mInputBuffers.mPtrSfA, data.mInputBuffers.mPtrPerTokenSfA, data.mInputBuffers.mPtrB, + data.mInputBuffers.mPtrSfB, data.mInputBuffers.mPtrPerTokenSfB, data.mOutputBuffers.mPtrC, reinterpret_cast(data.mInputBuffers.mPtrScaleC), data.mOutputBuffers.mPtrSfC, reinterpret_cast(data.mInputBuffers.mPtrScaleGate), reinterpret_cast(dRowMax), reinterpret_cast(dRowMaxBars)); @@ -407,7 +406,9 @@ int32_t GemmGatedActInterface::run(GemmGatedActConfig const& config, void* works // Run the kernel. auto result = trtllm::gen::launchKernel((void*) &kernelParams, cudaStream, config.mSharedMemSize, cuFunction, - block3, grid3, cluster3, config.mOptions.mGridWaitForPrimary); + block3, grid3, cluster3, + config.mOptions.mGridWaitForPrimaryEarlyExit | config.mOptions.mGridWaitForPrimaryA + | config.mOptions.mGridWaitForPrimaryB); if (result != CUDA_SUCCESS) { return -1; diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/GemmGatedActOptions.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/GemmGatedActOptions.h index a577756f1f4c..b23efd27746e 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/GemmGatedActOptions.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/GemmGatedActOptions.h @@ -117,6 +117,14 @@ inline bool checkAndUpdateGemmGatedActOptions( TLLM_CHECK_ERROR(hiddenSize % 256 == 0, "Output hidden size must be a multiple of 256"); } + if (options.mDtypeC == tg::Dtype::E2m1 || options.mDtypeC == tg::Dtype::MxE4m3) + { + int const outHiddenSize = (options.mTransposeMmaOutput ? options.mM : options.mN) / 2; + int const hiddenGranularity = 4 * tg::dtypeNumEltsPerSf(options.mDtypeC); + TLLM_CHECK_ERROR(outHiddenSize % hiddenGranularity == 0, "Output hidden size (", outHiddenSize, + ") must be a multiple of ", hiddenGranularity, " for block-scaled outputs."); + } + auto isValid = gemm::checkAndUpdateGemmOptions(options, isBlackwell, /* tpGrpSize */ 1, updateOptions); diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/GemmOptions.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/GemmOptions.h index 13273d2a184a..24624ee0aa2e 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/GemmOptions.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/GemmOptions.h @@ -50,7 +50,7 @@ void printArgs(T first, Args... args) #define TLLM_LOG_ERROR(...) TLLM_CHECK_ERROR(false, __VA_ARGS__) -#define TLLM_CHECK_ERROR_FMT(...) TLLM_CHECK_ERROR(false, __VA_ARGS__) +#define TLLM_CHECK_ERROR_FMT(cond, ...) TLLM_CHECK_ERROR(cond, __VA_ARGS__) #define TLLM_CHECK_WARNING(cond, ...) \ if (!(cond)) \ @@ -100,18 +100,24 @@ struct GemmOptions tg::Dtype mDtypeC{tg::Dtype::Void}; // Whether to enable early exit. bool mEnablesEarlyExit{false}; + // Whether to enable early exit. + bool mEnablesDelayedEarlyExit{false}; // Whether to enable the global PTX knobs for guiding the compiler optimizations. bool mEnablesGlobalPtxKnobs{true}; // Tile size for the epilogue in M dimension. int mEpilogueTileM{128}; // Tile size for the epilogue in N dimension. int mEpilogueTileN{32}; - // Whether A or B load task triggers grid dependency controls. - bool mGridDepTriggerA{false}; - // Whether we should trigger a grid dependency. - bool mGridTriggerSecondary{false}; - // Whether we should wait on a grid dependency. - bool mGridWaitForPrimary{false}; + // Whether load task A triggers the next grid. + bool mGridTriggerSecondaryA{false}; + // Whether load task B triggers the next grid. + bool mGridTriggerSecondaryB{false}; + // Whether the loads that check for an early exit should wait on a grid dependency. + bool mGridWaitForPrimaryEarlyExit{true}; + // Whether the load of A should wait on a grid dependency. + bool mGridWaitForPrimaryA{true}; + // Whether the load of B should wait on a grid dependency. + bool mGridWaitForPrimaryB{true}; // Whether to hoist the mbarrier try_waits (e.g., mma.prodAcq, smemAb.consWait) in the MMA task. bool mHoistMmaTaskTryWaits{false}; // The K dimension of GEMM. @@ -139,10 +145,15 @@ struct GemmOptions int mNumSlicesForSliceK{1}; // The depth of the mainloop pipeline. int mNumStages{2}; - // The depth of the mma pipeline. + // The depth of the mma pipeline. Equals numStagesMmaWithinWorkTile * numStagesMmaAcrossWorkTile. int mNumStagesMma{1}; - // The depth of the work id pipeline. - int mNumStagesWorkId{2}; + // The depth of the mma pipeline within work tile. Only GmemC classes with "WithAccInReg" suffix + // are allowed to be greater than 1. + int mNumStagesMmaWithinWorkTile{-1}; + // The depth of the mma pipeline across work tiles in the persistent loop. + int mNumStagesMmaAcrossWorkTile{-1}; + // The depth of the work id pipeline and the work throttle pipeline. + int mNumStagesWorkId{3}; // Whether to output debug tensors. bool mOutputDebugTensors{false}; // Reorder rows/cols in the A matrix for the better memory accesses in the M-major epilogue. @@ -163,6 +174,10 @@ struct GemmOptions bool mUseUnrollLoop2xForMma{true}; // Use custom MMA schedule optimized for low-latency. bool mUseCustomMmaSchedule{false}; + // The purpose of hoisting trywaits is to opportunistically peek at the availability of the next + // k-block. It benefits when the next k-block is already available and thus sustaining the + // momentum, but it adds latency to the first k-block for smaller k-loop. + bool mUseHoistTryWaitForCustomMmaSchedule{false}; // Use DeepSeek Fp8. bool mUseDeepSeekFp8{false}; // Apply per-token scales from A @@ -256,12 +271,15 @@ inline std::string dumpOptions(GemmOptions const& options) << "trtllm::gen::Dtype(" << static_cast(options.mDtypeC) << ")" << "," << std::endl; ss << "mEnablesEarlyExit=" << options.mEnablesEarlyExit << "," << std::endl; + ss << "mEnablesDelayedEarlyExit=" << options.mEnablesDelayedEarlyExit << "," << std::endl; ss << "mEnablesGlobalPtxKnobs=" << options.mEnablesGlobalPtxKnobs << "," << std::endl; ss << "mEpilogueTileM=" << options.mEpilogueTileM << "," << std::endl; ss << "mEpilogueTileN=" << options.mEpilogueTileN << "," << std::endl; - ss << "mGridDepTriggerA=" << options.mGridDepTriggerA << "," << std::endl; - ss << "mGridTriggerSecondary=" << options.mGridTriggerSecondary << "," << std::endl; - ss << "mGridWaitForPrimary=" << options.mGridWaitForPrimary << "," << std::endl; + ss << "mGridTriggerSecondaryA=" << options.mGridTriggerSecondaryA << "," << std::endl; + ss << "mGridTriggerSecondaryB=" << options.mGridTriggerSecondaryB << "," << std::endl; + ss << "mGridWaitForPrimaryEarlyExit=" << options.mGridWaitForPrimaryEarlyExit << "," << std::endl; + ss << "mGridWaitForPrimaryA=" << options.mGridWaitForPrimaryA << "," << std::endl; + ss << "mGridWaitForPrimaryB=" << options.mGridWaitForPrimaryB << "," << std::endl; ss << "mHoistMmaTaskTryWaits=" << options.mHoistMmaTaskTryWaits << "," << std::endl; ss << "mK=" << options.mK << "," << std::endl; ss << "mKernelTraits={}" @@ -276,6 +294,8 @@ inline std::string dumpOptions(GemmOptions const& options) ss << "mNumSlicesForSliceK=" << options.mNumSlicesForSliceK << "," << std::endl; ss << "mNumStages=" << options.mNumStages << "," << std::endl; ss << "mNumStagesMma=" << options.mNumStagesMma << "," << std::endl; + ss << "mNumStagesMmaWithinWorkTile=" << options.mNumStagesMmaWithinWorkTile << "," << std::endl; + ss << "mNumStagesMmaAcrossWorkTile=" << options.mNumStagesMmaAcrossWorkTile << "," << std::endl; ss << "mNumStagesWorkId=" << options.mNumStagesWorkId << "," << std::endl; ss << "mOutputDebugTensors=" << options.mOutputDebugTensors << "," << std::endl; ss << "mUseShuffledMatrixA=" << options.mUseShuffledMatrixA << "," << std::endl; @@ -289,6 +309,7 @@ inline std::string dumpOptions(GemmOptions const& options) ss << "mTileK=" << options.mTileK << "," << std::endl; ss << "mUseUnrollLoop2xForMma=" << options.mUseUnrollLoop2xForMma << "," << std::endl; ss << "mUseCustomMmaSchedule=" << options.mUseCustomMmaSchedule << "," << std::endl; + ss << "mUseHoistTryWaitForCustomMmaSchedule=" << options.mUseHoistTryWaitForCustomMmaSchedule << "," << std::endl; ss << "mUseDeepSeekFp8=" << options.mUseDeepSeekFp8 << "," << std::endl; ss << "mUsePerTokenSfA=" << options.mUsePerTokenSfA << "," << std::endl; ss << "mUsePerTokenSfB=" << options.mUsePerTokenSfB << "," << std::endl; @@ -350,20 +371,41 @@ inline bool checkAndUpdateGemmOptions( } } - // NvFp4 constraints - if (options.mDtypeElt == tg::Dtype::E2m1) + // Constraints for NvFp4 and MxFp8. + if ((options.mDtypeElt == tg::Dtype::E2m1 || options.mDtypeElt == tg::Dtype::MxE4m3 + || options.mDtypeC == tg::Dtype::MxE4m3) + && options.mMmaM != 128) + { + // MMA M must be 128 when the input uses block scaling, or when the output is an Mx format. + int newTileM = 128 * divUp(options.mTileM, 128); + TLLM_LOG_WARNING("Unsupported MmaM (", options.mMmaM, ") for dtypeElt=", gemm::toString(options.mDtypeElt), + ", dtypeC=", gemm::toString(options.mDtypeC), ". Setting MmaM to 128 and TileM to ", newTileM); + if (updateOptions) + { + options.mMmaM = 128; + options.mTileM = newTileM; + } + else + { + return false; + } + } + if (options.mDtypeElt == tg::Dtype::E2m1 || options.mDtypeElt == tg::Dtype::MxE4m3) { - TLLM_CHECK_ERROR(isBlackwell, "FP4 is only supported on Blackwell"); + TLLM_CHECK_ERROR(isBlackwell, "Block scaling is only supported on Blackwell"); + TLLM_CHECK_ERROR(options.mSfLayoutB == tg::SfLayout::R128c4 || options.mSfLayoutB == tg::SfLayout::R8c4, "Only the 128x4 and 8x4 SF layouts are supported for B, got ", tg::sfLayoutToString(options.mSfLayoutB)); - if (options.mMmaK != 64) + + int const mmaK = (options.mDtypeElt == tg::Dtype::E2m1) ? 64 : 32; + if (options.mMmaK != mmaK) { - int newTileK = 64 * divUp(options.mTileK, 64); + int newTileK = mmaK * divUp(options.mTileK, mmaK); TLLM_LOG_WARNING("Unsupported MmaK (", options.mMmaK, ") for ", gemm::toString(options.mDtypeElt), - ". Setting MmaK to 64 and TileK to ", newTileK); + ". Setting MmaK to ", mmaK, " and TileK to ", newTileK); if (updateOptions) { - options.mMmaK = 64; + options.mMmaK = mmaK; options.mTileK = newTileK; } else @@ -371,33 +413,22 @@ inline bool checkAndUpdateGemmOptions( return false; } } - if (options.mMmaM != 128) - { - int newTileM = 128 * divUp(options.mTileM, 128); - TLLM_LOG_WARNING("Unsupported MmaM (", options.mMmaM, ") for ", gemm::toString(options.mDtypeElt), - ". Setting MmaM to 128 and TileM to ", newTileM); - if (updateOptions) - { - options.mMmaM = 128; - options.mTileM = newTileM; - } - else - { - return false; - } - } + // TileN must be a multiple of the number of rows per SF tile. - // TODO: relax this restriction. int const numSfTileRowsB = options.mSfLayoutB == tg::SfLayout::R128c4 ? 128 : 8; TLLM_CHECK_ERROR(options.mTileN % numSfTileRowsB == 0, "TileN (", options.mTileN, ") must be a multiple of ", numSfTileRowsB, " for B SF layout ", tg::sfLayoutToString(options.mSfLayoutB)); // The MMA N may only be smaller than 64 if it is equal to the tile N. TLLM_CHECK_ERROR(options.mMmaN >= 64 || options.mMmaN == options.mTileN, "MmaN (", options.mMmaN, ") must be >= 64 or equal to TileN (", options.mTileN, ") for ", gemm::toString(options.mDtypeElt)); + + int numEltsPerSf = tg::dtypeNumEltsPerSf(options.mDtypeElt); + TLLM_CHECK_ERROR(options.mTileK % (4 * numEltsPerSf) == 0, "TileK (", options.mTileK, + ") must be a multiple of ", (4 * numEltsPerSf), " for type ", gemm::toString(options.mDtypeElt)); } - if (options.mDtypeC == tg::Dtype::E2m1) + if (options.mDtypeC == tg::Dtype::E2m1 || options.mDtypeC == tg::Dtype::MxE4m3) { - TLLM_CHECK_ERROR(isBlackwell, "FP4 is only supported on Blackwell"); + TLLM_CHECK_ERROR(isBlackwell, "Block scaling is only supported on Blackwell"); TLLM_CHECK_ERROR(options.mSfLayoutC == tg::SfLayout::R128c4 || options.mSfLayoutC == tg::SfLayout::R8c4, "Only the 128x4 and 8x4 SF layouts are supported for C."); @@ -406,9 +437,11 @@ inline bool checkAndUpdateGemmOptions( numSfTileRowsC, " for C SF layout ", tg::sfLayoutToString(options.mSfLayoutC)); int const hiddenDim = options.mTransposeMmaOutput ? options.mM : options.mN; - TLLM_CHECK_ERROR(hiddenDim % 64 == 0, "Hidden dim (", hiddenDim, ") must be a multiple of 64 for FP4 outputs."); + int const hiddenGranularity = 4 * tg::dtypeNumEltsPerSf(options.mDtypeC); + TLLM_CHECK_ERROR(hiddenDim % hiddenGranularity == 0, "Hidden dim (", hiddenDim, ") must be a multiple of ", + hiddenGranularity, " for block-scaled outputs."); TLLM_CHECK_ERROR(!options.mTransposeMmaOutput || options.mUseShuffledMatrixA, - "Transposing FP4 outputs requires shuffled A."); + "Transposing block-scaled outputs requires shuffled A."); } // If dtypeC is unspecified (Dtype::Void), assign to the input dtype. @@ -589,10 +622,61 @@ inline bool checkAndUpdateGemmOptions( "CGA size must be equal to the number of slices in split-k"); } + // Maps numStagesMma to (stagesWithinWorkTile, stagesAcrossWorkTile) if not already set. + // If (-1, -1) -> (numStagesMma / min(2, numStagesMma), min(2, numStagesMma)) + // If ( m, -1) -> (m, numStagesMma / m) + // If (-1, n) -> (numStagesMma / n, n) + if (options.mNumStagesMmaWithinWorkTile == -1 && options.mNumStagesMmaAcrossWorkTile == -1) + { + if (updateOptions) + { + options.mNumStagesMmaAcrossWorkTile = std::min(2, options.mNumStagesMma); + options.mNumStagesMmaWithinWorkTile = options.mNumStagesMma / options.mNumStagesMmaAcrossWorkTile; + } + else + { + return false; + } + } + else if (options.mNumStagesMmaWithinWorkTile == -1) + { + if (updateOptions) + { + options.mNumStagesMmaWithinWorkTile = options.mNumStagesMma / options.mNumStagesMmaAcrossWorkTile; + } + else + { + return false; + } + } + else if (options.mNumStagesMmaAcrossWorkTile == -1) + { + if (updateOptions) + { + options.mNumStagesMmaAcrossWorkTile = options.mNumStagesMma / options.mNumStagesMmaWithinWorkTile; + } + else + { + return false; + } + } + // Check mma stages. + TLLM_CHECK_ERROR_FMT( + options.mNumStagesMmaWithinWorkTile * options.mNumStagesMmaAcrossWorkTile == options.mNumStagesMma + && options.mNumStagesMmaAcrossWorkTile <= 2, + "Condition numStagesMmaWithinWorkTile (%d) * numStagesMmaAcrossWorkTile " + "(%d) == numStagesMma (%d) && numStagesMmaAcrossWorkTile (%d) <= 2 must be " + "satisfied. Check arguments.", + options.mNumStagesMmaWithinWorkTile, options.mNumStagesMmaAcrossWorkTile, options.mNumStagesMma, + options.mNumStagesMmaAcrossWorkTile); + // Mma stage must be 1 for pre-Hopper. + TLLM_CHECK_ERROR( + isBlackwell || options.mNumStagesMma == 1, "Mma stage must be 1 for pre-Hopper. Found ", options.mNumStagesMma); // DeepSeek Fp8 if (!options.mUseDeepSeekFp8) { - TLLM_CHECK_ERROR(options.mNumStagesMma == 1, "Non-DeepSeekFp8 requires numStagesMma == 1"); + TLLM_CHECK_ERROR( + options.mNumStagesMmaWithinWorkTile == 1, "Non-DeepSeekFp8 requires numStagesMmaWithinWorkTile == 1"); } if (options.mUseDeepSeekFp8) { @@ -610,7 +694,8 @@ inline bool checkAndUpdateGemmOptions( auto hiddenDimPerEpilogueTile = options.mTransposeMmaOutput ? options.mEpilogueTileM : options.mEpilogueTileN; auto hiddenDimPerMma = options.mTransposeMmaOutput ? options.mMmaM : options.mMmaN; auto hiddenDimName = options.mTransposeMmaOutput ? "M" : "N"; - TLLM_CHECK_WARNING(options.mNumStagesMma > 1, "DeepSeekFp8 recommends >1 MMA accumulator stages."); + TLLM_CHECK_WARNING(options.mNumStagesMmaWithinWorkTile > 1, + "DeepSeekFp8 recommends setting \"-numStagesMmaWithinWorkTile 2\"."); // Update the number of stages of the MMA accumulator pipeline. TODO: enable by default for // deepseek. // options.mNumStagesMma = 2; @@ -708,6 +793,42 @@ inline bool checkAndUpdateGemmOptions( } } + if (options.mEnablesDelayedEarlyExit && options.mEnablesEarlyExit) + { + TLLM_LOG_WARNING( + "Only one of early exit and delayed early exit should be enabled. Disabling " + "delayed early exit"); + if (updateOptions) + { + options.mEnablesDelayedEarlyExit = false; + } + else + { + return false; + } + } + + // This check prevents the triggering of the secondary (PREEXIT) from executing before the wait + // for primary (ACQBULK). This could lead to the following confusing situation, which we want to + // avoid: + // + // Kernel 3 is written with the assumption that it can read the output of + // kernel 1 *without* ACQBULK and the output of kernel 2 *with* ACQBULK. + // However, when we allow PREEXIT and ACQBULK to be executed out of order, + // this is not guaranteed. + // + // Time: ----> + // + // Kernel 1: ----PREEXIT-----------FLUSH + // Kernel 2: -------PREEXIT----ACQBULK---FLUSH + // Kernel 3: Warp 0: ---- (!) Output of 1,2 is not yet visible ----------------------- + // Warp 1: ---- (!) We normally assume that 1 is visible is not yet visible- + // Warp 2: -------------------ACQBULK-- Kernel 1,2 output visible ---------- + TLLM_CHECK_ERROR((options.mGridWaitForPrimaryA || !options.mGridTriggerSecondaryA), + "A: If a task triggers a secondary kernel, it must also wait for primary kernel."); + TLLM_CHECK_ERROR((options.mGridWaitForPrimaryB || !options.mGridTriggerSecondaryB), + "B: If a task triggers a secondary kernel, it must also wait for primary kernel."); + if (updateOptions) { // Init kernel traits. diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/KernelMetaInfo.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/KernelMetaInfo.h index e26b51536447..e4f9b89c9353 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/KernelMetaInfo.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/KernelMetaInfo.h @@ -25,38 +25,102 @@ namespace kernels { // clang-format off -#define TLLM_GEN_COMMIT "3d9075a-dirty" +#define TLLM_GEN_COMMIT "23d32a5" #define TLLM_GEN_EXPORT_VERSION "0.0" -static constexpr size_t tllmGenGemmGatedActListLen = 9; +static constexpr size_t tllmGenGemmGatedActListLen = 13; #ifndef EXCLUDE_SM_100 +extern unsigned char GemmGatedActKernel_Bfloat16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin[]; extern unsigned char GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin[]; extern unsigned char GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin[]; extern unsigned char GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin[]; +extern unsigned char GemmGatedActKernel_E2m1_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin[]; extern unsigned char GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin[]; extern unsigned char GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin[]; extern unsigned char GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin[]; +extern unsigned char GemmGatedActKernel_Fp16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin[]; extern unsigned char GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin[]; extern unsigned char GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin[]; extern unsigned char GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin[]; +extern unsigned char GemmGatedActKernel_Fp32_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin[]; #endif // EXCLUDE_SM_100 #ifndef EXCLUDE_SM_100 +extern unsigned int GemmGatedActKernel_Bfloat16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin_len; extern unsigned int GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin_len; extern unsigned int GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin_len; extern unsigned int GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin_len; +extern unsigned int GemmGatedActKernel_E2m1_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin_len; extern unsigned int GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin_len; extern unsigned int GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin_len; extern unsigned int GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin_len; +extern unsigned int GemmGatedActKernel_Fp16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin_len; extern unsigned int GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin_len; extern unsigned int GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin_len; extern unsigned int GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin_len; +extern unsigned int GemmGatedActKernel_Fp32_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin_len; #endif // EXCLUDE_SM_100 static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { #ifndef EXCLUDE_SM_100 +{GemmGatedActKernel_Bfloat16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin, GemmGatedActKernel_Bfloat16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin_len, 86016, "gemmGatedActKernel_Bfloat16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a", 448, {{ /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 4 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056777) +, /* mDtypeElt */ trtllm::gen::Dtype(17826819) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 4 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mTileScheduler */ gemm::TileScheduler(0) + }, /* mActType */ gemmGatedAct::ActType(0) + }, gemm::SmVersion::Sm100a }, {GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin, GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin_len, 168960, "gemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a", 224, {{ /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) , /* mClusterDimX */ 1 , /* mClusterDimY */ 1 @@ -65,12 +129,15 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mDtypeElt */ trtllm::gen::Dtype(1050630) , /* mDtypeC */ trtllm::gen::Dtype(1052672) , /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 , /* mEnablesGlobalPtxKnobs */ 1 , /* mEpilogueTileM */ 128 , /* mEpilogueTileN */ 128 -, /* mGridDepTriggerA */ 0 -, /* mGridTriggerSecondary */ 0 -, /* mGridWaitForPrimary */ 0 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 0 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 1 +, /* mGridWaitForPrimaryB */ 1 , /* mHoistMmaTaskTryWaits */ 0 , /* mK */ 2048 , /* mKernelTraits */ {} @@ -84,7 +151,9 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mNumSlicesForSliceK */ 1 , /* mNumStages */ 2 , /* mNumStagesMma */ 1 -, /* mNumStagesWorkId */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 , /* mOutputDebugTensors */ 0 , /* mUseShuffledMatrixA */ 0 , /* mSliceK */ 0 @@ -95,6 +164,7 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mTileK */ 256 , /* mUseUnrollLoop2xForMma */ 1 , /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 , /* mUseDeepSeekFp8 */ 0 , /* mUsePerTokenSfA */ 0 , /* mUsePerTokenSfB */ 0 @@ -107,7 +177,7 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mTileScheduler */ gemm::TileScheduler(0) }, /* mActType */ gemmGatedAct::ActType(0) }, gemm::SmVersion::Sm100a }, -{GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin, GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin_len, 111616, "gemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a", 224, {{ /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +{GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin, GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin_len, 112640, "gemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a", 224, {{ /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) , /* mClusterDimX */ 1 , /* mClusterDimY */ 1 , /* mClusterDimZ */ 1 @@ -115,12 +185,15 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mDtypeElt */ trtllm::gen::Dtype(1050630) , /* mDtypeC */ trtllm::gen::Dtype(1052672) , /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 , /* mEnablesGlobalPtxKnobs */ 1 , /* mEpilogueTileM */ 128 , /* mEpilogueTileN */ 8 -, /* mGridDepTriggerA */ 1 -, /* mGridTriggerSecondary */ 1 -, /* mGridWaitForPrimary */ 1 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 , /* mHoistMmaTaskTryWaits */ 0 , /* mK */ 2048 , /* mKernelTraits */ {} @@ -134,7 +207,9 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mNumSlicesForSliceK */ 1 , /* mNumStages */ 3 , /* mNumStagesMma */ 1 -, /* mNumStagesWorkId */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 , /* mOutputDebugTensors */ 0 , /* mUseShuffledMatrixA */ 1 , /* mSliceK */ 0 @@ -145,6 +220,7 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mTileK */ 256 , /* mUseUnrollLoop2xForMma */ 1 , /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 , /* mUseDeepSeekFp8 */ 0 , /* mUsePerTokenSfA */ 0 , /* mUsePerTokenSfB */ 0 @@ -157,7 +233,7 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mTileScheduler */ gemm::TileScheduler(0) }, /* mActType */ gemmGatedAct::ActType(0) }, gemm::SmVersion::Sm100a }, -{GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin, GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin_len, 109568, "gemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a", 224, {{ /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +{GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin, GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin_len, 110592, "gemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a", 224, {{ /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) , /* mClusterDimX */ 1 , /* mClusterDimY */ 1 , /* mClusterDimZ */ 4 @@ -165,12 +241,15 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mDtypeElt */ trtllm::gen::Dtype(1050630) , /* mDtypeC */ trtllm::gen::Dtype(1052672) , /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 , /* mEnablesGlobalPtxKnobs */ 1 , /* mEpilogueTileM */ 128 , /* mEpilogueTileN */ 8 -, /* mGridDepTriggerA */ 1 -, /* mGridTriggerSecondary */ 1 -, /* mGridWaitForPrimary */ 1 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 , /* mHoistMmaTaskTryWaits */ 0 , /* mK */ 2048 , /* mKernelTraits */ {} @@ -184,7 +263,9 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mNumSlicesForSliceK */ 1 , /* mNumStages */ 3 , /* mNumStagesMma */ 1 -, /* mNumStagesWorkId */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 , /* mOutputDebugTensors */ 0 , /* mUseShuffledMatrixA */ 1 , /* mSliceK */ 0 @@ -195,6 +276,7 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mTileK */ 256 , /* mUseUnrollLoop2xForMma */ 1 , /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 , /* mUseDeepSeekFp8 */ 0 , /* mUsePerTokenSfA */ 0 , /* mUsePerTokenSfB */ 0 @@ -204,6 +286,62 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mSfLayoutA */ trtllm::gen::SfLayout(3) , /* mSfLayoutB */ trtllm::gen::SfLayout(3) , /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mTileScheduler */ gemm::TileScheduler(0) + }, /* mActType */ gemmGatedAct::ActType(0) + }, gemm::SmVersion::Sm100a }, +{GemmGatedActKernel_E2m1_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin, GemmGatedActKernel_E2m1_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin_len, 86016, "gemmGatedActKernel_E2m1_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a", 448, {{ /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 4 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056777) +, /* mDtypeElt */ trtllm::gen::Dtype(17826819) +, /* mDtypeC */ trtllm::gen::Dtype(17826819) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 4 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) , /* mTileScheduler */ gemm::TileScheduler(0) }, /* mActType */ gemmGatedAct::ActType(0) }, gemm::SmVersion::Sm100a }, @@ -215,12 +353,15 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mDtypeElt */ trtllm::gen::Dtype(1050630) , /* mDtypeC */ trtllm::gen::Dtype(1050630) , /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 , /* mEnablesGlobalPtxKnobs */ 1 , /* mEpilogueTileM */ 128 , /* mEpilogueTileN */ 128 -, /* mGridDepTriggerA */ 0 -, /* mGridTriggerSecondary */ 0 -, /* mGridWaitForPrimary */ 0 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 0 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 1 +, /* mGridWaitForPrimaryB */ 1 , /* mHoistMmaTaskTryWaits */ 0 , /* mK */ 2048 , /* mKernelTraits */ {} @@ -234,7 +375,9 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mNumSlicesForSliceK */ 1 , /* mNumStages */ 3 , /* mNumStagesMma */ 1 -, /* mNumStagesWorkId */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 , /* mOutputDebugTensors */ 0 , /* mUseShuffledMatrixA */ 0 , /* mSliceK */ 0 @@ -245,6 +388,7 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mTileK */ 256 , /* mUseUnrollLoop2xForMma */ 1 , /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 , /* mUseDeepSeekFp8 */ 0 , /* mUsePerTokenSfA */ 0 , /* mUsePerTokenSfB */ 0 @@ -257,7 +401,7 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mTileScheduler */ gemm::TileScheduler(0) }, /* mActType */ gemmGatedAct::ActType(0) }, gemm::SmVersion::Sm100a }, -{GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin, GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin_len, 110592, "gemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a", 224, {{ /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +{GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin, GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin_len, 111616, "gemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a", 224, {{ /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) , /* mClusterDimX */ 1 , /* mClusterDimY */ 1 , /* mClusterDimZ */ 1 @@ -265,12 +409,15 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mDtypeElt */ trtllm::gen::Dtype(1050630) , /* mDtypeC */ trtllm::gen::Dtype(1050630) , /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 , /* mEnablesGlobalPtxKnobs */ 1 , /* mEpilogueTileM */ 128 , /* mEpilogueTileN */ 8 -, /* mGridDepTriggerA */ 1 -, /* mGridTriggerSecondary */ 1 -, /* mGridWaitForPrimary */ 1 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 , /* mHoistMmaTaskTryWaits */ 0 , /* mK */ 2048 , /* mKernelTraits */ {} @@ -284,7 +431,9 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mNumSlicesForSliceK */ 1 , /* mNumStages */ 3 , /* mNumStagesMma */ 1 -, /* mNumStagesWorkId */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 , /* mOutputDebugTensors */ 0 , /* mUseShuffledMatrixA */ 1 , /* mSliceK */ 0 @@ -295,6 +444,7 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mTileK */ 256 , /* mUseUnrollLoop2xForMma */ 1 , /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 , /* mUseDeepSeekFp8 */ 0 , /* mUsePerTokenSfA */ 0 , /* mUsePerTokenSfB */ 0 @@ -307,7 +457,7 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mTileScheduler */ gemm::TileScheduler(0) }, /* mActType */ gemmGatedAct::ActType(0) }, gemm::SmVersion::Sm100a }, -{GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin, GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin_len, 109568, "gemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a", 224, {{ /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +{GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin, GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin_len, 110592, "gemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a", 224, {{ /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) , /* mClusterDimX */ 1 , /* mClusterDimY */ 1 , /* mClusterDimZ */ 4 @@ -315,12 +465,15 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mDtypeElt */ trtllm::gen::Dtype(1050630) , /* mDtypeC */ trtllm::gen::Dtype(1050630) , /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 , /* mEnablesGlobalPtxKnobs */ 1 , /* mEpilogueTileM */ 128 , /* mEpilogueTileN */ 8 -, /* mGridDepTriggerA */ 1 -, /* mGridTriggerSecondary */ 1 -, /* mGridWaitForPrimary */ 1 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 , /* mHoistMmaTaskTryWaits */ 0 , /* mK */ 2048 , /* mKernelTraits */ {} @@ -334,7 +487,9 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mNumSlicesForSliceK */ 1 , /* mNumStages */ 3 , /* mNumStagesMma */ 1 -, /* mNumStagesWorkId */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 , /* mOutputDebugTensors */ 0 , /* mUseShuffledMatrixA */ 1 , /* mSliceK */ 0 @@ -345,6 +500,7 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mTileK */ 256 , /* mUseUnrollLoop2xForMma */ 1 , /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 , /* mUseDeepSeekFp8 */ 0 , /* mUsePerTokenSfA */ 0 , /* mUsePerTokenSfB */ 0 @@ -354,6 +510,62 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mSfLayoutA */ trtllm::gen::SfLayout(3) , /* mSfLayoutB */ trtllm::gen::SfLayout(3) , /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mTileScheduler */ gemm::TileScheduler(0) + }, /* mActType */ gemmGatedAct::ActType(0) + }, gemm::SmVersion::Sm100a }, +{GemmGatedActKernel_Fp16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin, GemmGatedActKernel_Fp16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin_len, 86016, "gemmGatedActKernel_Fp16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a", 448, {{ /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 4 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056777) +, /* mDtypeElt */ trtllm::gen::Dtype(17826819) +, /* mDtypeC */ trtllm::gen::Dtype(1052680) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 4 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) , /* mTileScheduler */ gemm::TileScheduler(0) }, /* mActType */ gemmGatedAct::ActType(0) }, gemm::SmVersion::Sm100a }, @@ -365,12 +577,15 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mDtypeElt */ trtllm::gen::Dtype(1050630) , /* mDtypeC */ trtllm::gen::Dtype(1052680) , /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 , /* mEnablesGlobalPtxKnobs */ 1 , /* mEpilogueTileM */ 128 , /* mEpilogueTileN */ 128 -, /* mGridDepTriggerA */ 0 -, /* mGridTriggerSecondary */ 0 -, /* mGridWaitForPrimary */ 0 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 0 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 1 +, /* mGridWaitForPrimaryB */ 1 , /* mHoistMmaTaskTryWaits */ 0 , /* mK */ 2048 , /* mKernelTraits */ {} @@ -384,7 +599,9 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mNumSlicesForSliceK */ 1 , /* mNumStages */ 2 , /* mNumStagesMma */ 1 -, /* mNumStagesWorkId */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 , /* mOutputDebugTensors */ 0 , /* mUseShuffledMatrixA */ 0 , /* mSliceK */ 0 @@ -395,6 +612,7 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mTileK */ 256 , /* mUseUnrollLoop2xForMma */ 1 , /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 , /* mUseDeepSeekFp8 */ 0 , /* mUsePerTokenSfA */ 0 , /* mUsePerTokenSfB */ 0 @@ -407,7 +625,7 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mTileScheduler */ gemm::TileScheduler(0) }, /* mActType */ gemmGatedAct::ActType(0) }, gemm::SmVersion::Sm100a }, -{GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin, GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin_len, 111616, "gemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a", 224, {{ /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +{GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin, GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin_len, 112640, "gemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a", 224, {{ /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) , /* mClusterDimX */ 1 , /* mClusterDimY */ 1 , /* mClusterDimZ */ 1 @@ -415,12 +633,15 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mDtypeElt */ trtllm::gen::Dtype(1050630) , /* mDtypeC */ trtllm::gen::Dtype(1052680) , /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 , /* mEnablesGlobalPtxKnobs */ 1 , /* mEpilogueTileM */ 128 , /* mEpilogueTileN */ 8 -, /* mGridDepTriggerA */ 1 -, /* mGridTriggerSecondary */ 1 -, /* mGridWaitForPrimary */ 1 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 , /* mHoistMmaTaskTryWaits */ 0 , /* mK */ 2048 , /* mKernelTraits */ {} @@ -434,7 +655,9 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mNumSlicesForSliceK */ 1 , /* mNumStages */ 3 , /* mNumStagesMma */ 1 -, /* mNumStagesWorkId */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 , /* mOutputDebugTensors */ 0 , /* mUseShuffledMatrixA */ 1 , /* mSliceK */ 0 @@ -445,6 +668,7 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mTileK */ 256 , /* mUseUnrollLoop2xForMma */ 1 , /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 , /* mUseDeepSeekFp8 */ 0 , /* mUsePerTokenSfA */ 0 , /* mUsePerTokenSfB */ 0 @@ -457,7 +681,7 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mTileScheduler */ gemm::TileScheduler(0) }, /* mActType */ gemmGatedAct::ActType(0) }, gemm::SmVersion::Sm100a }, -{GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin, GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin_len, 109568, "gemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a", 224, {{ /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +{GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin, GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin_len, 110592, "gemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a", 224, {{ /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) , /* mClusterDimX */ 1 , /* mClusterDimY */ 1 , /* mClusterDimZ */ 4 @@ -465,12 +689,15 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mDtypeElt */ trtllm::gen::Dtype(1050630) , /* mDtypeC */ trtllm::gen::Dtype(1052680) , /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 , /* mEnablesGlobalPtxKnobs */ 1 , /* mEpilogueTileM */ 128 , /* mEpilogueTileN */ 8 -, /* mGridDepTriggerA */ 1 -, /* mGridTriggerSecondary */ 1 -, /* mGridWaitForPrimary */ 1 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 , /* mHoistMmaTaskTryWaits */ 0 , /* mK */ 2048 , /* mKernelTraits */ {} @@ -484,7 +711,9 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mNumSlicesForSliceK */ 1 , /* mNumStages */ 3 , /* mNumStagesMma */ 1 -, /* mNumStagesWorkId */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 , /* mOutputDebugTensors */ 0 , /* mUseShuffledMatrixA */ 1 , /* mSliceK */ 0 @@ -495,6 +724,7 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mTileK */ 256 , /* mUseUnrollLoop2xForMma */ 1 , /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 , /* mUseDeepSeekFp8 */ 0 , /* mUsePerTokenSfA */ 0 , /* mUsePerTokenSfB */ 0 @@ -504,6 +734,62 @@ static const gemmGatedAct::GemmGatedActConfig tllmGenGemmGatedActList[] = { , /* mSfLayoutA */ trtllm::gen::SfLayout(3) , /* mSfLayoutB */ trtllm::gen::SfLayout(3) , /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mTileScheduler */ gemm::TileScheduler(0) + }, /* mActType */ gemmGatedAct::ActType(0) + }, gemm::SmVersion::Sm100a }, +{GemmGatedActKernel_Fp32_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin, GemmGatedActKernel_Fp32_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin_len, 86016, "gemmGatedActKernel_Fp32_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a", 448, {{ /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 4 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056777) +, /* mDtypeElt */ trtllm::gen::Dtype(17826819) +, /* mDtypeC */ trtllm::gen::Dtype(1056777) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 4 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) , /* mTileScheduler */ gemm::TileScheduler(0) }, /* mActType */ gemmGatedAct::ActType(0) }, gemm::SmVersion::Sm100a }, diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/KernelParams.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/KernelParams.h index 80ecff7132ae..4a7bde2a1728 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/KernelParams.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/KernelParams.h @@ -38,7 +38,7 @@ CUtensorMap buildNdTmaDescriptor(tg::Dtype dtype, std::vector const& s { CUtensorMap desc{}; // The data type. - CUtensorMapDataType tmaDataFormat; + CUtensorMapDataType tmaDataFormat{CU_TENSOR_MAP_DATA_TYPE_FLOAT32}; if (dtype == tg::Dtype::E4m3) { tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_UINT8; @@ -66,7 +66,7 @@ CUtensorMap buildNdTmaDescriptor(tg::Dtype dtype, std::vector const& s } // The swizzle type. - CUtensorMapSwizzle swizzleType; + CUtensorMapSwizzle swizzleType{CU_TENSOR_MAP_SWIZZLE_NONE}; int32_t tileKSizeInBytes = (tileSizeK * tg::dtypeGetNumBits(dtype)) / /* bits */ 8; if ((tileKSizeInBytes % 128) == 0) { @@ -243,14 +243,14 @@ struct KernelParams // If transposeMmaOutput is false, shape is [K / 128, M]. // Otherwise, shape is [M / 128, K / 128]. // The rightmost dimension is contiguous in memory. - float const* ptrSfA; + void const* ptrSfA; // The scaling factors to dequantize B. // It is used when the DeepSeek FP8 recipe is enabled. Otherwise should be set to nullptr. // If transposeMmaOutput is false, shape is [N / 128, K / 128]. // Otherwise, shape is [K / 128, N]. // The rightmost dimension is contiguous in memory. - float const* ptrSfB; + void const* ptrSfB; // The per-token scaling factors from scale A. // @@ -366,10 +366,94 @@ struct KernelParams return std::make_tuple(shape, stride); } + // Create the TMA shape/stride for A/B block scaling factors. + template + static auto makeTmaShapeStrideSfAb(GemmOptions const& options, MatrixType matrixType, tg::SfLayout layout) + { + // The outer dimension. + auto numTokens = matrixType == MatrixType::MatrixA ? options.mM : options.mN; + // The inner dimension. + auto hiddenSize = options.mK; + // The outer tile dimension. + auto numTokensPerTile = matrixType == MatrixType::MatrixA ? options.mTileM : options.mTileN; + // The inner tile dimension. + auto hiddenSizePerTile = options.mTileK; + // Number of elements per scaling factor. + int32_t const numEltsPerSf = (options.mDtypeElt == tg::Dtype::E2m1) ? 16 : 32; + + switch (layout) + { + case tg::SfLayout::R128c4: + { + // The scaling factor tensor packs 128x4 tiles into contiguous 512B blocks. + // The 512B block maps to a 32x16B (32x128b) block in TMEM. + // See https://nvbugspro.nvidia.com/bug/4165523 + // + // Additionally, we have to meet constraints of TMA that the box dimensions are less + // than 256 and boxDim[0] is a multiple of 16B. + // + // The "logical" tensor is: [outer, inner / numEltsPerSf] + // The aforementioned format is: [⌈outer / 128⌉, inner / (4 * numEltsPerSf), 512] + // The shape we use for TMA is: [⌈outer / 128⌉, inner / (4 * numEltsPerSf), 2, 256] + + auto shape = std::vector{256, 2, static_cast(tg::ceilDiv(hiddenSize, numEltsPerSf * 4)), + static_cast(tg::ceilDiv(numTokens, 128))}; + + std::vector stride(shape.size()); + stride[0] = 1; + for (size_t i = 1; i < shape.size(); i++) + { + stride[i] = shape[i - 1] * stride[i - 1]; + } + + auto tileShapes + = std::vector{256, 2, static_cast(tg::ceilDiv(hiddenSizePerTile, numEltsPerSf * 4)), + static_cast(tg::ceilDiv(numTokensPerTile, 128))}; + + return std::make_tuple(shape, stride, tileShapes); + } + + case tg::SfLayout::R8c4: + { + // The scaling factor tensor packs 8x4 tiles into contiguous 32B blocks. + // + // As the inner dimension (k) is required to be a multiple of the tile size, we + // can reshape to use fewer read requests, if the tile dimensions allow. + // I.e., let's define r = min(⌈hiddenSizePerTile / (numEltsPerSf * 4)⌉, 8) + // + // The "logical" tensor is: [outer, inner / numEltsPerSf] + // The 8x4 SF layout is: [⌈outer / 128⌉, inner / (4 * numEltsPerSf), 32] + // The TMA tensor shape is: [⌈outer / 128⌉, inner / (4 * numEltsPerSf * r), r * 32] + + int const repeats = std::min(tg::ceilDiv(hiddenSizePerTile, numEltsPerSf * 4), 8); + + auto shape = std::vector{static_cast(repeats * 32), + static_cast(tg::ceilDiv(hiddenSize, numEltsPerSf * 4 * repeats)), + static_cast(tg::ceilDiv(numTokens, 8))}; + + std::vector stride(shape.size()); + stride[0] = 1; + for (size_t i = 1; i < shape.size(); i++) + { + stride[i] = shape[i - 1] * stride[i - 1]; + } + + auto tileShapes = std::vector{static_cast(repeats * 32), + static_cast(tg::ceilDiv(hiddenSizePerTile, numEltsPerSf * 4 * repeats)), + static_cast(tg::ceilDiv(numTokensPerTile, 8))}; + + return std::make_tuple(shape, stride, tileShapes); + } + + default: assert(false); + } + return std::make_tuple(std::vector{}, std::vector{}, std::vector{}); + } + // Setup the kernel parameters. template - static KernelParams setKernelParams(GemmOptions_ const& options, void const* ptrA, float const* ptrDqSfsA, - void const* ptrPerTokenSfA, void const* ptrB, float const* ptrDqSfsB, void const* ptrPerTokenSfB, void* ptrC, + static KernelParams setKernelParams(GemmOptions_ const& options, void const* ptrA, void const* ptrSfA, + void const* ptrPerTokenSfA, void const* ptrB, void const* ptrSfB, void const* ptrPerTokenSfB, void* ptrC, float const* ptrScaleC, void* ptrSfC, float const* ptrScaleGate, float* rowMax, uint32_t* rowMaxBars) { @@ -388,6 +472,23 @@ struct KernelParams params.tmaB = gemmGatedAct::buildNdTmaDescriptor( options.mDtypeElt, shapeB, strideB, options.mTileN, options.mTileK, const_cast(ptrB)); + if (options.mDtypeElt == tg::Dtype::E2m1 || options.mDtypeElt == tg::Dtype::MxE4m3) + { + tg::Dtype const dTypeSf = tg::dtypeGetBlockSfType(options.mDtypeElt); + + // Build TMA descriptor for gmem A block scaling factors. + auto [shapeSfA, strideSfA, tileShapesSfA] + = makeTmaShapeStrideSfAb(options, MatrixType::MatrixA, tg::SfLayout::R128c4); + params.tmaSfA + = gemm::buildSfTmaDescriptor(dTypeSf, shapeSfA, strideSfA, tileShapesSfA, const_cast(ptrSfA)); + + // Build TMA descriptor for gmem B block scaling factors. + auto [shapeSfB, strideSfB, tileShapesSfB] + = makeTmaShapeStrideSfAb(options, MatrixType::MatrixB, options.mSfLayoutB); + params.tmaSfB + = gemm::buildSfTmaDescriptor(dTypeSf, shapeSfB, strideSfB, tileShapesSfB, const_cast(ptrSfB)); + } + if (options.mUseTmaStore) { // Shape/stride for gmem tensor C. @@ -403,8 +504,8 @@ struct KernelParams params.ptrC = ptrC; // Set the dequantization factors for A and B when DeepSeek FP8 recipe is used. - params.ptrSfA = ptrDqSfsA; - params.ptrSfB = ptrDqSfsB; + params.ptrSfA = ptrSfA; + params.ptrSfB = ptrSfB; params.ptrSfC = ptrSfC; // Set the per-token scale factors for MetaFP8 or scale inputs diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/KernelTraits.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/KernelTraits.h index 2ed49f33e6b5..1c3d4581c497 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/KernelTraits.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/KernelTraits.h @@ -38,9 +38,11 @@ class MemAllocatorHelper MemAllocatorHelper() {} // Constructor to initialize chunk sizes, alignments, and reuse flags - MemAllocatorHelper(std::vector> const& sizes, std::vector const& reuse) + MemAllocatorHelper(std::vector> const& sizes, std::vector const& reuse, + std::vector const& names) : mNumBytesAndAlignmentPerSmemChunk(sizes) , mFirstChunkReuse(reuse) + , mSmemChunkNames(names) { } @@ -92,6 +94,23 @@ class MemAllocatorHelper return getOffsetBeforeChunk(static_cast(mNumBytesAndAlignmentPerSmemChunk.size())); } + // Returns the first chunk reuse flag for the ith chunk. + int getFirstChunkReuseFlag(int32_t ii) const + { + return mFirstChunkReuse[ii]; + } + + // Print the contents of this object. + void print() const + { + for (size_t ii = 0; ii < mNumBytesAndAlignmentPerSmemChunk.size(); ++ii) + { + printf("Chunk %zd %s: %d bytes, %d alignment, reuse %s, offset %d\n", ii, mSmemChunkNames[ii].c_str(), + mNumBytesAndAlignmentPerSmemChunk[ii].first, mNumBytesAndAlignmentPerSmemChunk[ii].second, + mFirstChunkReuse[ii] ? "true" : "false", getChunkOffset(ii)); + } + } + private: // Helper function to calculate padded size int32_t getSizePaddedToAlignment(int32_t size, int32_t alignment) const @@ -107,6 +126,8 @@ class MemAllocatorHelper std::vector> mNumBytesAndAlignmentPerSmemChunk; // Chunk reuse configuration. True at ith position means that ith chunk starts at smemOffset = 0. std::vector mFirstChunkReuse; + // Buffer names for inspection purposes. + std::vector mSmemChunkNames; }; //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -145,6 +166,8 @@ class KernelTraits std::vector> numBytesAndAlignmentPerSmemChunk; std::vector firstChunkReuseSmem; + // Buffer names for inspection purposes. + std::vector smemChunkNames; // LoadA { @@ -154,8 +177,8 @@ class KernelTraits auto const numBytesAlignmentLoadA = 1024; // loadA is already at first chunk. No need to reuse it. auto const reuseChunksSmemLoadA = false; - // Add info. + smemChunkNames.emplace_back("smemLoadA"); numBytesAndAlignmentPerSmemChunk.emplace_back( std::make_pair(numSmemBytesLoadA, numBytesAlignmentLoadA)); firstChunkReuseSmem.emplace_back(reuseChunksSmemLoadA); @@ -169,8 +192,8 @@ class KernelTraits auto const numBytesAlignmentLoadB = 1024; // No need to reuse the first chunk. auto const reuseChunksSmemLoadB = false; - // Add info. + smemChunkNames.emplace_back("smemLoadB"); numBytesAndAlignmentPerSmemChunk.emplace_back( std::make_pair(numSmemBytesLoadB, numBytesAlignmentLoadB)); firstChunkReuseSmem.emplace_back(reuseChunksSmemLoadB); @@ -192,6 +215,7 @@ class KernelTraits auto const reuseChunksSmemLoadB = false; // Add info. + smemChunkNames.emplace_back("smemBShuffle"); numBytesAndAlignmentPerSmemChunk.emplace_back( std::make_pair(numSmemBytesLoadB, numBytesAlignmentLoadB)); firstChunkReuseSmem.emplace_back(reuseChunksSmemLoadB); @@ -236,6 +260,7 @@ class KernelTraits = doesSplitKUseDsmem(splitK) && resIdx == 0 && !usePersistentScheduler; // Add info. + smemChunkNames.emplace_back("smemGmemC" + std::to_string(resIdx)); numBytesAndAlignmentPerSmemChunk.emplace_back( std::make_pair(numBytesSmemStoreC, numBytesAlignmentStoreC)); firstChunkReuseSmem.emplace_back(reuseFirstChunksSmemStoreC); @@ -252,6 +277,7 @@ class KernelTraits auto const numBytesAlignmentRowMax = 16; // Add info. + smemChunkNames.emplace_back("smemRowMax"); numBytesAndAlignmentPerSmemChunk.emplace_back( std::make_pair(numBytesSmemRowMax, numBytesAlignmentRowMax)); firstChunkReuseSmem.emplace_back(false); @@ -268,6 +294,7 @@ class KernelTraits auto const numBytesAlignmentTile = 16; // Add info. + smemChunkNames.emplace_back("smemSliceK"); numBytesAndAlignmentPerSmemChunk.emplace_back(std::make_pair(numBytesSmemTile, numBytesAlignmentTile)); firstChunkReuseSmem.emplace_back(false); } @@ -280,13 +307,41 @@ class KernelTraits // Number of bytes alignment for per-token scale factors auto const numBytesAlignmentPerTokenSf = 16; // Add info. + smemChunkNames.emplace_back("smemPerTokenSf"); numBytesAndAlignmentPerSmemChunk.emplace_back( std::make_pair(numBytesSmemPerTokenSf, numBytesAlignmentPerTokenSf)); firstChunkReuseSmem.emplace_back(false); } + // Per-block absolute maximum for multi-warp reduction. + { + // Number of bytes: number of epilogue warps * number of tile columns. + // TODO: avoid allocating this memory when it's not needed (it's only for MxFp8 + fusedAct) + auto const numBytesSmemBlockAmax = transposeMmaOutput ? 4 * tileN * sizeof(float) : 0; + // Number of bytes alignment. + auto const numBytesAlignmentBlockAmax = 16; + // Add info. + smemChunkNames.emplace_back("smemBlockAmax"); + numBytesAndAlignmentPerSmemChunk.emplace_back( + std::make_pair(numBytesSmemBlockAmax, numBytesAlignmentBlockAmax)); + firstChunkReuseSmem.emplace_back(false); + } + // Create SMEM helper object. - mSmemAllocatorHelper = MemAllocatorHelper(numBytesAndAlignmentPerSmemChunk, firstChunkReuseSmem); + mSmemAllocatorHelper + = MemAllocatorHelper(numBytesAndAlignmentPerSmemChunk, firstChunkReuseSmem, smemChunkNames); +#if 0 + // E.g., + // Chunk 0 smemLoadA: 32768 bytes, 1024 alignment, false, offset 0 + // Chunk 1 smemLoadB: 32768 bytes, 1024 alignment, false, offset 32768 + // Chunk 2 smemBShuffle: 0 bytes, 1024 alignment, false, offset 65536 + // Chunk 3 smemGmemC0: 65536 bytes, 1024 alignment, true, offset 0 + // Chunk 4 smemGmemC1: 65536 bytes, 1024 alignment, false, offset 65536 + // Chunk 5 smemRowMax: 512 bytes, 16 alignment, false, offset 131072 + // Chunk 6 smemSliceK: 0 bytes, 16 alignment, false, offset 131584 + // Chunk 7 smemPerTokenSf: 0 bytes, 16 alignment, false, offset 131584 + mSmemAllocatorHelper.print(); +#endif } // @@ -296,7 +351,7 @@ class KernelTraits { std::vector> numBytesAndAlignmentPerTmemChunk; std::vector firstChunkReuseTmem; - + std::vector tmemChunkNames; // Matrix D { // Number of columns for accumulators. @@ -308,6 +363,7 @@ class KernelTraits auto const reuseChunksTmemD = false; // Add info. + tmemChunkNames.emplace_back("tmemD"); numBytesAndAlignmentPerTmemChunk.emplace_back(std::make_pair(numTmemColsD, numColsAlignmentD)); firstChunkReuseTmem.emplace_back(reuseChunksTmemD); } @@ -324,6 +380,7 @@ class KernelTraits auto const reuseChunksTmemA = false; // Add info. + tmemChunkNames.emplace_back("tmemA"); numBytesAndAlignmentPerTmemChunk.emplace_back(std::make_pair(numTmemColsA, numColsAlignmentA)); firstChunkReuseTmem.emplace_back(reuseChunksTmemA); } @@ -341,6 +398,7 @@ class KernelTraits auto const reuseChunksTmemSfA = false; // Add info. + tmemChunkNames.emplace_back("tmemSfA"); numBytesAndAlignmentPerTmemChunk.emplace_back(std::make_pair(numTmemColsSfA, numColsAlignmentSfA)); firstChunkReuseTmem.emplace_back(reuseChunksTmemSfA); } @@ -356,12 +414,14 @@ class KernelTraits auto const reuseChunksTmemSfB = false; // Add info. + tmemChunkNames.emplace_back("tmemSfB"); numBytesAndAlignmentPerTmemChunk.emplace_back(std::make_pair(numTmemColsSfB, numColsAlignmentSfB)); firstChunkReuseTmem.emplace_back(reuseChunksTmemSfB); } // Create TMEM helper object. - mTmemAllocatorHelper = MemAllocatorHelper(numBytesAndAlignmentPerTmemChunk, firstChunkReuseTmem); + mTmemAllocatorHelper + = MemAllocatorHelper(numBytesAndAlignmentPerTmemChunk, firstChunkReuseTmem, tmemChunkNames); } } @@ -446,6 +506,21 @@ inline int32_t getSmemOffsetPerTokenSf(KernelTraits traits) return traits.mSmemAllocatorHelper.getChunkOffset(7); } +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetBlockAmax(KernelTraits traits) +{ + return traits.mSmemAllocatorHelper.getChunkOffset(8); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t isSmemAbRepurposedToGmemC(KernelTraits traits, int resIdx = 0) +{ + // Be conscious that the index (3 + resIdx) should match the index in getSmemOffsetGmemC(). + return traits.mSmemAllocatorHelper.getFirstChunkReuseFlag(3 + resIdx); +} + //////////////////////////////////////////////////////////////////////////////////////////////////// // // Starting address of each TMEM buffer. diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/TmaDescriptor.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/TmaDescriptor.h index 46ee0176078e..8d26c4b972bb 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/TmaDescriptor.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/TmaDescriptor.h @@ -41,8 +41,8 @@ inline CUtensorMap buildNdTmaDescriptor(tg::Dtype dtype, std::vector c { CUtensorMap desc{}; // The data type. - CUtensorMapDataType tmaDataFormat; - if (dtype == tg::Dtype::E4m3) + CUtensorMapDataType tmaDataFormat{CU_TENSOR_MAP_DATA_TYPE_FLOAT32}; + if (dtype == tg::Dtype::E4m3 || dtype == tg::Dtype::MxE4m3) { tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_UINT8; } @@ -64,7 +64,7 @@ inline CUtensorMap buildNdTmaDescriptor(tg::Dtype dtype, std::vector c } else { - std::cerr << "Unexpected dtype " << static_cast(dtype) << std::endl; + std::cerr << "buildNdTmaDescriptor: unexpected dtype " << static_cast(dtype) << std::endl; assert(false); } @@ -87,7 +87,7 @@ inline CUtensorMap buildNdTmaDescriptor(tg::Dtype dtype, std::vector c } else { - std::cerr << "Unexpected tileKSizeInBytes " << tileKSizeInBytes << std::endl; + std::cerr << "buildNdTmaDescriptor: unexpected tileKSizeInBytes " << tileKSizeInBytes << std::endl; assert(false); } } @@ -187,13 +187,13 @@ inline CUtensorMap buildSfTmaDescriptor(tg::Dtype dtype, std::vector c { CUtensorMap desc{}; CUtensorMapDataType tmaDataFormat; - if (dtype == tg::Dtype::E4m3) + if (dtype == tg::Dtype::E4m3 || dtype == tg::Dtype::UE8m0) { tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_UINT8; } else { - std::cerr << "Unexpected dtype " << static_cast(dtype) << std::endl; + std::cerr << "buildSfTmaDescriptor: unexpected dtype " << static_cast(dtype) << std::endl; assert(false); } diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Bfloat16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Bfloat16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp new file mode 100644 index 000000000000..5a1179b298f2 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Bfloat16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9434907544b528b0f46ab2472402b64bf13a3405745d2c1ce02b2f962bc7c774 +size 297623 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Bfloat16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Bfloat16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp new file mode 100644 index 000000000000..4df2e220d6e8 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Bfloat16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b00b1627dea02258b7b5d895de301efd75c742ac468a220865aea997a704fb4c +size 331481 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp index aa39ccb6ae38..e8d20c23648f 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:6f50aff19cada4c88c025dc24cc7961dcae19980dbe80b633ad7701bb37c79b6 -size 292515 +oid sha256:418acca7365bcedd4fa73d756169184484666ffb5bfee08a281946e6c4d6d94b +size 297745 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp index d45d6b6c3c35..e613213678cb 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:153e621afd809f5aa5a8f9fd344e7b2698079fae3bab6ccb9513bc65dbf0f814 -size 219529 +oid sha256:98ee8d43aa5f90e048eb334711b861de202f769457cc4db1703ef71d6c75ccf1 +size 222291 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp index 7c014bc33d11..abf57d4c8fc8 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Bfloat16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:7294c71c6a4933aa8d68957639cf18cf4b5b0f04d4fcb33693f468937f812615 -size 249193 +oid sha256:e85b4cb09dda8f57b068b53030f19cdd0552407381bbc5bd50e32490e18c5e16 +size 254473 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_E2m1_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_E2m1_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp new file mode 100644 index 000000000000..409920f5954e --- /dev/null +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_E2m1_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a98affa59323662168205c915685f87a95c2408fad1c406056f226ec8739333b +size 305853 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_E2m1_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_E2m1_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp new file mode 100644 index 000000000000..ee2186243f79 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_E2m1_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:460cf57675a1137a9cf0d851413404e2bf2ace1a52a6b427235f94ba1d580efa +size 342425 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp index 621443bb8798..72f329db0551 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:7c84cbf557bf39fc79489f4e87f3f05382b9a9890bbcfe102f45b24863fa240d -size 297935 +oid sha256:1dbdc8e54e2187ffb345d5a797c9a7039066fb746f253220121095af454924dd +size 300203 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp index b77e5cd18fea..bd8f1bf51e22 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:5e09fdd64abe108f305a4b03aa787793bb3e84b3c054af29caa05149a9c0662f -size 217893 +oid sha256:ff17b22e983563501c3339ee7160475e2086cada9e957afb9a1ff89e901b8446 +size 221445 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp index eec9c69e563f..a94e66978ca3 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_E4m3_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:e560ba86ae47c1c215dc0107717daf44639b4ddee3ed8dac2c09b810b9024d8d -size 249137 +oid sha256:ca2dbad0de7cd53b39225d96ceaa071e1b27ce3feff861691b7ada7bde9e9411 +size 254415 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp new file mode 100644 index 000000000000..cea5c48ce0e4 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8b99d98321f44aea8334cfa4dea7317fc0e7dd2eec9b972106bfae5e1bf740bf +size 296825 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp new file mode 100644 index 000000000000..304dcd84b479 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp16_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3ef2353b8bcd016f7f35df660d958cb1b310a944915f655d6984441b264b6855 +size 330685 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp index 4a2c257c2e3e..21625a8ff00a 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x128x256_epilogueTile128x128_mma128x128x32_cluster1x1x1_sm100a_cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:83dfb999e17c4b339689a4471c24fe318e48e002baabf1654e8d4f93d1693aab -size 290929 +oid sha256:3841f0321dbfe7785af9b2c10f4178b664bf6b8166b02ac670ba379868a292b8 +size 296947 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp index 1d3fb606c53c..e110963d8a88 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:a9ae4835c2cb293a9561919a0c984f6f392ed5a7ea16b2e479b2b585208574fd -size 217941 +oid sha256:4915aeb2ad55d39bc8dc15d9d528710452271ac3959b44c93fdee4fe625dfdb8 +size 221493 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp index 91b09b217fa8..0bc89b9de35c 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp16_E4m3_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x32_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:f85ff9f9eb7760cf5515abd5cde453d87e0d2953f564400cdd74d7d6ac2178cd -size 249185 +oid sha256:9416865d0b59145c1ed6db76342e013dfa21137679aac4b56a44cf862fec0f62 +size 254465 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp32_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp32_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp new file mode 100644 index 000000000000..c9201422eb5a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp32_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x1_transposeMmaOutput_sm100a_cubin.cpp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8570922d77439b30ef9f1bd0f1de2e48cd0688c52d4c5d7b48a8ab9df80678bb +size 297023 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp32_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp32_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp new file mode 100644 index 000000000000..108db9b8d148 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/cubins/GemmGatedActKernel_Fp32_E2m1_Fp32_tile128x8x256_epilogueTile128x8_mma128x8x64_cluster1x1x4_splitK4_transposeMmaOutput_sm100a_cubin.cpp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6d9b1bac0bdb963a4b5ae0d2c63c77495dcab658079b742b59ec9c9965c9fbea +size 331573 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/trtllm/gen/DtypeDecl.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/trtllm/gen/DtypeDecl.h index dd8eb6f3924c..a6892f12ca09 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/trtllm/gen/DtypeDecl.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/gemmGatedAct/trtllmGen_gatedAct_export/trtllm/gen/DtypeDecl.h @@ -67,13 +67,14 @@ enum class Dtype : uint32_t Int32 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 1u, /*bits*/ 32u, /*uid*/ 11u), Int64 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 1u, /*bits*/ 64u, /*uid*/ 12u), MxE2m1 = TLLM_ENCODE_DTYPE(/*block*/ 1u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 4u, /*uid*/ 13u), - UE8m0 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 0u, /*bits*/ 8u, /*uid*/ 14u), - UInt8 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 8u, /*uid*/ 15u), - UInt16 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 16u, /*uid*/ 16u), - UInt32 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 32u, /*uid*/ 17u), - UInt64 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 64u, /*uid*/ 18u), - UInt128 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 128u, /*uid*/ 19u), - Void = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 0u, /*uid*/ 20u), + MxE4m3 = TLLM_ENCODE_DTYPE(/*block*/ 1u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 8u, /*uid*/ 14u), + UE8m0 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 0u, /*bits*/ 8u, /*uid*/ 15u), + UInt8 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 8u, /*uid*/ 16u), + UInt16 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 16u, /*uid*/ 17u), + UInt32 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 32u, /*uid*/ 18u), + UInt64 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 64u, /*uid*/ 19u), + UInt128 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 128u, /*uid*/ 20u), + Void = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 0u, /*uid*/ 21u), // clang-format on #undef TLLM_ENCODE_DTYPE @@ -151,6 +152,7 @@ inline std::string dtypeToString(Dtype dtype) case Dtype::Int8: return "Int8"; case Dtype::Int32: return "Int32"; case Dtype::Int64: return "Int64"; + case Dtype::MxE4m3: return "MxE4m3"; case Dtype::UE8m0: return "UE8m0"; case Dtype::UInt8: return "UInt8"; case Dtype::UInt16: return "UInt16"; @@ -164,5 +166,44 @@ inline std::string dtypeToString(Dtype dtype) //////////////////////////////////////////////////////////////////////////////////////////////////// +inline Dtype dtypeEltType(Dtype dtype) +{ + switch (dtype) + { + case Dtype::MxE2m1: return Dtype::E2m1; + case Dtype::MxE4m3: return Dtype::E4m3; + default: return dtype; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int dtypeNumEltsPerSf(Dtype dtype) +{ + switch (dtype) + { + case Dtype::E2m1: return 16; + case Dtype::MxE2m1: + case Dtype::MxE4m3: return 32; + default: assert(false); return -1; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Returns the dtype of scaling factors, if applicable. +inline Dtype dtypeGetBlockSfType(Dtype dtype) +{ + switch (dtype) + { + case Dtype::E2m1: return Dtype::E4m3; + case Dtype::MxE2m1: + case Dtype::MxE4m3: return Dtype::UE8m0; + default: assert(false); return Dtype::Void; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + } // namespace gen } // namespace trtllm diff --git a/cpp/tensorrt_llm/runtime/ipcUtils.cpp b/cpp/tensorrt_llm/runtime/ipcUtils.cpp index d3f8041ded3e..23a7e28a4f27 100644 --- a/cpp/tensorrt_llm/runtime/ipcUtils.cpp +++ b/cpp/tensorrt_llm/runtime/ipcUtils.cpp @@ -149,7 +149,7 @@ AllReduceBuffers::AllReduceBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWi { auto const tpSize = worldConfig.getTensorParallelism(); mAllReduceCommPtrs = BufferManager::cpu( - ITensor::makeShape({static_cast(7) * tpSize + 2}), nvinfer1::DataType::kINT64); + ITensor::makeShape({static_cast(7) * tpSize + 3}), nvinfer1::DataType::kINT64); } else { @@ -168,7 +168,7 @@ AllReduceBuffers::AllReduceBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWi ? 1 // zero size is not allowed for IpcMemory::allocateIpcMemory. : (tpSize * tensorrt_llm::kernels::reduce_fusion::details::kLamportTokenNumThreshold * realHiddenSize * sizeof(half)); - auto const flagsSize = IpcMemory::FLAGS_SIZE * tpSize * 2; + auto const flagsSize = IpcMemory::FLAGS_SIZE * tpSize * 2 * tpSize; for (auto size : {bufferSize, bufferSize, flagsSize, flagsSize, lamportBufferSize, lamportBufferSize, lamportBufferSize}) @@ -177,13 +177,15 @@ AllReduceBuffers::AllReduceBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWi } mAllReduceCommPtrs - = BufferManager::cpu(ITensor::makeShape({static_cast(mIpcMemoryHandles.size()) * tpSize + 2}), + = BufferManager::cpu(ITensor::makeShape({static_cast(mIpcMemoryHandles.size()) * tpSize + 3}), nvinfer1::DataType::kINT64); auto commPtrs = BufferRange(*mAllReduceCommPtrs); - auto const CustomARFlagPtr = static_cast(mAllReduceCommPtrs->data(mAllReduceCommPtrs->getSize() - 1)); - auto const LamportFlagPtr = static_cast(mAllReduceCommPtrs->data(mAllReduceCommPtrs->getSize() - 2)); - *CustomARFlagPtr = 0; - *LamportFlagPtr = 0; + // Start from 1 since 0 represents released state for barrier at the beginning of the all_reduce. + // The last element is the barrier flag counter. + mFlagPtrs = manager.copyFrom(std::vector{1, 1, 0}, ITensor::makeShape({3}), MemoryType::kGPU); + commPtrs[mIpcMemoryHandles.size() * tpSize] = mFlagPtrs->data(); + commPtrs[mIpcMemoryHandles.size() * tpSize + 1] = mFlagPtrs->data(1); + commPtrs[mIpcMemoryHandles.size() * tpSize + 2] = mFlagPtrs->data(2); for (std::size_t memIdx = 0; memIdx < mIpcMemoryHandles.size(); memIdx++) { diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 62f4fac2de62..eeb1a400ddf0 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -59,6 +59,7 @@ add_library( fusedTopkSoftmax.cpp gatherTreeOp.cpp groupRmsNormOp.cpp + llama4MinLatency.cpp logitsBitmaskOp.cpp mambaConv1dOp.cpp moeOp.cpp diff --git a/cpp/tensorrt_llm/thop/llama4MinLatency.cpp b/cpp/tensorrt_llm/thop/llama4MinLatency.cpp new file mode 100644 index 000000000000..53873e3d27d7 --- /dev/null +++ b/cpp/tensorrt_llm/thop/llama4MinLatency.cpp @@ -0,0 +1,219 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/dataType.h" +#include "tensorrt_llm/common/opUtils.h" +#include "tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Bf16Bf16Gemm.h" +#include "tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Bf16Gemm.h" +#include "tensorrt_llm/kernels/llama4MinLatencyKernels/llama4Fp8Fp8GemmSwiGLU.h" +#include "tensorrt_llm/kernels/llama4MinLatencyKernels/llama4MinLatencyMoEOp.h" +#include "tensorrt_llm/runtime/torchUtils.h" +#include "tensorrt_llm/runtime/utils/mpiUtils.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include +#include + +#include +#include +#include + +namespace torch_ext +{ + +torch::Tensor llama4_bf16_bf16_gemm(torch::Tensor const& inputA, torch::Tensor const& inputB) +{ + CHECK_INPUT(inputA, c10::ScalarType::BFloat16); + CHECK_INPUT(inputB, c10::ScalarType::BFloat16); + + TORCH_CHECK(inputA.dim() == 2, "inputA must be 2D."); + TORCH_CHECK(inputB.dim() == 2, "inputB must be 2D."); + TORCH_CHECK(inputA.sizes()[1] == 5120, "inputA.size(1) must be 5120"); + TORCH_CHECK(inputB.sizes()[0] == 128, "inputB.size(0) must be 128"); + TORCH_CHECK(inputB.sizes()[1] == 5120, "inputB.size(1) must be 5120"); + + auto const num_tokens = inputA.sizes()[0]; + auto const hidden_out = inputB.sizes()[0]; + + auto stream = at::cuda::getCurrentCUDAStream(inputA.get_device()); + auto output = torch::empty( + {num_tokens, hidden_out}, torch::TensorOptions().dtype(c10::ScalarType::BFloat16).device(inputA.device())); + + tensorrt_llm::kernels::llama4_min_latency::llama4_bf16_bf16_gemm::llama4_bf16_bf16_gemm_op( + num_tokens, inputA.data_ptr(), inputB.data_ptr(), output.data_ptr(), stream); + + return output; +} + +torch::Tensor llama4_fp8_bf16_gemm(torch::Tensor const& inputA, torch::Tensor const& inputB, + torch::Tensor const& scaling_factor, torch::optional const& position_ids) +{ + CHECK_INPUT(inputA, c10::ScalarType::Float8_e4m3fn); + CHECK_INPUT(inputB, c10::ScalarType::Float8_e4m3fn); + CHECK_INPUT(scaling_factor, c10::ScalarType::Float); + + bool const has_position_ids = position_ids.has_value(); + if (has_position_ids) + { + CHECK_TH_CUDA(position_ids.value()); + CHECK_CONTIGUOUS(position_ids.value()); + if (position_ids.value().scalar_type() != c10::ScalarType::Long + && position_ids.value().scalar_type() != c10::ScalarType::Int) + { + TORCH_CHECK(false, "position_ids must be a Long or Int tensor."); + } + } + + bool const position_ids_int64 = has_position_ids && position_ids.value().dtype() == c10::ScalarType::Long; + + TORCH_CHECK(inputA.dim() == 2, "inputA must be 2D."); + TORCH_CHECK(inputB.dim() == 2, "inputB must be 2D."); + TORCH_CHECK(inputA.sizes()[1] == 5120, "inputA.size(1) must be 5120"); + TORCH_CHECK(inputB.sizes()[1] == 5120, "inputB.size(1) must be 5120"); + TORCH_CHECK(scaling_factor.dim() == 0, "scaling_factor must be a scalar tensor"); + + auto const num_tokens = inputA.sizes()[0]; + auto const hidden_in = inputA.sizes()[1]; + auto const hidden_out = inputB.sizes()[0]; + + auto stream = at::cuda::getCurrentCUDAStream(inputA.get_device()); + auto output = torch::empty( + {num_tokens, hidden_out}, torch::TensorOptions().dtype(c10::ScalarType::BFloat16).device(inputA.device())); + + tensorrt_llm::kernels::llama4_min_latency::llama4_fp8_bf16_gemm::llama4_fp8_bf16_gemm_op(inputA.data_ptr(), + inputB.data_ptr(), output.data_ptr(), scaling_factor.data_ptr(), + has_position_ids ? position_ids.value().data_ptr() : nullptr, position_ids_int64, num_tokens, hidden_in, + hidden_out, stream); + + return output; +} + +torch::Tensor llama4_fp8_fp8_gemm_swiglu(torch::Tensor const& inputA, torch::Tensor const& inputB, + torch::Tensor const& in_scale, torch::Tensor const& out_scale_inv) +{ + CHECK_INPUT(inputA, c10::ScalarType::Float8_e4m3fn); + CHECK_INPUT(inputB, c10::ScalarType::Float8_e4m3fn); + CHECK_INPUT(in_scale, c10::ScalarType::Float); + CHECK_INPUT(out_scale_inv, c10::ScalarType::Float); + + TORCH_CHECK(inputA.dim() == 2, "inputA must be 2D."); + TORCH_CHECK(inputB.dim() == 2, "inputB must be 2D."); + TORCH_CHECK(inputA.sizes()[1] == 5120, "inputA.size(1) must be 5120"); + TORCH_CHECK(inputB.sizes()[1] == 5120, "inputB.size(1) must be 5120"); + TORCH_CHECK(in_scale.dim() == 0, "in_scale must be a scalar tensor"); + TORCH_CHECK(out_scale_inv.dim() == 0, "out_scale_inv must be a scalar tensor"); + + auto stream = at::cuda::getCurrentCUDAStream(inputA.get_device()); + auto const num_tokens = inputA.sizes()[0]; + auto const hidden_in = inputA.sizes()[1]; + auto const hidden_out = inputB.sizes()[0] / 2; + auto output = torch::empty( + {num_tokens, hidden_out}, torch::TensorOptions().dtype(c10::ScalarType::Float8_e4m3fn).device(inputA.device())); + + tensorrt_llm::kernels::llama4_min_latency::llama4_fp8_fp8_gemm_swiglu::llama4_fp8_fp8_gemm_swiglu_op(num_tokens, + hidden_in, hidden_out, inputA.data_ptr(), inputB.data_ptr(), output.data_ptr(), in_scale.data_ptr(), + out_scale_inv.data_ptr(), stream); + + return output; +} + +torch::Tensor llama4_moe_tp8ep1_min_latency(torch::Tensor const& input, torch::Tensor const& router_logits, + torch::Tensor const& fc1_expert_weights, torch::Tensor const& fc2_expert_weights, + torch::optional> const& quant_scales) +{ + CHECK_INPUT(input, c10::ScalarType::Float8_e4m3fn) + CHECK_INPUT(router_logits, c10::ScalarType::BFloat16) + CHECK_INPUT(fc1_expert_weights, c10::ScalarType::Float8_e4m3fn) + CHECK_INPUT(fc2_expert_weights, c10::ScalarType::Float8_e4m3fn) + + TORCH_CHECK(input.dim() == 2, "input must be 2D."); + TORCH_CHECK(router_logits.dim() == 2, "router_logits must be 2D."); + TORCH_CHECK(input.sizes()[0] == router_logits.sizes()[0], "input and router_logits must have the same num tokens."); + + TORCH_CHECK(fc1_expert_weights.dim() == 3, "fc1_expert_weights must be 3D."); + TORCH_CHECK(fc2_expert_weights.dim() == 3, "fc2_expert_weights must be 3D."); + TORCH_CHECK(fc1_expert_weights.sizes()[0] == fc2_expert_weights.sizes()[0], + "fc1_expert_weights and fc2_expert_weights must have the same number of experts."); + TORCH_CHECK(fc1_expert_weights.sizes()[1] == fc2_expert_weights.sizes()[2] * 2, + "fc1_expert_weights inter size must be 2 times fc2_expert_weights inter size."); + TORCH_CHECK(router_logits.sizes()[0] == input.sizes()[0], "router_logits and input must have the same num tokens."); + TORCH_CHECK(router_logits.sizes()[1] == 128, "router_logits must have 128 experts."); + TORCH_CHECK(input.sizes()[1] == 5120, "input must have 5120 hidden size."); + TORCH_CHECK(fc1_expert_weights.sizes()[2] == 5120, "fc1_expert_weights must have 5120 hidden size."); + + int64_t num_rows = input.sizes()[0]; + int64_t hidden_size = fc2_expert_weights.sizes()[1]; + int64_t inter_size = fc2_expert_weights.sizes()[2]; + int64_t num_experts = fc2_expert_weights.sizes()[0]; + + TORCH_CHECK(quant_scales.has_value(), "Expecting quant scales for fp8 quantization"); + TORCH_CHECK(quant_scales.value().size() == 4, "Expecting 4 quant scales for fp8 quantization"); + + auto const fc1_dequant = quant_scales.value()[0]; + auto const fc2_quant = quant_scales.value()[1]; + auto const fc2_dequant = quant_scales.value()[2]; + + CHECK_INPUT(fc1_dequant, c10::ScalarType::Float); + CHECK_INPUT(fc2_quant, c10::ScalarType::Float); + CHECK_INPUT(fc2_dequant, c10::ScalarType::Float); + TORCH_CHECK(fc1_dequant.dim() == 1, "fc1 dequant must be 1D"); + TORCH_CHECK(fc2_quant.dim() == 0, "fc2 quant must be a scalar tensor"); + TORCH_CHECK(fc2_dequant.dim() == 1, "fc2 dequant must be 1D"); + TORCH_CHECK(fc1_dequant.sizes()[0] == num_experts, "fc1 dequant size must be (num_experts,)"); + TORCH_CHECK(fc2_dequant.sizes()[0] == num_experts, "fc2 dequant size must be (num_experts,)"); + + auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); + + std::vector fc2_input_shape = {num_rows, inter_size}; + auto fc2_input = torch::empty(fc2_input_shape, input.options().dtype(c10::ScalarType::Float8_e4m3fn)); + std::vector exp_idx_shape = {num_rows}; + auto exp_idx = torch::empty(exp_idx_shape, input.options().dtype(c10::ScalarType::Int)); + std::vector output_shape = {num_rows, hidden_size}; + auto output = torch::empty(output_shape, input.options().dtype(c10::ScalarType::BFloat16)); + + tensorrt_llm::kernels::llama4_min_latency::llama4_moe::run_moe_llama4_tp8ep1_min_latency(num_rows, num_experts, + input.const_data_ptr(), router_logits.const_data_ptr(), fc1_expert_weights.const_data_ptr(), + fc2_expert_weights.const_data_ptr(), static_cast(fc1_dequant.data_ptr()), + static_cast(fc2_quant.data_ptr()), static_cast(fc2_dequant.data_ptr()), + fc2_input.data_ptr(), static_cast(exp_idx.data_ptr()), output.data_ptr(), stream); + + return output; +} + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("llama4_bf16_bf16_gemm(Tensor inputA, Tensor inputB) -> Tensor"); + m.def( + "llama4_fp8_bf16_gemm(Tensor inputA, Tensor inputB, " + "Tensor scaling_factor, Tensor? position_ids=None) -> Tensor"); + m.def("llama4_fp8_fp8_gemm_swiglu(Tensor inputA, Tensor inputB, Tensor in_scale, Tensor out_scale_inv) -> Tensor"); + m.def( + "llama4_moe_tp8ep1_min_latency(Tensor input, Tensor router_logits, " + "Tensor fc1_expert_weights, Tensor fc2_expert_weights, " + "Tensor[]? quant_scales=None) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("llama4_bf16_bf16_gemm", &torch_ext::llama4_bf16_bf16_gemm); + m.impl("llama4_fp8_bf16_gemm", &torch_ext::llama4_fp8_bf16_gemm); + m.impl("llama4_fp8_fp8_gemm_swiglu", &torch_ext::llama4_fp8_fp8_gemm_swiglu); + m.impl("llama4_moe_tp8ep1_min_latency", &torch_ext::llama4_moe_tp8ep1_min_latency); +} + +} // namespace torch_ext diff --git a/cpp/tests/unit_tests/kernels/allReduce/allReduceKernelTest.cu b/cpp/tests/unit_tests/kernels/allReduce/allReduceKernelTest.cu index f4603338e9d5..ac85131a2d76 100644 --- a/cpp/tests/unit_tests/kernels/allReduce/allReduceKernelTest.cu +++ b/cpp/tests/unit_tests/kernels/allReduce/allReduceKernelTest.cu @@ -212,12 +212,26 @@ public: { } - void set_params(AllReduceParams& params) const + void set_params(AllReduceParams& params, nvinfer1::DataType dataType, int token_num, int hidden_size, + AllReduceFusionOp op) const { int world_size = world_config.getSize(); + int flag_offset = 0; + if (op == AllReduceFusionOp::RESIDUAL_RMS_NORM + && reduce_fusion::is_lamport_supported(dataType, token_num, hidden_size)) + { + flag_offset = 0; + } + else + { + flag_offset = 1; + } + params.barrier_flag_ptr = static_cast(buffers.mFlagPtrs->data(flag_offset)); + params.barrier_flag_counter_ptr = static_cast(buffers.mFlagPtrs->data(2)); for (int i = 0; i < world_size; ++i) { params.peer_comm_buffer_ptrs[i] = buffers.mIpcMemoryHandles[0].getCommPtrs()[i]; + params.peer_comm_buffer_ptrs[i + MAX_RANKS_PER_NODE] = buffers.mIpcMemoryHandles[1].getCommPtrs()[i]; params.fusion_params.lamport_peer_comm_buffer_ptrs[i] = buffers.mIpcMemoryHandles[4].getCommPtrs()[i]; params.fusion_params.lamport_peer_comm_buffer_ptrs[i + MAX_RANKS_PER_NODE] = buffers.mIpcMemoryHandles[5].getCommPtrs()[i]; @@ -263,6 +277,7 @@ bool test(Workspace const& workspace, int token_num, int hidden_size, bool has_b std::vector bias_buffer(hidden_size); std::vector inter_buffer(message_size); std::vector output_buffer(message_size); + std::vector barrier_flag_buffer{1}; float eps = 1e-6; random_fill(residual_buffer, -1, 1); random_fill(weight_buffer, -1, 1); @@ -301,8 +316,7 @@ bool test(Workspace const& workspace, int token_num, int hidden_size, bool has_b in.copy_from(input_buffer.data()); AllReduceParams params; - workspace.set_params(params); - params.barrier_flag = 0; + workspace.set_params(params, nvinfer1::DataType::kHALF, token_num, hidden_size, fusion_op); params.ranks_per_node = world_size; params.local_rank = rank; params.local_output_buffer_ptr = out.data(); @@ -329,13 +343,11 @@ bool test(Workspace const& workspace, int token_num, int hidden_size, bool has_b comm.barrier(); for (int i = 0; i < warmup; ++i) { - params.barrier_flag += 1; customAllReduce(params, nvinfer1::DataType::kHALF, runtime_strategy, config, fusion_op, s); } cudaEventRecord(begin, s); for (int i = 0; i < iter; ++i) { - params.barrier_flag += 1; customAllReduce(params, nvinfer1::DataType::kHALF, runtime_strategy, config, fusion_op, s); } cudaEventRecord(end, s); @@ -399,7 +411,8 @@ bool test_prepostnorm(Workspace const& workspace, int token_num, int hidden_size int message_size = token_num * hidden_size; int buffer_size = sizeof(half) * message_size; CudaBuffer in(buffer_size), out(buffer_size), residual(buffer_size), weight(hidden_size * sizeof(half)), - weight_pre_residual_norm(hidden_size * sizeof(half)), inter(buffer_size), bias(hidden_size * sizeof(half)); + weight_pre_residual_norm(hidden_size * sizeof(half)), inter(buffer_size), bias(hidden_size * sizeof(half)), + barrier_flag(sizeof(uint32_t)); std::vector input_buffer(message_size); std::vector residual_buffer(message_size); std::vector weight_buffer(hidden_size); @@ -407,6 +420,7 @@ bool test_prepostnorm(Workspace const& workspace, int token_num, int hidden_size std::vector bias_buffer(hidden_size); std::vector inter_buffer(message_size); std::vector output_buffer(message_size); + std::vector barrier_flag_buffer{1}; float eps = 1e-6; random_fill(residual_buffer, -1, 1); random_fill(weight_buffer, -1, 1); @@ -420,6 +434,7 @@ bool test_prepostnorm(Workspace const& workspace, int token_num, int hidden_size bias.copy_from(bias_buffer.data()); inter.copy_from(inter_buffer.data()); out.copy_from(output_buffer.data()); + barrier_flag.copy_from(barrier_flag_buffer.data()); auto& comm = mpi::MpiComm::world(); auto world_size = comm.getSize(); auto rank = comm.getRank(); @@ -447,8 +462,7 @@ bool test_prepostnorm(Workspace const& workspace, int token_num, int hidden_size in.copy_from(input_buffer.data()); AllReduceParams params; - workspace.set_params(params); - params.barrier_flag = 0; + workspace.set_params(params, nvinfer1::DataType::kHALF, token_num, hidden_size, fusion_op); params.ranks_per_node = world_size; params.local_rank = rank; params.local_output_buffer_ptr = out.data(); @@ -470,13 +484,11 @@ bool test_prepostnorm(Workspace const& workspace, int token_num, int hidden_size comm.barrier(); for (int i = 0; i < warmup; ++i) { - params.barrier_flag += 1; customAllReduce(params, nvinfer1::DataType::kHALF, runtime_strategy, config, fusion_op, s); } cudaEventRecord(begin, s); for (int i = 0; i < iter; ++i) { - params.barrier_flag += 1; customAllReduce(params, nvinfer1::DataType::kHALF, runtime_strategy, config, fusion_op, s); } cudaEventRecord(end, s); @@ -568,7 +580,6 @@ TEST(Kernel, AllReduce) // clang-format off std::vector configs{ AllReduceStrategyConfig(0), - AllReduceStrategyConfig::USE_MEMCPY, AllReduceStrategyConfig::PUSH_MODE }; std::vector ops{ diff --git a/tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py b/tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py index 744b1aa7f6ea..15e4b3750f2f 100644 --- a/tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py @@ -38,7 +38,7 @@ def flashinfer_silu_and_mul(x: torch.Tensor) -> torch.Tensor: @flashinfer_silu_and_mul.register_fake def _(x: torch.Tensor) -> torch.Tensor: - return torch.empty_like(x).chunk(2, dim=-1)[1] + return torch.empty_like(x).chunk(2, dim=-1)[1].contiguous() # Warp this into custom op since flashinfer provides default value for eps with would produce two different graphs depends on the eps value. @torch.library.custom_op("trtllm::flashinfer_rmsnorm", mutates_args=()) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index ef86ed06844c..1616bcf0cd0e 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -86,6 +86,14 @@ def __post_init__(self): self.is_generation = self.is_generation_model( self.pretrained_config.architectures) + @property + def fuse_pos_embd(self): + if self.attn_backend == 'TRTLLM': + return True + elif self.attn_backend == 'FLASHINFER': + return False + return False + @property def enable_flash_mla(self): if self.attn_backend == 'TRTLLM': diff --git a/tensorrt_llm/_torch/models/modeling_clip.py b/tensorrt_llm/_torch/models/modeling_clip.py index 178e750b0593..546375720bf4 100644 --- a/tensorrt_llm/_torch/models/modeling_clip.py +++ b/tensorrt_llm/_torch/models/modeling_clip.py @@ -232,4 +232,4 @@ def load_weights(self, weights: Dict): r'(.*?)mlp\.fc1(.*)': r'\1mlp.up_proj\2', r'(.*?)mlp\.fc2(.*)': r'\1mlp.down_proj\2', } - _load_weights_impl(self, weights, pattern_mapping) + _load_weights_impl(self, weights, params_map=pattern_mapping) diff --git a/tensorrt_llm/_torch/models/modeling_llama.py b/tensorrt_llm/_torch/models/modeling_llama.py index 36163e6b10d6..f80b3bf6f647 100644 --- a/tensorrt_llm/_torch/models/modeling_llama.py +++ b/tensorrt_llm/_torch/models/modeling_llama.py @@ -1,5 +1,5 @@ import copy -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple, Union import torch from PIL.Image import Image @@ -33,7 +33,8 @@ WeightsLoadingConfig) from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm -from ..speculative import Eagle3SpecMetadata, SpecMetadata +from ..speculative import SpecMetadata, get_spec_worker +from ..utils import Fp4QuantizedTensor from .modeling_multimodal_utils import fuse_input_embeds from .modeling_utils import (DecoderModel, DecoderModelForCausalLM, EagerFusionConfig, register_auto_model) @@ -129,17 +130,22 @@ def _get_attn_scale(position_ids: torch.Tensor) -> torch.Tensor: def _forward_nope( self, position_ids: Optional[torch.LongTensor], - hidden_states: torch.Tensor, + hidden_states: Union[torch.Tensor, Fp4QuantizedTensor], attn_metadata: AttentionMetadata, attention_mask: PredefinedAttentionMask = PredefinedAttentionMask. CAUSAL, mrope_config: Optional[dict] = None, all_reduce_params: Optional[AllReduceParams] = None, + skip_attn_scaling: bool = False, ): + qkv = self.qkv_proj(hidden_states) - q, k, v = self.split_qkv(qkv) - if self.attn_temperature_tuning: + + q, k, v = qkv, None, None + if self.attn_temperature_tuning and not skip_attn_scaling: + q, k, v = self.split_qkv(q, k, v) q = self._attention_scaling(q, position_ids) + out_scale = None if self.o_proj.has_fp8_qdq or self.o_proj.has_nvfp4 or self.o_proj.has_fp8_block_scales: out_scale = self.o_proj.inv_input_scale @@ -161,7 +167,7 @@ def _forward_nope( def forward( self, position_ids: Optional[torch.LongTensor], - hidden_states: torch.Tensor, + hidden_states: Union[torch.Tensor, Fp4QuantizedTensor], attn_metadata: AttentionMetadata, attention_mask: PredefinedAttentionMask = PredefinedAttentionMask. CAUSAL, @@ -233,6 +239,18 @@ def __init__( config = model_config.pretrained_config self.enable_attention_dp = model_config.mapping.enable_attention_dp self.top_k = top_k + + # Create shared_expert before experts because in min-latency mode the experts depend on the scaling factors of + # shared_expert. + self.shared_expert = GatedMLP( + hidden_size=hidden_size, + intermediate_size=shared_expert_intermediate_size, + bias=False, + dtype=dtype, + config=model_config, + overridden_tp_size=1 if self.enable_attention_dp else None, + reduce_output=False) + self.experts = FusedMoE( routing_method=Llama4RenormalizeMoeRoutingMethod(top_k), num_experts=num_experts, @@ -245,15 +263,6 @@ def __init__( model_config=model_config, apply_router_weight_on_input=True) - self.shared_expert = GatedMLP( - hidden_size=hidden_size, - intermediate_size=shared_expert_intermediate_size, - bias=False, - dtype=dtype, - config=model_config, - overridden_tp_size=1 if self.enable_attention_dp else None, - reduce_output=False) - self.router = Linear(hidden_size, num_experts, bias=False, @@ -324,6 +333,11 @@ def __init__( self.layer_idx = layer_idx self.is_quanted = model_config.quant_config and model_config.quant_config.quant_mode.has_any_quant( ) + self.is_fp8_quant = self.is_quanted and model_config.quant_config.quant_mode.has_fp8_qdq( + ) + self.is_nvfp4 = self.is_quanted and model_config.quant_config.quant_mode.has_nvfp4( + ) + self.enable_attention_dp = model_config.mapping.enable_attention_dp self.fusion_config = EagerFusionConfig() @@ -346,9 +360,10 @@ def __init__( aux_stream=aux_stream, attention_chunk_size=attention_chunk_size) - is_mlp_layer = (layer_idx + 1) % config.interleave_moe_layer_step != 0 + self.is_mlp_layer = (layer_idx + + 1) % config.interleave_moe_layer_step != 0 - if is_mlp_layer: + if self.is_mlp_layer: self.feed_forward = GatedMLP( hidden_size=config.hidden_size, intermediate_size=config.intermediate_size_mlp, @@ -387,13 +402,14 @@ def __init__( self.mapping = model_config.mapping self.all_reduce = AllReduce(self.mapping) self.next_layer_layernorm: RMSNorm = None + self.next_attn: LlamaAttention = None self.moe_allreduce = MoEAllReduce(self.mapping) def forward( self, position_ids: torch.LongTensor, - hidden_states: torch.Tensor, + hidden_states: Union[torch.Tensor, Fp4QuantizedTensor], attn_metadata: AttentionMetadata, residual: Optional[torch.Tensor], spec_metadata: Optional[SpecMetadata] = None, @@ -448,6 +464,7 @@ def forward( or self.mapping.tp_size == 1 or self.enable_attention_dp)), cutlass_min_latency_mode=cutlass_min_latency_mode, ) + if spec_metadata is not None: # We save the hidden states in the spec metadata here. In _prepare_draft_tokens, # PyExecutor will extract these from the model engine's spec metadata. @@ -613,10 +630,15 @@ def __init__( layer_idx=layer_idx, ) + if config.model_type == "llama4_text": + inter_size = config.intermediate_size_mlp + else: + inter_size = config.intermediate_size + self.mlp = GatedMLP( hidden_size=config.hidden_size, - intermediate_size=config.intermediate_size, - bias=config.mlp_bias, + intermediate_size=inter_size, + bias=getattr(config, "mlp_bias", False), dtype=config.torch_dtype, config=model_config, ) @@ -657,7 +679,6 @@ def forward( hidden_states, residual) hidden_states = self.mlp(hidden_states) - assert isinstance(spec_metadata, Eagle3SpecMetadata) # We save the hidden states in the spec metadata here. In _prepare_draft_tokens, # PyExecutor will extract these from the draft model engine's spec metadata. # They will be passed to the draft model engine on the next iteration. @@ -676,6 +697,7 @@ def __init__(self, model_config: ModelConfig[LlamaConfig]): self.padding_idx = config.pad_token_id self.num_hidden_layers = config.num_hidden_layers self.aux_stream = torch.cuda.Stream() + self.mapping = model_config.mapping if self.model_config.mapping.enable_attention_dp: self.embed_tokens = Embedding(config.vocab_size, @@ -691,8 +713,14 @@ def __init__(self, model_config: ModelConfig[LlamaConfig]): gather_output=True, ) + # If enable_min_latency is True, we will use min-latency mode. + DecoderLayerClass = Llama4DecoderLayer + if model_config.pytorch_backend_config.enable_min_latency: + from .modeling_llama_min_latency import Llama4MinLatencyDecoderLayer + DecoderLayerClass = Llama4MinLatencyDecoderLayer + self.layers = nn.ModuleList([ - Llama4DecoderLayer( + DecoderLayerClass( model_config, layer_idx, self.aux_stream, @@ -722,7 +750,8 @@ def forward( hidden_states = inputs_embeds residual = None - for decoder_layer in self.layers[:self.num_hidden_layers]: + for idx, decoder_layer in enumerate( + self.layers[:self.num_hidden_layers]): hidden_states, residual = decoder_layer( position_ids=position_ids, hidden_states=hidden_states, @@ -830,22 +859,73 @@ def __init__( hidden_size=model_config.pretrained_config.hidden_size, vocab_size=model_config.pretrained_config.vocab_size) + self.is_eagle3_one_model = hasattr( + model_config, "spec_config" + ) and model_config.spec_config is not None and model_config.spec_config.spec_dec_mode.is_eagle3_one_model( + ) + self.draft_model = None + if self.is_eagle3_one_model: + draft_config = ModelConfig.from_pretrained( + model_config.spec_config.draft_model_path, + trust_remote_code=True, + attn_backend=model_config.attn_backend, + moe_backend=model_config.moe_backend, + mapping=model_config.mapping) + draft_config.spec_config = model_config.spec_config + draft_config.max_num_tokens = model_config.max_num_tokens + draft_config.moe_max_num_tokens = model_config.moe_max_num_tokens + draft_config.quant_config.kv_cache_quant_algo = \ + model_config.quant_config.kv_cache_quant_algo + self.draft_model = Eagle3LlamaForCausalLM( + draft_config, model_config.pretrained_config.num_hidden_layers) + self.spec_worker = get_spec_worker(model_config.spec_config, + model_config.mapping) -@register_auto_model("Llama4ForCausalLM") -class Llama4ForCausalLM(DecoderModelForCausalLM[Llama4Model, Llama4Config]): - - def __init__( + def forward( self, - model_config: ModelConfig[Llama4Config], - ): - model_config = copy.copy(model_config) - architectures = model_config.pretrained_config.architectures - model_config.pretrained_config = model_config.pretrained_config.text_config - model_config.pretrained_config.architectures = architectures - super().__init__(Llama4Model(model_config), - config=model_config, - hidden_size=model_config.pretrained_config.hidden_size, - vocab_size=model_config.pretrained_config.vocab_size) + attn_metadata: AttentionMetadata, + input_ids: torch.LongTensor = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + return_context_logits: bool = False, + spec_metadata: Optional[SpecMetadata] = None, + lora_params: Optional[dict] = None, + **kwargs, + ) -> torch.Tensor: + hidden_states = self.model( + input_ids=input_ids, + attn_metadata=attn_metadata, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + lora_params=lora_params, + ) + + if self.draft_model is not None: + # get logits + logits = self.logits_processor.forward( + hidden_states[spec_metadata.gather_ids], + self.lm_head, + attn_metadata, + True, + ) + # get accepted tokens and next draft tokens + return self.spec_worker(input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + logits=logits, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=self.draft_model) + else: + logits = self.logits_processor.forward( + hidden_states, + self.lm_head, + attn_metadata, + return_context_logits, + ) + + return logits def infer_max_seq_len(self): if self.model_config.attn_backend.upper() != 'TRTLLM': @@ -857,23 +937,11 @@ def infer_max_seq_len(self): return super().infer_max_seq_len() def load_weights(self, weights: Dict): - new_weights = {} - for key, tensor in weights.items(): - if key.startswith("language_model."): - new_key = key[len("language_model."):] - new_weights[new_key] = tensor - else: - new_weights[key] = tensor + super().load_weights(weights, skip_modules=["draft_model"]) - super().load_weights(new_weights) - - for idx, layer in enumerate( - self.model.layers[:self.config.num_hidden_layers]): - if idx == self.config.num_hidden_layers - 1: - layer.next_layer_layernorm = self.model.norm - else: - layer.next_layer_layernorm = self.model.layers[ - idx + 1].input_layernorm + def load_draft_weights(self, weights: Dict): + self.draft_model.load_weights(weights) + self.draft_model.load_weights_from_target_model(self) class Llama4InputProcessor(InputProcessor): @@ -938,31 +1006,129 @@ def __call__( @register_auto_model("Llama4ForConditionalGeneration") @register_input_processor(Llama4InputProcessor) -class Llama4ForConditionalGeneration(Llama4ForCausalLM): +class Llama4ForConditionalGeneration(DecoderModelForCausalLM[Llama4Model, + Llama4Config]): + + def __init__( + self, + model_config: ModelConfig[Llama4Config], + ): + # TODO: figure out a better way to handle multimodality. + model_config = copy.copy(model_config) + architectures = model_config.pretrained_config.architectures + model_config.pretrained_config = model_config.pretrained_config.text_config + model_config.pretrained_config.architectures = architectures + super().__init__(Llama4Model(model_config), + config=model_config, + hidden_size=model_config.pretrained_config.hidden_size, + vocab_size=model_config.pretrained_config.vocab_size) + + self.is_eagle3_one_model = hasattr( + model_config, "spec_config" + ) and model_config.spec_config is not None and model_config.spec_config.spec_dec_mode.is_eagle3_one_model( + ) + self.draft_model = None + if self.is_eagle3_one_model: + draft_config = ModelConfig.from_pretrained( + model_config.spec_config.draft_model_path, + trust_remote_code=True, + attn_backend=model_config.attn_backend, + moe_backend=model_config.moe_backend, + mapping=model_config.mapping) + draft_config.spec_config = model_config.spec_config + draft_config.max_num_tokens = model_config.max_num_tokens + draft_config.moe_max_num_tokens = model_config.moe_max_num_tokens + draft_config.quant_config.kv_cache_quant_algo = \ + model_config.quant_config.kv_cache_quant_algo + self.draft_model = Eagle3LlamaForCausalLM( + draft_config, model_config.pretrained_config.num_hidden_layers) + self.spec_worker = get_spec_worker(model_config.spec_config, + model_config.mapping) - @torch.inference_mode() def forward( self, attn_metadata: AttentionMetadata, - input_ids: Optional[torch.LongTensor] = None, + input_ids: torch.LongTensor = None, position_ids: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, - return_context_logits: Optional[bool] = False, + return_context_logits: bool = False, spec_metadata: Optional[SpecMetadata] = None, **kwargs, ) -> torch.Tensor: - mm_embed = kwargs.get("multi_modal_data", []) - input_ids, inputs_embeds = fuse_input_embeds(self.model.embed_tokens, - input_ids, mm_embed) - logits = super().forward( - attn_metadata, - input_ids, - position_ids, - inputs_embeds, - spec_metadata=spec_metadata, - return_context_logits=return_context_logits, - ) - return logits + if self.is_eagle3_one_model: + hidden_states = self.model( + input_ids=input_ids, + attn_metadata=attn_metadata, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + ) + + if self.draft_model is not None: + # get logits + logits = self.logits_processor.forward( + hidden_states[spec_metadata.gather_ids], + self.lm_head, + attn_metadata, + True, + ) + # get accepted tokens and next draft tokens + return self.spec_worker(input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + logits=logits, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=self.draft_model) + else: + logits = self.logits_processor.forward( + hidden_states, + self.lm_head, + attn_metadata, + return_context_logits, + ) + + return logits + else: + mm_embed = kwargs.get("multi_modal_data", []) + input_ids, inputs_embeds = fuse_input_embeds( + self.model.embed_tokens, input_ids, mm_embed) + logits = super().forward( + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + spec_metadata=spec_metadata, + return_context_logits=return_context_logits) + return logits + + def infer_max_seq_len(self): + # TODO: implement chunked attention to support 10M context length + return 8192 + + def load_weights(self, weights: Dict): + new_weights = {} + for key, tensor in weights.items(): + if key.startswith("language_model."): + new_key = key[len("language_model."):] + new_weights[new_key] = tensor + else: + new_weights[key] = tensor + + super().load_weights(new_weights, skip_modules=["draft_model"]) + + for idx, layer in enumerate( + self.model.layers[:self.config.num_hidden_layers]): + if idx == self.config.num_hidden_layers - 1: + layer.next_layer_layernorm = self.model.norm + else: + layer.next_layer_layernorm = self.model.layers[ + idx + 1].input_layernorm + layer.next_attn = self.model.layers[idx + 1].self_attn + + def load_draft_weights(self, weights: Dict): + self.draft_model.load_weights(weights) + self.draft_model.load_weights_from_target_model(self) @register_auto_model("MistralForCausalLM") @@ -988,12 +1154,16 @@ def __init__( class Eagle3LlamaDraftModel(DecoderModel): - def __init__(self, model_config: ModelConfig[LlamaConfig]) -> None: + def __init__(self, + model_config: ModelConfig[LlamaConfig], + start_layer_idx: int = 0) -> None: super().__init__(model_config) config = model_config.pretrained_config + self.spec_config = model_config.spec_config self.dtype = config.torch_dtype self.hidden_size = config.hidden_size + self.mapping = model_config.mapping if hasattr(config, "target_hidden_size"): self.hidden_size_in = config.target_hidden_size @@ -1002,18 +1172,19 @@ def __init__(self, model_config: ModelConfig[LlamaConfig]) -> None: self.fc = Linear(self.hidden_size_in * 3, config.hidden_size, - bias=False, + bias=getattr(config, "bias", False), dtype=config.torch_dtype) - self.midlayer = Eagle3LlamaDecoderLayer(model_config, 0) + self.midlayer = Eagle3LlamaDecoderLayer(model_config, start_layer_idx) self.norm = RMSNorm(hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype) - self.d2t = nn.Parameter(torch.empty((config.draft_vocab_size, ), - dtype=torch.int64), - requires_grad=False) + if config.vocab_size != config.draft_vocab_size: + self.d2t = nn.Parameter(torch.empty((config.draft_vocab_size, ), + dtype=torch.int64), + requires_grad=False) if self.hidden_size_in != config.hidden_size: self.embed_tokens = Embedding( @@ -1059,8 +1230,11 @@ def forward( attn_metadata=attn_metadata, spec_metadata=spec_metadata) - hidden_states, _ = self.norm(hidden_states, residual) - return hidden_states + hidden_states, hidden_states_to_save = self.norm( + hidden_states, residual) + if self.spec_config.spec_dec_mode.is_eagle3(): + spec_metadata.maybe_capture_hidden_states(1, hidden_states_to_save) + return hidden_states, hidden_states_to_save @register_auto_model("EAGLE3LlamaForCausalLM") @@ -1070,9 +1244,10 @@ class Eagle3LlamaForCausalLM(DecoderModelForCausalLM[Eagle3LlamaDraftModel, def __init__( self, model_config: ModelConfig[LlamaConfig], + start_layer_idx: int = 0, ): super().__init__( - Eagle3LlamaDraftModel(model_config), + Eagle3LlamaDraftModel(model_config, start_layer_idx), config=model_config, hidden_size=model_config.pretrained_config.hidden_size, vocab_size=model_config.pretrained_config.draft_vocab_size) @@ -1088,7 +1263,7 @@ def forward( hidden_states: Optional[torch.Tensor] = None, **kwargs, ) -> torch.Tensor: - output = self.model( + output, _ = self.model( input_ids=input_ids, attn_metadata=attn_metadata, position_ids=position_ids, diff --git a/tensorrt_llm/_torch/models/modeling_llama_min_latency.py b/tensorrt_llm/_torch/models/modeling_llama_min_latency.py new file mode 100644 index 000000000000..32f6ae4a5145 --- /dev/null +++ b/tensorrt_llm/_torch/models/modeling_llama_min_latency.py @@ -0,0 +1,860 @@ +from collections.abc import Callable +from functools import partial +from typing import Dict, List, Optional, Tuple, Union + +import torch +from torch.nn import functional as F +from transformers import LlamaConfig + +from tensorrt_llm._torch.distributed import AllReduceFusionOp, AllReduceParams +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.quantization.utils.fp4_utils import ( + reorder_rows_for_gated_act_gemm, shuffle_matrix_a) + +from ...models.modeling_utils import QuantConfig +from ..attention_backend import AttentionMetadata +from ..attention_backend.interface import PredefinedAttentionMask +from ..model_config import ModelConfig +from ..modules.fused_moe import (BaseMoeRoutingMethod, FusedMoE, + FusedMoEQuantScalesFP8, + Llama4RenormalizeMoeRoutingMethod, + MoEWeightLoadingMode) +from ..modules.gated_mlp import GatedMLP, swiglu +from ..modules.linear import (Linear, TensorParallelMode, WeightMode, + WeightsLoadingConfig) +from ..modules.multi_stream_utils import maybe_execute_in_parallel +from ..speculative import SpecMetadata +from ..utils import Fp4QuantizedTensor +from .modeling_llama import Llama4Attention, Llama4DecoderLayer, Llama4MoE + +# Perf heuristics thresholds. +# Use routing gemv kernels when num_tokens <= 8. +MIN_LATENCY_ROUTING_GEMM_NUM_TOKENS = 8 +# Use QKV gemv kernel with fused attn scaling when num_tokens <= 8. +MIN_LATENCY_QKV_GEMM_ATTN_SCALING_NUM_TOKENS = 8 +# Use QKV gemv kernel when num_tokens <= 4. +MIN_LATENCY_QKV_GEMM_NUM_TOKENS = 4 +# Use FC2 trtllm-gen kernel when num_tokens <= 16. +MIN_LATENCY_FC2_NUM_TOKENS = 16 +# Use FC13 gemv kernel with fused swiglu when num_tokens < 4. +MIN_LATENCY_FC13_FUSED_GEMM_SWIGLU_NUM_TOKENS_GEMV = 4 +# Use FC13 trtllm-gen kernel with fused swiglu when 4 <= num_tokens <= 16. +MIN_LATENCY_FC13_FUSED_GEMM_SWIGLU_NUM_TOKENS_TRTLLM_GEN = 16 +# Use min-latency MoE kernels when num_tokens <= 8. +MIN_LATENCY_FUSED_MOE_NUM_TOKENS = 8 + + +class Llama4MinLatencyLinear(Linear): + """ + A wrapper around Linear because we may optionally use min-latency kernels depending on input shapes. + """ + + def __init__( + self, + in_features: int, + out_features: int, + bias: bool = True, + dtype: torch.dtype = None, + mapping: Optional[Mapping] = None, + tensor_parallel_mode: Optional[TensorParallelMode] = None, + gather_output: bool = False, + quant_config: Optional[QuantConfig] = None, + weights_loading_config: Optional[WeightsLoadingConfig] = None, + reduce_output: bool = True, + skip_create_weights_in_init: bool = False, + use_custom_cublas_mm: bool = False, + enable_fused_gemm_swiglu: bool = False, + enable_fused_gemm_attn_scaling: bool = False, + enable_trtllm_gen: bool = False, + post_load_weights_hook: Optional[Callable] = None, + ): + # First, initialize the base class. + super().__init__( + in_features, + out_features, + bias, + dtype, + mapping, + tensor_parallel_mode, + gather_output, + quant_config, + weights_loading_config, + reduce_output, + skip_create_weights_in_init, + use_custom_cublas_mm, + ) + + # Set min-latency specific attributes. + self.enable_fused_gemm_swiglu = enable_fused_gemm_swiglu + self.enable_fused_gemm_attn_scaling = enable_fused_gemm_attn_scaling + self.enable_trtllm_gen = enable_trtllm_gen + self.post_load_weights_hook = post_load_weights_hook + self.position_ids = None + + def load_weights(self, weights: List[Dict]): + + super().load_weights(weights) + + # After loading weights, calculate the combined scale (input_scale * weight_scale) for special kernels and + # trtllm-gen kernels. + if self.has_fp8_qdq: + self.combined_scale = self.input_scale * self.weight_scale + + # If this is gate_up_proj + swiglu and trtllm-gen kernels will be used, we need to reorder the weights + # for trtllm-gen gemm+swiglu kernels. + if self.enable_trtllm_gen and self.enable_fused_gemm_swiglu: + reordered_weight = self.weight.reshape(2, + self.out_features // 2, + self.in_features) + reordered_weight = torch.cat( + [reordered_weight[1], reordered_weight[0]], dim=0) + reordered_weight = reorder_rows_for_gated_act_gemm( + reordered_weight) + self.trtllm_gen_weight = shuffle_matrix_a( + reordered_weight.view(torch.uint8), + 128).view(torch.float8_e4m3fn) + + # Otherwise, if trtllm-gen kernels will be used, shuffle the weights. + elif self.enable_trtllm_gen: + self.trtllm_gen_weight = shuffle_matrix_a( + self.weight.view(torch.uint8), + 128).view(torch.float8_e4m3fn) + + if self.post_load_weights_hook is not None: + self.post_load_weights_hook(self) + + # Override apply_linear instead of forward so that we can reuse the AllReduce/AllGather logic in the parent class. + def apply_linear( + self, + input, + weight, + bias, + lora_params: Optional[dict] | None = None, + layer_idx: Optional[int] | None = None, + ) -> torch.Tensor: + + # Quantize the input if it is not already in float8_e4m3fn. + # We cannot do this when enable_fused_gemm_swiglu is True and num_tokens > 16 because the default path + # does not support FP8 input. + if self.has_fp8_qdq \ + and input.dtype != torch.float8_e4m3fn: + input, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor( + input, self.input_scale) + + # Use special BF16-input BF16-output gemm kernel for routing gemm when num_tokens <= 8. + if input.shape[0] <= MIN_LATENCY_ROUTING_GEMM_NUM_TOKENS \ + and self.in_features == 5120 \ + and self.out_features == 128 \ + and input.dtype == torch.bfloat16 \ + and self.weight.dtype == torch.bfloat16 \ + and not self.has_any_quant: + return torch.ops.trtllm.llama4_bf16_bf16_gemm(input, self.weight) + + # Use special FP8-input BF16-output gemm kernel for QKV-gemm. + if input.shape[0] <= MIN_LATENCY_QKV_GEMM_ATTN_SCALING_NUM_TOKENS \ + and self.in_features == 5120 \ + and self.out_features == 896 \ + and input.dtype == torch.float8_e4m3fn: + + # Check if we should fuse attn_scaling into QKV-gemm. + should_fuse_attn_scaling = self.enable_fused_gemm_attn_scaling \ + and self.position_ids is not None + + # When num_tokens < 4, we use the special QKV-gemm kernel. + # Or if attn_scaling is needed, we also use the special QKV-gemm kernel. + if input.shape[ + 0] <= MIN_LATENCY_QKV_GEMM_NUM_TOKENS or should_fuse_attn_scaling: + return torch.ops.trtllm.llama4_fp8_bf16_gemm( + input, + self.weight, + self.combined_scale, + self.position_ids, + ) + # Otherwise, we use the trtllm-gen QKV-gemm kernel. + else: + if not hasattr(self, "trtllm_gen_weight"): + raise ValueError('Expect trtllm_gen_weight to be set') + return torch.ops.trtllm.fp8_per_tensor_scaling_tllmg_gemm( + input, + self.trtllm_gen_weight, + global_scale=self.combined_scale, + out_dtype=torch.bfloat16, + low_latency_kernel=True, + gated_silu=False, + ) + + # Use trtllm-gen FP8-input BF16-output gemm kernel for FC2. + if input.shape[0] <= MIN_LATENCY_FC2_NUM_TOKENS \ + and self.in_features % 512 == 0 \ + and self.out_features == 5120 \ + and input.dtype == torch.float8_e4m3fn: + if not hasattr(self, "trtllm_gen_weight"): + raise ValueError('Expect trtllm_gen_weight to be set') + return torch.ops.trtllm.fp8_per_tensor_scaling_tllmg_gemm( + input, + self.trtllm_gen_weight, + global_scale=self.combined_scale, + out_dtype=torch.bfloat16, + low_latency_kernel=True, + gated_silu=False, + ) + + # Use special FP8-input FP8-output gemm+swiglu kernel for FC13+swiglu. + if self.enable_fused_gemm_swiglu and self.has_fp8_qdq: + # When num_tokens < 4, we use the special gemm+swiglu kernel. + if input.shape[ + 0] < MIN_LATENCY_FC13_FUSED_GEMM_SWIGLU_NUM_TOKENS_GEMV: + if not hasattr(self, "inv_output_scale"): + raise ValueError('Expect inv_output_scale to be set') + return torch.ops.trtllm.llama4_fp8_fp8_gemm_swiglu( + input, + self.weight, + self.combined_scale, + self.inv_output_scale, + ) + # When 4 <= num_tokens <= 16, we use the trtllm-gen gemm+swiglu kernel. + elif MIN_LATENCY_FC13_FUSED_GEMM_SWIGLU_NUM_TOKENS_GEMV <= input.shape[0] \ + and input.shape[0] <= MIN_LATENCY_FC13_FUSED_GEMM_SWIGLU_NUM_TOKENS_TRTLLM_GEN: + if not hasattr(self, "trtllm_gen_global_scale"): + raise ValueError('Expect trtllm_gen_global_scale to be set') + if not hasattr(self, "trtllm_gen_weight"): + raise ValueError('Expect trtllm_gen_weight to be set') + return torch.ops.trtllm.fp8_per_tensor_scaling_tllmg_gemm( + input, + self.trtllm_gen_weight, + global_scale=self.trtllm_gen_global_scale, + global_scale_gate=self.combined_scale, + out_dtype=torch.float8_e4m3fn, + low_latency_kernel=True, + gated_silu=True, + ) + + # If special gemm+swiglu kernel is not used and enable_fused_gemm_swiglu is True, we need to apply swiglu + # manually. + if self.enable_fused_gemm_swiglu: + intermediate = super().apply_linear(input, weight, bias, + lora_params, layer_idx) + return swiglu(intermediate) + + # Otherwise, call the default apply_linear method. + return super().apply_linear(input, weight, bias, lora_params, layer_idx) + + # Set the position_ids for the next call to apply_linear. + def set_position_ids(self, position_ids: Optional[torch.LongTensor] = None): + self.position_ids = position_ids + + +class Llama4MinLatencyGatedMLP(GatedMLP): + """ + A wrapper around GatedMLP so that we can pre-compute some scales needed by the special kernels. + """ + + def __init__(self, + *, + hidden_size: int, + intermediate_size: int, + bias: bool, + activation: Callable[[torch.Tensor], torch.Tensor] = F.silu, + dtype: Optional[torch.dtype] = None, + config: Optional[ModelConfig] = None, + overridden_tp_size: Optional[int] = None, + reduce_output: bool = True, + layer_idx: Optional[int] = None): + + # First, initialize the base class. + super().__init__(hidden_size=hidden_size, + intermediate_size=intermediate_size, + bias=bias, + activation=activation, + dtype=dtype, + config=config, + overridden_tp_size=overridden_tp_size, + reduce_output=reduce_output, + layer_idx=layer_idx) + + # Override gate_up_proj and down_proj with Llama4Linear if we want to use special kernels. + self.enable_fused_gemm_swiglu = False + if self.hidden_size == 5120 \ + and (self.intermediate_size == 16384 or self.intermediate_size == 8192) \ + and self.mapping.tp_size == 8 \ + and config.quant_config.quant_mode.has_fp8_qdq(): + + self.enable_fused_gemm_swiglu = True + self.gate_up_proj = Llama4MinLatencyLinear( + self.hidden_size, + self.intermediate_size * 2, + bias=bias, + dtype=dtype, + mapping=self.mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + weights_loading_config=WeightsLoadingConfig( + weight_mode=WeightMode.FUSED_GATE_UP_LINEAR), + quant_config=config.get_quant_config(), + reduce_output=reduce_output, + skip_create_weights_in_init=config.skip_create_weights_in_init, + enable_fused_gemm_swiglu=True, + enable_trtllm_gen=True, + ) + + # After loading both gate_up_proj and down_proj, we need to set the scales needed by the special kernels and by + # the trtllm-gen gemm+swiglu kernel. + def post_load_weights_hook(gate_up_proj, down_proj): + if gate_up_proj.has_fp8_qdq: + # For the special gemm+swiglu kernel, we need to set the inverse of the output scale, which is the inverse + # of down_proj's combined input scale. + gate_up_proj.inv_output_scale = 1.0 / down_proj.input_scale + # For the trtllm-gen gemm+swiglu kernel, we need to set the global scale, which is gate_up_proj's + # combined input scale times inv_output_scale. + gate_up_proj.trtllm_gen_global_scale = gate_up_proj.combined_scale * gate_up_proj.inv_output_scale + + self.down_proj = Llama4MinLatencyLinear( + self.intermediate_size, + self.hidden_size, + bias=bias, + dtype=dtype, + mapping=self.mapping, + tensor_parallel_mode=TensorParallelMode.ROW, + quant_config=config.get_quant_config(), + reduce_output=reduce_output, + skip_create_weights_in_init=config.skip_create_weights_in_init, + enable_trtllm_gen=True, + post_load_weights_hook=partial(post_load_weights_hook, + self.gate_up_proj), + ) + + def forward( + self, + x: Union[torch.Tensor, Fp4QuantizedTensor], + all_rank_num_tokens=None, + final_all_reduce_params: Optional[AllReduceParams] = None, + lora_params: Optional[dict] = None, + ) -> torch.Tensor: + # When gemm+swiglu is fused, we need to temporarily disable the activation function to avoid reapplying swiglu. + if self.enable_fused_gemm_swiglu: + orig_activation = self.activation + self.activation = None + + # Call the parent's forward method. + output = super().forward( + x, + all_rank_num_tokens, + final_all_reduce_params, + lora_params, + ) + + # Restore the original activation function. + if self.enable_fused_gemm_swiglu: + self.activation = orig_activation + + return output + + +class Llama4MinLatencyAttention(Llama4Attention): + + def __init__( + self, + model_config: ModelConfig[LlamaConfig], + layer_idx: Optional[int] = None, + use_qk_norm: bool = False, + nope_layer: bool = False, + attn_temperature_tuning: bool = True, + aux_stream: Optional[torch.cuda.Stream] = None, + attention_chunk_size: Optional[int] = None, + ): + + # First, initialize the base class. + super().__init__( + model_config=model_config, + layer_idx=layer_idx, + use_qk_norm=use_qk_norm, + nope_layer=nope_layer, + attn_temperature_tuning=attn_temperature_tuning, + aux_stream=aux_stream, + attention_chunk_size=attention_chunk_size, + ) + + # Then, enable min-latency QKV gemm when the sizes are as expected. + config = model_config.pretrained_config + tp_size = model_config.mapping.tp_size + self.enable_min_latency_qkv = False + self.enable_fused_gemm_attn_scaling = False + + if config.hidden_size == 5120 \ + and config.num_attention_heads == 40 \ + and config.num_key_value_heads == 8 \ + and tp_size == 8: + self.enable_min_latency_qkv = True + + # Decide whether to fuse attn_scaling into QKV gemm. + self.enable_fused_gemm_attn_scaling = self.attn_temperature_tuning \ + and self.floor_scale == 8192.0 \ + and self.attn_scale == 0.1 + + # When min-latency QKV gemm is enabled, override qkv_proj. + self.qkv_proj = Llama4MinLatencyLinear( + self.hidden_size, + tp_size * self.q_size + 2 * tp_size * self.kv_size, + bias=config.attention_bias, + dtype=config.torch_dtype, + mapping=model_config.mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + weights_loading_config=WeightsLoadingConfig( + weight_mode=WeightMode.FUSED_QKV_LINEAR), + quant_config=model_config.get_quant_config(), + skip_create_weights_in_init=model_config. + skip_create_weights_in_init, + enable_fused_gemm_attn_scaling=self. + enable_fused_gemm_attn_scaling, + enable_trtllm_gen=True, + ) + + def _forward_nope( + self, + position_ids: Optional[torch.LongTensor], + hidden_states: Union[torch.Tensor, Fp4QuantizedTensor], + attn_metadata: AttentionMetadata, + attention_mask: PredefinedAttentionMask = PredefinedAttentionMask. + CAUSAL, + mrope_config: Optional[dict] = None, + all_reduce_params: Optional[AllReduceParams] = None, + ): + # If we are going to use min-latency gemm+attn_scaling kernel, pass position_ids to QKV gemm and set + # skip_attn_scaling to True. + skip_attn_scaling = False + if self.enable_min_latency_qkv \ + and self.enable_fused_gemm_attn_scaling \ + and position_ids is not None \ + and hidden_states.shape[0] <= MIN_LATENCY_QKV_GEMM_ATTN_SCALING_NUM_TOKENS: + self.qkv_proj.set_position_ids(position_ids) + skip_attn_scaling = True + + return super()._forward_nope(position_ids, hidden_states, attn_metadata, + attention_mask, mrope_config, + all_reduce_params, skip_attn_scaling) + + +class Llama4MinLatencyFusedMoE(FusedMoE): + + def __init__( + self, + *, + routing_method: BaseMoeRoutingMethod, + num_experts: int, + hidden_size: int, + intermediate_size: int, + dtype: Optional[torch.dtype] = None, + reduce_results: bool = False, + model_config: ModelConfig = ModelConfig(), + aux_stream: torch.cuda.Stream = torch.cuda.Stream(), + weight_loading_mode: MoEWeightLoadingMode = MoEWeightLoadingMode. + VANILLA, + apply_router_weight_on_input: bool = False, + post_load_weights_hook: Optional[Callable] = None, + ): + + super().__init__( + routing_method=routing_method, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + dtype=dtype, + reduce_results=reduce_results, + model_config=model_config, + aux_stream=aux_stream, + weight_loading_mode=weight_loading_mode, + apply_router_weight_on_input=apply_router_weight_on_input, + ) + + self.post_load_weights_hook = post_load_weights_hook + + # Enable min-latency mode for Llama4 Maverick TP8 EP1. + self.enable_min_latency_fused_moe = False + if num_experts == 128 \ + and hidden_size == 5120 \ + and intermediate_size == 8192 \ + and model_config.quant_config.quant_mode.has_fp8_qdq() \ + and model_config.mapping.moe_tp_size == 8 \ + and model_config.mapping.moe_ep_size == 1 \ + and routing_method.top_k == 1 \ + and apply_router_weight_on_input: + self.enable_min_latency_fused_moe = True + + def load_weights(self, weights: List[Dict]): + super().load_weights(weights) + + if self.post_load_weights_hook: + self.post_load_weights_hook(self) + + def forward( + self, + x: Union[torch.Tensor, Fp4QuantizedTensor], + router_logits: torch.Tensor, + output_dtype: Optional[torch.dtype] = None, + x_high: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + + # Use special min-latency MoE kernels when num_tokens <= 8. + if self.enable_min_latency_fused_moe \ + and x.dtype == torch.float8_e4m3fn \ + and x.shape[0] <= MIN_LATENCY_FUSED_MOE_NUM_TOKENS: + assert hasattr(self, "min_latency_quant_scales" + ), "Expect min_latency_quant_scales to be set" + + return torch.ops.trtllm.llama4_moe_tp8ep1_min_latency( + x, router_logits, self.w3_w1_weight, self.w2_weight, + self.min_latency_quant_scales) + + # Default MoE implementation does not support FP8 input, so use high-precision one instead. + if x.dtype == torch.float8_e4m3fn: + if x_high is not None: + x = x_high + else: + raise ValueError( + "x_high is required when x.dtype is float8_e4m3fn in Llama4FusedMoE fallback path!" + ) + + return super().forward(x, + router_logits, + cutlass_min_latency_mode=False, + output_dtype=output_dtype) + + +class Llama4MinLatencyMoE(Llama4MoE): + + def __init__( + self, + *, + num_experts: int, + top_k: int, + hidden_size: int, + intermediate_size: int, + shared_expert_intermediate_size: int, + aux_stream: torch.cuda.Stream, + dtype: Optional[torch.dtype] = None, + tune_max_num_tokens: int = 8192, + model_config: ModelConfig = ModelConfig(), + ): + + # First, initialize the base class. + super().__init__( + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + shared_expert_intermediate_size=shared_expert_intermediate_size, + aux_stream=aux_stream, + dtype=dtype, + tune_max_num_tokens=tune_max_num_tokens, + model_config=model_config, + ) + + # Then, override modules with min-latency versions. + self.shared_expert = Llama4MinLatencyGatedMLP( + hidden_size=hidden_size, + intermediate_size=shared_expert_intermediate_size, + bias=False, + dtype=dtype, + config=model_config, + overridden_tp_size=1 if self.enable_attention_dp else None, + reduce_output=False) + + def post_load_weights_hook(shared_expert, experts): + # Set min-latency quant scales for routed experts if we plan to use min-latency MoE kernels. + # This is because the routed experts' input scale is after the score multiplication, so we must use the + # pre-score scaling input scale, which happens to be shared expert's input scale. + if experts.enable_min_latency_fused_moe and hasattr( + shared_expert.gate_up_proj, "input_scale"): + pre_score_scaling_input_scale = shared_expert.gate_up_proj.input_scale + experts.min_latency_quant_scales = FusedMoEQuantScalesFP8( + fc1_dequant=experts.fc31_dequant.data / + experts.fc31_input_dequant.data * + pre_score_scaling_input_scale, + fc2_quant=experts.fc2_quant, + fc2_dequant=experts.fc2_dequant, + fc1_input_dequant=pre_score_scaling_input_scale, + ) + + self.experts = Llama4MinLatencyFusedMoE( + routing_method=Llama4RenormalizeMoeRoutingMethod(top_k), + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + dtype=dtype, + reduce_results= + False, # In both low latency and max-throughput scenarios, FusedMoE needs not to do allreduce inside op. + weight_loading_mode=MoEWeightLoadingMode.FUSED_GATE_UP_PROJ, + model_config=model_config, + apply_router_weight_on_input=True, + post_load_weights_hook=partial(post_load_weights_hook, + self.shared_expert)) + + self.router = Llama4MinLatencyLinear( + hidden_size, + num_experts, + bias=False, + dtype=model_config.pretrained_config.torch_dtype, + quant_config=None) + + def compute_routed_output( + self, + hidden_states, + all_rank_num_tokens, + hidden_states_high: Optional[torch.Tensor] = None): + # Use high precision hidden states for routing gemm if it is provided. + hidden_states_routing = hidden_states_high if hidden_states_high is not None else hidden_states + router_logits = self.router.forward(hidden_states_routing) + routed_output = self.experts.forward( + hidden_states, + router_logits, + x_high=hidden_states_high, + ) + + return routed_output + + def forward( + self, + hidden_states: torch.Tensor, + all_rank_num_tokens=None, + final_all_reduce_params: Optional[AllReduceParams] = None, + # Optional input for routing gemm if experts and routing gemm require + # different precisions for input hidden states. + hidden_states_high: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + + fn0 = lambda: self.shared_expert(hidden_states) + fn1 = lambda: self.compute_routed_output( + hidden_states, all_rank_num_tokens, hidden_states_high) + shared_output, routed_output = maybe_execute_in_parallel( + fn0, fn1, self.moe_event[0], self.moe_event[1], self.aux_stream) + + assert shared_output.size() == routed_output.size( + ), f'unmatched tensor shape' + final_hidden_states = shared_output + routed_output + if not self.enable_attention_dp and self.mapping.tp_size > 1: + final_hidden_states = self.all_reduce( + final_hidden_states, all_reduce_params=final_all_reduce_params) + + return final_hidden_states + + +class Llama4MinLatencyDecoderLayer(Llama4DecoderLayer): + + def __init__( + self, + model_config: ModelConfig[LlamaConfig], + layer_idx: int, + aux_stream: Optional[torch.cuda.Stream] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + + # First, initialize the base class. + super().__init__(model_config, layer_idx, aux_stream) + + # Then, override modules with min-latency versions. + config = model_config.pretrained_config + nope_layer = config.no_rope_layers[layer_idx] == 0 + attention_chunk_size = getattr(config, "attention_chunk_size", + None) if not nope_layer else None + self.self_attn = Llama4MinLatencyAttention( + model_config, + layer_idx=layer_idx, + use_qk_norm=getattr(config, "use_qk_norm", False), + nope_layer=nope_layer, + attn_temperature_tuning=config.attn_temperature_tuning > 0, + aux_stream=aux_stream, + attention_chunk_size=attention_chunk_size) + + self.fusion_config.PRE_MLP_FUSION = False + self.fusion_config.POST_MLP_FUSION = False + self.fusion_config.PRE_MOE_FUSION = False + self.fusion_config.POST_MOE_FUSION = False + + if self.is_mlp_layer: + self.feed_forward = Llama4MinLatencyGatedMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size_mlp, + # Llama4 has no mlp_bias field. + bias=getattr(config, "mlp_bias", False), + dtype=config.torch_dtype, + config=model_config, + overridden_tp_size=1 if self.enable_attention_dp else None, + layer_idx=layer_idx) + + self.fusion_config.PRE_MLP_FUSION = model_config.mapping.has_tp() + self.fusion_config.POST_MLP_FUSION = model_config.mapping.has_tp() + else: + self.feed_forward = Llama4MinLatencyMoE( + num_experts=config.num_local_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + shared_expert_intermediate_size=config.intermediate_size, + model_config=model_config, + aux_stream=aux_stream, + dtype=config.torch_dtype) + + self.fusion_config.PRE_MOE_FUSION = model_config.mapping.has_tp() + self.fusion_config.POST_MOE_FUSION = model_config.mapping.has_tp() + + def forward( + self, + position_ids: torch.LongTensor, + hidden_states: Union[torch.Tensor, Fp4QuantizedTensor], + attn_metadata: AttentionMetadata, + residual: Optional[torch.Tensor], + spec_metadata: Optional[SpecMetadata] = None, + **kwargs, + ) -> torch.Tensor: + + num_tokens = hidden_states.shape[0] + + use_fp8_allreduce = num_tokens <= 128 and self.is_fp8_quant + use_fp4_allreduce = num_tokens <= 128 and self.is_nvfp4 + + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states = self.self_attn( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + all_reduce_params=AllReduceParams(enable_allreduce=not ( + self.fusion_config.PRE_MOE_FUSION + or self.fusion_config.PRE_MLP_FUSION + or self.mapping.tp_size == 1 or self.enable_attention_dp)), + **kwargs, + ) + + # Reserved if we need both high-precision and low-precision hidden states. + hidden_states_high = None + if self.fusion_config.PRE_MLP_FUSION and use_fp8_allreduce: + hidden_states, residual = self.all_reduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_FP8, + residual=residual, + norm_weight=self.post_attention_layernorm.weight, + scale=self.feed_forward.gate_up_proj.input_scale, + eps=self.post_attention_layernorm.variance_epsilon, + )) + elif self.fusion_config.PRE_MLP_FUSION and use_fp4_allreduce: + act_fp4, act_sf, residual = self.all_reduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4, + residual=residual, + norm_weight=self.post_attention_layernorm.weight, + scale=self.feed_forward.gate_up_proj.input_scale, + eps=self.post_attention_layernorm.variance_epsilon, + )) + hidden_states = Fp4QuantizedTensor(act_fp4, act_sf) + elif self.fusion_config.PRE_MOE_FUSION and use_fp8_allreduce: + # For pre-MoE all reduce in FP8 mode, we must output two variants of + # the tensor: one in high-precision (BF16) and one in FP8. + # This is because routing gemm requires BF16 input while shared + # expert and routed expert require FP8 input. + hidden_states_high, hidden_states, residual = self.all_reduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM_OUT_QUANT_FP8, + residual=residual, + norm_weight=self.post_attention_layernorm.weight, + scale=self.feed_forward.shared_expert.gate_up_proj. + input_scale, + eps=self.post_attention_layernorm.variance_epsilon, + )) + elif self.fusion_config.PRE_MOE_FUSION or self.fusion_config.PRE_MLP_FUSION: + hidden_states, residual = self.all_reduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=self.post_attention_layernorm.weight, + eps=self.post_attention_layernorm.variance_epsilon, + )) + else: + # Fully Connected + hidden_states, residual = self.post_attention_layernorm( + hidden_states, residual) + + # In eagle3 mode, we capture the value in the boundary of decoder layer. + # If fusing rms in the next layer, the value is not correct. Thus, if + # this layer will be captured, we should not fuse the rms in the next + # layer. + if spec_metadata is not None: + if spec_metadata.is_layer_capture(self.layer_idx): + self.fusion_config.POST_MOE_FUSION = False + self.fusion_config.POST_MLP_FUSION = False + + if self.is_mlp_layer: + hidden_states = self.feed_forward( + hidden_states, + all_rank_num_tokens=attn_metadata.all_rank_num_tokens, + final_all_reduce_params=AllReduceParams(enable_allreduce=not ( + self.fusion_config.POST_MLP_FUSION + or self.mapping.tp_size == 1 or self.enable_attention_dp)), + ) + else: + hidden_states = self.feed_forward( + hidden_states, + all_rank_num_tokens=attn_metadata.all_rank_num_tokens, + final_all_reduce_params=AllReduceParams(enable_allreduce=not ( + self.fusion_config.POST_MOE_FUSION + or self.mapping.tp_size == 1 or self.enable_attention_dp)), + hidden_states_high=hidden_states_high, + ) + + if spec_metadata is not None: + # We save the hidden states in the spec metadata here. In _prepare_draft_tokens, + # PyExecutor will extract these from the model engine's spec metadata. + # They will be passed to the draft model engine on the first draft iteration. + # TODO: can we support multiple model outputs instead? + spec_metadata.maybe_capture_hidden_states(self.layer_idx, + hidden_states, residual) + + needs_post_allreduce = self.fusion_config.POST_MOE_FUSION \ + or self.fusion_config.POST_MLP_FUSION + if needs_post_allreduce and self.next_layer_layernorm is not None: + if use_fp8_allreduce and self.next_attn is not None: + hidden_states, residual = self.all_reduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_FP8, + residual=residual, + norm_weight=self.next_layer_layernorm.weight, + scale=self.next_attn.qkv_proj.input_scale, + eps=self.next_layer_layernorm.variance_epsilon, + )) + elif use_fp4_allreduce and self.next_attn is not None: + act_fp4, act_sf, residual = self.all_reduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp. + RESIDUAL_RMS_NORM_QUANT_NVFP4, + residual=residual, + norm_weight=self.next_layer_layernorm.weight, + scale=self.next_attn.qkv_proj.input_scale, + eps=self.next_layer_layernorm.variance_epsilon, + )) + else: + hidden_states, residual = self.all_reduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=self.next_layer_layernorm.weight, + eps=self.next_layer_layernorm.variance_epsilon, + )) + elif self.next_layer_layernorm: + hidden_states, residual = self.next_layer_layernorm( + hidden_states, residual) + elif needs_post_allreduce: + hidden_states, residual = self.all_reduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.NONE, + residual=residual, + )) + + return hidden_states, residual diff --git a/tensorrt_llm/_torch/models/modeling_siglip.py b/tensorrt_llm/_torch/models/modeling_siglip.py index 37678cf89ae7..ea770ea0acfd 100644 --- a/tensorrt_llm/_torch/models/modeling_siglip.py +++ b/tensorrt_llm/_torch/models/modeling_siglip.py @@ -108,4 +108,4 @@ def load_weights(self, weights: Dict): r'(.*?)fc1(.*)': r'\1up_proj\2', r'(.*?)fc2(.*)': r'\1down_proj\2', } - _load_weights_impl(self, weights, pattern_mapping) + _load_weights_impl(self, weights, params_map=pattern_mapping) diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 4dd10c592469..3865701f98b9 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -528,8 +528,8 @@ def forward( return_context_logits, ) - def load_weights(self, weights: Dict): - _load_weights_impl(self, weights) + def load_weights(self, weights: Dict, skip_modules: List[str] = []): + _load_weights_impl(self, weights, skip_modules) def infer_max_seq_len(self) -> int: # Modified from tensorrt_llm/builder.py _init_max_seq_len @@ -642,6 +642,7 @@ def filter_weights(prefix, weights: Dict): def _load_weights_impl(model: Union[nn.Module, DecoderModelForCausalLM], weights: Dict, + skip_modules: List[str] = [], params_map: Optional[Dict[str, str]] = None): if not hasattr(model, 'model_config') or not isinstance( model.model_config, ModelConfig): @@ -666,6 +667,10 @@ def _load_weights_impl(model: Union[nn.Module, DecoderModelForCausalLM], for name, module in tqdm(list(model.named_modules()), desc="Loading weights"): if len(module._parameters) > 0: + # skip load weights if module is in skip_modules + if any(skip_module in name for skip_module in skip_modules): + continue + # skip load weights if tie word embeddings is enabled and layer is lm_head if model.config.tie_word_embeddings and name.startswith("lm_head"): continue diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 4c7c8a5b6182..32dcea9fff57 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -1,7 +1,7 @@ import math import weakref from enum import IntEnum -from typing import Optional, cast +from typing import Optional, Union, cast import torch from torch import nn @@ -16,7 +16,7 @@ from ..distributed import AllReduceParams from ..model_config import ModelConfig from ..peft.lora.layer import LoraLayer, LoraModuleType -from ..utils import get_model_extra_attrs +from ..utils import Fp4QuantizedTensor, get_model_extra_attrs from .linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig from .multi_stream_utils import maybe_execute_in_parallel from .rms_norm import RMSNorm @@ -223,7 +223,7 @@ def convert_qkv(self, q, k, v): def forward( self, position_ids: Optional[torch.LongTensor], - hidden_states: torch.Tensor, + hidden_states: Union[torch.Tensor, Fp4QuantizedTensor], attn_metadata: AttentionMetadata, attention_mask: PredefinedAttentionMask = PredefinedAttentionMask. CAUSAL, diff --git a/tensorrt_llm/_torch/modules/embedding.py b/tensorrt_llm/_torch/modules/embedding.py index 2f0ddcfd3bec..94a088c5d2cc 100644 --- a/tensorrt_llm/_torch/modules/embedding.py +++ b/tensorrt_llm/_torch/modules/embedding.py @@ -43,6 +43,8 @@ def __init__( elif tensor_parallel_mode == TensorParallelMode.COLUMN: local_out_features = math.ceil(num_embeddings / tp_size) self.padding_size = tp_size * local_out_features - num_embeddings + else: + self.padding_size = 0 super().__init__( local_in_features * tp_size, @@ -120,6 +122,43 @@ def get_masked_input_and_mask( return input_, ~vocab_mask.unsqueeze(-1) +# We use torch.compile() to fuse the tiny pointwise ops before all_reduce/all_gather for Embedding module. +@torch.compile(mode="max-autotune-no-cudagraphs") +def pre_comm_embedding_ops( + input_: torch.Tensor, + weight: torch.Tensor, + tp_size: int, + tp_rank: int, + tp_mode: TensorParallelMode, + vocab_start_index: int, + vocab_end_index: int, + gather_output: bool, + padding_size: int, +): + # Generate the mask for the input if needed. + if tp_size > 1: + if tp_mode == TensorParallelMode.COLUMN: + input_, input_mask = get_masked_input_and_mask( + input_, + vocab_start_index, + vocab_end_index, + ) + + # Get the embeddings. + output = F.embedding(input_, weight) + + # Mask or pad the output if needed. + if tp_size > 1: + if tp_mode == TensorParallelMode.COLUMN: + output.masked_fill_(input_mask, 0) + elif tp_mode == TensorParallelMode.ROW: + if gather_output: + if tp_rank == tp_size - 1 and padding_size > 0: + output = F.pad(output, (0, padding_size)) + + return output + + class Embedding(LMHead): """Embedding layer. @@ -155,29 +194,28 @@ def __init__( self.vocab_start_index = self.tp_rank * slice_width self.vocab_end_index = min((self.tp_rank + 1) * slice_width, num_embeddings) + else: + self.vocab_start_index = 0 + self.vocab_end_index = num_embeddings def forward(self, input): + # Run the ops before all_reduce/all_gather. + output = pre_comm_embedding_ops(input, self.weight, self.tp_size, + self.tp_rank, self.tp_mode, + self.vocab_start_index, + self.vocab_end_index, + self.gather_output, self.padding_size) + + # Run the all_reduce/all_gather. if self.tp_size > 1: if self.tp_mode == TensorParallelMode.COLUMN: - # Build the mask. - input, input_mask = get_masked_input_and_mask( - input, - self.vocab_start_index, - self.vocab_end_index, - ) - # Get the embeddings. - output = F.embedding(input, self.weight) - # Mask the output embedding. - if self.tp_size > 1: - if self.tp_mode == TensorParallelMode.COLUMN: - output.masked_fill_(input_mask, 0) # Reduce across all the model parallel GPUs. output = self.all_reduce(output) elif self.tp_mode == TensorParallelMode.ROW: if self.gather_output: - if self.tp_rank == self.tp_size - 1 and self.padding_size > 0: - output = F.pad(output, (0, self.padding_size)) + # Run allgather. output = allgather(output, self.mapping) + # Remove the padding. if self.padding_size > 0: output = output[..., :-self.padding_size] diff --git a/tensorrt_llm/_torch/modules/fused_moe.py b/tensorrt_llm/_torch/modules/fused_moe.py index 91f5ec42faed..6e5e120bf1f1 100755 --- a/tensorrt_llm/_torch/modules/fused_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe.py @@ -1036,6 +1036,7 @@ def _check_configs(self): if self.is_trtllm(): # trtllm_gen backend only support min-latency mode now + assert not self.apply_router_weight_on_input, "TRTLLM backend does not support applying router weight on input yet." assert not self.reduce_results assert self.quant_config and ( self.quant_config.quant_mode.has_nvfp4() @@ -1421,6 +1422,8 @@ def forward_chunk( assert token_selected_slots.dtype == torch.int32 if self.apply_router_weight_on_input: + assert self.routing_method.top_k == 1, "Current workaround only supports top-1 routing" + assert x.dtype != torch.float8_e4m3fn, "Current workaround for apply_router_weight_on_input does not support fp8 input" x = x * token_final_scales.to(x.dtype) # TODO: remove this once we have correct fusedmoe kernel ready token_final_scales = None @@ -1578,6 +1581,7 @@ def forward_cutlass( num_rows = sum(all_rank_num_tokens) else: num_rows = x.shape[0] + # in case of num_rows is larger than max_chunk_size, we need to split the input into multiple chunks num_chunks = (num_rows + self.moe_max_num_tokens - 1) // self.moe_max_num_tokens @@ -2204,9 +2208,18 @@ def load_expert_fc2_input_scale_nvfp4( dst_fc2_input_scale.copy_(w2_input_scale[...].reshape([])) for expert_id in range(self.num_experts): - w1_input_scale = weights[f"{expert_id}.w1.input_scale"] - w3_input_scale = weights[f"{expert_id}.w3.input_scale"] - w2_input_scale = weights[f"{expert_id}.w2.input_scale"] + if self.weight_loading_mode == MoEWeightLoadingMode.VANILLA: + w1_input_scale = weights[f"{expert_id}.w1.input_scale"] + w3_input_scale = weights[f"{expert_id}.w3.input_scale"] + w2_input_scale = weights[f"{expert_id}.w2.input_scale"] + elif self.weight_loading_mode == MoEWeightLoadingMode.FUSED_GATE_UP_PROJ: + w1_input_scale = weights["gate_up_proj_input_scale"] + w3_input_scale = weights["gate_up_proj_input_scale"] + w2_input_scale = weights["down_proj_input_scale"] + else: + raise NotImplementedError( + f"Unknown weight loading mode in MoE: {self.weight_loading_mode}" + ) load_expert_fc31_input_scale_nvfp4(w1_input_scale, w3_input_scale, tmp_fc31_input_scale[expert_id]) diff --git a/tensorrt_llm/_torch/modules/gated_mlp.py b/tensorrt_llm/_torch/modules/gated_mlp.py index 877884227877..a727cc93ab92 100644 --- a/tensorrt_llm/_torch/modules/gated_mlp.py +++ b/tensorrt_llm/_torch/modules/gated_mlp.py @@ -74,6 +74,7 @@ def __init__(self, reduce_output=False, skip_create_weights_in_init=config.skip_create_weights_in_init, ) + self.down_lora = LoraLayer([LoraModuleType.MLP_4H_TO_H], [self.hidden_size]) @@ -102,6 +103,16 @@ def __init__(self, [LoraModuleType.MLP_GATE_UP], [2 * self.intermediate_size // mapping.tp_size]) + def _apply_activation(self, x): + if self.activation == F.silu: + return swiglu(x) + elif self.activation == None: + return x + else: + raise NotImplementedError( + f"Activation {self.activation} not yet implemented for fused GatedMLP" + ) + def forward( self, x: Union[torch.Tensor, Fp4QuantizedTensor], @@ -114,18 +125,12 @@ def forward( return self.forward_lora(x, all_rank_num_tokens, final_all_reduce_params, lora_params) - if self.activation == F.silu: - h1 = self.gate_up_proj(x) - - h2 = swiglu(h1) - output = self.down_proj(h2, - all_reduce_params=final_all_reduce_params, - layer_idx=self.layer_idx) - return output - else: - raise NotImplementedError( - f"Activation {self.activation} not yet implemented for fused GatedMLP" - ) + h1 = self.gate_up_proj(x) + h2 = self._apply_activation(h1) + output = self.down_proj(h2, + all_reduce_params=final_all_reduce_params, + layer_idx=self.layer_idx) + return output def forward_lora( self, @@ -136,7 +141,6 @@ def forward_lora( ) -> torch.Tensor: assert lora_params is not None assert self.layer_idx is not None, "layer_idx is required for lora" - assert self.activation == F.silu h1 = self.gate_up_proj(x) @@ -149,7 +153,7 @@ def forward_lora( if h1_lora is not None: h1 = h1 + h1_lora - h2 = swiglu(h1) + h2 = self._apply_activation(h1) output = self.down_proj(h2, all_reduce_params=final_all_reduce_params, lora_params=lora_params, diff --git a/tensorrt_llm/_torch/pyexecutor/config.py b/tensorrt_llm/_torch/pyexecutor/config.py index 3515dab60f0e..533b21b05021 100644 --- a/tensorrt_llm/_torch/pyexecutor/config.py +++ b/tensorrt_llm/_torch/pyexecutor/config.py @@ -84,6 +84,9 @@ class PyTorchConfig: # from the model checkpoint. load_format: Union[str, LoadFormat] = 'auto' + # If true, enable min-latency mode. Currently only used for Llama4. + enable_min_latency: bool = False + EXETENDED_EXECUTOR_CONFIG_FIELDS = [ 'backend', diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 8dcfe25dfbb0..176f5a51c3fb 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -430,10 +430,11 @@ def __init__( self.previous_kv_lens_offsets_cuda = torch.zeros((batch_size, ), dtype=torch.int, device='cuda') - self.is_mtp = self.spec_config.spec_dec_mode.is_mtp() + self.without_logits = self.spec_config.spec_dec_mode.without_logits( + ) self.max_draft_len = spec_config.max_draft_tokens else: - self.is_mtp = False + self.without_logits = False self.max_draft_len = 0 self.iter_counter = 0 @@ -729,6 +730,7 @@ def _set_up_spec_metadata( return get_spec_metadata( self.spec_config, self.batch_size, + max_num_tokens=self.max_num_tokens, spec_resource_manager=spec_resource_manager) if self.spec_metadata is not None: @@ -736,6 +738,7 @@ def _set_up_spec_metadata( self.spec_metadata = get_spec_metadata( self.spec_config, self.batch_size, + max_num_tokens=self.max_num_tokens, spec_resource_manager=spec_resource_manager) return self.spec_metadata @@ -883,6 +886,7 @@ def _load_model(self, config = ModelConfig.from_pretrained(checkpoint_dir, trust_remote_code=True, **kwargs) + config.pytorch_backend_config = self.pytorch_backend_config config.spec_config = self.spec_config config.max_num_tokens = max_num_tokens config.moe_max_num_tokens = moe_max_num_tokens @@ -935,6 +939,12 @@ def init_meta_tensor(t: torch.Tensor): model.load_weights(weights) + if self.spec_config is not None and self.spec_config.spec_dec_mode.need_load_draft_weights( + ): + weights = load_weights(self.spec_config.draft_model_path, + self.mapping) + model.load_draft_weights(weights) + elif load_format == LoadFormat.DUMMY: initialize_dummy_weights(model) @@ -1069,7 +1079,7 @@ def _prepare_tp_inputs( if new_tensors_device is not None: # speculative decoding cases: [batch, 1 + draft_len], others: [batch] new_tokens_device = new_tensors_device.new_tokens - if self.is_mtp: + if self.without_logits: assert isinstance(new_tensors_device, SampleStateTensorsMTP) new_tokens_lens_device = new_tensors_device.new_tokens_lens # [batch] next_draft_tokens_device = new_tensors_device.next_draft_tokens # [batch, draft_len] @@ -1927,7 +1937,7 @@ def model_forward(self, **kwargs): def _forward_step(self, inputs: Dict[str, Any], gather_ids: Optional[torch.Tensor]) -> Dict[str, Any]: inputs = self._preprocess_inputs(inputs) - if self.is_mtp: + if self.without_logits: outputs = self.model_forward(**inputs) return outputs diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 4c4fc628c1ed..57c04cc1c20f 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1833,7 +1833,8 @@ def _prepare_draft_tokens(self, scheduled_requests: ScheduledRequests): self.resource_manager, extra_model_inputs=extra_model_inputs) - if spec_metadata.spec_dec_mode.is_eagle3(): + if spec_metadata.spec_dec_mode.is_eagle3() and hasattr( + self.draft_model_engine.model.model, 'd2t'): outputs['d2t'] = self.draft_model_engine.model.model.d2t.data sample_state = self._sample_async(draft_batch, outputs) @@ -1885,7 +1886,8 @@ def _pad_to_max_draft_tokens(): self.resource_manager, extra_model_inputs=extra_model_inputs) - if spec_metadata.spec_dec_mode.is_eagle3(): + if spec_metadata.spec_dec_mode.is_eagle3() and hasattr( + self.draft_model_engine.model.model, 'd2t'): outputs[ 'd2t'] = self.draft_model_engine.model.model.d2t.data sample_state = self._sample_async(draft_batch, outputs) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 3f5b8cffaf00..675c8a213fea 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -18,7 +18,7 @@ from ..attention_backend.interface import AttentionRuntimeFeatures from ..distributed import MPIDist -from ..speculative import Eagle3Config, NGramConfig, get_spec_resource_manager +from ..speculative import NGramConfig, get_spec_resource_manager from ._util import (create_kv_cache_manager, create_py_executor_instance, estimate_max_kv_cache_tokens, get_token_num_for_estimation, instantiate_sampler, is_mla) @@ -177,7 +177,9 @@ def create_py_executor(executor_config: ExecutorConfig, dist = MPIDist(mapping=mapping) spec_config = executor_config.speculative_config - has_draft_model_engine = isinstance(spec_config, Eagle3Config) + has_draft_model_engine = False + if spec_config is not None: + has_draft_model_engine = spec_config.spec_dec_mode.has_draft_model() has_ngram_drafter = isinstance(spec_config, NGramConfig) attn_runtime_features = AttentionRuntimeFeatures( @@ -186,6 +188,7 @@ def create_py_executor(executor_config: ExecutorConfig, has_speculative_draft_tokens=has_draft_model_engine or has_ngram_drafter, ) + logger.info("ATTENTION RUNTIME FEATURES: ", attn_runtime_features) mem_monitor = _ExecutorMemoryMonitor() with mem_monitor.observe_creation_stage( @@ -213,7 +216,7 @@ def create_py_executor(executor_config: ExecutorConfig, draft_spec_config.max_draft_tokens = 0 draft_model_engine = PyTorchModelEngine( - spec_config.eagle_weights_path, + spec_config.draft_model_path, pytorch_backend_config, batch_size=executor_config.max_batch_size, max_num_tokens=executor_config.max_num_tokens, diff --git a/tensorrt_llm/_torch/speculative/__init__.py b/tensorrt_llm/_torch/speculative/__init__.py index 93ff6463c863..b1fb18cf1d98 100644 --- a/tensorrt_llm/_torch/speculative/__init__.py +++ b/tensorrt_llm/_torch/speculative/__init__.py @@ -3,11 +3,11 @@ from .mtp import MTPConfig, MTPEagleWorker, MTPSpecMetadata, MTPWorker from .ngram import NGramConfig from .utils import (get_num_spec_layers, get_spec_decoder, get_spec_metadata, - get_spec_resource_manager) + get_spec_resource_manager, get_spec_worker) __all__ = [ - "SpecConfig", "SpecMetadata", "MTPConfig", "MTPWorker", "MTPEagleWorker", - "Eagle3Config", "Eagle3SpecMetadata", "MTPSpecMetadata", + "SpecConfig", "SpecMetadata", "MTPConfig", "MTPEagleWorker", + "MTPSpecMetadata", "MTPWorker", "Eagle3Config", "Eagle3SpecMetadata", "get_spec_metadata", "get_spec_resource_manager", "get_spec_decoder", - "get_num_spec_layers", "NGramConfig" + "get_num_spec_layers", "get_spec_worker", "NGramConfig" ] diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index b5bdd1994acc..67aec85c4447 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -3,25 +3,36 @@ from typing import Dict, List, Optional, Tuple import torch +from torch import nn +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping + +from ..attention_backend import AttentionMetadata from ..pyexecutor.sampler import SampleState, SampleStateTensors, TorchSampler from .interface import SpecConfig, SpecMetadata, SpeculativeDecodingMode +from .mtp import MTPSampler @dataclass class Eagle3Config(SpecConfig): spec_dec_name: str = "EAGLE3" - eagle_weights_path: Optional[str] = None num_layers: int = 0 hidden_size: int = 0 + eagle3_one_model: bool = True def __post_init__(self): - if self.eagle_weights_path is None: + if self.draft_model_path is None: raise ValueError("Path to EAGLE3 weights must be specified.") - self.spec_dec_mode = SpeculativeDecodingMode.from_string( - self.spec_dec_name) - self.num_extra_kv_tokens = 0 + if self.eagle3_one_model: + self.spec_dec_mode = SpeculativeDecodingMode.EAGLE3_ONE_MODEL + self.num_extra_kv_tokens = self.max_draft_tokens - 1 + else: + self.spec_dec_mode = SpeculativeDecodingMode.from_string( + self.spec_dec_name) + self.num_extra_kv_tokens = 0 + logger.info(f"EAGLE3 Config: {self}") def update_from_model_config(self, model_config): self.num_layers = model_config.num_hidden_layers @@ -66,6 +77,9 @@ def prepare(self): if not self.is_cuda_graph: self.hidden_states = [] + def is_layer_capture(self, layer_id: int): + return layer_id in self.layers_to_capture + def maybe_capture_hidden_states(self, layer_id: int, hidden_states: torch.Tensor, residual: torch.Tensor) -> None: @@ -130,3 +144,325 @@ def _batch_sample(self, scheduled_requests, model_outputs) -> SampleState: device=device, host=host, sampler_event=sampler_event) + + +@dataclass +class Eagle3OneModelSpecMetadata(SpecMetadata): + # The hidden states + hidden_states: Optional[torch.Tensor] = None + # The number of layers + num_layers: int = 0 + # The layers to be captured + layers_to_capture: Tuple[int, ...] = field(init=False) + # The hidden size of the hidden states + hidden_size: int = 0 + # The max number of tokens + max_num_tokens: int = 0 + # The dtype of the hidden states + dtype: torch.dtype = torch.bfloat16 + # The index of the batche inputs + batch_indices_cuda: Optional[torch.Tensor] = None + + def __post_init__(self): + if self.num_layers == 1: + self.layers_to_capture = (1, ) + else: + if self.num_layers <= 5: + raise ValueError("Not enough hidden layers for EAGLE") + + self.layers_to_capture = (1, self.num_layers // 2 - 1, + self.num_layers - 4) + self.hidden_states = torch.empty( + (self.max_num_tokens, + self.hidden_size * len(self.layers_to_capture)), + dtype=self.dtype, + device='cuda') + self.batch_indices_cuda = torch.empty( + [self.max_num_requests], + dtype=torch.int, + device='cuda', + ) + + def is_layer_capture(self, layer_id: int): + return layer_id in self.layers_to_capture + + def prepare(self): + assert self.request_ids is not None + # update batch indeices + num_seqs = len(self.request_ids) + batch_indices = torch.arange(num_seqs, + dtype=torch.int, + device='cpu', + pin_memory=True) + self.batch_indices_cuda[:num_seqs].copy_(batch_indices, + non_blocking=True) + self.num_tokens -= (self.num_generations) * self.max_draft_tokens + + def maybe_capture_hidden_states( + self, + layer_id: int, + hidden_states: torch.Tensor, + residual: Optional[torch.Tensor] = None) -> None: + for i, captured_layer_id in enumerate(self.layers_to_capture): + if captured_layer_id == layer_id: + num_tokens = hidden_states.shape[0] + to_save = hidden_states + residual if residual is not None else hidden_states + self.hidden_states[:num_tokens, i * self.hidden_size:(i + 1) * + self.hidden_size].copy_(to_save, + non_blocking=True) + break + + +class Eagle3Decoder(TorchSampler): + + def _batch_sample(self, scheduled_requests, model_outputs) -> SampleState: + logits = model_outputs["logits"] + new_tokens_device = torch.argmax(logits, dim=-1) + if "d2t" in model_outputs: + d2t = model_outputs["d2t"] + new_tokens_device = d2t[new_tokens_device] + new_tokens_device + new_tokens_host = new_tokens_device.to('cpu', non_blocking=True) + new_tensors_device = {"new_tokens_device": new_tokens_device} + new_tensors_host = {"new_tokens_host": new_tokens_host} + decoder_event = torch.cuda.Event() + decoder_event.record() + return SampleState(scheduled_requests=scheduled_requests, + logits=logits, + new_tensors_device=new_tensors_device, + new_tensors_host=new_tensors_host, + decoder_event=decoder_event) + + +class Eagle3OneModelDecoder(MTPSampler): + + def __init__(self, max_seq_len: int, config: Eagle3Config): + super().__init__(max_seq_len, None) + self.draft_len = config.max_draft_tokens + + +class Eagle3OneModelWorker(nn.Module): + + def __init__(self, spec_config: Eagle3Config, mapping: Mapping): + super().__init__() + self.spec_config = spec_config + self.max_draft_tokens = self.spec_config.max_draft_tokens + self.mapping = mapping + + @torch.compile(mode="max-autotune-no-cudagraphs") + def forward(self, input_ids, position_ids, hidden_states, logits, + attn_metadata, spec_metadata, draft_model): + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + + raw_logits = logits + + # Sample and accept tokens + accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( + logits, attn_metadata, spec_metadata) + + # Save the old attn_metadata and spec_metadata + if attn_metadata.is_cuda_graph: + seq_len = attn_metadata._seq_lens[:batch_size].clone() + seq_len_cuda = attn_metadata._seq_lens_cuda[:batch_size].clone() + + # Prepare inputs for the 1st draft model forward + position_ids = position_ids.squeeze(0) + last_tokens_idx = torch.cumsum( + attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1 + inputs = self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + last_tokens_idx=last_tokens_idx, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model) + + # Predict draft tokens + next_draft_tokens = [] + for i in range(self.max_draft_tokens): + hidden_states, hidden_states_to_save = draft_model.model(**inputs) + if i == 0: + start_ids_gen = (spec_metadata.batch_indices_cuda[:num_gens] * + (self.max_draft_tokens + 1)).long() + gather_ids_gen = (start_ids_gen + + num_accepted_tokens[num_contexts:] - 1 + + attn_metadata.num_ctx_tokens) + gather_ids = torch.concat( + [last_tokens_idx[:num_contexts], gather_ids_gen], dim=0) + else: + # All of the seq_len are 1, use batch_indices_cuda as gather_ids + gather_ids = spec_metadata.batch_indices_cuda[:batch_size] + logits = draft_model.logits_processor(hidden_states[gather_ids], + draft_model.lm_head, + attn_metadata, True) + new_draft_token = self.draft_decoder(logits, draft_model) + next_draft_tokens.append(new_draft_token) + # update inputs + hidden_states = hidden_states_to_save[gather_ids] + position_ids = inputs["position_ids"][gather_ids] + 1 + # update attn_metadata + if i == 0: + attn_metadata._seq_lens[:batch_size].fill_(1) + attn_metadata._seq_lens_cuda[:batch_size].fill_(1) + attn_metadata.on_update() + # cannot run generation if their is no kv cache + if inputs["attn_metadata"].kv_cache_manager is not None: + attn_metadata.host_request_types[:attn_metadata. + num_contexts].fill_(1) + attn_metadata.num_contexts = 0 + # update kv_lens_cuda + if hasattr(attn_metadata, 'kv_lens_cuda'): + attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= ( + self.max_draft_tokens - + num_accepted_tokens[num_contexts:]) + attn_metadata.kv_lens_cuda[:num_contexts] += 1 + elif hasattr(attn_metadata, 'kv_lens_cuda'): + attn_metadata.kv_lens_cuda[:batch_size] += 1 + # support attention dp + if spec_metadata.all_rank_num_tokens is not None: + spec_metadata.all_rank_num_tokens = spec_metadata.all_rank_num_seqs + inputs = { + "input_ids": new_draft_token, + "position_ids": position_ids, + "hidden_states": hidden_states, + "attn_metadata": attn_metadata, + "spec_metadata": spec_metadata, + } + next_draft_tokens = torch.stack(next_draft_tokens, dim=1) + + # restore attn_metadata to support cuda graph + if attn_metadata.is_cuda_graph: + attn_metadata._seq_lens[:batch_size].copy_(seq_len) + attn_metadata._seq_lens_cuda[:batch_size].copy_(seq_len_cuda) + attn_metadata.on_update() + + # prepare next new tokens to support overlap scheduler + next_new_tokens = accepted_tokens[ + spec_metadata.batch_indices_cuda[:batch_size], + num_accepted_tokens - 1].unsqueeze(1) + next_new_tokens = torch.concat([next_new_tokens, next_draft_tokens], + dim=1) + + return { + 'logits': raw_logits, + 'new_tokens': accepted_tokens, + 'new_tokens_lens': num_accepted_tokens, + 'next_draft_tokens': next_draft_tokens, + 'next_new_tokens': next_new_tokens, + } + + def sample_and_accept_draft_tokens( + self, + logits: torch.Tensor, + attn_metadata: AttentionMetadata, + spec_metadata: Eagle3OneModelSpecMetadata, + ): + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + + if logits.dim() == 1: + logits = logits.unsqueeze(0) + + # The return buffer + accepted_tokens = torch.empty((batch_size, (self.max_draft_tokens + 1)), + dtype=torch.int, + device=logits.device) + num_accepted_tokens = torch.ones(batch_size, + dtype=torch.int, + device=logits.device) + + # Do greedy sampling for the input logits + target_tokens = torch.argmax(logits, dim=-1) + + # context + accepted_tokens[:num_contexts, 0] = target_tokens[:num_contexts] + + # generation + gen_target_tokens = target_tokens[num_contexts:].reshape( + num_gens, self.max_draft_tokens + 1) + accepted_tokens[num_contexts:, :] = gen_target_tokens + draft_tokens = spec_metadata.draft_tokens.reshape( + num_gens, self.max_draft_tokens) + num_accepted_tokens[num_contexts:] += torch.cumprod(( + draft_tokens == gen_target_tokens[:, :self.max_draft_tokens]).int(), + dim=-1).sum(1) + + return accepted_tokens, num_accepted_tokens + + def draft_decoder( + self, + logits: torch.Tensor, + draft_model: nn.Module, + ): + ''' + Sampling draft tokens. + + Args: + logits: torch.Tensor + [num_tokens, vocab_size] + Logits produced by the draft model. + draft_model: nn.Module + The draft model. + + Returns: + draft_tokens: torch.Tensor + [batch_size * max_draft_tokens] + Draft token ids. Flattened. + ''' + + draft_tokens = torch.argmax(logits, dim=-1).type(torch.int32) + + # Apply d2t (offsets between draft model dictionary and main model dictionary). + if hasattr(draft_model.model, + "d2t") and draft_model.model.d2t is not None: + draft_tokens = draft_model.model.d2t[draft_tokens] + draft_tokens + + return draft_tokens + + def prepare_1st_drafter_inputs( + self, + input_ids: torch.LongTensor, + position_ids: torch.LongTensor, + last_tokens_idx: torch.LongTensor, + hidden_states: torch.Tensor, + accepted_tokens: torch.Tensor, + attn_metadata: AttentionMetadata, + spec_metadata: Eagle3OneModelSpecMetadata, + draft_model: nn.Module, + ): + num_contexts = attn_metadata.num_contexts + num_tokens = input_ids.shape[0] + + # prepare hidden states + hidden_size_up = spec_metadata.hidden_size * len( + spec_metadata.layers_to_capture) + hidden_states = spec_metadata.hidden_states[:num_tokens, : + hidden_size_up] + hidden_states = draft_model.apply_eagle3_fc(hidden_states) + + # context + input_ctx_ids = input_ids[:attn_metadata.num_ctx_tokens] + input_ids_ctx = torch.empty_like(input_ctx_ids, + dtype=torch.int32, + device="cuda") + input_ids_ctx[:-1].copy_(input_ctx_ids[1:]) + input_ids_ctx[ + last_tokens_idx[:num_contexts]] = accepted_tokens[:num_contexts, 0] + + # generation + input_ids_gen = accepted_tokens[num_contexts:, :].flatten() + + # get draft inputs + input_ids = torch.concat([input_ids_ctx, input_ids_gen], dim=0) + + return { + "input_ids": input_ids, + "position_ids": position_ids, + "hidden_states": hidden_states, + "attn_metadata": attn_metadata, + "spec_metadata": spec_metadata, + } diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 22e3c0b348d3..d3fc6d48e30d 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -15,6 +15,7 @@ class SpeculativeDecodingMode(IntEnum): MTP = auto() MTP_EAGLE = auto() EAGLE3 = auto() + EAGLE3_ONE_MODEL = auto() NGRAM = auto() NONE = auto() @@ -27,20 +28,36 @@ def is_mtp_eagle(self): def is_eagle3(self): return self == SpeculativeDecodingMode.EAGLE3 + def is_eagle3_one_model(self): + return self == SpeculativeDecodingMode.EAGLE3_ONE_MODEL + def is_ngram(self): return self == SpeculativeDecodingMode.NGRAM def is_none(self): return self == SpeculativeDecodingMode.NONE + def without_logits(self): + return self.is_mtp() or self.is_eagle3_one_model() + def needs_kv_cache_rewind(self): - return self.is_mtp() + return self.is_mtp() or self.is_eagle3_one_model() def support_overlap_scheduler(self): - return self.is_mtp() + return self.is_mtp() or self.is_eagle3_one_model() + + def has_draft_model(self): + return self.is_eagle3() + + def need_load_draft_weights(self): + """ + Whether the draft model and target model are in the same model engine, + and the draft model needs to load weights from the separate checkpoint. + """ + return self.is_eagle3_one_model() def has_spec_decoder(self): - return self.is_mtp() or self.is_eagle3() + return self.is_mtp() or self.is_eagle3() or self.is_eagle3_one_model() def extend_ctx(self, attention_backend: AttentionBackend): """ @@ -72,6 +89,8 @@ class SpecConfig: spec_dec_mode: SpeculativeDecodingMode = SpeculativeDecodingMode.NONE # The max number of draft tokens max_draft_tokens: int = 1024 + # The path to the draft model + draft_model_path: Optional[str] = None def __post_init__(self) -> None: self.spec_dec_mode = SpeculativeDecodingMode.from_string( diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 6f5ed5bc7aa8..d2c9e89cf776 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -226,7 +226,9 @@ class MTPSampler(TorchSampler): def __init__(self, max_seq_len: int, config: MTPConfig): super().__init__(max_seq_len, False) self.mapping = None - self.draft_len = config.num_nextn_predict_layers + self.draft_len = 0 + if config is not None: + self.draft_len = config.num_nextn_predict_layers def _draft_meet_max_token_stop_criteria(self, request: LlmRequest, num_tokens: int, beam_idx: int): diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 62fdfee7d63b..ed7e433cc891 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -1,10 +1,13 @@ -from .eagle3 import Eagle3Sampler, Eagle3SpecMetadata -from .mtp import MTPHiddenStatesManager, MTPSampler, MTPSpecMetadata +from .eagle3 import (Eagle3OneModelDecoder, Eagle3OneModelSpecMetadata, + Eagle3OneModelWorker, Eagle3Sampler, Eagle3SpecMetadata) +from .mtp import (MTPEagleWorker, MTPHiddenStatesManager, MTPSampler, + MTPSpecMetadata, MTPWorker) from .ngram import NGramPoolManager def get_spec_metadata(spec_config, max_num_requests, + max_num_tokens, spec_resource_manager=None): if spec_config.spec_dec_mode.is_mtp(): return MTPSpecMetadata( @@ -19,6 +22,14 @@ def get_spec_metadata(spec_config, max_num_requests=max_num_requests, num_layers=spec_config.num_layers, hidden_size=spec_config.hidden_size) + elif spec_config.spec_dec_mode.is_eagle3_one_model(): + return Eagle3OneModelSpecMetadata( + max_draft_tokens=spec_config.max_draft_tokens, + spec_dec_mode=spec_config.spec_dec_mode, + max_num_requests=max_num_requests, + num_layers=spec_config.num_layers, + hidden_size=spec_config.hidden_size, + max_num_tokens=max_num_tokens) else: return None @@ -44,8 +55,10 @@ def get_spec_resource_manager(spec_config, model_config, max_num_requests): def get_spec_decoder(max_seq_len, spec_config): if spec_config.spec_dec_mode.is_mtp(): return MTPSampler(max_seq_len, spec_config) - if spec_config.spec_dec_mode.is_eagle3(): + elif spec_config.spec_dec_mode.is_eagle3(): return Eagle3Sampler(max_seq_len) + elif spec_config.spec_dec_mode.is_eagle3_one_model(): + return Eagle3OneModelDecoder(max_seq_len, spec_config) else: return None @@ -53,5 +66,18 @@ def get_spec_decoder(max_seq_len, spec_config): def get_num_spec_layers(spec_config): if spec_config.spec_dec_mode.is_mtp(): return spec_config.num_nextn_predict_layers + elif spec_config.spec_dec_mode.is_eagle3_one_model(): + return 1 else: return 0 + + +def get_spec_worker(spec_config, mapping): + if spec_config.spec_dec_mode.is_mtp(): + return MTPWorker(spec_config) + elif spec_config.spec_dec_mode.is_mtp_eagle(): + return MTPEagleWorker(spec_config) + elif spec_config.spec_dec_mode.is_eagle3_one_model(): + return Eagle3OneModelWorker(spec_config, mapping) + else: + return None diff --git a/tensorrt_llm/bench/benchmark/throughput.py b/tensorrt_llm/bench/benchmark/throughput.py index 71797268de7f..189be0a6a520 100755 --- a/tensorrt_llm/bench/benchmark/throughput.py +++ b/tensorrt_llm/bench/benchmark/throughput.py @@ -220,6 +220,19 @@ required=False, help="Path where output should be written to.", ) +@optgroup.option( + "--enable_chunked_context", + is_flag=True, + default=False, + help="Enable chunking in prefill stage for enhanced throughput benchmark.", +) +@optgroup.option( + "--scheduler_policy", + type=click.Choice(["guaranteed_no_evict", "max_utilization"]), + default="guaranteed_no_evict", + help= + "KV cache scheduler policy: guaranteed_no_evict prevents request eviction, max_utilization optimizes for throughput.", +) @click.pass_obj def throughput_command( bench_env: BenchmarkEnvironment, @@ -316,6 +329,8 @@ def throughput_command( kv_cache_percent = params.pop("kv_cache_free_gpu_mem_fraction") beam_width = params.pop("beam_width") streaming: bool = params.pop("streaming") + enable_chunked_context: bool = params.pop("enable_chunked_context") + scheduler_policy: str = params.pop("scheduler_policy") # Update configuration with runtime options exec_settings["settings_config"]["kv_cache_percent"] = kv_cache_percent @@ -323,7 +338,8 @@ def throughput_command( exec_settings["settings_config"]["max_num_tokens"] = runtime_max_tokens exec_settings["settings_config"]["beam_width"] = beam_width exec_settings["settings_config"][ - "scheduler_policy"] = CapacitySchedulerPolicy.GUARANTEED_NO_EVICT + "scheduler_policy"] = CapacitySchedulerPolicy.GUARANTEED_NO_EVICT if scheduler_policy == "guaranteed_no_evict" else CapacitySchedulerPolicy.MAX_UTILIZATION + exec_settings["settings_config"]["chunking"] = enable_chunked_context # Dynamic runtime features. exec_settings["settings_config"]["dynamic_max_batch_size"] = True diff --git a/tensorrt_llm/bench/dataclasses/reporting.py b/tensorrt_llm/bench/dataclasses/reporting.py index 699c61654e81..36d3923cc63b 100755 --- a/tensorrt_llm/bench/dataclasses/reporting.py +++ b/tensorrt_llm/bench/dataclasses/reporting.py @@ -519,9 +519,9 @@ def report_statistics(self) -> None: ["minimum", "maximum", "average", "p50", "p90", "p95", "p99"]) perf_stats += ( - f"Average time-to-first-token [TTFT] (ms): {streaming['avg_ttft_ms']:.4f}\n" - f"Average time-per-output-token [TPOT] (ms): {streaming['avg_tpot_ms']:.4f}\n" - f"Per User Output Speed (tps/user): {streaming['token_output_speed_tok_s']:.4f}\n" + f"Average time-to-first-token [TTFT] (ms): {streaming['avg_ttft_ms']:.4f}\n" + f"Average time-per-output-token [TPOT] (ms): {streaming['avg_tpot_ms']:.4f}\n" + f"Per User Output Speed (tps/user): {streaming['token_output_speed_tok_s']:.4f}\n" "\n-- Per-Request Time-per-Output-Token [TPOT] Breakdown (ms)\n\n" f"{tpot_stats}\n" "\n-- Per-Request Time-to-First-Token [TTFT] Breakdown (ms) \n\n" diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 6d769a12ee3a..82c3bc8cf2db 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -239,6 +239,7 @@ class EagleDecodingConfig(DecodingBaseConfig): num_eagle_layers: Optional[int] = None max_non_leaves_per_layer: Optional[int] = None pytorch_eagle_weights_path: Optional[str] = None + eagle3_one_model: Optional[bool] = True @classmethod def from_dict(cls, data: dict): @@ -1204,8 +1205,10 @@ def _setup_speculative_config(self): from tensorrt_llm._torch.speculative import Eagle3Config self.speculative_config = Eagle3Config( max_draft_tokens=self.speculative_config.max_draft_len, - eagle_weights_path=self.speculative_config. - pytorch_eagle_weights_path) + draft_model_path=self.speculative_config. + pytorch_eagle_weights_path, + eagle3_one_model=self.speculative_config. + eagle3_one_model) elif isinstance(self.speculative_config, NGramDecodingConfig): self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.NGRAM assert self.backend == 'pytorch' @@ -1573,6 +1576,12 @@ class TorchLlmArgs(BaseLlmArgs): "How to load the model weights. By default, detect the weight type from the model checkpoint." ) + enable_min_latency: bool = Field( + default=False, + description= + "If true, enable min-latency mode. Currently only used for Llama4.", + ) + @field_validator('load_format', mode='before') @classmethod def convert_load_format(cls, v): @@ -1658,7 +1667,8 @@ def get_pytorch_backend_config(self) -> "PyTorchConfig": torch_compile_enable_userbuffers, autotuner_enabled=self.autotuner_enabled, enable_layerwise_nvtx_marker=self.enable_layerwise_nvtx_marker, - load_format=self.load_format) + load_format=self.load_format, + enable_min_latency=self.enable_min_latency) @field_validator('cuda_graph_max_batch_size') @classmethod diff --git a/tensorrt_llm/plugin/plugin.py b/tensorrt_llm/plugin/plugin.py index 44d68c758fdd..2792ba2d647b 100644 --- a/tensorrt_llm/plugin/plugin.py +++ b/tensorrt_llm/plugin/plugin.py @@ -681,7 +681,7 @@ class CustomAllReduceHelper: Then, each instance of allreduce will reference that tensor automatically. """ POINTERS_PER_RANK = 7 - POINTERS_OF_COUNTER = 2 + POINTERS_OF_COUNTER = 3 def __init__(self) -> None: self.workspace: Optional[Tensor] = None @@ -740,11 +740,11 @@ def allocate_workspace(mapping: Mapping, ipc_buffers_pong = IpcMemory(mapping, ipc_buffers_size, is_p2p_supported) ipc_barriers_in = IpcMemory( - mapping, IpcMemory.IPC_BARRIERS_SIZE_PER_GPU * mapping.tp_size * 2, - is_p2p_supported) + mapping, IpcMemory.IPC_BARRIERS_SIZE_PER_GPU * mapping.tp_size * 2 * + mapping.tp_size, is_p2p_supported) ipc_barriers_out = IpcMemory( - mapping, IpcMemory.IPC_BARRIERS_SIZE_PER_GPU * mapping.tp_size * 2, - is_p2p_supported) + mapping, IpcMemory.IPC_BARRIERS_SIZE_PER_GPU * mapping.tp_size * 2 * + mapping.tp_size, is_p2p_supported) lamport_buffers_size = 1 if force_deterministic else size * mapping.tp_size lamport_buffers_0 = IpcMemory(mapping, lamport_buffers_size, is_p2p_supported) @@ -762,16 +762,24 @@ def allocate_workspace(mapping: Mapping, lamport_buffers_size, ) buffers = [ - ipc_buffers_ping, ipc_buffers_pong, ipc_barriers_in, - ipc_barriers_out, lamport_buffers_0, lamport_buffers_1, - lamport_buffers_2 + ipc_buffers_ping, + ipc_buffers_pong, + ipc_barriers_in, + ipc_barriers_out, + lamport_buffers_0, + lamport_buffers_1, + lamport_buffers_2, + # Start from 1 since 0 represents released state for barrier at the beginning of the all_reduce. + # The last element is the barrier flag counter. + torch.tensor([1, 1, 0], dtype=torch.int64, device="cuda") ] return buffers, torch.tensor( ipc_buffers_ping.serialize() + ipc_buffers_pong.serialize() + ipc_barriers_in.serialize() + ipc_barriers_out.serialize() + lamport_buffers_0.serialize() + lamport_buffers_1.serialize() + - lamport_buffers_2.serialize() + [0] + [0], + lamport_buffers_2.serialize() + [buffers[-1].data_ptr()] + + [buffers[-1][1:].data_ptr()] + [buffers[-1][2:].data_ptr()], dtype=torch.int64, device="cpu") diff --git a/tests/integration/defs/test_e2e.py b/tests/integration/defs/test_e2e.py index daf1caed6b51..9404b8be3ed1 100644 --- a/tests/integration/defs/test_e2e.py +++ b/tests/integration/defs/test_e2e.py @@ -1884,6 +1884,7 @@ def test_ptp_quickstart_multimodal(llm_root, llm_venv, model_name, model_path, *accuracy_inputs[modality]["prompt"], "--media", *accuracy_inputs[modality]["media"], + "--disable_kv_cache_reuse", ] # NOTE # Qwen2-VL and Qwen2-5-VL model need larger max_num_tokens. @@ -1964,6 +1965,7 @@ def parse_output(text): functionality_inputs[modality]["prompt"], "--media", *functionality_inputs[modality]["media"], + "--disable_kv_cache_reuse", ], stdout=running_log)