[Performance]: Use int over list[int] as output_tokens to reduce GC overhead

[Jialin]

Proposal to improve performance

Currently, we're consistently using list[int] to represent output_tokens in ModelRunnerOutput which is very inefficient from GC prospective.

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!

max_batch_size of list will be created here

vllm/vllm/v1/worker/gpu_model_runner.py

Lines 4504 to 4517 in a38c1bf

	def _to_list(self, sampled_token_ids: torch.Tensor) -> list[list[int]]:
	# This is a short term mitigation for issue mentioned in
	# https://github.com/vllm-project/vllm/issues/22754.
	# `tolist` would trigger a cuda wise stream sync, which
	# would block other copy ops from other cuda streams.
	# A cuda event sync would avoid such a situation. Since
	# this is in the critical path of every single model
	# forward loop, this has caused perf issue for a disagg
	# setup.
	pinned = self.sampled_token_ids_pinned_cpu[: sampled_token_ids.shape[0]]
	pinned.copy_(sampled_token_ids, non_blocking=True)
	self.transfer_event.record()
	self.transfer_event.synchronize()
	return pinned.tolist()

Proposal #1 (Change OutputToken from list[int] to Union[int, list[int]])

#26368 shows very promising results of this proposal

• 19% throughput boost in infinite request rate scenario for facebook-125m

Proposal #2 (Semi-Hacky but simple)

Increase the GC0 threshold by max_batch_size which ensured GC0 won't be triggered right after the sample output tensor to list[list[int]] conversion. But I think this is hacky, and it sounds more like a short term mitigation over long term solution.

Report of performance regression

No response

Misc discussion on performance

No response

Your current environment (if you think it is necessary)

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[Jialin]

CC @yeqcharlotte @houseroad @WoosukKwon @njhill

[noooop]

If GC has a significant impact on output tokens, maybe we could pre-allocate memory for it, just like we manage the KV cache. lol

I’m not sure how much memory it would take to pre-allocate max_model_len * max_num_seqs * int64. If it’s only a few GBs, it’s not a crazy idea.

[Jialin]

I’m not sure how much memory it would take to pre-allocate max_model_len * max_num_seqs * int64. If it’s only a few GBs, it’s not a totally outrageous approach.

Thanks for the suggestions. But I don't think pre-allocation memory might be directly help here, as OutputToken are serialized / deserialized and pass across threads.
Also, when async scheduler is in place, we could have multiple pending to process SchedulerOutputs queued up. It might be hard for us to predict number of memory to preallocated ahead as well. So I feel a cleaner way is to directly update the underlying data to store output tokens.

But within EngineCore, the output tokens are already preallocated and the size should be very small. Would love to clarify the problems a bit here.

[noooop]

vllm/vllm/v1/worker/gpu_model_runner.py

Line 4517 in a38c1bf

return pinned.tolist()

After carefully reviewing the code, completely removing the tolist() at line L4517 and running code similar to the following somewhere.

pre_allocation_output_tokens_tensor_cpu[some_fancy_indexing] = sampled_token_ids_tensor_cpu 

In other words, find a way to eliminate the list[int] operation and transfer sampled_token_ids_tensor_cpu to Request._output_token_ids_tensor_cpu in vllm/v1/request.py

quite challenging

[Jialin]

pre_allocation_output_tokens_tensor_cpu[some_fancy_indexing] = sampled_token_ids_tensor_cpu 

I think we don't need to removing the tolist(). It's sufficient to ensure that, we don't create nested list (e.g. list[list[int]]). In our case, when spec decoding is off, the shape of sampled_token_ids_tensor_cpu is always (max_batch_size, 1). So we just need to squeeze(dim=-1) to ensure the returned value changed form list[list[int]] to list[int], which is actually the only critical change in this PR, while all other changes are compatibilities refactoring.

[noooop]

pre_allocation_output_tokens_tensor_cpu[some_fancy_indexing] = sampled_token_ids_tensor_cpu 

I think we don't need to removing the tolist(). It's sufficient to ensure that, we don't create nested list (e.g. list[list[int]]). In our case, when spec decoding is off, the shape of sampled_token_ids_tensor_cpu is always (max_batch_size, 1). So we just need to squeeze(dim=-1) to ensure the returned value changed form list[list[int]] to list[int], which is actually the only critical change in this PR, while all other changes are compatibilities refactoring.

Thank you for explaining. I think your method is very reasonable.