diff --git a/cpp/tensorrt_llm/kernels/moeUtilOp.cu b/cpp/tensorrt_llm/kernels/moeUtilOp.cu new file mode 100644 index 000000000000..85cd9bda6753 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moeUtilOp.cu @@ -0,0 +1,879 @@ +/* + * Copyright (c) 2019-2023, 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 "cutlass_kernels/include/moe_kernels.h" +#include "tensorrt_llm/common/cudaTypeUtils.cuh" +#include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h" +#include "tensorrt_llm/kernels/moeUtilOp.h" +#include "tensorrt_llm/kernels/quantization.cuh" + +#include +#include + +#include // For INT_MAX +#include +#include +#include +#include // For numeric_limits +#include + +#include +#include +#include + +#ifndef CUDART_VERSION +#error CUDART_VERSION Undefined! +#elif (CUDART_VERSION >= 11050) +#include +#include +#include +#include +#include +#else +#include "3rdparty/cub/cub.cuh" +#include "3rdparty/cub/device/device_radix_sort.cuh" +#include "3rdparty/cub/util_type.cuh" +#endif + +namespace cg = cooperative_groups; +using namespace tensorrt_llm::common; + +namespace tensorrt_llm::kernels +{ + +template +__global__ void fusedBuildExpertMapsSortFirstTokenKernel(int const* const token_selected_experts, + int* const unpermuted_token_selected_experts, int* const permuted_source_token_ids, + int64_t* const expert_first_token_offset, int64_t const num_tokens, int const experts_per_token, + int const start_expert, int const end_expert, int const num_experts_per_node) +{ + // Only using block wise collective so we can only have one block + assert(gridDim.x == 1); + + assert(start_expert <= end_expert); + assert(num_experts_per_node == (end_expert - start_expert)); + assert(end_expert <= num_experts_per_node); + assert(num_experts_per_node <= (1 << LOG2_NUM_EXPERTS)); + + int const token = blockIdx.x * BLOCK_SIZE + threadIdx.x; + + bool is_valid_token = token < num_tokens; + + // This is the masked expert id for this token + int local_token_selected_experts[EXPERTS_PER_TOKEN]; + // This is the final permuted rank of this token (ranked by selected expert) + int local_token_permuted_indices[EXPERTS_PER_TOKEN]; + + // Wait PDL before reading token_selected_experts +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + +// build expert map +// we need to populate expert ids for all threads, even if there are +// fewer tokens +#pragma unroll + for (int i = 0; i < EXPERTS_PER_TOKEN; i++) + { + int const expert + = is_valid_token ? token_selected_experts[token * EXPERTS_PER_TOKEN + i] : num_experts_per_node; + + // If the token is not valid, set the expert id to num_experts_per_node + 1 + // If expert is not in the current node, set it to num_experts_per_node + // If expert is in the current node, subtract start_expert to shift the range to [0, num_experts_per_node) + bool is_valid_expert = expert >= start_expert && expert < end_expert; + local_token_selected_experts[i] = !is_valid_token ? num_experts_per_node + 1 + : is_valid_expert ? (expert - start_expert) + : num_experts_per_node; + } + + // TODO: decompose cub's sort to expose the bucket starts, and just return + // that to elide the binary search + + // sort the expert map + using BlockRadixRank = cub::BlockRadixRank; + extern __shared__ unsigned char temp_storage[]; + auto& sort_temp = *reinterpret_cast(temp_storage); + + // Sanity check that the number of bins do correspond to the number of experts + static_assert(BlockRadixRank::BINS_TRACKED_PER_THREAD * BLOCK_SIZE >= (1 << LOG2_NUM_EXPERTS)); + assert(BlockRadixRank::BINS_TRACKED_PER_THREAD * BLOCK_SIZE >= num_experts_per_node); + + int local_expert_first_token_offset[BlockRadixRank::BINS_TRACKED_PER_THREAD]; + + cub::BFEDigitExtractor extractor(0, LOG2_NUM_EXPERTS); + BlockRadixRank(sort_temp).RankKeys( + local_token_selected_experts, local_token_permuted_indices, extractor, local_expert_first_token_offset); + +// We are done with compute, launch the dependent kernels while the stores are in flight +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + + // write to shared memory and global memory + if (is_valid_token) + { +#pragma unroll + for (int i = 0; i < EXPERTS_PER_TOKEN; i++) + { + unpermuted_token_selected_experts[token * EXPERTS_PER_TOKEN + i] = local_token_selected_experts[i]; + permuted_source_token_ids[local_token_permuted_indices[i]] = i * num_tokens + token; + } + } + +#pragma unroll + for (int expert_id = 0; expert_id < BlockRadixRank::BINS_TRACKED_PER_THREAD; expert_id++) + { + int out_expert_id = expert_id + token * BlockRadixRank::BINS_TRACKED_PER_THREAD; + if (out_expert_id < num_experts_per_node + 1) + { + expert_first_token_offset[out_expert_id] = local_expert_first_token_offset[expert_id]; + } + } +} + +template +bool fusedBuildExpertMapsSortFirstTokenDispatch(int const* token_selected_experts, + int* unpermuted_token_selected_experts, int* permuted_source_token_ids, int64_t* expert_first_token_offset, + int64_t const num_tokens, int const num_experts_per_node, int const experts_per_token, int const start_expert, + int const end_expert, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(num_experts_per_node == (end_expert - start_expert), + "num_experts_per_node must be equal to end_expert - start_expert"); + int const threads = BLOCK_SIZE; + int const blocks = (num_tokens + threads - 1) / threads; + TLLM_CHECK_WITH_INFO(blocks == 1, "Current implementation requires single block"); + + using BlockRadixRank = cub::BlockRadixRank; + size_t shared_size = sizeof(typename BlockRadixRank::TempStorage); + + cudaLaunchConfig_t config; + config.gridDim = blocks; + config.blockDim = threads; + config.dynamicSmemBytes = shared_size; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.numAttrs = 1; + config.attrs = attrs; + + auto kernel = &fusedBuildExpertMapsSortFirstTokenKernel; + + int device = 0; + int max_smem_per_block = 0; + check_cuda_error(cudaGetDevice(&device)); + check_cuda_error(cudaDeviceGetAttribute(&max_smem_per_block, cudaDevAttrMaxSharedMemoryPerBlockOptin, device)); + if (shared_size >= static_cast(max_smem_per_block)) + { + // This should mean that + // cudaFuncSetAttribute(cutlass::Kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size) + // wouldn't work. + return false; + } + + check_cuda_error(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_size)); + check_cuda_error(cudaLaunchKernelEx(&config, kernel, token_selected_experts, unpermuted_token_selected_experts, + permuted_source_token_ids, expert_first_token_offset, num_tokens, experts_per_token, start_expert, end_expert, + num_experts_per_node)); + + return true; +} + +template +bool fusedBuildExpertMapsSortFirstTokenBlockSize(int const* token_selected_experts, + int* unpermuted_token_selected_experts, int* permuted_source_token_ids, int64_t* expert_first_token_offset, + int64_t const num_tokens, int const num_experts_per_node, int const experts_per_token, int const start_expert, + int const end_expert, cudaStream_t stream) +{ + int const block_size = num_tokens; + if (num_tokens > 256) + { + TLLM_LOG_TRACE( + "Number of tokens %d is greater than 256, which is not supported for fused moe prologues", num_tokens); + return false; + } + + auto func = &fusedBuildExpertMapsSortFirstTokenDispatch<32, EXPERTS_PER_TOKEN, LOG2_NUM_EXPERTS>; + if (block_size > 32 && block_size <= 64) + { + func = &fusedBuildExpertMapsSortFirstTokenDispatch<64, EXPERTS_PER_TOKEN, LOG2_NUM_EXPERTS>; + } + else if (block_size > 64 && block_size <= 128) + { + func = &fusedBuildExpertMapsSortFirstTokenDispatch<128, EXPERTS_PER_TOKEN, LOG2_NUM_EXPERTS>; + } + else if (block_size > 128 && block_size <= 256) + { + func = &fusedBuildExpertMapsSortFirstTokenDispatch<256, EXPERTS_PER_TOKEN, LOG2_NUM_EXPERTS>; + } + + return func(token_selected_experts, unpermuted_token_selected_experts, permuted_source_token_ids, + expert_first_token_offset, num_tokens, num_experts_per_node, experts_per_token, start_expert, end_expert, + stream); +} + +template +bool fusedBuildExpertMapsSortFirstTokenBlockSize(int const* token_selected_experts, + int* unpermuted_token_selected_experts, int* permuted_source_token_ids, int64_t* expert_first_token_offset, + int64_t const num_tokens, int const num_experts_per_node, int const experts_per_token, int const start_expert, + int const end_expert, cudaStream_t stream) +{ + auto func = &fusedBuildExpertMapsSortFirstTokenBlockSize<1, LOG2_NUM_EXPERTS>; + switch (experts_per_token) + { + case 1: + { + func = &fusedBuildExpertMapsSortFirstTokenBlockSize<1, LOG2_NUM_EXPERTS>; + break; + } + case 2: + { + func = &fusedBuildExpertMapsSortFirstTokenBlockSize<2, LOG2_NUM_EXPERTS>; + break; + } + case 4: + { + func = &fusedBuildExpertMapsSortFirstTokenBlockSize<4, LOG2_NUM_EXPERTS>; + break; + } + case 6: + { + func = &fusedBuildExpertMapsSortFirstTokenBlockSize<6, LOG2_NUM_EXPERTS>; + break; + } + case 8: + { + func = &fusedBuildExpertMapsSortFirstTokenBlockSize<8, LOG2_NUM_EXPERTS>; + break; + } + default: + { + TLLM_LOG_TRACE("Top-K value %d does not have supported fused moe prologues", experts_per_token); + return false; + } + } + return func(token_selected_experts, unpermuted_token_selected_experts, permuted_source_token_ids, + expert_first_token_offset, num_tokens, num_experts_per_node, experts_per_token, start_expert, end_expert, + stream); +} + +bool fusedBuildExpertMapsSortFirstToken(int const* token_selected_experts, int* unpermuted_token_selected_experts, + int* permuted_source_token_ids, int64_t* expert_first_token_offset, int64_t const num_tokens, + int const num_experts_per_node, int const experts_per_token, int const start_expert, int const end_expert, + cudaStream_t stream) +{ + // We need enough bits to represent [0, num_experts_per_node+1] (inclusive) i.e. num_experts_per_node + 2 values + // This is floor(log2(num_experts_per_node+1)) + 1 + int expert_log = static_cast(log2(num_experts_per_node + 1)) + 1; + if (expert_log <= 9) + { + auto funcs = std::array{&fusedBuildExpertMapsSortFirstTokenBlockSize<1>, + &fusedBuildExpertMapsSortFirstTokenBlockSize<2>, &fusedBuildExpertMapsSortFirstTokenBlockSize<3>, + &fusedBuildExpertMapsSortFirstTokenBlockSize<4>, &fusedBuildExpertMapsSortFirstTokenBlockSize<5>, + &fusedBuildExpertMapsSortFirstTokenBlockSize<6>, &fusedBuildExpertMapsSortFirstTokenBlockSize<7>, + &fusedBuildExpertMapsSortFirstTokenBlockSize<8>, &fusedBuildExpertMapsSortFirstTokenBlockSize<9>}; + + return funcs[expert_log - 1](token_selected_experts, unpermuted_token_selected_experts, + permuted_source_token_ids, expert_first_token_offset, num_tokens, num_experts_per_node, experts_per_token, + start_expert, end_expert, stream); + } + TLLM_LOG_TRACE("Experts per node %d does not have supported fused moe prologues", num_experts_per_node); + return false; +} + +// ============================== Infer GEMM sizes ================================= +// TODO Could linear search be better for small # experts +template +__device__ inline int64_t findTotalEltsLessThanTarget(T const* sorted_indices, int64_t const arr_length, T const target) +{ + int64_t low = 0, high = arr_length - 1, target_location = -1; + while (low <= high) + { + int64_t mid = (low + high) / 2; + + if (sorted_indices[mid] >= target) + { + high = mid - 1; + } + else + { + low = mid + 1; + target_location = mid; + } + } + return target_location + 1; +} + +// Calculates the start offset of the tokens for a given expert. The last element is the total number of valid tokens +__global__ void computeExpertFirstTokenOffsetKernel(int const* sorted_experts, int64_t const sorted_experts_len, + int64_t const num_experts_per_node, int64_t* expert_first_token_offset) +{ + // First, compute the global tid. We only need 1 thread per expert. + int const expert = blockIdx.x * blockDim.x + threadIdx.x; + + // Note that expert goes [0, num_experts] (inclusive) because we want a count for the total number of active tokens + // at the end of the scan. + if (expert >= num_experts_per_node + 1) + { + return; + } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + expert_first_token_offset[expert] = findTotalEltsLessThanTarget(sorted_experts, sorted_experts_len, expert); +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +void computeExpertFirstTokenOffset(int const* sorted_indices, int const total_indices, int const num_experts_per_node, + int64_t* expert_first_token_offset, cudaStream_t stream) +{ + int const num_entries = num_experts_per_node + 1; + int const threads = std::min(1024, num_entries); + int const blocks = (num_entries + threads - 1) / threads; + + cudaLaunchConfig_t config; + config.gridDim = blocks; + config.blockDim = threads; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.numAttrs = 1; + config.attrs = attrs; + cudaLaunchKernelEx(&config, computeExpertFirstTokenOffsetKernel, sorted_indices, total_indices, + num_experts_per_node, expert_first_token_offset); +} + +template +using sizeof_bits = cutlass::sizeof_bits>::type>; + +// Function to safely offset an pointer that may contain sub-byte types (FP4/INT4) +template +__host__ __device__ constexpr T* safe_inc_ptr(T* ptr, size_t offset) +{ + constexpr int adjustment = (sizeof_bits::value < 8) ? (8 / sizeof_bits::value) : 1; + assert(offset % adjustment == 0 && "Attempt to offset index to sub-byte"); + return ptr + offset / adjustment; +} + +__host__ __device__ constexpr int64_t getOffsetFlatSFArray(int64_t expert_id, int64_t gemm_n, int64_t gemm_k) +{ + auto min_alignment = cutlass_kernels::TmaWarpSpecializedGroupedGemmInput::MinNumRowsAlignmentFP4; + int64_t rounded_gemm_n = cute::ceil_div(gemm_n, min_alignment) * min_alignment; + assert(gemm_k % cutlass_kernels::TmaWarpSpecializedGroupedGemmInput::BlockScaleVectorSize == 0); + return expert_id * rounded_gemm_n * gemm_k + / cutlass_kernels::TmaWarpSpecializedGroupedGemmInput::BlockScaleVectorSize; +} + +__host__ __device__ constexpr int64_t getOffsetActivationSF(int64_t expert_id, int64_t token_offset, int64_t gemm_k) +{ + auto min_alignment = cutlass_kernels::TmaWarpSpecializedGroupedGemmInput::MinNumRowsAlignmentFP4; + // This formulation ensures that sf_offset[i + 1] - sf_offset[i] >= token_offset[i + 1] - token_offset[i]. + int64_t sf_offset = (token_offset + expert_id * (min_alignment - 1)) / min_alignment * min_alignment; + assert(gemm_k % cutlass_kernels::TmaWarpSpecializedGroupedGemmInput::BlockScaleVectorSize == 0); + return sf_offset * gemm_k / cutlass_kernels::TmaWarpSpecializedGroupedGemmInput::BlockScaleVectorSize; +} + +constexpr static int NVFP4_VEC_SIZE = 16; + +template +__device__ uint32_t quantizePackedFP4Value(ComputeElem& post_act_val, float global_scale_val, + int64_t num_tokens_before_expert, int64_t expert_id, int64_t token_id, int64_t elem_idx, int64_t num_cols, + int64_t max_tokens_per_expert, cutlass_kernels::TmaWarpSpecializedGroupedGemmInput::ElementSF* act_sf_flat) +{ + static constexpr int CVT_FP4_NUM_THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD; + // Quantize the input to FP4 + static_assert(std::is_same_v || std::is_same_v); + static_assert(ComputeElem::kElements == CVT_FP4_ELTS_PER_THREAD); + PackedVec packed_vec{}; + for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) + { + packed_vec.elts[i].x = static_cast(post_act_val[i * 2 + 0]); + packed_vec.elts[i].y = static_cast(post_act_val[i * 2 + 1]); + } + + // We need to offset into the scaling factors for just this expert + auto act_sf_expert = act_sf_flat + getOffsetFlatSFArray(expert_id, max_tokens_per_expert, num_cols); + + // Use `token - num_tokens_before_expert` because we want this to be relative to the start of this expert + // auto sf_out + // = cvt_quant_to_fp4_get_sf_out_offset( + // std::nullopt /* batchIdx */, token_id - num_tokens_before_expert, elem_idx, std::nullopt /* numRows */, + // num_cols, act_sf_expert, FP4QuantizationSFLayout::SWIZZLED); + auto sf_out = cvt_quant_to_fp4_get_sf_out_offset(std::nullopt /* batchIdx */, token_id - num_tokens_before_expert, + elem_idx, std::nullopt /* numRows */, num_cols, act_sf_expert, FP4QuantizationSFLayout::SWIZZLED); + + // Do the conversion and set the output and scaling factor + constexpr bool UE8M0 = false; + auto res = cvt_warp_fp16_to_fp4(packed_vec, global_scale_val, sf_out); + return res; +} + +__device__ void writeSF(int64_t num_tokens_before_expert, int64_t expert_id, int64_t source_token_id, int64_t token_id, + int64_t elem_idx, int64_t num_cols, int64_t max_tokens_per_expert, + cutlass_kernels::TmaWarpSpecializedGroupedGemmInput::ElementSF* act_sf_flat, + cutlass_kernels::TmaWarpSpecializedGroupedGemmInput::ElementSF const* input_sf) +{ + static constexpr int CVT_FP4_NUM_THREADS_PER_SF = NVFP4_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD; + + // We need to offset into the scaling factors for just this expert + auto act_sf_expert = act_sf_flat + getOffsetActivationSF(expert_id, num_tokens_before_expert, num_cols); + + // Use `token - num_tokens_before_expert` because we want this to be relative to the start of this expert + auto sf_out = cvt_quant_to_fp4_get_sf_out_offset(std::nullopt /* batchIdx */, token_id - num_tokens_before_expert, + elem_idx, std::nullopt /* numRows */, num_cols, act_sf_expert, FP4QuantizationSFLayout::SWIZZLED); + if (sf_out) + { + auto const sf_in + = cvt_quant_to_fp4_get_sf_out_offset(std::nullopt /* batchIdx */, source_token_id, elem_idx, + std::nullopt /* numRows */, num_cols, + const_cast(input_sf), + FP4QuantizationSFLayout::SWIZZLED); + *sf_out = *sf_in; + } +} + +void generateTokenPermutation(int const* unpermuted_token_selected_experts, int const* unpermuted_source_token_ids, + int* permuted_token_selected_experts, int* permuted_source_token_ids, int64_t* expert_first_token_offset, + int64_t num_rows, int64_t num_experts_per_node, int64_t k, cutlass_kernels::CubKeyValueSorter& sorter, + void* sorter_ws, cudaStream_t stream) +{ + int64_t const expanded_num_rows = k * num_rows; + sorter.updateNumExperts(num_experts_per_node); + size_t const sorter_ws_size_bytes + = cutlass_kernels::pad_to_multiple_of_16(sorter.getWorkspaceSize(expanded_num_rows, num_experts_per_node)); + sorter.run((void*) sorter_ws, sorter_ws_size_bytes, unpermuted_token_selected_experts, + permuted_token_selected_experts, unpermuted_source_token_ids, permuted_source_token_ids, expanded_num_rows, + stream); + + sync_check_cuda_error(stream); + + // Upper bound on number of expanded rows + computeExpertFirstTokenOffset( + permuted_token_selected_experts, expanded_num_rows, num_experts_per_node, expert_first_token_offset, stream); +} + +/** + * Takes the input maps and prepares the expanded maps for the sort step + * @param unpermuted_token_selected_experts: Buffer of transformed expert ids masked for the current node, used as the + * keys for the sort + * @param unpermuted_source_token_ids: Buffer of unpermuted token ids that will be used to identify the source row for + * each expanded token, used as the values for the sort + */ +__global__ void buildExpertMapsKernel(int const* token_selected_experts, int* unpermuted_token_selected_experts, + int* unpermuted_source_token_ids, int64_t const num_tokens, int const experts_per_token, int const start_expert, + int const end_expert, int const num_experts_per_node) +{ + int const token = blockIdx.x * blockDim.x + threadIdx.x; + if (token >= num_tokens) + { + return; + } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + + for (int i = 0; i < experts_per_token; i++) + { + int const expert = token_selected_experts[token * experts_per_token + i]; + // If expert is not in the current node, set it to num_experts_per_node + // If expert is in the current node, subtract start_expert to shift the range to [0, num_experts_per_node) + bool is_valid_expert = expert >= start_expert && expert < end_expert; + unpermuted_token_selected_experts[token * experts_per_token + i] + = is_valid_expert ? (expert - start_expert) : num_experts_per_node; + unpermuted_source_token_ids[token * experts_per_token + i] = i * num_tokens + token; + } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +void buildExpertMaps(int const* token_selected_experts, int* unpermuted_token_selected_experts, + int* unpermuted_source_token_ids, int64_t const num_tokens, int const num_experts_per_node, + int const experts_per_token, int const start_expert, int const end_expert, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(num_experts_per_node == (end_expert - start_expert), + "num_experts_per_node must be equal to end_expert - start_expert"); + int const threads = std::min(int64_t(1024), num_tokens); + int const blocks = (num_tokens + threads - 1) / threads; + + cudaLaunchConfig_t config; + config.gridDim = blocks; + config.blockDim = threads; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.numAttrs = 1; + config.attrs = attrs; + cudaLaunchKernelEx(&config, buildExpertMapsKernel, token_selected_experts, unpermuted_token_selected_experts, + unpermuted_source_token_ids, num_tokens, experts_per_token, start_expert, end_expert, num_experts_per_node); +} + +// ========================== Permutation things ======================================= +template +__host__ __device__ constexpr static U arrayConvert(T const& input) +{ + using Type = typename U::Element; + static_assert(T::kElements == U::kElements); + U u; +#pragma unroll + for (int i = 0; i < U::kElements; i++) + { + u[i] = static_cast(input[i]); + } + return u; +} + +// Duplicated and permutes rows for MoE. In addition, reverse the permutation map to help with finalizing routing. + +// "expanded_x_row" simply means that the number of values is num_rows x k. It is "expanded" since we will have to +// duplicate some rows in the input matrix to match the dimensions. Duplicates will always get routed to separate +// experts in the end. + +// Note that the expanded_dest_row_to_expanded_source_row map referred to here has indices in the range (0, +// k*rows_in_input - 1). However, it is set up so that index 0, rows_in_input, 2*rows_in_input ... (k-1)*rows_in_input +// all map to row 0 in the original matrix. Thus, to know where to read in the source matrix, we simply take the modulus +// of the expanded index. + +constexpr static int EXPAND_THREADS_PER_BLOCK = 256; + +template +__global__ void expandInputRowsKernel(InputActivationsType const* unpermuted_input, + ExpandedActivationsType* permuted_output, float const* unpermuted_scales, float* permuted_scales, + int const* expanded_dest_row_to_expanded_source_row, int* expanded_source_row_to_expanded_dest_row, + int64_t const num_rows, int64_t const* num_dest_rows, int64_t const cols, int64_t k, + float const* fc1_act_global_scale, int64_t* expert_first_token_offset, + cutlass_kernels::TmaWarpSpecializedGroupedGemmInput::ElementSF* fc1_act_sf_flat, + cutlass_kernels::TmaWarpSpecializedGroupedGemmInput::ElementSF const* input_sf, int64_t num_experts_per_node) +{ +#ifdef ENABLE_FP4 + constexpr bool is_fp4 = std::is_same_v; + constexpr bool is_fp4_input = is_fp4 && std::is_same_v; + constexpr bool need_fp4_quant = is_fp4 && !std::is_same_v; +#else + constexpr bool is_fp4 = false; + constexpr bool is_fp4_input = false; + constexpr bool need_fp4_quant = false; +#endif + + static_assert(need_fp4_quant || std::is_same_v, + "Only FP4 quantization supports outputting a different format as part of the expansion"); + + // Reverse permutation map. + // I do this so that later, we can use the source -> dest map to do the k-way reduction and unpermuting. I need the + // reverse map for that reduction to allow each threadblock to do 1 k-way reduce without atomics later in MoE. 1 + // thread block will be responsible for all k summations. + int64_t const expanded_dest_row = blockIdx.x; +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + int64_t const expanded_source_row = expanded_dest_row_to_expanded_source_row[expanded_dest_row]; + if (threadIdx.x == 0) + { + assert(expanded_dest_row <= INT32_MAX); + expanded_source_row_to_expanded_dest_row[expanded_source_row] = static_cast(expanded_dest_row); + } + + if (!CHECK_SKIPPED || blockIdx.x < *num_dest_rows) + { + // Load 128-bits per thread + constexpr int64_t ELEM_PER_THREAD + = is_fp4 ? CVT_FP4_ELTS_PER_THREAD : (128 / sizeof_bits::value); + constexpr int64_t ELEM_PER_BYTE = is_fp4_input ? 2 : 1; + using DataElem + = std::conditional_t>; + using OutputElem = std::conditional_t; + + // Duplicate and permute rows + int64_t const source_k_rank = expanded_source_row / num_rows; + int64_t const source_row = expanded_source_row % num_rows; + + auto const* source_row_ptr + = reinterpret_cast(unpermuted_input + source_row * cols / ELEM_PER_BYTE); + // Cast first to handle when this is FP4 + auto* dest_row_ptr + = reinterpret_cast(permuted_output) + expanded_dest_row * cols / ELEM_PER_THREAD; + + int64_t const start_offset = threadIdx.x; + int64_t const stride = EXPAND_THREADS_PER_BLOCK; + int64_t const num_elems_in_col = cols / ELEM_PER_THREAD; + assert(cols % ELEM_PER_THREAD == 0); + + if constexpr (is_fp4) + { + int64_t expert = findTotalEltsLessThanTarget( + expert_first_token_offset, num_experts_per_node, (int64_t) expanded_dest_row + 1) + - 1; + float global_scale_val = fc1_act_global_scale ? *fc1_act_global_scale : 1.0f; + int64_t num_tokens_before_expert = expert_first_token_offset[expert]; + + for (int elem_index = start_offset; elem_index < num_elems_in_col; elem_index += stride) + { + auto in_vec = source_row_ptr[elem_index]; + if constexpr (need_fp4_quant) + { + auto res = quantizePackedFP4Value(in_vec, global_scale_val, + num_tokens_before_expert, expert, expanded_dest_row, elem_index, cols, num_rows, + fc1_act_sf_flat); + dest_row_ptr[elem_index] = res; + } + else + { + writeSF(num_tokens_before_expert, expert, source_row, expanded_dest_row, elem_index, cols, num_rows, + fc1_act_sf_flat, input_sf); + dest_row_ptr[elem_index] = in_vec; + } + } + } + else + { + for (int elem_index = start_offset; elem_index < num_elems_in_col; elem_index += stride) + { + dest_row_ptr[elem_index] = source_row_ptr[elem_index]; + } + } + + if (permuted_scales && threadIdx.x == 0) + { + int64_t const source_k_idx = source_row * k + source_k_rank; + permuted_scales[expanded_dest_row] = unpermuted_scales ? unpermuted_scales[source_k_idx] : 1.0f; + } + } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +template +void expandInputRowsKernelLauncher(InputActivationsType const* unpermuted_input, + ExpandedActivationsType* permuted_output, float const* unpermuted_scales, float* permuted_scales, + int const* expanded_dest_row_to_expanded_source_row, int* expanded_source_row_to_expanded_dest_row, + int64_t const num_rows, int64_t const* num_valid_tokens_ptr, int64_t const cols, int const k, + int const num_experts_per_node, float const* fc1_act_global_scale, int64_t* expert_first_token_offset, + cutlass_kernels::TmaWarpSpecializedGroupedGemmInput::ElementSF* fc1_act_sf_flat, + cutlass_kernels::TmaWarpSpecializedGroupedGemmInput::ElementSF const* input_sf, cudaStream_t stream) +{ + if (fc1_act_sf_flat) + { + check_cuda_error( + cudaMemsetAsync(fc1_act_sf_flat, 0x0, getOffsetFlatSFArray(num_experts_per_node, num_rows, cols), stream)); + } + + int64_t const blocks = num_rows * k; + int64_t const threads = EXPAND_THREADS_PER_BLOCK; + auto func = (num_valid_tokens_ptr != nullptr) + ? expandInputRowsKernel + : expandInputRowsKernel; + + cudaLaunchConfig_t config; + config.gridDim = blocks; + config.blockDim = threads; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.numAttrs = 1; + config.attrs = attrs; + cudaLaunchKernelEx(&config, func, unpermuted_input, permuted_output, unpermuted_scales, permuted_scales, + expanded_dest_row_to_expanded_source_row, expanded_source_row_to_expanded_dest_row, num_rows, + num_valid_tokens_ptr, cols, k, fc1_act_global_scale, expert_first_token_offset, fc1_act_sf_flat, input_sf, + num_experts_per_node); +} + +#define INSTANTIATE_EXPAND_INPUT_ROWS(InputActivationsType, ExpandedActivationsType) \ + template void expandInputRowsKernelLauncher( \ + InputActivationsType const* unpermuted_input, ExpandedActivationsType* permuted_output, \ + float const* unpermuted_scales, float* permuted_scales, int const* expanded_dest_row_to_expanded_source_row, \ + int* expanded_source_row_to_expanded_dest_row, int64_t const num_rows, int64_t const* num_valid_tokens_ptr, \ + int64_t const cols, int const k, int const num_experts_per_node, float const* fc1_act_global_scale, \ + int64_t* expert_first_token_offset, \ + cutlass_kernels::TmaWarpSpecializedGroupedGemmInput::ElementSF* fc1_act_sf_flat, \ + cutlass_kernels::TmaWarpSpecializedGroupedGemmInput::ElementSF const* input_sf, cudaStream_t stream); + +INSTANTIATE_EXPAND_INPUT_ROWS(half, half); +INSTANTIATE_EXPAND_INPUT_ROWS(float, float); +#ifdef ENABLE_BF16 +INSTANTIATE_EXPAND_INPUT_ROWS(__nv_bfloat16, __nv_bfloat16); +#endif + +enum class ScaleMode : int +{ + NO_SCALE = 0, + DEFAULT = 1, +}; + +constexpr static int FINALIZE_THREADS_PER_BLOCK = 256; + +template +using sizeof_bits = cutlass::sizeof_bits>::type>; + +// Final kernel to unpermute and scale +// This kernel unpermutes the original data, does the k-way reduction and performs the final skip connection. +template +__global__ void finalizeMoeRoutingKernel(GemmOutputType const* expanded_permuted_rows, + OutputType* reduced_unpermuted_output, ScaleBiasType const* bias, float const* scales, + int const* expanded_source_row_to_expanded_dest_row, int const* expert_for_source_row, int64_t const orig_cols, + int64_t const experts_per_token, int64_t const* num_valid_ptr) +{ + assert(orig_cols % 4 == 0); + int64_t const original_row = blockIdx.x; + int64_t const num_rows = gridDim.x; + auto const offset = original_row * orig_cols; + OutputType* reduced_row_ptr = reduced_unpermuted_output + offset; + + // Load 128-bits per thread, according to the smallest data type we read/write + constexpr int64_t FINALIZE_ELEM_PER_THREAD + = 128 / std::min(sizeof_bits::value, sizeof_bits::value); + + int64_t const start_offset = threadIdx.x; + int64_t const stride = FINALIZE_THREADS_PER_BLOCK; + int64_t const num_elems_in_col = orig_cols / FINALIZE_ELEM_PER_THREAD; + + using BiasElem = cutlass::Array; + using InputElem = cutlass::Array; + using OutputElem = cutlass::Array; + using ComputeElem = cutlass::Array; + auto const* bias_v = reinterpret_cast(bias); + auto const* expanded_permuted_rows_v = reinterpret_cast(expanded_permuted_rows); + auto* reduced_row_ptr_v = reinterpret_cast(reduced_row_ptr); + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + int64_t const num_valid = *num_valid_ptr; + +#pragma unroll + for (int elem_index = start_offset; elem_index < num_elems_in_col; elem_index += stride) + { + bool has_valid = false; + ComputeElem thread_output; + thread_output.fill(0); + for (int k_idx = 0; k_idx < experts_per_token; ++k_idx) + { + int64_t const expanded_original_row = original_row + k_idx * num_rows; + int64_t const expanded_permuted_row = expanded_source_row_to_expanded_dest_row[expanded_original_row]; + + int64_t const k_offset = original_row * experts_per_token + k_idx; + float const row_scale = (SCALE_MODE == ScaleMode::NO_SCALE) ? 1.f : scales[k_offset]; + + // Check after row_rescale has accumulated + if (CHECK_SKIPPED && expanded_permuted_row >= num_valid) + { + continue; + } + + auto const* expanded_permuted_rows_row_ptr + = expanded_permuted_rows_v + expanded_permuted_row * num_elems_in_col; + + int64_t const expert_idx = expert_for_source_row[k_offset]; + + auto const* bias_ptr = bias_v + expert_idx * num_elems_in_col; + ComputeElem bias_value; + if (bias) + { + bias_value = arrayConvert(bias_ptr[elem_index]); + } + else + { + bias_value.fill(0); + } + + ComputeElem expert_result + = arrayConvert(expanded_permuted_rows_row_ptr[elem_index]); + thread_output = thread_output + row_scale * (expert_result + bias_value); + has_valid = true; + } + + OutputElem output_elem = arrayConvert(thread_output); + reduced_row_ptr_v[elem_index] = output_elem; + } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +template +void finalizeMoeRoutingKernelLauncher(GemmOutputType const* expanded_permuted_rows, + OutputType* reduced_unpermuted_output, ScaleBiasType const* bias, float const* final_scales, + int const* expanded_source_row_to_expanded_dest_row, int const* expert_for_source_row, int64_t const num_rows, + int64_t const cols, int64_t const experts_per_token, int64_t const* num_valid_ptr, + cutlass_kernels::MOEParallelismConfig parallelism_config, cudaStream_t stream) +{ + int64_t const blocks = num_rows; + int64_t const threads = FINALIZE_THREADS_PER_BLOCK; + + // Only add bias on rank 0 for tensor parallelism + bool const is_rank_0 = parallelism_config.tp_rank == 0; + ScaleBiasType const* bias_ptr = is_rank_0 ? bias : nullptr; + + bool const check_skipped = num_valid_ptr != nullptr; + + ScaleMode scale_mode = final_scales ? ScaleMode::DEFAULT : ScaleMode::NO_SCALE; + + using FuncPtr + = decltype(&finalizeMoeRoutingKernel); + FuncPtr func_map[2][3] = { + { + &finalizeMoeRoutingKernel, + &finalizeMoeRoutingKernel, + }, + { + &finalizeMoeRoutingKernel, + &finalizeMoeRoutingKernel, + }, + }; + auto* const func = func_map[check_skipped][int(scale_mode)]; + + cudaLaunchConfig_t config; + config.gridDim = blocks; + config.blockDim = threads; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.numAttrs = 1; + config.attrs = attrs; + cudaLaunchKernelEx(&config, func, expanded_permuted_rows, reduced_unpermuted_output, bias_ptr, final_scales, + expanded_source_row_to_expanded_dest_row, expert_for_source_row, cols, experts_per_token, num_valid_ptr); +} + +#define INSTANTIATE_FINALIZE_MOE_ROUTING(OutputT, GemmOutputT, ScaleBiasT) \ + template void finalizeMoeRoutingKernelLauncher( \ + GemmOutputT const* expanded_permuted_rows, OutputT* reduced_unpermuted_output, ScaleBiasT const* bias, \ + float const* final_scales, int const* expanded_source_row_to_expanded_dest_row, \ + int const* expert_for_source_row, int64_t const num_rows, int64_t const cols, int64_t const experts_per_token, \ + int64_t const* num_valid_ptr, cutlass_kernels::MOEParallelismConfig parallelism_config, cudaStream_t stream); + +INSTANTIATE_FINALIZE_MOE_ROUTING(half, half, half); +INSTANTIATE_FINALIZE_MOE_ROUTING(float, float, float); +#ifdef ENABLE_BF16 +INSTANTIATE_FINALIZE_MOE_ROUTING(__nv_bfloat16, __nv_bfloat16, __nv_bfloat16); +#endif + +} // namespace tensorrt_llm::kernels diff --git a/cpp/tensorrt_llm/kernels/moeUtilOp.h b/cpp/tensorrt_llm/kernels/moeUtilOp.h new file mode 100644 index 000000000000..968067d615ce --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moeUtilOp.h @@ -0,0 +1,56 @@ +/* + * Copyright (c) 2019-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "cutlass_kernels/include/moe_kernels.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include +#include + +namespace tensorrt_llm::kernels +{ +bool fusedBuildExpertMapsSortFirstToken(int const* token_selected_experts, int* unpermuted_token_selected_experts, + int* permuted_source_token_ids, int64_t* expert_first_token_offset, int64_t const num_tokens, + int const num_experts_per_node, int const experts_per_token, int const start_expert, int const end_expert, + cudaStream_t stream); + +void buildExpertMaps(int const* token_selected_experts, int* unpermuted_token_selected_experts, + int* unpermuted_source_token_ids, int64_t const num_tokens, int const num_experts_per_node, + int const experts_per_token, int const start_expert, int const end_expert, cudaStream_t stream); + +void generateTokenPermutation(int const* unpermuted_token_selected_experts, int const* unpermuted_source_token_ids, + int* permuted_token_selected_experts, int* permuted_source_token_ids, int64_t* expert_first_token_offset, + int64_t num_rows, int64_t num_experts_per_node, int64_t k, cutlass_kernels::CubKeyValueSorter& sorter, + void* sorter_ws, cudaStream_t stream); + +template +void expandInputRowsKernelLauncher(InputActivationsType const* unpermuted_input, + ExpandedActivationsType* permuted_output, float const* unpermuted_scales, float* permuted_scales, + int const* expanded_dest_row_to_expanded_source_row, int* expanded_source_row_to_expanded_dest_row, + int64_t const num_rows, int64_t const* num_valid_tokens_ptr, int64_t const cols, int const k, + int const num_experts_per_node, float const* fc1_act_global_scale, int64_t* expert_first_token_offset, + cutlass_kernels::TmaWarpSpecializedGroupedGemmInput::ElementSF* fc1_act_sf_flat, + cutlass_kernels::TmaWarpSpecializedGroupedGemmInput::ElementSF const* input_sf, cudaStream_t stream); + +template +void finalizeMoeRoutingKernelLauncher(GemmOutputType const* expanded_permuted_rows, + OutputType* reduced_unpermuted_output, ScaleBiasType const* bias, float const* final_scales, + int const* expanded_source_row_to_expanded_dest_row, int const* expert_for_source_row, int64_t const num_rows, + int64_t const cols, int64_t const experts_per_token, int64_t const* num_valid_ptr, + cutlass_kernels::MOEParallelismConfig parallelism_config, cudaStream_t stream); + +} // namespace tensorrt_llm::kernels diff --git a/cpp/tensorrt_llm/kernels/quantization.cuh b/cpp/tensorrt_llm/kernels/quantization.cuh index 1ab2e48a0fb5..cbb89579edaf 100644 --- a/cpp/tensorrt_llm/kernels/quantization.cuh +++ b/cpp/tensorrt_llm/kernels/quantization.cuh @@ -275,7 +275,7 @@ __global__ void perTokenQuantization(QuantT* dst, T const* src, int64_t const nu // FP4 Quantization constexpr int CVT_FP4_ELTS_PER_THREAD = 8; -// constexpr int CVT_FP4_SF_VEC_SIZE = 16; +constexpr int CVT_FP4_SF_VEC_SIZE = 16; constexpr int CVT_FP4_THREADS_PER_WARP = 32; constexpr int CVT_FP8_TO_FP4_ELTS_PER_THREAD = 16; diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 33bdcf88ec24..bf6b2de3753a 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -65,6 +65,7 @@ add_library( logitsBitmaskOp.cpp mambaConv1dOp.cpp moeOp.cpp + moeUtilOp.cpp moeCommOp.cpp moeLoadBalanceOp.cpp fp8BlockScaleMoe.cpp diff --git a/cpp/tensorrt_llm/thop/moeUtilOp.cpp b/cpp/tensorrt_llm/thop/moeUtilOp.cpp new file mode 100644 index 000000000000..990922fdb60b --- /dev/null +++ b/cpp/tensorrt_llm/thop/moeUtilOp.cpp @@ -0,0 +1,447 @@ +/* + * Copyright (c) 2022-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/moeUtilOp.h" +#include "moe_gemm_kernels.h" +#include "tensorrt_llm/common/workspace.h" +#include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h" +#include "tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h" +#include "tensorrt_llm/runtime/torchUtils.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include + +namespace th = torch; +namespace tl = tensorrt_llm; +namespace tk = tensorrt_llm::kernels; + +namespace common = tensorrt_llm::common; +namespace kernels = tensorrt_llm::kernels; +namespace cutlass_kernels = tensorrt_llm::kernels::cutlass_kernels; + +namespace torch_ext +{ + +// input_activations: [num_tokens, hidden_size] +// input: token_topk_unpermuted_scales, [num_tokens, k] +// output: permuted_data_, [num_token * k, hidden_size] +// output: permuted_token_final_scales_, [num_tokens, k] +template +void runPermute(void const* input_activations_void, void const* input_sf_void, int const* token_selected_experts, + float const* token_final_scales, void const* fc1_expert_weights_void, void const* fc1_expert_biases_void, + tensorrt_llm::ActivationType fc1_activation_type, void const* fc2_expert_weights_void, + void const* fc2_expert_biases_void, cutlass_kernels::QuantParams quant_params, int64_t const num_rows, + int64_t const hidden_size, int const full_num_experts, int const experts_per_token, + int* unpermuted_token_selected_experts_, int* unpermuted_source_token_ids_, int* permuted_source_token_ids_, + int* permuted_token_selected_experts_, T* permuted_data_, char* sorter_ws_, int64_t* expert_first_token_offset_, + float* permuted_token_final_scales_, int* expanded_source_row_to_expanded_dest_row, + cutlass_kernels::MOEParallelismConfig parallelism_config, cutlass_kernels::CubKeyValueSorter sorter_, bool use_lora, + kernels::LoraParams& lora_params, bool use_fp8_block_scaling, bool min_latency_mode, + cutlass_kernels::MoeMinLatencyParams& min_latency_params, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(experts_per_token * full_num_experts <= std::numeric_limits::max(), + "experts_per_token * num_experts is too large"); + + auto const* input_activations = static_cast(input_activations_void); + auto const* input_sf = input_sf_void + ? reinterpret_cast(input_sf_void) + : nullptr; + int const num_experts_per_node = full_num_experts / parallelism_config.ep_size; + int start_expert = num_experts_per_node * parallelism_config.ep_rank; + int end_expert = start_expert + num_experts_per_node; + + bool const needs_num_valid = parallelism_config.ep_size > 1; + // Note: expert_first_token_offset_[num_experts_per_node] stores the total number of expanded tokens + int64_t const* num_valid_tokens_ptr = needs_num_valid ? expert_first_token_offset_ + num_experts_per_node : nullptr; + + bool use_w4afp8 = false; + bool fused_prologue_result = false; + if (!use_w4afp8) + { + // WAR: fusedBuildExpertMapsSortFirstToken kernel will lead to illegal memory access for W4AFP8 + // input: token_selected_experts, [num_tokens, k] + // output: unpermuted_token_selected_experts_, [num_tokens, k] + // output: permuted_source_token_ids_, [num_tokens, k] + // output: expert_first_token_offset_, [num_experts_per_node + 1] + fused_prologue_result = kernels::fusedBuildExpertMapsSortFirstToken(token_selected_experts, + unpermuted_token_selected_experts_, permuted_source_token_ids_, expert_first_token_offset_, num_rows, + num_experts_per_node, experts_per_token, start_expert, end_expert, stream); + } + if (!fused_prologue_result) + { + TLLM_LOG_TRACE("Falling back to unfused prologue"); + kernels::buildExpertMaps(token_selected_experts, unpermuted_token_selected_experts_, + unpermuted_source_token_ids_, num_rows, num_experts_per_node, experts_per_token, start_expert, end_expert, + stream); + sync_check_cuda_error(stream); + + kernels::generateTokenPermutation(unpermuted_token_selected_experts_, unpermuted_source_token_ids_, + permuted_token_selected_experts_, permuted_source_token_ids_, expert_first_token_offset_, num_rows, + num_experts_per_node, experts_per_token, sorter_, static_cast(sorter_ws_), stream); + } + sync_check_cuda_error(stream); + + // using ExpandedActivationsType = std::conditional_t; + using ExpandedActivationsType = T; + // input_activations: [num_tokens, hidden_size] + // output: permuted_data_, [num_token * k, hidden_size] + // input: token_topk_unpermuted_scales, [num_tokens, k] + // output: permuted_token_final_scales_, [num_tokens * k] + // input: permuted_source_token_ids_, [num_tokens, k] + // output: expanded_source_row_to_expanded_dest_row, [num_tokens, k] + float const* token_topk_unpermuted_scales = token_final_scales; + kernels::expandInputRowsKernelLauncher(input_activations, + reinterpret_cast(permuted_data_), token_topk_unpermuted_scales, + permuted_token_final_scales_, permuted_source_token_ids_, expanded_source_row_to_expanded_dest_row, num_rows, + num_valid_tokens_ptr, hidden_size, experts_per_token, num_experts_per_node, + quant_params.fp4.fc1.act_global_scale, expert_first_token_offset_, + /* fc1_fp4_act_scale_ */ nullptr, input_sf, stream); + sync_check_cuda_error(stream); +} + +std::tuple +moe_permute_op(torch::Tensor const& input, torch::Tensor const& token_selected_experts, + torch::optional token_final_scales, torch::Tensor const& fc1_expert_weights, + torch::Tensor const& fc2_expert_weights, torch::optional> quant_scales, + torch::optional input_sf, int64_t const num_experts_on_rank, int64_t const tp_size, + int64_t const tp_rank, int64_t const ep_size, int64_t const ep_rank, int64_t const cluster_size, + int64_t const cluster_rank, bool min_latency_mode, bool use_fp8_block_scaling) +{ + cutlass_kernels::CubKeyValueSorter sorter_; + + TORCH_CHECK(cluster_size == 1 && cluster_rank == 0, "smart_router is supported in min_latency mode"); + TORCH_CHECK(min_latency_mode == false, "min_latency_mode is not supported now"); + + CHECK_INPUT(token_selected_experts, at::ScalarType::Int) + if (token_final_scales) + { + CHECK_INPUT(token_final_scales.value(), at::ScalarType::Float) + } + + TORCH_CHECK(input.dim() == 2, "input must be 2D."); + TORCH_CHECK(token_selected_experts.dim() == 2, "token_selected_experts must be 2D."); + + TORCH_CHECK(input.sizes()[0] == token_selected_experts.sizes()[0], + "input and token_selected_experts must have the same num tokens."); + if (token_final_scales) + { + TORCH_CHECK(token_final_scales.value().dim() == 2, "token_selected_experts_probs must be 2D."); + TORCH_CHECK(input.sizes()[0] == token_final_scales.value().sizes()[0], + "input and token_selected_experts_probs must have the same num tokens."); + TORCH_CHECK(token_selected_experts.sizes()[1] == token_final_scales.value().sizes()[1], + "token_selected_experts and token_final_scales must have the same number of experts per token."); + } + + int experts_per_token = token_selected_experts.sizes()[1]; + int64_t num_rows = input.sizes()[0]; + int64_t hidden_size = input.sizes()[1]; + auto const num_experts_total = static_cast(num_experts_on_rank * ep_size); + auto parallelism_config = cutlass_kernels::MOEParallelismConfig(tp_size, tp_rank, ep_size, ep_rank); + auto activation_type = tensorrt_llm::ActivationType::Swiglu; + + int const num_experts_per_node = num_experts_on_rank; + auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); + size_t num_moe_inputs = experts_per_token * num_rows; + + auto unpermuted_token_selected_experts_tensor + = torch::empty({num_moe_inputs}, torch::dtype(torch::kInt32).device(torch::kCUDA).requires_grad(false)); + + auto unpermuted_source_token_ids_tensor + = torch::empty({num_moe_inputs}, torch::dtype(torch::kInt32).device(torch::kCUDA).requires_grad(false)); + + auto permuted_source_token_ids_tensor + = torch::empty({num_moe_inputs}, torch::dtype(torch::kInt32).device(torch::kCUDA).requires_grad(false)); + + auto permuted_token_selected_experts_tensor + = torch::empty({num_moe_inputs}, torch::dtype(torch::kInt32).device(torch::kCUDA).requires_grad(false)); + + auto permuted_data_tensor = torch::empty({num_moe_inputs, hidden_size}, input.options().requires_grad(false)); + + auto permuted_token_final_scales_tensor + = torch::empty({num_moe_inputs}, torch::dtype(torch::kFloat32).device(torch::kCUDA).requires_grad(false)); + + auto expert_first_token_offset_tensor = torch::empty( + {num_experts_per_node + 1}, torch::dtype(torch::kInt64).device(torch::kCUDA).requires_grad(false)); + + size_t const sorter_size = min_latency_mode + ? 0 + : cutlass_kernels::CubKeyValueSorter::getWorkspaceSize(num_rows * experts_per_token, num_experts_per_node); + auto sorter_ws_tensor + = torch::empty({sorter_size}, torch::dtype(torch::kChar).device(torch::kCUDA).requires_grad(false)); + + auto src_to_dest_map_tensor = torch::empty( + {experts_per_token * num_rows}, torch::dtype(torch::kInt32).device(torch::kCUDA).requires_grad(false)); + + cutlass_kernels::QuantParams quant_params{}; + cutlass_kernels::MoeMinLatencyParams min_latency_params{}; + + kernels::LoraParams lora_params{}; + + auto data_type = input.scalar_type(); + switch (data_type) + { + case torch::kFloat32: + runPermute(input.const_data_ptr(), input_sf.has_value() ? input_sf.value().const_data_ptr() : nullptr, + reinterpret_cast(token_selected_experts.const_data_ptr()), + token_final_scales.has_value() ? reinterpret_cast(token_final_scales.value().const_data_ptr()) + : nullptr, + /*fc1_expert_weights.const_data_ptr()*/ nullptr, nullptr, activation_type, + /*fc2_expert_weights.const_data_ptr()*/ nullptr, nullptr, quant_params, num_rows, hidden_size, + num_experts_total, static_cast(experts_per_token), + static_cast(unpermuted_token_selected_experts_tensor.data_ptr()), + static_cast(unpermuted_source_token_ids_tensor.data_ptr()), + static_cast(permuted_source_token_ids_tensor.data_ptr()), + static_cast(permuted_token_selected_experts_tensor.data_ptr()), + static_cast(permuted_data_tensor.data_ptr()), static_cast(sorter_ws_tensor.data_ptr()), + static_cast(expert_first_token_offset_tensor.data_ptr()), + static_cast(permuted_token_final_scales_tensor.data_ptr()), + static_cast(src_to_dest_map_tensor.data_ptr()), parallelism_config, sorter_, false, lora_params, + use_fp8_block_scaling, min_latency_mode, min_latency_params, stream); + break; + case torch::kBFloat16: + runPermute<__nv_bfloat16>(input.const_data_ptr(), + input_sf.has_value() ? input_sf.value().const_data_ptr() : nullptr, + reinterpret_cast(token_selected_experts.const_data_ptr()), + token_final_scales.has_value() ? reinterpret_cast(token_final_scales.value().const_data_ptr()) + : nullptr, + /*fc1_expert_weights.const_data_ptr()*/ nullptr, nullptr, activation_type, + /*fc2_expert_weights.const_data_ptr()*/ nullptr, nullptr, quant_params, num_rows, hidden_size, + num_experts_total, static_cast(experts_per_token), + static_cast(unpermuted_token_selected_experts_tensor.data_ptr()), + static_cast(unpermuted_source_token_ids_tensor.data_ptr()), + static_cast(permuted_source_token_ids_tensor.data_ptr()), + static_cast(permuted_token_selected_experts_tensor.data_ptr()), + static_cast<__nv_bfloat16*>(permuted_data_tensor.data_ptr()), + static_cast(sorter_ws_tensor.data_ptr()), + static_cast(expert_first_token_offset_tensor.data_ptr()), + static_cast(permuted_token_final_scales_tensor.data_ptr()), + static_cast(src_to_dest_map_tensor.data_ptr()), parallelism_config, sorter_, false, lora_params, + use_fp8_block_scaling, min_latency_mode, min_latency_params, stream); + break; + case torch::kHalf: + runPermute(input.const_data_ptr(), input_sf.has_value() ? input_sf.value().const_data_ptr() : nullptr, + reinterpret_cast(token_selected_experts.const_data_ptr()), + token_final_scales.has_value() ? reinterpret_cast(token_final_scales.value().const_data_ptr()) + : nullptr, + /*fc1_expert_weights.const_data_ptr()*/ nullptr, nullptr, activation_type, + /*fc2_expert_weights.const_data_ptr()*/ nullptr, nullptr, quant_params, num_rows, hidden_size, + num_experts_total, static_cast(experts_per_token), + static_cast(unpermuted_token_selected_experts_tensor.data_ptr()), + static_cast(unpermuted_source_token_ids_tensor.data_ptr()), + static_cast(permuted_source_token_ids_tensor.data_ptr()), + static_cast(permuted_token_selected_experts_tensor.data_ptr()), + static_cast(permuted_data_tensor.data_ptr()), static_cast(sorter_ws_tensor.data_ptr()), + static_cast(expert_first_token_offset_tensor.data_ptr()), + static_cast(permuted_token_final_scales_tensor.data_ptr()), + static_cast(src_to_dest_map_tensor.data_ptr()), parallelism_config, sorter_, false, lora_params, + use_fp8_block_scaling, min_latency_mode, min_latency_params, stream); + break; + default: + throw std::invalid_argument( + "Invalid dtype, only supports input tensor with float32, float16 and bfloat16 dtype"); + break; + } + return std::make_tuple(unpermuted_token_selected_experts_tensor, unpermuted_source_token_ids_tensor, + permuted_source_token_ids_tensor, permuted_token_selected_experts_tensor, permuted_data_tensor, + expert_first_token_offset_tensor, permuted_token_final_scales_tensor, src_to_dest_map_tensor); +} + +std::tuple run_moe_expand_op(torch::Tensor const& input, + torch::optional token_final_scales, torch::Tensor const& permuted_source_token_ids, + int64_t const num_rows, torch::Tensor& expert_first_token_offset_tensor, int64_t const hidden_size, + int64_t const experts_per_token, int64_t const num_experts_per_node, int64_t const tp_size, int64_t const tp_rank, + int64_t const ep_size, int64_t const ep_rank, bool use_fp8_block_scaling) +{ + auto parallelism_config = cutlass_kernels::MOEParallelismConfig(tp_size, tp_rank, ep_size, ep_rank); + + bool const needs_num_valid = parallelism_config.ep_size > 1; + int64_t const* num_valid_tokens_ptr = needs_num_valid + ? static_cast(expert_first_token_offset_tensor.data_ptr()) + num_experts_per_node + : nullptr; + + size_t num_moe_inputs + = use_fp8_block_scaling ? (experts_per_token * num_rows + 3) / 4 * 4 : experts_per_token * num_rows; + auto permuted_data_tensor = torch::empty({num_moe_inputs, hidden_size}, input.options().requires_grad(false)); + auto permuted_token_final_scales_tensor + = torch::empty({num_moe_inputs}, torch::dtype(torch::kFloat32).device(torch::kCUDA).requires_grad(false)); + auto expanded_source_row_to_expanded_dest_row + = torch::empty({num_moe_inputs}, torch::dtype(torch::kInt32).device(torch::kCUDA).requires_grad(false)); + + auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); + cutlass_kernels::QuantParams quant_params{}; + + float const* token_topk_unpermuted_scales = token_final_scales.has_value() + ? reinterpret_cast(token_final_scales.value().const_data_ptr()) + : nullptr; + auto data_type = input.scalar_type(); + switch (data_type) + { + case torch::kFloat32: + kernels::expandInputRowsKernelLauncher(static_cast(input.const_data_ptr()), + reinterpret_cast(permuted_data_tensor.data_ptr()), token_topk_unpermuted_scales, + static_cast(permuted_token_final_scales_tensor.data_ptr()), + static_cast(permuted_source_token_ids.const_data_ptr()), + static_cast(expanded_source_row_to_expanded_dest_row.data_ptr()), num_rows, num_valid_tokens_ptr, + hidden_size, experts_per_token, num_experts_per_node, quant_params.fp4.fc1.act_global_scale, + static_cast(expert_first_token_offset_tensor.data_ptr()), + /* fc1_fp4_act_scale_ */ nullptr, /*input_sf*/ nullptr, stream); + break; + case torch::kBFloat16: + kernels::expandInputRowsKernelLauncher<__nv_bfloat16, __nv_bfloat16>( + static_cast<__nv_bfloat16 const*>(input.const_data_ptr()), + reinterpret_cast<__nv_bfloat16*>(permuted_data_tensor.data_ptr()), token_topk_unpermuted_scales, + static_cast(permuted_token_final_scales_tensor.data_ptr()), + static_cast(permuted_source_token_ids.const_data_ptr()), + static_cast(expanded_source_row_to_expanded_dest_row.data_ptr()), num_rows, num_valid_tokens_ptr, + hidden_size, experts_per_token, num_experts_per_node, quant_params.fp4.fc1.act_global_scale, + static_cast(expert_first_token_offset_tensor.data_ptr()), + /* fc1_fp4_act_scale_ */ nullptr, /*input_sf*/ nullptr, stream); + break; + case torch::kHalf: + kernels::expandInputRowsKernelLauncher(static_cast(input.const_data_ptr()), + reinterpret_cast(permuted_data_tensor.data_ptr()), token_topk_unpermuted_scales, + static_cast(permuted_token_final_scales_tensor.data_ptr()), + static_cast(permuted_source_token_ids.const_data_ptr()), + static_cast(expanded_source_row_to_expanded_dest_row.data_ptr()), num_rows, num_valid_tokens_ptr, + hidden_size, experts_per_token, num_experts_per_node, quant_params.fp4.fc1.act_global_scale, + static_cast(expert_first_token_offset_tensor.data_ptr()), + /* fc1_fp4_act_scale_ */ nullptr, /*input_sf*/ nullptr, stream); + break; + default: + throw std::invalid_argument( + "Invalid dtype, only supports input tensor with float32, float16 and bfloat16 dtype"); + break; + } + return std::make_tuple( + permuted_data_tensor, permuted_token_final_scales_tensor, expanded_source_row_to_expanded_dest_row); +} + +template +void runMoEFinalizeScaleOp(UnfusedGemmOutputType const* const gemm2_output, + ScaleBiasType const* const fc2_expert_biases, float const* const unpermuted_final_scales, + int const* const expanded_source_row_to_expanded_dest_row, int const* const expert_for_source_row, + int64_t const* const num_valid_tokens_ptr, int64_t const num_rows, /*int64_t const expanded_num_rows,*/ + int64_t const hidden_size, /*int64_t const inter_size, int const num_experts_per_node,*/ + int64_t const experts_per_token, cutlass_kernels::MOEParallelismConfig parallelism_config, cudaStream_t stream, + OutputType* const final_output) +{ + kernels::finalizeMoeRoutingKernelLauncher( + static_cast(gemm2_output), final_output, fc2_expert_biases, + unpermuted_final_scales, expanded_source_row_to_expanded_dest_row, expert_for_source_row, num_rows, hidden_size, + experts_per_token, num_valid_tokens_ptr, parallelism_config, stream); +} + +torch::Tensor run_moe_finalize_scale_op(torch::Tensor const& gemm2_output, torch::Tensor const& fc2_expert_biases, + torch::Tensor const& unpermuted_final_scales, torch::Tensor const& expanded_source_row_to_expanded_dest_row, + torch::Tensor const& expert_for_source_row, torch::Tensor const& expert_first_token_offset_tensor, + int64_t const num_rows, int64_t const hidden_size, int64_t const experts_per_token, + int64_t const num_experts_per_node, int64_t const tp_size, int64_t const tp_rank, int64_t const ep_size, + int64_t const ep_rank) +{ + TORCH_CHECK(gemm2_output.dim() == 2, "gemm2_output must be 2D."); + TORCH_CHECK(unpermuted_final_scales.dim() == 2, "unpermuted_final_scales must be 2D."); + TORCH_CHECK( + expanded_source_row_to_expanded_dest_row.dim() == 1, "expanded_source_row_to_expanded_dest_row must be 1D."); + TORCH_CHECK(expert_for_source_row.dim() == 1, "expert_for_source_row must be 1D."); + TORCH_CHECK(expert_first_token_offset_tensor.dim() == 1, "expert_first_token_offset_tensor must be 1D."); + + TORCH_CHECK(gemm2_output.sizes()[0] == expert_for_source_row.sizes()[0], + "gemm2_output and expert_for_source_row must have the same expanded num tokens."); + TORCH_CHECK(unpermuted_final_scales.sizes()[0] == num_rows, "unpermuted_final_scales[0] should equal to num_rows."); + TORCH_CHECK(unpermuted_final_scales.sizes()[1] == experts_per_token, + "unpermuted_final_scales[1] should equal to experts_per_token."); + TORCH_CHECK(expert_for_source_row.sizes()[0] == gemm2_output.sizes()[0], + "expert_for_source_row and gemm2_output must have the same expanded num tokens."); + TORCH_CHECK(expert_first_token_offset_tensor.sizes()[0] == num_experts_per_node + 1, + "expert_first_token_offset_tensor[0] should equal to num_experts_per_node + 1."); + + auto parallelism_config = cutlass_kernels::MOEParallelismConfig(tp_size, tp_rank, ep_size, ep_rank); + + bool const needs_num_valid = parallelism_config.ep_size > 1; + int64_t const* num_valid_tokens_ptr = needs_num_valid + ? static_cast(expert_first_token_offset_tensor.const_data_ptr()) + num_experts_per_node + : nullptr; + + auto final_output = torch::empty({num_rows, hidden_size}, gemm2_output.options()); + + auto stream = at::cuda::getCurrentCUDAStream(gemm2_output.get_device()); + auto data_type = gemm2_output.scalar_type(); + switch (data_type) + { + case torch::kFloat32: + runMoEFinalizeScaleOp(static_cast(gemm2_output.const_data_ptr()), + // static_cast(fc2_expert_biases.const_data_ptr()), + nullptr, static_cast(unpermuted_final_scales.const_data_ptr()), + static_cast(expanded_source_row_to_expanded_dest_row.const_data_ptr()), + static_cast(expert_for_source_row.const_data_ptr()), num_valid_tokens_ptr, num_rows, + hidden_size, experts_per_token, parallelism_config, stream, static_cast(final_output.data_ptr())); + break; + case torch::kBFloat16: + runMoEFinalizeScaleOp<__nv_bfloat16, __nv_bfloat16, __nv_bfloat16>( + static_cast<__nv_bfloat16 const*>(gemm2_output.const_data_ptr()), + // static_cast<__nv_bfloat16 const*>(fc2_expert_biases.const_data_ptr()), + nullptr, static_cast(unpermuted_final_scales.const_data_ptr()), + static_cast(expanded_source_row_to_expanded_dest_row.const_data_ptr()), + static_cast(expert_for_source_row.const_data_ptr()), num_valid_tokens_ptr, num_rows, + hidden_size, experts_per_token, parallelism_config, stream, + static_cast<__nv_bfloat16*>(final_output.data_ptr())); + break; + case torch::kHalf: + runMoEFinalizeScaleOp(static_cast(gemm2_output.const_data_ptr()), + // static_cast(fc2_expert_biases.const_data_ptr()), + nullptr, static_cast(unpermuted_final_scales.const_data_ptr()), + static_cast(expanded_source_row_to_expanded_dest_row.const_data_ptr()), + static_cast(expert_for_source_row.const_data_ptr()), num_valid_tokens_ptr, num_rows, + hidden_size, experts_per_token, parallelism_config, stream, static_cast(final_output.data_ptr())); + break; + default: + throw std::invalid_argument( + "Invalid dtype, only supports input tensor with float32, float16 and bfloat16 dtype"); + break; + } + return final_output; +} + +} // namespace torch_ext + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "moe_permute_op(Tensor input, Tensor token_selected_experts, Tensor? token_final_scales, Tensor " + "fc1_expert_weights, Tensor fc2_expert_weights, Tensor[]? quant_scales, Tensor? input_sf, int " + "num_experts_on_rank, int tp_size, int tp_rank, int ep_size, int ep_rank, int cluster_size, int cluster_rank, " + "bool min_latency_mode, bool use_fp8_block_scaling)" + "-> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)"); + m.def( + "moe_finalize_scale_op(Tensor gemm2_output, Tensor fc2_expert_biases, Tensor unpermuted_final_scales, Tensor " + "expanded_source_row_to_expanded_dest_row, Tensor expert_for_source_row, Tensor " + "expert_first_token_offset_tensor, int num_rows, int hidden_size, int experts_per_token, int " + "num_experts_per_node, int tp_size, int tp_rank, int ep_size, int ep_rank)" + "-> (Tensor)"); + m.def( + "moe_expand_op(Tensor input, Tensor? token_final_scales, Tensor permuted_source_token_ids, int num_rows, " + "Tensor expert_first_token_offset_tensor, int hidden_size, int experts_per_token, int num_experts_per_node, " + "int tp_size, int tp_rank, int ep_size, int ep_rank, bool use_fp8_block_scaling)" + "-> (Tensor, Tensor, Tensor)"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("moe_permute_op", &torch_ext::moe_permute_op); + m.impl("moe_finalize_scale_op", &torch_ext::run_moe_finalize_scale_op); + m.impl("moe_expand_op", &torch_ext::run_moe_expand_op); +} diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 292605c55ad3..a7ddf557a24b 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -24,6 +24,7 @@ def _( trigger_completion_at_end, ): from tensorrt_llm.functional import AllReduceFusionOp + if op == int(AllReduceFusionOp.NONE): return [torch.empty_like(input)] elif op == int(AllReduceFusionOp.RESIDUAL_RMS_NORM): @@ -55,7 +56,7 @@ def _( else: return [torch.empty_like(input)] - #MNNVL Allreduce + # MNNVL Allreduce @torch.library.register_fake("trtllm::mnnvl_twoshot_allreduce") def _(input, buffer, buffer_flags, wait_for_results): output = input.new_empty(input.shape) @@ -68,9 +69,18 @@ def _(comm_buf, gamma, eps, residual, buffer_flags): return [output, residual_out] @torch.library.register_fake("trtllm::moe_allreduce") - def _(residual, norm_weight, device_num_experts, scale_input, - active_experts_token_input, token_input, workspace, rank, nranks, - eps): + def _( + residual, + norm_weight, + device_num_experts, + scale_input, + active_experts_token_input, + token_input, + workspace, + rank, + nranks, + eps, + ): norm_out = torch.empty_like(token_input) residual_out = torch.empty_like(residual) return [norm_out, residual_out] @@ -175,8 +185,10 @@ def _( output_shape, scale_shape = fp4_utils.get_fp4_shape( input.shape, sf_vec_size) - return (input.new_empty(output_shape, dtype=torch.uint8), - global_scale.new_empty(scale_shape, dtype=torch.uint8)) + return ( + input.new_empty(output_shape, dtype=torch.uint8), + global_scale.new_empty(scale_shape, dtype=torch.uint8), + ) @torch.library.register_fake("trtllm::moe_comm_prepare_indices") def _( @@ -210,9 +222,14 @@ def _( backward_recv_rank_local_indices = gathered_target_rank_ids.new_empty( backward_recv_rank_local_indices_shape, dtype=torch.int32) - return (local_gather_indices, send_rank_count_cum_sum, - send_rank_local_indices, recv_rank_count_cum_sum, - recv_rank_local_indices, backward_recv_rank_local_indices) + return ( + local_gather_indices, + send_rank_count_cum_sum, + send_rank_local_indices, + recv_rank_count_cum_sum, + recv_rank_local_indices, + backward_recv_rank_local_indices, + ) @torch.library.register_fake("trtllm::moe_local_gather") def _( @@ -282,8 +299,11 @@ def _(global_expert_token_count: torch.Tensor, enabled: torch.Tensor, pass @torch.library.register_fake("trtllm::moe_load_balance_routing") - def _(single_layer_load_balancer_ptr: int, - token_selected_experts: torch.Tensor, offset_by_ep_rank: bool): + def _( + single_layer_load_balancer_ptr: int, + token_selected_experts: torch.Tensor, + offset_by_ep_rank: bool, + ): return torch.empty_like(token_selected_experts) @torch.library.custom_op("trtllm::group_rms_norm_base", @@ -352,9 +372,15 @@ def _( @torch.library.register_fake( "trtllm::mtp_sampling_and_accepted_draft_tokens_op") - def _(logits: torch.Tensor, draft_tokens: torch.Tensor, - target_tokens: torch.Tensor, num_mtp_modules: int, batch_size: int, - num_context_request: int, vocab_size: int): + def _( + logits: torch.Tensor, + draft_tokens: torch.Tensor, + target_tokens: torch.Tensor, + num_mtp_modules: int, + batch_size: int, + num_context_request: int, + vocab_size: int, + ): return logits.new_empty((batch_size, num_mtp_modules + 1), dtype=torch.int32), logits.new_empty( (batch_size, ), dtype=torch.int32) @@ -397,3 +423,82 @@ def _( pad_slot_id: int, ) -> None: pass + + @torch.library.register_fake("trtllm::moe_permute_op") + def _( + input: torch.Tensor, + token_selected_experts: torch.Tensor, + token_final_scales: torch.Tensor, + fc1_expert_weights: torch.Tensor, + fc2_expert_weights: torch.Tensor, + quant_scales: List[torch.Tensor], + input_sf: Optional[torch.Tensor], + num_experts_per_node: int, + tp_size: int, + tp_rank: int, + ep_size: int, + ep_rank: int, + cluster_size: int, + cluster_rank: int, + min_latency_mode: bool, + use_fp8_block_scaling: bool, + ): + + experts_per_token = token_selected_experts.shape[1] + num_rows = input.shape[0] + hidden_size = input.shape[1] + + num_moe_inputs = experts_per_token * num_rows + + unpermuted_token_selected_experts_tensor = token_selected_experts.new_empty( + (num_moe_inputs, ), dtype=torch.int32) + unpermuted_source_token_ids_tensor = token_selected_experts.new_empty( + (num_moe_inputs, ), dtype=torch.int32) + permuted_source_token_ids_tensor = token_selected_experts.new_empty( + (num_moe_inputs, ), dtype=torch.int32) + permuted_token_selected_experts_tensor = token_selected_experts.new_empty( + (num_moe_inputs, ), dtype=torch.int32) + permuted_data_tensor = input.new_empty((num_moe_inputs, hidden_size), + dtype=torch.float32) + expert_first_token_offset_tensor = token_selected_experts.new_empty( + (num_experts_per_node + 1, ), dtype=torch.int64) + permuted_token_final_scales_tensor = token_selected_experts.new_empty( + (num_moe_inputs, ), dtype=torch.float32) + src_to_dest_map_tensor = token_selected_experts.new_empty( + (num_moe_inputs, ), dtype=torch.int32) + + return ( + unpermuted_token_selected_experts_tensor, + unpermuted_source_token_ids_tensor, + permuted_source_token_ids_tensor, + permuted_token_selected_experts_tensor, + permuted_data_tensor, + expert_first_token_offset_tensor, + permuted_token_final_scales_tensor, + src_to_dest_map_tensor, + ) + + @torch.library.register_fake("trtllm::moe_finalize_scale_op") + def _( + gemm2_output: torch.Tensor, + fc2_expert_biases: torch.Tensor, + unpermuted_final_scales: torch.Tensor, + expanded_source_row_to_expanded_dest_row: torch.Tensor, + expert_for_source_row: torch.Tensor, + expert_first_token_offset_tensor: torch.Tensor, + num_rows: int, + hidden_size: int, + experts_per_token: int, + num_experts_per_node: int, + tp_size: int, + tp_rank: int, + ep_size: int, + ep_rank: int, + ): + + return gemm2_output.new_empty((num_rows, hidden_size), + dtype=gemm2_output.dtype) + + +def fp8_quantize_1x128(input: torch.Tensor): + return torch.ops.trtllm.fp8_quantize_1x128(input) diff --git a/tensorrt_llm/_torch/modules/fused_moe/__init__.py b/tensorrt_llm/_torch/modules/fused_moe/__init__.py index bb8f047fecf1..f7d7623af3d7 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/__init__.py +++ b/tensorrt_llm/_torch/modules/fused_moe/__init__.py @@ -1,4 +1,5 @@ from .create_moe import create_moe, get_moe_cls +from .fused_moe_cute_dsl import CuteDslFusedMoE from .fused_moe_cutlass import CutlassFusedMoE from .fused_moe_trtllm_gen import TRTLLMGenFusedMoE from .fused_moe_vanilla import VanillaMoE @@ -14,6 +15,10 @@ SparseMixerMoeRoutingMethod, StaticMoeRoutingMethod) __all__ = [ + "VanillaMoE", + "CutlassFusedMoE", + "CuteDslFusedMoE", + "TRTLLMGenFusedMoE", "BaseMoeRoutingMethod", "create_moe", "CutlassFusedMoE", diff --git a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py index 2d0c4c00c8b0..a7abe25ae13b 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py @@ -6,6 +6,7 @@ from tensorrt_llm.models.modeling_utils import QuantConfig from ...model_config import ModelConfig +from .fused_moe_cute_dsl import CuteDslFusedMoE from .fused_moe_cutlass import CutlassFusedMoE from .fused_moe_trtllm_gen import TRTLLMGenFusedMoE from .fused_moe_vanilla import VanillaMoE @@ -16,10 +17,11 @@ def get_moe_cls( - model_config: ModelConfig, - routing_method: BaseMoeRoutingMethod, - dtype: Optional[torch.dtype] = None, - override_quant_config: Optional[QuantConfig] = None) -> Type[MoE]: + model_config: ModelConfig, + routing_method: BaseMoeRoutingMethod, + dtype: Optional[torch.dtype] = None, + override_quant_config: Optional[QuantConfig] = None, +) -> Type[MoE]: moe_backend = model_config.moe_backend quant_config = model_config.quant_config if override_quant_config is not None: @@ -28,6 +30,8 @@ def get_moe_cls( return CutlassFusedMoE elif moe_backend.upper() == "VANILLA": return VanillaMoE + elif moe_backend.upper() == "CUTEDSL": + return CuteDslFusedMoE elif moe_backend.upper() == "TRTLLM": if quant_config is not None and ( quant_config.quant_mode.has_fp8_block_scales() @@ -67,7 +71,9 @@ def create_moe( assert moe_cls == WideEPMoE, "MoE Load Balance is only supported in WideEPMoE now." if moe_cls == TRTLLMGenFusedMoE: - assert not apply_router_weight_on_input, "apply_router_weight_on_input is not supported in TRTLLMGenFusedMoE." + assert ( + not apply_router_weight_on_input + ), "apply_router_weight_on_input is not supported in TRTLLMGenFusedMoE." return moe_cls( routing_method=routing_method, @@ -109,7 +115,9 @@ def create_moe( layer_idx=layer_idx, ) elif moe_cls == VanillaMoE: - assert not apply_router_weight_on_input, "apply_router_weight_on_input is not supported in VanillaMoE." + assert ( + not apply_router_weight_on_input + ), "apply_router_weight_on_input is not supported in VanillaMoE." return moe_cls( routing_method=routing_method, @@ -122,5 +130,19 @@ def create_moe( weight_loading_mode=weight_loading_mode, apply_router_weight_on_input=apply_router_weight_on_input, ) + elif moe_cls == CuteDslFusedMoE: + return moe_cls( + 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, + layer_idx=layer_idx, + ) else: raise ValueError(f"Unsupported moe backend: {moe_cls}") diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py new file mode 100644 index 000000000000..8154cafb4406 --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py @@ -0,0 +1,271 @@ +import math +from typing import List, Optional, Union + +import torch +import torch.nn.functional as F + +from ...distributed import allgather +from ...model_config import ModelConfig +from ...utils import Fp4QuantizedTensor, disable_fp4_allgather, reswizzle_sf +from .fused_moe_cutlass import CutlassFusedMoE +from .quantization import MoEWeightLoadingMode +from .routing import BaseMoeRoutingMethod + + +def swiglu_fused_moe(x): + x, gate = x.chunk(2, dim=-1) + return F.silu(gate) * x + + +def cute_dsl_fp8_group_blockwise_gemm_ref( + a: torch.Tensor, + b: torch.Tensor, + a_sf: torch.Tensor, + b_sf: torch.Tensor, + offset_array: torch.Tensor, +) -> torch.Tensor: + m, k = a.shape[0], a.shape[1] + l, n, k = b.shape[0], b.shape[1], b.shape[2] + num_group, w_n, w_k = b_sf.shape[0], b_sf.shape[1], b_sf.shape[2] + + # Note: view(int8) will cause error. + a_tmp = a.as_strided((m, k, 1), (k, 1, m * k)) + b_tmp = b.permute(1, 2, 0) + + m_padded = (m + 3) // 4 * 4 + input_scale_tmp = a_sf[0:m_padded * w_k] + input_scale_tmp = input_scale_tmp.reshape(-1, m_padded) + input_scale_tmp = input_scale_tmp[:w_k, :m].contiguous().permute(1, 0) + input_scale_tmp = input_scale_tmp.as_strided((m, w_k, 1), (1, m, m * w_k)) + + weight_scale_tmp = b_sf.permute(1, 2, 0) + + def pad_and_multiply(scale, tensor): + cm, ck, _ = scale.shape + m, k, _ = tensor.shape + IsGroupWise = False + IsBlockWise = False + if ck == math.ceil(k / 128): + IsGroupWise = True + if cm == math.ceil(m / 128): + IsBlockWise = True + if not IsBlockWise and not IsGroupWise: + raise ValueError("Only support granularity = 128") + + k_idx = torch.arange(k, device=scale.device) + if IsGroupWise: + k_idx = k_idx // 128 + m_idx = torch.arange(m, device=scale.device) + if IsBlockWise: + m_idx = m_idx // 128 + expanded_scale = scale[m_idx[:, None], k_idx, :] + + result = expanded_scale * tensor + + return result + + updated_a = pad_and_multiply(input_scale_tmp, a_tmp.to(torch.float32)) + updated_b = pad_and_multiply(weight_scale_tmp, b_tmp.to(torch.float32)) + + ref = torch.zeros((m, n), device="cuda", dtype=torch.float32) + + len_offset_array = offset_array.shape[0] + for i in range(len_offset_array - 1): + start = offset_array[i] + end = offset_array[i + 1] + # assert start <= end, f"Invalid group boundaries: start={start} > end={end}" + ref[start:end, :] = torch.einsum("mk,nk->mn", updated_a[start:end, :, + 0], + updated_b[:, :, i]) + ref = ref.to(torch.bfloat16) + return ref + + +class CuteDslFusedMoE(CutlassFusedMoE): + """ + Python Flow of Fused Mixture of Experts (MoE) Layer. + + Args: + num_experts (int): Number of experts in the MoE layer. + top_k (int): Number of top experts to select for each input token. + hidden_size (int): Size of the hidden state. + intermediate_size (int): Size of the intermediate state. + aux_stream (Optional[torch.cuda.Stream]): Auxiliary CUDA stream to overlap chunks. + dtype (Optional[torch.dtype]): Data type for the weights. + reduce_results (bool): Whether to reduce the results across devices. + model_config (ModelConfig): Configuration object for the model. + + This backend is composed of multiple custom ops: + 1. moe_permute_op: permute the input tensor and the expert selected tensor. + 2. cute_dsl_fp8_group_blockwise_gemm_ref: a reference implementation of the cute_dsl_fp8_group_blockwise_gemm. + 3. moe_finalize_scale_op: finalize the scale of the output tensor. + """ + + 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: Optional[torch.cuda.Stream] = None, + weight_loading_mode: MoEWeightLoadingMode = MoEWeightLoadingMode. + VANILLA, + apply_router_weight_on_input: bool = False, + layer_idx: Optional[int] = 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, + layer_idx=layer_idx, + ) + + def forward_chunk( + self, + x: Union[torch.Tensor, Fp4QuantizedTensor], + router_logits: torch.Tensor, + output_dtype: Optional[torch.dtype] = None, + all_rank_num_tokens: Optional[List[int]] = None, + use_dp_padding: Optional[bool] = None, + ) -> torch.Tensor: + if isinstance(x, Fp4QuantizedTensor): + assert output_dtype is not None + output_dtype = output_dtype + else: + output_dtype = x.dtype + + # apply routing + token_selected_experts, token_final_scales = self.routing_method.apply( + router_logits) + assert token_selected_experts.shape[ + 1] == self.routing_method.experts_per_token + assert token_selected_experts.shape == token_final_scales.shape + assert token_selected_experts.shape[0] == router_logits.shape[0] + assert token_final_scales.dtype == torch.float32 + assert token_selected_experts.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 + + # quantize inputs + use_deepseek_fp8_block_scale = False + weight_dtype = self.w3_w1_weight.dtype + x_sf = None + if self.has_any_quant: + if self.has_fp8_qdq: + x, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor( + x, self.fc31_input_dequant) + elif self.has_deepseek_fp8_block_scales: + use_deepseek_fp8_block_scale = True + elif self.has_w4afp8: + weight_dtype = torch.quint4x2 + elif self.has_nvfp4 and not disable_fp4_allgather(): + if isinstance(x, Fp4QuantizedTensor): + x_row = x.shape[0] + # note: we use uint8 to store 2 fp4 values + x_col = x.shape[1] * 2 + x, x_sf = x.fp4_tensor, x.scaling_factor + else: + x_row = x.shape[0] + x_col = x.shape[1] + x, x_sf = torch.ops.trtllm.fp4_quantize( + x, self.fc31_input_scale, self.scaling_vector_size, + False) + else: + raise ValueError( + f"unsupported quantization mode: {self.quant_config.quant_mode}" + ) + + # gather inputs for attention dp + if self.use_dp and self.parallel_size > 1 and not disable_fp4_allgather( + ): + x, x_sf, token_selected_experts, token_final_scales = allgather( + [x, x_sf, token_selected_experts, token_final_scales], + self.mapping, + dim=0, + sizes=None if use_dp_padding else all_rank_num_tokens) + # Fp4 gemm has extra scaling factor + if x_sf is not None: + x_sf = reswizzle_sf(x_sf, x_row, x_col, + self.scaling_vector_size) + + ( + unpermuted_token_selected_experts_tensor, + unpermuted_source_token_ids_tensor, + permuted_source_token_ids_tensor, + permuted_token_selected_experts_tensor, + permuted_data_tensor, + expert_first_token_offset_tensor, + permuted_token_final_scales_tensor, + src_to_dest_map_tensor, + ) = torch.ops.trtllm.moe_permute_op( + x, + token_selected_experts, + token_final_scales, + None, # w3_w1_weight.view(weight_dtype), + None, # w2_weight.view(weight_dtype), + None, # quant_scales, + input_sf=x_sf, + num_experts_on_rank=self.expert_size_per_partition, + tp_size=self.tp_size, + tp_rank=self.tp_rank, + ep_size=self.ep_size, + ep_rank=self.ep_rank, + cluster_size=self.cluster_size, + cluster_rank=self.cluster_rank, + min_latency_mode=False, + use_fp8_block_scaling=use_deepseek_fp8_block_scale, + ) + + act_input_fp8, act_input_sf = torch.ops.trtllm.fp8_quantize_1x128( + permuted_data_tensor) + h1 = cute_dsl_fp8_group_blockwise_gemm_ref( + a=act_input_fp8, + b=self.w3_w1_weight.view(weight_dtype), + a_sf=act_input_sf, + b_sf=self.quant_scales[0], + offset_array=expert_first_token_offset_tensor, + ) + h2 = swiglu_fused_moe(h1) + act_input_fp8, act_input_sf = torch.ops.trtllm.fp8_quantize_1x128(h2) + h3 = cute_dsl_fp8_group_blockwise_gemm_ref( + a=act_input_fp8, + b=self.w2_weight.view(weight_dtype), + a_sf=act_input_sf, + b_sf=self.quant_scales[1], + offset_array=expert_first_token_offset_tensor, + ) + final_hidden_states = torch.ops.trtllm.moe_finalize_scale_op( + h3, + None, + token_final_scales, + src_to_dest_map_tensor, + unpermuted_token_selected_experts_tensor, + expert_first_token_offset_tensor, + x.shape[0], # num_rows + x.shape[1], # hidden_size + self.routing_method.top_k, + self.expert_size_per_partition, # num_experts_per_node + self.tp_size, + self.tp_rank, + self.ep_size, + self.ep_rank, + ) + + return final_hidden_states diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 391d9a7dd199..026dead9ef00 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -73,11 +73,14 @@ def test_chunked_prefill(self, attn_backend): pytorch_config = dict( attn_backend=attn_backend, # https://nvbugspro.nvidia.com/bug/5345391 - disable_overlap_scheduler=True) - llm = LLM(self.MODEL_PATH, - enable_chunked_prefill=True, - max_num_tokens=512, - **pytorch_config) + disable_overlap_scheduler=True, + ) + llm = LLM( + self.MODEL_PATH, + enable_chunked_prefill=True, + max_num_tokens=512, + **pytorch_config, + ) with llm: task = MMLU(self.MODEL_NAME) task.evaluate(llm) @@ -88,8 +91,8 @@ def test_chunked_prefill(self, attn_backend): [False, pytest.param(True, marks=skip_device_contain_gb200)]) @parametrize_with_ids("attn_backend", ["TRTLLM", "FLASHINFER"]) def test_bfloat16(self, attn_backend, torch_compile): - torch_compile_config = TorchCompileConfig( - enable_fullgraph=True) if torch_compile else None + torch_compile_config = (TorchCompileConfig( + enable_fullgraph=True) if torch_compile else None) pytorch_config = dict( torch_compile_config=torch_compile_config, cuda_graph_padding_enabled=torch_compile, @@ -117,8 +120,8 @@ def test_bfloat16_4gpus(self, tp_size, pp_size, attn_backend, "Pipeline parallel with torch.compile is not supported yet.\n" "Issue: Unfusing flashinfer_fused_add_rmsnorm causes outputs to be " "discarded at graph breaks.") - torch_compile_config = TorchCompileConfig( - enable_fullgraph=True) if torch_compile else None + torch_compile_config = (TorchCompileConfig( + enable_fullgraph=True) if torch_compile else None) pytorch_config = dict( torch_compile_config=torch_compile_config, cuda_graph_padding_enabled=torch_compile, @@ -126,10 +129,12 @@ def test_bfloat16_4gpus(self, tp_size, pp_size, attn_backend, attn_backend=attn_backend, disable_overlap_scheduler=torch_compile, ) - llm = LLM(self.MODEL_PATH, - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - **pytorch_config) + llm = LLM( + self.MODEL_PATH, + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + **pytorch_config, + ) with llm: task = MMLU(self.MODEL_NAME) task.evaluate(llm) @@ -144,8 +149,8 @@ def test_bfloat16_4gpus(self, tp_size, pp_size, attn_backend, @parametrize_with_ids("fp8kv", [False, True]) def test_fp8(self, fp8kv, attn_backend, torch_compile): quant_config = QuantConfig(QuantAlgo.FP8) - torch_compile_config = TorchCompileConfig( - enable_fullgraph=True) if torch_compile else None + torch_compile_config = (TorchCompileConfig( + enable_fullgraph=True) if torch_compile else None) pytorch_config = dict( torch_compile_config=torch_compile_config, cuda_graph_padding_enabled=torch_compile, @@ -159,7 +164,8 @@ def test_fp8(self, fp8kv, attn_backend, torch_compile): llm = LLM( f"{llm_models_root()}/llama-3.1-model/Llama-3.1-8B-Instruct-FP8", quant_config=quant_config, - **pytorch_config) + **pytorch_config, + ) assert llm.args.quant_config.quant_algo == QuantAlgo.FP8 if fp8kv: assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 @@ -185,8 +191,8 @@ def test_fp8_4gpus(self, tp_size, pp_size, fp8kv, attn_backend, "Issue: Unfusing flashinfer_fused_add_rmsnorm causes outputs to be " "discarded at graph breaks.") quant_config = QuantConfig(QuantAlgo.FP8) - torch_compile_config = TorchCompileConfig( - enable_fullgraph=True) if torch_compile else None + torch_compile_config = (TorchCompileConfig( + enable_fullgraph=True) if torch_compile else None) pytorch_config = dict( torch_compile_config=torch_compile_config, cuda_graph_padding_enabled=torch_compile, @@ -202,7 +208,8 @@ def test_fp8_4gpus(self, tp_size, pp_size, fp8kv, attn_backend, tensor_parallel_size=tp_size, pipeline_parallel_size=pp_size, quant_config=quant_config, - **pytorch_config) + **pytorch_config, + ) assert llm.args.quant_config.quant_algo == QuantAlgo.FP8 if fp8kv: assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 @@ -225,9 +232,11 @@ def test_fp8_llm_sampler(self): with llm: task = MMLU(self.MODEL_NAME) - task.evaluate(llm, - sampling_params=sampling_params, - extra_acc_spec="temperature=0.8,top_p=0.95") + task.evaluate( + llm, + sampling_params=sampling_params, + extra_acc_spec="temperature=0.8,top_p=0.95", + ) def test_eagle3(self): pytorch_config = dict( @@ -244,11 +253,13 @@ def test_eagle3(self): spec_config = EagleDecodingConfig(max_draft_len=draft_len, pytorch_weights_path=eagle_model_dir) - llm = LLM(model=target_model_dir, - **pytorch_config, - kv_cache_config=kv_cache_config, - speculative_config=spec_config, - build_config=None) + llm = LLM( + model=target_model_dir, + **pytorch_config, + kv_cache_config=kv_cache_config, + speculative_config=spec_config, + build_config=None, + ) with llm: task = MMLU(self.MODEL_NAME) @@ -268,32 +279,38 @@ def test_ngram(self): is_public_pool=True, ) - llm = LLM(model=self.MODEL_PATH, - **pytorch_config, - kv_cache_config=kv_cache_config, - speculative_config=spec_config) + llm = LLM( + model=self.MODEL_PATH, + **pytorch_config, + kv_cache_config=kv_cache_config, + speculative_config=spec_config, + ) with llm: task = MMLU(self.MODEL_NAME) task.evaluate(llm) def test_guided_decoding(self): - llm = LLM(self.MODEL_PATH, - guided_decoding_backend="xgrammar", - disable_overlap_scheduler=True, - use_cuda_graph=True) + llm = LLM( + self.MODEL_PATH, + guided_decoding_backend="xgrammar", + disable_overlap_scheduler=True, + use_cuda_graph=True, + ) with llm: task = JsonModeEval(self.MODEL_NAME) task.evaluate(llm) @pytest.mark.skip_less_device(4) def test_guided_decoding_4gpus(self): - llm = LLM(self.MODEL_PATH, - guided_decoding_backend="xgrammar", - disable_overlap_scheduler=True, - use_cuda_graph=True, - tensor_parallel_size=2, - pipeline_parallel_size=2) + llm = LLM( + self.MODEL_PATH, + guided_decoding_backend="xgrammar", + disable_overlap_scheduler=True, + use_cuda_graph=True, + tensor_parallel_size=2, + pipeline_parallel_size=2, + ) with llm: task = JsonModeEval(self.MODEL_NAME) task.evaluate(llm) @@ -360,7 +377,9 @@ def test_auto_dtype_tp8(self): @pytest.mark.skip_less_device(4) @pytest.mark.skip_device_not_contain(["H100", "H200", "B200"]) def test_fp8_tp4(self): - model_path = f"{llm_models_root()}/modelopt-hf-model-hub/Llama-3.3-70B-Instruct-fp8" + model_path = ( + f"{llm_models_root()}/modelopt-hf-model-hub/Llama-3.3-70B-Instruct-fp8" + ) with LLM(model_path, tensor_parallel_size=4) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.FP8 task = MMLU(self.MODEL_NAME) @@ -374,7 +393,9 @@ def test_fp8_tp4(self): @pytest.mark.skip_less_device(4) @pytest.mark.skip_device_not_contain(["B200"]) def test_nvfp4_tp4(self): - model_path = f"{llm_models_root()}/modelopt-hf-model-hub/Llama-3.3-70B-Instruct-fp4" + model_path = ( + f"{llm_models_root()}/modelopt-hf-model-hub/Llama-3.3-70B-Instruct-fp4" + ) with LLM(model_path, tensor_parallel_size=4) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 @@ -394,15 +415,19 @@ class TestLlama4MaverickInstruct(LlmapiAccuracyTestHarness): @skip_pre_blackwell @pytest.mark.skip_less_mpi_world_size(8) @parametrize_with_ids("cuda_graph", [False, True]) - @pytest.mark.parametrize("tp_size,pp_size,ep_size", [(8, 1, 1), (8, 1, 4), - (8, 1, 8)], - ids=["tp8", "tp8ep4", "tp8ep8"]) + @pytest.mark.parametrize( + "tp_size,pp_size,ep_size", + [(8, 1, 1), (8, 1, 4), (8, 1, 8)], + ids=["tp8", "tp8ep4", "tp8ep8"], + ) def test_auto_dtype(self, cuda_graph, tp_size, pp_size, ep_size): - with LLM(self.MODEL_PATH, - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - moe_expert_parallel_size=ep_size, - use_cuda_graph=cuda_graph) as llm: + with LLM( + self.MODEL_PATH, + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + moe_expert_parallel_size=ep_size, + use_cuda_graph=cuda_graph, + ) as llm: task = MMLU(self.MODEL_NAME) task.evaluate(llm) task = GSM8K(self.MODEL_NAME) @@ -416,15 +441,19 @@ class TestLlama4ScoutInstruct(LlmapiAccuracyTestHarness): @skip_pre_hopper @pytest.mark.skip_less_mpi_world_size(8) @parametrize_with_ids("cuda_graph", [False, True]) - @pytest.mark.parametrize("tp_size,pp_size,ep_size", [(8, 1, 1), (8, 1, 4), - (8, 1, 8)], - ids=["tp8", "tp8ep4", "tp8ep8"]) + @pytest.mark.parametrize( + "tp_size,pp_size,ep_size", + [(8, 1, 1), (8, 1, 4), (8, 1, 8)], + ids=["tp8", "tp8ep4", "tp8ep8"], + ) def test_auto_dtype(self, cuda_graph, tp_size, pp_size, ep_size): - with LLM(self.MODEL_PATH, - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - moe_expert_parallel_size=ep_size, - use_cuda_graph=cuda_graph) as llm: + with LLM( + self.MODEL_PATH, + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + moe_expert_parallel_size=ep_size, + use_cuda_graph=cuda_graph, + ) as llm: task = MMLU(self.MODEL_NAME) task.evaluate(llm) task = GSM8K(self.MODEL_NAME) @@ -466,7 +495,9 @@ def test_tp2(self): @pytest.mark.skip_less_device(2) @pytest.mark.skip_device_not_contain(["H100", "H200", "B200"]) def test_fp8_tp2(self): - model_path = f"{llm_models_root()}/modelopt-hf-model-hub/Mixtral-8x7B-Instruct-v0.1-fp8" + model_path = ( + f"{llm_models_root()}/modelopt-hf-model-hub/Mixtral-8x7B-Instruct-v0.1-fp8" + ) with LLM(model_path, tensor_parallel_size=2) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.FP8 assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 @@ -478,7 +509,9 @@ def test_fp8_tp2(self): @pytest.mark.skip_less_device(2) @pytest.mark.skip_device_not_contain(["B200"]) def test_nvfp4_tp2(self): - model_path = f"{llm_models_root()}/modelopt-hf-model-hub/Mixtral-8x7B-Instruct-v0.1-fp4" + model_path = ( + f"{llm_models_root()}/modelopt-hf-model-hub/Mixtral-8x7B-Instruct-v0.1-fp4" + ) with LLM(model_path, tensor_parallel_size=2) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 @@ -499,10 +532,17 @@ class TestDeepSeekV3Lite(LlmapiAccuracyTestHarness): @parametrize_with_ids( "torch_compile", [False, pytest.param(True, marks=skip_device_contain_gb200)]) - @parametrize_with_ids("attention_dp,cuda_graph,overlap_scheduler", - [(False, False, False), (True, False, False), - (False, True, False), (False, False, True), - (False, True, True), (True, True, True)]) + @parametrize_with_ids( + "attention_dp,cuda_graph,overlap_scheduler", + [ + (False, False, False), + (True, False, False), + (False, True, False), + (False, False, True), + (False, True, True), + (True, True, True), + ], + ) # Only Hopper and Blackwell MLA kernel supports MTP @parametrize_with_ids("mtp_nextn", [0, pytest.param(2, marks=skip_pre_hopper)]) @@ -513,9 +553,9 @@ def test_bfloat16(self, mtp_nextn, attention_dp, cuda_graph, if torch_compile and attention_dp: pytest.skip("https://nvbugs/5252559") kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.9) - torch_compile_config = TorchCompileConfig( - enable_fullgraph=True, - enable_piecewise_cuda_graph=cuda_graph) if torch_compile else None + torch_compile_config = (TorchCompileConfig( + enable_fullgraph=True, enable_piecewise_cuda_graph=cuda_graph) + if torch_compile else None) pytorch_config = dict( disable_overlap_scheduler=not overlap_scheduler, use_cuda_graph=cuda_graph, @@ -524,11 +564,13 @@ def test_bfloat16(self, mtp_nextn, attention_dp, cuda_graph, mtp_config = None if mtp_nextn > 0: mtp_config = MTPDecodingConfig(num_nextn_predict_layers=mtp_nextn) - llm = LLM(self.MODEL_PATH, - kv_cache_config=kv_cache_config, - **pytorch_config, - enable_attention_dp=attention_dp, - speculative_config=mtp_config) + llm = LLM( + self.MODEL_PATH, + kv_cache_config=kv_cache_config, + **pytorch_config, + enable_attention_dp=attention_dp, + speculative_config=mtp_config, + ) with llm: task = MMLU(self.MODEL_NAME) task.evaluate(llm) @@ -539,19 +581,36 @@ def test_bfloat16(self, mtp_nextn, attention_dp, cuda_graph, @parametrize_with_ids( "torch_compile", [False, pytest.param(True, marks=skip_device_contain_gb200)]) - @parametrize_with_ids("attention_dp,cuda_graph,overlap_scheduler", - [(False, False, False), (True, False, False), - (False, True, False), (False, False, True), - (False, True, True), (True, True, True)]) + @parametrize_with_ids( + "attention_dp,cuda_graph,overlap_scheduler", + [ + (False, False, False), + (True, False, False), + (False, True, False), + (False, False, True), + (False, True, True), + (True, True, True), + ], + ) # Only Hopper and Blackwell MLA kernel supports MTP @parametrize_with_ids("mtp_nextn", [0, pytest.param(2, marks=skip_pre_hopper)]) - @pytest.mark.parametrize("tp_size,pp_size,ep_size", [(4, 1, 1), (4, 1, 4), - (2, 2, 1), (1, 4, 1)], - ids=["tp4", "ep4", "tp2pp2", "pp4"]) - def test_bfloat16_4gpus(self, tp_size, pp_size, ep_size, mtp_nextn, - attention_dp, cuda_graph, overlap_scheduler, - torch_compile): + @pytest.mark.parametrize( + "tp_size,pp_size,ep_size", + [(4, 1, 1), (4, 1, 4), (2, 2, 1), (1, 4, 1)], + ids=["tp4", "ep4", "tp2pp2", "pp4"], + ) + def test_bfloat16_4gpus( + self, + tp_size, + pp_size, + ep_size, + mtp_nextn, + attention_dp, + cuda_graph, + overlap_scheduler, + torch_compile, + ): if torch_compile and mtp_nextn > 0: pytest.skip("https://nvbugs/5252313") if torch_compile and attention_dp: @@ -559,9 +618,9 @@ def test_bfloat16_4gpus(self, tp_size, pp_size, ep_size, mtp_nextn, if torch_compile and pp_size > 1: pytest.skip("PP with torch.compile is not supported yet.") kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.9) - torch_compile_config = TorchCompileConfig( - enable_fullgraph=True, - enable_piecewise_cuda_graph=cuda_graph) if torch_compile else None + torch_compile_config = (TorchCompileConfig( + enable_fullgraph=True, enable_piecewise_cuda_graph=cuda_graph) + if torch_compile else None) pytorch_config = dict( disable_overlap_scheduler=not overlap_scheduler, use_cuda_graph=cuda_graph, @@ -570,14 +629,16 @@ def test_bfloat16_4gpus(self, tp_size, pp_size, ep_size, mtp_nextn, mtp_config = None if mtp_nextn > 0: mtp_config = MTPDecodingConfig(num_nextn_predict_layers=mtp_nextn) - llm = LLM(self.MODEL_PATH, - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - moe_expert_parallel_size=ep_size, - kv_cache_config=kv_cache_config, - **pytorch_config, - enable_attention_dp=attention_dp, - speculative_config=mtp_config) + llm = LLM( + self.MODEL_PATH, + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + **pytorch_config, + enable_attention_dp=attention_dp, + speculative_config=mtp_config, + ) with llm: task = MMLU(self.MODEL_NAME) task.evaluate(llm) @@ -603,13 +664,14 @@ def test_fp8_block_scales(self, mtp, fp8kv, attention_dp, cuda_graph, if torch_compile and attention_dp: pytest.skip("https://nvbugs/5252559") kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.9) - torch_compile_config = TorchCompileConfig( - enable_fullgraph=True, - enable_piecewise_cuda_graph=cuda_graph) if torch_compile else None + torch_compile_config = (TorchCompileConfig( + enable_fullgraph=True, enable_piecewise_cuda_graph=cuda_graph) + if torch_compile else None) pytorch_config = dict( disable_overlap_scheduler=not overlap_scheduler, use_cuda_graph=cuda_graph, torch_compile_config=torch_compile_config, + moe_backend="CUTEDSL", ) quant_config = QuantConfig() @@ -626,12 +688,76 @@ def test_fp8_block_scales(self, mtp, fp8kv, attention_dp, cuda_graph, mtp_config = MTPDecodingConfig(num_nextn_predict_layers=mtp_nextn, use_mtp_vanilla=True) - llm = LLM(f"{llm_models_root()}/DeepSeek-V3-Lite/fp8", - kv_cache_config=kv_cache_config, - **pytorch_config, - quant_config=quant_config, - enable_attention_dp=attention_dp, - speculative_config=mtp_config) + llm = LLM( + f"{llm_models_root()}/DeepSeek-V3-Lite/fp8", + kv_cache_config=kv_cache_config, + **pytorch_config, + quant_config=quant_config, + enable_attention_dp=attention_dp, + speculative_config=mtp_config, + ) + + assert llm.args.quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES + if fp8kv: + assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 + + with llm: + # No need to run MMLU for fp8kv + if not fp8kv: + task = MMLU(self.MODEL_NAME) + task.evaluate(llm) + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + + @skip_no_hopper + @parametrize_with_ids("torch_compile", [False]) + @parametrize_with_ids( + "fp8kv,attention_dp,cuda_graph,overlap_scheduler", + [(False, False, False, False)], + ) + @parametrize_with_ids("mtp_nextn", [0]) + def test_cute_dsl_fp8_block_scales( + self, + mtp_nextn, + fp8kv, + attention_dp, + cuda_graph, + overlap_scheduler, + torch_compile, + ): + if torch_compile and mtp_nextn > 0: + pytest.skip("https://nvbugs/5252313") + if torch_compile and attention_dp: + pytest.skip("https://nvbugs/5252559") + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.9) + torch_compile_config = (TorchCompileConfig( + enable_fullgraph=True, enable_piecewise_cuda_graph=cuda_graph) + if torch_compile else None) + pytorch_config = dict( + disable_overlap_scheduler=not overlap_scheduler, + use_cuda_graph=cuda_graph, + torch_compile_config=torch_compile_config, + moe_backend="CUTEDSL", + ) + + quant_config = QuantConfig() + quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES + if fp8kv: + quant_config.kv_cache_quant_algo = QuantAlgo.FP8 + pytorch_config["kv_cache_dtype"] = "fp8" + + mtp_config = None + if mtp_nextn > 0: + mtp_config = MTPDecodingConfig(num_nextn_predict_layers=mtp_nextn) + + llm = LLM( + f"{llm_models_root()}/DeepSeek-V3-Lite/fp8", + kv_cache_config=kv_cache_config, + **pytorch_config, + quant_config=quant_config, + enable_attention_dp=attention_dp, + speculative_config=mtp_config, + ) assert llm.args.quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES if fp8kv: @@ -658,10 +784,12 @@ def test_fp8_block_scales_cuda_graph_padding(self, mtp_nextn): cuda_graph_max_batch_size=512, cuda_graph_padding_enabled=True, ) - llm = LLM(f"{llm_models_root()}/DeepSeek-V3-Lite/fp8", - kv_cache_config=kv_cache_config, - **pytorch_config, - speculative_config=mtp_config) + llm = LLM( + f"{llm_models_root()}/DeepSeek-V3-Lite/fp8", + kv_cache_config=kv_cache_config, + **pytorch_config, + speculative_config=mtp_config, + ) assert llm.args.quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES with llm: task = MMLU(self.MODEL_NAME) @@ -687,13 +815,15 @@ def test_fp8_block_scales_cuda_graph_padding_4gpus(self, mtp_nextn, quant_config = QuantConfig() quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES - llm = LLM(f"{llm_models_root()}/DeepSeek-V3-Lite/fp8", - tensor_parallel_size=4, - kv_cache_config=kv_cache_config, - **pytorch_config, - quant_config=quant_config, - enable_attention_dp=attention_dp, - speculative_config=mtp_config) + llm = LLM( + f"{llm_models_root()}/DeepSeek-V3-Lite/fp8", + tensor_parallel_size=4, + kv_cache_config=kv_cache_config, + **pytorch_config, + quant_config=quant_config, + enable_attention_dp=attention_dp, + speculative_config=mtp_config, + ) assert llm.args.quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES with llm: task = MMLU(self.MODEL_NAME) @@ -703,24 +833,38 @@ def test_fp8_block_scales_cuda_graph_padding_4gpus(self, mtp_nextn, @pytest.mark.skip_less_device(4) @skip_no_hopper + @parametrize_with_ids("torch_compile", [False]) @parametrize_with_ids( - "torch_compile", - [False, pytest.param(True, marks=skip_device_contain_gb200)]) - @parametrize_with_ids("fp8kv,attention_dp,cuda_graph,overlap_scheduler", - [(False, False, False, False), - (True, False, False, False), - (False, True, False, False), - (False, False, True, False), - (False, False, False, True), - (False, True, True, True), (True, False, True, True), - (True, True, True, True)]) + "fp8kv,attention_dp,cuda_graph,overlap_scheduler", + [ + (False, False, False, False), + (True, False, False, False), + (False, True, False, False), + (False, False, True, False), + (False, False, False, True), + (False, True, True, True), + (True, False, True, True), + (True, True, True, True), + ], + ) @parametrize_with_ids("mtp_nextn", [0, 2]) - @pytest.mark.parametrize("tp_size,pp_size,ep_size", [(4, 1, 1), (4, 1, 4), - (2, 2, 1), (1, 4, 1)], - ids=["tp4", "ep4", "tp2pp2", "pp4"]) - def test_fp8_block_scales_4gpus(self, tp_size, pp_size, ep_size, mtp_nextn, - fp8kv, attention_dp, cuda_graph, - overlap_scheduler, torch_compile): + @pytest.mark.parametrize( + "tp_size,pp_size,ep_size", + [(4, 1, 1), (4, 1, 4), (2, 2, 1), (1, 4, 1)], + ids=["tp4", "ep4", "tp2pp2", "pp4"], + ) + def test_fp8_block_scales_4gpus( + self, + tp_size, + pp_size, + ep_size, + mtp_nextn, + fp8kv, + attention_dp, + cuda_graph, + overlap_scheduler, + torch_compile, + ): if torch_compile and mtp_nextn > 0: pytest.skip("https://nvbugs/5252313") if torch_compile and attention_dp: @@ -728,9 +872,9 @@ def test_fp8_block_scales_4gpus(self, tp_size, pp_size, ep_size, mtp_nextn, if torch_compile and pp_size > 1: pytest.skip("PP with torch.compile is not supported yet.") kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.9) - torch_compile_config = TorchCompileConfig( - enable_fullgraph=True, - enable_piecewise_cuda_graph=cuda_graph) if torch_compile else None + torch_compile_config = (TorchCompileConfig( + enable_fullgraph=True, enable_piecewise_cuda_graph=cuda_graph) + if torch_compile else None) pytorch_config = dict( disable_overlap_scheduler=not overlap_scheduler, use_cuda_graph=cuda_graph, @@ -747,15 +891,93 @@ def test_fp8_block_scales_4gpus(self, tp_size, pp_size, ep_size, mtp_nextn, if mtp_nextn > 0: mtp_config = MTPDecodingConfig(num_nextn_predict_layers=mtp_nextn) - llm = LLM(f"{llm_models_root()}/DeepSeek-V3-Lite/fp8", - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - moe_expert_parallel_size=ep_size, - kv_cache_config=kv_cache_config, - **pytorch_config, - quant_config=quant_config, - enable_attention_dp=attention_dp, - speculative_config=mtp_config) + llm = LLM( + f"{llm_models_root()}/DeepSeek-V3-Lite/fp8", + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + **pytorch_config, + quant_config=quant_config, + enable_attention_dp=attention_dp, + speculative_config=mtp_config, + ) + + assert llm.args.quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES + if fp8kv: + assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 + + with llm: + # No need to run MMLU for fp8kv + if not fp8kv: + task = MMLU(self.MODEL_NAME) + task.evaluate(llm) + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + + @pytest.mark.skip_less_device(4) + @skip_no_hopper + @parametrize_with_ids("torch_compile", [False]) + @parametrize_with_ids( + "fp8kv,attention_dp,cuda_graph,overlap_scheduler", + [(False, False, False, False)], + ) + @parametrize_with_ids("mtp_nextn", [0]) + @pytest.mark.parametrize( + "tp_size,pp_size,ep_size", + [(4, 1, 1), (4, 1, 4), (2, 2, 1), (1, 4, 1)], + ids=["tp4", "ep4", "tp2pp2", "pp4"], + ) + def test_cute_dsl_fp8_block_scales_4gpus( + self, + tp_size, + pp_size, + ep_size, + mtp_nextn, + fp8kv, + attention_dp, + cuda_graph, + overlap_scheduler, + torch_compile, + ): + if torch_compile and mtp_nextn > 0: + pytest.skip("https://nvbugs/5252313") + if torch_compile and attention_dp: + pytest.skip("https://nvbugs/5252559") + if torch_compile and pp_size > 1: + pytest.skip("PP with torch.compile is not supported yet.") + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.9) + torch_compile_config = (TorchCompileConfig( + enable_fullgraph=True, enable_piecewise_cuda_graph=cuda_graph) + if torch_compile else None) + pytorch_config = dict( + disable_overlap_scheduler=not overlap_scheduler, + use_cuda_graph=cuda_graph, + torch_compile_config=torch_compile_config, + moe_backend="CUTEDSL", + ) + + quant_config = QuantConfig() + quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES + if fp8kv: + quant_config.kv_cache_quant_algo = QuantAlgo.FP8 + pytorch_config["kv_cache_dtype"] = "fp8" + + mtp_config = None + if mtp_nextn > 0: + mtp_config = MTPDecodingConfig(num_nextn_predict_layers=mtp_nextn) + + llm = LLM( + f"{llm_models_root()}/DeepSeek-V3-Lite/fp8", + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + **pytorch_config, + quant_config=quant_config, + enable_attention_dp=attention_dp, + speculative_config=mtp_config, + ) assert llm.args.quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES if fp8kv: @@ -785,15 +1007,18 @@ def test_fp8_block_scales_4gpus_static_eplb(self): eplb_config = MoeLoadBalancerConfig( num_slots=num_slots, initial_global_assignments=initial_global_assignments, - layer_updates_per_iter=0) + layer_updates_per_iter=0, + ) pytorch_backend_options = dict(use_cuda_graph=True, moe_load_balancer=eplb_config) - llm = LLM(f"{llm_models_root()}/DeepSeek-V3-Lite/fp8", - tensor_parallel_size=4, - moe_expert_parallel_size=4, - kv_cache_config=kv_cache_config, - **pytorch_backend_options, - enable_attention_dp=True) + llm = LLM( + f"{llm_models_root()}/DeepSeek-V3-Lite/fp8", + tensor_parallel_size=4, + moe_expert_parallel_size=4, + kv_cache_config=kv_cache_config, + **pytorch_backend_options, + enable_attention_dp=True, + ) with llm: task = MMLU(self.MODEL_NAME) task.evaluate(llm) @@ -804,25 +1029,38 @@ def test_fp8_block_scales_4gpus_static_eplb(self): @parametrize_with_ids( "torch_compile", [False, pytest.param(True, marks=skip_device_contain_gb200)]) - @parametrize_with_ids("fp8kv,attention_dp,cuda_graph,overlap_scheduler", - [(False, False, False, False), - (True, False, False, False), - (False, True, False, False), - (False, False, True, False), - (False, False, False, True), - (True, False, True, True), (True, True, True, True)]) + @parametrize_with_ids( + "fp8kv,attention_dp,cuda_graph,overlap_scheduler", + [ + (False, False, False, False), + (True, False, False, False), + (False, True, False, False), + (False, False, True, False), + (False, False, False, True), + (True, False, True, True), + (True, True, True, True), + ], + ) @parametrize_with_ids("mtp_nextn", [0, 2]) @parametrize_with_ids("moe_backend", ["CUTLASS", "TRTLLM"]) - def test_nvfp4(self, fp8kv, attention_dp, cuda_graph, overlap_scheduler, - torch_compile, mtp_nextn, moe_backend): + def test_nvfp4( + self, + fp8kv, + attention_dp, + cuda_graph, + overlap_scheduler, + torch_compile, + mtp_nextn, + moe_backend, + ): if torch_compile and mtp_nextn > 0: pytest.skip("https://nvbugs/5252313") if torch_compile and attention_dp: pytest.skip("https://nvbugs/5252559") kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.9) - torch_compile_config = TorchCompileConfig( - enable_fullgraph=True, - enable_piecewise_cuda_graph=cuda_graph) if torch_compile else None + torch_compile_config = (TorchCompileConfig( + enable_fullgraph=True, enable_piecewise_cuda_graph=cuda_graph) + if torch_compile else None) pytorch_config = dict( disable_overlap_scheduler=not overlap_scheduler, use_cuda_graph=cuda_graph, @@ -839,12 +1077,14 @@ def test_nvfp4(self, fp8kv, attention_dp, cuda_graph, overlap_scheduler, quant_config.kv_cache_quant_algo = QuantAlgo.FP8 pytorch_config["kv_cache_dtype"] = "fp8" - llm = LLM(f"{llm_models_root()}/DeepSeek-V3-Lite/nvfp4_moe_only_mtp", - kv_cache_config=kv_cache_config, - **pytorch_config, - quant_config=quant_config, - enable_attention_dp=attention_dp, - speculative_config=mtp_config) + llm = LLM( + f"{llm_models_root()}/DeepSeek-V3-Lite/nvfp4_moe_only_mtp", + kv_cache_config=kv_cache_config, + **pytorch_config, + quant_config=quant_config, + enable_attention_dp=attention_dp, + speculative_config=mtp_config, + ) assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 if fp8kv: @@ -863,21 +1103,38 @@ def test_nvfp4(self, fp8kv, attention_dp, cuda_graph, overlap_scheduler, @parametrize_with_ids( "torch_compile", [False, pytest.param(True, marks=skip_device_contain_gb200)]) - @parametrize_with_ids("fp8kv,attention_dp,cuda_graph,overlap_scheduler", - [(False, False, False, False), - (True, False, False, False), - (False, True, False, False), - (False, False, True, False), - (False, False, False, True), - (True, False, True, True), (True, True, True, True)]) - @pytest.mark.parametrize("tp_size,pp_size,ep_size", [(4, 1, 1), (4, 1, 4), - (2, 2, 1), (1, 4, 1)], - ids=["tp4", "ep4", "tp2pp2", "pp4"]) + @parametrize_with_ids( + "fp8kv,attention_dp,cuda_graph,overlap_scheduler", + [ + (False, False, False, False), + (True, False, False, False), + (False, True, False, False), + (False, False, True, False), + (False, False, False, True), + (True, False, True, True), + (True, True, True, True), + ], + ) + @pytest.mark.parametrize( + "tp_size,pp_size,ep_size", + [(4, 1, 1), (4, 1, 4), (2, 2, 1), (1, 4, 1)], + ids=["tp4", "ep4", "tp2pp2", "pp4"], + ) @parametrize_with_ids("mtp_nextn", [0, 2]) @parametrize_with_ids("moe_backend", ["CUTLASS", "TRTLLM"]) - def test_nvfp4_4gpus(self, fp8kv, attention_dp, cuda_graph, - overlap_scheduler, tp_size, pp_size, ep_size, - torch_compile, mtp_nextn, moe_backend): + def test_nvfp4_4gpus( + self, + fp8kv, + attention_dp, + cuda_graph, + overlap_scheduler, + tp_size, + pp_size, + ep_size, + torch_compile, + mtp_nextn, + moe_backend, + ): if torch_compile and mtp_nextn > 0: pytest.skip("https://nvbugs/5252313") if torch_compile and attention_dp: @@ -887,9 +1144,9 @@ def test_nvfp4_4gpus(self, fp8kv, attention_dp, cuda_graph, if not attention_dp and (tp_size > 1 or ep_size > 1): pytest.skip("https://nvbugs/5336321") kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.9) - torch_compile_config = TorchCompileConfig( - enable_fullgraph=True, - enable_piecewise_cuda_graph=cuda_graph) if torch_compile else None + torch_compile_config = (TorchCompileConfig( + enable_fullgraph=True, enable_piecewise_cuda_graph=cuda_graph) + if torch_compile else None) pytorch_config = dict( disable_overlap_scheduler=not overlap_scheduler, use_cuda_graph=cuda_graph, @@ -907,15 +1164,17 @@ def test_nvfp4_4gpus(self, fp8kv, attention_dp, cuda_graph, quant_config.kv_cache_quant_algo = QuantAlgo.FP8 pytorch_config["kv_cache_dtype"] = "fp8" - llm = LLM(f"{llm_models_root()}/DeepSeek-V3-Lite/nvfp4_moe_only_mtp", - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - moe_expert_parallel_size=ep_size, - kv_cache_config=kv_cache_config, - **pytorch_config, - quant_config=quant_config, - enable_attention_dp=attention_dp, - speculative_config=mtp_config) + llm = LLM( + f"{llm_models_root()}/DeepSeek-V3-Lite/nvfp4_moe_only_mtp", + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + **pytorch_config, + quant_config=quant_config, + enable_attention_dp=attention_dp, + speculative_config=mtp_config, + ) assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 if fp8kv: @@ -931,17 +1190,25 @@ def test_nvfp4_4gpus(self, fp8kv, attention_dp, cuda_graph, @parametrize_with_ids( "fp8kv,attention_dp,cuda_graph,overlap_scheduler", - [(False, False, False, False), - pytest.param(True, False, False, False, marks=skip_no_hopper), - (False, True, False, False), (False, False, True, False), - (False, False, False, True), (False, True, True, True), - pytest.param(True, True, True, True, marks=skip_no_hopper)]) + [ + (False, False, False, False), + pytest.param(True, False, False, False, marks=skip_no_hopper), + (False, True, False, False), + (False, False, True, False), + (False, False, False, True), + (False, True, True, True), + pytest.param(True, True, True, True, marks=skip_no_hopper), + ], + ) @parametrize_with_ids("mtp_nextn", [0, 2]) - @parametrize_with_ids("quant_dtype", [ - pytest.param("none", marks=skip_pre_hopper), - pytest.param("fp8", marks=skip_no_hopper), - pytest.param("nvfp4", marks=skip_pre_blackwell) - ]) + @parametrize_with_ids( + "quant_dtype", + [ + pytest.param("none", marks=skip_pre_hopper), + pytest.param("fp8", marks=skip_no_hopper), + pytest.param("nvfp4", marks=skip_pre_blackwell), + ], + ) def test_no_kv_cache_reuse(self, quant_dtype, mtp_nextn, fp8kv, attention_dp, cuda_graph, overlap_scheduler): if quant_dtype == "nvfp4" and mtp_nextn > 0: @@ -976,12 +1243,14 @@ def test_no_kv_cache_reuse(self, quant_dtype, mtp_nextn, fp8kv, quant_config.kv_cache_quant_algo = QuantAlgo.FP8 pytorch_config["kv_cache_dtype"] = "fp8" - llm = LLM(model_path, - kv_cache_config=kv_cache_config, - **pytorch_config, - quant_config=quant_config, - enable_attention_dp=attention_dp, - speculative_config=mtp_config) + llm = LLM( + model_path, + kv_cache_config=kv_cache_config, + **pytorch_config, + quant_config=quant_config, + enable_attention_dp=attention_dp, + speculative_config=mtp_config, + ) if quant_dtype == "fp8": assert llm.args.quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES @@ -1011,74 +1280,100 @@ class TestDeepSeekR1(LlmapiAccuracyTestHarness): "tp_size,pp_size,ep_size,mtp_nextn,fp8kv,attention_dp,cuda_graph,overlap_scheduler,max_batch_size,moe_backend", [ # Use a larger batch_size to speed up the tests - pytest.param(8, - 1, - 4, - 3, - False, - False, - True, - True, - 32, - "CUTLASS", - marks=pytest.mark.skip_less_device(8)), - pytest.param(8, - 1, - 4, - 3, - False, - False, - True, - True, - 32, - "TRTLLM", - marks=pytest.mark.skip_less_device(8)), - pytest.param(8, - 1, - 8, - 0, - True, - True, - True, - True, - 32, - "CUTLASS", - marks=pytest.mark.skip_less_device(8)), - pytest.param(8, - 1, - 1, - 0, - True, - True, - True, - True, - 32, - "CUTLASS", - marks=pytest.mark.skip_less_device(8)), - pytest.param(4, - 1, - 1, - 0, - True, - True, - True, - True, - 16, - "CUTLASS", - marks=pytest.mark.skip_less_device(4)), + pytest.param( + 8, + 1, + 4, + 3, + False, + False, + True, + True, + 32, + "CUTLASS", + marks=pytest.mark.skip_less_device(8), + ), + pytest.param( + 8, + 1, + 4, + 3, + False, + False, + True, + True, + 32, + "TRTLLM", + marks=pytest.mark.skip_less_device(8), + ), + pytest.param( + 8, + 1, + 8, + 0, + True, + True, + True, + True, + 32, + "CUTLASS", + marks=pytest.mark.skip_less_device(8), + ), + pytest.param( + 8, + 1, + 1, + 0, + True, + True, + True, + True, + 32, + "CUTLASS", + marks=pytest.mark.skip_less_device(8), + ), + pytest.param( + 4, + 1, + 1, + 0, + True, + True, + True, + True, + 16, + "CUTLASS", + marks=pytest.mark.skip_less_device(4), + ), ], ids=[ - "latency", "latency_trtllmgen", "throughput", "throughput_tp8", - "throughput_tp4" - ]) - def test_nvfp4_multi_gpus(self, tp_size, pp_size, ep_size, mtp_nextn, fp8kv, - attention_dp, cuda_graph, overlap_scheduler, - max_batch_size, moe_backend): + "latency", + "latency_trtllmgen", + "throughput", + "throughput_tp8", + "throughput_tp4", + ], + ) + def test_nvfp4_multi_gpus( + self, + tp_size, + pp_size, + ep_size, + mtp_nextn, + fp8kv, + attention_dp, + cuda_graph, + overlap_scheduler, + max_batch_size, + moe_backend, + ): kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.9) - pytorch_config = dict(disable_overlap_scheduler=not overlap_scheduler, - use_cuda_graph=cuda_graph, - moe_backend=moe_backend) + pytorch_config = dict( + disable_overlap_scheduler=not overlap_scheduler, + use_cuda_graph=cuda_graph, + moe_backend=moe_backend, + ) quant_config = QuantConfig() quant_config.quant_algo = QuantAlgo.NVFP4 @@ -1089,16 +1384,18 @@ def test_nvfp4_multi_gpus(self, tp_size, pp_size, ep_size, mtp_nextn, fp8kv, mtp_config = None if mtp_nextn > 0: mtp_config = MTPDecodingConfig(num_nextn_predict_layers=mtp_nextn) - llm = LLM(f"{llm_models_root()}/DeepSeek-R1/DeepSeek-R1-FP4", - max_batch_size=max_batch_size, - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - moe_expert_parallel_size=ep_size, - kv_cache_config=kv_cache_config, - **pytorch_config, - quant_config=quant_config, - enable_attention_dp=attention_dp, - speculative_config=mtp_config) + llm = LLM( + f"{llm_models_root()}/DeepSeek-R1/DeepSeek-R1-FP4", + max_batch_size=max_batch_size, + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + **pytorch_config, + quant_config=quant_config, + enable_attention_dp=attention_dp, + speculative_config=mtp_config, + ) assert llm.args.moe_backend == moe_backend assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 @@ -1119,12 +1416,24 @@ def test_nvfp4_multi_gpus(self, tp_size, pp_size, ep_size, mtp_nextn, fp8kv, @skip_pre_hopper @pytest.mark.parametrize( "tp_size,pp_size,ep_size,mtp_nextn,fp8kv,attention_dp,cuda_graph,overlap_scheduler,max_batch_size", - [(8, 1, 4, 3, False, False, True, True, 1), - (8, 1, 8, 0, True, True, True, True, 24)], - ids=["latency", "throughput"]) - def test_fp8_blockscale(self, tp_size, pp_size, ep_size, mtp_nextn, fp8kv, - attention_dp, cuda_graph, overlap_scheduler, - max_batch_size): + [ + (8, 1, 4, 3, False, False, True, True, 1), + (8, 1, 8, 0, True, True, True, True, 24), + ], + ids=["latency", "throughput"], + ) + def test_fp8_blockscale( + self, + tp_size, + pp_size, + ep_size, + mtp_nextn, + fp8kv, + attention_dp, + cuda_graph, + overlap_scheduler, + max_batch_size, + ): kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.9) pytorch_config = dict( disable_overlap_scheduler=not overlap_scheduler, @@ -1140,16 +1449,18 @@ def test_fp8_blockscale(self, tp_size, pp_size, ep_size, mtp_nextn, fp8kv, mtp_config = None if mtp_nextn > 0: mtp_config = MTPDecodingConfig(num_nextn_predict_layers=mtp_nextn) - llm = LLM(f"{llm_models_root()}/DeepSeek-R1/DeepSeek-R1", - max_batch_size=max_batch_size, - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - moe_expert_parallel_size=ep_size, - kv_cache_config=kv_cache_config, - **pytorch_config, - quant_config=quant_config, - enable_attention_dp=attention_dp, - speculative_config=mtp_config) + llm = LLM( + f"{llm_models_root()}/DeepSeek-R1/DeepSeek-R1", + max_batch_size=max_batch_size, + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + **pytorch_config, + quant_config=quant_config, + enable_attention_dp=attention_dp, + speculative_config=mtp_config, + ) assert llm.args.quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES if fp8kv: assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 @@ -1163,7 +1474,9 @@ def test_fp8_blockscale(self, tp_size, pp_size, ep_size, mtp_nextn, fp8kv, class TestMinitron4BBaseInstruct(LlmapiAccuracyTestHarness): MODEL_NAME = "nvidia/Nemotron-Mini-4B-Instruct" - MODEL_PATH = f"{llm_models_root()}/nemotron/nemotron-mini-4b-instruct_vfp8-fp8-bf16-export" + MODEL_PATH = ( + f"{llm_models_root()}/nemotron/nemotron-mini-4b-instruct_vfp8-fp8-bf16-export" + ) @skip_pre_ada def test_fp8_prequantized(self): @@ -1182,10 +1495,12 @@ def test_auto_dtype_tp8(self): kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.9) pytorch_config = dict() - with LLM(self.MODEL_PATH, - tensor_parallel_size=8, - kv_cache_config=kv_cache_config, - **pytorch_config) as llm: + with LLM( + self.MODEL_PATH, + tensor_parallel_size=8, + kv_cache_config=kv_cache_config, + **pytorch_config, + ) as llm: task = CnnDailymail(self.MODEL_NAME) task.evaluate(llm) @@ -1212,7 +1527,9 @@ def test_auto_dtype_tp2(self): @pytest.mark.skip_less_device(2) @pytest.mark.skip_device_not_contain(["H100", "B200"]) def test_fp8_prequantized_tp2(self): - model_path = f"{llm_models_root()}/nemotron-nas/Llama-3_3-Nemotron-Super-49B-v1-FP8" + model_path = ( + f"{llm_models_root()}/nemotron-nas/Llama-3_3-Nemotron-Super-49B-v1-FP8" + ) with LLM(model_path, tensor_parallel_size=2) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.FP8 task = MMLU(self.MODEL_NAME) @@ -1260,15 +1577,19 @@ class TestNemotronUltra(LlmapiAccuracyTestHarness): @pytest.mark.skip_less_device(8) @pytest.mark.skip_less_device_memory(140000) @parametrize_with_ids("cuda_graph", [False, True]) - @pytest.mark.parametrize("tp_size,pp_size,ep_size", [(8, 1, 1), (8, 1, 4), - (8, 1, 8)], - ids=["tp8", "tp8ep4", "tp8ep8"]) + @pytest.mark.parametrize( + "tp_size,pp_size,ep_size", + [(8, 1, 1), (8, 1, 4), (8, 1, 8)], + ids=["tp8", "tp8ep4", "tp8ep8"], + ) def test_auto_dtype(self, cuda_graph, tp_size, pp_size, ep_size): - with LLM(self.MODEL_PATH, - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - moe_expert_parallel_size=ep_size, - use_cuda_graph=cuda_graph) as llm: + with LLM( + self.MODEL_PATH, + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + moe_expert_parallel_size=ep_size, + use_cuda_graph=cuda_graph, + ) as llm: task = MMLU(self.MODEL_NAME) task.evaluate(llm) task = GSM8K(self.MODEL_NAME) @@ -1280,16 +1601,22 @@ def test_auto_dtype(self, cuda_graph, tp_size, pp_size, ep_size): @pytest.mark.skip_less_device(8) @pytest.mark.skip_device_not_contain(["H100", "B200"]) @parametrize_with_ids("cuda_graph", [False, True]) - @pytest.mark.parametrize("tp_size,pp_size,ep_size", [(8, 1, 1), (8, 1, 4), - (8, 1, 8)], - ids=["tp8", "tp8ep4", "tp8ep8"]) + @pytest.mark.parametrize( + "tp_size,pp_size,ep_size", + [(8, 1, 1), (8, 1, 4), (8, 1, 8)], + ids=["tp8", "tp8ep4", "tp8ep8"], + ) def test_fp8_prequantized(self, cuda_graph, tp_size, pp_size, ep_size): - model_path = f"{llm_models_root()}/nemotron-nas/Llama-3_1-Nemotron-Ultra-253B-v1-FP8" - with LLM(model_path, - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - moe_expert_parallel_size=ep_size, - use_cuda_graph=cuda_graph) as llm: + model_path = ( + f"{llm_models_root()}/nemotron-nas/Llama-3_1-Nemotron-Ultra-253B-v1-FP8" + ) + with LLM( + model_path, + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + moe_expert_parallel_size=ep_size, + use_cuda_graph=cuda_graph, + ) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.FP8 assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 task = MMLU(self.MODEL_NAME) @@ -1320,9 +1647,11 @@ def test_auto_dtype(self): @skip_pre_ada def test_reasoning_fp8_prequantized(self): kv_cache_config = KvCacheConfig(enable_block_reuse=False) - with LLM(f"{llm_models_root()}/Nemotron-H-8B-Reasoning-128K-FP8", - kv_cache_config=kv_cache_config, - max_batch_size=256) as llm: + with LLM( + f"{llm_models_root()}/Nemotron-H-8B-Reasoning-128K-FP8", + kv_cache_config=kv_cache_config, + max_batch_size=256, + ) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.FP8 assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 task = MMLU(self.MODEL_NAME) @@ -1337,7 +1666,7 @@ class TestQwen2_7BInstruct(LlmapiAccuracyTestHarness): EXTRA_EVALUATOR_KWARGS = dict( apply_chat_template=True, system_prompt= - "You are a helpful assistant, please summarize the article entered by the user with one or two sentences." + "You are a helpful assistant, please summarize the article entered by the user with one or two sentences.", ) def test_auto_dtype(self): @@ -1354,18 +1683,21 @@ class TestQwen3_8B(LlmapiAccuracyTestHarness): @pytest.mark.parametrize( "tp_size,pp_size,ep_size,attention_dp,cuda_graph,overlap_scheduler", [(1, 1, 1, False, True, True)], - ids=["latency"]) + ids=["latency"], + ) def test_fp8_block_scales(self, tp_size, pp_size, ep_size, attention_dp, cuda_graph, overlap_scheduler): pytorch_config = dict(disable_overlap_scheduler=not overlap_scheduler, use_cuda_graph=cuda_graph) - llm = LLM(f"{llm_models_root()}/Qwen3/Qwen3-8B-FP8", - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - moe_expert_parallel_size=ep_size, - **pytorch_config, - enable_attention_dp=attention_dp) + llm = LLM( + f"{llm_models_root()}/Qwen3/Qwen3-8B-FP8", + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + moe_expert_parallel_size=ep_size, + **pytorch_config, + enable_attention_dp=attention_dp, + ) with llm: task = CnnDailymail(self.MODEL_NAME) task.evaluate(llm) @@ -1375,18 +1707,21 @@ def test_fp8_block_scales(self, tp_size, pp_size, ep_size, attention_dp, @pytest.mark.parametrize( "tp_size,pp_size,ep_size,attention_dp,cuda_graph,overlap_scheduler", [(1, 1, 1, False, True, True)], - ids=["latency"]) + ids=["latency"], + ) def test_bf16(self, tp_size, pp_size, ep_size, attention_dp, cuda_graph, overlap_scheduler): pytorch_config = dict(disable_overlap_scheduler=not overlap_scheduler, use_cuda_graph=cuda_graph) - llm = LLM(f"{llm_models_root()}/Qwen3/Qwen3-8B", - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - moe_expert_parallel_size=ep_size, - **pytorch_config, - enable_attention_dp=attention_dp) + llm = LLM( + f"{llm_models_root()}/Qwen3/Qwen3-8B", + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + moe_expert_parallel_size=ep_size, + **pytorch_config, + enable_attention_dp=attention_dp, + ) with llm: task = CnnDailymail(self.MODEL_NAME) task.evaluate(llm) @@ -1402,18 +1737,21 @@ class TestQwen3_30B_A3B(LlmapiAccuracyTestHarness): @pytest.mark.parametrize( "tp_size,pp_size,ep_size,attention_dp,cuda_graph,overlap_scheduler", [(1, 1, 1, False, False, True)], - ids=["latency"]) + ids=["latency"], + ) def test_fp8_block_scales(self, tp_size, pp_size, ep_size, attention_dp, cuda_graph, overlap_scheduler): pytorch_config = dict(disable_overlap_scheduler=not overlap_scheduler, use_cuda_graph=cuda_graph) - llm = LLM(f"{llm_models_root()}/Qwen3/Qwen3-30B-A3B-FP8", - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - moe_expert_parallel_size=ep_size, - **pytorch_config, - enable_attention_dp=attention_dp) + llm = LLM( + f"{llm_models_root()}/Qwen3/Qwen3-30B-A3B-FP8", + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + moe_expert_parallel_size=ep_size, + **pytorch_config, + enable_attention_dp=attention_dp, + ) with llm: task = MMLU(self.MODEL_NAME) task.evaluate(llm) @@ -1424,7 +1762,8 @@ def test_fp8_block_scales(self, tp_size, pp_size, ep_size, attention_dp, @pytest.mark.parametrize( "tp_size,pp_size,ep_size,attention_dp,cuda_graph,overlap_scheduler", [(1, 1, 1, True, True, True)], - ids=["latency"]) + ids=["latency"], + ) def test_fp8(self, tp_size, pp_size, ep_size, attention_dp, cuda_graph, overlap_scheduler): pytorch_config = dict(disable_overlap_scheduler=not overlap_scheduler, @@ -1436,7 +1775,8 @@ def test_fp8(self, tp_size, pp_size, ep_size, attention_dp, cuda_graph, pipeline_parallel_size=pp_size, moe_expert_parallel_size=ep_size, **pytorch_config, - enable_attention_dp=attention_dp) + enable_attention_dp=attention_dp, + ) with llm: task = MMLU(self.MODEL_NAME) task.evaluate(llm) @@ -1453,8 +1793,10 @@ def test_fp8(self, tp_size, pp_size, ep_size, attention_dp, cuda_graph, (4, 1, 4, False, True, True, "TRTLLM"), ], ids=[ - "latency_moe_cutlass", "latency_moe_trtllm", - "4gpu_latency_moe_trtllm", "4gpu_latency_moe_cutlass" + "latency_moe_cutlass", + "latency_moe_trtllm", + "4gpu_latency_moe_trtllm", + "4gpu_latency_moe_cutlass", ], ) def test_nvfp4( @@ -1479,7 +1821,8 @@ def test_nvfp4( pipeline_parallel_size=pp_size, moe_expert_parallel_size=ep_size, **pytorch_config, - enable_attention_dp=attention_dp) + enable_attention_dp=attention_dp, + ) with llm: task = MMLU(self.MODEL_NAME) task.evaluate(llm) @@ -1494,18 +1837,21 @@ class TestQwen3_32B(LlmapiAccuracyTestHarness): @pytest.mark.parametrize( "tp_size,pp_size,ep_size,attention_dp,cuda_graph,overlap_scheduler", [(1, 1, 1, False, False, True)], - ids=["latency"]) + ids=["latency"], + ) def test_fp8_block_scales(self, tp_size, pp_size, ep_size, attention_dp, cuda_graph, overlap_scheduler): pytorch_config = dict(disable_overlap_scheduler=not overlap_scheduler, use_cuda_graph=cuda_graph) - llm = LLM(f"{llm_models_root()}/Qwen3/Qwen3-32B-FP8", - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - moe_expert_parallel_size=ep_size, - **pytorch_config, - enable_attention_dp=attention_dp) + llm = LLM( + f"{llm_models_root()}/Qwen3/Qwen3-32B-FP8", + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + moe_expert_parallel_size=ep_size, + **pytorch_config, + enable_attention_dp=attention_dp, + ) with llm: task = CnnDailymail(self.MODEL_NAME) task.evaluate(llm) @@ -1520,7 +1866,8 @@ class TestQwen3_235B_A22B(LlmapiAccuracyTestHarness): @pytest.mark.parametrize( "tp_size,pp_size,ep_size,attention_dp,cuda_graph,overlap_scheduler", [(8, 1, 8, True, True, True), (8, 1, 8, False, True, True)], - ids=["latency", "throughput_latency"]) + ids=["latency", "throughput_latency"], + ) def test_fp8(self, tp_size, pp_size, ep_size, attention_dp, cuda_graph, overlap_scheduler): pytorch_config = dict(disable_overlap_scheduler=not overlap_scheduler, @@ -1534,7 +1881,8 @@ def test_fp8(self, tp_size, pp_size, ep_size, attention_dp, cuda_graph, moe_expert_parallel_size=ep_size, **pytorch_config, enable_attention_dp=attention_dp, - kv_cache_config=kv_cache_config) + kv_cache_config=kv_cache_config, + ) with llm: task = MMLU(self.MODEL_NAME) task.evaluate(llm) @@ -1544,15 +1892,27 @@ def test_fp8(self, tp_size, pp_size, ep_size, attention_dp, cuda_graph, @skip_pre_blackwell @pytest.mark.parametrize( "tp_size,pp_size,ep_size,attention_dp,cuda_graph,overlap_scheduler,moe_backend", - [(8, 1, 8, True, True, True, "CUTLASS"), - (8, 1, 8, False, True, True, "TRTLLM")], + [ + (8, 1, 8, True, True, True, "CUTLASS"), + (8, 1, 8, False, True, True, "TRTLLM"), + ], ids=["latency_moe_cutlass", "latency_moe_trtllm"], ) - def test_nvfp4(self, tp_size, pp_size, ep_size, attention_dp, cuda_graph, - overlap_scheduler, moe_backend): - pytorch_config = dict(disable_overlap_scheduler=not overlap_scheduler, - use_cuda_graph=cuda_graph, - moe_backend=moe_backend) + def test_nvfp4( + self, + tp_size, + pp_size, + ep_size, + attention_dp, + cuda_graph, + overlap_scheduler, + moe_backend, + ): + pytorch_config = dict( + disable_overlap_scheduler=not overlap_scheduler, + use_cuda_graph=cuda_graph, + moe_backend=moe_backend, + ) kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.6) llm = LLM( @@ -1562,7 +1922,8 @@ def test_nvfp4(self, tp_size, pp_size, ep_size, attention_dp, cuda_graph, moe_expert_parallel_size=ep_size, **pytorch_config, enable_attention_dp=attention_dp, - kv_cache_config=kv_cache_config) + kv_cache_config=kv_cache_config, + ) with llm: task = MMLU(self.MODEL_NAME) task.evaluate(llm) @@ -1598,9 +1959,11 @@ class TestKanana_Instruct(LlmapiAccuracyTestHarness): @pytest.mark.skip_device_not_contain(["H20", "H100"]) def test_auto_dtype(self): "RCCA: https://nvbugspro.nvidia.com/bug/5310520" - pytorch_config = dict(duse_cuda_graph=True, - cuda_graph_padding_enabled=True, - cuda_graph_max_batch_size=384) + pytorch_config = dict( + duse_cuda_graph=True, + cuda_graph_padding_enabled=True, + cuda_graph_max_batch_size=384, + ) with LLM(self.MODEL_PATH, **pytorch_config, enable_attention_dp=True) as llm: task = MMLU(self.MODEL_NAME) diff --git a/tests/unittest/_torch/modules/test_fused_moe.py b/tests/unittest/_torch/modules/test_fused_moe.py index 3ffc94c03cee..f3c0d033f450 100644 --- a/tests/unittest/_torch/modules/test_fused_moe.py +++ b/tests/unittest/_torch/modules/test_fused_moe.py @@ -4,10 +4,12 @@ from typing import Dict, List, Optional from unittest import mock +import _torch.helpers import cloudpickle import pytest import torch import torch.nn as nn +from _torch.helpers import per_block_cast_to_fp8 from mpi4py import MPI from mpi4py.futures import MPIPoolExecutor from utils.util import (skip_neither_ada_nor_hopper_unittest, @@ -20,6 +22,8 @@ DefaultMoeRoutingMethod, RenormalizeMoeRoutingMethod, VanillaMoE, WideEPMoE) +from tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl import \ + CuteDslFusedMoE from tensorrt_llm._torch.modules.fused_moe.fused_moe_wide_ep import \ AlltoallMethodType from tensorrt_llm._torch.modules.gated_mlp import GatedMLP @@ -28,6 +32,7 @@ from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig cloudpickle.register_pickle_by_value(sys.modules[__name__]) +cloudpickle.register_pickle_by_value(_torch.helpers) MPI.pickle.__init__( cloudpickle.dumps, cloudpickle.loads, @@ -37,9 +42,13 @@ @pytest.mark.parametrize( "moe_cls, dtype, experts, RoutingMethodCls", - product([CutlassFusedMoE, VanillaMoE], [torch.float16, torch.bfloat16], - [3, 8, 512], - [DefaultMoeRoutingMethod, RenormalizeMoeRoutingMethod])) + product( + [CutlassFusedMoE, VanillaMoE], + [torch.float16, torch.bfloat16], + [3, 8, 512], + [DefaultMoeRoutingMethod, RenormalizeMoeRoutingMethod], + ), +) def test_fused_moe(moe_cls, dtype, experts, RoutingMethodCls, mapping=None): SEQ_LEN = 8 HIDDEN_SIZE = 64 @@ -82,12 +91,14 @@ def test_fused_moe(moe_cls, dtype, experts, RoutingMethodCls, mapping=None): with torch.inference_mode(), autotune(): fused_moe.forward(x, router_logits) - ref_fused_moe = RefGatedMLPFusedMoE(num_experts=NUM_EXPERTS, - routing_method=routing_method, - hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, - dtype=dtype, - model_config=ModelConfig()) + ref_fused_moe = RefGatedMLPFusedMoE( + num_experts=NUM_EXPERTS, + routing_method=routing_method, + hidden_size=HIDDEN_SIZE, + intermediate_size=INTERMEDIATE_SIZE, + dtype=dtype, + model_config=ModelConfig(), + ) ref_fused_moe.load_weights([weights]) ref_fused_moe.cuda() @@ -116,11 +127,18 @@ def test_fused_moe_multi_gpu(moe_cls, ep_size): with MPIPoolExecutor(max_workers=world_size) as executor: results = executor.map( test_fused_moe, - *zip(*[(moe_cls, torch.bfloat16, 512, DefaultMoeRoutingMethod, - Mapping(world_size=world_size, - tp_size=world_size, - moe_ep_size=ep_size, - moe_tp_size=world_size // ep_size))] * world_size), + *zip(*[( + moe_cls, + torch.bfloat16, + 512, + DefaultMoeRoutingMethod, + Mapping( + world_size=world_size, + tp_size=world_size, + moe_ep_size=ep_size, + moe_tp_size=world_size // ep_size, + ), + )] * world_size), ) for r in results: assert r is None @@ -128,11 +146,15 @@ def test_fused_moe_multi_gpu(moe_cls, ep_size): @pytest.mark.skipif(torch.cuda.device_count() < 4, reason="needs 4 GPUs to run this test") -@pytest.mark.parametrize("alltoall_method_type", [ - AlltoallMethodType.MNNVL, AlltoallMethodType.DeepEP, - AlltoallMethodType.DeepEPLowLatency -], - ids=lambda s: s.name) +@pytest.mark.parametrize( + "alltoall_method_type", + [ + AlltoallMethodType.MNNVL, + AlltoallMethodType.DeepEP, + AlltoallMethodType.DeepEPLowLatency, + ], + ids=lambda s: s.name, +) def test_fused_moe_alltoall(alltoall_method_type): world_size = 4 dtype = torch.bfloat16 @@ -144,12 +166,14 @@ def test_fused_moe_alltoall(alltoall_method_type): def per_rank_test_fused_moe_alltoall(job_id): routing_method = DefaultMoeRoutingMethod(top_k=TOP_K) - mapping = Mapping(world_size=world_size, - rank=mpi_rank(), - tp_size=world_size, - moe_ep_size=world_size, - moe_tp_size=1, - enable_attention_dp=True) + mapping = Mapping( + world_size=world_size, + rank=mpi_rank(), + tp_size=world_size, + moe_ep_size=world_size, + moe_tp_size=1, + enable_attention_dp=True, + ) torch.cuda.set_device(mapping.rank) torch.manual_seed(mapping.rank) @@ -208,12 +232,14 @@ def per_rank_test_fused_moe_alltoall(job_id): x, router_logits, all_rank_num_tokens=all_rank_num_tokens, - use_dp_padding=False) + use_dp_padding=False, + ) ref_output = ref_model.forward( x, router_logits, all_rank_num_tokens=all_rank_num_tokens, - use_dp_padding=False) + use_dp_padding=False, + ) # Evaluate outputs torch.testing.assert_close(output, @@ -254,16 +280,16 @@ def test_fused_moe_fp8(dtype): w3_weight = torch.randn((INTERMEDIATE_SIZE, HIDDEN_SIZE), dtype=dtype).cuda() - w1_weight_fp8, w1_weight_scale = torch.ops.tensorrt_llm.quantize_e4m3_per_tensor( - w1_weight) + w1_weight_fp8, w1_weight_scale = ( + torch.ops.tensorrt_llm.quantize_e4m3_per_tensor(w1_weight)) w1_weight_fp8 = w1_weight_fp8.view(torch.float8_e4m3fn).cuda() - w2_weight_fp8, w2_weight_scale = torch.ops.tensorrt_llm.quantize_e4m3_per_tensor( - w2_weight) + w2_weight_fp8, w2_weight_scale = ( + torch.ops.tensorrt_llm.quantize_e4m3_per_tensor(w2_weight)) w2_weight_fp8 = w2_weight_fp8.view(torch.float8_e4m3fn).cuda() - w3_weight_fp8, w3_weight_scale = torch.ops.tensorrt_llm.quantize_e4m3_per_tensor( - w3_weight) + w3_weight_fp8, w3_weight_scale = ( + torch.ops.tensorrt_llm.quantize_e4m3_per_tensor(w3_weight)) w3_weight_fp8 = w3_weight_fp8.view(torch.float8_e4m3fn).cuda() w1_input_scale = x_scale.cuda() @@ -288,7 +314,8 @@ def test_fused_moe_fp8(dtype): intermediate_size=INTERMEDIATE_SIZE, dtype=dtype, reduce_results=False, - model_config=ModelConfig(quant_config=quant_config)) + model_config=ModelConfig(quant_config=quant_config), + ) fused_moe.cuda() fused_moe.load_weights([weights]) @@ -302,7 +329,8 @@ def test_fused_moe_fp8(dtype): hidden_size=HIDDEN_SIZE, intermediate_size=INTERMEDIATE_SIZE, dtype=dtype, - model_config=ModelConfig(quant_config=quant_config)) + model_config=ModelConfig(quant_config=quant_config), + ) ref_fused_moe.load_weights([weights]) ref_fused_moe.cuda() with torch.inference_mode(): @@ -314,6 +342,195 @@ def test_fused_moe_fp8(dtype): torch.testing.assert_close(output, ref_output, rtol=1e-2, atol=0.1) +def set_tensor_value_2(x, num_row, num_cols): + # Create 2x2 base pattern matrix + pattern = torch.tensor([[0.2, -0.5], [-0.3, 0.1]], device=x.device) + + # Repeat pattern to cover entire matrix + repeated = pattern.repeat((num_row + 1) // 2, + (num_cols + 1) // 2)[:num_row, :num_cols] + + x.copy_(repeated) + + +def set_tensor_value_3(x, num_row, num_cols): + # Create 3x3 base pattern matrix + pattern = torch.tensor( + [[0.1, 0.21, 0.31], [0.3, 0.6, 0.1], [0.11, 0.51, 0.62]], + device=x.device) + + # Repeat pattern to cover entire matrix + repeated = pattern.repeat((num_row + 2) // 3, + (num_cols + 2) // 3)[:num_row, :num_cols] + + x.copy_(repeated) + + +def set_tensor_value_4(x, num_row, num_cols): + # Create 4x4 base pattern matrix + pattern = torch.tensor( + [ + [0.1, 0.21, 0.31, 0.41], + [0.3, 0.6, 0.1, 0.2], + [0.11, 0.51, 0.61, 0.71], + [0.11, 0.52, 0.62, 0.72], + ], + device=x.device, + ) + + # Repeat pattern to cover entire matrix + repeated = pattern.repeat((num_row + 3) // 4, + (num_cols + 3) // 4)[:num_row, :num_cols] + + x.copy_(repeated) + + +@skip_pre_hopper +@pytest.mark.parametrize( + "dtype, num_experts, seq_len, hidden_size, RoutingMethodCls", + product( + [torch.bfloat16], + [72], + [128, 256, 384, 512, 1024, 2048, 4096, 8192], + [2560], + [DefaultMoeRoutingMethod], + ), +) +def test_fused_moe_fp8_blockwise(dtype, + num_experts, + seq_len, + hidden_size, + RoutingMethodCls, + mapping=None): + SEQ_LEN = seq_len + HIDDEN_SIZE = hidden_size + INTERMEDIATE_SIZE = 1536 + NUM_EXPERTS = num_experts + TOP_K = 6 + + routing_method = RoutingMethodCls(top_k=TOP_K) + + mapping = mapping or Mapping() + mapping.rank = mpi_rank() + torch.cuda.set_device(mapping.rank) + torch.manual_seed(0) + torch.cuda.manual_seed(0) + + x = torch.randn((SEQ_LEN, HIDDEN_SIZE), dtype=dtype).cuda() + # Note: we use some special values init x and weight, otherwise the test will false positive failed. + set_tensor_value_2(x, SEQ_LEN, HIDDEN_SIZE) + + x = x.cuda() + router_logits = torch.randn((SEQ_LEN, NUM_EXPERTS), dtype=dtype).cuda() + + weights = {} + for expert_id in range(NUM_EXPERTS): + w1_weight = torch.randn((INTERMEDIATE_SIZE, HIDDEN_SIZE), + dtype=dtype).cuda() + w2_weight = torch.randn((HIDDEN_SIZE, INTERMEDIATE_SIZE), + dtype=dtype).cuda() + w3_weight = torch.randn((INTERMEDIATE_SIZE, HIDDEN_SIZE), + dtype=dtype).cuda() + set_tensor_value_3(w1_weight, INTERMEDIATE_SIZE, HIDDEN_SIZE) + set_tensor_value_4(w2_weight, HIDDEN_SIZE, INTERMEDIATE_SIZE) + set_tensor_value_3(w3_weight, INTERMEDIATE_SIZE, HIDDEN_SIZE) + + w1_weight_fp8, w1_weight_scale = per_block_cast_to_fp8(w1_weight) + w1_weight_fp8 = w1_weight_fp8.view(torch.float8_e4m3fn).cuda() + + w2_weight_fp8, w2_weight_scale = per_block_cast_to_fp8(w2_weight) + w2_weight_fp8 = w2_weight_fp8.view(torch.float8_e4m3fn).cuda() + + w3_weight_fp8, w3_weight_scale = per_block_cast_to_fp8(w3_weight) + w3_weight_fp8 = w3_weight_fp8.view(torch.float8_e4m3fn).cuda() + + weights[f"{expert_id}.w1.weight"] = w1_weight_fp8 + weights[f"{expert_id}.w2.weight"] = w2_weight_fp8 + weights[f"{expert_id}.w3.weight"] = w3_weight_fp8 + weights[f"{expert_id}.w1.weight_scale_inv"] = w1_weight_scale + weights[f"{expert_id}.w2.weight_scale_inv"] = w2_weight_scale + weights[f"{expert_id}.w3.weight_scale_inv"] = w3_weight_scale + weights[f"{expert_id}.w1.weight_scale"] = w1_weight_scale + weights[f"{expert_id}.w2.weight_scale"] = w2_weight_scale + weights[f"{expert_id}.w3.weight_scale"] = w3_weight_scale + + quant_config = QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES) + + fused_moe = CuteDslFusedMoE( + num_experts=NUM_EXPERTS, + routing_method=routing_method, + hidden_size=HIDDEN_SIZE, + intermediate_size=INTERMEDIATE_SIZE, + dtype=dtype, + reduce_results=True, + model_config=ModelConfig(quant_config=quant_config, mapping=mapping), + ) + fused_moe.cuda() + fused_moe.load_weights([weights]) + + fused_moe_origin = CutlassFusedMoE( + num_experts=NUM_EXPERTS, + routing_method=routing_method, + hidden_size=HIDDEN_SIZE, + intermediate_size=INTERMEDIATE_SIZE, + dtype=dtype, + reduce_results=True, + model_config=ModelConfig(quant_config=quant_config, mapping=mapping), + ) + fused_moe_origin.cuda() + fused_moe_origin.load_weights([weights]) + + ref_fused_moe = RefGatedMLPFusedMoE( + num_experts=NUM_EXPERTS, + routing_method=routing_method, + hidden_size=HIDDEN_SIZE, + intermediate_size=INTERMEDIATE_SIZE, + dtype=dtype, + model_config=ModelConfig(quant_config=quant_config), + ) + ref_fused_moe.load_weights([weights]) + ref_fused_moe.cuda() + + with torch.inference_mode(): + output = fused_moe.forward(x, router_logits) + output_origin = fused_moe_origin.forward(x, router_logits) + ref_output = ref_fused_moe.forward(x, router_logits) + + # compare + torch.cuda.synchronize() + torch.testing.assert_close(output_origin, output, rtol=1e-2, atol=0.1) + torch.testing.assert_close(output_origin, ref_output, rtol=1e-2, atol=0.1) + torch.testing.assert_close(output, ref_output, rtol=1e-2, atol=0.1) + return True + + +@pytest.mark.skipif(torch.cuda.device_count() < 4, + reason="needs 4 GPUs to run this test") +@pytest.mark.parametrize("ep_size", [1, 2, 4]) +@pytest.mark.parametrize("routing_method", [DefaultMoeRoutingMethod]) +def test_fused_moe_fp8_blockwise_multi_gpu(ep_size, routing_method): + world_size = 4 + with MPIPoolExecutor(max_workers=world_size) as executor: + results = executor.map( + test_fused_moe_fp8_blockwise, + *zip(*[( + torch.bfloat16, + 72, + 384, + 384, + routing_method, + Mapping( + world_size=world_size, + tp_size=world_size, + moe_ep_size=ep_size, + moe_tp_size=world_size // ep_size, + ), + )] * world_size), + ) + for r in results: + assert r is True + + @skip_pre_blackwell @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_fused_moe_nvfp4(dtype): @@ -351,18 +568,21 @@ def test_fused_moe_nvfp4(dtype): w1_weight_nvfp4, w1_sf_block = torch.ops.trtllm.fp4_quantize( w1_weight, w3_w1_global, SCALING_VECTOR_SIZE, False) - w1_sf_block_unswizzled = torch.ops.tensorrt_llm.nvfp4_block_scale_interleave_reverse( - w1_sf_block.cpu().view(INTERMEDIATE_SIZE, -1)) + w1_sf_block_unswizzled = ( + torch.ops.tensorrt_llm.nvfp4_block_scale_interleave_reverse( + w1_sf_block.cpu().view(INTERMEDIATE_SIZE, -1))) w2_weight_nvfp4, w2_sf_block = torch.ops.trtllm.fp4_quantize( w2_weight, w2_sf_global, SCALING_VECTOR_SIZE, False) - w2_sf_block_unswizzled = torch.ops.tensorrt_llm.nvfp4_block_scale_interleave_reverse( - w2_sf_block.cpu().view(HIDDEN_SIZE, -1)) + w2_sf_block_unswizzled = ( + torch.ops.tensorrt_llm.nvfp4_block_scale_interleave_reverse( + w2_sf_block.cpu().view(HIDDEN_SIZE, -1))) w3_weight_nvfp4, w3_sf_block = torch.ops.trtllm.fp4_quantize( w3_weight, w3_w1_global, SCALING_VECTOR_SIZE, False) - w3_sf_block_unswizzled = torch.ops.tensorrt_llm.nvfp4_block_scale_interleave_reverse( - w3_sf_block.cpu().view(INTERMEDIATE_SIZE, -1)) + w3_sf_block_unswizzled = ( + torch.ops.tensorrt_llm.nvfp4_block_scale_interleave_reverse( + w3_sf_block.cpu().view(INTERMEDIATE_SIZE, -1))) w1_input_scale = x_sf_global.cuda() w2_input_scale = x_sf_global.cuda() @@ -392,7 +612,8 @@ def test_fused_moe_nvfp4(dtype): intermediate_size=INTERMEDIATE_SIZE, dtype=dtype, reduce_results=False, - model_config=ModelConfig(quant_config=quant_config)) + model_config=ModelConfig(quant_config=quant_config), + ) fused_moe.load_weights([weights]) fused_moe.cuda() @@ -403,7 +624,8 @@ def test_fused_moe_nvfp4(dtype): hidden_size=HIDDEN_SIZE, intermediate_size=INTERMEDIATE_SIZE, dtype=dtype, - model_config=ModelConfig(quant_config=quant_config)) + model_config=ModelConfig(quant_config=quant_config), + ) ref_fused_moe.load_weights([weights]) ref_fused_moe.cuda() @@ -451,15 +673,15 @@ def test_fused_moe_w4afp8(dtype): 127, (INTERMEDIATE_SIZE, HIDDEN_SIZE // 2), dtype=torch.int8).cuda() - w1_scale = torch.randn( + w1_scale = (torch.randn( (INTERMEDIATE_SIZE, HIDDEN_SIZE // SCALING_GROUP_SIZE), - dtype=dtype).cuda() * affine_coeff - w2_scale = torch.randn( + dtype=dtype).cuda() * affine_coeff) + w2_scale = (torch.randn( (HIDDEN_SIZE, INTERMEDIATE_SIZE // SCALING_GROUP_SIZE), - dtype=dtype).cuda() * affine_coeff - w3_scale = torch.randn( + dtype=dtype).cuda() * affine_coeff) + w3_scale = (torch.randn( (INTERMEDIATE_SIZE, HIDDEN_SIZE // SCALING_GROUP_SIZE), - dtype=dtype).cuda() * affine_coeff + dtype=dtype).cuda() * affine_coeff) w1_input = torch.randn(1, dtype=torch.float32).cuda() * 0.02 w2_input = w1_input @@ -483,7 +705,8 @@ def test_fused_moe_w4afp8(dtype): intermediate_size=INTERMEDIATE_SIZE, dtype=dtype, reduce_results=False, - model_config=ModelConfig(quant_config=quant_config)) + model_config=ModelConfig(quant_config=quant_config), + ) fused_moe.load_weights([weights]) fused_moe.cuda() @@ -521,16 +744,16 @@ def ref(): p3 = weights[f"{e_idx}.w3.input_scale"].cuda() p3_p1 = max(p1, p3) - act = torch.clamp((act / p3_p1), -448.0, - 448.0).to(torch.float8_e4m3fn).to(dtype) + act = (torch.clamp((act / p3_p1), -448.0, + 448.0).to(torch.float8_e4m3fn).to(dtype)) w3_w1 = (w3_w1.float() * s3_s1.repeat_interleave(128, dim=0).float()).to(dtype) fc1 = torch.matmul(act, w3_w1) * p3_p1 fc1, gate = fc1.chunk(2, dim=-1) fc1 = fc1 * torch.nn.functional.silu(gate) - act = torch.clamp((fc1 / p2), -448.0, - 448.0).to(torch.float8_e4m3fn).to(dtype) + act = (torch.clamp((fc1 / p2), -448.0, + 448.0).to(torch.float8_e4m3fn).to(dtype)) w2 = (w2.float() * s2.repeat_interleave(128, dim=0).float()).to(dtype) fc2 = torch.matmul(act, w2) * p2 @@ -554,13 +777,15 @@ def ref(): class RefGatedMLPFusedMoE(nn.Module): - def __init__(self, - num_experts: int, - routing_method: BaseMoeRoutingMethod, - hidden_size: int, - intermediate_size: int, - dtype: Optional[torch.dtype] = None, - model_config: ModelConfig = ModelConfig()): + def __init__( + self, + num_experts: int, + routing_method: BaseMoeRoutingMethod, + hidden_size: int, + intermediate_size: int, + dtype: Optional[torch.dtype] = None, + model_config: ModelConfig = ModelConfig(), + ): super().__init__() self.num_experts = num_experts self.routing_method = routing_method @@ -599,8 +824,8 @@ def forward(self, hidden_states: torch.Tensor, expert_inputs = hidden_states[batch_idx] output = self.experts[expert_id](expert_inputs) - final_hidden_states[batch_idx] += routing_weights[ - batch_idx, nth_expert, None] * output.float() + final_hidden_states[batch_idx] += ( + routing_weights[batch_idx, nth_expert, None] * output.float()) final_hidden_states = final_hidden_states.reshape(hidden_states.shape) return final_hidden_states @@ -613,42 +838,50 @@ def load_weights(self, weights: List[Dict]): gate_up_proj_weights = [{}, {}] down_proj_weights = [{}] - gate_up_proj_weights[0]['weight'] = weights[f"{expert}.w1.weight"] - gate_up_proj_weights[1]['weight'] = weights[f"{expert}.w3.weight"] - down_proj_weights[0]['weight'] = weights[f"{expert}.w2.weight"] + gate_up_proj_weights[0]["weight"] = weights[f"{expert}.w1.weight"] + gate_up_proj_weights[1]["weight"] = weights[f"{expert}.w3.weight"] + down_proj_weights[0]["weight"] = weights[f"{expert}.w2.weight"] if self.quant_config and self.quant_config.quant_algo == QuantAlgo.FP8: - gate_up_proj_weights[0]['weight_scale'] = weights[ + gate_up_proj_weights[0]["weight_scale"] = weights[ f"{expert}.w1.weight_scale"] - gate_up_proj_weights[1]['weight_scale'] = weights[ + gate_up_proj_weights[1]["weight_scale"] = weights[ f"{expert}.w3.weight_scale"] - down_proj_weights[0]['weight_scale'] = weights[ + down_proj_weights[0]["weight_scale"] = weights[ f"{expert}.w2.weight_scale"] - gate_up_proj_weights[0]['input_scale'] = weights[ + gate_up_proj_weights[0]["input_scale"] = weights[ f"{expert}.w1.input_scale"] - gate_up_proj_weights[1]['input_scale'] = weights[ + gate_up_proj_weights[1]["input_scale"] = weights[ f"{expert}.w3.input_scale"] - down_proj_weights[0]['input_scale'] = weights[ + down_proj_weights[0]["input_scale"] = weights[ f"{expert}.w2.input_scale"] elif self.quant_config and self.quant_config.quant_algo == QuantAlgo.NVFP4: - gate_up_proj_weights[0]['weight_scale'] = weights[ + gate_up_proj_weights[0]["weight_scale"] = weights[ f"{expert}.w1.weight_scale"] - gate_up_proj_weights[1]['weight_scale'] = weights[ + gate_up_proj_weights[1]["weight_scale"] = weights[ f"{expert}.w3.weight_scale"] - down_proj_weights[0]['weight_scale'] = weights[ + down_proj_weights[0]["weight_scale"] = weights[ f"{expert}.w2.weight_scale"] - gate_up_proj_weights[0]['input_scale'] = weights[ + gate_up_proj_weights[0]["input_scale"] = weights[ f"{expert}.w1.input_scale"] - gate_up_proj_weights[1]['input_scale'] = weights[ + gate_up_proj_weights[1]["input_scale"] = weights[ f"{expert}.w3.input_scale"] - down_proj_weights[0]['input_scale'] = weights[ + down_proj_weights[0]["input_scale"] = weights[ f"{expert}.w2.input_scale"] - gate_up_proj_weights[0]['weight_scale_2'] = weights[ + gate_up_proj_weights[0]["weight_scale_2"] = weights[ f"{expert}.w1.weight_scale_2"] - gate_up_proj_weights[1]['weight_scale_2'] = weights[ + gate_up_proj_weights[1]["weight_scale_2"] = weights[ f"{expert}.w3.weight_scale_2"] - down_proj_weights[0]['weight_scale_2'] = weights[ + down_proj_weights[0]["weight_scale_2"] = weights[ f"{expert}.w2.weight_scale_2"] + elif (self.quant_config and self.quant_config.quant_algo + == QuantAlgo.FP8_BLOCK_SCALES): + gate_up_proj_weights[0]["weight_scale"] = weights[ + f"{expert}.w1.weight_scale"] + gate_up_proj_weights[1]["weight_scale"] = weights[ + f"{expert}.w3.weight_scale"] + down_proj_weights[0]["weight_scale"] = weights[ + f"{expert}.w2.weight_scale"] self.experts[expert].gate_up_proj.load_weights(gate_up_proj_weights) self.experts[expert].down_proj.load_weights(down_proj_weights)