DMI-SGLang-Integration

DMI monitoring integration for an unmodified official SGLang release, pinned to SGLang v0.5.19 (0bcd822377da7b5718e674eaf9c870d349424dd1). SGLang is not forked: the integration is an out-of-tree package that registers through SGLang's general-plugin entry point (sglang.srt.plugins) and its external model package mechanism, and speaks to DMI core only through the stable dmi.api.v1 facade (DMI >=1.2.0,<2.0).

Read docs/v0519-port.md for the port audit, the qualified compatibility cells, and the explicit limits before relying on a mode that is not listed there.

Install

SGLang 0.5.19 is not on PyPI (wheels stop at 0.5.10, which predates the plugin framework), so install it from the source tag. Build DMI core's native backend against the same PyTorch/CUDA as SGLang 0.5.19 (torch 2.13.0 / CUDA 13.0):

git clone --branch v0.5.19 --depth 1 https://github.com/sgl-project/sglang.git
python -m pip install -e sglang/python
python -m pip install -e /path/to/DMI            # DMI core, then `make -C native -j`
python -m pip install -e .                       # this integration

docs/env-setup.md records the exact resolution used for qualification (uv, torch constraints, no Rust extensions, sgl-kernel 0.4.6.post1, flashinfer 0.6.18).

How it plugs in

SGLang boundary (v0.5.19)	DMI action
sglang.srt.plugins entry point dmi	dmi_sglang_integration.plugin:register validates versions and registers the hooks below in every SGLang process
ModelRunner.load_model (around)	create the DMI engine/ring before weights load, remap hf_config.architectures to the DMI-hooked class, register dmi_sglang_integration.models, attach hooks before CUDA-graph capture
ModelRunner.forward (around)	record the raw ForwardBatch layout (rids, seq_lens_cpu, extend/prefix lengths); verify exactly one commit
DecodeCudaGraphRunner.can_run_graph (around)	preflight the padded step; force eager when it cannot fit the ring
DecodeCudaGraphRunner.load_batch / EagerRunner.load_batch (after)	commit metadata with padded / unpadded rows, before replay or forward
PrefillCudaGraphRunner.can_run_graph (around)	prefill always runs eagerly under DMI
Scheduler.dmi_stop_monitoring, Scheduler.dmi_assert_active	RPC targets for Engine.collective_rpc

Hooked models live in src/dmi_sglang_integration/models/ as EntryClass modules (DMIQwen3ForCausalLM for Qwen3ForCausalLM). They are verbatim copies of the upstream SGLang classes with HookPoints inserted; a source-level test pins the upstream basis.

Offline API

from dmi_sglang_integration import DMIEngine

engine = DMIEngine(
    model_path="Qwen/Qwen3-0.6B",
    mem_fraction_static=0.5,
    disable_prefill_cuda_graph=True,
    dmx_model_id="my-run",
    dmx_db_host="localhost",
    dmx_hook_selection="resid_pre,final_ln,token_ids,final_logits",
    dmx_ring_payload_mb=1024,
    dmx_ring_pinned_mb=1024,
)
try:
    outputs = engine.generate(
        prompt=["The capital of France is"],
        sampling_params={"temperature": 0, "max_new_tokens": 8},
        rid=["request-0"],
    )
    engine.dmi_stop_monitoring()                 # authoritative flush
    print(outputs[0]["text"], outputs[0]["dmi_internal"].available)
finally:
    engine.shutdown()

DMIEngine forwards every ServerArgs keyword to sglang.Engine, carries the dmx_* settings to the spawned worker processes as DMX_* environment variables, fails loudly if the plugin did not attach, tags each output with a lazy dmi_internal handle keyed by meta_info["id"], and flushes on dmi_stop_monitoring(). SGLang kills its workers on shutdown(), so call the flush first. Plain sglang.Engine works too when the DMX_* variables are set; SGLANG_PLUGINS=none or DMX_SGLANG_ENABLE=0 disables DMI for a baseline.

Environment variables: DMX_MODEL_ID, DMX_HOOK_SELECTION (sglang-full default = every hook except attention matrices), DMX_DB_HOST / DMX_DB_PORT / DMX_DB_DATABASE / DMX_DB_TABLE, DMX_RING_PAYLOAD_MB / DMX_RING_PINNED_MB (1024 default), DMX_RING_TASK_ENTRIES, DMX_DRAIN_FLUSH_TIMEOUT_US, DMX_GPU_PADDING_STRIP, DMX_CH_*, DMX_SGLANG_PRELOAD_LIBSTDCXX.

Tests

# CPU gates (no GPU, no ClickHouse)
python -m pytest tests -q -m 'not gpu'

# Bounded GPU qualification: stock vs monitored, separate processes, ClickHouse
CUDA_VISIBLE_DEVICES=0 DMI_V0519_PYTHON=.venv/bin/python bash tests/run_v0519_smoke.sh
DMI_V0519_MODEL=Qwen/Qwen3-8B DMI_V0519_MEM_FRACTION=0.8 DMI_V0519_CELLS=graph \
    CUDA_VISIBLE_DEVICES=0 DMI_V0519_PYTHON=.venv/bin/python bash tests/run_v0519_smoke.sh

Model support

Architecture	DMI class	Status
Qwen3ForCausalLM (dense Qwen3 family)	DMIQwen3ForCausalLM	GPU-qualified on Qwen3-0.6B and Qwen3-8B (see the port audit)

Qwen3MoeForCausalLM and every other architecture are rejected at the ModelRunner.load_model boundary, before DMI allocation and weight loading. Upstream device/distributed initialization may already have happened.

Unsupported (rejected at startup)

TP > 1, PP, DP / DP attention, speculative decoding, two-batch overlap, PD multiplexing and disaggregation, LoRA, diffusion LLM, LayerNorm sequence parallelism, torch.compile, the Transformers model implementation, the tc_piecewise decode CUDA-graph backend, non-generation models, and unknown architectures. Prefill CUDA graphs are forced eager rather than rejected. Streaming generate(stream=True) is not supported by DMIEngine.

Create engines sequentially; concurrent engine construction in one process is not supported (configuration transport and SGLang plugin state are process-wide). Parent environment overrides are restored after startup. A failed explicit dmi_stop_monitoring() remains failed on subsequent calls; do not treat shutdown or a repeated stop as evidence that a failed capture was persisted successfully.