mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
d0f7e4f69c
- requirements-neuron.txt: requirements.txt + dev deps without tilelang / tile-kernels (pure-torch drop-ins are vendored), torch 2.14 + torch-neuronx 2.14 from the Neuron repository. - build_vendor.sh: nothing to build for neuron; setup.py drops tilelang / tile-kernels from install_requires on neuron. - setup.py and get_platform(): an exported DEEPINFER_ACCELERATOR wins over device detection, so images build where no accelerator is visible. - docker/neuron/Dockerfile (+ per-Dockerfile ignore): on the native PyTorch Neuron DLC public.ecr.aws/neuron/torch-neuronx:2.14.0-neuronx-py314- sdk2.33.0-ubuntu26.04; build.sh with TORCH_DEVICE_BACKEND_AUTOLOAD=0 and an import self-check in Neuron CPU mode. - README: Trainium in supported hardware and a "DeepInfer on AWS Trainium" section (English + folded Chinese); the pip torch-neuronx wheel needs glibc >= 2.43, so on Ubuntu 24.04 layer the venv on the DLAMI env. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
178 lines
8.8 KiB
Markdown
178 lines
8.8 KiB
Markdown
# 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 |
|
||
| :---: | :---: |
|
||
| [](doc/images/prefill_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}},
|
||
}
|
||
```
|