Update mlx_worker to be async
diff --git a/fastchat/serve/mlx_worker.py b/fastchat/serve/mlx_worker.py
index bff5820ade..795529104e 100644
--- a/fastchat/serve/mlx_worker.py
+++ b/fastchat/serve/mlx_worker.py
@@ -1,13 +1,12 @@
 """
 A model worker using Apple MLX
 
-docs/mlx_integration.md
-
 https://github.com/ml-explore/mlx-examples/tree/main/llms
 
 Code based on vllm_worker https://github.com/lm-sys/FastChat/blob/main/fastchat/serve/vllm_worker.py
 
 You must install MLX python:
+
 pip install mlx-lm
 """
 
@@ -19,6 +18,7 @@
 import uuid
 
 from fastapi import FastAPI, Request, BackgroundTasks
+from fastapi.concurrency import run_in_threadpool
 from fastapi.responses import StreamingResponse, JSONResponse
 import uvicorn
 
@@ -119,10 +119,12 @@ async def generate_stream(self, params):
 
         finish_reason = "length"
 
-        for token, _ in zip(
-            generate_step(context_mlx, self.mlx_model, temperature),
-            range(max_new_tokens),
-        ):
+        iterator = await run_in_threadpool(
+            generate_step, context_mlx, self.mlx_model, temperature
+        )
+
+        for i in range(max_new_tokens):
+            token = await run_in_threadpool(next, iterator)
             if token == self.mlx_tokenizer.eos_token_id:
                 finish_reason = "stop"
                 break

Generate tokens asynchronously (using run_in_threadpool) to avoid blocking FastAPI thread.

<!-- Thank you for your contribution! -->

<!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. -->

## Why are these changes needed?

<!-- Please give a short summary of the change and the problem this solves. -->

## Related issue number (if applicable)

<!-- For example: "Closes #1234" -->

## Checks

- [x] I've run `format.sh` to lint the changes in this PR.
- [x] I've included any doc changes needed.
- [x] I've made sure the relevant tests are passing (if applicable).
