Performance optimized interleaved mode JetStream server

[JoeZijunZhou]

• Optimized TPU duty cycle (largest gap < 4ms)
• Optimized TTFT: dispatch prefill tasks ASAP w/o unnecessary blocking in CPU, keep backpressure to enforce insert ASAP, return first token ASAP.
• Optimized TPOT: properly enforce generate and detokenize task in sequential w/o unnecessary blocking in CPU.
• Optimized output token throughput: properly prioritize prefill and balancing TTFT and decode in high throughput situation.
• Tested with llama2-70b JetStream MaxText server on v5e-8 VM

[FanhaiLu1]

• Optimized TPU duty cycle (largest gap < 4ms)
• Optimized TTFT: dispatch prefill tasks ASAP w/o unnecessary blocking in CPU, keep backpressure to enforce insert ASAP, return first token ASAP.
• Optimized TPOT: properly enforce generate and detokenize task in sequential w/o unnecessary blocking in CPU.
• Optimized output token throughput: properly prioritize prefill and balancing TTFT and decode in high throughput situation.
• Tested with llama2-70b JetStream MaxText server on v5e-8 VM

Optimized TTFT and Optimized output token throughput are conflicted with each. Can we expose some parameter to tuning the two part?

[FanhaiLu1]

Can you elaborate more on why there is synchronization issue after 8 steps?

[JoeZijunZhou]

Set it to 8 as the detokenize thread. Too large or too small will cause performance issue.

[FanhaiLu1]

Thanks, is this PR ready to submit?

[FanhaiLu1]

I though we already had prefill detokenize thread. Do current jetstream (before this pr) always return prefill token (fist token) after first decode step?

[JoeZijunZhou]

We only had detokenize thread that combined prefill detokenize and decode detokenize. The problem is that we have jax.block_until_ready() blocking the thread waiting for the prefill token or decode token copy to host async, so putting them in 1 thread would make the TTFT slow. JetStream returns prefill token in prefill thread (after prefill step generating the first token).

[FanhaiLu1]

sounds good, thanks for sharing insights!

[JoeZijunZhou]

• Optimized TPU duty cycle (largest gap < 4ms)
• Optimized TTFT: dispatch prefill tasks ASAP w/o unnecessary blocking in CPU, keep backpressure to enforce insert ASAP, return first token ASAP.
• Optimized TPOT: properly enforce generate and detokenize task in sequential w/o unnecessary blocking in CPU.
• Optimized output token throughput: properly prioritize prefill and balancing TTFT and decode in high throughput situation.
• Tested with llama2-70b JetStream MaxText server on v5e-8 VM

Optimized TTFT and Optimized output token throughput are conflicted with each. Can we expose some parameter to tuning the two part?

Currently, prioritize prefills in interleaved mode, and apply correct JAX blocking for copy to host async to reduce wasted wait time. 1 more optimization to do is to ensure the result returns immediately when the return channel has the result (from orchestrator).