Add new falsh attention fp8 support on BMG (#613)

## Description
Add FP8 support to new flash attention api examples and fix the acc
issue that caused by reorder

---------

Signed-off-by: Chen, Xi2 <[email protected]>

Author: ClarkChin08

examples/06_bmg_flash_attention/06_xe_fmha_fwd.cpp

@@ -147,6 +147,18 @@ int main(int argc, const char **argv) {
 #else
   constexpr int PipelineStages = 2;
 #endif
-
-  return FMHAConfig<false, ShapeQK, ShapePV, ShapeOut, SubgroupLayoutQK, void, PipelineStages>::run(options);
+#ifdef IS_FLOAT_E5M2
+  using ElementQ = cutlass::float_e5m2_t;
+  using ElementK = cutlass::float_e5m2_t;
+  using ElementV = cutlass::float_e5m2_t;
+#elif defined(IS_FLOAT_E4M3)
+  using ElementQ = cutlass::float_e4m3_t;
+  using ElementK = cutlass::float_e4m3_t;
+  using ElementV = cutlass::float_e4m3_t;
+#else
+  using ElementQ = bfloat16_t;
+  using ElementK = bfloat16_t;
+  using ElementV = bfloat16_t;
+#endif
+  return FMHAConfig<false, ShapeQK, ShapePV, ShapeOut, SubgroupLayoutQK, void, PipelineStages, ElementQ, ElementK, ElementV>::run(options);
 }

examples/06_bmg_flash_attention/CMakeLists.txt

@@ -33,16 +33,27 @@ set(TEST_NO_PAGED "")
 set(TEST_PAGED "--use_paged_kv")
 
 foreach(HEAD_DIM 64 96 128 192)
+  foreach(INPUT_TYPE bfloat16_t float_e5m2_t float_e4m3_t)
+    cutlass_example_add_executable(
+      06_xe_fmha_fwd_prefill_${INPUT_TYPE}_hdim${HEAD_DIM}
+      06_xe_fmha_fwd.cpp
+    )
 
-  cutlass_example_add_executable(
-    06_xe_fmha_fwd_prefill_hdim${HEAD_DIM}
-    06_xe_fmha_fwd.cpp
-  )
+    cutlass_example_add_executable(
+      06_xe_fmha_fwd_decode_${INPUT_TYPE}_hdim${HEAD_DIM}
+      06_xe_fmha_fwd.cpp
+    )
+    if(INPUT_TYPE STREQUAL "bfloat16_t")
+      set(INPUT_MACRO "IS_BFLOAT16")
+    elseif(INPUT_TYPE STREQUAL "float_e5m2_t")
+      set(INPUT_MACRO "IS_FLOAT_E5M2")
+    elseif(INPUT_TYPE STREQUAL "float_e4m3_t")
+      set(INPUT_MACRO "IS_FLOAT_E4M3")
+    endif()
 
-  cutlass_example_add_executable(
-    06_xe_fmha_fwd_decode_hdim${HEAD_DIM}
-    06_xe_fmha_fwd.cpp
-  )
+    target_compile_definitions(06_xe_fmha_fwd_prefill_${INPUT_TYPE}_hdim${HEAD_DIM} PRIVATE HEAD_DIM=${HEAD_DIM} PREFILL SHOW_DIFF=1 INPUT_TYPE=${INPUT_TYPE} ${INPUT_MACRO})
+    target_compile_definitions(06_xe_fmha_fwd_decode_${INPUT_TYPE}_hdim${HEAD_DIM} PRIVATE HEAD_DIM=${HEAD_DIM} DECODE SHOW_DIFF=1 INPUT_TYPE=${INPUT_TYPE} ${INPUT_MACRO})
+  endforeach()
 
   cutlass_example_add_executable(
     06_bmg_prefill_attention_hdim${HEAD_DIM}
@@ -82,6 +93,4 @@ foreach(HEAD_DIM 64 96 128 192)
   target_compile_definitions(06_bmg_decode_attention_hdim${HEAD_DIM} PRIVATE HEAD_DIM=${HEAD_DIM})
   target_compile_definitions(06_bmg_prefill_attention_fp8_hdim${HEAD_DIM} PRIVATE HEAD_DIM=${HEAD_DIM})
   target_compile_definitions(06_bmg_decode_attention_fp8_hdim${HEAD_DIM} PRIVATE HEAD_DIM=${HEAD_DIM})
-  target_compile_definitions(06_xe_fmha_fwd_prefill_hdim${HEAD_DIM} PRIVATE HEAD_DIM=${HEAD_DIM} PREFILL SHOW_DIFF=1)
-  target_compile_definitions(06_xe_fmha_fwd_decode_hdim${HEAD_DIM} PRIVATE HEAD_DIM=${HEAD_DIM} DECODE SHOW_DIFF=1)
 endforeach()

examples/06_bmg_flash_attention/xe_fmha_fwd_runner.hpp

@@ -359,7 +359,7 @@ template <class FMHAKernel> struct ExampleRunner {
 
     // Check if output from CUTLASS kernel and reference kernel are equal or not
     bool passed = cutlass::reference::device::BlockCompareRelativelyEqual(block_ref_O.get(), block_O.get(),
-                                                                          block_O.size(), ElementO{0.005}, ElementO{0.005});
+                                                                          block_O.size(), ElementO{0.05}, ElementO{0.05});
 
     return passed;
   }
@@ -531,7 +531,11 @@ struct FMHAConfig {
 
   static constexpr int SGTileQ = get<0>(shape_div(TileShapeQK{}, shape(SubgroupLayoutQK{})))();
   using MMAOperation = cute::conditional_t<is_void_v<MMAOperation_>,
-                                           XE_DPAS_TT<cute::gcd(SGTileQ, 8), float, ElementQ>,
+                                           typename cute::conditional_t<
+                                               cute::is_same_v<ElementQ, cutlass::float_e5m2_t> || cute::is_same_v<ElementQ, cutlass::float_e4m3_t>,
+                                               XE_DPAS_TT<cute::gcd(SGTileQ, 8), float, half_t>,
+                                               XE_DPAS_TT<cute::gcd(SGTileQ, 8), float, ElementQ> 
+                                           >,
                                            MMAOperation_>;
   using SubgroupLayoutPV = cute::conditional_t<is_void_v<SubgroupLayoutPV_>,
                                                decltype(cutlass::fmha::collective::get_sg_layout_pv(SubgroupLayoutQK{})),

include/cute/atom/copy_traits_xe_2d.hpp

@@ -1216,7 +1216,7 @@ make_block_2d_prefetch(const Shape&, Stride<Strides...> const& stride, const XMo
   constexpr auto shape_y = get<YMode::value>(Shape{});
 
   // Try to retrieve whole cache lines (contiguous dimension = x)
-  constexpr auto width = cute::min(shape_x, 512 / sizeof_bits_v<ValType>);
+  constexpr auto width = cute::gcd(shape_x, 512 / sizeof_bits_v<ValType>);
 
   // Do a preliminary tiling to choose appropriate height.
   constexpr int n_sg_x = cute::gcd(SGCount, ceil_div(shape_x, width));