diff --git a/megablocks/build.toml b/megablocks/build.toml index 024eeeba3..09bcf8f16 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,33 @@ 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 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 = [ + "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 bf6ccf6e8..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, @@ -13,9 +16,14 @@ torch::Tensor cutlass_grouped_gemm_interface( int64_t K, int64_t num_experts, bool is_B_int4, - bool is_B_mxfp4) { - // Use XE2 cutlass kernel - return cutlass_grouped_gemm_xe2( + bool is_B_mxfp4, + bool is_B_mxfp8) { + // 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, @@ -26,5 +34,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..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; @@ -366,7 +368,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 + @@ -420,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 2d6c4b29a..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,29 +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) { - 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); -} \ 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 feb722c6e..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, @@ -11,4 +16,7 @@ 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); + +} // 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 be7386409..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 @@ -90,7 +91,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 +152,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 +173,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 +189,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, \ @@ -217,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 6a0a6d312..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 @@ -50,23 +50,26 @@ #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" +#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, + MEGABLOCKS_XE_TARGET_ARCH>>( sycl::nd_range<3>{global * local, local}, kernel_props, [=](auto) { MoE::MoEGEMM< GmemTiledCopyA, @@ -162,18 +166,30 @@ 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_mxfp4, + bool is_B_mxfp8) { + // Guards against the CRI translation unit losing -D__SYCL_TARGET_INTEL_GPU_CRI__. + 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(); @@ -206,6 +222,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 +318,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/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..9f98e65e7 --- /dev/null +++ b/megablocks/csrc_xpu/grouped_gemm/xe_2/xe2_target_ns.hpp @@ -0,0 +1,19 @@ +/***************************************************************************************** + * 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 +// 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 +#define MEGABLOCKS_XE_TARGET_ARCH 20 +#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 diff --git a/megablocks/torch-ext/megablocks/xpu_fused_moe.py b/megablocks/torch-ext/megablocks/xpu_fused_moe.py index 0440b484e..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: 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( @@ -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..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) -> 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); }