v3.9 (#2185)

* v3.8 update x

* fix blackwell gg

* doc change

* doc change

* doc change

---------

Co-authored-by: yuzhai <yuzhai@nvidia.com>
Co-authored-by: Haicheng Wu <haichengw@nvidia.com>
Co-authored-by: Haicheng Wu <57973641+hwu36@users.noreply.github.com>

tools/library/src/block_scaled_gemm_operation_3x.hpp

@@ -72,7 +72,7 @@ public:
   using CollectiveEpilogue = typename Operator::CollectiveEpilogue;
   using ThreadEpilogueOp = typename CollectiveEpilogue::ThreadEpilogueOp;
 
-  using Sm100BlkScaledConfig =  typename CollectiveMainloop::Sm100BlkScaledConfig;
+  using Sm1xxBlkScaledConfig =  typename CollectiveMainloop::Sm1xxBlkScaledConfig;
     
   static constexpr bool epilogue_scalefactor_generation = not cute::is_same_v<typename ThreadEpilogueOp::ElementBlockScaleFactor, void>;
   static constexpr int32_t SFD_VectorSize = epilogue_scalefactor_generation ? ThreadEpilogueOp::SFVecSize : SFVecSize;
@@ -315,8 +315,8 @@ protected:
         arguments->ldc, arguments->batch_stride_C);
     operator_args.epilogue.dD = operator_args.epilogue.dC;
 
-    operator_args.mainloop.layout_SFA = Sm100BlkScaledConfig::tile_atom_to_shape_SFA(operator_args.problem_shape);
-    operator_args.mainloop.layout_SFB = Sm100BlkScaledConfig::tile_atom_to_shape_SFB(operator_args.problem_shape);
+    operator_args.mainloop.layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(operator_args.problem_shape);
+    operator_args.mainloop.layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(operator_args.problem_shape);
 
     /* Query device SM count to pass onto the kernel as an argument, where needed */
     operator_args.hw_info.sm_count = arguments->sm_count;

tools/library/src/gemm_operation_3x.hpp

@@ -170,13 +170,12 @@ public:
   using CollectiveEpilogue = typename Operator::CollectiveEpilogue;
   using ThreadEpilogueOp = typename CollectiveEpilogue::ThreadEpilogueOp;
 
-  
   static constexpr bool IsRuntimeDataTypeA = cutlass::gemm::collective::detail::is_sm10x_runtime_f8f6f4<ElementA>();
 
   static constexpr bool IsRuntimeDataTypeB = cutlass::gemm::collective::detail::is_sm10x_runtime_f8f6f4<ElementB>();
 
   static_assert((IsRuntimeDataTypeA && IsRuntimeDataTypeB) ||
-                (!IsRuntimeDataTypeA && !IsRuntimeDataTypeB), 
+                (!IsRuntimeDataTypeA && !IsRuntimeDataTypeB),
                 "ElementA and ElementB in a GEMM kernel should be both runtime or both static.");
 
   static constexpr bool IsRuntimeDataType = IsRuntimeDataTypeA && IsRuntimeDataTypeB;
@@ -252,13 +251,13 @@ protected:
   };
 
   template<template<int, class, class> class Policy, int Stages, class ClusterShape, class KernelSchedule>
-  static constexpr bool is_mixed_dtype_mainloop_(Policy<Stages, ClusterShape, KernelSchedule> policy) {
+  static constexpr bool is_sm90_mixed_dtype_mainloop_(Policy<Stages, ClusterShape, KernelSchedule> policy) {
     return (cute::is_same_v<Policy<Stages, ClusterShape, KernelSchedule>,
                             cutlass::gemm::MainloopSm90TmaGmmaRmemAWarpSpecializedMixedInput<Stages, ClusterShape, KernelSchedule>>);
   }
 
   template <class DispatchPolicy>
-  static constexpr bool is_mixed_dtype_mainloop_(DispatchPolicy) {
+  static constexpr bool is_sm90_mixed_dtype_mainloop_(DispatchPolicy) {
     return false;
   }
 
@@ -380,7 +379,6 @@ protected:
       arguments->batch_count);
 
     // update arguments
-    
 
     if constexpr (IsRuntimeDataType) {
       using ArrayElementA = typename Operator::GemmKernel::CollectiveMainloop::ArrayElementA;
@@ -390,7 +388,7 @@ protected:
 
       std::unordered_map<RuntimeDatatype, cute::UMMA::MXF8F6F4Format> mapping = {
           {RuntimeDatatype::kE4M3, cute::UMMA::MXF8F6F4Format::E4M3},
-          {RuntimeDatatype::kE5M2, cute::UMMA::MXF8F6F4Format::E5M2}, 
+          {RuntimeDatatype::kE5M2, cute::UMMA::MXF8F6F4Format::E5M2},
           {RuntimeDatatype::kE3M2, cute::UMMA::MXF8F6F4Format::E3M2},
           {RuntimeDatatype::kE2M1, cute::UMMA::MXF8F6F4Format::E2M1}
       };
@@ -412,10 +410,9 @@ protected:
 
     }
     else {
-    
-    operator_args.mainloop.ptr_A = static_cast<ElementA const *>(arguments->A);
-    operator_args.mainloop.ptr_B = static_cast<ElementB const *>(arguments->B);
-    } 
+      operator_args.mainloop.ptr_A = static_cast<ElementA const *>(arguments->A);
+      operator_args.mainloop.ptr_B = static_cast<ElementB const *>(arguments->B);
+    }
     operator_args.epilogue.ptr_C = static_cast<ElementC const *>(arguments->C);
     operator_args.epilogue.ptr_D = static_cast<ElementD       *>(arguments->D);
 
@@ -441,11 +438,11 @@ protected:
     operator_args.epilogue.dD = operator_args.epilogue.dC;
 
     using MainloopPolicy = typename CollectiveMainloop::DispatchPolicy;
-    if constexpr(is_mixed_dtype_mainloop_(MainloopPolicy{})) {
-      int problem_m = arguments->problem_size.m();
-      int problem_n = arguments->problem_size.n();
-      int problem_k = arguments->problem_size.k();
-      int options_l = arguments->batch_count;
+    if constexpr(is_sm90_mixed_dtype_mainloop_(MainloopPolicy{})) {
+      const int problem_m = arguments->problem_size.m();
+      const int problem_n = arguments->problem_size.n();
+      const int problem_k = arguments->problem_size.k();
+      const int options_l = arguments->batch_count;
 
       constexpr Sm90MixedInputWiderOperand wider_operand =
         (cutlass::sizeof_bits<ElementA>::value > cutlass::sizeof_bits<ElementB>::value) ?
@@ -455,100 +452,118 @@ protected:
 
       constexpr bool has_scale = !std::is_same_v<typename CollectiveMainloop::ElementScale, void>;
       constexpr bool has_zero  = !std::is_same_v<typename CollectiveMainloop::ElementZero,  void>;
-      if constexpr(has_scale) {
-        int options_g = problem_k;
-        int scale_k = (problem_k + options_g - 1) / options_g;
 
-        constexpr bool is_A4B8 = (
-          cutlass::is_same_v<ElementA, cutlass::int4b_t> &&
-          (cutlass::is_same_v<ElementB, cutlass::float_e4m3_t> ||
-           cutlass::is_same_v<ElementB, cutlass::float_e5m2_t>));
-        constexpr bool is_A8B4 = (
-          cutlass::is_same_v<ElementB, cutlass::int4b_t> &&
-          (cutlass::is_same_v<ElementA, cutlass::float_e4m3_t> ||
-           cutlass::is_same_v<ElementA, cutlass::float_e5m2_t>));
-        constexpr bool is_int4_x_fp8 = is_A4B8 || is_A8B4;
+      const int options_g = problem_k;
+      const int scale_k = (problem_k + options_g - 1) / options_g;
 
-        // In int4 * fp8, ElementScale is a cutlass::Array, need to take out it's real element
-        using ElementScaleMainloop = typename CollectiveMainloop::ElementScale;
-        using ElementScale = typename UnderlyingElement<typename CollectiveMainloop::ElementScale>::type;
-        using StrideS = typename CollectiveMainloop::StrideScale;
-        // In ScaleOnly mode, we have allocated the same size of memory for arguments->Z and arguments->S
-        using ElementZero = std::conditional_t<
-          has_zero,
-          typename CollectiveMainloop::ElementZero,
-          ElementScale
-        >;
-        const int SZ_1st_dim = (wider_operand == Sm90MixedInputWiderOperand::A) ? problem_n : problem_m;
-        const size_t SZ_size = static_cast<size_t>(SZ_1st_dim * scale_k * options_l);
-        auto shape_SZ = cute::make_shape(SZ_1st_dim, scale_k, options_l);
-        ElementScale *ptr_S = static_cast<ElementScale *>(arguments->Scale);
-        ElementZero  *ptr_Z = static_cast<ElementZero  *>(arguments->Zero);
+      constexpr bool is_A4B8 = (
+        cutlass::is_same_v<ElementA, cutlass::int4b_t> &&
+        (cutlass::is_same_v<ElementB, cutlass::float_e4m3_t> ||
+         cutlass::is_same_v<ElementB, cutlass::float_e5m2_t>));
+      constexpr bool is_A8B4 = (
+        cutlass::is_same_v<ElementB, cutlass::int4b_t> &&
+        (cutlass::is_same_v<ElementA, cutlass::float_e4m3_t> ||
+         cutlass::is_same_v<ElementA, cutlass::float_e5m2_t>));
+      constexpr bool is_int4_x_fp8 = is_A4B8 || is_A8B4;
 
-        // 1. If arguments is initialized in profiler, S and Z needs to be allocated and filled
-        if (arguments->generate_scale_and_zero) {
-          // Need to fix max_dequant_val and min_dequant_val?
+      // If this is a convert-only kernel, we still need to generate dequantized A or B for verification,
+      // and in this case ElementScale is the same as ElementWide
+      // In int4 * fp8, ElementScale is a cutlass::Array, need to take out it's real element
+      using DummyElementScaleMainloop = std::conditional_t<
+        is_int4_x_fp8,
+        typename cutlass::Array<ElementWide, 8>,
+        ElementWide
+      >;
+      using ElementScaleMainloop = std::conditional_t<
+        has_scale,
+        typename CollectiveMainloop::ElementScale,
+        DummyElementScaleMainloop
+      >;
+      using ElementScale = std::conditional_t<
+        has_scale,
+        typename UnderlyingElement<typename CollectiveMainloop::ElementScale>::type,
+        ElementWide
+      >;
+      using StrideScale = typename CollectiveMainloop::StrideScale;
+      // In ScaleOnly mode, we have allocated the same size of memory for arguments->Z and arguments->S
+      using ElementZero = std::conditional_t<
+        has_zero,
+        typename CollectiveMainloop::ElementZero,
+        ElementScale
+      >;
+      const int SZ_1st_dim = (wider_operand == Sm90MixedInputWiderOperand::A) ? problem_n : problem_m;
+      const size_t SZ_size = static_cast<size_t>(SZ_1st_dim * scale_k * options_l);
+      auto shape_SZ = cute::make_shape(SZ_1st_dim, scale_k, options_l);
+      ElementScale *ptr_S = static_cast<ElementScale *>(arguments->Scale);
+      ElementZero  *ptr_Z = static_cast<ElementZero  *>(arguments->Zero);
+
+      // 1. If arguments is initialized in profiler, S and Z needs to be allocated and filled
+      if (arguments->generate_scale_and_zero) {
+        float scale_min = 1.0f, scale_max = 1.0f;
+        if constexpr(has_scale) {
           const float elt_max_f = float(cutlass::platform::numeric_limits<ElementScale>::max());
+          // Need to fix max_dequant_val and min_dequant_val?
           const float max_dequant_val = elt_max_f * 0.25f;
           const float min_dequant_val = 0.5f;
-          const float scale_max = max_dequant_val / elt_max_f;
-          const float scale_min = min_dequant_val / elt_max_f;
-          uint64_t seed = 2023;
-          cutlass::reference::device::BlockFillRandomUniform(
-            ptr_S, SZ_size, seed, ElementScale(scale_max), ElementScale(scale_min));
+          scale_max = max_dequant_val / elt_max_f;
+          scale_min = min_dequant_val / elt_max_f;
+        }
+        uint64_t seed = 2023;
+        cutlass::reference::device::BlockFillRandomUniform(
+          ptr_S, SZ_size, seed, ElementScale(scale_max), ElementScale(scale_min));
 
-          // In ScaleOnly mode, set Z as zero for generating dequantized A or B
-          const float zero_max = has_zero ?  2.0f : 0.0f;
-          const float zero_min = has_zero ? -2.0f : 0.0f;
-          cutlass::reference::device::BlockFillRandomUniform(
-            ptr_Z, SZ_size, seed, ElementZero(zero_max), ElementZero(zero_min));
-        }  // End of "if (arguments->generate_scale_and_zero)"
+        // In ScaleOnly mode, set Z as zero for generating dequantized A or B
+        const float zero_max = has_zero ?  2.0f : 0.0f;
+        const float zero_min = has_zero ? -2.0f : 0.0f;
+        cutlass::reference::device::BlockFillRandomUniform(
+          ptr_Z, SZ_size, seed, ElementZero(zero_max), ElementZero(zero_min));
+      }  // End of "if (arguments->generate_scale_and_zero)"
 
-        // 2. Generate the dequantized A or B for verification
-        if (arguments->generate_dequantized_AB) {
-          StrideS stride_SZ = cutlass::make_cute_packed_stride(StrideS{}, shape_SZ);
-          auto layout_SZ = cute::make_layout(shape_SZ, stride_SZ);
-          if constexpr(wider_operand == Sm90MixedInputWiderOperand::A) {
-            if constexpr(is_StrideB_Layout) {
-              // The generator only generates row-major A and col-major B at the moment
-              // Need a way to read out the actual layout of B later
-              using ActualLayoutB = cutlass::layout::ColumnMajor;
-              using ActualStrideB = cutlass::detail::TagToStrideB_t<ActualLayoutB>;
-              dequantize_encode_<ElementWide, ElementNarrow, ElementScaleMainloop, ActualStrideB, wider_operand, is_A8B4>(
-                operator_args, arguments, stream, problem_m, problem_k, options_l, options_g, ptr_S, ptr_Z, SZ_size, layout_SZ);
-            }
-            else {
-              using ActualStrideB = typename CollectiveMainloop::StrideB;
-              dequantize_encode_<ElementWide, ElementNarrow, ElementScaleMainloop, ActualStrideB, wider_operand, is_A8B4>(
-                operator_args, arguments, stream, problem_m, problem_k, options_l, options_g, ptr_S, ptr_Z, SZ_size, layout_SZ);
-            }
+      // 2. Generate the dequantized A or B for verification
+      if (arguments->generate_dequantized_AB) {
+        StrideScale stride_SZ = cutlass::make_cute_packed_stride(StrideScale{}, shape_SZ);
+        auto layout_SZ = cute::make_layout(shape_SZ, stride_SZ);
+        if constexpr(wider_operand == Sm90MixedInputWiderOperand::A) {
+          if constexpr(is_StrideB_Layout) {
+            // The generator only generates row-major A and col-major B at the moment
+            // Need a way to read out the actual layout of B later
+            using ActualLayoutB = cutlass::layout::ColumnMajor;
+            using ActualStrideB = cutlass::detail::TagToStrideB_t<ActualLayoutB>;
+            dequantize_encode_<ElementWide, ElementNarrow, ElementScaleMainloop, ActualStrideB, wider_operand, is_A8B4>(
+              operator_args, arguments, stream, problem_m, problem_k, options_l, options_g, ptr_S, ptr_Z, SZ_size, layout_SZ);
           }
           else {
-            if constexpr(is_StrideA_Layout) {
-              // The generator only generates row-major A and col-major B at the moment
-              // Need a way to read out the actual layout of A later
-              using ActualLayoutA = cutlass::layout::RowMajor;
-              using ActualStrideA = cutlass::detail::TagToStrideA_t<ActualLayoutA>;
-              dequantize_encode_<ElementWide, ElementNarrow, ElementScaleMainloop, ActualStrideA, wider_operand, is_A4B8>(
-                operator_args, arguments, stream, problem_m, problem_k, options_l, options_g, ptr_S, ptr_Z, SZ_size, layout_SZ);
-            }
-            else {
-              using ActualStrideA = typename CollectiveMainloop::StrideA;
-              dequantize_encode_<ElementWide, ElementNarrow, ElementScaleMainloop, ActualStrideA, wider_operand, is_A4B8>(
-                operator_args, arguments, stream, problem_m, problem_k, options_l, options_g, ptr_S, ptr_Z, SZ_size, layout_SZ);
-            }
-          }  // End of "if constexpr(wider_operand == Sm90MixedInputWiderOperand::A)"
-          arguments->dequantized_AB_ready[0] = true;
-        }  // End of "if (arguments->generate_dequantized_AB)"
+            using ActualStrideB = typename CollectiveMainloop::StrideB;
+            dequantize_encode_<ElementWide, ElementNarrow, ElementScaleMainloop, ActualStrideB, wider_operand, is_A8B4>(
+              operator_args, arguments, stream, problem_m, problem_k, options_l, options_g, ptr_S, ptr_Z, SZ_size, layout_SZ);
+          }
+        }
+        else {
+          if constexpr(is_StrideA_Layout) {
+            // The generator only generates row-major A and col-major B at the moment
+            // Need a way to read out the actual layout of A later
+            using ActualLayoutA = cutlass::layout::RowMajor;
+            using ActualStrideA = cutlass::detail::TagToStrideA_t<ActualLayoutA>;
+            dequantize_encode_<ElementWide, ElementNarrow, ElementScaleMainloop, ActualStrideA, wider_operand, is_A4B8>(
+              operator_args, arguments, stream, problem_m, problem_k, options_l, options_g, ptr_S, ptr_Z, SZ_size, layout_SZ);
+          }
+          else {
+            using ActualStrideA = typename CollectiveMainloop::StrideA;
+            dequantize_encode_<ElementWide, ElementNarrow, ElementScaleMainloop, ActualStrideA, wider_operand, is_A4B8>(
+              operator_args, arguments, stream, problem_m, problem_k, options_l, options_g, ptr_S, ptr_Z, SZ_size, layout_SZ);
+          }
+        }  // End of "if constexpr(wider_operand == Sm90MixedInputWiderOperand::A)"
+      }  // End of "if (arguments->generate_dequantized_AB)"
 
-        // 3. Put arguments in mainloop
+      // 3. Put Scale and Zero in mainloop
+      if constexpr(has_scale) {
         if constexpr(is_int4_x_fp8) {
           operator_args.mainloop.ptr_S = static_cast<ElementScaleMainloop const*>(arguments->packed_Scale);
         }
         else {
           operator_args.mainloop.ptr_S = static_cast<ElementScale const*>(arguments->Scale);
         }
-        operator_args.mainloop.dS = cutlass::make_cute_packed_stride(StrideS{}, shape_SZ);
+        operator_args.mainloop.dS = cutlass::make_cute_packed_stride(StrideScale{}, shape_SZ);
         operator_args.mainloop.group_size = options_g;
         if constexpr(has_zero) {
           operator_args.mainloop.ptr_Z = static_cast<ElementZero const*>(arguments->Zero);
@@ -580,7 +595,7 @@ protected:
         handle_shuffle_tensor_<ElementA, ActualStrideA, LayoutA_Reordered, LayoutAtomQuant, wider_operand>(
           operator_args, arguments, problem_m, problem_k, options_l);
       }
-    } // End of "if constexpr(is_mixed_dtype_mainloop_(MainloopPolicy{}))"
+    } // End of "if constexpr(is_sm90_mixed_dtype_mainloop_(MainloopPolicy{}))"
 
     /* Query device SM count and max active clusters to pass onto the kernel as an argument, where needed */
     operator_args.hw_info.sm_count = arguments->sm_count;
@@ -609,7 +624,6 @@ protected:
       operator_args.scheduler.splits = arguments->split_k_slices;
     }
 
-    
     if constexpr (Operator::ArchTag::kMinComputeCapability >= 100) {
       operator_args.hw_info.cluster_shape = dim3(
         arguments->cluster_shape.m(),
@@ -620,7 +634,6 @@ protected:
         arguments->cluster_shape_fallback.n(),
         arguments->cluster_shape_fallback.k());
     }
-    
     return status;
   }
 
@@ -632,7 +645,7 @@ public:
     GemmUniversalArguments const *arguments =
       static_cast<GemmUniversalArguments const *>(arguments_ptr);
     OperatorArguments args;
-    
+
     auto status = update_arguments_(args, arguments);
     if (status != Status::kSuccess) {
       return status;

tools/library/src/grouped_gemm_operation_3x.hpp

@@ -512,8 +512,8 @@ public:
       auto N = get<1>(shape);
       auto K = get<2>(shape);
 
-      auto layout_SFA = CollectiveMainloop::Sm100BlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(M, N, K, 1));
-      auto layout_SFB = CollectiveMainloop::Sm100BlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1));
+      auto layout_SFA = CollectiveMainloop::Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(M, N, K, 1));
+      auto layout_SFB = CollectiveMainloop::Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1));
       layout_SFA_host[group_idx] = layout_SFA;
       layout_SFB_host[group_idx] = layout_SFB;
     }

tools/library/src/reference/block_scaled_gemm_reference_operation.h

@@ -243,14 +243,14 @@ public:
     auto stride_c = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, L));
     auto stride_d = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, L));
 
-    using Sm100BlockScaledConfig = cutlass::detail::Sm100BlockScaledConfig<SFVecSize>;
+    using Sm1xxBlockScaledConfig = cutlass::detail::Sm1xxBlockScaledConfig<SFVecSize>;
     auto A = cute::make_tensor(detail::make_iterator(static_cast<ElementA const*>(args.A)),
         cute::make_layout(cute::make_shape(M, K, L), stride_a));
-    auto SfA = make_tensor(static_cast<ElementSFA const*>(args.SFA), Sm100BlockScaledConfig::tile_atom_to_shape_SFA(problem_shape_MNKL));
+    auto SfA = make_tensor(static_cast<ElementSFA const*>(args.SFA), Sm1xxBlockScaledConfig::tile_atom_to_shape_SFA(problem_shape_MNKL));
 
     auto B = cute::make_tensor(detail::make_iterator(static_cast<ElementB const*>(args.B)),
         cute::make_layout(cute::make_shape(N, K, L), stride_b));
-    auto SfB = make_tensor(static_cast<ElementSFB const*>(args.SFB), Sm100BlockScaledConfig::tile_atom_to_shape_SFB(problem_shape_MNKL));
+    auto SfB = make_tensor(static_cast<ElementSFB const*>(args.SFB), Sm1xxBlockScaledConfig::tile_atom_to_shape_SFB(problem_shape_MNKL));
 
     auto C = [&]() {
       if constexpr (not is_same_v<ElementC, void>) {
@@ -273,11 +273,11 @@ public:
 
     if constexpr (not is_same_v<ElementSFD, void>) { 
 
-      using Sm100BlockScaledOutputConfig = cutlass::detail::Sm100BlockScaledOutputConfig<
+      using Sm1xxBlockScaledOutputConfig= cutlass::detail::Sm1xxBlockScaledOutputConfig<
                                               EpilogueSFVecSize
                                             >;
 
-      auto SfD = cute::make_tensor(detail::make_iterator(static_cast<ElementSFD*>(args.SFD)), Sm100BlockScaledOutputConfig::tile_atom_to_shape_SFD(problem_shape_MNKL));
+      auto SfD = cute::make_tensor(detail::make_iterator(static_cast<ElementSFD*>(args.SFD)), Sm1xxBlockScaledOutputConfig::tile_atom_to_shape_SFD(problem_shape_MNKL));
 
       cutlass::reference::host::GettBlockScalingEpilogueParams<
           ElementCompute, ElementAccumulator, ElementCompute,

tools/library/src/sparse_gemm_operation_3x.hpp

@@ -36,13 +36,22 @@
 
 #include "cutlass/cutlass.h"
 #include "cutlass/detail/collective.hpp"
+#include "cutlass/array.h"
+#include "cutlass/array_subbyte.h"
 #include "cutlass/library/library.h"
-#include "cutlass/library/util.h"
 #include "cutlass/transform/kernel/sparse_gemm_compressor.hpp" // StructuredSparseCompressor
 #include "cutlass/transform/device/transform_universal_adapter.hpp" // TransformUniversalAdapter
 #include "cutlass/util/packed_stride.hpp"        // make_cute_packed_stride
 #include "gemm_operation_3x.hpp"
 #include "library_internal.h"
+#include "cutlass/gemm/dispatch_policy.hpp"
+#include "cutlass/util/packed_stride.hpp"
+#include "cutlass/util/mixed_dtype_utils.hpp"
+#include "cutlass/util/device_memory.h"
+#include "cutlass/util/reference/device/tensor_fill.h"
+#include "cutlass/util/reference/device/tensor_compare.h"
+#include "cute/tensor.hpp"
+#include <unordered_map>
 
 ///////////////////////////////////////////////////////////////////////////////////////////////////
 
@@ -77,6 +86,16 @@ public:
   using CollectiveEpilogue = typename Operator::CollectiveEpilogue;
   using ThreadEpilogueOp = typename CollectiveEpilogue::ThreadEpilogueOp;
 
+  static constexpr bool IsRuntimeDataTypeA = cutlass::gemm::collective::detail::is_sm10x_runtime_f8f6f4<ElementA>();
+
+  static constexpr bool IsRuntimeDataTypeB = cutlass::gemm::collective::detail::is_sm10x_runtime_f8f6f4<ElementB>();
+
+  static_assert((IsRuntimeDataTypeA && IsRuntimeDataTypeB) ||
+                (!IsRuntimeDataTypeA && !IsRuntimeDataTypeB),
+                "ElementA and ElementB in a GEMM kernel should be both runtime or both static.");
+
+  static constexpr bool IsRuntimeDataType = IsRuntimeDataTypeA && IsRuntimeDataTypeB;
+
   using ElementE = typename CollectiveMainloop::ElementE;
   using LayoutE = typename CollectiveMainloop::LayoutE;
   using SparseConfig = typename CollectiveMainloop::SparseConfig;
@@ -170,9 +189,41 @@ protected:
       arguments->batch_count);
 
     // update arguments
-    operator_args.mainloop.ptr_A = reinterpret_cast<ElementA const *>(device_a_compressed_ptr);
-    operator_args.mainloop.ptr_B = static_cast<ElementB const *>(arguments->B);
-    operator_args.mainloop.ptr_E = reinterpret_cast<ElementE const *>(device_e_ptr);
+
+    if constexpr (IsRuntimeDataType) {
+      using ArrayElementA = typename Operator::GemmKernel::CollectiveMainloop::ArrayElementA;
+      using ArrayElementB = typename Operator::GemmKernel::CollectiveMainloop::ArrayElementB;
+      operator_args.mainloop.ptr_A = static_cast<ArrayElementA const *>(device_a_compressed_ptr);
+      operator_args.mainloop.ptr_B = static_cast<ArrayElementB const *>(arguments->B);
+
+      std::unordered_map<RuntimeDatatype, cute::UMMA::MXF8F6F4Format> mapping = {
+          {RuntimeDatatype::kE4M3, cute::UMMA::MXF8F6F4Format::E4M3},
+          {RuntimeDatatype::kE5M2, cute::UMMA::MXF8F6F4Format::E5M2},
+          {RuntimeDatatype::kE3M2, cute::UMMA::MXF8F6F4Format::E3M2},
+          {RuntimeDatatype::kE2M1, cute::UMMA::MXF8F6F4Format::E2M1}
+      };
+
+      auto iter_runtime_a = mapping.find(arguments->runtime_input_datatype_a);
+      auto iter_runtime_b = mapping.find(arguments->runtime_input_datatype_b);
+
+      if (iter_runtime_a != mapping.end()) {
+          operator_args.mainloop.runtime_data_type_a = iter_runtime_a->second;
+      } else {
+        assert("invalid runtime argument for datatype A!");
+      }
+
+      if (iter_runtime_b != mapping.end()) {
+          operator_args.mainloop.runtime_data_type_b = iter_runtime_b->second;
+      } else {
+        assert("invalid runtime argument for datatype B!");
+      }
+
+    }
+    else {
+      operator_args.mainloop.ptr_A = static_cast<ElementA const *>(device_a_compressed_ptr);
+      operator_args.mainloop.ptr_B = static_cast<ElementB const *>(arguments->B);
+    }
+    operator_args.mainloop.ptr_E = static_cast<ElementE const *>(device_e_ptr);
     operator_args.epilogue.ptr_C = static_cast<ElementC const *>(arguments->C);
     operator_args.epilogue.ptr_D = static_cast<ElementD       *>(arguments->D);
 
@@ -184,7 +235,7 @@ protected:
         arguments->ldc, arguments->batch_stride_C);
     operator_args.epilogue.dD = operator_args.epilogue.dC;
 
-    /* Query device SM count to pass onto the kernel as an argument, where needed */
+    /* Query device SM count and max active clusters to pass onto the kernel as an argument, where needed */
     operator_args.hw_info.sm_count = arguments->sm_count;
     if constexpr (!std::is_const_v<decltype(operator_args.scheduler.max_swizzle_size)>) {
       operator_args.scheduler.max_swizzle_size = arguments->swizzle_size;
@@ -208,6 +259,16 @@ protected:
       operator_args.scheduler.splits = arguments->split_k_slices;
     }
 
+    if constexpr (Operator::ArchTag::kMinComputeCapability >= 100) {
+      operator_args.hw_info.cluster_shape = dim3(
+        arguments->cluster_shape.m(),
+        arguments->cluster_shape.n(),
+        arguments->cluster_shape.k());
+      operator_args.hw_info.cluster_shape_fallback = dim3(
+        arguments->cluster_shape_fallback.m(),
+        arguments->cluster_shape_fallback.n(),
+        arguments->cluster_shape_fallback.k());
+    }
     return status;
   }