TensorRT-LLM驱动DeepSeek性能极限-协同腾讯联合优化实践

Raccoon Liu : 腾讯大模型推理负责人
朱文熙 : 腾讯开悟平台研发负责人
王猛 : NVIDIA 高级加速计算专家

议程 (Agenda)

  • Introduction (引言)
  • PD serving (PD服务)
  • IFB
  • Summary (总结)

引言 (Introduction)

DeepSeek-R1: 中国自主研发的突破性AI大模型

  • 模型能力强

    • 在数学、代码、自然语言处理等任务上，性能对齐OpenAI o1正式版。

  • 训练开销低

    • 训练成本为同类模型的1/10。

  • 开源免费模型重构生态

    • 采用MIT License，可商用，可模型蒸馏。

  • 技术架构创新

    • Mixture of Experts (MoE)：动态激活特定领域“专家”执行任务，减少冗余计算，对比传统大模型效率提升3倍。
    • Multi-Head Latent Attention (MLA)：大量减少KV-Cache的显存占用量，性能对标传统MHA。
    • Reinforcement Learning (RL)-Driven Training：强化学习与合成数据结合，无需人类标注数据即可实现复杂推理能力。

  • 部署的挑战

    • 技术架构创新带来不同的算子优化。
    • Hopper硬件上的特定优化。
Page 4
Page 4

市场反响与应用

DS（DeepSeek）证明了基础模型创新对应用有极强的拉动作用。

  • DeepSeek发布1个月后，App端DAU冲高到5,000万甩开豆包。
  • 各家应用纷纷拥抱DS，“后DS”阶段更需要独特的竞争优势。
Page 5
Page 5

需求与优化策略

  • DS爆火后，司内AI应用需求激增

    • 公司内部算力需求（Hopper卡数）在25春节后相较于春节前增长了200+%。

  • 混元针对DS做定向优化，提供业内性价比最高的DS推理服务

    • 高推理效率背后，是混元团队的针对性优化策略：
      1. 针对性kernel优化并引入高效的DeepSeek开源算子 (DeepGEMM / FlashMLA / DeepEP)。
      2. 支持INT4量化 + MTP并行解码。
      3. 支持PD分离并新增Attention-DP + MOE-EP 并行策略，降低卡间通信量。
Page 6
Page 6

性能优化目标与成果

  • 性能优化目标

    • 限制TPOT（Time Per Output Token）在50ms，追求最高吞吐（QPM: Query Per Minute）。

  • 优化路径与成果

    • PD 模式 -> 202 QPM

      • Runtime optimization (xPy0/dymmy request/...)
      • Large EP
      • DeepEP
      • Load balancing
      • Two batch overlap

    • IFB 模式 -> 137 QPM

      • Kernel-level optimization: W4A8 (kernel opt/quant+permute fusion/ktile opt), DeepGEMM (jit/swap AB), Flash-MLA (FP8 Flash-MLA and FP8 KV-cache/sparse MLA/Hopper-Style FMHA)
      • Runtime-level optimization: MTP (eagle-based opt), Runtime optimization (CPU bubble/chunked moe for long context/kv-cache reuse)

  • 结论

    • 通过腾讯和英伟达的深入合作，达到了业界领先的性能。
Page 7
Page 7

PD (Prefill-Decode) Serving

PD分离：P&D各自选取最优部署方案

PD分离的核心思想是将Prefill（提示处理）和Decode（生成）两个阶段分开部署和优化。

  • PD分离: TPOT = Decode总时长 / Decode长度，可在Decode阶段组合更大的BatchSize。
Page 9
Page 9
  • TTFT (Time To First Token): 针对Prefill做优化，同时减少KVCache的长度，更多的显存空间留给cache reuse。
  • TPOT (Time Per Output Token): 不受Prefill的打断影响，打到满足SLO的更高有效batch，充分利用算力，彻底解决吐字卡顿问题。
  • 异构支持:
    • 支持Prefill、Decode并行方式最优化，充分发挥各种并行方式的优势。
    • 支持Prefill、Decode卡型异构，充分利用不同卡型算力、带宽优势。

PD系统本身不产生性能提升，精细化的分阶段优化方案是性能提升源泉。

多样的并行方式

PD架构支持多种并行方式，以应对不同阶段的瓶颈。

Page 10
Page 10

  • Tensor Parallel (张量并行)

    • 优势: 使用简单，计算并行度高。
    • 缺点: 权重无法TP切分时重复存储 (如GQA、MLA的KV_Head)；All-Reduce通信量大，不适合跨机部署。
    • 适用场景: 多机多卡并行推理广泛适用。

  • Expert Parallel (专家并行)

    • 优势: 有效提升并行度，降低单卡显存带宽压力。
    • 缺点: 需要处理好专家的动态负载均衡问题；大EP跨机通信依赖高性能RDMA网络带宽，跨机部署通常是Memory bound。
    • 适用场景: 更适合 “Decode” 阶段，减轻Memory bound。

  • Data Parallel (数据并行)

    • 优势: 避免重复计算 + 存储；有效结合EP做大并行策略。
    • 缺点: 依赖dummy request补齐以支持cuda graph；需要做好请求调度以及均衡。
    • 适应场景: 与TP结合，避免切分部分重复计算算子和存储。

多维混合并行策略：TP+ADP+EP

针对Prefill和Decode阶段的不同瓶颈，采用不同的混合并行策略。

  • 针对 Prefill: 计算瓶颈，主要采用大TP + 小EP，降低首字时耗，减少跨机通信量。
  • 针对 Decode: 访存瓶颈，主要采用大DP + 大EP，减少单卡访存压力，充分利用显卡算力，趋近Compute bound。
Page 11
Page 11

计算通信Overlap: 异步请求级别Overlap

  • 异步传输KVCache，与当前Iteration forward step overlap，不影响吞吐。
  • RDMA高效传输，单个Forward step周期内完成，ITL有保障。
  • 调度层面，CPU & GPU协同，适配P+投机采样。
Page 12
Page 12

计算通信Overlap: Layerwise传输计算Overlap

  • 核心思想: 层间计算与KVCache传输Overlap，有效降低长Seqlen场景cache传输时耗。
  • 策略选择:
    • 长文 (ISL>8k) 时选择Layerwise传输。
    • 短文 (ISL<8k) 时整体传输，避免大量小包，并做小包合并传输。
Page 13
Page 13

mPnD: 最大化Prefill集群吞吐

  • P节点策略: 结合Cache reuse、Chunk context的调度策略，确保更优的TTFT。
  • Chunk调度: CHUNK调度时的请求长度排序策略，同一批到达的请求，按长度排序优先调度短Prompt，整体TTFT达到最优。
    • 这能有效缓解请求队列拥塞，使CHUNK调度更均衡。
Page 14
Page 14

mPnD: 最大化Decode集群吞吐

  • D节点调度策略1: D节点结合输入/当前decode长度的调度策略。根据统计平均输出长度，优先按剩余槽位调度，根据平均输出长度剩余槽位决策是否扩容。
  • D节点调度策略2: D节点承担部分短Prefill的Workloads，不做分离调度。收益是减少不必要的KVCache传输（该策略仅在P&D同构切分部署时生效）。
Page 15
Page 15

PD mPnD：最大化整个服务系统

下图展示了mPnD系统的整体架构。系统分为接入层、推理协调器（Inference Coordinator）、预填充节点（Prefill Node）、解码节点（Decode Node）以及远程KV缓存中心（Remote KVCache Center）。请求流程、缓存传输和控制信息在各组件间流动。

Page 16
Page 16

性能表现：
右图比较了PD与非PD（non-PD）的性能。结果显示，在有效ITL（in-flight tokens limit）区间，PD的吞吐明显更优，大幅提升了Goodput。在20~25 tokens/s的用户请求区间，PD有效提升了30~40%的Goodput。

PD Large EP: 理想很丰满，现实略显骨感

  • 理想情况：随着专家并行（EP）规模和并发度的线性增加，处理耗时应随之下降。
  • 现实情况：实际测试中，耗时急剧上升。

如下表所示，从EP=8（并发128）增加到EP=16（并发256），虽然MoE部分的耗时从37.86ms下降到36.79ms，但整体（All）耗时从71.31ms显著增加到95.64ms。

Page 17
Page 17

结论与原因分析：
- 结论：MoE-Gemm部分的耗时下降符合预期。
- 原因：虽然显存带宽压力降低，但当矩阵维度M小于32时，耗时基本不变。右侧表格详细列出了不同参数下gemm1和gemm2的性能指标，验证了这一判断。

解决方案：DeepEP + EPLB (Large-EP哼哈二将)
下图展示了结合DeepEP和EPLB的Transformer Layer架构，其中MoE Layer采用320路专家并行。通过采用DeepEP的Dispatch+Combine通信模式，可以实现注重处理时延、小消息友好。

PD Large-EP哼将: DeepEP

本页介绍了DeepEP在TRMT GPU通信优化基座上的应用，旨在提升大规模专家并行（Large-EP）场景下的通信效率。

  • 常规内核 vs. 低延迟内核：通过两组实验数据对比了使用NVLink和RDMA转发的常规内核与使用纯RDMA的低延迟内核。数据显示，在不同专家并行（EP）和分派（Dispatch）规模下，纯RDMA的低延迟内核在带宽上表现更优。
  • TRMT GPU通信优化：DeepEP通过拓扑感知的多QP建链和iBGDA驱动的多信道并行传输，在RoCEv2网络环境下性能提升超过100%，在IB网络环境下性能提升30%，实现了突破性进展。
  • 腾讯增强版TRMT-DeepEP架构：下图展示了该架构，其核心在于通过iBGDA实现多Channel数据传输，并优化了Channel与QP的映射、WQE下发和Doorbell通知等环节，最终与RDMA网卡高效协同。
Page 18
Page 18

PD Large-EP哈将: EPLB

本页介绍了EPLB（Expert Parallelism Load Balancer），一个用于解决MoE中专家负载不均问题的负载均衡器。

  • 核心功能：
    • 专家负载调度系统以及冗余专家机制。
    • 解决AllToAll通信中的木桶原理短板问题。
    • 使MoE部分达成线性扩展能力。

下图展示了EPLB的工作原理，负载均衡器（Load Balancer）将不同的专家（Expert）动态分配到各个GPU上，以实现负载均衡。右侧的两个图表直观地对比了负载均衡前后的效果，均衡后各Layer的负载更加均匀。

Page 19
Page 19

PD Summary

本页对PD（PagedAttention with Decentralization） serving方案进行了总结。

  • 性能优化总结

    • 针对不同阶段的设计不同并发方式：Prefill采用TP+EP并行，Decode采用DP+EP并行。
    • 通过Overlap隐藏PD之间的调度、传输和计算耗时。
    • 在mPnD大规模PD系统下，通过智能调度优化保障系统高负载下的TTFT和TPOPT。
    • 通过DeepEP+EPLB达成Large EP模式下的超高吞吐。
    • 最终实测吞吐超过200+ QPM。
    • TODO:
      • 完善PD系统容灾能力，实现自愈能力。
      • 支持异构PD系统。
      • 持续优化PD系统下的kv store性能。

  • PD优势

    • TTFT和TPOT互不影响。
    • Prefill和Decode使用不同的并行策略。
    • 整体系统的吞吐天花板被持续拔高。

  • PD劣势

    • 多机部署、机器资源需求大，不方便管理。
    • 系统容灾能力更加复杂。

IFB (In-flight-Batching)

IFB Introduction

本节介绍In-flight-Batching (IFB) 技术。

  • In-flight-Batching(IFB)简介
    IFB通过动态聚合请求，显著提升硬件利用率与吞吐量，避免了因请求长度差异导致的资源浪费。下图对比了静态批处理（Static batching）和动态批处理（In-flight-Batching）的差异，后者能更有效地填充计算资源。

  • IFB部署DeepSeek-R1的特点

    • Prefill阶段和Decode阶段在相同的GPU节点上。
    • Prefill阶段和Decode阶段使用相同的并行策略。

  • IFB部署优势

    • 轻量化：易部署，W4A8的DeepSeek-R1模型使用单个Hopper节点即可部署。
    • 可迁移：IFB优化点可迁移到PD Serving。

下图展示了IFB中的Prefill阶段和Decode阶段。

Page 22
Page 22

IFB DeepGEMM

DeepSeek-R1架构与核心算子

模型的核心是Transformer Block，其中包含了Feed-Forward Network和Multi-Head Latent Attention (MLA)。下图展示了其详细结构，包括Router和Shared Expert的机制。

  • 数学表达与算子
    页面右侧给出了FFN和MLA的数学公式，并指出了其在IFB中对应的FP8优化算子：
    • FFN部分对应 deepgemm.grouped_gemm_fp8。
    • MLA部分对应 Flash-MLA 和 deepgemm.gemm_fp8。
Page 23
Page 23

参考资料: https://arxiv.org/pdf/2405.04434

优化点与性能

本页详细介绍了DeepGemm相对于TRT-LLM旧版实现的优化点。

  • DeepGemm优化点

    • JIT(Just-In-Time)生成Kernel: 运行时生成最佳的Schedule。
    • TMA指令写回结果矩阵。
    • TMA multicast进行输入广播。
    • 所有形状输入开启setmaxreg: 没有JIT时，使用寄存器数量太多，开启setmaxreg是负优化。
    • 启动了更大的Block Tile Size(e.g., 256x128)。
    • Thread Block Swizzle提高L2 Cache命中率。

  • 性能指标

    • 对比TRT-LLM的旧版实现，DeepGemm的平均性能提升了约1.5倍。
    • 在ISL/OSL=3.5K/1K时，Linear耗时占总体的约20%。

  • 集成到TRT-LLM

    • C++实现轻量化的JIT模块，运行时实时编译最佳Kernel。
    • 支持NVRT C，编译时间缩短10倍。
    • End-to-End性能：吞吐提升3-8%。

右侧图表展示了DeepGemm在不同配置下的性能表现。

Page 24
Page 24

IFB DeepGEMM-SwapAB

  • SwapAB

    • 模型推理中，输入到Linear的M维通常小于N维。
    • 计算FP8 Gemm时，WGMMA指令支持固定的M维度64，而N的维度范围可以从最小8到最大256 (Stride=9)。
    • 矩阵乘法的转置序法则: (AB)^T = B^T A^T。
    • M < 64：交换AB输入，避免WGMMA指令的冗余计算。
    • M >= 64：DeepGemm的原生实现。

  • 性能结果

    • 单元测试: M < 64时，对比DeepGemm原生实现平均有1.3倍加速比。
    • End-to-End: 并发<64时，吞吐提升3-5%。

下图展示了SwapAB的原理和性能测试结果。

Page 25
Page 25

IFB Flash-MLA

  • MLA (多头潜在注意力机制)

    • MLA是传统多头注意力机制（MHA）的增强版本，通过引入潜在向量（latent queries and keys）来降低计算资源消耗。
    • 其核心目标是解决长序列处理任务中的计算复杂度和内存瓶颈问题。

  • Flash-MLA

    • Flash-MLA是为Hopper GPU设计的高效MLA解码内核，专门针对可变长度序列服务进行了优化。
    • 其设计灵感来源于FlashAttention 2&3和Cutlass项目。
    • 核心特性:
      • 协作线程束组: 通过协同两个不同的工作线程束组（Warp Groups）实现并行加速。
      • 精细化流水线配置: 优化指令流水以最大化硬件利用率。
      • 复用TRT-LLM工作流: 兼容TensorRT-LLM的模型运行流程，并针对性地加入了对MLA专属的专属优化。

下图展示了不同注意力机制的对比以及Flash-MLA的流水线工作示意图。

Page 26
Page 26

参考资料: https://arxiv.org/pdf/2405.04434
参考资料: https://github.com/deepseek-ai/FlashMLA

FP8 Quantization

  • FP8 Flash-MLA 性能优势

    • 潜在KV缓存的压缩: 通过压缩潜在键值（KV）缓存，支持更大的批处理规模，从而实现更高的吞吐量。
    • 生成阶段加速: MLA内核通过FP8算术运算和减少的KV缓存内存访问来加速。
    • 量化损失: 在通用基准数据集上的评估显示，FP8量化具有近乎无损的模型效果。

  • 支持的功能

    • FP8 KV缓存压缩以及Q的动态量化。
    • Per-Tensor float8 e4m3量化。

  • 性能指标

    • 单元测试: FP8格式的Flash-MLA在Hopper GPU上的计算性能约为BF16的2倍。
    • 端到端测试: 在IFB压力测试中，各档吞吐量提升2-6%。

下表展示了不同量化策略在多个基准测试上的表现，右侧图表为性能对比。

Page 27
Page 27

Overhead Optimization

  • get_mla_metadata开销优化
    • 作用: get_mla_metadata是Flash-MLA的内置核心之一，用于计算沿k序列长度（seqlen）维度分区的flash decoding元数据。
    • 优化前: 每次调用get_mla_metadata会为每个注意力层引入额外延迟，而实际所需每层metadata的seqlen和blocksize都相同，所需的metadata可以复用。
    • 优化后: 每个推理step（而非逐层）调用一次get_mla_metadata，并将元数据传递给注意力算子。
    • 结论: 优化后调度开销降低2-3%。

下图展示了优化前后的流程对比以及端到端加速效果。

Page 28
Page 28

IFB Flash-MHA

  • Hopper-style Flash-MHA

    • Prefill阶段的Flash-MHA Kernel由Ampere-style升级为Hopper-style。

  • 性能优势

    • TMA: TMA是Hopper的一个硬件新增功能，它允许应用程序在GPU全局内存和共享内存之间同步传输张量。
    • WGMMA: 支持直接使用共享内存（SMEM）中的数据进行计算，并将结果累积到寄存器。

  • 结论

    • 端到端首字耗时（TTFT）减少2-8%。

下图对比了A100和H100架构下，MHA的数据传输方式，并展示了Hopper-style Flash-MHA的端到端性能提升。

Page 29
Page 29

IFB Flash-MHA: sparse

  • Sparse Attention

    • 通过稀疏化解决长文本处理延迟问题，可以大大减少长文本下的TTFT。

    • 模型效果人工评估:

      • 单轮问答评估：Sparse模型与base模型持平，得分为96.71。
      • 多轮问答评估：Sparse模型脏出base模型1.5%，得分为84.73（vs base模型83.82）。

    • 参考MInference并在Triton上开发了适配的MHA相关内核。

  • 性能指标

    • 如下表所示，在长文本场景（>16k token）中稀疏化带来显著优势，实现15-25%端到端加速。在seq_len为32768时，latency从2.501s降至0.310s，加速比达到8.08x。
Page 30
Page 30

参考论文: Minference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention (NeurIPS'24 spotlight, FoMo @ ICML'24)

IFB W4A8

W4A8 量化

  • 权重 量化：从FP8量化为INT4精度，通常作用于Linear层和MoE层。
  • 激活值 量化：从A16/BF16量化为FP8/INT8精度。

W4A8 vs FP8

  • 显存空间优势

    • 模型权重的显存空间几乎减少50%。
    • 留出更多KVCache和激活值等显存空间，允许更大的Batch Size和并发数。
    • 允许单机8卡Hopper部署，可以避免跨节点引入的额外通信开销。

  • 显存带宽优势

    • 在小EP场景下，相同token数下访存虽较大，通常是memory-bound。
    • 4bit权重能缓解显存带宽压力，带来更好的kernel性能。
Page 31 - Hopper单机八卡部署分析及MoE GEMM延迟图
Page 31 - Hopper单机八卡部署分析及MoE GEMM延迟图

W4A8内核高效实现

  • 激活值使用Per-Tensor取代Group-Wise量化，减少访存和计算。
  • 使用CUTLASS实现高效的W4A8 Grouped GEMMs内核。
  • 加入Profiler在runtime搜索不同shape下的性能最好的内核配置，如Ktile等。

性能结果

  • 单元测试：相比于DeepGEMM FP8实现，W4A8实现有1.2x至1.8x加速比。

量化改进与结果

  • 量化损失

    • 直接采用开源的方式，对整模型采用W4A8量化，会带来能力下降。

  • 量化改进

    • 使用AWQ算法 (Activation-aware Weight Quantization)，对FP8权重进行高效PTQ。
    • 权重 量化使用 128 Group-Wise 的细粒度策略。
    • 只针对权重占比主导地位的MOE部分进行量化，减少其他部分引入量化误差。
    • Activation保留具有更大范围的FP8精度，而不是转为常用的INT8。

  • 结果

    • 模型能力：在通用基准测评数据集中评估显示，W4A8达到几乎无损的效果。
    • 端到端性能：相较于FP8部署，整体吞吐(QPM)提升 60-98%。
Page 32 - LLM Benchmark评估及端到端吞吐图
Page 32 - LLM Benchmark评估及端到端吞吐图

Kernel 融合

  • Kernel介绍

    • Pre-Quant

      • AWQ算法中引入了Per-Channel的Scale，用于降低量化误差。
      • Pre-Quant操作对输入做Scale，并做Per-Tensor量化为FP8精度。

    • Expand-Row

      • MoE中每个Token会激活8个Experts，需要将其复制至不同Experts对应的输入位置。

  • 存在问题

    • 使用AWQ后，引入了额外的Pre-Quant操作，增加额外的耗时。
    • 由于输入量化过程放置于Pre-Quant，Expand-Row输出需从FP8改成BF16。

  • 优化手段

    • 将Pre-Quant融合至Expand-Row中，在写回Global Mem前，对输出进行Scale和量化。
    • 该融合不仅去除了重复的Global Mem读写，而且量化后写回可省去一半的访存量。

  • 性能结果

    • 单元测试：融合Kernel性能提升至原来的3.7x。
    • 端到端：优化后整体QPM 平均提升 3%。
Page 33 - Pre-Quant融合前后端到端吞吐变化图
Page 33 - Pre-Quant融合前后端到端吞吐变化图

IFB MTP (Multi-Token Prediction)

MTP 介绍

  • MTP旨在辅助DeepSeek-V3的训练，提升模型的准确性。
  • MTP也可以作为投机采样用于推理阶段，提升生成速度；而且MTP模块与主模型同步训练，接受率更高。
  • 基本原理是通过顺序预测额外的多个Token，并在每个预测深度保持完整的因果链。
Page 34 - MTP 架构图
Page 34 - MTP 架构图

参考资料：DeepSeek-V3 Technical Report https://arxiv.org/pdf/2412.19437

MTP 实现

  • Vallina模式

    • 支持多个MTP模块使用不同的权重。

  • Eagle模式

    • 只支持多个MTP模块共用一份权重，顺序执行以预测多个draft token。
    • 兼容其他TRT-LLM特性（如CUDA Graph/Overlap Scheduler/PD等），以提升性能。

验证与接收流程

  • 每个draft token都会通过main model获得对应的logits。
  • 基于贪心策略 (Top-1) 对这些logits进行采样，得到"expected" tokens。
  • 然后将这些"expected" tokens与draft token进行比较。

示例

  • Config: 5 MTP layers
  • Golden token generated by the previous target model: G
  • Draft tokens: ABCDE
Page 35 - MTP 验证与接收流程示例
Page 35 - MTP 验证与接收流程示例

MTP 部署

  • MTP接受率

    • MTP next_n = 1 接受率为1.7-1.8，与数据集相关。
    • MTP next_n = 2 接受率为2.1左右，在并发较低时，可进一步提升吞吐。

  • Relaxed acceptance

    • Strict acceptance

      • 默认开启，意味着draft token与main model基于top-1策略采样的token完全一致时，才予以接受。

    • Relaxed acceptance

      • 对于reasoning model（如DeepSeek R1），生成过程可能包含两个阶段：思考阶段和实际输出阶段。在思考阶段，只要draft token出现在候选集中即可被接受。该候选集是基于logits的TopN和概率阈值生成的。
      • TopN: 从logits中采样概率最高的TopN个token。
      • 概率阈值: 在TopN候选集中，仅保留概率大于top1's probability – delta B的token，构成最终的候选集合。
      • 接受率提升10个点，精度几乎无损。
Page 36 - Relaxed acceptance 对精度的影响
Page 36 - Relaxed acceptance 对精度的影响

参考资料：https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/deepseek_v3#multi-token-prediction-mtp

IFB CUDA Graph

  • 对于PyTorch workflow至关重要，即使是对于大模型和高吞吐量场景。

  • Decoding阶段有严重的 kernel launch bound问题

    • 更糟糕的是，CPU端代码执行和计算开销较大。
    • PyTorch的调用栈较深：Python -> C++。
    • TRT-LLM使用pybind进一步加剧了问题。

  • 结果导致kernel launch之间存在较大间隙。

解决方案：使用CUDA Graph捕获kernel launch过程，并在运行时replay。

  • 限制：模块及其子模块的forward()函数必须是graph-capturable。
    • 即不能有CPU与GPU的同步操作，也不能依赖CPU的计算数值。
Page 37 - CUDA Graph开启前后性能对比
Page 37 - CUDA Graph开启前后性能对比

IFB Overlap Scheduler

  • 然而，即使使用CUDA Graph，我们仍需在每个decoding步骤后进行同步。

    • 同步操作是为了判断是否已达到“<eos>”，从而决定是否终止解码循环。
    • 问题在于，解码步骤之间存在CPU端开销，导致GPU处于空闲状态。

  • 解决方案是将target model + verification logic + draft model 合并到一个CUDA Graph中，并在当前步骤仍在运行时启动下一解码步骤。

    • 代价是在“<eos>”已出现时，最后可能会多执行一个无效的解码步骤。
Page 38 - Overlap Scheduler 工作流程图
Page 38 - Overlap Scheduler 工作流程图

The credit for the overlap scheduler idea belong to the SGLang team: https://lmsys.org/blog/2024-12-04-sglang-v0-4/

IFB Runtime optimizations

  • 优化目标

    • 移除CPU端空隙。

  • 观察现象

    • 对于cudaMemcpyAsync H2D，如果CPU memory是pageable，会引入一个额外的cudaStreamSynchronize，导致overlap scheduler失效，出现巨大的CPU空隙。

  • 解决办法

    • 把CPU memory从pageable改成pinned。
Page 39 - 使用 Pinned Memory 前后性能对比
Page 39 - 使用 Pinned Memory 前后性能对比

A guide on good usage of non_blocking and pin_memory() in PyTorch

IFB Memory optimizations

  • 相比于TensorRT backend，PyTorch的显存管理更加灵活，比较隐晦和难以分析。

  • 更精准地计算kv-cache memory

    • 使用真实的pyExecutor估算runtime activation memory。
    • 先计算CUDA Graph memory pool size再分配kv-cache memory。
    • 反向capture CUDA Graph来复用memory pool。

  • Chunked MoE

    • 显著降低MoE layer所需workspace，尤其是64k/128k等超长文。

  • 修复显存泄漏

    • Reset planned states to avoid memory leak in TrtllmAttentionWrapper。

  • 复用显存

    • Reuse Rope params between decoder layers。
    • Reuse lm_head layer between main model and MTP modules。

IFB Misc kernel optimizations

  • TopK算子融合

    • DeepSeek V3/R1采用分组专家机制，首先筛选出得分较高的专家组，然后从这些组中选出排名前8的专家。
    • 在原生PyTorch实现中，相关操作涉及18个kernel。而通过算子融合后仅需2个kernel，能够实现约10倍的kernel-level加速。

  • 优化Quantize kernel

  • 使用Grouped RMSNorm

  • 使用Ring AllReduce 替代 Tree Allreduce

  • 设置冗余CUDA stream来overlap modules

    • MOE: shared and routed experts
    • MLA: Q norm, and KV norm parallel
    • MLA: Q up, and K rope/nope concat

IFB Summary

  • 性能优化总结

    • 集成了DeepGEMM和Flash-MLA，并在开源基础上进一步优化。
    • 将MoE无损量化为W4Afp8格式，实现单机部署，避免跨机通信开销。
    • 支持了MTP，并扩展到next_n=2以及relaxed acceptance rate。
    • 对PyTorch runtime做了多项优化，去除CPU端空隙。
    • 最终吞吐达到137 QPM。
    • TODO项
      • 开启AllReduce+RMS_Norm算子融合。
      • 开启torch.compile。

  • IFB优势

    • 单机部署，对机器资源需求小，方便管理。

  • IFB劣势

    • Prefill latency影响TPOT。
    • Prefill和Decode只能使用相同的并行策略。

总结 (Summary)

大模型Infra任重道远

  • PD场景

    • 优化：多维混合并行策略、调度 & 传输 & 计算Overlap、智能调度、Lage EP (DeepEP + EPLB)
    • 场景：公共API等大流量场景

  • IFB场景

    • 优化：DeepGEMM、Flash-MLA、MoE W4A8、MTP、Runtime优化
    • 场景：私有化部署等小流量场景

  • TODO

    • Kernel持续优化
    • TBD对计算、通信、调度的精细化Overlap
    • 异构PD系统进一步提升PD系统上限
    • 长文专项优化