[ Frontend ] Multiprocessing for OpenAI Server with zeromq

[robertgshaw2-neuralmagic]

RFC: #6797

SUMMARY:

• Use multiprocessing for OpenAI Server such that it does not contend with AsyncLLMEngine
• Use zeromq for networking between OAI process and AsyncLLMEngine
• Currently uses pickling for python serialization, to update this to use protobufs in separate PR
• Gives ~20% TPOT improvement for QPS 10 for llama-8b-instruct on H100 (550 input tokens, 150 output tokens)

TODO:

• enabling aborting
• enable do_log_stats
• error handling
• refactor to use single socket
• use open port
• turn on the other OAI endpoints
• add flag to disable mp
• tokenizer
• guided decoding
• tracing
• make CI green

PERFORMANCE:

• zeromq based protocol is delivering speedup in high QPS scenario:
• Llama-3-8B-Instruct, 1xH100, 550 input tokens | 150 output tokens, 10 QPS, chunked prefill

Metric	main	grpc	zmq + pickle (this PR)	zmq + protobuf
TPOT (mean)	70.74	64.15	57.35	54.54
TTFT (mean)	687.69	450.05	298.09	269.01

Clear win + clear we should move forward with zeromq

---

Benchmark script for reference:

MODEL="meta-llama/Meta-Llama-3-8B-Instruct"
TOTAL_SECONDS=60
QPS_RATES=("10")

for QPS in ${QPS_RATES[@]}; do
    NUM_PROMPTS=$((TOTAL_SECONDS * QPS))
    echo "===== RUNNING NUM_PROMPTS = $NUM_PROMPTS QPS = $QPS ====="

    python3 benchmarks/benchmark_serving.py \
        --model $MODEL \
        --dataset-name sonnet --sonnet-input-len 550 --sonnet-output-len 150 --dataset-path benchmarks/sonnet.txt \
        --num-prompts $NUM_PROMPTS --request-rate $QPS
done

---

FOLLOW UP WORK:

• Refactor how logits processors are handled. They should not belong to the API layer anymore, rather below the engine. Currently there are hacks around making it serializable.
• Switch to using a protobuf communication layer for the RPC Server
• Move detokenization to the OAI layer
• Profile using N uvicorn processes + expose a flag
• Support embedding models via RPC

[robertgshaw2-neuralmagic]

@joerunde why are these pass?

[robertgshaw2-neuralmagic]

LGTM over all. This is a good direction. Can we check one edge case?

• Send request generate 2K toks
• Drop the connection
• Check the generate coroutine is properly cancelled, and the abort it called, and there is no leakage of coroutines in both client and server (by printing out something...)

Do you want this in automation?

[robertgshaw2-neuralmagic]

TPOT speedup:

QPS	2	6	10	14
main	14.55	37.02	72.31	82.56
pr	13.34	31.88	61.03	75.18

Nice win + more to go in the follow ups (especially protobufs)

[simon-mo]

[youkaichao]

I'm trying to understand this. We create a socket every time we call this function?

[robertgshaw2-neuralmagic]

Yes. I believe this is a common paradigm for zeromq based clients

also, this function is mostly used in setup

[robertgshaw2-neuralmagic]

Also, this does not mean a new Unix socket is created. This is a zeromq socket. how zeromq handles this internally is opaque to us

[youkaichao]

why do we need a new env var here? isn't get_open_port() enough?

[robertgshaw2-neuralmagic]

I would be open to reverting this. Will just need to check how it interacts with tp

[youkaichao]

why do we copy this file instead of adding several new lines to add a new arg for the server, like @pytest.mark.parametrize ?

[robertgshaw2-neuralmagic]

Because we need to modify default server args which is a fixture for the whole file so unfortunately I think this has to be a separate file

[youkaichao]

Thanks for the great work! I just left some nit comments that I don't understand. If you have time, could you please briefly draw some diagrams to show how many processes do we have use and how do they interact with each other?