support different accumulator dtype combinations

Author: wuxun-zhang

include/cutlass/epilogue/collective/xe_array_epilogue.hpp

@@ -90,9 +90,8 @@ class CollectiveEpilogue<
   using DispatchPolicy = IntelXeXMX16Group;
   using CtaTileMNK = CtaTileMNK_;
   using FusionCallbacks = FusionCallbacks_;
-  using ElementC = ElementD_;
-  // simple use fp32 as accumulator dtype
-  using ElementAccumulator = float;
+  using ElementC = typename FusionCallbacks::ElementSource;
+  using ElementAccumulator = ElementC_;
   using StrideC = StrideC_;
   using InternalStrideC = cute::remove_pointer_t<StrideC>;
   using ElementD = ElementD_;

include/cutlass/epilogue/collective/xe_epilogue.hpp

@@ -89,10 +89,8 @@ class CollectiveEpilogue<
   using DispatchPolicy = IntelXeXMX16;
   using CtaTileMNK = CtaTileMNK_;
   using FusionCallbacks = FusionCallbacks_;
-  // assume C and D have same dtype
-  using ElementC = ElementD_;
-  // simple use fp32 as accumulator dtype
-  using ElementAccumulator = float;
+  using ElementC = typename FusionCallbacks::ElementSource;;
+  using ElementAccumulator = ElementC_;
   using StrideC = StrideC_;
   using ElementD = ElementD_;
   using StrideD = StrideD_;