Optimizations and fixes in QMoE CPU kernel #25642
Conversation
| // Generate test weights for symmetric quantization (zero point is 0) | ||
| std::vector<uint8_t> fc1_experts_weights(num_experts * fc1_weight_size_per_expert, 0x12); // 1,2 -> small positive weights | ||
| std::vector<uint8_t> fc2_experts_weights(num_experts * fc2_weight_size_per_expert, 0xFF); // -1,0 -> small mixed weights | ||
| // Generate test weights for symmetric quantization (zero point is 8 for 4-bit) |
There was a problem hiding this comment.
zero point is 0 for symmetric quantization?
There was a problem hiding this comment.
Yes it is 0.
Original range is -8 to 7, it is zero because original zero maps at 8
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| ort_dtype_quant_bits_tolerance_map = { | ||
| "FP32:0": (5e-3, 1e-3), | ||
| "FP16:0": (5e-2, 1e-3), | ||
| "FP16:4": (2.0, 8e-3), # Improved tolerance with symmetric quantization | ||
| "FP16:8": (1.5, 8e-3), # Improved tolerance with symmetric quantization | ||
| "FP16:4": (8.0, 0.15), # 4-bit quantization error tolerance - improved with bug fixes |
There was a problem hiding this comment.
Need lower the thresholds?
| onnx_dtype=self.onnx_dtype, | ||
| fc1_experts_weights=self.moe_experts_weight1, | ||
| fc2_experts_weights=self.moe_experts_weight2, | ||
| # Biases are not used in QMoE |
There was a problem hiding this comment.
This comment is not True. Kernel support biases, but this test did not.
| @@ -959,7 +1186,7 @@ def test_phi3_qmoe_parity_cpu(self, batch_size, sequence_length, quant_bits, use | |||
|
|
|||
| # Skip if the session creation failed | |||
| if phi3_moe.ort_sess is None: | |||
| self.skipTest("Failed to create ONNX Runtime session - CPU MoE operator not available") | |||
| self.skipTest("Failed to create ONNX Runtime session") | |||
There was a problem hiding this comment.
Do not skip test for this case. Instead throw an exception.
| // Use arena allocator for better memory management and reduced fragmentation | ||
| // This is especially beneficial for repeated kernel invocations | ||
| AllocatorPtr arena_allocator; | ||
| if (context->GetUseDeterministicCompute()) { | ||
| // For deterministic compute, use the standard temp allocator | ||
| arena_allocator = allocator; | ||
| } else { | ||
| // Try to get arena allocator for better performance | ||
| ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&arena_allocator)); | ||
| } |
There was a problem hiding this comment.
I do not think using arena will make it difference in deterministic.
| router_probs_float_ptr = unified_conversion_buffer.get() + input_size; | ||
|
|
||
| // Set up smart pointers with custom deleters to avoid double-free | ||
| input_float = IAllocatorUniquePtr<float>(input_float_ptr, [](float*) {}); |
There was a problem hiding this comment.
unified_conversion_buffer is already smart pointer. Why to wrap its memory into another smart pointer here?
|
|
||
| // Set up smart pointers with custom deleters to avoid double-free | ||
| input_float = IAllocatorUniquePtr<float>(input_float_ptr, [](float*) {}); | ||
| router_probs_float = IAllocatorUniquePtr<float>(router_probs_float_ptr, [](float*) {}); |
There was a problem hiding this comment.
Same here.
Please make sure this code path be tested.
| const size_t total_conversion_size = input_size + router_probs_size; | ||
|
|
||
| // Single allocation for input and router_probs conversion | ||
| auto unified_conversion_buffer = IAllocator::MakeUniquePtr<float>(arena_allocator, total_conversion_size); |
There was a problem hiding this comment.
This smarter pointer will be deleted after if scope. I think it needs to be moved to outer scope since its life cycle shall be kept till computation is done.
| if (is_decoding_scenario) { | ||
| // Decoding scenario: partition work by tokens * experts for better parallelization | ||
| // This allows multiple threads to work on different experts for the same token | ||
| double cost_per_token_expert = static_cast<double>(moe_params.hidden_size * moe_params.inter_size * 0.6); // Cost per token-expert pair |
There was a problem hiding this comment.
It should be 2 for 2 GEMM operations
| expert_idx, routing_weight, moe_params, is_swiglu, | ||
| dequant_fc1_weights, dequant_fc2_weights, | ||
| fc1_bias_float_ptr, fc2_bias_float_ptr, fc1_output_size); // Atomically accumulate results to avoid race conditions | ||
| float* token_result = output_float_ptr + token_idx * moe_params.hidden_size; |
There was a problem hiding this comment.
Data race bug.
For instance, Thread 1 might handle (token_idx=0, expert_idx=0) while Thread 2 handles (token_idx=0, expert_idx=1). Both threads will eventually attempt to write to the same memory location (output_float_ptr + token_idx * hidden_size) without any synchronization.
|
PrepackAndDequantizeWeights shall be called the standard prepack override funciton like onnxruntime/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc Lines 123 to 128 in 562760a It has ability to remove the existing weights (quantized weights) by setting is_packed=true to reduce memory. |
|
Performance Improvement Suggestion: Batch GEMM Operations Instead of Iterative Calls Problem: In QuantizedMoEImpl, the code loops through tokens and then through experts, performing a small matrix multiplication for each. This is highly inefficient. Recommendation: Refactor the logic to use a batched GEMM approach. This involves a "token shuffling" strategy: First, iterate through all tokens to determine which expert each token is routed to. |
apsonawane
force-pushed
the
asonawane/qmoe-fix
branch
from
15fbac0 to
a0251ea
Compare
This reverts commit a0251ea.
I wonder if we even need to "clump" tokens per expert but rather pass through pointers to the token rows per expert and let MLAS use the row pointers to gather data into its compute registers. This way we can avoid the "clump" and "unshuffle" operations altogether ? |
"clump" and "unshuffle" is just logically. In implementation, we can pass pointers. That's supported by the MLAS Batch GEMM API. |
This pull request focuses on optimizing the SwiGLU activation implementation in the MoE (Mixture of Experts) module and updating corresponding test cases to reflect these changes. The most important changes include performance improvements in the SwiGLU activation, removal of support for non-interleaved formats, and updates to test cases for quantized weights to ensure consistency.
SwiGLU Activation Optimizations:
ApplySwiGLUActivationby introducing vectorized computations, faster clamping, and efficient memory access. Removed support for non-interleaved formats, replacing it with aORT_NOT_IMPLEMENTEDerror. (onnxruntime/contrib_ops/cpu/moe/moe_utils.cc, onnxruntime/contrib_ops/cpu/moe/moe_utils.ccL29-R64)