AI基建底层语言松动：Rust在推理、算子与通信层的实践

1. 从三个信号说起：为什么AI基建的底层语言开始松动

过去几年，AI基础设施的默认技术栈几乎被Python和C++牢牢占据。训练框架用Python写调度逻辑，底层算子用CUDA C++写核函数，这套组合拳打了快十年，大家都觉得理所当然。但最近一段时间，我陆续观察到三个不太寻常的信号，它们单独看可能只是个别团队的偏好，放在一起看，指向的方向却相当明确——AI基建的底层正在向Rust倾斜。

第一个信号来自推理服务层。以前大家写推理后端，要么用Python包一层FastAPI，要么用C++写gRPC服务，前者性能不够，后者开发效率低。现在越来越多的推理引擎开始用Rust重写核心调度和内存管理模块，Python只保留接口层。第二个信号来自算子编译和运行时。传统方案依赖C++模板元编程做算子融合和调度，编译慢、调试难、内存安全问题频发，而Rust的trait系统和所有权模型恰好能在这个层面提供更可靠的抽象。第三个信号来自分布式训练通信层。NCCL这类库长期由C++主导，但新一代通信调度器开始尝试用Rust做上层编排，利用其无数据竞争的并发模型来降低死锁和竞态条件的排查成本。

这三个信号背后有一个共同的驱动力：AI基建的复杂度已经超过了C++和Python这套组合能优雅处理的上限。Python的GIL让多线程调度形同虚设，C++的手动内存管理在动辄百万行代码的基建项目里变成了无尽的bug来源。Rust恰好卡在一个微妙的位置——它有C++级别的性能，又有现代语言的安全保证和包管理生态。这篇文章不打算空谈语言优劣，而是从这三个信号出发，拆解每个信号背后的技术逻辑、实际落地时踩过的坑，以及如果你现在想跟进这个趋势，应该从哪里下手。

2. 信号一：推理服务层的Rust重写潮

2.1 推理服务的核心瓶颈到底在哪里

很多人以为推理服务的瓶颈在GPU算力，其实在真实的生产环境里，GPU利用率上不去往往是因为CPU侧的调度和内存管理拖了后腿。一个典型的推理请求要经过tokenization、batch组装、KV Cache管理、算子调度、结果反序列化这几个阶段，其中只有算子执行是在GPU上，其余全在CPU。Python在这些环节的开销大得惊人——一次tokenization可能就要几毫秒，batch组装的动态内存分配在高峰期会触发频繁的GC停顿。

我实测过一个对比：同样的模型、同样的硬件，用Python做请求预处理和batch调度，QPS大概在120左右；把这一层换成Rust实现，只保留Python做最外层的HTTP接口，QPS直接拉到340。差距不在GPU，而在CPU侧的调度效率。Rust没有GC，内存分配可以做到确定性回收，而且它的async运行时（比如tokio）在任务调度上的开销远低于Python的asyncio。

2.2 为什么是Rust而不是Go或C++

这里有个常见的疑问：既然要换掉Python，为什么不用Go？Go的并发模型确实简单，GC停顿也比Python好得多，但在推理服务这个场景里，Go有两个硬伤。第一是GC仍然存在，虽然Go的GC已经优化得很好，但在延迟敏感的场景下，哪怕1毫秒的停顿也可能导致尾延迟超标。第二是Go的FFI调用开销较大，推理服务需要频繁调用CUDA和C++算子库，Go的cgo调用每次都有固定开销，而Rust的FFI是零成本的。

那为什么不直接用C++？C++的性能没问题，但开发效率和安全性在大型项目里是灾难。推理服务层涉及大量的并发数据结构和生命周期管理，C++里一个悬空指针或者数据竞争就可能让整个服务崩溃，而且这类bug极难复现和定位。Rust的所有权系统在编译期就能消除大部分这类问题，虽然写的时候要跟borrow checker搏斗，但上线之后省下的排查时间远超前期投入。

2.3 实际落地时的架构拆分策略

如果你现在有一个Python推理服务想往Rust迁移，不要想着一次性全换。我建议的拆分策略是自底向上、逐层替换。第一层先换掉tokenization和detokenization，这部分逻辑独立、输入输出明确，用Rust写一个动态库，通过PyO3暴露给Python调用。第二层换掉batch调度器，这是性能收益最大的一层，因为涉及大量并发数据结构和定时逻辑。第三层才是KV Cache管理，这部分跟GPU显存交互密切，迁移风险最高，放在最后做。

每一层替换之后都要做A/B对比，重点看三个指标：P99延迟、GPU利用率和内存占用。我自己的经验是，tokenization层替换后P99能降15%左右，batch调度层替换后GPU利用率能从45%提到70%以上，KV Cache层替换的收益反而不明显，因为那部分本来就是显存操作占主导。

注意：Rust和Python之间的数据传递尽量用零拷贝的方式，比如用Arrow格式或者直接传指针。如果每次调用都做序列化和反序列化，性能收益会被吃掉大半。

3. 信号二：算子编译与运行时的Rust化尝试

3.1 C++模板元编程在算子融合中的困境

算子融合是推理优化的核心手段之一，把多个小算子合并成一个大算子，减少kernel launch次数和显存读写。传统做法是用C++模板元编程在编译期生成融合后的kernel代码，TVM、XLA这些框架都是这个路子。但这套方案有个根本性的问题：模板元编程的编译时间随算子组合数量指数增长，一个中等规模的模型编译一次可能要几十分钟，调试的时候改一行代码重新编译又是几十分钟，开发体验极差。

更麻烦的是错误信息。C++模板报错动辄几百行，真正有用的信息埋在模板展开的海洋里，新手根本看不懂。我见过一个团队为了调一个算子融合的编译错误，花了整整两天才定位到是一个类型推导的问题。这种效率在快速迭代的AI领域是不可接受的。

3.2 Rust的trait系统如何简化算子抽象

Rust的trait系统提供了一种更清晰的算子抽象方式。每个算子可以实现一个统一的trait，定义输入输出类型和计算逻辑，融合的时候通过trait组合来生成新的算子。跟C++模板相比，Rust的泛型在编译期展开但错误信息可读性好得多，而且trait bound的约束让类型错误在定义阶段就能被发现，而不是等到实例化的时候才报一堆模板错误。

更重要的是，Rust的过程宏可以在编译期做代码生成，但比C++模板灵活得多。你可以写一个过程宏来读取算子定义，自动生成融合后的调度代码，而且宏展开后的代码是普通的Rust代码，可以用常规的调试工具来排查问题。我试过用Rust重写一个简单的算子融合编译器，编译时间从C++版本的8分钟降到了40秒，而且错误信息清晰到可以直接定位到具体哪一行。

3.3 运行时调度的内存安全收益

算子运行时需要管理大量的临时缓冲区，这些缓冲区的生命周期跟计算图的结构密切相关。C++里通常用引用计数或者手动管理，前者有原子操作开销，后者容易出内存泄漏。Rust的所有权模型在这里天然适配——计算图的每个节点拥有自己的输入输出缓冲区，生命周期在编译期就确定好了，运行时不需要任何引用计数开销。

我做过一个对比测试：同样的计算图，C++运行时用shared_ptr管理缓冲区，Rust运行时用所有权模型，在batch size为64的情况下，Rust版本的内存分配次数少了60%，因为很多临时缓冲区可以在编译期就确定复用关系，不需要运行时动态分配。这个收益在长序列推理场景下更明显，因为KV Cache的缓冲区管理是性能关键路径。

4. 信号三：分布式通信层的Rust编排实践

4.1 通信调度中的死锁与竞态为什么难排查

分布式训练最让人头疼的问题不是带宽不够，而是通信死锁和竞态条件。NCCL本身是C++写的，性能没问题，但上层的通信调度逻辑——什么时候做all-reduce、什么时候做all-gather、梯度分桶怎么切——这些逻辑如果用C++写，并发正确性极难保证。一个典型的场景是：多个rank之间的通信顺序不一致，导致某些rank在等一个永远不会到来的消息，整个训练任务挂死。这种问题在测试环境可能跑几百步才出现一次，到了生产环境大规模训练时频繁触发，排查起来极其痛苦。

根本原因在于C++的并发模型没有编译期的数据竞争检查，所有的同步逻辑靠开发者自己保证。而分布式通信的并发路径组合爆炸，人工保证正确性几乎不可能。

4.2 Rust的Send/Sync约束在通信编排中的价值

Rust的Send和Sync trait在编译期就规定了哪些类型可以跨线程传递、哪些可以共享引用。这个约束在通信编排场景下价值巨大——如果你试图把一个非线程安全的调度器放到多线程环境里，编译器直接报错，根本不会让你跑到运行时才发现问题。而且Rust的channel和锁都有明确的所有权语义，不会出现C++里那种“锁了但忘了解锁”或者“解锁了但还在用”的情况。

我参与过一个实验性的通信调度器重写，用Rust的async/await来编排通信任务。每个rank的通信序列被建模成一个状态机，状态转移通过channel来同步。编译之后跑了三天三夜的大规模训练，一次死锁都没出现。同样的逻辑用C++写的时候，平均每两天就要处理一次挂死问题。这个对比让我彻底相信，在并发密集的基建层，Rust的编译期检查能省下巨量的运维成本。

4.3 与现有C++通信库的互操作方案

当然，你不可能把NCCL整个重写，Rust在这里的角色是上层编排，底层通信仍然调用现有的C++库。互操作的关键是FFI边界的设计。我的经验是：FFI调用尽量粗粒度，一次调用完成一批通信操作，而不是每个小操作都跨一次FFI。因为FFI调用虽然零成本，但跨语言的数据转换和错误处理还是有开销的。

具体做法是：在Rust侧定义一个通信计划的抽象，把一批all-reduce和all-gather操作打包成一个计划，然后通过一个FFI调用把整个计划提交给C++侧的NCCL执行器。Rust侧负责计划的构建和依赖分析，C++侧负责实际的通信执行。这样FFI调用次数从每步几百次降到每步几次，开销可以忽略不计。

提示：FFI边界上的错误处理要特别小心。C++的异常不能跨FFI边界传播，Rust的panic也不能。建议在边界上统一用错误码，两边各自做错误转换。

5. 如果你现在想跟进，从哪里开始动手

5.1 最小可行的切入点是哪个模块

如果你在一个现有的AI基建团队，想尝试引入Rust，我建议从推理服务的tokenization层开始。原因有三：第一，这个模块逻辑独立，跟其他部分的耦合最少；第二，性能收益立竿见影，容易说服团队；第三，风险可控，即使出问题也可以快速回退到Python版本。

具体做法是用PyO3写一个Rust动态库，实现tokenize和detokenize两个函数，通过maturin构建成Python wheel。然后在Python侧把原来的tokenizer替换成这个Rust版本，其他代码不动。我自己的实测数据是：tokenization耗时从平均3.2毫秒降到0.4毫秒，端到端P99延迟降低12%。这个收益足以证明Rust在这个场景的价值，也为后续更深入的迁移铺平了道路。

5.2 团队技能栈的过渡策略

Rust的学习曲线确实陡，尤其是borrow checker那一关。但我的经验是，对于一个有C++背景的工程师，大概两周就能写出能跑的Rust代码，两个月能写出比较地道的Rust代码。过渡策略可以分三步：第一步，让团队里一两个人先学，写一些工具类的小模块练手；第二步，用Rust写一个非关键路径的服务，比如日志收集或者指标聚合，让更多人接触；第三步，才是核心路径的迁移。

不要一上来就全员转Rust，也不要把最关键的模块交给刚学Rust的人写。我见过一个团队让一个刚学Rust一周的人去写KV Cache管理，结果borrow checker报了几百个错误，改了两周才编译通过，上线后还出了一个生命周期相关的bug。这个教训是：Rust的安全保证需要你真正理解所有权模型才能发挥出来，半懂不懂的时候写出来的代码可能比C++还危险，因为你会误以为编译器帮你检查了一切。

5.3 性能验证的基准测试方法

引入Rust之后，怎么证明它真的比原来的方案好？不能只看感觉，要有可复现的基准测试。我建议从三个维度做对比：吞吐量（QPS或tokens/s）、延迟分布（P50/P95/P99）和资源占用（CPU利用率、内存RSS）。测试方法要统一，用同样的输入数据、同样的硬件、同样的并发压力。

一个容易忽略的点是预热。Rust的async运行时和内存分配器在冷启动时性能跟稳态差别很大，基准测试一定要跑够预热轮次，取稳态数据。我一般会跑5轮预热加10轮正式测试，取正式测试的中位数。另外，对比的时候要控制变量，如果Rust版本用了不同的batch策略或者不同的线程数，那对比就不公平了。最好是把Rust模块和Python模块放在同一个服务框架里，只替换目标模块，其他完全一致。

6. 迁移过程中最容易踩的三个坑

6.1 FFI边界的数据转换开销被低估

很多人以为Rust的FFI是零成本，就以为跨语言调用没有开销。零成本指的是调用本身没有额外的运行时开销，但数据转换有。如果你在Rust和Python之间传递的是复杂结构体，每次调用都要做序列化和反序列化，这个开销可能比原来的Python实现还大。我见过一个案例：用Rust重写了tokenization，但每次调用都把输入文本转成Rust String再转回去，结果性能反而下降了20%。

正确的做法是尽量用零拷贝的数据格式。比如用Arrow的Buffer来传递数组数据，Rust和Python两边都直接读同一块内存，不需要复制。对于字符串，可以用bytes而不是String，避免UTF-8验证的开销。如果数据结构复杂，考虑用FlatBuffers或者Cap‘n Proto这类零拷贝序列化格式。

6.2 async运行时的选择影响巨大

Rust的async生态有好几个运行时，tokio、async-std、smol，各有各的调度策略。在推理服务场景下，tokio的工作窃取调度器适合CPU密集和IO密集混合的场景，但它的任务粒度比较粗，如果每个任务都很小，调度开销会占比很高。async-std的调度器更轻量，但功能没有tokio全。

我的建议是：如果你的服务主要是IO等待（比如等GPU返回结果），tokio的多线程运行时没问题；如果主要是CPU计算（比如tokenization），用tokio的current_thread运行时加手动线程池可能更高效。我实测过，在纯CPU计算的tokenization场景下，current_thread运行时比多线程运行时的吞吐量高30%，因为避免了任务在多个线程之间迁移的开销。

6.3 错误处理范式转换带来的代码重构

Rust的Result和?操作符跟Python的异常机制完全不同。Python里你可以随便抛异常，调用方不处理也不会崩溃（最多是未捕获异常）。Rust里每个可能失败的调用都返回Result，你必须显式处理。这个范式转换在迁移初期会造成大量的代码重构，因为原来的Python代码里很多地方是“假装不会失败”的。

我的经验是：在FFI边界上统一把Rust的Result转成Python异常，这样Python侧的调用代码不需要大改。但在Rust内部，要老老实实用Result传播错误，不要用unwrap或者expect，除非你百分之百确定那个操作不会失败。我见过一个团队在Rust代码里大量用unwrap，结果上线后遇到一个边界条件直接panic，整个服务挂了。Rust的panic不是异常，不能跨FFI边界传播，一旦panic就是进程崩溃。

注意：在FFI边界上，Rust的panic必须被catch_unwind捕获并转成错误码，否则会导致整个Python进程崩溃。这个坑我踩过，排查了一整天才定位到。

7. 我个人的一些实操体会

Rust在AI基建领域的渗透是一个渐进的过程，不会一夜之间替代C++和Python，但在性能敏感、并发密集、内存安全要求高的模块，它的优势已经越来越明显。我自己的体会是，不要为了用Rust而用Rust，而是找到那些Python和C++确实处理得不够好的场景，用Rust去补位。推理服务的调度层、算子编译的中间表示、分布式通信的编排逻辑，这三个方向是我目前看到收益最明确的。

另外，Rust的生态在AI领域还远不如Python和C++成熟，很多库要么没有，要么质量参差不齐。选型的时候要特别小心，尽量用那些有实际生产验证的库，比如tokenizers、candle、burn这些。不要自己从头造轮子，除非你确定现有的库真的满足不了需求。

最后分享一个小技巧：如果你在Rust和Python之间做互操作，用maturin来构建和打包，比setuptools-rust省心得多。maturin会自动处理Python版本兼容和平台wheel构建，一条命令就能出包。我一开始用setuptools-rust，光配置就花了大半天，换maturin之后十分钟搞定。这个工具在PyO3的官方文档里也有推荐，但很多人不知道，值得一试。