setup.py

#
# Copyright (C) 2026 Max Kleinegger (mkleinegger@gmail.com). All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import os
import re

import torch
import torch.utils.cpp_extension as torch_cpp_ext
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

setup_dir = os.path.dirname(os.path.realpath(__file__))
torch_version = torch.__version__
MINIMUM_CUDA_CAPABILITY = (8, 6)


def remove_unwanted_pytorch_nvcc_flags():
    REMOVE_NVCC_FLAGS = [
        "-D__CUDA_NO_HALF_OPERATORS__",
        "-D__CUDA_NO_HALF_CONVERSIONS__",
        "-D__CUDA_NO_BFLOAT16_CONVERSIONS__",
        "-D__CUDA_NO_HALF2_OPERATORS__",
    ]
    for flag in REMOVE_NVCC_FLAGS:
        try:
            torch_cpp_ext.COMMON_NVCC_FLAGS.remove(flag)
        except ValueError:
            pass


def get_cuda_arch_flags():
    # Do not add a -gencode flag here. CUDAExtension selects the capabilities of
    # the visible GPUs by default and honours TORCH_CUDA_ARCH_LIST for native and
    # cross-compiled builds.
    flags = [
        "--expt-relaxed-constexpr",
        "--use_fast_math",
        "-std=c++20",
        "-O3",
        "-DNDEBUG",
        "-Xcompiler",
        "-funroll-loops",
        "-Xcompiler",
        "-ffast-math",
        "-Xcompiler",
        "-finline-functions",
    ]
    return flags


def validate_cuda_targets():
    arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST")
    if arch_list:
        print(f"CUDA architecture targets: {arch_list} (TORCH_CUDA_ARCH_LIST)")
        return

    if not torch.cuda.is_available():
        raise RuntimeError(
            "CUDA is not available. Set TORCH_CUDA_ARCH_LIST (for example, "
            "'8.6') when cross-compiling without a visible GPU."
        )

    capabilities = sorted(
        {
            tuple(map(int, torch.cuda.get_device_capability(device)))
            for device in range(torch.cuda.device_count())
        }
    )
    unsupported = [
        capability
        for capability in capabilities
        if capability < MINIMUM_CUDA_CAPABILITY
    ]
    if unsupported:
        formatted = ", ".join(f"{major}.{minor}" for major, minor in unsupported)
        raise RuntimeError(
            "CUDA capability must be >= "
            f"{MINIMUM_CUDA_CAPABILITY[0]}.{MINIMUM_CUDA_CAPABILITY[1]}; "
            f"found {formatted}"
        )

    targets = [f"{capability[0]}.{capability[1]}" for capability in capabilities]
    formatted = ";".join(targets)
    os.environ["TORCH_CUDA_ARCH_LIST"] = formatted
    print(f"CUDA architecture targets: {formatted} (visible GPUs)")


if __name__ == "__main__":
    validate_cuda_targets()
    print(f"PyTorch version: {torch_version}")
    m = re.match(r"^(\d+)\.(\d+)", torch_version)
    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}")

    remove_unwanted_pytorch_nvcc_flags()
    setup(
        name="matgptq_cuda",
        version="0.1.0",
        author="Elvir Crncevic",
        author_email="elvircrn@gmail.com",
        description="Efficient CUDA kernel with support for 2-8 'nested' bit-widths.",
        packages=["matgptq_cuda"],
        package_dir={"matgptq_cuda": "inference_lib/matgptq_cuda"},
        ext_modules=[
            CUDAExtension(
                name="matgptq_cuda._CUDA",
                sources=[
                    "inference_lib/matgptq_cuda/csrc/bindings.cpp",
                    "inference_lib/matgptq_cuda/csrc/utils.cpp",
                    "inference_lib/matgptq_cuda/csrc/matgptq_gemm.cu",
                ],
                include_dirs=[
                    os.path.join(setup_dir, "inference_lib/matgptq_cuda/csrc/include"),
                ],
                extra_compile_args={
                    "cxx": ["-std=c++20"],
                    "nvcc": get_cuda_arch_flags(),
                },
                extra_link_args=[
                    "-lcudart",
                    "-lcuda",
                ],
            )
        ],
        cmdclass={"build_ext": BuildExtension},
    )