[Performance] Optimize e2e overheads: Reduce python allocations

[alexm-redhat]

This PR introduces a bunch of end-to-end overhead optimizations to reduce python object allocations/deallocations over scheduler iterations. In particular:

1. Avoid python object allocations for "InterDataForSeqGroup" objects: These objects have lot of fields, and most of the them were allocated dynamically. In this PR, these objects are pre-allocated and reused between runs. The pre-allocation is done per "number of sequences per group" (to support beam search), where most of the time we simply have one sequence per group.
2. Avoid python object allocations for "SequenceGroupMetadata" inside scheduler, in the same way as for (1).
3. Avoid python object allocations for "SequenceGroupToSample" inside SamplingMetadata.prepare(..)
4. Add an optimization to scheduler to be aware of "one sequence per group" case (since most of them we do not run beam search).
5. A bunch of misc small stuff to reduce object allocations

End-to-end throughput of Llama3 8B on 1xH100 is 24% faster with this PR. Command used:
python3 benchmark_throughput.py --model meta-llama/Meta-Llama-3.1-8B-Instruct --backend vllm --input-len 512 --output-len 256 --num-prompts 1000 --tensor-parallel 1

Main branch from 08/05/2024:
Throughput: 19.48 requests/s, 14962.18 tokens/s
This PR:
Throughput: 24.32 requests/s, 18123.51 tokens/s

[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 consists a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of default ones by unblocking the steps in your fast-check build on Buildkite UI.

Once the PR is approved and ready to go, please make sure to run full CI as it is required to merge (or just use auto-merge).

To run full CI, you can do one of these:

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

🚀

[comaniac]

This is one way to introduce incremental prepare input. Will review soon.

cc @youkaichao @zhuohan123

[comaniac]

I took a brief look and have the following impressions:

1. This optimization applies to continuous decoding steps. In other words, if requests in the batch are added or removed, the cache is invalid.
2. This optimization doesn't work with chunked prefill as it will schedule prefill and decode requests in the same batch.

Please correct me if I misunderstood anything. IIUC, It seems to me that this optimization should be done with multi-step worker, as it focuses exactly on the scenarios listed above. For general cases, this optimization introduces decoding specific branches (e.g., if no prefill tokens) and we should avoid that as possible.

[alexm-redhat]

@comaniac and I discussed the relation to multi-step. In general, a micro-benchmark shows that the benefit from flash-attn decode optimization in this PR is around 2-3%, and most of the benefit comes from python object allocation reductions. There are two possible ways to go about this:

1. Remove flash-attn optimization from this PR in favor of multi-step (since multi-step should cover flash-attn decode)
2. Proceed with flash-attn optimization until multi-step is fully enabled.

[alexm-redhat]

/ready

[njhill]

Thanks @alexm-neuralmagic this is great!

1. This optimization applies to continuous decoding steps. In other words, if requests in the batch are added or removed, the cache is invalid.

I think (1) isn't a concern since the proportion of steps in which reqs are added or removed is typically small.

2. This optimization doesn't work with chunked prefill as it will schedule prefill and decode requests in the same batch.

I didn't look closely enough yet but it presumably it will still work once in this mode once the prefill tokens have been exhausted i.e. once the batch returns to decode-only state?

1. Remove flash-attn optimization from this PR in favor of multi-step (since multi-step should cover flash-attn decode)
2. Proceed with flash-attn optimization until multi-step is fully enabled.

Given my comments above I think both parts of this would still make sense to include, but perhaps separate into two separate PRs anyhow?

[njhill]

@alexm-neuralmagic the speed-up from allocation reduction is really encouraging!

Added some comments from quick glance through, will try to look more closely soon.

[njhill]

could de-dup a bit here by just setting those three values as vars in the if/else and then having a single call to the constructor?

[alexm-redhat]

Thanks for the suggestion, applied this change in #7206

[njhill]

Could make this a bit more concise (same for others below):

Suggested change

	if self.is_single_seq:
	if self.seqs[0].status == status:
	return self.seqs
	else:
	return []
	else:
	return [seq for seq in self.seqs if seq.status == status]
	if self.is_single_seq:
	return self.seqs if self.seqs[0].status == status else []
	return [seq for seq in self.seqs if seq.status == status]

[alexm-redhat]

Good idea, changed all of them to the form you proposed.

[njhill]

I feel like the above could be simplified .. maybe collect all the things to be cleared and all the things to be zeroed and then have a single loop over the sequences? These lists themselves could possibly also be reused.

[alexm-redhat]

This one is harder to generalize because it uses different fields. I tried to do lists, but then you still need to separate to different fields and it complicates things again. I will see if I can make the code cleaner here.

[comaniac]

Please note that we were planning to use dataclass with kw_only=True, but it is supported in Python 3.10 so we fallback to the current implementation. We should eventually use dataclass tho.

[njhill]

The standard python way of doing this is:

input_tokens = [
    in_toks for inter_data in self.inter_data_list
    for in_toks in inter_data.input_tokens
]

I'd guess that this would be more efficient but can never be sure, would have to microbench it...

[alexm-redhat]

Yeah, this is what flatten_2d_lists(..) was doing, but extend() is faster.

[comaniac]

We should change flatten_2d_list to use extend then.

[njhill]

The prior code wasn't doing the same thing. Technically this isn't flattening 2d lists - it's a list of objects each having a list inter_tokens field. So flatten_2d_lists() isn't appropriate here anyhow. If you look at the prior code is calling flatten_2d_lists twice, which is probably the reason for the excess allocations and slower speed.

[comaniac]

Hmm yeah you're right. I was hoping there's a way to extract this logic to be a utility function, but seems not trivial.

[alexm-redhat]

I see

[njhill]

I ran a benchmark, I guess the extend version is still much faster ! (more than 2x)

[comaniac]

Thanks @alexm-neuralmagic this is great!

1. This optimization applies to continuous decoding steps. In other words, if requests in the batch are added or removed, the cache is invalid.

I think (1) isn't a concern since the proportion of steps in which reqs are added or removed is typically small.

Isn't we continue adding/removing decoding requests at every step as long as there are new requests coming?

2. This optimization doesn't work with chunked prefill as it will schedule prefill and decode requests in the same batch.

I didn't look closely enough yet but it presumably it will still work once in this mode once the prefill tokens have been exhausted i.e. once the batch returns to decode-only state?

I'm not sure if prefill requests exhausting is a reasonable assumption especially for high QPS.

In summary I feel your assumptions make sense in offline batching, but I'm not sure about online serving. Meanwhile, considering the code complexity this PR will introduce, I personally would prefer to isolate it to the scenario it fits into. Of course we could still take the general Python code optimization.

[njhill]

Thanks @alexm-neuralmagic this is great!

1. This optimization applies to continuous decoding steps. In other words, if requests in the batch are added or removed, the cache is invalid.

I think (1) isn't a concern since the proportion of steps in which reqs are added or removed is typically small.

Isn't we continue adding/removing decoding requests at every step as long as there are new requests coming?

Offline is getting more interest lately with SDG etc. It obviously depends on the particular workload in terms of input/output split sizes and request patterns but I think even in the online serving case the batch constituency doesn't change for a significant proportion of steps. We should probably add metrics for this if we don't have already :)

2. This optimization doesn't work with chunked prefill as it will schedule prefill and decode requests in the same batch.

I didn't look closely enough yet but it presumably it will still work once in this mode once the prefill tokens have been exhausted i.e. once the batch returns to decode-only state?

I'm not sure if prefill requests exhausting is a reasonable assumption especially for high QPS.

Again I think this is very workload dependent, what about for "Write me an essay about x" type use cases.

In summary I feel your assumptions make sense in offline batching, but I'm not sure about online serving. Meanwhile, considering the code complexity this PR will introduce, I personally would prefer to isolate it to the scenario it fits into. Of course we could still take the general Python code optimization.

I agree about taking care to minimze additional complexity.

[alexm-redhat]

@njhill @comaniac thanks for the review and comments. @njhill I think it is a good idea to separate the optimizations. I will remove the flash-attn from this PR and will focus on python optimizations first, to avoid a conflict with mutli-step worker PR from Will.

[comaniac]

Extract this common part to be an inner function?

[alexm-redhat]

Done

[comaniac]

I'm a bit worry about this because block_table and block_table_ids are not strongly associated. If someone updated block_table somewhere else without using the update function then these 2 attributes are mismatched. Can we make block_table and block_table_ids properties and update them with setters? In this way you also don't need to change other parts in block manager

[alexm-redhat]

This is a bit complicated to do since "block_tables" is a dict of seq_ids . I have introduced update_block_table() and append_block() object functions to modify these two variables together, and made sure that every place in the class that needs to modify block_tables, is using these two functions. I could also change the BlockTable type from a List[PhysicalTokenBlock] to a full class object, so it can hold the cached ids, like we did for block_manager_v2. Not sure what's better here.

[alexm-redhat]

@comaniac I was able to fix it by doing the same thing we did originally for block_manager_v2. By introducing a single class that holds the "ids" and overrides "list methods", the code in block_manager_v1 does not need to change anymore and is cleaner.

[comaniac]

1. I feel we could make this cache a common class, and use it like metadata=PyObjectCache(SequenceGroupMetadata).
2. Please add docstring to the class and methods.

[alexm-redhat]

Good idea! Introduced a single python object caching class and reused it in both cases.

[comaniac]

It may be more straightforward to use a single API for this, such as

Suggested change

	seq_group_metadata = seq_group_metadata_cache.get_object()
	seq_group_metadata.__init__(
	seq_group_metadata = seq_group_metadata_cache.new_object(

[alexm-redhat]

Refactored to a single API call in both cases

[comaniac]

Please note that we were planning to use dataclass with kw_only=True, but it is supported in Python 3.10 so we fallback to the current implementation. We should eventually use dataclass tho.

[comaniac]

Can we use a common cache class for this as well instead of introducing another similar class?

[alexm-redhat]

Refactored to a common class

[comaniac]

We should change flatten_2d_list to use extend then.

[njhill]

The list() is unnecessary here

[alexm-redhat]

Good catch, removed.

[njhill]

Can avoid allocating lists here, same below

Suggested change

	input_tokens.extend([0] * cuda_graph_pad_size)
	input_positions.extend([0] * cuda_graph_pad_size)
	input_tokens.extend(itertools.repeat(0, cuda_graph_pad_size))
	input_positions.extend(itertools.repeat(0, cuda_graph_pad_size))

[alexm-redhat]

Nice trick, did not knew about this one.

[comaniac]

LGTM! Only nits

[comaniac]

docstring plz.

[alexm-redhat]

Added

[comaniac]

Suggested change

	for _, cache in cls.inter_data_cache.items():
	for cache in cls.inter_data_cache.values():

[alexm-redhat]

nice, changed

[comaniac]

@youkaichao PTAL and let's try to merge this PR by today or tomorrow.

[youkaichao]

Suggested change

	self.block_tables[wait_seqs[0].seq_id] = block_table
	else:
	for seq in seq_group.get_seqs(status=SequenceStatus.WAITING):
	self.block_tables[seq.seq_id] = block_table.copy()
	self.block_tables[seq.seq_id] = block_table
	else:
	for seq in wait_seqs:
	self.block_tables[seq.seq_id] = block_table.copy()

does this make sense?

[alexm-redhat]

Addressed #7364

[youkaichao]

not related with this PR, but I think we can do it in a followup PR, to refactor the way we store the data, e.g. self.cached_data[seq_id].input_positions. Then we can have just one for-loop to iterate over seq_id, and set all fields in one pass.

[youkaichao]

thanks for the great optimization! my previous concern is multiple instances might share (and reuse) the same object. This is fixed now.

left one nit comment and one possible followup improvement. LGTM in general!

[Henson-zh]

Maybe cuda_graph_pad_size > 0:

[alexm-redhat]

working on fixing tests

[alexm-redhat]

fixed all of the failing tests, should be green now I hope.

[njhill]

Suggested change

	self._obj_cache = []
	for _ in range(128):
	self._obj_cache.append(self._obj_builder())
	self._obj_cache = [self._obj_builder() for _ in range(128)]

[njhill]

I ran a benchmark, I guess the extend version is still much faster ! (more than 2x)

[Henson-zh]

why no shortcut here?

[alexm-redhat]

Good catch, addressed in #7364

[njhill]

Thanks @alexm-neuralmagic for this, it's a huge speedup!

It looks like there's a still couple of unaddressed comments though? #7162 (comment) and #7162 (comment)

[alexm-redhat]

@njhill no problem, to address the leftover comments I have posted a small followup: #7364