Megablocks mxfp8 cri

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]

megablocks/csrc_xpu/grouped_gemm/grouped_gemm_interface.cpp

@@ -1,7 +1,10 @@
 #include "grouped_gemm_interface.h"
 #include <stdio.h>
+#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);
 }

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);
+    bool is_B_mxfp4,
+    bool is_B_mxfp8);

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<TB, float_e2m1_t>) {
+          } else if constexpr (std::is_same_v<TB, float_e2m1_t> ||
+                               std::is_same_v<TB, float_e4m3_t> ||
+                               std::is_same_v<TB, float_e5m2_t>) {
             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

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<Shape<_1, _4, _1>, Stride<_4, _1, _0>>;
 };
 
+}  // inline namespace MEGABLOCKS_XE_TARGET_NS
 }  // namespace MoE

megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.cpp

@@ -1,29 +1,5 @@
-#include <torch/all.h>
 #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<at::Tensor>& ptr_scales,
-    const c10::optional<at::Tensor>& 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);
-}
+// Xe20 image, used by pvc and bmg.
+template decltype(MoE::cutlass_grouped_gemm_xe2<20>) MoE::cutlass_grouped_gemm_xe2<20>;

megablocks/csrc_xpu/grouped_gemm/xe_2/grouped_gemm_xe2.h

@@ -1,5 +1,10 @@
 #include <torch/all.h>
 
+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 <int Arch>
 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);
+    bool is_B_mxfp4,
+    bool is_B_mxfp8);
+
+}  // namespace MoE

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 <typename T, char LayoutKind>
@@ -90,7 +91,12 @@ CUTE_DEVICE void MoEGEMM(
                                     (!std::is_same_v<ElementS, uint8_t>);
   static constexpr bool is_B_mxfp4 = (std::is_same_v<ElementB, uint8_t>) &&
                                      (std::is_same_v<ElementS, uint8_t>);
+  static constexpr bool is_B_mxfp8 =
+      (std::is_same_v<ElementB, cutlass::float_e4m3_t> ||
+       std::is_same_v<ElementB, cutlass::float_e5m2_t>) &&
+      (std::is_same_v<ElementS, uint8_t>);
   static constexpr bool is_B_4bits = std::is_same_v<ElementB, uint8_t>;
+  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<ElementS*>(Scales) + B_offset * 2 / group_size;
+    } else if constexpr (is_B_mxfp8) {
+      ptr_Scales_curr_batch =
+          const_cast<ElementS*>(Scales) +
+          static_cast<int64_t>(expert_id) *
+              static_cast<int64_t>(gemm_n) *
+              (static_cast<int64_t>(gemm_k) / group_size);
     }
     ElementBI* ptr_Bias_curr_batch = nullptr;
     if (Bias != static_cast<ElementBI*>(nullptr)) {
@@ -161,6 +173,9 @@ CUTE_DEVICE void MoEGEMM(
       } else if constexpr (is_B_mxfp4) {
         return make_moe_tensor<float_e2m1_t, actual_layout_of_B>(
             reinterpret_cast<float_e2m1_t*>(ptr_B_curr_batch), gemm_n, gemm_k);
+      } else if constexpr (is_B_mxfp8) {
+        return make_moe_tensor<ElementB, actual_layout_of_B>(
+            ptr_B_curr_batch, gemm_n, gemm_k);
       } else {
         return make_moe_tensor<ElementB, actual_layout_of_B>(
             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<GmemTiledCopyA, GmemTiledCopyB, GmemTiledCopyD, GroupSize>( \
       A_tensor,                                                             \
@@ -217,4 +232,5 @@ CUTE_DEVICE void MoEGEMM(
   }
 }
 
+}  // inline namespace MEGABLOCKS_XE_TARGET_NS
 }  // namespace MoE

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>;