PyTorch-CUDA-v2.6镜像中查看CUDA架构（sm

PyTorch-CUDA-v2.6镜像中查看CUDA架构（sm_XX）的命令

在深度学习工程实践中，一个看似简单却常常引发连锁问题的操作是：确认当前环境支持的 CUDA 架构（sm_XX）。尤其是在使用预构建的PyTorch-CUDA-v2.6镜像时，开发者常默认“既然能跑 PyTorch，那肯定适配我的 GPU”，结果在编译自定义算子或迁移模型时遭遇no kernel image is available for execution on the device这类令人头疼的错误。

这类问题的根本原因往往不是驱动没装好，也不是容器权限配置错误，而是——编译目标与实际硬件 compute capability 不匹配。而这一切，都可以通过一条简洁的 Python 脚本提前规避。

我们不妨从一次典型的失败经历说起。假设你在本地 RTX 3090（compute capability 8.6）上训练了一个高性能自定义 CUDA 算子，一切顺利。然后你将代码推送到 CI/CD 流水线，在 A100（sm_80）服务器上运行测试，突然报错：

CUDA error: no kernel image is available for execution on the device

奇怪的是，torch.cuda.is_available()返回True，设备也能列出，但就是无法执行内核。这时你就该意识到：问题不在运行时，而在编译时——你的 CUDA 扩展没有为sm_80生成对应的二进制代码。

要解决这个问题，第一步就是：准确获取当前 GPU 的 compute capability，并将其转换为标准的 sm_XX 格式。

PyTorch 提供了最直接的方式：

import torch if torch.cuda.is_available(): capability = torch.cuda.get_device_capability() major, minor = capability print(f"Compute Capability: {major}.{minor} → sm_{major}{minor}") else: print("CUDA not accessible.")

输出示例：

Compute Capability: 8.0 → sm_80

这个值告诉你，当前设备属于 Ampere 架构，应使用-gencode arch=compute_80,code=sm_80进行编译。如果你看到的是sm_86，那就是面向 GA102 核心的消费级卡（如 RTX 3090），需要单独处理。

但别忘了，这一步必须在目标运行环境中执行。也就是说，不能在本地查完就完事，而要在 Docker 容器里运行这段代码。因为不同环境下的可见 GPU 可能不同，甚至同一个镜像在不同机器上启动后识别到的架构也可能是不一样的。

启动容器的标准命令如下：

docker run --gpus all -it pytorch-cuda:v2.6 bash

进入后立即运行上述 Python 脚本，确保你拿到的是真实部署环境的数据。有些团队习惯把setup.py中的nvcc参数写死成sm_75或sm_80，殊不知新一批机器已经升级到 H100（sm_90），导致新硬件的优势完全无法发挥。

说到这里，不得不提一下 NVIDIA 的 compute capability 演进逻辑。它并不仅仅是版本号递增那么简单，每个 major 版本都代表了一次架构革新：

Compute Capability	架构	关键特性
sm_70 / sm_75	Volta / Turing	引入 Tensor Core（初代）、独立线程调度
sm_80 / sm_86	Ampere	第三代 Tensor Core、稀疏化支持、FP32 性能翻倍
sm_90	Hopper	Transformer Engine、异步内存拷贝增强

比如sm_86相比sm_80，虽然同属 Ampere，但在 FP32 吞吐和内存带宽上有显著优化。如果只为sm_80编译，即使能在 RTX 3090 上运行，也可能无法充分利用其全部性能潜力。

因此，最佳实践是在构建通用镜像时采用multi-arch 编译策略，即在一个扩展中包含多个sm_XX的二进制代码。你可以这样设置：

from torch.utils.cpp_extension import CUDAExtension CUDAExtension( name='my_custom_op', sources=['custom_op.cu'], extra_compile_args={ 'nvcc': [ '-gencode', 'arch=compute_80,code=sm_80', '-gencode', 'arch=compute_86,code=sm_86', '-gencode', 'arch=compute_75,code=sm_75' ] } )

这种方式会生成一个“fat binary”，虽然体积稍大，但兼容性极强，适合发布给多类型用户的库。

再深入一层，你可能会问：那 PyTorch 自己是怎么做的？其实官方预编译的 PyTorch 包通常只包含主流架构（如 sm_70、sm_75、sm_80），并不会覆盖所有变种。这也是为什么当你使用非常规硬件（如 Jetson 或老旧显卡）时，可能需要从源码重新编译 PyTorch。

此外，还有一个容易被忽视的点：容器内的 CUDA Toolkit 版本是否支持目标架构。例如，CUDA 11.8 开始正式支持 sm_90（Hopper），而更早版本则无法编译对应代码。尽管 PyTorch-CUDA-v2.6 镜像大概率基于 CUDA 12.x，但仍建议验证：

nvcc --version

同时检查驱动版本：

nvidia-smi

确保驱动版本 ≥ 所需 CUDA 工具包的要求。否则即便镜像中有nvcc，也无法调用正确的 runtime API。

对于希望自动化检测流程的团队，可以编写一个简单的诊断脚本check_cuda_arch.py：

#!/usr/bin/env python import torch import subprocess import sys def main(): if not torch.cuda.is_available(): print("❌ CUDA is not available. Check driver and nvidia-docker setup.") sys.exit(1) device = torch.cuda.current_device() name = torch.cuda.get_device_name(device) capability = torch.cuda.get_device_capability() sm = f"sm_{capability[0]}{capability[1]}" print(f"✅ GPU Detected: {name}") print(f" Compute Capability: {capability[0]}.{capability[1]}") print(f" SM Architecture: {sm}") # 建议使用的 NVCC 参数 gencode = f"-gencode arch=compute_{capability[0]}{capability[1]},code={sm}" print(f" Recommended NVCC flag: {gencode}") # 检查 nvcc 是否可用 try: result = subprocess.run(['nvcc', '--version'], capture_output=True, text=True) if result.returncode == 0: version_line = [l for l in result.stdout.split('\n') if 'release' in l][0] print(f" NVCC Version: {version_line.strip()}") else: print("❌ nvcc not found or failed to execute.") except FileNotFoundError: print("❌ nvcc not found in PATH.") if __name__ == "__main__": main()

放入镜像后一键运行即可完成完整诊断。

回到最初的问题场景：如何在PyTorch-CUDA-v2.6镜像中查看 CUDA 架构？答案其实很简单，但背后的工程意义远不止一条命令这么轻巧。它是连接开发、测试、部署三个环节的关键锚点。

很多团队直到上线前才发现模型在生产环境跑不动，追根溯源却发现只是少加了一个-gencode参数。这种低级但高代价的失误，完全可以通过标准化流程避免。

最终建议如下：