More MoE GEMMs

Author: sanchitintel

examples/12_bmg_moe_gemm_cute_interface/moe_gemms.hpp

@@ -160,8 +160,10 @@ template <class GmemTiledCopyA, class GmemTiledCopyB, class GmemTiledCopyD,
           class STensor, class DTensor, class TiledMMA,
           class = std::enable_if_t<
               !cute::is_void_v<typename STensor::element_type> &&
-              is_same_v<typename BTensor::element_type, float_e2m1_t> &&
-              is_same_v<typename STensor::element_type, float_ue8m0_t> &&
+              is_any_of_v<typename BTensor::element_type, float_e2m1_t,
+                          float_e4m3_t, float_e5m2_t, int4_t> &&
+              is_any_of_v<typename STensor::element_type, float_ue8m0_t, half_t,
+                          bfloat16_t> &&
               is_any_of_v<typename ATensor::element_type, bfloat16_t, half_t>>>
 CUTE_DEVICE void moe_gemm(ATensor const &A, // (M,K)
                           BTensor const &B, // (N,K)