編譯過深度學習框架的人十有八九都被“計算能力Compute Capability”這個東西卡過。要不就是編譯時刷出一排Unsupported gpu architecture要不就是同一個安裝包在這塊卡上跑得好好的換到另一塊卡上直接報no kernel image is available for execution on the device。我第一次在 caffe 上踩這個坑時硬生生折騰到凌晨兩點最后才發(fā)現(xiàn)問題就出在一個小數(shù)字上——GPU 計算能力。這篇就把這個數(shù)字的前因后果講透順手給出 caffe / pytorch / tensorflow 里設置計算能力的具體寫法和翻車應急預案。適合所有需要從源碼構建框架的開發(fā)者、算法工程師以及被 CUDA 環(huán)境逼到懷疑人生的人。1. 先搞懂“計算能力”這個數(shù)字它就是 GPU 的架構版本號1.1 數(shù)字拆解主版本號是“代”次版本號是“檔”計算能力Compute Capability寫出來長這樣8.6、7.5、9.0。它本質(zhì)上就是 GPU 硬件架構的版本號前一位代表大代際后一位代表同代際里的細分檔位。比如8.6意味著這是 Ampere 大架構下的消費卡檔位RTX 3090 就是 8.6。而 A100 同樣是 Ampere 架構卻是8.0。同一個架構代際內(nèi)不同芯片的指令集和硬件特性有差異所以才有8.0、8.6、8.9這種細分編號。在編譯參數(shù)里這個數(shù)字還會變成兩種寫法sm_86SASS 目標架構生成的就是這塊卡直接執(zhí)行的機器碼。compute_86PTX 目標架構生成的是中間表示后續(xù)可以由驅(qū)動現(xiàn)場編譯到具體 GPU。我見過不少新手把sm_75和compute_75當成一回事其實它們對應的是兩種不同形態(tài)的產(chǎn)物這也是后面絕大多數(shù)編譯坑的根源。1.2 一張表看懂各大架構的計算能力區(qū)間我把這些年主流的 GPU 架構按計算能力梳理了一下方便對號入座架構代際計算能力典型 GPU 型號大概年代Fermi2.0 / 2.1GTX 4xx / 5xx2010 左右Kepler3.0 / 3.5 / 3.7GTX 680 / Tesla K402012 左右Maxwell5.0 / 5.2 / 5.3GTX 750 / GTX 9802014 左右Pascal6.0 / 6.1Tesla P100 / GTX 1080 Ti2016 左右Volta7.0Tesla V1002017 左右Turing7.5RTX 2080 Ti / RTX 20602018 左右Ampere8.0 / 8.6A100 / RTX 30902020 左右Ada Lovelace8.9RTX 40902022 左右Hopper9.0H1002023 左右2024 年之后的新架構計算能力已經(jīng)突破 10.x后續(xù)還會有新編號具體以廠商官方規(guī)格頁為準。記住這張表的核心規(guī)律不同計算能力的 GPU它們的 SASS 機器碼是不通用的但 PTX 有更強的跨代兼容能力。1.3 計算能力決定了哪些實實在在的事情這個數(shù)字不只是一個“型號標簽”它背后綁定了一整套硬件特性能力。具體來說指令集版本不同sm_86 能用的某些匯編指令sm_75 是不認的。編譯器要針對具體計算能力生成對應指令。Tensor Core 等特性不同8.0 以上的 Ampere 架構對 Tensor Core 的支持方式、WMMAwarp matrix multiply-accumulate指令版本都不一樣新框架的某些算子會針對高版本計算能力做專門優(yōu)化。CUDA Toolkit 版本兼容范圍不同老版本的 CUDA Toolkit 不一定認識新計算能力。比如 CUDA 10.x 去編譯sm_86的代碼基本直接報錯因為那個年代還沒有 Ampere 的編譯器支持。運行時行為不同程序跑起來后驅(qū)動需要根據(jù) GPU 的計算能力加載對應內(nèi)核匹配不上就崩。所以在深度學習框架的編譯場景里計算能力直接決定了“你編出來的這個庫到底能跑在哪些 GPU 上”。2. 編譯框架時為什么要跟這個數(shù)字較勁2.1 GPU 代碼的“提前編譯”與“運行時兜底”很多框架PyTorch、TensorFlow、Caffe在安裝時都會把 CUDA 算子源碼提前編譯成二進制而不是每次運行都現(xiàn)場編譯。這個模式叫 AOTAhead Of Time提前編譯它和純 JITJust In Time運行時編譯不同。AOT 的問題是你得在編譯時告訴編譯器“請按哪一代 GPU 的架構來生成代碼”。如果你寫死了sm_86那生成的就是 RTX 30 系列認識的機器碼一旦把編譯產(chǎn)物拿到計算能力 7.5 的 RTX 2080 上跑驅(qū)動根本沒有辦法執(zhí)行這份 SASS 代碼因為它壓根不認識?，F(xiàn)代 CUDA 編譯有個兜底機制可以在生成 SASS 機器碼的同時也生成一份 PTX 中間代碼。PTX 像“跨架構的應急預案”當驅(qū)動發(fā)現(xiàn)當前 GPU 沒有匹配的 SASS 時會嘗試把 PTX 現(xiàn)場 JIT 成當前計算能力能跑的代碼。但前提是你在編譯時確實附上了 PTX很多人恰恰漏了這一手。2.2 參數(shù)設高、設低、干脆不設分別是啥后果我把三種常見情況可能引發(fā)的后果列一下設得太高例如只寫8.9在老 GPU 上直接no kernel image。編譯產(chǎn)物像是專門給新卡寫的一封信老卡連拆信的機會都沒有。設得太低例如只寫5.0幾乎所有 GPU 都能兼容但新架構的新指令全用不上。Tensor Core 加速等于擺設推理和訓練性能會明顯偏低。干脆不設置這里的風險最大。不同構建腳本的默認值五花八門有的默認只編當前機器上探測到的架構有的默認編一個非常保守的老架構。你在這臺機器上編譯時可能一切正常換到另一張卡上就跑不了。我見過一個最經(jīng)典的場景某人在一臺 RTX 4090 機器上編譯完 PyTorch看起來很順利部署到一臺 T4 推理機上直接崩。后來我把編譯日志翻出來發(fā)現(xiàn)TORCH_CUDA_ARCH_LIST是空的構建腳本默認抓到了 4090 的 8.9老卡當然顆粒無收。2.3 誰在讀取這個設置CMake、nvcc 與框架腳本的協(xié)作從底層邏輯看設置計算能力最終都會轉化成 nvcc 的-gencode參數(shù)。比如-gencode archcompute_86,codesm_86 -gencode archcompute_86,codecompute_86第一行生成針對 8.6 的 SASS第二行生成 8.6 版本的 PTX。不同框架只是把“如何拼這串參數(shù)”包裝成了變量PyTorch 用環(huán)境變量TORCH_CUDA_ARCH_LISTTensorFlow 用TF_CUDA_ARCHSCaffe 用 CMake 的CUDA_ARCH_BIN和CUDA_ARCH_PTX現(xiàn)代 CMake 工程用CMAKE_CUDA_ARCHITECTURES所以你理解了底層原理后再去看這些變量名就會覺得它們只是同一個東西的“外號”。3. 查自己 GPU 計算能力的三條路線3.1 命令行直接查最省事的方法是用 nvidia-smi 的查詢模式。普通nvidia-smi只能看到型號、顯存、驅(qū)動信息看不到計算能力但只要加兩個字段就能直接輸出nvidia-smi --query-gpuname,compute_cap --formatcsv輸出大概是這樣NVIDIA GeForce RTX 3090, 8.6這個辦法適合快速確認手頭機器到底是什么架構。3.2 官方規(guī)格表對號入座如果手頭沒有機器或者想查一批 GPU 的規(guī)格直接去 GPU 廠商的 CUDA 支持矩陣頁面。這個頁面會按計算能力從小到大列出所有支持 CUDA 的 GPU 型號。比如你想知道 Tesla T4 是多少搜一下就能看到是 7.5A100 是 8.0H100 是 9.0。查表的另一種方式是直接用 CUDA Toolkit 自帶的 deviceQuery 示例/usr/local/cuda/samples/1_Utilities/deviceQuery/deviceQuery運行后最后幾行會輸出CUDA Capability Major/Minor version number: 8.6。3.3 在代碼里動態(tài)獲取如果你要在程序里動態(tài)判斷 GPU 能力可以直接問 CUDA Runtime。C 語言寫法是用cudaGetDevicePropertiesPython 里如果用 PyTorch 則更簡潔import torch # 返回一個元組比如 (8, 6) print(torch.cuda.get_device_capability(0))這個我在寫需要自適應選 kernel 的代碼時經(jīng)常用先查計算能力再決定走哪個優(yōu)化分支避免在運行時調(diào)用不支持的算子。3.4 常用 GPU 速查參考表個人常備版我把自己平時頻繁用到的型號整理成了參考表編譯時填參數(shù)基本都用得上GPU 型號計算能力備注GTX 1080 Ti6.1老但還能跑Tesla T47.5推理常見RTX 2080 Ti7.5Turing 代表V1007.0Volta適合老 CUDA 環(huán)境A1008.0Ampere 數(shù)據(jù)中心卡RTX 30908.6Ampere 消費卡RTX 40908.9AdaH1009.0Hopper只要把這張表記在心里編譯時選架構就有了基本盤。4. PyTorch 源碼編譯TORCH_CUDA_ARCH_LIST 怎么填4.1 環(huán)境變量的填法與格式PyTorch 從源碼編譯時讀取的是環(huán)境變量TORCH_CUDA_ARCH_LIST。格式是空格分隔的列表export TORCH_CUDA_ARCH_LIST7.5 8.0 8.6編譯時它會逐個生成對應的 SASS 代碼。比如包含 7.5 和 8.6就等于給 nvcc 傳了兩組-gencode參數(shù)最終產(chǎn)出的動態(tài)庫里嵌入了兩套機器碼。這里有個容易忽略的點如果你是用官方編譯好的 PyTorch wheel 包是不需要管這個變量的官方包默認包含了常見計算能力的支持自己搞反而容易出問題。只有當你需要從源碼定制編譯或者要編譯自定義 CUDA 擴展算子時才需要真正關心它。4.2 多架構與“PTX”的組合策略在 PyTorch 的環(huán)境變量里單個架構后面可以加PTX例如export TORCH_CUDA_ARCH_LIST7.5 8.0 8.6PTX意思是為 8.6 生成 SASS 機器碼同時額外生成一份 compute_86 的 PTX 中間代碼。這個 PTX 是給“未來可能出現(xiàn)的更高計算能力 GPU”準備的。比如某臺機器是計算能力 8.9 的新卡驅(qū)動看到 8.6 的 SASS 沒法用但發(fā)現(xiàn)有 compute_86 的 PTX就會現(xiàn)場把它 JIT 成 8.9 能跑的代碼。雖然性能不一定是最優(yōu)但至少能跑起來。我建議的策略很直接把你手頭最老的目標卡、當前主力卡都寫上最后一個檔位加PTX兜底。例如開發(fā)機器是 RTX 3090但部署機器還有 RTX 2080 Tiexport TORCH_CUDA_ARCH_LIST7.5 8.6PTX這樣既照顧了老卡又保留了新卡的優(yōu)化空間。4.3 一次標準編譯實操示例一個比較典型的源碼編譯流程是這樣conda create -n build_env python3.10 conda activate build_env conda install cmake ninja export TORCH_CUDA_ARCH_LIST7.5 8.0 8.6PTX git clone --recursive https://github.com/pytorch/pytorch cd pytorch python setup.py develop編譯結束后可以用 CUDA 自帶工具檢查動態(tài)庫里到底嵌入了哪些架構cuobjdump --list-elf torch/lib/libtorch_cuda.so | grep sm_正常情況下會看到類似sm_75、sm_80、sm_86的字樣。這一步能幫你確認設置是否真的生效非常值得養(yǎng)成習慣。5. TensorFlow 源碼編譯TF_CUDA_ARCHS 與 bazel 參數(shù)5.1 官方 configure 階段和 bazel action_envTensorFlow 從源碼編譯用的是 bazel在 configure 階段會生成一份配置。其中控制計算能力的變量是TF_CUDA_ARCHS。老版本 TensorFlow 的 configure 腳本會讓你手動輸入架構例如Please specify a list of comma-separated CUDA compute capabilities you want to build with.你也可以在編譯時直接通過環(huán)境變量覆蓋TF_CUDA_ARCHS7.5 8.0 8.6 bazel build --configcuda //tensorflow/tools/pip_package:build_pip_package注意這里同樣是空格分隔。最終 bazel 會把架構列表解析成 nvcc 的 gencode 參數(shù)。5.2 老參數(shù) -arch / -gencode 其實還在用有些老教程里寫的不是TF_CUDA_ARCHS而是往 bazel 編譯選項中直接塞 nvcc 參數(shù)--copt-archsm_75或者更完整的--copt-gencodearchcompute_75,codesm_75 --copt-gencodearchcompute_75,codecompute_75這種寫法在特殊環(huán)境下依然有效因為它直接作用于 nvcc。不過維護成本高切到新架構要改一堆地方還是用TF_CUDA_ARCHS更省心。5.3 TF 場景下容易忽略的兩件事第一件TensorFlow 官方預編譯包的架構范圍是固定的。如果你是自己從源碼構建定制版別把架構范圍設得太窄否則模型在另一種 GPU 上加載時會直接報找不到內(nèi)核。第二件TensorFlow 對 PTX 的處理在各版本間有變化。新版框架總體上是按架構列表生成 SASS但對“未來架構兜底”的支持不像 PyTorch 那樣直白。如果拿不準穩(wěn)妥做法是把目標機器范圍內(nèi)所有可能用到的計算能力都列進去哪怕稍微增大一點編譯產(chǎn)物體積。我遇到過有人圖省事只寫了8.0結果 A100 上跑得好好的換到 T4 上模型加載直接掛。這就是典型的“只考慮當下機器、沒考慮部署環(huán)境”的翻車現(xiàn)場。6. Caffe 與通用 CMake 工程CMAKE_CUDA_ARCHITECTURES6.1 Caffe 的老式寫法CUDA_ARCH_BIN 和 CUDA_ARCH_PTXCaffe 的構建腳本有一套歷史悠久的變量體系CUDA_ARCH_BIN負責寫 SASS 要覆蓋哪些計算能力CUDA_ARCH_PTX負責寫 PTX 兜底檔位。編譯時這樣設cmake .. \ -DCUDA_ARCH_BIN75 86 \ -DCUDA_ARCH_PTX75這條命令的意思是生成 7.5 和 8.6 兩份 SASS 代碼同時為 7.5 準備一份 PTX。如果你的目標是只對當前 GPU 優(yōu)化不想管其他型號也可以把CUDA_ARCH_BIN設成你機器對應的計算能力數(shù)字。老項目的坑在于不同版本的 Caffe 對這個變量的解析格式不完全一致有的版本要求寫成7.5有的版本要求寫成75。遇到編譯不過時先檢查是不是數(shù)字格式的問題。6.2 現(xiàn)代 CMake 寫法CMAKE_CUDA_ARCHITECTURES如果你用的是較新的 CMake 和自定義 CUDA 工程推薦直接使用標準變量CMAKE_CUDA_ARCHITECTURES。這個變量從 CMake 3.18 開始支持分號分隔多個架構cmake .. -DCMAKE_CUDA_ARCHITECTURES75;80;86注意這里的分號和前面環(huán)境變量的空格分隔不一樣。忘了帶引號的話shell 可能把分號理解成命令分隔符會直接裂開。CMake 3.24 以上版本還支持幾個關鍵字native直接用當前機器探測到的計算能力適合只在本機部署。all支持所有已知架構靈活但編譯時間極其感人。all-major覆蓋每個大代際里的主力架構比 all 輕量不少。6.3 三個實用檔位native、多架構組合、全家族覆蓋根據(jù)實際用途我一般推薦三檔只在本機跑-DCMAKE_CUDA_ARCHITECTURESnative最省事編出來也只認當前卡。要換機器、架構范圍明確-DCMAKE_CUDA_ARCHITECTURES75;80;86速度適中兼容性兼顧。要交付給多人/多卡環(huán)境、架構不確定保守起見寫61;75;80;86;89;90。代價是編譯時間和動態(tài)庫體積都會上升換來的是“裝上就能跑”的踏實感。我記得 CMake 官方對這個變量的默認值還鬧過一陣子如果不顯式設置新版本 CMake 在啟用 CUDA 語言時會直接警告甚至報錯要求你必須顯式指定架構。這不是 CMake 矯情而是它發(fā)現(xiàn)太多人因為默認值不明確而在換機器后翻車了。7. 編譯和運行階段的高頻翻車現(xiàn)場7.1 編譯期Unsupported gpu architecture這個報錯通常長這樣Unsupported gpu architecture compute_86出現(xiàn)原因一般有兩個第一CUDA Toolkit 版本太老編譯器根本不知道sm_86/compute_86這個架構。比如老版本 CUDA 10.x 就認不出 8.6。解決辦法是升級 CUDA Toolkit或者把目標架構改成當前工具鏈認識的舊架構。第二CMake 的架構列表里有拼寫錯誤。比如手滑把86寫成了6CMake 會一臉懵地表示不認識。查一遍變量值里有沒有漏數(shù)字、多標點就好。7.2 編譯期ptxas fatal error我遇得比較多的是這種ptxas fatal : Unsupported version or architecture這個信息比上一條更底層。ptxas 是 CUDA 工具鏈里負責把 PTX 編譯成 SASS 的程序它對架構的接受范圍更敏感。最常見的場景是你設置了很高的計算能力編號例如寫了compute_90但實際 CUDA Toolkit 版本還不夠新ptxas 不認 9.0于是當場報錯。解決辦法也很直接查看工具鏈支持的架構上限把目標架構調(diào)低或者升級 CUDA Toolkit兩者必有一項。7.3 運行期no kernel image is available for execution on the device這個錯誤幾乎全是在換機器后才出現(xiàn)RuntimeError: CUDA error: no kernel image is available for execution on the device原因是當前 GPU 的計算能力在程序編譯的 SASS 列表里找不到對應項而且程序里也沒有附帶能 JIT 到當前架構的 PTX。比如用 RTX 4090 編譯出的庫只帶了 8.9 的 SASS放到 RTX 30908.6上跑8.9 的機器碼它不認又沒有 8.6 的 PTX 兜底自然直接崩。碰到這個報錯不要急著懷疑代碼邏輯先回去看編譯時架構列表范圍把目標部署機器的計算能力加進去重編一次。7.4 高頻問題速查表我把常見情況整理成表方便直接對照現(xiàn)象常見原因處理思路編譯報 Unsupported gpu architectureCUDA 版本太老 / 拼寫錯誤升級工具鏈或修正架構編號ptxas fatal: Unsupported version目標架構超出工具鏈支持范圍降檔或升級 CUDA Toolkit運行報 no kernel imageSASS 和 PTX 都與當前卡不匹配把目標架構加入編譯范圍重編CMake 報 CUDA_ARCHITECTURES 相關錯誤變量未設置或格式不對顯式指定架構列表用引號包裹換機器后部分算子極慢只編了最低兼容架構新卡特性沒啟用增加新計算能力檔位重新編譯排查這類問題時我還有個習慣先跑一行nvidia-smi --query-gpuname,compute_cap --formatcsv把當前機器的計算能力確認掉再去看編譯日志里實際生成了哪些架構兩步就能定位九成問題。8. 折騰這么多次我最后想說的幾件事第一件編譯框架前一定要先想清楚“這個庫到底要在哪些機器上跑”。只服務于本機直接native或者寫當前計算能力就行要交付部署就必須把目標機器范圍拉出來逐個核對架構列表。第二件多架構列表確實會讓編譯變慢、包體變大但這是值得的。一個包含75;80;86;89;90的庫體積可能多個幾百 MB換來的是在所有主流量級 GPU 上“裝好即跑”遠比在生產(chǎn)環(huán)境里重編框架省時間。第三件學會用cuobjdump --list-elf檢查編譯產(chǎn)物。不管是 PyTorch 的.so還是 Caffe 的動態(tài)庫編譯完以后花十秒看一下里面嵌了哪些 sm 架構就能確認設置到底有沒有生效。這一步看著不起眼卻幫我避過無數(shù)個“以為配好了其實是默認值在撐場面”的坑。最后再分享一個小經(jīng)驗團隊多人共用構建機時架構列表最好統(tǒng)一寫進構建腳本或 Dockerfile 里別依賴某個人手動 export。否則一個人改環(huán)境變量、另一個人不知道編譯出的包帶回各自的機器上互相甩鍋的場景你就等著吧。