[ispobock]

Motivation

As mentioned in #2271, the original triton decoding kernel has significant performance degradation on long context. We refactored the kernel and adapted the flash decoding implementation from lightllm. Currently, the long context speed decay has been alleviated a lot.

Benchmark

Tested for input 128, output 2048.

Triton (this PR) num_kv_splits=8: 150->138

$ python3 -m sglang.bench_offline_throughput --model meta-llama/Llama-3.1-8B-Instruct --dataset-name random --num-prompt 1 --random-input 128 --random-output 2048 --random-range 1 --attention-backend triton

[2024-12-08 08:37:24 TP0] Decode batch. #running-req: 1, #token: 194, token usage: 0.00, gen throughput (token/s): 150.66, #queue-req: 0

[2024-12-08 08:37:37 TP0] Decode batch. #running-req: 1, #token: 2154, token usage: 0.00, gen throughput (token/s): 138.86, #queue-req: 0

We can increase the --triton-attention-num-kv-splits to get better performance on long context.

Triton (this PR) num_kv_splits=16: 150->144

python3 -m sglang.bench_offline_throughput --model meta-llama/Llama-3.1-8B-Instruct --dataset-name random --num-prompt 1 --random-input 128 --random-output 2048 --random-range 1 --attention-backend triton --triton-attention-num-kv-splits 16

[2024-12-08 08:40:28 TP0] Decode batch. #running-req: 1, #token: 194, token usage: 0.00, gen throughput (token/s): 150.18, #queue-req: 0

[2024-12-08 08:40:42 TP0] Decode batch. #running-req: 1, #token: 2154, token usage: 0.00, gen throughput (token/s): 144.00, #queue-req: 0

Triton (main branch): 147->126

$ python3 -m sglang.bench_offline_throughput --model meta-llama/Llama-3.1-8B-Instruct --dataset-name random --num-prompt 1 --random-input 128 --random-output 2048 --random-range 1 --attention-backend triton

[2024-12-08 08:35:01 TP0] Decode batch. #running-req: 1, #token: 194, token usage: 0.00, gen throughput (token/s): 147.93, #queue-req: 0

[2024-12-08 08:35:15 TP0] Decode batch. #running-req: 1, #token: 2154, token usage: 0.00, gen throughput (token/s): 126.67, #queue-req: 0

Flashinfer: 143->143

$ python3 -m sglang.bench_offline_throughput --model meta-llama/Llama-3.1-8B-Instruct --dataset-name random --num-prompt 1 --random-input 128 --random-output 2048 --random-range 1

[2024-12-08 08:43:15 TP0] Decode batch. #running-req: 1, #token: 194, token usage: 0.00, gen throughput (token/s): 143.84, #queue-req: 0

[2024-12-08 08:43:29 TP0] Decode batch. #running-req: 1, #token: 2154, token usage: 0.00, gen throughput (token/s): 143.24, #queue-req: 0

[merrymercy]

remove this in the func signature of decode_attention_fwd?

[ispobock]

Sure. max_len_in_batch and triton_attention_reduce_in_fp32 may also need to be removed.

[merrymercy]

After this, we do not need to reduce the cuda graph max bs for deepseek models?

[ispobock]

Let me verify it.

[WANG-GH]

Hello, does issue #2271 still need further development? I saw this issue and would like to give it a try. However, I noticed that you are currently working on fixing it. Could you share the progress so far? Do you still need any help? @ispobock @merrymercy

[MLKoz2]

This issue didn't aim to solve the problem with speculative decoding using Triton, did it?
(MTP/EAGLE3 slows down after a few thousands generated tokens with triton kernel, but with flashinfer there is no slow down even after 20k+ generated tokens)