已安装CUDA及驱动但Python TVM检测不到CUDA设备问题求助

核心问题排查 #

首先运行以下代码确认当前TVM是否内置CUDA支持：

import tvm
print(tvm.support("cuda"))

如果输出为False，说明你使用的TVM版本编译时未开启CUDA支持，这是此类问题的最常见原因。

解决方案：编译带CUDA支持的TVM #

1. 前置环境校验 #

• 确认CUDA Toolkit已正确安装，且版本<=驱动支持的最高版本（你当前驱动支持最高CUDA 11.5）
• 确认CUDA环境变量已配置，在~/.bashrc中添加以下内容后执行source ~/.bashrc生效：

export PATH=/usr/local/cuda/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH

• 执行nvcc -V确认能正常输出版本信息，环境配置正确。

2. 源码编译TVM #

• 拉取TVM源码：

git clone --recursive https://github.com/apache/tvm.git
cd tvm
# 可切换到适配你环境的稳定版本分支，如v0.14.0

• 配置编译参数：

mkdir build
cp cmake/config.cmake build/
cd build

编辑build/config.cmake文件，修改以下配置：

• 将set(USE_CUDA OFF)改为set(USE_CUDA ON)
• 新增配置set(CMAKE_CUDA_ARCHITECTURES 75)，适配你使用的T4显卡（算力7.5）
• 编译TVM：

cmake ..
make -j$(nproc)

3. 配置Python TVM环境 #

在~/.bashrc中添加以下内容，替换为你自己的TVM源码路径后执行source ~/.bashrc生效：

export TVM_HOME=/path/to/your/tvm/source/code
export PYTHONPATH=$TVM_HOME/python:$TVM_HOME/topi/python:$PYTHONPATH

验证修复结果 #

重新运行你的测试代码：

import tvm
ctx = tvm.device("cuda", 0)
print("ctx: ", ctx.exist)

正常会输出True。

其他排查点 #

• 执行echo $CUDA_VISIBLE_DEVICES检查该环境变量是否被错误设置为无效值，可手动设置export CUDA_VISIBLE_DEVICES=0,1,2对应你设备上的3张NVIDIA显卡
• 确认运行Python的环境和配置TVM的环境一致，避免多虚拟环境冲突

内容的提问来源于stack exchange，提问作者si yan