[Sampler] Vectorized sampler

[Yard1]

Ready for initial review.

Using the modified throughput benchmark with python benchmarks/benchmark_throughput.py --backend vllm --dataset "./ShareGPT_V3_unfiltered_cleaned_split.json" --model meta-llama/Llama-2-13b-chat-hf --tokenizer meta-llama/Llama-2-13b-chat-hf --num-prompts=1000 and one A100-80, I get the following results:

baseline: Throughput: 3.70 requests/s, 1771.58 tokens/s
this PR: Throughput: 4.77 requests/s, 2280.04 tokens/s

which would be a 1.28x improvement.

The argmax outputs between this PR and master match. Random sampling is different, but that is expected (the text still makes sense and doesn't differ much from master).

[Yard1]

Updated benchmark to also apply top_p.

baseline: Throughput: 3.58 requests/s, 1710.60 tokens/s
PR: Throughput: 4.69 requests/s, 2243.01 tokens/s

1.31x improvement

[Yard1]

I have also compared argmax outputs of the benchmark on the ShareGPT dataset, and they are identical between master and this PR (meaning that there is no impact on correctness).

[Yard1]

Using modified benchmark and llama-2-7b on A100-80:

baseline: Throughput: 4.67 requests/s, 2233.30 tokens/s
PR: Throughput: 6.84 requests/s, 3272.51 tokens/s

x1.47 improvement

As expected, the improvement is bigger the smaller the model (and thus more time is spent in sampling).

[scv119]

benchmark llama70-b with A100-40 * 4, 512 context length and 128 generation, 500 requests

baseline: dur_s 133.22 tokens_per_s 1807.24 qps 3.75 successful_responses 500
PR: dur_s 114.26 tokens_per_s 2108.07 qps 4.38 successful_responses 500 

[WoosukKwon]

@Yard1, thanks for submitting the PR. Before I go deeper into the PR, could you please update the branch with the latest commit?

[WoosukKwon]

A dumb question: What is this variable used for?

[Yard1]

Allows us to alternate between greedy and stochastic sampling for requests

[WoosukKwon]

I guess this is to check if vLLM generated valid outputs, right? If so, I believe this ought to be done by our test code, rather than in the benchmarking code.

[Yard1]

Yeah, I will revert the changes to the benchmarking script later.

[scv119]

also consider add a test?

[scv119]

lgtm

[WoosukKwon]

@Yard1 I got an error when running examples/llm_engine_example.py:

File "/home/workspace/vllm/vllm/model_executor/layers/sampler.py", line 98, in forward
    assert len(temperatures) == non_greedy_logits.shape[0]
AssertionError

[Yard1]

@WoosukKwon Good catch! Just pushed a fix.

[WoosukKwon]

@Yard1 Thanks for the PR. It seems the performance improvement is quite significant! However, I'm a bit worried that this PR 1) includes unnecessary changes, and 2) complicates the sampler logic.

For 1), could you revert back the unnecessary changes? I pointed out a few places in the code.
For 2), do you have any idea to simplify the code? I think the complexity mainly comes from a lot of tensor/list indexing operations, and sort-gathers. Can you somehow reduce the use of them?

[WoosukKwon]

What is the benefit of creating p and k outside the function? It seems the change has no effect.

[Yard1]

right now it's necessary for the p's and k's to be indexed correctly

[Yard1]

@WoosukKwon Thanks for the feedback, let me see if I can reduce unnecessary changes.

As for the compelxity, the current code is optimized for memory and speed which sacrifices simplicty. I am not sure if it's possible to reduce the complexity without impacting performance. I am happy to add more comments and change the code structure for more readability.

[Yard1]

Closed in favor of #1048