为大语言模型推理提供高性能且灵活的模型内部可观测性

# 实现大语言模型推理中高性能且灵活的内部可观测性

来源: https://arxiv.org/html/2605.11093
Nengneng Yu* Sixian Xiong* Yibo Zhao Wei Wang Zaoxing Liu
马里兰大学帕克分校计算机科学系

###### 摘要

如今的推理时工作负载越来越依赖于及时访问模型内部状态。我们提出了 DMI-Lib，这是一种高速深度模型检查器，它将内部可观测性视为一等系统原语，通过从 Ring2（一种用于捕获和暂存张量的 GPU-CPU 内存抽象）构建的异步可观测性底层，将其与推理热路径解耦，并由策略控制的主机后端导出这些数据。DMI-Lib 允许在丰富的内部信号空间和多样化的推理后端上放置观察点，同时保持服务优化并遵守严格的 GPU 内存预算。我们的实验表明，DMI-Lib 在离线批量推理中仅产生 0.4%–6.8% 的开销，在中等在线服务中平均产生 6% 的开销，与具有类似可观测性功能现有基线相比，延迟开销降低了 2–15 倍。DMI-Lib 已在 https://github.com/ProjectDMX/DMI 开源。

## 1 引言

大语言模型（LLMs）正迅速成为新兴 AI 驱动应用的核心服务底层。随着基于 LLM 的部署规模扩大，模型内部执行数据的**可观测性**——不仅仅是记录输入提示和输出令牌——变得日益重要。越来越多的用例依赖于在模型生成预测时及时访问内部状态，例如，长期以来对 LLM 可解释性的追求[25 (https://arxiv.org/html/2505.11093#bib.bib3)]、通过操纵隐藏状态来引导模型输出的测试时对齐技术[37 (https://arxiv.org/html/2505.11093#bib.bib16), 8 (https://arxiv.org/html/2505.11093#bib.bib11)]，以及利用目标模型内部状态来加速草稿模型推理的投机解码[22 (https://arxiv.org/html/2505.11093#bib.bib15), 23 (https://arxiv.org/html/2505.11093#bib.bib58), 2 (https://arxiv.org/html/2505.11093#bib.bib14)]。此外，激活探针可以以比基于 LLM 的监控器低得多的成本监控高风险交互[24 (https://arxiv.org/html/2505.11093#bib.bib10)]，而最近的幻觉检测器利用隐藏状态的跨层动态，而不仅仅依赖输出[49 (https://arxiv.org/html/2505.11093#bib.bib12)]。

为了启用这些可观测性用例，研究人员一直在现有推理栈中逆向添加额外的提取逻辑。然而，这些努力属于两个令人不满意的类别：（1）基于 PyTorch 提供的扩展机制构建的**临时性**观察，PyTorch 是许多现代 LLM 的标准模型实现框架[31 (https://arxiv.org/html/2505.11093#bib.bib25)]；或（2）与推理引擎绑定的功能，例如引擎提供的 API 或输出接口，其访问语义由引擎本身定义和实现。这两种方法都没有提供完整的解决方案。

基于 PyTorch 扩展构建的方法继承了底层扩展表面的局限性：例如，`register_forward_hook` 通常只在模块边界处暴露结果，而面向 Python 的扩展可能会引入不可忽视的开销。相比之下，与推理引擎绑定的方法紧密耦合于引擎内部，因此其可访问性受限于引擎明确暴露的内容，并且在其他平台上缺乏可移植性。更根本的是，这些方法将可观测性视为与执行耦合的**辅助功能**，而不是具有自身资源和性能约束的系统问题。

在本文中，我们提出将**内部模型可观测性**视为一等公民。我们认为，今天实时内部状态访问缺乏采用并非反映用户需求不足：AI 从业者经常需要从模型中获取更丰富的信号，而是系统基础设施从未设计用于支持这一点的后果。构建一个高性能、实时的可观测性系统，如果成功，可以降低从业者理解和利用大型模型内部洞察力的门槛。理想情况下，这样的可观测性系统应满足三个关键目标：

- **广泛的观测覆盖率**：系统应暴露跨模型和推理栈的语义丰富且可扩展的内部观察点空间（例如，残差流、注意力模式、MLP 激活、KV-cache 片段和辅助控制信号），以便下游用户可以在不进行引擎特定修改的情况下表达多样化的分析和控制任务。
- **低性能开销**：系统应保持高性能模型服务栈的执行和内存契约，即使在连续监控下也只引入少量、可量化的减速，以便在线工作负载可以安全地启用内部状态访问，而不违反延迟或吞吐量目标。
- **硬件约束下的高灵活性**：系统应能够通过显式策略适应其数据采集覆盖范围和速率以适应异构硬件预算（例如，PCIe 带宽、GPU 内存余量、存储吞吐量），在提取资源限制允许的最大可操作内部信号的同时，确保稳定的服务行为。

为了满足上述要求，我们引入了 **DMI-Lib**，这是一个高性能、与模型和推理管道无关的**深度模型检查器（Deep Model Inspector）**，系统地收集、转换和存储任意内部状态以启用下游用例。通过将可观测性与推理热路径解耦，DMI-Lib 不仅支持现有用例如可解释性、测试时对齐和投机解码，还解锁了以前不切实际的新用例，如实时安全监控，其中内部状态的连续提取允许操作人员在生成完成之前检测恶意交互并触发提前停止[53 (https://arxiv.org/html/2505.11093#bib.bib40)]。

特别地，DMI-Lib 引入了 **HookPoint**，这是一种轻量级的、兼容 CUDA 图（CUDA graph）的收集原语，可以放置在 PyTorch 模型中的任意位置，以暴露多样化的内部状态，而无需修改推理引擎主干或依赖引擎特定的 API，从而实现广泛的观测覆盖率。为了保持服务性能，DMI-Lib 通过 **Ring²** 路由捕获的张量，这是一种 GPU-CPU 协同设计的暂存抽象，将数据移动保持在 CUDA 图和专用的 GPU 侧环形缓冲区内，然后通过数据导出器在主机上异步排出数据，在仅引入少量、可量化开销的情况下保持执行和内存契约。最后，为了在严格的硬件预算下稳健运行，DMI-Lib 将可重构的 hook 与运行时策略管理器相结合，强制执行完整或尽力而为的导出策略，适应互连和内存约束下的数据速率和保真度，同时保持服务延迟和稳定性。

我们通过由 1.1 万行 CUDA、C++ 和 Python 代码组成的实现，将 DMI-Lib 与 Hugging Face[44 (https://arxiv.org/html/2505.11093#bib.bib4)] 和 vLLM[20 (https://arxiv.org/html/2505.11093#bib.bib1)] 集成。我们的评估显示，DMI-Lib 在离线批量推理中仅产生 0.4–6.8% 的开销，在在线服务中约为 6%，与具有可比可观测性的各种基线相比，延迟降低了 2–15 倍。

本文的其余部分组织如下。第 2 节 (https://arxiv.org/html/2505.11093#S2) 提供 LLM 推理的背景并论证专用可观测性系统的必要性。第 3 节 (https://arxiv.org/html/2505.11093#S3) 到第 4 节 (https://arxiv.org/html/2505.11093#S4) 识别挑战并提出解决这些挑战的设计。第 5 节 (https://arxiv.org/html/2505.11093#S5) 到第 6 节 (https://arxiv.org/html/2505.11093#S6) 讨论实现并进行全面评估。论文在第 8 节 (https://arxiv.org/html/2505.11093#S8) 讨论相关工作，并在第 9 节 (https://arxiv.org/html/2505.11093#S9) 得出结论。

## 2 背景与动机

在本节中，我们首先描述 LLM 推理以及对高性能和灵活内部可观测性系统的需求。然后，我们讨论现有工具及其在启用实时监控能力方面的局限性。

### 2.1 Transformer 模型与推理基础设施

现代 LLM 建立在 Transformer 架构之上。在标准的密集、仅解码 Transformer 中，主要组件包括嵌入层、Transformer 层堆栈（每个层包含自注意力和 MLP 模块）以及带有归一化层的残差连接[41 (https://arxiv.org/html/2505.11093#bib.bib44)]。最近的模型通过架构变体扩展了这种密集设计，如混合专家（MoE）层[16 (https://arxiv.org/html/2505.11093#bib.bib47), 35 (https://arxiv.org/html/2505.11093#bib.bib46)] 和修改后的注意力或序列处理块[5 (https://arxiv.org/html/2505.11093#bib.bib29), 13 (https://arxiv.org/html/2505.11093#bib.bib48)]，进一步增加了架构多样性。

模型架构通常在 PyTorch[31 (https://arxiv.org/html/2505.11093#bib.bib25)] 中实现，它将神经网络组件表示为 `nn.Module` 对象。每个 `nn.Module` 编码计算图的一部分，例如注意力块、MLP 块或整个 Transformer 层，组合这些模块既定义了模型的结构，也定义了阶段之间的数据流。这种可执行图视图使开发者能够以统一的抽象表达复杂的模型逻辑，同时将低级内核选择和图优化留给框架。

LLM 推理系统[20 (https://arxiv.org/html/2505.11093#bib.bib1), 51 (https://arxiv.org/html/2505.11093#bib.bib43), 29 (https://arxiv.org/html/2505.11093#bib.bib51), 15 (https://arxiv.org/html/2505.11093#bib.bib52), 11 (https://arxiv.org/html/2505.11093#bib.bib53)] 为这些模型图提供执行基础设施。它们负责批量处理请求、调度令牌生成、管理 KV-cache 内存、组织设备缓冲区，并调用优化的内核来执行模型描述的计算。现代服务系统如 vLLM[20 (https://arxiv.org/html/2505.11093#bib.bib1)]、SGLang[51 (https://arxiv.org/html/2505.11093#bib.bib43)] 和 Hugging Face TGI[15 (https://arxiv.org/html/2505.11093#bib.bib52)] 将高层模型定义与服务时关注点解耦：其后端将调度、内存管理和内核编排实现为单独的层，可以以最小的集成努力跨快速演变的模型架构重用。

### 2.2 下游应用示例

随着 LLM 的广泛部署，对访问模型内部状态的需求正在迅速增长。我们讨论了像 DMI-Lib 这样的高性能模型可观测性系统的多个动机性下游应用。

#### 机制可解释性。
对于模型可解释性，机制可解释性已经从描述输入-输出行为转变为系统分析残差流、注意力模式、MLP 激活和跨层表示，以揭示模型行为背后的内部机制[9 (https://arxiv.org/html/2505.11093#bib.bib38), 33 (https://arxiv.org/html/2505.11093#bib.bib39)]。在模型执行期间收集和分析大量内部状态已成为模型理解的前沿方法。

#### 在线模型监控。
对于在线监控和可靠性，仅凭提示和最终输出已不足以捕捉风险信号。最近的工作表明，激活探针[19 (https://arxiv.org/html/2505.11093#bib.bib42)] 可以以比基于 LLM 的监控器低得多的成本检测高风险交互[24 (https://arxiv.org/html/2505.11093#bib.bib10)]。同时，一些工作利用跨层隐藏状态动态进行幻觉检测，而不仅仅依赖最终生成结果[49 (https://arxiv.org/html/2505.11093#bib.bib12)]。因此，内部状态成为在线监控和可靠性评估的直接信号源。

#### 推理时控制与对齐。
激活工程[53 (https://arxiv.org/html/2505.11093#bib.bib40), 40 (https://arxiv.org/html/2505.11093#bib.bib41), 48 (https://arxiv.org/html/2505.11093#bib.bib37)] 将推理时激活修改视为一种测试时控制形式；最近的安全对齐工作直接使用中间隐藏状态对候选响应进行排名[8 (https://arxiv.org/html/2505.11093#bib.bib11)]。与提示工程或离线微调不同，这些方法将内部状态置于推理路径的控制逻辑中，使得访问和操作中间表示成为新的系统需求。

#### 服务系统优化。
模型内部状态越来越多地用于服务优化。EAGLE-3[22 (https://arxiv.org/html/2505.11093#bib.bib15)] 将投机解码提升到特征级别，并利用内部特征提高推理效率；后续工作[23 (https://arxiv.org/html/2505.11093#bib.bib58)] 进一步使用导出的隐藏状态训练草稿模型并提高草稿接受率。

### 2.3 现有工具与局限性

在模型服务中，主流系统如 vLLM[20 (https://arxiv.org/html/2505.11093#bib.bib1)] 和 Triton 推理服务器[28 (https://arxiv.org/html/2505.11093#bib.bib60)] 专注于通过暴露 Prometheus 兼容的系统指标[32 (https://arxiv.org/html/2505.11093#bib.bib54)]（例如，GPU 和内存利用率）、结构化日志（例如，记录错误和系统操作）和分布式追踪来实现系统和硬件级别的可观测性。流行的工具如 Nsight Systems[27 (https://arxiv.org/html/2505.11093#bib.bib49)]、PyTorch Profiler[31 (https://arxiv.org/html/2505.11093#bib.bib25)] 和 NCCL RAS[30 (https://arxiv.org/html/2505.11093#bib.bib50)] 覆盖了内核、流、内存和通信级别的执行行为。系统和硬件级别的可观测性已广泛用于性能或可靠性优化[6 (https://arxiv.org/html/2505.11093#bib.bib34), 43 (https://arxiv.org/html/2505.11093#bib.bib36)]。

相比之下，我们的 DMI-Lib 旨在收集模型执行路径期间的内部中间数据。目前提供模型内部运行时可观测性的现有努力仍然稀少。Hugging Face[44 (https://arxiv.org/html/2505.11093#bib.bib4)] 生成实用程序只能返回内部状态的一个子集，包括 logits、attention 和隐藏状态。vLLM 和 SGLang 提供有限的内部提取或插件机制；TensorRT-LLM[29 (https://arxiv.org/html/2505.11093#bib.bib51)] 引入了用于检查的调试张量；而 PyTorch hooks[31 (https://arxiv.org/html/2505.11093#bib.bib25)] 虽然灵活，但官方文档也将其定位为离线调试和分析工具，表明在运行时会产生显著的性能开销。这些接口大多是预定义的模型输出、引擎特定的调试路径或旨在用于分析原型的仪器，而不是用于在线服务的统一和通用可观测性系统。

#### 总结。
随着下游应用数量的增加，适用于模型内部的正确可观测性工具在很大程度上缺失。如图 1 (https://arxiv.org/html/2505.11093#S3.F1) 所示，DMI-Lib 提供