[V1] Optimize block table transfer from CPU to GPU #11401
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]>
Signed-off-by: Woosuk Kwon <[email protected]>
| int* d_matrix_tgt = matrix_tgt.data_ptr<int>(); | ||
|
|
||
| // One thread block per row. | ||
| int blocks = n; |
There was a problem hiding this comment.
it seems this can easily oversubscribe GPU SMs.
| int length = matrix_diff[row_id * 2 + 1]; | ||
| int end = start + length; | ||
| int thread_idx = threadIdx.x; | ||
| for (int i = start + thread_idx; i < end; i += blockDim.x) { |
There was a problem hiding this comment.
most threads in the block would be idle, e.g. for decoding, there's only one or even no entry changes in the block table.
| self.block_table_diff_np[row_idx, 0] = start | ||
| # Move-and-append is not allowed. | ||
| assert self.block_table_diff_np[row_idx, 1] == 0 | ||
| self.block_table_diff_np[row_idx, 1] = num_blocks |
There was a problem hiding this comment.
for the non-uva case, we still need to keep track of the max-block-table-length, so that apply_diff only needs to copy max-block-table-length columns.
There was a problem hiding this comment.
Good point. The problem is, the memcpy API requires the data to be in contiguous memory space: https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__MEMORY.html#group__CUDART__MEMORY_1g85073372f776b4c4d5f89f7124b7bf79
So when the block table tensor has the shape [batch_size, max_model_len] and if we slice over the second dimension, then we have to call the memcpy API batch_size times instead of once.
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
No description provided.