[Do not merge][Async] Asynchronous DP coordination

[MatthewBonanni]

Purpose

When using async scheduling with MTP and DP, _get_valid_sampled_token_count causes a sync point. Following this sync point, DP coordination causes a GPU bubble. This PR creates and employs an asynchronous replacement for coordinate_batch_across_dp, allowing this communication to overlap with execution.

Test Plan

vllm serve deepseek-ai/DeepSeek-R1 \
    -dp 8 \
    --enable-expert-parallel \
    --no-enable-prefix-caching \
    --trust-remote-code \
    --max-model-len 4096 \
    --async-scheduling \
    --speculative-config '{"method": "mtp", "num_speculative_tokens": 1}' \
    --gpu-memory-utilization 0.8

with

vllm bench serve --model deepseek-ai/DeepSeek-R1 \
    --dataset-name random \
    --random-input-len 2 \
    --random-output-len 512 \
    --max-concurrency 2048 \
    --request-rate "inf" --num-prompts 8192

Test Result

Correctness

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9553|±  |0.0057|
|     |       |strict-match    |     5|exact_match|↑  |0.9545|±  |0.0057|

Performance

mean TPOT: 268.02 ms -> 260.72 = 2.7% speedup

Main:

============ Serving Benchmark Result ============
Successful requests:                     8192      
Failed requests:                         0         
Maximum request concurrency:             2048      
Benchmark duration (s):                  694.74    
Total input tokens:                      8192      
Total generated tokens:                  4194304   
Request throughput (req/s):              11.79     
Output token throughput (tok/s):         6037.25   
Peak output token throughput (tok/s):    27887.00  
Peak concurrent requests:                2200.00   
Total token throughput (tok/s):          6049.05   
---------------Time to First Token----------------
Mean TTFT (ms):                          20707.70  
Median TTFT (ms):                        3599.56   
P99 TTFT (ms):                           104962.31 
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          268.02    
Median TPOT (ms):                        271.88    
P99 TPOT (ms):                           497.44    
---------------Inter-token Latency----------------
Mean ITL (ms):                           545.86    
Median ITL (ms):                         237.47    
P99 ITL (ms):                            1079.73   
==================================================

PR:

============ Serving Benchmark Result ============
Successful requests:                     8192      
Failed requests:                         0         
Maximum request concurrency:             2048      
Benchmark duration (s):                  692.14    
Total input tokens:                      8192      
Total generated tokens:                  4194304   
Request throughput (req/s):              11.84     
Output token throughput (tok/s):         6059.89   
Peak output token throughput (tok/s):    10913.00  
Peak concurrent requests:                2086.00   
Total token throughput (tok/s):          6071.73   
---------------Time to First Token----------------
Mean TTFT (ms):                          24249.12  
Median TTFT (ms):                        6028.87   
P99 TTFT (ms):                           109039.20 
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          260.72    
Median TPOT (ms):                        263.83    
P99 TPOT (ms):                           449.06    
---------------Inter-token Latency----------------
Mean ITL (ms):                           484.87    
Median ITL (ms):                         240.19    
P99 ITL (ms):                            841.68    
==================================================

Profile

Main: Bubble evident

PR: No bubble

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request introduces asynchronous data parallelism (DP) coordination to overlap communication with computation, which is a solid approach to improve performance by reducing GPU idle time. The changes are well-structured, introducing start_dp_coordination_async and finish_dp_coordination functions to manage the asynchronous operation. However, I've identified two critical issues. One is a faulty assertion in dp_utils.py that could lead to runtime crashes when DP ranks have conflicting configurations. The other is a bug in gpu_model_runner.py where a padded token count is used for a check that expects an unpadded count, potentially leading to incorrect behavior. Addressing these issues is crucial for the correctness and stability of this feature.

[MatthewBonanni]

Holding off on this for now to evaluate a different approach

[MatthewBonanni]

Closing, superseded by #32951