[MLAS] Extend the fp16 direct C output to the 4 bit CompFp32 path

[blazingphoenix7]

Description

#29791 and #29864 gave the CompInt8 compute path a direct fp16 C output: each worker runs the kernel into a small per thread fp32 scratch and converts its strip, so a fp16 MatMulNBits never allocates a full M x N fp32 copy of the result. The CompFp32 path still did: a 4 bit fp16 model at accuracy_level 0, 1, 2 or 3, which includes the schema default of 0, computed into a full fp32 output buffer and then ran a separate conversion pass over it.

This closes that hole for the 4 bit CompFp32 path. SQ4BitGemm_CompFp32 gains the same strip conversion the CompInt8 wrappers use, in both of its shapes: the M == 1 GEMV loop converts each 128 column chunk after the kernel call, and the M > 1 loop runs its dequant strip SGEMM into the scratch and converts after the rows of a strip are done, bias included. Since both loops fully finish a strip before moving on (the SGEMM runs with ZeroMode over the whole K), the conversion slots in without touching the compute structure, and the fp16 values are bitwise identical to converting the fp32 result.

The shared helpers move above SQ4BitGemm_CompFp32 and lose the misleading CompInt8 prefix (StripCToFp16, GetStripCFp16Scratch) since they now serve both compute types. MlasQNBitGemmFp16DirectCOutputSupported takes the compute type so the operator can ask per path; the CompInt8 answer is unchanged, and CompFp32 answers true for 4 bit wherever its kernels exist. The operator side is the same three lines as before, now also taken when compute_type_ is SQNBIT_CompFp32: skip the fp32 result buffer, set CFp16, skip the final conversion pass.

Motivation and Context

Follow up to #29791/#29864, which noted the C epilogue is portable. accuracy_level=0 is what an export gets when the attribute is not set at all, so this is the default path for a 4 bit fp16 model on x64, and it was the last MatMulNBits configuration still paying the M x N fp32 output buffer plus a full extra conversion pass at compute time.

Validation: a new MLAS test (test_sqnbitgemm_fp16_c_fp32path.cpp) checks the fp16 output against the fp32 result converted with MlasConvertFloatToHalfBuffer for byte equality across BlkLen 16/32/64/128, both loop shapes (M == 1 and M > 1), column remainders against both the 128 and 32 column chunkings, symmetric and asymmetric B, bias on and off, and with and without a thread pool; all pass. An AddressSanitizer build of the full suite is clean. The existing fp16 operator tests (MatMulNBits.Float16_4b_Accuracy0) run through the new branches, verified with temporary probes on both branches during development and then removed. Full MLAS suite green.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[blazingphoenix7]

@hariharans29 Please take a look at your convenience. Thanks

[Copilot]

Pull request overview

This PR extends MLAS’s “direct fp16 C output” optimization to the 4-bit SQNBIT_CompFp32 MatMulNBits path so fp16 models on the default accuracy levels (0–3) no longer need to materialize an M x N fp32 output buffer and then run a separate conversion pass.

Changes:

• Add fp16-direct output support to SQ4BitGemm_CompFp32 by running each output strip into per-thread fp32 scratch and converting to fp16 immediately.
• Refactor the shared strip-conversion helpers to be compute-type agnostic (StripCToFp16, GetStripCFp16Scratch) and update the capability query API to take ComputeType.
• Add a dedicated unit test covering fp16-direct output correctness for the 4-bit CompFp32 path across block lengths, shapes, bias, symmetry, and threadpool usage.

Reviewed changes

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

Show a summary per file

File	Description
onnxruntime/core/mlas/lib/qnbitgemm.cpp	Implements fp16-direct strip conversion for 4-bit CompFp32 and generalizes shared fp16 output helpers.
onnxruntime/core/mlas/inc/mlas_qnbit.h	Updates the public MLAS declaration/docs for MlasQNBitGemmFp16DirectCOutputSupported to be compute-type aware.
onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc	Enables the operator to take the fp16-direct output path for both CompInt8 and 4-bit CompFp32 when supported.
onnxruntime/test/mlas/unittest/test_sqnbitgemm_fp16_quant_a.cpp	Updates existing test to use the new capability-query signature for CompInt8.
onnxruntime/test/mlas/unittest/test_sqnbitgemm_fp16_c_fp32path.cpp	New unit test validating fp16-direct output correctness for 4-bit CompFp32.

[hariharans29]

Review: PR #29901 — [MLAS] Extend the fp16 direct C output to the 4 bit CompFp32 path (head af0fe4e)

Author: @blazingphoenix7 (established MLAS contributor — PRs #29791, #29864, #29842, #29903 and issue #29862). 1 commit, 5 files. CI: 86 / 86 checks OK. Copilot AI reviewer: 0 comments (clean review). Reviewers: @hariharans29 requested, not yet approved.

Verdict: approve. Well-scoped, bit-exact, thoroughly tested follow-up that closes the last remaining performance hole in 4-bit fp16 MatMulNBits on x64. Extends the direct fp16 C output optimization from the CompInt8 path (PRs #29791, #29864) to the CompFp32 path, which is what accuracy_level=0 (the schema default) uses — so this takes effect on the default export configuration for 4-bit fp16 models. The three-line operator-side change and clean shared-helper rename are exactly what you'd want from a follow-up of this shape.

---

What it does

Before this PR:

• CompInt8 path (2/4/8-bit): each worker runs the kernel into per-thread fp32 scratch and converts its strip to fp16 immediately — no full M × N fp32 buffer.
• CompFp32 path (4-bit): still materialized a full M × N fp32 output, then ran a separate MlasConvertFloatToHalfBuffer pass over it.

This PR gives SQ4BitGemm_CompFp32 the same strip-conversion treatment in both loop shapes:

• M == 1 (GEMV branch): converts each 128-column chunk after the kernel call.
• M > 1 (dequant-strip SGEMM branch): runs the SGEMM with ZeroMode (fresh C = A·B_dequant, no accumulate) over the whole K into scratch, adds bias per SGEMM sub-strip, then converts the entire strip once the M rows are done.

Since both loops fully finish a strip before advancing (SGEMM with ZeroMode over the whole K), the conversion slots in without touching the compute structure. And because both paths feed the exact same fp32 bit pattern into the same MlasConvertFloatToHalfBuffer (round-to-nearest-even, deterministic), the fp16 output is bitwise identical to computing fp32 then converting.

---

The five files

File	Change
onnxruntime/core/mlas/inc/mlas_qnbit.h	MlasQNBitGemmFp16DirectCOutputSupported gains a MLAS_QNBIT_GEMM_COMPUTE_TYPE ComputeType parameter. Doc updated to say "Supported for CompInt8 (2, 4 and 8 bit) and for the 4 bit CompFp32 path."
onnxruntime/core/mlas/lib/qnbitgemm.cpp	Capability query rewritten to branch on ComputeType. Shared helpers renamed CompInt8StripToFp16 → StripCToFp16 and GetCompInt8Fp16Scratch → GetStripCFp16Scratch, moved above SQ4BitGemm_CompFp32. SQ4BitGemm_CompFp32 gains the fp16 branch in both M==1 and M>1 loops. Call sites in SQ4BitGemm_CompInt8, SQ8BitGemm_CompInt8, SQ2BitGemm_CompInt8 updated to the new helper names.
onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc	Three-line operator-side change: output_fp16_direct gate widens to (SQNBIT_CompInt8 || SQNBIT_CompFp32) && MlasQNBitGemmFp16DirectCOutputSupported(nbits_, compute_type_). Comment updated.
onnxruntime/test/mlas/unittest/test_sqnbitgemm_fp16_quant_a.cpp	Existing test's call to the capability query updated to pass SQNBIT_CompInt8.
onnxruntime/test/mlas/unittest/test_sqnbitgemm_fp16_c_fp32path.cpp	New — dedicated fp16-direct correctness test for the 4-bit CompFp32 path. 153 lines, 24 configurations.

---

Correctness — bit-exact by construction

Key property: the new fp16 path's output bit pattern equals the old (fp32-then-convert) path's output bit pattern, for every input.

• Both paths use the same fp32 kernels (SQ4BitGemmM1Kernel_CompFp32 for M==1, SQ4BitBlkDequantBForSgemm_CompFp32 + GemmFloatKernel/MlasSgemmKernelZero for M>1). Same fp32 bit pattern in the compute output.
• Both paths use MlasConvertFloatToHalfBuffer for the final fp32→fp16 conversion. That function is deterministic (round-to-nearest-even).
• Strip-boundary rounding is not an issue because each strip is converted independently and column-tiled writes don't overlap in the final CFp16 buffer (strips advance by +n in the column dimension and by whole ldc * RowsHandled in the row dimension).

The new test asserts this property directly with ASSERT_EQ(CFp16[i].val, CFp16Ref[i].val) — comparing the raw fp16 payload bytes. That's the strongest correctness check for a numerical transformation.

M > 1 bias handling — the inner sub-strip loop pattern:

while (RowsRemaining > 0) {
  auto RowsHandled = GetMlasPlatform().GemmFloatKernel(a_strip, dequant_b, c_strip,
                                                       K, RowsRemaining, CountN, lda, c_out_ldc, 1.f, true);
  if (bias) {
    AddBiasForGemm(bias, c_strip, RowsHandled, CountN, c_out_ldc);
  }
  c_strip += c_out_ldc * RowsHandled;
  a_strip += lda * RowsHandled;
  RowsRemaining -= RowsHandled;
}

Structurally identical to the pre-existing fp32 branch — just different local variable names (c_blk → c_strip, ldc → c_out_ldc because the fp16 branch writes into a tight-packed RangeCountM × CountN scratch). Bias is added to each SGEMM sub-strip's rows exactly once. ✓

Scratch mechanics (GetStripCFp16Scratch):

float* GetStripCFp16Scratch(size_t Elements) {
  static thread_local std::vector<float> c_scratch;
  if (c_scratch.size() < Elements) {
    c_scratch.resize(Elements);
  }
  return c_scratch.data();
}

thread_local — each thread gets its own scratch, no cross-thread contention. Grows monotonically to the largest strip that thread has computed. Standard MLAS pattern (matches the pre-PR GetCompInt8Fp16Scratch implementation exactly, which is now this same function under the new name). ✓

---

API surface — MlasQNBitGemmFp16DirectCOutputSupported signature change

Before:

bool MLASCALL MlasQNBitGemmFp16DirectCOutputSupported(size_t BlkBitWidth);

After:

bool MLASCALL MlasQNBitGemmFp16DirectCOutputSupported(size_t BlkBitWidth, MLAS_QNBIT_GEMM_COMPUTE_TYPE ComputeType);

The new required parameter forces every caller to be explicit about which compute path they're querying. That's the right choice — the pre-PR function was implicitly CompInt8-only (as its body comment "Wired for the 2, 4 and 8 bit CompInt8 compute ops" acknowledged), and any generalization needs to make the compute-type part of the contract.

Callers verified:

• contrib_ops/cpu/quantization/matmul_nbits.cc — passes compute_type_. ✓
• test/mlas/unittest/test_sqnbitgemm_fp16_quant_a.cpp — passes SQNBIT_CompInt8 (preserves pre-PR behavior). ✓
• test/mlas/unittest/test_sqnbitgemm_fp16_c_fp32path.cpp — passes SQNBIT_CompFp32 (new). ✓

MLAS is generally treated as an ORT-internal library (not part of the onnxruntime_c_api.h public boundary), so external callers of this symbol are rare. Non-issue for API stability.

The new function body cleanly separates the two compute paths:

if (ComputeType == SQNBIT_CompInt8) {
  const bool supported_bit_width = BlkBitWidth == 2 || BlkBitWidth == 4 || BlkBitWidth == 8;
  return supported_bit_width && Dispatch->QuantizeARowComputeBlkSum_CompInt8_Fp16 != nullptr;
}
if (ComputeType == SQNBIT_CompFp32) {
  // The strip conversion in SQ4BitGemm_CompFp32 is portable, so it only needs the
  // CompFp32 kernels themselves.
  return BlkBitWidth == 4 &&
         Dispatch->SQ4BitGemmM1Kernel_CompFp32 != nullptr &&
         Dispatch->SQ4BitBlkDequantBForSgemm_CompFp32 != nullptr;
}
return false;

CompInt8 answer unchanged. CompFp32 answer requires 4-bit + both fp32 dispatch entries — matches what the strip conversion actually uses in the compute path. ✓

---

Operator-side widening (matmul_nbits.cc)

The change is precisely three lines:

// BEFORE
const bool output_fp16_direct =
    compute_type_ == SQNBIT_CompInt8 && MlasQNBitGemmFp16DirectCOutputSupported(nbits_);

// AFTER
const bool output_fp16_direct =
    (compute_type_ == SQNBIT_CompInt8 || compute_type_ == SQNBIT_CompFp32) &&
    MlasQNBitGemmFp16DirectCOutputSupported(nbits_, compute_type_);

Comment above updated: "the compute path" instead of "the int8 path". The rest of the fp16-direct wiring (skipping the fp32 result buffer allocation, setting params.CFp16, skipping the final conversion pass) is untouched — it already worked for both paths, it was just gated behind the compute-type check. ✓

---

Test coverage — comprehensive

New test test_sqnbitgemm_fp16_c_fp32path.cpp runs the following matrix per BlkLen ∈ {16, 32, 64, 128}:

M	N	K	Symmetric	Bias	ThreadPool	Purpose
1	128	256	asym	no	no	M==1, exact 128 column chunk
1	130	129	sym	yes	no	M==1, 128 + 2-column remainder
4	32	128	asym	yes	no	M>1, exact 32 column chunk
7	129	257	sym	no	no	M>1, 32-chunk remainder
33	63	512	asym	yes	yes	M>1, remainder + TP
129	257	1031	sym	yes	yes	Large M>1, remainder + TP

24 configurations total (6 × 4 BlkLens). Coverage:

• Both loop shapes: M==1 GEMV kernel, M>1 dequant-strip SGEMM loop.
• Both column chunkings and their remainders: 128-column chunks for M==1 (exact + 2-col remainder), 32-column chunks for M>1 (exact + 1-col remainder + 31-col remainder in different-K contexts).
• Symmetric and asymmetric B: QuantBZeroPoint either present or nullptr.
• Bias on and off: exercises AddBiasForGemm in the M>1 branch and the bias argument in the M1 kernel call.
• With and without a thread pool: exercises both single-threaded strip iteration and MlasQNBitGemmBatch's parallelization.

Reference construction (bit-exact):

call_gemm(CFp32, nullptr);              // fp32 path (params.C = CFp32)
call_gemm(nullptr, CFp16);              // new fp16 direct path (params.CFp16 = CFp16)
MlasConvertFloatToHalfBuffer(CFp32, CFp16Ref, M * N);
// then: ASSERT_EQ(CFp16[i].val, CFp16Ref[i].val)

The reference is the fp32 result converted with the same MlasConvertFloatToHalfBuffer the compute path uses internally. Byte-level assertion. Nothing weaker would meaningfully test this optimization.

The PR description also notes:

• Existing operator test MatMulNBits.Float16_4b_Accuracy0 runs through the new branches (verified with dev-time probes).
• Full MLAS suite green.
• AddressSanitizer build of the full suite clean.

---

Minor observations (non-blocking)

1. assert(DataParams->PostProcessor == nullptr) in both new fp16 branches. Matches the existing CompInt8 pattern. In Release builds assert compiles away — if a direct MLAS caller (not going through the ORT operator) ever sets both CFp16 and PostProcessor, the postprocessor is silently skipped. A MLAS_ASSERT_ALWAYS-style hard check would be more defensive, but that's a broader-scope refactor and would touch the pre-PR CompInt8 sites too. Fair to leave as-is for consistency.
2. The else branch for unknown compute types returns false. So if a future compute type gets added (e.g., SQNBIT_CompBf16) without also updating this function, direct fp16 output is disabled for it. Silent safe fallback. Reasonable.
3. Helper name choice StripCToFp16 / GetStripCFp16Scratch. The old CompInt8 prefix was misleading now, but the new names are also a bit terse. "Strip" here means "column slice of the output" (128 columns for M==1, 32 columns for M>1). The comment above the helper is clear enough; the names don't need to explain themselves. Fine.
4. The float* c_blk = (C == nullptr) ? nullptr : C + n; guards in both loop shapes handle the to_fp16 case correctly — when to_fp16 is true, the C pointer isn't accessed on the hot path (the continue after the to_fp16 block skips the rest of the loop iteration). ✓
5. No M > 1 branch's RowsHandled correctness worry. The scratch is sized RangeCountM * CountN, and the sub-strip loop accumulates RowsHandled into RowsRemaining decrement until zero. Each sub-strip writes rows [c_strip, c_strip + RowsHandled) × [0, CountN) in the scratch, and the loop advances c_strip by exactly c_out_ldc * RowsHandled = CountN * RowsHandled bytes. No overwrites, no gaps. ✓
6. The MLAS suite passing under ASAN is a strong signal that the strip pointer arithmetic and scratch bounds are correct. ✓
7. CI 86/86 OK on the first spin — clean pipeline. Good.
8. Author track record. The five PRs from @blazingphoenix7 in this series ([MLAS] Run fp16 MatMulNBits through the int8 path without fp32 temporaries #29791, Widen the fp16 direct C output to 2 and 8 bit MatMulNBits CompInt8 #29864, [MLAS] Quantize fp16 activations directly on the AVX-512 MatMulNBits CompInt8 path #29842, [MLAS] Document SDE based validation of the ISA dispatch tiers #29903, and this one) all show high-quality MLAS work with careful bit-exactness reasoning, thorough test coverage, and clear commit messages. Signal is that this can be reviewed on the code, not on the author's history alone, but the history helps the reviewer trust the "AddressSanitizer clean" and "full suite green" claims.

---

Impact statement (why this matters)

accuracy_level=0 is the schema default for MatMulNBits. Any 4-bit fp16 model exported without explicitly setting accuracy_level lands on SQNBIT_CompFp32. So this optimization takes effect on the default configuration for 4-bit fp16 models on x64 — the most common case in practice. Before this PR, that default configuration paid an unnecessary M × N fp32 output buffer allocation plus a full extra MlasConvertFloatToHalfBuffer pass at every compute-time. After this PR, it doesn't.

The savings per MatMulNBits call:

• Peak memory: M × N × sizeof(float) bytes not allocated. For a 4096×4096 GEMM this is 64 MB per call (allocated + freed).
• Compute: one full pass over M × N fp32→fp16 conversions is eliminated (the per-strip conversion has the same total work but stays in cache).

For LLM decode passes with many MatMulNBits calls per token, this compounds meaningfully.

---

Bottom line

Textbook follow-up: the previous PRs (#29791, #29864) established the strip-conversion pattern and the operator-side fp16-direct wiring for CompInt8; this PR extends both to CompFp32 with a clean shared-helper refactor, a byte-exactness contract, and a dedicated test that covers both loop shapes and all four block lengths across the chunking-remainder boundaries. API surface change is minimal and self-explaining. The 24-config byte-equality test is the correct correctness bar. CI green, ASAN green, no reviewer comments.

Ready for @hariharans29 to approve and merge.