Delay GPU->CPU sync in sampling

[Yard1]

This PR preallocates tensors used to prune hidden states and to index the samples by sampling type allowing us to delay the GPU->CPU sync slightly longer (up until torch.multinomial in the sampler).

Should fix the slight performance regression introduced in #1309

Technically, we can delay the sync even further by reordering some more operations, but that's left for a future PR.

[Yard1]

cc @zhuohan123

[Yard1]

updating the PR!

[hanzhi713]

This would be a nice improvement to have. Since you're adding more h2ds, a small optimization would be to batch the h2ds together and make it async.

[Yard1]

Yes, we could certainly optimize that as well - I would like to do that in a followup.

[Yard1]

@zhuohan123 Updated, PTAL!

[zhuohan123]

LGTM! Thank you for your contribution! Left some small comments based on recent PRs. Let's refactor the bloated input_metadata in a future PR.

[Yard1]

Need to fix max_prompt_len in worker

[hanzhi713]

Any progress on this PR?

[Yard1]

@hanzhi713 I am on holidays currently but will wrap it up next week!

[Yard1]

Should be good now

[zhuohan123]

LGTM! Thank you for your contribution!

[gesanqiu]

@Yard1 @zhuohan123 Have you been test the sampling performance on AWQ model? I found that the AWQ model's first sampling process is much slower than FP16 model, but faster in the following tokens. I can't figure the exact bottleneck out :(

[zhaoyang-star]

@Yard1 @zhuohan123 Sorry for that I am not familar with selected_token_indices.

Assume we have 3 prompts. The length are 9, 8, 6. On the prompt (prefill) stage, selected_token_indices=[5, 14, 23]. But I think that selected_token_indices should be [8, 16, 22]. Below is a diagram, which is relatively more intuitive. I have been struggling to understand it here, and I would appreciate your interpretation.