diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index f4d0ea1cd..2faeb7025 100644 --- a/scripts/generate_wrappers.py +++ b/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 {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" ) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index fb022422c..9c30701f8 100644 --- a/src/CMakeLists.txt +++ b/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 diff --git a/src/base/cutlass_scaled_mm.h b/src/base/cutlass_scaled_mm.h new file mode 100644 index 000000000..be1d812a1 --- /dev/null +++ b/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 +#include +#include + +#include "operator.h" + +namespace infini::ops { + +class CutlassScaledMm : public Operator { + public: + CutlassScaledMm(const Tensor a, const Tensor b, const Tensor scale_a, + const Tensor scale_b, const DataType out_dtype, + std::optional 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::max() && + n_ <= std::numeric_limits::max() && + k_ <= std::numeric_limits::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 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_ diff --git a/src/native/cuda/nvidia/ops/cutlass_scaled_mm/cutlass.cu b/src/native/cuda/nvidia/ops/cutlass_scaled_mm/cutlass.cu new file mode 100644 index 000000000..824243a84 --- /dev/null +++ b/src/native/cuda/nvidia/ops/cutlass_scaled_mm/cutlass.cu @@ -0,0 +1,244 @@ +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "native/cuda/nvidia/ops/cutlass_scaled_mm/cutlass.h" +#include "native/cuda/nvidia/runtime_.h" + +namespace infini::ops { +namespace { + +class DeviceGuard { + public: + explicit DeviceGuard(int device_index) { + auto status{cudaGetDevice(&previous_device_)}; + assert(status == cudaSuccess && + "`CutlassScaledMm` failed to query the current CUDA device"); + + if (previous_device_ != device_index) { + status = cudaSetDevice(device_index); + assert(status == cudaSuccess && + "`CutlassScaledMm` failed to select the input CUDA device"); + restore_ = true; + } + } + + ~DeviceGuard() { + if (restore_) { + const auto status{cudaSetDevice(previous_device_)}; + assert(status == cudaSuccess && + "`CutlassScaledMm` failed to restore the CUDA device"); + } + } + + private: + int previous_device_{0}; + + bool restore_{false}; +}; + +template +struct GemmDefinition { + using ThreadblockShape = cutlass::gemm::GemmShape<64, 128, 64>; + using WarpShape = cutlass::gemm::GemmShape<32, 64, 64>; + using InstructionShape = cutlass::gemm::GemmShape<8, 8, 16>; + using OutputThreadMap = + cutlass::epilogue::threadblock::OutputTileThreadLayout< + ThreadblockShape, WarpShape, float, 4, 1>; + using Accumulator = cutlass::epilogue::threadblock::VisitorAccFetch; + using Scalar = cutlass::epilogue::threadblock::VisitorScalarBroadcast; + using PerToken = cutlass::epilogue::threadblock::VisitorColBroadcast< + OutputThreadMap, float, + cute::Stride, cute::Int<0>, cute::Int<0>>>; + using PerChannel = cutlass::epilogue::threadblock::VisitorRowBroadcast< + OutputThreadMap, float, + cute::Stride, cute::Int<1>, cute::Int<0>>>; + using ScaleA = std::conditional_t; + using ScaleB = std::conditional_t; + using Bias = cutlass::epilogue::threadblock::VisitorRowBroadcast< + OutputThreadMap, ElementOutput, + cute::Stride, cute::Int<1>, cute::Int<0>>>; + using ScaleBAccumulator = cutlass::epilogue::threadblock::VisitorCompute< + cutlass::multiplies, float, float, + cutlass::FloatRoundStyle::round_to_nearest>; + using ScaleBTree = + cutlass::epilogue::threadblock::Sm80EVT; + using ScaleAAndBias = cutlass::epilogue::threadblock::VisitorCompute< + cutlass::homogeneous_multiply_add, ElementOutput, float, + cutlass::FloatRoundStyle::round_to_nearest>; + using ComputeTree = + cutlass::epilogue::threadblock::Sm80EVT; + using OutputStride = cute::Stride, cute::Int<0>>; + using Output = cutlass::epilogue::threadblock::VisitorAuxStore< + OutputThreadMap, ElementOutput, + cutlass::FloatRoundStyle::round_to_nearest, OutputStride>; + using Epilogue = cutlass::epilogue::threadblock::Sm80EVT; + using Kernel = typename cutlass::gemm::kernel::DefaultGemmWithVisitor< + int8_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 16, + int8_t, cutlass::layout::ColumnMajor, cutlass::ComplexTransform::kNone, + 16, float, cutlass::layout::RowMajor, 4, int32_t, float, + cutlass::arch::OpClassTensorOp, cutlass::arch::Sm75, ThreadblockShape, + WarpShape, InstructionShape, Epilogue, + cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8>, 2, + cutlass::arch::OpMultiplyAddSaturate, 1>::GemmKernel; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +}; + +template +void Execute(const Tensor a, const Tensor b, const Tensor scale_a, + const Tensor scale_b, std::optional bias, Tensor out, + Tensor::Size m, Tensor::Size n, Tensor::Size k, Tensor::Stride lda, + Tensor::Stride ldb, Tensor::Stride ldo, void* stream) { + using Definition = GemmDefinition; + using Gemm = typename Definition::Gemm; + + const auto* a_ptr{reinterpret_cast(a.data())}; + const auto* b_ptr{reinterpret_cast(b.data())}; + const auto* scale_a_ptr{reinterpret_cast(scale_a.data())}; + const auto* scale_b_ptr{reinterpret_cast(scale_b.data())}; + const auto* bias_ptr{ + bias ? reinterpret_cast(bias->data()) : nullptr}; + auto* out_ptr{reinterpret_cast(out.data())}; + + typename Definition::ScaleA::Arguments scale_a_arguments{}; + if constexpr (kPerToken) { + scale_a_arguments.ptr_col = scale_a_ptr; + } else { + scale_a_arguments.scalar_ptrs[0] = scale_a_ptr; + } + + typename Definition::ScaleB::Arguments scale_b_arguments{}; + if constexpr (kPerChannel) { + scale_b_arguments.ptr_row = scale_b_ptr; + } else { + scale_b_arguments.scalar_ptrs[0] = scale_b_ptr; + } + + typename Definition::Bias::Arguments bias_arguments{}; + bias_arguments.ptr_row = bias_ptr; + + typename Definition::ScaleBTree::Arguments scale_b_tree_arguments{ + scale_b_arguments, {}, {}}; + typename Definition::ComputeTree::Arguments compute_arguments{ + scale_a_arguments, scale_b_tree_arguments, bias_arguments, {}}; + const typename Definition::OutputStride output_stride{ + static_cast(ldo), cute::Int<1>{}, cute::Int<0>{}}; + typename Definition::Output::Arguments output_arguments{out_ptr, + output_stride}; + typename Definition::Epilogue::Arguments epilogue_arguments{compute_arguments, + output_arguments}; + + typename Gemm::Arguments arguments{ + cutlass::gemm::GemmUniversalMode::kGemm, + {static_cast(m), static_cast(n), static_cast(k)}, + 1, + epilogue_arguments, + a_ptr, + b_ptr, + nullptr, + nullptr, + 0, + 0, + 0, + 0, + static_cast(lda), + static_cast(ldb), + static_cast(ldo), + static_cast(ldo)}; + + assert(Gemm::get_workspace_size(arguments) == 0 && + "`CutlassScaledMm` unexpectedly requires workspace"); + + Gemm gemm; + auto status{gemm.can_implement(arguments)}; + assert(status == cutlass::Status::kSuccess && + "CUTLASS cannot implement the requested `CutlassScaledMm` geometry"); + + status = gemm(arguments, nullptr, static_cast(stream)); + assert(status == cutlass::Status::kSuccess && + "CUTLASS failed to execute `CutlassScaledMm`"); + + const auto launch_status{cudaGetLastError()}; + assert(launch_status == cudaSuccess && + "`CutlassScaledMm` kernel launch failed"); +} + +template +void DispatchScaleLayout(const Tensor a, const Tensor b, const Tensor scale_a, + const Tensor scale_b, std::optional bias, + Tensor out, Tensor::Size m, Tensor::Size n, + Tensor::Size k, Tensor::Stride lda, Tensor::Stride ldb, + Tensor::Stride ldo, Tensor::Size scale_a_size, + Tensor::Size scale_b_size, void* stream) { + const auto per_token{scale_a_size != 1}; + const auto per_channel{scale_b_size != 1}; + + if (per_token && per_channel) { + Execute(a, b, scale_a, scale_b, bias, out, m, n, + k, lda, ldb, ldo, stream); + } else if (per_token) { + Execute(a, b, scale_a, scale_b, bias, out, m, n, + k, lda, ldb, ldo, stream); + } else if (per_channel) { + Execute(a, b, scale_a, scale_b, bias, out, m, n, + k, lda, ldb, ldo, stream); + } else { + Execute(a, b, scale_a, scale_b, bias, out, m, + n, k, lda, ldb, ldo, stream); + } +} + +} // namespace + +Operator::Operator( + const Tensor a, const Tensor b, const Tensor scale_a, const Tensor scale_b, + const DataType out_dtype, std::optional bias, Tensor out) + : CutlassScaledMm{a, b, scale_a, scale_b, out_dtype, bias, out}, + device_index_{a.device().index()} { + cudaDeviceProp properties{}; + const auto status{cudaGetDeviceProperties(&properties, device_index_)}; + assert(status == cudaSuccess && + "`CutlassScaledMm` failed to query the CUDA device"); + assert((properties.major > 7 || + (properties.major == 7 && properties.minor >= 5)) && + "`CutlassScaledMm` requires compute capability 7.5 or newer"); +} + +void Operator::operator()( + const Tensor a, const Tensor b, const Tensor scale_a, const Tensor scale_b, + const DataType out_dtype, std::optional bias, Tensor out) const { + assert(out_dtype == out_dtype_ && + "`CutlassScaledMm` output dtype changed after descriptor creation"); + assert(out.dtype() == out_dtype_ && + "`CutlassScaledMm` output tensor dtype changed after descriptor " + "creation"); + assert((!bias || bias->dtype() == out_dtype_) && + "`CutlassScaledMm` bias dtype changed after descriptor creation"); + DeviceGuard device_guard{device_index_}; + + switch (out_dtype_) { + case DataType::kFloat16: + DispatchScaleLayout( + a, b, scale_a, scale_b, bias, out, m_, n_, k_, lda_, ldb_, ldo_, + scale_a_size_, scale_b_size_, stream_); + return; + case DataType::kBFloat16: + DispatchScaleLayout( + a, b, scale_a, scale_b, bias, out, m_, n_, k_, lda_, ldb_, ldo_, + scale_a_size_, scale_b_size_, stream_); + return; + default: + assert(false && "`CutlassScaledMm` received an unsupported output dtype"); + } +} + +} // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/cutlass_scaled_mm/cutlass.h b/src/native/cuda/nvidia/ops/cutlass_scaled_mm/cutlass.h new file mode 100644 index 000000000..93a6361ec --- /dev/null +++ b/src/native/cuda/nvidia/ops/cutlass_scaled_mm/cutlass.h @@ -0,0 +1,28 @@ +#ifndef INFINI_OPS_NVIDIA_CUTLASS_SCALED_MM_CUTLASS_H_ +#define INFINI_OPS_NVIDIA_CUTLASS_SCALED_MM_CUTLASS_H_ + +#include + +#include "base/cutlass_scaled_mm.h" + +namespace infini::ops { + +template <> +class Operator + : public CutlassScaledMm { + public: + Operator(const Tensor a, const Tensor b, const Tensor scale_a, + const Tensor scale_b, const DataType out_dtype, + std::optional bias, Tensor out); + + void operator()(const Tensor a, const Tensor b, const Tensor scale_a, + const Tensor scale_b, const DataType out_dtype, + std::optional bias, Tensor out) const override; + + private: + int device_index_{0}; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_NVIDIA_CUTLASS_SCALED_MM_CUTLASS_H_ diff --git a/src/pybind11_utils.h b/src/pybind11_utils.h index 1571e2981..fe025198f 100644 --- a/src/pybind11_utils.h +++ b/src/pybind11_utils.h @@ -29,6 +29,14 @@ inline DataType DataTypeFromString(const std::string& name) { return kStringToDataType.at(name); } +inline DataType DataTypeFromPybind11Handle(py::handle obj) { + auto dtype_str{py::str(obj).cast()}; + const auto pos{dtype_str.find_last_of('.')}; + + return DataTypeFromString( + pos == std::string::npos ? dtype_str : dtype_str.substr(pos + 1)); +} + template inline Device::Type DeviceTypeFromString(const std::string& name) { static const auto kTorchNameToTypes{ @@ -124,10 +132,7 @@ inline Tensor TensorFromPybind11Handle(py::handle obj) { auto shape{obj.attr("shape").cast()}; - auto dtype_str{py::str(obj.attr("dtype")).cast()}; - auto pos{dtype_str.find_last_of('.')}; - auto dtype{DataTypeFromString( - pos == std::string::npos ? dtype_str : dtype_str.substr(pos + 1))}; + auto dtype{DataTypeFromPybind11Handle(obj.attr("dtype"))}; auto device{DeviceFromPybind11Handle(obj)}; diff --git a/tests/conftest.py b/tests/conftest.py index f0be7431f..d49cfd340 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -198,6 +198,7 @@ def _is_smoke_item(item): "tests/test_matmul.py": _is_smoke_matmul_case, "tests/test_mul.py": _is_smoke_mul_case, "tests/test_rms_norm.py": _is_smoke_rms_norm_case, + "tests/test_cutlass_scaled_mm.py": _is_smoke_cutlass_scaled_mm_case, "tests/test_swiglu.py": _is_smoke_swiglu_case, "tests/test_torch_ops.py": _is_smoke_torch_op_case, } @@ -383,6 +384,19 @@ def _is_smoke_rms_norm_case(params): ) +def _is_smoke_cutlass_scaled_mm_case(params): + return ( + params.get("m") == 16 + and params.get("n") == 32 + and params.get("k") == 64 + and params.get("per_token") is True + and params.get("per_channel") is True + and params.get("has_bias") is True + and params.get("padded") is False + and params.get("dtype") == torch.float16 + ) + + def _is_smoke_swiglu_case(params): cases = { ((13, 4), None, None, None), diff --git a/tests/test_cutlass_scaled_mm.py b/tests/test_cutlass_scaled_mm.py new file mode 100644 index 000000000..4f229c942 --- /dev/null +++ b/tests/test_cutlass_scaled_mm.py @@ -0,0 +1,180 @@ +import infini.ops +import pytest +import torch + +from tests.utils import Payload, empty_strided, get_stream, randint_strided + + +if not hasattr(infini.ops, "CutlassScaledMm"): + pytest.skip( + "`CutlassScaledMm` is not available on this platform", allow_module_level=True + ) + + +@pytest.mark.auto_act_and_assert +@pytest.mark.parametrize("m, n, k", ((16, 32, 64), (64, 128, 128))) +@pytest.mark.parametrize( + "per_token, per_channel", + ((False, False), (False, True), (True, False), (True, True)), +) +@pytest.mark.parametrize("has_bias", (False, True)) +@pytest.mark.parametrize("padded", (False, True)) +@pytest.mark.parametrize( + ("dtype", "rtol", "atol"), + ( + (torch.float16, 1e-2, 3e-1), + (torch.bfloat16, 1e-2, 3e-1), + ), +) +def test_cutlass_scaled_mm( + m, + n, + k, + per_token, + per_channel, + has_bias, + padded, + dtype, + device, + implementation_index, + rtol, + atol, +): + a_strides = (k + 16, 1) if padded else None + b_strides = (1, k + 16) if padded else (1, k) + out_strides = (n + 16, 1) if padded else None + a = randint_strided(-4, 5, (m, k), a_strides, dtype=torch.int8, device=device) + b = randint_strided(-4, 5, (k, n), b_strides, dtype=torch.int8, device=device) + + scale_a_shape = (m, 1) if per_token else (1,) + scale_b_shape = (1, n) if per_channel else (1,) + scale_a = torch.rand(scale_a_shape, dtype=torch.float32, device=device) + scale_b = torch.rand(scale_b_shape, dtype=torch.float32, device=device) + bias = torch.rand((n,), dtype=dtype, device=device) if has_bias else None + out = empty_strided((m, n), out_strides, dtype=dtype, device=device) + + return Payload( + _cutlass_scaled_mm, + _torch_cutlass_scaled_mm, + (a, b, scale_a, scale_b, dtype, bias, out), + {"implementation_index": implementation_index}, + rtol=rtol, + atol=atol, + ) + + +def _cutlass_scaled_mm( + a, + b, + scale_a, + scale_b, + out_dtype, + bias, + out, + *, + implementation_index, +): + infini.ops.cutlass_scaled_mm( + a, + b, + scale_a, + scale_b, + out_dtype, + bias, + out, + stream=get_stream(a.device), + implementation_index=implementation_index, + ) + + return out + + +def _torch_cutlass_scaled_mm( + a, + b, + scale_a, + scale_b, + out_dtype, + bias, + out, + *, + implementation_index, +): + del implementation_index + + result = torch.matmul(a.float(), b.float()) * scale_a * scale_b + + if bias is not None: + result = result + bias.float() + + out.copy_(result.to(out_dtype)) + + return out + + +def test_cutlass_scaled_mm_non_default_stream(device, implementation_index): + if device != "cuda": + pytest.skip("non-default CUDA streams require the NVIDIA backend") + + m, n, k = 16, 32, 64 + a = torch.randint(-4, 5, (m, k), dtype=torch.int8, device=device) + b = torch.randint(-4, 5, (n, k), dtype=torch.int8, device=device).t() + scale_a = torch.rand((m, 1), dtype=torch.float32, device=device) + scale_b = torch.rand((1, n), dtype=torch.float32, device=device) + out = torch.empty((m, n), dtype=torch.float16, device=device) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + + with torch.cuda.stream(stream): + _cutlass_scaled_mm( + a, + b, + scale_a, + scale_b, + out.dtype, + None, + out, + implementation_index=implementation_index, + ) + + stream.synchronize() + expected = torch.matmul(a.float(), b.float()) * scale_a * scale_b + torch.testing.assert_close(out, expected.to(out.dtype), rtol=1e-2, atol=3e-1) + + +def test_cutlass_scaled_mm_multi_gpu_device_guard(device, implementation_index): + if device != "cuda" or torch.cuda.device_count() < 2: + pytest.skip("multi-GPU device guard test requires two NVIDIA GPUs") + + original_device = torch.cuda.current_device() + + try: + torch.cuda.set_device(0) + target_device = torch.device("cuda:1") + m, n, k = 16, 32, 64 + a = torch.randint(-4, 5, (m, k), dtype=torch.int8, device=target_device) + b = torch.randint(-4, 5, (n, k), dtype=torch.int8, device=target_device).t() + scale_a = torch.rand((m, 1), dtype=torch.float32, device=target_device) + scale_b = torch.rand((1, n), dtype=torch.float32, device=target_device) + out = torch.empty((m, n), dtype=torch.float16, device=target_device) + stream = torch.cuda.Stream(device=target_device) + stream.wait_stream(torch.cuda.current_stream(target_device)) + + infini.ops.cutlass_scaled_mm( + a, + b, + scale_a, + scale_b, + out.dtype, + None, + out, + stream=stream.cuda_stream, + implementation_index=implementation_index, + ) + + assert torch.cuda.current_device() == 0 + stream.synchronize() + expected = torch.matmul(a.float(), b.float()) * scale_a * scale_b + torch.testing.assert_close(out, expected.to(out.dtype), rtol=1e-2, atol=3e-1) + finally: + torch.cuda.set_device(original_device) diff --git a/tests/test_generate_wrappers.py b/tests/test_generate_wrappers.py index 66835e1f2..dc00aa6c8 100644 --- a/tests/test_generate_wrappers.py +++ b/tests/test_generate_wrappers.py @@ -241,6 +241,49 @@ class Mul { assert 'py::arg("implementation_index") = py::none()' in text +_DTYPE_OP_SOURCE = """ +namespace infini::ops { + +struct Tensor {}; + +enum class DataType {}; + +template +class Operator {}; + +class DtypeOp : public Operator { + public: + DtypeOp(const Tensor input, const DataType out_dtype, Tensor out) {} + + virtual void operator()(const Tensor input, const DataType out_dtype, + Tensor out) const = 0; +}; + +} // namespace infini::ops +""" + + +def test_pybind_converts_data_type_arguments_from_torch_dtype(tmp_path, monkeypatch): + text = _generate_binding("dtype_op", tmp_path, monkeypatch, _DTYPE_OP_SOURCE) + + assert "py::object out_dtype" in text + assert "DataTypeFromPybind11Handle(out_dtype)" in text + assert 'py::arg("out_dtype")' in text + + +def test_legacy_c_converts_data_type_arguments_from_infini_dtype(tmp_path, monkeypatch): + module = _load_generator_module() + base_header = tmp_path / "dtype_op.h" + base_header.write_text(_DTYPE_OP_SOURCE) + monkeypatch.setattr(module, "_find_base_header", lambda op_name: base_header) + operator = module._OperatorExtractor()("dtype_op") + + source, header = module._generate_legacy_c(operator, ()) + + assert "const infiniDtype_t out_dtype" in header + assert "DataTypeFromInfiniDType(out_dtype)" in source + + def test_normalize_op_allowlist_accepts_spaces_and_commas(): module = _load_generator_module()