diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 6dcb758a234fe..aaa818194cf24 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -6426,11 +6426,14 @@ This version of the operator has been available since version 1 of the 'com.micr dequantized_weight[e, n, k] = float(weight[e, n, k]) * scale[e, n / block_size, k / block_size] Partial edge blocks are allowed. No zero points or activation scales are accepted. Without a positive block_size, FP8 uses the legacy per-expert fc*_global_scale inputs instead. - Block-scaled FP8 does not use global scales. Activations retain the input type (weight-only quantization). - On CUDA SM80-or-newer GPUs, SiLU/SwiGLU block-scaled FP8 uses fused weight-only GEMV/GEMM by default. + Block-scaled FP8 does not use global scales. By default activations retain the input type (weight-only quantization). + activation_quantization='dynamic_fp8' instead selects dynamic per-token K128 activation quantization, + with native FP8 dot products, FP32 block-scale accumulation and FP32 routed down-projection reduction. + On CUDA SM80-or-newer GPUs, weight-only SiLU/SwiGLU block-scaled FP8 uses fused GEMV/GEMM by default. These kernels decode and scale FP8 tiles on chip without materializing full FP16/BF16 weight buffers. Set ORT_ENABLE_FP8_FUSED=0 before session creation to select the compact dense-dequant fallback. - Other activation types and unsupported routing sizes retain the fallback. + Other activation types and unsupported routing sizes retain the fallback in weight-only mode. + Dynamic activation quantization rejects unsupported configurations instead of falling back to weight-only execution. The WebGPU block-FP8 kernel supports block_size=128 with float32 scales and rejects projections that exceed 32-bit shader addressing or the device's per-dimension dispatch limit. @@ -6482,6 +6485,8 @@ This version of the operator has been available since version 1 of the 'com.micr
Alpha parameter used in activation function.
activation_beta : float
Beta parameter used in activation function.
+
activation_quantization : string
+
'none' (default) retains weight-only execution. 'dynamic_fp8' quantizes consecutive groups of 128 features per token to saturating nearest-even E4M3FN using FP32 div.full normalization and scale max(amax(abs(x)),1e-10) * float32(1/448). Requires block-scaled FP8 weights, block_size=128, FP32 weight scales, FP32/FP16/BF16 inputs and 128-aligned projections. Gate/up outputs narrow to the input type; their FP32 SiLU/SwiGLU result is quantized directly for FC2. FC2 stays FP32, multiplies routing weights rounded to the input type, and reduces in FP32 before the final output cast. Biases, supplied activation/global scales, zero points, router_weights and nonstandard activation parameters are unsupported. CUDA requires native FP8 MMA (SM89+) and at most 65535 token/expert pairs.
activation_type : string
Activation function to use. Choose from relu, gelu, silu, swiglu and identity. CUDA also supports relu2 (squared ReLU). Default is relu
block_size : int
@@ -6499,7 +6504,7 @@ This version of the operator has been available since version 1 of the 'com.micr
normalize_routing_weights : int
Whether to normalize routing weights
quant_type : string
-
Quantization type: 'int' for integer quantization (default), 'fp4' for MXFP4 quantization, 'nvfp4' for NVFP4 quantization, 'fp8' for FP8 e4m3 weight-only quantization, or 'wfp4afp8' for MXFP4 weight with FP8 activation. When quant_type is 'fp4' or 'nvfp4', weights are stored in E2M1 FP4 format (2 values per byte), fc*_scales inputs contain the FP4 block scales, and fc*_global_scale inputs must be provided. 'fp4' uses Float8E8M0 block scales with block_size 32; 'nvfp4' uses Float8E4M3FN block scales with block_size 16. 'fp8' uses float8e4m3fn weights and, when block_size > 0, float32, float16, or bfloat16 square-block scales in fc*_scales instead of fc*_global_scale. The INT3 profile requires integer quantization.
+
Quantization type: 'int' for integer quantization (default), 'fp4' for MXFP4 quantization, 'nvfp4' for NVFP4 quantization, 'fp8' for FP8 e4m3 weights (weight-only unless activation_quantization is set), or 'wfp4afp8' for MXFP4 weight with FP8 activation. When quant_type is 'fp4' or 'nvfp4', weights are stored in E2M1 FP4 format (2 values per byte), fc*_scales inputs contain the FP4 block scales, and fc*_global_scale inputs must be provided. 'fp4' uses Float8E8M0 block scales with block_size 32; 'nvfp4' uses Float8E4M3FN block scales with block_size 16. 'fp8' uses float8e4m3fn weights and, when block_size > 0, float32, float16, or bfloat16 square-block scales in fc*_scales instead of fc*_global_scale. The INT3 profile requires integer quantization.
swiglu_fusion : int
0: not fused, 1: fused and interleaved. 2: fused and not interleaved.
swiglu_limit : float
@@ -9014,5 +9019,3 @@ No versioning maintained for experimental ops.
T : tensor(float)
Constrain input and output types to float32 tensors.
- - diff --git a/docs/contrib_ops/cuda/moe_qmoe.md b/docs/contrib_ops/cuda/moe_qmoe.md index 089f28a05f48e..a6472e8030d46 100644 --- a/docs/contrib_ops/cuda/moe_qmoe.md +++ b/docs/contrib_ops/cuda/moe_qmoe.md @@ -1411,7 +1411,46 @@ partial final tiles, and the decode scratch bound. Internal CUDA tests in `qmoe_fp8_compaction_test.cc` check expert remapping, unused slots, and selected weight/bias decoding for separate and fused layouts. -### 10.4 Kernel instantiation files +### 10.4 Dynamic block-FP8 activations + +Set `activation_quantization="dynamic_fp8"` to opt into W8A8 execution instead of +the default `activation_quantization="none"` weight-only path. This mode requires +`quant_type="fp8"`, `expert_weight_bits=8`, `block_size=128`, FP32 weight scale +grids, FP32/FP16/BF16 inputs, and hidden/intermediate dimensions divisible by 128. +The existing row-major FP8 weight layouts and separate gate/up or fused SwiGLU +layouts are unchanged. Standard SiLU/SwiGLU is supported; biases, zero points, +activation clipping/parameter overrides, `router_weights`, and supplied +activation/global scales are rejected. + +For each token and consecutive group of 128 features, compute the FP32 scale +`max(amax(abs(x)), 1e-10) * float32(1/448)`, then clamp `x / scale` to `[-448,448]` +and convert to E4M3FN with round-to-nearest-even. This quantization is performed +before gate/up projection and again directly on its FP32 SiLU/SwiGLU result. +The rounded FP32 reciprocal multiplication is intentional: replacing it with +a correctly rounded division can shift a scale by one ULP and change FP8 codes. +Normalization uses FP32 `div.full` before the nearest-even FP8 conversion, not +`div.rn`; the latter can move a midpoint by one ULP and select a different code. +Gate/up projection outputs round to the input type, but there is **no** +FP16/BF16 rounding between the activation product and the second quantizer. +For FP32 input, gate/up values and the final output remain FP32. CUDA's FLOAT +registration is restricted to FP8 weights with FP32 scales and requires this +dynamic mode; it does not claim existing FLOAT/integer CPU QMoE models or enable +FLOAT weight-only execution. +Native FP8 MMA accumulates each K128 dot in FP32, applies that group's activation +and weight scales, then adds it to the FP32 projection accumulator. +Down projection remains FP32, multiplies the selected routing coefficients +rounded to the input type, and reduces in top-k order in FP32 before casting +the final output. This is not equivalent to FP8 QDQ followed by A16 GEMM. + +CUDA requires SM89+ and a kernel compiled for SM89 or newer. The existing +`onnxruntime_USE_FP8_QMOE=ON` build option is required. Dynamic mode rejects +`ORT_ENABLE_FP8_FUSED=0` and invocations with more than 65535 token/expert pairs +instead of silently falling back to weight-only execution. Chunk longer +prefills at the graph boundary. Activation scratch contains FP8 values and FP32 +group scales; no dense expert-weight buffer is allocated. Debug dispatch is +`fp8_dynamic_mma` with `DequantWeightBytes=0`. + +### 10.5 Kernel instantiation files | File | Template | |------|----------| diff --git a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc index 782cfc5ff364e..9d666f3dfc9e0 100644 --- a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc +++ b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc @@ -922,6 +922,8 @@ template QMoECPU::QMoECPU(const OpKernelInfo& op_kernel_info) : OpKernel(op_kernel_info), MoEBaseCPU(op_kernel_info) { + ORT_ENFORCE(op_kernel_info.GetAttrOrDefault("activation_quantization", "none") == "none", + "CPU QMoE does not support dynamic FP8 activation quantization."); ORT_ENFORCE(activation_type_ != ActivationType::SwiGLU || swiglu_fusion_ == 1, "CPU QMoE only supports interleaved SwiGLU format. Please set swiglu_fusion=1."); ORT_ENFORCE(op_kernel_info.GetAttr("expert_weight_bits", &expert_weight_bits_).IsOK()); diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 5d349ddf1bade..f526de4323a2c 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -102,6 +102,9 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, MoE); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, MoE); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, QMoE); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, QMoE); +#if !defined(DISABLE_FLOAT8_TYPES) +class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, QMoE); +#endif class CUDA_MS_OP_TYPED_CLASS_NAME(1, float_float, MultiHeadAttention); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float_MLFloat16, MultiHeadAttention); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16_float, MultiHeadAttention); @@ -470,6 +473,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/moe/moe_quantization.cc b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc index 028761186bad7..adece3d959020 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc +++ b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc @@ -326,6 +326,19 @@ namespace cuda { REGISTER_KERNEL_TYPED(MLFloat16) REGISTER_KERNEL_TYPED(BFloat16) +#if !defined(DISABLE_FLOAT8_TYPES) +// Do not claim FLOAT/integer models that already execute on CPU. +ONNX_OPERATOR_TYPED_KERNEL_EX( + QMoE, kMSDomain, 1, float, kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .MayInplace(0, 0) + .TypeConstraint("T", DataTypeImpl::GetTensorType()) + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) + .TypeConstraint("T4", DataTypeImpl::GetTensorType()), + QMoE); +#endif + QMoE::QMoE(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info), MoEBase(op_kernel_info, GetDeviceProp()) { enable_kernel_debug_info_ = onnxruntime::ParseEnvironmentVariableWithDefault(kEnableQMoEKernelDebugInfo, 0) != 0; @@ -381,6 +394,20 @@ QMoE::QMoE(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info), MoE block_size_ = op_kernel_info.GetAttrOrDefault("block_size", -1); this->quant_type_ = op_kernel_info.GetAttrOrDefault("quant_type", "int"); + const auto activation_quantization = + op_kernel_info.GetAttrOrDefault("activation_quantization", "none"); + ORT_ENFORCE(activation_quantization == "none" || activation_quantization == "dynamic_fp8", + "activation_quantization must be 'none' or 'dynamic_fp8'."); + dynamic_fp8_ = activation_quantization == "dynamic_fp8"; + ORT_ENFORCE(!dynamic_fp8_ || + (quant_type_ == "fp8" && block_size_ == 128 && expert_weight_bits_ == 8 && !is_mixed_width && + sm_ >= 89 && !use_sparse_mixer_ && + (activation_type_ == onnxruntime::llm::kernels::cutlass_kernels::ActivationType::Silu || + activation_type_ == onnxruntime::llm::kernels::cutlass_kernels::ActivationType::Swiglu) && + activation_alpha_ == 1.0f && activation_beta_ == 0.0f && + swiglu_limit_ == std::numeric_limits::infinity()), + "QMoE dynamic_fp8 requires SM89+, block-scaled FP8 weights with block_size=128, " + "SiLU or standard SwiGLU, and softmax/top-k routing."); ORT_ENFORCE(quant_type_ == "int" || quant_type_ == "fp4" || quant_type_ == "nvfp4" || quant_type_ == "fp8" || quant_type_ == "wfp4afp8", "quant_type must be 'int', 'fp4', 'nvfp4', 'fp8', or 'wfp4afp8', but got '", quant_type_, "'"); @@ -451,6 +478,8 @@ QMoE::QMoE(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info), MoE #endif enable_fp8_fused_ = quant_type_ == "fp8" && sm_ >= 80 && onnxruntime::ParseEnvironmentVariableWithDefault("ORT_ENABLE_FP8_FUSED", 1) != 0; + ORT_ENFORCE(!dynamic_fp8_ || (enable_fp8_fused_ && QMoEDynamicFp8Supported()), + "QMoE dynamic_fp8 requires fused execution and a CUDA kernel compiled for SM89 or newer."); using namespace onnxruntime::llm::kernels::cutlass_kernels; @@ -462,6 +491,12 @@ QMoE::QMoE(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info), MoE bool is_fp16 = input_type == ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT16; #endif is_fp16_ = is_fp16; + ORT_ENFORCE(input_type != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT || dynamic_fp8_, + "CUDA QMoE FLOAT input requires activation_quantization='dynamic_fp8'."); + if (dynamic_fp8_) { + // Native W8A8 has its own projection kernels and never invokes the A16 runner. + return; + } if (has_int3_) { ORT_ENFORCE(quant_type_ == "int" && weights_prepacked_mode == 0, "INT3 CUDA QMoE requires quant_type='int' and weights_prepacked=0."); @@ -788,6 +823,16 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { const bool is_fp4_family = is_fp4 || is_nvfp4; const bool is_fp8 = (quant_type_ == "fp8"); const bool is_block_fp8 = is_fp8 && block_size_ > 0; + if (dynamic_fp8_) { + for (int index : {4, 7, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20}) { + ORT_RETURN_IF(context->Input(index) != nullptr, + "QMoE dynamic_fp8 does not accept bias, zero points, router_weights or supplied activation/global scales."); + } + for (int index : {3, 6, 9}) { + const auto* scales = context->Input(index); + ORT_RETURN_IF(scales && !scales->IsDataType(), "QMoE dynamic_fp8 requires FP32 weight scales."); + } + } const bool is_wfp4afp8 = (quant_type_ == "wfp4afp8"); const bool is_int = (quant_type_ == "int"); const bool is_mixed_width = fc1_expert_weight_bits_ != expert_weight_bits_ || @@ -971,6 +1016,8 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { "Block-scaled FP8 requires row-major expert weights [E,N,K]."); ORT_RETURN_IF_NOT(moe_params.hidden_size % 8 == 0 && moe_params.inter_size % 8 == 0, "Block-scaled FP8 CUDA requires hidden_size and inter_size to be multiples of 8."); + ORT_RETURN_IF(dynamic_fp8_ && (moe_params.hidden_size % 128 != 0 || moe_params.inter_size % 128 != 0), + "QMoE dynamic_fp8 requires hidden and intermediate sizes to be multiples of 128."); ORT_RETURN_IF_ERROR(moe_helper::CheckFp8BlockScale( context->Input(3), moe_params.num_experts, fc1_n, moe_params.hidden_size, block_size_, "fc1_scales")); ORT_RETURN_IF_ERROR(moe_helper::CheckFp8BlockScale( @@ -1059,6 +1106,10 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { ORT_RETURN_IF_NOT(k_ > 0 && k_ <= moe_params.num_experts, "QMoE requires 0 < k <= num_experts, got k=", k_, " and num_experts=", moe_params.num_experts); + if (dynamic_fp8_ && moe_params.num_rows == 0) { + context->Output(0, input->Shape()); + return Status::OK(); + } // The INT4/INT8 weight-only path stores B in the column-interleaved layout (ColumnMajorTileInterleave), // whose CUTLASS pitchlinear iterators require the GEMM reduction dim K to be a whole multiple of the @@ -1379,6 +1430,8 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { (kernel_activation_type == ActivationType::Silu || kernel_activation_type == ActivationType::Swiglu) && moe_params.num_rows <= 65535 / k_; + ORT_RETURN_IF(dynamic_fp8_ && !use_fp8_fused, + "QMoE dynamic_fp8 requires at most 65535 expanded routing rows; chunk the input tokens."); // Offline-packed SM80 bytes cannot be consumed by the dense dequantization fallback. ORT_RETURN_IF(preloaded_int2_or_mixed && !use_packed_int, "INT2 or mixed-width CUDA QMoE with weights_prepacked=1 requires a packed execution path " @@ -1651,15 +1704,27 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { const int fc1_n = inter * (fusion ? 2 : 1); const int max_routes = narrow(row_tile_plan.rows_per_tile * k_); const size_t element_size = input->DataType()->Size(); - const size_t projection_element_size = fusion ? element_size : sizeof(float); + const size_t projection_element_size = fusion || dynamic_fp8_ ? element_size : sizeof(float); auto row_map = GetScratchBuffer(max_routes, GetComputeStream(context)); auto sorted_experts = GetScratchBuffer(max_routes, GetComputeStream(context)); auto expert_offsets = GetScratchBuffer(num_experts + 1, GetComputeStream(context)); auto tile_offsets = GetScratchBuffer(num_experts + 1, GetComputeStream(context)); const size_t projection_bytes = SafeInt(max_routes) * fc1_n * projection_element_size; auto projection = GetScratchBuffer(projection_bytes, GetComputeStream(context)); - auto activation = GetScratchBuffer(SafeInt(max_routes) * inter * element_size, GetComputeStream(context)); - auto result = GetScratchBuffer(SafeInt(max_routes) * hidden * element_size, GetComputeStream(context)); + IAllocatorUniquePtr activation; + if (!dynamic_fp8_) { + activation = GetScratchBuffer(SafeInt(max_routes) * inter * element_size, GetComputeStream(context)); + } + auto result = GetScratchBuffer(SafeInt(max_routes) * hidden * + (dynamic_fp8_ ? sizeof(float) : element_size), + GetComputeStream(context)); + const size_t quantized_elements = SafeInt(dynamic_fp8_ ? max_routes : 0) * std::max(hidden, inter); + IAllocatorUniquePtr quantized; + IAllocatorUniquePtr activation_scales; + if (dynamic_fp8_) { + quantized = GetScratchBuffer(quantized_elements, GetComputeStream(context)); + activation_scales = GetScratchBuffer(quantized_elements / 128, GetComputeStream(context)); + } auto scale_type = [](const Tensor* tensor) { return tensor->IsDataType() ? 0 : (tensor->IsDataType() ? 1 : 2); }; @@ -1691,7 +1756,7 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { expert_indices, sorted_experts.get(), row_map.get(), unpermuted_row_to_permuted_row, expert_offsets.get(), counts.get(), cumulative_counts.get(), blocked_rows.get(), rows, num_experts, k_, 0, stream); } - const bool gemm = rows > 8; + const bool gemm = dynamic_fp8_ || rows > 8; if (gemm) { LaunchQMoEFp8ExpertTiles(expert_offsets.get(), tile_offsets.get(), num_experts, stream); } @@ -1716,38 +1781,67 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { params.block_size = narrow(block_size_); params.fusion = fusion; const auto* tile_input = static_cast(input->DataRaw()) + row_offset * hidden; - const auto* fc1_bias = fc1_experts_bias_optional ? static_cast(fc1_experts_bias_optional->DataRaw()) : nullptr; - auto project_and_activate = [&](auto* projected) { - LaunchQMoEFp8Projection(params, tile_input, projected, stream); - LaunchQMoEFp8Activation(projected, fc1_bias, static_cast(activation.get()), - sorted_experts.get(), routes, inter, fusion, - split_fp8_fc1 || !fusion ? 1.0f : activation_alpha_, - split_fp8_fc1 ? 0.0f : activation_beta_, - split_fp8_fc1 ? std::numeric_limits::infinity() : swiglu_limit_, stream); - }; - if (fusion) { - project_and_activate(static_cast(projection.get())); + if (dynamic_fp8_) { + LaunchQMoEDynamicFp8Quantize(tile_input, quantized.get(), activation_scales.get(), rows, hidden, 0, stream); + LaunchQMoEDynamicFp8Projection( + params, quantized.get(), activation_scales.get(), static_cast(projection.get()), + nullptr, nullptr, narrow(k_), stream); + LaunchQMoEDynamicFp8Quantize(static_cast(projection.get()), quantized.get(), + activation_scales.get(), routes, inter, fusion ? 2 : 1, stream); + params.weights = static_cast(fc2_experts_weights->DataRaw()); + params.scales = fc2_scale_tensor->DataRaw(); + params.up_weights = nullptr; + params.up_scales = nullptr; + params.row_to_unpermuted = nullptr; + params.n = hidden; + params.k = inter; + params.fusion = 0; + LaunchQMoEDynamicFp8Projection( + params, quantized.get(), activation_scales.get(), static_cast(result.get()), + expert_scales, row_map.get(), narrow(k_), stream); + LaunchQMoEDynamicFp8Finalize(static_cast(result.get()), + static_cast(output->MutableDataRaw()) + row_offset * hidden, + unpermuted_row_to_permuted_row, rows, hidden, narrow(k_), stream); + return; + } + if constexpr (!std::is_same_v) { + const auto* fc1_bias = fc1_experts_bias_optional ? static_cast(fc1_experts_bias_optional->DataRaw()) : nullptr; + auto project_and_activate = [&](auto* projected) { + LaunchQMoEFp8Projection(params, tile_input, projected, stream); + LaunchQMoEFp8Activation(projected, fc1_bias, static_cast(activation.get()), + sorted_experts.get(), routes, inter, fusion, + split_fp8_fc1 || !fusion ? 1.0f : activation_alpha_, + split_fp8_fc1 ? 0.0f : activation_beta_, + split_fp8_fc1 ? std::numeric_limits::infinity() : swiglu_limit_, stream); + }; + if (fusion) { + project_and_activate(static_cast(projection.get())); + } else { + // Match the dense SiLU epilogue: narrow only after bias and activation. + project_and_activate(static_cast(projection.get())); + } + params.weights = static_cast(fc2_experts_weights->DataRaw()); + params.scales = fc2_scale_tensor->DataRaw(); + params.scale_type = scale_type(fc2_scale_tensor); + params.up_weights = nullptr; + params.up_scales = nullptr; + params.row_to_unpermuted = nullptr; + params.n = hidden; + params.k = inter; + params.fusion = 0; + LaunchQMoEFp8Projection(params, static_cast(activation.get()), static_cast(result.get()), stream); + const auto* fc2_bias = fc2_experts_bias_optional ? static_cast(fc2_experts_bias_optional->DataRaw()) : nullptr; + ck::finalizeMoeRoutingKernelLauncher( + static_cast(result.get()), static_cast(output->MutableDataRaw()) + row_offset * hidden, + fc2_bias, expert_scales, unpermuted_row_to_permuted_row, row_map.get(), expert_indices, expert_offsets.get(), + rows, hidden, k_, num_experts, parallelism_config, false, stream); } else { - // Match the dense SiLU epilogue: narrow only after bias and activation. - project_and_activate(static_cast(projection.get())); + ORT_THROW("CUDA QMoE FLOAT input requires dynamic_fp8 execution."); } - params.weights = static_cast(fc2_experts_weights->DataRaw()); - params.scales = fc2_scale_tensor->DataRaw(); - params.scale_type = scale_type(fc2_scale_tensor); - params.up_weights = nullptr; - params.up_scales = nullptr; - params.row_to_unpermuted = nullptr; - params.n = hidden; - params.k = inter; - params.fusion = 0; - LaunchQMoEFp8Projection(params, static_cast(activation.get()), static_cast(result.get()), stream); - const auto* fc2_bias = fc2_experts_bias_optional ? static_cast(fc2_experts_bias_optional->DataRaw()) : nullptr; - ck::finalizeMoeRoutingKernelLauncher( - static_cast(result.get()), static_cast(output->MutableDataRaw()) + row_offset * hidden, - fc2_bias, expert_scales, unpermuted_row_to_permuted_row, row_map.get(), expert_indices, expert_offsets.get(), - rows, hidden, k_, num_experts, parallelism_config, false, stream); }; - if (is_fp16_) { + if (input->IsDataType()) { + execute(static_cast(nullptr)); + } else if (is_fp16_) { execute(static_cast(nullptr)); } else { execute(static_cast<__nv_bfloat16*>(nullptr)); @@ -1760,8 +1854,11 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { } if (enable_kernel_debug_info_) { const size_t activation_bytes = SafeInt(projection_bytes) + - SafeInt(max_routes) * (inter + hidden) * element_size; - PrintQMoEKernelDebugInfo(row_tile_plan.rows_per_tile > 8 ? "fp8_fused_gemm" : "fp8_fused_gemv", + SafeInt(dynamic_fp8_ ? 0 : max_routes) * inter * element_size + + SafeInt(max_routes) * hidden * + (dynamic_fp8_ ? sizeof(float) : element_size) + + quantized_elements + (quantized_elements / 128) * sizeof(float); + PrintQMoEKernelDebugInfo(dynamic_fp8_ ? "fp8_dynamic_mma" : (row_tile_plan.rows_per_tile > 8 ? "fp8_fused_gemm" : "fp8_fused_gemv"), moe_params.num_rows, row_tile_plan.rows_per_tile, final_tile_rows, max_routes, final_tile_rows * k_, 0, total_scratch_bytes + activation_bytes); std::cout << "QMoE FP8 ExpertCapacity=" << num_experts << " DequantWeightBytes=0" << std::endl; @@ -3208,7 +3305,7 @@ Status QMoE::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, bool& is_packed, PrePackedWeights* prepacked_weights) { ORT_UNUSED_PARAMETER(prepacked_weights); is_packed = false; - if (has_int3_) return Status::OK(); + if (has_int3_ || dynamic_fp8_) return Status::OK(); if (nvfp4_weights_row_major_ && (input_idx == 2 || input_idx == 5 || (use_fp4_dequant_fallback_ && (input_idx == 3 || input_idx == 6)))) { return Status::OK(); diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h index 5e091436c01ad..e04488e453592 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h +++ b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h @@ -117,6 +117,7 @@ class QMoE final : public CudaKernel, public MoEBase { // Dequantizes FP8 weights to FP16/BF16 scratch buffers before invoking the A16 MoE runner. bool use_fp8_dequant_fallback_ = false; bool enable_fp8_fused_ = false; + bool dynamic_fp8_ = false; // WFP4AFP8 (W4A8) requires SM100+ (Blackwell) block-scaled tensor ops. On older GPUs we // dequantize MXFP4 weights to FP16/BF16 and run the dense A16 MoE runner. bool use_wfp4afp8_dequant_fallback_ = false; diff --git a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu index 9e2655d7672f6..7c36012d76d86 100644 --- a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu +++ b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu @@ -1811,6 +1811,202 @@ void LaunchQMoEFp8Projection(const QMoEFp8ProjectionParams& params, CUDA_CALL_THROW(cudaGetLastError()); } +template +__global__ void QMoEDynamicFp8QuantizeKernel(const T* input, uint8_t* output, float* scales, + int cols, int activation) { + const int row = blockIdx.x; + const int group = blockIdx.y; + const int col = group * 128 + threadIdx.x; + const int64_t index = static_cast(row) * cols + col; + float value = static_cast(input[activation == 2 ? 2 * index : index]); + if (activation) { + const float up = activation == 2 ? static_cast(input[2 * index + 1]) : 1.0f; + value = value / (1.0f + expf(-value)) * up; + } + using Reduce = cub::BlockReduce; + __shared__ typename Reduce::TempStorage storage; + __shared__ float scale; + const float amax = BlockReduceMax(fabsf(value), storage); + if (threadIdx.x == 0) { + // Preserve reciprocal-multiply rounding: a one-ULP scale change can flip an FP8 tie. + scale = fmaxf(amax, 1.0e-10f) * (1.0f / 448.0f); + scales[static_cast(row) * (cols / 128) + group] = scale; + } + __syncthreads(); + float normalized; + // Match native FP8 exporters' FP32 division; div.rn can move an E4M3 midpoint by one ULP. + asm("div.full.f32 %0, %1, %2;" : "=f"(normalized) : "f"(value), "f"(scale)); + normalized = fminf(448.0f, fmaxf(-448.0f, normalized)); + output[index] = __nv_cvt_float_to_fp8(normalized, __NV_SATFINITE, __NV_E4M3); +} + +template +void LaunchQMoEDynamicFp8Quantize(const T* input, uint8_t* quantized, float* scales, + int rows, int cols, int activation, cudaStream_t stream) { + QMoEDynamicFp8QuantizeKernel<<>>( + input, quantized, scales, cols, activation); + CUDA_CALL_THROW(cudaGetLastError()); +} + +__device__ __forceinline__ int QMoEDynamicFp8WeightRow(const QMoEFp8ProjectionParams& params, int row) { + return params.up_weights ? row / 2 + : params.fusion == 2 ? row / 2 + (row % 2) * (params.n / 2) + : row; +} + +template +__global__ void QMoEDynamicFp8ProjectionKernel(QMoEFp8ProjectionParams params, + const uint8_t* input, const float* input_scales, OutputT* output, + const float* routing_weights, const int* routed_rows, int top_k) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 890 + const int tile = blockIdx.y; + if (tile >= params.tile_offsets[params.num_experts]) return; + int first = 0; + int last = params.num_experts; + while (first < last) { + const int middle = (first + last) / 2; + if (params.tile_offsets[middle + 1] <= tile) { + first = middle + 1; + } else { + last = middle; + } + } + const int expert = first; + const int64_t row_start = params.expert_offsets[expert] + (tile - params.tile_offsets[expert]) * 16; + const int64_t row_end = params.expert_offsets[expert + 1]; + const int lane = threadIdx.x % 32; + const int group = lane / 4; + const int thread = lane % 4; + const int column_start = blockIdx.x * 32 + (threadIdx.x / 32) * 8; + float result[4] = {}; + const int k_groups = params.k / 128; + for (int k_group = 0; k_group < k_groups; ++k_group) { + float dot[4] = {}; + for (int inner = 0; inner < 128; inner += 32) { + uint32_t a[4] = {}; + uint32_t b[2] = {}; + // PTX m16n8k32: four A registers cover two rows and two K half-tiles; + // two B registers cover one output column and both K half-tiles. + for (int half_k = 0; half_k < 2; ++half_k) { + for (int byte = 0; byte < 4; ++byte) { + const int k = k_group * 128 + inner + half_k * 16 + thread * 4 + byte; + for (int half_m = 0; half_m < 2; ++half_m) { + const int64_t row = row_start + group + half_m * 8; + if (row < row_end) { + const int source = params.row_to_unpermuted ? params.row_to_unpermuted[row] % params.num_rows + : static_cast(row); + a[half_k * 2 + half_m] |= + static_cast(input[static_cast(source) * params.k + k]) << (8 * byte); + } + } + const int col = column_start + group; + const bool up = params.up_weights && (col % 2); + const uint8_t* weights = up ? params.up_weights : params.weights; + const int weight_row = QMoEDynamicFp8WeightRow(params, col); + const int weight_n = params.up_weights ? params.n / 2 : params.n; + b[half_k] |= static_cast( + weights[(static_cast(expert) * weight_n + weight_row) * params.k + k]) + << (8 * byte); + } + } + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};" + : "+f"(dot[0]), "+f"(dot[1]), "+f"(dot[2]), "+f"(dot[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + } + for (int element = 0; element < 4; ++element) { + const int64_t row = row_start + group + (element / 2) * 8; + const int col = column_start + thread * 2 + element % 2; + if (row < row_end) { + const int source = params.row_to_unpermuted ? params.row_to_unpermuted[row] % params.num_rows + : static_cast(row); + const bool up = params.up_weights && (col % 2); + const auto* scales = static_cast(up ? params.up_scales : params.scales); + const int weight_n = params.up_weights ? params.n / 2 : params.n; + const int weight_row = QMoEDynamicFp8WeightRow(params, col); + const float weight_scale = + scales[(static_cast(expert) * (weight_n / 128) + weight_row / 128) * k_groups + k_group]; + const float activation_scale = input_scales[static_cast(source) * k_groups + k_group]; + result[element] += dot[element] * activation_scale * weight_scale; + } + } + } + for (int element = 0; element < 4; ++element) { + const int64_t row = row_start + group + (element / 2) * 8; + const int col = column_start + thread * 2 + element % 2; + if (row < row_end) { + float value = result[element]; + if (routing_weights) { + const int route = routed_rows[row]; + const int token = route % params.num_rows; + const int slot = route / params.num_rows; + // The router's coefficients narrow before multiplication, not its FP32 GEMM output. + value *= static_cast(static_cast(routing_weights[token * top_k + slot])); + } + output[row * params.n + col] = static_cast(value); + } + } +#else + asm volatile("trap;"); +#endif +} + +bool QMoEDynamicFp8Supported() { + cudaFuncAttributes attributes{}; + CUDA_CALL_THROW(cudaFuncGetAttributes(&attributes, QMoEDynamicFp8ProjectionKernel)); + return attributes.binaryVersion >= 89 && attributes.ptxVersion >= 89; +} + +template +void LaunchQMoEDynamicFp8Projection(const QMoEFp8ProjectionParams& params, + const uint8_t* input, const float* input_scales, OutputT* output, + const float* routing_weights, const int* routed_rows, int top_k, + cudaStream_t stream) { + const int tiles = std::min(params.expanded_rows, + std::min(params.num_experts, params.expanded_rows) + params.expanded_rows / 16); + QMoEDynamicFp8ProjectionKernel<<>>( + params, input, input_scales, output, routing_weights, routed_rows, top_k); + CUDA_CALL_THROW(cudaGetLastError()); +} + +template +__global__ void QMoEDynamicFp8FinalizeKernel(const float* input, T* output, const int* row_map, + int rows, int cols, int top_k) { + const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index >= static_cast(rows) * cols) return; + const int token = static_cast(index / cols); + const int col = static_cast(index % cols); + float sum = 0.0f; + for (int slot = 0; slot < top_k; ++slot) { + sum += input[static_cast(row_map[token + slot * rows]) * cols + col]; + } + output[index] = static_cast(sum); +} + +template +void LaunchQMoEDynamicFp8Finalize(const float* input, T* output, const int* unpermuted_to_permuted, + int rows, int cols, int top_k, cudaStream_t stream) { + const int64_t count = static_cast(rows) * cols; + QMoEDynamicFp8FinalizeKernel<<((count + 255) / 256), 256, 0, stream>>>( + input, output, unpermuted_to_permuted, rows, cols, top_k); + CUDA_CALL_THROW(cudaGetLastError()); +} + +#define INSTANTIATE_DYNAMIC_FP8(T) \ + template void LaunchQMoEDynamicFp8Quantize(const T*, uint8_t*, float*, int, int, int, cudaStream_t); \ + template void LaunchQMoEDynamicFp8Projection(const QMoEFp8ProjectionParams&, const uint8_t*, const float*, T*, \ + const float*, const int*, int, cudaStream_t); \ + template void LaunchQMoEDynamicFp8Finalize(const float*, T*, const int*, int, int, int, cudaStream_t); +INSTANTIATE_DYNAMIC_FP8(half) +INSTANTIATE_DYNAMIC_FP8(__nv_bfloat16) +INSTANTIATE_DYNAMIC_FP8(float) +#undef INSTANTIATE_DYNAMIC_FP8 +template void LaunchQMoEDynamicFp8Projection( + const QMoEFp8ProjectionParams&, const uint8_t*, const float*, float*, const float*, const int*, int, cudaStream_t); +template void LaunchQMoEDynamicFp8Projection<__nv_bfloat16, float>( + const QMoEFp8ProjectionParams&, const uint8_t*, const float*, float*, const float*, const int*, int, cudaStream_t); + template __global__ void QMoEFp8ActivationKernel(const InputT* input, const T* bias, T* output, const int* experts, int rows, int inter_size, int fusion, diff --git a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h index f6b34ad81ff8f..d84b8c17dcb57 100644 --- a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h +++ b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h @@ -311,6 +311,23 @@ void LaunchQMoEFp8Projection(const QMoEFp8ProjectionParams& params, bool QMoEFp8GemvSupportsNativeDecode(); bool QMoEFp8GemvKernelUsesNativeDecode(const void* kernel); +bool QMoEDynamicFp8Supported(); + +// activation: 0 = input, 1 = SiLU, 2 = SiLU(gate) * up with interleaved gate/up columns. +template +void LaunchQMoEDynamicFp8Quantize(const T* input, uint8_t* quantized, float* scales, + int rows, int cols, int activation, cudaStream_t stream); + +template +void LaunchQMoEDynamicFp8Projection(const QMoEFp8ProjectionParams& params, + const uint8_t* input, const float* input_scales, OutputT* output, + const float* routing_weights, const int* routed_rows, int top_k, + cudaStream_t stream); + +template +void LaunchQMoEDynamicFp8Finalize(const float* input, T* output, const int* unpermuted_to_permuted, + int rows, int cols, int top_k, cudaStream_t stream); + template void LaunchQMoEFp8Activation(const InputT* input, const T* bias, T* output, const int* experts, int rows, int inter_size, int fusion, diff --git a/onnxruntime/contrib_ops/webgpu/moe/qmoe.h b/onnxruntime/contrib_ops/webgpu/moe/qmoe.h index 2d74a2a55100f..aaa80da48cd29 100755 --- a/onnxruntime/contrib_ops/webgpu/moe/qmoe.h +++ b/onnxruntime/contrib_ops/webgpu/moe/qmoe.h @@ -19,6 +19,8 @@ using onnxruntime::webgpu::ComputeContext; class QMoE final : public MoE { public: QMoE(const OpKernelInfo& info) : MoE(info) { + ORT_ENFORCE(info.GetAttrOrDefault("activation_quantization", "none") == "none", + "WebGPU QMoE does not support dynamic FP8 activation quantization."); ORT_ENFORCE(info.GetAttr("expert_weight_bits", &expert_weight_bits_).IsOK()); ORT_ENFORCE(expert_weight_bits_ == 2 || expert_weight_bits_ == 4 || expert_weight_bits_ == 8, "expert_weight_bits must be 2, 4, or 8, but got ", expert_weight_bits_); diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index e05119bd05dfa..7a1567aa75c74 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -1491,11 +1491,14 @@ constexpr const char* qMoE_ver1_doc = R"DOC( dequantized_weight[e, n, k] = float(weight[e, n, k]) * scale[e, n / block_size, k / block_size] Partial edge blocks are allowed. No zero points or activation scales are accepted. Without a positive block_size, FP8 uses the legacy per-expert fc*_global_scale inputs instead. - Block-scaled FP8 does not use global scales. Activations retain the input type (weight-only quantization). - On CUDA SM80-or-newer GPUs, SiLU/SwiGLU block-scaled FP8 uses fused weight-only GEMV/GEMM by default. + Block-scaled FP8 does not use global scales. By default activations retain the input type (weight-only quantization). + activation_quantization='dynamic_fp8' instead selects dynamic per-token K128 activation quantization, + with native FP8 dot products, FP32 block-scale accumulation and FP32 routed down-projection reduction. + On CUDA SM80-or-newer GPUs, weight-only SiLU/SwiGLU block-scaled FP8 uses fused GEMV/GEMM by default. These kernels decode and scale FP8 tiles on chip without materializing full FP16/BF16 weight buffers. Set ORT_ENABLE_FP8_FUSED=0 before session creation to select the compact dense-dequant fallback. - Other activation types and unsupported routing sizes retain the fallback. + Other activation types and unsupported routing sizes retain the fallback in weight-only mode. + Dynamic activation quantization rejects unsupported configurations instead of falling back to weight-only execution. The WebGPU block-FP8 kernel supports block_size=128 with float32 scales and rejects projections that exceed 32-bit shader addressing or the device's per-dimension dispatch limit. @@ -1519,8 +1522,152 @@ constexpr const char* qMoE_ver1_doc = R"DOC( When swiglu_fusion=2, two GEMMs are fused, and g and l are concatenated on each row. )DOC"; +static void QMoEDynamicFp8Inference(ONNX_NAMESPACE::InferenceContext& ctx) { + using namespace ONNX_NAMESPACE; + const auto* quant_type = ctx.getAttribute("quant_type"); + if (!quant_type || quant_type->s() != "fp8" || getAttribute(ctx, "block_size", 0) != 128) { + fail_shape_inference("QMoE dynamic_fp8 requires quant_type='fp8' and block_size=128."); + } + const int64_t bits = getAttribute(ctx, "expert_weight_bits", 4); + for (const char* name : {"expert_weight_bits", "fc1_expert_weight_bits", "fc2_expert_weight_bits", "fc3_expert_weight_bits"}) { + if (getAttribute(ctx, name, bits) != 8) fail_shape_inference("QMoE dynamic_fp8 requires 8-bit expert weights."); + } + const auto* activation = ctx.getAttribute("activation_type"); + const int64_t fusion = getAttribute(ctx, "swiglu_fusion", 0); + const bool valid_activation = activation && + ((activation->s() == "silu" && fusion == 0) || + (activation->s() == "swiglu" && (fusion == 1 || fusion == 2))); + if (!valid_activation || + getAttribute(ctx, "use_sparse_mixer", 0) != 0) { + fail_shape_inference("QMoE dynamic_fp8 requires SiLU or fused SwiGLU and softmax/top-k routing."); + } + for (auto parameter : {std::make_pair("activation_alpha", 1.0f), std::make_pair("activation_beta", 0.0f)}) { + const auto* value = ctx.getAttribute(parameter.first); + if (value && value->f() != parameter.second) fail_shape_inference("QMoE dynamic_fp8 requires standard SiLU/SwiGLU."); + } + const auto* limit = ctx.getAttribute("swiglu_limit"); + if (limit && limit->f() != std::numeric_limits::infinity()) { + fail_shape_inference("QMoE dynamic_fp8 does not support activation clipping."); + } + auto input = [&](size_t index) -> const TypeProto* { + return index < ctx.getNumInputs() ? ctx.getInputType(index) : nullptr; + }; + auto type = [&](size_t index, int expected, bool required) { + const auto* value = input(index); + if (!value) { + if (required) fail_type_inference("QMoE dynamic_fp8 requires input ", index, "."); + return; + } + if (expected && value->has_tensor_type() && value->tensor_type().elem_type() && + value->tensor_type().elem_type() != expected) { + fail_type_inference("QMoE dynamic_fp8 input ", index, " has an incompatible tensor element type."); + } + }; + const auto* x = input(0); + const int element_type = x && x->has_tensor_type() ? x->tensor_type().elem_type() : 0; + if (element_type && element_type != TensorProto::FLOAT && + element_type != TensorProto::FLOAT16 && element_type != TensorProto::BFLOAT16) { + fail_type_inference("QMoE dynamic_fp8 requires FP32, FP16 or BF16 activations."); + } + type(1, element_type, true); + for (size_t index : {2, 5, 8}) type(index, TensorProto::FLOAT8E4M3FN, index != 8); + for (size_t index : {3, 6, 9}) type(index, TensorProto::FLOAT, index != 9 || input(8) != nullptr); + for (size_t index : {4, 7, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20}) { + if (input(index)) fail_shape_inference("QMoE dynamic_fp8 does not accept input ", index, "."); + } + if ((fusion && input(8)) || (!input(8) && input(9))) { + fail_shape_inference("QMoE dynamic_fp8 FC3 weights/scales require separate SiLU projections."); + } + auto shape = [&](size_t index, int rank) -> const TensorShapeProto* { + const auto* value = input(index); + if (!value || !value->has_tensor_type() || !value->tensor_type().has_shape()) return nullptr; + const auto& result = value->tensor_type().shape(); + if (rank && result.dim_size() != rank) { + fail_shape_inference("QMoE dynamic_fp8 input ", index, " has an invalid rank."); + } + return &result; + }; + auto dim = [](const TensorShapeProto* value, int axis) -> const TensorShapeProto::Dimension* { + return value ? &value->dim(axis) : nullptr; + }; + auto equal = [&](const TensorShapeProto::Dimension* a, const TensorShapeProto::Dimension* b, int divisor = 1) { + if (a && b && a->has_dim_value() && b->has_dim_value() && a->dim_value() != b->dim_value() / divisor) { + fail_shape_inference("QMoE dynamic_fp8 has incompatible expert, projection or scale dimensions."); + } + }; + const auto* xs = shape(0, 0); + if (xs && xs->dim_size() != 2 && xs->dim_size() != 3) { + fail_shape_inference("QMoE dynamic_fp8 input must have rank 2 or 3."); + } + const auto* router = shape(1, 2); + const auto* fc1 = shape(2, 3); + const auto* fc2 = shape(5, 3); + const auto* fc3 = shape(8, 3); + const int64_t top_k = getAttribute(ctx, "k", 1); + if (top_k <= 0 || (fc1 && fc1->dim(0).has_dim_value() && top_k > fc1->dim(0).dim_value())) { + fail_shape_inference("QMoE dynamic_fp8 requires 0 < k <= num_experts."); + } + if (xs) { + int64_t tokens = 1; + bool known_tokens = true; + for (int axis = 0; axis < xs->dim_size() - 1; ++axis) { + const auto& d = xs->dim(axis); + if (!d.has_dim_value()) { + known_tokens = false; + continue; + } + if (d.dim_value() < 0 || + (d.dim_value() > 0 && tokens > std::numeric_limits::max() / d.dim_value())) { + fail_shape_inference("QMoE dynamic_fp8 token dimensions are invalid or overflow int64."); + } + tokens *= d.dim_value(); + } + if (known_tokens) { + if (tokens > 65535 / top_k) fail_shape_inference("QMoE dynamic_fp8 supports at most 65535 token/expert pairs."); + const auto* router_tokens = dim(router, 0); + if (router_tokens && router_tokens->has_dim_value() && router_tokens->dim_value() != tokens) { + fail_shape_inference("QMoE dynamic_fp8 router token count must match the flattened input."); + } + } + } + for (const auto* weight : {fc1, fc2, fc3}) { + if (!weight) continue; + for (int axis = 0; axis < 3; ++axis) { + const auto& d = weight->dim(axis); + if (d.has_dim_value() && (d.dim_value() <= 0 || (axis > 0 && d.dim_value() % 128 != 0))) { + fail_shape_inference("QMoE dynamic_fp8 requires positive expert dimensions and projections aligned to 128."); + } + } + equal(dim(weight, 0), dim(fc1, 0)); + } + equal(dim(router, 1), dim(fc1, 0)); + equal(dim(xs, xs ? xs->dim_size() - 1 : 0), dim(fc1, 2)); + equal(dim(fc2, 1), dim(fc1, 2)); + equal(dim(fc2, 2), dim(fc1, 1), fusion ? 2 : 1); + if (fc1 && fusion && fc1->dim(1).has_dim_value() && fc1->dim(1).dim_value() % 256 != 0) { + fail_shape_inference("QMoE dynamic_fp8 fused intermediate size must be aligned to 128."); + } + equal(dim(fc3, 1), dim(fc1, 1)); + equal(dim(fc3, 2), dim(fc1, 2)); + for (auto indices : {std::make_pair(2, 3), std::make_pair(5, 6), std::make_pair(8, 9)}) { + const auto* weight = shape(indices.first, 3); + const auto* scale = shape(indices.second, 3); + for (int axis = 0; axis < 3; ++axis) equal(dim(scale, axis), dim(weight, axis), axis == 0 ? 1 : 128); + } +} + static void QMoETypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) { using namespace ONNX_NAMESPACE; + const auto* activation_quantization = ctx.getAttribute("activation_quantization"); + const std::string mode = activation_quantization ? activation_quantization->s() : "none"; + if (mode != "none" && mode != "dynamic_fp8") { + fail_shape_inference("QMoE activation_quantization must be 'none' or 'dynamic_fp8'."); + } + if (mode == "dynamic_fp8") { + propagateShapeAndTypeFromFirstInput(ctx); + QMoEDynamicFp8Inference(ctx); + return; + } const int64_t root_bits = getAttribute(ctx, "expert_weight_bits", 4); const int64_t fc1_bits = getAttribute(ctx, "fc1_expert_weight_bits", root_bits); const int64_t fc2_bits = getAttribute(ctx, "fc2_expert_weight_bits", root_bits); @@ -1717,7 +1864,7 @@ constexpr const char* qMoE_int3_doc = R"DOC( static OpSchema QMoESchema() { constexpr const char* quant_type_doc = "Quantization type: 'int' for integer quantization (default), 'fp4' for MXFP4 quantization, " - "'nvfp4' for NVFP4 quantization, 'fp8' for FP8 e4m3 weight-only quantization, " + "'nvfp4' for NVFP4 quantization, 'fp8' for FP8 e4m3 weights (weight-only unless activation_quantization is set), " "or 'wfp4afp8' for MXFP4 weight with FP8 activation. " "When quant_type is 'fp4' or 'nvfp4', weights are stored in E2M1 FP4 format (2 values per byte), " "fc*_scales inputs contain the FP4 block scales, and fc*_global_scale inputs must be provided. " @@ -1727,6 +1874,18 @@ static OpSchema QMoESchema() { "The INT3 profile requires integer quantization."; auto schema = OpSchema() .SetDoc(std::string(qMoE_ver1_doc) + qMoE_int3_doc) + .Attr("activation_quantization", + "'none' (default) retains weight-only execution. 'dynamic_fp8' quantizes consecutive " + "groups of 128 features per token to saturating nearest-even E4M3FN " + "using FP32 div.full normalization and scale " + "max(amax(abs(x)),1e-10) * float32(1/448). Requires block-scaled FP8 weights, block_size=128, " + "FP32 weight scales, FP32/FP16/BF16 inputs and 128-aligned projections. Gate/up outputs " + "narrow to the input type; their FP32 SiLU/SwiGLU result is quantized directly for FC2. " + "FC2 stays FP32, multiplies routing weights rounded to the input type, and reduces " + "in FP32 before the final output cast. Biases, supplied activation/global scales, " + "zero points, router_weights and nonstandard activation parameters are unsupported. " + "CUDA requires native FP8 MMA (SM89+) and at most 65535 token/expert pairs.", + AttributeProto::STRING, std::string("none")) .Attr("activation_type", "Activation function to use. Choose from relu, gelu, silu, swiglu and identity. CUDA also supports relu2 (squared ReLU). Default is relu", AttributeProto::STRING, 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..67d6b6f206a75 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 @@ -53,6 +53,120 @@ struct CudaBuffer { size_t bytes; }; +template +void TestDynamicFp8Quantization() { + std::array input{}; + input[128] = T(448.0f); + input[129] = T(-448.0f); + input[130] = T(1.0625f); + input[131] = T(1.1875f); + input[256] = T(0.65625f); + input[257] = T(57.0f / 32768.0f); + CudaBuffer source(sizeof(input)); + CudaBuffer quantized(input.size()); + CudaBuffer scales(3 * sizeof(float)); + source.Upload(input.data()); + qmoe::LaunchQMoEDynamicFp8Quantize(source.As(), quantized.As(), scales.As(), + 1, 384, 0, nullptr); + std::array codes{}; + std::array actual_scales{}; + quantized.Download(codes.data()); + scales.Download(actual_scales.data()); + EXPECT_FLOAT_EQ(actual_scales[0], 1.0e-10f / 448.0f); + EXPECT_FLOAT_EQ(actual_scales[1], 1.0f); + for (int col = 0; col < 128; ++col) EXPECT_EQ(codes[col], 0); + EXPECT_EQ(codes[128], 0x7e); + EXPECT_EQ(codes[129], 0xfe); + EXPECT_EQ(codes[130], 0x38); + EXPECT_EQ(codes[131], 0x3a); + // Replacing reciprocal multiplication with division shifts this scale by one ULP and flips the tie. + EXPECT_EQ(actual_scales[2], 0.65625f * (1.0f / 448.0f)); + EXPECT_EQ(codes[257], 0x39); +} + +TEST(CUDA_EP_Unittest, QMoEDynamicFp8QuantizationScalesAndNearestEven) { + TestDynamicFp8Quantization(); + TestDynamicFp8Quantization(); + TestDynamicFp8Quantization<__nv_bfloat16>(); +} + +TEST(CUDA_EP_Unittest, QMoEDynamicFp8QuantizesActivationWithoutBFloat16RoundTrip) { + std::array<__nv_bfloat16, 256> input{}; + input[0] = __nv_bfloat16(32.0f); + input[1] = __nv_bfloat16(14.0f); + input[2] = __nv_bfloat16(0.0078125f); + input[3] = __nv_bfloat16(0.25f); + CudaBuffer source(sizeof(input)); + CudaBuffer quantized(128); + CudaBuffer scales(sizeof(float)); + source.Upload(input.data()); + qmoe::LaunchQMoEDynamicFp8Quantize(source.As<__nv_bfloat16>(), quantized.As(), scales.As(), + 1, 128, 2, nullptr); + std::array codes{}; + quantized.Download(codes.data()); + // The FP32 product is just above the E4M3 zero/subnormal midpoint; BF16 would round to the midpoint. + EXPECT_EQ(codes[0], 0x7e); + EXPECT_EQ(codes[1], 0x01); +} + +TEST(CUDA_EP_Unittest, QMoEDynamicFp8FloatNormalizationPreservesMidpoint) { + std::array input{}; + input[0] = 0.783361554145813f; + input[1] = -0.20283471047878265f; + CudaBuffer source(sizeof(input)), quantized(input.size()), scales(sizeof(float)); + source.Upload(input.data()); + qmoe::LaunchQMoEDynamicFp8Quantize(source.As(), quantized.As(), scales.As(), + 1, 128, 0, nullptr); + std::array codes{}; + quantized.Download(codes.data()); + // div.full resolves this to -116 (the midpoint); div.rn yields the next float below -116. + EXPECT_EQ(codes[1], 0xee); +} + +template +void TestDynamicFp8RoutingWeightRounding() { + std::vector weights(128 * 128, 0x38); + std::vector input(128, 0x38); + const float scale = 1.0f; + const float coefficient = 0.1f; + const int route = 0; + const std::array tiles{0, 1}; + const std::array offsets{0, 1}; + CudaBuffer w(weights.size()), x(input.size()), s(sizeof(float)), r(sizeof(float)); + CudaBuffer map(sizeof(int)), tile_offsets(sizeof(tiles)), expert_offsets(sizeof(offsets)); + CudaBuffer output(128 * sizeof(float)); + w.Upload(weights.data()); + x.Upload(input.data()); + s.Upload(&scale); + r.Upload(&coefficient); + map.Upload(&route); + tile_offsets.Upload(tiles.data()); + expert_offsets.Upload(offsets.data()); + qmoe::QMoEFp8ProjectionParams params; + params.weights = w.As(); + params.scales = s.As(); + params.tile_offsets = tile_offsets.As(); + params.expert_offsets = expert_offsets.As(); + params.num_experts = params.num_rows = params.expanded_rows = 1; + params.n = params.k = params.block_size = 128; + qmoe::LaunchQMoEDynamicFp8Projection( + params, x.As(), s.As(), output.As(), r.As(), map.As(), 1, nullptr); + std::array actual{}; + output.Download(actual.data()); + const float expected = 128.0f * static_cast(T(coefficient)); + if constexpr (!std::is_same_v) { + EXPECT_NE(expected, 128.0f * coefficient); + } + for (float value : actual) EXPECT_FLOAT_EQ(value, expected); +} + +TEST(CUDA_EP_Unittest, QMoEDynamicFp8RoundsRoutingWeightsBeforeFp32Product) { + if (!qmoe::QMoEDynamicFp8Supported()) GTEST_SKIP() << "Native FP8 MMA requires an SM89+ kernel."; + TestDynamicFp8RoutingWeightRounding(); + TestDynamicFp8RoutingWeightRounding<__nv_bfloat16>(); + TestDynamicFp8RoutingWeightRounding(); +} + TEST(CUDA_EP_Unittest, QMoEFp8CompactExpertsPreservesRoutesAndClearsUnusedSlots) { constexpr int num_experts = 512; constexpr int capacity = 6; diff --git a/onnxruntime/test/contrib_ops/qmoe_shape_inference_test.cc b/onnxruntime/test/contrib_ops/qmoe_shape_inference_test.cc index d29cac6463c1f..8a53d702769f3 100644 --- a/onnxruntime/test/contrib_ops/qmoe_shape_inference_test.cc +++ b/onnxruntime/test/contrib_ops/qmoe_shape_inference_test.cc @@ -31,6 +31,8 @@ struct QMoEShapeCase { std::optional prepacked = 0; std::optional block = 32; std::optional quant_type = "int"; + std::optional activation_quantization; + int64_t top_k = 1; int64_t fusion = 1; std::string activation = "swiglu"; bool omit_activation = false; @@ -75,12 +77,14 @@ common::Status ResolveQMoE(const QMoEShapeCase& spec, auto& output = graph.GetOrCreateNodeArg("output", nullptr); auto& node = graph.AddNode("qmoe", "QMoE", "", inputs, {&output}, nullptr, kMSDomain); node.AddAttribute("expert_weight_bits", spec.root_bits); + node.AddAttribute("k", spec.top_k); if (spec.fc1_bits) node.AddAttribute("fc1_expert_weight_bits", *spec.fc1_bits); if (spec.fc2_bits) node.AddAttribute("fc2_expert_weight_bits", *spec.fc2_bits); if (spec.fc3_bits) node.AddAttribute("fc3_expert_weight_bits", *spec.fc3_bits); if (spec.prepacked) node.AddAttribute("weights_prepacked", *spec.prepacked); if (spec.block) node.AddAttribute("block_size", *spec.block); if (spec.quant_type) node.AddAttribute("quant_type", *spec.quant_type); + if (spec.activation_quantization) node.AddAttribute("activation_quantization", *spec.activation_quantization); node.AddAttribute("swiglu_fusion", spec.fusion); if (!spec.omit_activation) node.AddAttribute("activation_type", spec.activation); graph.SetOutputs({&output}); @@ -89,6 +93,49 @@ common::Status ResolveQMoE(const QMoEShapeCase& spec, return status; } +QMoEShapeCase DynamicFp8Case() { + QMoEShapeCase spec; + spec.root_bits = 8; + spec.fc1_bits.reset(); + spec.fc2_bits.reset(); + spec.fc3_bits.reset(); + spec.quant_type = "fp8"; + spec.activation_quantization = "dynamic_fp8"; + spec.block = 128; + spec.shapes = {{0, {2, 256}}, {1, {2, 3}}, {2, {3, 256, 256}}, {3, {3, 2, 2}}, {5, {3, 256, 128}}, {6, {3, 2, 1}}}; + spec.types = {{2, TensorProto::FLOAT8E4M3FN}, {5, TensorProto::FLOAT8E4M3FN}, {3, TensorProto::FLOAT}, {6, TensorProto::FLOAT}}; + return spec; +} + +TEST(QMoEShapeInferenceTest, DynamicFp8PreservesActivationShapeAndType) { + for (int element_type : {TensorProto::FLOAT, TensorProto::FLOAT16, TensorProto::BFLOAT16}) { + for (int64_t fusion : {0, 1, 2}) { + auto spec = DynamicFp8Case(); + spec.activation_type = element_type; + spec.fusion = fusion; + if (!fusion) { + spec.activation = "silu"; + spec.shapes[2][1] = 128; + spec.shapes[3][1] = 1; + spec.shapes[8] = spec.shapes[2]; + spec.shapes[9] = spec.shapes[3]; + spec.types[8] = TensorProto::FLOAT8E4M3FN; + spec.types[9] = TensorProto::FLOAT; + } + ASSERT_STATUS_OK(ResolveQMoE(spec, [&](const Node&, const NodeArg& output) { + EXPECT_EQ(output.TypeAsProto()->tensor_type().elem_type(), element_type); + ASSERT_NE(output.Shape(), nullptr); + EXPECT_EQ(output.Shape()->dim(1).dim_value(), 256); + })); + spec.symbolic_dimension = {{0, 0}}; + ASSERT_STATUS_OK(ResolveQMoE(spec)); + spec.symbolic_dimension.reset(); + spec.shapes[0] = {1, 2, 256}; + ASSERT_STATUS_OK(ResolveQMoE(spec)); + } + } +} + TEST(QMoEShapeInferenceTest, Int3BaselineCombinationsAndTypes) { for (int element_type : {TensorProto::FLOAT16, TensorProto::BFLOAT16}) { for (int64_t block : {32, 64, 128}) { @@ -292,6 +339,41 @@ TEST(QMoEShapeInferenceTest, OtherMicrosoftOperatorsStillResolveToVersion1) { #ifndef ORT_NO_EXCEPTIONS +TEST(QMoEShapeInferenceTest, RejectsInvalidDynamicFp8Metadata) { + using Mutation = std::pair>; + const std::vector mutations{ + {"mode", [](auto& s) { s.activation_quantization = "dynamic_int8"; }}, + {"quant type", [](auto& s) { s.quant_type = "int"; }}, + {"block", [](auto& s) { s.block = 64; }}, + {"bits", [](auto& s) { s.fc2_bits = 4; }}, + {"zero k", [](auto& s) { s.top_k = 0; }}, + {"large k", [](auto& s) { s.top_k = 4; }}, + {"activation", [](auto& s) { s.activation = "relu"; }}, + {"fusion", [](auto& s) { s.fusion = 0; }}, + {"input type", [](auto& s) { s.activation_type = TensorProto::DOUBLE; }}, + {"weight type", [](auto& s) { s.types[2] = TensorProto::UINT8; }}, + {"scale type", [](auto& s) { s.types[6] = TensorProto::FLOAT16; }}, + {"missing scale", [](auto& s) { s.shapes.erase(3); }}, + {"extra scale", [](auto& s) { s.shapes[17] = {1}; }}, + {"bias", [](auto& s) { s.shapes[7] = {3, 256}; }}, + {"rank", [](auto& s) { s.shapes[2] = {256, 256}; }}, + {"unaligned hidden", [](auto& s) { s.shapes[2][2] = 160; }}, + {"unaligned intermediate", [](auto& s) { s.shapes[2][1] = 128; }}, + {"experts", [](auto& s) { s.shapes[5][0] = 4; }}, + {"projection", [](auto& s) { s.shapes[5][2] = 256; }}, + {"scale shape", [](auto& s) { s.shapes[3][2] = 3; }}, + {"scale rank", [](auto& s) { s.shapes[6] = {3, 2}; }}, + {"router tokens", [](auto& s) { s.shapes[1][0] = 3; }}, + {"too many routes", [](auto& s) { s.shapes[0][0] = s.shapes[1][0] = 65536; }}, + }; + for (const auto& [name, mutate] : mutations) { + SCOPED_TRACE(name); + auto spec = DynamicFp8Case(); + mutate(spec); + EXPECT_FALSE(ResolveQMoE(spec).IsOK()); + } +} + TEST(QMoEShapeInferenceTest, RejectsInvalidInt3Metadata) { using Mutation = std::pair>; const std::vector mutations{ diff --git a/onnxruntime/test/python/transformers/test_qmoe_fp8_cuda.py b/onnxruntime/test/python/transformers/test_qmoe_fp8_cuda.py index b9dd703f3b6af..2b3473bcf104d 100644 --- a/onnxruntime/test/python/transformers/test_qmoe_fp8_cuda.py +++ b/onnxruntime/test/python/transformers/test_qmoe_fp8_cuda.py @@ -4,6 +4,7 @@ # license information. # -------------------------------------------------------------------------- +import json import os import unittest @@ -133,7 +134,9 @@ def create_fp8_moe_onnx_graph( return model.SerializeToString() -def create_block_fp8_moe_graph(tensors, top_k, onnx_dtype, block_size, fusion=0, normalize=1, initializers=False): +def create_block_fp8_moe_graph( + tensors, top_k, onnx_dtype, block_size, fusion=0, normalize=1, initializers=False, activation_quantization="none" +): """Use runtime weights so the official 512-expert shape needs no >2GB protobuf.""" names = [""] * 17 indices = { @@ -190,6 +193,8 @@ def create_block_fp8_moe_graph(tensors, top_k, onnx_dtype, block_size, fusion=0, quant_type="fp8", block_size=block_size, ) + if activation_quantization != "none": + node.attribute.append(helper.make_attribute("activation_quantization", activation_quantization)) graph = helper.make_graph( [node], "BlockFP8", @@ -365,9 +370,16 @@ def _session( initializers=False, row_tile_size=0, enable_cuda_graph=False, + activation_quantization="none", ): - dtype = TensorProto.BFLOAT16 if tensors["input"].dtype == torch.bfloat16 else TensorProto.FLOAT16 - model, input_types = create_block_fp8_moe_graph(tensors, top_k, dtype, block, fusion, normalize, initializers) + dtype = { + torch.float32: TensorProto.FLOAT, + torch.float16: TensorProto.FLOAT16, + torch.bfloat16: TensorProto.BFLOAT16, + }[tensors["input"].dtype] + model, input_types = create_block_fp8_moe_graph( + tensors, top_k, dtype, block, fusion, normalize, initializers, activation_quantization + ) options = onnxruntime.SessionOptions() options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL options.add_session_config_entry("session.disable_cpu_ep_fallback", "1") @@ -389,11 +401,16 @@ def _execute(cls, tensors, block=128, fusion=0, top_k=10, normalize=1, session=N name = graph_input.name tensor = tensors[name] binding.bind_input(name, "cuda", 0, input_types[name], tensor.shape, tensor.data_ptr()) - binding.bind_output("output", "cuda", 0, input_types["input"], output.shape, output.data_ptr()) + if output.numel() == 0: + binding.bind_output("output", "cuda", 0) + else: + binding.bind_output("output", "cuda", 0, input_types["input"], output.shape, output.data_ptr()) torch.cuda.synchronize() binding.synchronize_inputs() session.run_with_iobinding(binding) binding.synchronize_outputs() + if output.numel() == 0: + assert binding.get_outputs()[0].shape() == list(output.shape) return output @staticmethod @@ -716,12 +733,27 @@ def test_block_fp8_fused_gemm_partial_tiles(monkeypatch, dtype, fusion, row_tile reason="SM80+ CUDA FP8 QMoE required", ) @pytest.mark.parametrize("tokens", [1, 32]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_block_fp8_fused_cuda_graph_replay(capfd, monkeypatch, tokens, dtype): +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("activation_quantization", ["none", "dynamic_fp8"]) +def test_block_fp8_fused_cuda_graph_replay(capfd, monkeypatch, tokens, dtype, activation_quantization): + dynamic = activation_quantization == "dynamic_fp8" + if dtype == torch.float32 and not dynamic: + pytest.skip("CUDA FLOAT QMoE supports only dynamic_fp8") + if dynamic and torch.cuda.get_device_capability() < (8, 9): + pytest.skip("Native FP8 MMA requires SM89+") monkeypatch.setenv("ORT_ENABLE_FP8_FUSED", "1") monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") - tensors = TestQMoEBlockFP8._inputs(experts=32, tokens=tokens, dtype=dtype) - session, input_types = TestQMoEBlockFP8._session(tensors, enable_cuda_graph=True) + tensors = TestQMoEBlockFP8._inputs( + hidden=128 if dynamic else 160, + inter=128 if dynamic else 80, + experts=32, + tokens=tokens, + dtype=dtype, + scale_dtype=torch.float32 if dynamic else torch.bfloat16, + ) + session, input_types = TestQMoEBlockFP8._session( + tensors, enable_cuda_graph=True, activation_quantization=activation_quantization + ) output = torch.empty_like(tensors["input"]) binding = session.io_binding() for graph_input in session.get_inputs(): @@ -737,10 +769,272 @@ def test_block_fp8_fused_cuda_graph_replay(capfd, monkeypatch, tokens, dtype): torch.cuda.synchronize() session.run_with_iobinding(binding) binding.synchronize_outputs() - expected = TestQMoEBlockFP8._reference(tensors, 128, 0) + expected = ( + dynamic_fp8_reference(tensors, 0, 10) if dynamic else TestQMoEBlockFP8._reference(tensors, 128, 0) + ) torch.testing.assert_close(output.float(), expected.float(), atol=0.003, rtol=0.04) assert "DequantWeightBytes=0" in capfd.readouterr().out +def dynamic_fp8_reference(tensors, fusion, top_k, normalize=1): + """Keep scaled K128 partial sums and routed down outputs in FP32, not the input dtype.""" + x = tensors["input"] + probabilities = tensors["router_probs"].float().softmax(-1) + weights, experts = probabilities.topk(top_k, dim=-1) + if normalize: + weights = weights / weights.sum(-1, keepdim=True) + weights = weights.to(x.dtype).float() + output = torch.zeros_like(x, dtype=torch.float32) + + def quantize(value): + groups = value.float().reshape(value.shape[0], -1, 128) + scales = groups.abs().amax(-1).clamp(min=1e-10) / 448 + quantized = (groups * scales.reciprocal()[..., None]).clamp(-448, 448).to(torch.float8_e4m3fn).float() + return quantized.flatten(1), scales + + def project(quantized, scales, name, expert): + weight = tensors[f"{name}_weights"][expert].view(torch.float8_e4m3fn).float() + weight_scales = tensors[f"{name}_scales"][expert].float().repeat_interleave(128, 0) + result = torch.zeros(quantized.shape[0], weight.shape[0], device=x.device) + for group in range(weight.shape[1] // 128): + start = group * 128 + dot = quantized[:, start : start + 128] @ weight[:, start : start + 128].T + result += dot * scales[:, group, None] * weight_scales[None, :, group] + return result + + quantized, scales = quantize(x) + # Accumulate in top-k order, including narrowing the router's coefficients. + for slot in range(top_k): + routed = torch.zeros_like(x, dtype=torch.float32) + for expert in experts[:, slot].unique().tolist(): + rows = torch.where(experts[:, slot] == expert)[0] + gu = project(quantized[rows], scales[rows], "fc1", expert).to(x.dtype).float() + if fusion == 1: + gate, up = gu[:, 0::2], gu[:, 1::2] + elif fusion == 2: + gate, up = gu.chunk(2, -1) + else: + gate = gu + up = ( + project(quantized[rows], scales[rows], "fc3", expert).to(x.dtype).float() + if "fc3_weights" in tensors + else 1.0 + ) + hidden, hidden_scales = quantize(F.silu(gate) * up) + routed[rows] = project(hidden, hidden_scales, "fc2", expert) * weights[rows, slot, None] + output += routed + return output.to(x.dtype) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not has_fp8_qmoe or torch.cuda.get_device_capability() < (8, 9), + reason="SM89+ CUDA FP8 QMoE required", +) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("tokens", [1, 33]) +@pytest.mark.parametrize("fusion", [0, 1, 2]) +def test_dynamic_fp8_native_mma(capfd, monkeypatch, dtype, tokens, fusion): + monkeypatch.setenv("ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO", "1") + tensors = TestQMoEBlockFP8._inputs( + hidden=256, inter=128, experts=4, tokens=tokens, fusion=fusion, dtype=dtype, scale_dtype=torch.float32 + ) + tensors["input"][:, :128] *= 0.03125 + session = TestQMoEBlockFP8._session( + tensors, + fusion=fusion, + top_k=2, + row_tile_size=17, + initializers=tokens == 1, + activation_quantization="dynamic_fp8", + ) + actual = TestQMoEBlockFP8._execute(tensors, fusion=fusion, top_k=2, session=session) + expected = dynamic_fp8_reference(tensors, fusion, 2) + if dtype == torch.float32: + torch.testing.assert_close(actual, expected, atol=2e-6, rtol=2e-5) + else: + torch.testing.assert_close(actual.float(), expected.float(), atol=0.0003, rtol=0.02) + log = capfd.readouterr().out + assert "fp8_dynamic_mma" in log + assert "DequantWeightBytes=0" in log + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not has_fp8_qmoe or torch.cuda.get_device_capability() < (8, 9), + reason="SM89+ CUDA FP8 QMoE required", +) +@pytest.mark.parametrize("normalize", [0, 1]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_dynamic_fp8_zero_and_ungated(normalize, dtype): + tensors = TestQMoEBlockFP8._inputs( + hidden=128, inter=128, experts=4, tokens=0, dtype=dtype, scale_dtype=torch.float32 + ) + del tensors["fc3_weights"], tensors["fc3_scales"] + session = TestQMoEBlockFP8._session(tensors, top_k=2, normalize=normalize, activation_quantization="dynamic_fp8") + actual = TestQMoEBlockFP8._execute(tensors, top_k=2, session=session) + assert actual.shape == (0, 128) + tensors = TestQMoEBlockFP8._inputs( + hidden=128, inter=128, experts=4, tokens=2, dtype=dtype, scale_dtype=torch.float32 + ) + tensors["input"][0].zero_() + del tensors["fc3_weights"], tensors["fc3_scales"] + session = TestQMoEBlockFP8._session(tensors, top_k=2, normalize=normalize, activation_quantization="dynamic_fp8") + actual = TestQMoEBlockFP8._execute(tensors, top_k=2, session=session) + expected = dynamic_fp8_reference(tensors, 0, 2, normalize) + if dtype == torch.float32: + torch.testing.assert_close(actual, expected, atol=2e-6, rtol=2e-5) + else: + torch.testing.assert_close(actual.float(), expected.float(), atol=0.0003, rtol=0.02) + torch.testing.assert_close(actual[0], torch.zeros_like(actual[0]), atol=0, rtol=0) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not has_fp8_qmoe or torch.cuda.get_device_capability() < (8, 9), + reason="SM89+ CUDA FP8 QMoE required", +) +def test_dynamic_fp8_rejects_weight_only_fallback(monkeypatch): + tensors = TestQMoEBlockFP8._inputs(hidden=128, inter=128, experts=2, scale_dtype=torch.float32) + monkeypatch.setenv("ORT_ENABLE_FP8_FUSED", "0") + with pytest.raises(Exception, match="requires fused execution"): + TestQMoEBlockFP8._session(tensors, top_k=2, activation_quantization="dynamic_fp8") + + +@pytest.mark.skipif( + os.getenv("ORT_RUN_LARGE_FP8_QMOE_TEST") != "1" + or not torch.cuda.is_available() + or not has_fp8_qmoe + or torch.cuda.get_device_capability() < (8, 9), + reason="Set ORT_RUN_LARGE_FP8_QMOE_TEST=1 on an SM89+ FP8 QMoE build", +) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_dynamic_fp8_drama3_projection_shapes(dtype): + tensors = TestQMoEBlockFP8._inputs( + hidden=2048, inter=512, experts=256, tokens=3, dtype=dtype, scale_dtype=torch.float32 + ) + session = TestQMoEBlockFP8._session(tensors, top_k=8, activation_quantization="dynamic_fp8") + actual = TestQMoEBlockFP8._execute(tensors, top_k=8, session=session) + expected = dynamic_fp8_reference(tensors, 0, 8) + if dtype == torch.float32: + torch.testing.assert_close(actual, expected, atol=2e-6, rtol=2e-5) + else: + torch.testing.assert_close(actual.float(), expected.float(), atol=0.004, rtol=0.03) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not has_fp8_qmoe or torch.cuda.get_device_capability() < (8, 9), + reason="SM89+ CUDA FP8 QMoE required", +) +def test_dynamic_fp8_down_output_retains_fp32(): + tensors = TestQMoEBlockFP8._inputs( + hidden=128, inter=128, experts=2, tokens=1, dtype=torch.float16, scale_dtype=torch.float32 + ) + tensors["input"].zero_() + tensors["input"][:, 0] = 1 + tensors["router_probs"].zero_() + for name in ("fc1", "fc2", "fc3"): + tensors[f"{name}_scales"].fill_(1) + gate = torch.zeros(2, 128, 128, device=device) + up = torch.zeros_like(gate) + gate[:, :, 0] = 32 + up[:, :, 0] = 16 + down = torch.full_like(gate, 4) + down[1].neg_() + for name, value in (("fc1", gate), ("fc2", down), ("fc3", up)): + tensors[f"{name}_weights"] = value.to(torch.float8_e4m3fn).view(torch.uint8) + session = TestQMoEBlockFP8._session(tensors, top_k=2, activation_quantization="dynamic_fp8") + actual = TestQMoEBlockFP8._execute(tensors, top_k=2, session=session) + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual, dynamic_fp8_reference(tensors, 0, 2), atol=0, rtol=0) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not has_fp8_qmoe or torch.cuda.get_device_capability() < (8, 9), + reason="SM89+ CUDA FP8 QMoE required", +) +def test_dynamic_fp8_activation_does_not_round_to_bfloat16(): + tensors = TestQMoEBlockFP8._inputs( + hidden=128, inter=128, experts=1, tokens=1, dtype=torch.bfloat16, scale_dtype=torch.float32 + ) + tensors["input"].zero_() + tensors["input"][:, 0] = 1 + gate = torch.zeros(1, 128, 128, device=device) + up = torch.zeros_like(gate) + down = torch.zeros_like(gate) + gate[:, 0, 0], up[:, 0, 0] = 32, 14 + gate[:, 1, 0], up[:, 1, 0] = 0.0078125, 0.25 + down[:, :, 1] = 1 + for name, value in (("fc1", gate), ("fc2", down), ("fc3", up)): + tensors[f"{name}_weights"] = value.to(torch.float8_e4m3fn).view(torch.uint8) + tensors[f"{name}_scales"].fill_(1) + session = TestQMoEBlockFP8._session(tensors, top_k=1, activation_quantization="dynamic_fp8") + actual = TestQMoEBlockFP8._execute(tensors, top_k=1, session=session) + torch.testing.assert_close(actual, torch.full_like(actual, 2**-9), atol=0, rtol=0) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not has_fp8_qmoe or torch.cuda.get_device_capability() < (8, 9), + reason="SM89+ CUDA FP8 QMoE required", +) +def test_dynamic_fp8_float_gate_up_retains_fp32(): + tensors = TestQMoEBlockFP8._inputs( + hidden=128, inter=128, experts=1, tokens=1, dtype=torch.float32, scale_dtype=torch.float32 + ) + tensors["input"].zero_() + tensors["input"][:, 0] = 1 + weights = torch.zeros(1, 128, 128, device=device) + weights[:, :, 0] = 1 + for name in ("fc1", "fc2", "fc3"): + tensors[f"{name}_weights"] = weights.to(torch.float8_e4m3fn).view(torch.uint8) + tensors[f"{name}_scales"].fill_(1) + tensors["fc1_scales"].fill_(100000) + session = TestQMoEBlockFP8._session(tensors, top_k=1, activation_quantization="dynamic_fp8") + actual = TestQMoEBlockFP8._execute(tensors, top_k=1, session=session) + assert actual.dtype == torch.float32 + assert torch.isfinite(actual).all() and (actual > 65504).all() + torch.testing.assert_close(actual, dynamic_fp8_reference(tensors, 0, 1), atol=0, rtol=1e-6) + + +@pytest.mark.skipif(not torch.cuda.is_available() or not has_fp8_qmoe, reason="CUDA FP8 QMoE required") +def test_float_weight_only_fp8_is_not_enabled(): + tensors = TestQMoEBlockFP8._inputs(hidden=128, inter=128, experts=2, dtype=torch.float32, scale_dtype=torch.float32) + with pytest.raises(Exception, match="FLOAT input requires activation_quantization='dynamic_fp8'"): + TestQMoEBlockFP8._session(tensors, top_k=2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required to verify CPU fallback") +def test_float_integer_qmoe_still_runs_on_cpu(tmp_path): + tensors = TestQMoEBlockFP8._inputs( + hidden=128, inter=128, experts=2, tokens=2, dtype=torch.float32, scale_dtype=torch.float32 + ) + tensors["input"].zero_() + del tensors["fc3_weights"], tensors["fc3_scales"] + for name in ("fc1", "fc2"): + tensors[f"{name}_scales"] = torch.ones(2, 128, 1, device=device) + model_bytes, _ = create_block_fp8_moe_graph(tensors, 2, TensorProto.FLOAT, 128) + model = load_model_from_string(model_bytes) + for attr in model.graph.node[0].attribute: + if attr.name == "quant_type": + attr.s = b"int" + model.graph.node[0].attribute.append(helper.make_attribute("weights_prepacked", 0)) + for value in model.graph.input: + if value.name.endswith("_weights"): + value.type.tensor_type.elem_type = TensorProto.UINT8 + options = onnxruntime.SessionOptions() + options.enable_profiling = True + options.profile_file_prefix = str(tmp_path / "float_integer_qmoe") + session = onnxruntime.InferenceSession( + model.SerializeToString(), + options, + providers=[resolve_cuda_plugin_ep("CUDAExecutionProvider"), "CPUExecutionProvider"], + ) + actual = session.run(None, {name: tensor.cpu().numpy() for name, tensor in tensors.items()})[0] + numpy.testing.assert_array_equal(actual, numpy.zeros((2, 128), dtype=numpy.float32)) + with open(session.end_profiling()) as profile_file: + events = json.load(profile_file) + assert any( + event.get("args", {}).get("op_name") == "QMoE" and event["args"].get("provider") == "CPUExecutionProvider" + for event in events + ) + + if __name__ == "__main__": unittest.main()