The project's README from github.com/ProjectDMX/DMI. As one markdown file: README.md.
A decoupled, asynchronous AI-native data observation backend for high-speed LLM inference.
[!IMPORTANT] Seeking Research and Dev. Collaborations. We are actively looking for research collaborators to explore downstream applications built on DMI such as interpretability, speculative decoding, hallucination analysis, distillation, activation steering, and beyond. If you're interested, please contact us.
Project Status — research preview. DMI supports HuggingFace and official vLLM integrations across 20+ dense and MoE model families, including Qwen3, Llama, Gemma, Mistral, GPT-OSS, Phi, and Granite, with early support for emerging architectures such as Qwen3.6, Llama 4, DeepSeek V4 Flash, GLM-5.2, Kimi K3, and MiniMax-M2.7, plus Megatron-LM training. SGLang support is on the way. APIs may change. Contributions, bug reports, and feature requests are welcome.
👀Technical Report Available: https://arxiv.org/abs/2605.11093
We are working to make DMI useful across more backends, more models, and more stages of the model lifecycle.
DMI is a full-feature observability layer for LLM inference and training. It gives real-time access to any internal model state — residual streams, attention patterns, MLP outputs, KV-cache slices, logits — during inference or training, with minimal overhead.
Right now, DMI supports inference through HuggingFace Transformers and vLLM, and training through Megatron-LM. It captures internal tensors through CUDA-Graph–compatible hooks and streams them off the GPU via a dedicated ring buffer to a host-side drain that pushes into a queryable store (or drops them, for transport-only profiling).
If you're:
you need internal visibility without rewriting your model or slowing inference 10×. That's the gap DMI fills.
HookPoint — drop-in observation primitive. Place it anywhere in a PyTorch
model; works under CUDA Graphs and survives torch.compile.Ring² — GPU↔CPU co-designed staging. A dedicated GPU-side payload ring
isolates captured tensors from the KV-cache memory pool; an on-host meta ring
is drained asynchronously.output_hidden_states,
and register_forward_hook. See benchmarks.Captured internals explored in a Jupyter notebook -- attention patterns, residual-stream norms, per-token confidence, and top-k alternatives over one prompt through Qwen3-0.6B.
Source under examples/visualization/.
https://github.com/user-attachments/assets/7aaf73ce-a0e4-4953-ba99-dd78dd36ca52
[!NOTE] New: DMI vs. vLLM Hidden State Extraction. vLLM recently added a native way to save hidden states. We benchmarked it against DMI on Qwen3-4B prefill: DMI captures about 13× more tensor data per token while keeping extraction overhead 15–17× lower at batch sizes 16–32 in a matched extraction-only setup. Read the full comparison:
docs/dmi_vllm_ehs/dmi-vs-ehs.md.
Offline throughput — Qwen3-4B / Llama-3.1-8B / Qwen3-14B on ShareGPT and WildChat, normalized to vanilla HuggingFace (ideal, no observation = 1.0). Red × = out of memory.
Online serving (TPOT) — same models on vLLM, plotted against request rate. DMI tracks the no-monitor baseline; synchronous hook/debug baselines saturate at much lower request rates.
Full setup, additional results, and how to reproduce:
docs/benchmarks.md.
DMI-configurator turns a model plus a few clicks into a validated capture configuration, and estimates what it will cost before you run anything. It needs the optional web dependencies once:
pip install -e ".[ui]"
After that, from a checkout:
make ui
That opens a browser on http://127.0.0.1:8000. Point it at your own model
with make ui MODEL=./Qwen3-8B. make ui runs from the source tree without
installing DMI itself, but the [ui] extra is required either way — the
server raises a RuntimeError naming this command if it is missing.
With DMI installed, it is a normal command:
dmi ui ./Qwen3-8B
dmi ui accepts a model directory, a config.json, a Hugging Face model id,
or a DMI descriptor YAML, and picks the descriptor out of the current directory
when there is exactly one. Resolving a model directory or a bare Hugging Face
model id needs transformers (pip install -e ".[hf]"); pointing at a
config.json directly does not. See
the configurator plan for the design.
Start with the core installation guide, then choose the HuggingFace, vLLM, or Megatron-LM path depending on the runtime you want to inspect. The project currently supports installation from source. Use a separate environment and checkout for each backend, and install only the integration you need. The snippet below shows the minimal vLLM entry point. The integration checkout connects DMI to an unmodified official vLLM installation.
import os
# Required for the current vLLM integration
os.environ["VLLM_USE_V2_MODEL_RUNNER"] = "0"
from vllm import LLM, SamplingParams
llm = LLM(
model="Qwen/Qwen3-0.6B",
worker_cls="dmi_vllm_integration.worker.DMXGPUWorker",
additional_config={
"dmx_hook_selection": "vllm-full",
"dmx_null_mode": False,
"dmx_db_host": "", # active capture + transport; no persistence
},
)
for o in llm.generate(["The answer is"], SamplingParams(max_tokens=16)):
print(o.outputs[0].text)
# Internal states for every layer have traversed Ring² during the run.
# Configure a sink to persist them. Setting "dmx_null_mode": True
# disables DMI planning, metadata, and payload copying.
| Core installation | Install DMI from source and build the native backend |
| HuggingFace | Run HF generation, monitored generation, and offline benchmark scripts |
| vLLM | Run DMI through the vLLM offline API or vllm serve |
| Megatron-LM | Run DMI during Megatron-LM training |
DMI is an early research system from FrootLab at the University of Maryland, and we welcome contributions from users, researchers, and systems builders. Useful contributions include bug reports, documentation fixes, benchmark reproduction notes, new model integrations, and backend-specific improvements for HuggingFace, vLLM, or Megatron-LM.
@misc{yu2026enablingperformantflexiblemodelinternal,
title={Enabling Performant and Flexible Model-Internal Observability for LLM Inference},
author={Nengneng Yu and Sixian Xiong and Yibo Zhao and Wei Wang and Zaoxing Liu},
year={2026},
eprint={2605.11093},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2605.11093},
}
DMI is licensed under the Apache License 2.0. See the LICENSE file for details.