The last thing that kept the SM90 path on TileLang was the sparse MLA attention kernel. FlagGems
has a pure-Triton implementation of the same operator, and it turns out to be drop-in once two
differences are handled. With TileLang unimportable, the engine now runs end to end.
`h20_ops/sparse_mla.py` picks the backend: `DEEPINFER_SPARSE_MLA_BACKEND=tilelang|triton`, unset
means TileLang when it imports and Triton otherwise. `attn_sm90` imports from the dispatcher, so its
two entry points were the only call sites touched. Auto resolution catches `ModuleNotFoundError`
rather than testing for the package by name, because a missing TileLang fails inside
`sparse_mla_sm90`, not at the module lookup -- and it only swallows a *missing* dependency, so a
present-but-broken module still surfaces instead of being swapped out.
Vendored rather than depended upon: `flag_gems/__init__.py` registers ATen implementations
process-wide (`torch.library.Library("aten", "IMPL")` then `replace_customized_ops`). Importing the
package to reach one kernel would swap operators under the whole engine. The PyPI 5.0.2 file has no
`flag_gems` internal imports at all, so there was nothing to gain; it is 329 lines, accumulates in
fp32, and already guards the zero-width rope tail that V4.1 needs (RoPE folded inside the 512-wide
latent). The file header carries the source md5 and both deviations.
Two things it needed, both found by comparison rather than by reading:
* 5.0.2 guards each top-k block by its **position** in the index array (`if ck * BK <= max_col`).
Non-causal that is `SKV-1`, so when `topk > SKV` every block past the first `SKV` entries is
skipped. At SQ=16/SKV=128/topk=512 the output is off by 2.8e-2 -- and any prompt shorter than the
512-entry top-k hits it. Restored the data-dependent guard FlagGems master uses, marked in place.
* Its lse is a natural log; the TileLang one is log2. Unconverted, `_apply_sink` would be off by
1.4427x: plausible output, wrong numbers. The adapter converts, and the test asserts the ratio is
1.0 rather than 0.693 so it cannot regress quietly.
The kernel's operands have no batch axis, while decode hands it a different KV window per query.
Rather than loop, the adapter concatenates the windows and shifts query b's indices by b*seq_kv --
one launch, no extra memory, since the concatenated buffer is the tensor decode already built.
Verification ran in order, each step gating the next: 7 shapes against TileLang (`ALL_MATCH`, <=2
bf16 ULP, including `topk_length` and `attn_sink`), the decode call shape incl. the SWA+full
two-cache case (`ALL_MATCH`), the adapter contract (5 tests), the dispatcher and both `attn_sm90`
entry points (4 tests), then end to end with TileLang blocked: exit=0, 8/8 DP ranks agreeing, zero
kernel fallbacks, correct answers. A control run with TileLang present is unchanged.
Not token-identical, and the doc says so: prompt A shares 15 tokens then diverges, prompt B shares
1. A control run of the TileLang backend twice is bit-identical, so this is accumulation order, not
noise -- consistent with the 2.7% logits drift already recorded for the SM90 quantisation floor.
One TileLang dependency remains, and it is the deliberate one: `fused_sample`'s rejection-sampling
variant, which needs TileLang's Philox stream to reproduce the same tokens. So the pure-Triton run
disables speculative decoding (`DEEPINFER_USE_NEXTN_DECODE=0`).
|
||
|---|---|---|
| csrc | ||
| deep_cluster | ||
| deepinfer | ||
| doc | ||
| eval | ||
| examples | ||
| scripts | ||
| tests | ||
| vendor | ||
| .clang-format | ||
| .gitignore | ||
| .gitmodules | ||
| build.sh | ||
| CONTRIBUTING.md | ||
| H20-DEBUG-PLAN.md | ||
| H20-DRIFT-ANALYSIS.md | ||
| H20-PORTING-JOURNAL.md | ||
| H20-SM90-FINDINGS.md | ||
| H20-TRITON-PLAN.md | ||
| LICENSE | ||
| Makefile | ||
| pyproject.toml | ||
| README.md | ||
| requirements-ascend.txt | ||
| requirements-dev.txt | ||
| requirements-nvidia.txt | ||
| requirements.txt | ||
| setup.py | ||
| TODO.md | ||
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
Feature details and configuration
- Optimized kernels: FlashMLA, DeepGEMM, DeepEP, TileKernels, DeepSelect, and custom TileLang kernels.
- 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_DECODEcontrols 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=compactandDEEPINFER_YOCO_PREFILL_KEEP_TOKENS=140for 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=0and pass a fixedseq_seedtoSamplingParams(NVIDIA only). - Structured outputs: constrain generation to a regex, JSON schema, or EBNF grammar. Pass
so_args={"schema_type": ..., "schema": ...}toSamplingParams; see examples.
Supported hardware:
- NVIDIA SM100 GPUs
- Huawei Ascend 950 NPUs
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)
Build from source
Before building DeepInfer, please create and activate an empty Python 3.12 environment.
# Use `ascend` instead for the Ascend platform
# 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, 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.
python examples/offline_inference.py
See examples for details.
Performance benchmarks
| Prefill | Decode |
|---|---|
![]() |
![]() |
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.
Evaluation benchmarks
See eval/README.md.
License
DeepInfer is released under the Apache License 2.0. See LICENSE for details.
Acknowledgments
DeepInfer was originally built primarily on vLLM, with additional inspiration from SGLang and TensorRT-LLM. We thank their developers.
Citation
@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}},
}

