生物分子结构预测已经进入蛋白质组规模，目标不再是一次算准一个序列，而是让整份工作清单高效地流过流水线。NVIDIA BioNeMo Inference Runtime（BioIR）面向这一场景：在 NVIDIA GPU 上加速受支持的结构预测模型，同时保留开发者熟悉的 PyTorch 工作流。它使用优化过的内核，并在适用时启用 CUDA Graphs 加速模型执行；面对大批独立输入，还可以用 Ray 在单节点每块 GPU 上运行一套完整模型副本，提高整体吞吐。
这套运行时已有蛋白质组级别的实际应用。在 AlphaFold Database（AFDB）的近期扩容中，BioIR 加速了跨 4777 个蛋白质组的蛋白复合物结构生成，候选复合物约 3100 万个，最终有 181 万个作为高置信预测发布。NVIDIA 同时给出两种使用方式：端到端处理器把一条输入请求依次送过解析、分词、特征生成、GPU 推理与 PDB/mmCIF 写出；直接集成则允许开发者自行构造受支持的 torch.nn.Module，或在自有代码中复用其中部分模块。
吞吐的关键在于把三个层次的优化分开看。内核选择层会根据模型配置、GPU、数据类型与张量形状，在 BioIR 自定义内核、cuEquivariance 与 PyTorch 回退实现之间挑选；模块优化层对兼容模块启用 CUDA Graph 捕获；流水线扩展层则用 Ray 把完整模型副本放到不同 GPU 上，分发相互独立的输入。前两层缩短单次模型前向的时间，第三层通过让 CPU 阶段与 GPU 折叠重叠来提升清单吞吐，但它不会把一次前向拆到多块 GPU 上。
NVIDIA 用一组匹配基准量化了收益：1000 个人源二聚体靶点，序列总长均在 2800 残基以内，使用 3 次 recycle、200 个采样步、每个靶点 5 个扩散样本，硬件为 8 张 H100，对照实现是经 torch.compile 优化的开源 Boltz-2，输入、MSA、推理配方与 GPU 配置完全一致。BioIR 完成全部 1000 个靶点，每 GPU 小时产出 5.85 万成功折叠残基，开源实现为 2.02 万，按残基归一化的吞吐提升 2.90 倍；开源实现还在 29 个靶点上出现显存不足。
能耗估算同样按匹配口径给出。NVIDIA 把该基准线性外推到 100 万个可比靶点：按 8 张 H100 80GB 节点的额定功耗（TDP）计算，BioIR 约需 11MWh，公开实现约需 35MWh；按整机最大功耗计算则是 21MWh 对 64MWh。NVIDIA 明确说明，这些只是折叠环节的 IT 设备估算，不是实测电表读数，也不包含 PUE 等数据中心开销。
工程上需要注意的是测量口径。BioIR 的 model_inference_time 是经过 CUDA 同步的折叠模型前向时间，不含解析、分词、特征生成、后处理与写出；Ray 的副本数受 engine_stage.compute 与 num_gpus 约束，总算力不能超过可见 GPU 数。NVIDIA 建议在真实工作清单上分别用 1、2、4 个副本测量，并同时记录每小时完成结构数、GPU 与 CPU 利用率、峰值显存、完成率与重试次数，避免把模型前向的加速比当成端到端交付能力。
WeChat
Profile