Optimize Triton decoding kernel for long context

[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.