Optimizing DeepSeek-V4.1-Flash Decoding on 16 A800 Cards

TL;DR

• Hardware and Objectives: Two 8×A800 40GB machines (16 GPUs in total, Ampere SM80 architecture), running DeepSeek-V4.1-Flash (552B total parameters, FP4 expert).
• Topology and Memory Solution: Abandoning cross-machine TP16, adopting TP8 × PP2 (cross-machine communication wait time reduced from 20ms to 15μs); 203GB Engram knowledge base large table offloaded through Host CPU locked page memory, freeing up GPU memory for the 5.65 million Token fp8_ds_mla cache pool.
• Core Performance Gains: Rewriting Fused Markov sampling and Small-Batch Sparse MLA decoding core using Triton, achieving 24.6% reduction in Median ITL (32.93ms → 24.84ms) under concurrent C=8, and +16.8% increase in overall throughput (181.5 → 212.0 tok/s).
• Open-Source Code and Upstream: Independent open-source repository JimiZhou/deepseek-v4.1-flash-sm80-turbo, upstream PR submitted to vLLM official #66275.

I've been experimenting with DeepSeek-V4.1-Flash (552B total parameters, 1017 release) on two 8-card A800-40G nodes over the past few days.

The vLLM upstream repository released a specialized decode (low-latency decoding path) for this model, which I took a look at by checking the diff of PR #64893. The entire screen was filled with Hopper (SM90) and Blackwell-related content: WGMMA asynchronous matrix instructions, TMA (cp.async.bulk.tensor) tensor memory copy descriptors, Thread Block Cluster distributed shared memory, and even FP4 expert calculations using native hardware instructions.

However, many teams and university labs still have Ampere (A100-80G or A800-40G) as their largest available GPU. The SM80 architecture lacks all of these hardware features, and FP8 and FP4 cannot run native Tensor Core MMA. The community has been tracking Issue #57144 for a long time, and most people have been relying on general fallback branches to get the model "running," but performance is not ideal once concurrency or small batch latency is measured.

This article reviews our process of optimizing 16 A800 40G GPUs, from hardware topology and memory planning to rewriting decoding hotpaths using Triton and finally conducting real-world tests.

1. The First Account: Why Does Cross-Machine TP16 Always Crash on Two Machines?

When receiving two 8-card machines, many people's first reaction is to enable cross-machine TP16. Theoretically, 16 cards can share the 510GB model weight, minimizing the pressure on each card.

However, measuring communication latency reveals that TP16 is completely unworkable in the decoding stage.

We conducted experiments on our machines using 14KB (the size of the hidden vector for single-token decoding) and 16MB data to measure communication overhead:

Doing the math makes it clear:
DeepSeek-V4.1-Flash has 40 decoder layers. Each layer requires two AllReduce operations (one for attention output and one for MoE expert output), resulting in at least 60 to 80 AllReduce operations per forward pass.

If cross-machine TP16 is used:

80 次规约 × 250 μs 单次网络等待 = 20,000 μs = 20 毫秒

This means that for every generated token, the GPU hasn't even finished running a line of code, and just waiting for the package to come back on the network card takes 20 milliseconds. The actual calculation time for a single card in a small batch is only around 10 milliseconds. The time spent waiting for packages on the network accounts for more than 65% of the time, and the entire system is directly crippled by the small package transmission across the network.

Therefore, for dual-machine 16-card low-latency running, there is almost only one solution: TP8 × PP2.
Within a node, 8 cards are fully utilized to run TP8 via NVLink, and between nodes (from PP0 to PP1), there is only one point-to-point activation transmission (P2P Send/Recv) per forward step. For the same 14KB data, using RDMA single-way direct transmission only takes 15 microseconds.

The cross-network communication latency of the entire machine is reduced from 20 milliseconds to 15 microseconds, and the decoding time can return to being dominated by GPU computing power.

2. The second bill: How to fit 510GB of weights into 640GB of video memory?

16 cards with 40G each have a physical video memory limit of only 640GB for the entire cluster.
The disk weights of the DeepSeek-V4.1-Flash model are as large as 498GB, and the actual resident occupancy is close to 510GB after adding scale parameters and other constants.

If the full weights are loaded into video memory without optimization, considering the fixed overhead of the driver and PyTorch, the 16 cards will have less than 100GB of video memory left, with only 5 to 6GB per card. Once a long context or concurrent running is started, the KV Cache will have no place to be stored.

The key to breaking through this limitation lies in the 203.4 GB Engram knowledge library large table:

The Engram is essentially a static large table that matches historical memories based on N-gram hashing. Since both host machines have 1TB of physical system memory (with available memory exceeding 750GB), it is possible to utilize --engram-config '{"cpu_offload":true}' to lock the two large tables in the Host CPU's Pinned Memory. The GPU can asynchronously fetch the corresponding slices according to the hit index during forward calculation, directly reducing the video memory load by 203 GB.

After completing this step, each card only needs to bear half of the pipeline stage (20 layers) of TP8 slicing weights, with resident occupancy reduced to around 20.5 GB.

Deducting 2.2GB of CUDA Graph capture space and 2GB of communication workspace, each card has stably freed up 4.0 GB of video memory for the fp8_ds_mla cache pool. Calculated based on 584 bytes per compact record, the KV Cache of the entire cluster can accommodate 565 million tokens, which can support 5.4 times the concurrency even when running 1 million ultra-long contexts.

3. Hot path measurement: Where is the real bottleneck?

Many people's first reaction to optimizing decoding is to tweak Attention. We used torch.cuda.Event on A800 hardware to measure the hot path of each component in decoding, and the actual time distribution is as follows:

7**% of the GEMVBF16 computing power is relatively sufficient, and the time consumption is stable

The bottleneck is very clear:
In the case of small batch decoding, the original MoE implementation treats the 6 activated experts as independent small GEMMs for scheduling. When Batch=1 or 4, there are only two or three tokens assigned to each expert, or even just one. The 108 SMs on the A800 haven't even reached 5% utilization, and are all experiencing kernel launch scheduling queues, with intermediate tensors being repeatedly transferred between video memory and on-chip cache.

---

4. Rewriting the decoding operator on SM80

After clarifying the bottleneck and hardware constraints, the optimization goal is clear:
without Hopper features, merge small operators and keep data on the chip as much as possible.

4.1 Pure Triton's DSpark serial sampling fusion

The DSpark speculative draft mechanism takes four steps in the original process for each step of prediction:
adding bias (Bias Add) → calculating Argmax → injecting Gumbel noise sampling → looking up the next Embedding.

A single prediction takes 4 to 5 steps, with nearly 20 independent kernel launches in between, not only fragmenting the CUDA Graph, but also causing scheduling bubbles between the CPU and GPU.

We rewrote a Fused Markov sampling operator in dspark_fused.py using Triton:

1. Single-thread block on-chip calculation: directly calculate the base logits and bias within a single block, add Gumbel noise locally, and calculate the local maximum value within a 1024-width word table block using shared memory.
2. Atomic grid to victory:

# 让每个线程块跑完后用原子操作报到
tl.debug_barrier()
arrival = tl.atomic_add(counter_ptr + req, 1, sem="acq_rel", scope="gpu")
if arrival != num_blocks - 1:
    return  # 没到齐的块直接退出

# 最后一个到达的线程块充当 Leader，直接在片上决出全词表的全局最大值，
# 并当场写入 draft_tokens，顺手从 embed_weight 里查出下一步 GEMM 需要的 Markov 嵌入
winner = tl.where(first_nan < BLOCK_P, first_nan, first_max)
token = tl.load(partial_idx_ptr + req * BLOCK_P + winner, volatile=True).to(tl.int64)
tl.store(draft_tokens_ptr + req * draft_tokens_stride + step, token)
if HAS_NEXT:
    embed_row = tl.load(embed_weight_ptr + token * embed_weight_stride + r, mask=r < rank)
    tl.store(embed_ptr + req * embed_stride + r, embed_row, mask=r < rank)

This step reduces the original multiple scheduling of each draft to 1 time, and each validation step can save about 200 microseconds of Host scheduling overhead.

4.2 Specialized Sparse MLA Decode Core for Small Batches

FlashMLA is locked to SM90+, and Ampere can only use Triton. However, the upstream universal Triton operator is compatible with Prefill, which includes a large number of dynamic SWA block scheduling and boundary branch judgments.

We specialized the decoding scenario for Q=1 and M <= 16 in sparse_mla_decode.py:

• Fixed block size to B_M = 1, B_N = 64, B_HD = 512;
• Directly addressed reading of 584 bytes of compact fp8_ds_mla records;
• Completed Online Softmax accumulation in registers, avoiding global memory storage.

The test results show that under the condition of M=2, K=256, the output is compared with the FP32 precision of the PyTorch native SDPA, and the maximum error is only 0.00507, which is completely within the tolerance range of BF16.

4.3 Pipeline Parallelism and Multi-Modal Engineering Pitfalls

When connecting the executor to the full model, there are two representative error reports:

1. Loss of IntermediateTensors causes graph capture assertion failure:
   In the TP8 × PP2 architecture, the node in the PP0 stage cannot return the expanded regular Tensor like a single machine, and must package the hidden stream into IntermediateTensors to trigger point-to-point communication in the _head. If this packaging is missing, the Worker will throw an AssertionError when entering the FULL graph capture stage.
2. Missing parameters for multi-modal vision routing:
   DeepSeek-V4.1's MoE expert routing has a visual Sentinel marker, and the underlying routing logic strictly checks if input_ids exist. A specialized decode executor that only passes the hidden vector and loses input_ids will immediately throw a ValueError: DeepSeek V4 vision MoE routing requires input_ids.

5. End-to-End Measurement Comparison

After all optimizations, we mounted a new image on the cluster and ran a benchmark comparison with the original Stock Baseline under the same conditions (AB test, same Prompt input, character-by-character capture of SSE Delta streaming timestamp).

1. Core Concurrency and Throughput Performance

Concurrency Baseline Throughput (Stock) Optimized Throughput (Specialized) Throughput Change Trend
C = 1 48.13 tok/s 46.98 tok/s -2.3% (extremely low concurrency, basically flat)
C = 2 77.80 tok/s 74.07 tok/s Robust doubling expansion
C = 8 181.46 tok/s 211.98 tok/s Significant increase +16.8%
C = 16 267.32 tok/s 249.58 tok/s High concurrency saturation without cliff

2. Token Generation Delay Performance (Median ITL)

Concurrency Baseline Delay (Stock) Optimized Delay (Specialized) Delay Optimization Change
C = 1 18.16 ms 18.16 ms Precise alignment with baseline
C = 2 19.26 ms 20.07 ms Maintaining excellent low latency
C = 8 32.93 ms 24.84 ms Significant reduction -24.6%
C = 16 43.96 ms 44.84 ms Normal queuing delay

3. DSpark Speculative Decoding Verification Step Performance (C = 1 Single Concurrency)

Verification Step Metric Stock Baseline Native Specialized Decode Performance Description
Full model single-step time (T_step) 18.05 ms 18.38 ms Millisecond-level equivalent, stable
DSpark burst equivalent TPOT 6.60 ms / character 6.56 ms / character Equivalent to single-stream ~154 tokens/s (extremely responsive)

Several Conclusions:

1. Medium concurrency (C=8) benefits the most:
   When concurrency reaches 8, Median ITL decreases from 32.93 milliseconds to 24.84 milliseconds, a decrease of 24.6%; at the same time, the overall machine aggregated throughput increases from 181.5 to 212.0 tokens/s (an increase of 16.8%). In this load range, the effect of fusion operators eliminating Launch Gap and memory migration is most pronounced.
2. How to look at delay metrics:
   Under single concurrency, the full-model single-step verification time is around 18ms.

However, in scenarios where DSpark draft hit rates are higher, such as code completion or logical reasoning, the actual generation latency is 6.56 milliseconds/character (equivalent to 154 tokens/s for a single stream), resulting in an extremely responsive human-computer interaction experience.
3. Correctness Verification:
In greedy sampling (T=0) mode, the output results of algebraic derivation (3x + 15 = 42 => x = 9), code generation (is_prime), and JSON format extraction remain precisely consistent before and after reconstruction.

6. Conclusion

After going through this process, the greatest realization is that in large model inference engineering, while hardware generation gaps do exist, it does not mean that old architectures can only accept degraded versions.

By clarifying the network topology between machines, utilizing the host's large memory to free up GPU memory, and then using Triton to fuse and reduce Launch for actual hotspots, two old 8-card A800 machines can still run DeepSeek-V4.1, a new generation of ultra-large MoE models, stably and quickly.

Open-Source Repository and Upstream PR Link

• Independent Open-Source Repository: JimiZhou/deepseek-v4.1-flash-sm80-turbo
  (includes a complete set of reproducible deployment scripts, Triton kernel source code, desensitized benchmark test data, and detailed GPU memory budget documentation)
• vLLM Official Upstream PR: vllm-project/vllm #66275
  ([Feature][Spec Decode] Add Ampere (SM80) Native Specialized Decode for DeepSeek-V4.1-Flash)