feat: frequency, min_new_tokens, presence, and repetition penalties

[vhain]

Thank you for your contribution, we really appreciate it. The following instructions will help improve your pull request and make it easier to receive feedback. If there are any items you don't understand, don't worry. Just submit the pull request and ask the maintainers for help.

Motivation

Penalties are defined in SamplingParams, and the tensors being created during batch execution, but not yet implemented.

Also we can cut-down some milliseconds in ITL, and some VRAM usage, by not creating those unnecessary tensors.

Modification

• Implemented following penalties, similarly to they are in vllm:
  • frequency_penalty
  • min_new_tokens
  • presence_penalty
  • repetition_penalty
• Introduces following new sampling_parameters, and added them to OpenAI protocol in same name as appeared on vllm for compatibility:
  • stop_token_ids: similar to stop, but token IDs instead of str. This could be useful when used with min_new_tokens.
  • repetition_penalty
  • min_new_tokens
• Each penalizers are fully batch-compatible (filter, merge supported).
• Penalizers will do nothing (not even create a tensor) if they are not required.
• Performance implcations on using penalizer sampling_parameters are documented.
• Add ability for sglang.bench_serving to accept additional --request-body to specify sampling_params if desired.

Benchmark

Benchmarks were ran on single NVIDIA H100 SXM5

Baeline calculated on dc9d06d:

$ python3 -m sglang.bench_serving --backend sglang --port 8413 --dataset-name random --num-prompts 3000 --random-input 256 --random-output 512

============ Serving Benchmark Result ============
Backend:                                 sglang
Traffic request rate:                    inf
Successful requests:                     3000
Benchmark duration (s):                  67.50
Total input tokens:                      378633
Total generated tokens:                  775651
Total generated tokens (retokenized):    774862
Request throughput (req/s):              44.44
Input token throughput (tok/s):          5609.11
Output token throughput (tok/s):         11490.58
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   42141.26
Median E2E Latency (ms):                 45393.52
---------------Time to First Token----------------
Mean TTFT (ms):                          20705.28
Median TTFT (ms):                        14893.99
P99 TTFT (ms):                           49444.93
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          94.00
Median TPOT (ms):                        91.96
P99 TPOT (ms):                           232.66
---------------Inter-token Latency----------------
Mean ITL (ms):                           181.11
Median ITL (ms):                         62.46
P99 ITL (ms):                            441.66
==================================================

Calculated on this PR's branch (without any penalizers):

$ python3 -m sglang.bench_serving --backend sglang --port 8413 --dataset-name random --num-prompts 3000 --random-input 256 --random-output 512

============ Serving Benchmark Result ============
Backend:                                 sglang
Traffic request rate:                    inf
Successful requests:                     3000
Benchmark duration (s):                  66.11
Total input tokens:                      378633
Total generated tokens:                  775651
Total generated tokens (retokenized):    775118
Request throughput (req/s):              45.38
Input token throughput (tok/s):          5727.04
Output token throughput (tok/s):         11732.16
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   40881.94
Median E2E Latency (ms):                 43967.10
---------------Time to First Token----------------
Mean TTFT (ms):                          19884.75
Median TTFT (ms):                        14226.56
P99 TTFT (ms):                           47738.97
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          91.96
Median TPOT (ms):                        90.11
P99 TPOT (ms):                           308.54
---------------Inter-token Latency----------------
Mean ITL (ms):                           174.54
Median ITL (ms):                         58.56
P99 ITL (ms):                            440.18
==================================================

Additional benchmarks can be found in the updated sampling_params.md

Notes to Reviewers

This is my very first time writing Python code, and thus working with PyTorch (my industry backgrounds are primarily Golang, Rust, Java, Nodejs, ...). So please feel to add any comments if something looks off or inefficient.

Checklist

• 1. Ensure pre-commit pre-commit run --all-files or other linting tools are used to fix potential lint issues.
• 2. Confirm that modifications are covered by complete unit tests. If not, please add more unit tests for correctness.
• 3. Modify documentation as needed, such as docstrings or example tutorials.

[zhyncs]

@vhain Nice work!

[merrymercy]

Thanks for your contributions! The code quality is very high and it is well tested. Great work! I only have few minor style comments.

[Ying1123]

@vhain ready to merge?
@zhyncs can you take another look and approve this?

[Ying1123]

In the future, we plan to wrap all the logits processing and sampling into the Cuda graph as well.
Are these penalizers cuda-graph compatible?

[zhyncs]

overall LGTM

[zhyncs]

In the future, we plan to wrap all the logits processing and sampling into the Cuda graph as well. Are these penalizers cuda-graph compatible?

@vhain

[vhain]

@Ying1123 @zhyncs I'm still pretty new to this ground (started playing with LLM ~ month ago, started using sglang ~ 2 weeks ago, and this is my first ever Python/PyTorch code written). I don't think this penalizers are implemented based on CUDA grapu, but I will take a look on other CUDA graph implementations in sglang to figure it out. Hopefully I can pull up another PR to make it use CUDA graphs...

[vhain]

@Ying1123 @zhyncs Although I might need some extra time on what and how the CUDA graph works, all those penalizers are purely written in PyTorch (with device="cuda"). So it should be fairly simple to migrate..?

[merrymercy]

Yes, do not worry about it for now. We will refactor and start moving the sampling code to be within the CUDA graph. We will let you know later and ask for your help if needed.

[Xu-Chen]

We found that presence_penalty has a great influence on the speed of decode.
When set presence_penalty = 0.5 qwen-72b-gptq-int4 ( ~17 token/s on 4 * A100)，3.x slower than set presence_penalty=0.0 (~60 token/s on 4*A100).
cc @merrymercy @zhyncs

This can be reproduced in version 0.3.0, but does not occur in version 0.2.13.

[merrymercy]

@hnyls2002 Is this related to how the cuda graph is triggered?

[merrymercy]

@Xu-Chen Recently, we did some modifications to how the cuda graph is enabled. In 0.3.0, if you set penalty or regex, cuda graph will be disabled. This is a bug and we will fix it soon.

[hnyls2002]

@vhain @Xu-Chen I believe the performance slowdown problem will be fixed in #1392.

With llama3.1-8b on 1 A100, the decoding speed is 77 tokens per sec the same as without presence penalty, but on main the speed is 71 tokens per sec.

On 4 x A100, the decoding speed is 129 tokens per sec, while without presence penalty the speed is 134 tokens per sec, but on main the speed is 58 tokens per sec.

[hnyls2002]

@vhain @Xu-Chen @merrymercy
The penalties consume calculation in two phases:

• Cumulate output tokens or input tokens.
• Apply the penalties to the logits results.

The second phase was fully overlapped by

sglang/python/sglang/srt/model_executor/model_runner.py

Line 568 in 184e444

batch.sampling_info.update_penalties()

(this calculation can be executed parallel with model forward)

Also, the first phase can be overlapped, I have made some tries, but it failed.

dc16725
3273d3f

What I thought was to make the cumulate a lazy operation, only being done when batch size changes or before applying. This can save more time and make the inference with penalties as fast as in normal cases. If you are interested in this, you are welcome to contribute to this optimization.