diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index 4e630b5bb45d2..500e3e833f51c 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -1160,9 +1160,11 @@ if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND onnxruntime_BUILD_CUDA_EP_AS_P "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/moe_gemm/moe_gemm_kernels_fp16_uint2.cu" "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/moe/qmoe_kernels.cu" ) - if (onnxruntime_USE_FP4_QMOE) + if (TARGET onnxruntime_providers_cuda_plugin_llm_fp4) list(APPEND onnxruntime_providers_cuda_plugin_ut_impl_src "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu") + set_property(SOURCE "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc" + APPEND PROPERTY COMPILE_DEFINITIONS ORT_TEST_NVFP4_GEMV) endif() endif() diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 11ef53dd60d45..c7fe3de628179 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -179,6 +179,9 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, ArgMaxLastAxis); class CUDA_MS_OP_CLASS_NAME(1, EmbeddingRMSNorm); class CUDA_MS_OP_CLASS_NAME(1, RowSelect); class CUDA_MS_OP_CLASS_NAME(1, SigmoidGateMul); +#if !defined(DISABLE_FLOAT8_TYPES) +class CUDA_MS_OP_TYPED_CLASS_NAME(1, Float8E4M3FN_BFloat16, DequantizeLinear); +#endif class CUDA_MS_OP_CLASS_NAME(1, UntileHyperConnectionWeights); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, LinearAttentionGate); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, LinearAttentionGate); @@ -560,6 +563,9 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, +#if !defined(DISABLE_FLOAT8_TYPES) + BuildKernelCreateInfo, +#endif BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu index f436c59c98d85..dbe438cc16a7c 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu @@ -4,6 +4,7 @@ #include "contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h" #include +#include #include // Shared INT/FP4 profiled-shape thresholds. @@ -360,24 +361,33 @@ __global__ void MoeGemvFp4RawNPackedKernel( #endif } -template +template __global__ void MoeGemvFp4RawKPackedKernel( const T* act, const uint8_t* weight, const uint8_t* block_scales, const float* global_scales, const T* bias, T* out, const int64_t* expert_first_token_offset, const int* permuted_row_to_expert, int num_experts, int64_t weight_expert_stride, int64_t scale_expert_stride, int output_columns, int reduction_size, cutlass_kernels::ActivationParams activation_params, - const int* permuted_row_to_source_row, int num_rows) { + const int* permuted_row_to_source_row, int num_rows, + const int* unpermuted_row_to_permuted_row = nullptr, const float* routing_scales = nullptr, + int experts_per_token = 0) { #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) constexpr int kKLanes = KLanes; constexpr int kBlockSize = 16; - constexpr int kColsPerBlock = 128 / kKLanes; + constexpr int kColsPerBlock = FusedFinalize ? 8 : 128 / kKLanes; + constexpr int kReductionUnroll = FusedSwiGlu ? 2 : 1; static_assert(kKLanes == 4 || kKLanes == 8); constexpr unsigned kFullMask = 0xffffffffu; using PackedT = std::conditional_t, half2, __nv_bfloat162>; - const int row = static_cast(blockIdx.x); - int expert = permuted_row_to_expert != nullptr ? permuted_row_to_expert[row] : 0; + const int row = FusedFinalize + ? unpermuted_row_to_permuted_row[threadIdx.y * num_rows + blockIdx.x] + : static_cast(blockIdx.x); + int expert = row >= 0 && permuted_row_to_expert != nullptr ? permuted_row_to_expert[row] : -1; + if constexpr (!FusedFinalize) { + if (permuted_row_to_expert == nullptr) expert = 0; + } #pragma unroll 1 for (int expert_index = 0; expert_index < num_experts && permuted_row_to_expert == nullptr; ++expert_index) { if (row >= static_cast(expert_first_token_offset[expert_index + 1])) { @@ -386,8 +396,10 @@ __global__ void MoeGemvFp4RawKPackedKernel( } break; } - if (expert < 0 || expert >= num_experts) { - return; + const bool valid_expert = expert >= 0 && expert < num_experts; + if (!valid_expert) { + if constexpr (!FusedFinalize) return; + expert = 0; } const int column = static_cast(blockIdx.y) * kColsPerBlock + threadIdx.x / kKLanes; @@ -403,7 +415,9 @@ __global__ void MoeGemvFp4RawKPackedKernel( } } const int k_lane = threadIdx.x % kKLanes; - const int source_row = permuted_row_to_source_row ? permuted_row_to_source_row[row] % num_rows : row; + const int source_row = valid_expert + ? (permuted_row_to_source_row ? permuted_row_to_source_row[row] % num_rows : row) + : 0; const T* row_act = act + static_cast(source_row) * reduction_size; const int second_source_row = second_row ? (permuted_row_to_source_row ? permuted_row_to_source_row[row + 1] % num_rows : row + 1) @@ -415,8 +429,9 @@ __global__ void MoeGemvFp4RawKPackedKernel( const T* expert_bias = EnableBias ? bias + static_cast(expert) * output_columns : nullptr; float accumulators[ReuseExpertRows ? 2 : 1][2] = {}; +#pragma unroll(kReductionUnroll) for (int reduction_column = k_lane * kBlockSize; - column < output_columns && reduction_column < reduction_size; + valid_expert && column < output_columns && reduction_column < reduction_size; reduction_column += kKLanes * kBlockSize) { uint2 packed = *reinterpret_cast( expert_weight + static_cast(column) * (reduction_size / 2) + reduction_column / 2); @@ -457,10 +472,36 @@ __global__ void MoeGemvFp4RawKPackedKernel( for (int offset = kKLanes / 2; offset > 0; offset /= 2) { accumulator += __shfl_xor_sync(kFullMask, accumulator, offset, kKLanes); } - const float paired_accumulator = __shfl_xor_sync(kFullMask, accumulator, kKLanes); - if (k_lane == 0 && column < output_columns && column % 2 == 0) { - StoreMoeGemvFp4Pair( - accumulator, paired_accumulator, expert_bias, out, row, column, output_columns, expert, activation_params); + if constexpr (FusedFinalize) { + __shared__ float expert_results[16][kColsPerBlock]; + __shared__ bool expert_valid[16]; + if (threadIdx.x == 0) expert_valid[threadIdx.y] = valid_expert; + if (k_lane == 0) { + if constexpr (EnableBias) { + if (valid_expert && column < output_columns) accumulator += static_cast(expert_bias[column]); + } + const T rounded = static_cast(accumulator); + expert_results[threadIdx.y][threadIdx.x / kKLanes] = static_cast(rounded); + } + __syncthreads(); + if (threadIdx.y == 0 && threadIdx.x < kColsPerBlock) { + const int output_column = static_cast(blockIdx.y) * kColsPerBlock + threadIdx.x; + if (output_column < output_columns) { + float sum = 0.0f; + for (int route = 0; route < experts_per_token; ++route) { + if (expert_valid[route]) { + sum += routing_scales[blockIdx.x * experts_per_token + route] * expert_results[route][threadIdx.x]; + } + } + out[static_cast(blockIdx.x) * output_columns + output_column] = static_cast(sum); + } + } + } else { + const float paired_accumulator = __shfl_xor_sync(kFullMask, accumulator, kKLanes); + if (k_lane == 0 && column < output_columns && column % 2 == 0) { + StoreMoeGemvFp4Pair( + accumulator, paired_accumulator, expert_bias, out, row, column, output_columns, expert, activation_params); + } } if constexpr (ReuseExpertRows) { if (second_row) { @@ -867,6 +908,43 @@ __global__ void __launch_bounds__(Warps * 32, MinBlocks) MoeGemvFp4DecodeKernel( #endif } +template +bool try_launch_moe_gemv_nvfp4_fused_finalize( + const T* act, const uint8_t* weight, const uint8_t* block_scales, const float* global_scales, + const T* bias, T* out, const int* unpermuted_row_to_permuted_row, + const int* permuted_row_to_expert, const float* routing_scales, int num_experts, + int num_rows, int experts_per_token, int n, int k, cudaStream_t stream) { + if (num_rows < 1 || num_rows > 2 || experts_per_token < 1 || experts_per_token > 16 || + n <= 0 || n % 2 != 0 || k <= 0 || k % 16 != 0) { + return false; + } + const dim3 grid(static_cast(num_rows), static_cast((n + 7) / 8)); + const dim3 block(64, static_cast(experts_per_token)); + const int64_t weight_stride = static_cast(n) * k / 2; + const int64_t scale_stride = static_cast(n) * (k / 16); + const cutlass_kernels::ActivationParams activation; + if (bias != nullptr) { + MoeGemvFp4RawKPackedKernel<<>>( + act, weight, block_scales, global_scales, bias, out, nullptr, permuted_row_to_expert, + num_experts, weight_stride, scale_stride, n, k, activation, nullptr, num_rows, + unpermuted_row_to_permuted_row, routing_scales, experts_per_token); + } else { + MoeGemvFp4RawKPackedKernel<<>>( + act, weight, block_scales, global_scales, bias, out, nullptr, permuted_row_to_expert, + num_experts, weight_stride, scale_stride, n, k, activation, nullptr, num_rows, + unpermuted_row_to_permuted_row, routing_scales, experts_per_token); + } + CUDA_CALL_THROW(cudaGetLastError()); + return true; +} + +template bool try_launch_moe_gemv_nvfp4_fused_finalize( + const half*, const uint8_t*, const uint8_t*, const float*, const half*, half*, + const int*, const int*, const float*, int, int, int, int, int, cudaStream_t); +template bool try_launch_moe_gemv_nvfp4_fused_finalize<__nv_bfloat16>( + const __nv_bfloat16*, const uint8_t*, const uint8_t*, const float*, const __nv_bfloat16*, __nv_bfloat16*, + const int*, const int*, const float*, int, int, int, int, int, cudaStream_t); + template void LaunchMoeGemvFp4RawNPacked( const T* act, const uint8_t* weight, const uint8_t* block_scales, const float* global_scales, @@ -879,10 +957,11 @@ void LaunchMoeGemvFp4RawNPacked( *decode_kernel_launched = false; } constexpr int kThreads = 128; - constexpr int kColsPerBlock = 16; + constexpr int kColsPerBlock = FusedSwiGlu ? 8 : 16; const int64_t weight_expert_stride = n * k / 2; const int64_t scale_expert_stride = n * (k / 16); - const dim3 grid(static_cast(expanded_num_rows), static_cast((n + kColsPerBlock - 1) / kColsPerBlock)); + const int columns_per_block = weights_row_major ? kColsPerBlock : 16; + const dim3 grid(static_cast(expanded_num_rows), static_cast((n + columns_per_block - 1) / columns_per_block)); if (weights_row_major) { static const bool enable_expert_reuse = onnxruntime::ParseEnvironmentVariableWithDefault("ORT_NVFP4_FC1_EXPERT_REUSE", 0) == 1; diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h index 4081168fd07cc..19d9385520fb1 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h @@ -124,6 +124,15 @@ void launch_moe_gemv_fp4_symmetric( bool raw_n_packed, cudaStream_t stream, bool weights_row_major = false, bool decode_scale_eligible = false, bool* decode_kernel_launched = nullptr); +// Native K-packed NVFP4 FC2 with top-k reduction; false leaves unsupported shapes to the caller. +// The caller must use single-EP/single-TP routing. Each expert result is rounded to T before reduction. +template +bool try_launch_moe_gemv_nvfp4_fused_finalize( + const T* act, const uint8_t* weight, const uint8_t* block_scales, const float* global_scales, + const T* bias, T* out, const int* unpermuted_row_to_permuted_row, + const int* permuted_row_to_expert, const float* routing_scales, int num_experts, + int num_rows, int experts_per_token, int n, int k, cudaStream_t stream); + // Launches the MXFP4 MoE GEMV and fuses interleaved SwiGLU activation. // weight/scales/bias use raw FC1 output width n = 2 * inter_size // out is post-activation [expanded_num_rows, inter_size] diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu index 43bd9170cc772..77e5c2f6e9783 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu @@ -1608,8 +1608,8 @@ __global__ void finalizeMoeRoutingKernel(const GemmOutputType* expanded_permuted constexpr int64_t FINALIZE_ELEM_PER_THREAD = 128 / std::min(sizeof_bits::value, sizeof_bits::value); assert(orig_cols % FINALIZE_ELEM_PER_THREAD == 0); - const int64_t start_offset = threadIdx.x; - const int64_t stride = FINALIZE_THREADS_PER_BLOCK; + const int64_t start_offset = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; + const int64_t stride = static_cast(blockDim.x) * gridDim.y; const int64_t num_elems_in_col = orig_cols / FINALIZE_ELEM_PER_THREAD; using BiasElem = cutlass::Array; @@ -1628,7 +1628,7 @@ __global__ void finalizeMoeRoutingKernel(const GemmOutputType* expanded_permuted __shared__ int s_expert_id[STAGE_ROUTING ? FINALIZE_MAX_STAGED_TOPK : 1]; __shared__ float s_row_scale[STAGE_ROUTING ? FINALIZE_MAX_STAGED_TOPK : 1]; if constexpr (STAGE_ROUTING) { - for (int64_t k_idx = threadIdx.x; k_idx < experts_per_token; k_idx += FINALIZE_THREADS_PER_BLOCK) { + for (int64_t k_idx = threadIdx.x; k_idx < experts_per_token; k_idx += blockDim.x) { int64_t const k_offset = original_row * experts_per_token + k_idx; int const expert_id = token_selected_experts[k_offset] - start_expert_id; bool const is_local_expert = expert_id >= 0 && expert_id < num_experts_per_node; @@ -1915,6 +1915,13 @@ void finalizeMoeRoutingKernelLauncher(const GemmOutputType* expanded_permuted_ro } #undef LAUNCH_FINALIZE_ONE_ROW } else if (experts_per_token <= FINALIZE_MAX_STAGED_TOPK) { + if (num_rows <= 2) { + constexpr int kDecodeThreads = 32; + config.blockDim = dim3(kDecodeThreads); + config.gridDim = dim3(static_cast(num_rows), + static_cast((cols / kFinalizeElemPerThread + kDecodeThreads - 1) / + kDecodeThreads)); + } auto func = final_scales ? &finalizeMoeRoutingKernel : &finalizeMoeRoutingKernel; diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cc b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cc index 28ef396fc8cbf..f98f3a7232f6c 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cc +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cc @@ -48,6 +48,16 @@ MatMulBlockQuantizedFp8Weight::MatMulBlockQuantizedFp8Weight(const OpKernelInfo& max_dequant_scratch_bytes_(DequantScratchLimitBytes()), enable_deep_gemm_(ParseEnvironmentVariableWithDefault("ORT_FP8_MATMUL_DEEPGEMM", false)) { ORT_ENFORCE(block_size_ > 0, "block_size must be positive."); + if (info.GetAttrs("split_sizes", split_sizes_).IsOK()) { + ORT_ENFORCE(!split_sizes_.empty() && split_sizes_.size() <= Fp8GemvOutputSplit::kMaxOutputs, + "split_sizes requires one to eight widths."); + int64_t total = 0; + for (const auto width : split_sizes_) { + ORT_ENFORCE(width > 0 && width <= INT64_MAX - total, + "split_sizes must contain positive widths with no sum overflow."); + total += width; + } + } } template @@ -94,27 +104,75 @@ Status MatMulBlockQuantizedFp8Weight::ComputeImpl(OpKernelContext* context) cons TensorShape b_logical_shape({n, k}); ORT_RETURN_IF_ERROR(helper.Compute(a_shape, b_logical_shape, transa, transb)); - Tensor* Y = context->Output(0, helper.OutputShape()); - if (Y->Shape().Size() == 0) { - return Status::OK(); - } - const int m_i = SafeInt(helper.M()); const int n_i = SafeInt(helper.N()); const int k_i = SafeInt(helper.K()); + const bool use_gemv = m_i > 0 && + m_i <= MatMulBlockScaledFp8GemvMaxM(k_i, SafeInt(block_size_), GetDeviceProp()) && + (k_i % 16 == 0) && (block_size_ % 16 == 0); + const bool direct_split = !split_sizes_.empty() && use_gemv && + MatMulBlockScaledFp8GemvSupportsSplitOutput( + m_i, k_i, SafeInt(block_size_), GetDeviceProp()); + Fp8GemvOutputSplit split; + void* y_data; + IAllocatorUniquePtr y_scratch; + const auto output_shape = helper.OutputShape(); + if (split_sizes_.empty()) { + ORT_ENFORCE(context->OutputCount() == 1, "Without split_sizes exactly one output is required."); + y_data = context->Output(0, output_shape)->MutableDataRaw(); + } else { + ORT_ENFORCE(context->OutputCount() == static_cast(split_sizes_.size()), + "Output count must match split_sizes."); + int64_t total = 0; + for (const auto width : split_sizes_) { + total += width; + } + ORT_ENFORCE(total == n, "split_sizes must sum to the weight's N dimension."); + split.count = static_cast(split_sizes_.size()); + for (int i = 0; i < split.count; ++i) { + TensorShapeVector shape(output_shape.GetDims().begin(), output_shape.GetDims().end()); + shape.back() = split_sizes_[i]; + split.outputs[i] = context->Output(i, TensorShape(shape))->MutableDataRaw(); + split.widths[i] = SafeInt(split_sizes_[i]); + } + if (direct_split) { + y_data = split.outputs[0]; + } else { + y_scratch = GetScratchBuffer(SafeInt(output_shape.Size()), GetComputeStream(context)); + y_data = y_scratch.get(); + } + } + if (output_shape.Size() == 0) { + return Status::OK(); + } + const auto finish_output = [&]() -> Status { + if (split.count != 0 && !direct_split) { + size_t offset = 0; + for (int i = 0; i < split.count; ++i) { + const size_t width_bytes = static_cast(split.widths[i]) * sizeof(CudaT); + CUDA_RETURN_IF_ERROR(cudaMemcpy2DAsync( + split.outputs[i], width_bytes, static_cast(y_data) + offset, + static_cast(n_i) * sizeof(CudaT), width_bytes, m_i, + cudaMemcpyDeviceToDevice, Stream(context))); + offset += width_bytes; + } + } + return Status::OK(); + }; + if (k_i == 0) { - CUDA_RETURN_IF_ERROR(cudaMemsetAsync(Y->MutableDataRaw(), 0, Y->SizeInBytes(), Stream(context))); + CUDA_RETURN_IF_ERROR(cudaMemsetAsync(y_data, 0, SafeInt(output_shape.Size()) * sizeof(CudaT), Stream(context))); if (bias != nullptr) { ORT_RETURN_IF_ERROR(LaunchAddBiasBlockScaledFp8( - Y->MutableDataRaw(), + y_data, bias->DataRaw(), m_i, n_i, std::is_same::value, Stream(context))); } - return Status::OK(); + return finish_output(); } // Optional W8A8 activation path: statically quantize A to FP8 E4M3 and dequantize back so the @@ -123,9 +181,6 @@ Status MatMulBlockQuantizedFp8Weight::ComputeImpl(OpKernelContext* context) cons // // The decode GEMV absorbs this round trip in-register, so the standalone kernel (and its [M, K] // scratch buffer) is only materialized for the cuBLAS path below. - const bool use_gemv = m_i > 0 && - m_i <= MatMulBlockScaledFp8GemvMaxM(k_i, SafeInt(block_size_), GetDeviceProp()) && - (k_i % 16 == 0) && (block_size_ % 16 == 0); const bool fuse_act_qdq = use_gemv && !FusedFp8ActivationQdqDisabled(); #if defined(USE_DEEP_GEMM) @@ -154,14 +209,14 @@ Status MatMulBlockQuantizedFp8Weight::ComputeImpl(OpKernelContext* context) cons const size_t weight_offset = SafeInt(offset) * k_i; const size_t scale_offset = SafeInt(offset) * (k_i / 128); ORT_RETURN_IF_ERROR(LaunchMatMulFp8DeepGemm( - Y->MutableData() + offset, a_quant.get(), a_scales.get(), + static_cast(y_data) + offset, a_quant.get(), a_scales.get(), static_cast(b->DataRaw()) + weight_offset, b_scale->Data() + scale_offset, bias != nullptr ? bias->Data() + offset : nullptr, packed_b_scales.get(), accum.get(), m_i, rows, k_i, aligned_m, n_i, std::is_same::value, GetDeviceProp().multiProcessorCount, Stream(context))); } - return Status::OK(); + return finish_output(); } } #endif @@ -186,8 +241,8 @@ Status MatMulBlockQuantizedFp8Weight::ComputeImpl(OpKernelContext* context) cons // A fused warp-per-column kernel reads the FP8 weight directly, avoiding both the [N, K] // dequant scratch buffer and the cuBLAS GEMM (which is underutilized at M == 1). if (use_gemv) { - return LaunchMatMulBlockScaledFp8Gemv( - Y->MutableDataRaw(), + ORT_RETURN_IF_ERROR(LaunchMatMulBlockScaledFp8Gemv( + y_data, a_ptr, b->DataRaw(), b_scale->Data(), @@ -199,7 +254,8 @@ Status MatMulBlockQuantizedFp8Weight::ComputeImpl(OpKernelContext* context) cons SafeInt(block_size_), std::is_same::value, GetDeviceProp(), - Stream(context)); + Stream(context), direct_split ? split : Fp8GemvOutputSplit{})); + return finish_output(); } // Dequantize the FP8 weight into a scratch buffer of the activation type, then GEMM. The scratch @@ -243,7 +299,7 @@ Status MatMulBlockQuantizedFp8Weight::ComputeImpl(OpKernelContext* context) cons reinterpret_cast(a_ptr), helper.Lda(transa), &zero, - reinterpret_cast(Y->MutableDataRaw()) + n_offset, + static_cast(y_data) + n_offset, helper.Ldc(), GetDeviceProp(), UseTF32())); @@ -251,7 +307,7 @@ Status MatMulBlockQuantizedFp8Weight::ComputeImpl(OpKernelContext* context) cons if (bias != nullptr) { ORT_RETURN_IF_ERROR(LaunchAddBiasBlockScaledFp8( - Y->MutableDataRaw(), + y_data, bias->DataRaw(), m_i, n_i, @@ -259,7 +315,7 @@ Status MatMulBlockQuantizedFp8Weight::ComputeImpl(OpKernelContext* context) cons Stream(context))); } - return Status::OK(); + return finish_output(); #endif } diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu index d045519626b27..16ae3c0d5c949 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu @@ -514,7 +514,20 @@ struct Fp8GemvMma<__nv_bfloat16> { } }; -template +template +__device__ __forceinline__ T* Fp8OutputColumn(T* output, int col, int& stride, + const Fp8GemvOutputSplit& split) { + for (int i = 0; i < split.count; ++i) { + if (col < split.widths[i]) { + stride = split.widths[i]; + return static_cast(split.outputs[i]) + col; + } + col -= split.widths[i]; + } + return output + col; +} + +template __device__ __forceinline__ void Fp8MmaGemvBody(AType* __restrict__ output, const AType* __restrict__ input_a, const __nv_fp8_e4m3* __restrict__ input_b, @@ -525,7 +538,8 @@ __device__ __forceinline__ void Fp8MmaGemvBody(AType* __restrict__ output, int n, int k, int block_size, - int k_blocks) { + int k_blocks, + const Fp8GemvOutputSplit& split) { using Mma = Fp8GemvMma; const bool act_qdq = act_scale != nullptr; @@ -560,6 +574,7 @@ __device__ __forceinline__ void Fp8MmaGemvBody(AType* __restrict__ output, float accb[MTiles][4] = {}; // unscaled, within the current K block int cur_kb = -1; +#pragma unroll 2 for (int wi = warp; wi < windows; wi += KSplit) { const int k0 = wi << 6; const int kb = k0 / block_size; @@ -670,23 +685,31 @@ __device__ __forceinline__ void Fp8MmaGemvBody(AType* __restrict__ output, // M tile. const float bias_lo = (bias != nullptr && lo_ok) ? to_float(bias[col_lo]) : 0.f; const float bias_hi = (bias != nullptr && hi_ok) ? to_float(bias[col_hi]) : 0.f; + int stride_lo = n; + int stride_hi = n; + AType* output_lo = output + col_lo; + AType* output_hi = output + col_hi; + if constexpr (SplitOutput) { + output_lo = lo_ok ? Fp8OutputColumn(output, col_lo, stride_lo, split) : nullptr; + output_hi = hi_ok ? Fp8OutputColumn(output, col_hi, stride_hi, split) : nullptr; + } #pragma unroll for (int mt = 0; mt < MTiles; ++mt) { const int row = (t << 1) + (mt << 3); if (row < m) { if (lo_ok) { - output[static_cast(row) * n + col_lo] = from_float(acc[mt][0] + bias_lo); + output_lo[static_cast(row) * stride_lo] = from_float(acc[mt][0] + bias_lo); } if (hi_ok) { - output[static_cast(row) * n + col_hi] = from_float(acc[mt][2] + bias_hi); + output_hi[static_cast(row) * stride_hi] = from_float(acc[mt][2] + bias_hi); } } if (row + 1 < m) { if (lo_ok) { - output[static_cast(row + 1) * n + col_lo] = from_float(acc[mt][1] + bias_lo); + output_lo[static_cast(row + 1) * stride_lo] = from_float(acc[mt][1] + bias_lo); } if (hi_ok) { - output[static_cast(row + 1) * n + col_hi] = from_float(acc[mt][3] + bias_hi); + output_hi[static_cast(row + 1) * stride_hi] = from_float(acc[mt][3] + bias_hi); } } } @@ -695,6 +718,12 @@ __device__ __forceinline__ void Fp8MmaGemvBody(AType* __restrict__ output, // Two entry points over one body. The pinned one carries a residency hint; see // `Fp8MmaGemvPinsResidency` for when the launcher picks it and why the plain one has to stay. // clang-format off +#if CUDA_VERSION >= 12000 +// Avoid a per-thread local copy when the inlined store helper references this descriptor. +#define ORT_FP8_SPLIT_PARAM const __grid_constant__ Fp8GemvOutputSplit split +#else +#define ORT_FP8_SPLIT_PARAM const Fp8GemvOutputSplit split +#endif #define ORT_FP8_MMA_GEMV_PARAMS \ AType* __restrict__ output, \ const AType* __restrict__ input_a, \ @@ -702,24 +731,25 @@ __device__ __forceinline__ void Fp8MmaGemvBody(AType* __restrict__ output, const float* __restrict__ weight_scale, \ const AType* __restrict__ bias, \ const float* __restrict__ act_scale, \ - int m, int n, int k, int block_size, int k_blocks + int m, int n, int k, int block_size, int k_blocks, ORT_FP8_SPLIT_PARAM #define ORT_FP8_MMA_GEMV_ARGS \ - output, input_a, input_b, weight_scale, bias, act_scale, m, n, k, block_size, k_blocks + output, input_a, input_b, weight_scale, bias, act_scale, m, n, k, block_size, k_blocks, split // clang-format on -template +template __global__ void MatMulBlockScaledFp8MmaGemvKernel(ORT_FP8_MMA_GEMV_PARAMS) { - Fp8MmaGemvBody(ORT_FP8_MMA_GEMV_ARGS); + Fp8MmaGemvBody(ORT_FP8_MMA_GEMV_ARGS); } -template +template __global__ __launch_bounds__(32 * KSplit, 3) void MatMulBlockScaledFp8MmaGemvKernelPinned(ORT_FP8_MMA_GEMV_PARAMS) { - Fp8MmaGemvBody(ORT_FP8_MMA_GEMV_ARGS); + Fp8MmaGemvBody(ORT_FP8_MMA_GEMV_ARGS); } #undef ORT_FP8_MMA_GEMV_ARGS #undef ORT_FP8_MMA_GEMV_PARAMS +#undef ORT_FP8_SPLIT_PARAM // Kill switch for A/B testing the tensor-core path against the FMA path in the same binary. bool Fp8GemvMmaEnabled() { @@ -915,6 +945,20 @@ bool Fp8MmaGb10TuningEnabled() { return enabled; } +bool MatMulBlockScaledFp8GemvSupportsSplitOutput(int m, int k, int block_size, + const cudaDeviceProp& device_prop) { +#if !defined(DISABLE_FLOAT8_TYPES) && defined(CUDA_VERSION) && CUDA_VERSION >= 11080 + return device_prop.major >= 8 && m > 0 && m <= kFp8MmaGemvTileM && + k >= 256 && k % 64 == 0 && block_size % 64 == 0 && Fp8GemvMmaEnabled(); +#else + ORT_UNUSED_PARAMETER(m); + ORT_UNUSED_PARAMETER(k); + ORT_UNUSED_PARAMETER(block_size); + ORT_UNUSED_PARAMETER(device_prop); + return false; +#endif +} + static Status LaunchMatMulBlockScaledFp8GemvImpl(void* y, const void* a, const void* b_fp8, @@ -928,7 +972,8 @@ static Status LaunchMatMulBlockScaledFp8GemvImpl(void* y, bool is_bf16, const cudaDeviceProp& device_prop, cudaStream_t stream, - bool enable_gb10_ksplit32) { + bool enable_gb10_ksplit32, + const Fp8GemvOutputSplit& split) { #if !defined(DISABLE_FLOAT8_TYPES) && defined(CUDA_VERSION) && CUDA_VERSION >= 11080 if (m <= 0 || n <= 0 || k <= 0) { return Status::OK(); @@ -941,6 +986,19 @@ static Status LaunchMatMulBlockScaledFp8GemvImpl(void* y, "MatMulBlockQuantizedFp8Weight GEMV requires block_size divisible by 16, got ", block_size, "."); const int k_blocks = (k + block_size - 1) / block_size; const auto* b = reinterpret_cast(b_fp8); + ORT_RETURN_IF_NOT(split.count >= 0 && split.count <= Fp8GemvOutputSplit::kMaxOutputs, + "FP8 GEMV split output count must be between zero and eight."); + if (split.count != 0) { + int64_t total = 0; + for (int i = 0; i < split.count; ++i) { + ORT_RETURN_IF_NOT(split.widths[i] > 0 && split.outputs[i] != nullptr, + "FP8 GEMV split outputs require positive widths and non-null buffers."); + total += split.widths[i]; + } + ORT_RETURN_IF_NOT(total == n, "FP8 GEMV split widths must sum to N."); + } + ORT_RETURN_IF_NOT(split.count == 0 || MatMulBlockScaledFp8GemvSupportsSplitOutput(m, k, block_size, device_prop), + "Direct split output requires the tensor-core GEMV path with M <= 32."); if (m > kFp8MmaGemvTileM) { ORT_RETURN_IF_NOT(device_prop.major >= 8 && k % 64 == 0 && k >= 256 && @@ -950,12 +1008,12 @@ static Status LaunchMatMulBlockScaledFp8GemvImpl(void* y, const size_t element_size = is_bf16 ? sizeof(__nv_bfloat16) : sizeof(half); ORT_RETURN_IF_ERROR(LaunchMatMulBlockScaledFp8GemvImpl( y, a, b_fp8, weight_scale, bias, act_scale, kFp8MmaGemvTileM, n, k, block_size, - is_bf16, device_prop, stream, false)); + is_bf16, device_prop, stream, false, {})); return LaunchMatMulBlockScaledFp8GemvImpl( static_cast(y) + static_cast(kFp8MmaGemvTileM) * n * element_size, static_cast(a) + static_cast(kFp8MmaGemvTileM) * k * element_size, b_fp8, weight_scale, bias, act_scale, m - kFp8MmaGemvTileM, n, k, block_size, - is_bf16, device_prop, stream, false); + is_bf16, device_prop, stream, false, {}); } // Tensor-core path (SM80+). Beats the FMA kernel at every M on H200: 1.06-1.23x at M == 1 and @@ -982,28 +1040,35 @@ static Status LaunchMatMulBlockScaledFp8GemvImpl(void* y, n, k_split, mtiles, device_prop.multiProcessorCount, device_prop.major, device_prop.minor); const auto launch_mma = [&]() { const dim3 mma_threads{32, KSplit}; + const auto launch_variant = [&]() { #define ORT_FP8_LAUNCH_MMA(kernel_name) \ do { \ if (is_bf16) { \ - kernel_name<<>>( \ + kernel_name<<>>( \ reinterpret_cast<__nv_bfloat16*>(y), reinterpret_cast(a), b, \ weight_scale, reinterpret_cast(bias), act_scale, m, n, k, \ - block_size, k_blocks); \ + block_size, k_blocks, split); \ } else { \ - kernel_name<<>>( \ + kernel_name<<>>( \ reinterpret_cast(y), reinterpret_cast(a), b, \ weight_scale, reinterpret_cast(bias), act_scale, m, n, k, \ - block_size, k_blocks); \ + block_size, k_blocks, split); \ } \ } while (0) - if constexpr (KSplit == 16 && MTiles == 1) { - if (pin_residency) { - ORT_FP8_LAUNCH_MMA(MatMulBlockScaledFp8MmaGemvKernelPinned); - return; + if constexpr (KSplit == 16 && MTiles == 1) { + if (pin_residency) { + ORT_FP8_LAUNCH_MMA(MatMulBlockScaledFp8MmaGemvKernelPinned); + return; + } } - } - ORT_FP8_LAUNCH_MMA(MatMulBlockScaledFp8MmaGemvKernel); + ORT_FP8_LAUNCH_MMA(MatMulBlockScaledFp8MmaGemvKernel); #undef ORT_FP8_LAUNCH_MMA + }; + if (split.count) { + launch_variant.template operator()(); + } else { + launch_variant.template operator()(); + } }; // Only 1, 2 and 4 row tiles are instantiated; an M of 17..24 rounds up to 4 and masks the // remainder, which costs nothing next to the weight traffic it shares. @@ -1109,6 +1174,7 @@ static Status LaunchMatMulBlockScaledFp8GemvImpl(void* y, ORT_UNUSED_PARAMETER(device_prop); ORT_UNUSED_PARAMETER(stream); ORT_UNUSED_PARAMETER(enable_gb10_ksplit32); + ORT_UNUSED_PARAMETER(split); return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "MatMulBlockQuantizedFp8Weight requires CUDA 11.8 or later."); #endif } @@ -1125,10 +1191,11 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y, int block_size, bool is_bf16, const cudaDeviceProp& device_prop, - cudaStream_t stream) { + cudaStream_t stream, + const Fp8GemvOutputSplit& split) { return LaunchMatMulBlockScaledFp8GemvImpl( y, a, b_fp8, weight_scale, bias, act_scale, m, n, k, block_size, - is_bf16, device_prop, stream, true); + is_bf16, device_prop, stream, true, split); } } // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h index 7128b86971dcc..b54a37da08a87 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h @@ -4,6 +4,7 @@ #pragma once #include +#include #include "core/providers/cuda/cuda_kernel.h" @@ -31,10 +32,21 @@ class MatMulBlockQuantizedFp8Weight final : public onnxruntime::cuda::CudaKernel Status ComputeImpl(OpKernelContext* context) const; int64_t block_size_; + std::vector split_sizes_; size_t max_dequant_scratch_bytes_; bool enable_deep_gemm_; }; +struct Fp8GemvOutputSplit { + static constexpr int kMaxOutputs = 8; + void* outputs[kMaxOutputs]{}; + int widths[kMaxOutputs]{}; + int count = 0; +}; + +bool MatMulBlockScaledFp8GemvSupportsSplitOutput(int m, int k, int block_size, + const cudaDeviceProp& device_prop); + #if defined(USE_DEEP_GEMM) Status LaunchPrepareMatMulFp8DeepGemm(void* a_quant, float* a_scales, const void* a, const float* a_scale, int m, int k, int aligned_m, @@ -83,6 +95,8 @@ Status LaunchQuantizeDequantizeActivationFp8(void* a_out, // directly (no [N, K] dequant buffer). a is [M, K] activation (FP16/BF16), b_fp8 is [N, K] // FP8 E4M3, weight_scale is [N, ceil(K/block_size)] fp32, bias is an optional [N] vector (may be // null). Output y is [M, N] in the activation type. Requires k % 16 == 0 and block_size % 16 == 0. +// A nonempty split descriptor routes tensor-core stores into consecutive column partitions; +// callers must check MatMulBlockScaledFp8GemvSupportsSplitOutput before using that mode. // device_prop selects the tensor-core (mma.m16n8k16) variant, which needs SM80+. // act_scale is an optional device fp32 scalar; when non-null the kernel applies the W8A8 // activation quantize/dequantize inline, bit-identically to LaunchQuantizeDequantizeActivationFp8, @@ -100,7 +114,8 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y, int block_size, bool is_bf16, const cudaDeviceProp& device_prop, - cudaStream_t stream); + cudaStream_t stream, + const Fp8GemvOutputSplit& split = {}); // Largest M selected for GEMV dispatch. The tensor-core default is conservative because the // crossover with dequantize + cuBLAS depends on the matrix shape; ORT_FP8_GEMV_MAX_M can raise the diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc index 7ab59ba44abd2..fb4110492546b 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc +++ b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc @@ -1496,11 +1496,8 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { std::array runner_tile_configs{}; size_t runner_tile_config_count = 0; size_t workspace_size = 0; - // The packed INT GEMV path allocates its own scratch below and returns before reaching the - // dense grouped-GEMM tile loop, so profiling/sizing the dense runner here would be pure - // overhead (mutex, two dense-tactic profiling launches, and an unused large workspace - // allocation). Skip it entirely for that path; workspace_size stays 0. - if (!use_packed_int && !use_fp8_fused) { + // Standalone projection paths return before runMoe and do not consume its tactics or workspace. + if (!use_packed_int && !use_fp8_fused && !use_nvfp4_packed_prefill && !run_fp4_gemv) { std::lock_guard profiler_lock(mGemmProfilerMutex); // Profiling launches grouped-GEMM kernels, records/synchronizes CUDA events, and @@ -2465,7 +2462,10 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { IAllocatorUniquePtr blocked_expert_counts; IAllocatorUniquePtr blocked_expert_counts_cumsum; IAllocatorUniquePtr blocked_row_to_unpermuted_row; - if (!ck::fusedBuildExpertMapsSortFirstToken( + if (!TryLaunchSmallDecodeExpertMaps( + expert_indices, p_r2u, unpermuted_row_to_permuted_row, p_exp, p_efto, + static_cast(num_rows), num_experts, static_cast(k_), stream) && + !ck::fusedBuildExpertMapsSortFirstToken( expert_indices, p_r2u, unpermuted_row_to_permuted_row, p_exp, p_efto, num_rows, num_experts, static_cast(k_), 0, num_experts, stream)) { const int64_t num_tokens_per_block = ck::computeNumTokensPerBlock(num_rows, num_experts); @@ -2637,6 +2637,16 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { } launch_fc1(fc1_config); + if (use_raw_nvfp4_gemv && nvfp4_weights_row_major_ && + parallelism_config.tp_size == 1 && parallelism_config.ep_size == 1 && + gemv::try_launch_moe_gemv_nvfp4_fused_finalize( + static_cast(p_fc1_buf.get()), gemv_fc2_weight, gemv_fc2_raw_block_scales, + gemv_fc2_raw_global_scales, static_cast(fc2_bias), + static_cast(output->MutableDataRaw()), unpermuted_row_to_permuted_row, + p_exp, expert_scales, num_experts, static_cast(num_rows), static_cast(k_), + static_cast(hidden), static_cast(inter), stream)) { + return; + } launch_fc2(fc2_config); ck::finalizeMoeRoutingKernelLauncher( static_cast(p_fc2_buf.get()), static_cast(output->MutableDataRaw()), diff --git a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu index 3c91fe4b7a967..3b2aa8f23a365 100644 --- a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu +++ b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu @@ -223,6 +223,42 @@ __global__ void SoftmaxTopKMergeKernel(const T* logits, float* topk_scales, int* } } +__global__ void SmallDecodeExpertMapsKernel( + const int* indices, int* row_to_source, int* source_to_row, int* sorted_experts, + int64_t* expert_offsets, int rows, int experts, int k) { + __shared__ int selected[32]; + const int expanded = rows * k; + if (threadIdx.x < expanded) selected[threadIdx.x] = indices[threadIdx.x]; + __syncthreads(); + if (threadIdx.x < expanded) { + const int original = threadIdx.x; + const int expert = selected[original]; + int rank = 0; + for (int i = 0; i < expanded; ++i) { + rank += selected[i] < expert || (selected[i] == expert && i < original); + } + const int source = (original % k) * rows + original / k; + row_to_source[rank] = source; + source_to_row[source] = rank; + sorted_experts[rank] = expert; + } + for (int expert = threadIdx.x; expert <= experts; expert += blockDim.x) { + int count = 0; + for (int i = 0; i < expanded; ++i) count += selected[i] < expert; + expert_offsets[expert] = count; + } +} + +bool TryLaunchSmallDecodeExpertMaps( + const int* indices, int* row_to_source, int* source_to_row, int* sorted_experts, + int64_t* expert_offsets, int rows, int experts, int k, cudaStream_t stream) { + if (rows < 1 || rows > 2 || experts < 1 || experts > 512 || k < 1 || k > 16 || k > experts) return false; + SmallDecodeExpertMapsKernel<<<1, 128, 0, stream>>>( + indices, row_to_source, source_to_row, sorted_experts, expert_offsets, rows, experts, k); + CUDA_CALL_THROW(cudaGetLastError()); + return true; +} + // Warp-bitonic softmax + top-k for num_experts <= 32. Each warp handles one // row, with lane `l` owning expert `l`. The whole softmax reduction and the // sort are done with warp shuffles (no shared memory). This is the fastest path @@ -1547,8 +1583,15 @@ void LaunchQMoECompactExperts( CUDA_CALL_THROW(cudaGetLastError()); } -template -__device__ __forceinline__ T QMoEFp8Weight(const QMoEFp8ProjectionParams& params, int expert, int row, int col) { +struct QMoEFp8WeightRow { + const uint8_t* weights; + const void* scales; + int scale_type; + int64_t scale_offset; +}; + +__device__ __forceinline__ QMoEFp8WeightRow ResolveQMoEFp8WeightRow( + const QMoEFp8ProjectionParams& params, int expert, int row) { const uint8_t* weights = params.weights; const void* scales = params.scales; int scale_type = params.scale_type; @@ -1566,13 +1609,23 @@ __device__ __forceinline__ T QMoEFp8Weight(const QMoEFp8ProjectionParams& params } const int scale_n = (source_n + params.block_size - 1) / params.block_size; const int scale_k = (params.k + params.block_size - 1) / params.block_size; - const int64_t scale_index = (static_cast(expert) * scale_n + row / params.block_size) * scale_k + - col / params.block_size; - const int64_t weight_index = (static_cast(expert) * source_n + row) * params.k + col; - const float scale = scale_type == 0 ? static_cast(scales)[scale_index] - : scale_type == 1 ? __half2float(static_cast(scales)[scale_index]) - : __bfloat162float(static_cast(scales)[scale_index]); - return static_cast(DecodeFloat8E4M3FN(weights[weight_index]) * scale); + const int64_t scale_offset = (static_cast(expert) * scale_n + row / params.block_size) * scale_k; + const int64_t weight_offset = (static_cast(expert) * source_n + row) * params.k; + return {weights + weight_offset, scales, scale_type, scale_offset}; +} + +__device__ __forceinline__ float QMoEFp8RowScale(const QMoEFp8WeightRow& row, int block) { + const int64_t index = row.scale_offset + block; + return row.scale_type == 0 ? static_cast(row.scales)[index] + : row.scale_type == 1 ? __half2float(static_cast(row.scales)[index]) + : __bfloat162float(static_cast(row.scales)[index]); +} + +template +__device__ __forceinline__ T QMoEFp8Weight(const QMoEFp8ProjectionParams& params, int expert, int row, int col) { + const auto weight_row = ResolveQMoEFp8WeightRow(params, expert, row); + const float scale = QMoEFp8RowScale(weight_row, col / params.block_size); + return static_cast(DecodeFloat8E4M3FN(weight_row.weights[col]) * scale); } template @@ -1584,10 +1637,27 @@ __global__ void QMoEFp8GemvKernel(QMoEFp8ProjectionParams params, const T* input return; } const int input_row = params.row_to_unpermuted ? params.row_to_unpermuted[row] % params.num_rows : row; + const auto weight_row = ResolveQMoEFp8WeightRow(params, expert, output_col); + const T* activation = input + static_cast(input_row) * params.k; float sum = 0.0f; - for (int col = threadIdx.x % 32; col < params.k; col += 32) { - sum = fmaf(static_cast(input[static_cast(input_row) * params.k + col]), - static_cast(QMoEFp8Weight(params, expert, output_col, col)), sum); + if (params.block_size == 128) { + for (int block_begin = 0; block_begin < params.k; block_begin += 128) { + const float scale = QMoEFp8RowScale(weight_row, block_begin / 128); +#pragma unroll + for (int offset = 0; offset < 128; offset += 32) { + const int col = block_begin + threadIdx.x % 32 + offset; + if (col < params.k) { + const T weight = static_cast(DecodeFloat8E4M3FN(weight_row.weights[col]) * scale); + sum = fmaf(static_cast(activation[col]), static_cast(weight), sum); + } + } + } + } else { + for (int col = threadIdx.x % 32; col < params.k; col += 32) { + const T weight = static_cast(DecodeFloat8E4M3FN(weight_row.weights[col]) * + QMoEFp8RowScale(weight_row, col / params.block_size)); + sum = fmaf(static_cast(activation[col]), static_cast(weight), sum); + } } sum = WarpReduceSum(sum); if (threadIdx.x % 32 == 0) { @@ -1982,6 +2052,16 @@ void LaunchQMoEDequantizeNvfp4WeightsImpl( bool weights_row_major) { ORT_ENFORCE(bias == nullptr || output_bias != nullptr, "QMoE NVFP4 bias gathering requires an output buffer."); constexpr int block = 256; + const int64_t row_major_grid_y = + (static_cast(k) + block * kQMoEDequantizeFp4VecK - 1) / + (block * kQMoEDequantizeFp4VecK); + if (weights_row_major && k % 16 == 0 && row_major_grid_y <= 65535 && num_experts <= 65535) { + const dim3 grid(n, onnxruntime::narrow(row_major_grid_y), num_experts); + QMoEDequantizeNvfp4RowMajorWeightsVecKernel<<>>( + packed_weights, block_scales, global_scales, output, n, k, compact_to_expert, bias, output_bias); + CUDA_CALL_THROW(cudaGetLastError()); + return; + } if (!weights_row_major && !compact_to_expert && !bias && QMoEDequantizeFp4VecApplies<16>(num_experts, n, k)) { const dim3 tile_block(kQMoEDequantizeFp4TileK / kQMoEDequantizeFp4VecK, kQMoEDequantizeFp4TileN); const dim3 tile_grid((n + kQMoEDequantizeFp4TileN - 1) / kQMoEDequantizeFp4TileN, diff --git a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h index 00a49bddbe72d..c28cc343f4ccf 100644 --- a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h +++ b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h @@ -12,6 +12,10 @@ namespace onnxruntime { namespace contrib { namespace cuda { +bool TryLaunchSmallDecodeExpertMaps( + const int* indices, int* row_to_source, int* source_to_row, int* sorted_experts, + int64_t* expert_offsets, int rows, int experts, int k, cudaStream_t stream); + void LaunchSoftmaxTopK( const float* logits, float* topk_scales, diff --git a/onnxruntime/contrib_ops/cuda/quantization/quantize_dequantize_linear.cc b/onnxruntime/contrib_ops/cuda/quantization/quantize_dequantize_linear.cc index f612a9b6a6f54..40f54aba301cf 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/quantize_dequantize_linear.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/quantize_dequantize_linear.cc @@ -34,7 +34,8 @@ REGISTER_Q_KERNEL_TYPED(uint8_t, MLFloat16) kCudaExecutionProvider, \ (*KernelDefBuilder::Create()) \ .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("T2", DataTypeImpl::GetTensorType()), \ + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("TS", DataTypeImpl::GetTensorType()), \ DequantizeLinear); REGISTER_DQ_KERNEL_TYPED(int8_t, MLFloat16) diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index c38c4d1e1936e..477b8b2b6dd69 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -3191,9 +3191,13 @@ are `A_scaled = fp8_e4m3(A / a_scale) * a_scale` (W8A8). Otherwise, A retains it precision (weight-only W8A16). The operator multiplies the activation by the transpose of B_scaled and adds the optional bias. -The output has shape [..., N] and the same element type as A.)DOC") +The output has shape [..., N] and the same element type as A. When split_sizes is provided, +return up to eight consecutive last-dimension partitions instead, equivalent to Split(Y, split_sizes). +The partition widths must be positive and sum to N.)DOC") .Attr("block_size", "Number of consecutive K values that share one weight scale. Default 128.", AttributeProto::INT, static_cast(128)) + .Attr("split_sizes", "Optional last-dimension output widths (one to eight positive sizes).", + AttributeProto::INTS, OPTIONAL_VALUE) .Input(0, "A", "Row-major FP16/BF16 activation of shape [..., K].", "T") .Input(1, "B", "Row-major FP8 E4M3 weight of shape [N, K].", "T1") .Input(2, "b_scale", "Per-block FP32 weight scales of shape [N, ceil(K / block_size)].", "T2") @@ -3203,13 +3207,31 @@ The output has shape [..., N] and the same element type as A.)DOC") "its FP16/BF16 precision.", "T2", OpSchema::Optional) .Input(4, "bias", "Optional bias of shape [N].", "T", OpSchema::Optional) - .Output(0, "Y", "Output of shape [..., N] in the activation type.", "T") + .Output(0, "Y", "Output(s) of shape [..., N] or [..., split_sizes[i]] in the activation type.", + "T", OpSchema::Variadic) .TypeConstraint("T", {"tensor(float16)", "tensor(bfloat16)"}, "Constrain activation, bias and output to FP16 or BF16.") .TypeConstraint("T1", {"tensor(float8e4m3fn)"}, "Constrain weight to FP8 E4M3.") .TypeConstraint("T2", {"tensor(float)"}, "Constrain scales to FP32.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - propagateElemTypeFromInputToOutput(ctx, 0, 0); + const auto* split = ctx.getAttribute("split_sizes"); + if ((split == nullptr && ctx.getNumOutputs() != 1) || + (split != nullptr && (split->ints_size() < 1 || split->ints_size() > 8 || + static_cast(split->ints_size()) != ctx.getNumOutputs()))) { + fail_shape_inference("Output count must match split_sizes, or be one when split_sizes is absent."); + } + int64_t split_total = 0; + if (split != nullptr) { + for (const auto width : split->ints()) { + if (width <= 0 || width > INT64_MAX - split_total) { + fail_shape_inference("split_sizes must contain positive widths with no sum overflow."); + } + split_total += width; + } + } + for (size_t i = 0; i < ctx.getNumOutputs(); ++i) { + propagateElemTypeFromInputToOutput(ctx, 0, i); + } if (!hasNInputShapes(ctx, 2)) { return; } @@ -3227,7 +3249,16 @@ The output has shape [..., N] and the same element type as A.)DOC") *output_shape.add_dim() = a_shape.dim(i); } *output_shape.add_dim() = b_shape.dim(0); - updateOutputShape(ctx, 0, output_shape); + if (split != nullptr && b_shape.dim(0).has_dim_value() && + split_total != b_shape.dim(0).dim_value()) { + fail_shape_inference("split_sizes must sum to the weight's N dimension."); + } + for (size_t i = 0; i < ctx.getNumOutputs(); ++i) { + if (split != nullptr) { + output_shape.mutable_dim(output_shape.dim_size() - 1)->set_dim_value(split->ints(static_cast(i))); + } + updateOutputShape(ctx, i, output_shape); + } })); static void MatmulWithQuantWeightShapeInference(ONNX_NAMESPACE::InferenceContext& ctx, diff --git a/onnxruntime/core/graph/contrib_ops/quantization_defs.cc b/onnxruntime/core/graph/contrib_ops/quantization_defs.cc index 70cd688832b02..46bfebf9ec693 100644 --- a/onnxruntime/core/graph/contrib_ops/quantization_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/quantization_defs.cc @@ -182,12 +182,16 @@ ONNX_MS_OPERATOR_SET_SCHEMA( })); static const char* DequantizeLinear_ver1_doc = R"DOC( -The linear dequantization operator. It consumes a quantized data, a scale, a zero point and computes the full precision data. -The dequantization formula is y = (x - x_zero_point) * x_scale. -Scale and zero point must have same shape. They must be either scalar (per tensor) or 1-D tensor (per 'axis').)DOC"; +The linear dequantization operator. It consumes quantized data, a scale, and an optional zero point. +The dequantization formula is y = (x - x_zero_point) * x_scale. Scale and zero point are scalar +or per-axis vectors. FP8 E4M3FN inputs can be dequantized directly to BF16 with FP32 scales.)DOC"; ONNX_MS_OPERATOR_SET_SCHEMA(DequantizeLinear, 1, OpSchema() + .Attr("output_dtype", + "Optional output element type. When omitted, the scale " + "element type determines the output type.", + AttributeProto::INT, false) .Attr("axis", "The axis along which same quantization parameters are applied. It's optional." "If it's not specified, it means per-tensor quantization and input 'x_scale' and " @@ -199,25 +203,36 @@ ONNX_MS_OPERATOR_SET_SCHEMA(DequantizeLinear, 1, .Input(1, "x_scale", "Scale for input 'x'. It can be a scalar, which means a per-tensor/layer " "dequantization, or a 1-D tensor for per-axis dequantization.", - "T2") + "TS") .Input(2, "x_zero_point", "Zero point for input 'x'. Shape must match x_scale. It's optional. " "Zero point is 0 when it's not specified.", "T1", OpSchema::Optional) .Output(0, "y", "N-D full precision output tensor. It has same shape as input 'x'.", "T2") - .TypeConstraint("T1", {"tensor(int8)", "tensor(uint8)", "tensor(int16)", - "tensor(uint16)", "tensor(int32)", "tensor(int4)", - "tensor(uint4)"}, + .TypeConstraint("T1", {"tensor(int8)", "tensor(uint8)", "tensor(int16)", "tensor(uint16)", "tensor(int32)", "tensor(int4)", "tensor(uint4)", "tensor(float8e4m3fn)"}, "Constrain 'x' and 'x_zero_point' to 8-bit integer tensors, " - "16-bit integer tensors, or 32-bit signed integer tensors.") - .TypeConstraint("T2", {"tensor(float16)", "tensor(float)"}, - "Constrain 'y', 'x_scale' to float tensors.") + "16-bit integer tensors, 32-bit signed integer tensors, or " + "FP8 E4M3FN tensors.") + .TypeConstraint("T2", {"tensor(float16)", "tensor(float)", "tensor(bfloat16)"}, + "Constrain the output element type.") + .TypeConstraint("TS", {"tensor(float16)", "tensor(float)", "tensor(bfloat16)"}, + "Constrain the scale element type.") .SetDoc(DequantizeLinear_ver1_doc) .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - auto y_type = ctx.getOutputType(0); - // only float is supported - y_type->mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto::FLOAT); + const int64_t output_dtype = + getAttribute(ctx, "output_dtype", static_cast(0)); + if (output_dtype != 0) { + if (output_dtype != ONNX_NAMESPACE::TensorProto::FLOAT && + output_dtype != ONNX_NAMESPACE::TensorProto::FLOAT16 && + output_dtype != ONNX_NAMESPACE::TensorProto::BFLOAT16) { + fail_shape_inference( + "DequantizeLinear: output_dtype must be float, float16, or bfloat16"); + } + updateOutputElemType(ctx, 0, static_cast(output_dtype)); + } else if (ctx.getInputType(1) != nullptr) { + propagateElemTypeFromInputToOutput(ctx, 1, 0); + } if (!hasInputShape(ctx, 0)) return; @@ -852,16 +867,16 @@ ONNX_MS_OPERATOR_SET_SCHEMA( {"tensor(uint8)", "tensor(int8)"}, "Constrain input and output types to 8 bit signed and unsigned tensors.") .TypeAndShapeInferenceFunction([](InferenceContext& ctx) { - propagateElemTypeFromInputToOutput(ctx, 1, 0); - if (hasNInputShapes(ctx, 9)) { - std::vector shapes; - shapes.push_back(&ctx.getInputType(0)->tensor_type().shape()); - shapes.push_back(&ctx.getInputType(1)->tensor_type().shape()); - shapes.push_back(&ctx.getInputType(4)->tensor_type().shape()); - multidirectionalBroadcastShapeInference( - shapes, *ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape()); - } - })); + propagateElemTypeFromInputToOutput(ctx, 1, 0); + if (hasNInputShapes(ctx, 9)) { + std::vector shapes; + shapes.push_back(&ctx.getInputType(0)->tensor_type().shape()); + shapes.push_back(&ctx.getInputType(1)->tensor_type().shape()); + shapes.push_back(&ctx.getInputType(4)->tensor_type().shape()); + multidirectionalBroadcastShapeInference( + shapes, *ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape()); + } +})); ONNX_MS_OPERATOR_SET_SCHEMA( QGemm, 1, diff --git a/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc index c47c029a58168..3888cd665e71e 100644 --- a/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc +++ b/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc @@ -11,6 +11,9 @@ #include #include "contrib_ops/cuda/moe/qmoe_kernels.h" +#ifdef ORT_TEST_NVFP4_GEMV +#include "contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h" +#endif #include "core/providers/cuda/cuda_common.h" #include "core/platform/env_var_utils.h" @@ -23,6 +26,9 @@ namespace test { namespace { namespace qmoe = contrib::cuda; +#ifdef ORT_TEST_NVFP4_GEMV +namespace gemv = llm::kernels::moe_gemv; +#endif struct CudaBuffer { ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(CudaBuffer); @@ -53,6 +59,242 @@ struct CudaBuffer { size_t bytes; }; +#ifdef ORT_TEST_NVFP4_GEMV +template +void TestRoutingMaps(int rows, int experts, int k, bool normalize, bool tied) { + std::vector logits(rows * experts); + for (int i = 0; i < rows * experts; ++i) { + logits[i] = T(tied ? 0.0f : float((i * 73 + i / experts * 11) % 101 - 50) / 16.0f); + } + const int expanded = rows * k; + CudaBuffer input(logits.size() * sizeof(T)); + CudaBuffer scales(expanded * sizeof(float)), indices(expanded * sizeof(int)); + CudaBuffer r2u(expanded * sizeof(int)), u2r(expanded * sizeof(int)), sorted(expanded * sizeof(int)); + CudaBuffer offsets((experts + 1) * sizeof(int64_t)); + input.Upload(logits.data()); + qmoe::LaunchSoftmaxTopK(input.As(), scales.As(), indices.As(), + rows, experts, k, normalize, nullptr); + ASSERT_TRUE(qmoe::TryLaunchSmallDecodeExpertMaps( + indices.As(), r2u.As(), u2r.As(), sorted.As(), + offsets.As(), rows, experts, k, nullptr)); + std::vector expected_indices(expanded), actual_r2u(expanded), + actual_u2r(expanded), actual_sorted(expanded); + std::vector actual_offsets(experts + 1); + indices.Download(expected_indices.data()); + r2u.Download(actual_r2u.data()); + u2r.Download(actual_u2r.data()); + sorted.Download(actual_sorted.data()); + offsets.Download(actual_offsets.data()); + std::vector order(expanded); + for (int i = 0; i < expanded; ++i) order[i] = i; + std::stable_sort(order.begin(), order.end(), [&](int a, int b) { + return expected_indices[a] < expected_indices[b]; + }); + for (int rank = 0; rank < expanded; ++rank) { + const int original = order[rank]; + const int source = (original % k) * rows + original / k; + EXPECT_EQ(actual_r2u[rank], source); + EXPECT_EQ(actual_u2r[source], rank); + EXPECT_EQ(actual_sorted[rank], expected_indices[original]); + } + for (int expert = 0; expert <= experts; ++expert) { + EXPECT_EQ(actual_offsets[expert], + std::count_if(expected_indices.begin(), expected_indices.end(), + [&](int selected) { return selected < expert; })); + } +} + +TEST(CUDA_EP_Unittest, QMoENvfp4SmallDecodeMapsPreserveStableOrder) { + for (int rows : {1, 2}) { + for (int experts : {16, 257, 512}) { + for (int k : {1, 10, 16}) { + for (bool normalize : {false, true}) { + for (bool tied : {false, true}) { + TestRoutingMaps(rows, experts, k, normalize, tied); + TestRoutingMaps<__nv_bfloat16>(rows, experts, k, normalize, tied); + } + } + } + } + } + EXPECT_FALSE(qmoe::TryLaunchSmallDecodeExpertMaps( + nullptr, nullptr, nullptr, nullptr, nullptr, 3, 512, 10, nullptr)); + EXPECT_FALSE(qmoe::TryLaunchSmallDecodeExpertMaps( + nullptr, nullptr, nullptr, nullptr, nullptr, 2, 513, 10, nullptr)); + EXPECT_FALSE(qmoe::TryLaunchSmallDecodeExpertMaps( + nullptr, nullptr, nullptr, nullptr, nullptr, 0, 512, 10, nullptr)); + EXPECT_FALSE(qmoe::TryLaunchSmallDecodeExpertMaps( + nullptr, nullptr, nullptr, nullptr, nullptr, 2, 512, 17, nullptr)); + EXPECT_FALSE(qmoe::TryLaunchSmallDecodeExpertMaps( + nullptr, nullptr, nullptr, nullptr, nullptr, 2, 8, 10, nullptr)); +} + +template +void TestNvfp4KPackedGateUp(int k, bool use_bias, bool varied_scales = false, int n = 48) { + constexpr int experts = 3; + constexpr int rows = 4; + constexpr std::array routes{2, 0, 2, 1}; + constexpr std::array source_rows{1, 0, 0, 1}; + constexpr std::array globals{0.25f, 0.125f, 0.5f}; + constexpr std::array magnitudes{0, 0.5f, 1, 1.5f, 2, 3, 4, 6}; + std::vector activations(2 * k), bias(experts * n); + std::vector weights(experts * n * k / 2), scales(experts * n * k / 16); + for (size_t i = 0; i < activations.size(); ++i) { + activations[i] = T(float(int(i % 11) - 5) / 64.0f); + } + for (size_t i = 0; i < weights.size(); ++i) weights[i] = static_cast(i * 37 + 19); + for (size_t i = 0; i < scales.size(); ++i) { + scales[i] = varied_scales ? static_cast(i % 127) : (i % 2 ? 0x38 : 0x30); + } + for (size_t i = 0; i < bias.size(); ++i) bias[i] = T(float(int(i % 7) - 3) / 16.0f); + CudaBuffer a(activations.size() * sizeof(T)), w(weights.size()), s(scales.size()); + CudaBuffer g(sizeof(globals)), b(bias.size() * sizeof(T)), r(sizeof(routes)), sources(sizeof(source_rows)); + CudaBuffer output(rows * n / 2 * sizeof(T)); + a.Upload(activations.data()); + w.Upload(weights.data()); + s.Upload(scales.data()); + g.Upload(globals.data()); + b.Upload(bias.data()); + r.Upload(routes.data()); + sources.Upload(source_rows.data()); + llm::kernels::cutlass_kernels::ActivationParams activation; + activation.alpha = 1.0f; + activation.beta = 0.25f; + activation.limit = 7.0f; + gemv::launch_moe_gemv_fp4_symmetric_interleaved_swiglu( + a.As(), w.As(), nullptr, s.As(), g.As(), + use_bias ? b.As() : nullptr, output.As(), nullptr, r.As(), + experts, rows, n / 2, k, 16, 90, activation, gemv::MoeGemvConfig::kDefault, + false, true, sources.As(), 2, nullptr, true); + CUDA_CALL_THROW(cudaGetLastError()); + std::vector actual(rows * n / 2); + output.Download(actual.data()); + for (int row = 0; row < rows; ++row) { + for (int col = 0; col < n; col += 2) { + float sums[2]{}; + for (int pair = 0; pair < 2; ++pair) { + const int column = col + pair; + for (int d = 0; d < k; ++d) { + const uint8_t packed = weights[(routes[row] * n + column) * (k / 2) + d / 2]; + const int code = (d % 2 ? packed >> 4 : packed) & 15; + const uint8_t scale_code = scales[(routes[row] * n + column) * (k / 16) + d / 16]; + const int exponent = scale_code >> 3; + const int mantissa = scale_code & 7; + const float block_scale = exponent == 0 ? std::ldexp(float(mantissa), -9) + : std::ldexp(1.0f + mantissa * 0.125f, exponent - 7); + const float rounded_scale = static_cast(T(block_scale * globals[routes[row]])); + const float weight = static_cast(T(magnitudes[code & 7] * (code & 8 ? -1 : 1) * rounded_scale)); + sums[pair] += weight * static_cast(activations[source_rows[row] * k + d]); + } + if (use_bias) sums[pair] += static_cast(bias[routes[row] * n + column]); + } + const float gate = std::min(sums[0], activation.limit); + const float up = std::clamp(sums[1], -activation.limit, activation.limit) + activation.beta; + const float expected = static_cast(T(gate / (1.0f + std::exp(-gate)) * up)); + EXPECT_NEAR(static_cast(actual[row * n / 2 + col / 2]), expected, + 0.001f + std::abs(expected) * 0.01f) + << "row=" << row << ", column=" << col << ", K=" << k << ", bias=" << use_bias; + } + } +} + +TEST(CUDA_EP_Unittest, QMoENvfp4KPackedGateUpTiling) { + for (int k : {144, 2560}) { + for (bool use_bias : {false, true}) { + TestNvfp4KPackedGateUp(k, use_bias); + TestNvfp4KPackedGateUp<__nv_bfloat16>(k, use_bias); + } + } +} + +TEST(CUDA_EP_Unittest, QMoENvfp4KPackedGateUpScaleRange) { + TestNvfp4KPackedGateUp(2560, true, true); + TestNvfp4KPackedGateUp<__nv_bfloat16>(2560, true, true); +} + +TEST(CUDA_EP_Unittest, QMoENvfp4KPackedGateUpPartialTiles) { + for (int n : {2, 10}) { + for (int k : {16, 272}) { + for (bool use_bias : {false, true}) { + TestNvfp4KPackedGateUp(k, use_bias, false, n); + TestNvfp4KPackedGateUp<__nv_bfloat16>(k, use_bias, false, n); + } + } + } +} + +template +void TestNvfp4FusedFinalize(int num_rows, int topk, int k, bool use_bias) { + constexpr int n = 50, experts = 3; + const int expanded = num_rows * topk; + std::vector activations(expanded * k), biases(experts * n); + std::vector weights(experts * n * k / 2), scales(experts * n * k / 16, 0x30); + std::vector globals{0.3f, 0.17f, 0.51f}, routing(expanded); + std::vector mapping(expanded), expert_ids(expanded); + for (size_t i = 0; i < activations.size(); ++i) activations[i] = T(float(int(i % 11) - 5) / 64.0f); + for (size_t i = 0; i < weights.size(); ++i) weights[i] = static_cast(i * 37 + 19); + for (size_t i = 0; i < biases.size(); ++i) biases[i] = T(float(int(i % 7) - 3) / 16.0f); + for (int i = 0; i < expanded; ++i) { + mapping[i] = (i * 7) % expanded; + expert_ids[i] = i % experts; + routing[i] = float(i % topk + 1) / float(topk * (topk + 1) / 2); + } + if (topk > 1) mapping[expanded - 1] = -1; + if (expanded > 1) expert_ids[expanded / 2] = experts; + CudaBuffer a(activations.size() * sizeof(T)), w(weights.size()), s(scales.size()); + CudaBuffer g(globals.size() * sizeof(float)), b(biases.size() * sizeof(T)); + CudaBuffer m(mapping.size() * sizeof(int)), e(expert_ids.size() * sizeof(int)), r(routing.size() * sizeof(float)); + CudaBuffer reference(expanded * n * sizeof(T)), output(num_rows * n * sizeof(T)); + a.Upload(activations.data()); + w.Upload(weights.data()); + s.Upload(scales.data()); + g.Upload(globals.data()); + b.Upload(biases.data()); + m.Upload(mapping.data()); + e.Upload(expert_ids.data()); + r.Upload(routing.data()); + gemv::launch_moe_gemv_fp4_symmetric( + a.As(), w.As(), nullptr, s.As(), g.As(), + use_bias ? b.As() : nullptr, reference.As(), nullptr, e.As(), + experts, expanded, n, k, 16, 90, gemv::MoeGemvConfig::kDefault, false, true, nullptr, true); + ASSERT_TRUE(gemv::try_launch_moe_gemv_nvfp4_fused_finalize( + a.As(), w.As(), s.As(), g.As(), + use_bias ? b.As() : nullptr, output.As(), m.As(), e.As(), r.As(), + experts, num_rows, topk, n, k, nullptr)); + CUDA_CALL_THROW(cudaGetLastError()); + std::vector reference_values(expanded * n), actual(num_rows * n); + reference.Download(reference_values.data()); + output.Download(actual.data()); + for (int row = 0; row < num_rows; ++row) { + for (int col = 0; col < n; ++col) { + float sum = 0.0f; + for (int rank = 0; rank < topk; ++rank) { + const int permuted = mapping[rank * num_rows + row]; + if (permuted >= 0 && expert_ids[permuted] < experts) { + sum += routing[row * topk + rank] * static_cast(reference_values[permuted * n + col]); + } + } + const float expected = static_cast(T(sum)); + EXPECT_NEAR(static_cast(actual[row * n + col]), expected, 1e-5f + std::abs(expected) * 0.01f) + << "row=" << row << ", col=" << col << ", topk=" << topk << ", bias=" << use_bias; + } + } +} + +TEST(CUDA_EP_Unittest, QMoENvfp4FusedFinalizePreservesRoutingAndRounding) { + for (int rows : {1, 2}) { + for (int topk : {1, 4, 10, 16}) { + for (int k : {144, 640}) { + for (bool use_bias : {false, true}) { + TestNvfp4FusedFinalize(rows, topk, k, use_bias); + TestNvfp4FusedFinalize<__nv_bfloat16>(rows, topk, k, use_bias); + } + } + } + } +} +#endif + TEST(CUDA_EP_Unittest, QMoEFp8CompactExpertsPreservesRoutesAndClearsUnusedSlots) { constexpr int num_experts = 512; constexpr int capacity = 6; diff --git a/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc index 2173434746ce9..9d6797761e3a0 100644 --- a/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc +++ b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc @@ -234,6 +234,22 @@ TEST(FusedTensorOpsTest, SigmoidGateMulSiluEmptyAndInvalidAttributes) { } } +#if !defined(DISABLE_FLOAT8_TYPES) +TEST(FusedTensorOpsTest, DequantizeLinearFp8ToBf16) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + + OpTester test("DequantizeLinear", 1, kMSDomain); + test.AddAttribute("axis", 1); + test.AddAttribute("output_dtype", ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16); + test.AddInput("weight", {1, 8}, std::vector(8, Float8E4M3FN(1.0f))); + test.AddInput("scale", {8}, std::vector(8, 2.0f)); + test.AddOptionalInputEdge(); + test.AddOutput("weight_bf16", {1, 8}, ToTensorType(std::vector(8, 2.0f))); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} +#endif + TEST(FusedTensorOpsTest, UntileHyperConnectionWeights) { constexpr int64_t kInputElements = 160ll * 20480; constexpr int64_t kVectorCount = kInputElements / 16; diff --git a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc index 9dd9b0e722a29..0695e717ea6eb 100644 --- a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include +#include #include "gtest/gtest.h" #include "test/common/cuda_op_test_utils.h" @@ -69,6 +70,121 @@ std::vector MakeConstRowWeight(const std::vector& row_value } } // namespace +template +void TestFp8SplitOutputs() { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "CUDA device is required for split-output FP8 matmul."; + } + for (const auto& [m, k] : {std::pair{0, 256}, {1, 256}, {2, 256}, {2, 320}, {2, 1088}, {9, 256}, {32, 256}, {33, 256}, {65, 256}, {2, 16}, {2, 30}, {2, 0}}) { + for (const bool activation_scale : {false, true}) { + SCOPED_TRACE("M=" + std::to_string(m) + " K=" + std::to_string(k) + + " W8A8=" + std::to_string(activation_scale)); + constexpr int64_t n = 17; + constexpr int64_t block_size = 64; + const int64_t k_blocks = (k + block_size - 1) / block_size; + const std::vector widths = m == 2 && k == 256 + ? std::vector{1, 1, 1, 1, 1, 1, 1, 10} + : (m == 33 ? std::vector{17} + : std::vector{1, 7, 9}); + std::vector a; + std::vector b; + std::vector scales; + std::vector bias; + for (int64_t row = 0; row < m; ++row) { + for (int64_t i = 0; i < k; ++i) { + a.emplace_back(static_cast((row + i) % 5 - 2) / 4.0f); + } + } + for (int64_t col = 0; col < n; ++col) { + bias.emplace_back(static_cast(col - 8) / 8.0f); + for (int64_t i = 0; i < k; ++i) { + b.emplace_back(static_cast((col + i) % 3 - 1)); + } + for (int64_t kb = 0; kb < k_blocks; ++kb) { + scales.push_back(static_cast(1 + (col + kb) % 2) / 4.0f); + } + } + TensorShapeVector a_shape = m == 1 ? TensorShapeVector{k} : TensorShapeVector{m, k}; + if (m == 9) { + a_shape = {3, 3, k}; + } + OpTester test("MatMulBlockQuantizedFp8Weight", 1, kMSDomain); + test.AddAttribute("block_size", block_size); + test.AddAttribute>("split_sizes", widths); + test.AddInput("A", a_shape, a); + test.AddInput("B", {n, k}, b); + test.AddInput("b_scale", {n, k_blocks}, scales); + if (activation_scale) { + test.AddInput("a_scale", {}, {0.5f}); + } else { + test.AddOptionalInputEdge(); + } + if (activation_scale) { + test.AddInput("bias", {n}, bias); + } + int64_t start = 0; + for (size_t output = 0; output < widths.size(); ++output) { + std::vector expected; + for (int64_t row = 0; row < m; ++row) { + for (int64_t col = start; col < start + widths[output]; ++col) { + float sum = activation_scale ? static_cast(bias[col]) : 0.0f; + for (int64_t i = 0; i < k; ++i) { + sum += static_cast(a[row * k + i]) * b[col * k + i].ToFloat() * + scales[col * k_blocks + i / block_size]; + } + expected.emplace_back(sum); + } + } + auto shape = a_shape; + shape.back() = widths[output]; + test.AddOutput(("Y" + std::to_string(output)).c_str(), shape, expected); + start += widths[output]; + } + test.SetOutputTolerance(0.0f, 0.0f); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + } + } +} + +TEST(MatMulBlockQuantizedFp8WeightOpTest, SplitOutputsFp16) { + TestFp8SplitOutputs(); +} + +TEST(MatMulBlockQuantizedFp8WeightOpTest, SplitOutputsBFloat16) { + TestFp8SplitOutputs(); +} + +TEST(MatMulBlockQuantizedFp8WeightOpTest, RejectsInvalidSplitSizes) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "CUDA device is required for split-output FP8 matmul."; + } + for (const std::vector& widths : {std::vector{}, {0, 17}, {-1, 18}, {8, 8}, {1, 1, 1, 1, 1, 1, 1, 1, 9}}) { + OpTester test("MatMulBlockQuantizedFp8Weight", 1, kMSDomain); + test.AddAttribute>("split_sizes", widths); + test.AddInput("A", {1, 16}, std::vector(16)); + test.AddInput("B", {17, 16}, std::vector(17 * 16)); + test.AddInput("b_scale", {17, 1}, std::vector(17, 1.0f)); + for (size_t i = 0; i < std::max(1, widths.size()); ++i) { + const int64_t width = i < widths.size() ? std::max(0, widths[i]) : 17; + test.AddOutput(("Y" + std::to_string(i)).c_str(), {1, width}, std::vector(width)); + } + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectFailure, "split_sizes", {}, nullptr, &providers); + } + OpTester count_test("MatMulBlockQuantizedFp8Weight", 1, kMSDomain); + count_test.AddAttribute>("split_sizes", {8, 9}); + count_test.AddInput("A", {1, 16}, std::vector(16)); + count_test.AddInput("B", {17, 16}, std::vector(17 * 16)); + count_test.AddInput("b_scale", {17, 1}, std::vector(17, 1.0f)); + count_test.AddOutput("Y", {1, 17}, std::vector(17)); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + count_test.Run(OpTester::ExpectResult::kExpectFailure, "split_sizes", {}, nullptr, &providers); +} + TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvTensorCoreKSplitSelection) { struct Case { int n; diff --git a/onnxruntime/test/contrib_ops/moe_test.cc b/onnxruntime/test/contrib_ops/moe_test.cc index 4ed7e9e4a3968..507044ded4329 100644 --- a/onnxruntime/test/contrib_ops/moe_test.cc +++ b/onnxruntime/test/contrib_ops/moe_test.cc @@ -2651,8 +2651,8 @@ template static void RunQMoEPackedPrefillRoutingTest(int64_t num_rows, int64_t hidden_size, int64_t inter_size, int64_t top_k, bool skewed, bool scale_initializers = true, int64_t row_tile_size = 0, int64_t block_size = 64, - int64_t fc1_bits = 2, int64_t fc2_bits = 4) { - constexpr int64_t num_experts = 8; + int64_t fc1_bits = 2, int64_t fc2_bits = 4, + int64_t num_experts = 8) { constexpr float alpha = 1.5f; constexpr float beta = 0.125f; constexpr float limit = 0.5f; @@ -2665,7 +2665,7 @@ static void RunQMoEPackedPrefillRoutingTest(int64_t num_rows, int64_t hidden_siz return 0.125f * static_cast(expert + 1) + 0.015625f * static_cast(column % 5) + 0.03125f * static_cast(block); }; - const auto make_weights = [](int64_t rows, int64_t columns, int64_t bits, bool interleaved) { + const auto make_weights = [num_experts](int64_t rows, int64_t columns, int64_t bits, bool interleaved) { const int64_t pack_size = 8 / bits; const uint8_t zero = static_cast(1u << (bits - 1)); std::vector weights(static_cast(num_experts * rows * columns / pack_size), 0); @@ -2793,6 +2793,7 @@ TEST(MoETest, QMoETest_MixedWidthCudaPackedPrefillRouting) { if (!HasCudaEnvironment(800)) { GTEST_SKIP() << "SM80 or later is required for packed INT2 prefill."; } + ScopedEnvironmentVariables scoped_env_vars{ {{"ORT_DISABLE_MOE_GEMV", "1"}, {"ORT_ENABLE_QMOE_INT2_PREFILL", "1"}}}; for (int64_t top_k : {1, 2, 4}) { @@ -2806,6 +2807,20 @@ TEST(MoETest, QMoETest_MixedWidthCudaPackedPrefillRouting) { } } +TEST(MoETest, QMoETest_CudaColumnTiledFinalize) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "SM80 or later is required for packed INT2 decode."; + } + for (int64_t rows : {1, 2}) { + for (int64_t hidden : {2560, 2624}) { + RunQMoEPackedPrefillRoutingTest(rows, hidden, 128, 10, false, true, 0, 64, 2, 4, 16); +#if defined(ENABLE_BF16) + RunQMoEPackedPrefillRoutingTest(rows, hidden, 128, 10, false, true, 0, 64, 2, 4, 16); +#endif + } + } +} + TEST(MoETest, QMoETest_MixedWidthCudaPackedPrefillRuntimeScales) { if (!HasCudaEnvironment(800)) { GTEST_SKIP() << "SM80 or later is required for packed INT2 prefill.";