[MoE] Add Xe2 FP8 W8A16 grouped GEMM support

[CaoE]

Summary

Add FP8 E4M3 MoE support for Xe2/BMG with BF16 activations and outputs.

The existing use_fp8_w8a8=True API remains compatible, but dispatches to W8A16 because Xe2 has no native FP8 x FP8 MMA and the evaluated W8A8 path was slower.

Changes

• Add Xe2 FP8 W8A16 grouped GEMM with:

  • per-expert scalar and 128x128 block scales
  • optional expert bias
  • SiLU, GELU, ReLU2, GPT-OSS SwiGLU, and DeepSeek-V4 gated activation
  • AOT and runtime JIT support

• Use a static grid-stride scheduler for grouped GEMMs with small routed
  work or short-K scalar scales:

  • total_m <= num_experts, or
  • scalar-scale GEMMs with K <= 128.

  This avoids dynamic tile-stealing atomics and improves short-GEMM latency
  by about 2.2%-3.6%.

• Add a specialized small-batch input-preparation path for FP8 MoE.
  For workloads with a small number of routed rows, one workgroup performs the expert histogram, prefix computation, stable route ranking, and activation scatter, avoiding the generic multi-kernel preparation path.

Performance

Measured on Intel Arc Pro B60 with PyTorch 2.12.0+xpu using XPU events.

Small-batch input preparation

Shape (M, TopK, E, H)	Generic / Optimized time ratio
(1, 8, 256, 2048)	2.28x
(2, 8, 128, 2048)	2.28x
(4, 8, 128, 2048)	1.61x
(5, 8, 128, 2048)	1.40x
(8, 8, 128, 128)	2.26x

End-to-end MoE

Scalar FP8 weight scales

Shape (H, I, E, TopK)	Tokens	BF16/FP8 time ratio
(7168, 512, 256, 8)	1	1.76x
(7168, 512, 256, 8)	32	2.15x
(7168, 512, 256, 8)	2048	1.58x
(2048, 256, 512, 10)	1	1.18x
(2048, 256, 512, 10)	32	1.92x
(2048, 256, 512, 10)	2048	1.38x

128x128 block FP8 weight scales

Shape (H, I, E, TopK)	Tokens	BF16/FP8 time ratio
(7168, 512, 256, 8)	1	1.73x
(7168, 512, 256, 8)	32	2.15x
(7168, 512, 256, 8)	2048	1.58x
(2048, 256, 512, 10)	1	1.18x
(2048, 256, 512, 10)	32	1.90x
(2048, 256, 512, 10)	2048	1.32x

[Copilot]

Pull request overview

This PR adds Xe2/BMG support for MoE checkpoints with FP8 (E4M3) expert weights by introducing a BF16-activation / FP8-weight grouped GEMM (W8A16) backend, wiring it through the Torch extension + Python fused_experts API, and providing both AOT and SYCL JIT launch paths. It also adds a small-routed fast-path input preparation kernel and extensive correctness/contract tests and benchmarks around the new FP8 flow.

Changes:

• Add Xe2 FP8(E4M3) weight-only grouped GEMM kernel (scalar + 128×128 block scales, optional bias) with AOT + JIT dispatch.
• Integrate FP8 weight-only execution into python/sgl_kernel/moe.py::fused_experts, including activation variants and small-routed input preparation.
• Add new tests and benchmarks covering API validation, kernel correctness, workspace caching, and performance.

Reviewed changes

Copilot reviewed 19 out of 19 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
tests/test_moe_prepare_input.py	Adds unit tests for small-route prepare selector and new prepare_moe_input_small op correctness/error handling.
tests/test_moe_gemm.py	Adds FP8 E4M3 W8A16 reference quant/dequant helpers and end-to-end FP8 MoE + grouped GEMM tests (scalar + block scales, activations, bias).
tests/test_moe_fused_experts_workspace.py	Updates workspace cache tests to cover the new view-cache behavior.
src/torch_extension_sycl.cc	Registers new ops: moe_grouped_mm_nt_xe20_fp8_w8a16 and prepare_moe_input_small.
src/sycl/MoEPrepareInputs.cpp	Implements prepare_moe_input_small SYCL kernel and its validation/device capability checks.
src/sycl/kernels/moe/xe20/fp8/moe_mainloop.hpp	Introduces FP8-weight/BF16-activation mainloop (scalar + block-scale handling) for Xe2.
src/sycl/kernels/moe/xe20/fp8/moe_kernel.hpp	Implements grouped GEMM kernel wrapper for FP8 weights, including scheduler modes and scale layout handling.
src/sycl/GroupGemmFp8W8A16Xe20LauncherInstance.cpp.in	Adds templated per-tile launcher instance for JIT compilation and entry symbol.
src/sycl/GroupGemmFp8W8A16Xe20.cpp	Adds dispatcher + validation for FP8 W8A16 grouped GEMM, including tile selection and static scheduler rules.
src/jit/moe_jit.h	Declares JIT launch API for FP8 W8A16 grouped GEMM.
src/jit/moe_jit.cpp	Adds FP8 W8A16 tile selection + JIT compilation cache plumbing for runtime kernels.
src/GroupGemmFp8W8A16Xe20.cmake	Adds AOT instantiation generation for the FP8 W8A16 grouped GEMM tile/scale/bias matrix.
src/CMakeLists.txt	Wires new FP8 grouped GEMM sources into the build (and MOE exclusion filter).
src/BuildOnLinux.cmake	Ensures FP8 instance bundles are included/linked correctly for AOT and JIT configurations.
python/sgl_kernel/moe.py	Adds FP8 weight-only MoE path, FP8 scale validation, small-route prepare selection, and workspace view caching.
include/sgl_kernel_ops.h	Exposes C++ API for moe_grouped_mm_nt_xe20_fp8_w8a16 and prepare_moe_input_small.
benchmark/bench_moe_fp8_w8a16_grouped_gemm.py	Adds op-level benchmark comparing SGL scalar/block FP8 paths (and optional vLLM reference).
benchmark/bench_fused_moe_fp8.py	Adds small opt-in end-to-end FP8 MoE benchmark with preloading of FP8 AOT instances.
.isort.cfg	Adds sgl_kernel classification to isort config.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[airMeng]

can you reuse W4A16 files since they are the same?

[CaoE]

Thank you for your comments. I will try to reuse the w4a16 implementation.

[airMeng]

Shall we request the FWK side to do the padding? You can compare the performance w/ and w/o padding

@jmunetong for awareness

[CaoE]

Thanks for your comments. Use real m size instead. No significant performance changes were observed w/ and w/o padding (~2% grouped gemm time).

[airMeng]

Can you try to reuse W4A16 code with several constexpr dispatch?