Support DP MLA

[ispobock]

Motivation

Support data parallel on MLA for DeepSeek model to reduce replicated KV cache.

python -m sglang.launch_server --model-path neuralmagic/DeepSeek-Coder-V2-Instruct-FP8 --trust-remote-code --tp 8 --dp 8 --enable-dp-attention

Modifications

• Add --enable-dp-attention option. When it is turned on, DP and TP share the same workers.
• Add IDLE forward mode for workers that do not have sequence to forward but need TP sync with other workers.
• Implement model forward with DP attention + TP MoE.

Performance

Compared to the main branch, this PR improves the prefill throughput by 20% and the decode throughput by 67% for DeepSeek-V2 model on 8*H100.

DP+TP (this PR):

• prefill: 21658.78 toks/s
• decode: 11174.62 toks/s

TP (main branch):

• prefill: 17941.92 toks/s
• decode: 6656.75 toks/s

Reproduce:

# DP+TP
python -m sglang.launch_server --model-path neuralmagic/DeepSeek-Coder-V2-Instruct-FP8 --disable-radix-cache --trust-remote-code --tp 8 --dp 8 --enable-dp-attention
# TP
python -m sglang.launch_server --model-path neuralmagic/DeepSeek-Coder-V2-Instruct-FP8 --disable-radix-cache --trust-remote-code --tp 8

# bench prefill
python3 -m sglang.bench_serving --backend sglang --dataset-name random --random-input 512 --random-output 1 --random-range-ratio 1 --num-prompts 10000
# bench decode
python3 -m sglang.bench_serving --backend sglang --dataset-name random --random-input 1 --random-output 512 --random-range-ratio 1 --num-prompts 10000

TODO

• Compatible with cuda graph: Support cuda graph for DP attention #2061
• Compatible with overlap mode

[fengyang95]

@ispobock How much performance improvement is expected? Is it mainly in throughput or latency?

[ispobock]

How much performance improvement is expected? Is it mainly in throughput or latency?

@fengyang95 There is an issue with dp 8. I will test the performance once the issue fixed. It's mainly for throughput.

[fengyang95]

• Compatible with cuda graph
• Compatible with overlap mode

@ispobock hi, when will support for cuda graph be planned? It is critical for latency improvement.

[ispobock]

when will support for cuda graph be planned?

I will support it soon. The code is almost done and needs some tests.

[fengyang95]

when will support for cuda graph be planned?

I will support it soon. The code is almost done and needs some tests.

@ispobock How much additional VRAM would this require approximately?

[ispobock]

How much additional VRAM would this require approximately?

@fengyang95 For the 236B V2 model, if use DP attention, the total weights will take ~570G, so prefer to use FP8 quantized model for better performance.

[merrymercy]

Great work!

[fengyang95]

@ispobock Does this support W4A16? My VRAM is very limited, and even using fp8 VRAM is not enough.

[ispobock]

Does this support W4A16?

@fengyang95 AWQ is supported in #2364.