diff --git a/qutlass/csrc/bindings.cpp b/qutlass/csrc/bindings.cpp index 600c8f3..974109d 100644 --- a/qutlass/csrc/bindings.cpp +++ b/qutlass/csrc/bindings.cpp @@ -308,7 +308,7 @@ std::tuple fusedQuantizeMxAbsMax(torch::Tensor con } else if(HAD_GS==64){ fusedQuantizeMxAbsMaxHad64_host(OUT, OUT_sf, A, B); } else if(HAD_GS==128){ -#if TARGET_CUDA_ARCH == 100 +#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 fusedQuantizeMxAbsMax_host_sm100(OUT, OUT_sf, A, B, global_scale); @@ -399,7 +399,7 @@ std::tuple fusedQuantizeNvAbsMax(torch::Tensor con } else if(HAD_GS==64){ fusedQuantizeNvAbsMaxHad64_host(OUT, OUT_sf, A, B, global_scale); } else if(HAD_GS==128){ -#if TARGET_CUDA_ARCH == 100 +#if TARGET_CUDA_ARCH == 100 || TARGET_CUDA_ARCH == 101 || TARGET_CUDA_ARCH == 110 fusedQuantizeNvAbsMax_host_sm100(OUT, OUT_sf, A, B, global_scale); #elif TARGET_CUDA_ARCH == 120 fusedQuantizeNvAbsMaxHad128_host(OUT, OUT_sf, A, B, global_scale); diff --git a/qutlass/csrc/fused_quantize_mx_sm100.cu b/qutlass/csrc/fused_quantize_mx_sm100.cu index 5d946f9..c0d1016 100644 --- a/qutlass/csrc/fused_quantize_mx_sm100.cu +++ b/qutlass/csrc/fused_quantize_mx_sm100.cu @@ -205,7 +205,7 @@ void fusedQuantizeMxAbsMax_host_sm100(torch::Tensor& D, torch::Tensor const& B, torch::Tensor const& global_scale) { -#if TARGET_CUDA_ARCH == 100 +#if TARGET_CUDA_ARCH == 100 || TARGET_CUDA_ARCH == 101 || TARGET_CUDA_ARCH == 110 int32_t M = A.numel() / 128; int32_t N = B.size(1); int32_t K = 128; diff --git a/qutlass/csrc/fused_quantize_nv_sm100.cu b/qutlass/csrc/fused_quantize_nv_sm100.cu index 4dc3b9c..b696066 100644 --- a/qutlass/csrc/fused_quantize_nv_sm100.cu +++ b/qutlass/csrc/fused_quantize_nv_sm100.cu @@ -205,7 +205,7 @@ void fusedQuantizeNvAbsMax_host_sm100(torch::Tensor& D, torch::Tensor const& B, torch::Tensor const& global_scale) { -#if TARGET_CUDA_ARCH == 100 +#if TARGET_CUDA_ARCH == 100 || TARGET_CUDA_ARCH == 101 || TARGET_CUDA_ARCH == 110 int32_t M = A.numel() / 128; int32_t N = B.size(1); int32_t K = 128; diff --git a/qutlass/csrc/gemm.cu b/qutlass/csrc/gemm.cu index 57ebcc2..a19383a 100644 --- a/qutlass/csrc/gemm.cu +++ b/qutlass/csrc/gemm.cu @@ -200,7 +200,7 @@ void matmul_host_mxf4_bf16_tn(torch::Tensor& D, using LayoutBTag = cutlass::layout::ColumnMajor; static constexpr int AlignmentB = 128; -#if TARGET_CUDA_ARCH == 100 //TODO: improve tuning +#if TARGET_CUDA_ARCH == 100 || TARGET_CUDA_ARCH == 101 || TARGET_CUDA_ARCH == 110 //TODO: improve tuning using ArchTag = cutlass::arch::Sm100; if(m<=16){ using MmaTileShape = Shape<_128,_128,_256>; @@ -276,7 +276,7 @@ void matmul_host_nvf4_bf16_tn(torch::Tensor& D, using LayoutBTag = cutlass::layout::ColumnMajor; static constexpr int AlignmentB = 32; -#if TARGET_CUDA_ARCH == 100 //TODO: improve tuning +#if TARGET_CUDA_ARCH == 100 || TARGET_CUDA_ARCH == 101 || TARGET_CUDA_ARCH == 110 //TODO: improve tuning using ArchTag = cutlass::arch::Sm100; if(m<=16){ using MmaTileShape = Shape<_128,_128,_256>; @@ -354,7 +354,7 @@ void matmul_host_mxf8_bf16_tn(torch::Tensor& D, using LayoutBTag = cutlass::layout::ColumnMajor; static constexpr int AlignmentB = 16; -#if TARGET_CUDA_ARCH == 100 +#if TARGET_CUDA_ARCH == 100 || TARGET_CUDA_ARCH == 101 || TARGET_CUDA_ARCH == 110 using ArchTag = cutlass::arch::Sm100; if(m<=8192){ @@ -414,7 +414,7 @@ void matmul_host_mxf8_bf16_nn(torch::Tensor& D, using LayoutBTag = cutlass::layout::ColumnMajor; static constexpr int AlignmentB = 16; -#if TARGET_CUDA_ARCH == 100 +#if TARGET_CUDA_ARCH == 100 || TARGET_CUDA_ARCH == 101 || TARGET_CUDA_ARCH == 110 using ArchTag = cutlass::arch::Sm100; if(m<=8192){ diff --git a/setup.py b/setup.py index 3fd3495..53d048a 100644 --- a/setup.py +++ b/setup.py @@ -57,6 +57,21 @@ def get_cuda_arch_flags(): "arch=compute_120a,code=sm_120a", "-gencode", "arch=compute_100a,code=sm_100a", + ] + # Thor / GB10 (SM 10.1, renamed SM 11.0 in CUDA 13.0) shares the SM100 + # tcgen05 block-scaled fp4 path. Only relevant when building on such a + # device (this is a per-machine build), so gate on the local capability + # rather than the CPU arch -- GB200/GH200 are also aarch64 but not Thor. + # Emit its SASS under whichever name the active toolkit uses: sm_110a on + # CUDA >= 13.0, sm_101a on <= 12.9. + cuda_ver = torch.version.cuda + if cc in (101, 110) and cuda_ver is not None: + cuda_major, cuda_minor = (int(v) for v in cuda_ver.split(".")[:2]) + if (cuda_major, cuda_minor) >= (13, 0): + flags += ["-gencode", "arch=compute_110a,code=sm_110a"] + else: + flags += ["-gencode", "arch=compute_101a,code=sm_101a"] + flags += [ "--expt-relaxed-constexpr", "--use_fast_math", "-std=c++17",