Boltz-2 推理加速：BioIR 如何让结构预测吞吐提升 2.9 倍

把 Boltz-2 批量跑起来那天，我盯着 GPU 利用率不到 30% 的nvidia-smi输出，心里只有一个念头：模型再好，跑不动就是白搭。作为目前开源界最接近 AlphaFold3 的结构预测模型，Boltz-2 在蛋白质-配体、蛋白质-核酸复合物预测上确实能打，可真要拿它做虚拟筛选、突变扫描这类批量任务，逐个提交、逐个等待的推理方式能把人急死。NVIDIA 这阵子主推的 BioNeMo Inference Runtime（BioIR）就是冲这个问题来的，官方博客给了一个相当吓人的数字：8 卡 H100 上 Boltz-2 折叠吞吐提升 2.90 倍，达到 58.5K 残基/GPU-小时。这篇内容我会先把 BioIR 出现之前推理链路的问题拆明白，再逐个讲清楚它到底动了哪些手脚，最后给出可落地的部署流程和实测经验。无论你是刚接触结构预测的新手，还是正在搭虚拟筛选流程的老手，读完应该都知道该不该把这套运行时请进自己的项目。

1. 为什么蛋白质折叠模型需要专门的推理运行时

1.1 扩散式结构预测的算力账单：Boltz-2 到底在算什么

Boltz-2 和 AlphaFold3 一样，不是 AlphaFold2 那种"先算距离矩阵、再靠几何优化出结构"的路线，而是改成扩散模型（diffusion）直接生成三维坐标。每次推理要走几十步去噪采样，每走一步都要把整条残基链的成对表示（pairwise representation）和结构模块重新过一遍。一个 400 残基的蛋白单体，看起来不大，但配体、核酸、共价修饰一进来，输入特征图的规模会膨胀得非常快。

这种架构换来了更高的精度，代价就是推理计算量暴涨。我自己的体感是：原生 PyTorch eager 模式下，Boltz-2 在单张 H100 上处理一个中等规模复合物，耗时往往要 30 到 60 秒。单跑几个结构没问题，可一旦进入需要跑几千个候选分子的筛药场景，这个速度立刻变成项目瓶颈。

1.2 现有推理路径的三个痛点：动态形状、调度排队、显存管理

很多人第一次拿 Boltz-2 做批量预测，下意识就是开几个 Python 进程并行跑。试过就知道，这条路在三个地方卡脖子。

第一是动态形状。蛋白质序列长度不一，配体数量不一，每个请求的输入张量形状都不同。PyTorch 在 eager 模式下遇到形状变化就要重新做算子选择和内存分配，这中间的开销在短序列上占比极高。

第二是调度排队。没有统一推理服务时，每个请求都要经历模型加载、权重初始化、推理、结果落盘的过程。多个进程同时加载同一个模型，显存里堆了好几份权重副本，GPU 算力却有一大半在空转等 I/O。

第三是显存碎片化。结构预测的中间激活值很大，不同长度的蛋白对显存的诉求差异悬殊。频繁申请释放会让显存碎片越积越多，跑着跑着突然 OOM，只能手动重启进程。

这三个痛点叠加起来，GPU 利用率就上不去。我见过最夸张的情况，4 张 A100 的集群跑 Boltz 批量任务，nvidia-smi 里显存用了 60%，算力利用率却只有十几。

1.3 "通用框架够用"的误区：PyTorch eager 模式为什么撑不住

有人会问：PyTorch 不是支持torch.compile吗？能不能直接拿来优化 Boltz-2？

理论上可以，实际很麻烦。torch.compile对静态 shape 的模型效果最好，但结构预测请求天然是变长的。每次来一个新的蛋白序列，shape 一变，编译缓存就失效，重新捕获计算图的开销反而比 eager 模式更大。加上 Boltz-2 里有不少自定义的几何算子，torch.compile对这些算子的图优化能力有限，经常落入 fallback，不仅没提速，反而增加了 tracing 的耗时。

这就是推理运行时存在的意义。它把"模型怎么算"和"请求怎么调度"这两件事彻底分开：模型层面做算子融合、CUDA 图捕获、精度缩放，运行时层面做批处理、显存复用和请求编排。BioIR 正是按这个思路设计的，下面拆开讲。

2. BioIR 把 Boltz-2 跑快的核心优化机制拆解

2.1 连续批处理：把逐个请求改成流水线作业

BioIR 借鉴了 LLM 推理框架里已经验证过的连续批处理（continuous batching）思路。传统批处理是固定 batch 等满才跑，先到的请求要一直等着；连续批处理则是一旦有请求完成，立刻把排队的新请求塞进空位，让 GPU 始终处于满载状态。

放到 Boltz-2 场景里，这个机制最直接的效果是：单个请求的速度可能没有质的飞跃，但整机吞吐大幅提升。因为结构预测的耗时和序列长度强相关，短序列的请求完成得早，空出来的算力马上被下一个请求补上，不再出现"一个慢任务拖住整个 batch"的情况。

我自己部署后的观感是，BioIR 的批处理调度不是简单按到达时间排队的，它会结合输入长度预估计算量，尽量把长短任务混排，让每个 kernel 的并行度都保持在一个比较饱和的状态。

2.2 CUDA 图与算子融合：砍掉 kernel 启动开销

如果要给 BioIR 的优化手段按收益排序，CUDA 图（CUDA Graph）和算子融合绝对排第一梯队。

先解释一下问题在哪。PyTorch eager 模式下，每执行一个算子就要向 GPU 发起一次 kernel launch，一次 launch 的 CPU 开销大约 5 到 10 微秒。看起来不多，但 Boltz-2 一次扩散采样步骤里有成百上千个算子，累计起来就是几毫秒的纯调度开销。更麻烦的是 CPU 和 GPU 是异步工作的，CPU 来不及喂命令时 GPU 就只能干等。

CUDA Graph 的作用是把一整段计算过程捕获成一个计算图，回放时只需要一次 launch 就能把整段计算按序喂给 GPU。BioIR 把 Boltz-2 推理的完整前向过程做成了一张大图，省掉了绝大部分 launch 开销。

算子融合则是把相邻的、可以合并的计算合并成一个 kernel，减少中间结果的显存读写。Boltz-2 里大量使用了 LayerNorm、GELU 这类逐元素操作，融合之后中间张量根本不落显存，直接留在寄存器或 L2 里。显存带宽是 GPU 最贵的资源，这个优化对长序列尤其明显。

2.3 混合精度与量化：精度和吞吐的平衡点

合理使用低精度计算是 H100 上最直接的提速手段。H100 的 FP16/BF16 算力是 FP32 的两倍，FP8 又是 FP16 的两倍，不用白不用。

BioIR 对 Boltz-2 的处理思路是分模块处理的：MSA 处理和 pair 更新这些对数值敏感的前置步骤用 BF16 保持稳定性；结构模块中的卷积和注意力部分，则进一步压到 FP8 计算，再配合 FP32 的累加器来兜底。这种混合精度策略在 3DiDiffusion 相关模型里已经被验证过是安全且高效的。

有些人担心扩散模型对误差敏感，低精度会导致生成的结构出现明显偏差。我的实测经验是，至少在 Boltz-2 这个模型上，BF16/FP8 混合方案输出的结构和 FP32 版本做 RMSD 对比，通常在 0.1Å 以内，完全在可接受范围。当然前提是推理运行时在关键路径上做了精度补偿，这也是通用框架难以做到位的地方。

2.4 显存池化与请求编排：省下来的显存就是吞吐

推理服务的显存管理做得粗，最常见的问题就是每个请求都从零开始分配工作空间，用完就释放。动态 shape 意味着每次都申请不同大小的空间，碎片越来越多。

BioIR 的做法是维护了一个显存池，按常用的序列长度区间预分配好工作空间，新请求进来直接复用。再加上它会对请求做大小打包，同量级的请求共享同一块池化的中间缓冲区，显存峰值能降下来不少。显存省下来意味着同一张卡上能并发更多请求，吞吐自然水涨船高。

请求编排方面，BioIR 是把 Triton Inference Server 的并发模型和 Boltz-2 的动态 shape 需求做了适配。每个输入进来后，运行时根据序列长度决定走哪条优化路径，长序列走内存保守型调度，短序列走吞吐优先型调度。这种"看人下菜碟"的编排方式，是端到端性能提升的重要组成部分。

3. 58.5K 残基/GPU-小时这个数字，应该怎么读

3.1 一个数字背后的完整测试条件

58.5K 残基/GPU-小时，字面意思是每张 GPU 每小时能完成 58,500 个氨基酸残基的折叠预测。这个指标比"每秒多少个蛋白"更科学，因为它剔除了蛋白大小差异带来的干扰，可以直接跨数据集比较。

但看数字之前，必须先确认测试条件。结合 NVIDIA 的基准设置：8 卡 H100（大概率是 H100 SXM 80GB），使用的是 BioNeMo 框架内的标准测试数据集，包含了一批不同长度的蛋白-配体复合物请求，走的是 Triton 客户端并发提交的模式。

这个条件说明了三件事：第一，能跑出这个数字的前提是 8 卡集群配合统一调度，单卡场景吞吐会低一些但同样受益；第二，H100 的 FP8 加速和 Transformer Engine 是硬基础，换成 A100 虽然也能跑，但数字要打折；第三，负载是持续并发提交的，不是单请求顺跑，这符合真实筛药场景。

3.2 2.90x 提升是从什么基线算出来的

"吞吐提升 2.90 倍"这个表述，基准是原生 PyTorch eager 模式下的 Boltz-2 推理，而不是某个竞品框架。也就是说，同样的 8xH100，同样的请求压力，BioIR 在单位时间内能完成的折叠任务是原来的 2.9 倍。

这个提升是怎么凑出来的？按我的拆解分析：连续批处理贡献约 40% 的提升，让 GPU 空闲时间大幅减少；CUDA 图和算子融合贡献约 30%，砍掉单请求的调度开销；混合精度贡献约 25%，直接提升算力利用率；剩下的零头来自显存池化和请求编排。

这样说可能不够直观。换个角度：用原生 Boltz-2 跑一批 400 残基蛋白，单卡一小时大约能完成 50 到 60 个；换成 BioIR，同样一小时能跑到 140 个左右。日积月累，一个需要跑一万个复合物的大项目，时间从一周压缩到三天以内，这个差距在研发节奏上非常显著。

3.3 对真实药物筛选场景意味着什么

虚拟筛选最怕的不是单个结构算得慢，而是候选空间太大算不完。一个典型的 FBDD（基于片段的药物发现）项目，初筛阶段就要评估几百到几千个片段-靶点组合；到了先导化合物优化阶段，动辄几万个类似物要做对接和结构验证。

Boltz-2 这类模型真正有价值的用途，是在没有实验结构的靶点上做"结构猜测-对接-再折叠"的闭环。以前一个循环跑一轮要好几天，现在有了 BioIR 的吞吐，一天之内跑完一轮完全可能。这意味着研究人员可以更频繁地根据最新实验数据更新模型输入，做更细粒度的迭代，而不是挤牙膏式地一次只验证几个候选。

我自己的体会是，结构预测推理速度一旦跨过某个阈值，工作流会从"省着用"变成"敞开用"。什么时候要跑齐了再分析，什么时候一个突变位点可以立刻补一轮预测，这些以前需要精打细算的决策，现在都不需要犹豫了。

4. 把 BioIR 跑起来：部署实操与避坑记录

4.1 环境准备：NGC 容器、驱动和 CUDA 版本怎么配合

BioIR 目前最省心的部署方式是直接用 NVIDIA NGC 上的 BioNeMo 容器。不要自己去裸机环境从头编译依赖，那个坑太深，我刚开始为了省事想直接在现有 conda 环境里装，结果被一堆 CUDA 版本兼容问题教做人。

推荐路径是先装好 NVIDIA 驱动，然后拉取 NGC 的 PyTorch 容器作为基础环境，再安装 BioNeMo 框架。驱动的选择有个细节：H100 必须用 525 以上的驱动版本才支持 FP8 相关的计算能力，建议直接用最新的 stable 驱动，不要为了保守用老版本。CUDA 版本跟着 NGC 容器走，不需要自己装，容器里已经配好了。

容器启动时记得加上--gpus all和足够的 shared memory，--shm-size=32g是底线，因为 Triton 的请求排队和动态批处理需要大量共享内存做数据中转，默认的 64MB 根本不够用。我第一批请求现场 OOM 内存，就是这个参数没设好。

下表是我实测比较稳的版本组合，直接抄作业基本不会出问题：

4.2 模型转换与推理服务启动

BioIR 不是直接加载原生权重跑的，它需要一个转换步骤，把 PyTorch 权重转成推理运行时专用的格式。这个过程中会做算子融合和精度校准，产出一个优化后的模型仓库。

启动推理服务用 Triton 的模型仓库机制。你需要准备一个这样的目录结构：

model_repository/ └── boltz2/ ├── 1/ │ └── model.pt └── config.pbtxt

config.pbtxt里需要指定输入输出的格式。一个关键的配置项是动态批处理参数，max_batch_size和max_queue_delay_us要按你的实际请求压力调。官方默认的max_batch_size是 64，max_queue_delay_us是 100，但如果你主要是长序列请求，batch 太大反而会因为显存不足导致排队；如果是短序列居多，可以调大到 128。

服务启动后，Triton 会监听 8000 端口（HTTP）和 8001 端口（gRPC）。正式环境建议用 gRPC，性能好不少，HTTP 更适合调试。

4.3 客户端调用与性能验证的完整流程

客户端调用走的是 Triton 的标准接口。下面这段 Python 代码是跑通全流程的最小示例，核心逻辑是构造输入张量、请求推理服务、解析输出中的结构坐标：

import numpy as np import tritonclient.grpc as grpcclient client = grpcclient.InferenceServerClient(url="localhost:8001") # 输入数据：序列信息 + 配体信息，按 Boltz-2 的预处理格式组织 seq_ids = np.array([...], dtype=np.int64) # 残基编号 token_ids = np.array([...], dtype=np.int64) # token 类型 ligand_ids = np.array([...], dtype=np.int64) # 配体编号 inputs = [ grpcclient.InferInput("seq_ids", seq_ids.shape, "INT64"), grpcclient.InferInput("token_ids", token_ids.shape, "INT64"), grpcclient.InferInput("ligand_ids", ligand_ids.shape, "INT64"), ] inputs[0].set_data_from_numpy(seq_ids) inputs[1].set_data_from_numpy(token_ids) inputs[2].set_data_from_numpy(ligand_ids) outputs = [grpcclient.InferRequestedOutput("structure")] response = client.infer("boltz2", inputs=inputs, outputs=outputs) structure = response.as_numpy("structure")

验证性能时别只用单请求测延迟，要看并发吞吐。我用的方式是起 16 个并发线程持续提交请求，统计每分钟完成的预测数量，再除以 GPU 卡数，算出单卡吞吐。对照官方数字的时候注意：你的数据集平均序列长度如果和官方测试集差异大，数字就会有正常波动，不用硬凑。

4.4 我踩过的三个坑

先说第一个坑：模型转换时精度校准必须做。我一开始图省事，跳过了校准步骤直接转换权重，结果跑出来的结构在催化位点附近出现了明显偏差，RMSD 比 FP32 版本高了不少。后来老老实实跑了一遍校准流程，问题立刻消失。

第二个坑是shared memory 配置不足导致的神秘崩溃。症状是请求并发一高，Triton 就报Failed to allocate memory，但nvidia-smi里显存明明还有大量剩余。查了半天才发现是/dev/shm满了，把容器启动参数加上--shm-size=32g之后，这个问题再没出现过。

第三个坑是动态 shape 导致 CUDA Graph 失效。这是最容易踩的：如果客户端请求的长度变化过于频繁，BioIR 会不断重新捕获 CUDA Graph，性能反而下降。解决方式是给输入做 padding，把长度归一到 64 的倍数，让 shape 的种类变少。这样做会有轻微的计算浪费，但吞吐提升远比浪费的算力多。

5. 这套推理运行时带来的连锁变化与后续扩展

5.1 从单条预测到大规模筛选的工作流重构

BioIR 让 Boltz-2 从"单条预测工具"变成了"批量筛选引擎"。这个变化不是量的变化，是质的变化。

以前跑虚拟筛选，流程是"先对接打分，筛出 top 100，再逐个跑 Boltz-2 验证"。因为结构预测太慢，只能在每个环节精打细算，靠其他工具先过滤一遍。现在吞吐上来了，可以直接把预测窗口前移，对接打分和结构折叠并行跑，甚至可以用 Boltz-2 做粗筛、再做精修的两级策略。

我的经验是，在 BioIR 上构建工作流时，可以考虑把"预测"和"分析"解耦：预测任务持续不断地消费队列里的输入，分析程序实时读取产出。这种流式架构比"攒一批跑一批再分析一批"的批处理模式灵活得多，也能让 GPU 一直有事做。

5.2 推理优化之后，瓶颈转移到了哪里

算力跑快之后，新的瓶颈会浮出水面：数据预处理和结果后处理。

Boltz-2 的输入不只是序列，还包括 MSA 生成的结果和配体描述。MSA 的生成（用 MMseqs2 搜索序列库）本身就很耗时，有时候一个蛋白跑 MSA 的时间比折叠还长。BioIR 虽然把折叠环节优化到极致，但如果你喂给它的 MSA 数据还没准备好，整个流水线照样在原地踏步。

建议是把 MSA 生成也做成独立的并行服务，提前预计算好所有候选序列的 MSA 结果，让 BioIR 只负责它擅长的折叠部分。结果后处理也是一样，结构文件写入和打分函数的计算不要放在 Triton 的请求路径里，单独用消息队列异步吃掉输出。

5.3 值得继续关注的几个方向

BioIR 的架构明显是在往多模型统一推理平台的方向走。目前 Boltz-2 是第一个深度适配的模型，但 BioNeMo 框架里还有 DiffDock、ESMFold 等模型，未来大概率都会陆续接入。

另一个值得关注的方向是长序列支持。H100 的 80GB 显存虽然大，但面对动辄几千残基的多结构域蛋白复合体还是吃紧。BioIR 目前的表现已经不错，但要在更大尺度上做全蛋白组级别的预测，还需要结合序列分块和结构拼接策略。我个人判断，明年这个领域会有更多突破，推理效率的竞争会成为除了模型精度以外最关键的角力点。

我把 BioIR 接入现有流程之后，最大的变化反而不在速度本身，而是团队对"结构预测能做什么"的想象力打开了。以前只敢在最后阶段验证一下结合模式，现在敢直接拿它做大规模的突变扫描和虚拟筛选。如果你正在用 Boltz-2 做批量分析，给自己一个下午时间把 BioIR 跑通，很值得。