DeepInfer-Lab

Rebuilding Inference Systems from First Principles

English · 简体中文

DeepInfer-Lab is an open experimental project for rediscovering how high-performance inference systems are designed. Starting from a simple, correct runtime, I progressively rebuild the core components of DeepInfer and investigate why each design decision becomes necessary, which bottleneck it solves, and how it changes the system as a whole.

The goal is to reconstruct the design process through experiments: understand where every improvement comes from, explain the remaining performance gap, and explore whether a different design could work better.

Experimental method

Every optimization should answer four questions:

Question	What the experiment should explain
Why?	Which compute, memory, communication or scheduling bottleneck motivates the change?
How?	What changes in the algorithm, data layout, kernel or runtime, and why should it help?
Performance?	What happens to kernel time, end-to-end throughput, latency and memory use?
Trade-offs?	Under which workloads and hardware conditions does it help or hurt, and what does it cost?

Each iteration starts with a correct baseline and a profile, then tests a specific design hypothesis. After the change, I check correctness, measure both kernel and end-to-end performance, and profile again to find the next bottleneck. A faster kernel can move the bottleneck elsewhere. Attention, GEMM, communication, state management and runtime scheduling must therefore be studied together inside a working inference engine.

The experiment record is part of the project: the motivation, implementation choices, correctness checks, benchmark configuration, profiles, results and unsuccessful attempts. Follow the Issues, performance reports and benchmark data to reproduce, verify and challenge the findings. Future experiments will follow the same approach, including publishing negative results and open questions.

Current focus

• Model: DeepSeek-V4.1-Flash. The runtime includes the full V4.1 architecture; the tiny random checkpoint provides a smaller development and correctness-testing setup.
• Measurement: real-model prefill on 4 × NVIDIA B200. The prefill research will keep this GPU configuration fixed, matching the official benchmark's GPU configuration, to study gains from algorithms, kernels and system design. Workload and measurement details are stated with each result.
• Implementation: a readable runtime that can be inspected, changed and measured at every step. It grows as the experiments require; the research is organized around design questions and measured effects.

Quick start

git clone https://github.com/algebra-MCX/DeepInfer-Lab.git
cd DeepInfer-Lab
uv run examples/offline_inference.py

• Environment: on first run uv builds .venv from uv.lock (PyTorch 2.14 + CUDA 13.0, NVIDIA driver ≥ 580) and reuses it afterwards.
• No uv: install it with curl -LsSf https://astral.sh/uv/install.sh | sh, or skip it: pip install -e . && python examples/offline_inference.py.
• Model: tiny-random/deepseek-v4.1-flash-1017 is downloaded automatically. It has 68M parameters and every V4.1 architectural feature, but its weights are random, so the generated text is meaningless.
• Hardware: any CUDA GPU with 8 GB of VRAM; without a GPU only the CPU tests run (see Hardware).

Key results

• A measured optimization history on fixed hardware: prefill throughput increased from 1,084.8 to 4,315.9 tok/s/GPU (3.98×) across the documented 4 × B200 optimization rounds. This is the improvement over this project's own baseline; the comparison with official DeepInfer is reported separately below. Follow report A, report B1 and report C1 for the design decisions and changing bottlenecks.
• The real DeepSeek-V4.1-Flash on 4 × B200, benchmarked the way DeepInfer benchmarks itself. EP4 × DP4, 131k context; request counts, warm-up and measurement windows and metrics match the official run_sweep.sh prefill, and the plot comes straight out of the official plot_pareto.py. The parallel layer is ~40 lines (parallel.py). → Performance
• Checked token by token against DeepSeek's official reference implementation. Quantization and candidate selection match the official TileLang kernels bit for bit; the whole engine stays within 0.05 of the official inference/model.py logits in every scheduling mode: batching, ZOS, chunked prefill and preemption. → Correctness
• The full V4.1 architecture: sliding window + compressed KV, the two-level sparse indexer, hyper-connections, Engram n-gram memory, MoE, and the FP8 / FP4 round trips of the attention state; weights stay in the checkpoint's FP8 / FP4 format.
• DeepInfer's engine design: paged KV blocks + per-sequence state slots, separate prefill / decode steps (chunked prefill, preemption), zero-overhead scheduling (ZOS), deterministic sampling independent of batch composition, request cancellation and stop tokens.

Performance

Real DeepSeek-V4.1-Flash, single-node 4 × B200 prefill

Paired run on 2026-10-22 (with the three optimizations of reports A / B1 / C1): the same four GPUs, checkpoint and input tokens, EP4 / DP4 / CP1, 131,072-token context with 126,976 cached tokens, and 4,096 uncached input tokens per step per rank. The window contains 10 complete steady-state steps.

Engine	Input throughput (tok/s/GPU)	Mean step (ms)
DeepInfer-Lab (current)	4,315.9	949.0
Official DeepInfer + single-node NVLink adapter	36,170.9	113.2

DeepInfer-Lab now reaches 11.9% of this official configuration (an 8.4× gap). This is one paired run, not a three-run median or a result with every optimization enabled. DeepInfer-Lab uses the current validated runtime, including grouped GEMMs, quantization, HC, RMSNorm and SwiGLU changes plus three new optimizations: sparse attention on tensor cores, fused indexer scoring, and MoE combine over local rows with in-place dequantization. Its optional non-expert weight cache is disabled. Both engines use full replay and real routing, with no redundant experts, speculation or decode CUDA graphs. Only prefill is measured; KV memory budgets are not forced to match.

Optimizations and the remaining gap: this result includes sparse attention on tensor cores (report A), fused indexer scoring (report B1) and MoE combine over local rows (report C1). The two largest items now are attention (about 432 ms, 182 ms of it key concatenation) and MoE (410 ms). Per-stage numbers and run instructions are in the benchmark documentation. Tiny-model single-GPU output throughput is a different workload and metric and is not comparable to this table.

Correctness

tests/test_v41.py runs DeepSeek's official reference implementation (inference/model.py, shipped with the V4.1-Flash checkpoint) on the same weights.

Check	Result
FP8 / FP4 fake quant vs. the official TileLang kernels	Bit-exact
Two-level candidate block selection	Bit-exact
Sinkhorn mixing and sparse attention vs. the official kernels	Match within bf16 tolerance
Engine end to end: 4 prompts (37–520 tokens) in one batch; ZOS on / off; chunked prefill 128 and 200; preemption	Logits within 0.05 of the reference at every step (~0.02 max observed); greedy tokens identical

Tokens are compared up to the first step where the reference's own top-2 margin is within tolerance: random-weight logits are flat, and past a near tie both sides compare noise.

pytest -m "not gpu"   # CPU only: scheduler, block manager
pytest                # full suite: needs a GPU + TileLang

The CPU CI runs only the CUDA-free checks; without a GPU the GPU tests show as skipped, which is not a pass.

Architecture

flowchart TB
    A["LLM.generate()"] -->|add_request| B["LLMEngine.step()"]
    subgraph CPU["CPU: scheduling"]
        S["Scheduler<br/>waiting / running queues<br/>prefill or decode per step<br/>chunked prefill, preemption"]
        BM["BlockManager x2<br/>KV blocks, 128 tokens each<br/>state slots, one per sequence"]
        S --> BM
    end
    subgraph GPU["GPU: execution"]
        R["ModelRunner<br/>batch construction, ZOS remap"]
        M["DeepseekV41<br/>n_layers x Block"]
        SA["Sampler<br/>Gumbel-max, seed + position"]
        C[("V41Cache<br/>blocks: compressed KV, indexer keys<br/>slots: window ring, compressor tail")]
        R --> M --> SA
        M <--> C
    end
    B -->|"1. schedule step N"| S
    B -->|"2. execute step N, async"| R
    SA -.->|"3. tokens of step N, read during step N+1"| B

Loading

Inside one V4.1 attention layer:

flowchart LR
    X["hidden x"] --> QR["qr = q_norm(wq_a x)"]
    QR --> Q["q = wq_b qr<br/>n_heads query heads"]
    X --> KV["kv = kv_norm(wkv x)<br/>1 KV head, FP8 round trip"]
    KV --> W[("window ring<br/>last 128 positions")]
    X --> CP["Compressor<br/>pools every ratio tokens"]
    CP --> CKV[("compressed KV, FP4")]
    CP --> IK[("indexer keys, FP4")]
    QR --> IDX["Indexer<br/>top-k compressed entries"]
    IK --> IDX
    Q --> ATT["sparse attention<br/>+ attention sink"]
    W --> ATT
    CKV --> ATT
    IDX --> ATT
    ATT --> O["grouped low-rank output<br/>wo_a, wo_b"]

Loading

Only KV-source layers run a compressor and only index-source layers run an indexer; the layers in between read the source's cache and reuse its top-k choice.

How one step works

step N (enable_zos=True)
  CPU  schedule(N)        decisions use reservations only (cached / in-flight token counts)
  CPU  execute(N)         enqueue inputs, forward, sampling and copy-back; an in-flight input
                          token is gathered on the GPU from step N-1's output buffer (remap)
  CPU  process(N-1)       wait for step N-1, append tokens, finish and free sequences
  GPU  ... step N-1 ... | ... step N ...

V4.1 has two kinds of attention state, and the scheduler allocates both:

State	Lives in	Written by
Compressed KV, indexer keys	Paged KV blocks, 128 tokens each	KV-source layers, one entry per ratio tokens
Sliding-window ring, compressor tail	One state slot per sequence	Every layer / KV-source layers

A preempted sequence releases both and is recomputed later; deterministic sampling reproduces the discarded token exactly.

Multiple GPUs: EP × DP

Every GPU runs its own scheduler and batch; attention and other dense layers are replicated. Routed experts are split evenly across GPUs and the Engram tables by rows: these layers gather the tokens of all GPUs, apply only their local shard, and reduce-scatter the results back. In a single process the communication is a no-op.

Usage

from deepinfer_lab import LLM, SamplingParams

llm = LLM("tiny-random/deepseek-v4.1-flash-1017", max_model_len=4096, kv_cache_memory_gb=2.0)
outputs = llm.generate(
    prompt_token_ids=[[100, 200, 300]],
    sampling_params=SamplingParams(temperature=0.7, top_p=0.9, max_tokens=16, seq_seed=42),
)
print(outputs[0].token_ids)

Option	Meaning
max_num_seqs, max_num_batched_tokens, max_model_len	Batch limits
chunked_prefill, prefill_attempt_interval	Prefill scheduling (DeepInfer's default interval is 70)
kv_cache_memory_gb / gpu_memory_utilization	Cache budget: absolute, or a fraction of free memory
weight_cache_memory_gb	BF16 weight-cache budget for non-expert FP8 projections, in GiB; 0 disables it by default
enable_zos	Zero-overhead scheduling
SamplingParams(seq_seed=...)	Same input + same seed = same output, whatever the batch
SamplingParams(stop_token_ids=...)	Extra stop tokens; LLMEngine.abort_request() safely cancels an in-flight request

With spare memory, LLM(..., weight_cache_memory_gb=0.02) avoids dequantizing weights on every projection. The cache selects projections that fit the budget in module order, retains the original quantized weights, and excludes routed experts and Engram tables. This setting adds 10.56 MiB on the tiny model. It is allocated before the automatic KV budget is calculated, so it reduces the memory available for KV. The benchmark option is bench.py --weight-cache-gb 0.02; gains depend on the model and workload, so caching stays disabled by default.

Reading the code

Question	Start here
How does a request move through the runtime?	docs/TECHNICAL.md, sections 2–5
Where does V4.1's attention state live?	models/deepseek_v41/model.py and section 3 of the technical notes
How do prefill, decode and preemption fit together?	engine/scheduler.py
How is the cache budget computed?	worker/model_runner.py
How is the model split across GPUs?	models/deepseek_v41/parallel.py, MoE and EngramTable
How is correctness checked?	tests/test_v41.py and tools/v41_reference.py

Where each part comes from in DeepInfer:

DeepInfer-Lab	DeepInfer
engine/block_manager.py	csrc/engine/block_manager.hpp
engine/scheduler.py	csrc/engine/scheduler.hpp (_schedule_prefill, _schedule_decode)
engine/llm_engine.py	deepinfer/engine/llm_engine.py (step)
worker/model_runner.py	deepinfer/worker/worker.py (execute_model_step), csrc/engine/zos.cpp
layers/sampler.py	tilelang_ops/random_sampling.py
models/deepseek_v41/	DeepSeek's inference/model.py, rewritten for paged, batched execution
bench_prefill.py	examples/benchmark.py + examples/sweep/run_sweep.sh prefill
tools/v41_reference.py	Golden logits from the reference implementation

The runtime design explains DeepInfer's core boundaries and the order this project implements them in; the technical notes walk through the current code.

Hardware

Use	Configuration	Notes
CPU tests	2 vCPUs, 4 GB RAM, 2 GB disk	Scheduler and block manager; the level the CI uses
Tiny checkpoint (quick start)	Any CUDA GPU with 8 GB VRAM (16 GB recommended), 16 GB RAM, 10 GB disk	uv.lock pins the CUDA 13.0 build of PyTorch, which needs NVIDIA driver ≥ 580
Real DeepSeek-V4.1-Flash	Verified: 4 × NVIDIA B200, ~510 GB disk for the checkpoint	EP4 × DP4, ~137 GiB resident per GPU (weights + 12 GiB KV cache); does not fit on one GPU

There is no CPU path for model execution. Check the license of a real checkpoint before using it.

Roadmap

Core kernel research (planned)

I plan to reimplement FlashMLA, DeepGEMM and DeepEP from scratch, starting with simple, correct research implementations and integrating them into the runtime one component at a time. Each stage will be profiled, optimized and evaluated at both kernel and end-to-end levels.

• FlashMLA: investigate attention data movement, KV access and tiling, and measure how they affect the full attention path.
• DeepGEMM: investigate low-precision GEMM, grouped expert computation and data preparation, including numerical and memory costs.
• DeepEP: investigate expert dispatch/combine, communication and overlap with computation inside the MoE execution path.

These are future research directions. The existing kernels and validated runtime capabilities are described above; rebuilding these components will produce additional experiments and comparisons.

Runtime and model milestones

• DeepSeek-V4.1 model and an official-reference regression suite
• Paged blocks + state slots, chunked prefill, preemption, ZOS, deterministic sampling
• Real V4.1-Flash on multiple GPUs with EP × DP, prefill benchmarked the official way
• Memory for long chunks (≥ 32,768 tokens per GPU): slice the indexer and attention output by query
• Decoder bounded replay (YoCo consumer layers replay only the last 140 tokens during prefill)
• Grouped expert GEMMs and CUDA graphs for decode; remove forward host synchronizations so ZOS overlaps
• DSpark speculative decoding with confidence-scheduled verification
• Cost-model prefill chunking (DeepInfer Q&A)
• End-to-end batch-invariant inference
• Image input (V4.1 ships a ViT; the checkpoint already contains the vision weights)

Contributing and discussion

Feedback, design critiques, profiling investigations, reproducibility checks and code contributions are welcome. I'd also like to organize informal technical study sessions around architecture, kernels, measurements and trade-offs. You can start with the Issues or contact me through my GitHub profile.

I also contribute to vLLM's Model Runner V2 (#63372, #63476) and share updates about my vLLM Promotion Challenge on Xiaohongshu.

Future collaboration and decode research

With enough community interest, I'd like to explore a collaboration with the Inferact team and seek access to additional GPU nodes, ideally a 64-GPU SM100 cluster, to extend the research into decode optimization. Collaboration and cluster access are future possibilities; this is currently an independent educational and experimental project. Prefill research will continue on the fixed 4 × B200 setup.

Project update (2026-10-25): This project will continue to receive updates. I've been feeling unwell recently, so I'm taking a few days off. To keep this repository's commit history clean, I've moved most of my personal optimization code to a private repository. Development will continue, and updates to this public repository will resume after this short break. Thank you for your patience and support.

License

Apache 2.0. Parts of models/deepseek_v41/ are adapted from DeepSeek's MIT-licensed reference implementation; see THIRD_PARTY_NOTICES.

Acknowledgments

Design follows DeepInfer and DeepSeek's V4.1 reference implementation; the project layout is inspired by nano-vllm.