[MoE] Add a bf16xfp32 gemm kernel for router gemm

[jjmiao1]

Purpose

Add an XPU operator gemm_bf16xfp32 that emulates an FP32-precision GEMM (BF16 activation × FP32 weight) using two BF16 GEMM passes (DualGemm). It targets the MoE router GateLinear in models such as Hunyuan-V3 and MiniMax-M3, where the gate weight is kept in FP32 for the numerical stability of routing. Currently this runs as FP32 F.linear: the BF16 activation is upcast to FP32 and fed to a oneDNN FP32 matmul. The execution on the FP32 pipe (~12.3 TFLOP/s on B60) is much slower than the BF16 (~98.5 TFLOP/s on B60). gemm_bf16xfp32 decomposes the FP32 weight into two BF16 components and combines them in the epilogue, recovering FP32-class accuracy (absolute error ~1e-5) while cutting latency (4–7.7×) on prefill;

How it works

The FP32 weight W is split into two BF16 tensors:

• W_high = bf16(W)
• W_low = bf16((W - fp32(W_high)) / scale) (default scale = 1/256)

DualGemm The kernel computes D = X @ W_high + (X @ W_low) * scale, sharing the A (activation) load and fusing the correction in the epilogue. The DualGemm collective is adapted from the dual_gemm example in sycl-tla: it runs the two BF16 GEMMs from a single kernel so A is loaded once and the W_low correction is applied in the shared epilogue. W_high + W_low * scale reproduces W to ~FP32 precision, so the result can match FP32 F.linear to FP32-class accuracy.

Split-K for decode At small M (decode) the default tile leaves most Xe-cores idle. For those shapes the kernel splits the K — each slice computes a partial [M,N], and a small fused kernel sums the partials.

Code Changes

• csrc/xpu/gemm_bf16xfp32/ — kernel interface + implementation (three tile policies: small-M / medium-M / default, dispatched by M; Split-K + fused reduction for decode).
• vllm_xpu_kernels/gemm_bf16xfp32.py — Python API: split_fp32_weight, gemm_bf16xfp32.
• tests/test_gemm_bf16xfp32.py — correctness vs FP32 F.linear.
• benchmark/benchmark_gemm_bf16xfp32.py — latency / roofline benchmark vs FP32 F.linear.
• CMakeLists.txt / setup.py — build wiring for gemm_bf16xfp32_xe_2.

Usage

from vllm_xpu_kernels.gemm_bf16xfp32 import split_fp32_weight, gemm_bf16xfp32

w_high, w_low = split_fp32_weight(weight_fp32)          # [N,K] -> two [K,N] bf16 
out = gemm_bf16xfp32(x_bf16, w_high, w_low)             # == F.linear(x, weight)

Note: Weight preparation is a one-time, offline step.split_fp32_weight decomposes the FP32 router weight into the two BF16 components (W_high, W_low); If integrated into upstream vLLM, this will run once at load time in the layer's process_weights_after_loading — the FP32 router weight will split there and W_high / W_low be cached.

Test Plan

• Correctness: pytest tests/test_gemm_bf16xfp32.py — compares kernel output against FP32 F.linear over the M × N × K sweep. Element-wise assert_close(rtol=1e-3, atol=1e-4) plus aggregate bounds (max_abs < 1e-4, mean_rel < 1e-4) and a dual-bf16 ≤ naive-bf16 invariant.
• Performance: python benchmark/benchmark_gemm_bf16xfp32.py on Intel Arc Pro B60 (BMG / Xe2) — reports device + wall latency, TFLOP/s / MFU / GB/s / MBU, and speedup vs FP32.
  Hardware: Intel Arc Pro B60 (BMG, Xe2). Peaks: BF16 98.5 TFLOP/s, FP32 vector 12.28 TFLOP/s, BW 456 GB/s.
  Baseline: FP32 F.linear (oneDNN) — the correctness-equivalent op this kernel replaces.
  Sweep: M ∈ {2 … 16384}, N ∈ {128, 192, 256}, K ∈ {3072, 4096, 6144}. 100 iters / 20 warmup.
  Measurement: ker dev / fp32 dev are on-device kernel time; ker wall / fp32 wall are end-to-end eager wall-clock; dev speedup / wall speedup are fp32 / kernel. TFLOP/s counts both BF16 passes (4·M·N·K); MFU vs BF16 peak; GB/s / MBU from IO traffic (A + W_high + W_low + D) vs BW peak.

Test Result

• pytest tests/test_gemm_bf16xfp32.py: all pass
  Accuracy is measured against the FP32 (oneDNN) baseline across the full sweep: absolute error ~1e-5 (worst 2.5e-5), mean relative ~2e-5.
• Benchmark vs FP32 F.linear (see tables below)

Representative MNK— Hunyuan-V3 router gemm shape (N = 192 experts, K = 4096 hidden), full M sweep:

     M    N     K |  ker dev  ker wall |   TFLOP/s     MFU      GB/s     MBU |  fp32 dev  fp32 wall | dev speedup wall speedup
--------------------------------------------------------------------------------------------------------------------------------
     2  192  4096 |     6.77     15.53 |      0.93    0.9%     467.2  102.4% |      7.92      42.09 |       1.17x        2.71x
     8  192  4096 |     6.70     15.40 |      3.76    3.8%     480.6  105.4% |     18.43      39.85 |       2.75x        2.59x
    16  192  4096 |     6.64     15.65 |      7.58    7.7%     495.4  108.7% |     18.87      40.16 |       2.84x        2.57x
    32  192  4096 |     6.67     15.47 |     15.09   15.3%     514.5  112.8% |     19.31      40.30 |       2.90x        2.61x
    64  192  4096 |     9.33     15.31 |     21.58   21.9%     398.7   87.4% |     21.34      39.44 |       2.29x        2.58x
   128  192  4096 |    12.07     15.37 |     33.36   33.9%     355.6   78.0% |     29.50      42.18 |       2.44x        2.74x
   256  192  4096 |    20.51     21.25 |     39.26   39.9%     265.2   58.2% |     54.43      57.78 |       2.65x        2.72x
   512  192  4096 |    31.20     31.86 |     51.61   52.4%     247.8   54.3% |    107.51     110.42 |       3.45x        3.47x
  1024  192  4096 |    60.14     64.00 |     53.56   54.4%     204.9   44.9% |    218.61     221.87 |       3.64x        3.47x
  2048  192  4096 |   126.81    124.54 |     50.80   51.6%     169.5   37.2% |    438.95     442.29 |       3.46x        3.55x
  4096  192  4096 |   176.12    176.40 |     73.16   74.3%     226.2   49.6% |    877.21     880.30 |       4.98x        4.99x
  8192  192  4096 |   297.13    295.30 |     86.73   88.0%     257.6   56.5% |   1748.78    1750.21 |       5.89x        5.93x
 16384  192  4096 |   587.07    584.28 |     87.79   89.1%     255.4   56.0% |   3489.41    3487.94 |       5.94x        5.97x

Representative MNK — MiniMax-M3 router shape (N = 128 experts, K = 6144 hidden), full M sweep:

     M    N     K |  ker dev  ker wall |   TFLOP/s     MFU      GB/s     MBU |  fp32 dev  fp32 wall | dev speedup wall speedup
--------------------------------------------------------------------------------------------------------------------------------
     2  128  6144 |     7.27     15.44 |      0.87    0.9%     436.5   95.7% |      8.34      42.63 |       1.15x        2.76x
     8  128  6144 |     7.17     15.46 |      3.51    3.6%     453.3   99.4% |     11.65      41.74 |       1.63x        2.70x
    16  128  6144 |     7.17     15.42 |      7.02    7.1%     467.6  102.5% |     11.95      41.61 |       1.67x        2.70x
    32  128  6144 |     7.25     15.84 |     13.88   14.1%     490.2  107.5% |     17.27      44.22 |       2.38x        2.79x
    64  128  6144 |    10.07     15.39 |     20.00   20.3%     393.9   86.4% |     23.07      41.63 |       2.29x        2.71x
   128  128  6144 |    11.91     15.57 |     33.80   34.3%     401.6   88.1% |     32.74      41.77 |       2.75x        2.68x
   256  128  6144 |    21.24     22.25 |     37.92   38.5%     302.4   66.3% |     61.81      65.49 |       2.91x        2.94x
   512  128  6144 |    38.27     38.90 |     42.09   42.7%     253.5   55.6% |    128.84     132.33 |       3.37x        3.40x
  1024  128  6144 |    46.60     47.07 |     69.12   70.2%     348.7   76.5% |    279.65     283.43 |       6.00x        6.02x
  2048  128  6144 |   142.01    142.97 |     45.37   46.1%     206.7   45.3% |    577.69     581.25 |       4.07x        4.07x
  4096  128  6144 |   341.38    337.97 |     37.74   38.3%     162.8   35.7% |   1411.30    1415.87 |       4.13x        4.19x
  8192  128  6144 |   536.50    533.94 |     48.03   48.8%     201.3   44.1% |   2333.43    2334.42 |       4.35x        4.37x
 16384  128  6144 |   915.95    907.54 |     56.27   57.1%     232.4   51.0% |   4533.84    4528.51 |       4.95x        4.99x

Benchmark Results Summary

• Prefill (large M) — compute-bound. The kernel runs on the BF16 pipe (98.5 TFLOP/s) while FP32 F.linear is stuck on the FP32 vector pipe (12.28 TFLOP/s). Result: 4–7.7× device speedup, reaching up to ~89% MFU.
• Decode (small M) — bandwidth-bound. Two BF16 weight passes move the same bytes as one FP32 weight, so the device-time floor is FP32 parity.

[Copilot]

Pull request overview

This PR adds a new Intel XPU (Xe2) custom operator, gemm_bf16xfp32, to accelerate MoE router GEMMs where activations are BF16 but router weights remain FP32 for routing stability, by emulating FP32-class accuracy via a dual BF16 GEMM (DualGemm) plus a scaled correction term.

Changes:

• Adds an Xe2 DualGemm kernel (policy-dispatched by M, with split-K + reduction for small-M decode shapes).
• Wires the kernel into the build + PyTorch dispatcher (torch.ops._xpu_C.gemm_bf16xfp32).
• Adds Python helpers, a correctness test against FP32 F.linear, and a benchmark script.

Reviewed changes

Copilot reviewed 14 out of 14 changed files in this pull request and generated 10 comments.

Show a summary per file

File	Description
vllm_xpu_kernels/gemm_bf16xfp32.py	Python API for weight splitting and calling the new op.
tests/test_gemm_bf16xfp32.py	Correctness sweep vs FP32 F.linear.
tests/register_ops.py	Test helper wrapper for invoking the op via torch.ops._xpu_C.
benchmark/benchmark_gemm_bf16xfp32.py	Performance benchmark comparing to FP32 F.linear.
csrc/xpu/torch_bindings.cpp	Registers the new op schema + XPU implementation.
csrc/xpu/gemm_bf16xfp32/gemm_bf16xfp32_interface.h	Declares the top-level C++ interface.
csrc/xpu/gemm_bf16xfp32/gemm_bf16xfp32_interface.cpp	Dispatches to Xe2 implementation (and errors on non-Xe2).
csrc/xpu/gemm_bf16xfp32/xe_2/CMakeLists.txt	Adds the Xe2 kernel library target.
csrc/xpu/gemm_bf16xfp32/xe_2/gemm_bf16xfp32_xe2.h	Declares the Xe2 entrypoint.
csrc/xpu/gemm_bf16xfp32/xe_2/gemm_bf16xfp32_xe2.cpp	Thin wrapper calling the Xe2 implementation.
csrc/xpu/gemm_bf16xfp32/xe_2/gemm_bf16xfp32_xe2_impl.hpp	Core Xe2 DualGemm kernel + split-K reduction + policy dispatch.
csrc/xpu/gemm_bf16xfp32/xe_2/bf16xfp32_epilogue.hpp	Custom epilogue combining the two accumulator results.
setup.py	Installs the additional Xe2 kernel library artifact.
CMakeLists.txt	Adds kernel library build + links it into _xpu_C.

---

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