PyTorch torch.backends 完全指南统一管理 CUDA、cuDNN、cuBLAS 与 DSL 后端的开关与调优【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch本文系统讲解 PyTorch 中torch.backends模块的职责边界与全部子后端cuda、cudnn、cusparselt、mha、mps、mkl、mkldnn、opt_einsum、python_native等的可配置项并结合 torch/backends 下的源码实现说明每个标志位的底层 setter/getter 机制、flags()上下文管理器的冻结语义以及 cuBLAS/cuSOLVER 库选择、cuFFT plan 缓存、SDPA 内核开关等高级调优手段。读完后你可以独立完成为训练/推理环境确定性地配置各后端行为按设备粒度管理 cuFFT plan 缓存覆盖 BLAS/Linalg 库选择策略以及批量禁用 Triton/CuteDSL 等 Python DSL 算子实现。一、torch.backends 是什么一个后端开关面板torch.backends负责控制 PyTorch 所支持的各种后端的运行时行为。从 docs/source/backends.md 的官方定义来看它把十几个后端开关集中到一个命名空间下torch.backends.cpu/torch.backends.cudatorch.backends.cudnn/torch.backends.cusparselt/torch.backends.mha/torch.backends.miopentorch.backends.mpsApple Metal/torch.backends.mkl/torch.backends.mkldnnoneDNNtorch.backends.nnpack/torch.backends.openmp/torch.backends.opt_einsumtorch.backends.python_native/torch.backends.xeon/torch.backends.quantized/torch.backends.kleidiai这些子模块在 torch/backends/init.py 中被统一 import 并挂到torch.backends命名空间下因此用户代码里可以直接写torch.backends.cudnn.benchmark True这样的属性赋值。这种写法之所以可行源于一个精巧的实现技巧在 torch/backends/init.py 中定义了ContextProp描述符和PropModule模块替换机制。各子模块如cudnn、opt_einsum在被 import 后会用sys.modules[__name__] XxxModule(...)把模块对象本身替换成带有描述符属性的代理模块。于是torch.backends.cudnn.enabled True实际触发的是ContextProp.__set__它再调用 C 侧的torch._C._set_cudnn_enabled把状态写入 PyTorch 的核心全局标志。值得注意的是其中的冻结机制torch/backends/init.py 定义了__allow_nonbracketed_mutation_flag与disable_global_flags()。PyTorch 的测试框架会调用它在测试套件运行期间禁止裸赋值防止某个测试忘记恢复全局状态而污染后续测试此时必须改用torch.backends.cudnn.flags(...)上下文管理器。普通用户的脚本不受影响但理解这一点有助于看懂报错信息 not allowed to set ... flags after disable_global_flags。二、torch.backends.cudacuBLAS 精度、工作区与 SDPA 内核torch.backends.cuda是覆盖面最广的子后端源码位于 torch/backends/cuda/init.py实现事实如下2.1 is_built 与 matmul 模块is_built()返回 PyTorch 是否编译了 CUDA 支持。注意它不代表CUDA 当前可用只代表如果运行在有可用 CUDA 驱动的机器上就能用实现是一行torch._C._has_cudatorch/backends/cuda/init.py。matmul属性这是一个cuBLASModule实例torch/backends/cuda/init.py文档中列出的属性在源码里逐一映射到 C setter/getter属性含义allow_tf32控制 Ampere 及以上 GPU 的矩阵乘法是否可使用 TensorFloat-32 张量核文档标注该 API 将被弃用新的精度控制推荐fp32_precisionallow_fp16_reduced_precision_reductionfp16 GEMM 是否允许低精度归约如 fp16 累加类型allow_fp16_reduced_precision_reduction_split_k只读报告 cuBLASLt 分派 fp16 GEMM 时是否允许 split-K 启发式allow_bf16_reduced_precision_reduction/..._split_kbf16 GEMM 的同款开关prefer_cublaslt_grouped_gemm支持的 grouped GEMM 是否优先走 cuBLASLtfp32_precision新的统一 fp32 精度控制源码中经_set_fp32_precision_setter(cuda, matmul, ...)生效torch/backends/cuda/init.py从源码看一个有意思的细节allow_fp16_reduced_precision_reduction等属性支持元组赋值。cuBLASModule._parse_reduction_settingtorch/backends/cuda/init.py会把True/False或(allow_reduced_precision, allow_splitk)两个布尔元组解析后一起写入 C 侧即可以同时独立控制是否允许低精度归约与是否允许 split-K。2.2 库选择覆盖preferred_blas_library / preferred_linalg_library / preferred_rocm_fa_library这三个函数允许覆盖 PyTorch 在多个 BLAS/线代库之间的默认启发式均为实验性标志源码 docstring 明确标注 experimental and subject to changepreferred_blas_library(backend)在 cuBLAS / cuBLASLt / CKROCm-only之间选择支持字符串cublas、cublaslt、ck、default及torch._C._BlasBackend枚举ROCm 构建下hipblas、hipblaslt是前两者的别名torch/backends/cuda/init.py。无参调用返回当前偏好。环境变量TORCH_BLAS_PREFER_CUBLASLT1可全局初始化为 cuBLASLt 偏好但可被后续函数调用覆盖。preferred_linalg_library(backend)在 cuSOLVER 与 MAGMA 之间选择cusolver/magma/default覆盖torch.linalg.inv、torch.linalg.cholesky、torch.linalg.lu、torch.linalg.qr、torch.linalg.eigh、torch.linalg.svd等 CUDA 线代算子的库选择TORCH_LINALG_PREFER_CUSOLVER1可全局预置。preferred_rocm_fa_library(backend)[ROCm-only]在 AOTriton 与 composable_kernel 之间切换 Flash Attention 后端TORCH_ROCM_FA_PREFER_CK1全局预置torch/backends/cuda/init.py。一个实用注意点源码 docstring 明示当偏好某个库时若该库没有实现所调用的算子仍会退回其他库偏好不是排他指定。2.3 cuFFT plan 缓存cufft_plan_cache是cuFFTPlanCacheManager单例torch/backends/cuda/init.py为每个 CUDA 设备维护一个cuFFTPlanCacheimport torch # 查询设备 0 当前 plan 数量只读 print(torch.backends.cuda.cufft_plan_cache[0].size) # 调整设备 0 缓存容量 torch.backends.cuda.cufft_plan_cache[0].max_size 32 # 清空当前设备的 plan 缓存 torch.backends.cuda.cufft_plan_cache.clear()从源码结构看torch/backends/cuda/init.pycufft_plan_cache支持直接当单个缓存用torch.backends.cuda.cufft_plan_cache.max_size会透明地解析到当前设备按索引访问cufft_plan_cache[i]时索引越界会抛出带device_count()信息的RuntimeError。2.4 工作区大小cublas_workspace_size / cublaslt_workspace_size / blas_workspace_size这三个函数都是无参查询、传参设置的双模式 APIcublas_workspace_size(sizeNone)查询/设置 cuBLAS 工作区字节数设置后优先于CUBLAS_WORKSPACE_CONFIG环境变量cublaslt_workspace_size(sizeNone)同理优先于CUBLASLT_WORKSPACE_SIZE环境变量blas_workspace_size(sizeNone, backendNone)便捷封装backend缺省时按当前preferred_blas_library解析解析到 CK 会抛RuntimeErrorCK 无工作区概念。源码 docstring 特别强调两点变更是惰性生效的——只有之后新建的 handle 才会拿到新工作区且开源 CUDA 构建默认开启TORCH_CUBLASLT_UNIFIED_WORKSPACEcuBLASLt 工作区会复用 cuBLAS 的同一块分配并以其大小为上限因此单独调大 cuBLASLt 工作区不一定增加显存占用torch/backends/cuda/init.py。2.5 SDPA 内核开关与 sdp_kernel 的弃用torch.backends.cuda还暴露了 scaled dot product attention 的各内核开关全部为 beta 标志flash_sdp_enabled()/enable_flash_sdp(enabled)mem_efficient_sdp_enabled()/enable_mem_efficient_sdp(enabled)math_sdp_enabled()/enable_math_sdp(enabled)cudnn_sdp_enabled()/enable_cudnn_sdp(enabled)fp16_bf16_reduction_math_sdp_allowed()/allow_fp16_bf16_reduction_math_sdp(enabled)可用性与可执行性检查is_flash_attention_available()、can_use_flash_attention(params, debugFalse)、can_use_efficient_attention(params, debugFalse)、can_use_cudnn_attention(params, debugFalse)、is_ck_sdpa_available()参数容器SDPAParams由torch._C._SDPAParams重导出__module__被改写为torch.backends.cudatorch/backends/cuda/init.pycan_use_*系列均要求传入该结构。需要特别注意的演进方向torch.backends.cuda.sdp_kernel(enable_flash, enable_math, enable_mem_efficient, enable_cudnn)这个上下文管理器已被标记弃用deprecatedFutureWarning源码中直接转发到新的torch.nn.attention.sdpa_kernel官方提示应迁移到新签名torch/backends/cuda/init.py。新写法示例from torch.nn.attention import sdpa_kernel, SDPBackend with sdpa_kernel([SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION]): out nn.functional.scaled_dot_product_attention(q, k, v)can_use_flash_attention(params, debugTrue)中的debug参数会在返回 False 时通过logging.warn输出具体原因是排查为什么我的 SDPA 没走 Flash 路径的实用入口。相关行为可在测试 test/test_native_mha.py 中找到大量覆盖用例。三、torch.backends.cudnn卷积后端最常用的一组开关torch.backends.cudnn的实现在 torch/backends/cudnn/init.py其中CudnnModuletorch/backends/cudnn/init.py暴露的属性与文档一一对应且比文档多暴露了两项源码级能力开关类型作用enabledbool是否启用 cuDNN/MIOpenROCm 构建下此模块同时控制 MIOpen源码注释明确 to globally disable CuDNN/MIOpenallow_tf32boolAmpere GPU 上 cuDNN 卷积能否使用 TF32 张量核同样标注将被弃用deterministicboolTrue 时 cuDNN 仅使用确定性卷积算法与torch.use_deterministic_algorithms联动benchmarkboolTrue 时 cuDNN 基准测试多个卷积算法并挑选最快者benchmark_limitintbenchmarkTrue时最多尝试的算法数设为 0 表示尝试所有可用算法仅影响经 cuDNN v8 API 分派的卷积depthwise_kernelstr源码额外暴露auto等控制 depthwise 卷积内核选择fp32_precision/convstr新的 fp32 精度统一控制conv是绑定到(cuda, conv)的_FP32Precision对象典型用法import torch torch.backends.cudnn.benchmark True # 形状固定时显著加速卷积 torch.backends.cudnn.deterministic False # 或者用上下文管理器临时覆盖退出后自动恢复 with torch.backends.cudnn.flags(benchmarkTrue, benchmark_limit10, enabledTrue, deterministicFalse): train_one_step(model, data)cudnn.flags(enabledFalse, benchmarkFalse, benchmark_limit10, deterministicFalse, allow_tf32True, fp32_precisionnone, depthwise_kernelauto)上下文管理器会先调用set_flags保存原值退出时用set_flags(*orig_flags)恢复torch/backends/cudnn/init.py。此外文档列出的两个函数torch.backends.cudnn.version()返回 cuDNN 版本号如91100不可用时返回None。源码中_init()还会做版本兼容性检查主版本不同直接不兼容cuDNN 7 起次版本向后兼容运行时次版本 ≥ 编译次版本即可MIOpen 则做严格相等检查不兼容时可设PYTORCH_SKIP_CUDNN_COMPATIBILITY_CHECK1跳过否则抛出含LD_LIBRARY_PATH诊断信息的RuntimeError。cuDNN ≥ 9.11 且设备 SM 7.5 时也会直接报错torch/backends/cudnn/init.py。torch.backends.cudnn.is_available()等价于torch._C._has_cudnn即编译时是否带 cuDNN/MIOpen。torch.backends.cudnn.rnnRNN 专用子模块对应 cuDNN RNN 的确定性与算法选择参见文档 docs/source/cudnn_rnn_determinism.md 与 docs/source/cudnn_persistent_rnn.rst。四、torch.backends.cusparselt、mha 与 miopen4.1 torch.backends.cusparselt面向 cuSPARSELt 稀疏矩阵乘库torch/backends/cusparselt/init.py 提供三个查询函数version()返回 cuSPARSELt 版本不可用时None内部同样带_init()初始化逻辑is_available()当前是否可用get_max_alg_id()返回 cuSPARSELt 算法 ID 上限用于约束稀疏 GEMM 的算法选择参数。4.2 torch.backends.mhaMulti-Head Attention 的 fastpath 开关源码在 torch/backends/mha/init.pyimport torch print(torch.backends.mha.get_fastpath_enabled()) # 当前是否启用 fastpath torch.backends.mha.set_fastpath_enabled(False) # 关闭后回退常规 MHA 路径行为验证可参考 test/test_native_mha.py。4.3 torch.backends.miopenROCm 平台的卷积后端暴露immediate布尔属性True 时 MIOpen 使用 Immediate Mode不建立/查询数据库每次实时选算法。该模块同样有set_flags/flags上下文管理器torch/backends/miopen/init.py。五、torch.backends.mps / mkl / mkldnn / nnpack / openmp非 CUDA 平台后端5.1 torch.backends.mpsApple Metaltorch/backends/mps/init.py 提供设备信息查询is_available()Metal 当前是否可用torch._C._mps_is_available()is_built()PyTorch 是否编译了 MPS 支持get_core_count()MPS 加速器核心数get_name()MPS 设备名称is_macos13_or_newer(minor0)/is_macos_or_newer(major, minor)系统版本判断用于在 macOS 上按特性门控 MPS 相关逻辑。5.2 torch.backends.mkl 与 torch.backends.mkldnn两者是可用性 详细日志的组合torch.backends.mkl.is_available()MKL 是否可用torch.backends.mkl.verbose(level)上下文管理器临时打印 MKL 内核选择信息with torch.backends.mkl.verbose(): y x wtorch.backends.mkldnn.is_available()oneDNN旧称 MKL-DNN是否可用torch.backends.mkldnn.verbose(level)同款日志上下文管理器。源码中verbose是带__enter__/__exit__的类torch/backends/mkl/init.py、torch/backends/mkldnn/init.py进入时置位全局日志级别、退出时恢复。行为测试见 test/test_mkldnn.py 与 test/test_mkldnn_fusion.py。5.3 torch.backends.nnpack移动端/嵌入式 CPU 卷积后端。除is_available()外还提供set_flags(_enabled)与flags(enabledFalse)上下文管理器torch/backends/nnpack/init.py写法与 cudnn 一致。5.4 torch.backends.openmp只有一个is_available()用于判断 PyTorch 是否以 OpenMP 构建对照 test/test_openmp.py。5.5 torch.backends.quantized 与 torch.backends.xeon补充模块文档末尾以py:module占位的三个模块在源码中也有实际能力torch.backends.quantized暴露engine当前量化工具链可读写与supported_engines只读列表由torch._C._supported_qengines()填充QuantizedEngine代理模块见 torch/backends/quantized/init.pytorch.backends.xeon.run_cpu.create_args为 Intel Xeon 专用基准进程torch/backends/xeon/run_cpu.py构造 argparse 参数含内存分配器、多实例KMP/IOMP参数组_add_memory_allocator_params、_add_multi_instance_params、_add_kmp_iomp_paramstorch/backends/xeon/run_cpu.pytorch.backends.kleidiai.is_available()Arm KleidiAI 加速库可用性查询。六、torch.backends.opt_einsumeinsum 收缩路径优化文档对它的说明是enabled默认True控制torch.einsum是否借助 opt_einsum 计算最优收缩路径strategy指定路径搜索策略auto默认也支持greedy与optimal后者对输入张量数量是阶乘级开销。opt_einsum 不可用时torch.einsum回退为从左到右收缩。torch/backends/opt_einsum/init.py 的实现补充了几个文档未展开的细节is_available()带lru_cache本质是包是否 import 成功docstring 给出安装方式pip install torch[opt-einsum]或pip install opt-einsumget_opt_einsum()直接返回包对象或None防御性 setter在未安装 opt_einsum 时把enabled设为True或设置strategy都会抛ValueError并解释回退行为torch/backends/opt_einsum/init.py——这避免了设了开关但毫无效果的静默陷阱strategy的合法值在源码中被硬校验为[auto, greedy, optimal]且strategy描述符仅在is_available()为真时才创建即未安装时该属性根本不存在。七、torch.backends.python_nativeDSL 算子的精细化总闸这是文档中篇幅最大的部分torch.backends.python_native让用户控制由torch._native中 DSL领域特定语言实现的算子可按 DSL 整体开关也可按单个算子/分派键开关且全部支持上下文管理器。实现在 torch/backends/python_native/init.py底层依赖torch._native.dsl_registry.dsl_registry与torch._native.registry的过滤状态_dsl_names、_op_symbols、_dispatch_keys三个集合。7.1 模块级函数与属性get_dsl_operations(dsl_name)查询某个 DSL 注册的算子列表disable_operations(*ops)/enable_operations(*ops)跨所有 DSL 禁用/恢复指定算子按名字如scaled_mmdisable_dispatch_keys(*keys)/enable_dispatch_keys(*keys)按 dispatch key 粒度开关operations_disabled查询当前被禁用的算子集合available_dsls运行时依赖已满足的 DSL 名字列表all_dsls的子集all_dsls所有已注册 DSL 的名字列表无论依赖是否满足。7.2 DSL Controller以 triton、cutedsl 为例每个已注册 DSL 会自动生成一个DSLController提供文档中的四个属性与三个方法属性/方法说明nameDSL 名称available运行时依赖是否满足查registry.is_dsl_availableenabled读写属性False禁用该 DSL 全部算子True重新启用测试框架冻结状态下设置会抛RuntimeErrorversionDSL 运行时版本不可用时为Nonedisable()/enable()整体禁用/启用调用deregister_op_overrides/reenable_op_overridesdisabled()上下文管理器退出时按进入前的状态恢复文档给出的用法示例可直接运行import torch.backends.python_native as pn # 查询可用 DSL print(pn.available_dsls) # [triton, cutedsl] # 禁用全部 Triton 算子 pn.triton.enabled False # 临时禁用 CuteDSL 算子 with pn.cutedsl.disabled(): result model(input) # CuteDSL ops disabled # 跨所有 DSL 禁用特定算子 pn.disable_operations(scaled_mm, _flash_attention_forward) # 查询某个 DSL 的算子 triton_ops pn.get_dsl_operations(triton)从源码结构看DSLController.enabled的实现是该 DSL 名是否不在过滤集合中self._dsl_name not in filter_state._dsl_namesdisable()则通过dsl_module.deregister_op_overrides()把该 DSL 的算子覆盖从注册表中摘除enable()用reenable_op_overrides(enable_dsl_names...)精准恢复。因此 DSL 集合是可扩展的——新增 DSL 无需改动torch.backends.python_native本身控制器通过注册表动态解析registry.get_dsl_module。对应测试位于 test/python_native 目录。八、实践小结一个典型 CUDA 环境的后端配置示例综合以上各节以下示例把本文涉及的torch.backends能力串起来适用于 CUDA 构建环境cufft_plan_cache与 SDPA 开关部分要求 CUDA 可用preferred_linalg_library等在 ROCm 构建下另有ck/hipblaslt选项import torch # 1. cuDNN 卷积调优形状固定的推理/训练 torch.backends.cudnn.benchmark True torch.backends.cudnn.benchmark_limit 10 # 2. BLAS 库选择与工作区 print(torch.backends.cuda.preferred_blas_library()) # 查询当前 torch.backends.cuda.preferred_blas_library(cublaslt) # 覆盖偏好 torch.backends.cuda.cublas_workspace_size(32 * 1024 * 1024) # 32 MiB # 3. fp16 GEMM允许低精度归约但限制 split-K元组赋值 torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction (True, False) # 4. cuFFT plan 缓存限制设备 0 的容量 torch.backends.cuda.cufft_plan_cache[0].max_size 16 # 5. SDPA 内核选择推荐新 APItorch.backends.cuda.sdp_kernel 已弃用 from torch.nn.attention import sdpa_kernel, SDPBackend with sdpa_kernel([SDPBackend.FLASH_ATTENTION]): out torch.nn.functional.scaled_dot_product_attention(q, k, v) # 6. 临时切换 einsum 策略 with torch.backends.opt_einsum.flags(strategygreedy): y torch.einsum(ijk,jkl-il, a, b)所有配置项的权威描述见 docs/source/backends.md实现分别在 torch/backends 各子模块中行为回归主要由 test/test_native_mha.py、test/test_mkldnn.py、test/test_openmp.py、test/backends/xeon 与 test/python_native 等测试守护。理解torch.backends的核心在于两点其一属性赋值背后是ContextProp对 C 全局状态的读写flags()上下文管理器提供了带自动恢复的安全写法其二随着fp32_precision统一精度 API、torch.nn.attention.sdpa_kernel等新机制落地部分旧开关如matmul.allow_tf32、cuda.sdp_kernel已进入弃用轨道新项目应优先采用新接口。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
PyTorch torch.backends 完全指南：统一管理 CUDA、cuDNN、cuBLAS 与 DSL 后端的开关与调优
PyTorch torch.backends 完全指南统一管理 CUDA、cuDNN、cuBLAS 与 DSL 后端的开关与调优【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch本文系统讲解 PyTorch 中torch.backends模块的职责边界与全部子后端cuda、cudnn、cusparselt、mha、mps、mkl、mkldnn、opt_einsum、python_native等的可配置项并结合 torch/backends 下的源码实现说明每个标志位的底层 setter/getter 机制、flags()上下文管理器的冻结语义以及 cuBLAS/cuSOLVER 库选择、cuFFT plan 缓存、SDPA 内核开关等高级调优手段。读完后你可以独立完成为训练/推理环境确定性地配置各后端行为按设备粒度管理 cuFFT plan 缓存覆盖 BLAS/Linalg 库选择策略以及批量禁用 Triton/CuteDSL 等 Python DSL 算子实现。一、torch.backends 是什么一个后端开关面板torch.backends负责控制 PyTorch 所支持的各种后端的运行时行为。从 docs/source/backends.md 的官方定义来看它把十几个后端开关集中到一个命名空间下torch.backends.cpu/torch.backends.cudatorch.backends.cudnn/torch.backends.cusparselt/torch.backends.mha/torch.backends.miopentorch.backends.mpsApple Metal/torch.backends.mkl/torch.backends.mkldnnoneDNNtorch.backends.nnpack/torch.backends.openmp/torch.backends.opt_einsumtorch.backends.python_native/torch.backends.xeon/torch.backends.quantized/torch.backends.kleidiai这些子模块在 torch/backends/init.py 中被统一 import 并挂到torch.backends命名空间下因此用户代码里可以直接写torch.backends.cudnn.benchmark True这样的属性赋值。这种写法之所以可行源于一个精巧的实现技巧在 torch/backends/init.py 中定义了ContextProp描述符和PropModule模块替换机制。各子模块如cudnn、opt_einsum在被 import 后会用sys.modules[__name__] XxxModule(...)把模块对象本身替换成带有描述符属性的代理模块。于是torch.backends.cudnn.enabled True实际触发的是ContextProp.__set__它再调用 C 侧的torch._C._set_cudnn_enabled把状态写入 PyTorch 的核心全局标志。值得注意的是其中的冻结机制torch/backends/init.py 定义了__allow_nonbracketed_mutation_flag与disable_global_flags()。PyTorch 的测试框架会调用它在测试套件运行期间禁止裸赋值防止某个测试忘记恢复全局状态而污染后续测试此时必须改用torch.backends.cudnn.flags(...)上下文管理器。普通用户的脚本不受影响但理解这一点有助于看懂报错信息 not allowed to set ... flags after disable_global_flags。二、torch.backends.cudacuBLAS 精度、工作区与 SDPA 内核torch.backends.cuda是覆盖面最广的子后端源码位于 torch/backends/cuda/init.py实现事实如下2.1 is_built 与 matmul 模块is_built()返回 PyTorch 是否编译了 CUDA 支持。注意它不代表CUDA 当前可用只代表如果运行在有可用 CUDA 驱动的机器上就能用实现是一行torch._C._has_cudatorch/backends/cuda/init.py。matmul属性这是一个cuBLASModule实例torch/backends/cuda/init.py文档中列出的属性在源码里逐一映射到 C setter/getter属性含义allow_tf32控制 Ampere 及以上 GPU 的矩阵乘法是否可使用 TensorFloat-32 张量核文档标注该 API 将被弃用新的精度控制推荐fp32_precisionallow_fp16_reduced_precision_reductionfp16 GEMM 是否允许低精度归约如 fp16 累加类型allow_fp16_reduced_precision_reduction_split_k只读报告 cuBLASLt 分派 fp16 GEMM 时是否允许 split-K 启发式allow_bf16_reduced_precision_reduction/..._split_kbf16 GEMM 的同款开关prefer_cublaslt_grouped_gemm支持的 grouped GEMM 是否优先走 cuBLASLtfp32_precision新的统一 fp32 精度控制源码中经_set_fp32_precision_setter(cuda, matmul, ...)生效torch/backends/cuda/init.py从源码看一个有意思的细节allow_fp16_reduced_precision_reduction等属性支持元组赋值。cuBLASModule._parse_reduction_settingtorch/backends/cuda/init.py会把True/False或(allow_reduced_precision, allow_splitk)两个布尔元组解析后一起写入 C 侧即可以同时独立控制是否允许低精度归约与是否允许 split-K。2.2 库选择覆盖preferred_blas_library / preferred_linalg_library / preferred_rocm_fa_library这三个函数允许覆盖 PyTorch 在多个 BLAS/线代库之间的默认启发式均为实验性标志源码 docstring 明确标注 experimental and subject to changepreferred_blas_library(backend)在 cuBLAS / cuBLASLt / CKROCm-only之间选择支持字符串cublas、cublaslt、ck、default及torch._C._BlasBackend枚举ROCm 构建下hipblas、hipblaslt是前两者的别名torch/backends/cuda/init.py。无参调用返回当前偏好。环境变量TORCH_BLAS_PREFER_CUBLASLT1可全局初始化为 cuBLASLt 偏好但可被后续函数调用覆盖。preferred_linalg_library(backend)在 cuSOLVER 与 MAGMA 之间选择cusolver/magma/default覆盖torch.linalg.inv、torch.linalg.cholesky、torch.linalg.lu、torch.linalg.qr、torch.linalg.eigh、torch.linalg.svd等 CUDA 线代算子的库选择TORCH_LINALG_PREFER_CUSOLVER1可全局预置。preferred_rocm_fa_library(backend)[ROCm-only]在 AOTriton 与 composable_kernel 之间切换 Flash Attention 后端TORCH_ROCM_FA_PREFER_CK1全局预置torch/backends/cuda/init.py。一个实用注意点源码 docstring 明示当偏好某个库时若该库没有实现所调用的算子仍会退回其他库偏好不是排他指定。2.3 cuFFT plan 缓存cufft_plan_cache是cuFFTPlanCacheManager单例torch/backends/cuda/init.py为每个 CUDA 设备维护一个cuFFTPlanCacheimport torch # 查询设备 0 当前 plan 数量只读 print(torch.backends.cuda.cufft_plan_cache[0].size) # 调整设备 0 缓存容量 torch.backends.cuda.cufft_plan_cache[0].max_size 32 # 清空当前设备的 plan 缓存 torch.backends.cuda.cufft_plan_cache.clear()从源码结构看torch/backends/cuda/init.pycufft_plan_cache支持直接当单个缓存用torch.backends.cuda.cufft_plan_cache.max_size会透明地解析到当前设备按索引访问cufft_plan_cache[i]时索引越界会抛出带device_count()信息的RuntimeError。2.4 工作区大小cublas_workspace_size / cublaslt_workspace_size / blas_workspace_size这三个函数都是无参查询、传参设置的双模式 APIcublas_workspace_size(sizeNone)查询/设置 cuBLAS 工作区字节数设置后优先于CUBLAS_WORKSPACE_CONFIG环境变量cublaslt_workspace_size(sizeNone)同理优先于CUBLASLT_WORKSPACE_SIZE环境变量blas_workspace_size(sizeNone, backendNone)便捷封装backend缺省时按当前preferred_blas_library解析解析到 CK 会抛RuntimeErrorCK 无工作区概念。源码 docstring 特别强调两点变更是惰性生效的——只有之后新建的 handle 才会拿到新工作区且开源 CUDA 构建默认开启TORCH_CUBLASLT_UNIFIED_WORKSPACEcuBLASLt 工作区会复用 cuBLAS 的同一块分配并以其大小为上限因此单独调大 cuBLASLt 工作区不一定增加显存占用torch/backends/cuda/init.py。2.5 SDPA 内核开关与 sdp_kernel 的弃用torch.backends.cuda还暴露了 scaled dot product attention 的各内核开关全部为 beta 标志flash_sdp_enabled()/enable_flash_sdp(enabled)mem_efficient_sdp_enabled()/enable_mem_efficient_sdp(enabled)math_sdp_enabled()/enable_math_sdp(enabled)cudnn_sdp_enabled()/enable_cudnn_sdp(enabled)fp16_bf16_reduction_math_sdp_allowed()/allow_fp16_bf16_reduction_math_sdp(enabled)可用性与可执行性检查is_flash_attention_available()、can_use_flash_attention(params, debugFalse)、can_use_efficient_attention(params, debugFalse)、can_use_cudnn_attention(params, debugFalse)、is_ck_sdpa_available()参数容器SDPAParams由torch._C._SDPAParams重导出__module__被改写为torch.backends.cudatorch/backends/cuda/init.pycan_use_*系列均要求传入该结构。需要特别注意的演进方向torch.backends.cuda.sdp_kernel(enable_flash, enable_math, enable_mem_efficient, enable_cudnn)这个上下文管理器已被标记弃用deprecatedFutureWarning源码中直接转发到新的torch.nn.attention.sdpa_kernel官方提示应迁移到新签名torch/backends/cuda/init.py。新写法示例from torch.nn.attention import sdpa_kernel, SDPBackend with sdpa_kernel([SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION]): out nn.functional.scaled_dot_product_attention(q, k, v)can_use_flash_attention(params, debugTrue)中的debug参数会在返回 False 时通过logging.warn输出具体原因是排查为什么我的 SDPA 没走 Flash 路径的实用入口。相关行为可在测试 test/test_native_mha.py 中找到大量覆盖用例。三、torch.backends.cudnn卷积后端最常用的一组开关torch.backends.cudnn的实现在 torch/backends/cudnn/init.py其中CudnnModuletorch/backends/cudnn/init.py暴露的属性与文档一一对应且比文档多暴露了两项源码级能力开关类型作用enabledbool是否启用 cuDNN/MIOpenROCm 构建下此模块同时控制 MIOpen源码注释明确 to globally disable CuDNN/MIOpenallow_tf32boolAmpere GPU 上 cuDNN 卷积能否使用 TF32 张量核同样标注将被弃用deterministicboolTrue 时 cuDNN 仅使用确定性卷积算法与torch.use_deterministic_algorithms联动benchmarkboolTrue 时 cuDNN 基准测试多个卷积算法并挑选最快者benchmark_limitintbenchmarkTrue时最多尝试的算法数设为 0 表示尝试所有可用算法仅影响经 cuDNN v8 API 分派的卷积depthwise_kernelstr源码额外暴露auto等控制 depthwise 卷积内核选择fp32_precision/convstr新的 fp32 精度统一控制conv是绑定到(cuda, conv)的_FP32Precision对象典型用法import torch torch.backends.cudnn.benchmark True # 形状固定时显著加速卷积 torch.backends.cudnn.deterministic False # 或者用上下文管理器临时覆盖退出后自动恢复 with torch.backends.cudnn.flags(benchmarkTrue, benchmark_limit10, enabledTrue, deterministicFalse): train_one_step(model, data)cudnn.flags(enabledFalse, benchmarkFalse, benchmark_limit10, deterministicFalse, allow_tf32True, fp32_precisionnone, depthwise_kernelauto)上下文管理器会先调用set_flags保存原值退出时用set_flags(*orig_flags)恢复torch/backends/cudnn/init.py。此外文档列出的两个函数torch.backends.cudnn.version()返回 cuDNN 版本号如91100不可用时返回None。源码中_init()还会做版本兼容性检查主版本不同直接不兼容cuDNN 7 起次版本向后兼容运行时次版本 ≥ 编译次版本即可MIOpen 则做严格相等检查不兼容时可设PYTORCH_SKIP_CUDNN_COMPATIBILITY_CHECK1跳过否则抛出含LD_LIBRARY_PATH诊断信息的RuntimeError。cuDNN ≥ 9.11 且设备 SM 7.5 时也会直接报错torch/backends/cudnn/init.py。torch.backends.cudnn.is_available()等价于torch._C._has_cudnn即编译时是否带 cuDNN/MIOpen。torch.backends.cudnn.rnnRNN 专用子模块对应 cuDNN RNN 的确定性与算法选择参见文档 docs/source/cudnn_rnn_determinism.md 与 docs/source/cudnn_persistent_rnn.rst。四、torch.backends.cusparselt、mha 与 miopen4.1 torch.backends.cusparselt面向 cuSPARSELt 稀疏矩阵乘库torch/backends/cusparselt/init.py 提供三个查询函数version()返回 cuSPARSELt 版本不可用时None内部同样带_init()初始化逻辑is_available()当前是否可用get_max_alg_id()返回 cuSPARSELt 算法 ID 上限用于约束稀疏 GEMM 的算法选择参数。4.2 torch.backends.mhaMulti-Head Attention 的 fastpath 开关源码在 torch/backends/mha/init.pyimport torch print(torch.backends.mha.get_fastpath_enabled()) # 当前是否启用 fastpath torch.backends.mha.set_fastpath_enabled(False) # 关闭后回退常规 MHA 路径行为验证可参考 test/test_native_mha.py。4.3 torch.backends.miopenROCm 平台的卷积后端暴露immediate布尔属性True 时 MIOpen 使用 Immediate Mode不建立/查询数据库每次实时选算法。该模块同样有set_flags/flags上下文管理器torch/backends/miopen/init.py。五、torch.backends.mps / mkl / mkldnn / nnpack / openmp非 CUDA 平台后端5.1 torch.backends.mpsApple Metaltorch/backends/mps/init.py 提供设备信息查询is_available()Metal 当前是否可用torch._C._mps_is_available()is_built()PyTorch 是否编译了 MPS 支持get_core_count()MPS 加速器核心数get_name()MPS 设备名称is_macos13_or_newer(minor0)/is_macos_or_newer(major, minor)系统版本判断用于在 macOS 上按特性门控 MPS 相关逻辑。5.2 torch.backends.mkl 与 torch.backends.mkldnn两者是可用性 详细日志的组合torch.backends.mkl.is_available()MKL 是否可用torch.backends.mkl.verbose(level)上下文管理器临时打印 MKL 内核选择信息with torch.backends.mkl.verbose(): y x wtorch.backends.mkldnn.is_available()oneDNN旧称 MKL-DNN是否可用torch.backends.mkldnn.verbose(level)同款日志上下文管理器。源码中verbose是带__enter__/__exit__的类torch/backends/mkl/init.py、torch/backends/mkldnn/init.py进入时置位全局日志级别、退出时恢复。行为测试见 test/test_mkldnn.py 与 test/test_mkldnn_fusion.py。5.3 torch.backends.nnpack移动端/嵌入式 CPU 卷积后端。除is_available()外还提供set_flags(_enabled)与flags(enabledFalse)上下文管理器torch/backends/nnpack/init.py写法与 cudnn 一致。5.4 torch.backends.openmp只有一个is_available()用于判断 PyTorch 是否以 OpenMP 构建对照 test/test_openmp.py。5.5 torch.backends.quantized 与 torch.backends.xeon补充模块文档末尾以py:module占位的三个模块在源码中也有实际能力torch.backends.quantized暴露engine当前量化工具链可读写与supported_engines只读列表由torch._C._supported_qengines()填充QuantizedEngine代理模块见 torch/backends/quantized/init.pytorch.backends.xeon.run_cpu.create_args为 Intel Xeon 专用基准进程torch/backends/xeon/run_cpu.py构造 argparse 参数含内存分配器、多实例KMP/IOMP参数组_add_memory_allocator_params、_add_multi_instance_params、_add_kmp_iomp_paramstorch/backends/xeon/run_cpu.pytorch.backends.kleidiai.is_available()Arm KleidiAI 加速库可用性查询。六、torch.backends.opt_einsumeinsum 收缩路径优化文档对它的说明是enabled默认True控制torch.einsum是否借助 opt_einsum 计算最优收缩路径strategy指定路径搜索策略auto默认也支持greedy与optimal后者对输入张量数量是阶乘级开销。opt_einsum 不可用时torch.einsum回退为从左到右收缩。torch/backends/opt_einsum/init.py 的实现补充了几个文档未展开的细节is_available()带lru_cache本质是包是否 import 成功docstring 给出安装方式pip install torch[opt-einsum]或pip install opt-einsumget_opt_einsum()直接返回包对象或None防御性 setter在未安装 opt_einsum 时把enabled设为True或设置strategy都会抛ValueError并解释回退行为torch/backends/opt_einsum/init.py——这避免了设了开关但毫无效果的静默陷阱strategy的合法值在源码中被硬校验为[auto, greedy, optimal]且strategy描述符仅在is_available()为真时才创建即未安装时该属性根本不存在。七、torch.backends.python_nativeDSL 算子的精细化总闸这是文档中篇幅最大的部分torch.backends.python_native让用户控制由torch._native中 DSL领域特定语言实现的算子可按 DSL 整体开关也可按单个算子/分派键开关且全部支持上下文管理器。实现在 torch/backends/python_native/init.py底层依赖torch._native.dsl_registry.dsl_registry与torch._native.registry的过滤状态_dsl_names、_op_symbols、_dispatch_keys三个集合。7.1 模块级函数与属性get_dsl_operations(dsl_name)查询某个 DSL 注册的算子列表disable_operations(*ops)/enable_operations(*ops)跨所有 DSL 禁用/恢复指定算子按名字如scaled_mmdisable_dispatch_keys(*keys)/enable_dispatch_keys(*keys)按 dispatch key 粒度开关operations_disabled查询当前被禁用的算子集合available_dsls运行时依赖已满足的 DSL 名字列表all_dsls的子集all_dsls所有已注册 DSL 的名字列表无论依赖是否满足。7.2 DSL Controller以 triton、cutedsl 为例每个已注册 DSL 会自动生成一个DSLController提供文档中的四个属性与三个方法属性/方法说明nameDSL 名称available运行时依赖是否满足查registry.is_dsl_availableenabled读写属性False禁用该 DSL 全部算子True重新启用测试框架冻结状态下设置会抛RuntimeErrorversionDSL 运行时版本不可用时为Nonedisable()/enable()整体禁用/启用调用deregister_op_overrides/reenable_op_overridesdisabled()上下文管理器退出时按进入前的状态恢复文档给出的用法示例可直接运行import torch.backends.python_native as pn # 查询可用 DSL print(pn.available_dsls) # [triton, cutedsl] # 禁用全部 Triton 算子 pn.triton.enabled False # 临时禁用 CuteDSL 算子 with pn.cutedsl.disabled(): result model(input) # CuteDSL ops disabled # 跨所有 DSL 禁用特定算子 pn.disable_operations(scaled_mm, _flash_attention_forward) # 查询某个 DSL 的算子 triton_ops pn.get_dsl_operations(triton)从源码结构看DSLController.enabled的实现是该 DSL 名是否不在过滤集合中self._dsl_name not in filter_state._dsl_namesdisable()则通过dsl_module.deregister_op_overrides()把该 DSL 的算子覆盖从注册表中摘除enable()用reenable_op_overrides(enable_dsl_names...)精准恢复。因此 DSL 集合是可扩展的——新增 DSL 无需改动torch.backends.python_native本身控制器通过注册表动态解析registry.get_dsl_module。对应测试位于 test/python_native 目录。八、实践小结一个典型 CUDA 环境的后端配置示例综合以上各节以下示例把本文涉及的torch.backends能力串起来适用于 CUDA 构建环境cufft_plan_cache与 SDPA 开关部分要求 CUDA 可用preferred_linalg_library等在 ROCm 构建下另有ck/hipblaslt选项import torch # 1. cuDNN 卷积调优形状固定的推理/训练 torch.backends.cudnn.benchmark True torch.backends.cudnn.benchmark_limit 10 # 2. BLAS 库选择与工作区 print(torch.backends.cuda.preferred_blas_library()) # 查询当前 torch.backends.cuda.preferred_blas_library(cublaslt) # 覆盖偏好 torch.backends.cuda.cublas_workspace_size(32 * 1024 * 1024) # 32 MiB # 3. fp16 GEMM允许低精度归约但限制 split-K元组赋值 torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction (True, False) # 4. cuFFT plan 缓存限制设备 0 的容量 torch.backends.cuda.cufft_plan_cache[0].max_size 16 # 5. SDPA 内核选择推荐新 APItorch.backends.cuda.sdp_kernel 已弃用 from torch.nn.attention import sdpa_kernel, SDPBackend with sdpa_kernel([SDPBackend.FLASH_ATTENTION]): out torch.nn.functional.scaled_dot_product_attention(q, k, v) # 6. 临时切换 einsum 策略 with torch.backends.opt_einsum.flags(strategygreedy): y torch.einsum(ijk,jkl-il, a, b)所有配置项的权威描述见 docs/source/backends.md实现分别在 torch/backends 各子模块中行为回归主要由 test/test_native_mha.py、test/test_mkldnn.py、test/test_openmp.py、test/backends/xeon 与 test/python_native 等测试守护。理解torch.backends的核心在于两点其一属性赋值背后是ContextProp对 C 全局状态的读写flags()上下文管理器提供了带自动恢复的安全写法其二随着fp32_precision统一精度 API、torch.nn.attention.sdpa_kernel等新机制落地部分旧开关如matmul.allow_tf32、cuda.sdp_kernel已进入弃用轨道新项目应优先采用新接口。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考