AI Infra 看起来是一套全新的技术体系,但它面对的核心问题并不陌生:数据能否放得下、计算能否跑得快、通信能否隐藏、服务能否兼顾延迟和吞吐。
真正发生变化的是约束条件。传统后端系统围绕 CPU、内存和网络 I/O 设计,AI 系统则更多围绕 GPU 算力、显存容量和设备间通信展开。理解这种变化,比孤立地记住某个框架或优化技巧更重要。
硬件中心从 CPU 转向 GPU
传统在线服务主要处理请求解析、业务逻辑、数据库访问和网络通信。系统吞吐通常通过多线程、异步 I/O、缓存、分片和横向扩展来提升。
大模型的主要负载则是高密度的张量计算。以自回归语言模型为例,每生成一个 token,都要执行大量矩阵运算,并频繁读取模型权重。CPU 的计算吞吐和内存带宽难以高效承担这类任务,因此核心计算与通信逐渐转移到 GPU,CPU 更多负责调度、数据准备和设备控制。
这带来三个直接变化:
- 内存容量问题转化为显存容量问题;
- CPU 线程并发转化为 GPU 上的大规模并行计算;
- 普通网络通信转化为对低延迟、高带宽设备互联的依赖。
GPU 的计算速度很快,因此等待通信的代价也被放大。传统服务能够容忍的毫秒级抖动,在大规模 GPU 集群中可能造成明显的算力浪费。这也是当前 AI 集群倾向于使用高性能 GPU、专用互联和集中式部署的重要原因。
从长期看,这种形态未必是终点。历史上的“去 IOE”通过通用硬件和分布式软件降低了基础设施成本。AI Infra 也可能经历类似过程,但在硬件、编译器、通信协议和软件生态成熟前,不能简单假设通用分布式方案可以直接替代高度集成的 GPU 集群。
AI Infra 的基础软件栈
AI 应用的核心流程可以粗略分为训练和推理:训练从数据中调整模型参数,推理则使用训练后的参数处理新输入。
PyTorch 屏蔽了什么
深度学习框架的作用类似传统后端中的微服务框架。后者封装服务发现、序列化和网络通信,前者则封装张量运算、自动微分、计算图和设备调度。
借助 PyTorch,开发者通常只需描述模型结构、前向计算和优化目标,无须为每个数学操作手写 GPU 程序。它降低了模型开发门槛,但没有消除显存、通信和调度等系统问题。
为什么仍然需要 GPU 编程
通用框架算子不一定能表达最合适的计算过程。某些模型结构如果只由现有算子拼接,可能产生多余的中间结果、访存或计算。这时需要融合算子或编写自定义内核。
CUDA 内核的优化需要理解线程组织、内存层次和执行模型,学习成本较高。Triton 提供了更接近 Python 的开发方式,并能与 PyTorch 生态结合,适合作为自定义 GPU 算子的入门工具。不过,Triton 降低的是表达门槛,而不是性能分析本身的难度。
Python 为什么成为主语言
现代模型部署不仅要执行静态网络,还要管理 KV Cache、模型并行、MoE 路由、动态控制流和自定义算子。模型与运行时的边界因此变得模糊,完整导出到独立的 C++ 推理图越来越困难。
这意味着 AI Infra 工程师不仅要会调用 Python 接口,还需要理解 Python、PyTorch、编译器和底层 GPU 运行时之间的关系。
模型训练:先存得下,再算得快
大模型训练不能只看参数文件大小。训练时显存通常还要容纳:
- 模型参数;
- 梯度;
- 优化器状态;
- 前向传播产生、供反向传播使用的激活值;
- 通信缓冲区和框架运行时开销。
因此,一个勉强能装入模型权重的设备,通常无法直接完成模型训练。
激活值是动态的显存压力
模型参数的大小相对固定,激活值则会随批大小、序列长度、隐藏维度、层数和注意力实现变化。标准注意力如果显式保存完整注意力矩阵,会引入与序列长度平方相关的存储开销;但不能把所有 LLM 激活显存都笼统描述为 (O(N^2))。FlashAttention、激活重计算等技术也会改变实际占用。
分析训练 OOM 时,应该使用 PyTorch Profiler 或显存快照观察各阶段分配,而不是只按参数量估算。
常见缓解手段包括减小 batch、混合精度训练、激活检查点、优化器状态分片,以及把模型拆到多张 GPU 上。
模型并行对应传统系统的分片
当单卡或单机无法容纳模型时,需要把参数、张量或网络层拆到不同 GPU。常见方法包括数据并行、张量并行、流水线并行和专家并行,实际系统通常会组合使用。
这种思路与数据库分片相似:本地容量问题得到缓解,但跨分片通信和全局协调随之增加。拆分粒度并非越细越好,关键是让计算收益覆盖通信成本。
用通信计算重叠减少等待
增加 GPU 数量不会自动带来线性加速。分布式训练中的梯度同步、参数交换和中间结果传输都可能让 GPU 等待网络。
GPU stream 可以组织多个相互独立的操作队列。将计算和通信安排到不同 stream,并合理切分任务,就有机会在计算当前数据时传输下一阶段所需的数据。这与传统后端用异步 I/O 隐藏等待时间的思路一致。
重叠并不是免费的:任务依赖、显存生命周期和通信粒度都会影响效果。优化前应先通过时间线确认瓶颈究竟来自计算、通信还是负载不均。
模型推理:延迟与吞吐的权衡
训练成本通常集中发生,推理成本则会随请求量持续增长。在线推理因此同时关心两个目标:降低用户感知延迟,以及提高单位 GPU 时间内处理的请求数。
CUDA Graph 减少调度开销
普通执行方式中,CPU 会反复向 GPU 发起内核、内存拷贝等操作。单个操作的启动成本不高,但模型推理会重复执行大量小操作,累计开销可能变得明显。
CUDA Graph 可以预先捕获一组 GPU 操作及其依赖关系,再整体提交和重复执行,减少 CPU 与 GPU 之间的调度开销。它类似把多个远程操作封装成一次调用。
它更适合形状和执行路径相对稳定的工作负载。输入形状频繁变化或动态控制流较多时,需要额外设计图的复用策略。
KV Cache 用空间换时间
自回归模型生成新 token 时,历史 token 对应的 Key 和 Value 不会改变。缓存这些结果,可以避免每一步都重新计算完整上下文。
KV Cache 显著降低了重复计算,却会占用大量显存,而且占用随并发数和上下文长度增长。推理系统因此需要进一步处理缓存分页、复用、淘汰和调度问题。vLLM 等框架正是围绕这些问题提供实现。
流式响应优化用户感知
如果完整生成无法足够快,可以先优化首 token 时间(TTFT),在得到首个 token 或音频帧后立即返回,并持续传输后续结果。
流式响应没有减少模型的总计算量,但缩短了用户看到反馈前的等待时间。在线服务需要分别观测首 token 延迟、逐 token 延迟和端到端耗时,不能只看单一平均值。
批处理提高 GPU 利用率
单个请求往往不足以占满 GPU。批处理将多个请求合并为更大的张量计算,可以提高计算单元利用率,类似 Redis MGET 用一次请求替代多次网络往返。
固定批处理必须等待一批请求全部完成。LLM 请求的输入和输出长度差异很大,短请求完成后仍可能被长请求拖住,导致资源空闲。
连续批处理允许已完成的请求及时离开,并把新请求动态加入当前批次。它把推理优化从简单的“合并请求”升级为调度问题,需要同时考虑吞吐、公平性、显存容量和尾延迟。
传统工程经验如何迁移
AI Infra 不是对传统基础设施经验的否定,而是把相同的系统问题放到了新的硬件约束下:
| 传统 Infra | AI Infra | 共同目标 |
|---|---|---|
| 数据分片 | 模型并行、状态分片 | 突破单机容量 |
| 异步 I/O | 通信计算重叠 | 隐藏等待时间 |
| RPC 合并、批量读取 | 推理批处理 | 摊薄固定开销 |
| 工作窃取、动态调度 | 连续批处理 | 减少计算资源空闲 |
| 缓存 | KV Cache | 用空间换时间 |
| 减少网络往返 | CUDA Graph | 降低调度与交互开销 |
| 流式传输 | token 流式响应 | 改善用户感知延迟 |
这些类比适合帮助建立直觉,但不能替代具体分析。GPU 的执行模型、显存层次和集合通信都有独立约束,照搬 CPU 时代的实现方式通常得不到理想结果。
阶段性理解
学习 AI Infra 可以先抓住一条主线:任何优化最终都在处理计算、存储、通信和调度之间的关系。
面对具体问题时,可以依次追问:
- 当前限制来自算力、带宽、容量还是调度?
- 时间花在有效计算上,还是花在访存、通信和等待上?
- 能否通过缓存、批处理、并行或流水线改变瓶颈?
- 优化吞吐后,显存占用和尾延迟是否恶化?
- 性能结论是否由 profiler 和线上指标验证?
进入 AI Infra ,不仅需要掌握某个 AI 框架,而是能够识别瓶颈、拆分状态、隐藏等待并设计可观测的系统,还需要补齐 GPU 执行模型、深度学习运行时和分布式训练推理的具体机制。