From cc540c2dc85663f52c73f3e2f4f69e01f2db80a7 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Tue, 23 Jun 2026 20:14:37 +0000 Subject: [PATCH 01/13] migrated kernels used in vllm to use the torch stable abi Signed-off-by: Chris Leonard --- qutlass/csrc/bindings.cpp | 755 +++++++++--------- qutlass/csrc/fused_quantize_mx.cu | 82 +- qutlass/csrc/fused_quantize_mx_mask.cu | 48 +- qutlass/csrc/fused_quantize_mx_sm100.cu | 58 +- qutlass/csrc/fused_quantize_nv.cu | 118 ++- qutlass/csrc/fused_quantize_nv_sm100.cu | 58 +- qutlass/csrc/gemm.cu | 134 ++-- qutlass/csrc/gemm_ada.cu | 55 +- qutlass/csrc/include/common.h | 35 +- .../thread/linear_combination_quant.h | 4 - qutlass/csrc/include/fused_quantize_host.h | 190 ++--- qutlass/csrc/include/gemm.h | 61 +- 12 files changed, 747 insertions(+), 851 deletions(-) diff --git a/qutlass/csrc/bindings.cpp b/qutlass/csrc/bindings.cpp index 974109d..cad881c 100644 --- a/qutlass/csrc/bindings.cpp +++ b/qutlass/csrc/bindings.cpp @@ -14,389 +14,450 @@ * limitations under the License. */ -#include -#include -#include -#include -#include - -#ifndef QUTLASS_DISABLE_PYBIND -#include -#endif +#include +#include -#include -#include +#include +#include #include #include "include/gemm.h" #include "include/fused_quantize_host.h" -#include "include/backward_host.h" -namespace QUTLASS { +namespace { + +using torch::stable::Tensor; +using torch::headeronly::DeviceType; +using torch::headeronly::ScalarType; + +struct TensorArg { + const Tensor& tensor; + const char* name; + int pos; +}; + +inline std::ostream& operator<<(std::ostream& out, const TensorArg& t) { + if (t.pos == 0) { + out << '\'' << t.name << '\''; + } else { + out << "argument #" << t.pos << " '" << t.name << '\''; + } + return out; +} -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) -{ - 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"); +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"; + } +} - 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())); +inline void append_device_type_string(std::ostream& out, DeviceType type) { + out << device_type_name(type); +} - matmul_host_mxf4_bf16_tn(OUT, A, B, A_sf, B_sf, alpha); +inline void append_device_string(std::ostream& out, const Tensor& t) { + const auto device = t.device(); + append_device_type_string(out, device.type()); + if (device.has_index()) { + out << ':' << device.index(); + } +} - return OUT; +inline void check_contiguous(const char* op, const TensorArg& t) { + if (t.tensor.is_contiguous()) { + return; + } + std::ostringstream oss; + oss << "Expected contiguous tensor, but got non-contiguous tensor for " << t + << " (while checking arguments for " << op << ")"; + STD_TORCH_CHECK(false, oss.str()); } -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) -{ - 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"); +inline void check_all_contiguous(const char* op, + std::initializer_list args) { + for (const auto& arg : args) { + check_contiguous(op, arg); + } +} - 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())); +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()); + } +} - matmul_host_nvf4_bf16_tn(OUT, A, B, A_sf, B_sf, alpha); +inline void check_same_gpu(const char* op, + const TensorArg& t1, + const TensorArg& t2) { + const bool t1_cuda = t1.tensor.device().is_cuda(); + const bool t2_cuda = t2.tensor.device().is_cuda(); + if (!t1_cuda || !t2_cuda) { + std::ostringstream oss; + if (!t1_cuda) { + oss << "Tensor for " << t1 << " is on CPU, "; + } + if (!t2_cuda) { + oss << "Tensor for " << t2 << " is on CPU, "; + } + oss << "but expected " << ((t1_cuda && t2_cuda) ? "them" : "it") + << " to be on GPU (while checking arguments for " << op << ')'; + STD_TORCH_CHECK(false, oss.str()); + } + const int d1 = t1.tensor.get_device_index(); + const int d2 = t2.tensor.get_device_index(); + 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()); +} - return OUT; +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); + } + } } -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) +} // namespace + +namespace QUTLASS { + +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_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_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 C = torch::empty({M, N}, torch::dtype(torch::kBFloat16).device(A.device())); + auto OUT = torch::stable::new_empty(A, {M, N}, ScalarType::BFloat16); - matmul_host_ada_mxf4_bf16_tn(A, B, A_sf, B_sf, C, alpha); + matmul_host_mxf4_bf16_tn(OUT, A, B, A_sf, B_sf, alpha); - return C; + return OUT; } -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_nvf4_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_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_mxf8_bf16_tn(OUT, A, B, A_sf, B_sf, alpha); + matmul_host_nvf4_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_ada_mxf4_bf16_tn(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}, + 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"); - 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"); - - uint32_t M = A.size(1); + 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 C = torch::stable::new_empty(A, {M, N}, ScalarType::BFloat16); - matmul_host_mxf8_bf16_nn(OUT, A, B, A_sf, B_sf, alpha); + matmul_host_ada_mxf4_bf16_tn(A, B, A_sf, B_sf, C, alpha); - return OUT; + return C; } -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) +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); } -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){ #if TARGET_CUDA_ARCH == 100 || TARGET_CUDA_ARCH == 101 || TARGET_CUDA_ARCH == 110 - 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 + auto global_scale = + torch::stable::new_zeros(A, {1}, ScalarType::Float); //FIXME: add input global_scale to interface for consistency 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){ #if TARGET_CUDA_ARCH == 100 || TARGET_CUDA_ARCH == 101 || TARGET_CUDA_ARCH == 110 @@ -405,136 +466,34 @@ std::tuple fusedQuantizeNvAbsMax(torch::Tensor con 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) -{ - int err = backward_t_bf16_cuda( - x.data_ptr(), - h.data_ptr(), - xh_e2m1.data_ptr(), - xh_e8m0.data_ptr(), - x.size(-1), - x.size(-2), - x.numel() / (x.size(-2) * x.size(-1)), - at::cuda::getCurrentCUDAStream(h.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(), - 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()) - ); -} - -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(), - 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()) - ); +} // 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("fusedQuantizeMxQuest(Tensor A, Tensor R, Tensor OUT, Tensor OUT_sf) -> (Tensor, Tensor)"); + ops.def("fusedQuantizeMxQuestWithMask(Tensor A, Tensor R, Tensor OUT, Tensor OUT_sf, Tensor OUT_mask) -> (Tensor, 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)"); } -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(), - x_fp4.size(0), - x_fp4.size(1) * 2, - x_fp8.data_ptr(), - shared_exps.data_ptr(), - at::cuda::getCurrentCUDAStream(x_fp4.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"); +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("fusedQuantizeMxQuest", TORCH_BOX(&QUTLASS::fusedQuantizeMxQuest)); + ops.impl("fusedQuantizeMxQuestWithMask", TORCH_BOX(&QUTLASS::fusedQuantizeMxQuestWithMask)); + ops.impl("fusedQuantizeMxAbsMax", TORCH_BOX(&QUTLASS::fusedQuantizeMxAbsMax)); + ops.impl("fusedQuantizeNvQuest", TORCH_BOX(&QUTLASS::fusedQuantizeNvQuest)); + ops.impl("fusedQuantizeNvAbsMax", TORCH_BOX(&QUTLASS::fusedQuantizeNvAbsMax)); } - -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)); -} - -//====== 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"); -} -#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 c0d1016..0e6d20d 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 || TARGET_CUDA_ARCH == 101 || TARGET_CUDA_ARCH == 110 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 b696066..08aefb0 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 || TARGET_CUDA_ARCH == 101 || TARGET_CUDA_ARCH == 110 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 a19383a..49cf407 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/common.h b/qutlass/csrc/include/common.h index 3caa1cd..923a5c5 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..b08284d 100644 --- a/qutlass/csrc/include/gemm.h +++ b/qutlass/csrc/include/gemm.h @@ -14,41 +14,40 @@ * limitations under the License. */ - #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 From 1e4f8390d802ab3478c6b56104d1f0dfc95b54fe Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Tue, 23 Jun 2026 21:48:01 +0000 Subject: [PATCH 02/13] moved new binding.cpp helper functions to bindings_utils.h to make git diff cleaner and easier to read Signed-off-by: Chris Leonard --- qutlass/csrc/bindings.cpp | 135 +----------------------- qutlass/csrc/include/bindings_utils.h | 141 ++++++++++++++++++++++++++ 2 files changed, 143 insertions(+), 133 deletions(-) create mode 100644 qutlass/csrc/include/bindings_utils.h diff --git a/qutlass/csrc/bindings.cpp b/qutlass/csrc/bindings.cpp index cad881c..2f3b74d 100644 --- a/qutlass/csrc/bindings.cpp +++ b/qutlass/csrc/bindings.cpp @@ -17,148 +17,17 @@ #include #include -#include -#include #include +#include "include/bindings_utils.h" #include "include/gemm.h" #include "include/fused_quantize_host.h" -namespace { +namespace QUTLASS { using torch::stable::Tensor; -using torch::headeronly::DeviceType; using torch::headeronly::ScalarType; -struct TensorArg { - const Tensor& tensor; - const char* name; - int pos; -}; - -inline std::ostream& operator<<(std::ostream& out, const TensorArg& t) { - if (t.pos == 0) { - 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_type_string(std::ostream& out, DeviceType type) { - out << device_type_name(type); -} - -inline void append_device_string(std::ostream& out, const Tensor& t) { - const auto device = t.device(); - append_device_type_string(out, device.type()); - if (device.has_index()) { - out << ':' << device.index(); - } -} - -inline void check_contiguous(const char* op, const TensorArg& t) { - if (t.tensor.is_contiguous()) { - return; - } - std::ostringstream oss; - oss << "Expected contiguous tensor, but got non-contiguous tensor for " << t - << " (while checking arguments for " << op << ")"; - STD_TORCH_CHECK(false, oss.str()); -} - -inline void check_all_contiguous(const char* op, - std::initializer_list args) { - for (const auto& arg : args) { - check_contiguous(op, arg); - } -} - -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_cuda = t1.tensor.device().is_cuda(); - const bool t2_cuda = t2.tensor.device().is_cuda(); - if (!t1_cuda || !t2_cuda) { - std::ostringstream oss; - if (!t1_cuda) { - oss << "Tensor for " << t1 << " is on CPU, "; - } - if (!t2_cuda) { - oss << "Tensor for " << t2 << " is on CPU, "; - } - oss << "but expected " << ((t1_cuda && t2_cuda) ? "them" : "it") - << " to be on GPU (while checking arguments for " << op << ')'; - STD_TORCH_CHECK(false, oss.str()); - } - const int d1 = t1.tensor.get_device_index(); - const int d2 = t2.tensor.get_device_index(); - 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 - -namespace QUTLASS { - Tensor matmul_mxf4_bf16_tn(Tensor const& A, Tensor const& B, Tensor const& A_sf, diff --git a/qutlass/csrc/include/bindings_utils.h b/qutlass/csrc/include/bindings_utils.h new file mode 100644 index 0000000..c97cf92 --- /dev/null +++ b/qutlass/csrc/include/bindings_utils.h @@ -0,0 +1,141 @@ +#pragma once + +#include +#include +#include + +#include +#include +#include + +namespace { + +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) { + 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_type_string(std::ostream& out, DeviceType type) { + out << device_type_name(type); +} + +inline void append_device_string(std::ostream& out, const Tensor& t) { + const auto device = t.device(); + append_device_type_string(out, device.type()); + if (device.has_index()) { + out << ':' << device.index(); + } +} + +inline void check_contiguous(const char* op, const TensorArg& t) { + if (t.tensor.is_contiguous()) { + return; + } + std::ostringstream oss; + oss << "Expected contiguous tensor, but got non-contiguous tensor for " << t + << " (while checking arguments for " << op << ")"; + STD_TORCH_CHECK(false, oss.str()); +} + +inline void check_all_contiguous(const char* op, + std::initializer_list args) { + for (const auto& arg : args) { + check_contiguous(op, arg); + } +} + +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_cuda = t1.tensor.device().is_cuda(); + const bool t2_cuda = t2.tensor.device().is_cuda(); + if (!t1_cuda || !t2_cuda) { + std::ostringstream oss; + if (!t1_cuda) { + oss << "Tensor for " << t1 << " is on CPU, "; + } + if (!t2_cuda) { + oss << "Tensor for " << t2 << " is on CPU, "; + } + oss << "but expected " << ((t1_cuda && t2_cuda) ? "them" : "it") + << " to be on GPU (while checking arguments for " << op << ')'; + STD_TORCH_CHECK(false, oss.str()); + } + const int d1 = t1.tensor.get_device_index(); + const int d2 = t2.tensor.get_device_index(); + 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 From a673d54c662f952f1395fe9a739b722b4307fe84 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Wed, 24 Jun 2026 15:32:57 +0000 Subject: [PATCH 03/13] small updates to bindings_utils.h to intent easier to understand. Signed-off-by: Chris Leonard --- qutlass/csrc/include/bindings_utils.h | 45 +++++++++++++-------------- 1 file changed, 21 insertions(+), 24 deletions(-) diff --git a/qutlass/csrc/include/bindings_utils.h b/qutlass/csrc/include/bindings_utils.h index c97cf92..e16bfd4 100644 --- a/qutlass/csrc/include/bindings_utils.h +++ b/qutlass/csrc/include/bindings_utils.h @@ -10,6 +10,9 @@ 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; @@ -21,6 +24,8 @@ struct TensorArg { 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 << '\''; @@ -51,32 +56,24 @@ inline const char* device_type_name(DeviceType type) { } } -inline void append_device_type_string(std::ostream& out, DeviceType type) { - out << device_type_name(type); -} - inline void append_device_string(std::ostream& out, const Tensor& t) { const auto device = t.device(); - append_device_type_string(out, device.type()); + out << device_type_name(device.type()); if (device.has_index()) { out << ':' << device.index(); } } -inline void check_contiguous(const char* op, const TensorArg& t) { - if (t.tensor.is_contiguous()) { - return; - } - std::ostringstream oss; - oss << "Expected contiguous tensor, but got non-contiguous tensor for " << t - << " (while checking arguments for " << op << ")"; - STD_TORCH_CHECK(false, oss.str()); -} - inline void check_all_contiguous(const char* op, std::initializer_list args) { for (const auto& arg : args) { - check_contiguous(op, arg); + 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()); } } @@ -97,22 +94,22 @@ inline void check_device_type_cuda(const char* op, inline void check_same_gpu(const char* op, const TensorArg& t1, const TensorArg& t2) { - const bool t1_cuda = t1.tensor.device().is_cuda(); - const bool t2_cuda = t2.tensor.device().is_cuda(); - if (!t1_cuda || !t2_cuda) { + 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_cuda) { + if (t1_cpu) { oss << "Tensor for " << t1 << " is on CPU, "; } - if (!t2_cuda) { + if (t2_cpu) { oss << "Tensor for " << t2 << " is on CPU, "; } - oss << "but expected " << ((t1_cuda && t2_cuda) ? "them" : "it") + 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 int d1 = t1.tensor.get_device_index(); - const int d2 = t2.tensor.get_device_index(); + const auto d1 = t1.tensor.device(); + const auto d2 = t2.tensor.device(); if (d1 == d2) { return; } From c1c4f9362965e6029b33b9ea25f6df16ee6af883 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Tue, 7 Jul 2026 14:11:46 +0000 Subject: [PATCH 04/13] very minor formatting fixes Signed-off-by: Chris Leonard --- qutlass/csrc/bindings.cpp | 5 +++-- qutlass/csrc/include/common.h | 2 +- qutlass/csrc/include/gemm.h | 1 + 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/qutlass/csrc/bindings.cpp b/qutlass/csrc/bindings.cpp index 2f3b74d..37c91a7 100644 --- a/qutlass/csrc/bindings.cpp +++ b/qutlass/csrc/bindings.cpp @@ -233,10 +233,11 @@ std::tuple fusedQuantizeMxAbsMax(Tensor const& A, fusedQuantizeMxAbsMax_host(OUT, OUT_sf, A, B); } 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 || TARGET_CUDA_ARCH == 101 || TARGET_CUDA_ARCH == 110 + // FIXME: add input global_scale to interface for consistency auto global_scale = - torch::stable::new_zeros(A, {1}, ScalarType::Float); //FIXME: add input global_scale to interface for consistency + 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); diff --git a/qutlass/csrc/include/common.h b/qutlass/csrc/include/common.h index 923a5c5..8b96c81 100644 --- a/qutlass/csrc/include/common.h +++ b/qutlass/csrc/include/common.h @@ -31,7 +31,7 @@ /** * Panic wrapper for unwinding CUDA runtime errors */ - #define CUDA_CHECK(status) \ +#define CUDA_CHECK(status) \ { \ cudaError_t error = status; \ STD_TORCH_CHECK(error == cudaSuccess, cudaGetErrorString(error)); \ diff --git a/qutlass/csrc/include/gemm.h b/qutlass/csrc/include/gemm.h index b08284d..245c0c4 100644 --- a/qutlass/csrc/include/gemm.h +++ b/qutlass/csrc/include/gemm.h @@ -14,6 +14,7 @@ * limitations under the License. */ + #pragma once #include From 1a61494d0440b4d6bf31476ed9c6f2c1ac05c33f Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Thu, 9 Jul 2026 13:11:36 +0000 Subject: [PATCH 05/13] added the functions that I removed from bindings.cpp back, but made them use torch abi stable APIs. Also, moved functions that are not used by vllm behind a guard QUTLASS_MINIMAL_BUILD, similar to how they were behind the QUTLASS_DISABLE_PYBIND guard before Signed-off-by: Chris Leonard --- qutlass/csrc/bindings.cpp | 166 +++++++++++++++++++++++++++++++++++++- 1 file changed, 164 insertions(+), 2 deletions(-) diff --git a/qutlass/csrc/bindings.cpp b/qutlass/csrc/bindings.cpp index 37c91a7..df679fc 100644 --- a/qutlass/csrc/bindings.cpp +++ b/qutlass/csrc/bindings.cpp @@ -22,6 +22,7 @@ #include "include/bindings_utils.h" #include "include/gemm.h" #include "include/fused_quantize_host.h" +#include "include/backward_host.h" namespace QUTLASS { @@ -136,6 +137,84 @@ Tensor matmul_ada_mxf4_bf16_tn(Tensor const& A, return C; } +Tensor matmul_mxf8_bf16_tn(Tensor const& A, + Tensor const& B, + Tensor const& A_sf, + Tensor const& B_sf, + Tensor const& alpha) +{ + 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::stable::new_empty(A, {M, N}, ScalarType::BFloat16); + + matmul_host_mxf8_bf16_tn(OUT, A, B, A_sf, B_sf, alpha); + + return OUT; +} + +Tensor matmul_mxf8_bf16_nn(Tensor const& A, + Tensor const& B, + Tensor const& A_sf, + Tensor const& B_sf, + Tensor const& alpha) +{ + 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::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(Tensor const& A, Tensor const& B, Tensor& OUT, @@ -344,26 +423,109 @@ std::tuple fusedQuantizeNvAbsMax(Tensor const& A, return std::make_tuple(OUT, OUT_sf); } +void backward_t_bf16(Tensor const& x, + Tensor const& h, + Tensor& xh_e2m1, + Tensor& xh_e8m0) +{ + 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)), + get_current_cuda_stream(h.get_device_index())); +} + +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)), + get_current_cuda_stream(h.get_device_index())); +} + +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.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(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.mutable_data_ptr(), + shared_exps.mutable_data_ptr(), + get_current_cuda_stream(x_fp4.get_device_index())); +} + } // 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("fusedQuantizeMxQuestWithMask(Tensor A, Tensor R, Tensor OUT, Tensor OUT_sf, Tensor OUT_mask) -> (Tensor, 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 } 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("fusedQuantizeMxQuestWithMask", TORCH_BOX(&QUTLASS::fusedQuantizeMxQuestWithMask)); 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 } From 7e79915427c73ec8cf1e85af24e829164619dd30 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Tue, 14 Jul 2026 15:00:35 +0000 Subject: [PATCH 06/13] Guard QUTLASS_MINIMAL_BUILD op defs and restore fusedQuantizeNv for vLLM --- qutlass/csrc/bindings.cpp | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/qutlass/csrc/bindings.cpp b/qutlass/csrc/bindings.cpp index df679fc..9cc44ac 100644 --- a/qutlass/csrc/bindings.cpp +++ b/qutlass/csrc/bindings.cpp @@ -251,6 +251,7 @@ std::tuple fusedQuantizeMxQuest(Tensor const& A, return std::make_tuple(OUT, OUT_sf); } +#ifndef QUTLASS_MINIMAL_BUILD std::tuple fusedQuantizeMxQuestWithMask( Tensor const& A, Tensor const& B, @@ -286,6 +287,7 @@ std::tuple fusedQuantizeMxQuestWithMask( return std::make_tuple(OUT, OUT_sf, OUT_mask); } +#endif // QUTLASS_MINIMAL_BUILD std::tuple fusedQuantizeMxAbsMax(Tensor const& A, Tensor const& B, @@ -423,6 +425,20 @@ std::tuple fusedQuantizeNvAbsMax(Tensor const& A, return std::make_tuple(OUT, OUT_sf); } +// Legacy vLLM entry point (pre-Quest/AbsMax split). Same dispatch as +// fusedQuantizeNvAbsMax: Quest=false kernels for 16/32/64, SM100 special +// path at 128, Had128 on SM120. +// Adding this to run unit test on vllm, +// will need to figure out how to resolve this before merging. +std::tuple fusedQuantizeNv(Tensor const& A, + Tensor const& B, + Tensor& OUT, + Tensor& OUT_sf, + Tensor const& global_scale) { + return fusedQuantizeNvAbsMax(A, B, OUT, OUT_sf, global_scale); +} + +#ifndef QUTLASS_MINIMAL_BUILD void backward_t_bf16(Tensor const& x, Tensor const& h, Tensor& xh_e2m1, @@ -488,6 +504,7 @@ void mxfp4_transpose_mxfp8(Tensor const& x_fp4, shared_exps.mutable_data_ptr(), get_current_cuda_stream(x_fp4.get_device_index())); } +#endif // QUTLASS_MINIMAL_BUILD } // namespace QUTLASS @@ -501,6 +518,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_qutlass_C, ops) { 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)"); + ops.def("fusedQuantizeNv(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) -> ()"); @@ -520,6 +538,7 @@ STABLE_TORCH_LIBRARY_IMPL(_qutlass_C, CUDA, ops) { ops.impl("fusedQuantizeMxAbsMax", TORCH_BOX(&QUTLASS::fusedQuantizeMxAbsMax)); ops.impl("fusedQuantizeNvQuest", TORCH_BOX(&QUTLASS::fusedQuantizeNvQuest)); ops.impl("fusedQuantizeNvAbsMax", TORCH_BOX(&QUTLASS::fusedQuantizeNvAbsMax)); + ops.impl("fusedQuantizeNv", TORCH_BOX(&QUTLASS::fusedQuantizeNv)); #ifndef QUTLASS_MINIMAL_BUILD ops.impl("fusedQuantizeMxQuestWithMask", TORCH_BOX(&QUTLASS::fusedQuantizeMxQuestWithMask)); ops.impl("backward_t_bf16", TORCH_BOX(&QUTLASS::backward_t_bf16)); From 650f171c0d148bce9cc7952f6ea2dc744d9fb433 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Wed, 15 Jul 2026 18:16:46 +0000 Subject: [PATCH 07/13] removed temporary fusedQuantizeNv wrapper Signed-off-by: Chris Leonard --- qutlass/csrc/bindings.cpp | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/qutlass/csrc/bindings.cpp b/qutlass/csrc/bindings.cpp index 9cc44ac..b074672 100644 --- a/qutlass/csrc/bindings.cpp +++ b/qutlass/csrc/bindings.cpp @@ -425,19 +425,6 @@ std::tuple fusedQuantizeNvAbsMax(Tensor const& A, return std::make_tuple(OUT, OUT_sf); } -// Legacy vLLM entry point (pre-Quest/AbsMax split). Same dispatch as -// fusedQuantizeNvAbsMax: Quest=false kernels for 16/32/64, SM100 special -// path at 128, Had128 on SM120. -// Adding this to run unit test on vllm, -// will need to figure out how to resolve this before merging. -std::tuple fusedQuantizeNv(Tensor const& A, - Tensor const& B, - Tensor& OUT, - Tensor& OUT_sf, - Tensor const& global_scale) { - return fusedQuantizeNvAbsMax(A, B, OUT, OUT_sf, global_scale); -} - #ifndef QUTLASS_MINIMAL_BUILD void backward_t_bf16(Tensor const& x, Tensor const& h, @@ -518,7 +505,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_qutlass_C, ops) { 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)"); - ops.def("fusedQuantizeNv(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) -> ()"); @@ -538,7 +524,6 @@ STABLE_TORCH_LIBRARY_IMPL(_qutlass_C, CUDA, ops) { ops.impl("fusedQuantizeMxAbsMax", TORCH_BOX(&QUTLASS::fusedQuantizeMxAbsMax)); ops.impl("fusedQuantizeNvQuest", TORCH_BOX(&QUTLASS::fusedQuantizeNvQuest)); ops.impl("fusedQuantizeNvAbsMax", TORCH_BOX(&QUTLASS::fusedQuantizeNvAbsMax)); - ops.impl("fusedQuantizeNv", TORCH_BOX(&QUTLASS::fusedQuantizeNv)); #ifndef QUTLASS_MINIMAL_BUILD ops.impl("fusedQuantizeMxQuestWithMask", TORCH_BOX(&QUTLASS::fusedQuantizeMxQuestWithMask)); ops.impl("backward_t_bf16", TORCH_BOX(&QUTLASS::backward_t_bf16)); From d27426342b5c9093432811b58934fc891565d7c1 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Wed, 15 Jul 2026 19:52:11 +0000 Subject: [PATCH 08/13] Wire standalone pip build for stable ABI and dispatch via torch.ops._qutlass_C Signed-off-by: Chris Leonard --- qutlass/__init__.py | 32 +++++++++++++++-------------- qutlass/csrc/bindings.cpp | 5 +++++ qutlass/csrc/include/registration.h | 18 ++++++++++++++++ setup.py | 15 +++++++++++--- 4 files changed, 52 insertions(+), 18 deletions(-) create mode 100644 qutlass/csrc/include/registration.h 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 b074672..5560655 100644 --- a/qutlass/csrc/bindings.cpp +++ b/qutlass/csrc/bindings.cpp @@ -533,3 +533,8 @@ STABLE_TORCH_LIBRARY_IMPL(_qutlass_C, CUDA, ops) { ops.impl("mxfp4_transpose_mxfp8", TORCH_BOX(&QUTLASS::mxfp4_transpose_mxfp8)); #endif } + +#ifndef QUTLASS_MINIMAL_BUILD +#include "include/registration.h" +REGISTER_EXTENSION(_CUDA) +#endif 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/setup.py b/setup.py index 53d048a..aeb2075 100644 --- a/setup.py +++ b/setup.py @@ -150,14 +150,23 @@ 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}, + cmdclass={"build_ext": BuildExtension.with_options(no_python_abi_suffix=True)}, ) From 00028965c541e60b69bd48dea912c0cc91ec5363 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Wed, 15 Jul 2026 19:59:04 +0000 Subject: [PATCH 09/13] Fix pip editable build: use plain BuildExtension instead of with_options Signed-off-by: Chris Leonard --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index aeb2075..11121b1 100644 --- a/setup.py +++ b/setup.py @@ -168,5 +168,5 @@ def third_party_cmake(): py_limited_api=True, ) ], - cmdclass={"build_ext": BuildExtension.with_options(no_python_abi_suffix=True)}, + cmdclass={"build_ext": BuildExtension}, ) From 6412cb019de9d6916d7c4e513c830c1df99d1381 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Wed, 15 Jul 2026 20:08:48 +0000 Subject: [PATCH 10/13] Remove legacy torch/extension.h from quartet_bwd_sm120.cu Signed-off-by: Chris Leonard --- qutlass/csrc/quartet_bwd_sm120.cu | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) 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; } From 5e0f439d6f8d6d3f6e21e31952df55da94d3ec1f Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Fri, 17 Jul 2026 13:22:38 +0000 Subject: [PATCH 11/13] Require PyTorch >= 2.11 to match TORCH_TARGET_VERSION, and tag abi3 wheels with cp39 via python_requires and bdist_wheel py_limited_api so installers accept the package across supported CPython versions. Signed-off-by: Chris Leonard --- setup.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/setup.py b/setup.py index 11121b1..5a52f17 100644 --- a/setup.py +++ b/setup.py @@ -116,8 +116,8 @@ def third_party_cmake(): if not m: raise RuntimeError(f"Cannot parse PyTorch version '{torch_version}'") major, minor = map(int, m.groups()) - if major < 2 or (major == 2 and minor < 7): - raise RuntimeError(f"PyTorch version must be >= 2.7, but found {torch_version}") + if major < 2 or (major == 2 and minor < 11): + raise RuntimeError(f"PyTorch version must be >= 2.11, but found {torch_version}") third_party_cmake() remove_unwanted_pytorch_nvcc_flags() @@ -128,6 +128,8 @@ def third_party_cmake(): author_email="Roberto.LopezCastro@ist.ac.at", description="CUTLASS-Powered Quantized BLAS for Deep Learning.", packages=find_packages(), + python_requires=">=3.9", + options={"bdist_wheel": {"py_limited_api": "cp39"}}, ext_modules=[ CUDAExtension( name="qutlass._CUDA", From b11d30ef25c5687f1f236dd24ba35dc99bcd7094 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Fri, 17 Jul 2026 13:44:34 +0000 Subject: [PATCH 12/13] Move the Python floor out of setup.py so setuptools honors it under the [project] table and no longer crashes SpecifierSet on None during uv/pip editable builds. Signed-off-by: Chris Leonard --- pyproject.toml | 1 + setup.py | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 394ffa7..ffd9a60 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,6 +8,7 @@ version = "0.2.0" description = "qutlass" authors = [{name = "Roberto L. Castro", email = "Roberto.LopezCastro@ist.ac.at"}] license = {text = "Apache-2.0"} +requires-python = ">=3.9" dependencies = [] [tool.setuptools] diff --git a/setup.py b/setup.py index 5a52f17..f890768 100644 --- a/setup.py +++ b/setup.py @@ -128,7 +128,6 @@ def third_party_cmake(): author_email="Roberto.LopezCastro@ist.ac.at", description="CUTLASS-Powered Quantized BLAS for Deep Learning.", packages=find_packages(), - python_requires=">=3.9", options={"bdist_wheel": {"py_limited_api": "cp39"}}, ext_modules=[ CUDAExtension( From 72869c6df67d28f5013784fb917072fa397a34e8 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Fri, 17 Jul 2026 14:00:51 +0000 Subject: [PATCH 13/13] Define TORCH_TARGET_MAJOR/MINOR once, encode TORCH_TARGET_VERSION from them, and use that both for the setup.py version check and the -D compile flag so the bound no longer drifts from the hardcoded 0x020B macro. Signed-off-by: Chris Leonard --- setup.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/setup.py b/setup.py index f890768..c6c4483 100644 --- a/setup.py +++ b/setup.py @@ -50,6 +50,12 @@ def detect_cc(): cc = detect_cc() +# 2.11 -> 0x020b000000000000ULL +TORCH_TARGET_MAJOR = 2 +TORCH_TARGET_MINOR = 11 +TORCH_TARGET_VERSION = (TORCH_TARGET_MAJOR << 56) | (TORCH_TARGET_MINOR << 48) +TORCH_TARGET_VERSION_MACRO = f"0x{TORCH_TARGET_VERSION:016X}ULL" + def get_cuda_arch_flags(): flags = [ @@ -116,8 +122,11 @@ def third_party_cmake(): if not m: raise RuntimeError(f"Cannot parse PyTorch version '{torch_version}'") major, minor = map(int, m.groups()) - if major < 2 or (major == 2 and minor < 11): - raise RuntimeError(f"PyTorch version must be >= 2.11, but found {torch_version}") + if (major, minor) < (TORCH_TARGET_MAJOR, TORCH_TARGET_MINOR): + raise RuntimeError( + f"PyTorch version must be >= {TORCH_TARGET_MAJOR}.{TORCH_TARGET_MINOR} " + f"but found {torch_version}" + ) third_party_cmake() remove_unwanted_pytorch_nvcc_flags() @@ -154,11 +163,11 @@ def third_party_cmake(): "cxx": [ "-std=c++17", "-DUSE_CUDA", - "-DTORCH_TARGET_VERSION=0x020B000000000000ULL", + f"-DTORCH_TARGET_VERSION={TORCH_TARGET_VERSION_MACRO}", ], "nvcc": [ "-DUSE_CUDA", - "-DTORCH_TARGET_VERSION=0x020B000000000000ULL", + f"-DTORCH_TARGET_VERSION={TORCH_TARGET_VERSION_MACRO}", *get_cuda_arch_flags(), ], },