cpu: x64: matmul: add MXFP4 (f4_e2m1 + e8m0) weights decompression for f32 matmul by amakarev · Pull Request #5645 · uxlfoundation/oneDNN

amakarev · 2026-07-21T16:24:15Z

Adds MXFP4 (f4_e2m1 weights + e8m0 group-32 scales) weights decompression for the f32 x64 brgemm matmul primitive on avx512_core.

Two decompression paths, selected by problem size:

Fused (M ≤ 4): inline f4 → f32 decode with e8m0 scale application inside the brgemm microkernel. Bypasses buffer_b entirely; one brgemm call covers exactly one K-group of 32.
Copy_b (M > 4): standard pre-decompression pass into buffer_b, then native f32 GEMM.

Copilot

Pull request overview

Adds MXFP4 (f4_e2m1 weights + e8m0 group-32 scales) support to the x64 BRGEMM-based matmul path by introducing a fused in-kernel decode path for small M and extending the existing copy-B decompression pipeline for larger problems.

Changes:

Extend brgemm matmul configuration to recognize f32:f4_e2m1:f32 + e8m0 scales and to route between fused decode (no buffer_b) vs copy-B decompression.
Add fused FP4 decode + e8m0 scale preload support to the x64 JIT brgemm kernel and plumb an optional weights-scales pointer through the execute API.
Add reorder/test plumbing for f4 storage and add benchdnn harness inputs for MXFP4 decompression coverage.

Reviewed changes

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

Show a summary per file

File	Description
tests/benchdnn/inputs/matmul/shapes_llm_mxfp4	Adds an LLM-oriented batch of MXFP4 matmul shapes for benchdnn.
tests/benchdnn/inputs/matmul/harness_matmul_decompression	Adds a benchdnn option set to exercise MXFP4 weights decompression.
src/cpu/x64/matmul/brgemm_matmul.cpp	Enables f32+f4 weights path and configures brgemm descriptors for fused FP4 decode (scale handling, packed offsets).
src/cpu/x64/matmul/brgemm_matmul_utils.hpp	Adds config flags for f32+f4 weights and generic “packed elems per byte” handling.
src/cpu/x64/matmul/brgemm_matmul_utils.cpp	Implements f32+f4 detection, ISA gating, buffer_b selection, and packed-element plumbing.
src/cpu/x64/matmul/brgemm_matmul_copy_utils.cpp	Extends copy-B decompression to handle f4_e2m1 and e8m0 scales; refactors 4-bit masking helpers.
src/cpu/x64/brgemm/jit_brgemm_kernel.cpp	Adds fused FP4 decode path in the microkernel (LUT/permd/masks + inline scaling).
src/cpu/x64/brgemm/brgemm.hpp	Extends execute APIs with an optional `ptr_wei_scales` argument.
src/cpu/x64/brgemm/brgemm.cpp	Plumbs `ptr_wei_scales` into `brgemm_kernel_params_t` on execution.
src/cpu/x64/brgemm/brgemm_utils.cpp	Updates datatype support and blocking rules for fused FP4 decode.
src/cpu/x64/brgemm/brgemm_types.hpp	Adds `is_f4_fused_decompress` flag and ensures post-op/store paths remain available.
src/cpu/reorder/simple_reorder.hpp	Adds a reference any→any reorder for f4_e2m1→f4_e2m1 (byte copy / index remap).
src/cpu/reorder/cpu_reorder_regular_fp4.cpp	Registers the new f4_e2m1→f4_e2m1 reorder implementation.
src/cpu/platform.cpp	Exposes f4_e2m1 datatype availability on x64 AVX-512 Core.
src/cpu/matmul/ref_matmul.hpp	Expands reference matmul datatype acceptance to include f4_e2m1.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

xuxinzen · 2026-07-22T19:44:20Z

-              && attr.fpmath_.apply_to_int_)
+    , f32_with_f4_wei_dt(everyone_is(f32, bgmmc.src_dt, bgmmc.dst_dt)
+              && bgmmc.wei_dt == data_type::f4_e2m1)
+    , weights_decompression_support(f32_with_f4_wei_dt


Do we have any restriction on fpmath for f32 with f4?

No restriction on fpmath is needed here. The MXFP4 pipeline decompresses f4_e2m1 * e8m0_scale → f32 and runs the FMA in f32; fpmath mode has no effect on either side of that.

xuxinzen · 2026-07-22T19:44:40Z

            && IMPLICATION(bm_conf_utils.is_f16_fp8(),
                    one_of(isa, avx512_core_amx_fp16, avx10_2))
+            && IMPLICATION(
+                    bm_conf_utils.is_f32_with_f4_wei(), isa == avx512_core)


should we use is_superset(isa, avx512_core)?

is_f4_fused_decompress is a zmm-non-AMX code path. Widening the check to is_superset(isa, avx512_core) would also match avx512_core_amx/avx512_core_fp16/avx10_2 etc., but the fused f4 nibble decode is only wired up in jit_brgemm_kernel under is_f4_fused_decompress_non_amx() (i.e. is_zmm && !is_tmm). On AMX-capable ISAs it wouldn't take effect.

I agree with @amakarev's conclusion but I believe the main reason for the check is that f32 is only supported by brgemm_matmul_t<avx512_core> (for machines with avx512 support). Everything else is the consequences of that.

densamoilov · 2026-07-27T14:33:03Z

@@ -0,0 +1,16 @@
+128x256x1536:128x1536x512_n"deepseek_r1_1*61"


Is this file unused?

densamoilov · 2026-07-27T15:24:12Z

            && IMPLICATION(bm_conf_utils.is_f16_fp8(),
                    one_of(isa, avx512_core_amx_fp16, avx10_2))
+            && IMPLICATION(
+                    bm_conf_utils.is_f32_with_f4_wei(), isa == avx512_core)


I agree with @amakarev's conclusion but I believe the main reason for the check is that f32 is only supported by brgemm_matmul_t<avx512_core> (for machines with avx512 support). Everything else is the consequences of that.

densamoilov · 2026-07-27T15:27:46Z

-    const int rd_unroll = 4;
-    const data_type_t rd_block_dt = get_mac_emu_data_type(
-            brg->dt_a, brg->isa_impl, brg->isa_impl != avx2_vnni_2);
+    const int rd_unroll = brg->is_f4_fused_decompress ? 32 : 4;


can you add a comment documenting this choice (32) and relevant constrains (e.g. when and if it can be changed).

As far as I understand fusion is a performance optimization. Did you see any performance benefits from this optimization?

densamoilov · 2026-07-27T15:35:32Z

            || bgmmc.wei_tag == adbc;
    bgmmc.use_buffer_b = bm_conf_utils.use_buffer_b();

+    // Fused f4 decode applies scales inside the brgemm kernel, so don't fold


The term decode is not really defined in oneDNN. We use decompression and quantization or a shorthand for them. I'd recommend to stick with the well known vocabulary.

densamoilov · 2026-07-27T15:47:01Z

 dim_t jit_brgemm_kernel_t<Wmm>::ldb_B_offset(
        dim_t ld_block2, bool is_tail) const noexcept {
+    if (brg.is_f4_fused_decompress_non_amx())
+        return (is_tail) ? brg.ldb_tail * brg.ld_step / 2


nit: since typesize_B is always defined as 1 for f4 data type so we could reuse the same formula by introducing b_elems_per_byte. Here and in other places.
const dim_t b_elems_per_byte = brg.is_f4_fused_decompress_non_amx() ? 2 : 1;

densamoilov · 2026-07-27T15:55:42Z

            = bm_conf_utils.is_bf16_fp8() || bm_conf_utils.is_f16_fp8();
    bgmmc.with_wei_decompression = bm_conf_utils.with_weights_decompression();
+    // The usage of this variable hardcodes the assumption
+    // there's only supported value > 1, which is 2. If this


?

Suggested change

// there's only supported value > 1, which is 2. If this

// there's only one supported value > 1, which is 2. If this

densamoilov · 2026-07-27T16:09:25Z

        if (bgmmc.is_xf16_fp8) return true;
        if (bgmmc.is_bf16_with_int_wei) return true;
+        if (bgmmc.is_f32_with_f4_wei) {
+            constexpr dim_t fused_M_threshold = 4;


Can you please add a justification for this threshold value?

densamoilov · 2026-07-27T16:12:20Z

-    const int reserved_regs_;
    const bool is_wei_grouped_over_k_;

+    struct vmm_idx_t {


Can you please elaborate on why you chose to create a structure instead of member functions we usually use for this purpose?

densamoilov · 2026-07-27T16:16:27Z

    bool is_xf16_fp8 = false;
    bool is_int4_weights = false;
-    bool is_f4_via_convert = false;
+    int wei_packed_elems_per_byte = 1;


I think the default value should be 0 as we expect this struct to be initialized before it's used.

densamoilov · 2026-07-27T16:21:48Z

+        if (bgmmc.is_f32_with_f4_wei) {
+            constexpr dim_t fused_M_threshold = 4;
+            const bool fused_eligible = !bgmmc.is_amx
+                    && is_superset(bgmmc.isa, avx512_core)


Do we actually need to check this here? I think if it's false then it doesn't matter as we'll return unimplemented anyway.

Same question about other conditions.

amakarev requested review from a team as code owners July 21, 2026 16:24

github-actions Bot added platform:cpu-x64 Intel64/AMD64 processors. Codeowner: @oneapi-src/onednn-cpu-x64 component:tests Codeowner: @oneapi-src/onednn-arch component:common labels Jul 21, 2026

amakarev mentioned this pull request Jul 21, 2026

cpu: x64: matmul: add MXFP4 weights support #4922

Closed

amakarev force-pushed the amakarev/matmul/f32_mxfp4 branch 2 times, most recently from cc0881c to 4d09e73 Compare July 21, 2026 16:41

amakarev requested a review from Copilot July 21, 2026 16:41

Copilot started reviewing on behalf of amakarev July 21, 2026 16:42 View session

Copilot AI reviewed Jul 21, 2026

View reviewed changes

Comment thread src/cpu/x64/brgemm/jit_brgemm_kernel.cpp Outdated

Comment thread src/cpu/x64/matmul/brgemm_matmul_utils.cpp Outdated

amakarev force-pushed the amakarev/matmul/f32_mxfp4 branch 2 times, most recently from 1882c04 to f383a96 Compare July 21, 2026 17:32

inteldimitrius reviewed Jul 22, 2026

View reviewed changes

xuxinzen reviewed Jul 22, 2026

View reviewed changes

georgen117 reviewed Jul 22, 2026

View reviewed changes

Comment thread src/cpu/x64/matmul/brgemm_matmul_utils.cpp

inteldimitrius approved these changes Jul 23, 2026

View reviewed changes

georgen117 approved these changes Jul 24, 2026

View reviewed changes

amakarev added 5 commits July 27, 2026 07:55

cpu: x64: matmul: add MXFP4 weights decompression for f32 matmul

a47d65b

benchdnn: matmul: switch mxfp4 harness to per_tensor scale policy

05f4157

cpu: x64: brgemm: reuse per-K wei scales path for fused f4 decode

6b9fd6f

cpu: x64: brgemm: drop ptr_wei_scales from brgemm_kernel_execute

6de2ff6

cpu: x64: matmul: restore is_int4_weights assignment

4cbc656

amakarev force-pushed the amakarev/matmul/f32_mxfp4 branch from 737b8a9 to 4cbc656 Compare July 27, 2026 14:57

densamoilov reviewed Jul 27, 2026

View reviewed changes

		@@ -0,0 +1,16 @@
		128x256x1536:128x1536x512_n"deepseek_r1_1*61"

	// there's only supported value > 1, which is 2. If this
	// there's only one supported value > 1, which is 2. If this

Uh oh!

Conversation

amakarev commented Jul 21, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants