onednn-fp8-gemm

oneDNN Quantized GEMM on Xe2 (Lunar Lake/LNL, Battlemage/BMG)

Specialized knowledge for implementing FP16/BF16 x FP8 and FP16 x INT4 quantized GEMM using oneDNN's matmul primitive on Intel Xe2 (Lunar Lake/LNL, Battlemage/BMG) discrete GPU.

---

Quick-Reference Rules (must follow every time)

Supported Configurations on BMG (oneDNN 2025.2)

KEY DIFFERENCE FROM PTL: On BMG with oneDNN 2025.2, both FP16xFP8 and BF16xFP8 have optimized JIT kernels. On PTL, only FP16xFP8 had JIT; BF16xFP8 fell back to ocl:ref:any.

Layout Convention

• A: [M, K] row-major (format_tag::ab) -- FP16 or BF16 activations
• B: logical [K, N], physical [N, K] (format_tag::ba) -- FP8_E4M3 or U4 weights
• scale: [N] FP32 for FP8; [n_groups, N] FP16 for INT4
• C: [M, N] row-major (format_tag::ab) -- same dtype as A

CRITICAL: set_scales_mask vs set_scales

Always check implementation string after creating primitive_desc:

std::string impl = pd.impl_info_str();
if (impl.find("ref") != std::string::npos) {
    printf("WARNING: Fallback to slow reference implementation!\n");
}

---

Pattern 1: FP16/BF16 x FP8_E4M3 with Per-N Scale

C[M,N] = A[M,K] x dequant(B[K,N])  where dequant(B) = B_fp8 * scale[N]

Memory Descriptors

using dt = dnnl::memory::data_type;
using ft = dnnl::memory::format_tag;

dnnl::memory::desc a_md({M, K}, dt::f16, ft::ab);     // or dt::bf16
dnnl::memory::desc b_md({K, N}, dt::f8_e4m3, ft::ba); // logical [K,N], physical [N,K]
dnnl::memory::desc scale_md({N}, dt::f32, ft::a);
dnnl::memory::desc c_md({M, N}, dt::f16, ft::ab);     // same dtype as A

Primitive Attributes

dnnl::primitive_attr attr;
attr.set_scales_mask(DNNL_ARG_WEIGHTS, 2);  // per-N, MUST use set_scales_mask
attr.set_fpmath_mode(dnnl::fpmath_mode::any, true);

Execute

dnnl::matmul::primitive_desc pd(eng, a_md, b_md, c_md, attr);
dnnl::matmul prim(pd);

std::unordered_map<int, dnnl::memory> args = {
    {DNNL_ARG_SRC,                            dnnl::memory(a_md, eng, ptr_A)},
    {DNNL_ARG_WEIGHTS,                        dnnl::memory(b_md, eng, ptr_B)},
    {DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS, dnnl::memory(scale_md, eng, ptr_scales)},
    {DNNL_ARG_DST,                            dnnl::memory(c_md, eng, ptr_C)},
};

// Optional bias: [N] in same dtype as A
if (has_bias) {
    dnnl::memory::desc bias_md({N}, dt::f16, ft::a);
    args.insert({DNNL_ARG_BIAS, dnnl::memory(bias_md, eng, ptr_bias)});
}

prim.execute(s, args);

FP8 E4M3 Format

• 1 sign + 4 exponent + 3 mantissa, bias=7, no inf
• Max value: +/-448, min normal: 2^(-6)
• Per-N absmax quantization: scale[n] = max(|B[n,:]|) / 448

---

Pattern 2: FP16 x INT4 (U4) with Block-wise Scale + Zero-point

C[M,N] = A[M,K] x dequant(B[K,N])
where dequant(B) = (B_u4 - zero_point) * scale[group, N]

Memory Descriptors

int block_size = 64;  // quantization group size along K
int n_groups = K / block_size;

dnnl::memory::desc a_md({M, K}, dt::f16, ft::ab);
dnnl::memory::desc w_md({K, N}, dt::u4, ft::ba);  // 4-bit unsigned, transposed
dnnl::memory::desc scale_md({n_groups, N}, dt::f16, ft::ba);  // block-wise
dnnl::memory::desc zp_md({1}, dt::u8, ft::a);  // scalar zero point
dnnl::memory::desc c_md({M, N}, dt::f16, ft::ab);

Primitive Attributes

dnnl::primitive_attr attr;
// Block-wise scales: mask = (1<<0)|(1<<1) = 3, group = {block_size, 1}
attr.set_scales(DNNL_ARG_WEIGHTS, (1 << 1) + (1 << 0), {block_size, 1}, dt::f16);
attr.set_zero_points(DNNL_ARG_WEIGHTS, 0, {}, dt::u8);
attr.set_fpmath_mode(dnnl::fpmath_mode::any, true);

Execute

std::unordered_map<int, dnnl::memory> args = {
    {DNNL_ARG_SRC,                                dnnl::memory(a_md, eng, ptr_A)},
    {DNNL_ARG_WEIGHTS,                            dnnl::memory(w_md, eng, ptr_W)},
    {DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS,     dnnl::memory(scale_md, eng, ptr_scales)},
    {DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_WEIGHTS, dnnl::memory(zp_md, eng, ptr_zp)},
    {DNNL_ARG_DST,                                dnnl::memory(c_md, eng, ptr_C)},
};
prim.execute(s, args);

INT4 Weight Packing (PyTorch side)

# Pack two U4 values per byte: low nibble first
packed[...] = qweight[:, 0::2]           # even indices -> low nibble
packed[...] |= qweight[:, 1::2] << 4     # odd indices -> high nibble

INT4 Per-group Absmax Quantization

grouped = weight.reshape(-1, group_size)
max_vals = grouped.abs().max(-1).values
scales = max_vals / 7.0                   # map to [1, 15] with zp=8
qweight = clamp(round(grouped / scales.unsqueeze(-1) + 8), 0, 15).to(uint8)

---

Pattern 3: FP16 x FP8_E4M3 with Block-wise Scale (1D along K)

C[M,N] = A[M,K] x dequant(B[K,N])
where dequant(B[gk*bk+k, n]) = B_fp8[gk*bk+k, n] * scale[gk, n]

Block-wise along K with per-N granularity within each block. Same pattern as INT4 (Pattern 2) but with FP8 weights and FP32 scales.

Memory Descriptors

int block_k = 128;  // block size along K
int n_groups_k = K / block_k;

dnnl::memory::desc a_md({M, K}, dt::f16, ft::ab);
dnnl::memory::desc b_md({K, N}, dt::f8_e4m3, ft::ba);  // logical [K,N], physical [N,K]
dnnl::memory::desc scale_md({n_groups_k, N}, dt::f32, ft::ab);  // block-wise FP32
dnnl::memory::desc c_md({M, N}, dt::f16, ft::ab);

Primitive Attributes

dnnl::primitive_attr attr;
// mask=3 = (1<<0)|(1<<1): per both K and N dims of logical B[K,N]
// groups={block_k, 1}: block along K, per-N within each block
attr.set_scales(DNNL_ARG_WEIGHTS, (1 << 1) + (1 << 0),
                {block_k, 1}, dnnl::memory::data_type::f32);
attr.set_fpmath_mode(dnnl::fpmath_mode::any, true);

Emulating 2D Block Quantization (block_k x block_n)

True 2D block scales [K/bk, N/bn] are NOT supported on oneDNN v3.7 (2025.2) -- primitive creation fails. Requires oneDNN v3.9+.

Workaround: Quantize with 2D blocks, then expand compact scales to 1D format:

# Quantize weight [N, K] with 2D blocks (128x128)
# scales_compact: [K//128, N//128]  -- one scale per block

# Expand to 1D format: repeat each N-group scale 128 times
# scales_1d: [K//128, N]  -- what oneDNN v3.7 accepts
scales_1d = scales_compact.repeat_interleave(128, dim=1).contiguous()

# Call oneDNN with expanded scales
out = sycl_kernels.onednn_w8a16_fp8_block(x, qweight, scales_1d, block_k=128)

Scale memory overhead: compact [32, 32] = 4 KB → expanded [32, 4096] = 512 KB (still small vs weight tensor ~16 MB).

Known Issues (oneDNN v3.7 / 2025.2)

• FP16 block-wise: WORKS -- jit:gemm:any, rel_rms = 0.0004
• BF16 block-wise: BROKEN -- selects jit:gemm:any but produces garbage output (rel_rms ~164). Use FP16 only for block-wise FP8.
• 2D native {128, 128} groups: primitive creation fails, needs oneDNN v3.9+

---

Scale Mask Cheat Sheet

• Mask=0: per-tensor (single scale for entire weight)
• Mask=1: per-K (per input channel)
• Mask=2: per-N (per output channel) -- most common for weight quantization
• Mask=3: per both dims = block-wise (groups along K, per-N)

Note: FP8 per-N MUST use set_scales_mask, NOT set_scales. FP8 block-wise uses set_scales with dt::f32. INT4 block-wise uses set_scales with dt::f16.

---

Required Headers and Initialization

#include <sycl/sycl.hpp>
#include <oneapi/dnnl/dnnl.hpp>
#include <oneapi/dnnl/dnnl_sycl.hpp>

sycl::queue q(sycl::gpu_selector_v);
dnnl::engine eng = dnnl::sycl_interop::make_engine(q.get_device(), q.get_context());
dnnl::stream s = dnnl::sycl_interop::make_stream(eng, q);

In PyTorch context: auto q = c10::xpu::getCurrentXPUStream(device.index()).queue();

---

Build Commands

Standalone C++ (no PyTorch)

set TORCH_INC=<conda_env>\Lib\site-packages\torch\include
set DNNL_LIB=<conda_env>\Library\lib
icpx test.cpp -fsycl -o test.exe -I"%TORCH_INC%" -L"%DNNL_LIB%" -ldnnl -std=c++20

Where <conda_env> provides oneapi/dnnl/dnnl.hpp via torch's include and dnnl.lib via Library/lib.

Within PyTorch extension (CMake)

find_package(IntelSYCL REQUIRED)
find_package(Torch REQUIRED)
target_link_libraries(${MODULE} torch torch_python dnnl)

---

Performance Results (BMG, oneDNN 2025.2)

Pure C++ ULT (no PyTorch overhead)

PyTorch Wrapper (lightx2v_kernel_xpu)

BMG peak: 135 TFLOPS FP16 XMX. FP16xFP8 reaches ~96% in pure C++, ~83% through PyTorch.

---

Correctness Testing

• Reference: FP32 CPU with per-N dequantization (FP8) or per-group dequantization (INT4)
• Expected rel_rms: < 0.01 for FP8, < 0.05 for INT4 (higher due to 4-bit quantization)
• Check for NaN before computing error
• All tests verified on BMG: rel_rms 0.0002-0.0034 for FP8, 0.0003 for INT4

---

Sample Code

---

Reference Files