DMI: Observability Infrastructure for Large Language Model Inference

DMI provides real-time internal state observation capabilities for LLM inference. Through the HookPoint and Ring² architectures, it captures key internal states such as attention patterns, residual streams, and KV caches without modifying the model or significantly reducing performance.

With the widespread application of large language models in critical business scenarios, focusing only on input and output is far from sufficient. Developers and researchers need to deeply understand the internal behavior of models to address the following challenges:

Hallucination Detection and Debugging: When a model generates hallucinations, its internal attention distribution often shows abnormal patterns. By observing these internal states, potential hallucination risks can be identified before output generation.

Interpretability Research: Understanding how models "think" is the core of AI safety research. Information such as attention patterns, hidden state evolution, and MLP activations is crucial for explaining model decisions.

Activation Steering and Behavior Correction: By monitoring internal states in real time, activation steering techniques can be implemented to adjust model behavior without retraining, such as enhancing or suppressing specific types of responses.

Speculative Decoding Optimization: Advanced decoding strategies require access to the internal states of the target model to generate high-quality draft tokens.

Long Text Generation Monitoring: Attention collapse is a common problem when generating long texts, which requires real-time monitoring to detect and mitigate.

The design philosophy of DMI is to provide an asynchronous observation mechanism decoupled from the inference engine, which neither modifies the model architecture nor significantly affects inference performance. Its core architecture includes two key components:

HookPoint is the basic building block of DMI and can be inserted anywhere in a PyTorch model. Its design meets the following key requirements:

• CUDA Graph Compatible: In modern inference engines, CUDA Graph is used to reduce CPU overhead. HookPoint is specially designed to work properly in a CUDA Graph environment.

• torch.compile Friendly: PyTorch 2.0's compilation optimization can significantly improve inference performance. HookPoint is compatible with torch.compile and does not sacrifice the benefits of compilation optimization due to observation requirements.

• Plug-and-Play: Developers only need to add HookPoint to the model definition without modifying the core logic of the inference engine.

Ring² is an innovative data transmission architecture of DMI, specially designed to efficiently transfer GPU internal states to the host side:

GPU-side Payload Ring: A dedicated ring buffer is maintained in GPU memory to store captured tensor data. This buffer is isolated from the KV cache memory pool to avoid mutual interference.

Host-side Meta Ring: A corresponding metadata ring buffer is maintained in CPU memory to asynchronously receive data from the GPU.

This double-layer design enables true asynchronous observation—GPU can continue to perform inference while data transmission proceeds in the background without blocking forward propagation.

DMI can capture various key internal states during the inference process of large language models:

Residual Streams: Input and output states of each layer, reflecting the transmission and transformation of information in the model.

Attention Patterns: Attention weight matrices, revealing which parts of the input sequence the model focuses on when processing the current token.

MLP Outputs: Activation values of the feed-forward network, containing the factual knowledge and reasoning patterns stored by the model.

KV Cache Slices: States of the key-value cache, which are crucial for understanding long text generation and context maintenance.

Logits Distribution: Probability distribution of the output layer, which can be used to analyze the model's confidence and uncertainty.

All these data can be accessed through a unified API and support real-time streaming to query storage or visualization tools.

DMI currently supports two mainstream large language model inference backends:

For users using HuggingFace Transformers, DMI provides a lightweight generation wrapper. Users only need to specify DMI-related configuration options when creating the model to automatically enable internal state capture.