From 6f888044747ac292e8165fbfc36e25f4eec236ea Mon Sep 17 00:00:00 2001 From: Yuri Khrustalev Date: Wed, 16 Sep 2026 09:44:38 -0400 Subject: [PATCH 01/11] perf: run block-wise int4/int8 QMoE CPU experts on the MLAS QNBit GEMM kernels The CPU QMoE kernel dequantized block-wise experts to fp32 and ran SGEMM, one expert per thread, so decode (top_k active experts) left most of the pool idle. Experts are now pre-packed once for the MatMulNBits kernels (no re-quantization) and each expert GEMM runs on them, sequentially with the whole pool when fewer experts are active than threads. ORT_QMOE_CPU_QNBIT_GEMM selects fp32 (default) or int8 activations, or disables the path. --- .../cpu/moe/moe_quantization_cpu.cc | 366 +++++++++++++++++- .../cpu/moe/moe_quantization_cpu.h | 21 + onnxruntime/test/contrib_ops/moe_test.cc | 239 ++++++++++++ 3 files changed, 610 insertions(+), 16 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc index fad6c5aecf512..17321c8935858 100644 --- a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc +++ b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc @@ -15,6 +15,7 @@ #include "core/framework/tensor_type_and_shape.h" #include "core/util/math.h" #include "core/platform/env_var_utils.h" +#include "core/util/thread_utils.h" #include "contrib_ops/cpu/moe/moe_utils.h" #include "contrib_ops/cpu/moe/moe_helper.h" @@ -121,6 +122,15 @@ namespace contrib { constexpr const char* kUseMlasQ4GemmMoe = "ORT_USE_MLAS_Q4_GEMM_MOE"; +// Controls the MLAS QNBit GEMM (MatMulNBits kernel) path for block-wise 4/8-bit experts without +// zero points. Unset or "fp32": fp32 activations where the platform has that kernel (4-bit); 8-bit +// has no fp32 variant and keeps the dequantize path. "int8": int8 activations for both (the +// MatMulNBits accuracy_level=4 configuration). "0": disabled. Any other value is an error. +constexpr const char* kQMoEQNBitGemmEnv = "ORT_QMOE_CPU_QNBIT_GEMM"; + +// Tag appended to the prepacked shape buffer to mark the QNBit-packed layout (see PrePackQNBitExperts). +constexpr int64_t kQNBitPackedLayoutTag = 0x514E4249; // 'QNBI' + template void DequantizeBlockWithMlas(const uint8_t* quantized_data, const TScale* scales, @@ -610,6 +620,13 @@ Status QMoECPU::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr all return Status::OK(); } + if ((input_idx == 2 || input_idx == 5) && tensor.Shape().NumDimensions() == 3) { + ORT_RETURN_IF_ERROR(PrePackQNBitExperts(tensor, input_idx, alloc, is_packed, prepacked_weights)); + if (is_packed) { + return Status::OK(); + } + } + // Only support PrePack for FC1 (2) and FC2 (5) weights. // 4-bit uses the existing unpacked-transposed/Q4 cache path. // 2-bit block-wise uses a direct LUT GEMM packed cache when supported. @@ -810,6 +827,35 @@ Status QMoECPU::UseSharedPrePackedBuffers(std::vector& prepa return Status::OK(); }; + // QNBit-packed layout: [packed experts, shape buffer ending with kQNBitPackedLayoutTag]. + if ((input_idx == 2 || input_idx == 5) && prepacked_buffers.size() == 2 && prepacked_buffer_sizes.size() == 2 && + prepacked_buffer_sizes[1] >= 5 * sizeof(int64_t)) { + const int64_t* buffer_data = static_cast(prepacked_buffers[1].get()); + if (buffer_data[0] == 3 && buffer_data[4] == kQNBitPackedLayoutTag) { + ORT_RETURN_IF_NOT(use_qnbit_gemm_, "QMoE prepacked weights use the QNBit layout but this kernel has it disabled."); + TensorShape& shape = (input_idx == 2) ? fc1_shape_ : fc2_shape_; + ORT_RETURN_IF_ERROR(parse_shape(shape)); + const int64_t num_experts = shape[0]; + const int64_t rows = shape[1]; + const int64_t cols = shape[2] * (8 / expert_weight_bits_); + const Tensor* scales_tensor = nullptr; + ORT_RETURN_IF_NOT(QNBitGemmEligible(input_idx, num_experts, rows, cols, &scales_tensor), + "QMoE prepacked weights use the QNBit layout but the node is not eligible for it."); + QNBitPackedExperts& packed = (input_idx == 2) ? qnbit_fc1_ : qnbit_fc2_; + // Already initialized when this kernel packed the weights itself (the session-state prepack + // pass always runs PrePack before handing the buffers back here). + if (packed.packed_size_per_expert == 0) { + ORT_RETURN_IF_ERROR(InitQNBitPacked(packed, num_experts, rows, cols, *scales_tensor, + Info().GetAllocator(OrtMemType::OrtMemTypeDefault))); + } + ORT_RETURN_IF_NOT(prepacked_buffer_sizes[0] == packed.packed_size_per_expert * static_cast(num_experts), + "QMoE prepacked QNBit buffer size does not match the expert shape."); + packed.packed = std::move(prepacked_buffers[0]); + used_shared_buffers = true; + return Status::OK(); + } + } + if (expert_weight_bits_ == 2) { if ((input_idx == 2 || input_idx == 5) && !prepacked_buffers.empty()) { if (input_idx == 2) { @@ -889,6 +935,201 @@ QMoECPU::QMoECPU(const OpKernelInfo& op_kernel_info) use_mlas_q4_gemm_ = true; use_mlas_q4_gemm_overridden_ = false; } + + // QNBit GEMM path policy. The MatMulNBits kernels consume the block-wise QMoE encoding directly + // (no re-quantization), have NEON/AVX2/AVX512 backends, and thread a single GEMM over N, so they + // are preferred for block-wise 4/8-bit experts without zero points. Int8 activations are lossier + // than the dequantize+SGEMM path they replace and QMoE has no accuracy_level attribute to opt in + // per model, so like MatMulNBits without accuracy_level=4 they are opt-in only: the default is + // fp32 activations where MLAS has that kernel (4-bit) and the existing path otherwise (8-bit). + bool allow_int8_compute = false; + bool allow_fp32_compute = true; + const auto qnbit_gemm_env = ParseEnvironmentVariable(kQMoEQNBitGemmEnv); + if (qnbit_gemm_env.has_value()) { + if (*qnbit_gemm_env == "0") { + use_qnbit_gemm_ = false; + } else if (*qnbit_gemm_env == "fp32") { + // Same as unset; accepted so the policy can be stated explicitly. + } else if (*qnbit_gemm_env == "int8") { + allow_int8_compute = true; + allow_fp32_compute = false; + } else { + ORT_THROW("Unsupported value for ", kQMoEQNBitGemmEnv, ": \"", *qnbit_gemm_env, + "\" (expected \"0\", \"fp32\" or \"int8\")."); + } + } + if (use_qnbit_gemm_ && (expert_weight_bits_ == 4 || expert_weight_bits_ == 8) && block_size_ > 0) { + const size_t nbits = static_cast(expert_weight_bits_); + const size_t blk = static_cast(block_size_); + if (allow_fp32_compute && MlasIsQNBitGemmAvailable(nbits, blk, SQNBIT_CompFp32)) { + qnbit_compute_type_ = SQNBIT_CompFp32; + } else if (allow_int8_compute && MlasIsQNBitGemmAvailable(nbits, blk, SQNBIT_CompInt8)) { + qnbit_compute_type_ = SQNBIT_CompInt8; + } else { + use_qnbit_gemm_ = false; + } + } else { + use_qnbit_gemm_ = false; + } +} + +template +bool QMoECPU::QNBitGemmEligible(int input_idx, int64_t num_experts, int64_t rows, int64_t cols, + const Tensor** scales_out) const { + *scales_out = nullptr; + if (!use_qnbit_gemm_ || (input_idx != 2 && input_idx != 5)) { + return false; + } + if (block_size_ <= 0 || (cols % block_size_) != 0) { + return false; + } + + const int scales_idx = (input_idx == 2) ? 3 : 6; + const int zp_idx = (input_idx == 2) ? 11 : 12; + const auto& input_defs = Info().node().InputDefs(); + const bool has_zp_input = zp_idx < static_cast(input_defs.size()) && input_defs[zp_idx]->Exists(); + if (has_zp_input) { + return false; + } + + const Tensor* scales_tensor = nullptr; + if (!Info().TryGetConstantInput(scales_idx, &scales_tensor) || scales_tensor == nullptr) { + return false; + } + const auto& scales_dims = scales_tensor->Shape().GetDims(); + const int64_t blocks_per_row = cols / block_size_; + if (scales_dims.size() != 3 || scales_dims[0] != num_experts || scales_dims[1] != rows || + scales_dims[2] != blocks_per_row) { + return false; // row-wise (per-channel) scales keep the existing paths + } + + const size_t packed_size = MlasQNBitGemmPackQuantBDataSize( + static_cast(rows), static_cast(cols), static_cast(expert_weight_bits_), + static_cast(block_size_), /*HasZeroPoint*/ false, qnbit_compute_type_, + &mlas_backend_kernel_selector_config_); + if (packed_size == 0) { + return false; + } + + *scales_out = scales_tensor; + return true; +} + +template +Status QMoECPU::InitQNBitPacked(QNBitPackedExperts& packed, int64_t num_experts, int64_t rows, int64_t cols, + const Tensor& scales, AllocatorPtr alloc) { + packed.packed_size_per_expert = MlasQNBitGemmPackQuantBDataSize( + static_cast(rows), static_cast(cols), static_cast(expert_weight_bits_), + static_cast(block_size_), /*HasZeroPoint*/ false, qnbit_compute_type_, + &mlas_backend_kernel_selector_config_); + ORT_RETURN_IF(packed.packed_size_per_expert == 0, "MLAS QNBit GEMM packing is not available for this expert shape."); + packed.scales_packed = MlasQNBitGemmScalesPacked( + static_cast(cols), static_cast(expert_weight_bits_), static_cast(block_size_), + qnbit_compute_type_, /*HasZeroPoint*/ false, &mlas_backend_kernel_selector_config_); + + if constexpr (std::is_same_v) { + // The MLAS kernels take fp32 scales; convert the constant fp16 scales once. + const size_t scales_count = static_cast(num_experts * rows * (cols / block_size_)); + packed.scales_fp32 = IAllocator::MakeUniquePtr(alloc, scales_count, true); + MlasConvertHalfToFloatBuffer(scales.Data(), packed.scales_fp32.get(), scales_count); + } else { + ORT_UNUSED_PARAMETER(num_experts); + ORT_UNUSED_PARAMETER(scales); + } + return Status::OK(); +} + +template +Status QMoECPU::PrePackQNBitExperts(const Tensor& tensor, int input_idx, AllocatorPtr alloc, + /*out*/ bool& is_packed, /*out*/ PrePackedWeights* prepacked_weights) { + const auto& shape = tensor.Shape(); + const int64_t num_experts = shape[0]; + const int64_t rows = shape[1]; + const int64_t pack_unit = 8 / expert_weight_bits_; + const int64_t cols = shape[2] * pack_unit; + + const Tensor* scales_tensor = nullptr; + if (!QNBitGemmEligible(input_idx, num_experts, rows, cols, &scales_tensor)) { + return Status::OK(); + } + + QNBitPackedExperts& packed = (input_idx == 2) ? qnbit_fc1_ : qnbit_fc2_; + ORT_RETURN_IF_ERROR(InitQNBitPacked(packed, num_experts, rows, cols, *scales_tensor, alloc)); + + const float* scales_fp32 = nullptr; + if constexpr (std::is_same_v) { + scales_fp32 = packed.scales_fp32.get(); + } else { + scales_fp32 = scales_tensor->Data(); + } + + const size_t nbits = static_cast(expert_weight_bits_); + const size_t blk = static_cast(block_size_); + const size_t per_expert = packed.packed_size_per_expert; + const size_t total_packed_size = per_expert * static_cast(num_experts); + auto packed_buffer = IAllocator::MakeUniquePtr(alloc, total_packed_size, true); + // The pack routines need not write every byte; zero so the prepacked-weights hash is deterministic. + std::memset(packed_buffer.get(), 0, total_packed_size); + + const size_t qdata_stride = static_cast(rows * shape[2]); + const size_t scale_stride = static_cast(rows * (cols / block_size_)); + const std::byte* qdata = static_cast(tensor.DataRaw()); + std::byte* dst = static_cast(packed_buffer.get()); + + // Same two-step sequence as MatMulNBits::PrePack: pack the quantized data, then (where the + // platform folds them into B) finalize the scales / block sums with a nullptr QuantBData call. + bool finalize_scales = (qnbit_compute_type_ == SQNBIT_CompInt8); +#if !defined(MLAS_TARGET_AMD64_IX86) + finalize_scales = finalize_scales && (nbits == 8 || packed.scales_packed); +#endif + + // Experts pack independently into disjoint regions, so spread them over a load-time pool + // (the session pool is not reachable from PrePack; same approach as MatMulNBits::PrePack). + OrtThreadPoolParams pack_tp_params; + pack_tp_params.thread_pool_size = Env::Default().GetNumPhysicalCpuCores(); + pack_tp_params.allow_spinning = false; + pack_tp_params.auto_set_affinity = false; + auto pack_tp = concurrency::CreateThreadPool(&Env::Default(), pack_tp_params, concurrency::ThreadPoolType::INTRA_OP); + concurrency::ThreadPool::TrySimpleParallelFor(pack_tp.get(), narrow(num_experts), [&](std::ptrdiff_t e) { + std::byte* expert_dst = dst + static_cast(e) * per_expert; + const float* expert_scales = scales_fp32 + static_cast(e) * scale_stride; + MlasQNBitGemmPackQuantBData(static_cast(rows), static_cast(cols), nbits, blk, qnbit_compute_type_, + qdata + static_cast(e) * qdata_stride, expert_dst, expert_scales, + /*HasZeroPoint*/ false, nullptr, nullptr, &mlas_backend_kernel_selector_config_); + if (finalize_scales) { + MlasQNBitGemmPackQuantBData(static_cast(rows), static_cast(cols), nbits, blk, qnbit_compute_type_, + nullptr, expert_dst, expert_scales, + /*HasZeroPoint*/ false, nullptr, nullptr, &mlas_backend_kernel_selector_config_); + } + }); + + if (input_idx == 2) { + fc1_shape_ = shape; + } else { + fc2_shape_ = shape; + } + + if (prepacked_weights != nullptr) { + prepacked_weights->buffers_.push_back(std::move(packed_buffer)); + prepacked_weights->buffer_sizes_.push_back(total_packed_size); + + // Shape buffer: rank, dims, then the layout tag so UseSharedPrePackedBuffers can tell this + // layout apart from the legacy unpacked-transposed one (which stores rank + dims only). + auto dims = shape.GetDims(); + const size_t shape_size = (2 + dims.size()) * sizeof(int64_t); + auto shape_buffer = IAllocator::MakeUniquePtr(alloc, shape_size); + int64_t* buffer_data = static_cast(shape_buffer.get()); + buffer_data[0] = static_cast(dims.size()); + memcpy(buffer_data + 1, dims.data(), dims.size() * sizeof(int64_t)); + buffer_data[1 + dims.size()] = kQNBitPackedLayoutTag; + prepacked_weights->buffers_.push_back(std::move(shape_buffer)); + prepacked_weights->buffer_sizes_.push_back(shape_size); + } else { + packed.packed = std::move(packed_buffer); + } + + is_packed = true; + return Status::OK(); } template @@ -896,10 +1137,10 @@ Status QMoECPU::Compute(OpKernelContext* context) const { const ComputeInputs inputs{ context->Input(0), context->Input(1), - ((packed_fc1_ != nullptr) || (packed_fc1_lut_cache_ != nullptr)) ? nullptr : context->Input(2), + ((packed_fc1_ != nullptr) || (packed_fc1_lut_cache_ != nullptr) || (qnbit_fc1_.packed != nullptr)) ? nullptr : context->Input(2), context->Input(3), context->Input(4), - ((packed_fc2_ != nullptr) || (packed_fc2_lut_cache_ != nullptr)) ? nullptr : context->Input(5), + ((packed_fc2_ != nullptr) || (packed_fc2_lut_cache_ != nullptr) || (qnbit_fc2_.packed != nullptr)) ? nullptr : context->Input(5), context->Input(6), context->Input(7), context->Input(8), @@ -911,8 +1152,8 @@ Status QMoECPU::Compute(OpKernelContext* context) const { context->Input(14), }; - const bool has_fc1_prepacked = (packed_fc1_ != nullptr) || (packed_fc1_lut_cache_ != nullptr); - const bool has_fc2_prepacked = (packed_fc2_ != nullptr) || (packed_fc2_lut_cache_ != nullptr); + const bool has_fc1_prepacked = (packed_fc1_ != nullptr) || (packed_fc1_lut_cache_ != nullptr) || (qnbit_fc1_.packed != nullptr); + const bool has_fc2_prepacked = (packed_fc2_ != nullptr) || (packed_fc2_lut_cache_ != nullptr) || (qnbit_fc2_.packed != nullptr); const TensorShape* fc1_shape_ptr = has_fc1_prepacked ? &fc1_shape_ : (inputs.fc1_experts_weights ? &inputs.fc1_experts_weights->Shape() : nullptr); const TensorShape* fc2_shape_ptr = has_fc2_prepacked ? &fc2_shape_ : (inputs.fc2_experts_weights ? &inputs.fc2_experts_weights->Shape() : nullptr); @@ -1134,16 +1375,22 @@ Status QMoECPU::ComputeCommon(OpKernelContext* context, const ComputeInputs& if (!tokens.empty()) ++num_active_experts; } - // Parallelize the expert loop across the active experts, one expert (batch) per thread. This is - // the key lever for decode (batch=seq=1): only top_k experts are active, and each per-expert GEMM - // is effectively a GEMV (M==1) that MLAS does not thread internally, so giving the whole pool to - // the inner op would leave most cores idle. Spreading active experts across threads instead keeps - // the cores busy. Nested parallelism (outer expert loop + inner GEMM/dequant on the same pool) can - // livelock ORT's Eigen pool (see PR #29081), so only one level may use the session pool tp: when - // the expert loop runs multi-threaded the inner ops run serially (inner_tp == nullptr); when a - // single expert is active the loop is serial and the inner GEMM gets the full pool (inner_tp == tp). + // Nested parallelism (outer expert loop + inner GEMM/dequant on the same pool) can livelock ORT's + // Eigen pool (see PR #29081), so exactly one level uses the session pool tp: either the expert + // loop is multi-threaded and the inner ops run serially (inner_tp == nullptr), or the expert loop + // is serial and the inner ops get the full pool (inner_tp == tp). Which level to thread depends + // on the GEMM path: + // - dequantize+SGEMM / LUT paths: one active expert per thread. For decode (M==1) the per-expert + // GEMM is a GEMV that MLAS does not thread internally, so this is the only way to keep the + // cores busy. + // - QNBit path: the MatMulNBits kernels thread a single GEMM over N (and M), so when fewer experts + // are active than there are threads (decode: top_k experts) the experts run sequentially and + // every GEMM gets the whole pool instead of leaving most of it idle. // This must be decided BEFORE the per-thread workspaces below, which are sized by num_expert_threads. int num_expert_threads = std::max(1, std::min(num_active_experts, max_expert_threads)); + if (qnbit_fc1_.packed != nullptr && qnbit_fc2_.packed != nullptr && num_active_experts < max_expert_threads) { + num_expert_threads = 1; + } concurrency::ThreadPool* inner_tp = (num_expert_threads > 1) ? nullptr : tp; auto thread_local_outputs_ptr = IAllocator::MakeUniquePtr(allocator, static_cast(num_expert_threads) * output_buffer_size); @@ -1163,8 +1410,8 @@ Status QMoECPU::ComputeCommon(OpKernelContext* context, const ComputeInputs& const size_t C1_size = align_size(static_cast(max_tokens_per_expert) * static_cast(fc1_out_features)); const size_t A2_size = align_size(static_cast(max_tokens_per_expert) * static_cast(inter_size)); const size_t C2_size = align_size(static_cast(max_tokens_per_expert) * static_cast(hidden_size)); - const size_t B1_dequant_size = align_size(static_cast(fc1_out_features) * static_cast(hidden_size)); - const size_t B2_dequant_size = align_size(static_cast(hidden_size) * static_cast(inter_size)); + const size_t B1_dequant_size = (qnbit_fc1_.packed != nullptr) ? 0 : align_size(static_cast(fc1_out_features) * static_cast(hidden_size)); + const size_t B2_dequant_size = (qnbit_fc2_.packed != nullptr) ? 0 : align_size(static_cast(hidden_size) * static_cast(inter_size)); const size_t workspace_elements_per_thread = A1_size + C1_size + A2_size + C2_size + B1_dequant_size + B2_dequant_size; @@ -1181,8 +1428,10 @@ Status QMoECPU::ComputeCommon(OpKernelContext* context, const ComputeInputs& const bool is_fc1_block_wise = (fc1_scales_dims.size() == 3 && fc1_scales_dims[2] > 1); const bool is_fc2_block_wise = (fc2_scales_dims.size() == 3 && fc2_scales_dims[2] > 1); - const uint8_t* fc1_weights_data = (packed_fc1_ != nullptr || packed_fc1_lut_cache_ != nullptr) ? nullptr : fc1_experts_weights->template Data(); - const uint8_t* fc2_weights_data = (packed_fc2_ != nullptr || packed_fc2_lut_cache_ != nullptr) ? nullptr : fc2_experts_weights->template Data(); + const bool use_qnbit_fc1 = (qnbit_fc1_.packed != nullptr); + const bool use_qnbit_fc2 = (qnbit_fc2_.packed != nullptr); + const uint8_t* fc1_weights_data = (packed_fc1_ != nullptr || packed_fc1_lut_cache_ != nullptr || use_qnbit_fc1) ? nullptr : fc1_experts_weights->template Data(); + const uint8_t* fc2_weights_data = (packed_fc2_ != nullptr || packed_fc2_lut_cache_ != nullptr || use_qnbit_fc2) ? nullptr : fc2_experts_weights->template Data(); const T* fc1_scales_data = fc1_scales->template Data(); const T* fc2_scales_data = fc2_scales->template Data(); const T* fc1_bias_data = fc1_experts_bias ? fc1_experts_bias->template Data() : nullptr; @@ -1297,6 +1546,57 @@ Status QMoECPU::ComputeCommon(OpKernelContext* context, const ComputeInputs& } } + // One expert GEMM on the MLAS QNBit (MatMulNBits) kernels: C[M, N] = A[M, K] * dequant(B_expert)^T + bias. + const size_t qnbit_bits = static_cast(expert_weight_bits_); + const size_t qnbit_blk = static_cast(std::max(block_size_, 1)); + // The CompInt8 kernels need a per-GEMM workspace (quantized A); size it for the largest expert + // batch once per thread rather than allocating per expert. + size_t qnbit_workspace_per_thread = 0; + if (use_qnbit_fc1) { + qnbit_workspace_per_thread = std::max(qnbit_workspace_per_thread, MlasQNBitGemmBatchWorkspaceSize( + static_cast(max_tokens_per_expert), static_cast(fc1_out_features), static_cast(hidden_size), + 1, qnbit_bits, qnbit_blk, /*HasZeroPoint*/ false, qnbit_compute_type_, &mlas_backend_kernel_selector_config_)); + } + if (use_qnbit_fc2) { + qnbit_workspace_per_thread = std::max(qnbit_workspace_per_thread, MlasQNBitGemmBatchWorkspaceSize( + static_cast(max_tokens_per_expert), static_cast(hidden_size), static_cast(inter_size), + 1, qnbit_bits, qnbit_blk, /*HasZeroPoint*/ false, qnbit_compute_type_, &mlas_backend_kernel_selector_config_)); + } + IAllocatorUniquePtr qnbit_workspace_ptr; + std::byte* qnbit_workspace = nullptr; + if (qnbit_workspace_per_thread > 0) { + qnbit_workspace_ptr = IAllocator::MakeUniquePtr(allocator, static_cast(num_expert_threads) * qnbit_workspace_per_thread, true); + qnbit_workspace = qnbit_workspace_ptr.get(); + } + auto run_qnbit_gemm = [&](const QNBitPackedExperts& packed, int64_t expert_idx, + const float* A, size_t M, size_t N, size_t K, + const T* expert_scales, const float* bias, float* C, + std::byte* gemm_workspace, concurrency::ThreadPool* gemm_tp) { + const float* scales = nullptr; + if (!packed.scales_packed) { + if constexpr (std::is_same_v) { + scales = packed.scales_fp32.get() + static_cast(expert_idx) * N * (K / qnbit_blk); + } else { + scales = reinterpret_cast(expert_scales); + } + } + const std::byte* packed_b = static_cast(packed.packed.get()) + + static_cast(expert_idx) * packed.packed_size_per_expert; + + MLAS_QNBIT_GEMM_DATA_PARAMS params; + params.A = A; + params.lda = K; + params.QuantBDataWorkspace = packed_b; + params.PackedQuantBData = packed_b; + params.QuantBScale = scales; + params.QuantBZeroPoint = nullptr; + params.Bias = bias; + params.C = C; + params.ldc = N; + MlasQNBitGemmBatch(M, N, K, 1, qnbit_bits, qnbit_blk, qnbit_compute_type_, ¶ms, + gemm_workspace, gemm_tp, &mlas_backend_kernel_selector_config_); + }; + std::vector> expert_workload; for (int64_t i = 0; i < num_experts; ++i) { const size_t token_count = expert_token_map[static_cast(i)].size(); @@ -1329,6 +1629,9 @@ Status QMoECPU::ComputeCommon(OpKernelContext* context, const ComputeInputs& float* thread_bias1_buffer = bias_conversion_buffers + static_cast(thread_id) * (static_cast(fc1_out_features) + static_cast(hidden_size)); float* thread_bias2_buffer = thread_bias1_buffer + static_cast(fc1_out_features); + std::byte* thread_qnbit_workspace = (qnbit_workspace == nullptr) + ? nullptr + : (qnbit_workspace + static_cast(thread_id) * qnbit_workspace_per_thread); for (int64_t expert_idx : expert_batch) { bool fc2_bias_added_by_mlas = false; @@ -1414,6 +1717,21 @@ Status QMoECPU::ComputeCommon(OpKernelContext* context, const ComputeInputs& CanUseMlasQ4Gemm(expert_weight_bits_, is_fc1_block_wise ? block_size_ : 0, fc1_out_features, hidden_size, q_type))); + if (use_qnbit_fc1) { + float* fc1_bias_float = nullptr; + if (has_fc1_bias) { + const T* B1_bias = fc1_bias_data + expert_idx * fc1_out_features; + if constexpr (std::is_same_v) { + MlasConvertHalfToFloatBuffer(reinterpret_cast(B1_bias), thread_bias1_buffer, static_cast(fc1_out_features)); + } else { + std::memcpy(thread_bias1_buffer, B1_bias, static_cast(fc1_out_features) * sizeof(float)); + } + fc1_bias_float = thread_bias1_buffer; + } + run_qnbit_gemm(qnbit_fc1_, expert_idx, A1, m, n, k, fc1_scales_ptr, fc1_bias_float, C1, thread_qnbit_workspace, inner_tp); + goto fc1_gemm_done; + } + if (can_use_fc1_lut_gemm && TryRunLutGemm(A1, C1, fc1_weights_data, fc1_direct_lut_cache_ptr != nullptr @@ -1658,6 +1976,22 @@ Status QMoECPU::ComputeCommon(OpKernelContext* context, const ComputeInputs& CanUseMlasQ4Gemm(expert_weight_bits_, is_fc2_block_wise ? block_size_ : 0, hidden_size, inter_size, q_type2))); + if (use_qnbit_fc2) { + float* fc2_bias_float = nullptr; + if (has_fc2_bias) { + const T* B2_bias = fc2_bias_data + expert_idx * hidden_size; + if constexpr (std::is_same_v) { + MlasConvertHalfToFloatBuffer(reinterpret_cast(B2_bias), thread_bias2_buffer, static_cast(hidden_size)); + } else { + std::memcpy(thread_bias2_buffer, B2_bias, static_cast(hidden_size) * sizeof(float)); + } + fc2_bias_float = thread_bias2_buffer; + } + run_qnbit_gemm(qnbit_fc2_, expert_idx, A2, m2, n2, k2, fc2_scales_ptr, fc2_bias_float, C2, thread_qnbit_workspace, inner_tp); + fc2_bias_added_by_mlas = true; + goto fc2_gemm_done; + } + if (can_use_fc2_lut_gemm && TryRunLutGemm(A2, C2, fc2_weights_data, fc2_direct_lut_cache_ptr != nullptr diff --git a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h index 228cc2c7b012f..68e7c72a9cbe2 100644 --- a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h +++ b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h @@ -60,6 +60,22 @@ class QMoECPU final : public OpKernel, public MoEBaseCPU { void ApplyActivationVectorized(float* data, int64_t size) const; + // Expert weights pre-packed for the MLAS QNBit GEMM kernels (the MatMulNBits kernels). Used for + // block-wise 4/8-bit experts without zero points; each expert's [rows, cols] matrix is packed + // independently and stored back to back. + struct QNBitPackedExperts { + IAllocatorUniquePtr packed; // num_experts * packed_size_per_expert bytes + size_t packed_size_per_expert{0}; + IAllocatorUniquePtr scales_fp32; // fp32 copy of the [E, rows, cols/block_size] scales (T == MLFloat16 only) + bool scales_packed{false}; // scales are folded into `packed`; pass no QuantBScale at compute time + }; + bool QNBitGemmEligible(int input_idx, int64_t num_experts, int64_t rows, int64_t cols, + const Tensor** scales_out) const; + Status InitQNBitPacked(QNBitPackedExperts& packed, int64_t num_experts, int64_t rows, int64_t cols, + const Tensor& scales, AllocatorPtr alloc); + Status PrePackQNBitExperts(const Tensor& tensor, int input_idx, AllocatorPtr alloc, + /*out*/ bool& is_packed, /*out*/ PrePackedWeights* prepacked_weights); + int64_t expert_weight_bits_; int64_t block_size_; bool use_mlas_q4_gemm_{false}; @@ -75,6 +91,11 @@ class QMoECPU final : public OpKernel, public MoEBaseCPU { IAllocatorUniquePtr packed_fc1_mlas_cache_; IAllocatorUniquePtr packed_fc2_mlas_cache_; + + bool use_qnbit_gemm_{true}; + MLAS_QNBIT_GEMM_COMPUTE_TYPE qnbit_compute_type_{SQNBIT_CompFp32}; + QNBitPackedExperts qnbit_fc1_; + QNBitPackedExperts qnbit_fc2_; }; } // namespace contrib diff --git a/onnxruntime/test/contrib_ops/moe_test.cc b/onnxruntime/test/contrib_ops/moe_test.cc index 1cdd22f789816..b5cd18deaa299 100644 --- a/onnxruntime/test/contrib_ops/moe_test.cc +++ b/onnxruntime/test/contrib_ops/moe_test.cc @@ -1,11 +1,18 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include +#include +#include +#include +#include + #include "gtest/gtest.h" #include "core/session/onnxruntime_session_options_config_keys.h" #include "test/common/tensor_op_test_utils.h" #include "test/common/cuda_op_test_utils.h" #include "test/providers/provider_test_utils.h" +#include "test/util/include/scoped_env_vars.h" namespace onnxruntime { namespace test { @@ -2120,6 +2127,238 @@ TEST(MoETest, QMoETest_CPU_Int4_MLAS) { #endif } +// Block-wise 4/8-bit SwiGLU QMoE on CPU with non-trivial weights. Block-wise symmetric experts +// without zero points take the MLAS QNBit GEMM (MatMulNBits kernel) path; the expected output is an +// fp32 reference computed here from the dequantized weights and the kernel's routing (softmax over +// the top-k router logits). Weights, scales and biases are constant initializers so that PrePack runs; +// as graph inputs the kernel would silently fall back to the dequantize+SGEMM path. +namespace { +struct QMoEBlockWiseCase { + int num_rows; + int num_experts; + int hidden_size; + int inter_size; + int block_size; + int top_k; + bool with_bias; + int bits{4}; +}; + +template +void RunQMoECpuBlockWiseSwiGLU(const QMoEBlockWiseCase& c, float tolerance, + bool share_prepacked_weights_across_sessions = false) { + const int num_rows = c.num_rows, num_experts = c.num_experts, hidden_size = c.hidden_size; + const int inter_size = c.inter_size, block_size = c.block_size, top_k = c.top_k, bits = c.bits; + const int fc1_rows = 2 * inter_size; // interleaved [gate, up] rows + const int fc1_blocks = hidden_size / block_size; + const int fc2_blocks = inter_size / block_size; + const int pack = 8 / bits; + const float code_zero = (bits == 4) ? 8.0f : 128.0f; + constexpr bool is_fp16 = std::is_same_v; + + uint32_t state = 0x12345678u; + auto next_uniform = [&state]() { + state = state * 1664525u + 1013904223u; + return static_cast((state >> 8) & 0xFFFF) / 65535.0f; // [0, 1] + }; + auto next_code = [&state, bits]() { + state = state * 1664525u + 1013904223u; + return static_cast((state >> 12) & ((1u << bits) - 1)); + }; + // The kernel converts T inputs to fp32, so the reference must see the T-rounded values. + auto round_to_t = [](float v) { return is_fp16 ? MLFloat16(v).ToFloat() : v; }; + + std::vector input(static_cast(num_rows * hidden_size)); + for (auto& v : input) v = round_to_t(2.0f * next_uniform() - 1.0f); + std::vector router_probs(static_cast(num_rows * num_experts)); + for (auto& v : router_probs) v = round_to_t(4.0f * next_uniform() - 2.0f); + + // Weights are stored as [E, N, K/pack] uint8 with the lower k in the low bits; scales as [E, N, K/block]. + std::vector fc1_codes(static_cast(num_experts * fc1_rows * hidden_size)); + std::vector fc2_codes(static_cast(num_experts * hidden_size * inter_size)); + for (auto& v : fc1_codes) v = next_code(); + for (auto& v : fc2_codes) v = next_code(); + auto pack_codes = [pack, bits](const std::vector& codes) { + std::vector packed(codes.size() / pack); + for (size_t i = 0; i < packed.size(); ++i) { + uint8_t byte = 0; + for (int j = 0; j < pack; ++j) byte |= static_cast(codes[i * pack + j] << (j * bits)); + packed[i] = byte; + } + return packed; + }; + const std::vector fc1_weights = pack_codes(fc1_codes); + const std::vector fc2_weights = pack_codes(fc2_codes); + std::vector fc1_scales(static_cast(num_experts * fc1_rows * fc1_blocks)); + std::vector fc2_scales(static_cast(num_experts * hidden_size * fc2_blocks)); + // Signed scales, as produced by MatMulNBits-style quantizers; scaled so the weight range is + // comparable for 4 and 8 bit codes. + const float scale_unit = 8.0f / code_zero; + for (auto& v : fc1_scales) v = round_to_t((next_uniform() < 0.5f ? -1.0f : 1.0f) * scale_unit * (0.01f + 0.03f * next_uniform())); + for (auto& v : fc2_scales) v = round_to_t((next_uniform() < 0.5f ? -1.0f : 1.0f) * scale_unit * (0.01f + 0.03f * next_uniform())); + std::vector fc1_bias, fc2_bias; + if (c.with_bias) { + fc1_bias.resize(static_cast(num_experts * fc1_rows)); + fc2_bias.resize(static_cast(num_experts * hidden_size)); + for (auto& v : fc1_bias) v = round_to_t(0.2f * next_uniform() - 0.1f); + for (auto& v : fc2_bias) v = round_to_t(0.2f * next_uniform() - 0.1f); + } + + auto dequant = [&](const std::vector& codes, const std::vector& scales, int e, int n, int k, + int rows, int cols, int blocks) { + const float scale = scales[static_cast((e * rows + n) * blocks + k / block_size)]; + return (static_cast(codes[static_cast((e * rows + n) * cols + k)]) - code_zero) * scale; + }; + + std::vector expected(static_cast(num_rows * hidden_size), 0.0f); + for (int t = 0; t < num_rows; ++t) { + const float* logits = router_probs.data() + t * num_experts; + std::vector> sorted; + for (int e = 0; e < num_experts; ++e) sorted.emplace_back(logits[e], e); + std::sort(sorted.begin(), sorted.end(), [](const auto& a, const auto& b) { return a.first > b.first; }); + float sum_exp = 0.0f; + std::vector weights(static_cast(top_k)); + for (int j = 0; j < top_k; ++j) { + weights[j] = std::exp(sorted[j].first - sorted[0].first); + sum_exp += weights[j]; + } + for (int j = 0; j < top_k; ++j) { + const int e = sorted[j].second; + const float w = weights[j] / sum_exp; + const float* x = input.data() + t * hidden_size; + std::vector h(static_cast(fc1_rows)); + for (int n = 0; n < fc1_rows; ++n) { + float acc = c.with_bias ? fc1_bias[static_cast(e * fc1_rows + n)] : 0.0f; + for (int k = 0; k < hidden_size; ++k) acc += x[k] * dequant(fc1_codes, fc1_scales, e, n, k, fc1_rows, hidden_size, fc1_blocks); + h[n] = acc; + } + std::vector a(static_cast(inter_size)); + for (int i = 0; i < inter_size; ++i) { + const float gate = h[2 * i], up = h[2 * i + 1]; + a[i] = gate / (1.0f + std::exp(-gate)) * up; + } + for (int n = 0; n < hidden_size; ++n) { + float acc = c.with_bias ? fc2_bias[static_cast(e * hidden_size + n)] : 0.0f; + for (int k = 0; k < inter_size; ++k) acc += a[k] * dequant(fc2_codes, fc2_scales, e, n, k, hidden_size, inter_size, fc2_blocks); + expected[static_cast(t * hidden_size + n)] += w * acc; + } + } + } + + OpTester tester("QMoE", 1, onnxruntime::kMSDomain); + tester.AddAttribute("k", top_k); + tester.AddAttribute("activation_type", "swiglu"); + tester.AddAttribute("swiglu_fusion", 1); + tester.AddAttribute("activation_alpha", 1.0f); + tester.AddAttribute("activation_beta", 0.0f); + tester.AddAttribute("normalize_routing_weights", 1); + tester.AddAttribute("expert_weight_bits", bits); + tester.AddAttribute("block_size", block_size); + + auto add_t_input = [&](const char* name, std::vector dims, const std::vector& values, bool is_initializer) { + if constexpr (is_fp16) { + tester.AddInput(name, dims, ToFloat16(values), is_initializer); + } else { + tester.AddInput(name, dims, values, is_initializer); + } + }; + add_t_input("input", {num_rows, hidden_size}, input, false); + add_t_input("router_probs", {num_rows, num_experts}, router_probs, false); + tester.AddInput("fc1_experts_weights", {num_experts, fc1_rows, hidden_size / pack}, fc1_weights, true); + add_t_input("fc1_scales", {num_experts, fc1_rows, fc1_blocks}, fc1_scales, true); + if (c.with_bias) { + add_t_input("fc1_experts_bias", {num_experts, fc1_rows}, fc1_bias, true); + } else { + tester.AddOptionalInputEdge(); + } + tester.AddInput("fc2_experts_weights", {num_experts, hidden_size, inter_size / pack}, fc2_weights, true); + add_t_input("fc2_scales", {num_experts, hidden_size, fc2_blocks}, fc2_scales, true); + if (c.with_bias) { + add_t_input("fc2_experts_bias", {num_experts, hidden_size}, fc2_bias, true); + } else { + tester.AddOptionalInputEdge(); + } + if constexpr (is_fp16) { + tester.AddOutput("output", {num_rows, hidden_size}, ToFloat16(expected)); + } else { + tester.AddOutput("output", {num_rows, hidden_size}, expected); + } + tester.SetOutputTolerance(tolerance); + + auto cpu_ep = []() { + std::vector> eps; + eps.push_back(DefaultCpuExecutionProvider()); + return eps; + }; + + if (!share_prepacked_weights_across_sessions) { + auto eps = cpu_ep(); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); + return; + } + + // Two sessions sharing the fc1/fc2 initializers: the second must adopt the first session's packed + // buffers through UseSharedPrePackedBuffers (the tagged QNBit layout) and still match the reference. + OrtValue fc1_value, fc2_value; + Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape({num_experts, fc1_rows, hidden_size / pack}), + const_cast(fc1_weights.data()), OrtMemoryInfo(CPU, OrtAllocatorType::OrtDeviceAllocator), fc1_value); + Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape({num_experts, hidden_size, inter_size / pack}), + const_cast(fc2_weights.data()), OrtMemoryInfo(CPU, OrtAllocatorType::OrtDeviceAllocator), fc2_value); + SessionOptions so; + ASSERT_STATUS_OK(so.AddInitializer("fc1_experts_weights", &fc1_value)); + ASSERT_STATUS_OK(so.AddInitializer("fc2_experts_weights", &fc2_value)); + tester.EnableSharingOfPrePackedWeightsAcrossSessions(); + + size_t prepacked_session_1 = 0, prepacked_session_2 = 0, shared = 0; + { + auto eps = cpu_ep(); + tester.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps, {}, &prepacked_session_1, &shared); + ASSERT_EQ(shared, static_cast(0)); + } + ASSERT_EQ(prepacked_session_1, tester.GetNumPrePackedWeightsShared()); + if (prepacked_session_1 == 0) { + return; // platform without the QNBit kernels: nothing was prepacked, nothing to share + } + { + auto eps = cpu_ep(); + tester.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps, {}, &prepacked_session_2, &shared); + ASSERT_EQ(prepacked_session_1, prepacked_session_2); + ASSERT_EQ(shared, prepacked_session_2); + } +} +} // namespace + +TEST(MoETest, QMoETest_CPU_Int4_BlockWise_SwiGLU_Decode) { + // Single token: the QNBit path runs the selected experts sequentially, each GEMM threaded internally. + RunQMoECpuBlockWiseSwiGLU({1, 8, 64, 64, 32, 2, false}, 0.01f); +} + +TEST(MoETest, QMoETest_CPU_Int4_BlockWise_SwiGLU_Prefill) { + RunQMoECpuBlockWiseSwiGLU({37, 4, 128, 96, 32, 2, false}, 0.01f); +} + +TEST(MoETest, QMoETest_CPU_Int4_BlockWise_SwiGLU_Bias_Block64) { + RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 64, 3, true}, 0.01f); +} + +TEST(MoETest, QMoETest_CPU_Int4_BlockWise_SwiGLU_Fp16) { + RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true}, 0.02f); +} + +TEST(MoETest, QMoETest_CPU_Int8_BlockWise_SwiGLU) { + // 8-bit has no fp32-activation QNBit kernel, so this runs the int8-activation kernels where available. + RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true, 8}, 0.05f); +} + +TEST(MoETest, QMoETest_CPU_Int4_BlockWise_SwiGLU_Int8Activations) { + ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_QMOE_CPU_QNBIT_GEMM", "int8"}}}; + RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true}, 0.05f); +} + +TEST(MoETest, QMoETest_CPU_Int4_BlockWise_SwiGLU_SharedPrePackedWeights) { + RunQMoECpuBlockWiseSwiGLU({3, 4, 64, 64, 32, 2, true}, 0.01f, /*share_prepacked_weights_across_sessions*/ true); +} + TEST(MoETest, QMoETest_CPU_Int8_MLAS) { #ifdef USE_MLAS // Skip this test if we're not testing CPU execution provider From 33b3684d8367144b7868a07ff58f08087fb3559c Mon Sep 17 00:00:00 2001 From: Yuri Khrustalev Date: Wed, 16 Sep 2026 11:21:08 -0400 Subject: [PATCH 02/11] refactor: tighten the QMoE CPU QNBit prepack and GEMM plumbing Make int8 activations opt-in (ORT_QMOE_CPU_QNBIT_GEMM=int8), matching MatMulNBits without accuracy_level=4; 8-bit experts keep the dequantize path by default. Record the compute type in the tagged prepacked shape buffer and reject shared buffers packed for another compute type. Share one shape-buffer writer across the three prepack layouts, return the packed size from the eligibility check instead of recomputing it, pass fp32 scales to the GEMM lambda by expert index, and only spin up the load-time pack pool when there is more than one expert. Split the 8-bit test into default (dequantize) and int8-activation cases. Co-Authored-By: Claude Fable 5.1 --- .../cpu/moe/moe_quantization_cpu.cc | 142 +++++++++--------- .../cpu/moe/moe_quantization_cpu.h | 6 +- onnxruntime/test/contrib_ops/moe_test.cc | 7 +- 3 files changed, 80 insertions(+), 75 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc index 17321c8935858..2afab3dc1fc81 100644 --- a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc +++ b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc @@ -131,6 +131,21 @@ constexpr const char* kQMoEQNBitGemmEnv = "ORT_QMOE_CPU_QNBIT_GEMM"; // Tag appended to the prepacked shape buffer to mark the QNBit-packed layout (see PrePackQNBitExperts). constexpr int64_t kQNBitPackedLayoutTag = 0x514E4249; // 'QNBI' +// Prepacked shape buffer: rank, dims, then `trailer` (empty for the legacy layouts; the QNBit layout +// appends kQNBitPackedLayoutTag and the compute type it was packed for). +static void PushPrePackedShapeBuffer(const TensorShape& shape, gsl::span trailer, + const AllocatorPtr& alloc, PrePackedWeights& prepacked_weights) { + const auto dims = shape.GetDims(); + const size_t shape_size = (1 + dims.size() + trailer.size()) * sizeof(int64_t); + auto shape_buffer = IAllocator::MakeUniquePtr(alloc, shape_size); + int64_t* buffer_data = static_cast(shape_buffer.get()); + buffer_data[0] = static_cast(dims.size()); + std::copy(dims.begin(), dims.end(), buffer_data + 1); + std::copy(trailer.begin(), trailer.end(), buffer_data + 1 + dims.size()); + prepacked_weights.buffers_.push_back(std::move(shape_buffer)); + prepacked_weights.buffer_sizes_.push_back(shape_size); +} + template void DequantizeBlockWithMlas(const uint8_t* quantized_data, const TScale* scales, @@ -699,19 +714,7 @@ Status QMoECPU::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr all prepacked_weights->buffers_.push_back(std::move(lut_cache_buffer)); prepacked_weights->buffer_sizes_.push_back(cache_size); is_packed = true; - - auto dims = shape.GetDims(); - size_t rank_bytes = sizeof(int64_t); - size_t dims_bytes = dims.size() * sizeof(int64_t); - size_t shape_size = rank_bytes + dims_bytes; - - auto shape_buffer = IAllocator::MakeUniquePtr(alloc, shape_size); - int64_t* buffer_data = static_cast(shape_buffer.get()); - *buffer_data = static_cast(dims.size()); - memcpy(buffer_data + 1, dims.data(), dims_bytes); - - prepacked_weights->buffers_.push_back(std::move(shape_buffer)); - prepacked_weights->buffer_sizes_.push_back(shape_size); + PushPrePackedShapeBuffer(shape, {}, alloc, *prepacked_weights); return Status::OK(); } @@ -738,20 +741,7 @@ Status QMoECPU::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr all prepacked_weights->buffers_.push_back(std::move(packed_buffer)); prepacked_weights->buffer_sizes_.push_back(packed_size); is_packed = true; - - // Pack Shape (Buffer 1) - auto dims = shape.GetDims(); - size_t rank_bytes = sizeof(int64_t); - size_t dims_bytes = dims.size() * sizeof(int64_t); - size_t shape_size = rank_bytes + dims_bytes; - - auto shape_buffer = IAllocator::MakeUniquePtr(alloc, shape_size); - int64_t* buffer_data = static_cast(shape_buffer.get()); - *buffer_data = static_cast(dims.size()); - memcpy(buffer_data + 1, dims.data(), dims_bytes); - - prepacked_weights->buffers_.push_back(std::move(shape_buffer)); - prepacked_weights->buffer_sizes_.push_back(shape_size); + PushPrePackedShapeBuffer(shape, {}, alloc, *prepacked_weights); // Try build MLAS Q4 cache if scales are available if (use_mlas_q4_gemm_) { @@ -827,25 +817,33 @@ Status QMoECPU::UseSharedPrePackedBuffers(std::vector& prepa return Status::OK(); }; - // QNBit-packed layout: [packed experts, shape buffer ending with kQNBitPackedLayoutTag]. + // QNBit-packed layout: [packed experts, shape buffer = rank, dims, kQNBitPackedLayoutTag, compute type]. if ((input_idx == 2 || input_idx == 5) && prepacked_buffers.size() == 2 && prepacked_buffer_sizes.size() == 2 && - prepacked_buffer_sizes[1] >= 5 * sizeof(int64_t)) { + prepacked_buffer_sizes[1] >= 6 * sizeof(int64_t)) { const int64_t* buffer_data = static_cast(prepacked_buffers[1].get()); if (buffer_data[0] == 3 && buffer_data[4] == kQNBitPackedLayoutTag) { ORT_RETURN_IF_NOT(use_qnbit_gemm_, "QMoE prepacked weights use the QNBit layout but this kernel has it disabled."); + // The packed bytes depend on the compute type (block sums, folded scales), and equal sizes do + // not imply equal layouts, so a buffer packed under another ORT_QMOE_CPU_QNBIT_GEMM setting or + // saved by another build must be rejected rather than silently used. + ORT_RETURN_IF_NOT(buffer_data[5] == static_cast(qnbit_compute_type_), + "QMoE prepacked weights were packed for QNBit compute type ", buffer_data[5], + " but this kernel uses ", static_cast(qnbit_compute_type_), "."); TensorShape& shape = (input_idx == 2) ? fc1_shape_ : fc2_shape_; ORT_RETURN_IF_ERROR(parse_shape(shape)); const int64_t num_experts = shape[0]; const int64_t rows = shape[1]; const int64_t cols = shape[2] * (8 / expert_weight_bits_); const Tensor* scales_tensor = nullptr; - ORT_RETURN_IF_NOT(QNBitGemmEligible(input_idx, num_experts, rows, cols, &scales_tensor), + size_t packed_size_per_expert = 0; + ORT_RETURN_IF_NOT(QNBitGemmEligible(input_idx, num_experts, rows, cols, &scales_tensor, &packed_size_per_expert), "QMoE prepacked weights use the QNBit layout but the node is not eligible for it."); QNBitPackedExperts& packed = (input_idx == 2) ? qnbit_fc1_ : qnbit_fc2_; - // Already initialized when this kernel packed the weights itself (the session-state prepack - // pass always runs PrePack before handing the buffers back here). + // Skip re-initialization (and the fp16 scales re-conversion) when this kernel already packed + // the weights itself: the session-state prepack pass always runs PrePack before handing the + // buffers back here. if (packed.packed_size_per_expert == 0) { - ORT_RETURN_IF_ERROR(InitQNBitPacked(packed, num_experts, rows, cols, *scales_tensor, + ORT_RETURN_IF_ERROR(InitQNBitPacked(packed, packed_size_per_expert, num_experts, rows, cols, *scales_tensor, Info().GetAllocator(OrtMemType::OrtMemTypeDefault))); } ORT_RETURN_IF_NOT(prepacked_buffer_sizes[0] == packed.packed_size_per_expert * static_cast(num_experts), @@ -975,8 +973,9 @@ QMoECPU::QMoECPU(const OpKernelInfo& op_kernel_info) template bool QMoECPU::QNBitGemmEligible(int input_idx, int64_t num_experts, int64_t rows, int64_t cols, - const Tensor** scales_out) const { + const Tensor** scales_out, size_t* packed_size_per_expert_out) const { *scales_out = nullptr; + *packed_size_per_expert_out = 0; if (!use_qnbit_gemm_ || (input_idx != 2 && input_idx != 5)) { return false; } @@ -1012,17 +1011,14 @@ bool QMoECPU::QNBitGemmEligible(int input_idx, int64_t num_experts, int64_t r } *scales_out = scales_tensor; + *packed_size_per_expert_out = packed_size; return true; } template -Status QMoECPU::InitQNBitPacked(QNBitPackedExperts& packed, int64_t num_experts, int64_t rows, int64_t cols, - const Tensor& scales, AllocatorPtr alloc) { - packed.packed_size_per_expert = MlasQNBitGemmPackQuantBDataSize( - static_cast(rows), static_cast(cols), static_cast(expert_weight_bits_), - static_cast(block_size_), /*HasZeroPoint*/ false, qnbit_compute_type_, - &mlas_backend_kernel_selector_config_); - ORT_RETURN_IF(packed.packed_size_per_expert == 0, "MLAS QNBit GEMM packing is not available for this expert shape."); +Status QMoECPU::InitQNBitPacked(QNBitPackedExperts& packed, size_t packed_size_per_expert, int64_t num_experts, + int64_t rows, int64_t cols, const Tensor& scales, AllocatorPtr alloc) { + packed.packed_size_per_expert = packed_size_per_expert; packed.scales_packed = MlasQNBitGemmScalesPacked( static_cast(cols), static_cast(expert_weight_bits_), static_cast(block_size_), qnbit_compute_type_, /*HasZeroPoint*/ false, &mlas_backend_kernel_selector_config_); @@ -1049,12 +1045,13 @@ Status QMoECPU::PrePackQNBitExperts(const Tensor& tensor, int input_idx, Allo const int64_t cols = shape[2] * pack_unit; const Tensor* scales_tensor = nullptr; - if (!QNBitGemmEligible(input_idx, num_experts, rows, cols, &scales_tensor)) { + size_t packed_size_per_expert = 0; + if (!QNBitGemmEligible(input_idx, num_experts, rows, cols, &scales_tensor, &packed_size_per_expert)) { return Status::OK(); } QNBitPackedExperts& packed = (input_idx == 2) ? qnbit_fc1_ : qnbit_fc2_; - ORT_RETURN_IF_ERROR(InitQNBitPacked(packed, num_experts, rows, cols, *scales_tensor, alloc)); + ORT_RETURN_IF_ERROR(InitQNBitPacked(packed, packed_size_per_expert, num_experts, rows, cols, *scales_tensor, alloc)); const float* scales_fp32 = nullptr; if constexpr (std::is_same_v) { @@ -1085,11 +1082,14 @@ Status QMoECPU::PrePackQNBitExperts(const Tensor& tensor, int input_idx, Allo // Experts pack independently into disjoint regions, so spread them over a load-time pool // (the session pool is not reachable from PrePack; same approach as MatMulNBits::PrePack). - OrtThreadPoolParams pack_tp_params; - pack_tp_params.thread_pool_size = Env::Default().GetNumPhysicalCpuCores(); - pack_tp_params.allow_spinning = false; - pack_tp_params.auto_set_affinity = false; - auto pack_tp = concurrency::CreateThreadPool(&Env::Default(), pack_tp_params, concurrency::ThreadPoolType::INTRA_OP); + std::unique_ptr pack_tp; + if (num_experts > 1) { + OrtThreadPoolParams pack_tp_params; + pack_tp_params.thread_pool_size = Env::Default().GetNumPhysicalCpuCores(); + pack_tp_params.allow_spinning = false; + pack_tp_params.auto_set_affinity = false; + pack_tp = concurrency::CreateThreadPool(&Env::Default(), pack_tp_params, concurrency::ThreadPoolType::INTRA_OP); + } concurrency::ThreadPool::TrySimpleParallelFor(pack_tp.get(), narrow(num_experts), [&](std::ptrdiff_t e) { std::byte* expert_dst = dst + static_cast(e) * per_expert; const float* expert_scales = scales_fp32 + static_cast(e) * scale_stride; @@ -1113,17 +1113,11 @@ Status QMoECPU::PrePackQNBitExperts(const Tensor& tensor, int input_idx, Allo prepacked_weights->buffers_.push_back(std::move(packed_buffer)); prepacked_weights->buffer_sizes_.push_back(total_packed_size); - // Shape buffer: rank, dims, then the layout tag so UseSharedPrePackedBuffers can tell this - // layout apart from the legacy unpacked-transposed one (which stores rank + dims only). - auto dims = shape.GetDims(); - const size_t shape_size = (2 + dims.size()) * sizeof(int64_t); - auto shape_buffer = IAllocator::MakeUniquePtr(alloc, shape_size); - int64_t* buffer_data = static_cast(shape_buffer.get()); - buffer_data[0] = static_cast(dims.size()); - memcpy(buffer_data + 1, dims.data(), dims.size() * sizeof(int64_t)); - buffer_data[1 + dims.size()] = kQNBitPackedLayoutTag; - prepacked_weights->buffers_.push_back(std::move(shape_buffer)); - prepacked_weights->buffer_sizes_.push_back(shape_size); + // The layout tag lets UseSharedPrePackedBuffers tell this layout apart from the legacy + // unpacked-transposed one (rank + dims only); the compute type guards against a buffer + // packed for different kernels. + const int64_t trailer[] = {kQNBitPackedLayoutTag, static_cast(qnbit_compute_type_)}; + PushPrePackedShapeBuffer(shape, trailer, alloc, *prepacked_weights); } else { packed.packed = std::move(packed_buffer); } @@ -1568,18 +1562,24 @@ Status QMoECPU::ComputeCommon(OpKernelContext* context, const ComputeInputs& qnbit_workspace_ptr = IAllocator::MakeUniquePtr(allocator, static_cast(num_expert_threads) * qnbit_workspace_per_thread, true); qnbit_workspace = qnbit_workspace_ptr.get(); } - auto run_qnbit_gemm = [&](const QNBitPackedExperts& packed, int64_t expert_idx, + // fp32 [E, N, K/block] scales for the kernels: the scales input itself for T == float, the copy + // converted at PrePack for T == MLFloat16. + const float* qnbit_fc1_scales = nullptr; + const float* qnbit_fc2_scales = nullptr; + if constexpr (std::is_same_v) { + qnbit_fc1_scales = qnbit_fc1_.scales_fp32.get(); + qnbit_fc2_scales = qnbit_fc2_.scales_fp32.get(); + } else { + qnbit_fc1_scales = fc1_scales_data; + qnbit_fc2_scales = fc2_scales_data; + } + auto run_qnbit_gemm = [&](const QNBitPackedExperts& packed, const float* scales_base, int64_t expert_idx, const float* A, size_t M, size_t N, size_t K, - const T* expert_scales, const float* bias, float* C, + const float* bias, float* C, std::byte* gemm_workspace, concurrency::ThreadPool* gemm_tp) { - const float* scales = nullptr; - if (!packed.scales_packed) { - if constexpr (std::is_same_v) { - scales = packed.scales_fp32.get() + static_cast(expert_idx) * N * (K / qnbit_blk); - } else { - scales = reinterpret_cast(expert_scales); - } - } + const float* scales = packed.scales_packed + ? nullptr + : scales_base + static_cast(expert_idx) * N * (K / qnbit_blk); const std::byte* packed_b = static_cast(packed.packed.get()) + static_cast(expert_idx) * packed.packed_size_per_expert; @@ -1728,7 +1728,7 @@ Status QMoECPU::ComputeCommon(OpKernelContext* context, const ComputeInputs& } fc1_bias_float = thread_bias1_buffer; } - run_qnbit_gemm(qnbit_fc1_, expert_idx, A1, m, n, k, fc1_scales_ptr, fc1_bias_float, C1, thread_qnbit_workspace, inner_tp); + run_qnbit_gemm(qnbit_fc1_, qnbit_fc1_scales, expert_idx, A1, m, n, k, fc1_bias_float, C1, thread_qnbit_workspace, inner_tp); goto fc1_gemm_done; } @@ -1987,7 +1987,7 @@ Status QMoECPU::ComputeCommon(OpKernelContext* context, const ComputeInputs& } fc2_bias_float = thread_bias2_buffer; } - run_qnbit_gemm(qnbit_fc2_, expert_idx, A2, m2, n2, k2, fc2_scales_ptr, fc2_bias_float, C2, thread_qnbit_workspace, inner_tp); + run_qnbit_gemm(qnbit_fc2_, qnbit_fc2_scales, expert_idx, A2, m2, n2, k2, fc2_bias_float, C2, thread_qnbit_workspace, inner_tp); fc2_bias_added_by_mlas = true; goto fc2_gemm_done; } diff --git a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h index 68e7c72a9cbe2..97801837da4e7 100644 --- a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h +++ b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h @@ -70,9 +70,9 @@ class QMoECPU final : public OpKernel, public MoEBaseCPU { bool scales_packed{false}; // scales are folded into `packed`; pass no QuantBScale at compute time }; bool QNBitGemmEligible(int input_idx, int64_t num_experts, int64_t rows, int64_t cols, - const Tensor** scales_out) const; - Status InitQNBitPacked(QNBitPackedExperts& packed, int64_t num_experts, int64_t rows, int64_t cols, - const Tensor& scales, AllocatorPtr alloc); + const Tensor** scales_out, size_t* packed_size_per_expert_out) const; + Status InitQNBitPacked(QNBitPackedExperts& packed, size_t packed_size_per_expert, int64_t num_experts, + int64_t rows, int64_t cols, const Tensor& scales, AllocatorPtr alloc); Status PrePackQNBitExperts(const Tensor& tensor, int input_idx, AllocatorPtr alloc, /*out*/ bool& is_packed, /*out*/ PrePackedWeights* prepacked_weights); diff --git a/onnxruntime/test/contrib_ops/moe_test.cc b/onnxruntime/test/contrib_ops/moe_test.cc index b5cd18deaa299..dc22e11936eed 100644 --- a/onnxruntime/test/contrib_ops/moe_test.cc +++ b/onnxruntime/test/contrib_ops/moe_test.cc @@ -2346,7 +2346,12 @@ TEST(MoETest, QMoETest_CPU_Int4_BlockWise_SwiGLU_Fp16) { } TEST(MoETest, QMoETest_CPU_Int8_BlockWise_SwiGLU) { - // 8-bit has no fp32-activation QNBit kernel, so this runs the int8-activation kernels where available. + // 8-bit stays on the dequantize path by default (no fp32-activation QNBit kernel). + RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true, 8}, 0.01f); +} + +TEST(MoETest, QMoETest_CPU_Int8_BlockWise_SwiGLU_Int8Activations) { + ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_QMOE_CPU_QNBIT_GEMM", "int8"}}}; RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true, 8}, 0.05f); } From 385cc3f62693c879e5f8557178a21efe37d0187c Mon Sep 17 00:00:00 2001 From: Yuri Khrustalev Date: Wed, 16 Sep 2026 11:36:30 -0400 Subject: [PATCH 03/11] feat: QMoE accuracy_level attribute, zero points and a fallback warning on the CPU QNBit path Add an accuracy_level attribute to QMoE with the MatMulNBits meaning so a model can opt into int8 activations (4) instead of relying on the ORT_QMOE_CPU_QNBIT_GEMM environment variable, which now only overrides the attribute. Pack constant block-wise zero points into the QNBit layout (QMoE's per-expert zero point layout is the MatMulNBits one) so asymmetric checkpoints take the same kernels. Log once when a block-wise node cannot use the path and why, since the dequantize fallback is an order of magnitude slower for decode. Record the zero point flag next to the compute type in the prepacked shape buffer, and cover accuracy_level and zero points (4-bit fp32 and int8, 8-bit int8) in the tests. Co-Authored-By: Claude Fable 5.1 --- docs/ContribOperators.md | 2 + .../cpu/moe/moe_quantization_cpu.cc | 151 ++++++++++++------ .../cpu/moe/moe_quantization_cpu.h | 17 +- .../core/graph/contrib_ops/contrib_defs.cc | 5 + onnxruntime/test/contrib_ops/moe_test.cc | 75 +++++++-- 5 files changed, 182 insertions(+), 68 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 7659698814171..baf781bb5744f 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -5675,6 +5675,8 @@ This version of the operator has been available since version 1 of the 'com.micr #### Attributes
+
accuracy_level : int
+
Minimum accuracy level of the expert GEMMs, with the MatMulNBits meaning: 0 (default) or 1 keeps fp32 activations; 4 allows int8 activations (int8 dot-product kernels) for block-wise 4/8-bit experts on CPU. Other values are treated as 0.
activation_alpha : float
Alpha parameter used in activation function.
activation_beta : float
diff --git a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc index 2afab3dc1fc81..9efe2f51be461 100644 --- a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc +++ b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc @@ -15,6 +15,7 @@ #include "core/framework/tensor_type_and_shape.h" #include "core/util/math.h" #include "core/platform/env_var_utils.h" +#include "core/common/logging/logging.h" #include "core/util/thread_utils.h" #include "contrib_ops/cpu/moe/moe_utils.h" #include "contrib_ops/cpu/moe/moe_helper.h" @@ -122,17 +123,17 @@ namespace contrib { constexpr const char* kUseMlasQ4GemmMoe = "ORT_USE_MLAS_Q4_GEMM_MOE"; -// Controls the MLAS QNBit GEMM (MatMulNBits kernel) path for block-wise 4/8-bit experts without -// zero points. Unset or "fp32": fp32 activations where the platform has that kernel (4-bit); 8-bit -// has no fp32 variant and keeps the dequantize path. "int8": int8 activations for both (the -// MatMulNBits accuracy_level=4 configuration). "0": disabled. Any other value is an error. +// Overrides the node's accuracy_level for the MLAS QNBit GEMM (MatMulNBits kernel) path of +// block-wise 4/8-bit experts. Unset: the attribute decides. "fp32": fp32 activations only +// (accuracy_level 0; 8-bit has no fp32 variant and keeps the dequantize path). "int8": int8 +// activations (accuracy_level 4). "0": path disabled. Any other value is an error. constexpr const char* kQMoEQNBitGemmEnv = "ORT_QMOE_CPU_QNBIT_GEMM"; // Tag appended to the prepacked shape buffer to mark the QNBit-packed layout (see PrePackQNBitExperts). constexpr int64_t kQNBitPackedLayoutTag = 0x514E4249; // 'QNBI' // Prepacked shape buffer: rank, dims, then `trailer` (empty for the legacy layouts; the QNBit layout -// appends kQNBitPackedLayoutTag and the compute type it was packed for). +// appends kQNBitPackedLayoutTag, the compute type it was packed for and whether zero points were folded in). static void PushPrePackedShapeBuffer(const TensorShape& shape, gsl::span trailer, const AllocatorPtr& alloc, PrePackedWeights& prepacked_weights) { const auto dims = shape.GetDims(); @@ -817,9 +818,9 @@ Status QMoECPU::UseSharedPrePackedBuffers(std::vector& prepa return Status::OK(); }; - // QNBit-packed layout: [packed experts, shape buffer = rank, dims, kQNBitPackedLayoutTag, compute type]. + // QNBit-packed layout: [packed experts, shape buffer = rank, dims, kQNBitPackedLayoutTag, compute type, has zero point]. if ((input_idx == 2 || input_idx == 5) && prepacked_buffers.size() == 2 && prepacked_buffer_sizes.size() == 2 && - prepacked_buffer_sizes[1] >= 6 * sizeof(int64_t)) { + prepacked_buffer_sizes[1] >= 7 * sizeof(int64_t)) { const int64_t* buffer_data = static_cast(prepacked_buffers[1].get()); if (buffer_data[0] == 3 && buffer_data[4] == kQNBitPackedLayoutTag) { ORT_RETURN_IF_NOT(use_qnbit_gemm_, "QMoE prepacked weights use the QNBit layout but this kernel has it disabled."); @@ -834,16 +835,18 @@ Status QMoECPU::UseSharedPrePackedBuffers(std::vector& prepa const int64_t num_experts = shape[0]; const int64_t rows = shape[1]; const int64_t cols = shape[2] * (8 / expert_weight_bits_); - const Tensor* scales_tensor = nullptr; - size_t packed_size_per_expert = 0; - ORT_RETURN_IF_NOT(QNBitGemmEligible(input_idx, num_experts, rows, cols, &scales_tensor, &packed_size_per_expert), + QNBitEligibility eligibility; + ORT_RETURN_IF_NOT(QNBitGemmEligible(input_idx, num_experts, rows, cols, eligibility), "QMoE prepacked weights use the QNBit layout but the node is not eligible for it."); + ORT_RETURN_IF_NOT(buffer_data[6] == (eligibility.zero_points != nullptr ? 1 : 0), + "QMoE prepacked weights were packed ", buffer_data[6] ? "with" : "without", + " zero points but this node has ", eligibility.zero_points != nullptr ? "them." : "none."); QNBitPackedExperts& packed = (input_idx == 2) ? qnbit_fc1_ : qnbit_fc2_; // Skip re-initialization (and the fp16 scales re-conversion) when this kernel already packed // the weights itself: the session-state prepack pass always runs PrePack before handing the // buffers back here. if (packed.packed_size_per_expert == 0) { - ORT_RETURN_IF_ERROR(InitQNBitPacked(packed, packed_size_per_expert, num_experts, rows, cols, *scales_tensor, + ORT_RETURN_IF_ERROR(InitQNBitPacked(packed, eligibility, num_experts, rows, cols, Info().GetAllocator(OrtMemType::OrtMemTypeDefault))); } ORT_RETURN_IF_NOT(prepacked_buffer_sizes[0] == packed.packed_size_per_expert * static_cast(num_experts), @@ -936,18 +939,19 @@ QMoECPU::QMoECPU(const OpKernelInfo& op_kernel_info) // QNBit GEMM path policy. The MatMulNBits kernels consume the block-wise QMoE encoding directly // (no re-quantization), have NEON/AVX2/AVX512 backends, and thread a single GEMM over N, so they - // are preferred for block-wise 4/8-bit experts without zero points. Int8 activations are lossier - // than the dequantize+SGEMM path they replace and QMoE has no accuracy_level attribute to opt in - // per model, so like MatMulNBits without accuracy_level=4 they are opt-in only: the default is - // fp32 activations where MLAS has that kernel (4-bit) and the existing path otherwise (8-bit). - bool allow_int8_compute = false; + // are preferred for block-wise 4/8-bit experts. Int8 activations are lossier than the + // dequantize+SGEMM path they replace, so as with MatMulNBits they need accuracy_level=4: the + // default is fp32 activations where MLAS has that kernel (4-bit) and the existing path otherwise + // (8-bit). The environment variable overrides the attribute. + accuracy_level_ = op_kernel_info.GetAttrOrDefault("accuracy_level", 0); + bool allow_int8_compute = (accuracy_level_ == 4); bool allow_fp32_compute = true; const auto qnbit_gemm_env = ParseEnvironmentVariable(kQMoEQNBitGemmEnv); if (qnbit_gemm_env.has_value()) { if (*qnbit_gemm_env == "0") { use_qnbit_gemm_ = false; } else if (*qnbit_gemm_env == "fp32") { - // Same as unset; accepted so the policy can be stated explicitly. + allow_int8_compute = false; } else if (*qnbit_gemm_env == "int8") { allow_int8_compute = true; allow_fp32_compute = false; @@ -973,64 +977,85 @@ QMoECPU::QMoECPU(const OpKernelInfo& op_kernel_info) template bool QMoECPU::QNBitGemmEligible(int input_idx, int64_t num_experts, int64_t rows, int64_t cols, - const Tensor** scales_out, size_t* packed_size_per_expert_out) const { - *scales_out = nullptr; - *packed_size_per_expert_out = 0; + QNBitEligibility& out) const { + out = QNBitEligibility{}; if (!use_qnbit_gemm_ || (input_idx != 2 && input_idx != 5)) { return false; } - if (block_size_ <= 0 || (cols % block_size_) != 0) { + if (block_size_ <= 0) { + return false; + } + if ((cols % block_size_) != 0) { + out.ineligible_reason = "the input feature size is not a multiple of block_size"; return false; } const int scales_idx = (input_idx == 2) ? 3 : 6; const int zp_idx = (input_idx == 2) ? 11 : 12; const auto& input_defs = Info().node().InputDefs(); - const bool has_zp_input = zp_idx < static_cast(input_defs.size()) && input_defs[zp_idx]->Exists(); - if (has_zp_input) { - return false; - } + const int64_t blocks_per_row = cols / block_size_; const Tensor* scales_tensor = nullptr; if (!Info().TryGetConstantInput(scales_idx, &scales_tensor) || scales_tensor == nullptr) { + out.ineligible_reason = "the scales are not a constant initializer"; return false; } const auto& scales_dims = scales_tensor->Shape().GetDims(); - const int64_t blocks_per_row = cols / block_size_; if (scales_dims.size() != 3 || scales_dims[0] != num_experts || scales_dims[1] != rows || scales_dims[2] != blocks_per_row) { return false; // row-wise (per-channel) scales keep the existing paths } + const bool has_zp_input = zp_idx < static_cast(input_defs.size()) && input_defs[zp_idx]->Exists(); + const Tensor* zp_tensor = nullptr; + if (has_zp_input) { + if (!Info().TryGetConstantInput(zp_idx, &zp_tensor) || zp_tensor == nullptr) { + out.ineligible_reason = "the zero points are not a constant initializer"; + return false; + } + // The kernels read the MatMulNBits zero point layout, [rows, ceil(blocks/pack)] per expert with + // the even block in the low nibble, which is QMoE's block-wise layout. With a single block per + // row QMoE switches to the row-wise [ceil(rows/pack)] layout instead, so that case stays out. + const int64_t zp_pack = 8 / expert_weight_bits_; + const auto& zp_dims = zp_tensor->Shape().GetDims(); + if (blocks_per_row < 2 || zp_dims.size() != 3 || zp_dims[0] != num_experts || zp_dims[1] != rows || + zp_dims[2] != (blocks_per_row + zp_pack - 1) / zp_pack) { + out.ineligible_reason = "the zero points are not in the block-wise [num_experts, rows, blocks/pack] layout"; + return false; + } + } + const size_t packed_size = MlasQNBitGemmPackQuantBDataSize( static_cast(rows), static_cast(cols), static_cast(expert_weight_bits_), - static_cast(block_size_), /*HasZeroPoint*/ false, qnbit_compute_type_, + static_cast(block_size_), has_zp_input, qnbit_compute_type_, &mlas_backend_kernel_selector_config_); if (packed_size == 0) { + out.ineligible_reason = "MLAS has no QNBit packing for this shape on this platform"; return false; } - *scales_out = scales_tensor; - *packed_size_per_expert_out = packed_size; + out.scales = scales_tensor; + out.zero_points = zp_tensor; + out.packed_size_per_expert = packed_size; return true; } template -Status QMoECPU::InitQNBitPacked(QNBitPackedExperts& packed, size_t packed_size_per_expert, int64_t num_experts, - int64_t rows, int64_t cols, const Tensor& scales, AllocatorPtr alloc) { - packed.packed_size_per_expert = packed_size_per_expert; +Status QMoECPU::InitQNBitPacked(QNBitPackedExperts& packed, const QNBitEligibility& eligibility, + int64_t num_experts, int64_t rows, int64_t cols, AllocatorPtr alloc) { + packed.packed_size_per_expert = eligibility.packed_size_per_expert; + packed.has_zero_point = (eligibility.zero_points != nullptr); packed.scales_packed = MlasQNBitGemmScalesPacked( static_cast(cols), static_cast(expert_weight_bits_), static_cast(block_size_), - qnbit_compute_type_, /*HasZeroPoint*/ false, &mlas_backend_kernel_selector_config_); + qnbit_compute_type_, packed.has_zero_point, &mlas_backend_kernel_selector_config_); if constexpr (std::is_same_v) { // The MLAS kernels take fp32 scales; convert the constant fp16 scales once. const size_t scales_count = static_cast(num_experts * rows * (cols / block_size_)); packed.scales_fp32 = IAllocator::MakeUniquePtr(alloc, scales_count, true); - MlasConvertHalfToFloatBuffer(scales.Data(), packed.scales_fp32.get(), scales_count); + MlasConvertHalfToFloatBuffer(eligibility.scales->Data(), packed.scales_fp32.get(), scales_count); } else { ORT_UNUSED_PARAMETER(num_experts); - ORT_UNUSED_PARAMETER(scales); } return Status::OK(); } @@ -1044,21 +1069,31 @@ Status QMoECPU::PrePackQNBitExperts(const Tensor& tensor, int input_idx, Allo const int64_t pack_unit = 8 / expert_weight_bits_; const int64_t cols = shape[2] * pack_unit; - const Tensor* scales_tensor = nullptr; - size_t packed_size_per_expert = 0; - if (!QNBitGemmEligible(input_idx, num_experts, rows, cols, &scales_tensor, &packed_size_per_expert)) { + QNBitEligibility eligibility; + if (!QNBitGemmEligible(input_idx, num_experts, rows, cols, eligibility)) { + // The dequantize+SGEMM fallback is an order of magnitude slower for decode, so say why once. + if (eligibility.ineligible_reason != nullptr && !qnbit_fallback_logged_) { + qnbit_fallback_logged_ = true; + LOGS_DEFAULT(WARNING) << "QMoE node '" << Info().node().Name() << "': block-wise expert weights (input " + << input_idx << ") cannot use the MLAS QNBit GEMM kernels because " + << eligibility.ineligible_reason << "; falling back to dequantize + SGEMM."; + } return Status::OK(); } QNBitPackedExperts& packed = (input_idx == 2) ? qnbit_fc1_ : qnbit_fc2_; - ORT_RETURN_IF_ERROR(InitQNBitPacked(packed, packed_size_per_expert, num_experts, rows, cols, *scales_tensor, alloc)); + ORT_RETURN_IF_ERROR(InitQNBitPacked(packed, eligibility, num_experts, rows, cols, alloc)); const float* scales_fp32 = nullptr; if constexpr (std::is_same_v) { scales_fp32 = packed.scales_fp32.get(); } else { - scales_fp32 = scales_tensor->Data(); + scales_fp32 = eligibility.scales->Data(); } + const uint8_t* zp_data = packed.has_zero_point ? eligibility.zero_points->Data() : nullptr; + const size_t zp_stride = packed.has_zero_point + ? static_cast(eligibility.zero_points->Shape()[1] * eligibility.zero_points->Shape()[2]) + : 0; const size_t nbits = static_cast(expert_weight_bits_); const size_t blk = static_cast(block_size_); @@ -1093,13 +1128,14 @@ Status QMoECPU::PrePackQNBitExperts(const Tensor& tensor, int input_idx, Allo concurrency::ThreadPool::TrySimpleParallelFor(pack_tp.get(), narrow(num_experts), [&](std::ptrdiff_t e) { std::byte* expert_dst = dst + static_cast(e) * per_expert; const float* expert_scales = scales_fp32 + static_cast(e) * scale_stride; + const uint8_t* expert_zp = (zp_data == nullptr) ? nullptr : zp_data + static_cast(e) * zp_stride; MlasQNBitGemmPackQuantBData(static_cast(rows), static_cast(cols), nbits, blk, qnbit_compute_type_, qdata + static_cast(e) * qdata_stride, expert_dst, expert_scales, - /*HasZeroPoint*/ false, nullptr, nullptr, &mlas_backend_kernel_selector_config_); + packed.has_zero_point, expert_zp, nullptr, &mlas_backend_kernel_selector_config_); if (finalize_scales) { MlasQNBitGemmPackQuantBData(static_cast(rows), static_cast(cols), nbits, blk, qnbit_compute_type_, nullptr, expert_dst, expert_scales, - /*HasZeroPoint*/ false, nullptr, nullptr, &mlas_backend_kernel_selector_config_); + packed.has_zero_point, expert_zp, nullptr, &mlas_backend_kernel_selector_config_); } }); @@ -1114,9 +1150,10 @@ Status QMoECPU::PrePackQNBitExperts(const Tensor& tensor, int input_idx, Allo prepacked_weights->buffer_sizes_.push_back(total_packed_size); // The layout tag lets UseSharedPrePackedBuffers tell this layout apart from the legacy - // unpacked-transposed one (rank + dims only); the compute type guards against a buffer - // packed for different kernels. - const int64_t trailer[] = {kQNBitPackedLayoutTag, static_cast(qnbit_compute_type_)}; + // unpacked-transposed one (rank + dims only); the compute type and zero point flag guard + // against a buffer packed for different kernels. + const int64_t trailer[] = {kQNBitPackedLayoutTag, static_cast(qnbit_compute_type_), + packed.has_zero_point ? int64_t{1} : int64_t{0}}; PushPrePackedShapeBuffer(shape, trailer, alloc, *prepacked_weights); } else { packed.packed = std::move(packed_buffer); @@ -1547,14 +1584,16 @@ Status QMoECPU::ComputeCommon(OpKernelContext* context, const ComputeInputs& // batch once per thread rather than allocating per expert. size_t qnbit_workspace_per_thread = 0; if (use_qnbit_fc1) { - qnbit_workspace_per_thread = std::max(qnbit_workspace_per_thread, MlasQNBitGemmBatchWorkspaceSize( + const size_t fc1_workspace = MlasQNBitGemmBatchWorkspaceSize( static_cast(max_tokens_per_expert), static_cast(fc1_out_features), static_cast(hidden_size), - 1, qnbit_bits, qnbit_blk, /*HasZeroPoint*/ false, qnbit_compute_type_, &mlas_backend_kernel_selector_config_)); + 1, qnbit_bits, qnbit_blk, qnbit_fc1_.has_zero_point, qnbit_compute_type_, &mlas_backend_kernel_selector_config_); + qnbit_workspace_per_thread = std::max(qnbit_workspace_per_thread, fc1_workspace); } if (use_qnbit_fc2) { - qnbit_workspace_per_thread = std::max(qnbit_workspace_per_thread, MlasQNBitGemmBatchWorkspaceSize( + const size_t fc2_workspace = MlasQNBitGemmBatchWorkspaceSize( static_cast(max_tokens_per_expert), static_cast(hidden_size), static_cast(inter_size), - 1, qnbit_bits, qnbit_blk, /*HasZeroPoint*/ false, qnbit_compute_type_, &mlas_backend_kernel_selector_config_)); + 1, qnbit_bits, qnbit_blk, qnbit_fc2_.has_zero_point, qnbit_compute_type_, &mlas_backend_kernel_selector_config_); + qnbit_workspace_per_thread = std::max(qnbit_workspace_per_thread, fc2_workspace); } IAllocatorUniquePtr qnbit_workspace_ptr; std::byte* qnbit_workspace = nullptr; @@ -1573,13 +1612,17 @@ Status QMoECPU::ComputeCommon(OpKernelContext* context, const ComputeInputs& qnbit_fc1_scales = fc1_scales_data; qnbit_fc2_scales = fc2_scales_data; } - auto run_qnbit_gemm = [&](const QNBitPackedExperts& packed, const float* scales_base, int64_t expert_idx, + auto run_qnbit_gemm = [&](const QNBitPackedExperts& packed, const float* scales_base, + const uint8_t* zp_base, int64_t zp_expert_stride, int64_t expert_idx, const float* A, size_t M, size_t N, size_t K, const float* bias, float* C, std::byte* gemm_workspace, concurrency::ThreadPool* gemm_tp) { const float* scales = packed.scales_packed ? nullptr : scales_base + static_cast(expert_idx) * N * (K / qnbit_blk); + const uint8_t* zero_points = packed.has_zero_point + ? zp_base + static_cast(expert_idx * zp_expert_stride) + : nullptr; const std::byte* packed_b = static_cast(packed.packed.get()) + static_cast(expert_idx) * packed.packed_size_per_expert; @@ -1589,7 +1632,7 @@ Status QMoECPU::ComputeCommon(OpKernelContext* context, const ComputeInputs& params.QuantBDataWorkspace = packed_b; params.PackedQuantBData = packed_b; params.QuantBScale = scales; - params.QuantBZeroPoint = nullptr; + params.QuantBZeroPoint = zero_points; params.Bias = bias; params.C = C; params.ldc = N; @@ -1728,7 +1771,8 @@ Status QMoECPU::ComputeCommon(OpKernelContext* context, const ComputeInputs& } fc1_bias_float = thread_bias1_buffer; } - run_qnbit_gemm(qnbit_fc1_, qnbit_fc1_scales, expert_idx, A1, m, n, k, fc1_bias_float, C1, thread_qnbit_workspace, inner_tp); + run_qnbit_gemm(qnbit_fc1_, qnbit_fc1_scales, fc1_zp_data, fc1_zp_expert_stride, expert_idx, A1, m, n, k, + fc1_bias_float, C1, thread_qnbit_workspace, inner_tp); goto fc1_gemm_done; } @@ -1987,7 +2031,8 @@ Status QMoECPU::ComputeCommon(OpKernelContext* context, const ComputeInputs& } fc2_bias_float = thread_bias2_buffer; } - run_qnbit_gemm(qnbit_fc2_, qnbit_fc2_scales, expert_idx, A2, m2, n2, k2, fc2_bias_float, C2, thread_qnbit_workspace, inner_tp); + run_qnbit_gemm(qnbit_fc2_, qnbit_fc2_scales, fc2_zp_data, fc2_zp_expert_stride, expert_idx, A2, m2, n2, k2, + fc2_bias_float, C2, thread_qnbit_workspace, inner_tp); fc2_bias_added_by_mlas = true; goto fc2_gemm_done; } diff --git a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h index 97801837da4e7..0993e7c6eb2e5 100644 --- a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h +++ b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h @@ -68,11 +68,20 @@ class QMoECPU final : public OpKernel, public MoEBaseCPU { size_t packed_size_per_expert{0}; IAllocatorUniquePtr scales_fp32; // fp32 copy of the [E, rows, cols/block_size] scales (T == MLFloat16 only) bool scales_packed{false}; // scales are folded into `packed`; pass no QuantBScale at compute time + bool has_zero_point{false}; // packed with the constant zero points; pass the expert's slice at compute time + }; + // Result of QNBitGemmEligible. `ineligible_reason` is set only when a block-wise node misses the + // path for a reason worth surfacing (row-wise scales and a disabled path leave it null). + struct QNBitEligibility { + const Tensor* scales{nullptr}; + const Tensor* zero_points{nullptr}; // null when the node has no zero points input + size_t packed_size_per_expert{0}; + const char* ineligible_reason{nullptr}; }; bool QNBitGemmEligible(int input_idx, int64_t num_experts, int64_t rows, int64_t cols, - const Tensor** scales_out, size_t* packed_size_per_expert_out) const; - Status InitQNBitPacked(QNBitPackedExperts& packed, size_t packed_size_per_expert, int64_t num_experts, - int64_t rows, int64_t cols, const Tensor& scales, AllocatorPtr alloc); + QNBitEligibility& out) const; + Status InitQNBitPacked(QNBitPackedExperts& packed, const QNBitEligibility& eligibility, int64_t num_experts, + int64_t rows, int64_t cols, AllocatorPtr alloc); Status PrePackQNBitExperts(const Tensor& tensor, int input_idx, AllocatorPtr alloc, /*out*/ bool& is_packed, /*out*/ PrePackedWeights* prepacked_weights); @@ -92,7 +101,9 @@ class QMoECPU final : public OpKernel, public MoEBaseCPU { IAllocatorUniquePtr packed_fc1_mlas_cache_; IAllocatorUniquePtr packed_fc2_mlas_cache_; + int64_t accuracy_level_{0}; bool use_qnbit_gemm_{true}; + bool qnbit_fallback_logged_{false}; MLAS_QNBIT_GEMM_COMPUTE_TYPE qnbit_compute_type_{SQNBIT_CompFp32}; QNBitPackedExperts qnbit_fc1_; QNBitPackedExperts qnbit_fc2_; diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index ce531a56cbf7b..76e426b2f5b98 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -1505,6 +1505,11 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Attr("activation_beta", "Beta parameter used in activation function.", AttributeProto::FLOAT, 0.0f) + .Attr("accuracy_level", + "Minimum accuracy level of the expert GEMMs, with the MatMulNBits meaning: 0 (default) or 1 keeps " + "fp32 activations; 4 allows int8 activations (int8 dot-product kernels) for block-wise 4/8-bit " + "experts on CPU. Other values are treated as 0.", + AttributeProto::INT, static_cast(0)) .Attr("block_size", "Size of each quantization block along the K (input feature) dimension. " "Must be power of two and ≥ 16 (e.g., 16, 32, 64, 128). " diff --git a/onnxruntime/test/contrib_ops/moe_test.cc b/onnxruntime/test/contrib_ops/moe_test.cc index dc22e11936eed..b7fc54c9a63c3 100644 --- a/onnxruntime/test/contrib_ops/moe_test.cc +++ b/onnxruntime/test/contrib_ops/moe_test.cc @@ -2127,11 +2127,12 @@ TEST(MoETest, QMoETest_CPU_Int4_MLAS) { #endif } -// Block-wise 4/8-bit SwiGLU QMoE on CPU with non-trivial weights. Block-wise symmetric experts -// without zero points take the MLAS QNBit GEMM (MatMulNBits kernel) path; the expected output is an -// fp32 reference computed here from the dequantized weights and the kernel's routing (softmax over -// the top-k router logits). Weights, scales and biases are constant initializers so that PrePack runs; -// as graph inputs the kernel would silently fall back to the dequantize+SGEMM path. +// Block-wise 4/8-bit SwiGLU QMoE on CPU with non-trivial weights. Block-wise experts with constant +// scales (and, if present, constant zero points) take the MLAS QNBit GEMM (MatMulNBits kernel) path; +// the expected output is an fp32 reference computed here from the dequantized weights and the +// kernel's routing (softmax over the top-k router logits). Weights, scales, zero points and biases +// are constant initializers so that PrePack runs; as graph inputs the kernel would fall back to the +// dequantize+SGEMM path. namespace { struct QMoEBlockWiseCase { int num_rows; @@ -2142,6 +2143,8 @@ struct QMoEBlockWiseCase { int top_k; bool with_bias; int bits{4}; + bool with_zero_points{false}; + int accuracy_level{0}; }; template @@ -2154,6 +2157,7 @@ void RunQMoECpuBlockWiseSwiGLU(const QMoEBlockWiseCase& c, float tolerance, const int fc2_blocks = inter_size / block_size; const int pack = 8 / bits; const float code_zero = (bits == 4) ? 8.0f : 128.0f; + const int code_max = (1 << bits) - 1; constexpr bool is_fp16 = std::is_same_v; uint32_t state = 0x12345678u; @@ -2189,6 +2193,28 @@ void RunQMoECpuBlockWiseSwiGLU(const QMoEBlockWiseCase& c, float tolerance, }; const std::vector fc1_weights = pack_codes(fc1_codes); const std::vector fc2_weights = pack_codes(fc2_codes); + // Zero points: one code per block, stored [E, N, ceil(blocks/pack)] with the lower block in the + // low bits (the MatMulNBits layout); kept away from the code range ends so weights stay signed. + std::vector fc1_zp_codes, fc2_zp_codes, fc1_zero_points, fc2_zero_points; + if (c.with_zero_points) { + fc1_zp_codes.resize(static_cast(num_experts * fc1_rows * fc1_blocks)); + fc2_zp_codes.resize(static_cast(num_experts * hidden_size * fc2_blocks)); + for (auto& v : fc1_zp_codes) v = static_cast(1 + static_cast(next_uniform() * (code_max - 2))); + for (auto& v : fc2_zp_codes) v = static_cast(1 + static_cast(next_uniform() * (code_max - 2))); + auto pack_zero_points = [pack, bits](const std::vector& codes, int rows_total, int blocks) { + const int packed_blocks = (blocks + pack - 1) / pack; + std::vector packed(static_cast(rows_total * packed_blocks), 0); + for (int r = 0; r < rows_total; ++r) { + for (int b = 0; b < blocks; ++b) { + packed[static_cast(r * packed_blocks + b / pack)] |= + static_cast(codes[static_cast(r * blocks + b)] << ((b % pack) * bits)); + } + } + return packed; + }; + fc1_zero_points = pack_zero_points(fc1_zp_codes, num_experts * fc1_rows, fc1_blocks); + fc2_zero_points = pack_zero_points(fc2_zp_codes, num_experts * hidden_size, fc2_blocks); + } std::vector fc1_scales(static_cast(num_experts * fc1_rows * fc1_blocks)); std::vector fc2_scales(static_cast(num_experts * hidden_size * fc2_blocks)); // Signed scales, as produced by MatMulNBits-style quantizers; scaled so the weight range is @@ -2204,10 +2230,12 @@ void RunQMoECpuBlockWiseSwiGLU(const QMoEBlockWiseCase& c, float tolerance, for (auto& v : fc2_bias) v = round_to_t(0.2f * next_uniform() - 0.1f); } - auto dequant = [&](const std::vector& codes, const std::vector& scales, int e, int n, int k, - int rows, int cols, int blocks) { - const float scale = scales[static_cast((e * rows + n) * blocks + k / block_size)]; - return (static_cast(codes[static_cast((e * rows + n) * cols + k)]) - code_zero) * scale; + auto dequant = [&](const std::vector& codes, const std::vector& scales, + const std::vector& zp_codes, int e, int n, int k, int rows, int cols, int blocks) { + const size_t block_index = static_cast((e * rows + n) * blocks + k / block_size); + const float scale = scales[block_index]; + const float zero = c.with_zero_points ? static_cast(zp_codes[block_index]) : code_zero; + return (static_cast(codes[static_cast((e * rows + n) * cols + k)]) - zero) * scale; }; std::vector expected(static_cast(num_rows * hidden_size), 0.0f); @@ -2229,7 +2257,7 @@ void RunQMoECpuBlockWiseSwiGLU(const QMoEBlockWiseCase& c, float tolerance, std::vector h(static_cast(fc1_rows)); for (int n = 0; n < fc1_rows; ++n) { float acc = c.with_bias ? fc1_bias[static_cast(e * fc1_rows + n)] : 0.0f; - for (int k = 0; k < hidden_size; ++k) acc += x[k] * dequant(fc1_codes, fc1_scales, e, n, k, fc1_rows, hidden_size, fc1_blocks); + for (int k = 0; k < hidden_size; ++k) acc += x[k] * dequant(fc1_codes, fc1_scales, fc1_zp_codes, e, n, k, fc1_rows, hidden_size, fc1_blocks); h[n] = acc; } std::vector a(static_cast(inter_size)); @@ -2239,7 +2267,7 @@ void RunQMoECpuBlockWiseSwiGLU(const QMoEBlockWiseCase& c, float tolerance, } for (int n = 0; n < hidden_size; ++n) { float acc = c.with_bias ? fc2_bias[static_cast(e * hidden_size + n)] : 0.0f; - for (int k = 0; k < inter_size; ++k) acc += a[k] * dequant(fc2_codes, fc2_scales, e, n, k, hidden_size, inter_size, fc2_blocks); + for (int k = 0; k < inter_size; ++k) acc += a[k] * dequant(fc2_codes, fc2_scales, fc2_zp_codes, e, n, k, hidden_size, inter_size, fc2_blocks); expected[static_cast(t * hidden_size + n)] += w * acc; } } @@ -2254,6 +2282,9 @@ void RunQMoECpuBlockWiseSwiGLU(const QMoEBlockWiseCase& c, float tolerance, tester.AddAttribute("normalize_routing_weights", 1); tester.AddAttribute("expert_weight_bits", bits); tester.AddAttribute("block_size", block_size); + if (c.accuracy_level != 0) { + tester.AddAttribute("accuracy_level", c.accuracy_level); + } auto add_t_input = [&](const char* name, std::vector dims, const std::vector& values, bool is_initializer) { if constexpr (is_fp16) { @@ -2278,6 +2309,13 @@ void RunQMoECpuBlockWiseSwiGLU(const QMoEBlockWiseCase& c, float tolerance, } else { tester.AddOptionalInputEdge(); } + if (c.with_zero_points) { + tester.AddOptionalInputEdge(); // fc3_experts_weights + tester.AddOptionalInputEdge(); // fc3_scales + tester.AddOptionalInputEdge(); // fc3_experts_bias + tester.AddInput("fc1_zero_points", {num_experts, fc1_rows, (fc1_blocks + pack - 1) / pack}, fc1_zero_points, true); + tester.AddInput("fc2_zero_points", {num_experts, hidden_size, (fc2_blocks + pack - 1) / pack}, fc2_zero_points, true); + } if constexpr (is_fp16) { tester.AddOutput("output", {num_rows, hidden_size}, ToFloat16(expected)); } else { @@ -2356,8 +2394,21 @@ TEST(MoETest, QMoETest_CPU_Int8_BlockWise_SwiGLU_Int8Activations) { } TEST(MoETest, QMoETest_CPU_Int4_BlockWise_SwiGLU_Int8Activations) { + // accuracy_level=4 opts into int8 activations, as for MatMulNBits. + RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true, 4, false, 4}, 0.05f); +} + +TEST(MoETest, QMoETest_CPU_Int4_BlockWise_SwiGLU_ZeroPoints) { + RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true, 4, true}, 0.01f); +} + +TEST(MoETest, QMoETest_CPU_Int4_BlockWise_SwiGLU_ZeroPoints_Int8Activations) { + RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true, 4, true, 4}, 0.05f); +} + +TEST(MoETest, QMoETest_CPU_Int8_BlockWise_SwiGLU_ZeroPoints_Int8Activations) { ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_QMOE_CPU_QNBIT_GEMM", "int8"}}}; - RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true}, 0.05f); + RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true, 8, true}, 0.05f); } TEST(MoETest, QMoETest_CPU_Int4_BlockWise_SwiGLU_SharedPrePackedWeights) { From b0361313c140e7fb34c340e044561073b6c9ae3f Mon Sep 17 00:00:00 2001 From: Yuri Khrustalev Date: Wed, 16 Sep 2026 11:42:24 -0400 Subject: [PATCH 04/11] fix: dependent template keyword for Tensor::Data calls in the QMoE QNBit prepack Co-Authored-By: Claude Fable 5.1 --- onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc index 9efe2f51be461..7f2e3878b5286 100644 --- a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc +++ b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc @@ -1053,7 +1053,7 @@ Status QMoECPU::InitQNBitPacked(QNBitPackedExperts& packed, const QNBitEligib // The MLAS kernels take fp32 scales; convert the constant fp16 scales once. const size_t scales_count = static_cast(num_experts * rows * (cols / block_size_)); packed.scales_fp32 = IAllocator::MakeUniquePtr(alloc, scales_count, true); - MlasConvertHalfToFloatBuffer(eligibility.scales->Data(), packed.scales_fp32.get(), scales_count); + MlasConvertHalfToFloatBuffer(eligibility.scales->template Data(), packed.scales_fp32.get(), scales_count); } else { ORT_UNUSED_PARAMETER(num_experts); } @@ -1088,9 +1088,9 @@ Status QMoECPU::PrePackQNBitExperts(const Tensor& tensor, int input_idx, Allo if constexpr (std::is_same_v) { scales_fp32 = packed.scales_fp32.get(); } else { - scales_fp32 = eligibility.scales->Data(); + scales_fp32 = eligibility.scales->template Data(); } - const uint8_t* zp_data = packed.has_zero_point ? eligibility.zero_points->Data() : nullptr; + const uint8_t* zp_data = packed.has_zero_point ? eligibility.zero_points->template Data() : nullptr; const size_t zp_stride = packed.has_zero_point ? static_cast(eligibility.zero_points->Shape()[1] * eligibility.zero_points->Shape()[2]) : 0; From c54c674c3eb037d44e04dfb599bcfc23325acbb0 Mon Sep 17 00:00:00 2001 From: Yuri Khrustalev Date: Wed, 16 Sep 2026 14:28:38 -0400 Subject: [PATCH 05/11] style: clang-format the QMoE QNBit packed-experts struct --- onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h index 0993e7c6eb2e5..d0a43d0064da0 100644 --- a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h +++ b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h @@ -64,7 +64,7 @@ class QMoECPU final : public OpKernel, public MoEBaseCPU { // block-wise 4/8-bit experts without zero points; each expert's [rows, cols] matrix is packed // independently and stored back to back. struct QNBitPackedExperts { - IAllocatorUniquePtr packed; // num_experts * packed_size_per_expert bytes + IAllocatorUniquePtr packed; // num_experts * packed_size_per_expert bytes size_t packed_size_per_expert{0}; IAllocatorUniquePtr scales_fp32; // fp32 copy of the [E, rows, cols/block_size] scales (T == MLFloat16 only) bool scales_packed{false}; // scales are folded into `packed`; pass no QuantBScale at compute time From 28ec438a59a6487a62e6f6a0440d83a5ecc6bc33 Mon Sep 17 00:00:00 2001 From: Yuri Khrustalev Date: Wed, 16 Sep 2026 16:17:56 -0400 Subject: [PATCH 06/11] fix: address QMoE QNBit review findings Use int8 activations where MLAS has no fp32 kernel for the bit width (8-bit), so 8-bit experts reach the QNBit path by default as documented; explicit fp32 still forces the previous behavior. Guard the packed-buffer size with SafeInt and size the load-time pack pool by the expert count. --- .../cpu/moe/moe_quantization_cpu.cc | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc index 7f2e3878b5286..19c312fd2c1d6 100644 --- a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc +++ b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc @@ -941,17 +941,19 @@ QMoECPU::QMoECPU(const OpKernelInfo& op_kernel_info) // (no re-quantization), have NEON/AVX2/AVX512 backends, and thread a single GEMM over N, so they // are preferred for block-wise 4/8-bit experts. Int8 activations are lossier than the // dequantize+SGEMM path they replace, so as with MatMulNBits they need accuracy_level=4: the - // default is fp32 activations where MLAS has that kernel (4-bit) and the existing path otherwise - // (8-bit). The environment variable overrides the attribute. + // default is fp32 activations where MLAS has that kernel (4-bit) and int8 where it does not + // (8-bit, which has no fp32 variant). The environment variable overrides the attribute. accuracy_level_ = op_kernel_info.GetAttrOrDefault("accuracy_level", 0); bool allow_int8_compute = (accuracy_level_ == 4); bool allow_fp32_compute = true; + bool fp32_compute_forced = false; const auto qnbit_gemm_env = ParseEnvironmentVariable(kQMoEQNBitGemmEnv); if (qnbit_gemm_env.has_value()) { if (*qnbit_gemm_env == "0") { use_qnbit_gemm_ = false; } else if (*qnbit_gemm_env == "fp32") { allow_int8_compute = false; + fp32_compute_forced = true; } else if (*qnbit_gemm_env == "int8") { allow_int8_compute = true; allow_fp32_compute = false; @@ -963,9 +965,12 @@ QMoECPU::QMoECPU(const OpKernelInfo& op_kernel_info) if (use_qnbit_gemm_ && (expert_weight_bits_ == 4 || expert_weight_bits_ == 8) && block_size_ > 0) { const size_t nbits = static_cast(expert_weight_bits_); const size_t blk = static_cast(block_size_); + // MLAS has no fp32-activation kernel for every bit width (8-bit is int8 only). Where fp32 is + // unavailable, int8 is the only way onto these kernels, so take it unless fp32 was forced. + const bool int8_is_only_option = !MlasIsQNBitGemmAvailable(nbits, blk, SQNBIT_CompFp32) && !fp32_compute_forced; if (allow_fp32_compute && MlasIsQNBitGemmAvailable(nbits, blk, SQNBIT_CompFp32)) { qnbit_compute_type_ = SQNBIT_CompFp32; - } else if (allow_int8_compute && MlasIsQNBitGemmAvailable(nbits, blk, SQNBIT_CompInt8)) { + } else if ((allow_int8_compute || int8_is_only_option) && MlasIsQNBitGemmAvailable(nbits, blk, SQNBIT_CompInt8)) { qnbit_compute_type_ = SQNBIT_CompInt8; } else { use_qnbit_gemm_ = false; @@ -1098,7 +1103,7 @@ Status QMoECPU::PrePackQNBitExperts(const Tensor& tensor, int input_idx, Allo const size_t nbits = static_cast(expert_weight_bits_); const size_t blk = static_cast(block_size_); const size_t per_expert = packed.packed_size_per_expert; - const size_t total_packed_size = per_expert * static_cast(num_experts); + const size_t total_packed_size = SafeInt(per_expert) * static_cast(num_experts); auto packed_buffer = IAllocator::MakeUniquePtr(alloc, total_packed_size, true); // The pack routines need not write every byte; zero so the prepacked-weights hash is deterministic. std::memset(packed_buffer.get(), 0, total_packed_size); @@ -1118,9 +1123,10 @@ Status QMoECPU::PrePackQNBitExperts(const Tensor& tensor, int input_idx, Allo // Experts pack independently into disjoint regions, so spread them over a load-time pool // (the session pool is not reachable from PrePack; same approach as MatMulNBits::PrePack). std::unique_ptr pack_tp; - if (num_experts > 1) { + const int pack_threads = narrow(std::min(num_experts, Env::Default().GetNumPhysicalCpuCores())); + if (pack_threads > 1) { OrtThreadPoolParams pack_tp_params; - pack_tp_params.thread_pool_size = Env::Default().GetNumPhysicalCpuCores(); + pack_tp_params.thread_pool_size = pack_threads; pack_tp_params.allow_spinning = false; pack_tp_params.auto_set_affinity = false; pack_tp = concurrency::CreateThreadPool(&Env::Default(), pack_tp_params, concurrency::ThreadPoolType::INTRA_OP); From d026ecb0da5f9c850aba83df6e79d3dc24277ba4 Mon Sep 17 00:00:00 2001 From: Yuri Khrustalev Date: Wed, 16 Sep 2026 16:23:39 -0400 Subject: [PATCH 07/11] test: cover the 8-bit QNBit default and its forced-fp32 fallback 8-bit experts now run with int8 activations by default, so the block-wise 8-bit case carries their quantization error; match the tolerance of the other int8 cases and add a case that forces fp32 back onto the dequantize path. --- onnxruntime/test/contrib_ops/moe_test.cc | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/onnxruntime/test/contrib_ops/moe_test.cc b/onnxruntime/test/contrib_ops/moe_test.cc index b7fc54c9a63c3..d68e94a98a954 100644 --- a/onnxruntime/test/contrib_ops/moe_test.cc +++ b/onnxruntime/test/contrib_ops/moe_test.cc @@ -2384,7 +2384,15 @@ TEST(MoETest, QMoETest_CPU_Int4_BlockWise_SwiGLU_Fp16) { } TEST(MoETest, QMoETest_CPU_Int8_BlockWise_SwiGLU) { - // 8-bit stays on the dequantize path by default (no fp32-activation QNBit kernel). + // MLAS has no fp32-activation QNBit kernel for 8 bits, so this runs with int8 activations by + // default and carries their quantization error; the tolerance matches the other int8 cases. + RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true, 8}, 0.05f); +} + +TEST(MoETest, QMoETest_CPU_Int8_BlockWise_SwiGLU_ForcedFp32FallsBackToDequantize) { + // Forcing fp32 activations keeps 8-bit experts on the dequantize path, which is exact enough + // for the tight tolerance. + ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_QMOE_CPU_QNBIT_GEMM", "fp32"}}}; RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true, 8}, 0.01f); } From 61173ff78c76a4d2072c2555e8b8dde297a8f830 Mon Sep 17 00:00:00 2001 From: Yuri Khrustalev Date: Wed, 16 Sep 2026 16:50:34 -0400 Subject: [PATCH 08/11] fix: document the 8-bit int8 default, SafeInt the shared-buffer size check, assert pre-packing in 8-bit QNBit tests Co-Authored-By: Claude Fable 5.1 --- docs/ContribOperators.md | 2 +- .../cpu/moe/moe_quantization_cpu.cc | 16 ++++++------ .../cpu/moe/moe_quantization_cpu.h | 4 +-- .../core/graph/contrib_ops/contrib_defs.cc | 7 +++--- onnxruntime/test/contrib_ops/moe_test.cc | 25 +++++++++++++++---- 5 files changed, 36 insertions(+), 18 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index baf781bb5744f..003c48980ad57 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -5676,7 +5676,7 @@ This version of the operator has been available since version 1 of the 'com.micr
accuracy_level : int
-
Minimum accuracy level of the expert GEMMs, with the MatMulNBits meaning: 0 (default) or 1 keeps fp32 activations; 4 allows int8 activations (int8 dot-product kernels) for block-wise 4/8-bit experts on CPU. Other values are treated as 0.
+
Minimum accuracy level of the expert GEMMs on CPU, with the MatMulNBits meaning. For block-wise 4-bit experts, 0 (default) or 1 keeps fp32 activations and 4 allows int8 activations (int8 dot-product kernels). Block-wise 8-bit experts have no fp32 kernel and use int8 activations at every level. Other values are treated as 0.
activation_alpha : float
Alpha parameter used in activation function.
activation_beta : float
diff --git a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc index 19c312fd2c1d6..49c62557d82bb 100644 --- a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc +++ b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc @@ -124,9 +124,10 @@ namespace contrib { constexpr const char* kUseMlasQ4GemmMoe = "ORT_USE_MLAS_Q4_GEMM_MOE"; // Overrides the node's accuracy_level for the MLAS QNBit GEMM (MatMulNBits kernel) path of -// block-wise 4/8-bit experts. Unset: the attribute decides. "fp32": fp32 activations only -// (accuracy_level 0; 8-bit has no fp32 variant and keeps the dequantize path). "int8": int8 -// activations (accuracy_level 4). "0": path disabled. Any other value is an error. +// block-wise 4/8-bit experts. Unset: the attribute decides for 4-bit; 8-bit has no fp32 kernel +// and uses int8 activations at every level. "fp32": fp32 activations only, so 8-bit keeps the +// dequantize path. "int8": int8 activations (accuracy_level 4). "0": path disabled. Any other +// value is an error. constexpr const char* kQMoEQNBitGemmEnv = "ORT_QMOE_CPU_QNBIT_GEMM"; // Tag appended to the prepacked shape buffer to mark the QNBit-packed layout (see PrePackQNBitExperts). @@ -849,7 +850,7 @@ Status QMoECPU::UseSharedPrePackedBuffers(std::vector& prepa ORT_RETURN_IF_ERROR(InitQNBitPacked(packed, eligibility, num_experts, rows, cols, Info().GetAllocator(OrtMemType::OrtMemTypeDefault))); } - ORT_RETURN_IF_NOT(prepacked_buffer_sizes[0] == packed.packed_size_per_expert * static_cast(num_experts), + ORT_RETURN_IF_NOT(prepacked_buffer_sizes[0] == SafeInt(packed.packed_size_per_expert) * static_cast(num_experts), "QMoE prepacked QNBit buffer size does not match the expert shape."); packed.packed = std::move(prepacked_buffers[0]); used_shared_buffers = true; @@ -940,9 +941,10 @@ QMoECPU::QMoECPU(const OpKernelInfo& op_kernel_info) // QNBit GEMM path policy. The MatMulNBits kernels consume the block-wise QMoE encoding directly // (no re-quantization), have NEON/AVX2/AVX512 backends, and thread a single GEMM over N, so they // are preferred for block-wise 4/8-bit experts. Int8 activations are lossier than the - // dequantize+SGEMM path they replace, so as with MatMulNBits they need accuracy_level=4: the - // default is fp32 activations where MLAS has that kernel (4-bit) and int8 where it does not - // (8-bit, which has no fp32 variant). The environment variable overrides the attribute. + // dequantize+SGEMM path they replace, so for 4-bit they need accuracy_level=4 as in MatMulNBits + // and the default is fp32 activations. 8-bit has no fp32 kernel; unlike MatMulNBits (which + // dequantizes at level 0) it uses int8 activations at every level, since that is the only way + // onto these kernels. The environment variable overrides the attribute. accuracy_level_ = op_kernel_info.GetAttrOrDefault("accuracy_level", 0); bool allow_int8_compute = (accuracy_level_ == 4); bool allow_fp32_compute = true; diff --git a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h index d0a43d0064da0..69f1e8fa2dbfe 100644 --- a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h +++ b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h @@ -61,8 +61,8 @@ class QMoECPU final : public OpKernel, public MoEBaseCPU { void ApplyActivationVectorized(float* data, int64_t size) const; // Expert weights pre-packed for the MLAS QNBit GEMM kernels (the MatMulNBits kernels). Used for - // block-wise 4/8-bit experts without zero points; each expert's [rows, cols] matrix is packed - // independently and stored back to back. + // block-wise 4/8-bit experts with optional constant zero points; each expert's [rows, cols] + // matrix is packed independently and stored back to back. struct QNBitPackedExperts { IAllocatorUniquePtr packed; // num_experts * packed_size_per_expert bytes size_t packed_size_per_expert{0}; diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 76e426b2f5b98..1731e26bab239 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -1506,9 +1506,10 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "Beta parameter used in activation function.", AttributeProto::FLOAT, 0.0f) .Attr("accuracy_level", - "Minimum accuracy level of the expert GEMMs, with the MatMulNBits meaning: 0 (default) or 1 keeps " - "fp32 activations; 4 allows int8 activations (int8 dot-product kernels) for block-wise 4/8-bit " - "experts on CPU. Other values are treated as 0.", + "Minimum accuracy level of the expert GEMMs on CPU, with the MatMulNBits meaning. For block-wise 4-bit " + "experts, 0 (default) or 1 keeps fp32 activations and 4 allows int8 activations (int8 dot-product kernels). " + "Block-wise 8-bit experts have no fp32 kernel and use int8 activations at every level. " + "Other values are treated as 0.", AttributeProto::INT, static_cast(0)) .Attr("block_size", "Size of each quantization block along the K (input feature) dimension. " diff --git a/onnxruntime/test/contrib_ops/moe_test.cc b/onnxruntime/test/contrib_ops/moe_test.cc index d68e94a98a954..34421b20bb1ae 100644 --- a/onnxruntime/test/contrib_ops/moe_test.cc +++ b/onnxruntime/test/contrib_ops/moe_test.cc @@ -13,6 +13,7 @@ #include "test/common/cuda_op_test_utils.h" #include "test/providers/provider_test_utils.h" #include "test/util/include/scoped_env_vars.h" +#include "core/mlas/inc/mlas_qnbit.h" namespace onnxruntime { namespace test { @@ -2145,6 +2146,10 @@ struct QMoEBlockWiseCase { int bits{4}; bool with_zero_points{false}; int accuracy_level{0}; + // Expected number of pre-packed expert weights (0 or 2), or -1 to skip the check. The output + // check alone cannot tell the QNBit path from the dequantize fallback, so 8-bit cases (where the + // fallback never pre-packs) also assert this; the 4-bit fallback pre-packs too, so it cannot. + int expected_prepacked{-1}; }; template @@ -2331,7 +2336,12 @@ void RunQMoECpuBlockWiseSwiGLU(const QMoEBlockWiseCase& c, float tolerance, if (!share_prepacked_weights_across_sessions) { auto eps = cpu_ep(); - tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); + SessionOptions so; + size_t prepacked = 0, shared = 0; + tester.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps, {}, &prepacked, &shared); + if (c.expected_prepacked >= 0) { + ASSERT_EQ(prepacked, static_cast(c.expected_prepacked)); + } return; } @@ -2383,22 +2393,27 @@ TEST(MoETest, QMoETest_CPU_Int4_BlockWise_SwiGLU_Fp16) { RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true}, 0.02f); } +// 8-bit experts pre-pack only on the QNBit path, so the pre-pack count proves which path ran. +static int ExpectedInt8QNBitPrepacks(int block_size) { + return MlasIsQNBitGemmAvailable(8, static_cast(block_size), SQNBIT_CompInt8) ? 2 : 0; +} + TEST(MoETest, QMoETest_CPU_Int8_BlockWise_SwiGLU) { // MLAS has no fp32-activation QNBit kernel for 8 bits, so this runs with int8 activations by // default and carries their quantization error; the tolerance matches the other int8 cases. - RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true, 8}, 0.05f); + RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true, 8, false, 0, ExpectedInt8QNBitPrepacks(32)}, 0.05f); } TEST(MoETest, QMoETest_CPU_Int8_BlockWise_SwiGLU_ForcedFp32FallsBackToDequantize) { // Forcing fp32 activations keeps 8-bit experts on the dequantize path, which is exact enough - // for the tight tolerance. + // for the tight tolerance and never pre-packs. ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_QMOE_CPU_QNBIT_GEMM", "fp32"}}}; - RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true, 8}, 0.01f); + RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true, 8, false, 0, 0}, 0.01f); } TEST(MoETest, QMoETest_CPU_Int8_BlockWise_SwiGLU_Int8Activations) { ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_QMOE_CPU_QNBIT_GEMM", "int8"}}}; - RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true, 8}, 0.05f); + RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true, 8, false, 0, ExpectedInt8QNBitPrepacks(32)}, 0.05f); } TEST(MoETest, QMoETest_CPU_Int4_BlockWise_SwiGLU_Int8Activations) { From acebeebc8d36bb9f11cc0c35ed654e5a502b2ef9 Mon Sep 17 00:00:00 2001 From: Yuri Khrustalev Date: Wed, 16 Sep 2026 21:32:32 -0400 Subject: [PATCH 09/11] fix: prefer int8 QNBit kernels for accuracy_level=4 and keep the per-call workspace arena-backed accuracy_level=4 (or the int8 env override) now selects the int8 kernels ahead of fp32 when both are available, as MatMulNBits does; fp32 stays the fallback when int8 is missing and the explicit fp32 override still pins it. A regression test compares raw outputs so level 4 is proven to run the int8 kernels rather than merely meeting the looser tolerance. The QNBit scratch buffer no longer uses the reserve allocation path. Co-Authored-By: Claude Fable 5.1 --- .../cpu/moe/moe_quantization_cpu.cc | 18 ++++---- onnxruntime/test/contrib_ops/moe_test.cc | 41 ++++++++++++++++++- 2 files changed, 51 insertions(+), 8 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc index 49c62557d82bb..7186cd2ae0721 100644 --- a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc +++ b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc @@ -967,13 +967,16 @@ QMoECPU::QMoECPU(const OpKernelInfo& op_kernel_info) if (use_qnbit_gemm_ && (expert_weight_bits_ == 4 || expert_weight_bits_ == 8) && block_size_ > 0) { const size_t nbits = static_cast(expert_weight_bits_); const size_t blk = static_cast(block_size_); - // MLAS has no fp32-activation kernel for every bit width (8-bit is int8 only). Where fp32 is - // unavailable, int8 is the only way onto these kernels, so take it unless fp32 was forced. - const bool int8_is_only_option = !MlasIsQNBitGemmAvailable(nbits, blk, SQNBIT_CompFp32) && !fp32_compute_forced; - if (allow_fp32_compute && MlasIsQNBitGemmAvailable(nbits, blk, SQNBIT_CompFp32)) { - qnbit_compute_type_ = SQNBIT_CompFp32; - } else if ((allow_int8_compute || int8_is_only_option) && MlasIsQNBitGemmAvailable(nbits, blk, SQNBIT_CompInt8)) { + const bool fp32_available = MlasIsQNBitGemmAvailable(nbits, blk, SQNBIT_CompFp32); + const bool int8_available = MlasIsQNBitGemmAvailable(nbits, blk, SQNBIT_CompInt8); + // int8 is taken when asked for (accuracy_level 4 or the env override), like MatMulNBits, and + // also where MLAS has no fp32 kernel for the bit width (8-bit) unless fp32 was forced, since + // int8 is then the only way onto these kernels. fp32 is the fallback whenever it is allowed. + const bool want_int8 = allow_int8_compute || (!fp32_available && !fp32_compute_forced); + if (want_int8 && int8_available) { qnbit_compute_type_ = SQNBIT_CompInt8; + } else if (allow_fp32_compute && fp32_available) { + qnbit_compute_type_ = SQNBIT_CompFp32; } else { use_qnbit_gemm_ = false; } @@ -1606,7 +1609,8 @@ Status QMoECPU::ComputeCommon(OpKernelContext* context, const ComputeInputs& IAllocatorUniquePtr qnbit_workspace_ptr; std::byte* qnbit_workspace = nullptr; if (qnbit_workspace_per_thread > 0) { - qnbit_workspace_ptr = IAllocator::MakeUniquePtr(allocator, static_cast(num_expert_threads) * qnbit_workspace_per_thread, true); + // Arena-backed (no reserve): this is per-call scratch that should be recycled between runs. + qnbit_workspace_ptr = IAllocator::MakeUniquePtr(allocator, static_cast(num_expert_threads) * qnbit_workspace_per_thread); qnbit_workspace = qnbit_workspace_ptr.get(); } // fp32 [E, N, K/block] scales for the kernels: the scales input itself for T == float, the copy diff --git a/onnxruntime/test/contrib_ops/moe_test.cc b/onnxruntime/test/contrib_ops/moe_test.cc index 34421b20bb1ae..3458f1157956d 100644 --- a/onnxruntime/test/contrib_ops/moe_test.cc +++ b/onnxruntime/test/contrib_ops/moe_test.cc @@ -2152,9 +2152,12 @@ struct QMoEBlockWiseCase { int expected_prepacked{-1}; }; +// `output_out`, when set, receives the raw fp32 output of the run so callers can compare the +// results of different kernel policies bitwise (T == float only). template void RunQMoECpuBlockWiseSwiGLU(const QMoEBlockWiseCase& c, float tolerance, - bool share_prepacked_weights_across_sessions = false) { + bool share_prepacked_weights_across_sessions = false, + std::vector* output_out = nullptr) { const int num_rows = c.num_rows, num_experts = c.num_experts, hidden_size = c.hidden_size; const int inter_size = c.inter_size, block_size = c.block_size, top_k = c.top_k, bits = c.bits; const int fc1_rows = 2 * inter_size; // interleaved [gate, up] rows @@ -2342,6 +2345,13 @@ void RunQMoECpuBlockWiseSwiGLU(const QMoEBlockWiseCase& c, float tolerance, if (c.expected_prepacked >= 0) { ASSERT_EQ(prepacked, static_cast(c.expected_prepacked)); } + if (output_out != nullptr) { + static_assert(std::is_same_v, "raw output capture is only wired for float"); + const auto fetches = tester.GetFetches(); + ASSERT_EQ(fetches.size(), static_cast(1)); + const Tensor& out = fetches[0].Get(); + output_out->assign(out.Data(), out.Data() + out.Shape().Size()); + } return; } @@ -2421,6 +2431,35 @@ TEST(MoETest, QMoETest_CPU_Int4_BlockWise_SwiGLU_Int8Activations) { RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true, 4, false, 4}, 0.05f); } +TEST(MoETest, QMoETest_CPU_Int4_BlockWise_SwiGLU_AccuracyLevel4SelectsInt8) { + // accuracy_level=4 must pick the int8 kernels even where the fp32 ones are also available. The + // output check alone cannot tell the two apart (both meet the int8 tolerance), so compare raw + // outputs: level 4 must be bitwise identical to the forced-int8 run (same kernels, each output + // column owned by one thread) and differ from the forced-fp32 run (int8 activation quantization + // perturbs every output). + if (!MlasIsQNBitGemmAvailable(4, 32, SQNBIT_CompInt8) || !MlasIsQNBitGemmAvailable(4, 32, SQNBIT_CompFp32)) { + GTEST_SKIP() << "platform lacks one of the 4-bit QNBit compute types"; + } + const QMoEBlockWiseCase base{5, 4, 128, 64, 32, 2, true}; + std::vector level4, forced_fp32, forced_int8; + { + QMoEBlockWiseCase c = base; + c.accuracy_level = 4; + RunQMoECpuBlockWiseSwiGLU(c, 0.05f, false, &level4); + } + { + ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_QMOE_CPU_QNBIT_GEMM", "fp32"}}}; + RunQMoECpuBlockWiseSwiGLU(base, 0.01f, false, &forced_fp32); + } + { + ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_QMOE_CPU_QNBIT_GEMM", "int8"}}}; + RunQMoECpuBlockWiseSwiGLU(base, 0.05f, false, &forced_int8); + } + ASSERT_FALSE(level4.empty()); + EXPECT_EQ(level4, forced_int8); + EXPECT_NE(level4, forced_fp32); +} + TEST(MoETest, QMoETest_CPU_Int4_BlockWise_SwiGLU_ZeroPoints) { RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true, 4, true}, 0.01f); } From 0d244bcb25089b2ba230d70c819f94c0de50c296 Mon Sep 17 00:00:00 2001 From: Yuri Khrustalev Date: Wed, 16 Sep 2026 21:33:36 -0400 Subject: [PATCH 10/11] test: capture QMoE outputs for both element types in the block-wise helper Co-Authored-By: Claude Fable 5.1 --- onnxruntime/test/contrib_ops/moe_test.cc | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/onnxruntime/test/contrib_ops/moe_test.cc b/onnxruntime/test/contrib_ops/moe_test.cc index 3458f1157956d..3cdaec36e575e 100644 --- a/onnxruntime/test/contrib_ops/moe_test.cc +++ b/onnxruntime/test/contrib_ops/moe_test.cc @@ -2152,8 +2152,8 @@ struct QMoEBlockWiseCase { int expected_prepacked{-1}; }; -// `output_out`, when set, receives the raw fp32 output of the run so callers can compare the -// results of different kernel policies bitwise (T == float only). +// `output_out`, when set, receives the run's output as fp32 so callers can compare the results of +// different kernel policies bitwise. template void RunQMoECpuBlockWiseSwiGLU(const QMoEBlockWiseCase& c, float tolerance, bool share_prepacked_weights_across_sessions = false, @@ -2346,11 +2346,18 @@ void RunQMoECpuBlockWiseSwiGLU(const QMoEBlockWiseCase& c, float tolerance, ASSERT_EQ(prepacked, static_cast(c.expected_prepacked)); } if (output_out != nullptr) { - static_assert(std::is_same_v, "raw output capture is only wired for float"); const auto fetches = tester.GetFetches(); ASSERT_EQ(fetches.size(), static_cast(1)); const Tensor& out = fetches[0].Get(); - output_out->assign(out.Data(), out.Data() + out.Shape().Size()); + const T* data = out.Data(); + output_out->resize(static_cast(out.Shape().Size())); + for (size_t i = 0; i < output_out->size(); ++i) { + if constexpr (is_fp16) { + (*output_out)[i] = data[i].ToFloat(); + } else { + (*output_out)[i] = data[i]; + } + } } return; } From bc5f446b427b7cb886234a88bae97ecce933a728 Mon Sep 17 00:00:00 2001 From: Yuri Khrustalev Date: Wed, 16 Sep 2026 21:43:44 -0400 Subject: [PATCH 11/11] test: pin the remaining QMoE QNBit policy cells Levels other than 4 and the fp32 override must match the 4-bit default bitwise; the 0 override keeps the previous paths (never pre-packing 8-bit experts); the shared-prepack case asserts both weights pre-packed instead of returning early. Co-Authored-By: Claude Fable 5.1 --- onnxruntime/test/contrib_ops/moe_test.cc | 42 ++++++++++++++++++++++-- 1 file changed, 39 insertions(+), 3 deletions(-) diff --git a/onnxruntime/test/contrib_ops/moe_test.cc b/onnxruntime/test/contrib_ops/moe_test.cc index 3cdaec36e575e..08cd6bd463980 100644 --- a/onnxruntime/test/contrib_ops/moe_test.cc +++ b/onnxruntime/test/contrib_ops/moe_test.cc @@ -2381,9 +2381,9 @@ void RunQMoECpuBlockWiseSwiGLU(const QMoEBlockWiseCase& c, float tolerance, ASSERT_EQ(shared, static_cast(0)); } ASSERT_EQ(prepacked_session_1, tester.GetNumPrePackedWeightsShared()); - if (prepacked_session_1 == 0) { - return; // platform without the QNBit kernels: nothing was prepacked, nothing to share - } + // Both expert weights pre-pack on every platform (the legacy 4-bit path pre-packs too), so a + // count of 2 proves the sharing round trip below exercises real buffers. + ASSERT_EQ(prepacked_session_1, static_cast(2)); { auto eps = cpu_ep(); tester.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps, {}, &prepacked_session_2, &shared); @@ -2467,6 +2467,42 @@ TEST(MoETest, QMoETest_CPU_Int4_BlockWise_SwiGLU_AccuracyLevel4SelectsInt8) { EXPECT_NE(level4, forced_fp32); } +TEST(MoETest, QMoETest_CPU_Int4_BlockWise_SwiGLU_PolicyMatrix) { + // Cells of the documented policy that the tolerance checks cannot tell apart: runs that must + // select the same kernels produce bitwise-identical outputs. accuracy_level values other than 4 + // behave as 0, and the fp32 override is the 4-bit default. + if (!MlasIsQNBitGemmAvailable(4, 32, SQNBIT_CompFp32)) { + GTEST_SKIP() << "platform lacks the 4-bit fp32 QNBit kernel"; + } + const QMoEBlockWiseCase base{5, 4, 128, 64, 32, 2, true}; + std::vector by_default, level1, forced_fp32; + RunQMoECpuBlockWiseSwiGLU(base, 0.01f, false, &by_default); + { + QMoEBlockWiseCase c = base; + c.accuracy_level = 1; + RunQMoECpuBlockWiseSwiGLU(c, 0.01f, false, &level1); + } + { + ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_QMOE_CPU_QNBIT_GEMM", "fp32"}}}; + RunQMoECpuBlockWiseSwiGLU(base, 0.01f, false, &forced_fp32); + } + ASSERT_FALSE(by_default.empty()); + EXPECT_EQ(level1, by_default); + EXPECT_EQ(forced_fp32, by_default); +} + +TEST(MoETest, QMoETest_CPU_Int4_BlockWise_SwiGLU_DisabledByEnv) { + // "0" keeps the previous paths; the output still matches the exact reference. + ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_QMOE_CPU_QNBIT_GEMM", "0"}}}; + RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true}, 0.01f); +} + +TEST(MoETest, QMoETest_CPU_Int8_BlockWise_SwiGLU_DisabledByEnv) { + // "0" keeps the previous paths, which never pre-pack 8-bit experts. + ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_QMOE_CPU_QNBIT_GEMM", "0"}}}; + RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true, 8, false, 0, 0}, 0.01f); +} + TEST(MoETest, QMoETest_CPU_Int4_BlockWise_SwiGLU_ZeroPoints) { RunQMoECpuBlockWiseSwiGLU({5, 4, 128, 64, 32, 2, true, 4, true}, 0.01f); }