v3.8.0 update (#2082)

* 3.8 update

* fix Markus' name

---------

Co-authored-by: yuzhai <yuzhai@nvidia.com>

tools/library/include/cutlass/library/library.h

@@ -43,7 +43,8 @@
                     computational overhead
 */
 
-#pragma once
+#ifndef CUTLASS_LIBRARY_LIBRARY_H
+#define CUTLASS_LIBRARY_LIBRARY_H
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
@@ -103,7 +104,7 @@ public:
     void *device_workspace = nullptr,
     cudaStream_t stream = nullptr) const = 0;
 
-  // Originally designed for metadata, but should be useful for FP8/6/4 too. 
+  // Originally designed for metadata, but should be useful for FP8/6/4 too.  
   virtual Status initialize_with_profiler_workspace(
     void const *configuration,
     void *host_workspace,
@@ -282,6 +283,11 @@ struct GemmUniversalConfiguration {
   int device_count{1};
 };
 
+enum class Sm90MixedInputWiderOperand {
+  A = 0,
+  B = 1
+};
+
 struct GemmUniversalArguments {
   // NOTE: these are replicated for 3.0 interfaces
   gemm::GemmCoord problem_size{};
@@ -317,6 +323,18 @@ struct GemmUniversalArguments {
   int swizzle_size{1};
   int split_k_slices{1};
 
+  // For mixed input dtype kernels
+  bool is_mixed_dtype{false};
+  Sm90MixedInputWiderOperand wider_operand{Sm90MixedInputWiderOperand::B};
+  bool generate_scale_and_zero{false};
+  bool generate_dequantized_AB{false};
+  bool *dequantized_AB_ready{nullptr};  // Carry the info back to gemm_operation_profiler.cu
+  void *Scale{nullptr};                 // Scale tensor
+  void *Zero{nullptr};                  // Zero tensor
+  void *dequantized_AB{nullptr};        // Dequantized A or B tensor for verification
+  void *encoded_AB{nullptr};            // Encoded A or B in int4 x fp8 or shuffle
+  void *packed_Scale{nullptr};          // Packed scale for int4 * fp8
+
   int device_index{0};
   
   bool use_pdl{false};
@@ -472,12 +490,16 @@ struct GemmPlanarComplexArrayArguments {
 
 struct GemmGroupedConfiguration {
   int problem_count{0};
-  int threadblock_count{0};
+  // GemmGroupedConfiguration is passed to initialize(), which
+  // is responsible for allocating the device-side stride storage.
+  int64_t* lda;
+  int64_t* ldb;
+  int64_t* ldc;
 };
 
 struct GemmGroupedArguments {
-
-  gemm::GemmCoord *problem_sizes{nullptr};
+  int problem_count{};
+  gemm::GemmCoord* problem_sizes{nullptr};
 
   void * ptr_A{nullptr};
   void * ptr_B{nullptr};
@@ -493,6 +515,18 @@ struct GemmGroupedArguments {
   void const *beta{nullptr};
   ScalarPointerMode pointer_mode{};
   bool use_pdl{false};
+
+  gemm::GemmCoord cluster_shape{};          
+  gemm::GemmCoord cluster_shape_fallback{}; 
+
+  // these should really be in the configuration but staying consistent with GEMM
+  int sm_count{0};
+  // The user is responsible for allocating storage for problem sizes.
+  // Since GemmGroupedArguments is used by both the 2.x and 3.x APIs, we
+  // unfortunately need to have both options in this struct, and the
+  // underlying operation uses the one it needs.
+  cute::Shape<int, int, int>* problem_sizes_3x;
+  cute::Shape<int, int, int>* problem_sizes_3x_host;
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
@@ -880,3 +914,5 @@ struct ReductionArguments {
 } // namespace cutlass
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
+
+#endif

tools/library/include/cutlass/library/types.h

@@ -142,7 +142,7 @@ enum class Provider {
 /// Enumeration indicating the kind of operation
 enum class OperationKind {
   kGemm,
-  kBlockScaledGemm,     
+  kBlockScaledGemm,
   kRankK,
   kRank2K,
   kTrmm,
@@ -152,6 +152,7 @@ enum class OperationKind {
   kEqGemm,
   kSparseGemm,
   kReduction,
+  kGroupedGemm,
   kInvalid
 };
 
@@ -270,7 +271,6 @@ enum class RuntimeDatatype {
   kStatic,
   kE4M3,
   kE5M2,
-  
   kE3M2,
   kE2M3,
   kE2M1,

tools/library/include/cutlass/library/util.h

@@ -34,7 +34,8 @@
   \brief Utilities accompanying the CUTLASS library for interacting with Library types.
 */
 
-#pragma once
+#ifndef CUTLASS_LIBRARY_UTIL_H
+#define CUTLASS_LIBRARY_UTIL_H
 
 #include "cutlass/cutlass.h"
 #include "cutlass/library/library.h"
@@ -213,6 +214,63 @@ bool cast_from_double(std::vector<uint8_t> &bytes, NumericTypeID type, double sr
 
 NumericTypeID dynamic_datatype_to_id(RuntimeDatatype type); 
 
+#define CUDA_CHECK(call)                                                                           \
+  do {                                                                                             \
+    cudaError_t err = (call);                                                                      \
+    if (err != cudaSuccess) {                                                                      \
+      std::cerr << "CUDA Error: " << cudaGetErrorString(err) << " in " << __func__ << " at "       \
+                << __FILE__ << ":" << __LINE__ << std::endl;                                       \
+      return Status::kInvalid;                                                                     \
+    }                                                                                              \
+  } while (0)
+
+// RAII CUDA buffer container
+class CudaBuffer {
+public:
+  CudaBuffer() : size_(0), d_ptr_(nullptr) {}
+
+  explicit CudaBuffer(size_t size) : size_(size), d_ptr_(nullptr) {
+    cudaError_t err = cudaMalloc(&d_ptr_, size_);
+    if (err != cudaSuccess) {
+      throw std::runtime_error("cudaMalloc failed: " + std::string(cudaGetErrorString(err)));
+    }
+  }
+
+  ~CudaBuffer() {
+    if (d_ptr_) {
+      cudaFree(d_ptr_);
+    }
+  }
+
+  CudaBuffer(CudaBuffer const&) = delete;
+  CudaBuffer& operator=(CudaBuffer const&) = delete;
+
+  CudaBuffer(CudaBuffer&& other) noexcept : size_(other.size_), d_ptr_(other.d_ptr_) {
+    other.d_ptr_ = nullptr;
+    other.size_ = 0;
+  }
+
+  CudaBuffer& operator=(CudaBuffer&& other) noexcept {
+    if (this != &other) {
+      if (d_ptr_) {
+        cudaFree(d_ptr_);
+      }
+      d_ptr_ = other.d_ptr_;
+      size_ = other.size_;
+      other.d_ptr_ = nullptr;
+      other.size_ = 0;
+    }
+    return *this;
+  }
+
+  void* data() const noexcept { return d_ptr_; }
+  size_t size() const noexcept { return size_; }
+
+private:
+  size_t size_;
+  void* d_ptr_;
+};
+
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 } // namespace library
@@ -220,3 +278,4 @@ NumericTypeID dynamic_datatype_to_id(RuntimeDatatype type);
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
+#endif