onnxruntime

c267f8e2 - Optimize MatMulNBits 2-bit + float zero_point CPU dequantization with multi-threaded kernel (#28589)

Optimize MatMulNBits 2-bit + float zero_point CPU dequantization with multi-threaded kernel (#28589) ### Description Replace the naive single-threaded scalar loop for 2-bit dequantization with float/MLFloat16 zero points with a multi-threaded kernel (`DequantizeBlockwise2Bits`) that: - **Parallelizes via `TrySimpleParallelFor`** — distributes work across all intra-op threads (previously single-threaded) - **Processes 16 elements per iteration** — one `uint32_t` = 16 packed 2-bit values, reducing per-element overhead - **Hoists scale/zp lookups** — all 16 elements share a block, so scale and zero_point are loaded once per batch Follows the same threading pattern as the existing 4-bit `DequantizeBlockwise` path for consistency. **Files changed:** - `matmul_nbits_impl.h` — declare `DequantizeBlockwise2Bits` - `matmul_nbits_impl.cc` — implement `Dequantize2BitsKernel` + `DequantizeBlockwise2Bits` with instantiations for `<float,float>` and `<float,MLFloat16>` - `matmul_nbits.cc` — replace naive loops in both `MatMulNBits<float>` and `MatMulNBits<MLFloat16>` `ComputeBUnpacked` ### Motivation and Context The `bits=2` + float zero_point path (added in #28354) was flagged with `// !!!!!!!!!!!!!! naive implementation, need to be optimized !!!!!!!!!!!!!!`. It ran ~20× slower than the `bits=4` MLAS path because it was a tight scalar `for n × for k` loop with no threading — the entire N×K dequantization ran on a single core before calling `MlasGemmBatch`. With 8 intra-op threads this should recover most of that gap. ### Benchmark Results Tested on a 96-core x86_64 Linux machine, ORT 1.27.0 CPU Release build, using typical LLM matrix shapes with `block_size=128` and float zero points. #### Multi-thread speedup (2-bit dequantization, 1 thread → 8 threads) | Shape (M×K×N) | 1-thread (ms) | 8-thread (ms) | Speedup | |---|---|---|---| | 1×4096×4096 | 41.0 | 8.5 | **4.84×** | | 32×4096×4096 | 47.9 | 8.8 | **5.46×** | | 1×4096×11008 | 120.7 | 24.2 | **4.99×** | | 32×4096×11008 | 146.8 | 28.2 | **5.21×** | | 1×11008×4096 | 119.2 | 24.5 | **4.87×** | | 32×11008×4096 | 154.4 | 28.2 | **5.47×** | | 1×1024×1024 | 1.18 | 0.16 | **7.61×** | #### 2-bit vs 4-bit comparison (ratio = 2-bit / 4-bit; <1.0 means 2-bit is faster) | Shape (M×K×N) | Threads | 4-bit (ms) | 2-bit (ms) | Ratio | |---|---|---|---|---| | 1×4096×4096 | 1 | 52.0 | 41.0 | **0.79×** | | 1×4096×4096 | 8 | 9.4 | 8.5 | **0.90×** | | 1×4096×11008 | 1 | 141.6 | 120.7 | **0.85×** | | 1×4096×11008 | 8 | 26.8 | 24.2 | **0.90×** | | 1×11008×4096 | 1 | 141.2 | 119.2 | **0.84×** | | 1×11008×4096 | 8 | 26.6 | 24.5 | **0.92×** | | 32×4096×4096 | 1 | 56.1 | 47.9 | **0.85×** | | 32×4096×4096 | 8 | 9.6 | 8.8 | **0.92×** | | 1×1024×1024 | 1 | 1.66 | 1.18 | **0.71×** | **Key findings:** - Multi-threading delivers **4.8–7.6× speedup** with 8 threads across all LLM shapes - 2-bit is now **10–30% faster** than 4-bit (ratio 0.71–0.93×), due to fewer bytes read from memory - The original ~20× regression (issue #28552) is fully resolved --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: tianleiwu <30328909+tianleiwu@users.noreply.github.com> Co-authored-by: Tianlei Wu <tlwu@microsoft.com>