[CPU] Add opt-in MatMulNBits FP32 throughput mode

[Mustapha Jaber (mustjab)]

Description

Adds an opt-in, session-wide CPU MatMulNBits throughput mode.

Set mlas.qnbit.force_fp32=1 to make eligible x86/x64 float-input, 4-bit, block-size-32, accuracy-level-4 nodes use the existing CompFp32 path instead of CompInt8 for the entire session.

• One compute type and one packed-weight layout are used for every batch shape.
• Dynamic scales and zero points are supported.
• Cross-session prepacked-weight sharing is supported.
• The option is disabled by default, so existing behavior is unchanged.

Motivation and Context

For Marian translation, CompInt8 provides strong single-sentence latency, but CompFp32 scales much better at large batches. Shape-dependent switching improved throughput but changed the numerical path between serial and batched execution, increasing token mismatches. This version intentionally selects CompFp32 once at session construction so all batch sizes use the same numerical path.

The intended deployment model is two sessions:

• default level-4 session for latency-sensitive requests;
• mlas.qnbit.force_fp32=1 session for throughput/batched requests.

End-to-end validation

Rebuilt ONNX Runtime GenAI from landed main commit 53f2379f (PR #2368), including Extensions pin 60687b12, against this ORT build.

Five alternating 97-sentence baseline/throughput process pairs used the same GenAI 0.16 binary, ORT DLL, graphs, corpus, and seed:

Batch	Baseline E2E	FP32 mode E2E	E2E gain	Baseline generation	FP32 mode generation	Generation gain	FP32 mismatches
1	3.691 s	6.308 s	0.59x	2.213 s	3.537 s	0.63x	0/97
5	2.495 s	3.510 s	0.71x	1.322 s	2.208 s	0.60x	0/97
10	2.391 s	2.824 s	0.85x	1.188 s	1.607 s	0.74x	0/97
20	1.949 s	1.828 s	1.07x	0.957 s	0.996 s	0.96x	0/97
32	1.925 s	1.654 s	1.16x	0.931 s	0.875 s	1.06x	0/97

Batch-32 scaling improves from 1.92x to 3.81x E2E and from 2.38x to 4.04x generation. Serial/batch parity improves from the existing 2/97 level-4 mismatches to 0/97 at every batch size in FP32 mode.

The tradeoff is explicit: FP32 mode is about 1.7x slower at batch 1, so it is intended for a separate throughput-oriented session rather than as a default.

Isolated encoder-session private memory decreased from approximately 89.3 MB to 60.6 MB because only the CompFp32 packed layout is retained.

Validation

• onnxruntime_provider_test.exe --gtest_filter=MatMulNBits.*
  • 26 passed
  • 3 expected architecture-specific skips
• Coverage includes small and batched FP32-mode execution, dynamic scales, dynamic zero points, asymmetric zero points, bias, and cross-session prepacked-weight sharing.
• GenAI library and Python extension built successfully against the updated ORT. The aggregate GenAI build later stopped linking unit_tests.exe on four existing logging symbols.
• git diff --check passes.

[Copilot]

Pull request overview

This PR adds an opt-in “adaptive compute” mode for CPU MatMulNBits (accuracy-level-4) that dynamically selects between existing CompInt8 and CompFp32 MLAS paths based on total row count (M * batch_count), aiming to improve throughput on large square W4/block-32 workloads while preserving current behavior by default.

Changes:

• Added a new session config key mlas.enable_qnbit_adaptive_compute and plumbed it into the CPU MatMulNBits kernel to retain both prepacked layouts and switch compute type at runtime.
• Extended prepacked-weight sharing to include the additional CompFp32 packed-B buffer when adaptive mode is enabled.
• Added/updated unit tests to cover the adaptive boundary behavior, dynamic scales/zero-points metadata, and cross-session prepack sharing with adaptive enabled.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
onnxruntime/test/contrib_ops/matmul_nbits_prepack_sharing_test_util.h	Extends the sharing-test helper signature to allow enabling adaptive mode.
onnxruntime/test/contrib_ops/matmul_nbits_prepack_sharing_test_util.cc	Enables the adaptive session option in sharing tests when requested.
onnxruntime/test/contrib_ops/matmul_4bits_test.cc	Adds adaptive compute tests (boundary + dynamic metadata) and an adaptive prepack-sharing test; allows scales to be non-initializers.
onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc	Implements adaptive mode gating, dual-layout prepack + sharing, and runtime compute-type selection.
include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h	Adds the public session option key and documentation for adaptive compute.

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

Suppressed comments (4)

onnxruntime/test/contrib_ops/matmul_4bits_test.cc:507

• The skip message mentions "Adaptive QNBit compute", but this test is validating the mlas.qnbit.force_fp32 session option (dynamic scales/zero points). Updating the message makes it clearer why the test is skipped on non-x86/x64 builds.

#if !defined(MLAS_TARGET_AMD64_IX86)
  GTEST_SKIP() << "Adaptive QNBit compute is currently supported only on x86/x64.";
#else

onnxruntime/test/contrib_ops/matmul_4bits_test.cc:746

• The skip message mentions "Adaptive QNBit compute", but this test is specifically about forced CompFp32 + prepacked-weight sharing via mlas.qnbit.force_fp32. Adjusting the message avoids implying adaptive/shape-dependent selection.

#if !defined(MLAS_TARGET_AMD64_IX86)
  GTEST_SKIP() << "Adaptive QNBit compute is currently supported only on x86/x64.";
#else

onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc:124

• When both mlas.qnbit.force_fp32=1 and mlas.use_lut_gemm=1 are set, prefer_lut_gemm_ can still be enabled and the kernel will take the LUT prepack/compute path, effectively bypassing the forced-CompFp32 intent for eligible accuracy-level-4 nodes. This makes the new session option behavior dependent on an unrelated LUT setting and can defeat the “force fp32” guarantee.

        compute_type_{info.GetConfigOptions().GetConfigEntry(kOrtSessionOptionsMlasQNBitForceFp32) == "1" &&
                              kQNBitForceFp32Supported && std::is_same_v<T1, float> &&
                              nbits_ == 4 && block_size_ == 32 &&
                              info.GetAttr<int64_t>("accuracy_level") == static_cast<int64_t>(Level4) &&
                              MlasIsQNBitGemmAvailable(nbits_, block_size_, SQNBIT_CompFp32)

onnxruntime/test/contrib_ops/matmul_4bits_test.cc:488

• The skip message mentions "Adaptive QNBit compute", but this test is exercising the session-wide mlas.qnbit.force_fp32 mode. The message should describe the actual feature being skipped to avoid confusion when diagnosing CI skips.

This issue also appears in the following locations of the same file:

• line 505
• line 744

#if !defined(MLAS_TARGET_AMD64_IX86)
  GTEST_SKIP() << "Adaptive QNBit compute is currently supported only on x86/x64.";
#else

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

Suppressed comments (1)

onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc:120

• This disables LUT GEMM for every float MatMulNBits node whenever the session option is set, even though the option only applies to eligible W4/block-32/level-4 nodes. LUT currently supports W2 nodes, so a mixed throughput session will turn an explicitly requested LUT node into the unpacked dequantization fallback. Gate this suppression with the same eligibility predicate used for compute_type_ (or leave LUT enabled for ineligible nodes).

                         info.GetConfigOptions().GetConfigEntry(kOrtSessionOptionsMlasQNBitForceFp32) != "1" &&