[Core] Optimize Async + Multi-step (#8050)

2024-09-03 14:50:29 -04:00
parent 95a178f861
commit 6d646d08a2
8 changed files with 325 additions and 247 deletions
--- a/tests/multi_step/test_correctness_async_llm.py
+++ b/tests/multi_step/test_correctness_async_llm.py
@ -103,13 +103,13 @@ async def test_multi_step(
        model,
        server_args + distributed_args,
        num_logprobs,
-        max_wait_seconds=3 * 240)
+        max_wait_seconds=5 * 240)
    test_completions = await completions_with_server_args(
        prompts,
        model,
        ms_server_args + distributed_args,
        num_logprobs,
-        max_wait_seconds=3 * 240)
+        max_wait_seconds=5 * 240)

    # Assert multi-step scheduling produces identical tokens
    # to single-step scheduling.