Skip to content

Optimize Triton decoding kernel for long context #2394

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 7 commits into from
Dec 8, 2024
Merged

Optimize Triton decoding kernel for long context #2394

merged 7 commits into from
Dec 8, 2024

Conversation

ispobock
Copy link
Collaborator

@ispobock ispobock commented Dec 8, 2024 •
edited
Loading

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
merrymercy reviewed Dec 8, 2024
View reviewed changes
python/sglang/srt/layers/attention/triton_ops/decode_attention.py
@@ -705,10 +650,10 @@ def decode_attention_fwd(
o,
req_to_token,
b_req_idx,
b_start_loc,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

remove this in the func signature of decode_attention_fwd?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

python/sglang/srt/layers/attention/triton_backend.py
forward_batch.batch_size,
self.num_head,
self.num_kv_splits,
self.v_head_dim + 1,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let me verify it.

@merrymercy merrymercy merged commit 7dc66fc into sgl-project:main Dec 8, 2024
0 of 14 checks passed
timethink pushed a commit to timethink/sglang that referenced this pull request Mar 9, 2025
@ispobock
Optimize Triton decoding kernel for long context (sgl-project#2394)
aff4fc1
@WANG-GH
Copy link
Contributor

WANG-GH commented Mar 11, 2025

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@merrymercy merrymercy merrymercy approved these changes

@Ying1123 Ying1123 Awaiting requested review from Ying1123 Ying1123 is a code owner

@zhyncs zhyncs Awaiting requested review from zhyncs zhyncs is a code owner

@hnyls2002 hnyls2002 Awaiting requested review from hnyls2002

@ByronHsu ByronHsu Awaiting requested review from ByronHsu

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@ispobock @WANG-GH @merrymercy