Optimize DeepSeek-V4.1 on NVIDIA Ampere with DeepInfer-SM80

What Is DeepInfer-SM80?

DeepInfer-SM80 is a collection of optimized Python, Triton, and TileLang modules for running DeepSeek-V4.1 on NVIDIA Ampere GPUs through the vLLM V1 engine. It ports core ideas from DeepSeek’s DeepInfer, which targets NVIDIA Blackwell, and redesigns the critical execution path around the capabilities of Ampere hardware.

Instead of relying on Blackwell-only features such as tcgen05 warp-group instructions, Tensor Memory, hardware UE8M0 scale instructions, and 228 KB of shared memory per streaming multiprocessor, the project uses Ampere’s mma.sync.aligned.m16n8k16 Tensor Cores, SM-aware scheduling, integer-bitshift UE8M0 decoding, and YOCO shared-KV reuse.

The project supports data-center SM80 GPUs such as the A100 and A800, as well as SM86 devices including the RTX 3090, RTX 3090 Ti, RTX A6000, A40, A10, and A30. Its kernels use no more than 48 KB of shared memory per CTA, allowing them to run within the lower shared-memory limits of SM86 devices.

Why Use It?

DeepInfer-SM80 focuses on the expensive operations involved in DeepSeek-V4.1 prefill, sparse attention, KV-cache access, residual routing, and long-context inference. The repository reports the following results from a production configuration using six GA100 64 GB GPUs, TP2 × PP3, DSpark speculative decoding with K=5, and a maximum context length of 1,048,576 tokens:

• 16K cold prefill: 4,037.0 tokens per second, up from 3,349.4 tokens per second.
• 64K cold prefill: 4,506.8 tokens per second, up from 3,627.3 tokens per second.
• 64K time to first token: 10.7 seconds, down from 13.2 seconds.
• Four-stream speculative decode: 170.5 aggregate tokens per second, compared with 163.8 tokens per second.
• Single-stream Python-code decode: 160.8 tokens per second, compared with 151.8 tokens per second.
• GPU-resident KV capacity: 4,648,480 tokens, with no additional VRAM overhead relative to the baseline.

These measurements describe the repository’s tested hardware and serving configuration. Results on other Ampere GPUs, model configurations, workloads, or vLLM environments may differ.

Key Optimizations

Unified 512-Dimensional Sparse MLA

DeepSeek-V4.1’s DeepSelect sparse attention stores a 448-dimensional NoPE latent in float8_e4m3fn, using UE8M0 power-of-two block scales, together with a 64-dimensional BF16 RoPE suffix. DeepInfer-SM80 loads the complete 512-dimensional tile and processes it through a single Tensor Core dot-product path rather than handling the NoPE and RoPE portions separately.

The kernel replaces an exp2-based UE8M0 conversion with an IEEE-754 integer bitshift:

(scale_u8.to(tl.int32) << 23).to(tl.float32, bitcast=True)

For valid UE8M0 exponents, this produces bit-identical FP32 scale values while moving the work from transcendental operations to the integer pipeline.

The tuned configuration uses BLOCK_N=64, num_warps=4, and num_stages=2. Reducing register pressure allows two CTAs to reside on each SM80 streaming multiprocessor. When num_kv_splits == 1, the kernel writes directly to the final output and avoids launching a separate reduction stage.

Single-CTA Top-K and SWA Index Combination

The optimized _combine_topk_swa_indices_kernel replaces a sequence of PyTorch allocations and operations with one Triton kernel using one CTA per query token. At T=1536, the measured latency falls from 1,083.72 microseconds to 48.88 microseconds, a 22.17× speedup.

To avoid inter-warp write races, each output row is divided into three disjoint regions:

1. The compressed Top-K indices occupy [0, topk_len).
2. The sliding-window attention indices occupy [topk_len, topk_len + swa_len).
3. The remaining columns are padded with -1.

The README reports bit-exact output across all tested sequence lengths and compression ratios of 0, 1, and 2.

Faster KV-Cache Dequantization and Gathering

The dequantize_and_gather_k_cache path uses the same bitshift-based UE8M0 conversion. It also writes through the destination strides supplied to the Triton kernel, eliminating intermediate torch.empty allocations and subsequent .copy_() operations.

The measured speedup is 4.02× at a 16K context and 7.64× at a 48K context, with bit-exact reference matches in the reported tests.

SM80-Tuned mHC Residual Routing

DeepSeek-V4.1 performs four-stream manifold-constrained hyper-connection routing at 83 boundaries per forward pass. DeepInfer-SM80 adds an aligned (6, 8, 128) split geometry for the model’s hidden size of 5,120. Each of the eight slices contains 640 elements, which is divisible by 128.

The project also lowers the small-matrix threshold from 32 to 16 so token counts of at least 16 can use the Tensor Core T.gemm path. The reported mhc_fused_post_pre_delayed speedup is 2.89× at T=12 and 1.30× at T=6.

YOCO Metadata Reuse

In DeepSeek-V4.1’s You-Only-Cache-Once architecture, layers 21 through 39 read from layer 20’s full-resolution KV cache. These consumers share the same expanded sequence lengths, block table, and compression ratio.

DeepInfer-SM80 stores the ragged K metadata in the vLLM forward context when processing layer 20. Reusing that data for layers 21 through 39 eliminates 19 redundant metadata-building GPU kernel chains per prefill chunk.

Before Installation

You need a vLLM V1 environment that will run DeepSeek-V4.1 on a supported NVIDIA Ampere GPU. The repository supplies a unified patch and six replacement modules, but the README does not prescribe exact CUDA, Python, Triton, TileLang, or vLLM version numbers. Use a vLLM checkout compatible with the patch and your existing DeepSeek-V4.1 deployment.

The supported hardware groups are:

• SM80 data-center GPUs: A100 40 GB, A100 80 GB, GA100 64 GB, and A800 40 GB or 80 GB.
• SM86 workstation and consumer GPUs: RTX 3090, RTX 3090 Ti, RTX A6000, A40, A10, and A30.

The software is distributed under the Apache License 2.0.

Installation Option A: Apply the Unified Patch

This method applies all six modified files to a vLLM source checkout. Open a terminal in the root directory of that checkout and run:

git apply /path/to/Deepinfer-SM80/patches/vllm-sm80-deepinfer.patch

Replace /path/to/Deepinfer-SM80 with the location of the cloned or downloaded project. The patch contains the sparse MLA, cache-gather, Top-K plus SWA, mHC, warmup, and YOCO metadata-reuse changes.

Installation Option B: Install the Drop-In Modules

If replacing files directly is more appropriate for your environment, copy the project’s vLLM module tree over the installed vLLM package:

cp -r vllm_sm80_deepinfer/vllm/* /path/to/site-packages/vllm/

Replace /path/to/site-packages/vllm/ with the actual vLLM package directory. The same files can also be bind-mounted as read-only overlays in a Docker deployment.

The drop-in tree modifies these functional areas:

• model_executor/kernels/mhc/ for TileLang mHC execution and warmup.
• models/deepseek_v4_1/amd/rocm.py for the Triton Top-K and SWA combiner.
• models/deepseek_v4_1/common/ops/cache_utils.py for KV-cache gathering.
• models/deepseek_v4_1/nvidia/model.py for YOCO metadata reuse.
• v1/attention/ops/rocm_aiter_mla_sparse.py for unified sparse MLA kernels.

Although two optimized files retain rocm or amd in their paths, the repository explicitly uses them as part of its Ampere-targeted vLLM integration. Preserve the provided directory structure when applying the drop-in files.

Basic Usage: Validate the Kernels

After installing the patch or drop-in modules, run the supplied isolated microbenchmarks on the target A100, A800, RTX 3090, RTX A6000, or other supported Ampere GPU:

python3 benchmarks/microbench_sm80_kernels.py
python3 benchmarks/microbench_sm80_part2.py
python3 benchmarks/microbench_sm80_part3.py

The benchmark files cover the sparse MLA and KV-gather kernels, TileLang mHC kernels, and the Top-K plus SWA combiner. Their purpose is to measure local kernel latency and verify the optimized outputs against baseline references.

Generate the repository’s benchmark charts with:

python3 benchmarks/generate_charts.py

The complete benchmark dataset is also available in benchmarks/results/sm80_benchmark_summary.json.

Understanding Numerical Validation

Several optimized paths are reported as bit-exact:

• _combine_topk_swa_indices_kernel
• dequantize_and_gather_k_cache
• _sparse_attn_decode_kernel
• mhc_post_with_sqrsum

Other Tensor Core paths can differ slightly because BF16 accumulation or MMA execution order changes. The reported maximum differences are below 1.9e-3 for the measured fused mHC configurations and below 4.8e-3 for the measured sparse-attention prefill configuration.

Prompt-logprob negative log-likelihood results remained within the reported run-to-run noise floor. When evaluating a deployment, inspect both performance and numerical output rather than assuming every optimized operator must be bit-identical.

Advanced Deployment Tips

Use TP2 × PP3 for the Tested Six-GPU Production Layout

For six 64 GB or 80 GB Ampere GPUs, the repository recommends TP2 × PP3 with a 14,14,12 layer partition. In the tested configuration, this arrangement provided 160.8 to 188.8 tokens per second for selected single-stream code and long-generation workloads, 4,506.8 tokens per second for 64K cold prefill, and capacity for approximately 4.65 million GPU-resident KV tokens.

This capacity corresponds to approximately 4.43 contexts of one million tokens under the benchmark configuration.

Understand the Alternatives

• PP6: offers roughly 4.5 to 4.7 million tokens of capacity, but the README reports substantially slower single-stream decode because tensor parallelism is one and requests cross six pipeline stages.
• Asymmetric TP2 plus TP4: increases estimated capacity to roughly six million tokens, but splitting the MoE intermediate size of 2,304 into 576-wide TP4 partitions hurts Marlin tile efficiency.

Consult docs/05_TOPOLOGY_1M_CONTEXT_ANALYSIS_SM80.md in the repository before selecting a topology for a one-million-token deployment.

Benchmark the Actual Workload

The largest gains come from frequently invoked operators, but end-to-end improvement depends on prompt length, concurrency, speculative decoding, topology, and GPU model. Run the isolated benchmarks first, then compare complete serving workloads using the same prompts, context sizes, and concurrency levels as your production application.

Preserve the Direct-Output Fast Path

The sparse MLA decode optimization skips a second reduction launch when num_kv_splits == 1. If you modify attention scheduling or split behavior, remeasure the kernel because the direct-output path is an important part of the reported decode improvement.

Keep YOCO Metadata in the Forward Context

The layer 20 metadata cache is valuable because all layer 21–39 consumers share its KV source. Changes to model execution or forward-context handling should preserve this reuse; otherwise, the engine may rebuild identical ragged metadata for every consumer layer.

Conclusion

DeepInfer-SM80 provides a focused path for accelerating DeepSeek-V4.1 on NVIDIA Ampere without adding KV-cache VRAM overhead. Its improvements combine Ampere-aware sparse attention, bitshift UE8M0 decoding, direct KV-cache output, a race-free Top-K and SWA combiner, aligned mHC routing, and YOCO metadata reuse.

Start by applying the unified patch or installing the six drop-in modules. Run all three microbenchmark suites to confirm latency and numerical behavior on your hardware, then evaluate the optimized vLLM environment with your real serving workload. For a six-GPU, one-million-context deployment, use the repository’s tested TP2 × PP3 configuration as the primary reference point.