# DeepInfer

DeepInfer is the high-performance inference framework that powers DeepSeek's large-scale production services. This open-source release provides offline inference with the core features needed for serving at scale.

Supported models:

- DeepSeek-V4.1-Flash

Supported features:

- Optimized kernels, model and sampler runtime
- Parallelism: Attention DP / balanced CP, MoE EP + EPLB
- DSpark and confidence-scheduled verification
- SWA bounded replay
- Asynchronous scheduling (zero-overhead batch scheduler)
- Batch-invariant inference
- Deterministic sampling
- Structured outputs

<details>
<summary>Feature details and configuration</summary>

- **Optimized kernels:** [FlashMLA](https://github.com/deepseek-ai/FlashMLA), [DeepGEMM](https://github.com/deepseek-ai/DeepGEMM), [DeepEP](https://github.com/deepseek-ai/DeepEP), [TileKernels](https://github.com/deepseek-ai/TileKernels), [DeepSelect](https://github.com/deepseek-ai/DeepSelect), and custom [TileLang kernels](deepinfer/model_executor/layers/tilelang_ops).
- **DSpark and confidence-scheduled verification:** accelerate decoding with draft tokens and use confidence estimates to schedule verification. Set `DEEPINFER_ENABLE_MTP_VARLEN=1`; `DEEPINFER_USE_NEXTN_DECODE` controls the maximum draft length (default: 5).
- **Decoder SWA bounded replay:** reduce decoder prefill work by replaying only a bounded suffix. Set `DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE=compact` and `DEEPINFER_YOCO_PREFILL_KEEP_TOKENS=140` for the default V4.1 settings; the retained suffix must cover the SWA window plus the maximum backtracking length.
- **Asynchronous scheduling:** overlap CPU batch scheduling with device execution. Enabled by default.
- **Batch-invariant inference:** keep model results identical across different batch compositions. Enable with `DEEPINFER_BATCH_INVARIANT=1` (NVIDIA only).
- **Deterministic sampling:** produce identical outputs for identical input tokens and sampling parameters. Set `DEEPINFER_BATCH_INVARIANT=1 DEEPINFER_ENABLE_MTP_VARLEN=0` and pass a fixed `seq_seed` to `SamplingParams` (NVIDIA only).
- **Structured outputs:** constrain generation to a regex, JSON schema, or EBNF grammar. Pass `so_args={"schema_type": ..., "schema": ...}` to `SamplingParams`; see [examples](examples/_structured_output.py).

</details>

Supported hardware:

- NVIDIA SM100 GPUs
- Huawei Ascend 950 NPUs
- AWS Trainium2 (native torch-neuronx, basic mode — see [DeepInfer on AWS Trainium](#deepinfer-on-aws-trainium))

Required for a complete inference service (not included in this release):

- A KV cache system
- VE/Prefill/Decode disaggregation (most required engine functionality is already implemented)
- Load balancing across data-parallel replicas for both prefill and decode

## Getting started

### Prerequisites

- Linux
- GCC >= 13.0
- Python 3.12
- CUDA 13 (NVIDIA only)
- CANN 9.2 (Ascend only)
- Neuron SDK 2.33 with torch-neuronx 2.14 (Trainium only)

### Build from source

Before building DeepInfer, please create and activate an empty Python 3.12 environment.

``` bash
# Use `ascend` instead for the Ascend platform, `neuron` for AWS Trainium
# Keep it exported in the environment where you run DeepInfer
export DEEPINFER_ACCELERATOR=nvidia
./build.sh
```

### Offline inference

The default checkpoint, [`deepseek-ai/DeepSeek-V4.1-Flash-1017`](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash-1017), is
downloaded into the Hugging Face cache on the first run; set `DEEPINFER_HF_MODEL_DIR=/path/to/DeepSeek-V4.1-Flash-1017` to
run from a local copy instead.

``` bash
python examples/offline_inference.py
```

See [examples](examples/README.md) for details.

## DeepInfer on AWS Trainium

DeepInfer runs on Trainium2 through the native PyTorch `neuron` device (torch-neuronx 2.14, no torch-xla),
structured like the Ascend port: `PlatformNeuron`, the `tilelang_ops/neuron` kernel backend and pure-torch
drop-ins for DeepGEMM / FlashMLA / DeepEP / DeepSelect / TileKernels in `deepinfer/neuron_ops/vendor`.
Design notes and conventions: [deepinfer/neuron_ops/NEURON_PORT.md](deepinfer/neuron_ops/NEURON_PORT.md).

**Status — basic mode.** Eager execution; DSpark/MTP, zero-overhead scheduling, CP, EPLB, YOCO bounded replay and
VL are off. EP > 1 runs on real neuron collectives. A tiny checkpoint in the real FP8/FP4 layout matches the
official reference implementation within its own noise floor, on CPU and on a NeuronCore
([examples/neuron](examples/neuron/README.md)). Performance work (static graphs, NKI kernels) and the remaining
features are tracked in the [roadmap](https://github.com/whn09/DeepInfer/issues/19).

On a Neuron DLAMI (trn2.48xlarge, SDK 2.33):

``` bash
bash tests/neuron/setup_trn2.sh          # venv on the DLAMI torch-neuronx env, build, tiny checkpoint + golden
bash tests/neuron/run_all.sh cpu         # host-only tests
CORES=0-7 bash tests/neuron/run_all.sh device
```

or build by hand in a venv layered on the DLAMI's torch-neuronx environment (the `torch-neuronx` wheel on the
Neuron pip repository needs glibc >= 2.43, newer than Ubuntu 24.04's; the DLAMI ships a compatible build):

``` bash
python3.12 -m venv venv && echo "import site; site.addsitedir('/opt/aws_neuronx_venv_torch_neuronx_2_14_0/lib/python3.12/site-packages')" \
  > venv/lib/python3.12/site-packages/zz_neuron.pth && source venv/bin/activate
export PATH=$PATH:/opt/aws_neuronx_venv_torch_neuronx_2_14_0/bin   # neuronx-cc
export DEEPINFER_ACCELERATOR=neuron
./build.sh                               # requirements-neuron.txt (torch / torch-neuronx already satisfied) + C++ build
python examples/neuron/offline_inference_neuron.py --model_name /path/to/DeepSeek-V4.1-Flash-1017 --ep 64
```

<details>
<summary>中文</summary>

DeepInfer 通过原生 PyTorch 的 `neuron` 设备（torch-neuronx 2.14，不依赖 torch-xla）在 Trainium2 上运行，结构与
Ascend 移植一致：`PlatformNeuron`、`tilelang_ops/neuron` kernel 后端，以及 `deepinfer/neuron_ops/vendor` 中
DeepGEMM / FlashMLA / DeepEP / DeepSelect / TileKernels 的纯 torch 实现。设计说明与约定见
[deepinfer/neuron_ops/NEURON_PORT.md](deepinfer/neuron_ops/NEURON_PORT.md)。

**现状：basic mode。** 逐算子即时执行；DSpark/MTP、零开销调度、CP、EPLB、YOCO 有界重放和 VL 均关闭。EP > 1 使用
真实的 neuron 集合通信。按真实 FP8/FP4 布局生成的小 checkpoint 在 CPU 和 NeuronCore 上都与官方参考实现一致（误差在
参考实现自身的噪声范围内，见 [examples/neuron](examples/neuron/README.md)）。性能优化（静态图、NKI kernel）和其余特性
在 [roadmap](https://github.com/whn09/DeepInfer/issues/19) 中跟踪。

在 Neuron DLAMI（trn2.48xlarge，SDK 2.33）上：

``` bash
bash tests/neuron/setup_trn2.sh          # 基于 DLAMI 的 torch-neuronx 环境建 venv、编译、生成小 checkpoint 和真值
bash tests/neuron/run_all.sh cpu         # 只用 host 的测试
CORES=0-7 bash tests/neuron/run_all.sh device
```

或者在叠加于 DLAMI torch-neuronx 环境之上的 venv 里手动构建（Neuron pip 源上的 `torch-neuronx` wheel 需要
glibc >= 2.43，比 Ubuntu 24.04 的版本新；DLAMI 自带的是兼容的构建）：

``` bash
python3.12 -m venv venv && echo "import site; site.addsitedir('/opt/aws_neuronx_venv_torch_neuronx_2_14_0/lib/python3.12/site-packages')" \
  > venv/lib/python3.12/site-packages/zz_neuron.pth && source venv/bin/activate
export PATH=$PATH:/opt/aws_neuronx_venv_torch_neuronx_2_14_0/bin   # neuronx-cc
export DEEPINFER_ACCELERATOR=neuron
./build.sh                               # requirements-neuron.txt（torch / torch-neuronx 已满足）+ C++ 编译
python examples/neuron/offline_inference_neuron.py --model_name /path/to/DeepSeek-V4.1-Flash-1017 --ep 64
```

</details>

## Performance benchmarks

| Prefill | Decode |
| :---: | :---: |
| [![Prefill: non-cached input throughput vs. latency](doc/images/prefill_pareto.png)](doc/images/prefill_pareto.png) | [![Decode: output throughput vs. per-user speed](doc/images/decode_pareto.png)](doc/images/decode_pareto.png) |

Measured on 4 SM100 GPUs for prefill and 64 for decode — each within one scale-up domain. To reproduce the benchmarks, see [run_sweep.sh](examples/sweep/run_sweep.sh).

## Evaluation benchmarks

See [eval/README.md](eval/README.md).

## License

DeepInfer is released under the Apache License 2.0. See [LICENSE](LICENSE) for details.

## Acknowledgments

DeepInfer was originally built primarily on [vLLM](https://github.com/vllm-project/vllm), with additional inspiration from [SGLang](https://github.com/sgl-project/sglang) and [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM). We thank their developers.

## Citation

```bibtex
@misc{deepinfer2026,
  title = {DeepInfer: High-performance, large-scale inference framework for DeepSeek models},
  author = {DeepSeek-AI},
  year = {2026},
  publisher = {GitHub},
  howpublished = {\url{https://github.com/deepseek-ai/DeepInfer}},
}
```