[V1][Perf] Reduce scheduling overhead in model runner after cuda sync

[youngkent]

We do some runner bookkeeping CPU operations after decoding iteration. We could parallelize some bookkeeping work while waiting on cuda sync. After the cuda sync, we only need to do simple and fast updates.

The change should reduce scheduling overhead between decode iterations by ~20%. (See attached gpu trace)

Before the optimization,

After the optimization,

E2E latency benchmark, ran

VLLM_USE_V1=1 python3 benchmarks/benchmark_latency.py --model "/data/users/ktong/llama/llm_8b_oss" --tensor-parallel-size 1 --input_len 1000 --batch_size 32

Output (1-2% e2e latency reduction):
Avg latency: 2.338167402730323 seconds
10% percentile latency: 2.3207896508742123 seconds
25% percentile latency: 2.3264574960339814 seconds
50% percentile latency: 2.3333765944698825 seconds
75% percentile latency: 2.343035737867467 seconds
90% percentile latency: 2.3567665563430635 seconds
99% percentile latency: 2.3934816433605737 seconds

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can do one of these:

• Add ready label to the PR
• Enable auto-merge.

🚀

[youkaichao]

is this necessary?
iirc, @tlrmchlsmth use List[int] because they are cheaper to serialize, and would benefit tensor parallel case, where we need to pass them across processes.

[tlrmchlsmth]

This is true — I didn’t look at how it impacts the non-TP case though

[robertgshaw2-redhat]

The ModelRunnerOutput is what we serialize for TP, we don't serialize the SamplerOutput directly, so this is not a concern

[tlrmchlsmth]

Ah, yep that's right -- I did change this line in #9856, but that was just downstream of changing sampled_token_ids to a List in the ModelRunnerOutput. This looks good to me since that's left as-is!

[robertgshaw2-redhat]

Wow, great idea. Im going to run some perfomance analysis on this tomorrow.

[WoosukKwon]

LGTM! Thanks for discovering and fixing this!

[WoosukKwon]

Just for a record: If top-p or top-k sampling is used (with the FlashInfer kernel), CPU-GPU synchronization happens inside the sampler at

vllm/vllm/v1/sample/ops/topk_topp_sampler.py

Lines 193 to 194 in 324960a

	# NOTE: CPU-GPU synchronization happens here.
	if not success.all():

[robertgshaw2-redhat]

Do you think we can avoid this in a follow up PR?

[WoosukKwon]

I don't think we can. This is a fundamental limitation of the kernel (or the algorithm itself). The rejection sampling method cannot 100% guarantee the success.

[robertgshaw2-redhat]

It might be faster to do sampler_output.sampled_token_ids.cpu()and thensampler_output.sampled_token_ids[i].item()` in the inner loop.

[WoosukKwon]

In my experience, item() took considerable time so should be avoided.

[WoosukKwon]

@youngkent Thanks for the PR! This change helps vLLM's performance noticeably.