[V1] Use more persistent buffers to optimize input preparation overheads #11111
Conversation
Signed-off-by: Woosuk Kwon <[email protected]>
|
👋 Hi! Thank you for contributing to the vLLM project. 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:
🚀 |
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
WoosukKwon
requested review from
robertgshaw2-neuralmagic,
njhill,
ywang96,
comaniac and
alexm-neuralmagic
as code owners
|
Nice idea! |
| dtype=torch.int32, | ||
| device="cpu", | ||
| pin_memory=self.pin_memory) | ||
| self.slot_mapping_np = self.slot_mapping_cpu.numpy() |
There was a problem hiding this comment.
Does the resulting numpy here shares the memory buffer of the source tensor?
There was a problem hiding this comment.
The returned ndarray and the tensor will share their storage, so changes to the tensor will be reflected in the ndarray and vice versa.
Yes. That's the trick here :)
…ads (vllm-project#11111) Signed-off-by: Woosuk Kwon <[email protected]> Signed-off-by: Akshat Tripathi <[email protected]>
…ads (vllm-project#11111) Signed-off-by: Woosuk Kwon <[email protected]>
…ads (vllm-project#11111) Signed-off-by: Woosuk Kwon <[email protected]>
This PR simplifies the input preparation code further while optimizing it by utilizing more persistent buffers. Creating new tensors can introduce considerable overhead for small-batch inputs, so persistent buffers effectively reduce latency.