[Core] Performance: Use list[np.ndarray] instead of list[list[int]] for output tokens for GC optimization

[Jialin]

Purpose

The default setup of GC is (700, 10, 10) which means

• if allocated_obj-deallocated_obj>=700 in generation 0, GC0 will be triggered
• GC1 is triggered after 10 GC0
• GC2 is triggered after 10 GC1
  In this scenario, large batch size scenarios (small models) each batch could be as large as 1024, which means GC0 will be triggered per decode cycle, GC1 will triggered per 10 decode cycle and GC2 per 100 decode cycle, which is very inefficient!

In this PR, we change output tokens from list[list[int]] to list[np.ndarray] to cut down objects counts from <batch_size> to 1, which would significantly reduce GC overhead especially for large batch size use cases.

Test Plan & Test Result

• Test 1: facebook/125m TP1 Input 48 Output 2000
• Test 2: facebook/125m TP1 Input 48 Output 500
• Test 3: Llama3 8B TP2 Input 48 Output 500

Request Throughput Change

• Test 1: facebook/125m Long Output: 30+% (from 10.04 request/s -> 13.11 request/s)
• Test 2: facebook/125m Short Output: 5+% (from 48.95 request/s -> 51.66 request/s)
• Test 3: Llama3 8B: 6+% (from 28.93 request/s -> 30.70 request/s)

With the change GC improved significantly.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[Jialin]

CC @yeqcharlotte @houseroad @WoosukKwon @njhill for RFC on the proposal and high level code changes, but the PR might not be completely ready yet

• Callers might not be updated completely
• Unit tests are still WIP

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

[gemini-code-assist]

Code Review

This pull request introduces a clever optimization to reduce garbage collection overhead by changing the representation of single output tokens from a list to an integer. The use of Union[int, list[int]] is a good way to handle both single and multiple token outputs efficiently. The changes are consistently applied across the codebase. My only feedback is to add unit tests for the new utility functions to ensure the correctness and robustness of this performance-critical change.

[Jialin]

Resolve #26369

[Jialin]

Gentle nudge @houseroad for the review :P

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Jialin.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[yeqcharlotte]

thanks for the change. hhh it's hard to keep classes around to avoid GC. wondering what's the actual gain we see from bigger models? we can play with deepseek v3, qwen-235b etc. and try it for both throughput and latency sensitive use cases.

wondering if you folks have strong opinion for this @heheda12345 @njhill

[zhuohan123]

I am a bit hesitant on this PR. I feel it's over optimization. To achieve low latency, we probably should invest in:

• Replace data structures with numpy arrays
• Async scheduling
• If Python overhead is actually that serious, we probably should seriously think about rewriting the scheduler in C++.

I am worried that optimizations like the one in this PR can pile up and make the code much harder for people to understand without actually improving the perf very significantly.

[Jialin]

I am a bit hesitant on this PR. I feel it's over optimization. To achieve low latency, we probably should invest in:

• Replace data structures with numpy arrays
• Async scheduling
• If Python overhead is actually that serious, we probably should seriously think about rewriting the scheduler in C++.

I am worried that optimizations like the one in this PR can pile up and make the code much harder for people to understand without actually improving the perf very significantly.

Thanks @zhuohan123 for your suggestions and I totally aligned with the tradeoff between code complexity and performance. But would love to better clarify my change in this PR.

• Motivation: reduce GC collect frequency by using list[int] for single element scenarios, but not due to to performance of list[int] itself. PR summary also had more details about the rational.
• Improvements: we've included a bit more validations in the PR summary, overall, we believe the improvement is significant for small models and large batch size scenarios (facebook/opt125m 30+% throughput improvement and Llama 8B 6+% throughput improvements)
• Alternatives:
  • Replace with numpy arrays: That's a great call and I didn't think of earlier. I think it's cleaner also won't involve in GC and more friendly for both non-spec decoding and spec decoding scenarios. But that also required some code changes, so would love to get more of your opinions in this change before investing more time on this direction.
  • Async scheduling won't help in this scenario, as long as the data structure stay as is, GC would kick in regardless and pause the whole process.
  • Just rewrite scheduler in C++ also might not help, as long as we still use python interface and use list[int] to pass single output tokens around, GC would still happen.

Next step: @zhuohan123 per the new data and new clarification I provided, do you feel that's something we should worth to further investigate? If yes, I could try to explore the np.array suggestion more. Appreciate that.

[Jialin]

we can play with deepseek v3, qwen-235b etc. and try it for both throughput and latency sensitive use cases.

@yeqcharlotte To be very honest, this change might only be helpful for small models and large batch size scenarios. So the effectiveness could be small to large models like deepseek v3 and qwen 235b :/

[DarkLight1337]

V1 Test others has been failing since this PR: https://buildkite.com/vllm/ci/builds/39126/steps/canvas?sid=019a84f2-ee8a-46bd-9845-3bc144e4cf4b

[Jialin]

V1 Test others has been failing since this PR: https://buildkite.com/vllm/ci/builds/39126/steps/canvas?sid=019a84f2-ee8a-46bd-9845-3bc144e4cf4b

Sorry for the inconvenience. I'll take a look later in the day for a quick fix forward. If needed, please feel free to revert for early unblock.

[DarkLight1337]

The fix seems to be straightforward enough, opened #28771

[WoosukKwon]

@Jialin @DarkLight1337 @njhill A dumb question: Can we further optimize the data structure by using two numpy arrays, sampled_tokens: int32[num_reqs, max_num_generated_tokens] and num_sampled_tokens: int32[num_reqs]?

[Jialin]

@Jialin @DarkLight1337 @njhill A dumb question: Can we further optimize the data structure by using two numpy arrays, sampled_tokens: int32[num_reqs, max_num_generated_tokens] and num_sampled_tokens: int32[num_reqs]?

@WoosukKwon Thanks for the suggestion! But I think there're some trade off of your proposal, and I might not be the most critical GC bottleneck as of now.

Purely from GC prospective, as numpy array don't even go through GC, the originally approach would only introduce one object (the list that contains sampled_tokens) per decode batch.
Your proposal could even remove the object involved as well (that's the sweet part :) )

However, I feel there're a few downside:

• memory overhead due to we need to pad all lanes to max_decode_tokens
• we might need to introduce more handy utils to manipulate the sample tokens (so the maintenance cost is also lifting)

Recently I think I found a pretty good way (with tracemalloc) to analyze the code and objects that introduce the largest GC overhead. I'll update gc_utils.py accordingly to make it accessible to others. With that we could evaluate the GC bottleneck easily and optimize them in order, then we could optimize GC costs more objectively (given that all the optimization might unfortunately had tradeoff, we should choose to optimize critical ones only and have a way to justify the criticality of the issue).

WDTY? :)