Fix starvation with async server and interleaving optimization

[JoeZijunZhou]

• Make server async to handle high throughput from client (the number of requests a sync server can handle is limited by server side thread number, since each thread handles a request in blocking manner).
• Add AsyncMultiFuture data structure for the retuern_channel to stream response tokens asynchronously (threadsafe ensures no token drop in async server)
• Using blocking fixed size queue to block and yield threads efficiently
  • blocking for get and put operations for generate and detokenize queue
  • blocking for prefill queue get operation
  • unblocking for prefill queue put operation (unbound queue size to enqueue all requests from client)
• Optimized interleaving prefill, insert, and generate
  • When decode slots are empty, yield to prefill, then insert the prefill results as many as possible to the decode slots
  • When decode slots are full, block insert and prefill, keep decoding and detokenizing until some deocde completes and returns available slot.
  • generate queue size set to 3/1 decode batch size
    • get decode slots saturated fast (larger generate queue size ensures it has enough requests to insert to slots)
    • avoid OOM when both decode slots and generate queue are both saturated (can't set generate queue size too large, otherwise too many prefill results enqueued in generate queue)

[rwitten]

This CR cannot merge. It is driving contention (probably GIL) causing stalls.
https://xprof.corp.google.com/trace_viewer/rwitten-7572422995686954905?host_index=0

[rwitten]

We cannot merge these changes because they drive a duty cycle migration.

[vipannalla]

Where is this module used? I can't find any references.

[JoeZijunZhou]

In orchestrator, the return_channel, https://github.com/google/JetStream/pull/13/files/09e0e659ba821e14c5595dbfb3c8ec620725d865#diff-753df909960695fd6aa61b95c4762c4edc50a298007d40de7096a7e49e0e1215R130

[vipannalla]

Can you please add a xprof with the new changes so we can compare against the GIL-contention xprof that Rafi provided and rule-out any contention issues?

[JoeZijunZhou]

Can you please add a xprof with the new changes so we can compare against the GIL-contention xprof that Rafi provided and rule-out any contention issues?

Added a xprof in https://docs.google.com/document/d/1PWje3lup0ZHjtgbcWQbK8r1nNwgKE2rHBdVLbWU2aas/edit?tab=t.0#heading=h.f39bxdhl45e2 (Since here is public, just shared in the internal doc).

[vipannalla]

Thanks for stacktrace, fix looks good. I had a quick question about detokenize queue

[vipannalla]

How did you comeup with 8? Is that hardware dependent -- will this work for all configs (v5e-8, v5e-1 etc?)

[JoeZijunZhou]

The JET has an investigation on GIL contention issue and detokenize queue size, and figured out that this size can avoid synchronization issue. It's not dependent to hardware.

[vipannalla]

LGTM

Issue resolved and code reviewed.