[Bugfix] Separate target and draft scheduling budgets

[slippersss]

Purpose

This PR aims to fix #52922. The idea is simply to track target and draft token budgets separately, since these two stages may process different numbers of tokens for the same request.

We define draft input cost model for each request as below:

num_draft_input_tokens = extend_target_batch * num_target_tokens + fixed_overhead

Here, extend_target_batch indicates whether the draft input batch extends the scheduled target batch and fixed_overhead represents the additional per-request draft input size.

Take mtp/eagle3/dflash/dspark/peagle/pard as examples:

• For mtp/eagle3, at draft step 0, the draft will forward as much tokens as the target without any other cost. So extend_target_batch=True and fixed_overhead=0. At later draft steps, the draft will forward just one token per request, which does not exceed first step batch size.
• For dflash/dspark, the draft will forward anchor/bonus and mask tokens, independent of the target. So extend_target_batch=False and fixed_overhead=num_speculative_tokens+offset.
• For peagle/pard, the draft will forward both target tokens and mask tokens. So extend_target_batch=True and fixed_overhead=num_speculative_tokens+offset.

Note that offset is -1 or 0 or +1, depending on how each method arranges anchor/bonus and mask tokens. Detailed values have been listed in speculative config.

This PR is still a work in progress. Feedbacks and suggestions are welcome.

Test Plan

• serve script

vllm serve Qwen/Qwen3-8B \
    --host 127.0.0.1 \
    --port 60006 \
    --gpu-memory-utilization 0.95 \
    --tensor-parallel-size 2 \
    --max-num-seqs 16 \
    --max-num-batched-tokens 128 \
    --max-model-len 4096 \
    --speculative-config '{
        "method":"dspark",
        "num_speculative_tokens":7,
        "model":"deepseek-ai/dspark_qwen3_8b_block7"
    }'

• bench script

vllm bench serve \
    --backend vllm \
    --model Qwen/Qwen3-8B \
    --host 127.0.0.1 \
    --port 60006 \
    --seed 0 \
    --dataset-name random \
    --num-prompts 4096 \
    --random-input-len 64 \
    --random-output-len 512 \
    --ignore-eos \
    --max-concurrency 4096

Test Result

• without this PR

============ Serving Benchmark Result ============
Successful requests:                     4096      
Failed requests:                         0         
Maximum request concurrency:             4096      
Benchmark duration (s):                  716.16    
Total input tokens:                      262144    
Total generated tokens:                  2097152   
Request throughput (req/s):              5.72      
Output token throughput (tok/s):         2928.34   
Peak output token throughput (tok/s):    975.00    
Peak concurrent requests:                4096.00   
Total token throughput (tok/s):          3294.38   
---------------Time to First Token----------------
Mean TTFT (ms):                          359777.36 
Median TTFT (ms):                        361221.78 
P99 TTFT (ms):                           707431.84 
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          2.92      
Median TPOT (ms):                        2.75      
P99 TPOT (ms):                           9.37      
---------------Inter-token Latency----------------
Mean ITL (ms):                           9.38      
Median ITL (ms):                         9.24      
P99 ITL (ms):                            10.76     
---------------Speculative Decoding---------------
Acceptance rate (%):                     31.85     
Acceptance length:                       3.23      
Drafts:                                  651352    
Draft tokens:                            4559464   
Accepted tokens:                         1451963   
Per-position acceptance (%):
  Position 0:                            62.35     
  Position 1:                            45.30     
  Position 2:                            33.04     
  Position 3:                            26.10     
  Position 4:                            22.14     
  Position 5:                            17.98     
  Position 6:                            16.01     
==================================================

• with this PR

============ Serving Benchmark Result ============
Successful requests:                     4096      
Failed requests:                         0         
Maximum request concurrency:             4096      
Benchmark duration (s):                  586.69    
Total input tokens:                      262144    
Total generated tokens:                  2097152   
Request throughput (req/s):              6.98      
Output token throughput (tok/s):         3574.55   
Peak output token throughput (tok/s):    1172.00   
Peak concurrent requests:                4096.00   
Total token throughput (tok/s):          4021.36   
---------------Time to First Token----------------
Mean TTFT (ms):                          295025.49 
Median TTFT (ms):                        295759.69 
P99 TTFT (ms):                           577136.67 
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          4.20      
Median TPOT (ms):                        3.89      
P99 TPOT (ms):                           13.54     
---------------Inter-token Latency----------------
Mean ITL (ms):                           13.53     
Median ITL (ms):                         13.53     
P99 ITL (ms):                            14.85     
---------------Speculative Decoding---------------
Acceptance rate (%):                     31.96     
Acceptance length:                       3.24      
Drafts:                                  649766    
Draft tokens:                            4548362   
Accepted tokens:                         1453501   
Per-position acceptance (%):
  Position 0:                            62.02     
  Position 1:                            45.15     
  Position 2:                            33.16     
  Position 3:                            26.35     
  Position 4:                            22.32     
  Position 5:                            18.30     
  Position 6:                            16.40     
==================================================

By removing the overly conservative input budget constraint in scheduling, DSpark can reach the concurrency configured by the user, improving output throughput from 2928.34 to 3574.55 tok/s (+22.1%).

---

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.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[drslark]

I think it is useful.