Fix: PrefillAdder.add_chunked_req with negative rem_total_tokens with pp

[strgrb]

Motivation

I need to run Ring-1T and Ling-1T model with tp8pp4, and met up with an error

The error may differ, but new-token is negative every time. It does not appear without pp.
This pr try to fix this problem.

Finally I found self.rem_total_tokens become negative in PrefillAdder.add_chunked_req , it's computed by available_and_evictable - self.rem_total_token_offset , and budgets for running requests are added to self.rem_total_token_offset .
After merging of last_batch and running_batch, this budget will increase, and if this increase to an amount larger than available size, PrefillAdder.add_chunked_req will calculate a negative extend input len.

Modifications

Since chunked_req's req_pool_idx is freed in Scheduler.get_next_batch_to_run , we should check remaining tokens excluding budgets for running requests here, and avoid freeing chunked_req's req_pool_idx. Since budget is not enough for decoding now, it's time for decoding, and chunked_req is hanged, waiting for budget is enough.

Accuracy Tests

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.
• Work with maintainers to merge your PR. See the PR Merge Process

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[ShangmingCai]

Will review this PR today. @strgrb Can you test whether this bug is still happening in #11852?

cc: @XucSh @whybeyoung

[strgrb]

Will review this PR today. @strgrb Can you test whether this bug is still happening in #11852?

cc: @XucSh @whybeyoung

OK, I'll try it.

[ShangmingCai]

Changes don't look like it is related to PP, more likely related to bybrid memory.

CC: @xiezhq-hermann @yizhang2077

[strgrb]

Changes don't look like it is related to PP, more likely related to bybrid memory.

CC: @xiezhq-hermann @yizhang2077

It's not related to hybrid memory, it's just for budget calculation where hybrid is used. The real reason is about budget, i.e. budget for chunked request is negative with pp situation.

[strgrb]

@ShangmingCai Should I move this budget check logic to PrefillAdder , which rem_total_tokens() can be used directly?

[ShangmingCai]

@ShangmingCai Should I move this budget check logic to PrefillAdder , which rem_total_tokens() can be used directly?

@strgrb If it is related to pp only, maybe you could try reverting this commit locally: #13144. If the bug still exists, then it might not be related to the pp. Also, you can try changing the attention backend to test whether this is a bug of flashinfer?

[xiezhq-hermann]

let's get this problem fixed after this refactoring to simplify the logics