diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 7659698814171..003c48980ad57 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 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 fad6c5aecf512..7186cd2ae0721 100644 --- a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc +++ b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc @@ -15,6 +15,8 @@ #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" @@ -121,6 +123,31 @@ 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 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). +constexpr int64_t kQNBitPackedLayoutTag = 0x514E4249; // 'QNBI' + +// Prepacked shape buffer: rank, dims, then `trailer` (empty for the legacy layouts; the QNBit layout +// 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(); + 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, @@ -610,6 +637,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. @@ -682,19 +716,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(); } @@ -721,20 +743,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_) { @@ -810,6 +819,45 @@ Status QMoECPU::UseSharedPrePackedBuffers(std::vector& prepa return Status::OK(); }; + // 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] >= 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."); + // 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_); + 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, eligibility, num_experts, rows, cols, + Info().GetAllocator(OrtMemType::OrtMemTypeDefault))); + } + 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; + return Status::OK(); + } + } + if (expert_weight_bits_ == 2) { if ((input_idx == 2 || input_idx == 5) && !prepacked_buffers.empty()) { if (input_idx == 2) { @@ -889,6 +937,241 @@ 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. Int8 activations are lossier than the + // 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; + 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; + } 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_); + 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; + } + } else { + use_qnbit_gemm_ = false; + } +} + +template +bool QMoECPU::QNBitGemmEligible(int input_idx, int64_t num_experts, int64_t rows, int64_t cols, + QNBitEligibility& out) const { + out = QNBitEligibility{}; + if (!use_qnbit_gemm_ || (input_idx != 2 && input_idx != 5)) { + return false; + } + 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 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(); + 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_), 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; + } + + out.scales = scales_tensor; + out.zero_points = zp_tensor; + out.packed_size_per_expert = packed_size; + return true; +} + +template +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_, 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(eligibility.scales->template Data(), packed.scales_fp32.get(), scales_count); + } else { + ORT_UNUSED_PARAMETER(num_experts); + } + 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; + + 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, 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 = eligibility.scales->template Data(); + } + 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; + + 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 = 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); + + 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). + std::unique_ptr pack_tp; + 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 = 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); + } + 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, + 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, + packed.has_zero_point, expert_zp, 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); + + // The layout tag lets UseSharedPrePackedBuffers tell this layout apart from the legacy + // 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); + } + + is_packed = true; + return Status::OK(); } template @@ -896,10 +1179,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 +1194,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 +1417,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 +1452,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 +1470,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 +1588,70 @@ 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) { + 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, 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) { + 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, 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; + if (qnbit_workspace_per_thread > 0) { + // 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 + // 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, + 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; + + 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 = zero_points; + 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 +1684,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 +1772,22 @@ 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_, 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; + } + if (can_use_fc1_lut_gemm && TryRunLutGemm(A1, C1, fc1_weights_data, fc1_direct_lut_cache_ptr != nullptr @@ -1658,6 +2032,23 @@ 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_, 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; + } + 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..69f1e8fa2dbfe 100644 --- a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h +++ b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.h @@ -60,6 +60,31 @@ 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 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}; + 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, + 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); + int64_t expert_weight_bits_; int64_t block_size_; bool use_mlas_q4_gemm_{false}; @@ -75,6 +100,13 @@ 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_; }; } // namespace contrib diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index ce531a56cbf7b..1731e26bab239 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -1505,6 +1505,12 @@ 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 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. " "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 1cdd22f789816..08cd6bd463980 100644 --- a/onnxruntime/test/contrib_ops/moe_test.cc +++ b/onnxruntime/test/contrib_ops/moe_test.cc @@ -1,11 +1,19 @@ // 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" +#include "core/mlas/inc/mlas_qnbit.h" namespace onnxruntime { namespace test { @@ -2120,6 +2128,398 @@ TEST(MoETest, QMoETest_CPU_Int4_MLAS) { #endif } +// 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; + int num_experts; + int hidden_size; + int inter_size; + int block_size; + int top_k; + bool with_bias; + 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}; +}; + +// `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, + 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 + 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; + const int code_max = (1 << bits) - 1; + 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); + // 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 + // 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, + 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); + 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, fc1_zp_codes, 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, fc2_zp_codes, 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); + 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) { + 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 (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 { + 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(); + 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)); + } + if (output_out != nullptr) { + const auto fetches = tester.GetFetches(); + ASSERT_EQ(fetches.size(), static_cast(1)); + const Tensor& out = fetches[0].Get(); + 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; + } + + // 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()); + // 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); + 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); +} + +// 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, 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 and never pre-packs. + ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_QMOE_CPU_QNBIT_GEMM", "fp32"}}}; + 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, false, 0, ExpectedInt8QNBitPrepacks(32)}, 0.05f); +} + +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_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_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); +} + +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, 8, 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