From 2fe6123afbccaddc35c639beb726c37a5f0e86e4 Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Mon, 27 Jul 2026 19:48:52 -0700 Subject: [PATCH 1/6] megablocks: add MXFP8 support to the XPU grouped GEMM 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. --- .../grouped_gemm/grouped_gemm_interface.cpp | 6 +- .../grouped_gemm/grouped_gemm_interface.h | 3 +- .../csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp | 4 +- .../grouped_gemm/xe_2/grouped_gemm_xe2.cpp | 6 +- .../grouped_gemm/xe_2/grouped_gemm_xe2.h | 3 +- .../grouped_gemm/xe_2/grouped_gemm_xe2.hpp | 16 ++++- .../xe_2/grouped_gemm_xe2_interface.hpp | 62 ++++++++++++++++++- .../torch-ext/megablocks/xpu_fused_moe.py | 34 ++++++---- megablocks/torch-ext/torch_binding.cpp | 2 +- 9 files changed, 113 insertions(+), 23 deletions(-) diff --git a/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.cpp b/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.cpp index bf6ccf6e8..b97adaa6a 100644 --- a/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.cpp +++ b/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.cpp @@ -13,7 +13,8 @@ torch::Tensor cutlass_grouped_gemm_interface( int64_t K, int64_t num_experts, bool is_B_int4, - bool is_B_mxfp4) { + bool is_B_mxfp4, + bool is_B_mxfp8) { // Use XE2 cutlass kernel return cutlass_grouped_gemm_xe2( ptr_A, @@ -26,5 +27,6 @@ torch::Tensor cutlass_grouped_gemm_interface( K, num_experts, is_B_int4, - is_B_mxfp4); + is_B_mxfp4, + is_B_mxfp8); } diff --git a/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.h b/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.h index b25b71d12..7dfcd1eb2 100644 --- a/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.h +++ b/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.h @@ -11,4 +11,5 @@ torch::Tensor cutlass_grouped_gemm_interface( int64_t K, int64_t num_experts, bool is_B_int4, - bool is_B_mxfp4); \ No newline at end of file + bool is_B_mxfp4, + bool is_B_mxfp8); \ No newline at end of file diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp index af916201c..dbe1e09e1 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp @@ -366,7 +366,9 @@ CUTE_DEVICE void xe_gemm_4bits( scale = Scales [(n_tile_start + n_sg_start + sg_local_n) * group_num + group_idx]; - } else if constexpr (std::is_same_v) { + } else if constexpr (std::is_same_v || + std::is_same_v || + std::is_same_v) { uint32_t scale_u32 = Scales [(n_tile_start + n_sg_start + sg_local_n) * group_num + diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.cpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.cpp index 2d6c4b29a..d2a2f9714 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.cpp +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.cpp @@ -13,7 +13,8 @@ torch::Tensor cutlass_grouped_gemm_xe2( int64_t K, int64_t num_experts, bool is_B_int4, - bool is_B_mxfp4) { + bool is_B_mxfp4, + bool is_B_mxfp8) { return MoE::cutlass_grouped_gemm_xe2_impl( ptr_A, ptr_B, @@ -25,5 +26,6 @@ torch::Tensor cutlass_grouped_gemm_xe2( K, num_experts, is_B_int4, - is_B_mxfp4); + is_B_mxfp4, + is_B_mxfp8); } \ No newline at end of file diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.h b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.h index feb722c6e..ea2d9a362 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.h +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.h @@ -11,4 +11,5 @@ torch::Tensor cutlass_grouped_gemm_xe2( int64_t K, int64_t num_experts, bool is_B_int4, - bool is_B_mxfp4); \ No newline at end of file + bool is_B_mxfp4, + bool is_B_mxfp8); \ No newline at end of file diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.hpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.hpp index be7386409..0e8f6d458 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.hpp +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.hpp @@ -90,7 +90,12 @@ CUTE_DEVICE void MoEGEMM( (!std::is_same_v); static constexpr bool is_B_mxfp4 = (std::is_same_v) && (std::is_same_v); + static constexpr bool is_B_mxfp8 = + (std::is_same_v || + std::is_same_v) && + (std::is_same_v); static constexpr bool is_B_4bits = std::is_same_v; + static constexpr bool is_B_block_scaled = is_B_4bits || is_B_mxfp8; auto item = sycl::ext::oneapi::this_work_item::get_nd_item<3>(); auto wg_tile = mma.tile_mnk(); @@ -146,6 +151,12 @@ CUTE_DEVICE void MoEGEMM( if constexpr (is_B_4bits) { ptr_Scales_curr_batch = const_cast(Scales) + B_offset * 2 / group_size; + } else if constexpr (is_B_mxfp8) { + ptr_Scales_curr_batch = + const_cast(Scales) + + static_cast(expert_id) * + static_cast(gemm_n) * + (static_cast(gemm_k) / group_size); } ElementBI* ptr_Bias_curr_batch = nullptr; if (Bias != static_cast(nullptr)) { @@ -161,6 +172,9 @@ CUTE_DEVICE void MoEGEMM( } else if constexpr (is_B_mxfp4) { return make_moe_tensor( reinterpret_cast(ptr_B_curr_batch), gemm_n, gemm_k); + } else if constexpr (is_B_mxfp8) { + return make_moe_tensor( + ptr_B_curr_batch, gemm_n, gemm_k); } else { return make_moe_tensor( ptr_B_curr_batch, gemm_n, gemm_k); @@ -174,7 +188,7 @@ CUTE_DEVICE void MoEGEMM( int m_coord = (group_m_id - pre_tiles); auto tile_coord = make_coord(m_coord, n_coord, _, 0); - if constexpr (is_B_4bits) { + if constexpr (is_B_block_scaled) { #define XE_GEMM_4BITS_CALLER(GroupSize) \ xe_gemm_4bits( \ A_tensor, \ diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp index 6a0a6d312..5ab94efdd 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp @@ -50,11 +50,11 @@ #include "cutlass/tensor_ref.h" #include "cutlass/util/GPU_Clock.hpp" #include "cutlass/util/device_memory.h" +#include "cutlass/util/sycl_event_manager.hpp" #include "cutlass/util/initialize_block.hpp" #include "cutlass/util/reference/device/gemm_complex.h" #include "cutlass/util/reference/device/tensor_compare.h" #include "cutlass/util/reference/host/tensor_fill.h" -#include "cutlass/util/sycl_event_manager.hpp" #include "gemm_xe2_policy.hpp" #include "grouped_gemm_xe2.hpp" @@ -173,7 +173,8 @@ at::Tensor cutlass_grouped_gemm_xe2_impl( int64_t K, int64_t num_experts, bool is_B_int4, - bool is_B_mxfp4) { + bool is_B_mxfp4, + bool is_B_mxfp8) { auto& dpcpp_queue = at::xpu::getCurrentXPUStream(ptr_A.device().index()).queue(); auto A_dtype = ptr_A.dtype(); @@ -206,6 +207,10 @@ at::Tensor cutlass_grouped_gemm_xe2_impl( if (is_B_int4 || is_B_mxfp4) { B_K = ptr_B.size(2) * 2; B_N = ptr_B.size(1); + } else if (is_B_mxfp8) { + // MXFP8: weights stored as [E, N, K] (no packing), dtype is fp8 + B_K = ptr_B.size(2); + B_N = ptr_B.size(1); } int D_total_M = ptr_D.size(0); @@ -298,6 +303,59 @@ at::Tensor cutlass_grouped_gemm_xe2_impl( W4A16LauncherCallER(policy); } #undef W4A16LauncherCallER + } else if (is_B_mxfp8) { + // MXFP8: block-scaled FP8 with e8m0 scales + TORCH_CHECK(ptr_scales.has_value(), "mxfp8 grouped gemm must have scales"); + TORCH_CHECK(ptr_scales->is_contiguous(), "ptr_scales must be contiguous"); + TORCH_CHECK( + ptr_scales->dim() == 3, + "ptr_scales of mxfp8 must be 3D [num_experts, N, K // group_size]"); + TORCH_CHECK( + ptr_scales->size(0) == num_experts, + "ptr_scales.size(0) of mxfp8 must match num_experts"); + TORCH_CHECK( + ptr_scales->size(1) == N, + "ptr_scales.size(1) of mxfp8 must match N"); + TORCH_CHECK( + K % ptr_scales->size(2) == 0, + "K must be divisible by ptr_scales.size(2) for mxfp8"); + int group_num = ptr_scales->size(2); + group_size = K / group_num; + + TORCH_CHECK( + group_size == 32 || group_size == 64 || group_size == 128, + "mxfp8 group_size must be 32, 64 or 128"); + +#define MXFP8LauncherCallER(policy) \ + if (B_dtype == at::kFloat8_e4m3fn && A_dtype == at::kBFloat16) { \ + using scalar_t = bfloat16_t; \ + MoEGEMMLauncherCallER( \ + 'R', 'C', policy, scalar_t, float_e4m3_t, uint8_t); \ + } else if (B_dtype == at::kFloat8_e4m3fn && A_dtype == at::kHalf) { \ + using scalar_t = half_t; \ + MoEGEMMLauncherCallER( \ + 'R', 'C', policy, scalar_t, float_e4m3_t, uint8_t); \ + } else if (B_dtype == at::kFloat8_e5m2 && A_dtype == at::kBFloat16) { \ + using scalar_t = bfloat16_t; \ + MoEGEMMLauncherCallER( \ + 'R', 'C', policy, scalar_t, float_e5m2_t, uint8_t); \ + } else if (B_dtype == at::kFloat8_e5m2 && A_dtype == at::kHalf) { \ + using scalar_t = half_t; \ + MoEGEMMLauncherCallER( \ + 'R', 'C', policy, scalar_t, float_e5m2_t, uint8_t); \ + } + + if (A_avg_M <= 32) { + using policy = w8a16_policy_m_16; + MXFP8LauncherCallER(policy); + } else if (A_avg_M <= 128) { + using policy = w8a16_policy_m_32; + MXFP8LauncherCallER(policy); + } else { + using policy = w8a16_policy; + MXFP8LauncherCallER(policy); + } +#undef MXFP8LauncherCallER } else if (is_weight_fp8) { TORCH_CHECK(ptr_scales.has_value(), "w8a16 grouped gemm must have scales"); TORCH_CHECK(ptr_scales->is_contiguous(), "ptr_scales must be contiguous"); diff --git a/megablocks/torch-ext/megablocks/xpu_fused_moe.py b/megablocks/torch-ext/megablocks/xpu_fused_moe.py index 0440b484e..cf527bf0e 100644 --- a/megablocks/torch-ext/megablocks/xpu_fused_moe.py +++ b/megablocks/torch-ext/megablocks/xpu_fused_moe.py @@ -26,7 +26,7 @@ def _register_if_available(op_name, fn): _register_if_available( "cutlass_grouped_gemm_interface", - lambda ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, expert_first_token_offset, N, K, num_experts, is_B_int4, is_B_mxfp4: ptr_D, + lambda ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, expert_first_token_offset, N, K, num_experts, is_B_int4, is_B_mxfp4, is_B_mxfp8: ptr_D, ) _register_if_available( @@ -106,12 +106,13 @@ def exclusive_prefix_sum(arr): K=k, num_experts=num_experts, is_B_int4=False, - is_B_mxfp4=False) + is_B_mxfp4=False, + is_B_mxfp8=False) def cutlass_grouped_gemm_xe2(input_A, input_B, scales, bias, output, num_rows_per_expert, n, k, num_experts, is_B_int4, - is_B_mxfp4): + is_B_mxfp4, is_B_mxfp8=False): expert_first_token_offset = torch.cat([ torch.tensor([0], dtype=num_rows_per_expert.dtype, @@ -129,7 +130,8 @@ def cutlass_grouped_gemm_xe2(input_A, input_B, scales, bias, output, K=k, num_experts=num_experts, is_B_int4=is_B_int4, - is_B_mxfp4=is_B_mxfp4) + is_B_mxfp4=is_B_mxfp4, + is_B_mxfp8=is_B_mxfp8) def ceilDiv(a, b): @@ -206,7 +208,8 @@ def xpu_fused_moe(hidden_states, ep_size=1, is_fp8=False, is_int4=False, - is_mxfp4=False): + is_mxfp4=False, + is_mxfp8=False): ''' hidden_states: [num_rows, hidden_size] w13: [num_experts, 2*inter_size, hidden_size] @@ -228,6 +231,7 @@ def xpu_fused_moe(hidden_states, num_experts: int is_int4: bool is_mxfp4: bool + is_mxfp8: bool ''' # Resolve DTensors to local tensors before passing to custom ops @@ -266,7 +270,7 @@ def xpu_fused_moe(hidden_states, # 4bits support [E, N, K] # other types [E, K, N] - if not is_int4 and not is_mxfp4: + if not is_int4 and not is_mxfp4 and not is_mxfp8: if not hasattr(w13, 'xpu_fused_moe'): if needs_transpose: w13.data = w13.transpose(-1, -2).contiguous() @@ -389,7 +393,7 @@ def config_ws(name, size): ########### gemm1 ################## input_B = w13 - if not is_fp8 and not is_int4 and not is_mxfp4: + if not is_fp8 and not is_int4 and not is_mxfp4 and not is_mxfp8: ops.cutlass_grouped_gemm_interface( ptr_A=gemm1_input, ptr_B=input_B, @@ -401,7 +405,8 @@ def config_ws(name, size): K=hidden_size, num_experts=num_experts_per_node, is_B_int4=is_int4, - is_B_mxfp4=is_mxfp4) + is_B_mxfp4=is_mxfp4, + is_B_mxfp8=is_mxfp8) else: ops.cutlass_grouped_gemm_interface( ptr_A=gemm1_input, @@ -414,7 +419,8 @@ def config_ws(name, size): K=hidden_size, num_experts=num_experts_per_node, is_B_int4=is_int4, - is_B_mxfp4=is_mxfp4) + is_B_mxfp4=is_mxfp4, + is_B_mxfp8=is_mxfp8) # act act_output = torch.empty((num_moe_inputs, inter_size), @@ -435,7 +441,7 @@ def config_ws(name, size): gemm2_output = torch.empty((num_moe_inputs, hidden_size), dtype=hidden_states.dtype, device=hidden_states.device) - if not is_fp8 and not is_int4 and not is_mxfp4: + if not is_fp8 and not is_int4 and not is_mxfp4 and not is_mxfp8: ops.cutlass_grouped_gemm_interface( ptr_A=input_A, ptr_B=input_B, @@ -447,7 +453,8 @@ def config_ws(name, size): K=inter_size, num_experts=num_experts_per_node, is_B_int4=is_int4, - is_B_mxfp4=is_mxfp4) + is_B_mxfp4=is_mxfp4, + is_B_mxfp8=is_mxfp8) else: ops.cutlass_grouped_gemm_interface( ptr_A=input_A, @@ -460,7 +467,8 @@ def config_ws(name, size): K=inter_size, num_experts=num_experts_per_node, is_B_int4=is_int4, - is_B_mxfp4=is_mxfp4) + is_B_mxfp4=is_mxfp4, + is_B_mxfp8=is_mxfp8) ops.moe_gather(output, gemm2_output, topk_weights, permuted_row_to_unpermuted_row, @@ -609,6 +617,7 @@ def forward(self, x: torch.Tensor) -> tuple: is_fp8 = getattr(self.experts, "is_fp8", False) is_int4 = getattr(self.experts, "is_int4", False) is_mxfp4 = getattr(self.experts, "is_mxfp4", False) + is_mxfp8 = getattr(self.experts, "is_mxfp8", False) w13_scales = getattr(self.experts, "gate_up_proj_scales", None) w2_scales = getattr(self.experts, "down_proj_scales", None) @@ -650,6 +659,7 @@ def forward(self, x: torch.Tensor) -> tuple: is_fp8=is_fp8, is_int4=is_int4, is_mxfp4=is_mxfp4, + is_mxfp8=is_mxfp8, ) # All-reduce across EP group to combine partial expert outputs diff --git a/megablocks/torch-ext/torch_binding.cpp b/megablocks/torch-ext/torch_binding.cpp index 8ea4db20c..fff510eab 100644 --- a/megablocks/torch-ext/torch_binding.cpp +++ b/megablocks/torch-ext/torch_binding.cpp @@ -263,7 +263,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "cutlass_grouped_gemm_interface(Tensor ptr_A, Tensor ptr_B, " "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"); ops.impl("cutlass_grouped_gemm_interface", torch::kXPU, &cutlass_grouped_gemm_interface); } From 725a6030b84a79e7105a9d7869201b6d39e1bb23 Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Tue, 28 Jul 2026 01:47:05 -0700 Subject: [PATCH 2/6] megablocks: build the XPU grouped GEMM twice and dispatch CRI at runtime 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. --- megablocks/build.toml | 27 +++++++++++++++ .../grouped_gemm/grouped_gemm_interface.cpp | 11 +++++-- .../csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp | 3 ++ .../grouped_gemm/xe_2/gemm_xe2_policy.hpp | 3 ++ .../grouped_gemm/xe_2/grouped_gemm_xe2.cpp | 30 ++--------------- .../grouped_gemm/xe_2/grouped_gemm_xe2.h | 9 ++++- .../grouped_gemm/xe_2/grouped_gemm_xe2.hpp | 2 ++ .../xe_2/grouped_gemm_xe2_cri.cpp | 6 ++++ .../xe_2/grouped_gemm_xe2_interface.hpp | 29 +++++++++++----- .../grouped_gemm/xe_2/xe2_target_ns.hpp | 16 +++++++++ megablocks/csrc_xpu/xpu_features.hpp | 33 +++++++++++++++++++ 11 files changed, 130 insertions(+), 39 deletions(-) create mode 100644 megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_cri.cpp create mode 100644 megablocks/csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp create mode 100644 megablocks/csrc_xpu/xpu_features.hpp diff --git a/megablocks/build.toml b/megablocks/build.toml index 024eeeba3..327142582 100644 --- a/megablocks/build.toml +++ b/megablocks/build.toml @@ -91,6 +91,7 @@ depends = [ src = [ "csrc_xpu/utils.h", "csrc_xpu/dispatch_utils.h", + "csrc_xpu/xpu_features.hpp", "csrc_xpu/core/registration.h", "csrc_xpu/activation.cpp", "csrc_xpu/activation.h", @@ -111,6 +112,32 @@ src = [ "csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp", "csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp", "csrc_xpu/grouped_gemm/xe_2/gemm_xe2_policy.hpp", + "csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp", +] + +# The Xe35 block-scaled DPAS path in sycl-tla is gated behind +# __SYCL_TARGET_INTEL_GPU_CRI__, which the compiler never sets on its own, so +# build a CRI-only TU with the macro forced on and spir64_gen as the sole +# target. The link step AOT-lowers that image to cri; pvc/bmg keep the Xe20 TU +# above and the host dispatches by device IP version at runtime. +[kernel.megablocks_xpu_cri] +backend = "xpu" +depends = [ + "torch", + "sycl_tla", +] +sycl-flags = [ + "-fsycl-targets=spir64_gen", + "-D__SYCL_TARGET_INTEL_GPU_CRI__=1", +] +src = [ + "csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_cri.cpp", + "csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.h", + "csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.hpp", + "csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp", + "csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp", + "csrc_xpu/grouped_gemm/xe_2/gemm_xe2_policy.hpp", + "csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp", ] [kernel.megablocks_cpu] diff --git a/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.cpp b/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.cpp index b97adaa6a..1962a8609 100644 --- a/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.cpp +++ b/megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.cpp @@ -1,7 +1,10 @@ #include "grouped_gemm_interface.h" #include +#include "../xpu_features.hpp" #include "xe_2/grouped_gemm_xe2.h" +using megablocks::xpu::XPUFeatures; + torch::Tensor cutlass_grouped_gemm_interface( torch::Tensor ptr_A, torch::Tensor ptr_B, @@ -15,8 +18,12 @@ torch::Tensor cutlass_grouped_gemm_interface( bool is_B_int4, bool is_B_mxfp4, bool is_B_mxfp8) { - // Use XE2 cutlass kernel - return cutlass_grouped_gemm_xe2( + // The grouped GEMMs are built twice: an Xe20 image for pvc/bmg and an Xe35 + // image for CRI. Both cover the same dtypes, so pick the variant by device. + auto gemm = XPUFeatures::isXe35(ptr_A.device().index()) + ? MoE::cutlass_grouped_gemm_xe2<35> + : MoE::cutlass_grouped_gemm_xe2<20>; + return gemm( ptr_A, ptr_B, ptr_scales, diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp index dbe1e09e1..7ab38637d 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2.hpp @@ -44,11 +44,13 @@ #include "cutlass/util/reference/device/gemm_complex.h" #include "cutlass/util/reference/device/tensor_compare.h" #include "cutlass/util/reference/host/tensor_fill.h" +#include "xe2_target_ns.hpp" #pragma clang diagnostic ignored "-Wpass-failed" #pragma clang diagnostic ignored "-Wdeprecated-declarations" namespace MoE { +inline namespace MEGABLOCKS_XE_TARGET_NS { using namespace cute; @@ -422,4 +424,5 @@ CUTE_DEVICE void xe_gemm_4bits( copy(copy_c, tCrC_final_sg_tensor, tCgC); } +} // inline namespace MEGABLOCKS_XE_TARGET_NS } // namespace MoE diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2_policy.hpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2_policy.hpp index 57f320725..aa53aeeba 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2_policy.hpp +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/gemm_xe2_policy.hpp @@ -2,8 +2,10 @@ #include "cute/atom/mma_atom.hpp" #include "cutlass/numeric_types.h" +#include "xe2_target_ns.hpp" namespace MoE { +inline namespace MEGABLOCKS_XE_TARGET_NS { using namespace cute; class xe_gemm_policy_base { @@ -85,4 +87,5 @@ class w4a16_policy_m_32 : public xe_gemm_policy_base { using SGLayout = Layout, Stride<_4, _1, _0>>; }; +} // inline namespace MEGABLOCKS_XE_TARGET_NS } // namespace MoE \ No newline at end of file diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.cpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.cpp index d2a2f9714..298d2fffa 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.cpp +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.cpp @@ -1,31 +1,5 @@ -#include #include "grouped_gemm_xe2.h" #include "grouped_gemm_xe2_interface.hpp" -torch::Tensor cutlass_grouped_gemm_xe2( - torch::Tensor ptr_A, - torch::Tensor ptr_B, - const c10::optional& ptr_scales, - const c10::optional& ptr_bias, - torch::Tensor ptr_D, - torch::Tensor expert_first_token_offset, - int64_t N, - int64_t K, - int64_t num_experts, - bool is_B_int4, - bool is_B_mxfp4, - bool is_B_mxfp8) { - return MoE::cutlass_grouped_gemm_xe2_impl( - ptr_A, - ptr_B, - ptr_scales, - ptr_bias, - ptr_D, - expert_first_token_offset, - N, - K, - num_experts, - is_B_int4, - is_B_mxfp4, - is_B_mxfp8); -} \ No newline at end of file +// Xe20 image, used by pvc and bmg. +template decltype(MoE::cutlass_grouped_gemm_xe2<20>) MoE::cutlass_grouped_gemm_xe2<20>; diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.h b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.h index ea2d9a362..aba36a0db 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.h +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.h @@ -1,5 +1,10 @@ #include +namespace MoE { + +// Instantiated once per target architecture: 20 for pvc/bmg, 35 for CRI. The +// two instantiations live in separate translation units with different flags. +template torch::Tensor cutlass_grouped_gemm_xe2( torch::Tensor ptr_A, torch::Tensor ptr_B, @@ -12,4 +17,6 @@ torch::Tensor cutlass_grouped_gemm_xe2( int64_t num_experts, bool is_B_int4, bool is_B_mxfp4, - bool is_B_mxfp8); \ No newline at end of file + bool is_B_mxfp8); + +} // namespace MoE diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.hpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.hpp index 0e8f6d458..21c374894 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.hpp +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.hpp @@ -45,6 +45,7 @@ #pragma clang diagnostic ignored "-Wdeprecated-declarations" namespace MoE { +inline namespace MEGABLOCKS_XE_TARGET_NS { using namespace cute; template @@ -231,4 +232,5 @@ CUTE_DEVICE void MoEGEMM( } } +} // inline namespace MEGABLOCKS_XE_TARGET_NS } // namespace MoE diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_cri.cpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_cri.cpp new file mode 100644 index 000000000..fd9da9b3e --- /dev/null +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_cri.cpp @@ -0,0 +1,6 @@ +#include "grouped_gemm_xe2.h" +#include "grouped_gemm_xe2_interface.hpp" + +// Xe35 image. Built with -D__SYCL_TARGET_INTEL_GPU_CRI__=1 and +// -fsycl-targets=spir64_gen so sycl-tla emits the CRI block-scaled DPAS path. +template decltype(MoE::cutlass_grouped_gemm_xe2<35>) MoE::cutlass_grouped_gemm_xe2<35>; diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp index 5ab94efdd..95307a20a 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp @@ -58,15 +58,18 @@ #include "gemm_xe2_policy.hpp" #include "grouped_gemm_xe2.hpp" +#include "xe2_target_ns.hpp" #pragma clang diagnostic ignored "-Wpass-failed" #pragma clang diagnostic ignored "-Wdeprecated-declarations" namespace MoE { +inline namespace MEGABLOCKS_XE_TARGET_NS { using namespace cute; -// type tag to define a unique sycl kernel name -template +// type tag to define a unique sycl kernel name. The trailing arch tag keeps the +// CRI (Xe35) and generic (Xe20) device images from sharing a kernel name. +template class GemmCuteName; template < @@ -135,7 +138,8 @@ void MoEGEMMLauncher( ElementD, layoutA, layoutB, - policy>>( + policy, + SYCL_INTEL_TARGET>>( sycl::nd_range<3>{global * local, local}, kernel_props, [=](auto) { MoE::MoEGEMM< GmemTiledCopyA, @@ -162,19 +166,28 @@ void MoEGEMMLauncher( EventManager::getInstance().addEvent(event); } -at::Tensor cutlass_grouped_gemm_xe2_impl( - at::Tensor& ptr_A, - at::Tensor& ptr_B, +} // inline namespace MEGABLOCKS_XE_TARGET_NS + +using namespace cute; + +// Entry point for one architecture, instantiated as <20> by the pvc/bmg +// translation unit and as <35> by the CRI one. +template +torch::Tensor cutlass_grouped_gemm_xe2( + torch::Tensor ptr_A, + torch::Tensor ptr_B, const c10::optional& ptr_scales, const c10::optional& ptr_bias, - at::Tensor& ptr_D, - at::Tensor& expert_first_token_offset, + torch::Tensor ptr_D, + torch::Tensor expert_first_token_offset, int64_t N, int64_t K, int64_t num_experts, bool is_B_int4, bool is_B_mxfp4, bool is_B_mxfp8) { + // Guards against the CRI translation unit losing -D__SYCL_TARGET_INTEL_GPU_CRI__. + static_assert(Arch == SYCL_INTEL_TARGET, "Arch must match SYCL_INTEL_TARGET"); auto& dpcpp_queue = at::xpu::getCurrentXPUStream(ptr_A.device().index()).queue(); auto A_dtype = ptr_A.dtype(); diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp new file mode 100644 index 000000000..23ed0af3c --- /dev/null +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp @@ -0,0 +1,16 @@ +/***************************************************************************************** + * Copyright (C) 2026 Intel Corporation, All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + ****************************************************************************************/ + +#pragma once + +// These headers are compiled twice, once per target architecture, and sycl-tla +// expands them into different code each time. Keeping the two expansions in +// separate namespaces stops the linker from merging symbols that share a name +// but not a body. +#if defined(__SYCL_TARGET_INTEL_GPU_CRI__) +#define MEGABLOCKS_XE_TARGET_NS xe35 +#else +#define MEGABLOCKS_XE_TARGET_NS xe20 +#endif diff --git a/megablocks/csrc_xpu/xpu_features.hpp b/megablocks/csrc_xpu/xpu_features.hpp new file mode 100644 index 000000000..ad7a8bfc4 --- /dev/null +++ b/megablocks/csrc_xpu/xpu_features.hpp @@ -0,0 +1,33 @@ +/***************************************************************************************** + * Copyright (C) 2026 Intel Corporation, All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + ****************************************************************************************/ + +// XPU Device Feature Detection for MegaBlocks + +#pragma once + +#include +#include + +namespace megablocks { +namespace xpu { + +// XPU device feature detection, keyed on the device IP version reported by the +// driver (pvc 12.x, bmg 20.x, CRI 35.x). +class XPUFeatures { + public: + // True when the device is Xe35 (CRI) or newer, which selects the kernels + // built from the CRI-only translation unit. + static bool isXe35(c10::DeviceIndex device) { + return ipVersionMajor(device) >= 35; + } + + // Major component of the device IP version, e.g. 20 for "20.1.0". + static int ipVersionMajor(c10::DeviceIndex device) { + return std::atoi(at::xpu::getDeviceProperties(device)->version.c_str()); + } +}; + +} // namespace xpu +} // namespace megablocks From f3d3ac7af2693afecdee97b2d742e9b23c9671f7 Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Tue, 11 Aug 2026 03:31:32 +0000 Subject: [PATCH 3/6] megablocks: correct the CRI TU comment 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. --- megablocks/build.toml | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/megablocks/build.toml b/megablocks/build.toml index 327142582..09bcf8f16 100644 --- a/megablocks/build.toml +++ b/megablocks/build.toml @@ -115,11 +115,12 @@ src = [ "csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp", ] -# The Xe35 block-scaled DPAS path in sycl-tla is gated behind -# __SYCL_TARGET_INTEL_GPU_CRI__, which the compiler never sets on its own, so -# build a CRI-only TU with the macro forced on and spir64_gen as the sole -# target. The link step AOT-lowers that image to cri; pvc/bmg keep the Xe20 TU -# above and the host dispatches by device IP version at runtime. +# The Xe35 code paths in sycl-tla (Xe35 block 2D cache controls, native fcvt for +# e4m3 -> bf16 reorder) are gated behind SYCL_INTEL_TARGET == 35, which cutlass.h +# only derives from __SYCL_TARGET_INTEL_GPU_CRI__. The compiler never sets that +# macro on its own, so build a CRI-only TU with it forced on and spir64_gen as +# the sole target. The link step AOT-lowers that image to cri; pvc/bmg keep the +# Xe20 TU above and the host dispatches by device IP version at runtime. [kernel.megablocks_xpu_cri] backend = "xpu" depends = [ From f6b068bd8b380f39650237d27070407351e82ee0 Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Tue, 25 Aug 2026 06:28:42 +0000 Subject: [PATCH 4/6] megablocks: default is_B_mxfp8 to False in the op schema 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. --- megablocks/torch-ext/megablocks/xpu_fused_moe.py | 2 +- megablocks/torch-ext/torch_binding.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/megablocks/torch-ext/megablocks/xpu_fused_moe.py b/megablocks/torch-ext/megablocks/xpu_fused_moe.py index cf527bf0e..48f66d35b 100644 --- a/megablocks/torch-ext/megablocks/xpu_fused_moe.py +++ b/megablocks/torch-ext/megablocks/xpu_fused_moe.py @@ -26,7 +26,7 @@ def _register_if_available(op_name, fn): _register_if_available( "cutlass_grouped_gemm_interface", - lambda ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, expert_first_token_offset, N, K, num_experts, is_B_int4, is_B_mxfp4, is_B_mxfp8: ptr_D, + lambda ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, expert_first_token_offset, N, K, num_experts, is_B_int4, is_B_mxfp4, is_B_mxfp8=False: ptr_D, ) _register_if_available( diff --git a/megablocks/torch-ext/torch_binding.cpp b/megablocks/torch-ext/torch_binding.cpp index fff510eab..bb98753da 100644 --- a/megablocks/torch-ext/torch_binding.cpp +++ b/megablocks/torch-ext/torch_binding.cpp @@ -263,7 +263,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "cutlass_grouped_gemm_interface(Tensor ptr_A, Tensor ptr_B, " "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, bool is_B_mxfp8) -> Tensor"); + "bool is_B_int4, bool is_B_mxfp4, bool is_B_mxfp8=False) -> Tensor"); ops.impl("cutlass_grouped_gemm_interface", torch::kXPU, &cutlass_grouped_gemm_interface); } From 7d6f3da0323cb38bd5c19e3bd9aced074fd76eb5 Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Tue, 25 Aug 2026 09:43:42 +0000 Subject: [PATCH 5/6] megablocks: derive the XPU arch tag from the CRI macro, not SYCL_INTEL_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. --- .../grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp | 6 ++++-- megablocks/csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp | 7 +++++++ 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp index 95307a20a..a2ba7f923 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp @@ -139,7 +139,7 @@ void MoEGEMMLauncher( layoutA, layoutB, policy, - SYCL_INTEL_TARGET>>( + MEGABLOCKS_XE_TARGET_ARCH>>( sycl::nd_range<3>{global * local, local}, kernel_props, [=](auto) { MoE::MoEGEMM< GmemTiledCopyA, @@ -187,7 +187,9 @@ torch::Tensor cutlass_grouped_gemm_xe2( bool is_B_mxfp4, bool is_B_mxfp8) { // Guards against the CRI translation unit losing -D__SYCL_TARGET_INTEL_GPU_CRI__. - static_assert(Arch == SYCL_INTEL_TARGET, "Arch must match SYCL_INTEL_TARGET"); + static_assert( + Arch == MEGABLOCKS_XE_TARGET_ARCH, + "Arch must match MEGABLOCKS_XE_TARGET_ARCH"); auto& dpcpp_queue = at::xpu::getCurrentXPUStream(ptr_A.device().index()).queue(); auto A_dtype = ptr_A.dtype(); diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp index 23ed0af3c..b893d4d06 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp @@ -9,8 +9,15 @@ // expands them into different code each time. Keeping the two expansions in // separate namespaces stops the linker from merging symbols that share a name // but not a body. +// +// MEGABLOCKS_XE_TARGET_ARCH names the same split numerically, for the arch +// guard and the SYCL kernel name tag. Do not use sycl-tla's SYCL_INTEL_TARGET +// for this: only 0.9.2 and later rewrite it to 20/35 in cutlass.h, while older +// releases leave it at the 1 that the build system defines. #if defined(__SYCL_TARGET_INTEL_GPU_CRI__) #define MEGABLOCKS_XE_TARGET_NS xe35 +#define MEGABLOCKS_XE_TARGET_ARCH 35 #else #define MEGABLOCKS_XE_TARGET_NS xe20 +#define MEGABLOCKS_XE_TARGET_ARCH 20 #endif From cdfab705a341e99081ed6d9439df75132f128804 Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Tue, 25 Aug 2026 09:46:00 +0000 Subject: [PATCH 6/6] megablocks: trim the Xe target macro comment --- megablocks/csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/megablocks/csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp b/megablocks/csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp index b893d4d06..9f98e65e7 100644 --- a/megablocks/csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp @@ -9,13 +9,9 @@ // expands them into different code each time. Keeping the two expansions in // separate namespaces stops the linker from merging symbols that share a name // but not a body. -// -// MEGABLOCKS_XE_TARGET_ARCH names the same split numerically, for the arch -// guard and the SYCL kernel name tag. Do not use sycl-tla's SYCL_INTEL_TARGET -// for this: only 0.9.2 and later rewrite it to 20/35 in cutlass.h, while older -// releases leave it at the 1 that the build system defines. #if defined(__SYCL_TARGET_INTEL_GPU_CRI__) #define MEGABLOCKS_XE_TARGET_NS xe35 +// Not sycl-tla's SYCL_INTEL_TARGET: only 0.9.2+ rewrites that to 20/35. #define MEGABLOCKS_XE_TARGET_ARCH 35 #else #define MEGABLOCKS_XE_TARGET_NS xe20