基于 Tenstorrent 硬件的 Boltz-2 蛋白质结构预测

基于 Tenstorrent 硬件的 Boltz-2 蛋白质结构预测

24分钟 ·
播放数5
·
评论数0

这份源代码转录自 Moritz Thüning 在 FOSDEM 2026 上的演讲,主要介绍了他在 Tenstorrent 硬件上移植并优化 Boltz-2 生物分子预测模型的成果。Boltz-2 作为 AlphaFold 3 的开源替代方案,通过预测蛋白质结构与配体亲和力,成为药物研发领域的重要工具。Thüning 详细讲解了如何利用 Black Hole 处理器的大容量 SRAM 和混精度计算来提升性能,并集成了 Flash Attention 2 以优化复杂的计算模块。实验数据表明,该项目在性价比和特定序列长度的推理速度上均优于 NVIDIA 的消费级显卡。最终,作者展示了模型能实时生成高精度的蛋白质折叠结构,为开源科学计算提供了高效的硬件加速方案。

超越 AlphaFold 3:开源 Boltz-2 与 Tenstorrent 硬件如何重塑生物 AI 的算力版图?

1. 引言:从“Thesis-to-Tenstorrent”的职业飞跃

在 FOSDEM 2026 的演讲台上,Moritz Thüning 的经历完美诠释了当代 AI 开发者的“进化路径”。从 2024 年 11 月的一个大学“兴趣项目”(TT-Boltz),到 2025 年 4 月完成原型并以此撰写毕业论文,再到最终加入 AI 芯片巨头 Tenstorrent 全职推动该项目——这种“论文即产品、爱好即职业”的跨越,正是由 AI 民主化浪潮所驱动的。

这不仅是一个励志故事,更是一场硬核的技术演习。当个人开发者能够利用创新的非 GPU 硬件,在自家的“Quiet Box”(桌面工作站)上挑战原本由 Google DeepMind 垄断的生物技术高地时,算力的权力结构已经发生了不可逆的松动。

2. 突破点一:开源 Boltz-2 对 AlphaFold 3 的“生态位”替代

尽管 AlphaFold 2 曾因解决蛋白质折叠的“50 年难题”而荣获诺贝尔奖,但其继任者 AlphaFold 3 严苛的非商业许可证却在生物技术社区引发了强烈的反弹。正是在这种背景下,完全遵循 MIT 协议的 Boltz-2 迅速崛起,成为行业事实上的新标准。

商业自由与“结合亲和力”的双重优势 Boltz-2 的核心杀手锏不仅仅是开源带来的商业自由,更是其在功能上的超越。它不仅能预测蛋白质的静态 3D 结构,还集成了关键的“结合亲和力预测”(Affinity Prediction)。这项功能允许实验室精准预测小分子(配体)与目标蛋白质的结合强度,这是药物筛选中决定“药效”的最核心指标。

“实际上,Boltz 已经成为了行业标准,所有的生物技术实验室都在使用 Boltz,而不再使用 AlphaFold 了。” —— Moritz Thüning

深度分析: 在生物 AI 领域,单纯的预测精度已不再是唯一的护城河。Boltz-2 的成功证明,当一个具备高度实用性(如亲和力预测)的模型遇到无门槛的开源协议时,它能迅速在学术界和工业界之间建立起强大的协同效应。这种“生态位”的更替,标志着生物 AI 研发正式进入了“商用透明化”时代。

3. 突破点二:软件定义的 SRAM 刮擦板——Tenstorrent 的底层降维

在硬件底层,Tenstorrent 的 Black Hole P150 展现了与传统 GPU 截然不同的计算哲学。Black Hole 不仅仅是一个算力单元,它是一个由 Tensix 核心组成的复杂网格,每个核心配备了 5 个 RISC-V 处理器来精确掌控数据流。

显式内存控制(Explicit Control)对抗透明缓存 GPU 依赖复杂的硬件自动管理缓存(Transparent Caches),开发者往往难以预测数据在显存层级间的移动。而 Tenstorrent 提供了 210MB 的超大容量 SRAM,并将其作为“显式受控刮擦板”(Scratchpad)。开发者可以利用软件精确定义数据何时驻留、何时交换。

深度分析: 这种“软件定义硅片”的能力带来了极高的确定性。在处理 Boltz-2 这种内存密集型模型时,通过显式控制 SRAM,开发者可以大幅减少对 GDDR6 外部带宽的依赖,让计算算子在 SRAM 内部实现极致的“核融合”。这不仅仅是 TFLOPS 的较量,更是“SRAM 局域性”对“HBM/GDDR 延迟”的一场算力效率阻击战。

4. 突破点三:物理学的归纳偏置——攻克“三次方复杂度”的挑战

Boltz-2 架构中最具挑战性的部分是“对演化器”(Pairformer)中的三角形注意力(Triangle Attention)和三角形乘法。这不仅是数学上的精妙,更是将物理世界的本质规律编码进 AI 模型。

归纳偏置的工程化实现 三角形注意力的设计初衷是确保原子间的空间距离遵循物理学中的“三角形不等式”。这种“归纳偏置”(Inductive Bias)虽然提升了预测精度,但也带来了灾难性的计算开销:其时间复杂度是序列长度的三次方(Cubic Complexity),而内存激活值的复杂度则是二次方(Quadratic)。

“这感觉就像是在工程化归纳偏置……他们采用常规的机器学习操作,但结合了来自物理学等领域的要求,并将其工程化到模型中。” —— Moritz Thüning

深度分析: 这里隐藏着一个关键的性能陷阱:由于“对表示”(Pair Representation)引入了额外的维度,激活值(Activations)而非权重(Weights)成为了内存瓶颈。Tenstorrent 的优势在于其极高的“SRAM-to-Compute”比例,通过对激活值进行“分块处理”(Chunking),可以在受限的内存空间内优雅地处理 N^2 级别的内存压力。

5. 突破点四:性价比神话——为什么 RTX 5090 输掉了“每美元收益”

在性能实测中,Moritz 将 Tenstorrent Black Hole 与市面上最强劲的消费级显卡 RTX 5090 进行了对比。结果显示,对于生物制药实验室而言,单纯追求峰值算力可能并非最优解。

内核级优化:Flash Attention 2 与 Batch Broadcasting 为了在 Black Hole 上榨取极致性能,Moritz 团队实施了一项关键的工程创新:他们在 Flash Attention 2 内核中直接添加了对“批量广播”(Batch Broadcasting)的支持,以处理三角形偏置。此外,通过使用混合精度 Block FP8(这是 NVIDIA 在此场景下尚未普及的格式),Tenstorrent 成功在小型蛋白质序列上跑赢了 RTX 5090。

深度分析: 虽然在超过 1000 个氨基酸的长序列中,Black Hole 目前仍存在性能波动(Moritz 坦言这可能是一个待优化的 Bug),但在“每美元每小时预测结构数”这一商业核心指标上,Tenstorrent 展现出了压倒性的优势。这证明了在垂直领域,针对特定算法(如三角形注意力)进行深度内核定制和算术强度优化,足以抵消通用显卡的规模优势。

6. 突破点五:从毫秒级折叠中见证生命的诞生

在 FOSDEM 现场的演示(Demo)中,Boltz-2 展现了令人屏息的效率:从生成多序列比对(MSA)、加载模块到执行扩散模型(Diffusion Module),整个过程流畅而迅速。

实时化的药物发现路径 随着扩散模块的运行,屏幕上的氨基酸序列在几秒钟内便“折叠”成了精确的 3D 蛋白质结构。与冷冻电镜(Cryo-EM)数年的实验周期相比,这种毫秒级的计算反馈让研究人员能够实时观察配体与蛋白的结合信心分数(Confidence Scores)。

深度分析: 这种“实时性”是革命性的。它意味着药物研发的漏斗可以被极度放大——从每天筛选几个化合物到每天筛选数百万个。当这种算力被部署到由 32 颗处理器组成的 Galaxy 服务器集群时,人类对抗病毒和绝症的速度将被重新定义。

7. 结语:算力、开源与治愈的距离

Moritz Thüning 从一个学生开发者成长为 Tenstorrent 工程师的历程,实际上是新一代 AI 基础设施生态的缩影。当开源模型(Boltz-2)打破了许可证的枷锁,当创新硬件(Tenstorrent)打破了内存墙与显存成本的垄断,创新的门槛正在以前所未有的速度崩塌。

我们正处于一个奇点:算力正在变得像电力一样触手可及且专业化。

思考题: 如果药物研发的算力成本降低 10 倍,而我们的预测效率提升 100 倍,我们离治愈下一种毁灭人类的绝症,究竟还有多远?