[PD] Respect sampling_params.max_new_tokens when PD disaggregation is activated by ShangmingCai · Pull Request #7147 · sgl-project/sglang

ShangmingCai · 2025-06-13T08:55:25Z

Motivation

This fixes a bug when PD disaggregation is activated and sampling_params.max_new_tokens <= 1

Before:

(base) root@iZ2zed0tiwrdk9tucyw2p9Z:~# curl -X POST http://127.0.0.1:8000/generate -H "Content-Type: application/json" -d '{
  "text": "You are a helpful assistant.\n",
  "sampling_params": {
    "temperature": 0,
    "max_new_tokens": 1
  }, "stream": true
}'
data: {"text":"Okay","meta_info":{"id":"05a3d0597bd642c991a293960228482d","finish_reason":null,"prompt_tokens":6,"completion_tokens":1,"cached_tokens":0}}

data: {"text":"Okay,","meta_info":{"id":"05a3d0597bd642c991a293960228482d","finish_reason":{"type":"length","length":1},"prompt_tokens":6,"completion_tokens":2,"cached_tokens":0,"e2e_latency":0.33734869956970215}}

data: [DONE]

(base) root@iZ2zed0tiwrdk9tucyw2p9Z:~# curl -X POST http://127.0.0.1:8000/generate -H "Content-Type: application/json" -d '{
  "text": "You are a helpful assistant.\n",
  "sampling_params": {
    "temperature": 0,
    "max_new_tokens": 1
  }, "stream": false
}'
{"text":"Okay,","meta_info":{"id":"87603e4404c446dc8bc4e3f69a1e6063","finish_reason":{"type":"length","length":1},"prompt_tokens":6,"completion_tokens":2,"cached_tokens":0,"e2e_latency":0.10740494728088379}}

The completion_tokens don't match the length required in the sampling params or in the finish reason.

After:

(base) root@iZ2zed0tiwrdk9tucyw2p9Z:~# curl -X POST http://127.0.0.1:8000/generate -H "Content-Type: application/json" -d '{
  "text": "You are a helpful assistant.\n",
  "sampling_params": {
    "temperature": 0,
    "max_new_tokens": 1
  }, "stream": true 
}'
data: {"text":"Okay","meta_info":{"id":"13d7bc930ff74b65b8048e52527e8aea","finish_reason":{"type":"length","length":1},"prompt_tokens":6,"completion_tokens":1,"cached_tokens":0,"e2e_latency":0.3873429298400879}}

data: [DONE]

(base) root@iZ2zed0tiwrdk9tucyw2p9Z:~# curl -X POST http://127.0.0.1:8000/generate -H "Content-Type: application/json" -d '{
  "text": "You are a helpful assistant.\n",
  "sampling_params": {
    "temperature": 0,
    "max_new_tokens": 1
  }, "stream": false
}'
{"text":"Okay","meta_info":{"id":"c8bbea881e054239b75b9c6431e3e2a4","finish_reason":{"type":"length","length":1},"prompt_tokens":6,"completion_tokens":1,"cached_tokens":0,"e2e_latency":0.18650269508361816}}

Modifications

The reason behind this is that we get the first token ids from the prefill node and return the first two tokens when the decode instance finishes the first forward pass. We didn't want to cause a weird waiting time and introduce a concept of Time to Second Token, but we still need to respect the sampling_params.max_new_tokens. This fix is simple by skipping the second token in the streaming output. We should not use req.check_finished() since it will cause a memory leak when disaggregation PD is on.

CC: @ByronHsu @acelyc111 @fzyzcjy @whybeyoung

Checklist

Format your code according to the Code Formatting with Pre-Commit.
Add unit tests as outlined in the Running Unit Tests.
Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

…tivated Signed-off-by: Shangming Cai <caishangming@linux.alibaba.com>

gemini-code-assist

Summary of Changes

Hello @ShangmingCai, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request fixes a specific issue in the system's handling of max_new_tokens when Prefill/Decode disaggregation is enabled, particularly for requests with a maximum token limit of 1 or less. The core change ensures the system correctly respects the requested output length and reports the accurate finish reason under these conditions.

Highlights

Bug Fix for max_new_tokens: I've addressed a bug where the reported completion_tokens and finish_reason were incorrect when Prefill/Decode (PD) disaggregation was active and sampling_params.max_new_tokens was set to 1 or less. The system was incorrectly generating and reporting two tokens instead of respecting the requested limit.
Respecting Length Constraint: The changes ensure that when in decode disaggregation mode and max_new_tokens <= 1, the system now correctly stops after the specified number of tokens (usually 1) and sets the finish_reason to 'length'. This prevents streaming of extra tokens and ensures accurate output length.
Tokenizer Manager Handling: I've added logic in the tokenizer manager to handle requests that are finished early due to this length constraint in decode disaggregation mode. This prevents unnecessary error logging related to missing request states when the request is intentionally completed.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configureGemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request aims to fix a bug where sampling_params.max_new_tokens was not correctly respected when PD disaggregation is active, especially for small values of max_new_tokens. The changes involve adjusting the logic in scheduler_output_processor_mixin.py to control token output and correctly set the finish reason, and modifying tokenizer_manager.py to prevent erroneous error logging in early finish scenarios.

My review focuses on ensuring the conditional logic precisely matches the intended scope described in the comments and PR description, clarifying potential discrepancies, suggesting a minor code simplification, and enhancing robustness against a potential runtime error. The changes appear to correctly address the core issue demonstrated in the examples.

ByronHsu

looks a bit hacky. i wonder if we can just skip the first decode step?

ShangmingCai · 2025-06-14T03:10:37Z

looks a bit hacky. i wonder if we can just skip the first decode step?

Sure. Let's fix this at the source.

[PD] Respect sampling_params.max_new_tokens when disaggregation is ac…

f63cba8

…tivated Signed-off-by: Shangming Cai <caishangming@linux.alibaba.com>

ShangmingCai requested review from Ying1123, hnyls2002, merrymercy and xiezhq-hermann as code owners June 13, 2025 08:55

gemini-code-assist Bot reviewed Jun 13, 2025

View reviewed changes

ByronHsu reviewed Jun 14, 2025

View reviewed changes

ShangmingCai closed this Jun 14, 2025

ShangmingCai mentioned this pull request Jun 27, 2025

[PD] Respect sampling_params.max_new_tokens when PD disaggregation is activated #7598

Merged

6 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[PD] Respect sampling_params.max_new_tokens when PD disaggregation is activated#7147

[PD] Respect sampling_params.max_new_tokens when PD disaggregation is activated#7147
ShangmingCai wants to merge 1 commit intosgl-project:mainfrom
kvcache-ai:fix_pd_max_new_tokens

ShangmingCai commented Jun 13, 2025

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

ByronHsu left a comment

Uh oh!

ShangmingCai commented Jun 14, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

ShangmingCai commented Jun 13, 2025

Motivation

Modifications

Checklist

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

ByronHsu left a comment

Choose a reason for hiding this comment

Uh oh!

ShangmingCai commented Jun 14, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants