Error installing SageAttention2 on Ubuntu

[tahseensheik]

I had spun up an Ubuntu Cloud PC which has CUDA 12.4 installed on it. When I try to install SageAttention2, I get the following error:

root@8e882ecbcede:/workspace/ComfyUI/SageAttention# python setup.py install

running install
/usr/local/lib/python3.10/dist-packages/setuptools/_distutils/cmd.py:66: SetuptoolsDeprecationWarning: setup.py install is deprecated.
!!

    ********************************************************************************
    Please avoid running ``setup.py`` directly.
    Instead, use pypa/build, pypa/installer or other
    standards-based tools.

    See https://blog.ganssle.io/articles/2021/10/setup-py-deprecated.html for details.
    ********************************************************************************

!!
self.initialize_options()
/usr/local/lib/python3.10/dist-packages/setuptools/_distutils/cmd.py:66: EasyInstallDeprecationWarning: easy_install command is deprecated.
!!

    ********************************************************************************
    Please avoid running ``setup.py`` and ``easy_install``.
    Instead, use pypa/build, pypa/installer or other
    standards-based tools.

    See https://github.com/pypa/setuptools/issues/917 for details.
    ********************************************************************************

!!
self.initialize_options()
running bdist_egg
running egg_info
writing sageattention.egg-info/PKG-INFO
writing dependency_links to sageattention.egg-info/dependency_links.txt
writing top-level names to sageattention.egg-info/top_level.txt
/usr/local/lib/python3.10/dist-packages/torch/utils/cpp_extension.py:497: UserWarning: Attempted to use ninja as the BuildExtension backend but we could not find ninja.. Falling back to using the slow distutils backend.
warnings.warn(msg.format('we could not find ninja.'))
reading manifest file 'sageattention.egg-info/SOURCES.txt'
adding license file 'LICENSE'
writing manifest file 'sageattention.egg-info/SOURCES.txt'
installing library code to build/bdist.linux-x86_64/egg
running install_lib
running build_py
creating build/lib.linux-x86_64-cpython-310/sageattention
copying sageattention/init.py -> build/lib.linux-x86_64-cpython-310/sageattention
copying sageattention/core.py -> build/lib.linux-x86_64-cpython-310/sageattention
copying sageattention/quant.py -> build/lib.linux-x86_64-cpython-310/sageattention
creating build/lib.linux-x86_64-cpython-310/sageattention/triton
copying sageattention/triton/init.py -> build/lib.linux-x86_64-cpython-310/sageattention/triton
copying sageattention/triton/attn_qk_int8_block_varlen.py -> build/lib.linux-x86_64-cpython-310/sageattention/triton
copying sageattention/triton/attn_qk_int8_per_block.py -> build/lib.linux-x86_64-cpython-310/sageattention/triton
copying sageattention/triton/attn_qk_int8_per_block_causal.py -> build/lib.linux-x86_64-cpython-310/sageattention/triton
copying sageattention/triton/attn_qk_int8_per_block_causal_varlen.py -> build/lib.linux-x86_64-cpython-310/sageattention/triton
copying sageattention/triton/quant_per_block.py -> build/lib.linux-x86_64-cpython-310/sageattention/triton
copying sageattention/triton/quant_per_block_varlen.py -> build/lib.linux-x86_64-cpython-310/sageattention/triton
copying sageattention/triton/quant_per_thread.py -> build/lib.linux-x86_64-cpython-310/sageattention/triton
running build_ext
/usr/local/lib/python3.10/dist-packages/torch/utils/cpp_extension.py:426: UserWarning: There are no x86_64-linux-gnu-g++ version bounds defined for CUDA version 12.4
warnings.warn(f'There are no {compiler_name} version bounds defined for CUDA version {cuda_str_version}')
building 'sageattention.qattn' extension
x86_64-linux-gnu-g++ -Wno-unused-result -Wsign-compare -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -g -fwrapv -O2 -fPIC -I/usr/local/lib/python3.10/dist-packages/torch/include -I/usr/local/lib/python3.10/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.10/dist-packages/torch/include/TH -I/usr/local/lib/python3.10/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.10 -c csrc/qattn/pybind.cpp -o build/temp.linux-x86_64-cpython-310/csrc/qattn/pybind.o -g -O3 -fopenmp -lgomp -std=c++17 -DENABLE_BF16 -D_GLIBCXX_USE_CXX11_ABI=0 -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE="gcc" -DPYBIND11_STDLIB="libstdcpp" -DPYBIND11_BUILD_ABI="cxxabi1011" -DTORCH_EXTENSION_NAME=qattn -D_GLIBCXX_USE_CXX11_ABI=0
/usr/local/cuda/bin/nvcc -I/usr/local/lib/python3.10/dist-packages/torch/include -I/usr/local/lib/python3.10/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.10/dist-packages/torch/include/TH -I/usr/local/lib/python3.10/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.10 -c csrc/qattn/qk_int_sv_f16_buffer_cuda.cu -o build/temp.linux-x86_64-cpython-310/csrc/qattn/qk_int_sv_f16_buffer_cuda.o -D__CUDA_NO_HALF_OPERATORS -D__CUDA_NO_HALF_CONVERSIONS -D__CUDA_NO_BFLOAT16_CONVERSIONS_ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr --compiler-options '-fPIC' -O3 -std=c++17 -U__CUDA_NO_HALF_OPERATORS__ -U__CUDA_NO_HALF_CONVERSIONS__ --use_fast_math --threads=8 -Xptxas=-v -diag-suppress=174 -D_GLIBCXX_USE_CXX11_ABI=0 -gencode arch=compute_89,code=sm_89 -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE="_gcc" -DPYBIND11_STDLIB="_libstdcpp" -DPYBIND11_BUILD_ABI="_cxxabi1011" -DTORCH_EXTENSION_NAME=_qattn -D_GLIBCXX_USE_CXX11_ABI=0
In file included from csrc/qattn/qk_int_sv_f16_buffer_cuda.cu:17:
csrc/qattn/../utils.cuh:36:37: warning: backslash-newline at end of file
36 | #define CHECK_LASTDIM_CONTIGUOUS(x)
|
In file included from csrc/qattn/qk_int_sv_f16_buffer_cuda.cu:17:
csrc/qattn/../utils.cuh:36:37: warning: backslash-newline at end of file
36 | #define CHECK_LASTDIM_CONTIGUOUS(x)
|
src/tcmalloc.cc:333] Attempt to free invalid pointer 0x7f093381e000
Aborted (core dumped)
error: command '/usr/local/cuda/bin/nvcc' failed with exit code 255