[CPU] Implement MXFP4 Gemm kernels for intel AMX to support GPT OSS series. (#14385)

2026-03-29 23:44:12 -07:00
parent ed01e1d5d6
commit af62bd9486
6 changed files with 555 additions and 92 deletions
@@ -45,7 +45,7 @@ namespace {
    }                                                                    \
  }()

-// dispatch: bfloat16, float16, int8_t, fp8_e4m3
+// dispatch: bfloat16, float16, int8_t, fp8_e4m3, uint8_t(mxfp4/int4)
 #define CPU_DISPATCH_PACKED_TYPES(TYPE, ...)                     \
  [&] {                                                          \
    switch (TYPE) {                                              \
@@ -65,6 +65,10 @@ namespace {
        using packed_t = at::Float8_e4m3fn;                      \
        return __VA_ARGS__();                                    \
      }                                                          \
+      case at::ScalarType::Byte: {                               \
+        using packed_t = uint8_t;                                \
+        return __VA_ARGS__();                                    \
+      }                                                          \
      default:                                                   \
        TORCH_CHECK(false, "Unsupported floating data type.\n"); \
    }                                                            \