freenode
AI & ML

PyTorch CUDA auto-detect drops Blackwell accelerated arch flags

Plain sm_120 targets alone leave NVFP4 and block-scaled MMA unavailable on new NVIDIA GPUs.

PyTorch’s automatic CUDA architecture detection has been omitting accelerated compute-capability suffixes, so Blackwell GPUs were built only for plain sm_120 and not sm_120a. That gap breaks NVFP4 and block-scaled matrix-multiply features that need the accelerated variant.

When TORCH_CUDA_ARCH_LIST is unset, auto-detect reads the GPU in the machine and historically emitted only the base architecture flag. User extensions compiled under that default therefore lacked the cubins required for the newer instruction paths.

A proposed fix adds the accelerated architecture as an extra target beside the plain one, rather than replacing it. The same dual-target pattern is already used when PyTorch builds its own CUDA sources. Because auto-detect only ever aims at the exact GPU present, including the “a” form does not raise the multi-arch compatibility issues that matter for broader binary builds.

The approach was checked on an RTX 5080: with the environment variable unset, both sm_120 and sm_120a flags are produced, and the resulting extension binary contains both cubins. Existing arch-flag coverage for older generations continues to pass.