Processing math: 9%

大模型推理优化笔记

总体框架

来源:ZOMI

GPU Architecture

体系结构(以 Hopper 为例)

  • 一个 GPU 包含 132 个 SM,80GB Global Memory(位于 HBM3(High Bandwidth Memory) 芯片,此外还有常量内存和纹理内存),50MB L2 Cache(跨 SM 共享);

  • 一个 SM 中有 128 个 CUDA Core(具体见后文参数表截图),4 个 Sub Core(每个 Sub Core有一个 Warp Scheduler 和一个 Dispatch Unit 共同控制指令发射),4 个 Tensor Core,1 个 TMA,228KB Shared Memory(L1 缓存,几百 K),256KB Register File,以及一般不需要考虑的 LSU 和 L0/L1 Instruction Cache;

  • 每个 CUDA Core 有整数、浮点运算单元,寄存器;寄存器放不下的放在 Local Memory (线程独占,但实际存放在 Global Memory 或 L2 Cache 中);

  • Tensor Core:

    • Ampere(A100) 4×4×16 FP16, BF16, TF32, INT8;
    • Hopper(H100) 8×8×16 FP8, FP16, BF16, TF32, INT8
    • Blackwell(GB200) 16×16×16 FP4, FP8, FP16, BF16
    image.png
  • NCCL:

    • NVLink:高速 GPU 互联总线(对应 CPU 的 PCIe)与通信协议,支持统一内存访问(Unified Memory Access);
    • NVSwitch:NVLink 专用交换芯片(Switch Fabric),允许 All-to-All Communication,解决直连拓扑的限制;
    • Broadcast:从一个 GPU 向所有 GPU 发送数据,用于模型参数广播(Distributed Training);
    • Scatter:将一个 GPU 数据复制到多个 GPU,用于参数分配;
    • Reduce:将所有 GPU 数据汇总累加到一个 GPU,用于参数收集;
    • All-Reduce:计算所有 GPU 的数据和,并同步到每个 GPU,用于梯度平均(Data Parallel);
    • Gather:将所有 GPU 数据汇总存储到一个 GPU;
    • All-Gather:所有 GPU 共享所有 GPU 的数据,用于模型并行(Tensor Parallel);
    • ReduceScatter:Reduce + Scatter 结合,用于高效梯度传输;

Hopper 架构参数

Blackwell 架构参数

CUDA

  • 入门1入门2官方文档部分官方文档索引内存访问
  • 一个 Grid 对应一次 Kernel Launch;每个 Block 与 唯一一个 SM 对应(一个 SM 可以同时处理多个 Block);线程调度的基本单位是 Warp,一个 Warp 32 个线程;
  • 通过 WMMA API 或 mma.sync PTX(asm volatile) 指令异步调用 Tensor Core(或调用 cuBLAS 和 CUTLASS 库,后者定制化程度更高),分配调度单位是 Warp;

Transformer 优化

  • 常见推理优化指标:

    • Throughput(吞吐量):单位时间内处理 token 数 / 请求数;
    • TPS(Tokens Per Second)/ TPOT(Time Per Output Token);
    • TTFT(Time to First Token)/ First-Token Latency;
    • MFU(Model FLOPs Utilization)模型算力利用率 / HFU(Hardware FLOPs Utilization)硬件算力利用率;
  • 分布式训练:

    • 数据并行(Data Parallelism):将模型复制 n 份,输入不同的小批量数据(mini-batch)并行训练,最后 All-Reduce 合并梯度更新;扩展:DeepSpeed ZeRO-1(优化器状态拆分)/ 2(梯度信息拆分)/ 3(参数动态加载);
    • 流水线并行 Pipeline Parallelism:每个 GPU 负责几层,同时与前后 GPU 交互数据(后、前向传播),相当于纵向切割;需要尽可能减少 Bubble;
    • 张量并行(Tensor Parallelism):单个算子拆到多个 GPU 上,相当于横向切割;与流水线并行统称模型并行(Model Parallelism);
    • LoRA:冻结其余参数,只微调 Attention 权重;使用低秩矩阵(将一个大矩阵拆成两个秩更低的矩阵之和)分解减少存储需求;
  • Attention 相关公式:

    • 前向传播公式:

      a(0)=xz(l)=W(l)a(l1)+b(l),l=1,2,,La(l)=f(l)(z(l)),l=1,2,,LL=L(a(L),y)

    • 反向传播公式( 为逐元素相乘)

      \begin{aligned} \boldsymbol{\delta}^{(L)} &= \nabla_{\mathbf{a}^{(L)}} \mathcal{L} \odot f^{(L){\prime}}\left(\mathbf{z}^{(L)}\right) \\ \boldsymbol{\delta}^{(l)} &= \left(\mathbf{W}^{(l+1)T} \boldsymbol{\delta}^{(l+1)}\right) \odot f^{(l){\prime}}\left(\mathbf{z}^{(l)}\right), \quad l=L-1,\ldots,1 \\ \frac{\partial L}{\partial \mathbf{W}^{(l)}} &= \boldsymbol{\delta}^{(l)} \, \mathbf{a}^{(l-1)T} \\ \frac{\partial L}{\partial \mathbf{b}^{(l)}} &= \boldsymbol{\delta}^{(l)} \end{aligned}

    • Self-Attention( \mathbf{X} \in \mathbb{R}^{n \times d}, \ \mathbf{W}^{Q}, \mathbf{W}^{K}, \mathbf{W}^{V} \in \mathbb{R}^{d \times d_k},\ d_k=d/head): 其它优化:

      \begin{aligned} \mathbf{Q} &= \mathbf{X} \mathbf{W}^{Q}, \\ \mathbf{K} &= \mathbf{X} \mathbf{W}^{K}, \\ \mathbf{V} &= \mathbf{X} \mathbf{W}^{V} \end{aligned}

      \mathbf{A} = \text{softmax}(\frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d_k}})\mathbf{V}, \\ a_{ij} = \frac{\exp\left(\frac{q_i \cdot k_j}{\sqrt{d_k}}\right)}{\sum_{j{\prime}=1}^{n} \exp\left(\frac{q_i \cdot k_{j{\prime}}}{\sqrt{d_k}}\right)}\mathbf{V}

    • Multi-Head Attention:

      \text{head}_i = \text{Attention}(Q_i, K_i, V_i) = \text{softmax} \left( \frac{Q_i K_i^T}{\sqrt{d_k}} \right) V_i

      \text{MultiHead}(Q, K, V) = \text{Concat} ( \text{head}_1, \text{head}_2, \cdots, \text{head}_h ) W^O

  • Attention 优化:

    • KV-Cache:每次新加一个 token(上次的输出), Q 新增一行, QK^T 新增一行(由 Q 新增那行与整个 K 共同算得), A 新增一行(由 QK^T 的最新行与整个 V 共同算得);一般分至少 30% 显存用于 KV-Cache;

    • KV-Cache Reuse:常用 KV Cache 常驻显存,适合 Long System Prompt 场景;

    • PagedAttention:由于序列增长速度快,需要类似 OS 的内存分配机制,设一个块表(Block Table),用于从连续的逻辑地址映射到不连续的物理地址;GIF 演示

    • FlashAttention:合理利用 GPU 内存结构,减少访存次数;计算复杂度保持 O(n^2) 但不存储 O(n^2) 的权重矩阵;尽量将 GMEM 的访存移到 SMEM 中;

    • LinearAttention:利用核函数(Kernel)和低秩分解将时空复杂度均从 QK^TO(n^2) 降到 \phi(Q) \cdot \phi(K)^T = KQ^TO(n^2)

    • StreamingAttention:类似 Sliding Window Attention,将全局注意力限制在常数级别的固定窗口内以达到加速,又尽可能捕捉到全局信息;

  • 其它优化:

    • Continuous Batching(In-flight Batching):因此当一个 Batch 中某个回答结束后,立即加入新的回答以填满 Batch;
    • Chunked Prefill / Context:将长 Prompts 拆成多个 Chunks,按批次计算,允许多个请求共享资源;提高吞吐量、降低推理延迟、节省计算资源,避免某个 Generation Phase 的 token 由于进来一个超长 Context 而被阻塞;对显存大小要求更高、访存代价更高;
    • Grouped-Query Attention:多个 Head 共享一组 Key 和 Value;Multi-Query Attention:所有 Head 共享一组 Key 和 Value;
    • Split-K:矩阵乘法维度分割;
    • 模型压缩:量化(Quantization)、蒸馏(Distillation)、剪枝(Pruning)、二值化(Binarization);
    • 训练部分:Training Resiliency
    • Batch GEMM:允许多个形状相同的矩阵对并行计算,在单一 Kernel 内完成;Grouped GEMM 允许不同形状,限制更少;

量化

  • 混合精度训练:

  • DeepSeek-V3 中的混合精度:

  • 作用:保精度加速、省显存;指标:性能、精度、显存;

  • TRT-LLM 会将量化(Quant)、重量化(ReQuant)和反量化(DeQuant)均作为算子;

  • 量化参数:缩放因子(scale)和零点(zero-point);

  • 各类方法:

    • 训练时量化(Quant Aware Training):插入伪量化算子(Fake Quant,作用通常是检测数据极差以确定缩放因子、阶梯化高精度数据),作为高精度模型一部分参与微调,然后导出为低精度;这些伪量化算子在推理时直接做常量折叠;
    • 训练后量化(Post-Training Quantization):离线静态、动态量化:用一些实际数据分布模拟输入用于校准(Calibration),使用 KL 散度(计算真实激活值分布与预测分布的相对熵差异)计算量化参数;静态量化参数固定,动态量化则会根据输入确定 scale 和 zero-point;目前推理引擎多使用静态;
    • 激活感知量化(Activation-aware Weight Quantization):静态,根据不同通道在前向传播中的影响(离群值、关键 Channel 等),针对性调整策略,仅量化影响小的参数;
    • GPT-Q:静态,采用梯度下降,追求全局最优化,可以追求极低比特 4bit;
  • 量化粒度:

    1. Per-token,Per-channel,Per-tensor;
    2. Block(块级量化) + Double(双重量化) + Dynamic Quant(根据输入动态量化);

NVIDIA AI Software

Cutlass

  • CUDA 矩阵乘法Cutlass 入门Tensor Core 利用

  • 可以做 Kernel Fusion,直接成果 Fused Multi-Head Attention;

  • 流水优化:在前一个 Thread 做运算的同时做后一个 Thread 的 Data Loading(Prefetch),从而掩盖访存 Latency;

    • Prologue:Launch,Load Activation,Load Weights;
    • Main Loop(Tensor Core 使用阶段);
    • Epilogue:Store Activation;
  • Cutlass 2.x:针对 Pre-Hopper archs(Ada sm_89 及之前);

    • 计算完成后,在 Shared Memory 中做 Epilogue 和 Alignment,对齐后使用更大访存 API 写回 Global Memory;
    1
    2
    3
    4
    using ThreadBlockShape = cutlass::gemm::GemmShape<128, 128, 64>; 
    using ShapeMMAWarp = cutlass::gemm::GemmShape<64, 64, 64>;
    using ShapeMMAThread = cutlass::gemm::GemmShape<16, 8, 16>;
    using EpilogueOp = cutlass::epilogue::thread::LinearCombination<...>;
  • Cutlass 3.x:增加 Hopper sm_90a kernels;

    • 只需要指定 BlockShaper 和 ClusterShape,WarpShape 会自动推断;

    • 分阶段使用 using 指定 Shape 和操作而非一次性全塞进 cutlass::gemm::device::Gemm 中;

    • 自动根据 Alignment 是否符合 TMA 要求决定使用 UTMALDG 还是 LDGSTS 异步指令,前者只需一个或两个 Warp 做异步 Data Loading,其余 Warp 只做 Tensor Core 计算(在所需数据加载完之后);

    • Kernel Scheduler: Warp Specialized + Pingpong;使用 Barrier 做同步;

    • 绿色代表 Producer Warps(Prologue / Epilogue),蓝色代表 TC Warps(Tensor Core Calculation);

  • CuTe:Cutlass 3.0 提出;

    • Layout<Shape, Stride>:将 n-D 逻辑坐标映射到 1-D 物理索引;Shape 指定逻辑矩阵形状,Stride 按行、列、层的顺序指定相邻两值在物理索引的距离;

    • Tensor<Ptr, Layout>:Layout + 对应物理空间指针;同一个 Ptr 可以组合不同 Layout;

    • 关键数据传输指令:Copy 和 TiledCopy;

TensorRT-LLM

  • TensorRT 不支持多卡、未对 LLM 做针对性适配;

  • FasterTransformer 太上层、C++ 搭建、未发布为产品,不再维护;

  • TRT-LLM 结合了 FasterTransformer 的性能和 TRT 的迭代速度;

  • 对非 LLM 模型可以直接使用 TensorRT Engine 搭建,再与 TRT-LLM Engine 结合成 Pipeline;也支持用户自定义模型架构;

  • 显存结构:Weight、Internal Activation Tensors(Activation)、IO Tensors(KV-Cache);max_num_tokens 参数决定了 Activation 的维度,从而影响剩余的 KV-Cache 大小;

  • 支持量化:

    • FP16 / BF16;

    • Weight-only(INT8 / INT4);

    • W4A16(AWQ / GPTQ);

    • W8A8 SmoothQuant(INT8);

    • FP8(Hopper / Ada Lovelace);

    • KV Cache Quantization(INT8 / FP8);

  • 架构流程:

    • Convert Checkpoint:ONNX 等模型导入,反序列化,转换成计算图和 Inference IR;在此指定 TP、PP Size 和 Data Type;
    • Define Model Architecture:预设各类主流 LLM;
    • Load weight & Build engine:构建一个独立的执行引擎;根据当前 GPU 做 Self-Attention 优化、混合精度量化和反量化、算子融合与图优化、Kernel 级别流水线优化;NVIDIA AMMO 低精度量化校准库;
    • Execution:推理,Summarize,MMLU 评估;

Megatron-LM

  • 多种并行模式混合的分布式训练框架,使用多种并行;
  • 混合精度矩乘:使用 FP16 前、反向传播但主权重仍为 FP32;通过 Loss Scaling(计算梯度前将损失 L 乘上缩放因子 S,更新时再除以 S)避免数值下溢;
  • 引入 Context Parallelism 和 Sequence Parallelism;

Triton

  • 多框架(PyTorch、TensorFlow)、多种请求(实时、批处理、流式)、多种硬件(CPU、GPU)、跨平台的统一推理服务;

  • 可与 K8S 结合以实现 MLOps,提供性能评估等功能;

  • 动态批处理:多个请求打包成一个 Batch(PD 可以被打进一个 Batch);

  • 一卡多用:提供请求调度功能;

  • 版本管理:方便 A/B 测试、在线更新与回滚、日志与监控;

  • 模型仓库、配置文件;

  • API:RESTful 与 gRPC 与客户端交互;

NeMo

  • 在 Triton 基础上新增端到端云上训练、微调与部署:同样包含 Megatron-LM 中带有 Transformer Engine 的核心功能 Megatron-Core;

  • 推理:NeMo Framework Inference 即 TRT-LLM 开发与运行时库;

  • 多模态、结构化数据处理(Curation);简化各项配置;

  • 检测硬件,自动选择合适的超参(Batch Size、Data Type 等);

  • SFT、RLHF 等后训练支持;

  • 企业级安全性;

NIM

  • 与 Triton 同层级或更高层级的版本管理、动态批处理、负载均衡、监控与管理;
  • 原生支持容器化与 K8S 云原生部署,开箱即用(OOTB),更易扩展,更适配多机部署;
  • 与 Triton 一样后端使用 TensorRT-LLM;
  • 解耦【应用层代码】与【模型代码、权重、Infrastructure】;
  • NIM Pod 可与 KServer 等结合;

其它推理框架

  • 推理框架常见架构:

    • Inference IR:将 TensorFlow、ONNX 等格式转换为统一中间表示;
    • Pre Optimization:公共表达式消除、死代码优化、代数简化;
    • Optimization:算子融合(典型:Conv2d+BN+RELU ⇒ CBR)、冗余算子消除、算子替换、常量折叠;
    • Post Optimization:数据格式转换、内存布局优化、重复算子合并;
  • MNN:

    • 针对移动端和嵌入式的轻量推理框架,支持 iOS 和支持 POSIX 的所有系统;

    • 无依赖,深度裁剪:静态库体积 iOS 几 M,安卓几百 K;

    • 组成部分:Converter(前端多种格式转换 → 图优化、算子融合) ⇒ Interpreter(Engine 加载转换后的模型,Backends 支持 Metal、OpenCL、Vulkan 等多种实现)

    • 其它优化:预推理(评估不同计算策略的资源需求,预分配资源),手写汇编低精度运算库、关键算子;

  • MoonCake

    • PD 分离:解耦 Prefill(Context Phase,Compute Bound,计算密集、带宽打不满)和 Decode(Generation Phase,Memory Bound,访存密集、算力不满)两阶段(即分布式内存系统 + 流处理系统);
    • 构建全局 KVCache Pool,实现以 Cache 为中心的调度(KVCache-centric Scheduling);
  • vLLM:

    • PagedAttention:除了解决显存分配问题外,可以更方便地实现部分 Cache 的共用,特别是共用前缀的情况;
    • Continuous Batching,将任务分为 waiting、running、swapped 三类,计算结束的立即释放;