feat(nvidia): add cutlass_scaled_mm operator

scripts/generate_wrappers.py

@@ -512,6 +512,12 @@ def _find_params_with_defaults(op_name):
     return mapping
 
 
+def _is_data_type_spelling(spelling):
+    spelling = spelling.replace("const ", "").replace("&", "").strip()
+
+    return spelling.rsplit("::", maxsplit=1)[-1] == "DataType"
+
+
 def _generate_pybind11(operator):
     optional_tensor_params = _find_optional_tensor_params(operator.name)
     optional_non_tensor_params = _find_optional_non_tensor_params(operator.name)
@@ -546,6 +552,9 @@ def _is_vector_tensor(arg):
     def _is_vector_int64(arg):
         return arg.spelling in vector_int64_params
 
+    def _is_data_type(arg):
+        return _is_data_type_spelling(arg.type.spelling)
+
     def _generate_params(node):
         parts = []
 
@@ -563,6 +572,8 @@ def _generate_params(node):
                 )
             elif _is_vector_int64(arg):
                 parts.append(f"const std::vector<int64_t> {arg.spelling}")
+            elif _is_data_type(arg):
+                parts.append(f"py::object {arg.spelling}")
             else:
                 param = arg.type.spelling.replace("const Tensor", "py::object").replace(
                     "Tensor", "py::object"
@@ -584,6 +595,8 @@ def _generate_arguments(node):
                 args.append(f"VectorTensorFromPybind11Handle({arg.spelling})")
             elif "Tensor" in arg.type.spelling:
                 args.append(f"TensorFromPybind11Handle({arg.spelling})")
+            elif _is_data_type(arg):
+                args.append(f"DataTypeFromPybind11Handle({arg.spelling})")
             else:
                 args.append(arg.spelling)
 
@@ -930,16 +943,28 @@ def _handle_tensor(spelling):
         def _handle_std_optional(spelling):
             return _unwrap_std_optional(spelling)
 
+        def _handle_data_type(spelling):
+            if not _is_data_type_spelling(spelling):
+                return spelling
+
+            prefix = "const " if spelling.strip().startswith("const ") else ""
+
+            return f"{prefix}infiniDtype_t"
+
         return ", ".join(
-            f"{_handle_std_optional(_handle_tensor(arg.type.spelling))} {arg.spelling}"
+            f"{_handle_data_type(_handle_std_optional(_handle_tensor(arg.type.spelling)))} {arg.spelling}"
             for arg in arguments
         )
 
     def _generate_arguments(node, is_data=False):
         return ", ".join(
-            _generate_tensor_caster(arg.spelling, is_data=is_data)
-            if "Tensor" in arg.type.spelling
-            else arg.spelling
+            f"DataTypeFromInfiniDType({arg.spelling})"
+            if _is_data_type_spelling(arg.type.spelling)
+            else (
+                _generate_tensor_caster(arg.spelling, is_data=is_data)
+                if "Tensor" in arg.type.spelling
+                else arg.spelling
+            )
             for arg in node.get_arguments()
             if arg.spelling != "handle" and arg.spelling != "stream"
         )

src/CMakeLists.txt

@@ -67,6 +67,44 @@ if(WITH_CPU)
 endif()
 
 if(WITH_NVIDIA)
+    if(NOT TARGET nvidia::cutlass::cutlass)
+        find_package(NvidiaCutlass 4.4.1 EXACT CONFIG QUIET)
+    endif()
+
+    if(NOT TARGET nvidia::cutlass::cutlass)
+        include(FetchContent)
+        cmake_policy(PUSH)
+        if(POLICY CMP0135)
+            cmake_policy(SET CMP0135 NEW)
+        endif()
+        if(POLICY CMP0169)
+            cmake_policy(SET CMP0169 OLD)
+        endif()
+
+        FetchContent_Declare(
+            cutlass
+            URL https://github.com/NVIDIA/cutlass/archive/087c84df83d254b5fb295a7a408f1a1d554085cf.tar.gz
+            URL_HASH SHA256=99340f7a8c87f5295b91aecf8295df6693ad954807672c3869a7fb3d0084f709
+        )
+        # Populate only to avoid CUTLASS 4.4.1's CMake 3.19 requirement and install rules.
+        FetchContent_GetProperties(cutlass)
+        if(NOT cutlass_POPULATED)
+            FetchContent_Populate(cutlass)
+        endif()
+        cmake_policy(POP)
+
+        if(NOT EXISTS "${cutlass_SOURCE_DIR}/include/cutlass/cutlass.h")
+            message(FATAL_ERROR
+                "CUTLASS source directory `${cutlass_SOURCE_DIR}` does not contain "
+                "`include/cutlass/cutlass.h`.")
+        endif()
+
+        add_library(infiniops_cutlass_headers INTERFACE)
+        target_include_directories(infiniops_cutlass_headers INTERFACE
+            "${cutlass_SOURCE_DIR}/include")
+        add_library(nvidia::cutlass::cutlass ALIAS infiniops_cutlass_headers)
+    endif()
+
     set(NVIDIA_PATTERNS
         "native/cuda/*.cc"
         "native/cuda/*.cpp"
@@ -78,10 +116,20 @@ if(WITH_NVIDIA)
 
     file(GLOB_RECURSE NVIDIA_SOURCES CONFIGURE_DEPENDS ${NVIDIA_PATTERNS})
 
+    # The CUTLASS INT8 Tensor Core kernel needs SM75 device compilation. Scope
+    # the extra code objects to this provider source so other CUDA operators
+    # retain the caller/compiler architecture configuration unchanged.
+    set(_cutlass_scaled_mm_source
+        "${CMAKE_CURRENT_SOURCE_DIR}/native/cuda/nvidia/ops/cutlass_scaled_mm/cutlass.cu")
+    set_source_files_properties("${_cutlass_scaled_mm_source}" PROPERTIES
+        COMPILE_OPTIONS
+            "--generate-code=arch=compute_75,code=sm_75;--generate-code=arch=compute_75,code=compute_75")
+
     enable_language(CUDA)
 
     target_compile_definitions(infiniops PUBLIC WITH_NVIDIA=1)
     target_sources(infiniops PRIVATE ${NVIDIA_SOURCES})
+    target_link_libraries(infiniops PRIVATE nvidia::cutlass::cutlass)
 
     find_package(CUDAToolkit REQUIRED)
     target_link_libraries(infiniops PUBLIC CUDA::cudart CUDA::cublas CUDA::cublasLt CUDA::cuda_driver)
@@ -419,9 +467,14 @@ set(INFINI_OPS_OPS "" CACHE STRING
 set(INFINI_OPS_SMOKE_BUILD OFF CACHE BOOL
     "Build only the smoke-test operator subset")
 set(_infini_ops_smoke_ops
-    add mul cast cat gemm matmul linear rms_norm swiglu causal_softmax abs clamp exp)
+    add mul cast cat gemm matmul linear rms_norm swiglu
+    causal_softmax abs clamp exp)
 set(_infini_ops_smoke_torch_ops abs clamp exp)
 
+if(WITH_NVIDIA)
+    list(APPEND _infini_ops_smoke_ops cutlass_scaled_mm)
+endif()
+
 if(INFINI_OPS_SMOKE_BUILD)
     if(NOT INFINI_OPS_OPS)
         set(INFINI_OPS_OPS "${_infini_ops_smoke_ops}" CACHE STRING

src/base/cutlass_scaled_mm.h

@@ -0,0 +1,111 @@
+#ifndef INFINI_OPS_BASE_CUTLASS_SCALED_MM_H_
+#define INFINI_OPS_BASE_CUTLASS_SCALED_MM_H_
+
+#include <cassert>
+#include <limits>
+#include <optional>
+
+#include "operator.h"
+
+namespace infini::ops {
+
+class CutlassScaledMm : public Operator<CutlassScaledMm> {
+ public:
+  CutlassScaledMm(const Tensor a, const Tensor b, const Tensor scale_a,
+                  const Tensor scale_b, const DataType out_dtype,
+                  std::optional<Tensor> bias, Tensor out)
+      : m_{a.ndim() == 2 ? a.size(0) : 0},
+        n_{b.ndim() == 2 ? b.size(1) : 0},
+        k_{a.ndim() == 2 ? a.size(1) : 0},
+        lda_{a.ndim() == 2 ? a.stride(0) : 0},
+        ldb_{b.ndim() == 2 ? b.stride(1) : 0},
+        ldo_{out.ndim() == 2 ? out.stride(0) : 0},
+        scale_a_size_{scale_a.numel()},
+        scale_b_size_{scale_b.numel()},
+        out_dtype_{out_dtype} {
+    assert(a.ndim() == 2 && b.ndim() == 2 && out.ndim() == 2 &&
+           "`CutlassScaledMm` requires 2D matrices");
+    assert(a.dtype() == DataType::kInt8 && b.dtype() == DataType::kInt8 &&
+           "`CutlassScaledMm` requires int8 matrix inputs");
+    assert((out_dtype_ == DataType::kFloat16 ||
+            out_dtype_ == DataType::kBFloat16) &&
+           "`CutlassScaledMm` requires float16 or bfloat16 output");
+    assert(out.dtype() == out_dtype_ &&
+           "`CutlassScaledMm` requires `out_dtype` to match the output dtype");
+    assert(a.size(1) == b.size(0) && out.size(0) == a.size(0) &&
+           out.size(1) == b.size(1) &&
+           "`CutlassScaledMm` matrix shapes are incompatible");
+    assert(m_ > 0 && n_ > 0 && k_ > 0 &&
+           "`CutlassScaledMm` requires non-empty matrices");
+    assert(a.stride(1) == 1 && b.stride(0) == 1 && out.stride(1) == 1 &&
+           "`CutlassScaledMm` requires row-major `a` and `out` and "
+           "column-major `b`");
+    assert(
+        lda_ >= k_ && ldb_ >= k_ && ldo_ >= n_ &&
+        "`CutlassScaledMm` matrix strides must cover their logical dimensions");
+    assert(k_ % 16 == 0 && n_ % 16 == 0 && lda_ % 16 == 0 && ldb_ % 16 == 0 &&
+           ldo_ % 16 == 0 &&
+           "`CutlassScaledMm` requires 16-element aligned matrix dimensions");
+    assert(scale_a.dtype() == DataType::kFloat32 &&
+           scale_b.dtype() == DataType::kFloat32 &&
+           "`CutlassScaledMm` requires float32 scales");
+    assert(scale_a.IsContiguous() && scale_b.IsContiguous() &&
+           "`CutlassScaledMm` requires contiguous scales");
+    const auto scale_a_is_per_token{
+        scale_a.ndim() == 2 && scale_a.size(0) == m_ && scale_a.size(1) == 1};
+    const auto scale_b_is_per_channel{
+        scale_b.ndim() == 2 && scale_b.size(0) == 1 && scale_b.size(1) == n_};
+    assert(
+        (scale_a_size_ == 1 || scale_a_is_per_token) &&
+        (scale_b_size_ == 1 || scale_b_is_per_channel) &&
+        "`CutlassScaledMm` scales must be scalar, per-token, or per-channel");
+    const auto same_device_as_a = [&](const Tensor tensor) {
+      return tensor.device().type() == a.device().type() &&
+             tensor.device().index() == a.device().index();
+    };
+    assert(same_device_as_a(b) && same_device_as_a(scale_a) &&
+           same_device_as_a(scale_b) && same_device_as_a(out) &&
+           "`CutlassScaledMm` tensors must be on the same device");
+    assert(m_ <= std::numeric_limits<int>::max() &&
+           n_ <= std::numeric_limits<int>::max() &&
+           k_ <= std::numeric_limits<int>::max() &&
+           "`CutlassScaledMm` matrix dimensions exceed CUTLASS limits");
+
+    if (bias) {
+      assert(bias->ndim() == 1 && bias->numel() == n_ &&
+             "`CutlassScaledMm` bias must have shape `[n]`");
+      assert(bias->dtype() == out_dtype_ && bias->IsContiguous() &&
+             "`CutlassScaledMm` bias must match the output dtype and be "
+             "contiguous");
+      assert(same_device_as_a(*bias) &&
+             "`CutlassScaledMm` bias must be on the output device");
+    }
+  }
+
+  virtual void operator()(const Tensor a, const Tensor b, const Tensor scale_a,
+                          const Tensor scale_b, const DataType out_dtype,
+                          std::optional<Tensor> bias, Tensor out) const = 0;
+
+ protected:
+  Tensor::Size m_{0};
+
+  Tensor::Size n_{0};
+
+  Tensor::Size k_{0};
+
+  Tensor::Stride lda_{0};
+
+  Tensor::Stride ldb_{0};
+
+  Tensor::Stride ldo_{0};
+
+  Tensor::Size scale_a_size_{0};
+
+  Tensor::Size scale_b_size_{0};
+
+  DataType out_dtype_;
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_BASE_CUTLASS_SCALED_MM_H_