moe: enable native Xe2 path for MXFP8 and block-FP8

[krisclarkdev]

Purpose

Enable the native Xe2 grouped-GEMM MoE path for MXFP8 (FP8 weights + E8M0 block scales, group=32) and block-FP8, instead of defaulting to the slow Python ref_fused_moe expert loop.

Scope / pairing (do not split)

This PR ships both pieces together (fork #9 path flip + kernels):

1. Path selection (fork add xpu op grouped topk #9): stop forcing MXFP8 / block-FP8 onto Python ref_fused_moe; default native with escape hatches below.
2. Kernels: Xe2 grouped-GEMM MXFP8 E8M0 B-scales (W8A16 activations today) + block-FP8 one-time weight dequant → W16A16; fused-moe accuracy tests.

A path-flip-only change without the kernel/dispatch work is incorrect and is not a supported split. All A/B numbers below are for this combined change vs VLLM_XPU_FUSED_MOE_USE_REF=1.

• MXFP8: native W8A16 grouped GEMM with E8M0 B-scales (activations dequantized in Python today). Default on; disable with VLLM_XPU_FUSED_MOE_NATIVE_MXFP8=0 or force ref with VLLM_XPU_FUSED_MOE_USE_REF=1.
• Block-FP8: one-time weight dequant at init → W16A16 native GEMM. Default on; disable with VLLM_XPU_FUSED_MOE_NATIVE_BLOCK_FP8=0.
• Detect MXFP8 scales as uint8 or float8_e8m0fnu.
• Xe2 grouped-GEMM: MXFP8 B-scale dispatch + fused-moe accuracy tests.

Serve note: loading HF MXFP8 MoE (weight_block_size=[1,32]) also needs vLLM to select the XPU MXFP8 MoE backend and transpose block scales with weights to [E,K/g,N] after the XPU [E,N,K]→[E,K,N] weight transpose. That wiring lives in vLLM (not this PR).

Test Plan

.venv/bin/python -m pytest \
  tests/fused_moe/test_grouped_gemm.py::test_xe_grouped_gemm_mxfp8 \
  tests/fused_moe/test_fused_moe.py::test_fused_moe_mxfp8 \
  tests/fused_moe/test_fused_moe.py::test_fused_moe_fp8block \
  -v

Micro + E2E A/B (HAL B70):

bash ab_artifacts/run_pb_e2e_full.sh
# equivalent: MODEL=.../Qwen1.5-MoE-A2.7B-Chat-MXFP8-RTN PHASES=ref,native \
#   bash ab_artifacts/e2e_hotpath_model_ab.sh

• Grouped-GEMM MXFP8 unit accuracy
• Fused MoE MXFP8 + block-FP8 unit accuracy
• Micro A/B: native vs USE_REF for MXFP8 and block-FP8 (ab_moe_native_micro.py / pb_e2e_micro.txt)
• E2E serve A/B on Arc Pro B70 (full bench scrape: tok/s, req/s, TTFT/TPOT/ITL)
• Upstream CI / reviewer smoke on XPU

Test Result

Prior E2E pack (PR open)

Hardware: Intel Arc Pro B70 · Commit: 97c29ca
Model: Qwen1.5-MoE-A2.7B-Chat-MXFP8-RTN (~15 GiB, FP8 + [1,32] / ue8m0)
Recipe: random, 24 prompts, 256→128, max_model_len=2048, gpu_util=0.85, --enforce-eager
Failed requests: 0 / 0 (24/24 both sides)

Side	Env	Median tok/s	Runs (tok/s)
REF	VLLM_XPU_FUSED_MOE_USE_REF=1	26.19	26.19 / 27.19 / 25.46
Native	VLLM_XPU_FUSED_MOE_NATIVE_MXFP8=1	611.0	603.34 / 613.88 / 611.0

Δ tok/s: +2233% vs REF → E2E KEEP (≥10% bar).

Metric (measure 2)	REF	Native
Median TTFT (ms)	3089.73	238.90
Median TPOT (ms)	865.37	37.50

Refreshed full scrape (P-B, 2026-07-25)

Source of truth: ab_artifacts/hal_results/PB_E2E_FULL_VERDICT.{md,json}.

Hardware: Intel Arc Pro B70 · Measured commit: 97c29ca · PR tip: 97c29ca (aligned; no remount/remeasure needed)
Model: Qwen1.5-MoE-A2.7B-Chat-MXFP8-RTN
Recipe: same as prior pack (random, 24 prompts, 256→128, max_model_len=2048, gpu_util=0.85, --enforce-eager)
Compare: combined path selection + kernels vs USE_REF=1
Failed requests: 0 / 0

Micro (op median ms)

Recipe	Ref ms	Native ms	Speedup
MXFP8	3.463	0.210	+1547% (~16.5×)
block-FP8	3.170	0.162	+1857% (~19.6×)

E2E serve (median of 3 timed runs after warmup)

Metric	REF	Native	Δ
Output tok/s	25.1	607.9	+2322%
Total tok/s	n/a	n/a	n/a (not reported by bench scrape)
Request throughput (req/s)	0.20	4.75	+2275%
Median TTFT (ms)	3338.39	248.27	+92.6% improve
Mean TTFT (ms)	3307.27	242.11	+92.7% improve
Median TPOT (ms)	937.46	37.81	+96.0% improve
Mean TPOT (ms)	937.71	37.85	+96.0% improve
Median ITL (ms)	947.70	37.96	+96.0% improve

Per-run output tok/s: REF [25.23, 25.10, 24.35] · Native [607.90, 614.93, 600.60]
E2E KEEP (≥10% tok/s): yes.

The refreshed pack is within noise of the prior +2233% pack (same order of magnitude; both KEEP). Prior numbers are retained above for continuity.

Claim: combined path flip (fork #9) + Xe2 kernels vs the previous Python expert-loop REF on the same checkpoint. Not a claim vs W8A16, not vs CUDA/Marlin MXFP8 backends. Device W8A8 is out of scope for this PR.

Serve logs confirmed XPUMxFp8LinearKernel + Using 'XPU' MxFp8 MoE backend on both sides.

(Optional) Documentation Update

None in-tree. Escape hatches: VLLM_XPU_FUSED_MOE_NATIVE_MXFP8, VLLM_XPU_FUSED_MOE_NATIVE_BLOCK_FP8, VLLM_XPU_FUSED_MOE_USE_REF.

[krisclarkdev]

Refreshed P-B A/B artifact index (HAL B70, 2026-07-25)

Combined path selection + native MXFP8/block-FP8 kernels vs VLLM_XPU_FUSED_MOE_USE_REF=1 (same pairing as this PR; do not split).

Pack	Location
Verdict (MD)	ab_artifacts/hal_results/PB_E2E_FULL_VERDICT.md
Verdict (JSON)	ab_artifacts/hal_results/PB_E2E_FULL_VERDICT.json
Micro log	ab_artifacts/hal_results/pb_e2e_micro.txt
E2E AB JSON	ab_artifacts/hal_results/e2e_hotpath_Qwen1.5-MoE-A2.7B-Chat-MXFP8-RTN_ab.json
Bench logs	ab_artifacts/hal_results/e2e_hotpath_bench_Qwen1.5-MoE-A2.7B-Chat-MXFP8-RTN_{ref,native}_{warmup,m1,m2,m3}.log

Headline (refreshed): micro MXFP8 +1547% / block-FP8 +1857%; E2E output tok/s 25.1 → 607.9 (+2322%), KEEP yes. Prior open-PR pack (26.19 → 611, +2233%) retained in the description for continuity.

[mayuyuace]

Please provide below data:

1. grouped gemm performance data like tflops and bandwidth for different shapes.
2. gsm8k results to compare ref path and kernels path.

[mayuyuace]

Please also add UT for block fp8.

[krisclarkdev]

Added test_xe_grouped_gemm_block_fp8 here (FP8 weights + float32 2D scales [E, K/128, N/128] vs a dequant+matmul gold) plus test_fused_moe.py::test_fused_moe_fp8block. Both pass on Arc Pro B70 — see the follow-up comment with the full 28-passed run and data.

[mayuyuace]

Weights data type should not be changed.
Casting to fp32 will spend more memory.
You should realize really block fp8 kernels instead of casting in memory.

[krisclarkdev]

Fixed in 8a4dee4: the fp32 weight promotion is removed. Block-FP8 weights stay float8_e4m3fn in memory and the float32 2D block scales are consumed inside the Xe2 grouped GEMM (W8A16, group_size=128), so no weight-sized fp32 buffer is allocated. Accuracy vs the ref path is cosine ≈ 1.0 (see follow-up comment).

[krisclarkdev]

Refreshed P-B A/B artifact index (HAL B70, 2026-07-25)

Combined path selection (fork #9) + native MXFP8/block-FP8 kernels vs VLLM_XPU_FUSED_MOE_USE_REF=1 (same pairing as this PR; do not split).

Pack	Location
Verdict (MD)	ab_artifacts/hal_results/PB_E2E_FULL_VERDICT.md
Verdict (JSON)	ab_artifacts/hal_results/PB_E2E_FULL_VERDICT.json
Micro log	ab_artifacts/hal_results/pb_e2e_micro.txt
E2E AB JSON	ab_artifacts/hal_results/e2e_hotpath_Qwen1.5-MoE-A2.7B-Chat-MXFP8-RTN_ab.json
Bench logs	ab_artifacts/hal_results/e2e_hotpath_bench_Qwen1.5-MoE-A2.7B-Chat-MXFP8-RTN_{ref,native}_{warmup,m1,m2,m3}.log

SHA: measured on 97c29ca; PR tip is 97c29ca (aligned).

Headline (refreshed): micro MXFP8 +1547% / block-FP8 +1857%; E2E output tok/s 25.1 → 607.9 (+2322%), KEEP yes. Prior open-PR pack (26.19 → 611, +2233%) retained in the description for continuity.

[kfojcik-intel]

Hi @krisclarkdev, I wasn't aware of this PR when I opened #490.

I noticed your block FP8 approach dequantizes weights to bf16/fp16 at init time and then runs the standard W16A16 grouped GEMM. In #490 I took a different approach that keeps weights in FP8 at runtime and uses a partial-accumulator kernel:

Memory: weights stay as 1-byte FP8 — no doubling of weight memory from init-time dequant
Bandwidth: reads FP8 from global memory (half the bytes vs. bf16)
Compute: uses XE2 native FP8 DPAS instructions directly, accumulates partial results per K-block, then applies the float32 block scale at group boundaries
For MXFP8 (E8M0 uint8 scales, gs=32) your pre-multiply-B approach makes sense and there's no overlap — that part is orthogonal.
There will be some merge conflicts in gemm_xe2.hpp and grouped_gemm_xe2.hpp depending on which lands first. Happy to coordinate — e.g. I can rebase on top of yours once it merges, or we can discuss combining into a single PR if that's preferred.

Maybe @mayuyuace can help to decide

[kfojcik-intel]

@krisclarkdev, I see you just pushed 8a4dee4 which keeps FP8 weights in memory and applies the float32 block scale in-kernel — nice, that eliminates the init-time dequant. A couple of observations on the kernel strategy:

Your implementation routes block FP8 through the existing pre-multiply-B path in xe_gemm_4bits — meaning the FP8 weight gets converted to float, multiplied by the block scale, then cast to bf16 before entering DPAS. This effectively makes it W16A16 compute with FP8 storage.

In #490 I use a partial-accumulator approach instead:

The FP8 weight goes directly into DPAS as-is (bf16 × FP8 → fp32, native W8A16 hardware path)
Partial fp32 results accumulate across K-tiles within a group
At each group boundary, the partial accumulator is multiplied by the block scale and flushed to the final accumulator
The advantage is that XE2 FP8 DPAS processes 2× elements per cycle compared to bf16 DPAS, since FP8 operands are half the width. You also avoid the fp8→float→scale→bf16 conversion overhead on every K-tile.

Worth benchmarking both to see which wins in practice — the partial accumulator adds a branch + scale multiply every group_size/tile_k tiles, but the 2× DPAS throughput should more than compensate.

[krisclarkdev]

Review follow-up: no fp32 weight cast, block-FP8 UT, and requested data

All three review points addressed. Measured on Intel Arc Pro B70, oneAPI 2026.0 / torch 2.13.0+xpu, PR tip 8a4dee4.

1. Weights stay FP8 (no fp32 cast in memory)

The fp32 weight promotion in fused_moe_interface.py is removed. Block-FP8 now runs a real in-kernel path: FP8 weights [E,K,N] are kept as-is and the 2D float32 block scales [E, K/128, N/128] are consumed inside the Xe2 grouped GEMM (W8A16 policy, group_size=128), matching how MXFP8 already consumed E8M0 block scales. No extra weight-sized fp32 buffer is allocated.

• csrc/xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp: detect FP8 B + float32 3D scale → block-FP8, launch W8A16 with ElementS=float.
• grouped_gemm_xe2.hpp / gemm_xe2.hpp: route group_size==128 to the block path; load float 2D scale Scales[k_block*(N/128) + n_block]; skip the 1D E8M0 prefetch when scales are float.

2. Block-FP8 UT

Added tests/fused_moe/test_grouped_gemm.py::test_xe_grouped_gemm_block_fp8 (FP8 weights + float32 2D scales vs a dequant+matmul gold), plus test_fused_moe.py::test_fused_moe_fp8block asserting the weights are not promoted and stay float8_e4m3fn. Full MoE accuracy suite on B70:

tests/fused_moe/test_grouped_gemm.py::test_xe_grouped_gemm_mxfp8       8 passed
tests/fused_moe/test_grouped_gemm.py::test_xe_grouped_gemm_block_fp8   8 passed
tests/fused_moe/test_fused_moe.py::test_fused_moe_mxfp8               8 passed
tests/fused_moe/test_fused_moe.py::test_fused_moe_fp8block            4 passed
============================= 28 passed in 23.52s =============================

3a. Grouped-GEMM performance (TFLOPS / bandwidth)

Median of 50 iters, 8 experts (tokens split round-robin), 2*M*N*K flops. Read+write bytes = FP8 B (1B) + float scales + bf16 A/D.

Recipe	M	N	K	E	ms	TFLOPS	GB/s
mxfp8	64	1024	1024	8	0.072	1.86	123
block_fp8	64	1024	1024	8	0.054	2.47	159
mxfp8	128	1024	1024	8	0.070	3.81	130
block_fp8	128	1024	1024	8	0.053	5.03	167
mxfp8	256	2048	1024	8	0.077	13.94	245
block_fp8	256	2048	1024	8	0.057	18.80	321
mxfp8	512	2048	2048	8	0.180	23.91	216
block_fp8	512	2048	2048	8	0.121	35.54	313
mxfp8	1024	4096	2048	8	0.337	51.02	243
block_fp8	1024	4096	2048	8	0.233	73.81	343
mxfp8	2048	4096	4096	8	1.060	64.84	162
block_fp8	2048	4096	4096	8	0.767	89.57	219

Harness: ab_artifacts/bench_grouped_gemm_tflops.py; raw JSON ab_artifacts/hal_results/grouped_gemm_tflops.json.

3b. Accuracy: ref path vs kernels path

Full GSM8K serve needs the vLLM-side MXFP8 MoE serve wiring (backend selection + block-scale transpose), which lives in vLLM and is out of scope for this kernels PR; that path is not available in the current test image. As the equivalent in-repo accuracy check, identical activations + weights are fed through the kernels path (XpuFusedMoe native Xe2 grouped GEMM) and the ref path (per-expert dequant → bf16/fp16 matmul), reporting divergence over several MoE shapes:

Recipe	M	H	E	cosine	rel-err mean	rel-err p99
mxfp8	64	256	8	0.99999	0.0103	0.111
mxfp8	256	512	8	0.99999	0.0127	0.169
mxfp8	512	1024	8	1.00000	0.0019	0.027
fp8block	64	256	8	0.99997	0.0196	0.176
fp8block	256	512	8	0.99997	0.0243	0.293
fp8block	512	1024	8	1.00000	0.0031	0.044

Cosine ≈ 1.0 across shapes; residual per-element error sits at the FP8 quantization noise floor, i.e. the kernels path matches the ref path. Harness: ab_artifacts/moe_accuracy_ref_vs_kernels.py; JSON ab_artifacts/hal_results/moe_accuracy_ref_vs_kernels.json.

If you'd like a true end-to-end GSM8K number rather than the layer-level comparison, I can run it once the vLLM MXFP8 MoE serve wiring is in place (happy to point at the wiring branch), but that is a separate vLLM change from this kernels PR.

[krisclarkdev]

Hi @kfojcik-intel — thanks, no worries on the overlap. You're right that 8a4dee4 scales the B fragment every k-tile; your partial-accumulator (one scale per tiles_per_group) is the cleaner structure and should win at compute-heavy shapes.

One correction, since it affects what we'd measure: I don't think #490 hits native FP8 DPAS. The MMA atom is still XE_DPAS_TT<8, float, ElementA_non_CV> (grouped_gemm_xe2_interface.hpp:93) and your diff leaves it and reorder(tBrB, tCrB) unchanged, so B is still converted FP8→bf16 into the fragment before cute::gemm. Both PRs are bf16 DPAS with FP8 storage; the difference is where the scale lands, not operand width. The 2× would need a mixed-precision atom, which neither of us has wired up.

Minor: your block-scale prefetch Int<SG_N / group_size>{} truncates to Int<0> at group_size=128 (SG_N is 64 for w8a16_policy, 16 for the m_16/m_32 policies) — looks like it moves zero elements there. Intended?

The PRs are complementary: yours routes MXFP8 to the ref path so #490 is block-FP8 only, while this one does both, and it carries the TFLOPS/bandwidth sweep and accuracy-vs-ref data mayuyuace asked for. Your group_size handling (32/64/128/256 + square check) is more general than my hard-coded 128, so I'd take yours there. Happy to rebase/combine however @mayuyuace prefers — and if you push the #490 kernel I'll run it through bench_grouped_gemm_tflops.py on the B70 for an apples-to-apples number.

[kfojcik-intel]

@krisclarkdev
You're absolutely right on both points -- thanks for catching those.

MMA atom: Correct, reorder(tBrB, tCrB) already converts FP8→bf16 before DPAS -- so the difference is ALU overhead only, not 2× DPAS throughput. I stand corrected.

Prefetch: Fixed - removed the block FP8 scale prefetch. I still need to think this over but I think that SG_N / group_size yields at most 2 with current policies, which I believe is below make_block_2d_prefetch's minimum shape requirement. The scale tensor is only a few KB per expert so L2 handles it fine. All 32 test cases pass on BMG for me now.

The updated kernel is at version #490 -- you can download it. I'll keep you updated with my performance comparison soon.
Would be great to see your numbers too if you get a chance to run it.

[krisclarkdev]

Required companion of vLLM FA-in-graph / XPU graphs PR: vllm#50038.

Please treat kernels #485–#489 as part of the graphs stack alongside vllm #48677 (torch/image) and #49813 (softcap/MXFP8).

[mayuyuace]

@kfojcik-intel
Based on your comments, I would like to merge this PR firstly.
Could you update your optimization like dynamic group size for block fp8 in next PR?

[mayuyuace]

@krisclarkdev
Please fix pre-commit.

[krisclarkdev]

Fixed pre-commit in 8e86676 (amended tip): clang-format on the Xe2 headers, isort/ruff on fused_moe_interface.py and the fused_moe test (E501 / SIM103). Local pre-commit run --all-files --hook-stage manual is clean.

[krisclarkdev]

Correction: tip is d7cb8ae (DCO-aligned amend of the formatting commit). Local verification:

pre-commit run --all-files --hook-stage manual

All hooks passed. CI self-hosted runs show action_required (likely workflow approval) — happy to re-trigger once approved.

[mayuyuace]

@jikunshang
CI has passed.
Please help review this PR.

[jikunshang]

LGTM. thanks for contributing!