Adding basic kv-cache transfer to vllm v1

[mrn3088]

A working implementation for vllm to do kvcache transfer between prefill and decode engine.
Tested on the examples/offline_inference.py.
To run it, using the following commands in two terminal:

Prefill Engine

vllm/examples# VLLM_PORT=47651 python3 offline_inference.py --dist-factor 2 --rank 0 --local-rank 0 --role prefill --max-tokens 1

Decode Engine

vllm/examples# VLLM_PORT=47651 python3 offline_inference.py --dist-factor 2 --rank 1 --local-rank 1 --role decode --max-tokens 16

The first process will execute the model for one step (the prefill step), and send the hidden_state and kv_cache to the second process to complete the following decode.

For now, at least it's working.

Decode Engine Output, seems correct

Processed prompts: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:01<00:00,  3.18it/s, est. speed input: 20.65 toks/s, output: 50.82 toks/s]
Prompt: 'Hello, my name is', Generated text: " Joel. I'm from Massachusetts and live in Melbourne, Australia.\nI'm"
Prompt: 'The president of the United States is', Generated text: ' about to be arrested in Europe for allegedly meddling in the 2016 election.\n\n'
Prompt: 'The capital of France is', Generated text: ' becoming a state of chaos with a significant urban and industrial boom. France’'
Prompt: 'The future of AI is', Generated text: ' not as simple as you think, and you have to understand it in order to'

Things to pay attention:

1. The current implementation has not been tested on other examples. Especially online examples. I'll check this later today.
2. Current kv transfer is not efficient. Ideally should only do one send/recv that transfer all kv_cache along with hidden_state.
3. What about more engines? Currently, the TP and PP is disabled since they aren't compatible yet.
4. Lots of hardcoded stuffs...

[Jocn2020]

Nice work on the rank hack!
Quick comment from me is currently you send the entire kvcache. Next step we want to do is just sending the kvcache of specific requests' block ids which you can find in scheduler_output.scheduled_new_reqs or scheduler_output.scheduled_resumed_reqs

[github-actions]

This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you!