AI Infra 工程师的日常工作
CUDA 算子开发与优化
这类工作最”底层”，直接和 GPU 硬件打交道。典型任务包括：为 Attention 计算编写高效的 CUDA kernel（如实现或改进 FlashAttention）、优化矩阵乘法（GEMM）以充分利用 Tensor Core、将多个小算子融合为一个大 kernel 减少显存读写。
和 Softmax 的计算与显存访问模式，如果你不知道 Self-Attention 的计算流程，就无法理解 FlashAttention 在做什么。
分布式训练
推理部署与优化
CUDA与算子优化
基础概念
- Transformer：神经网络模型结构，一种思想概念，纯模型数学结构、算法逻辑。核心是Self-Attention
- PyTorch：Python深度学习工具框架，写代码、搭建、训练、调度模型，能实现 Transformer
- CUDA：NVIDIA专属本地GPU计算软件平台，打通软件与NVIDIA GPU，让GPU执行大规模数学/矩阵计算。仅支持NVIDIA显卡，不支持CPU、AMD显卡、Intel显卡、Apple M芯片、TPU，无N卡电脑，CUDA完全无法生效，只能CPU运行模型
- GPU：图形处理器，海量轻量CUDA核心，高吞吐，适合简单、重复、大规模并行计算（矩阵乘法、Transformer运算）
- CPU：中央处理器，少量高性能核心（几核-几十核），并行算力极低，跑大模型训练极慢，适合复杂逻辑、分支判断、串行任务
- TPU：- Google自研AI专用芯片，对标GPU
- HBM：高带宽内存
- NVLink：NVIDIA 自研的高速互联总线，用来把多张 NVIDIA GPU 直接连在一起，让多卡之间高速交换显存数据。既是硬件，也是配套通信协议。
MHA、MQA、GQA、MLA 都只是【注意力模块（Attention）的变种】，本身不包含 FFN；完整 Transformer 层一定是：Attention + FFN 两个组件一起搭配使用。
Attention 变种：
- MHA：Multi-Head Attention 多头注意力。每个 Query 头，都有独立一套 K、V。优点：表达能力最强，不同头独立捕捉不同语义关系。缺点：KV Cache 体积最大；长上下文推理时，读取 KV 极度消耗显存带宽，容易触发访存瓶颈
- MQA：Multi-Query Attention 多查询注意力。所有 Query 头共用同一套 K、V；Q 头仍然保持独立。优点：KV Cache 极小，显存带宽开销大幅下降，解码速度快。缺点：精度下降明显；全部头共用一套 KV，表达能力损失大
- GQA：Grouped-Query Attention 分组查询注意力。把全部 Query 头分成若干组；同一组内所有 Q 头共享一套 K、V；不同分组 KV 相互独立，是目前绝大多数开源大模型的标配。优点：MHA 与 MQA 的折中方案。精度几乎接近 MHA，KV Cache 显著减少；平衡推理速度与模型能力。缺点：依然是「按头分组共享」，上下文极长时 KV 体积还是很大
- MLA：Multi-Head Latent Attention 多头潜在注意力，DeepSeek模型。
- 将 K、V 映射到更低维度的潜空间（latent vector），缓存压缩后的低维 KV
- 计算注意力前，再把低维 KV 上投影恢复维度
- RoPE 位置信息单独计算，不和压缩后的 KV 绑定，解决压缩破坏位置编码的痛点。
优点：KV Cache 压缩比例远高于 GQA；长上下文下显存 / 带宽收益巨大；精度损失很小。 缺点：架构复杂、训练成本更高，推理引擎实现难度更大
-
MoE：Mixture of Experts，混合专家，就是把单层 FFN 替换成多个独立小 FFN（专家 Expert），每次只激活少数几个专家参与计算。能在不显著增加单 token 计算量前提下，大幅提升模型总参数量与能力。
-
损失函数 Loss ：对模型每一个参数（weight/bias）的偏导数。表示：这个参数往哪个方向调整，能让 Loss 变小、模型预测更准。越低代表模型在训练数据上拟合越好；但 Loss 不是越低就一定代表模型效果越强。如果 Loss 过低，大概率发生过拟合
Memory Wall 显存带宽
GPU 显存和 GPU 计算核心之间，每秒能传输的数据总量。单位 GB/s。 类比： GPU 计算核心 = 工厂车间工人 显存 = 原料仓库 显存带宽 = 仓库到车间的传送带粗细 带宽越大，传送带越粗，每秒运送原料越多。
存储层次结构
寄存器 > 共享内存 > L1/L2 Cache > HBM > 主机内存。写 kernel 的核心就是在这个层次结构中最大化数据复用，减少对慢速存储的访问
CUDA 编程基础
编程模型：Grid / Block / Thread 层级，线程索引计算 内存模型：全局内存、共享内存、寄存器、常量内存的特性与用法 Warp：32 个线程组成的最小调度单位，GPU 每次下达指令都是以 Warp 为单位 核心直觉：“内存访问模式决定运行速度”
常见算子实现与优化
Reduce：并行归约的多种实现与优化（Warp Shuffle、多级归约） GEMM：矩阵乘法的分块、向量化、Shared Memory Tiling、利用 Tensor Core Softmax：Online normalizer calculation，高效 Softmax kernel 算子融合：将多个小算子合并为一个 kernel，减少全局内存读写
Attention 算子
FlashAttention V1/V2：Memory-aware 的精确 Attention 实现，通过 tiling 减少 HBM 访问——好比把一张大桌子上的拼图分成小块，每次只搬一小块到手边拼好再搬下一块，而不是把所有碎片一股脑倒出来占满桌面 FlashAttention-3：在 Hopper 架构上进一步拉高利用率 Flash-Decoding / FlashDecoding++：面向 Decode 阶段的 Attention 加速 FlashInfer：可定制 Attention 引擎，面向 Serving 的可组合格式与异构 KV 存储 PagedAttention CUDA Kernel：vLLM 中 PagedAttention 的底层实现
AI 编译器
Triton：OpenAI 开源的 GPU 编程语言，大幅降低高效算子编写门槛 TVM / XLA：计算图优化与代码生成 torch.compile：PyTorch 2.x 的编译模式，理解 Graph Break 与性能收益
分布式训练
分布式训练要解决的核心问题是任务怎么拆、怎么传、怎么汇总。
数据并行是把同一本书复印多份、每人抄不同章节的内容然后汇总；张量并行是把每一页拆成几列、每人只抄自己那几列；流水线并行则是第一个人抄完第一章就传给第二个人继续，自己接着抄下一批。
优化器
优化器是训练的”舵手”——模型算出梯度之后，怎么更新参数全靠它。
数据并行
DP（DataParallel）：最基础的数据并行，单进程多卡 DDP（DistributedDataParallel）：多进程数据并行，理解 AllReduce 梯度同步 FSDP（Fully Sharded Data Parallel）：PyTorch 原生的 ZeRO-3 实现
模型并行（3D 并行）
张量并行（TP）：将矩阵乘法沿特定维度切分到多卡，通信密集，通常限于单机 流水线并行（PP）：将模型不同层切分到不同机器，像流水线一样传递数据 序列并行（SP）：沿序列维度切分，与 TP 配合减少激活显存
ZeRO （Zero Redundancy Optimizer）
分布式训练的显存优化算法，核心思路：把【模型权重、梯度、优化器状态】这三块训练显存大头，拆分分散到多张 GPU 上，而不是每张 GPU 都完整保存全部副本。
ZeRO-1：优化器状态切分 ZeRO-2：优化器状态 + 梯度切分 ZeRO-3：优化器状态 + 梯度 + 参数切分（用通信换显存）
推理与部署
LLM 推理的两阶段：Prefill（处理输入）与 Decode（逐 token 生成）
KV Cache把已经算过的中间结果记下来，后续生成每个新 token 时就不用从头算起。
关键性能指标：TTFT（首 token 延迟）、TPOT（每 token 延迟）、吞吐量（token/s）、P50/P95 尾延迟
端到端部署：拿到一个 7B 模型，分别用 vLLM 和 SGLang 部署成 OpenAI 兼容 API，用相同的压测脚本（固定并发、输入输出长度）跑出 TTFT、TPOT、throughput 三个指标的对比表
推理框架
vLLM、SGLang、TensorRT-LLM
量化
把模型参数（权重、激活值）从高精度浮点数，转换成更低比特宽度的数据类型，用更少 bit 表示数值，目的降低显存占用、减少访存数据量、加速推理；代价是引入少量精度损失。
常见精度规格（面试必记） FP32（32bit 浮点数）：原始训练默认，占用最大，推理几乎不用 FP16 / BF16（16bit）：主流训练精度；很多模型原生推理精度 INT8（8bit 整数）：常用量化精度 INT4（4bit 整数）：当前开源 LLM 最常用，压缩幅度大 INT2（2bit）：压缩极强，精度损失通常很大，使用较少
Prefill/Decode 解耦
Prefill 与 Decode 混合 batching 造成资源耦合与互扰，导致尾延迟爆炸。这就像餐厅里让同一个厨师既负责快速出小炒（Decode 低延迟）又负责慢炖大菜（Prefill 重计算），互相拖累；解耦就是把快餐区和慢炖区分开，各配专属厨师。