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硬件上的特定优化。

市场反响与应用

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

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

需求与优化策略

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

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

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

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

性能优化目标与成果

• 性能优化目标

  • 限制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)

• 结论

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

PD (Prefill-Decode) Serving

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

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

• PD分离: TPOT = Decode总时长 / Decode长度，可在Decode阶段组合更大的BatchSize。

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

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

多样的并行方式

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

• 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。

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

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

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

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

mPnD: 最大化Prefill集群吞吐

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

mPnD: 最大化Decode集群吞吐

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

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

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

性能表现：
右图比较了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。

结论与原因分析：
- 结论：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网卡高效协同。

---

PD Large-EP哈将: EPLB

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

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

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

---

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阶段。

---

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。

参考资料: 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在不同配置下的性能表现。

---

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的原理和性能测试结果。

---

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的流水线工作示意图。

参考资料: 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%。

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

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%。

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

---

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的端到端性能提升。

---

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。

参考论文: 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性能。

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%。

---

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%。

---

IFB MTP (Multi-Token Prediction)

MTP 介绍

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

参考资料：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

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个点，精度几乎无损。

参考资料：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的计算数值。

---

IFB Overlap Scheduler

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

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

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

  • 代价是在“<eos>”已出现时，最后可能会多执行一个无效的解码步骤。

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。

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系统上限
  • 长文专项优化