From 3cacc85629e239016716bab59a8d59ca833db9b9 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Mon, 20 Jul 2026 18:40:32 +0800 Subject: [PATCH 1/5] build(nvidia): add CUTLASS dependency --- .github/ci_config.yml | 2 +- .gitmodules | 4 ++++ src/CMakeLists.txt | 9 +++++++++ third_party/cutlass | 1 + 4 files changed, 15 insertions(+), 1 deletion(-) create mode 160000 third_party/cutlass diff --git a/.github/ci_config.yml b/.github/ci_config.yml index 65b2b208f..fa5ce2d25 100644 --- a/.github/ci_config.yml +++ b/.github/ci_config.yml @@ -15,7 +15,7 @@ platforms: BASE_IMAGE: nvcr.io/nvidia/pytorch:24.10-py3 SKIP_APT: "1" PIP_INDEX_URL: https://pypi.tuna.tsinghua.edu.cn/simple - setup: pip install scikit-build-core pybind11 libclang pyyaml && rm -rf /tmp/infinirt-src /tmp/infinirt-build /tmp/infinirt-prefix && git clone --depth 1 ${INFINI_RT_GIT_URL:-https://gh-proxy.com/https://github.com/InfiniTensor/InfiniRT.git} /tmp/infinirt-src && cmake -S /tmp/infinirt-src -B /tmp/infinirt-build -DWITH_CPU=ON -DWITH_NVIDIA=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/tmp/infinirt-prefix && cmake --build /tmp/infinirt-build -j$(nproc) && cmake --install /tmp/infinirt-build && pip install .[dev] --no-build-isolation --config-settings=cmake.define.INFINI_OPS_SMOKE_BUILD=ON --config-settings=cmake.define.INFINI_RT_ROOT=/tmp/infinirt-prefix + setup: git submodule update --init --checkout --depth 1 third_party/cutlass && pip install scikit-build-core pybind11 libclang pyyaml && rm -rf /tmp/infinirt-src /tmp/infinirt-build /tmp/infinirt-prefix && git clone --depth 1 ${INFINI_RT_GIT_URL:-https://gh-proxy.com/https://github.com/InfiniTensor/InfiniRT.git} /tmp/infinirt-src && cmake -S /tmp/infinirt-src -B /tmp/infinirt-build -DWITH_CPU=ON -DWITH_NVIDIA=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/tmp/infinirt-prefix && cmake --build /tmp/infinirt-build -j$(nproc) && cmake --install /tmp/infinirt-build && pip install .[dev] --no-build-isolation --config-settings=cmake.define.INFINI_OPS_SMOKE_BUILD=ON --config-settings=cmake.define.INFINI_RT_ROOT=/tmp/infinirt-prefix jobs: gpu: type: unittest diff --git a/.gitmodules b/.gitmodules index be99e8a8f..bd4db8160 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,7 @@ [submodule ".ci"] path = .ci url = https://github.com/InfiniTensor/ci.git +[submodule "third_party/cutlass"] + path = third_party/cutlass + url = https://github.com/NVIDIA/cutlass.git + ignore = untracked diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index fb022422c..ba6e41b4e 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -67,6 +67,14 @@ if(WITH_CPU) endif() if(WITH_NVIDIA) + set(_cutlass_include_dir + "${PROJECT_SOURCE_DIR}/third_party/cutlass/include") + if(NOT EXISTS "${_cutlass_include_dir}/cutlass/cutlass.h") + message(FATAL_ERROR + "CUTLASS is required by the NVIDIA backend. Run " + "`git submodule update --init --checkout third_party/cutlass`.") + endif() + set(NVIDIA_PATTERNS "native/cuda/*.cc" "native/cuda/*.cpp" @@ -82,6 +90,7 @@ if(WITH_NVIDIA) target_compile_definitions(infiniops PUBLIC WITH_NVIDIA=1) target_sources(infiniops PRIVATE ${NVIDIA_SOURCES}) + target_include_directories(infiniops PRIVATE "${_cutlass_include_dir}") find_package(CUDAToolkit REQUIRED) target_link_libraries(infiniops PUBLIC CUDA::cudart CUDA::cublas CUDA::cublasLt CUDA::cuda_driver) diff --git a/third_party/cutlass b/third_party/cutlass new file mode 160000 index 000000000..087c84df8 --- /dev/null +++ b/third_party/cutlass @@ -0,0 +1 @@ +Subproject commit 087c84df83d254b5fb295a7a408f1a1d554085cf From 2683781dbe695f55f886fc4cdaaad93d2297db8c Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Mon, 20 Jul 2026 19:33:53 +0800 Subject: [PATCH 2/5] feat(nvidia): add CUTLASS scaled mm operator --- src/CMakeLists.txt | 16 +- src/base/scaled_mm.h | 104 ++++++++ .../cuda/nvidia/ops/scaled_mm/cutlass.cu | 237 ++++++++++++++++++ .../cuda/nvidia/ops/scaled_mm/cutlass.h | 31 +++ tests/conftest.py | 14 ++ tests/test_scaled_mm.py | 155 ++++++++++++ 6 files changed, 556 insertions(+), 1 deletion(-) create mode 100644 src/base/scaled_mm.h create mode 100644 src/native/cuda/nvidia/ops/scaled_mm/cutlass.cu create mode 100644 src/native/cuda/nvidia/ops/scaled_mm/cutlass.h create mode 100644 tests/test_scaled_mm.py diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index ba6e41b4e..90ba61919 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -86,6 +86,15 @@ 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/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) @@ -428,9 +437,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 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/scaled_mm.h b/src/base/scaled_mm.h new file mode 100644 index 000000000..20b30d8c4 --- /dev/null +++ b/src/base/scaled_mm.h @@ -0,0 +1,104 @@ +#ifndef INFINI_OPS_BASE_SCALED_MM_H_ +#define INFINI_OPS_BASE_SCALED_MM_H_ + +#include +#include +#include + +#include "operator.h" + +namespace infini::ops { + +class ScaledMm : public Operator { + public: + ScaledMm(const Tensor a, const Tensor b, const Tensor scale_a, + const Tensor scale_b, 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 && + "`ScaledMm` requires 2D matrices"); + assert(a.dtype() == DataType::kInt8 && b.dtype() == DataType::kInt8 && + "`ScaledMm` requires int8 matrix inputs"); + assert((out_dtype_ == DataType::kFloat16 || + out_dtype_ == DataType::kBFloat16) && + "`ScaledMm` requires float16 or bfloat16 output"); + assert(a.size(1) == b.size(0) && out.size(0) == a.size(0) && + out.size(1) == b.size(1) && + "`ScaledMm` matrix shapes are incompatible"); + assert(m_ > 0 && n_ > 0 && k_ > 0 && + "`ScaledMm` requires non-empty matrices"); + assert(a.stride(1) == 1 && b.stride(0) == 1 && out.stride(1) == 1 && + "`ScaledMm` requires row-major `a` and `out` and column-major `b`"); + assert(lda_ >= k_ && ldb_ >= k_ && ldo_ >= n_ && + "`ScaledMm` matrix strides must cover their logical dimensions"); + assert(k_ % 16 == 0 && n_ % 16 == 0 && lda_ % 16 == 0 && ldb_ % 16 == 0 && + ldo_ % 16 == 0 && + "`ScaledMm` requires 16-element aligned matrix dimensions"); + assert(scale_a.dtype() == DataType::kFloat32 && + scale_b.dtype() == DataType::kFloat32 && + "`ScaledMm` requires float32 scales"); + assert(scale_a.IsContiguous() && scale_b.IsContiguous() && + "`ScaledMm` 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) && + "`ScaledMm` 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) && + "`ScaledMm` tensors must be on the same device"); + assert(m_ <= std::numeric_limits::max() && + n_ <= std::numeric_limits::max() && + k_ <= std::numeric_limits::max() && + "`ScaledMm` matrix dimensions exceed CUTLASS limits"); + + if (bias) { + assert(bias->ndim() == 1 && bias->numel() == n_ && + "`ScaledMm` bias must have shape `[n]`"); + assert(bias->dtype() == out_dtype_ && bias->IsContiguous() && + "`ScaledMm` bias must match the output dtype and be contiguous"); + assert(same_device_as_a(*bias) && + "`ScaledMm` bias must be on the output device"); + } + } + + virtual void operator()(const Tensor a, const Tensor b, const Tensor scale_a, + const Tensor scale_b, 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_SCALED_MM_H_ diff --git a/src/native/cuda/nvidia/ops/scaled_mm/cutlass.cu b/src/native/cuda/nvidia/ops/scaled_mm/cutlass.cu new file mode 100644 index 000000000..3ca4212ba --- /dev/null +++ b/src/native/cuda/nvidia/ops/scaled_mm/cutlass.cu @@ -0,0 +1,237 @@ +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "native/cuda/nvidia/ops/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 `ScaledMm` geometry"); + + status = gemm(arguments, nullptr, static_cast(stream)); + assert(status == cutlass::Status::kSuccess && + "CUTLASS failed to execute `ScaledMm`"); + + 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 + +CutlassScaledMm::CutlassScaledMm(const Tensor a, const Tensor b, + const Tensor scale_a, const Tensor scale_b, + std::optional bias, Tensor out) + : ScaledMm{a, b, scale_a, scale_b, 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 CutlassScaledMm::operator()(const Tensor a, const Tensor b, + const Tensor scale_a, const Tensor scale_b, + std::optional bias, Tensor out) const { + 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/scaled_mm/cutlass.h b/src/native/cuda/nvidia/ops/scaled_mm/cutlass.h new file mode 100644 index 000000000..4fae417e9 --- /dev/null +++ b/src/native/cuda/nvidia/ops/scaled_mm/cutlass.h @@ -0,0 +1,31 @@ +#ifndef INFINI_OPS_NVIDIA_SCALED_MM_CUTLASS_H_ +#define INFINI_OPS_NVIDIA_SCALED_MM_CUTLASS_H_ + +#include + +#include "base/scaled_mm.h" + +namespace infini::ops { + +class CutlassScaledMm : public ScaledMm { + public: + CutlassScaledMm(const Tensor a, const Tensor b, const Tensor scale_a, + const Tensor scale_b, std::optional bias, Tensor out); + + void operator()(const Tensor a, const Tensor b, const Tensor scale_a, + const Tensor scale_b, std::optional bias, + Tensor out) const override; + + private: + int device_index_{0}; +}; + +template <> +class Operator : public CutlassScaledMm { + public: + using CutlassScaledMm::CutlassScaledMm; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_NVIDIA_SCALED_MM_CUTLASS_H_ diff --git a/tests/conftest.py b/tests/conftest.py index f0be7431f..5816e2497 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_scaled_mm.py": _is_smoke_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_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_scaled_mm.py b/tests/test_scaled_mm.py new file mode 100644 index 000000000..21a3856c0 --- /dev/null +++ b/tests/test_scaled_mm.py @@ -0,0 +1,155 @@ +import infini.ops +import pytest +import torch + +from tests.utils import Payload, empty_strided, get_stream, randint_strided + + +if not hasattr(infini.ops, "ScaledMm"): + pytest.skip("`ScaledMm` 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_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( + _scaled_mm, + _torch_scaled_mm, + (a, b, scale_a, scale_b, bias, out), + {"implementation_index": implementation_index}, + rtol=rtol, + atol=atol, + ) + + +def _scaled_mm(a, b, scale_a, scale_b, bias, out, *, implementation_index): + infini.ops.scaled_mm( + a, + b, + scale_a, + scale_b, + bias, + out, + stream=get_stream(a.device), + implementation_index=implementation_index, + ) + + return out + + +def _torch_scaled_mm(a, b, scale_a, scale_b, 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_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): + _scaled_mm( + a, + b, + scale_a, + scale_b, + 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_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.scaled_mm( + a, + b, + scale_a, + scale_b, + 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) From c60b1f3f0b0c580eccba34ab3fb28546b04eb226 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 21 Jul 2026 10:21:15 +0800 Subject: [PATCH 3/5] fix(nvidia): rename scaled_mm to cutlass_scaled_mm --- src/CMakeLists.txt | 4 +- src/base/{scaled_mm.h => cutlass_scaled_mm.h} | 54 ++++++++++--------- .../cutlass.cu | 20 +++---- .../nvidia/ops/cutlass_scaled_mm/cutlass.h | 27 ++++++++++ .../cuda/nvidia/ops/scaled_mm/cutlass.h | 31 ----------- tests/conftest.py | 4 +- ...scaled_mm.py => test_cutlass_scaled_mm.py} | 28 +++++----- 7 files changed, 86 insertions(+), 82 deletions(-) rename src/base/{scaled_mm.h => cutlass_scaled_mm.h} (61%) rename src/native/cuda/nvidia/ops/{scaled_mm => cutlass_scaled_mm}/cutlass.cu (93%) create mode 100644 src/native/cuda/nvidia/ops/cutlass_scaled_mm/cutlass.h delete mode 100644 src/native/cuda/nvidia/ops/scaled_mm/cutlass.h rename tests/{test_scaled_mm.py => test_cutlass_scaled_mm.py} (86%) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 90ba61919..852f5ce98 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -90,7 +90,7 @@ if(WITH_NVIDIA) # 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/scaled_mm/cutlass.cu") + "${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") @@ -442,7 +442,7 @@ set(_infini_ops_smoke_ops set(_infini_ops_smoke_torch_ops abs clamp exp) if(WITH_NVIDIA) - list(APPEND _infini_ops_smoke_ops scaled_mm) + list(APPEND _infini_ops_smoke_ops cutlass_scaled_mm) endif() if(INFINI_OPS_SMOKE_BUILD) diff --git a/src/base/scaled_mm.h b/src/base/cutlass_scaled_mm.h similarity index 61% rename from src/base/scaled_mm.h rename to src/base/cutlass_scaled_mm.h index 20b30d8c4..e69598015 100644 --- a/src/base/scaled_mm.h +++ b/src/base/cutlass_scaled_mm.h @@ -1,5 +1,5 @@ -#ifndef INFINI_OPS_BASE_SCALED_MM_H_ -#define INFINI_OPS_BASE_SCALED_MM_H_ +#ifndef INFINI_OPS_BASE_CUTLASS_SCALED_MM_H_ +#define INFINI_OPS_BASE_CUTLASS_SCALED_MM_H_ #include #include @@ -9,10 +9,10 @@ namespace infini::ops { -class ScaledMm : public Operator { +class CutlassScaledMm : public Operator { public: - ScaledMm(const Tensor a, const Tensor b, const Tensor scale_a, - const Tensor scale_b, std::optional bias, Tensor out) + CutlassScaledMm(const Tensor a, const Tensor b, const Tensor scale_a, + const Tensor scale_b, 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}, @@ -23,55 +23,59 @@ class ScaledMm : public Operator { scale_b_size_{scale_b.numel()}, out_dtype_{out.dtype()} { assert(a.ndim() == 2 && b.ndim() == 2 && out.ndim() == 2 && - "`ScaledMm` requires 2D matrices"); + "`CutlassScaledMm` requires 2D matrices"); assert(a.dtype() == DataType::kInt8 && b.dtype() == DataType::kInt8 && - "`ScaledMm` requires int8 matrix inputs"); + "`CutlassScaledMm` requires int8 matrix inputs"); assert((out_dtype_ == DataType::kFloat16 || out_dtype_ == DataType::kBFloat16) && - "`ScaledMm` requires float16 or bfloat16 output"); + "`CutlassScaledMm` requires float16 or bfloat16 output"); assert(a.size(1) == b.size(0) && out.size(0) == a.size(0) && out.size(1) == b.size(1) && - "`ScaledMm` matrix shapes are incompatible"); + "`CutlassScaledMm` matrix shapes are incompatible"); assert(m_ > 0 && n_ > 0 && k_ > 0 && - "`ScaledMm` requires non-empty matrices"); + "`CutlassScaledMm` requires non-empty matrices"); assert(a.stride(1) == 1 && b.stride(0) == 1 && out.stride(1) == 1 && - "`ScaledMm` requires row-major `a` and `out` and column-major `b`"); - assert(lda_ >= k_ && ldb_ >= k_ && ldo_ >= n_ && - "`ScaledMm` matrix strides must cover their logical dimensions"); + "`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 && - "`ScaledMm` requires 16-element aligned matrix dimensions"); + "`CutlassScaledMm` requires 16-element aligned matrix dimensions"); assert(scale_a.dtype() == DataType::kFloat32 && scale_b.dtype() == DataType::kFloat32 && - "`ScaledMm` requires float32 scales"); + "`CutlassScaledMm` requires float32 scales"); assert(scale_a.IsContiguous() && scale_b.IsContiguous() && - "`ScaledMm` requires contiguous scales"); + "`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) && - "`ScaledMm` scales must be scalar, per-token, or per-channel"); + 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) && - "`ScaledMm` tensors must be on the same device"); + "`CutlassScaledMm` tensors must be on the same device"); assert(m_ <= std::numeric_limits::max() && n_ <= std::numeric_limits::max() && k_ <= std::numeric_limits::max() && - "`ScaledMm` matrix dimensions exceed CUTLASS limits"); + "`CutlassScaledMm` matrix dimensions exceed CUTLASS limits"); if (bias) { assert(bias->ndim() == 1 && bias->numel() == n_ && - "`ScaledMm` bias must have shape `[n]`"); + "`CutlassScaledMm` bias must have shape `[n]`"); assert(bias->dtype() == out_dtype_ && bias->IsContiguous() && - "`ScaledMm` bias must match the output dtype and be contiguous"); + "`CutlassScaledMm` bias must match the output dtype and be " + "contiguous"); assert(same_device_as_a(*bias) && - "`ScaledMm` bias must be on the output device"); + "`CutlassScaledMm` bias must be on the output device"); } } @@ -101,4 +105,4 @@ class ScaledMm : public Operator { } // namespace infini::ops -#endif // INFINI_OPS_BASE_SCALED_MM_H_ +#endif // INFINI_OPS_BASE_CUTLASS_SCALED_MM_H_ diff --git a/src/native/cuda/nvidia/ops/scaled_mm/cutlass.cu b/src/native/cuda/nvidia/ops/cutlass_scaled_mm/cutlass.cu similarity index 93% rename from src/native/cuda/nvidia/ops/scaled_mm/cutlass.cu rename to src/native/cuda/nvidia/ops/cutlass_scaled_mm/cutlass.cu index 3ca4212ba..f91fb6cbe 100644 --- a/src/native/cuda/nvidia/ops/scaled_mm/cutlass.cu +++ b/src/native/cuda/nvidia/ops/cutlass_scaled_mm/cutlass.cu @@ -9,7 +9,7 @@ #include #include -#include "native/cuda/nvidia/ops/scaled_mm/cutlass.h" +#include "native/cuda/nvidia/ops/cutlass_scaled_mm/cutlass.h" #include "native/cuda/nvidia/runtime_.h" namespace infini::ops { @@ -161,11 +161,11 @@ void Execute(const Tensor a, const Tensor b, const Tensor scale_a, Gemm gemm; auto status{gemm.can_implement(arguments)}; assert(status == cutlass::Status::kSuccess && - "CUTLASS cannot implement the requested `ScaledMm` geometry"); + "CUTLASS cannot implement the requested `CutlassScaledMm` geometry"); status = gemm(arguments, nullptr, static_cast(stream)); assert(status == cutlass::Status::kSuccess && - "CUTLASS failed to execute `ScaledMm`"); + "CUTLASS failed to execute `CutlassScaledMm`"); const auto launch_status{cudaGetLastError()}; assert(launch_status == cudaSuccess && @@ -199,10 +199,10 @@ void DispatchScaleLayout(const Tensor a, const Tensor b, const Tensor scale_a, } // namespace -CutlassScaledMm::CutlassScaledMm(const Tensor a, const Tensor b, - const Tensor scale_a, const Tensor scale_b, - std::optional bias, Tensor out) - : ScaledMm{a, b, scale_a, scale_b, bias, out}, +Operator::Operator( + const Tensor a, const Tensor b, const Tensor scale_a, const Tensor scale_b, + std::optional bias, Tensor out) + : CutlassScaledMm{a, b, scale_a, scale_b, bias, out}, device_index_{a.device().index()} { cudaDeviceProp properties{}; const auto status{cudaGetDeviceProperties(&properties, device_index_)}; @@ -213,9 +213,9 @@ CutlassScaledMm::CutlassScaledMm(const Tensor a, const Tensor b, "`CutlassScaledMm` requires compute capability 7.5 or newer"); } -void CutlassScaledMm::operator()(const Tensor a, const Tensor b, - const Tensor scale_a, const Tensor scale_b, - std::optional bias, Tensor out) const { +void Operator::operator()( + const Tensor a, const Tensor b, const Tensor scale_a, const Tensor scale_b, + std::optional bias, Tensor out) const { DeviceGuard device_guard{device_index_}; switch (out_dtype_) { 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..a4f7ef462 --- /dev/null +++ b/src/native/cuda/nvidia/ops/cutlass_scaled_mm/cutlass.h @@ -0,0 +1,27 @@ +#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, std::optional bias, Tensor out); + + void operator()(const Tensor a, const Tensor b, const Tensor scale_a, + const Tensor scale_b, 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/native/cuda/nvidia/ops/scaled_mm/cutlass.h b/src/native/cuda/nvidia/ops/scaled_mm/cutlass.h deleted file mode 100644 index 4fae417e9..000000000 --- a/src/native/cuda/nvidia/ops/scaled_mm/cutlass.h +++ /dev/null @@ -1,31 +0,0 @@ -#ifndef INFINI_OPS_NVIDIA_SCALED_MM_CUTLASS_H_ -#define INFINI_OPS_NVIDIA_SCALED_MM_CUTLASS_H_ - -#include - -#include "base/scaled_mm.h" - -namespace infini::ops { - -class CutlassScaledMm : public ScaledMm { - public: - CutlassScaledMm(const Tensor a, const Tensor b, const Tensor scale_a, - const Tensor scale_b, std::optional bias, Tensor out); - - void operator()(const Tensor a, const Tensor b, const Tensor scale_a, - const Tensor scale_b, std::optional bias, - Tensor out) const override; - - private: - int device_index_{0}; -}; - -template <> -class Operator : public CutlassScaledMm { - public: - using CutlassScaledMm::CutlassScaledMm; -}; - -} // namespace infini::ops - -#endif // INFINI_OPS_NVIDIA_SCALED_MM_CUTLASS_H_ diff --git a/tests/conftest.py b/tests/conftest.py index 5816e2497..d49cfd340 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -198,7 +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_scaled_mm.py": _is_smoke_scaled_mm_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, } @@ -384,7 +384,7 @@ def _is_smoke_rms_norm_case(params): ) -def _is_smoke_scaled_mm_case(params): +def _is_smoke_cutlass_scaled_mm_case(params): return ( params.get("m") == 16 and params.get("n") == 32 diff --git a/tests/test_scaled_mm.py b/tests/test_cutlass_scaled_mm.py similarity index 86% rename from tests/test_scaled_mm.py rename to tests/test_cutlass_scaled_mm.py index 21a3856c0..07bd0624e 100644 --- a/tests/test_scaled_mm.py +++ b/tests/test_cutlass_scaled_mm.py @@ -5,8 +5,10 @@ from tests.utils import Payload, empty_strided, get_stream, randint_strided -if not hasattr(infini.ops, "ScaledMm"): - pytest.skip("`ScaledMm` is not available on this platform", allow_module_level=True) +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 @@ -24,7 +26,7 @@ (torch.bfloat16, 1e-2, 3e-1), ), ) -def test_scaled_mm( +def test_cutlass_scaled_mm( m, n, k, @@ -52,8 +54,8 @@ def test_scaled_mm( out = empty_strided((m, n), out_strides, dtype=dtype, device=device) return Payload( - _scaled_mm, - _torch_scaled_mm, + _cutlass_scaled_mm, + _torch_cutlass_scaled_mm, (a, b, scale_a, scale_b, bias, out), {"implementation_index": implementation_index}, rtol=rtol, @@ -61,8 +63,8 @@ def test_scaled_mm( ) -def _scaled_mm(a, b, scale_a, scale_b, bias, out, *, implementation_index): - infini.ops.scaled_mm( +def _cutlass_scaled_mm(a, b, scale_a, scale_b, bias, out, *, implementation_index): + infini.ops.cutlass_scaled_mm( a, b, scale_a, @@ -76,7 +78,9 @@ def _scaled_mm(a, b, scale_a, scale_b, bias, out, *, implementation_index): return out -def _torch_scaled_mm(a, b, scale_a, scale_b, bias, out, *, implementation_index): +def _torch_cutlass_scaled_mm( + a, b, scale_a, scale_b, bias, out, *, implementation_index +): del implementation_index result = torch.matmul(a.float(), b.float()) * scale_a * scale_b @@ -89,7 +93,7 @@ def _torch_scaled_mm(a, b, scale_a, scale_b, bias, out, *, implementation_index) return out -def test_scaled_mm_non_default_stream(device, implementation_index): +def test_cutlass_scaled_mm_non_default_stream(device, implementation_index): if device != "cuda": pytest.skip("non-default CUDA streams require the NVIDIA backend") @@ -103,7 +107,7 @@ def test_scaled_mm_non_default_stream(device, implementation_index): stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): - _scaled_mm( + _cutlass_scaled_mm( a, b, scale_a, @@ -118,7 +122,7 @@ def test_scaled_mm_non_default_stream(device, implementation_index): torch.testing.assert_close(out, expected.to(out.dtype), rtol=1e-2, atol=3e-1) -def test_scaled_mm_multi_gpu_device_guard(device, implementation_index): +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") @@ -136,7 +140,7 @@ def test_scaled_mm_multi_gpu_device_guard(device, implementation_index): stream = torch.cuda.Stream(device=target_device) stream.wait_stream(torch.cuda.current_stream(target_device)) - infini.ops.scaled_mm( + infini.ops.cutlass_scaled_mm( a, b, scale_a, From a9b65dd74409a07c238ad023d98fafe705ce0492 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 21 Jul 2026 12:47:18 +0800 Subject: [PATCH 4/5] fix(nvidia): align cutlass_scaled_mm out_dtype --- scripts/generate_wrappers.py | 33 ++++++++++++-- src/base/cutlass_scaled_mm.h | 11 +++-- .../nvidia/ops/cutlass_scaled_mm/cutlass.cu | 13 ++++-- .../nvidia/ops/cutlass_scaled_mm/cutlass.h | 7 +-- src/pybind11_utils.h | 13 ++++-- tests/test_cutlass_scaled_mm.py | 29 +++++++++++-- tests/test_generate_wrappers.py | 43 +++++++++++++++++++ 7 files changed, 127 insertions(+), 22 deletions(-) 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/base/cutlass_scaled_mm.h b/src/base/cutlass_scaled_mm.h index e69598015..be1d812a1 100644 --- a/src/base/cutlass_scaled_mm.h +++ b/src/base/cutlass_scaled_mm.h @@ -12,7 +12,8 @@ namespace infini::ops { class CutlassScaledMm : public Operator { public: CutlassScaledMm(const Tensor a, const Tensor b, const Tensor scale_a, - const Tensor scale_b, std::optional bias, Tensor out) + 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}, @@ -21,7 +22,7 @@ class CutlassScaledMm : public Operator { ldo_{out.ndim() == 2 ? out.stride(0) : 0}, scale_a_size_{scale_a.numel()}, scale_b_size_{scale_b.numel()}, - out_dtype_{out.dtype()} { + 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 && @@ -29,6 +30,8 @@ class CutlassScaledMm : public Operator { 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"); @@ -80,8 +83,8 @@ class CutlassScaledMm : public Operator { } virtual void operator()(const Tensor a, const Tensor b, const Tensor scale_a, - const Tensor scale_b, std::optional bias, - Tensor out) const = 0; + const Tensor scale_b, const DataType out_dtype, + std::optional bias, Tensor out) const = 0; protected: Tensor::Size m_{0}; diff --git a/src/native/cuda/nvidia/ops/cutlass_scaled_mm/cutlass.cu b/src/native/cuda/nvidia/ops/cutlass_scaled_mm/cutlass.cu index f91fb6cbe..824243a84 100644 --- a/src/native/cuda/nvidia/ops/cutlass_scaled_mm/cutlass.cu +++ b/src/native/cuda/nvidia/ops/cutlass_scaled_mm/cutlass.cu @@ -201,8 +201,8 @@ void DispatchScaleLayout(const Tensor a, const Tensor b, const Tensor scale_a, Operator::Operator( const Tensor a, const Tensor b, const Tensor scale_a, const Tensor scale_b, - std::optional bias, Tensor out) - : CutlassScaledMm{a, b, scale_a, scale_b, bias, out}, + 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_)}; @@ -215,7 +215,14 @@ Operator::Operator( void Operator::operator()( const Tensor a, const Tensor b, const Tensor scale_a, const Tensor scale_b, - std::optional bias, Tensor out) const { + 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_) { diff --git a/src/native/cuda/nvidia/ops/cutlass_scaled_mm/cutlass.h b/src/native/cuda/nvidia/ops/cutlass_scaled_mm/cutlass.h index a4f7ef462..93a6361ec 100644 --- a/src/native/cuda/nvidia/ops/cutlass_scaled_mm/cutlass.h +++ b/src/native/cuda/nvidia/ops/cutlass_scaled_mm/cutlass.h @@ -12,11 +12,12 @@ class Operator : public CutlassScaledMm { public: Operator(const Tensor a, const Tensor b, const Tensor scale_a, - const Tensor scale_b, std::optional bias, Tensor out); + 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, std::optional bias, - Tensor out) const override; + const Tensor scale_b, const DataType out_dtype, + std::optional bias, Tensor out) const override; private: int device_index_{0}; 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/test_cutlass_scaled_mm.py b/tests/test_cutlass_scaled_mm.py index 07bd0624e..4f229c942 100644 --- a/tests/test_cutlass_scaled_mm.py +++ b/tests/test_cutlass_scaled_mm.py @@ -56,19 +56,30 @@ def test_cutlass_scaled_mm( return Payload( _cutlass_scaled_mm, _torch_cutlass_scaled_mm, - (a, b, scale_a, scale_b, bias, out), + (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, bias, out, *, implementation_index): +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), @@ -79,7 +90,15 @@ def _cutlass_scaled_mm(a, b, scale_a, scale_b, bias, out, *, implementation_inde def _torch_cutlass_scaled_mm( - a, b, scale_a, scale_b, bias, out, *, implementation_index + a, + b, + scale_a, + scale_b, + out_dtype, + bias, + out, + *, + implementation_index, ): del implementation_index @@ -88,7 +107,7 @@ def _torch_cutlass_scaled_mm( if bias is not None: result = result + bias.float() - out.copy_(result.to(out.dtype)) + out.copy_(result.to(out_dtype)) return out @@ -112,6 +131,7 @@ def test_cutlass_scaled_mm_non_default_stream(device, implementation_index): b, scale_a, scale_b, + out.dtype, None, out, implementation_index=implementation_index, @@ -145,6 +165,7 @@ def test_cutlass_scaled_mm_multi_gpu_device_guard(device, implementation_index): b, scale_a, scale_b, + out.dtype, None, out, stream=stream.cuda_stream, 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() From f19b65463406f487663822cb38c7b725a457f1bb Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 21 Jul 2026 15:22:53 +0800 Subject: [PATCH 5/5] build(nvidia): fetch CUTLASS with CMake --- .github/ci_config.yml | 2 +- .gitmodules | 4 ---- src/CMakeLists.txt | 44 ++++++++++++++++++++++++++++++++++++------- third_party/cutlass | 1 - 4 files changed, 38 insertions(+), 13 deletions(-) delete mode 160000 third_party/cutlass diff --git a/.github/ci_config.yml b/.github/ci_config.yml index fa5ce2d25..65b2b208f 100644 --- a/.github/ci_config.yml +++ b/.github/ci_config.yml @@ -15,7 +15,7 @@ platforms: BASE_IMAGE: nvcr.io/nvidia/pytorch:24.10-py3 SKIP_APT: "1" PIP_INDEX_URL: https://pypi.tuna.tsinghua.edu.cn/simple - setup: git submodule update --init --checkout --depth 1 third_party/cutlass && pip install scikit-build-core pybind11 libclang pyyaml && rm -rf /tmp/infinirt-src /tmp/infinirt-build /tmp/infinirt-prefix && git clone --depth 1 ${INFINI_RT_GIT_URL:-https://gh-proxy.com/https://github.com/InfiniTensor/InfiniRT.git} /tmp/infinirt-src && cmake -S /tmp/infinirt-src -B /tmp/infinirt-build -DWITH_CPU=ON -DWITH_NVIDIA=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/tmp/infinirt-prefix && cmake --build /tmp/infinirt-build -j$(nproc) && cmake --install /tmp/infinirt-build && pip install .[dev] --no-build-isolation --config-settings=cmake.define.INFINI_OPS_SMOKE_BUILD=ON --config-settings=cmake.define.INFINI_RT_ROOT=/tmp/infinirt-prefix + setup: pip install scikit-build-core pybind11 libclang pyyaml && rm -rf /tmp/infinirt-src /tmp/infinirt-build /tmp/infinirt-prefix && git clone --depth 1 ${INFINI_RT_GIT_URL:-https://gh-proxy.com/https://github.com/InfiniTensor/InfiniRT.git} /tmp/infinirt-src && cmake -S /tmp/infinirt-src -B /tmp/infinirt-build -DWITH_CPU=ON -DWITH_NVIDIA=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/tmp/infinirt-prefix && cmake --build /tmp/infinirt-build -j$(nproc) && cmake --install /tmp/infinirt-build && pip install .[dev] --no-build-isolation --config-settings=cmake.define.INFINI_OPS_SMOKE_BUILD=ON --config-settings=cmake.define.INFINI_RT_ROOT=/tmp/infinirt-prefix jobs: gpu: type: unittest diff --git a/.gitmodules b/.gitmodules index bd4db8160..be99e8a8f 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,7 +1,3 @@ [submodule ".ci"] path = .ci url = https://github.com/InfiniTensor/ci.git -[submodule "third_party/cutlass"] - path = third_party/cutlass - url = https://github.com/NVIDIA/cutlass.git - ignore = untracked diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 852f5ce98..9c30701f8 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -67,12 +67,42 @@ if(WITH_CPU) endif() if(WITH_NVIDIA) - set(_cutlass_include_dir - "${PROJECT_SOURCE_DIR}/third_party/cutlass/include") - if(NOT EXISTS "${_cutlass_include_dir}/cutlass/cutlass.h") - message(FATAL_ERROR - "CUTLASS is required by the NVIDIA backend. Run " - "`git submodule update --init --checkout third_party/cutlass`.") + 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 @@ -99,7 +129,7 @@ if(WITH_NVIDIA) target_compile_definitions(infiniops PUBLIC WITH_NVIDIA=1) target_sources(infiniops PRIVATE ${NVIDIA_SOURCES}) - target_include_directories(infiniops PRIVATE "${_cutlass_include_dir}") + 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) diff --git a/third_party/cutlass b/third_party/cutlass deleted file mode 160000 index 087c84df8..000000000 --- a/third_party/cutlass +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 087c84df83d254b5fb295a7a408f1a1d554085cf