diff --git a/qutlass/__init__.py b/qutlass/__init__.py index e704007..7adce8b 100644 --- a/qutlass/__init__.py +++ b/qutlass/__init__.py @@ -15,12 +15,14 @@ # import torch -import qutlass._CUDA +import qutlass._CUDA # noqa: F401 - registers torch.ops._qutlass_C from qutlass.utils import get_padded_shape_mx, get_padded_shape_nv, pad_to_block from typing import Literal import warnings +qutlass_CUDA = torch.ops._qutlass_C + try: from flashinfer import mm_fp4 @@ -38,7 +40,7 @@ def matmul_mxf4_bf16_tn( backend: Literal["cutlass", "flashinfer"] = "cutlass", ) -> torch.Tensor: if backend == "cutlass": - return qutlass._CUDA.matmul_mxf4_bf16_tn(a, b, a_sf, b_sf, alpha) + return qutlass_CUDA.matmul_mxf4_bf16_tn(a, b, a_sf, b_sf, alpha) elif backend == "flashinfer": if not _HAS_FLASHINFER: raise ImportError( @@ -81,7 +83,7 @@ def matmul_ada_mxf4_bf16_tn( b_sf: torch.Tensor, alpha: torch.Tensor, ) -> torch.Tensor: - return qutlass._CUDA.matmul_ada_mxf4_bf16_tn(a, b, a_sf, b_sf, alpha) + return qutlass_CUDA.matmul_ada_mxf4_bf16_tn(a, b, a_sf, b_sf, alpha) def matmul_nvf4_bf16_tn( @@ -93,7 +95,7 @@ def matmul_nvf4_bf16_tn( backend: Literal["cutlass", "flashinfer"] = "cutlass", ) -> torch.Tensor: if backend == "cutlass": - return qutlass._CUDA.matmul_nvf4_bf16_tn(a, b, a_sf, b_sf, alpha) + return qutlass_CUDA.matmul_nvf4_bf16_tn(a, b, a_sf, b_sf, alpha) elif backend == "flashinfer": if not _HAS_FLASHINFER: raise ImportError( @@ -134,14 +136,14 @@ def matmul_mxf8_bf16_tn(a: torch.Tensor, block_scale_a: torch.Tensor, block_scale_b: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor: - return qutlass._CUDA.matmul_mxf8_bf16_tn(a, b, block_scale_a, block_scale_b, alpha) + return qutlass_CUDA.matmul_mxf8_bf16_tn(a, b, block_scale_a, block_scale_b, alpha) def matmul_mxf8_bf16_nn(a: torch.Tensor, b: torch.Tensor, block_scale_a: torch.Tensor, block_scale_b: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor: - return qutlass._CUDA.matmul_mxf8_bf16_nn(a, b, block_scale_a, block_scale_b, alpha) + return qutlass_CUDA.matmul_mxf8_bf16_nn(a, b, block_scale_a, block_scale_b, alpha) def fusedQuantizeMx( @@ -165,15 +167,15 @@ def fusedQuantizeMx( clip_mask = torch.empty( *a.shape[:-1], a.size(-1) // 8, dtype=torch.uint8, device=a.device ) - return qutlass._CUDA.fusedQuantizeMxQuestWithMask( + return qutlass_CUDA.fusedQuantizeMxQuestWithMask( a, b, xh_e2m1, xh_e8m0, clip_mask ) else: - return qutlass._CUDA.fusedQuantizeMxQuest(a, b, xh_e2m1, xh_e8m0) + return qutlass_CUDA.fusedQuantizeMxQuest(a, b, xh_e2m1, xh_e8m0) elif method == "abs_max": if return_mask: raise ValueError("return_mask is only supported for method 'quest'") - return qutlass._CUDA.fusedQuantizeMxAbsMax(a, b, xh_e2m1, xh_e8m0) + return qutlass_CUDA.fusedQuantizeMxAbsMax(a, b, xh_e2m1, xh_e8m0) else: raise ValueError(f"invalid method {method!r}, must be 'quest' or 'abs_max'") @@ -194,9 +196,9 @@ def fusedQuantizeNv( ) if method == "quest": - return qutlass._CUDA.fusedQuantizeNvQuest(a, b, xh_e2m1, xh_e4m3, global_scale) + return qutlass_CUDA.fusedQuantizeNvQuest(a, b, xh_e2m1, xh_e4m3, global_scale) elif method == "abs_max": - return qutlass._CUDA.fusedQuantizeNvAbsMax(a, b, xh_e2m1, xh_e4m3, global_scale) + return qutlass_CUDA.fusedQuantizeNvAbsMax(a, b, xh_e2m1, xh_e4m3, global_scale) else: raise ValueError(f"invalid method {method!r}, must be 'quest' or 'abs_max'") @@ -236,7 +238,7 @@ def backward_t_bf16( and xh_e8m0.is_contiguous() ) - qutlass._CUDA.backward_t_bf16(x, h, xh_e2m1, xh_e8m0) + qutlass_CUDA.backward_t_bf16(x, h, xh_e2m1, xh_e8m0) return xh_e2m1, xh_e8m0 @@ -275,7 +277,7 @@ def backward_qt_bf16( and xh_e8m0.is_contiguous() ) - qutlass._CUDA.backward_qt_bf16(x_e2m1, x_e8m0, h, alpha, xh_e2m1, xh_e8m0) + qutlass_CUDA.backward_qt_bf16(x_e2m1, x_e8m0, h, alpha, xh_e2m1, xh_e8m0) return xh_e2m1, xh_e8m0 @@ -286,7 +288,7 @@ def backward_bf16_square_double_mxfp8(x_bf16: torch.Tensor) -> tuple[torch.Tenso row_scales = torch.empty(x_bf16.shape[0], x_bf16.shape[1] // 32, device=x_bf16.device, dtype=torch.float8_e8m0fnu) column_scales = torch.empty(x_bf16.shape[1], x_bf16.shape[0] // 32, device=x_bf16.device, dtype=torch.float8_e8m0fnu) - qutlass._CUDA.backward_bf16_square_double_mxfp8(x_bf16, x_fp8, row_scales, column_scales) + qutlass_CUDA.backward_bf16_square_double_mxfp8(x_bf16, x_fp8, row_scales, column_scales) return x_fp8, row_scales, column_scales @@ -303,6 +305,6 @@ def mxfp4_transpose_mxfp8(x_fp4: torch.Tensor, scales: torch.Tensor) -> tuple[to x_fp8 = torch.empty(x_fp4.shape[1] * 2, x_fp4.shape[0], device=x_fp4.device, dtype=torch.float8_e4m3fn) shared_exps = torch.empty(x_fp4.shape[1] * 2, x_fp4.shape[0] // 32, device=x_fp4.device, dtype=torch.float8_e8m0fnu) - qutlass._CUDA.mxfp4_transpose_mxfp8(x_fp4, scales, x_fp8, shared_exps) + qutlass_CUDA.mxfp4_transpose_mxfp8(x_fp4, scales, x_fp8, shared_exps) return x_fp8, shared_exps \ No newline at end of file diff --git a/qutlass/csrc/bindings.cpp b/qutlass/csrc/bindings.cpp index 600c8f3..26e5a3f 100644 --- a/qutlass/csrc/bindings.cpp +++ b/qutlass/csrc/bindings.cpp @@ -14,527 +14,527 @@ * limitations under the License. */ -#include -#include -#include -#include -#include - -#ifndef QUTLASS_DISABLE_PYBIND -#include -#endif +#include +#include -#include -#include #include +#include "include/bindings_utils.h" #include "include/gemm.h" #include "include/fused_quantize_host.h" #include "include/backward_host.h" namespace QUTLASS { -torch::Tensor matmul_mxf4_bf16_tn(torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& A_sf, - torch::Tensor const& B_sf, - torch::Tensor const& alpha) +using torch::stable::Tensor; +using torch::headeronly::ScalarType; + +Tensor matmul_mxf4_bf16_tn(Tensor const& A, + Tensor const& B, + Tensor const& A_sf, + Tensor const& B_sf, + Tensor const& alpha) { - torch::checkAllContiguous("matmul_mxf4_bf16_tn", {{A, "A", 0}, - {B, "B", 1}, - {A_sf, "A_sf", 2}, - {B_sf, "B_sf", 3}}); - torch::checkDeviceType("matmul_mxf4_bf16_tn", {A, B, A_sf, B_sf, alpha}, at::DeviceType::CUDA); - torch::checkAllSameGPU("matmul_mxf4_bf16_tn", {{A, "A", 0}, - {B, "B", 1}, - {A_sf, "A_sf", 2}, - {B_sf, "B_sf", 3}, - {alpha, "alpha", 4}}); - TORCH_CHECK(A.scalar_type() == at::kByte, "A must be uint8"); - TORCH_CHECK(B.scalar_type() == at::kByte, "B must be uint8"); - TORCH_CHECK(A_sf.scalar_type() == at::kFloat8_e8m0fnu, "A_sf must be float8_e8m0fnu"); - TORCH_CHECK(B_sf.scalar_type() == at::kFloat8_e8m0fnu, "B_sf must be float8_e8m0fnu"); - TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "A and B must be 2D"); - TORCH_CHECK(A.size(1) == B.size(1), "Inner dimensions must match for A @ B.T"); - TORCH_CHECK(A.size(1) >= 32, "A K-dim must be >= 32"); - TORCH_CHECK(B.size(1) >= 32, "B K-dim must be >= 32"); + check_all_contiguous("matmul_mxf4_bf16_tn", {{A, "A", 0}, + {B, "B", 1}, + {A_sf, "A_sf", 2}, + {B_sf, "B_sf", 3}}); + check_device_type_cuda("matmul_mxf4_bf16_tn", {A, B, A_sf, B_sf, alpha}); + check_all_same_gpu("matmul_mxf4_bf16_tn", {{A, "A", 0}, + {B, "B", 1}, + {A_sf, "A_sf", 2}, + {B_sf, "B_sf", 3}, + {alpha, "alpha", 4}}); + STD_TORCH_CHECK(A.scalar_type() == ScalarType::Byte, "A must be uint8"); + STD_TORCH_CHECK(B.scalar_type() == ScalarType::Byte, "B must be uint8"); + STD_TORCH_CHECK(A_sf.scalar_type() == ScalarType::Float8_e8m0fnu, + "A_sf must be float8_e8m0fnu"); + STD_TORCH_CHECK(B_sf.scalar_type() == ScalarType::Float8_e8m0fnu, + "B_sf must be float8_e8m0fnu"); + STD_TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "A and B must be 2D"); + STD_TORCH_CHECK(A.size(1) == B.size(1), "Inner dimensions must match for A @ B.T"); + STD_TORCH_CHECK(A.size(1) >= 32, "A K-dim must be >= 32"); + STD_TORCH_CHECK(B.size(1) >= 32, "B K-dim must be >= 32"); uint32_t M = A.size(0); uint32_t N = B.size(0); - auto OUT = torch::empty({M, N}, torch::dtype(torch::kBFloat16).device(A.device())); + auto OUT = torch::stable::new_empty(A, {M, N}, ScalarType::BFloat16); matmul_host_mxf4_bf16_tn(OUT, A, B, A_sf, B_sf, alpha); return OUT; } -torch::Tensor matmul_nvf4_bf16_tn(torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& A_sf, - torch::Tensor const& B_sf, - torch::Tensor const& alpha) +Tensor matmul_nvf4_bf16_tn(Tensor const& A, + Tensor const& B, + Tensor const& A_sf, + Tensor const& B_sf, + Tensor const& alpha) { - torch::checkAllContiguous("matmul_nvf4_bf16_tn", {{A, "A", 0}, - {B, "B", 1}, - {A_sf, "A_sf", 2}, - {B_sf, "B_sf", 3}}); - torch::checkDeviceType("matmul_nvf4_bf16_tn", {A, B, A_sf, B_sf, alpha}, at::DeviceType::CUDA); - torch::checkAllSameGPU("matmul_nvf4_bf16_tn", {{A, "A", 0}, - {B, "B", 1}, - {A_sf, "A_sf", 2}, - {B_sf, "B_sf", 3}, - {alpha, "alpha", 4}}); - TORCH_CHECK(A.scalar_type() == at::kByte, "A must be uint8"); - TORCH_CHECK(B.scalar_type() == at::kByte, "B must be uint8"); - TORCH_CHECK(A_sf.scalar_type() == at::kFloat8_e4m3fn, "A_sf must be float8_e4m3fn"); - TORCH_CHECK(B_sf.scalar_type() == at::kFloat8_e4m3fn, "B_sf must be float8_e4m3fn"); - TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "A and B must be 2D"); - TORCH_CHECK(A.size(1) == B.size(1), "Inner dimensions must match for A @ B.T"); - TORCH_CHECK(A.size(1) >= 16, "A K-dim must be >= 16"); - TORCH_CHECK(B.size(1) >= 16, "B K-dim must be >= 16"); + check_all_contiguous("matmul_nvf4_bf16_tn", {{A, "A", 0}, + {B, "B", 1}, + {A_sf, "A_sf", 2}, + {B_sf, "B_sf", 3}}); + check_device_type_cuda("matmul_nvf4_bf16_tn", {A, B, A_sf, B_sf, alpha}); + check_all_same_gpu("matmul_nvf4_bf16_tn", {{A, "A", 0}, + {B, "B", 1}, + {A_sf, "A_sf", 2}, + {B_sf, "B_sf", 3}, + {alpha, "alpha", 4}}); + STD_TORCH_CHECK(A.scalar_type() == ScalarType::Byte, "A must be uint8"); + STD_TORCH_CHECK(B.scalar_type() == ScalarType::Byte, "B must be uint8"); + STD_TORCH_CHECK(A_sf.scalar_type() == ScalarType::Float8_e4m3fn, + "A_sf must be float8_e4m3fn"); + STD_TORCH_CHECK(B_sf.scalar_type() == ScalarType::Float8_e4m3fn, + "B_sf must be float8_e4m3fn"); + STD_TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "A and B must be 2D"); + STD_TORCH_CHECK(A.size(1) == B.size(1), "Inner dimensions must match for A @ B.T"); + STD_TORCH_CHECK(A.size(1) >= 16, "A K-dim must be >= 16"); + STD_TORCH_CHECK(B.size(1) >= 16, "B K-dim must be >= 16"); uint32_t M = A.size(0); uint32_t N = B.size(0); - auto OUT = torch::empty({M, N}, torch::dtype(torch::kBFloat16).device(A.device())); + auto OUT = torch::stable::new_empty(A, {M, N}, ScalarType::BFloat16); matmul_host_nvf4_bf16_tn(OUT, A, B, A_sf, B_sf, alpha); return OUT; } -torch::Tensor matmul_ada_mxf4_bf16_tn(torch::Tensor const&A, - torch::Tensor const&B, - torch::Tensor const&A_sf, - torch::Tensor const&B_sf, - torch::Tensor const& alpha) +Tensor matmul_ada_mxf4_bf16_tn(Tensor const& A, + Tensor const& B, + Tensor const& A_sf, + Tensor const& B_sf, + Tensor const& alpha) { - torch::checkAllContiguous("matmul_ada_mxf4_bf16_tn", {{A, "A", 0}, - {B, "B", 1}, - {A_sf, "A_sf", 2}, - {B_sf, "B_sf", 3}}); - torch::checkDeviceType("matmul_ada_mxf4_bf16_tn", {A, B, A_sf, B_sf, alpha}, at::DeviceType::CUDA); - torch::checkAllSameGPU("matmul_ada_mxf4_bf16_tn", {{A, "A", 0}, - {B, "B", 1}, - {A_sf, "A_sf", 2}, - {B_sf, "B_sf", 3}, - {alpha, "alpha", 4}}); - TORCH_CHECK(A.scalar_type() == at::kByte, "A must be uint8"); - TORCH_CHECK(B.scalar_type() == at::kByte, "B must be uint8"); - TORCH_CHECK(A_sf.scalar_type() == at::kFloat8_e8m0fnu, "A_sf must be float8_e8m0fnu"); - TORCH_CHECK(B_sf.scalar_type() == at::kFloat8_e8m0fnu, "B_sf must be float8_e8m0fnu"); - TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "A and B must be 2D"); - TORCH_CHECK(A.size(1) == B.size(1), "Inner dimensions must match for A @ B.T"); - TORCH_CHECK(A.size(1) >= 32, "A K-dim must be >= 32"); - TORCH_CHECK(B.size(1) >= 32, "B K-dim must be >= 32"); + check_all_contiguous("matmul_ada_mxf4_bf16_tn", {{A, "A", 0}, + {B, "B", 1}, + {A_sf, "A_sf", 2}, + {B_sf, "B_sf", 3}}); + check_device_type_cuda("matmul_ada_mxf4_bf16_tn", {A, B, A_sf, B_sf, alpha}); + check_all_same_gpu("matmul_ada_mxf4_bf16_tn", {{A, "A", 0}, + {B, "B", 1}, + {A_sf, "A_sf", 2}, + {B_sf, "B_sf", 3}, + {alpha, "alpha", 4}}); + STD_TORCH_CHECK(A.scalar_type() == ScalarType::Byte, "A must be uint8"); + STD_TORCH_CHECK(B.scalar_type() == ScalarType::Byte, "B must be uint8"); + STD_TORCH_CHECK(A_sf.scalar_type() == ScalarType::Float8_e8m0fnu, + "A_sf must be float8_e8m0fnu"); + STD_TORCH_CHECK(B_sf.scalar_type() == ScalarType::Float8_e8m0fnu, + "B_sf must be float8_e8m0fnu"); + STD_TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "A and B must be 2D"); + STD_TORCH_CHECK(A.size(1) == B.size(1), "Inner dimensions must match for A @ B.T"); + STD_TORCH_CHECK(A.size(1) >= 32, "A K-dim must be >= 32"); + STD_TORCH_CHECK(B.size(1) >= 32, "B K-dim must be >= 32"); uint32_t M = A.size(0); uint32_t N = B.size(0); - auto C = torch::empty({M, N}, torch::dtype(torch::kBFloat16).device(A.device())); + auto C = torch::stable::new_empty(A, {M, N}, ScalarType::BFloat16); matmul_host_ada_mxf4_bf16_tn(A, B, A_sf, B_sf, C, alpha); return C; } -torch::Tensor matmul_mxf8_bf16_tn(torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& A_sf, - torch::Tensor const& B_sf, - torch::Tensor const& alpha) +Tensor matmul_mxf8_bf16_tn(Tensor const& A, + Tensor const& B, + Tensor const& A_sf, + Tensor const& B_sf, + Tensor const& alpha) { - torch::checkAllContiguous("matmul_mxf8_bf16_tn", {{A, "A", 0}, - {B, "B", 1}, - {A_sf, "A_sf", 2}, - {B_sf, "B_sf", 3}, - {alpha, "alpha", 4}}); - torch::checkDeviceType("matmul_mxf8_bf16_tn", {A, B, A_sf, B_sf, alpha}, at::DeviceType::CUDA); - - torch::checkAllSameGPU("matmul_mxf8_bf16_tn", {{A, "A", 0}, - {B, "B", 1}, - {A_sf, "A_sf", 2}, - {B_sf, "B_sf", 3}, - {alpha, "alpha", 4}}); - - TORCH_CHECK(A.scalar_type() == at::kFloat8_e4m3fn, "A must be float8_e4m3fn"); - TORCH_CHECK(B.scalar_type() == at::kFloat8_e4m3fn, "B must be float8_e4m3fn"); - TORCH_CHECK(A_sf.scalar_type() == at::kFloat8_e8m0fnu, "A_sf must be float8_e8m0fnu"); - TORCH_CHECK(B_sf.scalar_type() == at::kFloat8_e8m0fnu, "B_sf must be float8_e8m0fnu"); - TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "A and B must be 2D"); - TORCH_CHECK(A.size(1) == B.size(1), "Inner dimensions must match for A @ B.T"); - TORCH_CHECK(A.size(1) >= 32, "A K-dim must be >= 32"); - TORCH_CHECK(B.size(1) >= 32, "B K-dim must be >= 32"); + check_all_contiguous("matmul_mxf8_bf16_tn", {{A, "A", 0}, + {B, "B", 1}, + {A_sf, "A_sf", 2}, + {B_sf, "B_sf", 3}, + {alpha, "alpha", 4}}); + check_device_type_cuda("matmul_mxf8_bf16_tn", {A, B, A_sf, B_sf, alpha}); + check_all_same_gpu("matmul_mxf8_bf16_tn", {{A, "A", 0}, + {B, "B", 1}, + {A_sf, "A_sf", 2}, + {B_sf, "B_sf", 3}, + {alpha, "alpha", 4}}); + STD_TORCH_CHECK(A.scalar_type() == ScalarType::Float8_e4m3fn, + "A must be float8_e4m3fn"); + STD_TORCH_CHECK(B.scalar_type() == ScalarType::Float8_e4m3fn, + "B must be float8_e4m3fn"); + STD_TORCH_CHECK(A_sf.scalar_type() == ScalarType::Float8_e8m0fnu, + "A_sf must be float8_e8m0fnu"); + STD_TORCH_CHECK(B_sf.scalar_type() == ScalarType::Float8_e8m0fnu, + "B_sf must be float8_e8m0fnu"); + STD_TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "A and B must be 2D"); + STD_TORCH_CHECK(A.size(1) == B.size(1), "Inner dimensions must match for A @ B.T"); + STD_TORCH_CHECK(A.size(1) >= 32, "A K-dim must be >= 32"); + STD_TORCH_CHECK(B.size(1) >= 32, "B K-dim must be >= 32"); uint32_t M = A.size(0); uint32_t N = B.size(0); - auto OUT = torch::empty({M, N}, torch::dtype(torch::kBFloat16).device(A.device())); + auto OUT = torch::stable::new_empty(A, {M, N}, ScalarType::BFloat16); matmul_host_mxf8_bf16_tn(OUT, A, B, A_sf, B_sf, alpha); return OUT; } -torch::Tensor matmul_mxf8_bf16_nn(torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& A_sf, - torch::Tensor const& B_sf, - torch::Tensor const& alpha) +Tensor matmul_mxf8_bf16_nn(Tensor const& A, + Tensor const& B, + Tensor const& A_sf, + Tensor const& B_sf, + Tensor const& alpha) { - torch::checkAllContiguous("matmul_mxf8_bf16_nn", {{A, "A", 0}, - {B, "B", 1}, - {A_sf, "A_sf", 2}, - {B_sf, "B_sf", 3}, - {alpha, "alpha", 4}}); - torch::checkDeviceType("matmul_mxf8_bf16_nn", {A, B, A_sf, B_sf, alpha}, at::DeviceType::CUDA); - - torch::checkAllSameGPU("matmul_mxf8_bf16_nn", {{A, "A", 0}, - {B, "B", 1}, - {A_sf, "A_sf", 2}, - {B_sf, "B_sf", 3}, - {alpha, "alpha", 4}}); - - TORCH_CHECK(A.scalar_type() == at::kFloat8_e4m3fn, "A must be float8_e4m3fn"); - TORCH_CHECK(B.scalar_type() == at::kFloat8_e4m3fn, "B must be float8_e4m3fn"); - TORCH_CHECK(A_sf.scalar_type() == at::kFloat8_e8m0fnu, "A_sf must be float8_e8m0fnu"); - TORCH_CHECK(B_sf.scalar_type() == at::kFloat8_e8m0fnu, "B_sf must be float8_e8m0fnu"); - TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "A and B must be 2D"); - TORCH_CHECK(A.size(0) == B.size(1), "Inner dimensions must match for A.T @ B.T"); - TORCH_CHECK(A.size(0) >= 32, "A K-dim must be >= 32"); - TORCH_CHECK(B.size(1) >= 32, "B K-dim must be >= 32"); + check_all_contiguous("matmul_mxf8_bf16_nn", {{A, "A", 0}, + {B, "B", 1}, + {A_sf, "A_sf", 2}, + {B_sf, "B_sf", 3}, + {alpha, "alpha", 4}}); + check_device_type_cuda("matmul_mxf8_bf16_nn", {A, B, A_sf, B_sf, alpha}); + check_all_same_gpu("matmul_mxf8_bf16_nn", {{A, "A", 0}, + {B, "B", 1}, + {A_sf, "A_sf", 2}, + {B_sf, "B_sf", 3}, + {alpha, "alpha", 4}}); + STD_TORCH_CHECK(A.scalar_type() == ScalarType::Float8_e4m3fn, + "A must be float8_e4m3fn"); + STD_TORCH_CHECK(B.scalar_type() == ScalarType::Float8_e4m3fn, + "B must be float8_e4m3fn"); + STD_TORCH_CHECK(A_sf.scalar_type() == ScalarType::Float8_e8m0fnu, + "A_sf must be float8_e8m0fnu"); + STD_TORCH_CHECK(B_sf.scalar_type() == ScalarType::Float8_e8m0fnu, + "B_sf must be float8_e8m0fnu"); + STD_TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "A and B must be 2D"); + STD_TORCH_CHECK(A.size(0) == B.size(1), "Inner dimensions must match for A.T @ B.T"); + STD_TORCH_CHECK(A.size(0) >= 32, "A K-dim must be >= 32"); + STD_TORCH_CHECK(B.size(1) >= 32, "B K-dim must be >= 32"); uint32_t M = A.size(1); uint32_t N = B.size(0); - auto OUT = torch::empty({M, N}, torch::dtype(torch::kBFloat16).device(A.device())); + auto OUT = torch::stable::new_empty(A, {M, N}, ScalarType::BFloat16); matmul_host_mxf8_bf16_nn(OUT, A, B, A_sf, B_sf, alpha); return OUT; } -std::tuple fusedQuantizeMxQuest(torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor& OUT, - torch::Tensor& OUT_sf) +std::tuple fusedQuantizeMxQuest(Tensor const& A, + Tensor const& B, + Tensor& OUT, + Tensor& OUT_sf) { - torch::checkAllContiguous("fusedQuantizeMxQuest", {{A, "A", 0}, - {B, "B", 1}, - {OUT, "OUT", 2}, - {OUT_sf, "OUT_sf", 3}}); - torch::checkDeviceType("fusedQuantizeMxQuest", {A, B, OUT, OUT_sf}, at::DeviceType::CUDA); - torch::checkAllSameGPU("fusedQuantizeMxQuest", {{A, "A", 0}, - {B, "B", 1}, - {OUT, "OUT", 2}, - {OUT_sf, "OUT_sf", 3}}); - TORCH_CHECK(A.scalar_type() == at::kBFloat16, "A must be bf16"); - TORCH_CHECK(B.scalar_type() == at::kBFloat16, "B must be bf16"); - TORCH_CHECK(B.size(0) == B.size(1), "Rotation matrix must be square"); + check_all_contiguous("fusedQuantizeMxQuest", {{A, "A", 0}, + {B, "B", 1}, + {OUT, "OUT", 2}, + {OUT_sf, "OUT_sf", 3}}); + check_device_type_cuda("fusedQuantizeMxQuest", {A, B, OUT, OUT_sf}); + check_all_same_gpu("fusedQuantizeMxQuest", {{A, "A", 0}, + {B, "B", 1}, + {OUT, "OUT", 2}, + {OUT_sf, "OUT_sf", 3}}); + STD_TORCH_CHECK(A.scalar_type() == ScalarType::BFloat16, "A must be bf16"); + STD_TORCH_CHECK(B.scalar_type() == ScalarType::BFloat16, "B must be bf16"); + STD_TORCH_CHECK(B.size(0) == B.size(1), "Rotation matrix must be square"); uint32_t HAD_GS = B.size(0); - TORCH_CHECK((A.numel()%HAD_GS)==0, "A must be divisible by", HAD_GS); + STD_TORCH_CHECK((A.numel() % HAD_GS) == 0, "A must be divisible by", HAD_GS); - if(HAD_GS==32){ + if (HAD_GS == 32) { fusedQuantizeMxQuest_host(OUT, OUT_sf, A, B); - } else if(HAD_GS==64){ + } else if (HAD_GS == 64) { fusedQuantizeMxQuestHad64_host(OUT, OUT_sf, A, B); - } else if(HAD_GS==128){ + } else if (HAD_GS == 128) { fusedQuantizeMxQuestHad128_host(OUT, OUT_sf, A, B); } else { - TORCH_CHECK(false, - "Unsupported rotation size ", HAD_GS, - "; expected 32, 64, or 128."); + STD_TORCH_CHECK(false, + "Unsupported rotation size ", HAD_GS, + "; expected 32, 64, or 128."); } return std::make_tuple(OUT, OUT_sf); } -std::tuple fusedQuantizeMxQuestWithMask( - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor& OUT, - torch::Tensor& OUT_sf, - torch::Tensor& OUT_mask) +#ifndef QUTLASS_MINIMAL_BUILD +std::tuple fusedQuantizeMxQuestWithMask( + Tensor const& A, + Tensor const& B, + Tensor& OUT, + Tensor& OUT_sf, + Tensor& OUT_mask) { - torch::checkAllContiguous("fusedQuantizeMxQuestWithMask", {{A, "A", 0}, - {B, "B", 1}, - {OUT, "OUT", 2}, - {OUT_sf, "OUT_sf", 3}, - {OUT_mask, "OUT_mask", 4}}); - torch::checkDeviceType("fusedQuantizeMxQuestWithMask", {A, B, OUT, OUT_sf, OUT_mask}, at::DeviceType::CUDA); - torch::checkAllSameGPU("fusedQuantizeMxQuestWithMask", {{A, "A", 0}, - {B, "B", 1}, - {OUT, "OUT", 2}, - {OUT_sf, "OUT_sf", 3}, - {OUT_mask, "OUT_mask", 4}}); - TORCH_CHECK(A.scalar_type() == at::kBFloat16, "A must be bf16"); - TORCH_CHECK(B.scalar_type() == at::kBFloat16, "B must be bf16"); - TORCH_CHECK(B.size(0) == B.size(1), "Rotation matrix must be square"); + check_all_contiguous("fusedQuantizeMxQuestWithMask", {{A, "A", 0}, + {B, "B", 1}, + {OUT, "OUT", 2}, + {OUT_sf, "OUT_sf", 3}, + {OUT_mask, "OUT_mask", 4}}); + check_device_type_cuda("fusedQuantizeMxQuestWithMask", {A, B, OUT, OUT_sf, OUT_mask}); + check_all_same_gpu("fusedQuantizeMxQuestWithMask", {{A, "A", 0}, + {B, "B", 1}, + {OUT, "OUT", 2}, + {OUT_sf, "OUT_sf", 3}, + {OUT_mask, "OUT_mask", 4}}); + STD_TORCH_CHECK(A.scalar_type() == ScalarType::BFloat16, "A must be bf16"); + STD_TORCH_CHECK(B.scalar_type() == ScalarType::BFloat16, "B must be bf16"); + STD_TORCH_CHECK(B.size(0) == B.size(1), "Rotation matrix must be square"); uint32_t HAD_GS = B.size(0); - TORCH_CHECK((A.numel()%HAD_GS)==0, "A must be divisible by", HAD_GS); + STD_TORCH_CHECK((A.numel() % HAD_GS) == 0, "A must be divisible by", HAD_GS); - if(HAD_GS==32){ + if (HAD_GS == 32) { fusedQuantizeMxQuestWithMask_host(OUT, OUT_sf, OUT_mask, A, B); } else { - TORCH_CHECK(false, - "Unsupported rotation size ", HAD_GS, - "; expected 32."); + STD_TORCH_CHECK(false, + "Unsupported rotation size ", HAD_GS, + "; expected 32."); } return std::make_tuple(OUT, OUT_sf, OUT_mask); } +#endif // QUTLASS_MINIMAL_BUILD -std::tuple fusedQuantizeMxAbsMax(torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor& OUT, - torch::Tensor& OUT_sf) +std::tuple fusedQuantizeMxAbsMax(Tensor const& A, + Tensor const& B, + Tensor& OUT, + Tensor& OUT_sf) { - torch::checkAllContiguous("fusedQuantizeMxAbsMax", {{A, "A", 0}, - {B, "B", 1}, - {OUT, "OUT", 2}, - {OUT_sf, "OUT_sf", 3}}); - torch::checkDeviceType("fusedQuantizeMxAbsMax", {A, B, OUT, OUT_sf}, at::DeviceType::CUDA); - torch::checkAllSameGPU("fusedQuantizeMxAbsMax", {{A, "A", 0}, - {B, "B", 1}, - {OUT, "OUT", 2}, - {OUT_sf, "OUT_sf", 3}}); - TORCH_CHECK(A.scalar_type() == at::kBFloat16, "A must be bf16"); - TORCH_CHECK(B.scalar_type() == at::kBFloat16, "B must be bf16"); - TORCH_CHECK(B.size(0) == B.size(1), "Rotation matrix must be square"); + check_all_contiguous("fusedQuantizeMxAbsMax", {{A, "A", 0}, + {B, "B", 1}, + {OUT, "OUT", 2}, + {OUT_sf, "OUT_sf", 3}}); + check_device_type_cuda("fusedQuantizeMxAbsMax", {A, B, OUT, OUT_sf}); + check_all_same_gpu("fusedQuantizeMxAbsMax", {{A, "A", 0}, + {B, "B", 1}, + {OUT, "OUT", 2}, + {OUT_sf, "OUT_sf", 3}}); + STD_TORCH_CHECK(A.scalar_type() == ScalarType::BFloat16, "A must be bf16"); + STD_TORCH_CHECK(B.scalar_type() == ScalarType::BFloat16, "B must be bf16"); + STD_TORCH_CHECK(B.size(0) == B.size(1), "Rotation matrix must be square"); uint32_t HAD_GS = B.size(0); - TORCH_CHECK((A.numel()%HAD_GS)==0, "A must be divisible by", HAD_GS); + STD_TORCH_CHECK((A.numel() % HAD_GS) == 0, "A must be divisible by", HAD_GS); - if(HAD_GS==32){ + if (HAD_GS == 32) { fusedQuantizeMxAbsMax_host(OUT, OUT_sf, A, B); - } else if(HAD_GS==64){ + } else if (HAD_GS == 64) { fusedQuantizeMxAbsMaxHad64_host(OUT, OUT_sf, A, B); - } else if(HAD_GS==128){ + } else if (HAD_GS == 128) { #if TARGET_CUDA_ARCH == 100 - auto opts = torch::TensorOptions().dtype(torch::kFloat).device(A.device()); - auto global_scale = torch::tensor(0.0f, opts); //FIXME: add input global_scale to interface for consistency + // FIXME: add input global_scale to interface for consistency + auto global_scale = + torch::stable::new_zeros(A, {1}, ScalarType::Float); fusedQuantizeMxAbsMax_host_sm100(OUT, OUT_sf, A, B, global_scale); #elif TARGET_CUDA_ARCH == 120 fusedQuantizeMxAbsMaxHad128_host(OUT, OUT_sf, A, B); #endif - } else { - TORCH_CHECK(false, - "Unsupported rotation size ", HAD_GS, - "; expected 32, 64, or 128."); + STD_TORCH_CHECK(false, + "Unsupported rotation size ", HAD_GS, + "; expected 32, 64, or 128."); } return std::make_tuple(OUT, OUT_sf); } -std::tuple fusedQuantizeNvQuest(torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor& OUT, - torch::Tensor& OUT_sf, - torch::Tensor const& global_scale) +std::tuple fusedQuantizeNvQuest(Tensor const& A, + Tensor const& B, + Tensor& OUT, + Tensor& OUT_sf, + Tensor const& global_scale) { - torch::checkAllContiguous("fusedQuantizeNvQuest", {{A, "A", 0}, + check_all_contiguous("fusedQuantizeNvQuest", {{A, "A", 0}, {B, "B", 1}, {OUT, "OUT", 2}, {OUT_sf, "OUT_sf", 3}}); - torch::checkDeviceType("fusedQuantizeNvQuest", {A, B, OUT, OUT_sf, global_scale}, at::DeviceType::CUDA); - torch::checkAllSameGPU("fusedQuantizeNvQuest", {{A, "A", 0}, - {B, "B", 1}, - {OUT, "OUT", 2}, - {OUT_sf, "OUT_sf", 3}, - {global_scale, "global_scale", 4}}); - TORCH_CHECK(A.scalar_type() == at::kBFloat16, "A must be bf16"); - TORCH_CHECK(B.scalar_type() == at::kBFloat16, "B must be bf16"); - TORCH_CHECK(global_scale.scalar_type() == at::kFloat, "global_scale must be float"); - TORCH_CHECK(global_scale.dim() == 1 && global_scale.size(0) == 1, "global_scale must be a scalar"); - TORCH_CHECK(B.size(0) == B.size(1), "Rotation matrix must be square"); + check_device_type_cuda("fusedQuantizeNvQuest", + {A, B, OUT, OUT_sf, global_scale}); + check_all_same_gpu("fusedQuantizeNvQuest", {{A, "A", 0}, + {B, "B", 1}, + {OUT, "OUT", 2}, + {OUT_sf, "OUT_sf", 3}, + {global_scale, "global_scale", 4}}); + STD_TORCH_CHECK(A.scalar_type() == ScalarType::BFloat16, "A must be bf16"); + STD_TORCH_CHECK(B.scalar_type() == ScalarType::BFloat16, "B must be bf16"); + STD_TORCH_CHECK(global_scale.scalar_type() == ScalarType::Float, + "global_scale must be float"); + STD_TORCH_CHECK(global_scale.dim() == 1 && global_scale.size(0) == 1, + "global_scale must be a scalar"); + STD_TORCH_CHECK(B.size(0) == B.size(1), "Rotation matrix must be square"); uint32_t HAD_GS = B.size(0); - TORCH_CHECK((A.numel()%HAD_GS)==0, "A must be divisible by", HAD_GS); + STD_TORCH_CHECK((A.numel() % HAD_GS) == 0, "A must be divisible by", HAD_GS); - if(HAD_GS==16){ + if (HAD_GS == 16) { fusedQuantizeNvQuest_host(OUT, OUT_sf, A, B, global_scale); - } else if(HAD_GS==32){ + } else if (HAD_GS == 32) { fusedQuantizeNvQuestHad32_host(OUT, OUT_sf, A, B, global_scale); - } else if(HAD_GS==64){ + } else if (HAD_GS == 64) { fusedQuantizeNvQuestHad64_host(OUT, OUT_sf, A, B, global_scale); - } else if(HAD_GS==128){ + } else if (HAD_GS == 128) { fusedQuantizeNvQuestHad128_host(OUT, OUT_sf, A, B, global_scale); } else { - TORCH_CHECK(false, - "Unsupported rotation size ", HAD_GS, - "; expected 16, 32, 64, or 128."); + STD_TORCH_CHECK(false, + "Unsupported rotation size ", HAD_GS, + "; expected 16, 32, 64, or 128."); } return std::make_tuple(OUT, OUT_sf); } -std::tuple fusedQuantizeNvAbsMax(torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor& OUT, - torch::Tensor& OUT_sf, - torch::Tensor const& global_scale) +std::tuple fusedQuantizeNvAbsMax(Tensor const& A, + Tensor const& B, + Tensor& OUT, + Tensor& OUT_sf, + Tensor const& global_scale) { - torch::checkAllContiguous("fusedQuantizeNvAbsMax", {{A, "A", 0}, - {B, "B", 1}, - {OUT, "OUT", 2}, - {OUT_sf, "OUT_sf", 3}}); - torch::checkDeviceType("fusedQuantizeNvAbsMax", {A, B, OUT, OUT_sf, global_scale}, at::DeviceType::CUDA); - torch::checkAllSameGPU("fusedQuantizeNvAbsMax", {{A, "A", 0}, - {B, "B", 1}, - {OUT, "OUT", 2}, - {OUT_sf, "OUT_sf", 3}, - {global_scale, "global_scale", 4}}); - TORCH_CHECK(A.scalar_type() == at::kBFloat16, "A must be bf16"); - TORCH_CHECK(B.scalar_type() == at::kBFloat16, "B must be bf16"); - TORCH_CHECK(global_scale.scalar_type() == at::kFloat, "global_scale must be float"); - TORCH_CHECK(global_scale.dim() == 1 && global_scale.size(0) == 1, "global_scale must be a scalar"); - TORCH_CHECK(B.size(0) == B.size(1), "Rotation matrix must be square"); + check_all_contiguous("fusedQuantizeNvAbsMax", {{A, "A", 0}, + {B, "B", 1}, + {OUT, "OUT", 2}, + {OUT_sf, "OUT_sf", 3}}); + check_device_type_cuda("fusedQuantizeNvAbsMax", {A, B, OUT, OUT_sf, global_scale}); + check_all_same_gpu("fusedQuantizeNvAbsMax", {{A, "A", 0}, + {B, "B", 1}, + {OUT, "OUT", 2}, + {OUT_sf, "OUT_sf", 3}, + {global_scale, "global_scale", 4}}); + STD_TORCH_CHECK(A.scalar_type() == ScalarType::BFloat16, "A must be bf16"); + STD_TORCH_CHECK(B.scalar_type() == ScalarType::BFloat16, "B must be bf16"); + STD_TORCH_CHECK(global_scale.scalar_type() == ScalarType::Float, + "global_scale must be float"); + STD_TORCH_CHECK(global_scale.dim() == 1 && global_scale.size(0) == 1, + "global_scale must be a scalar"); + STD_TORCH_CHECK(B.size(0) == B.size(1), "Rotation matrix must be square"); uint32_t HAD_GS = B.size(0); - TORCH_CHECK((A.numel()%HAD_GS)==0, "A must be divisible by", HAD_GS); + STD_TORCH_CHECK((A.numel() % HAD_GS) == 0, "A must be divisible by", HAD_GS); - if(HAD_GS==16){ + if (HAD_GS == 16) { fusedQuantizeNvAbsMax_host(OUT, OUT_sf, A, B, global_scale); - } else if(HAD_GS==32){ + } else if (HAD_GS == 32) { fusedQuantizeNvAbsMaxHad32_host(OUT, OUT_sf, A, B, global_scale); - } else if(HAD_GS==64){ + } else if (HAD_GS == 64) { fusedQuantizeNvAbsMaxHad64_host(OUT, OUT_sf, A, B, global_scale); - } else if(HAD_GS==128){ + } else if (HAD_GS == 128) { #if TARGET_CUDA_ARCH == 100 fusedQuantizeNvAbsMax_host_sm100(OUT, OUT_sf, A, B, global_scale); #elif TARGET_CUDA_ARCH == 120 fusedQuantizeNvAbsMaxHad128_host(OUT, OUT_sf, A, B, global_scale); #endif } else { - TORCH_CHECK(false, - "Unsupported rotation size ", HAD_GS, - "; expected 16, 32, 64, or 128."); + STD_TORCH_CHECK(false, + "Unsupported rotation size ", HAD_GS, + "; expected 16, 32, 64, or 128."); } return std::make_tuple(OUT, OUT_sf); } -void backward_t_bf16(const torch::Tensor& x, - const torch::Tensor& h, - torch::Tensor& xh_e2m1, - torch::Tensor& xh_e8m0) +#ifndef QUTLASS_MINIMAL_BUILD +void backward_t_bf16(Tensor const& x, + Tensor const& h, + Tensor& xh_e2m1, + Tensor& xh_e8m0) { - int err = backward_t_bf16_cuda( - x.data_ptr(), - h.data_ptr(), - xh_e2m1.data_ptr(), - xh_e8m0.data_ptr(), + backward_t_bf16_cuda( + x.const_data_ptr(), + h.const_data_ptr(), + xh_e2m1.mutable_data_ptr(), + xh_e8m0.mutable_data_ptr(), x.size(-1), x.size(-2), x.numel() / (x.size(-2) * x.size(-1)), - at::cuda::getCurrentCUDAStream(h.device().index()) - ); + get_current_cuda_stream(h.get_device_index())); } -void backward_qt_bf16(const torch::Tensor& x_e2m1, - const torch::Tensor& x_e8m0, - const torch::Tensor& h, - const torch::Tensor& alpha, - torch::Tensor& xh_e2m1, - torch::Tensor& xh_e8m0) { - int err = backward_qt_bf16_cuda( - x_e2m1.data_ptr(), - x_e8m0.data_ptr(), - h.data_ptr(), - alpha.data_ptr(), - xh_e2m1.data_ptr(), - xh_e8m0.data_ptr(), +void backward_qt_bf16(Tensor const& x_e2m1, + Tensor const& x_e8m0, + Tensor const& h, + Tensor const& alpha, + Tensor& xh_e2m1, + Tensor& xh_e8m0) +{ + backward_qt_bf16_cuda( + x_e2m1.const_data_ptr(), + x_e8m0.const_data_ptr(), + h.const_data_ptr(), + alpha.const_data_ptr(), + xh_e2m1.mutable_data_ptr(), + xh_e8m0.mutable_data_ptr(), x_e2m1.size(-1) * 2, x_e2m1.size(-2), x_e2m1.numel() / (x_e2m1.size(-2) * x_e2m1.size(-1)), - at::cuda::getCurrentCUDAStream(h.device().index()) - ); + get_current_cuda_stream(h.get_device_index())); } -void backward_bf16_square_double_mxfp8(const torch::Tensor& x_bf16, - torch::Tensor& x_fp8, - torch::Tensor& row_scales, - torch::Tensor& column_scales) { - int err = backward_bf16_square_double_mxfp8_cuda( - x_bf16.data_ptr(), +void backward_bf16_square_double_mxfp8(Tensor const& x_bf16, + Tensor& x_fp8, + Tensor& row_scales, + Tensor& column_scales) +{ + backward_bf16_square_double_mxfp8_cuda( + x_bf16.const_data_ptr(), x_bf16.size(0), x_bf16.size(1), - x_fp8.data_ptr(), - row_scales.data_ptr(), - column_scales.data_ptr(), - at::cuda::getCurrentCUDAStream(x_bf16.device().index()) - ); + x_fp8.mutable_data_ptr(), + row_scales.mutable_data_ptr(), + column_scales.mutable_data_ptr(), + get_current_cuda_stream(x_bf16.get_device_index())); } -void mxfp4_transpose_mxfp8(const torch::Tensor& x_fp4, - const torch::Tensor& scales, - torch::Tensor& x_fp8, - torch::Tensor& shared_exps) { - int err = mxfp4_transpose_mxfp8_cuda( - x_fp4.data_ptr(), - scales.data_ptr(), +void mxfp4_transpose_mxfp8(Tensor const& x_fp4, + Tensor const& scales, + Tensor& x_fp8, + Tensor& shared_exps) +{ + mxfp4_transpose_mxfp8_cuda( + x_fp4.const_data_ptr(), + scales.const_data_ptr(), x_fp4.size(0), x_fp4.size(1) * 2, - x_fp8.data_ptr(), - shared_exps.data_ptr(), - at::cuda::getCurrentCUDAStream(x_fp4.device().index()) - ); + x_fp8.mutable_data_ptr(), + shared_exps.mutable_data_ptr(), + get_current_cuda_stream(x_fp4.get_device_index())); } - - -TORCH_LIBRARY(_qutlass_C, m) { - m.def("matmul_mxf4_bf16_tn(Tensor A, Tensor B, Tensor A_sf, Tensor B_sf, Tensor alpha) -> Tensor"); - m.def("matmul_nvf4_bf16_tn(Tensor A, Tensor B, Tensor A_sf, Tensor B_sf, Tensor alpha) -> Tensor"); - m.def("matmul_ada_mxf4_bf16_tn(Tensor A, Tensor B, Tensor A_sf, Tensor B_sf, Tensor alpha) -> Tensor"); - - m.def("fusedQuantizeMxQuest(Tensor A, Tensor R, Tensor OUT, Tensor OUT_sf) -> (Tensor, Tensor)"); - m.def("fusedQuantizeMxQuestWithMask(Tensor A, Tensor R, Tensor OUT, Tensor OUT_sf, Tensor OUT_mask) -> (Tensor, Tensor, Tensor)"); - m.def("fusedQuantizeMxAbsMax(Tensor A, Tensor R, Tensor OUT, Tensor OUT_sf) -> (Tensor, Tensor)"); - m.def("fusedQuantizeNvQuest(Tensor A, Tensor R, Tensor OUT, Tensor OUT_sf, Tensor global_scale) -> (Tensor, Tensor)"); - m.def("fusedQuantizeNvAbsMax(Tensor A, Tensor R, Tensor OUT, Tensor OUT_sf, Tensor global_scale) -> (Tensor, Tensor)"); - - //m.def("backward_t_bf16(Tensor x_e2m1, Tensor x_e8m0, Tensor h, float alpha, Tensor xh_e2m1, Tensor xh_e8m0) -> void"); - //m.def("backward_qt_bf16(Tensor x, Tensor h, Tensor xh_e2m1, Tensor xh_e8m0) -> void"); +#endif // QUTLASS_MINIMAL_BUILD + +} // namespace QUTLASS + +STABLE_TORCH_LIBRARY_FRAGMENT(_qutlass_C, ops) { + ops.def("matmul_mxf4_bf16_tn(Tensor A, Tensor B, Tensor A_sf, Tensor B_sf, Tensor alpha) -> Tensor"); + ops.def("matmul_nvf4_bf16_tn(Tensor A, Tensor B, Tensor A_sf, Tensor B_sf, Tensor alpha) -> Tensor"); + ops.def("matmul_ada_mxf4_bf16_tn(Tensor A, Tensor B, Tensor A_sf, Tensor B_sf, Tensor alpha) -> Tensor"); + ops.def("matmul_mxf8_bf16_tn(Tensor A, Tensor B, Tensor A_sf, Tensor B_sf, Tensor alpha) -> Tensor"); + ops.def("matmul_mxf8_bf16_nn(Tensor A, Tensor B, Tensor A_sf, Tensor B_sf, Tensor alpha) -> Tensor"); + ops.def("fusedQuantizeMxQuest(Tensor A, Tensor R, Tensor OUT, Tensor OUT_sf) -> (Tensor, Tensor)"); + ops.def("fusedQuantizeMxAbsMax(Tensor A, Tensor R, Tensor OUT, Tensor OUT_sf) -> (Tensor, Tensor)"); + ops.def("fusedQuantizeNvQuest(Tensor A, Tensor R, Tensor OUT, Tensor OUT_sf, Tensor global_scale) -> (Tensor, Tensor)"); + ops.def("fusedQuantizeNvAbsMax(Tensor A, Tensor R, Tensor OUT, Tensor OUT_sf, Tensor global_scale) -> (Tensor, Tensor)"); +#ifndef QUTLASS_MINIMAL_BUILD + ops.def("fusedQuantizeMxQuestWithMask(Tensor A, Tensor R, Tensor OUT, Tensor OUT_sf, Tensor OUT_mask) -> (Tensor, Tensor, Tensor)"); + ops.def("backward_t_bf16(Tensor x, Tensor h, Tensor xh_e2m1, Tensor xh_e8m0) -> ()"); + ops.def("backward_qt_bf16(Tensor x_e2m1, Tensor x_e8m0, Tensor h, Tensor alpha, Tensor xh_e2m1, Tensor xh_e8m0) -> ()"); + ops.def("backward_bf16_square_double_mxfp8(Tensor x_bf16, Tensor x_fp8, Tensor row_scales, Tensor column_scales) -> ()"); + ops.def("mxfp4_transpose_mxfp8(Tensor x_fp4, Tensor scales, Tensor x_fp8, Tensor shared_exps) -> ()"); +#endif } -TORCH_LIBRARY_IMPL(_qutlass_C, CUDA, m) { - m.impl("matmul_mxf4_bf16_tn", TORCH_FN(QUTLASS::matmul_mxf4_bf16_tn)); - m.impl("matmul_nvf4_bf16_tn", TORCH_FN(QUTLASS::matmul_nvf4_bf16_tn)); - m.impl("matmul_ada_mxf4_bf16_tn", TORCH_FN(QUTLASS::matmul_ada_mxf4_bf16_tn)); - - m.impl("fusedQuantizeMxQuest", TORCH_FN(QUTLASS::fusedQuantizeMxQuest)); - m.impl("fusedQuantizeMxQuestWithMask", TORCH_FN(QUTLASS::fusedQuantizeMxQuestWithMask)); - m.impl("fusedQuantizeMxAbsMax", TORCH_FN(QUTLASS::fusedQuantizeMxAbsMax)); - m.impl("fusedQuantizeNvQuest", TORCH_FN(QUTLASS::fusedQuantizeNvQuest)); - m.impl("fusedQuantizeNvAbsMax", TORCH_FN(QUTLASS::fusedQuantizeNvAbsMax)); - - //m.impl("backward_t_bf16", TORCH_FN(QUTLASS::backward_t_bf16)); - //m.impl("backward_qt_bf16", TORCH_FN(QUTLASS::backward_qt_bf16)); +STABLE_TORCH_LIBRARY_IMPL(_qutlass_C, CUDA, ops) { + ops.impl("matmul_mxf4_bf16_tn", TORCH_BOX(&QUTLASS::matmul_mxf4_bf16_tn)); + ops.impl("matmul_nvf4_bf16_tn", TORCH_BOX(&QUTLASS::matmul_nvf4_bf16_tn)); + ops.impl("matmul_ada_mxf4_bf16_tn", TORCH_BOX(&QUTLASS::matmul_ada_mxf4_bf16_tn)); + ops.impl("matmul_mxf8_bf16_tn", TORCH_BOX(&QUTLASS::matmul_mxf8_bf16_tn)); + ops.impl("matmul_mxf8_bf16_nn", TORCH_BOX(&QUTLASS::matmul_mxf8_bf16_nn)); + ops.impl("fusedQuantizeMxQuest", TORCH_BOX(&QUTLASS::fusedQuantizeMxQuest)); + ops.impl("fusedQuantizeMxAbsMax", TORCH_BOX(&QUTLASS::fusedQuantizeMxAbsMax)); + ops.impl("fusedQuantizeNvQuest", TORCH_BOX(&QUTLASS::fusedQuantizeNvQuest)); + ops.impl("fusedQuantizeNvAbsMax", TORCH_BOX(&QUTLASS::fusedQuantizeNvAbsMax)); +#ifndef QUTLASS_MINIMAL_BUILD + ops.impl("fusedQuantizeMxQuestWithMask", TORCH_BOX(&QUTLASS::fusedQuantizeMxQuestWithMask)); + ops.impl("backward_t_bf16", TORCH_BOX(&QUTLASS::backward_t_bf16)); + ops.impl("backward_qt_bf16", TORCH_BOX(&QUTLASS::backward_qt_bf16)); + ops.impl("backward_bf16_square_double_mxfp8", + TORCH_BOX(&QUTLASS::backward_bf16_square_double_mxfp8)); + ops.impl("mxfp4_transpose_mxfp8", TORCH_BOX(&QUTLASS::mxfp4_transpose_mxfp8)); +#endif } -//====== pybind ====== - -#define DEFINE_pybind(name) m.def(#name, &name, #name); - -#ifndef QUTLASS_DISABLE_PYBIND -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m -) -{ - m.def("matmul_mxf4_bf16_tn", &matmul_mxf4_bf16_tn, "matmul_mxf4_bf16_tn"); - m.def("matmul_ada_mxf4_bf16_tn", &matmul_ada_mxf4_bf16_tn, "matmul_ada_mxf4_bf16_tn"); - m.def("matmul_nvf4_bf16_tn", &matmul_nvf4_bf16_tn, "matmul_nvf4_bf16_tn"); - m.def("matmul_mxf8_bf16_tn", &matmul_mxf8_bf16_tn, "matmul_mxf8_bf16_tn"); - m.def("matmul_mxf8_bf16_nn", &matmul_mxf8_bf16_nn, "matmul_mxf8_bf16_nn"); - - m.def("fusedQuantizeMxQuest", &QUTLASS::fusedQuantizeMxQuest, "fusedQuantizeMxQuest"); - m.def("fusedQuantizeMxQuestWithMask", &QUTLASS::fusedQuantizeMxQuestWithMask, "fusedQuantizeMxQuestWithMask"); - m.def("fusedQuantizeMxAbsMax", &QUTLASS::fusedQuantizeMxAbsMax, "fusedQuantizeMxAbsMax"); - m.def("fusedQuantizeNvQuest", &QUTLASS::fusedQuantizeNvQuest, "fusedQuantizeNvQuest"); - m.def("fusedQuantizeNvAbsMax", &QUTLASS::fusedQuantizeNvAbsMax, "fusedQuantizeNvAbsMax"); - - m.def("backward_t_bf16", &backward_t_bf16, "backward_t_bf16"); - m.def("backward_qt_bf16", &backward_qt_bf16, "backward_qt_bf16"); - m.def("backward_bf16_square_double_mxfp8", &backward_bf16_square_double_mxfp8, "backward_bf16_square_double_mxfp8"); - m.def("mxfp4_transpose_mxfp8", &mxfp4_transpose_mxfp8, "mxfp4_transpose_mxfp8"); -} +#ifndef QUTLASS_MINIMAL_BUILD +#include "include/registration.h" +REGISTER_EXTENSION(_CUDA) #endif -} \ No newline at end of file diff --git a/qutlass/csrc/fused_quantize_mx.cu b/qutlass/csrc/fused_quantize_mx.cu index bb6b83c..090384c 100644 --- a/qutlass/csrc/fused_quantize_mx.cu +++ b/qutlass/csrc/fused_quantize_mx.cu @@ -14,16 +14,6 @@ * limitations under the License. */ -#include -#include -#include -#include -#include - -#ifndef QUTLASS_DISABLE_PYBIND -#include -#endif - #include #include "cutlass/cutlass.h" @@ -78,12 +68,12 @@ struct GemmRunner { GemmRunner() { } bool run( - torch::Tensor &out, - torch::Tensor &out_sf, - torch::Tensor const&x, - torch::Tensor const&y, + torch::stable::Tensor &out, + torch::stable::Tensor &out_sf, + torch::stable::Tensor const& x, + torch::stable::Tensor const& y, int32_t M, int32_t N, int32_t K, - torch::Device device) + torch::stable::Device device) { using GemmCoord = cutlass::gemm::GemmCoord; @@ -93,16 +83,16 @@ struct GemmRunner { {static_cast(M), static_cast(N), static_cast(K)}, - {(cutlass::bfloat16_t *)x.data_ptr(), K}, - {(cutlass::bfloat16_t *)y.data_ptr(), N}, - {(cutlass::float_e2m1_t *)out.data_ptr(), N}, - {(cutlass::float_e2m1_t *)out.data_ptr(), N}, - {(cutlass::float_ue8m0_t *)out_sf.data_ptr(), M}, + {static_cast(x.const_data_ptr()), K}, + {static_cast(y.const_data_ptr()), N}, + {static_cast(out.mutable_data_ptr()), N}, + {static_cast(out.mutable_data_ptr()), N}, + {static_cast(out_sf.mutable_data_ptr()), M}, cutlass::bfloat16_t(0) //TODO (later): float }; - const at::cuda::OptionalCUDAGuard device_guard(device_of(x)); - cudaStream_t stream = at::cuda::getCurrentCUDAStream(device.index()); + const torch::stable::accelerator::DeviceGuard device_guard(x.get_device_index()); + cudaStream_t stream = get_current_cuda_stream(device.index()); CUTLASS_CHECK(gemmOp.initialize(arguments, nullptr, stream)); @@ -114,10 +104,10 @@ struct GemmRunner { }; -void fusedQuantizeMxQuest_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B) +void fusedQuantizeMxQuest_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B) { int32_t M = A.numel() / 32; int32_t N = B.size(1); @@ -131,10 +121,10 @@ void fusedQuantizeMxQuest_host(torch::Tensor& D, bool result = runGemm.run(D, D_sf, A, B, M, N, K, A.device()); } -void fusedQuantizeMxAbsMax_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B) +void fusedQuantizeMxAbsMax_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B) { int32_t M = A.numel() / 32; int32_t N = B.size(1); @@ -148,10 +138,10 @@ void fusedQuantizeMxAbsMax_host(torch::Tensor& D, bool result = runGemm.run(D, D_sf, A, B, M, N, K, A.device()); } -void fusedQuantizeMxQuestHad64_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B) +void fusedQuantizeMxQuestHad64_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B) { int32_t M = A.numel() / 64; int32_t N = B.size(1); @@ -165,10 +155,10 @@ void fusedQuantizeMxQuestHad64_host(torch::Tensor& D, bool result = runGemm.run(D, D_sf, A, B, M, N, K, A.device()); } -void fusedQuantizeMxAbsMaxHad64_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B) +void fusedQuantizeMxAbsMaxHad64_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B) { int32_t M = A.numel() / 64; int32_t N = B.size(1); @@ -182,10 +172,10 @@ void fusedQuantizeMxAbsMaxHad64_host(torch::Tensor& D, bool result = runGemm.run(D, D_sf, A, B, M, N, K, A.device()); } -void fusedQuantizeMxQuestHad128_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B) +void fusedQuantizeMxQuestHad128_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B) { int32_t M = A.numel() / 128; int32_t N = B.size(1); @@ -199,10 +189,10 @@ void fusedQuantizeMxQuestHad128_host(torch::Tensor& D, bool result = runGemm.run(D, D_sf, A, B, M, N, K, A.device()); } -void fusedQuantizeMxAbsMaxHad128_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B) +void fusedQuantizeMxAbsMaxHad128_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B) { int32_t M = A.numel() / 128; int32_t N = B.size(1); diff --git a/qutlass/csrc/fused_quantize_mx_mask.cu b/qutlass/csrc/fused_quantize_mx_mask.cu index 42a45ea..d237381 100644 --- a/qutlass/csrc/fused_quantize_mx_mask.cu +++ b/qutlass/csrc/fused_quantize_mx_mask.cu @@ -14,16 +14,6 @@ * limitations under the License. */ -#include -#include -#include -#include -#include - -#ifndef QUTLASS_DISABLE_PYBIND -#include -#endif - #include #include "cutlass/cutlass.h" @@ -76,13 +66,13 @@ struct GemmRunner { GemmRunner() { } bool run( - torch::Tensor &out, - torch::Tensor &out_sf, - torch::Tensor &out_mask, - torch::Tensor const&x, - torch::Tensor const&y, + torch::stable::Tensor &out, + torch::stable::Tensor &out_sf, + torch::stable::Tensor &out_mask, + torch::stable::Tensor const& x, + torch::stable::Tensor const& y, int32_t M, int32_t N, int32_t K, - torch::Device device) + torch::stable::Device device) { using GemmCoord = cutlass::gemm::GemmCoord; @@ -92,17 +82,17 @@ struct GemmRunner { {static_cast(M), static_cast(N), static_cast(K)}, - {(cutlass::bfloat16_t *)x.data_ptr(), K}, - {(cutlass::bfloat16_t *)y.data_ptr(), N}, - {(cutlass::float_e2m1_t *)out.data_ptr(), N}, - {(cutlass::float_e2m1_t *)out.data_ptr(), N}, - {(cutlass::float_ue8m0_t *)out_sf.data_ptr(), M}, - {(uint8_t *)out_mask.data_ptr(), M}, //FIXME: bfloat16_t + {static_cast(x.const_data_ptr()), K}, + {static_cast(y.const_data_ptr()), N}, + {static_cast(out.mutable_data_ptr()), N}, + {static_cast(out.mutable_data_ptr()), N}, + {static_cast(out_sf.mutable_data_ptr()), M}, + {static_cast(out_mask.mutable_data_ptr()), M}, //FIXME: bfloat16_t cutlass::bfloat16_t(0) //TODO (later): float }; - const at::cuda::OptionalCUDAGuard device_guard(device_of(x)); - cudaStream_t stream = at::cuda::getCurrentCUDAStream(device.index()); + const torch::stable::accelerator::DeviceGuard device_guard(x.get_device_index()); + cudaStream_t stream = get_current_cuda_stream(device.index()); CUTLASS_CHECK(gemmOp.initialize(arguments, nullptr, stream)); @@ -113,11 +103,11 @@ struct GemmRunner { }; -void fusedQuantizeMxQuestWithMask_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor& D_mask, - torch::Tensor const& A, - torch::Tensor const& B) +void fusedQuantizeMxQuestWithMask_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor& D_mask, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B) { int32_t M = A.numel() / 32; int32_t N = B.size(1); diff --git a/qutlass/csrc/fused_quantize_mx_sm100.cu b/qutlass/csrc/fused_quantize_mx_sm100.cu index 5d946f9..4fb72ff 100644 --- a/qutlass/csrc/fused_quantize_mx_sm100.cu +++ b/qutlass/csrc/fused_quantize_mx_sm100.cu @@ -14,16 +14,6 @@ * limitations under the License. */ -#include -#include -#include -#include -#include - -#ifndef QUTLASS_DISABLE_PYBIND -#include -#endif - #include "cutlass/cutlass.h" #include "cute/tensor.hpp" #include "cutlass/tensor_ref.h" @@ -128,11 +118,11 @@ constexpr bool IsBlockScaleSupported = FusionOp::IsBlockScaleSupported; using SfdOutputCfg = cutlass::detail::Sm1xxBlockScaledOutputConfig; using LayoutSFD = typename SfdOutputCfg::LayoutSF; -typename Gemm::Arguments args_from_options(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale, +typename Gemm::Arguments args_from_options(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale, int32_t M, int32_t N, int32_t K) { using ElementA = typename Gemm::ElementA; @@ -155,30 +145,30 @@ typename Gemm::Arguments args_from_options(torch::Tensor& D, cutlass::gemm::GemmUniversalMode::kGemm, {M, N, K, 1}, { - static_cast(A.data_ptr()), stride_A, - static_cast(B.data_ptr()), stride_B}, + static_cast(A.const_data_ptr()), stride_A, + static_cast(B.const_data_ptr()), stride_B}, { {1.f, 0.f}, nullptr, stride_C, - static_cast(D.data_ptr()), stride_D + static_cast(D.mutable_data_ptr()), stride_D } }; if constexpr (IsBlockScaleSupported) { - arguments.epilogue.thread.block_scale_factor_ptr = static_cast(D_sf.data_ptr()); - arguments.epilogue.thread.norm_constant_ptr = static_cast(global_scale.data_ptr());; + arguments.epilogue.thread.block_scale_factor_ptr = static_cast(D_sf.mutable_data_ptr()); + arguments.epilogue.thread.norm_constant_ptr = static_cast(global_scale.const_data_ptr()); } return arguments; } -void runGemm(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale, +void runGemm(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale, int32_t M, int32_t N, int32_t K, - torch::Device device) + torch::stable::Device device) { Gemm gemm; @@ -189,8 +179,8 @@ void runGemm(torch::Tensor& D, cutlass::device_memory::allocation workspace(workspace_size); - const at::cuda::OptionalCUDAGuard device_guard(device_of(A)); - cudaStream_t stream = at::cuda::getCurrentCUDAStream(device.index()); + const torch::stable::accelerator::DeviceGuard device_guard(A.get_device_index()); + cudaStream_t stream = get_current_cuda_stream(device.index()); CUTLASS_CHECK(gemm.can_implement(arguments)); @@ -199,11 +189,11 @@ void runGemm(torch::Tensor& D, CUTLASS_CHECK(gemm.run(arguments, workspace.get(), stream)); } -void fusedQuantizeMxAbsMax_host_sm100(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale) +void fusedQuantizeMxAbsMax_host_sm100(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale) { #if TARGET_CUDA_ARCH == 100 int32_t M = A.numel() / 128; @@ -212,7 +202,7 @@ void fusedQuantizeMxAbsMax_host_sm100(torch::Tensor& D, runGemm(D, D_sf, A, B, global_scale, M, N, K, A.device()); #else - TORCH_CHECK(false, "Unsupported CUDA arch"); + STD_TORCH_CHECK(false, "Unsupported CUDA arch"); #endif } diff --git a/qutlass/csrc/fused_quantize_nv.cu b/qutlass/csrc/fused_quantize_nv.cu index 3f3a96d..c4a7a71 100644 --- a/qutlass/csrc/fused_quantize_nv.cu +++ b/qutlass/csrc/fused_quantize_nv.cu @@ -14,16 +14,6 @@ * limitations under the License. */ -#include -#include -#include -#include -#include - -#ifndef QUTLASS_DISABLE_PYBIND -#include -#endif - #include #include "cutlass/cutlass.h" @@ -78,13 +68,13 @@ struct GemmRunner { GemmRunner() { } bool run( - torch::Tensor &out, - torch::Tensor &out_sf, - torch::Tensor const&x, - torch::Tensor const&y, + torch::stable::Tensor &out, + torch::stable::Tensor &out_sf, + torch::stable::Tensor const& x, + torch::stable::Tensor const& y, int32_t M, int32_t N, int32_t K, - torch::Device device, - torch::Tensor const& global_scale) + torch::stable::Device device, + torch::stable::Tensor const& global_scale) { using GemmCoord = cutlass::gemm::GemmCoord; @@ -94,17 +84,17 @@ struct GemmRunner { {static_cast(M), static_cast(N), static_cast(K)}, - {(cutlass::bfloat16_t *)x.data_ptr(), K}, - {(cutlass::bfloat16_t *)y.data_ptr(), N}, - {(cutlass::float_e2m1_t *)out.data_ptr(), N}, - {(cutlass::float_e2m1_t *)out.data_ptr(), N}, - {(cutlass::float_ue4m3_t *)out_sf.data_ptr(), M}, - static_cast(global_scale.data_ptr()), + {static_cast(x.const_data_ptr()), K}, + {static_cast(y.const_data_ptr()), N}, + {static_cast(out.mutable_data_ptr()), N}, + {static_cast(out.mutable_data_ptr()), N}, + {static_cast(out_sf.mutable_data_ptr()), M}, + static_cast(const_cast(global_scale.const_data_ptr())), cutlass::bfloat16_t(0) //TODO (later): float }; - const at::cuda::OptionalCUDAGuard device_guard(device_of(x)); - cudaStream_t stream = at::cuda::getCurrentCUDAStream(device.index()); + const torch::stable::accelerator::DeviceGuard device_guard(x.get_device_index()); + cudaStream_t stream = get_current_cuda_stream(device.index()); CUTLASS_CHECK(gemmOp.initialize(arguments, nullptr, stream)); @@ -116,11 +106,11 @@ struct GemmRunner { }; -void fusedQuantizeNvQuest_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale) +void fusedQuantizeNvQuest_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale) { int32_t M = A.numel() / 16; int32_t N = B.size(1); @@ -134,11 +124,11 @@ void fusedQuantizeNvQuest_host(torch::Tensor& D, bool result = runGemm.run(D, D_sf, A, B, M, N, K, A.device(), global_scale); } -void fusedQuantizeNvQuestHad32_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale) +void fusedQuantizeNvQuestHad32_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale) { int32_t M = A.numel() / 32; int32_t N = B.size(1); @@ -152,11 +142,11 @@ void fusedQuantizeNvQuestHad32_host(torch::Tensor& D, bool result = runGemm.run(D, D_sf, A, B, M, N, K, A.device(), global_scale); } -void fusedQuantizeNvQuestHad64_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale) +void fusedQuantizeNvQuestHad64_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale) { int32_t M = A.numel() / 64; int32_t N = B.size(1); @@ -170,11 +160,11 @@ void fusedQuantizeNvQuestHad64_host(torch::Tensor& D, bool result = runGemm.run(D, D_sf, A, B, M, N, K, A.device(), global_scale); } -void fusedQuantizeNvQuestHad128_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale) +void fusedQuantizeNvQuestHad128_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale) { int32_t M = A.numel() / 128; int32_t N = B.size(1); @@ -188,11 +178,11 @@ void fusedQuantizeNvQuestHad128_host(torch::Tensor& D, bool result = runGemm.run(D, D_sf, A, B, M, N, K, A.device(), global_scale); } -void fusedQuantizeNvAbsMax_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale) +void fusedQuantizeNvAbsMax_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale) { int32_t M = A.numel() / 16; int32_t N = B.size(1); @@ -206,11 +196,11 @@ void fusedQuantizeNvAbsMax_host(torch::Tensor& D, bool result = runGemm.run(D, D_sf, A, B, M, N, K, A.device(), global_scale); } -void fusedQuantizeNvAbsMaxHad32_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale) +void fusedQuantizeNvAbsMaxHad32_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale) { int32_t M = A.numel() / 32; int32_t N = B.size(1); @@ -224,11 +214,11 @@ void fusedQuantizeNvAbsMaxHad32_host(torch::Tensor& D, bool result = runGemm.run(D, D_sf, A, B, M, N, K, A.device(), global_scale); } -void fusedQuantizeNvAbsMaxHad64_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale) +void fusedQuantizeNvAbsMaxHad64_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale) { int32_t M = A.numel() / 64; int32_t N = B.size(1); @@ -242,11 +232,11 @@ void fusedQuantizeNvAbsMaxHad64_host(torch::Tensor& D, bool result = runGemm.run(D, D_sf, A, B, M, N, K, A.device(), global_scale); } -void fusedQuantizeNvAbsMaxHad128_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale) +void fusedQuantizeNvAbsMaxHad128_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale) { int32_t M = A.numel() / 128; int32_t N = B.size(1); diff --git a/qutlass/csrc/fused_quantize_nv_sm100.cu b/qutlass/csrc/fused_quantize_nv_sm100.cu index 4dc3b9c..de133a9 100644 --- a/qutlass/csrc/fused_quantize_nv_sm100.cu +++ b/qutlass/csrc/fused_quantize_nv_sm100.cu @@ -14,16 +14,6 @@ * limitations under the License. */ -#include -#include -#include -#include -#include - -#ifndef QUTLASS_DISABLE_PYBIND -#include -#endif - #include "cutlass/cutlass.h" #include "cute/tensor.hpp" @@ -128,11 +118,11 @@ constexpr bool IsBlockScaleSupported = FusionOp::IsBlockScaleSupported; using SfdOutputCfg = cutlass::detail::Sm1xxBlockScaledOutputConfig; using LayoutSFD = typename SfdOutputCfg::LayoutSF; -typename Gemm::Arguments args_from_options_nv(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale, +typename Gemm::Arguments args_from_options_nv(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale, int32_t M, int32_t N, int32_t K) { using ElementA = typename Gemm::ElementA; @@ -155,30 +145,30 @@ typename Gemm::Arguments args_from_options_nv(torch::Tensor& D, cutlass::gemm::GemmUniversalMode::kGemm, {M, N, K, 1}, { - static_cast(A.data_ptr()), stride_A, - static_cast(B.data_ptr()), stride_B}, + static_cast(A.const_data_ptr()), stride_A, + static_cast(B.const_data_ptr()), stride_B}, { {1.f, 0.f}, nullptr, stride_C, - static_cast(D.data_ptr()), stride_D + static_cast(D.mutable_data_ptr()), stride_D } }; if constexpr (IsBlockScaleSupported) { - arguments.epilogue.thread.block_scale_factor_ptr = static_cast(D_sf.data_ptr()); - arguments.epilogue.thread.norm_constant_ptr = static_cast(global_scale.data_ptr());; + arguments.epilogue.thread.block_scale_factor_ptr = static_cast(D_sf.mutable_data_ptr()); + arguments.epilogue.thread.norm_constant_ptr = static_cast(global_scale.const_data_ptr()); } return arguments; } -void runGemmNv(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale, +void runGemmNv(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale, int32_t M, int32_t N, int32_t K, - torch::Device device) + torch::stable::Device device) { Gemm gemm; @@ -189,8 +179,8 @@ void runGemmNv(torch::Tensor& D, cutlass::device_memory::allocation workspace(workspace_size); - const at::cuda::OptionalCUDAGuard device_guard(device_of(A)); - cudaStream_t stream = at::cuda::getCurrentCUDAStream(device.index()); + const torch::stable::accelerator::DeviceGuard device_guard(A.get_device_index()); + cudaStream_t stream = get_current_cuda_stream(device.index()); CUTLASS_CHECK(gemm.can_implement(arguments)); @@ -199,11 +189,11 @@ void runGemmNv(torch::Tensor& D, CUTLASS_CHECK(gemm.run(arguments, workspace.get(), stream)); } -void fusedQuantizeNvAbsMax_host_sm100(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale) +void fusedQuantizeNvAbsMax_host_sm100(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale) { #if TARGET_CUDA_ARCH == 100 int32_t M = A.numel() / 128; @@ -212,7 +202,7 @@ void fusedQuantizeNvAbsMax_host_sm100(torch::Tensor& D, runGemmNv(D, D_sf, A, B, global_scale, M, N, K, A.device()); #else - TORCH_CHECK(false, "Unsupported CUDA arch"); + STD_TORCH_CHECK(false, "Unsupported CUDA arch"); #endif } diff --git a/qutlass/csrc/gemm.cu b/qutlass/csrc/gemm.cu index 57ebcc2..4612262 100644 --- a/qutlass/csrc/gemm.cu +++ b/qutlass/csrc/gemm.cu @@ -14,16 +14,6 @@ * limitations under the License. */ -#include -#include -#include -#include -#include - -#ifndef QUTLASS_DISABLE_PYBIND -#include -#endif - #include "cutlass/cutlass.h" #include "cutlass/gemm/collective/collective_builder.hpp" #include "cutlass/epilogue/collective/collective_builder.hpp" @@ -99,12 +89,12 @@ struct FpGemm { template typename Gemm::Arguments args_from_options( - at::Tensor& D, - at::Tensor const& A, - at::Tensor const& B, - at::Tensor const& A_sf, - at::Tensor const& B_sf, - torch::Tensor const& alpha, + torch::stable::Tensor& D, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& A_sf, + torch::stable::Tensor const& B_sf, + torch::stable::Tensor const& alpha, int M, int N, int K) { using ElementA = typename Gemm::ElementA; @@ -136,31 +126,31 @@ typename Gemm::Arguments args_from_options( cutlass::gemm::GemmUniversalMode::kGemm, {M, N, K, 1}, { - static_cast(A.data_ptr()), stride_A, - static_cast(B.data_ptr()), stride_B, - static_cast(A_sf.data_ptr()), layout_SFA, - static_cast(B_sf.data_ptr()), layout_SFB}, + static_cast(A.const_data_ptr()), stride_A, + static_cast(B.const_data_ptr()), stride_B, + static_cast(A_sf.const_data_ptr()), layout_SFA, + static_cast(B_sf.const_data_ptr()), layout_SFB}, { {}, - static_cast(D.data_ptr()), stride_D, - static_cast(D.data_ptr()), stride_D + static_cast(D.const_data_ptr()), stride_D, + static_cast(D.mutable_data_ptr()), stride_D } }; auto& fusion_args = arguments.epilogue.thread; - fusion_args.alpha_ptr = static_cast(alpha.data_ptr()); + fusion_args.alpha_ptr = static_cast(alpha.const_data_ptr()); return arguments; } template -void runGemm(at::Tensor& D, - at::Tensor const& A, - at::Tensor const& B, - at::Tensor const& A_sf, - at::Tensor const& B_sf, - torch::Tensor const& alpha, +void runGemm(torch::stable::Tensor& D, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& A_sf, + torch::stable::Tensor const& B_sf, + torch::stable::Tensor const& alpha, int M, int N, int K, - torch::Device device) + torch::stable::Device device) { Gemm gemm; @@ -171,8 +161,8 @@ void runGemm(at::Tensor& D, cutlass::device_memory::allocation workspace(workspace_size); - const at::cuda::OptionalCUDAGuard device_guard(device_of(A)); - cudaStream_t stream = at::cuda::getCurrentCUDAStream(device.index()); + const torch::stable::accelerator::DeviceGuard device_guard(A.get_device_index()); + cudaStream_t stream = get_current_cuda_stream(device.index()); CUTLASS_CHECK(gemm.can_implement(arguments)); @@ -181,16 +171,16 @@ void runGemm(at::Tensor& D, CUTLASS_CHECK(gemm.run(arguments, workspace.get(), stream)); } -void matmul_host_mxf4_bf16_tn(torch::Tensor& D, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& A_sf, - torch::Tensor const& B_sf, - torch::Tensor const& alpha) +void matmul_host_mxf4_bf16_tn(torch::stable::Tensor& D, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& A_sf, + torch::stable::Tensor const& B_sf, + torch::stable::Tensor const& alpha) { - auto const m = A.sizes()[0]; - auto const n = B.sizes()[0]; - auto const k = A.sizes()[1] * 2; + auto const m = A.size(0); + auto const n = B.size(0); + auto const k = A.size(1) * 2; using ElementA = cutlass::mx_float4_t; using LayoutATag = cutlass::layout::RowMajor; @@ -253,20 +243,20 @@ void matmul_host_mxf4_bf16_tn(torch::Tensor& D, >(D, A, B, A_sf, B_sf, alpha, m, n, k, A.device()); } #else - TORCH_CHECK(false, "Unsupported CUDA arch"); + STD_TORCH_CHECK(false, "Unsupported CUDA arch"); #endif } -void matmul_host_nvf4_bf16_tn(torch::Tensor& D, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& A_sf, - torch::Tensor const& B_sf, - torch::Tensor const& alpha) +void matmul_host_nvf4_bf16_tn(torch::stable::Tensor& D, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& A_sf, + torch::stable::Tensor const& B_sf, + torch::stable::Tensor const& alpha) { - auto const m = A.sizes()[0]; - auto const n = B.sizes()[0]; - auto const k = A.sizes()[1] * 2; + auto const m = A.size(0); + auto const n = B.size(0); + auto const k = A.size(1) * 2; using ElementA = cutlass::nv_float4_t; using LayoutATag = cutlass::layout::RowMajor; @@ -330,21 +320,21 @@ void matmul_host_nvf4_bf16_tn(torch::Tensor& D, >(D, A, B, A_sf, B_sf, alpha, m, n, k, A.device()); } #else - TORCH_CHECK(false, "Unsupported CUDA arch"); + STD_TORCH_CHECK(false, "Unsupported CUDA arch"); #endif } -void matmul_host_mxf8_bf16_tn(torch::Tensor& D, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& A_sf, - torch::Tensor const& B_sf, - torch::Tensor const& alpha) +void matmul_host_mxf8_bf16_tn(torch::stable::Tensor& D, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& A_sf, + torch::stable::Tensor const& B_sf, + torch::stable::Tensor const& alpha) { - auto const m = A.sizes()[0]; - auto const n = B.sizes()[0]; - auto const k = A.sizes()[1]; + auto const m = A.size(0); + auto const n = B.size(0); + auto const k = A.size(1); using ElementA = cutlass::mx_float8_t; using LayoutATag = cutlass::layout::RowMajor; @@ -391,20 +381,20 @@ void matmul_host_mxf8_bf16_tn(torch::Tensor& D, ElementB, LayoutBTag, AlignmentB>::Gemm, cutlass::float_ue8m0_t >(D, A, B, A_sf, B_sf, alpha, m, n, k, A.device()); #else - TORCH_CHECK(false, "Unsupported CUDA arch"); + STD_TORCH_CHECK(false, "Unsupported CUDA arch"); #endif } -void matmul_host_mxf8_bf16_nn(torch::Tensor& D, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& A_sf, - torch::Tensor const& B_sf, - torch::Tensor const& alpha) +void matmul_host_mxf8_bf16_nn(torch::stable::Tensor& D, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& A_sf, + torch::stable::Tensor const& B_sf, + torch::stable::Tensor const& alpha) { - auto const m = A.sizes()[1]; - auto const n = B.sizes()[0]; - auto const k = A.sizes()[0]; + auto const m = A.size(1); + auto const n = B.size(0); + auto const k = A.size(0); using ElementA = cutlass::mx_float8_t; using LayoutATag = cutlass::layout::ColumnMajor; @@ -439,6 +429,6 @@ void matmul_host_mxf8_bf16_nn(torch::Tensor& D, >(D, A, B, A_sf, B_sf, alpha, m, n, k, A.device()); } #else - TORCH_CHECK(false, "Unsupported CUDA arch"); + STD_TORCH_CHECK(false, "Unsupported CUDA arch"); #endif } \ No newline at end of file diff --git a/qutlass/csrc/gemm_ada.cu b/qutlass/csrc/gemm_ada.cu index 3801115..0b2e954 100644 --- a/qutlass/csrc/gemm_ada.cu +++ b/qutlass/csrc/gemm_ada.cu @@ -14,16 +14,6 @@ * limitations under the License. */ -#include -#include -#include -#include -#include - -#ifndef QUTLASS_DISABLE_PYBIND -#include -#endif - #include #include @@ -37,13 +27,13 @@ #include template -void qutlass_matmul_mxf4_v1(torch::Tensor const&input, - torch::Tensor const&weight, - torch::Tensor const&input_sf, - torch::Tensor const&weight_sf, - torch::Tensor &out, - torch::Tensor const& alpha, - torch::Device device) +void qutlass_matmul_mxf4_v1(torch::stable::Tensor const& input, + torch::stable::Tensor const& weight, + torch::stable::Tensor const& input_sf, + torch::stable::Tensor const& weight_sf, + torch::stable::Tensor &out, + torch::stable::Tensor const& alpha, + torch::stable::Device device) { auto M = input.size(0); auto N = weight.size(0); @@ -87,29 +77,30 @@ void qutlass_matmul_mxf4_v1(torch::Tensor const&input, cutlass::gemm::GemmCoord problem_size(M, N, K); cutlass::TensorRef input_ref( - static_cast(input.data_ptr()), + static_cast(const_cast(input.const_data_ptr())), LayoutInputA::packed(input_size)); cutlass::TensorRef weight_ref( - static_cast(weight.data_ptr()), + static_cast(const_cast(weight.const_data_ptr())), LayoutInputB::packed(weight_size)); cutlass::TensorRef out_ref( - static_cast(out.data_ptr()), + static_cast(out.mutable_data_ptr()), LayoutOutput::packed(output_size)); typename Gemm::Arguments arguments{ problem_size, input_ref, - reinterpret_cast(input_sf.data_ptr()), + reinterpret_cast(input_sf.const_data_ptr()), weight_ref, - reinterpret_cast(weight_sf.data_ptr()), + reinterpret_cast(weight_sf.const_data_ptr()), out_ref, out_ref, {}, 1 }; - arguments.epilogue.alpha_ptr = static_cast(alpha.data_ptr()); + arguments.epilogue.alpha_ptr = + static_cast(alpha.const_data_ptr()); Gemm gemm_op; @@ -118,20 +109,20 @@ void qutlass_matmul_mxf4_v1(torch::Tensor const&input, CUTLASS_CHECK(gemm_op.can_implement(arguments)); - const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); - cudaStream_t stream = at::cuda::getCurrentCUDAStream(device.index()); + const torch::stable::accelerator::DeviceGuard device_guard(input.get_device_index()); + cudaStream_t stream = get_current_cuda_stream(device.index()); CUTLASS_CHECK(gemm_op.initialize(arguments, workspace.get(), stream)); CUTLASS_CHECK(gemm_op(stream)); } -void matmul_host_ada_mxf4_bf16_tn(torch::Tensor const&input, - torch::Tensor const&weight, - torch::Tensor const&input_sf, - torch::Tensor const&weight_sf, - torch::Tensor &out, - torch::Tensor const& alpha) +void matmul_host_ada_mxf4_bf16_tn(torch::stable::Tensor const& input, + torch::stable::Tensor const& weight, + torch::stable::Tensor const& input_sf, + torch::stable::Tensor const& weight_sf, + torch::stable::Tensor &out, + torch::stable::Tensor const& alpha) { using TileShape = typename cutlass::gemm::GemmShape<16, 16, 256>; using WarpShape = typename cutlass::gemm::GemmShape<16, 16, 256>; @@ -140,6 +131,6 @@ void matmul_host_ada_mxf4_bf16_tn(torch::Tensor const&input, #if TARGET_CUDA_ARCH == 120 qutlass_matmul_mxf4_v1(input, weight, input_sf, weight_sf, out, alpha, input.device()); #else - TORCH_CHECK(false, "matmul_ada_mxf4_bf16_tn was optimized for sm120. For other architectures, please use matmul_mxf4_bf16_tn instead"); + STD_TORCH_CHECK(false, "matmul_ada_mxf4_bf16_tn was optimized for sm120. For other architectures, please use matmul_mxf4_bf16_tn instead"); #endif } \ No newline at end of file diff --git a/qutlass/csrc/include/bindings_utils.h b/qutlass/csrc/include/bindings_utils.h new file mode 100644 index 0000000..e16bfd4 --- /dev/null +++ b/qutlass/csrc/include/bindings_utils.h @@ -0,0 +1,138 @@ +#pragma once + +#include +#include +#include + +#include +#include +#include + +namespace { + +// methods in here are meant to mirror some of the methods +// in PyTorch's aten/src/ATen/TensorUtils.cpp file + +using torch::stable::Tensor; +using torch::headeronly::DeviceType; + +struct TensorArg { + const Tensor& tensor; + const char* name; + int pos; +}; + +inline std::ostream& operator<<(std::ostream& out, const TensorArg& t) { + if (t.pos == 0) { + // 0 is distinguished; it usually indicates 'self' or the return + // tensor + out << '\'' << t.name << '\''; + } else { + out << "argument #" << t.pos << " '" << t.name << '\''; + } + return out; +} + +inline const char* device_type_name(DeviceType type) { + switch (type) { + case DeviceType::CPU: + return "cpu"; + case DeviceType::CUDA: + return "cuda"; + case DeviceType::HIP: + return "hip"; + case DeviceType::XLA: + return "xla"; + case DeviceType::MPS: + return "mps"; + case DeviceType::XPU: + return "xpu"; + case DeviceType::Meta: + return "meta"; + case DeviceType::PrivateUse1: + return "privateuseone"; + default: + return "unknown"; + } +} + +inline void append_device_string(std::ostream& out, const Tensor& t) { + const auto device = t.device(); + out << device_type_name(device.type()); + if (device.has_index()) { + out << ':' << device.index(); + } +} + +inline void check_all_contiguous(const char* op, + std::initializer_list args) { + for (const auto& arg : args) { + if (arg.tensor.is_contiguous()) { + continue; + } + std::ostringstream oss; + oss << "Expected contiguous tensor, but got non-contiguous tensor for " << arg + << " (while checking arguments for " << op << ")"; + STD_TORCH_CHECK(false, oss.str()); + } +} + +inline void check_device_type_cuda(const char* op, + std::initializer_list tensors) { + for (const auto& tensor : tensors) { + if (tensor.device().is_cuda()) { + continue; + } + std::ostringstream oss; + oss << "Expected tensor to have cuda DeviceType, but got tensor with " + << device_type_name(tensor.device().type()) + << " DeviceType (while checking arguments for " << op << ")"; + STD_TORCH_CHECK(false, oss.str()); + } +} + +inline void check_same_gpu(const char* op, + const TensorArg& t1, + const TensorArg& t2) { + const bool t1_cpu = t1.tensor.device().is_cpu(); + const bool t2_cpu = t2.tensor.device().is_cpu(); + if (t1_cpu || t2_cpu) { + std::ostringstream oss; + if (t1_cpu) { + oss << "Tensor for " << t1 << " is on CPU, "; + } + if (t2_cpu) { + oss << "Tensor for " << t2 << " is on CPU, "; + } + oss << "but expected " << ((t1_cpu && t2_cpu) ? "them" : "it") + << " to be on GPU (while checking arguments for " << op << ')'; + STD_TORCH_CHECK(false, oss.str()); + } + const auto d1 = t1.tensor.device(); + const auto d2 = t2.tensor.device(); + if (d1 == d2) { + return; + } + std::ostringstream oss; + oss << "Expected tensor for " << t1 << " to have the same device as tensor for " + << t2 << "; but device "; + append_device_string(oss, t1.tensor); + oss << " does not equal "; + append_device_string(oss, t2.tensor); + oss << " (while checking arguments for " << op << ")"; + STD_TORCH_CHECK(false, oss.str()); +} + +inline void check_all_same_gpu(const char* op, + std::initializer_list args) { + const TensorArg* first = nullptr; + for (const auto& arg : args) { + if (first == nullptr) { + first = &arg; + } else { + check_same_gpu(op, *first, arg); + } + } +} + +} // namespace diff --git a/qutlass/csrc/include/common.h b/qutlass/csrc/include/common.h index 3caa1cd..8b96c81 100644 --- a/qutlass/csrc/include/common.h +++ b/qutlass/csrc/include/common.h @@ -1,6 +1,5 @@ #pragma once -#pragma once #include #include #include @@ -8,31 +7,43 @@ #include #include - -#include #include +#include +#include +#include +#include +#include +#include + #include "cutlass/cutlass.h" /** * Helper function for checking CUTLASS errors */ -#define CUTLASS_CHECK(status) \ - { \ - cutlass::Status error = status; \ - TORCH_CHECK(error == cutlass::Status::kSuccess, \ - cutlassGetStatusString(error)); \ +#define CUTLASS_CHECK(status) \ + { \ + cutlass::Status error = status; \ + STD_TORCH_CHECK(error == cutlass::Status::kSuccess, \ + cutlassGetStatusString(error)); \ } /** * Panic wrapper for unwinding CUDA runtime errors */ -#define CUDA_CHECK(status) \ - { \ - cudaError_t error = status; \ - TORCH_CHECK(error == cudaSuccess, cudaGetErrorString(error)); \ +#define CUDA_CHECK(status) \ + { \ + cudaError_t error = status; \ + STD_TORCH_CHECK(error == cudaSuccess, cudaGetErrorString(error)); \ } +inline cudaStream_t get_current_cuda_stream(int32_t device_index) { + void* stream_ptr = nullptr; + TORCH_ERROR_CODE_CHECK( + aoti_torch_get_current_cuda_stream(device_index, &stream_ptr)); + return reinterpret_cast(stream_ptr); +} + inline int get_cuda_max_shared_memory_per_block_opt_in(int const device) { int max_shared_mem_per_block_opt_in = 0; cudaDeviceGetAttribute(&max_shared_mem_per_block_opt_in, diff --git a/qutlass/csrc/include/cutlass_extensions/epilogue/thread/linear_combination_quant.h b/qutlass/csrc/include/cutlass_extensions/epilogue/thread/linear_combination_quant.h index e85a99e..6f23348 100644 --- a/qutlass/csrc/include/cutlass_extensions/epilogue/thread/linear_combination_quant.h +++ b/qutlass/csrc/include/cutlass_extensions/epilogue/thread/linear_combination_quant.h @@ -39,10 +39,6 @@ */ #pragma once -#ifndef QUTLASS_DISABLE_PYBIND -#include -#endif - #include "cutlass/array.h" #include "cutlass/cutlass.h" #include "cutlass/epilogue/thread/linear_combination_params.h" diff --git a/qutlass/csrc/include/fused_quantize_host.h b/qutlass/csrc/include/fused_quantize_host.h index 708afdc..5f72670 100644 --- a/qutlass/csrc/include/fused_quantize_host.h +++ b/qutlass/csrc/include/fused_quantize_host.h @@ -19,100 +19,100 @@ namespace QUTLASS { -void fusedQuantizeMxQuest_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B); - -void fusedQuantizeMxQuestWithMask_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor& D_mask, - torch::Tensor const& A, - torch::Tensor const& B); - -void fusedQuantizeMxAbsMax_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B); - -void fusedQuantizeMxQuestHad64_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B); - -void fusedQuantizeMxAbsMaxHad64_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B); - -void fusedQuantizeMxQuestHad128_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B); - -void fusedQuantizeMxAbsMaxHad128_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B); - -void fusedQuantizeNvQuest_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale); - -void fusedQuantizeNvQuestHad32_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale); - -void fusedQuantizeNvQuestHad64_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale); - -void fusedQuantizeNvQuestHad128_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale); - -void fusedQuantizeNvAbsMax_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale); - -void fusedQuantizeNvAbsMaxHad32_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale); - -void fusedQuantizeNvAbsMaxHad64_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale); - -void fusedQuantizeNvAbsMaxHad128_host(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale); - -void fusedQuantizeMxAbsMax_host_sm100(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale); - -void fusedQuantizeNvAbsMax_host_sm100(torch::Tensor& D, - torch::Tensor& D_sf, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& global_scale); +void fusedQuantizeMxQuest_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B); + +void fusedQuantizeMxQuestWithMask_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor& D_mask, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B); + +void fusedQuantizeMxAbsMax_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B); + +void fusedQuantizeMxQuestHad64_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B); + +void fusedQuantizeMxAbsMaxHad64_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B); + +void fusedQuantizeMxQuestHad128_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B); + +void fusedQuantizeMxAbsMaxHad128_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B); + +void fusedQuantizeNvQuest_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale); + +void fusedQuantizeNvQuestHad32_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale); + +void fusedQuantizeNvQuestHad64_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale); + +void fusedQuantizeNvQuestHad128_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale); + +void fusedQuantizeNvAbsMax_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale); + +void fusedQuantizeNvAbsMaxHad32_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale); + +void fusedQuantizeNvAbsMaxHad64_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale); + +void fusedQuantizeNvAbsMaxHad128_host(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale); + +void fusedQuantizeMxAbsMax_host_sm100(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale); + +void fusedQuantizeNvAbsMax_host_sm100(torch::stable::Tensor& D, + torch::stable::Tensor& D_sf, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& global_scale); } // namespace QUTLASS diff --git a/qutlass/csrc/include/gemm.h b/qutlass/csrc/include/gemm.h index 05cf323..245c0c4 100644 --- a/qutlass/csrc/include/gemm.h +++ b/qutlass/csrc/include/gemm.h @@ -18,37 +18,37 @@ #pragma once #include -void matmul_host_mxf4_bf16_tn(torch::Tensor& D, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& A_sf, - torch::Tensor const& B_sf, - torch::Tensor const& alpha); +void matmul_host_mxf4_bf16_tn(torch::stable::Tensor& D, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& A_sf, + torch::stable::Tensor const& B_sf, + torch::stable::Tensor const& alpha); -void matmul_host_ada_mxf4_bf16_tn(torch::Tensor const& input, - torch::Tensor const& weight, - torch::Tensor const& input_sf, - torch::Tensor const& weight_sf, - torch::Tensor &out, - torch::Tensor const& alpha); +void matmul_host_ada_mxf4_bf16_tn(torch::stable::Tensor const& input, + torch::stable::Tensor const& weight, + torch::stable::Tensor const& input_sf, + torch::stable::Tensor const& weight_sf, + torch::stable::Tensor &out, + torch::stable::Tensor const& alpha); -void matmul_host_nvf4_bf16_tn(torch::Tensor& D, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& A_sf, - torch::Tensor const& B_sf, - torch::Tensor const& alpha); +void matmul_host_nvf4_bf16_tn(torch::stable::Tensor& D, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& A_sf, + torch::stable::Tensor const& B_sf, + torch::stable::Tensor const& alpha); -void matmul_host_mxf8_bf16_tn(torch::Tensor& D, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& A_sf, - torch::Tensor const& B_sf, - torch::Tensor const& alpha); +void matmul_host_mxf8_bf16_tn(torch::stable::Tensor& D, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& A_sf, + torch::stable::Tensor const& B_sf, + torch::stable::Tensor const& alpha); -void matmul_host_mxf8_bf16_nn(torch::Tensor& D, - torch::Tensor const& A, - torch::Tensor const& B, - torch::Tensor const& A_sf, - torch::Tensor const& B_sf, - torch::Tensor const& alpha); \ No newline at end of file +void matmul_host_mxf8_bf16_nn(torch::stable::Tensor& D, + torch::stable::Tensor const& A, + torch::stable::Tensor const& B, + torch::stable::Tensor const& A_sf, + torch::stable::Tensor const& B_sf, + torch::stable::Tensor const& alpha); \ No newline at end of file diff --git a/qutlass/csrc/include/registration.h b/qutlass/csrc/include/registration.h new file mode 100644 index 0000000..fc83b13 --- /dev/null +++ b/qutlass/csrc/include/registration.h @@ -0,0 +1,18 @@ +#pragma once + +#include + +#define _CONCAT(A, B) A##B +#define CONCAT(A, B) _CONCAT(A, B) + +#define _STRINGIFY(A) #A +#define STRINGIFY(A) _STRINGIFY(A) + +// REGISTER_EXTENSION allows the shared library to be loaded and initialized +// via python's import statement. +#define REGISTER_EXTENSION(NAME) \ + PyMODINIT_FUNC CONCAT(PyInit_, NAME)() { \ + static struct PyModuleDef module = {PyModuleDef_HEAD_INIT, \ + STRINGIFY(NAME), nullptr, 0, nullptr}; \ + return PyModule_Create(&module); \ + } diff --git a/qutlass/csrc/quartet_bwd_sm120.cu b/qutlass/csrc/quartet_bwd_sm120.cu index 6127ba6..8d51c0b 100644 --- a/qutlass/csrc/quartet_bwd_sm120.cu +++ b/qutlass/csrc/quartet_bwd_sm120.cu @@ -22,10 +22,6 @@ #include #include -#ifndef QUTLASS_DISABLE_PYBIND -#include -#endif - #include template @@ -457,7 +453,7 @@ int backward_t_bf16_cuda(const void* x_ptr, size_b ); #else - TORCH_CHECK(false, "Unsupported CUDA arch"); + STD_TORCH_CHECK(false, "Unsupported CUDA arch"); #endif return 0; @@ -491,7 +487,7 @@ int backward_qt_bf16_cuda(const void* x_e2m1_ptr, size_b ); #else - TORCH_CHECK(false, "Unsupported CUDA arch"); + STD_TORCH_CHECK(false, "Unsupported CUDA arch"); #endif return 0; } diff --git a/setup.py b/setup.py index 3fd3495..25af548 100644 --- a/setup.py +++ b/setup.py @@ -135,13 +135,22 @@ def third_party_cmake(): ], define_macros=[("TARGET_CUDA_ARCH", str(cc))], extra_compile_args={ - "cxx": ["-std=c++17"], - "nvcc": get_cuda_arch_flags(), + "cxx": [ + "-std=c++17", + "-DUSE_CUDA", + "-DTORCH_TARGET_VERSION=0x020B000000000000ULL", + ], + "nvcc": [ + "-DUSE_CUDA", + "-DTORCH_TARGET_VERSION=0x020B000000000000ULL", + *get_cuda_arch_flags(), + ], }, extra_link_args=[ "-lcudart", "-lcuda", ], + py_limited_api=True, ) ], cmdclass={"build_ext": BuildExtension},