Repository navigation
Megablocks mxfp8 cri - #1066
Megablocks mxfp8 cri#1066
Conversation
Port the block-scaled FP8 path from the Intel reference (607c031) so MoE
weights can be stored as float8_e4m3fn/float8_e5m2 with E8M0 block scales,
alongside the existing int4 and MXFP4 paths.
MXFP8 reuses the generic block-scaled kernel (xe_gemm_4bits) rather than
adding a new one. What is new: the is_B_mxfp8 type trait, the E8M0 scale
decode branch, the 3D [E, N, K/group_size] scale pointer offset, the
unpacked [E, N, K] weight layout, and the MXFP8LauncherCallER dispatch.
The dispatch is mutually exclusive at compile time (is_B_mxfp8 requires
ElementB in {e4m3, e5m2} *and* ElementS == uint8_t, while w8a16-fp8 uses a
float ElementS and int4 a uint8_t ElementB), so is_B_block_scaled reduces
to the previous is_B_4bits for every pre-existing path.
Validated on the CRI simulator against an fp32 CPU reference: the three
MXFP8 cases (e4m3, e5m2, e4m3+bias) land in the same residual range as the
unquantized w16a16 baseline, and the w16a16 / w8a16-fp8 / int4 / MXFP4
paths are unchanged.
The Xe35 paths in sycl-tla (2D transpose/prefetch copies, mxfp8 fcvt dequant) are gated behind __SYCL_TARGET_INTEL_GPU_CRI__, which the compiler never defines on its own. Passing it globally would leak into the pvc and bmg images, so add a CRI-only component that compiles the grouped GEMM a second time with the macro forced on and spir64_gen as its sole target; the link step AOT-lowers that image to cri. Both images cover the same dtypes, so the host picks between them by reading the device IP version (pvc 12.x, bmg 20.x, cri 35.x) through a new XPUFeatures helper, mirroring the existing CPUFeatures convention. The two compilations expand the shared headers into different code, so wrap them in an inline namespace keyed on the macro and add an arch tag to GemmCuteName. Without this the launchers mangle identically in both translation units and the linker silently keeps one body. Collapse the former pass-through wrapper into a single templated entry point that static_asserts its Arch against SYCL_INTEL_TARGET. Verified on the built extension: Li35E kernels appear only in the cri device images, pvc and bmg carry none, and the BMG test suite still passes 134/134.
The comment claimed this TU makes sycl-tla emit a block-scaled DPAS path. That is not what happens: bdpas requires both operands to share a low-precision type and to carry e8m0 scale factors, while these kernels are W4A16/W8A16. What the macro actually gates is SYCL_INTEL_TARGET == 35, which cutlass.h derives from __SYCL_TARGET_INTEL_GPU_CRI__, and that selects the Xe35 block 2D cache controls and the native fcvt lowering of the e4m3 -> bf16 reorder.
| "Tensor? ptr_scales, Tensor? ptr_bias, Tensor! ptr_D, " | ||
| "Tensor expert_first_token_offset, int N, int K, int num_experts, " | ||
| "bool is_B_int4, bool is_B_mxfp4) -> Tensor"); | ||
| "bool is_B_int4, bool is_B_mxfp4, bool is_B_mxfp8) -> Tensor"); |
There was a problem hiding this comment.
This might warrant a version bump in build.toml Not 100% sure though, do we expect that people use this op directly?
There was a problem hiding this comment.
Good question. The op isn't meant to be used directly — inside the package it is only
reached through xpu_fused_moe.py (cutlass_grouped_gemm / cutlass_grouped_gemm_xe2,
plus the register_fake shim), and nothing from that module is re-exported in
megablocks/__init__.py. Consumers go through the MegaBlocksMoeMLP layer, whose
signature is unchanged.
To be safe for anyone calling the op directly anyway, I gave the new argument a default:
bool is_B_int4, bool is_B_mxfp4, bool is_B_mxfp8=False
The schema change is now purely additive — the previous 11-argument call still resolves,
so no caller breaks and a v2 branch isn't needed. A major bump would force every
consumer to move get_kernel(..., version=1) to 2 for an argument they never pass.
|
/kernel-bot build megablocks |
|
Build request processed. Command: Dispatched (1):
|
The grouped GEMM op gained an is_B_mxfp8 argument, which made the registered schema incompatible with the previous call signature. Giving the argument a default keeps the change purely additive, so existing callers of the op keep working and the kernel does not need a major version bump.
|
/kernel-bot build megablocks |
|
Build request processed. Command: Dispatched (1):
|
…L_TARGET The arch guard compared the template parameter against SYCL_INTEL_TARGET, which only sycl-tla 0.9.2 and later rewrite to 20/35 in cutlass.h. On sycl-tla 0.8 it keeps the value 1 that the build system defines, so both translation units failed the assertion and the xpu20253 variant did not build. Derive MEGABLOCKS_XE_TARGET_ARCH from __SYCL_TARGET_INTEL_GPU_CRI__ instead, next to the namespace macro that already keys off it. That also keeps the SYCL kernel name tag distinct between the two units on older sycl-tla, where SYCL_INTEL_TARGET would have labelled both images 1.
|
/kernel-bot build megablocks |
|
Build request processed. Command: Dispatched (1):
|
|
/kernel-bot merge-and-upload megablocks |
|
Build request processed. Command: Merge result: PR is already merged. Continuing with build/upload. Dispatched (1):
Hub uploads: |
2 participants
Depends on the build-side changes in
huggingface/kernels#734.
MXFP8
Block-scaled FP8 (
e4m3/e5m2weights withe8m0scales, group size 32/64/128)now goes through the same block-scaled kernel path as the existing MXFP4 support,
so it costs a dtype branch and a scale-offset rule rather than a new kernel.
Weights are
[E, N, K]and unpacked, unlike the 4-bit path.This adds an
is_B_mxfp8argument tocutlass_grouped_gemm_interface, so the opschema changes.
Why the GEMM is built twice
sycl-tla selects its Xe35 code paths on
__SYCL_TARGET_INTEL_GPU_CRI__, and thecompiler never defines it on its own. The macro reaches through the headers, so
a unit built with it is an Xe35 unit throughout and cannot also cover pvc and
bmg.
The grouped GEMM is therefore compiled twice from the same sources — an Xe20 unit
for pvc/bmg and a CRI-only unit with the macro forced on and
spir64_genas thesole target — and the host picks between them by the device IP version reported by
the driver. Both units cover the same dtypes, so nothing is lost on older GPUs.
Two things follow from compiling the same headers twice:
(
xe20/xe35), otherwise the linker merges symbols that share a name but nota body;
collide.
A
static_assertties the instantiation toSYCL_INTEL_TARGET, so a build thatloses the macro fails instead of silently producing an Xe20 kernel labelled Xe35.
Validation
Built against
kernels#734(torch213-cxx11-xpu20260-x86_64-linux) and checkednumerics on BMG hardware and the CRI TBX simulator. Each case runs only the
kernel on XPU; data generation, the fp32 reference, dequantization and error
statistics stay on CPU.
relative error 2.6e-03.
The dispatch was checked rather than assumed: on the simulator the driver reports
IP version 35.11.0, the runtime loads the CRI AOT image instead of falling back
to JIT, and the kernel that launches is
MoE::xe35::GemmCuteName<..., 35>.One thing worth stating: MXFP8 here is a software dequantization feeding a plain
dpas. The Xe35 block-scaledbdpasinstructions are not used, so that remainsavailable as a later optimization.